From 7cd299dd1ee66d63b6afbf0c3fec5b555593bcd9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= Date: Wed, 15 Jul 2026 02:05:50 +0900 Subject: [PATCH] Add the classic three-round match layer Wraps the single-round engine rather than changing it: engine.py is the rules oracle the TypeScript client is differential-tested against, and the round itself is unchanged by the match. Only two rules live up here, both from the Kosmos rulebook -- three rounds decided on the summed total, and "the player who has more points begins" the next one. That is not alternating, and it is not what reset_from_order hardcoded, so it takes a first_player argument now. The rulebook says nothing about an exact tie, so the starter falls back to a coin flip. A deterministic tie-break would give one seat a standing edge in symmetric self-play and the agent would learn to steer for it. Round one needs no special case: carry is (0, 0) there, so the tie branch already yields the coin, which is exactly the rulebook's arbitrary "oldest player begins". All randomness -- three deals and three coins -- is drawn in match_reset and stored in the state, so match_step stays deterministic and needs no PRNG key threaded through every rollout, eval, and gate body. It also makes a mirrored match (same deals, seats swapped, same coins) a pure seat relabel, which the antithetic pairing later depends on. Tests cover the deck clock (44 deck draws a round, discard draws extend it), carry banking each round exactly once, the start-player rule across all three branches, a fair round-one coin, mirror symmetry, and that the running total does not jump across a round boundary -- the last one matters because potential shaping will be built on it. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh --- src/lost_cities_jax/engine.py | 7 +- src/lost_cities_jax/match.py | 139 +++++++++++++++++++++ tests/lost_cities_jax/test_match.py | 179 ++++++++++++++++++++++++++++ 3 files changed, 323 insertions(+), 2 deletions(-) create mode 100644 src/lost_cities_jax/match.py create mode 100644 tests/lost_cities_jax/test_match.py diff --git a/src/lost_cities_jax/engine.py b/src/lost_cities_jax/engine.py index 033bad8..1a916e2 100644 --- a/src/lost_cities_jax/engine.py +++ b/src/lost_cities_jax/engine.py @@ -36,11 +36,14 @@ def reset(rng: jax.Array) -> State: return reset_from_order(deck_order) -def reset_from_order(deck_order: jax.Array) -> State: +def reset_from_order(deck_order: jax.Array, first_player: jax.Array | int = 0) -> State: """Return an initial state using an explicit 60-card deck permutation. The first eight cards are dealt to player 0 and the next eight to player 1. The remaining cards are drawn in array order starting at ``draw_ptr == 16``. + + ``first_player`` moves first. Rounds two and three of a classic match are led + by whoever is ahead on points, so the match layer sets this per round. """ deck_order = jnp.asarray(deck_order, dtype=jnp.int8) @@ -60,7 +63,7 @@ def reset_from_order(deck_order: jax.Array) -> State: col_len=jnp.zeros((2, N_COLORS), dtype=jnp.int8), pile=jnp.full((N_COLORS, MAX_PILE_SIZE), NO_CARD, dtype=jnp.int8), pile_len=jnp.zeros((N_COLORS,), dtype=jnp.int8), - to_move=jnp.asarray(0, dtype=jnp.int8), + to_move=jnp.asarray(first_player, dtype=jnp.int8), just_discarded=jnp.asarray(NO_CARD, dtype=jnp.int8), step_count=jnp.asarray(0, dtype=jnp.int32), done=jnp.asarray(False), diff --git a/src/lost_cities_jax/match.py b/src/lost_cities_jax/match.py new file mode 100644 index 0000000..7e5cfa0 --- /dev/null +++ b/src/lost_cities_jax/match.py @@ -0,0 +1,139 @@ +"""Classic Lost Cities: three rounds, scores summed, highest total wins. + +The single-round engine in ``engine.py`` is left alone -- it is the rules oracle +the TypeScript client is differential-tested against. This module wraps it. + +Two rules only exist at the match level, both from the Kosmos rulebook: + +- "If after three games you have the highest overall score, you win." +- "The player who has more points begins" the next game. Not alternating. The + rulebook says nothing about an exact tie, so we flip a coin: any deterministic + tie-break would hand one seat a standing edge in symmetric self-play, and the + agent would learn to steer for it. + +All randomness is drawn in ``match_reset`` and stored in the state, so +``match_step`` stays deterministic and needs no PRNG key. That keeps the step +signature clean through every rollout/eval body, and it means a mirrored match +(same three deals, seats swapped, same coin flips) is just a seat relabel. +""" + +from __future__ import annotations + +from typing import NamedTuple + +import jax +import jax.numpy as jnp + +from lost_cities_jax.engine import board_score, legal_action_mask, reset_from_order, step +from lost_cities_jax.types import N_CARDS, State + +N_ROUNDS = 3 + + +class MatchState(NamedTuple): + round: State + """The round being played right now.""" + + deck_orders: jax.Array + """(N_ROUNDS, N_CARDS) -- every deal of the match, shuffled up front.""" + + coin_flips: jax.Array + """(N_ROUNDS,) tie-breaking starters, used only when the scores are level.""" + + round_idx: jax.Array + carry: jax.Array + """(2,) points banked by each player in the rounds already finished.""" + + done: jax.Array + + +def match_reset(rng: jax.Array) -> MatchState: + deck_key, coin_key = jax.random.split(rng) + deck_keys = jax.random.split(deck_key, N_ROUNDS) + deck_orders = jax.vmap( + lambda key: jax.random.permutation(key, jnp.arange(N_CARDS, dtype=jnp.int32), axis=0) + )(deck_keys).astype(jnp.int8) + coin_flips = jax.random.bernoulli(coin_key, 0.5, shape=(N_ROUNDS,)).astype(jnp.int8) + return match_reset_from(deck_orders, coin_flips) + + +def match_reset_from(deck_orders: jax.Array, coin_flips: jax.Array) -> MatchState: + """Build a match from explicit deals and coin flips (mirrored eval, fixtures).""" + + deck_orders = jnp.asarray(deck_orders, dtype=jnp.int8) + coin_flips = jnp.asarray(coin_flips, dtype=jnp.int8) + carry = jnp.zeros((2,), dtype=jnp.int32) + round_idx = jnp.asarray(0, dtype=jnp.int32) + return MatchState( + round=reset_from_order(deck_orders[0], first_player=starting_player(carry, 0, coin_flips)), + deck_orders=deck_orders, + coin_flips=coin_flips, + round_idx=round_idx, + carry=carry, + done=jnp.asarray(False), + ) + + +def starting_player(carry: jax.Array, round_idx: jax.Array, coin_flips: jax.Array) -> jax.Array: + """Whoever has banked more points leads; level scores fall back to the coin. + + Round one needs no special case: ``carry`` is (0, 0) there, so the tie branch + already picks the coin flip, which is exactly the rulebook's arbitrary + "oldest player begins". + """ + + lead = carry[0] - carry[1] + coin = coin_flips[round_idx].astype(jnp.int8) + ahead = jnp.where(lead > 0, jnp.int8(0), jnp.int8(1)) + return jnp.where(lead == 0, coin, ahead) + + +def match_score(state: MatchState) -> jax.Array: + """(2,) running totals: rounds already banked plus the board in play.""" + + return state.carry + board_score(state.round).astype(jnp.int32) + + +def match_legal_action_mask(state: MatchState) -> jax.Array: + return legal_action_mask(state.round) & ~state.done + + +def match_step(state: MatchState, action: jax.Array) -> tuple[MatchState, jax.Array, jax.Array]: + """Play one ply. Rolls into the next round when the deck runs out. + + Reward is zero until the third round ends, then it is each player's match + total. Intermediate rounds pay nothing: only the sum decides the match. + """ + + played, _, _ = step(state.round, action) + round_over = played.done & ~state.done + is_final_round = state.round_idx >= N_ROUNDS - 1 + + advance = round_over & ~is_final_round + finished = round_over & is_final_round + + banked = state.carry + board_score(played).astype(jnp.int32) + next_carry = jnp.where(advance, banked, state.carry) + next_round_idx = jnp.where(advance, state.round_idx + 1, state.round_idx) + + # Dealt eagerly every ply; only kept on the plies that actually roll over. + fresh = reset_from_order( + state.deck_orders[next_round_idx], + first_player=starting_player(next_carry, next_round_idx, state.coin_flips), + ) + next_round = jax.tree_util.tree_map( + lambda new, old: jnp.where(advance, new, old), fresh, played + ) + + next_state = MatchState( + round=next_round, + deck_orders=state.deck_orders, + coin_flips=state.coin_flips, + round_idx=next_round_idx, + carry=next_carry, + done=state.done | finished, + ) + + total = state.carry + board_score(played).astype(jnp.int32) + reward = jnp.where(finished, total.astype(jnp.float32), jnp.zeros((2,), dtype=jnp.float32)) + return next_state, reward, next_state.done diff --git a/tests/lost_cities_jax/test_match.py b/tests/lost_cities_jax/test_match.py new file mode 100644 index 0000000..b704c38 --- /dev/null +++ b/tests/lost_cities_jax/test_match.py @@ -0,0 +1,179 @@ +"""Phase 1: the classic three-round match layer.""" + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from lost_cities_jax.engine import board_score +from lost_cities_jax.match import ( + N_ROUNDS, + MatchState, + match_legal_action_mask, + match_reset, + match_reset_from, + match_score, + match_step, + starting_player, +) +from lost_cities_jax.opponents import random_legal_action +from lost_cities_jax.types import DECK_DRAWS, DRAW_DECK, N_CARDS + +MATCH_STEP_CAP = 1400 # 3 rounds x the engine's 400-ply safety cap, with slack + + +def _decks(seed: int) -> jnp.ndarray: + rng = np.random.default_rng(seed) + return jnp.asarray( + np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)]), dtype=jnp.int8 + ) + + +def _play_match(state: MatchState, key: jax.Array): + """Drive a match to completion with random legal play, recording each ply.""" + plies = [] + while not bool(state.done): + key, step_key = jax.random.split(key) + action = random_legal_action(state.round, state.round.to_move, step_key) + assert bool(match_legal_action_mask(state)[action]) + prev = state + state, reward, _ = match_step(state, action) + plies.append((prev, int(action), state, np.asarray(reward))) + assert len(plies) < MATCH_STEP_CAP, "match failed to terminate" + return state, plies + + +# --- match structure ------------------------------------------------------- + + +def test_match_plays_exactly_three_rounds_and_then_ends(): + state = match_reset(jax.random.PRNGKey(0)) + final, plies = _play_match(state, jax.random.PRNGKey(1)) + + assert bool(final.done) + assert int(final.round_idx) == N_ROUNDS - 1 + # round_idx advances 0 -> 1 -> 2 and stops. + seen = {int(before.round_idx) for before, _, _, _ in plies} + assert seen == {0, 1, 2} + + +def test_reward_is_paid_only_when_the_third_round_ends(): + state = match_reset(jax.random.PRNGKey(2)) + final, plies = _play_match(state, jax.random.PRNGKey(3)) + + rewards = np.stack([reward for _, _, _, reward in plies]) + paid = np.flatnonzero(np.any(rewards != 0.0, axis=1)) + # Intermediate rounds bank points into carry but pay nothing. + assert paid.tolist() == [len(plies) - 1] + assert np.allclose(rewards[-1], np.asarray(match_score(final), dtype=np.float32)) + + +def test_each_round_runs_44_deck_draws_plus_one_ply_per_discard_draw(): + """The deck clock: only deck draws end a round, so discard draws extend it.""" + state = match_reset(jax.random.PRNGKey(4)) + _, plies = _play_match(state, jax.random.PRNGKey(5)) + + per_round: dict[int, list[int]] = {0: [], 1: [], 2: []} + for before, action, _, _ in plies: + per_round[int(before.round_idx)].append(action % 6) + + for round_idx, draws in per_round.items(): + deck_draws = sum(1 for src in draws if src == DRAW_DECK) + discard_draws = len(draws) - deck_draws + assert deck_draws == DECK_DRAWS, f"round {round_idx} drew {deck_draws} deck cards" + assert len(draws) == DECK_DRAWS + discard_draws + + +# --- carry ----------------------------------------------------------------- + + +def test_carry_banks_each_finished_round_exactly_once(): + state = match_reset(jax.random.PRNGKey(6)) + final, plies = _play_match(state, jax.random.PRNGKey(7)) + + deltas = [ + np.asarray(after.carry) - np.asarray(before.carry) + for before, _, after, _ in plies + if int(after.round_idx) != int(before.round_idx) + ] + + # Two roll-overs for three rounds, and carry is exactly their sum. + assert len(deltas) == N_ROUNDS - 1 + assert np.array_equal(np.asarray(final.carry), sum(deltas)) + + # The third round is still on the board, not yet banked. + board = np.asarray(board_score(final.round)).astype(np.int32) + assert np.array_equal(np.asarray(match_score(final)), np.asarray(final.carry) + board) + + +def test_match_score_does_not_jump_across_a_round_boundary(): + """Phi = carry + board diff must be continuous, or PBRS gets a free kick.""" + state = match_reset(jax.random.PRNGKey(8)) + _, plies = _play_match(state, jax.random.PRNGKey(9)) + + for before, _, after, _ in plies: + rolled_over = int(after.round_idx) != int(before.round_idx) + if not rolled_over: + continue + # The finished round's board is folded into carry and the new board is + # empty, so the running total is unchanged by the roll-over itself. + assert np.array_equal(np.asarray(after.carry), np.asarray(match_score(after))) + assert int(np.asarray(board_score(after.round)).sum()) == 0 + + +# --- the start-player rule ------------------------------------------------- + + +@pytest.mark.parametrize( + ("carry", "coin", "expected"), + [ + ((60, 10), 1, 0), # p0 ahead -> p0 leads, coin ignored + ((10, 60), 0, 1), # p1 ahead -> p1 leads, coin ignored + ((30, 30), 0, 0), # level -> coin + ((30, 30), 1, 1), # level -> coin + ], +) +def test_the_player_with_more_points_begins(carry, coin, expected): + coin_flips = jnp.asarray([coin, coin, coin], dtype=jnp.int8) + starter = starting_player(jnp.asarray(carry, dtype=jnp.int32), jnp.int32(1), coin_flips) + assert int(starter) == expected + + +def test_round_two_is_actually_led_by_whoever_is_ahead(): + state = match_reset(jax.random.PRNGKey(10)) + _, plies = _play_match(state, jax.random.PRNGKey(11)) + + for before, _, after, _ in plies: + if int(after.round_idx) == int(before.round_idx): + continue + carry = np.asarray(after.carry) + if carry[0] == carry[1]: + continue # coin flip, covered above + leader = 0 if carry[0] > carry[1] else 1 + assert int(after.round.to_move) == leader + + +def test_round_one_start_player_is_a_fair_coin(): + starts = [int(match_reset(jax.random.PRNGKey(seed)).round.to_move) for seed in range(400)] + share = sum(starts) / len(starts) + assert 0.4 < share < 0.6, f"round-1 starter is skewed: {share:.2f}" + + +# --- mirrored matches ------------------------------------------------------ + + +def test_mirroring_the_seats_negates_the_match_score(): + """Same deals, same coins, seats swapped: the result must flip sign exactly.""" + decks = _decks(21) + coins = jnp.asarray([0, 1, 0], dtype=jnp.int8) + + # Swapping seats == swapping the two dealt hands and the coin bits. + swapped = decks.at[:, :16].set(jnp.concatenate([decks[:, 8:16], decks[:, :8]], axis=1)) + flipped_coins = (1 - coins).astype(jnp.int8) + + a, _ = _play_match(match_reset_from(decks, coins), jax.random.PRNGKey(30)) + b, _ = _play_match(match_reset_from(swapped, flipped_coins), jax.random.PRNGKey(30)) + + # Both matches see identical information; only the seat labels differ, so the + # per-seat totals must be each other's mirror image. + assert np.array_equal(np.asarray(match_score(a)), np.asarray(match_score(b))[::-1])