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

π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):

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

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

# 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