#!/usr/bin/env python from __future__ import annotations import argparse import hashlib 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.chart_features import CHART_FEATURE_MODES, OBJECT_LAYOUT_EMBED_DIM, OBSERVATION_EMBED_DIM # noqa: E402 from cil.metrics import macro_micro_summary, pairwise_causal_dominance_ece # noqa: E402 from scripts.eval_dominance_selector import _DominanceScorer, _chart_map, _first_train_seed, _rows # noqa: E402 from scripts.eval_ctt_generated_rollout import load_chart_items # noqa: E402 BASIC_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", ] FEATURE_NAMES = BASIC_FEATURE_NAMES CONTEXT_HASH_WIDTH = 8 SOURCE_EVIDENCE_NAMES = [ "source_chart_found", "source_positive_count_log", "source_negative_count_log", "source_nonbase_count_log", "source_positive_rate", "source_best_delta_utility", "source_mean_delta_utility", "source_best_utility", "source_mean_utility", "source_best_success", "source_mean_success", "source_best_progress", "source_mean_progress", "source_positive_safety_known_rate", "source_positive_unsafe_rate_known", "generated_to_source_positive_min_rms", "generated_to_source_negative_min_rms", "generated_source_pos_closer_than_neg", ] CHART_COMPAT_NAMES = [ "source_chart_feature_found", "target_chart_feature_norm", "source_chart_feature_norm", "target_source_chart_cosine", "target_source_chart_rms", "target_source_chart_linf", "target_source_chart_absmean", "target_source_chart_dot_per_dim", "target_base_action_norm", "source_base_action_norm", "target_source_base_cosine", "target_source_base_rms", "target_source_obs_available", "target_source_obs_cosine", "target_source_obs_rms", "target_obs_norm", "source_obs_norm", "target_source_obj_available", "target_source_obj_cosine", "target_source_obj_rms", "target_obj_norm", "source_obj_norm", ] SCORE_SHAPE_NAMES = [ "candidate_score_rank_fraction", "candidate_score_softmax_prob", "candidate_score_gap_to_best", "candidate_score_gap_to_second_best", "candidate_score_gap_to_prev_higher", "candidate_score_gap_to_next_lower", "candidate_score_percentile", "candidate_score_top_margin", ] BUNDLE_CONSENSUS_NAMES = [ "bundle_num_candidates_log", "bundle_neighbor_count_r020", "bundle_neighbor_count_r040", "bundle_neighbor_frac_r020", "bundle_neighbor_frac_r040", "bundle_mean_peer_rms", "bundle_min_peer_rms", "bundle_medoid_rms", "bundle_is_medoid", "bundle_peer_score_mean_r020", "bundle_peer_score_max_r020", "bundle_peer_score_mean_r040", "bundle_peer_score_max_r040", "bundle_unique_source_count_r040", "bundle_unique_source_frac_r040", "bundle_unique_task_count_r040", "bundle_same_source_frac_r040", "bundle_score_density_r040", "bundle_rank_density_r040", ] FEATURE_SET_CHOICES = ( "basic", "tangent", "context", "context_tangent", "score_context", "bundle_consensus", "score_bundle_consensus", "context_bundle_consensus", "context_tangent_bundle_consensus", "source_evidence", "tangent_source_evidence", "context_source_evidence", "context_tangent_source_evidence", "chart_compat", "chart_tangent_compat", "score_chart_compat", "score_context_chart_compat", "chart_bundle_consensus", "score_chart_bundle_consensus", "chart_source_compat", "chart_tangent_source_compat", "chart_source_bundle_consensus", ) 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( "--source-index", type=Path, default=None, help=( "Train split chart index used for source-evidence features. " "Defaults to --calibration-target-index." ), ) 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( "--feature-set", choices=FEATURE_SET_CHOICES, default="basic", help="Deployment-visible feature family for candidate-level dominance fitting.", ) parser.add_argument( "--selector-chart-feature-mode", choices=CHART_FEATURE_MODES, default="base_context_obs_obj", help=( "Chart feature mode used only for selector chart-compatibility " "features. These maps are loaded without hidden outcomes." ), ) parser.add_argument( "--target", choices=("utility_margin", "success", "success_weighted_margin"), default="utility_margin", help=( "Calibration target. success_weighted_margin fits utility margin plus " "candidate success to prioritize the lexicographic success/progress utility." ), ) parser.add_argument( "--success-bonus", type=float, default=1.0, help=( "Bonus multiplier for candidate_success when --target=success_weighted_margin. " "Default preserves the original +1 success bonus." ), ) parser.add_argument( "--threshold-scope", choices=("global", "task"), default="global", help=( "Calibrate one global execute/fallback threshold or a Mondrian " "threshold per visible task_id bucket using calibration rows only." ), ) parser.add_argument( "--fit-objective", choices=("pointwise", "pairwise", "hybrid_pairwise"), default="pointwise", help=( "Fit the linear utility proxy from candidate-level margins, " "within-chart pairwise causal comparisons, or both. Pairwise " "comparisons are built from calibration rows only." ), ) parser.add_argument( "--pairwise-weight", type=float, default=1.0, help="Relative weight for pairwise rows when --fit-objective=hybrid_pairwise.", ) parser.add_argument("--bootstrap-samples", type=int, default=1000) parser.add_argument( "--no-markdown-report", action="store_true", help="Do not write report.md; useful when the workspace is kept README-only.", ) 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") if args.pairwise_weight <= 0.0: parser.error("--pairwise-weight must be positive") if args.success_bonus < 0.0: parser.error("--success-bonus must be non-negative") 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())) chart_feature_mode = scorer.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, ) source_index_path = _resolve_index_path(args.source_index or args.calibration_target_index) source_evidence, source_index = ( _source_evidence_map(source_index_path) if _uses_source_evidence(args.feature_set) else ({}, {}) ) selector_source_charts: dict[str, Any] = {} selector_source_index: dict[str, Any] = {} selector_calibration_charts: dict[str, Any] = {} selector_eval_charts: dict[str, Any] = {} selector_calibration_index: dict[str, Any] = {} selector_eval_index: dict[str, Any] = {} if _uses_chart_compat(args.feature_set): selector_calibration_charts, selector_calibration_index = _selector_chart_map( args.calibration_target_index, chart_feature_mode=args.selector_chart_feature_mode, ) selector_eval_charts, selector_eval_index = _selector_chart_map( args.eval_target_index, chart_feature_mode=args.selector_chart_feature_mode, ) selector_source_charts, selector_source_index = _selector_chart_map( source_index_path, chart_feature_mode=args.selector_chart_feature_mode, ) calibration_dataset = _candidate_dataset( calibration_rows, calibration_charts, scorer=scorer, k=args.k, feature_set=args.feature_set, target=args.target, source_evidence=source_evidence, selector_target_charts=selector_calibration_charts, selector_source_charts=selector_source_charts, selector_chart_feature_mode=args.selector_chart_feature_mode, success_bonus=args.success_bonus, ) eval_dataset = _candidate_dataset( eval_rows, eval_charts, scorer=scorer, k=args.k, feature_set=args.feature_set, target=args.target, source_evidence=source_evidence, selector_target_charts=selector_eval_charts, selector_source_charts=selector_source_charts, selector_chart_feature_mode=args.selector_chart_feature_mode, success_bonus=args.success_bonus, ) best = _fit_select_ridge( calibration_dataset, lambdas=lambdas, threshold_scope=args.threshold_scope, fit_objective=args.fit_objective, pairwise_weight=args.pairwise_weight, ) eval_predictions = _linear_predictions(eval_dataset, best["weights"], best["mean"], best["std"]) calibration_predictions = _linear_predictions( calibration_dataset, best["weights"], best["mean"], best["std"], ) eval_pairwise = _pairwise_calibration_summary(eval_dataset, eval_predictions) calibration_pairwise = _pairwise_calibration_summary(calibration_dataset, calibration_predictions) eval_cases = _evaluate_predictions( eval_dataset, eval_predictions, tau=best["tau"], include_pairwise_calibration=True, pairwise_calibration=eval_pairwise, ) calibration_cases = _evaluate_predictions( calibration_dataset, calibration_predictions, tau=best["tau"], include_pairwise_calibration=True, pairwise_calibration=calibration_pairwise, ) 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_set": args.feature_set, "selector_chart_feature_mode": ( args.selector_chart_feature_mode if _uses_chart_compat(args.feature_set) else None ), "target": args.target, "success_bonus": args.success_bonus, "fit_objective": args.fit_objective, "pairwise_weight": args.pairwise_weight, "threshold_scope": args.threshold_scope, "chart_feature_mode": chart_feature_mode, "feature_names": _feature_names(args.feature_set), "ridge_lambdas": lambdas, "selected_lambda": best["lambda"], "tau": best["tau"], "fit_design": best["fit_design"], "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), "source_index": str(source_index_path) if _uses_source_evidence(args.feature_set) else None, "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"), "source_content_hash": source_index.get("content_hash"), "source_split_hash": source_index.get("split_hash"), "selector_source_content_hash": selector_source_index.get("content_hash"), "selector_source_split_hash": selector_source_index.get("split_hash"), "selector_calibration_target_content_hash": selector_calibration_index.get("content_hash"), "selector_calibration_target_split_hash": selector_calibration_index.get("split_hash"), "selector_eval_target_content_hash": selector_eval_index.get("content_hash"), "selector_eval_target_split_hash": selector_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": _summary_with_pairwise(calibration_cases, calibration_pairwise), "eval_summary": _summary_with_pairwise(eval_cases, eval_pairwise), "pairwise_causal_calibration": { "calibration": _pairwise_calibration_global(calibration_pairwise), "eval": _pairwise_calibration_global(eval_pairwise), }, "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") _write_report_artifact(out_dir, metrics, no_markdown_report=args.no_markdown_report) (out_dir / "train.log").write_text( "trained ridge dominance calibrator on calibration measured rows only\n" f"fit_objective={args.fit_objective}\n" f"selected_lambda={best['lambda']}\n" f"tau={_format_tau(best['tau'])}\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, feature_set: str = "basic", target: str = "utility_margin", source_evidence: dict[str, dict[str, Any]] | None = None, selector_target_charts: dict[str, Any] | None = None, selector_source_charts: dict[str, Any] | None = None, selector_chart_feature_mode: str = "base_context_obs_obj", success_bonus: float = 1.0, ) -> dict[str, Any]: source_evidence = source_evidence or {} selector_target_charts = selector_target_charts or {} selector_source_charts = selector_source_charts or {} 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") selector_target_chart = selector_target_charts.get(chart_id) 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] source_task_ids = row.get("candidate_source_task_ids", [])[:k] source_chart_ids = row.get("candidate_source_chart_ids", [])[: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 score_shape = _score_shape_matrix(scores) bundle_consensus = _bundle_consensus_matrix(tangents, scores, source_chart_ids, source_task_ids) for candidate_index, score in enumerate(scores): source_chart_id = str(source_chart_ids[candidate_index]) if candidate_index < len(source_chart_ids) else "" tangent = np.asarray( tangents[candidate_index] if candidate_index < len(tangents) else [], dtype=float, ) feature = _candidate_feature( score=score, base_score=base_score, score_mean=score_mean, score_std=score_std, candidate_index=candidate_index, candidate_type=candidate_types[candidate_index] if candidate_index < len(candidate_types) else "", context={ "target_task_id": row.get("task_id", ""), "instruction": row.get("instruction", ""), "source_task_id": source_task_ids[candidate_index] if candidate_index < len(source_task_ids) else "", }, tangent=tangent, source_evidence=_source_evidence_feature( source_evidence.get( source_chart_id ), tangent=tangent, ), chart_compat=_chart_compat_feature( selector_target_chart, selector_source_charts.get(source_chart_id), chart_feature_mode=selector_chart_feature_mode, ), score_shape=score_shape[candidate_index], bundle_consensus=bundle_consensus[candidate_index], num_candidates=len(scores), feature_set=feature_set, ) 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", [])] target_margin = utilities[candidate_index] - base_utility samples.append( { "row_index": row_index, "candidate_index": candidate_index, "feature": feature, "target_margin": _target_value( target, utility_margin=target_margin, candidate_success=successes[candidate_index], success_bonus=success_bonus, ), "measured_utility_margin": target_margin, "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 _feature_names(feature_set: str) -> list[str]: if feature_set == "basic": return list(BASIC_FEATURE_NAMES) context_names = [ *[f"target_task_hash_{index:02d}" for index in range(CONTEXT_HASH_WIDTH)], *[f"source_task_hash_{index:02d}" for index in range(CONTEXT_HASH_WIDTH)], *[f"instruction_hash_{index:02d}" for index in range(CONTEXT_HASH_WIDTH)], "source_target_same_task", "instruction_chars_per_128", "instruction_words_per_32", ] tangent_names = [ *[f"tangent_{index:02d}" for index in range(21)], *[f"abs_tangent_{index:02d}" for index in range(21)], ] names = list(BASIC_FEATURE_NAMES) if _uses_score_shape(feature_set): names.extend(SCORE_SHAPE_NAMES) if _uses_context(feature_set): names.extend(context_names) if _uses_tangent(feature_set): names.extend(tangent_names) if _uses_source_evidence(feature_set): names.extend(SOURCE_EVIDENCE_NAMES) if _uses_chart_compat(feature_set): names.extend(CHART_COMPAT_NAMES) if _uses_bundle_consensus(feature_set): names.extend(BUNDLE_CONSENSUS_NAMES) if feature_set in FEATURE_SET_CHOICES: return names raise ValueError(f"unknown feature_set: {feature_set}") def _candidate_feature( *, score: float, base_score: float, score_mean: float, score_std: float, candidate_index: int, candidate_type: Any, tangent: np.ndarray, num_candidates: int, feature_set: str, context: dict[str, Any] | None = None, source_evidence: np.ndarray | None = None, chart_compat: np.ndarray | None = None, score_shape: np.ndarray | None = None, bundle_consensus: np.ndarray | None = None, ) -> np.ndarray: tangent = np.asarray(tangent, dtype=float).reshape(-1) if tangent.size < 21: tangent = np.pad(tangent, (0, 21 - tangent.size)) elif tangent.size > 21: tangent = tangent[:21] 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 basic = np.asarray( [ 1.0, float(score), float(score) - float(base_score), (float(score) - float(score_mean)) / float(score_std), float(candidate_index), float(candidate_index == 0), _source_rank(candidate_type), tangent_rms, tangent_linf, float(num_candidates), ], dtype=float, ) if feature_set == "basic": return basic parts = [basic] if _uses_score_shape(feature_set): if score_shape is None: score_shape = np.zeros(len(SCORE_SHAPE_NAMES), dtype=float) parts.append(np.asarray(score_shape, dtype=float).reshape(-1)) if _uses_context(feature_set): parts.append(_context_feature(context or {})) if _uses_tangent(feature_set): parts.extend([tangent.astype(float), np.abs(tangent).astype(float)]) if _uses_source_evidence(feature_set): if source_evidence is None: source_evidence = np.zeros(len(SOURCE_EVIDENCE_NAMES), dtype=float) parts.append(np.asarray(source_evidence, dtype=float).reshape(-1)) if _uses_chart_compat(feature_set): if chart_compat is None: chart_compat = np.zeros(len(CHART_COMPAT_NAMES), dtype=float) parts.append(np.asarray(chart_compat, dtype=float).reshape(-1)) if _uses_bundle_consensus(feature_set): if bundle_consensus is None: bundle_consensus = np.zeros(len(BUNDLE_CONSENSUS_NAMES), dtype=float) parts.append(np.asarray(bundle_consensus, dtype=float).reshape(-1)) if feature_set in FEATURE_SET_CHOICES: return np.concatenate(parts) raise ValueError(f"unknown feature_set: {feature_set}") def _uses_context(feature_set: str) -> bool: return feature_set in { "context", "context_tangent", "context_source_evidence", "context_tangent_source_evidence", "score_context", "score_context_chart_compat", "context_bundle_consensus", "context_tangent_bundle_consensus", } def _uses_score_shape(feature_set: str) -> bool: return feature_set in { "score_context", "score_chart_compat", "score_context_chart_compat", "score_bundle_consensus", "score_chart_bundle_consensus", } def _uses_tangent(feature_set: str) -> bool: return feature_set in { "tangent", "context_tangent", "tangent_source_evidence", "context_tangent_source_evidence", "chart_tangent_compat", "chart_tangent_source_compat", "context_tangent_bundle_consensus", } def _uses_source_evidence(feature_set: str) -> bool: return feature_set in { "source_evidence", "tangent_source_evidence", "context_source_evidence", "context_tangent_source_evidence", "chart_source_compat", "chart_tangent_source_compat", "chart_source_bundle_consensus", } def _uses_chart_compat(feature_set: str) -> bool: return feature_set in { "chart_compat", "chart_tangent_compat", "score_chart_compat", "score_context_chart_compat", "chart_bundle_consensus", "score_chart_bundle_consensus", "chart_source_compat", "chart_tangent_source_compat", "chart_source_bundle_consensus", } def _uses_bundle_consensus(feature_set: str) -> bool: return feature_set in { "bundle_consensus", "score_bundle_consensus", "context_bundle_consensus", "context_tangent_bundle_consensus", "chart_bundle_consensus", "score_chart_bundle_consensus", "chart_source_bundle_consensus", } def _selector_chart_map(index_path: Path, *, chart_feature_mode: str) -> tuple[dict[str, Any], dict[str, Any]]: charts, index = load_chart_items( _resolve_index_path(index_path), max_charts=None, require_positive=False, include_hidden=False, include_metadata=True, chart_feature_mode=chart_feature_mode, ) return {chart.chart_id: chart for chart in charts}, index def _source_evidence_map(index_path: Path) -> tuple[dict[str, dict[str, Any]], dict[str, Any]]: index_path = _resolve_index_path(index_path) index = json.loads(index_path.read_text()) if not index.get("include_outcomes", False): raise SystemExit(f"{index_path} must include train outcomes for source evidence") grouped: dict[str, dict[str, Any]] = {} for shard in index.get("shards", []): shard_path = index_path.parent / shard["path"] with np.load(shard_path, allow_pickle=False) as data: chart_ids = data["chart_id"] labels = data["label"] is_base = data["is_base_branch"] tangents = data["spline_tangent_code"] utilities = data["utility"] delta_utilities = data["delta_utility"] outcomes = data["outcome_vector"] for row in range(chart_ids.shape[0]): if bool(is_base[row]): continue chart_id = str(chart_ids[row]) item = grouped.setdefault( chart_id, { "positive_tangents": [], "negative_tangents": [], "num_nonbase": 0, "positive_delta_utilities": [], "positive_utilities": [], "positive_success": [], "positive_progress": [], "positive_safety": [], }, ) item["num_nonbase"] += 1 label = str(labels[row]) if label == "positive": item["positive_tangents"].append(tangents[row].astype("float32")) item["positive_delta_utilities"].append(float(delta_utilities[row])) item["positive_utilities"].append(float(utilities[row])) outcome = np.asarray(outcomes[row], dtype=float) item["positive_success"].append(float(outcome[0]) if outcome.size > 0 else math.nan) item["positive_progress"].append(float(outcome[1]) if outcome.size > 1 else math.nan) item["positive_safety"].append(float(outcome[3]) if outcome.size > 3 else math.nan) elif label == "negative": item["negative_tangents"].append(tangents[row].astype("float32")) return grouped, index def _source_evidence_feature(source: dict[str, Any] | None, *, tangent: np.ndarray) -> np.ndarray: if not source: return np.zeros(len(SOURCE_EVIDENCE_NAMES), dtype=float) positives = np.asarray(source.get("positive_tangents") or [], dtype=float).reshape(-1, 21) negatives = np.asarray(source.get("negative_tangents") or [], dtype=float).reshape(-1, 21) positive_count = int(positives.shape[0]) negative_count = int(negatives.shape[0]) nonbase_count = int(source.get("num_nonbase") or (positive_count + negative_count)) positive_delta = _clean_array(source.get("positive_delta_utilities") or []) positive_utility = _clean_array(source.get("positive_utilities") or []) positive_success = _clean_array(source.get("positive_success") or []) positive_progress = _clean_array(source.get("positive_progress") or []) positive_safety_raw = np.asarray(source.get("positive_safety") or [], dtype=float) known_safety = positive_safety_raw[np.isfinite(positive_safety_raw)] tangent = np.asarray(tangent, dtype=float).reshape(-1) if tangent.size < 21: tangent = np.pad(tangent, (0, 21 - tangent.size)) elif tangent.size > 21: tangent = tangent[:21] pos_dist = _min_rms_distance(tangent, positives) neg_dist = _min_rms_distance(tangent, negatives) return np.asarray( [ 1.0, math.log1p(positive_count), math.log1p(negative_count), math.log1p(nonbase_count), positive_count / max(1.0, float(nonbase_count)), _safe_max(positive_delta), _safe_mean(positive_delta), _safe_max(positive_utility), _safe_mean(positive_utility), _safe_max(positive_success), _safe_mean(positive_success), _safe_max(positive_progress), _safe_mean(positive_progress), float(known_safety.size) / max(1.0, float(positive_count)), _safe_mean(known_safety), pos_dist, neg_dist, float(pos_dist < neg_dist) if math.isfinite(pos_dist) and math.isfinite(neg_dist) else 0.0, ], dtype=float, ) def _chart_compat_feature( target_chart: Any | None, source_chart: Any | None, *, chart_feature_mode: str, ) -> np.ndarray: if target_chart is None or source_chart is None: return np.zeros(len(CHART_COMPAT_NAMES), dtype=float) target_feature = np.asarray(getattr(target_chart, "feature", []), dtype=float).reshape(-1) source_feature = np.asarray(getattr(source_chart, "feature", []), dtype=float).reshape(-1) width = min(target_feature.size, source_feature.size) if width == 0: return np.zeros(len(CHART_COMPAT_NAMES), dtype=float) target_feature = target_feature[:width] source_feature = source_feature[:width] feature_diff = target_feature - source_feature target_base = np.asarray(getattr(target_chart, "base_action", []), dtype=float).reshape(-1) source_base = np.asarray(getattr(source_chart, "base_action", []), dtype=float).reshape(-1) base_width = min(target_base.size, source_base.size) target_base = target_base[:base_width] source_base = source_base[:base_width] base_diff = target_base - source_base if base_width else np.asarray([], dtype=float) target_segments = _chart_feature_segments(target_chart, chart_feature_mode=chart_feature_mode) source_segments = _chart_feature_segments(source_chart, chart_feature_mode=chart_feature_mode) target_obs, source_obs = target_segments.get("obs"), source_segments.get("obs") target_obj, source_obj = target_segments.get("obj"), source_segments.get("obj") obs_available = float(target_obs is not None and source_obs is not None) obj_available = float(target_obj is not None and source_obj is not None) obs_diff = ( np.asarray(target_obs, dtype=float) - np.asarray(source_obs, dtype=float) if obs_available else np.asarray([], dtype=float) ) obj_diff = ( np.asarray(target_obj, dtype=float) - np.asarray(source_obj, dtype=float) if obj_available else np.asarray([], dtype=float) ) return np.asarray( [ 1.0, _rms(target_feature), _rms(source_feature), _cosine(target_feature, source_feature), _rms(feature_diff), _linf(feature_diff), _absmean(feature_diff), float(np.dot(target_feature, source_feature) / max(1, width)), _rms(target_base), _rms(source_base), _cosine(target_base, source_base), _rms(base_diff), obs_available, _cosine(target_obs, source_obs) if obs_available else 0.0, _rms(obs_diff), _rms(target_obs) if obs_available else 0.0, _rms(source_obs) if obs_available else 0.0, obj_available, _cosine(target_obj, source_obj) if obj_available else 0.0, _rms(obj_diff), _rms(target_obj) if obj_available else 0.0, _rms(source_obj) if obj_available else 0.0, ], dtype=float, ) def _chart_feature_segments(chart: Any, *, chart_feature_mode: str) -> dict[str, np.ndarray]: feature = np.asarray(getattr(chart, "feature", []), dtype=float).reshape(-1) base_action = np.asarray(getattr(chart, "base_action", []), dtype=float).reshape(-1) offset = min(base_action.size, feature.size) if chart_feature_mode == "base": return {} context_width = 2 * CONTEXT_HASH_WIDTH + 2 offset = min(feature.size, offset + context_width) segments: dict[str, np.ndarray] = {} if chart_feature_mode in {"base_context_obs", "base_context_obs_obj"}: end = min(feature.size, offset + OBSERVATION_EMBED_DIM) if end - offset == OBSERVATION_EMBED_DIM: segments["obs"] = feature[offset:end] offset = end if chart_feature_mode in {"base_context_obj", "base_context_obs_obj"}: end = min(feature.size, offset + OBJECT_LAYOUT_EMBED_DIM) if end - offset == OBJECT_LAYOUT_EMBED_DIM: segments["obj"] = feature[offset:end] return segments def _rms(values: Any) -> float: array = np.asarray(values, dtype=float).reshape(-1) return float(np.sqrt(np.mean(array * array))) if array.size else 0.0 def _linf(values: Any) -> float: array = np.asarray(values, dtype=float).reshape(-1) return float(np.max(np.abs(array))) if array.size else 0.0 def _absmean(values: Any) -> float: array = np.asarray(values, dtype=float).reshape(-1) return float(np.mean(np.abs(array))) if array.size else 0.0 def _cosine(left: Any, right: Any) -> float: left_array = np.asarray(left, dtype=float).reshape(-1) right_array = np.asarray(right, dtype=float).reshape(-1) width = min(left_array.size, right_array.size) if width == 0: return 0.0 left_array = left_array[:width] right_array = right_array[:width] denom = float(np.linalg.norm(left_array) * np.linalg.norm(right_array)) if denom <= 1.0e-12: return 0.0 return float(np.dot(left_array, right_array) / denom) def _clean_array(values: Any) -> np.ndarray: array = np.asarray(values, dtype=float).reshape(-1) return array[np.isfinite(array)] def _safe_mean(values: np.ndarray) -> float: return float(values.mean()) if values.size else 0.0 def _safe_max(values: np.ndarray) -> float: return float(values.max()) if values.size else 0.0 def _min_rms_distance(tangent: np.ndarray, candidates: np.ndarray) -> float: if candidates.size == 0: return 0.0 diff = candidates - tangent.reshape(1, -1) return float(np.sqrt(np.mean(diff * diff, axis=1)).min()) def _score_shape_matrix(scores: list[float]) -> np.ndarray: """Deployment-visible row-relative score features for each candidate.""" score_array = np.asarray(scores, dtype=float).reshape(-1) if score_array.size == 0: return np.zeros((0, len(SCORE_SHAPE_NAMES)), dtype=float) order = sorted(range(score_array.size), key=lambda index: (-float(score_array[index]), index)) ranks = np.zeros(score_array.size, dtype=float) for rank, index in enumerate(order): ranks[index] = float(rank) sorted_scores = score_array[order] best = float(sorted_scores[0]) second = float(sorted_scores[1]) if sorted_scores.size > 1 else best denom = max(1.0, float(score_array.size - 1)) shifted = score_array - float(np.max(score_array)) exp_scores = np.exp(np.clip(shifted, -60.0, 60.0)) softmax = exp_scores / max(float(exp_scores.sum()), 1.0e-12) rows: list[list[float]] = [] for index, score in enumerate(score_array): rank = int(ranks[index]) prev_higher = sorted_scores[rank - 1] if rank > 0 else score next_lower = sorted_scores[rank + 1] if rank + 1 < sorted_scores.size else score percentile = float(np.mean(score_array <= score)) rows.append( [ float(rank) / denom, float(softmax[index]), float(score - best), float(score - second), float(score - prev_higher), float(score - next_lower), percentile, float(best - second), ] ) return np.asarray(rows, dtype=float) def _bundle_consensus_matrix( tangents: Any, scores: list[float], source_chart_ids: Any, source_task_ids: Any, ) -> np.ndarray: """Deployment-visible CTT bundle self-consistency features. These features are computed only from the generated transported tangents, their inference-time scores, and train-source identifiers already present in the candidate row. They deliberately do not inspect target positive or negative tangent sets, measured candidate utilities, or hidden outcomes. """ num_candidates = len(scores) if num_candidates == 0: return np.zeros((0, len(BUNDLE_CONSENSUS_NAMES)), dtype=float) tangent_matrix = _candidate_tangent_matrix(tangents, num_candidates) score_array = np.asarray(scores, dtype=float).reshape(-1) if score_array.size < num_candidates: score_array = np.pad(score_array, (0, num_candidates - score_array.size)) score_array = score_array[:num_candidates] source_ids = _string_list(source_chart_ids, num_candidates) task_ids = _string_list(source_task_ids, num_candidates) diff = tangent_matrix[:, None, :] - tangent_matrix[None, :, :] distances = np.sqrt(np.mean(diff * diff, axis=2)) nonself = ~np.eye(num_candidates, dtype=bool) peer_denominator = max(1.0, float(num_candidates - 1)) peer_distances = np.where(nonself, distances, np.nan) mean_peer = np.nanmean(peer_distances, axis=1) if num_candidates > 1 else np.zeros(num_candidates) min_peer = np.nanmin(peer_distances, axis=1) if num_candidates > 1 else np.zeros(num_candidates) mean_peer = np.nan_to_num(mean_peer, nan=0.0, posinf=0.0, neginf=0.0) min_peer = np.nan_to_num(min_peer, nan=0.0, posinf=0.0, neginf=0.0) medoid_index = int(np.argmin(mean_peer)) if num_candidates else 0 medoid_distances = distances[:, medoid_index] if num_candidates else np.zeros(0) rows: list[list[float]] = [] for index in range(num_candidates): peers_020 = [j for j in range(num_candidates) if j != index and distances[index, j] <= 0.20] peers_040 = [j for j in range(num_candidates) if j != index and distances[index, j] <= 0.40] group_040 = [index, *peers_040] unique_sources = {source_ids[j] for j in group_040 if source_ids[j]} unique_tasks = {task_ids[j] for j in group_040 if task_ids[j]} same_source_peers = [ j for j in peers_040 if source_ids[index] and source_ids[j] == source_ids[index] ] score_mean_020, score_max_020 = _score_stats(score_array, peers_020) score_mean_040, score_max_040 = _score_stats(score_array, peers_040) density_weights = np.exp(-np.clip(distances[index], 0.0, 10.0) / 0.40) density_weights[index] = 0.0 score_density = float(np.dot(density_weights, score_array) / max(1.0e-12, density_weights.sum())) rank_density = float(sum(1.0 for j in peers_040 if score_array[j] >= score_array[index])) rows.append( [ math.log1p(num_candidates), float(len(peers_020)), float(len(peers_040)), float(len(peers_020)) / peer_denominator, float(len(peers_040)) / peer_denominator, float(mean_peer[index]), float(min_peer[index]), float(medoid_distances[index]), float(index == medoid_index), score_mean_020, score_max_020, score_mean_040, score_max_040, float(len(unique_sources)), float(len(unique_sources)) / max(1.0, float(len(group_040))), float(len(unique_tasks)), float(len(same_source_peers)) / max(1.0, float(len(peers_040))), score_density, rank_density / max(1.0, float(len(peers_040))), ] ) output = np.asarray(rows, dtype=float) return np.nan_to_num(output, nan=0.0, posinf=0.0, neginf=0.0) def _candidate_tangent_matrix(tangents: Any, num_candidates: int) -> np.ndarray: rows: list[np.ndarray] = [] tangent_list = list(tangents or []) for index in range(num_candidates): tangent = np.asarray(tangent_list[index] if index < len(tangent_list) else [], dtype=float).reshape(-1) if tangent.size < 21: tangent = np.pad(tangent, (0, 21 - tangent.size)) elif tangent.size > 21: tangent = tangent[:21] rows.append(tangent.astype(float, copy=False)) return np.stack(rows, axis=0) def _string_list(values: Any, length: int) -> list[str]: raw = list(values or []) return [str(raw[index]) if index < len(raw) else "" for index in range(length)] def _score_stats(scores: np.ndarray, indices: list[int]) -> tuple[float, float]: if not indices: return 0.0, 0.0 values = scores[indices] return float(values.mean()), float(values.max()) def _context_feature(context: dict[str, Any]) -> np.ndarray: target_task = str(context.get("target_task_id", "")) source_task = str(context.get("source_task_id", "")) instruction = str(context.get("instruction", "")) return np.asarray( [ *_stable_hash_features(target_task, CONTEXT_HASH_WIDTH), *_stable_hash_features(source_task, CONTEXT_HASH_WIDTH), *_stable_hash_features(instruction.lower(), CONTEXT_HASH_WIDTH), float(bool(target_task) and target_task == source_task), min(len(instruction) / 128.0, 4.0), min(len(instruction.split()) / 32.0, 4.0), ], dtype=float, ) def _stable_hash_features(text: str, width: int) -> list[float]: digest = hashlib.sha256(text.encode("utf-8")).digest() return [((digest[index] / 127.5) - 1.0) for index in range(width)] def _target_value( target: str, *, utility_margin: float, candidate_success: float, success_bonus: float = 1.0, ) -> float: if target == "utility_margin": return float(utility_margin) if target == "success": return float(candidate_success) if target == "success_weighted_margin": return float(utility_margin) + float(success_bonus) * float(candidate_success) raise ValueError(f"unknown target: {target}") def _fit_select_ridge( dataset: dict[str, Any], *, lambdas: list[float], threshold_scope: str = "global", fit_objective: str = "pointwise", pairwise_weight: float = 1.0, ) -> dict[str, Any]: samples = dataset["samples"] if not samples: raise ValueError("cannot fit learned dominance selector without candidates") if fit_objective not in {"pointwise", "pairwise", "hybrid_pairwise"}: raise ValueError(f"unknown fit_objective: {fit_objective}") point_x = np.stack([sample["feature"] for sample in samples], axis=0) mean = np.zeros(point_x.shape[1], dtype=float) std = np.ones(point_x.shape[1], dtype=float) mean[1:] = point_x[:, 1:].mean(axis=0) std[1:] = point_x[:, 1:].std(axis=0) + 1.0e-6 x_norm, y, fit_design = _fit_design_matrix( dataset, fit_objective=fit_objective, pairwise_weight=pairwise_weight, mean=mean, std=std, ) if x_norm.shape[0] == 0: raise ValueError("fit objective produced no training rows") candidate_x_norm = _normalized_candidate_features(dataset, mean=mean, std=std) fit_design = { **fit_design, "num_candidate_rows": int(candidate_x_norm.shape[0]), "num_fit_rows": int(x_norm.shape[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 = candidate_x_norm @ weights tau, selection = _choose_thresholds( dataset, predictions, threshold_scope=threshold_scope, ) 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, "fit_design": fit_design, } assert best is not None return best def _normalized_candidate_features( dataset: dict[str, Any], *, mean: np.ndarray, std: np.ndarray, ) -> np.ndarray: x = np.stack([sample["feature"] for sample in dataset["samples"]], axis=0) x_norm = (x - mean) / std x_norm[:, 0] = 1.0 return x_norm def _fit_design_matrix( dataset: dict[str, Any], *, fit_objective: str, pairwise_weight: float, mean: np.ndarray, std: np.ndarray, ) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]: point_x = np.stack([sample["feature"] for sample in dataset["samples"]], axis=0) point_y = np.asarray([float(sample["target_margin"]) for sample in dataset["samples"]], dtype=float) point_x_norm = (point_x - mean) / std point_x_norm[:, 0] = 1.0 if fit_objective == "pointwise": return point_x_norm, point_y, { "fit_objective": fit_objective, "num_pointwise_rows": int(point_x.shape[0]), "num_pairwise_rows": 0, "pairwise_weight": 0.0, } pair_x, pair_y = _pairwise_design_matrix(dataset) pair_x_norm = pair_x / std pair_x_norm[:, 0] = 0.0 if fit_objective == "pairwise": return pair_x_norm, pair_y, { "fit_objective": fit_objective, "num_pointwise_rows": 0, "num_pairwise_rows": int(pair_x.shape[0]), "pairwise_weight": 1.0, } if fit_objective == "hybrid_pairwise": scale = math.sqrt(float(pairwise_weight)) x = np.concatenate([point_x_norm, pair_x_norm * scale], axis=0) y = np.concatenate([point_y, pair_y * scale], axis=0) return x, y, { "fit_objective": fit_objective, "num_pointwise_rows": int(point_x.shape[0]), "num_pairwise_rows": int(pair_x.shape[0]), "pairwise_weight": float(pairwise_weight), } raise ValueError(f"unknown fit_objective: {fit_objective}") def _pairwise_design_matrix(dataset: dict[str, Any]) -> tuple[np.ndarray, np.ndarray]: x_rows: list[np.ndarray] = [] y_rows: list[float] = [] for sample_indices in dataset["by_row"].values(): for left_pos, left_index in enumerate(sample_indices): left = dataset["samples"][left_index] left_target = float(left["target_margin"]) for right_index in sample_indices[left_pos + 1 :]: right = dataset["samples"][right_index] right_target = float(right["target_margin"]) target_delta = left_target - right_target if abs(target_delta) <= 1.0e-9: continue feature_delta = np.asarray(left["feature"], dtype=float) - np.asarray( right["feature"], dtype=float ) x_rows.append(feature_delta) y_rows.append(target_delta) x_rows.append(-feature_delta) y_rows.append(-target_delta) if not x_rows: raise ValueError("pairwise objective requires at least one non-tied within-row comparison") return np.stack(x_rows, axis=0), np.asarray(y_rows, dtype=float) def _choose_thresholds( dataset: dict[str, Any], predictions: np.ndarray, *, threshold_scope: str, ) -> tuple[float | dict[str, float], dict[str, float | None]]: if threshold_scope == "global": return _choose_tau(dataset, predictions) if threshold_scope != "task": raise ValueError(f"unknown threshold_scope: {threshold_scope}") global_tau, _global_summary = _choose_tau(dataset, predictions) task_taus: dict[str, float] = {"__global__": float(global_tau)} for task_id, row_indices in _rows_by_task(dataset).items(): subset = _subset_dataset_rows(dataset, row_indices) subset_predictions = _predictions_for_subset(dataset, predictions, subset) task_tau, _task_summary = _choose_tau(subset, subset_predictions) task_taus[task_id] = float(task_tau) cases = _evaluate_predictions(dataset, predictions, tau=task_taus) return task_taus, _simple_summary(cases) 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 | dict[str, float], include_pairwise_calibration: bool = False, ) -> list[dict[str, Any]]: predictions = _linear_predictions(dataset, weights, mean, std) return _evaluate_predictions( dataset, predictions, tau=tau, include_pairwise_calibration=include_pairwise_calibration, ) def _linear_predictions( dataset: dict[str, Any], weights: np.ndarray, mean: np.ndarray, std: np.ndarray, ) -> np.ndarray: x = np.stack([sample["feature"] for sample in dataset["samples"]], axis=0) x_norm = (x - mean) / std x_norm[:, 0] = 1.0 return x_norm @ weights def _evaluate_predictions( dataset: dict[str, Any], predictions: np.ndarray, *, tau: float | dict[str, float], include_pairwise_calibration: bool = False, pairwise_calibration: dict[str, Any] | None = None, ) -> list[dict[str, Any]]: samples = dataset["samples"] if include_pairwise_calibration and pairwise_calibration is None: pairwise_calibration = _pairwise_calibration_summary(dataset, predictions) pairwise_rows = (pairwise_calibration or {}).get("rows", {}) 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]) row_tau = _tau_for_sample(sample, tau) execute = predicted_margin > row_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"]) row = { "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": row_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, } if include_pairwise_calibration: calibration = pairwise_rows.get(row_index) or pairwise_rows.get(str(row_index), {}) row.update(_pairwise_calibration_scalars(calibration)) rows.append(row) return rows def _pairwise_calibration_summary( dataset: dict[str, Any], predictions: np.ndarray, *, n_bins: int = 10, ) -> dict[str, Any]: if len(predictions) != len(dataset["samples"]): raise ValueError("predictions must align with dataset samples") bins = [ { "count": 0, "accuracy_sum": 0.0, "confidence_sum": 0.0, "lower": index / n_bins, "upper": (index + 1) / n_bins, } for index in range(n_bins) ] by_row: dict[int, dict[str, Any]] = {} total_pairs = 0 correct_sum = 0.0 confidence_sum = 0.0 samples = dataset["samples"] for row_index, sample_indices in sorted(dataset["by_row"].items()): row_scores = [float(predictions[index]) for index in sample_indices] row_utilities = [float(samples[index]["candidate_utility"]) for index in sample_indices] row_metrics = pairwise_causal_dominance_ece(row_scores, row_utilities, n_bins=n_bins) by_row[int(row_index)] = row_metrics row_pairs = int(row_metrics.get("num_pairs") or 0) if row_pairs <= 0: continue total_pairs += row_pairs correct_sum += float(row_metrics.get("accuracy") or 0.0) * row_pairs confidence_sum += float(row_metrics.get("mean_confidence") or 0.0) * row_pairs for index, row_bin in enumerate(row_metrics.get("bins", [])): if index >= len(bins): break count = int(row_bin.get("count") or 0) bins[index]["count"] += count bins[index]["accuracy_sum"] += float(row_bin.get("accuracy") or 0.0) * count bins[index]["confidence_sum"] += float(row_bin.get("confidence") or 0.0) * count ece = 0.0 rendered_bins: list[dict[str, float | int]] = [] for bucket in bins: count = int(bucket["count"]) accuracy = bucket["accuracy_sum"] / count if count else 0.0 confidence = bucket["confidence_sum"] / count if count else 0.0 if total_pairs: ece += (count / total_pairs) * abs(accuracy - confidence) rendered_bins.append( { "lower": float(bucket["lower"]), "upper": float(bucket["upper"]), "count": count, "accuracy": accuracy, "confidence": confidence, "abs_gap": abs(accuracy - confidence), } ) return { "n_bins": int(n_bins), "num_rows": len(dataset["by_row"]), "ece": ece if total_pairs else math.nan, "num_pairs": int(total_pairs), "accuracy": correct_sum / total_pairs if total_pairs else math.nan, "mean_confidence": confidence_sum / total_pairs if total_pairs else math.nan, "bins": rendered_bins, "rows": by_row, } def _pairwise_calibration_scalars(calibration: dict[str, Any]) -> dict[str, float]: return { "pairwise_causal_calibration_ece": _finite_or_nan(calibration.get("ece")), "pairwise_causal_calibration_pairs": float(calibration.get("num_pairs") or 0), "pairwise_causal_calibration_accuracy": _finite_or_nan(calibration.get("accuracy")), "pairwise_causal_calibration_confidence": _finite_or_nan( calibration.get("mean_confidence") ), } def _pairwise_calibration_global(calibration: dict[str, Any]) -> dict[str, Any]: return {key: value for key, value in calibration.items() if key != "rows"} def _summary_with_pairwise( rows: list[dict[str, Any]], pairwise_calibration: dict[str, Any], ) -> dict[str, float | None]: summary = _simple_summary(rows) summary.update(_pairwise_calibration_scalars(pairwise_calibration)) return summary def _tau_for_sample(sample: dict[str, Any], tau: float | dict[str, float]) -> float: if isinstance(tau, dict): return float(tau.get(str(sample.get("task_id", "")), tau.get("__global__", 0.0))) return float(tau) def _rows_by_task(dataset: dict[str, Any]) -> dict[str, list[int]]: grouped: dict[str, list[int]] = defaultdict(list) for row_index, sample_indices in dataset["by_row"].items(): if not sample_indices: continue task_id = str(dataset["samples"][sample_indices[0]].get("task_id", "unknown")) grouped[task_id].append(int(row_index)) return grouped def _subset_dataset_rows(dataset: dict[str, Any], row_indices: list[int]) -> dict[str, Any]: wanted = set(int(row) for row in row_indices) old_to_new_sample: dict[int, int] = {} samples: list[dict[str, Any]] = [] by_row: dict[int, list[int]] = {} for old_row in sorted(wanted): if old_row not in dataset["by_row"]: continue new_row = len(by_row) by_row[new_row] = [] for old_sample_index in dataset["by_row"][old_row]: old_to_new_sample[old_sample_index] = len(samples) sample = dict(dataset["samples"][old_sample_index]) sample["row_index"] = new_row by_row[new_row].append(len(samples)) samples.append(sample) return { "samples": samples, "by_row": by_row, "num_rows": len(by_row), "_old_to_new_sample": old_to_new_sample, } def _predictions_for_subset( dataset: dict[str, Any], predictions: np.ndarray, subset: dict[str, Any], ) -> np.ndarray: old_to_new = subset.get("_old_to_new_sample", {}) ordered_old_indices = sorted(old_to_new, key=lambda old: old_to_new[old]) return np.asarray([float(predictions[old]) for old in ordered_old_indices], dtype=float) 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 _finite_or_nan(value: Any) -> float: return float(value) if isinstance(value, (int, float)) and math.isfinite(float(value)) else math.nan 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}{lrrrrrrrrr}", "\\toprule", "Rows & Coverage & Fallback & Base succ. & Selected succ. & Oracle succ. & OutcomePTR & Succ. support gap & Succ. selector gap & Cal. ECE \\\\", "\\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'))} & " f"{_fmt(summary.get('pairwise_causal_calibration_ece'))} \\\\", "\\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: `{_format_tau(metrics['tau'])}`", f"Threshold scope: `{metrics.get('threshold_scope', 'global')}`", f"Fit objective: `{metrics.get('fit_objective', 'pointwise')}`", f"Feature set: `{metrics['feature_set']}`", f"Target: `{metrics['target']}`", "", "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 | Calibration ECE |", "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", 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"{_fmt(calibration.get('pairwise_causal_calibration_ece'))} |", 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'))} | " f"{_fmt(summary.get('pairwise_causal_calibration_ece'))} |", "", "This is a selector diagnostic over already measured candidates, not a new rollout.", ] return "\n".join(lines) def _write_report_artifact( out_dir: Path, metrics: dict[str, Any], *, no_markdown_report: bool = False, ) -> None: report_path = out_dir / "report.md" if no_markdown_report: if report_path.exists(): report_path.unlink() return report_path.write_text(_report(metrics) + "\n") 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(_resolve_index_path(args.calibration_target_index)), "eval_input": _sha256(args.eval_input), "eval_target_index": _sha256(_resolve_index_path(args.eval_target_index)), } if getattr(args, "source_index", None) is not None or _uses_chart_compat(args.feature_set): hashes["source_index"] = _sha256( _resolve_index_path(args.source_index or args.calibration_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(_resolve_index_path(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 _resolve_index_path(path: Path) -> Path: return path / "index.json" if path.is_dir() else path 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 _format_tau(value: Any) -> str: if isinstance(value, dict): return json.dumps({key: round(float(val), 6) for key, val in sorted(value.items())}) if isinstance(value, (int, float)) and math.isfinite(float(value)): return f"{float(value):.6f}" return str(value) 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())