Buckets:
| """cwm3/tests/test_mask_v.py — the vehicle branch as a MISSING modality (variant `ue_noV`; pre-registration docs/decisions.md 2026-09-18T03:34:28Z). | |
| (a) registry: I1 variant `ue_noV` = {interactions: [], mask_V: true}; every other I1 variant carries no `mask_V`; train_cwm3_joint.build_model forwards the kwarg only then. | |
| (b) model: with mask_V the U / E / F / P / K outputs are BITWISE invariant to arbitrary V inputs (X replaced by noise / huge values), out['trunks']['z_v'] == 0 exactly, | |
| and the V-derived outputs (V emissions, R head) depend on the mode embedding only (identical for two batches with the same modes and different X). | |
| (c) prep_joint: with pc['V']['mask_V'] the V input X is all zero, the V-valid mask vv is all False, the V targets' masks (M, zVm) and the R aux mask are all False, | |
| joint_loss is finite (masked means over empty masks -> 0), the V chain parts are exactly 0; v_valid_np returns all False. | |
| (d) default off: CWM3Joint(mask_V=False) and prep_joint without the key reproduce the unchanged model / batch bitwise (identity with the pre-existing path). | |
| (e) forbidden tokens: none of the 14 openpilot DMS identifiers in the new / edited code. | |
| Run: cd cwm3 && python -m pytest tests/test_mask_v.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import re | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| import yaml # noqa: E402 | |
| torch.set_num_threads(4) | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train", CWM3 / "tests", CWM3 / "common", CWM3 / "eval", CWM3.parent / "s12_wm"): | |
| sys.path.insert(0, str(p)) | |
| import batch as bt # noqa: E402 | |
| import cwm3_dataset as ds # noqa: E402 | |
| from models.cwm3 import CWM3Joint, joint_loss # noqa: E402 | |
| from test_no_dms_inputs import FORBIDDEN # noqa: E402 | |
| from test_zero_init_identity import _batch, _fresh_branches, KEYS, CFG # noqa: E402 | |
| FILES = ["models/cwm3.py", "train/train_cwm3_joint.py", "train/stages_joint.yaml", "tests/test_mask_v.py"] | |
| def _store(seed=3): | |
| raw = ds.synthetic_raw(64, seed=seed, mode=0, need=ds.ALL_STREAMS) | |
| raw["mode"][::2] = 1 | |
| bank = ds.synthetic_bank() | |
| st = bt.compute_stats(raw, {"mu": np.array([0.0, 12.0, 0.0], np.float32), "sd": np.array([70.0, 8.0, 0.7], np.float32)}) | |
| store = bt.GpuStore(raw, KEYS, "cpu", bank) | |
| return raw, store, st | |
| # ----------------------------------------------------------------------------- (a) registry | |
| def test_variant_registered_and_forwarded(): | |
| reg = yaml.safe_load((CWM3 / "train" / "stages_joint.yaml").read_text())["stages"]["I1"]["variants"] | |
| assert reg["ue_noV"] == {"interactions": [], "mask_V": True} | |
| assert all(not v.get("mask_V", False) for k, v in reg.items() if k != "ue_noV"), "mask_V must be exclusive to ue_noV" | |
| import train_cwm3_joint as TJ | |
| cfg = TJ.stage_cfg("I1", "ue_noV") | |
| assert cfg["mask_V"] is True and cfg["interactions"] == [] and cfg["recipe"] == TJ.stage_cfg("I1", "all_minus_road2drv")["recipe"] | |
| assert not TJ.stage_cfg("I1", "all_minus_road2drv").get("mask_V", False) | |
| pc = TJ.apply_mask_v({"V": {"can_hist": "4s"}, "U": {}, "E": {}, "R": {}}, cfg) | |
| assert pc["V"]["mask_V"] is True | |
| pc0 = TJ.apply_mask_v({"V": {"can_hist": "4s"}, "U": {}, "E": {}, "R": {}}, TJ.stage_cfg("I1", "none")) | |
| assert "mask_V" not in pc0["V"] | |
| raw = {"t": np.zeros(7), "x8_ok": np.ones(7, bool)} | |
| assert not TJ.v_valid_np(raw, pc).any() and TJ.v_valid_np(raw, pc0).all() | |
| # ----------------------------------------------------------------------------- (b) model invariance | |
| def test_masked_model_ignores_every_v_input(): | |
| V, U, E = _fresh_branches() | |
| b, tg = _batch() | |
| m = CWM3Joint(V, U, E, interactions=(), decoder="direct", k_head=True, seed=0, mask_V=True).eval() | |
| with torch.no_grad(): | |
| o1 = m(b, return_trunks=True) | |
| b2 = {k: dict(v) for k, v in b.items()} | |
| b2["V"] = dict(b["V"], X=torch.randn_like(b["V"]["X"]) * 1e3) | |
| o2 = m(b2, return_trunks=True) | |
| b3 = {k: dict(v) for k, v in b.items()} | |
| b3["V"] = dict(b["V"], X=torch.full_like(b["V"]["X"], float("nan"))) | |
| o3 = m(b3, return_trunks=True) | |
| assert torch.equal(o1["trunks"]["z_v"], torch.zeros_like(o1["trunks"]["z_v"])), "z_v must be exact zeros" | |
| for k in ("U", "E", "F", "P", "V", "R"): | |
| assert torch.equal(o1[k], o2[k]) and torch.equal(o1[k], o3[k]), f"{k} depends on the masked V input" | |
| for k in ("h_u", "z_e", "z_vf"): | |
| assert torch.equal(o1["trunks"][k], o2["trunks"][k]) | |
| assert torch.isfinite(o1["V"]).all() and torch.isfinite(o1["R"]).all() | |
| # the V-side outputs are a function of the mode embedding only: rows with the same mode give the same V emissions (batched-GEMM round-off ~1e-7 allowed) | |
| modes = b["V"]["mode"] | |
| for md in modes.unique(): | |
| rows = (modes == md).nonzero().flatten() | |
| assert torch.allclose(o1["V"][rows], o1["V"][rows[:1]].expand_as(o1["V"][rows]), atol=1e-5, rtol=0) | |
| assert torch.allclose(o1["R"][rows], o1["R"][rows[:1]].expand_as(o1["R"][rows]), atol=1e-5, rtol=0) | |
| # ----------------------------------------------------------------------------- (c) batch prep + loss | |
| def test_prep_joint_masks_v_targets_and_loss_is_finite(): | |
| import train_cwm3_joint as TJ | |
| raw, store, st = _store() | |
| st = dict(st) | |
| st.update(TJ.aux_stats(TJ.synthetic_aux(dict(raw), 0))) # p0_base / r0_base for the detached aux chains | |
| raw2 = TJ.synthetic_aux(raw, 0) | |
| store = bt.GpuStore(raw2, [k for k in TJ.KEYS if k not in TJ.HOST_KEYS], "cpu", ds.synthetic_bank()) | |
| for k in TJ.HOST_KEYS: | |
| if k in raw2: | |
| store.t[k] = TJ.HostTensor(raw2[k], "cpu") | |
| pc = {"V": {"can_hist": "4s", "mask_V": True}, "U": {"cabin": "none"}, "E": {"tokens": "pooled", "n_slots": 16}, "R": {"tokens": "region", "n_slots": 16}} | |
| idx = torch.arange(16) | |
| b, tg, vv = TJ.prep_joint(store, idx, st, pc, True, False) | |
| assert not vv.any() and torch.equal(b["V"]["X"], torch.zeros_like(b["V"]["X"])) | |
| assert not tg["V"]["M"].any() and not tg["V"]["zVm"].any() | |
| assert "R" in tg and not tg["R"]["m"].any(), "the R aux chain must be masked out on masked-V rows" | |
| assert "P" in tg and tg["P"]["m"].any(), "the P aux chain (E trunk) is unaffected" | |
| V, U, E = _fresh_branches() | |
| m = CWM3Joint(V, U, E, interactions=(), decoder="direct", k_head=True, seed=0, mask_V=True).eval() | |
| with torch.no_grad(): | |
| out = m(b) | |
| loss, parts = joint_loss(out, tg, {"hz": 0.5, "K": 0.5, "F": 0.1}) | |
| assert torch.isfinite(loss) | |
| for c in ("gas", "brake", "steer_phase", "resp"): | |
| assert float(parts[f"chain_{c}"]) == 0.0, f"V chain part {c} must vanish under the all-False mask" | |
| assert float(parts["rchain_hard"]) == 0.0 | |
| assert float(parts["chain_hands_on"]) > 0 and float(parts["chain_ttc"]) > 0 | |
| # the unmasked prep of the same rows is the pre-existing path (X standardised, vv all True) | |
| pc0 = {"V": {"can_hist": "4s"}, "U": {"cabin": "none"}, "E": {"tokens": "pooled", "n_slots": 16}, "R": {"tokens": "region", "n_slots": 16}} | |
| b0, tg0, vv0 = TJ.prep_joint(store, idx, st, pc0, True, False) | |
| assert vv0.all() and not torch.equal(b0["V"]["X"], torch.zeros_like(b0["V"]["X"])) and tg0["V"]["M"].any() | |
| # ----------------------------------------------------------------------------- (d) default off = unchanged model | |
| def test_default_off_is_identity(): | |
| V, U, E = _fresh_branches() | |
| b, _ = _batch() | |
| m0 = CWM3Joint(V, U, E, interactions=(), decoder="direct", k_head=True, seed=0).eval() | |
| V2, U2, E2 = _fresh_branches() | |
| m1 = CWM3Joint(V2, U2, E2, interactions=(), decoder="direct", k_head=True, seed=0, mask_V=False).eval() | |
| with torch.no_grad(): | |
| o0, o1 = m0(b, return_trunks=True), m1(b, return_trunks=True) | |
| for k in ("V", "U", "E", "F", "P", "R"): | |
| assert torch.equal(o0[k], o1[k]) | |
| assert not torch.equal(o0["trunks"]["z_v"], torch.zeros_like(o0["trunks"]["z_v"])), "the unmasked trunk is not zero" | |
| assert m0.mask_V is False and m1.mask_V is False | |
| # ----------------------------------------------------------------------------- (e) forbidden identifiers | |
| def test_no_forbidden_tokens(): | |
| pat = re.compile("|".join(re.escape(t) for t in FORBIDDEN)) | |
| for f in FILES: | |
| txt = (CWM3 / f).read_text() | |
| assert not pat.search(txt), f"forbidden DMS identifier in {f}" | |
Xet Storage Details
- Size:
- 8.76 kB
- Xet hash:
- 8f48c2f767ac95e4a48b0b1a4f345f316f9a651270a48dfb00a971858df0a3c9
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.