pi0.5: Optimized for AMD ROCm
pi0.5 (PI05Policy, Physical Intelligence) is a vision-language-action policy from Hugging Face LeRobot for 6-DOF robot arm control. This repository packages inference for robot arm action prediction using PyTorch, exported and validated for AMD ROCm so it runs efficiently on AMD GPUs.
This is based on the implementation of pi0.5 found here. This repository contains configurations and scripts optimized for AMD® ROCm™ platforms. You can use the pi0.5 AMD scripts to reproduce results or export with custom configurations. More details on model performance can be found here.
Task Overview
Task: Robot arm action prediction (vision-language-action)
Dataset: BlankHead/so101_redcube_greencloth_3cams (LeRobot format)
Output metrics: MAE, RMSE (per-joint and per-episode)
Model variants: No
MODEL_SIZEvariants — the pipeline auto-detects PEFT/LoRA checkpoints viaadapter_config.jsonand merges the adapter, patchingaction_dimto the checkpoint's value.
PyTorch note: GPU only (ROCm — BF16); weights are streamed directly from disk into GPU VRAM in bf16, bypassing any CPU copy. A GPU with ≥32 GB VRAM is required — no CPU or NPU (VitisAI) path is available, since loading the ~14 GB fp32 weights on CPU would require >48 GB of system RAM. Requires Hugging Face auth (
HF_TOKEN) since the tokenizer pulls the gatedgoogle/paligemma-3b-pt-224repo.
AMD ROCm Optimization
This model export has been adapted and validated for AMD Instinct™ / Radeon™ GPUs running ROCm. Key points:
- Validated backend: PyTorch (native ROCm HIP kernels), BF16 precision, GPU only.
- No code changes required versus the upstream pi0.5 implementation — only environment/runtime configuration differs.
- No CPU or NPU fallback path is available for this model.
| Runtime | Precision | Backend | Hardware | Notes |
|---|---|---|---|---|
| PyTorch | BF16 | HIP (ROCm) | AMD Instinct™ / Radeon™ GPU (≥32 GB VRAM) | Direct-to-GPU weight streaming; no CPU/NPU path |
Getting Started
For setup instructions, evaluation scripts, and custom configuration options, see the pi0.5 on GitHub.
Model Details
Model Type: Vision-language-action policy for robot arm control
Base Model: lerobot/pi05_base (PI05Policy)
Model Stats:
- Precision tested: BF16 (GPU only)
- VRAM requirement: ≥32 GB (bf16 weight streaming, ~14 GB on-disk fp32 checkpoint)
- Configurable runtime knobs:
rtc_config.enabled(Real-Time Chunking),num_steps(flow-matching denoising passes per chunk),n_action_steps(actions consumed per chunk)
Accuracy Pipeline
Open-loop offline evaluation is fully implemented: make evaluate-gpu runs inference on recorded dataset episodes and computes per-joint and per-episode MAE / RMSE against the recorded ground-truth actions. Lower is better for both metrics.
Metrics Explained
| Metric | Description |
|---|---|
| MAE | Mean Absolute Error — average absolute difference between predicted and ground-truth joint positions across all timesteps. Lower is better. |
| RMSE | Root Mean Squared Error — penalizes large deviations more heavily than MAE. Lower is better. |
Accuracy Results
Published Results — filled from runs/eval/<dataset_tag>/loss.json:
| Dataset | Mode | Average MAE | Average RMSE |
|---|---|---|---|
| BlankHead/so101_redcube_greencloth_3cams (1 episode) | select_action | 19.5420 | 25.2683 |
| HarikrishnaVydana/Lerobot-redcube-wite-bg-pickdrop-v3 (1 episode) | chunked_rtc | 19.2629 | 32.3233 |
Dig Deeper
Want to explore the full evaluation scripts, config options, and other AMD-optimized model examples?
📂 View the full project on GitHub
The GitHub repository includes:
- Setup and prerequisites for ROCm environments
- Direct-to-GPU bf16 weight streaming and PEFT/LoRA adapter merging
- Open-loop dataset evaluation with trajectory plots and comparison videos
- Latency benchmarking with Chrome trace output