#!/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())