aailabkaist/pi05_recovery_2task_expert_5k

π0.5 · SO-101 pick & place (2-task) — expert-only, 5,000 steps

A π0.5 (pi05) vision-language-action policy fine-tuned on an SO-101 arm for a two-task cup manipulation setup: pick a cup up from the table and place it onto a blue circle.

Only the action expert was trained. The PaliGemma VLM (2.923 B params) is frozen (train_expert_only=true) — see configuration_pi05.py: "Freeze entire VLM, train only action expert and projections".

• Base model: lerobot/pi05_base — the last pre-relative-actions revision, matching LeRobot 0.4.x
• Dataset: aailabkaist/so101_recovery_2task (749 episodes · 151,255 frames · 30 fps · 8 operators)
• Checkpoint: step 5,000 of a 10,000-step run (batch 64 → 320,000 samples ≈ 2.1 epochs)
• Framework: LeRobot
• Robot: SO-101, 6-DoF, two RGB cameras (front, wrist, 480×640)
• Siblings: expert 10K · full 5K · full 10K — see the two-task collection

Task and data

Two instructions, both used at training time verbatim:

pick up the cup near the blue circle
place the cup on the blue circle

Pick-episode cup start positions (394 episodes, per-operator colors, anonymized). The cups form a ring around the blue circle rather than covering the workspace uniformly.

Place-episode transport vectors, start → placed (350 of 355 place episodes — 5 excluded where the cup is fully occluded by the gripper in the start frame). Endpoints converge tightly on the blue circle; the release target is nearly constant across the dataset.

Training

Trainable parameters	0.693 B (19.2 %)
Frozen	PaliGemma VLM, 2.923 B
Steps	5,000 (run stopped early at 10K; config said 20K)
Batch size	64
Precision	bfloat16, gradient_checkpointing=true
Optimizer	AdamW · lr 2.5e-5 · wd 0.01 · betas (0.9, 0.95) · grad-clip 1.0
Schedule	cosine decay with warmup — auto-scaled by LeRobot (20K steps < 30K decay): warmup 1,000→666, decay 30,000→20,000; peak 2.5e-5, LR at this checkpoint ≈2.2e-5 (configured floor 2.5e-6 never reached — run stopped at 10K)
Chunk	chunk_size=50, n_action_steps=50, n_obs_steps=1
Seed	1000
Hardware	1× NVIDIA RTX PRO 6000 Blackwell (96 GB)
VRAM / speed	28.2 GB · 3.93 s/step
Loss at this checkpoint	0.102

Interactive training curves for this run — loss · LR · grad-norm · epochs (this variant's W&B run · project overview):

Both runs used identical data, batch size and dtype — the only difference is what was trainable, so the two curves are directly comparable. expert-only is essentially flat after warmup (~step 1K): 0.129 → 0.087 over the remaining ~3.8 epochs; full fine-tuning keeps descending. The run was stopped at 10K rather than the configured 20K on the strength of these curves.

Usage

from lerobot.policies.pi05.modeling_pi05 import PI05Policy

policy = PI05Policy.from_pretrained("aailabkaist/pi05_recovery_2task_expert_5k")

Or serve it for async inference:

python -m lerobot.async_inference.policy_server --host=0.0.0.0 --port=8080 --fps=30
python -m lerobot.async_inference.robot_client \
    --server_address=<host>:8080 \
    --policy_type=pi05 --pretrained_name_or_path=aailabkaist/pi05_recovery_2task_expert_5k \
    --task="pick up the cup near the blue circle" \
    --actions_per_chunk=50 --chunk_size_threshold=0.2 \
    --aggregate_fn_name=weighted_average --fps=30

---

A single-task variant of this setup — one instruction covering the whole pick-and-place motion, with a plain white cup — lives in the SO-101 · cup → blue circle (single-task) collection.