67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from lost_cities_jax import legal_action_mask, reset_from_order, score, step
|
|
from tests.lost_cities_jax.conftest import game_count
|
|
from tests.lost_cities_jax.helpers import (
|
|
assert_state_consistent,
|
|
assert_states_equal,
|
|
shuffled_order,
|
|
)
|
|
|
|
JIT_MASK = jax.jit(legal_action_mask)
|
|
JIT_STEP = jax.jit(step)
|
|
|
|
|
|
@pytest.mark.slow
|
|
def test_random_trajectories_preserve_state_invariants(pytestconfig):
|
|
games = game_count(
|
|
pytestconfig,
|
|
env_name="LOST_CITIES_JAX_PROPERTY_GAMES",
|
|
local=50,
|
|
ci=1_000,
|
|
full=20_000,
|
|
)
|
|
rng = random.Random(777)
|
|
|
|
for game_idx in range(games):
|
|
state = reset_from_order(jnp.asarray(shuffled_order(20_000_000 + game_idx), dtype=jnp.int8))
|
|
assert_state_consistent(state)
|
|
|
|
for _ in range(500):
|
|
if bool(state.done):
|
|
after_done, reward, done = JIT_STEP(state, jnp.int32(0))
|
|
assert_states_equal(after_done, state)
|
|
np.testing.assert_array_equal(np.asarray(reward), np.zeros(2, dtype=np.float32))
|
|
assert bool(done) is True
|
|
break
|
|
|
|
mask = np.asarray(JIT_MASK(state), dtype=bool)
|
|
legal_actions = np.flatnonzero(mask)
|
|
assert len(legal_actions) > 0
|
|
action = int(legal_actions[rng.randrange(len(legal_actions))])
|
|
state, reward, done = JIT_STEP(state, jnp.int32(action))
|
|
assert_state_consistent(state)
|
|
if bool(done):
|
|
np.testing.assert_array_equal(np.asarray(reward), np.asarray(score(state)))
|
|
else:
|
|
pytest.fail(f"game did not terminate within guard loop: game={game_idx}")
|
|
|
|
|
|
def test_forced_termination_scores_current_board():
|
|
state = reset_from_order(jnp.arange(60, dtype=jnp.int8))._replace(
|
|
step_count=jnp.asarray(399, dtype=jnp.int32)
|
|
)
|
|
action = int(np.flatnonzero(np.asarray(legal_action_mask(state), dtype=bool))[0])
|
|
next_state, reward, done = step(state, jnp.int32(action))
|
|
|
|
assert bool(done) is True
|
|
assert int(next_state.step_count) == 400
|
|
np.testing.assert_array_equal(np.asarray(reward), np.asarray(score(next_state)))
|