Buckets:
| """CWM3-D-kp hybrid with a DSTformer-style SPATIAL attention block in front of the part tokenizer (coordinator request | |
| 2026-09-12; run family `cwm3d_kp_spatial_s*_v2`). | |
| Per observed frame the 136 joints (5 channels each: (Δx,Δy)/SCALE to the last observed pose, x/W, y/H, conf) pass through | |
| `blocks` pre-norm self-attention blocks over the joint axis (MotionBERT-kp's spatial stream, d_s = 32, 4 heads); the | |
| attended per-joint features are CONCATENATED to the raw 5 channels and the part tokenizer reads 5 + d_s channels per joint | |
| (hands share the tokenizer as before). Everything after the tokenizer -- the causal temporal-part trunk, the horizon head | |
| and the anchor mixture -- is the unchanged CWM3D code (subclass; only `self.tok` is replaced), so the zero-init identity | |
| (pred == zero-velocity) and the no-future-leak guarantee are inherited and re-tested in aide/tests/test_hardening.py. | |
| """ | |
| from __future__ import annotations | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| import torch.nn as nn | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.models.cwm3d import CWM3D # noqa: E402 | |
| from aide.models.heads import N_CH, joint_channels # noqa: E402 | |
| from aide.models.motionbert_kp import Attn # noqa: E402 | |
| class SpatialBlock(nn.Module): | |
| def __init__(self, d, heads, ffn): | |
| super().__init__() | |
| self.attn = Attn(d, heads) | |
| self.ln = nn.LayerNorm(d) | |
| self.ff = nn.Sequential(nn.Linear(d, ffn), nn.GELU(), nn.Linear(ffn, d)) | |
| def forward(self, x): # [N, J, d] | |
| x = self.attn(x) | |
| return x + self.ff(self.ln(x)) | |
| class SpatialPartTokenizer(nn.Module): | |
| """Spatial attention over the joints of each frame, then the 4 part tokens (body, face, lhand, rhand) from | |
| [raw 5 channels, attended d_s channels] per joint. Same output contract as heads.PartTokenizer: [B,T,4,d].""" | |
| def __init__(self, d=256, d_s=32, heads=4, blocks=2): | |
| super().__init__() | |
| self.d_s = d_s | |
| self.inp = nn.Linear(N_CH, d_s) | |
| self.joint_emb = nn.Parameter(torch.zeros(P.J, d_s)) | |
| nn.init.normal_(self.joint_emb, std=0.02) | |
| self.blocks = nn.ModuleList(SpatialBlock(d_s, heads, 4 * d_s) for _ in range(blocks)) | |
| self.ln_s = nn.LayerNorm(d_s) | |
| c = N_CH + d_s | |
| self.body = nn.Linear(len(P.PARTS["body"]) * c, d) | |
| self.face = nn.Linear(len(P.PARTS["face"]) * c, d) | |
| self.hand = nn.Linear(len(P.PARTS["lhand"]) * c, d) | |
| self.side = nn.Parameter(torch.zeros(2, d)) | |
| self.part_emb = nn.Parameter(torch.zeros(len(P.PARTS), d)) | |
| self.time_emb = nn.Parameter(torch.zeros(P.T_OBS, d)) | |
| nn.init.normal_(self.side, std=0.02) | |
| nn.init.normal_(self.part_emb, std=0.02) | |
| nn.init.normal_(self.time_emb, std=0.02) | |
| self.idx = {k: torch.as_tensor(v, dtype=torch.long) for k, v in P.PARTS.items()} | |
| def forward(self, obs): # obs [B,T,J,3] -> [B,T,4,d] | |
| c = joint_channels(obs) # [B,T,J,5] | |
| B, T = c.shape[:2] | |
| s = self.inp(c).reshape(B * T, P.J, self.d_s) + self.joint_emb[None] | |
| for blk in self.blocks: | |
| s = blk(s) # spatial attention within each frame (frames independent) | |
| s = self.ln_s(s).view(B, T, P.J, self.d_s) | |
| f = torch.cat([c, s], -1) # [B,T,J,5+d_s] | |
| g = lambda k: f[:, :, self.idx[k].to(f.device)].flatten(2) # noqa: E731 | |
| tok = torch.stack([self.body(g("body")), self.face(g("face")), | |
| self.hand(g("lhand")) + self.side[0], self.hand(g("rhand")) + self.side[1]], 2) | |
| return tok + self.part_emb[None, None] + self.time_emb[None, :T, None] | |
| class CWM3DSpatial(CWM3D): | |
| name = "cwm3d_kp_spatial" | |
| def __init__(self, d_s=32, s_heads=4, s_blocks=2, **kw): | |
| super().__init__(**kw) | |
| self.tok = SpatialPartTokenizer(self.d, d_s, s_heads, s_blocks) | |
| def build(name: str, **kw) -> nn.Module: | |
| if name == "cwm3d_kp_spatial": | |
| return CWM3DSpatial(**kw) | |
| raise ValueError(name) | |
| if __name__ == "__main__": | |
| m = CWM3DSpatial() | |
| x = torch.rand(2, P.T_OBS, P.J, 3) * torch.tensor([P.W, P.H, 1.0]) | |
| o = m(x) | |
| print("params", sum(p.numel() for p in m.parameters()), "zero-init == ZV:", | |
| torch.allclose(o["pred"], x[:, -1:, :, :2].expand_as(o["pred"]))) | |
Xet Storage Details
- Size:
- 4.49 kB
- Xet hash:
- e0326ed39ae5e2399f5c6cecbdff174985ca1b553abb319003222a77e0fcd8a2
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.