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

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

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

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:

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

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.