Buckets:
| """CWM3-D unit tests: zero-init == zero-velocity, causal mask, recursive rollout, ablation flags, | |
| no-future leak, loss terms, DCT basis. | |
| python -m pytest aide/tests/test_cwm3d.py -q | |
| """ | |
| from __future__ import annotations | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.models import cwm3d, heads, losses # noqa: E402 | |
| def _obs(B=3, T=P.T_OBS, seed=0): | |
| g = torch.Generator().manual_seed(seed) | |
| x = torch.rand(B, T, P.J, 3, generator=g) * torch.tensor([P.W, P.H, 1.0]) | |
| return x | |
| def _perturb(model, std=0.02, seed=1): | |
| g = torch.Generator().manual_seed(seed) | |
| with torch.no_grad(): | |
| for p in model.parameters(): | |
| p.add_(std * torch.randn(p.shape, generator=g)) | |
| def test_zero_init_equals_zero_velocity_all_variants(): | |
| x = _obs() | |
| zv = x[:, -1:, :, :2].expand(-1, P.T_PRED, -1, -1) | |
| for name in ("cwm3d", "cwm3d_noanchor", "cwm3d_nounc", "cwm3d_recursive", "cwm3d_dct", "cwm3d_vis"): | |
| torch.manual_seed(0) | |
| m = cwm3d.build(name).eval() | |
| with torch.no_grad(): | |
| o = m(x) | |
| assert torch.equal(o["pred"], zv) or torch.allclose(o["pred"], zv, atol=1e-4), name | |
| assert o["pred"].shape == (3, P.T_PRED, P.J, 2) | |
| if name != "cwm3d_nounc": | |
| assert o["log_b"].shape == (3, P.T_PRED, P.J) and torch.allclose(o["log_b"], torch.full_like(o["log_b"], np.log(30.0))) | |
| if name not in ("cwm3d_noanchor",): | |
| assert torch.allclose(o["anchor_w"], torch.full_like(o["anchor_w"], 1 / 3)) | |
| def test_causal_mask_blocks_future_frames_only(): | |
| m = heads.causal_part_mask(P.T_OBS, 4) | |
| assert m.shape == (20, 20) | |
| for i in range(20): | |
| for j in range(20): | |
| assert bool(m[i, j]) == (j // 4 > i // 4) | |
| m2 = heads.causal_part_mask(P.T_OBS, 4, 4) | |
| assert m2.shape == (40, 40) and not m2[7, 0] and m2[0, 8] | |
| def test_encoder_is_causal_across_frames(): | |
| """Changing frame t of the input must not change part tokens of frames < t.""" | |
| torch.manual_seed(0) | |
| m = cwm3d.CWM3D().eval() | |
| _perturb(m) | |
| x = _obs() | |
| with torch.no_grad(): | |
| p1, _ = m.encode(x) | |
| x2 = x.clone() | |
| x2[:, 3:] += 50.0 # perturb frames 3,4 (and, through Δ-features, the anchor) | |
| # Δ features are relative to the LAST frame, so perturbing the last frame changes all tokens; | |
| # perturb frames 1..3 only and keep frame 4 fixed to test causality of the attention itself. | |
| x3 = x.clone() | |
| x3[:, 2:4] += 50.0 | |
| p3, _ = m.encode(x3) | |
| assert torch.allclose(p1[:, :2], p3[:, :2], atol=1e-4) # frames 0,1 unaffected by frames 2,3 | |
| assert not torch.allclose(p1[:, 2:], p3[:, 2:]) | |
| def test_no_future_leak_and_finite(): | |
| torch.manual_seed(0) | |
| x10 = torch.rand(4, P.N_SAMPLE, P.J, 3) * torch.tensor([P.W, P.H, 1.0]) | |
| xp = x10.clone() | |
| xp[:, P.T_OBS:] = float("nan") | |
| for name in ("cwm3d", "cwm3d_recursive"): | |
| m = cwm3d.build(name).eval() | |
| _perturb(m) | |
| with torch.no_grad(): | |
| a = m(x10[:, :P.T_OBS])["pred"] | |
| b = m(xp[:, :P.T_OBS])["pred"] | |
| assert torch.isfinite(a).all() and torch.equal(a, b), name | |
| try: | |
| m(xp) | |
| raise RuntimeError("accepted 10 frames") | |
| except AssertionError: | |
| pass | |
| def test_recursive_differs_from_chunked_after_perturbation_but_same_at_init(): | |
| x = _obs() | |
| torch.manual_seed(0) | |
| mc = cwm3d.CWM3D(rollout="chunked").eval() | |
| mr = cwm3d.CWM3D(rollout="recursive").eval() | |
| mr.load_state_dict(mc.state_dict()) | |
| with torch.no_grad(): | |
| assert torch.allclose(mc(x)["pred"], mr(x)["pred"]) | |
| _perturb(mc); mr.load_state_dict(mc.state_dict()) | |
| with torch.no_grad(): | |
| pc, pr = mc(x)["pred"], mr(x)["pred"] | |
| assert torch.allclose(pc[:, 0], pr[:, 0], atol=1e-3) # first horizon identical by construction | |
| assert not torch.allclose(pc[:, 1:], pr[:, 1:]) | |
| def test_anchor_mixture_gate_one_is_convex_mixture(): | |
| x = _obs() | |
| anc = heads.anchors(x) | |
| B = x.shape[0] | |
| ho = torch.zeros(B, P.T_PRED, P.J, heads.HorizonHead.N_C) | |
| ho[..., heads.HorizonHead.C_GATE] = 1.0 | |
| ho[..., 4] = 50.0 # logit for mean_obs -> weight ~1 | |
| mix = heads.AnchorMixture() | |
| pred, _ = mix(ho, anc) | |
| assert torch.allclose(pred, anc[:, 1:2].expand_as(pred), atol=1e-3) | |
| ho[..., 4] = 0.0 | |
| ho[..., 5] = 50.0 # anti-velocity anchor | |
| pred, _ = mix(ho, anc) | |
| assert torch.allclose(pred, anc[:, 2:3].expand_as(pred), atol=1e-3) | |
| exp = x[:, -1, :, :2] - 0.1 * (x[:, -1, :, :2] - x[:, -2, :, :2]) | |
| assert torch.allclose(anc[:, 2], exp) | |
| def test_dct_basis_orthonormal(): | |
| C = heads.dct_matrix(5) | |
| assert torch.allclose(C @ C.t(), torch.eye(5), atol=1e-5) | |
| def test_losses_terms(): | |
| x = _obs() | |
| fut = torch.rand(3, P.T_PRED, P.J, 2) * torch.tensor([P.W, P.H]) | |
| out = {"pred": fut.clone(), "log_b": torch.zeros(3, P.T_PRED, P.J), "last": fut[:, :1]} | |
| tot, terms = losses.compute_losses(out, fut) | |
| assert abs(float(terms["euc"])) < 1e-2 and abs(float(terms["bone"])) < 1e-6 | |
| assert float(terms["nll"]) == 0.0 # |d|/1 + 2*0 | |
| # jerk of a constant trajectory (last == all preds) is zero | |
| out2 = {"pred": fut[:, :1].expand(-1, P.T_PRED, -1, -1), "last": fut[:, :1]} | |
| _, t2 = losses.compute_losses(out2, fut) | |
| assert float(t2["jerk"]) < 1e-3 # float32 cancellation noise only | |
| # plain tensor path | |
| tot3, t3 = losses.compute_losses(fut + 3.0, fut, plain_loss="l1") | |
| assert abs(float(tot3) - 3.0) < 1e-6 | |
| def test_param_count_in_plan_range(): | |
| n = sum(p.numel() for p in cwm3d.CWM3D().parameters()) | |
| assert 3_000_000 < n < 8_000_000, n # plan: ≈5.5 M | |
| if __name__ == "__main__": | |
| for k, f in list(globals().items()): | |
| if k.startswith("test_") and callable(f): | |
| f() | |
| print("ok", k) | |
Xet Storage Details
- Size:
- 6.11 kB
- Xet hash:
- 9bca74c65a5f202660b34bc803728aeae57e498e95b6e6d4f4f6f064022f19ea
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.