Buckets:
| """Guards for CWM3-D-kp "K-v2" (aide/models/cwm3d_v2.py; plan A.1, 2026-09-14): | |
| 1. zero-init exactness: EVERY flag combination predicts exactly the zero-velocity pose (|pred - last| <= 1e-6 px) on random inputs and on | |
| real val clips, in eval AND train mode; anchor_w = 1/3, gate = 0, log_b = log 30, haz_P == the prior for every clip (default and TRAIN prior); | |
| 2. no-future leak: NaN-poisoned future -> pred and haz_logits bit-identical; a 10-frame input is rejected; the last observed frame matters; | |
| 3. causality: part tokens of frame t do not depend on frames > t (Δ-to-last channel fixed); the spatial mixer never crosses frames; | |
| 4. hazard targets / losses: synthetic onset / risk cases, invariance to observed-frame noise, zero-risk BCE = 0, equality with the BATON | |
| chain_bce on-branch and forward_marginal (models/heads.py), hazard_stats <-> hazard_targets consistency on the TRAIN cache; | |
| 5. modulation isolation: HazardGate is the identity at zero parameters; beta / gamma touch the gate only, kappa the mean only, log_b and the | |
| anchor logits never; haz_detach cuts the geometry gradient into the hazard head; | |
| 6. parameter counts: full in [2.3 M, 3.2 M], head + haz + mod < 0.2 M, K's head variant > 4.5 M; | |
| 7. trunk-only init (train_kp.load_init_ckpt) copies exactly TRUNK_PREFIXES and keeps the zero-velocity start; pretrain_cwm3d_baton.build_model | |
| == train_kp.build_model for the same seed; | |
| 8. sampler contract: TrainSampler(hazard_tau) keeps batch() a 2-tuple and emits well-formed haz_onset / haz_risk from batch_extra. | |
| python -m pytest aide/tests/test_cwm3d_v2.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import itertools # noqa: E402 | |
| import math # noqa: E402 | |
| import sys # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| torch.set_num_threads(4) | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.models import cwm3d_v2 as V2, heads, losses as LS # noqa: E402 | |
| from aide.train import pretrain_cwm3d_baton as PRE # noqa: E402 | |
| from aide.train import train_kp as TK # noqa: E402 | |
| from aide.train.train_kp import TrainSampler, load_init_ckpt, seed_all, split_obs_fut # noqa: E402 | |
| from models import heads as BH # noqa: E402 (BATON chain_bce / forward_marginal) | |
| TOL = 1e-6 | |
| PROTO_IDX = P.sample_idx(P.N_FRAMES_STD) | |
| def _random_frames(n=6, seed=0): | |
| g = torch.Generator().manual_seed(seed) | |
| return torch.rand(n, P.N_SAMPLE, P.J, 3, generator=g) * torch.tensor([P.W, P.H, 1.0]) | |
| def _real_frames(n=32, split="val"): | |
| if not P.CACHE_NPZ.exists(): | |
| return None | |
| c = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| return torch.as_tensor(c["S10"][c["split"] == split][:n]).float() | |
| def _train_S10(): | |
| if not P.CACHE_NPZ.exists(): | |
| return None | |
| c = np.load(P.CACHE_NPZ, allow_pickle=False) | |
| return c["S"][c["split"] == "train"][:, PROTO_IDX] | |
| def _perturb(model, std=0.02, seed=1, skip=()): | |
| g = torch.Generator().manual_seed(seed) | |
| with torch.no_grad(): | |
| for n, p in model.named_parameters(): | |
| if not n.startswith(tuple(skip)): | |
| p.add_(std * torch.randn(p.shape, generator=g)) | |
| def _combos(): | |
| """Every flag combination of the model (joint_time only matters with the lean head; s_blocks 0 = LN-only mixer).""" | |
| for spatial, lean, hazard, mod, detach in itertools.product((True, False), repeat=5): | |
| if not hazard and (mod or detach): | |
| continue | |
| for jt in (("all", "last") if lean else ("all",)): | |
| yield dict(spatial=spatial, lean_head=lean, hazard=hazard, haz_mod=mod, haz_detach=detach, joint_time=jt) | |
| yield dict(s_blocks=0) | |
| yield dict(layers=1, d=64, ffn=128, d_j=16, s_heads=2, head_hidden=32, d_h=8, haz_hidden=16, k_head_hidden=64) | |
| # ----------------------------------------------------------------------------- 1. zero-init exactness | |
| def _zero_init_check(frames): | |
| obs, _ = split_obs_fut(frames) | |
| zv = obs[:, -1:, :, :2].expand(-1, P.T_PRED, -1, -1) | |
| n = 0 | |
| for kw in _combos(): | |
| torch.manual_seed(0) | |
| m = V2.CWM3Dv2(**kw) | |
| for mode in ("eval", "train"): | |
| getattr(m, mode)() | |
| with torch.no_grad(): | |
| o = m(obs) | |
| assert o["pred"].shape == zv.shape and torch.isfinite(o["pred"]).all(), kw | |
| assert (o["pred"] - zv).abs().max().item() <= TOL, (kw, mode, (o["pred"] - zv).abs().max().item()) | |
| assert torch.allclose(o["anchor_w"], torch.full_like(o["anchor_w"], 1 / 3)), kw | |
| assert torch.equal(o["anchor_gate"], torch.zeros_like(o["anchor_gate"])), kw | |
| assert torch.allclose(o["log_b"], torch.full_like(o["log_b"], math.log(30.0))), kw | |
| if m.haz is not None: | |
| assert o["haz_logits"].shape == (len(obs), 4, P.T_PRED) and o["haz_P"].shape == (len(obs), 4, P.T_PRED), kw | |
| assert torch.allclose(o["haz_P"], m.haz.prior_P[None].expand_as(o["haz_P"]), atol=1e-5), kw # P == prior for every clip | |
| assert torch.allclose(m.haz.prior_P[:, 0], torch.full((4,), 0.5)) and torch.allclose(m.haz.prior_P[:, -1], torch.full((4,), 1 - 0.5 ** 5)) | |
| else: | |
| assert "haz_P" not in o and "haz_logits" not in o | |
| n += 1 | |
| assert n >= 20, n | |
| def test_zero_init_equals_zero_velocity_random(): | |
| _zero_init_check(_random_frames()) | |
| def test_zero_init_equals_zero_velocity_real_clips(): | |
| fr = _real_frames() | |
| if fr is not None: | |
| _zero_init_check(fr) | |
| def test_zero_init_with_train_prior_and_hazard_head_constant(): | |
| """set_hazard_prior (TRAIN statistics) writes the bias so that P == prior_P for every clip; ZV still exact; buffers in the state dict.""" | |
| S10 = _train_S10() | |
| if S10 is None: | |
| S10 = _random_frames(64, seed=3).numpy() | |
| tau, prior = LS.hazard_stats(S10, 0.75) | |
| assert tau.shape == (4,) and prior.shape == (4, P.T_PRED) and (np.diff(prior, axis=1) >= 0).all() and (prior[:, 0] > 0).all() | |
| torch.manual_seed(0) | |
| m = V2.CWM3Dv2().eval() | |
| m.set_hazard_prior(tau, prior) | |
| frames = torch.cat([_random_frames(), torch.as_tensor(S10[:16]).float()]) | |
| obs, _ = split_obs_fut(frames) | |
| with torch.no_grad(): | |
| o = m(obs) | |
| zv = obs[:, -1:, :, :2].expand_as(o["pred"]) | |
| assert (o["pred"] - zv).abs().max().item() <= TOL | |
| assert torch.allclose(o["haz_P"], torch.as_tensor(prior, dtype=torch.float32)[None].expand_as(o["haz_P"]), atol=1e-5) | |
| assert torch.allclose(m.haz.tau, torch.as_tensor(tau, dtype=torch.float32)) | |
| sd = m.state_dict() | |
| for k in ("haz.tau", "haz.prior_P", "haz.prior_lam", "haz.bias"): | |
| assert k in sd, k | |
| # the per-step prior hazards reproduce the cumulative prior through the chain | |
| assert torch.allclose(1 - torch.cumprod(1 - m.haz.prior_lam, -1), m.haz.prior_P, atol=1e-5) | |
| # ----------------------------------------------------------------------------- 2. no-future leak | |
| def test_model_never_reads_future_frames(): | |
| frames = _random_frames() | |
| fr = _real_frames(8) | |
| if fr is not None: | |
| frames = torch.cat([frames, fr], 0) | |
| poison = frames.clone() | |
| poison[:, P.T_OBS:] = float("nan") | |
| obs_c, fut_c = split_obs_fut(frames) | |
| obs_p, fut_p = split_obs_fut(poison) | |
| assert torch.equal(obs_c, obs_p) and torch.isnan(fut_p).all() and torch.isfinite(fut_c).all() | |
| for kw in ({}, {"spatial": False}, {"lean_head": False}, {"hazard": False}, {"haz_mod": False}, {"joint_time": "last"}): | |
| torch.manual_seed(0) | |
| m = V2.CWM3Dv2(**kw).eval() | |
| _perturb(m) | |
| with torch.no_grad(): | |
| a, b = m(obs_c), m(obs_p) | |
| assert torch.isfinite(a["pred"]).all() and torch.equal(a["pred"], b["pred"]), kw | |
| if m.haz is not None: | |
| assert torch.isfinite(a["haz_logits"]).all() and torch.equal(a["haz_logits"], b["haz_logits"]) and torch.equal(a["haz_P"], b["haz_P"]), kw | |
| try: | |
| with torch.no_grad(): | |
| m(poison) | |
| raise RuntimeError(f"{kw} accepted {P.N_SAMPLE} frames") | |
| except AssertionError: | |
| pass | |
| shifted = obs_c.clone() | |
| shifted[:, -1, :, :2] += 30.0 | |
| with torch.no_grad(): | |
| c = m(shifted) | |
| assert not torch.allclose(a["pred"], c["pred"]), kw | |
| if m.haz is not None: | |
| assert not torch.allclose(a["haz_logits"], c["haz_logits"]), kw | |
| # ----------------------------------------------------------------------------- 3. causality | |
| def test_encoder_causal_across_observed_frames_and_spatial_within_frame(): | |
| torch.manual_seed(0) | |
| m = V2.CWM3Dv2().eval() | |
| _perturb(m) | |
| obs, _ = split_obs_fut(_random_frames()) | |
| x2 = obs.clone() | |
| x2[:, 2:P.T_OBS - 1, :, :2] += 40.0 # frames 2, 3 move; frame 4 (the Δ reference) fixed | |
| with torch.no_grad(): | |
| p1, u1 = m.encode(obs) | |
| p2, u2 = m.encode(x2) | |
| assert torch.allclose(p1[:, :2], p2[:, :2], atol=1e-5) and not torch.allclose(p1[:, 2:], p2[:, 2:]) | |
| assert not torch.allclose(u1, u2) # the joint read-out ("all") sees frames 2, 3 | |
| # attention mask: same frame fully connected, later frames blocked | |
| mk = m.mask | |
| L = P.T_OBS * 4 | |
| frame = torch.arange(L) // 4 | |
| assert torch.equal(mk, frame[None, :] > frame[:, None]) | |
| # spatial mixer + pooling act within a frame: tokens of frame 0 unchanged when frame 1 moves (Δ-to-last unchanged for frame 0) | |
| x3 = obs.clone() | |
| x3[:, 1, :, :2] += 25.0 | |
| with torch.no_grad(): | |
| s1, s3 = m.spatial(m.jtok(obs)), m.spatial(m.jtok(x3)) | |
| q1, q3 = m.pool(s1), m.pool(s3) | |
| assert torch.allclose(s1[:, 0], s3[:, 0], atol=1e-5) and not torch.allclose(s1[:, 1], s3[:, 1]) | |
| assert torch.allclose(q1[:, 0], q3[:, 0], atol=1e-5) and not torch.allclose(q1[:, 1], q3[:, 1]) | |
| # "last" read-out depends on the last frame only | |
| torch.manual_seed(0) | |
| ml = V2.CWM3Dv2(joint_time="last").eval() | |
| _perturb(ml) | |
| with torch.no_grad(): | |
| _, ul1 = ml.encode(obs) | |
| _, ul2 = ml.encode(x2) | |
| assert torch.allclose(ul1, ul2, atol=1e-5) | |
| # attention pooling is a proper masked softmax: one part's pooled token ignores the other parts' joints | |
| x4 = obs.clone() | |
| x4[:, :, P.PARTS["face"], :2] += 33.0 | |
| with torch.no_grad(): | |
| s4 = m.spatial(m.jtok(x4)) | |
| att1 = torch.einsum("pd,btjd->btpj", m.pool.q, s1) * m.pool.scale | |
| att1 = att1.masked_fill(~m.pool.mask[None, None], float("-inf")).softmax(-1) | |
| assert torch.equal(att1[..., 0, P.PARTS["face"]], torch.zeros_like(att1[..., 0, P.PARTS["face"]])) # body query -> zero weight on face | |
| assert torch.allclose(att1.sum(-1), torch.ones_like(att1.sum(-1)), atol=1e-5) | |
| # ----------------------------------------------------------------------------- 4. hazard targets and losses | |
| def _synthetic(jump_h, B=3, part="rhand", amp=300.0): | |
| """Static clip with the given part's joints jumping by `amp` px from future step jump_h (1-based; 0 = never).""" | |
| g = torch.Generator().manual_seed(11) | |
| frames = torch.rand(B, P.N_SAMPLE, P.J, 3, generator=g) * torch.tensor([P.W / 2, P.H / 2, 1.0]) + 100.0 | |
| frames[:, 1:, :, :2] = frames[:, :1, :, :2] # perfectly static | |
| if jump_h: | |
| frames[:, P.T_OBS + jump_h - 1:, P.PARTS[part], 0] += amp | |
| return frames | |
| def test_hazard_targets_synthetic_cases(): | |
| tau = torch.tensor([78.5, 32.9, 109.1, 135.2]) | |
| p = P.PART_NAMES.index("rhand") | |
| frames = _synthetic(2) # jump at future frame index 6 -> h = 2 | |
| obs, fut = split_obs_fut(frames) | |
| on, rk = LS.hazard_targets(obs[:, -1, :, :2], fut, tau) | |
| assert on.shape == rk.shape == (3, 4, P.T_PRED) and on.dtype == fut.dtype | |
| assert torch.equal(on[:, p], torch.tensor([[0., 1., 0., 0., 0.]]).expand(3, -1)) | |
| assert torch.equal(rk[:, p], torch.tensor([[1., 1., 0., 0., 0.]]).expand(3, -1)) | |
| others = [q for q in range(4) if q != p] | |
| assert torch.equal(on[:, others], torch.zeros_like(on[:, others])) and torch.equal(rk[:, others], torch.ones_like(rk[:, others])) | |
| # jump at h = 1 -> onset at h = 1, risk only at h = 1; never -> no onset, risk everywhere | |
| obs1, fut1 = split_obs_fut(_synthetic(1)) | |
| on1, rk1 = LS.hazard_targets(obs1[:, -1, :, :2], fut1, tau) | |
| assert torch.equal(on1[:, p], torch.tensor([[1., 0., 0., 0., 0.]]).expand(3, -1)) and torch.equal(rk1[:, p], torch.tensor([[1., 0., 0., 0., 0.]]).expand(3, -1)) | |
| on0, rk0 = LS.hazard_targets(split_obs_fut(_synthetic(0))[0][:, -1, :, :2], split_obs_fut(_synthetic(0))[1], tau) | |
| assert torch.equal(on0, torch.zeros_like(on0)) and torch.equal(rk0, torch.ones_like(rk0)) | |
| # a transient excursion that returns still counts as onset once (first exceedance), risk stays 0 afterwards | |
| fr = _synthetic(0) | |
| fr[:, P.T_OBS + 2, P.PARTS["body"], 1] += 500.0 # h = 3 only | |
| o3, r3 = LS.hazard_targets(split_obs_fut(fr)[0][:, -1, :, :2], split_obs_fut(fr)[1], tau) | |
| b = P.PART_NAMES.index("body") | |
| assert torch.equal(o3[:, b], torch.tensor([[0., 0., 1., 0., 0.]]).expand(3, -1)) and torch.equal(r3[:, b], torch.tensor([[1., 1., 1., 0., 0.]]).expand(3, -1)) | |
| # structural invariants | |
| for o_, r_ in ((on, rk), (on1, rk1), (o3, r3)): | |
| assert (o_ <= r_).all() and (r_[..., 1:] <= r_[..., :-1]).all() and (o_.sum(-1) <= 1).all() | |
| def test_hazard_targets_ignore_observed_frame_noise_and_use_the_trainer_cut(): | |
| tau = torch.tensor([78.5, 32.9, 109.1, 135.2]) | |
| frames = _synthetic(3) | |
| noisy = frames.clone() | |
| noisy[:, :P.T_OBS - 1, :, :2] += 50.0 * torch.randn(3, P.T_OBS - 1, P.J, 2, generator=torch.Generator().manual_seed(5)) | |
| noisy[:, P.T_OBS - 1, :, :2] += 0.5 * torch.randn(3, P.J, 2, generator=torch.Generator().manual_seed(6)) # sub-threshold jitter of the reference frame | |
| a = LS.hazard_targets(split_obs_fut(frames)[0][:, -1, :, :2], split_obs_fut(frames)[1], tau) | |
| b = LS.hazard_targets(split_obs_fut(noisy)[0][:, -1, :, :2], split_obs_fut(noisy)[1], tau) | |
| assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]) | |
| # targets depend on the future block: poisoning it changes them (so they can never be computed from obs alone) | |
| poison = frames.clone() | |
| poison[:, P.T_OBS:] = float("nan") | |
| c = LS.hazard_targets(split_obs_fut(poison)[0][:, -1, :, :2], split_obs_fut(poison)[1], tau) | |
| assert not torch.equal(a[0], c[0]) | |
| def test_hazard_bce_matches_baton_chain_and_cumulative_matches_forward_marginal(): | |
| g = torch.Generator().manual_seed(2) | |
| B, H = 64, P.T_PRED | |
| logits = torch.randn(B, 4, H, generator=g) * 3 | |
| ex = (torch.rand(B, 4, H, generator=g) < 0.3).float() | |
| cum = torch.cummax(ex, -1).values | |
| prev = torch.cat([torch.zeros_like(cum[..., :1]), cum[..., :-1]], -1) | |
| onset, risk = ex * (1 - prev), 1 - prev | |
| l_v2 = LS.hazard_bce(logits, onset, risk) | |
| # BATON two-state chain: z_prev = prev, z = cum, m = risk (the off branch is masked out) -> masked mean == the on-branch NLL | |
| l_bt = BH.masked_mean(BH.chain_bce(logits, torch.zeros_like(logits), prev, cum, risk), risk) | |
| assert torch.allclose(l_v2, l_bt, atol=1e-5), (l_v2.item(), l_bt.item()) | |
| # per-element equality of the on-branch too (rtol: chain_bce evaluates log(1 - sigmoid(x)) in float32, which loses ~1e-4 relative | |
| # precision for |x| > 6, the cancellation BCEWithLogits avoids; both are the same formula) | |
| per = torch.nn.functional.binary_cross_entropy_with_logits(logits, onset, reduction="none") * risk | |
| assert torch.allclose(per, BH.chain_bce(logits, torch.zeros_like(logits), prev, cum, risk), rtol=1e-4, atol=1e-6) | |
| small = logits.clamp(-4, 4) | |
| assert torch.allclose(torch.nn.functional.binary_cross_entropy_with_logits(small, onset, reduction="none") * risk, | |
| BH.chain_bce(small, torch.zeros_like(small), prev, cum, risk), atol=1e-6) | |
| # zero risk -> loss exactly 0 (and finite) | |
| assert LS.hazard_bce(logits, onset, torch.zeros_like(risk)).item() == 0.0 | |
| # cumulative == forward_marginal with h_off = 0, z0 = 0 | |
| Pc = V2.HazardHead.cumulative(logits) | |
| lam = torch.sigmoid(logits).clamp(BH.EPS_H, 1 - BH.EPS_H) | |
| for p in range(4): | |
| fm = BH.forward_marginal(lam[:, p], torch.zeros(B, H), torch.zeros(B)) | |
| assert torch.allclose(Pc[:, p], fm, atol=1e-6) | |
| assert (Pc[..., 1:] >= Pc[..., :-1]).all() and (Pc > 0).all() and (Pc <= 1).all() # == 1.0 possible in float32 when every lam is clamped near 1 | |
| # lam_from_P inverts cumulative | |
| lam2 = V2.HazardHead.lam_from_P(Pc) | |
| assert torch.allclose(1 - torch.cumprod(1 - lam2, -1), Pc, atol=1e-5) | |
| def test_hazard_stats_consistent_with_targets_on_train_cache(): | |
| S10 = _train_S10() | |
| if S10 is None: | |
| return | |
| tau, prior = LS.hazard_stats(S10, 0.75) | |
| x = torch.as_tensor(S10) | |
| on, rk = LS.hazard_targets(x[:, P.T_OBS - 1, :, :2], x[:, P.T_OBS:, :, :2], tau) | |
| cum = torch.cummax(on, -1).values.numpy() | |
| assert np.allclose(cum.mean(0), prior, atol=1e-6) | |
| # q = 0.75 -> exactly ~25 % of the clips exceed tau_p AT h = 5 (the quantile's own horizon); the cumulative prior is >= that | |
| for p_, k in enumerate(P.PART_NAMES): | |
| d5 = np.linalg.norm(S10[:, -1, P.PARTS[k], :2].mean(1) - S10[:, P.T_OBS - 1, P.PARTS[k], :2].mean(1), axis=-1) | |
| assert abs((d5 > tau[p_]).mean() - 0.25) < 2.0 / len(S10) and prior[p_, -1] >= 0.25 - 1e-6, k | |
| assert np.allclose(tau, [78.5, 32.9, 109.1, 135.2], atol=0.1), tau # plan A.1 numbers (TRAIN split, protocol grid) | |
| tau6, _ = LS.hazard_stats(S10, 0.6) | |
| assert (tau6 < tau).all() | |
| def test_rank_auc(): | |
| assert TK.rank_auc([0.1, 0.4, 0.35, 0.8], [0, 0, 1, 1]) == 0.75 | |
| assert TK.rank_auc([1, 2, 3, 4], [0, 0, 1, 1]) == 1.0 and TK.rank_auc([4, 3, 2, 1], [0, 0, 1, 1]) == 0.0 | |
| assert TK.rank_auc([1, 1, 1, 1], [0, 1, 0, 1]) == 0.5 and math.isnan(TK.rank_auc([1, 2], [1, 1])) | |
| # ----------------------------------------------------------------------------- 5. modulation isolation | |
| def test_hazard_gate_isolation_and_detach(): | |
| torch.manual_seed(0) | |
| m = V2.CWM3Dv2().eval() | |
| _perturb(m, skip=("mod.",)) # non-trivial head output and hazards; gate parameters still zero | |
| obs, _ = split_obs_fut(_random_frames()) | |
| with torch.no_grad(): | |
| parts, u = m.encode(obs) | |
| z = parts[:, -1] | |
| ho = m.head(u, z) | |
| Pc = m.haz.cumulative(m.haz(z)) | |
| assert torch.equal(m.mod(ho, Pc, m.haz.prior_P), ho) # identity at zero parameters (bit-exact) | |
| C = heads.HorizonHead | |
| for name, ch in (("gamma", C.C_GATE), ("beta", C.C_GATE), ("kappa", C.C_MU)): | |
| getattr(m.mod, name).fill_(0.7) | |
| ho2 = m.mod(ho, Pc, m.haz.prior_P) | |
| getattr(m.mod, name).zero_() | |
| touched = torch.zeros(C.N_C, dtype=torch.bool) | |
| touched[ch] = True | |
| assert torch.equal(ho2[..., ~touched], ho[..., ~touched]), name # log_b, anchor logits, the other channel untouched | |
| assert not torch.equal(ho2[..., ch], ho[..., ch]), name | |
| # with P == prior (Pc = 0) gamma and kappa are inert, beta still scales the gate through P itself | |
| prior = m.haz.prior_P | |
| Pp = prior[None].expand_as(Pc).clone() | |
| m.mod.gamma.fill_(1.0); m.mod.kappa.fill_(1.0) | |
| assert torch.allclose(m.mod(ho, Pp, prior), ho, atol=1e-6) | |
| m.mod.gamma.zero_(); m.mod.kappa.zero_(); m.mod.beta.fill_(1.0) | |
| ho3 = m.mod(ho, Pp, prior) | |
| poj = V2.part_of_joint() | |
| assert torch.allclose(ho3[..., C.C_GATE], ho[..., C.C_GATE] * (1 - prior[poj].t()[None]), atol=1e-6) | |
| m.mod.beta.zero_() | |
| # haz_detach: geometry gradient does not reach the hazard head when detached, does otherwise | |
| for detach, expect_grad in ((True, False), (False, True)): | |
| torch.manual_seed(0) | |
| md = V2.CWM3Dv2(haz_detach=detach) | |
| _perturb(md, skip=("mod.",)) | |
| with torch.no_grad(): | |
| md.mod.beta.fill_(0.5); md.mod.kappa.fill_(0.5) | |
| out = md(obs) | |
| out["pred"].sum().backward() | |
| gw = md.haz.net[-1].weight.grad | |
| has = gw is not None and gw.abs().sum().item() > 0 | |
| assert has == expect_grad, (detach, has) | |
| assert md.head.mlp[-1].weight.grad.abs().sum().item() > 0 # the head is always trained by geometry | |
| # ----------------------------------------------------------------------------- 6. parameter counts | |
| def test_param_counts(): | |
| full = V2.CWM3Dv2() | |
| n = V2.param_breakdown(full) | |
| assert 2.3e6 < n["total"] < 3.2e6, n | |
| assert n["head"] + n["haz"] + n["mod"] < 0.2e6, n | |
| assert n["mod"] == 15 and n["tok"] == 150784 and n["blocks"] == 3 * sum(p.numel() for p in heads.Block(256, 4, 1024).parameters()) | |
| nk = V2.param_breakdown(V2.CWM3Dv2(lean_head=False)) | |
| assert nk["total"] > 4.5e6 and nk["head"] == sum(p.numel() for p in heads.HorizonHead(256, 512).parameters()), nk | |
| ns = V2.param_breakdown(V2.CWM3Dv2(spatial=False)) | |
| assert "spatial" not in ns and "pool" not in ns and ns["total"] == n["total"] - n["spatial"] - n["pool"] | |
| nh = V2.param_breakdown(V2.CWM3Dv2(hazard=False)) | |
| assert "haz" not in nh and "mod" not in nh and nh["total"] == n["total"] - n["haz"] - n["mod"] | |
| assert V2.param_breakdown(V2.CWM3Dv2(haz_mod=False))["total"] == n["total"] - 15 | |
| # ----------------------------------------------------------------------------- 7. trunk-only init and pretrain / train_kp seed equality | |
| def test_trunk_init_loads_trunk_only_and_keeps_zero_velocity(tmp_path): | |
| obs, _ = split_obs_fut(_random_frames()) | |
| torch.manual_seed(1) | |
| src = V2.CWM3Dv2() | |
| _perturb(src, std=0.05, seed=2) # a "pretrained" model: every tensor differs from the seed init | |
| src.set_hazard_prior(np.array([10.0, 20.0, 30.0, 40.0]), np.tile(np.linspace(0.1, 0.5, P.T_PRED), (4, 1))) | |
| torch.save({"model": src.state_dict(), "epoch": 3}, tmp_path / "v2.pt") | |
| torch.manual_seed(7) | |
| dst = V2.CWM3Dv2() | |
| before = {k: v.clone() for k, v in dst.state_dict().items()} | |
| rep = load_init_ckpt(dst, str(tmp_path / "v2.pt"), "trunk", "cpu") | |
| after, ssd = dst.state_dict(), src.state_dict() | |
| trunk = tuple(dst.TRUNK_PREFIXES) | |
| assert trunk == ("tok.", "jtok.", "spatial.", "pool.", "jtime.", "blocks.", "ln_f.") | |
| for k in after: | |
| if k.startswith(trunk): | |
| assert torch.equal(after[k], ssd[k]), k | |
| else: | |
| assert torch.equal(after[k], before[k]), k # head / haz / mod / mask kept at the seed init | |
| assert rep["parts"] == "trunk" and rep["n_skipped_shape"] == 0 | |
| assert rep["n_loaded"] == sum(k.startswith(trunk) for k in after) > 0 | |
| assert all(k.startswith(("head.", "haz.", "mod.", "mask")) for k in rep["kept_keys"]), rep["kept_keys"] | |
| with torch.no_grad(): | |
| o = dst.eval()(obs) | |
| assert (o["pred"] - obs[:, -1:, :, :2]).abs().max().item() <= TOL # zero-velocity start of the fine-tune | |
| assert torch.allclose(o["haz_P"], dst.haz.prior_P[None].expand_as(o["haz_P"]), atol=1e-5) | |
| # full load of the same checkpoint reproduces the source exactly (buffers included) -> nothing is silently shape-skipped | |
| torch.manual_seed(9) | |
| full = V2.CWM3Dv2() | |
| rep2 = load_init_ckpt(full, str(tmp_path / "v2.pt"), "all", "cpu") | |
| assert rep2["n_skipped_shape"] == 0 and rep2["n_missing_or_kept"] == 0 | |
| for k, v in full.state_dict().items(): | |
| assert torch.equal(v, ssd[k]), k | |
| def test_pretrain_build_model_matches_train_kp_seed_init(): | |
| for argv in (["--model", "cwm3d_v2"], ["--model", "cwm3d_v2", "--no-lean-head", "--haz-mod", "0", "--v2-layers", "2"]): | |
| a = TK.get_args(argv) | |
| seed_all(123) | |
| m_pre, files = PRE.build_model(a) | |
| seed_all(123) | |
| m_tr, files_tr = TK.build_model("cwm3d_v2", a) | |
| assert type(m_pre) is V2.CWM3Dv2 and m_pre.TRUNK_PREFIXES == m_tr.TRUNK_PREFIXES | |
| sp, st = m_pre.state_dict(), m_tr.state_dict() | |
| assert list(sp) == list(st) | |
| for k in sp: | |
| assert torch.equal(sp[k], st[k]), k | |
| assert any(Path(f).name == "cwm3d_v2.py" for f in files) and any(Path(f).name == "cwm3d_v2.py" for f in files_tr) | |
| for need in ("cwm3d_v2.py", "heads.py", "cwm3d_spatial.py", "losses.py"): | |
| assert any(Path(f).name == need for f in files_tr), need | |
| # a different seed gives different trunk weights (the equality above is not trivial) | |
| seed_all(1) | |
| m1, _ = PRE.build_model(TK.get_args(["--model", "cwm3d_v2"])) | |
| seed_all(2) | |
| m2, _ = PRE.build_model(TK.get_args(["--model", "cwm3d_v2"])) | |
| assert not torch.equal(m1.tok.body.weight, m2.tok.body.weight) | |
| def test_flags_map_to_kwargs_and_k_layers_flag_is_not_reused(): | |
| a = TK.get_args(["--model", "cwm3d_v2", "--layers", "9", "--no-v2-spatial", "--no-lean-head", "--no-hazard", "--v2-joint-time", "last", | |
| "--v2-layers", "2", "--v2-dj", "24", "--head-hidden", "128", "--haz-detach", "0"]) | |
| kw = V2.kwargs_from_args(a) | |
| assert kw["layers"] == 2 and kw["spatial"] is False and kw["lean_head"] is False and kw["hazard"] is False and kw["joint_time"] == "last" | |
| assert kw["d_j"] == 24 and kw["k_head_hidden"] == 128 and kw["haz_detach"] is False and kw["dropout"] == 0.0 | |
| m = V2.build("cwm3d_v2", **kw) | |
| assert len(m.blocks) == 2 and m.spatial is None and m.haz is None and m.mod is None and isinstance(m.head, heads.HorizonHead) | |
| d = V2.kwargs_from_args(TK.get_args(["--model", "cwm3d_v2"])) | |
| assert d == {**V2.V2_DEFAULTS, "dropout": 0.0} | |
| a0 = TK.get_args(["--model", "cwm3d_v2"]) | |
| assert (a0.noise_px, a0.joint_drop, a0.translate, a0.scale) == (2.0, 0.1, 20.0, 0.05) # hardening-D recipe defaults apply to cwm3d_v2 | |
| assert a0.haz_weight == 1.0 and a0.haz_quantile == 0.75 and a0.haz_mod == 1 and a0.haz_detach == 1 | |
| try: | |
| V2.build("cwm3d_v3") | |
| raise RuntimeError("unknown name accepted") | |
| except ValueError: | |
| pass | |
| # ----------------------------------------------------------------------------- 8. sampler contract | |
| def test_train_sampler_hazard_targets_contract(): | |
| S45 = np.zeros((4, P.N_FRAMES_STD, P.J, 3), np.float32) + 500.0 | |
| S45[..., 2] = 1.0 | |
| S45[:, :, P.PARTS["rhand"], 0] += 12.0 * np.arange(P.N_FRAMES_STD)[None, :, None] # right hand drifts 12 px / frame, the rest static | |
| tau = np.array([78.5, 32.9, 109.1, 135.2]) | |
| smp = TrainSampler(S45, True, seed=0, noise_px=2.0, joint_drop=0.1, translate=20.0, scale=0.05, hazard_tau=tau) | |
| out = smp.batch(np.arange(4), torch.device("cpu")) | |
| assert isinstance(out, tuple) and len(out) == 2 # batch() contract unchanged | |
| obs, fut, ex = smp.batch_extra(np.arange(4), torch.device("cpu")) | |
| assert set(ex) >= {"haz_onset", "haz_risk", "fut_conf"} | |
| on, rk = ex["haz_onset"], ex["haz_risk"] | |
| assert on.shape == rk.shape == (4, 4, P.T_PRED) and on.dtype == fut.dtype | |
| assert (on <= rk).all() and (rk[..., 1:] <= rk[..., :-1]).all() and (on.sum(-1) <= 1).all() | |
| rh = P.PART_NAMES.index("rhand") | |
| assert on[:, rh].sum(-1).eq(1).all() and on[:, rh, 0].eq(0).all() # 12 px/frame x >= 8 frames > 135 px by h = 2 at the latest, never at h = 1 (<= 60 px) | |
| others = [q for q in range(4) if q != rh] | |
| assert on[:, others].eq(0).all() and rk[:, others].eq(1).all() # static parts never exceed | |
| # recomputing from the returned obs / fut reproduces the emitted targets (the sampler uses the SAME cut) | |
| on2, rk2 = LS.hazard_targets(obs[:, -1, :, :2], fut, tau) | |
| assert torch.equal(on, on2) and torch.equal(rk, rk2) | |
| # without hazard_tau nothing is emitted (byte-identical behaviour for every other model) | |
| smp0 = TrainSampler(S45, True, seed=0, noise_px=2.0, joint_drop=0.1, translate=20.0, scale=0.05) | |
| _, _, ex0 = smp0.batch_extra(np.arange(4), torch.device("cpu")) | |
| assert "haz_onset" not in ex0 and "haz_risk" not in ex0 | |
| # same seed -> same obs / fut with and without hazard_tau (the targets consume no randomness) | |
| kw = dict(noise_px=2.0, joint_drop=0.1, translate=20.0, scale=0.05) | |
| torch.manual_seed(0); o_a, f_a, _ = TrainSampler(S45, True, seed=0, **kw).batch_extra(np.arange(4), torch.device("cpu")) | |
| torch.manual_seed(0); o_b, f_b, _ = TrainSampler(S45, True, seed=0, hazard_tau=tau, **kw).batch_extra(np.arange(4), torch.device("cpu")) | |
| assert torch.equal(o_a, o_b) and torch.equal(f_a, f_b) | |
| def test_hazard_loss_trains_the_hazard_head_towards_the_targets(): | |
| """One optimiser step on the BCE lowers it and moves haz_P off the prior for the batch (the loss is wired to the right tensors).""" | |
| torch.manual_seed(0) | |
| m = V2.CWM3Dv2() | |
| tau = np.array([78.5, 32.9, 109.1, 135.2]) | |
| m.set_hazard_prior(tau, np.tile(np.linspace(0.1, 0.4, P.T_PRED), (4, 1))) | |
| frames = torch.cat([_synthetic(2, B=4), _synthetic(0, B=4)]) | |
| obs, fut = split_obs_fut(frames) | |
| on, rk = LS.hazard_targets(obs[:, -1, :, :2], fut, tau) | |
| opt = torch.optim.SGD(m.parameters(), lr=0.5) | |
| losses = [] | |
| for _ in range(5): | |
| out = m(obs) | |
| l = LS.hazard_bce(out["haz_logits"], on, rk) | |
| opt.zero_grad(); l.backward(); opt.step() | |
| losses.append(l.item()) | |
| assert losses[-1] < losses[0] | |
| with torch.no_grad(): | |
| out = m(obs) | |
| assert not torch.allclose(out["haz_P"], m.haz.prior_P[None].expand_as(out["haz_P"]), atol=1e-3) | |
| assert torch.equal(out["pred"], obs[:, -1:, :, :2].expand_as(out["pred"])) # BCE alone never moves the (zero) geometry head / gate | |
| if __name__ == "__main__": | |
| for k, f in list(globals().items()): | |
| if k.startswith("test_") and callable(f): | |
| f() if "tmp_path" not in f.__code__.co_varnames else f(Path("/tmp")) | |
| print("ok", k) | |
Xet Storage Details
- Size:
- 29.9 kB
- Xet hash:
- 5fcd9b956f495dc8e87ad71c15672cae2a48cd33b1e89cc9a5e2c340830dc22b
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.