workspace/scripts/build_ctt_rollout_comparison.py

#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import math
import re
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from cil.metrics import (  # noqa: E402
    any_unsafe,
    normalized_causal_action_regret,
    outcome_safety_violation,
    safety_label_coverage,
    unsafe_rate,
)
from scripts.eval_metrics import main as eval_metrics_main  # noqa: E402


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Aggregate measured CTT rollout runs.")
    parser.add_argument(
        "--run-glob",
        default="runs/ctt_residual_rollout_val69_seed*",
        help="Glob for run directories containing measured_candidates.json.",
    )
    parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_val_rollout_comparison"))
    parser.add_argument("--k", type=int, default=8)
    parser.add_argument("--bootstrap-samples", type=int, default=1000)
    parser.add_argument(
        "--no-markdown-report",
        action="store_true",
        help="Do not write report.md; persistent prose lives in README.md.",
    )
    args = parser.parse_args(argv)

    run_dirs = [
        path
        for path in sorted(Path().glob(args.run_glob))
        if (path / "measured_candidates.json").exists()
    ]
    if not run_dirs:
        raise SystemExit(f"no measured rollout runs found for glob {args.run_glob!r}")

    combined_rows: list[dict[str, Any]] = []
    payloads = []
    for run_dir in run_dirs:
        payload = json.loads((run_dir / "measured_candidates.json").read_text())
        payloads.append(payload)
        train_seed = _seed_from_path(run_dir)
        for row in payload.get("rows", []):
            item = dict(row)
            item["train_seed"] = train_seed
            item["seed"] = f"{row.get('seed', 'unknown')}/train{train_seed}"
            combined_rows.append(item)

    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    combined = {
        "report_type": "ctt_measured_rollout_comparison",
        "schema_version": 1,
        "k": args.k,
        "run_dirs": [str(path) for path in run_dirs],
        "train_seeds": [_seed_from_path(path) for path in run_dirs],
        "num_rows": len(combined_rows),
        "data_hash": _first(payloads, "target_content_hash"),
        "split_hash": _first(payloads, "target_split_hash"),
        "source_content_hash": _first(payloads, "source_content_hash"),
        "target_content_hash": _first(payloads, "target_content_hash"),
        "target_split_hash": _first(payloads, "target_split_hash"),
        "rows": combined_rows,
    }
    combined_path = out_dir / "combined_measured_candidates.json"
    combined_path.write_text(json.dumps(combined, indent=2, sort_keys=True) + "\n")
    metric_args = [
        "--input",
        str(combined_path),
        "--out-dir",
        str(out_dir / "measured_metrics"),
        "--mode",
        "measured",
        "--k",
        str(args.k),
        "--bootstrap-samples",
        str(args.bootstrap_samples),
    ]
    if args.no_markdown_report:
        metric_args.append("--no-markdown-report")
    eval_metrics_main(metric_args)
    metrics = json.loads((out_dir / "measured_metrics" / "metrics.json").read_text())
    (out_dir / "metrics.json").write_text(json.dumps(_summary_payload(combined, metrics), indent=2, sort_keys=True) + "\n")
    (out_dir / "table.tex").write_text((out_dir / "measured_metrics" / "table.tex").read_text())
    report_path = out_dir / "report.md"
    if args.no_markdown_report:
        report_path.unlink(missing_ok=True)
    else:
        report_path.write_text(_report(combined, metrics) + "\n")
    for filename in ("metrics_by_task.json", "metrics_by_seed.json"):
        shutil.copyfile(out_dir / "measured_metrics" / filename, out_dir / filename)
    (out_dir / "command.txt").write_text(
        "python scripts/build_ctt_rollout_comparison.py " + " ".join(sys.argv[1:]) + "\n"
    )
    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
    (out_dir / "data_hash.txt").write_text(str(combined.get("target_content_hash") or "") + "\n")
    (out_dir / "split_hash.txt").write_text(str(combined.get("target_split_hash") or "") + "\n")
    (out_dir / "source_data_hash.txt").write_text(str(combined.get("source_content_hash") or "") + "\n")
    (out_dir / "train.log").write_text("comparison artifact; source runs trained/evaluated separately\n")
    (out_dir / "eval.log").write_text(
        "\n".join(f"{run_dir}: measured_candidates.json" for run_dir in combined["run_dirs"]) + "\n"
    )
    print(json.dumps({"out_dir": str(out_dir), "runs": len(run_dirs), "rows": len(combined_rows)}, indent=2))
    return 0


def _summary_payload(combined: dict[str, Any], metrics: dict[str, Any]) -> dict[str, Any]:
    return {
        "report_type": "ctt_measured_rollout_comparison",
        "k": combined["k"],
        "run_dirs": combined["run_dirs"],
        "train_seeds": combined["train_seeds"],
        "num_rows": combined["num_rows"],
        "data_hash": combined.get("target_content_hash"),
        "split_hash": combined.get("target_split_hash"),
        "summary": metrics.get("summary", {}),
        "success_summary": _success_summary(combined.get("rows", []), k=int(combined["k"])),
        "source_content_hash": combined.get("source_content_hash"),
        "target_content_hash": combined.get("target_content_hash"),
        "target_split_hash": combined.get("target_split_hash"),
    }


def _report(combined: dict[str, Any], metrics: dict[str, Any]) -> str:
    summary = metrics.get("summary", {})
    split = _split_name(combined)
    lines = [
        f"# CTT {split.title()} Measured Rollout Comparison",
        "",
        f"Runs: `{len(combined['run_dirs'])}`",
        f"Rows: `{combined['num_rows']}`",
        f"K: `{combined['k']}`",
        "",
        "| Metric | N | Micro mean | 95% CI |",
        "| --- | ---: | ---: | ---: |",
    ]
    for name, payload in sorted(summary.items()):
        micro = payload.get("micro", {})
        lines.append(
            f"| {name} | {micro.get('n', 0)} | {_fmt(micro.get('mean'))} | "
            f"[{_fmt(micro.get('low'))}, {_fmt(micro.get('high'))}] |"
        )
    lines.append("")
    success = _success_summary(combined.get("rows", []), k=int(combined["k"]))
    lines.extend(
        [
            "| Success/Utility | Mean |",
            "| --- | ---: |",
            f"| base_success_rate | {_fmt(success.get('base_success_rate'))} |",
            f"| selected_success_rate | {_fmt(success.get('selected_success_rate'))} |",
            f"| proposal_oracle_success_rate | {_fmt(success.get('proposal_oracle_success_rate'))} |",
            f"| hidden_chart_oracle_success_rate | {_fmt(success.get('hidden_chart_oracle_success_rate'))} |",
            f"| selected_success_gain_over_base | {_fmt(success.get('selected_success_gain_over_base'))} |",
            f"| proposal_oracle_success_gain_over_base | {_fmt(success.get('proposal_oracle_success_gain_over_base'))} |",
            f"| success_support_gap | {_fmt(success.get('success_support_gap'))} |",
            f"| success_selector_gap | {_fmt(success.get('success_selector_gap'))} |",
            f"| base_utility_mean | {_fmt(success.get('base_utility_mean'))} |",
            f"| selected_utility_mean | {_fmt(success.get('selected_utility_mean'))} |",
            f"| proposal_oracle_utility_mean | {_fmt(success.get('proposal_oracle_utility_mean'))} |",
            f"| hidden_chart_oracle_utility_mean | {_fmt(success.get('hidden_chart_oracle_utility_mean'))} |",
            f"| ncar_to_proposal_oracle | {_fmt(success.get('ncar_to_proposal_oracle'))} |",
            f"| ncar_to_hidden_chart_oracle | {_fmt(success.get('ncar_to_hidden_chart_oracle'))} |",
            f"| utility_support_gap_fraction_to_hidden | {_fmt(success.get('utility_support_gap_fraction_to_hidden'))} |",
            f"| utility_selector_gap_fraction_to_hidden | {_fmt(success.get('utility_selector_gap_fraction_to_hidden'))} |",
            f"| generated_safety_label_coverage | {_fmt(success.get('generated_safety_label_coverage'))} |",
            f"| generated_unsafe_rate_known | {_fmt(success.get('generated_unsafe_rate_known'))} |",
            f"| any_generated_unsafe_known | {_fmt(success.get('any_generated_unsafe_known'))} |",
            f"| selected_safety_label_known_rate | {_fmt(success.get('selected_safety_label_known_rate'))} |",
            f"| selected_unsafe_rate_known | {_fmt(success.get('selected_unsafe_rate_known'))} |",
            f"| proposal_oracle_safety_label_known_rate | {_fmt(success.get('proposal_oracle_safety_label_known_rate'))} |",
            f"| proposal_oracle_unsafe_rate_known | {_fmt(success.get('proposal_oracle_unsafe_rate_known'))} |",
            f"| base_safety_label_known_rate | {_fmt(success.get('base_safety_label_known_rate'))} |",
            f"| base_unsafe_rate_known | {_fmt(success.get('base_unsafe_rate_known'))} |",
            "",
        ]
    )
    lines.append("These are measured generated-candidate rollouts, not PPTC proxies.")
    lines.append("")
    lines.append("Run dirs:")
    for run_dir in combined["run_dirs"]:
        lines.append(f"- `{run_dir}`")
    return "\n".join(lines)


def _seed_from_path(path: Path) -> str:
    match = re.search(r"seed(\d+)", path.name)
    return match.group(1) if match else path.name


def _split_name(combined: dict[str, Any]) -> str:
    text = " ".join(str(item) for item in combined.get("run_dirs", []))
    if "test" in text:
        return "test"
    if "val" in text or "validation" in text:
        return "validation"
    return "measured"


def _first(payloads: list[dict[str, Any]], key: str) -> Any:
    for payload in payloads:
        value = payload.get(key)
        if value:
            return value
    return None


def _success_summary(rows: list[dict[str, Any]], *, k: int) -> dict[str, Any]:
    base_success = []
    selected_success = []
    oracle_success = []
    base_utility = []
    selected_utility = []
    oracle_utility = []
    hidden_oracle_utility = []
    hidden_oracle_success = []
    ncar_to_proposal_oracle = []
    ncar_to_hidden_chart_oracle = []
    utility_support_gap_fraction = []
    utility_selector_gap_fraction = []
    success_support_gap = []
    success_selector_gap = []
    selected_success_gain = []
    proposal_oracle_success_gain = []
    restore_errors = []
    generated_safety_coverage = []
    generated_unsafe = []
    any_generated_unsafe = []
    selected_safety_known = []
    selected_unsafe = []
    proposal_oracle_safety_known = []
    proposal_oracle_unsafe = []
    base_safety_known = []
    base_unsafe = []
    for row in rows:
        generated_utilities = [float(value) for value in row.get("generated_utilities", [])[:k]]
        generated_success = [bool(value) for value in row.get("candidate_success", [])[:k]]
        candidate_outcomes = _outcome_list(row.get("candidate_outcomes", []))[:k]
        selected_index = int(row.get("selected_index", 0))
        selected_success_value: float | None = None
        proposal_oracle_success_value: float | None = None
        base_success_value: float | None = None
        base_outcome = row.get("base_outcome")
        if isinstance(base_outcome, dict):
            safety = outcome_safety_violation(base_outcome)
            base_safety_known.append(float(safety is not None))
            if safety is not None:
                base_unsafe.append(float(safety))
        if candidate_outcomes:
            generated_safety_coverage.append(safety_label_coverage(candidate_outcomes, k=k))
            unsafe = unsafe_rate(candidate_outcomes, k=k)
            if unsafe is not None:
                generated_unsafe.append(unsafe)
            any_unsafe_value = any_unsafe(candidate_outcomes, k=k)
            if any_unsafe_value is not None:
                any_generated_unsafe.append(any_unsafe_value)
            if selected_index < len(candidate_outcomes):
                safety = outcome_safety_violation(candidate_outcomes[selected_index])
                selected_safety_known.append(float(safety is not None))
                if safety is not None:
                    selected_unsafe.append(float(safety))
            if generated_utilities:
                oracle_index = max(
                    range(len(generated_utilities)),
                    key=lambda index: generated_utilities[index],
                )
                if oracle_index < len(candidate_outcomes):
                    safety = outcome_safety_violation(candidate_outcomes[oracle_index])
                    proposal_oracle_safety_known.append(float(safety is not None))
                    if safety is not None:
                        proposal_oracle_unsafe.append(float(safety))
        if "base_success" in row:
            base_success_value = float(bool(row["base_success"]))
            base_success.append(base_success_value)
        if selected_index < len(generated_success):
            selected_success_value = float(generated_success[selected_index])
            selected_success.append(selected_success_value)
        if generated_success:
            proposal_oracle_success_value = float(any(generated_success))
            oracle_success.append(proposal_oracle_success_value)
        if "base_utility" in row:
            base_utility_value = float(row["base_utility"])
            base_utility.append(base_utility_value)
        else:
            base_utility_value = None
        if selected_index < len(generated_utilities):
            selected_utility_value = generated_utilities[selected_index]
            selected_utility.append(selected_utility_value)
        else:
            selected_utility_value = None
        if generated_utilities:
            proposal_oracle_utility_value = max(generated_utilities)
            oracle_utility.append(proposal_oracle_utility_value)
        else:
            proposal_oracle_utility_value = None
        if (
            base_utility_value is not None
            and selected_utility_value is not None
            and proposal_oracle_utility_value is not None
        ):
            _append_stable_ncar(
                ncar_to_proposal_oracle,
                proposal_oracle_utility_value,
                selected_utility_value,
                base_utility_value,
            )
        hidden = [float(value) for value in row.get("hidden_chart_utilities", [])]
        if hidden:
            hidden_oracle_utility_value = max(hidden)
            hidden_oracle_utility.append(hidden_oracle_utility_value)
            if (
                base_utility_value is not None
                and selected_utility_value is not None
                and proposal_oracle_utility_value is not None
            ):
                _append_stable_ncar(
                    ncar_to_hidden_chart_oracle,
                    hidden_oracle_utility_value,
                    selected_utility_value,
                    base_utility_value,
                )
                hidden_gap = abs(hidden_oracle_utility_value - base_utility_value)
                if hidden_gap > 0.0:
                    utility_support_gap_fraction.append(
                        max(0.0, hidden_oracle_utility_value - proposal_oracle_utility_value)
                        / hidden_gap
                    )
                    utility_selector_gap_fraction.append(
                        max(0.0, proposal_oracle_utility_value - selected_utility_value)
                        / hidden_gap
                    )
            hidden_success_value = float(any(value >= 1.0 for value in hidden))
            hidden_oracle_success.append(hidden_success_value)
            if proposal_oracle_success_value is not None:
                success_support_gap.append(
                    max(0.0, hidden_success_value - proposal_oracle_success_value)
                )
            if (
                proposal_oracle_success_value is not None
                and selected_success_value is not None
            ):
                success_selector_gap.append(
                    max(0.0, proposal_oracle_success_value - selected_success_value)
                )
        if base_success_value is not None and selected_success_value is not None:
            selected_success_gain.append(selected_success_value - base_success_value)
        if base_success_value is not None and proposal_oracle_success_value is not None:
            proposal_oracle_success_gain.append(
                proposal_oracle_success_value - base_success_value
            )
        if "restore_error" in row:
            restore_errors.append(float(row["restore_error"]))
    return {
        "base_success_rate": _mean(base_success),
        "selected_success_rate": _mean(selected_success),
        "proposal_oracle_success_rate": _mean(oracle_success),
        "hidden_chart_oracle_success_rate": _mean(hidden_oracle_success),
        "selected_success_gain_over_base": _mean(selected_success_gain),
        "proposal_oracle_success_gain_over_base": _mean(proposal_oracle_success_gain),
        "success_support_gap": _mean(success_support_gap),
        "success_selector_gap": _mean(success_selector_gap),
        "base_utility_mean": _mean(base_utility),
        "selected_utility_mean": _mean(selected_utility),
        "proposal_oracle_utility_mean": _mean(oracle_utility),
        "hidden_chart_oracle_utility_mean": _mean(hidden_oracle_utility),
        "ncar_to_proposal_oracle": _mean(ncar_to_proposal_oracle),
        "ncar_to_hidden_chart_oracle": _mean(ncar_to_hidden_chart_oracle),
        "utility_support_gap_fraction_to_hidden": _mean(utility_support_gap_fraction),
        "utility_selector_gap_fraction_to_hidden": _mean(utility_selector_gap_fraction),
        "generated_safety_label_coverage": _mean(generated_safety_coverage),
        "generated_unsafe_rate_known": _mean(generated_unsafe),
        "any_generated_unsafe_known": _mean(any_generated_unsafe),
        "selected_safety_label_known_rate": _mean(selected_safety_known),
        "selected_unsafe_rate_known": _mean(selected_unsafe),
        "proposal_oracle_safety_label_known_rate": _mean(proposal_oracle_safety_known),
        "proposal_oracle_unsafe_rate_known": _mean(proposal_oracle_unsafe),
        "base_safety_label_known_rate": _mean(base_safety_known),
        "base_unsafe_rate_known": _mean(base_unsafe),
        "max_restore_error": max(restore_errors) if restore_errors else None,
    }


def _outcome_list(value: Any) -> list[dict[str, Any]]:
    if not isinstance(value, list):
        return []
    return [item for item in value if isinstance(item, dict)]


def _mean(values: list[float]) -> float | None:
    clean = [float(value) for value in values if math.isfinite(float(value))]
    return sum(clean) / len(clean) if clean else None


def _append_stable_ncar(
    output: list[float],
    oracle_utility: float,
    selected_utility: float,
    base_utility: float,
    *,
    min_denominator: float = 1.0e-3,
) -> None:
    if abs(float(oracle_utility) - float(base_utility)) <= min_denominator:
        return
    output.append(
        normalized_causal_action_regret(
            oracle_utility,
            selected_utility,
            base_utility,
        )
    )


def _fmt(value: Any) -> str:
    if not isinstance(value, (int, float)) or not math.isfinite(float(value)):
        return "n/a"
    return f"{float(value):.4f}"


def _run(command: list[str]) -> str:
    try:
        return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return ""


if __name__ == "__main__":
    raise SystemExit(main())