Plan JAX PPO static-opponent ladder

2026-07-04 22:14:37 +09:00
parent b37b841eef
commit 768f74693d
1 changed files with 340 additions and 0 deletions
@@ -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/<run>/latest \
  --opponent discard_only \
  --shuffle-bank-seed 20260704 \
  --games 10000 \
  --duplicate \
  --output /mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents/<run>/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.