Buckets:
| """Guards for the AIDE architecture-attribution experiments (docs/decisions.md 2026-09-14): | |
| * Driver-WM-kp starts exactly at zero-velocity, rejects 10-frame inputs, depends on the observed block only (keypoints AND exterior | |
| features of frames >= T_OBS never reach the output), and its internal stream uses cwm3d's PartTokenizer; | |
| * `train_kp.load_init_ckpt(..., parts="trunk")` on MotionBERT-kp / ST-GCN-kp loads exactly the declared trunk tensors and leaves the | |
| read-out head at its seed init, so a trunk-initialised fine-tune still starts at zero-velocity; | |
| * `pretrain_cwm3d_baton.build_model` reproduces train_kp's initial weights for the same seed (motionbert_kp / stgcn_kp) and the | |
| cwm3d path is unchanged. | |
| python -m pytest aide/tests/test_arch_attribution.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import argparse # noqa: E402 | |
| import sys # noqa: E402 | |
| from pathlib import Path # 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, driverwm_kp, driverwm_repro, motionbert_kp, stgcn_kp # noqa: E402 | |
| from aide.models.heads import PartTokenizer # noqa: E402 | |
| from aide.train import pretrain_cwm3d_baton as PRE # noqa: E402 | |
| from aide.train.train_kp import load_init_ckpt, seed_all, split_obs_fut # noqa: E402 | |
| D = driverwm_repro.D_FEAT | |
| 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): | |
| 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_driverwm_kp_zero_velocity_at_init_and_shapes(): | |
| x, v = _frames() | |
| obs, _ = split_obs_fut(x) | |
| torch.manual_seed(0) | |
| m = driverwm_kp.DriverWMKP().eval() | |
| with torch.no_grad(): | |
| y = m(obs, v[:, :P.T_OBS]) | |
| assert y.shape == (3, P.T_PRED, P.J, 2) | |
| assert torch.equal(y, obs[:, -1:, :, :2].expand_as(y)) | |
| assert isinstance(m.tok, PartTokenizer) and m.D == len(P.PARTS) * 256 | |
| def test_driverwm_kp_rejects_future_frames_and_is_causal_in_obs_only(): | |
| x, v = _frames() | |
| obs, _ = split_obs_fut(x) | |
| torch.manual_seed(0) | |
| m = driverwm_kp.DriverWMKP().eval() | |
| _perturb(m) | |
| with torch.no_grad(): | |
| try: | |
| m(x, v) # 10 frames must be rejected | |
| assert False, "10-frame input accepted" | |
| except AssertionError as e: | |
| if "10-frame input accepted" in str(e): | |
| raise | |
| y = m(obs, v[:, :P.T_OBS]) | |
| assert not torch.equal(y, obs[:, -1:, :, :2].expand_as(y)) # perturbed model moves | |
| # changing the FUTURE keypoints / features of the 10-frame block cannot change the output (only obs is handed in) | |
| x2, v2 = x.clone(), v.clone() | |
| x2[:, P.T_OBS:] += 100.0 | |
| v2[:, P.T_OBS:] += 3.0 | |
| obs2, _ = split_obs_fut(x2) | |
| y2 = m(obs2, v2[:, :P.T_OBS]) | |
| assert torch.equal(y, y2) | |
| # the exterior features DO reach the output (the external stream is live), the in-cabin view does NOT (replaced by keypoints) | |
| v3 = v.clone(); v3[:, :P.T_OBS, driverwm_repro.EXT_VIEWS] += 1.0 | |
| assert not torch.allclose(y, m(obs, v3[:, :P.T_OBS])) | |
| v4 = v.clone(); v4[:, :P.T_OBS, driverwm_repro.IN_VIEW] += 1.0 | |
| assert torch.equal(y, m(obs, v4[:, :P.T_OBS])) | |
| def test_trunk_init_loads_trunk_only_and_keeps_zero_velocity(tmp_path): | |
| x, _ = _frames() | |
| obs, _ = split_obs_fut(x) | |
| for mod, build in ((motionbert_kp, lambda: motionbert_kp.build("motionbert_kp", dropout=0.0)), | |
| (stgcn_kp, lambda: stgcn_kp.build("stgcn_kp", dropout=0.0))): | |
| torch.manual_seed(1) | |
| src = build() | |
| _perturb(src, std=0.05, seed=2) # a "pretrained" model: every tensor differs from the seed init | |
| ck = tmp_path / f"{src.name}.pt" | |
| torch.save({"model": src.state_dict(), "epoch": 3}, ck) | |
| torch.manual_seed(7) | |
| dst = build() | |
| before = {k: v.clone() for k, v in dst.state_dict().items()} | |
| rep = load_init_ckpt(dst, str(ck), "trunk", "cpu") | |
| after = dst.state_dict() | |
| trunk = tuple(dst.TRUNK_PREFIXES) | |
| for k in after: | |
| if k.startswith(trunk): | |
| assert torch.equal(after[k], src.state_dict()[k]), k # trunk taken from the checkpoint | |
| else: | |
| assert torch.equal(after[k], before[k]), k # head 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) | |
| assert all(k.startswith(("temp_out.", "out.")) for k in rep["kept_keys"]) | |
| with torch.no_grad(): | |
| y = dst.eval()(obs) | |
| assert torch.equal(y, obs[:, -1:, :, :2].expand_as(y)) # zero-velocity start of the fine-tune | |
| def test_pretrain_build_model_matches_train_kp_seed_init(): | |
| for name, mod in (("motionbert_kp", motionbert_kp), ("stgcn_kp", stgcn_kp)): | |
| a = argparse.Namespace(model=name, d=256, layers=4, heads=4, ffn=1024, head_hidden=512) | |
| seed_all(123) | |
| m_pre, files = PRE.build_model(a) | |
| seed_all(123) | |
| m_tr = mod.build(name, dropout=0.0) # train_kp.build_model's call for these names | |
| for (k, v), (k2, v2) in zip(m_pre.state_dict().items(), m_tr.state_dict().items()): | |
| assert k == k2 and torch.equal(v, v2), k | |
| assert any(Path(f).name == f"{name}.py" for f in files) | |
| a = argparse.Namespace(model="cwm3d", d=256, layers=4, heads=4, ffn=1024, head_hidden=512) | |
| seed_all(42) | |
| m_pre, _ = PRE.build_model(a) | |
| seed_all(42) | |
| m_ref = cwm3d.build("cwm3d", d=256, layers=4, heads=4, ffn=1024, head_hidden=512, dropout=0.0) | |
| for (k, v), (k2, v2) in zip(m_pre.state_dict().items(), m_ref.state_dict().items()): | |
| assert k == k2 and torch.equal(v, v2), k | |
Xet Storage Details
- Size:
- 6.28 kB
- Xet hash:
- df9cefc3ddd67cfc5ffe628548ee642a7192720e313ca7155b95632fba6230bc
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.