Buckets:
| """cwm3/eval/predict_u0_rows.py — run a U0 base run's BEST model on SPECIFIC TRAIN rows (the warning head's label rows) and write its gaze-chain outputs | |
| (gaze/robustness agent, 2026-09-12; coordinator: warning head + gaze features). The U0 `*_gz*` layouts exist for test / val only; the learned warning head | |
| is fitted on 979 TRAIN label rows (results/runs/I1_<main>_s<k>_final/label_rows.npz, data/labels.py), so the gaze chain has to be evaluated on those rows once. | |
| CPU (a few thousand rows; the data load dominates), the model / batch code is the trainer's own (train/train_cwm3.model_kwargs, train/batch.prep_U, unchanged). | |
| python eval/predict_u0_rows.py --run U0_kin_can_noK_s0_kcv2_clip8_ln_gzfc --rows-from results/runs/I1_all_minus_road2drv_s0_final/label_rows.npz | |
| -> results/runs/<run>/pred_gaze_{regime}_train_labelrows.npz: rows (targets-file rows), uid, t, gaze_off_onset_p f32[n,4] (taus 5 / 10 / 15 / 25 steps), | |
| gaze_off_pi f32[n,25], gaze_off_h_on / h_off f32[n,25], gaze_off_z0 f32[n], gaze_off_z0valid bool[n] (anchor state defined AND causal driver box present), | |
| gaze_off_taus, best_ckpt, source_rows. NaN on rows without a valid anchor state. No openpilot DMS field is read. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| os.environ.setdefault("MKL_NUM_THREADS", "4") | |
| import argparse # noqa: E402 | |
| import json # noqa: E402 | |
| import sys # noqa: E402 | |
| import time # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| EXPLORE = CWM3.parent | |
| for p in (CWM3, CWM3 / "data", CWM3 / "common", CWM3 / "train", CWM3 / "eval", EXPLORE / "s12_wm"): | |
| sys.path.insert(0, str(p)) | |
| import cwm3_dataset as ds # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| import train_cwm3 as T # noqa: E402 (model_kwargs / need_streams, unchanged) | |
| from models import build_branch # noqa: E402 | |
| from models import heads as hd # noqa: E402 | |
| from runmeta import write_run_json # noqa: E402 | |
| TAUS = (5, 10, 15, 25) | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--run", required=True, help="U0 run dir name under results/runs (a `u_gaze_chain` variant)") | |
| ap.add_argument("--rows-from", required=True, help="label_rows.npz of the joint run (keys manual / engaged = TRAIN targets-file rows)") | |
| ap.add_argument("--regimes", default="manual,engaged") | |
| ap.add_argument("--split", default="train") | |
| ap.add_argument("--force", action="store_true") | |
| a = ap.parse_args() | |
| import torch | |
| torch.set_num_threads(4) | |
| t0 = time.time() | |
| rd = CWM3 / "results" / "runs" / a.run | |
| cfg = json.loads((rd / "config.json").read_text())["cfg"] | |
| assert cfg["branch"] == "U" and cfg["recipe"].get("u_gaze_chain"), f"{a.run}: not a gaze-chain U0 run" | |
| seed = int(json.loads((rd / "config.json").read_text())["args"]["seed"]) | |
| ds.set_pack_version(cfg.get("pack_version", "v1")) | |
| ds.set_kin_source(cfg.get("kin_source", "cwm2"), cfg.get("kin_version", "v1")) | |
| best = (rd / "BEST").read_text().strip() | |
| z = torch.load(rd / best, map_location="cpu", weights_only=False) | |
| model = build_branch("U", seed, **T.model_kwargs(cfg)) | |
| model.load_state_dict(z["model"]) | |
| model.eval() | |
| st = {k: v for k, v in np.load(rd / "stats.npz").items()} | |
| lr = np.load(a.rows_from) | |
| need = T.need_streams(cfg) | |
| outs = [] | |
| for regime in a.regimes.split(","): | |
| op = rd / f"pred_gaze_{regime}_{a.split}_labelrows.npz" | |
| if op.exists() and not a.force: | |
| print(f"have {op}") | |
| outs.append(op) | |
| continue | |
| rows = np.sort(lr[regime].astype(np.int64)) | |
| raw = ds.load_regimes([regime], a.split, need, verbose=print) | |
| sub = ds.subset(raw, rows) | |
| del raw | |
| n = len(rows) | |
| assert len(sub["t"]) == n | |
| store = bt.GpuStore(sub, bt.KEYS["U"], "cpu") | |
| idx = torch.arange(n) | |
| with torch.no_grad(): | |
| b, tg = bt.prep_U(store, idx, st, cfg, True) | |
| o = model(b) | |
| assert "G" in o and "zG0" in tg | |
| hon, hoff = hd.hazards(o["G"].float().reshape(n, 25, 1, 2)) | |
| hon, hoff = hon[..., 0].clamp(1e-6, 1 - 1e-6), hoff[..., 0].clamp(1e-6, 1 - 1e-6) | |
| z0 = tg["zG0"].float() | |
| pi = hd.forward_marginal(hon, hoff, z0) | |
| onset = torch.stack([1 - torch.prod(1 - hon[:, :tau], 1) for tau in TAUS], 1) | |
| row_ok = sub["kin_row_ok"].astype(bool) if "kin_row_ok" in sub else np.ones(n, bool) | |
| z0v = tg["zG0valid"].numpy() & row_ok | |
| onset, pi, hon, hoff, z0 = (x.numpy().astype(np.float32) for x in (onset, pi, hon, hoff, z0)) | |
| for arr in (onset, pi, hon, hoff): | |
| arr[~z0v] = np.nan | |
| z0[~z0v] = np.nan | |
| np.savez_compressed(op, rows=rows, uid=sub["uid"], t=sub["t"], gaze_off_onset_p=onset, gaze_off_pi=pi, gaze_off_h_on=hon, gaze_off_h_off=hoff, gaze_off_z0=z0, | |
| gaze_off_z0valid=z0v, gaze_off_taus=np.array(TAUS), best_ckpt=np.array(best), run=np.array(a.run), source_rows=np.array(str(a.rows_from)), | |
| regime=np.array(regime), split=np.array(a.split)) | |
| print(f"{op.name}: {n} rows, anchor state defined {z0v.mean():.3f}, gaze_off@t {np.nanmean(z0):.3f}, onset_p@3s mean {np.nanmean(onset[:, 2]):.3f} ({time.time() - t0:.0f}s)", flush=True) | |
| outs.append(op) | |
| del sub, store | |
| write_run_json(str(rd), "predict_u0_rows", [Path(__file__), CWM3 / "train" / "batch.py", CWM3 / "models" / "branches.py"], [rd / best, Path(a.rows_from)], seed=seed, t_start=t0, | |
| input_listing=False, extra={"run": a.run, "rows_from": str(a.rows_from), "outputs": [str(p) for p in outs], "dms_fields_used": []}) | |
| print("PREDICT_U0_ROWS_DONE", flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 6.29 kB
- Xet hash:
- 36539c46fd4de7e2369848dd1bf958fdda8d82a9d4f25040251e09318f7ea4a8
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.