Buckets:
| """Building blocks of CWM3-D (plan §4): part tokenizer, chunked horizon head, anchor mixture. | |
| Coordinates are pixels; the per-joint input channels are | |
| (Δx, Δy)/SCALE relative to the LAST observed pose, x/W, y/H, conf (5 channels) | |
| """ | |
| from __future__ import annotations | |
| import math | |
| 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 | |
| SCALE = 200.0 | |
| N_CH = 5 | |
| ANCHORS = ("last", "mean_obs", "anti_velocity") | |
| ANTI_BETA = 0.1 | |
| # Pre-registered HM h=5 arms (docs/decisions.md 2026-09-17T07:00:22Z, part (A)): "the 42 hand keypoints" = PARTS lhand + rhand (HALPE 94-135) | |
| HAND_JOINTS = tuple(int(j) for j in list(P.PARTS["lhand"]) + list(P.PARTS["rhand"])) | |
| assert len(HAND_JOINTS) == 42 | |
| def joint_channels(obs: torch.Tensor) -> torch.Tensor: | |
| """obs [B,T,J,3] -> [B,T,J,5].""" | |
| xy = obs[..., :2] | |
| rel = (xy - xy[:, -1:]) / SCALE | |
| ab = torch.stack([xy[..., 0] / P.W, xy[..., 1] / P.H], -1) | |
| return torch.cat([rel, ab, obs[..., 2:3]], -1) | |
| def anchors(obs: torch.Tensor) -> torch.Tensor: | |
| """obs [B,T,J,3] -> anchor poses [B,3,J,2]: last, mean(obs), last - beta*(last - prev).""" | |
| xy = obs[..., :2] | |
| last, prev = xy[:, -1], xy[:, -2] | |
| return torch.stack([last, xy.mean(1), last - ANTI_BETA * (last - prev)], 1) | |
| class PartTokenizer(nn.Module): | |
| """4 tokens per frame (body, face, lhand, rhand). Hands share one Linear + a side embedding.""" | |
| def __init__(self, d=256): | |
| super().__init__() | |
| self.body = nn.Linear(len(P.PARTS["body"]) * N_CH, d) | |
| self.face = nn.Linear(len(P.PARTS["face"]) * N_CH, d) | |
| self.hand = nn.Linear(len(P.PARTS["lhand"]) * N_CH, d) | |
| self.side = nn.Parameter(torch.zeros(2, d)) # lhand / rhand | |
| 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): # -> [B,T,4,d] | |
| c = joint_channels(obs) # [B,T,J,5] | |
| B, T = c.shape[:2] | |
| g = lambda k: c[:, :, self.idx[k].to(c.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] | |
| def causal_part_mask(T: int, n_part: int, n_extra_per_frame: int = 0) -> torch.Tensor: | |
| """Boolean attention mask [L,L] (True = blocked) for tokens ordered frame-major: | |
| same-frame tokens fully connected, cross-frame strictly causal (later frames blocked).""" | |
| per = n_part + n_extra_per_frame | |
| frame = torch.arange(T * per) // per | |
| return frame[None, :] > frame[:, None] | |
| class Block(nn.Module): | |
| def __init__(self, d, heads, ffn, dropout=0.0): | |
| super().__init__() | |
| self.ln1 = nn.LayerNorm(d) | |
| self.attn = nn.MultiheadAttention(d, heads, dropout=dropout, batch_first=True) | |
| self.ln2 = nn.LayerNorm(d) | |
| self.ff = nn.Sequential(nn.Linear(d, ffn), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn, d)) | |
| self.drop = nn.Dropout(dropout) | |
| def forward(self, x, mask): | |
| h = self.ln1(x) | |
| x = x + self.drop(self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]) | |
| return x + self.drop(self.ff(self.ln2(x))) | |
| def dct_matrix(n: int) -> torch.Tensor: | |
| """Orthonormal DCT-II basis [n,n]; x = C^T c.""" | |
| k = torch.arange(n).float() | |
| C = torch.cos(math.pi / n * (k[None, :] + 0.5) * k[:, None]) * math.sqrt(2.0 / n) | |
| C[0] /= math.sqrt(2.0) | |
| return C | |
| class HorizonHead(nn.Module): | |
| """Chunked horizon head: the 4 part tokens of the last observed frame -> per part an MLP | |
| d -> hidden -> H*J_p*C. Channels per joint & horizon: μx, μy, log_b, 3 anchor logits, gate. | |
| Hands share one MLP (their tokens already carry the side embedding). | |
| Output layer zero-initialised so that μ = 0, logits = 0, gate = 0 at init (`zero_init=True`, the default; `zero_init=False` keeps | |
| PyTorch's default init of the output layers — formulation ablation `--formulation abs`, docs/decisions.md 2026-09-15).""" | |
| C_MU, C_LOGB, C_LOGIT, C_GATE = slice(0, 2), 2, slice(3, 6), 6 | |
| N_C = 7 | |
| def __init__(self, d=256, hidden=512, H=P.T_PRED, dct=False, zero_init=True): | |
| super().__init__() | |
| self.H, self.dct = H, dct | |
| self.mlp = nn.ModuleDict({ | |
| k: nn.Sequential(nn.LayerNorm(d), nn.Linear(d, hidden), nn.GELU(), | |
| nn.Linear(hidden, H * len(P.PARTS[k]) * self.N_C)) | |
| for k in ("body", "face", "lhand")}) | |
| if zero_init: | |
| for m in self.mlp.values(): | |
| nn.init.zeros_(m[-1].weight) | |
| nn.init.zeros_(m[-1].bias) | |
| self.register_buffer("C", dct_matrix(H)) | |
| def forward(self, tok_last): # tok_last [B,4,d] -> [B,H,J,N_C] | |
| outs = [] | |
| for i, k in enumerate(P.PART_NAMES): | |
| mlp = self.mlp["lhand" if k in ("lhand", "rhand") else k] | |
| o = mlp(tok_last[:, i]).view(-1, self.H, len(P.PARTS[k]), self.N_C) | |
| outs.append(o) | |
| o = torch.cat(outs, 2) # parts are contiguous in HALPE order | |
| if self.dct: # μ parametrised by DCT coefficients over horizons | |
| mu = torch.einsum("hk,bkjc->bhjc", self.C.t(), o[..., self.C_MU]) | |
| o = torch.cat([mu, o[..., 2:]], -1) | |
| return o | |
| class AnchorMixture(nn.Module): | |
| """pred = last + gate ⊙ (Σ_k softmax(logits)_k a_k − last) + SCALE·μ. | |
| With μ = logits = gate = 0 (zero-init head) this is exactly zero-velocity; gate = 1 recovers a | |
| pure convex anchor mixture. `enabled=False` -> pred = last + SCALE·μ (ablation).""" | |
| def __init__(self, enabled=True, log_b0=math.log(30.0)): | |
| super().__init__() | |
| self.enabled = enabled | |
| self.log_b0 = log_b0 | |
| def forward(self, head_out, anc): # head_out [B,H,J,7], anc [B,3,J,2] | |
| last = anc[:, 0:1] # [B,1,J,2] | |
| mu = SCALE * head_out[..., HorizonHead.C_MU] | |
| pred = last + mu | |
| if self.enabled: | |
| w = torch.softmax(head_out[..., HorizonHead.C_LOGIT], -1) # [B,H,J,3] | |
| mix = torch.einsum("bhjk,bkjc->bhjc", w, anc) | |
| pred = pred + head_out[..., HorizonHead.C_GATE, None] * (mix - last) | |
| log_b = head_out[..., HorizonHead.C_LOGB] + self.log_b0 | |
| return pred, log_b | |
| class HorizonGain(nn.Module): | |
| """Pre-registered arms (b) / (c) of the HM h=5 experiments (docs/decisions.md 2026-09-17T07:00:22Z, part (A)): a learned per-horizon gain | |
| on the ANCHORED residual of AnchorMixture, pred = last + g ⊙ (pred_mix − last) — i.e. the fixed SCALE (and the anchor term) are | |
| multiplied by g_h. g = 1 + δ with δ zero-initialised, so (i) the model is exactly zero-velocity at init as before, (ii) a gain | |
| model with δ = 0 predicts what the unchanged model predicts (up to float rounding of last + (pred − last)), and (iii) AdamW weight | |
| decay pulls g towards 1, not towards 0. mode "learned": δ [H] (5 scalars, arm b); mode "parts": δ [H, 2] (5 × 2 scalars, | |
| column 0 = the 94 other joints, column 1 = the 42 hand joints HAND_JOINTS, arm c). The joint→group index is a NON-persistent | |
| buffer, so the only new state-dict key is `hgain.delta`; the default model (horizon_gain=None in CWM3D) does not instantiate | |
| this module at all (its state dict, init and forward are byte-identical to the pre-change code — aide/tests/test_hm_h5_arms.py).""" | |
| MODES = ("learned", "parts") | |
| def __init__(self, mode: str, H: int = P.T_PRED): | |
| super().__init__() | |
| assert mode in self.MODES, mode | |
| self.mode, self.H = mode, H | |
| self.delta = nn.Parameter(torch.zeros(H) if mode == "learned" else torch.zeros(H, 2)) | |
| grp = torch.zeros(P.J, dtype=torch.long) | |
| grp[list(HAND_JOINTS)] = 1 | |
| self.register_buffer("group", grp, persistent=False) | |
| def gain(self) -> torch.Tensor: # -> [H, J] | |
| g = 1.0 + self.delta | |
| if self.mode == "learned": | |
| return g[:, None].expand(self.H, P.J) | |
| return g[:, self.group] | |
| def forward(self, pred, last): # pred [B,H,J,2], last [B,1,J,2] -> [B,H,J,2] | |
| return last + self.gain()[None, :, :, None] * (pred - last) | |
| def report(self) -> dict: | |
| """learned gains for the run json (g = 1 + δ).""" | |
| g = (1.0 + self.delta.detach()).cpu() | |
| if self.mode == "learned": | |
| return {"mode": self.mode, "g_h": [round(float(x), 5) for x in g]} | |
| return {"mode": self.mode, "g_h_other": [round(float(x), 5) for x in g[:, 0]], "g_h_hands": [round(float(x), 5) for x in g[:, 1]]} | |
| class AbsoluteReadout(nn.Module): | |
| """Formulation ablation (docs/decisions.md 2026-09-15 "Formulation ablation", EXPLANATORY, not for adoption): the head's μ channels | |
| are ABSOLUTE normalised coordinates, pred = μ ⊙ (W, H) px — no residual to the last pose, no ×SCALE, no anchor mixture (the anchor | |
| logits / gate channels are unused); log_b as AnchorMixture. Same call signature as AnchorMixture (the anchors are accepted and | |
| ignored); `enabled = False` so CWM3D emits no anchor outputs. The (W, H) buffer is non-persistent: the state dict keys are exactly | |
| those of the residual model.""" | |
| enabled = False | |
| def __init__(self, log_b0=math.log(30.0)): | |
| super().__init__() | |
| self.log_b0 = log_b0 | |
| self.register_buffer("wh", torch.tensor([float(P.W), float(P.H)]), persistent=False) | |
| def forward(self, head_out, anc=None): # head_out [B,H,J,7] -> pred [B,H,J,2] px, log_b [B,H,J] | |
| pred = head_out[..., HorizonHead.C_MU] * self.wh | |
| log_b = head_out[..., HorizonHead.C_LOGB] + self.log_b0 | |
| return pred, log_b | |
Xet Storage Details
- Size:
- 10.4 kB
- Xet hash:
- 12bd39e584c6194da42c8c5f5c324471a328aabc69699695cb5b2dce1ae59f99
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.