"""DinoVolumeSceneV3 — factorized dot-product volumetric model. The 2026-06-12 follow-up to ``DinoVolumeSceneVolumetricMultiView``. The big architectural moves: 1. **No per-voxel MLP.** Scoring is a pure 3-term dot product. The per-pixel feature ``F_k(y, x)`` is split into three sub-features: ``[F_k^z (16) | F_k^t (16) | F_k^s (F_s=32)]``. The query for a voxel at (view k, y, x, z, t) is ``[PE_z(z) | PE_t(t) | spatial_q(t)]``. Score factors additively in (z, t): score(k, y, x, z, t) = F_k^z(y,x)·PE_z(z) + F_k^t(y,x)·PE_t(t) + F_k^s(y,x)·spatial_q(t) 2. **High resolution (default P=256, Z=256).** The factored form makes the full (T, NZ, P, P) volume never need to be materialized — peak tensor is the (B, N, P, P, Z) Z-LSE precompute at ~134 MB fp16, total memory ~100x less than the naive dense form. 3. **Sum aggregation = product-of-experts across views.** Each view's voxel slab is scored INDEPENDENTLY against the shared query (using only that view's own pixel feature at the source pixel — no per-voxel cross-projection). Slabs are concatenated and a single joint softmax over (view, y, x, z) per t. Sum of additive scores = product of probabilities, giving multi-view AND semantics. 4. **Sin/cos PE for height.** ``PE_z`` is a Fourier feature encoding of the world Z value (meters), not the bin index. Generalizes across z_ranges (no embedding-retrain needed when the task changes Z extent) and starts informative from step 0. 5. **No eef_feat in the global query.** Just concat-CLS from both views + past_n proprio (grip, world-Z) encoded. 6. **Grip / rotation via argmax voxel + concat-of-both-views.** Find the argmax voxel per timestep, project its world XYZ into BOTH views, bilinear-sample each view's full per-pixel feature, concat, flatten over T, and feed a 4-layer flat MLP that emits T·(n_grip + n_rot) logits. Train via the canonical entry: ``train.py --v3 [...]``. Loss via ``v3_losses`` (factorized CE over the joint volume). Decode via ``decode_v3`` in ``lib/inference.py``. """ 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 FiveLayerCNN from .model import (DEFAULT_DINO_REPO, DEFAULT_DINO_WEIGHTS, sin_pe) def _fourier_pe_meters(z_lo: float, z_hi: float, Z: int, dim: int) -> torch.Tensor: """Sin/cos PE of the physical world-Z bin centres in metres. Each bin's encoding is ``[sin(z·π·2^0), cos(z·π·2^0), ..., sin(z·π·2^(K-1)), cos(z·π·2^(K-1))]`` for K = dim//2 octaves. Z is normalized to roughly [0, 1] over the working range so the lowest frequency is one cycle per range. """ assert dim % 2 == 0 K = dim // 2 centres = torch.tensor( [z_lo + (i + 0.5) * (z_hi - z_lo) / Z for i in range(Z)], dtype=torch.float32) # Normalise z into roughly [0, 1] so the lowest frequency spans the range. span = float(z_hi - z_lo) z_norm = (centres - z_lo) / max(span, 1e-6) # (Z,) freqs = (2.0 ** torch.arange(K, dtype=torch.float32)) # (K,) args = z_norm[:, None] * freqs[None, :] * math.pi # (Z, K) pe = torch.cat([torch.sin(args), torch.cos(args)], dim=-1) # (Z, dim) return pe def _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. the wrist cam moves with the EE so z_hi changes per frame). Returns ``(B, Z)`` in metres. """ B = z_hi_b.shape[0] 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) # (B,) return z_lo_t.unsqueeze(1) + idx.unsqueeze(0) * span.unsqueeze(1) # (B, Z) def _per_view_fourier_pe(z_lo: float, z_hi_global: float, bin_centers_b: torch.Tensor, dim: int) -> torch.Tensor: """Per-batch sin/cos PE of bin centres in metres. Normalisation is by the GLOBAL (z_lo, z_hi_global) span — NOT per-sample — so a voxel at z=0.15m produces the same PE regardless of which view (or batch) sampled it. The wrist view's bins concentrate in low z; they hit a subset of the global PE range, but with the same basis as the scene view at matching physical heights. This is what gives cross-view consistency. 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) pe = torch.cat([torch.sin(args), torch.cos(args)], dim=-1) # (B, Z, dim) return pe 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, T, 3) into one cam, bilinear-sample features at the projected pixel. Returns (B, T, F). Out-of-frustum voxels get zero features via ``padding_mode='zeros'``. """ 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) cam = torch.einsum("bij,bvj->bvi", T_w2c, xyz_h)[..., :3] z = cam[..., 2:3].clamp(min=1e-3) uv = torch.einsum("bij,bvj->bvi", K_in, cam) / z norm_u = (uv[..., 0] / float(img_size)) * 2.0 - 1.0 norm_v = (uv[..., 1] / float(img_size)) * 2.0 - 1.0 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, T, 2) sampled = F.grid_sample(F_map, grid, mode="bilinear", padding_mode="zeros", align_corners=False) # (B, F, 1, T) return sampled.squeeze(2).transpose(1, 2) # (B, T, F) class DinoVolumeSceneV3(nn.Module): """Factorized multi-view dot-product volumetric model. Returns ``out`` containing factored score components rather than a fully-materialized (B, T, V) volume tensor. The downstream loss function ``v3_losses`` computes the joint-volume CE without ever materialising the dense volume. """ def __init__( self, *, views: list[str] = ("scene", "wrist"), n_window: int = 32, n_height_bins: int = 256, pred_size: int = 256, F_z: int = 16, F_t: int = 16, F_s: int = 32, 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.4, img_size: int = 504, past_n: int = 8, past_enc_dim: int = 64, temporal_mlp_hidden: int = 512, cam_margin_m: float = 0.01, dino_repo: str = DEFAULT_DINO_REPO, dino_weights: str = DEFAULT_DINO_WEIGHTS, **_unused, ): super().__init__() self.views = list(views) self.n_views = len(self.views) self.n_window = int(n_window) self.n_z = int(n_height_bins) self.pred_size = int(pred_size) self.F_z = int(F_z) self.F_t = int(F_t) self.F_s = int(F_s) self.F_total = self.F_z + self.F_t + self.F_s self.n_gripper_bins = int(n_gripper_bins) self.n_rot_clusters = int(n_rot_clusters) self.z_lo = float(z_lo); self.z_hi = float(z_hi) self.img_size = int(img_size) self.past_n = int(past_n) self.past_enc_dim = int(past_enc_dim) self.d_model = int(d_model) self.cam_margin_m = float(cam_margin_m) # Backbone — DINOv3 ViT-S/16+ (same as other models). self.backbone = torch.hub.load( os.environ.get("DINO_REPO_DIR", dino_repo), "dinov3_vits16plus", source="local", weights=os.environ.get("DINO_WEIGHTS_PATH", dino_weights), ) embed_dim = self.backbone.embed_dim self.embed_dim = embed_dim # Per-view per-pixel feature head: DINO patch → F_total → bilinear-up # → 5-layer refine CNN. Output (B, F_total, P, P). self.feat_head = nn.Conv2d(embed_dim, self.F_total, 1) self.refine_cnn = FiveLayerCNN(self.F_total, kernel=3) # Fixed sin/cos PEs. pe_z = _fourier_pe_meters(self.z_lo, self.z_hi, self.n_z, self.F_z) self.register_buffer("pe_z", pe_z) # (Z, F_z) self.register_buffer("pe_t", sin_pe(self.n_window, self.F_t)) # (T, F_t) # Bin centres in metres — kept for downstream decode (xyz lookup). bin_centres_z = torch.tensor( [self.z_lo + (i + 0.5) * (self.z_hi - self.z_lo) / self.n_z for i in range(self.n_z)], dtype=torch.float32) self.register_buffer("bin_centers_z", bin_centres_z) # Past-trajectory encoder (matches the existing multi-view model's spec). 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 # Global query path: concat-CLS + past_enc → input_proj → global MLP # → q_head produces T spatial queries of dim F_s. global_in_dim = self.n_views * embed_dim if self.past_n > 0: global_in_dim += self.past_enc_dim self.input_proj = nn.Linear(global_in_dim, self.d_model) hidden = int(self.d_model * mlp_hidden_mult) layers: list[nn.Module] = [nn.Linear(self.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, self.n_window * self.F_s) # Grip / rot head: gather full F_total feature from BOTH views at # the chosen voxel's world XYZ, concat, stack over T, flatten, # 4-layer MLP → grip + rot bin logits. in_dim = self.n_window * self.n_views * self.F_total out_dim = self.n_window * (self.n_gripper_bins + self.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), ) # ──────────────────────────────────────────────────────────────────── # Backbone + per-pixel feature head # ──────────────────────────────────────────────────────────────────── def _patch_features(self, rgb: torch.Tensor): out = self.backbone.forward_features(rgb) cls = out["x_norm_clstoken"] toks = out["x_norm_patchtokens"] P_dino = int(round(math.sqrt(toks.shape[1]))) patch = toks.transpose(1, 2).reshape(rgb.size(0), -1, P_dino, P_dino) return patch, cls def _pixel_features(self, rgb: torch.Tensor): patch, cls = self._patch_features(rgb) F0 = self.feat_head(patch) F1 = F.interpolate(F0, size=(self.pred_size, self.pred_size), mode="bilinear", align_corners=False) F_pix = self.refine_cnn(F1) # (B, F_total, P, P) return F_pix, cls @torch.no_grad() 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. Ray geometry depends only on (K_in, T_w2c, bin_centers) — none of which receive gradients — so we skip autograd to drop the graph and free the (B, P, P, Z, 3) intermediates after the forward pass. """ 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) 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_centres = bin_centers_b.to(rz.dtype) # (B, Z) lam = ((z_centres.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) # ──────────────────────────────────────────────────────────────────── # Forward # ──────────────────────────────────────────────────────────────────── def forward(self, rgb: list[torch.Tensor], K_in: list[torch.Tensor], T_w2c: list[torch.Tensor], target_xyz: torch.Tensor | None = None, past_grip: torch.Tensor | None = None, past_z: torch.Tensor | None = None, border_mask_px: int = 0, start_pix: torch.Tensor | None = None, **_kwargs) -> dict: """Returns a dict with factored score components. The full per-voxel logits are NEVER materialized — downstream callers use the factored form to compute loss / argmax cheaply. Required: - ``rgb``: list of (B, 3, S, S) per view (N) - ``K_in`` list of (B, 3, 3) - ``T_w2c`` list of (B, 4, 4) Optional: - ``target_xyz`` (B, T, 3): if provided (training), the grip/rot path gathers features at the GT-nearest voxel. Otherwise it gathers at argmax voxel (inference). - ``past_grip`` (B, past_n), ``past_z`` (B, past_n) """ assert isinstance(rgb, (list, tuple)) and len(rgb) == self.n_views B = rgb[0].shape[0] T = self.n_window N = self.n_views P = self.pred_size Z = self.n_z # 1. Per-view backbone + per-pixel feature heads. view_feats: list[torch.Tensor] = [] # each (B, F_total, P, P) view_cls: list[torch.Tensor] = [] for k in range(N): F_pix, cls = self._pixel_features(rgb[k]) view_feats.append(F_pix) view_cls.append(cls) # 2a. Per-view adaptive z range: same n_z bins, but each view's # z_hi is capped at min(global_z_hi, cam_z_k - margin). Wrist cam # rides the EE so its cam_z varies per frame; bins concentrate # below it. Scene cam is high and static so its z_hi stays = global. view_bin_centers: list[torch.Tensor] = [] view_pe_z: list[torch.Tensor] = [] 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))) # Clamp: z_hi must be > z_lo + 1 cm. Below that the voxel grid # collapses; fall back to a thin slab at z_lo + 1cm. z_hi_k = torch.clamp(z_hi_k, min=float(self.z_lo) + 0.01) bin_b = _per_view_bin_centers(self.z_lo, z_hi_k, self.n_z) # (B, Z) pe_b = _per_view_fourier_pe( self.z_lo, self.z_hi, bin_b, self.F_z) # (B, Z, F_z) view_bin_centers.append(bin_b) view_pe_z.append(pe_b) # 2b. Per-view voxel positions for downstream decode + GT lookup. view_voxel_pos: list[torch.Tensor] = [] for k in range(N): view_voxel_pos.append(self._unproject_grid( K_in[k], T_w2c[k], view_bin_centers[k])) # 3. Global query path: concat-CLS + past_enc → spatial query (B, T, F_s). global_in = [torch.cat(view_cls, dim=-1)] if self.past_encoder is not None: assert past_grip is not None and past_z is not None, \ "past_n > 0 requires past_grip / past_z" past_flat = torch.cat([past_grip, past_z], dim=-1) global_in.append(self.past_encoder(past_flat)) global_state = self.input_proj(torch.cat(global_in, dim=-1)) hglob = self.global_mlp(global_state) spatial_q = self.q_head(hglob).reshape(B, T, self.F_s) # 4. Split per-pixel features into (F_z, F_t, F_s) sub-features. view_Fz = [v[:, :self.F_z] for v in view_feats] # (B, F_z, P, P) view_Ft = [v[:, self.F_z:self.F_z + self.F_t] for v in view_feats] view_Fs = [v[:, self.F_z + self.F_t:] for v in view_feats] # 5. Apply optional border mask in the spatial scoring. The simplest # way to enforce it without materializing the volume is to zero out # the corresponding (y, x) cells of F_z / F_t / F_s (so their dot # products vanish, which is < +∞ → contributes ~0 to LSE). if border_mask_px > 0 and border_mask_px * 2 < P: keep = torch.zeros(P, P, dtype=torch.bool, device=view_feats[0].device) keep[border_mask_px:P - border_mask_px, border_mask_px:P - border_mask_px] = True mask_yx = keep.view(1, 1, P, P).float() view_Fz = [f * mask_yx for f in view_Fz] view_Ft = [f * mask_yx for f in view_Ft] view_Fs = [f * mask_yx for f in view_Fs] # 6. Find the chosen voxel per (B, T) for grip/rot feature gather. # - At training (target_xyz given): nearest-3D voxel. # - At inference (no target): argmax via factored scoring. chosen_view, chosen_z, chosen_y, chosen_x = self._chosen_voxel_indices( view_voxel_pos, view_Fz, view_Ft, view_Fs, spatial_q, target_xyz, view_pe_z) # 6b. ALSO always compute the argmax-based prediction. This is what # the viz should plot (pred vs GT). During inference chosen_* and # pred_* coincide; during training they differ because chosen_* # snaps to the GT voxel. pred_view, pred_z, pred_y, pred_x = self._chosen_voxel_indices( view_voxel_pos, view_Fz, view_Ft, view_Fs, spatial_q, target_xyz=None, view_pe_z=view_pe_z) # 7. Gather chosen voxel world XYZ. b_idx = torch.arange(B, device=rgb[0].device).view(B, 1).expand(B, T) t_idx = torch.arange(T, device=rgb[0].device).view(1, T).expand(B, T) all_pos = torch.stack(view_voxel_pos, dim=1) # (B, N, Z, P, P, 3) chosen_xyz = all_pos[b_idx, chosen_view, chosen_z, chosen_y, chosen_x] # (B, T, 3) pred_xyz = all_pos[b_idx, pred_view, pred_z, pred_y, pred_x] # (B, T, 3) # 8. Grip + rotation: bilinear-sample the FULL F_total feature from # BOTH views at chosen_xyz, concat over views, flatten over T, MLP. feat_per_view: list[torch.Tensor] = [] for k in range(N): feat_per_view.append(_sample_view_feat( view_feats[k], chosen_xyz, K_in[k], T_w2c[k], self.img_size)) cat_feat = torch.cat(feat_per_view, dim=-1) # (B, T, N*F_total) flat_in = cat_feat.reshape(B, T * N * self.F_total) flat_out = self.temporal_mlp(flat_in) 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:] return { # Factored score components — loss / argmax use these directly. "view_Fz": view_Fz, # list of (B, F_z, P, P) "view_Ft": view_Ft, # list of (B, F_t, P, P) "view_Fs": view_Fs, # list of (B, F_s, P, P) # Per-view PE_z + bin_centers (z range is per-view, per-frame). "view_pe_z": view_pe_z, # list of (B, Z, F_z) "view_bin_centers": view_bin_centers, # list of (B, Z) "pe_t": self.pe_t, # (T, F_t) "spatial_q": spatial_q, # (B, T, F_s) "voxel_positions": all_pos, # (B, N, Z, P, P, 3) "chosen_view": chosen_view, # (B, T) "chosen_z": chosen_z, # (B, T) "chosen_y": chosen_y, # (B, T) "chosen_x": chosen_x, # (B, T) "chosen_xyz_world": chosen_xyz, # (B, T, 3) # Argmax-based prediction (== chosen_* at inference, differs # in training where chosen_* snaps to GT). Use this for viz. "pred_view": pred_view, # (B, T) "pred_z": pred_z, # (B, T) "pred_y": pred_y, # (B, T) "pred_x": pred_x, # (B, T) "pred_xyz_world": pred_xyz, # (B, T, 3) "grip_logits": grip_logits, # (B, T, n_grip) "rot_logits": rot_logits, # (B, T, n_rot) "views": list(self.views), "view_feat_maps": view_feats, # list of (B, F_total, P, P) — for viz "view_cls": view_cls, "n_z": Z, "pred_size": P, "n_views": N, } # ──────────────────────────────────────────────────────────────────── # Chosen-voxel index lookup (GT or argmax via factored scoring) # ──────────────────────────────────────────────────────────────────── @torch.no_grad() def _chosen_voxel_indices(self, view_voxel_pos, view_Fz, view_Ft, view_Fs, spatial_q, target_xyz, view_pe_z): """Returns four (B, T) long tensors: (view, z, y, x) per timestep. Training (target_xyz given): nearest-3D voxel to GT XYZ across all view slabs. Inference (target_xyz None): argmax of factored score per t. """ B = view_voxel_pos[0].shape[0] N = len(view_voxel_pos) Z = self.n_z; P = self.pred_size device = view_voxel_pos[0].device if target_xyz is not None: # Nearest 3D voxel. Concat positions: (B, N, Z, P, P, 3) → (B, V, 3). # ChUNK over T to avoid the (B, T, V) intermediate — at # P=Z=256, N=2, V=33M and a single (B, V) tensor is already # ~260 MB at B=2. all_pos = torch.stack(view_voxel_pos, dim=1) V_total = N * Z * P * P pos_flat = all_pos.reshape(B, V_total, 3) T_ = int(target_xyz.shape[1]) v2 = (pos_flat * pos_flat).sum(-1) # (B, V) idx_per_t = [] for t in range(T_): tgt_t = target_xyz[:, t] # (B, 3) dot_t = torch.einsum("bc,bvc->bv", tgt_t, pos_flat) # (B, V) t2_t = (tgt_t * tgt_t).sum(-1, keepdim=True) # (B, 1) d2_t = t2_t + v2 - 2 * dot_t # (B, V) idx_per_t.append(d2_t.argmin(dim=1)) # (B,) idx = torch.stack(idx_per_t, dim=1) # (B, T) else: # Argmax via factored scoring. # A(B, view, P, P, Z) = einsum F_z * PE_z. Per-view to bound memory. # For each view: F_z (B, F_z, P, P), PE_z (Z, F_z) → (B, P, P, Z) # Find A_max(B, view, P, P) and z_argmax(B, view, P, P). all_A_max = [] # list of (B, P, P) all_z_argmax = [] # list of (B, P, P) for k in range(N): A_k = torch.einsum("bfyx,bzf->byxz", view_Fz[k], view_pe_z[k]) A_max_k, z_argmax_k = A_k.max(dim=-1) # (B, P, P) all_A_max.append(A_max_k); all_z_argmax.append(z_argmax_k) A_max = torch.stack(all_A_max, dim=1) # (B, N, P, P) z_argmax = torch.stack(all_z_argmax, dim=1) # (B, N, P, P) # B(view, y, x, t) = F_t · PE_t + F_s · spatial_q. Per view. all_B = [] T_ = self.n_window for k in range(N): Bk = (torch.einsum("bfyx,tf->byxt", view_Ft[k], self.pe_t) + torch.einsum("bfyx,btf->byxt", view_Fs[k], spatial_q)) all_B.append(Bk) # (B, P, P, T) B_t = torch.stack(all_B, dim=1) # (B, N, P, P, T) # Combined per-(view, y, x, t): A_max + B_t. combined = A_max.unsqueeze(-1) + B_t # (B, N, P, P, T) comb_flat = combined.permute(0, 4, 1, 2, 3).reshape(B, T_, N * P * P) argmax_flat = comb_flat.argmax(dim=2) # (B, T) v_idx = argmax_flat // (P * P) yx = argmax_flat % (P * P) y_idx = yx // P x_idx = yx % P # Pull z from z_argmax at the chosen (view, y, x). b_idx = torch.arange(B, device=device).view(B, 1).expand(B, T_) z_idx = z_argmax[b_idx, v_idx, y_idx, x_idx] return v_idx, z_idx, y_idx, x_idx # Decode the flat index. v_idx = idx // (Z * P * P) zyx = idx % (Z * P * P) z_idx = zyx // (P * P) yx = zyx % (P * P) y_idx = yx // P x_idx = yx % P return v_idx, z_idx, y_idx, x_idx # ──────────────────────────────────────────────────────────────────────── # Factorized loss (CE over joint volume via LSE decomposition) # ──────────────────────────────────────────────────────────────────────── def v3_losses(out: dict, target_xyz: torch.Tensor, target_grip: torch.Tensor, target_rot: torch.Tensor) -> dict: """Factored CE on the joint volume + CE on grip + CE on rot. Score per voxel (view, y, x, z) at timestep t decomposes as ``A(view, y, x, z) + B(view, y, x, t)`` because (z) and (t) live in additive sub-spaces. The LSE over the joint volume becomes ``LSE_{view,y,x}[B + LSE_z A]``, so we never materialize a (B, T, V) tensor. """ view_Fz = out["view_Fz"]; view_Ft = out["view_Ft"]; view_Fs = out["view_Fs"] view_pe_z = out["view_pe_z"]; pe_t = out["pe_t"]; spatial_q = out["spatial_q"] all_pos = out["voxel_positions"] # (B, N, Z, P, P, 3) B = spatial_q.shape[0]; T = spatial_q.shape[1] N, Z, P = out["n_views"], out["n_z"], out["pred_size"] # GT voxel per (B, T) — nearest-3D across all view slabs. Chunk # over T to avoid the (B, T, V) intermediate (V can be ~33M at # P=Z=256, N=2). pos_flat = all_pos.reshape(B, N * Z * P * P, 3) v2 = (pos_flat * pos_flat).sum(-1) # (B, V) gt_idx_per_t = [] for t in range(T): tgt_t = target_xyz[:, t] # (B, 3) dot_t = torch.einsum("bc,bvc->bv", tgt_t, pos_flat) # (B, V) t2_t = (tgt_t * tgt_t).sum(-1, keepdim=True) # (B, 1) d2_t = t2_t + v2 - 2 * dot_t # (B, V) gt_idx_per_t.append(d2_t.argmin(dim=1)) # (B,) gt_idx = torch.stack(gt_idx_per_t, dim=1) # (B, T) gt_view = gt_idx // (Z * P * P) zyx = gt_idx % (Z * P * P) gt_z = zyx // (P * P); yx = zyx % (P * P) gt_y = yx // P; gt_x = yx % P # Numerator: score at GT voxel. # Need F_z, F_t, F_s gathered at (gt_view, gt_y, gt_x) per (B, T). # Then dot with PE_z(gt_z), PE_t(t), spatial_q[t]. b_idx = torch.arange(B, device=spatial_q.device).view(B, 1).expand(B, T) t_idx = torch.arange(T, device=spatial_q.device).view(1, T).expand(B, T) F_z_stack = torch.stack(view_Fz, dim=1) # (B, N, F_z, P, P) F_t_stack = torch.stack(view_Ft, dim=1) # (B, N, F_t, P, P) F_s_stack = torch.stack(view_Fs, dim=1) # (B, N, F_s, P, P) pe_z_stack = torch.stack(view_pe_z, dim=1) # (B, N, Z, F_z) f_z_at_gt = F_z_stack[b_idx, gt_view, :, gt_y, gt_x] # (B, T, F_z) f_t_at_gt = F_t_stack[b_idx, gt_view, :, gt_y, gt_x] # (B, T, F_t) f_s_at_gt = F_s_stack[b_idx, gt_view, :, gt_y, gt_x] # (B, T, F_s) pe_z_at_gt = pe_z_stack[b_idx, gt_view, gt_z] # (B, T, F_z) pe_t_at_gt = pe_t[t_idx] # (B, T, F_t) score_gt_z = (f_z_at_gt * pe_z_at_gt).sum(-1) # (B, T) score_gt_t = (f_t_at_gt * pe_t_at_gt).sum(-1) score_gt_s = (f_s_at_gt * spatial_q).sum(-1) score_gt = score_gt_z + score_gt_t + score_gt_s # (B, T) # Denominator: LSE over (view, y, x, z) per t. # Compute LSE_z(A) once per (view, y, x), then sum with B(view, y, x, t). A_lse_z_per_view = [] # list of (B, P, P) for k in range(N): A_k = torch.einsum("bfyx,bzf->byxz", view_Fz[k], view_pe_z[k]) # (B, P, P, Z) A_lse_z_per_view.append(torch.logsumexp(A_k, dim=-1)) del A_k A_lse_z = torch.stack(A_lse_z_per_view, dim=1) # (B, N, P, P) # B(view, y, x, t) per view. B_per_view = [] for k in range(N): Bk = (torch.einsum("bfyx,tf->byxt", view_Ft[k], pe_t) + torch.einsum("bfyx,btf->byxt", view_Fs[k], spatial_q)) B_per_view.append(Bk) # (B, P, P, T) B_t = torch.stack(B_per_view, dim=1) # (B, N, P, P, T) # Combined (B, N, P, P, T) → LSE over (N, P, P) per t. combined = A_lse_z.unsqueeze(-1) + B_t # (B, N, P, P, T) comb_flat = combined.permute(0, 4, 1, 2, 3).reshape(B, T, N * P * P) lse_voxel_per_t = torch.logsumexp(comb_flat, dim=2) # (B, T) loss_volume = -(score_gt - lse_voxel_per_t).mean() # Grip + rot. grip_logits = out["grip_logits"]; rot_logits = out["rot_logits"] loss_grip = F.cross_entropy( grip_logits.reshape(B * T, -1), target_grip.reshape(B * T)) loss_rot = F.cross_entropy( rot_logits.reshape(B * T, -1), target_rot.reshape(B * T)) return { "loss/volume": loss_volume, "loss/grip": loss_grip, "loss/rot": loss_rot, "loss/total": loss_volume + loss_grip + loss_rot, } # ──────────────────────────────────────────────────────────────────────── # View-dominance & per-view marginal heatmap helpers (for the rerun viz) # ──────────────────────────────────────────────────────────────────────── @torch.no_grad() def v3_view_dominance(out: dict) -> torch.Tensor: """Returns (B, N, T): the per-view log-marginal at each timestep, log p(view k | t) = LSE_{y,x,z} score(k, y, x, z, t) − LSE_global(t). A softmax across the N dim gives p(view | t), which is what to plot in the dominance bar. """ view_Fz = out["view_Fz"]; view_Ft = out["view_Ft"]; view_Fs = out["view_Fs"] view_pe_z = out["view_pe_z"]; pe_t = out["pe_t"]; spatial_q = out["spatial_q"] B = spatial_q.shape[0]; T = spatial_q.shape[1] N, Z, P = out["n_views"], out["n_z"], out["pred_size"] # Per-view LSE_z A. Alse = [] for k in range(N): A_k = torch.einsum("bfyx,bzf->byxz", view_Fz[k], view_pe_z[k]) Alse.append(torch.logsumexp(A_k, dim=-1)) A_lse_z = torch.stack(Alse, dim=1) # (B, N, P, P) # Per-view B(view, y, x, t). B_per_view = [] for k in range(N): Bk = (torch.einsum("bfyx,tf->byxt", view_Ft[k], pe_t) + torch.einsum("bfyx,btf->byxt", view_Fs[k], spatial_q)) B_per_view.append(Bk) B_t = torch.stack(B_per_view, dim=1) # (B, N, P, P, T) combined = A_lse_z.unsqueeze(-1) + B_t # (B, N, P, P, T) # Per-view LSE over (y, x). lse_view_t = torch.logsumexp(combined.flatten(2, 3), dim=2) # (B, N, T) # Subtract global LSE so softmax across N gives a proper distribution. lse_global = torch.logsumexp(lse_view_t, dim=1, keepdim=True) # (B, 1, T) return lse_view_t - lse_global # (B, N, T) @torch.no_grad() def v3_z_marginal(out: dict) -> torch.Tensor: """Returns (B, Z, T) log-marginal log p(z | t) — sums per-view per-yx score mass into a per-z distribution at each timestep. Computed chunked-over-t so the (B, N, P, P, Z) intermediate never materialises. """ view_Fz = out["view_Fz"]; view_Ft = out["view_Ft"]; view_Fs = out["view_Fs"] view_pe_z = out["view_pe_z"]; pe_t = out["pe_t"]; spatial_q = out["spatial_q"] B_ = spatial_q.shape[0]; T = spatial_q.shape[1] N, Z, P = out["n_views"], out["n_z"], out["pred_size"] # A_k(B, P, P, Z) = einsum F_z · PE_z per view. A_per_view = [torch.einsum("bfyx,bzf->byxz", view_Fz[k], view_pe_z[k]) for k in range(N)] # Per-t loop: compose A + B_t per view, LSE over (y, x), LSE over views. out_z = [] for t in range(T): per_view_lse = [] for k in range(N): Bk_t = (torch.einsum("bfyx,f->byx", view_Ft[k], pe_t[t]) + torch.einsum("bfyx,bf->byx", view_Fs[k], spatial_q[:, t])) combined = A_per_view[k] + Bk_t.unsqueeze(-1) # (B, P, P, Z) # LSE over (y, x) → (B, Z) per_view_lse.append(torch.logsumexp( combined.reshape(B_, P * P, Z), dim=1)) # LSE over views → (B, Z) z_lse = torch.logsumexp(torch.stack(per_view_lse, dim=1), dim=1) out_z.append(z_lse) z_logits = torch.stack(out_z, dim=2) # (B, Z, T) # Subtract global LSE so softmax over Z gives proper p(z | t). z_logits = z_logits - torch.logsumexp(z_logits, dim=1, keepdim=True) return z_logits @torch.no_grad() def v3_per_view_heatmaps_all_t(out: dict) -> list[torch.Tensor]: """Per-view (B, T, P, P) log-marginal over z at every timestep. Each entry is a proper log-prob: global LSE over (N, y, x) is subtracted per t so a softmax over (view, y, x) at fixed t sums to 1. This is the multi-t analogue of ``v3_per_view_heatmaps`` — suitable input to ``marginal_heatmap_grid``. """ view_Fz = out["view_Fz"]; view_Ft = out["view_Ft"]; view_Fs = out["view_Fs"] view_pe_z = out["view_pe_z"]; pe_t = out["pe_t"]; spatial_q = out["spatial_q"] B_ = spatial_q.shape[0]; T = spatial_q.shape[1] N, Z, P = out["n_views"], out["n_z"], out["pred_size"] # LSE_z A_k(B, P, P) per view. A_lse_per_view = [] for k in range(N): A_k = torch.einsum("bfyx,bzf->byxz", view_Fz[k], view_pe_z[k]) A_lse_per_view.append(torch.logsumexp(A_k, dim=-1)) # (B, P, P) del A_k # B_k(B, P, P, T) per view. B_per_view = [] for k in range(N): Bk = (torch.einsum("bfyx,tf->byxt", view_Ft[k], pe_t) + torch.einsum("bfyx,btf->byxt", view_Fs[k], spatial_q)) B_per_view.append(Bk) # (B, P, P, T) # Combined per view: A_lse + B → (B, P, P, T) per_view_score = [A_lse_per_view[k].unsqueeze(-1) + B_per_view[k] for k in range(N)] # Joint LSE over (N, y, x) per t. stacked = torch.stack(per_view_score, dim=1) # (B, N, P, P, T) lse_per_t = torch.logsumexp(stacked.flatten(1, 3), dim=1, keepdim=True) # (B, 1, T) # Return list of (B, T, P, P) per view. out_list = [] for k in range(N): log_p = per_view_score[k] - lse_per_t.view(B_, 1, 1, T) out_list.append(log_p.permute(0, 3, 1, 2).contiguous()) # (B, T, P, P) return out_list @torch.no_grad() def v3_per_view_heatmaps(out: dict, t_index: int = 0) -> list[torch.Tensor]: """Per-view (P, P) marginal-over-z probability map at one timestep. Returns list of (B, P, P) tensors in log-prob space (subtract global LSE so the joint over views and yx is a proper distribution). """ view_Fz = out["view_Fz"]; view_Ft = out["view_Ft"]; view_Fs = out["view_Fs"] view_pe_z = out["view_pe_z"]; pe_t = out["pe_t"]; spatial_q = out["spatial_q"] B = spatial_q.shape[0] N, Z, P = out["n_views"], out["n_z"], out["pred_size"] t = int(t_index) Alse = [] Bt = [] for k in range(N): A_k = torch.einsum("bfyx,bzf->byxz", view_Fz[k], view_pe_z[k]) Alse.append(torch.logsumexp(A_k, dim=-1)) # (B, P, P) Bk_t = (torch.einsum("bfyx,f->byx", view_Ft[k], pe_t[t]) + torch.einsum("bfyx,bf->byx", view_Fs[k], spatial_q[:, t])) Bt.append(Bk_t) A_lse_z = torch.stack(Alse, dim=1) # (B, N, P, P) B_at_t = torch.stack(Bt, dim=1) # (B, N, P, P) per_view_yx_score = A_lse_z + B_at_t # (B, N, P, P) # Global LSE over (N, P, P). lse_global = torch.logsumexp(per_view_yx_score.flatten(1, 3), dim=1, keepdim=True) log_p = per_view_yx_score.flatten(1, 3) - lse_global log_p = log_p.reshape(B, N, P, P) return [log_p[:, k] for k in range(N)]