Buckets:
| """Pre-registered HM h=5 arms (docs/decisions.md 2026-09-17T07:00:22Z, part (A)) — additive-flag guards. | |
| 1. DEFAULTS ARE THE OLD CODE: `--hand-weight 1 --horizon-gain none` (the defaults) give a cwm3d model whose state-dict keys, sha256 over every | |
| tensor at seed 0 and forward outputs (init + fixed perturbation) equal the tensors saved from the UNMODIFIED code | |
| (`aide/tests/fixtures/formulation_abs_prechange.pt`); compute_losses without / with jw=None is the same code path; joint_weights(1.0) is None. | |
| 2. Hand weight (arms a1 / a2): the 42 hand keypoints (HALPE 94-135) get w, the other 94 joints 1, NOT renormalised; the Euclidean term equals | |
| (Σ_other e + w Σ_hands e) / (B·H·J); NLL / bone / jerk terms are bit-identical to the default; combines multiplicatively with --h-weights. | |
| 3. Horizon gain (arms b / c): g = 1 + δ, δ zero-init -> exact zero-velocity at init, same predictions as the unchanged model for the same trunk / | |
| head weights; only `hgain.delta` is added to the state dict; δ_h changes horizon h only ("learned") / hand joints of horizon h only ("parts"); | |
| the gain has a gradient and is tracked by the EMA; `--init-parts trunk` keeps δ at 0 (reported as kept); non-cwm3d models and the abs | |
| formulation reject the flag; `--horizon-gain-parts` == `--horizon-gain parts`. | |
| python -m pytest aide/tests/test_hm_h5_arms.py -q | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| 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 # noqa: E402 | |
| from aide.models import losses as LS # noqa: E402 | |
| from aide.train.train_kp import EMA, build_model, get_args, load_init_ckpt, split_obs_fut # noqa: E402 | |
| REF = Path(__file__).resolve().parent / "fixtures" / "formulation_abs_prechange.pt" | |
| HANDS = list(heads.HAND_JOINTS) | |
| OTHER = [j for j in range(P.J) if j not in heads.HAND_JOINTS] | |
| 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 _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 _perturbed_base(seed=0): | |
| torch.manual_seed(seed) | |
| m = cwm3d.build("cwm3d").eval() | |
| _perturb(m) | |
| return m | |
| # ----------------------------------------------------------------------------- 1. defaults == old code | |
| def test_hand_joints_are_the_42_hand_keypoints(): | |
| assert len(HANDS) == 42 and HANDS == list(range(94, 136)) and len(OTHER) == 94 | |
| assert HANDS == list(P.PARTS["lhand"]) + list(P.PARTS["rhand"]) | |
| def test_default_flags_and_default_model_byte_identical_to_prechange_code(): | |
| a = get_args(["--model", "cwm3d", "--no-lock"]) | |
| assert a.hand_weight == 1.0 and a.horizon_gain == "none" and a.horizon_gain_parts is False | |
| assert LS.joint_weights(a.hand_weight) is None and LS.joint_weights(None) is None and LS.joint_weights(1) is None | |
| ref = torch.load(REF, weights_only=False) | |
| rec = ref["models"]["cwm3d"] | |
| obs, _ = split_obs_fut(ref["frames"]) | |
| for builder in (lambda: build_model("cwm3d", a)[0], lambda: cwm3d.build("cwm3d"), lambda: cwm3d.build("cwm3d", horizon_gain=None), | |
| lambda: cwm3d.build("cwm3d", horizon_gain="none")): | |
| torch.manual_seed(0) | |
| m = builder().eval() | |
| assert m.hgain is None | |
| assert [(k, tuple(v.shape)) for k, v in m.state_dict().items()] == [(k, tuple(s)) for k, s in rec["state_keys"]] | |
| assert _sha(m) == rec["sha_init"] | |
| with torch.no_grad(): | |
| o = m(obs) | |
| assert torch.equal(o["pred"], rec["pred_init"]) | |
| _perturb(m) | |
| assert _sha(m) == rec["sha_perturbed"] | |
| with torch.no_grad(): | |
| o = m(obs) | |
| assert torch.equal(o["pred"], rec["pred_perturbed"]) | |
| assert torch.equal(o["log_b"], rec["log_b_perturbed"]) | |
| assert sorted(k for k, v in o.items() if torch.is_tensor(v)) == rec["extra_keys"] | |
| def test_compute_losses_default_path_unchanged(): | |
| m = _perturbed_base() | |
| x = _frames() | |
| obs, fut = split_obs_fut(x) | |
| with torch.no_grad(): | |
| out = m(obs) | |
| w = {"euc": 1.0, "nll": 0.1, "bone": 0.05, "jerk": 0.01} | |
| t0, terms0 = LS.compute_losses(out, fut, w) | |
| t1, terms1 = LS.compute_losses(out, fut, w, "euc", sw=None, jw=None) | |
| assert torch.equal(t0, t1) and all(torch.equal(terms0[k], terms1[k]) for k in terms0) and set(terms0) == {"euc", "nll", "bone", "jerk"} | |
| t2, terms2 = LS.compute_losses(out, fut, w, jw=torch.ones(P.J)) # all-ones joint weights == unweighted (float rounding only) | |
| assert torch.allclose(t2, t0, rtol=1e-6, atol=1e-5) and torch.allclose(terms2["euc"], terms0["euc"], rtol=1e-6, atol=1e-5) | |
| # tensor-output models: jw=None -> plain(); jw ones == plain | |
| t3, _ = LS.compute_losses(out["pred"], fut, w, "euc") | |
| t4, _ = LS.compute_losses(out["pred"], fut, w, "euc", jw=torch.ones(P.J)) | |
| assert torch.equal(t3, terms0["euc"]) and torch.allclose(t4, t3, rtol=1e-6, atol=1e-5) | |
| # ----------------------------------------------------------------------------- 2. hand weight | |
| def test_hand_weight_weights_the_42_hand_joints_in_the_euclidean_term_only(hw): | |
| jw = LS.joint_weights(hw) | |
| assert jw.shape == (P.J,) and torch.all(jw[HANDS] == hw) and torch.all(jw[OTHER] == 1.0) | |
| m = _perturbed_base() | |
| x = _frames() | |
| obs, fut = split_obs_fut(x) | |
| with torch.no_grad(): | |
| out = m(obs) | |
| w = {"euc": 1.0, "nll": 0.1, "bone": 0.05, "jerk": 0.01} | |
| _, terms0 = LS.compute_losses(out, fut, w) | |
| total, terms = LS.compute_losses(out, fut, w, jw=jw) | |
| e = torch.sqrt(((out["pred"] - fut) ** 2).sum(-1) + 1e-6) # [B,H,J] | |
| expected = (e[..., OTHER].sum() + hw * e[..., HANDS].sum()) / e.numel() # other joints 1, hands w, denominator unchanged | |
| assert torch.allclose(terms["euc"], expected, rtol=1e-6, atol=1e-5) | |
| assert terms["euc"] > terms0["euc"] # not renormalised -> larger than the plain mean | |
| for k in ("nll", "bone", "jerk"): | |
| assert torch.equal(terms[k], terms0[k]), k | |
| assert torch.allclose(total, terms["euc"] + 0.1 * terms["nll"] + 0.05 * terms["bone"] + 0.01 * terms["jerk"], rtol=1e-6) | |
| # combines multiplicatively with the (existing) horizon sample weights | |
| sw = LS.sample_weights(fut, [1, 1, 1, 2, 3]) | |
| _, terms_sw = LS.compute_losses(out, fut, w, sw=sw, jw=jw) | |
| assert torch.allclose(terms_sw["euc"], (e * sw * jw[None, None, :]).mean(), rtol=1e-6, atol=1e-5) | |
| with pytest.raises(AssertionError): | |
| LS.compute_losses(out, fut, w, "l1", jw=jw) | |
| def test_hand_weight_gradient_hits_hand_channels_harder(): | |
| """the hand-joint gradient of the Euclidean term scales with w (sanity that the weight reaches the loss).""" | |
| x = _frames() | |
| obs, fut = split_obs_fut(x) | |
| pred = (obs[:, -1:, :, :2] + 5.0).expand(-1, P.T_PRED, -1, -1).clone().requires_grad_(True) | |
| LS.euclidean(pred, fut).backward() | |
| g1 = pred.grad.clone(); pred.grad = None | |
| LS.euclidean(pred, fut, w=LS.joint_weights(4.0)[None, None, :]).backward() | |
| g4 = pred.grad | |
| assert torch.allclose(g4[:, :, HANDS], 4.0 * g1[:, :, HANDS], rtol=1e-5, atol=1e-8) | |
| assert torch.allclose(g4[:, :, OTHER], g1[:, :, OTHER], rtol=1e-5, atol=1e-8) | |
| # ----------------------------------------------------------------------------- 3. horizon gain | |
| def test_horizon_gain_is_identity_at_init_and_keeps_zero_velocity(mode): | |
| x = _frames() | |
| obs, _ = split_obs_fut(x) | |
| torch.manual_seed(0) | |
| base = cwm3d.build("cwm3d").eval() | |
| torch.manual_seed(0) | |
| g = cwm3d.build("cwm3d", horizon_gain=mode).eval() | |
| assert g.hgain is not None and g.hgain.mode == mode | |
| assert set(g.state_dict()) == set(base.state_dict()) | {"hgain.delta"} | |
| assert tuple(g.hgain.delta.shape) == ((P.T_PRED,) if mode == "learned" else (P.T_PRED, 2)) and torch.all(g.hgain.delta == 0) | |
| assert tuple(g.hgain.gain().shape) == (P.T_PRED, P.J) and torch.all(g.hgain.gain() == 1.0) | |
| zv = obs[:, -1:, :, :2].expand(-1, P.T_PRED, -1, -1) | |
| with torch.no_grad(): | |
| assert torch.equal(g(obs)["pred"], zv) # exact zero-velocity start (as the adopted model) | |
| _perturb(base) | |
| g.load_state_dict(base.state_dict(), strict=False) # same trunk + head weights, δ = 0 | |
| with torch.no_grad(): | |
| ob, og = base(obs), g(obs) | |
| assert torch.allclose(og["pred"], ob["pred"], atol=1e-3) # last + 1·(pred − last): float rounding only (px) | |
| assert torch.equal(og["log_b"], ob["log_b"]) and torch.equal(og["anchor_w"], ob["anchor_w"]) | |
| # recursive rollout also passes through the gain (identity at δ = 0) | |
| torch.manual_seed(0) | |
| rg = cwm3d.build("cwm3d_recursive", horizon_gain=mode).eval() | |
| rg.load_state_dict(base.state_dict(), strict=False) | |
| torch.manual_seed(0) | |
| rb = cwm3d.build("cwm3d_recursive").eval() | |
| rb.load_state_dict(base.state_dict()) | |
| with torch.no_grad(): | |
| assert torch.allclose(rg(obs)["pred"], rb(obs)["pred"], atol=1e-3) | |
| def test_horizon_gain_learned_scales_the_anchored_residual_of_its_horizon_only(): | |
| x = _frames() | |
| obs, _ = split_obs_fut(x) | |
| base = _perturbed_base() | |
| torch.manual_seed(0) | |
| g = cwm3d.build("cwm3d", horizon_gain="learned").eval() | |
| g.load_state_dict(base.state_dict(), strict=False) | |
| with torch.no_grad(): | |
| g.hgain.delta[2] = 0.5 # g_3 = 1.5 | |
| pb, pg = base(obs)["pred"], g(obs)["pred"] | |
| last = obs[:, -1:, :, :2] | |
| for h in range(P.T_PRED): | |
| if h == 2: | |
| assert torch.allclose(pg[:, h] - last[:, 0], 1.5 * (pb[:, h] - last[:, 0]), rtol=1e-4, atol=1e-2) | |
| assert not torch.allclose(pg[:, h], pb[:, h], atol=1e-2) | |
| else: | |
| assert torch.allclose(pg[:, h], pb[:, h], atol=1e-3), h | |
| assert g.hgain.report() == {"mode": "learned", "g_h": [1.0, 1.0, 1.5, 1.0, 1.0]} | |
| def test_horizon_gain_parts_scales_hand_joints_of_its_horizon_only(): | |
| x = _frames() | |
| obs, _ = split_obs_fut(x) | |
| base = _perturbed_base() | |
| torch.manual_seed(0) | |
| g = cwm3d.build("cwm3d", horizon_gain="parts").eval() | |
| g.load_state_dict(base.state_dict(), strict=False) | |
| with torch.no_grad(): | |
| g.hgain.delta[4, 1] = 0.5 # hands at h=5: g = 1.5; everything else 1 | |
| pb, pg = base(obs)["pred"], g(obs)["pred"] | |
| last = obs[:, -1, :, :2] | |
| assert torch.allclose(pg[:, 4, HANDS] - last[:, HANDS], 1.5 * (pb[:, 4, HANDS] - last[:, HANDS]), rtol=1e-4, atol=1e-2) | |
| assert not torch.allclose(pg[:, 4, HANDS], pb[:, 4, HANDS], atol=1e-2) | |
| assert torch.allclose(pg[:, 4, OTHER], pb[:, 4, OTHER], atol=1e-3) | |
| assert torch.allclose(pg[:, :4], pb[:, :4], atol=1e-3) | |
| gn = g.hgain.gain() | |
| assert torch.all(gn[4, HANDS] == 1.5) and torch.all(gn[4, OTHER] == 1.0) and torch.all(gn[:4] == 1.0) | |
| with torch.no_grad(): | |
| g.hgain.delta[1, 0] = -0.25 # other joints at h=2: g = 0.75 | |
| r = g.hgain.report() | |
| assert r["mode"] == "parts" and r["g_h_hands"] == [1.0, 1.0, 1.0, 1.0, 1.5] and r["g_h_other"] == [1.0, 0.75, 1.0, 1.0, 1.0] | |
| def test_horizon_gain_has_gradient_and_is_tracked_by_ema(): | |
| x = _frames() | |
| obs, fut = split_obs_fut(x) | |
| torch.manual_seed(0) | |
| g = cwm3d.build("cwm3d", horizon_gain="learned") | |
| _perturb(g) | |
| with torch.no_grad(): | |
| g.hgain.delta.zero_() | |
| ema = EMA(g, 0.5) | |
| assert "hgain.delta" in dict(ema.shadow.named_parameters()) | |
| opt = torch.optim.AdamW(g.parameters(), lr=1e-2, weight_decay=0.05) | |
| loss, _ = LS.compute_losses(g(obs), fut) | |
| opt.zero_grad(); loss.backward() | |
| assert g.hgain.delta.grad is not None and torch.isfinite(g.hgain.delta.grad).all() and g.hgain.delta.grad.abs().sum() > 0 | |
| opt.step(); ema.update(g) | |
| assert not torch.all(g.hgain.delta == 0) | |
| assert torch.allclose(ema.shadow.hgain.delta, 0.5 * g.hgain.delta) # EMA(decay .5) of (0, δ) | |
| def test_train_kp_flags_build_the_arms_and_trunk_init_keeps_the_gain(tmp_path): | |
| a = get_args(["--model", "cwm3d", "--no-lock", "--hand-weight", "2", "--horizon-gain", "learned"]) | |
| assert a.hand_weight == 2.0 and a.horizon_gain == "learned" | |
| torch.manual_seed(0) | |
| m, _ = build_model("cwm3d", a) | |
| assert m.hgain is not None and m.hgain.mode == "learned" | |
| a2 = get_args(["--model", "cwm3d", "--no-lock", "--horizon-gain-parts"]) | |
| assert a2.horizon_gain == "parts" and a2.hand_weight == 1.0 | |
| m2, _ = build_model("cwm3d", a2) | |
| assert m2.hgain.mode == "parts" | |
| a3 = get_args(["--model", "cwm3d", "--no-lock", "--hand-weight", "4"]) | |
| m3, _ = build_model("cwm3d", a3) | |
| assert m3.hgain is None and LS.joint_weights(a3.hand_weight)[94] == 4.0 | |
| # trunk-only init from a checkpoint of the unchanged model (the adopted pretraining has no gain tensor): trunk loaded, δ kept at 0 | |
| base = _perturbed_base() | |
| ck = tmp_path / "ckpt_best.pt" | |
| torch.save({"model": base.state_dict(), "epoch": 4}, ck) | |
| rep = load_init_ckpt(m, ck, "trunk") | |
| assert rep["n_loaded"] == 59 and rep["n_skipped_shape"] == 0 and rep["n_skipped_part"] == 20 and rep["n_missing_or_kept"] == 21 | |
| assert "hgain.delta" in rep["kept_keys"] and torch.all(m.hgain.delta == 0) | |
| for k, v in base.state_dict().items(): | |
| if k.startswith(("tok.", "blocks.", "ln_f.")): | |
| assert torch.equal(m.state_dict()[k], v), k | |
| obs, _ = split_obs_fut(_frames()) | |
| with torch.no_grad(): | |
| assert torch.equal(m.eval()(obs)["pred"], obs[:, -1:, :, :2].expand(-1, P.T_PRED, -1, -1)) # head zero-init -> ZV start | |
| # rejected combinations | |
| with pytest.raises(ValueError): | |
| build_model("simlpe_kp", get_args(["--model", "simlpe_kp", "--no-lock", "--horizon-gain", "learned"])) | |
| with pytest.raises(ValueError): | |
| build_model("motionbert_kp", get_args(["--model", "motionbert_kp", "--no-lock", "--horizon-gain-parts"])) | |
| with pytest.raises(AssertionError): | |
| cwm3d.build("cwm3d", formulation="abs", horizon_gain="learned") | |
| with pytest.raises(AssertionError): | |
| cwm3d.build("cwm3d", horizon_gain="bogus") | |
Xet Storage Details
- Size:
- 15 kB
- Xet hash:
- b82da9f92a0032a65f75873f1b4776c160918499a35d48f8a635f60241595934
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.