README.md

---
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.