Add dominance-calibrated CTT selector diagnostics

workspace/scripts/eval_dominance_selector.py

@@ -0,0 +1,493 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import subprocess
+import sys
+from collections import defaultdict
+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))
+
+import torch  # noqa: E402
+
+from cil.metrics import macro_micro_summary  # noqa: E402
+from cil.models import CTTConfig, ChartEncoder, TangentNormalizer, UtilityEnergy  # noqa: E402
+from scripts.eval_ctt_generated_rollout import load_chart_items  # noqa: E402
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Calibrate a causal-dominance fallback rule on measured generated "
+            "candidate rollouts and evaluate it on a held-out measured rollout set."
+        )
+    )
+    parser.add_argument("--calibration-input", type=Path, required=True)
+    parser.add_argument("--calibration-target-index", type=Path, required=True)
+    parser.add_argument("--eval-input", type=Path, required=True)
+    parser.add_argument("--eval-target-index", type=Path, required=True)
+    parser.add_argument(
+        "--checkpoint-template",
+        default="runs/ctt_residual_full_seed{seed}/model.pt",
+        help="Template used to load the train-seed utility-energy checkpoint.",
+    )
+    parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_dominance_val_to_test"))
+    parser.add_argument("--alpha", type=float, default=0.1)
+    parser.add_argument(
+        "--tau",
+        default="auto",
+        help=(
+            "Dominance threshold. Use a float, or 'auto' to choose the threshold "
+            "that maximizes selected success on the calibration split after "
+            "conformal residual subtraction."
+        ),
+    )
+    parser.add_argument("--k", type=int, default=8)
+    parser.add_argument("--bootstrap-samples", type=int, default=1000)
+    args = parser.parse_args(argv)
+
+    if not 0.0 < args.alpha < 1.0:
+        parser.error("--alpha must be in (0, 1)")
+    if args.k <= 0:
+        parser.error("--k must be positive")
+
+    out_dir = args.out_dir
+    out_dir.mkdir(parents=True, exist_ok=True)
+    _write_provenance(out_dir, args)
+
+    calibrator = _DominanceScorer(args.checkpoint_template)
+    calibration_rows = _rows(json.loads(args.calibration_input.read_text()))
+    eval_rows = _rows(json.loads(args.eval_input.read_text()))
+    calibration_charts, calibration_index = _chart_map(args.calibration_target_index)
+    eval_charts, eval_index = _chart_map(args.eval_target_index)
+
+    calibration_cases = [
+        _dominance_case(row, calibration_charts, scorer=calibrator, k=args.k)
+        for row in calibration_rows
+    ]
+    residual_quantile = _conformal_quantile(
+        [abs(case["measured_margin"] - case["predicted_margin"]) for case in calibration_cases],
+        alpha=args.alpha,
+    )
+    tau = (
+        _choose_tau(calibration_cases, residual_quantile=residual_quantile)
+        if args.tau == "auto"
+        else float(args.tau)
+    )
+
+    evaluated_cases = [
+        _evaluate_case(
+            _dominance_case(row, eval_charts, scorer=calibrator, k=args.k),
+            residual_quantile=residual_quantile,
+            tau=tau,
+        )
+        for row in eval_rows
+    ]
+    calibration_eval_cases = [
+        _evaluate_case(case, residual_quantile=residual_quantile, tau=tau)
+        for case in calibration_cases
+    ]
+    metric_names = sorted(
+        {
+            key
+            for row in evaluated_cases
+            for key, value in row.items()
+            if key not in {"chart_id", "task_id", "seed", "train_seed"}
+            and isinstance(value, (int, float))
+            and math.isfinite(float(value))
+        }
+    )
+    summary = {
+        name: macro_micro_summary(
+            evaluated_cases,
+            name,
+            bootstrap_samples=args.bootstrap_samples,
+            confidence=0.95,
+        )
+        for name in metric_names
+    }
+    metrics = {
+        "report_type": "dominance_calibrated_selector_eval",
+        "schema_version": 1,
+        "k": args.k,
+        "alpha": args.alpha,
+        "tau": tau,
+        "tau_mode": args.tau,
+        "residual_quantile": residual_quantile,
+        "calibration_input": str(args.calibration_input),
+        "eval_input": str(args.eval_input),
+        "calibration_target_content_hash": calibration_index.get("content_hash"),
+        "calibration_target_split_hash": calibration_index.get("split_hash"),
+        "eval_target_content_hash": eval_index.get("content_hash"),
+        "eval_target_split_hash": eval_index.get("split_hash"),
+        "num_calibration_rows": len(calibration_cases),
+        "num_eval_rows": len(evaluated_cases),
+        "calibration_summary": _simple_summary(calibration_eval_cases),
+        "eval_summary": _simple_summary(evaluated_cases),
+        "summary": summary,
+        "rows": evaluated_cases,
+    }
+    (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n")
+    (out_dir / "metrics_by_task.json").write_text(
+        json.dumps(_group_means(evaluated_cases, "task_id", metric_names), indent=2, sort_keys=True)
+        + "\n"
+    )
+    (out_dir / "metrics_by_seed.json").write_text(
+        json.dumps(_group_means(evaluated_cases, "seed", metric_names), indent=2, sort_keys=True)
+        + "\n"
+    )
+    (out_dir / "table.tex").write_text(_table(metrics) + "\n")
+    (out_dir / "report.md").write_text(_report(metrics) + "\n")
+    (out_dir / "train.log").write_text(
+        "fit conformal residual quantile and tau on calibration measured rows only\n"
+        f"calibration_input={args.calibration_input}\n"
+        f"residual_quantile={residual_quantile:.6f}\n"
+        f"tau={tau:.6f}\n"
+    )
+    (out_dir / "eval.log").write_text(
+        "evaluated calibrated fallback rule on held-out measured rollout rows\n"
+        f"eval_input={args.eval_input}\n"
+        f"num_eval_rows={len(evaluated_cases)}\n"
+    )
+    print(json.dumps({"out_dir": str(out_dir), "tau": tau, "rows": len(evaluated_cases)}, indent=2))
+    return 0
+
+
+class _DominanceScorer:
+    def __init__(self, checkpoint_template: str) -> None:
+        self.checkpoint_template = checkpoint_template
+        self._models: dict[str, tuple[ChartEncoder, UtilityEnergy, TangentNormalizer, CTTConfig]] = {}
+
+    def base_score(self, row: dict[str, Any], chart: Any) -> float:
+        if "base_predicted_score" in row:
+            return float(row["base_predicted_score"])
+        seed = str(row.get("train_seed", "0"))
+        encoder, utility_energy, normalizer, config = self._model(seed)
+        with torch.no_grad():
+            feature = torch.as_tensor(chart.feature, dtype=torch.float32).unsqueeze(0)
+            z_chart = encoder(feature)
+            zero = torch.zeros((1, config.tangent_dim), dtype=torch.float32)
+            zero_norm = normalizer.transform(zero)
+            return float(utility_energy(z_chart, zero_norm).squeeze(0).item())
+
+    def _model(self, seed: str) -> tuple[ChartEncoder, UtilityEnergy, TangentNormalizer, CTTConfig]:
+        if seed in self._models:
+            return self._models[seed]
+        checkpoint_path = Path(self.checkpoint_template.format(seed=seed))
+        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+        config = CTTConfig(**checkpoint["config"])
+        encoder = ChartEncoder(config.chart_feature_dim, output_dim=config.chart_dim)
+        utility_energy = UtilityEnergy(chart_dim=config.chart_dim, tangent_dim=config.tangent_dim)
+        encoder.load_state_dict(checkpoint["chart_encoder"])
+        utility_energy.load_state_dict(checkpoint["utility_energy"])
+        normalizer = TangentNormalizer.from_dict(checkpoint["normalizer"])
+        encoder.eval()
+        utility_energy.eval()
+        self._models[seed] = (encoder, utility_energy, normalizer, config)
+        return self._models[seed]
+
+
+def _dominance_case(
+    row: dict[str, Any],
+    charts: dict[str, Any],
+    *,
+    scorer: _DominanceScorer,
+    k: int,
+) -> dict[str, Any]:
+    generated_utilities = [float(value) for value in row.get("generated_utilities", [])[:k]]
+    candidate_success = [float(bool(value)) for value in row.get("candidate_success", [])[:k]]
+    predicted_scores = [float(value) for value in row.get("predicted_scores", [])[:k]]
+    if not generated_utilities or not predicted_scores:
+        raise ValueError("dominance evaluation requires generated utilities and predicted scores")
+    chart_id = str(row.get("chart_id", row.get("group_id", "")))
+    if chart_id not in charts:
+        raise KeyError(f"chart_id {chart_id!r} not found in target index")
+    top_index = max(range(len(predicted_scores)), key=lambda index: predicted_scores[index])
+    base_score = scorer.base_score(row, charts[chart_id])
+    base_utility = float(row["base_utility"])
+    base_success = float(bool(row.get("base_success", False)))
+    selected_generated_utility = generated_utilities[top_index]
+    selected_generated_success = candidate_success[top_index]
+    proposal_oracle_utility = max(generated_utilities)
+    proposal_oracle_success = float(any(candidate_success))
+    hidden = [float(value) for value in row.get("hidden_chart_utilities", [])]
+    hidden_oracle_utility = max(hidden) if hidden else math.nan
+    hidden_oracle_success = float(any(value >= 1.0 for value in hidden)) if hidden else math.nan
+    predicted_margin = predicted_scores[top_index] - base_score
+    measured_margin = selected_generated_utility - base_utility
+    return {
+        "chart_id": chart_id,
+        "task_id": str(row.get("task_id", "unknown")),
+        "seed": str(row.get("seed", "unknown")),
+        "train_seed": str(row.get("train_seed", "unknown")),
+        "top_index": top_index,
+        "base_predicted_score": base_score,
+        "top_predicted_score": predicted_scores[top_index],
+        "predicted_margin": predicted_margin,
+        "measured_margin": measured_margin,
+        "base_utility": base_utility,
+        "base_success": base_success,
+        "top_generated_utility": selected_generated_utility,
+        "top_generated_success": selected_generated_success,
+        "proposal_oracle_utility": proposal_oracle_utility,
+        "proposal_oracle_success": proposal_oracle_success,
+        "hidden_chart_oracle_utility": hidden_oracle_utility,
+        "hidden_chart_oracle_success": hidden_oracle_success,
+        "outcome_ptr": float(any(value > base_utility for value in generated_utilities)),
+    }
+
+
+def _evaluate_case(case: dict[str, Any], *, residual_quantile: float, tau: float) -> dict[str, Any]:
+    lcb = float(case["predicted_margin"]) - float(residual_quantile)
+    execute_generated = lcb > float(tau)
+    selected_utility = (
+        float(case["top_generated_utility"]) if execute_generated else float(case["base_utility"])
+    )
+    selected_success = (
+        float(case["top_generated_success"]) if execute_generated else float(case["base_success"])
+    )
+    proposal_oracle_utility = float(case["proposal_oracle_utility"])
+    proposal_oracle_success = float(case["proposal_oracle_success"])
+    hidden_utility = float(case["hidden_chart_oracle_utility"])
+    hidden_success = float(case["hidden_chart_oracle_success"])
+    output = dict(case)
+    output.update(
+        {
+            "lcb_margin": lcb,
+            "execute_generated": float(execute_generated),
+            "fallback_to_base": float(not execute_generated),
+            "coverage": float(execute_generated),
+            "fallback_rate": float(not execute_generated),
+            "selected_utility": selected_utility,
+            "selected_success": selected_success,
+            "selected_utility_gain_over_base": selected_utility - float(case["base_utility"]),
+            "selected_success_gain_over_base": selected_success - float(case["base_success"]),
+            "selector_regret": max(0.0, proposal_oracle_utility - selected_utility),
+            "branch_car": max(0.0, proposal_oracle_utility - selected_utility),
+            "success_selector_gap": max(0.0, proposal_oracle_success - selected_success),
+            "support_gap": max(0.0, hidden_utility - proposal_oracle_utility)
+            if math.isfinite(hidden_utility)
+            else math.nan,
+            "success_support_gap": max(0.0, hidden_success - proposal_oracle_success)
+            if math.isfinite(hidden_success)
+            else math.nan,
+            "success_total_car_to_hidden": max(0.0, hidden_success - selected_success)
+            if math.isfinite(hidden_success)
+            else math.nan,
+        }
+    )
+    return output
+
+
+def _choose_tau(cases: list[dict[str, Any]], *, residual_quantile: float) -> float:
+    candidates = sorted({float(case["predicted_margin"]) - float(residual_quantile) for case in cases})
+    thresholds = [min(candidates, default=0.0) - 1.0, *candidates, max(candidates, default=0.0) + 1.0]
+    best_tau = thresholds[0]
+    best_key: tuple[float, float, float] | None = None
+    for tau in thresholds:
+        evaluated = [_evaluate_case(case, residual_quantile=residual_quantile, tau=tau) for case in cases]
+        summary = _simple_summary(evaluated)
+        # Maximize selected success, then selected utility, then coverage.
+        key = (
+            float(summary.get("selected_success", 0.0) or 0.0),
+            float(summary.get("selected_utility", 0.0) or 0.0),
+            float(summary.get("coverage", 0.0) or 0.0),
+        )
+        if best_key is None or key > best_key:
+            best_key = key
+            best_tau = tau
+    return float(best_tau)
+
+
+def _conformal_quantile(values: list[float], *, alpha: float) -> float:
+    clean = sorted(float(value) for value in values if math.isfinite(float(value)))
+    if not clean:
+        raise ValueError("cannot calibrate dominance without residuals")
+    index = min(len(clean) - 1, max(0, math.ceil((1.0 - alpha) * (len(clean) + 1)) - 1))
+    return clean[index]
+
+
+def _chart_map(index_path: Path) -> tuple[dict[str, Any], dict[str, Any]]:
+    charts, index = load_chart_items(
+        index_path,
+        max_charts=None,
+        require_positive=True,
+        include_hidden=True,
+        include_metadata=True,
+    )
+    return {chart.chart_id: chart for chart in charts}, index
+
+
+def _rows(payload: Any) -> list[dict[str, Any]]:
+    rows = payload.get("rows", payload) if isinstance(payload, dict) else payload
+    if not isinstance(rows, list):
+        raise SystemExit("input must be a JSON list or object with rows")
+    return rows
+
+
+def _simple_summary(rows: list[dict[str, Any]]) -> dict[str, float | None]:
+    keys = [
+        "base_success",
+        "selected_success",
+        "proposal_oracle_success",
+        "hidden_chart_oracle_success",
+        "selected_success_gain_over_base",
+        "coverage",
+        "fallback_rate",
+        "outcome_ptr",
+        "success_support_gap",
+        "success_selector_gap",
+        "base_utility",
+        "selected_utility",
+        "proposal_oracle_utility",
+        "hidden_chart_oracle_utility",
+        "support_gap",
+        "selector_regret",
+    ]
+    return {key: _mean([row.get(key) for row in rows]) for key in keys}
+
+
+def _group_means(
+    rows: list[dict[str, Any]],
+    key: str,
+    metric_names: list[str],
+) -> dict[str, dict[str, float]]:
+    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for row in rows:
+        grouped[str(row.get(key, "unknown"))].append(row)
+    output: dict[str, dict[str, float]] = {}
+    for group, group_rows in sorted(grouped.items()):
+        payload: dict[str, float] = {}
+        for metric in metric_names:
+            value = _mean([row.get(metric) for row in group_rows])
+            if value is not None:
+                payload[metric] = value
+        output[group] = payload
+    return output
+
+
+def _mean(values: list[Any]) -> float | None:
+    clean = [
+        float(value)
+        for value in values
+        if isinstance(value, (int, float)) and math.isfinite(float(value))
+    ]
+    return sum(clean) / len(clean) if clean else None
+
+
+def _table(metrics: dict[str, Any]) -> str:
+    summary = metrics["eval_summary"]
+    lines = [
+        "% Auto-generated by scripts/eval_dominance_selector.py",
+        "\\begin{tabular}{lrrrrrrrr}",
+        "\\toprule",
+        "Rows & Coverage & Fallback & Base succ. & Selected succ. & Oracle succ. & OutcomePTR & Succ. support gap & Succ. selector gap \\\\",
+        "\\midrule",
+        f"{metrics['num_eval_rows']} & {_fmt(summary.get('coverage'))} & "
+        f"{_fmt(summary.get('fallback_rate'))} & {_fmt(summary.get('base_success'))} & "
+        f"{_fmt(summary.get('selected_success'))} & {_fmt(summary.get('proposal_oracle_success'))} & "
+        f"{_fmt(summary.get('outcome_ptr'))} & {_fmt(summary.get('success_support_gap'))} & "
+        f"{_fmt(summary.get('success_selector_gap'))} \\\\",
+        "\\bottomrule",
+        "\\end{tabular}",
+    ]
+    return "\n".join(lines)
+
+
+def _report(metrics: dict[str, Any]) -> str:
+    summary = metrics["eval_summary"]
+    calibration = metrics["calibration_summary"]
+    lines = [
+        "# Dominance-Calibrated CTT Selector",
+        "",
+        f"Calibration rows: `{metrics['num_calibration_rows']}`",
+        f"Eval rows: `{metrics['num_eval_rows']}`",
+        f"Alpha: `{metrics['alpha']}`",
+        f"Residual quantile: `{metrics['residual_quantile']:.6f}`",
+        f"Tau: `{metrics['tau']:.6f}` (`{metrics['tau_mode']}`)",
+        "",
+        "The threshold is fit on calibration rows only. Eval outcomes are used only for reporting.",
+        "",
+        "| Split | Coverage | Fallback | Base success | Selected success | Proposal oracle | OutcomePTR | Success support gap | Success selector gap |",
+        "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |",
+        f"| calibration | {_fmt(calibration.get('coverage'))} | {_fmt(calibration.get('fallback_rate'))} | "
+        f"{_fmt(calibration.get('base_success'))} | {_fmt(calibration.get('selected_success'))} | "
+        f"{_fmt(calibration.get('proposal_oracle_success'))} | {_fmt(calibration.get('outcome_ptr'))} | "
+        f"{_fmt(calibration.get('success_support_gap'))} | {_fmt(calibration.get('success_selector_gap'))} |",
+        f"| eval | {_fmt(summary.get('coverage'))} | {_fmt(summary.get('fallback_rate'))} | "
+        f"{_fmt(summary.get('base_success'))} | {_fmt(summary.get('selected_success'))} | "
+        f"{_fmt(summary.get('proposal_oracle_success'))} | {_fmt(summary.get('outcome_ptr'))} | "
+        f"{_fmt(summary.get('success_support_gap'))} | {_fmt(summary.get('success_selector_gap'))} |",
+        "",
+        "This is a calibrated fallback diagnostic. It is not a final safety claim because unsafe-contact labels are not measured yet.",
+    ]
+    return "\n".join(lines)
+
+
+def _write_provenance(out_dir: Path, args: argparse.Namespace) -> None:
+    (out_dir / "config.yaml").write_text(
+        "\n".join(f"{key}: {value}" for key, value in sorted(vars(args).items())) + "\n"
+    )
+    (out_dir / "command.txt").write_text(
+        "python scripts/eval_dominance_selector.py " + " ".join(sys.argv[1:]) + "\n"
+    )
+    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
+    hashes = {
+        "calibration_input": _sha256(args.calibration_input),
+        "calibration_target_index": _sha256(args.calibration_target_index),
+        "eval_input": _sha256(args.eval_input),
+        "eval_target_index": _sha256(args.eval_target_index),
+    }
+    (out_dir / "data_hash.txt").write_text(json.dumps(hashes, indent=2, sort_keys=True) + "\n")
+    (out_dir / "split_hash.txt").write_text(
+        json.dumps(
+            {
+                "calibration_target_index": _index_hash(args.calibration_target_index),
+                "eval_target_index": _index_hash(args.eval_target_index),
+            },
+            indent=2,
+            sort_keys=True,
+        )
+        + "\n"
+    )
+
+
+def _index_hash(path: Path) -> dict[str, Any]:
+    payload = json.loads(path.read_text())
+    return {
+        "split": payload.get("split"),
+        "content_hash": payload.get("content_hash"),
+        "split_hash": payload.get("split_hash"),
+        "retrieval_index_allowed": payload.get("retrieval_index_allowed"),
+    }
+
+
+def _sha256(path: Path) -> str:
+    import hashlib
+
+    h = hashlib.sha256()
+    h.update(path.read_bytes())
+    return h.hexdigest()
+
+
+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())