1#!/usr/bin/env python2"""Stage D: fit the value controller on logged counterfactual labels; calibrate on a held-out partition; tune3(lambda, kappa) on calibration ROLLOUTS (never on test); save the locked checkpoint.45Usage: python -m training.train_controller --config configs/experiments/controller_pilot.yaml [--device mps]6"""7from __future__ import annotations89import argparse, json, math, sys, time10from pathlib import Path1112import numpy as np13import torch14import torch.nn.functional as F1516sys.path.insert(0, str(Path(__file__).resolve().parents[1]))17from prisme.losses import controller_losses, token_nll # noqa: E40218from prisme.flops import action_cost_proxy # noqa: E40219from training.common import Run, load_checkpoint, load_yaml, make_adamw, pick_device, save_checkpoint, seed_all, sync # noqa: E40220from training.train_editor import build_model # noqa: E402212223def regenerate_state(model, cfg, rec, device):24 x = rec["tokens"].to(device).long()25 ctx = model.encode(x)26 S = ctx.S027 B, T = x.shape28 r_pre, R_total = rec["r_pre"], rec["R_total"]29 for r in range(r_pre):30 a = rec["prefix"][:, :, r].to(device).long()31 S, _ = model.apply_action(S, ctx, a, r, R_total - r - 1)32 step = torch.full((B, T), r_pre, device=device, dtype=torch.long)33 remaining = torch.full((B, T), R_total - r_pre - 1, device=device, dtype=torch.long)34 return ctx, S, step, remaining353637def preds_for_candidates(model, ctx, S, step, remaining, cand):38 p = model.controller(ctx.ctrl_in, S, step, remaining, ctx.entropy) # [B,T,A]39 return {k: v.gather(-1, cand) for k, v in p.items()}, p404142@torch.no_grad()43def evaluate_labels(model, cfg, recs, device):44 """Regression/classification diagnostics + oracle regret on audit batches."""45 model.eval()46 hub, sgn, brier, n = 0.0, 0.0, 0.0, 047 regret, opp, top1, n_aud = 0.0, 0.0, 0.0, 048 ece_bins = torch.zeros(10); ece_cnt = torch.zeros(10); ece_conf = torch.zeros(10)49 for rec in recs:50 ctx, S, step, remaining = regenerate_state(model, cfg, rec, device)51 cand = rec["cand"].to(device).long(); g = rec["gain"].to(device); v = rec["valid"].to(device)52 pc, pall = preds_for_candidates(model, ctx, S, step, remaining, cand)53 vv = v.unsqueeze(-1).expand_as(g)54 hub += float((F.huber_loss(pc["g_hat"], g, delta=0.1, reduction="none") * vv).sum()); n += int(vv.sum())55 sgn += float((((pc["g_hat"] > 0) == (g > 0)).float() * vv).sum())56 ph = torch.sigmoid(pc["harm_logit"]); harm = (g < 0).float()57 brier += float((((ph - harm) ** 2) * vv).sum())58 b = (ph * 10).clamp(max=9.999).long()[vv]59 ece_bins.index_add_(0, b.cpu(), harm[vv].cpu()); ece_conf.index_add_(0, b.cpu(), ph[vv].cpu()); ece_cnt.index_add_(0, b.cpu(), torch.ones_like(b, dtype=torch.float).cpu())60 if "all_gains" in rec:61 ag = rec["all_gains"].to(device) # [B,T,A]62 best = ag.max(-1).values.clamp(min=0) # one-step oracle value (STOP if nothing helps)63 q = pall["g_hat"]64 chosen = q.argmax(-1); chosen_q = q.max(-1).values65 realised = torch.where(chosen_q > 0, ag.gather(-1, chosen.unsqueeze(-1)).squeeze(-1), torch.zeros_like(best))66 regret += float(((best - realised) * v).sum()); opp += float((best * v).sum())67 top1 += float(((chosen == ag.argmax(-1)).float() * v).sum()); n_aud += int(v.sum())68 ece = float((ece_cnt / ece_cnt.sum().clamp(min=1) * (ece_bins / ece_cnt.clamp(min=1) - ece_conf / ece_cnt.clamp(min=1)).abs()).sum())69 out = {"huber": hub / max(n, 1), "sign_acc": sgn / max(n, 1), "harm_brier": brier / max(n, 1), "harm_ece": ece, "n_labels": n}70 if n_aud:71 out.update({"oracle_opportunity_per_state": opp / n_aud, "regret_per_state": regret / n_aud,72 "captured_fraction": 1 - regret / max(opp, 1e-9), "top1_match": top1 / n_aud, "n_audit_states": n_aud})73 return out747576@torch.no_grad()77def rollout_cost_quality(model, cfg, recs, device, cost, lam, kap, max_budget):78 """Deployed-controller rollout (teacher forced) on calibration batches: mean loss, mean edits, mean ops ratio."""79 tot, n, edits, ops = 0.0, 0, 0.0, 0.080 base_ops = 1.081 for rec in recs:82 x = rec["tokens"].to(device).long(); y = rec["targets"].to(device).long()83 out = model.forward_policy(x, max_budget, cost, lambda_cost=lam, kappa_risk=kap)84 nll = token_nll(out["logits"], y); v = (y != -100)85 tot += float(nll[v].sum()); n += int(v.sum()); edits += float(out["n_edits"][v].float().sum())86 if out["actions"] is not None:87 a = out["actions"][v] # [n, R]88 ops += float(torch.where(a < cfg.n_actions, cost[a.clamp(max=cfg.n_actions - 1)], torch.zeros_like(cost[0]).expand_as(a).float()).sum())89 return {"loss": tot / max(n, 1), "edits_per_token": edits / max(n, 1), "extra_ops_ratio": ops / max(n, 1)}909192def main(argv=None) -> int:93 ap = argparse.ArgumentParser(description=__doc__)94 ap.add_argument("--config", required=True)95 ap.add_argument("--device", default="auto")96 ap.add_argument("--seed", type=int, default=None)97 ap.add_argument("--run-root", default="experiments/runs")98 args = ap.parse_args(argv)99 exp = load_yaml(args.config)100 seed = args.seed if args.seed is not None else exp.get("seed", 1)101 seed_all(seed); device = pick_device(args.device)102 cfg, model = build_model(exp)103 model = model.to(device)104 meta = load_checkpoint(exp["checkpoint"], model, strict=False)105 for n_, p in model.named_parameters():106 p.requires_grad_(n_.startswith("controller"))107 train = torch.load(exp["labels"]["train"], weights_only=False)108 calib = torch.load(exp["labels"]["calib"], weights_only=False)109 assert train["checkpoint_sha256"] == calib["checkpoint_sha256"] == meta.get("weights_sha256"), "labels were generated by a different editor checkpoint (stale labels)"110 tr = exp["train"]111 run = Run(exp.get("name", "controller"), {"experiment": exp, "model": cfg.to_dict()}, root=args.run_root, seed=seed)112 run.write_manifest(device=str(device), checkpoint_parent=exp["checkpoint"], labels_train=exp["labels"]["train"], labels_calib=exp["labels"]["calib"])113 opt = make_adamw(model.controller.parameters(), tr["lr"], tr.get("weight_decay", 0.01))114 recs = train["records"]; rng = np.random.default_rng(seed)115 alpha, beta, gamma = tr.get("alpha", 1.0), tr.get("beta", 0.5), tr.get("gamma", 0.1)116 model.train()117 best_state, best_score, best_step = None, -1e9, -1118 def score_of(ev):119 return -ev["regret_per_state"] if "regret_per_state" in ev else ev["sign_acc"]120 for step in range(tr["steps"]):121 rec = recs[int(rng.integers(0, len(recs)))]122 with torch.no_grad():123 ctx, S, st, rem = regenerate_state(model, cfg, rec, device)124 cand = rec["cand"].to(device).long(); g = rec["gain"].to(device); v = rec["valid"].to(device).unsqueeze(-1).expand_as(g)125 pc, pall = preds_for_candidates(model, ctx, S, st, rem, cand)126 L = controller_losses({k: x_[v] for k, x_ in pc.items()}, g[v], torch.ones_like(g[v], dtype=torch.bool), tr.get("huber_delta", 0.1))127 loss = alpha * L["gain"] + beta * L["harm"] + gamma * L["nll"]128 # ---- within-state ranking losses (D12): pairwise on the 4 logged candidates, listwise on exhaustive audits129 w_pair, w_list = tr.get("pairwise", 0.0), tr.get("listwise", 0.0)130 if w_pair > 0:131 gi, gj = g.unsqueeze(-1), g.unsqueeze(-2) # [B,T,4,4]132 pi, pj = pc["g_hat"].unsqueeze(-1), pc["g_hat"].unsqueeze(-2)133 wgt = (gi - gj).abs() * rec["valid"].to(device).float()[..., None, None]134 lp = F.binary_cross_entropy_with_logits((pi - pj) / tr.get("rank_temp", 0.02), (gi > gj).float(), weight=wgt, reduction="sum") / wgt.sum().clamp(min=1e-6)135 loss = loss + w_pair * lp; L["pair"] = lp136 if w_list > 0 and "all_gains" in rec:137 ag = rec["all_gains"].to(device) # [B,T,A]138 tgt = torch.softmax(ag / tr.get("rank_temp", 0.02), -1)139 ll = -(tgt * torch.log_softmax(pall["g_hat"] / tr.get("rank_temp", 0.02), -1)).sum(-1)140 ll = (ll * rec["valid"].to(device).float()).sum() / rec["valid"].sum().clamp(min=1)141 loss = loss + w_list * ll; L["list"] = ll142 opt.zero_grad(set_to_none=True); loss.backward()143 torch.nn.utils.clip_grad_norm_(model.controller.parameters(), 1.0); opt.step()144 if step % tr.get("log_every", 50) == 0 or step == tr["steps"] - 1:145 run.log(step=step, loss=float(loss), gain=float(L["gain"]), harm=float(L["harm"]), nll=float(L["nll"]))146 print(f"step {step} loss {float(loss):.5f} gain {float(L['gain']):.5f} harm {float(L['harm']):.4f}", flush=True)147 if (step + 1) % tr.get("eval_every", 500) == 0 or step == tr["steps"] - 1:148 ev = evaluate_labels(model, cfg, calib["records"], device); model.train()149 run.log(step=step, **{f"calib_{k}": v_ for k, v_ in ev.items()}); print(" calib:", {k: round(v_, 4) for k, v_ in ev.items()}, flush=True)150 if score_of(ev) > best_score:151 best_score, best_step = score_of(ev), step + 1152 best_state = {k: v_.detach().cpu().clone() for k, v_ in model.controller.state_dict().items()}153 if best_state is not None: # early stopping on the calibration partition (never on test)154 model.controller.load_state_dict(best_state)155 run.log(event="restored_best_calibration_checkpoint", step=best_step, score=best_score)156 print(f"restored controller from step {best_step} (calibration score {best_score:.5f})", flush=True)157 model.eval()158 ev_train = evaluate_labels(model, cfg, recs[: min(len(recs), 40)], device)159 ev_calib = evaluate_labels(model, cfg, calib["records"], device)160 # ---- (lambda, kappa) sweep by calibration rollouts ------------------------------------------------161 cost = torch.tensor(action_cost_proxy(cfg, tr["seq_len"]), device=device)162 sweep = []163 cal_sub = calib["records"][: exp.get("sweep_batches", 10)]164 for lam in exp.get("lambda_grid", [0.0, 0.01, 0.03, 0.1]):165 for kap in exp.get("kappa_grid", [0.0, 0.5, 1.0]):166 r = rollout_cost_quality(model, cfg, cal_sub, device, cost, lam, kap, exp.get("max_budget", 8))167 sweep.append({"lambda": lam, "kappa": kap, **r}); print("sweep", sweep[-1], flush=True)168 b0 = rollout_cost_quality(model, cfg, cal_sub, device, cost, 0.0, 0.0, 0)169 sel_rule = exp.get("selection_rule", "max_loss_reduction_per_ops")170 for s in sweep:171 s["loss_reduction"] = b0["loss"] - s["loss"]172 s["reduction_per_ops"] = s["loss_reduction"] / max(s["extra_ops_ratio"], 1e-9)173 if sel_rule == "max_loss_reduction_per_ops":174 cands = [s for s in sweep if s["loss_reduction"] > 0] or sweep175 best = max(cands, key=lambda s: s["reduction_per_ops"])176 else:177 best = min(sweep, key=lambda s: s["loss"])178 cfg.lambda_cost, cfg.kappa_risk = best["lambda"], best["kappa"]179 sha = save_checkpoint(run.dir / "checkpoint", model, cfg.to_dict(), {"run_id": run.run_id, "parent": exp["checkpoint"], "selected": best, "locked": True})180 run.finish(eval_train=ev_train, eval_calib=ev_calib, sweep=sweep, baseline_budget0=b0, selected=best, checkpoint_sha256=sha, selection_rule=sel_rule)181 print("selected", best, "\nrun", run.run_id)182 return 0183184185if __name__ == "__main__":186 raise SystemExit(main())187