Hoshipu/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

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

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.