scripts/eval_dominance_selector.py

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

import argparse
import json
import math
import os
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

try:
    torch.set_num_threads(int(os.environ.get("DOVLA_TORCH_THREADS", "1")))
except (RuntimeError, ValueError):
    pass


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(
        "--score-source",
        choices=("row", "checkpoint"),
        default="row",
        help=(
            "Use row predicted_scores from the measured rollout, or recompute "
            "candidate scores from the checkpoint utility-energy model."
        ),
    )
    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, score_source=args.score_source)
    calibration_rows = _rows(json.loads(args.calibration_input.read_text()))
    eval_rows = _rows(json.loads(args.eval_input.read_text()))
    chart_feature_mode = calibrator.chart_feature_mode(_first_train_seed(calibration_rows + eval_rows))
    calibration_charts, calibration_index = _chart_map(
        args.calibration_target_index,
        chart_feature_mode=chart_feature_mode,
    )
    eval_charts, eval_index = _chart_map(
        args.eval_target_index,
        chart_feature_mode=chart_feature_mode,
    )

    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,
        "score_source": args.score_source,
        "chart_feature_mode": chart_feature_mode,
        "checkpoint_template": args.checkpoint_template,
        "residual_quantile": residual_quantile,
        "calibration_input": str(args.calibration_input),
        "eval_input": str(args.eval_input),
        "data_hash": eval_index.get("content_hash"),
        "split_hash": eval_index.get("split_hash"),
        "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, *, score_source: str = "row") -> None:
        if score_source not in {"row", "checkpoint"}:
            raise ValueError("score_source must be 'row' or 'checkpoint'")
        self.checkpoint_template = checkpoint_template
        self.score_source = score_source
        self._models: dict[str, tuple[ChartEncoder, UtilityEnergy, TangentNormalizer, int]] = {}
        self._feature_modes: dict[str, str] = {}
        self._encoded_chart_cache: dict[tuple[str, str], torch.Tensor] = {}
        self._base_score_cache: dict[tuple[str, str], float] = {}

    def chart_feature_mode(self, seed: str) -> str:
        # Row-scored rollouts may still need the checkpoint utility model to
        # compute the base-action score. Load once so chart maps use the same
        # feature dimensionality as that checkpoint (for example base_context_obs).
        self._model(seed)
        return self._feature_modes.get(seed, "base")

    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"))
        cache_key = (seed, str(chart.chart_id))
        if cache_key in self._base_score_cache:
            return self._base_score_cache[cache_key]
        _encoder, utility_energy, normalizer, tangent_dim = self._model(seed)
        z_chart = self._encoded_chart(seed, chart)
        with torch.inference_mode():
            zero = torch.zeros((1, tangent_dim), dtype=torch.float32)
            zero_norm = normalizer.transform(zero)
            score = float(utility_energy(z_chart, zero_norm).squeeze(0).item())
        self._base_score_cache[cache_key] = score
        return score

    def candidate_scores(self, row: dict[str, Any], chart: Any, *, k: int) -> list[float]:
        if self.score_source == "row":
            return [float(value) for value in row.get("predicted_scores", [])[:k]]
        seed = str(row.get("train_seed", "0"))
        encoder, utility_energy, normalizer, tangent_dim = self._model(seed)
        tangents = row.get("generated_tangents", [])[:k]
        if not tangents:
            return []
        z_chart = self._encoded_chart(seed, chart)
        with torch.inference_mode():
            xi = torch.as_tensor(tangents, dtype=torch.float32)
            if xi.ndim != 2:
                xi = xi.reshape(1, -1)
            if xi.shape[1] < tangent_dim:
                pad = torch.zeros((xi.shape[0], tangent_dim - xi.shape[1]), dtype=xi.dtype)
                xi = torch.cat([xi, pad], dim=1)
            elif xi.shape[1] > tangent_dim:
                xi = xi[:, :tangent_dim]
            xi_norm = normalizer.transform(xi)
            z = z_chart.repeat(xi_norm.shape[0], 1)
            return [float(value) for value in utility_energy(z, xi_norm).detach().cpu().tolist()]

    def _encoded_chart(self, seed: str, chart: Any) -> torch.Tensor:
        cache_key = (seed, str(chart.chart_id))
        if cache_key in self._encoded_chart_cache:
            return self._encoded_chart_cache[cache_key]
        encoder, _utility_energy, _normalizer, _tangent_dim = self._model(seed)
        with torch.inference_mode():
            feature = torch.as_tensor(chart.feature, dtype=torch.float32).unsqueeze(0)
            z_chart = encoder(feature)
        self._encoded_chart_cache[cache_key] = z_chart
        return z_chart

    def _model(self, seed: str) -> tuple[ChartEncoder, UtilityEnergy, TangentNormalizer, int]:
        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)
        if "config" in checkpoint:
            config = CTTConfig(**checkpoint["config"])
            chart_feature_dim = config.chart_feature_dim
            chart_dim = config.chart_dim
            tangent_dim = config.tangent_dim
        else:
            chart_feature_dim = int(checkpoint["feature_dim"])
            chart_dim = int(checkpoint.get("chart_dim", 64))
            tangent_dim = int(checkpoint["tangent_dim"])
        self._feature_modes[seed] = str(checkpoint.get("chart_feature_mode", "base"))
        encoder = ChartEncoder(chart_feature_dim, output_dim=chart_dim)
        utility_energy = UtilityEnergy(chart_dim=chart_dim, tangent_dim=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, tangent_dim)
        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]]
    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")
    predicted_scores = scorer.candidate_scores(row, charts[chart_id], k=k)
    if not generated_utilities or not predicted_scores:
        raise ValueError("dominance evaluation requires generated utilities and predicted scores")
    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, *, chart_feature_mode: str = "base") -> 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,
        chart_feature_mode=chart_feature_mode,
    )
    return {chart.chart_id: chart for chart in charts}, index


def _first_train_seed(rows: list[dict[str, Any]]) -> str:
    for row in rows:
        if row.get("train_seed") is not None:
            return str(row.get("train_seed"))
    return "0"


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())