Buckets:
| """cwm3/tests/test_can_gate.py — the OPTIONAL BATON-side CAN arms on the driver chains (plan Part A §A.5; docs/decisions.md 2026-09-15T01:48:39Z). | |
| (a) identity at init (pattern tests/test_zero_init_identity.py): DriverBranch(can_gate=True) built with the survivor kwargs at seed k holds the parent's parameters | |
| bitwise (CanReleaseGate is created LAST and consumes no RNG; its a = b = 0) and, given the same weights, its chain logits, its gated hazards (heads.hazards_gated) | |
| and its loss (heads.driver_loss_gated through batch.loss_U) are bitwise the parent's (heads.hazards / driver_loss): 2g == 1 exactly. | |
| (b) causality: the gate reads Kx[:, 30:40, 12:15] only — perturbing the FUTURE CAN (the targets' t+0.2..t+5.0 s rows, raw "Y"), any history step older than 1 s, or | |
| the kinematics / valid / vEgo / aEgo / Δt columns leaves the activity and 2g bitwise unchanged; perturbing steer at the anchor step (= t) moves 2g once a != 0. | |
| (c) formula: CanReleaseGate.activity == the numpy formula act = (1/9) Σ_{k=31..39} Σ_{c in {12,13,14}} |Kx[k,c] − Kx[k−1,c]| of the decisions row; hazards_gated == | |
| clamp(hazards · 2g) for arbitrary 2g. | |
| (d) not a dead init: one gradient step on the chain loss moves (a, b) and 2g away from 1. | |
| (e) nocan: batch.prep_U under recipe `u_nocan` -> Kx[..., 12:17] exactly 0; the 13 other columns, mode and every target equal the parent's; the parent GRU consumes it. | |
| (f) registry: the two stages.yaml variants exist and differ from the survivor by their one recipe flag only; train_cwm3.model_kwargs forwards can_gate for the | |
| cangate variant only and is unchanged for the survivor / nocan; need_streams identical. | |
| Synthetic data only (never a result). Run: cd cwm3 && python -m pytest tests/test_can_gate.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| 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 / "common", CWM3 / "eval", CWM3.parent / "s12_wm"): | |
| sys.path.insert(0, str(p)) | |
| import cwm3_dataset as ds # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| from models import build_branch # noqa: E402 | |
| from models import heads as hd # noqa: E402 | |
| from models.branches import CanReleaseGate, DriverBranch # noqa: E402 | |
| SURVIVOR = {"hid": 128, "hidden": 128, "dropout": 0.0, "cabin": "none", "in_ln": True} # train_cwm3.model_kwargs of kin_can_noK_clip8_ln | |
| CFG_U = {"cabin": "none", "kclip": 8.0, "recipe": {"kclip": 8, "u_in_ln": True}} | |
| TS = {"mu": np.array([0.0, 12.0, 0.0], np.float32), "sd": np.array([70.0, 8.0, 0.7], np.float32)} | |
| B = 32 | |
| def _raw(seed=3, n=96): | |
| raw = ds.synthetic_raw(n, seed=seed, mode=0, need=ds.ALL_STREAMS) | |
| raw["mode"][::2] = 1 | |
| return raw | |
| def _batch(raw, cfg=CFG_U, n=B, stats=None): | |
| st = stats or bt.compute_stats(raw, TS) | |
| store = bt.GpuStore(raw, bt.KEYS["U"], "cpu") | |
| b, tg = bt.prep_U(store, torch.arange(n), st, cfg) | |
| return b, tg, st | |
| def _pair(seed=0): | |
| parent = build_branch("U", seed, **SURVIVOR) | |
| gated = build_branch("U", seed, **SURVIVOR, can_gate=True) | |
| return parent, gated | |
| # ----------------------------------------------------------------------------- (a) identity at init | |
| def test_gate_params_zero_and_parent_weights_bitwise(): | |
| for seed in (0, 1, 2): | |
| parent, gated = _pair(seed) | |
| sp, sg = parent.state_dict(), gated.state_dict() | |
| assert set(sg) - set(sp) == {"can_gate.a", "can_gate.b"} | |
| assert all(torch.equal(sp[k], sg[k]) for k in sp), "the gated branch does not hold the parent's paired init" | |
| assert float(gated.can_gate.a) == 0.0 and float(gated.can_gate.b) == 0.0 | |
| assert sum(p.numel() for p in gated.parameters()) == sum(p.numel() for p in parent.parameters()) + 2 | |
| def test_identity_at_init_outputs_hazards_loss(): | |
| raw = _raw() | |
| b, tg, _ = _batch(raw) | |
| for seed in (0, 1): | |
| parent, gated = _pair(seed) | |
| parent.eval() | |
| gated.eval() | |
| with torch.no_grad(): | |
| op, og = parent(b), gated(b) | |
| assert set(op) == {"U"} and set(og) == {"U", "HG"} | |
| assert torch.equal(og["U"], op["U"]), "chain logits differ at init" | |
| assert og["HG"].shape == (B, 1) and torch.equal(og["HG"], torch.ones_like(og["HG"])), "2g != 1 at init" | |
| hp = hd.hazards(hd.split_U(op["U"])["chains"]) | |
| hg = hd.hazards_gated(hd.split_U(og["U"])["chains"], og["HG"]) | |
| assert torch.equal(hp[0], hg[0]) and torch.equal(hp[1], hg[1]), "gated hazards differ from the parent's at init" | |
| lp, pp = bt.loss_U(op, tg, {"K": 0.0}) | |
| lg, pg = bt.loss_U(og, tg, {"K": 0.0}) | |
| assert torch.equal(lp, lg) and set(pp) == set(pg) and all(torch.equal(pp[k], pg[k]) for k in pp), "loss differs at init" | |
| lp2, _ = hd.driver_loss(op["U"], tg, lam_K=0.5) | |
| lg2, _ = hd.driver_loss_gated(og["U"], tg, og["HG"], lam_K=0.5) | |
| assert torch.equal(lp2, lg2) | |
| # ----------------------------------------------------------------------------- (b) causality | |
| def test_gate_is_causal_and_reads_only_the_last_second_of_control_can(): | |
| raw = _raw() | |
| b, _, st = _batch(raw) | |
| kx = b["Kx"] | |
| gate = CanReleaseGate() | |
| with torch.no_grad(): | |
| gate.a.fill_(1.0) | |
| gate.b.fill_(0.3) | |
| h0 = gate(kx) | |
| assert not torch.equal(h0, torch.ones_like(h0)) # a != 0: the gate is live | |
| for name, sl in (("history older than 1 s", (slice(None), slice(0, 30), slice(12, 15))), | |
| ("kinematics + valid bit", (slice(None), slice(None), slice(0, 12))), | |
| ("vEgo / aEgo / dt", (slice(None), slice(None), slice(15, 18)))): | |
| k2 = kx.clone() | |
| k2[sl] += 3.0 | |
| assert torch.equal(CanReleaseGate.activity(k2), CanReleaseGate.activity(kx)), name | |
| assert torch.equal(gate(k2), h0), name | |
| k3 = kx.clone() | |
| k3[:, 39, 12] += 1.0 # steer at the anchor step (= t): history, may move the gate | |
| assert not torch.equal(gate(k3), h0) | |
| # the FUTURE CAN (targets Y = t+0.2..t+5.0 s: the only CAN after t in the store) never reaches Kx, hence never the gate | |
| raw2 = {k: (v.copy() if isinstance(v, np.ndarray) else v) for k, v in raw.items()} | |
| rng = np.random.default_rng(0) | |
| raw2["Y"] = (rng.normal(0, 5, raw2["Y"].shape)).astype(np.float32) | |
| raw2["Y"][:, :, 1:3] = (rng.random(raw2["Y"][:, :, 1:3].shape) < 0.5).astype(np.float32) | |
| b2, _, _ = _batch(raw2, stats=st) | |
| assert torch.equal(b2["Kx"], kx), "future CAN reached Kx" | |
| with torch.no_grad(): | |
| assert torch.equal(gate(b2["Kx"]), h0) | |
| # ----------------------------------------------------------------------------- (c) formulas | |
| def test_activity_formula_matches_decisions_row(): | |
| raw = _raw(seed=11) | |
| b, _, _ = _batch(raw) | |
| kx = b["Kx"].numpy().astype(np.float32) | |
| c = kx[:, 30:40, 12:15] | |
| ref = np.abs(np.diff(c, axis=1)).sum(-1).mean(-1) # (1/9) Σ_{k=31..39} Σ_c |Kx[k,c] − Kx[k−1,c]| | |
| act = CanReleaseGate.activity(b["Kx"]) | |
| assert act.shape == (B, 1) | |
| np.testing.assert_allclose(act[:, 0].numpy(), ref, rtol=1e-6, atol=1e-6) | |
| assert (act >= 0).all() and float(act.max()) > 0 | |
| def test_hazards_gated_formula(): | |
| torch.manual_seed(0) | |
| chains = torch.randn(6, 25, 2, 2) * 4 | |
| g2 = torch.tensor([[0.2], [0.5], [1.0], [1.5], [1.9], [2.0]]) | |
| h_on, h_off = hd.hazards_gated(chains, g2) | |
| lam_on, lam_off = hd.hazards(chains) | |
| assert torch.equal(h_on, (lam_on * g2[:, :, None]).clamp(hd.EPS_H, 1 - hd.EPS_H)) | |
| assert torch.equal(h_off, (lam_off * g2[:, :, None]).clamp(hd.EPS_H, 1 - hd.EPS_H)) | |
| assert torch.equal(h_on[2], lam_on[2]) and torch.equal(h_off[2], lam_off[2]) # 2g == 1 -> identity row | |
| assert (h_on > 0).all() and (h_on < 1).all() and (h_off > 0).all() and (h_off < 1).all() | |
| # ----------------------------------------------------------------------------- (d) trainable | |
| def test_zero_init_is_trainable(): | |
| raw = _raw(seed=5) | |
| b, tg, _ = _batch(raw) | |
| _, gated = _pair(0) | |
| gated.train() | |
| loss, _ = bt.loss_U(gated(b), tg, {"K": 0.0}) | |
| loss.backward() | |
| assert gated.can_gate.a.grad is not None and float(gated.can_gate.a.grad.abs()) > 0 | |
| assert gated.can_gate.b.grad is not None and float(gated.can_gate.b.grad.abs()) > 0 | |
| with torch.no_grad(): | |
| for p in (gated.can_gate.a, gated.can_gate.b): | |
| p -= 1.0 * p.grad | |
| gated.eval() | |
| with torch.no_grad(): | |
| hg = gated(b)["HG"] | |
| assert not torch.equal(hg, torch.ones_like(hg)) and torch.isfinite(hg).all() and (hg > 0).all() and (hg < 2).all() | |
| # ----------------------------------------------------------------------------- (e) nocan | |
| def test_nocan_zeroes_exactly_the_five_can_columns(): | |
| raw = _raw(seed=7) | |
| b, tg, st = _batch(raw) | |
| cfg_n = dict(CFG_U, recipe=dict(CFG_U["recipe"], u_nocan=True)) | |
| bn, tgn, _ = _batch(raw, cfg=cfg_n, stats=st) | |
| assert bn["Kx"].shape == b["Kx"].shape == (B, 40, 18) | |
| assert torch.equal(bn["Kx"][..., 12:17], torch.zeros(B, 40, 5)), "nocan CAN columns are not exactly zero" | |
| assert not torch.equal(b["Kx"][..., 12:17], torch.zeros(B, 40, 5)), "the parent's CAN columns are trivially zero (bad fixture)" | |
| assert torch.equal(bn["Kx"][..., :12], b["Kx"][..., :12]) and torch.equal(bn["Kx"][..., 17:], b["Kx"][..., 17:]) | |
| assert torch.equal(bn["mode"], b["mode"]) and set(tgn) == set(tg) and all(torch.equal(tgn[k], tg[k]) for k in tg) | |
| parent = build_branch("U", 0, **SURVIVOR) | |
| with torch.no_grad(): | |
| assert parent(bn)["U"].shape == (B, 25, hd.D_U) # same width: the parent GRU consumes the nocan batch | |
| # ----------------------------------------------------------------------------- (f) registry | |
| def test_stage_registry_and_model_kwargs(): | |
| import train_cwm3 as T | |
| base = T.stage_cfg("U0", "kin_can_noK_clip8_ln") | |
| nocan = T.stage_cfg("U0", "kin_can_noK_clip8_ln_nocan") | |
| gate = T.stage_cfg("U0", "kin_can_noK_clip8_ln_cangate") | |
| def strip(c, key): | |
| c = {k: v for k, v in c.items() if k != "variant"} | |
| c["recipe"] = {k: v for k, v in c["recipe"].items() if k != key} | |
| return c | |
| assert "u_nocan" not in base["recipe"] and "u_can_gate" not in base["recipe"] | |
| assert nocan["recipe"]["u_nocan"] is True and "u_can_gate" not in nocan["recipe"] and strip(nocan, "u_nocan") == strip(base, None) | |
| assert gate["recipe"]["u_can_gate"] is True and "u_nocan" not in gate["recipe"] and strip(gate, "u_can_gate") == strip(base, None) | |
| kb, kn, kg = T.model_kwargs(base), T.model_kwargs(nocan), T.model_kwargs(gate) | |
| assert kb == kn == SURVIVOR | |
| assert kg == dict(SURVIVOR, can_gate=True) | |
| assert T.need_streams(gate) == T.need_streams(base) == T.need_streams(nocan) == frozenset({"X", "Kh", "kin", "Y"}) | |
| m = build_branch("U", 0, **kg) | |
| assert isinstance(m, DriverBranch) and m.can_gate is not None and build_branch("U", 0, **kb).can_gate is None | |
| raw = _raw(seed=9) | |
| b, _, _ = _batch(raw, cfg=dict(CFG_U, recipe=gate["recipe"])) | |
| m.eval() | |
| with torch.no_grad(): | |
| o = m(b) | |
| assert o["HG"].shape == (B, 1) and torch.equal(o["HG"], torch.ones(B, 1)) | |
Xet Storage Details
- Size:
- 11.7 kB
- Xet hash:
- 4be3383cd9881f49f26eb38316eb3fb21e6a64cfb56d03ad608ecd9d77ed2233
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.