Record full JAX differential verification

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",
 ]