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.