"""Baseline: DINOv3 CLS → global 3D regression (no spatial reasoning). Same DINOv3 backbone + 5× AdaLN-Zero transformer as ``model.py``, but: - Input is only the **CLS token** (no per-pixel features, no eef lookup). - Position head is a **direct (T, 3) regression** of EE world XYZ, MSE loss. - Grip + rot heads are unchanged (CE on bins, K-means quat). Purpose: comparison baseline for the PARA factorized-volume formulation. Same input + same auxiliary heads → only the spatial reasoning differs. """ from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F from .model import (AdaLNZeroBlock, DEFAULT_DINO_REPO, DEFAULT_DINO_WEIGHTS, sin_pe) class DinoCLSXYZBaseline(nn.Module): """CLS-only baseline: (B, embed) → AdaLN-Zero → (xyz, grip_logits, rot_logits).""" def __init__( self, n_window: int = 16, d_model: int = 256, d_cond: int = 128, d_sin_t: int = 48, n_blocks: int = 5, n_gripper_bins: int = 32, n_rot_clusters: int = 64, dino_repo: str = DEFAULT_DINO_REPO, dino_weights: str = DEFAULT_DINO_WEIGHTS, ): super().__init__() self.n_window = n_window self.d_model = d_model self.n_gripper_bins = n_gripper_bins self.n_rot_clusters = n_rot_clusters import os 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 self.register_buffer("t_sin", sin_pe(n_window, d_sin_t)) self.t_cond_proj = nn.Linear(d_sin_t, d_cond) self.input_proj = nn.Linear(embed_dim, d_model) self.blocks = nn.ModuleList([ AdaLNZeroBlock(d_model, d_cond, mlp_ratio=4) for _ in range(n_blocks) ]) self.final_norm = nn.LayerNorm(d_model) self.xyz_head = nn.Linear(d_model, 3) self.grip_head = nn.Linear(d_model, n_gripper_bins) self.rot_head = nn.Linear(d_model, n_rot_clusters) def _encode(self, rgb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Run DINOv3 once and return (CLS, patch grid).""" import math out = self.backbone.forward_features(rgb) cls = out["x_norm_clstoken"] # (B, D) toks = out["x_norm_patchtokens"] # (B, P*P, D) P = int(round(math.sqrt(toks.shape[1]))) patches = toks.transpose(1, 2).reshape(toks.shape[0], -1, P, P) # (B, D, P, P) return cls, patches def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor = None) -> dict: B = rgb.size(0); T = self.n_window cls, patches = self._encode(rgb) q_in = self.input_proj(cls) # (B, d_model) q_in_bt = q_in.unsqueeze(1).expand(B, T, self.d_model).reshape(B * T, -1) cond_t = self.t_cond_proj(self.t_sin) # (T, d_cond) cond_bt = cond_t.unsqueeze(0).expand(B, T, -1).reshape(B * T, -1) h = q_in_bt for blk in self.blocks: h = blk(h, cond_bt) h = self.final_norm(h) penult = h.view(B, T, self.d_model) return { "xyz_pred": self.xyz_head(penult), # (B, T, 3) "grip_logits": self.grip_head(penult), # (B, T, n_grip) "rot_logits": self.rot_head(penult), # (B, T, n_rot) "cls": cls, # for PCA-like viz "patch_features": patches, # (B, D, P, P) for DINO PCA } class DinoCLSXYZBaselineMultiView(nn.Module): """Multi-view CLS → global XYZ regression baseline. Same DINOv3 backbone (shared across views) + optional DA3-style cross-view attention + concat-all-views CLS as ``DinoVolumeSceneVolumetricMultiView``, but with a flat (T, 3) XYZ regression head and CE grip/rot heads — no voxel grid, no per-pixel feature volume. Same auxiliary structure as ``DinoCLSXYZBaseline``, just multi-view-encoded. Built to ablate the per-voxel volumetric scoring against a global-regression baseline at matched encoder capacity. """ def __init__( self, *, views: list[str], n_window: int = 32, d_model: int = 256, d_cond: int = 128, d_sin_t: int = 48, n_blocks: int = 5, n_gripper_bins: int = 32, n_rot_clusters: int = 64, cross_view_layers: int = 0, cross_view_heads: int = 6, dino_repo: str = DEFAULT_DINO_REPO, dino_weights: str = DEFAULT_DINO_WEIGHTS, ): super().__init__() self.views = list(views) self.n_views = len(views) self.n_window = n_window self.d_model = d_model self.n_gripper_bins = n_gripper_bins self.n_rot_clusters = n_rot_clusters import os, math 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 self._sqrt = math # Sin/cos time conditioning — same as single-view baseline. self.register_buffer("t_sin", sin_pe(n_window, d_sin_t)) self.t_cond_proj = nn.Linear(d_sin_t, d_cond) # Concat-all-views CLS → d_model. self.input_proj = nn.Linear(self.n_views * embed_dim, d_model) self.blocks = nn.ModuleList([ AdaLNZeroBlock(d_model, d_cond, mlp_ratio=4) for _ in range(n_blocks) ]) self.final_norm = nn.LayerNorm(d_model) # Optional cross-view patch attention. When > 0, we re-derive # each view's "CLS" by mean-pooling the cross-attended patch # tokens — keeps the regression head still CLS-only while # letting cross-attention fold scene/wrist info together. self.cross_view_layers = int(cross_view_layers) if self.cross_view_layers > 0 and self.n_views > 1: from .model_volumetric_multiview import CrossViewAttnStack self.cross_view_attn = CrossViewAttnStack( n_layers=self.cross_view_layers, embed_dim=embed_dim, n_heads=cross_view_heads, ) else: self.cross_view_attn = None self.xyz_head = nn.Linear(d_model, 3) self.grip_head = nn.Linear(d_model, n_gripper_bins) self.rot_head = nn.Linear(d_model, n_rot_clusters) def forward(self, rgb: list[torch.Tensor], start_pix=None, K_in=None, T_w2c=None, **_unused) -> dict: """``rgb`` is a list of (B, 3, S, S) tensors, one per view in ``self.views`` order. ``start_pix`` / ``K_in`` / ``T_w2c`` are accepted but unused — the baseline is image-only.""" import math assert isinstance(rgb, (list, tuple)) and len(rgb) == self.n_views B = rgb[0].size(0); T = self.n_window N = self.n_views view_patches: list[torch.Tensor] = [] view_cls: list[torch.Tensor] = [] for k in range(N): out_k = self.backbone.forward_features(rgb[k]) cls = out_k["x_norm_clstoken"] # (B, D) toks = out_k["x_norm_patchtokens"] # (B, P*P, D) P = int(round(math.sqrt(toks.shape[1]))) patch = toks.transpose(1, 2).reshape(B, -1, P, P) view_patches.append(patch) view_cls.append(cls) if self.cross_view_attn is not None: view_patches = self.cross_view_attn(view_patches) # Refresh CLS from the now cross-view-fused patches via mean pool. view_cls = [p.flatten(2).mean(dim=-1) for p in view_patches] cls_concat = torch.cat(view_cls, dim=-1) # (B, N*D) q_in = self.input_proj(cls_concat) # (B, d_model) q_in_bt = q_in.unsqueeze(1).expand(B, T, self.d_model).reshape(B * T, -1) cond_t = self.t_cond_proj(self.t_sin) cond_bt = cond_t.unsqueeze(0).expand(B, T, -1).reshape(B * T, -1) h = q_in_bt for blk in self.blocks: h = blk(h, cond_bt) h = self.final_norm(h) penult = h.view(B, T, self.d_model) return { "xyz_pred": self.xyz_head(penult), # (B, T, 3) "grip_logits": self.grip_head(penult), "rot_logits": self.rot_head(penult), "cls": view_cls[0], # primary CLS for viz "views": self.views, # Per-view DINO patch grids for downstream PCA viz. Same shape # convention as v4's `view_feat_maps`: a list of (B, D, P, P). "view_feat_maps": view_patches, "patch_features": view_patches[0], # primary view convenience alias } def baseline_losses( out: dict, target_xyz: torch.Tensor, # (B, T, 3) world EE position in meters target_grip: torch.Tensor, # (B, T) target_rot: torch.Tensor, # (B, T) xyz_weight: float = 10.0, # MSE on meters; small absolute scale ) -> dict: """MSE on world XYZ + CE on grip + CE on rot, summed.""" B, T = target_grip.shape loss_xyz = F.mse_loss(out["xyz_pred"], target_xyz) 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/xyz_mse": loss_xyz, "loss/grip": loss_grip, "loss/rot": loss_rot, "loss/total": xyz_weight * loss_xyz + loss_grip + loss_rot, }