Buckets:
| """cwm3/train/train_trunk_warning_head.py — EXPERIMENT B (owner question 2026-09-18; pre-registration docs/decisions.md 2026-09-18T03:34:28Z): a Driver-WM-style | |
| 'warning head on the trunk' — a small MLP trained DIRECTLY on the detached trunk features of a frozen joint model, with the warranted-warning labels of the HGB read-out. | |
| python train/train_trunk_warning_head.py --run I1_all_minus_road2drv_s0_final_udwm --extract # step 1 (GPU): trunk caches, both regimes x test / val + TRAIN label rows | |
| python train/train_trunk_warning_head.py --run I1_all_minus_road2drv_s0_final_udwm --fit # step 2 (CPU): heads, gold, layouts | |
| Step 1 (results/runs/TRK_cache/trunks_<run>_{regime}_{split}.npz; never a result): the frozen BEST checkpoint forwards every row of the (regime, split) pack under the | |
| prediction path (bf16 autocast, no grad; train/train_cwm3_joint.prep_joint + CWM3Joint(return_trunks=True)) -> feat f16 [N, 640] = the DETACHED post-interaction trunk | |
| vector [z_v 256 ; h_u 128 ; z_e 128 ; z_vf 128], uid / t / mode witness; the TRAIN label rows are the `rows` of the run's pred_{regime}_train_labelrows.npz (the rows the | |
| HGB read-out is trained on). Witness: max |recomputed U-chain hazard − saved pred pack hazard| on every full split (expected <= one f16 ulp). | |
| Step 2: inputs = feat (f32) + the regime bit (641); labels = the 979 clean TRAIN consensus rows of data/labels.train_label_table (the SAME rows / order as the HGB read-out; | |
| the clean gold never enters a fit). HEAD (pre-registered) = train_warning_head.make_model("mlp", ...) = 0-imputation + StandardScaler + MLPClassifier(hidden (64,), relu, | |
| adam, alpha in GRID["mlp"] {1e-3, 1e-2, 1e-1}, early_stopping validation_fraction 0.15, max_iter 500, random_state --seed) — 641*64 + 64 + 64 + 1 = 41,153 parameters | |
| (asserted <= 50,000); selection = mean AUROC over the IDENTICAL device-grouped 5-fold CV (train_warning_head.cv_auroc, GroupKFold on dongle); refit on all labels; clean-val | |
| temperature (warning_gold.temperature_fit). L2 logistic regression (GRID["lr"]) = DESCRIPTIVE linear-probe control. Gold clean-test AUROC / PR-AUC with device-cluster CIs | |
| (warning_gold.cluster_boot) -> gold_quick.csv; layouts results/runs/W_TRK_<run>/pred_layout_{regime}_{split}.npz (arm abl_W_trk; columns W_trunk_mlp64_t15_w10_n979, | |
| W_trunk_lr_t15_w10_n979; w_mask all True; uid / t = the run's layout rows) for eval/alert_timing_eval.py --extra-scores and the paired gold script. | |
| Nothing here reads an openpilot DMS field; no gradient reaches any trunk (features are cached numpy arrays). | |
| """ | |
| from __future__ import annotations | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| os.environ.setdefault("MKL_NUM_THREADS", "4") | |
| import argparse # noqa: E402 | |
| import copy # noqa: E402 | |
| import json # noqa: E402 | |
| import sys # noqa: E402 | |
| import time # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| from types import SimpleNamespace # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import pandas as pd # noqa: E402 | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train", CWM3 / "eval", EXPLORE / "s12_wm"): | |
| if str(p) not in sys.path: | |
| sys.path.insert(0, str(p)) | |
| from runmeta import write_run_json, code_hash # noqa: E402 | |
| RUNS = CWM3 / "results" / "runs" | |
| CACHE = RUNS / "TRK_cache" | |
| REGIMES = ("manual", "engaged") | |
| SPLITS = ("test", "val") | |
| TRUNK_KEYS = ("z_v", "h_u", "z_e", "z_vf") | |
| D_FEAT = 640 | |
| TAU, WIN = 15, 10 | |
| COLS = {"mlp": f"W_trunk_mlp64_t{TAU}_w{WIN}_n979", "lr": f"W_trunk_lr_t{TAU}_w{WIN}_n979"} | |
| MAX_PARAMS = 50_000 | |
| PREREG = "2026-09-18T03:34:28Z" | |
| CODE_FILES = [CWM3 / "train" / "train_trunk_warning_head.py", CWM3 / "train" / "train_warning_head.py", CWM3 / "models" / "cwm3.py", CWM3 / "train" / "batch.py", | |
| CWM3 / "train" / "train_cwm3_joint.py"] | |
| def log(msg, out=None): | |
| line = f"[{time.strftime('%H:%M:%S')}] {msg}" | |
| print(line, flush=True) | |
| if out: | |
| Path(out).mkdir(parents=True, exist_ok=True) | |
| with open(Path(out) / "train.log", "a") as f: | |
| f.write(line + "\n") | |
| # ----------------------------------------------------------------------------- pure helpers (unit-tested without data) | |
| def assemble(feat, mode): | |
| """cached trunk features f16 [N, 640] + regime bit [N] -> the head's input f32 [N, 641] (trunk order z_v, h_u, z_e, z_vf; the regime bit LAST).""" | |
| feat = np.asarray(feat) | |
| assert feat.ndim == 2 and feat.shape[1] == D_FEAT, feat.shape | |
| return np.concatenate([feat.astype(np.float32), np.asarray(mode, np.float32).reshape(-1, 1)], 1) | |
| def mlp_param_count(clf): | |
| """trainable parameters of a fitted sklearn MLPClassifier (weights + biases).""" | |
| return int(sum(w.size for w in clf.coefs_) + sum(b.size for b in clf.intercepts_)) | |
| def posmap_rows(cache_rows, wanted_rows): | |
| """positions inside a label-row cache (its `rows` = targets-file rows) for the wanted targets-file rows (every wanted row must be cached).""" | |
| pm = {int(r): i for i, r in enumerate(np.asarray(cache_rows))} | |
| return np.array([pm[int(r)] for r in wanted_rows], np.int64) | |
| def write_layout(path, uid, t, cols, meta): | |
| """results/runs/W_TRK_<run>/pred_layout_{regime}_{split}.npz in the W-layout contract (meta json, uid, t, w_names, w f16, w_mask all True).""" | |
| wn = list(cols) | |
| w = np.stack([np.asarray(cols[k], np.float32) for k in wn], 1).astype(np.float16) | |
| assert w.shape == (len(uid), len(wn)) and len(t) == len(uid) | |
| np.savez_compressed(path, meta=np.array(json.dumps(meta)), uid=np.asarray(uid), t=np.asarray(t), w_names=np.array(wn), w=w, w_mask=np.ones(len(uid), bool)) | |
| return path | |
| def cache_path(run, regime, split): | |
| return CACHE / f"trunks_{run}_{regime}_{split}.npz" | |
| # ----------------------------------------------------------------------------- step 1: extraction (GPU) | |
| def load_frozen(run_name, device, work_out): | |
| """the frozen joint run -> ctx (model in eval on device, cfg / stats / batch cfgs / bank); every data-layer log line goes to work_out.""" | |
| import train_cwm3_joint as TJ | |
| import cwm3_dataset as ds | |
| run = RUNS / run_name | |
| assert (run / "TRAIN_DONE").exists() and (run / "PRED_DONE").exists(), f"{run}: not a finished joint run" | |
| c = json.loads((run / "config.json").read_text()) | |
| cfg, a = c["cfg"], c["args"] | |
| args_model = SimpleNamespace(**a) | |
| args_model.out = str(run) | |
| args_model.cpu, args_model.no_amp, args_model.synthetic, args_model.max_anchors = False, False, 0, 0 | |
| args_data = copy.copy(args_model) | |
| args_data.out = str(work_out) | |
| ds.set_pack_version(cfg.get("pack_version", "v2")) | |
| ds.set_kin_source(cfg["kin_source"], cfg["kin_version"]) | |
| if cfg.get("u_kpf"): | |
| ds.set_kpf_dir(cfg["u_kpf"]) | |
| info = TJ.base_info_numpy(cfg, int(a["seed"])) | |
| st = {k: v for k, v in np.load(run / "stats.npz").items()} | |
| bank = TJ.load_bank(args_model, cfg, info) | |
| model, best = TJ.load_best_model(args_model, cfg, device, info) | |
| for p in model.parameters(): | |
| p.requires_grad_(False) | |
| pc = TJ.apply_mask_v(TJ.prep_cfgs(info), cfg) | |
| return {"model": model, "cfg": cfg, "st": st, "pc": pc, "bank": bank, "args": args_data, "info": info, "best": best, "model_id": f"{run_name}:{best}", "run": run, | |
| "TJ": TJ, "seed": int(a["seed"])} | |
| def extract_rows(ctx, raw, device): | |
| """raw dict (one regime) -> feat f16 [N, 640], hazards of the two U chains f32 [N, 25, 2, 2] (on / off) for the witness.""" | |
| import torch | |
| from models import heads as hd | |
| TJ, model, st, pc = ctx["TJ"], ctx["model"], ctx["st"], ctx["pc"] | |
| N = len(raw["t"]) | |
| store = TJ.make_store(raw, device, ctx["bank"]) | |
| region_keys = model.tok_embed is not None | |
| feat = np.zeros((N, D_FEAT), np.float16) | |
| hz = np.zeros((N, 25, 2, 2), np.float32) | |
| bs = 4096 | |
| with torch.no_grad(): | |
| for i in range(0, N, bs): | |
| idx = torch.arange(i, min(i + bs, N), device=device) | |
| sl = slice(i, i + len(idx)) | |
| b, _, _ = TJ.prep_joint(store, idx, st, pc, False, region_keys) | |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(device == "cuda")): | |
| o = model(b, return_trunks=True) | |
| tr = o["trunks"] | |
| feat[sl] = torch.cat([tr[k].float() for k in TRUNK_KEYS], 1).cpu().numpy().astype(np.float16) | |
| u = hd.split_U(o["U"]) | |
| hon, hoff = hd.hazards(u["chains"]) | |
| hz[sl] = torch.stack([hon, hoff], -1).float().cpu().numpy() | |
| del store | |
| if device == "cuda": | |
| torch.cuda.empty_cache() | |
| return feat, hz | |
| def extract(run_name, device, out, only=None): | |
| import cwm3_dataset as ds | |
| ctx = load_frozen(run_name, device, out) | |
| TJ = ctx["TJ"] | |
| log(f"frozen {ctx['model_id']} loaded on {device}; mask_V {bool(ctx['cfg'].get('mask_V', False))}", out) | |
| CACHE.mkdir(parents=True, exist_ok=True) | |
| witness = {} | |
| jobs = [(r, s) for s in SPLITS for r in REGIMES] + [(r, "train_labelrows") for r in REGIMES] | |
| for regime, split in jobs: | |
| if only and f"{regime}_{split}" not in only: | |
| continue | |
| pz = cache_path(run_name, regime, split) | |
| if pz.exists(): | |
| log(f"cache exists: {pz.name}", out) | |
| continue | |
| t0 = time.time() | |
| args = copy.copy(ctx["args"]) | |
| if split == "train_labelrows": | |
| lp = ctx["run"] / f"pred_{regime}_train_labelrows.npz" | |
| with np.load(lp, allow_pickle=True) as z: | |
| rows, uid_l, t_l = z["rows"].astype(np.int64), z["uid"], z["t"] | |
| full = TJ.load_raw(args, ctx["cfg"], "train", subset=False, regimes=[regime]) | |
| raw = ds.subset(full, rows) | |
| del full | |
| assert np.array_equal(raw["uid"], uid_l) and np.allclose(raw["t"], t_l), f"{lp}: label-row uid / t witness" | |
| extra = {"rows": rows} | |
| else: | |
| raw = TJ.load_raw(args, ctx["cfg"], split, subset=False, regimes=[regime]) | |
| extra = {} | |
| feat, hz = extract_rows(ctx, raw, device) | |
| pk = ctx["run"] / (f"pred_{regime}_{split}.npz" if split != "train_labelrows" else f"pred_{regime}_train_labelrows.npz") | |
| w = {} | |
| if pk.exists(): | |
| with np.load(pk, allow_pickle=True) as z: | |
| assert len(z["hz_on"]) == len(raw["t"]) | |
| w = {"max_abs_diff_hz_on_U": float(np.abs(z["hz_on"][:, :, 4:6].astype(np.float32) - hz[..., 0]).max()), | |
| "max_abs_diff_hz_off_U": float(np.abs(z["hz_off"][:, :, 4:6].astype(np.float32) - hz[..., 1]).max())} | |
| witness[f"{regime}_{split}"] = w | |
| np.savez(pz, feat=feat, uid=raw["uid"], t=raw["t"], mode=raw["mode"].astype(np.int8), dongle=raw["dongle"], model_id=np.array(ctx["model_id"]), | |
| trunk_keys=np.array(TRUNK_KEYS), witness=np.array(json.dumps(w)), **extra) | |
| log(f"extracted {regime}/{split}: {len(raw['t']):,} rows -> {pz.name} in {time.time() - t0:.0f}s; witness {w}", out) | |
| del raw, feat, hz | |
| (out / "EXTRACT_DONE").write_text(json.dumps({"run": run_name, "model_id": ctx["model_id"], "witness": witness}, indent=1)) | |
| return witness | |
| # ----------------------------------------------------------------------------- step 2: heads (CPU) | |
| def fit(run_name, out, n_boot, seed, val_recall): | |
| import train_warning_head as TW | |
| WG, LB, WJ = TW.WG, TW.LB, TW.WJ | |
| run = RUNS / run_name | |
| run_cfg = json.loads((run / "config.json").read_text()) | |
| seed_run = int(run_cfg["args"]["seed"]) | |
| base_ckpts = json.loads((run / "base_ckpts.json").read_text()) if (run / "base_ckpts.json").exists() else {} | |
| t0 = time.time() | |
| # ---- training table (the HGB read-out's rows and order) | |
| tab = LB.train_label_table(sources=("consensus",)) | |
| Xs, ys, gs = [], [], [] | |
| for regime in REGIMES: | |
| with np.load(cache_path(run_name, regime, "train_labelrows"), allow_pickle=True) as z: | |
| sub = tab[tab.regime == regime] | |
| pos = posmap_rows(z["rows"], sub.row.to_numpy().astype(np.int64)) | |
| Xs.append(assemble(z["feat"][pos], z["mode"][pos])) | |
| ys.append(sub.y.to_numpy().astype(int)) | |
| gs.append(sub.dongle.to_numpy()) | |
| X, y, groups = np.concatenate(Xs), np.concatenate(ys), np.concatenate(gs) | |
| assert len(y) == 979 and X.shape[1] == D_FEAT + 1, (X.shape, len(y)) | |
| log(f"train labels: {len(y)} rows / {int(y.sum())} pos / {len(np.unique(groups))} devices; inputs {X.shape[1]} (trunk {D_FEAT} + regime bit)", out) | |
| # ---- gold anchors (evaluation / temperature only) | |
| EP = LB.episode_labels() | |
| lab = EP[EP.label != "unsure"].reset_index(drop=True) | |
| packs = {(r, s): WG.load_pack(r, s) for r in REGIMES for s in SPLITS} | |
| legacy = WG.merge_legacy(seed_run) | |
| G = WG.build_gold(lab, packs, legacy, "legacy_first") | |
| assert not set(G[G.clean_test].idx) & set(tab.idx), "gold test episode inside the training table" | |
| gt, gv = G.clean_test.to_numpy(), G.clean_val.to_numpy() | |
| yg, dong_g = G.y.to_numpy().astype(int), G.dongle.to_numpy() | |
| full = {} | |
| for (r, s) in packs: | |
| with np.load(cache_path(run_name, r, s), allow_pickle=True) as z: | |
| assert np.array_equal(z["uid"], packs[(r, s)]["uid"]), f"{r}/{s}: cache rows != targets rows" | |
| full[(r, s)] = (assemble(z["feat"], z["mode"]), z["uid"], z["t"]) | |
| Xg = np.full((len(G), X.shape[1]), np.nan, np.float32) | |
| for (r, s), g in G.groupby(["regime", "split"]): | |
| Xg[g.index.to_numpy()] = full[(r, s)][0][g.row.to_numpy().astype(np.int64)] | |
| log(f"gold: clean_test {int(gt.sum())} rows / {int(yg[gt].sum())} pos / {len(np.unique(dong_g[gt]))} dev; clean_val {int(gv.sum())} / {int(yg[gv].sum())}", out) | |
| # ---- CV selection (the read-out's GroupKFold), refit, gold, temperature, layouts | |
| cv_rows, gold_rows, selected, fits = [], [], {}, {} | |
| for kind in ("mlp", "lr"): | |
| best = None | |
| for hp0 in TW.GRID[kind]: | |
| hp = dict(hp0) | |
| m, folds = TW.cv_auroc(kind, hp, X, y, groups, 5, seed) | |
| cv_rows.append({"model": kind, "hp": json.dumps(hp0), "cv_auroc": m, "folds": json.dumps([round(f, 4) for f in folds])}) | |
| log(f" cv {kind} {hp0}: AUROC {m:.4f}", out) | |
| if best is None or m > best[0]: | |
| best = (m, hp) | |
| fits[kind] = TW.make_model(kind, best[1], seed).fit(X, y) | |
| n_par = mlp_param_count(fits[kind][-1]) if kind == "mlp" else int(fits[kind][-1].coef_.size + fits[kind][-1].intercept_.size) | |
| if kind == "mlp": | |
| assert n_par <= MAX_PARAMS, f"MLP head has {n_par} parameters > {MAX_PARAMS}" | |
| selected[kind] = {"hp": best[1], "cv_auroc": best[0], "n_params": n_par, "n_inputs": int(X.shape[1]), "column": COLS[kind], "role": "pre-registered head" if kind == "mlp" else "descriptive linear probe"} | |
| s_all = fits[kind].predict_proba(Xg)[:, 1] | |
| au, ci = TW.boot_auroc(yg[gt], s_all[gt], dong_g[gt], n_boot, 0) | |
| b = WG.cluster_boot(yg[gt], {"s": s_all[gt]}, dong_g[gt], n_boot, 0)["s"]["prauc"] | |
| au_v = WG.fast_auroc(yg[gv], s_all[gv]) if yg[gv].min() != yg[gv].max() else np.nan | |
| Tm, pinned = WG.temperature_fit(s_all[gv], yg[gv].astype(float)) | |
| gold_rows.append({"model": kind, "column": COLS[kind], "hp": json.dumps(best[1]), "cv_auroc": best[0], "n_params": n_par, "gold_test_auroc": au, "ci_lo": ci[0], "ci_hi": ci[1], | |
| "gold_test_prauc": WG.fast_ap(yg[gt], s_all[gt]), "prauc_ci_lo": WG.pct(b)[0], "prauc_ci_hi": WG.pct(b)[1], "n": int(gt.sum()), "n_pos": int(yg[gt].sum()), | |
| "clean_val_auroc": au_v, "temperature": Tm, "T_pinned": pinned, "n_labels": int(len(y)), "run": run_name, "model_id": ""}) | |
| log(f" {kind}: {n_par:,} params; gold clean_test AUROC {au:.4f} [{ci[0]:.3f}, {ci[1]:.3f}] PR-AUC {gold_rows[-1]['gold_test_prauc']:.3f} (val {au_v:.3f}, T {Tm:.2f})", out) | |
| pd.DataFrame(cv_rows).to_csv(out / "cv.csv", index=False) | |
| pd.DataFrame(gold_rows).to_csv(out / "gold_quick.csv", index=False) | |
| (out / "selected.json").write_text(json.dumps({"run": run_name, "seed_run": seed_run, "heads": selected, "n_labels": int(len(y)), "n_pos": int(y.sum()), "base_ckpts": base_ckpts, | |
| "prereg": PREREG, "trunk_keys": list(TRUNK_KEYS), "d_feat": D_FEAT}, indent=1)) | |
| # ---- layouts on every test / val row | |
| wout = RUNS / f"W_TRK_{run_name}" | |
| wout.mkdir(parents=True, exist_ok=True) | |
| layouts = [] | |
| for (r, s), (Xrs, uid, t) in full.items(): | |
| lay = np.load(run / f"pred_layout_{r}_{s}.npz", allow_pickle=True) | |
| assert np.array_equal(lay["uid"], uid) and np.allclose(lay["t"].astype(float), np.asarray(t, float)), f"{r}/{s}: run layout rows != cache rows" | |
| meta = json.loads(str(lay["meta"])) | |
| meta.update({"arm": "abl_W_trk", "score_source": "independent_marginals", "model_id": f"W_TRK_{run_name}:" + "/".join(f"{k}={COLS[k]}" for k in fits), | |
| "warning_head": {"kind": "trunk MLP head (Experiment B)", "heads": selected, "labels": "n979", "n_labels": int(len(y)), "base_ckpts": base_ckpts.get("base_ckpts", {}), | |
| "inputs": "detached post-interaction trunk [z_v; h_u; z_e; z_vf] + regime bit", "prereg": PREREG}, | |
| "code_sha256": code_hash(CODE_FILES), | |
| "notes": "Experiment B: small MLP (<= 50k params) on the frozen joint model's detached trunk features, trained on the 979 clean TRAIN consensus labels with the HGB read-out's folds"}) | |
| cols = {COLS[k]: fits[k].predict_proba(Xrs)[:, 1].astype(np.float32) for k in fits} | |
| layouts.append(str(write_layout(wout / f"pred_layout_{r}_{s}.npz", lay["uid"], lay["t"], cols, meta))) | |
| log(f" wrote {layouts[-1]} ({len(uid)} rows, {len(cols)} columns)", out) | |
| write_run_json(out, "train_trunk_warning_head", CODE_FILES, [cache_path(run_name, r, s) for r in REGIMES for s in SPLITS + ("train_labelrows",)], seed=seed, t_start=t0, | |
| input_listing=False, extra={"run": run_name, "seed_run": seed_run, "heads": selected, "gold_quick": gold_rows, "layout_dir": str(wout), "layouts": layouts, | |
| "n_labels": int(len(y)), "n_pos": int(y.sum()), "prereg": PREREG, "dms_fields_used": []}) | |
| (out / "FIT_DONE").write_text(json.dumps({"layout_dir": str(wout), "heads": selected}, indent=1)) | |
| log(f"DONE in {time.time() - t0:.0f}s -> {out}", out) | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--run", required=True, help="frozen joint run dir name under results/runs") | |
| ap.add_argument("--extract", action="store_true") | |
| ap.add_argument("--fit", action="store_true") | |
| ap.add_argument("--only", default="", help="extract only these comma-separated <regime>_<split> jobs (debug)") | |
| ap.add_argument("--n-boot", type=int, default=1000) | |
| ap.add_argument("--seed", type=int, default=0, help="head random_state (the read-out's default 0)") | |
| ap.add_argument("--val-recall", type=float, default=0.5) | |
| ap.add_argument("--out-root", default=str(CWM3 / "results" / "warning_head")) | |
| ap.add_argument("--cpu", action="store_true", help="tests only") | |
| a = ap.parse_args() | |
| stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()) | |
| out = Path(a.out_root) / f"TRK_{a.run}_{stamp}" | |
| out.mkdir(parents=True, exist_ok=True) | |
| t0 = time.time() | |
| if a.extract: | |
| import torch | |
| torch.set_num_threads(4) | |
| device = "cuda" if torch.cuda.is_available() and not a.cpu else "cpu" | |
| w = extract(a.run, device, out, only=set(a.only.split(",")) if a.only else None) | |
| write_run_json(out, "trunk_extract", CODE_FILES, [], seed=0, t_start=t0, extra={"run": a.run, "witness": w, "cache": str(CACHE), "device": device, "prereg": PREREG, "dms_fields_used": []}) | |
| if a.fit: | |
| fit(a.run, out, a.n_boot, a.seed, a.val_recall) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 20.3 kB
- Xet hash:
- 72cb39798e423c8174108547520e9fe3a38fded6fd46cc49947060839e646d7f
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.