Buckets:
| """cwm3/tests/test_u_driverwm.py — the Driver-WM-like U trunk (models/u_driverwm.py; owner experiment C, pre-registration docs/decisions.md 2026-09-17T07:00:22Z). | |
| (a) registry: variant `kin_can_noK_clip8_ln_udwm` = the FINAL U recipe (kclip 8, u_in_ln, lambda K 0) + `arch: driverwm`; train_cwm3.model_kwargs and | |
| models.cwm3._model_kwargs forward `arch`; base_cfg_summary records it; the reference variant carries no `arch` (nothing changes for old runs). | |
| (b) input contract / output shape: the same prep_U batch (Kx [B,40,18] + mode) -> out {"U": [B,25,26]} exactly like the GRU DriverBranch. | |
| (c) causal (no future): per-step encoder state j is invariant to perturbations of steps > j; the state at t does react to the last step. | |
| (d) paired init: in_ln / mode / head start from the GRU baseline's seeded values (constructed first); only the encoder differs. | |
| (e) parameter counts: encoder 21,185,664 (tok 60,416 + pre_int 12,596,224 + f_int 8,395,776 + proj 133,248) vs GRU 56,832; NOT parameter-matched (stated). | |
| (f) forbidden tokens: no openpilot DMS field name in the new files. | |
| (g) joint contract: CWM3Joint with the DWM U (synthetic V / E) is identity at init (U chain logits, V, E, F bitwise the branches'), d_u = 128, trunk_U uses encode_trunk. | |
| (h) checkpoint round trip: a BEST checkpoint of the DWM base reloads through models.cwm3.load_base_branch with the right architecture (strict state dict). | |
| Run: cd cwm3 && python -m pytest tests/test_u_driverwm.py -q | |
| """ | |
| from __future__ import annotations | |
| import json | |
| 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 | |
| torch.set_num_threads(4) | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train", CWM3 / "tests"): | |
| sys.path.insert(0, str(p)) | |
| import batch as bt # noqa: E402 | |
| import cwm3_dataset as ds # noqa: E402 | |
| from models import build_branch # noqa: E402 | |
| from models.branches import DriverBranch # noqa: E402 | |
| from models.cwm3 import CWM3Joint, _model_kwargs, base_cfg_summary, load_base_branch # noqa: E402 | |
| from models.u_driverwm import D_FF, D_INT, HEADS, DriverBranchDWM, n_params # noqa: E402 | |
| from test_no_dms_inputs import FORBIDDEN # noqa: E402 | |
| from test_zero_init_identity import _assert_identity, _batch # noqa: E402 | |
| import train_cwm3 as T # noqa: E402 | |
| VARIANT, REF = "kin_can_noK_clip8_ln_udwm", "kin_can_noK_clip8_ln" | |
| U_CFG = {"cabin": "none", "kclip": 8.0} | |
| NEW_FILES = ["models/u_driverwm.py", "train/queue_udwm_u0.txt", "train/queue_udwm_gate.txt", "train/queue_udwm.txt", "train/run_queue_udwm.sh", | |
| "eval/run_joint_evals_final_udwm.sh", "eval/report_final_udwm.py", "train/run_udwm_chain.sh", "tests/test_u_driverwm.py"] | |
| def _cfg(variant=VARIANT): | |
| return T.stage_cfg("U0", variant) | |
| def _build(seed=0, variant=VARIANT): | |
| return build_branch("U", seed, **T.model_kwargs(_cfg(variant))) | |
| def _ubatch(n=6, seed=3): | |
| b, _ = _batch(n=n, seed=seed, u_cfg=U_CFG) | |
| return b | |
| # ----------------------------------------------------------------------------- (a) registry | |
| def test_variant_registered_and_forwarded(): | |
| c, r = _cfg(), _cfg(REF) | |
| assert c["arch"] == "driverwm" and "arch" not in r | |
| for k in ("batch", "rows_per_epoch", "max_epochs", "patience", "lr", "weight_decay", "sched", "val_rows", "kclip", "u_in_ln"): | |
| assert c["recipe"][k] == r["recipe"][k], k # the FINAL U recipe verbatim | |
| assert c["lambda"] == r["lambda"] == {"K": 0.0} and c["model"] == r["model"] and c["cabin"] == "none" | |
| kw, kw_ref = T.model_kwargs(c), T.model_kwargs(r) | |
| assert kw.pop("arch") == "driverwm" and kw == kw_ref # only the arch key differs | |
| kw2 = _model_kwargs(c) | |
| assert kw2["arch"] == "driverwm" and kw2["in_ln"] is True and "arch" not in _model_kwargs(r) | |
| assert base_cfg_summary(c)["arch"] == "driverwm" and "arch" not in base_cfg_summary(r) | |
| # ----------------------------------------------------------------------------- (b) contract | |
| def test_input_contract_output_shape(): | |
| m = _build() | |
| assert isinstance(m, DriverBranchDWM) and m.in_ln is not None and not hasattr(m, "gru") | |
| b = _ubatch() | |
| assert tuple(b["U"]["Kx"].shape) == (6, 40, 18) # the GRU U's input exactly (prep_U, kclip 8) | |
| out = m(b["U"]) | |
| assert set(out) == {"U"} and tuple(out["U"].shape) == (6, 25, 26) and torch.isfinite(out["U"]).all() | |
| ref = _build(variant=REF) | |
| assert tuple(ref(b["U"])["U"].shape) == tuple(out["U"].shape) | |
| # ----------------------------------------------------------------------------- (c) causal | |
| def test_causal_no_future(): | |
| m = _build().eval() | |
| x = torch.randn(4, 40, 18) | |
| with torch.no_grad(): | |
| z = m.enc.steps(x) | |
| for j in (0, 10, 25, 38): | |
| x2 = x.clone() | |
| x2[:, j + 1:] += torch.randn_like(x2[:, j + 1:]) * 3.0 | |
| z2 = m.enc.steps(x2) | |
| assert torch.allclose(z[:, :j + 1], z2[:, :j + 1], atol=1e-5), f"state <= {j} changed by future steps" | |
| assert not torch.allclose(z[:, j + 1:], z2[:, j + 1:], atol=1e-3), f"states > {j} did not react" | |
| h, h_last = m.enc(x), None | |
| x3 = x.clone(); x3[:, -1] += 3.0 | |
| assert not torch.allclose(h, m.enc(x3), atol=1e-4), "h_u ignores the last (current) step" | |
| assert bool(m.enc.mask[0, 1]) and not bool(m.enc.mask[1, 0]) and not bool(m.enc.mask.diagonal().any()) | |
| # ----------------------------------------------------------------------------- (d) paired init | |
| def test_paired_init_shared_modules(): | |
| m, ref = _build(seed=0), _build(seed=0, variant=REF) | |
| assert isinstance(ref, DriverBranch) and not isinstance(ref, DriverBranchDWM) | |
| for name in ("in_ln", "mode", "head"): | |
| a, b = getattr(m, name).state_dict(), getattr(ref, name).state_dict() | |
| assert set(a) == set(b) and all(torch.equal(a[k], b[k]) for k in a), f"{name} not paired with the GRU baseline" | |
| m1 = _build(seed=1) | |
| assert not torch.equal(m1.head[0].weight, m.head[0].weight) and not torch.equal(m1.enc.tok.weight, m.enc.tok.weight) | |
| # ----------------------------------------------------------------------------- (e) parameter counts | |
| def test_parameter_counts(): | |
| m = _build() | |
| ai = m.arch_info | |
| assert (ai["d_int"], ai["heads"], ai["d_ff"], ai["layers"], ai["transition"]) == (D_INT, HEADS, D_FF, 1, True) == (1024, 4, 4096, 1, True) | |
| assert ai["encoder_params"] == n_params(m.enc) == 21_185_664 and ai["reference_encoder_params"] == 56_832 and ai["parameter_matched"] is False | |
| assert ai["encoder_parts"] == {"tok": 60_416, "pre_int": 12_596_224, "f_int": 8_395_776, "proj": 133_248} | |
| assert n_params(m) == 21_288_146 and n_params(_build(variant=REF)) == 159_314 | |
| # the published width of the AIDE port: pre_int == TransformerEncoderLayer(1024, 4, 4096) | |
| assert m.enc.pre_int.self_attn.embed_dim == 1024 and m.enc.pre_int.self_attn.num_heads == 4 and m.enc.pre_int.linear1.out_features == 4096 and m.enc.pre_int.norm_first | |
| # ----------------------------------------------------------------------------- (f) forbidden tokens | |
| def test_forbidden_tokens_absent(): | |
| pat = re.compile("|".join(map(re.escape, FORBIDDEN))) | |
| hits = [] | |
| for rel in NEW_FILES: | |
| p = CWM3 / rel | |
| if not p.exists(): | |
| continue | |
| for i, line in enumerate(p.read_text().splitlines(), 1): | |
| if pat.search(line) and "FORBIDDEN" not in line: | |
| hits.append(f"{rel}:{i}") | |
| assert not hits, hits | |
| src = (CWM3 / "models" / "u_driverwm.py").read_text() | |
| assert "weak" not in src.replace("weak_", "") and "dms" not in src.lower().replace("no openpilot dms", "") | |
| # ----------------------------------------------------------------------------- (g) joint contract | |
| def test_joint_identity_with_dwm_u(): | |
| V = build_branch("V", 0, d=256, hidden=256, dropout=0.1) | |
| E = build_branch("E", 0, tokens="pooled", n_slots=16, d=128, n_layers=2, n_heads=4, d_ff=256, hidden=256, dropout=0.1) | |
| U = _build() | |
| m = CWM3Joint(V, U, E, interactions=("env2veh", "drv2veh"), decoder="direct", k_head=True, seed=0) | |
| b = _ubatch(n=8) | |
| o = _assert_identity(m, b) | |
| assert o["trunks"]["h_u"].shape == (8, 128) and m.drv2veh is not None and m.road2drv is None | |
| with torch.no_grad(): | |
| kx = b["U"]["Kx"].float() | |
| kx = torch.cat([U.in_ln(kx[..., :12]), kx[..., 12:]], -1) | |
| assert torch.equal(o["trunks"]["h_u"], U.encode_trunk(kx)) # trunk_U used encode_trunk on the LN-applied Kx | |
| rep = m.param_report() | |
| assert rep["U"]["enc"] == 21_185_664 and rep["U"]["total"] == 21_288_146 and "gru" not in rep["U"] | |
| assert len(list(m.base_parameters())) == len(list(V.parameters())) + len(list(U.parameters())) + len(list(E.parameters())) | |
| # ----------------------------------------------------------------------------- (h) checkpoint round trip | |
| def test_checkpoint_roundtrip(tmp_path): | |
| m = _build(seed=2) | |
| cfg = _cfg() | |
| cfg.update({"kin_source": "causal_routes", "kin_version": "v2", "stats_version": 2, "kclip": 8.0, "pack_version": "v1", "mask_routes": ""}) | |
| torch.save({"cfg": cfg, "seed": 2, "model": m.state_dict(), "stats": {"xmu": np.zeros(5, np.float32)}}, tmp_path / "ckpt_ep001_b00000.pt") | |
| (tmp_path / "BEST").write_text("ckpt_ep001_b00000.pt") | |
| m2, c2, st, best = load_base_branch(tmp_path) | |
| assert isinstance(m2, DriverBranchDWM) and best == "ckpt_ep001_b00000.pt" and c2["arch"] == "driverwm" | |
| a, b = m.state_dict(), m2.state_dict() | |
| assert set(a) == set(b) and all(torch.equal(a[k], b[k]) for k in a) | |
| assert base_cfg_summary(c2)["arch"] == "driverwm" and base_cfg_summary(c2)["kclip"] == 8.0 | |
| x = _ubatch(n=3)["U"] | |
| with torch.no_grad(): | |
| assert torch.equal(m.eval()(x)["U"], m2.eval()(x)["U"]) | |
Xet Storage Details
- Size:
- 10.2 kB
- Xet hash:
- 8e88fbd9355a0874134914bbd53620fd7d67c81bffb271bd1ffe9568caaea8ba
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.