diff --git a/scripts/altair_vs_borealis.py b/scripts/altair_vs_borealis.py new file mode 100644 index 0000000..a3b46c6 --- /dev/null +++ b/scripts/altair_vs_borealis.py @@ -0,0 +1,201 @@ +"""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()