b601-v2-40k

b601-v2-40k — π₀.₅ Seeed B601 "pick red cube", 40k total steps

Inference-only JAX/Orbax checkpoints (params/ + assets/). train_state/ has been stripped, so these cannot be used with scripts/train.py --resume; they are for serving and for warm-starting.

Contents

Directory names are the raw Orbax step numbers of the continuation run, which restarts its counter at 0. They are not cumulative step counts; the mapping above is the cumulative figure. "40k" in the repo name refers to 29,999 + 9,999 = 39,998 ≈ 40k total optimizer steps.

Provenance

Warm-started from LGG100/b601-v2 step 29999 and trained a further 10,000 steps.

This was a warm-start, not a resume: LGG100/b601-v2 ships only params/, so the Adam moments and the step counter from the first 30k run were unavailable. The optimizer state was therefore rebuilt from scratch and the LR schedule restarted from step 0. Same applies to anyone continuing from this repo.

Because the schedule restarted, it was retuned so as not to jolt an already-converged model — the first run used openpi's default CosineDecaySchedule (peak 2.5e-5, decay_steps 30,000), ending at ~2.5e-6:

lr_schedule = CosineDecaySchedule(
    warmup_steps=500, peak_lr=5e-6, decay_steps=10_000, decay_lr=5e-7
)

Training loss over the continuation: 6e-4 (step 0) → 3e-4 (step 9.9k), i.e. the model was already converged and these 10k steps are refinement, not new learning. Relative to step 29999, 52 of 61 parameter leaves moved; the 9 unchanged are the frozen Gemma-2B base. Max abs weight delta at 9999: ~7.3e-4 (LLM LoRA), ~1.6e-3 (action expert).

Config

openpi TrainConfig name pi05_b601_pick_red_cube_llm_lora_cont10k, matching the parameter tree of pi05_b601_pick_red_cube_llm_lora:

• Pi0Config(pi05=True, action_horizon=30, paligemma_variant="gemma_2b_lora")
• LoRA (rank 16) on the PaliGemma LLM only — its 2.5B base weights are frozen.
• Action expert (gemma_300m) and the SigLIP vision tower are fully fine-tuned (openpi's get_freeze_filter only matches .*llm.*, so vision is dense-trained).
• Trainable ≈872.8M / 3.381B (25.8%). ema_decay=None, batch_size=32.

Data

50 episodes / 37,932 frames, LeRobot v2.1, 30 fps, two cameras (top, wrist, 480×640, AV1). The third (right-wrist) view is zero-filled and masked out by B601Inputs.

State and action are 7-dim absolute joint positions in degrees (shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_yaw, wrist_roll, gripper). Actions are absolute, so no delta-action transform is applied.

Single task prompt, required verbatim (prompt_from_task=True):

Pick up the red cube and place it completely inside the tray.

⚠️ observation.state and action do not share a coordinate convention: shoulder_pan, shoulder_lift, wrist_roll and gripper are sign-flipped, and the gripper uses a different unit (state spans [-270, 0], action spans [0, 57]). The follower→leader conversion is not part of this model; a real-robot client must do it. State normalization is min-max, so an out-of-range state saturates π₀.₅'s 256-bin state discretization.

Norm stats

<step>/assets/b601/pick_red_cube_v1/norm_stats.json are the stats these weights were trained with — use them as-is. Recomputing them from the dataset yields slightly different values and would feed the model an input scaling it was not fit to.

Serving

uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_b601_pick_red_cube_llm_lora \
    --policy.dir=<local>/9999 \
    --policy.default-prompt="Pick up the red cube and place it completely inside the tray."