You can deploy this model and start earning money today!

Kavin60606/
pi0-fast-piper-h-plate-pick

Name: pi0-fast-piper-h-plate-pick
Availability: InStock

openpi

robotics

vision-language-action

Views

No views yet

π0-FAST LoRA Fine-Tune — AgileX Piper-H Plate-Pick

End-to-end LoRA fine-tune of π0-FAST (Physical Intelligence's autoregressive VLA) on a custom AgileX Piper-H plate-pick task, trained with openpi (JAX) on a single A100 80GB.

This repo contains every checkpoint produced during the 3-stage training pipeline (warmup → main → refine), plus all configs and scripts needed to reproduce.

TL;DR


Model	`pi0_fast` w/ `paligemma_variant="gemma_2b_lora"`
Trainable	~50M LoRA params (SigLIP + Gemma 2B fully frozen)
Warm-start	`gs://openpi-assets/checkpoints/pi0_fast_libero/params`
Action space	7-dim (right-arm joints 1–6 + right gripper); left arm held static at `park_q` by controller
Cameras	front (`base_0_rgb`) + right wrist (`left_wrist_0_rgb`); ego clips zero-fill the wrist slot
Task language	`"pick up the plate"` (constant across all 414 episodes)
Action horizon	10 steps @ 20 fps
Total train steps	1k warmup + 4k main (stopped early on val plateau) + 2k refine (stopped early on val regression) = 7k effective
Wall time	~6 hr on 1× A100 80GB
Best checkpoint	`main/4000/` — lowest val/loss = 1.5797 (refine never recovered)

Why this repo exists

New embodiment (AgileX Piper-H) — not in pi0-FAST pretraining.
Single-task in-distribution eval (same teleop environment).
Small dataset (≈ 2 hours teleop + 3.6 min egocentric augmentation).
LoRA path chosen to keep the LIBERO tabletop-pick prior intact while learning the new joint mapping with ~50M trainable parameters.

Datasets

All datasets converted to LeRobot v2.1, sliced to 7-dim action, trimmed at first gripper-close + 60-frame buffer (3s @ 20 fps), task string standardized to "pick up the plate".

Dataset	Source	Eps	Frames	Duration
`single_pick_v21`	`PranayTest/02062026_single_plate_pick` (HF)	113	40,633	33.9 min
`multi_pick_v21`	`PranayTest/02062026_multi_plate_pick` (HF)	198	90,827	75.7 min
`piper_h_ego_103`	EgoDex (Apple) + EPIC-KITCHENS-100 + EGTEA + Charades-Ego, retargeted via HaWoR + IK to Piper-H joint space, dominant-arm always mapped to right slot	103	4,304	3.6 min
Total		414	135,764	~113 min

Trim logic (matches single_pick_only README):

python

CLOSED_THRESH = 0.3
BUFFER = 60  # frames
# action[:, 12] = right_gripper (1.0 open, 0.0 closed)
def find_trim_frame(grip):
    for i in range(1, len(grip)):
        if grip[i] < CLOSED_THRESH:
            return min(i + BUFFER, len(grip) - 1)
    return len(grip) - 1

Action / state slice: from the source 14-dim bimanual vector we keep [0,1,2,3,4,5,12] = right joints 1–6 + right gripper. Left arm is parked at park_q by the robot controller at deploy, not predicted by the policy.

Cameras: source datasets have front, right, left. We drop left and keep front + right. Egocentric clips never had wrist cams → right slot is zero-filled (one-hot constant the model learns to ignore).

Training pipeline

Three chained stages, each registered in src/openpi/training/config.py. Stages share weights via weight_loader=CheckpointWeightLoader(...) pointing to the previous stage's saved params.

Stage	Steps run	Mix (single / multi / ego)	save@	eval@	Result
`piper_h_warmup`	1,000 (planned 1,000)	0.65 / 0.35 / 0.00	500	500	train 1.49 / val 1.75 at step 999
`piper_h_main`	4,000 (early-stopped from planned 11,500)	0.55 / 0.35 / 0.10	2,000	2,000	train 1.17 / val 1.5797 ← best
`piper_h_refine`	2,000 (early-stopped from planned 2,500)	0.60 / 0.37 / 0.03	1,000	1,000	train 1.04 / val 1.6126 at step 2000 (val regressed further → clear overfit)

Why main stopped at 4,000: val gap widened from 0.23 (step 2k) → 0.41 (step 4k) while train kept falling — classic plateau signal. Continuing to 11.5k would have risked further val degradation without improvement. Refine was launched from main/4000 to see if the lower-ego mix could pull val down further; it did not.

Decision rule applied

gap (val − train) < 0.3 → healthy, continue
gap 0.3–0.5 → watch; stop if val plateaus or rises
gap > 0.5 → hard stop, use best-by-val ckpt
Val rising for 3 consecutive evals → hard stop

Loss curves (all stages)

Combined view (warmup → main → refine on cumulative-step axis)

Warmup (log-y, init loss = 10.37)

Main (early-stopped at step 4,000 on val plateau)

Refine (early-stopped at step 2,000 on val regression)

Train loss (every 100 steps)

Step	Stage	train/loss
0	warmup	10.37
500	warmup	1.91
999	warmup	1.49
1,100 (= main step 100)	main	1.81 (mix-shift transient)
main 200	main	1.64
main 500	main	1.57
main 1,000	main	1.49
main 2,000	main	1.35
main 3,000	main	1.23
main 4,000	main	1.17 ← best ckpt
refine 100	refine	1.13
refine 500	refine	1.09
refine 1,000	refine	1.10
refine 1,500	refine	≈ 1.08
refine 2,000	refine	1.04 (final, run stopped)

Val loss (held-out 25 eps: 10 single + 15 multi, seed=42)

Step	val/loss	train	gap
warmup 500	2.1259	1.91	0.22
warmup 999*	1.7529	1.49	0.27
main 2,000	1.5854	1.35	0.23
main 4,000	1.5797 ← best	1.17	0.41
refine 1,000	1.5862 (+0.0065 ⚠)	1.10	0.49
refine 2,000	1.6126 (+0.0329 ⚠⚠)	1.04	0.58 ❌

*Step 999 val computed post-hoc via scripts/eval_ckpt.py since openpi's val eval block fires on step % eval_interval == 0 and 999 is not a multiple of 500.

Both train and val fell smoothly through ~step 2,000. Past step 2k, val flattened ~1.58 while train kept falling → diminishing generalization → stop.

Checkpoints in this repo

All checkpoints saved in Orbax PyTree format. Each is ~8 GB (full state including frozen backbone + LoRA delta + optimizer state).

Path in repo	Stage	Step	train/loss	val/loss	Use case
`checkpoints/warmup/500/`	warmup	500	1.91	2.13	early-warmup snapshot, low quality
`checkpoints/warmup/999/`	warmup	999	1.49	1.75	end of warmup; init for main
`checkpoints/main/2000/`	main	2,000	1.35	1.59	mid-main, healthy gap
`checkpoints/main/4000/`	main	4,000	1.17	1.58	★ best by val, recommended for deployment
`checkpoints/refine/1000/`	refine	1,000	1.10	1.586	refine mid; val regressed vs main/4000
`checkpoints/refine/2000/`	refine	2,000	1.04	1.613 ❌	refine final (run stopped early); train↓ val↑ = clear overfit

For deployment use checkpoints/main/4000/ unless evidence emerges that a later refine ckpt is better on the real robot.

How to load

python

from openpi.training import config as _cfg
from openpi.policies import policy_config
from openpi.shared import download

config = _cfg.get_config("piper_h_main")
ckpt_dir = download.maybe_download("Kavin60606/pi0-fast-piper-h-plate-pick/checkpoints/main/4000")
policy = policy_config.create_trained_policy(config, ckpt_dir)

example = {
    "observation/image": <HxWx3 front camera, uint8>,
    "observation/wrist_image": <HxWx3 right wrist camera, uint8>,
    "observation/state": <7-dim joint+gripper, float32>,
    "prompt": "pick up the plate",
}
action_chunk = policy.infer(example)["actions"]   # shape (10, 7)

At deploy, the left arm is held static at park_q by the robot controller — the policy only predicts the 7-dim right-arm chunk.

Reproduce locally

You need the same openpi commit + patched files (3 files modified to support multi-LeRobot mixing, held-out val, etc.). Files included in this repo under scripts/.

Required environment

Python 3.11
JAX with CUDA 12 support
uv sync against openpi @ main (we used the head as of June 2026)
lerobot pinned to 0cf86487... (openpi's pin)
A100 80GB (or any 80GB GPU)

Patched files (all included in this repo)

scripts/config.py            ← openpi/src/openpi/training/config.py
scripts/data_loader.py       ← openpi/src/openpi/training/data_loader.py
scripts/train.py             ← openpi/scripts/train.py (adds val eval block)
scripts/piper_policy.py      ← openpi/src/openpi/policies/piper_policy.py (new)
scripts/teleop_to_v21.py     ← v3→v2.1 converter w/ pyav decode + h264 encode + gripper-close trim + 7-dim slice
scripts/ego_to_lerobot.py    ← ego→LeRobot v2.1 converter, zero-fill right wrist
scripts/eval_ckpt.py         ← post-hoc train+val loss for any checkpoint

Pipeline

bash

# 1. Convert teleop datasets (~30 min CPU each w/ h264 + pyav)
python scripts/teleop_to_v21.py /path/to/single_pick /path/to/single_pick_v21 local/piper_single_pick_v21
python scripts/teleop_to_v21.py /path/to/multi_pick  /path/to/multi_pick_v21  local/piper_multi_pick_v21

# 2. Convert ego clips (~3 min)
python scripts/ego_to_lerobot.py   # paths hard-coded inside

# 3. Symlink under merged_root for MultiLeRobotDataset
mkdir -p /root/datasets/merged_root/local
ln -snf /path/to/single_pick_v21 /root/datasets/merged_root/local/piper_single_pick_v21
ln -snf /path/to/multi_pick_v21  /root/datasets/merged_root/local/piper_multi_pick_v21
ln -snf /path/to/matched_103_lerobot /root/datasets/merged_root/local/piper_h_ego_103

# 4. Norm stats (teleop only, 5000 sub-sampled frames, ~1.5 min)
uv run scripts/compute_norm_stats.py --config-name=piper_h_warmup --max-frames=5000

# 5. Train (3 stages, chained via weight_loader)
uv run scripts/train.py piper_h_warmup --exp-name=run1 --overwrite
uv run scripts/train.py piper_h_main   --exp-name=run1 --overwrite   # stopped early at step 4000 in our run
uv run scripts/train.py piper_h_refine --exp-name=run1 --overwrite   # weight_loader points to main/4000

Hyperparameters

Param	Value
`paligemma_variant`	`"gemma_2b_lora"`
`action_dim`	7
`action_horizon`	10
`max_token_len`	180
`dtype`	bfloat16
`batch_size`	32
`learning_rate` (peak)	2.5e-5
`lr_schedule`	cosine
`warmup_steps`	1,000
`optimizer`	AdamW (clip_gradient_norm=1.0)
`ema_decay`	None
`gradient_checkpointing`	true
`fsdp_devices`	1
`freeze_filter`	`Pi0FASTConfig(...).get_freeze_filter()` → freezes all non-LoRA params in `.llm.`

Validation setup

Split: 10 single + 15 multi episodes held out (deterministic, seed=42). 9,569 val frames.
Train sampler: WeightedRandomSampler over the full MultiLeRobotDataset; val frame indices zeroed so they're never drawn in training.
Val sampler: torch.utils.data.Subset of the same MultiLeRobotDataset, plain shuffle.
Eval pass: 50 batches × bs=32 = 1,600 val samples per eval, ~2.5 min on A100.
Metric logged: val/loss = mean cross-entropy of FAST action tokens (same metric as train/loss, no grad).

WandB runs

Stage	Run	Project
warmup	1myk72cf	`pi0-fast-mrfood`
main	prlijwmm	`pi0-fast-mrfood`
refine	wdrms755	`pi0-fast-mrfood`

Honest limitations

No physical-rollout evaluation in this repo. Loss numbers do not directly translate to success rate; only val/loss is reported.
Sim-to-real gap: training frames are recorded video at 20 fps; real robot has actuator latency, frame jitter, plate-position variance. Expect 10–20% SR drop vs val/loss-implied performance.
Egocentric IK quality is only ~43% before post-processing. The 10% mix in main and 3% mix in refine treat ego as visual augmentation; trajectories from ego clips should not be expected to drive real-robot grasping.
LoRA over-conservatism: trainable capacity is 50M / 3.3B = 1.5% of model. If first-deploy SR is poor, the next lever is unfreezing the last 4 Gemma layers (still mostly frozen backbone).
One run, one seed. No multi-seed averaging.

Files in this repo

checkpoints/
  warmup/
    500/   (~8 GB Orbax)
    999/   (~8 GB)
  main/
    2000/  (~8 GB)
    4000/  (~8 GB)  ← recommended
  refine/
    1000/  (~8 GB)
    2000/  (~8 GB)  ← final (run stopped early on val regression)
scripts/
  config.py
  data_loader.py
  train.py
  piper_policy.py
  teleop_to_v21.py
  ego_to_lerobot.py
  eval_ckpt.py
README.md   (this file)

Citation / acknowledgements

Base model: π₀-FAST (Physical Intelligence)
Codebase: openpi (Physical-Intelligence/openpi)
LIBERO warm-start: gs://openpi-assets/checkpoints/pi0_fast_libero
Teleop data: PranayTest/02062026_{single,multi}_plate_pick
Egocentric data: EgoDex (Apple), EPIC-KITCHENS-100, EGTEA Gaze+, Charades-Ego — retargeted via HaWoR + IK
Training compute: 1× A100 80GB on Vast.ai

License: Apache 2.0