behavior1k-task16-moving-boxes

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

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

Final training metrics (step 14999)

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

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.

(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):

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.