coorl-lost-cities/scripts/altair_vs_borealis.py

"""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()