yqi19/pi0_finetune

openpi — π0 / π0.5 / π0-FAST ManiSkill Training Guide

This fork adds ManiSkill fine-tuning configs for three model variants. See src/openpi/training/config.py for all training configs.

---

1. Environment Setup

cd /path/to/openpi

# Clean install (run once, or after pulling updates)
rm -rf .venv
export UV_LINK_MODE=copy
GIT_LFS_SKIP_SMUDGE=1 uv sync --index-url https://pypi.tuna.tsinghua.edu.cn/simple
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .
uv pip install datasets==3.0.0

# Patch transformers (required for pi0 PyTorch backend)
cp -r ./src/openpi/models_pytorch/transformers_replace/* \
      .venv/lib/python3.11/site-packages/transformers/

---

2. Data Preparation

Training expects a LeRobot-format dataset at a local path.

/your/dataset/lerobot_dataset/
├── meta/
│   ├── info.json
│   └── stats.json          ← generated by compute_norm_stats.py
├── data/
│   └── chunk-000/
│       └── episode_*.parquet
└── videos/
    └── ...

Step 1 — Compute normalization stats (run once per dataset):

uv run scripts/compute_norm_stats.py pi0_maniskill \
    --data.repo_id /path/to/lerobot_dataset

Replace pi0_maniskill with the config name matching your model variant (see §3).

---

3. Model Variants & Configs

Model	Config name	Architecture	Base checkpoint
π0.5	pi05_maniskill	Flow-matching (large VLM)	gs://…/pi05_base
π0	pi0_maniskill	Flow-matching (smaller)	gs://…/pi0_base
π0-FAST	pi0_fast_maniskill	Autoregressive (FAST tokenizer)	gs://…/pi0_fast_base

Key differences from pi05_maniskill:

pi0_maniskill — removes pi05=True, uses pi0_base weights:

model=pi0_config.Pi0Config(action_horizon=10, discrete_state_input=False),
weight_loader=CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"),
pytorch_weight_path="/root/workspace/openpi/ckpt/pi0-libero-base",

pi0_fast_maniskill — uses Pi0FASTConfig and pi0_fast_base weights:

model=pi0_fast.Pi0FASTConfig(action_dim=8, action_horizon=10, max_token_len=180),
weight_loader=CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_fast_base/params"),

---

4. Training

Start a new run

# π0.5 (default, recommended)
bash run1.sh <exp_name> /path/to/lerobot_dataset

# π0 (flow-matching, smaller base model)
uv run torchrun --standalone --nnodes=1 --nproc_per_node=8 \
    scripts/train_pytorch.py pi0_maniskill \
    --exp_name <exp_name> \
    --data.repo_id /path/to/lerobot_dataset

# π0-FAST (autoregressive)
uv run torchrun --standalone --nnodes=1 --nproc_per_node=8 \
    scripts/train_pytorch.py pi0_fast_maniskill \
    --exp_name <exp_name> \
    --data.repo_id /path/to/lerobot_dataset

• Checkpoints saved to: checkpoints/<config_name>/<exp_name>/
• --nproc_per_node = number of GPUs (default 8; adjust to your node)
• Logs written to run/<exp_name>.log

Key hyperparameters (edit in config.py)

Parameter	Default	Notes
num_train_steps	18,000	Total training steps
save_interval	6,000	Checkpoint frequency
batch_size	256	Reduce if OOM
peak_lr	5e-5	Learning rate

---

5. Resume Training

# Resume with 10,000 more steps (default)
bash resume_run.sh <exp_name> /path/to/lerobot_dataset

# Or with custom step count
EXTRA_STEPS=5000 SAVE_INTERVAL=2000 bash resume_run.sh <exp_name> /path/to/dataset

For π0 / π0-FAST, edit resume_run.sh and change pi05_maniskill to the correct config name.

---

6. Inference

Start the policy server:

uv run scripts/serve_policy.py --env MANISKILL --port 8000 \
    policy:checkpoint \
    --policy.config pi0_maniskill \
    --policy.dir checkpoints/pi0_maniskill/<exp_name>/<step>

Run evaluation:

cd examples/maniskill_conflict
HOST=127.0.0.1 PORT=8000 bash run_inference.sh

---

7. Git Workflow

# Push changes to your fork
git add src/openpi/training/config.py
git commit -m "Add pi0/pi0_fast maniskill configs"
git push origin main

# Pull upstream updates
git fetch upstream
git merge upstream/main

---

Notes

• 8 GPUs (A100/H100 80GB) recommended for batch_size=256; reduce batch size for fewer GPUs
• π0-FAST does not use pytorch_weight_path — downloads JAX checkpoint from GCS automatically
• Dataset path in config (repo_id) must be updated to point to your actual data directory