added training against checkpoint

scripts/evaluate.py

@@ -1,49 +1,33 @@
 """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`.
+Override defaults with OmegaConf arguments such as `num_games=200 agent0=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.
+`agent0=model:ppo_model.eqx`. Games are run in a single vmapped batch of
+`num_games` parallel envs.
 """
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
-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
+from general_bots_training.opponents import make_opponent
 
 
 @dataclass
 class EvalConfig:
     grid_dims: tuple[int, int] = (21, 21)
-    truncation: int = 500
+    truncation: int = 1200
     num_games: int = 100
-    agent0: AgentConfig = field(default_factory=lambda: AgentConfig(kind="random"))
-    agent1: AgentConfig = field(default_factory=lambda: AgentConfig(kind="expander"))
+    agent0: str = "random"
+    agent1: str = "expander"
     seed: int = 0
 
 
@@ -53,47 +37,13 @@ def load_config(args: list[str] | None = None) -> DictConfig:
     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)
+    key, env_key = jrandom.split(key, 2)
 
-    agent0 = make_agent(config.agent0, agent0_key)
-    agent1 = make_agent(config.agent1, agent1_key)
-    label0, label1 = agent_label(config.agent0), agent_label(config.agent1)
+    agent0 = make_opponent(config.agent0)
+    agent1 = make_opponent(config.agent1)
+    label0, label1 = config.agent0, config.agent1
 
     env = GeneralsEnv(grid_dims=tuple(config.grid_dims), truncation=config.truncation)
     pool, _ = env.reset(env_key)