From 768f74693d86692bf7294320fdc02391a07e77da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= Date: Sat, 4 Jul 2026 22:14:37 +0900 Subject: [PATCH] Plan JAX PPO static-opponent ladder --- docs/plans/jax-ppo-static-opponents.md | 340 +++++++++++++++++++++++++ 1 file changed, 340 insertions(+) create mode 100644 docs/plans/jax-ppo-static-opponents.md diff --git a/docs/plans/jax-ppo-static-opponents.md b/docs/plans/jax-ppo-static-opponents.md new file mode 100644 index 0000000..b23f856 --- /dev/null +++ b/docs/plans/jax-ppo-static-opponents.md @@ -0,0 +1,340 @@ +# Plan: JAX PPO Static-Opponent Ladder + +**Status:** Ready to implement. +**Owner:** Codex implements; operator reviews training gates. +**Scope:** A compact PPO training stack on top of `lost_cities_jax`, using GPU +via optional CUDA JAX execution. + +## Goal + +Build the first learning layer above the JAX Lost Cities rules engine: + +1. A batched rollout driver using `batched_reset`, `batched_step`, + `batched_legal_mask`, and `batched_obs`. +2. Three static pure-JAX opponent policies: + `discard_only`, `heuristic_balanced`, and `heuristic_cautious`. +3. A 3 x 512 MLP actor-critic with 96-action policy logits and scalar value. +4. Standard PPO training against one static opponent at a time. +5. Duplicate evaluation on a fixed 10,000-deck shuffle bank. +6. Orbax checkpoints, metrics logging, and a CLI: + `lost-cities-jax-ppo train --config ...`. + +The near-term objective is not league self-play. It is to pass the +static-opponent diagnostic ladder and preserve enough artifacts that a later +self-play failure can be localized cleanly. + +## Non-Goals + +- No league self-play, snapshot pool, Elo, or opponent matchmaking in this + phase. +- No MCTS, CFR, search, or neural opponent ensemble. +- No multi-round match wrapper. +- No rule variants, expanded colors, or extra players. +- No large binary checkpoints committed to git. Generated artifacts are stored + outside the code repo, with small manifests and summaries committed. + +## GPU Execution + +Keep the project dependency portable. Do not make CUDA wheels mandatory in +`pyproject.toml`. + +Use optional CUDA JAX for training and GPU benchmarks: + +```bash +flock -n .compute.lock uv run --with 'jax[cuda12]' lost-cities-jax-ppo train \ + --config configs/jax_ppo/discard-only.yaml +``` + +The CPU path must still work for tests: + +```bash +uv run pytest -q tests/lost_cities_jax +``` + +Current benchmark evidence on RTX 3090: + +```text +backend=gpu +batch_size=8192 +steps=256 +steps_per_sec=3956598.38 +``` + +## Implementation Shape + +Prefer a small number of files while keeping testable boundaries: + +```text +src/lost_cities_jax/ + ppo.py # config, network, rollout, PPO update, train loop + opponents.py # pure-JAX static opponent policy functions + eval.py # duplicate evaluation and Wilson intervals + ppo_cli.py # argparse CLI +configs/jax_ppo/ + discard-only.yaml + balanced.yaml + cautious.yaml +``` + +If the PPO implementation remains readable in one main file, keep it there. +Split only when a file becomes hard to test or review. + +Dependencies likely needed: + +- `flax` for model modules and train state. +- `optax` for Adam and PPO losses. +- `orbax-checkpoint` for checkpointing. + +Add them through `uv add`, not pip/conda/poetry. + +## Static Opponents + +All opponent policies are pure JAX functions: + +```python +policy_fn(state: State, player: jax.Array) -> jax.Array # int32 action +``` + +They must sample no Python-side randomness. If tie-breaking needs randomness, +pass a JAX key explicitly: + +```python +policy_fn(state, player, rng) -> action +``` + +Policies: + +- `discard_only`: always discard a legal hand slot and draw from deck when + legal. This opponent should make free expedition building easy. +- `heuristic_balanced`: prefer legal plays that improve expedition prospects, + avoid obviously toxic openings, draw useful discard tops when available. +- `heuristic_cautious`: stricter opening threshold, fewer negative expedition + commitments, more conservative discard/draw behavior. + +Opponent policies are not learning targets. They are diagnostic fixtures. + +## Rollout Driver + +Run 8192 games in parallel by default. + +The learner controls one fixed seat per rollout batch. The opponent occupies +the other seat. Because games alternate turns, each environment step chooses: + +- learner action from the actor policy when `state.to_move == learner_seat`; +- static opponent action otherwise. + +Done states remain no-op through the engine, so the rollout can keep a fixed +time axis. Use `MAX_STEPS == 400` as the scan length unless a shorter config +value is explicitly introduced. + +The first milestone is random-policy rollout with the full dashboard. This is +not throwaway; it defines the baseline canaries. + +Log canaries: + +- episode return; +- score difference; +- game length distribution; +- `max_steps` termination rate; +- `play_action_rate`; +- opened color count; +- positive expedition count per game; +- average entropy under legal-action masking. + +Record the random-policy baseline in `README.md` before PPO training starts. + +## PPO Details + +Network: + +- Input: `OBS_DIM` observation vector. +- Body: MLP 512 -> 512 -> 512, ReLU. +- Policy head: 96 logits. +- Value head: scalar. +- Illegal actions are masked to a large negative value before sampling and + before log-prob/loss computation. + +Training defaults: + +```yaml +ppo: + batch_games: 8192 + rollout_steps: 400 + gamma: 1.0 + gae_lambda: 0.95 + clip_epsilon: 0.2 + entropy_coef: 0.01 + value_coef: 0.5 + max_grad_norm: 0.5 + learning_rate: 0.0003 + epochs: 4 + minibatches: 8 +reward: + terminal_scale: 50.0 + potential_shaping_initial: 1.0 + potential_shaping_final: 0.0 + potential_shaping_anneal_steps: 5_000_000 +``` + +Terminal reward: + +```text +tanh((learner_score - opponent_score) / terminal_scale) +``` + +Potential shaping: + +```text +coef(t) * (board_score_diff_after - board_score_diff_before) +``` + +The shaping coefficient anneals linearly from `initial` to `final`. Expose all +four shaping fields in config. This shaping exists to prevent the previous +play-action-rate collapse by giving immediate credit for board progress; its +schedule is a controlled experiment variable, not a hidden constant. + +## Evaluation Gates + +Evaluation uses a fixed shuffle bank: + +- Generate 10,000 explicit `deck_order` permutations from a fixed seed. +- For each deck, play twice: + - learner as player 0, opponent as player 1; + - opponent as player 0, learner as player 1. +- Aggregate duplicate-pair results. + +Report: + +- win rate; +- Wilson confidence interval; +- mean score difference; +- mean game length; +- positive expedition count per game; +- opened colors; +- `play_action_rate`. + +Gates: + +| Gate | Opponent | Pass condition | Diagnosis if failed | +| --- | --- | --- | --- | +| 1 | `discard_only` | win rate >= 90% and >= 2 positive expeditions/game | Audit reward pipeline, observations, and masks. Self-play is irrelevant. | +| 2 | `heuristic_balanced` | mean score difference > 0 | Learner works; tune shaping schedule and batch/variance. | +| 3 | `heuristic_cautious` | mean score difference > 0 | Same as gate 2; static ladder still not cleared. | + +Only after all three gates pass should league self-play begin. + +## Artifact Policy + +Generated training output stays out of git. Use: + +```text +/mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents/ +``` + +Per gate, preserve: + +- checkpoint directory; +- resolved config; +- duplicate evaluation JSON; +- metrics JSONL; +- short summary markdown with command, git commit, seed, and pass/fail result. + +Commit only small, durable metadata to the code repo: + +- config files; +- evaluator/training code; +- README baseline and gate summaries; +- optional dated summary under docs/reports. + +If binary artifact tracking is later required, add DVC/Git LFS explicitly +instead of committing large checkpoint files directly. + +## Work Order + +### Phase 1: Rollout Dashboard + +1. Add static opponent policies. +2. Add random-policy batched rollout. +3. Log canary metrics: + returns, game length, max-step rate, `play_action_rate`, opened colors, + positive expeditions. +4. Run CPU tests and one GPU random rollout. +5. Record the random baseline in README. + +Exit criteria: + +- deterministic smoke rollout passes; +- canary metrics look finite and stable; +- `play_action_rate` and game length are logged before any PPO code is judged. + +### Phase 2: PPO Against `discard_only` + +1. Add MLP actor-critic and masked action sampling. +2. Add GAE and PPO update. +3. Add Orbax checkpoints and resume. +4. Add CLI and the discard-only config under configs/jax_ppo. +5. Train on GPU. + +First-run checks: + +- with shaping enabled, `play_action_rate` rises above random baseline within + the first few hundred thousand environment steps; +- average game length converges near natural 44-ply games; +- `max_steps` termination rate goes to zero. + +Then anneal shaping and run duplicate evaluation for gate 1. + +### Phase 3: Static Ladder + +1. Reuse the same config and checkpoint flow. +2. Change only the opponent config for `heuristic_balanced`. +3. Train/evaluate until gate 2 passes or diagnostics point to variance. +4. Repeat for `heuristic_cautious`. + +No self-play work starts before gate 3 passes. + +## Commands + +Smoke rollout: + +```bash +uv run lost-cities-jax-ppo rollout-smoke --config configs/jax_ppo/discard-only.yaml +``` + +GPU training in tmux: + +```bash +tmux new-session -s coolrl-jax-ppo-discard \ + -c /home/coolguy/dev/coolrl-lost-cities \ + "flock -n .compute.lock uv run --with 'jax[cuda12]' lost-cities-jax-ppo train \ + --config configs/jax_ppo/discard-only.yaml" +``` + +Duplicate evaluation: + +```bash +flock -n .compute.lock uv run --with 'jax[cuda12]' lost-cities-jax-ppo eval \ + --checkpoint /mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents//latest \ + --opponent discard_only \ + --shuffle-bank-seed 20260704 \ + --games 10000 \ + --duplicate \ + --output /mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents//eval_duplicate.json +``` + +## Risks + +- The engine is fast on GPU, but Python logging/evaluation can dominate if + metrics are copied every step. Aggregate in JAX and transfer per rollout. +- Static heuristics can accidentally become too weak or too strong. Keep them + deterministic and versioned by config. +- Potential shaping can teach score-chasing artifacts if it never anneals. + Treat the coefficient schedule as part of the experiment identity. +- Checkpoint volume can grow quickly. Store large artifacts under `/mnt/2tbhdd` + and keep only `latest` plus gate checkpoints unless a run is explicitly + archived. + +## Current Next Action + +Implement Phase 1: static opponents plus random-policy batched rollout with the +full dashboard, then record the random baseline before adding PPO updates.