Buckets:
| """Roadmap supplement B (AIDE hardening) guards: | |
| * `CWM3D.encode(ext_scale=1.0)` is BIT-identical to the plain pass; 0 / 2 change the output (P0-C); | |
| * `gate_override` clamps the traffic gate (0 removes the pathway, 1 / 2 force / over-inject) (P0-C); | |
| * `shift_ext` is a pure, edge-padded re-indexing of the OBSERVED exterior features (never reads beyond frame 4) and | |
| `drop_view` touches one view only (P0-C); | |
| * the late-horizon refinement head is the identity at initialisation and only ever moves h = 4, 5 (P3-B3); | |
| * MotionBERT-kp and ST-GCN-kp start exactly at zero-velocity, reject 10-frame inputs and depend on the observed | |
| block only (P1-D); | |
| * `--motion-oversample` probabilities follow p ∝ 1 + k·s/mean(s) (P2-B1). | |
| python -m pytest aide/tests/test_hardening.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| 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.eval import interventions as I # noqa: E402 | |
| from aide.eval import metrics as M # noqa: E402 | |
| from aide.models import cwm3d, cwm3d_spatial, motionbert_kp, stgcn_kp # noqa: E402 | |
| from aide.train.train_kp import split_obs_fut # noqa: E402 | |
| D = cwm3d.VIS_DIM | |
| def _frames(B=3, seed=0): | |
| g = torch.Generator().manual_seed(seed) | |
| x = torch.rand(B, P.N_SAMPLE, P.J, 3, generator=g) * torch.tensor([P.W, P.H, 1.0]) | |
| v = torch.randn(B, P.N_SAMPLE, len(P.VIEWS), D, generator=g) | |
| return x, v | |
| def _perturb(model, std=0.02, seed=1, only=None): | |
| g = torch.Generator().manual_seed(seed) | |
| with torch.no_grad(): | |
| for n, p in model.named_parameters(): | |
| if only is None or n.startswith(only): | |
| p.add_(std * torch.randn(p.shape, generator=g)) | |
| def test_ext_scale_one_is_bit_identical_and_others_differ(): | |
| x, v = _frames() | |
| obs, _ = split_obs_fut(x) | |
| torch.manual_seed(0) | |
| m = cwm3d.CWM3D(use_visual=True, use_semantic=True, sem_input="attn").eval() | |
| _perturb(m) | |
| with torch.no_grad(): | |
| a = m(obs, vis=v[:, :P.T_OBS]) | |
| b = m(obs, vis=v[:, :P.T_OBS], ext_scale=1.0) | |
| c = m(obs, vis=v[:, :P.T_OBS], ext_scale=None) | |
| z0 = m(obs, vis=v[:, :P.T_OBS], ext_scale=0.0) | |
| z2 = m(obs, vis=v[:, :P.T_OBS], ext_scale=2.0) | |
| for k in ["pred"] + [f"logits_{t}" for t in P.TASKS]: | |
| assert torch.equal(a[k], b[k]) and torch.equal(a[k], c[k]), k # bit-identical, not allclose | |
| assert not torch.allclose(a["pred"], z0["pred"]) and not torch.allclose(a["pred"], z2["pred"]) | |
| assert not torch.allclose(z0["pred"], z2["pred"]) | |
| # ext_scale = 0 removes the exterior tokens entirely: the in-cabin view still matters, the exterior views do not | |
| v_ext_changed = v.clone() | |
| v_ext_changed[:, :, I.EXT] += 3.0 | |
| with torch.no_grad(): | |
| z0b = m(obs, vis=v_ext_changed[:, :P.T_OBS], ext_scale=0.0)["pred"] | |
| assert torch.allclose(z0["pred"], z0b, atol=1e-4) | |
| # recursive rollout accepts the kwargs too | |
| torch.manual_seed(0) | |
| mr = cwm3d.CWM3D(use_visual=True, rollout="recursive").eval() | |
| _perturb(mr) | |
| with torch.no_grad(): | |
| r1 = mr(obs, vis=v[:, :P.T_OBS])["pred"] | |
| r2 = mr(obs, vis=v[:, :P.T_OBS], ext_scale=1.0)["pred"] | |
| r0 = mr(obs, vis=v[:, :P.T_OBS], ext_scale=0.0)["pred"] | |
| assert torch.equal(r1, r2) and not torch.allclose(r1, r0) | |
| def test_gate_override_clamps_the_traffic_gate(): | |
| x, v = _frames() | |
| obs, _ = split_obs_fut(x) | |
| torch.manual_seed(0) | |
| m = cwm3d.CWM3D(use_visual=True, traffic_gate=True).eval() | |
| _perturb(m) | |
| with torch.no_grad(): | |
| a = m(obs, vis=v[:, :P.T_OBS])["pred"] | |
| g0 = m(obs, vis=v[:, :P.T_OBS], gate_override=0.0)["pred"] | |
| g1 = m(obs, vis=v[:, :P.T_OBS], gate_override=1.0)["pred"] | |
| g2 = m(obs, vis=v[:, :P.T_OBS], gate_override=2.0)["pred"] | |
| none = m(obs, vis=v[:, :P.T_OBS], gate_override=None)["pred"] | |
| assert torch.equal(a, none) | |
| assert not torch.allclose(g0, g1) and not torch.allclose(g1, g2) and not torch.allclose(a, g2) | |
| # gate = 0: the cross-attention output is discarded -> the last-frame driver tokens equal the plain trunk tokens | |
| with torch.no_grad(): | |
| p_gate0, _ = m.encode(obs, v[:, :P.T_OBS], gate_override=0.0) | |
| m.gate, gate = None, m.gate | |
| p_nogate, _ = m.encode(obs, v[:, :P.T_OBS]) | |
| m.gate = gate | |
| assert torch.allclose(p_gate0, p_nogate, atol=1e-5) | |
| # a model without a gate rejects the override silently only when None | |
| torch.manual_seed(0) | |
| m2 = cwm3d.CWM3D(use_visual=True).eval() | |
| with torch.no_grad(): | |
| assert torch.equal(m2(obs, vis=v[:, :P.T_OBS])["pred"], m2(obs, vis=v[:, :P.T_OBS], gate_override=None)["pred"]) | |
| def test_shift_and_drop_view_read_observed_frames_only(): | |
| _, v = _frames() | |
| v_obs = v[:, :P.T_OBS] | |
| for name, k in I.SHIFT_STEPS.items(): | |
| s = I.shift_ext(v_obs, k) | |
| for t in range(P.T_OBS): | |
| assert torch.equal(s[:, t, I.EXT], v_obs[:, max(t - k, 0), I.EXT]), (name, t) # edge padding | |
| incar = P.VIEWS.index("incar") | |
| assert torch.equal(s[:, :, incar], v_obs[:, :, incar]) | |
| # NaN-poisoned future never leaks (the function only ever receives the observed block) | |
| vp = v.clone() | |
| vp[:, P.T_OBS:] = float("nan") | |
| assert torch.equal(I.shift_ext(vp[:, :P.T_OBS], k), s) | |
| try: | |
| I.shift_ext(v, k) # 10 frames must be rejected | |
| raise RuntimeError("shift_ext accepted 10 frames") | |
| except AssertionError: | |
| pass | |
| for view in ("front", "left", "right"): | |
| d = I.drop_view(v_obs, view) | |
| vi = P.VIEWS.index(view) | |
| assert (d[:, :, vi] == 0).all() | |
| others = [i for i in range(len(P.VIEWS)) if i != vi] | |
| assert torch.equal(d[:, :, others], v_obs[:, :, others]) | |
| z = I.zero_ext(v_obs) | |
| assert (z[:, :, I.EXT] == 0).all() and torch.equal(z[:, :, P.VIEWS.index("incar")], v_obs[:, :, P.VIEWS.index("incar")]) | |
| def test_refine_late_is_identity_at_init_and_moves_only_late_horizons(): | |
| x, v = _frames() | |
| obs, _ = split_obs_fut(x) | |
| torch.manual_seed(0) | |
| base = cwm3d.CWM3D(use_visual=True).eval() | |
| _perturb(base) | |
| torch.manual_seed(0) | |
| ref = cwm3d.CWM3D(use_visual=True, refine_late=True).eval() | |
| missing, unexpected = ref.load_state_dict(base.state_dict(), strict=False) | |
| assert not unexpected and all(k.startswith("refine.") for k in missing) | |
| with torch.no_grad(): | |
| pb = base(obs, vis=v[:, :P.T_OBS])["pred"] | |
| pr = ref(obs, vis=v[:, :P.T_OBS]) | |
| assert torch.equal(pb, pr["pred"]) # identity at init (zero-init output layer) | |
| assert torch.equal(pr["refine_delta_px"], torch.zeros_like(pr["refine_delta_px"])) | |
| # zero-init identity with ZV still holds for a fresh model | |
| torch.manual_seed(0) | |
| fresh = cwm3d.CWM3D(refine_late=True).eval() | |
| with torch.no_grad(): | |
| o = fresh(obs)["pred"] | |
| assert (o - obs[:, -1:, :, :2]).abs().max() <= 1e-6 | |
| # perturbing ONLY the refinement head leaves h <= 3 bit-identical and changes h = 4, 5 | |
| _perturb(ref, std=0.05, only="refine.") | |
| with torch.no_grad(): | |
| pr2 = ref(obs, vis=v[:, :P.T_OBS])["pred"] | |
| assert torch.equal(pr2[:, :3], pb[:, :3]) and not torch.allclose(pr2[:, 3:], pb[:, 3:]) | |
| # the early horizons feed the head detached: no gradient flows from the refinement loss into h <= 3 via that path | |
| ref.train() | |
| out = ref(obs, vis=v[:, :P.T_OBS]) | |
| loss = out["pred"][:, 3:].pow(2).mean() | |
| loss.backward() | |
| assert all(torch.isfinite(p.grad).all() for p in ref.refine.parameters() if p.grad is not None) | |
| def test_kp_baselines_zero_init_and_no_future_leak(): | |
| x, _ = _frames() | |
| obs, _ = split_obs_fut(x) | |
| zv = obs[:, -1:, :, :2].expand(-1, P.T_PRED, -1, -1) | |
| poison = x.clone() | |
| poison[:, P.T_OBS:] = float("nan") | |
| obs_p, _ = split_obs_fut(poison) | |
| for build, name in ((motionbert_kp.build, "motionbert_kp"), (motionbert_kp.build, "motionbert_kp_2b"), (stgcn_kp.build, "stgcn_kp")): | |
| torch.manual_seed(0) | |
| m = build(name).eval() | |
| with torch.no_grad(): | |
| o = m(obs) | |
| assert o.shape == zv.shape and torch.equal(o, zv), name # exactly zero-velocity at init | |
| _perturb(m) | |
| with torch.no_grad(): | |
| a, b = m(obs), m(obs_p) | |
| assert torch.isfinite(a).all() and torch.equal(a, b), name | |
| try: | |
| with torch.no_grad(): | |
| m(poison) | |
| raise RuntimeError(f"{name} accepted 10 frames") | |
| except AssertionError: | |
| pass | |
| shifted = obs.clone() | |
| shifted[:, -1, :, :2] += 30.0 | |
| with torch.no_grad(): | |
| c = m(shifted) | |
| assert not torch.allclose(a, c), name | |
| n = sum(p.numel() for p in m.parameters()) | |
| assert 2e5 < n < 5e6, (name, n) | |
| def test_cwm3d_spatial_hybrid_zero_init_and_no_future_leak(): | |
| x, _ = _frames() | |
| obs, _ = split_obs_fut(x) | |
| torch.manual_seed(0) | |
| m = cwm3d_spatial.build("cwm3d_kp_spatial").eval() | |
| with torch.no_grad(): | |
| o = m(obs) | |
| assert (o["pred"] - obs[:, -1:, :, :2]).abs().max() <= 1e-6 # anchor mixture + zero-init head inherited | |
| assert torch.allclose(o["anchor_w"], torch.full_like(o["anchor_w"], 1 / 3)) | |
| _perturb(m) | |
| poison = x.clone() | |
| poison[:, P.T_OBS:] = float("nan") | |
| obs_p, _ = split_obs_fut(poison) | |
| with torch.no_grad(): | |
| a, b = m(obs)["pred"], m(obs_p)["pred"] | |
| assert torch.isfinite(a).all() and torch.equal(a, b) | |
| try: | |
| with torch.no_grad(): | |
| m(poison) | |
| raise RuntimeError("hybrid accepted 10 frames") | |
| except AssertionError: | |
| pass | |
| # the spatial block acts within a frame: joints of frame 0 do not see frame 1 (tokens of frame 0 unchanged when frame 1 moves) | |
| x2 = obs.clone() | |
| x2[:, 1, :, :2] += 25.0 | |
| with torch.no_grad(): | |
| t1, t2 = m.tok(obs), m.tok(x2) | |
| assert torch.allclose(t1[:, 0], t2[:, 0], atol=1e-5) and not torch.allclose(t1[:, 1], t2[:, 1]) | |
| n = sum(p.numel() for p in m.parameters()) | |
| assert 5e6 < n < 1e7, n | |
| def test_motion_oversample_probabilities(): | |
| from aide.train import train_kp as T | |
| rng = np.random.default_rng(0) | |
| S45 = np.zeros((10, P.N_FRAMES_STD, P.J, 3), np.float32) | |
| for i in range(10): # clip i moves i px per frame (x only) | |
| S45[i, :, :, 0] = np.arange(P.N_FRAMES_STD)[:, None] * i | |
| s = M.motion_score(S45[:, T.PROTO_IDX][:, P.T_OBS:]) | |
| assert np.allclose(s, np.arange(10) * 5.0) # protocol future gaps are 5 frames | |
| k = 2.0 | |
| w = 1.0 + k * s / s.mean() | |
| p = w / w.sum() | |
| assert abs(p.sum() - 1) < 1e-12 and p[0] < p[-1] and abs(p[-1] / p[0] - (1 + k * s[-1] / s.mean())) < 1e-9 | |
| draw = rng.choice(10, size=20000, replace=True, p=p) | |
| freq = np.bincount(draw, minlength=10) / 20000 | |
| assert np.abs(freq - p).max() < 0.02 | |
| def test_dn_mpjpe_hm_matches_definition(): | |
| rng = np.random.default_rng(3) | |
| g = np.zeros((40, P.T_PRED, P.J, 3)) | |
| g[..., 0] = rng.uniform(0, P.W, (40, P.T_PRED, P.J)) | |
| g[..., 1] = rng.uniform(0, P.H, (40, P.T_PRED, P.J)) | |
| g[..., 2] = 1.0 | |
| pred = g[..., :2] + rng.normal(0, 20, (40, P.T_PRED, P.J, 2)) | |
| out, meta = M.evaluate(pred, g, n_boot=100) | |
| assert abs(out["dn_mpjpe_hm"]["value"] - 100.0 * out["mpjpe_hm"]["value"] / P.DIAG) < 1e-9 | |
| assert abs(out["dn_mpjpe"]["value"] - 100.0 * out["mpjpe_all"]["value"] / P.DIAG) < 1e-9 | |
| assert out["dn_mpjpe_hm"]["n"] == meta["hm_n"] == P.hm_top(40) | |
| 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:
- 11.9 kB
- Xet hash:
- 9d9ca57da544b8efc607726043961b13e801c81f49b6fa72343065938d84821a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.