123 lines
4.5 KiB
Python
123 lines
4.5 KiB
Python
"""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=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=model:ppo_model.eqx`. Games are run in a single vmapped batch of
|
|
`num_games` parallel envs.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import jax.random as jrandom
|
|
from generals.core import game
|
|
from generals.core.env import GeneralsEnv
|
|
from omegaconf import DictConfig, OmegaConf
|
|
|
|
from general_bots_training.opponents import make_opponent
|
|
|
|
|
|
@dataclass
|
|
class EvalConfig:
|
|
grid_dims: tuple[int, int] = (21, 21)
|
|
truncation: int = 1200
|
|
num_games: int = 100
|
|
agent0: str = "random"
|
|
agent1: str = "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
|
|
|
|
|
|
def main(config: DictConfig):
|
|
key = jrandom.PRNGKey(config.seed)
|
|
key, env_key = jrandom.split(key, 2)
|
|
|
|
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)
|
|
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())
|