Add JAX engine verification tests

reference/__init__.py

@@ -0,0 +1 @@
+"""Reference implementations used by tests."""

reference/lost_cities_ref.py

@@ -0,0 +1,247 @@
+"""Independent pure-Python Lost Cities rules reference.
+
+This module intentionally uses ordinary Python containers rather than the JAX
+state representation. It is deterministic under an explicit ``deck_order`` and
+the shared flat action encoding.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from random import Random
+
+N_PLAYERS = 2
+N_COLORS = 5
+CARDS_PER_COLOR = 12
+N_CARDS = N_COLORS * CARDS_PER_COLOR
+HAND_SIZE = 8
+INITIAL_DEAL = N_PLAYERS * HAND_SIZE
+MAX_STEPS = 400
+N_ACTIONS = HAND_SIZE * 2 * 6
+
+LOC_DECK = 0
+LOC_P0_HAND = 1
+LOC_P1_HAND = 2
+LOC_P0_BOARD = 3
+LOC_P1_BOARD = 4
+LOC_DISCARD = 5
+
+NO_CARD = -1
+PLAY = 0
+DISCARD = 1
+DRAW_DECK = 0
+
+
+@dataclass
+class RefState:
+    deck_order: list[int]
+    draw_ptr: int
+    card_loc: list[int]
+    hand_public: list[bool]
+    board: list[list[list[int]]]
+    piles: list[list[int]]
+    to_move: int
+    just_discarded: int
+    step_count: int
+    done: bool
+
+
+def reset(seed: int | None = None) -> RefState:
+    rng = Random(seed)
+    order = list(range(N_CARDS))
+    rng.shuffle(order)
+    return reset_from_order(order)
+
+
+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]:
+        card_loc[card] = LOC_P0_HAND
+    for card in order[HAND_SIZE:INITIAL_DEAL]:
+        card_loc[card] = LOC_P1_HAND
+    return RefState(
+        deck_order=order,
+        draw_ptr=INITIAL_DEAL,
+        card_loc=card_loc,
+        hand_public=[False] * N_CARDS,
+        board=[[[] for _ in range(N_COLORS)] for _ in range(N_PLAYERS)],
+        piles=[[] for _ in range(N_COLORS)],
+        to_move=0,
+        just_discarded=NO_CARD,
+        step_count=0,
+        done=False,
+    )
+
+
+def clone_state(state: RefState) -> RefState:
+    return RefState(
+        deck_order=list(state.deck_order),
+        draw_ptr=state.draw_ptr,
+        card_loc=list(state.card_loc),
+        hand_public=list(state.hand_public),
+        board=[[list(col) for col in player] for player in state.board],
+        piles=[list(pile) for pile in state.piles],
+        to_move=state.to_move,
+        just_discarded=state.just_discarded,
+        step_count=state.step_count,
+        done=state.done,
+    )
+
+
+def decode_action(action: int) -> tuple[int, int, int]:
+    action = int(action)
+    hand_slot = action // 12
+    rem = action % 12
+    place_type = rem // 6
+    draw_source = rem % 6
+    return hand_slot, place_type, draw_source
+
+
+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)
+
+
+def legal_action_mask(state: RefState) -> list[bool]:
+    mask = [False] * N_ACTIONS
+    if state.done:
+        return mask
+
+    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
+
+
+def step(state: RefState, action: int) -> 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]:
+        return clone_state(state), [0.0, 0.0], state.done
+
+    next_state = clone_state(state)
+    player = next_state.to_move
+    hand_slot, place_type, draw_source = decode_action(action)
+    card = hand_cards(next_state, player)[hand_slot]
+    color = card_color(card)
+
+    next_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
+    else:
+        next_state.piles[color].append(card)
+        next_state.card_loc[card] = LOC_DISCARD
+        next_state.just_discarded = card
+
+    if draw_source == DRAW_DECK:
+        drawn = next_state.deck_order[next_state.draw_ptr]
+        next_state.draw_ptr += 1
+        public = False
+    else:
+        src = draw_source - 1
+        drawn = next_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
+    )
+    next_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
+
+
+def can_play_card(state: RefState, player: int, card: int) -> bool:
+    column = state.board[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 rank(card) > top_rank
+
+
+def board_score(state: RefState) -> list[float]:
+    return [float(sum(score_column(column) for column in player)) for player in state.board]
+
+
+def score(state: RefState) -> list[float]:
+    return board_score(state)
+
+
+def score_column(column: list[int]) -> int:
+    if not column:
+        return 0
+    handshakes = sum(1 for card in column if is_handshake(card))
+    rank_sum = sum(rank(card) for card in column if not is_handshake(card))
+    value = (rank_sum - 20) * (1 + handshakes)
+    if len(column) >= 8:
+        value += 20
+    return value
+
+
+def card_color(card: int) -> int:
+    return int(card) // CARDS_PER_COLOR
+
+
+def card_slot(card: int) -> int:
+    return int(card) % CARDS_PER_COLOR
+
+
+def is_handshake(card: int) -> bool:
+    return card_slot(card) < 3
+
+
+def rank(card: int) -> int:
+    return card_slot(card) - 1
+
+
+def _can_draw_after_place(state: RefState, card: int, place_type: int, draw_source: int) -> bool:
+    if draw_source == DRAW_DECK:
+        return True
+    src = draw_source - 1
+    same_discard_pile = place_type == DISCARD and card_color(card) == src
+    after_len = len(state.piles[src]) + int(same_discard_pile)
+    if after_len == 0:
+        return False
+    after_top = card if same_discard_pile else state.piles[src][-1]
+    just_discarded = card if place_type == DISCARD else state.just_discarded
+    return after_top != just_discarded
+
+
+__all__ = [
+    "N_ACTIONS",
+    "RefState",
+    "board_score",
+    "clone_state",
+    "decode_action",
+    "hand_cards",
+    "legal_action_mask",
+    "reset",
+    "reset_from_order",
+    "score",
+    "step",
+]