Add the altair-vs-borealis measurement driver

The script behind runs/jax-ppo-match/altair_vs_borealis.json: duplicate, seat-swapped play between the two models across single-deal and three-round conditions, reporting win rate with a Wilson interval and mean margin. Kept so the 2x2 result is reproducible. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh
2026-07-15 18:11:27 +09:00
parent c8fa89dab9
commit 774faf2b51
1 changed files with 201 additions and 0 deletions
@@ -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()