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