checkpoint

src/general_bots_training/network.py

@@ -0,0 +1,156 @@
+"""Policy-value network and observation encoding for Generals.io PPO."""
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+
+
+def obs_to_tensor(obs) -> jnp.ndarray:
+    """Encode an Observation into a (C, H, W) float32 tensor for the network.
+
+    Armies and army-counts are log-normalized; land counts are divided by the
+    number of cells; scalar values are broadcast to spatial planes so a plain
+    conv stack can consume them. Grid-size agnostic.
+    """
+    H, W = obs.armies.shape
+    armies = jnp.log1p(obs.armies.astype(jnp.float32)) / jnp.log(50.0)
+
+    def bcast(scalar):
+        return jnp.broadcast_to(scalar.astype(jnp.float32), (H, W))
+
+    own_land = bcast(obs.owned_land_count / (H * W))
+    own_army = bcast(jnp.log1p(obs.owned_army_count.astype(jnp.float32)) / jnp.log(50.0))
+    opp_land = bcast(obs.opponent_land_count / (H * W))
+    opp_army = bcast(jnp.log1p(obs.opponent_army_count.astype(jnp.float32)) / jnp.log(50.0))
+    timestep = bcast(jnp.log1p(obs.timestep.astype(jnp.float32)) / jnp.log(1201.0))
+
+    return jnp.stack(
+        [
+            armies,
+            obs.generals.astype(jnp.float32),
+            obs.castles.astype(jnp.float32),
+            obs.mountains.astype(jnp.float32),
+            obs.neutral_cells.astype(jnp.float32),
+            obs.owned_cells.astype(jnp.float32),
+            obs.opponent_cells.astype(jnp.float32),
+            obs.fog_cells.astype(jnp.float32),
+            obs.structures_in_fog.astype(jnp.float32),
+            own_land,
+            own_army,
+            opp_land,
+            opp_army,
+            timestep,
+        ],
+        axis=0,
+    )
+
+
+class PolicyValueNetwork(eqx.Module):
+    """Conv policy-value network.
+
+    Action layout: 4 full-move directions and 4 half-move (split) directions
+    per source cell, plus one global pass action. Invalid moves are masked to
+    -1e9; pass is always available.
+
+    The value head uses global average pooling over the spatial grid, so the
+    network works for any board size without reshaping linear layers.
+    """
+
+    conv1: eqx.nn.Conv2d
+    conv2: eqx.nn.Conv2d
+    conv3: eqx.nn.Conv2d
+    conv4: eqx.nn.Conv2d
+    policy_conv: eqx.nn.Conv2d
+    value_conv: eqx.nn.Conv2d
+    value_linear1: eqx.nn.Linear
+    value_linear2: eqx.nn.Linear
+
+    def __init__(self, key, in_channels: int = 14, channels=(32, 32, 32, 16)):
+        keys = jrandom.split(key, 8)
+        self.conv1 = eqx.nn.Conv2d(in_channels, channels[0], kernel_size=3, padding=1, key=keys[0])
+        self.conv2 = eqx.nn.Conv2d(channels[0], channels[1], kernel_size=3, padding=1, key=keys[1])
+        self.conv3 = eqx.nn.Conv2d(channels[1], channels[2], kernel_size=3, padding=1, key=keys[2])
+        self.conv4 = eqx.nn.Conv2d(channels[2], channels[3], kernel_size=3, padding=1, key=keys[3])
+        # 9 = 4 dirs (full) + 4 dirs (half) + 1 pass
+        self.policy_conv = eqx.nn.Conv2d(channels[3], 9, kernel_size=1, key=keys[4])
+        self.value_conv = eqx.nn.Conv2d(channels[3], 4, kernel_size=1, key=keys[5])
+        self.value_linear1 = eqx.nn.Linear(4, 64, key=keys[6])
+        self.value_linear2 = eqx.nn.Linear(64, 1, key=keys[7])
+
+    def _features(self, obs):
+        x = jax.nn.relu(self.conv1(obs))
+        x = jax.nn.relu(self.conv2(x))
+        x = jax.nn.relu(self.conv3(x))
+        x = jax.nn.relu(self.conv4(x))
+        return x
+
+    def _value_from_features(self, feat):
+        v = jax.nn.relu(self.value_conv(feat))  # (4, H, W)
+        v = v.mean(axis=(1, 2))  # (4,) global average pool
+        v = jax.nn.relu(self.value_linear1(v))
+        return self.value_linear2(v)[0]
+
+    def value(self, obs):
+        return self._value_from_features(self._features(obs))
+
+    def policy_value(self, obs, mask):
+        """Return deterministic masked policy logits and the critic value."""
+        features = self._features(obs)
+        return self._policy_logits(features, mask), self._value_from_features(features)
+
+    def _policy_logits(self, feat, mask):
+        """Return 8*H*W move logits plus one global pass logit."""
+        logits = self.policy_conv(feat)  # (9, H, W)
+        mask_t = jnp.transpose(mask, (2, 0, 1))  # (4, H, W)
+        penalty = (1.0 - mask_t) * -1e9
+        move_penalty = jnp.concatenate([penalty, penalty], axis=0)
+        move_logits = (logits[:8] + move_penalty).reshape(-1)
+        pass_logit = jnp.mean(logits[8])[None]
+        return jnp.concatenate([move_logits, pass_logit])
+
+    def __call__(self, obs, mask, key, action=None):
+        """
+        Args:
+            obs: (C, H, W) tensor.
+            mask: (H, W, 4) valid-move mask.
+            key: PRNG key (only used when action is None).
+            action: if given, evaluate logprob of this [pass,row,col,dir,split];
+                otherwise sample an action.
+
+        Returns:
+            (action, value, logprob, entropy)
+        """
+        logits, value = self.policy_value(obs, mask)
+
+        H, W = mask.shape[:2]
+        cells = H * W
+
+        if action is None:
+            idx = jrandom.categorical(key, logits)
+        else:
+            is_pass, row, col, direction, is_half = action
+            encoded_dir = jnp.where(is_half > 0, direction + 4, direction)
+            move_idx = encoded_dir * cells + row * W + col
+            idx = jnp.where(is_pass > 0, 8 * cells, move_idx)
+
+        log_probs = jax.nn.log_softmax(logits)
+        logprob = log_probs[idx]
+        probs = jax.nn.softmax(logits)
+        entropy = -jnp.sum(probs * log_probs)
+
+        if action is None:
+            is_pass = idx == 8 * cells
+            move_idx = jnp.minimum(idx, 8 * cells - 1)
+            direction = move_idx // cells
+            position = move_idx % cells
+            row = jnp.where(is_pass, 0, position // W)
+            col = jnp.where(is_pass, 0, position % W)
+            is_half = (~is_pass) & (direction >= 4)
+            actual_dir = jnp.where(is_pass, 0, jnp.where(is_half, direction - 4, direction))
+            action = jnp.array(
+                [is_pass.astype(jnp.int32), row, col, actual_dir, is_half.astype(jnp.int32)],
+                dtype=jnp.int32,
+            )
+
+        return action, value, logprob, entropy