Add brains/sovereign-v3/model.py

brains/sovereign-v3/model.py

@@ -1,6 +1,6 @@
 """
-SOVEREIGN v3 — ~272M-param return-conditioned spatial policy for
-territorial.io (last-survivor elimination).
+SOVEREIGN v3 — ~290M-param return-conditioned spatial policy with MCTS
+planning for territorial.io (last-survivor elimination).
 
 = DESIGN SUMMARY (full reasoning: ARCHITECTURE.md, training: train_notes.md)
 Size mandate: >=200M, <400M (user-specified). Four modules:
@@ -115,11 +115,21 @@ class SovereignConfig:
     aux_seg_classes: int = 4      # water / neutral / me / enemy
     dropout_mlp: float = 0.1
     grad_ckpt: bool = False       # checkpoint the 4 biggest stages
+    # -- planning / search block (MuZero-lite) ------------------------------
+    use_search: bool = True       # latent encode/predict/dynamics for MCTS
+    lat_c: int = 32               # compressed spatial channels in the latent
+    lat_spatial: int = 8          # latent map size (8x8)
+    dyn_hidden: int = 1024        # dynamics MLP hidden width
+    act_embed: int = 512          # action embedding width
 
     @property
     def ctx_full_dim(self) -> int:
         return self.ctx_dim + self.rtg_dim
 
+    @property
+    def latent_dim(self) -> int:
+        return self.cortex_d + self.lat_c * self.lat_spatial * self.lat_spatial
+
 
 # --------------------------------------------------------------------------
 # Blocks
@@ -431,6 +441,32 @@ class Sovereign(nn.Module):
             self.mem = nn.GRU(pool_dim, c.mem_hidden, 1, batch_first=True)
             self.mem_proj = nn.Linear(c.mem_hidden, c.cortex_d)
 
+        # -- planning block (MuZero-lite): latent encode/predict/dynamics ---
+        # Trained on CONSECUTIVE shard frames (rgb_t, a_t -> rgb_{t+1}):
+        #   encode   -> z  (B, latent_dim)
+        #   prediction(z) -> prior logits over (kind,cell), value, win
+        #   dynamics(z, a) -> z_next, reward_hat
+        # Drives MCTS at decision time (see search.py). ~17M params.
+        self.planning = None
+        if c.use_search:
+            m = nn.Module()
+            m.lat_proj = nn.Sequential(
+                nn.Conv2d(c.enc_ch[-1], c.lat_c, 1, bias=False),
+                nn.BatchNorm2d(c.lat_c), nn.SiLU())
+            m.a_emb_kind = nn.Embedding(c.num_kinds, 64)
+            m.a_emb_cell = nn.Embedding(c.grid_default * c.grid_default, 256)
+            m.a_emb_pct = nn.Linear(1, 128)
+            m.a_emb_out = nn.Linear(64 + 256 + 128, c.act_embed)
+            m.dyn_mlp = nn.Sequential(
+                nn.Linear(c.latent_dim + c.act_embed, c.dyn_hidden),
+                nn.SiLU(), nn.Linear(c.dyn_hidden, c.latent_dim))
+            m.reward_head = nn.Linear(c.latent_dim, 1)
+            m.pred_prior = nn.Linear(
+                c.latent_dim, c.num_kinds * c.grid_default * c.grid_default)
+            m.pred_value = nn.Linear(c.latent_dim, 1)
+            m.pred_win = nn.Linear(c.latent_dim, 2)
+            self.planning = m
+
         self.apply(_init_weights)
 
     # -- helpers ------------------------------------------------------------
@@ -468,6 +504,64 @@ class Sovereign(nn.Module):
                                  align_corners=False)
         return logg
 
+    # -- planning interface (used by search.py) -----------------------------
+
+    def encode(self, rgb: torch.Tensor, nums: Optional[torch.Tensor] = None,
+               rtg: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Compressed latent state z (B, latent_dim) for MCTS + dynamics
+        training. Runs the vision spine + cortex (no decoder)."""
+        assert self.planning is not None, "planning block disabled"
+        cfg = self.cfg
+        B, C, H, W = rgb.shape
+        assert H == W == cfg.map_size, \
+            f"expected {cfg.map_size}x{cfg.map_size}, got {H}x{W}"
+        if nums is None:
+            nums = torch.zeros(B, cfg.ctx_dim, device=rgb.device)
+        ctx_full = self._ctx_full(nums, rtg)
+        x = rgb
+        if cfg.use_coord:
+            x = torch.cat([x, self.coords.expand(B, -1, -1, -1)], dim=1)
+        x = self.stem(x)
+        prev = None
+        for i, stage in enumerate(self.enc):
+            prev = x
+            x = self._staged(stage, x) if i >= len(self.enc) - 2 else stage(x)
+        # x = e5 output, prev = e4 output (cortex fine tokens)
+        bot = self.film(x, ctx_full)
+        _, g = self.cortex(bot, prev)
+        lat_map = F.adaptive_avg_pool2d(
+            self.planning.lat_proj(bot),
+            (cfg.lat_spatial, cfg.lat_spatial)).flatten(1)
+        return torch.cat([g, lat_map], dim=1)      # (B, latent_dim)
+
+    def prediction(self, z: torch.Tensor
+                   ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """MuZero-lite prediction net: prior logits over (kind,cell) joint
+        actions (B, 3*g*g), value (B,), win logits (B,2)."""
+        assert self.planning is not None, "planning block disabled"
+        prior = self.planning.pred_prior(z)
+        value = self.planning.pred_value(z).squeeze(-1)
+        win = self.planning.pred_win(z)
+        return prior, value, win
+
+    def dynamics(self, z: torch.Tensor, kind: torch.Tensor,
+                 cell: torch.Tensor,
+                 pct: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """One-step latent transition: (z, a) -> (z_next, reward_hat)."""
+        assert self.planning is not None, "planning block disabled"
+        cfg = self.cfg
+        B = z.shape[0]
+        ka = self.planning.a_emb_kind(kind)                          # (B,64)
+        ca = self.planning.a_emb_cell(
+            cell.clamp(0, cfg.grid_default * cfg.grid_default - 1))  # (B,256)
+        pa = self.planning.a_emb_pct(
+            pct.clamp(0.0, 1.0).to(z.dtype).view(B, 1))             # (B,128)
+        a = self.planning.a_emb_out(torch.cat([ka, ca, pa], dim=1))  # (B,512)
+        dz = self.planning.dyn_mlp(torch.cat([z, a], dim=1))
+        z_next = z + dz                                              # residual
+        r_hat = self.planning.reward_head(z_next).squeeze(-1)
+        return z_next, r_hat
+
     # -- forward ------------------------------------------------------------
 
     def forward(self, rgb: torch.Tensor, nums: Optional[torch.Tensor] = None,
@@ -783,8 +877,9 @@ def param_breakdown(net: Sovereign) -> dict:
     if net.cfg.use_cell_value:
         heads += n(net.cell_value)
     mem = (n(net.mem) + n(net.mem_proj)) if net.mem is not None else 0
+    planning = n(net.planning) if net.planning is not None else 0
     return dict(spine=spine, cortex=cortex, decoder=decoder, heads=heads,
-                mem=mem)
+                mem=mem, planning=planning)
 
 
 def make_model(cfg: Optional[SovereignConfig] = None,
@@ -803,7 +898,8 @@ if __name__ == "__main__":
     bd = param_breakdown(m)
     print(f"params: {total/1e6:.2f}M (spine {bd['spine']/1e6:.1f}M | "
           f"cortex {bd['cortex']/1e6:.1f}M | decoder {bd['decoder']/1e6:.1f}M "
-          f"| heads {bd['heads']/1e6:.1f}M | mem {bd['mem']/1e6:.1f}M)")
+          f"| heads {bd['heads']/1e6:.1f}M | mem {bd['mem']/1e6:.1f}M | "
+          f"planning {bd['planning']/1e6:.1f}M)")
     H = cfg.map_size
     rgb = torch.rand(2, 3, H, H)
     nums = torch.rand(2, 8)
@@ -821,6 +917,14 @@ if __name__ == "__main__":
                    for k, v in a.items() if k not in ("kind_probs",)})
     a2 = m.act(rgb, nums, rtg=rtg, state=a["state"])
     print("mem carry ok:", a2["state"].shape)
+    # planning block smoke (encode -> predict -> dynamics round trip)
+    z = m.encode(rgb, nums, rtg=rtg)
+    prior_z, v_z, win_z = m.prediction(z)
+    z2, r_hat = m.dynamics(z, torch.zeros(2, dtype=torch.long),
+                           torch.zeros(2, dtype=torch.long),
+                           torch.full((2,), 0.25))
+    print("planning:", prior_z.shape, v_z.shape, win_z.shape, z2.shape,
+          r_hat.shape)
     # no-rtg path (contract default)
     _, _, _, _ = m(rgb, nums)
     print("no-rtg contract path OK")