Buckets:
| #!/usr/bin/env python3 | |
| """cwm3/baselines_dms/chain_style.py — LAYER 2a of the DMS-baseline comparison (plan §1a (iii) / §2 第二层 (a), 2026-09-14): map a baseline's channels onto the | |
| THREE attention channels of the logged openpilot monitor chain exactly as docs/dms_att_chain.md §1.4 / §5 defines them, train the SAME two-state hazard chain | |
| (cwm2/baselines/legacy_cockpitwm/train_att.py recipe) on the baseline's own history, and compose `att@3s` with the SAME cwm2/compose.legacy_compose att_only | |
| (tau = 15 steps = 3 s, no shrink) -> pred layout results/runs/abl_dms_<b>_chain_s<k>/pred_layout_{regime}_{split}.npz for the UNCHANGED evaluators. | |
| MAPPING TABLE (legacy channel j -> label rule of build_targets_s15.att_states -> this module; every channel of every extractor is classified; CHANNEL_ROLES) | |
| j=0 `distracted` legacy: openpilot's composite instantaneous flag (head-pose deviation beyond the DM pose thresholds OR eye-closure; version-dependent constants, | |
| docs §5.2) -> here: |circ_dev(head_yaw)| > 20 deg OR |circ_dev(head_pitch)| > 15 deg (decisions.md 2026-09-14 fixed-rule constants) | |
| OR eye_closed > 0.5 (closed-probability channels averaged; open-ratio channels: ratio < 0.2, the OpenFace-3 SUMMARISE_SPEC constant); | |
| no hysteresis (the legacy labels have none, docs §1.4); NaN when the head-yaw channel is NaN (face invalid). | |
| A baseline without an eye-closure output defines channel 0 from head pose alone -> meta channels_undefined += "distracted.eye_closure". | |
| j=1 `gaze_off` legacy: abs(logged head-yaw field) > 0.4 rad = 22.918 deg on the HEAD yaw (despite the name; no per-driver offset) -> here: |circ_dev(head_yaw_deg)| > 22.918. | |
| j=2 `hands_off` legacy: pose-derived `hands_on_wheel < 0.5` from cabin body keypoints (NOT a DMS field; missing pose -> hands ON -> 0) -> here, for EVERY arm | |
| identically: our own causal pose channel min_hand_rest_dist (results/targets_kin_causal_routes_v2, Kh_c channel 10; hands off <=> >= WHEEL_TH 0.10, | |
| train/warning_features_ext.WHEEL_TH = s8_pose/posefeat WHEEL_TH) as the INPUT history and the pack's hands_on state (S0_c / S_c[:, :, 0]) as A0 / A | |
| (NaN = no causal box -> hands ON -> 0, the legacy rule). Face-only baselines cannot observe hands, so the shared channel is what docs §5.4 calls | |
| the "shared input" option; it is listed in meta (shared_inputs = ["hands_off (pose, targets_kin_causal_routes_v2)"]) and is the SAME for all arms. | |
| baseline | head yaw (j=0 pose, j=1) | head pitch (j=0) | eye closure (j=0) | unused channels (never enter the chain labels) | |
| ov_gaze | ov_head_yaw | ov_head_pitch | ov_eye_closed_l / _r (closed_p > 0.5) | ov_face_conf, ov_head_roll, ov_gaze_{x,y,z,yaw,pitch} (eye gaze; legacy j=1 is HEAD yaw) | |
| mp_face | mp_head_yaw | mp_head_pitch | mp_blink_l / _r (closed_p > 0.5) | mp_face_conf, mp_head_roll, mp_look_*, mp_squint_*, mp_wide_*, mp_ear_* (EAR kept out: blink prob is the closure output) | |
| openface3 | of3_head_yaw | of3_head_pitch | of3_eye_open_l / _r (open_ratio < 0.2) | of3_face_conf, of3_head_roll, of3_gaze_*, of3_au_*, of3_emo_* | |
| l2cs_gaze | l2cs_yaw (DEVIATION: eye-gaze yaw stands in for head yaw; L2CS has no head pose) | l2cs_pitch (deviation, same) | none -> "distracted.eye_closure" undefined | l2cs_face_conf, l2cs_box_w | |
| videomae2 | — (embedding + crop geometry only: no head-pose / eye-state channel -> chain-style arm NOT APPLICABLE, nothing written) | |
| ov_action | — (encoder embedding + 9 action logits: same -> NOT APPLICABLE) | |
| R0 lmface_rule | our own HALPE face pack (gaze_off state, eye_open) is the R0 rule arm (make_r0_layout.py); a chain-style arm on it would re-use CockpitWM's own | |
| face pack as a "baseline" and is not produced here. | |
| The INPUT history of the face chain (`ds` stream analogue) is the baseline's full per-frame channel set at t-3.9..t (10 Hz, 40 steps; NaN -> 0 as legacy :37), the | |
| legacy `ds` being all 7 DMS columns. Rows enter (train / predict) iff the cache has a frame at t, >= 50 % of the 40 history steps are valid (the legacy 4 s history `ok` rule of | |
| s15_cockpit/materialize_d*.py:36-38, docs §1.3) and A0 of channels 0 / 1 is defined (docs §1.5 keep rule; the crop row_ok term is dropped, see deviations). | |
| PROTOCOL (docs §5.3 "matched protocol"; recipe = train_att.train_chain verbatim: dilated TCN 128 / kernel 5 / dilations 1 2 4 8, 64-d step embedding, trunk 256, | |
| sigmoid on / off heads, two-state transition NLL, AdamW 3e-4 wd 0.01 cosine, <= 30 epochs, batch 1024, 300 k rows / epoch with flip rows up-weighted, early stop | |
| patience 6 on the val NLL, per regime, seed 0; channel routing of warn_compose.att_marginals: distracted / gaze_off <- the `ds`-only chain, hands_off <- the | |
| chain with the dynamics stream; compose.marginal -> pA[N,25,3]; legacy_compose(pA, None, 15)["att_only"]). | |
| DEVIATIONS from train_att.py, all documented in meta["deviations"]: (1) train_att.train_chain hard-codes `dev = "cuda"` and `tcn(7)` (7 DMS columns) — the replica | |
| below takes the channel count and the device as parameters (architecture otherwise identical; test_baselines_dms_layers asserts the state-dict shapes equal | |
| score_val_marginals.build_net for cin 7); (2) the legacy `cab` stream (s12 route-median crop feature, a NON-causal ROI) is NOT given to any baseline arm: the | |
| hands chain uses streams "ds,dyn" (ds = the 2 shared hands channels); (3) the loss mask is per channel (M & isfinite(A_c)) instead of all-channels-finite, because | |
| the hands channel is always defined while channels 0 / 1 may be NaN; (4) A0 uses the S0 tolerance (latest valid frame <= t within 0.15 s) instead of the frame at t. | |
| No-GPU references written beside the matched column: `persist_att3s` (train_att.persistence_hazards on the baseline's train labels; docs §5.3 (i)) and | |
| `state_att3s` = 1 - (1 - p_inatt(A0))^15 (docs §5.3 (ii)). | |
| OUTPUT columns dms_<b>_chain_att3s (primary; matched protocol), dms_<b>_chain_persist_att3s, dms_<b>_chain_state_att3s; w_mask = the keep rule; meta: mapping, | |
| thresholds, channels_undefined, shared_inputs, streams, recipe, deviations, checkpoints (sha256, best val NLL), dms_fields_used []. | |
| Chain checkpoints: results/baselines_dms/chain/ckpt_<b>_ds_{regime}_s<k>.pt and ckpt_shared_hands_ds-dyn_{regime}_s<k>.pt (append-only; the shared hands chain is | |
| trained once and reused by every arm). | |
| Nothing here reads a video, a model (only the baseline's cache files) or any openpilot field. | |
| python baselines_dms/chain_style.py --baseline ov_gaze [--seed 0] [--device cuda] # train (GPU) + layouts | |
| python baselines_dms/chain_style.py --dry-run # CPU: synthetic caches, 1 epoch, layout in a temp dir | |
| python baselines_dms/chain_style.py --print-mapping | |
| (2026-09-15, additive) --note <text>: inert launch note; results/baselines_dms/logs/layer1_tierD.sh trains the four face arms on the full train grid after the tier-D | |
| extraction, GPU when its launch gate allows (free VRAM >= 6 + 4 GB, host trainers < 3, DMS trainers < 2) else --device cpu. | |
| """ | |
| 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 | |
| PKG = Path(__file__).resolve().parent | |
| CWM3 = PKG.parent | |
| EXPLORE = CWM3.parent | |
| CWM2 = EXPLORE / "cwm2" | |
| LEGACY = CWM2 / "baselines" / "legacy_cockpitwm" | |
| S12R = EXPLORE / "s12_wm" / "results" | |
| for _p in (PKG, CWM3 / "data", CWM3 / "train", CWM3 / "common", CWM2, CWM2 / "eval", LEGACY): | |
| if str(_p) not in sys.path: | |
| sys.path.insert(0, str(_p)) | |
| import cache_schema as cs # noqa: E402 | |
| import summarise as SM # noqa: E402 | |
| import dataset_adapter as DA # noqa: E402 | |
| import layout as LY # noqa: E402 | |
| from compose import legacy_compose, marginal # noqa: E402 (cwm2/compose.py — the ONE composition implementation) | |
| H, TAU, HIST, HZ = 25, 15, 40, 10 | |
| STEP10 = 2 | |
| REGIMES = ("manual", "engaged") | |
| KIN_V2 = CWM3 / "results" / "targets_kin_causal_routes_v2" | |
| CHAIN_ROOT = cs.RESULTS_ROOT / "chain" | |
| ACH = ["distracted", "gaze_off", "hands_off"] | |
| TH = {"gaze_off_head_yaw_rad": 0.4, "gaze_off_head_yaw_deg": float(np.degrees(0.4)), "distracted_pose_yaw_deg": 20.0, "distracted_pose_pitch_deg": 15.0, | |
| "eye_closed_p": 0.5, "eye_open_ratio": 0.2, "hands_wheel_th": 0.10, "hist_valid_min": 0.5, "a0_tol_s": 0.15, | |
| "sources": {"gaze_off": "build_targets_s15.att_states :52 abs(head-yaw field) > 0.4 rad (head yaw)", "distracted_pose": "docs/decisions.md 2026-09-14 fixed-rule constants 20 / 15 deg", | |
| "eye_closed_p": "decisions.md closed_th 0.5", "eye_open_ratio": "extractors/openface3.SUMMARISE_SPEC (Soukupova & Cech 2016 EAR 0.2)", | |
| "hands_wheel_th": "train/warning_features_ext.WHEEL_TH = s8_pose/posefeat.py:23 WHEEL_TH 0.10", "hist_valid_min": "s15_cockpit/materialize_d*.py:36-38 legacy history ok rule (>= 50 % of the 40 steps valid)"}} | |
| RECIPE = {"tcn_dim": 128, "kernel": 5, "dilations": (1, 2, 4, 8), "step_emb": 64, "trunk": 256, "lr": 3e-4, "weight_decay": 0.01, "sched": "cosine(T_max=epochs)", "epochs": 30, | |
| "batch": 1024, "rows_per_epoch": 300_000, "flip_upweight": True, "patience": 6, "grad_clip": 1.0, "dropout": 0.1, "source": "cwm2/baselines/legacy_cockpitwm/train_att.py train_chain"} | |
| DEVIATIONS = ["train_att.train_chain hard-codes dev='cuda' and tcn(7): replica parameterises the ds channel count and the device (architecture identical)", | |
| "legacy `cab` stream (s12 route-median crop feature, non-causal ROI) not given to any baseline arm: hands chain streams ds,dyn", | |
| "loss mask per channel (M & isfinite(A_c)) instead of all-channels-finite", | |
| "A0 = latest valid frame <= t within 0.15 s (S0 tolerance) instead of exactly the frame at t"] | |
| SHARED_INPUTS = ["hands_off (our causal pose channel min_hand_rest_dist + hands_on state, results/targets_kin_causal_routes_v2; identical for every arm)"] | |
| # ----------------------------------------------------------------------------- mapping (every extractor channel classified) | |
| SPEC = { | |
| "ov_gaze": {"yaw": "ov_head_yaw", "pitch": "ov_head_pitch", "closed": ["ov_eye_closed_l", "ov_eye_closed_r"], "closed_kind": "closed_p", "note": "Intel OMZ head-pose + open-closed-eye"}, | |
| "mp_face": {"yaw": "mp_head_yaw", "pitch": "mp_head_pitch", "closed": ["mp_blink_l", "mp_blink_r"], "closed_kind": "closed_p", "note": "MediaPipe face transform pose + eyeBlink blendshapes"}, | |
| "openface3": {"yaw": "of3_head_yaw", "pitch": "of3_head_pitch", "closed": ["of3_eye_open_l", "of3_eye_open_r"], "closed_kind": "open_ratio", "note": "OpenFace 3 head pose + EAR"}, | |
| "l2cs_gaze": {"yaw": "l2cs_yaw", "pitch": "l2cs_pitch", "closed": None, "closed_kind": None, "note": "DEVIATION: eye-gaze angles stand in for head pose (L2CS has no head-pose output); no eye-closure output"}, | |
| "videomae2": None, | |
| "ov_action": None, | |
| } | |
| def _roles(): | |
| R = {} | |
| R["ov_gaze"] = {"ov_face_conf": "unused.confidence", "ov_head_yaw": "distracted.pose_yaw+gaze_off.head_yaw", "ov_head_pitch": "distracted.pose_pitch", "ov_head_roll": "unused.roll", | |
| **{f"ov_gaze_{k}": "unused.eye_gaze (legacy j=1 is head yaw)" for k in ("x", "y", "z", "yaw", "pitch")}, "ov_eye_closed_l": "distracted.eye_closed_p", "ov_eye_closed_r": "distracted.eye_closed_p"} | |
| R["mp_face"] = {"mp_face_conf": "unused.confidence", "mp_blink_l": "distracted.eye_closed_p", "mp_blink_r": "distracted.eye_closed_p", | |
| **{f"mp_look_{d}_{s}": "unused.eye_gaze_blendshape" for d in ("in", "out", "up", "down") for s in ("l", "r")}, | |
| **{f"mp_{k}_{s}": "unused.eyelid_blendshape" for k in ("squint", "wide") for s in ("l", "r")}, | |
| "mp_head_yaw": "distracted.pose_yaw+gaze_off.head_yaw", "mp_head_pitch": "distracted.pose_pitch", "mp_head_roll": "unused.roll", "mp_ear_l": "unused.eye_open_ratio (blink prob used)", "mp_ear_r": "unused.eye_open_ratio (blink prob used)"} | |
| R["openface3"] = {"of3_face_conf": "unused.confidence", "of3_gaze_yaw": "unused.eye_gaze", "of3_gaze_pitch": "unused.eye_gaze", "of3_head_yaw": "distracted.pose_yaw+gaze_off.head_yaw", | |
| "of3_head_pitch": "distracted.pose_pitch", "of3_head_roll": "unused.roll", **{f"of3_au_{i:02d}": "unused.action_unit" for i in range(1, 9)}, | |
| "of3_eye_open_l": "distracted.eye_open_ratio", "of3_eye_open_r": "distracted.eye_open_ratio", | |
| **{f"of3_emo_{e}": "unused.emotion (excluded)" for e in ("neutral", "happy", "sad", "surprise", "fear", "disgust", "anger", "contempt")}} | |
| R["l2cs_gaze"] = {"l2cs_face_conf": "unused.confidence", "l2cs_yaw": "distracted.pose_yaw+gaze_off.head_yaw (DEVIATION: eye-gaze yaw)", "l2cs_pitch": "distracted.pose_pitch (DEVIATION: eye-gaze pitch)", "l2cs_box_w": "unused.geometry"} | |
| R["videomae2"] = {"vmae2_crop_w": "geometry (n/a)", "vmae2_crop_h": "geometry (n/a)", "vmae2_lb_scale": "geometry (n/a)", "emb": "embedding (chain-style n/a)"} | |
| R["ov_action"] = {**{f"ova_enc_{i:03d}": "embedding (chain-style n/a)" for i in range(512)}, **{f"ova_logit_{i}": "logit (chain-style n/a)" for i in range(9)}} | |
| return R | |
| CHANNEL_ROLES = _roles() | |
| def applicable(baseline): | |
| return SPEC.get(baseline) is not None | |
| def channels_undefined(baseline): | |
| sp = SPEC[baseline] | |
| out = [] if sp.get("closed") else ["distracted.eye_closure"] | |
| return out | |
| # ----------------------------------------------------------------------------- channel states (docs §5.3 step 1) | |
| def _circ(x, ref=0.0): | |
| return SM.circ_dev(x, ref) if hasattr(SM, "circ_dev") else ((np.asarray(x, np.float32) - ref + 180.0) % 360.0) - 180.0 | |
| def channel_states(G, channels, spec, th=TH): | |
| """G f32[M,C] (NaN invalid) -> A f32[M,2] (distracted, gaze_off) in {0, 1, NaN}: NaN where the head-yaw channel is NaN (face invalid), as the legacy labels.""" | |
| ci = {c: j for j, c in enumerate(channels)} | |
| yaw = np.abs(_circ(G[:, ci[spec["yaw"]]])) | |
| valid = np.isfinite(yaw) | |
| pose = yaw > th["distracted_pose_yaw_deg"] | |
| if spec.get("pitch"): | |
| pitch = np.abs(_circ(G[:, ci[spec["pitch"]]])) | |
| pose |= np.nan_to_num(pitch) > th["distracted_pose_pitch_deg"] | |
| if spec.get("closed"): | |
| with np.errstate(all="ignore"): | |
| e = np.nanmean(np.stack([G[:, ci[c]] for c in spec["closed"]], 1), 1) | |
| closed = (e > th["eye_closed_p"]) if spec["closed_kind"] == "closed_p" else (e < th["eye_open_ratio"]) | |
| pose |= np.nan_to_num(closed.astype(np.float32)) > 0.5 | |
| gz = yaw > th["gaze_off_head_yaw_deg"] | |
| A = np.stack([pose, gz], 1).astype(np.float32) | |
| A[~valid] = np.nan | |
| return A | |
| def route_arrays(rc, anchors_t, spec, hist=HIST, tol_s=TH["a0_tol_s"]): | |
| """one route cache + anchor times -> dict(D f32[n,40,C] NaN invalid, Dvalid bool[n,40], A0 f32[n,2], A f32[n,25,2], M bool[n,25], row_ok bool[n]).""" | |
| D = SM.dense_series(rc) | |
| G, V, F, k0, M = D["G"], D["V"], D["F"], D["k0"], D["M"] | |
| Af = channel_states(G, rc.channels, spec) | |
| idx = SM.window_index(k0, M, anchors_t, hist) | |
| Dh, Dv = SM.gather(G, V, idx) | |
| ka = SM.anchor_steps(k0, anchors_t) | |
| inside = (ka >= 0) & (ka < M) | |
| row_ok = inside.copy() | |
| row_ok[inside] = F[ka[inside]] | |
| i0 = SM.s0_index(k0, M, V, anchors_t, tol_s) | |
| A0 = np.where((i0 >= 0)[:, None], Af[np.clip(i0, 0, max(M - 1, 0))], np.nan).astype(np.float32) if M else np.full((len(ka), 2), np.nan, np.float32) | |
| kf = ka[:, None] + (np.arange(1, H + 1) * STEP10)[None, :] # t + 0.2 .. t + 5.0 s (build_targets_s15.sample_route) | |
| Mf = (kf >= 0) & (kf < M) | |
| A = np.where(Mf[..., None], Af[np.clip(kf, 0, max(M - 1, 0))] if M else np.nan, np.nan).astype(np.float32) | |
| return {"D": Dh, "Dvalid": Dv, "A0": A0, "A": A, "M": Mf, "row_ok": row_ok} | |
| def hands_arrays(regime, split, rows=None, kin_dir=KIN_V2, th=TH["hands_wheel_th"]): | |
| """the SHARED hands_off channel from our causal pose pack -> dict(Dh f32[n,40,2] (hands_off state 0/1, min_hand_rest_dist; invalid -> 0 / NaN->0 at input), | |
| A0 f32[n] (1 - hands_on at t; NaN -> 0 = hands on), A f32[n,25] (1 - hands_on at t+0.2..t+5.0; NaN -> 0), M bool[n,25] (kin EM of the hands chain & row_ok)).""" | |
| with np.load(kin_dir / f"kinhist_causal_{regime}_{split}.npz", allow_pickle=True) as z: | |
| sel = slice(None) if rows is None else rows | |
| md = z["Kh_c"][sel][:, :, 10].astype(np.float32) | |
| kv = z["Khvalid"][sel] | |
| ok = z["row_ok"][sel].astype(bool) | |
| md = np.where(kv & ok[:, None], md, np.nan) | |
| hoff = np.where(np.isfinite(md), (md >= th).astype(np.float32), 0.0).astype(np.float32) | |
| Dh = np.stack([hoff, np.nan_to_num(md)], -1).astype(np.float32) | |
| with np.load(kin_dir / f"kin_causal_{regime}_{split}.npz", allow_pickle=True) as z: | |
| S0 = z["S0_c"][sel][:, 0].astype(np.float32) | |
| S = z["S_c"][sel][:, :, 0].astype(np.float32) | |
| EM = z["EM_c"][sel][:, :, 0].astype(bool) # the hands_on chain mask (train/batch.prep_U zUm: EM[:, :, 0]) | |
| ok2 = z["row_ok"][sel].astype(bool) | |
| A0 = np.where(np.isfinite(S0) & ok2, 1.0 - S0, 0.0).astype(np.float32) | |
| A = np.where(np.isfinite(S) & ok2[:, None], 1.0 - S, 0.0).astype(np.float32) | |
| return {"Dh": Dh, "A0": A0, "A": A, "M": EM & ok2[:, None], "row_ok": ok & ok2} | |
| def assemble(baseline, regime, split, cache_dir=None, spec=None, rows=None, log=print): | |
| """all rows of (regime, split): baseline D / A / A0 from the caches (rows of routes without a cache stay undefined) + shared hands + dyn X -> dict for train / predict.""" | |
| spec = spec or SPEC[baseline] | |
| cache_dir = Path(cache_dir) if cache_dir else cs.cache_dir(baseline) | |
| with np.load(S12R / f"targets_{regime}_{split}.npz", allow_pickle=True) as z: | |
| uid, t, X, Mtg, dongle = z["uid"], z["t"], z["X"].astype(np.float32), z["M"].astype(bool), z["dongle"] | |
| N = len(uid) | |
| C = None | |
| D = Dv = None | |
| A0 = np.full((N, 3), np.nan, np.float32) | |
| A = np.full((N, H, 3), np.nan, np.float32) | |
| Mm = np.zeros((N, H), bool) | |
| row_ok = np.zeros(N, bool) | |
| channels = None | |
| n_dec = 0 | |
| for u, rws in SM._route_blocks(uid): | |
| p = cache_dir / f"{u}.npz" | |
| if not p.exists(): | |
| continue | |
| rc = cs.RouteCache.load(p) | |
| if channels is None: | |
| channels = list(rc.channels) | |
| C = rc.C | |
| D = np.full((N, HIST, C), np.nan, np.float16) | |
| Dv = np.zeros((N, HIST), bool) | |
| assert list(rc.channels) == channels, f"{u}: channels differ" | |
| ra = route_arrays(rc, t[rws].astype(np.float64), spec) | |
| D[rws], Dv[rws], A0[rws, :2], A[rws, :, :2], Mm[rws], row_ok[rws] = ra["D"].astype(np.float16), ra["Dvalid"], ra["A0"], ra["A"], ra["M"], ra["row_ok"] | |
| n_dec += 1 | |
| if channels is None: | |
| raise FileNotFoundError(f"no route cache of {baseline} for {regime}/{split} under {cache_dir}") | |
| hd = hands_arrays(regime, split) | |
| A0[:, 2], A[:, :, 2] = hd["A0"], hd["A"] | |
| Mh = hd["M"] & Mtg | |
| Mm &= Mtg | |
| hist_share = Dv.mean(1) | |
| keep = row_ok & (hist_share >= TH["hist_valid_min"]) & np.isfinite(A0[:, :2]).all(1) | |
| log(f"{baseline} {regime}/{split}: N={N:,} decoded routes {n_dec} row_ok {row_ok.mean():.4f} keep {keep.mean():.4f} (hist>=50% {(hist_share >= 0.5).mean():.4f}) C={C}") | |
| return {"uid": uid, "t": t, "dongle": dongle, "X": X, "D": D, "Dvalid": Dv, "Dh": hd["Dh"], "A0": A0, "A": A, "M": Mm, "Mh": Mh, "keep": keep, "row_ok": row_ok, "channels": channels} | |
| # ----------------------------------------------------------------------------- the legacy chain (train_att.train_chain replica; cin / device parameterised) | |
| def build_net(streams, cin_ds, dim=128, cab_dim=4224): | |
| import torch | |
| import torch.nn as nn | |
| def tcn(cin, dim=dim): | |
| return nn.ModuleDict({"p": nn.Conv1d(cin, dim, 1), "b": nn.ModuleList([nn.Sequential(nn.Conv1d(dim, dim, 5, padding=4 * dd, dilation=dd), nn.GELU(), nn.Dropout(0.1)) for dd in (1, 2, 4, 8)])}) | |
| def run_tcn(m, x): | |
| h = m["p"](x.transpose(1, 2)) | |
| for blk in m["b"]: | |
| h = h + blk(h)[:, :, :h.shape[2]] | |
| return h[:, :, -1] | |
| class Net(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.ds = tcn(cin_ds) | |
| self.dyn = tcn(5) if "dyn" in streams else None | |
| self.cab = (nn.Sequential(nn.LayerNorm(cab_dim), nn.Linear(cab_dim, 128), nn.GELU(), nn.Dropout(0.1)) if "cab" in streams else None) | |
| dz = 128 + (128 if self.dyn else 0) + (128 if self.cab is not None else 0) | |
| self.step = nn.Embedding(H, 64) | |
| self.trunk = nn.Sequential(nn.Linear(dz + 64, 256), nn.GELU(), nn.Dropout(0.1)) | |
| self.on = nn.Linear(256, 3) | |
| self.off = nn.Linear(256, 3) | |
| def forward(self, D, X, Vc=None): | |
| zs = [run_tcn(self.ds, D)] | |
| if self.dyn is not None: | |
| zs.append(run_tcn(self.dyn, X)) | |
| if self.cab is not None: | |
| zs.append(self.cab(Vc)) | |
| z = torch.cat(zs, 1) | |
| B = z.shape[0] | |
| u = torch.cat([z[:, None].expand(B, H, -1), self.step.weight[None].expand(B, -1, -1)], -1) | |
| u = self.trunk(u) | |
| return torch.sigmoid(self.on(u)), torch.sigmoid(self.off(u)) | |
| return Net() | |
| def train_chain(tr, va, streams, seed=0, epochs=RECIPE["epochs"], device="cuda", rows_per_epoch=RECIPE["rows_per_epoch"], batch=RECIPE["batch"], patience=RECIPE["patience"], log=print): | |
| """tr / va: dict(D f32[n,40,cin], X f32[n,40,5], A f32[n,25,3] (NaN allowed), M bool[n,25,3] per-channel mask, A0 f32[n,3]) -> predict(d) -> (h_on, h_off) [n,25,3]; | |
| the train_att.train_chain loop with the channel count / device parameterised (see DEVIATIONS).""" | |
| import torch | |
| torch.manual_seed(seed) | |
| np.random.seed(seed) | |
| cin = tr["D"].shape[-1] | |
| xmu = tr["X"].reshape(-1, 5).mean(0) | |
| xsd = tr["X"].reshape(-1, 5).std(0) + 1e-6 | |
| def T(d): | |
| return (torch.from_numpy(np.ascontiguousarray(np.nan_to_num(d["D"].astype(np.float32)))), torch.from_numpy(np.ascontiguousarray(((d["X"] - xmu) / xsd).astype(np.float32))), | |
| torch.from_numpy(np.nan_to_num(d["A"], nan=0.0).astype(np.float32)), torch.from_numpy(d["M"].astype(np.float32)), torch.from_numpy(np.nan_to_num(d["A0"]).astype(np.float32))) | |
| Dtr, Xtr, Atr, Mtr, A0tr = T(tr) | |
| Dva, Xva, Ava, Mva, A0va = T(va) | |
| net = build_net(streams, cin).to(device) | |
| opt = torch.optim.AdamW(net.parameters(), lr=RECIPE["lr"], weight_decay=RECIPE["weight_decay"]) | |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs) | |
| z0 = np.nan_to_num(tr["A0"])[:, None, :] | |
| flips = (np.nan_to_num(np.abs(np.diff(np.concatenate([z0, np.nan_to_num(tr["A"])], 1), axis=1)), nan=0) > 0.5).any((1, 2)) | |
| w = np.where(flips, (len(flips) - flips.sum()) / max(flips.sum(), 1), 1.0) | |
| if not np.isfinite(w.sum()) or w.sum() <= 0: # degenerate case (every row flips -> all weights 0; the legacy formula has no guard): uniform sampling | |
| w = np.ones_like(w) | |
| p_s = w / w.sum() | |
| def loss_fn(hon, hoff, A, M, A0): | |
| zprev = torch.cat([A0[:, None, :], A[:, :-1, :]], 1) | |
| hon = hon.clamp(1e-5, 1 - 1e-5) | |
| hoff = hoff.clamp(1e-5, 1 - 1e-5) | |
| l = -((1 - zprev) * (A * hon.log() + (1 - A) * (1 - hon).log()) + zprev * ((1 - A) * hoff.log() + A * (1 - hoff).log())) | |
| return (l * M).sum() / M.sum().clamp_min(1) | |
| best, best_state, pat, hist = 1e9, None, 0, [] | |
| for ep in range(epochs): | |
| net.train() | |
| idx = np.random.choice(len(Dtr), size=min(len(Dtr), rows_per_epoch), replace=False, p=p_s) | |
| for i in range(0, len(idx), batch): | |
| b = idx[i:i + batch] | |
| hon, hoff = net(Dtr[b].to(device), Xtr[b].to(device)) | |
| loss = loss_fn(hon, hoff, Atr[b].to(device), Mtr[b].to(device), A0tr[b].to(device)) | |
| opt.zero_grad() | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(net.parameters(), RECIPE["grad_clip"]) | |
| opt.step() | |
| sched.step() | |
| net.eval() | |
| with torch.no_grad(): | |
| tot, cnt = 0.0, 0 | |
| for i in range(0, len(Dva), 8192): | |
| s = slice(i, i + 8192) | |
| hon, hoff = net(Dva[s].to(device), Xva[s].to(device)) | |
| n_ = min(8192, len(Dva) - i) | |
| tot += float(loss_fn(hon, hoff, Ava[s].to(device), Mva[s].to(device), A0va[s].to(device))) * n_ | |
| cnt += n_ | |
| v = tot / max(cnt, 1) | |
| hist.append({"epoch": ep, "val_nll": v}) | |
| log(f" ep {ep}: val {v:.4f}") | |
| if v < best - 1e-4: | |
| best, best_state, pat = v, {k: t.detach().cpu().clone() for k, t in net.state_dict().items()}, 0 | |
| else: | |
| pat += 1 | |
| if pat >= patience: | |
| break | |
| net.load_state_dict(best_state) | |
| net.eval() | |
| def predict(te): | |
| Dte, Xte = np.nan_to_num(te["D"].astype(np.float32)), ((te["X"] - xmu) / xsd).astype(np.float32) | |
| hon = np.zeros((len(Dte), H, 3), np.float32) | |
| hoff = np.zeros_like(hon) | |
| with torch.no_grad(): | |
| for i in range(0, len(Dte), 8192): | |
| s = slice(i, min(i + 8192, len(Dte))) | |
| a, b = net(torch.from_numpy(Dte[s]).to(device), torch.from_numpy(Xte[s]).to(device)) | |
| hon[s], hoff[s] = a.cpu().numpy(), b.cpu().numpy() | |
| return hon, hoff | |
| predict.best_val, predict.state, predict.xmu, predict.xsd, predict.history, predict.streams, predict.cin = best, best_state, xmu, xsd, hist, streams, cin | |
| return predict | |
| def save_ckpt(pred, path): | |
| import torch | |
| path = Path(path) | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| if not path.exists(): # append-only, as train_att :208-211 | |
| torch.save({"state": pred.state, "xmu": pred.xmu, "xsd": pred.xsd, "best_val": pred.best_val, "streams": pred.streams, "cin": pred.cin, "history": pred.history, "recipe": RECIPE}, path) | |
| return path | |
| def load_predictor(path, device="cpu"): | |
| import torch | |
| ck = torch.load(path, map_location="cpu", weights_only=False) | |
| net = build_net(ck["streams"], int(ck["cin"])) | |
| net.load_state_dict(ck["state"]) | |
| net.to(device).eval() | |
| xmu, xsd = ck["xmu"], ck["xsd"] | |
| def predict(te): | |
| Dte, Xte = np.nan_to_num(te["D"].astype(np.float32)), ((te["X"] - xmu) / xsd).astype(np.float32) | |
| hon = np.zeros((len(Dte), H, 3), np.float32) | |
| hoff = np.zeros_like(hon) | |
| with torch.no_grad(): | |
| for i in range(0, len(Dte), 8192): | |
| s = slice(i, min(i + 8192, len(Dte))) | |
| a, b = net(torch.from_numpy(Dte[s]).to(device), torch.from_numpy(Xte[s]).to(device)) | |
| hon[s], hoff[s] = a.cpu().numpy(), b.cpu().numpy() | |
| return hon, hoff | |
| predict.best_val, predict.streams, predict.cin = float(ck["best_val"]), ck["streams"], int(ck["cin"]) | |
| return predict | |
| def persistence_hazards(A, A0, M): | |
| """train_att.persistence_hazards on the assembled labels (imported when importable, local copy otherwise) -> (h_on[25,3], h_off[25,3]).""" | |
| tr = {"A": A, "A0": np.nan_to_num(A0), "M": M.all(-1) if M.ndim == 3 else M} | |
| try: | |
| import train_att as TA | |
| return TA.persistence_hazards(tr) | |
| except Exception: | |
| h_on = np.zeros((H, 3), np.float32) | |
| h_off = np.zeros_like(h_on) | |
| for j in range(3): | |
| z = tr["A"][:, :, j] | |
| z0 = tr["A0"][:, j] | |
| zprev = np.concatenate([z0[:, None], z[:, :-1]], 1) | |
| v = np.isfinite(z) & tr["M"] | |
| for k in range(H): | |
| m = v[:, k] | |
| h_on[k, j] = ((zprev[:, k] < 0.5) & (z[:, k] > 0.5) & m).sum() / max(((zprev[:, k] < 0.5) & m).sum(), 1) | |
| h_off[k, j] = ((zprev[:, k] > 0.5) & (z[:, k] < 0.5) & m).sum() / max(((zprev[:, k] > 0.5) & m).sum(), 1) | |
| return h_on, h_off | |
| # ----------------------------------------------------------------------------- composition (docs §5.3 steps 3-4) | |
| def compose_att(hon_face, hoff_face, hon_hands, hoff_hands, A0): | |
| """warn_compose.att_marginals routing (j=0,1 <- face chain, j=2 <- hands chain) + compose.marginal -> pA[N,25,3]; att_only = legacy_compose(pA, None, 15).""" | |
| N = len(A0) | |
| pA = np.zeros((N, H, 3), np.float32) | |
| for j, (hon, hoff) in enumerate(((hon_face, hoff_face), (hon_face, hoff_face), (hon_hands, hoff_hands))): | |
| pA[:, :, j] = marginal(hon[:, :, j], hoff[:, :, j], np.nan_to_num(A0[:, j])) | |
| out = legacy_compose(pA.astype(np.float32), None, TAU, None) | |
| return pA, out["att_only"].astype(np.float32), out["p_inatt"].astype(np.float32) | |
| def state_att(A0, tau=TAU): | |
| p = 1 - np.prod(1 - np.nan_to_num(A0), -1) | |
| return (1 - (1 - p) ** tau).astype(np.float32) | |
| def _chain_inputs(d, rows, which): | |
| if which == "face": | |
| return {"D": d["D"][rows].astype(np.float32), "X": d["X"][rows], "A": d["A"][rows], "A0": d["A0"][rows], | |
| "M": np.stack([d["M"][rows] & np.isfinite(d["A"][rows][:, :, 0]), d["M"][rows] & np.isfinite(d["A"][rows][:, :, 1]), np.zeros_like(d["M"][rows])], -1)} | |
| return {"D": d["Dh"][rows], "X": d["X"][rows], "A": d["A"][rows], "A0": d["A0"][rows], | |
| "M": np.stack([np.zeros_like(d["Mh"][rows]), np.zeros_like(d["Mh"][rows]), d["Mh"][rows]], -1)} | |
| def run_baseline(baseline, seed=0, device="cuda", cache_dir=None, out_root=LY.RUNS, chain_root=CHAIN_ROOT, splits=("test", "val"), epochs=RECIPE["epochs"], | |
| rows_per_epoch=RECIPE["rows_per_epoch"], batch=RECIPE["batch"], spec=None, regimes=REGIMES, log=print): | |
| """train (or load) the face chain per regime + the shared hands chain, predict test / val, write results/runs/abl_dms_<b>_chain_s<seed>/pred_layout_*.npz.""" | |
| if not applicable(baseline) and spec is None: | |
| log(f"{baseline}: chain-style arm not applicable (no head-pose / eye-state channel; CHANNEL_ROLES) — nothing written") | |
| return [] | |
| spec = spec or SPEC[baseline] | |
| arm = f"abl_dms_{baseline}_chain" | |
| written = [] | |
| ckpts = {} | |
| for regime in regimes: | |
| tr = assemble(baseline, regime, "train", cache_dir, spec, log=log) | |
| va = assemble(baseline, regime, "val", cache_dir, spec, log=log) | |
| rows_tr, rows_va = np.flatnonzero(tr["keep"]), np.flatnonzero(va["keep"]) | |
| assert len(rows_tr) >= 2 * batch and len(rows_va) >= 64, f"{baseline} {regime}: too few rows (train {len(rows_tr)}, val {len(rows_va)})" | |
| # face chain (streams ds): distracted / gaze_off | |
| ck_face = Path(chain_root) / f"ckpt_{baseline}_ds_{regime}_s{seed}.pt" | |
| if ck_face.exists(): | |
| pred_face = load_predictor(ck_face, device) | |
| log(f" face chain: loaded {ck_face} (val NLL {pred_face.best_val:.4f})") | |
| else: | |
| pred_face = train_chain(_chain_inputs(tr, rows_tr, "face"), _chain_inputs(va, rows_va, "face"), "ds", seed, epochs, device, rows_per_epoch, batch, log=log) | |
| save_ckpt(pred_face, ck_face) | |
| # shared hands chain (streams ds,dyn on the 2 pose channels): identical for every arm; rows = every row with a causal box (not baseline-dependent) | |
| ck_hands = Path(chain_root) / f"ckpt_shared_hands_ds-dyn_{regime}_s{seed}.pt" | |
| if ck_hands.exists(): | |
| pred_hands = load_predictor(ck_hands, device) | |
| log(f" hands chain: loaded {ck_hands} (val NLL {pred_hands.best_val:.4f})") | |
| else: | |
| hr_tr = np.flatnonzero(hands_arrays(regime, "train")["row_ok"] if cache_dir is None else tr["keep"]) | |
| hr_va = np.flatnonzero(hands_arrays(regime, "val")["row_ok"] if cache_dir is None else va["keep"]) | |
| pred_hands = train_chain(_chain_inputs(tr, hr_tr, "hands"), _chain_inputs(va, hr_va, "hands"), "ds,dyn", seed, epochs, device, rows_per_epoch, batch, log=log) | |
| save_ckpt(pred_hands, ck_hands) | |
| ckpts[regime] = {"face": {"path": str(ck_face), "sha256": cs.sha256_file(ck_face), "val_nll": pred_face.best_val, "cin": pred_face.cin, "streams": "ds"}, | |
| "hands": {"path": str(ck_hands), "sha256": cs.sha256_file(ck_hands), "val_nll": pred_hands.best_val, "cin": pred_hands.cin, "streams": "ds,dyn"}} | |
| Mtr3 = np.stack([tr["M"] & np.isfinite(tr["A"][:, :, 0]), tr["M"] & np.isfinite(tr["A"][:, :, 1]), tr["Mh"]], -1)[rows_tr] | |
| ph_on, ph_off = persistence_hazards(np.nan_to_num(tr["A"][rows_tr]), tr["A0"][rows_tr], Mtr3) | |
| for split in splits: | |
| te = va if split == "val" else assemble(baseline, regime, split, cache_dir, spec, log=log) | |
| rows = np.flatnonzero(te["keep"]) | |
| N = len(te["uid"]) | |
| cols = {f"dms_{baseline}_chain_att3s": np.full(N, np.nan, np.float32), f"dms_{baseline}_chain_persist_att3s": np.full(N, np.nan, np.float32), f"dms_{baseline}_chain_state_att3s": np.full(N, np.nan, np.float32)} | |
| if len(rows): | |
| fi, hi = _chain_inputs(te, rows, "face"), _chain_inputs(te, rows, "hands") | |
| hon_f, hoff_f = pred_face(fi) | |
| hon_h, hoff_h = pred_hands(hi) | |
| A0 = te["A0"][rows] | |
| _, att, _ = compose_att(hon_f, hoff_f, hon_h, hoff_h, A0) | |
| n = len(rows) | |
| _, att_p, _ = compose_att(np.broadcast_to(ph_on, (n, H, 3)).copy(), np.broadcast_to(ph_off, (n, H, 3)).copy(), np.broadcast_to(ph_on, (n, H, 3)).copy(), np.broadcast_to(ph_off, (n, H, 3)).copy(), A0) | |
| cols[f"dms_{baseline}_chain_att3s"][rows] = np.clip(att, 0, 1) | |
| cols[f"dms_{baseline}_chain_persist_att3s"][rows] = np.clip(att_p, 0, 1) | |
| cols[f"dms_{baseline}_chain_state_att3s"][rows] = state_att(A0) | |
| meta = {"baseline": {"id": baseline, "kind": "chain-style (matched protocol): legacy two-state hazard chain on the baseline's own channel history", "spec": spec, "channel_roles": CHANNEL_ROLES.get(baseline, {}), | |
| "channels_in": te["channels"], "thresholds": TH, "channels_undefined": channels_undefined(baseline), "shared_inputs": SHARED_INPUTS, "recipe": RECIPE, "deviations": DEVIATIONS, | |
| "routing": "distracted, gaze_off <- ds-only chain; hands_off <- ds,dyn chain (warn_compose.att_marginals)", "composition": "cwm2/compose.legacy_compose(pA, None, 15)['att_only'], no shrink", | |
| "checkpoints": ckpts[regime], "n_train_rows": int(len(rows_tr)), "n_val_rows": int(len(rows_va)), "keep_rule": "cache frame at t & >= 50 % valid history steps & A0[0:2] defined"}, | |
| "roi": "fullframe_causal", "roi_streams": ["face", "pose"], "driver_input_shift_s": 0.0, "notes": "DMS-baseline chain-style arm; comparator only, never an input to any model"} | |
| p = LY.write_layout(arm, regime, seed, split, te["uid"], te["t"], cols, meta_extra=meta, out_root=out_root, score_source="legacy_chain", | |
| model_id=f"{arm}:{ck_face.name}+{ck_hands.name}", mask=te["keep"], code_files=[Path(__file__), CWM2 / "compose.py"]) | |
| written.append(p) | |
| log(f" wrote {p} keep {te['keep'].mean():.4f} att3s mean {np.nanmean(cols[f'dms_{baseline}_chain_att3s']):.4f}") | |
| LY.mark_done(arm, seed, out_root, regimes, splits) | |
| return written | |
| # ----------------------------------------------------------------------------- dry run (CPU, synthetic caches, 1 epoch) | |
| def dry_run(keep_dir=None, log=print): | |
| import tempfile | |
| root = Path(keep_dir) if keep_dir else Path(tempfile.mkdtemp(prefix="dms_chain_dryrun_")) | |
| b, regime = "synthetic", "manual" | |
| spec = {"yaw": "yaw_deg", "pitch": "pitch_deg", "closed": ["closed_p"], "closed_kind": "closed_p", "note": "synthetic"} | |
| cdir = root / "cache" / b | |
| with np.load(S12R / f"targets_{regime}_val.npz", allow_pickle=True) as z: | |
| uid, t = z["uid"], z["t"] | |
| routes = [str(u) for u in np.unique(uid)[:3]] | |
| for i, u in enumerate(routes): | |
| ts = t[uid == u].astype(np.float64) | |
| cs.synthetic_route_cache(uid=u, t0=max(0.0, ts.min() - 8.2), t1=ts.max() + 5.2, channels=("yaw_deg", "pitch_deg", "closed_p", "conf"), seed=i, meta={"baseline": b}).save(cdir / f"{u}.npz") | |
| # train == val == the manual/val rows of the 3 synthetic routes (dry run only: the trainer sees a few hundred rows, 1 epoch, CPU) | |
| d = assemble(b, regime, "val", cdir, spec, log=log) | |
| rows = np.flatnonzero(d["keep"]) | |
| assert len(rows) >= 64, len(rows) | |
| t0 = time.time() | |
| pf = train_chain(_chain_inputs(d, rows, "face"), _chain_inputs(d, rows, "face"), "ds", 0, epochs=1, device="cpu", rows_per_epoch=256, batch=64, log=log) | |
| ph = train_chain(_chain_inputs(d, rows, "hands"), _chain_inputs(d, rows, "hands"), "ds,dyn", 0, epochs=1, device="cpu", rows_per_epoch=256, batch=64, log=log) | |
| hon_f, hoff_f = pf(_chain_inputs(d, rows, "face")) | |
| hon_h, hoff_h = ph(_chain_inputs(d, rows, "hands")) | |
| pA, att, p_in = compose_att(hon_f, hoff_f, hon_h, hoff_h, d["A0"][rows]) | |
| ref = legacy_compose(pA, None, TAU, None)["att_only"] | |
| assert np.allclose(att, ref, atol=1e-6) and np.all(att >= 0) and np.all(att <= 1) | |
| cols = {f"dms_{b}_chain_att3s": np.full(len(uid), np.nan, np.float32)} | |
| cols[f"dms_{b}_chain_att3s"][rows] = att | |
| p = LY.write_layout(f"abl_dms_{b}_chain", regime, 0, "val", uid, t, cols, meta_extra={"baseline": {"id": b, "dry_run": True, "thresholds": TH, "deviations": DEVIATIONS}}, out_root=root / "runs", | |
| score_source="legacy_chain", mask=d["keep"], code_files=[Path(__file__)]) | |
| m = LY.validate_layout(p) | |
| rep = {"root": str(root), "routes": routes, "n_keep": int(len(rows)), "face_val_nll": pf.best_val, "hands_val_nll": ph.best_val, "att3s_mean": float(att.mean()), "layout": str(p), "arm": m["arm"], | |
| "wall_s": round(time.time() - t0, 1), "deviations": DEVIATIONS} | |
| print(json.dumps(rep, indent=1)) | |
| return rep | |
| def main(argv=None): | |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) | |
| ap.add_argument("--baseline", default="") | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--device", default="cuda") | |
| ap.add_argument("--cache-dir", default="") | |
| ap.add_argument("--dry-run", action="store_true") | |
| ap.add_argument("--keep-dir", default="") | |
| ap.add_argument("--print-mapping", action="store_true") | |
| ap.add_argument("--note", default="", help="free-text launch note (inert; the layer-1 launcher passes gate_token=train_cwm3 on GPU launches so that every lane's trainer gate `pgrep -f 'python.*train_cwm3'` counts this process)") | |
| a = ap.parse_args(argv) | |
| if a.print_mapping: | |
| print(json.dumps({"thresholds": TH, "spec": SPEC, "channel_roles": CHANNEL_ROLES, "deviations": DEVIATIONS, "shared_inputs": SHARED_INPUTS}, indent=1)) | |
| if a.dry_run: | |
| dry_run(a.keep_dir or None) | |
| return 0 | |
| if a.baseline: | |
| run_baseline(a.baseline, a.seed, a.device, a.cache_dir or None) | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |
Xet Storage Details
- Size:
- 40 kB
- Xet hash:
- ca44a941efa04cdfc02fb0da0be63ce2f2ea234b50c58256145401cabe34af67
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.