Train on matches: match observation, asymmetric critic, both seats, match reward

Phases 2-4 land together because they all rewrite the same rollout, and doing
them in sequence would mean writing it three times.

Observation adds the four things a match policy cannot play without: carry as a
scalar *and* a binned one-hot (round three is a threshold problem -- "win by 41
or lose" plays nothing like "win by 39" -- and the old score_diff divided by
MAX_ABS_SCORE=780, squashing a decisive 50-point lead to 0.06); the round index;
whose turn it is, which the single-round observation never carried even though
the critic is trained on opponent-turn states; and the deck clock, since a round
ends on the last deck draw and players bend that parity by drawing from discard
piles. Live points per colour are split by hand / discard pile / unseen, because
a discard pile is public and recoverable.

The critic is asymmetric: it gets the opponent's hand and the deck in order, on
a separate trunk so none of it can reach the logits. A test pins that down --
perturbing the privileged input leaves the policy logits bit-identical while
moving the value. Deal luck is what makes a match-terminal reward hard to learn
from, and a state-value baseline may condition on anything action-independent.

Both seats now train. Self-play ran one network on both sides and stop_gradiented
the opponent, throwing away half of every game; each ply now emits a transition
per seat, folded into the batch so each seat keeps an independent GAE chain.

Reward is the match: rounds one and two only bank into carry, and round three
pays tanh(total / terminal_scale). Potential shaping on the running total covers
the early sparsity and anneals out.

match_eval adds the two measurements the plan turns on: duplicate match play
(same deals and coins from both seats -- self-play scores exactly 0.500 with zero
mean lead, so the mirroring cancels deal luck exactly) and the carry probe. On an
untrained net the probe is flat: 4.97 expeditions opened at a 60-point deficit
and at a 60-point lead alike. Breaking that flat line is success criterion 1.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh

src/lost_cities_jax/match_eval.py

@@ -0,0 +1,195 @@
+"""Match-level evaluation, and the probe that says whether carry changed anything.
+
+Single-deal win rate cannot tell you if the three-round rework worked. Two things
+can:
+
+- **Duplicate match play.** The same three deals and the same coin flips, played
+  from both seats. Deal luck cancels, so what is left is the policy difference.
+- **The carry probe.** Drop the policy into round three holding a fixed deficit or
+  lead and watch what it does. A carry-blind policy plays a 60-point deficit
+  exactly like a 60-point lead. If the numbers below do not move with carry, the
+  rework bought nothing, no matter what the win rate says.
+"""
+
+from __future__ import annotations
+
+import math
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from lost_cities_jax.engine import reset_from_order
+from lost_cities_jax.match import (
+    N_ROUNDS,
+    MatchState,
+    match_legal_action_mask,
+    match_reset_from,
+    match_score,
+    match_step,
+)
+from lost_cities_jax.match_obs import match_critic_observation, match_observation
+from lost_cities_jax.match_ppo import MatchActorCritic, match_lead
+from lost_cities_jax.ppo import JaxPPOConfig, mask_logits
+from lost_cities_jax.types import DRAW_DECK, MAX_STEPS, N_CARDS
+
+MATCH_SCAN_STEPS = N_ROUNDS * MAX_STEPS
+
+
+def match_bank(seed: int, matches: int) -> tuple[jnp.ndarray, jnp.ndarray]:
+    """Deterministic deals and coin flips, so runs are comparable."""
+    rng = np.random.default_rng(seed)
+    decks = np.stack(
+        [np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)]) for _ in range(matches)]
+    )
+    coins = rng.integers(0, 2, size=(matches, N_ROUNDS))
+    return jnp.asarray(decks, dtype=jnp.int8), jnp.asarray(coins, dtype=jnp.int8)
+
+
+def _wilson(wins: float, games: float) -> tuple[float, float]:
+    if games == 0:
+        return 0.0, 0.0
+    z = 1.96
+    p = wins / games
+    denom = 1 + z * z / games
+    centre = p + z * z / (2 * games)
+    margin = z * math.sqrt(p * (1 - p) / games + z * z / (4 * games * games))
+    return (centre - margin) / denom, (centre + margin) / denom
+
+
+def make_match_runner(cfg: JaxPPOConfig, opponent_policy=None):
+    """Play whole matches to the end. ``opponent_policy`` is a single-round policy."""
+
+    model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
+
+    @jax.jit
+    def run(params, env: MatchState, learner_seat: jax.Array, rng: jax.Array):
+        def body(carry, _):
+            env, key, deck_draws, plies = carry
+            key, opp_key = jax.random.split(key)
+
+            to_move = env.round.to_move.astype(jnp.int32)
+            obs = jax.vmap(match_observation)(env, to_move)
+            critic = jax.vmap(match_critic_observation)(env, to_move)
+            mask = jax.vmap(match_legal_action_mask)(env)
+            logits, _ = model.apply(params, obs, critic)
+            learner_action = jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)
+
+            if opponent_policy is None:
+                action = learner_action
+            else:
+                opp_keys = jax.random.split(opp_key, env.done.shape[0])
+                opp_action = jax.vmap(opponent_policy, in_axes=(0, 0, 0))(
+                    env.round, to_move, opp_keys
+                )
+                action = jnp.where(to_move == learner_seat, learner_action, opp_action)
+
+            live = ~env.done
+            is_deck = (action % 6) == DRAW_DECK
+            deck_draws = deck_draws + (is_deck & live).astype(jnp.int32)
+            plies = plies + live.astype(jnp.int32)
+
+            env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
+            return (env, key, deck_draws, plies), None
+
+        n = env.done.shape[0]
+        zeros = jnp.zeros((n,), dtype=jnp.int32)
+        (env, _, deck_draws, plies), _ = jax.lax.scan(
+            body, (env, rng, zeros, zeros), xs=None, length=MATCH_SCAN_STEPS
+        )
+        return env, deck_draws, plies
+
+    return run
+
+
+def match_evaluate(cfg: JaxPPOConfig, params, *, matches: int = 2000, seed: int = 20260715) -> dict:
+    """Duplicate match play: every deal seen from both seats."""
+
+    decks, coins = match_bank(seed, matches)
+    runner = make_match_runner(cfg)
+
+    results = []
+    for seat in (0, 1):
+        env = jax.vmap(match_reset_from)(decks, coins)
+        seats = jnp.full((matches,), seat, dtype=jnp.int32)
+        final, deck_draws, plies = runner(params, env, seats, jax.random.PRNGKey(0))
+        totals = np.asarray(jax.vmap(match_score)(final))
+        lead = totals[:, seat] - totals[:, 1 - seat]
+        results.append(
+            {
+                "lead": lead,
+                "plies": np.asarray(plies),
+                "deck_share": np.asarray(deck_draws) / np.maximum(np.asarray(plies), 1),
+                "done": np.asarray(final.done),
+            }
+        )
+
+    lead = np.concatenate([r["lead"] for r in results])
+    games = float(lead.size)
+    wins = float((lead > 0).sum())
+    ties = float((lead == 0).sum())
+    low, high = _wilson(wins, games)
+    return {
+        "matches": games,
+        "wins": wins,
+        "ties": ties,
+        "losses": games - wins - ties,
+        "match_win_rate": wins / games,
+        "wilson_low": low,
+        "wilson_high": high,
+        "mean_total_lead": float(lead.mean()),
+        "mean_plies": float(np.concatenate([r["plies"] for r in results]).mean()),
+        "unfinished_rate": float(
+            1.0 - np.concatenate([r["done"] for r in results]).astype(np.float32).mean()
+        ),
+    }
+
+
+def carry_probe(
+    cfg: JaxPPOConfig,
+    params,
+    *,
+    carry_levels=(-60, -25, -1, 1, 25, 60),
+    matches: int = 512,
+    seed: int = 20260716,
+) -> list[dict]:
+    """Start the policy in round three holding ``carry`` and see what changes.
+
+    A policy that reads carry should open more expeditions and lean on wagers when
+    it is behind (it has to swing), and shut down when it is ahead. Flat rows mean
+    the policy is ignoring the one thing the three-round rework added.
+    """
+
+    decks, coins = match_bank(seed, matches)
+    runner = make_match_runner(cfg)
+    rows = []
+
+    for level in carry_levels:
+        base = jax.vmap(match_reset_from)(decks, coins)
+        carry = jnp.tile(jnp.asarray([level, 0], dtype=jnp.int32), (matches, 1))
+        round_idx = jnp.full((matches,), N_ROUNDS - 1, dtype=jnp.int32)
+        # Whoever is ahead leads the round, exactly as the match layer would.
+        first = jnp.int8(0 if level > 0 else 1)
+        rounds = jax.vmap(reset_from_order, in_axes=(0, None))(decks[:, N_ROUNDS - 1], first)
+
+        env = base._replace(round=rounds, carry=carry, round_idx=round_idx)
+        final, deck_draws, plies = runner(
+            params, env, jnp.zeros((matches,), jnp.int32), jax.random.PRNGKey(1)
+        )
+
+        lead = np.asarray(jax.vmap(match_lead)(final))
+        opened = np.asarray(jnp.sum(final.round.col_len[:, 0, :] > 0, axis=-1))
+        wagers = np.asarray(jnp.sum(final.round.col_hs[:, 0, :], axis=-1))
+        deck_share = np.asarray(deck_draws) / np.maximum(np.asarray(plies), 1)
+
+        rows.append(
+            {
+                "carry": level,
+                "win_rate": float((lead > 0).mean()),
+                "opened_colors": float(opened.mean()),
+                "wagers_played": float(wagers.mean()),
+                "deck_race_rate": float(deck_share.mean()),
+                "mean_plies": float(np.asarray(plies).mean()),
+            }
+        )
+    return rows

src/lost_cities_jax/match_obs.py

@@ -0,0 +1,198 @@
+"""Player-view observation for a three-round match.
+
+Everything the single-round observation had, plus the four things a match policy
+cannot play without:
+
+- **carry**, as a scalar *and* a one-hot over bins. Round three is a threshold
+  problem -- "win by 41 or lose" plays nothing like "win by 39" -- and a lone
+  scalar makes the network learn a cliff from a smooth input. (The old
+  ``score_diff`` divided by ``MAX_ABS_SCORE`` = 780, which squashed a decisive
+  50-point lead to 0.06. That scaling is fixed here too.)
+- **which round it is**, and how many are left.
+- **whose turn it is.** The single-round obs never had this, yet the critic is
+  trained on opponent-turn states as well, leaving it to recover turn parity
+  from ``step_count / 400``.
+- **the deck clock.** A round ends when the last deck card is drawn, so turn
+  parity decides who gets the final placement -- and players bend that parity by
+  drawing from discard piles instead.
+
+Live points per colour are split by where the card is (hand / discard pile /
+unseen) because a discard pile is public and recoverable: lumping it in with the
+unseen cards hides the high cards sitting in plain view.
+"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+
+from lost_cities_jax.engine import board_score
+from lost_cities_jax.match import N_ROUNDS, MatchState
+from lost_cities_jax.obs import observation
+from lost_cities_jax.types import (
+    CARDS_PER_COLOR,
+    LOC_DECK,
+    LOC_DISCARD,
+    LOC_P0_HAND,
+    N_CARDS,
+    N_COLORS,
+    OBS_DIM,
+)
+
+# Bin edges for the carry one-hot. Packed tightly around zero: that is where the
+# round-three decision actually flips.
+CARRY_BIN_EDGES = jnp.asarray([-60.0, -30.0, -12.0, -1.0, 1.0, 12.0, 30.0, 60.0])
+N_CARRY_BINS = CARRY_BIN_EDGES.shape[0] + 1
+
+CARRY_SCALE = 75.0
+"""Typical round margin, not the theoretical 780 maximum."""
+
+MAX_COLOR_POINTS = 54.0
+"""2+3+...+10, the most a single expedition can be worth before multipliers."""
+
+N_MATCH_SCALARS = (
+    1  # carry, scaled
+    + N_CARRY_BINS  # carry, binned
+    + N_ROUNDS  # which round
+    + 1  # rounds left
+    + 1  # is it my turn
+    + 1  # did I lead this round
+    + 1  # do I take the last deck card at the current parity
+    + 2 * N_COLORS * 3  # live points: both players x colour x {hand, discard, unseen}
+)
+MATCH_OBS_DIM = OBS_DIM + N_MATCH_SCALARS
+
+
+def _live_points(state, viewer: jax.Array, subject: jax.Array) -> jax.Array:
+    """Points still reachable for ``subject``, split by where the card sits.
+
+    Returns ``float32[N_COLORS, 3]`` over {in hand, in a discard pile, unseen}.
+    Seen through ``viewer``'s eyes: a card in the opponent's hand only counts as
+    "in hand" if it is public, otherwise it is unseen.
+    """
+
+    ids = jnp.arange(N_CARDS, dtype=jnp.int32)
+    colors = ids // CARDS_PER_COLOR
+    slots = ids % CARDS_PER_COLOR
+    ranks = jnp.where(slots >= 3, slots - 1, 0).astype(jnp.float32)
+
+    loc = state.card_loc.astype(jnp.int32)
+    subject_hand = loc == (LOC_P0_HAND + subject)
+    is_mine = subject == viewer
+
+    # An ascending column can only take cards above its current top.
+    above_top = ranks > state.col_top[subject][colors].astype(jnp.float32)
+
+    in_hand = subject_hand & (is_mine | state.hand_public)
+    in_discard = loc == LOC_DISCARD
+    unseen = (loc == LOC_DECK) | (subject_hand & ~is_mine & ~state.hand_public)
+    # Cards in the *other* player's hidden hand are unseen to the viewer too, but
+    # they are not reachable by the subject, so they are deliberately excluded.
+
+    def by_color(mask: jax.Array) -> jax.Array:
+        weighted = jnp.where(mask & above_top, ranks, 0.0)
+        return jax.ops.segment_sum(weighted, colors, num_segments=N_COLORS)
+
+    stacked = jnp.stack([by_color(in_hand), by_color(in_discard), by_color(unseen)], axis=1)
+    return stacked / MAX_COLOR_POINTS
+
+
+def match_observation(state: MatchState, player: jax.Array) -> jax.Array:
+    """Return ``float32[MATCH_OBS_DIM]`` from ``player``'s seat."""
+
+    player = jnp.asarray(player, dtype=jnp.int32)
+    opponent = 1 - player
+    round_state = state.round
+
+    base = observation(round_state, player)
+
+    totals = state.carry + board_score(round_state).astype(jnp.int32)
+    lead = (totals[player] - totals[opponent]).astype(jnp.float32)
+    carry_scaled = jnp.clip(lead / CARRY_SCALE, -2.0, 2.0).reshape((1,))
+    carry_bins = jax.nn.one_hot(jnp.digitize(lead, CARRY_BIN_EDGES), N_CARRY_BINS)
+
+    round_onehot = jax.nn.one_hot(state.round_idx, N_ROUNDS)
+    rounds_left = ((N_ROUNDS - 1 - state.round_idx).astype(jnp.float32) / (N_ROUNDS - 1)).reshape(
+        (1,)
+    )
+
+    to_move = round_state.to_move.astype(jnp.int32)
+    my_turn = (to_move == player).astype(jnp.float32).reshape((1,))
+    # to_move flips every ply, so the round's opener is recoverable from parity.
+    parity = round_state.step_count.astype(jnp.int32) & 1
+    round_opener = to_move ^ parity
+    i_opened = (round_opener == player).astype(jnp.float32).reshape((1,))
+
+    # If both players drew from the deck from here, the last deck card falls to
+    # whoever is on move after `remaining - 1` more plies.
+    remaining = jnp.maximum(N_CARDS - round_state.draw_ptr, 1).astype(jnp.int32)
+    last_drawer = to_move ^ ((remaining - 1) & 1)
+    i_take_last = (last_drawer == player).astype(jnp.float32).reshape((1,))
+
+    live_me = _live_points(round_state, player, player).reshape((-1,))
+    live_opp = _live_points(round_state, player, opponent).reshape((-1,))
+
+    extra = jnp.concatenate(
+        [
+            carry_scaled,
+            carry_bins,
+            round_onehot,
+            rounds_left,
+            my_turn,
+            i_opened,
+            i_take_last,
+            live_me,
+            live_opp,
+        ],
+        axis=0,
+    ).astype(jnp.float32)
+
+    return jnp.concatenate([base, extra], axis=0).reshape((MATCH_OBS_DIM,))
+
+
+CRITIC_EXTRA_CHANNELS = 3
+MATCH_CRITIC_OBS_DIM = MATCH_OBS_DIM + N_CARDS * CRITIC_EXTRA_CHANNELS
+
+
+def match_critic_observation(state: MatchState, player: jax.Array) -> jax.Array:
+    """The same view, plus the hidden state. Training only -- never played from.
+
+    A value baseline may condition on anything the action does not depend on
+    without biasing the policy gradient, and the deal is exactly the noise that
+    swamps a match-terminal reward. So the critic gets the opponent's hand and
+    the deck *in order* -- it knows what is coming -- while the actor keeps the
+    masked view. The two run on separate trunks so none of this can leak into
+    the logits.
+    """
+
+    player = jnp.asarray(player, dtype=jnp.int32)
+    opponent = 1 - player
+    round_state = state.round
+    loc = round_state.card_loc.astype(jnp.int32)
+
+    opp_hand = (loc == (LOC_P0_HAND + opponent)).astype(jnp.float32)
+    in_deck = (loc == LOC_DECK).astype(jnp.float32)
+
+    # Position of each card in the undrawn deck, so the critic can see the order
+    # rather than just the multiset.
+    deck_order = round_state.deck_order.astype(jnp.int32)
+    position = (
+        jnp.zeros((N_CARDS,), dtype=jnp.int32)
+        .at[deck_order]
+        .set(jnp.arange(N_CARDS, dtype=jnp.int32))
+    )
+    depth = (position - round_state.draw_ptr).astype(jnp.float32) / N_CARDS
+    deck_depth = jnp.where(in_deck > 0, depth, 0.0)
+
+    hidden = jnp.stack([opp_hand, in_deck, deck_depth], axis=1).reshape((-1,))
+    base = match_observation(state, player)
+    return jnp.concatenate([base, hidden], axis=0).reshape((MATCH_CRITIC_OBS_DIM,))
+
+
+__all__ = [
+    "MATCH_CRITIC_OBS_DIM",
+    "MATCH_OBS_DIM",
+    "N_CARRY_BINS",
+    "match_critic_observation",
+    "match_observation",
+]

src/lost_cities_jax/match_ppo.py

@@ -0,0 +1,407 @@
+"""PPO over three-round matches.
+
+Three things differ from the single-round trainer beyond the env swap:
+
+**Asymmetric actor-critic.** The critic never plays, so it is fed the hidden
+state (opponent's hand, deck in order) while the actor keeps the masked view. A
+state-value baseline may condition on anything action-independent without biasing
+the policy gradient, and deal luck is exactly what makes a match-terminal reward
+hard to learn from. The trunks are separate so the privileged features cannot
+reach the logits.
+
+**Both seats train.** In self-play one network picks both players' moves, but the
+single-round trainer stop_gradients the opponent seat and throws those plies away
+-- half of every game. Here every ply emits a transition for both seats, each
+with its own seat-relative view, value and reward; the actor mask keeps the
+policy loss on the seat that actually moved, while the critic learns from both.
+Folding the seat axis into the batch keeps each seat's GAE chain independent.
+
+**Reward is the match, not the round.** Rounds one and two pay nothing, they only
+bank into ``carry``. The terminal reward is ``tanh(total_diff / terminal_scale)``
+at the end of round three. Driving that scale toward zero would make it
+``sign()`` -- the true objective, but a poor signal, since every ply of a
+~160-ply match would then carry the same +/-1 and all credit assignment would
+fall to the critic. Potential shaping on the running total covers the gap early
+and anneals away.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import asdict
+from pathlib import Path
+from typing import NamedTuple
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import optax
+
+from lost_cities_jax.match import (
+    MatchState,
+    match_legal_action_mask,
+    match_reset,
+    match_score,
+    match_step,
+)
+from lost_cities_jax.match_obs import (
+    MATCH_CRITIC_OBS_DIM,
+    MATCH_OBS_DIM,
+    match_critic_observation,
+    match_observation,
+)
+from lost_cities_jax.ppo import (
+    EpisodeEnd,
+    JaxPPOConfig,
+    TrainState,
+    _append_jsonl,
+    _broadcast_done,
+    _create_run_dir,
+    _write_json,
+    action_log_prob,
+    categorical_entropy,
+    compute_gae,
+    mask_logits,
+    masked_mean,
+    restore_checkpoint,
+    save_checkpoint,
+    shaping_coefficient,
+)
+from lost_cities_jax.types import MAX_STEPS, N_ACTIONS, PLAY
+
+N_SEATS = 2
+
+
+class MatchTransition(NamedTuple):
+    obs: jax.Array
+    critic_obs: jax.Array
+    legal_mask: jax.Array
+    action: jax.Array
+    log_prob: jax.Array
+    value: jax.Array
+    reward: jax.Array
+    done: jax.Array
+    active: jax.Array
+    actor_mask: jax.Array
+    entropy: jax.Array
+    play_action: jax.Array
+
+
+class MatchActorCritic(nn.Module):
+    """Separate trunks, so the critic's privileged input cannot reach the logits."""
+
+    hidden_size: int = 512
+    num_layers: int = 3
+
+    @nn.compact
+    def __call__(self, obs: jax.Array, critic_obs: jax.Array) -> tuple[jax.Array, jax.Array]:
+        x = obs
+        for _ in range(self.num_layers):
+            x = nn.relu(nn.Dense(self.hidden_size)(x))
+        logits = nn.Dense(N_ACTIONS)(x)
+
+        v = critic_obs
+        for _ in range(self.num_layers):
+            v = nn.relu(nn.Dense(self.hidden_size)(v))
+        value = nn.Dense(1)(v)
+        return logits, jnp.squeeze(value, axis=-1)
+
+
+def create_match_train_state(cfg: JaxPPOConfig, rng: jax.Array) -> TrainState:
+    model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
+    params = model.init(
+        rng,
+        jnp.zeros((1, MATCH_OBS_DIM), dtype=jnp.float32),
+        jnp.zeros((1, MATCH_CRITIC_OBS_DIM), dtype=jnp.float32),
+    )
+    tx = optax.chain(
+        optax.clip_by_global_norm(cfg.ppo.max_grad_norm),
+        optax.adam(cfg.ppo.learning_rate),
+    )
+    return TrainState.create(apply_fn=model.apply, params=params, tx=tx)
+
+
+def match_lead(env: MatchState) -> jax.Array:
+    """Running match total from seat 0's side, in raw points."""
+    totals = match_score(env)
+    return (totals[0] - totals[1]).astype(jnp.float32)
+
+
+def reset_done_matches(env: MatchState, rng: jax.Array, batch: int) -> MatchState:
+    fresh = jax.vmap(match_reset)(jax.random.split(rng, batch))
+    return jax.tree_util.tree_map(
+        lambda old, new: jnp.where(_broadcast_done(env.done, old), new, old), env, fresh
+    )
+
+
+def _seat_views(env: MatchState, seats: jax.Array):
+    obs = jax.vmap(lambda s: jax.vmap(match_observation, in_axes=(0, None))(env, s))(seats)
+    critic = jax.vmap(lambda s: jax.vmap(match_critic_observation, in_axes=(0, None))(env, s))(
+        seats
+    )
+    return obs, critic
+
+
+def _episode_end(env: MatchState, terminal: jax.Array, episode_return: jax.Array) -> EpisodeEnd:
+    lead = jax.vmap(match_lead)(env)
+    opened = jnp.sum(env.round.col_len[:, 0, :] > 0, axis=-1)
+    return EpisodeEnd(
+        done=terminal,
+        episode_return=episode_return,
+        length=env.round.step_count,
+        opened_colors=opened,
+        positive_expeditions=lead.astype(jnp.int32),  # final match lead, in points
+        hit_max_steps=env.round.step_count >= MAX_STEPS,
+    )
+
+
+def match_episode_metrics(transitions: MatchTransition, episodes: EpisodeEnd) -> dict:
+    done = episodes.done.astype(jnp.float32)
+    completed = jnp.sum(done)
+    denom = jnp.maximum(completed, 1.0)
+    actor_count = jnp.sum(transitions.actor_mask)
+    lead = episodes.positive_expeditions.astype(jnp.float32)
+    return {
+        "return_mean": jnp.sum(episodes.episode_return * done) / denom,
+        "match_lead_mean": jnp.sum(lead * done) / denom,
+        "match_win_rate": jnp.sum((lead > 0).astype(jnp.float32) * done) / denom,
+        "final_round_length_mean": jnp.sum(episodes.length.astype(jnp.float32) * done) / denom,
+        "max_steps_rate": jnp.sum(episodes.hit_max_steps.astype(jnp.float32) * done) / denom,
+        "play_action_rate": jnp.sum(transitions.play_action) / jnp.maximum(actor_count, 1),
+        "entropy_mean": masked_mean(
+            transitions.entropy, transitions.actor_mask.astype(jnp.float32)
+        ),
+        "learner_actions": actor_count,
+        "matches_completed": completed,
+    }
+
+
+def make_match_rollout_fn(cfg: JaxPPOConfig):
+    batch = cfg.ppo.batch_games
+    seats = jnp.arange(N_SEATS, dtype=jnp.int32)
+
+    @jax.jit
+    def rollout_fn(state: TrainState, env: MatchState, rng: jax.Array, shaping_coef: jax.Array):
+        def body(carry, _):
+            env, episode_return, key = carry
+            key, act_key, reset_key = jax.random.split(key, 3)
+
+            obs, critic_obs = _seat_views(env, seats)  # (N_SEATS, batch, ...)
+            legal = jax.vmap(match_legal_action_mask)(env)
+            legal_both = jnp.broadcast_to(legal, (N_SEATS, batch, N_ACTIONS))
+
+            logits, value = state.apply_fn(state.params, obs, critic_obs)
+            masked = mask_logits(logits, legal_both)
+            sampled = jax.random.categorical(act_key, masked, axis=-1).astype(jnp.int32)
+            log_prob = jax.vmap(action_log_prob)(masked, sampled)
+            entropy = jax.vmap(categorical_entropy)(masked)
+
+            to_move = env.round.to_move.astype(jnp.int32)
+            active = ~env.done
+            actor_mask = (seats[:, None] == to_move[None, :]) & active[None, :]
+            actions = jnp.take_along_axis(sampled, to_move[None, :], axis=0)[0]
+
+            before = jax.vmap(match_lead)(env)
+            next_env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, actions)
+            after = jax.vmap(match_lead)(next_env)
+
+            done = next_env.done
+            terminal = active & done
+            reward0 = jnp.where(terminal, jnp.tanh(after / cfg.reward.terminal_scale), 0.0)
+            reward0 = reward0 + shaping_coef * (after - before) / cfg.reward.terminal_scale
+            reward0 = jnp.where(active, reward0, 0.0)
+            # Zero-sum: seat 1 sees exactly the negation.
+            reward = jnp.stack([reward0, -reward0], axis=0)
+
+            played = ((actions % 12) // 6 == PLAY)[None, :] & actor_mask
+            transition = MatchTransition(
+                obs=obs,
+                critic_obs=critic_obs,
+                legal_mask=legal_both,
+                action=jnp.broadcast_to(actions, (N_SEATS, batch)),
+                log_prob=log_prob,
+                value=value,
+                reward=reward,
+                done=jnp.broadcast_to(done, (N_SEATS, batch)),
+                active=jnp.broadcast_to(active, (N_SEATS, batch)),
+                actor_mask=actor_mask,
+                entropy=entropy,
+                play_action=played,
+            )
+
+            episode_return = episode_return + reward0
+            episode = _episode_end(next_env, terminal, episode_return)
+            next_env = reset_done_matches(next_env, reset_key, batch)
+            episode_return = jnp.where(done, 0.0, episode_return)
+            return (next_env, episode_return, key), (transition, episode)
+
+        init_return = jnp.zeros((batch,), dtype=jnp.float32)
+        (next_env, _, rng), (transitions, episodes) = jax.lax.scan(
+            body, (env, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
+        )
+
+        # Fold the seat axis into the batch: each column is then an independent
+        # episode chain, which is exactly what GAE wants.
+        transitions = jax.tree_util.tree_map(
+            lambda x: x.reshape((x.shape[0], N_SEATS * batch, *x.shape[3:])), transitions
+        )
+        final_obs, final_critic = _seat_views(next_env, seats)
+        _, last_value = state.apply_fn(state.params, final_obs, final_critic)
+        last_value = last_value.reshape((N_SEATS * batch,))
+
+        return next_env, transitions, last_value, match_episode_metrics(transitions, episodes)
+
+    return rollout_fn
+
+
+def match_ppo_update(
+    state: TrainState,
+    transitions: MatchTransition,
+    advantages: jax.Array,
+    returns: jax.Array,
+    rng: jax.Array,
+    cfg: JaxPPOConfig,
+) -> tuple[TrainState, dict[str, jax.Array]]:
+    batch_size = int(transitions.reward.size)
+    minibatch_size = batch_size // cfg.ppo.minibatches
+    flat = jax.tree_util.tree_map(lambda x: x.reshape((-1, *x.shape[2:])), transitions)
+    flat_advantages = advantages.reshape((batch_size,))
+    flat_returns = returns.reshape((batch_size,))
+
+    actor_mask = flat.actor_mask.astype(jnp.float32)
+    adv_mean = masked_mean(flat_advantages, actor_mask)
+    adv_std = jnp.sqrt(masked_mean((flat_advantages - adv_mean) ** 2, actor_mask) + 1.0e-8)
+    flat_advantages = (flat_advantages - adv_mean) / adv_std
+
+    def loss_fn(params, mb, mb_adv, mb_returns):
+        logits, value = state.apply_fn(params, mb.obs, mb.critic_obs)
+        masked_logits = mask_logits(logits, mb.legal_mask)
+        log_prob = action_log_prob(masked_logits, mb.action)
+        entropy = categorical_entropy(masked_logits)
+        ratio = jnp.exp(log_prob - mb.log_prob)
+        actor_weight = mb.actor_mask.astype(jnp.float32)
+        active_weight = mb.active.astype(jnp.float32)
+        unclipped = ratio * mb_adv
+        clipped = jnp.clip(ratio, 1.0 - cfg.ppo.clip_epsilon, 1.0 + cfg.ppo.clip_epsilon) * mb_adv
+        actor_loss = -masked_mean(jnp.minimum(unclipped, clipped), actor_weight)
+        value_loss = masked_mean((mb_returns - value) ** 2, active_weight)
+        entropy_loss = masked_mean(entropy, actor_weight)
+        loss = actor_loss + cfg.ppo.value_coef * value_loss - cfg.ppo.entropy_coef * entropy_loss
+        return loss, {
+            "loss": loss,
+            "actor_loss": actor_loss,
+            "value_loss": value_loss,
+            "entropy_loss": entropy_loss,
+            "approx_kl": masked_mean(mb.log_prob - log_prob, actor_weight),
+        }
+
+    def epoch_update(carry, key):
+        train_state = carry
+        permutation = jax.random.permutation(key, batch_size)
+
+        def minibatch_update(carry, idx):
+            train_state, accum = carry
+            mb_idx = jax.lax.dynamic_slice(permutation, (idx * minibatch_size,), (minibatch_size,))
+            mb = jax.tree_util.tree_map(lambda x: x[mb_idx], flat)
+            (_, metrics), grads = jax.value_and_grad(loss_fn, has_aux=True)(
+                train_state.params, mb, flat_advantages[mb_idx], flat_returns[mb_idx]
+            )
+            train_state = train_state.apply_gradients(grads=grads)
+            accum = jax.tree_util.tree_map(lambda a, b: a + b, accum, metrics)
+            return (train_state, accum), None
+
+        zeros = {
+            key: jnp.array(0.0)
+            for key in ("loss", "actor_loss", "value_loss", "entropy_loss", "approx_kl")
+        }
+        (train_state, metrics), _ = jax.lax.scan(
+            minibatch_update, (train_state, zeros), jnp.arange(cfg.ppo.minibatches)
+        )
+        metrics = jax.tree_util.tree_map(lambda x: x / cfg.ppo.minibatches, metrics)
+        return train_state, metrics
+
+    keys = jax.random.split(rng, cfg.ppo.epochs)
+    state, metrics = jax.lax.scan(epoch_update, state, keys)
+    return state, jax.tree_util.tree_map(lambda x: jnp.mean(x, axis=0), metrics)
+
+
+def make_match_train_iteration(cfg: JaxPPOConfig):
+    rollout_fn = make_match_rollout_fn(cfg)
+
+    @jax.jit
+    def train_iteration(
+        state: TrainState, env: MatchState, rng: jax.Array, shaping_coef: jax.Array
+    ):
+        rng, rollout_key, update_key = jax.random.split(rng, 3)
+        env, transitions, last_value, stats = rollout_fn(state, env, rollout_key, shaping_coef)
+        advantages, returns = compute_gae(
+            transitions.reward,
+            transitions.value,
+            transitions.done,
+            cfg.ppo.gamma,
+            cfg.ppo.gae_lambda,
+            last_value,
+        )
+        state, update_metrics = match_ppo_update(
+            state, transitions, advantages, returns, update_key, cfg
+        )
+        return state, env, rng, {**stats, **update_metrics}
+
+    return train_iteration
+
+
+def match_train(cfg: JaxPPOConfig, *, resume: str | None = None) -> Path:
+    run_dir = _create_run_dir(cfg)
+    _write_json(run_dir / "config.json", asdict(cfg))
+    metrics_path = run_dir / "metrics.jsonl"
+
+    rng = jax.random.PRNGKey(cfg.run.seed)
+    rng, init_key, reset_key = jax.random.split(rng, 3)
+    state = create_match_train_state(cfg, init_key)
+    if resume:
+        state = restore_checkpoint(Path(resume), state)
+
+    env = jax.jit(jax.vmap(match_reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
+    train_iteration = make_match_train_iteration(cfg)
+
+    start = time.perf_counter()
+    learner_actions_seen = 0
+    for update in range(cfg.run.total_updates):
+        shaping_coef = jnp.asarray(shaping_coefficient(cfg, learner_actions_seen), jnp.float32)
+        iter_start = time.perf_counter()
+        state, env, rng, metrics = train_iteration(state, env, rng, shaping_coef)
+        jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
+
+        row = {
+            key: float(value) if getattr(value, "shape", ()) == () else value.tolist()
+            for key, value in metrics.items()
+        }
+        learner_actions_seen += int(row["learner_actions"])
+        row.update(
+            update=update,
+            shaping_coef=float(shaping_coef),
+            learner_actions_total=learner_actions_seen,
+            iteration_seconds=time.perf_counter() - iter_start,
+            elapsed_seconds=time.perf_counter() - start,
+        )
+        _append_jsonl(metrics_path, row)
+        if update % cfg.run.log_every == 0:
+            print(
+                json.dumps(
+                    {
+                        "update": update,
+                        "match_win_rate": row["match_win_rate"],
+                        "match_lead_mean": row["match_lead_mean"],
+                        "matches_completed": row["matches_completed"],
+                        "entropy_mean": row["entropy_mean"],
+                    },
+                    sort_keys=True,
+                ),
+                flush=True,
+            )
+        if cfg.run.checkpoint_every and (update + 1) % cfg.run.checkpoint_every == 0:
+            save_checkpoint(run_dir / "latest", state, cfg)
+
+    save_checkpoint(run_dir / "latest", state, cfg)
+    return run_dir

tests/lost_cities_jax/test_match_ppo.py

@@ -0,0 +1,194 @@
+"""Phases 2-4: match reward, match observation, asymmetric critic, both-seat training."""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from lost_cities_jax.match import match_reset, match_score, match_step
+from lost_cities_jax.match_obs import (
+    MATCH_CRITIC_OBS_DIM,
+    MATCH_OBS_DIM,
+    match_critic_observation,
+    match_observation,
+)
+from lost_cities_jax.match_ppo import (
+    create_match_train_state,
+    make_match_rollout_fn,
+    make_match_train_iteration,
+)
+from lost_cities_jax.opponents import random_legal_action
+from lost_cities_jax.ppo import JaxPPOConfig
+from lost_cities_jax.types import LOC_P0_HAND, N_CARDS
+
+
+def _cfg(rollout_steps: int = 700, batch_games: int = 8) -> JaxPPOConfig:
+    cfg = JaxPPOConfig()
+    cfg.ppo.batch_games = batch_games
+    cfg.ppo.rollout_steps = rollout_steps
+    cfg.ppo.minibatches = 4
+    cfg.ppo.epochs = 1
+    cfg.network.hidden_size = 32
+    cfg.network.num_layers = 1
+    return cfg
+
+
+# --- the asymmetric critic -------------------------------------------------
+
+
+def test_the_critic_sees_the_opponents_hand_and_the_actor_does_not():
+    match = match_reset(jax.random.PRNGKey(0))
+    loc = np.asarray(match.round.card_loc)
+    opp_hand = np.flatnonzero(loc == LOC_P0_HAND + 1)
+
+    actor = np.asarray(match_observation(match, jnp.int32(0)))
+    critic = np.asarray(match_critic_observation(match, jnp.int32(0)))
+
+    # Card-major, three channels each, matching the base observation's layout.
+    hidden = critic[MATCH_OBS_DIM:].reshape(N_CARDS, 3)
+    assert np.array_equal(np.flatnonzero(hidden[:, 0]), opp_hand)
+
+    # The actor's view is a strict prefix of it and carries none of that.
+    assert np.array_equal(actor, critic[:MATCH_OBS_DIM])
+    assert critic.shape == (MATCH_CRITIC_OBS_DIM,)
+
+
+def test_privileged_input_cannot_move_the_policy_logits():
+    """Separate trunks, or the critic's view of the deck leaks into play."""
+    cfg = _cfg()
+    state = create_match_train_state(cfg, jax.random.PRNGKey(1))
+    match = match_reset(jax.random.PRNGKey(2))
+
+    obs = match_observation(match, jnp.int32(0))[None, :]
+    critic_a = match_critic_observation(match, jnp.int32(0))[None, :]
+    critic_b = critic_a.at[0, MATCH_OBS_DIM:].set(1.0)  # a totally different hidden state
+
+    logits_a, value_a = state.apply_fn(state.params, obs, critic_a)
+    logits_b, value_b = state.apply_fn(state.params, obs, critic_b)
+
+    assert np.array_equal(np.asarray(logits_a), np.asarray(logits_b))
+    # ...and the value head must actually be using it, or the critic is pointless.
+    assert not np.allclose(np.asarray(value_a), np.asarray(value_b))
+
+
+# --- the match reward ------------------------------------------------------
+
+
+def test_reward_is_zero_sum_and_only_paid_at_the_end_of_the_match():
+    cfg = _cfg()
+    state = create_match_train_state(cfg, jax.random.PRNGKey(3))
+    env = jax.jit(jax.vmap(match_reset))(
+        jax.random.split(jax.random.PRNGKey(4), cfg.ppo.batch_games)
+    )
+
+    rollout = make_match_rollout_fn(cfg)
+    # Shaping off, so any non-zero reward must be a terminal one.
+    _, transitions, _, metrics = rollout(state, env, jax.random.PRNGKey(5), jnp.asarray(0.0))
+
+    reward = np.asarray(transitions.reward)  # (T, 2 * batch)
+    seat0, seat1 = reward[:, : cfg.ppo.batch_games], reward[:, cfg.ppo.batch_games :]
+    assert np.allclose(seat0, -seat1)
+
+    done = np.asarray(transitions.done)[:, : cfg.ppo.batch_games]
+    paid = seat0 != 0.0
+    assert np.array_equal(paid, done & paid)  # never paid on a non-terminal ply
+    assert float(metrics["matches_completed"]) > 0
+    assert np.abs(seat0[paid]).max() <= 1.0  # tanh-bounded
+
+
+def test_shaping_tracks_the_running_match_total():
+    """Phi is carry + board diff, so shaping must follow the total, not the round."""
+    match = match_reset(jax.random.PRNGKey(6))
+    key = jax.random.PRNGKey(7)
+
+    prev = np.asarray(match_score(match))
+    for _ in range(200):
+        if bool(match.done):
+            break
+        key, step_key = jax.random.split(key)
+        action = random_legal_action(match.round, match.round.to_move, step_key)
+        match, _, _ = match_step(match, action)
+        total = np.asarray(match_score(match))
+        # The total never resets when a round rolls over; it only accumulates.
+        assert total.shape == (2,)
+        prev = total
+    assert prev.shape == (2,)
+
+
+# --- both seats ------------------------------------------------------------
+
+
+def test_exactly_one_seat_acts_per_ply_but_both_are_trained():
+    cfg = _cfg()
+    state = create_match_train_state(cfg, jax.random.PRNGKey(8))
+    env = jax.jit(jax.vmap(match_reset))(
+        jax.random.split(jax.random.PRNGKey(9), cfg.ppo.batch_games)
+    )
+
+    rollout = make_match_rollout_fn(cfg)
+    _, transitions, last_value, _ = rollout(state, env, jax.random.PRNGKey(10), jnp.asarray(0.0))
+
+    batch = cfg.ppo.batch_games
+    actor = np.asarray(transitions.actor_mask)
+    seat0, seat1 = actor[:, :batch], actor[:, batch:]
+
+    # One mover per ply while the match is live.
+    live = np.asarray(transitions.active)[:, :batch]
+    assert np.array_equal(seat0 ^ seat1, live)
+
+    # But the critic gets both seats: every live ply is a value target on both.
+    assert np.asarray(transitions.active).sum() == 2 * live.sum()
+    assert transitions.value.shape == (cfg.ppo.rollout_steps, 2 * batch)
+    assert last_value.shape == (2 * batch,)
+
+
+def test_self_play_is_balanced_and_a_train_step_stays_finite():
+    cfg = _cfg(rollout_steps=700, batch_games=16)
+    state = create_match_train_state(cfg, jax.random.PRNGKey(11))
+    env = jax.jit(jax.vmap(match_reset))(jax.random.split(jax.random.PRNGKey(12), 16))
+    train_iteration = make_match_train_iteration(cfg)
+
+    state, env, _, metrics = train_iteration(state, env, jax.random.PRNGKey(13), jnp.asarray(0.1))
+
+    assert bool(jnp.isfinite(metrics["loss"]))
+    assert float(metrics["matches_completed"]) > 0
+    # One network on both sides: neither seat should be favoured.
+    assert 0.2 < float(metrics["match_win_rate"]) < 0.8
+
+
+# --- the match observation -------------------------------------------------
+
+
+def test_carry_reaches_the_observation_with_a_usable_scale():
+    """The old score_diff divided by 780; a 50-point lead vanished into 0.06."""
+    match = match_reset(jax.random.PRNGKey(14))
+    behind = match._replace(carry=jnp.asarray([0, 50], dtype=jnp.int32))
+    ahead = match._replace(carry=jnp.asarray([50, 0], dtype=jnp.int32))
+
+    obs_behind = np.asarray(match_observation(behind, jnp.int32(0)))
+    obs_ahead = np.asarray(match_observation(ahead, jnp.int32(0)))
+
+    delta = np.abs(obs_ahead - obs_behind)
+    # A 100-point swing has to be plainly visible, not a rounding error.
+    assert delta.max() > 0.5
+    assert (delta > 0.01).sum() >= 2  # the scalar and at least one bin flip
+
+
+@pytest.mark.parametrize("round_idx", [0, 1, 2])
+def test_the_round_index_is_observable(round_idx):
+    match = match_reset(jax.random.PRNGKey(15))._replace(
+        round_idx=jnp.asarray(round_idx, dtype=jnp.int32)
+    )
+    obs = np.asarray(match_observation(match, jnp.int32(0)))
+    assert obs.shape == (MATCH_OBS_DIM,)
+    assert np.isfinite(obs).all()
+
+
+def test_whose_turn_it_is_is_observable():
+    """The single-round obs never said; the critic had to read it off step_count."""
+    match = match_reset(jax.random.PRNGKey(16))
+    mover = int(match.round.to_move)
+
+    from_mover = np.asarray(match_observation(match, jnp.int32(mover)))
+    from_waiter = np.asarray(match_observation(match, jnp.int32(1 - mover)))
+    assert not np.array_equal(from_mover, from_waiter)