Train a robotics AI model
How to train a robotics AI model with a dataset?
You just recorded a dataset and now you want to train an AI model to control your robot. Here’s how you can do it.
PLB/Orange-brick-in-black-box
You dataset is saved in the LeRobot format and uploaded to your Hugging Face account. Any training script compatible with this dataset format can be used to train your AI model.
For now, we support the following models:
- ACT (Action Chunking Transformer) trained using LeRobot training scripts
- pi0 trained using openpi by Physical intelligence
ACT (Action Chunking Transformer)
On the device where you want to run the training, install the Phosphobot package:
-
Clone the LeRobot repository and install it:
-
(Optional) If you want to use Weights & Biases for tracking training metrics, log in with:
-
Run the training script with the following command in the
lerobot
repository (Set —device=mps for Apple Silicon (Mac M1/M2),cuda
if you have an NVIDIA GPU orcpu
if you don’t have a GPU). Ensure that yourlerobot
virtual environment is activated. -
Your trained model is saved in
lerobot/outputs/train/
.
Pi0 (Pi-Zero)
We added support for the SO-100 arm in this fork of the openpi repository. For a full fine-tuning of the model, you will need at least 70GB of GPU memory (we recommend using a A100 (80GB) / H100).
Install UV if you don’t have it already:
Setup the environment using uv:
(Optional) If you want to use Weights & Biases for tracking training metrics, log in with:
By default, our training config uses the PLB/Orange-brick-in-black-box
SO-100 dataset.
Make sure you dataset has a correct task description for each task in the dataset (in our dataset, it’s “Put the orange brick in the black pot”).
Edit the config file src/openpi/training/config.py
to change the dataset to your own.
Then, you can train your model using the following command:
Once training is done, push your model to Hugging Face:
Next steps
Test the model you just trained on your robot. See the Use AI models page for more information.
Use an AI model
Let a trained AI model control your robot
Was this page helpful?