pi05_droid_overlay_v1
pi05_droid_overlay_v1 is a fine-tune of
Physical Intelligence's pi05-DROID
on
brandonyang/droid_overlay_full: a
3,776-episode / 995,769-frame DROID-format dataset where the
exterior_image_1_left camera has
retrieval-trajectory overlays baked into every frame (rainbow transition segments,
square/circle gripper-event markers, direction arrows; matches the
visualize_retrieval_results.ipynb notebook style).
Files
| Path | Size | Purpose |
|---|
params/ | ~12 GB | Orbax PyTree checkpoint of the model weights (after EMA). Sufficient for inference. |
assets/droid/norm_stats.json | small | Joint / gripper / action normalization statistics (reused from the original gs://openpi-assets/checkpoints/pi05_droid/assets/droid). |
_CHECKPOINT_METADATA | 426 B | Orbax checkpoint manifest. |
The training-state directory (train_state/, ~30 GB) is intentionally not uploaded — it is
needed only to resume training from this exact optimizer state. For inference, fine-tuning from
these weights, or evaluation, only params/ and assets/ are required.
Training summary
| |
|---|
| Base model | pi05_droid from gs://openpi-assets/checkpoints/pi05_droid/params |
| Train config | pi05_droid_finetune (src/openpi/training/config.py), unchanged |
| Dataset | brandonyang/droid_overlay_full (LeRobot v2.1, 3,776 ep, 995,769 frames) |
| Steps | 20,000 (final checkpoint at step 19,999) |
| Batch size | 32 (8 per device with 4-way FSDP) |
| Sharding | fsdp_devices=4 (model sharded across all 4 GPUs, no replication) |
| Hardware | 4× NVIDIA A100 80GB PCIe (GPUs 6–9) |
| Wall time | ~16.5 hours |
| Wandb | https://wandb.ai/brandonyang/openpi/runs/n48h7bmc |
End-of-run training metrics (around step 19,900):
| Metric | Value |
|---|
loss | ~0.013 (down from ~0.016 mid-run) |
grad_norm | ~0.11 (stable throughout) |
param_norm | ~1941 (slow monotonic growth) |
Norm stats are reused from the upstream pi05_droid since action/state spaces and joint/gripper
ranges match standard DROID — only the exterior_image_1_left pixel content differs.
Use with openpi
Point the CheckpointWeightLoader at this repo's params/ directory (downloaded via
huggingface_hub) and pair with the standard pi05_droid_finetune config:
1from openpi.training import config as _config
2from openpi.training import weight_loaders
3from huggingface_hub import snapshot_download
4
5ckpt_dir = snapshot_download("brandonyang/pi05_droid_overlay_v1")
6
7cfg = _config.get_config("pi05_droid_finetune")
8cfg = dataclasses.replace(
9 cfg,
10 weight_loader=weight_loaders.CheckpointWeightLoader(f"{ckpt_dir}/params"),
11 data=dataclasses.replace(
12 cfg.data,
13 assets=dataclasses.replace(cfg.data.assets, assets_dir=ckpt_dir),
14 ),
15)
Or, for the standard serving path:
1uv run scripts/serve_policy.py policy:checkpoint \
2 --policy.config=pi05_droid_finetune \
3 --policy.dir=$(huggingface-cli download brandonyang/pi05_droid_overlay_v1)
What "trajectory overlay" means
Each frame's exterior_image_1_left has the projected 2D paths of 1–5 retrieved
nearest-neighbor DROID episodes drawn on top of the raw camera image. Within each retrieved
trajectory's gripper-close → gripper-open window, segments are drawn with a vivid rainbow
gradient at α=0.80; outside that window the line fades to α=0.10 against a per-trajectory
base-color gradient. Square markers denote gripper_open events and circles denote
gripper_close; small filled-triangle arrows annotate motion direction inside the
manipulation window.
The model sees this overlay as a normal RGB input — no auxiliary loss or special token is
added. It is up to the downstream policy / fine-tune to learn whether (and how) to attend to it.
