legacy: copy task12_da3_large_gtdepth_newbank_step19999 -> legacy/task12_da3_large_gtdepth_newbank_step19999

.gitattributes

@@ -2444,3 +2444,16 @@ legacy/task12_da3_giant_gtdepth_v2_1p5x_step19999/params/ocdbt.process_0/d/c516e
 legacy/task12_da3_giant_gtdepth_v2_1p5x_step19999/params/ocdbt.process_0/d/d0f1752779afa33444765ae859c7b43a filter=lfs diff=lfs merge=lfs -text
 legacy/task12_da3_giant_gtdepth_v2_1p5x_step19999/params/ocdbt.process_0/d/d9b2ee49d7408c6a61fee9c154994b45 filter=lfs diff=lfs merge=lfs -text
 legacy/task12_da3_giant_gtdepth_v2_1p5x_step19999/params/ocdbt.process_0/d/e73a1948be3256b94e98e9ea525d818c filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/assets/IliaLarchenko/behavior_224_rgb/norm_stats.json filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/2608e0b98840e07c0bddb547c83352b5 filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/26312df420be0ded6f7d2968b9eb31e4 filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/32e999b158e705cc43e62af68ed2c026 filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/35232fd09a73f92d9d541123f677d7b6 filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/589f02b3a7ab6adaca31058554766308 filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/58cb578d8033aedd34caa72dbd8345ae filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/5d963bbe5fe91155634586f828e91671 filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/60c607b33450be6346df3980b439ad2c filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/742cb389e1ad457ef05645486d8511fa filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/7de207fe9d87adc11cbbd47a23474c55 filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/8e22384919e66ee762de8aa29a643f1e filter=lfs diff=lfs merge=lfs -text
+legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/dc52e7b2f8d7e8b35764f64109dfd990 filter=lfs diff=lfs merge=lfs -text

legacy/task12_da3_large_gtdepth_newbank_step19999/README.md

@@ -0,0 +1,217 @@
+# Task-12 DA3-Large + ground-truth metric depth policy
+
+This directory documents the final task-12 policy trained as
+`b1k_task12_large_gtdepth_newbank`. It is a PiBehavior policy with:
+
+- frozen backbone: `depth-anything/DA3-Large-1.1`;
+- four DA3 feature taps: layers `11,15,19,23`;
+- 1024-channel, 18x18 features from three camera views;
+- **dataset/simulator ground-truth metric depth replacing DA3-Large's non-metric prediction**;
+- K/V-split spatial bank with locality and cross-view fusion;
+- no live batch-mean bank centering, so normal batch-size-1 policy serving does not zero the bank.
+
+Task: 2026 task index 12 / 2025 task embedding index 12, `preparing_lunch_box`.
+
+## Contents on Hugging Face
+
+The uploaded checkpoint folder contains:
+
+- `params/`: Orbax inference parameters;
+- `assets/`: policy normalization statistics and FAST tokenizer;
+- `code/`: the exact relevant model, loader, extractor, policy-loading, and serving source;
+- `reference/large_gt_spatial_inputs.py`: evaluation-side helper for constructing the spatial fields;
+- `reference/modernbert_b1k_tasks.pkl`: frozen task-language features used by the spatial bank;
+- this README.
+
+`train_state/` is deliberately omitted because it is optimizer state needed only to resume training.
+
+## Required model-build flags
+
+The stock PiBehavior configuration has no full DA3 branch. Set these variables **before importing
+JAX or calling `scripts/train_2026.py::build_config()`**:
+
+```bash
+export USE_DA3_FULL=1
+export USE_VGGT=0
+export DA3_CHANNELS=1024
+export DA3_GRID_H=18
+export DA3_GRID_W=18
+
+export DA3_KV_SPLIT=1
+export DA3_BANK_CENTER=0
+export DA3_PERC_LOCALITY=1
+export DA3_CROSS_VIEW=1
+export DA3_CROSS_VIEW_DEPTH=2
+export DA3_DEPTH_DROPOUT=0.5
+export DA3_BTE_QUERY=1
+
+export DA3_QK_NORM=1
+export DA3_PERC_NORM_FINAL=1
+export DA3_PERC_NORM_OUT=0
+export DA3_POS_EMB_SCALE=0.25
+export DA3_LOGIT_GAIN_INIT=3.0
+export DA3_INJ_GAIN_MAX=8.0
+export DA3_PERC_LOGIT_GAIN=1
+export DA3_PERC_GAIN_INIT=3.0
+export DA3_PERC_GAIN_MAX=8.0
+
+export DA3_SCALE=1.0
+export DA3_INIT_STD=0.01
+export DA3_LR_GROUPS=1
+
+export DA3_MODEL_NAME=depth-anything/DA3-Large-1.1
+export DA3_OUT_LAYERS=11,15,19,23
+export B1K_USE_GT_DEPTH=1
+```
+
+`DA3_DEPTH_DROPOUT=0.5` must match the stored configuration but is inactive during inference because
+the bank builder receives no training dropout RNG.
+
+Load strictly:
+
+```python
+params = restore_params(checkpoint_dir / "params", dtype=jnp.bfloat16)
+model = train_config.model.load(params, remove_extra_params=False)
+```
+
+Do not allow extra parameters to be silently removed. A stock configuration with `da3=None` can
+otherwise discard the entire spatial subtree and appear to load successfully.
+
+## Required per-timestep inputs
+
+View order is fixed:
+
+```text
+0 = zed/head
+1 = left wrist RealSense
+2 = right wrist RealSense
+```
+
+For every policy call, construct:
+
+| Observation field | Shape | Meaning |
+|---|---:|---|
+| `da3_features` | `[B,4,3,1024,18,18]` | DA3-Large taps, bf16 bits stored as `uint16` |
+| `da3_ray` | `[B,3,3,18,18]` | DA3 camera-frame unit rays |
+| `da3_depth` | `[B,3,1,18,18]` | **GT optical-axis Z-depth in metres** |
+| `camera_extrinsics` | `[B,3,4,4]` | robot-to-camera transform in OpenCV convention |
+| `lang_feat` | `[B,32,1024]` | ModernBERT task feature |
+| `lang_mask` | `[B,32]` | corresponding valid-token mask |
+
+The normal PiBehavior image, image-mask, state, task, and subtask fields are still required.
+
+## Spatial input procedure
+
+Use `reference/large_gt_spatial_inputs.py::build_spatial_inputs`:
+
+1. Arrange the three RGB images in head/left/right order.
+2. Supply matching robot-to-camera OpenCV extrinsics.
+3. Rescale camera intrinsics to the 252x252 DA3 raster.
+4. Run the frozen DA3-Large extractor to obtain features and rays.
+5. Ignore the extractor's third output—DA3-Large depth is non-metric.
+6. Convert the matching GT linear depth images to metres.
+7. Area-resize each native depth map directly to 18x18.
+8. Put this tensor into `Observation.da3_depth`.
+
+The 2D patch coordinates are generated internally by `SpatialBankBuilder` on the same 18x18 grid.
+The address/key path is:
+
+```text
+2D patch position + Plücker ray + view identity
+```
+
+The payload/value path is:
+
+```text
+DA3-Large latent feature + encoded log GT depth
+```
+
+Cross-view fusion then uses the supplied camera extrinsics. Depth does not independently alter the
+ray or 2D position; correctness comes from all tensors referring to the same camera, timestamp, and
+patch cell.
+
+## Depth conventions—critical
+
+Training used BEHAVIOR-1K `observation.depth_linear.*` videos:
+
+- native `gray12le` values are millimetres stored in `uint16`;
+- convert with `native.astype(float32) / 1000.0`;
+- do **not** request `gray16le`, which left-shifts the 12-bit samples by four bits;
+- depth is optical-axis Z-depth, not Euclidean range along the ray;
+- invalid/nonpositive depth should be handled consistently with the simulator before log-depth;
+- area interpolation was used from the native square raster to 18x18.
+
+If the evaluation environment emits metres already, do not divide again.
+
+## Camera convention
+
+Training camera pose metadata stored the camera pose in the robot frame. It was converted to
+robot-to-camera and then from OmniGibson OpenGL optical axes to OpenCV:
+
+```python
+GL_TO_CV = diag(1, -1, -1, 1)
+robot_to_camera_cv = GL_TO_CV @ inverse(camera_pose_in_robot)
+```
+
+If the evaluation system already supplies robot-to-camera OpenCV extrinsics, use them directly.
+Applying this conversion twice is incorrect.
+
+Training intrinsics used the OmniGibson camera ratio:
+
+```text
+fx = fy = image_width * 17.0 / 20.995
+cx = cy = image_width / 2
+```
+
+Real intrinsics supplied by the simulator are preferable. They must describe the image after
+rescaling to 252x252.
+
+## Language/task conditioning
+
+Use the `preparing_lunch_box` entry from `reference/modernbert_b1k_tasks.pkl`. The checkpoint uses
+task embedding index 12. Do not substitute a different task index even if the 2026 dataset has been
+reordered.
+
+## Serving integration
+
+The existing `scripts/serve_b1k.py` and `b1k.policies.policy_config` do not automatically run the
+frozen DA3 extractor or attach GT depth. The evaluation system must insert the spatial-input
+procedure before `PiBehaviorPolicy.infer`/`model.sample_actions`.
+
+At batch-size-1 serving:
+
+- keep `DA3_BANK_CENTER=0`;
+- do not introduce a population-mean centering workaround;
+- do not switch centering on;
+- ensure all six required spatial fields are present and finite.
+
+## Sanity gates before simulator evaluation
+
+For one batch-size-1 observation:
+
+```text
+features shape = (1,4,3,1024,18,18), dtype uint16
+ray shape      = (1,3,3,18,18), norm approximately 1
+depth shape    = (1,3,1,18,18), finite positive metres
+extrinsics     = (1,3,4,4), finite rigid transforms
+banks          = main (1,128,1024), left/right (1,96,1024)
+```
+
+Banks must be nonzero and change when the observation changes. Load failure, missing spatial fields,
+wrong feature dtype, view-order mismatch, millimetres passed as metres, or `DA3_BANK_CENTER=1` can
+invalidate evaluation without an obvious server crash.
+
+## Training recipe and result
+
+- Dataset: task 12, 200 episodes, 1,643,060 valid action-horizon frame samples.
+- Initialization: 50-task PiBehavior meta checkpoint.
+- Steps: 20,000.
+- Batch: 126.
+- Topology: one DA3-Large extractor GPU plus three JAX training GPUs.
+- Spatial depth dropout: 0.5.
+- Live bank centering: disabled.
+- Final checkpoint: step 19999. Final logged metric at step 19950: action loss `0.0621`, total loss
+  `0.0972`, gradient norm `0.4168`. The adjacent-batch action loss at step 19900 was `0.0611`;
+  normal batch-to-batch noise remains.
+- A causal correct/shuffled/mean/zeroed spatial-bank diagnostic is run separately after training;
+  its result is not required to load the checkpoint and should not be inferred from training loss.

legacy/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_step19999/assets/IliaLarchenko/behavior_224_rgb/fast_tokenizer/special_tokens_map.json

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

legacy/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_step19999/code/launch/launch_b1k_newbank.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+# NEW-BANK run (2026-07-24): K/V-split + moved bank_token_embeds + PERCEIVER LOCALITY + CROSS-VIEW 3D
+# FUSION, all spatial addons training from step 0. Fresh spatial (new arch), base warm-started from META.
+# Targets the confirmed weakness: geometry was load-bearing but NOT robustly per-sample specific
+# (multi-batch shuffle-damage ~0). Cross-view fusion (integrated 3D) + locality (per-location detail)
+# aim to make it discriminative. Judge by MULTI-BATCH shuffle-damage separating from zero.
+#
+# LR (user-specified 2026-07-24, plateau schedules):
+#   spatial (core & geom TOGETHER): 1e-4 from step 0, FLAT to 10k, cosine to 1e-5 (1/10) at 20k. No ramp.
+#   vlm (base+action expert): frozen 0-1k; 1e-6 @1k ramp to 1e-5 by 5k; flat to 10k; cosine to 1e-6 @20k.
+# 4 GPUs (HOLDOUT=2 => 2 train + 2 extract), BS=96 (falls back if OOM), 20k steps.
+set -uo pipefail
+GPUS="${1:-4,5,6,7}"; ACT="${2:-clean_up_your_desk}"
+export EXP_OVERRIDE=${EXP_OVERRIDE:-b1k_da3_newbank_desk}
+export B1K_INIT_PARAMS=${B1K_INIT_PARAMS:-/work/jack/behavior1k/checkpoints/behavior_50t_checkpoint/params}
+
+# --- deployable bank architecture ---
+# Keep the structural K/V split, but do NOT subtract the live batch mean. Live batch centering makes
+# every bank exactly zero when policy serving uses batch size 1. Existing centered checkpoints need a
+# frozen population mean on the evaluation side; new deployable training runs keep centering off.
+export DA3_KV_SPLIT=1 DA3_BANK_CENTER=0 DA3_SCALE=1.0 DA3_INIT_STD=0.01
+export DA3_PERC_LOCALITY=1 DA3_CROSS_VIEW=1 DA3_CROSS_VIEW_DEPTH=2
+export DA3_DEPTH_DROPOUT=0.5
+export DA3_AUX_GEOM_HEAD=0 DA3_DEPTH_TARGET_ONLY=0
+
+# --- spatial (core & geom together): flat 1e-4 from step 0, decay@10k to 1/10 ---
+# override the base launcher's geom delay/ramp defaults (500/2500) so geom == core, from step 0.
+# Phase *timings* (DECAY_START) are overridable so a shortened run can scale each phase by the same
+# relative amount (e.g. 20k run -> DECAY_START=5000); LR *magnitudes* (PEAK/END) stay fixed.
+export LR_CORE_PEAK=1.0e-4 LR_CORE_END=1.0e-5 LR_CORE_DECAY_START=${LR_CORE_DECAY_START:-10000} LR_CORE_DELAY=0 LR_CORE_RAMP=0
+export LR_GEOM_PEAK=1.0e-4 LR_GEOM_END=1.0e-5 LR_GEOM_DECAY_START=${LR_GEOM_DECAY_START:-10000} LR_GEOM_DELAY=0 LR_GEOM_RAMP=0
+# --- vlm: frozen 0-1k, ramp 1e-6->1e-5 by 5k, flat, decay@10k to 1/10 ---
+export LR_VLM_DELAY=${LR_VLM_DELAY:-1000} LR_VLM_RAMP=${LR_VLM_RAMP:-4000} LR_VLM_RAMP_START=1.0e-6
+export LR_VLM_PEAK=1.0e-5 LR_VLM_END=1.0e-6 LR_VLM_DECAY_START=${LR_VLM_DECAY_START:-10000}
+export WARMUP=1  # unused (all groups use plateau schedules) but keep >0
+
+# --- tuned qk-norm regime ---
+export DA3_QK_NORM=1 DA3_PERC_NORM_FINAL=1 DA3_PERC_NORM_OUT=0 DA3_POS_EMB_SCALE=0.25
+export DA3_LOGIT_GAIN_INIT=3.0 DA3_INJ_GAIN_MAX=8.0 DA3_PERC_GAIN_INIT=3.0 DA3_PERC_GAIN_MAX=8.0
+export DA3_SPATIAL_EPS=1e-16 DA3_SPATIAL_WD=1e-4
+
+export BS=${BS:-96} STEPS=${STEPS:-20000} LR_DECAY=${LR_DECAY:-20000} HOLDOUT=${HOLDOUT:-2}
+exec bash /work/jack/launch_b1k_v2.sh "$GPUS" "$ACT"

legacy/task12_da3_large_gtdepth_newbank_step19999/code/launch/launch_b1k_task12_large_gtdepth.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+# Task-12 experiment: DA3-Large-1.1 features + dataset GT metric depth.
+# Default is the exactly matched 2 extractor + 2 JAX training, BS=128 topology.
+# Optional throughput topology: HOLDOUT=1 BS=126 (1 extractor + 3 training), but do not use that for
+# the strict controlled comparison because it changes batch size and optimizer sample exposure.
+set -uo pipefail
+export B1K_2026_ROOT=/work/jack/behavior1k/data/behavior_2026_tasks0_49
+export EXP_OVERRIDE=${EXP_OVERRIDE:-b1k_task12_large_gtdepth_newbank}
+export DA3_MODEL_NAME=depth-anything/DA3-Large-1.1
+export DA3_OUT_LAYERS=11,15,19,23
+export DA3_CHANNELS=1024
+export B1K_USE_GT_DEPTH=1
+export BS=${BS:-128} STEPS=${STEPS:-20000} LR_DECAY=${LR_DECAY:-20000} HOLDOUT=${HOLDOUT:-2}
+export LR_CORE_DECAY_START=${LR_CORE_DECAY_START:-5000}
+export LR_GEOM_DECAY_START=${LR_GEOM_DECAY_START:-5000}
+export LR_VLM_DELAY=${LR_VLM_DELAY:-500} LR_VLM_RAMP=${LR_VLM_RAMP:-2000}
+export LR_VLM_DECAY_START=${LR_VLM_DECAY_START:-5000}
+exec bash /work/jack/launch_b1k_newbank.sh "${1:-0,1,2,3}" preparing_lunch_box

legacy/task12_da3_large_gtdepth_newbank_step19999/code/launch/launch_b1k_v2.sh

@@ -0,0 +1,114 @@
+#!/usr/bin/env bash
+# Task-#2+ launcher: all pipeline optimizations + hyperparameters tuned from the task-#1 run.
+# Arg1 = physical GPU list. 8 GPUs => 4 train + 4 dedicated extraction; 4 GPUs => inline (2 jobs in parallel).
+# Arg2 = activity name (default clean_up_your_desk). EXP name is derived from it.
+#
+# DELIBERATELY SEPARATE from launch_b1k_robopro.sh so the in-flight task-#1 run (and its babysitter,
+# which would relaunch on a crash) keeps its original recipe untouched.
+set -uo pipefail
+GPUS="${1:-0,1,2,3,4,5,6,7}"
+ACT="${2:-clean_up_your_desk}"
+cd /work/jack/behavior-1k-solution
+NGPU=$(echo "$GPUS" | tr ',' '\n' | grep -c .)
+
+# ---- hyperparameters (tuned from task #1) --------------------------------------------------------
+# Task #1 evidence: action_loss 0.0678 @34k x BS128 (= 4.35M samples) already matched the no-DA3
+# baseline's converged 0.069; 0.0654 @38k (4.86M). Target ~4.6M samples with the LR schedule MATCHED
+# to the budget (task #1 decayed over 50k, so its LR was still high mid-run -- a matched schedule
+# reaches comparable loss in far fewer steps).
+# BS default is derived AFTER the topology is known (see below) so it is always divisible by the
+# training-GPU count. A hardcoded 196 crashed an 8-GPU relaunch with
+# "Batch size 196 must be divisible by the number of devices 8" -- which would have tripped the
+# babysitter circuit breaker on any crash recovery.
+BS_EXPLICIT="${BS:-}"
+STEPS=${STEPS:-25000}              # 25000 x 196 = 4.90M samples (task#1 ref: 0.0654 action @4.86M)
+LR_DECAY=${LR_DECAY:-$STEPS}       # schedule matched to budget
+WARMUP=${WARMUP:-500}              # vlm/core ramp 0 -> peak by step 500
+LR_VLM_PEAK=${LR_VLM_PEAK:-4.0e-5};  LR_VLM_END=${LR_VLM_END:-2.0e-6}
+LR_CORE_PEAK=${LR_CORE_PEAK:-7.0e-4}; LR_CORE_END=${LR_CORE_END:-5.0e-5}
+LR_GEOM_PEAK=${LR_GEOM_PEAK:-7.0e-4}; LR_GEOM_END=${LR_GEOM_END:-1.0e-4}
+# Geometry group gets a DELAYED, SLOW phase-in: held at exactly 0 for the first 500 steps (while the
+# pretrained trunk does its normal warmup), then ramped 0 -> peak over 2.5k steps before cosine decay.
+# Geometry modules are sensitive to being driven hard from random init.
+LR_GEOM_DELAY=${LR_GEOM_DELAY:-500}
+LR_GEOM_RAMP=${LR_GEOM_RAMP:-2500}
+
+# ---- topology (unified: works at ANY GPU count, split-vs-inline decided by HOLDOUT) ---------------
+# HOLDOUT=0: all visible GPUs both train and extract inline (extraction competes with training).
+# HOLDOUT=k>0: hold the first k GPUs OUT of the training mesh for DEDICATED extraction; train on the
+#   remaining NGPU-k. Extraction runs continuously (never pauses for the step) and fills the prefetch
+#   queue, so training on the rest is fed steadily instead of bursting. Measured best on 8 GPUs at 4+4
+#   (1:1 extract:train). The same 1:1 ratio halves cleanly to 4 GPUs as 2+2.
+HOLDOUT=${HOLDOUT:-0}; QDEPTH=${QDEPTH:-6}; FWD=${FWD:-32}
+if [ "$NGPU" -ge 8 ]; then NW=${NW:-32}; else NW=${NW:-16}; fi
+if [ "$HOLDOUT" -gt 0 ]; then
+  EXTRACT_DEVS=$(seq -s, -f 'cuda:%g' 0 $((HOLDOUT-1)))
+  GMODE="${NGPU}-GPU: $((NGPU-HOLDOUT)) train + $HOLDOUT dedicated extract"
+else
+  EXTRACT_DEVS=$(seq -s, -f 'cuda:%g' 0 $((NGPU-1)))
+  GMODE="${NGPU}-GPU: all train, inline extraction"
+fi
+# NUMA pin: GPUs 0-3 sit on node 0 (CPUs 0-59,120-179), GPUs 4-7 on node 1 (60-119,180-239).
+# Pinning the process (dataloader workers inherit the affinity) keeps memory local via Linux
+# first-touch, avoiding cross-socket traffic on every batch. Skipped if the run spans both nodes.
+_g_first=$(echo "$GPUS" | cut -d, -f1); _g_last=$(echo "$GPUS" | tr ',' '\n' | tail -1)
+if [ "$_g_last" -le 3 ] 2>/dev/null;   then PIN="taskset -c 0-59,120-179"; PINNOTE="node0"
+elif [ "$_g_first" -ge 4 ] 2>/dev/null; then PIN="taskset -c 60-119,180-239"; PINNOTE="node1"
+else PIN=""; PINNOTE="none (spans both sockets)"; fi
+
+TRAIN_GPUS=$((NGPU - HOLDOUT))
+# 32 samples/GPU is the measured throughput peak (16->21.3, 32->27.8, 64->25.1 samp/s).
+BS=${BS_EXPLICIT:-$((32 * TRAIN_GPUS))}
+[ $((BS % TRAIN_GPUS)) -eq 0 ] || { echo "ERROR: BS=$BS not divisible by $TRAIN_GPUS training GPUs"; exit 1; }
+
+EXP="${EXP_OVERRIDE:-b1k_da3_v2_${ACT}}"
+CKDIR=outputs/checkpoints/pi_behavior_b1k_fast/$EXP
+if ls "$CKDIR" 2>/dev/null | grep -qE '^[0-9]+$'; then RES=1; OVW=0; MODE="RESUME from $(ls "$CKDIR"|grep -E '^[0-9]+$'|sort -n|tail -1)"; else RES=0; OVW=1; MODE="FRESH"; fi
+LOG=/work/jack/${EXP}_train.log
+echo "[$(date +%H:%M)] launching $EXP ($MODE, $GMODE) BS=$BS STEPS=$STEPS on GPUs $GPUS" | tee -a /work/jack/b1k_v2_babysitter.log
+
+CUDA_VISIBLE_DEVICES="$GPUS" \
+XLA_PYTHON_CLIENT_MEM_FRACTION=0.80 XLA_PYTHON_CLIENT_ALLOCATOR=platform \
+XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_all_reduce_combine_threshold_bytes=8388608 --xla_gpu_enable_highest_priority_async_stream=true" \
+JAX_COMPILATION_CACHE_DIR=/work/jack/.jax_compile_cache \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+PYTHONUNBUFFERED=1 \
+B1K_HOLDOUT_GPU0=$HOLDOUT \
+B1K_EXTRACT_DEVICES=$EXTRACT_DEVS B1K_DECODE_THREADS=1 \
+B1K_SHARED_DECODE=1 B1K_DLPACK=1 B1K_TORCH_COLLATE=1 \
+B1K_REMAT=${B1K_REMAT:-full} \
+B1K_DA3_FWD_CHUNK=${B1K_DA3_FWD_CHUNK:-$FWD} \
+B1K_EXTRACT_QUEUE=$QDEPTH \
+RAYON_NUM_THREADS=1 OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \
+NUMEXPR_NUM_THREADS=1 POLARS_MAX_THREADS=1 TOKENIZERS_PARALLELISM=false \
+HF_HOME=/work/jack/da3xvla_workspace/hf_cache HF_HUB_OFFLINE=1 \
+B1K_2026_ROOT=${B1K_2026_ROOT:-/work/jack/behavior1k/data/behavior_2026_task29_42} \
+USE_DA3_FULL=1 DA3_INIT_STD=${DA3_INIT_STD:-0.01} DA3_LOGIT_GAIN=1 DA3_SCALE=${DA3_SCALE:-2.0} DA3_LR_GROUPS=1 \
+ DA3_BANK_CENTER=${DA3_BANK_CENTER:-0} \
+ DA3_AUX_GEOM_HEAD=${DA3_AUX_GEOM_HEAD:-0} DA3_AUX_GEOM_WEIGHT=${DA3_AUX_GEOM_WEIGHT:-0.0} \
+ DA3_DEPTH_TARGET_ONLY=${DA3_DEPTH_TARGET_ONLY:-0} \
+ DA3_KV_SPLIT=${DA3_KV_SPLIT:-0} DA3_DEPTH_DROPOUT=${DA3_DEPTH_DROPOUT:-0.0} \
+ DA3_PERC_LOCALITY=${DA3_PERC_LOCALITY:-0} DA3_CROSS_VIEW=${DA3_CROSS_VIEW:-0} DA3_CROSS_VIEW_DEPTH=${DA3_CROSS_VIEW_DEPTH:-2} \
+ LR_VLM_RAMP_START=${LR_VLM_RAMP_START:-0.0} \
+ DA3_SPATIAL_EPS=${DA3_SPATIAL_EPS:-1e-16} \
+ DA3_QK_NORM=${DA3_QK_NORM:-1} DA3_PERC_NORM_FINAL=${DA3_PERC_NORM_FINAL:-1} DA3_POS_EMB_SCALE=${DA3_POS_EMB_SCALE:-0.25} \
+ DA3_SPATIAL_WD=${DA3_SPATIAL_WD:-1e-4} \
+ DA3_PERC_LOGIT_GAIN=1 DA3_PERC_NORM_OUT=${DA3_PERC_NORM_OUT:-1} \
+ DA3_LOGIT_GAIN_INIT=${DA3_LOGIT_GAIN_INIT:-3.0} DA3_INJ_GAIN_MAX=${DA3_INJ_GAIN_MAX:-8.0} \
+ DA3_PERC_GAIN_INIT=${DA3_PERC_GAIN_INIT:-3.0} DA3_PERC_GAIN_MAX=${DA3_PERC_GAIN_MAX:-8.0} \
+LR_VLM_PEAK=$LR_VLM_PEAK LR_VLM_END=$LR_VLM_END \
+LR_CORE_PEAK=$LR_CORE_PEAK LR_CORE_END=$LR_CORE_END \
+LR_GEOM_PEAK=$LR_GEOM_PEAK LR_GEOM_END=$LR_GEOM_END \
+LR_GEOM_DELAY=$LR_GEOM_DELAY LR_GEOM_RAMP=$LR_GEOM_RAMP \
+LR_VLM_DELAY=${LR_VLM_DELAY:-0} LR_VLM_RAMP=${LR_VLM_RAMP:-0} \
+LR_VLM_CYCLE=${LR_VLM_CYCLE:-0} LR_VLM_FREEZE=${LR_VLM_FREEZE:-1000} LR_VLM_COADAPT=${LR_VLM_COADAPT:-500} LR_VLM_CYCLE_END=${LR_VLM_CYCLE_END:-10000} \
+LR_CORE_DELAY=${LR_CORE_DELAY:-0} LR_CORE_RAMP=${LR_CORE_RAMP:-0} LR_CORE_RAMP_START=${LR_CORE_RAMP_START:-0.0} \
+LR_VLM_DECAY_START=${LR_VLM_DECAY_START:-0} LR_CORE_DECAY_START=${LR_CORE_DECAY_START:-0} LR_GEOM_DECAY_START=${LR_GEOM_DECAY_START:-0} \
+B1K_ACTIVITIES=$ACT \
+B1K_INIT_PARAMS=${B1K_INIT_PARAMS:-/work/jack/behavior1k/checkpoints/behavior_50t_checkpoint/params} \
+EXP=$EXP \
+STEPS=$STEPS BS=$BS FLOW=15 NW=$NW SHUFFLE=1 RESUME=$RES OVERWRITE=$OVW \
+SAVE_INTERVAL=2000 KEEP_PERIOD=${KEEP_PERIOD:-10000} LOG_INTERVAL=50 \
+LR_WARMUP=$WARMUP LR_DECAY_STEPS=$LR_DECAY \
+  nohup $PIN .venv/bin/python scripts/train_2026.py >> "$LOG" 2>&1 &
+echo "LAUNCHED pid $! ($MODE, $GMODE) BS=$BS STEPS=$STEPS LR_DECAY=$LR_DECAY numa=$PINNOTE -> $LOG"

legacy/task12_da3_large_gtdepth_newbank_step19999/code/openpi/src/openpi/training/data_loader.py

@@ -0,0 +1,676 @@
+from collections.abc import Iterator, Sequence
+import logging
+import multiprocessing
+import os
+import queue
+import threading
+import typing
+from typing import Literal, Protocol, SupportsIndex, TypeVar
+
+import jax
+import jax.numpy as jnp
+import lerobot.datasets.lerobot_dataset as lerobot_dataset
+import numpy as np
+import torch
+
+import openpi.models.model as _model
+import openpi.training.config as _config
+from openpi.training.droid_rlds_dataset import DroidRldsDataset
+import openpi.transforms as _transforms
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class Dataset(Protocol[T_co]):
+    """Interface for a dataset with random access."""
+
+    def __getitem__(self, index: SupportsIndex) -> T_co:
+        raise NotImplementedError("Subclasses of Dataset should implement __getitem__.")
+
+    def __len__(self) -> int:
+        raise NotImplementedError("Subclasses of Dataset should implement __len__.")
+
+
+class IterableDataset(Protocol[T_co]):
+    """Interface for an iterable dataset."""
+
+    def __iter__(self) -> Iterator[T_co]:
+        raise NotImplementedError("Subclasses of IterableDataset should implement __iter__.")
+
+    def __len__(self) -> int:
+        raise NotImplementedError("Subclasses of Dataset should implement __len__.")
+
+
+class DataLoader(Protocol[T_co]):
+    """Interface for a data loader."""
+
+    def data_config(self) -> _config.DataConfig:
+        """Get the data config for this data loader."""
+        raise NotImplementedError("Subclasses of DataLoader should implement data_config.")
+
+    def __iter__(self) -> Iterator[T_co]:
+        raise NotImplementedError("Subclasses of DataLoader should implement __iter__.")
+
+
+class TransformedDataset(Dataset[T_co]):
+    def __init__(self, dataset: Dataset, transforms: Sequence[_transforms.DataTransformFn]):
+        self._dataset = dataset
+        self._transform = _transforms.compose(transforms)
+
+    def __getitem__(self, index: SupportsIndex) -> T_co:
+        return self._transform(self._dataset[index])
+
+    def __len__(self) -> int:
+        return len(self._dataset)
+
+
+class IterableTransformedDataset(IterableDataset[T_co]):
+    def __init__(
+        self,
+        dataset: IterableDataset,
+        transforms: Sequence[_transforms.DataTransformFn],
+        *,
+        is_batched: bool = False,
+    ):
+        self._dataset = dataset
+        self._transform = _transforms.compose(transforms)
+        self._is_batched = is_batched
+
+    def __iter__(self):
+        for sample in self._dataset:
+            if self._is_batched:
+                # Transforms are designed to be applied to individual samples. So we need to split the batch into
+                # individual samples and apply the transform to each sample individually.
+                batch_size = next(v.shape[0] for v in sample.values())
+
+                # Split batch into individual samples using tree_map
+                individual_samples = [jax.tree.map(lambda x: x[i], sample) for i in range(batch_size)]  # noqa: B023
+
+                # Transform each sample
+                transformed = [self._transform(s) for s in individual_samples]
+
+                # Recombine batch with tree_map
+                yield jax.tree.map(lambda *x: np.stack(x, axis=0), *transformed)
+            else:
+                yield self._transform(sample)
+
+    def __len__(self) -> int:
+        return len(self._dataset)
+
+
+class FakeDataset(Dataset):
+    def __init__(self, model_config: _model.BaseModelConfig, num_samples: int):
+        self._num_samples = num_samples
+        self._observation_spec, self._action_spec = model_config.inputs_spec()
+
+    def __getitem__(self, index: SupportsIndex) -> dict:
+        rng = jax.random.key(index.__index__())
+
+        def make_from_spec(spec: jax.ShapeDtypeStruct):
+            nonlocal rng
+            rng, data_rng = jax.random.split(rng)
+            # Remove the batch dimension.
+            shape = spec.shape[1:]
+            if spec.dtype == jnp.float32:
+                return jax.random.uniform(data_rng, shape=shape, minval=-1.0, maxval=1.0)
+            if spec.dtype == jnp.int32:
+                return jax.random.randint(data_rng, shape=shape, minval=0, maxval=2048)
+            return jnp.zeros(shape=shape, dtype=spec.dtype)
+
+        observation = jax.tree.map(make_from_spec, self._observation_spec)
+        action = jax.tree.map(make_from_spec, self._action_spec)
+
+        return {
+            **observation.to_dict(),
+            "actions": action,
+        }
+
+    def __len__(self) -> int:
+        return self._num_samples
+
+
+def create_behavior_dataset(data_config: _config.DataConfig, action_horizon: int) -> Dataset:
+    """Create a dataset for training."""
+    from omnigibson.learning.datas.lerobot_dataset import BehaviorLeRobotDataset
+    
+    dataset = BehaviorLeRobotDataset(
+        repo_id=data_config.repo_id,
+        root=data_config.behavior_dataset_root,
+        tasks=["turning_on_radio"],
+        modalities=["rgb"],
+        local_only=True,
+        delta_timestamps={
+            key: [t / 30.0 for t in range(action_horizon)] for key in data_config.action_sequence_keys
+        },
+        episodes=data_config.episodes_index,
+        chunk_streaming_using_keyframe=True,
+        shuffle=True,
+    )
+
+    if data_config.prompt_from_task:
+        dataset = TransformedDataset(dataset, [_transforms.PromptFromLeRobotTask(dataset.meta.tasks)])
+
+    return dataset
+
+
+def create_torch_dataset(
+    data_config: _config.DataConfig, action_horizon: int, model_config: _model.BaseModelConfig
+) -> Dataset:
+    """Create a dataset for training."""
+    repo_id = data_config.repo_id
+    if repo_id is None:
+        raise ValueError("Repo ID is not set. Cannot create dataset.")
+    if repo_id == "fake":
+        return FakeDataset(model_config, num_samples=1024)
+
+    dataset_meta = lerobot_dataset.LeRobotDatasetMetadata(repo_id)
+    dataset = lerobot_dataset.LeRobotDataset(
+        data_config.repo_id,
+        delta_timestamps={
+            key: [t / dataset_meta.fps for t in range(action_horizon)] for key in data_config.action_sequence_keys
+        },
+        episodes=data_config.episodes_index,
+    )
+
+    if data_config.prompt_from_task:
+        dataset = TransformedDataset(dataset, [_transforms.PromptFromLeRobotTask(dataset_meta.tasks)])
+
+    return dataset
+
+
+def create_rlds_dataset(
+    data_config: _config.DataConfig,
+    action_horizon: int,
+    batch_size: int,
+    *,
+    shuffle: bool = False,
+) -> Dataset:
+    # At the moment, we only support DROID for RLDS datasets.
+    return DroidRldsDataset(
+        data_dir=data_config.rlds_data_dir,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        action_chunk_size=action_horizon,
+        action_space=data_config.action_space,
+        filter_dict_path=data_config.filter_dict_path,
+    )
+
+
+def transform_dataset(dataset: Dataset, data_config: _config.DataConfig, *, skip_norm_stats: bool = False) -> Dataset:
+    """Transform the dataset by applying the data transforms."""
+    norm_stats = {}
+    if data_config.repo_id != "fake" and not skip_norm_stats:
+        if data_config.norm_stats is None:
+            raise ValueError(
+                "Normalization stats not found. "
+                "Make sure to run `scripts/compute_norm_stats.py --config-name=<your-config>`."
+            )
+        norm_stats = data_config.norm_stats
+
+    return TransformedDataset(
+        dataset,
+        [
+            *data_config.repack_transforms.inputs,
+            *data_config.data_transforms.inputs,
+            _transforms.Normalize(norm_stats, use_quantiles=data_config.use_quantile_norm),
+            *data_config.model_transforms.inputs,
+        ],
+    )
+
+
+def transform_iterable_dataset(
+    dataset: IterableDataset,
+    data_config: _config.DataConfig,
+    *,
+    skip_norm_stats: bool = False,
+    is_batched: bool = False,
+) -> IterableDataset:
+    """Transform the dataset by applying the data transforms."""
+    norm_stats = {}
+    if data_config.repo_id != "fake" and not skip_norm_stats:
+        if data_config.norm_stats is None:
+            raise ValueError(
+                "Normalization stats not found. "
+                "Make sure to run `scripts/compute_norm_stats.py --config-name=<your-config>`."
+            )
+        norm_stats = data_config.norm_stats
+
+    return IterableTransformedDataset(
+        dataset,
+        [
+            *data_config.repack_transforms.inputs,
+            *data_config.data_transforms.inputs,
+            _transforms.Normalize(norm_stats, use_quantiles=data_config.use_quantile_norm),
+            *data_config.model_transforms.inputs,
+        ],
+        is_batched=is_batched,
+    )
+
+
+def create_data_loader(
+    config: _config.TrainConfig,
+    *,
+    sharding: jax.sharding.Sharding | None = None,
+    shuffle: bool = False,
+    num_batches: int | None = None,
+    skip_norm_stats: bool = False,
+    framework: Literal["jax", "pytorch"] = "jax",
+) -> DataLoader[tuple[_model.Observation, _model.Actions]]:
+    """Create a data loader for training.
+
+    Args:
+        config: The training configuration.
+        sharding: The sharding to use for the data loader (JAX only).
+        shuffle: Whether to shuffle the data.
+        num_batches: Determines the number of batches to return.
+        skip_norm_stats: Whether to skip data normalization.
+        framework: The framework to use ("jax" or "pytorch").
+    """
+    data_config = config.data.create(config.assets_dirs, config.model)
+    logging.info(f"data_config: {data_config}")
+
+    if data_config.rlds_data_dir is not None:
+        return create_rlds_data_loader(
+            data_config,
+            action_horizon=config.model.action_horizon,
+            batch_size=config.batch_size,
+            sharding=sharding,
+            shuffle=shuffle,
+            num_batches=num_batches,
+            skip_norm_stats=skip_norm_stats,
+            framework=framework,
+        )
+    return create_torch_data_loader(
+        data_config,
+        model_config=config.model,
+        action_horizon=config.model.action_horizon,
+        batch_size=config.batch_size,
+        sharding=sharding,
+        shuffle=shuffle,
+        num_batches=num_batches,
+        num_workers=config.num_workers,
+        seed=config.seed,
+        skip_norm_stats=skip_norm_stats,
+        framework=framework,
+    )
+
+
+def create_behavior_data_loader(
+    config: _config.TrainConfig,
+    *,
+    sharding: jax.sharding.Sharding | None = None,
+    shuffle: bool = False,
+    num_batches: int | None = None,
+    skip_norm_stats: bool = False,
+) -> DataLoader[tuple[_model.Observation, _model.Actions]]:
+    data_config = config.data.create(config.assets_dirs, config.model)
+    dataset = create_behavior_dataset(data_config, action_horizon=config.model.action_horizon)
+    dataset = transform_dataset(dataset, data_config, skip_norm_stats=skip_norm_stats)
+
+    data_loader = TorchDataLoader(
+        dataset,
+        local_batch_size=config.batch_size // jax.process_count(),
+        sharding=sharding,
+        shuffle=shuffle,
+        num_batches=num_batches,
+        num_workers=config.num_workers,
+        seed=config.seed,
+    )
+    
+    return DataLoaderImpl(data_config, data_loader)
+
+
+def create_torch_data_loader(
+    data_config: _config.DataConfig,
+    model_config: _model.BaseModelConfig,
+    action_horizon: int,
+    batch_size: int,
+    *,
+    sharding: jax.sharding.Sharding | None = None,
+    skip_norm_stats: bool = False,
+    shuffle: bool = False,
+    num_batches: int | None = None,
+    num_workers: int = 0,
+    seed: int = 0,
+    framework: str = "jax",
+) -> DataLoader[tuple[_model.Observation, _model.Actions]]:
+    """Create a data loader for training.
+
+    Args:
+        data_config: The data configuration.
+        action_horizon: The action horizon.
+        batch_size: The batch size.
+        sharding: The sharding to use for the data loader. If None, the data loader will
+            use a single device sharding.
+        skip_norm_stats: Whether to skip data normalization.
+        shuffle: Whether to shuffle the data.
+        num_batches: Determines the number of batches to return. If the number exceeds the
+            number of batches in the dataset, the data loader will loop over the dataset.
+            If not provided, will iterate over the dataset indefinitely.
+        num_workers: The number of worker processes to use. If zero, the data loader will
+            execute in the main process.
+        seed: The seed to use for shuffling the data.
+    """
+    dataset = create_torch_dataset(data_config, action_horizon, model_config)
+    dataset = transform_dataset(dataset, data_config, skip_norm_stats=skip_norm_stats)
+
+    # Use TorchDataLoader for both frameworks
+    # For PyTorch DDP, create DistributedSampler and divide batch size by world size
+    # For JAX, divide by process count
+    sampler = None
+    if framework == "pytorch":
+        if torch.distributed.is_initialized():
+            sampler = torch.utils.data.distributed.DistributedSampler(
+                dataset,
+                num_replicas=torch.distributed.get_world_size(),
+                rank=torch.distributed.get_rank(),
+                shuffle=shuffle,
+                drop_last=True,
+            )
+            local_batch_size = batch_size // torch.distributed.get_world_size()
+        else:
+            local_batch_size = batch_size
+    else:
+        local_batch_size = batch_size // jax.process_count()
+
+    logging.info(f"local_batch_size: {local_batch_size}")
+    data_loader = TorchDataLoader(
+        dataset,
+        local_batch_size=local_batch_size,
+        sharding=None if framework == "pytorch" else sharding,
+        shuffle=(sampler is None and shuffle),  # Don't shuffle if using sampler
+        sampler=sampler,
+        num_batches=num_batches,
+        num_workers=num_workers,
+        seed=seed,
+        framework=framework,
+    )
+
+    return DataLoaderImpl(data_config, data_loader)
+
+
+def create_rlds_data_loader(
+    data_config: _config.DataConfig,
+    action_horizon: int,
+    batch_size: int,
+    *,
+    sharding: jax.sharding.Sharding | None = None,
+    skip_norm_stats: bool = False,
+    shuffle: bool = False,
+    num_batches: int | None = None,
+    framework: str = "jax",
+) -> DataLoader[tuple[_model.Observation, _model.Actions]]:
+    """Create an RLDS data loader for training.
+
+    Note: This data loader requires some extra dependencies -- see examples/droid/README_train.md
+
+    Args:
+        data_config: The data configuration.
+        action_horizon: The action horizon.
+        batch_size: The batch size.
+        sharding: The sharding to use for the data loader. If None, the data loader will
+            use a single device sharding.
+        skip_norm_stats: Whether to skip data normalization.
+        shuffle: Whether to shuffle the data.
+        num_batches: Determines the number of batches to return. If the number exceeds the
+            number of batches in the dataset, the data loader will loop over the dataset.
+            If not provided, will iterate over the dataset indefinitely.
+    """
+    if framework == "pytorch":
+        raise NotImplementedError("PyTorch RLDS data loader is not supported yet")
+    dataset = create_rlds_dataset(data_config, action_horizon, batch_size, shuffle=shuffle)
+    dataset = transform_iterable_dataset(dataset, data_config, skip_norm_stats=skip_norm_stats, is_batched=True)
+
+    data_loader = RLDSDataLoader(
+        dataset,
+        sharding=sharding,
+        num_batches=num_batches,
+    )
+
+    return DataLoaderImpl(data_config, data_loader)
+
+
+class TorchDataLoader:
+    """Torch data loader implementation."""
+
+    def __init__(
+        self,
+        dataset,
+        local_batch_size: int,
+        *,
+        sharding: jax.sharding.Sharding | None = None,
+        shuffle: bool = False,
+        sampler: torch.utils.data.Sampler | None = None,
+        num_batches: int | None = None,
+        num_workers: int = 0,
+        seed: int = 0,
+        framework: str = "jax",
+        batch_transform=None,
+    ):
+        """Create a PyTorch data loader.
+
+        Args:
+            dataset: The dataset to load.
+            local_batch_size: The local batch size for each process.
+            sharding: The sharding to use for the data loader.
+            shuffle: Whether to shuffle the data.
+            num_batches: If provided, determines the number of returned batches. If the
+                number is larger than the number of batches in the dataset, the data loader
+                will loop over the dataset. If not provided, will iterate over the dataset
+                indefinitely.
+            num_workers: The number of worker processes to use. If zero, the data loader will
+                execute in the main process.
+            seed: The seed to use for shuffling the data.
+        """
+        if jax.process_count() > 1:
+            raise NotImplementedError("Data loading with multiple processes is not supported.")
+
+        if len(dataset) < local_batch_size:
+            raise ValueError(f"Local batch size ({local_batch_size}) is larger than the dataset size ({len(dataset)}).")
+
+        # Store sharding - None for PyTorch, JAX sharding for JAX
+        self._sharding = sharding
+        if sharding is None and framework == "jax":
+            # Use data parallel sharding by default for JAX only.
+            self._sharding = jax.sharding.NamedSharding(
+                jax.sharding.Mesh(jax.devices(), ("B",)),
+                jax.sharding.PartitionSpec("B"),
+            )
+        self._num_batches = num_batches
+        # optional per-batch hook applied to the (numpy) batch before sharding — used for inline
+        # DA3 feature extraction (runs the frozen GIANT on the batch's raw frames on GPU).
+        self._batch_transform = batch_transform
+
+        mp_context = None
+        if num_workers > 0:
+            mp_context = multiprocessing.get_context("spawn")
+
+        generator = torch.Generator()
+        generator.manual_seed(seed)
+        data_loader_kwargs = {}
+        if num_workers > 0:
+            prefetch_factor = int(os.environ.get("B1K_PREFETCH_FACTOR", "2"))
+            if prefetch_factor > 0:
+                data_loader_kwargs["prefetch_factor"] = prefetch_factor
+            if _TORCH_COLLATE:
+                # Only meaningful for torch tensors; pins in the loader's pin thread so the
+                # extractor's H2D can be a true async DMA (non_blocking) instead of a pageable copy.
+                data_loader_kwargs["pin_memory"] = True
+        self._data_loader = torch.utils.data.DataLoader(
+            typing.cast(torch.utils.data.Dataset, dataset),
+            batch_size=local_batch_size,
+            shuffle=(sampler is None and shuffle),  # Don't shuffle if using sampler
+            sampler=sampler,
+            num_workers=num_workers,
+            multiprocessing_context=mp_context,
+            persistent_workers=num_workers > 0,
+            collate_fn=_collate_fn,
+            worker_init_fn=_worker_init_fn,
+            drop_last=True,
+            generator=generator,
+            **data_loader_kwargs,
+        )
+
+    @property
+    def torch_loader(self) -> torch.utils.data.DataLoader:
+        return self._data_loader
+
+    def _transformed_batches(self):
+        """Yields batch_transform-ed batches, looping over the dataset indefinitely.
+
+        When a batch_transform is set (inline DA3 GPU extraction), it runs ONE BATCH AHEAD in a
+        background thread, so the frozen DA3 forward for batch N+1 overlaps the train step for
+        batch N instead of serializing with it (torch GPU kernels release the GIL).
+        """
+        def epochs():
+            epoch = 0
+            sampler = getattr(self._data_loader, "sampler", None)
+            while True:
+                if hasattr(sampler, "set_epoch"):
+                    sampler.set_epoch(epoch)
+                yield from self._data_loader
+                epoch += 1
+
+        if self._batch_transform is None:
+            yield from epochs()
+            return
+
+        # Depth = how many extracted batches can be buffered ahead of training. Default 2 (unchanged).
+        # On a dedicated-extraction-GPU setup (extract on spare GPUs, train on others) bump this via
+        # B1K_EXTRACT_QUEUE=3-4 so extraction runs further ahead and smooths over per-batch variance.
+        q: queue.Queue = queue.Queue(maxsize=int(os.environ.get("B1K_EXTRACT_QUEUE", "2")))
+
+        def producer():
+            try:
+                for raw in epochs():
+                    b = self._batch_transform(raw)
+                    # Stage the REMAINING host fields (base RGB / actions / state) onto the training
+                    # GPUs here in the producer instead of on the main thread at consume time. Same
+                    # make_array, just moved off the train loop's critical path so the host->GPU copy
+                    # overlaps the in-flight train step. (DA3 fields are already jax.Array via DLPack.)
+                    if self._sharding is not None:
+                        b = jax.tree.map(
+                            lambda x: x if isinstance(x, jax.Array)
+                            else jax.make_array_from_process_local_data(self._sharding, _host_np(x)), b)
+                    q.put(b)
+            except BaseException as e:  # noqa: BLE001 — propagate any failure to the consumer
+                q.put(e)
+
+        threading.Thread(target=producer, daemon=True, name="batch-transform-prefetch").start()
+        while True:
+            item = q.get()
+            if isinstance(item, BaseException):
+                raise item
+            yield item
+
+    def __iter__(self):
+        num_items = 0
+        batches = self._transformed_batches()
+        while True:
+            # Check BEFORE pulling: the generator prefetches, and pulling past num_batches would
+            # trigger (and discard) a whole extra batch across the epoch boundary.
+            if self._num_batches is not None and num_items >= self._num_batches:
+                return
+            batch = next(batches)
+            num_items += 1
+            # For JAX, convert to sharded arrays; for PyTorch, return torch tensors.
+            # Fields already placed as jax.Array (DLPack GPU->GPU handoff) pass through untouched.
+            if self._sharding is not None:
+                yield jax.tree.map(
+                    lambda x: x if isinstance(x, jax.Array)
+                    else jax.make_array_from_process_local_data(self._sharding, _host_np(x)), batch)
+            else:
+                yield jax.tree.map(torch.as_tensor, batch)
+
+
+_TORCH_COLLATE = os.environ.get("B1K_TORCH_COLLATE") == "1"
+
+
+def _host_np(x):
+    """torch CPU tensor -> numpy (zero-copy view); pass anything else through unchanged."""
+    return x.numpy() if isinstance(x, torch.Tensor) else x
+
+
+def _collate_fn(items):
+    """Collate the batch elements into batched arrays.
+
+    With B1K_TORCH_COLLATE=1 the batch is built as torch tensors instead of numpy. The DataLoader
+    hands torch tensors worker->main through SHARED MEMORY (fd passing) rather than pickling and
+    copying them, which is a real saving at ~131MB/batch, and it makes pin_memory (hence async
+    non_blocking H2D) possible. Values are bit-identical either way.
+    """
+    if _TORCH_COLLATE:
+        def _stack(*xs):
+            a = [np.asarray(x) for x in xs]
+            if a[0].dtype.kind in "biufc":  # numeric -> torch tensor (shared-memory IPC + pinnable)
+                return torch.stack([torch.as_tensor(v) for v in a], 0)
+            return np.stack(a, axis=0)      # strings/objects have no torch equivalent; keep numpy
+        return jax.tree.map(_stack, *items)
+    # Make sure to convert to numpy arrays before stacking since some of the incoming elements
+    # may be JAX arrays.
+    return jax.tree.map(lambda *xs: np.stack([np.asarray(x) for x in xs], axis=0), *items)
+
+
+def _worker_init_fn(worker_id: int) -> None:
+    """Tell JAX inside the worker process not to preallocate the GPU memory."""
+    # NOTE: This is called after jax is imported inside the worker process. This
+    # means that this approach will not work for selecting the backend.
+    os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+    os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"
+
+
+class RLDSDataLoader:
+    """Shallow wrapper around the DROID data loader to make it compatible with openpi.
+
+    All batching already happens in the DROID dataset, so we don't need to do anything here.
+    """
+
+    def __init__(
+        self,
+        dataset: DroidRldsDataset,
+        *,
+        sharding: jax.sharding.Sharding | None = None,
+        num_batches: int | None = None,
+    ):
+        self._dataset = dataset
+        self._num_batches = num_batches
+
+        if jax.process_count() > 1:
+            raise NotImplementedError("Data loading with multiple processes is not supported.")
+
+        if sharding is None:
+            # Use data parallel sharding by default.
+            sharding = jax.sharding.NamedSharding(
+                jax.sharding.Mesh(jax.devices(), ("B",)),
+                jax.sharding.PartitionSpec("B"),
+            )
+
+        self._sharding = sharding
+        self._num_batches = num_batches
+
+    def __iter__(self):
+        num_items = 0
+        while True:
+            data_iter = iter(self._dataset)
+            while True:
+                if self._num_batches is not None and num_items >= self._num_batches:
+                    return
+                try:
+                    batch = next(data_iter)
+                except StopIteration:
+                    break  # We've exhausted the dataset. Create a new iterator and start over.
+                num_items += 1
+                yield jax.tree.map(lambda x: jax.make_array_from_process_local_data(self._sharding, x), batch)
+
+
+class DataLoaderImpl(DataLoader):
+    def __init__(self, data_config: _config.DataConfig, data_loader: TorchDataLoader | RLDSDataLoader):
+        self._data_config = data_config
+        self._data_loader = data_loader
+
+    def data_config(self) -> _config.DataConfig:
+        return self._data_config
+
+    def __iter__(self):
+        for batch in self._data_loader:
+            yield _model.Observation.from_dict(batch), batch["actions"]

legacy/task12_da3_large_gtdepth_newbank_step19999/code/scripts/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/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_step19999/code/src/b1k/policies/b1k_policy.py

@@ -0,0 +1,156 @@
+"""B1K Policy Transforms
+
+Transforms BEHAVIOR-1K observations to model format.
+
+Reference: https://github.com/wensi-ai/openpi/blob/behavior/src/openpi/policies/b1k_policy.py
+"""
+
+import dataclasses
+
+import einops
+import numpy as np
+
+from openpi import transforms
+from openpi.models import model as _model
+try:
+    # 2025 path: 256-dim R1Pro proprio indices live in OmniGibson (eval stack).
+    from omnigibson.learning.utils.eval_utils import PROPRIOCEPTION_INDICES
+except ImportError:
+    # 2026 path does its own 61-dim extraction (see b1k.training.b1k_2026) and
+    # never calls extract_state_from_proprio, so OmniGibson isn't required.
+    PROPRIOCEPTION_INDICES = None
+
+
+def make_b1k_example() -> dict:
+    """Creates a random input example for the Droid policy."""
+    return {
+        "observation/egocentric_camera": np.random.randint(256, size=(224, 224, 3), dtype=np.uint8),
+        "observation/wrist_image_left": np.random.randint(256, size=(224, 224, 3), dtype=np.uint8),
+        "observation/wrist_image_right": np.random.randint(256, size=(224, 224, 3), dtype=np.uint8),
+        "observation/joint_position": np.random.rand(23),
+        "prompt": "do something",
+    }
+
+def extract_state_from_proprio(proprio_data):
+    """
+    We assume perfect correlation for the two gripper fingers.
+    """
+    # extract joint position
+    base_qvel = proprio_data[..., PROPRIOCEPTION_INDICES["R1Pro"]["base_qvel"]]  # 3
+    trunk_qpos = proprio_data[..., PROPRIOCEPTION_INDICES["R1Pro"]["trunk_qpos"]]  # 4
+    arm_left_qpos = proprio_data[..., PROPRIOCEPTION_INDICES["R1Pro"]["arm_left_qpos"]]  #  7
+    arm_right_qpos = proprio_data[..., PROPRIOCEPTION_INDICES["R1Pro"]["arm_right_qpos"]]  #  7
+    
+    # Extract raw gripper widths and normalize to [-1, 1] to match action space
+    left_gripper_raw = proprio_data[..., PROPRIOCEPTION_INDICES["R1Pro"]["gripper_left_qpos"]].sum(axis=-1, keepdims=True)
+    right_gripper_raw = proprio_data[..., PROPRIOCEPTION_INDICES["R1Pro"]["gripper_right_qpos"]].sum(axis=-1, keepdims=True)
+    
+    # Normalize gripper widths from [0, 0.1] to [-1, 1] 
+    # Based on statistics: physical range is [0, 0.1], action range is [-1, 1]
+    # Formula: normalized = 2 * (raw / max_width) - 1
+    MAX_GRIPPER_WIDTH = 0.1  # From statistics q99 values
+    left_gripper_width = 2.0 * (left_gripper_raw / MAX_GRIPPER_WIDTH) - 1.0
+    right_gripper_width = 2.0 * (right_gripper_raw / MAX_GRIPPER_WIDTH) - 1.0
+
+    # Original baseline uses incorrect order for the state
+    return np.concatenate([
+        base_qvel,
+        trunk_qpos,
+        arm_left_qpos,
+        left_gripper_width,    # Now normalized [-1, 1]
+        arm_right_qpos,
+        right_gripper_width,   # Now normalized [-1, 1]
+    ], axis=-1)
+
+
+def _parse_image(image) -> np.ndarray:
+    image = np.asarray(image)
+    if np.issubdtype(image.dtype, np.floating):
+        image = (255 * image).astype(np.uint8)
+    if image.shape[0] == 3:
+        image = einops.rearrange(image, "c h w -> h w c")
+    return image
+
+
+@dataclasses.dataclass(frozen=True)
+class B1kInputs(transforms.DataTransformFn):
+    # Determines which model will be used (not actually used in B1K, kept for compatibility)
+    model_type: _model.ModelType | str = _model.ModelType.PI0
+
+    def __call__(self, data: dict) -> dict:
+
+        proprio_data = data["observation/state"]
+        # extract joint position
+        state = extract_state_from_proprio(proprio_data)
+        if "actions" in data:
+            action =  data["actions"]
+
+        # Possibly need to parse images to uint8 (H,W,C) since LeRobot automatically
+        # stores as float32 (C,H,W), gets skipped for policy inference
+        base_image = _parse_image(data["observation/egocentric_camera"])
+        wrist_image_left = _parse_image(data["observation/wrist_image_left"])
+        wrist_image_right = _parse_image(data["observation/wrist_image_right"])
+
+        # For B1K, always use 3 cameras (base, left_wrist, right_wrist)
+        names = ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
+        images = (base_image, wrist_image_left, wrist_image_right)
+        image_masks = (np.True_, np.True_, np.True_)
+
+        inputs = {
+            "state": state,
+            "image": dict(zip(names, images, strict=True)),
+            "image_mask": dict(zip(names, image_masks, strict=True)),
+        }
+
+        if "actions" in data:
+            inputs["actions"] = action
+
+        if "prompt" in data:
+            inputs["prompt"] = data["prompt"]
+            
+        # Preserve task_index for PI_BEHAVIOR model
+        if "task_index" in data:
+            inputs["task_index"] = data["task_index"]
+            
+        # Preserve tokenized_prompt for PI_BEHAVIOR model
+        if "tokenized_prompt" in data:
+            inputs["tokenized_prompt"] = data["tokenized_prompt"]
+        if "tokenized_prompt_mask" in data:
+            inputs["tokenized_prompt_mask"] = data["tokenized_prompt_mask"]
+            
+        # Preserve subtask_state for PI_BEHAVIOR model
+        if "subtask_state" in data:
+            inputs["subtask_state"] = data["subtask_state"]
+            
+        # Preserve timestamp and episode_index for subtask state computation
+        if "timestamp" in data:
+            inputs["timestamp"] = data["timestamp"]
+        if "episode_index" in data:
+            inputs["episode_index"] = data["episode_index"]
+            
+        # Preserve initial_actions for inpainting
+        if "initial_actions" in data:
+            initial_actions = data["initial_actions"]
+            # Pad initial_actions from 23 dimensions to 32 dimensions (model's action_dim)
+            if initial_actions.shape[-1] < 32:
+                padding_dim = 32 - initial_actions.shape[-1]
+                padding = np.zeros(initial_actions.shape[:-1] + (padding_dim,))
+                initial_actions = np.concatenate([initial_actions, padding], axis=-1)
+            inputs["initial_actions"] = initial_actions
+
+        return inputs
+
+
+@dataclasses.dataclass(frozen=True)
+class B1kOutputs(transforms.DataTransformFn):
+    def __call__(self, data: dict) -> dict:
+        # Return actions (truncated to 23 dims) and preserve subtask predictions
+        result = {"actions": np.asarray(data["actions"][:, :23])}
+        
+        # Preserve subtask prediction fields for PI_BEHAVIOR models
+        if "subtask_logits" in data:
+            result["subtask_logits"] = data["subtask_logits"]
+        if "predicted_stage" in data:
+            result["predicted_stage"] = data["predicted_stage"]
+            
+        return result

legacy/task12_da3_large_gtdepth_newbank_step19999/code/src/b1k/policies/pi_behavior_policy.py

@@ -0,0 +1,125 @@
+"""Minimal Policy subclass for PiBehavior models that handles tuple unpacking.
+
+Reference: https://github.com/PhysicalIntelligence/openpi/blob/behavior/src/openpi/policies/policy.py
+"""
+
+import time
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from typing_extensions import override
+
+from openpi.policies.policy import Policy
+from b1k.models.observation import Observation  # Use our custom Observation with FAST fields
+
+
+class PiBehaviorPolicy(Policy):
+    """Policy for PiBehavior models - only difference is unpacking the tuple return.
+    
+    PiBehavior.sample_actions() returns (actions, subtask_logits) instead of just actions.
+    This minimal subclass unpacks the tuple before output transforms are applied.
+    """
+    
+    @override
+    def infer(self, obs: dict, *, noise: np.ndarray | None = None, initial_actions: np.ndarray | None = None) -> dict:
+        """Infer with PiBehavior-specific tuple unpacking.
+        
+        Identical to parent Policy.infer() except:
+        1. Accepts initial_actions parameter for rolling inpainting
+        2. Unpacks (actions, subtask_logits) tuple before output transforms
+        """
+        # Reuse all parent logic for input processing
+        inputs = jax.tree.map(lambda x: x, obs)
+        inputs = self._input_transform(inputs)
+        inputs = jax.tree.map(lambda x: jnp.asarray(x)[np.newaxis, ...], inputs)
+        self._rng, sample_rng = jax.random.split(self._rng)
+
+        # Prepare sample_kwargs
+        sample_kwargs = dict(self._sample_kwargs)
+        if noise is not None:
+            noise = jnp.asarray(noise)
+            if noise.ndim == 2:
+                noise = noise[None, ...]
+            sample_kwargs["noise"] = noise
+            
+        if initial_actions is not None:
+            # Create training-format observation batch for proper transform processing
+            training_obs = {}
+            
+            # Map evaluation keys to training keys if needed
+            if "observation/state" in obs:
+                training_obs["observation/state"] = obs["observation/state"]
+            elif "state" in obs:
+                training_obs["observation/state"] = obs["state"]
+            
+            if "observation/egocentric_camera" in obs:
+                training_obs["observation/egocentric_camera"] = obs["observation/egocentric_camera"]
+            elif "image" in obs and "base_0_rgb" in obs["image"]:
+                training_obs["observation/egocentric_camera"] = obs["image"]["base_0_rgb"]
+                
+            if "observation/wrist_image_left" in obs:
+                training_obs["observation/wrist_image_left"] = obs["observation/wrist_image_left"]
+            elif "image" in obs and "left_wrist_0_rgb" in obs["image"]:
+                training_obs["observation/wrist_image_left"] = obs["image"]["left_wrist_0_rgb"]
+                
+            if "observation/wrist_image_right" in obs:
+                training_obs["observation/wrist_image_right"] = obs["observation/wrist_image_right"]
+            elif "image" in obs and "right_wrist_0_rgb" in obs["image"]:
+                training_obs["observation/wrist_image_right"] = obs["image"]["right_wrist_0_rgb"]
+                
+            # Copy any other keys that might be needed (tokenized_prompt, subtask_state, etc.)
+            for key in obs:
+                if key not in training_obs and key not in ["image", "state"]:
+                    training_obs[key] = obs[key]
+            
+            initial_batch = {
+                **training_obs,  # Include all observation data in training format
+                "actions": initial_actions  # Add initial_actions as the actions field
+            }
+            
+            # Apply the full input transform pipeline (delta transforms + normalization)
+            transformed_batch = self._input_transform(initial_batch)
+            normalized_initial_actions = transformed_batch["actions"]
+            
+            # Convert to JAX and add batch dim
+            initial_actions = jnp.asarray(normalized_initial_actions)
+            if initial_actions.ndim == 2:
+                initial_actions = initial_actions[None, ...]
+            sample_kwargs["initial_actions"] = initial_actions
+
+        observation = Observation.from_dict(inputs)
+        start_time = time.monotonic()
+        
+        # ONLY DIFFERENCE: Unpack tuple return from PiBehavior.sample_actions
+        actions, subtask_logits = self._sample_actions(sample_rng, observation, **sample_kwargs)
+        
+        outputs = {
+            "state": inputs["state"],
+            "actions": actions,  # Now an array, not a tuple!
+            "subtask_logits": subtask_logits,
+        }
+        
+        model_time = time.monotonic() - start_time
+        
+        # Convert to numpy (same as parent)
+        outputs = {
+            k: np.asarray(v[0, ...]) if isinstance(v, (jnp.ndarray, np.ndarray)) else v
+            for k, v in outputs.items()
+        }
+
+        # Apply output transforms (now works because actions is an array)
+        outputs = self._output_transform(outputs)
+        
+        # Add convenience field
+        outputs["predicted_stage"] = int(np.argmax(outputs["subtask_logits"]))
+        
+        outputs["policy_timing"] = {
+            "infer_ms": model_time * 1000,
+        }
+        return outputs
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        return self._metadata

legacy/task12_da3_large_gtdepth_newbank_step19999/code/src/b1k/policies/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/task12_da3_large_gtdepth_newbank_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/task12_da3_large_gtdepth_newbank_step19999/code/src/b1k/training/b1k_da3.py

@@ -0,0 +1,416 @@
+"""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 glob
+
+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",
+)
+DEPTH_SRCS = {
+    "observation.images.rgb.head": "observation.depth_linear.zed_link_camera_0",
+    "observation.images.rgb.left_wrist": "observation.depth_linear.left_realsense_link_camera_0",
+    "observation.images.rgb.right_wrist": "observation.depth_linear.right_realsense_link_camera_0",
+}
+
+
+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)
+        if os.environ.get("B1K_USE_GT_DEPTH") == "1":
+            import pandas as pd
+            depth_meta = pd.concat(
+                [pd.read_parquet(f) for f in sorted(glob.glob(
+                    os.path.join(self.root, "meta", "episodes", "**", "*.parquet"), recursive=True
+                ))],
+                ignore_index=True,
+            ).set_index("episode_index")
+            for rec in self.episodes:
+                row = depth_meta.loc[rec["episode_index"]]
+                rec["depth_video"] = {}
+                rec["depth_from_ts"] = {}
+                for dst, src in DEPTH_SRCS.items():
+                    rec["depth_video"][dst] = os.path.join(
+                        self.root, "videos", src,
+                        f"chunk-{int(row[f'videos/{src}/chunk_index']):03d}",
+                        f"file-{int(row[f'videos/{src}/file_index']):03d}.mp4",
+                    )
+                    rec["depth_from_ts"][dst] = float(row[f"videos/{src}/from_timestamp"])
+                    if not os.path.exists(rec["depth_video"][dst]):
+                        raise FileNotFoundError(
+                            f"GT depth requested but missing: {rec['depth_video'][dst]}"
+                        )
+            logger.info("GT metric depth enabled for %d episodes", len(self.episodes))
+
+    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 _decode_gt_depth(self, path: str, ts: float, grid_hw: tuple[int, int]) -> np.ndarray:
+        """Decode BEHAVIOR-1K gray12le linear depth and align it to the DA3 patch grid.
+
+        Native gray12le samples are metric millimeters (not gray16-scaled values).  The depth and
+        RGB videos share the same square camera raster/FOV; resizing native depth directly to the
+        (H/14,W/14) token grid therefore aligns each depth cell with the corresponding DA3 patch.
+        Area resampling averages metric Z-depth over the same image support represented by a token.
+        """
+        import cv2
+
+        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
+        # Do not request gray16le: swscale left-shifts the native 12-bit samples by four bits.
+        depth_m = frame.to_ndarray().astype(np.float32) / 1000.0
+        gh, gw = grid_hw
+        depth_grid = cv2.resize(depth_m, (gw, gh), interpolation=cv2.INTER_AREA)
+        return depth_grid[None]  # [1,gh,gw]
+
+    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, gt_depth = [], [], [], []
+        use_gt_depth = os.environ.get("B1K_USE_GT_DEPTH") == "1"
+        grid_hw = (h // 14, w // 14)
+        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)
+            if use_gt_depth:
+                depth_ts = rec["depth_from_ts"][dst] + t / self.fps
+                gt_depth.append(self._decode_gt_depth(
+                    rec["depth_video"][dst], depth_ts, grid_hw
+                ))
+
+        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
+        if use_gt_depth:
+            item["gt_metric_depth"] = np.stack(gt_depth, 0).astype(np.float32)  # [V,1,gh,gw], meters
+        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:
+        model_name = os.environ.get(
+            "DA3_MODEL_NAME", "depth-anything/DA3NESTED-GIANT-LARGE-1.1"
+        )
+        out_layers = tuple(int(x) for x in os.environ.get(
+            "DA3_OUT_LAYERS", "19,26,33,39"
+        ).split(","))
+        extractor = _ex.DA3InlineExtractor(
+            model_name=model_name, out_layers=out_layers, da3_hw=da3_hw,
+            forward_chunk=_fchunk, devices=_devices,
+        )
+        logger.info("Using DA3 extractor model=%s out_layers=%s", model_name, out_layers)
+
+    # 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):
+        use_gt_depth = os.environ.get("B1K_USE_GT_DEPTH") == "1"
+        if _dlpack_ok():
+            import torch
+            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]
+            if use_gt_depth:
+                gt = batch.pop("gt_metric_depth")
+                bounds = [round(i * int(gt.shape[0]) / len(parts)) for i in range(len(parts) + 1)]
+                replaced = []
+                for k, part in enumerate(parts):
+                    d = gt[bounds[k]:bounds[k + 1]]
+                    if not isinstance(d, torch.Tensor):
+                        d = torch.as_tensor(d)
+                    d = d.to(extractor.devices[k], dtype=torch.float32, non_blocking=True)
+                    replaced.append((part[0], part[1], d))
+                parts = replaced
+
+            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"])
+            if use_gt_depth:
+                gt = batch.pop("gt_metric_depth")
+                if hasattr(gt, "detach"):
+                    gt = gt.detach().cpu().numpy()
+                outs = (outs[0], outs[1], np.asarray(gt, dtype=np.float32))
+            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/task12_da3_large_gtdepth_newbank_step19999/code/src/b1k/training/config.py

@@ -0,0 +1,438 @@
+"""Training configuration for BEHAVIOR-1K challenge.
+
+Reference: https://github.com/Physical-Intelligence/openpi
+"""
+
+import abc
+from collections.abc import Sequence
+import dataclasses
+import difflib
+import logging
+import os
+import pathlib
+from typing import Any, Literal, List, Protocol, TypeAlias
+
+import etils.epath as epath
+import flax.nnx as nnx
+from typing_extensions import override
+import tyro
+
+# Import from OpenPI
+import openpi.models.model as _model
+import openpi.models.pi0_config as pi0_config
+import openpi.policies.aloha_policy as aloha_policy
+import openpi.policies.droid_policy as droid_policy
+import openpi.policies.libero_policy as libero_policy
+import openpi.shared.download as _download
+import openpi.training.droid_rlds_dataset as droid_rlds_dataset
+import openpi.training.optimizer as _optimizer
+import openpi.transforms as _transforms
+
+# Import from B1K custom modules
+from b1k.models import pi_behavior_config
+from b1k.policies import b1k_policy
+from b1k.shared import normalize as _normalize
+from b1k.training import weight_loaders
+from b1k import transforms as b1k_transforms
+
+ModelType: TypeAlias = _model.ModelType
+Filter: TypeAlias = nnx.filterlib.Filter
+
+
+@dataclasses.dataclass(frozen=True)
+class AssetsConfig:
+    """Determines the location of assets (e.g., norm stats) that will be used to set up the data pipeline.
+
+    These assets will be replicated inside the checkpoint under the `assets/asset_id` directory.
+    """
+    # Assets directory. If not provided, the config assets_dirs will be used.
+    assets_dir: str | None = None
+    # Asset id. If not provided, the repo id will be used.
+    asset_id: str | None = None
+
+
+@dataclasses.dataclass(frozen=True)
+class DataConfig:
+    # LeRobot repo id. If None, fake data will be created.
+    repo_id: str | None = None
+    # Directory within the assets directory containing the data assets.
+    asset_id: str | None = None
+    # Contains precomputed normalization stats. If None, normalization will not be performed.
+    norm_stats: dict[str, _transforms.NormStats] | None = None
+
+    # Used to adopt the inputs from a dataset specific format to a common format
+    repack_transforms: _transforms.Group = dataclasses.field(default_factory=_transforms.Group)
+    # Data transforms, typically include robot specific transformations.
+    data_transforms: _transforms.Group = dataclasses.field(default_factory=_transforms.Group)
+    # Model specific transforms. Will be applied after the data is normalized.
+    model_transforms: _transforms.Group = dataclasses.field(default_factory=_transforms.Group)
+    
+    # If true, will use quantile normalization. Otherwise, normal z-score normalization will be used.
+    use_quantile_norm: bool = False
+    # If true, will use per-timestamp normalization for actions
+    use_per_timestamp_norm: bool = False
+
+    # Names of keys that will be used by the data loader to generate the action sequence.
+    action_sequence_keys: Sequence[str] = ("actions",)
+
+    # If true, will use the LeRobot dataset task to define the prompt (not used for PI_BEHAVIOR).
+    prompt_from_task: bool = False
+
+    # Only used for RLDS data loader.
+    rlds_data_dir: str | None = None
+
+    # Only used for B1K data loader.
+    behavior_dataset_root: str | None = None
+
+    # Action space for DROID dataset.
+    action_space: droid_rlds_dataset.DroidActionSpace | None = None
+    # Path to the data filter file for DROID dataset
+    filter_dict_path: str | None = None
+
+    # Episodes index to use for training 
+    episodes_index: List[int] | None = None
+
+
+class GroupFactory(Protocol):
+    def __call__(self, model_config: _model.BaseModelConfig) -> _transforms.Group:
+        """Create a group."""
+
+
+@dataclasses.dataclass(frozen=True)
+class ModelTransformFactory(GroupFactory):
+    """Creates model transforms for B1K."""
+
+    default_prompt: str | None = None  # Not used (task embeddings instead)
+
+    def __call__(self, model_config: _model.BaseModelConfig) -> _transforms.Group:
+        return _transforms.Group(
+            inputs=[
+                _transforms.ResizeImages(224, 224),
+                b1k_transforms.ComputeSubtaskStateFromMeta(dataset=None),
+                b1k_transforms.TaskIndexToTaskId(),
+                _transforms.PadStatesAndActions(model_config.action_dim),
+            ],
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class DataConfigFactory(abc.ABC):
+    # The LeRobot repo id.
+    repo_id: str = tyro.MISSING
+    # Determines how the assets will be loaded.
+    assets: AssetsConfig = dataclasses.field(default_factory=AssetsConfig)
+    # Base config that will be updated by the factory.
+    base_config: tyro.conf.Suppress[DataConfig | None] = None
+
+    @abc.abstractmethod
+    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
+        """Create a data config."""
+
+    def create_base_config(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
+        repo_id = self.repo_id if self.repo_id is not tyro.MISSING else None
+        asset_id = self.assets.asset_id or repo_id
+        return dataclasses.replace(
+            self.base_config or DataConfig(),
+            repo_id=repo_id,
+            asset_id=asset_id,
+            norm_stats=self._load_norm_stats(epath.Path(self.assets.assets_dir or assets_dirs), asset_id),
+            use_quantile_norm=False,  # Always use z-score normalization for B1K
+        )
+
+    def _load_norm_stats(self, assets_dir: epath.Path, asset_id: str | None) -> dict[str, _transforms.NormStats] | None:
+        if asset_id is None:
+            return None
+        try:
+            data_assets_dir = str(assets_dir / asset_id)
+            norm_stats = _normalize.load(_download.maybe_download(data_assets_dir))
+            logging.info(f"Loaded norm stats from {data_assets_dir}")
+            return norm_stats
+        except FileNotFoundError:
+            logging.info(f"Norm stats not found in {data_assets_dir}, skipping.")
+        return None
+
+
+@dataclasses.dataclass(frozen=True)
+class LeRobotB1KDataConfig(DataConfigFactory):
+    """Data configuration for BEHAVIOR-1K dataset."""
+
+    action_sequence_keys: Sequence[str] = ("action",)
+    use_delta_joint_actions: bool = False
+    
+    # FAST auxiliary tokenization (only for PI_BEHAVIOR with use_fast_auxiliary)
+    use_fast_tokenization: bool = False
+
+    @override
+    def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
+        # Repack transforms for B1K observations
+        repack_mapping = {
+            "observation/egocentric_camera": "observation.images.rgb.head",
+            "observation/wrist_image_left": "observation.images.rgb.left_wrist",
+            "observation/wrist_image_right": "observation.images.rgb.right_wrist",
+            "observation/state": "observation.state",
+            "actions": "action",
+            "task_index": "task_index",  # Always preserve task_index
+            "timestamp": "timestamp",    # Preserve timestamp for subtask state computation
+            "episode_index": "episode_index",  # Preserve episode_index for episode length lookup
+            "index": "index",           # Preserve index
+        }
+            
+        repack_transform = _transforms.Group(
+            inputs=[_transforms.RepackTransform(repack_mapping)]
+        )
+
+        # Prepare data for policy training
+        data_transforms = _transforms.Group(
+            inputs=[b1k_policy.B1kInputs(model_type=model_config.model_type)],
+            outputs=[b1k_policy.B1kOutputs()],
+        )
+
+        # Delta action transforms
+        if self.use_delta_joint_actions:
+            delta_action_mask = _transforms.make_bool_mask(-3, 3, -1, 7, -1, 7, -1)
+        else:
+            delta_action_mask = _transforms.make_bool_mask(-23)
+        
+        data_transforms = data_transforms.push(
+            inputs=[_transforms.DeltaActions(delta_action_mask)],
+            outputs=[_transforms.AbsoluteActions(delta_action_mask)],
+        )
+
+        # Model transforms (subtask state, task ID, padding)
+        model_transforms = ModelTransformFactory()(model_config)
+        
+        # FAST tokenization (if enabled for PI_BEHAVIOR)
+        if self.use_fast_tokenization and hasattr(model_config, 'use_fast_auxiliary') and model_config.use_fast_auxiliary:
+            asset_id = self.assets.asset_id or self.repo_id
+            tokenizer_path = assets_dirs / asset_id / "fast_tokenizer"
+            
+            # Get base config to access norm_stats
+            base_config = self.create_base_config(assets_dirs, model_config)
+            
+            # Only add transform if tokenizer directory exists
+            if tokenizer_path.exists():
+                model_transforms = model_transforms.push(
+                    inputs=[b1k_transforms.TokenizeFASTActions(
+                        tokenizer_path=str(tokenizer_path),
+                        encoded_dim_ranges=model_config.get_fast_dim_ranges(),
+                        max_fast_tokens=model_config.max_fast_tokens,
+                        norm_stats=base_config.norm_stats,
+                        use_per_timestamp=base_config.use_per_timestamp_norm,
+                    )],
+                )
+            else:
+                logging.warning(
+                    f"FAST tokenizer not found at {tokenizer_path}. "
+                    "FAST auxiliary training will be disabled (inference mode)."
+                )
+
+        return dataclasses.replace(
+            self.create_base_config(assets_dirs, model_config),
+            repack_transforms=repack_transform,
+            data_transforms=data_transforms,
+            model_transforms=model_transforms,
+            action_sequence_keys=self.action_sequence_keys,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class TrainConfig:
+    # Name of the config. Must be unique. Will be used to reference this config.
+    name: tyro.conf.Suppress[str]
+    # Project name.
+    project_name: str = "B1K"
+    # Experiment name. Will be used to name the metadata and checkpoint directories.
+    exp_name: str = tyro.MISSING
+
+    # Defines the model config (PI_BEHAVIOR only for B1K).
+    model: _model.BaseModelConfig = dataclasses.field(default_factory=pi_behavior_config.PiBehaviorConfig)
+
+    # A weight loader can optionally load (possibly partial) weights from disk after the model is initialized.
+    weight_loader: weight_loaders.WeightLoader = dataclasses.field(default_factory=weight_loaders.NoOpWeightLoader)
+
+    # Note: PyTorch support removed - JAX only
+
+    lr_schedule: _optimizer.LRScheduleConfig = dataclasses.field(default_factory=_optimizer.CosineDecaySchedule)
+    optimizer: _optimizer.OptimizerConfig = dataclasses.field(default_factory=_optimizer.AdamW)
+    ema_decay: float | None = 0.99
+
+    # Specifies which weights should be frozen.
+    freeze_filter: tyro.conf.Suppress[Filter] = dataclasses.field(default_factory=nnx.Nothing)
+
+    # Determines the data to be trained on.
+    data: DataConfigFactory = dataclasses.field(default_factory=LeRobotB1KDataConfig)
+
+    # Base directory for config assets (e.g., norm stats).
+    assets_base_dir: str = "./assets"
+    # Base directory for checkpoints.
+    checkpoint_base_dir: str = "./checkpoints"
+
+    # Random seed that will be used by random generators during training.
+    seed: int | None = None
+    # Global batch size.
+    batch_size: int = 32
+    # Number of workers to use for the data loader.
+    num_workers: int = 2
+    # Number of train steps (batches) to run.
+    num_train_steps: int = 30_000
+
+    # How often (in steps) to log training metrics.
+    log_interval: int = 100
+    # How often (in steps) to save checkpoints.
+    save_interval: int = 1000
+    # If set, any existing checkpoints matching step % keep_period == 0 will not be deleted.
+    keep_period: int | None = 5000
+
+    # If true, will overwrite the checkpoint directory if it already exists.
+    overwrite: bool = False
+    # If true, will resume training from the last checkpoint.
+    resume: bool = False
+
+    # If true, will enable wandb logging.
+    wandb_enabled: bool = True
+
+    # Used to pass metadata to the policy server.
+    policy_metadata: dict[str, Any] | None = None
+
+    # FSDP configuration for model sharding across devices.
+    fsdp_devices: int = 1
+    
+    # Validation configuration
+    val_log_interval: int = 100
+    val_batch_size: int | None = None
+    val_num_batches: int = 10
+    val_repo_id: str | None = None
+    val_episodes_index: List[int] | None = None
+    
+    # Number of flow matching samples per training step
+    num_flow_samples: int = 1
+
+    @property
+    def assets_dirs(self) -> pathlib.Path:
+        """Get the assets directory for this config."""
+        return (pathlib.Path(self.assets_base_dir) / self.name).resolve()
+
+    @property
+    def checkpoint_dir(self) -> pathlib.Path:
+        """Get the checkpoint directory for this config."""
+        if not self.exp_name:
+            raise ValueError("--exp_name must be set")
+        return (pathlib.Path(self.checkpoint_base_dir) / self.name / self.exp_name).resolve()
+
+    @property
+    def trainable_filter(self) -> nnx.filterlib.Filter:
+        """Get the filter for the trainable parameters."""
+        return nnx.All(nnx.Param, nnx.Not(self.freeze_filter))
+
+    def __post_init__(self) -> None:
+        if self.resume and self.overwrite:
+            raise ValueError("Cannot resume and overwrite at the same time.")
+
+
+# B1K Training Configurations
+_CONFIGS = [
+    TrainConfig(
+        name="pi_behavior_b1k_fast",
+        exp_name="openpi",
+        project_name="B1K",
+        model=pi_behavior_config.PiBehaviorConfig(
+            action_horizon=30,
+            action_dim=32,
+            use_correlated_noise=True,
+            correlation_beta=0.5,
+            # FAST auxiliary training
+            use_fast_auxiliary=True,
+            fast_loss_weight=0.05,
+            fast_encoded_dims="0:6,7:23",  # Encode 22 dimensions
+            fast_vocab_size=1024,
+            max_fast_tokens=200,
+            use_kv_transform=True,
+            use_knowledge_insulation=False,
+            subtask_loss_weight=0.1,
+            freeze_vision_backbone=True,
+        ),
+        data=LeRobotB1KDataConfig(
+            repo_id="IliaLarchenko/behavior_224_rgb",
+            base_config=DataConfig(
+                prompt_from_task=False,  # No text prompts for PI_BEHAVIOR
+                behavior_dataset_root="~/data/behavior_224_rgb",
+                use_per_timestamp_norm=True,  # Enable per-timestamp normalization
+            ),
+            use_delta_joint_actions=True,
+            use_fast_tokenization=True,  # Enable FAST tokenization in data pipeline
+        ),
+        lr_schedule=_optimizer.CosineDecaySchedule(
+            warmup_steps=1000,
+            peak_lr=1e-4,
+            decay_steps=20_000,
+            decay_lr=1e-5,
+        ),
+        num_flow_samples=15,
+        weight_loader=weight_loaders.PiBehaviorWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
+        num_train_steps=200_000,
+        assets_base_dir="./outputs/assets",
+        checkpoint_base_dir="./outputs/checkpoints",
+        num_workers=80,
+        save_interval=500,
+        keep_period=2000,
+    ),
+    TrainConfig(
+        name="pi_behavior_b1k_stage_only",
+        exp_name="stage_only",
+        project_name="B1K",
+        model=pi_behavior_config.PiBehaviorConfig(
+            action_horizon=30,
+            action_dim=32,
+            use_correlated_noise=False,
+            correlation_beta=0.5,
+            # Keep only the task/stage conditioning path on top of pi0.5.
+            use_fast_auxiliary=False,
+            use_kv_transform=False,
+            use_knowledge_insulation=False,
+            subtask_loss_weight=0.1,
+            freeze_vision_backbone=True,
+        ),
+        data=LeRobotB1KDataConfig(
+            repo_id="IliaLarchenko/behavior_224_rgb",
+            base_config=DataConfig(
+                prompt_from_task=False,
+                behavior_dataset_root="~/data/behavior_224_rgb",
+                use_per_timestamp_norm=True,
+            ),
+            use_delta_joint_actions=True,
+            use_fast_tokenization=False,
+        ),
+        lr_schedule=_optimizer.CosineDecaySchedule(
+            warmup_steps=1000,
+            peak_lr=1e-4,
+            decay_steps=20_000,
+            decay_lr=1e-5,
+        ),
+        num_flow_samples=15,
+        weight_loader=weight_loaders.PiBehaviorWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
+        num_train_steps=200_000,
+        assets_base_dir="./outputs/assets",
+        checkpoint_base_dir="./outputs/checkpoints",
+        num_workers=80,
+        save_interval=500,
+        keep_period=2000,
+    ),
+]
+
+if len({config.name for config in _CONFIGS}) != len(_CONFIGS):
+    raise ValueError("Config names must be unique.")
+_CONFIGS_DICT = {config.name: config for config in _CONFIGS}
+
+
+def cli() -> TrainConfig:
+    return tyro.extras.overridable_config_cli({k: (k, v) for k, v in _CONFIGS_DICT.items()})
+
+
+def get_config(config_name: str) -> TrainConfig:
+    """Get a config by name."""
+    if config_name not in _CONFIGS_DICT:
+        closest = difflib.get_close_matches(config_name, _CONFIGS_DICT.keys(), n=1, cutoff=0.0)
+        closest_str = f" Did you mean '{closest[0]}'? " if closest else ""
+        raise ValueError(f"Config '{config_name}' not found.{closest_str}")
+
+    return _CONFIGS_DICT[config_name]

legacy/task12_da3_large_gtdepth_newbank_step19999/code/src/b1k/training/da3_extractor.py

@@ -0,0 +1,212 @@
+"""Inline DA3-GIANT feature extractor (PyTorch, runs in the openpi venv).
+
+Produces the SAME feats/ray/depth as the offline cache, but on-the-fly at train time — so
+no precached features are needed (essential for datasets where caching is infeasible, e.g.
+b1k at 30 Hz would be ~3 PB). Frozen, no_grad; the output tensors are handed to the JAX
+model via dlpack.
+
+depth_anything_3.api transitively imports rendering/SfM utils (moviepy/gsplat/pycolmap/
+trimesh/evo) that (a) aren't needed for feature extraction and (b) pin numpy<2 (conflicts
+with openpi's numpy 2.x). We stub those modules so nothing gets installed/downgraded.
+"""
+
+import concurrent.futures
+import contextlib
+import importlib.util
+import sys
+import types
+
+import numpy as np
+import torch
+
+# The frozen extractor is pure GPU work; torch's CPU thread pools only add GIL/scheduler contention
+# with the JAX train loop in the same process (~1000 threads observed). Pin them to 1.
+try:
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)  # only settable before any parallel work; ignore if already set
+except Exception:  # noqa: BLE001
+    pass
+
+
+def _to_dev(a, device, dtype=None):
+    """Host array/tensor -> device tensor. A torch tensor (pinned by the DataLoader when
+    B1K_TORCH_COLLATE=1) is copied with non_blocking=True so the H2D is an async DMA that overlaps
+    compute, instead of a synchronous pageable copy. Numpy input keeps the original behavior."""
+    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)
+
+_DA3_SRC = "/work/jack/projects/Depth-Anything-3/src"
+_GEOSTACK = "/work/jack/da3xvla_src/DA3-XVLA-cache/models/da3_for_geostack.py"
+_IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 1, 3, 1, 1)
+_IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 1, 3, 1, 1)
+
+
+class _AutoStub(types.ModuleType):
+    """Module stub that returns a dummy callable for any non-dunder attribute + acts as a package."""
+
+    __path__: list = []
+
+    def __getattr__(self, name):
+        if name.startswith("__"):
+            raise AttributeError(name)
+        return lambda *a, **k: None
+
+
+def _install_stubs():
+    for name in ("moviepy", "moviepy.editor", "gsplat", "pycolmap", "trimesh",
+                 "depth_anything_3.utils.export", "depth_anything_3.utils.pose_align"):
+        sys.modules.setdefault(name, _AutoStub(name))
+    if hasattr(sys.modules["moviepy"], "__dict__"):
+        sys.modules["moviepy"].editor = sys.modules["moviepy.editor"]
+    if _DA3_SRC not in sys.path:
+        sys.path.insert(0, _DA3_SRC)
+
+
+def _load_da3_class():
+    _install_stubs()
+    spec = importlib.util.spec_from_file_location("_da3_for_geostack", _GEOSTACK)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod.DA3LargeForGeoStack
+
+
+def rescale_intrinsics(intr: np.ndarray, src_hw, dst_hw) -> np.ndarray:
+    """Rescale pixel K [.,3,3] from src (H,W) to dst (H,W). fx,cx by W-ratio; fy,cy by H-ratio."""
+    sh, sw = src_hw
+    dh, dw = dst_hw
+    rw, rh = dw / sw, dh / sh
+    out = np.array(intr, dtype=np.float32, copy=True)
+    out[..., 0, 0] *= rw
+    out[..., 0, 2] *= rw
+    out[..., 1, 1] *= rh
+    out[..., 1, 2] *= rh
+    return out
+
+
+class DA3InlineExtractor:
+    """Frozen DA3-GIANT posed multi-view extractor, REPLICATED one-per-GPU.
+
+    Each visible GPU holds its own frozen DA3-GIANT copy and runs the forward on ONLY its slice of the
+    batch, all GPUs concurrently (one thread per device; CUDA kernels are async per device, and torch
+    releases the GIL during them). This removes the single-GPU serial bottleneck of the old design, so
+    inline extraction scales with GPU count to match the JAX data-parallel training step.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "depth-anything/DA3NESTED-GIANT-LARGE-1.1",
+        out_layers=(19, 26, 33, 39),
+        da3_hw=(252, 336),
+        devices=None,
+        forward_chunk: int = 16,
+    ):
+        DA3 = _load_da3_class()
+        if devices is None:
+            # Default SINGLE-GPU (cuda:0): the multi-replica path is correct but thread-based, and the
+            # Python GIL serializes the DA3 forward's many kernel launches (~18% concurrency efficiency),
+            # so replicating across GPUs doesn't speed it up — it only wastes memory. Pass `devices`
+            # explicitly (e.g. for a future multiprocess extractor) to override.
+            devices = ["cuda:0" if torch.cuda.is_available() else "cpu"]
+        self.devices = list(devices)
+        self.da3_hw = tuple(da3_hw)
+        self.forward_chunk = int(forward_chunk)
+        self.replicas = []
+        for dev in self.devices:
+            m = (
+                DA3(model_name=model_name, out_layers=tuple(out_layers),
+                    da3_input_h=da3_hw[0], da3_input_w=da3_hw[1], patch_size=14, use_bf16=True)
+                .to(dev)
+                .eval()
+            )
+            for p in m.parameters():
+                p.requires_grad_(False)
+            self.replicas.append(m)
+        self._pool = concurrent.futures.ThreadPoolExecutor(max_workers=max(1, len(self.devices)))
+
+    def _preprocess(self, images: np.ndarray, device) -> torch.Tensor:
+        """images: [B,V,H,W,3] uint8 (or float [0,1]) -> [B,V,3,252,336] ImageNet-normalized on `device`."""
+        x = _to_dev(images, device)
+        if x.dtype == torch.uint8:
+            x = x.float() / 255.0
+        elif x.max() > 1.5:  # already float but in [0,255]
+            x = x.float() / 255.0
+        x = x.permute(0, 1, 4, 2, 3)  # [B,V,3,H,W]
+        x = torch.nn.functional.interpolate(
+            x.flatten(0, 1), size=self.da3_hw, mode="bicubic", align_corners=False, antialias=True
+        ).view(*x.shape[:2], 3, *self.da3_hw)
+        x = (x - _IMAGENET_MEAN.to(x)) / _IMAGENET_STD.to(x)
+        return x
+
+    def _run_shard(self, di: int, images: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray,
+                   return_torch: bool = False):
+        """Run replica `di` over its (pre-sliced) shard, in chunks to bound activation memory.
+        return_torch=True keeps the result as torch tensors ON cuda:di (no host round-trip) for the
+        DLPack handoff; default returns numpy (host) as before."""
+        dev = self.devices[di]
+        dev_idx = int(dev.split(":")[1]) if ":" in dev else None
+        replica = self.replicas[di]
+        chunk = self.forward_chunk
+        fc, rc, dc = [], [], []
+        ctx = torch.cuda.device(dev_idx) if dev_idx is not None else contextlib.nullcontext()
+        with ctx, torch.no_grad():  # set current device so implicit-device tensors land on the right GPU
+            for i in range(0, images.shape[0], chunk):
+                x = self._preprocess(images[i : i + chunk], dev)
+                e = _to_dev(extrinsics[i : i + chunk], dev, torch.float32)
+                k = _to_dev(intrinsics[i : i + chunk], dev, torch.float32)
+                out = replica.forward_multi_view(x, extrinsics=e, intrinsics=k)
+                # Ship feats as bf16 BITS (uint16): the model casts to bf16 anyway (see
+                # Pi0._compute_banks), so this is numerically identical to shipping f32 while
+                # halving the GPU->CPU->GPU transfer and the collate copies.
+                feats = torch.stack(list(out["feats"]), dim=1).to(torch.bfloat16)  # [b,4,V,C,h,w]
+                ray = out["ray"].float()  # [b,V,3,h,w]
+                depth = out["depth"]
+                if depth is None:
+                    depth = torch.zeros(ray.shape[0], ray.shape[1], 1, ray.shape[3], ray.shape[4], device=ray.device)
+                if return_torch:
+                    fc.append(feats.view(torch.uint16)); rc.append(ray); dc.append(depth.float())
+                else:
+                    fc.append(feats.view(torch.uint16).cpu().numpy())
+                    rc.append(ray.cpu().numpy())
+                    dc.append(depth.float().cpu().numpy())
+        if return_torch:
+            return torch.cat(fc, 0), torch.cat(rc, 0), torch.cat(dc, 0)  # torch tensors on cuda:di
+        return np.concatenate(fc, 0), np.concatenate(rc, 0), np.concatenate(dc, 0)
+
+    def extract_shards_torch(self, images: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray):
+        """Like extract() but returns per-shard torch GPU tensors (feats,ray,depth) each on cuda:di,
+        WITHOUT the host round-trip. The caller hands them to JAX via DLPack (GPU->GPU over NVLink).
+        Returns a list of len(devices) tuples, shard k on self.devices[k]."""
+        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], True)
+        return [futs[di].result() for di in sorted(futs)]
+
+    def extract(self, images: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray):
+        """images [B,V,H,W,3]; extrinsics [B,V,4,4] w2c; intrinsics [B,V,3,3] AT 252x336.
+
+        Splits the batch across all replicas/GPUs and runs the forwards concurrently.
+        Returns numpy: feats [B,4,V,1536,18,24] f32, ray [B,V,3,18,24] f32, depth [B,V,1,18,24] f32.
+        """
+        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])
+        parts = [futs[di].result() for di in sorted(futs)]
+        return (
+            np.concatenate([p[0] for p in parts], axis=0),
+            np.concatenate([p[1] for p in parts], axis=0),
+            np.concatenate([p[2] for p in parts], axis=0),
+        )

legacy/task12_da3_large_gtdepth_newbank_step19999/code/src/b1k/training/weight_loaders.py

@@ -0,0 +1,88 @@
+"""Weight loaders for PI_BEHAVIOR model initialization from Pi05 checkpoints.
+
+Reference: https://github.com/Physical-Intelligence
+"""
+
+import dataclasses
+import logging
+import re
+
+import flax.traverse_util
+import numpy as np
+
+import openpi.models.model as _model
+import openpi.shared.array_typing as at
+import openpi.shared.download as download
+
+# Re-export base loaders from OpenPI
+from openpi.training.weight_loaders import (
+    WeightLoader,
+    NoOpWeightLoader,
+    CheckpointWeightLoader,
+    _merge_params,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass(frozen=True)
+class PiBehaviorWeightLoader(WeightLoader):
+    """Loads checkpoints for PI_BEHAVIOR model.
+    
+    Automatically detects:
+    - Pi05 checkpoint: Loads weights, preserves new PI_BEHAVIOR parameters
+    - PI_BEHAVIOR checkpoint: Loads all weights directly
+    """
+
+    params_path: str
+
+    def load(self, params: at.Params) -> at.Params:
+        # Load checkpoint as host arrays so it can be initialized onto the
+        # current device topology, even if the released checkpoint used more GPUs.
+        params_path = download.maybe_download(self.params_path)
+        restored = _model.restore_params(params_path, restore_type=np.ndarray)
+        
+        # Handle nested 'params' key (from some checkpoint formats)
+        if isinstance(restored, dict) and "params" in restored:
+            loaded_params = restored["params"]
+        else:
+            loaded_params = restored
+        
+        # Remove 'value' suffixes (from nnx.State format)
+        flat_params = flax.traverse_util.flatten_dict(loaded_params)
+        if all(kp[-1] == "value" for kp in flat_params if len(kp) > 0):
+            flat_params = {kp[:-1]: v for kp, v in flat_params.items()}
+            loaded_params = flax.traverse_util.unflatten_dict(flat_params)
+        
+        # Detect checkpoint type
+        has_task_embeddings = 'task_embeddings' in loaded_params
+        
+        if has_task_embeddings:
+            # Loading PI_BEHAVIOR checkpoint. New optional adapters are allowed
+            # to remain at their initialized values when absent from older checkpoints.
+            logging.info("Loading PI_BEHAVIOR checkpoint")
+            missing_regex = ".*spatial_action_xattn.*|.*spatial_bank_builder.*|.*spatial_inject.*"
+            return _merge_params(loaded_params, params, missing_regex=missing_regex)
+        else:
+            # Loading Pi05 checkpoint - preserve new PI_BEHAVIOR-specific parameters
+            logging.info("Loading Pi05 checkpoint (new PI_BEHAVIOR parameters will use random init)")
+            
+            # These parameters are NEW in PI_BEHAVIOR (not in Pi05), so keep them from params (random init)
+            missing_regex = (
+                ".*task_embeddings.*|"
+                ".*stage_pred_from_vlm.*|"
+                ".*task_stage_embeddings.*|"
+                ".*gate_sincos.*|"
+                ".*gate_task_stage.*|"
+                ".*gate_task.*|"
+                ".*fusion_layer.*|"
+                ".*stage_projection.*|"
+                ".*task_subtask_fusion.*|"
+                ".*fast_token_embedding.*|"
+                ".*fast_token_proj.*|"
+                ".*kv_transform.*|"
+                ".*spatial_action_xattn.*|"
+                ".*spatial_bank_builder.*|"
+                ".*spatial_inject.*"
+            )
+            return _merge_params(loaded_params, params, missing_regex=missing_regex)

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/2608e0b98840e07c0bddb547c83352b5

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/26312df420be0ded6f7d2968b9eb31e4

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/32e999b158e705cc43e62af68ed2c026

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/35232fd09a73f92d9d541123f677d7b6

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/589f02b3a7ab6adaca31058554766308

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/58cb578d8033aedd34caa72dbd8345ae

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/5d963bbe5fe91155634586f828e91671

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/60c607b33450be6346df3980b439ad2c

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/742cb389e1ad457ef05645486d8511fa

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/7de207fe9d87adc11cbbd47a23474c55

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

legacy/task12_da3_large_gtdepth_newbank_step19999/params/ocdbt.process_0/d/8e22384919e66ee762de8aa29a643f1e

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