Files
aituner/runs/active-intervention-v0/train_policy.py

319 lines
12 KiB
Python

#!/usr/bin/env python3
"""Train and audit outcome-only versus telemetry action-response policies."""
from __future__ import annotations
import argparse
import hashlib
import importlib.util
import json
import math
import os
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any, Mapping, Sequence
HERE = Path(__file__).resolve().parent
def _load_model():
spec = importlib.util.spec_from_file_location(
"active_intervention_model", HERE / "model.py"
)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
MODEL = _load_model()
REGULARIZATION = 10.0
MINIMUM_MARGIN = 0.02
CONFIDENCE_Z = 1.0
def sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for chunk in iter(lambda: source.read(1 << 20), b""):
digest.update(chunk)
return digest.hexdigest()
def atomic_json(path: Path, payload: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
temporary.write_text(
json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
)
os.replace(temporary, path)
def grouped(
examples: Sequence[Mapping[str, Any]], key: str
) -> dict[str, list[Mapping[str, Any]]]:
result: dict[str, list[Mapping[str, Any]]] = defaultdict(list)
for example in examples:
result[str(example[key])].append(example)
return dict(result)
def evaluate_grouped_cv(
examples: Sequence[Mapping[str, Any]],
*,
include_telemetry: bool,
holdout_key: str,
) -> dict[str, Any]:
holdouts = grouped(examples, holdout_key)
decision_rows = []
for held_out, test_examples in sorted(holdouts.items()):
training = [example for example in examples if str(example[holdout_key]) != held_out]
if len({str(example["decision_id"]) for example in training}) < 2:
continue
model = MODEL.fit_ridge(
training,
include_telemetry=include_telemetry,
regularization=REGULARIZATION,
)
for decision_id, candidates in sorted(grouped(test_examples, "decision_id").items()):
predictions = []
for candidate in candidates:
source = candidate["source"]
action = candidate["action"]
noop = (
int(action["target_mns"]) == int(source["mns"])
and int(action["target_mbbt"]) == int(source["mbbt"])
)
if noop:
prediction = 0.0
else:
names, vector = MODEL.feature_vector(
candidate, include_telemetry=include_telemetry
)
if tuple(names) != model.feature_names:
raise ValueError("cross-validation feature schema mismatch")
prediction = max(-1.0, min(1.0, model.predict(vector)))
predictions.append(
{
"action_id": str(candidate["action"]["id"]),
"prediction": prediction,
"real": float(candidate["target_normalized_goodput"]),
}
)
predictions.sort(key=lambda row: (-row["prediction"], row["action_id"]))
selected = predictions[0]
oracle = max(row["real"] for row in predictions)
regret = 1.0 - selected["real"] / oracle if oracle > 0 else 0.0
best_actions = {
row["action_id"] for row in predictions if math.isclose(row["real"], oracle)
}
decision_rows.append(
{
"holdout": held_out,
"decision_id": decision_id,
"selected_action": selected["action_id"],
"best_actions": sorted(best_actions),
"correct": selected["action_id"] in best_actions,
"selected_real": selected["real"],
"oracle_real": oracle,
"regret": regret,
"predictions": predictions,
}
)
if not decision_rows:
return {"status": "INSUFFICIENT_GROUPS", "decisions": []}
regrets = [float(row["regret"]) for row in decision_rows]
return {
"status": "VALID",
"holdout_key": holdout_key,
"decision_n": len(decision_rows),
"correct_n": sum(bool(row["correct"]) for row in decision_rows),
"accuracy": sum(bool(row["correct"]) for row in decision_rows) / len(decision_rows),
"mean_regret": sum(regrets) / len(regrets),
"max_regret": max(regrets),
"decisions": decision_rows,
}
def paired_delta(outcome: Mapping[str, Any], telemetry: Mapping[str, Any]) -> dict[str, Any]:
if outcome.get("status") != "VALID" or telemetry.get("status") != "VALID":
return {"status": "INSUFFICIENT_GROUPS"}
outcome_by_id = {row["decision_id"]: row for row in outcome["decisions"]}
telemetry_by_id = {row["decision_id"]: row for row in telemetry["decisions"]}
common = sorted(set(outcome_by_id) & set(telemetry_by_id))
rows = []
for decision_id in common:
before = outcome_by_id[decision_id]
after = telemetry_by_id[decision_id]
rows.append(
{
"decision_id": decision_id,
"outcome_action": before["selected_action"],
"telemetry_action": after["selected_action"],
"action_changed": before["selected_action"] != after["selected_action"],
"regret_delta": float(after["regret"]) - float(before["regret"]),
"telemetry_corrected": (not before["correct"]) and bool(after["correct"]),
"telemetry_harmed": bool(before["correct"]) and (not after["correct"]),
}
)
return {
"status": "VALID",
"decision_n": len(rows),
"action_changed_n": sum(row["action_changed"] for row in rows),
"corrected_n": sum(row["telemetry_corrected"] for row in rows),
"harmed_n": sum(row["telemetry_harmed"] for row in rows),
"mean_regret_delta": (
sum(float(row["regret_delta"]) for row in rows) / len(rows) if rows else 0.0
),
"rows": rows,
}
def build_policy(dataset_path: Path) -> dict[str, Any]:
dataset = json.loads(dataset_path.read_text(encoding="utf-8"))
if dataset.get("status") != "VALID" or dataset["sanity"]["red_flags"]:
raise ValueError("training dataset is not valid")
examples = dataset["examples"]
phases = sorted({str(example["phase"]) for example in examples}, key=float)
phase_results = {}
incremental_candidates = []
for phase in phases:
selected = [example for example in examples if str(example["phase"]) == phase]
outcome_cv = evaluate_grouped_cv(
selected, include_telemetry=False, holdout_key="repetition"
)
telemetry_cv = evaluate_grouped_cv(
selected, include_telemetry=True, holdout_key="repetition"
)
outcome_regime = evaluate_grouped_cv(
selected, include_telemetry=False, holdout_key="regime"
)
telemetry_regime = evaluate_grouped_cv(
selected, include_telemetry=True, holdout_key="regime"
)
delta = paired_delta(outcome_cv, telemetry_cv)
outcome_models = MODEL.fit_jackknife_ensemble(
selected,
include_telemetry=False,
regularization=REGULARIZATION,
)
telemetry_models = MODEL.fit_jackknife_ensemble(
selected,
include_telemetry=True,
regularization=REGULARIZATION,
)
incremental = bool(
delta.get("status") == "VALID"
and int(delta["corrected_n"]) >= 1
and int(delta["harmed_n"]) == 0
and float(delta["mean_regret_delta"]) < -1e-12
and float(telemetry_cv["max_regret"]) <= 0.05
)
if incremental:
incremental_candidates.append(phase)
phase_results[phase] = {
"cutoff_s": float(selected[0]["cutoff_s"]),
"outcome_only": {
"leave_repetition_out": outcome_cv,
"leave_regime_out": outcome_regime,
"models": MODEL.models_to_json(outcome_models),
},
"telemetry": {
"leave_repetition_out": telemetry_cv,
"leave_regime_out": telemetry_regime,
"models": MODEL.models_to_json(telemetry_models),
},
"paired_incremental": delta,
"incremental_gate": incremental,
}
selected_phase = incremental_candidates[0] if incremental_candidates else phases[-1]
status = (
"RETROSPECTIVE_INCREMENTAL_SIGNAL"
if incremental_candidates
else "NO_RETROSPECTIVE_INCREMENTAL_SIGNAL"
)
target_values = [float(example["target_normalized_goodput"]) for example in examples]
effect_values = [
float(example["target_delta_normalized_goodput"]) for example in examples
]
invariants = {
"four_phases": len(phases) == 4,
"targets_bounded": all(0.0 <= value <= 1.0 for value in target_values),
"targets_not_all_identical": len(set(target_values)) > 1,
"effects_bounded": all(-1.0 <= value <= 1.0 for value in effect_values),
"effects_not_all_identical": len(set(effect_values)) > 1,
"models_present_every_phase": all(
phase_results[phase][mode]["models"]
for phase in phases
for mode in ("outcome_only", "telemetry")
),
}
red_flags = [name for name, passed in invariants.items() if not passed]
if red_flags:
raise RuntimeError(f"policy sanity failed: {red_flags}")
return {
"schema": "active-intervention-policy-v0",
"status": status,
"training": {
"dataset": str(dataset_path),
"dataset_sha256": sha256_file(dataset_path),
"examples": len(examples),
"decisions": len({example["decision_id"] for example in examples}),
"regularization": REGULARIZATION,
"confidence_z": CONFIDENCE_Z,
"minimum_margin": MINIMUM_MARGIN,
},
"measurement_policy": {
"selected_phase": selected_phase,
"selected_cutoff_s": phase_results[selected_phase]["cutoff_s"],
"selection_reason": (
"earliest phase passing the frozen incremental gate"
if incremental_candidates
else "no incremental phase; retain full measurement for exploratory held-out test"
),
},
"phases": phase_results,
"sanity": {
"invariants": invariants,
"red_flags": red_flags,
"target_normalized_goodput": {
"n": len(target_values),
"min": min(target_values),
"max": max(target_values),
"distinct_n": len(set(target_values)),
},
"target_delta_normalized_goodput": {
"n": len(effect_values),
"min": min(effect_values),
"max": max(effect_values),
"distinct_n": len(set(effect_values)),
},
},
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
policy = build_policy(args.dataset)
atomic_json(args.output, policy)
print(
json.dumps(
{
"status": policy["status"],
"measurement_policy": policy["measurement_policy"],
"sanity": policy["sanity"],
},
sort_keys=True,
)
)
if __name__ == "__main__":
main()