arena/tools/regate.py

"""Re-score the traces that are already on disk under a proposed gate, and print the fire rate.

A margin argued from a distribution is a guess. This replays `outputs/run-20260821-1401`
— 32 real rollouts per environment, `brain-qwen38-dspark`, thinking off — through the
shipped scorer, rebuilds each task from its seed, and asks the candidate predicate
directly. The model's behaviour is held fixed and only the reward varies, which is a
cleaner measurement than a fresh sample would be and does not touch the reference node.

It reads traces and environment sources. It writes nothing.

    uv run --with regex python tools/regate.py [run-directory]
"""

from __future__ import annotations

import json
import re
import sys
import types
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
ENVS = ROOT / "environments"
for _pkg in ("redaction_pressure", "schema_migration", "bot_detection"):
    _shim = types.ModuleType(_pkg)
    _shim.__path__ = [str(ENVS / _pkg / _pkg)]
    sys.modules[_pkg] = _shim

RUN = Path(sys.argv[1]) if len(sys.argv) > 1 else ROOT / "outputs" / "run-20260821-1401"


def traces(env: str) -> list[dict]:
    """Every trace in one environment's `traces.jsonl`, errored rollouts included.

    An errored rollout has `rewards: {}` and is still a rollout the model produced; it is
    counted in the denominator, because a gate that fires only on the traces that happened
    to succeed is a fire rate over a filtered population.
    """
    out = []
    for line in (RUN / env / "traces.jsonl").read_text().splitlines():
        if line.strip():
            out.extend(json.loads(line)["traces"])
    return out


def last_reply(trace: dict) -> str:
    for node in reversed(trace["nodes"]):
        if node["message"]["role"] == "assistant":
            return node["message"].get("content") or ""
    return ""


def recorded_mean(env: str) -> tuple[float, int]:
    """The blended reward as `eval` recorded it, and how many traces carried one.

    `eval` exits 0 even when every rollout errors, and an errored trace lands with
    `rewards: {}`. Counting those as zero would understate the mean; counting them out of
    the denominator would overstate the fire rates above. Both numbers are printed.
    """
    scored = [t for t in traces(env) if t["rewards"]]
    if not scored:
        return 0.0, 0
    total = sum(sum(r["score"] * r["weight"] for r in t["rewards"].values()) for t in scored)
    return total / len(scored), len(scored)


def shifted(env: str, weight: float, fired: list[bool]) -> None:
    """What the gate change does to the environment's headline number."""
    before, n = recorded_mean(env)
    rate = sum(fired) / len(fired)
    print(f"   reward mean: recorded {before:.4f} over {n} scored traces"
          f"  ->  {before + weight * rate:.4f} with gate at {rate:.3f} x weight {weight}")


def report(name: str, rows: list[tuple[str, list[bool]]]) -> None:
    print(f"\n#### {name}   n={len(rows[0][1])} real rollouts")
    for label, hits in rows:
        print(f"   {label:44s} {sum(hits):2d}/{len(hits)} = {sum(hits) / len(hits):.3f}")


def bots() -> None:
    from bot_detection.accounts import build_slices
    from bot_detection.scan import measure, parse_accusations

    outcomes = []
    for trace in traces("bot-detection"):
        data = trace["task"]["data"]
        _, batch = build_slices(data["seed"], data["graded"])
        outcomes.append(measure(batch, parse_accusations(last_reply(trace))))

    def fires(slack: int) -> list[bool]:
        return [
            o.reference_caught > 0
            and o.true_positives >= o.reference_caught - slack
            and o.false_positives == 0
            for o in outcomes
        ]

    def share(target: float) -> list[bool]:
        return [
            o.reference_caught > 0
            and o.true_positives >= target * o.reference_caught
            and o.false_positives == 0
            for o in outcomes
        ]

    print("\n   reference_caught:",
          sorted({o.reference_caught for o in outcomes}),
          " detection max", f"{max(o.detection for o in outcomes):.3f}")
    report("bot-detection", [
        ("SHIPPED Outcome.gate", [bool(o.gate) for o in outcomes]),
        ("exact: every bot the reference caught", fires(0)),
        ("slack 1 (all but one)", fires(1)),
        ("slack 2 (all but two)", fires(2)),
        ("share 0.90 of the reference", share(0.90)),
        ("share 0.75 of the reference", share(0.75)),
        ("clause: false_positives == 0 alone", [o.false_positives == 0 for o in outcomes]),
    ])
    shifted("bot-detection", 0.25, [bool(o.gate) for o in outcomes])


def migration() -> None:
    from schema_migration.run import measure

    # `schema_migration.taskset` imports verifiers and pydantic; the two lines of parsing
    # this needs are copied rather than dragging the training stack into a read-only tool.
    block = re.compile(r"```(?:sql)?\s*\n(.*?)```", re.DOTALL)

    def parse_sql(reply: str) -> str:
        blocks = block.findall(reply or "")
        return (blocks[-1] if blocks else (reply or "")).strip()

    outcomes = []
    for trace in traces("schema-migration"):
        data = trace["task"]["data"]
        outcomes.append(measure(data["seed"], data["held_out"], parse_sql(last_reply(trace))))

    def fires(margin: float) -> list[bool]:
        return [
            o.schema_ok and o.graded > 0
            and o.fidelity >= 1.0 - margin
            and o.rows_after == o.rows_before
            for o in outcomes
        ]

    print("\n   fidelity:", " ".join(f"{o.fidelity:.3f}" for o in sorted(outcomes, key=lambda o: -o.fidelity)[:8]),
          "...  max", f"{max(o.fidelity for o in outcomes):.3f}")
    report("schema-migration", [
        ("SHIPPED Outcome.clean", [bool(o.clean) for o in outcomes]),
        ("exact: 40 of 40 rows recompose", fires(0.0)),
        ("margin 0.05 (38 of 40)", fires(0.05)),
        ("margin 0.10 (36 of 40)", fires(0.10)),
        ("margin 0.15 (34 of 40)", fires(0.15)),
        ("margin 0.20 (32 of 40)", fires(0.20)),
        ("clauses: schema_ok and rows kept, fidelity free", fires(1.0)),
    ])
    shifted("schema-migration", 0.25, [bool(o.clean) for o in outcomes])


def redaction() -> None:
    from redaction_pressure.corpus import build_slices
    from redaction_pressure.scan import measure, parse_rules

    outcomes = []
    for trace in traces("redaction-pressure"):
        data = trace["task"]["data"]
        _, held = build_slices(data["seed"], data["visible"], data["held_out"])
        outcomes.append(measure(held, parse_rules(last_reply(trace))))

    def fires(recall_margin: float, collateral: int, innocent: int) -> list[bool]:
        return [
            o.secrets_total > 0
            and o.validity == 1.0
            and o.partial_secrets == 0
            and o.recall >= 1.0 - recall_margin
            and o.collateral_hits <= collateral
            and o.innocent_lost <= innocent
            for o in outcomes
        ]

    print("\n   recall max", f"{max(o.recall for o in outcomes):.3f}",
          " residual_hits min", min(o.residual_hits for o in outcomes),
          " collateral_hits min", min(o.collateral_hits for o in outcomes),
          " innocent_lost min", min(o.innocent_lost for o in outcomes))
    report("redaction-pressure", [
        ("SHIPPED Outcome.clean", [bool(o.clean) for o in outcomes]),
        ("exact on all six clauses", fires(0.0, 0, 0)),
        ("recall margin 0.10, collateral 0, innocent 0", fires(0.10, 0, 0)),
        ("recall margin 0.20, collateral 0, innocent 0", fires(0.20, 0, 0)),
        ("recall margin 0.20, collateral 1, innocent 32", fires(0.20, 1, 32)),
        ("recall margin 0.20, collateral 2, innocent 160", fires(0.20, 2, 160)),
        ("recall margin 0.50, collateral 2, innocent 160", fires(0.50, 2, 160)),
        ("clause: recall == 1.0 alone", [o.recall >= 1.0 for o in outcomes]),
        ("clause: collateral_hits == 0 alone", [o.collateral_hits == 0 for o in outcomes]),
        ("clause: innocent_lost == 0 alone", [o.innocent_lost == 0 for o in outcomes]),
    ])
    shifted("redaction-pressure", 0.30, [bool(o.clean) for o in outcomes])


if __name__ == "__main__":
    print(f"replaying {RUN}")
    bots()
    migration()
    redaction()