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>
This commit is contained in:
2026-05-07 22:15:09 +09:00
co-authored by Codex
parent 7c20d53103
commit e5ba247fcc
6 changed files with 268 additions and 21 deletions
+259
View File
@@ -0,0 +1,259 @@
# Plan: Trainer-side AMP (Automatic Mixed Precision)
**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.
## Goal
Wire `run.use_amp` into the Deep CFR trainer's optimization phases (`_train_advantage` and `_train_strategy`) using `torch.autocast` + `torch.amp.GradScaler`. Cut `advantage_train_seconds + strategy_train_seconds` (~7.1s/iter combined, ~40% of a non-eval iteration on the inspected default run) without regressing the learning curve.
## Non-goals
- Do not change traversal worker behavior. Traversal workers run on CPU today; AMP does not apply.
- Do not change the inference server's forward path. The server has its own `inference_server.use_amp` flag (eval-only AMP, no GradScaler). Trainer AMP and server AMP are independent.
- Do not change replay buffer dtype. Samples remain float32 in shared memory and on the host side; only the trainer's forward/backward switches to mixed precision.
- Do not change network dtype, parameter dtype, or optimizer dtype. AMP only autocasts forward; parameters and master weights remain fp32.
- Do not implement bf16 (no GradScaler needed) as the default. fp16 is the primary target; bf16 is a follow-up.
## Success criteria
1. With `--set run.use_amp=true` on `configs/deep_cfr/default.yaml` running on CUDA, `advantage_train_seconds + strategy_train_seconds` drops by at least 15% averaged over 20 non-eval iterations on `home` (RTX 3090) compared to `--set run.use_amp=false` baseline. **Realistic upper bound:** ~30% on `home`. If measured speedup is below 5%, document the result in `docs/performance.md` and leave the flag default-off (analogous to the `torch.compile` regression).
2. With `--set run.use_amp=true`, `loss/advantage` and `loss/strategy` trajectories track the fp32 baseline within seed noise over at least 50 iterations on `default.yaml`. No NaN/Inf appears in `loss/*` rows of `metrics.jsonl`.
3. Eval-side win-rate trajectories (`eval/<opponent>/win_rate0`) under AMP are indistinguishable from the fp32 baseline at iteration 50, 100, and 200 within the noise band of a single-seed comparison.
4. `--set run.use_amp=false` (the default) produces byte-identical training to `main` for the same seed: no AMP-related code path runs.
5. `run.use_amp=true` on CPU is a documented no-op (CUDA not available → skip autocast/scaler) and does not crash. Same for non-CUDA `run.device`.
6. All existing tests pass. New unit test covering the AMP code path passes.
## Key files (current)
- `src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py`:
- `DeepCFRTrainer.__init__` (lines ~187-258) — instantiates networks, optimizers; this is where the `GradScaler` should be created.
- `DeepCFRTrainer._train_advantage` (lines ~892-948) — advantage forward+backward+step. AMP wrapping target.
- `DeepCFRTrainer._train_strategy` (lines ~950-995) — strategy forward+backward+step. AMP wrapping target.
- `DeepCFRTrainer._batch_tensors` (lines ~853-883) — tensors stay float32 (do not cast inputs).
- `src/coolrl_lost_cities/games/classic/deep_cfr/config.py`:
- `RunConfig.use_amp: bool = False` (line 25) — flag already exists.
- `src/coolrl_lost_cities/games/classic/deep_cfr/networks.py` — no changes; `DeepCFRMLP` and `ColorSharedNetwork` work under autocast as-is.
- `tests/games/classic/test_deep_cfr_trainer.py` — extend with AMP smoke test.
## New surface
No new files. All changes live in `trainer.py` (and one optional config addition for AMP dtype selection). The plan extends `RunConfig` minimally if we want bf16 selectability; otherwise `run.use_amp` alone is sufficient and dtype is fp16 by default.
Optional `RunConfig` extension (deferred — only add if Step 5 measurement motivates it):
```python
class RunConfig(StrictModel):
# ... existing fields ...
use_amp: bool = False
amp_dtype: str = "float16" # or "bfloat16"
```
## Where AMP code goes
### Construction (in `__init__`)
After optimizers are created:
```python
self._amp_enabled = bool(self.config.run.use_amp) and self.device.type == "cuda"
self._amp_dtype = torch.float16 # bf16 deferred
self._scaler = torch.amp.GradScaler("cuda", enabled=self._amp_enabled)
```
Single shared `GradScaler` across both advantage and strategy networks is fine — it tracks one global loss-scale that adjusts based on observed Inf/NaN gradients and applies to whichever optimizer it is told to step. Sharing the scaler matches PyTorch's recommended pattern for multi-network training and avoids two competing scale schedules.
### Advantage train loop (`_train_advantage`)
Replace the forward + backward + step block with:
```python
with torch.autocast(device_type="cuda", dtype=self._amp_dtype, enabled=self._amp_enabled):
pred = network(x)
diff = (pred - y).masked_fill(~legal, 0.0)
if self.config.training_weighting.mode == "none":
loss = diff.square().sum() / legal.sum().clamp_min(1)
elif self.config.training_weighting.mode == "lcfr":
...
else:
...
optimizer.zero_grad(set_to_none=True)
self._scaler.scale(loss).backward()
if self.config.optimization.grad_clip > 0.0:
self._scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(network.parameters(), self.config.optimization.grad_clip)
self._scaler.step(optimizer)
self._scaler.update()
losses.append(float(loss.detach().float().cpu()))
```
Notes:
- `legal.sum().clamp_min(1)` and the weighted-mean denominators are int/float reductions that stay fp32 — `autocast` only casts ops that are on the autocast-allow-list. The division is safe.
- `loss.detach().float().cpu()` — explicit `.float()` to avoid logging an fp16 NaN that prints surprisingly.
- The `unscale_` before `clip_grad_norm_` is mandatory; otherwise the clip threshold is applied to scaled gradients. This is the standard PyTorch AMP idiom.
### Strategy train loop (`_train_strategy`)
Same pattern. The `log_softmax` over masked logits remains numerically stable under fp16 because `masked_fill(~legal, finfo.min)` uses fp32-min before autocast can downcast — verify this; if autocast downcasts the mask fill value, replace with `torch.finfo(self._amp_dtype if self._amp_enabled else torch.float32).min` or apply the mask after autocast. Safer: apply `masked_fill` *outside* the autocast block (operating on the post-cast logits is fine — `masked_fill` is allow-listed).
### Disabled path
When `self._amp_enabled is False`, `torch.autocast(..., enabled=False)` is a true no-op (does not change dispatcher state), and `GradScaler(enabled=False)` makes `scale`, `unscale_`, `step`, and `update` all forward to plain optimizer behavior. So the same code runs in both modes — no branching needed in the hot loop.
## Flag semantics
| `run.use_amp` | `run.device` resolves to | Behavior |
| --- | --- | --- |
| `false` (default) | any | No autocast, no scaler. Behavior identical to current `main`. |
| `true` | `cuda` | Autocast(fp16) + GradScaler active in advantage/strategy train loops. |
| `true` | `cpu` | Logs a one-time warning "AMP requested but device is CPU; running fp32." Trainer behaves as `false`. |
The warning is logged via `self.tracker.log_event(...)` once at run start, after `_amp_enabled` is computed.
Trainer-side AMP is fully independent of `inference_server.use_amp`. Both can be on, both off, or either alone:
- `inference_server.use_amp` → server-process `torch.autocast` around the inference forward inside `inference_mode()`. No GradScaler (no backward pass). Affects traversal forward latency only.
- `run.use_amp` → trainer-process autocast + GradScaler around `_train_advantage` and `_train_strategy` forward+backward+step. Affects optimization phases only.
## Numerical stability considerations
CFR regret targets can have wide dynamic range — both very small (regrets near zero for converged actions) and large in magnitude (early-iteration noisy estimates). fp16's 1e-5 to 6.5e4 representable range is narrow compared to fp32, so two failure modes are realistic:
1. **Gradient overflow → NaN.** Mitigated by `GradScaler`. The scaler observes Inf/NaN in unscaled gradients, skips the step, and halves the scale. Standard.
2. **Forward overflow in the squared-error loss.** `diff.square()` on fp16 inputs can overflow if `(pred - y)` magnitude exceeds ~256. This is upstream of GradScaler — the loss itself becomes Inf in the autocast region and the entire batch is wasted. If this happens, the scaler will skip the step but the loss-scale scheduler will not recover because the issue is in the forward, not the gradient.
Mitigations for (2):
- Compute the squared error in fp32 explicitly: cast `diff` to fp32 via `diff.float()` before `.square()` if measurement shows fp16 overflow on real CFR samples. This is cheap and confined to the trainer.
- Clamp `pred - y` to a wide-but-safe range before squaring: e.g. `diff.clamp(-128, 128)`. Only do this if measurement shows it's needed — clamping silently changes loss semantics.
- Default plan: do not pre-mitigate. Add a NaN/Inf guard (see below); if it triggers in the wild, switch to the explicit `.float()` cast in the loss computation.
### NaN/Inf guard
After computing `loss` in each train step:
```python
if not torch.isfinite(loss):
self._runtime_metrics["amp/nonfinite_loss_count"] = (
int(self._runtime_metrics.get("amp/nonfinite_loss_count", 0)) + 1
)
optimizer.zero_grad(set_to_none=True)
continue
```
This makes overflow visible in `metrics.jsonl` per iteration without aborting training. Codex should also expose `_scaler.get_scale()` as `amp/grad_scale` per iteration so we can see the scaler's scale schedule in W&B / metrics plots.
## Validation: A/B learning-curve test
After implementation, run two short trainings with the same seed:
```bash
# Baseline
uv run lost-cities-deep-cfr train \
--config configs/deep_cfr/default.yaml \
--keep \
--set run.max_iterations=100 \
--set run.use_amp=false \
--set run.experiment_name=amp-baseline
# AMP
uv run lost-cities-deep-cfr train \
--config configs/deep_cfr/default.yaml \
--keep \
--set run.max_iterations=100 \
--set run.use_amp=true \
--set run.experiment_name=amp-on
```
Compare with `lost-cities-deep-cfr analyze` on each, plus a side-by-side table of:
- `loss/advantage`, `loss/strategy` at iters {25, 50, 75, 100}
- `eval/<opponent>/win_rate0` at iters {50, 100} (eval_every=25 is default)
- `time/advantage_train_seconds`, `time/strategy_train_seconds`, `time/iteration_seconds` means over iters 5-100 (iter 1-4 dropped as warm-up)
- `amp/nonfinite_loss_count`, `amp/grad_scale` (AMP run only)
Acceptance: AMP win-rates within fp32 win-rates ± noise, training-phase wall clock down by ≥15%, no nonfinite losses, scaler scale stable (not collapsing toward 1).
## Implementation steps (ordered, each independently mergeable)
### Step 1: scaler + flag plumbing in `__init__`
- In `DeepCFRTrainer.__init__`, after the optimizer block, compute `self._amp_enabled`, `self._amp_dtype`, `self._scaler`.
- Log a one-time warning if `config.run.use_amp` is requested but `self.device.type != "cuda"`.
- Add an `amp/grad_scale` runtime metric emitted once per iteration in `run_iteration` (read `self._scaler.get_scale()` after train phases).
- No behavior change yet — scaler is created but unused.
### Step 2: wrap `_train_advantage`
- Add `with torch.autocast(...)` around the forward + loss computation.
- Replace `loss.backward()``self._scaler.scale(loss).backward()`.
- If `grad_clip > 0`: `self._scaler.unscale_(optimizer)` before `clip_grad_norm_`.
- Replace `optimizer.step()``self._scaler.step(optimizer)` + `self._scaler.update()`.
- Add the NaN/Inf guard described above; emit `amp/nonfinite_loss_count` to runtime metrics.
### Step 3: wrap `_train_strategy`
- Same pattern as Step 2.
- Verify `masked_fill(~legal, torch.finfo(torch.float32).min)` semantics under autocast. Either keep the mask fill outside autocast or use a dtype-aware `finfo`.
### Step 4: unit test
- Extend `tests/games/classic/test_deep_cfr_trainer.py`:
- `test_train_advantage_amp_smoke`: build a tiny trainer with `run.use_amp=True`, `device="cuda"` (skip if CUDA unavailable via `pytest.skip`). Seed memories with synthetic samples. Run one iteration. Assert no exception, finite `loss/advantage`, `amp/grad_scale` present in runtime metrics.
- `test_amp_cpu_falls_back`: `run.use_amp=True`, `device="cpu"`. Assert iteration runs without error and behaves as fp32 (e.g. compare loss to a fp32 reference run on the same seed).
### Step 5: bench + learning-curve A/B
- Add a short script `scripts/bench_amp_trainer.py` (mirrors `scripts/profile_gpu_forward.py` style):
- Construct trainer with `default.yaml`.
- Pre-fill advantage and strategy memories with `optimization.advantage_batch_size * advantage_updates_per_iteration` synthetic samples (so the train phases run end-to-end without needing a real traversal).
- Run `_train_advantage` and `_train_strategy` 20 times each under `use_amp=False` and `use_amp=True`. Drop the first 2 as warm-up. Print mean ms per call.
- Run the full A/B (Validation section above) and append a date-stamped subsection to `docs/performance.md` → "Experiments" with: hardware, iter-time delta, loss/win-rate parity table, NaN counts, scale schedule, decision (default-on / default-off / disabled).
### Step 6: documentation
- Update `docs/performance.md` § "AMP Status" to point to the new experiment subsection and remove the "treat as a no-op" line if AMP becomes default-on, or leave it and explain why if AMP measurement was a regression.
- No CLAUDE.md / AGENTS.md updates needed — `--set run.use_amp=true` already works syntactically.
## Definition of done
- `run.use_amp=true` on CUDA produces measurable speedup on advantage+strategy train phases on `home`, with the A/B learning curve showing parity within noise over 100 iterations.
- `run.use_amp=false` is byte-identical to `main`.
- AMP code path covered by tests; CPU fallback documented and tested.
- `metrics.jsonl` exposes `amp/grad_scale` and `amp/nonfinite_loss_count` so future runs are self-diagnosing.
- An experiment subsection is appended to `docs/performance.md` recording the speedup and the parity check.
- `uv run ruff check .` passes; `uv run pytest -q tests/games/classic/test_deep_cfr_trainer.py` passes.
## Risks and mitigations
- **Small-model regression (analogous to `torch.compile` 2026-05-07).** `DeepCFRMLP` at hidden=512, 3 layers is small. AMP overhead per call (cast/uncast, GradScaler bookkeeping) may exceed the kernel speedup at this size — the same dynamic that bit `torch.compile`. Mitigation: measure first; if speedup is below 5%, keep default off and document. Preserve the implementation on a branch (`experiments/amp-trainer`) for revisiting if `NetworkConfig.hidden_size` or `num_layers` increases.
- **fp16 overflow in `diff.square()`.** Mitigation: NaN/Inf guard logs the count; if observed, cast `diff` to fp32 before squaring inside the autocast region (`diff.float().square()`). PyTorch will not redowncast it.
- **GradScaler scale collapse.** If many consecutive steps overflow, the scaler can drop to scale=1 and stay there, defeating the point. Mitigation: log `amp/grad_scale` every iteration; if it stays ≤16 for >10 iterations, switch the loss-side cast (above) on.
- **Interaction with `clip_grad_norm_`.** Forgetting `unscale_` before clipping silently changes the effective clip threshold. Mitigation: explicit step in the implementation; covered by code review.
- **Determinism breakage.** Autocast can change op kernels and therefore reduction order; bit-exact reproducibility vs fp32 is not preserved. This is expected. Mitigation: A/B is on learning trajectories within seed noise, not on bit-identity.
- **Inference server contention.** If `inference_server.use_amp` and `run.use_amp` are both on, two autocast regions exist in two processes — they do not conflict. Trainer GradScaler does not affect server inference.
## Bench plan
Two artifacts:
1. `scripts/bench_amp_trainer.py` (new, small) — micro-bench the `_train_advantage` and `_train_strategy` methods in isolation under `use_amp=False` vs `True`. Reports per-call ms with mean + p50 + p95 over 20 runs. This is the fast feedback loop during implementation.
2. End-to-end run pair (above) — the real signal. 100-iter `default.yaml` baseline vs AMP, same seed, single GPU. Compare iter-time and learning curves.
Compare pattern matches the `torch.compile` experiment write-up in `docs/performance.md`. Use a similar table:
| | iter mean | adv+strat mean | adv+strat share | 1000-iter projection |
| --- | ---: | ---: | ---: | ---: |
| Baseline (use_amp=false) | TBD | TBD | TBD% | TBD h |
| AMP (use_amp=true) | TBD | TBD | TBD% | TBD h |
| Effect | TBD | TBD | TBD pp | TBD min |
## Out-of-scope follow-ups (do not start in this plan)
- bf16 dtype option (`amp_dtype: bfloat16`). bf16 sidesteps GradScaler entirely and avoids the fp16 overflow class. Add only if (a) hardware supports it efficiently (Ampere+) and (b) fp16 measurement shows scale collapse or frequent nonfinite losses.
- AMP on the inference server's training-side weight push path. Server already has `inference_server.use_amp` for forward; backward is not its job.
- AMP on evaluation forward. Evaluation runs `eval()` + no-grad; if eval forward becomes a bottleneck (see "Evaluation Optimization Options" #7 in `docs/performance.md`), wrap there separately.
- `torch.compile` retry. The 2026-05-07 regression was size-bound. Re-evaluate only after `NetworkConfig` grows substantially (per "Post-A Optimization Calculus" in `docs/performance.md`); coordinate with that work, not this plan.
- Larger model config experiment. AMP becomes meaningfully more useful at hidden≈1024 / layers≈6 per the post-A calculus. That is a separate model-architecture work item, not an AMP work item.
@@ -0,0 +1,209 @@
# Plan: Batched Traversal Inference Server (Priority #5, Option A)
**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.
## Goal
Replace per-worker single-state CPU policy forwards in Deep CFR traversal with a central GPU inference server that batches policy requests across all workers. Targets the dominant phase (`traversal_seconds` ≈ 60% of iteration time).
## Non-goals
- Do not modify the Cython traversal recursion structure.
- Do not modify the game engine, replay buffer, or training loop math.
- Do not implement Option C (single-process vectorized traversal). Keep it as future work.
- Do not require Cython `nogil`-cleanliness.
- Do not change the public CLI surface.
## Success criteria
1. With `traversal.inference_backend: server` enabled on `configs/deep_cfr/default.yaml`, end-to-end training produces eval-winrate trajectories indistinguishable (within seed noise) from the current `local` backend over at least 50 iterations on `home` hardware.
2. On `home` (6-core + RTX 3090), `traversal_seconds` decreases by at least 30% compared to the current run profile in `docs/performance.md`.
3. On `remote` (32-core + weak GPU), `traversal_seconds` decreases or stays within 10% of current; if it regresses more, fall back to `local` is the operator's choice — the plan still ships.
4. With `traversal.inference_backend: local`, behavior is byte-identical to current `main`.
5. All existing tests pass. New unit tests for the inference client/server round-trip pass.
## Key files (current)
- `src/coolrl_lost_cities/games/classic/deep_cfr/workers.py``run_traversal_worker_batch`. Spawns CPU networks per worker (`device = torch.device("cpu")` at line 62). This is the worker entry point that must learn about the inference server.
- `src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx` — recursive traversal. Two policy call sites:
- Lines 473475: `networks[player](x).squeeze(0).detach().cpu().numpy()` — advantage-net forward.
- Lines 550552: `self.strategy_network(x).squeeze(0).detach().cpu().numpy()` — strategy-net forward.
- Both are the seams that route through the inference client when the backend is `server`.
- `src/coolrl_lost_cities/games/classic/deep_cfr/networks.py``DeepCFRMLP`.
- `src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py` — main loop; spawns the worker pool and owns the trainer-side networks. Must also start/stop the inference server and push weights periodically.
- `src/coolrl_lost_cities/games/classic/deep_cfr/config.py` — config schema. Add `traversal.inference_backend` and an `inference_server` block.
## New files
- `src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py` — server process: owns models on GPU, drains the request queue, runs batched forward, writes responses.
- `src/coolrl_lost_cities/games/classic/deep_cfr/inference_client.py` — client used inside workers: claims a request slot, posts the encoded state, waits for the response. Provides a `forward(network_id, player, state)` API.
- `src/coolrl_lost_cities/games/classic/deep_cfr/inference_buffers.py` — shared-memory tensor pool: pre-allocated `[num_slots, input_dim]` request buffer and `[num_slots, action_size]` response buffer, plus per-slot ready events and a free-slot stack.
- `tests/games/classic/deep_cfr/test_inference_server.py` — unit/integration tests for the request path.
## Architecture
### Process layout
```
main process (trainer)
├─ inference server process (GPU)
│ - owns advantage networks (per-player), strategy network, league snapshots
│ - request buffer (shared mem): [num_slots, input_dim] float32
│ - response buffer (shared mem): [num_slots, action_size] float32
│ - control queue (mp.Queue): control messages only
│ - per-slot ready event (mp.Event[num_slots])
├─ traversal worker 1 ─┐
├─ ... ├─ inference clients post requests, wait on per-slot event
└─ traversal worker N ─┘
```
### Request protocol
A worker policy call becomes:
1. Client pops a free slot id from a shared `mp.Queue`-backed free-slot stack (small int).
2. Client writes the encoded state into `request_buffer[slot_id]` (shared memory; no pickle).
3. Client puts a `RequestMessage(slot_id, network_kind, player, network_index)` onto the request control queue. `network_kind ∈ {ADVANTAGE, STRATEGY, LEAGUE}`. Pickle cost is negligible (small struct of ints).
4. Server drains the request queue with a short batch window (`batch_window_us`, default 200μs) up to `max_batch` (default 256). Empty drain blocks on the queue with a small timeout.
5. Server stacks the requested rows from `request_buffer`, runs forward, writes outputs back to `response_buffer[slot_id]` for each request.
6. Server fires the per-slot ready event for each completed request.
7. Client wakes on its slot's event, reads `response_buffer[slot_id]`, copies to a local numpy array, returns the slot to the free-slot stack.
### Weight sync
- Trainer holds master weights. On a fixed cadence (`inference_server.weight_sync_every` iterations, default 1 — push every iter), trainer sends a `WeightUpdateMessage` containing `state_dict`s via `torch.multiprocessing` (auto-shares tensor storage; cheap once and copied into server's GPU model).
- Server applies `load_state_dict` and signals `weight_sync_complete`. Trainer waits before kicking off the next traversal iteration.
- League snapshots: passed via the same channel when the league updates. League list is small relative to per-iter cost.
### Backpressure & lifecycle
- Free-slot stack size = `num_slots` (default `max(64, 4 * num_workers * worker_chunk_size)`). Workers block on slot allocation if all slots in flight; this is the natural backpressure.
- Server shutdown: trainer puts a `Shutdown` sentinel on the control queue at training end. Server drains, exits.
- Crash isolation: if the server dies, workers will hang on their slot events. Trainer monitors the server process; on death, it raises and tears down the pool. No silent corruption.
### IPC choice
- `torch.multiprocessing` for weight passing (auto-shares tensors).
- `multiprocessing.shared_memory.SharedMemory` (numpy view) for request/response buffers — manual slot management. Pre-allocated once at startup; no per-call allocation.
- `multiprocessing.Queue` for control messages only (slot ids and small structs).
- `multiprocessing.Event` array for per-slot wakeups.
Rationale: see `docs/performance.md` § "IPC mechanism: multiprocessing + shared memory".
## Config schema
Extend `config.py`:
```python
@dataclass
class InferenceServerConfig:
enabled: bool = False # if False, behave like current code
device: str = "cuda" # server-side device
num_slots: int | None = None # None → auto: max(64, 4 * num_workers * worker_chunk_size)
max_batch: int = 256
batch_window_us: int = 200
weight_sync_every: int = 1 # iterations
use_amp: bool = False # eval-only AMP for forward (no grad)
@dataclass
class TraversalConfig:
# ... existing fields ...
inference_backend: Literal["local", "server"] = "local"
```
In `default.yaml`, leave `inference_backend: local` for now. Add an explicit `configs/deep_cfr/default_server.yaml` variant that flips it on for benchmarking.
## Worker integration
`run_traversal_worker_batch` (workers.py) currently:
- Builds CPU `DeepCFRMLP` instances and loads `state_dict`s.
- Passes them as positional `networks` into the `Traversal` Cython object.
When `inference_backend == "server"`:
- Skip building local networks. Instead, build an `InferenceClient` bound to the shared buffers and queues that the trainer wires in via `TraversalWorkerBatch`.
- Pass a small `NetworkProxy` object with the same call signature as the current network: `proxy(x: torch.Tensor) -> torch.Tensor`. Internally it converts to numpy, calls `client.forward(...)`, returns a torch tensor.
- Critically: the Cython traversal sites at `traversal.pyx:473-475` and `:550-552` should not need source changes if `NetworkProxy` is a callable returning a 2-D tensor. The existing `.squeeze(0).detach().cpu().numpy()` chain still works on the proxy's returned tensor (which can just be a CPU tensor wrapping the numpy result). **Verify this; if Cython has typed assumptions that reject a Python proxy, fall back to a thin Python helper invoked from the `.pyx` instead.**
Add to `TraversalWorkerBatch`:
- `inference_handles: InferenceClientHandles | None` — shared-memory names, queue handles, event arrays. None when `inference_backend == "local"`.
## Trainer integration
`trainer.py`:
1. On run start, if `inference_backend == "server"`: instantiate `InferenceServer` (spawns process), build `InferenceClientHandles`, push initial weights, wait for `weight_sync_complete`.
2. Per iteration: before `pool.starmap(run_traversal_worker_batch, ...)`, push fresh weights if `iter % weight_sync_every == 0`. Pass `inference_handles` into each `TraversalWorkerBatch`.
3. After traversal: same as today.
4. On run end: send shutdown sentinel; join the server process.
## Implementation steps (ordered, each independently mergeable)
### Step 1: shared-memory buffer module
- Create `inference_buffers.py` with `InferenceBuffers` class: pre-allocates request/response numpy arrays via `SharedMemory`, exposes `attach(name)` for child processes, has `release()` cleanup.
- Free-slot management: `mp.Queue` of slot ids, populated at startup with `range(num_slots)`.
- Per-slot ready events: `[mp.Event() for _ in range(num_slots)]`.
- Unit test: parent creates buffers, child attaches, writes a row, parent reads. Verify zero-copy semantics.
### Step 2: inference client
- Create `inference_client.py`. `InferenceClient.forward(network_kind, player, network_index, state_np: np.ndarray) -> np.ndarray`:
- Pop free slot, write state, post request, wait on event, copy response, return slot.
- Add a `NetworkProxy` callable that wraps `client.forward` to look like a `torch.nn.Module` for traversal call sites.
- Unit test: stub a server thread that echoes state*2; assert client gets the expected output.
### Step 3: inference server
- Create `inference_server.py`. Spawn-friendly entry function `run_inference_server(handles, model_config, control_queue, weight_queue)`.
- Owns models on `device`. Sets `eval()` and `inference_mode()`.
- Main loop: drain request queue with `batch_window_us` deadline, group by `(network_kind, network_index)`, run batched forward per group, scatter outputs to response buffer slots, fire events.
- Handles `WeightUpdateMessage` and `Shutdown`.
- Optional `use_amp`: wrap forward in `torch.autocast` when configured.
- Unit test: send N requests across multiple kinds; verify outputs match `model(stacked_input)`.
### Step 4: config + trainer wiring
- Extend `config.py` with `InferenceServerConfig` and `TraversalConfig.inference_backend`.
- Update `trainer.py` to start/stop server, push weights, attach handles to `TraversalWorkerBatch`.
- Update `workers.py`: when `inference_backend == "server"`, build proxies instead of CPU networks.
- Add `configs/deep_cfr/default_server.yaml` flipping `traversal.inference_backend: server`.
### Step 5: traversal call-site verification
- Run with `inference_backend: server` and a tiny config (1 worker, 1 traversal). Confirm the proxy is callable from `traversal.pyx:473-475` and `:550-552` without Cython type errors.
- If Cython rejects the proxy: refactor those two call sites to invoke a Python helper that takes `(networks, player, info_state)` and returns a numpy array. The helper picks `local` or `server` path by inspecting the object. This is a 2-line change per site.
### Step 6: integration tests
- `tests/games/classic/deep_cfr/test_inference_server.py`:
- End-to-end smoke: 1 iteration of training with `server` backend on CPU device; assert no crash, replay buffer populated.
- Equivalence: same seed, same initial weights, both backends → assert traversal samples match within numerical tolerance for at least 1 iteration. (May require `weight_sync_every = 1` and deterministic GPU forward; if exact match is fragile, accept distributional equivalence over 10 iterations.)
### Step 7: benchmarking
- Add `scripts/bench_inference_backend.py` (mirrors `scripts/profile_gpu_forward.py` style):
- Run 10 iterations on `default.yaml` with `inference_backend=local`.
- Run 10 iterations on `default_server.yaml`.
- Print iteration-mean and per-phase mean times.
- Document results in a new `docs/performance.md` experiment subsection (date-stamped).
## Risks and mitigations
- **Cython proxy incompatibility (Step 5).** Mitigation: pre-prototype with a 5-line Python script that imports `Traversal` and passes a stub callable in place of a `DeepCFRMLP`. If it works, the rest of the plan stands.
- **Weight sync staleness invalidates CFR.** Mitigation: default `weight_sync_every: 1` (every iteration). Only loosen after measuring.
- **GPU contention with eval.** Eval already runs on the trainer's device; the inference server adds another tenant on the same GPU. Mitigation: serialize eval and traversal phases (they already are sequential in the iteration loop). Document this in `inference_server.py`.
- **Deadlock on server crash.** Mitigation: trainer monitors server `process.is_alive()` between iterations and raises if dead. Workers waiting on events will be torn down with the pool.
- **Slot exhaustion under bursty load.** Mitigation: default `num_slots = 4 * num_workers * worker_chunk_size` to absorb bursts. Operators can tune.
- **Remote (weak GPU) regression.** If `traversal_seconds` regresses on remote, the operator can flip `inference_backend: local` and ship without it. The plan succeeds either way.
## What this plan explicitly does not do
- Does not implement async client (workers stay sync-blocking on slot events). Async client is a future optimization if IPC round-trip dominates after measurement.
- Does not implement `nogil` threading or Option C. See `docs/performance.md` for the deferred path.
- Does not optimize encoding (`policy_encoding_seconds`). If post-A measurements show encoding dominates, that is a separate work item.
## Out-of-scope follow-ups (do not start)
- Option C (vectorized traversal) — only re-evaluate if A's measurements show GPU forward is no longer on the critical path.
- TensorRT / `torch.compile` on the inference server's models — defer until A baseline numbers are collected.
- Batched encoding inside the server (compute encoding from raw game state on GPU) — only if `policy_encoding_seconds` becomes the new bottleneck.
@@ -0,0 +1,210 @@
# Plan: Cython Port of the Safe-Heuristic Bot Family
**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.
## Goal
Port `SafeHeuristicBot` (and its loose / strict parameterizations) from pure-Python `bots/heuristic.py` to Cython, consuming game state directly through the existing typed Cython `GameState` interface, so that the per-iteration evaluation cost dominated by `opponent_act_seconds` shrinks substantially without any change to bot decision behavior.
## Non-goals
- Do not change bot decision logic. The Cython port must be byte-identical to the Python implementation under fixed seeds for the full action sequence of every benchmarked seeded game.
- Do not port `RandomBot`, `PassiveDiscardBot`, or `NoisyPolicy`. Their per-iter cost is small (see `docs/performance.md` table: `random` 0.06s, `passive_discard` 0.00s, `noisy_safe` 0.57s).
- Do not change the public bot registry names or CLI surface.
- Do not touch the Cython traversal pipeline, encoding, training, or replay paths. This is a bot-only change.
- Do not redesign `LostCitiesPolicy` or the `PolicyInput` / `Snapshot` interfaces.
## Success criteria
1. **Equivalence (hard).** For a fixed corpus of seeded games (see "Equivalence strategy" below), the Cython implementation produces a byte-identical action id sequence to the Python implementation for each of `safe_heuristic`, `safe_heuristic_loose`, `safe_heuristic_strict`. Deviation in any single action fails CI.
2. **Per-opponent speedup.** On `configs/deep_cfr/default.yaml` evaluation, per-opponent `opponent_act_seconds` for each safe-heuristic variant decreases by **≥3×** versus the pre-port baseline measured on the same hardware. Concretely, target post-port values:
- `safe_heuristic`: ≤ 2.5s (from 7.57s)
- `safe_heuristic_loose`: ≤ 1.5s (from 4.63s)
- `safe_heuristic_strict`: ≤ 3.1s (from 9.22s)
3. **Eval wall-clock.** With the same `evaluation.games` and `evaluation.eval_every`, `eval/<variant>/elapsed_seconds` for each safe-heuristic variant decreases proportionally, and total `evaluation_seconds` for an iteration that runs the full opponent suite decreases meaningfully (the safe-heuristic opponents are the dominant tail per the performance doc).
4. **All existing tests pass**, including every `test_safe_heuristic_*` test in `tests/games/classic/test_bots.py`. New equivalence tests (see below) also pass.
5. **No behavioral regression in training.** A short training run on `default.yaml` with the Cython bots in eval reproduces the same eval-winrate trajectory (within seed noise) as the Python implementation over at least one full eval cadence.
## Independence claim
This work is **fully independent** of:
- the batched traversal inference server plan (`docs/plans/batched_traversal_inference_server.md`),
- the AMP wiring (`run.use_amp`),
- the `torch.compile` experiment branch (`experiments/torch-compile`).
It touches no traversal, training, or networks code. It can ship in parallel with any of those efforts. The eval pipeline already calls `bot.act(state)` per opponent turn; replacing that bot's implementation language is local to the bots package.
## Key files (current)
- `src/coolrl_lost_cities/games/classic/bots/heuristic.py` — pure-Python implementation. ~1160 lines. Fully read and analyzed. Contains:
- Module-level helpers `play_action`, `discard_action`, `draw_from_discard_action` (used by tests).
- `SafeHeuristicParams` (frozen dataclass, behavioral knobs).
- `DerivedHeuristicConfig` (frozen dataclass, derived constants).
- `derive_heuristic_config(config, params)``lru_cache(maxsize=64)` cached.
- `SafeHeuristicBot(LostCitiesPolicy)` — the policy itself, with `_act_card` / `_act_draw` and ~20 helper methods.
- `src/coolrl_lost_cities/games/classic/bots/base.py``legal_from_obs`, `first_legal`. The fallback path (non-`GameState` input) routes through these. The Cython port must preserve this fallback behavior unchanged.
- `src/coolrl_lost_cities/games/classic/bots/registry.py``BOT_REGISTRY`, `LOOSE_SAFE_HEURISTIC_PARAMS`, `STRICT_SAFE_HEURISTIC_PARAMS`. Currently constructs `SafeHeuristicBot` from the Python class.
- `src/coolrl_lost_cities/games/classic/game.pxd` — typed `GameState` declaration. The bot will consume this through Python attribute access (sufficient — see "Cython tactics" below) and via direct `cpdef` calls (`legal_card_mask`, `legal_draw_mask`, `score_diff`, `can_play_card`, etc.).
- `src/coolrl_lost_cities/games/classic/policy.py``LostCitiesPolicy` ABC and `PolicyInput` typedef.
- `setup.py` — Cython extension list.
- `tests/games/classic/test_bots.py` — existing safe-heuristic tests (`test_safe_heuristic_mirror_match_finishes`, `test_safe_heuristic_opponent_value_ignores_hidden_hand`, `test_safe_heuristic_started_expedition_value_ignores_invalid_lower_followup`, `test_safe_heuristic_draws_playable_discard_instead_of_deck`, `test_safe_heuristic_can_draw_discard_to_deny_opponent_when_losing`, `test_safe_heuristic_classic_self_play_opens_expeditions`, `test_safe_heuristic_avoids_opening_weak_fifth_color`, `test_safe_heuristic_prefers_followup_on_started_expedition`, `test_safe_heuristic_avoids_unopened_discard_draw_after_four_opens`).
## New files
- `src/coolrl_lost_cities/games/classic/bots/heuristic.pyx` — single Cython file containing the ported `SafeHeuristicBot` class plus `cdef` helpers. One file (not split per variant): the variants differ only in `SafeHeuristicParams` values, not in code.
- `src/coolrl_lost_cities/games/classic/bots/heuristic.pxd` — minimal typed declarations for the bot class and its hot helper signatures, so future extensions (or other Cython modules) can import typed entry points. Optional in Step 1; required in Step 3 if Cython call-site overhead from Python attribute lookup dominates measurement.
- `tests/games/classic/test_safe_heuristic_equivalence.py` — equivalence tests (see "Equivalence strategy").
## Files to touch
- `src/coolrl_lost_cities/games/classic/bots/heuristic.py`**keep as a thin wrapper** that re-exports `SafeHeuristicBot`, `SafeHeuristicParams`, `DerivedHeuristicConfig`, `derive_heuristic_config`, `play_action`, `discard_action`, `draw_from_discard_action`, `PLAY_OR_DISCARD_ACTIONS_PER_SLOT`, `DRAW_FROM_DECK_ACTION` from `heuristic.pyx`. This preserves all existing imports (`tests/games/classic/test_bots.py` imports `from coolrl_lost_cities.games.classic.bots.heuristic import draw_from_discard_action`). Do not delete the file. Do not duplicate logic.
- `src/coolrl_lost_cities/games/classic/bots/registry.py` — no functional change. Imports continue to resolve through the `heuristic.py` shim.
- `setup.py` — add the new extension:
```python
Extension(
"coolrl_lost_cities.games.classic.bots.heuristic",
["src/coolrl_lost_cities/games/classic/bots/heuristic.pyx"],
),
```
## Equivalence strategy
Behavior drift is the dominant risk. Equivalence tests gate the merge of every step.
### Test design
`tests/games/classic/test_safe_heuristic_equivalence.py` builds a fixed corpus of seeded `LostCitiesConfig` × seed pairs, plays full games to terminal, and asserts the Cython and Python bots emit byte-identical action ids at every turn.
Corpus:
- **Configs**: at minimum, the default `LostCitiesConfig()` from `tests/games/classic/helpers.py`. Add the small-tier config used by `test_safe_heuristic_started_expedition_value_ignores_invalid_lower_followup` if it exists, plus a tier where `n_handshakes == 0` (exercises the `state.config.n_handshakes <= 0` early return in `_best_handshake_play`).
- **Seeds**: `range(0, 50)` per config. Fifty full games per config gives broad coverage of decision branches without making CI slow.
- **Variants**: all three (`SafeHeuristicBot()`, `SafeHeuristicBot(LOOSE_SAFE_HEURISTIC_PARAMS)`, `SafeHeuristicBot(STRICT_SAFE_HEURISTIC_PARAMS)`). Loose and strict shift only the parameter dataclass, but each must be tested independently because their thresholds drive different branches in `_should_open_expedition`, `_best_handshake_play`, and the visible-draw support logic.
### Driver
Use a self-play harness (mirror match) that constructs two clone `GameState` instances per seed: one driven by the Python bot, one by the Cython bot. At each turn, both bots receive the same `GameState`. The test asserts:
```python
assert cython_action == python_action, (turn, seed, variant, state_summary)
```
If either bot diverges mid-game, the states diverge and subsequent comparisons are meaningless — so abort the comparison on first disagreement and report `(seed, variant, turn, action_py, action_cy)`.
### Branch coverage targets
The corpus must hit all decision paths. Each of these branches must be exercised by at least one (variant, seed, turn) tuple:
- `_act_card` → handshake play taken.
- `_act_card` → number play on a started expedition.
- `_act_card` → speculative open (`speculative_open=True`).
- `_act_card` → strong open (`strong_open=True`).
- `_act_card` → exceptional open (`exceptional_open=True`).
- `_act_card` → single-late open (`single_late_open=True`).
- `_act_card` → forced open path (no expeditions started, no normal play).
- `_act_card` → discard path with `unusable_discard_bonus`.
- `_act_card` → discard path with `discard_safety_bonus`.
- `_act_draw` → draw from deck.
- `_act_draw` → draw from discard (handshake top card).
- `_act_draw` → draw from discard (numeric top card).
- `_act_draw` → unopened-color discard penalty triggers (`opened_colors >= 4`).
- `_visible_draw_value` → exceptional-support short-circuit.
- Loose vs strict thresholds: at least one (seed, turn) where loose opens an expedition that strict declines, and one where strict's `late_open_block_threshold` blocks an opening that loose would take.
A coverage assertion at the end of the test sweep records which of the above branches fired (via a side-channel counter inside a debug build of the bot, or by analyzing the action stream); CI fails if any branch was not hit.
### Floating-point determinism
The bot uses Python `float` arithmetic with `max(...)` over `(value, action)` tuples. In Python, ties break by `action` ordering (since the value is the first element of the tuple). The Cython port **must use the same tie-break order**: stable `(value, action)` comparison, where ties prefer the lower-numbered action that came first in the iteration. Implementation: collect `(value, action)` pairs in the same iteration order as Python, then linear-scan for the maximum, replacing only on strict `>`. This matches Python's `max` semantics on a list-of-tuples.
All arithmetic must be `double` (matches Python `float`). Avoid `cdivision` floating-point divergence — there is essentially no division in the hot path, but `0.25 * number_sum` and similar must remain `double`.
## Cython-ization tactics
The hot loops are inside `_best_number_play`, `_best_discard`, `_visible_draw_value`, `_color_commitment`, `_public_color_commitment_for_opponent`, `_bonus_potential`, and `_opening_plan_value`. They iterate over hand cards (≤ ~10) and colors (≤ 6) per call, and the bot is called once per opponent turn (~30 turns/game × N games × N opponents).
### Tactics, in priority order
1. **Type the bot class as `cdef class SafeHeuristicBot`** with `cdef` helpers. Convert all `_<name>` methods to `cdef inline double _<name>(...)` (or `cdef int` for action ids) where possible. Keep `act` as `cpdef` so the Python registry can construct and call it.
2. **Type all hot locals** as `int` / `double` / `Py_ssize_t`. The dominant cost in the Python version is per-card list comprehensions building intermediate `list[Card]` objects; replace with explicit indexed loops over `state.hands[player]` and short fixed-size `cdef double` accumulators. No intermediate Python lists in the hot path.
3. **Cache `state.config` accessors once per `act` call.** `state.config.expedition_penalty`, `bonus_threshold`, `bonus_amount`, `n_colors`, `min_rank`, `max_rank`, `n_handshakes`, `n_ranks`, `hand_size`, `deck_size` — pull all of these into typed locals at the top of `act`.
4. **Cache `derive_heuristic_config` per `act`.** Already cached via `lru_cache`, but each lookup re-hashes the config. Pull the resulting `DerivedHeuristicConfig` into a typed local; access its fields once.
5. **Avoid Python-object operations in inner loops.** `Card` objects expose `color`, `rank`, `is_handshake`, `numeric_value(min_rank)`. In the Cython port, when the Python `Card` object is unavoidable (it is part of the public `GameState.hands` shape), bind its three relevant attributes to typed locals at the top of each loop iteration. Do not call `card.numeric_value(min_rank)` inside conditions repeatedly — compute it once per card per iteration.
6. **Direct typed access where possible.** `GameState.score_diff(player)` is `cpdef int` and `can_play_card` is a Python wrapper around `can_play_encoded_card` (`cpdef bint`). Where the bot calls `state.can_play_card(player, card)` for a `Card` object, the Cython port can encode the card once (`color * cards_per_color + (rank - min_rank)` per game.pyx encoding) and call `state.can_play_encoded_card(player, encoded)` directly. **Verify the encoding formula against `game.pyx` `_encode_card` before relying on it; otherwise keep the Python `can_play_card` call and accept the small overhead.**
7. **`cpdef list legal_card_mask(self)` returns a Python list of `bool`.** Type the binding as `list legal` in the Cython bot, and index it with typed `Py_ssize_t`. If profiling shows mask access dominates, convert callers to use `unified_legal_mask_np()` and read it as a typed numpy view — but only after Step 4 measurements show this is needed.
8. **Inline tiny helpers.** `_num`, `_late_penalty`, `_new_color_open_penalty` are one-liners. Inline them with `cdef inline`.
9. **Compiler directives.** Use the same directives as `setup.py` already applies project-wide: `boundscheck=False`, `wraparound=False`, `cdivision=True`, `initializedcheck=False`. The bot file inherits these from `setup.py`'s `compiler_directives` block — do not need to override per-file.
### What NOT to optimize
- Do not rewrite `derive_heuristic_config` to drop its `lru_cache`. It is called once per `act` and the cache hit is fast.
- Do not replace `GameState.hands[player]` with raw int-array access. The `Card` Python object boundary is the public API, and crossing it is what `_card_obj` already costs in Cython. Keep the boundary; just don't re-cross it inside tight loops.
- Do not memoize across calls. The bot is stateless; each `act` call gets a fresh state.
## Risks and mitigations
- **Behavior drift introduced silently.** Mitigation: equivalence tests (above) run in CI on every PR. They are the gate. Benchmarks are not run until equivalence is green.
- **Floating-point ordering difference between Python and Cython.** Mitigation: explicit `(value, action)` linear-scan max with strict `>` comparison; identical iteration order. Equivalence tests would surface any divergence.
- **`Card` object identity vs equality.** The Python code uses `is not` comparisons (`if other is not card`, `if followup is not card`, `if card is not exclude_card`). Mitigation: in Cython, preserve `is`-comparison semantics by using object-pointer identity (`PyObject*` compare) or by passing the slot index instead of the Card. Slot-index passing is preferred — it sidesteps identity altogether.
- **`max(candidates)[1]` Python tuple semantics.** Mitigation: documented above. Strict `>` linear scan.
- **Cached `derive_heuristic_config` shared between Python and Cython invocations.** Mitigation: the Cython port must import `derive_heuristic_config` from the same module path so the LRU cache is shared (or rebuild a Cython-side cache keyed identically). Easiest: keep `derive_heuristic_config` and `DerivedHeuristicConfig` in the same `.pyx` and re-export from the `.py` shim.
- **Build-system regression.** Adding an extension that fails to compile in CI on a system without the right toolchain. Mitigation: the project already builds three `.pyx` files (`game`, `cfr_math`, `encoding`, `traversal`); the toolchain is established. Add the new extension and run `uv run python -c "import coolrl_lost_cities.games.classic.bots.heuristic"` locally before pushing.
- **Performance regression on small tiers.** The bot may not actually be the bottleneck in micro-config evaluation. Mitigation: benchmarks are run on `default.yaml` (the contract surface), not on smoke configs.
## Implementation steps (each independently mergeable)
### Step 1: scaffolding (no behavior change)
- Add `bots/heuristic.pyx` containing the entire current `bots/heuristic.py` body verbatim, with no `cdef` or typing changes — just renamed.
- Convert `bots/heuristic.py` to a re-export shim (`from .heuristic_impl import *` style) — but to avoid a circular import via the `.pyx` module name, keep both as `heuristic.{py,pyx}` is not workable. Instead, name the new file `heuristic_cy.pyx` (Cython compiles to `heuristic_cy`), and have `heuristic.py` import the public symbols from `heuristic_cy`. Update `setup.py` accordingly.
- **Decision point:** if Cython supports a `.pyx` shadowing a `.py` in the same package, prefer that for a cleaner import path. Otherwise use the `heuristic_cy` naming. Resolve at Step 1 implementation time and update this plan inline.
- Run all existing tests to confirm zero behavior change.
- Merge.
### Step 2: type the base `SafeHeuristicBot`
- Convert `SafeHeuristicBot` to `cdef class`. Add typed locals to the hot helpers listed under "Cython tactics". Keep the data classes (`SafeHeuristicParams`, `DerivedHeuristicConfig`) as Python frozen dataclasses — they are not hot.
- Run the equivalence test suite (built in Step 4 — but for this step, run the existing `test_safe_heuristic_*` tests as a proxy gate; the full equivalence suite lands in Step 4).
- Merge.
### Step 3: type `cdef inline` helpers and inner loops
- Inline `_num`, `_late_penalty`, `_new_color_open_penalty`. Convert `_color_commitment`, `_public_color_commitment_for_opponent`, `_bonus_potential`, `_opening_plan_value`, `_visible_draw_value`, `_visible_open_support_value`, `_visible_number_can_help_open` to typed `cdef double` / `cdef bint` helpers.
- Replace list comprehensions with indexed loops over `state.hands[player]`, accumulating into typed scalars.
- Run equivalence suite. Merge only if green.
### Step 4: equivalence test suite
- Add `tests/games/classic/test_safe_heuristic_equivalence.py` with the corpus and branch-coverage assertions described under "Equivalence strategy".
- Wire into `pytest -q tests/games/classic/test_safe_heuristic_equivalence.py`. Confirm the test passes against the Step 3 build. If it surfaces drift, fix Step 3 before continuing.
- Merge.
- **Note on ordering vs Step 2/3:** ideally Step 4 lands before Step 2 so equivalence is a green gate the entire time. Recommended order: 1 → 4 (against the verbatim port, must be a no-op pass) → 2 → 3.
### Step 5: variant validation
- Loose and strict differ only by params, but exercise their threshold differences explicitly. Confirm `LOOSE_SAFE_HEURISTIC_PARAMS` and `STRICT_SAFE_HEURISTIC_PARAMS` in `registry.py` flow through the Cython implementation unchanged.
- Confirm `NoisyPolicy(SafeHeuristicBot(), RandomBot(seed))` still works (`NoisyPolicy` lives in `registry.py` and wraps the bot; the wrapper does not need porting).
- Add a small targeted test that constructs all three variants and runs one full game each.
- Merge.
### Step 6: benchmark and documentation
- Benchmark protocol: run `configs/deep_cfr/default.yaml` evaluation (full opponent suite, default `evaluation.games`) once on the pre-port commit and once on the post-port commit, on the same `home` machine, with no other GPU load. Record:
- `eval/safe_heuristic/elapsed_seconds`, `eval/safe_heuristic/opponent_act_seconds`,
- `eval/safe_heuristic_loose/elapsed_seconds`, `eval/safe_heuristic_loose/opponent_act_seconds`,
- `eval/safe_heuristic_strict/elapsed_seconds`, `eval/safe_heuristic_strict/opponent_act_seconds`,
- `evaluation_seconds` for the iteration.
- Confirm success criteria 2 and 3.
- Add a date-stamped subsection to `docs/performance.md` recording results, methodology, and any branch-coverage gaps surfaced during equivalence testing.
- Merge.
## Out-of-scope follow-ups (do not start)
- Porting `RandomBot`, `PassiveDiscardBot`, or `NoisyPolicy` to Cython. Their measured cost is small. Re-evaluate only if a future profile shows them on the critical path.
- Replacing `Card` Python objects with raw int encoding inside the bot interface. This would require touching `LostCitiesPolicy` and `PolicyInput` more broadly, expanding scope.
- Caching across `act` calls (e.g., remembering opened-color counts). The bot is stateless by design; introducing state risks correctness for marginal speedup.
- TensorRT / `torch.compile` on policy networks during eval. Tracked separately under the inference-server plan and `docs/performance.md` § "Post-A Optimization Calculus".
- Vectorizing eval to compute many game states' bot actions simultaneously. Substantial rework of the eval driver; out of scope here.