Buckets:
| """cwm3/probes/common.py — shared helpers of the latent-quality diagnostic study (docs/latent_quality.md). | |
| Everything here is CPU/GPU-agnostic numpy + a little torch; read-only w.r.t. results/runs, results/pack, results/banks and every | |
| upstream cache. Outputs of the study live under results/probes/ and results/tables/latent_*.md only. | |
| Probe rows: one seeded uniform row subset per (regime, split) — results/probes/rows_{regime}_{split}.npz — shared by EVERY | |
| representation and every probe (same rows, same split); a smaller seeded subset of those rows carries the raw 16x1408 region tensor. | |
| Targets (build_targets): CURRENT state at the anchor t and FUTURE onsets within 3 s, with the population rules of cwm2/eval/populations.py | |
| (drv_onset: inactive at t and EM valid over the 15 steps; d_onset / p_onset: inactive & valid at t, Dvalid & Moving & M over 15 steps; | |
| veh_onset: pedal chain inactive at t, M over 15 steps). No openpilot DMS field is read anywhere (inputs: CAN whitelist, pose-derived | |
| kinematics, V-JEPA2 latents, radar/lane/BADAS-derived pressure targets). | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import os | |
| import sys | |
| import time | |
| import zipfile | |
| from pathlib import Path | |
| import numpy as np | |
| from numpy.lib import format as npf | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| CWM2 = EXPLORE / "cwm2" | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| KIN = CWM2 / "results" / "targets_kin" | |
| DEM = CWM2 / "results" / "targets_demand" | |
| PACK = CWM3 / "results" / "pack" | |
| BANKS = CWM3 / "results" / "banks" | |
| PRESS = CWM3 / "results" / "targets_pressure" | |
| PCA_DIR = CWM3 / "results" / "pca" | |
| RUNS = CWM3 / "results" / "runs" | |
| FF = EXPLORE / "s12_wm" / "cache" / "frontfeat" | |
| S5F = EXPLORE / "s5_features" / "vjepa2g384" / "front" | |
| S5C = EXPLORE / "s5_features" / "vjepa2g384" / "cabin" | |
| OUT = CWM3 / "results" / "probes" | |
| TABLES = CWM3 / "results" / "tables" | |
| for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train", EXPLORE / "s12_wm"): | |
| sys.path.insert(0, str(p)) | |
| REGIMES, SPLITS = ("manual", "engaged"), ("train", "val", "test") | |
| ROWS_N = {"train": 50_000, "val": 15_000, "test": 40_000} # per regime | |
| REGION_N = {"train": 15_000, "val": 5_000, "test": 15_000} # rows that also carry the raw region tensor | |
| TAU3 = 15 # 3 s = 15 steps of 0.2 s | |
| P_FAMILIES = ["ttc", "ttc_drop", "cut_in", "lead_braking", "lane_excursion", "badas"] | |
| D_FAMILIES = ["steer_rate", "decel", "ttc", "ttc_drop", "lane", "badas"] | |
| DRV_EVENTS = ("hand_return", "hand_leave", "head_turn_onset") | |
| SEED = 20260910 | |
| def log(msg, path=None): | |
| line = f"[{time.strftime('%H:%M:%S')}] {msg}" | |
| print(line, flush=True) | |
| if path is not None: | |
| with open(path, "a") as f: | |
| f.write(line + "\n") | |
| def stamp(): | |
| return time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()) | |
| # ----------------------------------------------------------------------------- npz access | |
| def npz_memmap(path, key): | |
| """Read-only memmap of an UNCOMPRESSED (stored) .npy member of an .npz (results/banks/front_tok.npz).""" | |
| with zipfile.ZipFile(path) as zf: | |
| info = zf.getinfo(key + ".npy") | |
| assert info.compress_type == zipfile.ZIP_STORED, f"{path}:{key} is compressed; use iter_member" | |
| with open(path, "rb") as fh: | |
| fh.seek(info.header_offset) | |
| # local file header: 30 bytes + name + extra | |
| hdr = fh.read(30) | |
| n_name, n_extra = int.from_bytes(hdr[26:28], "little"), int.from_bytes(hdr[28:30], "little") | |
| data0 = info.header_offset + 30 + n_name + n_extra | |
| fh.seek(data0) | |
| version = npf.read_magic(fh) | |
| shape, fortran, dtype = npf._read_array_header(fh, version) | |
| off = fh.tell() | |
| return np.memmap(path, dtype=dtype, mode="r", offset=off, shape=shape, order="F" if fortran else "C") | |
| def iter_member(path, key, chunk=16384, max_bytes=64 << 20): | |
| """Stream a (compressed) .npy member chunk by chunk -> (row0, block); chunks are capped at max_bytes.""" | |
| with zipfile.ZipFile(path) as zf, zf.open(key + ".npy") as fh: | |
| version = npf.read_magic(fh) | |
| shape, fortran, dtype = npf._read_array_header(fh, version) | |
| assert not fortran and dtype != object | |
| rowsz = int(np.prod(shape[1:], dtype=np.int64)) * dtype.itemsize | |
| chunk = int(max(64, min(chunk, max_bytes // max(rowsz, 1)))) | |
| for i in range(0, shape[0], chunk): | |
| k = min(chunk, shape[0] - i) | |
| buf = fh.read(k * rowsz) | |
| assert len(buf) == k * rowsz, f"{path}:{key} short read at {i}" | |
| yield i, np.frombuffer(buf, dtype).reshape((k,) + tuple(shape[1:])) | |
| def gather_member_rows(path, key, want_pos): | |
| """rows `want_pos` (sorted positions inside the member) of a compressed member, streamed once -> array [len(want_pos), ...].""" | |
| want_pos = np.asarray(want_pos, np.int64) | |
| out = None | |
| for i, blk in iter_member(path, key): | |
| lo, hi = np.searchsorted(want_pos, i), np.searchsorted(want_pos, i + len(blk)) | |
| if hi > lo: | |
| if out is None: | |
| out = np.zeros((len(want_pos),) + blk.shape[1:], blk.dtype) | |
| out[lo:hi] = blk[want_pos[lo:hi] - i] | |
| return out | |
| # ----------------------------------------------------------------------------- probe rows | |
| def _seed(*parts): | |
| h = hashlib.sha256(("|".join(map(str, parts)) + f"|{SEED}").encode()).digest() | |
| return int.from_bytes(h[:8], "little") | |
| def probe_rows(regime, split, verbose=print): | |
| """-> dict {rows int64[n] (sorted anchor rows of targets_{regime}_{split}), region bool[n], N_file}. Created once, then reused.""" | |
| OUT.mkdir(parents=True, exist_ok=True) | |
| p = OUT / f"rows_{regime}_{split}.npz" | |
| if p.exists(): | |
| z = np.load(p) | |
| return {"rows": z["rows"], "region": z["region"], "N_file": int(z["N_file"])} | |
| with np.load(S12R / f"targets_{regime}_{split}.npz", allow_pickle=True) as z: | |
| N = len(z["t"]) | |
| rng = np.random.default_rng(_seed("rows", regime, split)) | |
| n = min(ROWS_N[split], N) | |
| rows = np.sort(rng.choice(N, n, replace=False)).astype(np.int64) | |
| rng2 = np.random.default_rng(_seed("region", regime, split)) | |
| reg = np.zeros(n, bool) | |
| reg[rng2.choice(n, min(REGION_N[split], n), replace=False)] = True | |
| np.savez(p, rows=rows, region=reg, N_file=np.array(N), seed=np.array(SEED), | |
| rule=np.array("uniform without replacement over the s12 anchor grid; region subset uniform over the probe rows")) | |
| verbose(f"[rows] {regime}/{split}: {n:,} of {N:,} rows ({reg.sum():,} with region) -> {p.name}") | |
| return {"rows": rows, "region": reg, "N_file": N} | |
| # ----------------------------------------------------------------------------- raw arrays for the probe rows | |
| def raw_path(regime, split): | |
| return OUT / f"raw_{regime}_{split}.npz" | |
| def load_raw(regime, split): | |
| """The compact per-row arrays of the probe rows (written by probes/gather_features.py).""" | |
| z = np.load(raw_path(regime, split), allow_pickle=True) | |
| return {k: z[k] for k in z.files} | |
| def feats_path(regime, split): | |
| return OUT / f"feats_{regime}_{split}.npz" | |
| # ----------------------------------------------------------------------------- targets | |
| def _first_onset(z0, z, em): | |
| 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 | |
| return inact0, onset, first | |
| def build_targets(raw, tau=TAU3): | |
| """-> OrderedDict name -> {y bool[n], m bool[n] (population membership), kind cur|fut, family}. Rules = cwm2 populations.""" | |
| import episodes as ep | |
| n = len(raw["t"]) | |
| T = {} | |
| M = raw["M"].astype(bool) | |
| S, S0, EM = raw["S"].astype(np.float32), raw["S0"].astype(np.float32), raw["EM"].astype(bool) & M[:, :, None] | |
| # ---- current | |
| T["cur_hands_on"] = dict(y=S0[:, 0] > 0.5, m=np.isfinite(S0[:, 0]), kind="cur", family="driver") | |
| T["cur_head_turned"] = dict(y=S0[:, 1] > 0.5, m=np.isfinite(S0[:, 1]), kind="cur", family="driver") | |
| for j, f in enumerate(P_FAMILIES): | |
| T[f"cur_P0_{f}"] = dict(y=raw["P0"][:, j].astype(bool), m=raw["P0valid"][:, j].astype(bool), kind="cur", family="pressure") | |
| v = raw["X"][:, -1, 3].astype(np.float32) | |
| for name, (lo, hi) in {"stopped": (-1e9, 1.0), "slow": (1.0, 12.0), "fast": (12.0, 1e9)}.items(): | |
| T[f"cur_vego_{name}"] = dict(y=(v >= lo) & (v < hi), m=np.isfinite(v), kind="cur", family="vehicle") | |
| gas0, brake0 = raw["X"][:, -1, 1] > 0.5, raw["X"][:, -1, 2] > 0.5 | |
| T["cur_brake"] = dict(y=brake0, m=np.ones(n, bool), kind="cur", family="vehicle") | |
| T["cur_gas"] = dict(y=gas0, m=np.ones(n, bool), kind="cur", family="vehicle") | |
| # ---- future driver onsets (P_drv) | |
| chains = [(S0[:, 0], S[:, :, 0], EM[:, :, 0]), (1 - S0[:, 0], 1 - S[:, :, 0], EM[:, :, 1]), (S0[:, 1], S[:, :, 1], EM[:, :, 2])] | |
| for ev, (z0, z, em) in zip(DRV_EVENTS, chains): | |
| inact0, onset, first = _first_onset(z0, z, em) | |
| m = inact0 & em[:, :tau].all(1) | |
| T[f"fut_{ev}"] = dict(y=(first >= 0) & (first < tau), m=m, kind="fut", family="driver") | |
| # ---- future pressure / demand onsets (P_dem rule on P and on D) | |
| for pref, Z, Zv, Z0, Z0v, mov, fams in (("P", raw["P"], raw["Pvalid"], raw["P0"], raw["P0valid"], raw["PMoving"], P_FAMILIES), | |
| ("D", raw["D"], raw["Dvalid"], raw["D0"], raw["D0valid"], raw.get("Moving", raw["PMoving"]), D_FAMILIES)): | |
| # (cwm3_dataset does not carry the per-step Moving of targets_demand; the pressure file's Moving is the same vEgo > 2 m/s gate on the same grid — asserted equal on manual/val by the audit) | |
| for j, f in enumerate(fams): | |
| valid = Zv[:, :, j].astype(bool) & mov.astype(bool) & M | |
| z = Z[:, :, j].astype(np.float32) | |
| z0 = Z0[:, j].astype(np.float32) | |
| zp = np.concatenate([z0[:, None], z[:, :-1]], 1) | |
| onset = (z > 0.5) & (zp < 0.5) & valid | |
| first = np.where(onset.any(1), onset.argmax(1), -1) | |
| inact0 = (z0 < 0.5) & Z0v[:, j].astype(bool) | |
| m = inact0 & valid[:, :tau].all(1) | |
| T[f"fut_{pref}_{f}"] = dict(y=(first >= 0) & (first < tau), m=m, kind="fut", family="pressure" if pref == "P" else "demand") | |
| # ---- future vehicle onsets (P_veh) | |
| for ev, ch in (("brake_on", "brakePressed"), ("gas_on", "gasPressed")): | |
| zp, z = ep.pedal_chain_targets(raw["Y"].astype(np.float32), raw["X"].astype(np.float32), ch) | |
| onset = (z > 0.5) & (zp < 0.5) & M | |
| first = np.where(onset.any(1), onset.argmax(1), -1) | |
| m = (zp[:, 0] < 0.5) & M[:, :tau].all(1) | |
| T[f"fut_{ev}"] = dict(y=(first >= 0) & (first < tau), m=m, kind="fut", family="vehicle") | |
| return T | |
| def target_table(T, dongle): | |
| rows = [] | |
| for k, v in T.items(): | |
| y, m = v["y"][v["m"]], v["m"] | |
| rows.append({"target": k, "kind": v["kind"], "family": v["family"], "n": int(m.sum()), "n_pos": int(y.sum()), | |
| "base_rate": float(y.mean()) if m.any() else np.nan, "n_dev_pos": int(len(np.unique(dongle[v["m"]][y])))}) | |
| return rows | |
| # ----------------------------------------------------------------------------- metrics | |
| def ap_weighted(y, s, w=None): | |
| y = np.asarray(y, bool) | |
| s = np.asarray(s, np.float64) | |
| w = np.ones(len(y)) if w is None else np.asarray(w, np.float64) | |
| order = np.argsort(-s, kind="mergesort") | |
| y, w = y[order], w[order] | |
| tp, fp = np.cumsum(w * y), np.cumsum(w * ~y) | |
| if tp[-1] <= 0 or fp[-1] <= 0: | |
| return np.nan | |
| prec = tp / np.maximum(tp + fp, 1e-12) | |
| rec = tp / tp[-1] | |
| drec = np.diff(np.r_[0.0, rec]) | |
| return float((prec * drec).sum()) | |
| def auroc(y, s): | |
| y = np.asarray(y, bool) | |
| s = np.asarray(s, np.float64) | |
| n1, n0 = y.sum(), (~y).sum() | |
| if n1 == 0 or n0 == 0: | |
| return np.nan | |
| r = np.empty(len(s)) | |
| order = np.argsort(s, kind="mergesort") | |
| ss = s[order] | |
| # average ranks for ties | |
| i = 0 | |
| while i < len(ss): | |
| j = i | |
| while j + 1 < len(ss) and ss[j + 1] == ss[i]: | |
| j += 1 | |
| r[order[i:j + 1]] = 0.5 * (i + j) + 1 | |
| i = j + 1 | |
| return float((r[y].sum() - n1 * (n1 + 1) / 2) / (n1 * n0)) | |
| def cluster_boot(y, s, groups, n_boot=500, seed=0, s_b=None): | |
| """device-cluster bootstrap of AP (and of AP_a - AP_b when s_b is given) -> dict.""" | |
| y = np.asarray(y, bool) | |
| uniq, dev = np.unique(groups, return_inverse=True) | |
| rng = np.random.default_rng(seed) | |
| ap = ap_weighted(y, s) | |
| res = {"ap": ap, "auroc": auroc(y, s), "n": int(len(y)), "n_pos": int(y.sum()), "n_dev": int(len(uniq))} | |
| if s_b is not None: | |
| res["ap_b"] = ap_weighted(y, s_b) | |
| res["delta"] = ap - res["ap_b"] | |
| ba, bd = np.full(n_boot, np.nan), np.full(n_boot, np.nan) | |
| 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 * y).sum() <= 0 or (w * ~y).sum() <= 0: | |
| continue | |
| ba[i] = ap_weighted(y, s, w) | |
| if s_b is not None: | |
| bd[i] = ba[i] - ap_weighted(y, s_b, w) | |
| res["ci"] = (float(np.nanpercentile(ba, 2.5)), float(np.nanpercentile(ba, 97.5))) | |
| if s_b is not None: | |
| ok = np.isfinite(bd) | |
| res["delta_ci"] = (float(np.nanpercentile(bd, 2.5)), float(np.nanpercentile(bd, 97.5))) | |
| res["frac_gt0"] = float((bd[ok] > 0).mean()) if ok.any() else np.nan | |
| return res | |
| # ----------------------------------------------------------------------------- standardisation + probes | |
| class Std: | |
| def __init__(self, X, clip=8.0): | |
| X = np.asarray(X, np.float32) | |
| self.mu = np.nanmean(X, 0).astype(np.float32) | |
| self.sd = (np.nanstd(X, 0) + 1e-6).astype(np.float32) | |
| self.clip = clip | |
| def __call__(self, X): | |
| Z = (np.asarray(X, np.float32) - self.mu) / self.sd | |
| return np.nan_to_num(np.clip(Z, -self.clip, self.clip), nan=0.0) | |
| def torch_device(): | |
| import torch | |
| if torch.cuda.is_available() and os.environ.get("PROBES_CPU", "0") != "1": | |
| try: | |
| free, total = torch.cuda.mem_get_info() | |
| if free / 2**30 >= 8.0: # need (<= 4 GB here) + 4 GB headroom | |
| return torch.device("cuda") | |
| except Exception: | |
| pass | |
| return torch.device("cpu") | |
| RIDGE_ALPHAS = (1e-3, 1e-2, 1e-1, 1.0, 10.0) | |
| def ridge_fit_predict(Xtr, Ytr, Mtr, Xva, Xte, alphas=RIDGE_ALPHAS, device=None): | |
| """Closed-form ridge, one solve per target on ITS masked rows (labels centred), lambda = alpha * trace(G)/d. | |
| Xtr/Xva/Xte standardised float32; Ytr [n,T] float (0/1), Mtr [n,T] bool. -> scores va [T, A, n_va], te [T, A, n_te].""" | |
| import torch | |
| dev = device or torch_device() | |
| Xt = torch.as_tensor(Xtr, device=dev) | |
| Xv = torch.as_tensor(Xva, device=dev) | |
| Xe = torch.as_tensor(Xte, device=dev) | |
| n, d = Xt.shape | |
| T = Ytr.shape[1] | |
| A = len(alphas) | |
| sva = np.zeros((T, A, len(Xva)), np.float32) | |
| ste = np.zeros((T, A, len(Xte)), np.float32) | |
| ones = torch.ones(n, 1, device=dev) | |
| Xb = torch.cat([Xt, ones], 1) # bias column | |
| for j in range(T): | |
| m = torch.as_tensor(Mtr[:, j], device=dev) | |
| if int(m.sum()) < 50: | |
| sva[j], ste[j] = np.nan, np.nan | |
| continue | |
| Xm = Xb[m] | |
| y = torch.as_tensor(Ytr[:, j], device=dev, dtype=torch.float32)[m] | |
| y = y - y.mean() | |
| G = Xm.T @ Xm | |
| b = Xm.T @ y | |
| tr = torch.trace(G[:d, :d]) / d | |
| evals, V = torch.linalg.eigh(G.double()) | |
| Vb = V.T @ b.double() | |
| for a, al in enumerate(alphas): | |
| lam = al * tr.double() | |
| w = (V @ (Vb / (evals + lam))).float() | |
| sva[j, a] = (Xv @ w[:d] + w[d]).cpu().numpy() | |
| ste[j, a] = (Xe @ w[:d] + w[d]).cpu().numpy() | |
| return sva, ste | |
| def mlp_fit_predict(Xtr, Ytr, Mtr, Xva, Yva, Mva, Xte, hidden=64, epochs=30, lr=1e-3, wd=1e-4, batch=1024, seed=0, device=None, | |
| dongle_va=None): | |
| """Multi-task 2-layer MLP (d -> hidden -> T), masked BCE; per-target early selection on VAL AP (a disjoint device set). | |
| -> test scores [T, n_te] at each target's best-val epoch, val scores [T, n_va], best epochs [T].""" | |
| import torch | |
| import torch.nn as nn | |
| dev = device or torch_device() | |
| torch.manual_seed(seed) | |
| n, d = Xtr.shape | |
| T = Ytr.shape[1] | |
| net = nn.Sequential(nn.Linear(d, hidden), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden, T)).to(dev) | |
| opt = torch.optim.AdamW(net.parameters(), lr=lr, weight_decay=wd) | |
| Xt = torch.as_tensor(Xtr, device=dev) | |
| Yt = torch.as_tensor(Ytr, device=dev, dtype=torch.float32) | |
| Mt = torch.as_tensor(Mtr, device=dev, dtype=torch.float32) | |
| Xv = torch.as_tensor(Xva, device=dev) | |
| Xe = torch.as_tensor(Xte, device=dev) | |
| # positive re-weighting per target (rare onsets): w_pos = n_neg / n_pos, capped | |
| npos = (Yt * Mt).sum(0) | |
| nneg = ((1 - Yt) * Mt).sum(0) | |
| wpos = torch.clamp(nneg / torch.clamp(npos, min=1.0), 1.0, 50.0) | |
| best_ap = np.full(T, -1.0) | |
| best_ep = np.zeros(T, int) | |
| best_te = np.zeros((T, len(Xte)), np.float32) | |
| best_va = np.zeros((T, len(Xva)), np.float32) | |
| rng = np.random.default_rng(seed) | |
| for ep in range(epochs): | |
| net.train() | |
| perm = torch.as_tensor(rng.permutation(n), device=dev) | |
| for i in range(0, n, batch): | |
| idx = perm[i:i + batch] | |
| logit = net(Xt[idx]) | |
| y, m = Yt[idx], Mt[idx] | |
| w = m * (1 + (wpos - 1) * y) | |
| loss = (nn.functional.binary_cross_entropy_with_logits(logit, y, reduction="none") * w).sum() / w.sum().clamp_min(1.0) | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| opt.step() | |
| net.eval() | |
| with torch.no_grad(): | |
| sv = torch.cat([net(Xv[i:i + 8192]) for i in range(0, len(Xv), 8192)]).cpu().numpy().T | |
| se = torch.cat([net(Xe[i:i + 8192]) for i in range(0, len(Xe), 8192)]).cpu().numpy().T | |
| for j in range(T): | |
| m = Mva[:, j] | |
| if m.sum() < 50 or Yva[m, j].sum() < 5: | |
| continue | |
| ap = ap_weighted(Yva[m, j], sv[j][m]) | |
| if np.isfinite(ap) and ap > best_ap[j]: | |
| best_ap[j], best_ep[j], best_te[j], best_va[j] = ap, ep, se[j], sv[j] | |
| return best_te, best_va, best_ep, best_ap | |
| def write_json(path, obj): | |
| Path(path).parent.mkdir(parents=True, exist_ok=True) | |
| Path(path).write_text(json.dumps(obj, indent=1, default=lambda o: o.tolist() if hasattr(o, "tolist") else str(o))) | |
| def md_table(header, rows): | |
| out = ["| " + " | ".join(header) + " |", "|" + "---|" * len(header)] | |
| for r in rows: | |
| out.append("| " + " | ".join(str(x) for x in r) + " |") | |
| return "\n".join(out) | |
| def f3(x): | |
| return "–" if x is None or (isinstance(x, float) and not np.isfinite(x)) else f"{x:.3f}" | |
| def ci3(v, lo, hi): | |
| if v is None or not np.isfinite(v): | |
| return "–" | |
| return f"{v:.3f} [{lo:.3f}, {hi:.3f}]" | |
Xet Storage Details
- Size:
- 19.2 kB
- Xet hash:
- 4f44d038e89917f6dd15c1ea98c6a7927674f1854cb5edea6f17c78d01d9f215
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.