brandonyang/pi05_droid_trajectory_overlay

pi0.5-DROID Fine-tuned on Trajectory Overlay Data

Fine-tuned pi05_droid model on DROID episodes with 2D gripper/object trajectory overlays rendered on exterior_image_1_left. The overlays were produced by a Gemini ER + SAM2 tracking pipeline (see dataset card).

Training Setup

Setting	Value
Base model	gs://openpi-assets/checkpoints/pi05_droid
Dataset	brandonyang/droid_1.0.1_trajectory_overlay (4,444 episodes)
Config	pi05_droid_finetune
Steps	20,000
Batch size	32
Hardware	4× A100 80GB
Total train time	~20.8 hours
Final loss	0.0064 (from initial 2.24)

Loss Curve

Step	Loss	Grad Norm
0	2.2385	18.17
1,000	0.0228	0.43
5,000	0.0130	0.23
10,000	0.0098	0.17
15,000	0.0083	0.13
19,900	0.0064	0.12

Usage

Serve as an openpi policy server:

uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_droid \
    --policy.dir=<local_path_to_this_checkpoint>

Or download to a local directory first:

from huggingface_hub import snapshot_download
snapshot_download("brandonyang/pi05_droid_trajectory_overlay", local_dir="./pi05_droid_trajectory_overlay")

Then feed DROID exterior_image_1_left observations with trajectory overlays drawn (same style as training data) to the policy.

Files

• params/ — model weights (12 GB)
• assets/ — norm stats + metadata (reused from base pi05_droid)
• _CHECKPOINT_METADATA — orbax checkpoint metadata

Optimizer state (train_state/) is not included — this is an inference-only checkpoint.

Related

• Dataset — the 4,444-episode training dataset
• WandB run — live training metrics