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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh

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),

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

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])