Add brains/sovereign-nano/nano_notes.md

brains/sovereign-nano/nano_notes.md

@@ -0,0 +1,147 @@
+# SOVEREIGN-nano — notes (no marketing)
+
+≤5M-param policy for territorial.io, built as the **fast half** of the
+SOVEREIGN system. Same contract as the v3 teacher; distilled from it;
+usable as a rollout engine and a CPU sanity baseline.
+
+---
+
+## 1. Why this exists
+
+| Role | Detail |
+|---|---|
+| Fast rollout engine | ~90× cheaper than the 290M teacher → the PPO/league/wake-sleep loop gets ~10× more sim experience per GPU-hour (the bottleneck of the whole project is rollout throughput, not model capacity). |
+| Distillation student | The teacher's wake-sleep loop (distill.py) distills into this. Contract is identical, so KL on the (B,3,1024) heatmaps is a direct op. |
+| CPU sanity baseline | ≤5M params, small activations — smoke-trainable on Kaggle CPU without touching the GPU quota. |
+
+## 2. Architecture
+
+| Module | Shapes | ~Params |
+|---|---|---|
+| stem + coords | 5→16 @128² | 1K |
+| e1–e4 (spine) | 16/32/64/128/192 @64²..8², 1–2 ResBlocks | ~1.7M |
+| FiLM | 9→192 scale/shift on the bottleneck | 3K |
+| decoder d0–d4 | 192→32 @16²..256² with e3/e2/e1 skips | ~810K |
+| heat mixer + SHARED tower | 32→48→48→48→3 @256² | ~150K (one tower emits all 3 kinds — 3× cheaper than per-kind towers) |
+| gate / seg / next_seg / threat / expand | 32→{3,4,4,1,1} @64² | 1K |
+| pooled-MLP global (cortex replacement) | 201→512 → kind 3 / value 1 / win 2 / econ 9 | ~640K |
+| pct Beta | (32+9+3)→128→128→2 | ~22K |
+| cell-value | 64→32→1 @32² | 2K |
+| GRU memory | 201→96→192 | ~105K |
+| **TOTAL** | printed by the smoke test | **≈3.3M — under the 5M cap** |
+
+Kept from v3 (scaled down): return-conditioning (rtg through FiLM + all
+heads), gate head (mask-free legality), intent mirrors (threat/expand),
+next_seg dynamics aux, win head, econ forecaster, Beta pct, cell-value.
+
+Dropped, honestly: the 77M transformer cortex → pooled-MLP global head;
+the MCTS planning block → OUT; SE channel attention → OUT; the 3072-wide
+MLPs → 512/256. Nano does not search — the teacher does, and the
+visit-target distillation happens on the teacher's side.
+
+## 3. Memory & speed (measured math, not vibes)
+
+Corrected FLOP derivation (hook-counted by `opt.py count_flops`; the
+first draft overestimated because the per-kind towers ran 3× the 256²
+convs — that waste was removed in the speed pass):
+
+| Block | MACs/frame |
+|---|---|
+| spine (stem + e1–e4) | ~0.65 G |
+| decoder (d0–d4) | ~1.7 G |
+| heat (mixer 48ch + SHARED tower @256²) | ~3.6 G (was ~10.4 G with per-kind towers) |
+| **Total** | **~6 G MACs ≈ 12 GFLOPs** |
+| with `heat_res=64` | heat ~0.34 G → total ≈ 2.7 G MACs ≈ 5.4 GFLOPs |
+
+| Item | Estimate (honest ±2× band) |
+|---|---|
+| Weights (fp32) | ~13 MB |
+| Grads + Adam | ~40 MB |
+| Activations (bs 16, 256², fp32) | ~300–500 MB peak |
+| T4 fp16 forward | ~3–8 ms · act ~5–12 ms |
+| CPU (8-core) forward | ~60–200 ms · with heat_res=64 ~30–90 ms |
+| CPU + quantize_cpu.py | Linear/GRU int8 only — convs stay fp32 (stated) |
+
+`bench.py` turns every one of these numbers into a measurement — run it
+on the pipeline machine and trust its output over this table. Against
+the 560 ms sim tick, T4 rollouts are ~50–100× cheaper than the teacher's
+forward+MCTS; CPU rollouts are viable with `heat_res`.
+
+## 3b. The speed stack (opt.py / bench.py / quantize_cpu.py)
+
+| Lever | How | Honest expected gain | Caveat |
+|---|---|---|---|
+| Shared heat tower (already in the model) | one tower emits 3 kind channels | ~3× on the heat block (was 65% of all FLOPs) | kept full-res 3×3 cross-talk |
+| `opt.fold_bn` | BN folded into convs at eval | ~10–20% + less sync | keep unfolded weights for training; smoke asserts exactness |
+| `torch.compile` | full-graph compile | 1.3–2× | one-time compile cost per shape |
+| fp16 autocast | T4-native half precision | 1.5–2× on T4 | rollouts only — training uses fp32 master |
+| `heat_res=64` | tower at 64², logits upsampled | ~2× total | blunts 256² heat detail — ablate before trusting |
+| channels_last (NHWC) | `opt.to_channels_last` + auto NCHW conversion at interpolate/grid_sample | ~5–15% | safe by construction; bench verifies |
+| `quantize_cpu.py` | dynamic int8 Linear/GRU | partial CPU win | convs stay fp32 — stated |
+| `bench.py` | measures all of the above | — | claims become measurements |
+
+## 4. The real shard schema (why v3 crashed, fixed in prep_shard)
+
+```
+rgb    (354,256,256,3) uint8   frames every 2 ticks
+lab    (354,256,256)   uint8   0 water / 1 neutral / 2 me / 3 enemy
+nums   (354,8) float32         [bal_log, me_frac, red, e1, e2, tick_frac, income, kills]
+kind   (354,) int64            0 expand / 1 attack / 2 bank
+cell   (354,) int64            32x32 grid idx 0..1023
+pct    (354,) float32          troops fraction
+logp   (354,) float32          TRUE behavior logprob (old shards -2.0)
+reward (708,) float32          PER TICK (2× frames), -1.3 death terminal
+alive  (4,)   int64            PER EPISODE
+lens   (4,)   int64            PER EPISODE TICK counts (frames = lens//2)
+```
+
+`prep_shard()` reads exactly this: per-frame reward = sum of the frame's
+2 ticks, per-episode returns over `lens//2` frames, episode-survival
+broadcast for win labels and weights. The smoke asserts
+`len(reward) == lens.sum()` and `sum(lens//2) == len(rgb)`.
+
+## 5. The six v3 review bugs — all fixed, all asserted
+
+| # | Bug | Fix in nano |
+|---|---|---|
+| 1 | Adapter assumed alive/reward/lens are per-frame | `prep_shard` reads the real schema (above); smoke asserts it |
+| 2 | `CE((B,3,g²), (B,))` crashes on the cell head | take the kind row first: `cell_logits[arange(B), kind]` |
+| 3 | seg/next_seg CE (B,H,W) silently broadcasts | `.mean(dim=(1,2))` per sample |
+| 4 | act() logprob/entropy broadcast to (B,B) | Beta params squeezed to (B,); **shape-asserted in act()** |
+| 5 | `F.categorical_entropy` does not exist | manual `-(p·log p).sum()` (`cat_entropy`) |
+| 6 | `torch.from_numpy` on 0-dim scalars | all prep outputs are ≥1-dim tensors |
+
+The smoke test runs forward + act(B=4) + full stage-A loss + backward on
+the **real** shard (via `$HF_TOKEN`) and prints which path it ran. It
+falls back to a schema-EXACT synthetic replica only when the shard cannot
+be fetched, and says so loudly.
+
+## 6. Honest limitations
+
+1. **No transformer cortex.** Pooled-MLP global heads cannot do the
+   multi-blob relational reasoning the v3 cortex does. Nano inherits
+   strategy through distillation, not by construction.
+2. **No search.** Nano is a reactive policy; the planning lives in the
+   teacher (search.py). Do not expect teacher-grade decision quality
+   from nano alone.
+3. **Sub-tick reward alignment is approximate.** Reward is per-tick and
+   frames are every 2 ticks; prep sums each frame's 2 ticks and
+   accumulates returns at the frame rate (γ=0.995 per frame). This is
+   the honest approximation — an exact sub-tick alignment needs the
+   pipeline's tick-level recorder.
+4. **Single-frame input.** No frame stacking (contract). Temporal signal
+   comes from the gate and the tiny GRU only.
+5. **5M cap forces choices.** Every head is narrower than ideal; the
+   width is justified by the rollout-engine role, not by accuracy.
+6. **CPU smoke at bs 16 only.** GPU-scale training of nano is untested
+   here; the recipe (AdamW 3e-4, same stage-A loss) is documented, not
+   benchmarked.
+
+## 7. Integration notes
+
+- Import: `from model_nano import Nano, NanoConfig, stage_a_loss`.
+- distill.py: replace `small_config()` with `NanoConfig` and import Nano.
+- Rollouts: `a = net.act(rgb, nums, rtg=rtg)` per tick; carry `a["state"]`
+  through each episode when `use_mem=True`.
+- CPU sanity: bs 8, lr 3e-4, 100 steps on the prep'd real shard, watch
+  `terms` fall.