Buckets:
| """Driver-WM reimplementation on frozen Qwen3-VL-2B features (arXiv 2605.05092v2; no official code). | |
| Port of /home/USER/Drive/dm_bench/loader/bench/models.py::DriverWM to the AIDE interface, made | |
| faithful to the paper where dm_bench deviated (assumptions C1-C11 in aide/docs/protocol.md): | |
| inputs vis [B,5,4,2048] fp16/32 pooled Qwen features of the 5 OBSERVED protocol frames, views in | |
| protocol.VIEWS order (front, left, right, incar). Keypoints are TARGETS only (never inputs). | |
| interface identity 2048-d latents + learned view embedding (Eq. 3-4); external = mean over the 3 | |
| out-cabin views (App. B.7); one causal temporal self-attention layer per stream (4 heads, | |
| d_ff 4096, pre-norm) as the pre-encoding (Sec. 3.3). | |
| rollout 5 closed-loop steps (Eq. 1, 5-7): m_t = Ctx(Z_int<=t, Z_ext<=t) (Perceiver-style, App. B.6: | |
| mean-pooled internal history -> MLP -> 4 latent queries in a 256-d bottleneck attending to the | |
| external prefix; output projected to a global vector and added to the last internal state), | |
| z~_{t+1} = z_t + f_int(z_t) (residual MLP), g_t = sigmoid(MLP_g(e_t)) (vector gate on the current | |
| pooled external latent, Eq. 18), z_{t+1} = (1-g) z~ + g m_t, e_{t+1} = e_t + f_ext(e_t). | |
| `override` replaces g by the scalar lambda_CA (Table 3 interventions do(lambda_CA = 0/1/2)). | |
| decoder D_skel: Linear(2048 -> 136 x 32) -> 2 ST-GCN blocks (spatial graph conv over the HALPE-136 | |
| adjacency protocol.EDGES_ALL, temporal conv k=3 over the 5 rolled-out steps, residual + LN) | |
| -> Linear(32 -> 2) -> sigmoid = normalised coordinates (Eq. 8); pixels = (x W, y H). No | |
| residual anchor (the paper decodes absolute coordinates; dm_bench's anchor is NOT used). | |
| heads DBR / DER from the last rolled-out internal state, TCR / VCR from the last rolled-out external | |
| state (Eq. 13), LayerNorm + Linear. | |
| loss Eq. 15 with lambda_lat = lambda_phys = 0.1, lambda_skel = lambda_aux = 1: | |
| L_lat^vel (Eq. 14; targets = the SAME internal pre-encoder run causally over the 10-frame | |
| in-cabin sequence [observed; GT future], detached), L_skel (Eq. 23: per-frame conf-masked MSE | |
| in normalised coordinates, conf >= 0.1, summed over the 5 steps), L_aux (4 CE), L_phys = L_bone | |
| (Eq. 25, body + hand edges, normalised coordinates) + L_smooth (Eq. 26, sequence [GT last | |
| observed; predicted 1..5]). Sums over t and squared vector norms as written; batch mean. | |
| `lat_norm="mean"` divides the latent term by D (sensitivity of the port, not the paper). | |
| Leak guard: forward() accepts exactly T_OBS frames; the future in-cabin features enter only loss(). | |
| """ | |
| from __future__ import annotations | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) | |
| from aide.data import protocol as P # noqa: E402 | |
| D_FEAT = 2048 | |
| EXT_VIEWS = [P.VIEWS.index(v) for v in ("front", "left", "right")] | |
| IN_VIEW = P.VIEWS.index("incar") | |
| _WH = torch.tensor([float(P.W), float(P.H)]) | |
| _EDGES = torch.as_tensor(P.EDGES_BODY_HANDS, dtype=torch.long) | |
| def causal_mask(L, device): | |
| return torch.triu(torch.ones(L, L, dtype=torch.bool, device=device), 1) | |
| def normalized_adjacency(J, edges): | |
| A = torch.zeros(J, J) | |
| for i, j in edges: | |
| A[i, j] = A[j, i] = 1.0 | |
| A = A + torch.eye(J) | |
| d = A.sum(1).pow(-0.5) | |
| return d[:, None] * A * d[None, :] | |
| def res_mlp(D, hidden): | |
| return nn.Sequential(nn.LayerNorm(D), nn.Linear(D, hidden), nn.GELU(), nn.Linear(hidden, D)) | |
| class STGCNBlock(nn.Module): | |
| def __init__(self, c, A, k_t=3): | |
| super().__init__() | |
| self.register_buffer("A", A) | |
| self.gcn = nn.Linear(c, c) | |
| self.tcn = nn.Conv2d(c, c, (k_t, 1), padding=(k_t // 2, 0)) | |
| self.ln = nn.LayerNorm(c) | |
| def forward(self, x): # [B,H,K,c] | |
| h = F.gelu(torch.einsum("kj,bhjc->bhkc", self.A, self.gcn(x))) | |
| h = self.tcn(h.permute(0, 3, 1, 2)).permute(0, 2, 3, 1) # temporal conv over the rollout steps | |
| return self.ln(x + h) | |
| class SkeletonDecoder(nn.Module): | |
| 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) | |
| def forward(self, z): # [B,H,D] -> [B,H,J,2] in (0,1) | |
| x = self.inp(z).view(z.shape[0], z.shape[1], self.J, self.c) | |
| for b in self.blocks: | |
| x = b(x) | |
| return torch.sigmoid(self.out(x)) | |
| class PerceiverCtx(nn.Module): | |
| def __init__(self, D, n_latents=4, b=256, heads=4): | |
| super().__init__() | |
| self.n, self.b = n_latents, b | |
| self.q = nn.Sequential(nn.LayerNorm(D), nn.Linear(D, n_latents * b)) | |
| self.kv_ln = nn.LayerNorm(D) | |
| self.k, self.v = nn.Linear(D, b), nn.Linear(D, b) | |
| self.attn = nn.MultiheadAttention(b, heads, batch_first=True) | |
| self.out = nn.Linear(n_latents * b, D) | |
| def forward(self, int_hist, ext_hist): # [B,t,D], [B,t,D] -> m_t [B,D] | |
| 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 DriverWMRepro(nn.Module): | |
| name = "driverwm_repro" | |
| def __init__(self, D=D_FEAT, heads=4, d_ff=4096, n_latents=4, ctx_bottleneck=256, gcn_ch=32, | |
| use_semantic=True, conf_min=0.1, dropout=0.0, lat_norm="sum"): | |
| super().__init__() | |
| assert lat_norm in ("sum", "mean") | |
| self.D, self.conf_min, self.lat_norm, self.use_semantic = D, conf_min, lat_norm, use_semantic | |
| self.view_emb = nn.Parameter(torch.zeros(len(P.VIEWS), D)) | |
| self.pos = nn.Parameter(torch.zeros(P.N_SAMPLE, D)) | |
| nn.init.normal_(self.view_emb, std=0.02) | |
| nn.init.normal_(self.pos, std=0.02) | |
| self.pre_int = nn.TransformerEncoderLayer(D, heads, d_ff, dropout=dropout, batch_first=True, norm_first=True) | |
| self.pre_ext = nn.TransformerEncoderLayer(D, heads, d_ff, dropout=dropout, batch_first=True, norm_first=True) | |
| self.f_int, self.f_ext = res_mlp(D, d_ff), res_mlp(D, d_ff) | |
| self.ctx = PerceiverCtx(D, n_latents, ctx_bottleneck, heads) | |
| self.gate = nn.Sequential(nn.LayerNorm(D), nn.Linear(D, D), nn.GELU(), nn.Linear(D, D)) | |
| self.dec = SkeletonDecoder(D, P.J, gcn_ch) | |
| if use_semantic: | |
| self.sem = nn.ModuleDict({t: nn.Sequential(nn.LayerNorm(D), nn.Linear(D, len(P.CLASSES[t]))) for t in P.TASKS}) | |
| # ------------------------------------------------------------------ streams | |
| def encode_int(self, f_in): # [B,t,D] -> causal latents [B,t,D] | |
| t = f_in.shape[1] | |
| z = f_in.float() + self.view_emb[IN_VIEW] + self.pos[:t] | |
| return self.pre_int(z, src_mask=causal_mask(t, z.device)) | |
| def encode_ext(self, f_out): # [B,t,3,D] -> pooled causal latents [B,t,D] | |
| t = f_out.shape[1] | |
| z = (f_out.float() + self.view_emb[EXT_VIEWS]).mean(2) + self.pos[:t] | |
| return self.pre_ext(z, src_mask=causal_mask(t, z.device)) | |
| def rollout(self, zi, ze, override=None): | |
| 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) # Eq. 5 (prefix <= t only) | |
| z_tilde = z_t + self.f_int(z_t) # Eq. 6 | |
| g = torch.sigmoid(self.gate(e_t)) # Eq. 7 / 18 | |
| 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, vis, override=None): | |
| assert vis.dim() == 4 and vis.shape[1] == P.T_OBS and vis.shape[2] == len(P.VIEWS), vis.shape | |
| f_in, f_out = vis[:, :, IN_VIEW], vis[:, :, EXT_VIEWS] | |
| zi, ze = self.encode_int(f_in), self.encode_ext(f_out) | |
| z_fut, e_fut, g = self.rollout(zi, ze, override) | |
| s = self.dec(z_fut) # [B,H,J,2] normalised | |
| out = {"pred": s * _WH.to(s.device), "pred_norm": s, "z_obs_last": zi[:, -1], "z_fut": z_fut, | |
| "e_fut": e_fut, "gate_mean": g, "f_in_obs": f_in.detach(), "last": None} | |
| if self.use_semantic: | |
| for t in P.TASKS: | |
| out[f"logits_{t}"] = self.sem[t](z_fut[:, -1] if t in ("dbr", "der") else e_fut[:, -1]) | |
| return out | |
| # ------------------------------------------------------------------ loss (training only) | |
| def loss(self, out, fut_xy, fut_conf, vis_fut, labels, lam=None): | |
| lam = {"lat": 0.1, "skel": 1.0, "aux": 1.0, "phys": 0.1, **(lam or {})} | |
| wh = _WH.to(fut_xy.device) | |
| s_hat, s = out["pred_norm"], fut_xy / wh | |
| m = (fut_conf >= self.conf_min).float() # [B,H,J] | |
| se = ((s_hat - s) ** 2).sum(-1) | |
| l_skel = ((se * m).sum(-1) / (m.sum(-1) + 1e-6)).sum(-1).mean() # Eq. 23, sum over t | |
| with torch.no_grad(): # Eq. 14 targets | |
| z_full = self.encode_int(torch.cat([out["f_in_obs"], vis_fut[:, :, IN_VIEW].float()], 1)) | |
| z_tgt = z_full[:, P.T_OBS - 1:] # [B,6,D] last obs + future | |
| z_hat = torch.cat([out["z_obs_last"][:, None], out["z_fut"]], 1) | |
| dv = (z_hat[:, 1:] - z_hat[:, :-1]) - (z_tgt[:, 1:] - z_tgt[:, :-1]) | |
| l_lat = dv.pow(2).sum(-1).sum(-1).mean() | |
| if self.lat_norm == "mean": | |
| l_lat = l_lat / self.D | |
| e = _EDGES.to(fut_xy.device) # Eq. 25 | |
| lp = torch.linalg.norm(s_hat[..., e[:, 0], :] - s_hat[..., e[:, 1], :], dim=-1) | |
| lg = torch.linalg.norm(s[..., e[:, 0], :] - s[..., e[:, 1], :], dim=-1) | |
| l_bone = (lp - lg).abs().sum((-1, -2)).mean() | |
| s0 = out["s_obs_last_norm"] if "s_obs_last_norm" in out else None # Eq. 26 | |
| seq_hat = s_hat if s0 is None else torch.cat([s0[:, None], s_hat], 1) | |
| seq_gt = s if s0 is None else torch.cat([s0[:, None], s], 1) | |
| d1 = ((seq_hat[:, 1:] - seq_hat[:, :-1]) - (seq_gt[:, 1:] - seq_gt[:, :-1])).pow(2).sum((-1, -2)).sum(-1).mean() | |
| d2 = (seq_hat[:, 2:] - 2 * seq_hat[:, 1:-1] + seq_hat[:, :-2]).pow(2).sum((-1, -2)).sum(-1).mean() | |
| l_smooth = d1 + d2 | |
| terms = {"skel": l_skel.detach(), "lat_vel": l_lat.detach(), "bone": l_bone.detach(), "smooth": l_smooth.detach()} | |
| total = lam["skel"] * l_skel + lam["lat"] * l_lat + lam["phys"] * (l_bone + l_smooth) | |
| if self.use_semantic and labels is not None: | |
| l_aux = sum(F.cross_entropy(out[f"logits_{t}"], labels[t]) for t in P.TASKS) | |
| terms["aux"] = l_aux.detach() | |
| total = total + lam["aux"] * l_aux | |
| return total, terms | |
| def build(name: str, **kw) -> nn.Module: | |
| if name == "driverwm_repro": | |
| return DriverWMRepro(**kw) | |
| raise ValueError(name) | |
| if __name__ == "__main__": | |
| m = DriverWMRepro() | |
| print("params", sum(p.numel() for p in m.parameters())) | |
| v = torch.randn(2, P.T_OBS, len(P.VIEWS), D_FEAT) | |
| o = m(v) | |
| print({k: tuple(x.shape) for k, x in o.items() if torch.is_tensor(x)}) | |
Xet Storage Details
- Size:
- 12.2 kB
- Xet hash:
- 371aa682c9ad63593643062edabd211c351a8d86e5a31e83feb3877e592eba5a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.