checkpoint

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