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

## Exact model definition (needed to load / eval)

The exact training & model code that produced this checkpoint is pushed here — check it out to get the
precise `PiBehaviorConfig` and the task-only model class (don't try to reconstruct from a generic B1a / pre-B1a branch,
the tables differ):

> **Repo:** https://github.com/markli1hoshipu/behavior-1k-solution
> **Branch:** `task-only-embeddings`  ·  **commit:** `5147ffb`
> Model class: `src/b1k/models/pi_behavior.py` · config: `src/b1k/training/config.py`
> (config name **`pi_behavior_b1k_task16_taskonly`**) · loader: `src/b1k/training/weight_loaders.py`

### `PiBehaviorConfig` values used

| field | value |
|---|---|
| `num_tasks` | 100 |
| `num_skills` | 34 |
| `task_embedding_dim` | 2048 |
| `action_dim` | 32 |
| `action_horizon` | 30 |
| `max_token_len` | 200 |

### Actual checkpoint param shapes (verify against these)

The **task-only** conditioning table is `task_only_embeddings (100, 1024)` — that is the one added for this recipe.
The base `task_embeddings (100, 2048)` is still present (inherited/expanded from the base checkpoint). There is
**no** `task_stage_embeddings`, **no** `stage_pred_from_vlm`, and **no** `skill_embeddings` table; only the
`skill_pred_from_vlm` head remains.

| param | shape |
|---|---|
| `task_only_embeddings/embedding` | **(100, 1024)** ← task-only conditioning |
| `task_embeddings/embedding` | (100, 2048) |
| `skill_pred_from_vlm/{kernel,bias}` | (2048, 34) / (34,) |
| `gate_task/{kernel,bias}` | (4096, 2048) / (2048,) |
| `gate_task_stage/{kernel,bias}` | (4096, 1024) / (1024,) |
| `gate_sincos/{kernel,bias}` | (4096, 1024) / (1024,) |
| `fusion_layer1/{kernel,bias}` | (4096, 4096) / (4096,) |
| `fusion_layer2/{kernel,bias}` | (4096, 2048) / (2048,) |
| `fast_token_embedding/embedding` | (1024, 2048) |
| `fast_token_proj/{kernel,bias}` | (2048, 1024) / (1024,) |
| `kv_transform/{k,v}_coeffs` · `{k,v}_bias` | (18, 18) · (18, 1, 256) |

(`gate_task_stage` is kept as a layer but has no `task_stage_embeddings` feeding it in this task-only variant.)

## How to load

Weights are plain Orbax params; load with openpi's restore utility (it strips sharding so you can re-shard freely):

```python
import numpy as np
import openpi.models.model as _model

params = _model.restore_params("path/to/params", restore_type=np.ndarray)  # the params/ dir of this repo
# -> params["task_only_embeddings"], ["task_embeddings"], the PaliGemma backbone, action expert, etc.
```

For rollout/eval: build the model from config `pi_behavior_b1k_task16_taskonly` on the branch above, 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 conditioning prompt is task 16's instruction
(`moving_boxes_to_storage`) from the challenge `meta/tasks.jsonl`; the task-only embedding is indexed by task_index 16.

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