add eval I/O contract README

pi_behavior_100t_meta_ft_task56_da3/README.md

@@ -1,15 +1,121 @@
-# PiBehavior — 100-task meta → task-56 DA3 finetune (BEHAVIOR-1K 2026)
+# PiBehavior DA3 — Evaluation I/O Contract
 
-Single-task DA3 finetune of **make_rose_centerpieces** (activity 56) from the 100-task meta checkpoint.
+This checkpoint is a **single-task DA3 finetune** of the PiBehavior VLA (BEHAVIOR-1K 2026), initialized from the 100-task meta (`meta100-1epoch/step69999`). It is a pi0-style flow-matching policy: PaliGemma **gemma_2b** VLM + **gemma_300m** action expert (bf16), `num_tasks=100`, `action_horizon=30`, `action_dim=32`, 3 RGB views @ 224², **plus a DA3 depth/spatial branch**.
 
-## Recipe
-- **Init:** 100-task meta (`meta100-1epoch/step69999`), plain PiBehaviorWeightLoader (100-wide table).
-- **Model:** PiBehavior (PaliGemma gemma_2b + gemma_300m, bf16), action_horizon=30, action_dim=32, num_tasks=100, correlated noise (beta=0.5), FAST-aux (w=0.05), frozen vision backbone, num_flow_samples=15.
-- **DA3 branch:** DA3-GIANT-1.1 (backbone-only, layers 19/27/33/39, 1536ch, 16x16), GT metric depth, camera-frame ray, spatial vec+FiLM inject (6 layers), **cross_view OFF**.
-- **Optimizer:** batch 128, FSDP 4-GPU, LR ramp 1e-8 -> 5e-5 (1k warmup) -> cosine 1e-6, **20,000 steps** (final ckpt 19999). qvel-fixed norm stats.
-- **Data:** b1k-224x224-gop8-fixed (224 RGB) + b1k-224-depth-4tasks (GT depth), 3 views.
+This doc is the exact **input/output contract** for your inference loop — what tensors to feed the model and how to decode what it returns. (It does not cover the sim wrapper; you have your own inference path.)
 
-## Contents
-- `params/` — Orbax bf16 params (load with PiBehaviorWeightLoader).
-- `assets/IliaLarchenko/behavior_224_rgb/` — qvel-fixed norm_stats.json + fast_tokenizer.
-- (optimizer train_state NOT included — inference/eval release.)
+> **`task_index` for this checkpoint** — set the per-step task id to the trained task:
+> `task54 = putting_away_toys` → **54** · `task56 = make_rose_centerpieces` → **56** · (use the id matching this repo's task).
+
+---
+
+## ⚠️ Read first: the DA3 branch is opt-in and OFF unless you feed it
+
+The model's DA3 branch (`spatial_bank_builder`) **contributes nothing unless the `da3_*` observation fields are populated.** Internally:
+
+```python
+# pi_behavior.py _compute_banks
+if self.spatial_bank_builder is None or observation.da3_features is None:
+    return None          # DA3 conditioning is skipped entirely
+```
+
+So if you feed only RGB+state+task (leaving `da3_features=None`), this DA3-trained checkpoint runs as an **RGB-only model** — a train/eval mismatch that will degrade it, since the base transformer was trained *with* DA3 conditioning active. **To evaluate DA3 faithfully you must populate the `da3_*` tensors** (§ Input B). If you want the RGB baseline instead, use the `noda3` checkpoint, not this one.
+
+---
+
+## 0. Model load
+
+```python
+import dataclasses, jax.numpy as jnp
+from b1k.training import config as _config
+import openpi.models.model as _model
+
+cfg = _config.get_config("pi_behavior_b1k_fast")
+# num_tasks and da3 are NOT stored in the checkpoint — you MUST reconstruct them
+# to match training or model.load() fails the strict shape check.
+model_cfg = dataclasses.replace(cfg.model, num_tasks=100, da3=<B1KDA3Config matching training>)
+cfg = dataclasses.replace(cfg, model=model_cfg)
+
+params = _model.restore_params(f"{ckpt}/params", dtype=jnp.bfloat16)
+model  = model_cfg.load(params)
+model.load_correlation_matrix(norm_stats)   # REQUIRED (use_correlated_noise=True)
+model.eval()
+```
+
+- **`num_tasks=100` is mandatory** — the default config is 50; loading the 100-row task-embedding table into a 50-task model fails the shape check.
+- **DA3 config must match training** so the `spatial_bank_builder` params exist to receive the weights. Flags for this checkpoint family:
+  `da3_channels=1536`, `out_layers=(19,27,33,39)`, `num_inject_layers=6`, `grid_hw=(16,16)`, `backbone_only=True`, `cam_frame_ray=True`, `spatial_inject_vlm=False`, `cross_view=False`, `use_perceiver=False`, `lang_fusion_depth=0`, `spatial_vec=True`, `spatial_film=True`, `spatial_use_da3=True`, `fourier_bands=10`, `qk_norm=True`, `inject_dropout=0.1`.
+- **Assets** (bundled under `assets/IliaLarchenko/behavior_224_rgb/`): `norm_stats.json` (contains the correlation matrix — required by `load_correlation_matrix`) and the FAST `fast_tokenizer/` (training-aux; **not used at inference**). `asset_id = "IliaLarchenko/behavior_224_rgb"`.
+
+---
+
+## Input A — base observation (always required)
+
+The model's `sample_actions(rng, Observation, num_steps=20)` consumes an `Observation` with these fields (batch dim `b`):
+
+| field | shape | dtype | notes |
+|---|---|---|---|
+| `images["base_0_rgb"]` | `[b,224,224,3]` | float32 | head (zed) camera |
+| `images["left_wrist_0_rgb"]` | `[b,224,224,3]` | float32 | left realsense |
+| `images["right_wrist_0_rgb"]` | `[b,224,224,3]` | float32 | right realsense |
+| `image_masks[*]` | `[b]` | bool | all `True` |
+| `state` | `[b,32]` | float32 | normalized proprio (see below) |
+| `tokenized_prompt` | `[b,2]` | int32 | **`[task_id (0–99), current_stage]`** — not text |
+| `tokenized_prompt_mask` | `[b,2]` | bool | `[True,True]` |
+| `subtask_state` | `[b]` | int32 | `= current_stage` (**required**, else the model raises) |
+| `fast_tokens` | — | — | **must be `None`** (inference raises if present) |
+
+**Images:** HWC, RGB order, values in **`[-1,1]`** (`Observation.from_dict` does `uint8/255*2 − 1` — if you build the `Observation` yourself, apply this). Resize/pad to 224² first.
+
+**State (32-d):** build the 23 real dims in this order, then zero-pad to 32:
+`[base_qvel(3), trunk_qpos(4), left_arm(7), left_gripper(1), right_arm(7), right_gripper(1)]`.
+Grippers: rescale raw `[0,0.1] → [-1,1]` via `2*(raw/0.1)-1`. Then **z-score normalize** with `norm_stats["state"]` (regular mean/std, `use_quantile_norm=False`): `(x-mean)/(std+1e-6)`. Dims 23:32 stay 0.
+
+**Task / stage:** `task_id` indexes the learned task table (must be a trained id, 0–99). `current_stage` is the subtask index; at t=0 use 0, then update it from the model's `subtask_logits` output (below). The model **always uses the stage you pass in** — it does not self-advance inside the forward.
+
+---
+
+## Input B — DA3 tensors (required to activate the depth branch)
+
+Populate these (else the DA3 branch is skipped — see the warning at top). These mirror what the **training** data pipeline produced with a frozen **`depth-anything/DA3-GIANT-1.1`** backbone run on the 3 RGB views; you must reproduce that at eval:
+
+| field | shape | dtype | notes |
+|---|---|---|---|
+| `da3_features` | `[b,4,3,1536,16,16]` | uint16 (bf16-bits) | DA3-GIANT backbone feats: 4 out-layers (19,27,33,39) × 3 views × 1536ch × 16×16 grid |
+| `da3_ray` | `[b,3,3,16,16]` | float32 | per-view intrinsic ray directions, **camera frame** (`cam_frame_ray=True`) |
+| `da3_depth` | `[b,3,1,16,16]` | float32 | **GT metric depth in METERS** per view — DA3-GIANT has no metric head, so this comes from the sim's depth render (`B1K_USE_GT_DEPTH=1`), downsampled to the 16×16 grid |
+| `camera_extrinsics` | `[b,3,4,4]` | float32 | per-view camera→world 4×4 |
+| `lang_feat` | `[b,32,1024]` | float32 | ModernBERT-large embedding of the task name |
+| `lang_mask` | `[b,32]` | bool | token mask for `lang_feat` |
+
+How to produce them (same as training): run DA3-GIANT-1.1 (via the `Depth-Anything-3` repo / `da3_for_geostack.py` wrapper) on the 3 RGB views to get backbone `feats` + intrinsic `rays`; supply the sim's GT depth as `da3_depth` (meters); provide camera extrinsics; and embed the task-name string with ModernBERT-large for `lang_feat`/`lang_mask` (a per-task cache is fine). `da3_features` is stored as the raw bf16 bit-pattern in uint16 — cast the bf16 feature tensor to uint16 bits, matching the training extractor.
+
+> If reproducing full DA3 extraction is impractical for your harness, note that the RGB-only fallback (leave `da3_*` = `None`) does **not** give a valid DA3 eval — prefer the `noda3` checkpoint for the RGB number and this one only when you can feed DA3.
+
+---
+
+## Output — decoding
+
+`sample_actions` returns a **tuple `(actions_raw, subtask_logits)`**.
+
+**Actions**
+- `actions_raw`: shape **`[b, 30, 32]`**, in **normalized** space (flow-matching sample; Euler-integrated, default **`num_steps=20`** denoise steps, correlated noise `beta=0.5`).
+- **De-normalize** with the per-timestamp action stats (`norm_stats["actions"]`, mean/std shape `[30,32]`):
+  `actions = actions_raw * (per_timestamp_std + 1e-6) + per_timestamp_mean`
+- **Keep the first 23 dims** (`actions[..., :23]`); dims 23:32 are zero-pad. Order matches the state:
+  `[base_vel(3), trunk(4), left_arm(7), left_gripper(1), right_arm(7), right_gripper(1)]`.
+  (Gripper action dims are in the normalized `[-1,1]` space — map back to your actuator range as your controller expects.)
+- Execute as an action chunk (receding horizon); `initial_actions` may be passed on the next call to seed correlated inpainting (enforced while flow-time `t > 0.3`).
+
+**Stage logits (for stage conditioning)**
+- `subtask_logits`: shape **`[b, 15]`** (`MAX_NUM_STAGES=15`), invalid stages for the task set to `-inf`.
+- `predicted_stage = int(argmax(subtask_logits))`. Feed the (optionally vote-smoothed) stage back as the next step's `subtask_state` / `tokenized_prompt[:,1]`, clamped to `TASK_NUM_STAGES[task_id]-1`.
+
+**FAST tokens:** never produced or consumed at inference — ignore.
+
+---
+
+## Quick reference
+- **Feed:** 3 RGB views `[-1,1]` + normalized 32-d state + `[task_id, stage]` + `subtask_state` **+ the 6 `da3_*` tensors** (to use DA3).
+- **Get back:** `[b,30,32]` normalized actions → de-norm per-timestamp → take `[:, :23]`; plus `[b,15]` stage logits.
+- **Denoise steps:** 20. **Correlated noise:** yes (matrix from `norm_stats`). **num_tasks:** 100. **FAST:** off.