95 lines
3.4 KiB
Python
95 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import random
|
|
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import pytest
|
|
|
|
from lost_cities_jax import legal_action_mask, reset_from_order, score, step
|
|
from reference import lost_cities_ref as ref
|
|
from tests.lost_cities_jax.conftest import game_count
|
|
from tests.lost_cities_jax.helpers import shuffled_order
|
|
|
|
JIT_MASK = jax.jit(legal_action_mask)
|
|
JIT_STEP = jax.jit(step)
|
|
JIT_SCORE = jax.jit(score)
|
|
|
|
|
|
@pytest.mark.slow
|
|
def test_jax_matches_reference_random_legal_policy(pytestconfig, tmp_path):
|
|
games = game_count(
|
|
pytestconfig,
|
|
env_name="LOST_CITIES_JAX_DIFF_GAMES",
|
|
local=100,
|
|
ci=100_000,
|
|
full=1_000_000,
|
|
)
|
|
policy_rng = random.Random(20260704)
|
|
|
|
for game_idx in range(games):
|
|
deck_order = shuffled_order(10_000_000 + game_idx)
|
|
jax_state = reset_from_order(jnp.asarray(deck_order, dtype=jnp.int8))
|
|
ref_state = ref.reset_from_order(deck_order)
|
|
actions: list[int] = []
|
|
|
|
for step_idx in range(500):
|
|
jax_mask = list(map(bool, JIT_MASK(jax_state).tolist()))
|
|
ref_mask = ref.legal_action_mask(ref_state)
|
|
if jax_mask != ref_mask:
|
|
_dump_failure(tmp_path, game_idx, step_idx, deck_order, actions)
|
|
diff = [
|
|
idx
|
|
for idx, pair in enumerate(zip(jax_mask, ref_mask, strict=False))
|
|
if pair[0] != pair[1]
|
|
]
|
|
pytest.fail(f"legal mask mismatch game={game_idx} step={step_idx} diff={diff[:20]}")
|
|
|
|
if bool(jax_state.done) != ref_state.done:
|
|
_dump_failure(tmp_path, game_idx, step_idx, deck_order, actions)
|
|
pytest.fail(f"done mismatch game={game_idx} step={step_idx}")
|
|
|
|
if ref_state.done:
|
|
jax_score = [float(value) for value in JIT_SCORE(jax_state).tolist()]
|
|
ref_score = ref.score(ref_state)
|
|
if jax_score != ref_score:
|
|
_dump_failure(tmp_path, game_idx, step_idx, deck_order, actions)
|
|
pytest.fail(
|
|
f"score mismatch game={game_idx} step={step_idx}: {jax_score} != {ref_score}"
|
|
)
|
|
break
|
|
|
|
legal_actions = [idx for idx, legal in enumerate(ref_mask) if legal]
|
|
action = legal_actions[policy_rng.randrange(len(legal_actions))]
|
|
actions.append(action)
|
|
jax_state, jax_reward, jax_done = JIT_STEP(jax_state, jnp.int32(action))
|
|
ref_state, ref_reward, ref_done = ref.step(ref_state, action)
|
|
if bool(jax_done) != ref_done or [float(v) for v in jax_reward.tolist()] != ref_reward:
|
|
_dump_failure(tmp_path, game_idx, step_idx, deck_order, actions)
|
|
pytest.fail(f"step result mismatch game={game_idx} step={step_idx}")
|
|
else:
|
|
_dump_failure(tmp_path, game_idx, 500, deck_order, actions)
|
|
pytest.fail(f"game did not finish within guard loop: game={game_idx}")
|
|
|
|
|
|
def _dump_failure(
|
|
tmp_path,
|
|
game_idx: int,
|
|
step_idx: int,
|
|
deck_order: list[int],
|
|
actions: list[int],
|
|
) -> None:
|
|
path = tmp_path / f"lost_cities_jax_diff_failure_{game_idx}_{step_idx}.json"
|
|
path.write_text(
|
|
json.dumps(
|
|
{
|
|
"game_idx": game_idx,
|
|
"step_idx": step_idx,
|
|
"deck_order": deck_order,
|
|
"actions": actions,
|
|
},
|
|
indent=2,
|
|
)
|
|
)
|