JackLiu0406/b1k-checkpoints

BEHAVIOR-1K 5-task meta-training (PiBehavior / pi0.5)

Continued meta-training of the public #1 BEHAVIOR-1K solution (IliaLarchenko/behavior_50t_checkpoint) on five tasks of the 2026 challenge set, two of which are outside the original 50-task space.

No DA3. This is the stock PiBehavior model (3.4 B params, gemma_2b VLM + gemma_300m action expert). No spatial branch, no depth.

Contents

folder	what
meta5_2026_224/step_19999	final model of the 20k-step run — use this
meta5_2026_224_ext/step_*	rolling checkpoint of an extension run past 20k (LR pinned at 1e-6; may not differ meaningfully)

Each folder holds params/, train_state/ (Adam moments, for resuming), assets/, _CHECKPOINT_METADATA.

---

READ THIS FIRST: task indices are the identity map

The model uses a 100-task space, and the 2026 dataset task index is the model task index. There is no remapping.

dataset task  0 (turning_on_radio)        -> model index  0
dataset task  5 (setting_mousetraps)      -> model index  5
dataset task 40 (make_microwave_popcorn)  -> model index 40
dataset task 76 (dispose_of_glass)        -> model index 76
dataset task 77 (installing_a_modem)      -> model index 77

Do not compress the new tasks into indices 50/51 or otherwise renumber them. The conditioning tables were widened to 100 rows, and tasks 76/77 own rows 76/77.

Verify with the repo's own helper rather than trusting any hardcoded table:

from b1k.training.b1k_2026 import build_task_index_maps
name2model, idx2026_to_model, name2idx = build_task_index_maps(ROOT_2026, TASK_DATA_JSON)
assert idx2026_to_model[77] == 77

Which tasks were actually trained

task	name	in upstream 50?	stages	episodes / frames
0	turning_on_radio	yes — embedding inherited	—	200 / 429,928
5	setting_mousetraps	yes — embedding inherited	—	200 / 2,039,222
40	make_microwave_popcorn	yes — embedding inherited	—	200 / 647,555
76	dispose_of_glass	no — new row, random init	11	200 / 1,839,066
77	installing_a_modem	no — new row, random init	5	200 / 482,597

Only 76 and 77 started from random embeddings; 0/5/40 kept their pretrained ones. Rows 50-99 other than 76 and 77 were never trained — they are still at random init. Evaluating any task index in 50..99 except 76/77 will produce garbage.

Conditioning tables

table	shape	notes
task_embeddings.embedding	(100, 2048)	rows 0-49 from upstream, 76/77 trained, other 50+ rows random
task_stage_embeddings.embedding	(1120, 1024)	rows 0-595 from upstream

TASK_NUM_STAGES has 100 entries summing to 1120. MAX_NUM_STAGES = 15 is unchanged (stage_pred_from_vlm emits 15 logits masked per task). TASK_STAGE_OFFSETS is a cumsum over TASK_NUM_STAGES — never insert a task below index 50, or every inherited stage row silently re-points.

Stage counts follow clip(ceil(mean_episode_frames / 900), 5, 15) (reproduces all 50 upstream entries exactly).

Set the model config to the wide table:

model = dataclasses.replace(config.model, num_tasks=100)

There is also a validation guard in b1k/transforms.py that rejects task_index >= 50; it must compare against len(TASK_NUM_STAGES) instead, or every sample from tasks 76/77 is silently dropped.

Loading the checkpoint (important)

params/ were saved from a 4-device mesh. Orbax embeds a NamedSharding per array bound to that mesh, so loading on a host with a different GPU count fails with:

ValueError: sharding passed to deserialization should be specified, concrete
and an instance of `jax.sharding.Sharding`. Got None

This happens regardless of --fsdp_devices. Restore to host numpy instead and let JAX shard afterwards:

with ocp.PyTreeCheckpointer() as ckptr:
    meta = ckptr.metadata(params_path)
    # orbax >= 0.12 wraps this in StepMetadata, which is not subscriptable
    meta = getattr(meta, "item_metadata", None) or meta
    restore_args = jax.tree.map(lambda _: ocp.RestoreArgs(restore_type=np.ndarray), meta)
    restored = ckptr.restore(params_path, restore_args=restore_args)

PiBehaviorExpandTasksWeightLoader already does this.

Assets — use as shipped, do not regenerate

assets/ contains norm_stats.json and fast_tokenizer/, both byte-identical to upstream behavior_50t_checkpoint.

• Do not retrain the FAST tokenizer. It transfers to the new tasks: 0.000% alphabet overflow, no clipping, round-trip MAE on new tasks matching in-distribution controls. It operates on DCT-quantised delta actions, so if you re-verify, feed it deltas.
• norm_stats are the upstream 2025 stats. The pretrained weights were optimised against exactly this normalisation, so keep them. Caveat: the 2026 data carries the fixed robot-frame base_qvel while these stats describe the 2025 world-frame convention, so state dims 0:3 are normalised against a stale distribution. All other state dims and all 32 action dims are correct. Per-channel action_loss_base_vel_* stayed in line with the arm channels during training, so this did not appear to hurt — but it is the first thing to look at if base-velocity behaviour is off at eval.

Data format

Trained on the 2026 challenge demos (LeRobot v3.0) with RGB video pre-resized to 224x224. The model resizes to 224 internally via resize_with_pad, and all cameras are square (480x480 wrist, 720x720 head), so pre-resizing is numerically equivalent and ~10x cheaper to decode.

Camera key mapping:

2026 dataset	model
observation.rgb.zed_link_camera_0	observation.images.rgb.head
observation.rgb.left_realsense_link_camera_0	observation.images.rgb.left_wrist
observation.rgb.right_realsense_link_camera_0	observation.images.rgb.right_wrist

Note lerobot 0.3.4 (what openpi pins) cannot read v3.0 — it raises ForwardCompatibilityError. Use the repo's b1k_2026.py v3 loader.

Training setup

init	IliaLarchenko/behavior_50t_checkpoint
hardware	4x GB300, batch 1024 (256/GPU), fsdp_devices=4
steps	20,000 (final checkpoint at 19,999)
LR	1e-6 →1e-4 @1k, cosine →1e-5 @10k, hold to 15k, cosine →1e-6 @20k
loss	0.7071 → 0.0341
grad_norm	2.76 → 0.053
subtask_accuracy	0.64 → 0.9994
fast_accuracy	~0.88

Sanity checks at eval

1. Confirm task_embeddings is (100, 2048) and task_stage_embeddings is (1120, 1024).
2. subtask_accuracy should be ≈0.99 on the trained tasks. If it sits near chance (~1/num_stages) your TASK_STAGE_OFFSETS are wrong — almost always caused by renumbering tasks below index 50.
3. Only evaluate task indices 0, 5, 40, 76, 77. Other indices in 50..99 are untrained.