checkpoint

src/general_bots_training/__init__.py

@@ -0,0 +1,22 @@
+from .mcts import MCTSConfig, ParticlePUCT
+from .network import PolicyValueNetwork, obs_to_tensor
+from .opponents import Opponent, SelfPlayOpponent, StaticOpponent, make_opponent
+from .ppo import compute_advantages_and_returns, compute_gae, make_train_epoch, ppo_loss
+from .rollout import make_collect_rollout, make_rollout_step
+
+__all__ = [
+    "PolicyValueNetwork",
+    "MCTSConfig",
+    "ParticlePUCT",
+    "obs_to_tensor",
+    "compute_gae",
+    "compute_advantages_and_returns",
+    "make_train_epoch",
+    "ppo_loss",
+    "make_collect_rollout",
+    "make_rollout_step",
+    "make_opponent",
+    "StaticOpponent",
+    "SelfPlayOpponent",
+    "Opponent",
+]

src/general_bots_training/mcts.py

@@ -0,0 +1,337 @@
+"""Observation-safe, simultaneous-action PUCT for competition inference.
+
+The search never receives the evaluator's hidden GameState. It samples conservative
+full-state determinizations from the current fog observation, selects our root move
+with PUCT, samples opponent replies from the same policy, and applies the exact
+competition build/deathtouch transition.
+"""
+
+import time
+from dataclasses import dataclass
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import numpy as np
+from generals.core import game
+from generals.core.action import compute_valid_move_mask
+from generals.core.game import GameState
+from generals.core.observation import Observation
+from generals.modifiers import build_castles, deathtouch
+
+from .network import PolicyValueNetwork, obs_to_tensor
+
+PASS_ACTION = np.array([1, 0, 0, 0, 0], dtype=np.int32)
+
+
+@dataclass(frozen=True)
+class MCTSConfig:
+    time_budget_ms: float = 125.0
+    max_simulations: int = 128
+    rollout_depth: int = 2
+    top_k: int = 20
+    opponent_top_k: int = 12
+    max_build_actions: int = 3
+    c_puct: float = 1.5
+    value_scale: float = 5.0
+
+
+@jax.jit
+def competition_move_mask(obs: Observation) -> jnp.ndarray:
+    """Legal move mask that conservatively blocks unresolved fog structures."""
+    blocked = obs.mountains | obs.structures_in_fog
+    return compute_valid_move_mask(obs.armies, obs.owned_cells, blocked)
+
+
+@jax.jit
+def build_cost_grid(obs: Observation) -> jnp.ndarray:
+    """Exact own-castle build costs derivable from a fog observation."""
+    structures = ((obs.castles | obs.generals) & obs.owned_cells).astype(jnp.int32)
+    height, width = structures.shape
+    radius = 6
+    padded = jnp.pad(structures, radius)
+    costs = jnp.full((height, width), 35, dtype=jnp.int32)
+    for row_offset in range(-radius, radius + 1):
+        for col_offset in range(-radius, radius + 1):
+            surcharge = 14 - 2 * (abs(row_offset) + abs(col_offset))
+            if surcharge > 0:
+                shifted = padded[
+                    radius + row_offset : radius + row_offset + height,
+                    radius + col_offset : radius + col_offset + width,
+                ]
+                costs = costs + surcharge * shifted
+    return costs
+
+
+@jax.jit
+def valid_build_mask(obs: Observation) -> jnp.ndarray:
+    costs = build_cost_grid(obs)
+    plain_owned = obs.owned_cells & ~obs.generals & ~obs.castles
+    return plain_owned & (obs.armies >= costs)
+
+
+def _decode_policy_index(index: int, height: int, width: int) -> np.ndarray:
+    cells = height * width
+    if index == 8 * cells:
+        return PASS_ACTION.copy()
+    encoded_direction, position = divmod(index, cells)
+    row, col = divmod(position, width)
+    split = int(encoded_direction >= 4)
+    direction = encoded_direction - 4 if split else encoded_direction
+    return np.array([0, row, col, direction, split], dtype=np.int32)
+
+
+def observation_from_wire(
+    timestep: int,
+    owned_land_count: int,
+    owned_army_count: int,
+    opponent_land_count: int,
+    opponent_army_count: int,
+    type_grid,
+    owner_grid,
+    army_grid,
+) -> Observation:
+    """Convert a competition wire frame into the network's Observation type."""
+    types = jnp.asarray(type_grid, dtype=jnp.int32)
+    owners = jnp.asarray(owner_grid, dtype=jnp.int32)
+    armies = jnp.asarray(army_grid, dtype=jnp.int32)
+    visible = (types != 0) & (types != 5)
+    return Observation(
+        armies=armies,
+        generals=types == 4,
+        castles=types == 3,
+        mountains=types == 2,
+        neutral_cells=visible & (owners == 0) & (types != 2),
+        owned_cells=owners == 1,
+        opponent_cells=owners == 2,
+        fog_cells=types == 0,
+        structures_in_fog=types == 5,
+        owned_land_count=jnp.int32(owned_land_count),
+        owned_army_count=jnp.int32(owned_army_count),
+        opponent_land_count=jnp.int32(opponent_land_count),
+        opponent_army_count=jnp.int32(opponent_army_count),
+        timestep=jnp.int32(timestep),
+    )
+
+
+def sample_determinization(
+    obs: Observation, player_index: int, rng: np.random.Generator
+) -> GameState:
+    """Sample a conservative hidden state consistent with visible cells/totals."""
+    armies = np.asarray(obs.armies, dtype=np.int32).copy()
+    owned = np.asarray(obs.owned_cells, dtype=bool)
+    visible_opponent = np.asarray(obs.opponent_cells, dtype=bool)
+    mountains = np.asarray(obs.mountains | obs.structures_in_fog, dtype=bool)
+    castles = np.asarray(obs.castles, dtype=bool)
+    generals = np.asarray(obs.generals, dtype=bool).copy()
+    fog = np.asarray(obs.fog_cells, dtype=bool) & ~mountains
+
+    opponent = visible_opponent.copy()
+    hidden_candidates = np.argwhere(fog & ~owned)
+    missing_land = max(0, int(obs.opponent_land_count) - int(opponent.sum()))
+    if len(hidden_candidates):
+        chosen = hidden_candidates[
+            rng.choice(
+                len(hidden_candidates),
+                size=min(missing_land, len(hidden_candidates)),
+                replace=False,
+            )
+        ]
+        opponent[chosen[:, 0], chosen[:, 1]] = True
+
+    visible_enemy_general = generals & opponent
+    if not visible_enemy_general.any():
+        candidates = np.argwhere((opponent | fog) & ~owned & ~mountains)
+        if len(candidates):
+            row, col = candidates[rng.integers(len(candidates))]
+            generals[row, col] = True
+            opponent[row, col] = True
+
+    hidden_opponent = opponent & ~visible_opponent
+    remaining_army = max(0, int(obs.opponent_army_count) - int(armies[visible_opponent].sum()))
+    hidden_cells = np.argwhere(hidden_opponent)
+    if len(hidden_cells):
+        base = min(remaining_army, len(hidden_cells))
+        armies[hidden_opponent] = 0
+        armies[hidden_cells[:base, 0], hidden_cells[:base, 1]] = 1
+        remaining_army -= base
+        if remaining_army:
+            allocations = rng.multinomial(
+                remaining_army, np.full(len(hidden_cells), 1 / len(hidden_cells))
+            )
+            armies[hidden_cells[:, 0], hidden_cells[:, 1]] += allocations.astype(np.int32)
+
+    ownership_relative = np.stack([owned, opponent])
+    if player_index == 1:
+        ownership = ownership_relative[::-1]
+    else:
+        ownership = ownership_relative
+    passable = ~mountains
+    neutral = passable & ~ownership[0] & ~ownership[1]
+
+    general_positions = []
+    for absolute_player in range(2):
+        positions = np.argwhere(generals & ownership[absolute_player])
+        if len(positions):
+            general_positions.append(positions[0])
+        else:
+            fallback = np.argwhere(ownership[absolute_player] & passable)
+            general_positions.append(fallback[0] if len(fallback) else np.array([0, 0]))
+
+    return GameState(
+        armies=jnp.asarray(armies),
+        ownership=jnp.asarray(ownership),
+        ownership_neutral=jnp.asarray(neutral),
+        generals=jnp.asarray(generals),
+        castles=jnp.asarray(castles),
+        mountains=jnp.asarray(mountains),
+        passable=jnp.asarray(passable),
+        general_positions=jnp.asarray(general_positions, dtype=jnp.int32),
+        time=jnp.asarray(obs.timestep, dtype=jnp.int32),
+        winner=jnp.int32(-1),
+        pool_idx=jnp.int32(0),
+    )
+
+
+@jax.jit
+def competition_step(state: GameState, actions: jnp.ndarray):
+    state, actions = build_castles.apply_build_actions(state, actions)
+    return deathtouch.step(state, actions, turn=800)
+
+
+class ParticlePUCT:
+    """Deadline-bounded root PUCT with policy-guided simultaneous rollouts."""
+
+    def __init__(
+        self,
+        network: PolicyValueNetwork,
+        player_index: int,
+        config: MCTSConfig = MCTSConfig(),
+        seed: int = 0,
+    ):
+        self.network = network
+        self.player_index = player_index
+        self.config = config
+        self.rng = np.random.default_rng(seed)
+        self._infer = eqx.filter_jit(
+            lambda network, tensor, mask: network.policy_value(tensor, mask)
+        )
+
+    def warmup(self, height: int, width: int) -> None:
+        """Compile shape-dependent network and competition transition kernels."""
+        grid = jnp.zeros((height, width), dtype=jnp.int32)
+        grid = grid.at[0, 0].set(1).at[height - 1, width - 1].set(2)
+        state = game.create_initial_state(grid)
+        state = state._replace(
+            armies=state.armies.at[0, 0].set(100).at[height - 1, width - 1].set(100),
+            time=jnp.int32(100),
+        )
+        obs = game.get_observation(state, self.player_index)
+        self.policy_value(obs)
+        actions = jnp.stack([jnp.asarray(PASS_ACTION), jnp.asarray(PASS_ACTION)])
+        warmed_state, _ = competition_step(state, actions)
+        jax.block_until_ready(warmed_state)
+        candidates, _ = self.candidates(obs)
+        self._simulate(obs, candidates[0])
+        self._simulate(obs, candidates[min(1, len(candidates) - 1)])
+
+    def policy_value(self, obs: Observation):
+        mask = competition_move_mask(obs)
+        logits, value = self._infer(self.network, obs_to_tensor(obs), mask)
+        return np.asarray(logits), float(value), mask
+
+    def candidates(self, obs: Observation, top_k: int | None = None):
+        logits, _, _ = self.policy_value(obs)
+        height, width = obs.armies.shape
+        valid_indices = np.flatnonzero(logits > -1e8)
+        count = min(top_k or self.config.top_k, len(valid_indices))
+        ranked = np.argsort(logits[valid_indices])[::-1][:count]
+        indices = valid_indices[ranked]
+        selected_logits = logits[indices]
+        selected_logits = selected_logits - selected_logits.max()
+        priors = np.exp(selected_logits)
+        actions = [_decode_policy_index(int(index), height, width) for index in indices]
+
+        build_mask = np.asarray(valid_build_mask(obs))
+        build_positions = np.argwhere(build_mask)
+        if len(build_positions):
+            costs = np.asarray(build_cost_grid(obs))
+            armies = np.asarray(obs.armies)
+            scores = np.array([armies[r, c] - costs[r, c] for r, c in build_positions])
+            order = np.argsort(scores)[::-1][: self.config.max_build_actions]
+            build_prior = max(float(priors.sum()) * 0.05, 1e-3)
+            for position_index in order:
+                row, col = build_positions[position_index]
+                actions.append(np.array([2, row, col, 0, 0], dtype=np.int32))
+                priors = np.append(priors, build_prior)
+
+        priors = priors / priors.sum()
+        return actions, priors
+
+    def _sample_policy_action(self, obs: Observation, top_k: int) -> np.ndarray:
+        actions, priors = self.candidates(obs, top_k)
+        return actions[int(self.rng.choice(len(actions), p=priors))]
+
+    def _leaf_value(self, state: GameState, info) -> float:
+        if bool(info.is_done):
+            winner = int(info.winner)
+            return 0.0 if winner < 0 else (1.0 if winner == self.player_index else -1.0)
+        obs = game.get_observation(state, self.player_index)
+        _, value, _ = self.policy_value(obs)
+        return float(np.tanh(value / self.config.value_scale))
+
+    def _simulate(self, root_obs: Observation, root_action: np.ndarray) -> float:
+        state = sample_determinization(root_obs, self.player_index, self.rng)
+        info = game.get_info(state)
+        our_action = root_action
+        for _ in range(self.config.rollout_depth):
+            opponent_index = 1 - self.player_index
+            opponent_obs = game.get_observation(state, opponent_index)
+            opponent_action = self._sample_policy_action(opponent_obs, self.config.opponent_top_k)
+            joint_actions = [None, None]
+            joint_actions[self.player_index] = our_action
+            joint_actions[opponent_index] = opponent_action
+            state, info = competition_step(state, jnp.asarray(joint_actions, dtype=jnp.int32))
+            jax.block_until_ready(state)
+            if bool(info.is_done):
+                break
+            our_obs = game.get_observation(state, self.player_index)
+            our_action = self._sample_policy_action(our_obs, self.config.top_k)
+        return self._leaf_value(state, info)
+
+    def search(self, obs: Observation) -> tuple[np.ndarray, dict[str, float]]:
+        started_at = time.perf_counter()
+        deadline = started_at + self.config.time_budget_ms / 1000.0
+        actions, priors = self.candidates(obs)
+        visits = np.zeros(len(actions), dtype=np.int32)
+        value_sums = np.zeros(len(actions), dtype=np.float64)
+        simulations = 0
+        estimated_simulation_seconds = 0.0
+
+        while (
+            simulations < self.config.max_simulations
+            and time.perf_counter() + estimated_simulation_seconds < deadline
+        ):
+            total_visits = max(1, int(visits.sum()))
+            q_values = np.divide(
+                value_sums,
+                visits,
+                out=np.zeros_like(value_sums),
+                where=visits > 0,
+            )
+            scores = q_values + self.config.c_puct * priors * np.sqrt(total_visits) / (1 + visits)
+            action_index = int(np.argmax(scores))
+            simulation_started_at = time.perf_counter()
+            value = self._simulate(obs, actions[action_index])
+            simulation_elapsed = time.perf_counter() - simulation_started_at
+            estimated_simulation_seconds = max(estimated_simulation_seconds, simulation_elapsed)
+            visits[action_index] += 1
+            value_sums[action_index] += value
+            simulations += 1
+
+        selected = int(np.argmax(visits)) if simulations else int(np.argmax(priors))
+        return actions[selected], {
+            "simulations": float(simulations),
+            "selected_visits": float(visits[selected]),
+            "elapsed_budget_ms": self.config.time_budget_ms,
+        }

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

src/general_bots_training/opponents.py

@@ -0,0 +1,44 @@
+"""Interfaces and factories for stateless, JAX-compatible opponents."""
+
+from typing import Protocol
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+from generals.agents import Agent, ExpanderAgent, HunterAgent, RandomAgent
+from generals.core.observation import Observation
+
+
+class StaticOpponent(Protocol):
+    """Opponent policy that carries no per-environment mutable state."""
+
+    def act(self, observation: Observation, key: jnp.ndarray) -> jnp.ndarray: ...
+
+
+@dataclass(frozen=True)
+class SelfPlayOpponent:
+    """Marker selecting the current training network as player 1."""
+
+
+Opponent = StaticOpponent | SelfPlayOpponent
+
+
+OPPONENT_TYPES: dict[str, type[Agent]] = {
+    "random": RandomAgent,
+    "expander": ExpanderAgent,
+    "hunter": HunterAgent,
+}
+
+
+def make_opponent(name: str) -> Opponent:
+    """Create an opponent strategy by configuration name."""
+    normalized_name = name.lower().replace("-", "_")
+    if normalized_name == "self_play":
+        return SelfPlayOpponent()
+
+    try:
+        opponent_type = OPPONENT_TYPES[normalized_name]
+    except KeyError as error:
+        choices = ", ".join([*sorted(OPPONENT_TYPES), "self_play"])
+        raise ValueError(f"unknown opponent {name!r}; choose one of: {choices}") from error
+    return opponent_type()

src/general_bots_training/ppo.py

@@ -0,0 +1,158 @@
+"""PPO: GAE, clipped surrogate loss, and an optax training step."""
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+
+
+@jax.jit
+def compute_gae(rewards, values, next_value, dones, gamma=0.99, lam=0.95):
+    """Generalized Advantage Estimation.
+
+    Args:
+        rewards: (T, N) per-step rewards.
+        values: (T, N) critic values of the states the actions were taken from.
+        next_value: (N,) bootstrap value for the state observed *after* the
+            last collected step. Should be 0 if the last step was terminal.
+        dones: (T, N) True when the step ended an episode (terminated or
+            truncated). The bootstrap is zeroed on done steps.
+    """
+    T, N = rewards.shape
+    values_with_bootstrap = jnp.concatenate([values, next_value[None, :]], axis=0)
+
+    def gae_step(carry, inputs):
+        last_adv = carry
+        reward, value, next_value, done = inputs
+        nonterminal = 1.0 - done
+        delta = reward + gamma * next_value * nonterminal - value
+        adv = delta + gamma * lam * nonterminal * last_adv
+        return adv, adv
+
+    # Process in reverse time order.
+    inputs = (
+        rewards[::-1],
+        values[::-1],
+        values_with_bootstrap[1:][::-1],
+        dones[::-1],
+    )
+    _, advantages_rev = jax.lax.scan(gae_step, jnp.zeros(N), inputs)
+    return advantages_rev[::-1]
+
+
+def compute_advantages_and_returns(rewards, values, next_value, dones, gamma=0.99, lam=0.95):
+    """Return normalized policy advantages and unnormalized critic targets."""
+    raw_advantages = compute_gae(rewards, values, next_value, dones, gamma, lam)
+    returns = raw_advantages + values
+    advantages = (raw_advantages - raw_advantages.mean()) / (raw_advantages.std() + 1e-8)
+    return advantages, returns
+
+
+def batch_forward(network, obs, mask, action):
+    """Run the network on a batch of samples, returning per-sample outputs.
+
+    The network is the vmapped callable, so its array leaves are batched along
+    axis 0 alongside the data (same pattern that works in rollout.py). `action`
+    is passed as a non-batched (None) positional arg to `__call__`.
+    """
+    return jax.vmap(network, in_axes=(0, 0, None, 0))(obs, mask, None, action)
+
+
+def ppo_loss(
+    network,
+    obs,
+    mask,
+    action,
+    old_logprob,
+    advantage,
+    return_,
+    clip=0.2,
+    value_coef=0.5,
+    entropy_coef=0.01,
+):
+    # obs/mask/action are batched along axis 0; the network is the vmapped
+    # callable so its weights are batched too.
+    _, value, logprob, entropy = batch_forward(network, obs, mask, action)
+
+    ratio = jnp.exp(logprob - old_logprob)
+    clipped = jnp.clip(ratio, 1 - clip, 1 + clip) * advantage
+    policy_loss = -jnp.minimum(ratio * advantage, clipped)
+
+    value_loss = value_coef * (value - return_) ** 2
+    entropy_loss = -entropy_coef * entropy
+
+    return jnp.mean(policy_loss + value_loss + entropy_loss)
+
+
+def make_train_epoch(
+    optimizer,
+    minibatch_size: int,
+    clip: float = 0.2,
+    value_coef: float = 0.5,
+    entropy_coef: float = 0.01,
+):
+    """Return a function (network, opt_state, batch, key) -> (network, opt_state, loss).
+
+    Minibatches are taken from the flattened (T*N) buffer with a fresh shuffle
+    per epoch. The last incomplete minibatch is dropped to avoid recompilation.
+    """
+
+    @eqx.filter_value_and_grad
+    def loss_fn(network, minibatch):
+        obs, mask, action, old_logprob, advantage, return_ = minibatch
+        return ppo_loss(
+            network,
+            obs,
+            mask,
+            action,
+            old_logprob,
+            advantage,
+            return_,
+            clip=clip,
+            value_coef=value_coef,
+            entropy_coef=entropy_coef,
+        )
+
+    @eqx.filter_jit
+    def train_epoch(network, opt_state, batch, key):
+        obs, mask, actions, old_logprobs, advantages, returns = batch
+        # Flatten (T, N, ...) -> (T*N, ...). Each leaf may have a different
+        # rank, so reshape preserves the per-sample trailing dims.
+        obs = obs.reshape(-1, *obs.shape[2:])
+        mask = mask.reshape(-1, *mask.shape[2:])
+        actions = actions.reshape(-1, *actions.shape[2:])
+        old_logprobs = old_logprobs.reshape(-1)
+        advantages = advantages.reshape(-1)
+        returns = returns.reshape(-1)
+        bs = obs.shape[0]
+        perm = jax.random.permutation(key, bs)
+        obs = obs[perm]
+        mask = mask[perm]
+        actions = actions[perm]
+        old_logprobs = old_logprobs[perm]
+        advantages = advantages[perm]
+        returns = returns[perm]
+
+        num_complete = bs // minibatch_size
+        if num_complete == 0:
+            raise ValueError("minibatch_size must not exceed the flattened rollout size")
+        used = num_complete * minibatch_size
+        minibatches = (
+            obs[:used].reshape(num_complete, minibatch_size, *obs.shape[1:]),
+            mask[:used].reshape(num_complete, minibatch_size, *mask.shape[1:]),
+            actions[:used].reshape(num_complete, minibatch_size, *actions.shape[1:]),
+            old_logprobs[:used].reshape(num_complete, minibatch_size),
+            advantages[:used].reshape(num_complete, minibatch_size),
+            returns[:used].reshape(num_complete, minibatch_size),
+        )
+
+        def update_step(carry, minibatch):
+            network, opt_state = carry
+            loss, grads = loss_fn(network, minibatch)
+            updates, opt_state = optimizer.update(grads, opt_state, network)
+            network = eqx.apply_updates(network, updates)
+            return (network, opt_state), loss
+
+        (network, opt_state), losses = jax.lax.scan(update_step, (network, opt_state), minibatches)
+        return network, opt_state, losses.mean()
+
+    return train_epoch

src/general_bots_training/rollout.py

@@ -0,0 +1,134 @@
+"""Jitted rollout collection against static opponents or the current policy."""
+
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+from generals.core import game
+from generals.core.action import compute_valid_move_mask
+from generals.core.env import GeneralsEnv
+from generals.core.rewards import composite_reward_fn
+
+from .network import obs_to_tensor
+from .opponents import Opponent, SelfPlayOpponent
+
+
+def _encode_observations(observations):
+    obs_arrays = jax.vmap(obs_to_tensor)(observations)
+    masks = jax.vmap(
+        lambda obs: compute_valid_move_mask(obs.armies, obs.owned_cells, obs.mountains)
+    )(observations)
+    return obs_arrays, masks
+
+
+def _policy_actions(network, observations, keys):
+    obs_arrays, masks = _encode_observations(observations)
+    actions, values, logprobs, _ = jax.vmap(network, in_axes=(0, 0, 0, None))(
+        obs_arrays, masks, keys, None
+    )
+    return obs_arrays, masks, actions, values, logprobs
+
+
+def make_rollout_step(env: GeneralsEnv, opponent: Opponent):
+    """Build one vectorized rollout step.
+
+    Static-opponent transitions have batch axis N. Self-play transitions have
+    batch axis 2N: player 0 trajectories followed by player 1 trajectories.
+    """
+    step_env = jax.vmap(env.step, in_axes=(0, 0, None))
+    get_obs = game.get_full_observation if env.perfect_info else game.get_observation
+
+    def step(states, pool, network, key):
+        num_envs = states.armies.shape[0]
+        obs_p0 = jax.vmap(lambda state: get_obs(state, 0))(states)
+        obs_p1 = jax.vmap(lambda state: get_obs(state, 1))(states)
+
+        key, p0_key, p1_key = jrandom.split(key, 3)
+        keys_p0 = jrandom.split(p0_key, num_envs)
+        obs_arr_p0, masks_p0, actions_p0, values_p0, logprobs_p0 = _policy_actions(
+            network, obs_p0, keys_p0
+        )
+
+        keys_p1 = jrandom.split(p1_key, num_envs)
+        if isinstance(opponent, SelfPlayOpponent):
+            obs_arr_p1, masks_p1, actions_p1, values_p1, logprobs_p1 = _policy_actions(
+                network, obs_p1, keys_p1
+            )
+        else:
+            actions_p1 = jax.vmap(opponent.act)(obs_p1, keys_p1)
+
+        actions = jnp.stack([actions_p0, actions_p1], axis=1)
+        timesteps, new_states = step_env(states, actions, pool)
+
+        # Use the pre-auto-reset terminal state for reward shaping.
+        obs_p0_post = jax.vmap(lambda state: get_obs(state, 0))(timesteps.last_state)
+        rewards_p0 = jax.vmap(composite_reward_fn)(obs_p0, actions_p0, obs_p0_post)
+        dones = timesteps.terminated | timesteps.truncated
+        winners = timesteps.info.winner
+
+        next_obs_p0 = jax.vmap(lambda state: get_obs(state, 0))(new_states)
+        next_obs_arr_p0 = jax.vmap(obs_to_tensor)(next_obs_p0)
+
+        if isinstance(opponent, SelfPlayOpponent):
+            obs_p1_post = jax.vmap(lambda state: get_obs(state, 1))(timesteps.last_state)
+            rewards_p1 = jax.vmap(composite_reward_fn)(obs_p1, actions_p1, obs_p1_post)
+            next_obs_p1 = jax.vmap(lambda state: get_obs(state, 1))(new_states)
+            next_obs_arr_p1 = jax.vmap(obs_to_tensor)(next_obs_p1)
+            winners_p1 = jnp.where(winners < 0, winners, 1 - winners)
+
+            transition = dict(
+                obs=jnp.concatenate([obs_arr_p0, obs_arr_p1]),
+                mask=jnp.concatenate([masks_p0, masks_p1]),
+                action=jnp.concatenate([actions_p0, actions_p1]),
+                logprob=jnp.concatenate([logprobs_p0, logprobs_p1]),
+                value=jnp.concatenate([values_p0, values_p1]),
+                reward=jnp.concatenate([rewards_p0, rewards_p1]),
+                done=jnp.concatenate([dones, dones]),
+                winner=jnp.concatenate([winners, winners_p1]),
+                player=jnp.concatenate([jnp.zeros_like(winners), jnp.ones_like(winners)]),
+            )
+            next_obs_array = jnp.concatenate([next_obs_arr_p0, next_obs_arr_p1])
+        else:
+            transition = dict(
+                obs=obs_arr_p0,
+                mask=masks_p0,
+                action=actions_p0,
+                logprob=logprobs_p0,
+                value=values_p0,
+                reward=rewards_p0,
+                done=dones,
+                winner=winners,
+                player=jnp.zeros_like(winners),
+            )
+            next_obs_array = next_obs_arr_p0
+
+        return new_states, transition, (key, next_obs_array)
+
+    return step
+
+
+def make_collect_rollout(env: GeneralsEnv, num_steps: int, opponent: Opponent):
+    """Collect a rollout against a static opponent or in two-sided self-play."""
+    step_fn = make_rollout_step(env, opponent)
+    get_obs = game.get_full_observation if env.perfect_info else game.get_observation
+    self_play = isinstance(opponent, SelfPlayOpponent)
+
+    @jax.jit
+    def collect(states, pool, network, key):
+        def body(carry, _):
+            states, pool, network, key, _previous_next_obs = carry
+            states, transition, (key, next_obs) = step_fn(states, pool, network, key)
+            return (states, pool, network, key, next_obs), transition
+
+        initial_p0 = jax.vmap(obs_to_tensor)(jax.vmap(lambda state: get_obs(state, 0))(states))
+        if self_play:
+            initial_p1 = jax.vmap(obs_to_tensor)(jax.vmap(lambda state: get_obs(state, 1))(states))
+            initial_next_obs = jnp.concatenate([initial_p0, initial_p1])
+        else:
+            initial_next_obs = initial_p0
+
+        (states, _, _, key, last_next_obs), transitions = jax.lax.scan(
+            body, (states, pool, network, key, initial_next_obs), None, length=num_steps
+        )
+        return states, transitions, (key, last_next_obs)
+
+    return collect