#!/usr/bin/env python from __future__ import annotations import argparse import hashlib import json import math import subprocess import sys from collections import defaultdict from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) import numpy as np # noqa: E402 from sklearn.ensemble import ( # noqa: E402 HistGradientBoostingClassifier, HistGradientBoostingRegressor, RandomForestRegressor, ) from cil.chart_features import CHART_FEATURE_MODES # noqa: E402 from cil.metrics import macro_micro_summary # noqa: E402 from scripts.eval_dominance_selector import ( # noqa: E402 _DominanceScorer, _chart_map, _first_train_seed, _rows, ) from scripts.eval_learned_dominance_selector import ( # noqa: E402 FEATURE_SET_CHOICES, _candidate_dataset, _evaluate_predictions, _feature_names, _group_means, _pairwise_calibration_global, _pairwise_calibration_summary, _resolve_index_path, _selector_chart_map, _simple_summary, _source_evidence_map, _summary_with_pairwise, _uses_chart_compat, _uses_source_evidence, ) def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( description=( "Train a nonlinear train-only dominance selector on measured CTT " "calibration rollouts and evaluate it on held-out measured 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_nonlinear_dominance_train_to_test"), ) parser.add_argument("--k", type=int, default=8) parser.add_argument( "--feature-set", choices=FEATURE_SET_CHOICES, default="context_tangent", ) 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", "positive_margin"), default="positive_margin", ) parser.add_argument( "--model-types", default="hgb_classifier,hgb_regressor,rf_regressor", help="Comma-separated candidates from hgb_classifier,hgb_regressor,rf_regressor.", ) parser.add_argument("--selection-frac", type=float, default=0.35) parser.add_argument("--seed", type=int, default=0) 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") if not 0.05 <= args.selection_frac <= 0.75: parser.error("--selection-frac must be in [0.05, 0.75]") model_types = [item.strip() for item in args.model_types.split(",") if item.strip()] allowed = {"hgb_classifier", "hgb_regressor", "rf_regressor"} if not model_types or any(item not in allowed for item in model_types): parser.error(f"--model-types must be comma-separated values from {sorted(allowed)}") 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, ) dataset_target = ( "utility_margin" if args.target == "positive_margin" else args.target ) calibration_dataset = _candidate_dataset( calibration_rows, calibration_charts, scorer=scorer, k=args.k, feature_set=args.feature_set, target=dataset_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, ) eval_dataset = _candidate_dataset( eval_rows, eval_charts, scorer=scorer, k=args.k, feature_set=args.feature_set, target=dataset_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, ) fit_rows, select_rows = _split_rows( calibration_dataset, selection_frac=args.selection_frac, seed=args.seed, ) fit_dataset = _subset_dataset(calibration_dataset, fit_rows) select_dataset = _subset_dataset(calibration_dataset, select_rows) if not fit_dataset["samples"] or not select_dataset["samples"]: raise SystemExit("calibration split produced an empty fit or selection set") best = _fit_select_model( fit_dataset, select_dataset, model_types=model_types, target=args.target, seed=args.seed, ) eval_predictions = _predict(best["model"], eval_dataset, model_type=best["model_type"]) fit_predictions = _predict(best["model"], fit_dataset, model_type=best["model_type"]) select_predictions = _predict(best["model"], select_dataset, model_type=best["model_type"]) eval_pairwise = _pairwise_calibration_summary(eval_dataset, eval_predictions) fit_pairwise = _pairwise_calibration_summary(fit_dataset, fit_predictions) select_pairwise = _pairwise_calibration_summary(select_dataset, select_predictions) eval_cases = _evaluate_predictions( eval_dataset, eval_predictions, tau=best["tau"], include_pairwise_calibration=True, pairwise_calibration=eval_pairwise, ) fit_cases = _evaluate_predictions( fit_dataset, fit_predictions, tau=best["tau"], include_pairwise_calibration=True, pairwise_calibration=fit_pairwise, ) select_cases = _evaluate_predictions( select_dataset, select_predictions, tau=best["tau"], include_pairwise_calibration=True, pairwise_calibration=select_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": "nonlinear_dominance_selector_eval", "schema_version": 1, "k": args.k, "feature_set": args.feature_set, "feature_names": _feature_names(args.feature_set), "target": args.target, "model_types": model_types, "selected_model_type": best["model_type"], "selected_model_params": best["params"], "tau": best["tau"], "selection_frac": args.selection_frac, "seed": args.seed, "chart_feature_mode": chart_feature_mode, "selector_chart_feature_mode": ( args.selector_chart_feature_mode if _uses_chart_compat(args.feature_set) else None ), "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_fit_rows": fit_dataset["num_rows"], "num_selection_rows": select_dataset["num_rows"], "num_eval_rows": len(eval_rows), "num_fit_candidates": len(fit_dataset["samples"]), "num_selection_candidates": len(select_dataset["samples"]), "num_eval_candidates": len(eval_dataset["samples"]), "model_selection": best["selection"], "fit_summary": _summary_with_pairwise(fit_cases, fit_pairwise), "selection_summary": _summary_with_pairwise(select_cases, select_pairwise), "eval_summary": _summary_with_pairwise(eval_cases, eval_pairwise), "pairwise_causal_calibration": { "fit": _pairwise_calibration_global(fit_pairwise), "selection": _pairwise_calibration_global(select_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 nonlinear dominance selector on calibration-fit rows only\n" f"selected_model_type={best['model_type']}\n" f"selected_model_params={json.dumps(best['params'], sort_keys=True)}\n" f"tau={best['tau']:.6f}\n" f"calibration_input={args.calibration_input}\n" ) (out_dir / "eval.log").write_text( "selected model/tau on held-out calibration-selection rows; evaluated 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), "model_type": best["model_type"], "tau": best["tau"], "selected_success": metrics["eval_summary"]["selected_success"], }, indent=2, ) ) return 0 def _split_rows( dataset: dict[str, Any], *, selection_frac: float, seed: int, ) -> tuple[list[int], list[int]]: row_indices = sorted(dataset["by_row"]) scored = [] for row_index in row_indices: sample = dataset["samples"][dataset["by_row"][row_index][0]] key = f"{sample['chart_id']}|{sample['seed']}|{seed}" digest = hashlib.sha256(key.encode("utf-8")).digest() value = int.from_bytes(digest[:8], "big") / float(2**64 - 1) scored.append((value, row_index)) selection = [row for value, row in scored if value < selection_frac] fit = [row for value, row in scored if value >= selection_frac] if not selection or not fit: cutoff = max(1, min(len(scored) - 1, int(round(len(scored) * selection_frac)))) selection = [row for _value, row in scored[:cutoff]] fit = [row for _value, row in scored[cutoff:]] return fit, selection def _subset_dataset(dataset: dict[str, Any], row_indices: list[int]) -> dict[str, Any]: wanted = set(int(row) for row in row_indices) samples = [] by_row: dict[int, list[int]] = {} row_map: dict[int, int] = {} for old_row in sorted(wanted): if old_row not in dataset["by_row"]: continue new_row = len(row_map) row_map[old_row] = new_row by_row[new_row] = [] for old_sample_index in dataset["by_row"][old_row]: 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)} def _fit_select_model( fit_dataset: dict[str, Any], select_dataset: dict[str, Any], *, model_types: list[str], target: str, seed: int, ) -> dict[str, Any]: x_fit = np.stack([sample["feature"] for sample in fit_dataset["samples"]], axis=0) best: dict[str, Any] | None = None for model_type, params in _model_grid(model_types, seed=seed): y_fit = _target_array(fit_dataset, target=target, model_type=model_type) if model_type == "hgb_classifier" and len(set(int(value) for value in y_fit)) < 2: continue model = _make_model(model_type, params) model.fit(x_fit, y_fit) predictions = _predict(model, select_dataset, model_type=model_type) tau, selection = _choose_tau(select_dataset, predictions) key = ( float(selection.get("selected_success") or 0.0), float(selection.get("selected_utility") or 0.0), float(selection.get("coverage") or 0.0), -float(_model_complexity(model_type, params)), ) if best is None or key > best["key"]: best = { "key": key, "model": model, "model_type": model_type, "params": params, "tau": tau, "selection": selection, } if best is None: raise ValueError("could not fit any nonlinear dominance model") return best def _target_array(dataset: dict[str, Any], *, target: str, model_type: str) -> np.ndarray: if model_type == "hgb_classifier": if target == "success": return np.asarray( [float(sample["candidate_success"]) for sample in dataset["samples"]], dtype=int, ) return np.asarray( [float(sample["measured_utility_margin"] > 0.0) for sample in dataset["samples"]], dtype=int, ) if target == "positive_margin": return np.asarray( [float(sample["measured_utility_margin"] > 0.0) for sample in dataset["samples"]], dtype=float, ) if target == "success": return np.asarray( [float(sample["candidate_success"]) for sample in dataset["samples"]], dtype=int, ) return np.asarray([float(sample["target_margin"]) for sample in dataset["samples"]], dtype=float) def _model_grid(model_types: list[str], *, seed: int) -> list[tuple[str, dict[str, Any]]]: grid: list[tuple[str, dict[str, Any]]] = [] for model_type in model_types: if model_type == "hgb_classifier": for max_iter in (40, 80): for max_leaf_nodes in (7, 15): grid.append( ( model_type, { "learning_rate": 0.05, "max_iter": max_iter, "max_leaf_nodes": max_leaf_nodes, "l2_regularization": 0.01, "random_state": seed, }, ) ) elif model_type == "hgb_regressor": for max_iter in (40, 80): for max_leaf_nodes in (7, 15): grid.append( ( model_type, { "learning_rate": 0.05, "max_iter": max_iter, "max_leaf_nodes": max_leaf_nodes, "l2_regularization": 0.01, "random_state": seed, }, ) ) elif model_type == "rf_regressor": for max_depth in (3, 5): grid.append( ( model_type, { "n_estimators": 128, "max_depth": max_depth, "min_samples_leaf": 8, "random_state": seed, "n_jobs": 1, }, ) ) else: raise ValueError(f"unknown model_type: {model_type}") return grid def _make_model(model_type: str, params: dict[str, Any]) -> Any: if model_type == "hgb_classifier": return HistGradientBoostingClassifier(**params) if model_type == "hgb_regressor": return HistGradientBoostingRegressor(**params) if model_type == "rf_regressor": return RandomForestRegressor(**params) raise ValueError(f"unknown model_type: {model_type}") def _predict(model: Any, dataset: dict[str, Any], *, model_type: str) -> np.ndarray: x = np.stack([sample["feature"] for sample in dataset["samples"]], axis=0) if model_type == "hgb_classifier": return np.asarray(model.predict_proba(x)[:, 1], dtype=float) return np.asarray(model.predict(x), dtype=float) 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 _model_complexity(model_type: str, params: dict[str, Any]) -> float: if model_type == "rf_regressor": return float(params.get("n_estimators", 0) * params.get("max_depth", 1)) return float(params.get("max_iter", 0) * params.get("max_leaf_nodes", 1)) def _table(metrics: dict[str, Any]) -> str: summary = metrics["eval_summary"] lines = [ "% Auto-generated by scripts/eval_nonlinear_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: fit = metrics["fit_summary"] selection = metrics["selection_summary"] eval_summary = metrics["eval_summary"] lines = [ "# Nonlinear Train-Calibrated CTT Selector", "", f"Calibration rows: `{metrics['num_calibration_rows']}`", f"Fit rows: `{metrics['num_fit_rows']}`", f"Selection rows: `{metrics['num_selection_rows']}`", f"Eval rows: `{metrics['num_eval_rows']}`", f"Selected model: `{metrics['selected_model_type']}`", f"Selected params: `{json.dumps(metrics['selected_model_params'], sort_keys=True)}`", f"Tau: `{metrics['tau']:.6f}`", f"Feature set: `{metrics['feature_set']}`", f"Target: `{metrics['target']}`", "", "The model is fit on calibration-fit rows, and model/tau are selected on held-out calibration-selection 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"| fit | {_fmt(fit.get('coverage'))} | {_fmt(fit.get('fallback_rate'))} | " f"{_fmt(fit.get('base_success'))} | {_fmt(fit.get('selected_success'))} | " f"{_fmt(fit.get('proposal_oracle_success'))} | {_fmt(fit.get('outcome_ptr'))} | " f"{_fmt(fit.get('success_support_gap'))} | {_fmt(fit.get('success_selector_gap'))} | " f"{_fmt(fit.get('pairwise_causal_calibration_ece'))} |", f"| selection | {_fmt(selection.get('coverage'))} | {_fmt(selection.get('fallback_rate'))} | " f"{_fmt(selection.get('base_success'))} | {_fmt(selection.get('selected_success'))} | " f"{_fmt(selection.get('proposal_oracle_success'))} | {_fmt(selection.get('outcome_ptr'))} | " f"{_fmt(selection.get('success_support_gap'))} | {_fmt(selection.get('success_selector_gap'))} | " f"{_fmt(selection.get('pairwise_causal_calibration_ece'))} |", f"| eval | {_fmt(eval_summary.get('coverage'))} | {_fmt(eval_summary.get('fallback_rate'))} | " f"{_fmt(eval_summary.get('base_success'))} | {_fmt(eval_summary.get('selected_success'))} | " f"{_fmt(eval_summary.get('proposal_oracle_success'))} | {_fmt(eval_summary.get('outcome_ptr'))} | " f"{_fmt(eval_summary.get('success_support_gap'))} | {_fmt(eval_summary.get('success_selector_gap'))} | " f"{_fmt(eval_summary.get('pairwise_causal_calibration_ece'))} |", "", "This is a train-calibrated 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_nonlinear_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_source_evidence(args.feature_set) 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), "source_index": _index_hash(args.source_index or args.calibration_target_index) if ( getattr(args, "source_index", None) is not None or _uses_source_evidence(args.feature_set) or _uses_chart_compat(args.feature_set) ) else None, }, 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: 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())