Add brains/sovereign-v3/train_notes.md

brains/sovereign-v3/train_notes.md

@@ -1,29 +1,27 @@
-# SOVEREIGN v3 — Training Doctrine (~272M dense)
+# SOVEREIGN v3 — Training Doctrine (≈290M dense + MCTS planning)
 
 Exact recipe for **2× T4 (16 GB)**, shared shard format
-`rl/shards_v2/*.npz`. Memory ≈9.4 GB/GPU at bs 16 + fp16 (see
-ARCHITECTURE.md §3). Everything is implemented by `model.py` +
-`data_adapters.py`.
+`rl/shards_v2/*.npz`. Memory ≈9.5 GB/GPU at bs 16 + fp16 (ARCHITECTURE.md
+§3). Everything below is implemented: `model.py`, `data_adapters.py`,
+`search.py`, `league.py`, `distill.py`.
 
 ---
 
-## Stage A — Return-conditioned filtered BC + dynamics (offline, ~2–3 epochs)
+## Stage A — Return-conditioned filtered BC + world model (offline, ~2–3 epochs)
 
 ### A.1 Sample preparation (`ShardPrepper`)
 
-1. **Returns / RTG**: per-episode `R_t = r_t + γR_{t+1}`, γ=0.995;
-   +2 survive / +5 win terminal bonuses; z-scored per shard. **rtg = R̄** is
-   the conditioning input; `rtg_next` shifts within the episode (last
-   frame repeats).
+1. **Returns / RTG**: per-episode `R_t = r_t + γR_{t+1}`, γ=0.995; +2
+   survive / +5 win terminal bonuses; z-scored per shard. `rtg` conditions
+   every head; `rtg_next` shifts within episodes.
 2. **Weights** `w = clip(exp(R̄/τ), 0.2, 8)` τ=1; ×4 survivor episodes;
-   ×2 ±2-frame kill windows; mean-normalized. (AWR soft filtering.)
-3. **Free-label targets**: `lab64`, `lab64_next` (majority-vote 64×64),
-   `threat`/`expand` intent mirrors, `gate_mask`/`gate_valid`,
-   `win_lab` (survived=1), `nums_next`.
+   ×2 ±2-frame kill windows; mean-normalized (AWR soft filtering).
+3. **Free-label targets**: `lab64`, `lab64_next`, `threat`/`expand`
+   intent mirrors, `gate_mask`/`gate_valid`, `win_lab`, `nums_next`,
+   `rgb_next` (consecutive frames), `reward`.
 4. **Arena filtering**: drop >85%-water episodes.
-5. Frame stacks / sequence windows: opt-in (`stack=`, `build_sequences`).
 
-### A.2 Loss (exact)
+### A.2 Policy loss (exact)
 
 ```
 L = w · [ λk CE(kind; ls=0.05)
@@ -33,81 +31,89 @@ L = w · [ λk CE(kind; ls=0.05)
        + λw CE(win_logits, win_lab; ls=0.05)
        + λg BCE(gate, gate_mask)·1[valid]
        + λs CE(seg64, lab64)
-       + λd CE(next_seg, lab64_next)              # dynamics
-       + λt BCE(threat, threat) + λx BCE(expand, expand)   # intent mirrors
+       + λd CE(next_seg, lab64_next)
+       + λt BCE(threat) + λx BCE(expand)
        + λe smoothL1(econ, cat(nums_next, rtg_next))
        − λent (H(kind)+H(cell)+H(pct)) ]
 λ = (0.5, 1.0, 0.25, 0.1, 0.1, 0.5, 0.5, 0.4, 0.2, 0.2, 0.05, 0.003)
 ```
 
-**Value option (IQL-expectile)**: for the value head only, replace Huber
-with the τ=0.9 expectile loss `|τ−1[δ<0]|·δ²` — more robust to the weak
-farmer buffer. Both are supported; ablation #6 decides.
+Value option: τ=0.9 expectile instead of Huber (IQL-style robustness on
+the weak buffer).
 
-### A.3 Optimizer & schedule
+### A.3 World-model loss (the new thing — trains the search block)
 
-- AdamW lr **3e-4 → 3e-5** cosine, **300-step warmup** (cortex), wd 1e-4
-  (Linear only), grad clip 5.0, **AMP fp16**, bs **16/GPU × accum 2**.
-- 2–3 epochs. Smoke test first: `python model.py` (prints the 272M
-  breakdown), `python data_adapters.py` (adapter shapes).
+```
+z_t     = model.encode(rgb_t, nums_t, rtg_t)            # stopgrad target
+ẑ_{t+1} = model.dynamics(z_t, kind_t, cell_t, pct_t)
+prior̂, v̂, win̂ = model.prediction(z_t)
+z_{t+1} = model.encode(rgb_next_t, ...).detach()
+
+L_wm = λz MSE(ẑ_{t+1}, z_{t+1})               # latent transition
+     + λr Huber(r̂, reward_t)                  # reward model
+     + λp KL(softmax(prior̂), softmax(cell_logits_t).detach())   # prior consistency
+     + λv̂ Huber(v̂, R̄_t) + λŵ CE(win̂, win_lab_t)
+λ_wm = (0.5, 0.1, 0.3, 0.1, 0.1)
+```
+
+This is MuZero-lite: representation + dynamics + prediction trained
+entirely on labels the shards already contain (consecutive frames +
+reward). Implement with `encode()/prediction()/dynamics()` + a small loop
+over the loader (10 lines; both loss blocks share the same batch).
+
+### A.4 Optimizer & schedule
+
+- AdamW lr **3e-4 → 3e-5** cosine, **300-step warmup** (cortex), wd 1e-4,
+  grad clip 5.0, **AMP fp16**, bs **16/GPU × accum 2**. 2–3 epochs.
+- Smoke tests first: `python model.py`, `python data_adapters.py`,
+  `python search.py`, `python league.py`, `python distill.py`.
 
 ---
 
-## Stage B — RPO: return-conditioned PPO (online, in sim)
+## Stage B — RPO + AlphaZero-lite policy improvement (online, in sim)
 
 ### B.1 Rollouts
 
 - `model.act(..., rtg=rtg_actual)` per tick; store `logprob`, `value`,
-  `win_prob`. **8 envs/GPU × 256 ticks** on CPU workers (2×T4 ⇒ ~2048
-  samples/round).
-- Memory mode (optional): carry `state` through each episode
-  (truncated-BPTT); the env wrapper sketch:
-
-```python
-# env wrapper sketch (mem mode)
-s = None
-for step in range(ep_len):
-    rgb_t, nums_t = env.obs()
-    a = model.act(rgb_t, nums_t, rtg=rtg_est, state=s)
-    s = a["state"]            # (1, B, 768)
-    obs, r, done = env.step(a["kind"], a["cell"], a["pct"])
-```
+  `win_prob`. **8 envs/GPU × 256 ticks** on CPU workers.
+- **Search rollouts (the planning loop)**: every 4th action, call
+  `Searcher(model, sims=64, max_depth=12).search(rgb, nums, rtg, mask)`
+  and execute the search-chosen action. Store the **visit distribution**
+  `visit_probs` alongside the transition (single-sample search; batch over
+  envs by running search per env — 64 sims × 12 depth ≈ 20–40 ms on T4,
+  well inside the 560 ms tick).
+
+### B.2 RPO loss (unchanged core)
+
+`A_t = R_t − V(s_t, R_t)` (rtg-consistent critic), GAE(0.995, 0.95),
+z-normalize, clip 0.2 on the action logprob, 0.5×clipped Huber value,
+win-head CE on outcomes, entropy 0.01→0.002. 4 epochs/round, minibatch
+64/GPU, AdamW 3e-4 + cosine tail, grad clip 5.0, **fp16 + fp32 master**.
 
-### B.2 RPO specifics
+### B.3 Search-target distillation (AlphaZero's policy improvement)
 
-- **RTG-consistent advantages**: `A_t = R_t − V(s_t, R_t)` (value is
-  rtg-conditioned — that is the whole point; the critic never has to
-  guess the aspiration).
-- GAE(γ=0.995, λ=0.95) over `A_t`; z-normalize; clip ε=0.2 on the full
-  action logprob.
-- Value loss: 0.5×clipped Huber (double clip ±0.2 in return space) +
-  **win head CE on episode outcomes** + optional cell-value auxiliary
-  (Huber vs per-cell utility bootstrap — PPO-friendly spatial critic).
-- 4 epochs/round, minibatch 64/GPU, AdamW lr 3e-4 + cosine tail, grad
-  clip 5.0, **AMP fp16 + fp32 master**.
-- **KL-to-pretrain guard β=0.1** (frozen Stage-A prior); roll-back if mean
-  KL > 0.15 for 3 rounds (β→0.02). ICLR 2025 evidence: pretrain-init +
-  guard, no offline replay retention — do NOT mix shards into the PPO
-  batches.
+On rollout states where search ran, add:
 
-### B.3 Aspiration at play time
+```
+L_search = CE(kind_logits, kind_from_search)
+         + KL(softmax(cell_logits[·,kind]), visit_probs_kind)
+         + Huber(value, q_search)
+λ_search = 0.5  (weighted by search-depth, first N=500k transitions)
+```
 
-Eval/deploy with `rtg = +1.5` z-units (above-survival ambition; tune over
-{0.5, 1.0, 1.5, 2.0} — ablation #1). The BC stage never saw rtg>2σ, so
-stay inside the trained range.
+This lifts the policy toward the planner's improved distribution —
+exactly the loop AlphaZero proved, scaled to a 3×1024 action space.
 
-### B.4 Reward shaping (flag the sim edit)
+### B.4 League-lite self-play
 
-`+0.004/tick alive, +0.4/kill, +2 survive timeout, +8 win`. Keep
-intermediates small; the win is the judged objective.
+Use `league.py`: after each eval round, `update_elo` vs the pool; when the
+current agent beats the pool (`maybe_promote`, threshold 0.55), freeze a
+snapshot. Rollouts run vs **7 bots + 2 frozen past selves** picked by
+`pick_opponents(2)` (alternating weakest/strongest — curriculum pressure).
 
-### B.5 League-lite (fictitious self-play, cheap)
+### B.5 KL guard
 
-Keep the top-3 checkpoints by win-rate. Every 10 rounds, evaluate the
-current policy against the pool; promote if it wins >55%. Once promoted,
-the sim runs our policy vs 7 bots + 2 frozen past selves — the only
-"league" a 2×T4 budget can afford, and it directly targets the weakness
-of the last-survivor objective (myopic bot-exploit policies).
+β=0.1 KL-to-Stage-A-prior; roll-back if mean KL > 0.15 for 3 rounds.
 
 ---
 
@@ -115,11 +121,12 @@ of the last-survivor objective (myopic bot-exploit policies).
 
 - 100+ episodes vs **9 mixed bots**, fixed seeds; metric **win-rate +
   mean rank**; save best on rank; report binomial CI.
-- **Wake-sleep distill**: after PPO plateaus, distill the 272M policy into
-  a ~30M dense student (same heads, distilled heat/kind/pct/value with KL;
-  `return_all=True` provides all teacher targets). Use the student for
-  2× faster rollouts and richer league self-play data; the big teacher
-  trains from it. Repeat (cap 3 cycles).
+- **Wake-sleep** (`distill.py`): when PPO plateaus, distill the teacher
+  into the ~35M student (`small_config()`), including any stored
+  search targets. The student rolls out 2–3× faster → richer league data
+  → teacher trains again. Cap 3 cycles.
+- **Search eval**: measure win-rate with search ON vs OFF (the pure
+  planning ablation — it isolates how much "thinking" is worth).
 
 ## Curriculum
 
@@ -131,31 +138,33 @@ Epoch 1 of Stage A on early-game frames (tick_frac < 0.35) → full data →
 | # | Question | Change | Resolves |
 |---|---|---|---|
 | 1 | Aspiration rtg at play time | {0.5, 1.0, 1.5, 2.0} | the RPO knob |
-| 2 | Return conditioning itself | rtg≡0 (frozen) vs conditioned | Bet 1's real value |
+| 2 | Return conditioning itself | rtg≡0 vs conditioned | Bet 1's value |
 | 3 | Intent mirrors | λt=λx=0 vs default | opponent modeling |
 | 4 | Dynamics head | λd=0 vs default | world-model value |
-| 5 | Memory | GRU off vs on (seq BPTT) | temporal credit |
-| 6 | Value objective | Huber vs τ=0.9 expectile | robustness |
-| 7 | Win head | on/off in Stage B loss | endgame behavior |
+| 5 | **Search on/off** | Searcher vs policy-only at eval | how much MCTS is worth |
+| 6 | **Search sims** | {16, 64, 128, 256} | compute vs strength |
+| 7 | Value objective | Huber vs τ=0.9 expectile | robustness |
 | 8 | League-lite | pool vs no pool | self-play payoff |
 | 9 | Wake-sleep | 0 vs 3 distill cycles | scaling loop |
+| 10 | Search-target distillation | λ_search ∈ {0, 0.5, 1.0} | policy-improvement op |
 
 ## Expected compute (2× T4, fp16)
 
-- Stage A: ~8–14 h for the current GB-scale shard set (272M, bs 16/accum 2).
-- Stage B: ~2–3 updates/min at minibatch 128 → one full experiment per
-  ~1–2 days. That is the honest cost of the mandate; the sanctioned trims
-  (in order): `head_hidden` 3072→2048, `cortex_layers` 6→4, `enc_ch` top
-  832→640. Each keeps the total >200M.
+- Stage A: ~10–16 h (≈290M incl. world-model losses).
+- Stage B: ~2–3 updates/min; search on 8 envs every 4th tick adds
+  ~0.3 s/env/step of T4 time. One full experiment ≈ 1–2 days.
+- Sanctioned trims (in order): `head_hidden` 3072→2048,
+  `cortex_layers` 6→4, `enc_ch` top 832→640, `sims` 64→32.
 
 ## Risk register
 
 | Risk | Mitigation |
 |---|---|
-| fp16 NaN in attention | fp32 master, grad clip 5, warmup; fall back fp32 for Stage A |
-| Unlearning the BC prior | KL guard + roll-back (B.2); no offline replay retention |
-| Aspiration OOD (rtg>2σ) | clamp rtg to ±2.5; ablation #1 |
-| Threat/expand labels are imbalanced | BCE weights 8:1 positive:negative |
-| GRU rollout complexity | default off; env sketch included |
-| Activation memory at bs 16 | `grad_ckpt=True` halves the heavy stages |
-| Color-identity prior (red=me) | sim-specific; documented |
+| fp16 NaN in attention | fp32 master, grad clip 5, warmup |
+| Unlearning the BC prior | KL guard + roll-back (B.2) |
+| Search too slow in rollouts | sims 32, top_k 8, depth 8; distill student |
+| World model collapses (z→ẑ≈z) | residual dynamics + λz; monitor MSE curve |
+| Threat/expand label imbalance | BCE weights 8:1 positive:negative |
+| Aspiration OOD (rtg>2σ) | clamp rtg to ±2.5 |
+| League pool drifts stale | promotion threshold 0.55 + periodic re-eval |
+| GRU rollout complexity | default off; env sketch in README |