legacy: copy vggt_newbank_boxing_gloves_step19999 -> legacy/vggt_newbank_boxing_gloves_step19999

.gitattributes

@@ -2497,3 +2497,16 @@ legacy/task12_da3_large_gtdepth_v2_pointmap_step29999_ext30k/params/ocdbt.proces
 legacy/task12_da3_large_gtdepth_v2_pointmap_step29999_ext30k/params/ocdbt.process_0/d/d98a3fad37dee54708c107db1c8fb9ec filter=lfs diff=lfs merge=lfs -text
 legacy/task12_da3_large_gtdepth_v2_pointmap_step29999_ext30k/params/ocdbt.process_0/d/e253ce1315c04daa9ed15b1288857241 filter=lfs diff=lfs merge=lfs -text
 legacy/task12_da3_large_gtdepth_v2_pointmap_step29999_ext30k/params/ocdbt.process_0/d/ffd84587e396735aab5be6c968b420b4 filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/assets/IliaLarchenko/behavior_224_rgb/norm_stats.json filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/d/a46900611b6ae8d52cc9c73344a60cd5 filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/79393094e7346483b9ad835f750e221d filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/91a085ae6e1894bb048d726fc0eb5095 filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/b4e08dc278ff9559ca81bab2466ad7eb filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/b6d0beee53fde61e4a4cad79fe82d2bb filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/c89bf5c18befe2a4541c9878cd67c634 filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/d35b7454f6d0270967e4a343cf9a0532 filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/df61096e2ebc9bed5ca155e01070f51d filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/eca7e002c9a0d1b4044e00ed6a6b7c72 filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/f054a053cf6d0352891dfa0126e38abb filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/f4c80a2f2a49d019fd150fe66b20b6f6 filter=lfs diff=lfs merge=lfs -text
+legacy/vggt_newbank_boxing_gloves_step19999/params/ocdbt.process_0/d/fd79ec0eed15dfaa789ab7d2b3cd12ea filter=lfs diff=lfs merge=lfs -text

legacy/vggt_newbank_boxing_gloves_step19999/README.md

@@ -0,0 +1,100 @@
+# VGGT-Omega newbank policy — how to run / serve / eval
+
+This bundle contains **everything needed to load, serve, and evaluate** the checkpoint
+`JackLiu0406/b1k-checkpoints/vggt_newbank_boxing_gloves_step19999`. Read this + the `code/` files and you
+will know exactly how to run the model.
+
+The policy is the 2025 BEHAVIOR-1K winner (PiBehavior = pi0.5 JAX fork: `gemma_2b` VLM + `gemma_300m`
+action expert) with the **DA3-GIANT geometry backbone REPLACED by VGGT-Omega-1B**, feeding an *enriched*
+spatial bank (adds depth-confidence, camera pose-encoding, and camera/register global tokens). Base init =
+the 50-task meta checkpoint (`behavior-1k/2025-challenge-demos`). Task = `clean_boxing_gloves`, **step
+19999/20000 (FINAL checkpoint, action_loss ≈ 0.041)**.
+
+---
+
+## 0. TL;DR to run
+
+1. Get the code into a `behavior-1k-solution` checkout (file map in §3) and put `third_party/vggt_omega`
+   on `sys.path`.
+2. Get the VGGT weights: `JackLiu0406/vggt-omega-1b` → `vggt_omega_1b_512.pt` (cached under `$HF_HOME`).
+3. Export the flags in §2 so `build_config()` attaches the VGGT-enriched `da3` config to the model.
+4. Load `params/` from this checkpoint with **`remove_extra_params=False`** (see §5 — critical).
+5. At every timestep, run the **VGGT extractor** on the 3 RGB views (§4) to produce the 6 spatial fields,
+   put them on the `Observation`, and call `sample_actions`. Camera extrinsics/intrinsics come from the
+   sim (GT); rays are computed analytically from the intrinsics.
+
+---
+
+## 1. Checkpoint contents
+- `params/` — inference weights (orbax, bf16). This is what you serve.
+- `assets/` — norm stats (`IliaLarchenko/behavior_224_rgb/norm_stats.json`), tokenizer.
+- (No `train_state/` — that's only for resuming training.)
+
+## 2. Model-build flags (MUST match training or the checkpoint won't load / will mis-serve)
+`scripts/train_2026.py::build_config()` reads these env vars and does
+`model = dataclasses.replace(model, da3=B1KDA3Config(...))`. Base config has `da3=None`; **without these
+the VGGT spatial branch is absent and `params/` won't match.**
+```
+USE_DA3_FULL=1  USE_VGGT=1  VGGT_PROCESS_RES=256
+DA3_CHANNELS=2048  DA3_GRID_H=16  DA3_GRID_W=16
+DA3_USE_DEPTH_CONF=1  DA3_USE_POSE_ENC=1  DA3_USE_CAM_TOKENS=1  DA3_CAM_TOKEN_DIM=2048
+DA3_FEAT_INPUT_NORM=1
+DA3_KV_SPLIT=1  DA3_BANK_CENTER=1  DA3_PERC_LOCALITY=1  DA3_CROSS_VIEW=1  DA3_CROSS_VIEW_DEPTH=2
+DA3_DEPTH_DROPOUT=0.5  DA3_BTE_QUERY=1
+DA3_QK_NORM=1  DA3_PERC_NORM_FINAL=1  DA3_PERC_NORM_OUT=0  DA3_POS_EMB_SCALE=0.25
+DA3_LOGIT_GAIN_INIT=3.0  DA3_INJ_GAIN_MAX=8.0  DA3_PERC_LOGIT_GAIN=1  DA3_PERC_GAIN_INIT=3.0  DA3_PERC_GAIN_MAX=8.0
+DA3_SCALE=1.0  DA3_INIT_STD=0.01  DA3_LR_GROUPS=1
+```
+(`depth_dropout` is inference-inert — it only fires when a dropout rng is passed, which serving does not.)
+
+## 3. File map — where each `code/` file goes in `behavior-1k-solution`
+| bundle file | repo path | what it is |
+|---|---|---|
+| `src/b1k/training/vggt_extractor.py` | same | **NEW** — VGGTInlineExtractor (the backbone). Loads VGGT-Omega, outputs the 6 fields |
+| `src/b1k/models/spatial_da3.py` | same | the enriched `SpatialBankBuilder` (depth_conf / pose_enc / cam_tokens / feat_input_norm) |
+| `src/b1k/models/pi_behavior_config.py` | same | `B1KDA3Config` — all the fields the flags in §2 set |
+| `src/b1k/models/observation.py` | same | `Observation` — the spatial fields incl. `da3_depth_conf/pose_enc/cam_tokens` |
+| `src/b1k/models/pi_behavior.py` | same | `_compute_banks()` (passes the 6 fields to the builder) + the layer-12–17 injection |
+| `src/b1k/training/b1k_da3.py` | same | loader + `batch_transform` (extractor selection via `USE_VGGT`, 6-field plumbing) |
+| `src/b1k/training/b1k_2026.py` | same | `BehaviorV3Dataset` + `da3_fields()` (RGB decode, GT extrinsics, intrinsics from FOCAL_RATIO) |
+| `scripts/train_2026.py` | same | `build_config()` — env → config |
+| `third_party/vggt_omega/` | `<repo>/third_party/vggt_omega` | the VGGT-Omega package (put its parent on sys.path) |
+| `reference/serve_b1k.py`, `reference/eval_b1k_wrapper.py`, `reference/policy_config.py` | — | the serve entrypoints for reference; see §5 |
+
+## 4. The VGGT extractor (`vggt_extractor.py`) — what it produces
+`VGGTInlineExtractor(process_res=256).extract(images[B,V,H,W,3], extrinsics[B,V,4,4], intrinsics[B,V,3,3])`
+returns a 6-tuple (field order matters — `b1k_da3.py` maps them to observation keys):
+1. `da3_features` `[B,4,V,2048,16,16]` bf16-as-uint16 — 4 VGGT aggregator taps (blocks 4/11/17/23), patch tokens
+2. `da3_ray` `[B,V,3,16,16]` — **analytic** camera-frame unit ray dirs from the intrinsics (VGGT has no ray head)
+3. `da3_depth` `[B,V,1,16,16]` — VGGT dense_head depth, pooled to grid (VGGT-PREDICTED, not GT)
+4. `da3_depth_conf` `[B,V,1,16,16]` — VGGT depth confidence
+5. `da3_pose_enc` `[B,V,9]` — VGGT camera_head pose enc (trans3+quat4+fov2)
+6. `da3_cam_tokens` `[B,V,17,2048]` — camera(1)+register(16) global tokens
+Views V=3 in order **(main=zed/head, left wrist, right wrist)**. Feeds `[0,1]` RGB; the aggregator does its
+own ImageNet/ResNet renorm. patch_size 16, so `process_res` must be a multiple of 16.
+NOTE: a warmup pass runs per replica in `__init__` (single-threaded) — needed because `torch.linalg` lazy
+wrappers race across the per-device extraction threads; and rays are analytic (no `torch.linalg.inv`).
+
+## 5. Serving / eval integration (CRITICAL correctness note)
+- Build the policy with `da3 = B1KDA3Config(...)` per §2, and load with
+  `train_config.model.load(restore_params(dir/'params'), remove_extra_params=False)`. **`remove_extra_params`
+  must be False** so a param mismatch RAISES instead of silently dropping the whole `da3.*` subtree — a
+  stock serve (which builds `da3=None`) would otherwise load a lobotomized base model with no spatial branch
+  and give meaningless rollouts. (`reference/policy_config.py` is where the current DA3 serve loads params.)
+- Per timestep: decode the 3 RGB views at 256, run the VGGT extractor (§4), and set the six
+  `observation.da3_*` fields + `observation.camera_extrinsics` (GT from sim). `da3_features` ship as uint16
+  bf16-bits; `_compute_banks` (pi_behavior.py) bitcasts them. Then `model.sample_actions(...)`.
+- The bank the action expert attends to is **145/113/113 tokens** for main/left/right (K + 17 cam tokens).
+
+## 6. Data flow (how the enriched bank is built — see `spatial_da3.py::SpatialBankBuilder.__call__`)
+- **payload (attention values)** = fused VGGT feats (4-tap, LayerNorm'd via `feat_input_norm`) + depth_emb
+  (log VGGT depth) + conf_emb (log VGGT depth_conf)
+- **address (attention keys)** = pos_emb(×0.25) + view_emb + ray_emb (world Plücker from analytic ray + GT extrinsics)
+- K/V split → perceiver (locality-biased, K queries) → cross-view fusion (adds cam-pose feat + VGGT pose_enc)
+  → language fusion (ModernBERT task tokens) → bank-center → append projected VGGT cam/register tokens
+- injected into action-expert layers 12–17 via cross-attention.
+
+## 7. Sanity checks before trusting an eval
+- A forward should give a finite `action_loss` (~0.041 for this final ckpt on `clean_boxing_gloves`).
+- **Specificity**: zeroing the bank should HURT loss (measured +41%±32 at step 2000 → the spatial branch is
+  load-bearing). If zeroing the bank does nothing, the spatial branch was dropped — recheck §5.

legacy/vggt_newbank_boxing_gloves_step19999/assets/IliaLarchenko/behavior_224_rgb/fast_tokenizer/metadata.json

@@ -0,0 +1,25 @@
+{
+  "vocab_size": 1024,
+  "scale": 10.0,
+  "encoded_dims": "0:6,7:23",
+  "encoded_dim_ranges": [
+    [
+      0,
+      6
+    ],
+    [
+      7,
+      23
+    ]
+  ],
+  "total_encoded_dims": 22,
+  "action_horizon": 30,
+  "num_training_chunks": 5935465,
+  "compression_stats": {
+    "compression_ratio": 3.644254501482549,
+    "mean_token_length": 181.107,
+    "p99_token_length": 658.0,
+    "min_token_length": 35.0,
+    "max_token_length": 660.0
+  }
+}

legacy/vggt_newbank_boxing_gloves_step19999/assets/IliaLarchenko/behavior_224_rgb/fast_tokenizer/processing_action_tokenizer.py

@@ -0,0 +1,158 @@
+import logging
+from typing import ClassVar
+
+import numpy as np
+from scipy.fft import dct
+from scipy.fft import idct
+from tokenizers import ByteLevelBPETokenizer
+from tokenizers.trainers import BpeTrainer
+from transformers import PreTrainedTokenizerFast
+from transformers.processing_utils import ProcessorMixin
+
+
+class UniversalActionProcessor(ProcessorMixin):
+    attributes: ClassVar[list[str]] = ["bpe_tokenizer"]
+    bpe_tokenizer_class: str = "AutoTokenizer"
+
+    def __init__(
+        self,
+        bpe_tokenizer: PreTrainedTokenizerFast,
+        scale: float = 10,
+        vocab_size: int = 1024,
+        min_token: int = 0,
+        *,
+        action_dim: int | None = None,
+        time_horizon: int | None = None,
+    ):
+        self.scale = scale
+        self.vocab_size = vocab_size
+        self.min_token = min_token
+
+        # Action horizon and dimension needed during decoding. These can be specified
+        # in three ways (in order of priority):
+        # 1. passed in as kwargs to decode()
+        # 2. in the constructor
+        # 3. cached from the last time decode() was called
+        self.time_horizon = time_horizon
+        self.action_dim = action_dim
+        self.called_time_horizon = time_horizon
+        self.called_action_dim = action_dim
+
+        super().__init__(bpe_tokenizer)
+
+    def __call__(self, action_chunk: np.array) -> np.array:
+        assert action_chunk.ndim <= 3, "Only 3 dimensions supported: [batch, timesteps, action_dim]"
+        if action_chunk.ndim == 2:
+            action_chunk = action_chunk[None, ...]
+
+        # Cache the time horizon and action dimension for decoding
+        self.called_time_horizon = action_chunk.shape[-2]
+        self.called_action_dim = action_chunk.shape[-1]
+
+        dct_coeff = dct(action_chunk, axis=1, norm="ortho")
+        dct_coeff = np.around(dct_coeff * self.scale)
+        tokens = []
+        for elem in dct_coeff:
+            token_str = "".join(map(chr, np.maximum(elem.flatten() - self.min_token, 0).astype(int)))
+            tokens.append(self.bpe_tokenizer(token_str)["input_ids"])
+        return tokens
+
+    def decode(
+        self,
+        tokens: list[list[int]],
+        *,
+        time_horizon: int | None = None,
+        action_dim: int | None = None,
+    ) -> np.array:
+        self.time_horizon = time_horizon or self.time_horizon or self.called_time_horizon
+        self.action_dim = action_dim or self.action_dim or self.called_action_dim
+
+        # Cache the time horizon and action dimension for the next call
+        self.called_time_horizon = self.time_horizon
+        self.called_action_dim = self.action_dim
+
+        assert (
+            self.time_horizon is not None and self.action_dim is not None
+        ), "Tokenizer not initialized, call encode() once or pass in time_horizon and action_dim."
+
+        decoded_actions = []
+        for token in tokens:
+            try:
+                decoded_tokens = self.bpe_tokenizer.decode(token)
+                decoded_dct_coeff = np.array(list(map(ord, decoded_tokens))) + self.min_token
+                decoded_dct_coeff = decoded_dct_coeff.reshape(-1, self.action_dim)
+                assert (
+                    decoded_dct_coeff.shape
+                    == (
+                        self.time_horizon,
+                        self.action_dim,
+                    )
+                ), f"Decoded DCT coefficients have shape {decoded_dct_coeff.shape}, expected ({self.time_horizon}, {self.action_dim})"
+            except Exception as e:
+                print(f"Error decoding tokens: {e}")
+                print(f"Tokens: {token}")
+                decoded_dct_coeff = np.zeros((self.time_horizon, self.action_dim))
+            decoded_actions.append(idct(decoded_dct_coeff / self.scale, axis=0, norm="ortho"))
+        return np.stack(decoded_actions)
+
+    @classmethod
+    def fit(
+        cls,
+        action_data: list[np.array],
+        scale: float = 10,
+        vocab_size: int = 1024,
+        *,
+        time_horizon: int | None = None,
+        action_dim: int | None = None,
+    ) -> "UniversalActionProcessor":
+        # Run DCT over all inputs
+        dct_tokens = [dct(a, axis=0, norm="ortho").flatten() for a in action_data]
+
+        # Quantize and find min token
+        max_token = int(np.around(np.concatenate(dct_tokens) * scale).max())
+        min_token = int(np.around(np.concatenate(dct_tokens) * scale).min())
+        min_vocab_size = max_token - min_token
+
+        assert (
+            min_vocab_size <= vocab_size
+        ), f"Vocab size {vocab_size} is too small for the range of tokens {min_vocab_size}"
+        if min_vocab_size + 100 > vocab_size:
+            logging.warning(
+                f"Initial alphabet size {min_vocab_size} is almost as large as the vocab"
+                f"size {vocab_size}, consider increasing vocab size"
+            )
+
+        # Make token iterator for BPE training
+        def _token_iter():
+            for tokens in dct_tokens:
+                rounded_tokens = np.around(tokens * scale) - min_token
+                rounded_tokens = rounded_tokens.astype(int)
+                string = "".join(map(chr, rounded_tokens))
+                yield string
+
+        # Train BPE tokenizer
+        bpe = ByteLevelBPETokenizer()
+
+        # Set up the entire range of possible tokens as the initial alphabet
+        alphabet = [chr(i) for i in range(max_token - min_token + 1)]
+        trainer = BpeTrainer(
+            vocab_size=vocab_size,
+            min_frequency=2,
+            show_progress=True,
+            special_tokens=[],
+            initial_alphabet=alphabet,
+            max_token_length=10000,
+        )
+
+        # Train the inner tokenizer (don't use ByteLevelBPETokenizer.train_from_iterator()
+        # because it doesn't support custom alphabets)
+        bpe._tokenizer.train_from_iterator(_token_iter(), trainer=trainer)
+
+        return cls(
+            PreTrainedTokenizerFast(tokenizer_object=bpe, clean_up_tokenization_spaces=False),
+            scale=scale,
+            vocab_size=vocab_size,
+            min_token=min_token,
+            time_horizon=time_horizon,
+            action_dim=action_dim,
+        )

legacy/vggt_newbank_boxing_gloves_step19999/assets/IliaLarchenko/behavior_224_rgb/fast_tokenizer/processor_config.json

@@ -0,0 +1,11 @@
+{
+  "action_dim": 22,
+  "auto_map": {
+    "AutoProcessor": "processing_action_tokenizer.UniversalActionProcessor"
+  },
+  "min_token": -55,
+  "processor_class": "UniversalActionProcessor",
+  "scale": 10.0,
+  "time_horizon": 30,
+  "vocab_size": 1024
+}

legacy/vggt_newbank_boxing_gloves_step19999/assets/IliaLarchenko/behavior_224_rgb/fast_tokenizer/special_tokens_map.json

@@ -0,0 +1 @@
+{}

legacy/vggt_newbank_boxing_gloves_step19999/assets/IliaLarchenko/behavior_224_rgb/fast_tokenizer/tokenizer_config.json

@@ -0,0 +1,11 @@
+{
+  "added_tokens_decoder": {},
+  "auto_map": {
+    "AutoProcessor": "processing_action_tokenizer.UniversalActionProcessor"
+  },
+  "clean_up_tokenization_spaces": false,
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "processor_class": "UniversalActionProcessor",
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}

legacy/vggt_newbank_boxing_gloves_step19999/assets/IliaLarchenko/behavior_224_rgb/norm_stats.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccd14a0210fc59b2d2726ba599cc0c4b81347395dd60d2a15b334b28ed15a80b
+size 18009212

legacy/vggt_newbank_boxing_gloves_step19999/code/reference/eval_b1k_wrapper.py

@@ -0,0 +1,307 @@
+"""B1K policy wrapper with action compression, rolling inpainting, and stage voting."""
+
+import logging
+import numpy as np
+import torch
+import dataclasses
+from collections import deque
+
+from openpi_client.base_policy import BasePolicy
+from openpi_client.image_tools import resize_with_pad
+from b1k.policies.b1k_policy import extract_state_from_proprio
+from b1k.models.pi_behavior_config import TASK_NUM_STAGES
+from b1k.shared.correction_rules import apply_correction_rules, check_gripper_variation
+from omnigibson.learning.utils.eval_utils import PROPRIOCEPTION_INDICES
+
+logger = logging.getLogger(__name__)
+
+RESIZE_SIZE = 224
+
+
+@dataclasses.dataclass
+class B1KWrapperConfig:
+    """Configuration for B1K policy wrapper execution parameters."""
+    actions_to_execute: int = 26
+    actions_to_keep: int = 4
+    execute_in_n_steps: int = 20
+    history_len: int = 3
+    votes_to_promote: int = 2
+    time_threshold_inpaint: float = 0.3
+    num_steps: int = 20
+    apply_eval_tricks: bool = True
+
+
+class B1KPolicyWrapper():
+    """B1K policy wrapper for PI_BEHAVIOR models with action compression, rolling inpainting, and stage voting."""
+    
+    def __init__(
+        self, 
+        policy: BasePolicy,
+        text_prompt: str = "PI_BEHAVIOR model (task-conditioned)",  # Not used, kept for compatibility
+        action_horizon: int = 30,
+        task_id: int | None = None,
+        config: B1KWrapperConfig = None,
+        checkpoint_switcher = None,
+    ) -> None:
+        self.base_policy = policy
+        self.policy = policy
+        self.checkpoint_switcher = checkpoint_switcher
+        self.text_prompt = text_prompt
+        self.action_horizon = action_horizon
+        self.config = config if config is not None else B1KWrapperConfig()
+        
+        # Validate configuration
+        if self.config.actions_to_execute + self.config.actions_to_keep > self.action_horizon:
+            raise ValueError(
+                f"actions_to_execute + actions_to_keep exceeds action_horizon"
+            )
+        
+        # PI_BEHAVIOR specific (always True for B1K)
+        self.task_id = task_id
+        self.current_stage = 0
+        self.prediction_history = deque([], maxlen=self.config.history_len)
+        
+        # Control loop variables
+        self.last_actions = None
+        self.action_index = 0
+        self.step_count = 0
+        self.prediction_count = 0
+        self.next_initial_actions = None
+    
+    def reset(self):
+        """Reset policy state."""
+        self.policy.reset()
+        self.last_actions = None
+        self.action_index = 0
+        self.step_count = 0
+        self.prediction_count = 0
+        self.next_initial_actions = None
+        self.current_stage = 0
+        self.prediction_history.clear()
+        logger.info(f"Policy reset - Task ID: {self.task_id}, Action horizon: {self.action_horizon}")
+    
+    def _handle_task_change(self, new_task_id):
+        """Handle task ID change by switching checkpoint and resetting state."""
+        if self.task_id != new_task_id:
+            old_task_id = self.task_id
+            self.task_id = new_task_id
+            
+            logger.info(f"🔄 Task change detected: {old_task_id} → {new_task_id} (max stages: {TASK_NUM_STAGES[new_task_id]})")
+            
+            if self.checkpoint_switcher:
+                new_policy = self.checkpoint_switcher.get_policy_for_task(new_task_id)
+                if new_policy is not self.policy:
+                    logger.info(f"📦 Switching checkpoint: task {old_task_id} → {new_task_id}")
+                    self.base_policy = new_policy
+                    self.policy = new_policy
+                    self.policy.reset()
+            
+            self.current_stage = 0
+            self.prediction_history.clear()
+            self.last_actions = None
+            self.action_index = 0
+            self.next_initial_actions = None
+
+    def process_obs(self, obs: dict) -> dict:
+        """Process observation to match model input format."""
+        prop_state = obs["robot_r1::proprio"]
+        
+        head_original = obs["robot_r1::robot_r1:zed_link:Camera:0::rgb"][..., :3]
+        left_original = obs["robot_r1::robot_r1:left_realsense_link:Camera:0::rgb"][..., :3]
+        right_original = obs["robot_r1::robot_r1:right_realsense_link:Camera:0::rgb"][..., :3]
+        
+        # Resize images
+        head_resized = resize_with_pad(head_original, RESIZE_SIZE, RESIZE_SIZE)
+        left_resized = resize_with_pad(left_original, RESIZE_SIZE, RESIZE_SIZE)
+        right_resized = resize_with_pad(right_original, RESIZE_SIZE, RESIZE_SIZE)
+        
+        return {
+            "observation/egocentric_camera": head_resized,
+            "observation/wrist_image_left": left_resized,
+            "observation/wrist_image_right": right_resized,
+            "observation/state": prop_state,
+            "prompt": self.text_prompt,
+        }
+    
+    def update_current_stage(self, predicted_subtask_logits):
+        """Update current stage using majority voting."""
+        if self.task_id is None:
+            return
+            
+        max_stage = TASK_NUM_STAGES[self.task_id] - 1
+        predicted_stage = int(np.argmax(predicted_subtask_logits))
+        
+        if predicted_stage > max_stage:
+            predicted_stage = max_stage
+        
+        self.prediction_history.append(predicted_stage)
+        
+        if len(self.prediction_history) == self.config.history_len:
+            next_stage = self.current_stage + 1
+            
+            if next_stage <= max_stage:
+                votes_for_next = sum(1 for pred in self.prediction_history if pred == next_stage)
+                votes_to_skip = sum(1 for pred in self.prediction_history if pred == next_stage + 1)
+                votes_to_go_back = sum(1 for pred in self.prediction_history if pred == self.current_stage - 1)
+                
+                if votes_for_next >= self.config.votes_to_promote:
+                    old_stage = self.current_stage
+                    self.current_stage = next_stage
+                    self.prediction_history.clear()
+                    logger.info(f"⬆️  Stage advanced: {old_stage} → {self.current_stage} (task {self.task_id}, step {self.step_count})")
+                elif votes_to_skip == self.config.history_len:
+                    old_stage = self.current_stage
+                    self.current_stage = next_stage
+                    self.prediction_history.clear()
+                    logger.info(f"⏭️  Stage skipped: {old_stage} → {self.current_stage} (task {self.task_id}, step {self.step_count})")
+                elif votes_to_go_back == self.config.history_len and self.current_stage > 0:
+                    old_stage = self.current_stage
+                    self.current_stage -= 1
+                    self.prediction_history.clear()
+                    logger.info(f"⬅️  Stage went back: {old_stage} → {self.current_stage} (task {self.task_id}, step {self.step_count})")
+    
+    def prepare_batch_for_pi_behavior(self, batch):
+        """Prepare batch for PI_BEHAVIOR model by adding task_id and current_stage."""
+        task_id = self.task_id if self.task_id is not None else -1
+        batch_copy = batch.copy()
+        if "prompt" in batch_copy:
+            del batch_copy["prompt"]
+        
+        batch_copy["tokenized_prompt"] = np.array([task_id, self.current_stage], dtype=np.int32)
+        batch_copy["tokenized_prompt_mask"] = np.array([True, True], dtype=bool)
+        batch_copy["subtask_state"] = np.array(self.current_stage, dtype=np.int32)
+        
+        return batch_copy
+    
+    def _interpolate_actions(self, actions, target_steps):
+        """Interpolate actions using cubic spline."""
+        from scipy.interpolate import interp1d
+        
+        original_indices = np.linspace(0, len(actions)-1, len(actions))
+        target_indices = np.linspace(0, len(actions)-1, target_steps)
+        
+        interpolated = np.zeros((target_steps, actions.shape[1]))
+        for dim in range(actions.shape[1]):
+            f = interp1d(original_indices, actions[:, dim], kind='cubic')
+            interpolated[:, dim] = f(target_indices)
+        
+        return interpolated
+
+    def act(self, obs: dict) -> torch.Tensor:
+        """Main action function."""
+        
+        # Extract task_id from observations
+        if "task_id" in obs:
+            new_task_id = int(obs["task_id"][0])
+            self._handle_task_change(new_task_id)
+        
+        raw_state = obs["robot_r1::proprio"]
+        current_state = extract_state_from_proprio(raw_state)
+        
+        # Check if we need new actions
+        if self.last_actions is None or self.action_index >= self.config.execute_in_n_steps:
+            
+            # Process observation
+            model_input = self.process_obs(obs)
+            model_input = self.prepare_batch_for_pi_behavior(model_input)
+            
+            # Add rolling inpainting if available
+            if self.next_initial_actions is not None and ("initial_actions" not in model_input or model_input["initial_actions"] is None):
+                model_input["initial_actions"] = self.next_initial_actions
+            
+            # Get prediction
+            if "initial_actions" in model_input and model_input["initial_actions"] is not None:
+                output = self.policy.infer(model_input, initial_actions=model_input["initial_actions"])
+            else:
+                output = self.policy.infer(model_input)
+            
+            actions = output["actions"]
+            
+            # Ensure correct shape
+            if len(actions.shape) == 3:
+                actions = actions[0]
+            if actions.shape[1] > 23:
+                actions = actions[:, :23]
+            
+            # Apply eval tricks if enabled
+            should_compress = self.config.execute_in_n_steps < self.config.actions_to_execute
+            
+            if self.config.apply_eval_tricks:
+                if self.task_id is not None:
+                    actions_before = actions.copy()
+                    actions, corrected_stage = apply_correction_rules(
+                        self.task_id, self.current_stage, current_state, actions
+                    )
+                    
+                    # Log if stage was corrected
+                    if corrected_stage != self.current_stage:
+                        logger.info(f"🔧 Correction rule: Stage corrected {self.current_stage} → {corrected_stage} (task {self.task_id}, step {self.step_count})")
+                        self.current_stage = corrected_stage
+                        self.prediction_history.clear()
+                    
+                    # Log if actions were modified
+                    if not np.allclose(actions_before, actions, rtol=1e-3):
+                        max_diff = np.max(np.abs(actions_before - actions))
+                        logger.info(f"🔧 Correction rule: Actions modified (max diff: {max_diff:.4f}, task {self.task_id}, stage {self.current_stage})")
+                
+                if should_compress:
+                    has_high_variation, mean_var, max_var = check_gripper_variation(
+                        actions, self.config.actions_to_execute
+                    )
+                    if has_high_variation:
+                        should_compress = False
+                        logger.info(f"🔧 Gripper variation: Compression disabled (mean: {mean_var:.4f}, max: {max_var:.4f})")
+            
+            # Determine execution parameters
+            actions_to_execute = self.config.actions_to_execute if should_compress else self.config.execute_in_n_steps
+            execute_steps = self.config.execute_in_n_steps
+            
+            # Save actions for next inpainting (before compression)
+            inpainting_start = actions_to_execute
+            inpainting_end = inpainting_start + self.config.actions_to_keep
+            
+            if len(actions) >= inpainting_end:
+                self.next_initial_actions = actions[inpainting_start:inpainting_end].copy()
+            else:
+                self.next_initial_actions = None
+            
+            # Extract and compress actions
+            self.last_actions = actions[:actions_to_execute].copy()
+            
+            if should_compress:
+                compressed_actions = self._interpolate_actions(self.last_actions, execute_steps)
+                compression_factor = actions_to_execute / execute_steps
+                compressed_actions[:, :3] *= compression_factor  # Scale velocities
+                self.last_actions = compressed_actions
+            
+            self.action_index = 0
+            self.prediction_count += 1
+            
+            # Log prediction details (at lower frequency, every 10 predictions)
+            if self.prediction_count % 10 == 0:
+                compression_status = f"compressed {actions_to_execute}→{execute_steps}" if should_compress else f"uncompressed ({execute_steps})"
+                logger.info(f"🎯 Prediction #{self.prediction_count} | Actions: {compression_status} | Inpainting: {self.next_initial_actions is not None}")
+            
+            # Update stage based on model predictions
+            if "subtask_logits" in output:
+                self.update_current_stage(output["subtask_logits"])
+        
+        # Get current action from sequence
+        if self.action_index >= len(self.last_actions):
+            self.action_index = 0
+            
+        current_action = self.last_actions[self.action_index]
+        self.action_index += 1
+        self.step_count += 1
+        
+        # Log progress every 100 steps
+        if self.step_count % 100 == 0:
+            logger.info(f"📊 Step {self.step_count} | Task: {self.task_id} | Stage: {self.current_stage}/{TASK_NUM_STAGES[self.task_id]-1} | Predictions: {self.prediction_count}")
+        
+        # Convert to torch tensor
+        action_tensor = torch.from_numpy(current_action).float()
+        if len(action_tensor) > 23:
+            action_tensor = action_tensor[:23]
+        
+        return action_tensor
+

legacy/vggt_newbank_boxing_gloves_step19999/code/reference/policy_config.py

@@ -0,0 +1,119 @@
+"""Policy configuration for B1K - loads checkpoints and creates policies.
+
+Exact copy of openpi.policies.policy_config but imports b1k.models.pi_behavior.PiBehavior.
+"""
+
+import logging
+import os
+import pathlib
+from typing import Any
+
+import numpy as np
+import jax.numpy as jnp
+
+import openpi.models.model as _model
+import openpi.policies.policy as _policy
+import openpi.shared.download as download
+import openpi.transforms as transforms
+
+# Import B1K-specific modules
+from b1k.models.pi_behavior import PiBehavior
+from b1k.policies.pi_behavior_policy import PiBehaviorPolicy
+from b1k.training import checkpoints as _checkpoints
+from b1k.training import config as _config
+from b1k import transforms as b1k_transforms
+from b1k.transforms_normalize import NormalizeWithPerTimestamp, UnnormalizeWithPerTimestamp
+
+
+def create_trained_policy(
+    train_config: _config.TrainConfig,
+    checkpoint_dir: pathlib.Path | str,
+    *,
+    repack_transforms: transforms.Group | None = None,
+    sample_kwargs: dict[str, Any] | None = None,
+    default_prompt: str | None = None,
+    norm_stats: dict[str, transforms.NormStats] | None = None,
+    pytorch_device: str | None = None,
+) -> _policy.Policy:
+    """Create a policy from a trained checkpoint - EXACT COPY from openpi with b1k imports."""
+    repack_transforms = repack_transforms or transforms.Group()
+    checkpoint_dir = download.maybe_download(str(checkpoint_dir))
+
+    # Detect PyTorch model
+    is_pytorch = (checkpoint_dir / "pytorch_model.safetensors").exists() or (checkpoint_dir / "pytorch_model.pt").exists()
+    
+    if is_pytorch:
+        raise NotImplementedError("PyTorch inference not supported in b1k")
+    
+    # JAX model loading - load directly as bfloat16 to save memory (12GB vs 24GB)
+    model = train_config.model.load(_model.restore_params(checkpoint_dir / "params", dtype=jnp.bfloat16))
+    
+    # Get data config
+    data_config = train_config.data.create(train_config.assets_dirs, train_config.model)
+    
+    # Load norm stats if not provided
+    if norm_stats is None:
+        if data_config.asset_id is None:
+            raise ValueError("Asset id is required to load norm stats.")
+        norm_stats = _checkpoints.load_norm_stats(checkpoint_dir / "assets", data_config.asset_id)
+    
+    # Load correlation matrix for PiBehavior models
+    if isinstance(model, PiBehavior):
+        if norm_stats is None:
+            raise ValueError("PiBehavior requires norm_stats but none found.")
+        model.load_correlation_matrix(norm_stats)
+        logging.info("Loaded correlation matrix for inference")
+    
+    # Determine the device for PyTorch (not used for b1k but kept for compatibility)
+    if is_pytorch and pytorch_device is None:
+        try:
+            import torch
+            pytorch_device = "cuda" if torch.cuda.is_available() else "cpu"
+        except ImportError:
+            pytorch_device = "cpu"
+    
+    # For PI_BEHAVIOR models during inference, skip training-specific transforms
+    model_transforms_inputs = []
+    for transform in data_config.model_transforms.inputs:
+        # Skip training-specific transforms during inference
+        if isinstance(transform, (b1k_transforms.ComputeSubtaskStateFromMeta, b1k_transforms.TaskIndexToTaskId, b1k_transforms.TokenizeFASTActions)):
+            continue
+        model_transforms_inputs.append(transform)
+    
+    # Build input transform pipeline (skip data_config.repack_transforms - has 'actions' mapping for training)
+    input_transforms = [
+        *repack_transforms.inputs,
+        transforms.InjectDefaultPrompt(default_prompt),
+        *data_config.data_transforms.inputs,
+        NormalizeWithPerTimestamp(norm_stats, use_quantiles=data_config.use_quantile_norm, use_per_timestamp=data_config.use_per_timestamp_norm),
+        *model_transforms_inputs,
+    ]
+    
+    # Build output transform pipeline
+    output_transforms = [
+        *data_config.model_transforms.outputs,
+        UnnormalizeWithPerTimestamp(norm_stats, use_quantiles=data_config.use_quantile_norm, use_per_timestamp=data_config.use_per_timestamp_norm),
+        *data_config.data_transforms.outputs,
+        *repack_transforms.outputs,
+    ]
+    
+    # Use custom PiBehaviorPolicy for PiBehavior models (handles tuple unpacking)
+    if isinstance(model, PiBehavior):
+        return PiBehaviorPolicy(
+            model,
+            transforms=input_transforms,
+            output_transforms=output_transforms,
+            sample_kwargs=sample_kwargs,
+            metadata=train_config.policy_metadata,
+        )
+    else:
+        return _policy.Policy(
+            model,
+            transforms=input_transforms,
+            output_transforms=output_transforms,
+            sample_kwargs=sample_kwargs,
+            metadata=train_config.policy_metadata,
+            is_pytorch=is_pytorch,
+            pytorch_device=pytorch_device if is_pytorch else "cpu",
+        )
+

legacy/vggt_newbank_boxing_gloves_step19999/code/reference/serve_b1k.py

@@ -0,0 +1,190 @@
+import dataclasses
+import enum
+import logging
+import os
+import pathlib
+import socket
+
+import numpy as np
+import tyro
+
+# Set JAX memory allocation before importing JAX (can be overridden by env vars)
+os.environ.setdefault('XLA_PYTHON_CLIENT_MEM_FRACTION', '0.5')  # Use 50% of GPU memory
+os.environ.setdefault('XLA_PYTHON_CLIENT_ALLOCATOR', 'platform')  # Platform allocator
+
+from omnigibson.learning.utils.network_utils import WebsocketPolicyServer
+from omnigibson.learning.datas import BehaviorLerobotDatasetMetadata
+
+from openpi.policies import policy as _policy
+
+# Import B1K-specific modules
+from b1k.policies import policy_config as _policy_config  # Use our custom policy_config
+from b1k.policies.checkpoint_switcher import CheckpointSwitcher
+from b1k.shared.eval_b1k_wrapper import B1KPolicyWrapper, B1KWrapperConfig
+from b1k.training import config as _config
+
+
+class EnvMode(enum.Enum):
+    # Not used, just kept for compatibility
+    ALOHA = "aloha"
+    ALOHA_SIM = "aloha_sim"
+    DROID = "droid"
+    LIBERO = "libero"
+
+
+@dataclasses.dataclass
+class Checkpoint:
+    """Load a policy from a trained checkpoint."""
+    config: str
+    dir: str
+
+
+@dataclasses.dataclass
+class Default:
+    """Use the default policy for the given environment."""
+
+
+@dataclasses.dataclass
+class Args:
+    """Arguments for the serve_policy script."""
+
+    # Environment to serve the policy for. This is only used when serving default policies.
+    env: EnvMode = EnvMode.ALOHA_SIM
+
+    # If provided, will be used in case the "prompt" key is not present in the data, or if the model doesn't have a default prompt.
+    default_prompt: str | None = None
+    
+    # For PI_BEHAVIOR models: task ID (0-49) instead of text prompt
+    task_id: int | None = None
+
+    # Dataset root, used to retrieve the prompt of the task if taskname is not None.
+    dataset_root: str | None = "/scr/behavior/2025-challenge-demos"
+    # If provided, will be used to retrieve the prompt of the task, otherwise use turning_on_radio as default.
+    task_name: str | None = None
+
+    # Port to serve the policy on.
+    port: int = 8000
+    # Record the policy's behavior for debugging.
+    record: bool = False
+
+    # Specifies how to load the policy. If not provided, the default policy for the environment will be used.
+    policy: Checkpoint | Default = dataclasses.field(default_factory=Default)
+    
+    # B1K Wrapper execution parameters
+    actions_to_execute: int = 26
+    actions_to_keep: int = 4
+    execute_in_n_steps: int = 20
+    history_len: int = 3
+    votes_to_promote: int = 2
+    time_threshold_inpaint: float = 0.3
+    num_steps: int = 20
+    apply_eval_tricks: bool = True  # Enable correction rules and gripper variation checks
+    
+    # Multi-checkpoint support for PI_BEHAVIOR models (optional)
+    task_checkpoint_mapping: str | None = None  # Path to task-checkpoint mapping JSON file
+
+
+def create_policy(args: Args) -> _policy.Policy:
+    """Create a policy from the given arguments."""
+    sample_kwargs = {"num_steps": args.num_steps}
+    return _policy_config.create_trained_policy(
+        _config.get_config(args.policy.config), 
+        args.policy.dir, 
+        default_prompt=args.default_prompt,
+        sample_kwargs=sample_kwargs
+    )
+
+
+def main(args: Args) -> None:
+    # B1K only supports PI_BEHAVIOR models (task embeddings, no text prompts)
+    config = _config.get_config(args.policy.config)
+    
+    # PI_BEHAVIOR model setup
+    if args.task_id is not None:
+        logging.info(f"Using PI_BEHAVIOR model with task_id: {args.task_id}")
+        task_id = args.task_id
+    else:
+        logging.info(f"Using PI_BEHAVIOR model - task_id will be extracted from observations")
+        task_id = None
+    
+    # Placeholder prompt for PI_BEHAVIOR (not actually used by model)
+    prompt = "PI_BEHAVIOR model (task-conditioned)"
+    logging.info(f"Using prompt: {prompt}")
+
+    # Load initial/default policy
+    policy = create_policy(args)
+    policy_metadata = policy.metadata
+
+    # Create checkpoint switcher if mapping file provided
+    checkpoint_switcher = None
+    if args.task_checkpoint_mapping:
+        logging.info(f"Multi-checkpoint mode enabled: {args.task_checkpoint_mapping}")
+        
+        sample_kwargs = {"num_steps": args.num_steps}
+        
+        try:
+            checkpoint_switcher = CheckpointSwitcher(
+                config_path=args.task_checkpoint_mapping,
+                training_config=config,
+                sample_kwargs=sample_kwargs
+            )
+            logging.info("Checkpoint switcher initialized - will switch checkpoints based on task_id")
+        except Exception as e:
+            logging.error(f"Failed to initialize checkpoint switcher: {e}")
+            raise
+    else:
+        logging.info("Single checkpoint mode - using one checkpoint for all tasks")
+
+    # Record the policy's behavior.
+    if args.record:
+        policy = _policy.PolicyRecorder(policy, "policy_records")
+
+    # Create wrapper configuration
+    wrapper_config = B1KWrapperConfig(
+        actions_to_execute=args.actions_to_execute,
+        actions_to_keep=args.actions_to_keep,
+        execute_in_n_steps=args.execute_in_n_steps,
+        history_len=args.history_len,
+        votes_to_promote=args.votes_to_promote,
+        time_threshold_inpaint=args.time_threshold_inpaint,
+        num_steps=args.num_steps,
+        apply_eval_tricks=args.apply_eval_tricks,
+    )
+    
+    logging.info(f"Wrapper config: execute={wrapper_config.actions_to_execute}, keep={wrapper_config.actions_to_keep}, steps={wrapper_config.execute_in_n_steps}, num_steps={wrapper_config.num_steps}")
+    
+    if wrapper_config.apply_eval_tricks:
+        logging.info("Eval tricks ENABLED - correction rules and gripper variation checks active")
+    else:
+        logging.info("Eval tricks DISABLED (default behavior)")
+
+    # Create B1K wrapper with PI_BEHAVIOR-specific features
+    policy = B1KPolicyWrapper(
+        policy, 
+        text_prompt=prompt,  # Not used by PI_BEHAVIOR, kept for compatibility
+        task_id=task_id,
+        config=wrapper_config,
+        checkpoint_switcher=checkpoint_switcher
+    )
+    
+    if checkpoint_switcher:
+        logging.info("Multi-checkpoint mode: checkpoints will switch based on task_id from observations")
+    else:
+        logging.info("Rolling inpainting enabled: will use initial_actions from input batch when provided")
+
+    hostname = socket.gethostname()
+    local_ip = socket.gethostbyname(hostname)
+    logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip)
+
+    server = WebsocketPolicyServer(
+        policy=policy,
+        host="0.0.0.0",
+        port=args.port,
+        metadata=policy_metadata,
+    )
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, force=True)
+    main(tyro.cli(Args))

legacy/vggt_newbank_boxing_gloves_step19999/code/scripts/train_2026.py

@@ -0,0 +1,180 @@
+"""Train / correctness-gate the 2025 PiBehavior model on the 2026 v3 subset.
+
+Reuses scripts/train.py's main() but swaps the data loader for the v3 reader
+(b1k.training.b1k_2026) and initializes weights from a 2025 checkpoint.
+
+Env knobs:
+  B1K_2026_ROOT   dataset root (default: checkpoint_3's 13-task subset)
+  B1K_INIT_PARAMS 2025 checkpoint params dir to init from (default: checkpoint_3)
+  B1K_ACTIVITIES  comma-separated activity ids to train/evaluate (default: ck3 subset)
+  B1K_BASE_CONFIG training config to clone (default: pi_behavior_b1k_fast)
+  USE_DA3_SPATIAL Enable precomputed DA3 spatial-token adapter (default 0)
+  DA3_SPATIAL_TOKENS Number of precomputed DA3 tokens per sample (default 320)
+  DA3_SPATIAL_DIM Feature width of each DA3 token (default 1024)
+  DA3_SPATIAL_HEADS Cross-attention heads (default 8)
+  DA3_SPATIAL_SCALE Residual scale for the DA3 adapter (default 1.0)
+  BS              global batch size (must be divisible by #devices; default 16)
+  FSDP_DEVICES    number of devices for FSDP sharding (default: repo config)
+  NW              dataloader workers (default 24)
+  STEPS           num_train_steps (default 40 — a gate, not a full run)
+  FLOW            num_flow_samples (default 4; paper uses 15)
+  LR_WARMUP       cosine LR warmup steps (default: repo config)
+  LR_PEAK         cosine peak LR (default: repo config)
+  LR_DECAY_STEPS  cosine decay steps (default: repo config)
+  LR_DECAY        cosine final LR (default: repo config)
+  SAVE_INTERVAL   checkpoint interval (default: disabled during gates)
+  KEEP_PERIOD     checkpoint keep period (default: repo config)
+  RESUME          resume existing checkpoint directory (default 0)
+  OVERWRITE       overwrite checkpoint directory (default 1 unless RESUME=1)
+  LOG_INTERVAL    metric logging interval (default 10)
+  SKIP_IMAGE_LOG  skip first-batch image logging (default 1)
+"""
+import logging
+import os
+import sys
+import json
+import dataclasses
+
+# JAX-friendly + headless defaults
+os.environ.setdefault("WANDB_MODE", "disabled")
+os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.9")
+os.environ.setdefault("SKIP_IMAGE_LOG", "1")
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))  # so `import train` (scripts/train.py) works
+
+from b1k.training import config as _config
+from b1k.training import data_loader as _data_loader
+from b1k.training import weight_loaders
+from b1k.training.b1k_2026 import create_v3_behavior_data_loader
+from b1k.training.b1k_da3 import create_v3_behavior_da3_loader
+
+ROOT = os.environ.get("B1K_2026_ROOT", "/work/jack/behavior1k/data/behavior_2026_ck3")
+INIT_PARAMS = os.environ.get("B1K_INIT_PARAMS", "/work/jack/behavior1k/checkpoints/checkpoint_3/params")
+BASE_CONFIG = os.environ.get("B1K_BASE_CONFIG", "pi_behavior_b1k_fast")
+TASK_DATA_JSON = os.environ.get("B1K_TASK_DATA_JSON", "/work/jack/behavior1k/task_data.json")
+if os.environ.get("B1K_ACTIVITIES"):
+    ACTIVITIES = [x.strip() for x in os.environ["B1K_ACTIVITIES"].split(",") if x.strip()]
+else:
+    ACTIVITIES = json.load(open("/work/jack/behavior1k/subset_plan.json"))["ck3_names"]
+
+
+def _v3_loader(config, *, sharding=None, shuffle=False, num_batches=None, skip_norm_stats=False):
+    if "SHUFFLE" in os.environ:
+        shuffle = bool(int(os.environ["SHUFFLE"]))
+    if getattr(config.model, "da3", None) is not None and config.model.da3.enabled:
+        return create_v3_behavior_da3_loader(
+            config, ROOT, ACTIVITIES, TASK_DATA_JSON,
+            lang_cache=os.environ.get("DA3_LANG_CACHE", "/work/jack/behavior1k/modernbert_b1k_tasks.pkl"),
+            sharding=sharding, shuffle=shuffle, num_workers=config.num_workers, seed=config.seed or 0,
+        )
+    return create_v3_behavior_data_loader(
+        config, ROOT, ACTIVITIES, TASK_DATA_JSON,
+        sharding=sharding, shuffle=shuffle, num_workers=config.num_workers, seed=config.seed or 0,
+    )
+
+
+# Swap the loader everywhere main() reaches it.
+_data_loader.create_behavior_data_loader = _v3_loader
+import train  # scripts/train.py — defines main()
+train._data_loader.create_behavior_data_loader = _v3_loader
+
+
+def build_config() -> _config.TrainConfig:
+    c = _config.get_config(BASE_CONFIG)
+    model = c.model
+    lr_schedule = c.lr_schedule
+    if any(k in os.environ for k in ("LR_WARMUP", "LR_PEAK", "LR_DECAY_STEPS", "LR_DECAY")):
+        lr_schedule = _config._optimizer.CosineDecaySchedule(
+            warmup_steps=int(os.environ.get("LR_WARMUP", str(lr_schedule.warmup_steps))),
+            peak_lr=float(os.environ.get("LR_PEAK", str(lr_schedule.peak_lr))),
+            decay_steps=int(os.environ.get("LR_DECAY_STEPS", str(lr_schedule.decay_steps))),
+            decay_lr=float(os.environ.get("LR_DECAY", str(lr_schedule.decay_lr))),
+        )
+    if bool(int(os.environ.get("USE_DA3_FULL", "0"))):
+        from b1k.models.pi_behavior_config import B1KDA3Config
+        model = dataclasses.replace(
+            model,
+            da3=B1KDA3Config(
+                spatial_scale=float(os.environ.get("DA3_SCALE", "2.0")),
+                spatial_init_std=float(os.environ.get("DA3_INIT_STD", "0.01")),
+                attn_logit_gain=bool(int(os.environ.get("DA3_LOGIT_GAIN", "1"))),
+                # Gains RETUNED 2026-07-22 for the qk_norm regime. With QK-norm the logits are
+                # O(1), so these act as a real temperature: measured eff-tokens-attended of 324 is
+                # gain 1 -> 202 (avg-pool), 2 -> 64, 3 -> 20, 4 -> 9, 8 -> 2.5, 16 -> 1.5 (one-hot).
+                # The old 32/8 defaults were calibrated for UNBOUNDED logits and are one-hot here.
+                attn_logit_gain_init=float(os.environ.get("DA3_LOGIT_GAIN_INIT", "3.0")),
+                attn_logit_gain_max=float(os.environ.get("DA3_INJ_GAIN_MAX", "8.0")),
+                perceiver_logit_gain=bool(int(os.environ.get("DA3_PERC_LOGIT_GAIN", "1"))),
+                perceiver_logit_gain_init=float(os.environ.get("DA3_PERC_GAIN_INIT", "3.0")),
+                perceiver_logit_gain_max=float(os.environ.get("DA3_PERC_GAIN_MAX", "8.0")),
+                perceiver_norm_attn_out=bool(int(os.environ.get("DA3_PERC_NORM_OUT", "1"))),
+                qk_norm=bool(int(os.environ.get("DA3_QK_NORM", "1"))),
+                perceiver_norm_out=bool(int(os.environ.get("DA3_PERC_NORM_FINAL", "1"))),
+                pos_emb_scale=float(os.environ.get("DA3_POS_EMB_SCALE", "0.25")),
+                bank_center=bool(int(os.environ.get("DA3_BANK_CENTER", "0"))),
+                aux_geom_head=bool(int(os.environ.get("DA3_AUX_GEOM_HEAD", "0"))),
+                aux_geom_weight=float(os.environ.get("DA3_AUX_GEOM_WEIGHT", "0.0")),
+                depth_target_only=bool(int(os.environ.get("DA3_DEPTH_TARGET_ONLY", "0"))),
+                kv_split=bool(int(os.environ.get("DA3_KV_SPLIT", "0"))),
+                depth_dropout=float(os.environ.get("DA3_DEPTH_DROPOUT", "0.0")),
+                perc_locality=bool(int(os.environ.get("DA3_PERC_LOCALITY", "0"))),
+                cross_view=bool(int(os.environ.get("DA3_CROSS_VIEW", "0"))),
+                cross_view_depth=int(os.environ.get("DA3_CROSS_VIEW_DEPTH", "2")),
+                bank_token_embed_query=bool(int(os.environ.get("DA3_BTE_QUERY", "1"))),
+                da3_channels=int(os.environ.get("DA3_CHANNELS", "1536")),
+                grid_hw=(int(os.environ.get("DA3_GRID_H", "18")), int(os.environ.get("DA3_GRID_W", "18"))),
+                use_depth_conf=bool(int(os.environ.get("DA3_USE_DEPTH_CONF", "0"))),
+                use_pose_enc=bool(int(os.environ.get("DA3_USE_POSE_ENC", "0"))),
+                use_cam_tokens=bool(int(os.environ.get("DA3_USE_CAM_TOKENS", "0"))),
+                cam_token_dim=int(os.environ.get("DA3_CAM_TOKEN_DIM", "2048")),
+                feat_input_norm=bool(int(os.environ.get("DA3_FEAT_INPUT_NORM", "0"))),
+            ),
+        )
+    if bool(int(os.environ.get("USE_DA3_SPATIAL", "0"))):
+        model = dataclasses.replace(
+            model,
+            use_spatial_action_cross_attention=True,
+            spatial_num_tokens=int(os.environ.get("DA3_SPATIAL_TOKENS", "320")),
+            spatial_token_dim=int(os.environ.get("DA3_SPATIAL_DIM", "1024")),
+            spatial_num_heads=int(os.environ.get("DA3_SPATIAL_HEADS", "8")),
+            spatial_residual_scale=float(os.environ.get("DA3_SPATIAL_SCALE", "1.0")),
+        )
+    # HARD-FREEZE the base: train ONLY the spatial branch. freeze_filter matches everything that is
+    # NOT spatial, so trainable_filter = All(Param, Not(freeze)) resolves to spatial-only. Because the
+    # train step restricts BOTH the grad (nnx.DiffState) and the optimizer state (tx.init) to the
+    # trainable filter, this (a) prunes the entire base backward -- true "no backward weights" -- and
+    # (b) never allocates Adam moments for the ~3.4B base params (~30GB freed -> room for a bigger BS).
+    _extra = {}
+    if bool(int(os.environ.get("DA3_FREEZE_BASE_HARD", "0"))):
+        import flax.nnx as _nnx
+        import openpi.shared.nnx_utils as _nnxu
+        _spatial = _nnxu.PathRegex(r".*(spatial_bank_builder|spatial_inject).*")
+        _extra["freeze_filter"] = _nnx.Not(_spatial)   # freeze all non-spatial params
+        logging.info("DA3 HARD FREEZE: training ONLY spatial params (base grad + base Adam state skipped)")
+
+    return dataclasses.replace(
+        c,
+        exp_name=os.environ.get("EXP", "v3_ck3_gate"),
+        model=model,
+        lr_schedule=lr_schedule,
+        weight_loader=weight_loaders.PiBehaviorWeightLoader(INIT_PARAMS),
+        wandb_enabled=False,
+        **_extra,
+        overwrite=bool(int(os.environ.get("OVERWRITE", "0" if os.environ.get("RESUME", "0") == "1" else "1"))),
+        resume=bool(int(os.environ.get("RESUME", "0"))),
+        batch_size=int(os.environ.get("BS", "16")),
+        fsdp_devices=int(os.environ.get("FSDP_DEVICES", str(c.fsdp_devices))),
+        num_workers=int(os.environ.get("NW", "24")),
+        num_train_steps=int(os.environ.get("STEPS", "40")),
+        num_flow_samples=int(os.environ.get("FLOW", "4")),
+        log_interval=int(os.environ.get("LOG_INTERVAL", "10")),
+        save_interval=int(os.environ.get("SAVE_INTERVAL", "10000000")),   # disabled during gates by default
+        keep_period=int(os.environ.get("KEEP_PERIOD", str(c.keep_period))),
+        seed=0,
+        assets_base_dir="./outputs/assets",
+        checkpoint_base_dir="./outputs/checkpoints",
+    )
+
+
+if __name__ == "__main__":
+    train.main(build_config())

legacy/vggt_newbank_boxing_gloves_step19999/code/src/b1k/models/observation.py

@@ -0,0 +1,176 @@
+"""Observation class and preprocessing with FAST auxiliary fields support.
+
+Based on openpi with FAST fields added for PI_BEHAVIOR model.
+
+Reference: https://github.com/wensi-ai/openpi/blob/behavior/src/openpi/models/model.py
+"""
+
+from collections.abc import Sequence
+from typing import Generic, TypeVar
+import dataclasses
+
+import augmax
+from flax import struct
+import jax
+import jax.numpy as jnp
+import numpy as np
+import torch
+
+from openpi.shared import image_tools
+from openpi.shared import array_typing as at
+
+ArrayT = TypeVar("ArrayT", bound=jax.Array | torch.Tensor | np.ndarray)
+
+IMAGE_KEYS = (
+    "base_0_rgb",
+    "left_wrist_0_rgb",
+    "right_wrist_0_rgb",
+)
+IMAGE_RESOLUTION = (224, 224)
+
+
+@at.typecheck
+@struct.dataclass
+class Observation(Generic[ArrayT]):
+    """Observation with FAST auxiliary fields."""
+    
+    images: dict[str, at.Float[ArrayT, "*b h w c"]]
+    image_masks: dict[str, at.Bool[ArrayT, "*b"]]
+    state: at.Float[ArrayT, "*b s"]
+    tokenized_prompt: at.Int[ArrayT, "*b l"] | None = None
+    tokenized_prompt_mask: at.Bool[ArrayT, "*b l"] | None = None
+    token_ar_mask: at.Int[ArrayT, "*b l"] | None = None
+    token_loss_mask: at.Bool[ArrayT, "*b l"] | None = None
+    
+    fast_tokens: at.Int[ArrayT, "*b t"] | None = None
+    fast_token_mask: at.Bool[ArrayT, "*b t"] | None = None
+    spatial_tokens: at.Float[ArrayT, "*b n d"] | None = None
+    spatial_token_mask: at.Bool[ArrayT, "*b n"] | None = None
+    # DA3 inline spatial inputs (frozen DA3-GIANT features + geometry; consumed by the trainable
+    # bank builder in PiBehavior). da3_features ship as raw bits: uint16=bf16 bits, uint8=fp8 bytes.
+    da3_features: at.Num[ArrayT, "*b dl v dc gh gw"] | None = None
+    da3_ray: at.Float[ArrayT, "*b v three gh gw"] | None = None
+    da3_depth: at.Float[ArrayT, "*b v one gh gw"] | None = None
+    # VGGT-Omega enrichments (None on the DA3 path): per-patch depth confidence, per-view pose
+    # encoding (trans3+quat4+fov2), and camera+register global tokens.
+    da3_depth_conf: at.Float[ArrayT, "*b v one gh gw"] | None = None
+    da3_pose_enc: at.Float[ArrayT, "*b v pe"] | None = None
+    da3_cam_tokens: at.Float[ArrayT, "*b v ct cd"] | None = None
+    camera_extrinsics: at.Float[ArrayT, "*b v four four2"] | None = None
+    lang_feat: at.Float[ArrayT, "*b lt ld"] | None = None
+    lang_mask: at.Bool[ArrayT, "*b lt"] | None = None
+
+    @classmethod
+    def from_dict(cls, data: at.PyTree[ArrayT]) -> "Observation[ArrayT]":
+        """Convert dict to Observation."""
+        if ("tokenized_prompt" in data) != ("tokenized_prompt_mask" in data):
+            raise ValueError("tokenized_prompt and tokenized_prompt_mask must be provided together.")
+        
+        # Convert uint8 images to float32 [-1, 1]
+        for key in data["image"]:
+            if data["image"][key].dtype == np.uint8:
+                data["image"][key] = data["image"][key].astype(np.float32) / 255.0 * 2.0 - 1.0
+            elif hasattr(data["image"][key], "dtype") and data["image"][key].dtype == torch.uint8:
+                data["image"][key] = data["image"][key].to(torch.float32).permute(0, 3, 1, 2) / 255.0 * 2.0 - 1.0
+        
+        return cls(
+            images=data["image"],
+            image_masks=data["image_mask"],
+            state=data["state"],
+            tokenized_prompt=data.get("tokenized_prompt"),
+            tokenized_prompt_mask=data.get("tokenized_prompt_mask"),
+            token_ar_mask=data.get("token_ar_mask"),
+            token_loss_mask=data.get("token_loss_mask"),
+            fast_tokens=data.get("fast_tokens"),
+            fast_token_mask=data.get("fast_token_mask"),
+            spatial_tokens=data.get("spatial_tokens"),
+            spatial_token_mask=data.get("spatial_token_mask"),
+            da3_features=data.get("da3_features"),
+            da3_ray=data.get("da3_ray"),
+            da3_depth=data.get("da3_depth"),
+            da3_depth_conf=data.get("da3_depth_conf"),
+            da3_pose_enc=data.get("da3_pose_enc"),
+            da3_cam_tokens=data.get("da3_cam_tokens"),
+            camera_extrinsics=data.get("camera_extrinsics"),
+            lang_feat=data.get("lang_feat"),
+            lang_mask=data.get("lang_mask"),
+        )
+
+    def to_dict(self) -> at.PyTree[ArrayT]:
+        """Convert Observation to dict."""
+        result = dataclasses.asdict(self)
+        result["image"] = result.pop("images")
+        result["image_mask"] = result.pop("image_masks")
+        return result
+
+
+def preprocess_observation(
+    rng: at.KeyArrayLike | None,
+    observation: Observation,
+    *,
+    train: bool = False,
+    image_keys: Sequence[str] = IMAGE_KEYS,
+    image_resolution: tuple[int, int] = IMAGE_RESOLUTION,
+) -> Observation:
+    """Preprocess observations with image augmentation and FAST fields preservation."""
+    if not set(image_keys).issubset(observation.images):
+        raise ValueError(f"images dict missing keys: expected {image_keys}, got {list(observation.images)}")
+
+    batch_shape = observation.state.shape[:-1]
+
+    out_images = {}
+    for key in image_keys:
+        image = observation.images[key]
+        if image.shape[1:3] != image_resolution:
+            image = image_tools.resize_with_pad(image, *image_resolution)
+
+        if train:
+            # Convert from [-1, 1] to [0, 1] for augmax
+            image = image / 2.0 + 0.5
+
+            transforms = []
+            if "wrist" not in key:
+                height, width = image.shape[1:3]
+                transforms += [
+                    augmax.RandomCrop(int(width * 0.95), int(height * 0.95)),
+                    augmax.Resize(width, height),
+                    augmax.Rotate((-5, 5)),
+                ]
+            transforms += [
+                augmax.ColorJitter(brightness=0.3, contrast=0.4, saturation=0.5),
+            ]
+            sub_rngs = jax.random.split(rng, image.shape[0])
+            image = jax.vmap(augmax.Chain(*transforms))(sub_rngs, image)
+
+            # Back to [-1, 1]
+            image = image * 2.0 - 1.0
+
+        out_images[key] = image
+
+    # Obtain masks
+    out_masks = {}
+    for key in out_images:
+        if key not in observation.image_masks:
+            out_masks[key] = jnp.ones(batch_shape, dtype=jnp.bool)
+        else:
+            out_masks[key] = jnp.asarray(observation.image_masks[key])
+
+    return Observation(
+        images=out_images,
+        image_masks=out_masks,
+        state=observation.state,
+        tokenized_prompt=observation.tokenized_prompt,
+        tokenized_prompt_mask=observation.tokenized_prompt_mask,
+        token_ar_mask=observation.token_ar_mask,
+        token_loss_mask=observation.token_loss_mask,
+        fast_tokens=getattr(observation, 'fast_tokens', None),
+        fast_token_mask=getattr(observation, 'fast_token_mask', None),
+        spatial_tokens=getattr(observation, 'spatial_tokens', None),
+        da3_features=getattr(observation, 'da3_features', None),
+        da3_ray=getattr(observation, 'da3_ray', None),
+        da3_depth=getattr(observation, 'da3_depth', None),
+        camera_extrinsics=getattr(observation, 'camera_extrinsics', None),
+        lang_feat=getattr(observation, 'lang_feat', None),
+        lang_mask=getattr(observation, 'lang_mask', None),
+        spatial_token_mask=getattr(observation, 'spatial_token_mask', None),
+    )

legacy/vggt_newbank_boxing_gloves_step19999/code/src/b1k/models/pi_behavior.py

@@ -0,0 +1,1327 @@
+"""The main model for BEHAVIOR-1K challenge.
+
+Based on Pi0.5 implementation from PhysicalIntelligence/openpi
+"""
+
+import logging
+import pathlib
+
+import einops
+import flax.linen as nn
+import flax.nnx as nnx
+import flax.nnx.bridge as nnx_bridge
+import jax
+import jax.numpy as jnp
+from typing_extensions import override
+
+from openpi.models import model as _model
+from openpi.models import gemma as _gemma
+from b1k.models import spatial_da3 as _spatial_da3
+from openpi.models import siglip as _siglip
+from openpi.models.pi0 import make_attn_mask, posemb_sincos
+from openpi.shared import array_typing as at
+
+# Import from our custom modules
+from b1k.models import pi_behavior_config
+from b1k.models.observation import Observation, preprocess_observation
+from b1k.models.pi_behavior_config import (
+    TASK_NUM_STAGES, 
+    MAX_NUM_STAGES, 
+    TOTAL_TASK_STAGE_EMBEDDINGS, 
+    TASK_STAGE_OFFSETS
+)
+
+logger = logging.getLogger("b1k")
+
+
+class KVCacheTransform(nnx.Module):
+    """Transforms prefix KV cache by mixing across layers.
+    
+    Each destination layer's K and V become learnable linear combinations
+    of all source layers' K and V, plus a bias term. This allows the action
+    expert to attend to learned combinations of VLM layers rather than being
+    forced to attend layer-by-layer.
+    
+    Initialized as identity transform (k_coeffs = I, bias = 0) so the model
+    starts with the same behavior as without transformation.
+    """
+    
+    def __init__(self, num_layers: int, head_dim: int, num_kv_heads: int, rngs: nnx.Rngs):
+        # K transformation: [dest_layer, src_layer]
+        # Initialize as identity so transformation is initially a no-op
+        self.k_coeffs = nnx.Param(jnp.eye(num_layers, dtype=jnp.float32))
+        
+        # K bias: [layer, num_kv_heads, head_dim]
+        # Initialize as zeros
+        self.k_bias = nnx.Param(jnp.zeros((num_layers, num_kv_heads, head_dim), dtype=jnp.float32))
+        
+        # V transformation (independent from K)
+        self.v_coeffs = nnx.Param(jnp.eye(num_layers, dtype=jnp.float32))
+        self.v_bias = nnx.Param(jnp.zeros((num_layers, num_kv_heads, head_dim), dtype=jnp.float32))
+    
+    def __call__(self, kv_cache: tuple[jnp.ndarray, jnp.ndarray]) -> tuple[jnp.ndarray, jnp.ndarray]:
+        """Transform KV cache by mixing across layers.
+        
+        Args:
+            kv_cache: Tuple of (cache_k, cache_v) where each has shape
+                     [num_layers, batch, seq_len, num_kv_heads, head_dim]
+        
+        Returns:
+            Transformed (k_new, v_new) with same shape and dtype as input
+        """
+        cache_k, cache_v = kv_cache
+        # Shape: [layers, batch, seq_len, num_kv_heads, head_dim]
+        
+        # Preserve original dtype (important for bfloat16 training)
+        original_dtype = cache_k.dtype
+        
+        # Transform K: each destination layer is a weighted combination of all source layers
+        # k_new[dest] = sum_src(k_coeffs[dest, src] * cache_k[src]) + k_bias[dest]
+        # Einsum: [dest, src] @ [src, batch, seq, heads, dim] -> [dest, batch, seq, heads, dim]
+        k_new = jnp.einsum('ds,sbtkh->dbtkh', self.k_coeffs.value, cache_k)
+        k_new = k_new + self.k_bias.value[:, None, None, :, :]  # Add bias
+        
+        # Transform V (same operation, independent parameters)
+        v_new = jnp.einsum('ds,sbtkh->dbtkh', self.v_coeffs.value, cache_v)
+        v_new = v_new + self.v_bias.value[:, None, None, :, :]
+        
+        # Cast back to original dtype
+        k_new = k_new.astype(original_dtype)
+        v_new = v_new.astype(original_dtype)
+        
+        return (k_new, v_new)
+
+
+class SpatialActionCrossAttention(nnx.Module):
+    """Residual cross-attention from action tokens to precomputed DA3 spatial tokens."""
+
+    def __init__(
+        self,
+        action_width: int,
+        spatial_width: int,
+        num_heads: int,
+        rngs: nnx.Rngs,
+    ):
+        if action_width % num_heads != 0:
+            raise ValueError(f"action_width={action_width} must be divisible by num_heads={num_heads}")
+
+        self.num_heads = num_heads
+        self.head_dim = action_width // num_heads
+        self.q_proj = nnx.Linear(action_width, action_width, use_bias=False, rngs=rngs)
+        self.k_proj = nnx.Linear(spatial_width, action_width, use_bias=False, rngs=rngs)
+        self.v_proj = nnx.Linear(spatial_width, action_width, use_bias=False, rngs=rngs)
+        self.out_proj = nnx.Linear(
+            action_width,
+            action_width,
+            kernel_init=nn.initializers.zeros,
+            bias_init=nn.initializers.zeros,
+            rngs=rngs,
+        )
+
+    def __call__(
+        self,
+        action_tokens: jnp.ndarray,
+        spatial_tokens: jnp.ndarray,
+        spatial_token_mask: jnp.ndarray | None = None,
+        *,
+        residual_scale: float = 1.0,
+    ) -> jnp.ndarray:
+        original_dtype = action_tokens.dtype
+        q = self.q_proj(action_tokens)
+        k = self.k_proj(spatial_tokens.astype(action_tokens.dtype))
+        v = self.v_proj(spatial_tokens.astype(action_tokens.dtype))
+
+        q = einops.rearrange(q, "b t (h d) -> b h t d", h=self.num_heads)
+        k = einops.rearrange(k, "b s (h d) -> b h s d", h=self.num_heads)
+        v = einops.rearrange(v, "b s (h d) -> b h s d", h=self.num_heads)
+
+        logits = jnp.einsum("bhtd,bhsd->bhts", q, k, preferred_element_type=jnp.float32)
+        logits = logits * (self.head_dim ** -0.5)
+
+        if spatial_token_mask is not None:
+            big_neg = -2.3819763e38
+            logits = jnp.where(spatial_token_mask[:, None, None, :], logits, big_neg)
+
+        probs = jax.nn.softmax(logits, axis=-1).astype(original_dtype)
+        context = jnp.einsum("bhts,bhsd->bhtd", probs, v)
+        context = einops.rearrange(context, "b h t d -> b t (h d)")
+        delta = self.out_proj(context).astype(original_dtype)
+        return action_tokens + residual_scale * delta
+
+
+class PiBehavior(_model.BaseModel):
+    def __init__(self, config: pi_behavior_config.PiBehaviorConfig, rngs: nnx.Rngs):
+        super().__init__(config.action_dim, config.action_horizon, config.max_token_len)
+        
+        # Store config for later use
+        self.config = config
+        
+        paligemma_config = _gemma.get_config(config.paligemma_variant)
+        action_expert_config = _gemma.get_config(config.action_expert_variant)
+        
+        # Initialize Gemma models with AdaRMS (Pi05 style)
+        spatial_inject = getattr(config, "da3", None) is not None and config.da3.enabled
+        self.da3_cfg = getattr(config, "da3", None)
+        llm = nnx_bridge.ToNNX(
+            _gemma.Module(
+                configs=[paligemma_config, action_expert_config],
+                embed_dtype=config.dtype,
+                adarms=True,
+                spatial_inject=spatial_inject,
+                num_spatial_layers=config.da3.num_inject_layers if spatial_inject else 6,
+                spatial_scale=config.da3.spatial_scale if spatial_inject else 2.0,
+                spatial_init_std=config.da3.spatial_init_std if spatial_inject else 0.0,
+                spatial_logit_gain=config.da3.attn_logit_gain if spatial_inject else False,
+                spatial_logit_gain_init=config.da3.attn_logit_gain_init if spatial_inject else 1.0,
+                spatial_logit_gain_max=config.da3.attn_logit_gain_max if spatial_inject else 0.0,
+                spatial_qk_norm=config.da3.qk_norm if spatial_inject else False,
+            )
+        )
+        llm.lazy_init(rngs=rngs, method="init", use_adarms=[False, True])
+        
+        # Initialize vision model
+        img = nnx_bridge.ToNNX(
+            _siglip.Module(
+                num_classes=paligemma_config.width,
+                variant="So400m/14",
+                pool_type="none",
+                scan=True,
+                dtype_mm=config.dtype,
+            )
+        )
+        img.lazy_init(next(iter(config.fake_obs().images.values())), train=False, rngs=rngs)
+        
+        self.PaliGemma = nnx.Dict(llm=llm, img=img)
+
+        # DA3 spatial-language bank builder (trainable; frozen DA3 runs inline in the data pipeline).
+        self.spatial_bank_builder = None
+        if spatial_inject:
+            d = config.da3
+            self.spatial_bank_builder = _spatial_da3.SpatialBankBuilder(
+                hidden_dim=d.hidden_dim,
+                da3_channels=d.da3_channels,
+                num_layers=d.da3_layers,
+                grid_hw=d.grid_hw,
+                lang_dim=d.lang_dim,
+                num_heads=d.num_heads,
+                lang_fusion_depth=d.lang_fusion_depth,
+                perceiver_query_std=d.perceiver_query_std,
+                qk_norm=d.qk_norm,
+                perceiver_norm_out=d.perceiver_norm_out,
+                pos_emb_scale=d.pos_emb_scale,
+                perceiver_logit_gain=d.perceiver_logit_gain,
+                perceiver_logit_gain_init=d.perceiver_logit_gain_init,
+                perceiver_logit_gain_max=d.perceiver_logit_gain_max,
+                perceiver_norm_attn_out=d.perceiver_norm_attn_out,
+                bank_token_embed=d.bank_token_embed,
+                bank_center=d.bank_center,
+                aux_geom_head=d.aux_geom_head,
+                depth_target_only=d.depth_target_only,
+                kv_split=d.kv_split,
+                depth_dropout=d.depth_dropout,
+                perc_locality=d.perc_locality,
+                cross_view=d.cross_view,
+                cross_view_depth=d.cross_view_depth,
+                bank_token_embed_query=d.bank_token_embed_query,
+                use_depth_conf=d.use_depth_conf,
+                use_pose_enc=d.use_pose_enc,
+                use_cam_tokens=d.use_cam_tokens,
+                cam_token_dim=d.cam_token_dim,
+                pose_enc_dim=d.pose_enc_dim,
+                feat_input_norm=d.feat_input_norm,
+                rngs=rngs,
+            )
+        
+        # KV cache transformation for cross-layer attention
+        # Allows action expert to attend to learned combinations of VLM layers
+        if config.use_kv_transform:
+            self.kv_transform = KVCacheTransform(
+                num_layers=paligemma_config.depth,
+                head_dim=paligemma_config.head_dim,
+                num_kv_heads=paligemma_config.num_kv_heads,
+                rngs=rngs
+            )
+        else:
+            self.kv_transform = None
+        
+        # Task embeddings table - trainable embeddings for each task
+        self.task_embeddings = nnx.Embed(
+            num_embeddings=config.num_tasks,
+            features=config.task_embedding_dim,
+            rngs=rngs,
+        )
+        
+        # Stage predictor - predicts stage from VLM output of base task token
+        # Outputs MAX_NUM_STAGES logits, but invalid stages are masked per task
+        self.stage_pred_from_vlm = nnx.Linear(paligemma_config.width, MAX_NUM_STAGES, rngs=rngs)
+        
+        # Task + subtask fusion layers
+        # Combines task embedding + cos/sin encoded subtask state
+        self.subtask_encoding_dim = config.task_embedding_dim // 2  # Half of task embedding dim (1024)
+        
+        # Task-specific stage embeddings (one per stage per task)
+        # Total embeddings = sum of stages across all tasks (596 for 5-15 stages per task)
+        self.task_stage_embeddings = nnx.Embed(
+            num_embeddings=TOTAL_TASK_STAGE_EMBEDDINGS,
+            features=self.subtask_encoding_dim,
+            rngs=rngs,
+        )
+        
+        # Gated fusion layers
+        # Input: task_embedding + sincos + task_stage_emb = task_dim + 2*subtask_dim
+        fusion_input_dim = config.task_embedding_dim + 2 * self.subtask_encoding_dim
+        
+        # Gate networks to learn how to combine different signals
+        self.gate_sincos = nnx.Linear(fusion_input_dim, self.subtask_encoding_dim, rngs=rngs)
+        self.gate_task_stage = nnx.Linear(fusion_input_dim, self.subtask_encoding_dim, rngs=rngs)
+        self.gate_task = nnx.Linear(fusion_input_dim, config.task_embedding_dim, rngs=rngs)
+        
+        # Fusion networks to create multiple conditioned vectors
+        self.fusion_layer1 = nnx.Linear(fusion_input_dim, config.task_embedding_dim * 2, rngs=rngs)
+        self.fusion_layer2 = nnx.Linear(config.task_embedding_dim * 2, config.task_embedding_dim, rngs=rngs)
+        
+        # Additional projection for stage-dominant representation (2 signals now)
+        self.stage_projection = nnx.Linear(2 * self.subtask_encoding_dim, config.task_embedding_dim, rngs=rngs)
+        
+        # Pi05 style layers
+        self.action_in_proj = nnx.Linear(config.action_dim, action_expert_config.width, rngs=rngs)
+        self.time_mlp_in = nnx.Linear(action_expert_config.width, action_expert_config.width, rngs=rngs)
+        self.time_mlp_out = nnx.Linear(action_expert_config.width, action_expert_config.width, rngs=rngs)
+        self.action_out_proj = nnx.Linear(action_expert_config.width, config.action_dim, rngs=rngs)
+        if config.use_spatial_action_cross_attention:
+            self.spatial_action_xattn = SpatialActionCrossAttention(
+                action_width=action_expert_config.width,
+                spatial_width=config.spatial_token_dim,
+                num_heads=config.spatial_num_heads,
+                rngs=rngs,
+            )
+            logger.info(
+                "DA3 spatial action cross-attention enabled: tokens=%s dim=%s heads=%s scale=%s",
+                config.spatial_num_tokens,
+                config.spatial_token_dim,
+                config.spatial_num_heads,
+                config.spatial_residual_scale,
+            )
+        else:
+            self.spatial_action_xattn = None
+
+        # Correlated noise generation
+        # Initialize as NNX Intermediate (excluded from checkpoints, loaded from norm_stats)
+        # Full correlation matrix with beta shrinkage for robustness
+        flat_dim = config.action_horizon * config.action_dim
+        self.action_correlation_cholesky = nnx.Intermediate(
+            jnp.eye(flat_dim),  # Identity matrix as placeholder
+        )
+        self.correlation_loaded = False  # Track if correlation matrix has been loaded
+        self.use_correlated_noise = config.use_correlated_noise
+        self.correlation_beta = config.correlation_beta  # Shrinkage parameter for regularization
+        
+        # Inpainting cache: stores precomputed matrices for simple correlation-based inpainting
+        # Key: num_inpainted_steps (length of inpainted sequence)
+        # Value: dict with {O_indices, U_indices, Sigma_UO_SOOinv}
+        self.inpainting_cache = {}
+        
+        # FAST auxiliary training components
+        if config.use_fast_auxiliary:
+            # FAST embedding layer (vocab_size → paligemma_width)
+            # Use paligemma width (2048) to match other prefix tokens
+            self.fast_token_embedding = nnx.Embed(
+                num_embeddings=config.fast_vocab_size,
+                features=paligemma_config.width,
+                rngs=rngs
+            )
+            
+            # FAST projection head (paligemma_width → vocab_size)
+            self.fast_token_proj = nnx.Linear(
+                paligemma_config.width,
+                config.fast_vocab_size,
+                rngs=rngs
+            )
+            
+            logger.info(f"FAST auxiliary enabled, vocab_size={config.fast_vocab_size}")
+
+        # This attribute gets automatically set by model.train() and model.eval().
+        self.deterministic = True
+
+    def _compute_banks(self, observation, return_aux=False, depth_drop_rng=None):
+        """Build the per-view DA3 spatial banks once per forward (timestep-independent; reused
+        across all flow samples / denoise steps). If return_aux, also return the aux geometry loss.
+        depth_drop_rng enables depth_dropout (training only); None = inference, depth kept."""
+        if self.spatial_bank_builder is None or getattr(observation, "da3_features", None) is None:
+            return (None, None) if return_aux else None
+        feats = observation.da3_features
+        # Features arrive as raw BITS to minimize host<->device transfer; decode on-device to bf16.
+        # uint16 = bf16 bits (inline extractor); uint8 = fp8-e4m3fn bytes (legacy cache).
+        if feats.dtype == jnp.uint8:
+            feats = jax.lax.bitcast_convert_type(feats, jnp.float8_e4m3fn).astype(jnp.bfloat16)
+        elif feats.dtype == jnp.uint16:
+            feats = jax.lax.bitcast_convert_type(feats, jnp.bfloat16)
+        else:
+            feats = feats.astype(jnp.bfloat16)
+        return self.spatial_bank_builder(
+            feats,
+            observation.da3_ray,
+            observation.da3_depth,
+            observation.camera_extrinsics,
+            observation.lang_feat,
+            observation.lang_mask,
+            return_aux=return_aux,
+            depth_drop_rng=depth_drop_rng,
+            depth_conf=getattr(observation, "da3_depth_conf", None),
+            pose_enc=getattr(observation, "da3_pose_enc", None),
+            cam_tokens=getattr(observation, "da3_cam_tokens", None),
+        )
+
+    def apply_spatial_action_conditioning(self, observation: Observation, action_tokens: jnp.ndarray) -> jnp.ndarray:
+        """Inject precomputed DA3 spatial tokens into action-token hidden states."""
+        if self.spatial_action_xattn is None or observation.spatial_tokens is None:
+            return action_tokens
+
+        return self.spatial_action_xattn(
+            action_tokens,
+            observation.spatial_tokens,
+            observation.spatial_token_mask,
+            residual_scale=self.config.spatial_residual_scale,
+        )
+
+    def encode_subtask_state(
+        self, 
+        subtask_state: at.Int[at.Array, " b"],
+        task_ids: at.Int[at.Array, " b"]
+    ) -> at.Float[at.Array, "b {self.subtask_encoding_dim}"]:
+        """Encode subtask state using cos/sin positional encoding, scaled per task.
+        
+        Args:
+            subtask_state: Current stage for each sample [B]
+            task_ids: Task ID for each sample [B]
+            
+        Returns:
+            Positional encodings scaled to [0, 1] range based on task-specific stage count [B, 1024]
+        """
+        # Get number of stages for each task in batch using JAX array indexing
+        # Convert tuple to JAX array inside function to avoid import-time device allocation
+        task_num_stages_array = jnp.array(TASK_NUM_STAGES, dtype=jnp.int32)
+        task_num_stages = task_num_stages_array[task_ids]  # [B] - JAX array indexing
+        
+        # Normalize: stage 0 → 0.0, last stage → 1.0 (per-task scaling)
+        # Add maximum to avoid division by zero for edge cases
+        normalized_state = subtask_state.astype(jnp.float32) / jnp.maximum(task_num_stages.astype(jnp.float32) - 1.0, 1.0)
+        
+        # Use cos/sin encoding similar to timestep encoding
+        return posemb_sincos(
+            normalized_state, 
+            self.subtask_encoding_dim, 
+            min_period=1e-3, 
+            max_period=1.0
+        )
+
+    def load_correlation_matrix(self, norm_stats: dict):
+        """Load full correlation matrix from normalization statistics and apply shrinkage.
+        
+        This should be called after model initialization when norm_stats are available.
+        Applies shrinkage regularization: S_reg = beta * S + (1-beta) * I for robustness.
+        
+        Args:
+            norm_stats: Dictionary containing normalization statistics (from normalize.load()),
+                       with 'actions' key containing NormStats with action_correlation_cholesky field.
+                       
+        Raises:
+            ValueError: If use_correlated_noise=True but correlation matrix is missing.
+            TypeError: If norm_stats structure is incorrect.
+        """
+        if not self.use_correlated_noise:
+            logger.info("Correlated noise disabled in config, skipping correlation matrix loading")
+            return
+        
+        # Validate norm_stats is a dict
+        if not isinstance(norm_stats, dict):
+            raise TypeError(
+                f"norm_stats must be a dict, got {type(norm_stats).__name__}. "
+                "Ensure norm_stats are loaded using openpi.shared.normalize.load()."
+            )
+        
+        # Check 'actions' key exists
+        if 'actions' not in norm_stats:
+            raise ValueError(
+                "use_correlated_noise=True but 'actions' key not found in norm_stats. "
+                f"Found keys: {list(norm_stats.keys())}. "
+                "Run compute_norm_stats.py with --correlation flag to generate correlation matrix."
+            )
+        
+        actions_stats = norm_stats['actions']
+        
+        # Extract correlation matrix (support both dict and attribute access for flexibility)
+        if isinstance(actions_stats, dict):
+            chol_matrix = actions_stats.get('action_correlation_cholesky')
+            access_method = "dict"
+        elif hasattr(actions_stats, 'action_correlation_cholesky'):
+            chol_matrix = actions_stats.action_correlation_cholesky
+            access_method = "attribute"
+        else:
+            raise TypeError(
+                f"norm_stats['actions'] has unexpected type {type(actions_stats).__name__} "
+                f"and cannot access 'action_correlation_cholesky'. "
+                "Ensure norm_stats are loaded using openpi.shared.normalize.load()."
+            )
+        
+        # Strict validation: correlation matrix must exist and be non-None
+        if chol_matrix is None:
+            raise ValueError(
+                "use_correlated_noise=True but 'action_correlation_cholesky' is None in norm_stats['actions']. "
+                "This means the correlation matrix was not computed during norm_stats generation. "
+                "Run compute_norm_stats.py with --correlation flag to generate correlation matrix."
+            )
+        
+        logger.info(f"Successfully accessed correlation matrix via {access_method} access")
+        
+        # Validate correlation matrix shape
+        expected_dim = self.action_horizon * self.action_dim
+        try:
+            L = jnp.array(chol_matrix)
+        except Exception as e:
+            raise ValueError(
+                f"Failed to convert action_correlation_cholesky to array: {e}. "
+                "The correlation matrix may be corrupted or in an invalid format."
+            )
+        
+        if L.ndim != 2 or L.shape[0] != L.shape[1]:
+            raise ValueError(
+                f"action_correlation_cholesky must be a square 2D matrix, got shape {L.shape}. "
+                f"Expected shape: ({expected_dim}, {expected_dim})"
+            )
+        
+        if L.shape[0] != expected_dim:
+            raise ValueError(
+                f"action_correlation_cholesky has wrong dimensions: {L.shape[0]}x{L.shape[0]}. "
+                f"Expected {expected_dim}x{expected_dim} (action_horizon={self.action_horizon} * action_dim={self.action_dim}). "
+                "This indicates the correlation matrix was computed for a different action space configuration."
+            )
+        
+        # Reconstruct covariance matrix from Cholesky
+        Sigma = L @ L.T
+        
+        # Apply shrinkage regularization: Σ_reg = beta * Σ + (1-beta) * I
+        beta = self.correlation_beta
+        logger.info(f"Applying shrinkage regularization with beta={beta:.2f}")
+        
+        Sigma_reg = beta * Sigma + (1 - beta) * jnp.eye(Sigma.shape[0])
+        
+        # Compute Cholesky decomposition of regularized covariance
+        try:
+            L_reg = jnp.linalg.cholesky(Sigma_reg)
+        except Exception as e:
+            raise RuntimeError(
+                f"Cholesky decomposition failed on regularized covariance: {e}. "
+                "This indicates the regularized correlation matrix is not positive definite. "
+                f"Current beta={beta:.2f}. Try decreasing correlation_beta closer to 0.0 for more shrinkage/regularization."
+            )
+        
+        # Update the Intermediate value
+        self.action_correlation_cholesky.value = L_reg
+        self.correlation_loaded = True
+        
+        logger.info(
+            f"✓ Loaded correlation matrix with shape {L_reg.shape} "
+            f"(beta={beta:.2f} shrinkage applied)"
+        )
+        logger.info(
+            f"  Memory usage: {L_reg.nbytes / 1024 / 1024:.2f} MB"
+        )
+    
+    def generate_correlated_noise(
+        self, 
+        rng: at.KeyArrayLike, 
+        batch_size: int,
+    ) -> at.Float[at.Array, "b {self.action_horizon} {self.action_dim}"]:
+        """Generate correlated noise matching action covariance structure.
+        
+        Uses full correlation matrix with optional beta shrinkage for robustness.
+        
+        Args:
+            rng: Random key for noise generation
+            batch_size: Number of noise samples to generate
+            
+        Returns:
+            Correlated noise with shape [batch_size, action_horizon, action_dim]
+            
+        Raises:
+            RuntimeError: If use_correlated_noise=True but correlation matrix not loaded.
+        """
+        if not self.use_correlated_noise:
+            # Independent Gaussian noise when correlated noise is disabled
+            return jax.random.normal(rng, (batch_size, self.action_horizon, self.action_dim))
+        
+        if not self.correlation_loaded:
+            raise RuntimeError(
+                "use_correlated_noise=True but correlation matrix is not loaded. "
+                "Ensure load_correlation_matrix() was called during model initialization. "
+                "Run compute_norm_stats.py with --correlation flag to generate correlation matrix."
+            )
+        
+        # Generate standard correlated noise using Cholesky decomposition
+        flat_dim = self.action_horizon * self.action_dim
+        standard_normal = jax.random.normal(rng, (batch_size, flat_dim))
+        correlated_flat = standard_normal @ self.action_correlation_cholesky.value.T
+        correlated_noise = correlated_flat.reshape(batch_size, self.action_horizon, self.action_dim)
+        return correlated_noise
+
+    def _precompute_correction_matrix(
+        self,
+        O_indices: at.Int[at.Array, " nO"],
+        U_indices: at.Int[at.Array, " nU"],
+    ) -> dict:
+        """Precompute matrix for correlation-aware inpainting correction.
+        
+        Computes Σ_{UO}Σ_{OO}^{-1} which propagates corrections from O to U
+        while preserving correlation structure.
+        
+        Args:
+            O_indices: Flat indices of inpainted dimensions [|O|]
+            U_indices: Flat indices of free dimensions [|U|]
+            
+        Returns:
+            Dictionary with {O_indices, U_indices, correction_matrix}
+            
+        Raises:
+            RuntimeError: If correlation matrix is not loaded
+        """
+        if not self.correlation_loaded:
+            raise RuntimeError(
+                "Cannot precompute correction matrix: correlation matrix not loaded. "
+                "Call load_correlation_matrix() first."
+            )
+        
+        L = self.action_correlation_cholesky.value
+        Sigma = L @ L.T  # Full covariance matrix [hd, hd]
+        
+        # Extract submatrices
+        Sigma_OO = Sigma[jnp.ix_(O_indices, O_indices)]  # [|O|, |O|]
+        Sigma_UO = Sigma[jnp.ix_(U_indices, O_indices)]  # [|U|, |O|]
+        
+        # Compute correction matrix: Σ_{UO} @ Σ_{OO}^{-1}
+        # This propagates corrections from O to U
+        eps_OO = 1e-6 * jnp.maximum(jnp.mean(jnp.diag(Sigma_OO)), 1.0)
+        Sigma_OO_reg = Sigma_OO + eps_OO * jnp.eye(Sigma_OO.shape[0])
+        
+        # Solve Σ_{OO}_reg @ X = Σ_{UO}.T for X, then transpose
+        correction_matrix = jax.scipy.linalg.solve(
+            Sigma_OO_reg, Sigma_UO.T, assume_a='pos'
+        ).T  # [|U|, |O|]
+        
+        return {
+            'O_indices': O_indices,
+            'U_indices': U_indices,
+            'correction_matrix': correction_matrix,  # Σ_{UO}Σ_{OO}^{-1}
+        }
+
+    def fuse_task_and_subtask(
+        self, task_embedding: at.Float[at.Array, "b d"], task_ids: at.Int[at.Array, " b"], subtask_state: at.Int[at.Array, " b"]
+    ) -> at.Float[at.Array, "b n d"]:
+        """Fuse task embedding with subtask state encoding using multiple representations.
+        
+        Returns multiple vectors that are differently conditioned by the subtask state:
+        1. Task-gated representation (task embedding modulated by subtask)
+        2. Balanced fusion (task + subtask combined)
+        3. Stage-dominant representation (subtask features projected to task space)
+        4. Pure stage representation (concatenated learned embeddings)
+        
+        All output representations have dimension 2048 (task_embedding_dim).
+        
+        Args:
+            task_embedding: Base task embedding [b, 2048]
+            task_ids: Task IDs for task-specific stage embeddings [b]
+            subtask_state: Subtask state indices [b]
+            
+        Returns:
+            Multiple fused embeddings [b, 4, 2048]
+        """
+        # Get subtask representations
+        sincos_encoding = self.encode_subtask_state(subtask_state, task_ids)  # [b, 1024]
+        
+        # Task-specific stage embedding with corrected indexing
+        # Use vectorized lookup: offset + stage for each task
+        # Convert tuple to JAX array inside function to avoid import-time device allocation
+        task_stage_offsets_array = jnp.array(TASK_STAGE_OFFSETS, dtype=jnp.int32)
+        task_stage_offsets = task_stage_offsets_array[task_ids]  # [b] - JAX array indexing
+        task_stage_idx = task_stage_offsets + subtask_state  # [b]
+        task_stage_embedding = self.task_stage_embeddings(task_stage_idx)  # [b, 1024]
+        
+        # Concatenate inputs for gating: task (2048) + sincos (1024) + task_stage (1024) = 4096
+        all_inputs = jnp.concatenate([
+            task_embedding,       # [b, 2048]
+            sincos_encoding,      # [b, 1024]
+            task_stage_embedding  # [b, 1024]
+        ], axis=-1)  # [b, 4096]
+        
+        # Learn gates for each component (sigmoid to get 0-1 scaling)
+        gate_sincos = nnx.sigmoid(self.gate_sincos(all_inputs))      # [b, 1024]
+        gate_task_stage = nnx.sigmoid(self.gate_task_stage(all_inputs))  # [b, 1024]
+        gate_task = nnx.sigmoid(self.gate_task(all_inputs))          # [b, 2048]
+        
+        # 1. Task-gated representation: task embedding modulated by subtask info [b, 2048]
+        task_gated = task_embedding * gate_task
+        
+        # 2. Balanced fusion: combine all signals through fusion network [b, 2048]
+        x = self.fusion_layer1(all_inputs)  # [b, 4096]
+        x = nnx.relu(x)
+        balanced_fusion = self.fusion_layer2(x)  # [b, 2048]
+        
+        # 3. Stage-dominant: weighted combination of stage signals, then project [b, 2048]
+        gated_stage_features = jnp.concatenate([
+            sincos_encoding * gate_sincos,        # [b, 1024]
+            task_stage_embedding * gate_task_stage  # [b, 1024]
+        ], axis=-1)  # [b, 2048]
+        stage_dominant = self.stage_projection(gated_stage_features)  # [b, 2048]
+        
+        # 4. Pure stage: concatenate the embeddings (already 2048) [b, 2048]
+        pure_stage = jnp.concatenate([sincos_encoding, task_stage_embedding], axis=-1)
+        
+        # Stack all four representations [b, 4, 2048]
+        fused_embeddings = jnp.stack([task_gated, balanced_fusion, stage_dominant, pure_stage], axis=1)
+        
+        return fused_embeddings
+
+    @at.typecheck
+    def embed_prefix(
+        self, 
+        obs: Observation
+    ) -> tuple[
+        at.Float[at.Array, "b s emb"], 
+        at.Bool[at.Array, "b s"], 
+        at.Bool[at.Array, " s"]
+    ]:
+        """
+        Embed prefix: images + task + state + FAST_tokens (if provided).
+        
+        Args:
+            obs: Observation (may include fast_tokens and fast_token_mask)
+            
+        Returns:
+            tokens, input_mask, ar_mask
+        """
+        input_mask = []
+        ar_mask = []
+        tokens = []
+        
+        # Embed images
+        image_token_list = []
+        # Respect freeze_vision_backbone config: if frozen, always use train=False
+        # If not frozen, use the model's training state (self.deterministic)
+        vision_train_mode = (not self.deterministic) and (not self.config.freeze_vision_backbone)
+        
+        for name in obs.images:
+            image_tokens, _ = self.PaliGemma.img(obs.images[name], train=vision_train_mode)
+            image_token_list.append(image_tokens)  # Store for subtask prediction
+
+            tokens.append(image_tokens)
+            input_mask.append(
+                einops.repeat(
+                    obs.image_masks[name],
+                    "b -> b s",
+                    s=image_tokens.shape[1],
+                )
+            )
+            # Image tokens attend to each other
+            ar_mask += [False] * image_tokens.shape[1]
+
+        # Add task embeddings with subtask state fusion
+        if obs.tokenized_prompt is not None:
+            # obs.tokenized_prompt now contains task_ids (shape: [batch_size, 2])
+            task_ids = obs.tokenized_prompt[:, 0]  # Extract task_id: [batch_size]
+            base_task_embedding = self.task_embeddings(task_ids)  # shape: [batch_size, embed_dim]
+            
+            # ALWAYS use the input subtask state - never use predicted state inside model
+            if obs.tokenized_prompt.shape[1] > 1:  # If we have [task_id, subtask_state]
+                subtask_state = obs.tokenized_prompt[:, 1]  # Use input subtask state
+            else:
+                raise ValueError("subtask_state must be provided in tokenized_prompt for PI_BEHAVIOR model")
+            
+            # Fuse task embedding with subtask state - returns [b, 4, d] with multiple representations
+            fused_task_embeddings = self.fuse_task_and_subtask(base_task_embedding, task_ids, subtask_state)
+            
+            # Create task token sequence: [base_task, task_gated, balanced_fusion, stage_dominant, pure_stage]
+            task_sequence = jnp.concatenate([
+                base_task_embedding[:, None, :],  # [b, 1, d] - base task token
+                fused_task_embeddings              # [b, 4, d] - stage-conditioned tokens
+            ], axis=1)  # [b, 5, d]
+            
+            tokens.append(task_sequence)
+            # All task tokens are valid
+            task_mask = jnp.ones((obs.tokenized_prompt.shape[0], 5), dtype=jnp.bool_)
+            input_mask.append(task_mask)
+            # Hierarchical attention: base task (False) then stage tokens (True, False, False, False)
+            # Base task attends to images bidirectionally
+            # Stage tokens attend to images+task but not vice versa
+            ar_mask += [False] + [True, False, False, False]
+            
+        # Add state as discrete tokens (Pi05 style)
+        # Discretize state into bins
+        discretized_state = jnp.digitize(obs.state, bins=jnp.linspace(-1, 1, 256 + 1)[:-1]) - 1
+        discretized_state = jnp.clip(discretized_state, 0, 255)  # Ensure valid range
+        
+        # Embed each dimension of the discretized state
+        state_tokens = []
+        for i in range(obs.state.shape[-1]):
+            state_dim_tokens = self.PaliGemma.llm(discretized_state[:, i:i+1], method="embed")
+            state_tokens.append(state_dim_tokens)
+        
+        if state_tokens:
+            state_tokens = jnp.concatenate(state_tokens, axis=1)  # shape: [batch_size, state_dim, embed_dim]
+            tokens.append(state_tokens)
+            input_mask.append(jnp.ones((obs.state.shape[0], obs.state.shape[-1]), dtype=jnp.bool_))
+            # State tokens have full bidirectional attention with all prefix tokens
+            # (images, task, stages, and other state tokens)
+            ar_mask += [False] * state_tokens.shape[1]
+        
+        # FAST tokens (from observation if provided)
+        if self.config.use_fast_auxiliary and obs.fast_tokens is not None:
+            fast_tokens = obs.fast_tokens  # [B, T]
+            fast_token_mask = obs.fast_token_mask  # [B, T]
+            
+            # Teacher forcing: shift right [BOS, tok0, tok1, ..., tok_{T-1}]
+            bos_token = jnp.zeros((fast_tokens.shape[0], 1), dtype=jnp.int32)
+            shifted_tokens = jnp.concatenate([bos_token, fast_tokens[:, :-1]], axis=1)
+            
+            # Shift mask too: [True, mask_0, mask_1, ..., mask_{T-1}]
+            bos_mask = jnp.ones((fast_tokens.shape[0], 1), dtype=jnp.bool_)
+            shifted_mask = jnp.concatenate([bos_mask, fast_token_mask[:, :-1]], axis=1)
+            
+            # Embed using FAST embedding layer (NOT Paligemma!)
+            fast_token_emb = self.fast_token_embedding(shifted_tokens)  # [B, T, D]
+            
+            tokens.append(fast_token_emb)
+            input_mask.append(shifted_mask)  # Use the actual token mask
+            # Causal for FAST: ALL tokens are causal (pure autoregressive)
+            ar_mask += [True] * shifted_tokens.shape[1]
+        
+        tokens = jnp.concatenate(tokens, axis=1)
+        input_mask = jnp.concatenate(input_mask, axis=1)
+        ar_mask = jnp.array(ar_mask)
+        return tokens, input_mask, ar_mask
+
+    @at.typecheck
+    def embed_suffix(
+        self, obs: Observation, noisy_actions: _model.Actions, timestep: at.Float[at.Array, " b"]
+    ) -> tuple[
+        at.Float[at.Array, "b s emb"],
+        at.Bool[at.Array, "b s"],
+        at.Bool[at.Array, " s"],
+        at.Float[at.Array, "b emb"],
+    ]:
+        input_mask = []
+        ar_mask = []
+        tokens = []
+        
+        # Pi05 style: no explicit state token in suffix (it's in prefix as discrete tokens)
+        
+        action_tokens = self.action_in_proj(noisy_actions)
+        # Embed timestep using sine-cosine positional encoding
+        time_emb = posemb_sincos(timestep, self.action_in_proj.out_features, min_period=4e-3, max_period=4.0)
+        
+        # Pi05 style: time MLP for adaRMS
+        time_emb = self.time_mlp_in(time_emb)
+        time_emb = nnx.swish(time_emb)
+        time_emb = self.time_mlp_out(time_emb)
+        time_emb = nnx.swish(time_emb)
+        action_expert_tokens = action_tokens
+        adarms_cond = time_emb
+        
+        tokens.append(action_expert_tokens)
+        input_mask.append(jnp.ones(action_expert_tokens.shape[:2], dtype=jnp.bool_))
+        
+        # image/task/state inputs do not attend to action tokens
+        ar_mask += [True] + ([False] * (self.action_horizon - 1))
+        
+        tokens = jnp.concatenate(tokens, axis=1)
+        input_mask = jnp.concatenate(input_mask, axis=1)
+        ar_mask = jnp.array(ar_mask)
+        return tokens, input_mask, ar_mask, adarms_cond
+
+    @override
+    def compute_loss(
+        self, rng: at.KeyArrayLike, observation: Observation, actions: _model.Actions, *, train: bool = False
+    ) -> at.Float[at.Array, "*b ah"]:
+        """Not used - we only use compute_detailed_loss() for training."""
+        raise NotImplementedError("Use compute_detailed_loss() instead")
+
+    @override
+    def compute_detailed_loss(
+        self, rng: at.KeyArrayLike, observation: Observation, actions: _model.Actions, *, train: bool = False, num_flow_samples: int = 1
+    ) -> dict[str, at.Float[at.Array, "*b"]]:
+        """
+        Compute detailed loss with multiple flow matching samples.
+        
+        Simplified approach using KV cache:
+        - Compute prefix KV cache once (with FAST tokens)
+        - Remove FAST tokens from cache (action expert doesn't attend to FAST)
+        - Process N flow samples independently, each reusing the same cached prefix
+        - Each sample has different noise and different time
+        - Average losses across samples
+        """
+        losses = {}
+
+        preprocess_rng, rng = jax.random.split(rng)
+        observation = preprocess_observation(preprocess_rng, observation, train=train)
+
+        batch_size = actions.shape[0]
+        
+        # 1. Embed prefix once (includes FAST tokens if provided in observation)
+        prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
+        
+        # 2. Compute prefix KV cache
+        prefix_attn_mask = make_attn_mask(prefix_mask, prefix_ar_mask)
+        positions_prefix = jnp.cumsum(prefix_mask, axis=1) - 1
+        (prefix_out, _), kv_cache_full = self.PaliGemma.llm(
+            [prefix_tokens, None],
+            mask=prefix_attn_mask,
+            positions=positions_prefix
+        )
+
+        # DA3 banks: timestep-independent, computed ONCE and closure-captured by the vmapped
+        # flow-sample fn (vmap broadcasts them across the N samples).
+        depth_drop_rng, rng = jax.random.split(rng)
+        spatial_banks, geom_aux_loss = self._compute_banks(
+            observation, return_aux=True, depth_drop_rng=depth_drop_rng if train else None
+        )
+
+        # 3. Predict stage from VLM output of base task token
+        # Base task token is the first token after all image tokens
+        # Image tokens all have ar_mask=False, task starts with ar_mask=False (base) then True (stage tokens)
+        # Structure: [images (all False)] [base_task (False)] [stages (True, False, False, False)]
+        # Find first True (first stage token), base task is at that index - 1
+        first_stage_token_idx = jnp.argmax(prefix_ar_mask)  # Returns index of first True
+        base_task_token_idx = first_stage_token_idx - 1
+        base_task_output = prefix_out[:, base_task_token_idx, :]
+        subtask_logits = self.stage_pred_from_vlm(base_task_output)  # [B, MAX_NUM_STAGES]
+        
+        # Mask out invalid stages for each task (vectorized JAX operations)
+        task_ids = observation.tokenized_prompt[:, 0]  # [B]
+        task_num_stages_array = jnp.array(TASK_NUM_STAGES, dtype=jnp.int32)
+        task_num_stages = task_num_stages_array[task_ids]  # [B] - JAX array indexing
+        stage_range = jnp.arange(MAX_NUM_STAGES)  # [15]
+        valid_mask = stage_range[None, :] < task_num_stages[:, None]  # [B, 15]
+        subtask_logits = jnp.where(valid_mask, subtask_logits, -jnp.inf)  # Mask invalid stages
+        
+        # 4. Extract FAST loss from prefix output (before removing from cache)
+        fast_loss_value = 0.0
+        fast_len = 0
+        fast_targets = observation.fast_tokens
+        fast_token_mask = observation.fast_token_mask
+        
+        if self.config.use_fast_auxiliary and fast_targets is not None:
+            fast_len = fast_targets.shape[1]
+            fast_start_idx = prefix_tokens.shape[1] - fast_len
+            fast_outputs = prefix_out[:, fast_start_idx:, :]  # [B, T, D]
+            
+            # Project to FAST vocab
+            fast_logits = self.fast_token_proj(fast_outputs)  # [B, T, vocab_size]
+            
+            # Cross-entropy loss with teacher forcing
+            pred_logits = fast_logits  # [B, T, vocab]
+            target_tokens = fast_targets  # [B, T]
+            loss_mask = fast_token_mask  # [B, T]
+            
+            log_probs = jax.nn.log_softmax(pred_logits, axis=-1)
+            target_log_probs = jnp.take_along_axis(
+                log_probs,
+                target_tokens[:, :, None],
+                axis=-1
+            ).squeeze(-1)  # [B, T]
+            
+            fast_token_loss = -target_log_probs  # [B, T]
+            
+            # Apply mask and normalize by number of valid tokens
+            masked_loss = fast_token_loss * loss_mask  # [B, T]
+            num_valid_tokens = jnp.maximum(jnp.sum(loss_mask, axis=-1), 1)  # [B]
+            losses["fast_loss"] = jnp.sum(masked_loss, axis=-1) / num_valid_tokens  # [B]
+            
+            # Accuracy (only on valid tokens)
+            pred_tokens = jnp.argmax(pred_logits, axis=-1)
+            correct = (pred_tokens == target_tokens) * loss_mask
+            losses["fast_accuracy"] = jnp.sum(correct, axis=-1) / num_valid_tokens
+            
+            fast_loss_value = self.config.fast_loss_weight * jnp.mean(losses["fast_loss"])
+        elif fast_targets is not None:
+            # FAST auxiliary is disabled but data contains FAST tokens
+            raise ValueError(
+                "use_fast_auxiliary=False but observation contains fast_tokens. "
+                "Either enable use_fast_auxiliary in config or ensure data doesn't contain fast_tokens."
+            )
+        
+        # 5. Remove FAST tokens from KV cache (action expert doesn't attend to FAST)
+        # KV cache shape: [layers, batch, seq_len, num_kv_heads, head_dim]
+        if fast_len > 0:
+            cache_k, cache_v = kv_cache_full
+            # Remove last fast_len tokens from sequence dimension
+            cache_k = cache_k[:, :, :-fast_len, :, :]
+            cache_v = cache_v[:, :, :-fast_len, :, :]
+            kv_cache_for_actions = (cache_k, cache_v)
+            prefix_len_for_actions = prefix_tokens.shape[1] - fast_len
+            # Truncate prefix mask and ar_mask for action expert
+            prefix_mask_for_actions = prefix_mask[:, :-fast_len]
+            prefix_ar_mask_for_actions = prefix_ar_mask[:-fast_len]
+        else:
+            kv_cache_for_actions = kv_cache_full
+            prefix_len_for_actions = prefix_tokens.shape[1]
+            prefix_mask_for_actions = prefix_mask
+            prefix_ar_mask_for_actions = prefix_ar_mask
+        
+        # 6. Knowledge insulation: stop gradients from action expert to VLM
+        # This must happen BEFORE kv_transform so transform still receives gradients
+        if self.config.use_knowledge_insulation:
+            kv_cache_for_actions = jax.tree.map(jax.lax.stop_gradient, kv_cache_for_actions)
+        
+        # 7. Transform KV cache (after stop_gradient, so it receives action expert gradients)
+        if self.kv_transform is not None:
+            kv_cache_for_actions = self.kv_transform(kv_cache_for_actions)
+        
+        # 8. Define single flow sample processing
+        def process_one_flow_sample(sample_rng):
+            """Process one flow sample using the original cached prefix."""
+            noise_rng, time_rng = jax.random.split(sample_rng)
+            
+            # Generate different noise and time for this sample
+            noise = self.generate_correlated_noise(noise_rng, batch_size)
+            time = jax.random.beta(time_rng, 1.5, 1, (batch_size,)) * 0.999 + 0.001
+            
+            # Compute noisy actions and target velocity
+            time_expanded = time[:, None, None]
+            x_t = time_expanded * noise + (1 - time_expanded) * actions
+            u_t = noise - actions
+            
+            # Embed suffix for this sample
+            suffix_tokens, suffix_mask, suffix_ar_mask, adarms_cond = self.embed_suffix(
+                observation, x_t, time
+            )
+            
+            # Build attention mask: suffix attends to prefix (without FAST) + itself
+            # When using KV cache, mask shape should be [batch, suffix_len, prefix_len + suffix_len]
+            suffix_attn_mask = make_attn_mask(suffix_mask, suffix_ar_mask)
+            prefix_attn_mask = einops.repeat(
+                prefix_mask_for_actions, "b p -> b s p", s=suffix_tokens.shape[1]
+            )
+            full_attn_mask = jnp.concatenate([prefix_attn_mask, suffix_attn_mask], axis=-1)
+            
+            # Positions for suffix start after cached prefix
+            suffix_positions = prefix_len_for_actions + jnp.cumsum(suffix_mask, axis=-1) - 1
+            
+            # Forward pass with cached prefix (discard returned cache - don't modify original!)
+            (_, suffix_out), _ = self.PaliGemma.llm(
+                [None, suffix_tokens],
+                mask=full_attn_mask,
+                positions=suffix_positions,
+                kv_cache=kv_cache_for_actions,  # Original cache, reused for all samples
+                adarms_cond=[None, adarms_cond],
+                banks=spatial_banks,
+            )
+            
+            # Compute velocity and loss
+            action_hidden = self.apply_spatial_action_conditioning(
+                observation,
+                suffix_out[:, -self.action_horizon:],
+            )
+            v_t = self.action_out_proj(action_hidden)
+            action_loss = jnp.square(v_t - u_t)  # [B, H, D]
+            
+            return action_loss
+        
+        # 9. Vectorize over N flow samples
+        # Disable type checking inside vmap (jaxtyping doesn't handle traced values well)
+        flow_rngs = jax.random.split(rng, num_flow_samples)
+        with at.disable_typechecking():
+            all_action_losses = jax.vmap(process_one_flow_sample)(flow_rngs)  # [N, B, H, D]
+        
+        # 10. Average over flow samples
+        action_loss = jnp.mean(all_action_losses, axis=0)  # [B, H, D]
+        
+        # 11. Build per-dimension action losses
+        # Base velocity (x,y,z)
+        losses["action_loss_base_vel_x"] = jnp.mean(action_loss[..., 0], axis=-1)
+        losses["action_loss_base_vel_y"] = jnp.mean(action_loss[..., 1], axis=-1)
+        losses["action_loss_base_vel_z"] = jnp.mean(action_loss[..., 2], axis=-1)
+        
+        # Trunk joints (4)
+        for i in range(4):
+            losses[f"action_loss_trunk_{i}"] = jnp.mean(action_loss[..., 3+i], axis=-1)
+            
+        # Left arm joints (7)
+        for i in range(7):
+            losses[f"action_loss_left_arm_{i}"] = jnp.mean(action_loss[..., 7+i], axis=-1)
+            
+        # Left gripper
+        losses["action_loss_left_gripper"] = jnp.mean(action_loss[..., 14], axis=-1)
+        
+        # Right arm joints (7)
+        for i in range(7):
+            losses[f"action_loss_right_arm_{i}"] = jnp.mean(action_loss[..., 15+i], axis=-1)
+            
+        # Right gripper
+        losses["action_loss_right_gripper"] = jnp.mean(action_loss[..., 22], axis=-1)
+
+        # Total action loss: mean over horizon (H) and action dims (D) -> [B]
+        losses["action_loss"] = jnp.mean(action_loss, axis=(-2, -1))
+
+        # 12. Add subtask loss during training
+        subtask_loss_value = 0.0
+        if train and observation.tokenized_prompt.shape[1] > 1:            
+            ground_truth_subtask = observation.tokenized_prompt[:, 1]
+            subtask_loss = -jax.nn.log_softmax(subtask_logits)[
+                jnp.arange(ground_truth_subtask.shape[0]), ground_truth_subtask
+            ]
+            losses["subtask_loss"] = jnp.mean(subtask_loss)
+            losses["subtask_accuracy"] = jnp.mean(
+                jnp.argmax(subtask_logits, axis=-1) == ground_truth_subtask
+            )
+            subtask_loss_value = self.config.subtask_loss_weight * jnp.mean(subtask_loss)
+        
+        # 12b. Aux geometry loss: force the perceiver tokens to carry per-sample geometry (log-depth).
+        geom_aux_value = 0.0
+        aux_w = getattr(getattr(self.config, "da3", None), "aux_geom_weight", 0.0)
+        if train and aux_w > 0 and geom_aux_loss is not None:
+            losses["geom_aux_loss"] = geom_aux_loss
+            geom_aux_value = aux_w * geom_aux_loss
+
+        # 13. Total loss
+        losses["total_loss"] = losses["action_loss"] + subtask_loss_value + fast_loss_value + geom_aux_value
+
+        return losses
+
+    @override
+    def sample_actions(
+        self,
+        rng: at.KeyArrayLike,
+        observation: Observation,
+        *,
+        num_steps: int | at.Int[at.Array, ""] = 20,
+        noise: at.Float[at.Array, "b ah ad"] | None = None,
+        initial_actions: at.Float[at.Array, "b n ad"] | None = None,
+        prefix_tokens: at.Float[at.Array, "b p emb"] | None = None,
+        prefix_mask: at.Bool[at.Array, "b p"] | None = None,
+        prefix_ar_mask: at.Bool[at.Array, "b p"] | None = None,
+    ) -> _model.Actions:
+        observation = preprocess_observation(None, observation, train=False)
+        # Note that we use the convention more common in diffusion literature, where t=1 is noise and t=0 is the target
+        # distribution. yes, this is the opposite of the pi0 paper, and I'm sorry.
+        dt = -1.0 / num_steps
+        batch_size = observation.state.shape[0]
+        
+        # Generate or constrain noise based on inpainting requirements
+        if initial_actions is not None:
+            # INPAINTING PATH: Construct constrained noise z that satisfies initial_actions
+            num_initial_actions = initial_actions.shape[1]
+            input_action_dim = initial_actions.shape[2]
+            
+            # Pad initial_actions to full model dimensions (32D) and action_horizon (30)
+            if input_action_dim < self.action_dim:
+                action_padding = jnp.zeros((batch_size, num_initial_actions, self.action_dim - input_action_dim))
+                initial_actions_full_dim = jnp.concatenate([initial_actions, action_padding], axis=2)
+            else:
+                initial_actions_full_dim = initial_actions[:, :, :self.action_dim]
+            
+            if num_initial_actions < self.action_horizon:
+                seq_padding = jnp.zeros((batch_size, self.action_horizon - num_initial_actions, self.action_dim))
+                initial_actions_padded = jnp.concatenate([initial_actions_full_dim, seq_padding], axis=1)
+            else:
+                initial_actions_padded = initial_actions_full_dim[:, :self.action_horizon]
+            
+            # Compute O and U indices for inpainting (JIT-safe: static list comprehensions)
+            flat_dim = self.action_horizon * self.action_dim
+            
+            # Build O_indices: first num_initial_actions timesteps, first input_action_dim dimensions
+            O_indices = jnp.array([
+                t * self.action_dim + d 
+                for t in range(num_initial_actions) 
+                for d in range(input_action_dim)
+            ], dtype=jnp.int32)
+            
+            # Build U_indices: all other indices (JIT-safe: static list comprehension)
+            # Python set operations happen at trace time (before JIT), so this is safe
+            O_set = {t * self.action_dim + d for t in range(num_initial_actions) for d in range(input_action_dim)}
+            U_indices = jnp.array([
+                i for i in range(flat_dim) if i not in O_set
+            ], dtype=jnp.int32)
+            
+            # Generate noise
+            rng, noise_rng = jax.random.split(rng)
+            
+            if self.correlation_loaded:
+                # CORRELATED NOISE: Sample with correlation matrix
+                noise = self.generate_correlated_noise(noise_rng, batch_size)
+            else:
+                # FALLBACK: Independent noise
+                noise = jax.random.normal(noise_rng, (batch_size, self.action_horizon, self.action_dim))
+            
+            # Extract fixed z_O and x0_O for constraint enforcement
+            noise_flat = noise.reshape(batch_size, flat_dim)
+            fixed_z_O = noise_flat[:, O_indices]  # [b, |O|] - fixed noise for inpainting
+            x0_O = initial_actions_padded.reshape(batch_size, flat_dim)[:, O_indices]  # [b, |O|] - target actions
+            
+            # Precompute correction matrix for correlation-aware inpainting
+            inpainting_cache = None
+            if self.correlation_loaded:
+                cache_key = (num_initial_actions, input_action_dim)
+                if cache_key not in self.inpainting_cache:
+                    logger.info(f"Computing correction matrix for {num_initial_actions} steps, {input_action_dim} dims...")
+                    self.inpainting_cache[cache_key] = self._precompute_correction_matrix(O_indices, U_indices)
+                inpainting_cache = self.inpainting_cache[cache_key]
+            
+        else:
+            # NO INPAINTING: Standard noise generation
+            if noise is None:
+                rng, noise_rng = jax.random.split(rng)
+                noise = self.generate_correlated_noise(noise_rng, batch_size)
+            
+            fixed_z_O = None
+            x0_O = None
+            O_indices = None
+            inpainting_cache = None
+
+        # Split RNG for step loop
+        rng, step_rng = jax.random.split(rng)
+
+        # Ensure FAST tokens are never used during inference
+        if observation.fast_tokens is not None:
+            raise ValueError(
+                "FAST tokens must not be provided during inference (sample_actions). "
+                "FAST tokens are only used during training for auxiliary loss. "
+                "Set observation.fast_tokens=None before calling sample_actions."
+            )
+
+        # Allow cache-generation callers to reuse the exact prefix embeddings
+        # they also pool as conditioning context. The default path is unchanged
+        # for policy inference and existing checkpoints.
+        supplied_prefix = (
+            prefix_tokens is not None,
+            prefix_mask is not None,
+            prefix_ar_mask is not None,
+        )
+        if any(supplied_prefix) and not all(supplied_prefix):
+            raise ValueError(
+                "prefix_tokens, prefix_mask, and prefix_ar_mask must be supplied together"
+            )
+        if prefix_tokens is None:
+            prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
+
+        # First fill KV cache with a forward pass of the prefix (no FAST tokens during inference)
+        prefix_attn_mask = make_attn_mask(prefix_mask, prefix_ar_mask)
+        positions = jnp.cumsum(prefix_mask, axis=1) - 1
+        (prefix_out, _), kv_cache = self.PaliGemma.llm([prefix_tokens, None], mask=prefix_attn_mask, positions=positions)
+
+        # DA3 banks: computed once, reused across all denoise steps.
+        spatial_banks = self._compute_banks(observation)
+        
+        # Predict stage from VLM output of base task token
+        # Find base task token position (same logic as in compute_detailed_loss)
+        first_stage_token_idx = jnp.argmax(prefix_ar_mask)  # Returns index of first True
+        base_task_token_idx = first_stage_token_idx - 1
+        base_task_output = prefix_out[:, base_task_token_idx, :]
+        subtask_logits = self.stage_pred_from_vlm(base_task_output)  # [B, MAX_NUM_STAGES]
+        
+        # Mask out invalid stages for each task (vectorized JAX operations)
+        task_ids = observation.tokenized_prompt[:, 0]  # [B]
+        task_num_stages_array = jnp.array(TASK_NUM_STAGES, dtype=jnp.int32)
+        task_num_stages = task_num_stages_array[task_ids]  # [B] - JAX array indexing
+        stage_range = jnp.arange(MAX_NUM_STAGES)  # [15]
+        valid_mask = stage_range[None, :] < task_num_stages[:, None]  # [B, 15]
+        subtask_logits = jnp.where(valid_mask, subtask_logits, -jnp.inf)
+        
+        # Transform KV cache for cross-layer attention
+        if self.kv_transform is not None:
+            kv_cache = self.kv_transform(kv_cache)
+        
+        def step(carry):
+            x_t, time, step_rng = carry
+            
+            # Use config value for time threshold
+            TIME_THRESHOLD_INPAINT = self.config.time_threshold_inpaint
+            
+            # Model forward pass
+            suffix_tokens, suffix_mask, suffix_ar_mask, adarms_cond = self.embed_suffix(
+                observation, x_t, jnp.broadcast_to(time, batch_size)
+            )
+            suffix_attn_mask = make_attn_mask(suffix_mask, suffix_ar_mask)
+            prefix_attn_mask = einops.repeat(prefix_mask, "b p -> b s p", s=suffix_tokens.shape[1])
+            full_attn_mask = jnp.concatenate([prefix_attn_mask, suffix_attn_mask], axis=-1)
+            assert full_attn_mask.shape == (
+                batch_size,
+                suffix_tokens.shape[1],
+                prefix_tokens.shape[1] + suffix_tokens.shape[1],
+            )
+            positions = jnp.sum(prefix_mask, axis=-1)[:, None] + jnp.cumsum(suffix_mask, axis=-1) - 1
+
+            (prefix_out, suffix_out), _ = self.PaliGemma.llm(
+                [None, suffix_tokens],
+                mask=full_attn_mask,
+                positions=positions,
+                kv_cache=kv_cache,
+                adarms_cond=[None, adarms_cond],
+                banks=spatial_banks,
+            )
+            assert prefix_out is None
+            action_hidden = self.apply_spatial_action_conditioning(
+                observation,
+                suffix_out[:, -self.action_horizon :],
+            )
+            v_t = self.action_out_proj(action_hidden)
+            
+            # Euler step: x_{t+dt} = x_t + dt * v_t
+            x_t_new = x_t + dt * v_t
+            
+            # Apply correlation-aware inpainting correction
+            # Only enforce when time > TIME_THRESHOLD_INPAINT (let model be free in final steps)
+            if fixed_z_O is not None:
+                time_new = time + dt
+                
+                def apply_correlated_correction(x):
+                    x_flat = x.reshape(batch_size, -1)
+                    
+                    # Compute desired state at O: x_t[O] = (1-t)*x0[O] + t*z_O
+                    x_desired_O = (1.0 - time_new) * x0_O + time_new * fixed_z_O  # [b, |O|]
+                    
+                    # Compute correction at O
+                    delta_O = x_desired_O - x_flat[:, O_indices]  # [b, |O|]
+                    
+                    # Apply hard constraint at O
+                    x_flat = x_flat.at[:, O_indices].set(x_desired_O)
+                    
+                    # If correlation matrix available, propagate correction to U
+                    if inpainting_cache is not None:
+                        correction_matrix = inpainting_cache['correction_matrix']  # [|U|, |O|]
+                        U_indices_cached = inpainting_cache['U_indices']
+                        
+                        # Compute correlated correction: δ_U = Σ_{UO}Σ_{OO}^{-1} @ δ_O
+                        delta_U = delta_O @ correction_matrix.T  # [b, |U|]
+                        
+                        # Skip if correction too large (indicates instability)
+                        max_correction = jnp.max(jnp.abs(delta_U))
+                        x_flat = jax.lax.cond(
+                            # Prevents exploding corrections in case of noisy out of distribution initial actions
+                            max_correction <= 1.0,
+                            lambda x: x.at[:, U_indices_cached].add(delta_U),
+                            lambda x: x,
+                            x_flat
+                        )
+                    
+                    # Sanity check: if Σ = I, correction_matrix = 0, so delta_U = 0 that is correct
+                    # If the correlation is 1 everywhere we will go to the flat prediction that is correct
+                    
+                    return x_flat.reshape(batch_size, self.action_horizon, self.action_dim)
+                
+                # Only apply correction when NEW time > threshold
+                x_t_new = jax.lax.cond(
+                    time_new > TIME_THRESHOLD_INPAINT,
+                    apply_correlated_correction,
+                    lambda x: x,
+                    x_t_new
+                )
+            
+            return x_t_new, time + dt, step_rng
+
+        def cond(carry):
+            x_t, time, step_rng = carry
+            # Robust to floating-point error
+            return time >= -dt / 2
+
+        x_0, _, _ = jax.lax.while_loop(cond, step, (noise, 1.0, step_rng))
+        
+        return x_0, subtask_logits

legacy/vggt_newbank_boxing_gloves_step19999/code/src/b1k/models/pi_behavior_config.py

@@ -0,0 +1,291 @@
+"""PI_BEHAVIOR Model Configuration
+
+Configuration for PI_BEHAVIOR model on BEHAVIOR-1K challenge.
+"""
+
+import dataclasses
+import json
+import pathlib
+from typing import TYPE_CHECKING
+
+import flax.nnx as nnx
+import jax
+import jax.numpy as jnp
+from typing_extensions import override
+
+from openpi.models import model as _model
+from openpi.models import gemma as _gemma
+from openpi.shared import array_typing as at
+import openpi.shared.nnx_utils as nnx_utils
+
+from b1k.models.observation import Observation
+
+if TYPE_CHECKING:
+    from b1k.models.pi_behavior import PiBehavior
+
+
+# Per-task stage counts (based on avg_episode_length / 900, capped between 5-15)
+# Use tuple for immutability and to avoid JAX device allocation at import time
+TASK_NUM_STAGES = (
+    5, 6, 15, 15, 14, 12, 9, 15, 10, 15,  # Tasks 0-9
+    7, 13, 10, 15, 15, 15, 15, 11, 13, 12,  # Tasks 10-19
+    14, 15, 9, 15, 15, 15, 15, 15, 15, 15,  # Tasks 20-29
+    11, 10, 10, 13, 5, 5, 14, 6, 8, 10,  # Tasks 30-39
+    5, 15, 8, 15, 12, 11, 9, 14, 15, 15,  # Tasks 40-49
+)
+
+MAX_NUM_STAGES = 15  # Maximum stages per task
+TOTAL_TASK_STAGE_EMBEDDINGS = sum(TASK_NUM_STAGES)  # 596 total embeddings
+
+# Cumulative offsets for indexing into task_stage_embeddings (as tuple)
+TASK_STAGE_OFFSETS = tuple([0] + [sum(TASK_NUM_STAGES[:i+1]) for i in range(len(TASK_NUM_STAGES) - 1)])
+
+
+@dataclasses.dataclass(frozen=True)
+class B1KDA3Config:
+    """DA3 spatial-language branch for PiBehavior (inline extraction; b1k cameras are square)."""
+
+    enabled: bool = True
+    num_views: int = 3  # zed head (main), left/right realsense (wrist branches)
+    da3_channels: int = 1536  # GIANT embed dim
+    da3_layers: int = 4  # out_layers (19, 26, 33, 39)
+    grid_hw: tuple[int, int] = (18, 18)  # 252x252 square DA3 input / patch 14
+    hidden_dim: int = 1024  # == action-expert width
+    lang_dim: int = 1024  # ModernBERT-large (task-name embeddings)
+    lang_max_len: int = 32
+    num_heads: int = 8
+    lang_fusion_depth: int = 2
+    num_inject_layers: int = 6  # last 6 of 18 action-expert blocks
+    spatial_scale: float = 2.0
+    # V2 "force-spatial-on" defaults (now that geometry is CORRECT). Zero-init lets the model learn to
+    # IGNORE spatial (image path fits first, no gradient left to turn the injection on). Nonzero init +
+    # per-head logit-gain keep the injection ACTIVE and the attention SHARP/learnable from step 0, so the
+    # model must account for the (now-sane) banks. This only hurt before because geometry was garbage.
+    spatial_init_std: float = 0.01
+    attn_logit_gain: bool = True
+    attn_logit_gain_init: float = 3.0   # retuned for qk_norm: 20 eff tokens of 324
+    attn_logit_gain_max: float = 8.0    # gain 8 -> 2.5 eff tokens; hard ceiling
+    bank_token_embed: bool = True
+    perceiver_query_std: float = 0.05
+    # Perceiver-collapse fixes. Default False preserves the arch of existing checkpoints.
+    #  root cause: random-init queries -> q.k ~ 0 -> near-uniform softmax over 432 patches
+    #  -> every query reads the same mean(V) AND dL/dQ,K is starved (~1/432) so queries never
+    #  train; the shared output (||.||~500) then swamps query identity (||q||~1.6) ~300:1.
+    perceiver_logit_gain: bool = False       # sharpen attention at init -> diverse reads + live Q/K grads
+    # --- 2026-07-22 attention-saturation fixes (see DA3_ATTENTION_SATURATION.md) ---
+    qk_norm: bool = True                 # per-head RMSNorm on Q,K before the dot product
+    perceiver_norm_out: bool = True      # LayerNorm the perceiver output (was amplifying x1900)
+    pos_emb_scale: float = 0.25          # constant pos_emb was rms 5.03 vs signal 4.38
+    perceiver_logit_gain_init: float = 3.0  # retuned for qk_norm (was 8 -> 2.5 eff tokens)
+    perceiver_logit_gain_max: float = 8.0
+    perceiver_norm_attn_out: bool = False    # LN attn-out before residual -> query identity survives
+    # --- 2026-07-23 constant-collapse fix (see b1k-da3-frozenbase-verdict) ---
+    # The bank was measured ~90% learned-constant (view/pos/lang/bank_token embeds) vs ~10% per-sample
+    # DA3 content; the frozen base latched onto the constant (net-harmful: zeroing the bank cut loss 92%)
+    # and never used geometry (shuffling banks across samples moved loss +0.2%). bank_center projects out
+    # the batch-mean so a constant injects EXACTLY zero -- only per-sample deviation survives, forcing the
+    # model to use geometry or nothing. NOTE: like batchnorm, needs bs>1; deploy at bs=1 needs an EMA of
+    # the mean (TODO) -- the current-batch projection is for the "does geometry get used" experiment.
+    bank_center: bool = False
+    # --- 2026-07-23 aux geometry loss ---
+    # Decode the perceiver token output back to per-patch log-depth (grid-pos queries attend the K
+    # perceiver tokens). MSE against the DA3 depth FORCES the perceiver output to carry per-sample
+    # geometry regardless of the action loss's incentive -- the guaranteed fix for "geometry unused".
+    aux_geom_head: bool = False      # build the decoder head
+    aux_geom_weight: float = 0.0     # weight of the log-depth MSE in the total loss
+    # Zero the log-depth INPUT channel (ray7 ch 6) so depth is target-only. Without this the aux task
+    # is circular (depth in -> depth out, a trivial autoencoder); with it, predicting depth REQUIRES
+    # reading it out of the DA3 features. Shape-compatible (channel zeroed, not removed).
+    depth_target_only: bool = False
+    # --- 2026-07-23 K/V split (address/payload separation in the perceiver) ---
+    # payload (values) = DA3 latents + depth encoding; address (keys only) = pos_emb + ray_emb +
+    # view_emb. Addresses steer routing but are structurally excluded from the value stream, so an
+    # input-independent constant can no longer flow into (and dominate) the bank. depth_dropout
+    # zeroes the depth encoding for that fraction of training samples so the DA3 features must carry
+    # geometry redundantly. NOTE: kv_split changes the spatial arch (ray_mlp 7ch -> 6ch + depth_mlp);
+    # spatial params are NOT checkpoint-compatible across this flag.
+    kv_split: bool = False
+    depth_dropout: float = 0.0
+    # --- 2026-07-24 spatial-bank upgrades ---
+    # perc_locality: anchor each perceiver query to a grid region with a learnable -gamma*dist^2 logit
+    #   bias, so tokens are LOCAL descriptors (fixes over-averaging) instead of global scene means.
+    # cross_view: after the per-view perceivers, add a camera-pose embed and self-attend across the
+    #   concatenated view tokens so the three views fuse into one 3D scene (then split back per view).
+    perc_locality: bool = False
+    cross_view: bool = False
+    cross_view_depth: int = 2
+    bank_token_embed_query: bool = True  # False = old post-fusion placement (faithful eval of old ckpts)
+    # --- VGGT-Omega enrichments (v2): extra bank inputs harvested from the VGGT forward; all no-ops
+    # unless the loader is the VGGT extractor (which supplies da3_depth_conf/pose_enc/cam_tokens). ---
+    use_depth_conf: bool = False   # add VGGT depth confidence as a payload reliability channel
+    use_pose_enc: bool = False     # add VGGT pose encoding to the cross-view camera feature
+    use_cam_tokens: bool = False   # append VGGT camera+register tokens as global bank tokens
+    cam_token_dim: int = 2048      # channel width of da3_cam_tokens (VGGT 2*embed_dim)
+    pose_enc_dim: int = 9          # VGGT pose_enc width (trans3+quat4+fov2)
+    feat_input_norm: bool = False  # LayerNorm raw backbone feats before projection (tames VGGT outliers)
+
+
+@dataclasses.dataclass(frozen=True)
+class PiBehaviorConfig(_model.BaseModelConfig):
+    dtype: str = "bfloat16"
+    paligemma_variant: _gemma.Variant = "gemma_2b"
+    action_expert_variant: _gemma.Variant = "gemma_300m"
+
+    # Set the model specific defaults.
+    action_dim: int = 32
+    action_horizon: int = 30
+    max_token_len: int = 200  # Only used for compatibility, not for actual tokenization
+    
+    # Number of tasks in the behavior dataset
+    num_tasks: int = 50
+    # Task embedding dimension - will match the paligemma width
+    task_embedding_dim: int = None  # type: ignore
+    # Maximum number of subtask states across all tasks
+    max_num_subtask_states: int = MAX_NUM_STAGES
+    
+    # Path to task data JSON file for initialization
+    task_data_path: str = "b1k/BEHAVIOR-1K/docs/challenge/task_data.json"
+    
+    # Whether to use correlated noise matching action covariance structure
+    # Requires correlation matrix in norm_stats (computed by compute_norm_stats.py)
+    use_correlated_noise: bool = True
+    
+    # Shrinkage parameter for correlation regularization
+    # Applied as: S_regularized = beta * S + (1-beta) * I
+    # beta=1.0 means full correlation (no shrinkage)
+    # beta=0.7 means 70% correlation + 30% independence (recommended for robustness)
+    # beta=0.0 means independence (no correlation)
+    correlation_beta: float = 0.5
+    
+    # FAST auxiliary training configuration
+    use_fast_auxiliary: bool = False  # Enable FAST during training
+    fast_loss_weight: float = 0.1  # Weight for FAST loss (vs flow loss)
+    
+    # Action dimensions to encode with FAST (default: 0:6, 7:23 = 22 dims)
+    # Format: "0:6,7:23" or list of tuples [(0, 6), (7, 23)]
+    fast_encoded_dims: str | list[tuple[int, int]] = "0:6,7:23"
+    
+    # FAST tokenizer vocab size
+    fast_vocab_size: int = 1024
+    
+    # Max FAST tokens to predict (truncate if exceeded)
+    max_fast_tokens: int = 32
+    
+    # FAST tokenizer path (set during initialization, relative to assets_dir/asset_id)
+    fast_tokenizer_path: str | None = None
+    
+    # KV cache transformation for cross-layer attention between VLM and action expert
+    # Allows each action expert layer to attend to a learned combination of all VLM layers
+    use_kv_transform: bool = True
+    
+    # Knowledge insulation: stop action expert gradients from flowing to VLM backbone
+    # VLM trains on FAST tokens only, action expert on flow matching with frozen VLM features
+    # Implements approach from https://www.physicalintelligence.company/research/knowledge_insulation
+    use_knowledge_insulation: bool = True
+    
+    # Subtask/stage prediction auxiliary loss weight (relative to action loss)
+    # Higher values emphasize stage prediction accuracy at the expense of action quality
+    subtask_loss_weight: float = 0.1
+    
+    # Time threshold for inpainting during inference
+    # Stop enforcing inpainting constraint when t < threshold (let model be free in final steps)
+    time_threshold_inpaint: float = 0.3
+    
+    # Vision backbone finetuning control
+    freeze_vision_backbone: bool = True
+    
+    # DA3 spatial-language adapter. The DA3/ModernBERT branch is computed
+    # offline and supplied as tokens in Observation.spatial_tokens.
+    use_spatial_action_cross_attention: bool = False
+    spatial_token_dim: int = 1024
+    spatial_num_tokens: int = 320  # DA3 perc bank default: 128 + 96 + 96
+    spatial_num_heads: int = 8
+    spatial_residual_scale: float = 1.0
+
+    # Full DA3 spatial-language branch (supersedes the flat spatial_tokens adapter above):
+    # frozen DA3-GIANT runs INLINE in the data pipeline; the trainable bank builder + method-B
+    # cross-attention injection (action-expert layers 12-17) live in the model. Proven on RoboReal.
+    da3: "B1KDA3Config | None" = None
+
+    def __post_init__(self):
+        if self.task_embedding_dim is None:
+            paligemma_config = _gemma.get_config(self.paligemma_variant)
+            object.__setattr__(self, "task_embedding_dim", paligemma_config.width)
+    
+    def get_fast_dim_ranges(self) -> list[tuple[int, int]]:
+        """Parse fast_encoded_dims into list of ranges."""
+        if isinstance(self.fast_encoded_dims, str):
+            ranges = []
+            for range_str in self.fast_encoded_dims.split(','):
+                start, end = map(int, range_str.strip().split(':'))
+                ranges.append((start, end))
+            return ranges
+        return self.fast_encoded_dims
+    
+    def get_total_fast_dims(self) -> int:
+        """Get total number of dimensions encoded by FAST."""
+        return sum(end - start for start, end in self.get_fast_dim_ranges())
+
+    @property
+    @override
+    def model_type(self):
+        return "pi_behavior"
+
+    @override
+    def create(self, rng: at.KeyArrayLike) -> "PiBehavior":
+        from b1k.models.pi_behavior import PiBehavior
+
+        return PiBehavior(self, rngs=nnx.Rngs(rng))
+
+    @override
+    def inputs_spec(self, *, batch_size: int = 1) -> tuple["Observation", _model.Actions]:
+        image_spec = jax.ShapeDtypeStruct([batch_size, *_model.IMAGE_RESOLUTION, 3], jnp.float32)
+        image_mask_spec = jax.ShapeDtypeStruct([batch_size], jnp.bool_)
+
+        with at.disable_typechecking():
+            obs_kwargs = {
+                "images": {
+                    "base_0_rgb": image_spec,
+                    "left_wrist_0_rgb": image_spec,
+                    "right_wrist_0_rgb": image_spec,
+                },
+                "image_masks": {
+                    "base_0_rgb": image_mask_spec,
+                    "left_wrist_0_rgb": image_mask_spec,
+                    "right_wrist_0_rgb": image_mask_spec,
+                },
+                "state": jax.ShapeDtypeStruct([batch_size, self.action_dim], jnp.float32),
+                "tokenized_prompt": jax.ShapeDtypeStruct([batch_size, 2], jnp.int32),
+                "tokenized_prompt_mask": jax.ShapeDtypeStruct([batch_size, 2], bool),
+            }
+            
+            if self.use_fast_auxiliary:
+                obs_kwargs["fast_tokens"] = jax.ShapeDtypeStruct([batch_size, self.max_fast_tokens], jnp.int32)
+                obs_kwargs["fast_token_mask"] = jax.ShapeDtypeStruct([batch_size, self.max_fast_tokens], bool)
+            
+            if self.da3 is not None and self.da3.enabled:
+                d = self.da3
+                gh, gw = d.grid_hw
+                obs_kwargs["da3_features"] = jax.ShapeDtypeStruct(
+                    [batch_size, d.da3_layers, d.num_views, d.da3_channels, gh, gw], jnp.uint16
+                )
+                obs_kwargs["da3_ray"] = jax.ShapeDtypeStruct([batch_size, d.num_views, 3, gh, gw], jnp.float32)
+                obs_kwargs["da3_depth"] = jax.ShapeDtypeStruct([batch_size, d.num_views, 1, gh, gw], jnp.float32)
+                obs_kwargs["camera_extrinsics"] = jax.ShapeDtypeStruct([batch_size, d.num_views, 4, 4], jnp.float32)
+                obs_kwargs["lang_feat"] = jax.ShapeDtypeStruct([batch_size, d.lang_max_len, d.lang_dim], jnp.float32)
+                obs_kwargs["lang_mask"] = jax.ShapeDtypeStruct([batch_size, d.lang_max_len], bool)
+
+            if self.use_spatial_action_cross_attention:
+                obs_kwargs["spatial_tokens"] = jax.ShapeDtypeStruct(
+                    [batch_size, self.spatial_num_tokens, self.spatial_token_dim],
+                    jnp.float32,
+                )
+                obs_kwargs["spatial_token_mask"] = jax.ShapeDtypeStruct(
+                    [batch_size, self.spatial_num_tokens],
+                    bool,
+                )
+            
+            observation_spec = Observation(**obs_kwargs)
+        
+        action_spec = jax.ShapeDtypeStruct([batch_size, self.action_horizon, self.action_dim], jnp.float32)
+        return observation_spec, action_spec

legacy/vggt_newbank_boxing_gloves_step19999/code/src/b1k/models/spatial_da3.py

@@ -0,0 +1,593 @@
+"""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 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:[h,Lq,Lk] (broadcast over batch) additive logit bias, e.g. a locality prior.
+        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:
+            logits = logits + attn_bias[None].astype(logits.dtype)  # [1,h,Lq,Lk] broadcast over batch
+        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 _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,
+                 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))))
+        # 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):
+        # 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]
+        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 _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
+# ---------------------------------------------------------------------------
+
+_VIEWS = (("main", 0, 128), ("left", 1, 96), ("right", 2, 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,
+        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
+        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)
+        self.t5_projector = ProjLN(lang_dim, H, rngs=rngs)
+        # 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 = {
+            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, rngs=rngs)
+            for name, _, k in _VIEWS
+        }
+        self.lang_fusers = {name: LanguageFusionStack(H, lang_fusion_depth, num_heads,
+                                                      qk_norm=qk_norm, rngs=rngs)
+                            for name, _, _ in _VIEWS}
+        # (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)
+        # (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, k, H)) * 0.05) for name, _, k in _VIEWS}
+            if bank_token_embed
+            else None
+        )
+        # 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 __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)
+        lang_tokens = self.t5_projector(lang_feat)  # [b,Lt,H]
+        lang_pad = jnp.logical_not(lang_mask)  # True=pad
+        geos = {}
+        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]
+            view_emb = self.view_embed(jnp.asarray(vidx))[None, None, :]  # [1,1,H]
+            _bte = self.bank_token_embeds[name].value if self.bank_token_embeds is not None else None
+            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_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
+                geo = self.perceivers[name](payload, addr=addr, token_embed=tok_emb)  # [b,K,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,432,H]
+                geo = self.perceivers[name](spatial, token_embed=tok_emb)  # [b,K,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:
+            parts = []
+            for name, vidx, _k in _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
+                parts.append(geos[name] + cam[:, None, :])
+            x = self.cross_view_fusion(jnp.concatenate(parts, axis=1))  # [b, sum_k, H]
+            off = 0
+            for name, _vidx, k in _VIEWS:
+                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)  # [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
+        if return_aux:
+            aux = jnp.mean(jnp.stack(aux_losses)) if aux_losses else jnp.asarray(0.0, jnp.float32)
+            return banks, aux
+        return banks

legacy/vggt_newbank_boxing_gloves_step19999/code/src/b1k/training/b1k_2026.py

@@ -0,0 +1,414 @@
+"""BEHAVIOR-1K 2026 (LeRobot v3.0) data loading for the 2025 PiBehavior model.
+
+The 2025 winner's loader targets OmniGibson's v2.1 `BehaviorLeRobotDataset`
+(task-partitioned, 256-dim proprio). The 2026 challenge dataset is LeRobot
+**v3.0** (chunk-based `data/chunk-XXX/file-XXX.parquet`, RGB in videos, 61-dim
+proprio, 100 tasks). The pinned lerobot (v2.1) cannot read it, so this module
+provides a self-contained v3 reader that yields items in the exact dict format
+the existing transform pipeline expects, plus the two remaps needed:
+
+  * camera keys: zed_link -> head, left/right_realsense -> left/right_wrist
+  * task_index: 2026 index -> the 2025 index the checkpoint's task-embedding
+    table is keyed on (via activity name <-> task_data.json)
+
+Videos are decoded with PyAV (torchcodec's ffmpeg libs are absent here).
+"""
+from __future__ import annotations
+
+import dataclasses
+import functools
+import glob
+import hashlib
+import json
+import logging
+import os
+import time
+from collections import OrderedDict
+from typing import Dict, List, Optional
+
+import av
+import einops
+import numpy as np
+import pandas as pd
+import torch
+
+from openpi import transforms
+from openpi.models import model as _model
+
+logger = logging.getLogger("b1k.2026")
+
+# ---- 2026 R1Pro proprioception layout (61-dim), from BEHAVIOR-1K main
+# OmniGibson/omnigibson/eval/utils/eval_utils.py::PROPRIOCEPTION_INDICES["R1Pro"]
+PROP_2026 = {
+    "base_qvel": slice(0, 3),
+    "arm_left_qpos": slice(3, 10),
+    "gripper_left_qpos": slice(24, 26),
+    "arm_right_qpos": slice(28, 35),
+    "gripper_right_qpos": slice(49, 51),
+    "trunk_qpos": slice(53, 57),
+}
+MAX_GRIPPER_WIDTH = 0.1  # matches the 2025 gripper normalization
+
+
+def extract_state_2026(proprio: np.ndarray) -> np.ndarray:
+    """61-dim 2026 proprio -> 23-dim model state, in the SAME field order the
+    2025 `extract_state_from_proprio` produced (base_qvel, trunk, arm_left,
+    gripper_left, arm_right, gripper_right)."""
+    p = np.asarray(proprio, dtype=np.float32)
+    base_qvel = p[..., PROP_2026["base_qvel"]]                 # 3
+    trunk_qpos = p[..., PROP_2026["trunk_qpos"]]               # 4
+    arm_left = p[..., PROP_2026["arm_left_qpos"]]              # 7
+    arm_right = p[..., PROP_2026["arm_right_qpos"]]            # 7
+    lg = p[..., PROP_2026["gripper_left_qpos"]].sum(-1, keepdims=True)
+    rg = p[..., PROP_2026["gripper_right_qpos"]].sum(-1, keepdims=True)
+    lg = 2.0 * (lg / MAX_GRIPPER_WIDTH) - 1.0                  # -> [-1,1]
+    rg = 2.0 * (rg / MAX_GRIPPER_WIDTH) - 1.0
+    return np.concatenate([base_qvel, trunk_qpos, arm_left, lg, arm_right, rg], axis=-1)
+
+
+@dataclasses.dataclass(frozen=True)
+class B1kInputs2026(transforms.DataTransformFn):
+    """Same as b1k_policy.B1kInputs but with 2026 61-dim state extraction and
+    no OmniGibson dependency."""
+    model_type: object = _model.ModelType.PI0
+
+    def __call__(self, data: dict) -> dict:
+        state = extract_state_2026(data["observation/state"])
+
+        def _img(x):
+            x = np.asarray(x)
+            if np.issubdtype(x.dtype, np.floating):
+                x = (255 * x).astype(np.uint8)
+            if x.shape[0] == 3:
+                x = einops.rearrange(x, "c h w -> h w c")
+            return x
+
+        names = ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
+        imgs = (_img(data["observation/egocentric_camera"]),
+                _img(data["observation/wrist_image_left"]),
+                _img(data["observation/wrist_image_right"]))
+        out = {
+            "state": state,
+            "image": dict(zip(names, imgs, strict=True)),
+            "image_mask": {n: np.True_ for n in names},
+        }
+        for k in ("actions", "task_index", "timestamp", "episode_index",
+                  "tokenized_prompt", "tokenized_prompt_mask", "subtask_state"):
+            if k in data:
+                out[k] = data[k]
+        return out
+
+
+class _V3Meta:
+    """Minimal `dataset.meta` shim for ComputeSubtaskStateFromMeta.
+    `.episodes` maps episode_index -> {'length': int}."""
+    def __init__(self, episodes: Dict[int, dict]):
+        self.episodes = episodes
+
+
+# --------------------------------------------------------------------------- #
+# task-index remapping: 2026 index -> 2025 index (what the checkpoint knows)
+# --------------------------------------------------------------------------- #
+def build_task_index_maps(root_2026: str, task_data_json: str):
+    """Return (name->2025idx, 2026idx->2025idx, name->2026idx)."""
+    td = json.load(open(task_data_json))["tasks"]
+    name2025 = {t["id"]: i for i, t in enumerate(td)}          # activity -> 2025 idx
+    dt = pd.read_parquet(os.path.join(root_2026, "meta", "tasks.parquet"))
+    # tasks.parquet: index = activity name, column task_index (2026)
+    name2026 = {name: int(row["task_index"]) for name, row in dt.iterrows()}
+    idx2026_to_2025 = {name2026[n]: name2025[n] for n in name2026 if n in name2025}
+    return name2025, idx2026_to_2025, name2026
+
+
+class BehaviorV3Dataset(torch.utils.data.Dataset):
+    """LeRobot v3.0 reader for a subset of activities. Yields per-frame items in
+    the 2025-style LeRobot dict format (pre-repack keys), with task_index already
+    remapped to the 2025 index and RGB decoded to uint8 HWC."""
+
+    RGB_KEYS = OrderedDict([
+        ("observation.images.rgb.head", "observation.rgb.zed_link_camera_0"),
+        ("observation.images.rgb.left_wrist", "observation.rgb.left_realsense_link_camera_0"),
+        ("observation.images.rgb.right_wrist", "observation.rgb.right_realsense_link_camera_0"),
+    ])
+
+    def __init__(self, root: str, activities: List[str], action_horizon: int,
+                 task_data_json: str, seed: int = 0, parquet_cache: int = 16):
+        self.root = root
+        self.H = int(action_horizon)
+        self.fps = float(json.load(open(os.path.join(root, "meta", "info.json")))["fps"])
+        _, self.idx2026_to_2025, self.name2026 = build_task_index_maps(root, task_data_json)
+        acts = set(activities)
+
+        # episode metadata (filtered to our activities)
+        ep = pd.concat([pd.read_parquet(f) for f in sorted(
+            glob.glob(os.path.join(root, "meta", "episodes", "**", "*.parquet"), recursive=True))],
+            ignore_index=True)
+        ep["task0"] = ep["tasks"].apply(lambda v: v[0] if hasattr(v, "__len__") and not isinstance(v, str) else v)
+        ep = ep[ep["task0"].isin(acts)].reset_index(drop=True)
+
+        self.episodes: List[dict] = []
+        meta_eps: Dict[int, dict] = {}
+        samples: List[tuple] = []
+        for _, r in ep.iterrows():
+            E = int(r["episode_index"]); L = int(r["length"])
+            rec = {
+                "episode_index": E, "length": L, "task0": r["task0"],
+                "data": os.path.join(root, "data", f"chunk-{int(r['data/chunk_index']):03d}",
+                                     f"file-{int(r['data/file_index']):03d}.parquet"),
+                "video": {}, "from_ts": {},
+            }
+            for dst, src in self.RGB_KEYS.items():
+                rec["video"][dst] = os.path.join(
+                    root, "videos", src,
+                    f"chunk-{int(r[f'videos/{src}/chunk_index']):03d}",
+                    f"file-{int(r[f'videos/{src}/file_index']):03d}.mp4")
+                rec["from_ts"][dst] = float(r[f"videos/{src}/from_timestamp"])
+            ei = len(self.episodes)
+            self.episodes.append(rec)
+            meta_eps[E] = {"length": L}
+            # only frames with a full future action window
+            for t in range(max(1, L - self.H)):
+                samples.append((ei, t))
+        self.samples = samples
+        self.meta = _V3Meta(meta_eps)
+        self._pq_cache: "OrderedDict[str, pd.DataFrame]" = OrderedDict()
+        self._pq_cache_max = parquet_cache
+        self._video_cache: "OrderedDict[str, av.container.InputContainer]" = OrderedDict()
+        self._video_cache_max = int(os.environ.get("B1K_VIDEO_CACHE_SIZE", "12"))
+        self._decode_resize = int(os.environ.get("B1K_DECODE_RESIZE", "224"))
+        self._frame_cache_dir = os.environ.get("B1K_FRAME_CACHE_DIR")
+        self._frame_cache_max_bytes = int(float(os.environ.get("B1K_FRAME_CACHE_MAX_GB", "4")) * (1024 ** 3))
+        self._frame_cache_prune_every = max(1, int(os.environ.get("B1K_FRAME_CACHE_PRUNE_EVERY", "2048")))
+        self._frame_cache_checks = 0
+        self._frame_mem_cache: "OrderedDict[tuple[str, int], np.ndarray]" = OrderedDict()
+        self._frame_mem_cache_bytes = 0
+        self._frame_mem_cache_max_bytes = int(float(os.environ.get("B1K_FRAME_MEM_CACHE_GB", "0")) * (1024 ** 3))
+        self._frame_cache_touch_disk = os.environ.get("B1K_FRAME_CACHE_TOUCH", "0") == "1"
+        if self._frame_cache_dir:
+            os.makedirs(self._frame_cache_dir, exist_ok=True)
+            if os.environ.get("B1K_FRAME_CACHE_PRUNE_ON_INIT", "0") == "1":
+                self._prune_frame_cache(force=True)
+        logger.info("BehaviorV3Dataset: %d episodes, %d frame-samples, %d activities",
+                    len(self.episodes), len(self.samples), len(acts))
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["_pq_cache"] = OrderedDict()
+        state["_video_cache"] = OrderedDict()
+        state["_frame_mem_cache"] = OrderedDict()
+        state["_frame_mem_cache_bytes"] = 0
+        return state
+
+    def __del__(self):
+        for container in getattr(self, "_video_cache", {}).values():
+            try:
+                container.close()
+            except Exception:
+                pass
+
+    def __len__(self):
+        return len(self.samples)
+
+    def _episode_frames(self, rec) -> pd.DataFrame:
+        """Cached per-episode frame table (state, action, timestamp), sorted."""
+        key = rec["data"]
+        if key not in self._pq_cache:
+            df = pd.read_parquet(key, columns=["episode_index", "frame_index",
+                                               "observation.state", "action",
+                                               "timestamp", "task_index"])
+            self._pq_cache[key] = df
+            if len(self._pq_cache) > self._pq_cache_max:
+                self._pq_cache.popitem(last=False)
+        df = self._pq_cache[key]
+        sub = df[df["episode_index"] == rec["episode_index"]].sort_values("frame_index")
+        return sub
+
+    def _cached_container(self, path: str):
+        container = self._video_cache.get(path)
+        if container is not None:
+            self._video_cache.move_to_end(path)
+            return container
+
+        container = av.open(path)
+        # Cap ffmpeg decode threads per stream. HEVC's default thread_type=AUTO spawns up to ncores
+        # (240 here) threads PER container; with many workers x 6 containers/sample this blows past the
+        # kernel/cgroup thread ceiling ("can't start new thread"). A small fixed count is plenty since
+        # parallelism comes from the dataloader workers, not per-decode threads.
+        _dt = int(os.environ.get("B1K_DECODE_THREADS", "1"))
+        try:
+            vs0 = container.streams.video[0]
+            vs0.thread_count = _dt
+            vs0.thread_type = "NONE" if _dt <= 1 else "FRAME"
+        except Exception:
+            pass
+        self._video_cache[path] = container
+        if len(self._video_cache) > self._video_cache_max:
+            _, old = self._video_cache.popitem(last=False)
+            old.close()
+        return container
+
+    def _frame_cache_path(self, path: str, frame_idx: int) -> str | None:
+        if not self._frame_cache_dir:
+            return None
+        key = hashlib.blake2b(f"{path}|{frame_idx}|{self._decode_resize}".encode(), digest_size=16).hexdigest()
+        return os.path.join(self._frame_cache_dir, f"{key}.npy")
+
+    def _get_frame_mem_cache(self, key: tuple[str, int]) -> np.ndarray | None:
+        if self._frame_mem_cache_max_bytes <= 0:
+            return None
+        img = self._frame_mem_cache.get(key)
+        if img is None:
+            return None
+        self._frame_mem_cache.move_to_end(key)
+        return img
+
+    def _put_frame_mem_cache(self, key: tuple[str, int], img: np.ndarray) -> None:
+        if self._frame_mem_cache_max_bytes <= 0:
+            return
+        old = self._frame_mem_cache.pop(key, None)
+        if old is not None:
+            self._frame_mem_cache_bytes -= old.nbytes
+        self._frame_mem_cache[key] = img
+        self._frame_mem_cache_bytes += img.nbytes
+        while self._frame_mem_cache_bytes > self._frame_mem_cache_max_bytes and self._frame_mem_cache:
+            _, evicted = self._frame_mem_cache.popitem(last=False)
+            self._frame_mem_cache_bytes -= evicted.nbytes
+
+    def _prune_frame_cache(self, *, force: bool = False):
+        if not self._frame_cache_dir or self._frame_cache_max_bytes <= 0:
+            return
+        self._frame_cache_checks += 1
+        if not force and self._frame_cache_checks % self._frame_cache_prune_every:
+            return
+        files = []
+        total = 0
+        for p in glob.glob(os.path.join(self._frame_cache_dir, "*.npy")):
+            try:
+                st = os.stat(p)
+            except FileNotFoundError:
+                continue
+            total += st.st_size
+            files.append((st.st_mtime, st.st_size, p))
+        if total <= self._frame_cache_max_bytes:
+            return
+        for _, size, p in sorted(files):
+            try:
+                os.remove(p)
+                total -= size
+            except FileNotFoundError:
+                pass
+            if total <= int(self._frame_cache_max_bytes * 0.85):
+                break
+
+    def _decode_rgb(self, path: str, ts: float) -> np.ndarray:
+        frame_idx = int(round(ts * self.fps))
+        mem_key = (path, frame_idx)
+        cached_img = self._get_frame_mem_cache(mem_key)
+        if cached_img is not None:
+            return cached_img
+        cache_path = self._frame_cache_path(path, frame_idx)
+        if cache_path and os.path.exists(cache_path):
+            try:
+                if self._frame_cache_touch_disk:
+                    os.utime(cache_path, None)
+                img = np.load(cache_path)
+                self._put_frame_mem_cache(mem_key, img)
+                return img
+            except Exception:
+                try:
+                    os.remove(cache_path)
+                except FileNotFoundError:
+                    pass
+
+        container = self._cached_container(path)
+        vs = container.streams.video[0]
+        container.seek(int(max(0.0, ts) / vs.time_base), stream=vs, backward=True)
+        frame = None
+        for fr in container.decode(vs):
+            if fr.time is not None and fr.time >= ts - 1e-3:
+                frame = fr
+                break
+        if frame is None:  # ts past end — take last decoded
+            container.seek(int(max(0.0, ts) / vs.time_base), stream=vs, backward=True)
+            for fr in container.decode(vs):
+                frame = fr
+        if self._decode_resize > 0:
+            frame = frame.reformat(width=self._decode_resize, height=self._decode_resize, format="rgb24")
+        img = frame.to_ndarray(format="rgb24")
+        self._put_frame_mem_cache(mem_key, img)
+        if cache_path:
+            self._prune_frame_cache()
+            tmp = f"{cache_path}.{os.getpid()}.{time.time_ns()}.tmp"
+            try:
+                with open(tmp, "wb") as f:
+                    np.save(f, img)
+                os.replace(tmp, cache_path)
+            except Exception:
+                try:
+                    os.remove(tmp)
+                except FileNotFoundError:
+                    pass
+        return img  # HWC uint8
+
+    def __getitem__(self, i):
+        ei, t = self.samples[i]
+        rec = self.episodes[ei]
+        sub = self._episode_frames(rec)
+        states = np.stack(sub["observation.state"].to_numpy())       # [L,61]
+        actions = np.stack(sub["action"].to_numpy())                 # [L,23]
+        ts = float(sub["timestamp"].iloc[t])
+        act_win = actions[t:t + self.H]                              # [H,23]
+        item = {
+            "observation.state": states[t].astype(np.float32),       # raw 61-dim
+            "action": act_win.astype(np.float32),
+            "task_index": np.int64(self.idx2026_to_2025[int(sub["task_index"].iloc[t])]),
+            "timestamp": np.float32(ts),
+            "episode_index": np.int64(rec["episode_index"]),
+            "index": np.int64(i),
+        }
+        for dst in self.RGB_KEYS:
+            frame_ts = rec["from_ts"][dst] + t / self.fps
+            item[dst] = self._decode_rgb(rec["video"][dst], frame_ts)  # HWC uint8
+        return item
+
+
+def create_v3_behavior_data_loader(config, root_2026: str, activities: List[str],
+                                   task_data_json: str, *, sharding=None,
+                                   shuffle: bool = True, num_workers: Optional[int] = None,
+                                   seed: int = 0):
+    """Build a training data loader over the 2026 v3 subset, reusing the 2025
+    transform pipeline but with 2026 state extraction (B1kInputs2026)."""
+    import jax
+    import dataclasses as _dc
+    from b1k.policies import b1k_policy
+    from b1k.training.data_loader import transform_dataset, DataLoaderImpl
+    from openpi.training.data_loader import TorchDataLoader
+
+    data_config = config.data.create(config.assets_dirs, config.model)
+    # Swap the OmniGibson-dependent B1kInputs for the 2026 61-dim variant.
+    new_inputs = tuple(
+        B1kInputs2026(model_type=config.model.model_type)
+        if isinstance(x, b1k_policy.B1kInputs) else x
+        for x in data_config.data_transforms.inputs
+    )
+    data_config = _dc.replace(
+        data_config,
+        data_transforms=_dc.replace(data_config.data_transforms, inputs=new_inputs),
+    )
+
+    ds = BehaviorV3Dataset(root_2026, activities=activities,
+                           action_horizon=config.model.action_horizon,
+                           task_data_json=task_data_json, seed=seed)
+    ds = transform_dataset(ds, data_config)   # adds dataset-aware subtask + per-ts norm
+    loader = TorchDataLoader(
+        ds,
+        local_batch_size=config.batch_size // jax.process_count(),
+        sharding=sharding, shuffle=shuffle,
+        num_workers=config.num_workers if num_workers is None else num_workers,
+        seed=seed,
+    )
+    return DataLoaderImpl(data_config, loader)
+
+
+__all__ = ["BehaviorV3Dataset", "B1kInputs2026", "extract_state_2026",
+           "build_task_index_maps", "PROP_2026", "create_v3_behavior_data_loader"]

legacy/vggt_newbank_boxing_gloves_step19999/code/src/b1k/training/b1k_da3.py

@@ -0,0 +1,321 @@
+"""DA3 spatial inputs for the 2026 v3 pipeline.
+
+Extends BehaviorV3Dataset with per-frame DA3 inputs (3 cams at DA3 resolution + robot->cam
+OpenCV extrinsics + intrinsics + ModernBERT task-language), and provides a data-loader factory
+that runs the frozen DA3-GIANT extractor once per BATCH (GPU) via the loader's batch hook.
+
+Geometry (empirically calibrated against GT depth, see /work/jack/behavior1k/calib):
+  * robot2cam_pose[7] = [pos(3), quat_wxyz(4)] = the CAMERA POSE IN THE ROBOT FRAME,
+    already OpenCV-convention (+Z optical axis). robot->cam = inv(pose_matrix).
+  * intrinsics: fx = fy = W * 17.0/20.995 (OmniGibson VisionSensor defaults), cx=cy=W/2.
+"""
+import logging
+import os
+import pickle
+
+import numpy as np
+
+from b1k.training.b1k_2026 import BehaviorV3Dataset, B1kInputs2026
+
+logger = logging.getLogger(__name__)
+
+FOCAL_RATIO = 17.0 / 20.995  # OmniGibson VisionSensor default focal/aperture
+
+# dst rgb key -> pose parquet column (same camera)
+POSE_COLS = {
+    "observation.images.rgb.head": "observation.robot2cam_pose.zed_link_camera_0",
+    "observation.images.rgb.left_wrist": "observation.robot2cam_pose.left_realsense_link_camera_0",
+    "observation.images.rgb.right_wrist": "observation.robot2cam_pose.right_realsense_link_camera_0",
+}
+# view order MUST match the bank builder: 0=main(head), 1=left, 2=right
+VIEW_ORDER = (
+    "observation.images.rgb.head",
+    "observation.images.rgb.left_wrist",
+    "observation.images.rgb.right_wrist",
+)
+
+
+def quat_wxyz_to_R(q):
+    w, x, y, z = q / (np.linalg.norm(q) + 1e-12)
+    return np.array([
+        [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
+        [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
+        [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)],
+    ])
+
+
+# OmniGibson cameras use the OpenGL optical convention (-Z forward, +Y up); the DA3 pinhole
+# projection assumes OpenCV (+Z forward, +Y down). This diag(1,-1,-1) flips the camera Y/Z axes.
+# WITHOUT it, cross-view GT-depth reprojection is 0.00 (cameras point the wrong way); WITH it, 0.20+
+# (best of all 8 conventions), and the head camera lands at its true +1.56 m height. (calib/calibrate_v3.py)
+_GL2CV = np.diag([1.0, -1.0, -1.0, 1.0]).astype(np.float32)
+
+
+def pose7_to_robot2cam(p7: np.ndarray) -> np.ndarray:
+    """[pos3, quat_wxyz] camera-pose-in-robot-frame -> 4x4 robot->cam (OpenCV).
+
+    Convention (validated in calib/): quat=wxyz, pose is the CAMERA-IN-ROBOT transform so
+    robot->cam = inv(T), then GL->CV optical flip.
+    """
+    T = np.eye(4, dtype=np.float32)
+    T[:3, :3] = quat_wxyz_to_R(np.asarray(p7[3:], np.float64))
+    T[:3, 3] = p7[:3]
+    return (_GL2CV @ np.linalg.inv(T)).astype(np.float32)
+
+
+class BehaviorV3DA3Dataset(BehaviorV3Dataset):
+    """BehaviorV3Dataset + DA3 inputs (frames @ da3_hw, extrinsics, intrinsics, task language)."""
+
+    def __init__(self, *args, da3_hw=(252, 252), lang_cache: str | None = None, lang_max_len: int = 32, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._da3_hw = tuple(da3_hw)
+        self._lang_max_len = int(lang_max_len)
+        self._lang = None
+        if lang_cache:
+            with open(lang_cache, "rb") as f:
+                self._lang = pickle.load(f)
+            logger.info("DA3 lang cache: %d tasks from %s", len(self._lang), lang_cache)
+
+    def _episode_poses(self, rec):
+        """Cached per-episode pose table (the base reader's parquet cache omits pose columns)."""
+        import pandas as pd
+        if not hasattr(self, "_pose_cache"):
+            from collections import OrderedDict
+            self._pose_cache = OrderedDict()
+        key = rec["data"]
+        if key not in self._pose_cache:
+            df = pd.read_parquet(key, columns=["episode_index", "frame_index", *POSE_COLS.values()])
+            self._pose_cache[key] = df
+            if len(self._pose_cache) > 8:
+                self._pose_cache.popitem(last=False)
+        df = self._pose_cache[key]
+        return df[df["episode_index"] == rec["episode_index"]].sort_values("frame_index")
+
+    def _decode_native(self, path: str, ts: float):
+        """Decode ONE frame at NATIVE resolution, cached by (path, frame_idx). Returns (HWC uint8, native_w).
+        Shared by the base (224) and DA3 (252) decode so each frame is decoded ONCE, not twice.
+        Enabled only when B1K_SHARED_DECODE=1 (default off => original two-decode behavior)."""
+        if not hasattr(self, "_native_cache"):
+            from collections import OrderedDict
+            self._native_cache = OrderedDict()
+        fidx = int(round(ts * self.fps))
+        key = (path, fidx)
+        hit = self._native_cache.get(key)
+        if hit is not None:
+            self._native_cache.move_to_end(key)
+            return hit
+        container = self._cached_container(path)
+        vs = container.streams.video[0]
+        container.seek(int(max(0.0, ts) / vs.time_base), stream=vs, backward=True)
+        frame = None
+        for fr in container.decode(vs):
+            if fr.time is not None and fr.time >= ts - 1e-3:
+                frame = fr
+                break
+        if frame is None:
+            container.seek(int(max(0.0, ts) / vs.time_base), stream=vs, backward=True)
+            for fr in container.decode(vs):
+                frame = fr
+        out = (frame.to_ndarray(format="rgb24"), int(frame.width))  # native HWC uint8, no resize
+        self._native_cache[key] = out
+        if len(self._native_cache) > 12:  # a few cams x a couple frames in flight
+            self._native_cache.popitem(last=False)
+        return out
+
+    def _decode_rgb(self, path: str, ts: float) -> np.ndarray:
+        """Base VLM (224) frame. With shared-decode, resize from the single native decode (no 2nd decode)."""
+        if os.environ.get("B1K_SHARED_DECODE") != "1":
+            return super()._decode_rgb(path, ts)
+        import cv2
+        native, _ = self._decode_native(path, ts)
+        r = self._decode_resize
+        return cv2.resize(native, (r, r), interpolation=cv2.INTER_AREA) if r > 0 else native
+
+    def _decode_da3(self, path: str, ts: float) -> np.ndarray:
+        """Decode one frame at DA3 resolution (252). HWC uint8 + native width (for intrinsics)."""
+        if os.environ.get("B1K_SHARED_DECODE") == "1":
+            import cv2
+            native, native_w = self._decode_native(path, ts)  # reuses the base decode (no 2nd decode)
+            h, w = self._da3_hw
+            return cv2.resize(native, (w, h), interpolation=cv2.INTER_AREA), native_w
+        container = self._cached_container(path)
+        vs = container.streams.video[0]
+        container.seek(int(max(0.0, ts) / vs.time_base), stream=vs, backward=True)
+        frame = None
+        for fr in container.decode(vs):
+            if fr.time is not None and fr.time >= ts - 1e-3:
+                frame = fr
+                break
+        if frame is None:
+            container.seek(int(max(0.0, ts) / vs.time_base), stream=vs, backward=True)
+            for fr in container.decode(vs):
+                frame = fr
+        h, w = self._da3_hw
+        native_w = frame.width
+        frame = frame.reformat(width=w, height=h, format="rgb24")
+        img = frame.to_ndarray(format="rgb24")
+        return img, native_w
+
+    def _lang_entry(self, task_name: str):
+        if self._lang is None:
+            L = self._lang_max_len
+            return np.zeros((L, 1024), np.float32), np.zeros((L,), bool)
+        feat, mask = self._lang[task_name]
+        return np.asarray(feat, np.float32), np.asarray(mask, bool)
+
+    def da3_fields(self, i):
+        """Compute ONLY the DA3 input fields for sample i (attached AFTER the transform stack,
+        which constructs fresh dicts and would drop unknown keys)."""
+        item = {}
+        ei, t = self.samples[i]
+        rec = self.episodes[ei]
+        sub = self._episode_poses(rec)
+
+        h, w = self._da3_hw
+        imgs, extr, intr = [], [], []
+        for dst in VIEW_ORDER:
+            frame_ts = rec["from_ts"][dst] + t / self.fps
+            img, native_w = self._decode_da3(rec["video"][dst], frame_ts)
+            imgs.append(img)
+            p7 = np.asarray(sub[POSE_COLS[dst]].iloc[t], np.float64)
+            extr.append(pose7_to_robot2cam(p7))
+            f_native = FOCAL_RATIO * native_w
+            # native (square) -> da3_hw rescale: fx,cx scale by w/native_w; fy,cy by h/native_h(=native_w)
+            K = np.array([
+                [f_native * w / native_w, 0, (native_w / 2) * w / native_w],
+                [0, f_native * h / native_w, (native_w / 2) * h / native_w],
+                [0, 0, 1],
+            ], np.float32)
+            intr.append(K)
+
+        item["da3_images"] = np.stack(imgs, 0)  # [V,252,252,3] uint8
+        item["camera_extrinsics"] = np.stack(extr, 0)  # [V,4,4] robot->cam OpenCV
+        item["camera_intrinsics"] = np.stack(intr, 0)  # [V,3,3] @ da3_hw
+        lf, lm = self._lang_entry(rec["task0"])
+        item["lang_feat"] = lf
+        item["lang_mask"] = lm
+        return item
+
+
+class _AttachDA3Fields:
+    """Wraps the TRANSFORMED dataset; merges the raw dataset's DA3 fields into each sample."""
+
+    def __init__(self, transformed, raw: BehaviorV3DA3Dataset):
+        self._transformed = transformed
+        self._raw = raw
+
+    def __len__(self):
+        return len(self._transformed)
+
+    def __getitem__(self, i):
+        out = dict(self._transformed[i])
+        out.update(self._raw.da3_fields(i))
+        return out
+
+
+def create_v3_behavior_da3_loader(config, root_2026, activities, task_data_json, *,
+                                  lang_cache, sharding=None, shuffle=True,
+                                  num_workers=None, seed=0, da3_hw=(252, 252)):
+    """v3 loader with DA3 inputs + a per-batch frozen DA3-GIANT extraction hook (GPU)."""
+    import jax
+    import dataclasses as _dc
+    from b1k.policies import b1k_policy
+    from b1k.training.data_loader import transform_dataset, DataLoaderImpl
+    from openpi.training.data_loader import TorchDataLoader
+    from b1k.training import da3_extractor as _ex
+
+    data_config = config.data.create(config.assets_dirs, config.model)
+    new_inputs = tuple(
+        B1kInputs2026(model_type=config.model.model_type)
+        if isinstance(x, b1k_policy.B1kInputs) else x
+        for x in data_config.data_transforms.inputs
+    )
+    data_config = _dc.replace(
+        data_config, data_transforms=_dc.replace(data_config.data_transforms, inputs=new_inputs))
+
+    # VGGT-Omega extractor: decode/patchify at process_res (patch 16), not the DA3 252 grid.
+    _use_vggt = os.environ.get("USE_VGGT") == "1"
+    if _use_vggt:
+        da3_hw = (int(os.environ.get("VGGT_PROCESS_RES", "256")),) * 2
+
+    ds = BehaviorV3DA3Dataset(
+        root_2026, activities=activities, action_horizon=config.model.action_horizon,
+        task_data_json=task_data_json, seed=seed,
+        da3_hw=da3_hw, lang_cache=lang_cache,
+        lang_max_len=config.model.da3.lang_max_len,
+    )
+    tds = transform_dataset(ds, data_config)
+    tds = _AttachDA3Fields(tds, ds)
+
+    logger.info("Building inline DA3-GIANT extractor (da3_hw=%s) ...", da3_hw)
+    # Extraction devices: default single-GPU (cuda:0). Set B1K_EXTRACT_DEVICES to spread the frozen
+    # DA3-GIANT forward across GPUs (one replica per device, batch split, run concurrently) so the
+    # ~2.7s single-GPU extraction shrinks and better overlaps the JAX train step.
+    _dev_env = os.environ.get("B1K_EXTRACT_DEVICES", "").strip()
+    _devices = [d.strip() for d in _dev_env.split(",") if d.strip()] or None
+    _fchunk = int(os.environ.get("B1K_DA3_FWD_CHUNK", "16"))
+    logger.info("DA3 extractor: devices=%s forward_chunk=%d", _devices or ["cuda:0"], _fchunk)
+    if _use_vggt:
+        from b1k.training import vggt_extractor as _vex
+        extractor = _vex.VGGTInlineExtractor(process_res=da3_hw[0], forward_chunk=_fchunk, devices=_devices)
+        logger.info("Using VGGT-Omega extractor (process_res=%d, grid=%d)", da3_hw[0], da3_hw[0] // 16)
+    else:
+        extractor = _ex.DA3InlineExtractor(da3_hw=da3_hw, forward_chunk=_fchunk, devices=_devices)
+
+    # DLPack GPU->GPU handoff: skip the ~2.6s/batch host round-trip by moving extractor features
+    # straight from the extraction GPUs to the training GPUs over NVLink. Requires CUDA extraction
+    # devices whose count matches the training mesh size (contiguous batch split aligns 1:1).
+    # Holds the last few batches' torch source shards alive so the async NVLink copies (device_put)
+    # can never read freed memory — replaces a blocking block_until_ready that serialized the producer.
+    import collections as _collections
+    _keepalive = _collections.deque(maxlen=4)
+
+    def _dlpack_ok():
+        try:
+            m = getattr(sharding, "mesh", None)
+            return (os.environ.get("B1K_DLPACK") == "1" and m is not None
+                    and len(list(m.devices.flat)) == len(extractor.devices)
+                    and all(str(d).startswith("cuda") for d in extractor.devices))
+        except Exception:
+            return False
+
+    # Output field order MUST match the extractor's extract()/extract_shards_torch() tuple order.
+    _field_names = (["da3_features", "da3_ray", "da3_depth",
+                     "da3_depth_conf", "da3_pose_enc", "da3_cam_tokens"] if _use_vggt
+                    else ["da3_features", "da3_ray", "da3_depth"])
+
+    def batch_transform(batch):
+        if _dlpack_ok():
+            import jax
+            parts = extractor.extract_shards_torch(
+                batch["da3_images"], batch["camera_extrinsics"], batch["camera_intrinsics"])
+            tdevs = list(sharding.mesh.devices.flat)  # training devices, batch-chunk k -> tdevs[k]
+
+            def _asm(fi):  # assemble per-shard torch tensors (field fi) into one sharded jax array
+                js = [jax.device_put(jax.dlpack.from_dlpack(parts[k][fi]), tdevs[k]) for k in range(len(parts))]
+                gshape = (sum(int(s.shape[0]) for s in js),) + tuple(int(d) for d in js[0].shape[1:])
+                return jax.make_array_from_single_device_arrays(gshape, sharding, js)
+
+            for fi, nm in enumerate(_field_names):
+                batch[nm] = _asm(fi)
+            # Do NOT block here: the device_put queues behind the in-flight train step on the target
+            # GPUs, so blocking would serialize the producer with training (killing the overlap).
+            # Instead keep the torch source shards referenced for a few batches so the async NVLink
+            # copy can't read freed memory.
+            _keepalive.append(parts)
+        else:
+            outs = extractor.extract(
+                batch["da3_images"], batch["camera_extrinsics"], batch["camera_intrinsics"])
+            for nm, arr in zip(_field_names, outs):
+                batch[nm] = arr  # da3_features = uint16 bf16-bits; rest fp32
+        batch.pop("da3_images", None)
+        batch.pop("camera_intrinsics", None)
+        return batch
+
+    loader = TorchDataLoader(
+        tds, local_batch_size=config.batch_size // jax.process_count(),
+        sharding=sharding, shuffle=shuffle,
+        num_workers=config.num_workers if num_workers is None else num_workers,
+        seed=seed, batch_transform=batch_transform,
+    )
+    return DataLoaderImpl(data_config, loader)

legacy/vggt_newbank_boxing_gloves_step19999/code/src/b1k/training/vggt_extractor.py

@@ -0,0 +1,208 @@
+"""Inline VGGT-Omega feature/geometry extractor (PyTorch, runs in the openpi venv).
+
+Drop-in alternative to DA3InlineExtractor: same replicated-per-GPU design and the same
+extract()/extract_shards_torch() batch-splitting, but produces VGGT-Omega's richer outputs.
+
+Per forward we harvest EVERYTHING VGGT-Omega gives us and map it into the spatial bank:
+  - feats       [B,4,V,2048,gh,gw]  4 aggregator taps (blocks 4/11/17/23), patch tokens only
+  - ray         [B,V,3,gh,gw]       ANALYTIC camera-frame unit ray dirs from the intrinsics
+                                    (VGGT has no ray output; K^-1[u,v,1] is exact, better than a
+                                    predicted ray). world lift happens in the bank via extrinsics.
+  - depth       [B,V,1,gh,gw]       VGGT dense_head depth, pooled to the patch grid
+  - depth_conf  [B,V,1,gh,gw]       VGGT dense_head confidence (per-patch geometric reliability)
+  - pose_enc    [B,V,9]             VGGT camera_head pose encoding (trans3+quat4+fov2)
+  - cam_tokens  [B,V,17,2048]       camera(1)+register(16) global tokens
+
+Field order for extract()/extract_shards_torch(): (feats, ray, depth, depth_conf, pose_enc, cam_tokens).
+feats ship as bf16 BITS (uint16) like the DA3 path; the rest are fp32.
+"""
+
+import concurrent.futures
+import contextlib
+import os
+import sys
+
+import numpy as np
+import torch
+
+try:
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+except Exception:  # noqa: BLE001
+    pass
+
+# Vendored VGGT-Omega package (uses absolute `from vggt_omega...` imports internally).
+_VENDOR = "/work/jack/da3xvla_src/DA3-XVLA/third_party"
+if _VENDOR not in sys.path:
+    sys.path.insert(0, _VENDOR)
+
+
+def _to_dev(a, device, dtype=None):
+    if isinstance(a, torch.Tensor):
+        t = a.to(device, non_blocking=True)
+        return t.to(dtype) if dtype is not None and t.dtype != dtype else t
+    return torch.as_tensor(a, device=device, dtype=dtype)
+
+
+class VGGTInlineExtractor:
+    """Frozen VGGT-Omega-1B multi-view extractor, REPLICATED one-per-GPU (mirrors DA3InlineExtractor)."""
+
+    CACHED_LAYERS = (4, 11, 17, 23)  # aggregator.cached_layer_indices; 4 taps -> num_layers=4
+    PATCH_START = 17                 # 1 camera + 16 register tokens
+    OUT_CHANNELS = 2048              # 2 * embed_dim (frame + inter-frame concat)
+
+    def __init__(
+        self,
+        model_name: str = "JackLiu0406/vggt-omega-1b",
+        ckpt_filename: str = "vggt_omega_1b_512.pt",
+        process_res: int = 256,
+        devices=None,
+        forward_chunk: int = 8,
+    ):
+        from vggt_omega.models.vggt_omega import VGGTOmega
+        from huggingface_hub import hf_hub_download
+
+        if int(process_res) % 16 != 0:
+            raise ValueError(f"process_res={process_res} must be a multiple of patch_size=16")
+        self.process_res = int(process_res)
+        self.grid = self.process_res // 16
+        if devices is None:
+            devices = ["cuda:0" if torch.cuda.is_available() else "cpu"]
+        self.devices = list(devices)
+        self.forward_chunk = int(forward_chunk)
+
+        ckpt_path = hf_hub_download(repo_id=model_name, filename=ckpt_filename,
+                                    cache_dir=os.environ.get("HF_HOME", None))
+        sd = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        if isinstance(sd, dict) and "state_dict" in sd:
+            sd = sd["state_dict"]
+
+        self.replicas = []
+        for dev in self.devices:
+            m = VGGTOmega(patch_size=16, embed_dim=1024,
+                          enable_camera=True, enable_depth=True, enable_alignment=False)
+            m.load_state_dict(sd, strict=False)  # heads present; drops text_alignment_head keys
+            m = m.to(dev).eval()
+            for p in m.parameters():
+                p.requires_grad_(False)
+            self.replicas.append(m)
+        del sd  # free the ~4GB host checkpoint copy before the DataLoader spawns its worker processes
+        self._pool = concurrent.futures.ThreadPoolExecutor(max_workers=max(1, len(self.devices)))
+        self._warmup()
+
+    def _warmup(self):
+        """Run one forward per replica SINGLE-THREADED so every lazy op wrapper is initialized before
+        the concurrent per-device extraction threads run (avoids first-call races across threads)."""
+        V = 3
+        R = self.process_res
+        img = np.zeros((1, V, R, R, 3), np.uint8)
+        ext = np.tile(np.eye(4, dtype=np.float32), (1, V, 1, 1))
+        f = 0.9 * R
+        K = np.tile(np.array([[f, 0, R / 2], [0, f, R / 2], [0, 0, 1]], np.float32), (1, V, 1, 1))
+        for di in range(len(self.devices)):
+            self._run_shard(di, img, ext, K, return_torch=False)
+
+    def _preprocess(self, images: np.ndarray, device) -> torch.Tensor:
+        """images [B,V,H,W,3] uint8/float -> [B,V,3,process_res,process_res] in [0,1] (aggregator renorms)."""
+        x = _to_dev(images, device)
+        if x.dtype == torch.uint8:
+            x = x.float() / 255.0
+        elif x.max() > 1.5:
+            x = x.float() / 255.0
+        x = x.permute(0, 1, 4, 2, 3)  # [B,V,3,H,W]
+        b, v = x.shape[:2]
+        x = torch.nn.functional.interpolate(
+            x.flatten(0, 1), size=(self.process_res, self.process_res),
+            mode="bilinear", align_corners=False,
+        ).view(b, v, 3, self.process_res, self.process_res)
+        return x.clamp_(0.0, 1.0)
+
+    def _rays_from_intrinsics(self, intrinsics: np.ndarray, src_hw, device) -> torch.Tensor:
+        """intrinsics [B,V,3,3] at src_hw -> camera-frame unit ray dirs [B,V,3,grid,grid] (row-major).
+        Analytic pinhole back-projection: dir = normalize([(x-cx)/fx, (y-cy)/fy, 1]). Deliberately NO
+        torch.linalg.inv -- its lazy wrapper races across the per-device extraction THREADS on first use
+        ('lazy wrapper should be called at most once'); the pinhole inverse is closed-form anyway."""
+        g = self.grid
+        K = _to_dev(intrinsics, device, torch.float32)  # [B,V,3,3] at src_hw
+        sh, sw = src_hw
+        rw, rh = g / sw, g / sh  # rescale K from src pixels to the g x g grid
+        fx = (K[..., 0, 0] * rw)[..., None, None]  # [B,V,1,1]
+        cx = (K[..., 0, 2] * rw)[..., None, None]
+        fy = (K[..., 1, 1] * rh)[..., None, None]
+        cy = (K[..., 1, 2] * rh)[..., None, None]
+        ys, xs = torch.meshgrid(torch.arange(g, device=device, dtype=torch.float32) + 0.5,
+                                torch.arange(g, device=device, dtype=torch.float32) + 0.5,
+                                indexing="ij")  # [g,g] row-major: rows=y=h, cols=x=w (matches 'b (h w)')
+        dx = (xs[None, None] - cx) / fx  # [B,V,g,g]
+        dy = (ys[None, None] - cy) / fy
+        dz = torch.ones_like(dx)
+        dirs = torch.stack([dx, dy, dz], dim=2)  # [B,V,3,g,g]
+        return dirs / dirs.norm(dim=2, keepdim=True).clamp_min(1e-8)
+
+    def _run_shard(self, di, images, extrinsics, intrinsics, return_torch=False):
+        dev = self.devices[di]
+        dev_idx = int(dev.split(":")[1]) if ":" in dev else None
+        m = self.replicas[di]
+        chunk = self.forward_chunk
+        g = self.grid
+        src_hw = (int(images.shape[2]), int(images.shape[3]))  # [B,V,H,W,3]
+        acc = [[] for _ in range(6)]  # feats, ray, depth, depth_conf, pose_enc, cam_tokens
+        ctx = torch.cuda.device(dev_idx) if dev_idx is not None else contextlib.nullcontext()
+        with ctx, torch.no_grad():
+            for i in range(0, images.shape[0], chunk):
+                x = self._preprocess(images[i:i + chunk], dev)  # [b,V,3,R,R] in [0,1]
+                b, V = x.shape[:2]
+                amp = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+                with torch.autocast(device_type="cuda", dtype=amp):
+                    cached, patch_start = m.aggregator(x)   # list; patch_start == 17
+                # --- patch feature taps -> [b,4,V,C,g,g] ---
+                taps = []
+                for li in self.CACHED_LAYERS:
+                    tok = cached[li]                        # [b,V,ntok,C]
+                    pt = tok[:, :, patch_start:, :]         # [b,V,g*g,C]
+                    pt = pt.reshape(b, V, g, g, self.OUT_CHANNELS).permute(0, 1, 4, 2, 3)  # [b,V,C,g,g]
+                    taps.append(pt)
+                feats = torch.stack(taps, dim=1).to(torch.bfloat16)   # [b,4,V,C,g,g]
+                cam_tokens = cached[-1][:, :, :patch_start, :].float()  # [b,V,17,C]
+                # --- heads (fp32) ---
+                with torch.autocast(device_type="cuda", enabled=False):
+                    pose_enc = m.camera_head(cached, patch_token_start=patch_start).float()  # [b,V,9]
+                    depth, depth_conf = m.dense_head(cached, images=x, patch_token_start=patch_start)
+                # depth [b,V,R,R,1], depth_conf [b,V,R,R] -> pool to grid
+                depth = depth.float().squeeze(-1)                      # [b,V,R,R]
+                depth = torch.nn.functional.adaptive_avg_pool2d(
+                    depth.reshape(b * V, 1, *depth.shape[2:]), (g, g)).reshape(b, V, 1, g, g)
+                depth_conf = torch.nn.functional.adaptive_avg_pool2d(
+                    depth_conf.float().reshape(b * V, 1, *depth_conf.shape[2:]), (g, g)).reshape(b, V, 1, g, g)
+                ray = self._rays_from_intrinsics(intrinsics[i:i + chunk], src_hw, dev)  # [b,V,3,g,g]
+                outs = (feats.view(torch.uint16), ray, depth, depth_conf, pose_enc, cam_tokens)
+                if return_torch:
+                    for j, o in enumerate(outs):
+                        acc[j].append(o)
+                else:
+                    for j, o in enumerate(outs):
+                        acc[j].append(o.cpu().numpy())
+        if return_torch:
+            return tuple(torch.cat(a, 0) for a in acc)
+        return tuple(np.concatenate(a, 0) for a in acc)
+
+    def _split_run(self, images, extrinsics, intrinsics, return_torch):
+        b = int(images.shape[0]); nd = len(self.devices)
+        bounds = [round(i * b / nd) for i in range(nd + 1)]
+        futs = {}
+        for di in range(nd):
+            s, e = bounds[di], bounds[di + 1]
+            if s >= e:
+                continue
+            futs[di] = self._pool.submit(self._run_shard, di, images[s:e], extrinsics[s:e],
+                                         intrinsics[s:e], return_torch)
+        return [futs[di].result() for di in sorted(futs)]
+
+    def extract_shards_torch(self, images, extrinsics, intrinsics):
+        """Per-shard torch GPU tuples (feats,ray,depth,depth_conf,pose_enc,cam_tokens), shard k on devices[k]."""
+        return self._split_run(images, extrinsics, intrinsics, True)
+
+    def extract(self, images, extrinsics, intrinsics):
+        """Numpy (feats,ray,depth,depth_conf,pose_enc,cam_tokens), batch split across replicas."""
+        parts = self._split_run(images, extrinsics, intrinsics, False)
+        return tuple(np.concatenate([p[j] for p in parts], axis=0) for j in range(6))

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""VGGT-Omega inference package."""
+
+from .models import VGGTOmega
+
+__version__ = "0.0.1"
+
+__all__ = ["VGGTOmega", "__version__"]

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .vggt_omega import VGGTOmega
+
+__all__ = ["VGGTOmega"]

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/aggregator.py

@@ -0,0 +1,250 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+
+from vggt_omega.models.layers import Mlp, RopePositionEmbedding, SelfAttentionBlock
+from vggt_omega.models.layers.vision_transformer import DinoVisionTransformer
+
+
+_RESNET_MEAN = [0.485, 0.456, 0.406]
+_RESNET_STD = [0.229, 0.224, 0.225]
+
+
+class Aggregator(nn.Module):
+    """Alternating-attention encoder over video frames."""
+
+    def __init__(
+        self,
+        patch_size: int = 16,
+        embed_dim: int = 1024,
+        depth: int = 24,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        num_register_tokens: int = 16,
+        register_attention_block_indices: list[int] = [2, 6, 9, 14, 20],
+        cached_layer_indices: tuple[int, ...] = (4, 11, 17, 23),
+    ) -> None:
+        super().__init__()
+
+        self.patch_embed = _build_patch_embed(patch_size=patch_size, embed_dim=embed_dim)
+        self.rope_embed = RopePositionEmbedding(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            base=100,
+            normalize_coords="max",
+            dtype=torch.float32,
+        )
+
+        self.frame_blocks = nn.ModuleList(
+            [
+                SelfAttentionBlock(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    ffn_ratio=mlp_ratio,
+                    qkv_bias=True,
+                    proj_bias=True,
+                    ffn_bias=True,
+                    ffn_layer=Mlp,
+                    init_values=1e-5,
+                    use_qk_norm=True,
+                    mask_k_bias=True,
+                )
+                for _ in range(depth)
+            ]
+        )
+        self.inter_frame_blocks = nn.ModuleList(
+            [
+                SelfAttentionBlock(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    ffn_ratio=mlp_ratio,
+                    qkv_bias=True,
+                    proj_bias=True,
+                    ffn_bias=True,
+                    ffn_layer=Mlp,
+                    init_values=1e-5,
+                    use_qk_norm=True,
+                    mask_k_bias=True,
+                )
+                for _ in range(depth)
+            ]
+        )
+
+        self.depth = depth
+        self.patch_size = patch_size
+        self.cached_layer_indices = set(cached_layer_indices)
+        self.camera_token = nn.Parameter(torch.empty(1, 2, 1, embed_dim))
+        self.register_token = nn.Parameter(torch.empty(1, 2, num_register_tokens, embed_dim))
+        self.patch_token_start = 1 + num_register_tokens
+
+        self.inter_frame_attention_types = ["global"] * depth
+        for idx in register_attention_block_indices:
+            if idx < 0 or idx >= depth:
+                raise ValueError(f"register_attention_block_indices contains invalid block index {idx}")
+            self.inter_frame_attention_types[idx] = "register"
+
+        for name, value in (("_resnet_mean", _RESNET_MEAN), ("_resnet_std", _RESNET_STD)):
+            self.register_buffer(name, torch.FloatTensor(value).view(1, 1, 3, 1, 1), persistent=False)
+
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        nn.init.normal_(self.camera_token, std=1e-3)
+        nn.init.normal_(self.register_token, std=1e-3)
+
+    def forward(
+        self,
+        images: torch.Tensor,
+    ) -> tuple[list[torch.Tensor | None], int]:
+        batch_size, num_frames, num_channels, height, width = images.shape
+        if num_channels != 3:
+            raise ValueError(f"Expected 3 input channels, got {num_channels}")
+
+        images = (images - self._resnet_mean) / self._resnet_std
+        images = images.view(batch_size * num_frames, num_channels, height, width)
+
+        camera_token = slice_expand_and_flatten(self.camera_token, batch_size, num_frames)
+        register_token = slice_expand_and_flatten(self.register_token, batch_size, num_frames)
+
+        patch_tokens = self.patch_embed(images)
+        if isinstance(patch_tokens, dict):
+            patch_tokens = patch_tokens["x_norm_patchtokens"]
+
+        tokens = torch.cat([camera_token, register_token, patch_tokens], dim=1)
+        _, num_tokens, embed_dim = tokens.shape
+
+        patch_grid_size = (height // self.patch_size, width // self.patch_size)
+        with torch.no_grad():
+            rope_sin, rope_cos = self.rope_embed(H=patch_grid_size[0], W=patch_grid_size[1])
+            frame_rope = (
+                rope_sin.to(device=patch_tokens.device, dtype=torch.float32),
+                rope_cos.to(device=patch_tokens.device, dtype=torch.float32),
+            )
+
+        outputs = []
+        for block_idx in range(self.depth):
+            tokens, frame_tokens = self._run_frame_block(
+                tokens,
+                batch_size,
+                num_frames,
+                num_tokens,
+                embed_dim,
+                block_idx,
+                frame_rope,
+            )
+            tokens = self._run_inter_frame_attention_block(
+                tokens,
+                batch_size,
+                num_frames,
+                num_tokens,
+                embed_dim,
+                block_idx,
+                self.inter_frame_attention_types[block_idx],
+            )
+            if block_idx in self.cached_layer_indices:
+                outputs.append(torch.cat([frame_tokens, tokens], dim=-1))
+            else:
+                outputs.append(None)
+
+        return outputs, self.patch_token_start
+
+    def _run_frame_block(
+        self,
+        tokens: torch.Tensor,
+        batch_size: int,
+        num_frames: int,
+        num_tokens: int,
+        embed_dim: int,
+        block_idx: int,
+        rope_sincos: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        tokens = tokens.view(batch_size * num_frames, num_tokens, embed_dim)
+        tokens = self.frame_blocks[block_idx](tokens, rope_sincos)
+        return tokens, tokens.view(batch_size, num_frames, num_tokens, embed_dim)
+
+    def _run_inter_frame_attention_block(
+        self,
+        tokens: torch.Tensor,
+        batch_size: int,
+        num_frames: int,
+        num_tokens: int,
+        embed_dim: int,
+        block_idx: int,
+        attention_type: str,
+    ) -> torch.Tensor:
+        tokens = tokens.view(batch_size, num_frames, num_tokens, embed_dim)
+
+        if attention_type == "global":
+            tokens = tokens.view(batch_size, num_frames * num_tokens, embed_dim)
+            tokens = self.inter_frame_blocks[block_idx](tokens, None)
+            return tokens.view(batch_size, num_frames, num_tokens, embed_dim)
+
+        if attention_type != "register":
+            raise ValueError(f"Unknown inter-frame attention type: {attention_type}")
+
+        patch_token_start = self.patch_token_start
+        camera_and_register_tokens = tokens[:, :, :patch_token_start].reshape(
+            batch_size,
+            num_frames * patch_token_start,
+            embed_dim,
+        )
+        patch_tokens = tokens[:, :, patch_token_start:].reshape(
+            batch_size,
+            num_frames * (num_tokens - patch_token_start),
+            embed_dim,
+        )
+
+        camera_and_register_tokens = self.inter_frame_blocks[block_idx](camera_and_register_tokens, None)
+        tokens = torch.cat([camera_and_register_tokens, patch_tokens], dim=1)
+
+        camera_and_register_tokens = tokens[:, : num_frames * patch_token_start].view(
+            batch_size,
+            num_frames,
+            patch_token_start,
+            embed_dim,
+        )
+        patch_tokens = tokens[:, num_frames * patch_token_start :].view(
+            batch_size,
+            num_frames,
+            num_tokens - patch_token_start,
+            embed_dim,
+        )
+        return torch.cat([camera_and_register_tokens, patch_tokens], dim=2)
+
+
+def _build_patch_embed(patch_size: int, embed_dim: int) -> DinoVisionTransformer:
+    model = DinoVisionTransformer(
+        img_size=224,
+        patch_size=patch_size,
+        in_chans=3,
+        pos_embed_rope_base=100,
+        pos_embed_rope_normalize_coords="max",
+        pos_embed_rope_dtype="fp32",
+        embed_dim=embed_dim,
+        depth=24,
+        num_heads=16,
+        ffn_ratio=4,
+        qkv_bias=True,
+        drop_path_rate=0.0,
+        layerscale_init=1.0e-5,
+        norm_layer="layernormbf16",
+        ffn_layer="mlp",
+        ffn_bias=True,
+        proj_bias=True,
+        n_storage_tokens=4,
+        mask_k_bias=True,
+    )
+    model.init_weights()
+    return model
+
+
+def slice_expand_and_flatten(token_tensor: torch.Tensor, batch_size: int, num_frames: int) -> torch.Tensor:
+    first_frame_token = token_tensor[:, 0:1].expand(batch_size, 1, *token_tensor.shape[2:])
+    other_frame_tokens = token_tensor[:, 1:].expand(batch_size, num_frames - 1, *token_tensor.shape[2:])
+    tokens = torch.cat([first_frame_token, other_frame_tokens], dim=1)
+    return tokens.view(batch_size * num_frames, *tokens.shape[2:])

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/heads/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .camera_head import CameraHead
+from .dense_head import DenseHead
+from .text_alignment_head import TextAlignmentHead
+
+__all__ = ["CameraHead", "DenseHead", "TextAlignmentHead"]

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/heads/camera_head.py

@@ -0,0 +1,80 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vggt_omega.models.layers import SelfAttentionBlock
+
+
+class CameraHead(nn.Module):
+    """Camera head used by the released VGGT-Omega checkpoints."""
+
+    def __init__(self, dim_in: int = 2048) -> None:
+        super().__init__()
+
+        self.token_norm = nn.LayerNorm(dim_in, eps=1e-5)
+        # Head-local transformer blocks that mix camera and register tokens across frames.
+        self.trunk = nn.ModuleList(
+            [
+                SelfAttentionBlock(
+                    dim=dim_in,
+                    num_heads=16,
+                    ffn_ratio=4.0,
+                    qkv_bias=True,
+                    proj_bias=True,
+                    ffn_bias=True,
+                    init_values=1e-5,
+                    use_qk_norm=False,
+                    mask_k_bias=True,
+                )
+                for _ in range(4)
+            ]
+        )
+        self.trunk_norm = nn.LayerNorm(dim_in, eps=1e-5)
+        self.camera_branch = nn.Sequential(
+            nn.Linear(dim_in, dim_in // 2, bias=True),
+            nn.GELU(),
+            nn.Linear(dim_in // 2, 9, bias=True),
+        )
+
+    def forward(
+        self,
+        aggregated_tokens_list: list[torch.Tensor | None],
+        patch_token_start: int,
+    ) -> torch.Tensor:
+        tokens = aggregated_tokens_list[-1]
+        if tokens is None:
+            raise ValueError("Aggregator did not cache the final layer, which CameraHead needs.")
+        batch_size, num_frames, num_tokens, _ = tokens.shape
+
+        if patch_token_start is None:
+            raise ValueError("patch_token_start is required for CameraHead")
+        if patch_token_start > num_tokens:
+            raise ValueError(f"patch_token_start ({patch_token_start}) exceeds token length ({num_tokens})")
+
+        if tokens.dtype != torch.float32:
+            tokens = tokens.float()
+
+        camera_and_register_tokens = tokens[:, :, :patch_token_start]
+        camera_and_register_tokens = self.token_norm(camera_and_register_tokens)
+
+        camera_and_register_tokens = camera_and_register_tokens.reshape(batch_size, num_frames * patch_token_start, -1)
+        rope_sincos = None
+        for block in self.trunk:
+            camera_and_register_tokens = block(camera_and_register_tokens, rope_sincos)
+
+        camera_and_register_tokens = camera_and_register_tokens.reshape(batch_size, num_frames, patch_token_start, -1)
+        camera_tokens = self.trunk_norm(camera_and_register_tokens[:, :, 0])
+        return _apply_camera_activation(self.camera_branch(camera_tokens))
+
+
+def _apply_camera_activation(raw_camera: torch.Tensor) -> torch.Tensor:
+    translation = raw_camera[..., :3]
+    quaternion = raw_camera[..., 3:7]
+    fov = F.relu(raw_camera[..., 7:]) + 0.01
+    return torch.cat([translation, quaternion, fov], dim=-1)

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/heads/dense_head.py

@@ -0,0 +1,308 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Inspired by https://github.com/DepthAnything/Depth-Anything-V2
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import create_uv_grid, position_grid_to_embed
+
+
+class DenseHead(nn.Module):
+    """Dense prediction head used by the released VGGT-Omega checkpoints."""
+
+    def __init__(
+        self,
+        dim_in: int = 2048,
+        patch_size: int = 16,
+        features: int = 256,
+        out_channels: list[int] = [256, 512, 1024, 1024],
+        intermediate_layer_idx: list[int] = [4, 11, 17, 23],
+    ) -> None:
+        super().__init__()
+
+        if patch_size % 4 != 0:
+            raise ValueError(
+                "DenseHead expects patch_size divisible by 4 because the fused feature is decoded "
+                f"from 1/4 scale. Got patch_size={patch_size}."
+            )
+
+        self.patch_size = patch_size
+        self.intermediate_layer_idx = intermediate_layer_idx
+        self.final_shuffle_factor = patch_size // 4
+        self.norm = nn.LayerNorm(dim_in, eps=1e-5)
+
+        self.projects = nn.ModuleList(
+            [nn.Conv2d(in_channels=dim_in, out_channels=oc, kernel_size=1, stride=1, padding=0) for oc in out_channels]
+        )
+        self.resize_layers = nn.ModuleList(
+            [
+                _make_dense_resize_layer(channels=out_channels[0], resize_scale=4.0),
+                _make_dense_resize_layer(channels=out_channels[1], resize_scale=2.0),
+                _make_dense_resize_layer(channels=out_channels[2], resize_scale=1.0),
+                _make_dense_resize_layer(channels=out_channels[3], resize_scale=0.5),
+            ]
+        )
+
+        self.scratch = _make_scratch(out_channels, features)
+        self.scratch.stem_transpose = None
+        self.scratch.refinenet1 = _make_fusion_block(features)
+        self.scratch.refinenet2 = _make_fusion_block(features)
+        self.scratch.refinenet3 = _make_fusion_block(features)
+        self.scratch.refinenet4 = _make_fusion_block(features, has_residual=False)
+
+        self.proj = _make_prediction_head(
+            features,
+            self.final_shuffle_factor**2,
+        )
+        self.proj_conf = _make_prediction_head(
+            features,
+            self.final_shuffle_factor**2,
+        )
+        _init_small_conf_prediction_head(self.proj_conf)
+
+    def forward(
+        self,
+        aggregated_tokens_list: list[torch.Tensor | None],
+        images: torch.Tensor,
+        patch_token_start: int,
+        frames_chunk_size: int | None = 8,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if patch_token_start is None:
+            raise ValueError("patch_token_start is required for DenseHead")
+
+        _, num_frames, _, _, _ = images.shape
+
+        if frames_chunk_size is None or frames_chunk_size >= num_frames:
+            return self._forward_impl(aggregated_tokens_list, images, patch_token_start)
+
+        assert frames_chunk_size > 0
+
+        depth_chunks = []
+        depth_conf_chunks = []
+        for frames_start_idx in range(0, num_frames, frames_chunk_size):
+            frames_end_idx = min(frames_start_idx + frames_chunk_size, num_frames)
+            depth_chunk, depth_conf_chunk = self._forward_impl(
+                aggregated_tokens_list,
+                images,
+                patch_token_start,
+                frames_start_idx,
+                frames_end_idx,
+            )
+            depth_chunks.append(depth_chunk)
+            depth_conf_chunks.append(depth_conf_chunk)
+
+        return torch.cat(depth_chunks, dim=1), torch.cat(depth_conf_chunks, dim=1)
+
+    def _forward_impl(
+        self,
+        aggregated_tokens_list: list[torch.Tensor | None],
+        images: torch.Tensor,
+        patch_token_start: int,
+        frames_start_idx: int | None = None,
+        frames_end_idx: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if frames_start_idx is not None and frames_end_idx is not None:
+            images = images[:, frames_start_idx:frames_end_idx].contiguous()
+
+        batch_size, num_frames, _, height, width = images.shape
+        patch_h, patch_w = height // self.patch_size, width // self.patch_size
+
+        multi_scale_features = []
+        for feature_idx, layer_idx in enumerate(self.intermediate_layer_idx):
+            x = aggregated_tokens_list[layer_idx]
+            if x is None:
+                raise ValueError(f"Aggregator did not cache layer {layer_idx}, which DenseHead needs.")
+            x = x[:, :, patch_token_start:]
+            if frames_start_idx is not None and frames_end_idx is not None:
+                x = x[:, frames_start_idx:frames_end_idx]
+            if x.dtype != torch.float32:
+                x = x.float()
+
+            x = x.reshape(batch_size * num_frames, -1, x.shape[-1])
+            x = self.norm(x)
+            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
+            x = self.projects[feature_idx](x)
+            x = self._apply_pos_embed(x, width, height)
+            x = self.resize_layers[feature_idx](x)
+            multi_scale_features.append(x)
+
+        fused = self.scratch_forward(multi_scale_features)
+        fused = self._apply_pos_embed(fused, width, height)
+
+        depth_logits = self.proj(fused)
+        depth_logits = F.pixel_shuffle(depth_logits, self.final_shuffle_factor)
+        depth_logits = depth_logits.permute(0, 2, 3, 1)
+
+        confidence_logits = self.proj_conf(fused)
+        confidence_logits = F.pixel_shuffle(confidence_logits, self.final_shuffle_factor)
+        confidence_logits = confidence_logits.permute(0, 2, 3, 1).squeeze(-1)
+
+        depth = torch.exp(depth_logits)
+        depth_conf = 1.0 + torch.exp(confidence_logits)
+
+        depth = depth.view(batch_size, num_frames, *depth.shape[1:])
+        depth_conf = depth_conf.view(batch_size, num_frames, *depth_conf.shape[1:])
+
+        if depth.dtype != torch.float32 or depth_conf.dtype != torch.float32:
+            raise TypeError(f"DenseHead outputs must be fp32, got depth={depth.dtype}, conf={depth_conf.dtype}")
+
+        return depth, depth_conf
+
+    def _apply_pos_embed(self, x: torch.Tensor, width: int, height: int, ratio: float = 0.1) -> torch.Tensor:
+        patch_w = x.shape[-1]
+        patch_h = x.shape[-2]
+        pos_embed = create_uv_grid(patch_w, patch_h, aspect_ratio=width / height, dtype=x.dtype, device=x.device)
+        pos_embed = position_grid_to_embed(pos_embed, x.shape[1])
+        pos_embed = pos_embed * ratio
+        pos_embed = pos_embed.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1)
+        return x + pos_embed
+
+    def scratch_forward(self, features: list[torch.Tensor]) -> torch.Tensor:
+        layer_1, layer_2, layer_3, layer_4 = features
+
+        layer_1_rn = self.scratch.layer1_rn(layer_1)
+        layer_2_rn = self.scratch.layer2_rn(layer_2)
+        layer_3_rn = self.scratch.layer3_rn(layer_3)
+        layer_4_rn = self.scratch.layer4_rn(layer_4)
+
+        out = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
+        out = self.scratch.refinenet3(out, layer_3_rn, size=layer_2_rn.shape[2:])
+        out = self.scratch.refinenet2(out, layer_2_rn, size=layer_1_rn.shape[2:])
+        return self.scratch.refinenet1(out, layer_1_rn, size=layer_1_rn.shape[2:])
+
+
+def _make_dense_resize_layer(channels: int, resize_scale: float) -> nn.Module:
+    if resize_scale == 1.0:
+        return nn.Identity()
+
+    if resize_scale == 0.5:
+        return nn.Conv2d(
+            in_channels=channels,
+            out_channels=channels,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        )
+
+    upsample_scale = int(resize_scale)
+    return nn.ConvTranspose2d(
+        in_channels=channels,
+        out_channels=channels,
+        kernel_size=upsample_scale,
+        stride=upsample_scale,
+        padding=0,
+    )
+
+
+def _make_prediction_head(in_channels: int, out_channels: int) -> nn.Module:
+    return nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=True)
+
+
+def _init_small_conf_prediction_head(proj: nn.Module) -> None:
+    if not isinstance(proj, nn.Conv2d):
+        raise TypeError(f"Unsupported confidence projection layer: {type(proj)}")
+
+    nn.init.zeros_(proj.weight)
+    if proj.bias is None:
+        raise ValueError("Small confidence init requires a bias term for proj_conf")
+
+    # With expp1 confidence activation this starts from conf ~= 1.05.
+    nn.init.constant_(proj.bias, math.log(1.05 - 1.0))
+
+
+def _make_fusion_block(features: int, has_residual: bool = True) -> nn.Module:
+    return FeatureFusionBlock(
+        features,
+        nn.ReLU(inplace=False),
+        has_residual=has_residual,
+    )
+
+
+def _make_scratch(in_shape: list[int], out_shape: int) -> nn.Module:
+    scratch = nn.Module()
+    scratch.layer1_rn = nn.Conv2d(in_shape[0], out_shape, kernel_size=3, stride=1, padding=1, bias=False)
+    scratch.layer2_rn = nn.Conv2d(in_shape[1], out_shape, kernel_size=3, stride=1, padding=1, bias=False)
+    scratch.layer3_rn = nn.Conv2d(in_shape[2], out_shape, kernel_size=3, stride=1, padding=1, bias=False)
+    scratch.layer4_rn = nn.Conv2d(in_shape[3], out_shape, kernel_size=3, stride=1, padding=1, bias=False)
+    return scratch
+
+
+class ResidualConvUnit(nn.Module):
+    def __init__(self, features: int, activation: nn.Module) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True)
+        self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True)
+        self.activation = activation
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.activation(x)
+        out = self.conv1(out)
+        out = self.activation(out)
+        out = self.conv2(out)
+        return out + x
+
+
+class FeatureFusionBlock(nn.Module):
+    def __init__(self, features: int, activation: nn.Module, has_residual: bool = True) -> None:
+        super().__init__()
+        self.out_conv = nn.Conv2d(features, features, kernel_size=1, stride=1, padding=0, bias=True)
+        self.has_residual = has_residual
+        if has_residual:
+            self.resConfUnit1 = ResidualConvUnit(features, activation)
+        self.resConfUnit2 = ResidualConvUnit(features, activation)
+
+    def forward(self, x: torch.Tensor, residual: torch.Tensor | None = None, size: Tuple[int, int] | None = None) -> torch.Tensor:
+        output = x
+        if self.has_residual:
+            if residual is None:
+                raise ValueError("FeatureFusionBlock requires a residual tensor when has_residual=True")
+            output = output + self.resConfUnit1(residual)
+
+        output = self.resConfUnit2(output)
+        output = custom_interpolate(output, size=size, mode="bilinear", align_corners=True)
+        return self.out_conv(output)
+
+
+def custom_interpolate(
+    x: torch.Tensor,
+    size: Tuple[int, int] | None = None,
+    scale_factor: float | None = None,
+    mode: str = "bilinear",
+    align_corners: bool = True,
+) -> torch.Tensor:
+    if size is None:
+        if scale_factor is None:
+            raise ValueError("custom_interpolate requires either size or scale_factor")
+        size = (
+            int(x.shape[-2] * scale_factor),
+            int(x.shape[-1] * scale_factor),
+        )
+
+    if tuple(x.shape[-2:]) == tuple(size):
+        return x
+
+    int_max = 1610612736
+    input_elements = size[0] * size[1] * x.shape[0] * x.shape[1]
+    if input_elements <= int_max:
+        return F.interpolate(x, size=size, mode=mode, align_corners=align_corners)
+
+    chunks = torch.chunk(x, chunks=(input_elements // int_max) + 1, dim=0)
+    interpolated_chunks = [
+        F.interpolate(
+            chunk,
+            size=size,
+            mode=mode,
+            align_corners=align_corners,
+        )
+        for chunk in chunks
+    ]
+    return torch.cat(interpolated_chunks, dim=0).contiguous()

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/heads/text_alignment_head.py

@@ -0,0 +1,79 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vggt_omega.models.layers import SelfAttentionBlock
+
+
+class TextAlignmentHead(nn.Module):
+    """Read out a language-aligned sequence embedding from camera/register tokens."""
+
+    def __init__(self, dim_in: int = 2048) -> None:
+        super().__init__()
+        self.token_norm = nn.LayerNorm(dim_in, eps=1e-5)
+
+        self.language_token = nn.Parameter(torch.zeros(1, 1, dim_in))
+        nn.init.trunc_normal_(self.language_token, std=0.02)
+
+        self.readout_blocks = nn.ModuleList(
+            [
+                SelfAttentionBlock(
+                    dim=dim_in,
+                    num_heads=16,
+                    ffn_ratio=4.0,
+                    qkv_bias=True,
+                    proj_bias=True,
+                    ffn_bias=True,
+                    init_values=1e-5,
+                    use_qk_norm=False,
+                    mask_k_bias=True,
+                )
+                for _ in range(4)
+            ]
+        )
+        self.language_token_norm = nn.LayerNorm(dim_in, eps=1e-5)
+        self.embedding_projector = nn.Sequential(
+            nn.Linear(dim_in, dim_in // 2, bias=True),
+            nn.GELU(),
+            nn.LayerNorm(dim_in // 2, eps=1e-5),
+            nn.Linear(dim_in // 2, dim_in, bias=True),
+        )
+
+    def forward(
+        self,
+        aggregated_tokens_list: list[torch.Tensor | None],
+        patch_token_start: int,
+    ) -> dict[str, torch.Tensor]:
+        tokens = aggregated_tokens_list[-1]
+        if tokens is None:
+            raise ValueError("Aggregator did not cache the final layer, which TextAlignmentHead needs.")
+        if patch_token_start is None:
+            raise ValueError("patch_token_start is required for TextAlignmentHead")
+        if patch_token_start > tokens.shape[2]:
+            raise ValueError(f"patch_token_start ({patch_token_start}) exceeds token length ({tokens.shape[2]})")
+
+        if tokens.dtype != torch.float32:
+            tokens = tokens.float()
+
+        batch_size, num_frames, _, _ = tokens.shape
+        camera_and_register_tokens = tokens[:, :, :patch_token_start]
+        camera_and_register_tokens = self.token_norm(camera_and_register_tokens)
+        camera_and_register_tokens = camera_and_register_tokens.reshape(batch_size, num_frames * patch_token_start, -1)
+
+        language_token = self.language_token.expand(batch_size, -1, -1)
+        readout_tokens = torch.cat([language_token, camera_and_register_tokens], dim=1)
+        for block in self.readout_blocks:
+            readout_tokens = block(readout_tokens, None)
+
+        language_token = self.language_token_norm(readout_tokens[:, 0])
+        text_alignment_embedding = self.embedding_projector(language_token)
+        return {
+            "text_alignment_embedding": F.normalize(text_alignment_embedding, dim=-1),
+            "text_alignment_token": language_token,
+        }

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/heads/utils.py

@@ -0,0 +1,108 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+
+def position_grid_to_embed(pos_grid: torch.Tensor, embed_dim: int, omega_0: float = 100) -> torch.Tensor:
+    """
+    Convert 2D position grid (HxWx2) to sinusoidal embeddings (HxWxC)
+
+    Args:
+        pos_grid: Tensor of shape (H, W, 2) containing 2D coordinates
+        embed_dim: Output channel dimension for embeddings
+
+    Returns:
+        Tensor of shape (H, W, embed_dim) with positional embeddings
+    """
+    H, W, grid_dim = pos_grid.shape
+    assert grid_dim == 2
+    pos_flat = pos_grid.reshape(-1, grid_dim)  # Flatten to (H*W, 2)
+
+    # Process x and y coordinates separately
+    emb_x = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 0], omega_0=omega_0)  # [1, H*W, D/2]
+    emb_y = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 1], omega_0=omega_0)  # [1, H*W, D/2]
+
+    # Combine and reshape
+    emb = torch.cat([emb_x, emb_y], dim=-1)  # [1, H*W, D]
+
+    return emb.view(H, W, embed_dim)  # [H, W, D]
+
+
+def make_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, omega_0: float = 100) -> torch.Tensor:
+    """
+    This function generates a 1D positional embedding from a given grid using sine and cosine functions.
+
+    Args:
+    - embed_dim: The embedding dimension.
+    - pos: The position to generate the embedding from.
+
+    Returns:
+    - emb: The generated 1D positional embedding.
+    """
+    assert embed_dim % 2 == 0
+    device = pos.device
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32 if device.type == "mps" else torch.double, device=device)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / omega_0**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = torch.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = torch.sin(out)  # (M, D/2)
+    emb_cos = torch.cos(out)  # (M, D/2)
+
+    emb = torch.cat([emb_sin, emb_cos], dim=1)  # (M, D)
+    return emb.float()
+
+
+# Inspired by https://github.com/microsoft/moge
+
+
+def create_uv_grid(
+    width: int, height: int, aspect_ratio: float = None, dtype: torch.dtype = None, device: torch.device = None
+) -> torch.Tensor:
+    """
+    Create a normalized UV grid of shape (width, height, 2).
+
+    The grid spans horizontally and vertically according to an aspect ratio,
+    ensuring the top-left corner is at (-x_span, -y_span) and the bottom-right
+    corner is at (x_span, y_span), normalized by the diagonal of the plane.
+
+    Args:
+        width (int): Number of points horizontally.
+        height (int): Number of points vertically.
+        aspect_ratio (float, optional): Width-to-height ratio. Defaults to width/height.
+        dtype (torch.dtype, optional): Data type of the resulting tensor.
+        device (torch.device, optional): Device on which the tensor is created.
+
+    Returns:
+        torch.Tensor: A (width, height, 2) tensor of UV coordinates.
+    """
+    # Derive aspect ratio if not explicitly provided
+    if aspect_ratio is None:
+        aspect_ratio = float(width) / float(height)
+
+    # Compute normalized spans for X and Y
+    diag_factor = (aspect_ratio**2 + 1.0) ** 0.5
+    span_x = aspect_ratio / diag_factor
+    span_y = 1.0 / diag_factor
+
+    # Establish the linspace boundaries
+    left_x = -span_x * (width - 1) / width
+    right_x = span_x * (width - 1) / width
+    top_y = -span_y * (height - 1) / height
+    bottom_y = span_y * (height - 1) / height
+
+    # Generate 1D coordinates
+    x_coords = torch.linspace(left_x, right_x, steps=width, dtype=dtype, device=device)
+    y_coords = torch.linspace(top_y, bottom_y, steps=height, dtype=dtype, device=device)
+
+    # Create 2D meshgrid (width x height) and stack into UV
+    uu, vv = torch.meshgrid(x_coords, y_coords, indexing="xy")
+    uv_grid = torch.stack((uu, vv), dim=-1)
+
+    return uv_grid

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .attention import CausalSelfAttention, LinearKMaskedBias, SelfAttention
+from .block import CausalSelfAttentionBlock, SelfAttentionBlock
+from .ffn_layers import Mlp, SwiGLUFFN
+from .layer_scale import LayerScale
+from .patch_embed import PatchEmbed
+from .rms_norm import RMSNorm
+from .rope_position_encoding import RopePositionEmbedding
+
+__all__ = [
+    "CausalSelfAttention",
+    "CausalSelfAttentionBlock",
+    "LayerScale",
+    "LinearKMaskedBias",
+    "Mlp",
+    "PatchEmbed",
+    "RMSNorm",
+    "RopePositionEmbedding",
+    "SelfAttention",
+    "SelfAttentionBlock",
+    "SwiGLUFFN",
+]

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/attention.py

@@ -0,0 +1,184 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import math
+from typing import List, Tuple
+
+from torch import Tensor, nn
+import torch
+import torch.nn.functional as F
+
+from .utils import cat_keep_shapes, uncat_with_shapes
+
+
+# RoPE-related functions:
+def rope_rotate_half(x: Tensor) -> Tensor:
+    # x:   [ x0  x1  x2  x3  x4  x5]
+    # out: [-x3 -x4 -x5  x0  x1  x2]
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def rope_apply(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
+    # x:   [..., D], eg [x0,     x1,   x2,   x3,   x4,   x5]
+    # sin: [..., D], eg [sin0, sin1, sin2, sin0, sin1, sin2]
+    # cos: [..., D], eg [cos0, cos1, cos2, cos0, cos1, cos2]
+    return (x * cos) + (rope_rotate_half(x) * sin)
+
+
+class LinearKMaskedBias(nn.Linear):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        o = self.out_features
+        assert o % 3 == 0
+        if self.bias is not None:
+            self.register_buffer("bias_mask", torch.full_like(self.bias, fill_value=math.nan))
+
+    def forward(self, input: Tensor) -> Tensor:
+        masked_bias = self.bias * self.bias_mask.to(self.bias.dtype) if self.bias is not None else None
+        return F.linear(input, self.weight, masked_bias)
+
+
+class SelfAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        mask_k_bias: bool = False,
+        use_qk_norm: bool = False,
+        device=None,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+        # VGGT-Omega change: the aggregator checkpoint was trained with Q/K
+        # normalization, while upstream DINOv3 attention does not expose it.
+        self.use_qk_norm = use_qk_norm
+        if self.use_qk_norm:
+            self.q_norm = nn.LayerNorm(head_dim, eps=1e-5)
+            self.k_norm = nn.LayerNorm(head_dim, eps=1e-5)
+        else:
+            self.q_norm = None
+            self.k_norm = None
+
+        linear_class = LinearKMaskedBias if mask_k_bias else nn.Linear
+        self.qkv = linear_class(dim, dim * 3, bias=qkv_bias, device=device)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias, device=device)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def apply_rope(self, q: Tensor, k: Tensor, rope: Tensor | Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
+        # All operations will use the dtype of rope, the output is cast back to the dtype of q and k
+        q_dtype = q.dtype
+        k_dtype = k.dtype
+        sin, cos = rope
+        rope_dtype = sin.dtype
+        q = q.to(dtype=rope_dtype)
+        k = k.to(dtype=rope_dtype)
+        N = q.shape[-2]
+        prefix = N - sin.shape[-2]
+        assert prefix >= 0
+        q_prefix = q[:, :, :prefix, :]
+        q = rope_apply(q[:, :, prefix:, :], sin, cos)  # [B, head, hw, D//head]
+        q = torch.cat((q_prefix, q), dim=-2)  # [B, head, N, D//head]
+        k_prefix = k[:, :, :prefix, :]
+        k = rope_apply(k[:, :, prefix:, :], sin, cos)  # [B, head, hw, D//head]
+        k = torch.cat((k_prefix, k), dim=-2)  # [B, head, N, D//head]
+        q = q.to(dtype=q_dtype)
+        k = k.to(dtype=k_dtype)
+        return q, k
+
+    def forward(self, x: Tensor, attn_bias=None, rope: Tensor = None) -> Tensor:
+        qkv = self.qkv(x)
+        attn_v = self.compute_attention(qkv=qkv, attn_bias=attn_bias, rope=rope)
+        x = self.proj(attn_v)
+        x = self.proj_drop(x)
+        return x
+
+    def forward_list(self, x_list, attn_bias=None, rope_list=None) -> List[Tensor]:
+        assert len(x_list) == len(rope_list)  # should be enforced by the Block
+        x_flat, shapes, num_tokens = cat_keep_shapes(x_list)
+        qkv_flat = self.qkv(x_flat)
+        qkv_list = uncat_with_shapes(qkv_flat, shapes, num_tokens)
+        att_out = []
+        for _, (qkv, _, rope) in enumerate(zip(qkv_list, shapes, rope_list)):
+            att_out.append(self.compute_attention(qkv, attn_bias=attn_bias, rope=rope))
+        x_flat, shapes, num_tokens = cat_keep_shapes(att_out)
+        x_flat = self.proj(x_flat)
+        return uncat_with_shapes(x_flat, shapes, num_tokens)
+
+    def compute_attention(self, qkv: Tensor, attn_bias=None, rope=None) -> Tensor:
+        assert attn_bias is None
+        B, N, _ = qkv.shape
+        C = self.qkv.in_features
+
+        qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        q, k, v = torch.unbind(qkv, 2)
+        q, k, v = [t.transpose(1, 2) for t in [q, k, v]]
+        if self.use_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        if rope is not None:
+            q, k = self.apply_rope(q, k, rope)
+        x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+        x = x.transpose(1, 2)
+        return x.reshape([B, N, C])
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = attn_drop
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def init_weights(
+        self, init_attn_std: float | None = None, init_proj_std: float | None = None, factor: float = 1.0
+    ) -> None:
+        init_attn_std = init_attn_std or (self.dim**-0.5)
+        init_proj_std = init_proj_std or init_attn_std * factor
+        nn.init.normal_(self.qkv.weight, std=init_attn_std)
+        nn.init.normal_(self.proj.weight, std=init_proj_std)
+        if self.qkv.bias is not None:
+            nn.init.zeros_(self.qkv.bias)
+        if self.proj.bias is not None:
+            nn.init.zeros_(self.proj.bias)
+
+    def forward(self, x: Tensor, is_causal: bool = True) -> Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        q, k, v = torch.unbind(qkv, 2)
+        q, k, v = [t.transpose(1, 2) for t in [q, k, v]]
+        x = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=None, dropout_p=self.attn_drop if self.training else 0, is_causal=is_causal
+        )
+        x = x.transpose(1, 2).contiguous().view(B, N, C)
+        x = self.proj_drop(self.proj(x))
+        return x

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/block.py

@@ -0,0 +1,274 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from typing import Callable, List, Optional
+
+import torch
+from torch import Tensor, nn
+
+from .attention import CausalSelfAttention, SelfAttention
+from .ffn_layers import Mlp
+from .layer_scale import LayerScale  # , DropPath
+from .utils import cat_keep_shapes, uncat_with_shapes
+
+class SelfAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        ffn_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = SelfAttention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+        mask_k_bias: bool = False,
+        use_qk_norm: bool = False,
+        device=None,
+    ) -> None:
+        super().__init__()
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+            mask_k_bias=mask_k_bias,
+            # VGGT-Omega change: pass through Q/K normalization for the
+            # aggregator blocks trained with q_norm/k_norm parameters.
+            use_qk_norm=use_qk_norm,
+            device=device,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * ffn_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+            device=device,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    @staticmethod
+    def _maybe_index_rope(rope: tuple[Tensor, Tensor] | None, indices: Tensor) -> tuple[Tensor, Tensor] | None:
+        if rope is None:
+            return None
+
+        sin, cos = rope
+        assert sin.ndim == cos.ndim
+        if sin.ndim == 4:
+            # If the rope embedding has a batch dimension (is different for each batch element), index into it
+            return sin[indices], cos[indices]  # [batch, heads, patches, embed_dim]
+        else:
+            # No batch dimension, do not index
+            return sin, cos  # [heads, patches, embed_dim] or [patches, embed_dim]
+
+    def _forward(self, x: Tensor, rope=None) -> Tensor:
+        """
+        This is the reference implementation for a single tensor, matching what is done below for a list.
+        We call the list op on [x] instead of this function.
+        """
+        b, _, _ = x.shape
+        sample_subset_size = max(int(b * (1 - self.sample_drop_ratio)), 1)
+        residual_scale_factor = b / sample_subset_size
+
+        if self.training and self.sample_drop_ratio > 0.0:
+            indices_1 = (torch.randperm(b, device=x.device))[:sample_subset_size]
+
+            x_subset_1 = x[indices_1]
+            rope_subset = self._maybe_index_rope(rope, indices_1)
+            residual_1 = self.attn(self.norm1(x_subset_1), rope=rope_subset)
+
+            x_attn = torch.index_add(
+                x,
+                dim=0,
+                source=self.ls1(residual_1),
+                index=indices_1,
+                alpha=residual_scale_factor,
+            )
+
+            indices_2 = (torch.randperm(b, device=x.device))[:sample_subset_size]
+
+            x_subset_2 = x_attn[indices_2]
+            residual_2 = self.mlp(self.norm2(x_subset_2))
+
+            x_ffn = torch.index_add(
+                x_attn,
+                dim=0,
+                source=self.ls2(residual_2),
+                index=indices_2,
+                alpha=residual_scale_factor,
+            )
+        else:
+            x_attn = x + self.ls1(self.attn(self.norm1(x), rope=rope))
+            x_ffn = x_attn + self.ls2(self.mlp(self.norm2(x_attn)))
+
+        return x_ffn
+
+    def _forward_list(self, x_list: List[Tensor], rope_list=None) -> List[Tensor]:
+        """
+        This list operator concatenates the tokens from the list of inputs together to save
+        on the elementwise operations. Torch-compile memory-planning allows hiding the overhead
+        related to concat ops.
+        """
+        b_list = [x.shape[0] for x in x_list]
+        sample_subset_sizes = [max(int(b * (1 - self.sample_drop_ratio)), 1) for b in b_list]
+        residual_scale_factors = [b / sample_subset_size for b, sample_subset_size in zip(b_list, sample_subset_sizes)]
+
+        if self.training and self.sample_drop_ratio > 0.0:
+            indices_1_list = [
+                (torch.randperm(b, device=x.device))[:sample_subset_size]
+                for x, b, sample_subset_size in zip(x_list, b_list, sample_subset_sizes)
+            ]
+            x_subset_1_list = [x[indices_1] for x, indices_1 in zip(x_list, indices_1_list)]
+
+            if rope_list is not None:
+                rope_subset_list = [
+                    self._maybe_index_rope(rope, indices_1) for rope, indices_1 in zip(rope_list, indices_1_list)
+                ]
+            else:
+                rope_subset_list = rope_list
+
+            flattened, shapes, num_tokens = cat_keep_shapes(x_subset_1_list)
+            norm1 = uncat_with_shapes(self.norm1(flattened), shapes, num_tokens)
+            residual_1_list = self.attn.forward_list(norm1, rope_list=rope_subset_list)
+
+            x_attn_list = [
+                torch.index_add(
+                    x,
+                    dim=0,
+                    source=self.ls1(residual_1),
+                    index=indices_1,
+                    alpha=residual_scale_factor,
+                )
+                for x, residual_1, indices_1, residual_scale_factor in zip(
+                    x_list, residual_1_list, indices_1_list, residual_scale_factors
+                )
+            ]
+
+            indices_2_list = [
+                (torch.randperm(b, device=x.device))[:sample_subset_size]
+                for x, b, sample_subset_size in zip(x_list, b_list, sample_subset_sizes)
+            ]
+            x_subset_2_list = [x[indices_2] for x, indices_2 in zip(x_attn_list, indices_2_list)]
+            flattened, shapes, num_tokens = cat_keep_shapes(x_subset_2_list)
+            norm2_flat = self.norm2(flattened)
+            norm2_list = uncat_with_shapes(norm2_flat, shapes, num_tokens)
+
+            residual_2_list = self.mlp.forward_list(norm2_list)
+
+            x_ffn = [
+                torch.index_add(
+                    x_attn,
+                    dim=0,
+                    source=self.ls2(residual_2),
+                    index=indices_2,
+                    alpha=residual_scale_factor,
+                )
+                for x_attn, residual_2, indices_2, residual_scale_factor in zip(
+                    x_attn_list, residual_2_list, indices_2_list, residual_scale_factors
+                )
+            ]
+        else:
+            x_out = []
+            for x, rope in zip(x_list, rope_list):
+                x_attn = x + self.ls1(self.attn(self.norm1(x), rope=rope))
+                x_ffn = x_attn + self.ls2(self.mlp(self.norm2(x_attn)))
+                x_out.append(x_ffn)
+            x_ffn = x_out
+
+        return x_ffn
+
+    def forward(self, x_or_x_list, rope_or_rope_list=None) -> List[Tensor]:
+        if isinstance(x_or_x_list, Tensor):
+            # for reference:
+            # return self._forward(x_or_x_list, rope=rope_or_rope_list)
+            # in order to match implementations we call the list op:
+            return self._forward_list([x_or_x_list], rope_list=[rope_or_rope_list])[0]
+        elif isinstance(x_or_x_list, list):
+            if rope_or_rope_list is None:
+                rope_or_rope_list = [None for x in x_or_x_list]
+            # return [self._forward(x, rope=rope) for x, rope in zip(x_or_x_list, rope_or_rope_list)]
+            return self._forward_list(x_or_x_list, rope_list=rope_or_rope_list)
+        else:
+            raise AssertionError
+
+
+class CausalSelfAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        ffn_ratio: float = 4.0,
+        ls_init_value: Optional[float] = None,
+        is_causal: bool = True,
+        act_layer: Callable = nn.GELU,
+        norm_layer: Callable = nn.LayerNorm,
+        dropout_prob: float = 0.0,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.is_causal = is_causal
+        self.ls1 = LayerScale(dim, init_values=ls_init_value) if ls_init_value else nn.Identity()
+        self.attention_norm = norm_layer(dim)
+        self.attention = CausalSelfAttention(dim, num_heads, attn_drop=dropout_prob, proj_drop=dropout_prob)
+
+        self.ffn_norm = norm_layer(dim)
+        ffn_hidden_dim = int(dim * ffn_ratio)
+        self.feed_forward = Mlp(
+            in_features=dim,
+            hidden_features=ffn_hidden_dim,
+            drop=dropout_prob,
+            act_layer=act_layer,
+        )
+
+        self.ls2 = LayerScale(dim, init_values=ls_init_value) if ls_init_value else nn.Identity()
+
+    def init_weights(
+        self,
+        init_attn_std: float | None = None,
+        init_proj_std: float | None = None,
+        init_fc_std: float | None = None,
+        factor: float = 1.0,
+    ) -> None:
+        init_attn_std = init_attn_std or (self.dim**-0.5)
+        init_proj_std = init_proj_std or init_attn_std * factor
+        init_fc_std = init_fc_std or (2 * self.dim) ** -0.5
+        self.attention.init_weights(init_attn_std, init_proj_std)
+        self.attention_norm.reset_parameters()
+        nn.init.normal_(self.feed_forward.fc1.weight, std=init_fc_std)
+        nn.init.normal_(self.feed_forward.fc2.weight, std=init_proj_std)
+        self.ffn_norm.reset_parameters()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+
+        x_attn = x + self.ls1(self.attention(self.attention_norm(x), self.is_causal))
+        x_ffn = x_attn + self.ls2(self.feed_forward(self.ffn_norm(x_attn)))
+        return x_ffn

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/ffn_layers.py

@@ -0,0 +1,83 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from typing import Callable, List, Optional
+
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from .utils import cat_keep_shapes, uncat_with_shapes
+
+
+class ListForwardMixin(object):
+    def forward(self, x: Tensor):
+        raise NotImplementedError
+
+    def forward_list(self, x_list: List[Tensor]) -> List[Tensor]:
+        x_flat, shapes, num_tokens = cat_keep_shapes(x_list)
+        x_flat = self.forward(x_flat)
+        return uncat_with_shapes(x_flat, shapes, num_tokens)
+
+
+class Mlp(nn.Module, ListForwardMixin):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+        device=None,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias, device=device)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias, device=device)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class SwiGLUFFN(nn.Module, ListForwardMixin):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Optional[Callable[..., nn.Module]] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+        align_to: int = 8,
+        device=None,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        d = int(hidden_features * 2 / 3)
+        swiglu_hidden_features = d + (-d % align_to)
+        self.w1 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device)
+        self.w2 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device)
+        self.w3 = nn.Linear(swiglu_hidden_features, out_features, bias=bias, device=device)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1 = self.w1(x)
+        x2 = self.w2(x)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/layer_scale.py

@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from typing import Union
+
+import torch
+from torch import Tensor, nn
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+        device=None,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(torch.empty(dim, device=device))
+        self.init_values = init_values
+
+    def reset_parameters(self):
+        nn.init.constant_(self.gamma, self.init_values)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/patch_embed.py

@@ -0,0 +1,95 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import math
+from typing import Callable, Tuple, Union
+
+from torch import Tensor, nn
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+    Args:
+        img_size: Image size.
+        patch_size: Patch token size.
+        in_chans: Number of input image channels.
+        embed_dim: Number of linear projection output channels.
+        norm_layer: Normalization layer.
+    """
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Callable | None = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        _, _, H, W = x.shape
+        # patch_H, patch_W = self.patch_size
+        # assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+        # assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+        x = self.proj(x)  # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)  # B HW C
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
+        return x
+
+    def flops(self) -> float:
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+    def reset_parameters(self):
+        k = 1 / (self.in_chans * (self.patch_size[0] ** 2))
+        nn.init.uniform_(self.proj.weight, -math.sqrt(k), math.sqrt(k))
+        if self.proj.bias is not None:
+            nn.init.uniform_(self.proj.bias, -math.sqrt(k), math.sqrt(k))

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/rms_norm.py

@@ -0,0 +1,30 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import torch
+from torch import Tensor, nn
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def reset_parameters(self) -> None:
+        nn.init.constant_(self.weight, 1)
+
+    def _norm(self, x: Tensor) -> Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: Tensor) -> Tensor:
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/rope_position_encoding.py

@@ -0,0 +1,127 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import math
+from typing import Literal
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+
+
+# RoPE positional embedding with no mixing of coordinates (axial) and no learnable weights
+# Supports two parametrizations of the rope parameters: either using `base` or `min_period` and `max_period`.
+class RopePositionEmbedding(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        *,
+        num_heads: int,
+        base: float | None = 100.0,
+        min_period: float | None = None,
+        max_period: float | None = None,
+        normalize_coords: Literal["min", "max", "separate"] = "separate",
+        shift_coords: float | None = None,
+        jitter_coords: float | None = None,
+        rescale_coords: float | None = None,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+        assert embed_dim % (4 * num_heads) == 0
+        both_periods = min_period is not None and max_period is not None
+        if (base is None and not both_periods) or (base is not None and both_periods):
+            raise ValueError("Either `base` or `min_period`+`max_period` must be provided.")
+
+        D_head = embed_dim // num_heads
+        self.base = base
+        self.min_period = min_period
+        self.max_period = max_period
+        self.D_head = D_head
+        self.normalize_coords = normalize_coords
+        self.shift_coords = shift_coords
+        self.jitter_coords = jitter_coords
+        self.rescale_coords = rescale_coords
+
+        # Needs persistent=True because we do teacher.load_state_dict(student.state_dict()) to initialize the teacher
+        self.dtype = dtype  # Don't rely on self.periods.dtype
+        self.register_buffer(
+            "periods",
+            torch.empty(D_head // 4, device=device, dtype=dtype),
+            persistent=True,
+        )
+        self._init_weights()
+
+    def forward(self, *, H: int, W: int) -> tuple[Tensor, Tensor]:
+        device = self.periods.device
+        dtype = self.dtype
+        dd = {"device": device, "dtype": dtype}
+
+        # Prepare coords in range [-1, +1]
+        if self.normalize_coords == "max":
+            max_HW = max(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / max_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / max_HW  # [W]
+        elif self.normalize_coords == "min":
+            min_HW = min(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / min_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / min_HW  # [W]
+        elif self.normalize_coords == "separate":
+            coords_h = torch.arange(0.5, H, **dd) / H  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / W  # [W]
+        else:
+            raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}")
+        coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1)  # [H, W, 2]
+        coords = coords.flatten(0, 1)  # [HW, 2]
+        coords = 2.0 * coords - 1.0  # Shift range [0, 1] to [-1, +1]
+
+        # Shift coords by adding a uniform value in [-shift, shift]
+        if self.training and self.shift_coords is not None:
+            shift_hw = torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords)
+            coords += shift_hw[None, :]
+
+        # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter]
+        if self.training and self.jitter_coords is not None:
+            jitter_max = np.log(self.jitter_coords)
+            jitter_min = -jitter_max
+            jitter_hw = torch.empty(2, **dd).uniform_(jitter_min, jitter_max).exp()
+            coords *= jitter_hw[None, :]
+
+        # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale]
+        if self.training and self.rescale_coords is not None:
+            rescale_max = np.log(self.rescale_coords)
+            rescale_min = -rescale_max
+            rescale_hw = torch.empty(1, **dd).uniform_(rescale_min, rescale_max).exp()
+            coords *= rescale_hw
+
+        # Prepare angles and sin/cos
+        angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :]  # [HW, 2, D//4]
+        angles = angles.flatten(1, 2)  # [HW, D//2]
+        angles = angles.tile(2)  # [HW, D]
+        cos = torch.cos(angles)  # [HW, D]
+        sin = torch.sin(angles)  # [HW, D]
+
+        return (sin, cos)  # 2 * [HW, D]
+
+    def _init_weights(self):
+        device = self.periods.device
+        dtype = self.dtype
+        if self.base is not None:
+            periods = self.base ** (
+                2 * torch.arange(self.D_head // 4, device=device, dtype=dtype) / (self.D_head // 2)
+            )  # [D//4]
+        else:
+            base = self.max_period / self.min_period
+            exponents = torch.linspace(0, 1, self.D_head // 4, device=device, dtype=dtype)  # [D//4] range [0, 1]
+            periods = base**exponents  # range [1, max_period / min_period]
+            periods = periods / base  # range [min_period / max_period, 1]
+            periods = periods * self.max_period  # range [min_period, max_period]
+        self.periods.data = periods

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/utils.py

@@ -0,0 +1,136 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+import os
+import random
+import subprocess
+from typing import Callable, List, Optional, Tuple
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+
+logger = logging.getLogger("dinov3")
+
+
+def cat_keep_shapes(x_list: List[Tensor]) -> Tuple[Tensor, List[Tuple[int]], List[int]]:
+    shapes = [x.shape for x in x_list]
+    num_tokens = [x.select(dim=-1, index=0).numel() for x in x_list]
+    flattened = torch.cat([x.flatten(0, -2) for x in x_list])
+    return flattened, shapes, num_tokens
+
+
+def uncat_with_shapes(flattened: Tensor, shapes: List[Tuple[int]], num_tokens: List[int]) -> List[Tensor]:
+    outputs_splitted = torch.split_with_sizes(flattened, num_tokens, dim=0)
+    shapes_adjusted = [shape[:-1] + torch.Size([flattened.shape[-1]]) for shape in shapes]
+    outputs_reshaped = [o.reshape(shape) for o, shape in zip(outputs_splitted, shapes_adjusted)]
+    return outputs_reshaped
+
+
+def named_replace(
+    fn: Callable,
+    module: nn.Module,
+    name: str = "",
+    depth_first: bool = True,
+    include_root: bool = False,
+) -> nn.Module:
+    if not depth_first and include_root:
+        module = fn(module=module, name=name)
+    for child_name_o, child_module in list(module.named_children()):
+        child_name = ".".join((name, child_name_o)) if name else child_name_o
+        new_child = named_replace(
+            fn=fn,
+            module=child_module,
+            name=child_name,
+            depth_first=depth_first,
+            include_root=True,
+        )
+        setattr(module, child_name_o, new_child)
+
+    if depth_first and include_root:
+        module = fn(module=module, name=name)
+    return module
+
+
+def named_apply(
+    fn: Callable,
+    module: nn.Module,
+    name: str = "",
+    depth_first: bool = True,
+    include_root: bool = False,
+) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(
+            fn=fn,
+            module=child_module,
+            name=child_name,
+            depth_first=depth_first,
+            include_root=True,
+        )
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+def fix_random_seeds(seed: int = 31):
+    """
+    Fix random seeds.
+    """
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+
+def get_sha() -> str:
+    cwd = os.path.dirname(os.path.abspath(__file__))
+
+    def _run(command):
+        return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
+
+    sha = "N/A"
+    diff = "clean"
+    branch = "N/A"
+    try:
+        sha = _run(["git", "rev-parse", "HEAD"])
+        subprocess.check_output(["git", "diff"], cwd=cwd)
+        diff = _run(["git", "diff-index", "HEAD"])
+        diff = "has uncommited changes" if diff else "clean"
+        branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
+    except Exception:
+        pass
+    message = f"sha: {sha}, status: {diff}, branch: {branch}"
+    return message
+
+
+def get_conda_env() -> Tuple[Optional[str], Optional[str]]:
+    conda_env_name = os.environ.get("CONDA_DEFAULT_ENV")
+    conda_env_path = os.environ.get("CONDA_PREFIX")
+    return conda_env_name, conda_env_path
+
+
+def count_parameters(module: nn.Module) -> int:
+    c = 0
+    for m in module.parameters():
+        c += m.nelement()
+    return c
+
+
+def has_batchnorms(model: nn.Module) -> bool:
+    bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
+    for _, module in model.named_modules():
+        if isinstance(module, bn_types):
+            return True
+    return False

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/layers/vision_transformer.py

@@ -0,0 +1,424 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+from functools import partial
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn.init
+from torch import Tensor, nn
+
+from . import LayerScale, Mlp, PatchEmbed, RMSNorm, RopePositionEmbedding, SelfAttentionBlock, SwiGLUFFN
+from .utils import named_apply
+
+logger = logging.getLogger("dinov3")
+
+ffn_layer_dict = {
+    "mlp": Mlp,
+    "swiglu": SwiGLUFFN,
+    "swiglu32": partial(SwiGLUFFN, align_to=32),
+    "swiglu64": partial(SwiGLUFFN, align_to=64),
+    "swiglu128": partial(SwiGLUFFN, align_to=128),
+}
+
+norm_layer_dict = {
+    "layernorm": partial(nn.LayerNorm, eps=1e-6),
+    "layernormbf16": partial(nn.LayerNorm, eps=1e-5),
+    "rmsnorm": RMSNorm,
+}
+
+dtype_dict = {
+    "fp32": torch.float32,
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+}
+
+
+def init_weights_vit(module: nn.Module, name: str = ""):
+    if isinstance(module, nn.Linear):
+        torch.nn.init.trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+        if hasattr(module, "bias_mask") and module.bias_mask is not None:
+            o = module.out_features
+            module.bias_mask.fill_(1)
+            module.bias_mask[o // 3 : 2 * o // 3].fill_(0)
+    if isinstance(module, nn.LayerNorm):
+        module.reset_parameters()
+    if isinstance(module, LayerScale):
+        module.reset_parameters()
+    if isinstance(module, PatchEmbed):
+        module.reset_parameters()
+    if isinstance(module, RMSNorm):
+        module.reset_parameters()
+
+
+class DinoVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        img_size: int = 224,
+        patch_size: int = 16,
+        in_chans: int = 3,
+        pos_embed_rope_base: float = 100.0,
+        pos_embed_rope_min_period: float | None = None,
+        pos_embed_rope_max_period: float | None = None,
+        pos_embed_rope_normalize_coords: Literal["min", "max", "separate"] = "separate",
+        pos_embed_rope_shift_coords: float | None = None,
+        pos_embed_rope_jitter_coords: float | None = None,
+        pos_embed_rope_rescale_coords: float | None = None,
+        pos_embed_rope_dtype: str = "bf16",
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        ffn_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        drop_path_rate: float = 0.0,
+        layerscale_init: float | None = None,
+        norm_layer: str = "layernorm",
+        ffn_layer: str = "mlp",
+        ffn_bias: bool = True,
+        proj_bias: bool = True,
+        n_storage_tokens: int = 0,
+        mask_k_bias: bool = False,
+        untie_cls_and_patch_norms: bool = False,
+        untie_global_and_local_cls_norm: bool = False,
+        device: Any | None = None,
+        **ignored_kwargs,
+    ):
+        super().__init__()
+        if len(ignored_kwargs) > 0:
+            logger.warning(f"Ignored kwargs: {ignored_kwargs}")
+        del ignored_kwargs
+
+        norm_layer_cls = norm_layer_dict[norm_layer]
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.n_blocks = depth
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            flatten_embedding=False,
+        )
+
+        self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim, device=device))
+        self.n_storage_tokens = n_storage_tokens
+        if self.n_storage_tokens > 0:
+            self.storage_tokens = nn.Parameter(torch.empty(1, n_storage_tokens, embed_dim, device=device))
+        logger.info(f"using base={pos_embed_rope_base} for rope new")
+        logger.info(f"using min_period={pos_embed_rope_min_period} for rope new")
+        logger.info(f"using max_period={pos_embed_rope_max_period} for rope new")
+        logger.info(f"using normalize_coords={pos_embed_rope_normalize_coords} for rope new")
+        logger.info(f"using shift_coords={pos_embed_rope_shift_coords} for rope new")
+        logger.info(f"using rescale_coords={pos_embed_rope_rescale_coords} for rope new")
+        logger.info(f"using jitter_coords={pos_embed_rope_jitter_coords} for rope new")
+        logger.info(f"using dtype={pos_embed_rope_dtype} for rope new")
+        self.rope_embed = RopePositionEmbedding(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            base=pos_embed_rope_base,
+            min_period=pos_embed_rope_min_period,
+            max_period=pos_embed_rope_max_period,
+            normalize_coords=pos_embed_rope_normalize_coords,
+            shift_coords=pos_embed_rope_shift_coords,
+            jitter_coords=pos_embed_rope_jitter_coords,
+            rescale_coords=pos_embed_rope_rescale_coords,
+            dtype=dtype_dict[pos_embed_rope_dtype],
+            device=device,
+        )
+        logger.info(f"using {ffn_layer} layer as FFN")
+        ffn_layer_cls = ffn_layer_dict[ffn_layer]
+        ffn_ratio_sequence = [ffn_ratio] * depth
+        blocks_list = [
+            SelfAttentionBlock(
+                dim=embed_dim,
+                num_heads=num_heads,
+                ffn_ratio=ffn_ratio_sequence[i],
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=drop_path_rate,
+                norm_layer=norm_layer_cls,
+                act_layer=nn.GELU,
+                ffn_layer=ffn_layer_cls,
+                init_values=layerscale_init,
+                mask_k_bias=mask_k_bias,
+                device=device,
+            )
+            for i in range(depth)
+        ]
+
+        self.chunked_blocks = False
+        self.blocks = nn.ModuleList(blocks_list)
+
+        # This norm is applied to everything, or when untying, to patch and mask tokens.
+        self.norm = norm_layer_cls(embed_dim)
+
+        self.untie_cls_and_patch_norms = untie_cls_and_patch_norms
+        if untie_cls_and_patch_norms:
+            # When untying, this norm is applied to CLS tokens and registers.
+            self.cls_norm = norm_layer_cls(embed_dim)
+        else:
+            self.cls_norm = None
+
+        self.untie_global_and_local_cls_norm = untie_global_and_local_cls_norm
+        if untie_global_and_local_cls_norm:
+            # When untying, this norm is applied to local CLS tokens and registers.
+            # This norm is never used during eval.
+            self.local_cls_norm = norm_layer_cls(embed_dim)
+        else:
+            self.local_cls_norm = None
+        self.head = nn.Identity()
+        self.mask_token = nn.Parameter(torch.empty(1, embed_dim, device=device))
+
+    def init_weights(self):
+        self.rope_embed._init_weights()
+        nn.init.normal_(self.cls_token, std=0.02)
+        if self.n_storage_tokens > 0:
+            nn.init.normal_(self.storage_tokens, std=0.02)
+        nn.init.zeros_(self.mask_token)
+        named_apply(init_weights_vit, self)
+
+    def prepare_tokens_with_masks(self, x: Tensor, masks=None) -> Tuple[Tensor, Tuple[int]]:
+        x = self.patch_embed(x)
+        B, H, W, _ = x.shape
+        x = x.flatten(1, 2)
+
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+            cls_token = self.cls_token
+        else:
+            cls_token = self.cls_token + 0 * self.mask_token
+        if self.n_storage_tokens > 0:
+            storage_tokens = self.storage_tokens
+        else:
+            storage_tokens = torch.empty(
+                1,
+                0,
+                cls_token.shape[-1],
+                dtype=cls_token.dtype,
+                device=cls_token.device,
+            )
+
+        x = torch.cat(
+            [
+                cls_token.expand(B, -1, -1),
+                storage_tokens.expand(B, -1, -1),
+                x,
+            ],
+            dim=1,
+        )
+
+        return x, (H, W)
+
+    def forward_features_list(self, x_list: List[Tensor], masks_list: List[Tensor]) -> List[Dict[str, Tensor]]:
+        x = []
+        rope = []
+        for t_x, t_masks in zip(x_list, masks_list):
+            t2_x, hw_tuple = self.prepare_tokens_with_masks(t_x, t_masks)
+            x.append(t2_x)
+            rope.append(hw_tuple)
+        for _, blk in enumerate(self.blocks):
+            if self.rope_embed is not None:
+                rope_sincos = [self.rope_embed(H=H, W=W) for H, W in rope]
+            else:
+                rope_sincos = [None for r in rope]
+            x = blk(x, rope_sincos)
+        all_x = x
+        output = []
+        for idx, (x, masks) in enumerate(zip(all_x, masks_list)):
+            if self.untie_cls_and_patch_norms or self.untie_global_and_local_cls_norm:
+                if self.untie_global_and_local_cls_norm and self.training and idx == 1:
+                    # Assume second entry of list corresponds to local crops.
+                    # We only ever apply this during training.
+                    x_norm_cls_reg = self.local_cls_norm(x[:, : self.n_storage_tokens + 1])
+                elif self.untie_cls_and_patch_norms:
+                    x_norm_cls_reg = self.cls_norm(x[:, : self.n_storage_tokens + 1])
+                else:
+                    x_norm_cls_reg = self.norm(x[:, : self.n_storage_tokens + 1])
+                x_norm_patch = self.norm(x[:, self.n_storage_tokens + 1 :])
+            else:
+                x_norm = self.norm(x)
+                x_norm_cls_reg = x_norm[:, : self.n_storage_tokens + 1]
+                x_norm_patch = x_norm[:, self.n_storage_tokens + 1 :]
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm_cls_reg[:, 0],
+                    "x_storage_tokens": x_norm_cls_reg[:, 1:],
+                    "x_norm_patchtokens": x_norm_patch,
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward_features(self, x: Tensor | List[Tensor], masks: Optional[Tensor] = None) -> List[Dict[str, Tensor]]:
+        if isinstance(x, torch.Tensor):
+            return self.forward_features_list([x], [masks])[0]
+        else:
+            return self.forward_features_list(x, masks)
+
+    def _get_intermediate_layers_not_chunked(self, x: Tensor, n: int = 1) -> List[Tensor]:
+        x, (H, W) = self.prepare_tokens_with_masks(x)
+        # If n is an int, take the n last blocks. If it's a list, take them
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for i, blk in enumerate(self.blocks):
+            if self.rope_embed is not None:
+                rope_sincos = self.rope_embed(H=H, W=W)
+            else:
+                rope_sincos = None
+            x = blk(x, rope_sincos)
+            if i in blocks_to_take:
+                output.append(x)
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        *,
+        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
+        reshape: bool = False,
+        return_class_token: bool = False,
+        return_extra_tokens: bool = False,
+        norm: bool = True,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, ...]]]:
+        outputs = self._get_intermediate_layers_not_chunked(x, n)
+        if norm:
+            outputs_normed = []
+            for out in outputs:
+                if self.untie_cls_and_patch_norms:
+                    x_norm_cls_reg = self.cls_norm(out[:, : self.n_storage_tokens + 1])
+                    x_norm_patch = self.norm(out[:, self.n_storage_tokens + 1 :])
+                    outputs_normed.append(torch.cat((x_norm_cls_reg, x_norm_patch), dim=1))
+                else:
+                    outputs_normed.append(self.norm(out))
+            outputs = outputs_normed
+        class_tokens = [out[:, 0] for out in outputs]
+        extra_tokens = [out[:, 1 : self.n_storage_tokens + 1] for out in outputs]
+        outputs = [out[:, self.n_storage_tokens + 1 :] for out in outputs]
+        if reshape:
+            B, _, h, w = x.shape
+            outputs = [
+                out.reshape(B, h // self.patch_size, w // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+                for out in outputs
+            ]
+        if not return_class_token and not return_extra_tokens:
+            return tuple(outputs)
+        elif return_class_token and not return_extra_tokens:
+            return tuple(zip(outputs, class_tokens))
+        elif not return_class_token and return_extra_tokens:
+            return tuple(zip(outputs, extra_tokens))
+        elif return_class_token and return_extra_tokens:
+            return tuple(zip(outputs, class_tokens, extra_tokens))
+
+    def forward(self, *args, is_training: bool = True, **kwargs) -> List[Dict[str, Tensor]] | Tensor:
+        # VGGT-Omega change: the aggregator consumes DINOv3 patch-token
+        # features directly, so the default forward returns the feature dict.
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return self.head(ret["x_norm_clstoken"])
+
+
+def vit_small(patch_size=16, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        ffn_ratio=4,
+        **kwargs,
+    )
+    return model
+
+
+def vit_base(patch_size=16, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        ffn_ratio=4,
+        **kwargs,
+    )
+    return model
+
+
+def vit_large(patch_size=16, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        ffn_ratio=4,
+        **kwargs,
+    )
+    return model
+
+
+def vit_so400m(patch_size=16, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1152,
+        depth=27,
+        num_heads=18,
+        ffn_ratio=3.777777778,
+        **kwargs,
+    )
+    return model
+
+
+def vit_huge2(patch_size=16, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1280,
+        depth=32,
+        num_heads=20,
+        ffn_ratio=4,
+        **kwargs,
+    )
+    return model
+
+
+def vit_giant2(patch_size=16, **kwargs):
+    """
+    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1536,
+        depth=40,
+        num_heads=24,
+        ffn_ratio=4,
+        **kwargs,
+    )
+    return model
+
+
+def vit_7b(patch_size=16, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=4096,
+        depth=40,
+        num_heads=32,
+        ffn_ratio=3,
+        **kwargs,
+    )
+    return model

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/models/vggt_omega.py

@@ -0,0 +1,88 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import warnings
+
+import torch
+import torch.nn as nn
+
+from vggt_omega.models.aggregator import Aggregator
+from vggt_omega.models.heads import CameraHead, DenseHead, TextAlignmentHead
+
+
+class VGGTOmega(nn.Module):
+    """Minimal VGGT-Omega inference model for camera and depth prediction."""
+
+    def __init__(
+        self,
+        patch_size: int = 16,
+        embed_dim: int = 1024,
+        enable_camera: bool = True,
+        enable_depth: bool = True,
+        enable_alignment: bool = False,
+    ) -> None:
+        super().__init__()
+
+        self.aggregator = Aggregator(patch_size=patch_size, embed_dim=embed_dim)
+        _warn_if_rope_not_max(self.aggregator)
+        self.camera_head = CameraHead(dim_in=2 * embed_dim) if enable_camera else None
+        self.dense_head = DenseHead(dim_in=2 * embed_dim, patch_size=patch_size) if enable_depth else None
+        self.text_alignment_head = TextAlignmentHead(dim_in=2 * embed_dim) if enable_alignment else None
+
+    def forward(self, images: torch.Tensor) -> dict[str, torch.Tensor]:
+        if len(images.shape) == 4:
+            images = images.unsqueeze(0)
+
+        amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        with torch.autocast(device_type="cuda", dtype=amp_dtype):
+            aggregated_tokens_list, patch_token_start = self.aggregator(images)
+
+        final_tokens = aggregated_tokens_list[-1]
+        if final_tokens is None:
+            raise ValueError("Aggregator did not cache the final layer, which VGGTOmega needs.")
+
+        predictions = {
+            "camera_and_register_tokens": final_tokens[:, :, :patch_token_start].contiguous(),
+        }
+        with torch.autocast(device_type="cuda", enabled=False):
+            if self.camera_head is not None:
+                predictions["pose_enc"] = self.camera_head(
+                    aggregated_tokens_list,
+                    patch_token_start=patch_token_start,
+                )
+
+            if self.dense_head is not None:
+                depth, depth_conf = self.dense_head(
+                    aggregated_tokens_list,
+                    images=images,
+                    patch_token_start=patch_token_start,
+                )
+                predictions["depth"] = depth
+                predictions["depth_conf"] = depth_conf
+
+            if self.text_alignment_head is not None:
+                predictions.update(
+                    self.text_alignment_head(
+                        aggregated_tokens_list,
+                        patch_token_start=patch_token_start,
+                    )
+                )
+
+        if not self.training:
+            predictions["images"] = images
+        return predictions
+
+
+def _warn_if_rope_not_max(aggregator: nn.Module) -> None:
+    for name, module in (("aggregator.patch_embed", aggregator.patch_embed), ("aggregator", aggregator)):
+        rope_embed = getattr(module, "rope_embed", None)
+        normalize_coords = getattr(rope_embed, "normalize_coords", None)
+        if normalize_coords != "max":
+            warnings.warn(
+                f"{name} RoPE normalize_coords is {normalize_coords!r}; "
+                "the released VGGT-Omega checkpoint was trained with 'max'.",
+                stacklevel=2,
+            )

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/utils/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Camera and geometry utilities for VGGT-Omega."""

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/utils/geometry.py

@@ -0,0 +1,34 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+import torch
+
+
+def closed_form_inverse_se3(se3, R=None, T=None):
+    """Invert a batch of 3x4 or 4x4 SE(3) matrices."""
+    is_numpy = isinstance(se3, np.ndarray)
+
+    if se3.shape[-2:] != (4, 4) and se3.shape[-2:] != (3, 4):
+        raise ValueError(f"se3 must have shape (N, 4, 4) or (N, 3, 4), got {se3.shape}")
+
+    if R is None:
+        R = se3[:, :3, :3]
+    if T is None:
+        T = se3[:, :3, 3:]
+
+    if is_numpy:
+        R_t = np.transpose(R, (0, 2, 1))
+        top_right = -np.matmul(R_t, T)
+        inverted = np.tile(np.eye(4), (len(R), 1, 1))
+    else:
+        R_t = R.transpose(1, 2)
+        top_right = -torch.bmm(R_t, T)
+        inverted = torch.eye(4, device=R.device, dtype=R.dtype)[None].repeat(len(R), 1, 1)
+
+    inverted[:, :3, :3] = R_t
+    inverted[:, :3, 3:] = top_right
+    return inverted

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/utils/load_fn.py

@@ -0,0 +1,129 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import warnings
+
+import numpy as np
+import torch
+from PIL import Image
+from torchvision import transforms as TF
+
+
+def load_and_preprocess_images(image_path_list, mode="balanced", image_resolution=512, patch_size=16):
+    """Load images for VGGT-Omega inference.
+
+    `balanced` keeps the total token count close to image_resolution**2.
+    `max_size` resizes the longest side to image_resolution.
+    Both modes first center-crop extreme aspect ratios into [0.5, 2.0].
+    """
+    if len(image_path_list) == 0:
+        raise ValueError("At least 1 image is required")
+    if mode not in ["balanced", "max_size"]:
+        raise ValueError("Mode must be either 'balanced' or 'max_size'")
+    if image_resolution <= 0:
+        raise ValueError("image_resolution must be positive")
+    if patch_size <= 0:
+        raise ValueError("patch_size must be positive")
+    if image_resolution % patch_size != 0:
+        raise ValueError("image_resolution must be divisible by patch_size")
+
+    images = []
+    shapes = set()
+    to_tensor = TF.ToTensor()
+
+    for image_path in image_path_list:
+        image = _crop_to_supported_aspect_ratio(_load_rgb_image(image_path))
+        width, height = image.size
+        aspect_ratio = height / max(width, 1)
+
+        if mode == "balanced":
+            target_h, target_w = _balanced_target_shape(aspect_ratio, image_resolution, patch_size)
+        else:
+            target_h, target_w = _max_size_target_shape(aspect_ratio, image_resolution, patch_size)
+
+        image = image.resize((target_w, target_h), Image.Resampling.BICUBIC)
+        image = to_tensor(image)
+
+        shapes.add((image.shape[1], image.shape[2]))
+        images.append(image)
+
+    if len(shapes) > 1:
+        warnings.warn(f"Found images with different shapes: {shapes}; padding to a common size.", stacklevel=2)
+        images = _pad_images_to_common_size(images, shapes)
+
+    return torch.stack(images)
+
+
+def _load_rgb_image(image_path):
+    with Image.open(image_path) as image:
+        if image.mode == "RGBA":
+            background = Image.new("RGBA", image.size, (255, 255, 255, 255))
+            image = Image.alpha_composite(background, image)
+        return image.convert("RGB")
+
+
+def _crop_to_supported_aspect_ratio(image, min_aspect_ratio=0.5, max_aspect_ratio=2.0):
+    width, height = image.size
+    aspect_ratio = height / max(width, 1)
+
+    if aspect_ratio < min_aspect_ratio:
+        crop_width = min(width, max(1, int(round(height / min_aspect_ratio))))
+        left = max((width - crop_width) // 2, 0)
+        return image.crop((left, 0, left + crop_width, height))
+
+    if aspect_ratio > max_aspect_ratio:
+        crop_height = min(height, max(1, int(round(width * max_aspect_ratio))))
+        top = max((height - crop_height) // 2, 0)
+        return image.crop((0, top, width, top + crop_height))
+
+    return image
+
+
+def _balanced_target_shape(aspect_ratio, image_resolution, patch_size):
+    token_number = (image_resolution // patch_size) ** 2
+    w_patches = np.sqrt(token_number / aspect_ratio)
+    h_patches = token_number / w_patches
+    w_patches = max(1, int(np.round(w_patches)))
+    h_patches = max(1, int(np.round(h_patches)))
+    return h_patches * patch_size, w_patches * patch_size
+
+
+def _max_size_target_shape(aspect_ratio, image_resolution, patch_size):
+    if aspect_ratio >= 1.0:
+        height = image_resolution
+        width = _round_to_patch_multiple(image_resolution / aspect_ratio, patch_size)
+    else:
+        width = image_resolution
+        height = _round_to_patch_multiple(image_resolution * aspect_ratio, patch_size)
+    return height, width
+
+
+def _round_to_patch_multiple(value, patch_size):
+    return max(patch_size, int(np.round(float(value) / patch_size)) * patch_size)
+
+
+def _pad_images_to_common_size(images, shapes):
+    max_height = max(shape[0] for shape in shapes)
+    max_width = max(shape[1] for shape in shapes)
+
+    padded_images = []
+    for image in images:
+        h_padding = max_height - image.shape[1]
+        w_padding = max_width - image.shape[2]
+        if h_padding > 0 or w_padding > 0:
+            pad_top = h_padding // 2
+            pad_bottom = h_padding - pad_top
+            pad_left = w_padding // 2
+            pad_right = w_padding - pad_left
+            image = torch.nn.functional.pad(
+                image,
+                (pad_left, pad_right, pad_top, pad_bottom),
+                mode="constant",
+                value=1.0,
+            )
+        padded_images.append(image)
+
+    return padded_images

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/utils/pose_enc.py

@@ -0,0 +1,52 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+from .rotation import mat_to_quat, quat_to_mat
+
+
+def extri_intri_to_pose_encoding(extrinsics, intrinsics, image_size_hw):
+    """Convert camera extrinsics and intrinsics to VGGT-Omega pose encoding.
+
+    The released checkpoints use a 9D camera encoding:
+    translation (3), quaternion rotation (4), and vertical/horizontal FoV (2).
+    Extrinsics are camera-from-world matrices in OpenCV coordinates.
+    """
+    R = extrinsics[:, :, :3, :3]
+    T = extrinsics[:, :, :3, 3]
+
+    H, W = image_size_hw
+    quat = mat_to_quat(R)
+    fov_h = 2 * torch.atan((H / 2) / intrinsics[..., 1, 1])
+    fov_w = 2 * torch.atan((W / 2) / intrinsics[..., 0, 0])
+    return torch.cat([T, quat, fov_h[..., None], fov_w[..., None]], dim=-1).float()
+
+
+def encoding_to_camera(pose_encoding, image_size_hw, build_intrinsics=True):
+    """Decode VGGT-Omega pose encoding into extrinsics and intrinsics."""
+    T = pose_encoding[..., :3]
+    quat = pose_encoding[..., 3:7]
+    fov_h = pose_encoding[..., 7]
+    fov_w = pose_encoding[..., 8]
+
+    R = quat_to_mat(quat)
+    extrinsics = torch.cat([R, T[..., None]], dim=-1)
+
+    intrinsics = None
+    if build_intrinsics:
+        H, W = image_size_hw
+        fy = (H / 2.0) / torch.tan(fov_h / 2.0)
+        fx = (W / 2.0) / torch.tan(fov_w / 2.0)
+
+        intrinsics = torch.zeros(pose_encoding.shape[:2] + (3, 3), device=pose_encoding.device)
+        intrinsics[..., 0, 0] = fx
+        intrinsics[..., 1, 1] = fy
+        intrinsics[..., 0, 2] = W / 2
+        intrinsics[..., 1, 2] = H / 2
+        intrinsics[..., 2, 2] = 1.0
+
+    return extrinsics, intrinsics

legacy/vggt_newbank_boxing_gloves_step19999/code/third_party/vggt_omega/utils/rotation.py

@@ -0,0 +1,130 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Modified from PyTorch3D, https://github.com/facebookresearch/pytorch3d
+
+import torch
+import torch.nn.functional as F
+
+
+def quat_to_mat(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Quaternion Order: XYZW or say ijkr, scalar-last
+
+    Convert rotations given as quaternions to rotation matrices.
+    Args:
+        quaternions: quaternions with real part last,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    i, j, k, r = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def mat_to_quat(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to quaternions.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        quaternions with real part last, as tensor of shape (..., 4).
+        Quaternion Order: XYZW or say ijkr, scalar-last
+    """
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+
+    batch_dim = matrix.shape[:-2]
+    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(matrix.reshape(batch_dim + (9,)), dim=-1)
+
+    q_abs = _sqrt_positive_part(
+        torch.stack(
+            [1.0 + m00 + m11 + m22, 1.0 + m00 - m11 - m22, 1.0 - m00 + m11 - m22, 1.0 - m00 - m11 + m22], dim=-1
+        )
+    )
+
+    # we produce the desired quaternion multiplied by each of r, i, j, k
+    quat_by_rijk = torch.stack(
+        [
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
+        ],
+        dim=-2,
+    )
+
+    # We floor here at 0.1 but the exact level is not important; if q_abs is small,
+    # the candidate won't be picked.
+    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
+    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
+
+    # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
+    # forall i; we pick the best-conditioned one (with the largest denominator)
+    out = quat_candidates[F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :].reshape(batch_dim + (4,))
+
+    # Convert from rijk to ijkr
+    out = out[..., [1, 2, 3, 0]]
+
+    out = standardize_quaternion(out)
+
+    return out
+
+
+def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
+    """
+    Returns torch.sqrt(torch.max(0, x))
+    but with a zero subgradient where x is 0.
+    """
+    ret = torch.zeros_like(x)
+    positive_mask = x > 0
+    if torch.is_grad_enabled():
+        ret[positive_mask] = torch.sqrt(x[positive_mask])
+    else:
+        ret = torch.where(positive_mask, torch.sqrt(x), ret)
+    return ret
+
+
+def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part last,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions)

legacy/vggt_newbank_boxing_gloves_step19999/params/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 288, 4304], "chunk_shape": [27, 288, 4304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 4, 72, 1152], "chunk_shape": [27, 4, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [288, 2048], "chunk_shape": [288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [256, 3072], "chunk_shape": [256, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 2, 256, 2048], "chunk_shape": [18, 2, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 2, 256, 1024], "chunk_shape": [18, 2, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 2, 2048, 256], "chunk_shape": [18, 2, 2048, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 2, 1024, 256], "chunk_shape": [18, 2, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 512, 16384], "chunk_shape": [18, 2, 512, 16384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 256, 4096], "chunk_shape": [18, 2, 256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [18, 256, 3072], "chunk_shape": [18, 256, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [18, 256, 3072], "chunk_shape": [18, 256, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_branch.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_branch.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.k_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.k_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.k_rmsnorm.scale.value", "write_shape": [18, 32], "chunk_shape": [18, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.kv_norm.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.kv_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.log_gain.value", "write_shape": [18, 2], "chunk_shape": [18, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.out_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.out_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.q_norm.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.q_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.q_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.q_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.q_rmsnorm.scale.value", "write_shape": [18, 32], "chunk_shape": [18, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.v_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.left_xattn.v_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.k_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.k_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.k_rmsnorm.scale.value", "write_shape": [18, 32], "chunk_shape": [18, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.kv_norm.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.kv_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.log_gain.value", "write_shape": [18, 2], "chunk_shape": [18, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.out_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.out_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.q_norm.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.q_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.q_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.q_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.q_rmsnorm.scale.value", "write_shape": [18, 32], "chunk_shape": [18, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.v_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.main_xattn.v_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.merge_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.merge_proj.kernel.value", "write_shape": [18, 512, 1024], "chunk_shape": [18, 512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_branch.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_branch.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.k_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.k_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.k_rmsnorm.scale.value", "write_shape": [18, 32], "chunk_shape": [18, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.kv_norm.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.kv_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.log_gain.value", "write_shape": [18, 2], "chunk_shape": [18, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.out_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.out_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.q_norm.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.q_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.q_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.q_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.q_rmsnorm.scale.value", "write_shape": [18, 32], "chunk_shape": [18, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.v_proj.bias.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.spatial_inject_1.right_xattn.v_proj.kernel.value", "write_shape": [18, 256, 1024], "chunk_shape": [18, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_correlation_cholesky.value", "write_shape": [240, 960], "chunk_shape": [240, 960], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.fast_token_embedding.embedding.value", "write_shape": [256, 2048], "chunk_shape": [256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.fast_token_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.fast_token_proj.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.fusion_layer1.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.fusion_layer1.kernel.value", "write_shape": [1024, 4096], "chunk_shape": [1024, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.fusion_layer2.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.fusion_layer2.kernel.value", "write_shape": [1024, 2048], "chunk_shape": [1024, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.gate_sincos.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.gate_sincos.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.gate_task.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.gate_task.kernel.value", "write_shape": [1024, 2048], "chunk_shape": [1024, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.gate_task_stage.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.gate_task_stage.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.kv_transform.k_bias.value", "write_shape": [18, 1, 64], "chunk_shape": [18, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.kv_transform.k_coeffs.value", "write_shape": [18, 18], "chunk_shape": [18, 18], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.kv_transform.v_bias.value", "write_shape": [18, 1, 64], "chunk_shape": [18, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.kv_transform.v_coeffs.value", "write_shape": [18, 18], "chunk_shape": [18, 18], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.bank_token_embeds.left.value", "write_shape": [1, 24, 1024], "chunk_shape": [1, 24, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.bank_token_embeds.main.value", "write_shape": [1, 32, 1024], "chunk_shape": [1, 32, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.bank_token_embeds.right.value", "write_shape": [1, 24, 1024], "chunk_shape": [1, 24, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_in_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_in_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_pose_mlp.fc1.bias.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_pose_mlp.fc1.kernel.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_pose_mlp.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_pose_mlp.fc2.kernel.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_token_proj.fc1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_token_proj.fc1.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_token_proj.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_token_proj.fc2.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_token_proj.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cam_token_proj.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.conf_mlp.fc1.bias.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.conf_mlp.fc1.kernel.value", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.conf_mlp.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.conf_mlp.fc2.kernel.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.0.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.1.fc1.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.1.fc1.kernel.value", "write_shape": [256, 4096], "chunk_shape": [256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.1.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.1.fc2.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.1.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.0.1.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.0.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.1.fc1.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.1.fc1.kernel.value", "write_shape": [256, 4096], "chunk_shape": [256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.1.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.1.fc2.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.1.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.cross_view_fusion.blocks.1.1.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.depth_mlp.fc1.bias.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.depth_mlp.fc1.kernel.value", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.depth_mlp.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.depth_mlp.fc2.kernel.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.feat_in_norm.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.feat_in_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.0.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.1.fc1.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.1.fc1.kernel.value", "write_shape": [256, 4096], "chunk_shape": [256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.1.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.1.fc2.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.1.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.0.1.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.0.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.1.fc1.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.1.fc1.kernel.value", "write_shape": [256, 4096], "chunk_shape": [256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.1.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.1.fc2.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.1.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.left.layers.1.1.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.0.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.1.fc1.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.1.fc1.kernel.value", "write_shape": [256, 4096], "chunk_shape": [256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.1.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.1.fc2.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.1.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.0.1.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.0.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.1.fc1.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.1.fc1.kernel.value", "write_shape": [256, 4096], "chunk_shape": [256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.1.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.1.fc2.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.1.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.main.layers.1.1.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.0.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.1.fc1.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.1.fc1.kernel.value", "write_shape": [256, 4096], "chunk_shape": [256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.1.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.1.fc2.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.1.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.0.1.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.0.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.1.fc1.bias.value", "write_shape": [1024], "chunk_shape": [1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.1.fc1.kernel.value", "write_shape": [256, 4096], "chunk_shape": [256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.1.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.1.fc2.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.1.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.lang_fusers.right.layers.1.1.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_embed.value", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_fuse.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_fuse.kernel.value", "write_shape": [1024, 1024], "chunk_shape": [1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.0.fc1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.0.fc1.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.0.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.0.fc2.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.0.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.0.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.1.fc1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.1.fc1.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.1.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.1.fc2.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.1.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.1.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.2.fc1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.2.fc1.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.2.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.2.fc2.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.2.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.2.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.3.fc1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.3.fc1.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.3.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.3.fc2.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.3.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.layer_projectors.3.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.loc_log_gamma.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.mlp.fc1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.mlp.fc1.kernel.value", "write_shape": [256, 2048], "chunk_shape": [256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.mlp.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.mlp.fc2.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.mlp.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.mlp.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.out_ln.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.out_ln.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.query.value", "write_shape": [1, 24, 1024], "chunk_shape": [1, 24, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.log_gain.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.left.xattn.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.loc_log_gamma.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.mlp.fc1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.mlp.fc1.kernel.value", "write_shape": [256, 2048], "chunk_shape": [256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.mlp.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.mlp.fc2.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.mlp.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.mlp.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.out_ln.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.out_ln.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.query.value", "write_shape": [1, 32, 1024], "chunk_shape": [1, 32, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.log_gain.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.main.xattn.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.loc_log_gamma.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.mlp.fc1.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.mlp.fc1.kernel.value", "write_shape": [256, 2048], "chunk_shape": [256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.mlp.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.mlp.fc2.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.mlp.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.mlp.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.out_ln.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.out_ln.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.query.value", "write_shape": [1, 24, 1024], "chunk_shape": [1, 24, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.k_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.k_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.log_gain.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.out_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.q_ln.scale.value", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.q_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.attn.v_proj.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.perceivers.right.xattn.q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.pos2d_mlp.fc1.bias.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.pos2d_mlp.fc1.kernel.value", "write_shape": [2, 64], "chunk_shape": [2, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.pos2d_mlp.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.pos2d_mlp.fc2.kernel.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.pose_enc_mlp.fc1.bias.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.pose_enc_mlp.fc1.kernel.value", "write_shape": [9, 64], "chunk_shape": [9, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.pose_enc_mlp.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.pose_enc_mlp.fc2.kernel.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.ray_mlp.fc1.bias.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.ray_mlp.fc1.kernel.value", "write_shape": [6, 64], "chunk_shape": [6, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.ray_mlp.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.ray_mlp.fc2.kernel.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.t5_projector.fc1.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.t5_projector.fc1.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.t5_projector.fc2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.t5_projector.fc2.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.t5_projector.norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.t5_projector.norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.spatial_bank_builder.view_embed.embedding.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.stage_pred_from_vlm.bias.value", "write_shape": [15], "chunk_shape": [15], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.stage_pred_from_vlm.kernel.value", "write_shape": [512, 15], "chunk_shape": [512, 15], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.stage_projection.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.stage_projection.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.task_embeddings.embedding.value", "write_shape": [50, 512], "chunk_shape": [50, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.task_stage_embeddings.embedding.value", "write_shape": [149, 1024], "chunk_shape": [149, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}]}

legacy/vggt_newbank_boxing_gloves_step19999/params/d/a46900611b6ae8d52cc9c73344a60cd5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9f44289129b1bc2cde3993756c5441b0797f8f35aa3ac2307242ef4e1911c25
+size 639734