Buckets:
| """cwm3/eval/shift_twin.py — K5 identity-leak check for a pretrained-init U0 run vs its scratch twin (docs/weak_teacher.md §4). | |
| A = fine-tuned-from-pretrained-trunk run dir, B = scratch run dir (same seed, same recipe). Three statistics, all CPU except the latent | |
| extraction (GRU forward over a seeded train-row subsample, GPU if available): | |
| 1. Shifted-twin pair-AUROC (test split): for every positive row r of the event population (inactive at the anchor, EM & M valid on the | |
| first 15 steps, first onset < 3 s — kin_predictability.metric_rows) pick one twin r' on the SAME route with Δt ∈ [30, 60] s, y' = 0, | |
| |vEgo(t) − vEgo(t')| < 2 m/s, same regime (seed 0). Twins share camera / identity exactly and differ in behaviour; a model that | |
| encodes identity only scores them equally -> pair-AUROC 0.5. Score = 1 − Π_{k<15}(1 − h_k) from pred_{regime}_test.npz. | |
| Reported for A and B with a device-cluster bootstrap on the paired difference (n_boot 1000, seed 0). | |
| 2. Identity probe (train split): multinomial logistic regression from the frozen 128-d GRU latent to the device id, GroupKFold(5) by | |
| route, on a seeded subsample; accuracy for A and B vs the majority-class rate. Verdict flag: A − B > 5 points. | |
| 3. Device-mean-centred behaviour probe (train split): logistic regression latent -> y (first onset of the event within 3 s) with and | |
| without subtracting the per-device latent mean (GroupKFold by route); if the A − B PR-AUC gain vanishes under centring the gain | |
| rides on device identity. | |
| Output: results/shift_twin/<tag>_<regime>.json + a row per statistic appended to results/shift_twin.csv; run json. No DMS field, no | |
| weak/ import: reads model checkpoints, prediction files, the s12 / cwm2 kin target files. | |
| python eval/shift_twin.py --a results/runs/U0_kin_can_noK_dmspre_s0 --b results/runs/U0_kin_can_noK_s0 --regime engaged --tag dmspre_s0 | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import csv | |
| import json | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| CWM2 = EXPLORE / "cwm2" | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| KIN = CWM2 / "results" / "targets_kin" | |
| for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train"): | |
| sys.path.insert(0, str(p)) | |
| import cwm3_dataset as ds # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| from models import build_branch # noqa: E402 | |
| from runmeta import write_run_json # noqa: E402 | |
| from models import identity as idn # noqa: E402 (2026-09-11 opt-in: devnorm prep / extra store keys of the identity arms) | |
| import train_cwm3 as TC # noqa: E402 (model_kwargs: rebuilds a run's module layout incl. in_ln / adv head) | |
| KINC = CWM3 / "results" / "targets_kin_causal_routes" # K4 causal packs (opt-in --kin-source causal_routes --kin-version v1|v2) | |
| OUT = CWM3 / "results" / "shift_twin" | |
| CSV = CWM3 / "results" / "shift_twin.csv" | |
| EVENTS = {"hand_return": ("hands_on", "on"), "hand_leave": ("hands_on", "off"), "head_turn_onset": ("head_turned", "on")} | |
| TAU = 15 | |
| COLS = ["run_stamp", "tag", "regime", "event", "statistic", "a", "b", "value_a", "value_b", "delta", "delta_ci_lo", "delta_ci_hi", "n", "n_pairs_or_pos", "n_dev", "note"] | |
| def kin_population(regime, split, event, kin_source="cwm2", kin_version="v1"): | |
| """(inact0 & EM valid on first TAU steps) rows, y = first onset < TAU; plus dongle, uid, t, vEgo at the anchor. | |
| kin_source causal_routes (2026-09-11 opt-in): the K4 causal packs (S_c / S0_c / EM_c, rows without a causal box censored) — the targets the causal U0 runs are scored on.""" | |
| s = np.load(S12R / f"targets_{regime}_{split}.npz", allow_pickle=True) | |
| M = s["M"].astype(bool) | |
| if kin_source == "causal_routes": | |
| kdir = KINC if kin_version == "v1" else Path(str(KINC) + f"_{kin_version}") | |
| d = np.load(kdir / f"kin_causal_{regime}_{split}.npz", allow_pickle=True) | |
| assert np.array_equal(s["uid"], d["uid"]) and np.allclose(s["t"], d["t"]) | |
| ok = d["row_ok"].astype(bool) | |
| S, S0 = d["S_c"].astype(np.float32), d["S0_c"].astype(np.float32) | |
| S[~ok], S0[~ok] = np.nan, np.nan | |
| EM = d["EM_c"].astype(bool) & ok[:, None, None] & M[:, :, None] | |
| else: | |
| d = np.load(KIN / f"kin_{regime}_{split}.npz", allow_pickle=True) | |
| assert np.array_equal(s["uid"], d["uid"]) and np.allclose(s["t"], d["t"]) | |
| S, S0, EM = d["S"].astype(np.float32), d["S0"].astype(np.float32), d["EM"] & M[:, :, None] | |
| j = list(EVENTS).index(event) | |
| z0, z, em = [(S0[:, 0], S[:, :, 0], EM[:, :, 0]), (1 - S0[:, 0], 1 - S[:, :, 0], EM[:, :, 1]), (S0[:, 1], S[:, :, 1], EM[:, :, 2])][j] | |
| zp = np.concatenate([z0[:, None], z[:, :-1]], 1) | |
| onset = (np.nan_to_num(z, nan=0.0) > 0.5) & (np.nan_to_num(zp, nan=1.0) < 0.5) & em | |
| first = np.where(onset.any(1), onset.argmax(1), -1) | |
| inact0 = np.nan_to_num(zp[:, 0], nan=1.0) < 0.5 | |
| sel = np.flatnonzero(inact0 & em[:, :TAU].all(1)) | |
| y = (first[sel] >= 0) & (first[sel] < TAU) | |
| return {"sel": sel, "y": y, "dongle": s["dongle"][sel], "uid": s["uid"][sel], "t": s["t"][sel].astype(np.float64), "vego": s["X"][sel, -1, 3].astype(np.float32)} | |
| def onset_score(run_dir, regime, split, event, sel): | |
| """1 − Π_{k<TAU}(1 − h_k) from the raw prediction pack; when that pack has been archived (shipping watcher keeps only the layouts) the SAME quantity is read | |
| from pred_layout_{regime}_{split}.npz onset_p at tau = TAU (verified identical to f16 precision, max |Δ| 2.4e-4 on U0_kin_can_noK_s0_kcv2_clip8_ln_adv01).""" | |
| pz = Path(run_dir) / f"pred_{regime}_{split}.npz" | |
| if pz.exists(): | |
| z = np.load(pz, allow_pickle=True) | |
| chains = list(map(str, z["chain_names"])) | |
| ch, kind = EVENTS[event] | |
| c = chains.index(ch) | |
| h = (z["hz_off"] if kind == "off" else z["hz_on"])[sel][:, :TAU, c].astype(np.float64) | |
| return 1.0 - np.prod(1.0 - h, 1) | |
| sys.path.insert(0, str(CWM2 / "eval")) | |
| from pred_layout import load_pred | |
| tg = np.load(S12R / f"targets_{regime}_{split}.npz", allow_pickle=True) | |
| p = load_pred(Path(run_dir) / f"pred_layout_{regime}_{split}.npz", tg) | |
| e, ti = list(map(str, p["event_names"])).index(event), list(map(int, p["onset_taus"])).index(TAU) | |
| return np.nan_to_num(p["onset_p"][sel, ti, e].astype(np.float64), nan=0.0) | |
| def make_twins(pop, seed=0, dt_lo=30.0, dt_hi=60.0, dv=2.0): | |
| """for each positive row one negative twin on the same route, |Δt| in [dt_lo, dt_hi], |ΔvEgo| < dv. -> (pos_idx, neg_idx) into pop arrays.""" | |
| rng = np.random.default_rng(seed) | |
| order = np.argsort(pop["uid"], kind="stable") | |
| uid_s = pop["uid"][order] | |
| bounds = np.r_[0, np.flatnonzero(uid_s[1:] != uid_s[:-1]) + 1, len(order)] | |
| pos_i, neg_i = [], [] | |
| for a_, b_ in zip(bounds[:-1], bounds[1:]): | |
| rows = order[a_:b_] | |
| y = pop["y"][rows] | |
| P, Nn = rows[y], rows[~y] | |
| if not len(P) or not len(Nn): | |
| continue | |
| tN, vN = pop["t"][Nn], pop["vego"][Nn] | |
| for r in P: | |
| dt = np.abs(tN - pop["t"][r]) | |
| ok = (dt >= dt_lo) & (dt <= dt_hi) & (np.abs(vN - pop["vego"][r]) < dv) | |
| if ok.any(): | |
| cand = Nn[ok] | |
| pos_i.append(r) | |
| neg_i.append(cand[rng.integers(0, len(cand))]) | |
| return np.array(pos_i, np.int64), np.array(neg_i, np.int64) | |
| def pair_auroc(sp, sn): | |
| return float(np.mean((sp > sn) + 0.5 * (sp == sn))) | |
| def boot_pairs(sa_p, sa_n, sb_p, sb_n, dongle, n_boot=1000, seed=0): | |
| uniq, dev = np.unique(dongle, return_inverse=True) | |
| rng = np.random.default_rng(seed) | |
| d = np.full(n_boot, np.nan) | |
| ia, ib = (sa_p > sa_n) + 0.5 * (sa_p == sa_n), (sb_p > sb_n) + 0.5 * (sb_p == sb_n) | |
| for i in range(n_boot): | |
| cnt = np.bincount(rng.integers(0, len(uniq), len(uniq)), minlength=len(uniq)).astype(np.float64) | |
| w = cnt[dev] | |
| if w.sum() <= 0: | |
| continue | |
| d[i] = (w * ia).sum() / w.sum() - (w * ib).sum() / w.sum() | |
| return float(np.nanpercentile(d, 2.5)), float(np.nanpercentile(d, 97.5)), int(len(uniq)) | |
| def latents(run_dir, raw, st, device, bs=8192): | |
| """frozen GRU final hidden state [N,128] of the run's BEST checkpoint over raw rows (U inputs only).""" | |
| import torch | |
| cfg = json.loads((Path(run_dir) / "config.json").read_text())["cfg"] | |
| seed = int(json.loads((Path(run_dir) / "config.json").read_text())["args"]["seed"]) | |
| kw = TC.model_kwargs(cfg) if "recipe" in cfg else dict(cfg["model"], cabin=cfg.get("cabin", "none")) # identical kwargs for the old runs; adds in_ln / adv head for the new ones | |
| model = build_branch("U", seed, **kw).to(device) | |
| best = (Path(run_dir) / "BEST").read_text().strip() | |
| z = torch.load(Path(run_dir) / best, map_location=device, weights_only=False) | |
| model.load_state_dict(z["model"]) | |
| model.eval() | |
| store = bt.GpuStore(raw, ["X", "Kh", "mode", "slice_tag"] + idn.store_keys(cfg), device) | |
| prep = idn.prep_for(cfg) # prep_U, or the devnorm prep when the run's recipe asks for it | |
| N = len(raw["t"]) | |
| out = np.zeros((N, model.gru.hidden_size), np.float32) | |
| with torch.no_grad(): | |
| for i in range(0, N, bs): | |
| idx = torch.arange(i, min(i + bs, N), device=device) | |
| b, _ = prep(store, idx, st, cfg, need_targets=False) | |
| kx = b["Kx"].float() | |
| if getattr(model, "in_ln", None) is not None: # the trunk entry LayerNorm of the `u_in_ln` recipe (part of the GRU input path) | |
| kx = torch.cat([model.in_ln(kx[..., :model.N_KIN_IN]), kx[..., model.N_KIN_IN:]], -1) | |
| _, h = model.gru(kx) | |
| out[i:i + len(idx)] = h[-1].float().cpu().numpy() | |
| return out | |
| def probe_identity(Z, dongle, uid, seed=0): | |
| from sklearn.linear_model import LogisticRegression | |
| from sklearn.model_selection import GroupKFold | |
| from sklearn.preprocessing import StandardScaler | |
| y = np.unique(dongle, return_inverse=True)[1] | |
| acc = [] | |
| for tr, te in GroupKFold(5).split(Z, y, groups=uid): | |
| keep = np.isin(y[te], np.unique(y[tr])) | |
| sc = StandardScaler().fit(Z[tr]) | |
| clf = LogisticRegression(max_iter=300, C=1.0).fit(sc.transform(Z[tr]), y[tr]) | |
| acc.append(float((clf.predict(sc.transform(Z[te][keep])) == y[te][keep]).mean()) if keep.any() else np.nan) | |
| maj = float(np.bincount(y).max() / len(y)) | |
| return float(np.nanmean(acc)), maj | |
| def probe_behaviour(Z, y, dongle, uid, centre): | |
| from sklearn.linear_model import LogisticRegression | |
| from sklearn.metrics import average_precision_score | |
| from sklearn.model_selection import GroupKFold | |
| from sklearn.preprocessing import StandardScaler | |
| Zc = Z.copy() | |
| if centre: | |
| codes = np.unique(dongle, return_inverse=True)[1] | |
| for c in np.unique(codes): | |
| m = codes == c | |
| Zc[m] -= Zc[m].mean(0) | |
| aps = [] | |
| for tr, te in GroupKFold(5).split(Zc, y, groups=uid): | |
| if y[tr].sum() < 20 or y[te].sum() < 20: | |
| continue | |
| sc = StandardScaler().fit(Zc[tr]) | |
| clf = LogisticRegression(max_iter=300, C=1.0, class_weight="balanced").fit(sc.transform(Zc[tr]), y[tr]) | |
| aps.append(average_precision_score(y[te], clf.decision_function(sc.transform(Zc[te])))) | |
| return float(np.mean(aps)) if aps else float("nan") | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--a", required=True) | |
| ap.add_argument("--b", required=True) | |
| ap.add_argument("--regime", required=True, choices=["manual", "engaged"]) | |
| ap.add_argument("--events", default="hand_leave,head_turn_onset") | |
| ap.add_argument("--tag", default="") | |
| ap.add_argument("--n-boot", type=int, default=1000) | |
| ap.add_argument("--probe-rows", type=int, default=120000) | |
| ap.add_argument("--skip-probes", action="store_true") | |
| ap.add_argument("--kin-source", default="cwm2", choices=["cwm2", "causal_routes"], help="2026-09-11 opt-in: driver populations / probe labels from the K4 causal packs") | |
| ap.add_argument("--kin-version", default="v1", choices=["v1", "v2"]) | |
| a = ap.parse_args() | |
| import torch | |
| torch.set_num_threads(4) | |
| t0 = time.time() | |
| OUT.mkdir(parents=True, exist_ok=True) | |
| stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()) | |
| na, nb = Path(a.a).name, Path(a.b).name | |
| rows, res = [], {"a": a.a, "b": a.b, "regime": a.regime, "tag": a.tag, "stamp": stamp} | |
| def add(event, stat, va, vb, ci, n, npos, ndev, note=""): | |
| rows.append({"run_stamp": stamp, "tag": a.tag, "regime": a.regime, "event": event, "statistic": stat, "a": na, "b": nb, "value_a": va, "value_b": vb, | |
| "delta": va - vb, "delta_ci_lo": ci[0], "delta_ci_hi": ci[1], "n": n, "n_pairs_or_pos": npos, "n_dev": ndev, "note": note}) | |
| res.setdefault(event, {})[stat] = {"a": va, "b": vb, "delta": va - vb, "ci": list(ci), "n": n, "n_pairs_or_pos": npos, "n_dev": ndev, "note": note} | |
| print(f" {event:16s} {stat:24s} A={va:.4f} B={vb:.4f} Δ={va - vb:+.4f} [{ci[0]:+.4f},{ci[1]:+.4f}] n={n} pos/pairs={npos} dev={ndev} {note}", flush=True) | |
| # 1. shifted-twin pair-AUROC on test | |
| for ev in a.events.split(","): | |
| pop = kin_population(a.regime, "test", ev, a.kin_source, a.kin_version) | |
| sa, sb = onset_score(a.a, a.regime, "test", ev, pop["sel"]), onset_score(a.b, a.regime, "test", ev, pop["sel"]) | |
| pi, ni = make_twins(pop) | |
| if len(pi) < 100: | |
| print(f" {ev}: only {len(pi)} twin pairs -> skipped") | |
| continue | |
| va, vb = pair_auroc(sa[pi], sa[ni]), pair_auroc(sb[pi], sb[ni]) | |
| lo, hi, ndev = boot_pairs(sa[pi], sa[ni], sb[pi], sb[ni], pop["dongle"][pi], a.n_boot) | |
| add(ev, "twin_pair_auroc@3s", va, vb, (lo, hi), int(len(pop["sel"])), int(len(pi)), ndev, "same route, dt 30-60 s, |dvEgo|<2, y flips") | |
| # reference: plain AUROC-like pair statistic on random cross-route negatives (identity NOT shared) for context | |
| rng = np.random.default_rng(1) | |
| neg_all = np.flatnonzero(~pop["y"]) | |
| ri = rng.choice(neg_all, len(pi)) | |
| add(ev, "random_pair_auroc@3s", pair_auroc(sa[pi], sa[ri]), pair_auroc(sb[pi], sb[ri]), (float("nan"), float("nan")), int(len(pop["sel"])), int(len(pi)), ndev, | |
| "same positives vs random negatives (identity not shared)") | |
| # 2./3. probes on train latents | |
| if not a.skip_probes: | |
| device = "cpu" # GPU only with >= 4 GB free (shared GPU, no lock; the 120k-row GRU forward is cheap on CPU) | |
| if torch.cuda.is_available(): | |
| try: | |
| free = torch.cuda.mem_get_info()[0] / 2**30 | |
| device = "cuda" if free >= 4.0 else "cpu" | |
| except Exception: # noqa: BLE001 | |
| device = "cpu" | |
| ds.set_kin_source(a.kin_source, a.kin_version) | |
| cfgs = [json.loads((Path(r) / "config.json").read_text())["cfg"] for r in (a.a, a.b)] | |
| st = {k: v for k, v in np.load(Path(a.b) / "stats.npz").items()} | |
| if any(c.get("recipe", {}).get("u_devnorm") for c in cfgs): | |
| # the causal per-route statistics need every earlier row of the route: load the full pack, attach them, then take the SAME seeded subsample as the old path | |
| raw = ds.load_regimes([a.regime], "train", frozenset({"X", "Kh", "kin"}), verbose=lambda m: None) | |
| c = next(c for c in cfgs if c.get("recipe", {}).get("u_devnorm")) | |
| idn.add_devnorm(raw, st, float(c["recipe"].get("devnorm_n0", 50)), float(c["recipe"].get("devnorm_sd_floor", 0.2))) | |
| N_file = len(raw["t"]) | |
| if a.probe_rows and a.probe_rows < N_file: | |
| raw = ds.subset(raw, np.sort(np.random.default_rng(11).choice(N_file, a.probe_rows, replace=False))) | |
| else: | |
| raw = ds.load_regimes([a.regime], "train", frozenset({"X", "Kh", "kin"}), max_anchors=a.probe_rows, seed=11, verbose=lambda m: None) | |
| Za, Zb = latents(a.a, raw, st, device), latents(a.b, raw, st, device) | |
| acc_a, maj = probe_identity(Za, raw["dongle"], raw["uid"]) | |
| acc_b, _ = probe_identity(Zb, raw["dongle"], raw["uid"]) | |
| add("all", "identity_probe_acc", acc_a, acc_b, (float("nan"), float("nan")), int(len(raw["t"])), int(len(np.unique(raw["dongle"]))), int(len(np.unique(raw["dongle"]))), | |
| f"majority-class rate {maj:.4f}; GroupKFold(5) by route; flag if A-B > 0.05") | |
| res["identity_probe_majority"] = maj | |
| # behaviour probe with / without per-device centring (labels from the train kin file for the subsample rows) | |
| S, S0, EM, M = raw["S"].astype(np.float32), raw["S0"].astype(np.float32), raw["EM"], raw["M"] | |
| for ev in a.events.split(","): | |
| j = list(EVENTS).index(ev) | |
| z0, z, em = [(S0[:, 0], S[:, :, 0], EM[:, :, 0]), (1 - S0[:, 0], 1 - S[:, :, 0], EM[:, :, 1]), (S0[:, 1], S[:, :, 1], EM[:, :, 2])][j] | |
| em = em & M | |
| zp = np.concatenate([z0[:, None], z[:, :-1]], 1) | |
| onset = (np.nan_to_num(z, nan=0.0) > 0.5) & (np.nan_to_num(zp, nan=1.0) < 0.5) & em | |
| first = np.where(onset.any(1), onset.argmax(1), -1) | |
| sel = np.flatnonzero((np.nan_to_num(zp[:, 0], nan=1.0) < 0.5) & em[:, :TAU].all(1)) | |
| y = (first[sel] >= 0) & (first[sel] < TAU) | |
| for centre in (False, True): | |
| pa_ = probe_behaviour(Za[sel], y, raw["dongle"][sel], raw["uid"][sel], centre) | |
| pb_ = probe_behaviour(Zb[sel], y, raw["dongle"][sel], raw["uid"][sel], centre) | |
| add(ev, f"latent_probe_prauc{'_centred' if centre else ''}", pa_, pb_, (float("nan"), float("nan")), int(len(sel)), int(y.sum()), | |
| int(len(np.unique(raw["dongle"][sel]))), "per-device mean-centred latents" if centre else "raw latents") | |
| new = not CSV.exists() | |
| with open(CSV, "a", newline="") as f: | |
| w = csv.DictWriter(f, fieldnames=COLS) | |
| if new: | |
| w.writeheader() | |
| w.writerows(rows) | |
| (OUT / f"{a.tag or 'shift_twin'}_{a.regime}_{stamp}.json").write_text(json.dumps(res, indent=1, default=float)) | |
| inputs = [(Path(r) / f"pred_{a.regime}_test.npz") if (Path(r) / f"pred_{a.regime}_test.npz").exists() else (Path(r) / f"pred_layout_{a.regime}_test.npz") for r in (a.a, a.b)] | |
| write_run_json(OUT, f"shift_twin_{a.tag}_{a.regime}", [Path(__file__)], inputs, | |
| seed=0, t_start=t0, input_listing=False, extra={"tag": a.tag, "regime": a.regime, "n_rows": len(rows), "dms_fields_used": [], | |
| "kin_source": a.kin_source, "kin_version": a.kin_version}) | |
| print(f"done in {time.time() - t0:.0f}s; {len(rows)} rows -> {CSV}") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 19.1 kB
- Xet hash:
- 3ee7816cfa7f4775067f7b1ad831f2fc28f88cd4442a1ee536ca432f701a4586
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.