From 9ba5a07a9d53bf1fa34883efd6e7dc405a0d771b 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:23:50 +0900 Subject: [PATCH] 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 Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh --- src/lost_cities_jax/match_eval.py | 195 ++++++++++++ src/lost_cities_jax/match_obs.py | 198 ++++++++++++ src/lost_cities_jax/match_ppo.py | 407 ++++++++++++++++++++++++ tests/lost_cities_jax/test_match_ppo.py | 194 +++++++++++ 4 files changed, 994 insertions(+) create mode 100644 src/lost_cities_jax/match_eval.py create mode 100644 src/lost_cities_jax/match_obs.py create mode 100644 src/lost_cities_jax/match_ppo.py create mode 100644 tests/lost_cities_jax/test_match_ppo.py diff --git a/src/lost_cities_jax/match_eval.py b/src/lost_cities_jax/match_eval.py new file mode 100644 index 0000000..32857ca --- /dev/null +++ b/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 diff --git a/src/lost_cities_jax/match_obs.py b/src/lost_cities_jax/match_obs.py new file mode 100644 index 0000000..d839177 --- /dev/null +++ b/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", +] diff --git a/src/lost_cities_jax/match_ppo.py b/src/lost_cities_jax/match_ppo.py new file mode 100644 index 0000000..8456f1a --- /dev/null +++ b/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 diff --git a/tests/lost_cities_jax/test_match_ppo.py b/tests/lost_cities_jax/test_match_ppo.py new file mode 100644 index 0000000..cd31528 --- /dev/null +++ b/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)