From 72893250ec3bac5d8c6d3b62dcae44d6757196ce 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 20:39:01 +0900 Subject: [PATCH] Record full JAX differential verification --- README.md | 15 ++- reference/lost_cities_ref.py | 138 +++++++++++++-------- tests/lost_cities_jax/test_differential.py | 24 +++- 3 files changed, 120 insertions(+), 57 deletions(-) diff --git a/README.md b/README.md index b514728..ea345b0 100644 --- a/README.md +++ b/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 diff --git a/reference/lost_cities_ref.py b/reference/lost_cities_ref.py index 4b9ccbe..647a21f 100644 --- a/reference/lost_cities_ref.py +++ b/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", ] diff --git a/tests/lost_cities_jax/test_differential.py b/tests/lost_cities_jax/test_differential.py index d0ff5eb..82aa3f8 100644 --- a/tests/lost_cities_jax/test_differential.py +++ b/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("