--- license: apache-2.0 library_name: openpi tags: - robotics - behavior-1k - pi05 - vla - imitation-learning pipeline_tag: robotics --- # BEHAVIOR-1K · pi0.5 · Task 16 (moving_boxes_to_storage) A single-task fine-tune of the pi0.5 (π₀.₅) 3B Vision-Language-Action model (the **PiBehavior** variant) on **task 16 — `moving_boxes_to_storage`** from the BEHAVIOR-1K 2026 Challenge demonstration set. - **Architecture:** pi0.5 / PiBehavior (≈3B params, JAX/Flax, [openpi](https://github.com/Physical-Intelligence/openpi)) - **Conditioning:** **task-only embeddings** — one learned embedding per task (no per-skill/sub-stage table) - **Base checkpoint:** [`IliaLarchenko/behavior_submission`](https://huggingface.co/IliaLarchenko/behavior_submission) → `checkpoint_2` - **Checkpoint format:** Orbax (OCDBT) — `params/` holds the model weights - **Camera input:** RGB only (224×224, GOP8); depth channels removed ## Files in this repo | Path | What it is | |---|---| | `params/` | Trained model weights, step 14999 (Orbax OCDBT) | | `assets/behavior-1k/2026-challenge-demos/norm_stats.json` | Action/state normalization stats (**from the 100-task meta checkpoint**, *not* from checkpoint_2) | | `assets/behavior-1k/2026-challenge-demos/fast_tokenizer/` | FAST action tokenizer (same 100-task assets) | > **Important:** the normalization stats and FAST tokenizer are the **100-task ("meta100")** assets, deliberately kept > instead of the base checkpoint_2's own assets. Use the ones shipped here — mixing in a different `norm_stats.json` > will produce wrong action scaling. ## How it was trained | | | |---|---| | Task | `moving_boxes_to_storage` (task_index 16), 200 demo episodes, 2,919,245 frames | | Dataset | `behavior-1k/2026-challenge-demos` (LeRobot v3.0), RGB-only view | | Base weights | `IliaLarchenko/behavior_submission/checkpoint_2` (2025 arch: `task_embeddings (50,2048)`, `task_stage_embeddings`, `stage_pred_from_vlm`) | | Weight adaptation | `task_embeddings` expanded **50 → 100** (rows 0–49 copied, 50–99 random-init); `task_stage_embeddings` and `stage_pred_from_vlm` dropped; new `task_only_embeddings` and `skill_pred_from_vlm` random-init | | Hardware | 8 × NVIDIA H200 | | Sharding | FSDP, `fsdp_devices=8`, pure model-parallel (no data-parallel replica) | | Global batch | 512 (64 / GPU) | | Steps | 15,000 (≈2.6 epochs), ~19.5 h wall-clock, ~4.2 s/it | | Optimizer | AdamW, cosine schedule: warmup 1000 → peak LR 1e-4 → 1e-5 over 15,000 steps | | Framework | openpi (JAX/Flax nnx), `b1k` training stack | ### Final training metrics (step 14999) | metric | value | |---|---| | total loss | ≈ 0.018 | | FAST token accuracy | ≈ 0.95 | | subtask accuracy | ≈ 0.9998 | | grad norm | ≈ 0.045 | Loss trajectory: 0.71 (step 0) → 0.035 (step 7k) → 0.018 (step 15k). ## How to use This checkpoint plugs into the `openpi` / `b1k` (behavior-1k-solution) training & inference stack. Weights are stored as plain Orbax params; load them with openpi's restore utility (which strips sharding so you can re-shard freely): ```python import numpy as np import openpi.models.model as _model # Point at the params/ directory of this repo (downloaded locally) params = _model.restore_params("path/to/params", restore_type=np.ndarray) # params["task_only_embeddings"], the pi0.5 backbone, action expert, etc. ``` For rollout/eval, build the **PiBehavior task-only** model config, load these `params`, and supply the `norm_stats.json` + `fast_tokenizer` from `assets/` in this repo so inputs are normalized and actions are tokenized exactly as in training. The task prompt for conditioning is task 16's language instruction (`moving_boxes_to_storage`) from the challenge `meta/tasks.jsonl`. ## Notes & limitations - Single-task specialist: trained only on task 16. The task-only embedding for other task indices is either random-init (indices 50–99) or inherited from the base (0–49) and was **not** tuned here. - RGB-only: expects the 224×224 RGB observation layout; depth inputs are not used. - Trained on the 2026-challenge demo split; no held-out evaluation numbers are included in this card.