Buckets:
| """cwm3/tests/test_latent_aux.py — C4 low-dimensional multi-horizon latent auxiliary objective (models/latent_aux.py; train_cwm3_joint `latent_aux` variants), CPU, | |
| synthetic data (never a result). | |
| (a) lam = 0 is a TRUE no-op: stage_cfg('I1', 'all_minus_road2drv_lat0') builds the SAME model as the parent (no latent module, identical state_dict keys and | |
| parameter count), the forward has no "LAT" output, and joint_loss + add_latent_loss returns the parent's loss tensor bit for bit (same parts). | |
| (b) lam > 0: out["LAT"] is [B, 3, 16]; the total equals base + lam * L_aux exactly; parts lat_pred / lat_rec / lat_var are finite and none starts with chain_ | |
| (selection objective unchanged); the latent module has < 0.25 M trainable parameters and the EMA copy is buffers only (not in parameters()). | |
| (c) strict future / target construction: build_targets reads only the horizon indices — poisoning P / S at every other step and the road quantities of | |
| other horizons leaves s / v / row_ok bitwise unchanged; invalid components are zero in s and False in v; M gates every component; ttc_inv semantics. | |
| (d) gradient paths: the prediction term back-propagates into the TRUNKS (U GRU, V TCN, E encoder) and the interaction modules — the stop-gradient sits on | |
| the target side only (the EMA target has no grad; the online encoder receives no gradient from the prediction term, only from rec / var). | |
| (e) EMA update: buffers move to m * ema + (1 - m) * online; a collapsed latent gets variance penalty ~1, a unit-variance one ~0. | |
| (f) registry: the lat variants equal the parent except for `latent_aux`; lam 0 / 0.1 / 0.3. | |
| Run: cd cwm3 && python -m pytest tests/test_latent_aux.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 | |
| import yaml # 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.cwm3 import CWM3Joint, joint_loss # noqa: E402 | |
| from models import latent_aux as LA # noqa: E402 | |
| import train_cwm3_joint as TJ # noqa: E402 | |
| KEYS = sorted(set(bt.KEYS["V"]) | set(bt.KEYS["U"]) | set(bt.KEYS["E"])) | |
| CFG = {"V": {"can_hist": "4s"}, "U": {"cabin": "none"}, "E": {"tokens": "pooled", "n_slots": 16}, "R": {"tokens": "region", "n_slots": 16}} | |
| LAM = {"hz": 0.5, "K": 0.5, "F": 0.1} | |
| LAT = {"lam": 0.1, "d_z": 16, "d_hidden": 256, "d_enc": 64, "ema": 0.99, "var_coef": 1.0, "rec_coef": 1.0, "horizons": [5, 10, 15]} | |
| def _raw(n=64, seed=3): | |
| raw = ds.synthetic_raw(n, seed=seed, mode=0, need=ds.ALL_STREAMS) | |
| raw["mode"][::2] = 1 | |
| raw = TJ.synthetic_aux(raw, seed) | |
| raw = TJ.synthetic_latent(raw, {"latent_aux": LAT}, seed) | |
| return raw | |
| def _batch(n=16, seed=3, latent=True): | |
| raw = _raw(64, seed) | |
| 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)}) | |
| st.update(TJ.aux_stats(raw)) | |
| if latent: | |
| st.update(LA.latent_stats(raw["LYf"], raw["LYv"])) | |
| st["lat_horizons"] = np.asarray(LAT["horizons"], np.int64) | |
| store = TJ.make_store(raw, "cpu", bank) | |
| idx = torch.arange(n) | |
| b, tg, _ = TJ.prep_joint(store, idx, st, CFG, True, True) | |
| return b, tg, st, store | |
| def _model(seed=0, **kw): | |
| V = build_branch("V", seed, d=256, hidden=256, dropout=0.1) | |
| U = build_branch("U", seed, hid=128, hidden=128, dropout=0.0, cabin="none") | |
| E = build_branch("E", seed, tokens="pooled", n_slots=16, d=128, n_layers=2, n_heads=4, d_ff=256, hidden=256, dropout=0.1) | |
| m = CWM3Joint(V, U, E, decoder="direct", k_head=True, seed=seed, interactions=("env2veh", "drv2veh"), **kw) | |
| m.eval() | |
| return m | |
| # ----------------------------------------------------------------------------- (a) lam = 0 is the parent path | |
| def test_lam0_is_true_noop(): | |
| c0 = TJ.stage_cfg("I1", "all_minus_road2drv_lat0") | |
| cp = TJ.stage_cfg("I1", "all_minus_road2drv") | |
| assert not TJ.latent_enabled(c0) and TJ.latent_lam(c0) == 0.0 | |
| assert TJ.latent_enabled(TJ.stage_cfg("I1", "all_minus_road2drv_lat01")) | |
| assert {k: v for k, v in c0.items() if k != "latent_aux"} == {k: v for k, v in cp.items() if k != "variant"} | {"variant": "all_minus_road2drv_lat0"} | |
| m0, mp = _model(0), _model(0) | |
| assert m0.latent_aux is None and list(m0.state_dict()) == list(mp.state_dict()) | |
| assert m0.param_report()["latent_aux"] == 0 and m0.param_report()["total"] == mp.param_report()["total"] | |
| b, tg, _, _ = _batch(latent=False) | |
| assert "LAT" not in tg | |
| out = m0(b) | |
| assert "LAT" not in out | |
| loss, parts = joint_loss(out, tg, LAM) | |
| loss2, parts2 = TJ.add_latent_loss(m0, 0.0, out, tg, loss, dict(parts)) | |
| assert loss2 is loss and parts2 == parts | |
| # and with the latent targets present in the batch but no module: still the identity | |
| b, tg, _, _ = _batch(latent=True) | |
| assert "LAT" in tg | |
| out = m0(b) | |
| loss, parts = joint_loss(out, tg, LAM) | |
| loss3, _ = TJ.add_latent_loss(m0, 0.0, out, tg, loss, dict(parts)) | |
| assert torch.equal(loss3, loss) | |
| # ----------------------------------------------------------------------------- (b) lam > 0 | |
| def test_latent_forward_and_loss(): | |
| m = _model(0, latent_aux=LAT) | |
| assert m.latent_aux is not None | |
| b, tg, _, _ = _batch() | |
| out = m(b) | |
| assert out["LAT"].shape == (16, 3, 16) | |
| base, parts = joint_loss(out, tg, LAM) | |
| total, parts2 = TJ.add_latent_loss(m, LAT["lam"], out, tg, base, dict(parts)) | |
| la = parts2["lat_total"] | |
| assert torch.equal(total, base + LAT["lam"] * la) | |
| for k in ("lat_pred", "lat_rec", "lat_var"): | |
| assert torch.isfinite(parts2[k]) | |
| assert not any(k.startswith("chain_") for k in parts2 if k.startswith("lat_")) | |
| from models import heads as hd | |
| assert hd.selection_objective(parts2) == hd.selection_objective(parts) | |
| rep = m.param_report() | |
| assert 0 < rep["latent_aux"] < 250_000, rep["latent_aux"] | |
| assert rep["latent_aux_ema_buffers"] > 0 | |
| assert not any("ema_" in n for n, _ in m.named_parameters()) | |
| assert any("ema_" in n for n, _ in m.named_buffers()) | |
| assert set(m.state_dict()) >= {"latent_aux.ema_w1", "latent_aux.ema_b1", "latent_aux.ema_w2", "latent_aux.ema_b2"} | |
| assert rep["total"] <= 4_500_000 | |
| # ----------------------------------------------------------------------------- (c) strict future / target construction | |
| def test_build_targets_reads_only_horizons_and_masks(): | |
| b, tg, st, store = _batch() | |
| idx = torch.arange(16) | |
| args = (store.t["LYf"][idx], store.t["LYv"][idx], store.t["P"][idx], store.t["Pvalid"][idx], store.t["S"][idx].float(), store.t["M"][idx], st["lat_mu"], st["lat_sd"]) | |
| t0 = LA.build_targets(*args) | |
| assert t0["s"].shape == (16, 3, 13) and t0["v"].shape == (16, 3, 13) and t0["row_ok"].shape == (16, 3) | |
| # poison every non-horizon step of P / S and check nothing moves | |
| P, Pv, S = args[2].clone(), args[3].clone(), args[4].clone() | |
| keep = torch.zeros(25, dtype=torch.bool) | |
| keep[[4, 9, 14]] = True | |
| P[:, ~keep] = ~P[:, ~keep] | |
| Pv[:, ~keep] = ~Pv[:, ~keep] | |
| S[:, ~keep] = 1.0 - torch.nan_to_num(S[:, ~keep]) | |
| t1 = LA.build_targets(args[0], args[1], P, Pv, S, args[5], st["lat_mu"], st["lat_sd"]) | |
| for k in ("s", "v", "row_ok"): | |
| assert torch.equal(t0[k], t1[k]), k | |
| # invalid components are exactly 0 in s and False in v; M gates everything | |
| assert torch.all(t0["s"][~t0["v"]] == 0) | |
| Mz = torch.zeros_like(args[5]) | |
| tz = LA.build_targets(*args[:5], Mz, st["lat_mu"], st["lat_sd"]) | |
| assert not tz["v"].any() and not tz["row_ok"].any() and torch.all(tz["s"] == 0) | |
| # ttc_inv: approaching lead 20 m at -10 m/s -> 0.5; receding -> 0; clipped at 1 | |
| d = torch.tensor([20.0, 20.0, 2.0]); v = torch.tensor([-10.0, 3.0, -50.0]) | |
| assert torch.allclose(LA.ttc_inv(d, v), torch.tensor([0.5, 0.0, 1.0])) | |
| # validity: ttc channel valid iff the lead channel is | |
| Yv = torch.tensor([[True, False, True, False]]) | |
| assert LA.continuous_valid(Yv).tolist() == [[True, False, True, False, True]] | |
| def test_attach_latent_targets_witness_and_slice(tmp_path, monkeypatch): | |
| """the host-side slice: Yf [N,25,4] -> [N,K,4] at k-1; uid / t witness enforced.""" | |
| rng = np.random.default_rng(0) | |
| N = 40 | |
| raw = {"t": np.arange(N, dtype=np.float64), "uid": np.array([f"u{i}" for i in range(N)]), "regime_id": np.zeros(N, np.int64), "regime_row": np.arange(N)} | |
| Yf = rng.normal(size=(N, 25, 4)).astype(np.float32) | |
| Yv = rng.random((N, 25, 4)) < 0.7 | |
| np.savez(tmp_path / "road_dyn_manual_val.npz", Yf=Yf, Yvalid=Yv, uid=raw["uid"], t=raw["t"], target_channels=np.array(list(LA.CONT[:4]))) | |
| monkeypatch.setattr(TJ, "ROAD_DYN", tmp_path) | |
| out = TJ.attach_latent_targets(dict(raw), ["manual"], "val", (5, 10, 15)) | |
| assert out["LYf"].shape == (N, 3, 4) and np.array_equal(out["LYf"], Yf[:, [4, 9, 14]]) and np.array_equal(out["LYv"], Yv[:, [4, 9, 14]]) | |
| bad = dict(raw); bad["uid"] = bad["uid"].copy(); bad["uid"][3] = "zz" | |
| try: | |
| TJ.attach_latent_targets(bad, ["manual"], "val", (5, 10, 15)) | |
| assert False, "witness not enforced" | |
| except AssertionError as ex: | |
| assert "witness" in str(ex) | |
| # ----------------------------------------------------------------------------- (d) gradient paths | |
| def test_prediction_term_reaches_trunks_but_not_target_side(): | |
| m = _model(0, latent_aux=dict(LAT, var_coef=0.0, rec_coef=0.0)) | |
| b, tg, _, _ = _batch() | |
| out = m(b) | |
| la, parts = m.latent_aux.loss(out["LAT"], tg["LAT"]) | |
| assert torch.equal(la, parts["lat_pred"]) | |
| m.zero_grad(set_to_none=True) | |
| la.backward() | |
| def gnorm(prefix): | |
| return sum(float(p.grad.abs().sum()) for n, p in m.named_parameters() if n.startswith(prefix) and p.grad is not None) | |
| assert gnorm("U.gru") > 0 and gnorm("V.tcn") > 0 and gnorm("E.enc") > 0, "the auxiliary loss must reach the trunks (stated in the pre-registration)" | |
| assert gnorm("latent_aux.pred") > 0 | |
| assert all(p.grad is None or float(p.grad.abs().sum()) == 0 for n, p in m.named_parameters() if n.startswith(("latent_aux.enc", "latent_aux.dec"))), "stop-gradient on the target side" | |
| assert all(p.grad is None for n, p in m.named_parameters() if n.startswith(("dec_", "k_head", "p_head", "r_head"))) | |
| # rec + var train the online encoder / decoder, never the predictor or the trunks | |
| m2 = _model(0, latent_aux=LAT) | |
| out2 = m2(b) | |
| _, parts2 = m2.latent_aux.loss(out2["LAT"], tg["LAT"]) | |
| m2.zero_grad(set_to_none=True) | |
| (parts2["lat_rec"] + parts2["lat_var"]).backward() | |
| assert sum(float(p.grad.abs().sum()) for n, p in m2.named_parameters() if n.startswith("latent_aux.enc") and p.grad is not None) > 0 | |
| assert all(p.grad is None for n, p in m2.named_parameters() if not n.startswith(("latent_aux.enc", "latent_aux.dec"))) | |
| # ----------------------------------------------------------------------------- (e) EMA + variance term | |
| def test_ema_update_and_variance_term(): | |
| m = _model(0, latent_aux=LAT) | |
| la = m.latent_aux | |
| w0 = la.ema_w1.clone() | |
| with torch.no_grad(): | |
| la.enc[0].weight.add_(1.0) | |
| la.ema_update() | |
| assert torch.allclose(la.ema_w1, 0.99 * w0 + 0.01 * la.enc[0].weight.detach()) | |
| z_col = torch.zeros(64, 16) + 0.3 | |
| assert abs(float(la.variance_term(z_col)) - (1.0 - np.sqrt(1e-4))) < 1e-4 | |
| torch.manual_seed(0) | |
| z_unit = torch.randn(4096, 16) | |
| assert float(la.variance_term(z_unit)) < 0.03 | |
| assert float(la.variance_term(torch.randn(1, 16))) == 0.0 | |
| # the target is computed without grad | |
| b, tg, _, _ = _batch() | |
| x = la.enc_input(tg["LAT"]["s"], tg["LAT"]["v"]) | |
| assert not la.encode_ema(x).requires_grad | |
| # ----------------------------------------------------------------------------- (f) registry | |
| def test_stage_registry_lat_variants(): | |
| Y = yaml.safe_load((CWM3 / "train" / "stages_joint.yaml").read_text()) | |
| v = Y["stages"]["I1"]["variants"] | |
| parent = v["all_minus_road2drv"] | |
| for k, lam in (("all_minus_road2drv_lat0", 0.0), ("all_minus_road2drv_lat01", 0.1), ("all_minus_road2drv_lat03", 0.3)): | |
| assert {a: b_ for a, b_ in v[k].items() if a != "latent_aux"} == parent, f"{k} deviates from its parent recipe" | |
| la = v[k]["latent_aux"] | |
| assert la["lam"] == lam and la["d_z"] == 16 and la["horizons"] == [5, 10, 15] | |
| assert {a: b_ for a, b_ in la.items() if a != "lam"} == {a: b_ for a, b_ in v["all_minus_road2drv_lat01"]["latent_aux"].items() if a != "lam"} | |
| assert "latent_aux" not in parent | |
Xet Storage Details
- Size:
- 13.2 kB
- Xet hash:
- bb589d25aa67a484148f0c618d7f006700ee46ab5233a162107589c98b54b6c7
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.