State-of-the-art Machine Learning for real-world robotics

---

🤗 LeRobot aims to provide models, datasets, and tools for real-world robotics in PyTorch. The goal is to lower the barrier for entry to robotics so that everyone can contribute and benefit from sharing datasets and pretrained models.

🤗 LeRobot contains state-of-the-art approaches that have been shown to transfer to the real-world with a focus on imitation learning and reinforcement learning.

🤗 LeRobot already provides a set of pretrained models, datasets with human collected demonstrations, and simulated environments so that everyone can get started. In the coming weeks, the plan is to add more and more support for real-world robotics on the most affordable and capable robots out there.

🤗 LeRobot hosts pretrained models and datasets on this HuggingFace community page: huggingface.co/lerobot

Examples of pretrained models and environments

ACT policy on ALOHA env	TDMPC policy on SimXArm env	Diffusion policy on PushT env

Acknowledgment

• ACT policy and ALOHA environment are adapted from ALOHA
• Diffusion policy and Pusht environment are adapted from Diffusion Policy
• TDMPC policy and Simxarm environment are adapted from FOWM
• Abstractions and utilities for Reinforcement Learning come from TorchRL

Installation

Download our source code:

git clone https://github.com/huggingface/lerobot.git && cd lerobot

Create a virtual environment with Python 3.10 and activate it, e.g. with miniconda:

conda create -y -n lerobot python=3.10 && conda activate lerobot

Install 🤗 LeRobot:

python -m pip install .

For simulations, 🤗 LeRobot comes with gymnasium environments that can be installed as extras:

• aloha
• xarm
• pusht

For instance, to install 🤗 LeRobot with aloha and pusht, use:

python -m pip install ".[aloha, pusht]"

To use Weights and Biases for experiments tracking, log in with

wandb login

Walkthrough

.
├── lerobot
|   ├── configs          # contains hydra yaml files with all options that you can override in the command line
|   |   ├── default.yaml   # selected by default, it loads pusht environment and diffusion policy
|   |   ├── env            # various sim environments and their datasets: aloha.yaml, pusht.yaml, xarm.yaml
|   |   └── policy         # various policies: act.yaml, diffusion.yaml, tdmpc.yaml
|   ├── common           # contains classes and utilities
|   |   ├── datasets       # various datasets of human demonstrations: aloha, pusht, xarm
|   |   ├── envs           # various sim environments: aloha, pusht, xarm
|   |   └── policies       # various policies: act, diffusion, tdmpc
|   └── scripts                  # contains functions to execute via command line
|       ├── visualize_dataset.py  # load a dataset and render its demonstrations
|       ├── eval.py               # load policy and evaluate it on an environment
|       └── train.py              # train a policy via imitation learning and/or reinforcement learning
├── outputs               # contains results of scripts execution: logs, videos, model checkpoints
├── .github
|   └── workflows
|       └── test.yml      # defines install settings for continuous integration and specifies end-to-end tests
└── tests                 # contains pytest utilities for continuous integration

Visualize datasets

You can import our dataset class, download the data from the HuggingFace hub and use our rendering utilities:

""" Copy pasted from `examples/1_visualize_dataset.py` """
import os
from pathlib import Path

import lerobot
from lerobot.common.datasets.aloha import AlohaDataset
from lerobot.scripts.visualize_dataset import render_dataset

print(lerobot.available_datasets)
# >>> ['aloha_sim_insertion_human', 'aloha_sim_insertion_scripted', 'aloha_sim_transfer_cube_human', 'aloha_sim_transfer_cube_scripted', 'pusht', 'xarm_lift_medium']

# TODO(rcadene): remove DATA_DIR
dataset = AlohaDataset("pusht", root=Path(os.environ.get("DATA_DIR")))

video_paths = render_dataset(
    dataset,
    out_dir="outputs/visualize_dataset/example",
    max_num_episodes=1,
)
print(video_paths)
# ['outputs/visualize_dataset/example/episode_0.mp4']

Or you can achieve the same result by executing our script from the command line:

python lerobot/scripts/visualize_dataset.py \
env=pusht \
hydra.run.dir=outputs/visualize_dataset/example
# >>> ['outputs/visualize_dataset/example/episode_0.mp4']

Evaluate a pretrained policy

Check out example 2 to see how you can load a pretrained policy from HuggingFace hub, load up the corresponding environment and model, and run an evaluation.

Or you can achieve the same result by executing our script from the command line:

python lerobot/scripts/eval.py \
--hub-id lerobot/diffusion_policy_pusht_image \
eval_episodes=10 \
hydra.run.dir=outputs/eval/example_hub

After training your own policy, you can also re-evaluate the checkpoints with:

python lerobot/scripts/eval.py \
--config PATH/TO/FOLDER/config.yaml \
policy.pretrained_model_path=PATH/TO/FOLDER/weights.pth \
eval_episodes=10 \
hydra.run.dir=outputs/eval/example_dir

See python lerobot/scripts/eval.py --help for more instructions.

Train your own policy

You can import our dataset, environment, policy classes, and use our training utilities (if some data is missing, it will be automatically downloaded from HuggingFace hub): check out example 3. After you run this, you may want to revisit example 2 to evaluate your training output!

In general, you can use our training script to easily train any policy on any environment:

python lerobot/scripts/train.py \
env=aloha \
task=sim_insertion \
dataset_id=aloha_sim_insertion_scripted \
policy=act \
hydra.run.dir=outputs/train/aloha_act

Contribute

If you would like to contribute to 🤗 LeRobot, please check out our contribution guide.

Add a new dataset

To add a dataset to the hub, first login and use a token generated from huggingface settings with write access:

huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential

Then you can upload it to the hub with:

HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli upload $HF_USER/$DATASET data/$DATASET \
--repo-type dataset  \
--revision v1.0

You will need to set the corresponding version as a default argument in your dataset class:

  version: str | None = "v1.0",

See: lerobot/common/datasets/pusht.py

For instance, for lerobot/pusht, we used:

HF_USER=lerobot
DATASET=pusht

If you want to improve an existing dataset, you can download it locally with:

mkdir -p data/$DATASET
HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download ${HF_USER}/$DATASET \
--repo-type dataset \
--local-dir data/$DATASET \
--local-dir-use-symlinks=False \
--revision v1.0

Iterate on your code and dataset with:

DATA_DIR=data python train.py

Upload a new version (v2.0 or v1.1 if the changes are respectively more or less significant):

HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli upload $HF_USER/$DATASET data/$DATASET \
--repo-type dataset \
--revision v1.1 \
--delete "*"

Then you will need to set the corresponding version as a default argument in your dataset class:

  version: str | None = "v1.1",

See: lerobot/common/datasets/pusht.py

Finally, you might want to mock the dataset if you need to update the unit tests as well:

python tests/scripts/mock_dataset.py --in-data-dir data/$DATASET --out-data-dir tests/data/$DATASET

Add a pretrained policy

Once you have trained a policy you may upload it to the HuggingFace hub.

Firstly, make sure you have a model repository set up on the hub. The hub ID looks like HF_USER/REPO_NAME.

Secondly, assuming you have trained a policy, you need:

• config.yaml which you can get from the .hydra directory of your training output folder.
• model.pt which should be one of the saved models in the models directory of your training output folder (they won't be named model.pt but you will need to choose one).
• stats.pth which should point to the same file in the dataset directory (found in data/{dataset_name}).

To upload these to the hub, prepare a folder with the following structure (you can use symlinks rather than copying):

to_upload
    ├── config.yaml
    ├── model.pt
    └── stats.pth

With the folder prepared, run the following with a desired revision ID.

huggingface-cli upload $HUB_ID to_upload --revision $REVISION_ID

If you want this to be the default revision also run the following (don't worry, it won't upload the files again; it will just adjust the file pointers):

huggingface-cli upload $HUB_ID to_upload

See eval.py for an example of how a user may use your policy.

Improve your code with profiling

An example of a code snippet to profile the evaluation of a policy:

from torch.profiler import profile, record_function, ProfilerActivity

def trace_handler(prof):
    prof.export_chrome_trace(f"tmp/trace_schedule_{prof.step_num}.json")

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(
        wait=2,
        warmup=2,
        active=3,
    ),
    on_trace_ready=trace_handler
) as prof:
    with record_function("eval_policy"):
        for i in range(num_episodes):
            prof.step()
            # insert code to profile, potentially whole body of eval_policy function

python lerobot/scripts/eval.py \
--config outputs/pusht/.hydra/config.yaml \
pretrained_model_path=outputs/pusht/model.pt \
eval_episodes=7