Classic Lost Cities: three-round agent, and borealis in the browser

The trainer was learning a game nobody plays. Lost Cities is decided on three
rounds summed; the stack optimised a single deal. This branch builds the real
game, trains an agent on it, and ships it.

borealis beats altair -- the league policy that was on the web -- 0.6094 at the
real three-round game over 8192 duplicate matches (+22.0 points), at a matched
compute budget. It is also harder to farm: a from-scratch exploiter funded to the
same 131M reaches 0.6295 against altair and only 0.4657 against borealis.

What actually earned that, measured by ablation at a matched budget rather than
asserted:

  both-seat training   +17.1 points.  Self-play already played the opponent's
                       plies with the same network and the old trainer
                       stop_gradiented them away. Free doubling of the data.
  in-scan auto-reset   5.8x the learner actions per update at identical compute:
                       87% of every rollout was spent stepping already-done envs.
                       Exposed a latent zero-bootstrap in compute_gae.
  linear total reward  Beats tanh(total/scale) 0.5859. The two share an objective;
                       the linear one just pays it out every ply instead of once
                       per ~160, and risk attitude turns out to be worth under a
                       win-rate point in this game -- a policy told to gamble when
                       behind *loses* to a greedy clone.
  match observation    ~2.5 points. Mostly single-round defects: the observation
                       never said whose turn it was, and divided the score
                       difference by 780.
  privileged critic    +3.7 points, but only at scale -- ablated at 39.3M it
                       measured *negative*. An ablation run at a budget you do not
                       intend to ship can invert.

Also: rounds are independent (corr 0.004), so carry earns its place in the
observation through the start-player rule, not the objective.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh
This commit is contained in:
2026-07-15 07:10:07 +09:00
co-authored by Claude Opus 4.8
40 changed files with 4649 additions and 193 deletions
+7 -1
View File
@@ -15,6 +15,9 @@ src/**/*.c
# Local session / lock files # Local session / lock files
.compute.lock .compute.lock
.claude/ .claude/
# ...except checked-in project skills, which are documentation for the next session.
!web/.claude/
!web/.claude/**
# Rust build output # Rust build output
target/ target/
@@ -29,10 +32,13 @@ tools/julia/
.ruff_cache .ruff_cache
# Web dependencies and locally exported models. The verified public browser # Web dependencies and locally exported models. The verified public browser
# policy below is the one exception: it is deliberately served as a static asset. # policies below are the exception: they are deliberately served as static assets,
# and the deploy builds straight from the repo, so an ignored model ships as a 404.
web/node_modules/ web/node_modules/
web/*.tsbuildinfo web/*.tsbuildinfo
web/public/models/*.onnx web/public/models/*.onnx
web/public/models/*.json web/public/models/*.json
!web/public/models/jax-ppo.onnx !web/public/models/jax-ppo.onnx
!web/public/models/jax-ppo.json !web/public/models/jax-ppo.json
!web/public/models/borealis.onnx
!web/public/models/borealis.json
+50
View File
@@ -0,0 +1,50 @@
# Classic Lost Cities: three rounds, self-play.
#
# rollout_steps is deliberately shorter than a match. Early self-play stalls
# rounds badly (untrained matches run ~900 plies), and the env carries across
# updates, so a long match simply spans several rollouts -- the GAE truncation
# bootstrap keeps that unbiased. Sizing the scan to the worst-case match instead
# would blow up the rollout tensors, which are already doubled by training both
# seats and widened by the critic's 681-dim privileged view.
run:
experiment_name: match-selfplay
seed: 20260715
learner_seat: 0
total_updates: 300
log_every: 10
checkpoint_every: 50
artifact_root: runs/jax-ppo-match
opponent:
name: discard_only # unused: self-play
network:
hidden_size: 512
num_layers: 3
ppo:
batch_games: 512
rollout_steps: 256
gamma: 1.0
gae_lambda: 0.97 # a match is ~3x a round; 0.95 reaches too little of it
clip_epsilon: 0.2
entropy_coef: 0.01
value_coef: 0.5
max_grad_norm: 0.5
learning_rate: 0.0003
epochs: 4
minibatches: 32
reward:
terminal_scale: 50.0
# Counted in learner actions now, not padded scan steps. Held slightly above
# zero: the match-terminal signal alone is one bounded number per ~160 plies.
potential_shaping_initial: 1.0
potential_shaping_final: 0.05
potential_shaping_anneal_steps: 20000000
evaluation:
games: 2000
duplicate: true
shuffle_bank_seed: 20260715
batch_games: 512
File diff suppressed because one or more lines are too long
+43
View File
@@ -0,0 +1,43 @@
{
"_scheme": {
"naming": "Astronomical names, alphabetically ordered. The first letter is the generation; a new letter means the observation space broke, not that the model got better.",
"rule": "A codename never encodes quality. The record this replaces stored 'FINAL PPO', which stops meaning anything the moment there is a second final model.",
"identity": "The hash is the truth -- it is what actually played. The codename is for humans, and it is assigned here, not derived. Records should carry the hash; look the codename up.",
"next": "cygnus, deneb, ..."
},
"models": {
"altair": {
"hash": "e8241e305c01",
"hash_kind": "sha256 of web/public/models/jax-ppo.onnx",
"generation": "a",
"game": "single round",
"observation_size": 454,
"hidden_size": 512,
"num_layers": 3,
"trained": "self-play league, 122.6M learner actions",
"source": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/latest",
"deployed": "web/public/models/jax-ppo.onnx",
"displayed_as": "WASM · FINAL PPO",
"note": "Every game in data/human-play/game-records.jsonl (format v1, 111 games, 2026-07-14) was played against this model. The v1 schema has no model field -- it stores the on-screen label -- so this line is the record of what they played."
},
"borealis": {
"hash": "4ae613b010ca",
"hash_kind": "sha256 over the orbax checkpoint files (no ONNX export yet)",
"generation": "b",
"game": "three-round match (classic rules)",
"observation_size": 501,
"critic_observation_size": 681,
"hidden_size": 512,
"num_layers": 3,
"trained": "self-play, 131.1M learner actions, linear total-score reward, both seats, privileged critic",
"source": "runs/jax-ppo-match/2026-07-15_031529_match-scaled/latest",
"deployed": null,
"results": {
"vs_altair_3round": "0.6094 win rate [0.599, 0.620], +22.0 points, 8192 duplicate matches",
"exploitability": "a from-scratch exploiter funded to 131M reaches 0.4657 against it, and 0.6295 against altair -- lower bound, neither exploiter had plateaued"
}
}
}
}
+617
View File
@@ -0,0 +1,617 @@
# Plan: 클래식 3라운드 로스트시티 에이전트
**Status:** Ready to execute (Fable 2차 검토 반영 완료)
**Priority:** High — 현재 스택은 **단판(1라운드)** 게임을 학습하는데, 원작 로스트시티는
**3라운드 합산**이 승패를 가른다. 즉 지금 에이전트는 애초에 다른 게임을 배우고 있다.
**Background:** Fable 설계 검토 2회 (2026-07-15), 인간 대 AI 기보
(`data/human-play/game-records.jsonl`, 계속 증가 중)
## 목표
**원작 규칙 그대로의 로스트시티에서 가장 강한 에이전트.** 레거시 코드 호환은 고려하지 않는다.
env/obs/보상/학습 루프 재작성 모두 허용.
## 확정된 규칙 (룰북 원문 검증 완료)
코스모스 공식 영문 룰북(691820-02) 및 Board Game Arena 구현체로 교차 확인:
- 3라운드를 두고 **누적 총점**이 높은 쪽이 매치 승리.
- 1라운드 선공: "가장 나이 많은 사람" → 게임 내적으로는 **임의**.
- **2·3라운드 선공: 누적 점수가 더 많은 쪽** ("The player who has more points begins.")
— 번갈아 가는 것이 아니다.
- **정확한 동점일 때의 선공은 룰북에 없다.** 공식 룰북·BGA 모두 침묵.
**무작위(동전 던지기)로 정한다.** 결정론적 규칙은 대칭 제로섬 셀프플레이에
자리(seat) 비대칭을 주입해 착취 가능한 구멍이 된다.
점수 계산(`(랭크합 20) × (1 + 악수) + (8장 이상 +20)`)과 "방금 버린 카드는 즉시
회수 불가"(`just_discarded`)는 **현재 엔진이 이미 정확하다**
(`engine.py:244`, `engine.py:124`).
### 선공이 중요한 이유: 덱 시계
라운드는 **덱의 마지막 카드를 뽑는 순간** 끝난다. 카드는 놓은 뒤에 뽑으므로 마지막
카드를 뽑은 쪽은 그 카드를 쓰지 못한다. 총 턴 수가 홀수면 **선공이 한 장 더 놓는다**.
총 턴 수 = 44 + (버림패 드로우 횟수)이므로 **플레이어가 홀짝을 조작할 수 있다**
(실측 게임당 버림패 드로우 ≈6.6회). 선공권이 누적 점수에 달려 있으므로, 2라운드
마진에는 **점수를 넘어 3라운드 선공권이라는 추가 가치**가 붙는다.
### 측정된 사실: stalling은 정적 상대 착취이지 균형이 아니다 (Phase 1 사이징의 근거)
버림패 드로우는 덱을 줄이지 않으므로, **양쪽이 계속 버림패만 뒤지면 라운드가 끝나지
않는다.** 기존 학습 로그를 보면 이 착취가 실제로 학습된다:
| 학습 체제 | 라운드 길이 추이 | 400수 상한 도달 |
|---|---|---|
| 정적 휴리스틱 상대 (`balanced.yaml`) | 82 → **116수** (eval greedy **168수**) | 0.4~0.7% |
| **셀프플레이 (league)** | 71 → **53.6수로 수렴** | **0%** |
정적 약체 상대에게는 턴을 늘릴수록 원정 깊이·8장 보너스에서 강자가 이득이라 질질 끄는
것이 합리적이다. 그러나 **셀프플레이에서는 상대도 똑같이 끌 수 있어 상쇄되고, 라운드가
자연 길이(≈54수 = 44 덱 드로우 + ~10 버림패 드로우)로 수렴한다.**
**결론:** 우리의 목표 체제는 셀프플레이/league이므로 **3라운드 매치 ≈ 160수**로 잡으면
되고, 기존 `rollout_steps=400` 안에 여유롭게 들어간다. 단 두 가지를 유의한다:
1. **정적 상대로 3라운드를 학습시키면 매치가 350~500수까지 늘어날 수 있다** — 정적 상대
워밍업 단계를 쓴다면 스캔 길이를 따로 잡아야 한다.
2. **라운드 상한(~120수)은 그 자체가 착취 가능한 구멍이다.** 앞선 쪽이 상한까지 끌면
라운드가 유리한 상태로 얼어붙는다. `max_steps_rate`를 게이트 지표로 계속 감시할 것
(셀프플레이에서는 0이어야 정상).
## 현재 코드의 결함 (Fable 2회 교차 검증 — 8건 전부 사실 확인)
1. **env가 1라운드짜리**`State`에 라운드/누적 점수 없음 (`types.py:46`),
`to_move=0` 하드코딩 (`engine.py:63`).
2. **롤아웃의 87%가 죽은 스텝** — 한 라운드는 실측 평균 50.6수(최대 56)인데
`rollout_steps = MAX_STEPS = 400` (`ppo.py:66`)이고, 스캔 **안에는 리셋이 없다**
(`reset_done_envs`는 업데이트 사이에서만, `ppo.py:275`).
3. **GAE 부트스트랩 초기값이 0** (`ppo.py:691`). 지금은 모든 에피소드가 스캔 안에서 끝나
무해하나, **in-scan 리셋을 켜는 순간 편향**이 된다.
4. **shaping 어닐링 회계 버그**`update * batch_games * rollout_steps` (`ppo.py:222`)로
**죽은 스텝까지 세어서**, 5M 어닐링이 250 업데이트 중 2번째에 끝난다. league는 shaping을
아예 끈다 (`league.py:592`).
**"shaping은 불필요하다"는 과거 결론이 있다면 무효다. 켜진 적이 없다.**
5. **obs에 `to_move`가 없다.** 그런데 value loss는 `active` 마스크로 학습되어 (`ppo.py:624`)
**상대 차례 상태에서도** 가치를 맞추라고 요구한다. critic이 차례를 구분할 단서가
`step_count / 400`뿐 — MLP에게 부동소수점에서 홀짝을 뽑으라는 요구다.
6. **점수차 정규화가 ÷780** (`MAX_ABS_SCORE`, `obs.py:68`). ±50점 차가 ±0.06 → 사실상
안 보인다.
7. **셀프플레이에서 상대 자리 결정을 전부 버린다** (`stop_gradient`, `ppo.py:425`) —
의사결정의 절반을 낭비.
8. **평가가 단판 승률** (`gates.py`) — 3라운드 개편의 성패를 측정할 자가 없다.
## 핵심 설계 결정
### PRNG: 모든 무작위성을 `reset()`에서 미리 뽑는다 (필수)
라운드 전환 리셔플과 동점 코인플립을 `step()` 안에서 샘플링하면 `step`이 확률적이 되어
서명이 바뀌고, 모든 rollout/eval/gates 바디에 키를 꿰어야 하며 duplicate 미러링이 꼬인다.
**`reset()`에서 3라운드 덱 순서 전부(3×60), 동점 코인플립 비트 2개, 1라운드 선공 비트
1개를 미리 샘플링해 `State`에 저장한다.** `step()`은 결정론을 유지하고 라운드 전환은
다음 덱 슬라이스로 스위치만 한다. 덤: **미러 매치(같은 3딜 + 자리 교대 + 같은 코인플립
비트)가 공짜로 따라온다** — Phase 5의 antithetic 페어링이 그대로 성립.
### 라운드 전환 vs 매치 리셋의 분리
- **라운드 전환은 엔진 `step()` 내부**에서 (`done=False` 유지, carry 갱신, 덱 슬라이스 전환)
- **매치 done 리셋은 rollout 바디**에서 (in-scan auto-reset)
이러면 두 메커니즘이 깨끗이 분리된다. 단 위의 pre-sampled 덱 설계가 전제다.
- `State`**라운드 스텝 카운터와 매치 스텝 카운터를 둘 다** 둔다 (forced-done은 라운드
단위, obs 정규화는 라운드 상대 진행도 + round 원-핫).
- 중간 라운드 forced-done(라운드 상한 초과) 시 **자연 종료와 동일하게 라운드 종료 → 전환**.
### 보상 스케일에 대한 해소 (혼동 방지)
`sign()`에 가까운 작은 scale은 **틀린 선택**이다 — 150수에 걸쳐 모든 수가 동일한 ±1을
받아 크레딧 할당이 전부 critic에 떠넘겨지고, 5점 차 패배와 80점 차 패배의 그래디언트가
같아진다. **scale 30~50의 tanh가 절충점**이다: |마진| ≳ 60에서는 `sign`을 근사하면서
크레딧 그래디언트를 보존한다. 대신 5점 승과 80점 승을 여전히 구분하므로 순수 승률
최적화 대비 **마진 쪽으로 약간 왜곡**된다 — 이 잔여 왜곡은 Phase 5의 후기 fine-tune으로
scale을 낮춰서(예: 50 → 25) 제거한다.
명시적 risk 항은 **넣지 않는다.** carry + round를 obs에 넣은 종료 보상이 올바른 리스크
태도(뒤지면 도박, 앞서면 잠금)를 자동으로 유도한다.
## 체크리스트
### Phase 0a — 학습 루프 수정 (최우선, 단판 체제에서 검증 가능)
- [x] **GAE 부트스트랩 수정**: 잘린 에피소드를 `V(s_final)`로 부트스트랩 (`ppo.py:691`).
- [x] **in-scan auto-reset**: 롤아웃 스캔 안에서 done env 리셋 → 죽은 스텝 제거.
**실측(rollout_steps=400, batch_games=64, 3 업데이트):** main은 활성 17.3%
(learner 액션 2,221/업데이트), 수정 후 **활성 100% (12,842/업데이트) — 동일 연산량에
샘플 5.8배.** 활성 비율은 게임 길이에 비례하므로, 학습 초기 정책(~69수)에서 5.8배,
학습된 정책(~50수)이면 8배에 가까워진다. 3라운드 매치(~150수)가 되면 약 2.7배로
줄지만 그때는 스캔 전체가 실제 매치로 채워진다.
- [x] **메트릭 파이프라인 재작업 (필수, 놓치면 조용히 오염됨)**: `rollout_metrics`
`final_env`에서 통계를 읽고 `episode_return = jnp.sum(transitions.reward, axis=0)`
**env당 에피소드 1개를 가정**한다 (`ppo.py:1069`). in-scan 리셋 후에는 final_env가
에피소드 중간이고 여러 에피소드의 보상이 합산된다 → **metrics.jsonl 전체가 쓰레기가
된다.** done 경계에서 누적하는 에피소드 단위 집계로 바꿀 것.
- [x] **league assignments를 scan carry로**: `learner_seat`/`opponent_index`/`use_mirror`
업데이트 사이에서만 재샘플링된다 (`ppo.py:568`). in-scan 리셋을 켜면 리셋 시점에
스캔 내부에서 재샘플링해야 한다 — 안 하면 자리/상대가 에피소드 간 고정되어 편향.
- [x] **shaping 어닐링 회계 수정**: 패딩된 스캔 스텝이 아니라 **learner 액션 수** 기준
(`ppo.py:222`). `shaping_coef`는 호스트에서 계산해 jit 함수에 넘기므로
(`ppo.py:221`), 직전 업데이트의 액션 수를 호스트에서 누적하는 **1-업데이트 지연**
구조가 된다 (무해).
- [x] **회귀 게이트 — 통과 (단, 게이트 설계를 고쳐야 했다).**
**(1) "동일 learner-액션 예산"은 잘못된 게이트였다.** PPO 진척도는 샘플 수와
**옵티마이저 스텝 수** 둘 다에 달려 있다. 샘플을 맞추면(64 vs 250 업데이트) 새 코드는
그래디언트 스텝이 1/4이 되어 덜 수렴한다. 업데이트당 연산량은 동일하므로
**올바른 게이트는 250 대 250**(같은 컴퓨트·같은 그래디언트 스텝, 샘플만 5.8배)이다.
**(2) 동일 컴퓨트 게이트 결과 (2만판 duplicate eval):**
| 앵커 | 지표 | 베이스라인 | Phase 0a |
|---|---|---|---|
| **expert** (강함) | 승률 | 0.4525 | **0.5071** |
| **expert** | 평균 점수차 | 4.28 | **0.56** |
| balanced (약함) | 승률 | 0.9866 | 0.9292 |
| balanced | 평균 점수차 | 116.8 | **128.1** |
학습 곡선은 **모든 업데이트 지점에서** 새 코드가 위(최종 return 0.963 vs 0.935).
**(3) balanced 승률 하락은 전부 `MAX_STEPS=400` 아티팩트다.** 새 정책 2,048판 분해:
| | 판수 | 패배율 |
|---|---|---|
| 자연 종료 | 1,214 | **0.08%** |
| 400수 상한 도달 | 834 (40.7%) | **20.62%** |
**전체 패배의 99.4%가 상한에 부딪힌 판에서 발생.** 새 정책은 자연 종료 게임에서
**99.92% 승률**이다. 샘플이 5.8배라 stalling 착취를 더 깊이 배웠고(greedy 225수),
그 결과 40%가 인위적 벽에 박아 원정 미완성인 채 얼어붙는다.
**결론:** Phase 0a는 성공. expert(=stalling을 허용하지 않는 앵커)에서 45%→51%로
실제 실력이 올랐다.
### Phase 1 — 3라운드 env
**구현: `src/lost_cities_jax/match.py` (커밋 `7cd299d`).** 단판 엔진을 고치지 않고 **감쌌다**
`engine.py`는 TS 클라이언트와의 차분 테스트가 지키는 규칙 오라클이고, 라운드 자체는
매치가 생겨도 변하지 않는다. `MatchState`가 내부에 단판 `State`를 든다.
- [x] `MatchState`: `round`(단판 State), `deck_orders`(3×60), `coin_flips`(3,),
`round_idx`, `carry`(2, 플레이어별 누적 점수), `done`.
별도 매치 스텝 카운터는 두지 않았다 — 매치는 3라운드가 끝나면 종료하므로 종료 판정에
불필요하고, 라운드 진행도는 내부 `State.step_count`(라운드마다 리셋)가 이미 준다.
- [x] `match_step()`: 덱 소진 시 라운드 < 3이면 carry 갱신 → 다음 덱 슬라이스 → `round_idx += 1`,
`done` 유지. 라운드 == 3이면 `done = True`. **보상은 3라운드 끝에서만.**
- [x] **선공 규칙**: `carry[0] > carry[1]` → p0, `<` → p1, `==` → pre-sampled 코인플립.
1라운드는 특수 케이스가 **필요 없다** — carry가 (0,0)이라 동점 분기가 자동으로 코인을
뽑고, 그게 룰북의 "가장 나이 많은 사람"(= 임의)과 정확히 같다.
- [x] **PRNG**: 딜 3벌 + 코인 3개를 `match_reset`에서 전부 미리 뽑아 상태에 저장.
`match_step`은 결정론적이라 키를 rollout/eval/gates에 꿸 필요가 없고, **미러 매치가
자리 라벨 교체만으로 성립**한다 (Phase 5 antithetic 페어링이 여기 의존).
- [x] ~~라운드당 스텝 상한을 400 → ~120으로 분리~~ — **철회한다. 라운드 상한은 400을 유지.**
**이 항목의 원래 명분이 사라졌다.** 근거는 "3라운드(~360수)가 400스텝 스캔에 들어가게"
였는데, **Phase 0a의 GAE 절단 부트스트랩이 스캔 길이 제약을 없앴다.** 잘린 매치는
올바르게 부트스트랩되므로 스캔 길이는 이제 자유 파라미터다.
**그리고 상한을 낮추면 아티팩트가 악화된다.** Phase 0a 실측: 400수 벽에 부딪힌 판의
20.6%가 패배로 뒤집혔고 **전체 패배의 99.4%가 거기서 발생**했다. 상한을 120으로 내리면
**더 많은 판이 벽에 박고**, "앞선 쪽이 상한까지 끌어 라운드를 얼린다"는 착취가 훨씬
쉬워진다. 셀프플레이에서는 라운드가 54수로 수렴해 120이든 400이든 안 걸리므로,
**덜 걸리는 쪽을 두고 `max_steps_rate`를 감시**하는 것이 맞다.
- [ ] **학습 파이프라인 배선** (원래 계획서에서 누락된 항목): `ppo.py`의 rollout/eval,
`gates.py`, `league.py`가 전부 단판 `State`를 받는다. `MatchState`를 받도록 배선한다.
- 매치용 obs 래퍼: 우선 `observation(state.round, player)`로 carry-blind하게 연결
(carry/round 피처는 Phase 3에서 추가).
- `match_legal_action_mask`, `match_step`, `match_score`로 교체.
- 스캔 길이는 자유 파라미터가 되었으므로 셀프플레이 매치 길이(~160수) 기준으로 잡되,
잘려도 부트스트랩이 처리한다.
- [ ] **league 풀 재구축**: 기존 스냅샷 멤버는 obs 변경으로 전부 무효화된다 — Phase 3 이후로
미룬다 (obs가 확정되기 전에 풀을 다시 채우면 두 번 일한다).
- [x] **테스트 12개** (`tests/lost_cities_jax/test_match.py`): 라운드당 덱 드로우 = 44 +
버림패 드로우만큼 연장, `carry`가 각 라운드를 정확히 한 번씩 적립, 선공 규칙 3분기,
1라운드 코인 공정성(≈50/50), 라운드 경계에서 누적 점수 불연속 없음(PBRS 전제),
미러 매치 대칭성, 보상은 3라운드 끝에서만 지급.
### Phase 0b — 측정 자 (첫 장기 3라운드 학습 **전에** 착지)
- [ ] **매치 단위 평가**: 3딜 전부 미러링 + 자리 교대 + **동일 코인플립 비트**,
매치 승률 + Wilson CI, 평균 총 마진. `gates.py`/`league.py`의 단판 기준을 대체.
**`gates.py:884`, `gates.py:961``MAX_STEPS` 길이 단판 스캔 2곳 포팅 포함** —
exploiter 학습·평가 경로 전체가 3라운드로 가야 성공 기준 3을 잴 수 있다.
- [ ] **carry 조건부 프로브**: carry ∈ {60, 25, 1, +1, +25, +60}을 주입한 3라운드 시작
위치에서 승률·행동 변화(원정 개수, 악수 비율, 덱 레이스 비율) 측정.
- [ ] **shuffle bank 포맷 확장**: 덱과 함께 코인플립·선공 비트도 뽑도록. 안 그러면
`jax.random.PRNGKey(0)` 고정 eval(`ppo.py:1025`)의 재현성이 매치 정의와 얽힌다.
- [ ] **legacy obs 버전 보존**: 구 정책을 3라운드 매치에 투입해 베이스라인을 재려면 구
obs(454차원)로 추론해야 하는데, `checkpoint_policy_from_params`는 전역 `observation`
호출한다 (`ppo.py:802`) → obs 개편 후 구 체크포인트는 **로드 자체가 실패**한다.
policy 로더가 체크포인트별 obs 버전을 받도록 할 것. **성공 기준 2의 분모가 여기 달렸다.**
### Phase 2 — 보상
- [ ] 종료 보상: 3라운드 끝에서만 `tanh(총_점수차 / scale)`, **scale = 30~50**
(근거는 위 "보상 스케일에 대한 해소" 절).
- [ ] 마진 shaping 부활: `Φ(s) = carry + 현재 보드 점수차`, **작은 계수**(종료 보상 스케일의
0.05~0.2). 현재의 계수 1.0 원점수 shaping은 ±1 종료 보상보다 10~30배 크다.
learner 스텝 기준으로 어닐링하되 **후반까지 정확히 0으로 내리지 않는다.**
- [ ] γ=1.0에서 `Φ(s') Φ(s)`**올바른 PBRS 형태다** (γ 누락 아님 — 검토에서 확인).
### Phase 3 — 관측
- [ ] `carry`: **÷75 스칼라 + 구간 원-핫**(약 9구간). 3라운드 정책은 "1점만 더" 임계값
근처에서 급격히 꺾여야 한다. 기존 `score_diff`의 ÷780 정규화도 같이 고친다.
- [ ] `round_idx` 원-핫 + 남은 라운드 수.
- [ ] **`to_move` 비트** + "내가 이번 라운드 선공인가" + "현재 홀짝에서 마지막 덱 카드를
누가 뽑는가"(덱 시계).
- [ ] 색깔별 **살아있는 점수 3분할**: 내 `col_top` 위로 아직 나올 수 있는 점수를
(a) **내 손패**, (b) **버림패 더미**(공개돼 있고 회수 가능 — 빠뜨리기 쉬움),
(c) **미공개**(덱 ∪ 상대 은닉 손패)로 나눠 넣는다. 상대에 대해서도 동일
(상대 `col_top`은 공개). MLP가 7×60 채널에서 뽑아내기 어려운 비선형 집계이고,
모든 개시/연장/차단 판단을 좌우한다.
### Phase 4 — 학습 효율
- [ ] **전지적 critic (CTDE)**: critic에만 상대 손패 + 덱 구성을 준다. 행동과 무관한
정보이므로 정책 그래디언트를 편향시키지 않는다. 딜 운 분산을 정면으로 깎는다.
**가치 경로를 정책 트렁크에서 분리해야 한다** (현재 공유 트렁크, `ppo.py:104`) —
특권 정보가 정책 로짓으로 새면 안 된다.
→ 이 critic은 나중에 **PIMC 탐색의 리프 평가기로 그대로 재활용**된다.
- [ ] **양쪽 자리 학습**: `stop_gradient`된 상대 자리 전이(`ppo.py:425`)도 학습에 쓴다
(샘플 효율 2배). 같은 게임의 두 자리는 반상관이므로 같은 배치에 두고 advantage
정규화에 맡긴다.
- [ ] `gae_lambda` 재검토: 0.95는 150수 지평에서 너무 짧다(중반 수 직접 가중치 0.02).
전지적 critic이 있으면 유지, 없으면 0.97~0.99. **λ=1은 금지**(딜 분산).
- [ ] `entropy_coef` **스윕으로 재결정**. (주의: "0.01이 원점수 shaping 기준으로 잡혔다"는
추론은 **틀렸다** — league는 shaping을 끄고 학습했으므로 0.01은 이미 ±1 tanh 체제에서
동작해온 값이다. 다만 3라운드에서 보상 빈도가 1/150로 희석되고 작은 shaping이
추가되므로 재튜닝 자체는 타당하다. **근거 없이 10배 낮추면 과소탐색으로 직행한다.**)
### Phase 5 (나중) — 최종 강함
- [ ] **후기 fine-tune**: 학습 말미에 종료 보상 scale을 낮춰(50 → 25) 마진 왜곡을 제거하고
순수 승률 쪽으로 당긴다.
- [ ] 페어드 antithetic 딜(같은 3딜 + 자리 교대 + 같은 코인플립 비트)을 **학습에** 도입.
단순 포함이 아니라 **쌍으로 묶어** control variate로 써야 효과가 있다.
(Phase 1의 pre-sampled PRNG 설계 덕에 사실상 공짜.)
- [ ] MMD식 정규화 셀프플레이 (loss에 ~20줄) — 2인 제로섬 근사 내시 보험.
- [ ] **추론 시점 탐색 (PIMC / ISMCTS)** — 최종 강함의 가장 큰 이득. 로스트시티는 블러핑
경제가 없는 저기만성 불완전정보 게임이라 결정화 탐색이 잘 맞는다. raw net 대
net+search 맞대결로 측정.
- [ ] 네트워크 용량 A/B (512×3 → 1024×3 또는 residual) — **파이프라인 변경이 끝난 뒤에.**
## 범위 밖 (명시적 동결)
- **웹 클라이언트/ONNX는 별도 계획 전까지 레거시 단판 모델로 동결한다.** obs 개편 즉시
export 파이프라인(`scripts/export_jax_ppo_onnx.py:70`, manifest `observation_size: 454`),
TS obs 빌더, TS 단판 엔진이 전부 비호환이 된다. 이 선언이 없으면 실행 중 스코프가
웹 재작성으로 샌다.
- **Deep CFR 복귀** — 이미 BC 천장을 쳤고, 3라운드는 트리만 키운다. PPO+league를 학습
백본으로 유지한다.
- 레거시 호환을 위한 타협.
## 성공 기준
0. **(Phase 0a 게이트)** 루프 수정 후 단판 체제에서 동일 learner-액션 예산으로 기존 anchor
성적 재현 ≥ 동등. **없으면 auto-reset 버그가 3라운드 결과에 섞여 원인 분리가 불가능해진다.**
1. **carry 프로브에서 행동이 단조롭게 변한다** — carry 6개 수준에 걸쳐 원정 개수/악수 비율의
**단조 추세**(CI 포함). ("행동이 변한다"는 노이즈로도 통과 가능하므로 단조성으로 정의.)
2. **매치 승률이 구 정책(3라운드에 그대로 투입)을 이긴다** — 페어드 매치 **≥ 1만 쌍**,
매치 승률 **Wilson 하한 > 0.5** ** 평균 총 마진 **CI 하한 > 0**.
3. 착취자(exploiter) 승률이 악화되지 않는다.
4. **(anchor 비회귀)** 휴리스틱 anchor 상대 매치 승률·라운드당 마진이 구 정책 대비 악화되지
않는다. **기준 2의 구멍을 막는 항목**: carry-blind인 구 정책만 상대로 이기는 것은
**카드 플레이가 퇴보해도 carry 착취만으로 달성 가능**하다.
**단, anchor 선택에 주의:** Phase 0a에서 실증됐듯 **약한 anchor(`heuristic_balanced`)는
무한 stalling을 허용해서 신호를 오염시킨다** — 승률이 실력이 아니라 상한 도달률을 잰다.
**`heuristic_expert`처럼 stalling을 허용하지 않는 anchor**(게임이 자연 길이로 끝남)만
실력 지표로 쓰고, 약한 anchor는 **자연 종료 게임만 필터링해서** 보거나 `max_steps_rate`
함께 보고할 것.
---
## 실행 결과 (2026-07-15, 커밋 `9ba5a07`)
### 동작하는 것
- **셀프플레이가 stalling을 스스로 제거한다.** 매치가 146.8수(라운드 ≈49수)로 수렴하고
덱 레이스 비율 91%. Phase 0a에서 정적 상대가 225수까지 끌던 것과 대조된다.
- **미러 매치 평가가 정확하다.** 셀프플레이 duplicate 승률 0.4968, **평균 lead 정확히 0.0**
같은 3딜 + 같은 코인 + 자리 교대가 딜 운을 완전히 상쇄한다. 미완료 0%.
- **전지적 critic이 격리돼 있다.** 특권 입력을 흔들어도 정책 로짓은 비트 단위로 동일하고,
가치만 움직인다 (테스트로 고정).
### 성공 기준 1 — **약하게만 충족. 사실상 미달.**
carry 프로브(3라운드 시작 시 carry 주입)에서:
| | scale 50 | scale 25 | scale 12 |
|---|---|---|---|
| 악수 스프레드 (carry 60 → +60) | 0.19 | 0.08 | **0.31** |
| 원정 개수 | 5.00 고정 | 4.99 고정 | **5.00 고정** |
- **악수 사용은 6개 carry 수준 전체에서 단조**로 움직인다 (뒤지면 배수 베팅 ↑). 방향은 맞다.
- 그러나 **원정 개수는 carry와 무관하게 5.00에 붙어 있다.** 굳은 습관이지 조건부 플레이가 아니다.
- 효과 크기가 작다.
### 확정된 진단
1. **shaping 가설은 틀렸다.** `potential_shaping_final=0.0`으로 완전히 꺼도 프로브는 평평했다.
2. **원인은 `terminal_scale`이다.** carry = 60에서 `tanh((margin 60)/50)`의 인자는 현실적
마진 범위에서 **거의 선형**이고, 선형 구간에서 `E[tanh]` 최대화는 `E[margin]` 최대화와 같다
→ 도박할 이유가 없다. 위험 추구는 tanh가 **강하게 볼록한** 구간에서만 나오며 scale을 낮춰야
그 구간에 들어간다. scale 50 → 12에서 스프레드가 0.19 → 0.31로 커진 것이 이를 확인한다.
3. **프로브 설계 결함:** scale 12에서 `tanh(60/12) ≈ 1.0`이라 **±60은 포화 = 그래디언트 0**.
거기서 정책은 학습된 바가 없다. **의미 있는 측정 구간은 `|carry| ≲ 2 × terminal_scale`.**
프로브 수준을 scale에 맞춰 재설계해야 한다.
### 다음에 할 일
- [ ] **프로브 재설계**: carry 수준을 `terminal_scale`에 맞춰 잡는다 (포화 구간 측정 금지).
- [ ] **학습 중 carry 분포 확인**: 셀프플레이에서 3라운드 진입 시 carry가 실제로 얼마나
퍼지는가. 좁으면 정책이 큰 deficit을 본 적이 없다는 뜻이고, 그게 진짜 원인일 수 있다.
- [ ] **"항상 5색 개시"가 정상인지 검증**: 인간 기보에서 AI는 4.81, 인간은 4.19를 열었다.
5.00 고정은 의심스럽다. 셀프플레이 균형인지, 탐색 붕괴인지 (엔트로피 3.5 → 1.03).
- [ ] scale 스케줄(50 → 12 후기 fine-tune)이 처음부터 12로 학습하는 것보다 나은지 A/B.
---
## 판정: 보상은 **선형 총점**이다 (2026-07-15, Fable 3차 검토 + 실측)
사용자 제안("그냥 3판 총점이 크기만 하면 되는 것 아니냐")이 **맞았다. tanh 종료 보상은
오버엔지니어링이었다.**
### 근거 1 — 분해 논증 (실측)
보상이 총점에 선형이면 라운드가 독립이므로 3라운드 게임이 3개의 독립 단판으로 분해된다.
셀프플레이 2,048매치 실측: `corr(m1, m2) = 0.004`, `corr(m1+m2, m3) = 0.05`.
라운드를 잇는 유일한 고리인 **선공 어드밴티지는 +0.73 ± 0.84점** — 0과 구분 불가.
### 근거 2 — 리스크 태도는 값어치가 없다 (실측)
Fable이 리스크 태도를 직접 구현해 greedy 클론과 duplicate 6,144판 맞대결:
| 도박 정책 | 매치 승률 |
|---|---|
| 3R에서 20점 이상 뒤지면 온도 샘플링 | **0.482** (진다) |
| 40점 이상 뒤질 때만 | 0.498 (본전) |
**일부러 도박을 시켜도 지거나 본전.** 로스트시티의 분산 레버(한계 악수 ≈ Δσ +1.7에 마진
−2~3점)가 근본적으로 약해서, 분산을 사는 비용이 볼록성 이득을 먹는다.
carry 조건부 플레이의 가치 상한: **1승점 미만.**
### 근거 3 — 맞대결에서 단순한 쪽이 **이겼다**
duplicate 10,000판 (같은 3딜 + 같은 코인 + 자리 교대), 동일 컴퓨트 300 업데이트:
| A | B | A 승률 | A 평균 총점차 |
|---|---|---|---|
| **선형 총점** | tanh(총점/12) | **0.5859** (CI 0.5760.596) | **+20.3점** |
버려도 손해가 없는 게 아니라 **버리니 더 강해졌다.** 이유는 리스크가 아니라 **신호 밀도**다:
tanh는 ~150수 매치에 포화된 ±1 하나를 주고, 선형은 매 수마다 그 수가 총점차를 움직인 만큼을
준다. γ=1에서 후자의 합이 정확히 최종 총점차로 telescoping되므로 **목적함수는 동일한데 크레딧
할당만 150배 조밀**하다.
### 근거 4 — 탐색 붕괴가 아니었다
프로브가 평평했던 이유는 탐색 붕괴가 아니다. 샘플링 프로브에서도 원정 5.00 ± 0.05,
엔트로피 1.36나트(유효 행동 3.9개)로 정상 수렴. 그리고 **critic은 carry를 완벽히 읽고 있었다**
(3R 시작 가치 0.87 → +0.86 단조). 신호는 있었고, **정책이 그걸로 살 수 있는 물건이
없었을 뿐**이다.
### 최종 설계
```python
# match_ppo.py 롤아웃 바디
reward0 = jnp.where(active, (after - before) / cfg.reward.terminal_scale, 0.0)
```
- **성공 기준 1(carry 프로브 단조성)은 게이트에서 제거한다.** 최적 반응 자체가 이 게임에서
거의 존재하지 않는다는 것이 측정 결과다. 기준 2·4(맞대결 + anchor 비회귀)가 옳은 자다.
- `carry`는 obs에 **남긴다** (비용 0, 선공 규칙이 키로 쓰는 상태, critic이 잘 읽음).
### 목적함수와 무관하게 살아남은 것 (전부 순이득)
in-scan auto-reset + GAE 절단 부트스트랩(샘플 5.8배), 메트릭 재작업, **비대칭 CTDE critic**,
**양쪽 좌석 학습**, **duplicate 매치 평가**, pre-sampled PRNG/미러 설계, 그리고 match_obs의
단판 결함 수정분(`to_move`, 덱 시계, ÷780 → ÷75, 살아있는 점수 3분할).
**죽은 것은 tanh 종료 보상과 carry 구간 원-핫뿐이다.**
---
## 신구 대결: 진짜 3라운드 클래식에서 (2026-07-15)
**질문: "기존 학습 방법 대비 실제로 뭐가 나아졌나?"**
duplicate 매치 8,192판(같은 3딜 + 같은 코인 + 자리 교대), 진짜 3라운드 게임:
| 신(매치 스택) | 상대 | 신 승률 | 신 평균 총점차 | 상대 learner 액션 |
|---|---|---|---|---|
| 39.3M | 기존 baseline (정적 상대 학습) | 0.6077 | +19.5 | 104M |
| 39.3M | Phase 0a 게이트 (루프 수정만) | **0.5842** | +16.0 | **411M (10.5배)** |
| 39.3M | league (셀프플레이, 웹 배포판) | 0.3142 | 37.0 | 122.6M (3.1배) |
| **131M** | **league (셀프플레이, 웹 배포판)** | **0.6094** (CI 0.5990.620) | **+22.0** | 122.6M |
**결론:**
1. **샘플 효율이 크게 올랐다.** 39.3M짜리가 **411M(10.5배)짜리를 이긴다.**
2. **동일 예산에서 기존 최강(league)을 이긴다** — 131M vs 122.6M에서 0.6094.
3. 39.3M에서 league에 졌던 것(0.3142)은 **약해서가 아니라 예산이 1/3이어서**였다.
**남은 단서:** league는 **착취자(exploiter) 구조**를 포함한 런이다. 평균 강함은 우리가 이기지만
**exploitability는 아직 재지 않았다.** "덜 착취당한다"는 증명되지 않았다.
---
## 기여도 A/B (2026-07-15)
각 조각을 하나씩 끄고 동일 컴퓨트(300 업데이트, batch 512)로 학습시킨 뒤,
full 스택과 duplicate 매치 8,192판 맞대결. **0.5 미만 = 그 조각이 기여하고 있었다.**
| 제거한 것 | full 대비 승률 | 95% CI | 평균 총점차 | 판정 |
|---|---|---|---|---|
| **양쪽 좌석 학습** | **0.3317** | [0.322, 0.342] | 28.6 | **압도적 기여** |
| 매치 관측 (carry/to_move/덱시계/살아있는점수) | 0.4751 | [0.464, 0.486] | 4.0 | 유의하게 기여 |
| **전지적 critic (CTDE)** | **0.5160** | [0.505, 0.527] | +2.6 | **오히려 해가 된다** |
### 1. 양쪽 좌석 학습 — 가장 큰 기여
끄면 승률이 0.33으로 무너진다. 같은 업데이트 수에서 learner 액션이 **정확히 절반**이 되므로
상당 부분은 샘플 수 효과다 — 그러나 **그게 요점이다.** 셀프플레이에서 상대 좌석의 수는 같은
네트워크가 둔 것인데 기존 트레이너는 `stop_gradient`로 버렸다. **공짜로 데이터가 2배**가 된다.
### 2. 매치 관측 — 작지만 실재
0.4751 (CI 상한 0.486 < 0.5). `to_move`, 덱 시계, ÷780 → ÷75 스케일 수정, 살아있는 점수
3분할이 합쳐서 약 2.5%p 값어치. carry 자체는 (분해 논증대로) 거의 기여하지 않으므로,
이 이득의 대부분은 **단판 obs의 결함 수정분**으로 보인다.
### 3. 전지적 critic — **Fable의 최우선 권고가 틀렸다**
Fable은 이것을 "가장 큰 누락 아이디어"로 꼽았다. 실측은 **반대**다: 끄면 오히려
**0.5160 (CI 하한 0.505 > 0.5)** 으로 유의하게 **더 강해진다.**
그럴듯한 이유: 특권 critic은 덱을 알기 때문에 가치를 아주 잘 맞추지만, 그 결과
`V(전체상태) ≠ E[리턴 | 마스킹된 관측]`이 되어, **어드밴티지에 정책이 통제할 수 없는 성분**이
섞인다. 액터 입장에서 그건 분산 감소가 아니라 노이즈다. 부분관측 환경에서 비대칭 critic이
해가 되는 것은 알려진 현상이다.
**조치: 전지적 critic을 기본에서 끈다.** (PIMC 탐색의 리프 평가기로 재활용한다는 계획은
별개로 유효할 수 있으나, 정책 학습용으로는 손해다.)
### 종합: 실제로 값어치 있었던 것
1. **in-scan auto-reset + GAE 절단 부트스트랩** — 샘플 5.8배, expert 0.4525 → 0.5071
2. **양쪽 좌석 학습** — 데이터 2배, 단독으로 승률 0.33 → 0.5
3. **선형 총점 보상** — tanh 대비 0.5859 (사용자 제안)
4. **매치 관측 결함 수정** — 약 2.5%p
5. ~~전지적 critic~~**해가 된다. 끈다.**
### 정정: 전지적 critic의 가치는 **스케일에 따라 뒤집힌다**
위 A/B는 39.3M learner 액션 규모에서 돌렸다. 실전 규모(131M)에서 동일 설정으로 둘을 직접
맞붙이면 **결론이 반대로 나온다.**
| 규모 | critic OFF의 승률 (critic ON 상대) | 판정 |
|---|---|---|
| 39.3M (A/B 규모) | **0.5160** [0.505, 0.527] | OFF가 낫다 |
| **131M (실전 규모)** | **0.4633** [0.453, 0.474] | **ON이 낫다** |
**전지적 critic은 그것을 학습시킬 데이터가 충분해질 때 값어치를 한다.** 39.3M에서는 특권
정보를 쓰는 큰 가치 트렁크를 제대로 못 맞춰서 어드밴티지에 노이즈만 얹었고, 131M에서는
제대로 맞춰서 분산 감소가 실현된다.
**교훈: 배포할 예산이 아닌 규모에서 ablation을 돌리면 결론이 뒤집힐 수 있다.**
`privileged_critic` 기본값은 **ON으로 되돌린다.**
### 최종 순위 (진짜 3라운드 클래식, duplicate 8,192판)
| 모델 | learner 액션 | vs league |
|---|---|---|
| **매치 스택 (critic ON)** | 131M | **0.6094** (+22.0점) |
| 매치 스택 (critic OFF) | 131M | 0.5526 (+10.9점) |
| league (기존 최강, 웹 배포판) | 122.6M | — |
### 기여도 재측정: 동일 규모(131M)에서
39.3M A/B는 critic에게 가장 불리한 지점이었다. 두 조각을 **같은 131M 예산**에서 다시 쟀다:
| 조각 | full의 승률 (그 조각 뺀 버전 상대) | 순이득 |
|---|---|---|
| **양쪽 좌석 학습** | 0.6707 | **+17.1 %p** |
| **전지적 critic** | 0.5367 | **+3.7 %p** |
**양쪽 좌석이 약 4.6배 크다.** 그리고 성격이 다르다:
| 규모 | 양쪽 좌석 | 전지적 critic |
|---|---|---|
| 39.3M | +16.8 %p | **1.6 %p** |
| 131M | **+17.1 %p** | +3.7 %p |
**양쪽 좌석은 규모와 무관하게 안정적**(+17%p). **critic은 규모를 탄다.**
이유: 양쪽 좌석은 **데이터를 2배로** 만든다(65.5M → 131M, 같은 컴퓨트에서). critic은 데이터를
안 늘리고 **채점 정확도만** 올리는데, 코치가 배울 게 많아져서 데이터가 부족하면 오히려
엉터리 채점을 한다.
**Fable은 critic을 "가장 큰 누락 아이디어" 1순위로, 양쪽 좌석을 4순위로 꼽았다. 정확히 거꾸로였다.**
---
## Exploitability: 우리가 league보다 **덜 착취당한다** (2026-07-15)
**남겨뒀던 단서를 해소했다.** 평균 실력에서 이기는 것과 "파먹을 약점이 없는 것"은 다른
속성이다. 가위바위보에서 바위를 60% 내는 놈은 아무한테나 이길 수도 있지만, 그 습관을
알아챈 놈에게는 매번 진다.
**측정 방법:** 정책을 얼려놓고, **오직 그놈만 이기도록 특화된 새 정책을 처음부터 학습**시킨다
(`src/lost_cities_jax/exploit.py`). 착취자가 도달한 승률이 곧 그 정책이 못 막아낸 습관의 크기다.
**동일 착취자 예산 (250 업데이트 × batch 1024 = 32.5M learner 액션):**
| 얼려놓은 정책 | 착취자 승률 | 95% CI |
|---|---|---|
| **우리 매치 스택 (131M)** | **0.2278** | [0.219, 0.237] |
| league (웹 배포판, 122.6M) | **0.3213** | [0.311, 0.331] |
**league가 9.4%p 더 털린다.** 그리고 둘 다 0.5에서 한참 멀다 — 어느 쪽도 만만한 상대는 아니다.
### 최종 정리: 두 축 모두에서 이긴다
| 축 | 결과 |
|---|---|
| **평균 실력** (정면 대결) | 0.6094, +22.0점 ✅ |
| **약점 적음** (exploitability) | 0.2278 vs 0.3213 ✅ |
### 단서
1. **하한선이다.** 더 세거나 더 오래 학습한 착취자는 더 찾아낼 수 있다. **절대값이 아니라
동일 예산에서의 비교로만 의미가 있다.**
2. **착취자 예산이 목표의 약 1/4에 불과하다 (32.5M vs 131M).** 약한 공격자다.
따라서 **"우리는 22.8%밖에 안 털린다"고 말할 수 없다** — 그건 이 약한 공격자 기준일 뿐이다.
살아남는 주장은 **"동일 예산에서 league가 더 털린다"** 하나뿐이다.
착취자를 목표와 같은 예산(131M)으로 키워서 **순서가 유지되는지** 확인해야 한다.
3. league는 단판 정책이라 3라운드 게임에선 다소 제 물이 아니다 — 다만 carry가 무의미하다는
것이 이미 증명됐으므로 큰 불리함은 아니다.
**흥미로운 점:** 우리는 **순수 셀프플레이**이고 league는 **착취자 구조를 학습에 넣은** 런인데,
그런데도 우리가 덜 착취당한다. 착취자 구조가 exploitability를 낮춰줄 것이라는 기대가
이 게임에서는 확인되지 않았다.
### 착취자에 제대로 자금을 대고 재측정 (131M, 목표와 동일 예산)
앞선 32.5M 측정은 **공격자가 목표보다 4배 부족**했다. 예산을 목표와 맞춰 다시 쟀다.
실측 확인: 착취자 learner 액션 **130.4M** (우리 상대) / **131.1M** (league 상대).
| 얼려놓은 정책 | 착취자 승률 | 95% CI | 착취자 평균 마진 | 이전(32.5M) |
|---|---|---|---|---|
| **우리 매치 스택** (131.1M 학습) | **0.4657** | [0.455, 0.477] | **5.5점** | 0.2278 |
| league (122.6M 학습) | **0.6295** | [0.619, 0.640] | **+28.8점** | 0.3213 |
**CI가 전혀 겹치지 않고, 격차가 오히려 벌어졌다 (0.094 → 0.164).**
**질적으로 선을 넘었다:** 제대로 자금 댄 전담 공격자는 **league를 아예 이긴다**(0.63, +28.8점).
반면 **우리는 여전히 못 뚫는다**(0.47, 5.5점).
#### 우리에게 불리한 단서 (반드시 함께 읽을 것)
- **둘 다 1000 업데이트 끝까지 정체 없이 오르고 있었다. 이 수치도 여전히 하한선이다.**
- **우리 쪽 착취자의 말단 기울기가 더 가파르다** (+0.031 vs +0.016 / 100 업데이트).
league 착취자 곡선은 꺾이기 시작했는데 우리 쪽은 아직 갈 길이 남았다.
**공격자 예산을 크게 더 키우면 격차가 좁혀질 수 있고, 뒤집힐 가능성도 배제 못 한다.**
#### 구조적 교란 (순수한 "학습법 A vs B"가 아니다)
1. **league는 단판 정책이다.** carry도 매치 점수도 모른다. 더 털리는 것의 일부는
**학습 방법이 아니라 구조적 맹점** 때문일 수 있다.
2. **둘 다 greedy(argmax)로 뒀다.** 결정론적 정책은 **정의상 최대로 착취당한다.**
양쪽을 똑같이 대우했으니 비교는 공정하나, **절대값은 부풀려져 있다.**
#### 살아남는 주장 / 못 하는 주장
-**"동일 예산의 전담 공격자 앞에서 league는 뚫리고 우리는 안 뚫린다."**
- ❌ "우리 정책은 착취 불가능하다." — 하한선일 뿐이고, 우리 쪽 곡선은 아직 오르는 중이다.
+148
View File
@@ -0,0 +1,148 @@
#!/usr/bin/env python3
"""What is each piece of the match stack worth?
Train one variant per switch at identical compute, then play each against the full
stack in duplicate matches -- same three deals, same coins, both seats -- so deal
luck cancels and only the policy difference is left.
The variants do not share an observation shape (dropping the privileged critic or
the match features changes the input dims), so this carries its own head-to-head
rather than reusing match_eval's, which assumes one architecture.
"""
from __future__ import annotations
import json
from pathlib import Path
import jax
import jax.numpy as jnp
import numpy as np
from lost_cities_jax.match import (
match_legal_action_mask,
match_reset_from,
match_score,
match_step,
)
from lost_cities_jax.match_eval import MATCH_SCAN_STEPS, _wilson, match_bank
from lost_cities_jax.match_ppo import (
Ablation,
MatchActorCritic,
_seat_views,
create_match_train_state,
match_train,
)
from lost_cities_jax.ppo import load_config, mask_logits, restore_checkpoint
VARIANTS = [
Ablation(privileged_critic=False),
Ablation(both_seats=False),
Ablation(match_obs=False),
]
MATCHES = 4096
FULL_CHECKPOINT = Path("runs/jax-ppo-match/2026-07-15_030152_match-linear/latest")
def _load(cfg, checkpoint: Path, ablation: Ablation):
state = create_match_train_state(cfg, jax.random.PRNGKey(0), ablation)
return restore_checkpoint(checkpoint, state).params
def _existing_run(label: str) -> Path | None:
runs = sorted(Path("runs/jax-ppo-match").glob(f"*ablate-{label}"))
for run in reversed(runs):
if (run / "latest").exists():
return run
return None
def duel(cfg, params_a, ablation_a, params_b, ablation_b, *, matches: int) -> dict:
"""A vs B over duplicate matches. Each side sees the world its own way."""
decks, coins = match_bank(20260719, matches)
model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
seats = jnp.arange(2, dtype=jnp.int32)
@jax.jit
def run(env, a_seat):
def body(carry, _):
env, _unused = carry
to_move = env.round.to_move.astype(jnp.int32)
mask = jax.vmap(match_legal_action_mask)(env)
def act(params, ablation):
obs, critic = _seat_views(env, seats, ablation)
idx = to_move[None, :, None]
seat_obs = jnp.take_along_axis(obs, idx, axis=0)[0]
seat_critic = jnp.take_along_axis(critic, idx, axis=0)[0]
logits, _ = model.apply(params, seat_obs, seat_critic)
return jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)
action = jnp.where(
to_move == a_seat, act(params_a, ablation_a), act(params_b, ablation_b)
)
env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
return (env, _unused), None
(env, _), _ = jax.lax.scan(body, (env, jnp.int32(0)), xs=None, length=MATCH_SCAN_STEPS)
return env
leads = []
for seat in (0, 1):
env = jax.vmap(match_reset_from)(decks, coins)
final = run(env, jnp.full((matches,), seat, dtype=jnp.int32))
totals = np.asarray(jax.vmap(match_score)(final))
leads.append(totals[:, seat] - totals[:, 1 - seat])
lead = np.concatenate(leads)
games = float(lead.size)
wins = float((lead > 0).sum())
low, high = _wilson(wins, games)
return {
"matches": games,
"a_win_rate": wins / games,
"wilson_low": low,
"wilson_high": high,
"a_mean_lead": float(lead.mean()),
}
def main() -> None:
cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
full = _load(cfg, FULL_CHECKPOINT, Ablation())
rows = []
for ablation in VARIANTS:
label = ablation.label()
run_dir = _existing_run(label)
if run_dir is None:
print(f"\n===== training {label} =====", flush=True)
cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
cfg.run.experiment_name = f"ablate-{label}"
run_dir = match_train(cfg, ablation=ablation)
else:
print(f"\n===== reusing {run_dir} =====", flush=True)
cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
params = _load(cfg, run_dir / "latest", ablation)
# The ablated net plays seat A, the full stack answers. Below 0.5 means
# the piece we removed was carrying weight.
result = duel(cfg, params, ablation, full, Ablation(), matches=MATCHES)
rows.append({"ablation": label, **result})
print(f" {label}: win rate vs full = {result['a_win_rate']:.4f}", flush=True)
print("\n\n==================== ablation ====================")
print(f"{'removed':<26}{'win rate vs full':>18}{'95% CI':>22}{'mean lead':>12}")
print("-" * 78)
for row in rows:
ci = f"[{row['wilson_low']:.3f}, {row['wilson_high']:.3f}]"
print(
f"{row['ablation']:<26}{row['a_win_rate']:>18.4f}{ci:>22}{row['a_mean_lead']:>+12.1f}"
)
print("\nbelow 0.5 = the removed piece was carrying weight")
Path("runs/jax-ppo-match/ablation.json").write_text(json.dumps(rows, indent=2, sort_keys=True))
if __name__ == "__main__":
main()
+138
View File
@@ -0,0 +1,138 @@
#!/usr/bin/env python3
"""Export a match policy's actor trunk to ONNX for the browser.
Only the actor ships. The critic exists to grade moves during training and never
plays, so its trunk -- and the privileged view of the opponent's hand and the deck
that feeds it -- is dropped here rather than shipped and then not used. That also
means the exported graph physically cannot leak hidden state, which is a stronger
guarantee than promising not to call it.
MatchActorCritic lays the actor out as Dense_0..Dense_{num_layers} exactly as the
single-round model does, so the graph construction is the same; only the input
width and the checkpoint loader differ.
"""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
import jax
import jax.numpy as jnp
import numpy as np
from lost_cities_jax.match_obs import MATCH_CRITIC_OBS_DIM, MATCH_OBS_DIM
from lost_cities_jax.match_ppo import Ablation, MatchActorCritic, create_match_train_state
from lost_cities_jax.ppo import load_config, restore_checkpoint
from lost_cities_jax.types import N_ACTIONS
def build_argparser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--config", type=Path, default=Path("configs/jax_ppo/match-selfplay.yaml"))
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--codename", required=True, help="see data/models.json")
return parser
def export_model(checkpoint: Path, config: Path, output: Path, codename: str) -> None:
try:
import onnx
from onnx import TensorProto, helper, numpy_helper
from onnx.reference import ReferenceEvaluator
except ImportError as exc:
raise SystemExit("onnx is required; run with `uv run --with onnx ...`") from exc
cfg = load_config(config)
state = restore_checkpoint(
checkpoint, create_match_train_state(cfg, jax.random.PRNGKey(0), Ablation())
)
params = state.params
dense = params["params"]
nodes = []
initializers = []
previous = "obs"
for index in range(cfg.network.num_layers):
layer = dense[f"Dense_{index}"]
weight, bias = f"dense_{index}.weight", f"dense_{index}.bias"
initializers.extend(
[
numpy_helper.from_array(np.asarray(layer["kernel"], dtype=np.float32), weight),
numpy_helper.from_array(np.asarray(layer["bias"], dtype=np.float32), bias),
]
)
nodes.append(helper.make_node("Gemm", [previous, weight, bias], [f"dense_{index}.linear"]))
nodes.append(helper.make_node("Relu", [f"dense_{index}.linear"], [f"dense_{index}.relu"]))
previous = f"dense_{index}.relu"
actor = dense[f"Dense_{cfg.network.num_layers}"]
initializers.extend(
[
numpy_helper.from_array(np.asarray(actor["kernel"], dtype=np.float32), "actor.weight"),
numpy_helper.from_array(np.asarray(actor["bias"], dtype=np.float32), "actor.bias"),
]
)
nodes.append(helper.make_node("Gemm", [previous, "actor.weight", "actor.bias"], ["logits"]))
graph = helper.make_graph(
nodes,
f"coolrl-lost-cities-match-actor-{codename}",
[helper.make_tensor_value_info("obs", TensorProto.FLOAT, [None, MATCH_OBS_DIM])],
[helper.make_tensor_value_info("logits", TensorProto.FLOAT, [None, N_ACTIONS])],
initializer=initializers,
)
model = helper.make_model(
graph, producer_name="coolrl-lost-cities", opset_imports=[helper.make_opsetid("", 17)]
)
model.ir_version = 8
onnx.checker.check_model(model)
# The exported graph has to agree with the trained one, not merely load.
rng = np.random.default_rng(20260715)
sample = rng.normal(size=(8, MATCH_OBS_DIM)).astype(np.float32)
critic_stub = jnp.zeros((8, MATCH_CRITIC_OBS_DIM), dtype=jnp.float32)
flax_model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
expected, _ = flax_model.apply(params, jnp.asarray(sample), critic_stub)
actual = ReferenceEvaluator(model).run(None, {"obs": sample})[0]
np.testing.assert_allclose(actual, np.asarray(expected), rtol=2e-5, atol=2e-5)
np.testing.assert_array_equal(
np.argmax(actual, axis=1), np.argmax(np.asarray(expected), axis=1)
)
output.parent.mkdir(parents=True, exist_ok=True)
onnx.save(model, output)
model_bytes = output.read_bytes()
manifest = {
"format": "coolrl-lost-cities-match-onnx-v1",
"codename": codename,
"model_file": output.name,
"model_size_bytes": len(model_bytes),
"model_sha256": hashlib.sha256(model_bytes).hexdigest(),
"source_checkpoint": str(checkpoint),
"source_config": config.name,
"observation_size": MATCH_OBS_DIM,
"action_size": N_ACTIONS,
"hidden_size": cfg.network.hidden_size,
"num_layers": cfg.network.num_layers,
"dtype": "float32",
"validation_max_abs_error": float(np.max(np.abs(actual - np.asarray(expected)))),
}
output.with_suffix(".json").write_text(json.dumps(manifest, indent=2) + "\n")
print(f"exported {codename} -> {output} ({output.stat().st_size:,} bytes)")
print(
f" sha256 {manifest['model_sha256'][:12]} obs {MATCH_OBS_DIM} max err "
f"{manifest['validation_max_abs_error']:.2e}"
)
def main() -> None:
args = build_argparser().parse_args()
export_model(args.checkpoint, args.config, args.output, args.codename)
if __name__ == "__main__":
main()
+97
View File
@@ -0,0 +1,97 @@
#!/usr/bin/env python3
"""Generate match states and their observations so TypeScript can be checked against JAX.
The two observation builders must agree to the bit. A mismatch does not throw --
the ONNX policy consumes a wrong vector quite happily and plays worse for reasons
nobody can see. So the port is not trusted; it is checked.
States are drawn from real random play so the fixture covers the awkward parts:
mid-round, both seats to move, past a round roll-over, with a non-zero carry.
"""
from __future__ import annotations
import json
from pathlib import Path
import jax
import numpy as np
from lost_cities_jax.match import MatchState, match_reset_from, match_step
from lost_cities_jax.match_obs import MATCH_OBS_DIM, match_observation
from lost_cities_jax.opponents import random_legal_action
from lost_cities_jax.types import N_CARDS
OUTPUT = Path(__file__).resolve().parents[1] / "web" / "src" / "game" / "match-parity-fixture.json"
N_ROUNDS = 3
def match_json(match: MatchState) -> dict:
round_state = match.round
return {
"round": {
"deckOrder": np.asarray(round_state.deck_order).astype(int).tolist(),
"drawPtr": int(round_state.draw_ptr),
"cardLoc": np.asarray(round_state.card_loc).astype(int).tolist(),
"handPublic": np.asarray(round_state.hand_public).astype(bool).tolist(),
"colTop": np.asarray(round_state.col_top).astype(int).tolist(),
"colHandshakes": np.asarray(round_state.col_hs).astype(int).tolist(),
"colLength": np.asarray(round_state.col_len).astype(int).tolist(),
"piles": [
np.asarray(round_state.pile[color, : int(round_state.pile_len[color])])
.astype(int)
.tolist()
for color in range(5)
],
"toMove": int(round_state.to_move),
"stepCount": int(round_state.step_count),
"done": bool(round_state.done),
},
"deckOrders": np.asarray(match.deck_orders).astype(int).tolist(),
"coinFlips": np.asarray(match.coin_flips).astype(int).tolist(),
"roundIdx": int(match.round_idx),
"carry": np.asarray(match.carry).astype(int).tolist(),
"done": bool(match.done),
}
def main() -> None:
rng = np.random.default_rng(20260715)
key = jax.random.PRNGKey(7)
rows = []
for match_index in range(6):
decks = np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)])
coins = rng.integers(0, 2, size=(N_ROUNDS,))
match = match_reset_from(decks.astype(np.int8), coins.astype(np.int8))
# Sample the opening position and then every 17th ply, which lands in all
# three rounds and on both seats without hand-picking anything.
ply = 0
while not bool(match.done) and ply < 400:
if ply % 17 == 0 or ply == 0:
for player in (0, 1):
obs = np.asarray(match_observation(match, player), dtype=np.float64)
assert obs.shape == (MATCH_OBS_DIM,)
rows.append(
{
"match": match_json(match),
"player": player,
"observation": [round(float(v), 7) for v in obs],
}
)
key, step_key = jax.random.split(key)
action = int(random_legal_action(match.round, match.round.to_move, step_key))
match, _, _ = match_step(match, action)
ply += 1
del match_index
OUTPUT.write_text(
json.dumps({"format": "jax-web-match-parity-v1", "obsDim": MATCH_OBS_DIM, "rows": rows})
+ "\n"
)
print(f"wrote {len(rows)} rows -> {OUTPUT}")
if __name__ == "__main__":
main()
+124
View File
@@ -0,0 +1,124 @@
#!/usr/bin/env python3
"""Freeze each policy, train an exploiter against it on the same budget, compare.
Winning the head-to-head says a policy is strong on average. It does not say the
policy is hard to beat. This does: whatever the exploiter reaches is a habit the
frozen policy could not defend.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import jax
import numpy as np
from lost_cities_jax.exploit import (
match_frozen_policy,
single_round_frozen_policy,
train_exploiter,
)
from lost_cities_jax.match import match_reset_from, match_score
from lost_cities_jax.match_eval import MATCH_SCAN_STEPS, _wilson, match_bank
from lost_cities_jax.match_ppo import Ablation, create_match_train_state
from lost_cities_jax.ppo import create_train_state, load_config, restore_checkpoint
OURS = Path("runs/jax-ppo-match/2026-07-15_031529_match-scaled/latest")
LEAGUE = Path(
"/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/latest"
)
MATCHES = 4096
def _final_score(cfg, exploiter_params, frozen, matches: int) -> dict:
"""Play the trained exploiter against the frozen policy, both seats."""
import jax.numpy as jnp
from lost_cities_jax.exploit import match_frozen_policy as _mk
from lost_cities_jax.match import match_step
attacker = _mk(cfg, exploiter_params, Ablation())
decks, coins = match_bank(20260722, matches)
@jax.jit
def run(env, a_seat):
def body(carry, _):
env, _u = carry
action = jnp.where(
env.round.to_move.astype(jnp.int32) == a_seat,
attacker(env, a_seat),
frozen(env, 1 - a_seat),
)
env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
return (env, _u), None
(env, _), _ = jax.lax.scan(body, (env, jnp.int32(0)), xs=None, length=MATCH_SCAN_STEPS)
return env
leads = []
for seat in (0, 1):
env = jax.vmap(match_reset_from)(decks, coins)
final = run(env, jnp.full((matches,), seat, dtype=jnp.int32))
totals = np.asarray(jax.vmap(match_score)(final))
leads.append(totals[:, seat] - totals[:, 1 - seat])
lead = np.concatenate(leads)
games = float(lead.size)
wins = float((lead > 0).sum())
low, high = _wilson(wins, games)
return {
"exploiter_win_rate": wins / games,
"wilson_low": low,
"wilson_high": high,
"exploiter_mean_lead": float(lead.mean()),
"matches": games,
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--updates", type=int, default=250, help="exploiter training updates")
parser.add_argument("--batch-games", type=int, default=1024)
parser.add_argument("--tag", default="", help="suffix for the run dir, to keep runs apart")
args = parser.parse_args()
match_cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
old_cfg = load_config("configs/jax_ppo/balanced.yaml")
ours = restore_checkpoint(
OURS, create_match_train_state(match_cfg, jax.random.PRNGKey(0), Ablation())
).params
league = restore_checkpoint(LEAGUE, create_train_state(old_cfg, jax.random.PRNGKey(0))).params
targets = {
"ours (match stack, 131M)": match_frozen_policy(match_cfg, ours, Ablation()),
"league (web-deployed)": single_round_frozen_policy(old_cfg, league),
}
rows = []
for name, frozen in targets.items():
print(f"\n===== training an exploiter against: {name} =====", flush=True)
cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
cfg.ppo.batch_games = args.batch_games
cfg.run.total_updates = args.updates
cfg.run.log_every = max(1, args.updates // 10)
slug = name.split()[0] + args.tag
state = train_exploiter(cfg, frozen, Path(f"runs/jax-ppo-match/exploit-{slug}"))
result = _final_score(cfg, state.params, frozen, MATCHES)
rows.append({"target": name, **result})
print(f" exploiter reached {result['exploiter_win_rate']:.4f} vs {name}", flush=True)
print("\n\n============ exploitability (same exploiter budget) ============")
print(f"{'frozen policy':<28}{'exploiter win rate':>20}{'95% CI':>22}")
print("-" * 72)
for row in rows:
ci = f"[{row['wilson_low']:.3f}, {row['wilson_high']:.3f}]"
print(f"{row['target']:<28}{row['exploiter_win_rate']:>20.4f}{ci:>22}")
print("\nhigher = the frozen policy had more to farm. 0.5 = nothing found.")
Path(f"runs/jax-ppo-match/exploitability{args.tag}.json").write_text(json.dumps(rows, indent=2))
if __name__ == "__main__":
main()
+86
View File
@@ -0,0 +1,86 @@
#!/usr/bin/env python3
"""Serve the production web client and append completed games to JSONL."""
from __future__ import annotations
import argparse
import json
from http import HTTPStatus
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
MAX_RECORD_BYTES = 6_000_000
# v2 adds the opponent's identity (codename + hash) and the match layer -- three
# deals, the coin flips, and which round each move belongs to. v1 is still
# accepted because 111 games were recorded under it; they are all altair, which
# data/models.json records since v1 has nowhere to say so.
SUPPORTED_FORMATS = frozenset({"lost-cities-web-game-v1", "lost-cities-web-game-v2"})
def parse_record(body: bytes) -> dict[str, Any]:
if len(body) > MAX_RECORD_BYTES:
raise ValueError("record is too large")
value = json.loads(body)
if not isinstance(value, dict) or value.get("format") not in SUPPORTED_FORMATS:
raise ValueError("unsupported game record")
if not isinstance(value.get("gameId"), str) or not isinstance(value.get("moves"), list):
raise ValueError("invalid game record")
return value
def make_handler(dist: Path, output: Path):
seen_ids: set[str] = set()
if output.exists():
for line in output.read_text(encoding="utf-8").splitlines():
try:
game_id = json.loads(line).get("gameId")
if isinstance(game_id, str):
seen_ids.add(game_id)
except (json.JSONDecodeError, AttributeError):
continue
class Handler(SimpleHTTPRequestHandler):
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, directory=str(dist), **kwargs)
def do_POST(self) -> None: # noqa: N802
if self.path != "/api/game-records":
self.send_error(HTTPStatus.NOT_FOUND)
return
try:
length = int(self.headers.get("content-length", "0"))
record = parse_record(self.rfile.read(length))
except (ValueError, json.JSONDecodeError) as error:
self.send_error(HTTPStatus.BAD_REQUEST, str(error))
return
game_id = record["gameId"]
if game_id not in seen_ids:
output.parent.mkdir(parents=True, exist_ok=True)
with output.open("a", encoding="utf-8") as stream:
stream.write(
json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n"
)
seen_ids.add(game_id)
self.send_response(HTTPStatus.NO_CONTENT)
self.end_headers()
return Handler
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="0.0.0.0")
parser.add_argument("--port", type=int, default=5173)
parser.add_argument("--dist", type=Path, default=Path("web/dist"))
parser.add_argument("--output", type=Path, default=Path("data/human-play/game-records.jsonl"))
args = parser.parse_args()
server = ThreadingHTTPServer((args.host, args.port), make_handler(args.dist, args.output))
print(f"Serving {args.dist} on http://{args.host}:{args.port}", flush=True)
print(f"Writing game records to {args.output}", flush=True)
server.serve_forever()
if __name__ == "__main__":
main()
+62
View File
@@ -0,0 +1,62 @@
#!/usr/bin/env python3
"""Train the classic three-round agent by self-play, then measure it."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import jax
from lost_cities_jax.match_eval import carry_probe, match_evaluate
from lost_cities_jax.match_ppo import create_match_train_state, match_train
from lost_cities_jax.ppo import load_config, restore_checkpoint
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config", default="configs/jax_ppo/match-selfplay.yaml")
parser.add_argument("--eval-only", type=Path, default=None)
parser.add_argument("--matches", type=int, default=1024)
parser.add_argument("--set", action="append", default=[])
args = parser.parse_args()
cfg = load_config(args.config)
for override in args.set:
path, _, raw = override.partition("=")
section, _, field = path.partition(".")
target = getattr(cfg, section)
current = getattr(target, field)
value = type(current)(raw) if not isinstance(current, bool) else raw == "true"
setattr(target, field, value)
if args.eval_only is None:
run_dir = match_train(cfg)
checkpoint = run_dir / "latest"
else:
checkpoint = args.eval_only
state = create_match_train_state(cfg, jax.random.PRNGKey(0))
state = restore_checkpoint(Path(checkpoint), state)
result = match_evaluate(cfg, state.params, matches=args.matches)
print("\n== duplicate match eval (self-play) ==")
print(json.dumps(result, indent=2, sort_keys=True))
print("\n== carry probe: does round-three play react to the deficit? ==")
rows = carry_probe(cfg, state.params, matches=args.matches // 2)
header = f"{'carry':>7}{'win_rate':>10}{'opened':>9}{'wagers':>9}{'deck_race':>11}{'plies':>8}"
print(header)
for row in rows:
print(
f"{row['carry']:>7}{row['win_rate']:>10.3f}{row['opened_colors']:>9.2f}"
f"{row['wagers_played']:>9.2f}{row['deck_race_rate']:>11.3f}{row['mean_plies']:>8.0f}"
)
Path(checkpoint).parent.joinpath("carry_probe.json").write_text(
json.dumps({"eval": result, "probe": rows}, indent=2, sort_keys=True)
)
if __name__ == "__main__":
main()
+5 -2
View File
@@ -36,11 +36,14 @@ def reset(rng: jax.Array) -> State:
return reset_from_order(deck_order) return reset_from_order(deck_order)
def reset_from_order(deck_order: jax.Array) -> State: def reset_from_order(deck_order: jax.Array, first_player: jax.Array | int = 0) -> State:
"""Return an initial state using an explicit 60-card deck permutation. """Return an initial state using an explicit 60-card deck permutation.
The first eight cards are dealt to player 0 and the next eight to player 1. The first eight cards are dealt to player 0 and the next eight to player 1.
The remaining cards are drawn in array order starting at ``draw_ptr == 16``. The remaining cards are drawn in array order starting at ``draw_ptr == 16``.
``first_player`` moves first. Rounds two and three of a classic match are led
by whoever is ahead on points, so the match layer sets this per round.
""" """
deck_order = jnp.asarray(deck_order, dtype=jnp.int8) deck_order = jnp.asarray(deck_order, dtype=jnp.int8)
@@ -60,7 +63,7 @@ def reset_from_order(deck_order: jax.Array) -> State:
col_len=jnp.zeros((2, N_COLORS), dtype=jnp.int8), col_len=jnp.zeros((2, N_COLORS), dtype=jnp.int8),
pile=jnp.full((N_COLORS, MAX_PILE_SIZE), NO_CARD, dtype=jnp.int8), pile=jnp.full((N_COLORS, MAX_PILE_SIZE), NO_CARD, dtype=jnp.int8),
pile_len=jnp.zeros((N_COLORS,), dtype=jnp.int8), pile_len=jnp.zeros((N_COLORS,), dtype=jnp.int8),
to_move=jnp.asarray(0, dtype=jnp.int8), to_move=jnp.asarray(first_player, dtype=jnp.int8),
just_discarded=jnp.asarray(NO_CARD, dtype=jnp.int8), just_discarded=jnp.asarray(NO_CARD, dtype=jnp.int8),
step_count=jnp.asarray(0, dtype=jnp.int32), step_count=jnp.asarray(0, dtype=jnp.int32),
done=jnp.asarray(False), done=jnp.asarray(False),
+244
View File
@@ -0,0 +1,244 @@
"""How exploitable is a frozen policy?
Beating an opponent on average and being hard to beat are different properties.
A policy can win the head-to-head and still have a habit a dedicated opponent can
farm -- rock-paper-scissors throwing rock 60% of the time beats a lot of people
and loses every time to anyone who noticed.
So: freeze the policy, train a fresh one from scratch whose only job is to beat
*that* policy, and see how far it gets. Near 0.5 means there was nothing to find.
Well above it means there was.
PPO self-play carries no Nash guarantee, so this is not a formality. The number it
produces is a lower bound on exploitability -- a stronger exploiter might find
more -- so it is only meaningful compared against the same exploiter budget spent
on another policy.
"""
from __future__ import annotations
import json
import time
from collections.abc import Callable
from pathlib import Path
import jax
import jax.numpy as jnp
from lost_cities_jax.match import MatchState, match_legal_action_mask, match_reset, match_step
from lost_cities_jax.match_ppo import (
FULL,
N_SEATS,
Ablation,
MatchActorCritic,
MatchTransition,
_seat_views,
create_match_train_state,
match_lead,
match_ppo_update,
)
from lost_cities_jax.obs import observation
from lost_cities_jax.ppo import (
ActorCritic,
JaxPPOConfig,
TrainState,
_append_jsonl,
_broadcast_done,
action_log_prob,
categorical_entropy,
compute_gae,
mask_logits,
masked_mean,
)
from lost_cities_jax.types import PLAY
FrozenPolicy = Callable[[MatchState, jax.Array], jax.Array]
"""(env batch, seat batch) -> greedy actions. Whatever architecture, hidden here."""
def match_frozen_policy(cfg: JaxPPOConfig, params, ablation: Ablation = FULL) -> FrozenPolicy:
model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
seats = jnp.arange(N_SEATS, dtype=jnp.int32)
def policy(env: MatchState, seat: jax.Array) -> jax.Array:
obs, critic = _seat_views(env, seats, ablation)
idx = seat[None, :, None]
logits, _ = model.apply(
params,
jnp.take_along_axis(obs, idx, axis=0)[0],
jnp.take_along_axis(critic, idx, axis=0)[0],
)
mask = jax.vmap(match_legal_action_mask)(env)
return jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)
return policy
def single_round_frozen_policy(cfg: JaxPPOConfig, params) -> FrozenPolicy:
"""The old stack: it sees one round at a time and has never heard of carry."""
model = ActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
def policy(env: MatchState, seat: jax.Array) -> jax.Array:
logits, _ = model.apply(params, jax.vmap(observation)(env.round, seat))
mask = jax.vmap(match_legal_action_mask)(env)
return jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)
return policy
def _reset_done(env: MatchState, seat: jax.Array, rng: jax.Array, batch: int):
env_key, seat_key = jax.random.split(rng)
fresh = jax.vmap(match_reset)(jax.random.split(env_key, batch))
next_env = jax.tree_util.tree_map(
lambda old, new: jnp.where(_broadcast_done(env.done, old), new, old), env, fresh
)
# Redraw the exploiter's seat too, so it learns to beat the frozen policy from
# either side rather than only the one it started in.
fresh_seat = jax.random.bernoulli(seat_key, 0.5, shape=(batch,)).astype(jnp.int32)
return next_env, jnp.where(env.done, fresh_seat, seat)
def make_exploiter_rollout_fn(cfg: JaxPPOConfig, frozen: FrozenPolicy):
batch = cfg.ppo.batch_games
seats = jnp.arange(N_SEATS, dtype=jnp.int32)
ablation = FULL
@jax.jit
def rollout_fn(state: TrainState, env: MatchState, seat: jax.Array, rng: jax.Array):
def body(carry, _):
env, seat, key = carry
key, act_key, reset_key = jax.random.split(key, 3)
obs, critic_obs = _seat_views(env, seats, ablation)
idx = seat[None, :, None]
my_obs = jnp.take_along_axis(obs, idx, axis=0)[0]
my_critic = jnp.take_along_axis(critic_obs, idx, axis=0)[0]
legal = jax.vmap(match_legal_action_mask)(env)
logits, value = state.apply_fn(state.params, my_obs, my_critic)
masked = mask_logits(logits, legal)
mine = jax.random.categorical(act_key, masked, axis=-1).astype(jnp.int32)
theirs = jax.lax.stop_gradient(frozen(env, 1 - seat))
to_move = env.round.to_move.astype(jnp.int32)
active = ~env.done
my_turn = to_move == seat
actions = jnp.where(my_turn, mine, theirs)
before = jax.vmap(match_lead)(env)
next_env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, actions)
after = jax.vmap(match_lead)(next_env)
# match_lead is seat 0's view; flip it when the exploiter sits in 1.
sign = jnp.where(seat == 0, 1.0, -1.0)
reward = jnp.where(active, sign * (after - before) / cfg.reward.terminal_scale, 0.0)
actor_mask = active & my_turn
transition = MatchTransition(
obs=my_obs,
critic_obs=my_critic,
legal_mask=legal,
action=actions,
log_prob=action_log_prob(masked, mine),
value=value,
reward=reward,
done=next_env.done,
active=active,
actor_mask=actor_mask,
entropy=categorical_entropy(masked),
play_action=((actions % 12) // 6 == PLAY) & actor_mask,
)
done = next_env.done
final_lead = sign * jax.vmap(match_lead)(next_env)
won = (done & active & (final_lead > 0)).astype(jnp.float32)
finished = (done & active).astype(jnp.float32)
next_env, seat = _reset_done(next_env, seat, reset_key, batch)
return (next_env, seat, key), (transition, won, finished)
(next_env, seat, rng), (transitions, won, finished) = jax.lax.scan(
body, (env, seat, rng), xs=None, length=cfg.ppo.rollout_steps
)
obs, critic_obs = _seat_views(next_env, seats, ablation)
idx = seat[None, :, None]
_, last_value = state.apply_fn(
state.params,
jnp.take_along_axis(obs, idx, axis=0)[0],
jnp.take_along_axis(critic_obs, idx, axis=0)[0],
)
played = jnp.sum(finished)
metrics = {
"exploiter_win_rate": jnp.sum(won) / jnp.maximum(played, 1.0),
"matches_completed": played,
"learner_actions": jnp.sum(transitions.actor_mask),
"entropy_mean": masked_mean(
transitions.entropy, transitions.actor_mask.astype(jnp.float32)
),
}
return next_env, seat, transitions, last_value, metrics
return rollout_fn
def train_exploiter(cfg: JaxPPOConfig, frozen: FrozenPolicy, run_dir: Path) -> TrainState:
"""Train a fresh policy whose only goal is to beat ``frozen``."""
rng = jax.random.PRNGKey(cfg.run.seed)
rng, init_key, reset_key, seat_key = jax.random.split(rng, 4)
state = create_match_train_state(cfg, init_key, Ablation())
env = jax.jit(jax.vmap(match_reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
seat = jax.random.bernoulli(seat_key, 0.5, shape=(cfg.ppo.batch_games,)).astype(jnp.int32)
rollout_fn = make_exploiter_rollout_fn(cfg, frozen)
@jax.jit
def iteration(state, env, seat, rng):
rng, roll_key, update_key = jax.random.split(rng, 3)
env, seat, transitions, last_value, metrics = rollout_fn(state, env, seat, roll_key)
advantages, returns = compute_gae(
transitions.reward,
transitions.value,
transitions.done,
cfg.ppo.gamma,
cfg.ppo.gae_lambda,
last_value,
)
state, update_metrics = match_ppo_update(
state, transitions, advantages, returns, update_key, cfg
)
return state, env, seat, rng, {**metrics, **update_metrics}
run_dir.mkdir(parents=True, exist_ok=True)
metrics_path = run_dir / "metrics.jsonl"
start = time.perf_counter()
for update in range(cfg.run.total_updates):
state, env, seat, rng, metrics = iteration(state, env, seat, rng)
jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
row = {k: float(v) for k, v in metrics.items()}
row["update"] = update
row["elapsed_seconds"] = time.perf_counter() - start
_append_jsonl(metrics_path, row)
if update % cfg.run.log_every == 0:
print(
json.dumps(
{
"update": update,
"exploiter_win_rate": row["exploiter_win_rate"],
"matches": row["matches_completed"],
},
sort_keys=True,
),
flush=True,
)
return state
__all__ = [
"make_exploiter_rollout_fn",
"match_frozen_policy",
"single_round_frozen_policy",
"train_exploiter",
]
+139
View File
@@ -0,0 +1,139 @@
"""Classic Lost Cities: three rounds, scores summed, highest total wins.
The single-round engine in ``engine.py`` is left alone -- it is the rules oracle
the TypeScript client is differential-tested against. This module wraps it.
Two rules only exist at the match level, both from the Kosmos rulebook:
- "If after three games you have the highest overall score, you win."
- "The player who has more points begins" the next game. Not alternating. The
rulebook says nothing about an exact tie, so we flip a coin: any deterministic
tie-break would hand one seat a standing edge in symmetric self-play, and the
agent would learn to steer for it.
All randomness is drawn in ``match_reset`` and stored in the state, so
``match_step`` stays deterministic and needs no PRNG key. That keeps the step
signature clean through every rollout/eval body, and it means a mirrored match
(same three deals, seats swapped, same coin flips) is just a seat relabel.
"""
from __future__ import annotations
from typing import NamedTuple
import jax
import jax.numpy as jnp
from lost_cities_jax.engine import board_score, legal_action_mask, reset_from_order, step
from lost_cities_jax.types import N_CARDS, State
N_ROUNDS = 3
class MatchState(NamedTuple):
round: State
"""The round being played right now."""
deck_orders: jax.Array
"""(N_ROUNDS, N_CARDS) -- every deal of the match, shuffled up front."""
coin_flips: jax.Array
"""(N_ROUNDS,) tie-breaking starters, used only when the scores are level."""
round_idx: jax.Array
carry: jax.Array
"""(2,) points banked by each player in the rounds already finished."""
done: jax.Array
def match_reset(rng: jax.Array) -> MatchState:
deck_key, coin_key = jax.random.split(rng)
deck_keys = jax.random.split(deck_key, N_ROUNDS)
deck_orders = jax.vmap(
lambda key: jax.random.permutation(key, jnp.arange(N_CARDS, dtype=jnp.int32), axis=0)
)(deck_keys).astype(jnp.int8)
coin_flips = jax.random.bernoulli(coin_key, 0.5, shape=(N_ROUNDS,)).astype(jnp.int8)
return match_reset_from(deck_orders, coin_flips)
def match_reset_from(deck_orders: jax.Array, coin_flips: jax.Array) -> MatchState:
"""Build a match from explicit deals and coin flips (mirrored eval, fixtures)."""
deck_orders = jnp.asarray(deck_orders, dtype=jnp.int8)
coin_flips = jnp.asarray(coin_flips, dtype=jnp.int8)
carry = jnp.zeros((2,), dtype=jnp.int32)
round_idx = jnp.asarray(0, dtype=jnp.int32)
return MatchState(
round=reset_from_order(deck_orders[0], first_player=starting_player(carry, 0, coin_flips)),
deck_orders=deck_orders,
coin_flips=coin_flips,
round_idx=round_idx,
carry=carry,
done=jnp.asarray(False),
)
def starting_player(carry: jax.Array, round_idx: jax.Array, coin_flips: jax.Array) -> jax.Array:
"""Whoever has banked more points leads; level scores fall back to the coin.
Round one needs no special case: ``carry`` is (0, 0) there, so the tie branch
already picks the coin flip, which is exactly the rulebook's arbitrary
"oldest player begins".
"""
lead = carry[0] - carry[1]
coin = coin_flips[round_idx].astype(jnp.int8)
ahead = jnp.where(lead > 0, jnp.int8(0), jnp.int8(1))
return jnp.where(lead == 0, coin, ahead)
def match_score(state: MatchState) -> jax.Array:
"""(2,) running totals: rounds already banked plus the board in play."""
return state.carry + board_score(state.round).astype(jnp.int32)
def match_legal_action_mask(state: MatchState) -> jax.Array:
return legal_action_mask(state.round) & ~state.done
def match_step(state: MatchState, action: jax.Array) -> tuple[MatchState, jax.Array, jax.Array]:
"""Play one ply. Rolls into the next round when the deck runs out.
Reward is zero until the third round ends, then it is each player's match
total. Intermediate rounds pay nothing: only the sum decides the match.
"""
played, _, _ = step(state.round, action)
round_over = played.done & ~state.done
is_final_round = state.round_idx >= N_ROUNDS - 1
advance = round_over & ~is_final_round
finished = round_over & is_final_round
banked = state.carry + board_score(played).astype(jnp.int32)
next_carry = jnp.where(advance, banked, state.carry)
next_round_idx = jnp.where(advance, state.round_idx + 1, state.round_idx)
# Dealt eagerly every ply; only kept on the plies that actually roll over.
fresh = reset_from_order(
state.deck_orders[next_round_idx],
first_player=starting_player(next_carry, next_round_idx, state.coin_flips),
)
next_round = jax.tree_util.tree_map(
lambda new, old: jnp.where(advance, new, old), fresh, played
)
next_state = MatchState(
round=next_round,
deck_orders=state.deck_orders,
coin_flips=state.coin_flips,
round_idx=next_round_idx,
carry=next_carry,
done=state.done | finished,
)
total = state.carry + board_score(played).astype(jnp.int32)
reward = jnp.where(finished, total.astype(jnp.float32), jnp.zeros((2,), dtype=jnp.float32))
return next_state, reward, next_state.done
+250
View File
@@ -0,0 +1,250 @@
"""Match-level evaluation, and the probe that says whether carry changed anything.
Single-deal win rate cannot tell you if the three-round rework worked. Two things
can:
- **Duplicate match play.** The same three deals and the same coin flips, played
from both seats. Deal luck cancels, so what is left is the policy difference.
- **The carry probe.** Drop the policy into round three holding a fixed deficit or
lead and watch what it does. A carry-blind policy plays a 60-point deficit
exactly like a 60-point lead. If the numbers below do not move with carry, the
rework bought nothing, no matter what the win rate says.
"""
from __future__ import annotations
import math
import jax
import jax.numpy as jnp
import numpy as np
from lost_cities_jax.engine import reset_from_order
from lost_cities_jax.match import (
N_ROUNDS,
MatchState,
match_legal_action_mask,
match_reset_from,
match_score,
match_step,
)
from lost_cities_jax.match_obs import match_critic_observation, match_observation
from lost_cities_jax.match_ppo import MatchActorCritic, match_lead
from lost_cities_jax.ppo import JaxPPOConfig, mask_logits
from lost_cities_jax.types import DRAW_DECK, MAX_STEPS, N_CARDS
MATCH_SCAN_STEPS = N_ROUNDS * MAX_STEPS
def match_bank(seed: int, matches: int) -> tuple[jnp.ndarray, jnp.ndarray]:
"""Deterministic deals and coin flips, so runs are comparable."""
rng = np.random.default_rng(seed)
decks = np.stack(
[np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)]) for _ in range(matches)]
)
coins = rng.integers(0, 2, size=(matches, N_ROUNDS))
return jnp.asarray(decks, dtype=jnp.int8), jnp.asarray(coins, dtype=jnp.int8)
def _wilson(wins: float, games: float) -> tuple[float, float]:
if games == 0:
return 0.0, 0.0
z = 1.96
p = wins / games
denom = 1 + z * z / games
centre = p + z * z / (2 * games)
margin = z * math.sqrt(p * (1 - p) / games + z * z / (4 * games * games))
return (centre - margin) / denom, (centre + margin) / denom
def make_match_runner(cfg: JaxPPOConfig, opponent_policy=None):
"""Play whole matches to the end. ``opponent_policy`` is a single-round policy."""
model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
@jax.jit
def run(params, env: MatchState, learner_seat: jax.Array, rng: jax.Array):
def body(carry, _):
env, key, deck_draws, plies = carry
key, opp_key = jax.random.split(key)
to_move = env.round.to_move.astype(jnp.int32)
obs = jax.vmap(match_observation)(env, to_move)
critic = jax.vmap(match_critic_observation)(env, to_move)
mask = jax.vmap(match_legal_action_mask)(env)
logits, _ = model.apply(params, obs, critic)
learner_action = jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)
if opponent_policy is None:
action = learner_action
else:
opp_keys = jax.random.split(opp_key, env.done.shape[0])
opp_action = jax.vmap(opponent_policy, in_axes=(0, 0, 0))(
env.round, to_move, opp_keys
)
action = jnp.where(to_move == learner_seat, learner_action, opp_action)
live = ~env.done
is_deck = (action % 6) == DRAW_DECK
deck_draws = deck_draws + (is_deck & live).astype(jnp.int32)
plies = plies + live.astype(jnp.int32)
env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
return (env, key, deck_draws, plies), None
n = env.done.shape[0]
zeros = jnp.zeros((n,), dtype=jnp.int32)
(env, _, deck_draws, plies), _ = jax.lax.scan(
body, (env, rng, zeros, zeros), xs=None, length=MATCH_SCAN_STEPS
)
return env, deck_draws, plies
return run
def match_evaluate(cfg: JaxPPOConfig, params, *, matches: int = 2000, seed: int = 20260715) -> dict:
"""Duplicate match play: every deal seen from both seats."""
decks, coins = match_bank(seed, matches)
runner = make_match_runner(cfg)
results = []
for seat in (0, 1):
env = jax.vmap(match_reset_from)(decks, coins)
seats = jnp.full((matches,), seat, dtype=jnp.int32)
final, deck_draws, plies = runner(params, env, seats, jax.random.PRNGKey(0))
totals = np.asarray(jax.vmap(match_score)(final))
lead = totals[:, seat] - totals[:, 1 - seat]
results.append(
{
"lead": lead,
"plies": np.asarray(plies),
"deck_share": np.asarray(deck_draws) / np.maximum(np.asarray(plies), 1),
"done": np.asarray(final.done),
}
)
lead = np.concatenate([r["lead"] for r in results])
games = float(lead.size)
wins = float((lead > 0).sum())
ties = float((lead == 0).sum())
low, high = _wilson(wins, games)
return {
"matches": games,
"wins": wins,
"ties": ties,
"losses": games - wins - ties,
"match_win_rate": wins / games,
"wilson_low": low,
"wilson_high": high,
"mean_total_lead": float(lead.mean()),
"mean_plies": float(np.concatenate([r["plies"] for r in results]).mean()),
"unfinished_rate": float(
1.0 - np.concatenate([r["done"] for r in results]).astype(np.float32).mean()
),
}
def head_to_head(
cfg: JaxPPOConfig, params_a, params_b, *, matches: int = 4096, seed: int = 20260717
) -> dict:
"""Duplicate matches between two policies: same deals, same coins, both seats.
Every deal is played twice with the seats swapped, so deal luck cancels and
what is left is the difference between the two policies.
"""
decks, coins = match_bank(seed, matches)
model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
@jax.jit
def run(env: MatchState, a_seat: jax.Array):
def body(carry, _):
env, _unused = carry
to_move = env.round.to_move.astype(jnp.int32)
obs = jax.vmap(match_observation)(env, to_move)
critic = jax.vmap(match_critic_observation)(env, to_move)
mask = jax.vmap(match_legal_action_mask)(env)
logits_a, _ = model.apply(params_a, obs, critic)
logits_b, _ = model.apply(params_b, obs, critic)
act_a = jnp.argmax(mask_logits(logits_a, mask), axis=-1).astype(jnp.int32)
act_b = jnp.argmax(mask_logits(logits_b, mask), axis=-1).astype(jnp.int32)
action = jnp.where(to_move == a_seat, act_a, act_b)
env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
return (env, _unused), None
(env, _), _ = jax.lax.scan(body, (env, jnp.int32(0)), xs=None, length=MATCH_SCAN_STEPS)
return env
leads = []
for seat in (0, 1):
env = jax.vmap(match_reset_from)(decks, coins)
final = run(env, jnp.full((matches,), seat, dtype=jnp.int32))
totals = np.asarray(jax.vmap(match_score)(final))
leads.append(totals[:, seat] - totals[:, 1 - seat])
lead = np.concatenate(leads)
games = float(lead.size)
wins = float((lead > 0).sum())
ties = float((lead == 0).sum())
low, high = _wilson(wins, games)
return {
"matches": games,
"a_win_rate": wins / games,
"wilson_low": low,
"wilson_high": high,
"ties": ties,
"a_mean_lead": float(lead.mean()),
}
def carry_probe(
cfg: JaxPPOConfig,
params,
*,
carry_levels=(-60, -25, -1, 1, 25, 60),
matches: int = 512,
seed: int = 20260716,
) -> list[dict]:
"""Start the policy in round three holding ``carry`` and see what changes.
A policy that reads carry should open more expeditions and lean on wagers when
it is behind (it has to swing), and shut down when it is ahead. Flat rows mean
the policy is ignoring the one thing the three-round rework added.
"""
decks, coins = match_bank(seed, matches)
runner = make_match_runner(cfg)
rows = []
for level in carry_levels:
base = jax.vmap(match_reset_from)(decks, coins)
carry = jnp.tile(jnp.asarray([level, 0], dtype=jnp.int32), (matches, 1))
round_idx = jnp.full((matches,), N_ROUNDS - 1, dtype=jnp.int32)
# Whoever is ahead leads the round, exactly as the match layer would.
first = jnp.int8(0 if level > 0 else 1)
rounds = jax.vmap(reset_from_order, in_axes=(0, None))(decks[:, N_ROUNDS - 1], first)
env = base._replace(round=rounds, carry=carry, round_idx=round_idx)
final, deck_draws, plies = runner(
params, env, jnp.zeros((matches,), jnp.int32), jax.random.PRNGKey(1)
)
lead = np.asarray(jax.vmap(match_lead)(final))
opened = np.asarray(jnp.sum(final.round.col_len[:, 0, :] > 0, axis=-1))
wagers = np.asarray(jnp.sum(final.round.col_hs[:, 0, :], axis=-1))
deck_share = np.asarray(deck_draws) / np.maximum(np.asarray(plies), 1)
rows.append(
{
"carry": level,
"win_rate": float((lead > 0).mean()),
"opened_colors": float(opened.mean()),
"wagers_played": float(wagers.mean()),
"deck_race_rate": float(deck_share.mean()),
"mean_plies": float(np.asarray(plies).mean()),
}
)
return rows
+198
View File
@@ -0,0 +1,198 @@
"""Player-view observation for a three-round match.
Everything the single-round observation had, plus the four things a match policy
cannot play without:
- **carry**, as a scalar *and* a one-hot over bins. Round three is a threshold
problem -- "win by 41 or lose" plays nothing like "win by 39" -- and a lone
scalar makes the network learn a cliff from a smooth input. (The old
``score_diff`` divided by ``MAX_ABS_SCORE`` = 780, which squashed a decisive
50-point lead to 0.06. That scaling is fixed here too.)
- **which round it is**, and how many are left.
- **whose turn it is.** The single-round obs never had this, yet the critic is
trained on opponent-turn states as well, leaving it to recover turn parity
from ``step_count / 400``.
- **the deck clock.** A round ends when the last deck card is drawn, so turn
parity decides who gets the final placement -- and players bend that parity by
drawing from discard piles instead.
Live points per colour are split by where the card is (hand / discard pile /
unseen) because a discard pile is public and recoverable: lumping it in with the
unseen cards hides the high cards sitting in plain view.
"""
from __future__ import annotations
import jax
import jax.numpy as jnp
from lost_cities_jax.engine import board_score
from lost_cities_jax.match import N_ROUNDS, MatchState
from lost_cities_jax.obs import observation
from lost_cities_jax.types import (
CARDS_PER_COLOR,
LOC_DECK,
LOC_DISCARD,
LOC_P0_HAND,
N_CARDS,
N_COLORS,
OBS_DIM,
)
# Bin edges for the carry one-hot. Packed tightly around zero: that is where the
# round-three decision actually flips.
CARRY_BIN_EDGES = jnp.asarray([-60.0, -30.0, -12.0, -1.0, 1.0, 12.0, 30.0, 60.0])
N_CARRY_BINS = CARRY_BIN_EDGES.shape[0] + 1
CARRY_SCALE = 75.0
"""Typical round margin, not the theoretical 780 maximum."""
MAX_COLOR_POINTS = 54.0
"""2+3+...+10, the most a single expedition can be worth before multipliers."""
N_MATCH_SCALARS = (
1 # carry, scaled
+ N_CARRY_BINS # carry, binned
+ N_ROUNDS # which round
+ 1 # rounds left
+ 1 # is it my turn
+ 1 # did I lead this round
+ 1 # do I take the last deck card at the current parity
+ 2 * N_COLORS * 3 # live points: both players x colour x {hand, discard, unseen}
)
MATCH_OBS_DIM = OBS_DIM + N_MATCH_SCALARS
def _live_points(state, viewer: jax.Array, subject: jax.Array) -> jax.Array:
"""Points still reachable for ``subject``, split by where the card sits.
Returns ``float32[N_COLORS, 3]`` over {in hand, in a discard pile, unseen}.
Seen through ``viewer``'s eyes: a card in the opponent's hand only counts as
"in hand" if it is public, otherwise it is unseen.
"""
ids = jnp.arange(N_CARDS, dtype=jnp.int32)
colors = ids // CARDS_PER_COLOR
slots = ids % CARDS_PER_COLOR
ranks = jnp.where(slots >= 3, slots - 1, 0).astype(jnp.float32)
loc = state.card_loc.astype(jnp.int32)
subject_hand = loc == (LOC_P0_HAND + subject)
is_mine = subject == viewer
# An ascending column can only take cards above its current top.
above_top = ranks > state.col_top[subject][colors].astype(jnp.float32)
in_hand = subject_hand & (is_mine | state.hand_public)
in_discard = loc == LOC_DISCARD
unseen = (loc == LOC_DECK) | (subject_hand & ~is_mine & ~state.hand_public)
# Cards in the *other* player's hidden hand are unseen to the viewer too, but
# they are not reachable by the subject, so they are deliberately excluded.
def by_color(mask: jax.Array) -> jax.Array:
weighted = jnp.where(mask & above_top, ranks, 0.0)
return jax.ops.segment_sum(weighted, colors, num_segments=N_COLORS)
stacked = jnp.stack([by_color(in_hand), by_color(in_discard), by_color(unseen)], axis=1)
return stacked / MAX_COLOR_POINTS
def match_observation(state: MatchState, player: jax.Array) -> jax.Array:
"""Return ``float32[MATCH_OBS_DIM]`` from ``player``'s seat."""
player = jnp.asarray(player, dtype=jnp.int32)
opponent = 1 - player
round_state = state.round
base = observation(round_state, player)
totals = state.carry + board_score(round_state).astype(jnp.int32)
lead = (totals[player] - totals[opponent]).astype(jnp.float32)
carry_scaled = jnp.clip(lead / CARRY_SCALE, -2.0, 2.0).reshape((1,))
carry_bins = jax.nn.one_hot(jnp.digitize(lead, CARRY_BIN_EDGES), N_CARRY_BINS)
round_onehot = jax.nn.one_hot(state.round_idx, N_ROUNDS)
rounds_left = ((N_ROUNDS - 1 - state.round_idx).astype(jnp.float32) / (N_ROUNDS - 1)).reshape(
(1,)
)
to_move = round_state.to_move.astype(jnp.int32)
my_turn = (to_move == player).astype(jnp.float32).reshape((1,))
# to_move flips every ply, so the round's opener is recoverable from parity.
parity = round_state.step_count.astype(jnp.int32) & 1
round_opener = to_move ^ parity
i_opened = (round_opener == player).astype(jnp.float32).reshape((1,))
# If both players drew from the deck from here, the last deck card falls to
# whoever is on move after `remaining - 1` more plies.
remaining = jnp.maximum(N_CARDS - round_state.draw_ptr, 1).astype(jnp.int32)
last_drawer = to_move ^ ((remaining - 1) & 1)
i_take_last = (last_drawer == player).astype(jnp.float32).reshape((1,))
live_me = _live_points(round_state, player, player).reshape((-1,))
live_opp = _live_points(round_state, player, opponent).reshape((-1,))
extra = jnp.concatenate(
[
carry_scaled,
carry_bins,
round_onehot,
rounds_left,
my_turn,
i_opened,
i_take_last,
live_me,
live_opp,
],
axis=0,
).astype(jnp.float32)
return jnp.concatenate([base, extra], axis=0).reshape((MATCH_OBS_DIM,))
CRITIC_EXTRA_CHANNELS = 3
MATCH_CRITIC_OBS_DIM = MATCH_OBS_DIM + N_CARDS * CRITIC_EXTRA_CHANNELS
def match_critic_observation(state: MatchState, player: jax.Array) -> jax.Array:
"""The same view, plus the hidden state. Training only -- never played from.
A value baseline may condition on anything the action does not depend on
without biasing the policy gradient, and the deal is exactly the noise that
swamps a match-terminal reward. So the critic gets the opponent's hand and
the deck *in order* -- it knows what is coming -- while the actor keeps the
masked view. The two run on separate trunks so none of this can leak into
the logits.
"""
player = jnp.asarray(player, dtype=jnp.int32)
opponent = 1 - player
round_state = state.round
loc = round_state.card_loc.astype(jnp.int32)
opp_hand = (loc == (LOC_P0_HAND + opponent)).astype(jnp.float32)
in_deck = (loc == LOC_DECK).astype(jnp.float32)
# Position of each card in the undrawn deck, so the critic can see the order
# rather than just the multiset.
deck_order = round_state.deck_order.astype(jnp.int32)
position = (
jnp.zeros((N_CARDS,), dtype=jnp.int32)
.at[deck_order]
.set(jnp.arange(N_CARDS, dtype=jnp.int32))
)
depth = (position - round_state.draw_ptr).astype(jnp.float32) / N_CARDS
deck_depth = jnp.where(in_deck > 0, depth, 0.0)
hidden = jnp.stack([opp_hand, in_deck, deck_depth], axis=1).reshape((-1,))
base = match_observation(state, player)
return jnp.concatenate([base, hidden], axis=0).reshape((MATCH_CRITIC_OBS_DIM,))
__all__ = [
"MATCH_CRITIC_OBS_DIM",
"MATCH_OBS_DIM",
"N_CARRY_BINS",
"match_critic_observation",
"match_observation",
]
+480
View File
@@ -0,0 +1,480 @@
"""PPO over three-round matches.
Three things differ from the single-round trainer beyond the env swap:
**Asymmetric actor-critic.** The critic never plays, so it is fed the hidden
state (opponent's hand, deck in order) while the actor keeps the masked view. A
state-value baseline may condition on anything action-independent without biasing
the policy gradient, and deal luck is exactly what makes a match-terminal reward
hard to learn from. The trunks are separate so the privileged features cannot
reach the logits.
**Both seats train.** In self-play one network picks both players' moves, but the
single-round trainer stop_gradients the opponent seat and throws those plies away
-- half of every game. Here every ply emits a transition for both seats, each
with its own seat-relative view, value and reward; the actor mask keeps the
policy loss on the seat that actually moved, while the critic learns from both.
Folding the seat axis into the batch keeps each seat's GAE chain independent.
**Reward is the match total, paid densely.** Each ply pays the points by which it
moved the running match difference; at gamma=1 that telescopes to the final total
difference, so the objective is exactly the rulebook's -- score more across the
three rounds.
This used to be ``tanh(total / terminal_scale)`` at the end of round three, to
optimise P(win match) rather than expected total. That was over-engineered. The
two differ only in risk attitude, and risk attitude is worth almost nothing here:
a policy told to gamble when it trails by 20 going into round three *loses* to a
greedy clone over 6144 duplicate matches (0.482), and one that gambles only when
it trails by 40 breaks even (0.498). Lost Cities' variance levers -- a marginal
wager, a marginal expedition -- cost more expected margin than the convexity they
buy. Measured ceiling on the whole carry-conditioning idea: under one win-rate
point. The dense linear reward also gives credit assignment every ply instead of
one bounded number per ~160.
``carry`` stays in the observation: it costs nothing, it is the state the
start-player rule keys off, and the critic reads it cleanly.
"""
from __future__ import annotations
import json
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import NamedTuple
import flax.linen as nn
import jax
import jax.numpy as jnp
import optax
from lost_cities_jax.match import (
MatchState,
match_legal_action_mask,
match_reset,
match_score,
match_step,
)
from lost_cities_jax.match_obs import (
MATCH_CRITIC_OBS_DIM,
MATCH_OBS_DIM,
match_critic_observation,
match_observation,
)
from lost_cities_jax.obs import observation
from lost_cities_jax.ppo import (
EpisodeEnd,
JaxPPOConfig,
TrainState,
_append_jsonl,
_broadcast_done,
_create_run_dir,
_write_json,
action_log_prob,
categorical_entropy,
compute_gae,
mask_logits,
masked_mean,
restore_checkpoint,
save_checkpoint,
shaping_coefficient,
)
from lost_cities_jax.types import MAX_STEPS, N_ACTIONS, OBS_DIM, PLAY
N_SEATS = 2
@dataclass
class Ablation:
"""Switches for measuring what each piece of the match stack is worth."""
privileged_critic: bool = True
"""On: the critic also sees the opponent's hand and the deck in order.
Its worth flips with scale, so do not trust a cheap ablation here. Ablated at
39.3M learner actions, switching it OFF *won* 0.5160 [0.505, 0.527] -- the
obvious reading being that a critic fitting V(full state) rather than
E[return | masked obs] hands the actor advantage noise it cannot act on. But
head to head at 131M, both sides trained identically, OFF *loses*: 0.4633
[0.453, 0.474]. The privileged critic earns its keep once there is enough data
to fit it. Defaults on for that reason, and the small-scale result is kept as
a warning about ablating at a budget you do not intend to ship.
"""
both_seats: bool = True
"""Off: the opponent seat's plies are discarded, as the old trainer did."""
match_obs: bool = True
"""Off: the actor gets the plain single-round observation, carry and all the
rest of the match features withheld."""
def actor_obs_dim(self) -> int:
return MATCH_OBS_DIM if self.match_obs else OBS_DIM
def critic_obs_dim(self) -> int:
if not self.privileged_critic:
return self.actor_obs_dim()
return MATCH_CRITIC_OBS_DIM
def label(self) -> str:
parts = []
if self.privileged_critic:
parts.append("privileged_critic")
if not self.both_seats:
parts.append("no-both_seats")
if not self.match_obs:
parts.append("no-match_obs")
return "default" if not parts else "+".join(parts)
FULL = Ablation()
class MatchTransition(NamedTuple):
obs: jax.Array
critic_obs: jax.Array
legal_mask: jax.Array
action: jax.Array
log_prob: jax.Array
value: jax.Array
reward: jax.Array
done: jax.Array
active: jax.Array
actor_mask: jax.Array
entropy: jax.Array
play_action: jax.Array
class MatchActorCritic(nn.Module):
"""Separate trunks, so the critic's privileged input cannot reach the logits."""
hidden_size: int = 512
num_layers: int = 3
@nn.compact
def __call__(self, obs: jax.Array, critic_obs: jax.Array) -> tuple[jax.Array, jax.Array]:
x = obs
for _ in range(self.num_layers):
x = nn.relu(nn.Dense(self.hidden_size)(x))
logits = nn.Dense(N_ACTIONS)(x)
v = critic_obs
for _ in range(self.num_layers):
v = nn.relu(nn.Dense(self.hidden_size)(v))
value = nn.Dense(1)(v)
return logits, jnp.squeeze(value, axis=-1)
def create_match_train_state(
cfg: JaxPPOConfig, rng: jax.Array, ablation: Ablation = FULL
) -> TrainState:
model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
params = model.init(
rng,
jnp.zeros((1, ablation.actor_obs_dim()), dtype=jnp.float32),
jnp.zeros((1, ablation.critic_obs_dim()), dtype=jnp.float32),
)
tx = optax.chain(
optax.clip_by_global_norm(cfg.ppo.max_grad_norm),
optax.adam(cfg.ppo.learning_rate),
)
return TrainState.create(apply_fn=model.apply, params=params, tx=tx)
def match_lead(env: MatchState) -> jax.Array:
"""Running match total from seat 0's side, in raw points."""
totals = match_score(env)
return (totals[0] - totals[1]).astype(jnp.float32)
def reset_done_matches(env: MatchState, rng: jax.Array, batch: int) -> MatchState:
fresh = jax.vmap(match_reset)(jax.random.split(rng, batch))
return jax.tree_util.tree_map(
lambda old, new: jnp.where(_broadcast_done(env.done, old), new, old), env, fresh
)
def _seat_views(env: MatchState, seats: jax.Array, ablation: Ablation = FULL):
actor_fn = match_observation if ablation.match_obs else (lambda m, p: observation(m.round, p))
obs = jax.vmap(lambda s: jax.vmap(actor_fn, in_axes=(0, None))(env, s))(seats)
if not ablation.privileged_critic:
return obs, obs
critic = jax.vmap(lambda s: jax.vmap(match_critic_observation, in_axes=(0, None))(env, s))(
seats
)
return obs, critic
def _episode_end(env: MatchState, terminal: jax.Array, episode_return: jax.Array) -> EpisodeEnd:
lead = jax.vmap(match_lead)(env)
opened = jnp.sum(env.round.col_len[:, 0, :] > 0, axis=-1)
return EpisodeEnd(
done=terminal,
episode_return=episode_return,
length=env.round.step_count,
opened_colors=opened,
positive_expeditions=lead.astype(jnp.int32), # final match lead, in points
hit_max_steps=env.round.step_count >= MAX_STEPS,
)
def match_episode_metrics(transitions: MatchTransition, episodes: EpisodeEnd) -> dict:
done = episodes.done.astype(jnp.float32)
completed = jnp.sum(done)
denom = jnp.maximum(completed, 1.0)
actor_count = jnp.sum(transitions.actor_mask)
lead = episodes.positive_expeditions.astype(jnp.float32)
return {
"return_mean": jnp.sum(episodes.episode_return * done) / denom,
"match_lead_mean": jnp.sum(lead * done) / denom,
"match_win_rate": jnp.sum((lead > 0).astype(jnp.float32) * done) / denom,
"final_round_length_mean": jnp.sum(episodes.length.astype(jnp.float32) * done) / denom,
"max_steps_rate": jnp.sum(episodes.hit_max_steps.astype(jnp.float32) * done) / denom,
"play_action_rate": jnp.sum(transitions.play_action) / jnp.maximum(actor_count, 1),
"entropy_mean": masked_mean(
transitions.entropy, transitions.actor_mask.astype(jnp.float32)
),
"learner_actions": actor_count,
"matches_completed": completed,
}
def make_match_rollout_fn(cfg: JaxPPOConfig, ablation: Ablation = FULL):
batch = cfg.ppo.batch_games
seats = jnp.arange(N_SEATS, dtype=jnp.int32)
@jax.jit
def rollout_fn(state: TrainState, env: MatchState, rng: jax.Array, shaping_coef: jax.Array):
def body(carry, _):
env, episode_return, key = carry
key, act_key, reset_key = jax.random.split(key, 3)
obs, critic_obs = _seat_views(env, seats, ablation)
legal = jax.vmap(match_legal_action_mask)(env)
legal_both = jnp.broadcast_to(legal, (N_SEATS, batch, N_ACTIONS))
logits, value = state.apply_fn(state.params, obs, critic_obs)
masked = mask_logits(logits, legal_both)
sampled = jax.random.categorical(act_key, masked, axis=-1).astype(jnp.int32)
log_prob = jax.vmap(action_log_prob)(masked, sampled)
entropy = jax.vmap(categorical_entropy)(masked)
to_move = env.round.to_move.astype(jnp.int32)
active = ~env.done
actor_mask = (seats[:, None] == to_move[None, :]) & active[None, :]
if not ablation.both_seats:
# The old trainer kept one seat and stop_gradiented the other.
keep = (seats[:, None] == 0) & jnp.ones_like(active)[None, :]
actor_mask = actor_mask & keep
actions = jnp.take_along_axis(sampled, to_move[None, :], axis=0)[0]
before = jax.vmap(match_lead)(env)
next_env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, actions)
after = jax.vmap(match_lead)(next_env)
done = next_env.done
terminal = active & done
# Every ply pays the points by which it moved the match total. At
# gamma=1 that telescopes to the final total difference, so the
# objective is exactly "score more across the three rounds" -- just
# handed out densely instead of once, 150 plies later.
reward0 = jnp.where(active, (after - before) / cfg.reward.terminal_scale, 0.0)
# Zero-sum: seat 1 sees exactly the negation.
reward = jnp.stack([reward0, -reward0], axis=0)
played = ((actions % 12) // 6 == PLAY)[None, :] & actor_mask
transition = MatchTransition(
obs=obs,
critic_obs=critic_obs,
legal_mask=legal_both,
action=jnp.broadcast_to(actions, (N_SEATS, batch)),
log_prob=log_prob,
value=value,
reward=reward,
done=jnp.broadcast_to(done, (N_SEATS, batch)),
active=(
jnp.broadcast_to(active, (N_SEATS, batch))
if ablation.both_seats
else jnp.broadcast_to(active, (N_SEATS, batch)) & (seats[:, None] == 0)
),
actor_mask=actor_mask,
entropy=entropy,
play_action=played,
)
episode_return = episode_return + reward0
episode = _episode_end(next_env, terminal, episode_return)
next_env = reset_done_matches(next_env, reset_key, batch)
episode_return = jnp.where(done, 0.0, episode_return)
return (next_env, episode_return, key), (transition, episode)
init_return = jnp.zeros((batch,), dtype=jnp.float32)
(next_env, _, rng), (transitions, episodes) = jax.lax.scan(
body, (env, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
)
# Fold the seat axis into the batch: each column is then an independent
# episode chain, which is exactly what GAE wants.
transitions = jax.tree_util.tree_map(
lambda x: x.reshape((x.shape[0], N_SEATS * batch, *x.shape[3:])), transitions
)
final_obs, final_critic = _seat_views(next_env, seats, ablation)
_, last_value = state.apply_fn(state.params, final_obs, final_critic)
last_value = last_value.reshape((N_SEATS * batch,))
return next_env, transitions, last_value, match_episode_metrics(transitions, episodes)
return rollout_fn
def match_ppo_update(
state: TrainState,
transitions: MatchTransition,
advantages: jax.Array,
returns: jax.Array,
rng: jax.Array,
cfg: JaxPPOConfig,
) -> tuple[TrainState, dict[str, jax.Array]]:
batch_size = int(transitions.reward.size)
minibatch_size = batch_size // cfg.ppo.minibatches
flat = jax.tree_util.tree_map(lambda x: x.reshape((-1, *x.shape[2:])), transitions)
flat_advantages = advantages.reshape((batch_size,))
flat_returns = returns.reshape((batch_size,))
actor_mask = flat.actor_mask.astype(jnp.float32)
adv_mean = masked_mean(flat_advantages, actor_mask)
adv_std = jnp.sqrt(masked_mean((flat_advantages - adv_mean) ** 2, actor_mask) + 1.0e-8)
flat_advantages = (flat_advantages - adv_mean) / adv_std
def loss_fn(params, mb, mb_adv, mb_returns):
logits, value = state.apply_fn(params, mb.obs, mb.critic_obs)
masked_logits = mask_logits(logits, mb.legal_mask)
log_prob = action_log_prob(masked_logits, mb.action)
entropy = categorical_entropy(masked_logits)
ratio = jnp.exp(log_prob - mb.log_prob)
actor_weight = mb.actor_mask.astype(jnp.float32)
active_weight = mb.active.astype(jnp.float32)
unclipped = ratio * mb_adv
clipped = jnp.clip(ratio, 1.0 - cfg.ppo.clip_epsilon, 1.0 + cfg.ppo.clip_epsilon) * mb_adv
actor_loss = -masked_mean(jnp.minimum(unclipped, clipped), actor_weight)
value_loss = masked_mean((mb_returns - value) ** 2, active_weight)
entropy_loss = masked_mean(entropy, actor_weight)
loss = actor_loss + cfg.ppo.value_coef * value_loss - cfg.ppo.entropy_coef * entropy_loss
return loss, {
"loss": loss,
"actor_loss": actor_loss,
"value_loss": value_loss,
"entropy_loss": entropy_loss,
"approx_kl": masked_mean(mb.log_prob - log_prob, actor_weight),
}
def epoch_update(carry, key):
train_state = carry
permutation = jax.random.permutation(key, batch_size)
def minibatch_update(carry, idx):
train_state, accum = carry
mb_idx = jax.lax.dynamic_slice(permutation, (idx * minibatch_size,), (minibatch_size,))
mb = jax.tree_util.tree_map(lambda x: x[mb_idx], flat)
(_, metrics), grads = jax.value_and_grad(loss_fn, has_aux=True)(
train_state.params, mb, flat_advantages[mb_idx], flat_returns[mb_idx]
)
train_state = train_state.apply_gradients(grads=grads)
accum = jax.tree_util.tree_map(lambda a, b: a + b, accum, metrics)
return (train_state, accum), None
zeros = {
key: jnp.array(0.0)
for key in ("loss", "actor_loss", "value_loss", "entropy_loss", "approx_kl")
}
(train_state, metrics), _ = jax.lax.scan(
minibatch_update, (train_state, zeros), jnp.arange(cfg.ppo.minibatches)
)
metrics = jax.tree_util.tree_map(lambda x: x / cfg.ppo.minibatches, metrics)
return train_state, metrics
keys = jax.random.split(rng, cfg.ppo.epochs)
state, metrics = jax.lax.scan(epoch_update, state, keys)
return state, jax.tree_util.tree_map(lambda x: jnp.mean(x, axis=0), metrics)
def make_match_train_iteration(cfg: JaxPPOConfig, ablation: Ablation = FULL):
rollout_fn = make_match_rollout_fn(cfg, ablation)
@jax.jit
def train_iteration(
state: TrainState, env: MatchState, rng: jax.Array, shaping_coef: jax.Array
):
rng, rollout_key, update_key = jax.random.split(rng, 3)
env, transitions, last_value, stats = rollout_fn(state, env, rollout_key, shaping_coef)
advantages, returns = compute_gae(
transitions.reward,
transitions.value,
transitions.done,
cfg.ppo.gamma,
cfg.ppo.gae_lambda,
last_value,
)
state, update_metrics = match_ppo_update(
state, transitions, advantages, returns, update_key, cfg
)
return state, env, rng, {**stats, **update_metrics}
return train_iteration
def match_train(cfg: JaxPPOConfig, *, resume: str | None = None, ablation: Ablation = FULL) -> Path:
run_dir = _create_run_dir(cfg)
_write_json(run_dir / "config.json", asdict(cfg))
metrics_path = run_dir / "metrics.jsonl"
rng = jax.random.PRNGKey(cfg.run.seed)
rng, init_key, reset_key = jax.random.split(rng, 3)
state = create_match_train_state(cfg, init_key, ablation)
if resume:
state = restore_checkpoint(Path(resume), state)
env = jax.jit(jax.vmap(match_reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
train_iteration = make_match_train_iteration(cfg, ablation)
start = time.perf_counter()
learner_actions_seen = 0
for update in range(cfg.run.total_updates):
shaping_coef = jnp.asarray(shaping_coefficient(cfg, learner_actions_seen), jnp.float32)
iter_start = time.perf_counter()
state, env, rng, metrics = train_iteration(state, env, rng, shaping_coef)
jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
row = {
key: float(value) if getattr(value, "shape", ()) == () else value.tolist()
for key, value in metrics.items()
}
learner_actions_seen += int(row["learner_actions"])
row.update(
update=update,
shaping_coef=float(shaping_coef),
learner_actions_total=learner_actions_seen,
iteration_seconds=time.perf_counter() - iter_start,
elapsed_seconds=time.perf_counter() - start,
)
_append_jsonl(metrics_path, row)
if update % cfg.run.log_every == 0:
print(
json.dumps(
{
"update": update,
"match_win_rate": row["match_win_rate"],
"match_lead_mean": row["match_lead_mean"],
"matches_completed": row["matches_completed"],
"entropy_mean": row["entropy_mean"],
},
sort_keys=True,
),
flush=True,
)
if cfg.run.checkpoint_every and (update + 1) % cfg.run.checkpoint_every == 0:
save_checkpoint(run_dir / "latest", state, cfg)
save_checkpoint(run_dir / "latest", state, cfg)
return run_dir
+165 -87
View File
@@ -77,6 +77,14 @@ class PPOHyperConfig:
@dataclass @dataclass
class RewardConfig: class RewardConfig:
"""Terminal reward plus potential-based shaping on the score difference.
``potential_shaping_anneal_steps`` counts **learner actions**, not scan
steps. It used to be fed padded scan steps, which made a 5M-step anneal
expire inside the first two updates; any earlier conclusion that shaping
does not help was drawn with shaping effectively off.
"""
terminal_scale: float = 50.0 terminal_scale: float = 50.0
potential_shaping_initial: float = 1.0 potential_shaping_initial: float = 1.0
potential_shaping_final: float = 0.0 potential_shaping_final: float = 0.0
@@ -134,6 +142,23 @@ class Transition(NamedTuple):
play_action: jax.Array play_action: jax.Array
class EpisodeEnd(NamedTuple):
"""Per-step episode-completion records emitted by the rollout scan.
With in-scan auto-reset an env slot hosts several episodes per rollout, so
``final_env`` is mid-episode and summing rewards over the scan axis mixes
episodes. Every field here is masked by ``done``: read it only where
``done`` is set.
"""
done: jax.Array
episode_return: jax.Array
length: jax.Array
opened_colors: jax.Array
positive_expeditions: jax.Array
hit_max_steps: jax.Array
class EvalBatch(NamedTuple): class EvalBatch(NamedTuple):
wins: jax.Array wins: jax.Array
losses: jax.Array losses: jax.Array
@@ -218,14 +243,20 @@ def train(
train_iteration = make_train_iteration(cfg, opponent_policy) train_iteration = make_train_iteration(cfg, opponent_policy)
start = time.perf_counter() start = time.perf_counter()
# Anneal against learner actions taken, not scan steps issued: most scan
# steps used to be spent stepping already-done envs, which made the anneal
# finish within the first couple of updates. The count lags by one update
# because it comes back as a device metric.
learner_actions_seen = 0
for update in range(cfg.run.total_updates): for update in range(cfg.run.total_updates):
shaping_coef = shaping_coefficient( shaping_coef = shaping_coefficient(cfg, learner_actions_seen)
cfg, update * cfg.ppo.batch_games * cfg.ppo.rollout_steps
)
iter_start = time.perf_counter() iter_start = time.perf_counter()
state, env_state, rng, metrics = train_iteration(state, env_state, rng, shaping_coef) state, env_state, rng, metrics = train_iteration(state, env_state, rng, shaping_coef)
jax.tree_util.tree_leaves(metrics)[0].block_until_ready() jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
row = _metrics_to_row(metrics, update, shaping_coef, cfg) row = _metrics_to_row(metrics, update, shaping_coef, cfg)
learner_actions_seen += int(row["learner_actions"])
row["learner_actions_total"] = learner_actions_seen
row["iteration_seconds"] = time.perf_counter() - iter_start row["iteration_seconds"] = time.perf_counter() - iter_start
row["elapsed_seconds"] = time.perf_counter() - start row["elapsed_seconds"] = time.perf_counter() - start
_append_jsonl(metrics_path, row) _append_jsonl(metrics_path, row)
@@ -260,8 +291,10 @@ def make_train_iteration(cfg: JaxPPOConfig, opponent_policy):
def train_iteration( def train_iteration(
state: TrainState, env_state: State, rng: jax.Array, shaping_coef: jax.Array state: TrainState, env_state: State, rng: jax.Array, shaping_coef: jax.Array
): ):
rng, rollout_key, update_key, reset_key = jax.random.split(rng, 4) rng, rollout_key, update_key = jax.random.split(rng, 3)
env_state, transitions, rollout_metrics = rollout_fn( # The rollout resets finished envs in-scan, so env_state never comes
# back done and needs no reset here.
env_state, transitions, last_value, rollout_stats = rollout_fn(
state, env_state, rollout_key, shaping_coef state, env_state, rollout_key, shaping_coef
) )
advantages, returns = compute_gae( advantages, returns = compute_gae(
@@ -270,10 +303,10 @@ def make_train_iteration(cfg: JaxPPOConfig, opponent_policy):
transitions.done, transitions.done,
cfg.ppo.gamma, cfg.ppo.gamma,
cfg.ppo.gae_lambda, cfg.ppo.gae_lambda,
last_value,
) )
state, update_metrics = ppo_update(state, transitions, advantages, returns, update_key, cfg) state, update_metrics = ppo_update(state, transitions, advantages, returns, update_key, cfg)
env_state = reset_done_envs(env_state, reset_key, cfg.ppo.batch_games) return state, env_state, rng, {**rollout_stats, **update_metrics}
return state, env_state, rng, {**rollout_metrics, **update_metrics}
return train_iteration return train_iteration
@@ -299,8 +332,10 @@ def make_league_train_iteration(
rng: jax.Array, rng: jax.Array,
shaping_coef: jax.Array, shaping_coef: jax.Array,
): ):
rng, rollout_key, update_key, reset_key = jax.random.split(rng, 4) rng, rollout_key, update_key = jax.random.split(rng, 3)
env_state, transitions, rollout_metrics = rollout_fn( # The rollout resets finished envs and redraws their league assignments
# in-scan, so neither needs to be refreshed here.
env_state, assignments, transitions, last_value, rollout_stats = rollout_fn(
state, env_state, assignments, rollout_key, shaping_coef state, env_state, assignments, rollout_key, shaping_coef
) )
advantages, returns = compute_gae( advantages, returns = compute_gae(
@@ -309,17 +344,10 @@ def make_league_train_iteration(
transitions.done, transitions.done,
cfg.ppo.gamma, cfg.ppo.gamma,
cfg.ppo.gae_lambda, cfg.ppo.gae_lambda,
last_value,
) )
state, update_metrics = ppo_update(state, transitions, advantages, returns, update_key, cfg) state, update_metrics = ppo_update(state, transitions, advantages, returns, update_key, cfg)
env_state, assignments = reset_done_envs_with_assignments( return state, env_state, assignments, rng, {**rollout_stats, **update_metrics}
env_state,
assignments,
reset_key,
cfg.ppo.batch_games,
opponent_probs,
mirror_probability,
)
return state, env_state, assignments, rng, {**rollout_metrics, **update_metrics}
return train_iteration return train_iteration
@@ -328,11 +356,13 @@ def make_rollout_fn(cfg: JaxPPOConfig, opponent_policy):
learner = jnp.asarray(cfg.run.learner_seat, dtype=jnp.int32) learner = jnp.asarray(cfg.run.learner_seat, dtype=jnp.int32)
opponent = jnp.asarray(1 - cfg.run.learner_seat, dtype=jnp.int32) opponent = jnp.asarray(1 - cfg.run.learner_seat, dtype=jnp.int32)
learners = jnp.full((cfg.ppo.batch_games,), cfg.run.learner_seat, dtype=jnp.int32)
@jax.jit @jax.jit
def rollout_fn(state: TrainState, env_state: State, rng: jax.Array, shaping_coef: jax.Array): def rollout_fn(state: TrainState, env_state: State, rng: jax.Array, shaping_coef: jax.Array):
def body(carry, _): def body(carry, _):
env, key = carry env, episode_return, key = carry
key, learner_key, opponent_key = jax.random.split(key, 3) key, learner_key, opponent_key, reset_key = jax.random.split(key, 4)
obs = jax.vmap(observation, in_axes=(0, None))(env, learner) obs = jax.vmap(observation, in_axes=(0, None))(env, learner)
legal = jax.vmap(legal_action_mask)(env) legal = jax.vmap(legal_action_mask)(env)
logits, value = state.apply_fn(state.params, obs) logits, value = state.apply_fn(state.params, obs)
@@ -355,7 +385,8 @@ def make_rollout_fn(cfg: JaxPPOConfig, opponent_policy):
next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions) next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
after_diff = batch_score_diff(next_env, learner) after_diff = batch_score_diff(next_env, learner)
terminal_reward = jnp.tanh(after_diff / cfg.reward.terminal_scale) terminal_reward = jnp.tanh(after_diff / cfg.reward.terminal_scale)
terminal = active & next_env.done done = next_env.done
terminal = active & done
reward = jnp.where(terminal, terminal_reward, 0.0) reward = jnp.where(terminal, terminal_reward, 0.0)
reward = reward + shaping_coef * (after_diff - before_diff) reward = reward + shaping_coef * (after_diff - before_diff)
reward = jnp.where(active, reward, 0.0) reward = jnp.where(active, reward, 0.0)
@@ -369,19 +400,27 @@ def make_rollout_fn(cfg: JaxPPOConfig, opponent_policy):
log_prob=log_prob, log_prob=log_prob,
value=value, value=value,
reward=reward, reward=reward,
done=next_env.done, done=done,
active=active, active=active,
actor_mask=actor_mask, actor_mask=actor_mask,
entropy=entropy, entropy=entropy,
play_action=(place_type == PLAY) & actor_mask, play_action=(place_type == PLAY) & actor_mask,
) )
return (next_env, key), transition
(next_env, rng), transitions = jax.lax.scan( episode_return = episode_return + reward
body, (env_state, rng), xs=None, length=cfg.ppo.rollout_steps episode = _episode_end(next_env, learners, terminal, episode_return)
next_env = reset_done_envs(next_env, reset_key, cfg.ppo.batch_games)
episode_return = jnp.where(done, 0.0, episode_return)
return (next_env, episode_return, key), (transition, episode)
init_return = jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32)
(next_env, _, rng), (transitions, episodes) = jax.lax.scan(
body, (env_state, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
) )
metrics = rollout_metrics(transitions, next_env, learner) final_obs = jax.vmap(observation, in_axes=(0, None))(next_env, learner)
return next_env, transitions, metrics _, last_value = state.apply_fn(state.params, final_obs)
metrics = episode_metrics(transitions, episodes)
return next_env, transitions, last_value, metrics
return rollout_fn return rollout_fn
@@ -405,8 +444,8 @@ def make_league_rollout_fn(
shaping_coef: jax.Array, shaping_coef: jax.Array,
): ):
def body(carry, _): def body(carry, _):
env, key = carry env, assignments, episode_return, key = carry
key, learner_key, mirror_key, pool_key = jax.random.split(key, 4) key, learner_key, mirror_key, pool_key, reset_key = jax.random.split(key, 5)
learner = assignments.learner_seat.astype(jnp.int32) learner = assignments.learner_seat.astype(jnp.int32)
opponent = 1 - learner opponent = 1 - learner
@@ -449,7 +488,8 @@ def make_league_rollout_fn(
next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions) next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
after_diff = batch_score_diff_for_players(next_env, learner) after_diff = batch_score_diff_for_players(next_env, learner)
terminal_reward = jnp.tanh(after_diff / cfg.reward.terminal_scale) terminal_reward = jnp.tanh(after_diff / cfg.reward.terminal_scale)
terminal = active & next_env.done done = next_env.done
terminal = active & done
reward = jnp.where(terminal, terminal_reward, 0.0) reward = jnp.where(terminal, terminal_reward, 0.0)
reward = reward + shaping_coef * (after_diff - before_diff) reward = reward + shaping_coef * (after_diff - before_diff)
reward = jnp.where(active, reward, 0.0) reward = jnp.where(active, reward, 0.0)
@@ -463,19 +503,37 @@ def make_league_rollout_fn(
log_prob=log_prob, log_prob=log_prob,
value=value, value=value,
reward=reward, reward=reward,
done=next_env.done, done=done,
active=active, active=active,
actor_mask=actor_mask, actor_mask=actor_mask,
entropy=entropy, entropy=entropy,
play_action=(place_type == PLAY) & actor_mask, play_action=(place_type == PLAY) & actor_mask,
) )
return (next_env, key), transition
(next_env, rng), transitions = jax.lax.scan( episode_return = episode_return + reward
body, (env_state, rng), xs=None, length=cfg.ppo.rollout_steps # Scored against the seat that just played the episode out, not the
# freshly drawn one.
episode = _episode_end(next_env, learner, terminal, episode_return)
next_env, assignments = reset_done_envs_with_assignments(
next_env,
assignments,
reset_key,
cfg.ppo.batch_games,
opponent_probs,
mirror_probability,
)
episode_return = jnp.where(done, 0.0, episode_return)
return (next_env, assignments, episode_return, key), (transition, episode)
init_return = jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32)
(next_env, assignments, _, rng), (transitions, episodes) = jax.lax.scan(
body, (env_state, assignments, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
) )
metrics = rollout_metrics_for_players(transitions, next_env, assignments.learner_seat) final_learner = assignments.learner_seat.astype(jnp.int32)
return next_env, transitions, metrics final_obs = jax.vmap(observation, in_axes=(0, 0))(next_env, final_learner)
_, last_value = state.apply_fn(state.params, final_obs)
metrics = episode_metrics(transitions, episodes)
return next_env, assignments, transitions, last_value, metrics
return rollout_fn return rollout_fn
@@ -487,11 +545,12 @@ def random_rollout(cfg: JaxPPOConfig) -> dict:
env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games)) env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
learner = jnp.asarray(cfg.run.learner_seat, dtype=jnp.int32) learner = jnp.asarray(cfg.run.learner_seat, dtype=jnp.int32)
opponent = jnp.asarray(1 - cfg.run.learner_seat, dtype=jnp.int32) opponent = jnp.asarray(1 - cfg.run.learner_seat, dtype=jnp.int32)
learners = jnp.full((cfg.ppo.batch_games,), cfg.run.learner_seat, dtype=jnp.int32)
@jax.jit @jax.jit
def rollout(env_state: State, rng: jax.Array): def rollout(env_state: State, rng: jax.Array):
def body(carry, _): def body(carry, _):
env, key = carry env, episode_return, key = carry
key, learner_key, opponent_key = jax.random.split(key, 3) key, learner_key, opponent_key = jax.random.split(key, 3)
learner_keys = jax.random.split(learner_key, cfg.ppo.batch_games) learner_keys = jax.random.split(learner_key, cfg.ppo.batch_games)
opponent_keys = jax.random.split(opponent_key, cfg.ppo.batch_games) opponent_keys = jax.random.split(opponent_key, cfg.ppo.batch_games)
@@ -508,6 +567,7 @@ def random_rollout(cfg: JaxPPOConfig) -> dict:
next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions) next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
after_diff = batch_score_diff(next_env, learner) after_diff = batch_score_diff(next_env, learner)
reward = jnp.where(active, after_diff - before_diff, 0.0) reward = jnp.where(active, after_diff - before_diff, 0.0)
terminal = active & next_env.done
actor_mask = active & learner_turn actor_mask = active & learner_turn
place_type = (actions % 12) // 6 place_type = (actions % 12) // 6
dummy_obs = jnp.zeros((cfg.ppo.batch_games, OBS_DIM), dtype=jnp.float32) dummy_obs = jnp.zeros((cfg.ppo.batch_games, OBS_DIM), dtype=jnp.float32)
@@ -525,12 +585,15 @@ def random_rollout(cfg: JaxPPOConfig) -> dict:
entropy=jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32), entropy=jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32),
play_action=(place_type == PLAY) & actor_mask, play_action=(place_type == PLAY) & actor_mask,
) )
return (next_env, key), transition episode_return = episode_return + reward
episode = _episode_end(next_env, learners, terminal, episode_return)
return (next_env, episode_return, key), (transition, episode)
(next_env, rng), transitions = jax.lax.scan( init_return = jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32)
body, (env_state, rng), xs=None, length=cfg.ppo.rollout_steps (next_env, _, rng), (transitions, episodes) = jax.lax.scan(
body, (env_state, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
) )
return rng, rollout_metrics(transitions, next_env, learner) return rng, episode_metrics(transitions, episodes)
rng, metrics = rollout(env_state, rng) rng, metrics = rollout(env_state, rng)
jax.tree_util.tree_leaves(metrics)[0].block_until_ready() jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
@@ -679,7 +742,18 @@ def compute_gae(
dones: jax.Array, dones: jax.Array,
gamma: float, gamma: float,
gae_lambda: float, gae_lambda: float,
last_value: jax.Array | None = None,
) -> tuple[jax.Array, jax.Array]: ) -> tuple[jax.Array, jax.Array]:
"""GAE over a rollout that may end mid-episode.
``last_value`` is V(s_T) for the state the scan stopped on. Episodes that
are cut by the rollout boundary bootstrap from it; leaving it at zero would
tell the critic every truncated episode is worth nothing.
"""
if last_value is None:
last_value = jnp.zeros_like(values[-1])
def body(carry, x): def body(carry, x):
next_value, next_advantage = carry next_value, next_advantage = carry
reward, value, done = x reward, value, done = x
@@ -688,7 +762,7 @@ def compute_gae(
advantage = delta + gamma * gae_lambda * nonterminal * next_advantage advantage = delta + gamma * gae_lambda * nonterminal * next_advantage
return (value, advantage), advantage return (value, advantage), advantage
init = (jnp.zeros_like(values[-1]), jnp.zeros_like(values[-1])) init = (last_value, jnp.zeros_like(values[-1]))
_, advantages_rev = jax.lax.scan(body, init, (rewards[::-1], values[::-1], dones[::-1])) _, advantages_rev = jax.lax.scan(body, init, (rewards[::-1], values[::-1], dones[::-1]))
advantages = advantages_rev[::-1] advantages = advantages_rev[::-1]
return advantages, advantages + values return advantages, advantages + values
@@ -1066,56 +1140,60 @@ def make_policy_match_batch_fn(learner_policy, opponent_policy):
return eval_batch return eval_batch
def rollout_metrics( def _episode_end(
transitions: Transition, final_env: State, learner: jax.Array next_env: State,
) -> dict[str, jax.Array]: learners: jax.Array,
episode_return = jnp.sum(transitions.reward, axis=0) done: jax.Array,
episode_return: jax.Array,
) -> EpisodeEnd:
"""Snapshot the finished-episode stats of ``next_env`` before it is reset."""
learners = jnp.broadcast_to(learners.astype(jnp.int32), done.shape)
color_scores = jax.vmap(color_scores_for_player, in_axes=(0, 0))(next_env, learners)
batch_idx = jnp.arange(done.shape[0])
return EpisodeEnd(
done=done,
episode_return=episode_return,
length=next_env.step_count,
opened_colors=jnp.sum(next_env.col_len[batch_idx, learners, :] > 0, axis=-1),
positive_expeditions=jnp.sum(color_scores > 0, axis=-1),
hit_max_steps=next_env.step_count >= MAX_STEPS,
)
def episode_metrics(transitions: Transition, episodes: EpisodeEnd) -> dict[str, jax.Array]:
"""Aggregate over episodes that actually finished inside the rollout.
Averaging over ``final_env`` instead would sample envs mid-episode, and
summing rewards along the scan axis would add up several episodes per slot.
"""
done = episodes.done.astype(jnp.float32)
completed = jnp.sum(done)
denom = jnp.maximum(completed, 1.0)
actor_count = jnp.sum(transitions.actor_mask) actor_count = jnp.sum(transitions.actor_mask)
active_count = jnp.sum(transitions.active)
color_scores = jax.vmap(color_scores_for_player, in_axes=(0, None))(final_env, learner) returns = episodes.episode_return
opened = jnp.sum(final_env.col_len[:, learner, :] > 0, axis=-1) return_mean = jnp.sum(returns * done) / denom
positive = jnp.sum(color_scores > 0, axis=-1) return_var = jnp.sum(((returns - return_mean) ** 2) * done) / denom
return { return {
"return_mean": jnp.mean(episode_return), "return_mean": return_mean,
"return_std": jnp.std(episode_return), "return_std": jnp.sqrt(jnp.maximum(return_var, 0.0)),
"game_length_mean": jnp.mean(final_env.step_count.astype(jnp.float32)), "game_length_mean": jnp.sum(episodes.length.astype(jnp.float32) * done) / denom,
"game_length_max": jnp.max(final_env.step_count), "game_length_max": jnp.max(jnp.where(episodes.done, episodes.length, 0)),
"max_steps_rate": jnp.mean(final_env.step_count >= MAX_STEPS), "max_steps_rate": jnp.sum(episodes.hit_max_steps.astype(jnp.float32) * done) / denom,
"play_action_rate": jnp.sum(transitions.play_action) / jnp.maximum(actor_count, 1), "play_action_rate": jnp.sum(transitions.play_action) / jnp.maximum(actor_count, 1),
"opened_colors_mean": jnp.mean(opened.astype(jnp.float32)), "opened_colors_mean": jnp.sum(episodes.opened_colors.astype(jnp.float32) * done) / denom,
"positive_expeditions_mean": jnp.mean(positive.astype(jnp.float32)), "positive_expeditions_mean": (
jnp.sum(episodes.positive_expeditions.astype(jnp.float32) * done) / denom
),
"entropy_mean": masked_mean( "entropy_mean": masked_mean(
transitions.entropy, transitions.actor_mask.astype(jnp.float32) transitions.entropy, transitions.actor_mask.astype(jnp.float32)
), ),
"active_steps": active_count, "active_steps": jnp.sum(transitions.active),
"learner_actions": actor_count,
}
def rollout_metrics_for_players(
transitions: Transition, final_env: State, learners: jax.Array
) -> dict[str, jax.Array]:
episode_return = jnp.sum(transitions.reward, axis=0)
actor_count = jnp.sum(transitions.actor_mask)
active_count = jnp.sum(transitions.active)
color_scores = jax.vmap(color_scores_for_player, in_axes=(0, 0))(final_env, learners)
batch_idx = jnp.arange(learners.shape[0])
opened = jnp.sum(final_env.col_len[batch_idx, learners.astype(jnp.int32), :] > 0, axis=-1)
positive = jnp.sum(color_scores > 0, axis=-1)
return {
"return_mean": jnp.mean(episode_return),
"return_std": jnp.std(episode_return),
"game_length_mean": jnp.mean(final_env.step_count.astype(jnp.float32)),
"game_length_max": jnp.max(final_env.step_count),
"max_steps_rate": jnp.mean(final_env.step_count >= MAX_STEPS),
"play_action_rate": jnp.sum(transitions.play_action) / jnp.maximum(actor_count, 1),
"opened_colors_mean": jnp.mean(opened.astype(jnp.float32)),
"positive_expeditions_mean": jnp.mean(positive.astype(jnp.float32)),
"entropy_mean": masked_mean(
transitions.entropy, transitions.actor_mask.astype(jnp.float32)
),
"active_steps": active_count,
"learner_actions": actor_count, "learner_actions": actor_count,
"episodes_completed": completed,
} }
@@ -1179,11 +1257,11 @@ def _flatten_transitions(transitions: Transition) -> Transition:
return jax.tree_util.tree_map(lambda x: x.reshape((-1, *x.shape[2:])), transitions) return jax.tree_util.tree_map(lambda x: x.reshape((-1, *x.shape[2:])), transitions)
def shaping_coefficient(cfg: JaxPPOConfig, env_steps: int) -> float: def shaping_coefficient(cfg: JaxPPOConfig, learner_actions: int) -> float:
reward_cfg = cfg.reward reward_cfg = cfg.reward
if reward_cfg.potential_shaping_anneal_steps <= 0: if reward_cfg.potential_shaping_anneal_steps <= 0:
return reward_cfg.potential_shaping_final return reward_cfg.potential_shaping_final
progress = min(env_steps / reward_cfg.potential_shaping_anneal_steps, 1.0) progress = min(learner_actions / reward_cfg.potential_shaping_anneal_steps, 1.0)
return reward_cfg.potential_shaping_initial + progress * ( return reward_cfg.potential_shaping_initial + progress * (
reward_cfg.potential_shaping_final - reward_cfg.potential_shaping_initial reward_cfg.potential_shaping_final - reward_cfg.potential_shaping_initial
) )
+179
View File
@@ -0,0 +1,179 @@
"""Phase 1: the classic three-round match layer."""
import jax
import jax.numpy as jnp
import numpy as np
import pytest
from lost_cities_jax.engine import board_score
from lost_cities_jax.match import (
N_ROUNDS,
MatchState,
match_legal_action_mask,
match_reset,
match_reset_from,
match_score,
match_step,
starting_player,
)
from lost_cities_jax.opponents import random_legal_action
from lost_cities_jax.types import DECK_DRAWS, DRAW_DECK, N_CARDS
MATCH_STEP_CAP = 1400 # 3 rounds x the engine's 400-ply safety cap, with slack
def _decks(seed: int) -> jnp.ndarray:
rng = np.random.default_rng(seed)
return jnp.asarray(
np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)]), dtype=jnp.int8
)
def _play_match(state: MatchState, key: jax.Array):
"""Drive a match to completion with random legal play, recording each ply."""
plies = []
while not bool(state.done):
key, step_key = jax.random.split(key)
action = random_legal_action(state.round, state.round.to_move, step_key)
assert bool(match_legal_action_mask(state)[action])
prev = state
state, reward, _ = match_step(state, action)
plies.append((prev, int(action), state, np.asarray(reward)))
assert len(plies) < MATCH_STEP_CAP, "match failed to terminate"
return state, plies
# --- match structure -------------------------------------------------------
def test_match_plays_exactly_three_rounds_and_then_ends():
state = match_reset(jax.random.PRNGKey(0))
final, plies = _play_match(state, jax.random.PRNGKey(1))
assert bool(final.done)
assert int(final.round_idx) == N_ROUNDS - 1
# round_idx advances 0 -> 1 -> 2 and stops.
seen = {int(before.round_idx) for before, _, _, _ in plies}
assert seen == {0, 1, 2}
def test_reward_is_paid_only_when_the_third_round_ends():
state = match_reset(jax.random.PRNGKey(2))
final, plies = _play_match(state, jax.random.PRNGKey(3))
rewards = np.stack([reward for _, _, _, reward in plies])
paid = np.flatnonzero(np.any(rewards != 0.0, axis=1))
# Intermediate rounds bank points into carry but pay nothing.
assert paid.tolist() == [len(plies) - 1]
assert np.allclose(rewards[-1], np.asarray(match_score(final), dtype=np.float32))
def test_each_round_runs_44_deck_draws_plus_one_ply_per_discard_draw():
"""The deck clock: only deck draws end a round, so discard draws extend it."""
state = match_reset(jax.random.PRNGKey(4))
_, plies = _play_match(state, jax.random.PRNGKey(5))
per_round: dict[int, list[int]] = {0: [], 1: [], 2: []}
for before, action, _, _ in plies:
per_round[int(before.round_idx)].append(action % 6)
for round_idx, draws in per_round.items():
deck_draws = sum(1 for src in draws if src == DRAW_DECK)
discard_draws = len(draws) - deck_draws
assert deck_draws == DECK_DRAWS, f"round {round_idx} drew {deck_draws} deck cards"
assert len(draws) == DECK_DRAWS + discard_draws
# --- carry -----------------------------------------------------------------
def test_carry_banks_each_finished_round_exactly_once():
state = match_reset(jax.random.PRNGKey(6))
final, plies = _play_match(state, jax.random.PRNGKey(7))
deltas = [
np.asarray(after.carry) - np.asarray(before.carry)
for before, _, after, _ in plies
if int(after.round_idx) != int(before.round_idx)
]
# Two roll-overs for three rounds, and carry is exactly their sum.
assert len(deltas) == N_ROUNDS - 1
assert np.array_equal(np.asarray(final.carry), sum(deltas))
# The third round is still on the board, not yet banked.
board = np.asarray(board_score(final.round)).astype(np.int32)
assert np.array_equal(np.asarray(match_score(final)), np.asarray(final.carry) + board)
def test_match_score_does_not_jump_across_a_round_boundary():
"""Phi = carry + board diff must be continuous, or PBRS gets a free kick."""
state = match_reset(jax.random.PRNGKey(8))
_, plies = _play_match(state, jax.random.PRNGKey(9))
for before, _, after, _ in plies:
rolled_over = int(after.round_idx) != int(before.round_idx)
if not rolled_over:
continue
# The finished round's board is folded into carry and the new board is
# empty, so the running total is unchanged by the roll-over itself.
assert np.array_equal(np.asarray(after.carry), np.asarray(match_score(after)))
assert int(np.asarray(board_score(after.round)).sum()) == 0
# --- the start-player rule -------------------------------------------------
@pytest.mark.parametrize(
("carry", "coin", "expected"),
[
((60, 10), 1, 0), # p0 ahead -> p0 leads, coin ignored
((10, 60), 0, 1), # p1 ahead -> p1 leads, coin ignored
((30, 30), 0, 0), # level -> coin
((30, 30), 1, 1), # level -> coin
],
)
def test_the_player_with_more_points_begins(carry, coin, expected):
coin_flips = jnp.asarray([coin, coin, coin], dtype=jnp.int8)
starter = starting_player(jnp.asarray(carry, dtype=jnp.int32), jnp.int32(1), coin_flips)
assert int(starter) == expected
def test_round_two_is_actually_led_by_whoever_is_ahead():
state = match_reset(jax.random.PRNGKey(10))
_, plies = _play_match(state, jax.random.PRNGKey(11))
for before, _, after, _ in plies:
if int(after.round_idx) == int(before.round_idx):
continue
carry = np.asarray(after.carry)
if carry[0] == carry[1]:
continue # coin flip, covered above
leader = 0 if carry[0] > carry[1] else 1
assert int(after.round.to_move) == leader
def test_round_one_start_player_is_a_fair_coin():
starts = [int(match_reset(jax.random.PRNGKey(seed)).round.to_move) for seed in range(400)]
share = sum(starts) / len(starts)
assert 0.4 < share < 0.6, f"round-1 starter is skewed: {share:.2f}"
# --- mirrored matches ------------------------------------------------------
def test_mirroring_the_seats_negates_the_match_score():
"""Same deals, same coins, seats swapped: the result must flip sign exactly."""
decks = _decks(21)
coins = jnp.asarray([0, 1, 0], dtype=jnp.int8)
# Swapping seats == swapping the two dealt hands and the coin bits.
swapped = decks.at[:, :16].set(jnp.concatenate([decks[:, 8:16], decks[:, :8]], axis=1))
flipped_coins = (1 - coins).astype(jnp.int8)
a, _ = _play_match(match_reset_from(decks, coins), jax.random.PRNGKey(30))
b, _ = _play_match(match_reset_from(swapped, flipped_coins), jax.random.PRNGKey(30))
# Both matches see identical information; only the seat labels differ, so the
# per-seat totals must be each other's mirror image.
assert np.array_equal(np.asarray(match_score(a)), np.asarray(match_score(b))[::-1])
+200
View File
@@ -0,0 +1,200 @@
"""Phases 2-4: match reward, match observation, asymmetric critic, both-seat training."""
import jax
import jax.numpy as jnp
import numpy as np
import pytest
from lost_cities_jax.match import match_reset, match_score, match_step
from lost_cities_jax.match_obs import (
MATCH_CRITIC_OBS_DIM,
MATCH_OBS_DIM,
match_critic_observation,
match_observation,
)
from lost_cities_jax.match_ppo import (
Ablation,
create_match_train_state,
make_match_rollout_fn,
make_match_train_iteration,
)
from lost_cities_jax.opponents import random_legal_action
from lost_cities_jax.ppo import JaxPPOConfig
from lost_cities_jax.types import LOC_P0_HAND, N_CARDS
def _cfg(rollout_steps: int = 700, batch_games: int = 8) -> JaxPPOConfig:
cfg = JaxPPOConfig()
cfg.ppo.batch_games = batch_games
cfg.ppo.rollout_steps = rollout_steps
cfg.ppo.minibatches = 4
cfg.ppo.epochs = 1
cfg.network.hidden_size = 32
cfg.network.num_layers = 1
return cfg
# --- the asymmetric critic -------------------------------------------------
def test_the_critic_sees_the_opponents_hand_and_the_actor_does_not():
match = match_reset(jax.random.PRNGKey(0))
loc = np.asarray(match.round.card_loc)
opp_hand = np.flatnonzero(loc == LOC_P0_HAND + 1)
actor = np.asarray(match_observation(match, jnp.int32(0)))
critic = np.asarray(match_critic_observation(match, jnp.int32(0)))
# Card-major, three channels each, matching the base observation's layout.
hidden = critic[MATCH_OBS_DIM:].reshape(N_CARDS, 3)
assert np.array_equal(np.flatnonzero(hidden[:, 0]), opp_hand)
# The actor's view is a strict prefix of it and carries none of that.
assert np.array_equal(actor, critic[:MATCH_OBS_DIM])
assert critic.shape == (MATCH_CRITIC_OBS_DIM,)
def test_privileged_input_cannot_move_the_policy_logits():
"""Separate trunks, or the critic's view of the deck leaks into play.
The privileged critic is off by default -- it measured negative, see the
ablation in the plan -- but the isolation property still has to hold for
anyone who switches it on, and for its later reuse as a search evaluator.
"""
cfg = _cfg()
state = create_match_train_state(cfg, jax.random.PRNGKey(1), Ablation(privileged_critic=True))
match = match_reset(jax.random.PRNGKey(2))
obs = match_observation(match, jnp.int32(0))[None, :]
critic_a = match_critic_observation(match, jnp.int32(0))[None, :]
critic_b = critic_a.at[0, MATCH_OBS_DIM:].set(1.0) # a totally different hidden state
logits_a, value_a = state.apply_fn(state.params, obs, critic_a)
logits_b, value_b = state.apply_fn(state.params, obs, critic_b)
assert np.array_equal(np.asarray(logits_a), np.asarray(logits_b))
# ...and the value head must actually be using it, or the critic is pointless.
assert not np.allclose(np.asarray(value_a), np.asarray(value_b))
# --- the match reward ------------------------------------------------------
def test_reward_is_zero_sum():
cfg = _cfg()
state = create_match_train_state(cfg, jax.random.PRNGKey(3))
env = jax.jit(jax.vmap(match_reset))(
jax.random.split(jax.random.PRNGKey(4), cfg.ppo.batch_games)
)
rollout = make_match_rollout_fn(cfg)
_, transitions, _, metrics = rollout(state, env, jax.random.PRNGKey(5), jnp.asarray(1.0))
reward = np.asarray(transitions.reward) # (T, 2 * batch)
seat0, seat1 = reward[:, : cfg.ppo.batch_games], reward[:, cfg.ppo.batch_games :]
assert np.allclose(seat0, -seat1)
assert float(metrics["matches_completed"]) > 0
def test_the_dense_reward_telescopes_to_the_match_total():
"""The whole justification for paying every ply: at gamma=1 the sum is the total.
That is what makes "score more across three rounds" the objective, rather than
some shaped proxy for it.
"""
match = match_reset(jax.random.PRNGKey(6))
key = jax.random.PRNGKey(7)
scale = 50.0
paid = 0.0
while not bool(match.done):
key, step_key = jax.random.split(key)
action = random_legal_action(match.round, match.round.to_move, step_key)
before = np.asarray(match_score(match))
match, _, _ = match_step(match, action)
after = np.asarray(match_score(match))
paid += ((after[0] - after[1]) - (before[0] - before[1])) / scale
final = np.asarray(match_score(match))
total_lead = (final[0] - final[1]) / scale
# Every point banked along the way, and nothing else.
assert paid == pytest.approx(total_lead, abs=1e-4)
# --- both seats ------------------------------------------------------------
def test_exactly_one_seat_acts_per_ply_but_both_are_trained():
cfg = _cfg()
state = create_match_train_state(cfg, jax.random.PRNGKey(8))
env = jax.jit(jax.vmap(match_reset))(
jax.random.split(jax.random.PRNGKey(9), cfg.ppo.batch_games)
)
rollout = make_match_rollout_fn(cfg)
_, transitions, last_value, _ = rollout(state, env, jax.random.PRNGKey(10), jnp.asarray(0.0))
batch = cfg.ppo.batch_games
actor = np.asarray(transitions.actor_mask)
seat0, seat1 = actor[:, :batch], actor[:, batch:]
# One mover per ply while the match is live.
live = np.asarray(transitions.active)[:, :batch]
assert np.array_equal(seat0 ^ seat1, live)
# But the critic gets both seats: every live ply is a value target on both.
assert np.asarray(transitions.active).sum() == 2 * live.sum()
assert transitions.value.shape == (cfg.ppo.rollout_steps, 2 * batch)
assert last_value.shape == (2 * batch,)
def test_self_play_is_balanced_and_a_train_step_stays_finite():
cfg = _cfg(rollout_steps=700, batch_games=16)
state = create_match_train_state(cfg, jax.random.PRNGKey(11))
env = jax.jit(jax.vmap(match_reset))(jax.random.split(jax.random.PRNGKey(12), 16))
train_iteration = make_match_train_iteration(cfg)
state, env, _, metrics = train_iteration(state, env, jax.random.PRNGKey(13), jnp.asarray(0.1))
assert bool(jnp.isfinite(metrics["loss"]))
assert float(metrics["matches_completed"]) > 0
# One network on both sides: neither seat should be favoured.
assert 0.2 < float(metrics["match_win_rate"]) < 0.8
# --- the match observation -------------------------------------------------
def test_carry_reaches_the_observation_with_a_usable_scale():
"""The old score_diff divided by 780; a 50-point lead vanished into 0.06."""
match = match_reset(jax.random.PRNGKey(14))
behind = match._replace(carry=jnp.asarray([0, 50], dtype=jnp.int32))
ahead = match._replace(carry=jnp.asarray([50, 0], dtype=jnp.int32))
obs_behind = np.asarray(match_observation(behind, jnp.int32(0)))
obs_ahead = np.asarray(match_observation(ahead, jnp.int32(0)))
delta = np.abs(obs_ahead - obs_behind)
# A 100-point swing has to be plainly visible, not a rounding error.
assert delta.max() > 0.5
assert (delta > 0.01).sum() >= 2 # the scalar and at least one bin flip
@pytest.mark.parametrize("round_idx", [0, 1, 2])
def test_the_round_index_is_observable(round_idx):
match = match_reset(jax.random.PRNGKey(15))._replace(
round_idx=jnp.asarray(round_idx, dtype=jnp.int32)
)
obs = np.asarray(match_observation(match, jnp.int32(0)))
assert obs.shape == (MATCH_OBS_DIM,)
assert np.isfinite(obs).all()
def test_whose_turn_it_is_is_observable():
"""The single-round obs never said; the critic had to read it off step_count."""
match = match_reset(jax.random.PRNGKey(16))
mover = int(match.round.to_move)
from_mover = np.asarray(match_observation(match, jnp.int32(mover)))
from_waiter = np.asarray(match_observation(match, jnp.int32(1 - mover)))
assert not np.array_equal(from_mover, from_waiter)
+8 -3
View File
@@ -61,7 +61,9 @@ def tiny_config(tmp_path) -> JaxPPOConfig:
), ),
opponent=OpponentConfig(name="discard_only"), opponent=OpponentConfig(name="discard_only"),
network=NetworkConfig(hidden_size=32, num_layers=1), network=NetworkConfig(hidden_size=32, num_layers=1),
ppo=PPOHyperConfig(batch_games=8, rollout_steps=16, epochs=1, minibatches=2), # A round runs at least 44 plies (one per deck draw), so a shorter
# rollout would finish no episodes and leave the episode metrics empty.
ppo=PPOHyperConfig(batch_games=8, rollout_steps=80, epochs=1, minibatches=2),
) )
@@ -305,9 +307,12 @@ def test_gate3_checkpoint_duplicate_self_mirror_score_diff_is_zero():
def test_random_rollout_smoke(tmp_path): def test_random_rollout_smoke(tmp_path):
row = random_rollout(tiny_config(tmp_path)) row = random_rollout(tiny_config(tmp_path))
assert row["env_steps"] == 8 * 16 assert row["env_steps"] == 8 * 80
assert 0.0 <= row["play_action_rate"] <= 1.0 assert 0.0 <= row["play_action_rate"] <= 1.0
assert row["game_length_mean"] > 0.0 # Averaged over episodes that actually finished, so it must clear the
# 44-ply floor rather than report a mid-episode step count.
assert row["episodes_completed"] > 0
assert row["game_length_mean"] >= 44.0
def test_train_checkpoint_and_eval_smoke(tmp_path): def test_train_checkpoint_and_eval_smoke(tmp_path):
+188
View File
@@ -0,0 +1,188 @@
"""Phase 0a: GAE bootstrap, in-scan auto-reset, and episode-boundary metrics."""
import jax
import jax.numpy as jnp
import pytest
from lost_cities_jax.engine import reset
from lost_cities_jax.opponents import policy_by_name
from lost_cities_jax.ppo import (
EpisodeEnd,
JaxPPOConfig,
Transition,
compute_gae,
create_train_state,
episode_metrics,
make_rollout_fn,
shaping_coefficient,
)
def _cfg(**overrides) -> JaxPPOConfig:
cfg = JaxPPOConfig()
cfg.ppo.batch_games = 8
cfg.ppo.rollout_steps = 120
cfg.network.hidden_size = 16
cfg.network.num_layers = 1
for key, value in overrides.items():
setattr(cfg.run, key, value)
return cfg
# --- GAE bootstrap ---------------------------------------------------------
def test_gae_bootstraps_truncated_episode_from_last_value():
"""A rollout that ends mid-episode must carry V(s_T), not zero."""
rewards = jnp.zeros((3, 1))
values = jnp.zeros((3, 1))
dones = jnp.zeros((3, 1), dtype=bool) # never terminates inside the window
_, returns_zero = compute_gae(rewards, values, dones, gamma=1.0, gae_lambda=1.0)
_, returns_boot = compute_gae(
rewards, values, dones, gamma=1.0, gae_lambda=1.0, last_value=jnp.array([5.0])
)
# Without a bootstrap the truncated episode looks worthless.
assert jnp.allclose(returns_zero, 0.0)
# With it, every step inherits the tail value.
assert jnp.allclose(returns_boot, 5.0)
def test_gae_ignores_last_value_when_episode_terminates():
"""A terminal step must not bootstrap past the episode boundary."""
rewards = jnp.array([[0.0], [1.0]])
values = jnp.zeros((2, 1))
dones = jnp.array([[False], [True]])
_, returns = compute_gae(
rewards, values, dones, gamma=1.0, gae_lambda=1.0, last_value=jnp.array([99.0])
)
# Terminal step sees only its own reward; the step before it sees that too.
assert jnp.allclose(returns, jnp.array([[1.0], [1.0]]))
def test_gae_done_flag_cuts_credit_between_episodes():
"""With in-scan resets two episodes share a slot; credit must not leak."""
rewards = jnp.array([[1.0], [7.0]])
values = jnp.zeros((2, 1))
dones = jnp.array([[True], [False]]) # first step ends episode A
_, returns = compute_gae(rewards, values, dones, gamma=1.0, gae_lambda=1.0)
# Episode A keeps its own reward; episode B's +7 must not flow backwards.
assert jnp.allclose(returns[0], 1.0)
assert jnp.allclose(returns[1], 7.0)
# --- episode-boundary metrics ---------------------------------------------
def _episodes(done, episode_return, length) -> EpisodeEnd:
shape = jnp.asarray(done).shape
return EpisodeEnd(
done=jnp.asarray(done),
episode_return=jnp.asarray(episode_return, dtype=jnp.float32),
length=jnp.asarray(length, dtype=jnp.int32),
opened_colors=jnp.zeros(shape, dtype=jnp.int32),
positive_expeditions=jnp.zeros(shape, dtype=jnp.int32),
hit_max_steps=jnp.zeros(shape, dtype=bool),
)
def _transitions(steps: int, games: int) -> Transition:
zeros = jnp.zeros((steps, games), dtype=jnp.float32)
false = jnp.zeros((steps, games), dtype=bool)
return Transition(
obs=jnp.zeros((steps, games, 1)),
legal_mask=jnp.zeros((steps, games, 1), dtype=bool),
action=jnp.zeros((steps, games), dtype=jnp.int32),
log_prob=zeros,
value=zeros,
reward=zeros,
done=false,
active=~false,
actor_mask=~false,
entropy=zeros,
play_action=false,
)
def test_episode_metrics_only_counts_finished_episodes():
"""Unfinished episodes carry a partial return; they must not be averaged in."""
# Slot 0 finishes twice (returns 10 and 20); slot 1 never finishes.
done = jnp.array([[True, False], [False, False], [True, False]])
episode_return = jnp.array([[10.0, 3.0], [0.0, 6.0], [20.0, 9.0]])
length = jnp.array([[50, 1], [0, 2], [60, 3]])
metrics = episode_metrics(_transitions(3, 2), _episodes(done, episode_return, length))
assert float(metrics["episodes_completed"]) == 2.0
assert float(metrics["return_mean"]) == pytest.approx(15.0)
assert float(metrics["game_length_mean"]) == pytest.approx(55.0)
assert int(metrics["game_length_max"]) == 60
def test_episode_metrics_survives_a_rollout_with_no_completions():
done = jnp.zeros((2, 2), dtype=bool)
metrics = episode_metrics(
_transitions(2, 2), _episodes(done, jnp.zeros((2, 2)), jnp.zeros((2, 2)))
)
assert float(metrics["episodes_completed"]) == 0.0
assert float(metrics["return_mean"]) == 0.0 # guarded denominator, not a NaN
# --- in-scan auto-reset ----------------------------------------------------
def test_rollout_never_returns_a_done_env_and_completes_many_episodes():
"""The whole point of in-scan reset: no dead steps, several games per slot."""
cfg = _cfg()
rng = jax.random.PRNGKey(0)
rng, init_key, reset_key, roll_key = jax.random.split(rng, 4)
state = create_train_state(cfg, init_key)
env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
rollout_fn = make_rollout_fn(cfg, policy_by_name("discard_only"))
next_env, transitions, last_value, metrics = rollout_fn(
state, env_state, roll_key, jnp.asarray(0.0, dtype=jnp.float32)
)
# Every slot is mid-episode, never parked on a finished game.
assert not bool(jnp.any(next_env.done))
# 120 steps at ~50 plies a game means each of the 8 slots finished at least one.
assert float(metrics["episodes_completed"]) >= cfg.ppo.batch_games
# No step is wasted stepping an already-done env.
assert bool(jnp.all(transitions.active))
assert int(metrics["active_steps"]) == cfg.ppo.rollout_steps * cfg.ppo.batch_games
assert last_value.shape == (cfg.ppo.batch_games,)
def test_rollout_game_length_matches_the_rules_clock():
"""A round is 44 deck draws plus one turn per discard-pile draw."""
cfg = _cfg()
rng = jax.random.PRNGKey(1)
rng, init_key, reset_key, roll_key = jax.random.split(rng, 4)
state = create_train_state(cfg, init_key)
env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
rollout_fn = make_rollout_fn(cfg, policy_by_name("discard_only"))
_, _, _, metrics = rollout_fn(state, env_state, roll_key, jnp.asarray(0.0, dtype=jnp.float32))
# 44 deck draws is the floor; discard-pile draws only ever extend a round.
assert float(metrics["game_length_mean"]) >= 44.0
assert float(metrics["max_steps_rate"]) == 0.0
# --- shaping anneal accounting --------------------------------------------
def test_shaping_anneal_tracks_learner_actions_not_padded_steps():
cfg = _cfg()
cfg.reward.potential_shaping_initial = 1.0
cfg.reward.potential_shaping_final = 0.0
cfg.reward.potential_shaping_anneal_steps = 1_000
assert shaping_coefficient(cfg, 0) == pytest.approx(1.0)
assert shaping_coefficient(cfg, 500) == pytest.approx(0.5)
assert shaping_coefficient(cfg, 1_000) == pytest.approx(0.0)
assert shaping_coefficient(cfg, 10_000) == pytest.approx(0.0)
@@ -0,0 +1,34 @@
import json
import pytest
from scripts.serve_web_with_logs import parse_record
def test_parse_record_accepts_analysis_record():
record = {"format": "lost-cities-web-game-v1", "gameId": "game-1", "moves": []}
assert parse_record(json.dumps(record).encode()) == record
@pytest.mark.parametrize("record", [{}, {"format": "other"}, {"format": "lost-cities-web-game-v1"}])
def test_parse_record_rejects_invalid_input(record):
with pytest.raises(ValueError):
parse_record(json.dumps(record).encode())
def test_parse_record_accepts_the_match_schema():
# v2 carries which model played and the three-round match; v1 had neither.
record = {
"format": "lost-cities-web-game-v2",
"gameId": "game-2",
"opponent": {"codename": "borealis", "hash": "13a25243de1c"},
"mode": 3,
"moves": [],
}
assert parse_record(json.dumps(record).encode()) == record
def test_parse_record_still_accepts_v1():
# 111 games were recorded under v1; refusing them now would orphan them.
record = {"format": "lost-cities-web-game-v1", "gameId": "game-1", "moves": []}
assert parse_record(json.dumps(record).encode()) == record
+68
View File
@@ -0,0 +1,68 @@
---
name: verify
description: Drive the Lost Cities web client in a real browser to observe a change working — model loading, a full match, the result card, records being written.
---
# Verifying the web client
The client is where a silent bug hides best. `npm test` and `tsc` were both green
while the app was announcing the wrong winner and dropping every game record on
the floor. Neither throws. Run it.
## Build and serve
```bash
cd web && npm run build
cd .. && tmux new-session -d -s websrv \
"uv run python scripts/serve_web_with_logs.py --host 127.0.0.1 --port 5199 \
--dist web/dist --output /tmp/verify-records.jsonl"
```
Serve the **built dist**, not the vite dev server — the deploy builds from dist,
and the model is a static asset whose path only resolves there. Point `--output`
at a scratch file so verification runs never touch `data/human-play/`.
## Drive it
Playwright is deliberately **not** a dependency: the `playwright` package downloads
~114MB of Chromium on install, and `npm ci` runs in the deploy workflow. Install it
for the run and uninstall after.
```bash
cd web
npm i -D playwright --no-fund --no-audit && npx playwright install chromium
# ... drive ...
npm uninstall playwright
```
## Selectors that actually work
Found by dumping the DOM; guessing at them wasted two runs.
| What | Selector |
|---|---|
| a hand card | `button.card:not([disabled])` |
| play onto an expedition | `button.lane__zone--mine.is-target` |
| discard | `button.lane__discard.is-target` |
| draw from deck | `button.deck-stack:not([disabled])` |
| which model loaded | `.score-plaque--rival small` |
| match progress | `.round-strip` |
| final scores | `.result-card` |
A turn is three clicks: pick a card, choose where it goes, then draw. The place and
draw targets only appear **after** the card is selected, and only the legal ones are
enabled — so click the card first, then query.
The rival answers on a 620ms timer plus inference; ~200ms of slack between plies is
enough. A full three-round match runs ~140 plies, so budget a few minutes.
## Worth driving
- **A whole match, not one round.** The round roll-over is where carry banks, and it
is where the result card got it wrong.
- **Check the record actually saved.** A schema bump on the client silently 400s
against `scripts/serve_web_with_logs.py` until its allowlist is updated too.
- **A stale save in localStorage.** Bump the key on a schema change; a half-migrated
save is worse than a fresh deal.
- **The same seed twice.** In match mode the seed fixes all three deals and the coin
flips, so a match is a pure function of it.
+21
View File
@@ -0,0 +1,21 @@
# Web bug checklist
- [ ] Prevent cards from launching more than once during initial loading.
- Implementation and automated checks complete; user visual acceptance pending.
- [ ] Stack Hint, Undo, and Redo vertically at the bottom-right, aligned with the hand row.
- Implementation and automated checks complete; user visual acceptance pending.
- [ ] Restore the latest game timeline and review position from local storage.
- Implementation and automated checks complete; user reload acceptance pending.
- [ ] Keep Hint, Undo, Redo, and Play From Here in fixed vertical positions.
- Play From Here remains visible but disabled when unavailable; user visual acceptance pending.
- [ ] Automatically record completed human-versus-AI games for later analysis.
- JSONL includes the full deal, both hands before every move, actions, draws, scores, and final state.
- Local production output: `data/human-play/game-records.jsonl` (was `runs/tmp/`, which is gitignored and documented as disposable).
- [ ] Support game-record IDs when LAN HTTP does not expose `crypto.randomUUID`.
- Uses `crypto.getRandomValues` with a compatibility fallback; user acceptance pending.
Verification:
- `cd web && npm test`
- `cd web && npm run build`
- Visual acceptance is performed by the user; do not run browser automation.
+15
View File
@@ -0,0 +1,15 @@
{
"format": "coolrl-lost-cities-match-onnx-v1",
"codename": "borealis",
"model_file": "borealis.onnx",
"model_size_bytes": 3327043,
"model_sha256": "13a25243de1c0c4f27dcac348b072cc6ea8ffd283dc033b877a37099fd809415",
"source_checkpoint": "runs/jax-ppo-match/2026-07-15_031529_match-scaled/latest",
"source_config": "match-selfplay.yaml",
"observation_size": 501,
"action_size": 96,
"hidden_size": 512,
"num_layers": 3,
"dtype": "float32",
"validation_max_abs_error": 9.1552734375e-05
}
Binary file not shown.
+252 -64
View File
@@ -10,16 +10,35 @@ import {
currentHandSorted, currentHandSorted,
decodeAction, decodeAction,
encodeAction, encodeAction,
legalActionMask,
previewPlacement, previewPlacement,
resetFromOrder,
step,
} from "./game/engine"; } from "./game/engine";
import { deckOrderFromSeed, normalizeSeed, randomSeed } from "./game/random"; import {
import { DRAW_DECK, N_CARDS, PLAY, type GameState, type PlaceType } from "./game/types"; N_ROUNDS,
matchFromOrders,
matchLegalActionMask,
matchScore,
matchStep,
type MatchState,
} from "./game/match";
import { MODEL_CODENAME, MODEL_HASH } from "./model/policy";
import { matchFromSeed, normalizeSeed, randomSeed } from "./game/random";
import { parseSavedGame, SAVED_GAME_KEY, type SavedGame } from "./game/persistence";
import { DRAW_DECK, N_CARDS, PLAY, type PlaceType } from "./game/types";
import { fallbackHeuristicPolicy, loadPolicy, type Policy } from "./model/policy"; import { fallbackHeuristicPolicy, loadPolicy, type Policy } from "./model/policy";
import { useCardMotion } from "./ui/useCardMotion"; import { useCardMotion } from "./ui/useCardMotion";
/** One deal, or the classic three-round match decided on the summed total. */
export type Mode = 1 | 3;
const MODE_KEY = "lost-cities.mode";
function loadMode(): Mode {
const requested = new URLSearchParams(window.location.search).get("rounds");
if (requested === "3") return 3;
if (requested === "1") return 1;
return window.localStorage.getItem(MODE_KEY) === "3" ? 3 : 1;
}
interface Hint { interface Hint {
action: number; action: number;
text: string; text: string;
@@ -36,8 +55,9 @@ const EMPTY_SELECTION: Selection = { handSlot: null, placeType: null };
* makes a finished game reviewable ply by ply. * makes a finished game reviewable ply by ply.
*/ */
interface Frame { interface Frame {
state: GameState; state: MatchState;
selection: Selection; selection: Selection;
move?: { player: 0 | 1; action: number };
} }
// The hand row is centered, so the binding constraint is the score plaque // The hand row is centered, so the binding constraint is the score plaque
@@ -66,15 +86,12 @@ function handCardStep(viewportWidth: number, count: number): number {
return Math.max(minVisible, Math.min(relaxed, fitted)); return Math.max(minVisible, Math.min(relaxed, fitted));
} }
/** A `?seed=` in the URL loads that exact deal, so a game can be shared or replayed. */ /** A `?seed=` in the URL loads that exact deal, so a game can be shared or replayed.
function initialSeed(): string { * In match mode the seed fixes all three deals and the coin flips too, so the
const fromUrl = new URLSearchParams(window.location.search).get("seed"); * whole match is a pure function of it. */
const seed = fromUrl === null ? "" : normalizeSeed(fromUrl); function openingFrame(seed: string, mode: Mode): Frame {
return seed === "" ? randomSeed() : seed; const { deckOrders, coinFlips } = matchFromSeed(seed);
} return { state: matchFromOrders(deckOrders, coinFlips, mode), selection: EMPTY_SELECTION };
function openingFrame(seed: string): Frame {
return { state: resetFromOrder(deckOrderFromSeed(seed)), selection: EMPTY_SELECTION };
} }
/** Keep the address bar in sync so the current deal stays shareable. */ /** Keep the address bar in sync so the current deal stays shareable. */
@@ -84,29 +101,66 @@ function publishSeed(seed: string): void {
window.history.replaceState(null, "", url); window.history.replaceState(null, "", url);
} }
function newGameId(): string {
if (typeof crypto.randomUUID === "function") return crypto.randomUUID();
if (typeof crypto.getRandomValues === "function") {
const bytes = crypto.getRandomValues(new Uint8Array(16));
bytes[6] = (bytes[6] & 0x0f) | 0x40;
bytes[8] = (bytes[8] & 0x3f) | 0x80;
const hex = Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join("");
return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20)}`;
}
return `game-${Date.now()}-${Math.random().toString(36).slice(2)}`;
}
function initialGame(): Pick<SavedGame, "seed" | "frames" | "cursor" | "resultOpen" | "mode"> {
const requested = new URLSearchParams(window.location.search).get("seed");
const mode = loadMode();
const saved = parseSavedGame(window.localStorage.getItem(SAVED_GAME_KEY));
const requestedSeed = requested === null ? null : normalizeSeed(requested);
// A saved game only resumes into the mode it was played in.
if (saved && saved.mode === mode && (requestedSeed === null || requestedSeed === saved.seed)) {
return saved;
}
const seed = requestedSeed || randomSeed();
return { seed, mode, frames: [openingFrame(seed, mode)], cursor: 0, resultOpen: true };
}
function App() { function App() {
const [seed, setSeed] = useState<string>(initialSeed); const [initial] = useState(initialGame);
const [frames, setFrames] = useState<Frame[]>(() => [openingFrame(seed)]); const [seed, setSeed] = useState(initial.seed);
const [cursor, setCursor] = useState(0); const [mode, setMode] = useState<Mode>(initial.mode as Mode);
const [frames, setFrames] = useState<Frame[]>(initial.frames);
const [cursor, setCursor] = useState(initial.cursor);
const [seedDraft, setSeedDraft] = useState(""); const [seedDraft, setSeedDraft] = useState("");
const [resultOpen, setResultOpen] = useState(true); const [resultOpen, setResultOpen] = useState(initial.resultOpen);
const [policy, setPolicy] = useState<Policy | null>(null); const [policy, setPolicy] = useState<Policy | null>(null);
const [modelMessage, setModelMessage] = useState("LOADING FINAL PPO"); const [modelMessage, setModelMessage] = useState("LOADING BOREALIS");
const [thinking, setThinking] = useState(false); const [thinking, setThinking] = useState(false);
const [menuOpen, setMenuOpen] = useState(false); const [menuOpen, setMenuOpen] = useState(false);
const [viewportWidth, setViewportWidth] = useState(() => window.innerWidth); const [viewportWidth, setViewportWidth] = useState(() => window.innerWidth);
const [hint, setHint] = useState<Hint | null>(null); const [hint, setHint] = useState<Hint | null>(null);
const [hintPending, setHintPending] = useState(false); const [hintPending, setHintPending] = useState(false);
const generation = useRef(0); const generation = useRef(0);
const gameId = useRef(newGameId());
const startedAt = useRef(new Date().toISOString());
const loggedGame = useRef<string | null>(null);
const { cardRef, deckRef, overlayRef, resetMotion } = useCardMotion(); const { cardRef, deckRef, overlayRef, resetMotion } = useCardMotion();
const { state, selection } = frames[cursor]; const { state: match, selection } = frames[cursor];
// The board on screen is the round in play; the match is what decides the game.
const state = match.round;
const canUndo = cursor > 0; const canUndo = cursor > 0;
const canRedo = cursor < frames.length - 1; const canRedo = cursor < frames.length - 1;
const latestState = useRef(state); const latestMatch = useRef(match);
useEffect(() => { publishSeed(seed); }, [seed]); useEffect(() => { publishSeed(seed); }, [seed]);
useEffect(() => { latestState.current = state; }, [state]); useEffect(() => { latestMatch.current = match; }, [match]);
useEffect(() => { window.localStorage.setItem(MODE_KEY, String(mode)); }, [mode]);
useEffect(() => {
const saved: SavedGame = { version: 2, seed, mode, frames, cursor, resultOpen };
window.localStorage.setItem(SAVED_GAME_KEY, JSON.stringify(saved));
}, [cursor, frames, mode, resultOpen, seed]);
useEffect(() => { useEffect(() => {
function onResize() { setViewportWidth(window.innerWidth); } function onResize() { setViewportWidth(window.innerWidth); }
@@ -116,7 +170,7 @@ function App() {
const humanHand = useMemo(() => currentHandSorted(state, 0), [state]); const humanHand = useMemo(() => currentHandSorted(state, 0), [state]);
const opponentHand = useMemo(() => currentHandSorted(state, 1), [state]); const opponentHand = useMemo(() => currentHandSorted(state, 1), [state]);
const legal = useMemo(() => legalActionMask(state), [state]); const legal = useMemo(() => matchLegalActionMask(match), [match]);
const selectedCard = selection.handSlot === null ? null : humanHand[selection.handSlot]; const selectedCard = selection.handSlot === null ? null : humanHand[selection.handSlot];
const selectedColor = selectedCard === null ? null : cardColor(selectedCard); const selectedColor = selectedCard === null ? null : cardColor(selectedCard);
const displayState = useMemo( const displayState = useMemo(
@@ -125,14 +179,23 @@ function App() {
: previewPlacement(state, selection.handSlot, selection.placeType), : previewPlacement(state, selection.handSlot, selection.placeType),
[selection, state], [selection, state],
); );
const scores = useMemo(() => boardScore(displayState), [displayState]); /** This round's board. */
const roundScores = useMemo(() => boardScore(displayState), [displayState]);
/** Rounds already banked plus the board in play — the number that decides a match. */
const scores = useMemo(
(): [number, number] => [
match.carry[0] + roundScores[0],
match.carry[1] + roundScores[1],
],
[match.carry, roundScores],
);
useEffect(() => { useEffect(() => {
let cancelled = false; let cancelled = false;
loadPolicy().then(({ policy: loaded, warning }) => { loadPolicy().then(({ policy: loaded, warning }) => {
if (cancelled) return; if (cancelled) return;
setPolicy(loaded); setPolicy(loaded);
setModelMessage(warning ?? `${loaded.provider.toUpperCase()} · FINAL PPO`); setModelMessage(warning ?? `${loaded.provider.toUpperCase()} · BOREALIS`);
}); });
return () => { cancelled = true; }; return () => { cancelled = true; };
}, []); }, []);
@@ -149,24 +212,28 @@ function App() {
const rivalSuspended = canRedo; const rivalSuspended = canRedo;
useEffect(() => { useEffect(() => {
if (!policy || rivalSuspended || state.done || state.toMove !== 1) return; if (!policy || rivalSuspended || match.done || state.toMove !== 1) return;
const currentGeneration = generation.current; const currentGeneration = generation.current;
const timer = window.setTimeout(async () => { const timer = window.setTimeout(async () => {
setThinking(true); setThinking(true);
try { try {
let action: number; let action: number;
try { try {
action = await policy.action(state); action = await policy.action(match);
} catch (error) { } catch (error) {
console.error("AI action failed; falling back to heuristic policy", error); console.error("AI action failed; falling back to heuristic policy", error);
const fallback = fallbackHeuristicPolicy(); const fallback = fallbackHeuristicPolicy();
action = await fallback.action(state); action = await fallback.action(match);
if (generation.current !== currentGeneration) return; if (generation.current !== currentGeneration) return;
setPolicy(fallback); setPolicy(fallback);
setModelMessage("HEURISTIC FALLBACK (MODEL ERROR)"); setModelMessage("HEURISTIC FALLBACK (MODEL ERROR)");
} }
if (generation.current !== currentGeneration) return; if (generation.current !== currentGeneration) return;
pushFrame({ state: step(state, action), selection: EMPTY_SELECTION }); pushFrame({
state: matchStep(match, action),
selection: EMPTY_SELECTION,
move: { player: 1, action },
});
} catch (error) { } catch (error) {
console.error("AI action failed even with heuristic fallback", error); console.error("AI action failed even with heuristic fallback", error);
setModelMessage("MODEL INFERENCE ERROR"); setModelMessage("MODEL INFERENCE ERROR");
@@ -175,15 +242,19 @@ function App() {
} }
}, 620); }, 620);
return () => window.clearTimeout(timer); return () => window.clearTimeout(timer);
}, [policy, rivalSuspended, state]); }, [policy, rivalSuspended, match, state.toMove]);
/** Deal a game. Without a seed this rolls a fresh one; the same seed always /** Deal a game. Without a seed this rolls a fresh one; the same seed always
* reproduces the same deal, so `restart(seed)` also replays the current one. */ * reproduces the same deal, so `restart(seed)` also replays the current one. */
function restart(nextSeed: string = randomSeed()) { function restart(nextSeed: string = randomSeed(), nextMode: Mode = mode) {
generation.current += 1; generation.current += 1;
resetMotion(); resetMotion();
gameId.current = newGameId();
startedAt.current = new Date().toISOString();
loggedGame.current = null;
setSeed(nextSeed); setSeed(nextSeed);
setFrames([openingFrame(nextSeed)]); setMode(nextMode);
setFrames([openingFrame(nextSeed, nextMode)]);
setCursor(0); setCursor(0);
setSeedDraft(""); setSeedDraft("");
setResultOpen(true); setResultOpen(true);
@@ -217,8 +288,8 @@ function App() {
/** Ask the policy driving the rival what it would do in your seat. */ /** Ask the policy driving the rival what it would do in your seat. */
async function requestHint() { async function requestHint() {
if (!policy || state.done || state.toMove !== 0 || hintPending) return; if (!policy || match.done || state.toMove !== 0 || hintPending) return;
const position = state; const position = match;
setHintPending(true); setHintPending(true);
setMenuOpen(false); setMenuOpen(false);
try { try {
@@ -229,10 +300,10 @@ function App() {
const [best] = ranked; const [best] = ranked;
// The position can move on while inference runs — a hint for a stale // The position can move on while inference runs — a hint for a stale
// position would point at the wrong hand slot. // position would point at the wrong hand slot.
if (best === undefined || latestState.current !== position) return; if (best === undefined || latestMatch.current !== position) return;
setHint({ setHint({
action: best.action, action: best.action,
text: describeAction(position, best.action), text: describeAction(position.round, best.action),
probability: best.probability, probability: best.probability,
}); });
} catch (error) { } catch (error) {
@@ -246,21 +317,21 @@ function App() {
useEffect(() => { setHint(null); }, [cursor, frames]); useEffect(() => { setHint(null); }, [cursor, frames]);
function chooseCard(handSlot: number) { function chooseCard(handSlot: number) {
if (state.toMove !== 0 || state.done) return; if (state.toMove !== 0 || match.done) return;
const next = selection.handSlot !== handSlot || selection.placeType !== null const next = selection.handSlot !== handSlot || selection.placeType !== null
? { handSlot, placeType: null } ? { handSlot, placeType: null }
: EMPTY_SELECTION; : EMPTY_SELECTION;
pushFrame({ state, selection: next }); pushFrame({ state: match, selection: next });
} }
function choosePlace(placeType: PlaceType) { function choosePlace(placeType: PlaceType) {
if (selection.handSlot === null) return; if (selection.handSlot === null) return;
pushFrame({ state, selection: { ...selection, placeType } }); pushFrame({ state: match, selection: { ...selection, placeType } });
} }
function cancelPlace() { function cancelPlace() {
if (selection.handSlot === null || selection.placeType === null) return; if (selection.handSlot === null || selection.placeType === null) return;
pushFrame({ state, selection: { ...selection, placeType: null } }); pushFrame({ state: match, selection: { ...selection, placeType: null } });
} }
useEffect(() => { useEffect(() => {
@@ -282,7 +353,11 @@ function App() {
function commit(drawSource: number) { function commit(drawSource: number) {
if (!legalDraw(drawSource) || selection.handSlot === null || selection.placeType === null) return; if (!legalDraw(drawSource) || selection.handSlot === null || selection.placeType === null) return;
const action = encodeAction(selection.handSlot, selection.placeType, drawSource); const action = encodeAction(selection.handSlot, selection.placeType, drawSource);
pushFrame({ state: step(state, action), selection: EMPTY_SELECTION }); pushFrame({
state: matchStep(match, action),
selection: EMPTY_SELECTION,
move: { player: 0, action },
});
} }
// Cards actually drawn in the hand row (the selected card moves into the // Cards actually drawn in the hand row (the selected card moves into the
@@ -309,18 +384,83 @@ function App() {
const hintConfidence = hint !== null && hint.probability !== null const hintConfidence = hint !== null && hint.probability !== null
? `${Math.round(hint.probability * 100)}%` ? `${Math.round(hint.probability * 100)}%`
: null; : null;
const canHint = policy !== null && !state.done && state.toMove === 0; const canHint = policy !== null && !match.done && state.toMove === 0;
const outcome = scores[0] === scores[1] const outcome = scores[0] === scores[1]
? "DRAW" ? "DRAW"
: scores[0] > scores[1] ? "YOU WIN" : "THE RIVAL WINS"; : scores[0] > scores[1] ? "YOU WIN" : "THE RIVAL WINS";
useEffect(() => {
if (!match.done || loggedGame.current === gameId.current) return;
loggedGame.current = gameId.current;
const moves = frames.flatMap((frame, index) => {
if (!frame.move || index === 0) return [];
const beforeMatch = frames[index - 1].state;
const afterMatch = frame.state;
const before = beforeMatch.round;
const after = afterMatch.round;
const decoded = decodeAction(frame.move.action);
const beforeHand = currentHandSorted(before, frame.move.player);
// A move that ends a round is followed by a fresh deal, so the drawn card
// has to be read from the hand it was drawn into, not from the next round's.
const rolledOver = afterMatch.roundIdx !== beforeMatch.roundIdx;
const afterHand = rolledOver ? beforeHand : currentHandSorted(after, frame.move.player);
return [{
ply: before.stepCount + 1,
round: beforeMatch.roundIdx + 1,
player: frame.move.player === 0 ? "human" : "ai",
action: frame.move.action,
playedCard: beforeHand[decoded.handSlot],
placeType: decoded.placeType === PLAY ? "play" : "discard",
drawSource: decoded.drawSource === DRAW_DECK ? "deck" : `discard_${decoded.drawSource - 1}`,
drawnCard: rolledOver
? null
: afterHand.find((card) => !beforeHand.includes(card)) ?? null,
handsBefore: {
human: currentHandSorted(before, 0),
ai: currentHandSorted(before, 1),
},
scoresAfter: matchScore(afterMatch),
}];
});
const record = {
format: "lost-cities-web-game-v2",
gameId: gameId.current,
seed,
// Which model actually played. The old records stored the on-screen label,
// which stops identifying anything the moment there is a second model.
opponent: { codename: MODEL_CODENAME, hash: MODEL_HASH },
mode,
startedAt: startedAt.current,
finishedAt: new Date().toISOString(),
deckOrders: match.deckOrders.slice(0, mode),
coinFlips: match.coinFlips.slice(0, mode),
moves,
finalScores: scores,
roundScores: match.carry,
outcome,
finalState: match,
policy: modelMessage,
};
void fetch("./api/game-records", {
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify(record),
keepalive: true,
}).then((response) => {
if (!response.ok) throw new Error(`game log HTTP ${response.status}`);
}).catch((error) => {
loggedGame.current = null;
console.warn("Game record was not saved", error);
});
}, [frames, match, mode, modelMessage, outcome, scores, seed, state.done]);
// While the rival is paused on its own turn (or on the final position) nothing // While the rival is paused on its own turn (or on the final position) nothing
// will happen until the move is redone or play is resumed — say so. On your own // will happen until the move is redone or play is resumed — say so. On your own
// turn the normal prompt still applies: you can simply play on from here. // turn the normal prompt still applies: you can simply play on from here.
const reviewingRival = rivalSuspended && (state.done || state.toMove === 1); const reviewingRival = rivalSuspended && (match.done || state.toMove === 1);
const status = reviewingRival const status = reviewingRival
? `Reviewing ply ${state.stepCount} — redo, or play on from here` ? `Reviewing ply ${state.stepCount} — redo, or play on from here`
: state.done : match.done
? outcome ? outcome
: state.toMove === 1 : state.toMove === 1
? thinking ? "The rival is thinking ···" : "The rival's turn" ? thinking ? "The rival is thinking ···" : "The rival's turn"
@@ -350,6 +490,24 @@ function App() {
<div className="menu-popover"> <div className="menu-popover">
<button onClick={() => restart()}>NEW GAME</button> <button onClick={() => restart()}>NEW GAME</button>
<button onClick={() => restart(seed)}>REPLAY THIS SEED</button> <button onClick={() => restart(seed)}>REPLAY THIS SEED</button>
<div className="menu-modes" role="group" aria-label="Game length">
<button
className={mode === 1 ? "is-active" : ""}
onClick={() => restart(randomSeed(), 1)}
>
ONE DEAL
</button>
<button
className={mode === 3 ? "is-active" : ""}
onClick={() => restart(randomSeed(), 3)}
>
MATCH · 3 ROUNDS
</button>
</div>
<span>
A MATCH IS DECIDED ON THE SUMMED TOTAL · WHOEVER LEADS ON POINTS
OPENS THE NEXT ROUND
</span>
<label className="menu-seed"> <label className="menu-seed">
<span>DEAL SEED · {seed}</span> <span>DEAL SEED · {seed}</span>
<input <input
@@ -367,6 +525,24 @@ function App() {
)} )}
</div> </div>
{mode === N_ROUNDS && (
<div className="round-strip" aria-label={`Round ${match.roundIdx + 1} of ${N_ROUNDS}`}>
{Array.from({ length: N_ROUNDS }, (_, round) => (
<span
key={round}
className={
round < match.roundIdx ? "is-done" : round === match.roundIdx ? "is-live" : ""
}
>
R{round + 1}
</span>
))}
<em>
{match.carry[0]} : {match.carry[1]} BANKED
</em>
</div>
)}
<div className="table-center"> <div className="table-center">
<Board <Board
state={displayState} state={displayState}
@@ -398,14 +574,6 @@ function App() {
<p className={`turn-prompt ${thinking ? "turn-prompt--thinking" : ""} ${reviewingRival ? "turn-prompt--review" : ""}`}> <p className={`turn-prompt ${thinking ? "turn-prompt--thinking" : ""} ${reviewingRival ? "turn-prompt--review" : ""}`}>
{hint ? `Hint — ${hint.text}${hintConfidence ? ` (${hintConfidence})` : ""}` : status} {hint ? `Hint — ${hint.text}${hintConfidence ? ` (${hintConfidence})` : ""}` : status}
</p> </p>
<button
type="button"
className={`hint-button ${hint ? "is-active" : ""}`}
onClick={requestHint}
disabled={!canHint || hintPending}
>
{hintPending ? "THINKING…" : "HINT"}
</button>
</div> </div>
<section className="human-hand" aria-label="Your hand"> <section className="human-hand" aria-label="Your hand">
@@ -415,7 +583,7 @@ function App() {
key={card} key={card}
selected={selection.handSlot === slot} selected={selection.handSlot === slot}
hinted={hintedCard === card} hinted={hintedCard === card}
disabled={state.toMove !== 0 || state.done} disabled={state.toMove !== 0 || match.done}
onClick={() => chooseCard(slot)} onClick={() => chooseCard(slot)}
style={position === 0 ? undefined : { marginLeft: handCardMarginLeft }} style={position === 0 ? undefined : { marginLeft: handCardMarginLeft }}
innerRef={cardRef(card)} innerRef={cardRef(card)}
@@ -424,24 +592,41 @@ function App() {
</section> </section>
<section className="score-plaque score-plaque--human" aria-label={`Your score ${scores[0]}`}> <section className="score-plaque score-plaque--human" aria-label={`Your score ${scores[0]}`}>
<div><strong>YOU <i /></strong><small>{state.toMove === 0 && !state.done ? "YOUR TURN" : "EXPEDITION LEAD"}</small></div><b>{scores[0]}</b> <div><strong>YOU <i /></strong><small>{state.toMove === 0 && !match.done ? "YOUR TURN" : "EXPEDITION LEAD"}</small></div><b>{scores[0]}</b>
</section> </section>
<div className="history-bar"> <div className="control-stack">
<button
type="button"
className={`hint-button ${hint ? "is-active" : ""}`}
onClick={requestHint}
disabled={!canHint || hintPending}
>
{hintPending ? "THINKING…" : "HINT"}
</button>
<button type="button" onClick={undo} disabled={!canUndo} aria-label="Undo one action" title="Undo (←)"> <button type="button" onClick={undo} disabled={!canUndo} aria-label="Undo one action" title="Undo (←)">
<span>UNDO</span> <span>UNDO</span>
</button> </button>
<button type="button" onClick={redo} disabled={!canRedo} aria-label="Redo one action" title="Redo (→)"> <button type="button" onClick={redo} disabled={!canRedo} aria-label="Redo one action" title="Redo (→)">
<span>REDO</span> <span>REDO</span>
</button> </button>
{rivalSuspended && ( <button
<button type="button" className="history-bar__resume" onClick={resumeFromHere}> type="button"
PLAY FROM HERE className="control-stack__resume"
</button> onClick={resumeFromHere}
)} disabled={!rivalSuspended}
{state.done && !resultOpen && ( >
<button type="button" onClick={() => setResultOpen(true)}>SCORE</button> PLAY FROM HERE
)} </button>
<button
type="button"
className={`control-stack__score ${state.done && !resultOpen ? "is-available" : ""}`}
onClick={() => setResultOpen(true)}
disabled={!state.done || resultOpen}
aria-hidden={!state.done || resultOpen}
>
SCORE
</button>
</div> </div>
{/* Cards in flight back to the deck (undo of a draw) are re-parented here {/* Cards in flight back to the deck (undo of a draw) are re-parented here
@@ -453,6 +638,9 @@ function App() {
state={state} state={state}
outcome={outcome} outcome={outcome}
seed={seed} seed={seed}
carry={match.carry}
roundIdx={match.roundIdx}
totalRounds={mode}
onReview={() => setResultOpen(false)} onReview={() => setResultOpen(false)}
onPlayAgain={() => restart()} onPlayAgain={() => restart()}
/> />
+40 -4
View File
@@ -6,6 +6,15 @@ interface ResultCardProps {
state: GameState; state: GameState;
outcome: string; outcome: string;
seed: string; seed: string;
/**
* Rounds already banked, when this ends a match. The table below breaks down the
* board in front of you -- the final round -- but a match is decided on the sum,
* so the headline and the total have to carry the earlier rounds or they name
* the wrong winner.
*/
carry?: [number, number];
roundIdx?: number;
totalRounds?: number;
onReview: () => void; onReview: () => void;
onPlayAgain: () => void; onPlayAgain: () => void;
} }
@@ -14,17 +23,35 @@ function signed(value: number): string {
return value > 0 ? `+${value}` : String(value); return value > 0 ? `+${value}` : String(value);
} }
export function ResultCard({ state, outcome, seed, onReview, onPlayAgain }: ResultCardProps) { export function ResultCard({
state,
outcome,
seed,
carry = [0, 0],
roundIdx = 0,
totalRounds = 1,
onReview,
onPlayAgain,
}: ResultCardProps) {
const you = scoreBreakdown(state, 0); const you = scoreBreakdown(state, 0);
const rival = scoreBreakdown(state, 1); const rival = scoreBreakdown(state, 1);
const isMatch = totalRounds > 1;
const matchTotals: [number, number] = [carry[0] + you.total, carry[1] + rival.total];
return ( return (
<div className="result-overlay"> <div className="result-overlay">
<section className="result-card" aria-label="Final score"> <section className="result-card" aria-label="Final score">
<header className="result-card__head"> <header className="result-card__head">
<p>ROUND COMPLETE</p> <p>{isMatch ? `MATCH COMPLETE · ${totalRounds} ROUNDS` : "ROUND COMPLETE"}</p>
<h1>{outcome}</h1> <h1>{outcome}</h1>
<strong>{you.total} <i>:</i> {rival.total}</strong> <strong>{matchTotals[0]} <i>:</i> {matchTotals[1]}</strong>
{isMatch && (
<small className="result-card__banked">
ROUND {roundIdx + 1} · {signed(you.total)} : {signed(rival.total)}
{" · BANKED "}
{signed(carry[0])} : {signed(carry[1])}
</small>
)}
</header> </header>
<table className="result-table"> <table className="result-table">
@@ -69,12 +96,21 @@ export function ResultCard({ state, outcome, seed, onReview, onPlayAgain }: Resu
</tbody> </tbody>
<tfoot> <tfoot>
<tr> <tr>
<th scope="row">TOTAL</th> <th scope="row">{isMatch ? `ROUND ${roundIdx + 1}` : "TOTAL"}</th>
<td className="result-table__detail" /> <td className="result-table__detail" />
<td>{signed(you.total)}</td> <td>{signed(you.total)}</td>
<td className="result-table__detail" /> <td className="result-table__detail" />
<td>{signed(rival.total)}</td> <td>{signed(rival.total)}</td>
</tr> </tr>
{isMatch && (
<tr className="result-table__match">
<th scope="row">MATCH TOTAL</th>
<td className="result-table__detail" />
<td>{signed(matchTotals[0])}</td>
<td className="result-table__detail" />
<td>{signed(matchTotals[1])}</td>
</tr>
)}
</tfoot> </tfoot>
</table> </table>
+6 -2
View File
@@ -32,7 +32,11 @@ export function shuffledDeck(random: () => number = Math.random): number[] {
return deck; return deck;
} }
export function resetFromOrder(deckOrder: number[]): GameState { /**
* `firstPlayer` moves first. Rounds two and three of a classic match are led by
* whoever is ahead on points, so the match layer sets this per round.
*/
export function resetFromOrder(deckOrder: number[], firstPlayer: 0 | 1 | number = 0): GameState {
if (deckOrder.length !== N_CARDS || new Set(deckOrder).size !== N_CARDS) { if (deckOrder.length !== N_CARDS || new Set(deckOrder).size !== N_CARDS) {
throw new Error("deckOrder must be a permutation of 0..59"); throw new Error("deckOrder must be a permutation of 0..59");
} }
@@ -48,7 +52,7 @@ export function resetFromOrder(deckOrder: number[]): GameState {
colHandshakes: emptyMatrix(2, N_COLORS), colHandshakes: emptyMatrix(2, N_COLORS),
colLength: emptyMatrix(2, N_COLORS), colLength: emptyMatrix(2, N_COLORS),
piles: Array.from({ length: N_COLORS }, () => []), piles: Array.from({ length: N_COLORS }, () => []),
toMove: 0, toMove: (firstPlayer === 1 ? 1 : 0) as 0 | 1,
stepCount: 0, stepCount: 0,
done: false, done: false,
}; };
File diff suppressed because one or more lines are too long
+116
View File
@@ -0,0 +1,116 @@
import { boardScore, legalActionMask, resetFromOrder, step } from "./engine";
import type { GameState } from "./types";
export const N_ROUNDS = 3;
/**
* Classic Lost Cities: three rounds, scores summed, highest total wins.
*
* The single-round engine is left alone -- it is the rules oracle the Python
* engine is differential-tested against. Only two rules live up here, both from
* the Kosmos rulebook:
*
* "If after three games you have the highest overall score, you win."
* "The player who has more points begins" the next game -- not alternating.
*
* The rulebook says nothing about an exact tie, so the starter falls back to a
* coin flip, drawn up front with the deals. This mirrors src/lost_cities_jax/match.py;
* the two are checked against each other in match.test.ts.
*/
export interface MatchState {
round: GameState;
/** Every deal of the match, shuffled up front. */
deckOrders: number[][];
/** Tie-breaking starters, used only when the scores are level. */
coinFlips: number[];
roundIdx: number;
/** Points banked by each player in the rounds already finished. */
carry: [number, number];
/**
* 1 for a one-off deal, 3 for the classic match.
*
* The observation always reports the round index out of three regardless --
* that is the space the policy was trained on, and a one-off deal is simply its
* round one, played at a carry of zero.
*/
totalRounds: number;
done: boolean;
}
/**
* Whoever has banked more points leads; level scores fall back to the coin.
*
* Round one needs no special case: carry is (0, 0) there, so the tie branch
* already picks the coin flip, which is exactly the rulebook's arbitrary
* "oldest player begins".
*/
export function startingPlayer(
carry: readonly [number, number],
roundIdx: number,
coinFlips: readonly number[],
): number {
const lead = carry[0] - carry[1];
if (lead > 0) return 0;
if (lead < 0) return 1;
return coinFlips[roundIdx];
}
export function matchFromOrders(
deckOrders: number[][],
coinFlips: number[],
totalRounds: number = N_ROUNDS,
): MatchState {
const carry: [number, number] = [0, 0];
return {
round: resetFromOrder(deckOrders[0], startingPlayer(carry, 0, coinFlips)),
deckOrders,
coinFlips,
roundIdx: 0,
carry,
totalRounds,
done: false,
};
}
/** Running totals: rounds already banked plus the board in play. */
export function matchScore(state: MatchState): [number, number] {
const board = boardScore(state.round);
return [state.carry[0] + board[0], state.carry[1] + board[1]];
}
export function matchLegalActionMask(state: MatchState): boolean[] {
const mask = legalActionMask(state.round);
return state.done ? mask.map(() => false) : mask;
}
/**
* Play one ply. Rolls into the next round when the deck runs out.
*
* Rounds one and two pay nothing -- they only bank into `carry`. Only the sum
* decides the match.
*/
export function matchStep(state: MatchState, action: number): MatchState {
const played = step(state.round, action);
if (!played.done || state.done) {
return { ...state, round: played };
}
const board = boardScore(played);
const banked: [number, number] = [state.carry[0] + board[0], state.carry[1] + board[1]];
if (state.roundIdx >= state.totalRounds - 1) {
return { ...state, round: played, done: true };
}
const roundIdx = state.roundIdx + 1;
return {
...state,
round: resetFromOrder(
state.deckOrders[roundIdx],
startingPlayer(banked, roundIdx, state.coinFlips),
),
roundIdx,
carry: banked,
done: false,
};
}
+122
View File
@@ -0,0 +1,122 @@
import { describe, expect, it } from "vitest";
import fixture from "./match-parity-fixture.json";
import {
N_ROUNDS,
matchFromOrders,
matchLegalActionMask,
matchScore,
matchStep,
startingPlayer,
type MatchState,
} from "./match";
import { MATCH_OBS_DIM, matchObservation } from "./matchObservation";
import type { GameState } from "./types";
interface FixtureRow {
match: {
round: GameState;
deckOrders: number[][];
coinFlips: number[];
roundIdx: number;
carry: number[];
done: boolean;
};
player: number;
observation: number[];
}
const rows = fixture.rows as unknown as FixtureRow[];
function toMatch(row: FixtureRow): MatchState {
return {
round: row.match.round,
deckOrders: row.match.deckOrders,
coinFlips: row.match.coinFlips,
roundIdx: row.match.roundIdx,
carry: [row.match.carry[0], row.match.carry[1]],
// The fixture is generated from the JAX match, which is always three rounds.
totalRounds: N_ROUNDS,
done: row.match.done,
};
}
function firstLegal(match: MatchState): number {
const index = matchLegalActionMask(match).findIndex(Boolean);
if (index < 0) throw new Error("no legal action");
return index;
}
describe("match observation parity with JAX", () => {
it("has the dimension the exported model expects", () => {
expect(MATCH_OBS_DIM).toBe(fixture.obsDim);
expect(MATCH_OBS_DIM).toBe(501);
});
it("reproduces every fixture observation", () => {
expect(rows.length).toBeGreaterThan(100);
let worst = 0;
let worstAt = "";
for (const [index, row] of rows.entries()) {
const actual = matchObservation(toMatch(row), row.player);
expect(actual.length).toBe(row.observation.length);
for (let i = 0; i < actual.length; i += 1) {
const delta = Math.abs(actual[i] - row.observation[i]);
if (delta > worst) {
worst = delta;
worstAt = `row ${index}, feature ${i}`;
}
}
}
// A mismatch here throws nowhere: the ONNX policy consumes the wrong vector
// and plays worse for reasons nobody can see. So the bar is float32 round-off,
// not "close enough".
expect(worst, `largest disagreement at ${worstAt}`).toBeLessThan(1e-5);
});
it("covers positions past a round roll-over, with a real carry", () => {
expect(rows.some((row) => row.match.roundIdx > 0)).toBe(true);
expect(rows.some((row) => row.match.carry[0] !== row.match.carry[1])).toBe(true);
});
});
describe("match rules", () => {
it("plays exactly three rounds and then ends", () => {
let match = matchFromOrders(rows[0].match.deckOrders, rows[0].match.coinFlips);
const seen = new Set<number>();
for (let ply = 0; ply < 1400 && !match.done; ply += 1) {
seen.add(match.roundIdx);
match = matchStep(match, firstLegal(match));
}
expect(match.done).toBe(true);
expect(seen).toEqual(new Set([0, 1, 2]));
expect(match.roundIdx).toBe(N_ROUNDS - 1);
});
it("lets whoever has more points begin, and flips a coin when level", () => {
expect(startingPlayer([60, 10], 1, [1, 1, 1])).toBe(0); // ahead leads, coin ignored
expect(startingPlayer([10, 60], 1, [0, 0, 0])).toBe(1);
expect(startingPlayer([30, 30], 1, [0, 0, 0])).toBe(0); // level falls back to the coin
expect(startingPlayer([30, 30], 1, [1, 1, 1])).toBe(1);
});
it("keeps the running total continuous across a round boundary", () => {
let match = matchFromOrders(rows[0].match.deckOrders, rows[0].match.coinFlips);
for (let ply = 0; ply < 1400 && !match.done; ply += 1) {
const before = match.roundIdx;
const next = matchStep(match, firstLegal(match));
if (next.roundIdx !== before) {
// The finished round folds into carry and the fresh board is empty, so
// the roll-over itself moves nothing.
expect(matchScore(next)).toEqual([next.carry[0], next.carry[1]]);
}
match = next;
}
expect(match.done).toBe(true);
});
});
+140
View File
@@ -0,0 +1,140 @@
import { cardColor } from "./cards";
import { boardScore } from "./engine";
import { N_ROUNDS, type MatchState } from "./match";
import { observation } from "./observation";
import {
CARDS_PER_COLOR,
LOC_DECK,
LOC_DISCARD,
LOC_P0_HAND,
N_CARDS,
N_COLORS,
OBS_DIM,
} from "./types";
/**
* Player-view observation for a three-round match: the mirror of
* src/lost_cities_jax/match_obs.py.
*
* These two must agree to the bit. A mismatch does not throw -- the model happily
* consumes a wrong vector and plays worse for reasons nobody can see. That is why
* matchObservation.test.ts checks this against fixtures generated from the Python
* side rather than trusting the port.
*
* On top of the single-round observation this adds the four things a match policy
* cannot play without: carry (scalar *and* binned, because round three is a
* threshold problem); which round it is; whose turn it is; and the deck clock,
* since a round ends on the last deck draw and players bend that parity by drawing
* from discard piles. Live points per colour are split by hand / discard pile /
* unseen, because a discard pile is public and recoverable.
*/
/** Packed tightly around zero: that is where the round-three decision flips. */
const CARRY_BIN_EDGES = [-60, -30, -12, -1, 1, 12, 30, 60];
export const N_CARRY_BINS = CARRY_BIN_EDGES.length + 1;
/** A typical round margin, not the theoretical 780 maximum the old obs divided by. */
const CARRY_SCALE = 75;
/** 2+3+...+10, the most one expedition can be worth before multipliers. */
const MAX_COLOR_POINTS = 54;
const N_MATCH_SCALARS =
1 + N_CARRY_BINS + N_ROUNDS + 1 + 1 + 1 + 1 + 2 * N_COLORS * 3;
export const MATCH_OBS_DIM = OBS_DIM + N_MATCH_SCALARS;
function cardRank(card: number): number {
const slot = card % CARDS_PER_COLOR;
return slot >= 3 ? slot - 1 : 0;
}
/** numpy.digitize: the count of edges strictly below `value`. */
function digitize(value: number, edges: readonly number[]): number {
let index = 0;
while (index < edges.length && value >= edges[index]) index += 1;
return index;
}
/**
* Points still reachable for `subject`, split by where the card sits:
* [in hand, in a discard pile, unseen] per colour.
*
* Seen through `viewer`'s eyes -- a card in the opponent's hand only counts as "in
* hand" if it is public, otherwise it is unseen.
*/
function livePoints(state: MatchState, viewer: number, subject: number): number[] {
const round = state.round;
const isMine = subject === viewer;
const subjectHandLoc = LOC_P0_HAND + subject;
const out = Array.from({ length: N_COLORS }, () => [0, 0, 0]);
for (let card = 0; card < N_CARDS; card += 1) {
const color = cardColor(card);
const rank = cardRank(card);
// An ascending column can only take cards above its current top.
if (rank <= round.colTop[subject][color]) continue;
const loc = round.cardLoc[card];
const inSubjectHand = loc === subjectHandLoc;
if (inSubjectHand && (isMine || round.handPublic[card])) {
out[color][0] += rank;
} else if (loc === LOC_DISCARD) {
out[color][1] += rank;
} else if (loc === LOC_DECK || (inSubjectHand && !isMine && !round.handPublic[card])) {
// Cards in the *other* player's hidden hand are unseen to the viewer too,
// but the subject cannot reach them, so they are deliberately excluded.
out[color][2] += rank;
}
}
return out.flat().map((points) => points / MAX_COLOR_POINTS);
}
export function matchObservation(state: MatchState, player: number): Float32Array {
const opponent = 1 - player;
const round = state.round;
const base = observation(round, player);
const board = boardScore(round);
const lead =
state.carry[player] + board[player] - (state.carry[opponent] + board[opponent]);
const carryScaled = Math.min(2, Math.max(-2, lead / CARRY_SCALE));
const carryBins = Array.from({ length: N_CARRY_BINS }, (_, i) =>
Number(i === digitize(lead, CARRY_BIN_EDGES)),
);
const roundOneHot = Array.from({ length: N_ROUNDS }, (_, i) => Number(i === state.roundIdx));
const roundsLeft = (N_ROUNDS - 1 - state.roundIdx) / (N_ROUNDS - 1);
const toMove = round.toMove;
const myTurn = Number(toMove === player);
// toMove flips every ply, so the round's opener is recoverable from parity.
const roundOpener = toMove ^ (round.stepCount & 1);
const iOpened = Number(roundOpener === player);
// If both players drew from the deck from here, the last deck card falls to
// whoever is on move after `remaining - 1` more plies.
const remaining = Math.max(N_CARDS - round.drawPtr, 1);
const lastDrawer = toMove ^ ((remaining - 1) & 1);
const iTakeLast = Number(lastDrawer === player);
const extra = [
carryScaled,
...carryBins,
...roundOneHot,
roundsLeft,
myTurn,
iOpened,
iTakeLast,
...livePoints(state, player, player),
...livePoints(state, player, opponent),
];
if (extra.length !== N_MATCH_SCALARS) {
throw new Error(`match scalars ${extra.length} != ${N_MATCH_SCALARS}`);
}
const out = new Float32Array(MATCH_OBS_DIM);
out.set(base, 0);
out.set(Float32Array.from(extra), OBS_DIM);
return out;
}
+52
View File
@@ -0,0 +1,52 @@
import { describe, expect, it } from "vitest";
import { matchFromOrders } from "./match";
import { matchFromSeed } from "./random";
import { parseSavedGame } from "./persistence";
const { deckOrders, coinFlips } = matchFromSeed("abc");
const match = matchFromOrders(deckOrders, coinFlips, 3);
function savedGame(overrides: Record<string, unknown> = {}) {
return {
version: 2,
seed: "abc",
mode: 3,
frames: [{ state: match, selection: { handSlot: null, placeType: null } }],
cursor: 0,
resultOpen: true,
...overrides,
};
}
describe("saved game parsing", () => {
it("accepts a complete saved timeline", () => {
const saved = savedGame();
expect(parseSavedGame(JSON.stringify(saved))).toEqual(saved);
});
it("accepts a one-deal game", () => {
const saved = savedGame({
mode: 1,
frames: [
{ state: matchFromOrders(deckOrders, coinFlips, 1), selection: { handSlot: null, placeType: null } },
],
});
expect(parseSavedGame(JSON.stringify(saved))).toEqual(saved);
});
it("rejects corrupt and incompatible data", () => {
expect(parseSavedGame("not json")).toBeNull();
expect(parseSavedGame(JSON.stringify({ version: 2 }))).toBeNull();
expect(parseSavedGame(JSON.stringify(savedGame({ frames: [] })))).toBeNull();
expect(parseSavedGame(JSON.stringify(savedGame({ mode: 2 })))).toBeNull();
});
it("refuses a v1 save rather than guessing what it meant", () => {
// A v1 save is one round: no carry, no deals for rounds two and three, no
// coin flips. There is nothing honest to migrate it into, so it is dropped
// and a fresh game is dealt.
const v1 = { version: 1, seed: "abc", frames: [{ state: match.round, selection: {} }], cursor: 0, resultOpen: true };
expect(parseSavedGame(JSON.stringify(v1))).toBeNull();
});
});
+96
View File
@@ -0,0 +1,96 @@
import type { MatchState } from "./match";
import { N_CARDS, type GameState, type PlaceType } from "./types";
/**
* v2 because a saved game now holds a match, not a round. The key is bumped
* rather than migrated: a v1 save has no carry, no deals for rounds two and
* three, and no coin flips, so there is nothing honest to migrate it into. An
* unreadable save just deals a fresh game, which is the right failure.
*/
export const SAVED_GAME_KEY = "lost-cities-jax-ppo.game.v2";
export interface SavedFrame {
state: MatchState;
selection: { handSlot: number | null; placeType: PlaceType | null };
}
export interface SavedGame {
version: 2;
seed: string;
/** 1 for a one-off deal, 3 for the classic match. */
mode: number;
frames: SavedFrame[];
cursor: number;
resultOpen: boolean;
}
function isNumberArray(value: unknown, length?: number): value is number[] {
return Array.isArray(value) && (length === undefined || value.length === length) &&
value.every((item) => typeof item === "number" && Number.isFinite(item));
}
function isBooleanArray(value: unknown, length: number): value is boolean[] {
return Array.isArray(value) && value.length === length &&
value.every((item) => typeof item === "boolean");
}
function isMatrix(value: unknown, rows: number, columns: number): value is number[][] {
return Array.isArray(value) && value.length === rows &&
value.every((row) => isNumberArray(row, columns));
}
function isDeckOrder(value: unknown): value is number[] {
return isNumberArray(value, N_CARDS) && new Set(value).size === N_CARDS;
}
function isGameState(value: unknown): value is GameState {
if (typeof value !== "object" || value === null) return false;
const state = value as Partial<GameState>;
return isDeckOrder(state.deckOrder) &&
isNumberArray(state.cardLoc, N_CARDS) && isBooleanArray(state.handPublic, N_CARDS) &&
isMatrix(state.colTop, 2, 5) && isMatrix(state.colHandshakes, 2, 5) &&
isMatrix(state.colLength, 2, 5) && Array.isArray(state.piles) && state.piles.length === 5 &&
state.piles.every((pile) => isNumberArray(pile)) &&
typeof state.drawPtr === "number" && typeof state.stepCount === "number" &&
(state.toMove === 0 || state.toMove === 1) && typeof state.done === "boolean";
}
function isMatchState(value: unknown): value is MatchState {
if (typeof value !== "object" || value === null) return false;
const match = value as Partial<MatchState>;
return isGameState(match.round) &&
Array.isArray(match.deckOrders) && match.deckOrders.length === 3 &&
match.deckOrders.every(isDeckOrder) &&
isNumberArray(match.coinFlips, 3) &&
Number.isInteger(match.roundIdx) && match.roundIdx! >= 0 && match.roundIdx! < 3 &&
isNumberArray(match.carry, 2) &&
(match.totalRounds === 1 || match.totalRounds === 3) &&
typeof match.done === "boolean";
}
function isSavedGame(value: unknown): value is SavedGame {
if (typeof value !== "object" || value === null) return false;
const saved = value as Partial<SavedGame>;
if (saved.version !== 2 || typeof saved.seed !== "string" ||
(saved.mode !== 1 && saved.mode !== 3) ||
!Array.isArray(saved.frames) || saved.frames.length === 0 ||
!Number.isInteger(saved.cursor) || saved.cursor! < 0 || saved.cursor! >= saved.frames.length ||
typeof saved.resultOpen !== "boolean") return false;
return saved.frames.every((frame) => {
if (typeof frame !== "object" || frame === null || !isMatchState(frame.state)) return false;
const selection = frame.selection;
return typeof selection === "object" && selection !== null &&
(selection.handSlot === null || Number.isInteger(selection.handSlot)) &&
(selection.placeType === null || selection.placeType === 0 || selection.placeType === 1);
});
}
export function parseSavedGame(raw: string | null): SavedGame | null {
if (raw === null) return null;
try {
const value: unknown = JSON.parse(raw);
return isSavedGame(value) ? value : null;
} catch {
return null;
}
}
+15
View File
@@ -25,6 +25,21 @@ export function deckOrderFromSeed(seed: string): number[] {
return shuffledDeck(mulberry32(hashSeed(seed))); return shuffledDeck(mulberry32(hashSeed(seed)));
} }
/**
* Every deal and coin flip of a match, drawn up front from one seed.
*
* Drawing them all now — rather than shuffling again when a round rolls over —
* is what keeps a seed reproducible: the whole match is a pure function of it.
* The first deal is `deckOrderFromSeed`, so a match and a one-off deal on the
* same seed open on the same position.
*/
export function matchFromSeed(seed: string): { deckOrders: number[][]; coinFlips: number[] } {
const random = mulberry32(hashSeed(seed));
const deckOrders = [shuffledDeck(random), shuffledDeck(random), shuffledDeck(random)];
const coinFlips = [0, 1, 2].map(() => (random() < 0.5 ? 0 : 1));
return { deckOrders, coinFlips };
}
/** A fresh shareable seed, e.g. "k3f9qa". */ /** A fresh shareable seed, e.g. "k3f9qa". */
export function randomSeed(): string { export function randomSeed(): string {
return Math.floor(Math.random() * 36 ** 6).toString(36).padStart(6, "0"); return Math.floor(Math.random() * 36 ** 6).toString(36).padStart(6, "0");
+27 -14
View File
@@ -1,13 +1,25 @@
import * as ort from "onnxruntime-web/all"; import * as ort from "onnxruntime-web/all";
import { currentHandSorted, legalActionMask } from "../game/engine"; import { currentHandSorted } from "../game/engine";
import { observation } from "../game/observation"; import { matchLegalActionMask, type MatchState } from "../game/match";
import { DISCARD, DRAW_DECK, PLAY, type GameState } from "../game/types"; import { matchObservation } from "../game/matchObservation";
import { DISCARD, DRAW_DECK, PLAY } from "../game/types";
import { cardColor, cardRank, isHandshake } from "../game/cards"; import { cardColor, cardRank, isHandshake } from "../game/cards";
export type ExecutionProvider = "webgpu" | "wasm" | "heuristic"; export type ExecutionProvider = "webgpu" | "wasm" | "heuristic";
const MODEL_URL = `${import.meta.env.BASE_URL}models/jax-ppo.onnx`; /**
* borealis -- see data/models.json. Trained on the three-round match, so it takes
* the match view (carry, round, deck clock) rather than a bare round. A one-off
* deal is simply its round one at a carry of zero, which is a position it has seen
* a great many times.
*/
const MODEL_URL = `${import.meta.env.BASE_URL}models/borealis.onnx`;
/** Identity of what actually plays. Records carry the hash; the codename is for
* humans and is assigned in data/models.json, not derived. */
export const MODEL_CODENAME = "borealis";
export const MODEL_HASH = "13a25243de1c";
/** A legal action with the policy's confidence in it, if the policy has one. */ /** A legal action with the policy's confidence in it, if the policy has one. */
export interface RankedAction { export interface RankedAction {
@@ -18,16 +30,16 @@ export interface RankedAction {
export interface Policy { export interface Policy {
readonly provider: ExecutionProvider; readonly provider: ExecutionProvider;
/** Legal actions, best first. Drives both the rival's move and the hint. */ /** Legal actions, best first. Drives both the rival's move and the hint. */
rank(state: GameState): Promise<RankedAction[]>; rank(match: MatchState): Promise<RankedAction[]>;
action(state: GameState): Promise<number>; action(match: MatchState): Promise<number>;
} }
abstract class RankingPolicy implements Policy { abstract class RankingPolicy implements Policy {
abstract readonly provider: ExecutionProvider; abstract readonly provider: ExecutionProvider;
abstract rank(state: GameState): Promise<RankedAction[]>; abstract rank(match: MatchState): Promise<RankedAction[]>;
async action(state: GameState): Promise<number> { async action(match: MatchState): Promise<number> {
const [best] = await this.rank(state); const [best] = await this.rank(match);
if (best === undefined) throw new Error("state has no legal actions"); if (best === undefined) throw new Error("state has no legal actions");
return best.action; return best.action;
} }
@@ -54,12 +66,12 @@ class OnnxPolicy extends RankingPolicy {
super(); super();
} }
async rank(state: GameState): Promise<RankedAction[]> { async rank(match: MatchState): Promise<RankedAction[]> {
const obs = observation(state, state.toMove); const obs = matchObservation(match, match.round.toMove);
const result = await this.session.run({ obs: new ort.Tensor("float32", obs, [1, obs.length]) }); const result = await this.session.run({ obs: new ort.Tensor("float32", obs, [1, obs.length]) });
const logits = result.logits?.data; const logits = result.logits?.data;
if (!logits) throw new Error("ONNX model did not return a logits output"); if (!logits) throw new Error("ONNX model did not return a logits output");
const legal = legalActionMask(state); const legal = matchLegalActionMask(match);
return softmaxOverLegal(legal.map((_, action) => Number(logits[action])), legal); return softmaxOverLegal(legal.map((_, action) => Number(logits[action])), legal);
} }
} }
@@ -67,8 +79,9 @@ class OnnxPolicy extends RankingPolicy {
export class HeuristicPolicy extends RankingPolicy { export class HeuristicPolicy extends RankingPolicy {
readonly provider = "heuristic" as const; readonly provider = "heuristic" as const;
async rank(state: GameState): Promise<RankedAction[]> { async rank(match: MatchState): Promise<RankedAction[]> {
const legal = legalActionMask(state); const state = match.round;
const legal = matchLegalActionMask(match);
const hand = currentHandSorted(state); const hand = currentHandSorted(state);
const ranked = legal.flatMap((isLegal, action) => { const ranked = legal.flatMap((isLegal, action) => {
if (!isLegal) return []; if (!isLegal) return [];
+93 -15
View File
@@ -472,7 +472,7 @@ button.card:disabled { cursor: default; }
.deck-stack > span { margin-top: 5px; color: #7e817c; font: 700 9px Arial, sans-serif; letter-spacing: 0.2em; } .deck-stack > span { margin-top: 5px; color: #7e817c; font: 700 9px Arial, sans-serif; letter-spacing: 0.2em; }
.deck-stack.is-draw-target .card { box-shadow: 0 0 0 3px var(--gold), 7px 8px 0 #171b1e, 0 0 32px rgba(201, 163, 75, 0.3); } .deck-stack.is-draw-target .card { box-shadow: 0 0 0 3px var(--gold), 7px 8px 0 #171b1e, 0 0 32px rgba(201, 163, 75, 0.3); }
/* Insets keep the row clear of the history bar parked at the right. */ /* Insets keep the prompt clear of the controls parked at the right. */
.prompt-row { .prompt-row {
position: absolute; position: absolute;
right: 330px; right: 330px;
@@ -498,7 +498,6 @@ button.card:disabled { cursor: default; }
.turn-prompt--review { color: #9fb4c8; } .turn-prompt--review { color: #9fb4c8; }
.hint-button { .hint-button {
flex: 0 0 auto;
height: 32px; height: 32px;
padding: 0 14px; padding: 0 14px;
border: 1px solid var(--gold-dim); border: 1px solid var(--gold-dim);
@@ -610,23 +609,25 @@ button.card:disabled { cursor: default; }
.result-card button:hover { border-color: var(--gold); } .result-card button:hover { border-color: var(--gold); }
.result-card button.is-primary { border-color: var(--gold); color: var(--gold); } .result-card button.is-primary { border-color: var(--gold); color: var(--gold); }
/* Sits in the band above the hand: the hand row is centered and its width is /* Player controls share the hand row's bottom edge and stay in the mirrored
computed against the score plaques, so a wide bar down at the hand's level right-side keep-out reserved by the centered hand layout. */
would overlap the leftmost/rightmost cards. */ .control-stack {
.history-bar {
position: absolute; position: absolute;
right: 24px; right: 24px;
bottom: 196px; bottom: 30px;
z-index: 12; z-index: 12;
display: flex; display: flex;
align-items: center; width: 142px;
flex-direction: column;
align-items: stretch;
gap: 8px; gap: 8px;
} }
.history-bar button { .control-stack button {
height: 32px; height: 32px;
padding: 0 11px; padding: 0 11px;
display: flex; display: flex;
align-items: center; align-items: center;
justify-content: center;
gap: 6px; gap: 6px;
border: 1px solid var(--gold-dim); border: 1px solid var(--gold-dim);
border-radius: 9px; border-radius: 9px;
@@ -634,15 +635,17 @@ button.card:disabled { cursor: default; }
font: 700 13px/1 Georgia, serif; font: 700 13px/1 Georgia, serif;
cursor: pointer; cursor: pointer;
} }
.history-bar button span { font: 700 9px Arial, sans-serif; letter-spacing: 0.14em; } .control-stack button span { font: 700 9px Arial, sans-serif; letter-spacing: 0.14em; }
.history-bar button:hover:not(:disabled) { border-color: var(--gold); } .control-stack button:hover:not(:disabled) { border-color: var(--gold); }
.history-bar button:disabled { opacity: 0.3; cursor: default; } .control-stack button:disabled { opacity: 0.3; cursor: default; }
.history-bar__resume { .control-stack__resume {
border-color: var(--gold) !important; border-color: var(--gold) !important;
color: var(--gold); color: var(--gold);
font: 700 9px Arial, sans-serif; font: 700 9px Arial, sans-serif;
letter-spacing: 0.14em; letter-spacing: 0.14em;
} }
.control-stack__score { visibility: hidden; }
.control-stack__score.is-available { visibility: visible; }
@media (max-width: 1250px) { @media (max-width: 1250px) {
.table-center { width: 600px; } .table-center { width: 600px; }
@@ -651,7 +654,6 @@ button.card:disabled { cursor: default; }
.human-hand .card { width: 98px; height: 142px; margin-left: 7px; } .human-hand .card { width: 98px; height: 142px; margin-left: 7px; }
.human-hand { height: 146px; } .human-hand { height: 146px; }
.prompt-row { bottom: 177px; } .prompt-row { bottom: 177px; }
.history-bar { bottom: 172px; }
.table-center { bottom: 210px; } .table-center { bottom: 210px; }
/* Narrow viewports: the plaques are the keep-out that squeezes the hand row, /* Narrow viewports: the plaques are the keep-out that squeezes the hand row,
so they become compact chips (model-name line dropped) to give the cards so they become compact chips (model-name line dropped) to give the cards
@@ -675,7 +677,83 @@ button.card:disabled { cursor: default; }
.human-hand { bottom: 18px; height: 138px; } .human-hand { bottom: 18px; height: 138px; }
.human-hand .card { width: 94px; height: 134px; } .human-hand .card { width: 94px; height: 134px; }
.prompt-row { bottom: 165px; } .prompt-row { bottom: 165px; }
.history-bar { bottom: 160px; } .control-stack { bottom: 18px; }
.score-plaque--human { bottom: 9px; } .score-plaque--human { bottom: 9px; }
.deck-stack { transform: scale(0.86); } .deck-stack { transform: scale(0.86); }
} }
/* Match mode: which round is live, and what is already banked. The board only
ever shows the round in play, so without this the score plaques would be the
only hint that two more rounds are coming. */
.round-strip {
position: absolute;
top: 16px;
left: 50%;
transform: translateX(-50%);
display: flex;
align-items: center;
gap: 8px;
padding: 6px 12px;
border-radius: 999px;
background: rgba(18, 14, 10, 0.72);
border: 1px solid rgba(214, 188, 140, 0.22);
font-size: 11px;
letter-spacing: 0.12em;
color: rgba(214, 188, 140, 0.55);
z-index: 4;
pointer-events: none;
}
.round-strip span {
font-weight: 700;
}
.round-strip span.is-done {
color: rgba(214, 188, 140, 0.85);
}
.round-strip span.is-live {
color: #f0d9a6;
text-shadow: 0 0 10px rgba(240, 217, 166, 0.45);
}
.round-strip em {
font-style: normal;
margin-left: 4px;
padding-left: 10px;
border-left: 1px solid rgba(214, 188, 140, 0.22);
color: rgba(214, 188, 140, 0.75);
}
.menu-modes {
display: flex;
gap: 6px;
}
.menu-modes button {
flex: 1;
font-size: 10px;
}
.menu-modes button.is-active {
background: rgba(240, 217, 166, 0.16);
color: #f0d9a6;
border-color: rgba(240, 217, 166, 0.45);
}
/* A match's headline is the summed total; the table still breaks down the round
in front of you, so the two have to be told apart. */
.result-card__banked {
display: block;
margin-top: 6px;
font-size: 11px;
letter-spacing: 0.1em;
color: rgba(214, 188, 140, 0.55);
}
.result-table__match td,
.result-table__match th {
border-top: 1px solid rgba(214, 188, 140, 0.22);
color: #f0d9a6;
font-weight: 700;
}
+12 -1
View File
@@ -153,11 +153,22 @@ export function useCardMotion(): CardMotion {
continue; continue;
} }
// getBoundingClientRect includes a transform applied by WAAPI. A render
// while a deal/move is still in flight (for example when the policy
// finishes loading) must retain the intended destination instead of
// treating the animated visual position as a new layout and launching a
// second flight.
const previous = rects.current.get(card);
if (element.getAnimations().some((animation) => animation.playState === "running")) {
if (previous) next.set(card, previous);
continue;
}
const to = element.getBoundingClientRect(); const to = element.getBoundingClientRect();
next.set(card, to); next.set(card, to);
if (reduced) continue; if (reduced) continue;
const from = rects.current.get(card); const from = previous;
if (from) { if (from) {
animateMove(element, from, to); animateMove(element, from, to);
} else if (deckRect) { } else if (deckRect) {