Add learned dominance selector diagnostic

workspace/scripts/eval_learned_dominance_selector.py

@@ -0,0 +1,542 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import re
+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 numpy as np  # noqa: E402
+
+from cil.metrics import macro_micro_summary  # noqa: E402
+from scripts.eval_dominance_selector import _DominanceScorer, _chart_map, _rows  # noqa: E402
+
+
+FEATURE_NAMES = [
+    "bias",
+    "candidate_score",
+    "candidate_score_minus_base_score",
+    "row_score_z",
+    "candidate_index",
+    "is_top_score_candidate",
+    "source_chart_rank",
+    "tangent_rms_norm",
+    "tangent_linf_norm",
+    "num_candidates",
+]
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Train a lightweight dominance calibrator on measured calibration "
+            "rollouts and evaluate deployment-clean fallback on held-out rollouts."
+        )
+    )
+    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",
+    )
+    parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_learned_dominance_val_to_test"))
+    parser.add_argument("--k", type=int, default=8)
+    parser.add_argument("--ridge-lambdas", default="0,0.01,0.1,1,10,100")
+    parser.add_argument("--bootstrap-samples", type=int, default=1000)
+    args = parser.parse_args(argv)
+
+    if args.k <= 0:
+        parser.error("--k must be positive")
+    lambdas = [float(item.strip()) for item in args.ridge_lambdas.split(",") if item.strip()]
+    if not lambdas or any(value < 0.0 for value in lambdas):
+        parser.error("--ridge-lambdas must contain non-negative values")
+
+    out_dir = args.out_dir
+    out_dir.mkdir(parents=True, exist_ok=True)
+    _write_provenance(out_dir, args)
+
+    scorer = _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_dataset = _candidate_dataset(
+        calibration_rows,
+        calibration_charts,
+        scorer=scorer,
+        k=args.k,
+    )
+    eval_dataset = _candidate_dataset(
+        eval_rows,
+        eval_charts,
+        scorer=scorer,
+        k=args.k,
+    )
+    best = _fit_select_ridge(calibration_dataset, lambdas=lambdas)
+    eval_cases = _evaluate_dataset(eval_dataset, best["weights"], best["mean"], best["std"], tau=best["tau"])
+    calibration_cases = _evaluate_dataset(
+        calibration_dataset,
+        best["weights"],
+        best["mean"],
+        best["std"],
+        tau=best["tau"],
+    )
+    metric_names = sorted(
+        {
+            key
+            for row in eval_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(
+            eval_cases,
+            name,
+            bootstrap_samples=args.bootstrap_samples,
+            confidence=0.95,
+        )
+        for name in metric_names
+    }
+    metrics = {
+        "report_type": "learned_dominance_selector_eval",
+        "schema_version": 1,
+        "k": args.k,
+        "feature_names": FEATURE_NAMES,
+        "ridge_lambdas": lambdas,
+        "selected_lambda": best["lambda"],
+        "tau": best["tau"],
+        "weights": best["weights"].tolist(),
+        "feature_mean": best["mean"].tolist(),
+        "feature_std": best["std"].tolist(),
+        "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_rows),
+        "num_eval_rows": len(eval_rows),
+        "num_calibration_candidates": len(calibration_dataset["samples"]),
+        "num_eval_candidates": len(eval_dataset["samples"]),
+        "calibration_model_selection": best["selection"],
+        "calibration_summary": _simple_summary(calibration_cases),
+        "eval_summary": _simple_summary(eval_cases),
+        "summary": summary,
+        "rows": eval_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(eval_cases, "task_id", metric_names), indent=2, sort_keys=True)
+        + "\n"
+    )
+    (out_dir / "metrics_by_seed.json").write_text(
+        json.dumps(_group_means(eval_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(
+        "trained ridge dominance calibrator on calibration measured rows only\n"
+        f"selected_lambda={best['lambda']}\n"
+        f"tau={best['tau']:.6f}\n"
+        f"calibration_input={args.calibration_input}\n"
+    )
+    (out_dir / "eval.log").write_text(
+        "evaluated learned fallback rule on held-out measured rollout rows\n"
+        f"eval_input={args.eval_input}\n"
+        f"num_eval_rows={len(eval_rows)}\n"
+    )
+    print(
+        json.dumps(
+            {
+                "out_dir": str(out_dir),
+                "lambda": best["lambda"],
+                "tau": best["tau"],
+                "selected_success": metrics["eval_summary"]["selected_success"],
+            },
+            indent=2,
+        )
+    )
+    return 0
+
+
+def _candidate_dataset(
+    rows: list[dict[str, Any]],
+    charts: dict[str, Any],
+    *,
+    scorer: _DominanceScorer,
+    k: int,
+) -> dict[str, Any]:
+    samples: list[dict[str, Any]] = []
+    by_row: dict[int, list[int]] = defaultdict(list)
+    for row_index, row in enumerate(rows):
+        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")
+        base_score = scorer.base_score(row, charts[chart_id])
+        scores = [float(value) for value in row.get("predicted_scores", [])[:k]]
+        utilities = [float(value) for value in row.get("generated_utilities", [])[:k]]
+        successes = [float(bool(value)) for value in row.get("candidate_success", [])[:k]]
+        tangents = row.get("generated_tangents", [])[:k]
+        candidate_types = row.get("candidate_types", [])[:k]
+        if not scores or len(scores) != len(utilities):
+            continue
+        score_mean = sum(scores) / len(scores)
+        score_std = math.sqrt(sum((score - score_mean) ** 2 for score in scores) / len(scores)) + 1.0e-6
+        for candidate_index, score in enumerate(scores):
+            tangent = np.asarray(
+                tangents[candidate_index] if candidate_index < len(tangents) else [],
+                dtype=float,
+            )
+            tangent_rms = float(np.linalg.norm(tangent) / math.sqrt(max(1, tangent.size)))
+            tangent_linf = float(np.max(np.abs(tangent))) if tangent.size else 0.0
+            feature = np.asarray(
+                [
+                    1.0,
+                    score,
+                    score - base_score,
+                    (score - score_mean) / score_std,
+                    float(candidate_index),
+                    float(candidate_index == 0),
+                    _source_rank(candidate_types[candidate_index] if candidate_index < len(candidate_types) else ""),
+                    tangent_rms,
+                    tangent_linf,
+                    float(len(scores)),
+                ],
+                dtype=float,
+            )
+            sample_index = len(samples)
+            by_row[row_index].append(sample_index)
+            base_utility = float(row["base_utility"])
+            base_success = float(bool(row.get("base_success", False)))
+            hidden = [float(value) for value in row.get("hidden_chart_utilities", [])]
+            samples.append(
+                {
+                    "row_index": row_index,
+                    "candidate_index": candidate_index,
+                    "feature": feature,
+                    "target_margin": utilities[candidate_index] - base_utility,
+                    "candidate_utility": utilities[candidate_index],
+                    "candidate_success": successes[candidate_index],
+                    "base_utility": base_utility,
+                    "base_success": base_success,
+                    "proposal_oracle_utility": max(utilities),
+                    "proposal_oracle_success": float(any(successes)),
+                    "hidden_chart_oracle_utility": max(hidden) if hidden else math.nan,
+                    "hidden_chart_oracle_success": float(any(value >= 1.0 for value in hidden)) if hidden else math.nan,
+                    "outcome_ptr": float(any(value > base_utility for value in utilities)),
+                    "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")),
+                }
+            )
+    return {"samples": samples, "by_row": dict(by_row), "num_rows": len(rows)}
+
+
+def _fit_select_ridge(dataset: dict[str, Any], *, lambdas: list[float]) -> dict[str, Any]:
+    samples = dataset["samples"]
+    if not samples:
+        raise ValueError("cannot fit learned dominance selector without candidates")
+    x = np.stack([sample["feature"] for sample in samples], axis=0)
+    y = np.asarray([float(sample["target_margin"]) for sample in samples], dtype=float)
+    mean = np.zeros(x.shape[1], dtype=float)
+    std = np.ones(x.shape[1], dtype=float)
+    mean[1:] = x[:, 1:].mean(axis=0)
+    std[1:] = x[:, 1:].std(axis=0) + 1.0e-6
+    x_norm = (x - mean) / std
+    x_norm[:, 0] = 1.0
+    best: dict[str, Any] | None = None
+    for ridge_lambda in lambdas:
+        penalty = ridge_lambda * np.eye(x_norm.shape[1], dtype=float)
+        penalty[0, 0] = 0.0
+        weights = np.linalg.pinv(x_norm.T @ x_norm + penalty) @ (x_norm.T @ y)
+        predictions = x_norm @ weights
+        tau, selection = _choose_tau(dataset, predictions)
+        key = (
+            float(selection["selected_success"]),
+            float(selection["selected_utility"]),
+            float(selection["coverage"]),
+            -float(ridge_lambda),
+        )
+        if best is None or key > best["key"]:
+            best = {
+                "key": key,
+                "lambda": ridge_lambda,
+                "weights": weights,
+                "mean": mean,
+                "std": std,
+                "tau": tau,
+                "selection": selection,
+            }
+    assert best is not None
+    return best
+
+
+def _choose_tau(dataset: dict[str, Any], predictions: np.ndarray) -> tuple[float, dict[str, float | None]]:
+    candidates = sorted(float(value) for value in predictions)
+    thresholds = [min(candidates) - 1.0, *candidates, max(candidates) + 1.0]
+    best_tau = thresholds[0]
+    best_summary: dict[str, float | None] | None = None
+    best_key: tuple[float, float, float] | None = None
+    for tau in thresholds:
+        cases = _evaluate_predictions(dataset, predictions, tau=tau)
+        summary = _simple_summary(cases)
+        key = (
+            float(summary.get("selected_success") or 0.0),
+            float(summary.get("selected_utility") or 0.0),
+            float(summary.get("coverage") or 0.0),
+        )
+        if best_key is None or key > best_key:
+            best_key = key
+            best_tau = float(tau)
+            best_summary = summary
+    assert best_summary is not None
+    return best_tau, best_summary
+
+
+def _evaluate_dataset(
+    dataset: dict[str, Any],
+    weights: np.ndarray,
+    mean: np.ndarray,
+    std: np.ndarray,
+    *,
+    tau: float,
+) -> list[dict[str, Any]]:
+    x = np.stack([sample["feature"] for sample in dataset["samples"]], axis=0)
+    x_norm = (x - mean) / std
+    x_norm[:, 0] = 1.0
+    return _evaluate_predictions(dataset, x_norm @ weights, tau=tau)
+
+
+def _evaluate_predictions(dataset: dict[str, Any], predictions: np.ndarray, *, tau: float) -> list[dict[str, Any]]:
+    samples = dataset["samples"]
+    rows: list[dict[str, Any]] = []
+    for row_index, sample_indices in sorted(dataset["by_row"].items()):
+        best_index = max(sample_indices, key=lambda index: float(predictions[index]))
+        sample = samples[best_index]
+        predicted_margin = float(predictions[best_index])
+        execute = predicted_margin > float(tau)
+        selected_utility = (
+            float(sample["candidate_utility"]) if execute else float(sample["base_utility"])
+        )
+        selected_success = (
+            float(sample["candidate_success"]) if execute else float(sample["base_success"])
+        )
+        hidden_utility = float(sample["hidden_chart_oracle_utility"])
+        hidden_success = float(sample["hidden_chart_oracle_success"])
+        rows.append(
+            {
+                "chart_id": sample["chart_id"],
+                "task_id": sample["task_id"],
+                "seed": sample["seed"],
+                "train_seed": sample["train_seed"],
+                "selected_candidate_index": int(sample["candidate_index"]),
+                "predicted_margin": predicted_margin,
+                "tau": float(tau),
+                "execute_generated": float(execute),
+                "coverage": float(execute),
+                "fallback_rate": float(not execute),
+                "base_utility": float(sample["base_utility"]),
+                "base_success": float(sample["base_success"]),
+                "selected_utility": selected_utility,
+                "selected_success": selected_success,
+                "selected_utility_gain_over_base": selected_utility - float(sample["base_utility"]),
+                "selected_success_gain_over_base": selected_success - float(sample["base_success"]),
+                "proposal_oracle_utility": float(sample["proposal_oracle_utility"]),
+                "proposal_oracle_success": float(sample["proposal_oracle_success"]),
+                "hidden_chart_oracle_utility": hidden_utility,
+                "hidden_chart_oracle_success": hidden_success,
+                "outcome_ptr": float(sample["outcome_ptr"]),
+                "selector_regret": max(0.0, float(sample["proposal_oracle_utility"]) - selected_utility),
+                "success_selector_gap": max(0.0, float(sample["proposal_oracle_success"]) - selected_success),
+                "support_gap": max(0.0, hidden_utility - float(sample["proposal_oracle_utility"]))
+                if math.isfinite(hidden_utility)
+                else math.nan,
+                "success_support_gap": max(0.0, hidden_success - float(sample["proposal_oracle_success"]))
+                if math.isfinite(hidden_success)
+                else math.nan,
+            }
+        )
+    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 _source_rank(value: Any) -> float:
+    match = re.search(r"rank(\d+)", str(value))
+    return float(match.group(1)) if match else 0.0
+
+
+def _table(metrics: dict[str, Any]) -> str:
+    summary = metrics["eval_summary"]
+    lines = [
+        "% Auto-generated by scripts/eval_learned_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 = [
+        "# Learned Dominance-Calibrated CTT Selector",
+        "",
+        f"Calibration rows: `{metrics['num_calibration_rows']}`",
+        f"Eval rows: `{metrics['num_eval_rows']}`",
+        f"Selected ridge lambda: `{metrics['selected_lambda']}`",
+        f"Tau: `{metrics['tau']:.6f}`",
+        "",
+        "The ridge calibrator and threshold are fit on calibration measured 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 selector diagnostic over already measured candidates, not a new rollout.",
+    ]
+    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_learned_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())