From 30ccc3cf41b5c2c6ff294e2f30aaf0afb0fc8517 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= Date: Sat, 4 Jul 2026 19:42:05 +0900 Subject: [PATCH] Batch JAX differential verification --- reference/lost_cities_ref.py | 6 +- tests/lost_cities_jax/test_differential.py | 148 +++++++++++++++------ 2 files changed, 112 insertions(+), 42 deletions(-) diff --git a/reference/lost_cities_ref.py b/reference/lost_cities_ref.py index 45e7887..4b9ccbe 100644 --- a/reference/lost_cities_ref.py +++ b/reference/lost_cities_ref.py @@ -129,10 +129,12 @@ def legal_action_mask(state: RefState) -> list[bool]: return mask -def step(state: RefState, action: int) -> tuple[RefState, list[float], bool]: +def step( + state: RefState, action: int, *, validate: bool = True +) -> tuple[RefState, list[float], bool]: if state.done or action < 0 or action >= N_ACTIONS: return clone_state(state), [0.0, 0.0], state.done - if not legal_action_mask(state)[action]: + if validate and not legal_action_mask(state)[action]: return clone_state(state), [0.0, 0.0], state.done next_state = clone_state(state) diff --git a/tests/lost_cities_jax/test_differential.py b/tests/lost_cities_jax/test_differential.py index 97e4b06..d0ff5eb 100644 --- a/tests/lost_cities_jax/test_differential.py +++ b/tests/lost_cities_jax/test_differential.py @@ -1,20 +1,20 @@ from __future__ import annotations import json +import os 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 lost_cities_jax import batched_legal_mask, batched_reset_from_order, batched_step, score 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) +BATCHED_SCORE = jax.jit(jax.vmap(score)) @pytest.mark.slow @@ -26,51 +26,119 @@ 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")) 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 batch_start in range(0, games, batch_size): + current_batch = min(batch_size, games - batch_start) + deck_orders = [ + shuffled_order(10_000_000 + batch_start + idx) for idx in range(current_batch) + ] + jax_state = batched_reset_from_order(jnp.asarray(deck_orders, dtype=jnp.int8)) + ref_states = [ref.reset_from_order(deck_order) for deck_order in deck_orders] + action_histories: list[list[int]] = [[] for _ in range(current_batch)] 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]}") + jax_masks = np.asarray(batched_legal_mask(jax_state), dtype=bool) + jax_done = np.asarray(jax_state.done, dtype=bool) + actions = np.zeros((current_batch,), dtype=np.int32) + ref_rewards: list[list[float]] = [] + ref_done_values: list[bool] = [] + active = 0 - 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}" + 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: + _dump_failure( + tmp_path, + game_idx, + step_idx, + deck_orders[batch_idx], + action_histories[batch_idx], ) + diff = [ + idx + for idx, pair in enumerate( + zip(list(jax_masks[batch_idx]), 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_done[batch_idx]) != ref_state.done: + _dump_failure( + tmp_path, + game_idx, + step_idx, + deck_orders[batch_idx], + action_histories[batch_idx], + ) + pytest.fail(f"done mismatch game={game_idx} step={step_idx}") + + if ref_state.done: + ref_rewards.append([0.0, 0.0]) + ref_done_values.append(True) + continue + + active += 1 + legal_actions = [idx for idx, legal in enumerate(ref_mask) if legal] + action = legal_actions[policy_rng.randrange(len(legal_actions))] + actions[batch_idx] = action + action_histories[batch_idx].append(action) + ref_states[batch_idx], ref_reward, ref_done = ref.step( + ref_state, action, validate=False + ) + ref_rewards.append(ref_reward) + ref_done_values.append(ref_done) + + if active == 0: + jax_scores = np.asarray(BATCHED_SCORE(jax_state), dtype=np.float32) + for batch_idx, ref_state in enumerate(ref_states): + ref_score = np.asarray(ref.score(ref_state), dtype=np.float32) + if not np.array_equal(jax_scores[batch_idx], ref_score): + game_idx = batch_start + batch_idx + _dump_failure( + tmp_path, + game_idx, + step_idx, + deck_orders[batch_idx], + action_histories[batch_idx], + ) + pytest.fail( + f"score mismatch game={game_idx} step={step_idx}: " + f"{jax_scores[batch_idx].tolist()} != {ref_score.tolist()}" + ) 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}") + jax_state, jax_reward, jax_done_after = batched_step(jax_state, jnp.asarray(actions)) + jax_reward = np.asarray(jax_reward, dtype=np.float32) + jax_done_after = np.asarray(jax_done_after, dtype=bool) + for batch_idx, (ref_reward, ref_done) in enumerate( + zip(ref_rewards, ref_done_values, strict=False) + ): + if bool(jax_done_after[batch_idx]) != ref_done or not np.array_equal( + jax_reward[batch_idx], np.asarray(ref_reward, dtype=np.float32) + ): + game_idx = batch_start + batch_idx + _dump_failure( + tmp_path, + game_idx, + step_idx, + deck_orders[batch_idx], + action_histories[batch_idx], + ) + 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}") + for batch_idx, ref_state in enumerate(ref_states): + if not ref_state.done: + game_idx = batch_start + batch_idx + _dump_failure( + tmp_path, game_idx, 500, deck_orders[batch_idx], action_histories[batch_idx] + ) + pytest.fail(f"game did not finish within guard loop: game={game_idx}") def _dump_failure(