Step 7: Training — pi0.5 Command

Before running

• Environment: first open an instance and upload the dataset as described in Cloud GPU training environment setup, then come back to this page and carry out the two sections "Install the environment" and "Command line"
• Dataset: --dataset.root=~/lerobot_my_dataset_shake_hands in the command points to the handshake dataset collected in step 6. If you are training your own task, replace it with your own dataset name
• Output directory: if --output_dir already exists, first delete it with the sudo rm -rf line above, or use a new name
• You can check the curves on wandb at any time during training, see View real-time training curves on wandb

Reference documentation

https://github.com/huggingface/lerobot/blob/46e19ae579f80ce66211afafd1c3c649c569131f/docs/source/pi0.mdx

https://www.pi.website/blog/pi05

Recommended cloud GPU instance

Install the environment

cd lerobot
pip install -e ".[pi0]"

Command line

• Delete the files under output from a previously interrupted training run

sudo rm -rf output_lerobot_train/shake/pi05_A

• Train

lerobot-train \
    --dataset.repo_id=<username>/lerobot_my_dataset_shake_hands \
    --dataset.root=~/lerobot_my_dataset_shake_hands \
    --dataset.revision=v0.1.0 \
    --policy.type=pi05 \
    --output_dir=~/output_lerobot_train/shake/pi05_A \
    --job_name=shake_pi05_A \
    --policy.pretrained_path=lerobot/pi05_base \
    --policy.compile_model=true \
    --policy.gradient_checkpointing=true \
    --policy.dtype=bfloat16 \
    --policy.freeze_vision_encoder=false \
    --policy.train_expert_only=false \
    --steps=50000 \
    --policy.device=cuda \
    --policy.push_to_hub=false \
    --wandb.enable=true \
    --wandb.project=Lerobot_my_Project \
    --batch_size=8

Training does not formally begin until about 20 minutes after the command line is started

The model archive is about 5 GB, and 7 GB after extraction

Related Products

SO-ARM101 Developer Kit→