ykail/pi0_droid_fp32

π₀ - Droid

This is a PyTorch version of the PI0 pi0_droid model, converted from the original JAX/Flax implementation.

Model Details

• Architecture: PI0 (Vision-Language-Action model)
• Model Type: PI0
• Domain: DROID (robotic manipulation)
• Precision: 32-bit floating point (fp32)
• Vision Model: PaliGemma (gemma_2b)
• Action Expert: gemma_300m

Key Features

• Vision-Language-Action: Multimodal model combining vision, language, and action
• PaliGemma Backbone: Leverages PaliGemma for vision-language understanding
• Continuous State Input: Direct continuous state input processing

Conversion Details

This model was converted from JAX to PyTorch using the OpenPI conversion script:

python examples/convert_jax_model_to_pytorch.py \
    --checkpoint_dir /pi0_droid \
    --config_name pi0_droid \
    --output_path /pi0_droid/pytorch/fp32/ \
    --precision float32

Usage

from openpi.models_pytorch.pi0_pytorch import PI0Pytorch
import torch

# Load the model
model = PI0Pytorch.from_pretrained("pepijn223/pi0_droid_fp32")

# The model expects inputs in the format:
# - images: torch.Tensor of shape [batch, height, width, channels]  
# - text: tokenized text prompts
# - proprioceptive_state: robot state information (if applicable)

Model Architecture

The model consists of:

1. Vision Encoder: PaliGemma-based vision processing
2. Language Encoder: Text prompt understanding
3. Action Expert: Specialized network for action prediction
4. Integration Layer: Combines multimodal information for action output

Training Data

This model was trained on robotics datasets appropriate for its domain:

• DROID models: Trained on diverse robot manipulation data
• LIBERO models: Trained on diverse tabletop manipulation scenarios
• Base models: Trained on general robotics datasets

Limitations

• Model performance depends on similarity between deployment and training environments
• May require domain-specific fine-tuning for optimal performance
• Action space must match the trained action dimension (32)

Citation

If you use this model, please cite the original OpenPI work:

@article{openpi2024,
    title={Open-World Robotic Manipulation with Vision-Language-Action Models},
    author={Physical Intelligence},
    year={2024},
    url={https://github.com/Physical-Intelligence/openpi}
}

Original Repository

OpenPI GitHub Repository

License

This model follows the same license as the original OpenPI repository.