from __future__ import annotations import json from types import SimpleNamespace import numpy as np import pytest from scripts.eval_dominance_selector import ( _choose_tau, _conformal_quantile, _evaluate_case, _pairwise_calibration_summary as _score_pairwise_calibration_summary, ) from scripts.eval_learned_dominance_selector import ( _chart_compat_feature, _candidate_feature, _choose_thresholds, _evaluate_dataset, _evaluate_predictions, _feature_names, _fit_select_ridge, _pairwise_calibration_summary, _score_shape_matrix, _source_evidence_feature, _target_value, _write_report_artifact, ) from scripts.eval_nonlinear_dominance_selector import _split_rows, _subset_dataset from scripts.eval_nonlinear_dominance_selector import _write_provenance as _write_nonlinear_provenance def test_conformal_quantile_and_fallback_selection() -> None: quantile = _conformal_quantile([0.0, 0.2, 0.4, 0.6], alpha=0.25) assert quantile == 0.6 case = { "chart_id": "c0", "task_id": "pick", "seed": "0", "train_seed": "0", "predicted_margin": 0.5, "measured_margin": 0.4, "base_utility": 0.3, "base_success": 0.0, "top_generated_utility": 1.2, "top_generated_success": 1.0, "proposal_oracle_utility": 1.2, "proposal_oracle_success": 1.0, "hidden_chart_oracle_utility": 1.5, "hidden_chart_oracle_success": 1.0, "outcome_ptr": 1.0, } executed = _evaluate_case(case, residual_quantile=0.1, tau=0.0) assert executed["execute_generated"] == 1.0 assert executed["selected_success"] == 1.0 fallback = _evaluate_case(case, residual_quantile=0.6, tau=0.0) assert fallback["execute_generated"] == 0.0 assert fallback["fallback_to_base"] == 1.0 assert fallback["selected_utility"] == 0.3 def test_score_dominance_logs_pairwise_calibration_and_unsafe_execution() -> None: case = { "chart_id": "c0", "task_id": "pick", "seed": "0", "train_seed": "0", "predicted_scores": [3.0, 1.0, -1.0], "predicted_margin": 1.0, "measured_margin": 0.7, "base_utility": 0.2, "base_success": 0.0, "base_unsafe_known": 0.0, "top_generated_utility": 0.9, "top_generated_success": 1.0, "top_candidate_unsafe_known": 1.0, "candidate_safety_label_coverage": 1.0, "candidate_unsafe_rate_known": 1.0 / 3.0, "generated_utilities": [0.9, 0.4, 0.1], "proposal_oracle_utility": 0.9, "proposal_oracle_success": 1.0, "hidden_chart_oracle_utility": 1.0, "hidden_chart_oracle_success": 1.0, "outcome_ptr": 1.0, } calibration = _score_pairwise_calibration_summary([case]) executed = _evaluate_case( case, residual_quantile=0.0, tau=0.0, pairwise_calibration=calibration["rows"][0], ) fallback = _evaluate_case( case, residual_quantile=2.0, tau=0.0, pairwise_calibration=calibration["rows"][0], ) assert calibration["num_pairs"] == 3 assert executed["pairwise_causal_calibration_pairs"] == 3.0 assert executed["pairwise_causal_calibration_accuracy"] == 1.0 assert executed["unsafe_execution_label_known"] == 1.0 assert executed["unsafe_execution_known"] == 1.0 assert fallback["fallback_to_base"] == 1.0 assert fallback["unsafe_execution_label_known"] == 1.0 assert fallback["unsafe_execution_known"] == 0.0 def test_choose_tau_prefers_success_then_utility_then_coverage() -> None: cases = [ { "predicted_margin": 2.0, "base_utility": 0.5, "base_success": 0.0, "top_generated_utility": 1.0, "top_generated_success": 1.0, "proposal_oracle_utility": 1.0, "proposal_oracle_success": 1.0, "hidden_chart_oracle_utility": 1.0, "hidden_chart_oracle_success": 1.0, "outcome_ptr": 1.0, }, { "predicted_margin": -1.0, "base_utility": 0.8, "base_success": 1.0, "top_generated_utility": 0.1, "top_generated_success": 0.0, "proposal_oracle_utility": 0.1, "proposal_oracle_success": 0.0, "hidden_chart_oracle_utility": 1.0, "hidden_chart_oracle_success": 1.0, "outcome_ptr": 0.0, }, ] tau = _choose_tau(cases, residual_quantile=0.0) evaluated = [_evaluate_case(case, residual_quantile=0.0, tau=tau) for case in cases] assert sum(row["selected_success"] for row in evaluated) == 2.0 def test_learned_dominance_ridge_prefers_positive_candidate() -> None: def sample(row: int, candidate: int, score: float, margin: float, success: float) -> dict: feature = np.zeros(10, dtype=float) feature[0] = 1.0 feature[1] = score feature[2] = margin feature[4] = float(candidate) return { "row_index": row, "candidate_index": candidate, "feature": feature, "target_margin": success - 0.2, "candidate_utility": success, "candidate_success": success, "base_utility": 0.2, "base_success": 0.0, "proposal_oracle_utility": success, "proposal_oracle_success": success, "hidden_chart_oracle_utility": 1.0, "hidden_chart_oracle_success": 1.0, "outcome_ptr": success, "chart_id": f"c{row}", "task_id": "pick", "seed": "0", "train_seed": "0", } dataset = { "samples": [ sample(0, 0, 2.0, 2.0, 1.0), sample(0, 1, -1.0, -1.0, 0.0), sample(1, 0, 2.0, 2.0, 1.0), sample(1, 1, -1.0, -1.0, 0.0), ], "by_row": {0: [0, 1], 1: [2, 3]}, "num_rows": 2, } fit = _fit_select_ridge(dataset, lambdas=[0.1]) rows = _evaluate_dataset(dataset, fit["weights"], fit["mean"], fit["std"], tau=fit["tau"]) assert [row["selected_success"] for row in rows] == [1.0, 1.0] def test_learned_dominance_pairwise_objective_prefers_within_chart_winner() -> None: def sample(row: int, candidate: int, feature_value: float, success: float) -> dict: feature = np.zeros(10, dtype=float) feature[0] = 1.0 feature[1] = feature_value return { "row_index": row, "candidate_index": candidate, "feature": feature, "target_margin": success, "candidate_utility": success, "candidate_success": success, "base_utility": 0.0, "base_success": 0.0, "proposal_oracle_utility": 1.0, "proposal_oracle_success": 1.0, "hidden_chart_oracle_utility": 1.0, "hidden_chart_oracle_success": 1.0, "outcome_ptr": 1.0, "chart_id": f"c{row}", "task_id": "pick", "seed": "0", "train_seed": "0", } dataset = { "samples": [ sample(0, 0, 2.0, 1.0), sample(0, 1, -1.0, 0.0), sample(0, 2, 0.0, 0.0), sample(1, 0, 3.0, 1.0), sample(1, 1, -2.0, 0.0), sample(1, 2, 0.5, 0.0), ], "by_row": {0: [0, 1, 2], 1: [3, 4, 5]}, "num_rows": 2, } fit = _fit_select_ridge(dataset, lambdas=[0.1], fit_objective="pairwise") rows = _evaluate_dataset(dataset, fit["weights"], fit["mean"], fit["std"], tau=fit["tau"]) assert fit["fit_design"]["num_candidate_rows"] == 6 assert fit["fit_design"]["num_pairwise_rows"] == 8 assert fit["fit_design"]["num_fit_rows"] == 8 assert [row["selected_candidate_index"] for row in rows] == [0, 0] assert [row["selected_success"] for row in rows] == [1.0, 1.0] def test_task_scoped_threshold_can_use_visible_task_bucket() -> None: dataset = { "samples": [ { "row_index": 0, "candidate_index": 0, "feature": np.ones(2, dtype=float), "target_margin": 1.0, "candidate_utility": 1.0, "candidate_success": 1.0, "base_utility": 0.0, "base_success": 0.0, "proposal_oracle_utility": 1.0, "proposal_oracle_success": 1.0, "hidden_chart_oracle_utility": 1.0, "hidden_chart_oracle_success": 1.0, "outcome_ptr": 1.0, "chart_id": "a", "task_id": "task_execute", "seed": "0", "train_seed": "0", }, { "row_index": 1, "candidate_index": 0, "feature": np.ones(2, dtype=float), "target_margin": -1.0, "candidate_utility": 0.0, "candidate_success": 0.0, "base_utility": 1.0, "base_success": 1.0, "proposal_oracle_utility": 0.0, "proposal_oracle_success": 0.0, "hidden_chart_oracle_utility": 1.0, "hidden_chart_oracle_success": 1.0, "outcome_ptr": 0.0, "chart_id": "b", "task_id": "task_fallback", "seed": "0", "train_seed": "0", }, ], "by_row": {0: [0], 1: [1]}, "num_rows": 2, } predictions = np.asarray([0.6, 0.8], dtype=float) global_tau, _global_summary = _choose_thresholds( dataset, predictions, threshold_scope="global", ) task_tau, task_summary = _choose_thresholds( dataset, predictions, threshold_scope="task", ) global_rows = _evaluate_predictions(dataset, predictions, tau=global_tau) task_rows = _evaluate_predictions(dataset, predictions, tau=task_tau) assert np.mean([row["selected_success"] for row in global_rows]) == 0.5 assert task_summary["selected_success"] == 1.0 assert [row["selected_success"] for row in task_rows] == [1.0, 1.0] def test_pairwise_calibration_ece_is_logged_per_selector_row() -> None: def sample(row: int, candidate: int, utility: float) -> dict: return { "row_index": row, "candidate_index": candidate, "feature": np.asarray([1.0, float(candidate)]), "target_margin": utility, "candidate_utility": utility, "candidate_success": float(utility > 0.5), "base_utility": 0.0, "base_success": 0.0, "proposal_oracle_utility": 1.0, "proposal_oracle_success": 1.0, "hidden_chart_oracle_utility": 1.0, "hidden_chart_oracle_success": 1.0, "outcome_ptr": 1.0, "chart_id": f"c{row}", "task_id": "pick", "seed": "0", "train_seed": "0", } dataset = { "samples": [ sample(0, 0, 1.0), sample(0, 1, 0.0), sample(0, 2, 0.5), sample(1, 0, 1.0), sample(1, 1, 0.0), sample(1, 2, 0.5), ], "by_row": {0: [0, 1, 2], 1: [3, 4, 5]}, "num_rows": 2, } good_predictions = np.asarray([2.0, -2.0, 0.0, 2.0, -2.0, 0.0], dtype=float) bad_predictions = -good_predictions good = _pairwise_calibration_summary(dataset, good_predictions) bad = _pairwise_calibration_summary(dataset, bad_predictions) rows = _evaluate_predictions( dataset, good_predictions, tau=-10.0, include_pairwise_calibration=True, pairwise_calibration=good, ) assert good["num_pairs"] == 6 assert good["accuracy"] == 1.0 assert bad["accuracy"] == 0.0 assert rows[0]["pairwise_causal_calibration_pairs"] == 3.0 assert rows[0]["pairwise_causal_calibration_ece"] == pytest.approx( good["rows"][0]["ece"] ) def test_learned_dominance_tangent_features_and_success_weighted_target() -> None: tangent = np.arange(21, dtype=float) basic = _candidate_feature( score=1.0, base_score=0.25, score_mean=0.5, score_std=0.5, candidate_index=1, candidate_type="ctt_transport_rank3", tangent=tangent, num_candidates=8, feature_set="basic", ) expanded = _candidate_feature( score=1.0, base_score=0.25, score_mean=0.5, score_std=0.5, candidate_index=1, candidate_type="ctt_transport_rank3", tangent=tangent, num_candidates=8, feature_set="tangent", ) context = _candidate_feature( score=1.0, base_score=0.25, score_mean=0.5, score_std=0.5, candidate_index=1, candidate_type="ctt_transport_rank3", tangent=tangent, num_candidates=8, feature_set="context", context={ "target_task_id": "PickCube-v1", "source_task_id": "PickCube-v1", "instruction": "Pick up the cube.", }, ) context_tangent = _candidate_feature( score=1.0, base_score=0.25, score_mean=0.5, score_std=0.5, candidate_index=1, candidate_type="ctt_transport_rank3", tangent=tangent, num_candidates=8, feature_set="context_tangent", context={ "target_task_id": "PickCube-v1", "source_task_id": "PickCube-v1", "instruction": "Pick up the cube.", }, ) source_stats = { "positive_tangents": [np.zeros(21, dtype=float), np.ones(21, dtype=float)], "negative_tangents": [np.full(21, 4.0, dtype=float)], "num_nonbase": 4, "positive_delta_utilities": [0.4, 0.8], "positive_utilities": [1.2, 1.6], "positive_success": [1.0, 0.0], "positive_progress": [0.8, 0.9], "positive_safety": [0.0, 0.0], } source_evidence = _source_evidence_feature(source_stats, tangent=np.zeros(21, dtype=float)) context_tangent_source = _candidate_feature( score=1.0, base_score=0.25, score_mean=0.5, score_std=0.5, candidate_index=1, candidate_type="ctt_transport_rank3", tangent=tangent, num_candidates=8, feature_set="context_tangent_source_evidence", context={ "target_task_id": "PickCube-v1", "source_task_id": "PickCube-v1", "instruction": "Pick up the cube.", }, source_evidence=source_evidence, ) assert len(basic) == len(_feature_names("basic")) == 10 assert len(expanded) == len(_feature_names("tangent")) == 52 assert len(context) == len(_feature_names("context")) == 37 assert len(context_tangent) == len(_feature_names("context_tangent")) == 79 assert len(context_tangent_source) == len(_feature_names("context_tangent_source_evidence")) assert len(context_tangent_source) == 97 assert basic[2] == 0.75 assert basic[6] == 3.0 assert context[-3] == 1.0 assert source_evidence[0] == 1.0 assert source_evidence[1] == pytest.approx(np.log1p(2)) assert source_evidence[-1] == 1.0 assert _target_value("success_weighted_margin", utility_margin=0.2, candidate_success=1.0) == 1.2 assert _target_value( "success_weighted_margin", utility_margin=0.2, candidate_success=1.0, success_bonus=2.5, ) == 2.7 def test_score_shape_feature_set_and_markdown_suppression(tmp_path) -> None: score_shape = _score_shape_matrix([3.0, 1.0, 2.0]) feature = _candidate_feature( score=3.0, base_score=0.25, score_mean=2.0, score_std=1.0, candidate_index=0, candidate_type="ctt_transport_rank0", tangent=np.zeros(21, dtype=float), num_candidates=3, feature_set="score_context_chart_compat", context={ "target_task_id": "PickCube-v1", "source_task_id": "PickCube-v1", "instruction": "Pick up the cube.", }, chart_compat=np.ones(22, dtype=float), score_shape=score_shape[0], ) assert score_shape.shape == (3, 8) assert score_shape[0, 0] == 0.0 assert score_shape[0, 2] == 0.0 assert score_shape[0, 3] == 1.0 assert score_shape[1, 0] == 1.0 assert len(feature) == len(_feature_names("score_context_chart_compat")) assert len(feature) == 10 + 8 + 27 + 22 metrics = { "num_calibration_rows": 1, "num_eval_rows": 1, "selected_lambda": 0.0, "tau": 0.0, "threshold_scope": "global", "fit_objective": "pointwise", "feature_set": "basic", "target": "utility_margin", "eval_summary": {}, "calibration_summary": {}, } report_path = tmp_path / "report.md" report_path.write_text("old\n") _write_report_artifact(tmp_path, metrics, no_markdown_report=True) assert not report_path.exists() _write_report_artifact(tmp_path, metrics, no_markdown_report=False) assert report_path.exists() def test_chart_compat_feature_is_deployment_visible_and_fixed_width() -> None: base = np.arange(14, dtype=float).reshape(2, 7) / 10.0 context = np.linspace(-0.5, 0.5, 18) obs = np.linspace(0.0, 1.0, 32) obj = np.linspace(1.0, 2.0, 64) target = SimpleNamespace( base_action=base, feature=np.concatenate([base.reshape(-1), context, obs, obj]), ) source = SimpleNamespace( base_action=base + 0.1, feature=np.concatenate([base.reshape(-1) + 0.1, context, obs + 0.2, obj + 0.3]), ) compat = _chart_compat_feature( target, source, chart_feature_mode="base_context_obs_obj", ) feature = _candidate_feature( score=1.0, base_score=0.25, score_mean=0.5, score_std=0.5, candidate_index=1, candidate_type="ctt_transport_rank3", tangent=np.zeros(21, dtype=float), num_candidates=8, feature_set="chart_compat", chart_compat=compat, ) assert compat.shape == (22,) assert len(feature) == len(_feature_names("chart_compat")) == 32 assert compat[0] == 1.0 assert compat[12] == 1.0 assert compat[17] == 1.0 assert np.all(np.isfinite(compat)) assert np.all(_chart_compat_feature(target, None, chart_feature_mode="base_context_obs_obj") == 0.0) def test_nonlinear_dominance_split_is_row_disjoint() -> None: samples = [] by_row = {} for row in range(10): by_row[row] = [] for candidate in range(2): by_row[row].append(len(samples)) samples.append( { "row_index": row, "candidate_index": candidate, "feature": np.asarray([1.0, float(candidate)]), "target_margin": float(candidate), "measured_utility_margin": float(candidate), "candidate_success": float(candidate), "candidate_utility": float(candidate), "base_utility": 0.0, "base_success": 0.0, "proposal_oracle_utility": 1.0, "proposal_oracle_success": 1.0, "hidden_chart_oracle_utility": 1.0, "hidden_chart_oracle_success": 1.0, "outcome_ptr": 1.0, "chart_id": f"c{row}", "task_id": "pick", "seed": "0", "train_seed": "0", } ) dataset = {"samples": samples, "by_row": by_row, "num_rows": len(by_row)} fit_rows, select_rows = _split_rows(dataset, selection_frac=0.3, seed=7) assert set(fit_rows).isdisjoint(select_rows) assert set(fit_rows) | set(select_rows) == set(by_row) subset = _subset_dataset(dataset, select_rows) assert subset["num_rows"] == len(select_rows) assert len(subset["samples"]) == 2 * len(select_rows) def test_nonlinear_chart_compat_provenance_records_default_source_index(tmp_path) -> None: calibration_input = tmp_path / "calibration.json" eval_input = tmp_path / "eval.json" train_index = tmp_path / "train_index.json" test_index = tmp_path / "test_index.json" calibration_input.write_text("{}\n") eval_input.write_text("{}\n") train_index.write_text( '{"split":"train","content_hash":"train-content","split_hash":"train-split","retrieval_index_allowed":true}\n' ) test_index.write_text( '{"split":"test","content_hash":"test-content","split_hash":"test-split","retrieval_index_allowed":false}\n' ) args = SimpleNamespace( calibration_input=calibration_input, calibration_target_index=train_index, eval_input=eval_input, eval_target_index=test_index, source_index=None, checkpoint_template="runs/ctt_residual_base_context_obs_seed{seed}/model.pt", out_dir=tmp_path / "out", k=16, feature_set="chart_compat", selector_chart_feature_mode="base_context_obs", target="utility_margin", model_types="hgb_regressor", selection_frac=0.35, seed=0, bootstrap_samples=10, ) args.out_dir.mkdir() _write_nonlinear_provenance(args.out_dir, args) data_hash = json.loads((args.out_dir / "data_hash.txt").read_text()) split_hash = json.loads((args.out_dir / "split_hash.txt").read_text()) assert "source_index" in data_hash assert split_hash["source_index"]["split"] == "train" assert split_hash["source_index"]["retrieval_index_allowed"] is True