70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
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
|