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:
Claude Opus 4.8
+7
-1
@@ -15,6 +15,9 @@ src/**/*.c
|
||||
# Local session / lock files
|
||||
.compute.lock
|
||||
.claude/
|
||||
# ...except checked-in project skills, which are documentation for the next session.
|
||||
!web/.claude/
|
||||
!web/.claude/**
|
||||
|
||||
# Rust build output
|
||||
target/
|
||||
@@ -29,10 +32,13 @@ tools/julia/
|
||||
.ruff_cache
|
||||
|
||||
# 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/*.tsbuildinfo
|
||||
web/public/models/*.onnx
|
||||
web/public/models/*.json
|
||||
!web/public/models/jax-ppo.onnx
|
||||
!web/public/models/jax-ppo.json
|
||||
!web/public/models/borealis.onnx
|
||||
!web/public/models/borealis.json
|
||||
|
||||
@@ -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
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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.576–0.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.599–0.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는 뚫리고 우리는 안 뚫린다."**
|
||||
- ❌ "우리 정책은 착취 불가능하다." — 하한선일 뿐이고, 우리 쪽 곡선은 아직 오르는 중이다.
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -36,11 +36,14 @@ def reset(rng: jax.Array) -> State:
|
||||
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.
|
||||
|
||||
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``.
|
||||
|
||||
``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)
|
||||
@@ -60,7 +63,7 @@ def reset_from_order(deck_order: jax.Array) -> State:
|
||||
col_len=jnp.zeros((2, N_COLORS), 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),
|
||||
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),
|
||||
step_count=jnp.asarray(0, dtype=jnp.int32),
|
||||
done=jnp.asarray(False),
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
@@ -77,6 +77,14 @@ class PPOHyperConfig:
|
||||
|
||||
@dataclass
|
||||
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
|
||||
potential_shaping_initial: float = 1.0
|
||||
potential_shaping_final: float = 0.0
|
||||
@@ -134,6 +142,23 @@ class Transition(NamedTuple):
|
||||
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):
|
||||
wins: jax.Array
|
||||
losses: jax.Array
|
||||
@@ -218,14 +243,20 @@ def train(
|
||||
train_iteration = make_train_iteration(cfg, opponent_policy)
|
||||
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):
|
||||
shaping_coef = shaping_coefficient(
|
||||
cfg, update * cfg.ppo.batch_games * cfg.ppo.rollout_steps
|
||||
)
|
||||
shaping_coef = shaping_coefficient(cfg, learner_actions_seen)
|
||||
iter_start = time.perf_counter()
|
||||
state, env_state, rng, metrics = train_iteration(state, env_state, rng, shaping_coef)
|
||||
jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
|
||||
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["elapsed_seconds"] = time.perf_counter() - start
|
||||
_append_jsonl(metrics_path, row)
|
||||
@@ -260,8 +291,10 @@ def make_train_iteration(cfg: JaxPPOConfig, opponent_policy):
|
||||
def train_iteration(
|
||||
state: TrainState, env_state: State, rng: jax.Array, shaping_coef: jax.Array
|
||||
):
|
||||
rng, rollout_key, update_key, reset_key = jax.random.split(rng, 4)
|
||||
env_state, transitions, rollout_metrics = rollout_fn(
|
||||
rng, rollout_key, update_key = jax.random.split(rng, 3)
|
||||
# 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
|
||||
)
|
||||
advantages, returns = compute_gae(
|
||||
@@ -270,10 +303,10 @@ def make_train_iteration(cfg: JaxPPOConfig, opponent_policy):
|
||||
transitions.done,
|
||||
cfg.ppo.gamma,
|
||||
cfg.ppo.gae_lambda,
|
||||
last_value,
|
||||
)
|
||||
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_metrics, **update_metrics}
|
||||
return state, env_state, rng, {**rollout_stats, **update_metrics}
|
||||
|
||||
return train_iteration
|
||||
|
||||
@@ -299,8 +332,10 @@ def make_league_train_iteration(
|
||||
rng: jax.Array,
|
||||
shaping_coef: jax.Array,
|
||||
):
|
||||
rng, rollout_key, update_key, reset_key = jax.random.split(rng, 4)
|
||||
env_state, transitions, rollout_metrics = rollout_fn(
|
||||
rng, rollout_key, update_key = jax.random.split(rng, 3)
|
||||
# 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
|
||||
)
|
||||
advantages, returns = compute_gae(
|
||||
@@ -309,17 +344,10 @@ def make_league_train_iteration(
|
||||
transitions.done,
|
||||
cfg.ppo.gamma,
|
||||
cfg.ppo.gae_lambda,
|
||||
last_value,
|
||||
)
|
||||
state, update_metrics = ppo_update(state, transitions, advantages, returns, update_key, cfg)
|
||||
env_state, assignments = reset_done_envs_with_assignments(
|
||||
env_state,
|
||||
assignments,
|
||||
reset_key,
|
||||
cfg.ppo.batch_games,
|
||||
opponent_probs,
|
||||
mirror_probability,
|
||||
)
|
||||
return state, env_state, assignments, rng, {**rollout_metrics, **update_metrics}
|
||||
return state, env_state, assignments, rng, {**rollout_stats, **update_metrics}
|
||||
|
||||
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)
|
||||
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
|
||||
def rollout_fn(state: TrainState, env_state: State, rng: jax.Array, shaping_coef: jax.Array):
|
||||
def body(carry, _):
|
||||
env, key = carry
|
||||
key, learner_key, opponent_key = jax.random.split(key, 3)
|
||||
env, episode_return, key = carry
|
||||
key, learner_key, opponent_key, reset_key = jax.random.split(key, 4)
|
||||
obs = jax.vmap(observation, in_axes=(0, None))(env, learner)
|
||||
legal = jax.vmap(legal_action_mask)(env)
|
||||
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)
|
||||
after_diff = batch_score_diff(next_env, learner)
|
||||
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 = reward + shaping_coef * (after_diff - before_diff)
|
||||
reward = jnp.where(active, reward, 0.0)
|
||||
@@ -369,19 +400,27 @@ def make_rollout_fn(cfg: JaxPPOConfig, opponent_policy):
|
||||
log_prob=log_prob,
|
||||
value=value,
|
||||
reward=reward,
|
||||
done=next_env.done,
|
||||
done=done,
|
||||
active=active,
|
||||
actor_mask=actor_mask,
|
||||
entropy=entropy,
|
||||
play_action=(place_type == PLAY) & actor_mask,
|
||||
)
|
||||
return (next_env, key), transition
|
||||
|
||||
(next_env, rng), transitions = jax.lax.scan(
|
||||
body, (env_state, rng), xs=None, length=cfg.ppo.rollout_steps
|
||||
episode_return = episode_return + reward
|
||||
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)
|
||||
return next_env, transitions, metrics
|
||||
final_obs = jax.vmap(observation, in_axes=(0, None))(next_env, learner)
|
||||
_, last_value = state.apply_fn(state.params, final_obs)
|
||||
metrics = episode_metrics(transitions, episodes)
|
||||
return next_env, transitions, last_value, metrics
|
||||
|
||||
return rollout_fn
|
||||
|
||||
@@ -405,8 +444,8 @@ def make_league_rollout_fn(
|
||||
shaping_coef: jax.Array,
|
||||
):
|
||||
def body(carry, _):
|
||||
env, key = carry
|
||||
key, learner_key, mirror_key, pool_key = jax.random.split(key, 4)
|
||||
env, assignments, episode_return, key = carry
|
||||
key, learner_key, mirror_key, pool_key, reset_key = jax.random.split(key, 5)
|
||||
learner = assignments.learner_seat.astype(jnp.int32)
|
||||
opponent = 1 - learner
|
||||
|
||||
@@ -449,7 +488,8 @@ def make_league_rollout_fn(
|
||||
next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
|
||||
after_diff = batch_score_diff_for_players(next_env, learner)
|
||||
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 = reward + shaping_coef * (after_diff - before_diff)
|
||||
reward = jnp.where(active, reward, 0.0)
|
||||
@@ -463,19 +503,37 @@ def make_league_rollout_fn(
|
||||
log_prob=log_prob,
|
||||
value=value,
|
||||
reward=reward,
|
||||
done=next_env.done,
|
||||
done=done,
|
||||
active=active,
|
||||
actor_mask=actor_mask,
|
||||
entropy=entropy,
|
||||
play_action=(place_type == PLAY) & actor_mask,
|
||||
)
|
||||
return (next_env, key), transition
|
||||
|
||||
(next_env, rng), transitions = jax.lax.scan(
|
||||
body, (env_state, rng), xs=None, length=cfg.ppo.rollout_steps
|
||||
episode_return = episode_return + reward
|
||||
# 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)
|
||||
return next_env, transitions, metrics
|
||||
final_learner = assignments.learner_seat.astype(jnp.int32)
|
||||
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
|
||||
|
||||
@@ -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))
|
||||
learner = jnp.asarray(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
|
||||
def rollout(env_state: State, rng: jax.Array):
|
||||
def body(carry, _):
|
||||
env, key = carry
|
||||
env, episode_return, key = carry
|
||||
key, learner_key, opponent_key = jax.random.split(key, 3)
|
||||
learner_keys = jax.random.split(learner_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)
|
||||
after_diff = batch_score_diff(next_env, learner)
|
||||
reward = jnp.where(active, after_diff - before_diff, 0.0)
|
||||
terminal = active & next_env.done
|
||||
actor_mask = active & learner_turn
|
||||
place_type = (actions % 12) // 6
|
||||
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),
|
||||
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(
|
||||
body, (env_state, rng), xs=None, length=cfg.ppo.rollout_steps
|
||||
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
|
||||
)
|
||||
return rng, rollout_metrics(transitions, next_env, learner)
|
||||
return rng, episode_metrics(transitions, episodes)
|
||||
|
||||
rng, metrics = rollout(env_state, rng)
|
||||
jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
|
||||
@@ -679,7 +742,18 @@ def compute_gae(
|
||||
dones: jax.Array,
|
||||
gamma: float,
|
||||
gae_lambda: float,
|
||||
last_value: jax.Array | None = None,
|
||||
) -> 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):
|
||||
next_value, next_advantage = carry
|
||||
reward, value, done = x
|
||||
@@ -688,7 +762,7 @@ def compute_gae(
|
||||
advantage = delta + gamma * gae_lambda * nonterminal * next_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 = advantages_rev[::-1]
|
||||
return advantages, advantages + values
|
||||
@@ -1066,56 +1140,60 @@ def make_policy_match_batch_fn(learner_policy, opponent_policy):
|
||||
return eval_batch
|
||||
|
||||
|
||||
def rollout_metrics(
|
||||
transitions: Transition, final_env: State, learner: jax.Array
|
||||
) -> dict[str, jax.Array]:
|
||||
episode_return = jnp.sum(transitions.reward, axis=0)
|
||||
def _episode_end(
|
||||
next_env: State,
|
||||
learners: jax.Array,
|
||||
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)
|
||||
active_count = jnp.sum(transitions.active)
|
||||
color_scores = jax.vmap(color_scores_for_player, in_axes=(0, None))(final_env, learner)
|
||||
opened = jnp.sum(final_env.col_len[:, learner, :] > 0, axis=-1)
|
||||
positive = jnp.sum(color_scores > 0, axis=-1)
|
||||
|
||||
returns = episodes.episode_return
|
||||
return_mean = jnp.sum(returns * done) / denom
|
||||
return_var = jnp.sum(((returns - return_mean) ** 2) * done) / denom
|
||||
|
||||
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),
|
||||
"return_mean": return_mean,
|
||||
"return_std": jnp.sqrt(jnp.maximum(return_var, 0.0)),
|
||||
"game_length_mean": jnp.sum(episodes.length.astype(jnp.float32) * done) / denom,
|
||||
"game_length_max": jnp.max(jnp.where(episodes.done, episodes.length, 0)),
|
||||
"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),
|
||||
"opened_colors_mean": jnp.mean(opened.astype(jnp.float32)),
|
||||
"positive_expeditions_mean": jnp.mean(positive.astype(jnp.float32)),
|
||||
"opened_colors_mean": jnp.sum(episodes.opened_colors.astype(jnp.float32) * done) / denom,
|
||||
"positive_expeditions_mean": (
|
||||
jnp.sum(episodes.positive_expeditions.astype(jnp.float32) * done) / denom
|
||||
),
|
||||
"entropy_mean": masked_mean(
|
||||
transitions.entropy, transitions.actor_mask.astype(jnp.float32)
|
||||
),
|
||||
"active_steps": active_count,
|
||||
"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,
|
||||
"active_steps": jnp.sum(transitions.active),
|
||||
"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)
|
||||
|
||||
|
||||
def shaping_coefficient(cfg: JaxPPOConfig, env_steps: int) -> float:
|
||||
def shaping_coefficient(cfg: JaxPPOConfig, learner_actions: int) -> float:
|
||||
reward_cfg = cfg.reward
|
||||
if reward_cfg.potential_shaping_anneal_steps <= 0:
|
||||
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 * (
|
||||
reward_cfg.potential_shaping_final - reward_cfg.potential_shaping_initial
|
||||
)
|
||||
|
||||
@@ -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])
|
||||
@@ -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)
|
||||
@@ -61,7 +61,9 @@ def tiny_config(tmp_path) -> JaxPPOConfig:
|
||||
),
|
||||
opponent=OpponentConfig(name="discard_only"),
|
||||
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):
|
||||
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 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):
|
||||
|
||||
@@ -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
|
||||
@@ -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
@@ -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.
|
||||
@@ -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
@@ -10,16 +10,35 @@ import {
|
||||
currentHandSorted,
|
||||
decodeAction,
|
||||
encodeAction,
|
||||
legalActionMask,
|
||||
previewPlacement,
|
||||
resetFromOrder,
|
||||
step,
|
||||
} from "./game/engine";
|
||||
import { deckOrderFromSeed, normalizeSeed, randomSeed } from "./game/random";
|
||||
import { DRAW_DECK, N_CARDS, PLAY, type GameState, type PlaceType } from "./game/types";
|
||||
import {
|
||||
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 { 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 {
|
||||
action: number;
|
||||
text: string;
|
||||
@@ -36,8 +55,9 @@ const EMPTY_SELECTION: Selection = { handSlot: null, placeType: null };
|
||||
* makes a finished game reviewable ply by ply.
|
||||
*/
|
||||
interface Frame {
|
||||
state: GameState;
|
||||
state: MatchState;
|
||||
selection: Selection;
|
||||
move?: { player: 0 | 1; action: number };
|
||||
}
|
||||
|
||||
// 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));
|
||||
}
|
||||
|
||||
/** A `?seed=` in the URL loads that exact deal, so a game can be shared or replayed. */
|
||||
function initialSeed(): string {
|
||||
const fromUrl = new URLSearchParams(window.location.search).get("seed");
|
||||
const seed = fromUrl === null ? "" : normalizeSeed(fromUrl);
|
||||
return seed === "" ? randomSeed() : seed;
|
||||
}
|
||||
|
||||
function openingFrame(seed: string): Frame {
|
||||
return { state: resetFromOrder(deckOrderFromSeed(seed)), selection: EMPTY_SELECTION };
|
||||
/** A `?seed=` in the URL loads that exact deal, so a game can be shared or replayed.
|
||||
* In match mode the seed fixes all three deals and the coin flips too, so the
|
||||
* whole match is a pure function of it. */
|
||||
function openingFrame(seed: string, mode: Mode): Frame {
|
||||
const { deckOrders, coinFlips } = matchFromSeed(seed);
|
||||
return { state: matchFromOrders(deckOrders, coinFlips, mode), selection: EMPTY_SELECTION };
|
||||
}
|
||||
|
||||
/** 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);
|
||||
}
|
||||
|
||||
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() {
|
||||
const [seed, setSeed] = useState<string>(initialSeed);
|
||||
const [frames, setFrames] = useState<Frame[]>(() => [openingFrame(seed)]);
|
||||
const [cursor, setCursor] = useState(0);
|
||||
const [initial] = useState(initialGame);
|
||||
const [seed, setSeed] = useState(initial.seed);
|
||||
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 [resultOpen, setResultOpen] = useState(true);
|
||||
const [resultOpen, setResultOpen] = useState(initial.resultOpen);
|
||||
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 [menuOpen, setMenuOpen] = useState(false);
|
||||
const [viewportWidth, setViewportWidth] = useState(() => window.innerWidth);
|
||||
const [hint, setHint] = useState<Hint | null>(null);
|
||||
const [hintPending, setHintPending] = useState(false);
|
||||
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 { 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 canRedo = cursor < frames.length - 1;
|
||||
const latestState = useRef(state);
|
||||
const latestMatch = useRef(match);
|
||||
|
||||
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(() => {
|
||||
function onResize() { setViewportWidth(window.innerWidth); }
|
||||
@@ -116,7 +170,7 @@ function App() {
|
||||
|
||||
const humanHand = useMemo(() => currentHandSorted(state, 0), [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 selectedColor = selectedCard === null ? null : cardColor(selectedCard);
|
||||
const displayState = useMemo(
|
||||
@@ -125,14 +179,23 @@ function App() {
|
||||
: previewPlacement(state, selection.handSlot, selection.placeType),
|
||||
[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(() => {
|
||||
let cancelled = false;
|
||||
loadPolicy().then(({ policy: loaded, warning }) => {
|
||||
if (cancelled) return;
|
||||
setPolicy(loaded);
|
||||
setModelMessage(warning ?? `${loaded.provider.toUpperCase()} · FINAL PPO`);
|
||||
setModelMessage(warning ?? `${loaded.provider.toUpperCase()} · BOREALIS`);
|
||||
});
|
||||
return () => { cancelled = true; };
|
||||
}, []);
|
||||
@@ -149,24 +212,28 @@ function App() {
|
||||
const rivalSuspended = canRedo;
|
||||
|
||||
useEffect(() => {
|
||||
if (!policy || rivalSuspended || state.done || state.toMove !== 1) return;
|
||||
if (!policy || rivalSuspended || match.done || state.toMove !== 1) return;
|
||||
const currentGeneration = generation.current;
|
||||
const timer = window.setTimeout(async () => {
|
||||
setThinking(true);
|
||||
try {
|
||||
let action: number;
|
||||
try {
|
||||
action = await policy.action(state);
|
||||
action = await policy.action(match);
|
||||
} catch (error) {
|
||||
console.error("AI action failed; falling back to heuristic policy", error);
|
||||
const fallback = fallbackHeuristicPolicy();
|
||||
action = await fallback.action(state);
|
||||
action = await fallback.action(match);
|
||||
if (generation.current !== currentGeneration) return;
|
||||
setPolicy(fallback);
|
||||
setModelMessage("HEURISTIC FALLBACK (MODEL ERROR)");
|
||||
}
|
||||
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) {
|
||||
console.error("AI action failed even with heuristic fallback", error);
|
||||
setModelMessage("MODEL INFERENCE ERROR");
|
||||
@@ -175,15 +242,19 @@ function App() {
|
||||
}
|
||||
}, 620);
|
||||
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
|
||||
* 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;
|
||||
resetMotion();
|
||||
gameId.current = newGameId();
|
||||
startedAt.current = new Date().toISOString();
|
||||
loggedGame.current = null;
|
||||
setSeed(nextSeed);
|
||||
setFrames([openingFrame(nextSeed)]);
|
||||
setMode(nextMode);
|
||||
setFrames([openingFrame(nextSeed, nextMode)]);
|
||||
setCursor(0);
|
||||
setSeedDraft("");
|
||||
setResultOpen(true);
|
||||
@@ -217,8 +288,8 @@ function App() {
|
||||
|
||||
/** Ask the policy driving the rival what it would do in your seat. */
|
||||
async function requestHint() {
|
||||
if (!policy || state.done || state.toMove !== 0 || hintPending) return;
|
||||
const position = state;
|
||||
if (!policy || match.done || state.toMove !== 0 || hintPending) return;
|
||||
const position = match;
|
||||
setHintPending(true);
|
||||
setMenuOpen(false);
|
||||
try {
|
||||
@@ -229,10 +300,10 @@ function App() {
|
||||
const [best] = ranked;
|
||||
// The position can move on while inference runs — a hint for a stale
|
||||
// position would point at the wrong hand slot.
|
||||
if (best === undefined || latestState.current !== position) return;
|
||||
if (best === undefined || latestMatch.current !== position) return;
|
||||
setHint({
|
||||
action: best.action,
|
||||
text: describeAction(position, best.action),
|
||||
text: describeAction(position.round, best.action),
|
||||
probability: best.probability,
|
||||
});
|
||||
} catch (error) {
|
||||
@@ -246,21 +317,21 @@ function App() {
|
||||
useEffect(() => { setHint(null); }, [cursor, frames]);
|
||||
|
||||
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
|
||||
? { handSlot, placeType: null }
|
||||
: EMPTY_SELECTION;
|
||||
pushFrame({ state, selection: next });
|
||||
pushFrame({ state: match, selection: next });
|
||||
}
|
||||
|
||||
function choosePlace(placeType: PlaceType) {
|
||||
if (selection.handSlot === null) return;
|
||||
pushFrame({ state, selection: { ...selection, placeType } });
|
||||
pushFrame({ state: match, selection: { ...selection, placeType } });
|
||||
}
|
||||
|
||||
function cancelPlace() {
|
||||
if (selection.handSlot === null || selection.placeType === null) return;
|
||||
pushFrame({ state, selection: { ...selection, placeType: null } });
|
||||
pushFrame({ state: match, selection: { ...selection, placeType: null } });
|
||||
}
|
||||
|
||||
useEffect(() => {
|
||||
@@ -282,7 +353,11 @@ function App() {
|
||||
function commit(drawSource: number) {
|
||||
if (!legalDraw(drawSource) || selection.handSlot === null || selection.placeType === null) return;
|
||||
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
|
||||
@@ -309,18 +384,83 @@ function App() {
|
||||
const hintConfidence = hint !== null && hint.probability !== null
|
||||
? `${Math.round(hint.probability * 100)}%`
|
||||
: 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]
|
||||
? "DRAW"
|
||||
: 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
|
||||
// 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.
|
||||
const reviewingRival = rivalSuspended && (state.done || state.toMove === 1);
|
||||
const reviewingRival = rivalSuspended && (match.done || state.toMove === 1);
|
||||
const status = reviewingRival
|
||||
? `Reviewing ply ${state.stepCount} — redo, or play on from here`
|
||||
: state.done
|
||||
: match.done
|
||||
? outcome
|
||||
: state.toMove === 1
|
||||
? thinking ? "The rival is thinking ···" : "The rival's turn"
|
||||
@@ -350,6 +490,24 @@ function App() {
|
||||
<div className="menu-popover">
|
||||
<button onClick={() => restart()}>NEW GAME</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">
|
||||
<span>DEAL SEED · {seed}</span>
|
||||
<input
|
||||
@@ -367,6 +525,24 @@ function App() {
|
||||
)}
|
||||
</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">
|
||||
<Board
|
||||
state={displayState}
|
||||
@@ -398,14 +574,6 @@ function App() {
|
||||
<p className={`turn-prompt ${thinking ? "turn-prompt--thinking" : ""} ${reviewingRival ? "turn-prompt--review" : ""}`}>
|
||||
{hint ? `Hint — ${hint.text}${hintConfidence ? ` (${hintConfidence})` : ""}` : status}
|
||||
</p>
|
||||
<button
|
||||
type="button"
|
||||
className={`hint-button ${hint ? "is-active" : ""}`}
|
||||
onClick={requestHint}
|
||||
disabled={!canHint || hintPending}
|
||||
>
|
||||
{hintPending ? "THINKING…" : "HINT"}
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<section className="human-hand" aria-label="Your hand">
|
||||
@@ -415,7 +583,7 @@ function App() {
|
||||
key={card}
|
||||
selected={selection.handSlot === slot}
|
||||
hinted={hintedCard === card}
|
||||
disabled={state.toMove !== 0 || state.done}
|
||||
disabled={state.toMove !== 0 || match.done}
|
||||
onClick={() => chooseCard(slot)}
|
||||
style={position === 0 ? undefined : { marginLeft: handCardMarginLeft }}
|
||||
innerRef={cardRef(card)}
|
||||
@@ -424,24 +592,41 @@ function App() {
|
||||
</section>
|
||||
|
||||
<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>
|
||||
|
||||
<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 (←)">
|
||||
↶ <span>UNDO</span>
|
||||
</button>
|
||||
<button type="button" onClick={redo} disabled={!canRedo} aria-label="Redo one action" title="Redo (→)">
|
||||
↷ <span>REDO</span>
|
||||
</button>
|
||||
{rivalSuspended && (
|
||||
<button type="button" className="history-bar__resume" onClick={resumeFromHere}>
|
||||
PLAY FROM HERE
|
||||
</button>
|
||||
)}
|
||||
{state.done && !resultOpen && (
|
||||
<button type="button" onClick={() => setResultOpen(true)}>SCORE</button>
|
||||
)}
|
||||
<button
|
||||
type="button"
|
||||
className="control-stack__resume"
|
||||
onClick={resumeFromHere}
|
||||
disabled={!rivalSuspended}
|
||||
>
|
||||
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>
|
||||
|
||||
{/* Cards in flight back to the deck (undo of a draw) are re-parented here
|
||||
@@ -453,6 +638,9 @@ function App() {
|
||||
state={state}
|
||||
outcome={outcome}
|
||||
seed={seed}
|
||||
carry={match.carry}
|
||||
roundIdx={match.roundIdx}
|
||||
totalRounds={mode}
|
||||
onReview={() => setResultOpen(false)}
|
||||
onPlayAgain={() => restart()}
|
||||
/>
|
||||
|
||||
@@ -6,6 +6,15 @@ interface ResultCardProps {
|
||||
state: GameState;
|
||||
outcome: 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;
|
||||
onPlayAgain: () => void;
|
||||
}
|
||||
@@ -14,17 +23,35 @@ function signed(value: number): string {
|
||||
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 rival = scoreBreakdown(state, 1);
|
||||
const isMatch = totalRounds > 1;
|
||||
const matchTotals: [number, number] = [carry[0] + you.total, carry[1] + rival.total];
|
||||
|
||||
return (
|
||||
<div className="result-overlay">
|
||||
<section className="result-card" aria-label="Final score">
|
||||
<header className="result-card__head">
|
||||
<p>ROUND COMPLETE</p>
|
||||
<p>{isMatch ? `MATCH COMPLETE · ${totalRounds} ROUNDS` : "ROUND COMPLETE"}</p>
|
||||
<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>
|
||||
|
||||
<table className="result-table">
|
||||
@@ -69,12 +96,21 @@ export function ResultCard({ state, outcome, seed, onReview, onPlayAgain }: Resu
|
||||
</tbody>
|
||||
<tfoot>
|
||||
<tr>
|
||||
<th scope="row">TOTAL</th>
|
||||
<th scope="row">{isMatch ? `ROUND ${roundIdx + 1}` : "TOTAL"}</th>
|
||||
<td className="result-table__detail" />
|
||||
<td>{signed(you.total)}</td>
|
||||
<td className="result-table__detail" />
|
||||
<td>{signed(rival.total)}</td>
|
||||
</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>
|
||||
</table>
|
||||
|
||||
|
||||
@@ -32,7 +32,11 @@ export function shuffledDeck(random: () => number = Math.random): number[] {
|
||||
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) {
|
||||
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),
|
||||
colLength: emptyMatrix(2, N_COLORS),
|
||||
piles: Array.from({ length: N_COLORS }, () => []),
|
||||
toMove: 0,
|
||||
toMove: (firstPlayer === 1 ? 1 : 0) as 0 | 1,
|
||||
stepCount: 0,
|
||||
done: false,
|
||||
};
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -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,
|
||||
};
|
||||
}
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -25,6 +25,21 @@ export function deckOrderFromSeed(seed: string): number[] {
|
||||
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". */
|
||||
export function randomSeed(): string {
|
||||
return Math.floor(Math.random() * 36 ** 6).toString(36).padStart(6, "0");
|
||||
|
||||
+27
-14
@@ -1,13 +1,25 @@
|
||||
import * as ort from "onnxruntime-web/all";
|
||||
|
||||
import { currentHandSorted, legalActionMask } from "../game/engine";
|
||||
import { observation } from "../game/observation";
|
||||
import { DISCARD, DRAW_DECK, PLAY, type GameState } from "../game/types";
|
||||
import { currentHandSorted } from "../game/engine";
|
||||
import { matchLegalActionMask, type MatchState } from "../game/match";
|
||||
import { matchObservation } from "../game/matchObservation";
|
||||
import { DISCARD, DRAW_DECK, PLAY } from "../game/types";
|
||||
import { cardColor, cardRank, isHandshake } from "../game/cards";
|
||||
|
||||
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. */
|
||||
export interface RankedAction {
|
||||
@@ -18,16 +30,16 @@ export interface RankedAction {
|
||||
export interface Policy {
|
||||
readonly provider: ExecutionProvider;
|
||||
/** Legal actions, best first. Drives both the rival's move and the hint. */
|
||||
rank(state: GameState): Promise<RankedAction[]>;
|
||||
action(state: GameState): Promise<number>;
|
||||
rank(match: MatchState): Promise<RankedAction[]>;
|
||||
action(match: MatchState): Promise<number>;
|
||||
}
|
||||
|
||||
abstract class RankingPolicy implements Policy {
|
||||
abstract readonly provider: ExecutionProvider;
|
||||
abstract rank(state: GameState): Promise<RankedAction[]>;
|
||||
abstract rank(match: MatchState): Promise<RankedAction[]>;
|
||||
|
||||
async action(state: GameState): Promise<number> {
|
||||
const [best] = await this.rank(state);
|
||||
async action(match: MatchState): Promise<number> {
|
||||
const [best] = await this.rank(match);
|
||||
if (best === undefined) throw new Error("state has no legal actions");
|
||||
return best.action;
|
||||
}
|
||||
@@ -54,12 +66,12 @@ class OnnxPolicy extends RankingPolicy {
|
||||
super();
|
||||
}
|
||||
|
||||
async rank(state: GameState): Promise<RankedAction[]> {
|
||||
const obs = observation(state, state.toMove);
|
||||
async rank(match: MatchState): Promise<RankedAction[]> {
|
||||
const obs = matchObservation(match, match.round.toMove);
|
||||
const result = await this.session.run({ obs: new ort.Tensor("float32", obs, [1, obs.length]) });
|
||||
const logits = result.logits?.data;
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -67,8 +79,9 @@ class OnnxPolicy extends RankingPolicy {
|
||||
export class HeuristicPolicy extends RankingPolicy {
|
||||
readonly provider = "heuristic" as const;
|
||||
|
||||
async rank(state: GameState): Promise<RankedAction[]> {
|
||||
const legal = legalActionMask(state);
|
||||
async rank(match: MatchState): Promise<RankedAction[]> {
|
||||
const state = match.round;
|
||||
const legal = matchLegalActionMask(match);
|
||||
const hand = currentHandSorted(state);
|
||||
const ranked = legal.flatMap((isLegal, action) => {
|
||||
if (!isLegal) return [];
|
||||
|
||||
+93
-15
@@ -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.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 {
|
||||
position: absolute;
|
||||
right: 330px;
|
||||
@@ -498,7 +498,6 @@ button.card:disabled { cursor: default; }
|
||||
.turn-prompt--review { color: #9fb4c8; }
|
||||
|
||||
.hint-button {
|
||||
flex: 0 0 auto;
|
||||
height: 32px;
|
||||
padding: 0 14px;
|
||||
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.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
|
||||
computed against the score plaques, so a wide bar down at the hand's level
|
||||
would overlap the leftmost/rightmost cards. */
|
||||
.history-bar {
|
||||
/* Player controls share the hand row's bottom edge and stay in the mirrored
|
||||
right-side keep-out reserved by the centered hand layout. */
|
||||
.control-stack {
|
||||
position: absolute;
|
||||
right: 24px;
|
||||
bottom: 196px;
|
||||
bottom: 30px;
|
||||
z-index: 12;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
width: 142px;
|
||||
flex-direction: column;
|
||||
align-items: stretch;
|
||||
gap: 8px;
|
||||
}
|
||||
.history-bar button {
|
||||
.control-stack button {
|
||||
height: 32px;
|
||||
padding: 0 11px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
gap: 6px;
|
||||
border: 1px solid var(--gold-dim);
|
||||
border-radius: 9px;
|
||||
@@ -634,15 +635,17 @@ button.card:disabled { cursor: default; }
|
||||
font: 700 13px/1 Georgia, serif;
|
||||
cursor: pointer;
|
||||
}
|
||||
.history-bar button span { font: 700 9px Arial, sans-serif; letter-spacing: 0.14em; }
|
||||
.history-bar button:hover:not(:disabled) { border-color: var(--gold); }
|
||||
.history-bar button:disabled { opacity: 0.3; cursor: default; }
|
||||
.history-bar__resume {
|
||||
.control-stack button span { font: 700 9px Arial, sans-serif; letter-spacing: 0.14em; }
|
||||
.control-stack button:hover:not(:disabled) { border-color: var(--gold); }
|
||||
.control-stack button:disabled { opacity: 0.3; cursor: default; }
|
||||
.control-stack__resume {
|
||||
border-color: var(--gold) !important;
|
||||
color: var(--gold);
|
||||
font: 700 9px Arial, sans-serif;
|
||||
letter-spacing: 0.14em;
|
||||
}
|
||||
.control-stack__score { visibility: hidden; }
|
||||
.control-stack__score.is-available { visibility: visible; }
|
||||
|
||||
@media (max-width: 1250px) {
|
||||
.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 { height: 146px; }
|
||||
.prompt-row { bottom: 177px; }
|
||||
.history-bar { bottom: 172px; }
|
||||
.table-center { bottom: 210px; }
|
||||
/* 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
|
||||
@@ -675,7 +677,83 @@ button.card:disabled { cursor: default; }
|
||||
.human-hand { bottom: 18px; height: 138px; }
|
||||
.human-hand .card { width: 94px; height: 134px; }
|
||||
.prompt-row { bottom: 165px; }
|
||||
.history-bar { bottom: 160px; }
|
||||
.control-stack { bottom: 18px; }
|
||||
.score-plaque--human { bottom: 9px; }
|
||||
.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;
|
||||
}
|
||||
|
||||
@@ -153,11 +153,22 @@ export function useCardMotion(): CardMotion {
|
||||
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();
|
||||
next.set(card, to);
|
||||
if (reduced) continue;
|
||||
|
||||
const from = rects.current.get(card);
|
||||
const from = previous;
|
||||
if (from) {
|
||||
animateMove(element, from, to);
|
||||
} else if (deckRect) {
|
||||
|
||||
Reference in New Issue
Block a user