# Written by yams_any4d agent, 2026-07-11 (v2, per Cameron's spec 2026-07-11). """ V4Head v2: faithful multiview DinoVolumeSceneV4 on Any4D DiT features, instantiated PER GENERATED-VIDEO TIMESTEP. Spec (Cameron, 2026-07-11): - Multiview scene + wrist, exactly like the v4 model: per-view feature maps, per-view voxel sub-volumes concatenated (NZ = 2 x n_z), every voxel cross-view-samples BOTH views' maps, per-voxel MLP fusion. - NO temporal pooling: each of the 11 latent feature frames instantiates its own volume; each latent step predicts a 4-action chunk (flattened in time, v4-style q_head), combined 11x4 = 44 -> trimmed to 41 action steps. - Wrist voxel grid + sampling always use the FIRST frame's wrist pose (T_c2w = T_base_ee(0) @ hand_eye), matching v4's predict-chunk-from-current- view convention and staying deploy-consistent. Scene pose is the static rig calibration. - No past grip/z conditioning. - Production v4 hyperparameters mirrored from puget /tmp/launch_trains.sh: pred_size 64, n_height_bins 128 per view, grip range (0, 1), rot kmeans 256. Adaptive per-view z-range (z_hi = cam_z - margin) as in v4. Supervision from a4d_raw['action_gt_full'] (B, 41, 20) action_format_v2: [r_xyz(3), r_rot6d(6), l_xyz(3), l_rot6d(6), r_grip(1), l_grip(1)], right-arm base frame, rot6d = first two ROWS of R. NEVER read a4d_raw['action'] (zeroed under forecast task masks). """ import math import numpy as np import torch import torch.nn as nn import torch.nn.functional as F def rot6d_to_quat(r6: torch.Tensor) -> torch.Tensor: """(..., 6) rot6d (first two ROWS of R) -> (..., 4) wxyz unit quat, w >= 0.""" a, b = r6[..., 0:3], r6[..., 3:6] r1 = F.normalize(a, dim=-1) b = b - (r1 * b).sum(-1, keepdim=True) * r1 r2 = F.normalize(b, dim=-1) r3 = torch.cross(r1, r2, dim=-1) m = torch.stack([r1, r2, r3], dim=-2) m00, m11, m22 = m[..., 0, 0], m[..., 1, 1], m[..., 2, 2] w = 0.5 * torch.sqrt((1 + m00 + m11 + m22).clamp(min=1e-8)) x = 0.5 * torch.sqrt((1 + m00 - m11 - m22).clamp(min=1e-8)) y = 0.5 * torch.sqrt((1 - m00 + m11 - m22).clamp(min=1e-8)) z = 0.5 * torch.sqrt((1 - m00 - m11 + m22).clamp(min=1e-8)) x = torch.copysign(x, m[..., 2, 1] - m[..., 1, 2]) y = torch.copysign(y, m[..., 0, 2] - m[..., 2, 0]) z = torch.copysign(z, m[..., 1, 0] - m[..., 0, 1]) q = torch.stack([w, x, y, z], dim=-1) q = torch.where(q[..., 0:1] < 0, -q, q) return F.normalize(q, dim=-1) def rot6d_to_matrix(r6: torch.Tensor) -> torch.Tensor: """(..., 6) rot6d (first two ROWS of R) -> (..., 3, 3) with r1/r2/r3 as ROWS.""" a, b = r6[..., 0:3], r6[..., 3:6] r1 = F.normalize(a, dim=-1) b = b - (r1 * b).sum(-1, keepdim=True) * r1 r2 = F.normalize(b, dim=-1) r3 = torch.cross(r1, r2, dim=-1) return torch.stack([r1, r2, r3], dim=-2) class V4Head(nn.Module): def __init__( self, prep_path: str, d_in: int = 2048, feat_dim: int = 32, pred_size: int = 64, n_z: int = 128, # per view (production --n_height_bins 128) n_gripper_bins: int = 32, n_rot_clusters: int = 256, # production rot kmeans 256 height_enc_dim: int = 16, d_model: int = 256, t_out: int = 41, chunk: int = 4, # actions per latent feature frame (VAE temporal ratio) cam_margin_m: float = 0.01, use_wrist: bool = True, # False = scene-only derisk mode (2026-07-11): no wrist # sub-volume / sampling / cls; NZ = n_z. temporal_mode: str = 'per_step', # 'per_step' (11 volumes, 4-action chunks) or # 'stack_mlp' (Cameron 2026-07-11: concat Tl frames per # patch -> 4-layer res-MLP -> feat_dim -> ONE volume, # v4-style time-flattened queries) n_lat_frames: int = 11, # Tl expected by stack_mlp input proj stack_hidden: int = 512, input_channel_norm: bool = False, # per-sample, per-channel standardization of the # incoming DiT features (ONLINE feat-norm): raw DiT feats # carry a big per-channel DC bias + massive-activation # "sink" channels that make features content-blind; this # removes them without any offline extraction/stats pass. ): super().__init__() self.use_wrist = use_wrist self.temporal_mode = temporal_mode self.input_channel_norm = input_channel_norm prep = np.load(prep_path) for name in ['K_scene', 'T_w2c_scene', 'K_wrist', 'hand_eye', 'rot_centroids', 'grip_range', 'z_range', 'scene_save_wh', 'wrist_save_wh']: self.register_buffer(name, torch.from_numpy(prep[name]).float()) assert self.rot_centroids.shape[0] == n_rot_clusters self.feat_dim, self.pred_size, self.n_z = feat_dim, pred_size, n_z self.n_gripper_bins, self.n_rot_clusters = n_gripper_bins, n_rot_clusters self.height_enc_dim, self.t_out, self.chunk = height_enc_dim, t_out, chunk self.cam_margin_m = cam_margin_m self.n_lat = (t_out + chunk - 1) // chunk # 11 for 41/4 # NOTE: names must avoid peft LoRA target suffixes (q_proj etc.). # Shared across views + latent frames, like v4's shared feat_head/refine. self.vh_proj = nn.Conv2d(d_in, feat_dim, 1) if temporal_mode == 'stack_mlp': # time-flattened per-patch vector (Tl*D) -> res-MLP -> feat_dim self.vh_stack_in = nn.Linear(n_lat_frames * d_in, stack_hidden) self.vh_stack_blocks = nn.ModuleList([ nn.Sequential(nn.Linear(stack_hidden, stack_hidden), nn.GELU(), nn.Linear(stack_hidden, stack_hidden)) for _ in range(4)]) self.vh_stack_out = nn.Linear(stack_hidden, feat_dim) self.vh_refine = nn.Sequential( # FiveLayerCNN (model_hires) nn.Conv2d(feat_dim, feat_dim, 3, padding=1), nn.GELU(), nn.Conv2d(feat_dim, feat_dim, 3, padding=1), nn.GELU(), nn.Conv2d(feat_dim, feat_dim, 3, padding=1), nn.GELU(), nn.Conv2d(feat_dim, feat_dim, 3, padding=1), nn.GELU(), nn.Conv2d(feat_dim, feat_dim, 3, padding=1), ) _n_view_feats = 2 if use_wrist else 1 self.vh_voxel_mlp = nn.Sequential( # v4: concat views + height PE nn.Linear(_n_view_feats * feat_dim + height_enc_dim, 64), nn.GELU(), nn.Linear(64, 64), nn.GELU(), nn.Linear(64, feat_dim), ) # global query input: eef_feat + scene cls (+ wrist cls if use_wrist) self.vh_input = nn.Linear((2 + int(use_wrist)) * feat_dim, d_model) hidden = d_model * 2 self.vh_global = nn.Sequential( nn.Linear(d_model, hidden), nn.GELU(), nn.Linear(hidden, hidden), nn.GELU(), nn.Linear(hidden, hidden), nn.GELU(), ) T_full = self.n_lat * chunk # 44 self.vh_q = nn.Linear(hidden, T_full * feat_dim) self.vh_grip = nn.Linear(hidden, T_full * n_gripper_bins) self.vh_rot = nn.Linear(hidden, T_full * n_rot_clusters) # ------------------------------------------------------------------ def _scaled_K(self, K, save_wh, pixel_hw, device, dtype): (h, w) = pixel_hw K = K.clone().to(device).to(dtype) K[0] = K[0] * (w / float(save_wh[0])) K[1] = K[1] * (h / float(save_wh[1])) return K def _height_pe(self, bin_centers): """(..., Z) bin centres in metres -> (..., Z, He); GLOBAL-span normalization (v4: same PE basis at a given physical height regardless of view/batch).""" (z_lo, z_hi) = (float(self.z_range[0]), float(self.z_range[1])) z_norm = (bin_centers - z_lo) / max(z_hi - z_lo, 1e-6) Kf = self.height_enc_dim // 2 freqs = 2.0 ** torch.arange(Kf, device=bin_centers.device, dtype=bin_centers.dtype) args = z_norm.unsqueeze(-1) * freqs * math.pi return torch.cat([torch.sin(args), torch.cos(args)], dim=-1) def _unproject(self, K, T_w2c, bin_centers, pixel_hw): """Per-batch grid: K (3,3) or (B,3,3); T_w2c (B,4,4); bin_centers (B,Z). Returns (B, Z, P, P, 3) world voxel positions. fp32.""" (h, w) = pixel_hw P = self.pred_size B, Z = bin_centers.shape dev, dt = bin_centers.device, bin_centers.dtype if K.ndim == 2: K = K.unsqueeze(0).expand(B, 3, 3) T_c2w = torch.inverse(T_w2c) u = (torch.arange(P, device=dev, dtype=dt) + 0.5) * (w / P) v = (torch.arange(P, device=dev, dtype=dt) + 0.5) * (h / P) vv, uu = torch.meshgrid(v, u, indexing='ij') pix_h = torch.stack([uu, vv, torch.ones_like(uu)], dim=-1) # (P,P,3) ray_cam = torch.einsum('bij,hwj->bhwi', torch.inverse(K), pix_h) # (B,P,P,3) R, cam_o = T_c2w[:, :3, :3], T_c2w[:, :3, 3] ray_w = torch.einsum('bij,bhwj->bhwi', R, ray_cam) rz = ray_w[..., 2] rz = torch.where(rz.abs() < 1e-6, torch.full_like(rz, 1e-6), rz) lam = (bin_centers.view(B, Z, 1, 1) - cam_o[:, 2].view(B, 1, 1, 1)) / rz.unsqueeze(1) world = cam_o.view(B, 1, 1, 1, 3) + lam.unsqueeze(-1) * ray_w.unsqueeze(1) return world # (B,Z,P,P,3) @staticmethod def _project(K, T_w2c, xyz): """K (3,3)|(B,3,3), T_w2c (B,4,4), xyz (B,...,3) -> (uv (B,...,2), camz).""" B = T_w2c.shape[0] if K.ndim == 2: K = K.unsqueeze(0).expand(B, 3, 3) sh = xyz.shape pts = xyz.reshape(B, -1, 3) cam = torch.einsum('bij,bnj->bni', T_w2c[:, :3, :3], pts) + T_w2c[:, :3, 3].unsqueeze(1) z = cam[..., 2:3].clamp(min=1e-3) uvw = torch.einsum('bij,bnj->bni', K, cam) return (uvw[..., :2] / z).reshape(*sh[:-1], 2), cam[..., 2].reshape(*sh[:-1]) def _bins_for(self, cam_z): """(B,) cam heights -> per-batch adaptive bin centres (B, Z), v4 formula.""" (z_lo, z_hi) = (float(self.z_range[0]), float(self.z_range[1])) z_hi_b = torch.minimum(cam_z - self.cam_margin_m, torch.full_like(cam_z, z_hi)).clamp(min=z_lo + 0.01) idx = (torch.arange(self.n_z, device=cam_z.device, dtype=cam_z.dtype) + 0.5) / self.n_z return z_lo + idx.unsqueeze(0) * (z_hi_b - z_lo).unsqueeze(1) # ------------------------------------------------------------------ def forward(self, feats: dict, pixel_hw: tuple, action_raw: torch.Tensor): """ :param feats: {0: scene, 1: wrist} pre-final-layer DiT features, each (B, Tl, Hp, Wp, D). :param pixel_hw: (H, W) pixel dims of the (resized) video. :param action_raw: (B, 41, 20) UNNORMALIZED action_format_v2 GT. """ assert 0 in feats, 'head needs scene (v0) features' assert (1 in feats) or not self.use_wrist, 'use_wrist=True needs wrist (v1) features' dt = self.vh_proj.weight.dtype f0 = feats[0].to(dt) f1 = feats[1].to(dt) if self.use_wrist else None (B, Tl, Hp, Wp, D) = f0.shape P, Z, Fd = self.pred_size, self.n_z, self.feat_dim dev = f0.device act = action_raw.float() # ONLINE per-channel standardization (Cameron 2026-07-13: live frozen backbone at a # random diffusion timestep, no extraction pass). Standardize each channel over the # spatial+temporal token axes per sample -> kills the per-channel DC bias / sink-channel # scale that otherwise makes raw DiT features ~content-blind. Stateless (no buffers/BN), # so it is FSDP- and eval-safe and needs no precomputed stats. if self.input_channel_norm: def _cn(x): # x: (B, Tl, Hp, Wp, D) mu = x.mean(dim=(1, 2, 3), keepdim=True) sd = x.std(dim=(1, 2, 3), keepdim=True).clamp_min(1e-3) return (x - mu) / sd f0 = _cn(f0) if f1 is not None: f1 = _cn(f1) # ---- per-view, per-latent-frame feature maps (shared weights, v4 recipe: # 1x1 proj at token res -> bilinear up -> FiveLayerCNN at pred res) ---- def view_maps(x): tok = self.vh_proj(x.permute(0, 1, 4, 2, 3).reshape(B * Tl, D, Hp, Wp)) up = F.interpolate(tok, size=(P, P), mode='bilinear', align_corners=False) ref = self.vh_refine(up) return (tok.reshape(B, Tl, Fd, Hp, Wp), ref.reshape(B, Tl, Fd, P, P)) if self.temporal_mode == 'stack_mlp': def stack_maps(x): # (B,Tl,Hp,Wp,D) -> single fused map v = x.permute(0, 2, 3, 1, 4).reshape(B, Hp, Wp, Tl * D) h0 = self.vh_stack_in(v) for blk in self.vh_stack_blocks: h0 = h0 + blk(h0) # residual tok = self.vh_stack_out(h0).permute(0, 3, 1, 2) # (B,Fd,Hp,Wp) up0 = F.interpolate(tok, size=(P, P), mode='bilinear', align_corners=False) ref = self.vh_refine(up0) # singleton time dim keeps the rest of forward + viz unchanged return (tok.unsqueeze(1), ref.unsqueeze(1)) (tok0, maps0) = stack_maps(f0) (tok1, maps1) = stack_maps(f1) if self.use_wrist else (None, None) else: (tok0, maps0) = view_maps(f0) (tok1, maps1) = view_maps(f1) if self.use_wrist else (None, None) # ---- geometry (fp32) ---- K_s = self._scaled_K(self.K_scene, self.scene_save_wh, pixel_hw, dev, torch.float32) K_w = self._scaled_K(self.K_wrist, self.wrist_save_wh, pixel_hw, dev, torch.float32) T_w2c_s = self.T_w2c_scene.to(dev).float().unsqueeze(0).expand(B, 4, 4) # wrist pose from FIRST frame's EE pose (deploy-consistent, v4 convention) # action_format_v2 packs rot6d as the first two ROWS of the EE rotation # matrix (base<-ee), so the reconstructed row-matrix IS R_base_ee directly. # Verified via the wrist kp panel at launch; if mirrored, transpose here. R_ee0 = rot6d_to_matrix(act[:, 0, 3:9]).to(dev) # (B,3,3) T_be = torch.zeros(B, 4, 4, device=dev) T_be[:, :3, :3] = R_ee0 T_be[:, :3, 3] = act[:, 0, 0:3].to(dev) T_be[:, 3, 3] = 1.0 T_c2w_wrist = T_be @ self.hand_eye.to(dev).float().unsqueeze(0) T_w2c_w = torch.inverse(T_c2w_wrist) cam_z_s = torch.inverse(T_w2c_s)[:, 2, 3] cam_z_w = T_c2w_wrist[:, 2, 3] bins_s = self._bins_for(cam_z_s) # (B,Z) bins_w = self._bins_for(cam_z_w) grid_s = self._unproject(K_s, T_w2c_s, bins_s, pixel_hw) # (B,Z,P,P,3) if self.use_wrist: grid_w = self._unproject(K_w, T_w2c_w, bins_w, pixel_hw) voxel_pos = torch.cat([grid_s, grid_w], dim=1) # (B,NZ,P,P,3) NZ = 2 * Z else: voxel_pos = grid_s NZ = Z V = NZ * P * P vflat = voxel_pos.reshape(B, V, 3) # height PE: per-view adaptive bins, global-span basis, concat along NZ if self.use_wrist: pe = torch.cat([self._height_pe(bins_s), self._height_pe(bins_w)], dim=1) # (B,NZ,He) else: pe = self._height_pe(bins_s) pe_v = pe.unsqueeze(2).expand(B, NZ, P * P, self.height_enc_dim) \ .reshape(B, V, self.height_enc_dim).to(dt) # sampling coords are time-invariant (static grids): (B,1,V,2) per view (h, w) = pixel_hw def samp_grid(K, T_w2c): uv, camz = self._project(K, T_w2c, vflat) gx = uv[..., 0] / w * 2 - 1 gy = uv[..., 1] / h * 2 - 1 behind = camz <= 0 gx = torch.where(behind, torch.full_like(gx, 5.0), gx) gy = torch.where(behind, torch.full_like(gy, 5.0), gy) return torch.stack([gx, gy], dim=-1).unsqueeze(1).to(dt) # (B,1,V,2) g_s = samp_grid(K_s, T_w2c_s) g_w = samp_grid(K_w, T_w2c_w) if self.use_wrist else None # ---- per-latent-step volumes: cross-view sample + voxel MLP ---- voxel_feat = [] for l in range(maps0.shape[1]): s0 = F.grid_sample(maps0[:, l], g_s, mode='bilinear', padding_mode='zeros', align_corners=False).squeeze(2) parts = [s0.transpose(1, 2)] if self.use_wrist: s1 = F.grid_sample(maps1[:, l], g_w, mode='bilinear', padding_mode='zeros', align_corners=False).squeeze(2) parts.append(s1.transpose(1, 2)) vox_in = torch.cat(parts + [pe_v], dim=-1) voxel_feat.append(self.vh_voxel_mlp(vox_in)) # (B,V,Fd) # ---- global query path (latent-0 features, v4-style) ---- ee0_uv, _ = self._project(K_s, T_w2c_s, act[:, 0:1, 0:3].to(dev)) # (B,1,2) gx = (ee0_uv[:, 0, 0] / w * 2 - 1).clamp(-1, 1) gy = (ee0_uv[:, 0, 1] / h * 2 - 1).clamp(-1, 1) eef_feat = F.grid_sample(maps0[:, 0], torch.stack([gx, gy], -1).view(B, 1, 1, 2).to(dt), mode='bilinear', align_corners=False).reshape(B, Fd) # CLS analog from TOKEN-RES projections (v4 uses DINO CLS tokens — healthy # activation scale). Post-refine maps at init are ~100x smaller after spatial # mean, which killed hglob/q and stalled ALL heads (found via probe, run 6). cls0 = tok0[:, 0].mean(dim=(2, 3)) glob_in = [eef_feat, cls0] if self.use_wrist: glob_in.append(tok1[:, 0].mean(dim=(2, 3))) hglob = self.vh_global(self.vh_input(torch.cat(glob_in, dim=-1))) T_full = self.n_lat * self.chunk q = self.vh_q(hglob).reshape(B, T_full, Fd) if self.temporal_mode == 'stack_mlp': # ONE volume (voxel_feat has a single entry), all queries dot it (v4-style) vol = torch.einsum('btc,bvc->btv', q[:, :self.t_out], voxel_feat[0]) else: vol_chunks = [] for l in range(min(Tl, self.n_lat)): ql = q[:, l * self.chunk:(l + 1) * self.chunk] # (B,chunk,Fd) vol_chunks.append(torch.einsum('btc,bvc->btv', ql, voxel_feat[l])) vol = torch.cat(vol_chunks, dim=1)[:, :self.t_out] # (B,41,V) T = vol.shape[1] vol = vol.reshape(B, T, NZ, P, P) grip_logits = self.vh_grip(hglob).reshape(B, T_full, -1)[:, :T] rot_logits = self.vh_rot(hglob).reshape(B, T_full, -1)[:, :T] # ---- targets (no grad; per-sample voxel positions) ---- with torch.no_grad(): gt_xyz = act[:, :T, 0:3].to(dev) # (B,T,3) v_sq = (vflat ** 2).sum(-1) # (B,V) g_sq = (gt_xyz ** 2).sum(-1) # (B,T) dot = torch.einsum('btc,bvc->btv', gt_xyz, vflat) tgt_vol = (g_sq.unsqueeze(-1) + v_sq.unsqueeze(1) - 2 * dot).argmin(-1) (g_lo, g_hi) = (float(self.grip_range[0]), float(self.grip_range[1])) gn = (act[:, :T, 18] - g_lo) / max(g_hi - g_lo, 1e-6) tgt_grip = (gn * self.n_gripper_bins).clamp(0, self.n_gripper_bins - 1).long().to(dev) quat = rot6d_to_quat(act[:, :T, 3:9]).to(dev) sims = torch.einsum('btq,kq->btk', quat, self.rot_centroids.float()) tgt_rot = sims.abs().argmax(-1) losses = { 'v4h_vol': F.cross_entropy(vol.reshape(B * T, V).float(), tgt_vol.reshape(B * T)), 'v4h_grip': F.cross_entropy(grip_logits.reshape(B * T, -1).float(), tgt_grip.reshape(B * T)), 'v4h_rot': F.cross_entropy(rot_logits.reshape(B * T, -1).float(), tgt_rot.reshape(B * T)), } return { 'losses': losses, 'volume_logits': vol, 'grip_logits': grip_logits, 'rot_logits': rot_logits, 'voxel_positions': voxel_pos, 'tgt_vol': tgt_vol, 'tgt_grip': tgt_grip, 'tgt_rot': tgt_rot, 'gt_xyz': gt_xyz, 'pixel_hw': pixel_hw, 'scene_feats_per_frame': tok0, # (B,Tl,Fd,Hp,Wp) token-res PCA strip 'feat_map': maps0[:, 0], # (B,Fd,P,P) scene latent-0 (back-compat) 'feat_maps_all': maps0, # (B,Tl,Fd,P,P) ALL upsampled+refined scene maps **({'feat_map_wrist': maps1[:, 0], 'T_w2c_wrist': T_w2c_w, 'K_wrist_scaled': K_w} if self.use_wrist else {}), } # ------------------------------------------------------------------ def project_world_to_pix(self, xyz, pixel_hw): """Scene-cam projection for viz: (T,3) or (B,T,3) -> uv (mean rig calib).""" one = xyz.ndim == 2 pts = xyz.unsqueeze(0) if one else xyz K = self._scaled_K(self.K_scene, self.scene_save_wh, pixel_hw, pts.device, torch.float32) Tw = self.T_w2c_scene.to(pts.device).float().unsqueeze(0).expand(pts.shape[0], 4, 4) uv, _ = self._project(K, Tw, pts.float()) return uv[0] if one else uv @staticmethod def decode_pred_xyz(vol: torch.Tensor, voxel_pos: torch.Tensor): """vol (B,T,NZ,P,P) + per-sample voxel_pos (B,NZ,P,P,3) -> (B,T,3).""" (B, T) = vol.shape[:2] idx = vol.reshape(B, T, -1).argmax(-1) # (B,T) vflat = voxel_pos.reshape(B, -1, 3) return torch.gather(vflat, 1, idx.unsqueeze(-1).expand(B, T, 3))