|
|
--- |
|
|
license: mit |
|
|
tags: |
|
|
- robotics |
|
|
- pi-zero |
|
|
- diffusion |
|
|
- vision-language-action |
|
|
- aloha |
|
|
- manipulation |
|
|
- bolt-nut-sorting |
|
|
base_model: google/paligemma-3b-pt-224 |
|
|
library_name: openpi |
|
|
pipeline_tag: robotics |
|
|
--- |
|
|
|
|
|
# Pi-0 Bolt Nut Sort Model |
|
|
|
|
|
This is a Pi-0 (Pi-Zero) model trained for bolt and nut sorting tasks using the OpenPI framework. |
|
|
|
|
|
## Model Description |
|
|
|
|
|
- **Architecture**: Pi-0 (diffusion-based vision-language-action model) |
|
|
- **Base Model**: PaLiGemma 3B with SigLIP vision encoder |
|
|
- **Task**: Sorting bolts and nuts into separate baskets |
|
|
- **Robot**: Dual-arm ALOHA setup |
|
|
- **Action Space**: 14-DoF (7 per arm: 6 joints + 1 gripper) |
|
|
- **Training Steps**: 29,999 |
|
|
- **Action Horizon**: 50 steps |
|
|
- **Image Resolution**: 224x224 |
|
|
|
|
|
## Dataset |
|
|
|
|
|
Trained on the `naungth/pi0_bolt_nut_sort` dataset with the task instruction: |
|
|
"sort the bolts and the nuts into separate baskets" |
|
|
|
|
|
## Usage |
|
|
|
|
|
### With OpenPI |
|
|
|
|
|
```python |
|
|
from openpi.policies import policy_config |
|
|
from openpi.training import config |
|
|
|
|
|
# Load the model configuration |
|
|
config_name = "pi0_bns" |
|
|
train_config = config.get_config(config_name) |
|
|
|
|
|
# Create policy from your local checkpoint |
|
|
policy = policy_config.create_trained_policy( |
|
|
train_config, |
|
|
"path/to/checkpoint", |
|
|
default_prompt="sort the bolts and the nuts into separate baskets" |
|
|
) |
|
|
|
|
|
# Use for inference |
|
|
observation = { |
|
|
"images": { |
|
|
"cam_high": image_array, # [H, W, 3] uint8 |
|
|
"cam_left_wrist": left_wrist_image, # [H, W, 3] uint8 |
|
|
"cam_right_wrist": right_wrist_image, # [H, W, 3] uint8 |
|
|
}, |
|
|
"state": joint_positions, # [14] float32 |
|
|
"prompt": "sort the bolts and the nuts into separate baskets" |
|
|
} |
|
|
|
|
|
actions = policy.infer(observation)["actions"] # [50, 14] |
|
|
``` |
|
|
|
|
|
### With Policy Server |
|
|
|
|
|
```bash |
|
|
# Start the policy server |
|
|
uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_bns --policy.dir=path/to/checkpoint |
|
|
|
|
|
# Use with client |
|
|
from openpi_client import websocket_client_policy |
|
|
client = websocket_client_policy.WebsocketClientPolicy("localhost", 8000) |
|
|
actions = client.infer(observation) |
|
|
``` |
|
|
|
|
|
## Model Architecture |
|
|
|
|
|
- **Vision Encoder**: SigLIP-So400m/14 |
|
|
- **Language Model**: Gemma 2B + Gemma 300M (action expert) |
|
|
- **Training**: Diffusion-based action prediction |
|
|
- **Input**: Multi-camera RGB + proprioception + language instruction |
|
|
- **Output**: Future action sequence (50 timesteps) |
|
|
|
|
|
## Training Details |
|
|
|
|
|
- **Framework**: JAX/Flax with OpenPI |
|
|
- **Optimizer**: AdamW |
|
|
- **Base Checkpoint**: Pi-0 base model from Google |
|
|
- **Fine-tuning**: Task-specific fine-tuning on bolt nut sort data |
|
|
- **Normalization**: Dataset-specific state/action normalization |
|
|
|
|
|
## License |
|
|
|
|
|
MIT License |
|
|
|
|
|
## Citation |
|
|
|
|
|
If you use this model, please cite: |
|
|
|
|
|
```bibtex |
|
|
@article{pi0, |
|
|
title={Pi-Zero: A Diffusion-Based Policy for Robot Manipulation}, |
|
|
author={TODO: Add authors}, |
|
|
year={2024} |
|
|
} |
|
|
``` |
|
|
|
|
|
## Acknowledgments |
|
|
|
|
|
- Built using the [OpenPI](https://github.com/google-deepmind/openpi) framework |
|
|
- Based on the Pi-0 architecture |
|
|
- Training data from bolt nut sorting demonstrations |
|
|
|
