Train Your First Policy

Step-by-step guide to training your first robot learning policy with LeRobot

This tutorial walks you through training your first robot learning policy from scratch using LeRobot. We’ll train a Diffusion Policy on the PushT task, a popular benchmark for imitation learning.

Prerequisites

Make sure you have LeRobot installed:

pip install lerobot

For GPU acceleration (recommended):

pip install lerobot[gpu]

Training Steps

Using the CLI

For production training, use the lerobot-train CLI which includes advanced features:

lerobot-train \
  --policy.type=diffusion \
  --dataset.repo_id=lerobot/pusht \
  --policy.repo_id=your_username/my_first_policy \
  --output_dir=outputs/train/pusht_diffusion \
  --steps=10000 \
  --batch_size=64 \
  --log_freq=100 \
  --save_freq=5000 \
  --eval_freq=2500 \
  --policy.optimizer_lr=1e-4 \
  --policy.device=cuda \
  --num_workers=4

The CLI provides:

• Automatic checkpointing and resumption
• WandB integration for logging
• Distributed training support
• Evaluation during training
• Configuration management

Training with Your Own Data

To train on your own collected data:

Monitoring Training

Using WandB

Integrate with Weights & Biases for detailed logging:

lerobot-train \
  --policy.type=diffusion \
  --dataset.repo_id=lerobot/pusht \
  --use_wandb=true \
  --wandb_entity=your_username \
  --wandb_project=robot_learning \
  --wandb_run_name=pusht_diffusion_v1

Tensorboard

View training logs with Tensorboard:

# Training automatically logs to outputs/train/
tensorboard --logdir outputs/train/

Common Issues

Out of Memory

Reduce batch size or image resolution:

lerobot-train \
  --policy.type=diffusion \
  --dataset.repo_id=lerobot/pusht \
  --batch_size=32 \
  --training.image_size=128  # Down from 256

Slow Training

Increase number of dataloader workers:

lerobot-train \
  --policy.type=diffusion \
  --dataset.repo_id=lerobot/pusht \
  --num_workers=8 \
  --batch_size=64

Loss Not Decreasing

Check learning rate and try gradient clipping:

lerobot-train \
  --policy.type=diffusion \
  --dataset.repo_id=lerobot/pusht \
  --policy.optimizer_lr=5e-5 \
  --grad_clip_norm=10.0

Next Steps

• Evaluate Policies - Test your trained policy
• Multi-GPU Training - Scale up training
• PEFT Training - Fine-tune large models efficiently
• Imitation Learning - Learn about different approaches

Complete Example Script

See the complete training example at:

• examples/training/train_policy.py
• examples/tutorial/diffusion/diffusion_training_example.py