pi05_drift_libero

Pi0.5-Drift · LIBERO (100k)

Pi0.5 is a vision-language-action model from Physical Intelligence that decodes an action chunk by integrating a flow-matching ODE — 10 forward passes of the action expert per chunk. Pi0.5-Drift trains the same network with the one-step Drifting (DBPO) objective instead: a single forward pass maps noise directly to the chunk — 1 NFE, ~3× faster decode at batch 1, byte-identical weight layout.

• Concept & method: https://zuoxingdong.github.io/drift-vla/
• Policy plugin (required to load this model): lerobot_policy_pi05_drift
• Matched flow-matching baseline: zuoxingdong/pi05_fm_libero

Results

LIBERO, official lerobot-eval on lerobot==0.6.0, 3 eval seeds {1000, 1001, 1002}, 50 episodes/task, n_action_steps=10. Both models: identical recipe and 100k-step schedule; the training objective is the only difference.

Drift ≥ FM at every budget of the training sweep — 25k/50k/75k/100k: Drift 96.8 / 97.3 / 97.7 / 97.4 vs FM 96.2 / 95.9 / 96.7 / 96.9.

Evaluation

Needs a graphics-capable node for the MuJoCo/LIBERO renderer, and access to the gated Pi0.5 tokenizer google/paligemma-3b-pt-224 (request access, then hf auth login).

pip install "git+https://github.com/zuoxingdong/lerobot_policy_pi05_drift.git"
pip install "lerobot[libero]" "mujoco==3.3.2"

lerobot-eval \
  --policy.path=zuoxingdong/pi05_drift_libero \
  --policy.n_action_steps=10 \
  --env.type=libero \
  --env.task=libero_spatial \
  --eval.n_episodes=50 \
  --eval.batch_size=2 \
  --eval.use_async_envs=true \
  --seed=1000 \
  --output_dir=eval/pi05_drift_libero

--policy.n_action_steps=10 is load-bearing (closed-loop replanning; the chunk tail is discarded).

Optional KeyStone test-time selection (K one-step candidates, guarded cluster-medoid pick, ~zero added latency) — append to the eval command:

  --policy.test_time_samples=8 \
  --policy.test_time_clusters=2 \
  --policy.test_time_unimodal_tau=0.3

Note: measured on this checkpoint (config sweep, 3 seeds), K=8/C=2 is the only net-positive setting and the effect is suite-dependent — libero_object +0.7, libero_10 +0.7 (with much tighter seed variance), libero_goal −1.0; overall ≈ +0.2. Larger C or pure-medoid selection hurts. Treat it as a per-task knob rather than a blanket default.

Finetune on your data

The drift flags shown below are the plugin defaults (the winning recipe) — spelled out for clarity:

lerobot-train \
  --policy.type=pi05_drift \
  --policy.pretrained_path=zuoxingdong/pi05_drift_libero \
  --policy.use_drifting_loss=true \
  --policy.drifting_gen_per_label=8 \
  --policy.drifting_temperatures='[0.02, 0.05, 0.2]' \
  --policy.drifting_perdim_loss=true \
  --policy.drifting_per_timestep_loss=false \
  --policy.dtype=bfloat16 \
  --policy.gradient_checkpointing=true \
  --dataset.repo_id=lerobot/libero \
  --dataset.revision=a1aaacb7f6cd6ee5fb43120f673cebb0cfea7dd4 \
  --batch_size=56 \
  --steps=10000

The dataset shown is the one this model was trained on — lerobot/libero (LIBERO, LeRobot v3 format), pinned to the training revision. To finetune on your own data, swap in your --dataset.repo_id. Add --policy.freeze_vision_encoder=true --policy.train_expert_only=true for the ~4× cheaper expert-only variant (near-lossless for Drift).

Use from Python

import lerobot_policy_pi05_drift              # registers "pi05_drift" — import before loading
from lerobot.policies.factory import get_policy_class

policy = get_policy_class("pi05_drift").from_pretrained("zuoxingdong/pi05_drift_libero")
action = policy.select_action(batch)          # (B, action_dim); 1-NFE drift inference

Training details

• Base: lerobot/pi05_libero VLM with a freshly re-initialized action expert, full-parameter finetune on lerobot/libero (all four suites), MEAN_STD normalization, bf16, effective batch 56, LR 2.5e-5 cosine-decayed over 100k steps (this checkpoint = step 100k, schedule end), seed 1000.
• Stack: lerobot==0.6.0 + lerobot_policy_pi05_drift; the matched FM baseline differs only by use_drifting_loss=false.

License

Weights derive from Pi0.5 / PaliGemma and are distributed under the Gemma license, matching the upstream lerobot/pi05_libero release. The plugin code is Apache-2.0.