Batch JAX differential verification

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)

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,50 +26,118 @@ 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()))
+            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
+
+            for batch_idx, ref_state in enumerate(ref_states):
+                game_idx = batch_start + batch_idx
                 ref_mask = ref.legal_action_mask(ref_state)
-            if jax_mask != ref_mask:
-                _dump_failure(tmp_path, game_idx, step_idx, deck_order, actions)
+                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(jax_mask, ref_mask, strict=False))
+                        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]}")
+                    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)
+                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:
-                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)
+                    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}: {jax_score} != {ref_score}"
+                            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)
+            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)
+            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}")