Buckets:
| """Formulation ablation guards (docs/decisions.md 2026-09-15 "PRE-REGISTRATION ... Formulation ablation", EXPLANATORY, not for adoption). | |
| 1. The DEFAULT path (`--formulation residual`, the only value every existing run used) is byte-identical to the pre-change code: | |
| state-dict keys / shapes, sha256 over every tensor at seed 0, and the forward outputs at init and after a fixed perturbation equal the | |
| tensors saved from the UNMODIFIED motionbert_kp.py / cwm3d.py / heads.py (`aide/tests/fixtures/formulation_abs_prechange.pt`; its | |
| `source_sha256` records the pre-change files). `train_kp.build_model` with a default Namespace reproduces the same weights. | |
| 2. The `abs` path is NOT zero-velocity at init (default non-zero read-out init, no residual to the last pose, no anchors), still emits | |
| log_b (K's NLL term stays live), keeps the same state-dict keys, and reads out pred = out ⊙ (W, H) px independently of the last pose | |
| (checked with controlled read-out weights). | |
| 3. No-future leak holds for `abs`: NaN-poisoning frames >= T_OBS never changes the output, a 10-frame input is rejected, and the output | |
| depends on the observed block. | |
| 4. Every other model rejects `--formulation abs`; K's four loss terms (euc / nll / bone / jerk) are all active on the abs output dict and | |
| MB's tensor output gives euc only — exactly what the pre-registration says is kept. | |
| python -m pytest aide/tests/test_formulation_abs.py -q | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import math | |
| 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 pytest # 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, heads, losses, motionbert_kp # noqa: E402 | |
| from aide.train.train_kp import build_model, get_args, split_obs_fut # noqa: E402 | |
| REF = Path(__file__).resolve().parent / "fixtures" / "formulation_abs_prechange.pt" | |
| def _sha(m): | |
| h = hashlib.sha256() | |
| for k, v in sorted(m.state_dict().items()): | |
| h.update(k.encode()); h.update(str(tuple(v.shape)).encode()); h.update(v.detach().cpu().contiguous().numpy().tobytes()) | |
| return h.hexdigest() | |
| def _perturb(m, std=0.02, seed=1): | |
| g = torch.Generator().manual_seed(seed) | |
| with torch.no_grad(): | |
| for p in m.parameters(): | |
| p.add_(std * torch.randn(p.shape, generator=g)) | |
| def _pred(o): | |
| return o["pred"] if isinstance(o, dict) else o | |
| def _build(name, **kw): | |
| return (motionbert_kp if name.startswith("motionbert") else cwm3d).build(name, **kw) | |
| def _frames(n=4, 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=8): | |
| 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"] == "val"][:n]).float() | |
| def _abs_models(): | |
| torch.manual_seed(0) | |
| mb = motionbert_kp.build("motionbert_kp", formulation="abs").eval() | |
| torch.manual_seed(0) | |
| k = cwm3d.build("cwm3d", formulation="abs").eval() | |
| return mb, k | |
| # ----------------------------------------------------------------------------- 1. default path byte-identical | |
| def test_default_path_byte_identical_to_prechange_code(): | |
| ref = torch.load(REF, weights_only=False) | |
| obs, _ = split_obs_fut(ref["frames"]) | |
| assert set(ref["models"]) == {"motionbert_kp", "motionbert_kp_2b", "cwm3d", "cwm3d_noanchor", "cwm3d_nounc"} | |
| for name, rec in ref["models"].items(): | |
| for kw in ({}, {"formulation": "residual"}): # implicit default and the explicit default value | |
| torch.manual_seed(0) | |
| m = _build(name, **kw).eval() | |
| assert [(k, tuple(v.shape)) for k, v in m.state_dict().items()] == [(k, tuple(s)) for k, s in rec["state_keys"]], name | |
| assert _sha(m) == rec["sha_init"], (name, kw) | |
| with torch.no_grad(): | |
| o = m(obs) | |
| assert torch.equal(_pred(o), rec["pred_init"]), (name, kw) | |
| _perturb(m) | |
| assert _sha(m) == rec["sha_perturbed"], (name, kw) | |
| with torch.no_grad(): | |
| o = m(obs) | |
| assert torch.equal(_pred(o), rec["pred_perturbed"]), (name, kw) | |
| if isinstance(o, dict): | |
| assert sorted(k for k, v in o.items() if torch.is_tensor(v)) == rec["extra_keys"], name | |
| if rec["log_b_perturbed"] is not None: | |
| assert torch.equal(o["log_b"], rec["log_b_perturbed"]), name | |
| def test_train_kp_default_namespace_reproduces_prechange_weights(): | |
| ref = torch.load(REF, weights_only=False) | |
| for model in ("motionbert_kp", "cwm3d"): | |
| a = get_args(["--model", model, "--no-lock"]) | |
| assert a.formulation == "residual" | |
| torch.manual_seed(0) | |
| m, _ = build_model(model, a) | |
| assert _sha(m) == ref["models"][model]["sha_init"], model | |
| assert getattr(m, "formulation", "residual") == "residual" | |
| assert isinstance(cwm3d.CWM3D().mix, heads.AnchorMixture) and cwm3d.CWM3D().mix.enabled | |
| # ----------------------------------------------------------------------------- 2. abs path | |
| def test_abs_is_not_zero_velocity_at_init_and_keeps_state_keys(): | |
| obs, _ = split_obs_fut(_frames()) | |
| zv = obs[:, -1:, :, :2].expand(-1, P.T_PRED, -1, -1) | |
| mb, k = _abs_models() | |
| torch.manual_seed(0) | |
| mb_res = motionbert_kp.build("motionbert_kp").eval() | |
| torch.manual_seed(0) | |
| k_res = cwm3d.build("cwm3d").eval() | |
| for m, r in ((mb, mb_res), (k, k_res)): | |
| assert list(m.state_dict()) == list(r.state_dict()) # the (W, H) buffer is non-persistent | |
| with torch.no_grad(): | |
| p = _pred(m(obs)) | |
| assert p.shape == zv.shape and torch.isfinite(p).all() | |
| assert (p - zv).abs().max().item() > 1.0 # px: NOT the zero-velocity start | |
| assert mb.formulation == "abs" and not torch.equal(mb.out.weight, torch.zeros_like(mb.out.weight)) | |
| assert k.formulation == "abs" and isinstance(k.mix, heads.AbsoluteReadout) and k.mix.enabled is False | |
| for mlp in k.head.mlp.values(): | |
| assert not torch.equal(mlp[-1].weight, torch.zeros_like(mlp[-1].weight)) | |
| with torch.no_grad(): | |
| o = k(obs) | |
| assert "anchor_w" not in o and "anchor_gate" not in o | |
| assert o["log_b"] is not None and torch.isfinite(o["log_b"]).all() | |
| for model in ("motionbert_kp", "cwm3d"): # trainer path | |
| a = get_args(["--model", model, "--formulation", "abs", "--no-lock"]) | |
| torch.manual_seed(0) | |
| m, _ = build_model(model, a) | |
| assert m.formulation == "abs" | |
| with torch.no_grad(): | |
| p = _pred(m.eval()(obs)) | |
| assert (p - zv).abs().max().item() > 1.0 | |
| for name in ("cwm3d_recursive", "cwm3d_nounc"): # other cwm3d variants accept the flag too | |
| torch.manual_seed(0) | |
| with torch.no_grad(): | |
| p = cwm3d.build(name, formulation="abs").eval()(obs)["pred"] | |
| assert p.shape == zv.shape and torch.isfinite(p).all() | |
| def test_abs_readout_is_absolute_normalised_coordinates_times_WH(): | |
| obs, _ = split_obs_fut(_frames()) | |
| shifted = obs + torch.tensor([30.0, -20.0, 0.0]) # every observed pose moved (incl. the last one) | |
| mb, k = _abs_models() | |
| target = torch.tensor([0.25 * P.W, 0.75 * P.H]) | |
| with torch.no_grad(): | |
| mb.out.weight.zero_() | |
| mb.out.bias.copy_(torch.tensor([0.25, 0.75])) | |
| p, p2 = mb(obs), mb(shifted) | |
| assert torch.allclose(p, target.expand_as(p), atol=1e-3) | |
| assert torch.allclose(p2, target.expand_as(p2), atol=1e-3) # no residual to the last pose | |
| with torch.no_grad(): | |
| for name, mlp in k.head.mlp.items(): | |
| lin = mlp[-1] | |
| lin.weight.zero_() | |
| b = lin.bias.view(P.T_PRED, len(P.PARTS[name]), heads.HorizonHead.N_C) | |
| b.zero_(); b[..., 0] = 0.25; b[..., 1] = 0.75 | |
| o, o2 = k(obs), k(shifted) | |
| assert torch.allclose(o["pred"], target.expand_as(o["pred"]), atol=1e-3) | |
| assert torch.allclose(o2["pred"], o["pred"], atol=1e-3) | |
| assert torch.allclose(o["log_b"], torch.full_like(o["log_b"], math.log(30.0))) | |
| # the residual model with the same controlled bias moves WITH the last pose (contrast) | |
| torch.manual_seed(0) | |
| r = motionbert_kp.build("motionbert_kp").eval() | |
| with torch.no_grad(): | |
| r.out.bias.copy_(torch.tensor([0.25, 0.75])) | |
| q, q2 = r(obs), r(shifted) | |
| assert torch.allclose(q2 - q, torch.tensor([30.0, -20.0]).expand_as(q), atol=1e-3) | |
| # ----------------------------------------------------------------------------- 3. no future leak | |
| def test_abs_models_never_read_future_frames(): | |
| frames = _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 m in _abs_models(): | |
| _perturb(m) | |
| with torch.no_grad(): | |
| a, b = _pred(m(obs_c)), _pred(m(obs_p)) | |
| assert torch.isfinite(a).all() and torch.equal(a, b) | |
| try: | |
| with torch.no_grad(): | |
| m(poison) | |
| raise RuntimeError(f"{m.name} accepted {P.N_SAMPLE} frames") | |
| except AssertionError: | |
| pass | |
| shifted = obs_c.clone() | |
| shifted[:, -1, :, :2] += 30.0 | |
| with torch.no_grad(): | |
| c = _pred(m(shifted)) | |
| assert not torch.allclose(a, c) | |
| # ----------------------------------------------------------------------------- 4. flag scope and loss terms | |
| def test_other_models_reject_abs(): | |
| for model in ("simlpe_kp", "stgcn_kp", "cwm3d_kp_spatial", "driverwm_kp", "cwm3d_v2"): | |
| a = get_args(["--model", model, "--formulation", "abs", "--no-lock"]) | |
| with pytest.raises(ValueError): | |
| build_model(model, a) | |
| def test_abs_loss_terms_as_preregistered(): | |
| obs, fut = split_obs_fut(_frames()) | |
| mb, k = _abs_models() | |
| with torch.no_grad(): | |
| ok, omb = k(obs), mb(obs) | |
| total, terms = losses.compute_losses(ok, fut, losses.DEFAULT_W, "euc") | |
| assert set(terms) == {"euc", "nll", "bone", "jerk"} and torch.isfinite(total) # K: all four kept | |
| total, terms = losses.compute_losses(omb, fut, losses.DEFAULT_W, "euc") | |
| assert set(terms) == {"euc"} and torch.isfinite(total) # MB: tensor output -> euc only | |
| 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 kB
- Xet hash:
- 4f37e3cb382734cde090db3fb44217b4247c05f5615ec093a20bc4206e70c44d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.