Record full JAX differential verification

README.md

@@ -143,6 +143,17 @@ CI=1 uv run pytest -q tests/lost_cities_jax/test_differential.py
 uv run pytest -q tests/lost_cities_jax/test_differential.py --full
 ```
 
+Observed differential results on 2026-07-04 with CPU JAX backend:
+
+```text
+CI=1 ... test_differential.py
+1 passed in 247.61s (0:04:07)
+
+... test_differential.py --full
+1 passed in 2514.14s (0:41:54)
+elapsed=41:54.48
+```
+
 Throughput benchmark:
 
 ```bash
@@ -155,8 +166,8 @@ Measured on 2026-07-04 with CPU JAX backend:
 backend=cpu
 batch_size=8192
 steps=256
-elapsed_sec=4.944725
-steps_per_sec=424119.01
+elapsed_sec=4.655526
+steps_per_sec=450465.14
 ```
 
 ### DECISIONS.md

reference/lost_cities_ref.py

@@ -7,6 +7,7 @@ the shared flat action encoding.
 
 from __future__ import annotations
 
+from bisect import insort
 from dataclasses import dataclass
 from random import Random
 
@@ -38,7 +39,9 @@ class RefState:
     draw_ptr: int
     card_loc: list[int]
     hand_public: list[bool]
+    hands: list[list[int]]
     board: list[list[list[int]]]
+    col_top: list[list[int]]
     piles: list[list[int]]
     to_move: int
     just_discarded: int
@@ -56,16 +59,20 @@ def reset(seed: int | None = None) -> RefState:
 def reset_from_order(deck_order: list[int] | tuple[int, ...]) -> RefState:
     order = [int(card) for card in deck_order]
     card_loc = [LOC_DECK] * N_CARDS
-    for card in order[:HAND_SIZE]:
+    p0_hand = sorted(order[:HAND_SIZE])
+    p1_hand = sorted(order[HAND_SIZE:INITIAL_DEAL])
+    for card in p0_hand:
         card_loc[card] = LOC_P0_HAND
-    for card in order[HAND_SIZE:INITIAL_DEAL]:
+    for card in p1_hand:
         card_loc[card] = LOC_P1_HAND
     return RefState(
         deck_order=order,
         draw_ptr=INITIAL_DEAL,
         card_loc=card_loc,
         hand_public=[False] * N_CARDS,
+        hands=[p0_hand, p1_hand],
         board=[[[] for _ in range(N_COLORS)] for _ in range(N_PLAYERS)],
+        col_top=[[0 for _ in range(N_COLORS)] for _ in range(N_PLAYERS)],
         piles=[[] for _ in range(N_COLORS)],
         to_move=0,
         just_discarded=NO_CARD,
@@ -80,7 +87,9 @@ def clone_state(state: RefState) -> RefState:
         draw_ptr=state.draw_ptr,
         card_loc=list(state.card_loc),
         hand_public=list(state.hand_public),
+        hands=[list(hand) for hand in state.hands],
         board=[[list(col) for col in player] for player in state.board],
+        col_top=[list(player) for player in state.col_top],
         piles=[list(pile) for pile in state.piles],
         to_move=state.to_move,
         just_discarded=state.just_discarded,
@@ -101,32 +110,50 @@ def decode_action(action: int) -> tuple[int, int, int]:
 def hand_cards(state: RefState, player: int | None = None) -> list[int]:
     if player is None:
         player = state.to_move
-    hand_loc = LOC_P0_HAND + player
-    return sorted(card for card, loc in enumerate(state.card_loc) if loc == hand_loc)
+    return list(state.hands[player])
 
 
 def legal_action_mask(state: RefState) -> list[bool]:
-    mask = [False] * N_ACTIONS
+    bits = legal_action_bits(state)
+    return [bool(bits & (1 << action)) for action in range(N_ACTIONS)]
+
+
+def legal_action_bits(state: RefState) -> int:
+    bits = 0
     if state.done:
-        return mask
+        return bits
 
     player = state.to_move
-    hand = hand_cards(state, player)
-    for hand_slot in range(HAND_SIZE):
-        if hand_slot >= len(hand):
-            continue
-        card = hand[hand_slot]
-        for place_type in (PLAY, DISCARD):
-            if place_type == PLAY:
-                place_ok = can_play_card(state, player, card)
-            else:
-                place_ok = True
-            if not place_ok:
-                continue
-            for draw_source in range(6):
-                if _can_draw_after_place(state, card, place_type, draw_source):
-                    mask[hand_slot * 12 + place_type * 6 + draw_source] = True
-    return mask
+    hand = state.hands[player]
+    draw_bits = 1
+    for color, pile in enumerate(state.piles):
+        if pile and pile[-1] != state.just_discarded:
+            draw_bits |= 1 << (color + 1)
+
+    for hand_slot, card in enumerate(hand[:HAND_SIZE]):
+        base = hand_slot * 12
+        if can_play_card(state, player, card):
+            bits |= draw_bits << base
+
+        color = card_color(card)
+        discard_bits = draw_bits & ~(1 << (color + 1))
+        bits |= discard_bits << (base + 6)
+    return bits
+
+
+def nth_legal_action(bits: int, index: int) -> int:
+    """Return the ``index``-th set action bit from low to high."""
+
+    remaining = int(index)
+    action = 0
+    while bits:
+        if bits & 1:
+            if remaining == 0:
+                return action
+            remaining -= 1
+        action += 1
+        bits >>= 1
+    raise IndexError(index)
 
 
 def step(
@@ -138,51 +165,61 @@ def step(
         return clone_state(state), [0.0, 0.0], state.done
 
     next_state = clone_state(state)
-    player = next_state.to_move
+    _, reward, done = step_in_place(next_state, action, validate=False)
+    return next_state, reward, done
+
+
+def step_in_place(
+    state: RefState, action: int, *, validate: bool = True
+) -> tuple[RefState, list[float], bool]:
+    if state.done or action < 0 or action >= N_ACTIONS:
+        return state, [0.0, 0.0], state.done
+    if validate and not legal_action_mask(state)[action]:
+        return state, [0.0, 0.0], state.done
+
+    player = state.to_move
     hand_slot, place_type, draw_source = decode_action(action)
-    card = hand_cards(next_state, player)[hand_slot]
+    hand = state.hands[player]
+    card = hand.pop(hand_slot)
     color = card_color(card)
 
-    next_state.hand_public[card] = False
+    state.hand_public[card] = False
     if place_type == PLAY:
-        next_state.board[player][color].append(card)
-        next_state.card_loc[card] = LOC_P0_BOARD + player
-        next_state.just_discarded = NO_CARD
+        state.board[player][color].append(card)
+        if not is_handshake(card):
+            state.col_top[player][color] = rank(card)
+        state.card_loc[card] = LOC_P0_BOARD + player
     else:
-        next_state.piles[color].append(card)
-        next_state.card_loc[card] = LOC_DISCARD
-        next_state.just_discarded = card
+        state.piles[color].append(card)
+        state.card_loc[card] = LOC_DISCARD
 
     if draw_source == DRAW_DECK:
-        drawn = next_state.deck_order[next_state.draw_ptr]
-        next_state.draw_ptr += 1
+        drawn = state.deck_order[state.draw_ptr]
+        state.draw_ptr += 1
         public = False
     else:
         src = draw_source - 1
-        drawn = next_state.piles[src].pop()
+        drawn = state.piles[src].pop()
         public = True
 
-    next_state.card_loc[drawn] = LOC_P0_HAND + player
-    next_state.hand_public[drawn] = public
-    next_state.just_discarded = NO_CARD
-    next_state.step_count += 1
-    next_state.done = (draw_source == DRAW_DECK and next_state.draw_ptr >= N_CARDS) or (
-        next_state.step_count >= MAX_STEPS
+    insort(hand, drawn)
+    state.card_loc[drawn] = LOC_P0_HAND + player
+    state.hand_public[drawn] = public
+    state.just_discarded = NO_CARD
+    state.step_count += 1
+    state.done = (draw_source == DRAW_DECK and state.draw_ptr >= N_CARDS) or (
+        state.step_count >= MAX_STEPS
     )
-    next_state.to_move = 1 - player
+    state.to_move = 1 - player
 
-    reward = board_score(next_state) if next_state.done else [0.0, 0.0]
-    return next_state, reward, next_state.done
+    reward = board_score(state) if state.done else [0.0, 0.0]
+    return state, reward, state.done
 
 
 def can_play_card(state: RefState, player: int, card: int) -> bool:
-    column = state.board[player][card_color(card)]
+    top_rank = state.col_top[player][card_color(card)]
     if is_handshake(card):
-        return not any(not is_handshake(played) for played in column)
-    top_rank = 0
-    for played in column:
-        if not is_handshake(played):
-            top_rank = rank(played)
+        return top_rank == 0
     return rank(card) > top_rank
 
 
@@ -241,9 +278,12 @@ __all__ = [
     "clone_state",
     "decode_action",
     "hand_cards",
+    "legal_action_bits",
     "legal_action_mask",
+    "nth_legal_action",
     "reset",
     "reset_from_order",
     "score",
     "step",
+    "step_in_place",
 ]

tests/lost_cities_jax/test_differential.py

@@ -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,