To train an AI model for your robot, you need a robotics dataset. For that, you need to have first recorded a dataset.

Don’t have a dataset? Use one of ours: PLB/Orange-brick-in-black-box

​
Dataset format and training scripts

The standard dataset format for robotics is the LeRobot v2 format by HuggingFace. This format is compatible with the ACT (Action Chunking Transformer) model and other models. Every dataset recorded with phosphobot is LeRobot v2 compatible and pushed to the Hugging Face dataset hub with your token.

This means that any training script compatible with the LeRobot v2 dataset format can be used to train your AI model.

As of now, we sucessfuly trained the following models with the LeRobot v2 dataset format:

  • ACT (Action Chunking Transformer) trained using LeRobot training scripts
  • gr00t-n1 by NVIDIA
  • pi0 trained using openpi by Physical intelligence

As more models are released, we will update this list with relevant repos and tutorials.

​
Cloud training

To train your model, you can use Replicate. Replicate provides GPU instances to train AI models in one click.

Train ACT on Replicate

The ACT model mimicks one behavior at a time

Train gr00t-n1 on Replicate

The gr00t-N1 model is a promptable model by NVIDIA

​
How to train the ACT (Action Chunking Transformer) model with LeRobot?

  1. On the device where you want to run the training, install the phosphobot python module:

    pip install --upgrade phosphobot

  2. Clone the LeRobot repository and install it:

    git clone https://github.com/huggingface/lerobot.git
cd lerobot
pip install -e .  # Requires Python 3.10+

  3. (Optional) If you want to use Weights & Biases for tracking training metrics, log in with:

    wandb login

  4. Run the training script with the following command in the lerobot repository. Ensure that your lerobot virtual environment is activated.

You need a GPU with at least 16GB of memory to train the ACT model.

sudo python lerobot/scripts/train.py \
  --dataset.repo_id=<HF_USERNAME>/<DATASET_NAME> \
  --policy.type=<act or diffusion or tdmpc or vqbet> \
  --output_dir=outputs/train/phoshobot_test \
  --job_name=phosphobot_test \
  --device=mps \ # mps, cpu or cuda 
  --wandb.enable=true
  • For Apple Silicon (Mac M1, M2) set --device=mps
  • For NVIDIA GPU set --device=cuda
  • For CPU set --device=cpu
  1. Your trained model is saved in lerobot/outputs/train/.

​
How to train the Pi0 (Pi-Zero) model with the SO-100 robot arm?

To train the Pi0 model, you need to use the openpi repository with a few tweaks for the SO-100 robot arm compatibility. We added support for the SO-100 arm in this fork of the openpi repository.

  1. Clone our custom openpi repository and install it:
git clone https://github.com/phospho-app/openpi.git
cd openpi
  1. Install UV if you don’t have it already:
curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env
  1. Setup the environment using uv:
GIT_LFS_SKIP_SMUDGE=1 uv sync
  1. (Optional) If you want to use Weights & Biases for tracking training metrics, log in with:
# Get your auth token at https://wandb.ai/authorize
uv run wandb login
  1. Edit the config file src/openpi/training/config.py to change the dataset to your own.

By default, the 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. Eg: in the PLB/Orange-brick-in-black-box dataset, it’s “Put the orange brick in the black pot”
  1. Start the training:
uv run scripts/compute_norm_stats.py --config-name pi0_so100
# Main train command in nohup, it will take a while
XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py pi0_so100 --exp-name=pi0_so100_blackbox --overwrite

You need a GPU with at least 70GB of memory to train the pi0 model. We recommend using a A100 (80GB) or a H100.

  1. Once training is done, push your model to Hugging Face:
huggingface-cli login  # Follow the authentication steps
huggingface-cli repo create your-model-name --type model

# Push direclty using the cli
huggingface-cli upload USERNAME/your-model-name \
  ~/path/to/folder/ \
  --repo-type model

​
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?