Training a model on the LeRobot dataset

Background background

LeRobot is a project by huggingface that aims to provide models, datasets and tools for real-world robotics in PyTorch. This example shows how one can train a model on the pusht-dataset and visualize it's progress using rerun.

Run the code run-the-code

This is an external example, check the repository for more information.

To train the model as shown in the video, install git-lfs and clone the repository and then run the following code:

pip install -e '.[pusht]'
WANDB_MODE=offline python lerobot/scripts/train.py \
  hydra.run.dir=outputs/train/diffusion_pusht \
  hydra.job.name=diffusion_pusht \
  policy=diffusion \
  env=pusht \
  env.task=PushT-v0 \
  dataset_repo_id=lerobot/pusht \
  training.offline_steps=20000 \
  training.save_freq=5000 ++training.log_freq=50 \
  training.eval_freq=1500 \
  eval.n_episodes=50 \
  wandb.enable=true \
  wandb.disable_artifact=true \
  device=cuda

If you don't have CUDA installed you will have to change the last argument device=cuda to device=cpu or another device.