Upload model_v3.py with huggingface_hub

model_v3.py

@@ -17,10 +17,20 @@ policy could just read the bit. The bypass split:
 
 Latent: z_grid (B, latent_c, H/16, W/16). No global vector: the policy can
 pool spatially itself, and per-player/global state arrives via bypass.
+
+v3.1 additions (all off by default so old v3 checkpoints load unchanged):
+  - terrain_cond: the decoder consumes the 2 STATIC terrain planes (land,
+    magnitude) as free side-information at every scale, so the latent only
+    has to encode ownership relative to terrain. Fallout is dynamic state
+    and is never fed to the decoder.
+  - upsample_decoder: nearest-upsample + 3x3 conv stages (no checkerboard)
+    plus a full-resolution 3x3 refinement block before the classifier.
+  - latent_down: 8 or 16; latent grid at 1/8 or 1/16 resolution.
 """
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 from ae.units import STATIC_CLASSES
 
@@ -49,50 +59,133 @@ def deconv_block(c_in: int, c_out: int) -> nn.Sequential:
     )
 
 
+class UpsampleBlock(nn.Module):
+    """Nearest 2x upsample + 3x3 conv (no ConvTranspose checkerboard).
+
+    Optionally concatenates static terrain planes (at the post-upsample
+    resolution) before the conv.
+    """
+
+    def __init__(self, c_in: int, c_out: int, extra_c: int = 0):
+        super().__init__()
+        self.conv = conv_block(c_in + extra_c, c_out, stride=1)
+
+    def forward(self, x: torch.Tensor, extra: torch.Tensor | None = None):
+        x = F.interpolate(x, scale_factor=2, mode="nearest")
+        if extra is not None:
+            x = torch.cat([x, extra], dim=1)
+        return self.conv(x)
+
+
+STATIC_TERRAIN_C = 2  # land, magnitude (fallout is dynamic: never decoded from)
+
+
 class SpatialAE(nn.Module):
-    def __init__(self, latent_c: int = 64):
+    """Defaults preserve the original v3 architecture (old checkpoints load
+    with strict state_dict matching). v3.1 runs set terrain_cond=True and
+    upsample_decoder=True (and optionally latent_down=8)."""
+
+    def __init__(
+        self,
+        latent_c: int = 64,
+        terrain_cond: bool = False,
+        upsample_decoder: bool = False,
+        latent_down: int = 16,
+    ):
         super().__init__()
+        if latent_down not in (8, 16):
+            raise ValueError(f"latent_down must be 8 or 16, got {latent_down}")
+        if latent_down == 8 and not upsample_decoder:
+            raise ValueError("latent_down=8 requires the v3.1 upsample decoder")
+        if terrain_cond and not upsample_decoder:
+            raise ValueError("terrain_cond requires the v3.1 upsample decoder")
         self.latent_c = latent_c
+        self.terrain_cond = terrain_cond
+        self.upsample_decoder = upsample_decoder
+        self.latent_down = latent_down
         self.owner_emb = nn.Embedding(MAX_SLOTS, OWNER_EMB_DIM)
 
-        self.enc_stem = nn.Sequential(
+        stem = [
             conv_block(OWNER_EMB_DIM + TERRAIN_CHANNELS, 32, stride=1),
             conv_block(32, 64, stride=2),
             conv_block(64, 96, stride=2),
             conv_block(96, 128, stride=2),
-            conv_block(128, 128, stride=2),  # -> 1/16
-        )
+        ]
+        if latent_down == 16:
+            stem.append(conv_block(128, 128, stride=2))  # -> 1/16
+        self.enc_stem = nn.Sequential(*stem)
         self.enc_fuse = nn.Sequential(
             conv_block(128 + NUM_STATIC, 128, stride=1),
             nn.Conv2d(128, latent_c, kernel_size=1),
         )
 
-        self.dec_in = conv_block(latent_c, 128, stride=1)
-        self.dec_tiles = nn.Sequential(
-            deconv_block(128, 128),
-            deconv_block(128, 96),
-            deconv_block(96, 64),
-            deconv_block(64, 32),
-            nn.Conv2d(32, MAX_SLOTS, kernel_size=1),
-        )
-        # Static structure occupancy logits at 1/16 resolution.
+        cond_c = STATIC_TERRAIN_C if terrain_cond else 0
+        self.dec_in = conv_block(latent_c + cond_c, 128, stride=1)
+        if upsample_decoder:
+            chans = [128, 128, 96, 64, 32] if latent_down == 16 else [128, 96, 64, 32]
+            self.dec_up = nn.ModuleList(
+                UpsampleBlock(chans[i], chans[i + 1], extra_c=cond_c)
+                for i in range(len(chans) - 1)
+            )
+            self.dec_refine = conv_block(32 + cond_c, 32, stride=1)
+            self.dec_out = nn.Conv2d(32, MAX_SLOTS, kernel_size=1)
+        else:
+            self.dec_tiles = nn.Sequential(
+                deconv_block(128, 128),
+                deconv_block(128, 96),
+                deconv_block(96, 64),
+                deconv_block(64, 32),
+                nn.Conv2d(32, MAX_SLOTS, kernel_size=1),
+            )
+        # Static structure occupancy logits at latent resolution.
         self.dec_units = nn.Conv2d(128, NUM_STATIC, kernel_size=1)
 
     def encode(
         self,
         owners: torch.Tensor,  # (B, H, W) int64
         terrain: torch.Tensor,  # (B, 3, H, W)
-        static_planes: torch.Tensor,  # (B, NUM_STATIC, H/16, W/16)
+        static_planes: torch.Tensor,  # (B, NUM_STATIC, H/down, W/down)
     ) -> torch.Tensor:
         emb = self.owner_emb(owners).permute(0, 3, 1, 2)
         g = self.enc_stem(torch.cat([emb, terrain], dim=1))
         return self.enc_fuse(torch.cat([g, static_planes], dim=1))
 
-    def decode(self, z_grid: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        h = self.dec_in(z_grid)
-        return self.dec_tiles(h), self.dec_units(h)
+    def decode(
+        self,
+        z_grid: torch.Tensor,
+        terrain: torch.Tensor | None = None,  # (B, >=2, H, W); only [:, :2] used
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if not self.terrain_cond:
+            h = self.dec_in(z_grid)
+            if self.upsample_decoder:
+                x = h
+                for up in self.dec_up:
+                    x = up(x)
+                return self.dec_out(self.dec_refine(x)), self.dec_units(h)
+            return self.dec_tiles(h), self.dec_units(h)
+
+        if terrain is None:
+            raise ValueError("terrain_cond model needs terrain in decode()")
+        # Static side-information pyramid: full res, 1/2, 1/4, ... latent res.
+        static_t = terrain[:, :STATIC_TERRAIN_C]
+        pyramid = {1: static_t}
+        down = 2
+        while down <= self.latent_down:
+            pyramid[down] = F.avg_pool2d(static_t, kernel_size=down)
+            down *= 2
+
+        h = self.dec_in(torch.cat([z_grid, pyramid[self.latent_down]], dim=1))
+        x = h
+        scale = self.latent_down
+        for up in self.dec_up:
+            scale //= 2
+            x = up(x, pyramid[scale])
+        x = self.dec_refine(torch.cat([x, pyramid[1]], dim=1))
+        return self.dec_out(x), self.dec_units(h)
 
     def forward(self, owners, terrain, static_planes):
         z_grid = self.encode(owners, terrain, static_planes)
-        tile_logits, unit_logits = self.decode(z_grid)
+        tile_logits, unit_logits = self.decode(
+            z_grid, terrain if self.terrain_cond else None
+        )
         return tile_logits, unit_logits, z_grid