Buckets:
| """cwm3/weak/dms_teacher_pretrain.py — pretrain the U TRUNK (GRU(Kh+X -> 128) ⊕ mode embedding, exactly models.branches.DriverBranch as | |
| U0 kin_can_noK builds it) to predict the DMS weak-teacher soft targets of results/targets_dms_teacher/ (plan §5.2; docs/weak_teacher.md §3). | |
| *** This file and weak/dms_teacher_targets.py are the ONLY cwm3 code allowed to read DMS-derived data (here: the teacher targets | |
| face_yaw, face_pitch, eye_open, blink, distracted_soft, gaze_off_soft written by the single DMS reader). *** | |
| The trunk weights it saves (results/runs/U0_dmspre_trunk_s<k>/trunk.pt: gru + mode state_dicts) carry no DMS input: the fine-tuned | |
| U0 model (train/train_u0_from_init.py -> unchanged train/train_cwm3.py) sees exactly the kin_can_noK inputs (Kh, X, Δt, mode). | |
| Inputs (identical to U0): data/cwm3_dataset.load_regimes(["manual","engaged"], split, need={"X","Kh"}); train/batch.prep_U builds the | |
| [B,40,18] GRU input with the TRAIN-only statistics of train/batch.compute_stats (asserted equal to results/runs/U0_kin_can_noK_s0/stats.npz | |
| when present, so pretraining and fine-tuning share the standardisation). Kh is the route-median (non-causal ROI) cache as in every | |
| U0 run so far — the comparison pretrained-vs-scratch is like for like (K4 pending). | |
| Targets (train rows only; val rows only for early stopping; test rows never loaded): T0[N,6] at the anchor, T[N,25,6] at t+0.2..t+5.0, | |
| validity masks, route reliability weight w_route. Teacher head = flat_head(128 + 16 -> 128 -> 26 x 6) (Linear-GELU-Linear; discarded | |
| after pretraining). Loss = masked, weighted mean of SmoothL1 on (target - norm_mean)/norm_std for the 4 continuous channels and | |
| BCE-with-logits on the two soft probabilities; weights = valid x w_route x (1.0 anchor, 0.5 per future step). | |
| Recipe = the U0 recipe (AdamW 1e-3 / wd 1e-4, batch 4096, full pass per epoch, <= 12 epochs, patience 2 on the val loss, bf16, clip 1). | |
| Outputs results/runs/U0_dmspre_trunk_s<k>/: trunk.pt (best epoch), history.jsonl, BEST, stats.npz, val_metrics.json (per-channel MAE / | |
| BCE at the best epoch vs the constant train-mean predictor), run_dms_teacher_pretrain_<utc>.json (dms_fields_used, inputs sha256). | |
| GPU etiquette: no lock (coordinator 2026-09-10), threads <= 4, VRAM / RAM checked before CUDA init. | |
| python weak/dms_teacher_pretrain.py --seed 0 [--max-anchors N] [--epochs E] [--out DIR] | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| os.environ.setdefault("MKL_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"): | |
| sys.path.insert(0, str(p)) | |
| import cwm3_dataset as ds # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| from models import build_branch, params_by_module # noqa: E402 | |
| from models.branches import flat_head # noqa: E402 | |
| from runmeta import write_run_json, sha256_file # noqa: E402 | |
| DMS_DIR = CWM3 / "results" / "targets_dms_teacher" | |
| REF_STATS = CWM3 / "results" / "runs" / "U0_kin_can_noK_s0" / "stats.npz" | |
| DMS_FIELDS = ["face_yaw", "face_pitch", "face_roll", "blink", "eye", "isDistracted", "faceProb"] # lineage of the teacher targets | |
| REGIMES = ["manual", "engaged"] | |
| NEED = frozenset({"X", "Kh"}) | |
| RECIPE = dict(batch=4096, max_epochs=12, patience=2, lr=1e-3, weight_decay=1e-4, grad_clip=1.0, amp=True, w_anchor=1.0, w_future=0.5) | |
| MODEL = dict(hid=128, hidden=128, dropout=0.0, cabin="none") | |
| N_CONT = 4 # face_yaw, face_pitch, eye_open, blink (SmoothL1); then distracted_soft, gaze_off_soft (BCE) | |
| MIN_FREE_VRAM_GB, MIN_AVAIL_RAM_GB = 8.0, 14.0 | |
| def log(msg, out): | |
| line = f"[{time.strftime('%H:%M:%S')}] {msg}" | |
| print(line, flush=True) | |
| with open(Path(out) / "train.log", "a") as f: | |
| f.write(line + "\n") | |
| def load_teacher(split, raw): | |
| """DMS teacher targets for both regimes, concatenated in the raw's regime order; uid/t asserted equal to the raw rows.""" | |
| parts = {k: [] for k in ("T0", "T0valid", "T", "Tvalid", "w_route", "uid", "t")} | |
| norm = None | |
| for r in REGIMES: | |
| z = np.load(DMS_DIR / f"dms_{r}_{split}.npz", allow_pickle=True) | |
| for k in parts: | |
| parts[k].append(z[k]) | |
| assert list(map(str, z["channels"])) == ["face_yaw", "face_pitch", "eye_open", "blink", "distracted_soft", "gaze_off_soft"] | |
| norm = (z["norm_mean"].astype(np.float32), z["norm_std"].astype(np.float32)) | |
| cat = {k: np.concatenate(v) for k, v in parts.items()} | |
| idx = raw["regime_row"] | |
| reg = raw["regime_id"] | |
| off = np.cumsum([0] + [len(p) for p in parts["t"]]) | |
| rows = idx + off[reg.astype(np.int64)] # raw row -> row in the concatenated (manual, engaged) teacher arrays | |
| assert np.array_equal(cat["uid"][rows], raw["uid"]) and np.allclose(cat["t"][rows], raw["t"]), "teacher / data rows misaligned" | |
| out = {k: cat[k][rows] for k in ("T0", "T0valid", "T", "Tvalid", "w_route")} | |
| out["T0"] = out["T0"].astype(np.float32) | |
| out["T"] = out["T"].astype(np.float16) | |
| return out, norm | |
| def check_resources(out): | |
| import shutil | |
| import torch | |
| free_b, total_b = torch.cuda.mem_get_info() | |
| free_gb = free_b / 2**30 | |
| avail_gb = float([l for l in open("/proc/meminfo") if l.startswith("MemAvailable")][0].split()[1]) / 2**20 | |
| log(f"resources: VRAM free {free_gb:.1f} GB, RAM available {avail_gb:.1f} GB, RSS {ds.rss_gb():.1f} GB, disk free {shutil.disk_usage(out).free / 2**30:.0f} GB", out) | |
| assert free_gb >= MIN_FREE_VRAM_GB, f"free VRAM {free_gb:.1f} GB < {MIN_FREE_VRAM_GB}" | |
| assert avail_gb >= MIN_AVAIL_RAM_GB, f"available RAM {avail_gb:.1f} GB < {MIN_AVAIL_RAM_GB}" | |
| def teacher_loss(out, tg, idx, nm, ns, step_w): | |
| import torch | |
| import torch.nn.functional as F | |
| T0, T = tg["T0"][idx], tg["T"][idx].float() | |
| tgt = torch.cat([T0[:, None, :], T], 1) # [B,26,6] | |
| valid = torch.cat([tg["T0valid"][idx][:, None], tg["Tvalid"][idx]], 1) & torch.isfinite(tgt).all(-1) | |
| wt = valid.float() * tg["w_route"][idx][:, None] * step_w[None, :] # [B,26] | |
| tgt = torch.nan_to_num(tgt) | |
| reg = F.smooth_l1_loss(out[..., :N_CONT], (tgt[..., :N_CONT] - nm[:N_CONT]) / ns[:N_CONT], reduction="none") | |
| bce = F.binary_cross_entropy_with_logits(out[..., N_CONT:], tgt[..., N_CONT:].clamp(0, 1), reduction="none") | |
| per = torch.cat([reg, bce], -1) # [B,26,6] | |
| denom = wt.sum().clamp_min(1.0) | |
| loss = (per * wt[..., None]).sum() / (denom * per.shape[-1]) | |
| with torch.no_grad(): | |
| anc = (per[:, 0] * wt[:, :1]).sum(0) / wt[:, 0].sum().clamp_min(1.0) # per-channel anchor term | |
| fut = (per[:, 1:] * wt[:, 1:, None]).sum((0, 1)) / wt[:, 1:].sum().clamp_min(1.0) | |
| return loss, anc, fut, float(denom) | |
| class Teacher: | |
| def __init__(self, seed, device): | |
| import torch | |
| self.base = build_branch("U", seed, **MODEL).to(device) # torch.manual_seed(seed) inside -> same trunk init as U0 seed k | |
| torch.manual_seed(seed + 1) | |
| self.head = flat_head(MODEL["hid"] + 16, MODEL["hidden"], 26 * 6, 0.0).to(device) | |
| def parameters(self): | |
| return list(self.base.gru.parameters()) + list(self.base.mode.parameters()) + list(self.head.parameters()) | |
| def train(self): | |
| self.base.train(); self.head.train() | |
| def eval(self): | |
| self.base.eval(); self.head.eval() | |
| def forward(self, b): | |
| import torch | |
| _, h = self.base.gru(b["Kx"].float()) | |
| z = torch.cat([h[-1], self.base.mode(b["mode"].long())], 1) | |
| return self.head(z).view(-1, 26, 6).float() | |
| def trunk_state(self): | |
| return {"gru": {k: v.detach().cpu() for k, v in self.base.gru.state_dict().items()}, | |
| "mode": {k: v.detach().cpu() for k, v in self.base.mode.state_dict().items()}, | |
| "teacher_head": {k: v.detach().cpu() for k, v in self.head.state_dict().items()}} | |
| def evaluate(model, store, tg, rows, st, cfg, nm, ns, step_w, device, amp, bs=8192): | |
| """val loss (selection) + per-channel anchor / future MAE (raw units) and BCE, plus the constant train-mean predictor.""" | |
| import torch | |
| model.eval() | |
| tot, n = 0.0, 0.0 | |
| ch_sum = torch.zeros(2, 6, device=device) # [anchor/future, channel] weighted error sums | |
| ch_w = torch.zeros(2, device=device) | |
| base_sum = torch.zeros(2, 6, device=device) | |
| with torch.no_grad(): | |
| for i in range(0, len(rows), bs): | |
| idx = torch.as_tensor(rows[i:i + bs], device=device) | |
| b, _ = bt.prep_U(store, idx, st, cfg, need_targets=False) | |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp): | |
| out = model.forward(b) | |
| loss, _, _, denom = teacher_loss(out, tg, idx, nm, ns, step_w) | |
| tot += float(loss) * denom | |
| n += denom | |
| T0, T = tg["T0"][idx], tg["T"][idx].float() | |
| tgt = torch.cat([T0[:, None, :], T], 1) | |
| valid = torch.cat([tg["T0valid"][idx][:, None], tg["Tvalid"][idx]], 1) & torch.isfinite(tgt).all(-1) | |
| wt = valid.float() * (tg["w_route"][idx] > 0).float()[:, None] | |
| tgt = torch.nan_to_num(tgt) | |
| pred_raw = out[..., :N_CONT] * ns[:N_CONT] + nm[:N_CONT] | |
| err = torch.cat([(pred_raw - tgt[..., :N_CONT]).abs(), | |
| torch.nn.functional.binary_cross_entropy_with_logits(out[..., N_CONT:], tgt[..., N_CONT:].clamp(0, 1), reduction="none")], -1) | |
| base_p = torch.cat([(nm[:N_CONT] - tgt[..., :N_CONT]).abs().expand_as(tgt[..., :N_CONT]), | |
| torch.nn.functional.binary_cross_entropy(nm[N_CONT:].clamp(1e-4, 1 - 1e-4).expand_as(tgt[..., N_CONT:]), tgt[..., N_CONT:].clamp(0, 1), reduction="none")], -1) | |
| ch_sum[0] += (err[:, 0] * wt[:, :1]).sum(0) | |
| ch_sum[1] += (err[:, 1:] * wt[:, 1:, None]).sum((0, 1)) | |
| base_sum[0] += (base_p[:, 0] * wt[:, :1]).sum(0) | |
| base_sum[1] += (base_p[:, 1:] * wt[:, 1:, None]).sum((0, 1)) | |
| ch_w[0] += wt[:, 0].sum() | |
| ch_w[1] += wt[:, 1:].sum() | |
| ch = (ch_sum / ch_w[:, None].clamp_min(1)).cpu().numpy() | |
| bs_ = (base_sum / ch_w[:, None].clamp_min(1)).cpu().numpy() | |
| return tot / max(n, 1.0), ch, bs_ | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--max-anchors", type=int, default=0, help="seeded row subset (smoke only)") | |
| ap.add_argument("--epochs", type=int, default=0) | |
| ap.add_argument("--out", default="") | |
| ap.add_argument("--cpu", action="store_true") | |
| a = ap.parse_args() | |
| import torch | |
| torch.set_num_threads(4) | |
| out = Path(a.out) if a.out else CWM3 / "results" / "runs" / (f"U0_dmspre_trunk_s{a.seed}" + (f"_smoke_max{a.max_anchors}" if a.max_anchors else "")) | |
| out.mkdir(parents=True, exist_ok=True) | |
| if (out / "TRAIN_DONE").exists(): | |
| print(f"{out} already done") | |
| return | |
| t_start = time.time() | |
| epochs = a.epochs or RECIPE["max_epochs"] | |
| cfg = {"cabin": "none"} | |
| log(f"=== dms_teacher_pretrain seed {a.seed} out={out} max_anchors={a.max_anchors} epochs={epochs} ===", out) | |
| tr = ds.load_regimes(REGIMES, "train", NEED, max_anchors=a.max_anchors or None, seed=a.seed + 7, verbose=lambda m: log(m, out)) | |
| va = ds.load_regimes(REGIMES, "val", NEED, max_anchors=(max(a.max_anchors // 4, 64) if a.max_anchors else None), seed=a.seed + 7, verbose=lambda m: log(m, out)) | |
| tg_tr, norm = load_teacher("train", tr) | |
| tg_va, _ = load_teacher("val", va) | |
| zst = np.load(S12R / "target_stats.npz") | |
| st = bt.compute_stats(tr, {"mu": zst["mu"], "sd": zst["sd"]}) | |
| if REF_STATS.exists() and not a.max_anchors: | |
| ref = np.load(REF_STATS) | |
| for k in ("xmu", "xsd", "kmu", "ksd"): | |
| assert np.allclose(ref[k], st[k], rtol=1e-4, atol=1e-5), f"standardisation differs from the U0 baseline: {k}" | |
| log("standardisation statistics identical to results/runs/U0_kin_can_noK_s0/stats.npz", out) | |
| np.savez(out / "stats.npz", **st) | |
| log(f"loaded train {len(tr['t']):,} val {len(va['t']):,} rows; RSS {ds.rss_gb():.2f} GiB; teacher valid anchors train {tg_tr['T0valid'].mean():.3f} " | |
| f"w>0 {(tg_tr['w_route'] > 0).mean():.3f}", out) | |
| device = "cpu" if a.cpu else "cuda" | |
| if device == "cuda": | |
| check_resources(out) | |
| amp = device == "cuda" and RECIPE["amp"] | |
| store_tr = bt.GpuStore(tr, ["X", "Kh", "mode", "slice_tag"], device) | |
| store_va = bt.GpuStore(va, ["X", "Kh", "mode", "slice_tag"], device) | |
| to_dev = lambda d: {k: torch.from_numpy(v).to(device) for k, v in d.items()} | |
| tg_tr_d, tg_va_d = to_dev(tg_tr), to_dev(tg_va) | |
| del tr["X"], tr["Kh"], va["X"], va["Kh"] | |
| nm, ns = torch.as_tensor(norm[0], device=device), torch.as_tensor(norm[1], device=device) | |
| step_w = torch.cat([torch.full((1,), RECIPE["w_anchor"]), torch.full((25,), RECIPE["w_future"])]).to(device) | |
| model = Teacher(a.seed, device) | |
| n_par = {"gru": sum(p.numel() for p in model.base.gru.parameters()), "mode": sum(p.numel() for p in model.base.mode.parameters()), | |
| "teacher_head": sum(p.numel() for p in model.head.parameters())} | |
| opt = torch.optim.AdamW(model.parameters(), lr=RECIPE["lr"], weight_decay=RECIPE["weight_decay"]) | |
| rows_tr, rows_va = np.arange(len(tg_tr["w_route"])), np.arange(len(tg_va["w_route"])) | |
| bs = RECIPE["batch"] | |
| log(f"model: trunk params {n_par} ; gpu store {(store_tr.bytes() + store_va.bytes()) / 2**30:.2f} GiB ; device {device} amp {amp}", out) | |
| best, best_ep, patience = 1e18, -1, 0 | |
| for ep in range(epochs): | |
| t_ep = time.time() | |
| rng = np.random.default_rng(a.seed * 1000 + ep) | |
| order = rows_tr[rng.permutation(len(rows_tr))] | |
| model.train() | |
| tot, nb = 0.0, 0 | |
| for i in range(0, len(order), bs): | |
| idx = torch.as_tensor(order[i:i + bs], device=device) | |
| b, _ = bt.prep_U(store_tr, idx, st, cfg, need_targets=False) | |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp): | |
| outp = model.forward(b) | |
| loss, anc, fut, _ = teacher_loss(outp, tg_tr_d, idx, nm, ns, step_w) | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), RECIPE["grad_clip"]) | |
| opt.step() | |
| tot += float(loss.detach()) | |
| nb += 1 | |
| if nb % 100 == 0: | |
| log(f" ep {ep} b {nb} loss {tot / nb:.4f} anchor_ch={np.round(anc.cpu().numpy(), 3).tolist()}", out) | |
| v, ch, base = evaluate(model, store_va, tg_va_d, rows_va, st, cfg, nm, ns, step_w, device, amp) | |
| improved = v < best - 1e-5 | |
| if improved: | |
| best, best_ep, patience = v, ep, 0 | |
| torch.save({**model.trunk_state(), "epoch": ep, "val_loss": v, "seed": a.seed, "recipe": RECIPE, "model": MODEL, | |
| "channels": ["face_yaw", "face_pitch", "eye_open", "blink", "distracted_soft", "gaze_off_soft"], | |
| "norm_mean": norm[0], "norm_std": norm[1], "stats": {k: np.asarray(v_) for k, v_ in st.items()}, | |
| "val_metrics": {"anchor": ch[0].tolist(), "future": ch[1].tolist(), "const_anchor": base[0].tolist(), "const_future": base[1].tolist()}, | |
| "source": "weak/dms_teacher_pretrain.py", "dms_fields_used": DMS_FIELDS}, out / "trunk.pt") | |
| (out / "BEST").write_text(f"trunk.pt epoch {ep}\n") | |
| else: | |
| patience += 1 | |
| rec = {"epoch": ep, "train_loss": tot / max(nb, 1), "val_loss": v, "best": best, "best_epoch": best_ep, "val_anchor_err": ch[0].tolist(), | |
| "val_future_err": ch[1].tolist(), "const_anchor_err": base[0].tolist(), "const_future_err": base[1].tolist(), "lr": RECIPE["lr"], | |
| "wall_s": round(time.time() - t_ep, 1), "t_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())} | |
| with open(out / "history.jsonl", "a") as f: | |
| f.write(json.dumps(rec) + "\n") | |
| log(f"ep {ep}: train {rec['train_loss']:.4f} val {v:.4f} best {best:.4f}@{best_ep} anchor err {np.round(ch[0], 3).tolist()} " | |
| f"(const {np.round(base[0], 3).tolist()}) ({rec['wall_s']:.0f}s)", out) | |
| if patience >= RECIPE["patience"]: | |
| break | |
| z = torch.load(out / "trunk.pt", map_location="cpu", weights_only=False) | |
| vm = {"best_epoch": best_ep, "best_val_loss": best, "channels": z["channels"], "anchor_err": z["val_metrics"]["anchor"], "future_err": z["val_metrics"]["future"], | |
| "const_anchor_err": z["val_metrics"]["const_anchor"], "const_future_err": z["val_metrics"]["const_future"], | |
| "err_kind": ["MAE raw units"] * N_CONT + ["BCE"] * 2, "n_train_rows": int(len(rows_tr)), "n_val_rows": int(len(rows_va)), | |
| "train_w_gt0_frac": float((tg_tr["w_route"] > 0).mean()), "epochs_run": int(len(open(out / "history.jsonl").read().splitlines()))} | |
| (out / "val_metrics.json").write_text(json.dumps(vm, indent=1)) | |
| (out / "TRAIN_DONE").write_text(json.dumps({"best_epoch": best_ep, "best": best, "params": n_par, "smoke": bool(a.max_anchors)})) | |
| inputs = [DMS_DIR / f"dms_{r}_{s}.npz" for r in REGIMES for s in ("train", "val")] + [S12R / "target_stats.npz"] | |
| for r in REGIMES: | |
| for s in ("train", "val"): | |
| f = ds.files(r, s) | |
| inputs += [f["targets"], f["kinhist"]] | |
| write_run_json(out, "dms_teacher_pretrain", [Path(__file__), CWM3 / "models" / "branches.py", CWM3 / "train" / "batch.py", CWM3 / "data" / "cwm3_dataset.py"], | |
| inputs, seed=a.seed, t_start=t_start, | |
| extra={"stage": "U0_dmspre_trunk", "recipe": RECIPE, "model": MODEL, "params": n_par, "best_epoch": best_ep, "best_val_loss": best, | |
| "val_metrics": vm, "dms_fields_used": DMS_FIELDS, "dms_lineage": "teacher targets only (results/targets_dms_teacher); no DMS array is an input", | |
| "trunk_file": str(out / "trunk.pt"), "trunk_sha256": sha256_file(out / "trunk.pt"), "max_anchors": a.max_anchors, | |
| "note_kh": "Kh from the route-median (non-causal) pose cache, as every U0 run (K4 pending)", "gpu_lock": "none (coordinator 2026-09-10)"}) | |
| log(f"DONE best {best:.4f}@{best_ep} in {time.time() - t_start:.0f}s", out) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 18.7 kB
- Xet hash:
- 97a6b572ae1592d8e76f48075d6d77ea2eee0b104ee58ddde0edccf0e174f889
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.