tennyyyin/pi05-clean-spill-dagger

pi0.5 clean-spill policies — base + DAgger fine-tunes

Four pi0.5 checkpoints for the bimanual BimanualCleanUpSpill task ("pick up the knocked over cup, set it upright, and wipe up the spilled liquid with a towel"), for use with openpi.

Each folder holds params/ + assets/ only — no train_state/. These are for inference; you cannot resume training from them.

folder	what it is	steps	notes
base_pi05_clean_spill_v2_3999/	the base clean-spill policy, trained on demos	3,999	the policy both DAgger rounds continue from
r1_1000/	DAgger round 1 fine-tune	1,000	9 takes on 5 legacy anchors, mixed 50/50 with demos
r2_1000/	DAgger round 2 fine-tune, 1 GPU, batch 32	1,000	20 takes (r1's 9 + 11 on distractor anchors), 50/50 mix
r2_8gpu_250/	round 2, 8-GPU, batch 128	250	⚠️ only 250 steps — barely trained, included for comparison

r2_8gpu_250 saw ~32k samples (≈3.5 epochs) versus r2_1000's 32k at batch 32. Treat it as an early snapshot, not a finished model.

---

1. Robot computer setup (from scratch)

Needs an NVIDIA GPU, Python ≥3.11, and ~15 GB free per checkpoint.

# 1. uv (dependency manager)
curl -LsSf https://astral.sh/uv/install.sh | sh

# 2. THIS fork of openpi -- not upstream Physical-Intelligence/openpi.
#    The clean-spill configs live only here.
git clone --recurse-submodules https://github.com/tenny-yinyijun/openpi.git
cd openpi

# 3. deps. GIT_LFS_SKIP_SMUDGE=1 is required (openpi pulls LeRobot as a git dep).
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

If the machine fights you on CUDA, openpi ships a Docker path: docs/docker.md. You do not need system CUDA libraries — uv installs them.

2. Download the checkpoints

uv pip install -U "huggingface_hub[cli]"

# all four (~48 GB)
hf download tennyyyin/pi05-clean-spill-dagger --local-dir ./ckpts

# or just one (~12 GB)
hf download tennyyyin/pi05-clean-spill-dagger \
    --include "r2_1000/*" --local-dir ./ckpts

3. Serve a policy

uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_clean_spill \
    --policy.dir=./ckpts/r2_1000

⚠️ Use --policy.config=pi05_clean_spill for ALL FOUR checkpoints

Including the r1/r2 fine-tunes. --policy.config selects the model architecture and the observation packing, not the training recipe, and all four share Pi0Config(pi05=True, action_horizon=16) with the same 3-camera / 16-d-state layout. The training configs (pi05_clean_spill_dagger, ..._v2, ..._v2_8gpu) differ only in dataset, LR and step count — all irrelevant at inference. Two of them are not in the fork's committed history anyway, so asking for them will just fail.

Norm stats come from each checkpoint's own assets/ directory, which is why assets/ is included here. Don't point --policy.dir at a folder missing it, and don't recompute stats.

First inference compiles for ~30 s; after that expect ~150–250 ms per call.

4. Talk to it from the robot

The server is a websocket. Follow docs/remote_inference.md in the openpi repo — that is the README for this part. A minimal client is in examples/simple_client/, and the client package is packages/openpi-client (installable standalone on the robot, so the robot process does not need jax).

from openpi_client import websocket_client_policy
client = websocket_client_policy.WebsocketClientPolicy(host="localhost", port=8000)
action_chunk = client.infer(observation)["actions"]   # (16, 20)

The observation dict it expects

Keys, exactly (this is what openpi's repack transform consumes):

key	shape / type	meaning
observation/image	uint8 [224,224,3]	base camera (scene_right_0)
observation/left_wrist_image	uint8 [224,224,3]	left wrist (wrist_left_plus)
observation/right_wrist_image	uint8 [224,224,3]	right wrist (wrist_right_plus)
observation/state	float32 [16]	measured joint state (14 joints + 2 grippers)
prompt	str	see below

Prompt the models were trained with:

pick up the knocked over cup, set it upright, and wipe up the spilled liquid with a towel

Actions out

float32 [16, 20] — an action horizon of 16 at 10 Hz (1.6 s of motion). Actions are absolute cartesian poses (not deltas), 20-d, laid out:

[ R_xyz(0:3) | R_rot6d(3:9) | L_xyz(9:12) | L_rot6d(12:18) | gripR(18) | gripL(19) ]

Right arm first. State is joint-space, actions are task-space — there is no delta transform.

5. Gotchas

symptom	cause / fix
FileNotFoundError on norm stats, or wildly wrong actions	--policy.dir is missing its assets/ folder. Re-download that checkpoint.
unknown config pi05_clean_spill_dagger_v2	not in the fork's committed history — use pi05_clean_spill, per §3.
LeRobot video decode errors	export OPENPI_VIDEO_BACKEND=pyav (torchcodec's FFmpeg libs are often unloadable). Only matters if you load LeRobot datasets.
ENOMEM from avcodec_open2 with plenty of RAM free	mmap exhaustion, not memory. export MALLOC_ARENA_MAX=2.
GPU OOM at serve time	the policy alone is ~12 GB of params; give it a card with headroom.