Refresh optimization plans

Archive implemented AMP, Option A inference-server, and Cython heuristic plans. Add the active Option B interleaved traversal plan and update model-size/torch.compile plans to reflect the current traversal scheduling conclusion.

Co-Authored-By: Codex <codex@openai.com>

docs/plans/archive/amp_trainer.md

@@ -1,6 +1,6 @@
 # Plan: Trainer-side AMP (Automatic Mixed Precision)
 
-**Status:** Ready for implementation
+**Status:** Archived. Implemented and benchmarked on 2026-05-07; default remains off after smoke-config AMP regression. See `docs/performance.md` "AMP on trainer networks (2026-05-07, regression)".
 **Owner:** Codex
 **Background:** See `docs/performance.md` → "AMP Status" (currently a no-op flag) and "Post-A Optimization Calculus" (AMP becomes more meaningful only at larger model sizes; today's gains are bounded by the small `DeepCFRMLP` and the dominant traversal phase). Also note the `torch.compile` regression experiment in the same doc — small models punish low-level kernel optimizations because dispatch overhead outweighs fused-kernel gains. AMP can hit the same wall.
 

docs/plans/archive/batched_traversal_inference_server.md

@@ -1,6 +1,6 @@
 # Plan: Batched Traversal Inference Server (Priority #5, Option A)
 
-**Status:** Ready for implementation
+**Status:** Archived. Implemented and benchmarked on 2026-05-07 as Option A; end-to-end traversal regressed because sync-blocking traversal could not feed large batches. Superseded by `docs/plans/option_b_interleaved_traversal.md`.
 **Owner:** Codex
 **Background:** See `docs/performance.md` → "Batched Traversal Inference: Design Decision (2026-05-07)" for the A/B/C analysis and rationale. This plan implements Option A.
 

docs/plans/archive/cython_safe_heuristic_bots.md

@@ -1,6 +1,6 @@
 # Plan: Cython Port of the Safe-Heuristic Bot Family
 
-**Status:** Ready for implementation
+**Status:** Archived. Implemented as `heuristic_cy.pyx` with Python shim and equivalence tests.
 **Owner:** Codex
 **Background:** See `docs/performance.md` → "Evaluation Breakdown" and "Post-A Optimization Calculus". The safe-heuristic family dominates eval wall-clock through `opponent_act_seconds` (7.57s / 4.63s / 9.22s for `safe_heuristic` / `safe_heuristic_loose` / `safe_heuristic_strict` in the inspected eval row), not GPU forward. As denser eval becomes operationally useful (`eval_every: 5`, `evaluation.games: 1000`), this cost becomes a first-order wall-clock concern.
 

docs/plans/model_size_experiment.md

@@ -1,8 +1,14 @@
-# Plan: Model-Size Experiment (Keystone for AMP / compile / TRT Unblock)
+# Plan: Model-Size Experiment (Keystone for Model-Scale Optimizations)
 
 **Status:** Ready to execute
 **Owner:** Operator (runs grid on `home`); Codex (adds configs and runner script)
-**Background:** See `docs/performance.md` → "Post-A Optimization Calculus" and "Option A Bench Result and Structural Ceiling" for the sequencing rationale. AMP, `torch.compile`, TensorRT, and re-enabling Option A are all gated on the outcome of this experiment.
+**Background:** See `docs/performance.md` → "Post-A Optimization Calculus",
+"Option A Bench Result and Structural Ceiling", and "Clarifying the traversal
+bottleneck: sync policy boundary, not SIMD" for the sequencing rationale. AMP,
+`torch.compile`, and TensorRT are gated on the outcome of this experiment.
+Traversal batching itself is now tracked separately in
+`docs/plans/option_b_interleaved_traversal.md`; model-size growth does not fix
+the current sync-blocking traversal scheduler by itself.
 
 ## Goal
 
@@ -11,12 +17,15 @@ Identify a `network` config (hidden_size × num_layers) where:
 1. Learning-curve win-rates improve meaningfully over the current `default.yaml` baseline (hidden=512, layers=3) by iteration 200, AND
 2. Per-call forward time is large enough that AMP / `torch.compile` / TensorRT overhead is amortized — the documented threshold is `hidden_size >= 1024` or `num_layers >= 6`.
 
-The experiment must produce either a recommended new `network` config or a documented null result. Its output determines which downstream optimization plans ship next.
+The experiment must produce either a recommended new `network` config or a documented null result. Its output determines which model-scale optimization plans ship next.
 
 ## Non-goals
 
 - Do NOT implement AMP, `torch.compile`, or TensorRT here. Those are separate plans gated on this experiment's outcome.
-- Do NOT re-enable Option A (`traversal.inference_backend: server`). That gate is also conditional on the model size chosen here.
+- Do NOT re-enable Option A (`traversal.inference_backend: server`). Option A
+  has already been implemented and benchmarked; it is structurally capped by
+  sync-blocking traversal. Revisit it only after Option B-style traversal
+  interleaving can feed larger batches.
 - Do NOT change traversal, replay buffer, or evaluation architecture.
 - Do NOT change eval cadence (`eval_every`, `evaluation.games`) — those are separate variables.
 - Do NOT change encoding (`input_dim` stays 365).
@@ -78,7 +87,6 @@ Confirm that `NetworkConfig.hidden_size` and `NetworkConfig.num_layers` flow thr
 uv run python -c "
 from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP
 from coolrl_lost_cities.games.classic.deep_cfr.config import NetworkConfig
-import sum as _
 cfg_small = NetworkConfig(hidden_size=512, num_layers=3)
 cfg_large = NetworkConfig(hidden_size=1024, num_layers=6)
 m_small = DeepCFRMLP.from_config(365, 22, cfg_small)
@@ -282,7 +290,10 @@ Grid: hidden={512,768,1024,1536} × layers={3,4,6,8} subset. 200 iterations on h
 
 **Recommendation:** <see decision tree below>
 
-Cross-reference: AMP plan `docs/plans/amp_trainer.md`, compile plan `docs/plans/torch_compile.md`.
+Cross-reference: archived AMP implementation plan
+`docs/plans/archive/amp_trainer.md`, compile plan
+`docs/plans/torch_compile.md`, and Option B traversal plan
+`docs/plans/option_b_interleaved_traversal.md`.
 ```
 
 ## Decision tree
@@ -296,10 +307,13 @@ Apply this logic after the grid completes:
 **Action:**
 1. Recommend that config as the new `network` default.
 2. Update `default.yaml` `network` block in a follow-up commit.
-3. Trigger AMP re-measurement: run `scripts/bench_amp_trainer.py` with the new model size per `docs/plans/amp_trainer.md`.
+3. Trigger AMP re-measurement: run `scripts/bench_amp_trainer.py` with the new model size. The implementation exists and is default-off after the 2026-05-07 smoke regression; see `docs/plans/archive/amp_trainer.md` for the original implementation plan.
 4. Trigger `torch.compile` re-measurement per `docs/plans/torch_compile.md`.
-5. Trigger TensorRT evaluation on the inference server.
-6. Note that Option A (`inference_backend: server`) should be re-benchmarked at the new model size — the larger forward time directly reduces the IPC/GPU-forward ratio that caused A's regression.
+5. Trigger TensorRT evaluation for batched eval/inference surfaces.
+6. Do not treat this as sufficient to re-enable Option A. Larger models improve
+   the IPC/GPU-forward ratio, but the observed traversal ceiling is still
+   scheduling shape. Re-benchmark server inference only after Option B or a
+   comparable traversal interleaving path can feed bs=64+ batches.
 
 ### Branch B — no size unlocks the curve
 

docs/plans/option_b_interleaved_traversal.md

@@ -0,0 +1,226 @@
+# Plan: Option B Per-Worker Interleaved Traversal
+
+**Status:** Planning only; do not implement until this design is reviewed.
+**Owner:** Codex for prototype design and implementation; operator for long-run
+benchmarks on `home`.
+**Background:** Option A, the central traversal inference server, was implemented
+and benchmarked on 2026-05-07. It regressed traversal because the existing
+recursive worker path is sync-blocking and can only feed batches near the worker
+count, not the GPU-efficient bs=64+ regime. See `docs/performance.md`:
+"Option A Bench Result and Structural Ceiling" and "Clarifying the traversal
+bottleneck: sync policy boundary, not SIMD."
+
+## Goal
+
+Restructure traversal scheduling so each worker advances many traversal
+instances concurrently, suspends each instance at policy-needed states, batches
+the pending policy requests, runs one policy forward, and resumes the matching
+instances.
+
+The target is to turn the current policy-forward shape:
+
+```text
+one traversal -> policy request -> bs=1 forward -> resume
+```
+
+into:
+
+```text
+N traversal continuations -> collect policy requests -> bs=32..128 forward -> resume
+```
+
+without changing CFR math, game rules, replay sample semantics, or public
+training CLI behavior.
+
+## Why This Is The Next Optimization
+
+Microbench evidence from `experiments/traversal_policy_boundary/`:
+
+| Device | Component | median us/call |
+| --- | --- | ---: |
+| CPU | encode + legal | 3.10 |
+| CPU | push + pop | 0.15 |
+| CPU | policy boundary bs=1 | 111.50 |
+| CUDA | policy boundary bs=1 | 181.30 |
+| CUDA | torch forward bs=64 | 2.55 |
+
+The game mechanics and encoding are not the dominant cost. The dominant cost is
+the one-row Python/PyTorch policy boundary. Option A moved that boundary to a
+server process, but because every worker blocks waiting for one response, the
+server observed mean batches around 7-8 and regressed end-to-end. Option B is
+the first design that directly changes the scheduling shape.
+
+## Non-Goals
+
+- Do not re-enable `traversal.inference_backend: server` as the default.
+  Option A remains available but structurally capped until traversal can feed
+  larger batches.
+- Do not port traversal to Julia, C++, or a new game engine.
+- Do not change Deep CFR sampling math, regret targets, strategy-memory
+  location, weighting, or replay schema.
+- Do not change model architecture.
+- Do not implement TensorRT, `torch.compile`, or AMP in this plan.
+- Do not remove the existing recursive traversal path until the interleaved path
+  has parity and benchmark evidence.
+
+## Design Sketch
+
+The current Cython traversal is recursive and calls policy synchronously. Option
+B needs an explicit continuation representation so policy calls become yield
+points.
+
+One worker owns a fixed set of active traversal contexts:
+
+```text
+TraversalContext
+  GameState state
+  explicit stack frames replacing recursion
+  RNG state
+  traverser player
+  iteration
+  partial node/action values
+  pending info_state/legal mask
+  output advantage/strategy samples
+  TraversalStats
+```
+
+Worker loop:
+
+1. Initialize `K` traversal contexts from the worker's assigned seeds.
+2. Advance each context until it reaches one of:
+   - terminal/cutoff/done,
+   - needs policy forward,
+   - error.
+3. Collect pending policy requests into a batch, grouped by network target:
+   advantage player 0/1, strategy network, or league snapshot if enabled.
+4. Run batched forward for each group on the worker's selected inference device.
+5. Scatter logits/advantages back into the contexts.
+6. Resume contexts until all assigned traversals finish.
+7. Return the same `(stats, advantage_samples, strategy_samples)` shape as
+   `run_cython_traversal_batch`.
+
+The first prototype should keep one worker process and one GPU model copy per
+worker if `device=cuda`. That may duplicate VRAM across workers, so the initial
+benchmark can run with fewer workers and larger `interleave_width`. A later
+hybrid can combine Option B's continuation batching with Option A's central
+server if VRAM pressure dominates.
+
+## Config Surface
+
+Add only after the prototype proves parity:
+
+```yaml
+traversal:
+  scheduler: recursive        # recursive | interleaved
+  interleave_width: 64        # traversal contexts advanced per worker
+  interleave_max_batch: 128   # cap per forward group
+```
+
+Default remains `recursive`.
+
+## Implementation Phases
+
+### Phase 0: Design Spike
+
+- Trace current `_traverse` control flow and enumerate every value that must
+  survive across a policy yield point.
+- Decide whether to implement the explicit stack in Cython (`.pyx`) or as a
+  Python prototype first.
+- Identify exact parity surfaces:
+  `TraversalStats`, advantage samples, strategy samples, RNG sequence, and
+  terminal/cutoff behavior.
+
+Deliverable: short design note appended to this plan before code work starts.
+
+### Phase 1: Python Prototype, No Production Wiring
+
+- Add an experiment-only traversal prototype under `experiments/` that mimics
+  the current traversal semantics with explicit stacks.
+- Use small configs (`max_depth`, low traversal count) and compare samples/stats
+  to the recursive path under fixed seeds.
+- Measure realized batch size and scheduler overhead.
+
+Success gate: sample/stat parity on small deterministic fixtures, and realized
+policy batches materially above worker count.
+
+### Phase 2: Cython Prototype Behind Non-Default Flag
+
+- Add an interleaved traversal entry point beside the existing recursive one.
+- Keep the existing recursive path untouched and default.
+- Wire through `workers.py` only behind `traversal.scheduler: interleaved`.
+- Add focused tests for parity on smoke configs.
+
+Success gate: `uv run pytest -q tests/games/classic/test_deep_cfr_trainer.py`
+and new interleaved traversal tests pass.
+
+### Phase 3: Benchmark
+
+Benchmark against current `default.yaml`, eval/checkpoint disabled:
+
+```bash
+uv run lost-cities-deep-cfr train \
+  --config configs/deep_cfr/default.yaml \
+  --set run.max_iterations=10 \
+  --set checkpoint.save_latest=false \
+  --set checkpoint.save_every=0 \
+  --set evaluation.eval_every=0
+```
+
+Compare:
+
+- recursive baseline,
+- interleaved with `interleave_width` in `{16, 32, 64, 128}`,
+- worker counts in `{1, 2, 4, 8}` as VRAM allows.
+
+Metrics:
+
+- `iteration_seconds`
+- `traversal_seconds`
+- realized policy batch size mean/p50/p95/max
+- scheduler overhead if instrumented
+- `advantage_train_seconds`, `strategy_train_seconds` to confirm no unrelated
+  drift
+- sample counts and traversal stats
+
+Success gate: at least **1.5x traversal speedup** with no sample/stat parity
+failure. Stretch target: **3x traversal speedup** if realized batches reach the
+bs=64 regime without high scheduler overhead.
+
+## Risks
+
+- **State-machine complexity.** Recursive CFR control flow has many local
+  values. Mitigation: prototype with small depth and exhaustive parity before
+  optimizing.
+- **RNG drift.** Interleaving changes operation order. Mitigation: store RNG
+  state per traversal context and define parity against the recursive path only
+  where ordering is intentionally preserved. If exact ordering is impossible,
+  require distributional/sample-count parity and document the break.
+- **VRAM duplication.** Per-worker GPU models may not fit at larger model sizes.
+  Mitigation: start with fewer workers and larger interleave width; revisit a
+  central server only after Option B proves the scheduling benefit.
+- **Sample memory pressure.** More active contexts mean more pending samples.
+  Mitigation: stream completed samples out of contexts as soon as a traversal
+  finishes.
+- **Scheduler overhead cancels batching.** Mitigation: benchmark light and heavy
+  modes; record realized batch size and overhead explicitly.
+
+## Decision Tree
+
+- **Parity fails in Phase 1/2:** stop. Do not optimize. Document the exact
+  mismatch.
+- **Parity passes, realized batch remains <16:** Option B did not change the
+  structural ceiling enough. Reconsider Option C or a deeper traversal rewrite.
+- **Parity passes, realized batch >=64, speedup <1.5x:** batching worked but
+  non-forward work dominates. Keep recursive default and document.
+- **Parity passes, traversal speedup >=1.5x:** keep interleaved behind config,
+  run longer learning-curve A/B.
+- **Longer A/B is stable and speedup persists:** consider making
+  `traversal.scheduler: interleaved` the default.
+
+## Definition Of Done
+
+- Plan reviewed and Phase 0 design note added.
+- Prototype proves whether explicit continuation batching can preserve traversal
+  semantics.
+- Bench results are added to `docs/performance.md`.
+- Default behavior remains unchanged until parity and benchmark gates pass.

docs/plans/torch_compile.md

@@ -5,18 +5,19 @@
 **Background:** See `docs/performance.md`:
 - "Experiments → `torch.compile` on trainer networks (2026-05-07, regression)" — the prior attempt regressed iter time by +4.8% on the small default model. Implementation is preserved on branch `experiments/torch-compile`.
 - "Post-A Optimization Calculus (forward-looking, 2026-05-07)" — argues compile/TensorRT become meaningful only after model size grows out of the dispatch-bound regime and/or eval density rises. This plan honors that sequencing.
+- "Clarifying the traversal bottleneck: sync policy boundary, not SIMD" — the current traversal bottleneck is scheduling shape, not a compile-able model-kernel problem.
 
 This plan is the follow-up referenced in step 3 of "Recommended sequencing".
 
 ## Goal
 
-Re-enable `torch.compile` on the Deep CFR trainer's networks at a model size where kernel work amortizes compile dispatch overhead, and (optionally, secondary) on the inference server's eval-mode forward path introduced by `docs/plans/batched_traversal_inference_server.md`. The goal is iter-time speedup with **no learning-curve drift**.
+Re-enable `torch.compile` on the Deep CFR trainer's networks at a model size where kernel work amortizes compile dispatch overhead, and (optionally, secondary) on batched eval/inference-server forward paths. The goal is iter-time speedup with **no learning-curve drift**.
 
 ## Non-goals
 
 - Do not enable compile on the current 512-hidden / 3-layer `default.yaml`. The regression is already measured.
 - Do not compile the Cython traversal call-site networks (workers' CPU networks under `inference_backend: local`). Their per-call shape is `batch_size == 1` and dispatch-dominated; compile cannot help and recompilation triggers are higher-risk.
-- Do not introduce TensorRT here. TensorRT is a separate work item also gated on the inference server (see `docs/performance.md` § "Tooling split").
+- Do not introduce TensorRT here. TensorRT is a separate work item for batched inference/eval surfaces (see `docs/performance.md` § "Tooling split").
 - Do not change network architectures. This plan picks up whatever larger model the project settles on in step 2 of the post-A sequencing.
 - Do not change the public CLI surface. Compile toggles via config only.
 - Do not enable `mode="max-autotune"` by default. It is opt-in for benchmarking.
@@ -59,11 +60,17 @@ Wraps `advantage_networks[player]` and `strategy_network` at trainer constructio
 
 Why this target: the trainer's optimization step has fixed batch shapes (no dynamic shapes), runs many steps per iteration, and exercises forward+backward+optimizer — the regime where compile pays best when the kernel is large enough.
 
-### B. Inference-server forward (eval-mode, no grad)
+### B. Batched inference/eval forward (eval-mode, no grad)
 
-Optional follow-up. After Option A from the batched-traversal plan lands and is benchmarked, optionally compile the server's `model.forward` path with `mode="reduce-overhead"` or `mode="default"`. The server already uses `eval()` + `inference_mode()`. Batches are bounded by `max_batch` (default 256) but variable in size up to that — this introduces a dynamic-shape concern (see Risks).
+Optional follow-up. The Option A inference server has landed and benchmarked,
+but it is structurally capped by sync-blocking traversal and is not the default.
+Only compile the server's `model.forward` path after either (a) Option B-style
+interleaved traversal can feed meaningful batches, or (b) the target is
+evaluation, which already has batching. The server/eval path uses `eval()` +
+`inference_mode()`. Batches vary in size, which introduces a dynamic-shape
+concern (see Risks).
 
-Why this target is secondary: per `docs/performance.md` § "Why compile / TensorRT are negligible *today* but become meaningful later", the inference-server forward share of an iter is <1% post-A on the small model. Even a 50% forward speedup is iter-level negligible until the model grows. At that point compile and TensorRT compete for the same role; this plan covers compile, and the TensorRT plan (separate, future) covers TensorRT.
+Why this target is secondary: per `docs/performance.md` § "Why compile / TensorRT are negligible *today* but become meaningful later", small-model inference forward is not the iter-level limiter unless traversal can actually feed batched requests. At larger model sizes or denser evaluation, compile and TensorRT compete for the same role; this plan covers compile, and the TensorRT plan (separate, future) covers TensorRT.
 
 ## Compile mode selection
 
@@ -80,7 +87,7 @@ For the inference server (target B), `default` is the only safe mode initially.
 - `src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py` — owns trainer-side networks, optimizer, and the train loop. This is where target A's `torch.compile(...)` calls go, plus the `_clean_state_dict()` helper and the `_orig_mod`-routed `load_state_dict` shim. Reference implementation lives on branch `experiments/torch-compile` (commit `05cc02a`).
 - `src/coolrl_lost_cities/games/classic/deep_cfr/networks.py` — network classes. Not modified by this plan; compile wraps the constructed module from outside.
 - `src/coolrl_lost_cities/games/classic/deep_cfr/workers.py` — multiprocessing traversal worker entry. Reconstructs CPU networks from `state_dict`s. Must keep receiving cleaned (no-`_orig_mod.`) state dicts.
-- `src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py` — created by the batched-traversal plan. Target B wraps `model.forward` here. If the inference-server plan has not landed, target B is deferred.
+- `src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py` — created by the archived Option A plan. Target B wraps `model.forward` here only after Option B or eval batching provides large enough batches.
 - `src/coolrl_lost_cities/games/classic/deep_cfr/config.py` — extend with a `CompileConfig` block.
 
 ## Reference implementation (branch `experiments/torch-compile`)
@@ -186,7 +193,7 @@ Caveats:
 - Cherry-pick `experiments/torch-compile` (commit `05cc02a`) onto a fresh branch. Resolve any conflicts against current `main`.
 - Move the unconditional `torch.compile(...)` calls behind `config.compile.trainer.enabled`. Default is `false`.
 - Add the `CompileConfig` and `CompileTrainerConfig` dataclasses to `config.py`. Wire `--set compile.trainer.enabled=true` through.
-- Confirm `_clean_state_dict()` and the `_orig_mod`-routed `load_state_dict` paths are correctly invoked on every checkpoint save, every checkpoint load, every weight push to workers, and (post-A) every weight push to the inference server.
+- Confirm `_clean_state_dict()` and the `_orig_mod`-routed `load_state_dict` paths are correctly invoked on every checkpoint save, every checkpoint load, every weight push to workers, and every weight push to the inference server if target B is enabled.
 
 ### Step 2: tests for state-dict round-tripping
 
@@ -208,7 +215,7 @@ See "Bench plan" below. This is the gating step. If results do not clear success
 
 ### Step 6: (optional) inference-server forward compile (target B)
 
-- Only if the batched-traversal-inference server has landed (per `docs/plans/batched_traversal_inference_server.md`).
+- Only if Option B or evaluation batching can feed large enough batches to make the inference-server forward a meaningful target. The archived Option A server exists, but the sync-blocking traversal path did not feed large batches.
 - Add `CompileInferenceServerConfig` wiring. Add `torch.compile(...)` invocation inside the server process. Route weight sync through `_orig_mod` when enabled.
 - Add a small inference-server bench script (or extend `scripts/bench_inference_backend.py` if it has landed) to measure server-side `policy_network_seconds` with and without compile at the chosen larger model size.
 - Ship only if criterion 6 is met. Otherwise leave disabled.
@@ -247,7 +254,7 @@ For target B, bench `policy_network_seconds` from the inference-server side with
 - **Backward path not benefiting.** Compile's biggest theoretical wins on the trainer phase come from fusing forward+backward+optimizer. In practice on simple MLPs the optimizer is already fused (e.g. `torch.optim.Adam(..., fused=True)` if available); compile may add little on top. Mitigation: bench is the answer — if iter time does not move 5%, the plan does not merge. This is a real possibility.
 - **CFR variable-length traversal does not feed compile.** Confirmed: traversal call sites use `batch_size == 1` and run on CPU workers (or via the inference server, which is target B not target A). Target A only sees the trainer optimization batches, which are fixed-shape. So dynamic-shape concerns do not apply to target A.
 - **GPU contention with eval.** Eval already runs on the trainer's device. Compile increases peak memory during compile/autotune phases (especially `max-autotune`). Mitigation: compile happens once per process at startup; eval runs after warm-up. Bench with eval disabled to isolate iter time, then re-enable for the drift check.
-- **AMP interaction.** `run.use_amp` is currently a no-op (per `docs/performance.md`). If AMP is implemented before this plan ships, re-bench compile under AMP — the two interact and prior numbers do not transfer.
+- **AMP interaction.** Trainer AMP is implemented but default-off after the 2026-05-07 smoke regression. If the model-size experiment later makes AMP attractive, re-bench compile with and without AMP because the two interact and prior numbers do not transfer.
 
 ## Decision tree