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:

from openpi.training import config as _config
from openpi.training import weight_loaders
from huggingface_hub import snapshot_download

ckpt_dir = snapshot_download("brandonyang/pi05_droid_overlay_v1")

cfg = _config.get_config("pi05_droid_finetune")
cfg = dataclasses.replace(
    cfg,
    weight_loader=weight_loaders.CheckpointWeightLoader(f"{ckpt_dir}/params"),
    data=dataclasses.replace(
        cfg.data,
        assets=dataclasses.replace(cfg.data.assets, assets_dir=ckpt_dir),
    ),
)

Or, for the standard serving path:

uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_droid_finetune \
    --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.