T5-DA3-#2: eval README, source bundles, and the DA3 eval-input builder

T5-DA3-#2/README.md

@@ -0,0 +1,154 @@
+# T5-DA3-#2 — BEHAVIOR-1K task 5 (`setting_mousetraps`), pi0.5 + DA3 spatial injection
+
+Terminal checkpoint, **step 32999 of 33000**. Final training `action_loss` ≈ **0.022**.
+
+Fine-tuned from the 100-task meta checkpoint (`meta100-1epoch/step35000`) on task 5 only, with a
+frozen DA3-GIANT-1.1 backbone supplying per-patch geometry features that are cross-attended into the
+last 6 blocks of the action expert.
+
+---
+
+## ⚠ Read this first
+
+**The stock eval path does not support DA3.** `grep -rn "da3\|DA3" src/b1k/policies/ src/b1k/shared/`
+returns nothing. If you serve this checkpoint through `scripts/serve_b1k.py` unchanged:
+
+- `observation.da3_features` is `None`
+- `PiBehavior._compute_banks` returns `None` (`src/b1k/models/pi_behavior.py:372`)
+- the spatial injection is skipped **silently — no error, no warning**
+
+You get a policy running without the branch it was trained with. Use `code/da3_eval_inputs.py`
+(included here) to build the missing fields.
+
+**The architecture flags are not stored in the checkpoint.** They come from environment variables
+read at model-construction time. If they don't match, the parameter tree won't match either — some
+mismatches raise, others load partially. Export the block below before building the model.
+
+---
+
+## 1. Environment (required, exact)
+
+```bash
+# --- architecture: these DEFINE the parameter tree, not just behaviour ---
+export USE_DA3_FULL=1 B1K_USE_GT_DEPTH=1 B1K_GRASP_AUX=0
+export DA3_MODEL_NAME=depth-anything/DA3-GIANT-1.1
+export DA3_OUT_LAYERS=19,27,33,39 DA3_CHANNELS=1536
+export DA3_GRID_H=16 DA3_GRID_W=16 DA3_BACKBONE_ONLY=1
+export DA3_USE_PERCEIVER=0 DA3_LANG_DEPTH=0
+export DA3_SPATIAL_VEC=1 DA3_SPATIAL_FILM=1 DA3_SPATIAL_USE_DA3=1
+export DA3_FOURIER_BANDS=10 DA3_BANK_TOKEN_EMBED=1 DA3_BANK_CENTER=0
+export DA3_QK_NORM=1 DA3_LOGIT_GAIN=1 DA3_LOGIT_GAIN_INIT=3.0 DA3_INJ_GAIN_MAX=8.0
+export DA3_INIT_STD=0.01 DA3_SCALE=1.0 DA3_EMBED_GROUP=0
+
+# --- this run's architecture changes vs the earlier DA3-#1 checkpoints ---
+export DA3_SAMPLE_TOKEN_CENTER=1     # per-sample token centering on the final bank
+export DA3_FEAT_INPUT_NORM=1         # LayerNorm the raw DA3 taps before projection
+export DA3_TAP_TOKEN_CENTER=1        # per-tap token centering before layer_fuse
+export DA3_TAP_GAIN=1                # learnable per-tap scalar gain
+export DA3_CENTER_FUSED=1            # strip the latent DC inside SpatialConditioner
+export DA3_FILM_INIT_STD=0.01        # FiLM gamma/beta non-zero init
+export DA3_S1=1                      # merge_proj reads the ATTENTION OUTPUTS
+export DA3_INJ_NO_BIAS=1             # no bias on merge_proj / out_proj
+export DA3_INJECT_DROPOUT=0.0        # inference: always 0
+
+# --- task / assets ---
+export B1K_ACTIVITIES=setting_mousetraps B1K_TASK_SPACE=100
+export B1K_ASSETS_BASE=<dir containing assets/ from this repo folder>
+export DA3_LANG_CACHE=<path>/modernbert_b1k_tasks.pkl
+export B1K_TASK_DATA_JSON=<path>/task_data.json
+export DA3_FREEZE_BASE_HARD=0 DA3_LR_GROUPS=1 DA3_TASK5_STAGED=1
+```
+
+`DA3_S1=1` and `DA3_INJ_NO_BIAS=1` change the *shape* of the injection block. Loading this checkpoint
+with `DA3_S1=0` will not reproduce the trained function even if it loads.
+
+## 2. Files in this folder
+
+| path | what |
+|---|---|
+| `params/` | orbax params (13 GB). This is what you load. |
+| `assets/IliaLarchenko/behavior_224_rgb/norm_stats.json` | action/state normalization + the correlated-noise matrix |
+| `_CHECKPOINT_METADATA` | orbax metadata |
+| `code/b1k_code.tgz` | the full training/serving source at the exact commit state this was trained with (`src/`, `scripts/`, `openpi/src/`) |
+| `code/da3_geostack.tgz` | `da3_for_geostack.py`, the DA3-GIANT wrapper the extractor loads |
+| `code/b1k_support.tgz` | `modernbert_b1k_tasks.pkl` (task language features), `task_data.json` |
+| `code/da3_eval_inputs.py` | **the missing piece** — builds the DA3 observation fields at eval time |
+
+`train_state/` was deliberately not uploaded (another ~34 GB). Ask if you need to resume training.
+
+## 3. Inputs the model expects
+
+Standard b1k fields, built by the existing wrapper:
+
+| field | shape | notes |
+|---|---|---|
+| `image` | 3 × `[224,224,3]` uint8 | `base_0_rgb` (head/zed), `left_wrist_0_rgb`, `right_wrist_0_rgb` |
+| `state` | `[23]` | `extract_state_2026` of the 61-d proprio: base_qvel(3), trunk(4), left arm(7), left grip(1), right arm(7), right grip(1) |
+| `tokenized_prompt` | `[2]` | `[task_id, stage_id]`; task 5 = `setting_mousetraps`, in the 100-task space |
+
+DA3 fields — **you must add these** (see `code/da3_eval_inputs.py`):
+
+| field | shape | notes |
+|---|---|---|
+| `da3_features` | `[B,4,3,1536,16,16]` **uint16** | bf16 *bits*, not floats. The model bitcasts them back. taps 19/27/33/39, 3 views, 16×16 patch grid |
+| `da3_ray` | `[B,3,3,16,16]` f32 | unit camera rays through exact patch centres, from the intrinsics |
+| `da3_depth` | `[B,3,1,16,16]` f32 | metric **metres**, area-averaged over each patch's 14×14 pixel footprint |
+| `camera_extrinsics` | `[B,3,4,4]` f32 | robot→cam OpenCV, from `pose7_to_robot2cam` |
+| `lang_feat` / `lang_mask` | `[B,32,1024]` / `[B,32]` | ModernBERT features for the task string, from the pickle |
+
+View order is fixed everywhere: **head, left wrist, right wrist**. Depth must be *simulator/GT* depth
+at native camera resolution before the 16×16 area-average — the model was never trained on DA3's own
+predicted depth (`DA3_BACKBONE_ONLY=1` skips that head entirely). **Enable `depth_linear` on all three
+cameras in your OmniGibson env config**; without it this checkpoint cannot be evaluated faithfully.
+
+Intrinsics are derived, not measured: `fx = fy = (17.0/20.995) · 224`, `cx = cy = 112`.
+
+## 4. Output and how to decode it
+
+The model returns an action chunk of shape `[B, 30, 32]` in **normalized** space:
+
+1. Only the **first 23 dims are real**; 23–31 are padding and are unconstrained by the loss — discard them.
+2. Unnormalize with `assets/.../norm_stats.json` (`actions` block). Training used quantile
+   normalization; use the same path the existing b1k eval wrapper uses — do not hand-roll it.
+3. Field order matches the state: `base_qvel(3) | trunk(4) | left_arm(7) | left_gripper(1) | right_arm(7) | right_gripper(1)` = 23.
+4. Gripper dims are normalized against `MAX_GRIPPER_WIDTH = 0.1`.
+
+Sampling is flow-matching: `sample_actions(rng, obs, num_steps=10)` (10 denoising steps is what the
+probes used; the trainer used 15 flow samples for the loss, which is unrelated). It returns
+`(actions, subtask_logits)` — take `[0]`.
+
+## 5. Minimal loading sketch
+
+```python
+config = build_config()               # scripts/train_2026_da3_task5.py, with the env above set
+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(f"{ASSETS}/norm_stats.json"))["norm_stats"]["actions"]
+model.load_correlation_matrix({"actions": {k: np.asarray(v, np.float32) if isinstance(v, list) else v
+                                           for k, v in raw.items()}})
+```
+Point `B1K_INIT_PARAMS` at this folder's `params/`. `config.weight_loader` is
+`PiBehaviorWeightLoader`, which reads it directly.
+
+Sanity check before trusting any eval number: run one batch and confirm the injection is live —
+`model._compute_banks(obs)` must return a dict of three `[B,256,1024]` banks, **not None**. If it
+returns None, `da3_features` didn't reach the model and you are evaluating a different model.
+
+## 6. What changed vs DA3-#1 (context for interpreting results)
+
+DA3-#1 was trained with `init_std=0.05`, no token centering, and the original merge wiring. Measured
+problems, all fixed here:
+
+- its bank was **71–80% a fixed DC vector** (the same vector for ~98% of samples), so the injection
+  attention had almost nothing to discriminate on;
+- the injection had a **bank-free path** — `left_branch`/`right_branch`/`merge_proj` composed into a
+  learned linear map of the action stream that never touched the geometry, so ablating the branch
+  looked catastrophic while shuffling the geometry looked free;
+- the 0.05 init multiplied the residual stream by ~2× at step 0, costing ~7–8k steps of repair.
+
+At checkpoint 10000 of this run, shuffling geometry across the batch costs **3× more** than removing
+the injection entirely (+2.30% vs +0.78%) — i.e. the model reads per-sample geometry, and giving it
+the *wrong* scene is worse than giving it none. That is the property that makes correct DA3 inputs at
+eval time matter.

T5-DA3-#2/code/b1k_code.tgz

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

T5-DA3-#2/code/b1k_support.tgz

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

T5-DA3-#2/code/da3_eval_inputs.py

@@ -0,0 +1,88 @@
+"""Build the DA3 observation fields at EVAL time for T5-DA3-#2.
+
+The stock b1k serving path (`scripts/serve_b1k.py` -> `b1k.shared.eval_b1k_wrapper`) has NO DA3
+support: `grep -r da3 src/b1k/policies src/b1k/shared` is empty. If you run this checkpoint through
+it unchanged, `observation.da3_features` is None, `PiBehavior._compute_banks` returns None, and the
+spatial injection is skipped **silently** -- you get a degraded policy and no error. This module
+fills that gap by reproducing the training-time inputs exactly.
+
+Use:
+
+    from da3_eval_inputs import DA3EvalInputs
+    da3 = DA3EvalInputs(lang_cache="support/behavior1k/modernbert_b1k_tasks.pkl",
+                        task_name="setting_mousetraps", device="cuda:0")
+
+    extra = da3(rgb224=[head, left_wrist, right_wrist],      # 3x [224,224,3] uint8, VIEW_ORDER
+                depth_m=[d_head, d_left, d_right],           # 3x [H,W] float32 METRES, native res
+                pose7=[p_head, p_left, p_right])             # 3x [7] cam-in-robot, [xyz, quat wxyz]
+    # -> dict with da3_features, da3_ray, da3_depth, camera_extrinsics, lang_feat, lang_mask
+    # merge into the Observation you pass to the policy (batch axis added for you).
+
+Order is fixed: head (zed), left wrist (left realsense), right wrist (right realsense). Getting it
+wrong silently mispairs geometry with pixels.
+"""
+from __future__ import annotations
+
+import pickle
+
+import numpy as np
+
+FOCAL_RATIO = 17.0 / 20.995   # OmniGibson VisionSensor default; fx = fy = FOCAL_RATIO * width
+GRID = 16                      # 224 / patch 14
+DA3_HW = (224, 224)
+OUT_LAYERS = (19, 27, 33, 39)
+MODEL_NAME = "depth-anything/DA3-GIANT-1.1"
+
+
+class DA3EvalInputs:
+    def __init__(self, lang_cache: str, task_name: str, device: str = "cuda:0",
+                 model_name: str = MODEL_NAME, forward_chunk: int = 8):
+        from b1k.training.da3_extractor import DA3InlineExtractor
+
+        self.ex = DA3InlineExtractor(
+            model_name=model_name, out_layers=OUT_LAYERS, da3_hw=DA3_HW,
+            devices=[device], forward_chunk=forward_chunk,
+            backbone_only=True,     # DA3's own depth head is unused: we feed GT/sim depth
+            out_views=3,            # only the current frame's 3 cameras are read out
+        )
+        with open(lang_cache, "rb") as f:
+            lang = pickle.load(f)
+        feat, mask = lang[task_name]
+        self.lang_feat = np.asarray(feat, np.float32)
+        self.lang_mask = np.asarray(mask, bool)
+        k = np.array([[FOCAL_RATIO * DA3_HW[1], 0.0, DA3_HW[1] / 2.0],
+                      [0.0, FOCAL_RATIO * DA3_HW[0], DA3_HW[0] / 2.0],
+                      [0.0, 0.0, 1.0]], np.float32)
+        self.intr = np.tile(k, (3, 1, 1))
+
+    @staticmethod
+    def _depth_grid(depth_m: np.ndarray) -> np.ndarray:
+        """Native metric depth -> [1,16,16], area-averaged over each patch's pixel footprint.
+
+        INTER_AREA is a box mean, matching training (`b1k_da3._decode_gt_depth`). Depth and RGB share
+        the camera raster and FOV, so cell (i,j) covers exactly the pixels of patch token (i,j).
+        """
+        import cv2
+        d = np.asarray(depth_m, np.float32)
+        if d.ndim == 3:
+            d = d[..., 0]
+        return cv2.resize(d, (GRID, GRID), interpolation=cv2.INTER_AREA)[None]
+
+    def __call__(self, rgb224, depth_m, pose7) -> dict:
+        from b1k.training.b1k_da3 import pose7_to_robot2cam
+
+        assert len(rgb224) == len(depth_m) == len(pose7) == 3, "expect 3 views in VIEW_ORDER"
+        imgs = np.stack([np.asarray(x, np.uint8) for x in rgb224], 0)          # [3,224,224,3]
+        assert imgs.shape[1:3] == DA3_HW, f"pass the SAME 224x224 images the VLM gets, got {imgs.shape}"
+        extr = np.stack([pose7_to_robot2cam(np.asarray(p, np.float32)) for p in pose7], 0)  # [3,4,4]
+        depth = np.stack([self._depth_grid(d) for d in depth_m], 0).astype(np.float32)      # [3,1,16,16]
+
+        feats, ray, _ = self.ex.extract(imgs[None], extr[None], self.intr[None])
+        return {
+            "da3_features": feats,                    # uint16 = bf16 BITS, [1,4,3,1536,16,16]
+            "da3_ray": np.asarray(ray, np.float32),   # [1,3,3,16,16] unit rays at patch centres
+            "da3_depth": depth[None],                 # [1,3,1,16,16] metric metres
+            "camera_extrinsics": extr[None],          # [1,3,4,4] robot->cam OpenCV
+            "lang_feat": self.lang_feat[None],        # [1,32,1024] ModernBERT, per task
+            "lang_mask": self.lang_mask[None],        # [1,32]
+        }

T5-DA3-#2/code/da3_geostack.tgz

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