"""DA3VolumeScene — Depth-Anything-3 backbone + factored YX × Z × T volume head. Same outer interface as ``DinoVolumeScene`` in ``lib/model.py`` — forward returns ``volume_logits``, ``grip_logits``, ``rot_logits``, ``F_refined``, ``patch_features``. Drop-in for train.py / inference.py / deploy_panels.py. Three things DA3 brings over our DINOv3 baseline: - **Camera token** (post-DA3 pose pretraining): used as the global descriptor in place of CLS. DA3 wasn't trained with CLS as a meaningful signal — the camera token IS the global geometric descriptor that the model's own pose decoder consumes via ``feats[-1][1]``. We mirror that. - **DPT decoder**: high-res pixel features by fusing 4 intermediate ViT layers. We add a parallel "feat" branch deep-copied from the depth branch (``refinenet{1..4}``, ``output_conv1``, ``output_conv2``'s first conv + ReLU). Only the final 1×1 conv is fresh — re-targeted to emit ``feat_dim`` channels instead of depth. - **Geometry pretraining** (monocular depth + camera pose) gives the trunk inductive bias the DINOv3 SSL trunk lacks. Camera-token plumbing in DA3's ViT (see ``depth_anything_3/model/dinov2/ vision_transformer.py:300+``): position 0 of the token sequence becomes the cam token at layer ``alt_start`` and is updated by subsequent local+global attention blocks. ``feats[i][0]`` is the patch tokens (cam token + registers stripped), ``feats[i][1]`` is the post-block cam token at output layer i. For ``cat_token=True`` (default da3-small) the cam token has dim ``2 × embed_dim`` (concat of local + global states), matching head ``dim_in``. Volume scoring is unchanged from ``DinoVolumeScene`` — factorized ``score_yx + score_z + score_t`` with **raw dot products** (no L2-norm, no temperature; see ``vault/para/robot/yam/known-bugs.md`` #11). """ from __future__ import annotations import copy import os import sys import types from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F # DA3 setup: sys.path inject + stub the optional utils that pull in heavy # deps DA3 doesn't need at train/inference time (GLB writers, scipy-based # pose alignment). Same minimal monkey-patch the backbones agent uses in # libero/model_da3_volume_v3.py — fragility we can't avoid without vendoring. DA3_REPO_DIR = os.environ.get( "DA3_REPO_DIR", os.path.expanduser("~/lab/da3_repo/src") if os.path.exists(os.path.expanduser("~/lab/da3_repo/src")) else "/data/cameron/da3_repo/src", ) DA3_WEIGHTS_DEFAULT = os.environ.get( "DA3_WEIGHTS_PATH", os.path.expanduser("~/lab/da3_weights") if os.path.exists(os.path.expanduser("~/lab/da3_weights")) else "/data/cameron/da3_weights", ) if DA3_REPO_DIR not in sys.path: sys.path.insert(0, DA3_REPO_DIR) for _n in ("depth_anything_3.utils.export", "depth_anything_3.utils.pose_align"): if _n not in sys.modules: sys.modules[_n] = types.ModuleType(_n) sys.modules["depth_anything_3.utils.export"].export = lambda *a, **k: None sys.modules["depth_anything_3.utils.pose_align"].align_poses_umeyama = ( lambda *a, **k: None) sys.modules["depth_anything_3.utils.pose_align"].batch_align_poses_umeyama = ( lambda *a, **k: None) from depth_anything_3.api import DepthAnything3 # noqa: E402 from depth_anything_3.model.utils.head_utils import ( # noqa: E402 create_uv_grid, position_grid_to_embed, ) # Local helpers (sin_pe + AdaLN-Zero block + all_losses) are shared with # DinoVolumeScene — no duplication. sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from lib.model import AdaLNZeroBlock, sin_pe # noqa: E402, F401 class DA3VolumeScene(nn.Module): """DA3 backbone + cam-token global + parallel feat DPT branch + factored scoring. Defaults: ``n_height_bins=128`` (4× the DINOv3 baseline's 32 — finer height resolution for the same factorized cost), ``feat_dim=32``, ``pred_size=128``. ``z_range`` is auto-derived per-dataset in ``YamScene`` (set ``--z_lo/--z_hi`` unset on the CLI). """ def __init__( self, n_window: int = 16, n_height_bins: int = 128, pred_size: int = 128, feat_dim: int = 32, d_model: int = 256, d_cond: int = 128, d_sin_z: int = 48, d_sin_t: int = 48, n_blocks: int = 5, mlp_hidden: int = 128, # kept for CLI-arg parity with DinoVolumeScene n_gripper_bins: int = 32, n_rot_clusters: int = 64, da3_weights: str = DA3_WEIGHTS_DEFAULT, patch_size: int = 14, ): super().__init__() del mlp_hidden # unused; kept for parity with --mlp_hidden CLI self.n_window = n_window self.n_z = n_height_bins self.pred_size = pred_size self.feat_dim = feat_dim self.d_model = d_model self.d_sin_z = d_sin_z self.d_sin_t = d_sin_t self.patch_size = patch_size self.n_gripper_bins = n_gripper_bins self.n_rot_clusters = n_rot_clusters # ── DA3 backbone + DPT head extraction ────────────────────────── full = DepthAnything3.from_pretrained(da3_weights) self.backbone = full.model.backbone head = full.model.head embed_dim_eff = head.norm.normalized_shape[0] # 768 for da3-small (cat_token=True) self.embed_dim_eff = embed_dim_eff # Shared DPT trunk — modules referenced (not copied) from the DPT # head's depth branch. Gradients flow through them and the backbone; # DA3's geometry pretraining is the starting point, not a freeze. self.dpt_norm = head.norm self.dpt_projects = head.projects self.dpt_resize = head.resize_layers self.dpt_layer_rn = nn.ModuleList([ head.scratch.layer1_rn, head.scratch.layer2_rn, head.scratch.layer3_rn, head.scratch.layer4_rn, ]) # ── New "feat" branch: deep-copy depth branch, swap final 1×1 conv ─ # Mirror the depth branch's refinenets + output_conv1 + first conv # of output_conv2 exactly. Only the final 1×1 is re-targeted from # ``output_dim=2`` (depth + conf) to ``feat_dim`` channels. self.feat_refinenets = nn.ModuleList([ copy.deepcopy(head.scratch.refinenet1), copy.deepcopy(head.scratch.refinenet2), copy.deepcopy(head.scratch.refinenet3), copy.deepcopy(head.scratch.refinenet4), ]) self.feat_out_conv1 = copy.deepcopy(head.scratch.output_conv1) orig_conv2 = head.scratch.output_conv2 # Sequential(Conv, ReLU, Conv→2) head_features_2 = orig_conv2[2].in_channels # 32 across DA3 sizes self.feat_out_conv2_pre = nn.Sequential( copy.deepcopy(orig_conv2[0]), # Conv f//2 → 32, 3×3 copy.deepcopy(orig_conv2[1]), # ReLU ) self.feat_out_conv2_final = nn.Conv2d( head_features_2, feat_dim, kernel_size=1, stride=1, padding=0) nn.init.zeros_(self.feat_out_conv2_final.bias) # Small non-zero init so factored dot products aren't stuck at 0. nn.init.normal_(self.feat_out_conv2_final.weight, std=0.05) del full # drop cam_dec, gs head, etc. # ── PARA heads — same shape as DinoVolumeScene but feeding cam_token, not CLS ─ # cam_norm gives the raw cam token a LayerNorm before the MLP — DINOv3's # ``x_norm_clstoken`` is already normed; DA3's cam_token is the raw # ``out_x[:, :, 0]`` (the model's cam_dec adds its own norm). Mirror that. self.cam_norm = nn.LayerNorm(embed_dim_eff) self.input_proj = nn.Linear(feat_dim + embed_dim_eff, d_model) self.t_cond_proj = nn.Linear(d_sin_t, d_cond) self.blocks = nn.ModuleList([ AdaLNZeroBlock(d_model, d_cond, mlp_ratio=4) for _ in range(n_blocks) ]) self.final_norm = nn.LayerNorm(d_model) d_q = feat_dim + d_sin_z + d_sin_t self.q_head = nn.Linear(d_model, d_q) self.grip_head = nn.Linear(d_model, n_gripper_bins) self.rot_head = nn.Linear(d_model, n_rot_clusters) self.register_buffer("z_sin", sin_pe(n_height_bins, d_sin_z)) self.register_buffer("t_sin", sin_pe(n_window, d_sin_t)) @staticmethod def _add_pos_embed(x: torch.Tensor, W: int, H: int, ratio: float = 0.1) -> torch.Tensor: """UV positional embedding added to a feature map. Vendored from ``DualDPT._add_pos_embed`` so we don't depend on monkey-patching the DA3 head class.""" pw, ph = x.shape[-1], x.shape[-2] pe = create_uv_grid(pw, ph, aspect_ratio=W / H, dtype=x.dtype, device=x.device) pe = position_grid_to_embed(pe, x.shape[1]) * ratio pe = pe.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1) return x + pe def _backbone_features(self, rgb: torch.Tensor): """rgb: (B, 3, H, W). Returns (patch_per_stage, cam_token, ph, pw). Backbone runs in bf16 autocast; outputs cast back to fp32 for the head to keep numerics stable.""" B, _, H, W = rgb.shape x = rgb.unsqueeze(1) # (B, S=1, 3, H, W) ac_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 with torch.autocast(device_type=rgb.device.type, dtype=ac_dtype): feats_tuples, _ = self.backbone( x, cam_token=None, export_feat_layers=[], ref_view_strategy="saddle_balanced", ) ph, pw = H // self.patch_size, W // self.patch_size # feats_tuples[i] = (patch_tokens (B, 1, N, C), cam_token (B, 1, C)) patch_per_stage = [t[0].squeeze(1).float() for t in feats_tuples] cam_token = feats_tuples[-1][1].squeeze(1).float() # (B, C) return patch_per_stage, cam_token, ph, pw def _feat_dpt(self, patch_per_stage, ph: int, pw: int, H: int, W: int): """Shared DPT trunk + our deep-copied feat branch. Returns (F_scene (B, feat_dim, P, P), patch_last_grid (B, C, ph, pw)).""" resized = [] for i in range(4): t = patch_per_stage[i] # (B, N_patches, C) t = self.dpt_norm(t) B, N, C = t.shape t = t.permute(0, 2, 1).reshape(B, C, ph, pw) t = self.dpt_projects[i](t) t = self._add_pos_embed(t, W, H) t = self.dpt_resize[i](t) resized.append(t) l1_rn = self.dpt_layer_rn[0](resized[0]) l2_rn = self.dpt_layer_rn[1](resized[1]) l3_rn = self.dpt_layer_rn[2](resized[2]) l4_rn = self.dpt_layer_rn[3](resized[3]) # Feat fusion (mirrors DualDPT._fuse main path; our ModuleList index # 0..3 maps to refinenet1..4). out = self.feat_refinenets[3](l4_rn, size=l3_rn.shape[2:]) out = self.feat_refinenets[2](out, l3_rn, size=l2_rn.shape[2:]) out = self.feat_refinenets[1](out, l2_rn, size=l1_rn.shape[2:]) out = self.feat_refinenets[0](out, l1_rn) # default 2× upsample out = self.feat_out_conv1(out) # (B, 32, ~2*l1_size, ~) # Upsample DIRECTLY to pred_size — skip DA3's native H/down_ratio # output (which would be 504 at S=504, down_ratio=1) since we'd # just downsample it again. Saves ~16× compute on the final convs. out = F.interpolate(out, size=(self.pred_size, self.pred_size), mode="bilinear", align_corners=True) out = self._add_pos_embed(out, W, H) out = self.feat_out_conv2_pre(out) # (B, 32, P, P) F_scene = self.feat_out_conv2_final(out) # (B, feat_dim, P, P) # Last-stage patch grid for PCA viz parity with DinoVolumeScene. last = patch_per_stage[-1] # (B, N, C) Bp, Np, Cp = last.shape patch_grid = last.permute(0, 2, 1).reshape(Bp, Cp, ph, pw) return F_scene, patch_grid def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor) -> dict: """rgb (B, 3, S, S); start_pix (B, 2) xy in S-coords.""" B, _, S, _ = rgb.shape T, Z, P = self.n_window, self.n_z, self.pred_size patch_per_stage, cam_token, ph, pw = self._backbone_features(rgb) F_scene, patch_last = self._feat_dpt(patch_per_stage, ph, pw, S, S) # Sample F_scene at start_pix (S-coords → pred-grid coords). 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, feat_dim) # Cam token → norm → concat with eef_feat → d_model. cam_n = self.cam_norm(cam_token) q_in = self.input_proj(torch.cat([eef_feat, cam_n], dim=-1)) # (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) q_spatial = self.q_head(penult) q_F = q_spatial[..., :self.feat_dim] q_z = q_spatial[..., self.feat_dim:self.feat_dim + self.d_sin_z] q_t = q_spatial[..., self.feat_dim + self.d_sin_z:] # Raw dot products — NO L2-norm, NO temperature (see known-bugs.md #11). score_yx = torch.einsum("btc,bchw->bthw", q_F, F_scene) score_z = torch.einsum("btc,zc->btz", q_z, self.z_sin) score_t = torch.einsum("btc,tc->bt", q_t, self.t_sin) vol = (score_yx[:, :, None, :, :] + score_z[:, :, :, None, None] + score_t[:, :, None, None, None]) # (B, T, Z, P, P) return { "volume_logits": vol, "grip_logits": self.grip_head(penult), "rot_logits": self.rot_head(penult), "F_refined": F_scene, "patch_features": patch_last, } if __name__ == "__main__": import time device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"device = {device}") m = DA3VolumeScene(n_window=16, n_height_bins=128, pred_size=128, feat_dim=32, d_model=256, n_blocks=5).to(device).eval() n_t = sum(p.numel() for p in m.parameters() if p.requires_grad) print(f"Trainable: {n_t:,} ({n_t / 1e6:.1f}M)") rgb = torch.rand(1, 3, 504, 504, device=device) sp = torch.tensor([[252.0, 252.0]], device=device) with torch.no_grad(): t0 = time.time() out = m(rgb, sp) t1 = time.time() for k, v in out.items(): if hasattr(v, "shape"): print(f" {k}: {tuple(v.shape)}") print(f"forward = {(t1 - t0) * 1000:.1f} ms") if device.type == "cuda": print(f"peak: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")