|
|
--- |
|
|
license: apache-2.0 |
|
|
tags: |
|
|
- robotics |
|
|
- lerobot |
|
|
- pi0-fast |
|
|
- imitation-learning |
|
|
- vla |
|
|
datasets: |
|
|
- gpudad/so101_pick_cube_chunked |
|
|
base_model: |
|
|
- lerobot/pi0fast-base |
|
|
pipeline_tag: robotics |
|
|
--- |
|
|
|
|
|
# π₀-FAST SO101 Pick & Place |
|
|
|
|
|
A finetuned [π₀-FAST](https://huggingface.co/lerobot/pi0fast-base) model for pick and place tasks on the SO101 robot arm. |
|
|
|
|
|
## Model Details |
|
|
|
|
|
- **Base Model:** lerobot/pi0fast-base (3B parameters) |
|
|
- **Training Dataset:** [gpudad/so101_pick_cube_chunked](https://huggingface.co/datasets/gpudad/so101_pick_cube_chunked) |
|
|
- 10,990 episodes |
|
|
- 1,456,443 frames @ 30 FPS |
|
|
- 3 cameras: front, overhead, wrist (512x512) |
|
|
- 6-DOF action space |
|
|
- **Training Steps:** 10,000 (quick validation run) |
|
|
- **Final Loss:** 2.35 |
|
|
- **Hardware:** NVIDIA RTX 5090 (32GB VRAM) |
|
|
|
|
|
## Performance |
|
|
|
|
|
Tested on held-out samples from the dataset: |
|
|
|
|
|
| Metric | Value | |
|
|
|--------|-------| |
|
|
| Mean MAE | 0.079 | |
|
|
| Relative Error | ~2.6% of action range | |
|
|
| Best MAE | 0.0085 | |
|
|
|
|
|
## Usage |
|
|
|
|
|
```python |
|
|
from lerobot.policies.pi0_fast.modeling_pi0_fast import PI0FastPolicy |
|
|
from lerobot.processor.pipeline import PolicyProcessorPipeline |
|
|
|
|
|
# Load model |
|
|
policy = PI0FastPolicy.from_pretrained("gpudad/pi0fast-so101-pick-cube") |
|
|
policy.to("cuda") |
|
|
policy.eval() |
|
|
|
|
|
# Load processors |
|
|
preprocessor = PolicyProcessorPipeline.from_pretrained( |
|
|
"gpudad/pi0fast-so101-pick-cube", |
|
|
"policy_preprocessor.json" |
|
|
) |
|
|
postprocessor = PolicyProcessorPipeline.from_pretrained( |
|
|
"gpudad/pi0fast-so101-pick-cube", |
|
|
"policy_postprocessor.json" |
|
|
) |
|
|
|
|
|
# Run inference |
|
|
observation = { |
|
|
"observation.state": state_tensor, |
|
|
"observation.images.front": front_image, |
|
|
"observation.images.wrist": wrist_image, |
|
|
"observation.images.overhead": overhead_image, |
|
|
"task": "pick up the object and place it in the target location", |
|
|
} |
|
|
|
|
|
batch = preprocessor(observation) |
|
|
batch['observation.language.attention_mask'] = batch['observation.language.attention_mask'].bool() |
|
|
|
|
|
policy.reset() |
|
|
with torch.no_grad(): |
|
|
action = policy.select_action(batch) |
|
|
|
|
|
result = postprocessor({"action": action}) |
|
|
final_action = result["action"] |
|
|
``` |
|
|
|
|
|
## Training Configuration |
|
|
|
|
|
```yaml |
|
|
policy.type: pi0_fast |
|
|
policy.dtype: bfloat16 |
|
|
policy.gradient_checkpointing: true |
|
|
policy.chunk_size: 10 |
|
|
policy.n_action_steps: 10 |
|
|
batch_size: 4 |
|
|
optimizer_lr: 2.5e-5 |
|
|
scheduler_warmup_steps: 400 |
|
|
``` |
|
|
|
|
|
## Citation |
|
|
|
|
|
If you use this model, please cite the original π₀ paper and LeRobot: |
|
|
|
|
|
```bibtex |
|
|
@article{black2024pi0, |
|
|
title={$\pi_0$: A Vision-Language-Action Flow Model for General Robot Control}, |
|
|
author={Black, Kevin and Brown, Noah and Driess, Danny and others}, |
|
|
journal={arXiv preprint arXiv:2410.24164}, |
|
|
year={2024} |
|
|
} |
|
|
``` |
|
|
|
