Modelpublic

Hoshipu/behavior1k-task16-moving-boxes

sourceHugging Faceapache-2.0updated 22d agoView on Hugging Face

0likes

Model Card

BEHAVIOR-1K · pi0.5 · Task 16 (movingboxesto_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 checkpoint2's own assets. Use the ones shipped here — mixing in a different `normstats.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.