openpi_skill

Overview

This skill packages end-to-end workflows to fine-tune and serve Physical Intelligence OpenPI models (pi0, pi0-fast, pi0.5) using JAX or PyTorch backends. It covers blank-machine setup, computing normalization statistics, JAX and PyTorch training, JAX→PyTorch checkpoint conversion, and running low-latency WebSocket policy servers for robot inference. The content targets robot manipulation tasks across ALOHA, DROID, and LIBERO environments.

How this skill works

The skill provides config-driven recipes and scripts to prepare datasets, compute norm stats, launch training (JAX or PyTorch), convert checkpoints, and start a policy server. Serving exposes a WebSocket API for lightweight clients to send observations and receive action chunks for real-time control. It includes troubleshooting tips for norm stats, memory tuning, transformer patching, and cluster-oriented environment variables for HPC users.

When to use it

• Fine-tuning pi0, pi0-fast, or pi0.5 on custom LeRobot or RLDS-format datasets.
• Converting official JAX checkpoints to PyTorch for deployment or further PyTorch training.
• Running a low-latency policy inference server for robot or simulator integration.
• Debugging training failures: missing norm stats, OOMs, or diverging PyTorch runs.
• Preparing runs on HPC/Slurm where cache routing and resource flags are required.

Best practices

• Always compute normalization statistics whenever configs, transforms, or datasets change.
• Start from the closest preset config, then minimally edit dataset mapping and weight loader.
• Use XLA/JAX memory env vars and fsdp_devices to avoid OOMs for large JAX jobs.
• For PyTorch, apply the provided transformer patches and test conversion on a small checkpoint.
• Serve with a preset checkpoint first, validate with the example client, then integrate into robot code.

Example use cases

• Fine-tune pi0.5 on a custom DROID subset: compute norm stats → JAX train → serve checkpoint → validate.
• Convert a JAX pi0 checkpoint to PyTorch, apply transformer patches, and run distributed torchrun training.
• Deploy a pi0-fast policy server for low-latency robot control in ALOHA with a WebSocket client.
• Debug a failing training job by re-running compute_norm_stats, lowering batch size, or enabling fsdp sharding.
• Run inference on an HPC node: export HF_HOME/XDG_CACHE_HOME to scratch and start the server under srun.