Buckets:
| """Stage-A guarantees of CWM3-D (plan §4): | |
| 1. zero-init identity: a freshly built model (every variant) predicts EXACTLY the zero-velocity | |
| pose (|pred - last_obs| <= 1e-6 px) on random inputs and on real protocol clips from the cache; | |
| 2. no-future leak: the model is a pure function of the T_OBS observed frames -- NaN-poisoning the | |
| future block through the trainer's split never changes the output, a 10-frame input is | |
| rejected, perturbing observed frames changes the output while frames >= T_OBS cannot be seen; | |
| 3. causality inside the encoder: part tokens of frame t do not depend on frames > t (with the | |
| Δ-to-last channel held fixed). | |
| python -m pytest aide/tests/test_cwm3d_zero_init.py -q | |
| """ | |
| from __future__ import annotations | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| 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.train.train_kp import split_obs_fut # noqa: E402 | |
| VARIANTS = ("cwm3d", "cwm3d_noanchor", "cwm3d_nounc", "cwm3d_recursive", "cwm3d_dct", "cwm3d_vis") | |
| TOL = 1e-6 | |
| def _random_frames(n=6, seed=0): | |
| g = torch.Generator().manual_seed(seed) | |
| x = torch.rand(n, P.N_SAMPLE, P.J, 3, generator=g) * torch.tensor([P.W, P.H, 1.0]) | |
| return x | |
| def _real_frames(n=32): | |
| 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 _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 _zero_init_check(frames): | |
| obs, _ = split_obs_fut(frames) | |
| zv = obs[:, -1:, :, :2].expand(-1, P.T_PRED, -1, -1) | |
| for name in VARIANTS: | |
| torch.manual_seed(0) | |
| m = cwm3d.build(name).eval() | |
| with torch.no_grad(): | |
| o = m(obs) | |
| assert o["pred"].shape == zv.shape, name | |
| assert torch.isfinite(o["pred"]).all(), name | |
| assert (o["pred"] - zv).abs().max().item() <= TOL, (name, (o["pred"] - zv).abs().max().item()) | |
| # training mode too (dropout = 0, modality dropout only touches visual tokens) | |
| m.train() | |
| with torch.no_grad(): | |
| o2 = m(obs) | |
| assert (o2["pred"] - zv).abs().max().item() <= TOL, name | |
| def test_zero_init_equals_zero_velocity_random(): | |
| _zero_init_check(_random_frames()) | |
| def test_zero_init_equals_zero_velocity_real_clips(): | |
| fr = _real_frames() | |
| if fr is None: | |
| return | |
| _zero_init_check(fr) | |
| def test_zero_init_uncertainty_and_anchor_weights(): | |
| obs, _ = split_obs_fut(_random_frames()) | |
| torch.manual_seed(0) | |
| m = cwm3d.CWM3D().eval() | |
| with torch.no_grad(): | |
| o = m(obs) | |
| assert torch.allclose(o["log_b"], torch.full_like(o["log_b"], float(np.log(30.0)))) | |
| assert torch.allclose(o["anchor_w"], torch.full_like(o["anchor_w"], 1 / 3)) | |
| assert torch.equal(o["anchor_gate"], torch.zeros_like(o["anchor_gate"])) | |
| def test_anchor_mixture_reduces_to_zero_velocity_only_through_gate(): | |
| """With head output exactly zero the mixture is ZV; a non-zero gate moves towards the anchors.""" | |
| obs, _ = split_obs_fut(_random_frames()) | |
| anc = heads.anchors(obs) | |
| ho = torch.zeros(obs.shape[0], P.T_PRED, P.J, heads.HorizonHead.N_C) | |
| pred, _ = heads.AnchorMixture()(ho, anc) | |
| assert torch.equal(pred, anc[:, 0:1].expand_as(pred)) | |
| ho[..., heads.HorizonHead.C_GATE] = 1.0 | |
| pred, _ = heads.AnchorMixture()(ho, anc) | |
| assert torch.allclose(pred, anc.mean(1, keepdim=True).expand_as(pred), atol=1e-3) # uniform softmax | |
| def test_model_never_reads_future_frames(): | |
| frames = _random_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 name in VARIANTS: | |
| torch.manual_seed(0) | |
| m = cwm3d.build(name).eval() | |
| _perturb(m) | |
| with torch.no_grad(): | |
| a = m(obs_c)["pred"] | |
| b = m(obs_p)["pred"] | |
| assert torch.isfinite(a).all() and torch.equal(a, b), name | |
| # a 10-frame tensor (observed + future) must be rejected outright | |
| try: | |
| with torch.no_grad(): | |
| m(poison) | |
| raise RuntimeError(f"{name} accepted {P.N_SAMPLE} frames") | |
| except AssertionError: | |
| pass | |
| # the output does depend on the observed block (not a constant function) | |
| shifted = obs_c.clone() | |
| shifted[:, -1, :, :2] += 30.0 | |
| with torch.no_grad(): | |
| c = m(shifted)["pred"] | |
| assert not torch.allclose(a, c), name | |
| def test_encoder_causal_across_observed_frames(): | |
| torch.manual_seed(0) | |
| m = cwm3d.CWM3D().eval() | |
| _perturb(m) | |
| obs, _ = split_obs_fut(_random_frames()) | |
| x2 = obs.clone() | |
| x2[:, 2:P.T_OBS - 1, :, :2] += 40.0 # frames 2,3 change; frame 4 (the Δ reference) fixed | |
| with torch.no_grad(): | |
| p1, _ = m.encode(obs) | |
| p2, _ = m.encode(x2) | |
| assert torch.allclose(p1[:, :2], p2[:, :2], atol=1e-5) # frames 0,1 unaffected | |
| assert not torch.allclose(p1[:, 2:], p2[:, 2:]) | |
| def test_recursive_rollout_only_feeds_back_its_own_predictions(): | |
| """Recursive mode: horizon-1 output equals the chunked horizon-1 output (same weights), so the | |
| sliding window is built from predictions, never from ground truth.""" | |
| obs, _ = split_obs_fut(_random_frames()) | |
| torch.manual_seed(0) | |
| mc = cwm3d.CWM3D(rollout="chunked").eval() | |
| _perturb(mc) | |
| mr = cwm3d.CWM3D(rollout="recursive").eval() | |
| mr.load_state_dict(mc.state_dict()) | |
| with torch.no_grad(): | |
| pc, pr = mc(obs)["pred"], mr(obs)["pred"] | |
| assert torch.allclose(pc[:, 0], pr[:, 0], atol=1e-4) | |
| assert torch.isfinite(pr).all() | |
| 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:
- 6.26 kB
- Xet hash:
- 239cd8f921e49c9fb8b16d41e14fcf52095cc9c23daa08967547ba3ece914a7d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.