"""altair (single-round, gen a) vs borealis (3-round match, gen b). Full 2x2: {single deal, 3-round match} x {win rate + Wilson, mean margin + CI}. Both policies act on the SAME MatchState but from their own view: - borealis reads the full MatchState (match_observation + privileged critic). - altair reads only the round in play as a single-round State (observation). We play whole matches (duplicated: every deal-triple from both seats). From the same runs we harvest two scoring conventions: - single deal = the round-0 board score, snapshotted the ply round 0 rolls over (carry=0, round_idx=0 there, so it is an honest standalone deal). - 3-round match = match_score(final), the summed total. Margins are borealis-minus-altair. Self-play (borealis vs borealis) is run as a harness check: mirrored identical policies must give win rate 0.5 / margin 0.0. """ from __future__ import annotations import json import math from pathlib import Path import jax import jax.numpy as jnp import numpy as np from lost_cities_jax.match import ( MatchState, match_legal_action_mask, match_reset_from, match_score, match_step, ) from lost_cities_jax.match_eval import MATCH_SCAN_STEPS, _wilson, match_bank from lost_cities_jax.match_obs import match_critic_observation, match_observation from lost_cities_jax.match_ppo import Ablation, MatchActorCritic, create_match_train_state from lost_cities_jax.obs import observation from lost_cities_jax.ppo import ( ActorCritic, create_train_state, load_config, mask_logits, restore_checkpoint, ) ALTAIR_CKPT = Path( "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/latest" ) BOREALIS_CKPT = Path("runs/jax-ppo-match/2026-07-15_031529_match-scaled/latest") MATCHES = 4096 # -> 8192 duplicate games per cell SEED = 20260715 OUT = Path("runs/jax-ppo-match/altair_vs_borealis.json") borealis_cfg = load_config("configs/jax_ppo/match-selfplay.yaml") altair_cfg = load_config("configs/jax_ppo/balanced.yaml") ABL = Ablation() # privileged_critic=True, as borealis was trained borealis_st = restore_checkpoint( BOREALIS_CKPT, create_match_train_state(borealis_cfg, jax.random.PRNGKey(0), ABL) ) altair_st = restore_checkpoint(ALTAIR_CKPT, create_train_state(altair_cfg, jax.random.PRNGKey(0))) borealis_model = MatchActorCritic(borealis_cfg.network.hidden_size, borealis_cfg.network.num_layers) altair_model = ActorCritic(altair_cfg.network.hidden_size, altair_cfg.network.num_layers) def _borealis_action(env: MatchState, to_move, mask): obs = jax.vmap(match_observation)(env, to_move) crit = jax.vmap(match_critic_observation)(env, to_move) logits, _ = borealis_model.apply(borealis_st.params, obs, crit) return jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32) def _altair_action(env: MatchState, to_move, mask): obs = jax.vmap(observation)(env.round, to_move) logits, _ = altair_model.apply(altair_st.params, obs) return jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32) def _make_run(action_seat0, action_seat1): """Build a jitted full-match runner. ``action_seat0`` is the policy that plays when ``to_move == borealis_seat`` (i.e. borealis); ``action_seat1`` is the other policy (altair). Selection is by ``borealis_seat`` so mirrored orientation is a pure seat relabel. """ @jax.jit def run(env: MatchState, borealis_seat): def body(carry, _): env, r0_snap = carry to_move = env.round.to_move.astype(jnp.int32) mask = jax.vmap(match_legal_action_mask)(env) a0 = action_seat0(env, to_move, mask) a1 = action_seat1(env, to_move, mask) action = jnp.where(to_move == borealis_seat, a0, a1) was_r0 = env.round_idx == 0 nxt, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action) just_finished_r0 = was_r0 & (nxt.round_idx == 1) # nxt.carry == round-0 board score exactly on the roll-over ply. r0_snap = jnp.where(just_finished_r0[:, None], nxt.carry, r0_snap) return (nxt, r0_snap), None n = env.done.shape[0] r0_snap = jnp.zeros((n, 2), dtype=jnp.int32) (env, r0_snap), _ = jax.lax.scan(body, (env, r0_snap), xs=None, length=MATCH_SCAN_STEPS) return env, r0_snap return run def _summ(margin: np.ndarray) -> dict: """margin = borealis - altair, per duplicate game. Positive = borealis wins.""" n = int(margin.size) b_wins = int((margin > 0).sum()) a_wins = int((margin < 0).sum()) ties = int((margin == 0).sum()) lo, hi = _wilson(float(b_wins), float(n)) std = float(margin.std(ddof=1)) sem = std / math.sqrt(n) return { "n_duplicate_games": n, "borealis_wins": b_wins, "altair_wins": a_wins, "ties": ties, "borealis_win_rate": b_wins / n, "wilson_95": [lo, hi], "mean_margin_borealis_minus_altair": float(margin.mean()), "margin_std": std, "margin_sem": sem, "margin_95ci": [float(margin.mean() - 1.96 * sem), float(margin.mean() + 1.96 * sem)], } def _play(run, decks, coins, borealis_first: bool): """Duplicate play; returns (single_deal_margins, match_margins). ``borealis_first`` picks which policy is action_seat0 in the runner. When the two policies are identical (self-play) this must yield perfectly antisymmetric margins -> win rate 0.5, margin 0. """ single, match = [], [] for seat in (0, 1): env = jax.vmap(match_reset_from)(decks, coins) b_seat = jnp.full((MATCHES,), seat, dtype=jnp.int32) final, r0 = run(env, b_seat) r0 = np.asarray(r0) totals = np.asarray(jax.vmap(match_score)(final)) # borealis is at index ``seat``. single.append(r0[:, seat] - r0[:, 1 - seat]) match.append(totals[:, seat] - totals[:, 1 - seat]) return np.concatenate(single), np.concatenate(match) def main(): decks, coins = match_bank(SEED, MATCHES) # --- Harness check: borealis vs borealis (both seats borealis) --- run_self = _make_run(_borealis_action, _borealis_action) self_single, self_match = _play(run_self, decks, coins, True) self_check = { "single_deal": _summ(self_single), "three_round_match": _summ(self_match), } # --- Real comparison: altair vs borealis --- # seat0-slot = borealis (selected when to_move == borealis_seat), seat1 = altair run_av = _make_run(_borealis_action, _altair_action) av_single, av_match = _play(run_av, decks, coins, True) comparison = { "single_deal": _summ(av_single), "three_round_match": _summ(av_match), } result = { "meta": { "altair_ckpt": str(ALTAIR_CKPT), "borealis_ckpt": str(BOREALIS_CKPT), "matches_per_orientation": MATCHES, "duplicate_games_per_cell": 2 * MATCHES, "match_scan_steps": MATCH_SCAN_STEPS, "seed": SEED, "privileged_critic": ABL.privileged_critic, "policies": "greedy argmax (deterministic given the deal)", "margin_sign": "borealis total minus altair total", }, "harness_check_borealis_vs_borealis": self_check, "altair_vs_borealis": comparison, } OUT.parent.mkdir(parents=True, exist_ok=True) OUT.write_text(json.dumps(result, indent=2)) print(json.dumps(result, indent=2)) if __name__ == "__main__": main()