Buckets:
| """Guards for the pre-registered pretraining experiments P1-ext / P2 (docs/decisions.md 2026-09-15): | |
| * `pretrain_cwm3d_baton.build_model(--model driverwm_kp)` reproduces train_kp's initial weights for the same seed (build parity); | |
| * `DriverWMKP.TRUNK_PREFIXES` covers every tensor except the zero-initialised residual read-out `dec.out.*`, and | |
| `train_kp.load_init_ckpt(..., parts="trunk")` loads exactly those tensors while the fine-tune still starts at exact zero-velocity; | |
| * the pretraining path (forward without exterior features) depends on the observed block only — poisoning the future frames of the | |
| 10-frame block leaves the output bit-identical, 10-frame inputs are rejected, and the missing-exterior path equals explicit zero features; | |
| * `pretrain_cwm3d_baton.motion_scale`: --motion-scale-max 1.0 is an exact no-op; factor 2 doubles every displacement relative to the last | |
| observed pose while leaving that pose and the confidences fixed; random factors lie in [1, S] and are constant within a window. | |
| python -m pytest aide/tests/test_driverwm_pretrain.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 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.models import driverwm_kp, driverwm_repro # noqa: E402 | |
| from aide.train import pretrain_cwm3d_baton as PRE # noqa: E402 | |
| from aide.train import train_kp as TK # noqa: E402 | |
| from aide.train.train_kp import load_init_ckpt, seed_all, split_obs_fut # noqa: E402 | |
| D = driverwm_repro.D_FEAT | |
| READOUT = {"dec.out.weight", "dec.out.bias"} | |
| 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 _train_kp_args(seed): | |
| a = TK.get_args(["--model", "driverwm_kp", "--seed", str(seed), "--cpu", "--no-lock"]) | |
| assert a.dropout == 0.0 and a.visual is True | |
| return a | |
| def test_pretrain_build_model_matches_train_kp_seed_init_driverwm(): | |
| for seed in (123, 42): | |
| seed_all(seed) | |
| m_pre, files = PRE.build_model(argparse.Namespace(model="driverwm_kp")) | |
| seed_all(seed) | |
| m_tr, _ = TK.build_model("driverwm_kp", _train_kp_args(seed)) # train_kp's build for this seed (dropout 0) | |
| sd_pre, sd_tr = m_pre.state_dict(), m_tr.state_dict() | |
| assert list(sd_pre) == list(sd_tr) | |
| for k in sd_pre: | |
| assert torch.equal(sd_pre[k], sd_tr[k]), k | |
| assert isinstance(m_pre, driverwm_kp.DriverWMKP) and any(Path(f).name == "driverwm_kp.py" for f in files) | |
| # the other models are untouched by the new branch | |
| seed_all(7) | |
| m_c, _ = PRE.build_model(argparse.Namespace(model="cwm3d", d=256, layers=4, heads=4, ffn=1024, head_hidden=512)) | |
| assert sum(p.numel() for p in m_c.parameters()) == 5_771_449 | |
| def test_trunk_prefixes_cover_everything_but_the_zero_readout(): | |
| torch.manual_seed(0) | |
| m = driverwm_kp.DriverWMKP() | |
| keys = list(m.state_dict()) | |
| trunk = tuple(m.TRUNK_PREFIXES) | |
| kept = {k for k in keys if not k.startswith(trunk)} | |
| assert kept == READOUT, kept | |
| n_trunk = sum(k.startswith(trunk) for k in keys) | |
| assert n_trunk == len(keys) - 2 and n_trunk > 0 | |
| assert torch.count_nonzero(m.dec.out.weight) == 0 and torch.count_nonzero(m.dec.out.bias) == 0 | |
| # every prefix names at least one tensor (no dead prefix) and each tensor matches exactly one prefix | |
| for pre in trunk: | |
| assert any(k.startswith(pre) for k in keys), pre | |
| for k in keys: | |
| assert sum(k.startswith(pre) for pre in trunk) <= 1, k | |
| def test_trunk_init_loads_trunk_only_and_keeps_zero_velocity_driverwm(tmp_path): | |
| x, v = _frames() | |
| obs, _ = split_obs_fut(x) | |
| torch.manual_seed(1) | |
| src = driverwm_kp.DriverWMKP() | |
| _perturb(src, std=0.05, seed=2) # "pretrained": every tensor differs from the seed init | |
| ck = tmp_path / "driverwm_kp_pre.pt" | |
| torch.save({"model": src.state_dict(), "epoch": 3}, ck) | |
| torch.manual_seed(7) | |
| dst = driverwm_kp.DriverWMKP() | |
| before = {k: t.clone() for k, t in dst.state_dict().items()} | |
| rep = load_init_ckpt(dst, str(ck), "trunk", "cpu") | |
| after, ssd = dst.state_dict(), src.state_dict() | |
| trunk = tuple(dst.TRUNK_PREFIXES) | |
| for k in after: | |
| if k.startswith(trunk): | |
| assert torch.equal(after[k], ssd[k]), k # trunk taken from the checkpoint | |
| else: | |
| assert k in READOUT and torch.equal(after[k], before[k]), k # read-out kept at the (zero) seed init | |
| assert rep["parts"] == "trunk" and rep["n_skipped_shape"] == 0 and rep["n_skipped_part"] == 2 | |
| assert rep["n_loaded"] == len(after) - 2 and set(rep["kept_keys"]) == READOUT | |
| with torch.no_grad(): | |
| y_zero = dst.eval()(obs) # pretraining path: no exterior features | |
| y_vis = dst.eval()(obs, v[:, :P.T_OBS]) # AIDE path: exterior features present | |
| assert torch.equal(y_zero, obs[:, -1:, :, :2].expand_as(y_zero)) # exact zero-velocity start of the fine-tune | |
| assert torch.equal(y_vis, obs[:, -1:, :, :2].expand_as(y_vis)) | |
| def test_pretraining_path_ignores_future_and_equals_explicit_zero_exterior(): | |
| x, v = _frames() | |
| obs, _ = split_obs_fut(x) | |
| seed_all(123) | |
| m, _ = PRE.build_model(argparse.Namespace(model="driverwm_kp")) | |
| m.eval() | |
| _perturb(m) | |
| with torch.no_grad(): | |
| y = m(obs) # the pretraining script's call (vis omitted) | |
| assert not torch.equal(y, obs[:, -1:, :, :2].expand_as(y)) # perturbed model moves | |
| x2 = x.clone(); x2[:, P.T_OBS:] = float("nan") # poisoned future block | |
| obs2, fut2 = split_obs_fut(x2) | |
| assert torch.equal(obs, obs2) and torch.isnan(fut2).all() | |
| assert torch.equal(y, m(obs2)) | |
| z = obs.new_zeros(obs.shape[0], P.T_OBS, len(P.VIEWS), D) | |
| assert torch.equal(y, m(obs, z)) # missing-exterior path == explicit zero features | |
| assert not torch.allclose(y, m(obs, v[:, :P.T_OBS])) # the exterior stream is live when features are given | |
| try: | |
| m(x2) | |
| raise RuntimeError("10-frame input accepted") | |
| except AssertionError: | |
| pass | |
| # the same loss call as the pretraining loop reduces to the Euclidean term for the tensor output | |
| from aide.models import losses as LS | |
| _, fut = split_obs_fut(x) | |
| w = {"euc": 1.0, "nll": 0.1, "bone": 0.05, "jerk": 0.01} | |
| with torch.no_grad(): | |
| loss, terms = LS.compute_losses(m(obs), fut, w, "euc") | |
| assert set(terms) == {"euc"} and torch.isfinite(loss) | |
| def test_motion_scale_noop_and_doubling(): | |
| x, _ = _frames(B=4, seed=3) | |
| x0 = x.clone() | |
| y = PRE.motion_scale(x.clone(), 1.0) | |
| assert torch.equal(y, x0) # --motion-scale-max 1.0 (default) is an exact no-op | |
| y2 = PRE.motion_scale(x.clone(), 2.5, factors=torch.full((4,), 2.0)) | |
| last = x0[:, P.T_OBS - 1:P.T_OBS, :, :2] | |
| assert torch.equal(y2[:, P.T_OBS - 1, :, :2], last[:, 0]) # the last observed pose is fixed | |
| assert torch.equal(y2[..., 2], x0[..., 2]) # confidences untouched | |
| assert torch.allclose(y2[..., :2] - last, 2.0 * (x0[..., :2] - last), atol=1e-3) # every displacement doubled (future AND observed) | |
| fut_disp0 = (x0[:, P.T_OBS:, :, :2] - last).norm(dim=-1).mean() | |
| fut_disp2 = (y2[:, P.T_OBS:, :, :2] - last).norm(dim=-1).mean() | |
| assert abs(float(fut_disp2 / fut_disp0) - 2.0) < 1e-3 | |
| def test_motion_scale_random_factors_in_range_and_constant_per_window(): | |
| x, _ = _frames(B=64, seed=5) | |
| x0 = x.clone() | |
| torch.manual_seed(0) | |
| y = PRE.motion_scale(x.clone(), 2.5) | |
| last = x0[:, P.T_OBS - 1:P.T_OBS, :, :2] | |
| d0, d1 = x0[..., :2] - last, y[..., :2] - last | |
| mask = d0.abs() > 1.0 # avoid 0/0 on the anchor frame | |
| ratio = (d1[mask] / d0[mask]) | |
| assert float(ratio.min()) >= 1.0 - 1e-4 and float(ratio.max()) <= 2.5 + 1e-4 | |
| per_win = [float((d1[b][mask[b]] / d0[b][mask[b]]).std()) for b in range(64)] | |
| assert max(per_win) < 1e-3 # one factor per window | |
| assert float(np.std([float((d1[b][mask[b]] / d0[b][mask[b]]).mean()) for b in range(64)])) > 0.1 # factors do vary across windows | |
Xet Storage Details
- Size:
- 9.19 kB
- Xet hash:
- dfda1267d663434b0bedb2307c0f4bf6f0eff3a1eb7858d3994147b6facb4f4e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.