coolguy/coorl-lost-cities
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Record full JAX differential verification

Browse Source
This commit is contained in:
coolguy
2026-07-04 20:39:01 +09:00
parent 30ccc3cf41
commit 72893250ec
3 changed files with 120 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/lost_cities_jax/test_differential.py
+18 -6
View File
@@ -10,11 +10,13 @@ import numpy as np
import pytest
from lost_cities_jax import batched_legal_mask, batched_reset_from_order, batched_step, score
from lost_cities_jax.types import N_ACTIONS
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
BATCHED_SCORE = jax.jit(jax.vmap(score))
LOW_WORD_MASK = (1 << 64) - 1
@pytest.mark.slow
@@ -26,7 +28,7 @@ def test_jax_matches_reference_random_legal_policy(pytestconfig, tmp_path):
ci=100_000,
full=1_000_000,
)
batch_size = int(os.environ.get("LOST_CITIES_JAX_DIFF_BATCH", "512"))
batch_size = int(os.environ.get("LOST_CITIES_JAX_DIFF_BATCH", "8192"))
policy_rng = random.Random(20260704)
for batch_start in range(0, games, batch_size):
@@ -40,6 +42,7 @@ def test_jax_matches_reference_random_legal_policy(pytestconfig, tmp_path):
for step_idx in range(500):
jax_masks = np.asarray(batched_legal_mask(jax_state), dtype=bool)
jax_mask_low, jax_mask_high = _pack_masks_to_words(jax_masks)
jax_done = np.asarray(jax_state.done, dtype=bool)
actions = np.zeros((current_batch,), dtype=np.int32)
ref_rewards: list[list[float]] = []
@@ -48,8 +51,10 @@ def test_jax_matches_reference_random_legal_policy(pytestconfig, tmp_path):
for batch_idx, ref_state in enumerate(ref_states):
game_idx = batch_start + batch_idx
ref_mask = ref.legal_action_mask(ref_state)
if list(jax_masks[batch_idx]) != ref_mask:
ref_bits = ref.legal_action_bits(ref_state)
if int(jax_mask_low[batch_idx]) != (ref_bits & LOW_WORD_MASK) or int(
jax_mask_high[batch_idx]
) != (ref_bits >> 64):
_dump_failure(
tmp_path,
game_idx,
@@ -57,6 +62,7 @@ def test_jax_matches_reference_random_legal_policy(pytestconfig, tmp_path):
deck_orders[batch_idx],
action_histories[batch_idx],
)
ref_mask = ref.legal_action_mask(ref_state)
diff = [
idx
for idx, pair in enumerate(
@@ -84,11 +90,10 @@ def test_jax_matches_reference_random_legal_policy(pytestconfig, tmp_path):
continue
active += 1
legal_actions = [idx for idx, legal in enumerate(ref_mask) if legal]
action = legal_actions[policy_rng.randrange(len(legal_actions))]
action = ref.nth_legal_action(ref_bits, policy_rng.randrange(ref_bits.bit_count()))
actions[batch_idx] = action
action_histories[batch_idx].append(action)
ref_states[batch_idx], ref_reward, ref_done = ref.step(
ref_states[batch_idx], ref_reward, ref_done = ref.step_in_place(
ref_state, action, validate=False
)
ref_rewards.append(ref_reward)
@@ -141,6 +146,13 @@ def test_jax_matches_reference_random_legal_policy(pytestconfig, tmp_path):
pytest.fail(f"game did not finish within guard loop: game={game_idx}")
def _pack_masks_to_words(masks: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
packed = np.packbits(masks[:, :N_ACTIONS], axis=1, bitorder="little")
low = np.ascontiguousarray(packed[:, :8]).view("<u8").reshape((-1,))
high = np.ascontiguousarray(packed[:, 8:12]).view("<u4").reshape((-1,))
return low, high
def _dump_failure(
tmp_path,
game_idx: int,