Buckets:
| """cwm3/tests/test_kp_forecaster.py — LANE F1 keypoint residual forecaster + distillation injector (models/kp_forecaster.py, models/kpf_inject.py, | |
| train/train_kp_residual.py; R2 WP4 batch 2, pre-registration docs/decisions.md 2026-09-17T19:33:12Z). | |
| (a) observation offsets are all <= 0 (dense 20 / full 28 / 1 Hz 11 frames); the horizon ticks are all > 0. | |
| (b) CAUSAL NORMALISATION: `ref` is the newest valid position <= t per joint, `scale` the median shoulder width over the observed frames; a store whose FUTURE | |
| frames (ticks > 0) are poisoned yields bitwise identical observation tensors and forecaster outputs, while the gathered targets differ. | |
| (c) parameter match: the MotionBERT-architecture arm and the CockpitWM-form arm are within ±10 % of each other and both < 2 M; the forecaster is zero-velocity at | |
| init (every mode's residual = 0); loss parts are finite on synthetic batches; the distillation channel is all-zero for invalid anchors and sums to 1 over modes. | |
| (d) the joint model with the F1 injector (kpf_dim 10, zero-init) returns emissions BITWISE equal to the model without it on the same batch (random kpf input); | |
| after one optimiser step the injector's residual is non-zero; train/batch.prep_U passes `kpf` through only when the store has it. | |
| Synthetic data only (never a result). Run: cd cwm3 && python -m pytest tests/test_kp_forecaster.py -q | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| torch.set_num_threads(4) | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train", CWM3 / "common"): | |
| if str(p) not in sys.path: | |
| sys.path.insert(0, str(p)) | |
| from models import kp_forecaster as kf # noqa: E402 | |
| from models.kpf_inject import KpfInject # noqa: E402 | |
| from models.cwm3 import CWM3Joint # noqa: E402 | |
| import train_kp_residual as tkr # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| from tests.test_zero_init_identity import _batch, _fresh_branches # noqa: E402 | |
| def _synthetic_store(n_routes=3, n_ticks=400, seed=0, gap_at=None): | |
| """a CPU Store with a contiguous tick grid per route (optionally one missing tick).""" | |
| rng = np.random.default_rng(seed) | |
| kp, valid, lut, off, ln = [], [], [], [], [] | |
| o = 0 | |
| for r in range(n_routes): | |
| k = rng.uniform(0.2, 0.8, (n_ticks, kf.NJ, 3)).astype(np.float16) | |
| v = rng.random((n_ticks, kf.NJ)) > 0.1 | |
| table = np.arange(n_ticks, dtype=np.int64) + o | |
| if gap_at is not None: | |
| table[gap_at] = -1 | |
| kp.append(k); valid.append(v); lut.append(table); off.append(o); ln.append(n_ticks); o += n_ticks | |
| st = tkr.Store.__new__(tkr.Store) | |
| st.kp = torch.from_numpy(np.concatenate(kp)); st.valid = torch.from_numpy(np.concatenate(valid)); st.lut = torch.from_numpy(np.concatenate(lut)) | |
| st.lut_off = torch.tensor(off); st.lut_len = torch.tensor(ln); st.L = int(st.lut.numel()); st.device = torch.device("cpu"); st.n_frames = int(st.kp.shape[0]) | |
| return st | |
| def test_offsets_and_horizons_non_positive(): | |
| for kind, n in (("full", 28), ("dense", 20), ("1hz", 11)): | |
| off = kf.obs_offsets(kind) | |
| assert off.numel() == n and int(off.max()) == 0 and bool((off <= 0).all()) and bool((off[1:] > off[:-1]).all()) | |
| assert all(h > 0 for h in kf.HORIZON_TICKS) and kf.HORIZONS_S == (0.5, 1.0, 2.0, 3.0) | |
| def test_causal_normalise_ref_and_scale(): | |
| B, T = 2, 6 | |
| xy = torch.rand(B, T, kf.NJ, 2) | |
| valid = torch.ones(B, T, kf.NJ, dtype=torch.bool) | |
| valid[0, -1, 9] = False # left wrist invalid at t -> ref = frame t-1 | |
| valid[1, :, 10] = False # right wrist never valid -> ref_ok False, ref 0 | |
| xn, ref, ref_ok, scale = kf.causal_normalise(xy, valid, 0.12) | |
| assert torch.equal(ref[0, 0], xy[0, -1, 0]) and torch.equal(ref[0, 9], xy[0, -2, 9]) | |
| assert not ref_ok[1, 10] and torch.equal(ref[1, 10], torch.zeros(2)) | |
| sw = (xy[:, :, kf.L_SHO] - xy[:, :, kf.R_SHO]).norm(dim=-1) | |
| assert torch.allclose(scale, sw.median(1).values.clamp(kf.SCALE_MIN, kf.SCALE_MAX)) | |
| assert torch.equal(xn[1, :, 10], torch.zeros(T, 2)) # invalid steps are zero | |
| assert torch.allclose(xn[0, -1, 0], torch.zeros(2)) # the anchor frame of a valid joint is exactly the reference | |
| def test_future_poison_leaves_inputs_and_outputs_unchanged(): | |
| torch.manual_seed(0) | |
| st = _synthetic_store() | |
| model = kf.KpForecaster("mb", "full", d=32, depth=1, heads=2, mlp_ratio=1).eval() | |
| pos = torch.tensor([150, 200, 250]); route = torch.tensor([0, 1, 2]) | |
| a = {"pos": pos, "route": route, "S": torch.full((3, 3, 2), 0.5)} | |
| idx = torch.arange(3) | |
| with torch.no_grad(): | |
| obs1 = st.gather(pos, route, model.offsets) | |
| out1, has1, tot1, parts1, err1, fxy1, fv1 = tkr.run_batch(model, st, a, idx) | |
| # poison EVERY frame strictly after each anchor tick (global frames > pos within the route) -> observations must not change | |
| poison = st.kp.clone(); pv = st.valid.clone() | |
| for p_, r_ in zip(pos.tolist(), route.tolist()): | |
| lo = int(st.lut_off[r_]); hi = lo + int(st.lut_len[r_]) | |
| fut = st.lut[p_ + 1:hi]; fut = fut[fut >= 0] | |
| poison[fut] = torch.tensor([9.0, 9.0, 0.0]).half(); pv[fut] = ~pv[fut] | |
| st.kp, st.valid = poison, pv | |
| with torch.no_grad(): | |
| obs2 = st.gather(pos, route, model.offsets) | |
| out2, has2, tot2, parts2, err2, fxy2, fv2 = tkr.run_batch(model, st, a, idx) | |
| for x1, x2 in zip(obs1, obs2): | |
| assert torch.equal(x1, x2), "observation tensors read a future frame" | |
| for k in ("res", "mode_logits", "state_logits", "ref", "scale"): | |
| assert torch.equal(out1[k], out2[k]), f"forecaster output {k} changed under a poisoned future" | |
| assert not torch.equal(fxy1, fxy2), "the gathered targets must see the poison (sanity of the poison)" | |
| def test_gather_respects_route_bounds_and_gaps(): | |
| st = _synthetic_store(gap_at=205) | |
| pos = torch.tensor([5, 210]); route = torch.tensor([0, 0]) # anchor tick 5: the 1 Hz history falls before the route start -> invalid | |
| xy, conf, valid = st.gather(pos, route, kf.obs_offsets("full")) | |
| assert not valid[0, :8].any() and valid[0, 8:].any() | |
| assert not valid[1, kf.obs_offsets("full").tolist().index(-5)].any() # tick 205 is the gap -> that frame invalid | |
| fxy, _, fv = st.gather(torch.tensor([394]), torch.tensor([0]), torch.tensor(kf.HORIZON_TICKS)) # ticks 0..399: +5 -> 399 inside, +10 -> 404 outside | |
| assert fv[0, 0].any() and not fv[0, 1:].any() # only +0.5 s exists before the route end | |
| def test_param_match_zero_velocity_init_loss_and_channel(): | |
| torch.manual_seed(0) | |
| mb = kf.KpForecaster(**tkr.ARMS["mb_full"]); cw = kf.KpForecaster(**tkr.ARMS["cwm_full"]) | |
| n_mb, n_cw = kf.n_params(mb), kf.n_params(cw) | |
| assert n_mb < 2_000_000 and n_cw < 2_000_000 and abs(n_mb - n_cw) / max(n_mb, n_cw) <= 0.10, (n_mb, n_cw) | |
| for m in (mb, cw): | |
| m.eval() | |
| B, T = 5, m.n_obs | |
| xy = torch.rand(B, T, kf.NJ, 2); conf = torch.rand(B, T, kf.NJ); valid = torch.rand(B, T, kf.NJ) > 0.2 | |
| out = m(xy, conf, valid) | |
| assert out["res"].shape == (B, kf.H, kf.N_MODES, kf.NJ, 2) and torch.equal(out["res"], torch.zeros_like(out["res"])) # zero-velocity at init | |
| fxy = torch.rand(B, kf.H, kf.NJ, 2); fv = torch.rand(B, kf.H, kf.NJ) > 0.3 | |
| S = torch.tensor([[0.0, 1.0], [1.0, float("nan")], [0.0, 0.0]]).expand(B, -1, -1).clone() | |
| total, parts, err = kf.kp_loss(out, fxy, fv, S) | |
| assert torch.isfinite(total) and all(np.isfinite(v) for v in parts.values()) and err.shape == (B, kf.N_MODES) | |
| ok = torch.tensor([True, True, False, True, False]) | |
| ch = kf.kpf_channel(out, ok) | |
| assert ch.shape == (B, kf.KPF_DIM) and torch.equal(ch[~ok], torch.zeros(2, kf.KPF_DIM)) | |
| assert torch.allclose(ch[ok, 6:9].sum(1), torch.ones(3)) and torch.equal(ch[ok, 9], torch.ones(3)) | |
| def test_inject_identity_at_init_and_learns(): | |
| V, U, E = _fresh_branches() | |
| b, _ = _batch() | |
| b["U"]["kpf"] = torch.rand(len(b["U"]["mode"]), kf.KPF_DIM) | |
| ref = CWM3Joint(V, U, E, decoder="direct", k_head=True, seed=0).eval() | |
| V2, U2, E2 = _fresh_branches() | |
| m = CWM3Joint(V2, U2, E2, decoder="direct", k_head=True, seed=0, kpf_dim=kf.KPF_DIM).eval() | |
| assert m.kpf is not None and isinstance(m.kpf, KpfInject) and m.param_report()["kpf"] == kf.n_params(m.kpf) and kf.n_params(m.kpf) < 20_000 | |
| with torch.no_grad(): | |
| o1, o2 = ref(b), m(b) | |
| for k in ("V", "U", "E", "F", "P", "R"): | |
| assert torch.equal(o1[k], o2[k]), f"{k} differs at init with the zero-init injector" | |
| assert torch.equal(m.kpf(b["U"]["kpf"]), torch.zeros(len(b["U"]["mode"]), 128)) | |
| opt = torch.optim.SGD(m.kpf.parameters(), lr=0.1) | |
| m.train() | |
| out = m(b) | |
| out["U"].pow(2).mean().backward() | |
| opt.step() | |
| assert not torch.equal(m.kpf(b["U"]["kpf"]), torch.zeros(len(b["U"]["mode"]), 128)), "the injector never learns" | |
| def test_prep_U_passes_kpf_only_when_present(): | |
| class FakeStore: | |
| def __init__(self, with_kpf): | |
| n = 6 | |
| self.device = torch.device("cpu") | |
| self.t = {"Kh": torch.rand(n, 40, 11), "X": torch.rand(n, 40, 5), "mode": torch.zeros(n, dtype=torch.long)} | |
| if with_kpf: | |
| self.t["kpf"] = torch.rand(n, kf.KPF_DIM).half() | |
| def has(self, k): | |
| return k in self.t | |
| st = {"xmu": np.zeros(5, np.float32), "xsd": np.ones(5, np.float32), "kmu": np.zeros(11, np.float32), "ksd": np.ones(11, np.float32)} | |
| b, _ = bt.prep_U(FakeStore(True), torch.arange(6), st, {"cabin": "none"}, need_targets=False) | |
| assert "kpf" in b and b["kpf"].shape == (6, kf.KPF_DIM) and b["kpf"].dtype == torch.float32 and b["Kx"].shape == (6, 40, 18) | |
| b0, _ = bt.prep_U(FakeStore(False), torch.arange(6), st, {"cabin": "none"}, need_targets=False) | |
| assert "kpf" not in b0 | |
Xet Storage Details
- Size:
- 10.3 kB
- Xet hash:
- d434156216d8d10769848004d583c30426646fc4f664b8d7b084600d8983aca6
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.