"""DinoVolumeSceneV4 — production multi-view fusion volumetric head. This is the model used in the main PARA paper experiments (in-distribution + OOD pickplace + towel folding). It is v2 (per-voxel MLP fusion across cross-projected view features) with two precision-focused additions: 1. **Sin/cos PE replaces learned `height_emb`.** 2. **Adaptive per-view z range** keyed off each view's camera height. What follows is the paper-grade reference for the model. If you change anything material here, also update lib/V4_DESIGN.md and re-derive the loss/throughput numbers in the paper's training-setup section. ────────────────────────────────────────────────────────────────────── ARCHITECTURE (forward pass, top to bottom) ────────────────────────────────────────────────────────────────────── Inputs (training): rgb[k] : (B, 3, S, S) — view k's image at img_size×img_size K_in[k] : (B, 3, 3) — intrinsics scaled to img_size T_w2c[k] : (B, 4, 4) — world-to-cam for view k start_pix : (B, 2) — EE pixel in primary view (legacy) target_xyz : (B, T, 3) — GT world XYZ for per-voxel gather past_grip : (B, past_n) past_z : (B, past_n) — past EE world-Z (metres) Stage 1 — per-view DINOv3 backbone Each view's RGB independently runs through DINOv3 ViT-S/16+ (frozen-init with low LR via --backbone_lr_mult). Output per view: patch grid (B, embed_dim, ph, pw) where ph = pw = S / 16 (31×31 at S=504), and CLS (B, embed_dim). Stage 1a (optional) — cross-view attention (DA3-style) If cross_view_layers > 0, run alternating local+cross-view self-attention blocks over the concatenated per-view patch grids before the feat_head. Off by default in the v4 paper configs. Stage 1b — per-pixel feature head Per view: 1×1 conv embed_dim → feat_dim, bilinear-up to pred_size, then a 5-layer 3×3 refine CNN (FiveLayerCNN from lib/model_hires). Output per view: (B, feat_dim, pred_size, pred_size). feat_dim defaults to 32. Stage 2 — per-view ADAPTIVE z range *** v4 change *** Each view k computes per-batch z_hi_k from the cam height: cam_z_k = T_c2w_k[:, 2, 3] # (B,) z_hi_k = min(global_z_hi, cam_z_k − cam_margin_m) # (B,) z_lo_k = global_z_lo (always) bin_centers_k = linspace(z_lo_k, z_hi_k, n_z) cell-centred # (B, Z) The scene cam is high and static → z_hi_scene ≈ global_z_hi → no change. The wrist cam rides the EE → z_hi_wrist = (cam_z − margin) shrinks with the EE, concentrating Z bins where the wrist can actually see. Stage 2b — sin/cos height PE *** v4 change *** For each view k, compute a per-batch (B, Z, height_enc_dim) PE table: z_norm = (bin_centers_k − global_z_lo) / (global_z_hi − global_z_lo) args[i,k] = z_norm · π · 2^k for k = 0..(height_enc_dim/2 − 1) PE = concat([sin(args), cos(args)], dim=-1) Critical: normalisation uses the GLOBAL span (not per-view). This means PE at world-Z = 0.15 m is the SAME basis vector regardless of which view sampled it. The wrist view simply covers a SUBSET of the global PE range. → cross-view consistency: scene + wrist features at matching physical heights live in the same PE coordinate system. Stage 3 — per-view voxel grid (closed-form unprojection) For view k: build (B, Z, P, P, 3) world XYZ of each (y, x, z) voxel by back-projecting pixel (u, v) through K_in[k], intersecting with the world-Z plane at bin_centers_k[z]. No autograd through this — geometry is fixed per frame. Stage 4 — cross-view feature sampling per voxel Concatenate all views' voxel grids along the Z axis: voxel_pos : (B, NZ, P, P, 3) where NZ = N · n_z Flatten to (B, V=NZ·P·P, 3). For EACH view k: project each voxel's world XYZ into K_in[k] · T_w2c[k] → pixel (u, v) bilinear-sample view_feats[k] at that pixel Concatenate the N sampled feature vectors per voxel → (B, V, N·feat_dim). Stage 5 — per-voxel MLP fusion per_voxel_input = concat([all_view_feat, h_emb], dim=-1) where h_emb is the per-view PE_z broadcast over (P, P) and concatenated along the NZ axis (so each view's NZ slabs use that view's PE_z). Run a 3-layer MLP (default voxel_mlp_hidden=64) → (B, V, feat_dim). This MLP is the SECRET SAUCE: it learns nonlinear cross-view aggregation (e.g. "scene sees object boundary AND wrist sees fingertip → high score"). Stage 6 — global query MLP eef_feat = primary-view feature map sampled at start_pix (legacy signal; small impact at convergence). cls_fused = concat over views' CLS tokens (cls_fusion='concat'). past_enc = past_encoder([past_grip, past_z]) global_in = concat([eef_feat, cls_fused, past_enc]) q = q_head( global_mlp( input_proj(global_in) ) ).reshape(B, T, feat_dim) Stage 7 — dot-product score volume_logits = einsum("btc,bvc->btv", q, voxel_feat).reshape(B, T, NZ, P, P) Stage 8 — per-voxel grip/rot At training: take voxel feature at the GT-nearest voxel. At inference: at the argmax voxel of volume_logits. Flatten over T → 4-layer MLP → (B, T, n_grip + n_rot). Stage 9 — loss (lib/model_volumetric_multiview.multiview_losses) GT voxel = argmin world-distance from target_xyz to any voxel in the combined volume. Cross-entropy over the (NZ, P, P) joint, masked-fill -1e9 for border voxels at inference. Grip/rot are CE on per-voxel-MLP outputs. ────────────────────────────────────────────────────────────────────── KEY HYPERPARAMETERS (paper config) ────────────────────────────────────────────────────────────────────── Default for the YAM paper, pickplace_redo_slow + towel_folding_redo_6_11: --img_size 504 DINOv3 ViT-S/16+ standard --pred_size 64 per-view feature-map resolution (dropped from v2's 128 for memory; w/ wrist view we have enough redundancy.) --n_height_bins 128 per view (doubled from v2's 32/view = 64 total) → adaptive wrist gets ~1 mm bins, scene ~2 mm --n_window 32 chunk length --frame_stride 4 per-chunk subsampling --batch_size 4 fits 24 GB RTX 3090 at the above settings --mv_past_n 8 past-EE conditioning depth --lr 3e-4 --backbone_lr_mult 0.1 head 3e-4, DINO 3e-5 --lr_warmup 500 --lr_schedule cosine --max_iters 20000 v2 ran 10k; v4 benefits from 20k @ adaptive z cam_margin_m = 0.01 voxel grid stays this far below cam height ────────────────────────────────────────────────────────────────────── WHY THESE CHANGES MATTER (paper bullets) ────────────────────────────────────────────────────────────────────── a) Height precision. v2's learned `height_emb` is random-init from step 0 and needs many steps × bins of supervision before it discovers useful features. sin/cos PE is informative from step 0 and generalises across z_ranges without retrain (we used to hack around this with --refit_z_range). At z=0.15 m the PE is the same basis vector for both views — the model learns one notion of "this height" instead of two. b) Wrist resolution. The wrist cam rides the EE. With v2's fixed z_range = (0, 0.4) m and 32 bins/view, the wrist spent ~25 bins on heights ABOVE its own camera, which are physically unreachable from that view. v4 caps wrist z_hi at cam_z − 1 cm, so the wrist's full 128 bins land in (0, cam_z) → ~1 mm bin width when the EE is low. Combined with the shared PE basis this delivers sub-millimetre height inference without any decoder change. c) Architecture vs. cost. v4 has the SAME total voxel count as v2 (NZ · P² = 256 · 64² = 1.05 M vs. v2's 64 · 128² = 1.05 M). The uniform-baseline CE is therefore identical (log(V·T) ≈ 17.3), so loss numbers are directly comparable. v4 ran at 2.20 it/s on a 3090 vs. v2's ~2 it/s. ────────────────────────────────────────────────────────────────────── RELATIONSHIP TO OTHER VERSIONS ────────────────────────────────────────────────────────────────────── v1 single-view DinoVolumeScene (lib/model.py). v2 multi-view MLP fusion (lib/model_volumetric_multiview.py) — direct ancestor of v4; same forward shape, learned height_emb, fixed z range. v3 factored sum-of-views (lib/model_volumetric_v3.py). EXPLICITLY ABANDONED. The product-of-experts head has no cross-view nonlinear fusion → systematically worse than v2/v4 on real rollouts. Dead end. v4 THIS FILE. v2 + sin/cos PE + adaptive per-view z. Production model. ────────────────────────────────────────────────────────────────────── CKPT COMPATIBILITY ────────────────────────────────────────────────────────────────────── model_kind = "dino_volume_scene_v4". inference.py handles this kind explicitly. The bin_centers_z buffer kept on the module is the GLOBAL (unused at forward — per-view bin centres are computed from cam pose), and exists for legacy decode-path code that reads it. Original v2 docstring follows for historical reference: --- DinoVolumeSceneVolumetricMultiView — multi-view fusion volumetric model. For N input views (e.g. ``["scene"]``, ``["wrist"]``, or ``["scene", "wrist"]``): 1. Run each view's RGB through the (shared) DINOv3 backbone + feat_head → per-pixel feature map at ``pred_size × pred_size`` (default 128×128). 2. Build one world-space voxel grid PER view by unprojecting that view's pixel grid at Z height-bin centers. ``N × P × P × Z`` voxels total. 3. For every voxel in the combined grid: project its world XYZ into EACH view, bilinear-sample that view's feature map, and concatenate all N view-features. Voxels outside a view's frustum just use zero-padding from grid_sample with ``padding_mode='zeros'``. 4. Concatenate the learnable per-bucket height embedding (16-d) and run a 3-layer MLP per voxel → ``feat_dim``. 5. Flat global MLP produces per-T queries; dot product against the ``N×P×P×Z`` voxel-feature volume → ``volume_logits``. Loss target = argmin world-distance from GT-XYZ to any voxel (across all N sub-volumes). This is the multi-view generalization of ``DinoVolumeSceneVolumetric``. When ``len(views) == 1`` and ``views[0]`` is the scene cam, it's functionally identical to the single-view volumetric model with the exception that the voxel grid is N×P×P×Z = 1×P×P×Z = P×P×Z. """ from __future__ import annotations import math import torch import torch.nn as nn import torch.nn.functional as F from .model_hires import _DinoHiresBase from .model import DEFAULT_DINO_REPO, DEFAULT_DINO_WEIGHTS class _ViTBlock(nn.Module): """Pre-norm transformer block: norm → multi-head self-attention → norm → MLP. The attention call uses ``F.scaled_dot_product_attention`` (flash on H100 / A100 when available). No relative position bias — we lean on the DINO patch features already carrying spatial structure. """ def __init__(self, d: int, n_heads: int = 6, mlp_ratio: float = 4.0): super().__init__() assert d % n_heads == 0 self.n_heads = n_heads self.head_dim = d // n_heads self.norm1 = nn.LayerNorm(d) self.qkv = nn.Linear(d, 3 * d, bias=True) self.proj = nn.Linear(d, d, bias=True) self.norm2 = nn.LayerNorm(d) h = int(d * mlp_ratio) self.mlp = nn.Sequential(nn.Linear(d, h), nn.GELU(), nn.Linear(h, d)) def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, L, D) B, L, D = x.shape h = self.norm1(x) qkv = self.qkv(h).reshape(B, L, 3, self.n_heads, self.head_dim) \ .permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # (B, H, L, D/H) attn = F.scaled_dot_product_attention(q, k, v) attn = attn.transpose(1, 2).reshape(B, L, D) x = x + self.proj(attn) x = x + self.mlp(self.norm2(x)) return x class CrossViewAttnStack(nn.Module): """DA3-style alternating within-view / cross-view self-attention over the native DINO patch grid (~31×31 per view). Each block is a plain ViT block. Layer i is cross-view if ``i % 2 == 0`` else within-view — same "token reordering" idea Depth-Anything-3 (arXiv:2511.10647) uses to share information across views without introducing a separate cross-attention module. """ def __init__(self, n_layers: int, embed_dim: int, n_heads: int = 6, mlp_ratio: float = 4.0): super().__init__() self.n_layers = n_layers self.blocks = nn.ModuleList( [_ViTBlock(embed_dim, n_heads, mlp_ratio) for _ in range(n_layers)]) # is_cross[i] toggles which axis attention spans. Start with cross # so view-context lands early, then alternate. self.is_cross = [i % 2 == 0 for i in range(n_layers)] def forward(self, patches: list[torch.Tensor]) -> list[torch.Tensor]: """Input/output: list of (B, D, P, P) patch grids, one per view.""" N = len(patches) if N == 0: return patches B, D, P, _ = patches[0].shape # Stack into (B, N, P*P, D). x = torch.stack( [p.flatten(2).transpose(1, 2) for p in patches], dim=1) for i, blk in enumerate(self.blocks): if self.is_cross[i]: x = blk(x.reshape(B, N * P * P, D)).reshape(B, N, P * P, D) else: x = blk(x.reshape(B * N, P * P, D)).reshape(B, N, P * P, D) # Unstack back to per-view (B, D, P, P). return [x[:, k].transpose(1, 2).reshape(B, D, P, P) for k in range(N)] def _v4_per_view_bin_centers(z_lo: float, z_hi_b: torch.Tensor, Z: int) -> torch.Tensor: """Per-batch bin centres for a view whose z-range upper bound varies per sample (e.g. wrist cam moves with the EE). Returns ``(B, Z)`` in m. """ device, dtype = z_hi_b.device, z_hi_b.dtype idx = (torch.arange(Z, device=device, dtype=dtype) + 0.5) / Z # (Z,) z_lo_t = torch.full_like(z_hi_b, float(z_lo)) # (B,) span = (z_hi_b - z_lo_t).clamp(min=1e-3) return z_lo_t.unsqueeze(1) + idx.unsqueeze(0) * span.unsqueeze(1) # (B, Z) def _v4_fourier_pe_meters(z_lo: float, z_hi_global: float, bin_centers_b: torch.Tensor, dim: int) -> torch.Tensor: """Sin/cos PE of bin centres in metres, normalised by the GLOBAL (z_lo, z_hi_global) span — NOT per-sample — so PE at z=0.15m is the same regardless of which view or batch entry sampled it. Wrist concentrates its bins in low z; they hit a subset of the global PE range with the same basis as scene at matching physical heights. Input ``bin_centers_b`` is (B, Z) in metres. Returns (B, Z, dim). """ assert dim % 2 == 0 K = dim // 2 device, dtype = bin_centers_b.device, bin_centers_b.dtype span = max(float(z_hi_global - z_lo), 1e-6) z_norm = (bin_centers_b - z_lo) / span # (B, Z) freqs = (2.0 ** torch.arange(K, device=device, dtype=dtype)) # (K,) args = z_norm.unsqueeze(-1) * freqs * math.pi # (B, Z, K) return torch.cat([torch.sin(args), torch.cos(args)], dim=-1) # (B, Z, dim) class DinoVolumeSceneV4(_DinoHiresBase): """N-view world-space voxel grid with shared DINO backbone + per-voxel feature fusion across views.""" def __init__( self, *, views: list[str], # e.g. ["scene"], ["wrist"], ["scene", "wrist"] n_window: int = 32, n_height_bins: int = 64, pred_size: int = 128, feat_dim: int = 32, height_enc_dim: int = 16, voxel_mlp_hidden: int = 64, d_model: int = 256, mlp_layers: int = 3, mlp_hidden_mult: float = 2.0, n_gripper_bins: int = 32, n_rot_clusters: int = 64, z_lo: float = 0.0, z_hi: float = 0.5, img_size: int = 504, cls_fusion: str = "concat", # "concat" (all-view CLS into grip/rot head) or "primary" cross_view_layers: int = 0, # 0 disables; recommended 4-5 for true cross-view fusion cross_view_heads: int = 6, per_voxel_grip_rot: bool = False, # local grip/rot via voxel feature at GT (train) / argmax (infer) voxel temporal_head: str = "flat_mlp", # "flat_mlp" (default 2026-06-10) or "temporal_cnn" temporal_cnn_layers: int = 3, # only used when temporal_head=="temporal_cnn" temporal_cnn_kernel: int = 3, temporal_mlp_hidden: int = 512, # only used when temporal_head=="flat_mlp" past_n: int = 0, # # of past (grip, z) frames concatenated into the global input past_enc_dim: int = 64, # encoder output dim for the past trajectory cam_margin_m: float = 0.01, # v4: voxel grid stays this far below the cam (per-view z_hi cap) **_unused, ): super().__init__( n_window=n_window, n_height_bins=n_height_bins, pred_size=pred_size, feat_dim=feat_dim, d_model=d_model, d_cond=128, d_sin_z=48, d_sin_t=48, n_blocks=0, n_gripper_bins=n_gripper_bins, n_rot_clusters=n_rot_clusters, dino_repo=DEFAULT_DINO_REPO, dino_weights=DEFAULT_DINO_WEIGHTS, ) self.views = list(views) self.n_views = len(views) self.height_enc_dim = height_enc_dim self.z_lo = float(z_lo); self.z_hi = float(z_hi) self.img_size = int(img_size) self.cls_fusion = cls_fusion # Past-trajectory encoder. Receives past_n frames of (grip, z) → # past_n * 2 raw scalars → encoder MLP → past_enc_dim vector that # gets concatenated into the global query input. This is the # proprio signal Cameron added 2026-06-10 so the model knows # "current gripper is closed / EE is at height X" before predicting # the next chunk. self.past_n = int(past_n) self.past_enc_dim = int(past_enc_dim) if self.past_n > 0: self.past_encoder = nn.Sequential( nn.Linear(self.past_n * 2, self.past_enc_dim), nn.GELU(), nn.Linear(self.past_enc_dim, self.past_enc_dim), ) else: self.past_encoder = None # Override input_proj for both multi-CLS fusion AND optional past # conditioning. Input is always feat_dim + (cls fan-in) + (past fan-in). cls_fan = self.n_views * self.embed_dim if cls_fusion == "concat" else self.embed_dim past_fan = self.past_enc_dim if self.past_n > 0 else 0 in_dim = feat_dim + cls_fan + past_fan if cls_fusion not in ("concat", "primary"): raise ValueError(f"cls_fusion must be 'concat' or 'primary'; got {cls_fusion!r}") if in_dim != self.input_proj.in_features: self.input_proj = nn.Linear(in_dim, d_model) # DA3-style cross-view attention over native DINO patch grid. Only # built when (a) explicitly requested AND (b) there are 2+ views; # a single-view model has nothing to cross-attend with. self.cross_view_layers = int(cross_view_layers) if self.cross_view_layers > 0 and self.n_views > 1: self.cross_view_attn = CrossViewAttnStack( n_layers=self.cross_view_layers, embed_dim=self.embed_dim, n_heads=cross_view_heads, ) else: self.cross_view_attn = None # Per-voxel grip/rot heads. At train time we teacher-force by # gathering ``voxel_feat`` at the GT-nearest voxel per timestep; # at inference we gather at the argmax voxel. Either way we get # a (B, T, F) sequence which a small 1D CNN smooths along T # before two linear heads emit grip / rot bin logits. self.per_voxel_grip_rot = bool(per_voxel_grip_rot) self.temporal_head_kind = temporal_head if self.per_voxel_grip_rot: if temporal_head == "temporal_cnn": assert temporal_cnn_kernel % 2 == 1, \ "use an odd kernel so padding preserves T" pad = temporal_cnn_kernel // 2 cnn: list[nn.Module] = [] for _ in range(int(temporal_cnn_layers)): cnn.append(nn.Conv1d(feat_dim, feat_dim, temporal_cnn_kernel, padding=pad)) cnn.append(nn.GELU()) self.temporal_cnn = nn.Sequential(*cnn) self.pv_grip_head = nn.Linear(feat_dim, n_gripper_bins) self.pv_rot_head = nn.Linear(feat_dim, n_rot_clusters) self.temporal_mlp = None elif temporal_head == "flat_mlp": # 4-layer MLP. Input = T·F flattened, output = T·(n_grip+n_rot) # flattened then split. Gives full receptive field over T AND # absolute positional information for free (input channel # index encodes (t, f) pair) — no sin/cos PE needed. in_dim = n_window * feat_dim out_dim = n_window * (n_gripper_bins + n_rot_clusters) h = int(temporal_mlp_hidden) self.temporal_mlp = nn.Sequential( nn.Linear(in_dim, h), nn.GELU(), nn.Linear(h, h), nn.GELU(), nn.Linear(h, h), nn.GELU(), nn.Linear(h, out_dim), ) self.temporal_cnn = None self.pv_grip_head = None self.pv_rot_head = None else: raise ValueError( f"temporal_head must be 'flat_mlp' or 'temporal_cnn'; " f"got {temporal_head!r}") else: self.temporal_cnn = None self.temporal_mlp = None self.pv_grip_head = None self.pv_rot_head = None # v4: bin_centers_z buffer kept for inference fallback / decode # XYZ lookup; the actual per-view per-batch bin centres are # computed inside ``forward`` from the cam pose. v4 replaces the # learned height embedding with a sin/cos PE. Z = n_height_bins centers = torch.tensor( [self.z_lo + (i + 0.5) * (self.z_hi - self.z_lo) / Z for i in range(Z)], dtype=torch.float32) self.register_buffer("bin_centers_z", centers) self.cam_margin_m = float(cam_margin_m) # Per-voxel MLP: concat all view features + height emb. d_in = self.n_views * feat_dim + height_enc_dim h = voxel_mlp_hidden self.voxel_mlp = nn.Sequential( nn.Linear(d_in, h), nn.GELU(), nn.Linear(h, h), nn.GELU(), nn.Linear(h, feat_dim), ) # Flat global MLP (same as volumetric). T = n_window hidden = int(d_model * mlp_hidden_mult) layers: list[nn.Module] = [nn.Linear(d_model, hidden), nn.GELU()] for _ in range(mlp_layers - 1): layers.extend([nn.Linear(hidden, hidden), nn.GELU()]) self.global_mlp = nn.Sequential(*layers) self.q_head = nn.Linear(hidden, T * feat_dim) self.grip_head = nn.Linear(hidden, T * n_gripper_bins) self.rot_head = nn.Linear(hidden, T * n_rot_clusters) # ──────────────────────────────────────────────────────────────────── # Geometry helpers # ──────────────────────────────────────────────────────────────────── # ──────────────────────────────────────────────────────────────────── # Backbone hook (subclass swap point) # ──────────────────────────────────────────────────────────────────── def _per_view_patch_features(self, rgb: list[torch.Tensor]): """Returns (per-view patch grids, per-view global tokens). Default DINOv3 path: loops over views, each through the (shared) backbone independently — no cross-view interaction at the backbone level (use ``CrossViewAttnStack`` for that). Override in subclasses to swap the backbone (v4_da3 uses DA3's multi-view ViT with built-in cross-attention; v4_vlm uses InternVL's vision tower). Subclass must return: - ``view_patches``: list of N tensors of shape (B, embed_dim, ph, pw) - ``view_global``: list of N tensors of shape (B, embed_dim) with ``embed_dim`` matching ``self.embed_dim`` (so feat_head and input_proj line up). ph/pw can differ from DINOv3's 31×31; the downstream 5-layer refine CNN bilinear-ups to ``pred_size``. """ view_patches: list[torch.Tensor] = [] view_cls: list[torch.Tensor] = [] for k in range(self.n_views): patch, cls = self.patch_features(rgb[k]) view_patches.append(patch) view_cls.append(cls) return view_patches, view_cls def _unproject_grid(self, K_in: torch.Tensor, T_w2c: torch.Tensor, bin_centers_b: torch.Tensor) -> torch.Tensor: """For ONE camera, return (B, Z, P, P, 3) world voxel positions. ``bin_centers_b`` is (B, Z) per-batch z bin centres in metres — lets the wrist view shrink its z range to its own cam-Z each frame. """ B = K_in.shape[0] P, Z, S = self.pred_size, self.n_z, self.img_size device = K_in.device; dtype = K_in.dtype cell = S / P u_idx = torch.arange(P, device=device, dtype=dtype) v_idx = torch.arange(P, device=device, dtype=dtype) u_pix = (u_idx + 0.5) * cell v_pix = (v_idx + 0.5) * cell vv, uu = torch.meshgrid(v_pix, u_pix, indexing="ij") ones = torch.ones_like(uu) pix_h = torch.stack([uu, vv, ones], dim=-1) # (P, P, 3) K_inv = torch.inverse(K_in) ray_cam = torch.einsum("bij,hwj->bhwi", K_inv, pix_h) T_c2w = torch.inverse(T_w2c) R = T_c2w[:, :3, :3]; cam_o = T_c2w[:, :3, 3] ray_world = torch.einsum("bij,bhwj->bhwi", R, ray_cam) rz = ray_world[..., 2] cam_oz = cam_o[..., 2].unsqueeze(-1).unsqueeze(-1) rz_safe = torch.where(rz.abs() < 1e-6, torch.full_like(rz, 1e-6), rz) z_centers = bin_centers_b.to(rz.dtype) # (B, Z) lam = ((z_centers.view(B, 1, 1, Z) - cam_oz.unsqueeze(-1)) / rz_safe.unsqueeze(-1)) world = (cam_o.view(B, 1, 1, 1, 3) + lam.unsqueeze(-1) * ray_world.unsqueeze(3)) return world.permute(0, 3, 1, 2, 4).contiguous() # (B, Z, P, P, 3) @staticmethod def _sample_view_feat(F_map: torch.Tensor, xyz_world: torch.Tensor, K_in: torch.Tensor, T_w2c: torch.Tensor, img_size: int) -> torch.Tensor: """Project ``xyz_world`` (B, V, 3) into a camera and bilinear-sample features from ``F_map`` (B, F, P, P). Returns (B, V, F). Out-of-frustum voxels (negative depth OR projected pixel outside the image) get zero features via padding_mode='zeros'. """ B, V, _ = xyz_world.shape ones = torch.ones(B, V, 1, device=xyz_world.device, dtype=xyz_world.dtype) xyz_h = torch.cat([xyz_world, ones], dim=-1) # (B, V, 4) cam = torch.einsum("bij,bvj->bvi", T_w2c, xyz_h)[..., :3] # (B, V, 3) # Mask voxels with depth ≤ 0 (behind the camera) BEFORE the projective # divide — feature stays zero for those via grid_sample's border range. z = cam[..., 2:3].clamp(min=1e-3) # (B, V, 1) uv = torch.einsum("bij,bvj->bvi", K_in, cam) / z # (B, V, 3) # Convert pixel coords → grid_sample normalized [-1, 1]. S = float(img_size) norm_u = (uv[..., 0] / S) * 2.0 - 1.0 norm_v = (uv[..., 1] / S) * 2.0 - 1.0 # Zero-out voxels that are behind the camera by pushing their # sample coord well outside [-1, 1] so padding_mode='zeros' kicks in. behind = cam[..., 2] <= 0 norm_u = torch.where(behind, torch.full_like(norm_u, 5.0), norm_u) norm_v = torch.where(behind, torch.full_like(norm_v, 5.0), norm_v) grid = torch.stack([norm_u, norm_v], dim=-1).unsqueeze(1) # (B, 1, V, 2) sampled = F.grid_sample(F_map, grid, mode="bilinear", padding_mode="zeros", align_corners=False) # (B, F, 1, V) return sampled.squeeze(2).transpose(1, 2) # (B, V, F) # ──────────────────────────────────────────────────────────────────── # Forward # ──────────────────────────────────────────────────────────────────── def forward(self, rgb: list[torch.Tensor], start_pix: torch.Tensor, K_in: list[torch.Tensor], T_w2c: list[torch.Tensor], target_xyz: torch.Tensor | None = None, border_mask_px: int = 0, past_grip: torch.Tensor | None = None, past_z: torch.Tensor | None = None) -> dict: """Lists are length N (one entry per view, in ``self.views`` order). ``rgb[k]`` : (B, 3, S, S) — view k's image at img_size×img_size ``K_in[k]`` : (B, 3, 3) — view k's intrinsics, scaled to img_size ``T_w2c[k]`` : (B, 4, 4) — view k's world-to-cam pose ``start_pix`` : (B, 2) — EE pixel in PRIMARY view (views[0]) Returns: - ``volume_logits`` (B, T, N*Z, P, P) — flat-decode-compatible - ``voxel_positions`` (B, N*Z, P, P, 3) — for loss (nearest-3D) + decode - ``grip_logits``, ``rot_logits``, ``F_refined`` (primary view), ``patch_features`` (primary view) """ assert len(rgb) == self.n_views and len(K_in) == self.n_views assert len(T_w2c) == self.n_views B, _, S, _ = rgb[0].shape T, Z, P, F_d = self.n_window, self.n_z, self.pred_size, self.feat_dim N = self.n_views # 1. Per-view backbone → raw patch grids + global tokens. # Factored into a method so subclasses can swap the backbone # (v4_da3, v4_vlm) without duplicating the full forward pass. view_patches, view_cls = self._per_view_patch_features(rgb) # 1a. Optional DA3-style cross-view attention over the native DINO # patch grid. Runs only when --cross_view_layers > 0 and N>1. if self.cross_view_attn is not None: view_patches = self.cross_view_attn(view_patches) patch0 = view_patches[0] # 1b. Per-view feat_head → bilinear-up to pred_size → 5-layer # refine CNN. Same recipe as ``_pixel_features`` but split so we # could insert the cross-view stage in between. # v4: per-view adaptive z range computed from each view's cam pose. # Scene cam is high+static so cap is inactive; wrist cam rides the # EE so its z_hi drops to just under cam_z each frame. view_bin_centers: list[torch.Tensor] = [] # each (B, Z) view_pe_z: list[torch.Tensor] = [] # each (B, Z, height_enc_dim) for k in range(N): T_c2w_k = torch.inverse(T_w2c[k]) cam_z_k = T_c2w_k[:, 2, 3] # (B,) z_hi_k = torch.minimum( cam_z_k - float(self.cam_margin_m), torch.full_like(cam_z_k, float(self.z_hi))) z_hi_k = torch.clamp(z_hi_k, min=float(self.z_lo) + 0.01) bin_b = _v4_per_view_bin_centers(self.z_lo, z_hi_k, self.n_z) pe_b = _v4_fourier_pe_meters( self.z_lo, self.z_hi, bin_b, self.height_enc_dim) view_bin_centers.append(bin_b) view_pe_z.append(pe_b) view_feats: list[torch.Tensor] = [] view_voxel_pos: list[torch.Tensor] = [] for k in range(N): F0 = self.feat_head(view_patches[k]) F1 = F.interpolate(F0, size=(self.pred_size, self.pred_size), mode="bilinear", align_corners=False) F_scene = self.refine_cnn(F1) view_feats.append(F_scene) view_voxel_pos.append(self._unproject_grid( K_in[k], T_w2c[k], view_bin_centers[k])) # Fuse the CLS tokens that feed the grip/rot head. "concat" makes # the global MLP see every view's global descriptor; "primary" # falls back to the original behaviour (matches old ckpts). cls_fused = (torch.cat(view_cls, dim=-1) if self.cls_fusion == "concat" and N > 1 else view_cls[0]) # Combined voxel positions (B, N*Z, P, P, 3) — stack along Z. voxel_pos = torch.cat(view_voxel_pos, dim=1) NZ = N * Z V = NZ * P * P # 2. For each voxel (in the combined volume), sample features from # every view. The features for view k include its OWN unprojected # grid features (direct tile, no projection error) for that view's # slice [k*Z : (k+1)*Z], and projected-sampled features for the # other views' slices. # Flatten voxel_pos for sampling. voxel_flat = voxel_pos.reshape(B, V, 3) # (B, V, 3) per_view_feat: list[torch.Tensor] = [] # each (B, V, F) for k in range(N): per_view_feat.append(self._sample_view_feat( view_feats[k], voxel_flat, K_in[k], T_w2c[k], self.img_size)) all_view_feat = torch.cat(per_view_feat, dim=-1) # (B, V, N*F) # 3. v4: per-view per-batch sin/cos PE of the world-Z bin centres # (the centres differ per view because of the adaptive z range). # Stack into (B, NZ, h_enc) → broadcast over (P, P) → flat (B, V, h_enc). # view_pe_z[k] is (B, Z, h_enc); concat across views gives (B, NZ, h_enc). h_per_view = torch.cat(view_pe_z, dim=1) # (B, NZ, h_enc) h_enc = h_per_view.shape[-1] h_emb = (h_per_view.view(B, NZ, 1, 1, h_enc) .expand(B, NZ, P, P, h_enc) .reshape(B, V, h_enc)) # (B, V, h_enc) # 4. Voxel MLP: concat [features..., height_emb] → feat_dim. voxel_in = torch.cat([all_view_feat, h_emb], dim=-1) # (B, V, N*F + h_enc) voxel_feat = self.voxel_mlp(voxel_in) # (B, V, F) # 5. Flat global MLP query — eef sampled from PRIMARY view's feature map. F_prim = view_feats[0] scale = P / S sx = (start_pix[..., 0] * scale).long().clamp(0, P - 1) sy = (start_pix[..., 1] * scale).long().clamp(0, P - 1) b_idx = torch.arange(B, device=rgb[0].device) eef_feat = F_prim[b_idx, :, sy, sx] # (B, F) # Optional past-trajectory conditioning. Both tensors are (B, past_n). # We flatten + encode + concat into the global query input so the # MLP knows "current gripper / EE-height" before predicting next chunk. global_in = [eef_feat, cls_fused] if self.past_encoder is not None: assert past_grip is not None and past_z is not None, \ "past_n > 0 but past_grip / past_z not provided" past_flat = torch.cat([past_grip, past_z], dim=-1) # (B, 2*past_n) global_in.append(self.past_encoder(past_flat)) # (B, past_enc_dim) q_in = self.input_proj(torch.cat(global_in, dim=-1)) hglob = self.global_mlp(q_in) # (B, hidden) q = self.q_head(hglob).reshape(B, T, F_d) # (B, T, F) # 6. Dot-product score per voxel. vol = torch.einsum("btc,bvc->btv", q, voxel_feat) # (B, T, V) vol = vol.view(B, T, NZ, P, P) # (B, T, NZ, P, P) # Border mask applies BEFORE both the per-voxel grip/rot argmax # AND the downstream xyz lookup so the two stay consistent. # Training never passes a positive ``border_mask_px`` so the # volume loss still sees every voxel. if border_mask_px > 0 and border_mask_px * 2 < P: keep = torch.zeros(P, P, dtype=torch.bool, device=vol.device) keep[border_mask_px:P - border_mask_px, border_mask_px:P - border_mask_px] = True vol = vol.masked_fill(~keep.view(1, 1, 1, P, P), -1e9) grip_logits = self.grip_head(hglob).reshape(B, T, self.n_gripper_bins) rot_logits = self.rot_head(hglob).reshape(B, T, self.n_rot_clusters) # Per-voxel grip/rot path — overrides the global heads above when # enabled. Same teacher-force-at-GT / argmax-at-infer pattern the # original production PARA used for grip/rot, but now gathering # from the cross-view-fused per-voxel features. if self.per_voxel_grip_rot: vol_flat = vol.reshape(B, T, V) if target_xyz is not None: # Argmin world-distance per (B, T) → GT voxel index. # d² = |a|² + |b|² − 2 a·b (avoids materialising (B,T,V,3)). v_norm_sq = (voxel_flat * voxel_flat).sum(-1) # (B, V) t_norm_sq = (target_xyz * target_xyz).sum(-1) # (B, T) dot = torch.einsum("btc,bvc->btv", target_xyz, voxel_flat) dist_sq = (t_norm_sq.unsqueeze(-1) + v_norm_sq.unsqueeze(1) - 2 * dot) # (B, T, V) idx_t = dist_sq.argmin(dim=2) # (B, T) else: idx_t = vol_flat.argmax(dim=2) # (B, T) gather_idx = idx_t.unsqueeze(-1).expand(B, T, F_d) temp_feat = torch.gather(voxel_feat, 1, gather_idx) # (B, T, F) if self.temporal_head_kind == "flat_mlp": flat_in = temp_feat.reshape(B, T * F_d) flat_out = self.temporal_mlp(flat_in) # (B, T*(n_grip+n_rot)) gr = flat_out.reshape(B, T, self.n_gripper_bins + self.n_rot_clusters) grip_logits = gr[:, :, :self.n_gripper_bins] rot_logits = gr[:, :, self.n_gripper_bins:] else: # temporal_cnn x_tc = self.temporal_cnn(temp_feat.transpose(1, 2)).transpose(1, 2) grip_logits = self.pv_grip_head(x_tc) # (B, T, n_grip) rot_logits = self.pv_rot_head(x_tc) # (B, T, n_rot) return { "volume_logits": vol, "voxel_positions": voxel_pos, # (B, NZ, P, P, 3) "grip_logits": grip_logits, "rot_logits": rot_logits, "F_refined": F_prim, # primary view's "patch_features": patch0, "views": self.views, "view_feat_maps": view_feats, # list of (B, F, P, P) } def multiview_losses( out: dict, target_yx: torch.Tensor, target_z: torch.Tensor, target_grip: torch.Tensor, target_rot: torch.Tensor, target_xyz: torch.Tensor, ) -> dict: """CE loss over the FULL N×Z×P×P voxel set with target = argmin world distance from GT-XYZ. Same identity-trick as single-view to avoid materialising a per-voxel diff tensor.""" vol = out["volume_logits"] # (B, T, NZ, P, P) voxel_pos = out["voxel_positions"] # (B, NZ, P, P, 3) B, T, NZ, P, _ = vol.shape V = NZ * P * P voxel_flat = voxel_pos.reshape(B, V, 3) with torch.no_grad(): gt = target_xyz # (B, T, 3) voxel_sq = (voxel_flat ** 2).sum(dim=-1) # (B, V) gt_sq = (gt ** 2).sum(dim=-1) # (B, T) dot = torch.einsum("btc,bvc->btv", gt, voxel_flat) # (B, T, V) dist_sq = (gt_sq[..., None] + voxel_sq[:, None, :] - 2.0 * dot) target = dist_sq.argmin(dim=-1) vol_flat = vol.reshape(B * T, V) loss_vol = F.cross_entropy(vol_flat, target.reshape(B * T)) loss_grip = F.cross_entropy( out["grip_logits"].reshape(B * T, -1), target_grip.reshape(B * T)) loss_rot = F.cross_entropy( out["rot_logits"].reshape(B * T, -1), target_rot.reshape(B * T)) return { "loss/volume": loss_vol, "loss/grip": loss_grip, "loss/rot": loss_rot, "loss/total": loss_vol + loss_grip + loss_rot, }