Buckets:
| """Driver-WM-kp: the Driver-WM reproduction (aide/models/driverwm_repro.py, protocol.md C1-C16) whose INTERNAL (in-cabin) stream is | |
| fed by CWM3-D's keypoint part tokens (aide/models/heads.PartTokenizer) instead of the frozen Qwen3-VL in-cabin features, trained | |
| with the shared keypoint recipe (arch attribution experiment (2), docs/decisions.md 2026-09-14; design notes aide/docs/driverwm_kp_port.md). | |
| KEPT from the reproduction (classes imported from driverwm_repro so the code is literally the same) | |
| * external-context stream: mean over the 3 exterior views of the pooled Qwen3-VL features + view embedding + step embedding | |
| -> one causal pre-norm TransformerEncoderLayer(2048, 4 heads, d_ff 4096) (`pre_ext`), residual MLP transition `f_ext` | |
| (2048 -> 4096 -> 2048), 5 closed-loop steps e_{t+1} = e_t + f_ext(e_t) | |
| * Perceiver-style context m_t (mean-pooled internal history -> 4 latent queries in a 256-d bottleneck attending to the | |
| external prefix <= t, output projected back and added to the last internal state), gated causal injection | |
| z_{t+1} = (1 - g) (z_t + f_int(z_t)) + g m_t with g = sigmoid(MLP(e_t)) (vector gate; `override` clamps it as in Table 3) | |
| * ST-GCN skeleton decoder: Linear(D -> 136 x 32) -> 2 STGCNBlocks (normalised HALPE-136 adjacency protocol.EDGES_ALL, temporal | |
| conv k = 3 over the 5 rolled-out steps, residual, LayerNorm) -> Linear(32 -> 2) | |
| REPLACED | |
| * internal stream input: PartTokenizer(d = 256) -> 4 part tokens per observed frame (body / face / lhand / rhand; per joint | |
| (dx, dy)/SCALE to the last observed pose, x/W, y/H, conf; part + step embeddings — the tokenizer of cwm3d_kp, unchanged), | |
| concatenated to ONE latent per frame (D_int = 4 x 256 = 1024, the Driver-WM "one latent per step" structure); causal pre-norm | |
| TransformerEncoderLayer(1024, 4 heads, d_ff 4096) as the internal pre-encoder; f_int 1024 -> 4096 -> 1024; the context output, | |
| the gate output and the decoder input are D_int-wide (the external stream stays 2048-d). | |
| * decoder output: RESIDUAL over the last observed pose in pixels, pred = last + SCALE * Linear(32 -> 2)(.), the Linear | |
| ZERO-initialised -> the untrained model is exactly zero-velocity (aide/tests/test_arch_attribution.py); the reproduction's | |
| sigmoid absolute normalised coordinates (Eq. 8) are not used. | |
| DROPPED | |
| * L_lat^vel (its target was the frozen-feature encoder on the FUTURE in-cabin features, which the keypoint stream does not have), | |
| L_aux / semantic heads (the comparison target cwm3d_kp has none), the paper's L_skel / L_phys weighting: the model returns a plain | |
| tensor and is trained with train_kp's `--loss euc` exactly like MotionBERT-kp / ST-GCN-kp (shared recipe). | |
| forward(obs [B,5,136,3] px+conf, vis [B,5,4,2048] fp16/32 (only the 3 exterior views are read; the in-cabin view is ignored), | |
| override=None) -> pred [B,5,136,2] px. Leak guard: exactly T_OBS frames of obs and vis are accepted. | |
| """ | |
| 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.driverwm_repro import D_FEAT, EXT_VIEWS, STGCNBlock, causal_mask, normalized_adjacency, res_mlp # noqa: E402 | |
| from aide.models.heads import SCALE, PartTokenizer # noqa: E402 | |
| class PerceiverCtxKP(nn.Module): | |
| """driverwm_repro.PerceiverCtx with separate internal (D_int) and external (D_ext) widths.""" | |
| def __init__(self, D_int, D_ext, n_latents=4, b=256, heads=4): | |
| super().__init__() | |
| self.n, self.b = n_latents, b | |
| self.q = nn.Sequential(nn.LayerNorm(D_int), nn.Linear(D_int, n_latents * b)) | |
| self.kv_ln = nn.LayerNorm(D_ext) | |
| self.k, self.v = nn.Linear(D_ext, b), nn.Linear(D_ext, b) | |
| self.attn = nn.MultiheadAttention(b, heads, batch_first=True) | |
| self.out = nn.Linear(n_latents * b, D_int) | |
| def forward(self, int_hist, ext_hist): # [B,t,D_int], [B,t,D_ext] -> m_t [B,D_int] | |
| B = int_hist.shape[0] | |
| q = self.q(int_hist.mean(1)).view(B, self.n, self.b) | |
| kv = self.kv_ln(ext_hist) | |
| a, _ = self.attn(q, self.k(kv), self.v(kv), need_weights=False) | |
| return int_hist[:, -1] + self.out(a.reshape(B, -1)) | |
| class SkeletonDecoderResidual(nn.Module): | |
| """driverwm_repro.SkeletonDecoder with a zero-initialised residual read-out (no sigmoid): returns the residual in SCALE units.""" | |
| def __init__(self, D, J=P.J, c=32, blocks=2, edges=P.EDGES_ALL): | |
| super().__init__() | |
| self.J, self.c = J, c | |
| self.inp = nn.Linear(D, J * c) | |
| A = normalized_adjacency(J, edges) | |
| self.blocks = nn.ModuleList(STGCNBlock(c, A) for _ in range(blocks)) | |
| self.out = nn.Linear(c, 2) | |
| nn.init.zeros_(self.out.weight) # exactly zero-velocity at init | |
| nn.init.zeros_(self.out.bias) | |
| def forward(self, z): # [B,H,D] -> [B,H,J,2] | |
| x = self.inp(z).view(z.shape[0], z.shape[1], self.J, self.c) | |
| for b in self.blocks: | |
| x = b(x) | |
| return self.out(x) | |
| class DriverWMKP(nn.Module): | |
| name = "driverwm_kp" | |
| # train_kp --init-parts trunk (P2, docs/decisions.md 2026-09-15 "Driver-WM architecture on CockpitWM tokens + BATON pretraining"): the trunk is | |
| # EVERYTHING except the zero-initialised residual read-out `dec.out.*` — internal-stream tokenizer (tok.), exterior view / step embeddings | |
| # (view_emb, pos_ext), both causal pre-encoders (pre_int., pre_ext.), the transitions = dynamics core (f_int., f_ext.), the Perceiver context | |
| # (ctx.), the gate (gate.), the decoder input projection and ST-GCN blocks (dec.inp., dec.blocks.). Keeping dec.out at its zero init means a | |
| # trunk-initialised fine-tune still starts at exact zero-velocity (aide/tests/test_driverwm_pretrain.py). | |
| TRUNK_PREFIXES = ("tok.", "view_emb", "pos_ext", "pre_int.", "pre_ext.", "f_int.", "f_ext.", "ctx.", "gate.", "dec.inp.", "dec.blocks.") | |
| def __init__(self, d=256, D_ext=D_FEAT, heads=4, d_ff=4096, n_latents=4, ctx_bottleneck=256, gcn_ch=32, dropout=0.0): | |
| super().__init__() | |
| self.tok = PartTokenizer(d) # == cwm3d_kp's tokenizer (body / face / shared hand Linear + side emb) | |
| self.n_part = len(P.PARTS) | |
| self.D, self.D_ext = self.n_part * d, D_ext # internal latent 1024-d, external 2048-d | |
| self.view_emb = nn.Parameter(torch.zeros(len(P.VIEWS), D_ext)) | |
| self.pos_ext = nn.Parameter(torch.zeros(P.N_SAMPLE, D_ext)) | |
| nn.init.normal_(self.view_emb, std=0.02) | |
| nn.init.normal_(self.pos_ext, std=0.02) | |
| self.pre_int = nn.TransformerEncoderLayer(self.D, heads, d_ff, dropout=dropout, batch_first=True, norm_first=True) | |
| self.pre_ext = nn.TransformerEncoderLayer(D_ext, heads, d_ff, dropout=dropout, batch_first=True, norm_first=True) | |
| self.f_int, self.f_ext = res_mlp(self.D, d_ff), res_mlp(D_ext, d_ff) | |
| self.ctx = PerceiverCtxKP(self.D, D_ext, n_latents, ctx_bottleneck, heads) | |
| self.gate = nn.Sequential(nn.LayerNorm(D_ext), nn.Linear(D_ext, D_ext), nn.GELU(), nn.Linear(D_ext, self.D)) | |
| self.dec = SkeletonDecoderResidual(self.D, P.J, gcn_ch) | |
| # ------------------------------------------------------------------ streams | |
| def encode_int(self, obs): # [B,T,J,3] -> causal latents [B,T,D] | |
| B, T = obs.shape[:2] | |
| z = self.tok(obs).reshape(B, T, self.D) # 4 part tokens (with part + step embeddings) -> one latent per frame | |
| return self.pre_int(z, src_mask=causal_mask(T, z.device)) | |
| def encode_ext(self, f_out): # [B,T,3,D_ext] -> pooled causal latents [B,T,D_ext] (== reproduction) | |
| T = f_out.shape[1] | |
| z = (f_out.float() + self.view_emb[EXT_VIEWS]).mean(2) + self.pos_ext[:T] | |
| return self.pre_ext(z, src_mask=causal_mask(T, z.device)) | |
| def rollout(self, zi, ze, override=None): # == reproduction (Eq. 5-7 / 18), D_int-wide internal state | |
| int_hist, ext_hist = zi, ze | |
| z_t, e_t = zi[:, -1], ze[:, -1] | |
| zs, es, gs = [], [], [] | |
| for _ in range(P.T_PRED): | |
| m_t = self.ctx(int_hist, ext_hist) | |
| z_tilde = z_t + self.f_int(z_t) | |
| g = torch.sigmoid(self.gate(e_t)) | |
| if override is not None: | |
| g = torch.full_like(g, float(override)) | |
| z_next = (1.0 - g) * z_tilde + g * m_t | |
| e_next = e_t + self.f_ext(e_t) | |
| zs.append(z_next), es.append(e_next), gs.append(g.mean(-1)) | |
| int_hist = torch.cat([int_hist, z_next[:, None]], 1) | |
| ext_hist = torch.cat([ext_hist, e_next[:, None]], 1) | |
| z_t, e_t = z_next, e_next | |
| return torch.stack(zs, 1), torch.stack(es, 1), torch.stack(gs, 1) | |
| # ------------------------------------------------------------------ forward | |
| def forward(self, obs, vis=None, override=None): | |
| assert obs.shape[1] == P.T_OBS and obs.shape[2] == P.J, obs.shape | |
| B, T = obs.shape[:2] | |
| if vis is None: # shape tests only; training / evaluation always pass the exterior features | |
| vis = obs.new_zeros(B, T, len(P.VIEWS), self.D_ext) | |
| assert vis.dim() == 4 and vis.shape[1] == P.T_OBS and vis.shape[2] == len(P.VIEWS), vis.shape | |
| last = obs[:, -1:, :, :2] | |
| zi, ze = self.encode_int(obs), self.encode_ext(vis[:, :, EXT_VIEWS]) | |
| z_fut, _, _ = self.rollout(zi, ze, override) | |
| return last + SCALE * self.dec(z_fut) # residual over the last observed pose (shared recipe) | |
| def build(name: str, **kw) -> nn.Module: | |
| if name == "driverwm_kp": | |
| return DriverWMKP(**kw) | |
| raise ValueError(name) | |
| if __name__ == "__main__": | |
| m = DriverWMKP() | |
| print("params", sum(p.numel() for p in m.parameters())) | |
| x = torch.rand(2, P.T_OBS, P.J, 3) * torch.tensor([P.W, P.H, 1.0]) | |
| v = torch.randn(2, P.T_OBS, len(P.VIEWS), D_FEAT) | |
| y = m(x, v) | |
| print("out", tuple(y.shape), "init == zero-velocity:", torch.equal(y, x[:, -1:, :, :2].expand_as(y))) | |
Xet Storage Details
- Size:
- 10.2 kB
- Xet hash:
- e378c8c7690c0d49f8d4b5e98028575076d0d46bd1c3dac7fc87db2156658984
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.