diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c18dd8d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/ppo_model.eqx b/ppo_model.eqx index f0e1c62..1e5d698 100644 Binary files a/ppo_model.eqx and b/ppo_model.eqx differ diff --git a/ppo_model.eqx.6 b/ppo_model.eqx.6 new file mode 100644 index 0000000..4bcf7cb Binary files /dev/null and b/ppo_model.eqx.6 differ diff --git a/scripts/evaluate.py b/scripts/evaluate.py index d537dbf..b9c7b91 100644 --- a/scripts/evaluate.py +++ b/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) diff --git a/src/general_bots_training/__init__.py b/src/general_bots_training/__init__.py index 899572c..5c24ebc 100644 --- a/src/general_bots_training/__init__.py +++ b/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", ] diff --git a/src/general_bots_training/__pycache__/__init__.cpython-313.pyc b/src/general_bots_training/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 5b3e23f..0000000 Binary files a/src/general_bots_training/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/src/general_bots_training/__pycache__/mcts.cpython-313.pyc b/src/general_bots_training/__pycache__/mcts.cpython-313.pyc deleted file mode 100644 index d969103..0000000 Binary files a/src/general_bots_training/__pycache__/mcts.cpython-313.pyc and /dev/null differ diff --git a/src/general_bots_training/__pycache__/network.cpython-313.pyc b/src/general_bots_training/__pycache__/network.cpython-313.pyc deleted file mode 100644 index 4cb8cab..0000000 Binary files a/src/general_bots_training/__pycache__/network.cpython-313.pyc and /dev/null differ diff --git a/src/general_bots_training/__pycache__/opponents.cpython-313.pyc b/src/general_bots_training/__pycache__/opponents.cpython-313.pyc deleted file mode 100644 index 3616d0d..0000000 Binary files a/src/general_bots_training/__pycache__/opponents.cpython-313.pyc and /dev/null differ diff --git a/src/general_bots_training/__pycache__/ppo.cpython-313.pyc b/src/general_bots_training/__pycache__/ppo.cpython-313.pyc deleted file mode 100644 index 56eef53..0000000 Binary files a/src/general_bots_training/__pycache__/ppo.cpython-313.pyc and /dev/null differ diff --git a/src/general_bots_training/__pycache__/rollout.cpython-313.pyc b/src/general_bots_training/__pycache__/rollout.cpython-313.pyc deleted file mode 100644 index ad76e46..0000000 Binary files a/src/general_bots_training/__pycache__/rollout.cpython-313.pyc and /dev/null differ diff --git a/src/general_bots_training/__pycache__/train.cpython-313.pyc b/src/general_bots_training/__pycache__/train.cpython-313.pyc deleted file mode 100644 index 871538a..0000000 Binary files a/src/general_bots_training/__pycache__/train.cpython-313.pyc and /dev/null differ diff --git a/src/general_bots_training/opponents.py b/src/general_bots_training/opponents.py index ee5abe5..a77e136 100644 --- a/src/general_bots_training/opponents.py +++ b/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: - """Create an opponent strategy by configuration name.""" +def make_opponent(spec: str) -> Opponent: + """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"]) - raise ValueError(f"unknown opponent {name!r}; choose one of: {choices}") from error + choices = ", ".join([*sorted(OPPONENT_TYPES), "model", "self_play"]) + raise ValueError(f"unknown opponent {spec!r}; choose one of: {choices}") from error return opponent_type() diff --git a/tests/__pycache__/test_checkpoint.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_checkpoint.cpython-313-pytest-9.1.1.pyc deleted file mode 100644 index 3aa7e31..0000000 Binary files a/tests/__pycache__/test_checkpoint.cpython-313-pytest-9.1.1.pyc and /dev/null differ diff --git a/tests/__pycache__/test_mcts.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_mcts.cpython-313-pytest-9.1.1.pyc deleted file mode 100644 index ec155ac..0000000 Binary files a/tests/__pycache__/test_mcts.cpython-313-pytest-9.1.1.pyc and /dev/null differ diff --git a/tests/__pycache__/test_opponents.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_opponents.cpython-313-pytest-9.1.1.pyc deleted file mode 100644 index a5150c5..0000000 Binary files a/tests/__pycache__/test_opponents.cpython-313-pytest-9.1.1.pyc and /dev/null differ diff --git a/tests/__pycache__/test_self_play.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_self_play.cpython-313-pytest-9.1.1.pyc deleted file mode 100644 index 8aeda74..0000000 Binary files a/tests/__pycache__/test_self_play.cpython-313-pytest-9.1.1.pyc and /dev/null differ diff --git a/tests/__pycache__/test_train_config.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_train_config.cpython-313-pytest-9.1.1.pyc deleted file mode 100644 index 75beb00..0000000 Binary files a/tests/__pycache__/test_train_config.cpython-313-pytest-9.1.1.pyc and /dev/null differ diff --git a/tests/__pycache__/test_training.cpython-313-pytest-9.1.1.pyc b/tests/__pycache__/test_training.cpython-313-pytest-9.1.1.pyc deleted file mode 100644 index 0705288..0000000 Binary files a/tests/__pycache__/test_training.cpython-313-pytest-9.1.1.pyc and /dev/null differ diff --git a/tests/test_opponents.py b/tests/test_opponents.py index e553630..362870d 100644 --- a/tests/test_opponents.py +++ b/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")