"""DA3 spatial-language branch for pi0.5 (JAX/Flax nnx port of the X-VLA addon). Faithful reimplementation of the TRAINABLE modules from `DA3-XVLA-cache/models/spatial_language.py` (H=1024, GIANT C=1536, grid 18x24=432, perceiver tokens 128/96/96, 7-ch scale-aware ray, ModernBERT language fusion). The FROZEN DA3 backbone + ModernBERT run offline (features precached); this module consumes their outputs as arrays and produces per-view "banks" that are cross-attended into the action-expert's late blocks (see gemma.py `SpatialActionInjection`). Only the bank BUILDER lives here (nnx, a submodule of Pi0). The injection layer lives in gemma.py (linen, inside the action-expert scan). Both use identical X-VLA math. Reference math (verified by the understand-phase spec): - ResidualCrossAttention: out = q_hidden + scale * MHA(LN_q(q_hidden), LN_kv(kv), LN_kv(kv)) - MHA matches torch nn.MultiheadAttention: separate q/k/v/out Linears w/ bias, 1/sqrt(head_dim). - GELU is the tanh approximation everywhere; LayerNorm eps=1e-5. - Perceiver residual adds the RAW learned query (not the normalized one). - View order everywhere: 0=main/countertop, 1=left wrist, 2=right wrist. """ import math as _math import os as _os import einops import flax.nnx as nnx import jax import jax.numpy as jnp import openpi.shared.array_typing as at # --------------------------------------------------------------------------- # primitives # --------------------------------------------------------------------------- def _gelu(x): return nnx.gelu(x, approximate=True) # tanh approximation (matches torch GELU(approximate="tanh")) class MHACrossAttn(nnx.Module): """Multi-head cross-attention matching torch nn.MultiheadAttention math (no residual, no norm).""" def __init__(self, dim: int, num_heads: int, *, logit_gain: bool = False, logit_gain_init: float = 32.0, logit_gain_max: float = 16.0, qk_norm: bool = False, rngs: nnx.Rngs): assert dim % num_heads == 0 self.num_heads = num_heads self.head_dim = dim // num_heads self.q_proj = nnx.Linear(dim, dim, rngs=rngs) self.k_proj = nnx.Linear(dim, dim, rngs=rngs) self.v_proj = nnx.Linear(dim, dim, rngs=rngs) self.out_proj = nnx.Linear(dim, dim, rngs=rngs) # QK-NORM: per-head RMSNorm on Q and K BEFORE the dot product. Measured at step 10k without # it: raw |logit| reached 6653 (normal is O(1-10)), softmax saturated to one-hot # (entropy 0.007 vs uniform 5.78, effective tokens attended = 1.0/324, max prob 0.997). # A saturated softmax has a vanishing Jacobian, so the attention pattern then FREEZES and # cannot recover. Nothing else bounds logit scale here: q_proj/k_proj grow freely under the # high-LR 'core' group with weight_decay 1e-10. Normalizing Q,K to unit RMS caps # |q.k|/sqrt(head_dim) at O(1) structurally, no matter how large the projections get -- # which also makes logit_gain behave as the temperature it was meant to be. self.qk_norm = bool(qk_norm) if self.qk_norm: self.q_ln = nnx.RMSNorm(self.head_dim, rngs=rngs) self.k_ln = nnx.RMSNorm(self.head_dim, rngs=rngs) # Learnable per-head gain on the attention logits (same fix already used for the injection). # With random-init queries the q.k logits are ~0, so softmax over 432 patches is near-uniform; # that (a) makes every query read the SAME mean(V) and (b) starves dL/dQ,K (Jacobian ~1/432) # so the queries never train. exp(log_gain) with init 32 sharpens attention at init, which # both diversifies the per-query reads and unfreezes the Q/K gradients. self.logit_gain = bool(logit_gain) if self.logit_gain: # CLAMPED: exp(log_gain) is unbounded, and this param sits in the high-LR 'core' group. # Unclamped, a few large updates make exp(log_gain) blow up -> logits overflow -> NaN # (observed: gain 32 already gives max|logit| ~168 vs ~5 baseline). jnp.clip also zeroes # the gradient outside the range, so the parameter self-arrests instead of running away. self.log_gain = nnx.Param(jnp.full((num_heads,), jnp.log(jnp.asarray(logit_gain_init, jnp.float32)))) # plain Python math (NOT jnp): __init__ runs under jit tracing, so float(jnp...) # raises ConcretizationTypeError. This is a static constant, no tracing needed. self.log_gain_max = _math.log(max(float(logit_gain_max), 1.0)) def __call__(self, q, kv, key_pad_mask=None, kv_addr=None, attn_bias=None): # q:[b,Lq,d] kv:[b,Lk,d] key_pad_mask:[b,Lk] True=pad (ignored) # kv_addr:[b|1,Lk,d] optional ADDRESS stream (K/V split): added to the keys ONLY, so it steers # routing (which tokens each query reads) but is structurally excluded from the values -- an # input-independent address can never leak into the output and dilute per-sample content. # attn_bias: additive logit bias. [h,Lq,Lk] = same for every sample (e.g. the static grid # locality prior); [b,h,Lq,Lk] = PER-SAMPLE (e.g. the EE-anchored prior, whose anchors are the # wrist-camera centres and therefore move with the arms). h = self.num_heads Q = einops.rearrange(self.q_proj(q), "b l (h d) -> b h l d", h=h) k_in = kv if kv_addr is None else kv + kv_addr K = einops.rearrange(self.k_proj(k_in), "b l (h d) -> b h l d", h=h) V = einops.rearrange(self.v_proj(kv), "b l (h d) -> b h l d", h=h) if self.qk_norm: # bounds |q.k| structurally; see __init__ for the saturation evidence Q = self.q_ln(Q) K = self.k_ln(K) logits = jnp.einsum("bhqd,bhkd->bhqk", Q, K) * (self.head_dim**-0.5) if self.logit_gain: g = jnp.clip(self.log_gain.value, -self.log_gain_max, self.log_gain_max) logits = logits * jnp.exp(g)[None, :, None, None].astype(logits.dtype) if attn_bias is not None: # ndim 3 -> [h,Lq,Lk] shared across the batch; ndim 4 -> [b,h,Lq,Lk] already per-sample. ab = attn_bias[None] if attn_bias.ndim == 3 else attn_bias logits = logits + ab.astype(logits.dtype) if key_pad_mask is not None: logits = jnp.where(key_pad_mask[:, None, None, :], jnp.asarray(-1e30, logits.dtype), logits) probs = jax.nn.softmax(logits.astype(jnp.float32), axis=-1).astype(logits.dtype) ctx = jnp.einsum("bhqk,bhkd->bhqd", probs, V) ctx = einops.rearrange(ctx, "b h q d -> b q (h d)") return self.out_proj(ctx) class ResidualCrossAttn(nnx.Module): """Pre-LN residual cross-attention: out = q_hidden + scale * MHA(LN_q(q_hidden), LN_kv(kv)).""" def __init__(self, dim: int, num_heads: int, *, logit_gain: bool = False, logit_gain_init: float = 32.0, logit_gain_max: float = 16.0, norm_attn_out: bool = False, qk_norm: bool = False, rngs: nnx.Rngs): self.q_norm = nnx.LayerNorm(dim, epsilon=1e-5, rngs=rngs) self.kv_norm = nnx.LayerNorm(dim, epsilon=1e-5, rngs=rngs) self.attn = MHACrossAttn(dim, num_heads, logit_gain=logit_gain, logit_gain_init=logit_gain_init, logit_gain_max=logit_gain_max, qk_norm=qk_norm, rngs=rngs) # The residual adds the RAW query. If ||attn_out|| >> ||q|| (measured ~500 vs ~1.6, i.e. 300:1) # the shared attention output swamps per-query identity and every output collapses to # mlp(q_i + const) with cos ~ 1.0. Normalizing the attention output before the residual puts # the two terms on comparable scale, preserving query identity even if attention stays uniform. self.out_norm = nnx.LayerNorm(dim, epsilon=1e-5, rngs=rngs) if norm_attn_out else None def __call__(self, q_hidden, kv_hidden, key_pad_mask=None, residual_scale: float = 1.0, kv_addr=None, attn_bias=None): q = self.q_norm(q_hidden) kv = self.kv_norm(kv_hidden) # kv_addr bypasses kv_norm deliberately: the payload is normalized for stable value scale, # while the address keeps its own (MLP-output) scale as a routing bias on the keys. out = self.attn(q, kv, key_pad_mask=key_pad_mask, kv_addr=kv_addr, attn_bias=attn_bias) if self.out_norm is not None: out = self.out_norm(out) return q_hidden + residual_scale * out class ResidualMlp(nnx.Module): """Pre-LN residual MLP: x + Linear2(gelu(Linear1(LN(x)))).""" def __init__(self, dim: int, mlp_ratio: float, *, rngs: nnx.Rngs): hidden = int(dim * mlp_ratio) self.norm = nnx.LayerNorm(dim, epsilon=1e-5, rngs=rngs) self.fc1 = nnx.Linear(dim, hidden, rngs=rngs) self.fc2 = nnx.Linear(hidden, dim, rngs=rngs) def __call__(self, x): return x + self.fc2(_gelu(self.fc1(self.norm(x)))) class ProjLN(nnx.Module): """Linear(in->H) -> gelu -> Linear(H->H) -> LayerNorm(H). Used for layer projectors & t5_projector.""" def __init__(self, in_dim: int, dim: int, *, rngs: nnx.Rngs): self.fc1 = nnx.Linear(in_dim, dim, rngs=rngs) self.fc2 = nnx.Linear(dim, dim, rngs=rngs) self.norm = nnx.LayerNorm(dim, epsilon=1e-5, rngs=rngs) def __call__(self, x): return self.norm(self.fc2(_gelu(self.fc1(x)))) class Mlp2(nnx.Module): """Linear(in->hidden) -> gelu -> Linear(hidden->out). Used for ray_mlp & pos2d_mlp (no LN).""" def __init__(self, in_dim: int, hidden: int, out_dim: int, *, rngs: nnx.Rngs): self.fc1 = nnx.Linear(in_dim, hidden, rngs=rngs) self.fc2 = nnx.Linear(hidden, out_dim, rngs=rngs) def __call__(self, x): return self.fc2(_gelu(self.fc1(x))) def fourier_encode(x, num_bands: int): """[..., D] in [-1,1] -> [..., D*(1+2*num_bands)]: the raw value plus sin/cos at 2^k*pi. Raw low-dimensional coordinates through an MLP can only express smooth functions of position (spectral bias), so neighbouring patches collapse to near-identical embeddings. Lifting to a Fourier basis first makes nearby coordinates far apart in the high bands. Band count is chosen against real precision, not maximum precision: with a ~1.3 m half-range, 10 bands resolve ~5 mm, while GT depth is 1 mm-quantised and then area-averaged over a 14x14 patch (~1-5 cm of surface), so more bands would encode noise. """ freqs = (2.0 ** jnp.arange(num_bands)) * jnp.pi xb = x[..., None] * freqs # [..., D, K] enc = jnp.concatenate([jnp.sin(xb), jnp.cos(xb)], axis=-1) return jnp.concatenate([x, enc.reshape(*x.shape[:-1], -1)], axis=-1) class SpatialConditioner(nnx.Module): """token = LN( W [ s ; (1+gamma(s)) * fused + beta(s) ] ), s = MLP(spatial_vector).""" def __init__(self, dim: int, in_dim: int, *, film: bool = True, use_da3: bool = True, hidden: int = 256, rngs: nnx.Rngs): self.fc1 = nnx.Linear(in_dim, hidden, rngs=rngs) self.fc2 = nnx.Linear(hidden, dim, rngs=rngs) self.film = bool(film) self.use_da3 = bool(use_da3) if self.film: # zero-init so the token starts as the plain concat and geometry modulation ramps in self.gamma = nnx.Linear(dim, dim, rngs=rngs) self.beta = nnx.Linear(dim, dim, rngs=rngs) for lin in (self.gamma, self.beta): lin.kernel.value = jnp.zeros_like(lin.kernel.value) lin.bias.value = jnp.zeros_like(lin.bias.value) self.out = nnx.Linear(2 * dim, dim, rngs=rngs) self.norm = nnx.LayerNorm(dim, epsilon=1e-5, rngs=rngs) def __call__(self, svec, fused): s = self.fc2(_gelu(self.fc1(svec))) if not self.use_da3: f = jnp.zeros_like(s) # geometry-only ablation arm else: f = fused if self.film: f = (1.0 + self.gamma(s)) * f + self.beta(s) return self.norm(self.out(jnp.concatenate([s, f], axis=-1))) def _locality_dist2(num_queries, h, w): """Squared distance [K, h*w] between each query's tiled anchor and each patch position, both in a normalized [0,1]^2 grid (patches row-major to match _fuse_layers 'b (h w)').""" import numpy as _np ys, xs = _np.meshgrid(_np.linspace(0.0, 1.0, h), _np.linspace(0.0, 1.0, w), indexing="ij") patch = _np.stack([ys.ravel(), xs.ravel()], axis=-1) # [h*w, 2] ar = int(_np.ceil(_np.sqrt(num_queries))); ac = int(_np.ceil(num_queries / ar)) ay, ax = _np.meshgrid(_np.linspace(0.0, 1.0, ar), _np.linspace(0.0, 1.0, ac), indexing="ij") anch = _np.stack([ay.ravel(), ax.ravel()], axis=-1)[:num_queries] # [K, 2] return (((anch[:, None, :] - patch[None, :, :]) ** 2).sum(-1)).astype(_np.float32) # [K, h*w] class PerceiverDownsampler(nnx.Module): """432 grid tokens -> K learned-query tokens (single cross-attn + residual MLP).""" def __init__(self, dim: int, num_queries: int, num_heads: int, *, query_std: float = 0.02, logit_gain: bool = False, logit_gain_init: float = 32.0, logit_gain_max: float = 16.0, norm_attn_out: bool = False, qk_norm: bool = False, norm_out: bool = False, locality: bool = False, grid_hw: tuple = (18, 24), locality_gamma_init: float = 4.0, ee_anchor: bool = False, ee_query_frac: float = 0.25, ee_gamma_init: float = 4.0, rngs: nnx.Rngs): key = rngs.params() self.query = nnx.Param(jax.random.normal(key, (1, num_queries, dim)) * query_std) self.xattn = ResidualCrossAttn(dim, num_heads, logit_gain=logit_gain, logit_gain_init=logit_gain_init, logit_gain_max=logit_gain_max, norm_attn_out=norm_attn_out, qk_norm=qk_norm, rngs=rngs) self.mlp = ResidualMlp(dim, mlp_ratio=2.0, rngs=rngs) # LOCALITY: each query gets a fixed anchor tiling the grid; a learnable per-head gamma biases # the attention logits by -gamma*dist2 so each of the K tokens preferentially reads its own # neighborhood (a local descriptor) instead of a global average -- fixes over-averaging while # staying flexible (gamma can shrink toward global if content demands). self.locality = bool(locality) if self.locality: self._loc_nq = int(num_queries) # ints only (nnx rejects bare array attrs); self._loc_gh = (int(grid_hw[0]), int(grid_hw[1])) # dist2 is recomputed (static) in __call__ self.loc_log_gamma = nnx.Param(jnp.full((num_heads,), _math.log(max(locality_gamma_init, 1e-3)))) # EE ANCHORING: reserve the FIRST n_ee queries and re-anchor them onto the two end-effectors # (wrist-camera centres in the robot frame) instead of a fixed grid cell. Their bias is # -gamma_ee*||p_patch - p_ee||^2 in METRES, computed per sample, so these tokens always read the # geometry around the hands. The remaining K-n_ee queries keep the grid locality, so scene # context is not lost. Half the reserved slots track the left EE, half the right. self.ee_anchor = bool(ee_anchor) self._n_ee = min(int(round(num_queries * float(ee_query_frac))), num_queries) if ee_anchor else 0 if self.ee_anchor and self._n_ee > 0: self.ee_log_gamma = nnx.Param(jnp.full((num_heads,), _math.log(max(ee_gamma_init, 1e-3)))) # FIX 4: bound the perceiver output. Measured without it: the residual MLP amplified a # unit-rms input to rms 1790 (x1900). Nothing penalized that -- the injection's kv_norm makes # downstream scale irrelevant and weight_decay was 1e-10 -- so the block became an # unconstrained amplifier whose output was ~92% batch-constant. self.out_ln = nnx.LayerNorm(dim, epsilon=1e-5, rngs=rngs) if norm_out else None def __call__(self, tokens, addr=None, token_embed=None, ee_dist2=None): # tokens = PAYLOAD (per-sample content: DA3 latents + depth enc). addr = optional ADDRESS # stream (pos/ray/view annotations) -> keys only; see MHACrossAttn.kv_addr. # token_embed [1,K,H]: per-output-token identity added to the QUERY (not the final bank). This # shapes WHICH patches each of the K queries reads, so it produces per-token-distinct AND # per-sample-varying output -- unlike a post-hoc constant it survives bank-centering. b = tokens.shape[0] q = jnp.broadcast_to(self.query.value, (b, *self.query.value.shape[1:])) if token_embed is not None: q = q + token_embed bias = None if self.locality: gamma = jnp.exp(self.loc_log_gamma.value) # [h] >0 dist2 = jnp.asarray(_locality_dist2(self._loc_nq, *self._loc_gh)) # static const [K, Lk] bias = -gamma[:, None, None] * dist2[None] # [h, K, Lk] if self.ee_anchor and self._n_ee > 0 and ee_dist2 is not None: # ee_dist2 [b,2,Lk] = squared metric distance from every patch's 3D point to the left(0) # and right(1) end-effector. Alternate the reserved slots between the two hands. n_ee, Lk = self._n_ee, tokens.shape[1] h = self.xattn.attn.num_heads side = jnp.asarray([j % 2 for j in range(n_ee)]) # [n_ee] 0=left,1=right g_ee = jnp.exp(self.ee_log_gamma.value) # [h] sel = jnp.take(ee_dist2, side, axis=1) # [b,n_ee,Lk] ee_rows = -g_ee[None, :, None, None] * sel[:, None, :, :] # [b,h,n_ee,Lk] if bias is not None: rest = jnp.broadcast_to(bias[None, :, n_ee:, :], (b, h, self._loc_nq - n_ee, Lk)) else: rest = jnp.zeros((b, h, q.shape[1] - n_ee, Lk), ee_rows.dtype) bias = jnp.concatenate([ee_rows, rest.astype(ee_rows.dtype)], axis=2) # [b,h,K,Lk] z = self.xattn(q, tokens, residual_scale=1.0, kv_addr=addr, attn_bias=bias) # residual adds RAW q out = self.mlp(z) return self.out_ln(out) if self.out_ln is not None else out class LanguageFusionStack(nnx.Module): """N x [cross-attn(bank, lang) + residual-MLP], with language padding mask.""" def __init__(self, dim: int, depth: int, num_heads: int, *, qk_norm: bool = False, rngs: nnx.Rngs): self.layers = [ (ResidualCrossAttn(dim, num_heads, qk_norm=qk_norm, rngs=rngs), ResidualMlp(dim, mlp_ratio=4.0, rngs=rngs)) for _ in range(depth) ] def __call__(self, geo, lang_tokens, lang_pad_mask): for xattn, mlp in self.layers: geo = xattn(geo, lang_tokens, key_pad_mask=lang_pad_mask, residual_scale=1.0) geo = mlp(geo) return geo def _cam_pose_feat(ext): """Camera-pose feature [b,12] from a w2c extrinsic [b,4,4]: R_c2w flattened (9) + camera center (3). Gives the cross-view fusion the RELATIVE viewpoints so it can reason across cameras geometrically.""" R = ext[:, :3, :3] # R_w2c t = ext[:, :3, 3] Rc2w = jnp.swapaxes(R, -1, -2) center = -jnp.einsum("bij,bj->bi", Rc2w, t) # camera center in world return jnp.concatenate([Rc2w.reshape(ext.shape[0], 9), center], axis=-1).astype(jnp.float32) class CrossViewFusion(nnx.Module): """Self-attention over the CONCATENATED per-view tokens so the three views exchange 3D information (grounded by per-view camera pose), turning three separate 2.5D banks into one integrated scene.""" def __init__(self, dim: int, num_heads: int, depth: int, *, qk_norm: bool = False, rngs: nnx.Rngs): self.blocks = [ (ResidualCrossAttn(dim, num_heads, qk_norm=qk_norm, rngs=rngs), ResidualMlp(dim, mlp_ratio=4.0, rngs=rngs)) for _ in range(depth) ] def __call__(self, x): # x [b, N_total, H] for attn, mlp in self.blocks: x = attn(x, x, residual_scale=1.0) # self-attention (q == kv) x = mlp(x) return x # --------------------------------------------------------------------------- # geometry helpers # --------------------------------------------------------------------------- def compute_world_ray_6d(ray_local, ext_w2c): """ray_local [b,3,h,w] cam-local unit dir; ext_w2c [b,4,4] OpenCV world->cam. Returns [b,6,h,w] = concat([origin_world(camera center), dir_world]). """ R_w2c = ext_w2c[:, :3, :3] # [b,3,3] t_w2c = ext_w2c[:, :3, 3] # [b,3] R_c2w = jnp.swapaxes(R_w2c, -1, -2) pos_world = -jnp.einsum("bij,bj->bi", R_c2w, t_w2c) # [b,3] camera center in world b, _, h, w = ray_local.shape dir_world = jnp.einsum("bij,bjk->bik", R_c2w, ray_local.reshape(b, 3, h * w)).reshape(b, 3, h, w) origin = jnp.broadcast_to(pos_world[:, :, None, None], (b, 3, h, w)) return jnp.concatenate([origin, dir_world], axis=1) # [b,6,h,w] def _cam_center(ext): """Camera centre in the reference (robot) frame from a robot->cam extrinsic [b,4,4]: -R_c2w @ t. For the WRIST-mounted realsense cameras this is the end-effector position.""" R_c2w = jnp.swapaxes(ext[:, :3, :3], -1, -2) return -jnp.einsum("bij,bj->bi", R_c2w, ext[:, :3, 3]).astype(jnp.float32) # [b,3] def _info_nce(za, zb, temp: float): """Symmetric (CLIP-style) InfoNCE between two L2-normalised embeddings [b,D]. The positive pair is the diagonal (same sample); every other sample in the batch is a negative.""" za = za / (jnp.linalg.norm(za, axis=-1, keepdims=True) + 1e-6) zb = zb / (jnp.linalg.norm(zb, axis=-1, keepdims=True) + 1e-6) logits = jnp.einsum("id,jd->ij", za, zb) / temp # [b,b] labels = jnp.arange(logits.shape[0]) ce = lambda lg: -jnp.mean(jax.nn.log_softmax(lg, axis=-1)[labels, labels]) return 0.5 * (ce(logits) + ce(logits.T)) def _grid_coords(h: int, w: int): v = 2.0 * jnp.arange(h) / (h - 1) - 1.0 u = 2.0 * jnp.arange(w) / (w - 1) - 1.0 yy, xx = jnp.meshgrid(v, u, indexing="ij") return jnp.stack([xx, yy], axis=-1).reshape(1, h * w, 2) # [1,432,2] (x=u, y=v), row-major # --------------------------------------------------------------------------- # bank builder # --------------------------------------------------------------------------- # Perceiver bank token count per view (main/left/right). Scalable via DA3_PERC_TOKEN_MULT for # capacity experiments; default 1.0 keeps the canonical 128/96/96 so existing checkpoints load # unchanged. NOTE: a different count is an INCOMPATIBLE architecture (perceiver query / bank_token # / cross-view shapes change) -- only for FRESH training, never for resuming/serving a 128/96/96 ckpt. _PERC_TOKEN_MULT = float(_os.environ.get("DA3_PERC_TOKEN_MULT", "1.0")) def _perc_k(base: int) -> int: return max(1, int(round(base * _PERC_TOKEN_MULT))) _VIEWS = (("main", 0, _perc_k(128)), ("left", 1, _perc_k(96)), ("right", 2, _perc_k(96))) class SpatialBankBuilder(nnx.Module): """Cached DA3 (feats/ray/depth) + extrinsics + ModernBERT feats -> 3 per-view banks.""" def __init__( self, *, hidden_dim: int = 1024, da3_channels: int = 1536, num_layers: int = 4, grid_hw: tuple[int, int] = (18, 24), lang_dim: int = 1024, # ModernBERT-large last_hidden width (768) -> set by config num_heads: int = 8, lang_fusion_depth: int = 2, perceiver_query_std: float = 0.02, qk_norm: bool = False, perceiver_norm_out: bool = False, pos_emb_scale: float = 1.0, perceiver_logit_gain: bool = False, perceiver_logit_gain_init: float = 32.0, perceiver_logit_gain_max: float = 16.0, perceiver_norm_attn_out: bool = False, bank_token_embed: bool = False, bank_center: bool = False, aux_geom_head: bool = False, depth_target_only: bool = False, kv_split: bool = False, depth_dropout: float = 0.0, perc_locality: bool = False, cross_view: bool = False, cross_view_depth: int = 2, bank_token_embed_query: bool = True, use_depth_conf: bool = False, use_pose_enc: bool = False, use_cam_tokens: bool = False, cam_token_dim: int = 2048, pose_enc_dim: int = 9, feat_input_norm: bool = False, use_point_map: bool = False, depth_aware_crossview: bool = False, ee_anchor: bool = False, ee_query_frac: float = 0.25, ee_gamma_init: float = 4.0, ee_max_dist2: float = 25.0, infonce: bool = False, infonce_temp: float = 0.07, infonce_dim: int = 128, infonce_pool_k: int = 16, use_perceiver: bool = True, spatial_vec: bool = False, spatial_film: bool = True, spatial_use_da3: bool = True, fourier_bands: int = 10, point_centre: tuple = (1.032, 0.527, 1.064), point_centre_ee_l: tuple = (0.811, 0.278, 0.332), point_centre_ee_r: tuple = (0.268, 0.895, 0.154), point_scale: float = 1.30, point_max_depth: float = 5.0, rngs: nnx.Rngs, ): # placement of bank_token_embeds: True (new) = added to the perceiver QUERY (center-surviving); # False (old) = added POST-fusion (dead under bank_center). Set False to faithfully evaluate # checkpoints trained before the move (e.g. spatretrain/strongbase/kvsplit_desk). self._bte_query = bool(bank_token_embed_query) # K/V SPLIT (2026-07-23): separate ADDRESS from PAYLOAD instead of one additive sum. # payload (values) = DA3 latents + depth encoding -- what flows into the bank # address (keys) = pos_emb + ray_emb + view_emb -- where it is; routing only # In the summed design the constant "where" terms enter the value stream and, under broad # attention, average into an input-independent constant (the measured collapse). With the # split, addresses are structurally excluded from the output: constants can route, but only # per-sample content can flow. depth moves to the payload (per-sample geometry content); # the Plucker ray (camera geometry) stays as address. self.kv_split = bool(kv_split) # depth_dropout: during training, zero the depth encoding for this fraction of samples so the # bank cannot rely on the explicit depth channel alone -- the DA3 features must carry the # geometry too. Applied only when a dropout rng is passed (training); inference keeps depth. self.depth_dropout = float(depth_dropout) self.bank_center = bool(bank_center) # AUX GEOMETRY HEAD (2026-07-23): decode the PERCEIVER token output back to per-patch log-depth # (grid-position queries cross-attend to the K perceiver tokens). Supervised by the DA3 depth we # already have (ray_flat[...,6]), this FORCES the perceiver output to carry per-sample scene # geometry regardless of whether the action loss rewards it -- the guaranteed fix for the # "geometry read but unused" verdict. Shared across views; queries are the (constant) grid # positions so the prediction varies only through the per-sample perceiver tokens. self.aux_geom_head = bool(aux_geom_head) # depth TARGET-ONLY mode: zero the log-depth channel in the ray7 INPUT so depth is never given # to the network -- only used as the aux target. Without this the aux task is circular (depth # in -> depth out = a trivial autoencoder through the perceiver bottleneck, satisfiable without # reading the DA3 features at all). With it, the ONLY path to the target is extracting depth # from the DA3 features -> the aux loss forces genuine feature use. Plucker ray dirs (ch 0-5) # remain as input: they are camera geometry, not the answer. self.depth_target_only = bool(depth_target_only) H = hidden_dim self.hidden_dim = H # With the perceiver removed the bank IS the patch grid, so every view contributes # grid_h*grid_w tokens instead of its perceiver's K. self.use_perceiver = bool(use_perceiver) self._ntok = {name: (k if self.use_perceiver else grid_hw[0] * grid_hw[1]) for name, _, k in _VIEWS} self.num_layers = num_layers self.grid_hw = grid_hw # (a) per-tap projectors + layer embed + fuse self.layer_projectors = [ProjLN(da3_channels, H, rngs=rngs) for _ in range(num_layers)] self.layer_embed = nnx.Param(jax.random.normal(rngs.params(), (num_layers, H)) * 0.02) self.layer_fuse = nnx.Linear(num_layers * H, H, rngs=rngs) # (b) ray encoder. kv_split: Plucker-6 only (address) + separate depth encoder (payload). # legacy: scale-aware ray (Plucker-6 + log-depth = 7) summed into everything. if self.kv_split: self.ray_mlp = Mlp2(6, 256, H, rngs=rngs) self.depth_mlp = Mlp2(1, 256, H, rngs=rngs) else: self.ray_mlp = Mlp2(7, 256, H, rngs=rngs) # (c) 2D grid pos + per-view embedding self.pos2d_mlp = Mlp2(2, 256, H, rngs=rngs) self.view_embed = nnx.Embed(3, H, rngs=rngs) # (d) language projector (ModernBERT feat -> H). Only built when language fusion is on: # with lang_fusion_depth=0 it would be a 2.1 M dead branch receiving no gradient. self.lang_fusion_depth = int(lang_fusion_depth) self.t5_projector = ProjLN(lang_dim, H, rngs=rngs) if self.lang_fusion_depth > 0 else None # FIX 5: pos_emb is INPUT-INDEPENDENT and was measured at rms 5.03 vs the DA3-derived # signal's 4.38 -- the constant was LARGER than the content it annotates, diluting # per-sample diversity 0.474 -> 0.270 before the perceiver even ran. Scale it down so # position annotates content instead of dominating it. self.pos_emb_scale = float(pos_emb_scale) # (e) per-view perceiver + language fusion. perc_locality anchors each query to a grid region. self.perceivers = None if not self.use_perceiver else { name: PerceiverDownsampler(H, k, num_heads, query_std=perceiver_query_std, logit_gain=perceiver_logit_gain, logit_gain_init=perceiver_logit_gain_init, logit_gain_max=perceiver_logit_gain_max, norm_attn_out=perceiver_norm_attn_out, qk_norm=qk_norm, norm_out=perceiver_norm_out, locality=perc_locality, grid_hw=grid_hw, ee_anchor=ee_anchor, ee_query_frac=ee_query_frac, ee_gamma_init=ee_gamma_init, rngs=rngs) for name, _, k in _VIEWS } # EE anchoring needs the metric point map (to measure patch->hand distance in the robot frame). self.ee_anchor = bool(ee_anchor) self.ee_max_dist2 = float(ee_max_dist2) # InfoNCE: two small projection heads onto a shared unit sphere -- one reads the pooled bank, # the other the pooled metric point map. Cheap (a [b,D] x [D,b] matmul, D=128). self.infonce = bool(infonce) self.infonce_temp = float(infonce_temp) self.infonce_pool_k = int(infonce_pool_k) if self.infonce: self.nce_bank_proj = Mlp2(H * len(_VIEWS), 512, infonce_dim, rngs=rngs) self.nce_geom_proj = Mlp2(3 * infonce_pool_k * len(_VIEWS), 512, infonce_dim, rngs=rngs) self.lang_fusers = ({name: LanguageFusionStack(H, lang_fusion_depth, num_heads, qk_norm=qk_norm, rngs=rngs) for name, _, _ in _VIEWS} if self.lang_fusion_depth > 0 else None) # (e2) CROSS-VIEW 3D FUSION: after the per-view perceivers, add a camera-pose embed to each # view's tokens, concatenate, and self-attend so views exchange 3D info; then split back. self.cross_view = bool(cross_view) if self.cross_view: self.cam_pose_mlp = Mlp2(12, 256, H, rngs=rngs) self.cross_view_fusion = CrossViewFusion(H, num_heads, cross_view_depth, qk_norm=qk_norm, rngs=rngs) # --- VGGT-Omega enrichments (all gated; DA3 path leaves them off) --- # depth_conf: VGGT per-patch confidence -> a payload reliability channel (added to the values, # so the bank can down-weight geometry where VGGT is uncertain). self.use_depth_conf = bool(use_depth_conf) if self.use_depth_conf: self.conf_mlp = Mlp2(1, 256, H, rngs=rngs) # pose_enc: VGGT learned camera encoding (trans+quat+fov) -> added to the cross-view camera # feature (a learned pose signal alongside the hand-built R|t feature). self.use_pose_enc = bool(use_pose_enc) if self.use_pose_enc: self.pose_enc_mlp = Mlp2(pose_enc_dim, 256, H, rngs=rngs) # cam_tokens: VGGT camera+register global tokens -> projected and APPENDED to each view's final # bank (global scene/camera context the action expert can attend to). Appended after fusion so # they never disturb the perceiver locality grid or the cross-view token split. self.use_cam_tokens = bool(use_cam_tokens) if self.use_cam_tokens: # VGGT camera/register tokens carry ViT massive-activation outliers (absmax ~180); LayerNorm # the raw tokens BEFORE the projector so the projector weight-grads stay O(1) (else runaway). self.cam_in_norm = nnx.LayerNorm(cam_token_dim, epsilon=1e-5, rngs=rngs) self.cam_token_proj = ProjLN(cam_token_dim, H, rngs=rngs) # feat_input_norm: LayerNorm the raw backbone features before the layer projectors. DA3-GIANT # features are O(1) so this was unneeded; VGGT aggregator taps have outlier channels (absmax ~160) # that blow up the projector weight-grads (grad_norm 62 vs DA3's 0.77 -> NaN by step ~50). self.feat_input_norm = bool(feat_input_norm) if self.feat_input_norm: self.feat_in_norm = nnx.LayerNorm(da3_channels, epsilon=1e-5, rngs=rngs) # POINT MAP: unproject (ray + metric depth) -> per-patch camera-frame 3D coordinate, encode into # the PAYLOAD. Exploits exact GT metric depth: instead of a scalar log-depth the bank gets the # actual metric surface position, a strongly per-sample-discriminative geometry signal. self.use_point_map = bool(use_point_map) if self.use_point_map: self.point_mlp = Mlp2(3, 256, H, rngs=rngs) # DEPTH-AWARE CROSS-VIEW: give each perceiver token its world-frame 3D position (pooled from the # world point map to the locality-anchor grid), so cross-view self-attention can match tokens by # actual 3D correspondence, not just camera-pose identity. Requires cross_view. self.depth_aware_crossview = bool(depth_aware_crossview) if self.depth_aware_crossview: self.pos3d_mlp = Mlp2(3, 256, H, rngs=rngs) # (f) v2: learned per-token embedding added to each view's FINAL bank tokens. Guarantees # persistent cross-token diversity — the quantity that drives softmax gradients to the # injection's Q/K (shared content cancels in the softmax jacobian, so without this the # attention pattern barely trains; measured ~1000x slower than V/out in v1). self.bank_token_embeds = ( {name: nnx.Param(jax.random.normal(rngs.params(), (1, self._ntok[name], H)) * 0.05) for name, _, _ in _VIEWS} if bank_token_embed else None ) # ---- spatial conditioning (WHERE bound to WHAT) -------------------------------- # Normalisation constants are MEASURED (18,109 patch-points over 4 tasks x 6 frames x # 3 views): per-axis centre, ISOTROPIC scale. Isotropic on purpose -- per-axis scaling # would make 10 cm along x encode differently from 10 cm along z and destroy the metric # the network is meant to learn distances from. # NOTE these belong in the assets beside norm_stats.json and be shipped with the # checkpoint; defaults here are a stopgap. A train/eval mismatch silently shifts all # geometry, exactly like the base_qvel mismatch did. self.spatial_vec = bool(spatial_vec) self.fourier_bands = int(fourier_bands) self.point_centre = tuple(float(v) for v in point_centre) self.point_centre_ee_l = tuple(float(v) for v in point_centre_ee_l) self.point_centre_ee_r = tuple(float(v) for v in point_centre_ee_r) self.point_scale = float(point_scale) self.point_max_depth = float(point_max_depth) if self.spatial_vec: _pd = 3 * (1 + 2 * self.fourier_bands) # one Fourier-encoded 3-vector _in = 3 * _pd + 3 + 1 # p, p-eeL, p-eeR, dir_world, valid self.spatial_cond = SpatialConditioner( H, _in, film=spatial_film, use_da3=spatial_use_da3, rngs=rngs) # aux geometry decoder (shared across views): grid-pos query -> attend perceiver tokens -> log-depth if self.aux_geom_head: self.aux_q = nnx.Linear(H, H, rngs=rngs) self.aux_k = nnx.Linear(H, H, rngs=rngs) self.aux_v = nnx.Linear(H, H, rngs=rngs) self.aux_out = nnx.Linear(H, 1, rngs=rngs) def _fuse_layers(self, feats_v): # feats_v: [b, num_layers, C, h, w] -> [b, 432, H] b, L, C, h, w = feats_v.shape parts = [] for li in range(self.num_layers): flat = einops.rearrange(feats_v[:, li], "b c h w -> b (h w) c") # row-major if self.feat_input_norm: flat = self.feat_in_norm(flat) # tame VGGT outlier channels before projection p = self.layer_projectors[li](flat) + self.layer_embed.value[li][None, None, :] parts.append(p) return self.layer_fuse(jnp.concatenate(parts, axis=-1)) def _ray7(self, ray_v, depth_v, ext_v): # ray_v [b,3,h,w], depth_v [b,1,h,w], ext_v [b,4,4] -> [b,432,7] ray6 = compute_world_ray_6d(ray_v, ext_v) # [b,6,h,w] logd = jnp.log(jnp.clip(depth_v.astype(jnp.float32), a_min=1e-3)).astype(ray6.dtype) # [b,1,h,w] ray7 = jnp.concatenate([ray6, logd], axis=1) # [b,7,h,w] return einops.rearrange(ray7, "b c h w -> b (h w) c") def _point_maps(self, ray_v, depth_v, ext_v): """(unit cam ray, metric Z-depth, extrinsics) -> per-patch metric 3D points. Z-depth d and unit cam dir r: range along ray = d / r_z; p_cam = range*r; p_world = origin+range*dir_world. Returns (p_cam_flat [b,432,3], p_world_flat [b,432,3]).""" ray6 = compute_world_ray_6d(ray_v, ext_v) # [b,6,h,w] = [origin(cam center world), dir_world] origin, dir_world = ray6[:, 0:3], ray6[:, 3:6] r = ray_v.astype(jnp.float32) # [b,3,h,w] unit cam dir rng = depth_v.astype(jnp.float32) / jnp.clip(r[:, 2:3], a_min=0.1) # [b,1,h,w] range along ray p_cam = rng * r # [b,3,h,w] p_world = origin + rng * dir_world # [b,3,h,w] f = lambda t: einops.rearrange(t, "b c h w -> b (h w) c") return f(p_cam), f(p_world) def _norm_pt(self, p, centre): c = jnp.asarray(centre, p.dtype) return jnp.clip((p - c) / jnp.asarray(self.point_scale, p.dtype), -1.0, 1.0) def _spatial_vector(self, p_world, dir_world, depth_flat, ee_pos): """[b,P,3] robot-frame points (+ view dir, depth, EE positions) -> [b,P,in_dim]. Three Fourier-encoded position vectors: absolute, and relative to each hand. The hand-relative pair is the per-sample quantity that survives a static scene -- the gripper moves every frame even when nothing else does. dir_world is kept raw (already a bounded unit vector) because it says which side of the surface we are looking at, which `p` alone cannot express. Invalid depth zeroes the geometry but keeps the flag, so the network can tell "no measurement" from "at the origin". """ valid = ((depth_flat > 0.05) & (depth_flat < self.point_max_depth)).astype(p_world.dtype) K = self.fourier_bands geom = jnp.concatenate([ fourier_encode(self._norm_pt(p_world, self.point_centre), K), fourier_encode(self._norm_pt(p_world - ee_pos[:, 0:1, :], self.point_centre_ee_l), K), fourier_encode(self._norm_pt(p_world - ee_pos[:, 1:2, :], self.point_centre_ee_r), K), dir_world, ], axis=-1) * valid return jnp.concatenate([geom, valid], axis=-1) def _pool_pts(self, pw_flat, k): """world points [b, h*w, 3] -> [b, k, 3] pooled to the perceiver's locality-anchor grid (row-major).""" h, w = self.grid_hw pmap = einops.rearrange(pw_flat, "b (h w) c -> b c h w", h=h, w=w) ar = int(_math.ceil(_math.sqrt(k))); ac = int(_math.ceil(k / ar)) pooled = jax.image.resize(pmap.astype(jnp.float32), (pmap.shape[0], 3, ar, ac), method="linear") return einops.rearrange(pooled, "b c ar ac -> b (ar ac) c")[:, :k, :] # [b,k,3] def __call__(self, feats, ray, depth, extrinsics, lang_feat, lang_mask, return_aux: bool = False, depth_drop_rng=None, depth_conf=None, pose_enc=None, cam_tokens=None): # feats [b,L,V,C,h,w]; ray [b,V,3,h,w]; depth [b,V,1,h,w]; extrinsics [b,V,4,4] # lang_feat [b,Lt,lang_dim]; lang_mask [b,Lt] True=real token # return_aux: also return the aux geometry (log-depth reconstruction) loss (training only). # depth_drop_rng: training-only rng enabling depth_dropout (kv_split path); None = keep depth. h, w = self.grid_hw pos_emb = self.pos2d_mlp(_grid_coords(h, w).astype(feats.dtype)) # [1,432,H] if self.pos_emb_scale != 1.0: pos_emb = pos_emb * jnp.asarray(self.pos_emb_scale, pos_emb.dtype) # Language is a per-TASK constant: in a single-task fine-tune lang_feat is byte-identical # for every sample, so the whole fusion stack can only add a constant to the bank -- while # costing 63.5% of the builder's parameters. Skipped entirely when depth is 0. lang_tokens = self.t5_projector(lang_feat) if self.t5_projector is not None else None lang_pad = jnp.logical_not(lang_mask) # True=pad # ---- END-EFFECTOR POSITIONS (free, from the extrinsics) ---- # The left/right realsense cameras are WRIST-mounted, so each one's camera centre in the robot # frame is that arm's end-effector position (up to a fixed wrist->camera offset the network can # absorb). Verified on real data: the wrist cams travel metres per episode (std 0.15-0.21 m) # while the head cam is static (std 0.016-0.033 m, y identically 0). ee_pos = None if self.ee_anchor or self.infonce or self.spatial_vec: ee_pos = jnp.stack([_cam_center(extrinsics[:, vi]) for vi in (1, 2)], axis=1) # [b,2,3] geos = {} world_pts = {} aux_losses = [] for name, vidx, _k in _VIEWS: fused = self._fuse_layers(feats[:, :, vidx]) # [b,432,H] ray_flat = self._ray7(ray[:, vidx], depth[:, vidx], extrinsics[:, vidx]) # [b,432,7] p_cam_flat = None need_pts = (self.use_point_map or self.depth_aware_crossview or self.ee_anchor or self.infonce or self.spatial_vec) if need_pts: p_cam_flat, p_world_flat = self._point_maps(ray[:, vidx], depth[:, vidx], extrinsics[:, vidx]) if self.depth_aware_crossview or self.infonce: world_pts[name] = p_world_flat # [b,432,3] robot-frame metric points ee_d2 = None if self.ee_anchor and ee_pos is not None: # squared distance from every patch's 3D point to each hand, in the ROBOT frame. # Clamped: invalid/far depth would otherwise give a huge negative bias (-> -inf logits). d = p_world_flat[:, None, :, :] - ee_pos[:, :, None, :] # [b,2,432,3] ee_d2 = jnp.clip(jnp.sum(d * d, axis=-1), 0.0, self.ee_max_dist2) # [b,2,432] _bte = self.bank_token_embeds[name].value if self.bank_token_embeds is not None else None view_emb = self.view_embed(jnp.asarray(vidx))[None, None, :] # [1,1,H] if self.spatial_vec: # Resolve depth AGAINST the ray into a position, then let that position # condition the DA3 latent. Replaces the kv_split address/payload split, whose # whole purpose (keeping input-independent terms out of the values) is moot # once the value stream carries a per-sample position. _r6 = compute_world_ray_6d(ray[:, vidx], extrinsics[:, vidx]) _dirw = einops.rearrange(_r6[:, 3:6], "b c h w -> b (h w) c") _dfl = einops.rearrange(depth[:, vidx], "b c h w -> b (h w) c") # [b,P,1] _sv = self._spatial_vector(p_world_flat.astype(jnp.float32), _dirw.astype(jnp.float32), _dfl.astype(jnp.float32), ee_pos.astype(jnp.float32)) geo = self.spatial_cond(_sv.astype(feats.dtype), fused) if _bte is not None: geo = geo + _bte geos[name] = geo continue tok_emb = _bte if self._bte_query else None # into query (new) vs post-fusion (old) if self.kv_split: # K/V split: payload (values) = DA3 latents + depth enc; address (keys) = pos/ray/view. ray_emb = self.ray_mlp(ray_flat[..., :6].astype(feats.dtype)) # Plucker only [b,432,H] depth_emb = self.depth_mlp(ray_flat[..., 6:7].astype(feats.dtype)) # [b,432,H] if depth_drop_rng is not None and self.depth_dropout > 0.0: # per-sample: this fraction of the batch sees NO explicit depth channel, so the # DA3 features must carry the geometry for those samples (redundancy pressure). keep = jax.random.bernoulli( jax.random.fold_in(depth_drop_rng, vidx), 1.0 - self.depth_dropout, (depth_emb.shape[0], 1, 1), ) depth_emb = depth_emb * keep.astype(depth_emb.dtype) payload = fused + depth_emb # [b,432,H] if self.use_point_map and p_cam_flat is not None: payload = payload + self.point_mlp(p_cam_flat.astype(feats.dtype)) # metric 3D point if self.use_depth_conf and depth_conf is not None: conf_flat = einops.rearrange(depth_conf[:, vidx], "b c h w -> b (h w) c") # [b,432,1] conf_flat = jnp.log(jnp.clip(conf_flat.astype(feats.dtype), 1e-3)) # bound VGGT's exp-scaled conf payload = payload + self.conf_mlp(conf_flat) addr = view_emb + pos_emb + ray_emb # routing-only annotations if self.use_perceiver: geo = self.perceivers[name](payload, addr=addr, token_embed=tok_emb, ee_dist2=ee_d2) else: # No perceiver: the patch tokens ARE the bank. The address is dropped rather # than folded in -- it is exactly the input-independent term kv_split existed # to keep out of the values, and position now rides in the payload instead. geo = payload if tok_emb is None else payload + tok_emb # [b,P,H] else: if self.depth_target_only: # depth is a TARGET, never an input: zero ch 6 (log-depth) so the aux prediction # can only come from the DA3 features. Keeps ray_mlp's 7-ch shape (ckpt-compat). ray_in = ray_flat.at[..., 6].set(0.0) else: ray_in = ray_flat ray_emb = self.ray_mlp(ray_in.astype(feats.dtype)) # [b,432,H] spatial = fused + view_emb + pos_emb + ray_emb # [b,P,H] if self.use_perceiver: geo = self.perceivers[name](spatial, token_embed=tok_emb, ee_dist2=ee_d2) else: geo = spatial if tok_emb is None else spatial + tok_emb # [b,P,H] if return_aux and self.aux_geom_head: # grid-pos queries (constant) attend to this view's K perceiver tokens -> per-patch # log-depth. Prediction varies ONLY through geo, so a good fit REQUIRES geo to encode # per-sample geometry. MSE against the true DA3 log-depth (ray_flat channel 6). qh = jnp.broadcast_to(self.aux_q(pos_emb), (geo.shape[0], h * w, self.hidden_dim)) # [b,P,H] kh = self.aux_k(geo) # [b,K,H] vh = self.aux_v(geo) # [b,K,H] scale = jnp.sqrt(jnp.asarray(self.hidden_dim, qh.dtype)) attn = jax.nn.softmax(jnp.einsum("bph,bkh->bpk", qh, kh) / scale, axis=-1) # [b,P,K] pred_logd = self.aux_out(jnp.einsum("bpk,bkh->bph", attn, vh)) # [b,P,1] true_logd = ray_flat[..., 6:7].astype(pred_logd.dtype) # [b,P,1] aux_losses.append(jnp.mean(jnp.square(pred_logd - true_logd))) geos[name] = geo # [b,K,H] # ---- CROSS-VIEW 3D FUSION: views exchange info, grounded by camera pose ---- if self.cross_view: # depth-aware: per-token world-frame 3D positions (pooled to anchors), scene-centered so # the cross-view self-attention can match tokens across views by actual 3D correspondence. tok_xyz = None if self.depth_aware_crossview and world_pts: xyz_list = [self._pool_pts(world_pts[nm], self._ntok[nm]) for nm, _, _ in _VIEWS] center = jnp.mean(jnp.concatenate(xyz_list, axis=1), axis=1, keepdims=True) # [b,1,3] scene centroid tok_xyz = [xz - center for xz in xyz_list] parts = [] for i, (name, vidx, _k) in enumerate(_VIEWS): cam = self.cam_pose_mlp(_cam_pose_feat(extrinsics[:, vidx]).astype(feats.dtype)) # [b,H] if self.use_pose_enc and pose_enc is not None: cam = cam + self.pose_enc_mlp(pose_enc[:, vidx].astype(feats.dtype)) # learned VGGT pose tok = geos[name] + cam[:, None, :] if tok_xyz is not None: tok = tok + self.pos3d_mlp(tok_xyz[i].astype(feats.dtype)) # per-token world-3D position parts.append(tok) x = self.cross_view_fusion(jnp.concatenate(parts, axis=1)) # [b, sum_k, H] off = 0 for name, _vidx, _k in _VIEWS: k = self._ntok[name] geos[name] = x[:, off:off + k] off += k # ---- language fusion + bank-centering, per view ---- banks = {} for name, _vidx, _k in _VIEWS: bank = (self.lang_fusers[name](geos[name], lang_tokens, lang_pad) if self.lang_fusers is not None else geos[name]) # [b,K,H] # bank_token_embeds: new placement shapes the perceiver query (above); OLD placement adds it # here post-fusion (faithful eval of pre-move checkpoints; dead under bank_center as before). if self.bank_token_embeds is not None and not self._bte_query: bank = bank + self.bank_token_embeds[name].value if self.bank_center: # Project out the batch-mean (over the sharded batch axis => global mean under jit). # A purely-constant bank now injects zero; only per-sample deviation reaches the base, # so the model must use per-sample geometry or nothing. See bank_center in the config. bank = bank - jnp.mean(bank, axis=0, keepdims=True) if self.use_cam_tokens and cam_tokens is not None: # VGGT camera+register global tokens -> projected and appended (after all fusion, so the # perceiver locality grid and cross-view split are untouched). Centered for consistency. ct = self.cam_token_proj(self.cam_in_norm(cam_tokens[:, _vidx].astype(feats.dtype))) # [b,17,H] if self.bank_center: ct = ct - jnp.mean(ct, axis=0, keepdims=True) bank = jnp.concatenate([bank, ct], axis=1) # [b, K+17, H] banks[name] = bank # ---- INFONCE: make the bank IDENTIFY its own sample's geometry ---- # Shuffle-damage sat flat at +3-4% because nothing ever trained specificity -- it was only # measured. Here the pooled bank of sample i must beat every other sample in the batch at # matching sample i's pooled metric point map. Constant-ish banks score chance and are punished. nce = jnp.asarray(0.0, jnp.float32) if return_aux and self.infonce and world_pts: zb = jnp.concatenate([jnp.mean(banks[nm], axis=1) for nm, _, _ in _VIEWS], axis=-1) # [b,3H] zg = jnp.concatenate( [self._pool_pts(world_pts[nm], self.infonce_pool_k).reshape(zb.shape[0], -1) for nm, _, _ in _VIEWS], axis=-1) # [b, 3*3*pool_k] nce = _info_nce(self.nce_bank_proj(zb.astype(jnp.float32)), self.nce_geom_proj(zg.astype(jnp.float32)), self.infonce_temp) if return_aux: aux = jnp.mean(jnp.stack(aux_losses)) if aux_losses else jnp.asarray(0.0, jnp.float32) return banks, aux, nce return banks