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.

Don’t have a dataset? Use one of ours: 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:

pip install --upgrade phosphobot
  1. Clone the LeRobot repository and install it:

    git clone https://github.com/huggingface/lerobot.git
    cd lerobot
    pip install -e .  # Requires Python 3.10+
    
  2. (Optional) If you want to use Weights & Biases for tracking training metrics, log in with:

    wandb login
    
  3. 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 or cpu if you don’t have a GPU). Ensure that your lerobot virtual environment is activated.

    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 \
      --wandb.enable=true
    
  4. 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).

git clone https://github.com/phospho-app/openpi.git
cd openpi

Install UV if you don’t have it already:

curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env

Setup the environment using uv:

GIT_LFS_SKIP_SMUDGE=1 uv sync

(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

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:

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

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