JackLiu0406/meta-SFT-checkpoints

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

Continued meta-training of IliaLarchenko/behavior_50t_checkpoint on all 100 activities of the BEHAVIOR-1K 2026 challenge set.

No DA3. Stock PiBehavior (~3.83 B params, gemma_2b VLM + gemma_300m action expert). No spatial branch, no depth.

This repo is private. The public JackLiu0406/b1k-checkpoints holds the single-task finetunes and the 5-task meta5_2026_224* runs; the 100-task meta checkpoints live here.

Contents

folder	step	what
meta100-1epoch/step20000	20,000	mid-run, complete (params + train_state + assets)
meta100-1epoch/step35000	35,000	weights only — no train_state, cannot resume
meta100-1epoch/step45000	45,000	complete
meta100-1epoch/step60000	60,000	complete
meta100-1epoch/step69999	70,000	final model of epoch 1 — use this

The final checkpoint is named step69999, not step70000

train.py saves the last checkpoint at num_train_steps - 1, and save_state runs after the step counter increments. So the directory named 69999 contains train_state.step == 70000. It is the complete 70,000-step model — only the folder name is off by one.

Task space: all 100 activities are trained

Unlike the 5-task meta5_2026_224 run, every task index 0–99 was trained here. There is no "untrained rows" caveat. The 2026 dataset task index is the model task index — no remapping:

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

Conditioning tables: task_embeddings.embedding is (100, 2048), task_stage_embeddings.embedding is (1120, 1024). TASK_NUM_STAGES has 100 entries summing to 1120; MAX_NUM_STAGES = 15. Set the model config to the wide table:

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

Note b1k/transforms.py has a validation guard rejecting task_index >= 50; it must compare against len(TASK_NUM_STAGES) or every sample from tasks 50+ is silently dropped.

Assets — use as shipped

Each checkpoint's assets/IliaLarchenko/behavior_224_rgb/ contains norm_stats.json and fast_tokenizer/. Use them as they are.

These are the qvel-fixed stats, not the stale 2025 ones. Verified by full deep comparison: the embedded norm_stats.json is numerically identical to norm-stats-fixed/norm_stats.json in this repo — 0 differences across every field, including the 960×960 action_correlation_cholesky. The larger byte size is JSON float formatting only.

state.mean[0:3] = [ 0.044826,  0.001065, -0.000582]
state.std [0:3] = [ 0.108685,  0.060844,  0.133118]   <- robot-frame base_qvel

If you see state.std[0:3] ≈ [0.0094, 0.0087, 0.0137] you have loaded the 2025 world-frame stats and base velocity will be mis-scaled by ~10×. See norm-stats-fixed/README.md.

Do not retrain the FAST tokenizer — it transfers to the new tasks with 0.000% alphabet overflow. It operates on DCT-quantised delta actions.

Loading the checkpoint

params/ were saved from an 8-device mesh (fsdp_devices=8). 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

Restore to host numpy and let JAX shard afterwards:

with ocp.PyTreeCheckpointer() as ckptr:
    meta = ckptr.metadata(params_path)
    meta = getattr(meta, "item_metadata", None) or meta   # orbax >= 0.12 wraps in StepMetadata
    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.

Data

2026 challenge demos, LeRobot v3.0, RGB pre-resized to 224×224. Verified encoding: HEVC / libx265, 224×224, yuv420p.

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

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 (50-task → 100-row table, new rows random init)
hardware	8× B300, batch 3072 (384/GPU), fsdp_devices=8
steps	70,000 = 70,000 × 3072 = 215,040,000 samples over 210,316,774 = 1.022 epochs
LR	1e-7 →1e-4 @2k, cosine →1e-5 @10k, flat to 35k, cosine →5e-6 @70k
assets	assets_qvelfix (robot-frame base_qvel)
throughput	~9.4 s/it

Metrics

	step 0	step 60,000	step 66,880
action_loss	0.4620	0.0758	0.0713
total_loss	0.7738	0.1023	0.0978
fast_accuracy	0.7925	0.8378	0.8386
subtask_accuracy	0.6045	0.9964	0.9966

action_loss was still descending monotonically at the end of epoch 1 (power-law fit 2.173·s^-0.305, no detectable floor), but fast_accuracy has been flat at ~0.838 since step 25,000 — the discrete action-token predictions converged long before the regression loss.

Sanity checks at eval

1. task_embeddings is (100, 2048), task_stage_embeddings is (1120, 1024).
2. state.std[0:3] ≈ [0.1087, 0.0608, 0.1331] — if it is ~10× smaller you loaded the 2025 stats.
3. subtask_accuracy ≈ 0.99. Near chance (~1/num_stages) means TASK_STAGE_OFFSETS is wrong.
4. All 100 task indices are valid here.

Known caveat: train/eval pixel pipeline differ

Training pixels came from swscale-resized, x265-encoded, YUV420-subsampled 224×224 video. Eval feeds PIL BILINEAR through resize_with_pad on RGB. The two pipelines do not produce identical pixels, and the size of the resulting distribution shift has not been measured. If eval underperforms the loss curve, this is the first thing to quantify.

Epoch 2

Training continues from step69999 to step 140,000 (a second pass over the same data), with the LR resuming at exactly 5e-6 and cosine-decaying to 1e-6. The junction is C1-continuous: epoch 1's cosine ends with zero slope and epoch 2's begins with zero slope.