Buckets:
| """cwm3/eval/shift_twin_c1.py — K5 shifted-twin / identity check for the C1 transfer runs (E10 (c)), on the CAUSAL v2 kinematics populations. | |
| eval/shift_twin.py (unchanged, imported) builds its event populations from the OLD route-median kin targets and its latent extractor builds the | |
| DriverBranch without the survivor recipe's input LayerNorm. The C1 runs are trained and scored on the causal v2 packs with `u_in_ln`, so this | |
| thin driver re-uses shift_twin's statistics (make_twins, pair_auroc, boot_pairs, onset_score, probe_identity, probe_behaviour) and only replaces | |
| * the population source: results/targets_kin_causal_routes_v2/kin_causal_<regime>_test.npz (S_c / S0_c / EM_c), rows restricted to row_ok, and | |
| * the latent extractor: DriverBranch built with train_cwm3.model_kwargs(cfg) (in_ln honoured) and prep_U with the run's own kclip. | |
| Rows are appended to the same results/shift_twin.csv (same columns; note column says "causal v2 populations") and a json goes to results/shift_twin/. | |
| A = fine-tuned-from-AIDE-trunk run, B = its scratch twin (same seed, same recipe, same rows). No DMS field, no weak/ import. | |
| python eval/shift_twin_c1.py --a results/runs/U0_kin_can_noK_s0_kcv2_clip8_ln_aidepre --b results/runs/U0_kin_can_noK_s0_kcv2_clip8_ln --regime engaged --tag C1_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 | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train", CWM3 / "eval"): | |
| sys.path.insert(0, str(p)) | |
| import cwm3_dataset as ds # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| import shift_twin as STW # noqa: E402 (statistics re-used verbatim) | |
| import train_cwm3 as tc # noqa: E402 (model_kwargs: in_ln from the run cfg) | |
| from models import build_branch # noqa: E402 | |
| from runmeta import write_run_json # noqa: E402 | |
| TAU = STW.TAU | |
| EVENTS = STW.EVENTS | |
| def kin_population_causal(regime, split, event, version="v2"): | |
| """shift_twin.kin_population on the causal packs: rows with a causal box (row_ok), S_c / S0_c / EM_c keys.""" | |
| ds.set_kin_source("causal_routes", version) | |
| d = np.load(ds.kin_dir() / f"kin_causal_{regime}_{split}.npz", allow_pickle=True) | |
| s = np.load(S12R / f"targets_{regime}_{split}.npz", allow_pickle=True) | |
| assert np.array_equal(s["uid"], d["uid"]) and np.allclose(s["t"], d["t"]) | |
| M = s["M"].astype(bool) | |
| S, S0, EM = d["S_c"].astype(np.float32), d["S0_c"].astype(np.float32), d["EM_c"] & 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) & d["row_ok"].astype(bool)) | |
| 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 - prod(1 - h) over the first TAU steps, read from the run's pred_layout (the artefact kept after the raw prediction pack is archived: | |
| layout onset_p[:, tau index of 15 steps, event] is exactly shift_twin.onset_score's quantity); falls back to the raw pack when present.""" | |
| lay = Path(run_dir) / f"pred_layout_{regime}_{split}.npz" | |
| raw = Path(run_dir) / f"pred_{regime}_{split}.npz" | |
| if raw.exists(): | |
| return STW.onset_score(run_dir, regime, split, event, sel) | |
| z = np.load(lay, allow_pickle=True) | |
| ev = list(map(str, z["event_names"])) | |
| ti = list(map(int, z["onset_taus"])).index(TAU) | |
| return z["onset_p"][sel, ti, ev.index(event)].astype(np.float64) | |
| def latents(run_dir, raw, st, device, bs=8192): | |
| """frozen GRU final state [N,128] of the run's BEST checkpoint; the model is built exactly as the trainer builds it (in_ln, kclip from cfg).""" | |
| import torch | |
| cfgj = json.loads((Path(run_dir) / "config.json").read_text()) | |
| cfg, seed = cfgj["cfg"], int(cfgj["args"]["seed"]) | |
| model = build_branch("U", seed, **tc.model_kwargs(cfg)).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"], device) | |
| 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, _ = bt.prep_U(store, idx, st, cfg, need_targets=False) | |
| kx = b["Kx"].float() | |
| if model.in_ln is not None: | |
| 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 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,hand_return") | |
| 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("--kin-version", default="v2", choices=["v1", "v2"]) | |
| ap.add_argument("--skip-probes", action="store_true") | |
| a = ap.parse_args() | |
| import torch | |
| torch.set_num_threads(4) | |
| t0 = time.time() | |
| STW.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 | |
| for r in (a.a, a.b): | |
| assert (Path(r) / "PRED_DONE").exists() and (Path(r) / "TRAIN_DONE").exists(), f"{r} not finished" | |
| rows, res = [], {"a": a.a, "b": a.b, "regime": a.regime, "tag": a.tag, "stamp": stamp, "populations": f"causal {a.kin_version} (row_ok rows)"} | |
| note_pop = f"causal {a.kin_version} populations; " | |
| 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_pop + 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) | |
| for ev in a.events.split(","): | |
| pop = kin_population_causal(a.regime, "test", ev, 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 = STW.make_twins(pop) | |
| ok = np.isfinite(sa[pi]) & np.isfinite(sa[ni]) & np.isfinite(sb[pi]) & np.isfinite(sb[ni]) | |
| pi, ni = pi[ok], ni[ok] | |
| if len(pi) < 100: | |
| print(f" {ev}: only {len(pi)} twin pairs -> skipped") | |
| continue | |
| va, vb = STW.pair_auroc(sa[pi], sa[ni]), STW.pair_auroc(sb[pi], sb[ni]) | |
| lo, hi, ndev = STW.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") | |
| rng = np.random.default_rng(1) | |
| ri = rng.choice(np.flatnonzero(~pop["y"] & np.isfinite(sa) & np.isfinite(sb)), len(pi)) | |
| add(ev, "random_pair_auroc@3s", STW.pair_auroc(sa[pi], sa[ri]), STW.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)") | |
| if not a.skip_probes: | |
| free_gb = torch.cuda.mem_get_info()[0] / 2**30 if torch.cuda.is_available() else 0.0 | |
| device = "cuda" if free_gb >= 6.0 else "cpu" # shared GPU (no lock): the probe forward is tiny, fall back to CPU when VRAM is short | |
| print(f" probes on {device} (VRAM free {free_gb:.1f} GB)", flush=True) | |
| ds.set_kin_source("causal_routes", a.kin_version) | |
| raw = ds.load_regimes([a.regime], "train", frozenset({"X", "Kh", "kin"}), max_anchors=a.probe_rows, seed=11, verbose=lambda m: None) | |
| keep = np.flatnonzero(raw["kin_row_ok"].astype(bool)) if "kin_row_ok" in raw else np.arange(len(raw["t"])) | |
| raw = ds.subset(raw, keep) | |
| st = {k: v for k, v in np.load(Path(a.b) / "stats.npz").items()} | |
| Za, Zb = latents(a.a, raw, st, device), latents(a.b, raw, st, device) | |
| acc_a, maj = STW.probe_identity(Za, raw["dongle"], raw["uid"]) | |
| acc_b, _ = STW.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 | |
| 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_ = STW.probe_behaviour(Za[sel], y, raw["dongle"][sel], raw["uid"][sel], centre) | |
| pb_ = STW.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 STW.CSV.exists() | |
| with open(STW.CSV, "a", newline="") as f: | |
| w = csv.DictWriter(f, fieldnames=STW.COLS) | |
| if new: | |
| w.writeheader() | |
| w.writerows(rows) | |
| (STW.OUT / f"{a.tag or 'shift_twin_c1'}_{a.regime}_{stamp}.json").write_text(json.dumps(res, indent=1, default=float)) | |
| write_run_json(STW.OUT, f"shift_twin_c1_{a.tag}_{a.regime}", [Path(__file__), CWM3 / "eval" / "shift_twin.py"], | |
| [Path(a.a) / f"pred_layout_{a.regime}_test.npz", Path(a.b) / f"pred_layout_{a.regime}_test.npz"], seed=0, t_start=t0, input_listing=False, | |
| extra={"tag": a.tag, "regime": a.regime, "n_rows": len(rows), "kin_version": a.kin_version, "dms_fields_used": []}) | |
| print(f"done in {time.time() - t0:.0f}s; {len(rows)} rows -> {STW.CSV}") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 12 kB
- Xet hash:
- 6649fe85bb5fbba3a7f7fe23161d6e2d0753c2e3e450029a917bc979f7f9b959
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.