admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 10:16:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into user/michel-aractingi/2025_06_30_dataset_v3

Browse Source
This commit is contained in:
Michel Aractingi
2025-08-13 01:39:32 +02:00
committed by GitHub
parent 4048b02d4a 0878c6880f
commit 267a753eda
68 changed files with 7608 additions and 335 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/PULL_REQUEST_TEMPLATE.md
+1 -1
View File
@@ -30,7 +30,7 @@ pytest -sx tests/test_stuff.py::test_something
``` ```
```bash ```bash
python -m lerobot.scripts.train --some.option=true lerobot-train --some.option=true
``` ```
## SECTION TO REMOVE BEFORE SUBMITTING YOUR PR ## SECTION TO REMOVE BEFORE SUBMITTING YOUR PR
.github/workflows/nightly.yml
+2 -2
View File
@@ -29,8 +29,8 @@ on:
env: env:
UV_VERSION: "0.8.0" UV_VERSION: "0.8.0"
PYTHON_VERSION: "3.10" PYTHON_VERSION: "3.10"
DOCKER_IMAGE_NAME_CPU: huggingface/lerobot-gpu:latest DOCKER_IMAGE_NAME_CPU: huggingface/lerobot-cpu:latest
DOCKER_IMAGE_NAME_GPU: huggingface/lerobot-cpu:latest DOCKER_IMAGE_NAME_GPU: huggingface/lerobot-gpu:latest
# Ensures that only the latest commit is built, canceling older runs. # Ensures that only the latest commit is built, canceling older runs.
concurrency: concurrency:
.github/workflows/release.yml
+46 -8
View File
@@ -19,6 +19,11 @@ on:
tags: tags:
- 'v*.*.*' # Trigger on tags like v0.1.0, v1.0.0 - 'v*.*.*' # Trigger on tags like v0.1.0, v1.0.0
# Sets up the environment variables
env:
UV_VERSION: "0.8.0"
PYTHON_VERSION: "3.10"
jobs: jobs:
# This job builds the Python package and publishes it to PyPI # This job builds the Python package and publishes it to PyPI
build-and-publish: build-and-publish:
@@ -50,6 +55,7 @@ jobs:
VERSION_NUMBER=${VERSION#v} VERSION_NUMBER=${VERSION#v}
echo "tag_version=$VERSION_NUMBER" >> $GITHUB_OUTPUT echo "tag_version=$VERSION_NUMBER" >> $GITHUB_OUTPUT
- name: Check if version matches pyproject.toml - name: Check if version matches pyproject.toml
if: startsWith(github.ref, 'refs/tags/v') && !contains(github.ref, '-')
# zizmor: ignore[template-injection] # zizmor: ignore[template-injection]
run: | run: |
TAG_VERSION=${{ steps.extract_info.outputs.tag_version }} TAG_VERSION=${{ steps.extract_info.outputs.tag_version }}
@@ -86,13 +92,29 @@ jobs:
env: env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# zizmor: ignore[template-injection] # zizmor: ignore[template-injection]
run: gh release create ${{ github.ref_name }} --release-name "Release ${{ github.ref_name }}" --generate-notes ./dist/* run: |
gh release create ${{ github.ref_name }} \
--title "Release ${{ github.ref_name }}" \
--generate-notes \
--draft=$([[ "${{ github.ref_name }}" == *-* ]] && echo true || echo false) \
--prerelease=$([[ "${{ github.ref_name }}" == *-* ]] && echo true || echo false) \
./dist/*
- name: Publish to PyPI - name: Publish to TestPyPI for pre-releases
if: startsWith(github.ref, 'refs/tags/v') # True for tags like 'v0.2.0-rc1'
if: startsWith(github.ref, 'refs/tags/v') && contains(github.ref, '-')
uses: pypa/gh-action-pypi-publish@v1.12.4 # zizmor: ignore[unpinned-uses, use-trusted-publishing] uses: pypa/gh-action-pypi-publish@v1.12.4 # zizmor: ignore[unpinned-uses, use-trusted-publishing]
with: with:
password: ${{ secrets.PYPI_API_TOKEN }} repository-url: https://test.pypi.org/legacy/
verbose: true
print-hash: true
- name: Publish to PyPI
if: startsWith(github.ref, 'refs/tags/v') && !contains(github.ref, '-')
uses: pypa/gh-action-pypi-publish@v1.12.4 # zizmor: ignore[unpinned-uses, use-trusted-publishing]
with:
verbose: true
print-hash: true
# This job runs end-to-end tests on the release # This job runs end-to-end tests on the release
test-release: test-release:
@@ -119,15 +141,31 @@ jobs:
enable-cache: true enable-cache: true
version: ${{ env.UV_VERSION }} version: ${{ env.UV_VERSION }}
python-version: ${{ env.PYTHON_VERSION }} python-version: ${{ env.PYTHON_VERSION }}
- name: Create uv virtual environment
run: uv venv
- name: Install lerobot release - name: Install lerobot release
run: uv run pip install lerobot==${{ needs.build-and-publish.outputs.version }} # zizmor: ignore[template-injection] # zizmor: ignore[template-injection]
run: |
VERSION="${{ needs.build-and-publish.outputs.version }}"
if [[ "$VERSION" == *-* ]]; then
BASE_VERSION="${VERSION%%-*}"
echo "Installing pre-release version $BASE_VERSION from TestPyPI..."
uv pip install \
--index-url https://test.pypi.org/simple/ \
--extra-index-url https://pypi.org/simple \
--index-strategy unsafe-best-match \
"lerobot[all]==$BASE_VERSION"
else
echo "Installing release version $VERSION from PyPI..."
uv pip install "lerobot[all]==$VERSION"
fi
- name: Check lerobot version - name: Check lerobot version
run: uv run lerobot --version run: uv run python -c "import lerobot; print(lerobot.__version__)"
- name: Run end-to-end tests - name: Run end-to-end tests
run: uv run make test-end-to-end run: uv run make test-end-to-end
# TODO(Steven): Publish draft/pre-release and to test pypi # TODO(Steven): Publish draft/pre-release and to test pypi weekly
# TODO(Steven): Separate build and publish job
# TODO(Steven): Tag documentation with the same version as the package # TODO(Steven): Tag documentation with the same version as the package
Makefile
+9 -9
View File
@@ -44,7 +44,7 @@ test-end-to-end:
${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-eval ${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-eval
test-act-ete-train: test-act-ete-train:
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=act \ --policy.type=act \
--policy.dim_model=64 \ --policy.dim_model=64 \
--policy.n_action_steps=20 \ --policy.n_action_steps=20 \
@@ -68,12 +68,12 @@ test-act-ete-train:
--output_dir=tests/outputs/act/ --output_dir=tests/outputs/act/
test-act-ete-train-resume: test-act-ete-train-resume:
python -m lerobot.scripts.train \ lerobot-train \
--config_path=tests/outputs/act/checkpoints/000002/pretrained_model/train_config.json \ --config_path=tests/outputs/act/checkpoints/000002/pretrained_model/train_config.json \
--resume=true --resume=true
test-act-ete-eval: test-act-ete-eval:
python -m lerobot.scripts.eval \ lerobot-eval \
--policy.path=tests/outputs/act/checkpoints/000004/pretrained_model \ --policy.path=tests/outputs/act/checkpoints/000004/pretrained_model \
--policy.device=$(DEVICE) \ --policy.device=$(DEVICE) \
--env.type=aloha \ --env.type=aloha \
@@ -82,7 +82,7 @@ test-act-ete-eval:
--eval.batch_size=1 --eval.batch_size=1
test-diffusion-ete-train: test-diffusion-ete-train:
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=diffusion \ --policy.type=diffusion \
--policy.down_dims='[64,128,256]' \ --policy.down_dims='[64,128,256]' \
--policy.diffusion_step_embed_dim=32 \ --policy.diffusion_step_embed_dim=32 \
@@ -106,7 +106,7 @@ test-diffusion-ete-train:
--output_dir=tests/outputs/diffusion/ --output_dir=tests/outputs/diffusion/
test-diffusion-ete-eval: test-diffusion-ete-eval:
python -m lerobot.scripts.eval \ lerobot-eval \
--policy.path=tests/outputs/diffusion/checkpoints/000002/pretrained_model \ --policy.path=tests/outputs/diffusion/checkpoints/000002/pretrained_model \
--policy.device=$(DEVICE) \ --policy.device=$(DEVICE) \
--env.type=pusht \ --env.type=pusht \
@@ -115,7 +115,7 @@ test-diffusion-ete-eval:
--eval.batch_size=1 --eval.batch_size=1
test-tdmpc-ete-train: test-tdmpc-ete-train:
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=tdmpc \ --policy.type=tdmpc \
--policy.device=$(DEVICE) \ --policy.device=$(DEVICE) \
--policy.push_to_hub=false \ --policy.push_to_hub=false \
@@ -137,7 +137,7 @@ test-tdmpc-ete-train:
--output_dir=tests/outputs/tdmpc/ --output_dir=tests/outputs/tdmpc/
test-tdmpc-ete-eval: test-tdmpc-ete-eval:
python -m lerobot.scripts.eval \ lerobot-eval \
--policy.path=tests/outputs/tdmpc/checkpoints/000002/pretrained_model \ --policy.path=tests/outputs/tdmpc/checkpoints/000002/pretrained_model \
--policy.device=$(DEVICE) \ --policy.device=$(DEVICE) \
--env.type=xarm \ --env.type=xarm \
@@ -148,7 +148,7 @@ test-tdmpc-ete-eval:
test-smolvla-ete-train: test-smolvla-ete-train:
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=smolvla \ --policy.type=smolvla \
--policy.n_action_steps=20 \ --policy.n_action_steps=20 \
--policy.chunk_size=20 \ --policy.chunk_size=20 \
@@ -171,7 +171,7 @@ test-smolvla-ete-train:
--output_dir=tests/outputs/smolvla/ --output_dir=tests/outputs/smolvla/
test-smolvla-ete-eval: test-smolvla-ete-eval:
python -m lerobot.scripts.eval \ lerobot-eval \
--policy.path=tests/outputs/smolvla/checkpoints/000004/pretrained_model \ --policy.path=tests/outputs/smolvla/checkpoints/000004/pretrained_model \
--policy.device=$(DEVICE) \ --policy.device=$(DEVICE) \
--env.type=aloha \ --env.type=aloha \
README.md
+79 -172
View File
@@ -1,25 +1,21 @@
<p align="center"> <p align="center">
<picture> <img alt="LeRobot, Hugging Face Robotics Library" src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/lerobot-logo-thumbnail.png" width="100%">
<source media="(prefers-color-scheme: dark)" srcset="media/lerobot-logo-thumbnail.png">
<source media="(prefers-color-scheme: light)" srcset="media/lerobot-logo-thumbnail.png">
<img alt="LeRobot, Hugging Face Robotics Library" src="media/lerobot-logo-thumbnail.png" style="max-width: 100%;">
</picture>
<br/> <br/>
<br/> <br/>
</p> </p>
<div align="center"> <div align="center">
[![Tests](https://github.com/huggingface/lerobot/actions/workflows/nightly-tests.yml/badge.svg?branch=main)](https://github.com/huggingface/lerobot/actions/workflows/nightly-tests.yml?query=branch%3Amain) [![Tests](https://github.com/huggingface/lerobot/actions/workflows/nightly.yml/badge.svg?branch=main)](https://github.com/huggingface/lerobot/actions/workflows/nightly.yml?query=branch%3Amain)
[![Coverage](https://codecov.io/gh/huggingface/lerobot/branch/main/graph/badge.svg?token=TODO)](https://codecov.io/gh/huggingface/lerobot)
[![Python versions](https://img.shields.io/pypi/pyversions/lerobot)](https://www.python.org/downloads/) [![Python versions](https://img.shields.io/pypi/pyversions/lerobot)](https://www.python.org/downloads/)
[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/huggingface/lerobot/blob/main/LICENSE) [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/huggingface/lerobot/blob/main/LICENSE)
[![Status](https://img.shields.io/pypi/status/lerobot)](https://pypi.org/project/lerobot/) [![Status](https://img.shields.io/pypi/status/lerobot)](https://pypi.org/project/lerobot/)
[![Version](https://img.shields.io/pypi/v/lerobot)](https://pypi.org/project/lerobot/) [![Version](https://img.shields.io/pypi/v/lerobot)](https://pypi.org/project/lerobot/)
[![Examples](https://img.shields.io/badge/Examples-green.svg)](https://github.com/huggingface/lerobot/tree/main/examples) [![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-v2.1-ff69b4.svg)](https://github.com/huggingface/lerobot/blob/main/CODE_OF_CONDUCT.md)
[![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-v2.1%20adopted-ff69b4.svg)](https://github.com/huggingface/lerobot/blob/main/CODE_OF_CONDUCT.md)
[![Discord](https://dcbadge.vercel.app/api/server/C5P34WJ68S?style=flat)](https://discord.gg/s3KuuzsPFb) [![Discord](https://dcbadge.vercel.app/api/server/C5P34WJ68S?style=flat)](https://discord.gg/s3KuuzsPFb)
<!-- [![Coverage](https://codecov.io/gh/huggingface/lerobot/branch/main/graph/badge.svg?token=TODO)](https://codecov.io/gh/huggingface/lerobot) -->
</div> </div>
<h2 align="center"> <h2 align="center">
@@ -29,10 +25,10 @@
<div align="center"> <div align="center">
<img <img
src="media/hope_jr/hopejr.png?raw=true" src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/hope_jr/hopejr.png"
alt="HopeJR robot" alt="HopeJR robot"
title="HopeJR robot" title="HopeJR robot"
style="width: 60%;" width="60%"
/> />
<p><strong>Meet HopeJR – A humanoid robot arm and hand for dexterous manipulation!</strong></p> <p><strong>Meet HopeJR – A humanoid robot arm and hand for dexterous manipulation!</strong></p>
@@ -51,20 +47,12 @@
</h2> </h2>
<div align="center"> <div align="center">
<div style="display: flex; gap: 1rem; justify-content: center; align-items: center;" > <table>
<img <tr>
src="media/so101/so101.webp?raw=true" <td align="center"><img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/so101/so101.webp" alt="SO-101 follower arm" title="SO-101 follower arm" width="90%"/></td>
alt="SO-101 follower arm" <td align="center"><img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/so101/so101-leader.webp" alt="SO-101 leader arm" title="SO-101 leader arm" width="90%"/></td>
title="SO-101 follower arm" </tr>
style="width: 40%;" </table>
/>
<img
src="media/so101/so101-leader.webp?raw=true"
alt="SO-101 leader arm"
title="SO-101 leader arm"
style="width: 40%;"
/>
</div>
<p><strong>Meet the updated SO100, the SO-101 – Just €114 per arm!</strong></p> <p><strong>Meet the updated SO100, the SO-101 – Just €114 per arm!</strong></p>
<p>Train it in minutes with a few simple moves on your laptop.</p> <p>Train it in minutes with a few simple moves on your laptop.</p>
@@ -76,7 +64,7 @@
<p>Want to take it to the next level? Make your SO-101 mobile by building LeKiwi!</p> <p>Want to take it to the next level? Make your SO-101 mobile by building LeKiwi!</p>
<p>Check out the <a href="https://huggingface.co/docs/lerobot/lekiwi">LeKiwi tutorial</a> and bring your robot to life on wheels.</p> <p>Check out the <a href="https://huggingface.co/docs/lerobot/lekiwi">LeKiwi tutorial</a> and bring your robot to life on wheels.</p>
<img src="media/lekiwi/kiwi.webp?raw=true" alt="LeKiwi mobile robot" title="LeKiwi mobile robot" width="50%"> <img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/lekiwi/kiwi.webp" alt="LeKiwi mobile robot" title="LeKiwi mobile robot" width="50%">
</div> </div>
<br/> <br/>
@@ -99,9 +87,9 @@
<table> <table>
<tr> <tr>
<td><img src="media/gym/aloha_act.gif" width="100%" alt="ACT policy on ALOHA env"/></td> <td><img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/gym/aloha_act.gif" width="100%" alt="ACT policy on ALOHA env"/></td>
<td><img src="media/gym/simxarm_tdmpc.gif" width="100%" alt="TDMPC policy on SimXArm env"/></td> <td><img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/gym/simxarm_tdmpc.gif" width="100%" alt="TDMPC policy on SimXArm env"/></td>
<td><img src="media/gym/pusht_diffusion.gif" width="100%" alt="Diffusion policy on PushT env"/></td> <td><img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/gym/pusht_diffusion.gif" width="100%" alt="Diffusion policy on PushT env"/></td>
</tr> </tr>
<tr> <tr>
<td align="center">ACT policy on ALOHA env</td> <td align="center">ACT policy on ALOHA env</td>
@@ -110,23 +98,11 @@
</tr> </tr>
</table> </table>
### Acknowledgment
- The LeRobot team 🤗 for building SmolVLA [Paper](https://arxiv.org/abs/2506.01844), [Blog](https://huggingface.co/blog/smolvla).
- Thanks to Tony Zhao, Zipeng Fu and colleagues for open sourcing ACT policy, ALOHA environments and datasets. Ours are adapted from [ALOHA](https://tonyzhaozh.github.io/aloha) and [Mobile ALOHA](https://mobile-aloha.github.io).
- Thanks to Cheng Chi, Zhenjia Xu and colleagues for open sourcing Diffusion policy, Pusht environment and datasets, as well as UMI datasets. Ours are adapted from [Diffusion Policy](https://diffusion-policy.cs.columbia.edu) and [UMI Gripper](https://umi-gripper.github.io).
- Thanks to Nicklas Hansen, Yunhai Feng and colleagues for open sourcing TDMPC policy, Simxarm environments and datasets. Ours are adapted from [TDMPC](https://github.com/nicklashansen/tdmpc) and [FOWM](https://www.yunhaifeng.com/FOWM).
- Thanks to Antonio Loquercio and Ashish Kumar for their early support.
- Thanks to [Seungjae (Jay) Lee](https://sjlee.cc/), [Mahi Shafiullah](https://mahis.life/) and colleagues for open sourcing [VQ-BeT](https://sjlee.cc/vq-bet/) policy and helping us adapt the codebase to our repository. The policy is adapted from [VQ-BeT repo](https://github.com/jayLEE0301/vq_bet_official).
## Installation ## Installation
Download our source code: LeRobot works with Python 3.10+ and PyTorch 2.2+.
```bash ### Environment Setup
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`](https://docs.anaconda.com/free/miniconda/index.html): Create a virtual environment with Python 3.10 and activate it, e.g. with [`miniconda`](https://docs.anaconda.com/free/miniconda/index.html):
@@ -151,7 +127,18 @@ conda install ffmpeg -c conda-forge
> >
> - _[On Linux only]_ Install [ffmpeg build dependencies](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#GettheDependencies) and [compile ffmpeg from source with libsvtav1](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#libsvtav1), and make sure you use the corresponding ffmpeg binary to your install with `which ffmpeg`. > - _[On Linux only]_ Install [ffmpeg build dependencies](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#GettheDependencies) and [compile ffmpeg from source with libsvtav1](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#libsvtav1), and make sure you use the corresponding ffmpeg binary to your install with `which ffmpeg`.
Install 🤗 LeRobot: ### Install LeRobot 🤗
#### From Source
First, clone the repository and navigate into the directory:
```bash
git clone https://github.com/huggingface/lerobot.git
cd lerobot
```
Then, install the library in editable mode. This is useful if you plan to contribute to the code.
```bash ```bash
pip install -e . pip install -e .
@@ -172,6 +159,34 @@ For instance, to install 🤗 LeRobot with aloha and pusht, use:
pip install -e ".[aloha, pusht]" pip install -e ".[aloha, pusht]"
``` ```
### Installation from PyPI
**Core Library:**
Install the base package with:
```bash
pip install lerobot
```
_This installs only the default dependencies._
**Extra Features:**
To install additional functionality, use one of the following:
```bash
pip install 'lerobot[all]' # All available features
pip install 'lerobot[aloha,pusht]' # Specific features (Aloha & Pusht)
pip install 'lerobot[feetech]' # Feetech motor support
```
_Replace `[...]` with your desired features._
**Available Tags:**
For a full list of optional dependencies, see:
https://pypi.org/project/lerobot/
### Weights & Biases
To use [Weights and Biases](https://docs.wandb.ai/quickstart) for experiment tracking, log in with To use [Weights and Biases](https://docs.wandb.ai/quickstart) for experiment tracking, log in with
```bash ```bash
@@ -182,7 +197,7 @@ wandb login
### Visualize datasets ### Visualize datasets
Check out [example 1](./examples/1_load_lerobot_dataset.py) that illustrates how to use our dataset class which automatically downloads data from the Hugging Face hub. Check out [example 1](https://github.com/huggingface/lerobot/blob/main/examples/1_load_lerobot_dataset.py) that illustrates how to use our dataset class which automatically downloads data from the Hugging Face hub.
You can also locally visualize episodes from a dataset on the hub by executing our script from the command line: You can also locally visualize episodes from a dataset on the hub by executing our script from the command line:
@@ -212,7 +227,7 @@ Our script can also visualize datasets stored on a distant server. See `python -
A dataset in `LeRobotDataset` format is very simple to use. It can be loaded from a repository on the Hugging Face hub or a local folder simply with e.g. `dataset = LeRobotDataset("lerobot/aloha_static_coffee")` and can be indexed into like any Hugging Face and PyTorch dataset. For instance `dataset[0]` will retrieve a single temporal frame from the dataset containing observation(s) and an action as PyTorch tensors ready to be fed to a model. A dataset in `LeRobotDataset` format is very simple to use. It can be loaded from a repository on the Hugging Face hub or a local folder simply with e.g. `dataset = LeRobotDataset("lerobot/aloha_static_coffee")` and can be indexed into like any Hugging Face and PyTorch dataset. For instance `dataset[0]` will retrieve a single temporal frame from the dataset containing observation(s) and an action as PyTorch tensors ready to be fed to a model.
A specificity of `LeRobotDataset` is that, rather than retrieving a single frame by its index, we can retrieve several frames based on their temporal relationship with the indexed frame, by setting `delta_timestamps` to a list of relative times with respect to the indexed frame. For example, with `delta_timestamps = {"observation.image": [-1, -0.5, -0.2, 0]}` one can retrieve, for a given index, 4 frames: 3 "previous" frames 1 second, 0.5 seconds, and 0.2 seconds before the indexed frame, and the indexed frame itself (corresponding to the 0 entry). See example [1_load_lerobot_dataset.py](examples/1_load_lerobot_dataset.py) for more details on `delta_timestamps`. A specificity of `LeRobotDataset` is that, rather than retrieving a single frame by its index, we can retrieve several frames based on their temporal relationship with the indexed frame, by setting `delta_timestamps` to a list of relative times with respect to the indexed frame. For example, with `delta_timestamps = {"observation.image": [-1, -0.5, -0.2, 0]}` one can retrieve, for a given index, 4 frames: 3 "previous" frames 1 second, 0.5 seconds, and 0.2 seconds before the indexed frame, and the indexed frame itself (corresponding to the 0 entry). See example [1_load_lerobot_dataset.py](https://github.com/huggingface/lerobot/blob/main/examples/1_load_lerobot_dataset.py) for more details on `delta_timestamps`.
Under the hood, the `LeRobotDataset` format makes use of several ways to serialize data which can be useful to understand if you plan to work more closely with this format. We tried to make a flexible yet simple dataset format that would cover most type of features and specificities present in reinforcement learning and robotics, in simulation and in real-world, with a focus on cameras and robot states but easily extended to other types of sensory inputs as long as they can be represented by a tensor. Under the hood, the `LeRobotDataset` format makes use of several ways to serialize data which can be useful to understand if you plan to work more closely with this format. We tried to make a flexible yet simple dataset format that would cover most type of features and specificities present in reinforcement learning and robotics, in simulation and in real-world, with a focus on cameras and robot states but easily extended to other types of sensory inputs as long as they can be represented by a tensor.
@@ -266,12 +281,12 @@ Dataset can be uploaded/downloaded from the HuggingFace hub seamlessly. To work
### Evaluate a pretrained policy ### Evaluate a pretrained policy
Check out [example 2](./examples/2_evaluate_pretrained_policy.py) that illustrates how to download a pretrained policy from Hugging Face hub, and run an evaluation on its corresponding environment. Check out [example 2](https://github.com/huggingface/lerobot/blob/main/examples/2_evaluate_pretrained_policy.py) that illustrates how to download a pretrained policy from Hugging Face hub, and run an evaluation on its corresponding environment.
We also provide a more capable script to parallelize the evaluation over multiple environments during the same rollout. Here is an example with a pretrained model hosted on [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht): We also provide a more capable script to parallelize the evaluation over multiple environments during the same rollout. Here is an example with a pretrained model hosted on [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht):
```bash ```bash
python -m lerobot.scripts.eval \ lerobot-eval \
--policy.path=lerobot/diffusion_pusht \ --policy.path=lerobot/diffusion_pusht \
--env.type=pusht \ --env.type=pusht \
--eval.batch_size=10 \ --eval.batch_size=10 \
@@ -283,22 +298,22 @@ python -m lerobot.scripts.eval \
Note: After training your own policy, you can re-evaluate the checkpoints with: Note: After training your own policy, you can re-evaluate the checkpoints with:
```bash ```bash
python -m lerobot.scripts.eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model lerobot-eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
``` ```
See `python -m lerobot.scripts.eval --help` for more instructions. See `lerobot-eval --help` for more instructions.
### Train your own policy ### Train your own policy
Check out [example 3](./examples/3_train_policy.py) that illustrates how to train a model using our core library in python, and [example 4](./examples/4_train_policy_with_script.md) that shows how to use our training script from command line. Check out [example 3](https://github.com/huggingface/lerobot/blob/main/examples/3_train_policy.py) that illustrates how to train a model using our core library in python, and [example 4](https://github.com/huggingface/lerobot/blob/main/examples/4_train_policy_with_script.md) that shows how to use our training script from command line.
To use wandb for logging training and evaluation curves, make sure you've run `wandb login` as a one-time setup step. Then, when running the training command above, enable WandB in the configuration by adding `--wandb.enable=true`. To use wandb for logging training and evaluation curves, make sure you've run `wandb login` as a one-time setup step. Then, when running the training command above, enable WandB in the configuration by adding `--wandb.enable=true`.
A link to the wandb logs for the run will also show up in yellow in your terminal. Here is an example of what they look like in your browser. Please also check [here](./examples/4_train_policy_with_script.md#typical-logs-and-metrics) for the explanation of some commonly used metrics in logs. A link to the wandb logs for the run will also show up in yellow in your terminal. Here is an example of what they look like in your browser. Please also check [here](https://github.com/huggingface/lerobot/blob/main/examples/4_train_policy_with_script.md#typical-logs-and-metrics) for the explanation of some commonly used metrics in logs.
![](media/wandb.png) \<img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/wandb.png" alt="WandB logs example"\>
Note: For efficiency, during training every checkpoint is evaluated on a low number of episodes. You may use `--eval.n_episodes=500` to evaluate on more episodes than the default. Or, after training, you may want to re-evaluate your best checkpoints on more episodes or change the evaluation settings. See `python -m lerobot.scripts.eval --help` for more instructions. Note: For efficiency, during training every checkpoint is evaluated on a low number of episodes. You may use `--eval.n_episodes=500` to evaluate on more episodes than the default. Or, after training, you may want to re-evaluate your best checkpoints on more episodes or change the evaluation settings. See `lerobot-eval --help` for more instructions.
#### Reproduce state-of-the-art (SOTA) #### Reproduce state-of-the-art (SOTA)
@@ -306,7 +321,7 @@ We provide some pretrained policies on our [hub page](https://huggingface.co/ler
You can reproduce their training by loading the config from their run. Simply running: You can reproduce their training by loading the config from their run. Simply running:
```bash ```bash
python -m lerobot.scripts.train --config_path=lerobot/diffusion_pusht lerobot-train --config_path=lerobot/diffusion_pusht
``` ```
reproduces SOTA results for Diffusion Policy on the PushT task. reproduces SOTA results for Diffusion Policy on the PushT task.
@@ -315,26 +330,6 @@ reproduces SOTA results for Diffusion Policy on the PushT task.
If you would like to contribute to 🤗 LeRobot, please check out our [contribution guide](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md). If you would like to contribute to 🤗 LeRobot, please check out our [contribution guide](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md).
<!-- ### Add a new dataset
To add a dataset to the hub, you need to login using a write-access token, which can be generated from the [Hugging Face settings](https://huggingface.co/settings/tokens):
```bash
huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
```
Then point to your raw dataset folder (e.g. `data/aloha_static_pingpong_test_raw`), and push your dataset to the hub with:
```bash
python lerobot/scripts/push_dataset_to_hub.py \
--raw-dir data/aloha_static_pingpong_test_raw \
--out-dir data \
--repo-id lerobot/aloha_static_pingpong_test \
--raw-format aloha_hdf5
```
See `python lerobot/scripts/push_dataset_to_hub.py --help` for more instructions.
If your dataset format is not supported, implement your own in `lerobot/datasets/push_dataset_to_hub/${raw_format}_format.py` by copying examples like [pusht_zarr](https://github.com/huggingface/lerobot/blob/main/lerobot/datasets/push_dataset_to_hub/pusht_zarr_format.py), [umi_zarr](https://github.com/huggingface/lerobot/blob/main/lerobot/datasets/push_dataset_to_hub/umi_zarr_format.py), [aloha_hdf5](https://github.com/huggingface/lerobot/blob/main/lerobot/datasets/push_dataset_to_hub/aloha_hdf5_format.py), or [xarm_pkl](https://github.com/huggingface/lerobot/blob/main/lerobot/datasets/push_dataset_to_hub/xarm_pkl_format.py). -->
### Add a pretrained policy ### Add a pretrained policy
Once you have trained a policy you may upload it to the Hugging Face hub using a hub id that looks like `${hf_user}/${repo_name}` (e.g. [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht)). Once you have trained a policy you may upload it to the Hugging Face hub using a hub id that looks like `${hf_user}/${repo_name}` (e.g. [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht)).
@@ -351,34 +346,16 @@ To upload these to the hub, run the following:
huggingface-cli upload ${hf_user}/${repo_name} path/to/pretrained_model huggingface-cli upload ${hf_user}/${repo_name} path/to/pretrained_model
``` ```
See [eval.py](https://github.com/huggingface/lerobot/blob/main/lerobot/scripts/eval.py) for an example of how other people may use your policy. See [eval.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/eval.py) for an example of how other people may use your policy.
### Improve your code with profiling ### Acknowledgment
An example of a code snippet to profile the evaluation of a policy: - The LeRobot team 🤗 for building SmolVLA [Paper](https://arxiv.org/abs/2506.01844), [Blog](https://huggingface.co/blog/smolvla).
- Thanks to Tony Zhao, Zipeng Fu and colleagues for open sourcing ACT policy, ALOHA environments and datasets. Ours are adapted from [ALOHA](https://tonyzhaozh.github.io/aloha) and [Mobile ALOHA](https://mobile-aloha.github.io).
<!-- prettier-ignore-start --> - Thanks to Cheng Chi, Zhenjia Xu and colleagues for open sourcing Diffusion policy, Pusht environment and datasets, as well as UMI datasets. Ours are adapted from [Diffusion Policy](https://diffusion-policy.cs.columbia.edu) and [UMI Gripper](https://umi-gripper.github.io).
```python - Thanks to Nicklas Hansen, Yunhai Feng and colleagues for open sourcing TDMPC policy, Simxarm environments and datasets. Ours are adapted from [TDMPC](https://github.com/nicklashansen/tdmpc) and [FOWM](https://www.yunhaifeng.com/FOWM).
from torch.profiler import profile, record_function, ProfilerActivity - Thanks to Antonio Loquercio and Ashish Kumar for their early support.
- Thanks to [Seungjae (Jay) Lee](https://sjlee.cc/), [Mahi Shafiullah](https://mahis.life/) and colleagues for open sourcing [VQ-BeT](https://sjlee.cc/vq-bet/) policy and helping us adapt the codebase to our repository. The policy is adapted from [VQ-BeT repo](https://github.com/jayLEE0301/vq_bet_official).
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
```
<!-- prettier-ignore-end -->
## Citation ## Citation
@@ -386,83 +363,13 @@ If you want, you can cite this work with:
```bibtex ```bibtex
@misc{cadene2024lerobot, @misc{cadene2024lerobot,
author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Palma, Steven and Kooijmans, Pepijn and Aractingi, Michel and Shukor, Mustafa and Aubakirova, Dana and Russi, Martino and Capuano, Francesco and Pascale, Caroline and Choghari, Jade and Moss, Jess and Wolf, Thomas}, author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Palma, Steven and Kooijmans, Pepijn and Aractingi, Michel and Shukor, Mustafa and Aubakirova, Dana and Russi, Martino and Capuano, Francesco and Pascal, Caroline and Choghari, Jade and Moss, Jess and Wolf, Thomas},
title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch}, title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch},
howpublished = "\url{https://github.com/huggingface/lerobot}", howpublished = "\url{https://github.com/huggingface/lerobot}",
year = {2024} year = {2024}
} }
``` ```
Additionally, if you are using any of the particular policy architecture, pretrained models, or datasets, it is recommended to cite the original authors of the work as they appear below:
- [SmolVLA](https://arxiv.org/abs/2506.01844)
```bibtex
@article{shukor2025smolvla,
title={SmolVLA: A Vision-Language-Action Model for Affordable and Efficient Robotics},
author={Shukor, Mustafa and Aubakirova, Dana and Capuano, Francesco and Kooijmans, Pepijn and Palma, Steven and Zouitine, Adil and Aractingi, Michel and Pascal, Caroline and Russi, Martino and Marafioti, Andres and Alibert, Simon and Cord, Matthieu and Wolf, Thomas and Cadene, Remi},
journal={arXiv preprint arXiv:2506.01844},
year={2025}
}
```
- [Diffusion Policy](https://diffusion-policy.cs.columbia.edu)
```bibtex
@article{chi2024diffusionpolicy,
author = {Cheng Chi and Zhenjia Xu and Siyuan Feng and Eric Cousineau and Yilun Du and Benjamin Burchfiel and Russ Tedrake and Shuran Song},
title ={Diffusion Policy: Visuomotor Policy Learning via Action Diffusion},
journal = {The International Journal of Robotics Research},
year = {2024},
}
```
- [ACT or ALOHA](https://tonyzhaozh.github.io/aloha)
```bibtex
@article{zhao2023learning,
title={Learning fine-grained bimanual manipulation with low-cost hardware},
author={Zhao, Tony Z and Kumar, Vikash and Levine, Sergey and Finn, Chelsea},
journal={arXiv preprint arXiv:2304.13705},
year={2023}
}
```
- [TDMPC](https://www.nicklashansen.com/td-mpc/)
```bibtex
@inproceedings{Hansen2022tdmpc,
title={Temporal Difference Learning for Model Predictive Control},
author={Nicklas Hansen and Xiaolong Wang and Hao Su},
booktitle={ICML},
year={2022}
}
```
- [VQ-BeT](https://sjlee.cc/vq-bet/)
```bibtex
@article{lee2024behavior,
title={Behavior generation with latent actions},
author={Lee, Seungjae and Wang, Yibin and Etukuru, Haritheja and Kim, H Jin and Shafiullah, Nur Muhammad Mahi and Pinto, Lerrel},
journal={arXiv preprint arXiv:2403.03181},
year={2024}
}
```
- [HIL-SERL](https://hil-serl.github.io/)
```bibtex
@Article{luo2024hilserl,
title={Precise and Dexterous Robotic Manipulation via Human-in-the-Loop Reinforcement Learning},
author={Jianlan Luo and Charles Xu and Jeffrey Wu and Sergey Levine},
year={2024},
eprint={2410.21845},
archivePrefix={arXiv},
primaryClass={cs.RO}
}
```
## Star History ## Star History
[![Star History Chart](https://api.star-history.com/svg?repos=huggingface/lerobot&type=Timeline)](https://star-history.com/#huggingface/lerobot&Timeline) [![Star History Chart](https://api.star-history.com/svg?repos=huggingface/lerobot&type=Timeline)](https://star-history.com/#huggingface/lerobot&Timeline)
docs-requirements.txt
+3
View File
@@ -0,0 +1,3 @@
# docs-requirements.txt
hf-doc-builder @ git+https://github.com/huggingface/doc-builder.git@main
watchdog>=6.0.0
docs/README.md
+1 -1
View File
@@ -20,7 +20,7 @@ To generate the documentation, you first have to build it. Several packages are
you can install them with the following command, at the root of the code repository: you can install them with the following command, at the root of the code repository:
```bash ```bash
pip install -e ".[docs]" pip install -e . -r docs-requirements.txt
``` ```
You will also need `nodejs`. Please refer to their [installation page](https://nodejs.org/en/download) You will also need `nodejs`. Please refer to their [installation page](https://nodejs.org/en/download)
docs/source/cameras.mdx
+1 -1
View File
@@ -9,7 +9,7 @@ To instantiate a camera, you need a camera identifier. This identifier might cha
To find the camera indices of the cameras plugged into your system, run the following script: To find the camera indices of the cameras plugged into your system, run the following script:
```bash ```bash
python -m lerobot.find_cameras opencv # or realsense for Intel Realsense cameras lerobot-find-cameras opencv # or realsense for Intel Realsense cameras
``` ```
The output will look something like this if you have two cameras connected: The output will look something like this if you have two cameras connected:
docs/source/hilserl.mdx
+2 -2
View File
@@ -412,7 +412,7 @@ Example configuration for training the [reward classifier](https://huggingface.c
To train the classifier, use the `train.py` script with your configuration: To train the classifier, use the `train.py` script with your configuration:
```bash ```bash
python -m lerobot.scripts.train --config_path path/to/reward_classifier_train_config.json lerobot-train --config_path path/to/reward_classifier_train_config.json
``` ```
**Deploying and Testing the Model** **Deploying and Testing the Model**
@@ -458,7 +458,7 @@ The reward classifier will automatically provide rewards based on the visual inp
3. **Train the classifier**: 3. **Train the classifier**:
```bash ```bash
python -m lerobot.scripts.train --config_path src/lerobot/configs/reward_classifier_train_config.json lerobot-train --config_path src/lerobot/configs/reward_classifier_train_config.json
``` ```
4. **Test the classifier**: 4. **Test the classifier**:
docs/source/hope_jr.mdx
+11 -11
View File
@@ -19,7 +19,7 @@ pip install -e ".[hopejr]"
Before starting calibration and operation, you need to identify the USB ports for each HopeJR component. Run this script to find the USB ports for the arm, hand, glove, and exoskeleton: Before starting calibration and operation, you need to identify the USB ports for each HopeJR component. Run this script to find the USB ports for the arm, hand, glove, and exoskeleton:
```bash ```bash
python -m lerobot.find_port lerobot-find-port
``` ```
This will display the available USB ports and their associated devices. Make note of the port paths (e.g., `/dev/tty.usbmodem58760433331`, `/dev/tty.usbmodem11301`) as you'll need to specify them in the `--robot.port` and `--teleop.port` parameters when recording data, replaying episodes, or running teleoperation scripts. This will display the available USB ports and their associated devices. Make note of the port paths (e.g., `/dev/tty.usbmodem58760433331`, `/dev/tty.usbmodem11301`) as you'll need to specify them in the `--robot.port` and `--teleop.port` parameters when recording data, replaying episodes, or running teleoperation scripts.
@@ -31,7 +31,7 @@ Before performing teleoperation, HopeJR's limbs need to be calibrated. Calibrati
### 1.1 Calibrate Robot Hand ### 1.1 Calibrate Robot Hand
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--robot.type=hope_jr_hand \ --robot.type=hope_jr_hand \
--robot.port=/dev/tty.usbmodem58760432281 \ --robot.port=/dev/tty.usbmodem58760432281 \
--robot.id=blue \ --robot.id=blue \
@@ -81,7 +81,7 @@ Once you have set the appropriate boundaries for all joints, click "Save" to sav
### 1.2 Calibrate Teleoperator Glove ### 1.2 Calibrate Teleoperator Glove
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--teleop.type=homunculus_glove \ --teleop.type=homunculus_glove \
--teleop.port=/dev/tty.usbmodem11201 \ --teleop.port=/dev/tty.usbmodem11201 \
--teleop.id=red \ --teleop.id=red \
@@ -120,7 +120,7 @@ Once calibration is complete, the system will save the calibration to `/Users/yo
### 1.3 Calibrate Robot Arm ### 1.3 Calibrate Robot Arm
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--robot.type=hope_jr_arm \ --robot.type=hope_jr_arm \
--robot.port=/dev/tty.usbserial-1110 \ --robot.port=/dev/tty.usbserial-1110 \
--robot.id=white --robot.id=white
@@ -146,7 +146,7 @@ Use the calibration interface to set the range boundaries for each joint. Move e
### 1.4 Calibrate Teleoperator Exoskeleton ### 1.4 Calibrate Teleoperator Exoskeleton
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--teleop.type=homunculus_arm \ --teleop.type=homunculus_arm \
--teleop.port=/dev/tty.usbmodem11201 \ --teleop.port=/dev/tty.usbmodem11201 \
--teleop.id=black --teleop.id=black
@@ -178,7 +178,7 @@ Due to global variable conflicts in the Feetech middleware, teleoperation for ar
### Hand ### Hand
```bash ```bash
python -m lerobot.teleoperate \ lerobot-teleoperate \
--robot.type=hope_jr_hand \ --robot.type=hope_jr_hand \
--robot.port=/dev/tty.usbmodem58760432281 \ --robot.port=/dev/tty.usbmodem58760432281 \
--robot.id=blue \ --robot.id=blue \
@@ -194,7 +194,7 @@ python -m lerobot.teleoperate \
### Arm ### Arm
```bash ```bash
python -m lerobot.teleoperate \ lerobot-teleoperate \
--robot.type=hope_jr_arm \ --robot.type=hope_jr_arm \
--robot.port=/dev/tty.usbserial-1110 \ --robot.port=/dev/tty.usbserial-1110 \
--robot.id=white \ --robot.id=white \
@@ -214,7 +214,7 @@ Record, Replay and Train with Hope-JR is still experimental.
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings). This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
```bash ```bash
python -m lerobot.record \ lerobot-record \
--robot.type=hope_jr_hand \ --robot.type=hope_jr_hand \
--robot.port=/dev/tty.usbmodem58760432281 \ --robot.port=/dev/tty.usbmodem58760432281 \
--robot.id=right \ --robot.id=right \
@@ -236,7 +236,7 @@ python -m lerobot.record \
### Replay ### Replay
```bash ```bash
python -m lerobot.replay \ lerobot-replay \
--robot.type=hope_jr_hand \ --robot.type=hope_jr_hand \
--robot.port=/dev/tty.usbmodem58760432281 \ --robot.port=/dev/tty.usbmodem58760432281 \
--robot.id=right \ --robot.id=right \
@@ -248,7 +248,7 @@ python -m lerobot.replay \
### Train ### Train
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--dataset.repo_id=nepyope/hand_record_test_with_video_data \ --dataset.repo_id=nepyope/hand_record_test_with_video_data \
--policy.type=act \ --policy.type=act \
--output_dir=outputs/train/hopejr_hand \ --output_dir=outputs/train/hopejr_hand \
@@ -263,7 +263,7 @@ python -m lerobot.scripts.train \
This training run can be viewed as an example [here](https://wandb.ai/tino/lerobot/runs/rp0k8zvw?nw=nwusertino). This training run can be viewed as an example [here](https://wandb.ai/tino/lerobot/runs/rp0k8zvw?nw=nwusertino).
```bash ```bash
python -m lerobot.record \ lerobot-record \
--robot.type=hope_jr_hand \ --robot.type=hope_jr_hand \
--robot.port=/dev/tty.usbmodem58760432281 \ --robot.port=/dev/tty.usbmodem58760432281 \
--robot.id=right \ --robot.id=right \
docs/source/il_robots.mdx
+12 -12
View File
@@ -45,7 +45,7 @@ Note that the `id` associated with a robot is used to store the calibration file
<hfoptions id="teleoperate_so101"> <hfoptions id="teleoperate_so101">
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.teleoperate \ lerobot-teleoperate \
--robot.type=so101_follower \ --robot.type=so101_follower \
--robot.port=/dev/tty.usbmodem58760431541 \ --robot.port=/dev/tty.usbmodem58760431541 \
--robot.id=my_awesome_follower_arm \ --robot.id=my_awesome_follower_arm \
@@ -101,7 +101,7 @@ With `rerun`, you can teleoperate again while simultaneously visualizing the cam
<hfoptions id="teleoperate_koch_camera"> <hfoptions id="teleoperate_koch_camera">
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.teleoperate \ lerobot-teleoperate \
--robot.type=koch_follower \ --robot.type=koch_follower \
--robot.port=/dev/tty.usbmodem58760431541 \ --robot.port=/dev/tty.usbmodem58760431541 \
--robot.id=my_awesome_follower_arm \ --robot.id=my_awesome_follower_arm \
@@ -174,7 +174,7 @@ Now you can record a dataset. To record 5 episodes and upload your dataset to th
<hfoptions id="record"> <hfoptions id="record">
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.record \ lerobot-record \
--robot.type=so101_follower \ --robot.type=so101_follower \
--robot.port=/dev/tty.usbmodem585A0076841 \ --robot.port=/dev/tty.usbmodem585A0076841 \
--robot.id=my_awesome_follower_arm \ --robot.id=my_awesome_follower_arm \
@@ -294,7 +294,7 @@ dataset.push_to_hub()
#### Dataset upload #### Dataset upload
Locally, your dataset is stored in this folder: `~/.cache/huggingface/lerobot/{repo-id}`. At the end of data recording, your dataset will be uploaded on your Hugging Face page (e.g. https://huggingface.co/datasets/cadene/so101_test) that you can obtain by running: Locally, your dataset is stored in this folder: `~/.cache/huggingface/lerobot/{repo-id}`. At the end of data recording, your dataset will be uploaded on your Hugging Face page (e.g. `https://huggingface.co/datasets/${HF_USER}/so101_test`) that you can obtain by running:
```bash ```bash
echo https://huggingface.co/datasets/${HF_USER}/so101_test echo https://huggingface.co/datasets/${HF_USER}/so101_test
@@ -376,7 +376,7 @@ You can replay the first episode on your robot with either the command below or
<hfoptions id="replay"> <hfoptions id="replay">
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.replay \ lerobot-replay \
--robot.type=so101_follower \ --robot.type=so101_follower \
--robot.port=/dev/tty.usbmodem58760431541 \ --robot.port=/dev/tty.usbmodem58760431541 \
--robot.id=my_awesome_follower_arm \ --robot.id=my_awesome_follower_arm \
@@ -428,10 +428,10 @@ Your robot should replicate movements similar to those you recorded. For example
## Train a policy ## Train a policy
To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](../src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command: To train a policy to control your robot, use the [`lerobot-train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--dataset.repo_id=${HF_USER}/so101_test \ --dataset.repo_id=${HF_USER}/so101_test \
--policy.type=act \ --policy.type=act \
--output_dir=outputs/train/act_so101_test \ --output_dir=outputs/train/act_so101_test \
@@ -444,7 +444,7 @@ python -m lerobot.scripts.train \
Let's explain the command: Let's explain the command:
1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/so101_test`. 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/so101_test`.
2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../src/lerobot/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset. 2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
3. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon. 3. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
4. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`. 4. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
@@ -453,7 +453,7 @@ Training should take several hours. You will find checkpoints in `outputs/train/
To resume training from a checkpoint, below is an example command to resume from `last` checkpoint of the `act_so101_test` policy: To resume training from a checkpoint, below is an example command to resume from `last` checkpoint of the `act_so101_test` policy:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--config_path=outputs/train/act_so101_test/checkpoints/last/pretrained_model/train_config.json \ --config_path=outputs/train/act_so101_test/checkpoints/last/pretrained_model/train_config.json \
--resume=true --resume=true
``` ```
@@ -462,9 +462,9 @@ If you do not want to push your model to the hub after training use `--policy.pu
Additionally you can provide extra `tags` or specify a `license` for your model or make the model repo `private` by adding this: `--policy.private=true --policy.tags=\[ppo,rl\] --policy.license=mit` Additionally you can provide extra `tags` or specify a `license` for your model or make the model repo `private` by adding this: `--policy.private=true --policy.tags=\[ppo,rl\] --policy.license=mit`
#### Train using Collab #### Train using Google Colab
If your local computer doesn't have a powerful GPU you could utilize Google Collab to train your model by following the [ACT training notebook](./notebooks#training-act). If your local computer doesn't have a powerful GPU you could utilize Google Colab to train your model by following the [ACT training notebook](./notebooks#training-act).
#### Upload policy checkpoints #### Upload policy checkpoints
@@ -490,7 +490,7 @@ You can use the `record` script from [`lerobot/record.py`](https://github.com/hu
<hfoptions id="eval"> <hfoptions id="eval">
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.record \ lerobot-record \
--robot.type=so100_follower \ --robot.type=so100_follower \
--robot.port=/dev/ttyACM1 \ --robot.port=/dev/ttyACM1 \
--robot.cameras="{ up: {type: opencv, index_or_path: /dev/video10, width: 640, height: 480, fps: 30}, side: {type: intelrealsense, serial_number_or_name: 233522074606, width: 640, height: 480, fps: 30}}" \ --robot.cameras="{ up: {type: opencv, index_or_path: /dev/video10, width: 640, height: 480, fps: 30}, side: {type: intelrealsense, serial_number_or_name: 233522074606, width: 640, height: 480, fps: 30}}" \
docs/source/il_sim.mdx
+3 -3
View File
@@ -96,10 +96,10 @@ If you uploaded your dataset to the hub you can [visualize your dataset online](
## Train a policy ## Train a policy
To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](../src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command: To train a policy to control your robot, use the [`lerobot-train`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--dataset.repo_id=${HF_USER}/il_gym \ --dataset.repo_id=${HF_USER}/il_gym \
--policy.type=act \ --policy.type=act \
--output_dir=outputs/train/il_sim_test \ --output_dir=outputs/train/il_sim_test \
@@ -111,7 +111,7 @@ python -m lerobot.scripts.train \
Let's explain the command: Let's explain the command:
1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/il_gym`. 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/il_gym`.
2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../src/lerobot/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset. 2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
3. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon. 3. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
4. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`. 4. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
docs/source/installation.mdx
+39 -11
View File
@@ -1,15 +1,6 @@
# Installation # Installation
## Install LeRobot ## Environment Setup
Currently only available from source.
Download our source code:
```bash
git clone https://github.com/huggingface/lerobot.git
cd lerobot
```
Create a virtual environment with Python 3.10, using [`Miniconda`](https://docs.anaconda.com/miniconda/install/#quick-command-line-install) Create a virtual environment with Python 3.10, using [`Miniconda`](https://docs.anaconda.com/miniconda/install/#quick-command-line-install)
@@ -40,12 +31,49 @@ conda install ffmpeg -c conda-forge
> >
> - _[On Linux only]_ If you want to bring your own ffmpeg: Install [ffmpeg build dependencies](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#GettheDependencies) and [compile ffmpeg from source with libsvtav1](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#libsvtav1), and make sure you use the corresponding ffmpeg binary to your install with `which ffmpeg`. > - _[On Linux only]_ If you want to bring your own ffmpeg: Install [ffmpeg build dependencies](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#GettheDependencies) and [compile ffmpeg from source with libsvtav1](https://trac.ffmpeg.org/wiki/CompilationGuide/Ubuntu#libsvtav1), and make sure you use the corresponding ffmpeg binary to your install with `which ffmpeg`.
Install 🤗 LeRobot: ## Install LeRobot 🤗
### From Source
First, clone the repository and navigate into the directory:
```bash
git clone https://github.com/huggingface/lerobot.git
cd lerobot
```
Then, install the library in editable mode. This is useful if you plan to contribute to the code.
```bash ```bash
pip install -e . pip install -e .
``` ```
### Installation from PyPI
**Core Library:**
Install the base package with:
```bash
pip install lerobot
```
_This installs only the default dependencies._
**Extra Features:**
To install additional functionality, use one of the following:
```bash
pip install 'lerobot[all]' # All available features
pip install 'lerobot[aloha,pusht]' # Specific features (Aloha & Pusht)
pip install 'lerobot[feetech]' # Feetech motor support
```
_Replace `[...]` with your desired features._
**Available Tags:**
For a full list of optional dependencies, see:
https://pypi.org/project/lerobot/
### Troubleshooting ### Troubleshooting
If you encounter build errors, you may need to install additional dependencies: `cmake`, `build-essential`, and `ffmpeg libs`. If you encounter build errors, you may need to install additional dependencies: `cmake`, `build-essential`, and `ffmpeg libs`.
docs/source/koch.mdx
+5 -5
View File
@@ -31,7 +31,7 @@ pip install -e ".[dynamixel]"
To find the port for each bus servo adapter, run this script: To find the port for each bus servo adapter, run this script:
```bash ```bash
python -m lerobot.find_port lerobot-find-port
``` ```
<hfoptions id="example"> <hfoptions id="example">
@@ -98,7 +98,7 @@ For a visual reference on how to set the motor ids please refer to [this video](
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.setup_motors \ lerobot-setup-motors \
--robot.type=koch_follower \ --robot.type=koch_follower \
--robot.port=/dev/tty.usbmodem575E0031751 # <- paste here the port found at previous step --robot.port=/dev/tty.usbmodem575E0031751 # <- paste here the port found at previous step
``` ```
@@ -174,7 +174,7 @@ Do the same steps for the leader arm but modify the command or script accordingl
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.setup_motors \ lerobot-setup-motors \
--teleop.type=koch_leader \ --teleop.type=koch_leader \
--teleop.port=/dev/tty.usbmodem575E0031751 \ # <- paste here the port found at previous step --teleop.port=/dev/tty.usbmodem575E0031751 \ # <- paste here the port found at previous step
``` ```
@@ -211,7 +211,7 @@ Run the following command or API example to calibrate the follower arm:
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--robot.type=koch_follower \ --robot.type=koch_follower \
--robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
--robot.id=my_awesome_follower_arm # <- Give the robot a unique name --robot.id=my_awesome_follower_arm # <- Give the robot a unique name
@@ -249,7 +249,7 @@ Do the same steps to calibrate the leader arm, run the following command or API
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--teleop.type=koch_leader \ --teleop.type=koch_leader \
--teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
--teleop.id=my_awesome_leader_arm # <- Give the robot a unique name --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name
docs/source/lekiwi.mdx
+4 -4
View File
@@ -60,7 +60,7 @@ First, we will assemble the two SO100/SO101 arms. One to attach to the mobile ba
To find the port for each bus servo adapter, run this script: To find the port for each bus servo adapter, run this script:
```bash ```bash
python -m lerobot.find_port lerobot-find-port
``` ```
<hfoptions id="example"> <hfoptions id="example">
@@ -116,7 +116,7 @@ The instructions for configuring the motors can be found in the SO101 [docs](./s
You can run this command to setup motors for LeKiwi. It will first setup the motors for arm (id 6..1) and then setup motors for wheels (9,8,7) You can run this command to setup motors for LeKiwi. It will first setup the motors for arm (id 6..1) and then setup motors for wheels (9,8,7)
```bash ```bash
python -m lerobot.setup_motors \ lerobot-setup-motors \
--robot.type=lekiwi \ --robot.type=lekiwi \
--robot.port=/dev/tty.usbmodem58760431551 # <- paste here the port found at previous step --robot.port=/dev/tty.usbmodem58760431551 # <- paste here the port found at previous step
``` ```
@@ -174,7 +174,7 @@ The calibration process is very important because it allows a neural network tra
Make sure the arm is connected to the Raspberry Pi and run this script or API example (on the Raspberry Pi via SSH) to launch calibration of the follower arm: Make sure the arm is connected to the Raspberry Pi and run this script or API example (on the Raspberry Pi via SSH) to launch calibration of the follower arm:
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--robot.type=lekiwi \ --robot.type=lekiwi \
--robot.id=my_awesome_kiwi # <- Give the robot a unique name --robot.id=my_awesome_kiwi # <- Give the robot a unique name
``` ```
@@ -193,7 +193,7 @@ Then, to calibrate the leader arm (which is attached to the laptop/pc). Run the
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--teleop.type=so100_leader \ --teleop.type=so100_leader \
--teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
--teleop.id=my_awesome_leader_arm # <- Give the robot a unique name --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name
docs/source/policy_act_README.md
+14
View File
@@ -0,0 +1,14 @@
## Paper
https://tonyzhaozh.github.io/aloha
## Citation
```bibtex
@article{zhao2023learning,
title={Learning fine-grained bimanual manipulation with low-cost hardware},
author={Zhao, Tony Z and Kumar, Vikash and Levine, Sergey and Finn, Chelsea},
journal={arXiv preprint arXiv:2304.13705},
year={2023}
}
```
docs/source/policy_diffusion_README.md
+14
View File
@@ -0,0 +1,14 @@
## Paper
https://diffusion-policy.cs.columbia.edu
## Citation
```bibtex
@article{chi2024diffusionpolicy,
author = {Cheng Chi and Zhenjia Xu and Siyuan Feng and Eric Cousineau and Yilun Du and Benjamin Burchfiel and Russ Tedrake and Shuran Song},
title ={Diffusion Policy: Visuomotor Policy Learning via Action Diffusion},
journal = {The International Journal of Robotics Research},
year = {2024},
}
```
docs/source/policy_smolvla_README.md
+14
View File
@@ -0,0 +1,14 @@
## Paper
https://arxiv.org/abs/2506.01844
## Citation
```bibtex
@article{shukor2025smolvla,
title={SmolVLA: A Vision-Language-Action Model for Affordable and Efficient Robotics},
author={Shukor, Mustafa and Aubakirova, Dana and Capuano, Francesco and Kooijmans, Pepijn and Palma, Steven and Zouitine, Adil and Aractingi, Michel and Pascal, Caroline and Russi, Martino and Marafioti, Andres and Alibert, Simon and Cord, Matthieu and Wolf, Thomas and Cadene, Remi},
journal={arXiv preprint arXiv:2506.01844},
year={2025}
}
```
docs/source/policy_tdmpc_README.md
+14
View File
@@ -0,0 +1,14 @@
## Paper
https://www.nicklashansen.com/td-mpc/
## Citation
```bibtex
@inproceedings{Hansen2022tdmpc,
title={Temporal Difference Learning for Model Predictive Control},
author={Nicklas Hansen and Xiaolong Wang and Hao Su},
booktitle={ICML},
year={2022}
}
```
docs/source/policy_vqbet_README.md
+14
View File
@@ -0,0 +1,14 @@
## Paper
https://sjlee.cc/vq-bet/
## Citation
```bibtex
@article{lee2024behavior,
title={Behavior generation with latent actions},
author={Lee, Seungjae and Wang, Yibin and Etukuru, Haritheja and Kim, H Jin and Shafiullah, Nur Muhammad Mahi and Pinto, Lerrel},
journal={arXiv preprint arXiv:2403.03181},
year={2024}
}
```
docs/source/smolvla.mdx
+3 -3
View File
@@ -54,7 +54,7 @@ If you don't have a gpu device, you can train using our notebook on [![Google Co
Pass your dataset to the training script using `--dataset.repo_id`. If you want to test your installation, run the following command where we use one of the datasets we collected for the [SmolVLA Paper](https://huggingface.co/papers/2506.01844). Pass your dataset to the training script using `--dataset.repo_id`. If you want to test your installation, run the following command where we use one of the datasets we collected for the [SmolVLA Paper](https://huggingface.co/papers/2506.01844).
```bash ```bash
cd lerobot && python -m lerobot.scripts.train \ cd lerobot && lerobot-train \
--policy.path=lerobot/smolvla_base \ --policy.path=lerobot/smolvla_base \
--dataset.repo_id=${HF_USER}/mydataset \ --dataset.repo_id=${HF_USER}/mydataset \
--batch_size=64 \ --batch_size=64 \
@@ -73,7 +73,7 @@ cd lerobot && python -m lerobot.scripts.train \
Fine-tuning is an art. For a complete overview of the options for finetuning, run Fine-tuning is an art. For a complete overview of the options for finetuning, run
```bash ```bash
python -m lerobot.scripts.train --help lerobot-train --help
``` ```
<p align="center"> <p align="center">
@@ -97,7 +97,7 @@ Similarly for when recording an episode, it is recommended that you are logged i
Once you are logged in, you can run inference in your setup by doing: Once you are logged in, you can run inference in your setup by doing:
```bash ```bash
python -m lerobot.record \ lerobot-record \
--robot.type=so101_follower \ --robot.type=so101_follower \
--robot.port=/dev/ttyACM0 \ # <- Use your port --robot.port=/dev/ttyACM0 \ # <- Use your port
--robot.id=my_blue_follower_arm \ # <- Use your robot id --robot.id=my_blue_follower_arm \ # <- Use your robot id
docs/source/so100.mdx
+5 -5
View File
@@ -26,7 +26,7 @@ Unlike the SO-101, the motor connectors are not easily accessible once the arm i
To find the port for each bus servo adapter, run this script: To find the port for each bus servo adapter, run this script:
```bash ```bash
python -m lerobot.find_port lerobot-find-port
``` ```
<hfoptions id="example"> <hfoptions id="example">
@@ -93,7 +93,7 @@ For a visual reference on how to set the motor ids please refer to [this video](
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.setup_motors \ lerobot-setup-motors \
--robot.type=so100_follower \ --robot.type=so100_follower \
--robot.port=/dev/tty.usbmodem585A0076841 # <- paste here the port found at previous step --robot.port=/dev/tty.usbmodem585A0076841 # <- paste here the port found at previous step
``` ```
@@ -168,7 +168,7 @@ Do the same steps for the leader arm.
<hfoptions id="setup_motors"> <hfoptions id="setup_motors">
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.setup_motors \ lerobot-setup-motors \
--teleop.type=so100_leader \ --teleop.type=so100_leader \
--teleop.port=/dev/tty.usbmodem575E0031751 # <- paste here the port found at previous step --teleop.port=/dev/tty.usbmodem575E0031751 # <- paste here the port found at previous step
``` ```
@@ -568,7 +568,7 @@ Run the following command or API example to calibrate the follower arm:
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--robot.type=so100_follower \ --robot.type=so100_follower \
--robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
--robot.id=my_awesome_follower_arm # <- Give the robot a unique name --robot.id=my_awesome_follower_arm # <- Give the robot a unique name
@@ -606,7 +606,7 @@ Do the same steps to calibrate the leader arm, run the following command or API
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--teleop.type=so100_leader \ --teleop.type=so100_leader \
--teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
--teleop.id=my_awesome_leader_arm # <- Give the robot a unique name --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name
docs/source/so101.mdx
+5 -5
View File
@@ -162,7 +162,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
To find the port for each bus servo adapter, connect MotorBus to your computer via USB and power. Run the following script and disconnect the MotorBus when prompted: To find the port for each bus servo adapter, connect MotorBus to your computer via USB and power. Run the following script and disconnect the MotorBus when prompted:
```bash ```bash
python -m lerobot.find_port lerobot-find-port
``` ```
<hfoptions id="example"> <hfoptions id="example">
@@ -240,7 +240,7 @@ Connect the usb cable from your computer and the power supply to the follower ar
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.setup_motors \ lerobot-setup-motors \
--robot.type=so101_follower \ --robot.type=so101_follower \
--robot.port=/dev/tty.usbmodem585A0076841 # <- paste here the port found at previous step --robot.port=/dev/tty.usbmodem585A0076841 # <- paste here the port found at previous step
``` ```
@@ -316,7 +316,7 @@ Do the same steps for the leader arm.
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.setup_motors \ lerobot-setup-motors \
--teleop.type=so101_leader \ --teleop.type=so101_leader \
--teleop.port=/dev/tty.usbmodem575E0031751 # <- paste here the port found at previous step --teleop.port=/dev/tty.usbmodem575E0031751 # <- paste here the port found at previous step
``` ```
@@ -353,7 +353,7 @@ Run the following command or API example to calibrate the follower arm:
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--robot.type=so101_follower \ --robot.type=so101_follower \
--robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot --robot.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
--robot.id=my_awesome_follower_arm # <- Give the robot a unique name --robot.id=my_awesome_follower_arm # <- Give the robot a unique name
@@ -402,7 +402,7 @@ Do the same steps to calibrate the leader arm, run the following command or API
<hfoption id="Command"> <hfoption id="Command">
```bash ```bash
python -m lerobot.calibrate \ lerobot-calibrate \
--teleop.type=so101_leader \ --teleop.type=so101_leader \
--teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot --teleop.port=/dev/tty.usbmodem58760431551 \ # <- The port of your robot
--teleop.id=my_awesome_leader_arm # <- Give the robot a unique name --teleop.id=my_awesome_leader_arm # <- Give the robot a unique name
examples/4_train_policy_with_script.md
+14 -14
View File
@@ -62,7 +62,7 @@ By default, every field takes its default value specified in the dataclass. If a
Let's say that we want to train [Diffusion Policy](../src/lerobot/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this: Let's say that we want to train [Diffusion Policy](../src/lerobot/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--dataset.repo_id=lerobot/pusht \ --dataset.repo_id=lerobot/pusht \
--policy.type=diffusion \ --policy.type=diffusion \
--env.type=pusht --env.type=pusht
@@ -77,7 +77,7 @@ Let's break this down:
Let's see another example. Let's say you've been training [ACT](../src/lerobot/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with: Let's see another example. Let's say you've been training [ACT](../src/lerobot/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=act \ --policy.type=act \
--dataset.repo_id=lerobot/aloha_sim_insertion_human \ --dataset.repo_id=lerobot/aloha_sim_insertion_human \
--env.type=aloha \ --env.type=aloha \
@@ -90,7 +90,7 @@ We now want to train a different policy for aloha on another task. We'll change
Looking at the [`AlohaEnv`](../src/lerobot/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using: Looking at the [`AlohaEnv`](../src/lerobot/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=act \ --policy.type=act \
--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \ --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
--env.type=aloha \ --env.type=aloha \
@@ -127,7 +127,7 @@ Now, let's assume that we want to reproduce the run just above. That run has pro
We can then simply load the config values from this file using: We can then simply load the config values from this file using:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \ --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
--output_dir=outputs/train/act_aloha_transfer_2 --output_dir=outputs/train/act_aloha_transfer_2
``` ```
@@ -137,7 +137,7 @@ python -m lerobot.scripts.train \
Similarly to Hydra, we can still override some parameters in the CLI if we want to, e.g.: Similarly to Hydra, we can still override some parameters in the CLI if we want to, e.g.:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \ --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
--output_dir=outputs/train/act_aloha_transfer_2 --output_dir=outputs/train/act_aloha_transfer_2
--policy.n_action_steps=80 --policy.n_action_steps=80
@@ -148,7 +148,7 @@ python -m lerobot.scripts.train \
`--config_path` can also accept the repo_id of a repo on the hub that contains a `train_config.json` file, e.g. running: `--config_path` can also accept the repo_id of a repo on the hub that contains a `train_config.json` file, e.g. running:
```bash ```bash
python -m lerobot.scripts.train --config_path=lerobot/diffusion_pusht lerobot-train --config_path=lerobot/diffusion_pusht
``` ```
will start a training run with the same configuration used for training [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht) will start a training run with the same configuration used for training [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht)
@@ -160,7 +160,7 @@ Being able to resume a training run is important in case it crashed or aborted f
Let's reuse the command from the previous run and add a few more options: Let's reuse the command from the previous run and add a few more options:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=act \ --policy.type=act \
--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \ --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
--env.type=aloha \ --env.type=aloha \
@@ -179,7 +179,7 @@ INFO 2025-01-24 16:10:56 ts/train.py:263 Checkpoint policy after step 100
Now let's simulate a crash by killing the process (hit `ctrl`+`c`). We can then simply resume this run from the last checkpoint available with: Now let's simulate a crash by killing the process (hit `ctrl`+`c`). We can then simply resume this run from the last checkpoint available with:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \ --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
--resume=true --resume=true
``` ```
@@ -190,7 +190,7 @@ Another reason for which you might want to resume a run is simply to extend trai
You could double the number of steps of the previous run with: You could double the number of steps of the previous run with:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \ --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
--resume=true \ --resume=true \
--steps=200000 --steps=200000
@@ -224,7 +224,7 @@ In addition to the features currently in Draccus, we've added a special `.path`
For example, we could fine-tune a [policy pre-trained on the aloha transfer task](https://huggingface.co/lerobot/act_aloha_sim_transfer_cube_human) on the aloha insertion task. We can achieve this with: For example, we could fine-tune a [policy pre-trained on the aloha transfer task](https://huggingface.co/lerobot/act_aloha_sim_transfer_cube_human) on the aloha insertion task. We can achieve this with:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.path=lerobot/act_aloha_sim_transfer_cube_human \ --policy.path=lerobot/act_aloha_sim_transfer_cube_human \
--dataset.repo_id=lerobot/aloha_sim_insertion_human \ --dataset.repo_id=lerobot/aloha_sim_insertion_human \
--env.type=aloha \ --env.type=aloha \
@@ -270,7 +270,7 @@ We'll summarize here the main use cases to remember from this tutorial.
#### Train a policy from scratch – CLI #### Train a policy from scratch – CLI
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=act \ # <- select 'act' policy --policy.type=act \ # <- select 'act' policy
--env.type=pusht \ # <- select 'pusht' environment --env.type=pusht \ # <- select 'pusht' environment
--dataset.repo_id=lerobot/pusht # <- train on this dataset --dataset.repo_id=lerobot/pusht # <- train on this dataset
@@ -279,7 +279,7 @@ python -m lerobot.scripts.train \
#### Train a policy from scratch - config file + CLI #### Train a policy from scratch - config file + CLI
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--config_path=path/to/pretrained_model \ # <- can also be a repo_id --config_path=path/to/pretrained_model \ # <- can also be a repo_id
--policy.n_action_steps=80 # <- you may still override values --policy.n_action_steps=80 # <- you may still override values
``` ```
@@ -287,7 +287,7 @@ python -m lerobot.scripts.train \
#### Resume/continue a training run #### Resume/continue a training run
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--config_path=checkpoint/pretrained_model/ \ --config_path=checkpoint/pretrained_model/ \
--resume=true \ --resume=true \
--steps=200000 # <- you can change some training parameters --steps=200000 # <- you can change some training parameters
@@ -296,7 +296,7 @@ python -m lerobot.scripts.train \
#### Fine-tuning #### Fine-tuning
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.path=lerobot/act_aloha_sim_transfer_cube_human \ # <- can also be a local path to a checkpoint --policy.path=lerobot/act_aloha_sim_transfer_cube_human \ # <- can also be a local path to a checkpoint
--dataset.repo_id=lerobot/aloha_sim_insertion_human \ --dataset.repo_id=lerobot/aloha_sim_insertion_human \
--env.type=aloha \ --env.type=aloha \
examples/backward_compatibility/replay.py
+1 -1
View File
@@ -18,7 +18,7 @@ Replays the actions of an episode from a dataset on a robot.
Example: Example:
```shell ```shell
python -m lerobot.replay \ lerobot-replay \
--robot.type=so100_follower \ --robot.type=so100_follower \
--robot.port=/dev/tty.usbmodem58760431541 \ --robot.port=/dev/tty.usbmodem58760431541 \
--robot.id=black \ --robot.id=black \
pyproject.toml
+5 -6
View File
@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
[project] [project]
name = "lerobot" name = "lerobot"
version = "0.2.0" version = "0.3.4"
description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch" description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
readme = "README.md" readme = "README.md"
license = { text = "Apache-2.0" } license = { text = "Apache-2.0" }
@@ -68,15 +68,16 @@ dependencies = [
"einops>=0.8.0", "einops>=0.8.0",
"opencv-python-headless>=4.9.0", "opencv-python-headless>=4.9.0",
"av>=14.2.0", "av>=14.2.0",
"torch>=2.2.1",
"torchcodec>=0.2.1; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')",
"torchvision>=0.21.0",
"jsonlines>=4.0.0", "jsonlines>=4.0.0",
"packaging>=24.2", "packaging>=24.2",
"pynput>=1.7.7", "pynput>=1.7.7",
"pyserial>=3.5", "pyserial>=3.5",
"wandb>=0.20.0", "wandb>=0.20.0",
"torch>=2.2.1,<2.8.0", # TODO: Bumb dependency
"torchcodec>=0.2.1,<0.6.0; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')", # TODO: Bumb dependency
"torchvision>=0.21.0,<0.23.0", # TODO: Bumb dependency
"draccus==0.10.0", # TODO: Remove == "draccus==0.10.0", # TODO: Remove ==
"gymnasium>=0.29.1,<1.0.0", # TODO: Bumb dependency "gymnasium>=0.29.1,<1.0.0", # TODO: Bumb dependency
"rerun-sdk>=0.21.0,<0.23.0", # TODO: Bumb dependency "rerun-sdk>=0.21.0,<0.23.0", # TODO: Bumb dependency
@@ -125,7 +126,6 @@ hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.9", "lerobot[grpcio-dep]",
async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3"] async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3"]
# Development # Development
docs = ["hf-doc-builder @ git+https://github.com/huggingface/doc-builder.git@main", "watchdog >= 6.0.0"]
dev = ["pre-commit>=3.7.0", "debugpy>=1.8.1", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1"] dev = ["pre-commit>=3.7.0", "debugpy>=1.8.1", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1"]
test = ["pytest>=8.1.0", "pytest-timeout>=2.4.0", "pytest-cov>=5.0.0", "mock-serial>=0.0.1 ; sys_platform != 'win32'"] test = ["pytest>=8.1.0", "pytest-timeout>=2.4.0", "pytest-cov>=5.0.0", "mock-serial>=0.0.1 ; sys_platform != 'win32'"]
video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"] video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"]
@@ -147,7 +147,6 @@ all = [
"lerobot[smolvla]", "lerobot[smolvla]",
"lerobot[hilserl]", "lerobot[hilserl]",
"lerobot[async]", "lerobot[async]",
"lerobot[docs]",
"lerobot[dev]", "lerobot[dev]",
"lerobot[test]", "lerobot[test]",
"lerobot[video_benchmark]", "lerobot[video_benchmark]",
requirements-macos.txt
+625
View File
@@ -0,0 +1,625 @@
# This file is autogenerated by pip-compile with Python 3.10
# by the following command:
#
# pip-compile --output-file=requirements-macos.txt requirements.in
#
-e .[all]
# via -[all]
absl-py==2.3.1
# via
# dm-control
# dm-env
# dm-tree
# labmaze
# mujoco
accelerate==1.9.0
# via lerobot
aiohappyeyeballs==2.6.1
# via aiohttp
aiohttp==3.12.15
# via fsspec
aiosignal==1.4.0
# via aiohttp
annotated-types==0.7.0
# via pydantic
asttokens==3.0.0
# via stack-data
async-timeout==5.0.1
# via aiohttp
attrs==25.3.0
# via
# aiohttp
# dm-tree
# jsonlines
# rerun-sdk
av==15.0.0
# via lerobot
blinker==1.9.0
# via flask
certifi==2025.7.14
# via
# requests
# sentry-sdk
cffi==1.17.1
# via pymunk
cfgv==3.4.0
# via pre-commit
charset-normalizer==3.4.2
# via requests
click==8.2.1
# via
# flask
# wandb
cloudpickle==3.1.1
# via gymnasium
cmake==4.0.3
# via lerobot
cmeel==0.57.3
# via
# cmeel-assimp
# cmeel-boost
# cmeel-console-bridge
# cmeel-octomap
# cmeel-qhull
# cmeel-tinyxml2
# cmeel-urdfdom
# cmeel-zlib
# coal-library
# eigenpy
# eiquadprog
# pin
# placo
# rhoban-cmeel-jsoncpp
cmeel-assimp==5.4.3.1
# via coal-library
cmeel-boost==1.87.0.1
# via
# coal-library
# eigenpy
# eiquadprog
# pin
cmeel-console-bridge==1.0.2.3
# via cmeel-urdfdom
cmeel-octomap==1.10.0
# via coal-library
cmeel-qhull==8.0.2.1
# via coal-library
cmeel-tinyxml2==10.0.0
# via cmeel-urdfdom
cmeel-urdfdom==4.0.1
# via pin
cmeel-zlib==1.3.1
# via cmeel-assimp
coal-library==3.0.1
# via pin
contourpy==1.3.2
# via matplotlib
coverage[toml]==7.10.1
# via pytest-cov
cycler==0.12.1
# via matplotlib
datasets==3.6.0
# via lerobot
debugpy==1.8.15
# via lerobot
decorator==5.2.1
# via ipython
deepdiff==8.5.0
# via lerobot
diffusers==0.34.0
# via lerobot
dill==0.3.8
# via
# datasets
# multiprocess
distlib==0.4.0
# via virtualenv
dm-control==1.0.14
# via gym-aloha
dm-env==1.6
# via dm-control
dm-tree==0.1.9
# via
# dm-control
# dm-env
docopt==0.6.2
# via num2words
draccus==0.10.0
# via lerobot
dynamixel-sdk==3.7.31
# via lerobot
eigenpy==3.10.3
# via coal-library
einops==0.8.1
# via lerobot
eiquadprog==1.2.9
# via placo
exceptiongroup==1.3.0
# via
# ipython
# pytest
executing==2.2.0
# via stack-data
farama-notifications==0.0.4
# via gymnasium
feetech-servo-sdk==1.0.0
# via lerobot
filelock==3.18.0
# via
# datasets
# diffusers
# huggingface-hub
# torch
# transformers
# virtualenv
flask==3.1.1
# via lerobot
fonttools==4.59.0
# via matplotlib
frozenlist==1.7.0
# via
# aiohttp
# aiosignal
fsspec[http]==2025.3.0
# via
# datasets
# huggingface-hub
# torch
gitdb==4.0.12
# via gitpython
gitpython==3.1.45
# via wandb
glfw==2.9.0
# via
# dm-control
# mujoco
grpcio==1.73.1
# via
# grpcio-tools
# lerobot
grpcio-tools==1.73.1
# via lerobot
gym-aloha==0.1.1
# via lerobot
gym-hil==0.1.10
# via lerobot
gym-pusht==0.1.5
# via lerobot
gym-xarm==0.1.1
# via lerobot
gymnasium==0.29.1
# via
# gym-aloha
# gym-hil
# gym-pusht
# gym-xarm
# gymnasium-robotics
# lerobot
# pettingzoo
gymnasium-robotics==1.2.4
# via gym-xarm
hf-transfer==0.1.9
# via huggingface-hub
hf-xet==1.1.5
# via huggingface-hub
hidapi==0.14.0.post4
# via
# gym-hil
# lerobot
huggingface-hub[cli,hf-transfer]==0.34.3
# via
# accelerate
# datasets
# diffusers
# lerobot
# tokenizers
# transformers
identify==2.6.12
# via pre-commit
idna==3.10
# via
# requests
# yarl
imageio[ffmpeg]==2.37.0
# via
# gym-aloha
# gym-hil
# gymnasium-robotics
# lerobot
# scikit-image
imageio-ffmpeg==0.6.0
# via imageio
importlib-metadata==8.7.0
# via diffusers
iniconfig==2.1.0
# via pytest
inquirerpy==0.3.4
# via huggingface-hub
ipython==8.37.0
# via meshcat
ischedule==1.2.7
# via placo
itsdangerous==2.2.0
# via flask
jedi==0.19.2
# via ipython
jinja2==3.1.6
# via
# flask
# gymnasium-robotics
# torch
jsonlines==4.0.0
# via lerobot
kiwisolver==1.4.8
# via matplotlib
labmaze==1.0.6
# via dm-control
lazy-loader==0.4
# via scikit-image
lxml==6.0.0
# via dm-control
markupsafe==3.0.2
# via
# flask
# jinja2
# werkzeug
matplotlib==3.10.5
# via lerobot
matplotlib-inline==0.1.7
# via ipython
mergedeep==1.3.4
# via draccus
meshcat==0.3.2
# via placo
mock-serial==0.0.1
# via lerobot
mpmath==1.3.0
# via sympy
mujoco==2.3.7
# via
# dm-control
# gym-aloha
# gym-hil
# gym-xarm
# gymnasium-robotics
multidict==6.6.3
# via
# aiohttp
# yarl
multiprocess==0.70.16
# via datasets
mypy-extensions==1.1.0
# via typing-inspect
networkx==3.4.2
# via
# scikit-image
# torch
nodeenv==1.9.1
# via pre-commit
num2words==0.5.14
# via lerobot
numpy==2.2.6
# via
# accelerate
# cmeel-boost
# contourpy
# datasets
# diffusers
# dm-control
# dm-env
# dm-tree
# gymnasium
# gymnasium-robotics
# imageio
# labmaze
# matplotlib
# meshcat
# mujoco
# opencv-python
# opencv-python-headless
# pandas
# pettingzoo
# rerun-sdk
# scikit-image
# scipy
# shapely
# tifffile
# torchvision
# transformers
opencv-python==4.12.0.88
# via gym-pusht
opencv-python-headless==4.12.0.88
# via lerobot
orderly-set==5.5.0
# via deepdiff
packaging==25.0
# via
# accelerate
# datasets
# huggingface-hub
# lazy-loader
# lerobot
# matplotlib
# pytest
# scikit-image
# transformers
# wandb
pandas==2.3.1
# via
# datasets
# lerobot
parso==0.8.4
# via jedi
pettingzoo==1.24.3
# via gymnasium-robotics
pexpect==4.9.0
# via ipython
pfzy==0.3.4
# via inquirerpy
pillow==11.3.0
# via
# diffusers
# imageio
# matplotlib
# meshcat
# rerun-sdk
# scikit-image
# torchvision
pin==3.4.0
# via placo
placo==0.9.14
# via lerobot
platformdirs==4.3.8
# via
# virtualenv
# wandb
pluggy==1.6.0
# via
# pytest
# pytest-cov
pre-commit==4.2.0
# via lerobot
prompt-toolkit==3.0.51
# via
# inquirerpy
# ipython
propcache==0.3.2
# via
# aiohttp
# yarl
protobuf==6.31.0
# via
# dm-control
# grpcio-tools
# lerobot
# wandb
psutil==7.0.0
# via
# accelerate
# imageio
ptyprocess==0.7.0
# via pexpect
pure-eval==0.2.3
# via stack-data
pyarrow==21.0.0
# via
# datasets
# rerun-sdk
pycparser==2.22
# via cffi
pydantic==2.11.7
# via wandb
pydantic-core==2.33.2
# via pydantic
pygame==2.6.1
# via
# gym-hil
# gym-pusht
# lerobot
pygments==2.19.2
# via
# ipython
# pytest
pymunk==6.11.1
# via
# gym-pusht
# lerobot
pyngrok==7.2.12
# via meshcat
pynput==1.8.1
# via
# gym-hil
# lerobot
pyobjc-core==11.1
# via
# pyobjc-framework-applicationservices
# pyobjc-framework-cocoa
# pyobjc-framework-coretext
# pyobjc-framework-quartz
pyobjc-framework-applicationservices==11.1
# via pynput
pyobjc-framework-cocoa==11.1
# via
# pyobjc-framework-applicationservices
# pyobjc-framework-coretext
# pyobjc-framework-quartz
pyobjc-framework-coretext==11.1
# via pyobjc-framework-applicationservices
pyobjc-framework-quartz==11.1
# via
# pynput
# pyobjc-framework-applicationservices
# pyobjc-framework-coretext
pyopengl==3.1.9
# via
# dm-control
# mujoco
pyparsing==3.2.3
# via
# dm-control
# matplotlib
pyrealsense2-macosx==2.54.2
# via lerobot
pyserial==3.5
# via
# dynamixel-sdk
# feetech-servo-sdk
# lerobot
pytest==8.4.1
# via
# lerobot
# pytest-cov
# pytest-timeout
pytest-cov==6.2.1
# via lerobot
pytest-timeout==2.4.0
# via lerobot
python-dateutil==2.9.0.post0
# via
# matplotlib
# pandas
pytz==2025.2
# via pandas
pyyaml==6.0.2
# via
# accelerate
# datasets
# draccus
# huggingface-hub
# pre-commit
# pyngrok
# pyyaml-include
# transformers
# wandb
pyyaml-include==1.4.1
# via draccus
pyzmq==27.0.0
# via
# lerobot
# meshcat
regex==2025.7.34
# via
# diffusers
# transformers
requests==2.32.4
# via
# datasets
# diffusers
# dm-control
# huggingface-hub
# transformers
# wandb
rerun-sdk==0.22.1
# via lerobot
rhoban-cmeel-jsoncpp==1.9.4.9
# via placo
safetensors==0.5.3
# via
# accelerate
# diffusers
# lerobot
# transformers
scikit-image==0.25.2
# via
# gym-pusht
# lerobot
scipy==1.15.3
# via
# dm-control
# scikit-image
sentry-sdk==2.34.1
# via wandb
shapely==2.1.1
# via gym-pusht
six==1.17.0
# via
# pynput
# python-dateutil
smmap==5.0.2
# via gitdb
stack-data==0.6.3
# via ipython
sympy==1.14.0
# via torch
termcolor==3.1.0
# via lerobot
tifffile==2025.5.10
# via scikit-image
tokenizers==0.21.4
# via transformers
toml==0.10.2
# via draccus
tomli==2.2.1
# via
# cmeel
# coverage
# pytest
torch==2.7.1
# via
# accelerate
# lerobot
# torchvision
torchcodec==0.5
# via lerobot
torchvision==0.22.1
# via lerobot
tornado==6.5.1
# via meshcat
tqdm==4.67.1
# via
# datasets
# dm-control
# huggingface-hub
# transformers
traitlets==5.14.3
# via
# ipython
# matplotlib-inline
transformers==4.51.3
# via lerobot
typing-extensions==4.14.1
# via
# aiosignal
# exceptiongroup
# gymnasium
# huggingface-hub
# ipython
# multidict
# pydantic
# pydantic-core
# rerun-sdk
# torch
# typing-inspect
# typing-inspection
# wandb
typing-inspect==0.9.0
# via draccus
typing-inspection==0.4.1
# via pydantic
tzdata==2025.2
# via pandas
u-msgpack-python==2.8.0
# via meshcat
urllib3==2.5.0
# via
# requests
# sentry-sdk
virtualenv==20.32.0
# via pre-commit
wandb==0.21.0
# via lerobot
wcwidth==0.2.13
# via prompt-toolkit
werkzeug==3.1.3
# via flask
wrapt==1.17.2
# via dm-tree
xxhash==3.5.0
# via datasets
yarl==1.20.1
# via aiohttp
zipp==3.23.0
# via importlib-metadata
# The following packages are considered to be unsafe in a requirements file:
# setuptools
requirements-ubuntu.txt
+650
View File
@@ -0,0 +1,650 @@
#
# This file is autogenerated by pip-compile with Python 3.10
# by the following command:
#
# pip-compile --output-file=requirements-ubuntu.txt requirements.in
#
-e .[all]
# via -[all]
absl-py==2.3.1
# via
# dm-control
# dm-env
# dm-tree
# labmaze
# mujoco
accelerate==1.9.0
# via lerobot
aiohappyeyeballs==2.6.1
# via aiohttp
aiohttp==3.12.15
# via fsspec
aiosignal==1.4.0
# via aiohttp
annotated-types==0.7.0
# via pydantic
asttokens==3.0.0
# via stack-data
async-timeout==5.0.1
# via aiohttp
attrs==25.3.0
# via
# aiohttp
# dm-tree
# jsonlines
# rerun-sdk
av==15.0.0
# via lerobot
blinker==1.9.0
# via flask
certifi==2025.7.14
# via
# requests
# sentry-sdk
cffi==1.17.1
# via pymunk
cfgv==3.4.0
# via pre-commit
charset-normalizer==3.4.2
# via requests
click==8.2.1
# via
# flask
# wandb
cloudpickle==3.1.1
# via gymnasium
cmake==4.0.3
# via lerobot
cmeel==0.57.3
# via
# cmeel-assimp
# cmeel-boost
# cmeel-console-bridge
# cmeel-octomap
# cmeel-qhull
# cmeel-tinyxml2
# cmeel-urdfdom
# cmeel-zlib
# coal-library
# eigenpy
# eiquadprog
# pin
# placo
# rhoban-cmeel-jsoncpp
cmeel-assimp==5.4.3.1
# via coal-library
cmeel-boost==1.87.0.1
# via
# coal-library
# eigenpy
# eiquadprog
# pin
cmeel-console-bridge==1.0.2.3
# via cmeel-urdfdom
cmeel-octomap==1.10.0
# via coal-library
cmeel-qhull==8.0.2.1
# via coal-library
cmeel-tinyxml2==10.0.0
# via cmeel-urdfdom
cmeel-urdfdom==4.0.1
# via pin
cmeel-zlib==1.3.1
# via cmeel-assimp
coal-library==3.0.1
# via pin
contourpy==1.3.2
# via matplotlib
coverage[toml]==7.10.1
# via pytest-cov
cycler==0.12.1
# via matplotlib
datasets==3.6.0
# via lerobot
debugpy==1.8.15
# via lerobot
decorator==5.2.1
# via ipython
deepdiff==8.5.0
# via lerobot
diffusers==0.34.0
# via lerobot
dill==0.3.8
# via
# datasets
# multiprocess
distlib==0.4.0
# via virtualenv
dm-control==1.0.14
# via gym-aloha
dm-env==1.6
# via dm-control
dm-tree==0.1.9
# via
# dm-control
# dm-env
docopt==0.6.2
# via num2words
draccus==0.10.0
# via lerobot
dynamixel-sdk==3.7.31
# via lerobot
eigenpy==3.10.3
# via coal-library
einops==0.8.1
# via lerobot
eiquadprog==1.2.9
# via placo
evdev==1.9.2
# via pynput
exceptiongroup==1.3.0
# via
# ipython
# pytest
executing==2.2.0
# via stack-data
farama-notifications==0.0.4
# via gymnasium
feetech-servo-sdk==1.0.0
# via lerobot
filelock==3.18.0
# via
# datasets
# diffusers
# huggingface-hub
# torch
# transformers
# virtualenv
flask==3.1.1
# via lerobot
fonttools==4.59.0
# via matplotlib
frozenlist==1.7.0
# via
# aiohttp
# aiosignal
fsspec[http]==2025.3.0
# via
# datasets
# huggingface-hub
# torch
gitdb==4.0.12
# via gitpython
gitpython==3.1.45
# via wandb
glfw==2.9.0
# via
# dm-control
# mujoco
grpcio==1.73.1
# via
# grpcio-tools
# lerobot
grpcio-tools==1.73.1
# via lerobot
gym-aloha==0.1.1
# via lerobot
gym-hil==0.1.10
# via lerobot
gym-pusht==0.1.5
# via lerobot
gym-xarm==0.1.1
# via lerobot
gymnasium==0.29.1
# via
# gym-aloha
# gym-hil
# gym-pusht
# gym-xarm
# gymnasium-robotics
# lerobot
# pettingzoo
gymnasium-robotics==1.2.4
# via gym-xarm
hf-transfer==0.1.9
# via huggingface-hub
hf-xet==1.1.5
# via huggingface-hub
hidapi==0.14.0.post4
# via
# gym-hil
# lerobot
huggingface-hub[cli,hf-transfer]==0.34.3
# via
# accelerate
# datasets
# diffusers
# lerobot
# tokenizers
# transformers
identify==2.6.12
# via pre-commit
idna==3.10
# via
# requests
# yarl
imageio[ffmpeg]==2.37.0
# via
# gym-aloha
# gym-hil
# gymnasium-robotics
# lerobot
# scikit-image
imageio-ffmpeg==0.6.0
# via imageio
importlib-metadata==8.7.0
# via diffusers
iniconfig==2.1.0
# via pytest
inquirerpy==0.3.4
# via huggingface-hub
ipython==8.37.0
# via meshcat
ischedule==1.2.7
# via placo
itsdangerous==2.2.0
# via flask
jedi==0.19.2
# via ipython
jinja2==3.1.6
# via
# flask
# gymnasium-robotics
# torch
jsonlines==4.0.0
# via lerobot
kiwisolver==1.4.8
# via matplotlib
labmaze==1.0.6
# via dm-control
lazy-loader==0.4
# via scikit-image
lxml==6.0.0
# via dm-control
markupsafe==3.0.2
# via
# flask
# jinja2
# werkzeug
matplotlib==3.10.5
# via lerobot
matplotlib-inline==0.1.7
# via ipython
mergedeep==1.3.4
# via draccus
meshcat==0.3.2
# via placo
mock-serial==0.0.1
# via lerobot
mpmath==1.3.0
# via sympy
mujoco==2.3.7
# via
# dm-control
# gym-aloha
# gym-hil
# gym-xarm
# gymnasium-robotics
multidict==6.6.3
# via
# aiohttp
# yarl
multiprocess==0.70.16
# via datasets
mypy-extensions==1.1.0
# via typing-inspect
networkx==3.4.2
# via
# scikit-image
# torch
nodeenv==1.9.1
# via pre-commit
num2words==0.5.14
# via lerobot
numpy==2.2.6
# via
# accelerate
# cmeel-boost
# contourpy
# datasets
# diffusers
# dm-control
# dm-env
# dm-tree
# gymnasium
# gymnasium-robotics
# imageio
# labmaze
# matplotlib
# meshcat
# mujoco
# opencv-python
# opencv-python-headless
# pandas
# pettingzoo
# rerun-sdk
# scikit-image
# scipy
# shapely
# tifffile
# torchvision
# transformers
nvidia-cublas-cu12==12.6.4.1
# via
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
# torch
nvidia-cuda-cupti-cu12==12.6.80
# via torch
nvidia-cuda-nvrtc-cu12==12.6.77
# via torch
nvidia-cuda-runtime-cu12==12.6.77
# via torch
nvidia-cudnn-cu12==9.5.1.17
# via torch
nvidia-cufft-cu12==11.3.0.4
# via torch
nvidia-cufile-cu12==1.11.1.6
# via torch
nvidia-curand-cu12==10.3.7.77
# via torch
nvidia-cusolver-cu12==11.7.1.2
# via torch
nvidia-cusparse-cu12==12.5.4.2
# via
# nvidia-cusolver-cu12
# torch
nvidia-cusparselt-cu12==0.6.3
# via torch
nvidia-nccl-cu12==2.26.2
# via torch
nvidia-nvjitlink-cu12==12.6.85
# via
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
# torch
nvidia-nvtx-cu12==12.6.77
# via torch
opencv-python==4.12.0.88
# via gym-pusht
opencv-python-headless==4.12.0.88
# via lerobot
orderly-set==5.5.0
# via deepdiff
packaging==25.0
# via
# accelerate
# datasets
# huggingface-hub
# lazy-loader
# lerobot
# matplotlib
# pytest
# scikit-image
# transformers
# wandb
pandas==2.3.1
# via
# datasets
# lerobot
parso==0.8.4
# via jedi
pettingzoo==1.24.3
# via gymnasium-robotics
pexpect==4.9.0
# via ipython
pfzy==0.3.4
# via inquirerpy
pillow==11.3.0
# via
# diffusers
# imageio
# matplotlib
# meshcat
# rerun-sdk
# scikit-image
# torchvision
pin==3.4.0
# via placo
placo==0.9.14
# via lerobot
platformdirs==4.3.8
# via
# virtualenv
# wandb
pluggy==1.6.0
# via
# pytest
# pytest-cov
pre-commit==4.2.0
# via lerobot
prompt-toolkit==3.0.51
# via
# inquirerpy
# ipython
propcache==0.3.2
# via
# aiohttp
# yarl
protobuf==6.31.0
# via
# dm-control
# grpcio-tools
# lerobot
# wandb
psutil==7.0.0
# via
# accelerate
# imageio
ptyprocess==0.7.0
# via pexpect
pure-eval==0.2.3
# via stack-data
pyarrow==21.0.0
# via
# datasets
# rerun-sdk
pycparser==2.22
# via cffi
pydantic==2.11.7
# via wandb
pydantic-core==2.33.2
# via pydantic
pygame==2.6.1
# via
# gym-hil
# gym-pusht
# lerobot
pygments==2.19.2
# via
# ipython
# pytest
pymunk==6.11.1
# via
# gym-pusht
# lerobot
pyngrok==7.2.12
# via meshcat
pynput==1.8.1
# via
# gym-hil
# lerobot
pyopengl==3.1.9
# via
# dm-control
# mujoco
pyparsing==3.2.3
# via
# dm-control
# matplotlib
pyrealsense2==2.56.5.9235
# via lerobot
pyserial==3.5
# via
# dynamixel-sdk
# feetech-servo-sdk
# lerobot
pytest==8.4.1
# via
# lerobot
# pytest-cov
# pytest-timeout
pytest-cov==6.2.1
# via lerobot
pytest-timeout==2.4.0
# via lerobot
python-dateutil==2.9.0.post0
# via
# matplotlib
# pandas
python-xlib==0.33
# via pynput
pytz==2025.2
# via pandas
pyyaml==6.0.2
# via
# accelerate
# datasets
# draccus
# huggingface-hub
# pre-commit
# pyngrok
# pyyaml-include
# transformers
# wandb
pyyaml-include==1.4.1
# via draccus
pyzmq==27.0.0
# via
# lerobot
# meshcat
regex==2025.7.34
# via
# diffusers
# transformers
requests==2.32.4
# via
# datasets
# diffusers
# dm-control
# huggingface-hub
# transformers
# wandb
rerun-sdk==0.22.1
# via lerobot
rhoban-cmeel-jsoncpp==1.9.4.9
# via placo
safetensors==0.5.3
# via
# accelerate
# diffusers
# lerobot
# transformers
scikit-image==0.25.2
# via
# gym-pusht
# lerobot
scipy==1.15.3
# via
# dm-control
# scikit-image
sentry-sdk==2.34.1
# via wandb
shapely==2.1.1
# via gym-pusht
six==1.17.0
# via
# pynput
# python-dateutil
# python-xlib
smmap==5.0.2
# via gitdb
stack-data==0.6.3
# via ipython
sympy==1.14.0
# via torch
termcolor==3.1.0
# via lerobot
tifffile==2025.5.10
# via scikit-image
tokenizers==0.21.4
# via transformers
toml==0.10.2
# via draccus
tomli==2.2.1
# via
# cmeel
# coverage
# pytest
torch==2.7.1
# via
# accelerate
# lerobot
# torchvision
torchcodec==0.5
# via lerobot
torchvision==0.22.1
# via lerobot
tornado==6.5.1
# via meshcat
tqdm==4.67.1
# via
# datasets
# dm-control
# huggingface-hub
# transformers
traitlets==5.14.3
# via
# ipython
# matplotlib-inline
transformers==4.51.3
# via lerobot
triton==3.3.1
# via torch
typing-extensions==4.14.1
# via
# aiosignal
# exceptiongroup
# gymnasium
# huggingface-hub
# ipython
# multidict
# pydantic
# pydantic-core
# rerun-sdk
# torch
# typing-inspect
# typing-inspection
# wandb
typing-inspect==0.9.0
# via draccus
typing-inspection==0.4.1
# via pydantic
tzdata==2025.2
# via pandas
u-msgpack-python==2.8.0
# via meshcat
urllib3==2.5.0
# via
# requests
# sentry-sdk
virtualenv==20.32.0
# via pre-commit
wandb==0.21.0
# via lerobot
wcwidth==0.2.13
# via prompt-toolkit
werkzeug==3.1.3
# via flask
wrapt==1.17.2
# via dm-tree
xxhash==3.5.0
# via datasets
yarl==1.20.1
# via aiohttp
zipp==3.23.0
# via importlib-metadata
# The following packages are considered to be unsafe in a requirements file:
# setuptools
requirements.in
+9
View File
@@ -0,0 +1,9 @@
# requirements.in
# requirements-macos.txt was generated on macOS and is platform-specific (macOS 15.5 24F74 arm64).
# Darwin MacBook-Pro.local 24.5.0 Darwin Kernel Version 24.5.0: Tue Apr 22 19:54:43 PDT 2025; root:xnu-11417.121.6~2/RELEASE_ARM64_T8132 arm64
# requirements-ubuntu.txt was generated on Linux and is platform-specific (Ubuntu 24.04.2 LTS x86_64).
# Linux mlerobot-linux 6.14.0-27-generic #27~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC Tue Jul 22 17:38:49 UTC 2 x86_64 x86_64 x86_64 GNU/Linux
-e .[all]
src/lerobot/calibrate.py
+6 -2
View File
@@ -18,7 +18,7 @@ Helper to recalibrate your device (robot or teleoperator).
Example: Example:
```shell ```shell
python -m lerobot.calibrate \ lerobot-calibrate \
--teleop.type=so100_leader \ --teleop.type=so100_leader \
--teleop.port=/dev/tty.usbmodem58760431551 \ --teleop.port=/dev/tty.usbmodem58760431551 \
--teleop.id=blue --teleop.id=blue
@@ -82,5 +82,9 @@ def calibrate(cfg: CalibrateConfig):
device.disconnect() device.disconnect()
if __name__ == "__main__": def main():
calibrate() calibrate()
if __name__ == "__main__":
main()
src/lerobot/cameras/opencv/camera_opencv.py
+2 -3
View File
@@ -60,7 +60,7 @@ class OpenCVCamera(Camera):
or port changes, especially on Linux. Use the provided utility script to find or port changes, especially on Linux. Use the provided utility script to find
available camera indices or paths: available camera indices or paths:
```bash ```bash
python -m lerobot.find_cameras opencv lerobot-find-cameras opencv
``` ```
The camera's default settings (FPS, resolution, color mode) are used unless The camera's default settings (FPS, resolution, color mode) are used unless
@@ -165,8 +165,7 @@ class OpenCVCamera(Camera):
self.videocapture.release() self.videocapture.release()
self.videocapture = None self.videocapture = None
raise ConnectionError( raise ConnectionError(
f"Failed to open {self}." f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
f"Run `python -m lerobot.find_cameras opencv` to find available cameras."
) )
self._configure_capture_settings() self._configure_capture_settings()
src/lerobot/cameras/realsense/camera_realsense.py
+2 -3
View File
@@ -51,7 +51,7 @@ class RealSenseCamera(Camera):
Use the provided utility script to find available camera indices and default profiles: Use the provided utility script to find available camera indices and default profiles:
```bash ```bash
python -m lerobot.find_cameras realsense lerobot-find-cameras realsense
``` ```
A `RealSenseCamera` instance requires a configuration object specifying the A `RealSenseCamera` instance requires a configuration object specifying the
@@ -176,8 +176,7 @@ class RealSenseCamera(Camera):
self.rs_profile = None self.rs_profile = None
self.rs_pipeline = None self.rs_pipeline = None
raise ConnectionError( raise ConnectionError(
f"Failed to open {self}." f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
"Run `python -m lerobot.find_cameras realsense` to find available cameras."
) from e ) from e
self._configure_capture_settings() self._configure_capture_settings()
src/lerobot/configs/policies.py
+5 -4
View File
@@ -27,6 +27,7 @@ from huggingface_hub.constants import CONFIG_NAME
from huggingface_hub.errors import HfHubHTTPError from huggingface_hub.errors import HfHubHTTPError
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_STATE
from lerobot.optim.optimizers import OptimizerConfig from lerobot.optim.optimizers import OptimizerConfig
from lerobot.optim.schedulers import LRSchedulerConfig from lerobot.optim.schedulers import LRSchedulerConfig
from lerobot.utils.hub import HubMixin from lerobot.utils.hub import HubMixin
@@ -119,8 +120,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
@property @property
def robot_state_feature(self) -> PolicyFeature | None: def robot_state_feature(self) -> PolicyFeature | None:
for _, ft in self.input_features.items(): for ft_name, ft in self.input_features.items():
if ft.type is FeatureType.STATE: if ft.type is FeatureType.STATE and ft_name == OBS_STATE:
return ft return ft
return None return None
@@ -137,8 +138,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
@property @property
def action_feature(self) -> PolicyFeature | None: def action_feature(self) -> PolicyFeature | None:
for _, ft in self.output_features.items(): for ft_name, ft in self.output_features.items():
if ft.type is FeatureType.ACTION: if ft.type is FeatureType.ACTION and ft_name == ACTION:
return ft return ft
return None return None
src/lerobot/find_cameras.py
+6 -2
View File
@@ -20,7 +20,7 @@ Helper to find the camera devices available in your system.
Example: Example:
```shell ```shell
python -m lerobot.find_cameras lerobot-find-cameras
``` ```
""" """
@@ -286,7 +286,7 @@ def save_images_from_all_cameras(
print(f"Image capture finished. Images saved to {output_dir}") print(f"Image capture finished. Images saved to {output_dir}")
if __name__ == "__main__": def main():
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Unified camera utility script for listing cameras and capturing images." description="Unified camera utility script for listing cameras and capturing images."
) )
@@ -313,3 +313,7 @@ if __name__ == "__main__":
) )
args = parser.parse_args() args = parser.parse_args()
save_images_from_all_cameras(**vars(args)) save_images_from_all_cameras(**vars(args))
if __name__ == "__main__":
main()
src/lerobot/find_port.py
+6 -2
View File
@@ -18,7 +18,7 @@ Helper to find the USB port associated with your MotorsBus.
Example: Example:
```shell ```shell
python -m lerobot.find_port lerobot-find-port
``` ```
""" """
@@ -61,5 +61,9 @@ def find_port():
raise OSError(f"Could not detect the port. More than one port was found ({ports_diff}).") raise OSError(f"Could not detect the port. More than one port was found ({ports_diff}).")
if __name__ == "__main__": def main():
find_port() find_port()
if __name__ == "__main__":
main()
src/lerobot/motors/motors_bus.py
+2 -2
View File
@@ -222,7 +222,7 @@ class MotorsBus(abc.ABC):
A MotorsBus subclass instance requires a port (e.g. `FeetechMotorsBus(port="/dev/tty.usbmodem575E0031751"`)). A MotorsBus subclass instance requires a port (e.g. `FeetechMotorsBus(port="/dev/tty.usbmodem575E0031751"`)).
To find the port, you can run our utility script: To find the port, you can run our utility script:
```bash ```bash
python -m lerobot.find_port.py lerobot-find-port.py
>>> Finding all available ports for the MotorsBus. >>> Finding all available ports for the MotorsBus.
>>> ["/dev/tty.usbmodem575E0032081", "/dev/tty.usbmodem575E0031751"] >>> ["/dev/tty.usbmodem575E0032081", "/dev/tty.usbmodem575E0031751"]
>>> Remove the usb cable from your MotorsBus and press Enter when done. >>> Remove the usb cable from your MotorsBus and press Enter when done.
@@ -446,7 +446,7 @@ class MotorsBus(abc.ABC):
except (FileNotFoundError, OSError, serial.SerialException) as e: except (FileNotFoundError, OSError, serial.SerialException) as e:
raise ConnectionError( raise ConnectionError(
f"\nCould not connect on port '{self.port}'. Make sure you are using the correct port." f"\nCould not connect on port '{self.port}'. Make sure you are using the correct port."
"\nTry running `python -m lerobot.find_port`\n" "\nTry running `lerobot-find-port`\n"
) from e ) from e
@abc.abstractmethod @abc.abstractmethod
src/lerobot/policies/act/README.md
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../docs/source/policy_act_README.md
src/lerobot/policies/diffusion/README.md
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../docs/source/policy_diffusion_README.md
src/lerobot/policies/pi0/modeling_pi0.py
+97 -3
View File
@@ -30,7 +30,7 @@ pip install -e ".[pi0]"
Example of finetuning the pi0 pretrained model (`pi0_base` in `openpi`): Example of finetuning the pi0 pretrained model (`pi0_base` in `openpi`):
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.path=lerobot/pi0 \ --policy.path=lerobot/pi0 \
--dataset.repo_id=danaaubakirova/koch_test --dataset.repo_id=danaaubakirova/koch_test
``` ```
@@ -38,7 +38,7 @@ python -m lerobot.scripts.train \
Example of finetuning the pi0 neural network with PaliGemma and expert Gemma Example of finetuning the pi0 neural network with PaliGemma and expert Gemma
pretrained with VLM default parameters before pi0 finetuning: pretrained with VLM default parameters before pi0 finetuning:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=pi0 \ --policy.type=pi0 \
--dataset.repo_id=danaaubakirova/koch_test --dataset.repo_id=danaaubakirova/koch_test
``` ```
@@ -66,7 +66,8 @@ from lerobot.policies.pi0.paligemma_with_expert import (
PaliGemmaWithExpertModel, PaliGemmaWithExpertModel,
) )
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.utils.utils import get_safe_dtype from lerobot.policies.utils import log_model_loading_keys
from lerobot.utils.utils import get_safe_dtype, init_logging
def create_sinusoidal_pos_embedding( def create_sinusoidal_pos_embedding(
@@ -252,6 +253,99 @@ class PI0Policy(PreTrainedPolicy):
"""This should be called whenever the environment is reset.""" """This should be called whenever the environment is reset."""
self._action_queue = deque([], maxlen=self.config.n_action_steps) self._action_queue = deque([], maxlen=self.config.n_action_steps)
@classmethod
def _transform_state_dict_keys(cls, state_dict: dict) -> dict:
"""
Transform state dict keys to match expected model structure.
Transformations:
- model.paligemma_with_expert.paligemma.language_model.lm_head ->
model.paligemma_with_expert.paligemma.lm_head
- model.paligemma_with_expert.paligemma.language_model.model ->
model.paligemma_with_expert.paligemma.model.language_model
- model.paligemma_with_expert.paligemma.vision_tower ->
model.paligemma_with_expert.paligemma.model.vision_tower
- model.paligemma_with_expert.paligemma.multi_modal_projector ->
model.paligemma_with_expert.paligemma.model.multi_modal_projector
Also handles tied weights between lm_head.weight and
embed_tokens.weight.
"""
import re
transformed_dict = {}
transformations = [
(
re.compile(r"\.paligemma_with_expert\.paligemma\.language_model\.lm_head"),
".paligemma_with_expert.paligemma.lm_head",
),
(
re.compile(r"\.paligemma_with_expert\.paligemma\.language_model\.model"),
".paligemma_with_expert.paligemma.model.language_model",
),
(
re.compile(r"\.paligemma_with_expert\.paligemma\.vision_tower"),
".paligemma_with_expert.paligemma.model.vision_tower",
),
(
re.compile(r"\.paligemma_with_expert\.paligemma\.multi_modal_projector"),
".paligemma_with_expert.paligemma.model.multi_modal_projector",
),
]
for key, value in state_dict.items():
new_key = key
for pattern, replacement in transformations:
new_key = pattern.sub(replacement, new_key)
transformed_dict[new_key] = value
# Handle tied weights: lm_head.weight and embed_tokens.weight share memory
lm_head_key = None
embed_tokens_key = None
for key in transformed_dict:
if key.endswith(".paligemma_with_expert.paligemma.lm_head.weight"):
lm_head_key = key
elif key.endswith(".paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"):
embed_tokens_key = key
if lm_head_key and embed_tokens_key:
break
if lm_head_key and not embed_tokens_key:
embed_tokens_key = lm_head_key.replace(
".lm_head.weight", ".model.language_model.embed_tokens.weight"
)
transformed_dict[embed_tokens_key] = transformed_dict[lm_head_key]
elif embed_tokens_key and not lm_head_key:
lm_head_key = embed_tokens_key.replace(
".model.language_model.embed_tokens.weight", ".lm_head.weight"
)
transformed_dict[lm_head_key] = transformed_dict[embed_tokens_key]
return transformed_dict
@classmethod
def _load_as_safetensor(
cls, model: "PI0Policy", model_file: str, map_location: str, strict: bool
) -> "PI0Policy":
"""Override to apply key transformations before loading."""
from safetensors.torch import load_file
init_logging()
# Load the state dict from file safely
state_dict = load_file(model_file, device=map_location)
# Apply key transformations
transformed_state_dict = cls._transform_state_dict_keys(state_dict)
# Load the transformed state dict
msg = model.load_state_dict(transformed_state_dict, strict=strict)
# Log message
log_model_loading_keys(msg.missing_keys, msg.unexpected_keys)
return model
def get_optim_params(self) -> dict: def get_optim_params(self) -> dict:
return self.parameters() return self.parameters()
src/lerobot/policies/pi0fast/modeling_pi0fast.py
+2 -2
View File
@@ -25,14 +25,14 @@ Disclaimer: It is not expected to perform as well as the original implementation
Example of finetuning the pi0+FAST pretrained model (`pi0_fast_base` in `openpi`): Example of finetuning the pi0+FAST pretrained model (`pi0_fast_base` in `openpi`):
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.path=lerobot/pi0fast_base \ --policy.path=lerobot/pi0fast_base \
--dataset.repo_id=danaaubakirova/koch_test --dataset.repo_id=danaaubakirova/koch_test
``` ```
Example of training the pi0+FAST neural network with from scratch: Example of training the pi0+FAST neural network with from scratch:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=pi0fast \ --policy.type=pi0fast \
--dataset.repo_id=danaaubakirova/koch_test --dataset.repo_id=danaaubakirova/koch_test
``` ```
src/lerobot/policies/pretrained.py
+21 -12
View File
@@ -30,6 +30,7 @@ from torch import Tensor, nn
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig from lerobot.configs.train import TrainPipelineConfig
from lerobot.policies.utils import log_model_loading_keys
from lerobot.utils.hub import HubMixin from lerobot.utils.hub import HubMixin
T = TypeVar("T", bound="PreTrainedPolicy") T = TypeVar("T", bound="PreTrainedPolicy")
@@ -128,18 +129,26 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
@classmethod @classmethod
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T: def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
if packaging.version.parse(safetensors.__version__) < packaging.version.parse("0.4.3"): # Create base kwargs
load_model_as_safetensor(model, model_file, strict=strict) kwargs = {"strict": strict}
if map_location != "cpu":
logging.warning( # Add device parameter for newer versions that support it
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors." if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
" This means that the model is loaded on 'cpu' first and then copied to the device." kwargs["device"] = map_location
" This leads to a slower loading time."
" Please update safetensors to version 0.4.3 or above for improved performance." # Load the model with appropriate kwargs
) missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
model.to(map_location) log_model_loading_keys(missing_keys, unexpected_keys)
else:
safetensors.torch.load_model(model, model_file, strict=strict, device=map_location) # For older versions, manually move to device if needed
if "device" not in kwargs and map_location != "cpu":
logging.warning(
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
" This means that the model is loaded on 'cpu' first and then copied to the device."
" This leads to a slower loading time."
" Please update safetensors to version 0.4.3 or above for improved performance."
)
model.to(map_location)
return model return model
@abc.abstractmethod @abc.abstractmethod
src/lerobot/policies/smolvla/README.md
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../docs/source/policy_smolvla_README.md
src/lerobot/policies/smolvla/modeling_smolvla.py
+2 -2
View File
@@ -28,7 +28,7 @@ pip install -e ".[smolvla]"
Example of finetuning the smolvla pretrained model (`smolvla_base`): Example of finetuning the smolvla pretrained model (`smolvla_base`):
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.path=lerobot/smolvla_base \ --policy.path=lerobot/smolvla_base \
--dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \ --dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \
--batch_size=64 \ --batch_size=64 \
@@ -38,7 +38,7 @@ python -m lerobot.scripts.train \
Example of finetuning a smolVLA. SmolVLA is composed of a pretrained VLM, Example of finetuning a smolVLA. SmolVLA is composed of a pretrained VLM,
and an action expert. and an action expert.
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--policy.type=smolvla \ --policy.type=smolvla \
--dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \ --dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \
--batch_size=64 \ --batch_size=64 \
src/lerobot/policies/tdmpc/README.md
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../docs/source/policy_tdmpc_README.md
src/lerobot/policies/utils.py
+14
View File
@@ -14,6 +14,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import logging
from collections import deque from collections import deque
import torch import torch
@@ -71,3 +72,16 @@ def get_output_shape(module: nn.Module, input_shape: tuple) -> tuple:
with torch.inference_mode(): with torch.inference_mode():
output = module(dummy_input) output = module(dummy_input)
return tuple(output.shape) return tuple(output.shape)
def log_model_loading_keys(missing_keys: list[str], unexpected_keys: list[str]) -> None:
"""Log missing and unexpected keys when loading a model.
Args:
missing_keys (list[str]): Keys that were expected but not found.
unexpected_keys (list[str]): Keys that were found but not expected.
"""
if missing_keys:
logging.warning(f"Missing key(s) when loading model: {missing_keys}")
if unexpected_keys:
logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")
src/lerobot/policies/vqbet/README.md
Symlink
+1
View File
@@ -0,0 +1 @@
../../../../docs/source/policy_vqbet_README.md
src/lerobot/processor/__init__.py
+54
View File
@@ -0,0 +1,54 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from .device_processor import DeviceProcessor
from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor
from .observation_processor import VanillaObservationProcessor
from .pipeline import (
ActionProcessor,
DoneProcessor,
EnvTransition,
IdentityProcessor,
InfoProcessor,
ObservationProcessor,
ProcessorStep,
ProcessorStepRegistry,
RewardProcessor,
RobotProcessor,
TransitionKey,
TruncatedProcessor,
)
from .rename_processor import RenameProcessor
__all__ = [
"ActionProcessor",
"DeviceProcessor",
"DoneProcessor",
"EnvTransition",
"IdentityProcessor",
"InfoProcessor",
"NormalizerProcessor",
"UnnormalizerProcessor",
"ObservationProcessor",
"ProcessorStep",
"ProcessorStepRegistry",
"RenameProcessor",
"RewardProcessor",
"RobotProcessor",
"TransitionKey",
"TruncatedProcessor",
"VanillaObservationProcessor",
]
src/lerobot/processor/device_processor.py
+82
View File
@@ -0,0 +1,82 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
from typing import Any
import torch
from lerobot.configs.types import PolicyFeature
from lerobot.processor.pipeline import EnvTransition, TransitionKey
from lerobot.utils.utils import get_safe_torch_device
@dataclass
class DeviceProcessor:
"""Processes transitions by moving tensors to the specified device.
This processor ensures that all tensors in the transition are moved to the
specified device (CPU or GPU) before they are returned.
"""
device: torch.device = "cpu"
def __post_init__(self):
self.device = get_safe_torch_device(self.device)
self.non_blocking = "cuda" in str(self.device)
def __call__(self, transition: EnvTransition) -> EnvTransition:
# Create a copy of the transition
new_transition = transition.copy()
# Process observation tensors
observation = transition.get(TransitionKey.OBSERVATION)
if observation is not None:
new_observation = {
k: v.to(self.device, non_blocking=self.non_blocking) if isinstance(v, torch.Tensor) else v
for k, v in observation.items()
}
new_transition[TransitionKey.OBSERVATION] = new_observation
# Process action tensor
action = transition.get(TransitionKey.ACTION)
if action is not None and isinstance(action, torch.Tensor):
new_transition[TransitionKey.ACTION] = action.to(self.device, non_blocking=self.non_blocking)
# Process reward tensor
reward = transition.get(TransitionKey.REWARD)
if reward is not None and isinstance(reward, torch.Tensor):
new_transition[TransitionKey.REWARD] = reward.to(self.device, non_blocking=self.non_blocking)
# Process done tensor
done = transition.get(TransitionKey.DONE)
if done is not None and isinstance(done, torch.Tensor):
new_transition[TransitionKey.DONE] = done.to(self.device, non_blocking=self.non_blocking)
# Process truncated tensor
truncated = transition.get(TransitionKey.TRUNCATED)
if truncated is not None and isinstance(truncated, torch.Tensor):
new_transition[TransitionKey.TRUNCATED] = truncated.to(
self.device, non_blocking=self.non_blocking
)
return new_transition
def get_config(self) -> dict[str, Any]:
"""Return configuration for serialization."""
return {"device": self.device}
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
src/lerobot/processor/normalize_processor.py
+331
View File
@@ -0,0 +1,331 @@
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Any
import numpy as np
import torch
from torch import Tensor
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
def _convert_stats_to_tensors(stats: dict[str, dict[str, Any]]) -> dict[str, dict[str, Tensor]]:
"""Convert numpy arrays and other types to torch tensors."""
tensor_stats: dict[str, dict[str, Tensor]] = {}
for key, sub in stats.items():
tensor_stats[key] = {}
for stat_name, value in sub.items():
if isinstance(value, np.ndarray):
tensor_val = torch.from_numpy(value.astype(np.float32))
elif isinstance(value, torch.Tensor):
tensor_val = value.to(dtype=torch.float32)
elif isinstance(value, (int, float, list, tuple)):
tensor_val = torch.tensor(value, dtype=torch.float32)
else:
raise TypeError(f"Unsupported type for stats['{key}']['{stat_name}']: {type(value)}")
tensor_stats[key][stat_name] = tensor_val
return tensor_stats
@dataclass
@ProcessorStepRegistry.register(name="normalizer_processor")
class NormalizerProcessor:
"""Normalizes observations and actions in a single processor step.
This processor handles normalization of both observation and action tensors
using either mean/std normalization or min/max scaling to a [-1, 1] range.
For each tensor key in the stats dictionary, the processor will:
- Use mean/std normalization if those statistics are provided: (x - mean) / std
- Use min/max scaling if those statistics are provided: 2 * (x - min) / (max - min) - 1
The processor can be configured to normalize only specific keys by setting
the normalize_keys parameter.
"""
# Features and normalisation map are mandatory to match the design of normalize.py
features: dict[str, PolicyFeature]
norm_map: dict[FeatureType, NormalizationMode]
# Pre-computed statistics coming from dataset.meta.stats for instance.
stats: dict[str, dict[str, Any]] | None = None
# Explicit subset of keys to normalise. If ``None`` every key (except
# "action") found in ``stats`` will be normalised. Using a ``set`` makes
# membership checks O(1).
normalize_keys: set[str] | None = None
eps: float = 1e-8
_tensor_stats: dict[str, dict[str, Tensor]] = field(default_factory=dict, init=False, repr=False)
@classmethod
def from_lerobot_dataset(
cls,
dataset: LeRobotDataset,
features: dict[str, PolicyFeature],
norm_map: dict[FeatureType, NormalizationMode],
*,
normalize_keys: set[str] | None = None,
eps: float = 1e-8,
) -> NormalizerProcessor:
"""Factory helper that pulls statistics from a :class:`LeRobotDataset`.
The features and norm_map parameters are mandatory to match the design
pattern used in normalize.py.
"""
return cls(
features=features,
norm_map=norm_map,
stats=dataset.meta.stats,
normalize_keys=normalize_keys,
eps=eps,
)
def __post_init__(self):
# Handle deserialization from JSON config
if self.features and isinstance(list(self.features.values())[0], dict):
# Features came from JSON - need to reconstruct PolicyFeature objects
reconstructed_features = {}
for key, ft_dict in self.features.items():
reconstructed_features[key] = PolicyFeature(
type=FeatureType(ft_dict["type"]), shape=tuple(ft_dict["shape"])
)
self.features = reconstructed_features
if self.norm_map and isinstance(list(self.norm_map.keys())[0], str):
# norm_map came from JSON - need to reconstruct enum keys and values
reconstructed_norm_map = {}
for ft_type_str, norm_mode_str in self.norm_map.items():
reconstructed_norm_map[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
self.norm_map = reconstructed_norm_map
# Convert statistics once so we avoid repeated numpy→Tensor conversions
# during runtime.
self.stats = self.stats or {}
self._tensor_stats = _convert_stats_to_tensors(self.stats)
# Ensure *normalize_keys* is a set for fast look-ups and compare by
# value later when returning the configuration.
if self.normalize_keys is not None and not isinstance(self.normalize_keys, set):
self.normalize_keys = set(self.normalize_keys)
def _normalize_obs(self, observation):
if observation is None:
return None
# Decide which keys should be normalised for this call.
if self.normalize_keys is not None:
keys_to_norm = self.normalize_keys
else:
# Use feature map to skip action keys.
keys_to_norm = {k for k, ft in self.features.items() if ft.type is not FeatureType.ACTION}
processed = dict(observation)
for key in keys_to_norm:
if key not in processed or key not in self._tensor_stats:
continue
orig_val = processed[key]
tensor = (
orig_val.to(dtype=torch.float32)
if isinstance(orig_val, torch.Tensor)
else torch.as_tensor(orig_val, dtype=torch.float32)
)
stats = {k: v.to(tensor.device) for k, v in self._tensor_stats[key].items()}
if "mean" in stats and "std" in stats:
mean, std = stats["mean"], stats["std"]
processed[key] = (tensor - mean) / (std + self.eps)
elif "min" in stats and "max" in stats:
min_val, max_val = stats["min"], stats["max"]
processed[key] = 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
return processed
def _normalize_action(self, action):
if action is None or "action" not in self._tensor_stats:
return action
tensor = (
action.to(dtype=torch.float32)
if isinstance(action, torch.Tensor)
else torch.as_tensor(action, dtype=torch.float32)
)
stats = {k: v.to(tensor.device) for k, v in self._tensor_stats["action"].items()}
if "mean" in stats and "std" in stats:
mean, std = stats["mean"], stats["std"]
return (tensor - mean) / (std + self.eps)
if "min" in stats and "max" in stats:
min_val, max_val = stats["min"], stats["max"]
return 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
raise ValueError("Action stats must contain either ('mean','std') or ('min','max')")
def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = self._normalize_obs(transition.get(TransitionKey.OBSERVATION))
action = self._normalize_action(transition.get(TransitionKey.ACTION))
# Create a new transition with normalized values
new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = observation
new_transition[TransitionKey.ACTION] = action
return new_transition
def get_config(self) -> dict[str, Any]:
config = {
"eps": self.eps,
"features": {
key: {"type": ft.type.value, "shape": ft.shape} for key, ft in self.features.items()
},
"norm_map": {ft_type.value: norm_mode.value for ft_type, norm_mode in self.norm_map.items()},
}
if self.normalize_keys is not None:
# Serialise as a list for YAML / JSON friendliness
config["normalize_keys"] = sorted(self.normalize_keys)
return config
def state_dict(self) -> dict[str, Tensor]:
flat = {}
for key, sub in self._tensor_stats.items():
for stat_name, tensor in sub.items():
flat[f"{key}.{stat_name}"] = tensor
return flat
def load_state_dict(self, state: Mapping[str, Tensor]) -> None:
self._tensor_stats.clear()
for flat_key, tensor in state.items():
key, stat_name = flat_key.rsplit(".", 1)
self._tensor_stats.setdefault(key, {})[stat_name] = tensor
def reset(self):
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
@dataclass
@ProcessorStepRegistry.register(name="unnormalizer_processor")
class UnnormalizerProcessor:
"""Inverse normalisation for observations and actions.
Exactly mirrors :class:`NormalizerProcessor` but applies the inverse
transform.
"""
features: dict[str, PolicyFeature]
norm_map: dict[FeatureType, NormalizationMode]
stats: dict[str, dict[str, Any]] | None = None
_tensor_stats: dict[str, dict[str, Tensor]] = field(default_factory=dict, init=False, repr=False)
@classmethod
def from_lerobot_dataset(
cls,
dataset: LeRobotDataset,
features: dict[str, PolicyFeature],
norm_map: dict[FeatureType, NormalizationMode],
) -> UnnormalizerProcessor:
return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats)
def __post_init__(self):
# Handle deserialization from JSON config
if self.features and isinstance(list(self.features.values())[0], dict):
# Features came from JSON - need to reconstruct PolicyFeature objects
reconstructed_features = {}
for key, ft_dict in self.features.items():
reconstructed_features[key] = PolicyFeature(
type=FeatureType(ft_dict["type"]), shape=tuple(ft_dict["shape"])
)
self.features = reconstructed_features
if self.norm_map and isinstance(list(self.norm_map.keys())[0], str):
# norm_map came from JSON - need to reconstruct enum keys and values
reconstructed_norm_map = {}
for ft_type_str, norm_mode_str in self.norm_map.items():
reconstructed_norm_map[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
self.norm_map = reconstructed_norm_map
self.stats = self.stats or {}
self._tensor_stats = _convert_stats_to_tensors(self.stats)
def _unnormalize_obs(self, observation):
if observation is None:
return None
keys = [k for k, ft in self.features.items() if ft.type is not FeatureType.ACTION]
processed = dict(observation)
for key in keys:
if key not in processed or key not in self._tensor_stats:
continue
orig_val = processed[key]
tensor = (
orig_val.to(dtype=torch.float32)
if isinstance(orig_val, torch.Tensor)
else torch.as_tensor(orig_val, dtype=torch.float32)
)
stats = {k: v.to(tensor.device) for k, v in self._tensor_stats[key].items()}
if "mean" in stats and "std" in stats:
mean, std = stats["mean"], stats["std"]
processed[key] = tensor * std + mean
elif "min" in stats and "max" in stats:
min_val, max_val = stats["min"], stats["max"]
processed[key] = (tensor + 1) / 2 * (max_val - min_val) + min_val
return processed
def _unnormalize_action(self, action):
if action is None or "action" not in self._tensor_stats:
return action
tensor = (
action.to(dtype=torch.float32)
if isinstance(action, torch.Tensor)
else torch.as_tensor(action, dtype=torch.float32)
)
stats = {k: v.to(tensor.device) for k, v in self._tensor_stats["action"].items()}
if "mean" in stats and "std" in stats:
mean, std = stats["mean"], stats["std"]
return tensor * std + mean
if "min" in stats and "max" in stats:
min_val, max_val = stats["min"], stats["max"]
return (tensor + 1) / 2 * (max_val - min_val) + min_val
raise ValueError("Action stats must contain either ('mean','std') or ('min','max')")
def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = self._unnormalize_obs(transition.get(TransitionKey.OBSERVATION))
action = self._unnormalize_action(transition.get(TransitionKey.ACTION))
# Create a new transition with unnormalized values
new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = observation
new_transition[TransitionKey.ACTION] = action
return new_transition
def get_config(self) -> dict[str, Any]:
return {
"features": {
key: {"type": ft.type.value, "shape": ft.shape} for key, ft in self.features.items()
},
"norm_map": {ft_type.value: norm_mode.value for ft_type, norm_mode in self.norm_map.items()},
}
def state_dict(self) -> dict[str, Tensor]:
flat = {}
for key, sub in self._tensor_stats.items():
for stat_name, tensor in sub.items():
flat[f"{key}.{stat_name}"] = tensor
return flat
def load_state_dict(self, state: Mapping[str, Tensor]) -> None:
self._tensor_stats.clear()
for flat_key, tensor in state.items():
key, stat_name = flat_key.rsplit(".", 1)
self._tensor_stats.setdefault(key, {})[stat_name] = tensor
def reset(self):
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
src/lerobot/processor/observation_processor.py
+157
View File
@@ -0,0 +1,157 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
import einops
import numpy as np
import torch
from torch import Tensor
from lerobot.configs.types import PolicyFeature
from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
from lerobot.processor.pipeline import ObservationProcessor, ProcessorStepRegistry
@dataclass
@ProcessorStepRegistry.register(name="observation_processor")
class VanillaObservationProcessor(ObservationProcessor):
"""
Processes environment observations into the LeRobot format by handling both images and states.
Image processing:
- Converts channel-last (H, W, C) images to channel-first (C, H, W)
- Normalizes uint8 images ([0, 255]) to float32 ([0, 1])
- Adds a batch dimension if missing
- Supports single images and image dictionaries
State processing:
- Maps 'environment_state' to observation.environment_state
- Maps 'agent_pos' to observation.state
- Converts numpy arrays to tensors
- Adds a batch dimension if missing
"""
def _process_single_image(self, img: np.ndarray) -> Tensor:
"""Process a single image array."""
# Convert to tensor
img_tensor = torch.from_numpy(img)
# Add batch dimension if needed
if img_tensor.ndim == 3:
img_tensor = img_tensor.unsqueeze(0)
# Validate image format
_, h, w, c = img_tensor.shape
if not (c < h and c < w):
raise ValueError(f"Expected channel-last images, but got shape {img_tensor.shape}")
if img_tensor.dtype != torch.uint8:
raise ValueError(f"Expected torch.uint8 images, but got {img_tensor.dtype}")
# Convert to channel-first format
img_tensor = einops.rearrange(img_tensor, "b h w c -> b c h w").contiguous()
# Convert to float32 and normalize to [0, 1]
img_tensor = img_tensor.type(torch.float32) / 255.0
return img_tensor
def _process_observation(self, observation):
"""
Processes both image and state observations.
"""
processed_obs = observation.copy()
if "pixels" in processed_obs:
pixels = processed_obs.pop("pixels")
if isinstance(pixels, dict):
imgs = {f"{OBS_IMAGES}.{key}": img for key, img in pixels.items()}
else:
imgs = {OBS_IMAGE: pixels}
for imgkey, img in imgs.items():
processed_obs[imgkey] = self._process_single_image(img)
if "environment_state" in processed_obs:
env_state_np = processed_obs.pop("environment_state")
env_state = torch.from_numpy(env_state_np).float()
if env_state.dim() == 1:
env_state = env_state.unsqueeze(0)
processed_obs[OBS_ENV_STATE] = env_state
if "agent_pos" in processed_obs:
agent_pos_np = processed_obs.pop("agent_pos")
agent_pos = torch.from_numpy(agent_pos_np).float()
if agent_pos.dim() == 1:
agent_pos = agent_pos.unsqueeze(0)
processed_obs[OBS_STATE] = agent_pos
return processed_obs
def observation(self, observation):
return self._process_observation(observation)
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
"""Transforms feature keys to a standardized contract.
This method handles several renaming patterns:
- Exact matches (e.g., 'pixels' -> 'OBS_IMAGE').
- Prefixed exact matches (e.g., 'observation.pixels' -> 'OBS_IMAGE').
- Prefix matches (e.g., 'pixels.cam1' -> 'OBS_IMAGES.cam1').
- Prefixed prefix matches (e.g., 'observation.pixels.cam1' -> 'OBS_IMAGES.cam1').
- environment_state -> OBS_ENV_STATE,
- agent_pos -> OBS_STATE,
- observation.environment_state -> OBS_ENV_STATE,
- observation.agent_pos -> OBS_STATE
"""
exact_pairs = {
"pixels": OBS_IMAGE,
"environment_state": OBS_ENV_STATE,
"agent_pos": OBS_STATE,
}
prefix_pairs = {
"pixels.": f"{OBS_IMAGES}.",
}
for key in list(features.keys()):
matched_prefix = False
for old_prefix, new_prefix in prefix_pairs.items():
prefixed_old = f"observation.{old_prefix}"
if key.startswith(prefixed_old):
suffix = key[len(prefixed_old) :]
features[f"{new_prefix}{suffix}"] = features.pop(key)
matched_prefix = True
break
if key.startswith(old_prefix):
suffix = key[len(old_prefix) :]
features[f"{new_prefix}{suffix}"] = features.pop(key)
matched_prefix = True
break
if matched_prefix:
continue
for old, new in exact_pairs.items():
if key == old or key == f"observation.{old}":
if key in features:
features[new] = features.pop(key)
break
return features
src/lerobot/processor/pipeline.py
+1264
View File
File diff suppressed because it is too large Load Diff
src/lerobot/processor/rename_processor.py
+51
View File
@@ -0,0 +1,51 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass, field
from typing import Any
from lerobot.configs.types import PolicyFeature
from lerobot.processor.pipeline import (
ObservationProcessor,
ProcessorStepRegistry,
)
@dataclass
@ProcessorStepRegistry.register(name="rename_processor")
class RenameProcessor(ObservationProcessor):
"""Rename processor that renames keys in the observation."""
rename_map: dict[str, str] = field(default_factory=dict)
def observation(self, observation):
processed_obs = {}
for key, value in observation.items():
if key in self.rename_map:
processed_obs[self.rename_map[key]] = value
else:
processed_obs[key] = value
return processed_obs
def get_config(self) -> dict[str, Any]:
return {"rename_map": self.rename_map}
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
"""Transforms:
- Each key in the observation that appears in `rename_map` is renamed to its value.
- Keys not in `rename_map` remain unchanged.
"""
return {self.rename_map.get(k, k): v for k, v in features.items()}
src/lerobot/record.py
+7 -3
View File
@@ -18,7 +18,7 @@ Records a dataset. Actions for the robot can be either generated by teleoperatio
Example: Example:
```shell ```shell
python -m lerobot.record \ lerobot-record \
--robot.type=so100_follower \ --robot.type=so100_follower \
--robot.port=/dev/tty.usbmodem58760431541 \ --robot.port=/dev/tty.usbmodem58760431541 \
--robot.cameras="{laptop: {type: opencv, camera_index: 0, width: 640, height: 480}}" \ --robot.cameras="{laptop: {type: opencv, camera_index: 0, width: 640, height: 480}}" \
@@ -36,7 +36,7 @@ python -m lerobot.record \
Example recording with bimanual so100: Example recording with bimanual so100:
```shell ```shell
python -m lerobot.record \ lerobot-record \
--robot.type=bi_so100_follower \ --robot.type=bi_so100_follower \
--robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \
--robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \
@@ -389,5 +389,9 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
return dataset return dataset
if __name__ == "__main__": def main():
record() record()
if __name__ == "__main__":
main()
src/lerobot/replay.py
+7 -3
View File
@@ -18,7 +18,7 @@ Replays the actions of an episode from a dataset on a robot.
Examples: Examples:
```shell ```shell
python -m lerobot.replay \ lerobot-replay \
--robot.type=so100_follower \ --robot.type=so100_follower \
--robot.port=/dev/tty.usbmodem58760431541 \ --robot.port=/dev/tty.usbmodem58760431541 \
--robot.id=black \ --robot.id=black \
@@ -28,7 +28,7 @@ python -m lerobot.replay \
Example replay with bimanual so100: Example replay with bimanual so100:
```shell ```shell
python -m lerobot.replay \ lerobot-replay \
--robot.type=bi_so100_follower \ --robot.type=bi_so100_follower \
--robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \
--robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \
@@ -112,5 +112,9 @@ def replay(cfg: ReplayConfig):
robot.disconnect() robot.disconnect()
if __name__ == "__main__": def main():
replay() replay()
if __name__ == "__main__":
main()
src/lerobot/robots/viperx/README.md
+2 -2
View File
@@ -141,10 +141,10 @@ python lerobot/scripts/control_robot.py \
## Train a policy ## Train a policy
To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](../src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command: To train a policy to control your robot, use the [`lerobot-train`](../src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--dataset.repo_id=${HF_USER}/aloha_test \ --dataset.repo_id=${HF_USER}/aloha_test \
--policy.type=act \ --policy.type=act \
--output_dir=outputs/train/act_aloha_test \ --output_dir=outputs/train/act_aloha_test \
src/lerobot/scripts/eval.py
+7 -3
View File
@@ -21,7 +21,7 @@ You want to evaluate a model from the hub (eg: https://huggingface.co/lerobot/di
for 10 episodes. for 10 episodes.
``` ```
python -m lerobot.scripts.eval \ lerobot-eval \
--policy.path=lerobot/diffusion_pusht \ --policy.path=lerobot/diffusion_pusht \
--env.type=pusht \ --env.type=pusht \
--eval.batch_size=10 \ --eval.batch_size=10 \
@@ -32,7 +32,7 @@ python -m lerobot.scripts.eval \
OR, you want to evaluate a model checkpoint from the LeRobot training script for 10 episodes. OR, you want to evaluate a model checkpoint from the LeRobot training script for 10 episodes.
``` ```
python -m lerobot.scripts.eval \ lerobot-eval \
--policy.path=outputs/train/diffusion_pusht/checkpoints/005000/pretrained_model \ --policy.path=outputs/train/diffusion_pusht/checkpoints/005000/pretrained_model \
--env.type=pusht \ --env.type=pusht \
--eval.batch_size=10 \ --eval.batch_size=10 \
@@ -501,6 +501,10 @@ def eval_main(cfg: EvalPipelineConfig):
logging.info("End of eval") logging.info("End of eval")
if __name__ == "__main__": def main():
init_logging() init_logging()
eval_main() eval_main()
if __name__ == "__main__":
main()
src/lerobot/scripts/train.py
+5 -1
View File
@@ -287,6 +287,10 @@ def train(cfg: TrainPipelineConfig):
policy.push_model_to_hub(cfg) policy.push_model_to_hub(cfg)
if __name__ == "__main__": def main():
init_logging() init_logging()
train() train()
if __name__ == "__main__":
main()
src/lerobot/setup_motors.py
+6 -2
View File
@@ -18,7 +18,7 @@ Helper to set motor ids and baudrate.
Example: Example:
```shell ```shell
python -m lerobot.setup_motors \ lerobot-setup-motors \
--teleop.type=so100_leader \ --teleop.type=so100_leader \
--teleop.port=/dev/tty.usbmodem575E0031751 --teleop.port=/dev/tty.usbmodem575E0031751
``` ```
@@ -80,5 +80,9 @@ def setup_motors(cfg: SetupConfig):
device.setup_motors() device.setup_motors()
if __name__ == "__main__": def main():
setup_motors() setup_motors()
if __name__ == "__main__":
main()
src/lerobot/teleoperate.py
+7 -3
View File
@@ -18,7 +18,7 @@ Simple script to control a robot from teleoperation.
Example: Example:
```shell ```shell
python -m lerobot.teleoperate \ lerobot-teleoperate \
--robot.type=so101_follower \ --robot.type=so101_follower \
--robot.port=/dev/tty.usbmodem58760431541 \ --robot.port=/dev/tty.usbmodem58760431541 \
--robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \ --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 1920, height: 1080, fps: 30}}" \
@@ -32,7 +32,7 @@ python -m lerobot.teleoperate \
Example teleoperation with bimanual so100: Example teleoperation with bimanual so100:
```shell ```shell
python -m lerobot.teleoperate \ lerobot-teleoperate \
--robot.type=bi_so100_follower \ --robot.type=bi_so100_follower \
--robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \
--robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \
@@ -153,5 +153,9 @@ def teleoperate(cfg: TeleoperateConfig):
robot.disconnect() robot.disconnect()
if __name__ == "__main__": def main():
teleoperate() teleoperate()
if __name__ == "__main__":
main()
src/lerobot/templates/lerobot_modelcard_template.md
+2 -2
View File
@@ -44,7 +44,7 @@ Below is the short version on how to train and run inference/eval:
### Train from scratch ### Train from scratch
```bash ```bash
python -m lerobot.scripts.train \ lerobot-train \
--dataset.repo_id=${HF_USER}/<dataset> \ --dataset.repo_id=${HF_USER}/<dataset> \
--policy.type=act \ --policy.type=act \
--output_dir=outputs/train/<desired_policy_repo_id> \ --output_dir=outputs/train/<desired_policy_repo_id> \
@@ -59,7 +59,7 @@ _Writes checkpoints to `outputs/train/<desired_policy_repo_id>/checkpoints/`._
### Evaluate the policy/run inference ### Evaluate the policy/run inference
```bash ```bash
python -m lerobot.record \ lerobot-record \
--robot.type=so100_follower \ --robot.type=so100_follower \
--dataset.repo_id=<hf_user>/eval_<dataset> \ --dataset.repo_id=<hf_user>/eval_<dataset> \
--policy.path=<hf_user>/<desired_policy_repo_id> \ --policy.path=<hf_user>/<desired_policy_repo_id> \
src/lerobot/utils/robot_utils.py
+2 -3
View File
@@ -17,10 +17,9 @@ import time
def busy_wait(seconds): def busy_wait(seconds):
if platform.system() == "Darwin": if platform.system() == "Darwin" or platform.system() == "Windows":
# On Mac, `time.sleep` is not accurate and we need to use this while loop trick, # On Mac and Windows, `time.sleep` is not accurate and we need to use this while loop trick,
# but it consumes CPU cycles. # but it consumes CPU cycles.
# TODO(rcadene): find an alternative: from python 11, time.sleep is precise
end_time = time.perf_counter() + seconds end_time = time.perf_counter() + seconds
while time.perf_counter() < end_time: while time.perf_counter() < end_time:
pass pass
tests/conftest.py
+17
View File
@@ -19,6 +19,7 @@ import traceback
import pytest import pytest
from serial import SerialException from serial import SerialException
from lerobot.configs.types import FeatureType, PolicyFeature
from tests.utils import DEVICE from tests.utils import DEVICE
# Import fixture modules as plugins # Import fixture modules as plugins
@@ -69,3 +70,19 @@ def patch_builtins_input(monkeypatch):
print(text) print(text)
monkeypatch.setattr("builtins.input", print_text) monkeypatch.setattr("builtins.input", print_text)
@pytest.fixture
def policy_feature_factory():
"""PolicyFeature factory"""
def _pf(ft: FeatureType, shape: tuple[int, ...]) -> PolicyFeature:
return PolicyFeature(type=ft, shape=shape)
return _pf
def assert_contract_is_typed(features: dict[str, PolicyFeature]) -> None:
assert isinstance(features, dict)
assert all(isinstance(k, str) for k in features.keys())
assert all(isinstance(v, PolicyFeature) for v in features.values())
tests/policies/test_policies.py
+50
View File
@@ -27,11 +27,13 @@ from lerobot import available_policies
from lerobot.configs.default import DatasetConfig from lerobot.configs.default import DatasetConfig
from lerobot.configs.train import TrainPipelineConfig from lerobot.configs.train import TrainPipelineConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.constants import ACTION, OBS_STATE
from lerobot.datasets.factory import make_dataset from lerobot.datasets.factory import make_dataset
from lerobot.datasets.utils import cycle, dataset_to_policy_features from lerobot.datasets.utils import cycle, dataset_to_policy_features
from lerobot.envs.factory import make_env, make_env_config from lerobot.envs.factory import make_env, make_env_config
from lerobot.envs.utils import preprocess_observation from lerobot.envs.utils import preprocess_observation
from lerobot.optim.factory import make_optimizer_and_scheduler from lerobot.optim.factory import make_optimizer_and_scheduler
from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.policies.act.modeling_act import ACTTemporalEnsembler from lerobot.policies.act.modeling_act import ACTTemporalEnsembler
from lerobot.policies.factory import ( from lerobot.policies.factory import (
get_policy_class, get_policy_class,
@@ -366,6 +368,54 @@ def test_normalize(insert_temporal_dim):
unnormalize(output_batch) unnormalize(output_batch)
@pytest.mark.parametrize("multikey", [True, False])
def test_multikey_construction(multikey: bool):
"""
Asserts that multiple keys with type State/Action are correctly processed by the policy constructor,
preventing erroneous creation of the policy object.
"""
input_features = {
"observation.state": PolicyFeature(
type=FeatureType.STATE,
shape=(10,),
),
}
output_features = {
"action": PolicyFeature(
type=FeatureType.ACTION,
shape=(5,),
),
}
if multikey:
"""Simulates the complete state/action is constructed from more granular multiple
keys, of the same type as the overall state/action"""
input_features = {}
input_features["observation.state.subset1"] = PolicyFeature(type=FeatureType.STATE, shape=(5,))
input_features["observation.state.subset2"] = PolicyFeature(type=FeatureType.STATE, shape=(5,))
input_features["observation.state"] = PolicyFeature(type=FeatureType.STATE, shape=(10,))
output_features = {}
output_features["action.first_three_motors"] = PolicyFeature(type=FeatureType.ACTION, shape=(3,))
output_features["action.last_two_motors"] = PolicyFeature(type=FeatureType.ACTION, shape=(2,))
output_features["action"] = PolicyFeature(
type=FeatureType.ACTION,
shape=(5,),
)
config = ACTConfig(input_features=input_features, output_features=output_features)
state_condition = config.robot_state_feature == input_features[OBS_STATE]
action_condition = config.action_feature == output_features[ACTION]
assert state_condition, (
f"Discrepancy detected. Robot state feature is {config.robot_state_feature} but policy expects {input_features[OBS_STATE]}"
)
assert action_condition, (
f"Discrepancy detected. Action feature is {config.action_feature} but policy expects {output_features[ACTION]}"
)
@pytest.mark.parametrize( @pytest.mark.parametrize(
"ds_repo_id, policy_name, policy_kwargs, file_name_extra", "ds_repo_id, policy_name, policy_kwargs, file_name_extra",
[ [
tests/processor/test_batch_conversion.py
+282
View File
@@ -0,0 +1,282 @@
import torch
from lerobot.processor.pipeline import (
RobotProcessor,
TransitionKey,
_default_batch_to_transition,
_default_transition_to_batch,
)
def _dummy_batch():
"""Create a dummy batch using the new format with observation.* and next.* keys."""
return {
"observation.image.left": torch.randn(1, 3, 128, 128),
"observation.image.right": torch.randn(1, 3, 128, 128),
"observation.state": torch.tensor([[0.1, 0.2, 0.3, 0.4]]),
"action": torch.tensor([[0.5]]),
"next.reward": 1.0,
"next.done": False,
"next.truncated": False,
"info": {"key": "value"},
}
def test_observation_grouping_roundtrip():
"""Test that observation.* keys are properly grouped and ungrouped."""
proc = RobotProcessor([])
batch_in = _dummy_batch()
batch_out = proc(batch_in)
# Check that all observation.* keys are preserved
original_obs_keys = {k: v for k, v in batch_in.items() if k.startswith("observation.")}
reconstructed_obs_keys = {k: v for k, v in batch_out.items() if k.startswith("observation.")}
assert set(original_obs_keys.keys()) == set(reconstructed_obs_keys.keys())
# Check tensor values
assert torch.allclose(batch_out["observation.image.left"], batch_in["observation.image.left"])
assert torch.allclose(batch_out["observation.image.right"], batch_in["observation.image.right"])
assert torch.allclose(batch_out["observation.state"], batch_in["observation.state"])
# Check other fields
assert torch.allclose(batch_out["action"], batch_in["action"])
assert batch_out["next.reward"] == batch_in["next.reward"]
assert batch_out["next.done"] == batch_in["next.done"]
assert batch_out["next.truncated"] == batch_in["next.truncated"]
assert batch_out["info"] == batch_in["info"]
def test_batch_to_transition_observation_grouping():
"""Test that _default_batch_to_transition correctly groups observation.* keys."""
batch = {
"observation.image.top": torch.randn(1, 3, 128, 128),
"observation.image.left": torch.randn(1, 3, 128, 128),
"observation.state": [1, 2, 3, 4],
"action": "action_data",
"next.reward": 1.5,
"next.done": True,
"next.truncated": False,
"info": {"episode": 42},
}
transition = _default_batch_to_transition(batch)
# Check observation is a dict with all observation.* keys
assert isinstance(transition[TransitionKey.OBSERVATION], dict)
assert "observation.image.top" in transition[TransitionKey.OBSERVATION]
assert "observation.image.left" in transition[TransitionKey.OBSERVATION]
assert "observation.state" in transition[TransitionKey.OBSERVATION]
# Check values are preserved
assert torch.allclose(
transition[TransitionKey.OBSERVATION]["observation.image.top"], batch["observation.image.top"]
)
assert torch.allclose(
transition[TransitionKey.OBSERVATION]["observation.image.left"], batch["observation.image.left"]
)
assert transition[TransitionKey.OBSERVATION]["observation.state"] == [1, 2, 3, 4]
# Check other fields
assert transition[TransitionKey.ACTION] == "action_data"
assert transition[TransitionKey.REWARD] == 1.5
assert transition[TransitionKey.DONE]
assert not transition[TransitionKey.TRUNCATED]
assert transition[TransitionKey.INFO] == {"episode": 42}
assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
def test_transition_to_batch_observation_flattening():
"""Test that _default_transition_to_batch correctly flattens observation dict."""
observation_dict = {
"observation.image.top": torch.randn(1, 3, 128, 128),
"observation.image.left": torch.randn(1, 3, 128, 128),
"observation.state": [1, 2, 3, 4],
}
transition = {
TransitionKey.OBSERVATION: observation_dict,
TransitionKey.ACTION: "action_data",
TransitionKey.REWARD: 1.5,
TransitionKey.DONE: True,
TransitionKey.TRUNCATED: False,
TransitionKey.INFO: {"episode": 42},
TransitionKey.COMPLEMENTARY_DATA: {},
}
batch = _default_transition_to_batch(transition)
# Check that observation.* keys are flattened back to batch
assert "observation.image.top" in batch
assert "observation.image.left" in batch
assert "observation.state" in batch
# Check values are preserved
assert torch.allclose(batch["observation.image.top"], observation_dict["observation.image.top"])
assert torch.allclose(batch["observation.image.left"], observation_dict["observation.image.left"])
assert batch["observation.state"] == [1, 2, 3, 4]
# Check other fields are mapped to next.* format
assert batch["action"] == "action_data"
assert batch["next.reward"] == 1.5
assert batch["next.done"]
assert not batch["next.truncated"]
assert batch["info"] == {"episode": 42}
def test_no_observation_keys():
"""Test behavior when there are no observation.* keys."""
batch = {
"action": "action_data",
"next.reward": 2.0,
"next.done": False,
"next.truncated": True,
"info": {"test": "no_obs"},
}
transition = _default_batch_to_transition(batch)
# Observation should be None when no observation.* keys
assert transition[TransitionKey.OBSERVATION] is None
# Check other fields
assert transition[TransitionKey.ACTION] == "action_data"
assert transition[TransitionKey.REWARD] == 2.0
assert not transition[TransitionKey.DONE]
assert transition[TransitionKey.TRUNCATED]
assert transition[TransitionKey.INFO] == {"test": "no_obs"}
# Round trip should work
reconstructed_batch = _default_transition_to_batch(transition)
assert reconstructed_batch["action"] == "action_data"
assert reconstructed_batch["next.reward"] == 2.0
assert not reconstructed_batch["next.done"]
assert reconstructed_batch["next.truncated"]
assert reconstructed_batch["info"] == {"test": "no_obs"}
def test_minimal_batch():
"""Test with minimal batch containing only observation.* and action."""
batch = {"observation.state": "minimal_state", "action": "minimal_action"}
transition = _default_batch_to_transition(batch)
# Check observation
assert transition[TransitionKey.OBSERVATION] == {"observation.state": "minimal_state"}
assert transition[TransitionKey.ACTION] == "minimal_action"
# Check defaults
assert transition[TransitionKey.REWARD] == 0.0
assert not transition[TransitionKey.DONE]
assert not transition[TransitionKey.TRUNCATED]
assert transition[TransitionKey.INFO] == {}
assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
# Round trip
reconstructed_batch = _default_transition_to_batch(transition)
assert reconstructed_batch["observation.state"] == "minimal_state"
assert reconstructed_batch["action"] == "minimal_action"
assert reconstructed_batch["next.reward"] == 0.0
assert not reconstructed_batch["next.done"]
assert not reconstructed_batch["next.truncated"]
assert reconstructed_batch["info"] == {}
def test_empty_batch():
"""Test behavior with empty batch."""
batch = {}
transition = _default_batch_to_transition(batch)
# All fields should have defaults
assert transition[TransitionKey.OBSERVATION] is None
assert transition[TransitionKey.ACTION] is None
assert transition[TransitionKey.REWARD] == 0.0
assert not transition[TransitionKey.DONE]
assert not transition[TransitionKey.TRUNCATED]
assert transition[TransitionKey.INFO] == {}
assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
# Round trip
reconstructed_batch = _default_transition_to_batch(transition)
assert reconstructed_batch["action"] is None
assert reconstructed_batch["next.reward"] == 0.0
assert not reconstructed_batch["next.done"]
assert not reconstructed_batch["next.truncated"]
assert reconstructed_batch["info"] == {}
def test_complex_nested_observation():
"""Test with complex nested observation data."""
batch = {
"observation.image.top": {"image": torch.randn(1, 3, 128, 128), "timestamp": 1234567890},
"observation.image.left": {"image": torch.randn(1, 3, 128, 128), "timestamp": 1234567891},
"observation.state": torch.randn(7),
"action": torch.randn(8),
"next.reward": 3.14,
"next.done": False,
"next.truncated": True,
"info": {"episode_length": 200, "success": True},
}
transition = _default_batch_to_transition(batch)
reconstructed_batch = _default_transition_to_batch(transition)
# Check that all observation keys are preserved
original_obs_keys = {k for k in batch if k.startswith("observation.")}
reconstructed_obs_keys = {k for k in reconstructed_batch if k.startswith("observation.")}
assert original_obs_keys == reconstructed_obs_keys
# Check tensor values
assert torch.allclose(batch["observation.state"], reconstructed_batch["observation.state"])
# Check nested dict with tensors
assert torch.allclose(
batch["observation.image.top"]["image"], reconstructed_batch["observation.image.top"]["image"]
)
assert torch.allclose(
batch["observation.image.left"]["image"], reconstructed_batch["observation.image.left"]["image"]
)
# Check action tensor
assert torch.allclose(batch["action"], reconstructed_batch["action"])
# Check other fields
assert batch["next.reward"] == reconstructed_batch["next.reward"]
assert batch["next.done"] == reconstructed_batch["next.done"]
assert batch["next.truncated"] == reconstructed_batch["next.truncated"]
assert batch["info"] == reconstructed_batch["info"]
def test_custom_converter():
"""Test that custom converters can still be used."""
def to_tr(batch):
# Custom converter that modifies the reward
tr = _default_batch_to_transition(batch)
# Double the reward
reward = tr.get(TransitionKey.REWARD, 0.0)
new_tr = tr.copy()
new_tr[TransitionKey.REWARD] = reward * 2 if reward is not None else 0.0
return new_tr
def to_batch(tr):
batch = _default_transition_to_batch(tr)
return batch
processor = RobotProcessor(steps=[], to_transition=to_tr, to_output=to_batch)
batch = {
"observation.state": torch.randn(1, 4),
"action": torch.randn(1, 2),
"next.reward": 1.0,
"next.done": False,
}
result = processor(batch)
# Check the reward was doubled by our custom converter
assert result["next.reward"] == 2.0
assert torch.allclose(result["observation.state"], batch["observation.state"])
assert torch.allclose(result["action"], batch["action"])
tests/processor/test_normalize_processor.py
+628
View File
@@ -0,0 +1,628 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest.mock import Mock
import numpy as np
import pytest
import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.processor.normalize_processor import (
NormalizerProcessor,
UnnormalizerProcessor,
_convert_stats_to_tensors,
)
from lerobot.processor.pipeline import RobotProcessor, TransitionKey
def create_transition(
observation=None, action=None, reward=None, done=None, truncated=None, info=None, complementary_data=None
):
"""Helper to create an EnvTransition dictionary."""
return {
TransitionKey.OBSERVATION: observation,
TransitionKey.ACTION: action,
TransitionKey.REWARD: reward,
TransitionKey.DONE: done,
TransitionKey.TRUNCATED: truncated,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: complementary_data,
}
def test_numpy_conversion():
stats = {
"observation.image": {
"mean": np.array([0.5, 0.5, 0.5]),
"std": np.array([0.2, 0.2, 0.2]),
}
}
tensor_stats = _convert_stats_to_tensors(stats)
assert isinstance(tensor_stats["observation.image"]["mean"], torch.Tensor)
assert isinstance(tensor_stats["observation.image"]["std"], torch.Tensor)
assert torch.allclose(tensor_stats["observation.image"]["mean"], torch.tensor([0.5, 0.5, 0.5]))
assert torch.allclose(tensor_stats["observation.image"]["std"], torch.tensor([0.2, 0.2, 0.2]))
def test_tensor_conversion():
stats = {
"action": {
"mean": torch.tensor([0.0, 0.0]),
"std": torch.tensor([1.0, 1.0]),
}
}
tensor_stats = _convert_stats_to_tensors(stats)
assert tensor_stats["action"]["mean"].dtype == torch.float32
assert tensor_stats["action"]["std"].dtype == torch.float32
def test_scalar_conversion():
stats = {
"reward": {
"mean": 0.5,
"std": 0.1,
}
}
tensor_stats = _convert_stats_to_tensors(stats)
assert torch.allclose(tensor_stats["reward"]["mean"], torch.tensor(0.5))
assert torch.allclose(tensor_stats["reward"]["std"], torch.tensor(0.1))
def test_list_conversion():
stats = {
"observation.state": {
"min": [0.0, -1.0, -2.0],
"max": [1.0, 1.0, 2.0],
}
}
tensor_stats = _convert_stats_to_tensors(stats)
assert torch.allclose(tensor_stats["observation.state"]["min"], torch.tensor([0.0, -1.0, -2.0]))
assert torch.allclose(tensor_stats["observation.state"]["max"], torch.tensor([1.0, 1.0, 2.0]))
def test_unsupported_type():
stats = {
"bad_key": {
"mean": "string_value",
}
}
with pytest.raises(TypeError, match="Unsupported type"):
_convert_stats_to_tensors(stats)
# Helper functions to create feature maps and norm maps
def _create_observation_features():
return {
"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96)),
"observation.state": PolicyFeature(FeatureType.STATE, (2,)),
}
def _create_observation_norm_map():
return {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
FeatureType.STATE: NormalizationMode.MIN_MAX,
}
# Fixtures for observation normalisation tests using NormalizerProcessor
@pytest.fixture
def observation_stats():
return {
"observation.image": {
"mean": np.array([0.5, 0.5, 0.5]),
"std": np.array([0.2, 0.2, 0.2]),
},
"observation.state": {
"min": np.array([0.0, -1.0]),
"max": np.array([1.0, 1.0]),
},
}
@pytest.fixture
def observation_normalizer(observation_stats):
"""Return a NormalizerProcessor that only has observation stats (no action)."""
features = _create_observation_features()
norm_map = _create_observation_norm_map()
return NormalizerProcessor(features=features, norm_map=norm_map, stats=observation_stats)
def test_mean_std_normalization(observation_normalizer):
observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]),
}
transition = create_transition(observation=observation)
normalized_transition = observation_normalizer(transition)
normalized_obs = normalized_transition[TransitionKey.OBSERVATION]
# Check mean/std normalization
expected_image = (torch.tensor([0.7, 0.5, 0.3]) - 0.5) / 0.2
assert torch.allclose(normalized_obs["observation.image"], expected_image)
def test_min_max_normalization(observation_normalizer):
observation = {
"observation.state": torch.tensor([0.5, 0.0]),
}
transition = create_transition(observation=observation)
normalized_transition = observation_normalizer(transition)
normalized_obs = normalized_transition[TransitionKey.OBSERVATION]
# Check min/max normalization to [-1, 1]
# For state[0]: 2 * (0.5 - 0.0) / (1.0 - 0.0) - 1 = 0.0
# For state[1]: 2 * (0.0 - (-1.0)) / (1.0 - (-1.0)) - 1 = 0.0
expected_state = torch.tensor([0.0, 0.0])
assert torch.allclose(normalized_obs["observation.state"], expected_state, atol=1e-6)
def test_selective_normalization(observation_stats):
features = _create_observation_features()
norm_map = _create_observation_norm_map()
normalizer = NormalizerProcessor(
features=features, norm_map=norm_map, stats=observation_stats, normalize_keys={"observation.image"}
)
observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]),
}
transition = create_transition(observation=observation)
normalized_transition = normalizer(transition)
normalized_obs = normalized_transition[TransitionKey.OBSERVATION]
# Only image should be normalized
assert torch.allclose(normalized_obs["observation.image"], (torch.tensor([0.7, 0.5, 0.3]) - 0.5) / 0.2)
# State should remain unchanged
assert torch.allclose(normalized_obs["observation.state"], observation["observation.state"])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_device_compatibility(observation_stats):
features = _create_observation_features()
norm_map = _create_observation_norm_map()
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=observation_stats)
observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]).cuda(),
}
transition = create_transition(observation=observation)
normalized_transition = normalizer(transition)
normalized_obs = normalized_transition[TransitionKey.OBSERVATION]
assert normalized_obs["observation.image"].device.type == "cuda"
def test_from_lerobot_dataset():
# Mock dataset
mock_dataset = Mock()
mock_dataset.meta.stats = {
"observation.image": {"mean": [0.5], "std": [0.2]},
"action": {"mean": [0.0], "std": [1.0]},
}
features = {
"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96)),
"action": PolicyFeature(FeatureType.ACTION, (1,)),
}
norm_map = {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
FeatureType.ACTION: NormalizationMode.MEAN_STD,
}
normalizer = NormalizerProcessor.from_lerobot_dataset(mock_dataset, features, norm_map)
# Both observation and action statistics should be present in tensor stats
assert "observation.image" in normalizer._tensor_stats
assert "action" in normalizer._tensor_stats
def test_state_dict_save_load(observation_normalizer):
# Save state
state_dict = observation_normalizer.state_dict()
# Create new normalizer and load state
features = _create_observation_features()
norm_map = _create_observation_norm_map()
new_normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats={})
new_normalizer.load_state_dict(state_dict)
# Test that it works the same
observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])}
transition = create_transition(observation=observation)
result1 = observation_normalizer(transition)[TransitionKey.OBSERVATION]
result2 = new_normalizer(transition)[TransitionKey.OBSERVATION]
assert torch.allclose(result1["observation.image"], result2["observation.image"])
# Fixtures for ActionUnnormalizer tests
@pytest.fixture
def action_stats_mean_std():
return {
"mean": np.array([0.0, 0.0, 0.0]),
"std": np.array([1.0, 2.0, 0.5]),
}
@pytest.fixture
def action_stats_min_max():
return {
"min": np.array([-1.0, -2.0, 0.0]),
"max": np.array([1.0, 2.0, 1.0]),
}
def _create_action_features():
return {
"action": PolicyFeature(FeatureType.ACTION, (3,)),
}
def _create_action_norm_map_mean_std():
return {
FeatureType.ACTION: NormalizationMode.MEAN_STD,
}
def _create_action_norm_map_min_max():
return {
FeatureType.ACTION: NormalizationMode.MIN_MAX,
}
def test_mean_std_unnormalization(action_stats_mean_std):
features = _create_action_features()
norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessor(
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
)
normalized_action = torch.tensor([1.0, -0.5, 2.0])
transition = create_transition(action=normalized_action)
unnormalized_transition = unnormalizer(transition)
unnormalized_action = unnormalized_transition[TransitionKey.ACTION]
# action * std + mean
expected = torch.tensor([1.0 * 1.0 + 0.0, -0.5 * 2.0 + 0.0, 2.0 * 0.5 + 0.0])
assert torch.allclose(unnormalized_action, expected)
def test_min_max_unnormalization(action_stats_min_max):
features = _create_action_features()
norm_map = _create_action_norm_map_min_max()
unnormalizer = UnnormalizerProcessor(
features=features, norm_map=norm_map, stats={"action": action_stats_min_max}
)
# Actions in [-1, 1]
normalized_action = torch.tensor([0.0, -1.0, 1.0])
transition = create_transition(action=normalized_action)
unnormalized_transition = unnormalizer(transition)
unnormalized_action = unnormalized_transition[TransitionKey.ACTION]
# Map from [-1, 1] to [min, max]
# (action + 1) / 2 * (max - min) + min
expected = torch.tensor(
[
(0.0 + 1) / 2 * (1.0 - (-1.0)) + (-1.0), # 0.0
(-1.0 + 1) / 2 * (2.0 - (-2.0)) + (-2.0), # -2.0
(1.0 + 1) / 2 * (1.0 - 0.0) + 0.0, # 1.0
]
)
assert torch.allclose(unnormalized_action, expected)
def test_numpy_action_input(action_stats_mean_std):
features = _create_action_features()
norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessor(
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
)
normalized_action = np.array([1.0, -0.5, 2.0], dtype=np.float32)
transition = create_transition(action=normalized_action)
unnormalized_transition = unnormalizer(transition)
unnormalized_action = unnormalized_transition[TransitionKey.ACTION]
assert isinstance(unnormalized_action, torch.Tensor)
expected = torch.tensor([1.0, -1.0, 1.0])
assert torch.allclose(unnormalized_action, expected)
def test_none_action(action_stats_mean_std):
features = _create_action_features()
norm_map = _create_action_norm_map_mean_std()
unnormalizer = UnnormalizerProcessor(
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
)
transition = create_transition()
result = unnormalizer(transition)
# Should return transition unchanged
assert result == transition
def test_action_from_lerobot_dataset():
mock_dataset = Mock()
mock_dataset.meta.stats = {"action": {"mean": [0.0], "std": [1.0]}}
features = {"action": PolicyFeature(FeatureType.ACTION, (1,))}
norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD}
unnormalizer = UnnormalizerProcessor.from_lerobot_dataset(mock_dataset, features, norm_map)
assert "mean" in unnormalizer._tensor_stats["action"]
# Fixtures for NormalizerProcessor tests
@pytest.fixture
def full_stats():
return {
"observation.image": {
"mean": np.array([0.5, 0.5, 0.5]),
"std": np.array([0.2, 0.2, 0.2]),
},
"observation.state": {
"min": np.array([0.0, -1.0]),
"max": np.array([1.0, 1.0]),
},
"action": {
"mean": np.array([0.0, 0.0]),
"std": np.array([1.0, 2.0]),
},
}
def _create_full_features():
return {
"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96)),
"observation.state": PolicyFeature(FeatureType.STATE, (2,)),
"action": PolicyFeature(FeatureType.ACTION, (2,)),
}
def _create_full_norm_map():
return {
FeatureType.VISUAL: NormalizationMode.MEAN_STD,
FeatureType.STATE: NormalizationMode.MIN_MAX,
FeatureType.ACTION: NormalizationMode.MEAN_STD,
}
@pytest.fixture
def normalizer_processor(full_stats):
features = _create_full_features()
norm_map = _create_full_norm_map()
return NormalizerProcessor(features=features, norm_map=norm_map, stats=full_stats)
def test_combined_normalization(normalizer_processor):
observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]),
}
action = torch.tensor([1.0, -0.5])
transition = create_transition(
observation=observation,
action=action,
reward=1.0,
done=False,
truncated=False,
info={},
complementary_data={},
)
processed_transition = normalizer_processor(transition)
# Check normalized observations
processed_obs = processed_transition[TransitionKey.OBSERVATION]
expected_image = (torch.tensor([0.7, 0.5, 0.3]) - 0.5) / 0.2
assert torch.allclose(processed_obs["observation.image"], expected_image)
# Check normalized action
processed_action = processed_transition[TransitionKey.ACTION]
expected_action = torch.tensor([(1.0 - 0.0) / 1.0, (-0.5 - 0.0) / 2.0])
assert torch.allclose(processed_action, expected_action)
# Check other fields remain unchanged
assert processed_transition[TransitionKey.REWARD] == 1.0
assert not processed_transition[TransitionKey.DONE]
def test_processor_from_lerobot_dataset(full_stats):
# Mock dataset
mock_dataset = Mock()
mock_dataset.meta.stats = full_stats
features = _create_full_features()
norm_map = _create_full_norm_map()
processor = NormalizerProcessor.from_lerobot_dataset(
mock_dataset, features, norm_map, normalize_keys={"observation.image"}
)
assert processor.normalize_keys == {"observation.image"}
assert "observation.image" in processor._tensor_stats
assert "action" in processor._tensor_stats
def test_get_config(full_stats):
features = _create_full_features()
norm_map = _create_full_norm_map()
processor = NormalizerProcessor(
features=features, norm_map=norm_map, stats=full_stats, normalize_keys={"observation.image"}, eps=1e-6
)
config = processor.get_config()
expected_config = {
"normalize_keys": ["observation.image"],
"eps": 1e-6,
"features": {
"observation.image": {"type": "VISUAL", "shape": (3, 96, 96)},
"observation.state": {"type": "STATE", "shape": (2,)},
"action": {"type": "ACTION", "shape": (2,)},
},
"norm_map": {
"VISUAL": "MEAN_STD",
"STATE": "MIN_MAX",
"ACTION": "MEAN_STD",
},
}
assert config == expected_config
def test_integration_with_robot_processor(normalizer_processor):
"""Test integration with RobotProcessor pipeline"""
robot_processor = RobotProcessor([normalizer_processor])
observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]),
}
action = torch.tensor([1.0, -0.5])
transition = create_transition(
observation=observation,
action=action,
reward=1.0,
done=False,
truncated=False,
info={},
complementary_data={},
)
processed_transition = robot_processor(transition)
# Verify the processing worked
assert isinstance(processed_transition[TransitionKey.OBSERVATION], dict)
assert isinstance(processed_transition[TransitionKey.ACTION], torch.Tensor)
# Edge case tests
def test_empty_observation():
stats = {"observation.image": {"mean": [0.5], "std": [0.2]}}
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
transition = create_transition()
result = normalizer(transition)
assert result == transition
def test_empty_stats():
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats={})
observation = {"observation.image": torch.tensor([0.5])}
transition = create_transition(observation=observation)
result = normalizer(transition)
# Should return observation unchanged since no stats are available
assert torch.allclose(
result[TransitionKey.OBSERVATION]["observation.image"], observation["observation.image"]
)
def test_partial_stats():
"""If statistics are incomplete, the value should pass through unchanged."""
stats = {"observation.image": {"mean": [0.5]}} # Missing std / (min,max)
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
observation = {"observation.image": torch.tensor([0.7])}
transition = create_transition(observation=observation)
processed = normalizer(transition)[TransitionKey.OBSERVATION]
assert torch.allclose(processed["observation.image"], observation["observation.image"])
def test_missing_action_stats_no_error():
mock_dataset = Mock()
mock_dataset.meta.stats = {"observation.image": {"mean": [0.5], "std": [0.2]}}
features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))}
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
processor = UnnormalizerProcessor.from_lerobot_dataset(mock_dataset, features, norm_map)
# The tensor stats should not contain the 'action' key
assert "action" not in processor._tensor_stats
def test_serialization_roundtrip(full_stats):
"""Test that features and norm_map can be serialized and deserialized correctly."""
features = _create_full_features()
norm_map = _create_full_norm_map()
original_processor = NormalizerProcessor(
features=features, norm_map=norm_map, stats=full_stats, normalize_keys={"observation.image"}, eps=1e-6
)
# Get config (serialization)
config = original_processor.get_config()
# Create a new processor from the config (deserialization)
new_processor = NormalizerProcessor(
features=config["features"],
norm_map=config["norm_map"],
stats=full_stats,
normalize_keys=set(config["normalize_keys"]),
eps=config["eps"],
)
# Test that both processors work the same way
observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]),
}
action = torch.tensor([1.0, -0.5])
transition = create_transition(
observation=observation,
action=action,
reward=1.0,
done=False,
truncated=False,
info={},
complementary_data={},
)
result1 = original_processor(transition)
result2 = new_processor(transition)
# Compare results
assert torch.allclose(
result1[TransitionKey.OBSERVATION]["observation.image"],
result2[TransitionKey.OBSERVATION]["observation.image"],
)
assert torch.allclose(result1[TransitionKey.ACTION], result2[TransitionKey.ACTION])
# Verify features and norm_map are correctly reconstructed
assert new_processor.features.keys() == original_processor.features.keys()
for key in new_processor.features:
assert new_processor.features[key].type == original_processor.features[key].type
assert new_processor.features[key].shape == original_processor.features[key].shape
assert new_processor.norm_map == original_processor.norm_map
tests/processor/test_observation_processor.py
+486
View File
@@ -0,0 +1,486 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import pytest
import torch
from lerobot.configs.types import FeatureType
from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
from lerobot.processor import VanillaObservationProcessor
from lerobot.processor.pipeline import TransitionKey
from tests.conftest import assert_contract_is_typed
def create_transition(
observation=None, action=None, reward=None, done=None, truncated=None, info=None, complementary_data=None
):
"""Helper to create an EnvTransition dictionary."""
return {
TransitionKey.OBSERVATION: observation,
TransitionKey.ACTION: action,
TransitionKey.REWARD: reward,
TransitionKey.DONE: done,
TransitionKey.TRUNCATED: truncated,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: complementary_data,
}
def test_process_single_image():
"""Test processing a single image."""
processor = VanillaObservationProcessor()
# Create a mock image (H, W, C) format, uint8
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
observation = {"pixels": image}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that the image was processed correctly
assert "observation.image" in processed_obs
processed_img = processed_obs["observation.image"]
# Check shape: should be (1, 3, 64, 64) - batch, channels, height, width
assert processed_img.shape == (1, 3, 64, 64)
# Check dtype and range
assert processed_img.dtype == torch.float32
assert processed_img.min() >= 0.0
assert processed_img.max() <= 1.0
def test_process_image_dict():
"""Test processing multiple images in a dictionary."""
processor = VanillaObservationProcessor()
# Create mock images
image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
image2 = np.random.randint(0, 256, size=(48, 48, 3), dtype=np.uint8)
observation = {"pixels": {"camera1": image1, "camera2": image2}}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that both images were processed
assert "observation.images.camera1" in processed_obs
assert "observation.images.camera2" in processed_obs
# Check shapes
assert processed_obs["observation.images.camera1"].shape == (1, 3, 32, 32)
assert processed_obs["observation.images.camera2"].shape == (1, 3, 48, 48)
def test_process_batched_image():
"""Test processing already batched images."""
processor = VanillaObservationProcessor()
# Create a batched image (B, H, W, C)
image = np.random.randint(0, 256, size=(2, 64, 64, 3), dtype=np.uint8)
observation = {"pixels": image}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that batch dimension is preserved
assert processed_obs["observation.image"].shape == (2, 3, 64, 64)
def test_invalid_image_format():
"""Test error handling for invalid image formats."""
processor = VanillaObservationProcessor()
# Test wrong channel order (channels first)
image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8)
observation = {"pixels": image}
transition = create_transition(observation=observation)
with pytest.raises(ValueError, match="Expected channel-last images"):
processor(transition)
def test_invalid_image_dtype():
"""Test error handling for invalid image dtype."""
processor = VanillaObservationProcessor()
# Test wrong dtype
image = np.random.rand(64, 64, 3).astype(np.float32)
observation = {"pixels": image}
transition = create_transition(observation=observation)
with pytest.raises(ValueError, match="Expected torch.uint8 images"):
processor(transition)
def test_no_pixels_in_observation():
"""Test processor when no pixels are in observation."""
processor = VanillaObservationProcessor()
observation = {"other_data": np.array([1, 2, 3])}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Should preserve other data unchanged
assert "other_data" in processed_obs
np.testing.assert_array_equal(processed_obs["other_data"], np.array([1, 2, 3]))
def test_none_observation():
"""Test processor with None observation."""
processor = VanillaObservationProcessor()
transition = create_transition()
result = processor(transition)
assert result == transition
def test_serialization_methods():
"""Test serialization methods."""
processor = VanillaObservationProcessor()
# Test get_config
config = processor.get_config()
assert isinstance(config, dict)
# Test state_dict
state = processor.state_dict()
assert isinstance(state, dict)
# Test load_state_dict (should not raise)
processor.load_state_dict(state)
# Test reset (should not raise)
processor.reset()
def test_process_environment_state():
"""Test processing environment_state."""
processor = VanillaObservationProcessor()
env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
observation = {"environment_state": env_state}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that environment_state was renamed and processed
assert "observation.environment_state" in processed_obs
assert "environment_state" not in processed_obs
processed_state = processed_obs["observation.environment_state"]
assert processed_state.shape == (1, 3) # Batch dimension added
assert processed_state.dtype == torch.float32
torch.testing.assert_close(processed_state, torch.tensor([[1.0, 2.0, 3.0]]))
def test_process_agent_pos():
"""Test processing agent_pos."""
processor = VanillaObservationProcessor()
agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
observation = {"agent_pos": agent_pos}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that agent_pos was renamed and processed
assert "observation.state" in processed_obs
assert "agent_pos" not in processed_obs
processed_state = processed_obs["observation.state"]
assert processed_state.shape == (1, 3) # Batch dimension added
assert processed_state.dtype == torch.float32
torch.testing.assert_close(processed_state, torch.tensor([[0.5, -0.5, 1.0]]))
def test_process_batched_states():
"""Test processing already batched states."""
processor = VanillaObservationProcessor()
env_state = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
agent_pos = np.array([[0.5, -0.5], [1.0, -1.0]], dtype=np.float32)
observation = {"environment_state": env_state, "agent_pos": agent_pos}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that batch dimensions are preserved
assert processed_obs["observation.environment_state"].shape == (2, 2)
assert processed_obs["observation.state"].shape == (2, 2)
def test_process_both_states():
"""Test processing both environment_state and agent_pos."""
processor = VanillaObservationProcessor()
env_state = np.array([1.0, 2.0], dtype=np.float32)
agent_pos = np.array([0.5, -0.5], dtype=np.float32)
observation = {"environment_state": env_state, "agent_pos": agent_pos, "other_data": "keep_me"}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that both states were processed
assert "observation.environment_state" in processed_obs
assert "observation.state" in processed_obs
# Check that original keys were removed
assert "environment_state" not in processed_obs
assert "agent_pos" not in processed_obs
# Check that other data was preserved
assert processed_obs["other_data"] == "keep_me"
def test_no_states_in_observation():
"""Test processor when no states are in observation."""
processor = VanillaObservationProcessor()
observation = {"other_data": np.array([1, 2, 3])}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Should preserve data unchanged
np.testing.assert_array_equal(processed_obs, observation)
def test_complete_observation_processing():
"""Test processing a complete observation with both images and states."""
processor = VanillaObservationProcessor()
# Create mock data
image = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
observation = {
"pixels": image,
"environment_state": env_state,
"agent_pos": agent_pos,
"other_data": "preserve_me",
}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that image was processed
assert "observation.image" in processed_obs
assert processed_obs["observation.image"].shape == (1, 3, 32, 32)
# Check that states were processed
assert "observation.environment_state" in processed_obs
assert "observation.state" in processed_obs
# Check that original keys were removed
assert "pixels" not in processed_obs
assert "environment_state" not in processed_obs
assert "agent_pos" not in processed_obs
# Check that other data was preserved
assert processed_obs["other_data"] == "preserve_me"
def test_image_only_processing():
"""Test processing observation with only images."""
processor = VanillaObservationProcessor()
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
observation = {"pixels": image}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
assert "observation.image" in processed_obs
assert len(processed_obs) == 1
def test_state_only_processing():
"""Test processing observation with only states."""
processor = VanillaObservationProcessor()
agent_pos = np.array([1.0, 2.0], dtype=np.float32)
observation = {"agent_pos": agent_pos}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
assert "observation.state" in processed_obs
assert "agent_pos" not in processed_obs
def test_empty_observation():
"""Test processing empty observation."""
processor = VanillaObservationProcessor()
observation = {}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
assert processed_obs == {}
def test_equivalent_to_original_function():
"""Test that ObservationProcessor produces equivalent results to preprocess_observation."""
# Import the original function for comparison
from lerobot.envs.utils import preprocess_observation
processor = VanillaObservationProcessor()
# Create test data similar to what the original function expects
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
observation = {"pixels": image, "environment_state": env_state, "agent_pos": agent_pos}
# Process with original function
original_result = preprocess_observation(observation)
# Process with new processor
transition = create_transition(observation=observation)
processor_result = processor(transition)[TransitionKey.OBSERVATION]
# Compare results
assert set(original_result.keys()) == set(processor_result.keys())
for key in original_result:
torch.testing.assert_close(original_result[key], processor_result[key])
def test_equivalent_with_image_dict():
"""Test equivalence with dictionary of images."""
from lerobot.envs.utils import preprocess_observation
processor = VanillaObservationProcessor()
# Create test data with multiple cameras
image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
image2 = np.random.randint(0, 256, size=(48, 48, 3), dtype=np.uint8)
agent_pos = np.array([1.0, 2.0], dtype=np.float32)
observation = {"pixels": {"cam1": image1, "cam2": image2}, "agent_pos": agent_pos}
# Process with original function
original_result = preprocess_observation(observation)
# Process with new processor
transition = create_transition(observation=observation)
processor_result = processor(transition)[TransitionKey.OBSERVATION]
# Compare results
assert set(original_result.keys()) == set(processor_result.keys())
for key in original_result:
torch.testing.assert_close(original_result[key], processor_result[key])
def test_image_processor_feature_contract_pixels_to_image(policy_feature_factory):
processor = VanillaObservationProcessor()
features = {
"pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"keep": policy_feature_factory(FeatureType.ENV, (1,)),
}
out = processor.feature_contract(features.copy())
assert OBS_IMAGE in out and out[OBS_IMAGE] == features["pixels"]
assert "pixels" not in out
assert out["keep"] == features["keep"]
assert_contract_is_typed(out)
def test_image_processor_feature_contract_observation_pixels_to_image(policy_feature_factory):
processor = VanillaObservationProcessor()
features = {
"observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"keep": policy_feature_factory(FeatureType.ENV, (1,)),
}
out = processor.feature_contract(features.copy())
assert OBS_IMAGE in out and out[OBS_IMAGE] == features["observation.pixels"]
assert "observation.pixels" not in out
assert out["keep"] == features["keep"]
assert_contract_is_typed(out)
def test_image_processor_feature_contract_multi_camera_and_prefixed(policy_feature_factory):
processor = VanillaObservationProcessor()
features = {
"pixels.front": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"pixels.wrist": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"observation.pixels.rear": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"keep": policy_feature_factory(FeatureType.ENV, (7,)),
}
out = processor.feature_contract(features.copy())
assert f"{OBS_IMAGES}.front" in out and out[f"{OBS_IMAGES}.front"] == features["pixels.front"]
assert f"{OBS_IMAGES}.wrist" in out and out[f"{OBS_IMAGES}.wrist"] == features["pixels.wrist"]
assert f"{OBS_IMAGES}.rear" in out and out[f"{OBS_IMAGES}.rear"] == features["observation.pixels.rear"]
assert "pixels.front" not in out and "pixels.wrist" not in out and "observation.pixels.rear" not in out
assert out["keep"] == features["keep"]
assert_contract_is_typed(out)
def test_state_processor_feature_contract_environment_and_agent_pos(policy_feature_factory):
processor = VanillaObservationProcessor()
features = {
"environment_state": policy_feature_factory(FeatureType.STATE, (3,)),
"agent_pos": policy_feature_factory(FeatureType.STATE, (7,)),
"keep": policy_feature_factory(FeatureType.ENV, (1,)),
}
out = processor.feature_contract(features.copy())
assert OBS_ENV_STATE in out and out[OBS_ENV_STATE] == features["environment_state"]
assert OBS_STATE in out and out[OBS_STATE] == features["agent_pos"]
assert "environment_state" not in out and "agent_pos" not in out
assert out["keep"] == features["keep"]
assert_contract_is_typed(out)
def test_state_processor_feature_contract_prefixed_inputs(policy_feature_factory):
proc = VanillaObservationProcessor()
features = {
"observation.environment_state": policy_feature_factory(FeatureType.STATE, (2,)),
"observation.agent_pos": policy_feature_factory(FeatureType.STATE, (4,)),
}
out = proc.feature_contract(features.copy())
assert OBS_ENV_STATE in out and out[OBS_ENV_STATE] == features["observation.environment_state"]
assert OBS_STATE in out and out[OBS_STATE] == features["observation.agent_pos"]
assert "environment_state" not in out and "agent_pos" not in out
assert_contract_is_typed(out)
tests/processor/test_pipeline.py
+1919
View File
File diff suppressed because it is too large Load Diff
tests/processor/test_rename_processor.py
+467
View File
@@ -0,0 +1,467 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import tempfile
from pathlib import Path
import numpy as np
import torch
from lerobot.configs.types import FeatureType
from lerobot.processor import ProcessorStepRegistry, RenameProcessor, RobotProcessor, TransitionKey
from tests.conftest import assert_contract_is_typed
def create_transition(
observation=None, action=None, reward=None, done=None, truncated=None, info=None, complementary_data=None
):
"""Helper to create an EnvTransition dictionary."""
return {
TransitionKey.OBSERVATION: observation,
TransitionKey.ACTION: action,
TransitionKey.REWARD: reward,
TransitionKey.DONE: done,
TransitionKey.TRUNCATED: truncated,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: complementary_data,
}
def test_basic_renaming():
"""Test basic key renaming functionality."""
rename_map = {
"old_key1": "new_key1",
"old_key2": "new_key2",
}
processor = RenameProcessor(rename_map=rename_map)
observation = {
"old_key1": torch.tensor([1.0, 2.0]),
"old_key2": np.array([3.0, 4.0]),
"unchanged_key": "keep_me",
}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check renamed keys
assert "new_key1" in processed_obs
assert "new_key2" in processed_obs
assert "old_key1" not in processed_obs
assert "old_key2" not in processed_obs
# Check values are preserved
torch.testing.assert_close(processed_obs["new_key1"], torch.tensor([1.0, 2.0]))
np.testing.assert_array_equal(processed_obs["new_key2"], np.array([3.0, 4.0]))
# Check unchanged key is preserved
assert processed_obs["unchanged_key"] == "keep_me"
def test_empty_rename_map():
"""Test processor with empty rename map (should pass through unchanged)."""
processor = RenameProcessor(rename_map={})
observation = {
"key1": torch.tensor([1.0]),
"key2": "value2",
}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# All keys should be unchanged
assert processed_obs.keys() == observation.keys()
torch.testing.assert_close(processed_obs["key1"], observation["key1"])
assert processed_obs["key2"] == observation["key2"]
def test_none_observation():
"""Test processor with None observation."""
processor = RenameProcessor(rename_map={"old": "new"})
transition = create_transition()
result = processor(transition)
# Should return transition unchanged
assert result == transition
def test_overlapping_rename():
"""Test renaming when new names might conflict."""
rename_map = {
"a": "b",
"b": "c", # This creates a potential conflict
}
processor = RenameProcessor(rename_map=rename_map)
observation = {
"a": 1,
"b": 2,
"x": 3,
}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that renaming happens correctly
assert "a" not in processed_obs
assert processed_obs["b"] == 1 # 'a' renamed to 'b'
assert processed_obs["c"] == 2 # original 'b' renamed to 'c'
assert processed_obs["x"] == 3
def test_partial_rename():
"""Test renaming only some keys."""
rename_map = {
"observation.state": "observation.proprio_state",
"pixels": "observation.image",
}
processor = RenameProcessor(rename_map=rename_map)
observation = {
"observation.state": torch.randn(10),
"pixels": np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8),
"reward": 1.0,
"info": {"episode": 1},
}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check renamed keys
assert "observation.proprio_state" in processed_obs
assert "observation.image" in processed_obs
assert "observation.state" not in processed_obs
assert "pixels" not in processed_obs
# Check unchanged keys
assert processed_obs["reward"] == 1.0
assert processed_obs["info"] == {"episode": 1}
def test_get_config():
"""Test configuration serialization."""
rename_map = {
"old1": "new1",
"old2": "new2",
}
processor = RenameProcessor(rename_map=rename_map)
config = processor.get_config()
assert config == {"rename_map": rename_map}
def test_state_dict():
"""Test state dict (should be empty for RenameProcessor)."""
processor = RenameProcessor(rename_map={"old": "new"})
state = processor.state_dict()
assert state == {}
# Load state dict should work even with empty dict
processor.load_state_dict({})
def test_integration_with_robot_processor():
"""Test integration with RobotProcessor pipeline."""
rename_map = {
"agent_pos": "observation.state",
"pixels": "observation.image",
}
rename_processor = RenameProcessor(rename_map=rename_map)
pipeline = RobotProcessor([rename_processor])
observation = {
"agent_pos": np.array([1.0, 2.0, 3.0]),
"pixels": np.zeros((32, 32, 3), dtype=np.uint8),
"other_data": "preserve_me",
}
transition = create_transition(
observation=observation, reward=0.5, done=False, truncated=False, info={}, complementary_data={}
)
result = pipeline(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check renaming worked through pipeline
assert "observation.state" in processed_obs
assert "observation.image" in processed_obs
assert "agent_pos" not in processed_obs
assert "pixels" not in processed_obs
assert processed_obs["other_data"] == "preserve_me"
# Check other transition elements unchanged
assert result[TransitionKey.REWARD] == 0.5
assert result[TransitionKey.DONE] is False
def test_save_and_load_pretrained():
"""Test saving and loading processor with RobotProcessor."""
rename_map = {
"old_state": "observation.state",
"old_image": "observation.image",
}
processor = RenameProcessor(rename_map=rename_map)
pipeline = RobotProcessor([processor], name="TestRenameProcessor")
with tempfile.TemporaryDirectory() as tmp_dir:
# Save pipeline
pipeline.save_pretrained(tmp_dir)
# Check files were created
config_path = Path(tmp_dir) / "testrenameprocessor.json" # Based on name="TestRenameProcessor"
assert config_path.exists()
# No state files should be created for RenameProcessor
state_files = list(Path(tmp_dir).glob("*.safetensors"))
assert len(state_files) == 0
# Load pipeline
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir)
assert loaded_pipeline.name == "TestRenameProcessor"
assert len(loaded_pipeline) == 1
# Check that loaded processor works correctly
loaded_processor = loaded_pipeline.steps[0]
assert isinstance(loaded_processor, RenameProcessor)
assert loaded_processor.rename_map == rename_map
# Test functionality after loading
observation = {"old_state": [1, 2, 3], "old_image": "image_data"}
transition = create_transition(observation=observation)
result = loaded_pipeline(transition)
processed_obs = result[TransitionKey.OBSERVATION]
assert "observation.state" in processed_obs
assert "observation.image" in processed_obs
assert processed_obs["observation.state"] == [1, 2, 3]
assert processed_obs["observation.image"] == "image_data"
def test_registry_functionality():
"""Test that RenameProcessor is properly registered."""
# Check that it's registered
assert "rename_processor" in ProcessorStepRegistry.list()
# Get from registry
retrieved_class = ProcessorStepRegistry.get("rename_processor")
assert retrieved_class is RenameProcessor
# Create instance from registry
instance = retrieved_class(rename_map={"old": "new"})
assert isinstance(instance, RenameProcessor)
assert instance.rename_map == {"old": "new"}
def test_registry_based_save_load():
"""Test save/load using registry name instead of module path."""
processor = RenameProcessor(rename_map={"key1": "renamed_key1"})
pipeline = RobotProcessor([processor])
with tempfile.TemporaryDirectory() as tmp_dir:
# Save and load
pipeline.save_pretrained(tmp_dir)
# Verify config uses registry name
import json
with open(Path(tmp_dir) / "robotprocessor.json") as f: # Default name is "RobotProcessor"
config = json.load(f)
assert "registry_name" in config["steps"][0]
assert config["steps"][0]["registry_name"] == "rename_processor"
assert "class" not in config["steps"][0] # Should use registry, not module path
# Load should work
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir)
loaded_processor = loaded_pipeline.steps[0]
assert isinstance(loaded_processor, RenameProcessor)
assert loaded_processor.rename_map == {"key1": "renamed_key1"}
def test_chained_rename_processors():
"""Test multiple RenameProcessors in a pipeline."""
# First processor: rename raw keys to intermediate format
processor1 = RenameProcessor(
rename_map={
"pos": "agent_position",
"img": "camera_image",
}
)
# Second processor: rename to final format
processor2 = RenameProcessor(
rename_map={
"agent_position": "observation.state",
"camera_image": "observation.image",
}
)
pipeline = RobotProcessor([processor1, processor2])
observation = {
"pos": np.array([1.0, 2.0]),
"img": "image_data",
"extra": "keep_me",
}
transition = create_transition(observation=observation)
# Step through to see intermediate results
results = list(pipeline.step_through(transition))
# After first processor
assert "agent_position" in results[1][TransitionKey.OBSERVATION]
assert "camera_image" in results[1][TransitionKey.OBSERVATION]
# After second processor
final_obs = results[2][TransitionKey.OBSERVATION]
assert "observation.state" in final_obs
assert "observation.image" in final_obs
assert final_obs["extra"] == "keep_me"
# Original keys should be gone
assert "pos" not in final_obs
assert "img" not in final_obs
assert "agent_position" not in final_obs
assert "camera_image" not in final_obs
def test_nested_observation_rename():
"""Test renaming with nested observation structures."""
rename_map = {
"observation.images.left": "observation.camera.left_view",
"observation.images.right": "observation.camera.right_view",
"observation.proprio": "observation.proprioception",
}
processor = RenameProcessor(rename_map=rename_map)
observation = {
"observation.images.left": torch.randn(3, 64, 64),
"observation.images.right": torch.randn(3, 64, 64),
"observation.proprio": torch.randn(7),
"observation.gripper": torch.tensor([0.0]), # Not renamed
}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check renames
assert "observation.camera.left_view" in processed_obs
assert "observation.camera.right_view" in processed_obs
assert "observation.proprioception" in processed_obs
# Check unchanged key
assert "observation.gripper" in processed_obs
# Check old keys removed
assert "observation.images.left" not in processed_obs
assert "observation.images.right" not in processed_obs
assert "observation.proprio" not in processed_obs
def test_value_types_preserved():
"""Test that various value types are preserved during renaming."""
rename_map = {"old_tensor": "new_tensor", "old_array": "new_array", "old_scalar": "new_scalar"}
processor = RenameProcessor(rename_map=rename_map)
tensor_value = torch.randn(3, 3)
array_value = np.random.rand(2, 2)
observation = {
"old_tensor": tensor_value,
"old_array": array_value,
"old_scalar": 42,
"old_string": "hello",
"old_dict": {"nested": "value"},
"old_list": [1, 2, 3],
}
transition = create_transition(observation=observation)
result = processor(transition)
processed_obs = result[TransitionKey.OBSERVATION]
# Check that values and types are preserved
assert torch.equal(processed_obs["new_tensor"], tensor_value)
assert np.array_equal(processed_obs["new_array"], array_value)
assert processed_obs["new_scalar"] == 42
assert processed_obs["old_string"] == "hello"
assert processed_obs["old_dict"] == {"nested": "value"}
assert processed_obs["old_list"] == [1, 2, 3]
def test_feature_contract_basic_renaming(policy_feature_factory):
processor = RenameProcessor(rename_map={"a": "x", "b": "y"})
features = {
"a": policy_feature_factory(FeatureType.STATE, (2,)),
"b": policy_feature_factory(FeatureType.ACTION, (3,)),
"c": policy_feature_factory(FeatureType.ENV, (1,)),
}
out = processor.feature_contract(features.copy())
# Values preserved and typed
assert out["x"] == features["a"]
assert out["y"] == features["b"]
assert out["c"] == features["c"]
assert_contract_is_typed(out)
# Input not mutated
assert set(features) == {"a", "b", "c"}
def test_feature_contract_overlapping_keys(policy_feature_factory):
# Overlapping renames: both 'a' and 'b' exist. 'a'->'b', 'b'->'c'
processor = RenameProcessor(rename_map={"a": "b", "b": "c"})
features = {
"a": policy_feature_factory(FeatureType.STATE, (1,)),
"b": policy_feature_factory(FeatureType.STATE, (2,)),
}
out = processor.feature_contract(features)
assert set(out) == {"b", "c"}
assert out["b"] == features["a"] # 'a' renamed to'b'
assert out["c"] == features["b"] # 'b' renamed to 'c'
assert_contract_is_typed(out)
def test_feature_contract_chained_processors(policy_feature_factory):
# Chain two rename processors at the contract level
processor1 = RenameProcessor(rename_map={"pos": "agent_position", "img": "camera_image"})
processor2 = RenameProcessor(
rename_map={"agent_position": "observation.state", "camera_image": "observation.image"}
)
pipeline = RobotProcessor([processor1, processor2])
spec = {
"pos": policy_feature_factory(FeatureType.STATE, (7,)),
"img": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"extra": policy_feature_factory(FeatureType.ENV, (1,)),
}
out = pipeline.feature_contract(initial_features=spec)
assert set(out) == {"observation.state", "observation.image", "extra"}
assert out["observation.state"] == spec["pos"]
assert out["observation.image"] == spec["img"]
assert out["extra"] == spec["extra"]
assert_contract_is_typed(out)