Add source_code/ training snapshot + depth/ray provenance + reference distributions

pibehavior_da3_clean_up_your_desk_40k/source_code/README.md

@@ -0,0 +1,107 @@
+# source_code — training snapshot for `pibehavior_da3_clean_up_your_desk_40k`
+
+Exact source used to produce the checkpoint, so eval doesn't have to reconstruct architecture from
+Orbax metadata. **Read the "Answers" section first — it settles the train/eval mismatch questions.**
+
+---
+
+## Answers to the questions that matter
+
+### 1. Where do `da3_depth` and `da3_ray` come from? (the critical one)
+**Both are DA3 model predictions. Neither comes from the simulator.**
+
+- **`da3_depth` = DA3's own predicted METRIC depth.** From the Nested-GIANT `da3_metric` branch
+  (`out["depth"]`, `out["is_metric"]`) — see `da3/da3_for_geostack.py`.
+- **`da3_ray` = DA3's ray-head output**, captured by a forward hook on the final aux conv of DA3's
+  DualDPT. Direction-only, unit-normalized (verified: per-pixel L2 norm = 1.0000).
+- **The simulator's `depth_linear` is NEVER fed to the model.** It was used *only offline*, to
+  calibrate the extrinsics convention (that is how the GL2CV flip + wxyz + camera-in-robot ordering
+  were derived). **Do not feed `depth_linear` at eval** — run the same DA3 forward and use its outputs.
+- `use_ray_pose=False`: the extrinsics/intrinsics drive the **camera encoder (`cam_enc`)** that
+  conditions the features, *not* the ray head. This is why wrong poses corrupt the **features**
+  themselves rather than merely an auxiliary channel.
+
+⚠️ **Silent-zeros trap.** `da3_extractor.py` contains:
+```python
+depth = out["depth"]
+if depth is None:
+    depth = torch.zeros(...)     # silently all-zeros
+```
+If your wrapper fails to surface depth you get **zeros with no error**. We verified this fallback was
+**NOT** active in training (`ALL_ZERO=False`, `frac_zero=0.0000`). Assert on this at eval.
+
+### 2. Reference distributions — validate your eval against these
+Measured on **real b1k frames with real `robot2cam_pose`**, batch of 4, `clean_up_your_desk`:
+
+```
+da3_features (B,4,3,1536,18,18) uint16   # uint16 = bf16 BITS, not floats
+da3_ray      (B,3,3,18,18) float32   L2 norm per pixel = 1.0000
+da3_depth    (B,3,1,18,18) float32
+   min=0.2110  max=2.0765  mean=0.9492  std=0.3715
+   p1=0.309  p25=0.617  p50=0.979  p75=1.244  p99=1.782
+per-view depth mean:  head 1.310 (0.81–2.08) | left_wrist 0.723 (0.21–1.48) | right_wrist 0.815 (0.37–1.59)
+```
+**Sanity checks for your eval:**
+- depth must be **non-zero** and roughly **0.2–2.1 m**
+- **head depth > wrist depth** (head views the scene; wrists sit near manipulated objects).
+  If this ordering inverts, your **view order is wrong** (must be `head, left_wrist, right_wrist`).
+- ray L2 norm must be **1.0**. If not, you're feeding the wrong tensor.
+- On structureless input DA3 degenerates to near-constant depth (~0.93, std 0.006) — if you see that
+  on real frames, images aren't reaching the extractor.
+
+### 3. Bank tokens are 320 = 128 + 96 + 96
+`pi_behavior_config.py`: `spatial_num_tokens: int = 320  # DA3 perc bank default: 128 + 96 + 96`.
+(Not 32/24/24.)
+
+### 4. Feature wire format
+bf16 bit-packed into uint16 (torch and JAX share no bf16 interchange dtype):
+```python
+# torch: feats.to(torch.bfloat16).view(torch.uint16)
+# jax:   jax.lax.bitcast_convert_type(feats, jnp.bfloat16)
+```
+**Do not cast to float32** — the model branches on dtype (see `PiBehavior._compute_banks`).
+
+---
+
+## Contents
+
+```
+model/    pi_behavior.py            policy + _compute_banks + spatial action conditioning
+          pi_behavior_config.py     B1KDA3Config (spatial_num_tokens=320, init_std, logit_gain...)
+          spatial_da3.py            bank builder + cross-attention injection
+          observation.py            Observation dataclass (da3_features / da3_ray / da3_depth fields)
+training/ da3_extractor.py          DA3InlineExtractor  <-- produces the DA3 tensors
+          b1k_da3.py                dataset + da3_fields() (extrinsics/intrinsics) + batch_transform
+          b1k_2026.py               BehaviorV3Dataset (video decode)
+          train.py                  multi-group LR (vlm/core/geom) + train step
+          train_2026.py             entrypoint / env -> config
+          _da3_contrib_check.py     bank-ablation sanity check (RUN THIS, see below)
+openpi/   gemma.py                  spatial injection into last 6 of 18 action-expert layers
+          data_loader.py            collator + DLPack/async staging
+          optimizer.py              LR schedules
+          sharding.py               B1K_HOLDOUT_GPU0 mesh holdout
+          weight_loaders.py         PiBehaviorWeightLoader (param key merge)
+da3/      da3_for_geostack.py       DA3 wrapper: forward_multi_view, ray hook, depth surfacing
+launch/   launch_b1k_robopro.sh     EXACT env/hyperparameters that produced this checkpoint
+```
+
+**DA3 upstream:** `Depth-Anything-3` @ commit **`4173623`**, sources at `<repo>/src`, loaded via
+`_DA3_SRC` in `da3_extractor.py`. Weights: `depth-anything/DA3NESTED-GIANT-LARGE-1.1`,
+`out_layers=(19,26,33,39)`, `patch_size=14`, `use_bf16=True`, frozen (`no_grad()` + `eval()`).
+
+## Verify your integration before trusting results
+Run the **bank ablation** (`training/_da3_contrib_check.py`): compute action loss normally, then with
+the spatial banks zeroed. On this checkpoint the gap is large (at 14k it was 0.301 → 4.300, **+1329%**).
+**A near-zero gap means the spatial branch isn't actually engaged** — and note the model *silently
+falls back to base behavior* when `da3_features` is `None`, so a wiring bug looks like "DA3 doesn't
+help", not like a crash.
+
+## Training configuration (as shipped in `launch/`)
+```
+BS=128 FLOW=15 STEPS=50000 warmup=1000 decay=50000
+LR: vlm 2.5e-5->2.5e-6 | core 5e-4->5e-5 | geom 5e-4->1e-4   (DA3_LR_GROUPS=1)
+DA3: spatial_init_std=0.01  attn_logit_gain=1 (init 32.0)  spatial_scale=2.0
+da3_hw=(252,252) -> 18x18 grid; VLM images 224x224; views: head, left_wrist, right_wrist
+```
+`spatial_init_std` is deliberately **non-zero**: with zero-init the model learned to *ignore* the
+spatial branch entirely.

pibehavior_da3_clean_up_your_desk_40k/source_code/da3/da3_for_geostack.py

@@ -0,0 +1,639 @@
+# SPDX: same as DA3-XVLA repo
+"""
+DA3-Large frozen feature extractor for GeoStack-XVLA v2B.
+
+This wrapper around `depth_anything_3.api.DepthAnything3`:
+  * Loads DA3-Large-1.1 and freezes ALL parameters.
+  * Receives ImageNet-normalized RGB at any input dimensions (typically the
+    already-square 224×224 that the X-VLA processor produces for Florence).
+  * Bilinearly resizes to the configured (da3_input_h, da3_input_w) — typically
+    a 4:3-aspect rectangle like 252×336 that recovers the original aspect after
+    Florence's stretch. NO PADDING — both dims must be multiples of patch_size.
+  * Returns a list of N feature levels at the configured DPT taps (default
+    out_layers=[11, 19, 23] for DA3-Large = shallow, mid, deep), each as
+    [B, C, h_grid, w_grid] with (h_grid, w_grid) = (input_h/14, input_w/14).
+  * Returns the DA3 ray-head output (3D direction unit vector per pixel),
+    captured via a forward hook on `head.scratch.output_conv2_aux[-1]` and
+    resampled to (h_grid, w_grid).
+
+Spatial-bias coordinate alignment:
+  * Florence sees a stretched 224×224 → its token-grid (e.g. 7×7) maps linearly
+    back to the ORIGINAL image's normalized [0,1]² coords (no aspect correction
+    needed; the stretch is bijective per axis).
+  * DA3 sees the same content but at the aspect-correct dims (e.g. 252×336)
+    → its token-grid (e.g. 18×24) ALSO maps linearly to the ORIGINAL image's
+    normalized [0,1]² coords.
+  * Both grids therefore live in the same normalized coord space, and the
+    spatial-distance bias `-λ · ||pos_vlm - pos_da3||²` is well-defined.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Tuple
+
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+# DA3-Large fallback (used if hub model has no `backbone.out_layers` attribute).
+_DA3_LARGE_OUT_LAYERS = (11, 15, 19, 23)
+_DA3_LARGE_EMBED_DIM = 1024
+
+
+class DA3LargeForGeoStack(nn.Module):
+    """Frozen DA3-Large feature + ray-head extractor for GeoStack.
+
+    Usage:
+        da3 = DA3LargeForGeoStack(
+            model_name="depth-anything/DA3-Large-1.1",
+            out_layers=(11, 19, 23),
+            da3_input_h=252, da3_input_w=336,   # aspect-correct 4:3 (~1.05× of 240×320)
+        )
+        # pixel_values: [B, 3, H_in, W_in] in ImageNet-normalized space
+        # (typically H_in=W_in=224 — the X-VLA processor's stretched square).
+        out = da3(pixel_values)
+        # out["feats"]:        list[N] of [B, C, h_grid, w_grid]
+        # out["ray"]:          [B, 3, h_grid, w_grid]  (direction-only, unit-normalized)
+        # out["h_grid"]:       int
+        # out["w_grid"]:       int
+    """
+
+    def __init__(
+        self,
+        model_name: str = "depth-anything/DA3-Large-1.1",
+        out_layers: Tuple[int, ...] = (11, 19, 23),
+        da3_input_h: int = 252,
+        da3_input_w: int = 336,
+        patch_size: int = 14,
+        use_bf16: bool = True,
+    ):
+        super().__init__()
+        try:
+            from depth_anything_3.api import DepthAnything3
+        except Exception as exc:
+            raise ImportError(
+                "depth_anything_3 not installed; install third_party/Depth-Anything-3"
+            ) from exc
+
+        self.model_name = str(model_name)
+        self.out_layers = tuple(int(i) for i in out_layers)
+        self.patch_size = int(patch_size)
+        self.use_bf16 = bool(use_bf16)
+        self.da3_input_h = int(da3_input_h)
+        self.da3_input_w = int(da3_input_w)
+        if self.da3_input_h % self.patch_size != 0 or self.da3_input_w % self.patch_size != 0:
+            raise ValueError(
+                f"DA3 input dims must be multiples of patch_size={self.patch_size}; "
+                f"got ({self.da3_input_h}, {self.da3_input_w})"
+            )
+        self.h_grid = self.da3_input_h // self.patch_size      # e.g. 252/14 = 18
+        self.w_grid = self.da3_input_w // self.patch_size      # e.g. 336/14 = 24
+
+        # Load DA3 + freeze
+        self.model = DepthAnything3.from_pretrained(self.model_name)
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+        self.model.eval()
+
+        # Channel dim: introspect the loaded backbone rather than assuming
+        # DA3-Large. DA3-Large=1024, DA3NESTED-GIANT-LARGE-1.1=1536, etc.
+        # Downstream perceiver/projection is sized from self.embed_dim, so this
+        # MUST match the real feature width or K/V dims mismatch.
+        self.embed_dim = self._infer_embed_dim()
+
+        # Ray-head hook capture (after DA3's DualDPT output_conv2_aux).
+        self._ray_capture: Dict[str, torch.Tensor] = {}
+        self._register_ray_hook()
+
+        # v4.1 OPT-IN: capture depth output (main head) via a second hook on
+        # head.scratch.output_conv2. Off by default — only enabled when the
+        # depth-distillation aux head is configured.
+        self._depth_capture: Dict[str, torch.Tensor] = {}
+        self._depth_hook = None
+        self._capture_depth = False
+
+    # ---------- backbone feature-dim introspection (GIANT-safe) ----------
+
+    def _infer_embed_dim(self) -> int:
+        """Read the DINOv2 backbone hidden width from the loaded DA3 model.
+
+        Walks common attribute paths to a `.embed_dim`/`.num_features`; falls
+        back to inspecting a transformer block's LayerNorm, then to the
+        DA3-Large constant. DA3-Large=1024, DA3NESTED-GIANT-LARGE-1.1=1536.
+        """
+        obj = self.model
+        for path in (("model", "da3", "backbone", "pretrained"),
+                     ("da3", "backbone", "pretrained"),
+                     ("backbone", "pretrained"),
+                     ("backbone",)):
+            cur = obj
+            ok = True
+            for a in path:
+                if hasattr(cur, a):
+                    cur = getattr(cur, a)
+                else:
+                    ok = False
+                    break
+            if ok:
+                for a in ("embed_dim", "num_features", "hidden_size", "n_embd"):
+                    if hasattr(cur, a):
+                        try:
+                            return int(getattr(cur, a))
+                        except Exception:
+                            pass
+        # last-ditch: infer from a block's norm weight length
+        for n, p in self.model.named_parameters():
+            if "blocks.0." in n and n.endswith("norm1.weight") and p.ndim == 1:
+                return int(p.shape[0])
+        logger.warning("[da3_for_geostack] could not introspect embed_dim; "
+                       "falling back to DA3-Large=%d", int(_DA3_LARGE_EMBED_DIM))
+        return int(_DA3_LARGE_EMBED_DIM)
+
+    # ---------- post-load reload of DA3 weights ----------
+
+    def reload_pretrained_weights(self) -> None:
+        """Reload DA3-Large pretrained weights AFTER ``XVLA.from_pretrained``.
+
+        HF's ``from_pretrained`` treats every key absent from the loaded
+        checkpoint as a "missing" param and overwrites it via ``_init_weights``.
+        That zeroes/random-inits the DA3-Large weights we loaded inside
+        ``__init__``, producing garbage geometry features as soon as alpha > 0.
+
+        Call this once right after ``XVLA.from_pretrained(...)`` to restore the
+        actual DA3 weights. Same pattern as ``DA3InlineEncoder.reload_pretrained_weights``.
+        """
+        try:
+            from depth_anything_3.api import DepthAnything3
+        except Exception as exc:
+            raise ImportError("depth_anything_3 not installed") from exc
+        device = next(self.model.parameters()).device
+        dtype = next(self.model.parameters()).dtype
+        fresh = DepthAnything3.from_pretrained(self.model_name)
+        self.model.load_state_dict(fresh.state_dict(), strict=True)
+        self.model.to(device=device, dtype=dtype)
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+        self.model.eval()
+        # Re-attach the ray hook defensively (load_state_dict copies weights
+        # into existing modules so the hook target is preserved, but if any
+        # downstream change re-creates head modules this re-attach saves us).
+        try:
+            if hasattr(self, "_ray_hook"):
+                self._ray_hook.remove()
+        except Exception:
+            pass
+        self._register_ray_hook()
+
+    # ---------- head-path resolution (DA3-Large vs NESTED) ----------
+
+    def _find_head_scratch(self):
+        """Return the DPT ``head.scratch`` module across DA3 variants.
+
+        DA3-Large: ``self.model.model.head.scratch``.
+        DA3NESTED-GIANT-LARGE: the encoder is wrapped one level deeper as
+        ``self.model.model.da3.head.scratch`` (net children = da3 / da3_metric).
+        """
+        for path in (("model", "head", "scratch"),
+                     ("model", "da3", "head", "scratch"),
+                     ("da3", "head", "scratch"),
+                     ("head", "scratch")):
+            cur = self.model
+            ok = True
+            for a in path:
+                if hasattr(cur, a):
+                    cur = getattr(cur, a)
+                else:
+                    ok = False
+                    break
+            if ok:
+                return cur
+        return None
+
+    # ---------- ray-head hook ----------
+
+    def _register_ray_hook(self) -> None:
+        """Capture the ray-head pre-activation output via a forward hook on the
+        final aux conv in DA3's DualDPT. Output shape: [B*V, 7, H_ray, W_ray]
+        where channels are [dir_x, dir_y, dir_z, ori_x, ori_y, ori_z, conf].
+
+        DA3's ``head.scratch.output_conv2_aux`` may be an nn.Sequential or a
+        single Conv2d — hook whichever leaf module is at the tail.
+        """
+        scratch = self._find_head_scratch()
+        if scratch is None or not hasattr(scratch, "output_conv2_aux"):
+            raise RuntimeError(
+                "DA3LargeForGeoStack: could not locate head.scratch.output_conv2_aux "
+                "(this DA3 variant may not have a DualDPT head)"
+            )
+        aux = scratch.output_conv2_aux
+        # DA3's DualDPT defines output_conv2_aux as an nn.ModuleList of length
+        # aux_pyramid_levels (default 4). The forward call site uses
+        # `scratch.output_conv2_aux[-1](last_aux)` (dualdpt.py:255) — i.e., only
+        # the LAST module is actually called. Hook it.
+        if isinstance(aux, (nn.Sequential, nn.ModuleList)):
+            if len(aux) == 0:
+                raise RuntimeError("DA3LargeForGeoStack: output_conv2_aux is empty")
+            target = aux[-1]
+        else:
+            target = aux
+
+        def _hook(module, inputs, output):
+            self._ray_capture["ray_raw"] = output
+
+        self._ray_hook = target.register_forward_hook(_hook)
+
+    # ---------- depth-head hook (v4.1 opt-in) ----------
+
+    def enable_depth_capture(self, enabled: bool = True) -> None:
+        """Turn on capture of DA3's main depth output (from output_conv2).
+        Idempotent — safe to call multiple times.
+        """
+        if enabled and self._depth_hook is None:
+            scratch = self._find_head_scratch()
+            if scratch is None or not hasattr(scratch, "output_conv2"):
+                raise RuntimeError(
+                    "DA3LargeForGeoStack: could not locate head.scratch.output_conv2 for depth hook"
+                )
+            target = scratch.output_conv2
+            def _hook(module, inputs, output):
+                self._depth_capture["depth_raw"] = output
+            self._depth_hook = target.register_forward_hook(_hook)
+        elif not enabled and self._depth_hook is not None:
+            self._depth_hook.remove()
+            self._depth_hook = None
+        self._capture_depth = bool(enabled)
+
+    # ---------- forward ----------
+
+    @torch.no_grad()
+    def _da3_forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Run DA3 inner forward and return the aux dict.
+
+        x: [B, V, 3, H, W] in ImageNet-normalized space (already resized to
+        (da3_input_h, da3_input_w)).
+        """
+        inner = self.model.model
+        out = inner(
+            x, None, None,                          # no extrinsics/intrinsics
+            list(self.out_layers),                  # export these feat layers
+            False,                                  # infer_gs
+            False,                                  # use_ray_pose
+            "saddle_balanced",                       # ref_view_strategy
+        )
+        if isinstance(out, dict):
+            aux = out.get("aux") or {}
+        else:
+            aux = getattr(out, "aux", None) or {}
+        return aux
+
+    @torch.no_grad()
+    def _da3_forward_posed(
+        self,
+        x: torch.Tensor,          # [B, V, 3, H, W]
+        extrinsics: torch.Tensor,  # [B, V, 4, 4] OpenCV world-to-camera
+        intrinsics: torch.Tensor,  # [B, V, 3, 3] rescaled to (H, W)
+    ) -> Dict[str, torch.Tensor]:
+        """Joint posed multi-view DA3 forward: ALL views in one inner() call
+        with camera extrinsics + intrinsics, so DA3's camera encoder fuses pose
+        tokens into the backbone and returns per-view features that are
+        geometrically consistent ACROSS views. Mirrors da3_inline's posed call
+        (use_ray_pose=False; poses drive cam_enc, not the ray head)."""
+        inner = self.model.model
+        out = inner(
+            x,
+            extrinsics.to(device=x.device, dtype=torch.float32),
+            intrinsics.to(device=x.device, dtype=torch.float32),
+            list(self.out_layers),                  # export these feat layers
+            False,                                  # infer_gs
+            False,                                  # use_ray_pose
+            "saddle_balanced",                       # ref_view_strategy
+        )
+        # Nested GIANT returns a top-level metric-depth map (da3_metric branch):
+        # out["depth"] = [B,V,H,W], out["is_metric"]. Surface it alongside aux so
+        # the spatial tokens can carry ray + DEPTH (scale-aware ray).
+        if isinstance(out, dict):
+            aux = out.get("aux") or {}
+            depth = out.get("depth", None)
+        else:
+            aux = getattr(out, "aux", None) or {}
+            depth = getattr(out, "depth", None)
+        return aux, depth
+
+    @staticmethod
+    def _rescale_intrinsics(
+        K: torch.Tensor, src_hw: Tuple[int, int], dst_hw: Tuple[int, int]
+    ) -> torch.Tensor:
+        """Rescale pixel intrinsics [..., 3, 3] from (src_H, src_W) to
+        (dst_H, dst_W). fx, cx scale with the width ratio; fy, cy with height.
+        Matches da3_inline._batched_call so the pose math sees K aligned to the
+        actual tensor DA3 receives. Identity when src == dst."""
+        sh, sw = int(src_hw[0]), int(src_hw[1])
+        dh, dw = int(dst_hw[0]), int(dst_hw[1])
+        if (sh, sw) == (dh, dw):
+            return K.to(torch.float32)
+        sx = float(dw) / float(sw)   # width scale
+        sy = float(dh) / float(sh)   # height scale
+        K = K.clone().to(torch.float32)
+        K[..., 0, 0] *= sx   # fx
+        K[..., 1, 1] *= sy   # fy
+        K[..., 0, 2] *= sx   # cx
+        K[..., 1, 2] *= sy   # cy
+        return K
+
+    def _to_bchw(self, feat: torch.Tensor, B: int) -> torch.Tensor:
+        """Normalize an aux entry to [B*V, C, h, w]. V=1 for GeoStack."""
+        if feat.dim() == 6:                  # [B, V, s, h, w, C]
+            feat = feat.mean(dim=2)
+        if feat.dim() == 5:                  # [B, V, h, w, C]
+            return feat.permute(0, 1, 4, 2, 3).reshape(-1, feat.shape[4], feat.shape[2], feat.shape[3])
+        if feat.dim() == 4:                  # [B, V, N, C] or [B*V, C, h, w]
+            if feat.shape[0] == B:           # heuristic: [B, V, N, C]
+                B_, V_, N, C_ = feat.shape
+                s = max(1, int(round(N ** 0.5)))
+                return feat[:, :, : s * s].permute(0, 1, 3, 2).reshape(B_ * V_, C_, s, s)
+            return feat                       # already [B*V, C, h, w]
+        raise RuntimeError(f"unexpected DA3 feat tensor rank: {feat.dim()}, shape={tuple(feat.shape)}")
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        target_input_hw: Tuple[int, int] | None = None,
+    ) -> Dict[str, object]:
+        """Run DA3 forward on a batch with aspect-correct resize.
+
+        Args:
+            pixel_values: [B, 3, H_in, W_in] ImageNet-normalized float. Any input
+                resolution accepted; bilinearly resized to (h, w) before DA3 forward.
+            target_input_hw: optional (h, w) override for DA3 input size. When set,
+                this overrides the configured (da3_input_h, da3_input_w) — useful
+                for running the same DA3 at different resolutions per view
+                (e.g., main camera at 252×336, wrist at 168×224). Both dims must
+                be multiples of patch_size. The returned feats and ray use a
+                token grid derived from the target dims, NOT the configured ones.
+
+        Returns:
+            dict with keys:
+                feats:  List[torch.Tensor]  one per out_layer, each [B, C, h_grid, w_grid]
+                ray:    torch.Tensor        [B, 3, h_grid, w_grid], unit-normalized direction
+                h_grid: int
+                w_grid: int
+        """
+        if pixel_values.dim() != 4:
+            raise ValueError(f"pixel_values must be [B, 3, H, W]; got shape {tuple(pixel_values.shape)}")
+        B = pixel_values.shape[0]
+        device = pixel_values.device
+        x = pixel_values.to(torch.float32)
+
+        # Resolve input dims + corresponding token grid.
+        if target_input_hw is None:
+            in_h, in_w = self.da3_input_h, self.da3_input_w
+            h_grid, w_grid = self.h_grid, self.w_grid
+        else:
+            in_h, in_w = int(target_input_hw[0]), int(target_input_hw[1])
+            if in_h % self.patch_size != 0 or in_w % self.patch_size != 0:
+                raise ValueError(
+                    f"target_input_hw=({in_h},{in_w}) must be multiples of patch_size={self.patch_size}"
+                )
+            h_grid, w_grid = in_h // self.patch_size, in_w // self.patch_size
+
+        # Aspect-correct resize (bilinear). NO PADDING — input dims chosen so
+        # both are multiples of patch_size.
+        x = F.interpolate(x, size=(in_h, in_w), mode="bilinear", align_corners=False)
+        # Expand to V=1 view dim (DA3 expects [B, V, 3, H, W])
+        x = x.unsqueeze(1)
+
+        # Clear stale ray capture (and depth if enabled), run DA3
+        self._ray_capture.clear()
+        if self._capture_depth:
+            self._depth_capture.clear()
+        if self.use_bf16:
+            with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+                aux = self._da3_forward(x)
+        else:
+            aux = self._da3_forward(x)
+
+        feats, ray_at_feat, depth_at_feat = self._finalize(aux, B, h_grid, w_grid)
+        return {
+            "feats": feats,                # list of [B, C, h_grid, w_grid]
+            "ray": ray_at_feat,            # [B, 3, h_grid, w_grid]
+            "depth": depth_at_feat,        # [B, 1, h_grid, w_grid] or None
+            "h_grid": int(h_grid),
+            "w_grid": int(w_grid),
+        }
+
+    def _finalize(self, aux, B: int, h_grid: int, w_grid: int):
+        """Pull per-layer feats + ray + optional depth out of a completed DA3
+        forward. Returns tensors whose leading dim is B*V (V baked in by DA3's
+        aux layout): feats as list of [B*V, C, h_grid, w_grid], ray as
+        [B*V, 3, h_grid, w_grid], depth as [B*V, 1, h_grid, w_grid] or None.
+        The single-view ``forward`` (V=1) gets [B, ...]; the joint multi-view
+        path unflattens the leading dim back to (B, V).
+        """
+        # Extract per-level feats at native (h_grid, w_grid). DA3 may emit them
+        # at the backbone resolution already matching (h_grid, w_grid) since
+        # in_h/in_w = h_grid*patch × w_grid*patch.
+        feats: List[torch.Tensor] = []
+        for li in self.out_layers:
+            key = f"feat_layer_{li}"
+            fi = aux.get(key)
+            if fi is None:
+                raise RuntimeError(
+                    f"DA3LargeForGeoStack: aux missing key '{key}'; available={list(aux.keys())[:8]}"
+                )
+            fi = self._to_bchw(fi, B=B)  # [B*V, C, h, w]
+            if fi.shape[-1] != w_grid or fi.shape[-2] != h_grid:
+                # Defensive resize in case DA3 produced a different grid.
+                fi = F.interpolate(
+                    fi.float(), size=(h_grid, w_grid),
+                    mode="bilinear", align_corners=False,
+                )
+            feats.append(fi)
+
+        # Extract ray-head output (3D direction only, normalized)
+        ray_raw = self._ray_capture.get("ray_raw", None)
+        if ray_raw is None:
+            raise RuntimeError(
+                "DA3LargeForGeoStack: ray-head hook did not fire; DA3 forward did not run "
+                "the aux/ray branch. Check head.scratch.output_conv2_aux exists."
+            )
+        if ray_raw.dim() == 4:
+            ray_dir = ray_raw[:, :3, :, :]
+        elif ray_raw.dim() == 5:                          # [B, V, C, H, W]
+            ray_dir = ray_raw[:, :, :3, :, :].reshape(-1, 3, ray_raw.shape[-2], ray_raw.shape[-1])
+        else:
+            raise RuntimeError(f"Unexpected ray tensor rank: {ray_raw.dim()}, shape={tuple(ray_raw.shape)}")
+        ray_at_feat = F.interpolate(
+            ray_dir.float(), size=(h_grid, w_grid),
+            mode="bilinear", align_corners=False,
+        )
+        ray_at_feat = ray_at_feat / (ray_at_feat.norm(dim=1, keepdim=True) + 1e-6)
+        self._ray_capture.clear()
+
+        # Optional depth (v4.1) — extract DA3's main-head depth output at patch grid
+        depth_at_feat = None
+        if self._capture_depth:
+            depth_raw = self._depth_capture.get("depth_raw", None)
+            if depth_raw is not None:
+                # depth_raw: [B*V, C_out, H, W] or [B, V, C_out, H, W]
+                if depth_raw.dim() == 5:
+                    depth_raw = depth_raw.reshape(-1, depth_raw.shape[-3], depth_raw.shape[-2], depth_raw.shape[-1])
+                # First channel = raw depth logits (DPT convention)
+                depth_scalar = depth_raw[:, :1].float()
+                depth_at_feat = F.interpolate(
+                    depth_scalar, size=(h_grid, w_grid),
+                    mode="bilinear", align_corners=False,
+                )
+            self._depth_capture.clear()
+
+        return feats, ray_at_feat, depth_at_feat
+
+
+    # ---------- multi-view forward (v2C) ----------
+
+    def forward_multi_view(
+        self,
+        pixel_values: torch.Tensor,                              # [B, V, 3, H_in, W_in]
+        target_input_hw: Tuple[int, int] | None = None,
+        extrinsics: torch.Tensor | None = None,                  # [B, V, 4, 4] OpenCV w2c
+        intrinsics: torch.Tensor | None = None,                  # [B, V, 3, 3] @ (H_in, W_in)
+    ) -> Dict[str, object]:
+        """Process multiple views, returning per-view feats/rays.
+
+        Two modes, same output contract:
+          * Unposed (extrinsics/intrinsics=None): flatten [B,V]→[B*V] and run
+            each view through DA3 INDEPENDENTLY (legacy v2C wrist-depth path).
+          * Posed joint (both provided): feed ALL V views TOGETHER into one DA3
+            inner() call with camera extrinsics + intrinsics, so DA3's camera
+            encoder makes the per-view features geometrically consistent across
+            views. Downstream per-view routing/cross-attention is unchanged —
+            only the feature content becomes cross-view-aware.
+
+        Args:
+            pixel_values: [B, V, 3, H_in, W_in] ImageNet-normalized float.
+            target_input_hw: optional (h, w) override (e.g., 224×224 for wrists).
+            extrinsics: optional [B, V, 4, 4] OpenCV world-to-camera.
+            intrinsics: optional [B, V, 3, 3] calibrated to (H_in, W_in); DA3
+                resizes views to (in_h, in_w) internally so K is rescaled to match.
+
+        Returns:
+            dict with keys:
+                feats:  List[torch.Tensor]  one per out_layer, each [B, V, C, h_grid, w_grid]
+                ray:    torch.Tensor        [B, V, 3, h_grid, w_grid]
+                depth:  torch.Tensor | None [B, V, 1, h_grid, w_grid]
+                h_grid: int
+                w_grid: int
+        """
+        if pixel_values.dim() != 5:
+            raise ValueError(
+                f"forward_multi_view expects [B,V,3,H,W]; got {tuple(pixel_values.shape)}"
+            )
+        B, V = pixel_values.shape[:2]
+
+        # ---- Posed joint multi-view path ----
+        if extrinsics is not None and intrinsics is not None:
+            return self._forward_multi_view_posed(
+                pixel_values, extrinsics, intrinsics, target_input_hw
+            )
+
+        # ---- Legacy unposed per-view path ----
+        flat = pixel_values.flatten(0, 1)                        # [B*V, 3, H_in, W_in]
+        out_flat = self.forward(flat, target_input_hw=target_input_hw)
+        feats_flat = out_flat["feats"]                            # list of [B*V, C, h, w]
+        ray_flat = out_flat["ray"]                                # [B*V, 3, h, w]
+        h_grid = int(out_flat["h_grid"])
+        w_grid = int(out_flat["w_grid"])
+        feats = [f.unflatten(0, (B, V)) for f in feats_flat]      # list of [B, V, C, h, w]
+        ray = ray_flat.unflatten(0, (B, V))                       # [B, V, 3, h, w]
+        result = {
+            "feats": feats,
+            "ray": ray,
+            "h_grid": h_grid,
+            "w_grid": w_grid,
+        }
+        # v4.1: forward multi-view depth if the depth hook is on
+        depth_flat = out_flat.get("depth", None)
+        if depth_flat is not None:
+            result["depth"] = depth_flat.unflatten(0, (B, V))    # [B, V, 1, h, w]
+        else:
+            result["depth"] = None
+        return result
+
+    def _forward_multi_view_posed(
+        self,
+        pixel_values: torch.Tensor,           # [B, V, 3, H_in, W_in]
+        extrinsics: torch.Tensor,             # [B, V, 4, 4] OpenCV w2c
+        intrinsics: torch.Tensor,             # [B, V, 3, 3] @ (H_in, W_in)
+        target_input_hw: Tuple[int, int] | None = None,
+    ) -> Dict[str, object]:
+        """Joint posed multi-view DA3: one inner() call over all V views with
+        camera pose conditioning → cross-view-consistent per-view features.
+        Same return contract as forward_multi_view (unposed)."""
+        B, V, _, H_in, W_in = pixel_values.shape
+        device = pixel_values.device
+
+        # Resolve DA3 input dims + the resulting token grid.
+        if target_input_hw is None:
+            in_h, in_w = self.da3_input_h, self.da3_input_w
+            h_grid, w_grid = self.h_grid, self.w_grid
+        else:
+            in_h, in_w = int(target_input_hw[0]), int(target_input_hw[1])
+            if in_h % self.patch_size != 0 or in_w % self.patch_size != 0:
+                raise ValueError(
+                    f"target_input_hw=({in_h},{in_w}) must be multiples of patch_size={self.patch_size}"
+                )
+            h_grid, w_grid = in_h // self.patch_size, in_w // self.patch_size
+
+        # Aspect-correct resize of every view to (in_h, in_w).
+        x = pixel_values.to(torch.float32).flatten(0, 1)          # [B*V, 3, H_in, W_in]
+        if (H_in, W_in) != (in_h, in_w):
+            x = F.interpolate(x, size=(in_h, in_w), mode="bilinear", align_corners=False)
+        x = x.unflatten(0, (B, V))                                # [B, V, 3, in_h, in_w]
+
+        # Rescale intrinsics from the incoming tensor resolution to (in_h, in_w)
+        # so DA3's pose math sees K aligned with the pixels it actually gets.
+        K = self._rescale_intrinsics(intrinsics, (H_in, W_in), (in_h, in_w))
+
+        # Clear stale captures, run ONE joint posed DA3 forward.
+        self._ray_capture.clear()
+        if self._capture_depth:
+            self._depth_capture.clear()
+        if self.use_bf16:
+            with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+                aux, depth_full = self._da3_forward_posed(x, extrinsics, K)
+        else:
+            aux, depth_full = self._da3_forward_posed(x, extrinsics, K)
+
+        # Shared extraction — leading dim is B*V (V baked into DA3's aux layout).
+        feats_flat, ray_flat, depth_hook = self._finalize(aux, B, h_grid, w_grid)
+        feats = [f.unflatten(0, (B, V)) for f in feats_flat]      # list of [B, V, C, h, w]
+        ray = ray_flat.unflatten(0, (B, V))                       # [B, V, 3, h, w]
+
+        # Metric depth from the top-level output (preferred over the DPT hook):
+        # [B,V,H,W] → resize to the token grid → [B,V,1,h,w].
+        if depth_full is not None:
+            df = depth_full.float()
+            if df.dim() == 4:                                     # [B, V, H, W]
+                df = df.flatten(0, 1).unsqueeze(1)                # [B*V, 1, H, W]
+            elif df.dim() == 5:                                   # [B, V, 1, H, W]
+                df = df.flatten(0, 1)                             # [B*V, 1, H, W]
+            df = F.interpolate(df, size=(h_grid, w_grid), mode="bilinear", align_corners=False)
+            depth_grid = df.unflatten(0, (B, V))                  # [B, V, 1, h, w]
+        elif depth_hook is not None:
+            depth_grid = depth_hook.unflatten(0, (B, V))
+        else:
+            depth_grid = None
+
+        result: Dict[str, object] = {
+            "feats": feats,
+            "ray": ray,
+            "h_grid": int(h_grid),
+            "w_grid": int(w_grid),
+            "depth": depth_grid,                                  # [B, V, 1, h, w] metric
+        }
+        return result
+
+
+__all__ = ["DA3LargeForGeoStack"]

pibehavior_da3_clean_up_your_desk_40k/source_code/launch/launch_b1k_robopro.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+# Launch/resume the b1k DA3 robopro-exact run. Arg1 = physical GPU list (e.g. "4,5,6,7").
+# Auto-resumes from the latest checkpoint if one exists; else fresh from behavior_50t.
+set -uo pipefail
+GPUS="${1:-4,5,6,7}"
+cd /work/jack/behavior-1k-solution
+# Auto mode: 8 GPUs => dedicated extraction (hold out first 4 for the frozen DA3 forward, train on last 4);
+# 4 GPUs => inline (extraction shares the training GPUs). Extraction always pins cuda:0-3 (the first 4 given).
+NGPU=$(echo "$GPUS" | tr ',' '\n' | grep -c .)
+if [ "$NGPU" -ge 8 ]; then
+  # 8-GPU: 4 train + 4 extract. Training stays 4-way DP to MATCH the checkpoint's mesh — changing the
+  # #training-GPUs on resume reshapes the sharding and breaks restore, so we hold it fixed at 4.
+  HOLDOUT=4; QDEPTH=6; EXTRACT_DEVS="cuda:0,cuda:1,cuda:2,cuda:3"; BSZ=128; GMODE="8-GPU: 4 train + 4 extract"
+else
+  HOLDOUT=0; QDEPTH=2; EXTRACT_DEVS="cuda:0,cuda:1,cuda:2,cuda:3"; BSZ=128; GMODE="${NGPU}-GPU inline"
+fi
+LOG=/work/jack/b1k_da3_robopro_desk_train.log
+CKDIR=outputs/checkpoints/pi_behavior_b1k_fast/b1k_da3_robopro_desk
+# resume if a numeric checkpoint exists, else fresh
+if ls "$CKDIR" 2>/dev/null | grep -qE '^[0-9]+$'; then RES=1; OVW=0; MODE="RESUME from step $(ls "$CKDIR"|grep -E '^[0-9]+$'|sort -n|tail -1)"; else RES=0; OVW=1; MODE="FRESH"; fi
+echo "[$(date +%H:%M)] launching ($MODE, $GMODE) on GPUs $GPUS" >> /work/jack/b1k_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" \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+B1K_HOLDOUT_GPU0=$HOLDOUT \
+B1K_EXTRACT_DEVICES=$EXTRACT_DEVS B1K_DECODE_THREADS=1 \
+B1K_SHARED_DECODE=1 \
+B1K_DLPACK=1 \
+B1K_DA3_FWD_CHUNK=${B1K_DA3_FWD_CHUNK:-32} \
+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=/work/jack/behavior1k/data/behavior_2026_task29_42 \
+USE_DA3_FULL=1 DA3_INIT_STD=0.01 DA3_LOGIT_GAIN=1 DA3_SCALE=2.0 DA3_LR_GROUPS=1 \
+LR_VLM_PEAK=2.5e-5 LR_VLM_END=2.5e-6 LR_CORE_PEAK=5e-4 LR_CORE_END=5e-5 LR_GEOM_PEAK=5e-4 LR_GEOM_END=1e-4 \
+B1K_ACTIVITIES=clean_up_your_desk \
+B1K_INIT_PARAMS=/work/jack/behavior1k/checkpoints/behavior_50t_checkpoint/params \
+EXP=b1k_da3_robopro_desk \
+STEPS=50000 BS=$BSZ FLOW=15 NW=${NW:-32} SHUFFLE=1 RESUME=$RES OVERWRITE=$OVW \
+SAVE_INTERVAL=2000 KEEP_PERIOD=2000 LOG_INTERVAL=50 \
+LR_WARMUP=1000 LR_DECAY_STEPS=50000 \
+  nohup .venv/bin/python scripts/train_2026.py >> "$LOG" 2>&1 &
+echo "LAUNCHED pid $! ($MODE) on GPUs $GPUS"

pibehavior_da3_clean_up_your_desk_40k/source_code/model/observation.py

@@ -0,0 +1,168 @@
+"""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
+    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"),
+            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),
+    )

pibehavior_da3_clean_up_your_desk_40k/source_code/model/pi_behavior.py

@@ -0,0 +1,1270 @@
+"""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,
+            )
+        )
+        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,
+                bank_token_embed=d.bank_token_embed,
+                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):
+        """Build the per-view DA3 spatial banks once per forward (timestep-independent; reused
+        across all flow samples / denoise steps)."""
+        if self.spatial_bank_builder is None or getattr(observation, "da3_features", None) is None:
+            return 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,
+        )
+
+    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).
+        spatial_banks = self._compute_banks(observation)
+        
+        # 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)
+        
+        # 13. Total loss
+        losses["total_loss"] = losses["action_loss"] + subtask_loss_value + fast_loss_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,
+    ) -> _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."
+            )
+
+        # First fill KV cache with a forward pass of the prefix (no FAST tokens during inference)
+        prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
+        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

pibehavior_da3_clean_up_your_desk_40k/source_code/model/pi_behavior_config.py

@@ -0,0 +1,234 @@
+"""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 = 32.0
+    bank_token_embed: bool = True
+    perceiver_query_std: float = 0.05
+
+
+@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

pibehavior_da3_clean_up_your_desk_40k/source_code/model/spatial_da3.py

@@ -0,0 +1,266 @@
+"""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 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, *, 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)
+
+    def __call__(self, q, kv, key_pad_mask=None):
+        # q:[b,Lq,d]  kv:[b,Lk,d]  key_pad_mask:[b,Lk] True=pad (ignored)
+        h = self.num_heads
+        Q = einops.rearrange(self.q_proj(q), "b l (h d) -> b h l d", h=h)
+        K = einops.rearrange(self.k_proj(kv), "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)
+        logits = jnp.einsum("bhqd,bhkd->bhqk", Q, K) * (self.head_dim**-0.5)
+        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, *, 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, rngs=rngs)
+
+    def __call__(self, q_hidden, kv_hidden, key_pad_mask=None, residual_scale: float = 1.0):
+        q = self.q_norm(q_hidden)
+        kv = self.kv_norm(kv_hidden)
+        out = self.attn(q, kv, key_pad_mask=key_pad_mask)
+        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)))
+
+
+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, 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, rngs=rngs)
+        self.mlp = ResidualMlp(dim, mlp_ratio=2.0, rngs=rngs)
+
+    def __call__(self, tokens):
+        b = tokens.shape[0]
+        q = jnp.broadcast_to(self.query.value, (b, *self.query.value.shape[1:]))
+        z = self.xattn(q, tokens, residual_scale=1.0)  # residual adds RAW q (matches X-VLA)
+        return self.mlp(z)
+
+
+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, *, rngs: nnx.Rngs):
+        self.layers = [
+            (ResidualCrossAttn(dim, num_heads, 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
+
+
+# ---------------------------------------------------------------------------
+# 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,
+        bank_token_embed: bool = False,
+        rngs: nnx.Rngs,
+    ):
+        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) scale-aware ray (Plucker-6 + log-depth = 7)
+        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)
+        # (e) per-view perceiver + language fusion
+        self.perceivers = {
+            name: PerceiverDownsampler(H, k, num_heads, query_std=perceiver_query_std, rngs=rngs)
+            for name, _, k in _VIEWS
+        }
+        self.lang_fusers = {name: LanguageFusionStack(H, lang_fusion_depth, num_heads, rngs=rngs) for name, _, _ in _VIEWS}
+        # (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
+        )
+
+    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
+            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):
+        # 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
+        h, w = self.grid_hw
+        pos_emb = self.pos2d_mlp(_grid_coords(h, w).astype(feats.dtype))  # [1,432,H]
+        lang_tokens = self.t5_projector(lang_feat)  # [b,Lt,H]
+        lang_pad = jnp.logical_not(lang_mask)  # True=pad
+        banks = {}
+        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]
+            ray_emb = self.ray_mlp(ray_flat.astype(feats.dtype))  # [b,432,H]
+            view_emb = self.view_embed(jnp.asarray(vidx))[None, None, :]  # [1,1,H]
+            spatial = fused + view_emb + pos_emb + ray_emb  # [b,432,H]
+            geo = self.perceivers[name](spatial)  # [b,K,H]
+            bank = self.lang_fusers[name](geo, lang_tokens, lang_pad)  # [b,K,H]
+            if self.bank_token_embeds is not None:
+                bank = bank + self.bank_token_embeds[name].value
+            banks[name] = bank
+        return banks

pibehavior_da3_clean_up_your_desk_40k/source_code/openpi/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"]

pibehavior_da3_clean_up_your_desk_40k/source_code/openpi/gemma.py

@@ -0,0 +1,613 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Gemma adaptation for Pi, taken from big_vision.
+
+We follow this einsum axis naming convention:
+  B: batch
+  T: query length
+  S: k/v length
+  N: num query heads
+  K: num k/v heads
+  G: num query heads per k/v head
+  H: head dim
+  D: d_model ("features")
+"""
+
+from collections.abc import Sequence
+import dataclasses
+from typing import Literal, TypeAlias
+
+import os
+import einops
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+
+import openpi.models.lora as lora
+import openpi.shared.array_typing as at
+import openpi.training.sharding as sharding
+
+PALIGEMMA_VOCAB_SIZE = 257_152
+
+
+@dataclasses.dataclass
+class Config:
+    width: int
+    depth: int
+    mlp_dim: int
+    num_heads: int
+    num_kv_heads: int
+    head_dim: int
+    lora_configs: dict[str, lora.LoRAConfig] = dataclasses.field(default_factory=dict)
+
+
+Variant = Literal["dummy", "gemma_300m", "gemma_300m_lora", "gemma_2b", "gemma_2b_lora"]
+
+
+def get_config(variant: Variant) -> Config:
+    """Returns config for specified gemma variant."""
+    if variant == "dummy":
+        return Config(
+            width=64,
+            depth=4,
+            mlp_dim=128,
+            num_heads=8,
+            num_kv_heads=1,
+            head_dim=16,
+        )
+    if variant == "gemma_300m":
+        # 311M params
+        return Config(
+            width=1024,
+            depth=18,
+            mlp_dim=4096,
+            num_heads=8,
+            num_kv_heads=1,
+            head_dim=256,
+        )
+    if variant == "gemma_2b":
+        return Config(
+            width=2048,
+            depth=18,
+            mlp_dim=16_384,
+            num_heads=8,
+            num_kv_heads=1,
+            head_dim=256,
+        )
+    if variant == "gemma_2b_lora":
+        return Config(
+            width=2048,
+            depth=18,
+            mlp_dim=16_384,
+            num_heads=8,
+            num_kv_heads=1,
+            head_dim=256,
+            lora_configs={"attn": lora.LoRAConfig(rank=16, alpha=16.0), "ffn": lora.LoRAConfig(rank=16, alpha=16.0)},
+        )
+    if variant == "gemma_2b_lora_32":
+        return Config(
+            width=2048,
+            depth=18,
+            mlp_dim=16_384,
+            num_heads=8,
+            num_kv_heads=1,
+            head_dim=256,
+            lora_configs={"attn": lora.LoRAConfig(rank=32, alpha=32.0), "ffn": lora.LoRAConfig(rank=32, alpha=32.0)},
+        )
+    if variant == "gemma_300m_lora":
+        # 311M params
+        return Config(
+            width=1024,
+            depth=18,
+            mlp_dim=4096,
+            num_heads=8,
+            num_kv_heads=1,
+            head_dim=256,
+            lora_configs={"attn": lora.LoRAConfig(rank=32, alpha=32.0), "ffn": lora.LoRAConfig(rank=32, alpha=32.0)},
+        )
+    raise ValueError(f"Unknown variant: {variant}")
+
+
+@at.typecheck
+class RMSNorm(nn.Module):
+    @nn.compact
+    def __call__(self, x, cond):
+        dtype = x.dtype  # original dtype, could be half-precision
+        var = jnp.mean(jnp.square(x.astype(jnp.float32)), axis=-1, keepdims=True)  # compute variance in float32
+        normed_inputs = jnp.asarray(x * jnp.reciprocal(jnp.sqrt(var + 1e-06)))  # compute normalization in float32
+        if cond is None:
+            # regular RMSNorm
+            scale = self.param("scale", nn.initializers.zeros_init(), (x.shape[-1]))
+            normed_inputs = normed_inputs * (
+                1 + scale
+            )  # scale by learned parameter in float32 (matches Flax implementation)
+            return normed_inputs.astype(dtype), None  # return in original dtype
+
+        # adaptive RMSNorm
+        modulation = nn.Dense(x.shape[-1] * 3, kernel_init=nn.initializers.zeros, dtype=dtype)(cond)
+        scale, shift, gate = jnp.split(modulation[:, None, :], 3, axis=-1)
+        normed_inputs = normed_inputs * (1 + scale) + shift  # scale and shift in float32
+        return normed_inputs.astype(dtype), gate
+
+
+@at.typecheck
+class Embedder(nn.Module):
+    """Embedder module."""
+
+    vocab_size: int
+    embed_dim: int
+
+    def setup(self):
+        self.input_embedding_table = self.param(
+            "input_embedding",
+            nn.initializers.normal(),
+            (self.vocab_size, self.embed_dim),
+        )
+
+    def encode(self, x):
+        x = self.input_embedding_table[(x,)]
+        x *= jnp.sqrt(self.embed_dim).astype(x.dtype)
+        return x
+
+    def decode(self, x):
+        return jnp.dot(x, self.input_embedding_table.T)
+
+
+@at.typecheck
+class Attention(nn.Module):
+    """Attention module."""
+
+    configs: Sequence[Config]
+
+    @nn.compact
+    def __call__(self, xs, positions, attn_mask, kv_cache):
+        # all experts must share the same head dim, num heads, and num kv heads for self-attention to work
+        assert all(config.head_dim == self.configs[0].head_dim for config in self.configs)
+        assert all(config.num_heads == self.configs[0].num_heads for config in self.configs)
+        assert all(config.num_kv_heads == self.configs[0].num_kv_heads for config in self.configs)
+
+        dtype = next(x.dtype for x in xs if x is not None)  # original dtype, could be half-precision
+
+        qkvs = []
+        for i, (x, config) in enumerate(zip(xs, self.configs, strict=True)):
+            if x is None:
+                continue
+            if config.num_kv_heads == config.num_heads:
+                qkv_einsum = lora.Einsum(
+                    shape=(3, config.num_heads, config.width, config.head_dim),
+                    name=_name("qkv_einsum", i),
+                    init_fn=nn.initializers.lecun_normal(in_axis=-2, out_axis=-1, batch_axis=(0, 1)),
+                    lora_config=config.lora_configs.get("attn"),
+                )
+                qkvs.append(qkv_einsum("BSD,3KDH->3BSKH", x))
+            else:
+                q_einsum = lora.Einsum(
+                    shape=(config.num_heads, config.width, config.head_dim),
+                    name=_name("q_einsum", i),
+                    init_fn=nn.initializers.lecun_normal(in_axis=-2, out_axis=-1, batch_axis=(0,)),
+                    lora_config=config.lora_configs.get("attn"),
+                )
+                q = q_einsum("BTD,NDH->BTNH", x)
+                kv_einsum = lora.Einsum(
+                    shape=(2, config.num_kv_heads, config.width, config.head_dim),
+                    name=_name("kv_einsum", i),
+                    init_fn=nn.initializers.lecun_normal(in_axis=-2, out_axis=-1, batch_axis=(0, 1)),
+                    lora_config=config.lora_configs.get("attn"),
+                )
+                k, v = kv_einsum("BSD,2KDH->2BSKH", x)
+                qkvs.append((q, k, v))
+
+        q, k, v = (jnp.concatenate(y, axis=1) for y in zip(*qkvs, strict=True))
+
+        q = _apply_rope(q, positions=positions)
+        q *= self.configs[0].head_dim ** -0.5
+
+        k = _apply_rope(k, positions=positions)
+
+        # should still be half-precision here (if input was half-precision)
+        assert q.dtype == k.dtype == v.dtype == dtype
+
+        if kv_cache is not None:
+            cache_k, cache_v = kv_cache
+            k = jnp.concatenate([cache_k, k], axis=1)
+            v = jnp.concatenate([cache_v, v], axis=1)
+
+        q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=self.configs[0].num_kv_heads)
+        logits = jnp.einsum("BTKGH,BSKH->BKGTS", q, k, preferred_element_type=jnp.float32)
+
+        if attn_mask.shape != (q.shape[0], 1, q.shape[1], k.shape[1]):
+            raise ValueError(
+                f"Attention mask with shape {attn_mask.shape} but shapes for q and k are: {q.shape} and {k.shape}"
+            )
+
+        # big_neg = jnp.finfo(logits.dtype).min
+        big_neg = -2.3819763e38  # See gemma/modules.py
+        masked_logits = jnp.where(attn_mask[:, :, None, :, :], logits, big_neg)
+
+        probs = jax.nn.softmax(masked_logits, axis=-1).astype(dtype)
+
+        encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v)
+        encoded = einops.rearrange(encoded, "B T K G H -> B T (K G) H")
+
+        out = []
+        start = 0
+        for i, (x, config) in enumerate(zip(xs, self.configs, strict=True)):
+            if x is not None:
+                end = start + x.shape[1]
+                out_einsum = lora.Einsum(
+                    shape=(config.num_heads, config.head_dim, config.width),
+                    name=_name("attn_vec_einsum", i),
+                    init_fn=nn.initializers.lecun_normal(in_axis=(-3, -2), out_axis=-1),
+                    lora_config=config.lora_configs.get("attn"),
+                )
+                out.append(out_einsum("BTNH,NHD->BTD", encoded[:, start:end]))
+                start = end
+            else:
+                out.append(None)
+
+        return out, (k, v)
+
+
+@at.typecheck
+class FeedForward(nn.Module):
+    """Feed forward module."""
+
+    features: int
+    hidden_dim: int
+
+    @nn.compact
+    def __call__(self, x):
+        dtype = x.dtype  # original dtype, could be half-precision
+        w_gating = self.param(
+            "gating_einsum",
+            nn.initializers.lecun_normal(in_axis=-2, out_axis=-1, batch_axis=(0,)),
+            (2, self.features, self.hidden_dim),
+        ).astype(dtype)
+        ff_gate = jnp.dot(x, w_gating[0])
+        gate_value = nn.gelu(ff_gate)
+
+        ff1 = jnp.dot(x, w_gating[1])
+        activations = gate_value * ff1
+
+        w_linear = self.param(
+            "linear",
+            nn.initializers.lecun_normal(in_axis=-2, out_axis=-1),
+            (self.hidden_dim, self.features),
+        ).astype(dtype)
+        outputs = jnp.dot(activations, w_linear)
+        assert outputs.dtype == dtype
+        return outputs
+
+
+class _LinenCrossAttn(nn.Module):
+    """Pre-LN cross-attention. out_proj is zero-init by default (identity at step 0); out_init_std>0
+    gives it a small random init so q/k/v receive gradient from step 0 (with zero-init their grads
+    are zero until out_proj grows off zero, LoRA-style bootstrap)."""
+
+    d_model: int
+    num_heads: int
+    out_init_std: float = 0.0
+    logit_gain: bool = False  # learnable per-head gain on the attention logits (v2 fix)
+    logit_gain_init: float = 1.0  # initial gain value; >1 = sharper attention at init (bigger softmax
+    # jacobian => stronger, more consistent Q/K gradients; uniform attention scales them by ~1/K)
+
+    @nn.compact
+    def __call__(self, q_hidden, kv):
+        hd = self.d_model // self.num_heads
+        q = nn.LayerNorm(epsilon=1e-5, name="q_norm")(q_hidden)
+        kvn = nn.LayerNorm(epsilon=1e-5, name="kv_norm")(kv)
+        Q = einops.rearrange(nn.Dense(self.d_model, name="q_proj")(q), "b l (h d) -> b h l d", h=self.num_heads)
+        K = einops.rearrange(nn.Dense(self.d_model, name="k_proj")(kvn), "b l (h d) -> b h l d", h=self.num_heads)
+        V = einops.rearrange(nn.Dense(self.d_model, name="v_proj")(kvn), "b l (h d) -> b h l d", h=self.num_heads)
+        logits = jnp.einsum("bhqd,bhkd->bhqk", Q, K) * (hd**-0.5)
+        if self.logit_gain:
+            # exp-parameterized per-head gain, init logit_gain_init (v2 uses 32); multiplies dL/dQ,K so
+            # the attention pattern can actually train (v1 measured Q/K ~1000x slower than V/out).
+            log_gain = self.param(
+                "log_gain", nn.initializers.constant(jnp.log(self.logit_gain_init)), (self.num_heads,)
+            )
+            logits = logits * jnp.exp(log_gain)[None, :, None, None].astype(logits.dtype)
+        probs = jax.nn.softmax(logits.astype(jnp.float32), axis=-1).astype(Q.dtype)
+        ctx = einops.rearrange(jnp.einsum("bhqk,bhkd->bhqd", probs, V), "b h q d -> b q (h d)")
+        out_init = nn.initializers.normal(stddev=self.out_init_std) if self.out_init_std > 0 else nn.initializers.zeros
+        return nn.Dense(
+            self.d_model, name="out_proj", kernel_init=out_init, bias_init=nn.initializers.zeros
+        )(ctx)
+
+
+class SpatialActionInjection(nn.Module):
+    """X-VLA Method-B injection (main x-attn -> left/right branches -> merge), scaled residual.
+
+    init_std=0 (default): output projections zero-init => returns 0 (exact identity at step 0, pi05_base
+    reproduced bit-exactly). init_std>0: small normal init on the xattn out_projs and merge_proj => a
+    small nonzero spatial delta from step 0 and immediate gradient flow to all injection weights.
+    Returns the DELTA to add to the action-token stream (external per-layer gate selects it).
+    """
+
+    d_model: int
+    num_heads: int = 8
+    spatial_scale: float = 2.0
+    init_std: float = 0.0
+    logit_gain: bool = False
+    logit_gain_init: float = 1.0
+
+    @nn.compact
+    def __call__(self, h0, banks):
+        s = self.spatial_scale
+        std = self.init_std
+        g = self.logit_gain
+        gi = self.logit_gain_init
+        h1 = h0 + s * _LinenCrossAttn(self.d_model, self.num_heads, std, g, gi, name="main_xattn")(h0, banks["main"])
+        hL = nn.Dense(self.d_model, name="left_branch")(h1)
+        hR = nn.Dense(self.d_model, name="right_branch")(h1)
+        hL = hL + s * _LinenCrossAttn(self.d_model, self.num_heads, std, g, gi, name="left_xattn")(hL, banks["left"])
+        hR = hR + s * _LinenCrossAttn(self.d_model, self.num_heads, std, g, gi, name="right_xattn")(hR, banks["right"])
+        merge_init = nn.initializers.normal(stddev=std) if std > 0 else nn.initializers.zeros
+        merge = nn.Dense(
+            self.d_model, name="merge_proj", kernel_init=merge_init, bias_init=nn.initializers.zeros
+        )(jnp.concatenate([hL, hR], axis=-1))
+        out = h1 + s * merge
+        return out - h0
+
+
+@at.typecheck
+class Block(nn.Module):
+    """Transformer block."""
+
+    configs: tuple[Config, ...]
+
+    dropout: float = 0.0
+    dropout_bdims: tuple[int, ...] = ()
+    spatial_scale: float = 2.0
+    spatial_init_std: float = 0.0
+    spatial_logit_gain: bool = False
+    spatial_logit_gain_init: float = 1.0
+
+    @nn.compact
+    def __call__(  # noqa: FBT002
+        self, xs, kv_cache, positions, attn_mask, adarms_cond, deterministic=True, banks=None, layer_gate=None
+    ):
+        xs = sharding.activation_sharding_constraint(xs)
+        drop = nn.Dropout(self.dropout, self.dropout_bdims) if self.dropout else lambda x, _: x
+
+        attn = Attention(configs=self.configs, name="attn")
+
+        pre_attn = []
+        gates = []
+        for i, x in enumerate(xs):
+            if x is not None:
+                x, gate = RMSNorm(name=_name("pre_attention_norm", i))(x, adarms_cond[i])  # noqa: PLW2901
+            pre_attn.append(x)
+            gates.append(gate if x is not None else None)
+
+        pre_attn = sharding.activation_sharding_constraint(pre_attn)
+        post_attn, kv_cache = attn(pre_attn, positions, attn_mask, kv_cache)
+        post_attn = jax.tree.map(lambda x: drop(x, deterministic), post_attn)
+        post_attn = sharding.activation_sharding_constraint(post_attn)
+        xs = [_gated_residual(x, y, gate) for x, y, gate in zip(xs, post_attn, gates, strict=True)]
+        xs = sharding.activation_sharding_constraint(xs)
+
+        out = []
+        gates = []
+        for i, (x, config) in enumerate(zip(xs, self.configs, strict=True)):
+            if x is not None:
+                x, gate = RMSNorm(name=_name("pre_ffw_norm", i))(x, adarms_cond[i])  # noqa: PLW2901
+                x = lora.FeedForward(  # noqa: PLW2901
+                    features=config.width,
+                    hidden_dim=config.mlp_dim,
+                    name=_name("mlp", i),
+                    lora_config=config.lora_configs.get("ffn"),
+                )(x)
+            out.append(x)
+            gates.append(gate if x is not None else None)
+
+        out = sharding.activation_sharding_constraint(out)
+        out = jax.tree.map(lambda x: drop(x, deterministic), out)
+        xs = [_gated_residual(x, y, gate) for x, y, gate in zip(xs, out, gates, strict=True)]
+        xs = sharding.activation_sharding_constraint(xs)
+
+        # ---- DA3 spatial injection on the SUFFIX (action-expert, expert 1) stream ONLY ----
+        # layer_gate is a per-layer scalar (0 for early blocks, 1 for the last N); with init_std=0 the
+        # injection delta is 0 (zero-init output) so this is an exact no-op regardless of the gate.
+        if banks is not None and len(xs) > 1 and xs[1] is not None:
+            suf = xs[1]
+            delta = SpatialActionInjection(
+                d_model=self.configs[1].width,
+                spatial_scale=self.spatial_scale,
+                init_std=self.spatial_init_std,
+                logit_gain=self.spatial_logit_gain,
+                logit_gain_init=self.spatial_logit_gain_init,
+                name="spatial_inject_1",
+            )(suf, banks)
+            # cast back to the suffix (carry) dtype: the injection runs in fp32 but the scan carry is bf16
+            xs = [xs[0], (suf + layer_gate.astype(delta.dtype) * delta).astype(suf.dtype)]
+            xs = sharding.activation_sharding_constraint(xs)
+
+        return xs, kv_cache
+
+
+KVCache: TypeAlias = tuple[at.Float[at.Array, "l b _t _k _h"], at.Float[at.Array, "l b _t _v _h"]]
+
+
+@at.typecheck
+class Module(nn.Module):
+    """Transformer model, supporting a mixture of different weights for different tokens."""
+
+    configs: Sequence[Config]  # list of configs, one for each expert
+    embed_dtype: str
+
+    dropout: float = 0.0
+    dropout_bdims: tuple[int, ...] = ()  # Every float is dropped independently.
+    adarms: bool = False
+    # DA3 spatial injection: inject per-view banks into the SUFFIX (action-expert) stream at the
+    # last `num_spatial_layers` blocks, via a gated cross-attention inside the same scan.
+    spatial_inject: bool = False
+    num_spatial_layers: int = 6
+    spatial_scale: float = 2.0
+    spatial_init_std: float = 0.0
+    spatial_logit_gain: bool = False  # learnable per-head attention-logit gain in the injection (v2)
+    spatial_logit_gain_init: float = 1.0  # v2 load-bearing fix: init the per-head gain to 32
+
+    def setup(self):
+        # all experts must have the same depth
+        assert all(config.depth == self.configs[0].depth for config in self.configs)
+
+        self.embedder = Embedder(
+            vocab_size=PALIGEMMA_VOCAB_SIZE,
+            embed_dim=self.configs[0].width,  # embedder for first expert only
+            name="embedder",
+        )
+        # Rematerialization trades memory for compute: activations are recomputed in the backward
+        # pass instead of stored. `nothing_saveable` (the default here) is the MOST aggressive
+        # setting -- minimum memory, maximum recompute (~30% extra FLOPs). Relaxing it is
+        # NUMERICALLY IDENTICAL (it only changes how gradients are computed, not their values), so
+        # where VRAM allows, a lighter policy trains faster:
+        #   B1K_REMAT=full (default) -> nothing_saveable : least memory, slowest
+        #   B1K_REMAT=dots           -> dots_saveable    : keeps matmul outputs, much less recompute
+        #   B1K_REMAT=none           -> no remat         : most memory, fastest
+        _remat = os.environ.get("B1K_REMAT", "full").lower()
+        if _remat == "none":
+            block_cls = Block
+        else:
+            block_cls = nn.remat(
+                Block,
+                prevent_cse=False,
+                static_argnums=(5,),  # 0=self, 6=deterministic
+                policy=jax.checkpoint_policies.dots_saveable
+                if _remat == "dots"
+                else jax.checkpoint_policies.nothing_saveable,
+            )
+        # banks are broadcast (same for all blocks, computed once); layer_gate is scanned on axis 0
+        # so each block receives its own scalar gate[l] (0 for early blocks, 1 for the last N).
+        base_in_axes = (0, nn.broadcast, nn.broadcast, nn.broadcast, nn.broadcast)
+        in_axes = (*base_in_axes, nn.broadcast, 0) if self.spatial_inject else base_in_axes
+        self.layers = nn.scan(
+            block_cls,
+            variable_axes={"params": 0},
+            split_rngs={"params": True, "dropout": True},
+            in_axes=in_axes,
+            length=self.configs[0].depth,
+        )(
+            configs=self.configs,
+            dropout=self.dropout,
+            dropout_bdims=self.dropout_bdims,
+            spatial_scale=self.spatial_scale,
+            spatial_init_std=self.spatial_init_std,
+            spatial_logit_gain=self.spatial_logit_gain,
+            spatial_logit_gain_init=self.spatial_logit_gain_init,
+        )
+        self.final_norms = [RMSNorm(name=_name("final_norm", i)) for i in range(len(self.configs))]
+
+    @at.typecheck
+    def embed(self, tokens: at.Int[at.Array, "b t"]) -> at.Float[at.Array, "b t d"]:
+        return self.embedder.encode(tokens).astype(self.embed_dtype)
+
+    @at.typecheck
+    def __call__(
+        self,
+        # list of token arrays, one for each expert, or None if that expert should not be run
+        embedded: Sequence[at.Float[at.Array, "b _t _d"] | None],
+        positions: at.Int[at.Array, "b t"],
+        mask: at.Bool[at.Array, "b t s"],
+        adarms_cond: Sequence[at.Float[at.Array, "b _d"] | None] | None = None,
+        *,
+        kv_cache: KVCache | None = None,
+        deterministic: bool = True,
+        banks: dict | None = None,
+    ) -> tuple[Sequence[at.Float[at.Array, "b _t _d"] | None], KVCache]:
+        embedded = jax.tree.map(lambda e: e.astype(self.embed_dtype), embedded)
+        mask = jnp.asarray(mask)[:, None, :, :]
+        if adarms_cond is None:
+            adarms_cond = [None] * len(self.configs)
+
+        if self.spatial_inject:
+            depth = self.configs[0].depth
+            if banks is None:
+                # Prefix-only passes carry no banks. The scan's in_axes arity is fixed at build time,
+                # so feed DUMMY banks + a zero gate — the Block skips injection anyway when the
+                # suffix stream (xs[1]) is None, and the zero gate kills it otherwise.
+                b = next(e.shape[0] for e in embedded if e is not None)
+                w = self.configs[1].width
+                banks = {k: jnp.zeros((b, 1, w), dtype=jnp.dtype(self.embed_dtype)) for k in ("main", "left", "right")}
+                layer_gate = jnp.zeros((depth,), jnp.float32)
+            else:
+                # gate[l] = 1 for the last num_spatial_layers blocks, else 0 (per-layer scalar via scan axis 0)
+                layer_gate = (jnp.arange(depth) >= depth - self.num_spatial_layers).astype(jnp.float32)
+            embedded, kv_cache = self.layers(
+                embedded, kv_cache, positions, mask, adarms_cond, deterministic, banks, layer_gate
+            )
+        else:
+            embedded, kv_cache = self.layers(embedded, kv_cache, positions, mask, adarms_cond, deterministic)
+
+        assert all(e.dtype == jnp.dtype(self.embed_dtype) for e in embedded if e is not None)
+
+        return [
+            f(e, a)[0] if e is not None else e for f, e, a in zip(self.final_norms, embedded, adarms_cond, strict=True)
+        ], kv_cache
+
+    def init(self, use_adarms: Sequence[bool]):
+        """Convenience method for initializing all parameters, necessary due to the quirks of linen."""
+        self.embed(jnp.zeros((1, 1), dtype=jnp.int32))
+        # dummy banks so the spatial-injection params are materialized during lazy_init (bank dim
+        # == action-expert width == configs[1].width); 1 dummy token per view is enough.
+        banks = None
+        if self.spatial_inject:
+            w = self.configs[1].width
+            banks = {k: jnp.zeros((1, 1, w)) for k in ("main", "left", "right")}
+        self(
+            [jnp.zeros((1, 1, c.width)) for c in self.configs],
+            jnp.zeros((1, len(self.configs)), dtype=jnp.int32),
+            jnp.zeros((1, len(self.configs), len(self.configs)), dtype=bool),
+            adarms_cond=[jnp.zeros((1, c.width)) if u else None for u, c in zip(use_adarms, self.configs, strict=True)],
+            banks=banks,
+        )
+
+
+def _apply_rope(x, *, positions, max_wavelength=10_000):
+    """Applies RoPE positions [B, L] to x [B, L, H, D]."""
+    freq_exponents = (2.0 / x.shape[-1]) * jnp.arange(x.shape[-1] // 2, dtype=jnp.float32)
+    timescale = max_wavelength**freq_exponents
+    radians = positions[..., None] / timescale[None, None, :]
+    radians = radians[..., None, :]
+    assert radians.dtype == jnp.float32
+    # radians.shape = [...,L,1,d=D/2]
+    sin, cos = jnp.sin(radians), jnp.cos(radians)
+    x1, x2 = jnp.split(x, 2, axis=-1)
+    res = jnp.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1)
+    assert res.dtype == jnp.float32
+    # The original bigvision impl allows RoPE to upcast to float32. It is then immediately downcast again to the cache
+    # dtype when in inference mode (but not in training mode). I don't think any of this was intentional. Based on the
+    # original DeepMind impl, as well as the widely-used transformers impl, it is ok to always downcast back to bfloat16
+    # here.
+    return res.astype(x.dtype)
+
+
+def _name(name, i):
+    # we name layers like this because we want the first expert's weights to have no suffix (e.g., "attn"), so that they
+    # can be loaded seamlessly from the existing PaliGemma checkpoint. subsequent experts will have a suffix (e.g.,
+    # "attn_1") and their weights will be initialized from scratch. in practice, we only use two experts -- PaliGemma,
+    # and the action expert.
+    if i == 0:
+        return name
+    return f"{name}_{i}"
+
+
+def _gated_residual(x, y, gate):
+    assert (x is None) == (y is None)
+    if x is None:
+        return None
+    if gate is None:
+        return x + y
+    return x + y * gate

pibehavior_da3_clean_up_your_desk_40k/source_code/openpi/optimizer.py

@@ -0,0 +1,176 @@
+import dataclasses
+from typing import Protocol, runtime_checkable
+
+import jax.numpy as jnp
+import optax
+
+import openpi.shared.array_typing as at
+
+
+@runtime_checkable
+class LRScheduleConfig(Protocol):
+    def create(self) -> optax.Schedule: ...
+
+
+@dataclasses.dataclass(frozen=True)
+class CosineDecaySchedule(LRScheduleConfig):
+    """Cosine decay schedule with warmup."""
+
+    warmup_steps: int = 1_000
+    peak_lr: float = 2.5e-5
+    decay_steps: int = 30_000
+    decay_lr: float = 2.5e-6
+
+    def create(self) -> optax.Schedule:
+        return optax.warmup_cosine_decay_schedule(
+            init_value=self.peak_lr / (self.warmup_steps + 1),
+            peak_value=self.peak_lr,
+            warmup_steps=self.warmup_steps,
+            decay_steps=self.decay_steps,
+            end_value=self.decay_lr,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class DelayedRampCosineSchedule(LRScheduleConfig):
+    """Hold at exactly 0 for `delay_steps`, linearly ramp 0 -> peak over `ramp_steps`, then cosine
+    decay to `decay_lr` at `decay_steps` (total).
+
+    For the DA3 *geometry* group: the randomly-initialized ray-MLP is held frozen while the
+    pretrained trunk takes the initial adaptation, then phased in slowly instead of hitting peak LR
+    alongside everything else. Geometry modules are sensitive to being driven hard from init.
+    """
+
+    delay_steps: int = 500
+    ramp_steps: int = 10_000
+    peak_lr: float = 7e-4
+    decay_steps: int = 25_000
+    decay_lr: float = 1e-4
+
+    def create(self) -> optax.Schedule:
+        tail = max(1, self.decay_steps - self.delay_steps - self.ramp_steps)
+        return optax.join_schedules(
+            [
+                optax.constant_schedule(0.0),  # frozen: geometry group takes no updates
+                optax.linear_schedule(init_value=0.0, end_value=self.peak_lr, transition_steps=self.ramp_steps),
+                optax.cosine_decay_schedule(
+                    init_value=self.peak_lr, decay_steps=tail, alpha=self.decay_lr / self.peak_lr
+                ),
+            ],
+            [self.delay_steps, self.delay_steps + self.ramp_steps],
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class RsqrtDecaySchedule(LRScheduleConfig):
+    """Inverse square root decay schedule with warmup."""
+
+    warmup_steps: int = 1_000
+    peak_lr: float = 5e-5
+    timescale: float = 10_000
+
+    def create(self) -> optax.Schedule:
+        return optax.join_schedules(
+            [
+                optax.linear_schedule(
+                    init_value=self.peak_lr / (self.warmup_steps + 1),
+                    end_value=self.peak_lr,
+                    transition_steps=self.warmup_steps,
+                ),
+                lambda step: self.peak_lr / jnp.sqrt((self.timescale + step) / self.timescale),
+            ],
+            [self.warmup_steps],
+        )
+
+
+@runtime_checkable
+class OptimizerConfig(Protocol):
+    def create(
+        self,
+        lr: optax.ScalarOrSchedule,
+        weight_decay_mask: at.PyTree | None = None,
+    ) -> optax.GradientTransformation: ...
+
+
+@dataclasses.dataclass(frozen=True)
+class AdamW(OptimizerConfig):
+    """AdamW optimizer."""
+
+    b1: float = 0.9
+    b2: float = 0.95
+    eps: float = 1e-8
+    # Changing this to 0 can cause out-of-memory errors for some reason, so we set it to a negligible value.
+    weight_decay: float = 1e-10
+    clip_gradient_norm: float = 1.0
+
+    def create(
+        self,
+        lr: optax.ScalarOrSchedule,
+        weight_decay_mask: at.PyTree | None = None,
+    ) -> optax.GradientTransformation:
+        tx = optax.adamw(
+            lr, b1=self.b1, b2=self.b2, eps=self.eps, weight_decay=self.weight_decay, mask=weight_decay_mask
+        )
+
+        return optax.chain(optax.clip_by_global_norm(self.clip_gradient_norm), tx)
+
+
+@dataclasses.dataclass(frozen=True)
+class SGD(OptimizerConfig):
+    """SGD optimizer."""
+
+    lr: float = 5e-5
+    momentum: float = 0.9
+    nesterov: bool = False
+
+    def create(
+        self,
+        lr: optax.ScalarOrSchedule,
+        weight_decay_mask: at.PyTree | None = None,
+    ) -> optax.GradientTransformation:
+        assert weight_decay_mask is None, "Weight decay is not supported for SGD"
+        return optax.sgd(lr, momentum=self.momentum, nesterov=self.nesterov)
+
+
+def create_optimizer(
+    optimizer: OptimizerConfig, lr_schedule: LRScheduleConfig, weight_decay_mask: at.PyTree | None = None
+) -> optax.GradientTransformation:
+    lr = lr_schedule.create()
+    return optimizer.create(lr, weight_decay_mask=weight_decay_mask)
+
+
+def spatial_group_labels(params: at.PyTree) -> at.PyTree:
+    """Label each param leaf 'geom' | 'core' | 'vlm' by its keypath, for the DA3 multi-LR setup.
+
+    geom = the DA3 geometry encoder (ray MLP); core = the trainable spatial branch (bank builder +
+    cross-attn injection); vlm = the PiBehavior backbone (VLM + action expert + task/stage/kv heads).
+    """
+    import jax
+
+    def label(kp, _leaf):
+        s = jax.tree_util.keystr(kp)
+        if "ray_mlp" in s:
+            return "geom"
+        if "spatial_bank_builder" in s or "spatial_inject" in s:
+            return "core"
+        return "vlm"
+
+    return jax.tree_util.tree_map_with_path(label, params)
+
+
+def create_multi_group_optimizer(
+    optimizer: OptimizerConfig,
+    lr_schedules: dict[str, LRScheduleConfig],
+    label_fn,
+    weight_decay_mask: at.PyTree | None = None,
+) -> optax.GradientTransformation:
+    """Per-parameter-group optimizer: a distinct LR schedule per group via optax.multi_transform.
+
+    `lr_schedules` maps group name -> schedule; `label_fn(params)` returns a matching-structure tree of
+    group names. Every leaf must map to a group present in `lr_schedules`.
+    """
+    txs = {
+        group: optimizer.create(sched.create(), weight_decay_mask=weight_decay_mask)
+        for group, sched in lr_schedules.items()
+    }
+    return optax.multi_transform(txs, label_fn)

pibehavior_da3_clean_up_your_desk_40k/source_code/openpi/sharding.py

@@ -0,0 +1,114 @@
+import contextlib
+import logging
+import os
+
+import jax
+import numpy as np
+
+BATCH_AXIS = "batch"
+FSDP_AXIS = "fsdp"
+# In FSDP, we shard the data across both the batch and FSDP axes.
+DATA_AXIS = (BATCH_AXIS, FSDP_AXIS)
+
+
+class _MeshState:
+    active_mesh: jax.sharding.Mesh | None = None
+
+
+def make_mesh(num_fsdp_devices: int) -> jax.sharding.Mesh:
+    devices = jax.devices()
+    # B1K DA3 inline-extraction speedup: hold GPU 0 out of the training mesh so the frozen DA3-GIANT
+    # extractor (pinned to cuda:0) runs there WITHOUT contending with the JAX train step. The prefetch
+    # thread then overlaps extraction of batch N+1 with the training of batch N on GPUs 1..N.
+    holdout = int(os.environ.get("B1K_HOLDOUT_GPU0", "0"))
+    if holdout > 0:
+        devices = devices[holdout:]
+        logging.info("make_mesh: holding out %d device(s) for extraction; training on %d devices %s",
+                     holdout, len(devices), [d.id for d in devices])
+    n = len(devices)
+    if n % num_fsdp_devices != 0:
+        raise ValueError(
+            f"Number of training devices {n} must be divisible by the number of FSDP devices {num_fsdp_devices}."
+        )
+    mesh_shape = (n // num_fsdp_devices, num_fsdp_devices)
+    mesh_devices = np.array(devices).reshape(mesh_shape)
+    return jax.sharding.Mesh(mesh_devices, (BATCH_AXIS, FSDP_AXIS))
+
+
+@contextlib.contextmanager
+def set_mesh(mesh: jax.sharding.Mesh):
+    """Plumbing the mesh deep into the module tree is extremeley cumbersome; until the JAX team lands a better API, a
+    custom context manager like this one is the recommended way to maintain a reference to a global mesh. This is only used
+    in `activation_sharding_constraint` below."""
+    if _MeshState.active_mesh is not None:
+        raise ValueError("Cannot nest set_mesh context managers.")
+    _MeshState.active_mesh = mesh
+    try:
+        yield
+    finally:
+        _MeshState.active_mesh = None
+
+
+def activation_sharding_constraint(pytree):
+    if _MeshState.active_mesh is None:
+        return pytree
+    return jax.lax.with_sharding_constraint(
+        pytree, jax.sharding.NamedSharding(_MeshState.active_mesh, jax.sharding.PartitionSpec(DATA_AXIS))
+    )
+
+
+def fsdp_sharding(
+    pytree,
+    mesh: jax.sharding.Mesh,
+    *,
+    min_size_mbytes: int = 4,  # 4 MiB
+    log: bool = False,
+):
+    """Apply FSDP sharding to a pytree of arrays based on the mesh shape.
+
+    Args:
+        pytree: A pytree to be apply sharding specified by the mesh, note that only array types (eg. contains .shape attr)
+          will be considered for sharding.
+        mesh: The mesh being used for applying sharding on to pytree.
+        min_size_mbytes: The minimum size of the array in MiB to be considered for sharding, any array smaller than this
+          will be replicated.
+        log: If true, will log the sharding decisions for arrays that are being considered for sharding.
+
+    Returns:
+        The sharded pytree.
+    """
+    min_size_bytes = min_size_mbytes * 2**20
+
+    def _shard_arr(kp, array: jax.ShapeDtypeStruct):
+        # if fsdp is not actually going to be used, replicate everything to avoid extraneous logging
+        if mesh.shape[FSDP_AXIS] == 1:
+            return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+        # replicate scalar and vector arrays
+        if not hasattr(array, "shape"):
+            return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+        if len(array.shape) < 2:
+            return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+        # replicate small arrays
+        if (arr_size := np.prod(array.shape) * np.dtype(array.dtype).itemsize) < min_size_bytes:
+            return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+
+        # shard matrices and larger tensors along the largest axis that is divisible by the fsdp dimension
+        axes = np.argsort(array.shape)[::-1]
+        spec = [None] * len(axes)
+        for i in axes:
+            if array.shape[i] % mesh.shape[FSDP_AXIS] == 0:
+                if log:
+                    logging.info(
+                        f"Sharding {jax.tree_util.keystr(kp)} of shape {array.shape} ({arr_size / 2**20:.2f} MiB) along axis {i}"
+                    )
+                spec[i] = FSDP_AXIS
+                return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec))
+
+        # replicate if no valid sharding was found
+        if log:
+            logging.warning(
+                f"Could not find a valid sharding for {jax.tree_util.keystr(kp)} of shape {array.shape} with mesh of shape {mesh.shape}"
+            )
+        return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+
+    return jax.tree_util.tree_map_with_path(_shard_arr, pytree)

pibehavior_da3_clean_up_your_desk_40k/source_code/openpi/weight_loaders.py

@@ -0,0 +1,110 @@
+import dataclasses
+import logging
+import re
+from typing import Protocol, runtime_checkable
+
+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
+
+logger = logging.getLogger(__name__)
+
+
+@runtime_checkable
+class WeightLoader(Protocol):
+    def load(self, params: at.Params) -> at.Params:
+        """Loads the model weights.
+
+        Args:
+            params: Parameters of the model. This is a nested structure of array-like objects that
+                represent the model's parameters.
+
+        Returns:
+            Loaded parameters. The structure must be identical to `params`. If returning a subset of
+            the parameters the loader must merge the loaded parameters with `params`.
+        """
+
+
+@dataclasses.dataclass(frozen=True)
+class NoOpWeightLoader(WeightLoader):
+    def load(self, params: at.Params) -> at.Params:
+        return params
+
+
+@dataclasses.dataclass(frozen=True)
+class CheckpointWeightLoader(WeightLoader):
+    """Loads an entire set of weights from a checkpoint.
+
+    Compatible with:
+      trained checkpoints:
+        example: "./checkpoints/<config>/<exp>/<step>/params"
+      released checkpoints:
+        example: "gs://openpi-assets/checkpoints/<model>/params"
+    """
+
+    params_path: str
+
+    def load(self, params: at.Params) -> at.Params:
+        # We are loading np.ndarray and relying on the training code to properly convert and shard the params.
+        loaded_params = _model.restore_params(download.maybe_download(self.params_path), restore_type=np.ndarray)
+        # Add all missing LoRA weights.
+        return _merge_params(loaded_params, params, missing_regex=".*lora.*")
+
+
+@dataclasses.dataclass(frozen=True)
+class PaliGemmaWeightLoader(WeightLoader):
+    """Loads weights from the official PaliGemma checkpoint.
+
+    This will overwrite existing weights with similar names while keeping all extra weights intact.
+    This allows us to support the action expert which is used by the Pi0 model.
+    """
+
+    def load(self, params: at.Params) -> at.Params:
+        path = download.maybe_download(
+            "gs://vertex-model-garden-paligemma-us/paligemma/pt_224.npz", gs={"token": "anon"}
+        )
+        with path.open("rb") as f:
+            flat_params = dict(np.load(f, allow_pickle=False))
+        loaded_params = {"PaliGemma": flax.traverse_util.unflatten_dict(flat_params, sep="/")["params"]}
+        # Add all missing weights.
+        return _merge_params(loaded_params, params, missing_regex=".*")
+
+
+def _merge_params(loaded_params: at.Params, params: at.Params, *, missing_regex: str) -> at.Params:
+    """Merges the loaded parameters with the reference parameters.
+
+    Args:
+        loaded_params: The parameters to merge.
+        params: The reference parameters.
+        missing_regex: A regex pattern for all missing keys that should be merged from the reference parameters.
+
+    Returns:
+        A new dictionary with the merged parameters.
+    """
+    # Flatten with TUPLE keys (no sep-join): some modules (e.g. the DA3 spatial bank builder) store
+    # submodules in lists, which nnx encodes with INTEGER path elements — sep="/" would crash on
+    # `"/".join(path)`. The string path is built only for the regex match below.
+    flat_ref = flax.traverse_util.flatten_dict(params)
+    flat_loaded = flax.traverse_util.flatten_dict(loaded_params)
+
+    # First, take all weights that are a subset of the reference weights.
+    result = {}
+    for k, v in flat_loaded.items():
+        if k in flat_ref:
+            if v.dtype == flat_ref[k].dtype:
+                result[k] = v
+            else:
+                print(f"Warning: {k} has dtype {v.dtype} but reference has dtype {flat_ref[k].dtype}")
+                result[k] = v.astype(flat_ref[k].dtype)
+    flat_loaded.clear()
+
+    # Then, merge any missing weights as defined by the missing regex.
+    pattern = re.compile(missing_regex)
+    for k in {k for k in flat_ref if pattern.fullmatch("/".join(str(p) for p in k))}:
+        if k not in result:
+            result[k] = flat_ref[k]
+
+    return flax.traverse_util.unflatten_dict(result)

pibehavior_da3_clean_up_your_desk_40k/source_code/training/_da3_contrib_check.py

@@ -0,0 +1,63 @@
+"""Measure DA3 contribution on a trained checkpoint, on CORRECTED geometry.
+CPU-only (extractor on CPU) so it never touches the training GPUs.
+
+Generates one real batch through the fixed-geometry loader, loads the checkpoint, then compares
+loss WITH banks vs banks-zeroed (same noise). If delta ~ 0 => the model IGNORES spatial (the
+zero-init failure mode). If delta > 0 => DA3 is load-bearing.
+"""
+import os, sys, dataclasses, json
+import numpy as np, jax, jax.numpy as jnp, flax.nnx as nnx
+sys.path.insert(0, "scripts")
+
+CKPT = os.environ.get("CONTRIB_CKPT",
+    "/work/jack/behavior-1k-solution/outputs/checkpoints/pi_behavior_b1k_fast/b1k_da3_robopro_desk/14000/params")
+
+
+def main():
+    os.environ.update(USE_DA3_FULL="1", DA3_INIT_STD="0.01", DA3_LOGIT_GAIN="1",
+                      B1K_ACTIVITIES="clean_up_your_desk", B1K_EXTRACT_DEVICES="cpu",
+                      B1K_2026_ROOT="/work/jack/behavior1k/data/behavior_2026_task29_42",
+                      B1K_INIT_PARAMS=CKPT)
+    import train_2026 as t26
+    from b1k.training.b1k_da3 import create_v3_behavior_da3_loader
+
+    config = dataclasses.replace(t26.build_config(), batch_size=2, num_workers=0)
+    import jax as _j
+    sharding = _j.sharding.SingleDeviceSharding(_j.devices("cpu")[0])
+    print("[1] building CPU loader (extractor on CPU, corrected geometry)...", flush=True)
+    loader = create_v3_behavior_da3_loader(
+        config, os.environ["B1K_2026_ROOT"], ["clean_up_your_desk"], "/work/jack/behavior1k/task_data.json",
+        lang_cache="/work/jack/behavior1k/modernbert_b1k_tasks.pkl", sharding=sharding, shuffle=True, num_workers=0)
+    obs, actions = next(iter(loader))
+    print("[2] got one fixed-geometry batch; loading checkpoint", flush=True)
+
+    model = config.model.create(jax.random.key(0))
+    loaded = config.weight_loader.load(jax.tree.map(np.asarray, nnx.state(model).to_pure_dict()))
+    gd, st = nnx.split(model); st.replace_by_pure_dict(loaded); model = nnx.merge(gd, st)
+    raw = json.load(open("/work/jack/behavior1k/checkpoints/behavior_50t_checkpoint/assets/IliaLarchenko/behavior_224_rgb/norm_stats.json"))["norm_stats"]
+    model.load_correlation_matrix({"actions": {k: (np.asarray(v, np.float32) if isinstance(v, list) else v) for k, v in raw["actions"].items()}})
+
+    rng = jax.random.key(7); NFS = 8
+    ld_with = model.compute_detailed_loss(rng, obs, actions, train=False, num_flow_samples=NFS)
+    orig = model._compute_banks; model._compute_banks = lambda o: None
+    ld_wo = model.compute_detailed_loss(rng, obs, actions, train=False, num_flow_samples=NFS)
+    model._compute_banks = orig
+    aw, ao = float(jnp.mean(ld_with["action_loss"])), float(jnp.mean(ld_wo["action_loss"]))
+    tw, to = float(jnp.mean(ld_with["total_loss"])), float(jnp.mean(ld_wo["total_loss"]))
+    print(f"[3] action_loss WITH banks={aw:.5f}  WITHOUT(zeroed)={ao:.5f}  delta={ao-aw:+.5f} ({100*(ao-aw)/max(aw,1e-6):+.1f}%)", flush=True)
+    print(f"[3] total_loss  WITH banks={tw:.5f}  WITHOUT(zeroed)={to:.5f}  delta={to-tw:+.5f}", flush=True)
+    # inspect how far the injection out_proj has grown off zero-init
+    flat = nnx.state(model, nnx.Param).flat_state()
+    import math
+    def gnorm(sub):
+        return math.sqrt(sum(float((np.asarray(v.value)**2).sum()) for k,v in flat.items() if sub in "/".join(str(p) for p in k)))
+    print(f"[4] injection out_proj L2={gnorm('out_proj'):.4f}  merge_proj L2={gnorm('merge_proj'):.4f} "
+          f"(near 0 => injection stayed off = model ignoring spatial)", flush=True)
+    verdict = ("DA3 IGNORED (zero-init failure — switch to nonzero init)" if (ao-aw) < 0.005 else
+               "DA3 load-bearing" if (ao-aw) > 0.02 else "DA3 weakly contributing")
+    print(f"[VERDICT] {verdict}", flush=True)
+    print("CONTRIB CHECK DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()

pibehavior_da3_clean_up_your_desk_40k/source_code/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"]

pibehavior_da3_clean_up_your_desk_40k/source_code/training/b1k_da3.py

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

pibehavior_da3_clean_up_your_desk_40k/source_code/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),
+        )

pibehavior_da3_clean_up_your_desk_40k/source_code/training/train.py

@@ -0,0 +1,489 @@
+"""
+Training script for BEHAVIOR-1K solution.
+
+Based on https://github.com/PhysicalIntelligence/openpi/blob/behavior/openpi/scripts/train.py with custom modifications.
+"""
+
+import dataclasses
+import functools
+import logging
+import os
+import platform
+import time
+from typing import Any
+
+import etils.epath as epath
+import flax.nnx as nnx
+from flax.training import common_utils
+import flax.traverse_util as traverse_util
+import jax
+import jax.experimental
+import jax.numpy as jnp
+import numpy as np
+import optax
+import tqdm_loggable.auto as tqdm
+import wandb
+
+# Configure JAX memory allocation to prevent OOM errors
+os.environ.setdefault('XLA_PYTHON_CLIENT_MEM_FRACTION', '0.9')
+os.environ.setdefault('XLA_PYTHON_CLIENT_ALLOCATOR', 'platform')
+
+# Configure OpenBLAS to prevent thread creation errors
+os.environ.setdefault('OPENBLAS_NUM_THREADS', '16')
+os.environ.setdefault('MKL_NUM_THREADS', '16')
+
+import openpi.models.model as _model
+import openpi.shared.array_typing as at
+import openpi.shared.nnx_utils as nnx_utils
+import openpi.training.optimizer as _optimizer
+import openpi.training.sharding as sharding
+import openpi.training.utils as training_utils
+
+# Import B1K-specific modules
+from b1k.training import checkpoints as _checkpoints  # Use our custom checkpoints (not openpi's!)
+from b1k.training import config as _config
+from b1k.training import data_loader as _data_loader
+from b1k.training import weight_loaders as _weight_loaders
+from b1k.models.pi_behavior import PiBehavior
+from b1k.models.pi_behavior_config import PiBehaviorConfig
+from b1k.models.observation import Observation
+
+
+def init_logging():
+    """Custom logging format for better readability."""
+    level_mapping = {"DEBUG": "D", "INFO": "I", "WARNING": "W", "ERROR": "E", "CRITICAL": "C"}
+
+    class CustomFormatter(logging.Formatter):
+        def format(self, record):
+            record.levelname = level_mapping.get(record.levelname, record.levelname)
+            return super().format(record)
+
+    formatter = CustomFormatter(
+        fmt="%(asctime)s.%(msecs)03d [%(levelname)s] %(message)-80s (%(process)d:%(filename)s:%(lineno)s)",
+        datefmt="%H:%M:%S",
+    )
+
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    logger.handlers[0].setFormatter(formatter)
+
+
+def init_wandb(config: _config.TrainConfig, *, resuming: bool, log_code: bool = False, enabled: bool = True):
+    if not enabled:
+        wandb.init(mode="disabled")
+        return
+
+    ckpt_dir = config.checkpoint_dir
+    if not ckpt_dir.exists():
+        raise FileNotFoundError(f"Checkpoint directory {ckpt_dir} does not exist.")
+    if resuming:
+        run_id = (ckpt_dir / "wandb_id.txt").read_text().strip()
+        wandb.init(id=run_id, resume="must", project=config.project_name)
+    else:
+        wandb.init(
+            name=config.exp_name,
+            config=dataclasses.asdict(config),
+            project=config.project_name,
+        )
+        (ckpt_dir / "wandb_id.txt").write_text(wandb.run.id)
+
+    if log_code:
+        wandb.run.log_code(epath.Path(__file__).parent.parent)
+
+
+def _load_weights_and_validate(loader: _weight_loaders.WeightLoader, params_shape: at.Params) -> at.Params:
+    """Loads and validates the weights. Returns a loaded subset of the weights."""
+    loaded_params = loader.load(params_shape)
+    
+    # Filter out nnx.Intermediate fields from both sides (they're not params, excluded from checkpoints)
+    # This allows loading old checkpoints that didn't have these fields
+    def filter_intermediate_fields(params_dict):
+        flat = traverse_util.flatten_dict(params_dict)
+        # List of field names that are nnx.Intermediate (excluded from checkpoints)
+        intermediate_field_names = [
+            'action_correlation_cholesky',  # Legacy full correlation matrix
+            'L_spatial',                     # Separable spatial correlation
+            'L_temporal',                    # Separable temporal correlation
+            'cached_num_inpaint_actions',    # Conditional sampling cache
+            'cached_input_action_dim',       # Conditional sampling cache
+            'cached_Sigma_uo_Sigma_oo_inv',  # Conditional sampling cache
+            'cached_L_cond_free',            # Conditional sampling cache
+            'cached_Sigma_ou_Sigma_uu_inv',  # Conditional sampling cache
+            'cached_L_cond_inp',             # Conditional sampling cache
+        ]
+        filtered = {k: v for k, v in flat.items() 
+                   if not any(field in str(k) for field in intermediate_field_names)}
+        return traverse_util.unflatten_dict(filtered)
+    
+    # Validate loaded params structure  
+    params_shape_filtered = filter_intermediate_fields(params_shape)
+    loaded_params_filtered = filter_intermediate_fields(loaded_params)
+    at.check_pytree_equality(expected=params_shape_filtered, got=loaded_params_filtered, check_shapes=True, check_dtypes=True)
+    
+    # Remove jax.ShapeDtypeStruct and Intermediate fields from the loaded params
+    def should_exclude(k, v):
+        if isinstance(v, jax.ShapeDtypeStruct):
+            return True
+        # Exclude all intermediate fields
+        intermediate_field_names = [
+            'action_correlation_cholesky', 'L_spatial', 'L_temporal',
+            'cached_num_inpaint_actions', 'cached_input_action_dim',
+            'cached_Sigma_uo_Sigma_oo_inv', 'cached_L_cond_free',
+            'cached_Sigma_ou_Sigma_uu_inv', 'cached_L_cond_inp',
+        ]
+        return any(field in str(k) for field in intermediate_field_names)
+    
+    return traverse_util.unflatten_dict(
+        {k: v for k, v in traverse_util.flatten_dict(loaded_params).items() 
+         if not should_exclude(k, v)}
+    )
+
+
+@at.typecheck
+def init_train_state(
+    config: _config.TrainConfig, 
+    init_rng: at.KeyArrayLike, 
+    mesh: jax.sharding.Mesh, 
+    *, 
+    resume: bool,
+    norm_stats: dict | None = None
+) -> tuple[training_utils.TrainState, Any]:
+    # DA3 v2 multi-group LR: the fresh spatial modules (bank builder + injection) and the geometry
+    # ray-MLP need a much higher LR than the pretrained trunk, or they barely move off init. Keyed on
+    # the same env the DA3 entrypoint sets; falls back to the single-group optimizer otherwise.
+    if os.environ.get("USE_DA3_FULL") == "1" and os.environ.get("DA3_LR_GROUPS", "1") == "1":
+        base = config.lr_schedule
+        def _sched(peak, end):
+            return dataclasses.replace(base, peak_lr=float(peak), decay_lr=float(end))
+        lr_groups = {
+            # pretrained SigLIP + gemma trunk + action expert: openpi pi0.5 finetune recipe
+            "vlm": _sched(os.environ.get("LR_VLM_PEAK", 2.5e-5), os.environ.get("LR_VLM_END", 2.5e-6)),
+            # fresh spatial branch (bank builder + cross-attn injection)
+            "core": _sched(os.environ.get("LR_CORE_PEAK", 5e-4), os.environ.get("LR_CORE_END", 5e-5)),
+            # fresh geometry ray-MLP
+            "geom": _sched(os.environ.get("LR_GEOM_PEAK", 5e-4), os.environ.get("LR_GEOM_END", 1e-4)),
+        }
+        # Optional delayed/slow phase-in for the geometry group: hold at 0 for LR_GEOM_DELAY steps
+        # (while vlm/core do their normal warmup), then ramp 0->peak over LR_GEOM_RAMP steps before
+        # the usual cosine decay. Off unless LR_GEOM_RAMP is set, so existing recipes are unchanged.
+        _g_ramp = int(os.environ.get("LR_GEOM_RAMP", "0"))
+        if _g_ramp > 0:
+            lr_groups["geom"] = _optimizer.DelayedRampCosineSchedule(
+                delay_steps=int(os.environ.get("LR_GEOM_DELAY", "0")),
+                ramp_steps=_g_ramp,
+                peak_lr=float(os.environ.get("LR_GEOM_PEAK", 5e-4)),
+                decay_steps=base.decay_steps,
+                decay_lr=float(os.environ.get("LR_GEOM_END", 1e-4)),
+            )
+            logging.info("DA3 geom LR: delayed ramp (0 for %d steps, 0->peak over %d, then cosine to %s)",
+                         int(os.environ.get("LR_GEOM_DELAY", "0")), _g_ramp,
+                         os.environ.get("LR_GEOM_END", 1e-4))
+        logging.info("DA3 multi-group LR: vlm=%s core=%s geom=%s",
+                     lr_groups["vlm"].peak_lr, lr_groups["core"].peak_lr, lr_groups["geom"].peak_lr)
+        tx = _optimizer.create_multi_group_optimizer(
+            config.optimizer, lr_groups, _optimizer.spatial_group_labels, weight_decay_mask=None)
+    else:
+        tx = _optimizer.create_optimizer(config.optimizer, config.lr_schedule, weight_decay_mask=None)
+
+    def init(rng: at.KeyArrayLike, partial_params: at.Params | None = None) -> training_utils.TrainState:
+        rng, model_rng = jax.random.split(rng)
+        # initialize the model (and its parameters).
+        model = config.model.create(model_rng)
+        
+        # Load correlation matrix into PiBehavior models BEFORE creating graphdef
+        if isinstance(model, PiBehavior) and norm_stats is not None:
+            model.load_correlation_matrix(norm_stats)
+            logging.info("Loaded correlation matrix during model initialization")
+
+        # Merge the partial params into the model.
+        if partial_params is not None:
+            graphdef, state = nnx.split(model)
+            # This will produce an error if the partial params are not a subset of the state.
+            state.replace_by_pure_dict(partial_params)
+            model = nnx.merge(graphdef, state)
+
+        params = nnx.state(model)
+        params = nnx_utils.state_map(params, config.freeze_filter, lambda p: p.replace(p.value.astype(jnp.bfloat16)))
+        return training_utils.TrainState(
+            step=0,
+            params=params,
+            model_def=nnx.graphdef(model),
+            tx=tx,
+            opt_state=tx.init(params.filter(config.trainable_filter)),
+            ema_decay=config.ema_decay,
+            ema_params=None if config.ema_decay is None else params,
+        )
+
+    train_state_shape = jax.eval_shape(init, init_rng)
+    state_sharding = sharding.fsdp_sharding(train_state_shape, mesh, log=True)
+
+    if resume:
+        return train_state_shape, state_sharding
+
+    partial_params = _load_weights_and_validate(config.weight_loader, train_state_shape.params.to_pure_dict())
+    replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+
+    # Initialize the train state and mix in the partial params.
+    train_state = jax.jit(
+        init,
+        donate_argnums=(1,),  # donate the partial params buffer.
+        in_shardings=replicated_sharding,
+        out_shardings=state_sharding,
+    )(init_rng, partial_params)
+    
+    # Log KV transform coefficients for PiBehavior models
+    model = nnx.merge(train_state.model_def, train_state.params)
+    if isinstance(model, PiBehavior) and hasattr(model, 'kv_transform') and model.kv_transform is not None:
+        logging.info("KV Transform Coefficients (after loading):")
+        logging.info("=" * 80)
+        
+        k_coeffs = model.kv_transform.k_coeffs.value
+        v_coeffs = model.kv_transform.v_coeffs.value
+        
+        logging.info("K Coefficients (each layer attends to all VLM layers):")
+        for i in range(k_coeffs.shape[0]):
+            coeffs_str = ", ".join([f"{float(c):.2f}" for c in k_coeffs[i]])
+            logging.info(f"  Layer {i:2d}: [{coeffs_str}]")
+        
+        logging.info("")
+        logging.info("V Coefficients (each layer attends to all VLM layers):")
+        for i in range(v_coeffs.shape[0]):
+            coeffs_str = ", ".join([f"{float(c):.2f}" for c in v_coeffs[i]])
+            logging.info(f"  Layer {i:2d}: [{coeffs_str}]")
+        
+        logging.info("=" * 80)
+
+    return train_state, state_sharding
+
+
+@at.typecheck
+def train_step(
+    config: _config.TrainConfig,
+    rng: at.KeyArrayLike,
+    state: training_utils.TrainState,
+    batch: tuple[Observation, _model.Actions],
+) -> tuple[training_utils.TrainState, dict[str, at.Array]]:
+    model = nnx.merge(state.model_def, state.params)
+    model.train()
+
+    @at.typecheck  
+    def loss_fn(
+        model: PiBehavior, rng: at.KeyArrayLike, observation: Observation, actions: _model.Actions
+    ):
+        losses_dict = model.compute_detailed_loss(rng, observation, actions, train=True, num_flow_samples=config.num_flow_samples)
+        total_loss = jnp.mean(losses_dict["total_loss"])
+        return total_loss, losses_dict
+
+    train_rng = jax.random.fold_in(rng, state.step)
+    observation, actions = batch
+
+    # Filter out frozen params.
+    diff_state = nnx.DiffState(0, config.trainable_filter)
+    (loss, losses_dict), grads = nnx.value_and_grad(loss_fn, argnums=diff_state, has_aux=True)(model, train_rng, observation, actions)
+    
+    # Knowledge insulation gradient monitoring
+    if config.model.use_knowledge_insulation:
+        # Helper functions to identify parameter groups
+        def is_action_expert_param(path_str):
+            # Action expert parameters: 
+            # - Second LLM expert (300M params, marked with _1 suffix)
+            # - Action projections, time MLPs, kv_transform
+            return any(x in path_str for x in [
+                "_1",  # All second expert parameters
+                "action_in_proj",
+                "action_out_proj",
+                "time_mlp_in",
+                "time_mlp_out",
+                "kv_transform"
+            ])
+        
+        def is_vlm_param(path_str):
+            # VLM parameters: everything else (first expert, img, FAST, task modules)
+            return not is_action_expert_param(path_str)
+        
+        # Compute gradient norms for monitoring only (no scaling applied)
+        def compute_group_norm(grads_state, predicate):
+            """Compute norm for gradients matching predicate."""
+            flat_grads = []
+            for path, value in jax.tree_util.tree_flatten_with_path(grads_state.to_pure_dict())[0]:
+                path_str = "/".join(str(k) for k in path)
+                if predicate(path_str):
+                    if hasattr(value, 'value'):
+                        flat_grads.append(value.value if hasattr(value, 'value') else value)
+                    else:
+                        flat_grads.append(value)
+            
+            if flat_grads:
+                return jnp.sqrt(sum(jnp.sum(jnp.square(g)) for g in flat_grads))
+            return 0.0
+        
+        grad_norm_vlm = compute_group_norm(grads, is_vlm_param)
+        grad_norm_action = compute_group_norm(grads, is_action_expert_param)
+    else:
+        grad_norm_vlm = None
+        grad_norm_action = None
+
+    params = state.params.filter(config.trainable_filter)
+    updates, new_opt_state = state.tx.update(grads, state.opt_state, params)
+    new_params = optax.apply_updates(params, updates)
+
+    # Update the model in place and return the new full state.
+    nnx.update(model, new_params)
+    new_params = nnx.state(model)
+
+    new_state = dataclasses.replace(state, step=state.step + 1, params=new_params, opt_state=new_opt_state)
+    if state.ema_decay is not None:
+        new_state = dataclasses.replace(
+            new_state,
+            ema_params=jax.tree.map(
+                lambda old, new: state.ema_decay * old + (1 - state.ema_decay) * new, state.ema_params, new_params
+            ),
+        )
+
+    # Filter out params that aren't kernels.
+    kernel_params = nnx.state(
+        model,
+        nnx.All(
+            nnx.Param,
+            nnx.Not(nnx_utils.PathRegex(".*/(bias|scale|pos_embedding|input_embedding)")),
+            lambda _, x: x.value.ndim > 1,
+        ),
+    )
+    info = {
+        "loss": loss,
+        "grad_norm": optax.global_norm(grads),
+        "param_norm": optax.global_norm(kernel_params),
+    }
+    
+    # Add gradient norm breakdown for knowledge insulation monitoring
+    if grad_norm_vlm is not None:
+        info["grad_norm_vlm"] = grad_norm_vlm
+        info["grad_norm_action_expert"] = grad_norm_action
+
+    # Add detailed loss components to info
+    for key, value in losses_dict.items():
+        if isinstance(value, (float, int)) or (hasattr(value, 'ndim') and value.ndim == 0):
+            info[key] = value
+        else:
+            info[key] = jnp.mean(value)
+    return new_state, info
+
+
+def main(config: _config.TrainConfig):
+    init_logging()
+    logging.info(f"Running on: {platform.node()}")
+
+    if config.batch_size % jax.device_count() != 0:
+        raise ValueError(
+            f"Batch size {config.batch_size} must be divisible by the number of devices {jax.device_count()}."
+        )
+
+    jax.config.update("jax_compilation_cache_dir", str(epath.Path("~/.cache/jax").expanduser()))
+
+    # Generate random seed if not provided
+    seed = config.seed
+    if seed is None:
+        seed = int(time.time() * 1000) % (2**32)
+        logging.info(f"Using random seed for JAX RNG: {seed}")
+    
+    rng = jax.random.key(seed)
+    train_rng, init_rng = jax.random.split(rng)
+
+    mesh = sharding.make_mesh(config.fsdp_devices)
+    data_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(sharding.DATA_AXIS))
+    replicated_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+
+    checkpoint_manager, resuming = _checkpoints.initialize_checkpoint_dir(
+        config.checkpoint_dir,
+        keep_period=config.keep_period,
+        overwrite=config.overwrite,
+        resume=config.resume,
+    )
+    init_wandb(config, resuming=resuming, enabled=config.wandb_enabled)
+
+    data_loader = _data_loader.create_behavior_data_loader(
+        config,
+        sharding=data_sharding,
+        shuffle=True,
+    )
+
+    data_iter = iter(data_loader)
+    batch = next(data_iter)
+    logging.info(f"Initialized data loader:\n{training_utils.array_tree_to_info(batch)}")
+
+    if os.environ.get("SKIP_IMAGE_LOG", "0") != "1":
+        # Log images from first batch to sanity check.
+        images_to_log = [
+            wandb.Image(np.concatenate([np.array(img[i]) for img in batch[0].images.values()], axis=1))
+            for i in range(min(5, len(next(iter(batch[0].images.values())))))
+        ]
+        wandb.log({"camera_views": images_to_log}, step=0)
+
+    # Get norm_stats for correlation matrix loading
+    data_config = data_loader.data_config()
+    if data_config.norm_stats is None:
+        raise ValueError(
+            "norm_stats not found. Run compute_norm_stats.py to generate normalization statistics."
+        )
+    norm_stats = data_config.norm_stats
+
+    train_state, train_state_sharding = init_train_state(
+        config, init_rng, mesh, resume=resuming, norm_stats=norm_stats
+    )
+    jax.block_until_ready(train_state)
+    logging.info(f"Initialized train state:\n{training_utils.array_tree_to_info(train_state.params)}")
+
+    if resuming:
+        train_state = _checkpoints.restore_state(checkpoint_manager, train_state, data_loader)
+        
+        # Reload correlation matrix after restore
+        model = nnx.merge(train_state.model_def, train_state.params)
+        model.load_correlation_matrix(norm_stats)
+        logging.info("Reloaded correlation matrix after checkpoint restore")
+        train_state = dataclasses.replace(train_state, model_def=nnx.graphdef(model))
+
+    ptrain_step = jax.jit(
+        functools.partial(train_step, config),
+        in_shardings=(replicated_sharding, train_state_sharding, data_sharding),
+        out_shardings=(train_state_sharding, replicated_sharding),
+        donate_argnums=(1,),
+    )
+
+    start_step = int(train_state.step)
+    pbar = tqdm.tqdm(
+        range(start_step, config.num_train_steps),
+        initial=start_step,
+        total=config.num_train_steps,
+        dynamic_ncols=True,
+    )
+
+    infos = []
+    for step in pbar:
+        with sharding.set_mesh(mesh):
+            train_state, info = ptrain_step(train_rng, train_state, batch)
+        infos.append(info)
+        if step % config.log_interval == 0:
+            stacked_infos = common_utils.stack_forest(infos)
+            reduced_info = jax.device_get(jax.tree.map(jnp.mean, stacked_infos))
+            
+            # Create a concise console log with main metrics
+            main_metrics = {k: v for k, v in reduced_info.items() 
+                          if "loss" in k or "accuracy" in k or k in ["grad_norm", "param_norm", "grad_norm_vlm", "grad_norm_action_expert"]}
+            info_str = ", ".join(f"{k}={v:.4f}" for k, v in main_metrics.items())
+            pbar.write(f"Step {step}: {info_str}")
+            wandb.log(reduced_info, step=step)
+            infos = []
+        batch = next(data_iter)
+
+        skip_final_save = os.environ.get("SKIP_FINAL_SAVE", "0") == "1"
+        should_save_interval = step % config.save_interval == 0 and step > start_step
+        should_save_final = step == config.num_train_steps - 1 and not skip_final_save
+        if should_save_interval or should_save_final:
+            _checkpoints.save_state(checkpoint_manager, train_state, data_loader, step)
+
+    logging.info("Waiting for checkpoint manager to finish")
+    checkpoint_manager.wait_until_finished()
+
+
+if __name__ == "__main__":
+    main(_config.cli())

pibehavior_da3_clean_up_your_desk_40k/source_code/training/train_2026.py

@@ -0,0 +1,135 @@
+"""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 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"))),
+            ),
+        )
+    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")),
+        )
+    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,
+        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())