README.md

metadata

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)
• Conditioning: task-only embeddings — one learned embedding per task (no per-skill/sub-stage table)
• Base checkpoint: 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):

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.