"""DinoVolumeSceneVolumetric — world-space voxel grid via image-feature unprojection. Same backbone + feat_head as ``DinoVolumeSceneHires`` (DINOv3 ViT-S/16+ plus 1×1 conv → bilinear up → 5-layer refinement CNN). The new part: 1. **Unproject the image into a world-space voxel grid.** For each (u_idx, v_idx, z_idx) on a ``P × P × Z`` grid (default ``128 × 128 × 64``), ray-cast the pixel through the camera at height ``z`` to get a 3D world position. Per-voxel feature = the image feature at that (u, v) concatenated with a **learnable per-bucket height embedding** of ``z``. 2. **3-layer per-voxel MLP** refines the concatenated features back to ``feat_dim``. 3. **Flat global MLP** (same as ``DinoVolumeSceneHiresFlat``) produces a per-timestep query of dim ``feat_dim``. 4. **Spatial attention** = dot product of the per-T query against the ``P×P×Z`` feature volume — same einsum shape as the factorized model. 5. **Supervision = nearest grid voxel by 3D distance**: for each (b, t), the CE target is ``argmin_v ||voxel_pos[v] − gt_xyz[b, t]||``. No +1 slot. Pure grid-voxel classification, but the "nearest" is taken in world XYZ (not factorized per-axis) so it respects the camera geometry. Why: world-space grid makes future multi-view fusion (max- or weighted-sum pooling per voxel across cameras) a drop-in extension. The 3D-nearest target eliminates the GT-slot gradient asymmetry that caused the old +1 slot variant to ignore grid voxels for ~3k iters. """ from __future__ import annotations import math import os import torch import torch.nn as nn import torch.nn.functional as F from .model_hires import _DinoHiresBase, FiveLayerCNN from .model import DEFAULT_DINO_REPO, DEFAULT_DINO_WEIGHTS def _sin_pe_continuous(z: torch.Tensor, d: int) -> torch.Tensor: """Sin/cos PE for continuous values. ``z`` of any shape → ``(..., d)``. Uses the standard half-sin / half-cos split. Frequencies are log-spaced over [1, 1e4] — matches transformer-style PEs. """ half = d // 2 freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=z.device, dtype=torch.float32) / max(1, half)) args = z.unsqueeze(-1) * freqs # (..., half) return torch.cat([torch.sin(args), torch.cos(args)], dim=-1) class DinoVolumeSceneVolumetric(_DinoHiresBase): """Predict world-space voxel via per-voxel image-feature + height MLP. Inherits the backbone + ``_pixel_features`` from ``_DinoHiresBase``. Replaces the AdaLN-Zero stack with a single flat global MLP (cribbed from ``DinoVolumeSceneHiresFlat``) — every timestep's query comes from one MLP forward, then dot-products against the same per-sample voxel feature volume. """ def __init__( self, *, 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, dino_repo: str = DEFAULT_DINO_REPO, dino_weights: str = DEFAULT_DINO_WEIGHTS, ): # The parent allocates an AdaLN stack we don't use (n_blocks=0 # makes ``self.blocks`` an empty ModuleList). 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=dino_repo, dino_weights=dino_weights, ) 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) # Bin centers (Z,). Used by decode + viz for the world-z of each bin. 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) # Learnable per-bucket height embedding (Z, height_enc_dim). # Replaces the old sin/cos PE — gives the model direct control over # how heights map into voxel-feature space. self.height_emb = nn.Embedding(Z, height_enc_dim) nn.init.normal_(self.height_emb.weight, std=0.02) # Per-voxel MLP: (feat_dim + height_enc_dim) → feat_dim, 3 layers. d_in = 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 recipe as DinoVolumeSceneHiresFlat). # Override grip_head / rot_head from the parent — flat outputs the # whole T-trajectory in one shot. T = n_window embed_dim = self.embed_dim # input_proj was created by parent at (feat_dim + embed_dim) → d_model. 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 # --------------------------------------------------------------------- def _unproject_grid(self, K_in: torch.Tensor, T_w2c: torch.Tensor ) -> torch.Tensor: """Compute world-space voxel positions ``(B, Z, P, P, 3)``. ``K_in`` — (B, 3, 3) intrinsics already scaled to ``img_size``. ``T_w2c`` — (B, 4, 4) world-to-cam pose. For each grid cell ``(u_idx, v_idx, z_idx)`` (u indexes columns → pixel x; v indexes rows → pixel y; z indexes height bin), the pixel center in ``img_size`` coords is (u_pix, v_pix) = ((u_idx + 0.5) * S/P, (v_idx + 0.5) * S/P) We ray-cast that pixel through the camera and intersect with the plane ``world_z = bin_centers_z[z_idx]``. """ 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 # Pixel grid in img_size coords (cell centers). 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 # (P,) v_pix = (v_idx + 0.5) * cell vv, uu = torch.meshgrid(v_pix, u_pix, indexing="ij") # (P, P) ones = torch.ones_like(uu) pix_h = torch.stack([uu, vv, ones], dim=-1) # (P, P, 3) # Cam ray dir for each pixel: K^-1 · (u, v, 1). K_inv = torch.inverse(K_in) # (B, 3, 3) ray_cam = torch.einsum("bij,hwj->bhwi", K_inv, pix_h) # (B, P, P, 3) # Cam→world: pose has world-to-cam, invert. T_c2w = torch.inverse(T_w2c) # (B, 4, 4) R = T_c2w[:, :3, :3] # (B, 3, 3) cam_o = T_c2w[:, :3, 3] # (B, 3) ray_world = torch.einsum("bij,bhwj->bhwi", R, ray_cam) # (B, P, P, 3) # Solve cam_o + lambda * ray_world = (x, y, z=target_z). rz = ray_world[..., 2] # (B, P, P) cam_oz = cam_o[..., 2].unsqueeze(-1).unsqueeze(-1) # (B, 1, 1) # Avoid division by near-zero z components (rays nearly parallel # to the world ground plane → unprojection is unstable; clamp). rz_safe = torch.where(rz.abs() < 1e-6, torch.full_like(rz, 1e-6), rz) # For each z bin center, compute lambda (B, P, P, Z). z_centers = self.bin_centers_z.to(rz.dtype) # (Z,) # broadcast: (B, P, P, Z) lam = ((z_centers.view(1, 1, 1, Z) - cam_oz.unsqueeze(-1)) / rz_safe.unsqueeze(-1)) # World position: cam_o + lam * ray_world. We want axis order # (B, Z, P, P, 3) to match the existing volume_logits layout. world = (cam_o.view(B, 1, 1, 1, 3) + lam.unsqueeze(-1) * ray_world.unsqueeze(3)) # (B, P, P, Z, 3) return world.permute(0, 3, 1, 2, 4).contiguous() # (B, Z, P, P, 3) @staticmethod def _project_world_to_image(xyz_world: torch.Tensor, K_in: torch.Tensor, T_w2c: torch.Tensor) -> torch.Tensor: """Project world XYZ → image (u, v). Shapes: ``xyz_world`` (B, T, 3); ``K_in`` (B, 3, 3); ``T_w2c`` (B, 4, 4). Returns ``(B, T, 2)``. """ B, T, _ = xyz_world.shape ones = torch.ones(B, T, 1, device=xyz_world.device, dtype=xyz_world.dtype) xyz_h = torch.cat([xyz_world, ones], dim=-1) # (B, T, 4) cam = torch.einsum("bij,btj->bti", T_w2c, xyz_h)[..., :3] # (B, T, 3) uvw = torch.einsum("bij,btj->bti", K_in, cam) # (B, T, 3) uv = uvw[..., :2] / uvw[..., 2:3].clamp(min=1e-6) return uv # --------------------------------------------------------------------- # Forward # --------------------------------------------------------------------- def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor, K_in: torch.Tensor, T_w2c: torch.Tensor) -> dict: """``rgb`` (B, 3, S, S); ``start_pix`` (B, 2); ``K_in`` (B, 3, 3); ``T_w2c`` (B, 4, 4). Returns: - ``volume_logits``: (B, T, Z, P, P) — same layout as factorized model - ``voxel_positions``: (B, Z, P, P, 3) world-space (for loss + decode + viz) - ``grip_logits`` (B, T, n_grip), ``rot_logits`` (B, T, n_rot) - ``F_refined`` (B, feat_dim, P, P), ``patch_features`` """ B, _, S, _ = rgb.shape T, Z, P = self.n_window, self.n_z, self.pred_size F_scene, patch, cls = self._pixel_features(rgb) # (B, F, P, P) # 1. World-space voxel positions. voxel_pos = self._unproject_grid(K_in, T_w2c) # (B, Z, P, P, 3) # 2. Per-voxel features = image feat at (u, v) ⊕ learnable height emb(z_bin). feat_pixel = F_scene.permute(0, 2, 3, 1).unsqueeze(1) # (B, 1, P, P, F) feat_pixel = feat_pixel.expand(-1, Z, -1, -1, -1) # (B, Z, P, P, F) h_emb = self.height_emb.weight.view(1, Z, 1, 1, -1)\ .expand(B, Z, P, P, -1) # (B, Z, P, P, h_enc) voxel_in = torch.cat([feat_pixel, h_emb], dim=-1) # (B, Z, P, P, F+h_enc) voxel_feat = self.voxel_mlp(voxel_in) # (B, Z, P, P, F) # 3. Flat global MLP. 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.device) eef_feat = F_scene[b_idx, :, sy, sx] # (B, F) q_in = self.input_proj(torch.cat([eef_feat, cls], dim=-1)) # (B, d_model) hglob = self.global_mlp(q_in) # (B, hidden) q = self.q_head(hglob).reshape(B, T, self.feat_dim) # (B, T, F) # 4. Score: dot product of q against voxel features. vol = torch.einsum("btc,bzhwc->btzhw", q, voxel_feat) # (B, T, Z, P, P) # 5. Heads (flat — whole-trajectory at once). 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) return { "volume_logits": vol, "voxel_positions": voxel_pos, "grip_logits": grip_logits, "rot_logits": rot_logits, "F_refined": F_scene, "patch_features": patch, } def volumetric_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 grid voxels with target = nearest voxel by 3D distance. Computes per-(b, t) the flat voxel index ``v* = argmin_v ||voxel_pos[b, v] − target_xyz[b, t]||²`` and applies CE against ``volume_logits.reshape(B*T, V)``. No +1 slot: this gives the model a single per-(b, t) bin target, but the "nearest" is taken in world XYZ (not factorized per-axis) so it respects camera geometry. Uses the |a−b|² = |a|² + |b|² − 2 a·b identity to avoid materializing the (B, T, V, 3) per-voxel diff tensor (~600 MB for default sizes). """ vol = out["volume_logits"] # (B, T, Z, P, P) voxel_pos = out["voxel_positions"] # (B, Z, P, P, 3) B, T, Z, P, _ = vol.shape V = Z * P * P voxel_flat = voxel_pos.reshape(B, V, 3) # (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) # (B, T, V) target = dist_sq.argmin(dim=-1) # (B, T) 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, }