benchmark_rlvr_pool_batched.py

#!/usr/bin/env python3
"""Fast exact-engine matched pre/post evaluation for temporal RLVR models."""

from __future__ import annotations

import argparse
from collections import defaultdict
import json
import math
import multiprocessing
from pathlib import Path
import random
import time

import numpy as np
import torch

from benchmark_rlvr_pool import (
    OPPONENTS,
    PLAYER_WEIGHTS,
    _opponent_summaries,
    _stratum_summary,
    _weighted_summary,
    file_sha256,
)
from batched_pesten import (
    BatchedPesten,
    PHASE_SUIT,
    POLICY_RANDOM,
    POLICY_RL,
    POLICY_SMART,
    POLICY_SMARTER,
)


_POLICIES = None
_MLP_FACTORY = None


def _load_policy(path: Path):
    from temporal_puffer_v4 import load_temporal_policy_checkpoint

    policy, _ = load_temporal_policy_checkpoint(path)
    policy.eval()
    for parameter in policy.parameters():
        parameter.requires_grad_(False)
    return policy


def _initialize_worker(
    before: str,
    after: str,
    mlp: str,
    self_anchor: str,
) -> None:
    _initialize_multi_worker(
        {
            "before": before,
            "after": after,
            "self_anchor": self_anchor,
        },
        mlp,
    )


def _initialize_multi_worker(
    policy_paths: dict[str, str],
    mlp: str,
) -> None:
    global _POLICIES, _MLP_FACTORY
    from benchmark_agents import load_ppo_agent

    torch.set_num_threads(1)
    try:
        torch.set_num_interop_threads(1)
    except RuntimeError:
        pass
    loaded = {}
    cache = {}
    for name, raw_path in policy_paths.items():
        path = Path(raw_path)
        key = str(path.resolve())
        if key not in cache:
            cache[key] = _load_policy(path)
        loaded[name] = cache[key]
    _POLICIES = loaded
    _MLP_FACTORY = load_ppo_agent(Path(mlp), "cpu")


def _make_layout(task: dict):
    players = int(task["players"])
    blocks = int(task["blocks"])
    opponent = str(task["opponent"])
    variants = tuple(task.get("variants", ("before", "after")))
    policy_names = set(
        task.get("policy_names", ("before", "after", "self_anchor"))
    )
    if not variants or any(variant not in policy_names for variant in variants):
        raise ValueError("variants must be a non-empty subset of policy names")
    self_anchor_name = str(task.get("self_anchor_name", "self_anchor"))
    if opponent == "self_anchor" and self_anchor_name not in policy_names:
        raise ValueError("self anchor must name one initialized policy")
    records = []
    seeds = []
    roles = []
    for block in range(blocks):
        deal_seed = int(task["seed"]) + block
        for variant in variants:
            for candidate_seat in range(players):
                row_roles = []
                for seat in range(players):
                    if seat == candidate_seat:
                        row_roles.append(("temporal", variant))
                    elif opponent == "self_anchor":
                        row_roles.append(("temporal", self_anchor_name))
                    elif opponent == "mlp":
                        row_roles.append(("mlp", None))
                    elif opponent == "canonical_random":
                        row_roles.append(("canonical_random", None))
                    elif opponent == "smart_rule":
                        row_roles.append(("smart_rule", None))
                    elif opponent == "smarter_rule":
                        row_roles.append(("smarter_rule", None))
                    else:
                        raise ValueError(f"unknown opponent: {opponent}")
                records.append((block, variant, candidate_seat, deal_seed))
                seeds.append(deal_seed)
                roles.append(row_roles)
    return records, roles, np.asarray(seeds, dtype=np.uint32)


def _evaluate_task(task: dict) -> dict:
    if _POLICIES is None or _MLP_FACTORY is None:
        raise RuntimeError("batched benchmark worker was not initialized")
    from temporal_batched_speed import (
        SafeEventCursor,
        TemporalSeatRuntime,
        _batch_champion_actions,
        _batch_temporal_actions,
        _champion_has_card_decision,
        build_safe_decision,
    )

    players = int(task["players"])
    blocks = int(task["blocks"])
    records, roles, seeds = _make_layout(task)
    policy_kinds = np.full((len(records), players), POLICY_RL, dtype=np.int8)
    actor_keys = np.empty((len(records), players), dtype=object)
    for row, row_roles in enumerate(roles):
        for seat, (kind, policy_name) in enumerate(row_roles):
            if kind == "canonical_random":
                policy_kinds[row, seat] = POLICY_RANDOM
            elif kind == "smart_rule":
                policy_kinds[row, seat] = POLICY_SMART
            elif kind == "smarter_rule":
                policy_kinds[row, seat] = POLICY_SMARTER
            actor_keys[row, seat] = (kind, policy_name or "")
    engine = BatchedPesten(
        policy_kinds,
        actor_keys,
        seeds,
        max_turns=1000,
        record_public_events=True,
    )
    runtimes = {}
    suit_generators = {}
    for row, row_roles in enumerate(roles):
        for seat, (kind, policy_name) in enumerate(row_roles):
            if kind == "temporal":
                policy = _POLICIES[policy_name]
                runtimes[(row, seat)] = TemporalSeatRuntime(
                    game_index=row,
                    absolute_game_index=row,
                    seat=seat,
                    policy=policy,
                    generator=torch.Generator(device="cpu").manual_seed(
                        int(seeds[row]) * 2 + seat
                    ),
                    state=policy.initial_state(),
                    cursor=SafeEventCursor(),
                )
            elif kind == "canonical_random":
                suit_generators[row] = random.Random(int(seeds[row]))

    while True:
        active = engine.active_games()
        if len(active) == 0:
            break
        actions = {}
        temporal_groups = {}
        champion_games = []
        for row_value in active:
            row = int(row_value)
            seat = int(engine.current_player[row])
            kind, _ = roles[row][seat]
            if kind == "temporal":
                runtime = runtimes[(row, seat)]
                decision = build_safe_decision(
                    engine, row, seat, runtime.cursor
                )
                disable_strategic_pass = bool(
                    task.get(f"disable_pass_{roles[row][seat][1]}", False)
                )
                pass_only_one_card = bool(
                    task.get(
                        f"pass_only_one_card_{roles[row][seat][1]}", False
                    )
                )
                strategic_pass_max_hand_size = (
                    1 if pass_only_one_card else None
                )
                group = temporal_groups.setdefault(
                    (
                        id(runtime.policy),
                        disable_strategic_pass,
                        strategic_pass_max_hand_size,
                    ),
                    (
                        runtime.policy,
                        disable_strategic_pass,
                        strategic_pass_max_hand_size,
                        [],
                    ),
                )
                group[3].append((runtime, decision))
            elif kind == "mlp":
                if (
                    int(engine.phase[row]) != PHASE_SUIT
                    and not _champion_has_card_decision(engine, row)
                ):
                    actions[row] = 54
                else:
                    champion_games.append(row)
            elif kind == "canonical_random":
                if int(engine.phase[row]) == PHASE_SUIT:
                    actions[row] = 50 + suit_generators[row].randint(0, 3)
                else:
                    actions[row] = engine.random_action(row)
            elif kind in ("smart_rule", "smarter_rule"):
                actions[row] = engine.non_neural_action(row)
            else:
                raise AssertionError(f"unexpected role {kind}")

        for _, (
            policy,
            disable_strategic_pass,
            strategic_pass_max_hand_size,
            entries,
        ) in temporal_groups.items():
            chosen = _batch_temporal_actions(
                policy,
                entries,
                None,
                deterministic=True,
                disable_strategic_pass=disable_strategic_pass,
                strategic_pass_max_hand_size=(
                    strategic_pass_max_hand_size
                ),
            )
            for (runtime, _), action in zip(entries, chosen):
                actions[runtime.game_index] = action
        if champion_games:
            chosen = _batch_champion_actions(
                engine, champion_games, _MLP_FACTORY, None
            )
            actions.update(zip(champion_games, chosen))
        if len(actions) != len(active):
            raise RuntimeError("each active benchmark game needs one action")
        for row_value in active:
            row = int(row_value)
            engine.apply_action(row, actions[row])

    variants = tuple(task.get("variants", ("before", "after")))
    scores = {
        variant: np.zeros(blocks, dtype=np.float64) for variant in variants
    }
    win_shares = {
        variant: np.zeros(blocks, dtype=np.float64) for variant in variants
    }
    loss_shares = {
        variant: np.zeros(blocks, dtype=np.float64) for variant in variants
    }
    draw_shares = {
        variant: np.zeros(blocks, dtype=np.float64) for variant in variants
    }
    totals = {variant: defaultdict(int) for variant in variants}
    for row, (block, variant, candidate_seat, _) in enumerate(records):
        winner = int(engine.winners[row])
        totals[variant]["turns"] += int(engine.turn_counts[row])
        if winner < 0:
            totals[variant]["draws"] += 1
            scores[variant][block] += 1.0 / players
            draw_shares[variant][block] += 1.0
        elif winner == candidate_seat:
            totals[variant]["wins"] += 1
            scores[variant][block] += 1.0
            win_shares[variant][block] += 1.0
        else:
            totals[variant]["losses"] += 1
            loss_shares[variant][block] += 1.0
    for variant in scores:
        scores[variant] /= players
        win_shares[variant] /= players
        loss_shares[variant] /= players
        draw_shares[variant] /= players
    arrays = {}
    for variant in variants:
        arrays[f"{variant}_scores"] = scores[variant].tolist()
        arrays[f"{variant}_win_shares"] = win_shares[variant].tolist()
        arrays[f"{variant}_loss_shares"] = loss_shares[variant].tolist()
        arrays[f"{variant}_draw_shares"] = draw_shares[variant].tolist()
    return {
        "opponent": task["opponent"],
        "players": players,
        **arrays,
        "totals": {
            name: dict(values) for name, values in totals.items()
        },
    }


def _weighted_outcome_summary(strata: list[dict]) -> dict:
    """Aggregate matched win/loss/draw rates with the campaign mixture."""

    result = {}
    for outcome in ("win", "loss", "draw"):
        before = after = variance = 0.0
        for row in strata:
            players = int(row["players"])
            weight = PLAYER_WEIGHTS[players] / len(OPPONENTS)
            before_values = np.asarray(
                row[f"before_{outcome}_shares"], dtype=np.float64
            )
            after_values = np.asarray(
                row[f"after_{outcome}_shares"], dtype=np.float64
            )
            if len(before_values) != len(after_values) or not len(before_values):
                raise ValueError("outcome arrays must be non-empty matched blocks")
            differences = after_values - before_values
            before += weight * float(before_values.mean())
            after += weight * float(after_values.mean())
            if len(differences) > 1:
                variance += (
                    weight
                    * weight
                    * float(differences.var(ddof=1))
                    / len(differences)
                )
        delta = after - before
        standard_error = math.sqrt(variance)
        interval = [
            delta - 3.2905267314919255 * standard_error,
            delta + 3.2905267314919255 * standard_error,
        ]
        result[outcome] = {
            "before_rate": before,
            "after_rate": after,
            "delta_after_minus_before": delta,
            "delta_standard_error": standard_error,
            "delta_normal_ci99_9": interval,
        }
    result["loss_avoidance"] = {
        "delta_before_minus_after": -result["loss"][
            "delta_after_minus_before"
        ],
        "delta_normal_ci99_9": [
            -result["loss"]["delta_normal_ci99_9"][1],
            -result["loss"]["delta_normal_ci99_9"][0],
        ],
    }
    return result


def run_batched_pool_benchmark(
    before: Path,
    after: Path,
    mlp: Path,
    self_anchor: Path,
    *,
    blocks_per_stratum: int,
    workers: int,
    chunks_per_stratum: int,
    seed: int,
) -> dict:
    if min(blocks_per_stratum, workers, chunks_per_stratum) < 1:
        raise ValueError("benchmark counts must be positive")
    tasks = []
    chunk_size = math.ceil(blocks_per_stratum / chunks_per_stratum)
    for opponent_index, opponent in enumerate(OPPONENTS):
        for players in range(2, 9):
            stratum_seed = seed + opponent_index * 100_000_000 + players * 1_000_000
            for start in range(0, blocks_per_stratum, chunk_size):
                tasks.append(
                    {
                        "opponent": opponent,
                        "players": players,
                        "seed": stratum_seed + start,
                        "blocks": min(
                            chunk_size, blocks_per_stratum - start
                        ),
                    }
                )
    context = multiprocessing.get_context("spawn")
    started = time.perf_counter()
    with context.Pool(
        processes=min(workers, len(tasks)),
        initializer=_initialize_worker,
        initargs=(
            str(before),
            str(after),
            str(mlp),
            str(self_anchor),
        ),
    ) as pool:
        chunks = pool.map(_evaluate_task, tasks, chunksize=1)
    grouped = {}
    for chunk in chunks:
        key = (chunk["opponent"], chunk["players"])
        target = grouped.setdefault(
            key,
            {
                "opponent": chunk["opponent"],
                "players": chunk["players"],
                "before_scores": [],
                "after_scores": [],
                "before_win_shares": [],
                "after_win_shares": [],
                "before_loss_shares": [],
                "after_loss_shares": [],
                "before_draw_shares": [],
                "after_draw_shares": [],
                "totals": {
                    "before": defaultdict(int),
                    "after": defaultdict(int),
                },
            },
        )
        target["before_scores"].extend(chunk["before_scores"])
        target["after_scores"].extend(chunk["after_scores"])
        for variant in ("before", "after"):
            for outcome in ("win", "loss", "draw"):
                target[f"{variant}_{outcome}_shares"].extend(
                    chunk[f"{variant}_{outcome}_shares"]
                )
        for variant in ("before", "after"):
            for name, value in chunk["totals"][variant].items():
                target["totals"][variant][name] += value
    raw = []
    strata = []
    for key in sorted(grouped):
        row = grouped[key]
        row["totals"] = {
            variant: dict(values)
            for variant, values in row["totals"].items()
        }
        raw.append(row)
        strata.append(_stratum_summary(row))
    return {
        "strata": strata,
        "weighted_result": _weighted_summary(raw),
        "weighted_outcomes": _weighted_outcome_summary(raw),
        "opponent_results": _opponent_summaries(raw),
        "elapsed_seconds": time.perf_counter() - started,
    }


def run_multi_policy_benchmark(
    policies: dict[str, Path],
    mlp: Path,
    *,
    self_anchor_name: str,
    reference_name: str,
    blocks_by_stratum: dict[tuple[str, int], int],
    workers: int,
    chunks_per_stratum: int,
    seed: int,
    target_normalized_ci_half_width: float,
    normal_quantile: float = 3.2905267314919255,
    confidence_label: str = "99.9 percent",
) -> dict:
    """Evaluate several actors once and form paired contrasts to a reference."""

    if len(policies) < 2 or reference_name not in policies:
        raise ValueError("multi-policy evaluation needs a named reference")
    if self_anchor_name not in policies:
        raise ValueError("multi-policy evaluation needs a fixed self anchor")
    if workers < 1 or chunks_per_stratum < 1:
        raise ValueError("benchmark worker and chunk counts must be positive")
    if target_normalized_ci_half_width <= 0:
        raise ValueError("confidence-interval target must be positive")
    if normal_quantile <= 0 or not confidence_label:
        raise ValueError("normal confidence specification must be positive")
    expected_strata = {
        (opponent, players)
        for opponent in OPPONENTS
        for players in range(2, 9)
    }
    if set(blocks_by_stratum) != expected_strata:
        missing = sorted(expected_strata - set(blocks_by_stratum))
        extra = sorted(set(blocks_by_stratum) - expected_strata)
        raise ValueError(
            f"allocation must cover every benchmark stratum; "
            f"missing={missing}, extra={extra}"
        )
    if any(blocks < 2 for blocks in blocks_by_stratum.values()):
        raise ValueError("each benchmark stratum needs at least two blocks")

    candidate_names = tuple(
        name
        for name in policies
        if name != self_anchor_name
    )
    if reference_name not in candidate_names:
        raise ValueError("the fixed self anchor cannot also be the reference arm")
    policy_names = tuple(policies)
    tasks = []
    for opponent_index, opponent in enumerate(OPPONENTS):
        for players in range(2, 9):
            blocks = blocks_by_stratum[(opponent, players)]
            chunk_size = math.ceil(blocks / chunks_per_stratum)
            stratum_seed = seed + opponent_index * 100_000_000 + players * 1_000_000
            for start in range(0, blocks, chunk_size):
                tasks.append(
                    {
                        "opponent": opponent,
                        "players": players,
                        "seed": stratum_seed + start,
                        "blocks": min(chunk_size, blocks - start),
                        "variants": candidate_names,
                        "policy_names": policy_names,
                        "self_anchor_name": self_anchor_name,
                    }
                )

    context = multiprocessing.get_context("spawn")
    started = time.perf_counter()
    with context.Pool(
        processes=min(workers, len(tasks)),
        initializer=_initialize_multi_worker,
        initargs=(
            {name: str(path) for name, path in policies.items()},
            str(mlp),
        ),
    ) as pool:
        chunks = []
        for completed, chunk in enumerate(
            pool.imap(_evaluate_task, tasks, chunksize=1), start=1
        ):
            chunks.append(chunk)
            if completed == 1 or completed % 10 == 0 or completed == len(tasks):
                print(
                    json.dumps(
                        {
                            "kind": "multiarm_evaluation_progress",
                            "tasks_completed": completed,
                            "tasks_total": len(tasks),
                            "elapsed_seconds": time.perf_counter() - started,
                        },
                        sort_keys=True,
                    ),
                    flush=True,
                )

    grouped = {}
    for chunk in chunks:
        key = (chunk["opponent"], int(chunk["players"]))
        target = grouped.setdefault(
            key,
            {
                "opponent": chunk["opponent"],
                "players": int(chunk["players"]),
                "scores": {name: [] for name in candidate_names},
                "totals": {
                    name: defaultdict(int) for name in candidate_names
                },
            },
        )
        for name in candidate_names:
            target["scores"][name].extend(chunk[f"{name}_scores"])
            for metric, value in chunk["totals"][name].items():
                target["totals"][name][metric] += value

    pairwise = {}
    for candidate_name in candidate_names:
        if candidate_name == reference_name:
            continue
        raw = []
        strata = []
        for key in sorted(grouped):
            grouped_row = grouped[key]
            row = {
                "opponent": grouped_row["opponent"],
                "players": grouped_row["players"],
                "before_scores": grouped_row["scores"][reference_name],
                "after_scores": grouped_row["scores"][candidate_name],
                "totals": {
                    "before": dict(grouped_row["totals"][reference_name]),
                    "after": dict(grouped_row["totals"][candidate_name]),
                },
            }
            raw.append(row)
            strata.append(_stratum_summary(row))
        weighted = _weighted_summary(raw)
        half_width = (
            normal_quantile * weighted["delta_normalized_standard_error"]
        )
        selected_interval = [
            weighted["delta_normalized_win_share"] - half_width,
            weighted["delta_normalized_win_share"] + half_width,
        ]
        weighted["selected_normal_confidence"] = confidence_label
        weighted["selected_normal_quantile"] = normal_quantile
        weighted["delta_normalized_selected_normal_ci"] = selected_interval
        weighted["delta_normalized_selected_normal_ci_half_width"] = (
            half_width
        )
        weighted["positive_selected_normal"] = selected_interval[0] > 0
        weighted["negative_selected_normal"] = selected_interval[1] < 0
        if math.isclose(normal_quantile, 3.2905267314919255):
            weighted["delta_normalized_ci99_9_half_width"] = half_width
        elif math.isclose(normal_quantile, 2.5758293035489004):
            weighted["delta_normalized_ci99"] = selected_interval
            weighted["delta_normalized_ci99_half_width"] = half_width
        weighted["precision_target_half_width"] = (
            target_normalized_ci_half_width
        )
        weighted["precision_target_met"] = (
            half_width <= target_normalized_ci_half_width
        )
        pairwise[candidate_name] = {
            "reference": reference_name,
            "candidate": candidate_name,
            "strata": strata,
            "weighted_result": weighted,
            "opponent_results": _opponent_summaries(raw),
        }

    return {
        "pairwise": pairwise,
        "elapsed_seconds": time.perf_counter() - started,
        "tasks": len(tasks),
        "candidate_names": list(candidate_names),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--before", type=Path, required=True)
    parser.add_argument("--after", type=Path, required=True)
    parser.add_argument("--mlp", type=Path, required=True)
    parser.add_argument("--self-anchor", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--blocks-per-stratum", type=int, default=300)
    parser.add_argument("--workers", type=int, default=10)
    parser.add_argument("--chunks-per-stratum", type=int, default=3)
    parser.add_argument("--seed", type=int, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    before = args.before.resolve()
    after = args.after.resolve()
    mlp = args.mlp.resolve()
    self_anchor = (
        args.self_anchor.resolve() if args.self_anchor else before
    )
    for path in (before, after, mlp, self_anchor):
        if not path.is_file():
            raise FileNotFoundError(path)
    result = run_batched_pool_benchmark(
        before,
        after,
        mlp,
        self_anchor,
        blocks_per_stratum=args.blocks_per_stratum,
        workers=args.workers,
        chunks_per_stratum=args.chunks_per_stratum,
        seed=args.seed,
    )
    payload = {
        "schema_version": 1,
        "kind": "matched_rlvr_pool_evaluation_batched_exact_engine",
        "before": {"path": str(before), "sha256": file_sha256(before)},
        "after": {"path": str(after), "sha256": file_sha256(after)},
        "self_anchor": {
            "path": str(self_anchor),
            "sha256": file_sha256(self_anchor),
        },
        "mlp": {"path": str(mlp), "sha256": file_sha256(mlp)},
        "protocol": {
            "opponents": list(OPPONENTS),
            "player_weights": {
                str(key): value for key, value in PLAYER_WEIGHTS.items()
            },
            "bulk_2_3_4_fraction": sum(
                PLAYER_WEIGHTS[key] for key in (2, 3, 4)
            ),
            "tail_5_8_fraction": sum(
                PLAYER_WEIGHTS[key] for key in (5, 6, 7, 8)
            ),
            "blocks_per_stratum": args.blocks_per_stratum,
            "seat_rotations_per_block": "equal to player count",
            "matched_pre_post_seeds": True,
            "fixed_self_anchor": True,
            "deterministic_temporal_actions": True,
            "engine": "BatchedPesten exact-semantics evaluator",
            "confidence": (
                "paired normal and distribution-free Hoeffding 99.9% intervals"
            ),
            "seed": args.seed,
        },
        **result,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(
        json.dumps(payload["weighted_result"], indent=2, sort_keys=True),
        flush=True,
    )


if __name__ == "__main__":
    main()