added training against checkpoint

.gitignore

@@ -0,0 +1 @@
+__pycache__/

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
+`agent0=model:ppo_model.eqx`. Games are run in a single vmapped batch of
-vmapped batch of `num_games` parallel envs.
+`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 make_opponent
-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
+    truncation: int = 1200
     num_games: int = 100
-    agent0: AgentConfig = field(default_factory=lambda: AgentConfig(kind="random"))
+    agent0: str = "random"
-    agent1: AgentConfig = field(default_factory=lambda: AgentConfig(kind="expander"))
+    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)
+    agent0 = make_opponent(config.agent0)
-    agent1 = make_agent(config.agent1, agent1_key)
+    agent1 = make_opponent(config.agent1)
-    label0, label1 = agent_label(config.agent0), agent_label(config.agent1)
+    label0, label1 = config.agent0, config.agent1
 
     env = GeneralsEnv(grid_dims=tuple(config.grid_dims), truncation=config.truncation)
     pool, _ = env.reset(env_key)

src/general_bots_training/__init__.py

@@ -1,6 +1,6 @@
 from .mcts import MCTSConfig, ParticlePUCT
 from .network import PolicyValueNetwork, obs_to_tensor
-from .opponents import Opponent, SelfPlayOpponent, StaticOpponent, make_opponent
+from .opponents import ModelOpponent, 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
 
@@ -17,6 +17,7 @@ __all__ = [
     "make_rollout_step",
     "make_opponent",
     "StaticOpponent",
+    "ModelOpponent",
     "SelfPlayOpponent",
     "Opponent",
 ]

src/general_bots_training/opponents.py

@@ -4,10 +4,15 @@ from typing import Protocol
 
 from dataclasses import dataclass
 
+import equinox as eqx
 import jax.numpy as jnp
+import jax.random as jrandom
 from generals.agents import Agent, ExpanderAgent, HunterAgent, RandomAgent
+from generals.core.action import compute_valid_move_mask
 from generals.core.observation import Observation
 
+from .network import PolicyValueNetwork, obs_to_tensor
+
 
 class StaticOpponent(Protocol):
     """Opponent policy that carries no per-environment mutable state."""
@@ -20,6 +25,21 @@ class SelfPlayOpponent:
     """Marker selecting the current training network as player 1."""
 
 
+class ModelOpponent:
+    """Static opponent that plays a fixed policy loaded from an equinox checkpoint."""
+
+    def __init__(self, network: PolicyValueNetwork):
+        self._network = network
+
+    def act(self, observation: Observation, key: jnp.ndarray) -> jnp.ndarray:
+        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
+
+
 Opponent = StaticOpponent | SelfPlayOpponent
 
 
@@ -30,15 +50,36 @@ OPPONENT_TYPES: dict[str, type[Agent]] = {
 }
 
 
-def make_opponent(name: str) -> Opponent:
+def make_opponent(spec: str) -> Opponent:
-    """Create an opponent strategy by configuration name."""
+    """Create an opponent strategy from a single configuration string.
+
+    The spec is either a bare opponent name ("random", "expander", "hunter",
+    "self_play") or a ``type:checkpoint`` pair for a model opponent, e.g.
+    ``"model:ppo_model.eqx"``.
+    """
+    name, sep, checkpoint = spec.partition(":")
     normalized_name = name.lower().replace("-", "_")
     if normalized_name == "self_play":
         return SelfPlayOpponent()
 
+    if normalized_name == "model":
+        if not sep:
+            raise ValueError(
+                "opponent 'model' requires a checkpoint path, e.g. 'model:ppo_model.eqx'"
+            )
+        # The key only seeds the network structure; deserialization overwrites
+        # every leaf, so a fixed seed is sufficient.
+        network = PolicyValueNetwork(jrandom.PRNGKey(0), in_channels=14)
+        network = eqx.tree_deserialise_leaves(checkpoint, network)
+        return ModelOpponent(network)
+
+    if sep:
+        choices = ", ".join([*sorted(OPPONENT_TYPES), "model", "self_play"])
+        raise ValueError(f"opponent {spec!r} does not take a checkpoint; choose one of: {choices}")
+
     try:
         opponent_type = OPPONENT_TYPES[normalized_name]
     except KeyError as error:
-        choices = ", ".join([*sorted(OPPONENT_TYPES), "self_play"])
+        choices = ", ".join([*sorted(OPPONENT_TYPES), "model", "self_play"])
-        raise ValueError(f"unknown opponent {name!r}; choose one of: {choices}") from error
+        raise ValueError(f"unknown opponent {spec!r}; choose one of: {choices}") from error
     return opponent_type()

tests/test_opponents.py

@@ -1,7 +1,7 @@
 import pytest
 from generals.agents import ExpanderAgent, HunterAgent, RandomAgent
 
-from general_bots_training.opponents import SelfPlayOpponent, make_opponent
+from general_bots_training.opponents import ModelOpponent, SelfPlayOpponent, make_opponent
 
 
 @pytest.mark.parametrize(
@@ -28,3 +28,28 @@ def test_make_opponent_supports_self_play_aliases():
 def test_make_opponent_rejects_unknown_name():
     with pytest.raises(ValueError, match="unknown opponent"):
         make_opponent("turtle")
+
+
+def test_make_opponent_model_requires_checkpoint():
+    with pytest.raises(ValueError, match="requires a checkpoint path"):
+        make_opponent("model")
+
+
+def test_make_opponent_model_loads_checkpoint(tmp_path):
+    import equinox as eqx
+    import jax.random as jrandom
+
+    from general_bots_training.network import PolicyValueNetwork
+
+    network = PolicyValueNetwork(jrandom.PRNGKey(0), in_channels=14)
+    checkpoint_path = tmp_path / "model.eqx"
+    eqx.tree_serialise_leaves(checkpoint_path, network)
+
+    opponent = make_opponent(f"model:{checkpoint_path}")
+
+    assert isinstance(opponent, ModelOpponent)
+
+
+def test_make_opponent_rejects_checkpoint_on_non_model():
+    with pytest.raises(ValueError, match="does not take a checkpoint"):
+        make_opponent("random:some.eqx")