---
license: apache-2.0
tags:
  - robotics
  - openpi
  - pi0
  - vla
  - vision-language-action
  - trossen
  - aloha
  - lerobot
library_name: openpi
---

# Pi0 Fine-tuned for Trossen AI Stationary Robot (Simulation)

This is a fine-tuned [OpenPI](https://github.com/Physical-Intelligence/openpi) π₀ (pi0) Vision-Language-Action model for the Trossen AI stationary bimanual robot in simulation.

**Developed by [ANR Robot](https://anrrobot.com)**

## ⚠️ Requires Custom Fork

**This model requires a modified version of OpenPI.** You must use [anredlich/openpi](https://github.com/anredlich/openpi) fork, which adds the `adapt_trossen_to_pi` parameter to `aloha_policy.py`. The original Physical-Intelligence/openpi repository will NOT work with this checkpoint.

## Model Details

- **Author**: [ANR Robot](https://anrrobot.com)
- **Base Model**: `gs://openpi-assets/checkpoints/pi0_base`
- **Fine-tuning Method**: LoRA (Low-Rank Adaptation)
- **Task**: Cube transfer
- **Training Dataset**: [ANRedlich/trossen_ai_stationary_transfer_multi_cube_03](https://huggingface.co/datasets/ANRedlich/trossen_ai_stationary_transfer_multi_cube_03)
- **Training Steps**: 20,000
- **Framework**: JAX/Flax (OpenPI)

## ⚠️ Important: Model Params AND Norm Stats Must Match

This checkpoint contains both:
- **Model parameters** (`params/`) - the fine-tuned weights
- **Normalization statistics** (`assets/.../norm_stats.json`) - dataset-specific normalization

**You MUST use both from this checkpoint together.** Using mismatched params and norm_stats will cause incorrect behavior or inference failures.

## Installation

**Use the custom fork** (not the original OpenPI):
```bash
git clone --recurse-submodules https://github.com/anredlich/openpi.git
cd openpi
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .
```

## Download
```bash
huggingface-cli download ANRedlich/trossen_ai_stationary_real_pi03 \
    --local-dir ./checkpoints/hf_checkpoint \
    --local-dir-use-symlinks False
```

## Usage with OpenPI

### Step 1: Update `serve_policy.py` EnvMode Checkpoint

In your `serve_policy.py` (or wherever `EnvMode` checkpoints are defined), update the checkpoint directory:
```python
EnvMode.ALOHA_SIM_TROSSEN_AI_FINETUNE: Checkpoint(
    config="pi0_aloha_sim_trossen_ai_mem_finetune_v2",
    # Point to the downloaded HuggingFace checkpoint
    dir="./checkpoints/hf_checkpoint"
),
```

### Step 2: Update TrainConfig with AssetsConfig

In `src/openpi/training/config.py`, you **must** add an `AssetsConfig` to ensure the norm_stats are loaded from this checkpoint (not from a default location):
```python
TrainConfig(
    name="pi0_aloha_sim_trossen_ai_mem_finetune_v2",
    model=pi0_config.Pi0Config(
        paligemma_variant="gemma_2b_lora", 
        action_expert_variant="gemma_300m_lora"
    ),
    data=LeRobotAlohaDataConfig(
        repo_id="ANRedlich/trossen_ai_stationary_transfer_multi_cube_03",
        # ⚠️ CRITICAL: Add this AssetsConfig to load norm_stats from the checkpoint
        assets=AssetsConfig(
            assets_dir="./checkpoints/hf_checkpoint/assets",
            asset_id="ANRedlich/trossen_ai_stationary_transfer_multi_cube_03",
        ),
        default_prompt="Transfer cube",
        use_delta_joint_actions=False,
        adapt_to_pi=False,
        # ⚠️ CRITICAL: This parameter only exists in anredlich/openpi fork
        adapt_trossen_to_pi=True,
        repack_transforms=_transforms.Group(
            inputs=[
                _transforms.RepackTransform(
                    {
                        "images": {
                            "cam_high": "observation.images.cam_high",
                            "cam_low": "observation.images.cam_low",
                            "cam_left_wrist": "observation.images.cam_left_wrist",
                            "cam_right_wrist": "observation.images.cam_right_wrist",
                        },
                        "state": "observation.state",
                        "actions": "action",
                    }
                )
            ]
        ),
    ),       
    weight_loader=weight_loaders.CheckpointWeightLoader(
        "gs://openpi-assets/checkpoints/pi0_base/params"
    ),
    num_train_steps=20_000,
    freeze_filter=pi0_config.Pi0Config(
        paligemma_variant="gemma_2b_lora", 
        action_expert_variant="gemma_300m_lora"
    ).get_freeze_filter(),
    ema_decay=None,
),
```

### Step 3: Run Inference
```bash
cd ~/openpi/openpi
uv run scripts/serve_policy.py --env=ALOHA_SIM_TROSSEN_AI_FINETUNE
```

You should see in the logs:
```
INFO:absl:Restoring checkpoint from ./checkpoints/hf_checkpoint/params.
INFO:root:Loaded norm stats from ./checkpoints/hf_checkpoint/assets/ANRedlich/trossen_ai_stationary_transfer_multi_cube_03
```

**Verify both paths point to the downloaded checkpoint!**

## About `adapt_trossen_to_pi`

The `adapt_trossen_to_pi` parameter in `aloha_policy.py` handles coordinate transformations specific to the Trossen AI robot configuration. This modification is required for correct state/action normalization with Trossen hardware. This parameter is only available in the [anredlich/openpi](https://github.com/anredlich/openpi) fork.

## Checkpoint Structure
```
hf_checkpoint/
├── _CHECKPOINT_METADATA
├── params/                  # Model weights (Orbax format)
│   ├── array_metadatas/
│   ├── d/
│   ├── manifest.ocdbt
│   ├── _METADATA
│   ├── ocdbt.process_0/
│   └── _sharding
└── assets/                  # Normalization statistics
    └── ANRedlich/
        └── trossen_ai_stationary_sim_transfer_40mm_cube_13/
            └── norm_stats.json
```

## Citation

If you use this model, please cite the original OpenPI work:
```bibtex
@article{openpi2024,
    title={π₀: A Vision-Language-Action Flow Model for General Robot Control},
    author={Physical Intelligence},
    year={2024},
    url={https://www.physicalintelligence.company/blog/pi0}
}
```

## Contact

For questions about this model, visit [ANR Robot](https://anrrobot.com) or open an issue on the [GitHub repository](https://github.com/anredlich/openpi).

## License

Apache-2.0 (following the OpenPI repository license)