checkpoint

tests/test_checkpoint.py

@@ -0,0 +1,28 @@
+import runpy
+from pathlib import Path
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+
+SCRIPT_PATH = Path(__file__).parents[1] / "scripts" / "train.py"
+
+
+def test_initialize_network_restores_serialized_weights(tmp_path):
+    script = runpy.run_path(str(SCRIPT_PATH), run_name="train_script")
+    initialize_network = script["initialize_network"]
+
+    original = initialize_network(jrandom.PRNGKey(0))
+    checkpoint_path = tmp_path / "model.eqx"
+    eqx.tree_serialise_leaves(checkpoint_path, original)
+
+    restored = initialize_network(jrandom.PRNGKey(1), str(checkpoint_path))
+    original_leaves = jax.tree.leaves(eqx.filter(original, eqx.is_array))
+    restored_leaves = jax.tree.leaves(eqx.filter(restored, eqx.is_array))
+
+    assert len(original_leaves) == len(restored_leaves)
+    assert all(
+        jnp.array_equal(original_leaf, restored_leaf)
+        for original_leaf, restored_leaf in zip(original_leaves, restored_leaves, strict=True)
+    )

tests/test_mcts.py

@@ -0,0 +1,70 @@
+import jax.numpy as jnp
+import numpy as np
+from generals.core import game
+
+from general_bots_training.mcts import (
+    build_cost_grid,
+    competition_step,
+    observation_from_wire,
+    sample_determinization,
+    valid_build_mask,
+)
+
+
+def make_observation(armies=60):
+    types = np.ones((5, 5), dtype=np.int32)
+    owners = np.zeros((5, 5), dtype=np.int32)
+    army_grid = np.zeros((5, 5), dtype=np.int32)
+    types[2, 2] = 4
+    owners[2, 2] = 1
+    army_grid[2, 2] = armies
+    types[0, 0] = 0
+    return observation_from_wire(100, 1, armies, 1, 20, types, owners, army_grid)
+
+
+def test_wire_observation_does_not_treat_fog_as_neutral():
+    observation = make_observation()
+
+    assert bool(observation.fog_cells[0, 0])
+    assert not bool(observation.neutral_cells[0, 0])
+    assert bool(observation.neutral_cells[0, 1])
+
+
+def test_observation_build_cost_and_legality_match_rules():
+    observation = make_observation()
+    costs = build_cost_grid(observation)
+
+    assert int(costs[2, 2]) == 49
+    assert int(costs[2, 3]) == 47
+    assert not bool(valid_build_mask(observation)[2, 2])
+
+    armies = observation.armies.at[2, 3].set(47)
+    owned = observation.owned_cells.at[2, 3].set(True)
+    observation = observation._replace(armies=armies, owned_cells=owned)
+    assert bool(valid_build_mask(observation)[2, 3])
+
+
+def test_determinization_matches_observed_global_totals():
+    observation = make_observation()
+    state = sample_determinization(observation, 0, np.random.default_rng(0))
+
+    assert int(state.ownership[0].sum()) == 1
+    assert int(state.ownership[1].sum()) == 1
+    assert int((state.armies * state.ownership[1]).sum()) == 20
+    assert bool(state.ownership[0, 2, 2])
+
+
+def test_competition_step_applies_build_action():
+    grid = jnp.zeros((5, 5), dtype=jnp.int32).at[2, 2].set(1).at[4, 4].set(2)
+    state = game.create_initial_state(grid)
+    state = state._replace(
+        armies=state.armies.at[2, 3].set(60).at[4, 4].set(10),
+        ownership=state.ownership.at[0, 2, 3].set(True),
+        ownership_neutral=state.ownership_neutral.at[2, 3].set(False),
+    )
+    actions = jnp.array([[2, 2, 3, 0, 0], [1, 0, 0, 0, 0]], dtype=jnp.int32)
+
+    new_state, _ = competition_step(state, actions)
+
+    assert bool(new_state.castles[2, 3])
+    assert int(new_state.armies[2, 3]) == 13

tests/test_opponents.py

@@ -0,0 +1,30 @@
+import pytest
+from generals.agents import ExpanderAgent, HunterAgent, RandomAgent
+
+from general_bots_training.opponents import SelfPlayOpponent, make_opponent
+
+
+@pytest.mark.parametrize(
+    ("name", "expected_type"),
+    [
+        ("random", RandomAgent),
+        ("expander", ExpanderAgent),
+        ("hunter", HunterAgent),
+    ],
+)
+def test_make_opponent(name, expected_type):
+    assert isinstance(make_opponent(name), expected_type)
+
+
+def test_make_opponent_is_case_insensitive():
+    assert isinstance(make_opponent("HUNTER"), HunterAgent)
+
+
+def test_make_opponent_supports_self_play_aliases():
+    assert isinstance(make_opponent("self_play"), SelfPlayOpponent)
+    assert isinstance(make_opponent("self-play"), SelfPlayOpponent)
+
+
+def test_make_opponent_rejects_unknown_name():
+    with pytest.raises(ValueError, match="unknown opponent"):
+        make_opponent("turtle")

tests/test_self_play.py

@@ -0,0 +1,28 @@
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+from generals.core.env import GeneralsEnv
+
+from general_bots_training.network import PolicyValueNetwork
+from general_bots_training.opponents import SelfPlayOpponent
+from general_bots_training.rollout import make_collect_rollout
+
+
+def test_self_play_collects_both_player_perspectives():
+    key = jrandom.PRNGKey(0)
+    key, network_key, pool_key, state_key = jrandom.split(key, 4)
+    env = GeneralsEnv(grid_dims=(4, 4), truncation=20, pool_size=8)
+    pool, _ = env.reset(pool_key)
+    states = jax.vmap(env.init_state)(jrandom.split(state_key, 2))
+    states = states._replace(pool_idx=jnp.arange(2, dtype=states.pool_idx.dtype))
+    network = PolicyValueNetwork(network_key)
+
+    _, transitions, (_, last_next_obs) = make_collect_rollout(env, 1, SelfPlayOpponent())(
+        states, pool, network, key
+    )
+
+    assert transitions["obs"].shape == (1, 4, 14, 4, 4)
+    assert transitions["action"].shape == (1, 4, 5)
+    assert transitions["player"][0].tolist() == [0, 0, 1, 1]
+    assert transitions["done"].shape == (1, 4)
+    assert last_next_obs.shape == (4, 14, 4, 4)

tests/test_train_config.py

@@ -0,0 +1,38 @@
+import runpy
+from pathlib import Path
+
+import pytest
+from omegaconf.errors import ConfigKeyError
+
+SCRIPT_PATH = Path(__file__).parents[1] / "scripts" / "train.py"
+
+
+def load_config(args):
+    script = runpy.run_path(str(SCRIPT_PATH), run_name="train_script")
+    return script["load_config"](args)
+
+
+def test_cli_values_override_structured_defaults():
+    config = load_config(
+        [
+            "num_envs=32",
+            "lr=1e-4",
+            "grid_dims=[6,6]",
+            "checkpoint_path=models/test.eqx",
+            "opponent=hunter",
+            "resume_from=models/previous.eqx",
+        ]
+    )
+
+    assert config.num_envs == 32
+    assert config.lr == 1e-4
+    assert list(config.grid_dims) == [6, 6]
+    assert config.checkpoint_path == "models/test.eqx"
+    assert config.opponent == "hunter"
+    assert config.resume_from == "models/previous.eqx"
+    assert config.rollout_steps == 256
+
+
+def test_unknown_cli_key_is_rejected():
+    with pytest.raises(ConfigKeyError):
+        load_config(["unknown_option=1"])

tests/test_training.py

@@ -0,0 +1,53 @@
+import jax.numpy as jnp
+import jax.random as jrandom
+
+from general_bots_training.network import PolicyValueNetwork, obs_to_tensor
+from general_bots_training.ppo import compute_advantages_and_returns
+
+
+def test_returns_use_raw_advantages_before_policy_normalization():
+    rewards = jnp.array([[1.0, 3.0]])
+    values = jnp.zeros_like(rewards)
+    next_value = jnp.zeros(2)
+    dones = jnp.ones_like(rewards, dtype=bool)
+
+    advantages, returns = compute_advantages_and_returns(rewards, values, next_value, dones)
+
+    assert jnp.allclose(advantages, jnp.array([[-1.0, 1.0]]))
+    assert jnp.allclose(returns, rewards)
+
+
+def test_pass_is_one_global_action_and_round_trips():
+    network = PolicyValueNetwork(jrandom.PRNGKey(0))
+    obs = jnp.zeros((14, 4, 4))
+    mask = jnp.zeros((4, 4, 4), dtype=bool)
+
+    action, _, sampled_logprob, entropy = network(obs, mask, jrandom.PRNGKey(1))
+    _, _, evaluated_logprob, _ = network(obs, mask, jrandom.PRNGKey(2), action)
+
+    assert action.tolist() == [1, 0, 0, 0, 0]
+    assert jnp.allclose(sampled_logprob, evaluated_logprob)
+    assert jnp.isclose(entropy, 0.0)
+
+
+def test_observation_encoder_includes_normalized_timestep():
+    class Observation:
+        armies = jnp.zeros((2, 3), dtype=jnp.int32)
+        generals = jnp.zeros((2, 3), dtype=bool)
+        castles = jnp.zeros((2, 3), dtype=bool)
+        mountains = jnp.zeros((2, 3), dtype=bool)
+        neutral_cells = jnp.ones((2, 3), dtype=bool)
+        owned_cells = jnp.zeros((2, 3), dtype=bool)
+        opponent_cells = jnp.zeros((2, 3), dtype=bool)
+        fog_cells = jnp.zeros((2, 3), dtype=bool)
+        structures_in_fog = jnp.zeros((2, 3), dtype=bool)
+        owned_land_count = jnp.array(0)
+        owned_army_count = jnp.array(0)
+        opponent_land_count = jnp.array(0)
+        opponent_army_count = jnp.array(0)
+        timestep = jnp.array(1200)
+
+    encoded = obs_to_tensor(Observation())
+
+    assert encoded.shape == (14, 2, 3)
+    assert jnp.allclose(encoded[-1], 1.0)