Buckets:
| """cwm3/models/identity.py — opt-in modules of the INNOVATION experiments (docs/roadmap_2026-09-11.md §二 items 2–3; docs/innovations_report.md). | |
| Nothing here is active unless a stage recipe (train/stages.yaml) sets one of the keys below; every default reproduces the old trainer bit-for-bit | |
| (train/train_cwm3.py calls prep_for / store_keys / extra_losses / prepare_raws, which are no-ops for the old recipes). | |
| Innovation 2 — identity-invariant driver trunk U (baseline U0 kin_can_noK, causal v2 inputs, stats v2, kclip 8, u_in_ln = `_kcv2_clip8_ln`): | |
| recipe `u_adv: true`, lambda `adv` : gradient-reversal (DANN) device-classifier head Linear(128 -> n_train_devices) on the GRU final state. The head | |
| learns the dongle id of TRAIN rows (val / test devices are disjoint from train -> dev_id -1, ignored); the trunk receives | |
| the NEGATED gradient of that cross-entropy scaled by lambda adv, i.e. it is pushed towards a state from which the device | |
| cannot be read. Selection objective unchanged (val chain NLL; the term is logged as part `adv_dev`, never selected on). | |
| `adv_devices` (the class count) is recorded in cfg / config.json by prepare_raws so that resumed runs and evaluators | |
| rebuild the same module (train_cwm3.model_kwargs). | |
| recipe `u_devnorm: true` : per-device CAUSAL normalisation of the 11 kinematics channels. For every anchor row the mean / sd of the anchor | |
| kinematics K0 = Kh[:, -1] over the rows of the SAME (regime, route) with t <= anchor t (inclusive) — a route belongs to one | |
| dongle, so this is a per-device statistic restricted to the route — Bayesian-shrunk to the TRAIN-pool statistics kmu / ksd | |
| (pseudo-count `devnorm_n0`, default 50 anchors = 10 s of the 0.2 s grid) with an sd floor `devnorm_sd_floor` x ksd | |
| (default 0.2). Only the model INPUT changes (Kx[..., :11] standardised with the per-row statistics instead of kmu / ksd); | |
| the valid bit, the CAN channels, Δt and the K targets / k_mae keep the old path. Causal by construction: row i uses rows j | |
| of the same route with t_j <= t_i only (tests/test_identity_invariance.py perturbs later rows and checks invariance). | |
| Innovation 3 (b2) — state-supervised E dynamics (baseline E0 pooled_P, pack v2, stats v2, select_LF false): | |
| recipe `f_state_sup: true`, lambda `FS`: Linear(256 -> 10) read-out SHARED over the five horizons on the F head output F_hat(t+k), k = 1..5 (steps 4, 9, 14, | |
| 19, 24) -> the STATE at t+k: the six P active bits (targets_pressure P, mask Pvalid & M), R >= 1 and R >= 2 | |
| (targets_response R / Rvalid; grades none / mild / hard) and brake / gas pressed (targets Y, mask M); masked BCE with | |
| logits, weight lambda FS. Selection objective unchanged (D + P chain NLL); the term is logged as part `fs_state`. | |
| No openpilot DMS field is read anywhere (dongle id = the device string of the s12 targets file; kinematics = YOLO-pose; P / R = radar / lane / CAN rules). | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| RESP = CWM3 / "results" / "targets_response" | |
| P_CHAINS = ("ttc", "ttc_drop", "cut_in", "lead_braking", "lane_excursion", "badas") | |
| FS_NAMES = [f"P_{c}" for c in P_CHAINS] + ["R_any", "R_hard", "brake", "gas"] # 10 state bits read from F_hat(t+k) | |
| FS_STEPS = [4, 9, 14, 19, 24] # t+1 .. t+5 s on the 0.2 s grid | |
| DK = 11 | |
| # ----------------------------------------------------------------------------- gradient reversal (u_adv) | |
| class _GradReverse(torch.autograd.Function): | |
| def forward(ctx, x, lam): | |
| ctx.lam = float(lam) | |
| return x.view_as(x) | |
| def backward(ctx, g): | |
| return -ctx.lam * g, None | |
| def grad_reverse(x, lam=1.0): | |
| """identity in the forward pass, -lam * gradient in the backward pass (DANN).""" | |
| return _GradReverse.apply(x, lam) | |
| def adv_loss(logits, dev_id): | |
| """cross-entropy of the device head over rows whose device is a TRAIN device (dev_id >= 0); 0 when no such row is in the batch.""" | |
| m = dev_id >= 0 | |
| if not bool(m.any()): | |
| return logits.float().sum() * 0.0 | |
| return F.cross_entropy(logits[m].float(), dev_id[m].long()) | |
| def device_codes(dongle): | |
| """sorted unique TRAIN device strings -> the class index of each (the order is deterministic, recorded via adv_devices).""" | |
| return np.unique(np.asarray(dongle)) | |
| def add_device_ids(raw, codes): | |
| """raw['dev_id'] int64: class index of the row's dongle among the TRAIN codes, -1 for a device not in the train set.""" | |
| pos = np.searchsorted(codes, raw["dongle"]) | |
| pos = np.clip(pos, 0, len(codes) - 1) | |
| hit = codes[pos] == raw["dongle"] | |
| raw["dev_id"] = np.where(hit, pos, -1).astype(np.int64) | |
| return raw | |
| # ----------------------------------------------------------------------------- causal per-device normalisation (u_devnorm) | |
| def causal_devnorm(K0, uid_code, mode, t, kmu, ksd, n0=50.0, sd_floor=0.2): | |
| """K0 [N,11] anchor kinematics (NaN = invalid); groups = (mode, uid_code); expanding INCLUSIVE mean / sd along t within each group, shrunk to the pool | |
| statistics (kmu, ksd) with pseudo-count n0; sd floored at sd_floor * ksd. -> mu [N,11], sd [N,11] float32 (row order preserved).""" | |
| K0 = np.asarray(K0, np.float64) | |
| N, D = K0.shape | |
| uid_code, mode, t = np.asarray(uid_code), np.asarray(mode), np.asarray(t, np.float64) | |
| order = np.lexsort((t, uid_code, mode)) # sort by mode, then route, then time | |
| k = K0[order] | |
| fin = np.isfinite(k) | |
| kz = np.where(fin, k, 0.0) | |
| m_o, u_o = mode[order], uid_code[order] | |
| new = np.r_[True, (u_o[1:] != u_o[:-1]) | (m_o[1:] != m_o[:-1])] | |
| gid = np.cumsum(new) - 1 | |
| cs, cs2, cn = np.cumsum(kz, 0), np.cumsum(kz ** 2, 0), np.cumsum(fin, 0).astype(np.float64) | |
| starts = np.flatnonzero(new) | |
| off, off2, offn = np.zeros((len(starts), D)), np.zeros((len(starts), D)), np.zeros((len(starts), D)) | |
| if len(starts) > 1: | |
| off[1:], off2[1:], offn[1:] = cs[starts[1:] - 1], cs2[starts[1:] - 1], cn[starts[1:] - 1] | |
| S, S2, n = cs - off[gid], cs2 - off2[gid], cn - offn[gid] | |
| kmu, ksd = np.asarray(kmu, np.float64), np.asarray(ksd, np.float64) | |
| mu = (n0 * kmu + S) / (n0 + n) | |
| ex2 = (n0 * (ksd ** 2 + kmu ** 2) + S2) / (n0 + n) | |
| sd = np.maximum(np.sqrt(np.maximum(ex2 - mu ** 2, 0.0)), sd_floor * ksd) + 1e-6 | |
| out_mu, out_sd = np.empty((N, D), np.float32), np.empty((N, D), np.float32) | |
| out_mu[order], out_sd[order] = mu, sd | |
| return out_mu, out_sd | |
| def add_devnorm(raw, st, n0=50.0, sd_floor=0.2): | |
| """raw['Kdn_mu'], raw['Kdn_sd'] float32 [N,11] from the anchor kinematics Kh[:, -1] of the rows in `raw` (full split at prediction time; the train / | |
| val packs at training time; a probe SUBSAMPLE would see fewer earlier rows -> callers pass the full pack and subset afterwards).""" | |
| uid_code = raw["uid_code"] if "uid_code" in raw else np.unique(raw["uid"], return_inverse=True)[1] | |
| K0 = np.asarray(raw["Kh"][:, -1, :DK], np.float32) | |
| raw["Kdn_mu"], raw["Kdn_sd"] = causal_devnorm(K0, uid_code, np.asarray(raw["mode"], np.int64), raw["t"], np.asarray(st["kmu"]), np.asarray(st["ksd"]), n0, sd_floor) | |
| return raw | |
| def prep_U_devnorm(store, idx, st, cfg, need_targets=True): | |
| """train/batch.prep_U, then the 11 kinematics input channels re-standardised with the per-row causal statistics (same clamp).""" | |
| sys.path.insert(0, str(CWM3 / "train")) | |
| import batch as bt | |
| b, tg = bt.prep_U(store, idx, st, cfg, need_targets) | |
| Kh = store.t["Kh"][idx].float() | |
| mu, sd = store.t["Kdn_mu"][idx].float(), store.t["Kdn_sd"][idx].float() | |
| clip = bt.k_clip(cfg) | |
| Z = torch.nan_to_num(torch.clamp((Kh - mu[:, None, :]) / sd[:, None, :], -clip, clip)) | |
| b["Kx"] = torch.cat([Z, b["Kx"][..., DK:]], -1) | |
| return b, tg | |
| # ----------------------------------------------------------------------------- state supervision of F_hat (f_state_sup) | |
| def add_response(raw, split, synthetic=0, seed=0): | |
| """raw['R'] uint8 [N,25] (response grade 0 none / 1 mild / 2 hard at each future step), raw['Rvalid'] bool [N,25] from results/targets_response, row-aligned | |
| through raw['regime_row'] (position inside the regime file) with a uid / t witness; synthetic runs get random grades.""" | |
| N = len(raw["t"]) | |
| R, Rv = np.zeros((N, 25), np.uint8), np.zeros((N, 25), bool) | |
| if synthetic: | |
| rng = np.random.default_rng(seed) | |
| raw["R"], raw["Rvalid"] = rng.integers(0, 3, (N, 25)).astype(np.uint8), np.ones((N, 25), bool) | |
| return raw | |
| mode = np.asarray(raw["mode"], np.int64) | |
| rows_all = np.asarray(raw["regime_row"], np.int64) | |
| for mv, regime in ((0, "manual"), (1, "engaged")): | |
| m = mode == mv | |
| if not m.any(): | |
| continue | |
| rows = rows_all[m] | |
| with np.load(RESP / f"response_{regime}_{split}.npz", allow_pickle=True) as z: | |
| assert np.array_equal(z["uid"][rows], raw["uid"][m]) and np.allclose(z["t"][rows], raw["t"][m]), f"response_{regime}_{split}: uid/t witness" | |
| R[m], Rv[m] = z["R"][rows].astype(np.uint8), z["Rvalid"][rows].astype(bool) | |
| raw["R"], raw["Rvalid"] = R, Rv | |
| return raw | |
| def fs_targets(store, idx, tg): | |
| """state bits at the five horizons -> y float [B,5,10], m bool [B,5,10] (FS_NAMES order).""" | |
| M = store.t["M"][idx][:, FS_STEPS] | |
| zP = tg["zP"][:, FS_STEPS, :].float() # P active (prep_E: zP = P as float) | |
| mP = store.t["Pvalid"][idx][:, FS_STEPS, :] & M[:, :, None] | |
| R = store.t["R"][idx][:, FS_STEPS].long() | |
| Rv = store.t["Rvalid"][idx][:, FS_STEPS] & M | |
| yR = torch.stack([(R >= 1), (R >= 2)], -1).float() | |
| mR = torch.stack([Rv, Rv], -1) | |
| Y = store.t["Y"][idx][:, FS_STEPS].float() | |
| yV = torch.stack([(Y[..., 2] > 0.5), (Y[..., 1] > 0.5)], -1).float() # brake, gas pressed | |
| mV = torch.stack([M, M], -1) | |
| return torch.cat([zP, yR, yV], -1), torch.cat([mP, mR, mV], -1) | |
| def fs_loss(logits, y, m): | |
| """masked mean BCE-with-logits over the valid state bits.""" | |
| l = F.binary_cross_entropy_with_logits(logits.float(), y.float(), reduction="none") | |
| mf = m.float() | |
| return (l * mf).sum() / mf.sum().clamp_min(1.0) | |
| # ----------------------------------------------------------------------------- trainer hooks (no-ops for the old recipes) | |
| def store_keys(cfg): | |
| """extra GpuStore keys of the opt-in arms (GpuStore skips keys absent from the raw dict).""" | |
| rec = cfg.get("recipe", {}) | |
| ks = [] | |
| if rec.get("u_adv"): | |
| ks.append("dev_id") | |
| if rec.get("u_devnorm"): | |
| ks += ["Kdn_mu", "Kdn_sd"] | |
| if rec.get("f_state_sup"): | |
| ks += ["Y", "R", "Rvalid"] | |
| return ks | |
| def prep_for(cfg): | |
| """the batch-preparation function of a run cfg: prep_U_devnorm for `u_devnorm`, else train/batch.PREP[branch] (the old path).""" | |
| sys.path.insert(0, str(CWM3 / "train")) | |
| import batch as bt | |
| if cfg.get("branch") == "U" and cfg.get("recipe", {}).get("u_devnorm"): | |
| return prep_U_devnorm | |
| return bt.PREP[cfg["branch"]] | |
| def extra_losses(loss, parts, out, tg, store, idx, cfg): | |
| """adds the opt-in terms to the TRAINING loss (never to the selection objective): lambda adv * adv_dev (u_adv), lambda FS * fs_state (f_state_sup).""" | |
| rec, lam = cfg.get("recipe", {}), cfg.get("lambda", {}) | |
| if rec.get("u_adv") and "ADV" in out and store.has("dev_id"): | |
| la = adv_loss(out["ADV"], store.t["dev_id"][idx]) | |
| parts["adv_dev"] = la.detach() | |
| loss = loss + float(lam.get("adv", 0.1)) * la | |
| if rec.get("f_state_sup") and "FS" in out and store.has("R") and tg is not None and "zP" in tg: | |
| y, m = fs_targets(store, idx, tg) | |
| lf = fs_loss(out["FS"], y, m) | |
| parts["fs_state"] = lf.detach() | |
| loss = loss + float(lam.get("FS", 1.0)) * lf | |
| return loss | |
| def prepare_raws(args, cfg, tr, va, st, log=print): | |
| """host-side arrays of the opt-in arms on the TRAIN / VAL raws; -> True when cfg gained keys that config.json must record (adv_devices).""" | |
| rec = cfg.get("recipe", {}) | |
| changed = False | |
| if rec.get("u_adv"): | |
| codes = device_codes(tr["dongle"]) | |
| cfg["adv_devices"] = int(len(codes)) | |
| add_device_ids(tr, codes) | |
| add_device_ids(va, codes) | |
| log(f"u_adv: {len(codes)} train devices -> GRL device head; val rows with a train device: {(va['dev_id'] >= 0).mean():.3f} (expected 0: disjoint devices)") | |
| changed = True | |
| if rec.get("u_devnorm"): | |
| n0, fl = float(rec.get("devnorm_n0", 50)), float(rec.get("devnorm_sd_floor", 0.2)) | |
| for r in (tr, va): | |
| add_devnorm(r, st, n0, fl) | |
| log(f"u_devnorm: causal per-route statistics on train / val (n0 {n0:g}, sd floor {fl:g} x pool sd); train Kdn_sd / ksd median per channel " | |
| f"{np.round(np.nanmedian(tr['Kdn_sd'] / np.asarray(st['ksd'], np.float32)[None], 0), 2).tolist()}") | |
| if rec.get("f_state_sup"): | |
| add_response(tr, "train", synthetic=getattr(args, "synthetic", 0), seed=getattr(args, "seed", 0)) | |
| add_response(va, "val", synthetic=getattr(args, "synthetic", 0), seed=getattr(args, "seed", 0) + 1) | |
| log(f"f_state_sup: response grades attached (train R>=1 rate {(tr['R'] >= 1)[tr['Rvalid']].mean():.4f}, R>=2 {(tr['R'] >= 2)[tr['Rvalid']].mean():.4f}); FS bits {FS_NAMES}") | |
| return changed | |
| def prepare_pred_raws(args, cfg, raws, st): | |
| """the prediction packs need only the devnorm arrays (ADV / FS outputs are ignored at prediction time).""" | |
| rec = cfg.get("recipe", {}) | |
| if rec.get("u_devnorm"): | |
| for raw in raws.values(): | |
| add_devnorm(raw, st, float(rec.get("devnorm_n0", 50)), float(rec.get("devnorm_sd_floor", 0.2))) | |
| def restore_cfg(out_dir, cfg): | |
| """a resumed / re-invoked run: copy the keys prepare_raws recorded in an earlier invocation's config.json (adv_devices) into cfg.""" | |
| p = Path(out_dir) / "config.json" | |
| if not p.exists(): | |
| return cfg | |
| try: | |
| old = json.loads(p.read_text()).get("cfg", {}) | |
| except Exception: # noqa: BLE001 | |
| return cfg | |
| for k in ("adv_devices",): | |
| if k in old and k not in cfg: | |
| cfg[k] = old[k] | |
| return cfg | |
Xet Storage Details
- Size:
- 14.9 kB
- Xet hash:
- f98747c900bbd0577833a4b4a7ea6ebfe7d43732fcefb22a07643dc38522674d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.