Buckets:
| """cwm3/eval/write_pred_layout.py — a cwm3 branch's prediction file -> the cwm2 pred_layout.py contract (pred_layout_{regime}_{split}.npz). | |
| Fork of cwm2/eval/pred_adapter.py restricted to what ONE base branch emits (no sampled joint paths, no cross-leg endpoints, | |
| no warning columns): everything else is NaN / False / absent and listed in meta["absent_blocks"]. | |
| Input: <run_dir>/pred_{regime}_{split}.npz written by train/train_cwm3.py: hz_on / hz_off f16[N,25,C] (head hazards, all N rows | |
| of targets_{regime}_{split}.npz), valid bool[N] (rows the branch predicted, e.g. x8_ok for V can8s), z0 / z0valid, chain_names, | |
| plus per-branch extras: V -> y_mae[N,3], y_crps[N,3], pi_resp[N,25]; E -> f_cos[N,5]. | |
| Output blocks: | |
| onset_p[N,4,E] 1 - prod_{k<tau}(1 - h_on,k) for the events the branch emits (V: gas_on, brake_on, steer_phase_on, resp_on; | |
| U: hand_return = hands_on.on, hand_leave = hands_on.off, head_turn_onset = head_turned.on, head_return = head_turned.off; | |
| E: D_<family>) — NaN for every other EVENT_NAMES entry and for rows with valid = False | |
| h_on[N,25,E] the same per-step hazards (hazard-NLL cells) | |
| ep_p / ep_mask NaN / False (needs joint paths — not a base-branch quantity) w / w_names / w_mask empty / False | |
| nr_p[N,25] V only: P^(NR_{k,15}) = (1 - pi_k) * prod_{j=k+1..k+15}(1 - h_on_j) under the resp chain (exact Markov form; NaN where k+15 > 24) | |
| y_mae / y_crps V only (per-row means over valid steps, steer / vEgo / aEgo) f_cos[N,5] E only | |
| uniform_mask = valid (every predicted row is a uniform row: the branch predicts the whole split) | |
| meta arm cwm3_{V,U,E}, regime, seed, split, score_source hazard_closed_form, roi, model_id, code_sha256, inputs_digest, | |
| dms_fields_used [], external_priors, stage / variant, absent_blocks | |
| No openpilot DMS field is read (targets file only for the uid / t witness). | |
| python write_pred_layout.py --run-dir results/runs/U0_kin_can_s0 [--regimes manual,engaged] [--splits test,val] | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import sys | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| CWM2 = EXPLORE / "cwm2" | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| for p in (CWM2 / "eval", CWM3 / "common"): | |
| sys.path.insert(0, str(p)) | |
| import first_passage as fp # noqa: E402 | |
| from pred_layout import EVENT_NAMES, EP_NAMES, ONSET_TAUS # noqa: E402 | |
| from runmeta import code_hash, sha256_file # noqa: E402 | |
| T = 25 | |
| ARM_OF = {"V": "cwm3_V", "U": "cwm3_U", "E": "cwm3_E"} | |
| ROI_OF = {"V": "none", "U": "route-median (non-causal)", "E": "none"} # U: Kh / K / S come from the route-median crop pose cache | |
| EXTERNAL_PRIORS = {"V": [], "U": ["YOLO11x-pose: driver kinematics Kh/K/S (route-median ROI, non-causal)", "V-JEPA2 ViT-g 384: cabin PCA-256 (vc / crop variants)"], | |
| "E": ["V-JEPA2 ViT-g 384: front region tokens / pooled Vf (PCA) and F targets", "OpenBADAS: `badas` family of the D target"]} | |
| # (event name, chain name, on|off) | |
| EVENT_SRC = {"V": [("gas_on", "gas", "on"), ("brake_on", "brake", "on"), ("steer_phase_on", "steer_phase", "on"), ("resp_on", "resp", "on")], | |
| "U": [("hand_return", "hands_on", "on"), ("hand_leave", "hands_on", "off"), ("head_turn_onset", "head_turned", "on"), | |
| ("head_return", "head_turned", "off")], | |
| "E": [(f"D_{f}", f, "on") for f in ("steer_rate", "decel", "ttc", "ttc_drop", "lane", "badas")]} | |
| def write_layout(run_dir, regime, split, run_cfg, code_files, w_window=15, log=print): | |
| rd = Path(run_dir) | |
| pf = rd / f"pred_{regime}_{split}.npz" | |
| z = np.load(pf, allow_pickle=True) | |
| branch, seed = str(z["branch"]), int(z["seed"]) | |
| tg = np.load(S12R / f"targets_{regime}_{split}.npz", allow_pickle=True) | |
| N = len(tg["t"]) | |
| assert z["hz_on"].shape[0] == N, f"{pf}: {z['hz_on'].shape[0]} rows, targets {N}" | |
| valid = z["valid"].astype(bool) | |
| chains = [str(c) for c in z["chain_names"]] | |
| hz_on, hz_off = z["hz_on"].astype(np.float32), z["hz_off"].astype(np.float32) | |
| E = len(EVENT_NAMES) | |
| h = np.full((N, T, E), np.nan, np.float32) | |
| emitted = [] | |
| for ev, ch, kind in EVENT_SRC[branch]: | |
| e = EVENT_NAMES.index(ev) | |
| h[:, :, e] = (hz_on if kind == "on" else hz_off)[..., chains.index(ch)] | |
| emitted.append(ev) | |
| h[~valid] = np.nan | |
| onset_p = np.full((N, len(ONSET_TAUS), E), np.nan, np.float32) | |
| for ti, tau in enumerate(ONSET_TAUS): | |
| onset_p[:, ti, :] = 1 - np.prod(1 - h[:, :tau, :], 1) | |
| out = {"event_names": np.array(EVENT_NAMES), "onset_taus": np.array(ONSET_TAUS), "onset_p": onset_p.astype(np.float16), | |
| "h_on": h.astype(np.float16), "uniform_mask": valid.copy(), "uid": tg["uid"], "t": tg["t"], | |
| "ep_names": np.array(EP_NAMES), "ep_p": np.full((N, len(EP_NAMES)), np.nan, np.float16), "ep_mask": np.zeros((N, len(EP_NAMES)), bool), | |
| "w_names": np.array([], dtype="<U1"), "w": np.zeros((N, 0), np.float16), "w_mask": np.zeros(N, bool)} | |
| absent = ["ep_p (joint paths)", "w (warning scores)", "cooccur_p"] | |
| if branch == "V": | |
| out["y_mae"] = z["y_mae"].astype(np.float32) | |
| out["y_crps"] = z["y_crps"].astype(np.float32) | |
| # P^(NR_{k,w}) under the branch's own 2-state resp chain (exact Markov form, not per-step independence): | |
| # P(z_k = 0, z_{k+1..k+w} = 0) = (1 - pi_k) * prod_{j=k+1..k+w} (1 - h_on,j) (h_on = 0->1 hazard of the resp chain) | |
| # (the independence product prod (1 - pi_j) ignores the strong persistence of the response state: NoResp is 0.59 per step | |
| # while Resp is 0.26 per step; on V0_can4s_s0_rpe400k it gave nr_brier 0.344 vs 0.111 for can_only_direct's path form) | |
| pi_resp = z["pi_resp"].astype(np.float64) | |
| h_resp = hz_on[..., chains.index("resp")].astype(np.float64) | |
| lc = np.concatenate([np.zeros((N, 1)), np.cumsum(np.log1p(-np.clip(h_resp, 0, 1 - 1e-12)), 1)], 1) # lc[k] = sum_{j<k} log(1-h_j) | |
| nr = np.full((N, T), np.nan) | |
| for k in range(T): | |
| if k + w_window <= T - 1: | |
| nr[:, k] = (1 - pi_resp[:, k]) * np.exp(lc[:, k + w_window + 1] - lc[:, k + 1]) | |
| nr[~valid] = np.nan | |
| out["nr_p"] = nr.astype(np.float16) | |
| meta_nr = "chain_form: (1-pi_k) * prod_{j=k+1..k+w}(1-h_on_j) from the resp chain" | |
| for k in ("y_mae", "y_crps"): | |
| out[k][~valid] = np.nan | |
| else: | |
| absent += ["y_mae", "y_crps", "nr_p"] | |
| meta_nr = "" | |
| if branch == "E": | |
| fc = z["f_cos"].astype(np.float32) | |
| fc[~valid] = np.nan | |
| out["f_cos"] = fc | |
| else: | |
| absent.append("f_cos") | |
| p_block = "" | |
| if "hz_on_P" in z.files: | |
| # road-pressure chains of the E0 `*_P` variants (2026-09-11): EXTRA keys with the joint layout's names (p_names / p_onset_p / p_pi / p_z0, | |
| # + p_h_on) so the warning head consumes them unchanged; the 14 contract event columns are untouched. Column names P_<chain>. | |
| hon = np.clip(z["hz_on_P"].astype(np.float64), 1e-6, 1 - 1e-6) | |
| hoff = np.clip(z["hz_off_P"].astype(np.float64), 1e-6, 1 - 1e-6) | |
| zz0 = z["z0_P"].astype(np.float64) | |
| op = np.stack([1 - np.prod(1 - hon[:, :tau], 1) for tau in ONSET_TAUS], 1) # [N,4,C] | |
| pim = np.empty_like(hon) | |
| prev = zz0.copy() | |
| for k in range(hon.shape[1]): | |
| prev = prev * (1 - hoff[:, k]) + (1 - prev) * hon[:, k] | |
| pim[:, k] = prev | |
| for a in (op, pim, hon): | |
| a[~valid] = np.nan | |
| out.update(p_names=np.array([f"P_{str(c)}" for c in z["p_chain_names"]]), p_onset_p=op.astype(np.float16), p_h_on=hon.astype(np.float16), | |
| p_pi=pim.astype(np.float16), p_z0=zz0.astype(np.float16), p_z0valid=z["z0valid_P"].astype(bool)) | |
| p_block = "p_names / p_onset_p[N,4,6] / p_h_on[N,25,6] / p_pi / p_z0 (road-pressure chains, targets_pressure; extra keys, contract columns unchanged)" | |
| g_block = "" | |
| if "hz_on_G" in z.files: | |
| # gaze_off chain of the U0 `*_gz*` variants (2026-09-12): EXTRA keys for the warning head — gaze_off_onset_p[N,4] (P(first 0->1 within tau | z0), the 4 | |
| # ONSET_TAUS), gaze_off_pi[N,25] (forward marginal P(gaze_off = 1 at t+k) from z0), gaze_off_h_on / h_off[N,25], gaze_off_z0 / z0valid[N]; NaN on rows | |
| # the branch did not predict or whose anchor state is undefined. The 14 contract event columns are untouched (evaluators ignore these keys). | |
| hon = np.clip(z["hz_on_G"].astype(np.float64)[..., 0], 1e-6, 1 - 1e-6) | |
| hoff = np.clip(z["hz_off_G"].astype(np.float64)[..., 0], 1e-6, 1 - 1e-6) | |
| zz0 = z["z0_G"].astype(np.float64)[:, 0] | |
| z0v = z["z0valid_G"].astype(bool)[:, 0] | |
| op = np.stack([1 - np.prod(1 - hon[:, :tau], 1) for tau in ONSET_TAUS], 1) # [N,4] | |
| pim = np.empty_like(hon) | |
| prev = zz0.copy() | |
| for k in range(hon.shape[1]): | |
| prev = prev * (1 - hoff[:, k]) + (1 - prev) * hon[:, k] | |
| pim[:, k] = prev | |
| bad = ~(valid & z0v) | |
| for arr in (op, pim, hon, hoff): | |
| arr[bad] = np.nan | |
| zz0[bad] = np.nan | |
| out.update(gaze_off_onset_p=op.astype(np.float16), gaze_off_pi=pim.astype(np.float16), gaze_off_h_on=hon.astype(np.float16), gaze_off_h_off=hoff.astype(np.float16), | |
| gaze_off_z0=zz0.astype(np.float16), gaze_off_z0valid=(valid & z0v), gaze_off_taus=np.array(ONSET_TAUS)) | |
| g_block = ("gaze_off_onset_p[N,4] / gaze_off_pi[N,25] / gaze_off_h_on / gaze_off_h_off / gaze_off_z0 / gaze_off_z0valid (causal gaze_off chain, " | |
| "results/targets_gaze_causal_routes_v2; extra keys for the warning head, contract columns unchanged)") | |
| meta = {"arm": ARM_OF[branch], "regime": regime, "seed": seed, "split": split, "score_source": "hazard_closed_form", "n_samples": 0, | |
| "sample_seed": "", "roi": run_cfg.get("roi", ROI_OF[branch]), "roi_streams": run_cfg.get("roi_streams", []), "driver_input_shift_s": 0, | |
| "model_id": f"{run_cfg.get('stage', '')}_{run_cfg.get('variant', '')}_s{seed}:{str(z['best_ckpt'])}", | |
| "code_sha256": code_hash(code_files), "inputs_digest": sha256_file(pf), "external_priors": EXTERNAL_PRIORS[branch], | |
| "dms_fields_used": [], "hazard_kind": "head", "n_rows_predicted": int(valid.sum()), "n_rows_uniform": int(valid.sum()), | |
| "events_emitted": emitted, "absent_blocks": absent, "stage": run_cfg.get("stage", ""), "variant": run_cfg.get("variant", ""), | |
| "branch": branch, "cfg": run_cfg.get("cfg", {}), "nr_p_form": meta_nr, "p_block": p_block, "gaze_block": g_block, | |
| "notes": "cwm3 base branch: head hazards -> closed-form onset probabilities on every predicted row; NaN for events the branch " | |
| "does not model and for rows without a valid input history; no joint paths, no warning columns"} | |
| outp = rd / f"pred_layout_{regime}_{split}.npz" | |
| np.savez_compressed(outp, meta=np.array(json.dumps(meta)), **out) | |
| log(f"[layout] {outp.name}: N={N} predicted={int(valid.sum())} events={emitted} absent={absent}") | |
| return outp | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--run-dir", required=True) | |
| ap.add_argument("--regimes", default="manual,engaged") | |
| ap.add_argument("--splits", default="test,val") | |
| a = ap.parse_args() | |
| rd = Path(a.run_dir) | |
| run = json.loads(sorted(rd.glob("run_train_*_latest.json"))[-1].read_text()) if list(rd.glob("run_train_*_latest.json")) else {} | |
| code_files = [Path(__file__), CWM2 / "eval" / "first_passage.py", CWM2 / "eval" / "pred_layout.py"] + sorted((CWM3 / "models").glob("*.py")) | |
| for r in a.regimes.split(","): | |
| for s in a.splits.split(","): | |
| if (rd / f"pred_{r}_{s}.npz").exists(): | |
| write_layout(rd, r, s, run, code_files) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 12.1 kB
- Xet hash:
- d71e1dcaed3c47c55b9e4843cfe5fdeeac6f1f96bb63fa730cd18e8647da3ea4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.