added evaluation script

scripts/evaluate.py

@@ -0,0 +1,173 @@
+"""Evaluate Generals.io agents by playing them against each other.
+
+Run with `.venv/bin/python scripts/evaluate.py` or `uv run python scripts/evaluate.py`.
+Override defaults with OmegaConf arguments such as `num_games=200 agent0.kind=hunter`.
+
+Each agent is either one of the bundled JAX agents (`random`, `expander`,
+`hunter`) or a trained policy loaded from an equinox checkpoint via
+`agent0.kind=model agent0.checkpoint=ppo_model.eqx`. Games are run in a single
+vmapped batch of `num_games` parallel envs.
+"""
+
+from dataclasses import dataclass, field
+
+import equinox as eqx
+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.observation import Observation
+from omegaconf import DictConfig, OmegaConf
+
+from general_bots_training.network import PolicyValueNetwork, obs_to_tensor
+from general_bots_training.opponents import OPPONENT_TYPES, StaticOpponent
+
+
+@dataclass
+class AgentConfig:
+    """One side of the matchup.
+
+    kind: "random" | "expander" | "hunter" | "model"
+    checkpoint: path to an equinox checkpoint (only used when kind == "model")
+    """
+
+    kind: str = "random"
+    checkpoint: str | None = None
+
+
+@dataclass
+class EvalConfig:
+    grid_dims: tuple[int, int] = (21, 21)
+    truncation: int = 500
+    num_games: int = 100
+    agent0: AgentConfig = field(default_factory=lambda: AgentConfig(kind="random"))
+    agent1: AgentConfig = field(default_factory=lambda: AgentConfig(kind="expander"))
+    seed: int = 0
+
+
+def load_config(args: list[str] | None = None) -> DictConfig:
+    """Merge CLI overrides into the typed default evaluation configuration."""
+    defaults = OmegaConf.structured(EvalConfig)
+    return OmegaConf.merge(defaults, OmegaConf.from_cli(args))  # type: ignore
+
+
+class NetworkAgent:
+    """Wrap a PolicyValueNetwork in the stateless StaticOpponent interface."""
+
+    def __init__(self, network):
+        self._network = network
+
+    def act(self, observation: Observation, key):
+        obs_arr = obs_to_tensor(observation)
+        mask = compute_valid_move_mask(
+            observation.armies, observation.owned_cells, observation.mountains
+        )
+        action, _, _, _ = self._network(obs_arr, mask, key, None)
+        return action
+
+
+def make_agent(cfg: AgentConfig, key) -> StaticOpponent:
+    """Build an agent from its config. `key` seeds the network when needed."""
+    kind = cfg.kind.lower().replace("-", "_")
+    if kind == "model":
+        if cfg.checkpoint is None:
+            raise ValueError("agent kind 'model' requires a checkpoint path")
+        network = PolicyValueNetwork(key, in_channels=14)
+        network = eqx.tree_deserialise_leaves(cfg.checkpoint, network)
+        return NetworkAgent(network)
+    if kind in OPPONENT_TYPES:
+        return OPPONENT_TYPES[kind]()
+    choices = ", ".join([*sorted(OPPONENT_TYPES), "model"])
+    raise ValueError(f"unknown agent kind {cfg.kind!r}; choose one of: {choices}")
+
+
+def agent_label(cfg: AgentConfig) -> str:
+    return cfg.checkpoint if cfg.kind == "model" and cfg.checkpoint else cfg.kind
+
+
+def main(config: DictConfig):
+    key = jrandom.PRNGKey(config.seed)
+    key, agent0_key, agent1_key, env_key = jrandom.split(key, 4)
+
+    agent0 = make_agent(config.agent0, agent0_key)
+    agent1 = make_agent(config.agent1, agent1_key)
+    label0, label1 = agent_label(config.agent0), agent_label(config.agent1)
+
+    env = GeneralsEnv(grid_dims=tuple(config.grid_dims), truncation=config.truncation)
+    pool, _ = env.reset(env_key)
+    get_obs = game.get_full_observation if env.perfect_info else game.get_observation
+
+    n = config.num_games
+    key, init_key = jrandom.split(key)
+    init_keys = jrandom.split(init_key, n)
+    states = jax.vmap(env.init_state)(init_keys)
+    states = states._replace(pool_idx=jnp.arange(n, dtype=states.pool_idx.dtype) % env.pool_size)
+
+    step_env = jax.vmap(env.step, in_axes=(0, 0, None))
+
+    def run_episode(states, key):
+        """Play one full game per env; return final winner per env.
+
+        The env auto-resets on done and overwrites `state.winner` with -1, so we
+        capture the winner from `timestep.info` on the first done step per env and
+        latch it for the remainder of the scan.
+        """
+
+        def body(carry, _):
+            states, latched_winner, key = carry
+            key, p0_key, p1_key = jrandom.split(key, 3)
+            obs_p0 = jax.vmap(lambda s: get_obs(s, 0))(states)
+            obs_p1 = jax.vmap(lambda s: get_obs(s, 1))(states)
+            keys_p0 = jrandom.split(p0_key, n)
+            keys_p1 = jrandom.split(p1_key, n)
+            actions_p0 = jax.vmap(agent0.act)(obs_p0, keys_p0)
+            actions_p1 = jax.vmap(agent1.act)(obs_p1, keys_p1)
+            actions = jnp.stack([actions_p0, actions_p1], axis=1)
+            timesteps, states = step_env(states, actions, pool)
+            done = timesteps.terminated | timesteps.truncated
+            # On the first done, latch the winner (truncated games stay at -1 = draw).
+            new_winner = jnp.where(
+                done & (latched_winner < 0), timesteps.info.winner, latched_winner
+            )
+            return (states, new_winner, key), done
+
+        # Run a fixed number of steps equal to the truncation length, which
+        # guarantees every env has terminated or truncated at least once.
+        (states, latched_winner, key), _ = jax.lax.scan(
+            body, (states, jnp.full(n, -1, dtype=jnp.int32), key), None, length=config.truncation
+        )
+        return latched_winner
+
+    print("Generals.io evaluation")
+    print(OmegaConf.to_yaml(config, resolve=True).rstrip())
+    print(f"device: {jax.devices()[0]}")
+    print(f"agent0: {label0}")
+    print(f"agent1: {label1}")
+    print(f"games:  {n}")
+    print()
+    print("warming up (jit compile)...")
+    winners = run_episode(states, key)
+    jax.block_until_ready(winners)
+    print("warming up done\n")
+
+    print("playing...")
+    winners = run_episode(states, key)
+    jax.block_until_ready(winners)
+
+    wins0 = int(jnp.sum(winners == 0))
+    wins1 = int(jnp.sum(winners == 1))
+    draws = int(jnp.sum(winners < 0))
+
+    print("\nsummary")
+    print("-" * 40)
+    print(f"agent0 ({label0}): {wins0:4d} wins ({wins0 / n * 100:.1f}%)")
+    print(f"agent1 ({label1}): {wins1:4d} wins ({wins1 / n * 100:.1f}%)")
+    print(f"draws:              {draws:4d}      ({draws / n * 100:.1f}%)")
+    print("-" * 40)
+    print(f"total games: {n}")
+
+
+if __name__ == "__main__":
+    main(load_config())