Skip to main content

Pi0 and Pi0.5 Model Fine-tuning: Physical Intelligence Open Source Robotics Learning Platform

Pi0 and Pi0.5 are industry-leading Vision-Language-Action (VLA) model series developed by the Physical Intelligence team. These models are pre-trained on large-scale robotic data and can quickly adapt to specific robotic tasks and application scenarios through fine-tuning.

This guide is based on the official OpenPI framework from Physical Intelligence, providing detailed instructions on how to efficiently fine-tune pre-trained Pi0/Pi0.5 models on custom datasets.

Model Series Overview​

The OpenPI repository currently provides three types of models:

Model TypeDescriptionFeaturesUse Cases
π₀ (Pi0)Flow-matching based Vision-Language-Action modelFlow-based generation, high-quality action predictionComplex manipulation tasks, high precision requirements
π₀-FASTAutoregressive VLA based on FAST action tokenizerAutoregressive generation, fast inferenceReal-time control, low latency requirements
π₀.₅ (Pi0.5)Upgraded version of π₀ with better open-world generalizationKnowledge-isolated training, enhanced generalizationDiverse environments, cross-domain applications

Note: All models are pre-trained on 10,000+ hours of robotic data, providing base model checkpoints for fine-tuning.

Prerequisites and Environment Setup​

System Requirements​

Minimum Configuration:

  • Python 3.11+ (recommended to use uv package manager)
  • GPU: NVIDIA GPU required for training (recommended A100/H100)
  • Memory: 32GB+ system RAM
  • Storage: 100GB+ available disk space

Recommended Configuration:

  • Hardware: NVIDIA A100/H100 or multi-GPU setup
  • Storage: NVMe SSD solid-state drive
  • Network: Stable network connection for model and data downloads

Environment Installation and Configuration​

# Clone OpenPI official repository
git clone https://github.com/Physical-Intelligence/openpi.git
cd openpi

# Install dependencies using uv (automatically creates virtual environment)
uv sync

# Verify installation
uv run python -c "from openpi.policies import policy_config; print('OpenPI environment configured successfully!')"

Traditional pip Installation​

# Clone repository
git clone https://github.com/Physical-Intelligence/openpi.git
cd openpi

# Create virtual environment
python -m venv venv
source venv/bin/activate  # Linux/macOS
# or venv\Scripts\activate  # Windows

# Install dependencies
pip install -e .

# Verify installation
python -c "from openpi.policies import policy_config; print('OpenPI environment configured successfully!')"

Dataset Preparation​

Using IO Platform for Data Export​

The IO Data Platform provides complete data export functionality, allowing one-click export of annotated data to LeRobot format supported by OpenPI:

Export Process:

  1. Select Export Format: Choose LeRobot or HDF5 format export in the IO platform
  2. Filter Data: Filter required data based on project, time, quality level, and other dimensions
  3. Batch Export: Support batch export processing for large-scale datasets
  4. Format Conversion: Convert exported data to LeRobot format required by OpenPI

Export LeRobot Format Data

Data Format Requirements​

OpenPI supports multiple data formats. Datasets need to include:

  • Observation Data: Images, state information, etc.
  • Action Data: Action sequences executed by the robot
  • Language Instructions: Task descriptions (optional, for language-conditioned tasks)

Dataset Configuration​

OpenPI uses configuration files to define datasets and training parameters. You need to:

  1. Create Dataset Configuration: Define data paths, formats, etc. in configuration files
  2. Data Preprocessing: Compute normalization statistics
  3. Validate Data: Ensure data format is correct
# Compute normalization statistics for dataset (required step)
uv run scripts/compute_norm_stats.py <config_name>

# Example: Compute statistics for custom configuration
uv run scripts/compute_norm_stats.py my_custom_config

Supported Data Formats​

OpenPI supports various robotic data formats, including:

  • HuggingFace Datasets
  • LeRobot Format Data
  • HDF5 Format Data
  • Data exported from IO Platform (supported through format conversion)

Model Training Practice​

The IO Data Platform provides complete model training functionality, supporting training of various robotic learning models including OpenPI:

IO Platform Training Overview

Platform Training Advantages:

  • No-Code Training: Complete the entire training process through visual interface without programming experience
  • Flexible Computing Resources: Support private and public cloud computing resources, allocated on demand
  • Real-time Monitoring: Provide real-time monitoring of training metrics, model outputs, and system logs
  • Automated Management: Support checkpoint management, resume training, parameter adjustment, and other functions

Training Workflow​

  1. Data Preparation: Select exported datasets or upload external data Select Training Data Source

  2. Model Configuration: Select Pi0/Pi0.5 models and configure training parameters Training Parameter Configuration

  3. Training Monitoring: View training progress and model performance in real-time Training Detail Monitoring

  4. Model Export: Obtain model checkpoints after training completion Model Checkpoint Management

Local JAX Version Training​

If you need to train in a local environment, OpenPI primarily uses JAX for training, providing optimal performance and stability:

# Set JAX memory allocation (recommended)
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9

# Basic training command
uv run scripts/train.py <config_name> --exp_name <experiment_name>

# Example: Train Pi0 model
uv run scripts/train.py pi0_aloha_sim --exp_name my_pi0_experiment

# Multi-GPU training (using FSDP)
uv run scripts/train.py <config_name> --exp_name <experiment_name> --fsdp-devices <number_of_GPUs>

# Example: Train using 4 GPUs
uv run scripts/train.py pi0_aloha_sim --exp_name my_experiment --fsdp-devices 4

PyTorch Version Training​

OpenPI also supports PyTorch training, but requires environment configuration first:

PyTorch Environment Configuration​

# Install PyTorch support (requires replacing transformers files)
cp -r src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/

# Warning: This permanently modifies the transformers library and may affect other projects

PyTorch Training Commands​

# Single GPU training
uv run scripts/train_pytorch.py <config_name> --exp_name <experiment_name>

# Multi-GPU training (single node)
uv run torchrun --standalone --nnodes=1 --nproc_per_node=<number_of_GPUs> \
    scripts/train_pytorch.py <config_name> --exp_name <experiment_name>

# Example: Train using 2 GPUs
uv run torchrun --standalone --nnodes=1 --nproc_per_node=2 \
    scripts/train_pytorch.py pi0_aloha_sim --exp_name pytorch_experiment

Model Inference and Deployment​

The IO Data Platform provides complete model inference services, supporting full-scenario AI inference capabilities from cloud validation to edge deployment:

New Inference Service

Platform Inference Advantages:

  • One-Click Deployment: Deploy trained models as production-level inference services with one click
  • Multiple Testing Methods: Support simulation inference, MCAP file testing, offline edge deployment
  • Real-time Monitoring: Provide real-time monitoring of service status, resource usage, and performance metrics
  • Intelligent Adaptation: Automatically identify and adapt to different models' input/output requirements

Inference Service Features​

  1. Service Status Monitoring: View inference service running status and configuration information Inference Service Details

  2. Simulation Inference Testing: Use random data to quickly verify model functionality Simulation Inference Testing

  3. MCAP File Testing: Use real robotic data to verify inference effectiveness MCAP File Testing

  4. Offline Edge Deployment: Deploy inference services to robot's local GPU Offline Edge Deployment

Deployment Method Comparison​

Inference MethodUse CaseDescription
Simulation Inference TestingQuick VerificationUse random data or custom inputs to quickly verify model inference functionality and performance
MCAP File TestingReal Data VerificationUse recorded robotic demonstration data to verify model inference effectiveness in real scenarios
Offline Edge DeploymentProduction Environment ApplicationDeploy inference services to robot's local GPU for low-latency real-time control

Local Policy Server Deployment​

If you need to deploy in a local environment, OpenPI provides policy servers for model inference deployment:

# Start policy server (JAX version)
uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=<config_name> \
    --policy.dir=<checkpoint_path>

# Example: Deploy Pi0.5 model
uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_droid \
    --policy.dir=/path/to/trained/checkpoint

# PyTorch version deployment
uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi05_droid \
    --policy.dir=/path/to/converted/pytorch/checkpoint

Python API Usage​

from openpi.training import config as _config
from openpi.policies import policy_config

# Load configuration
config = _config.get_config("pi05_droid")
checkpoint_dir = "/path/to/trained/checkpoint"

# Create trained policy
policy = policy_config.create_trained_policy(config, checkpoint_dir)

# Execute inference
action_chunk = policy.infer(example)["actions"]

Available Configurations​

OpenPI provides multiple predefined configurations suitable for different robotic platforms and tasks:

Configuration NameDescriptionUse Case
pi0_aloha_simPi0 model, ALOHA simulation environmentALOHA robot simulation training
pi05_droidPi0.5 model, DROID datasetReal robot data training
debugDebug configuration, small datasetQuick testing and debugging

Custom Configuration​

You can create custom configuration files to adapt to specific datasets and training requirements:

# Add new configuration in src/openpi/training/config.py
@config_flags.DEFINE_config_file("config")
def get_config(config_name: str):
    # Modify based on existing configuration
    config = get_base_config()
    
    # Custom dataset path
    config.dataset.path = "/path/to/your/dataset"
    
    # Adjust training parameters
    config.training.learning_rate = 1e-4
    config.training.batch_size = 8
    
    return config

JAX vs PyTorch Version Comparison​

Precision Settings​

JAX Version:

  1. Inference: Most weights and computations use bfloat16, with a few computations using float32 for stability
  2. Training: Default mixed precision, weights and gradients use float32, activations and computations use bfloat16

PyTorch Version:

  1. Inference: Matches JAX version, mostly uses bfloat16
  2. Training: Supports full bfloat16 (default) or full float32, configured via pytorch_training_precision

Performance Comparison​

  • Inference Speed: After using torch.compile, PyTorch performance is comparable to JAX
  • Memory Usage: JAX version is typically more efficient
  • Stability: JAX version is more mature and stable

Model Conversion​

If you need to convert models between JAX and PyTorch:

# Convert JAX model to PyTorch
uv run examples/convert_jax_model_to_pytorch.py \
    --checkpoint_dir /path/to/jax/checkpoint \
    --config_name <config_name> \
    --output_path /path/to/pytorch/checkpoint

Common Problem Solving​

Memory and Performance Optimization​

  1. GPU Memory Insufficient:

    # Set JAX memory allocation ratio
    export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9

    # Use FSDP to reduce memory usage
    uv run scripts/train.py <config> --exp_name <name> --fsdp-devices <n>

  2. Training Speed Optimization:

    • Use SSD storage to improve data loading speed
    • Appropriately increase batch size and gradient accumulation
    • Consider disabling EMA (Exponential Moving Average) to save memory

  1. Normalization Statistics Error:

    # Recompute normalization statistics
    uv run scripts/compute_norm_stats.py <config_name>

  2. Dataset Download Failure:

    • Check network connection
    • For HuggingFace datasets, ensure you're logged in: huggingface-cli login

Environment Configuration Issues​

  1. Dependency Conflicts:

    # Clean virtual environment and reinstall
    rm -rf .venv
    uv sync

  2. CUDA Errors:

    • Verify NVIDIA driver installation is correct
    • For Docker, ensure nvidia-container-toolkit is installed
    • May need to uninstall system-level CUDA libraries to avoid conflicts

Training Monitoring and Evaluation​

Training Logs​

OpenPI training process automatically records detailed logs, including:

  • Loss Curves: Training loss change trends
  • Learning Rate Scheduling: Learning rate change situations
  • Resource Usage: GPU memory and computation utilization
  • Training Metrics: Various evaluation metrics

Checkpoint Management​

# Training checkpoints saved in experiment directory
ls experiments/<exp_name>/checkpoints/

# Resume training (continue from latest checkpoint)
uv run scripts/train.py <config> --exp_name <name> --resume

# Resume from specific checkpoint
uv run scripts/train.py <config> --exp_name <name> --resume --checkpoint_path <path>

Model Evaluation​

Use trained models for inference evaluation:

from openpi.policies import policy_config
from openpi.training import config as _config

# Load configuration and model
config = _config.get_config("pi05_droid")
policy = policy_config.create_trained_policy(config, checkpoint_dir)

# Execute inference
result = policy.infer(observation_data)
actions = result["actions"]

IO Platform Integrated Evaluation​

The IO Platform provides integrated support for OpenPI models, enabling:

  • Visual Comparison: Visual comparison between real data and model outputs
  • Joint-level Analysis: Prediction accuracy analysis for each joint
  • Task Success Rate Statistics: Task completion statistics after actual deployment

Evaluation Metrics​

Quantitative Metrics​

  • Action Accuracy: Error statistics between predicted and real actions
  • Trajectory Consistency: Smoothness and reasonableness of generated trajectories
  • Convergence Analysis: Loss function convergence during training process

Qualitative Evaluation​

  • Real Robot Verification: Test model performance on real robots
  • Generalization Ability: Adaptation ability in unseen scenarios
  • Robustness Testing: Stability under noise and interference

IO Platform Complete Workflow​

End-to-End Robotic Learning Process​

The IO Data Platform provides complete end-to-end solutions for OpenPI models, supporting the entire pipeline from data collection to model deployment:

Platform Integration Advantages​

Advantages compared to pure OpenPI development:

  1. Lower Technical Barrier: No deep learning or programming experience required, operate through visual interface
  2. Improved Development Efficiency: Full-process automation from data to model, significantly shortening development cycles
  3. Guaranteed Data Quality: Professional data annotation and quality control processes
  4. Simplified Deployment Process: One-click deployment to production environment, supporting multiple deployment methods
  5. Continuous Optimization Support: Provide model performance monitoring and continuous improvement mechanisms

Usage Recommendations​

Recommended Use Cases:

  • Beginners or Non-technical Users: Use IO platform's visual interface for complete process operation
  • Rapid Prototype Validation: R&D teams needing to quickly validate ideas and concepts
  • Production Environment Deployment: Solutions requiring stable and reliable production-level deployment
  • Large-scale Data Processing: Projects requiring processing of large amounts of robotic data

Local Development Scenarios:

  • Deep Customization Needs: Advanced users needing to modify model architectures or training algorithms
  • Research Experiments: Academic users conducting algorithm research and experimental comparisons
  • Offline Environment: Network-restricted or high-security requirement environments

Frequently Asked Questions​

Q: How to choose between JAX and PyTorch versions?​

A:

  • JAX Version: Recommended for production environments, more stable and mature
  • PyTorch Version: Suitable for scenarios requiring PyTorch ecosystem integration

Q: How to use Pi0.5 model in OpenPI?​

A: Pi0.5 model is fully supported in OpenPI, just use configuration pi05_droid or similar Pi0.5 configurations.

Q: How to handle memory insufficiency issues?​

A:

# Set JAX memory allocation
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9

# Use FSDP distributed training
uv run scripts/train.py <config> --exp_name <name> --fsdp-devices <n>

Q: What format is required for training data?​

A: OpenPI supports multiple data formats, including HuggingFace dataset format. Data exported from IO platform needs format conversion.

Q: How to monitor training progress?​

A: OpenPI automatically saves training logs and checkpoints to experiments/<exp_name>/ directory, you can monitor training progress through log files.

Q: How long does model fine-tuning take?​

A: Training time depends on:

  • Dataset size
  • Hardware configuration (number and model of GPUs)
  • Model complexity
  • Training configuration parameters

Typically takes several hours to several days on a single A100.

Q: How to integrate trained models into IO platform?​

A: The IO platform provides integrated support for OpenPI models, including model registration, deployment, and monitoring functions. Please consult the technical support team for specific integration methods.

Through this guide, you should be able to successfully use the OpenPI framework to train and deploy Pi0/Pi0.5 models. If you have any questions, please refer to the official documentation or seek technical support.