Compare commits

..

56 Commits

Author SHA1 Message Date
Michel Aractingi b2d7eecdb4 format docs 2025-08-05 10:36:22 +02:00
Michel Aractingi 0710f3a0f1 Remove dataset and mode from HilSerlEnvConfig to a GymManipulatorConfig to reduce verbose of configs during training 2025-08-05 10:35:54 +02:00
Michel Aractingi 9effc5214f Refactored hilserl config 2025-08-05 01:24:46 +02:00
Michel Aractingi b292dbbc55 change folder structure to reduce the size of gym_manip 2025-08-04 16:41:42 +02:00
Michel Aractingi f49280e89b RL works at this commit - fixed actor.py and bugs in gym_manipulator 2025-08-03 23:21:13 +02:00
pre-commit-ci[bot] ff38a51df9 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-02 17:07:19 +00:00
Michel Aractingi cfa672129e Refactored actor.py to use the pipeline 2025-08-02 19:06:56 +02:00
Michel Aractingi e6e1edfd74 Added the replay functionality with the pipeline 2025-08-02 17:57:27 +02:00
Michel Aractingi 384101731e Added the capability to record a dataset 2025-08-02 17:14:14 +02:00
Michel Aractingi 1fdbecad3c Migrate gym_manipulator to use the pipeline
Added get_teleop_events function to capture relevant events from teleop devices unrelated to actions
2025-08-01 20:20:13 +02:00
Pepijn 2c4e888c7f Feat/pipeline add feature contract (#1637)
* Add feature contract to pipelinestep and pipeline

* Add tests

* Add processor tests

* PR feedback

* encorperate pr feedback

* type in doc

* oops
2025-08-01 08:41:54 +02:00
Adil Zouitine 5ced72e6b8 docs(pipeline): Add clarification for repo name sanitization process 2025-08-01 08:41:54 +02:00
Adil Zouitine 907023f9f7 refactor(pipeline): Improve state file naming conventions for clarity and uniqueness
- Enhanced state file naming to include the processor's sanitized name, ensuring uniqueness when multiple processors are saved in the same directory.
- Updated tests to reflect changes in state file naming, verifying that filenames now include the processor name and step indices to prevent conflicts.
- Added a new test to validate state file naming when using multiple processors, ensuring distinct filenames for each processor's state files.
2025-08-01 08:41:54 +02:00
Adil Zouitine 4ba23ea029 feat(pipeline): Enhance configuration filename handling and state file naming
- Introduced support for custom configuration filenames in the `save_pretrained` method, allowing users to specify a filename instead of the default.
- Improved state file naming to include step indices, preventing conflicts when multiple processors of the same type are saved.
- Added automatic detection for configuration files when loading from a directory, with error handling for multiple files.
- Updated tests to validate new features, including custom filenames and automatic config detection.
2025-08-01 08:41:54 +02:00
Adil Zouitine 409ac0baca chore(doc): address pip install commant lerobot that not exist yet 2025-08-01 08:41:54 +02:00
Adil Zouitine 699363f9fc refactor(pipeline): Enhance state filename generation and profiling method
- Updated state filename generation to use the registry name when available, improving clarity in saved files.
- Modified the profile_steps method to include a warmup_runs parameter, allowing for more controlled performance profiling.
- Ensured consistent conditions during profiling by deep copying transitions for each run, enhancing accuracy in timing results.
2025-08-01 08:41:54 +02:00
Adil Zouitine ae7a54de57 refactor(pipeline): Utilize get_safe_torch_device for device assignment
- Replaced direct torch.device instantiation with get_safe_torch_device to ensure safe device handling.
- This change enhances code readability and maintains consistency in device management across the RobotProcessor class.
2025-08-01 08:41:54 +02:00
Adil Zouitine fb9139b882 chore(pipeline): Move _CFG_NAME along other class member 2025-08-01 08:41:54 +02:00
Adil Zouitine 9fe3a3fb17 feat(pipeline): Add __repr__ method to RobotProcessor for improved readability
- Implemented a __repr__ method in the RobotProcessor class to provide a clear string representation of the processor, including step names and optional parameters like name and seed.
- Added comprehensive tests to validate the __repr__ output for various scenarios, including empty processors, single and multiple steps, custom names, and seed values.
- Ensured that the representation handles long lists of steps with truncation for better readability.
2025-08-01 08:41:54 +02:00
Adil Zouitine 26cb9a24c3 refactor(pipeline): Clarify hook behavior and improve documentation
- Updated the RobotProcessor class to ensure hooks are strictly for observation and do not modify transitions, enhancing clarity and maintainability.
- Refactored hook registration methods to reflect the new behavior, ensuring they accept only functions that do not return modified transitions.
- Enhanced documentation to clearly outline the purpose of hooks and their execution semantics.
- Added tests to verify that hooks are not executed during the step_through method while ensuring they function correctly during the __call__ method.
2025-08-01 08:41:54 +02:00
Adil Zouitine 77106697c3 feat(pipeline): Add hook unregistration functionality and enhance documentation
- Implemented methods to unregister before, after, and reset hooks in the RobotProcessor class, allowing for more flexible hook management.
- Enhanced documentation to clarify hook execution semantics and the implications of modifying transitions within hooks.
- Added comprehensive tests to verify the correct behavior of hook registration and unregistration, including error handling for non-existent hooks.
2025-08-01 08:41:54 +02:00
Adil Zouitine 75bc44c166 refactor(observation_processor): Improve observation processing by using constants and simplifying pixel handling
- Introduced constants for observation keys to enhance readability.
- Streamlined the handling of the "pixels" key by copying observations first and processing images more clearly.
- Updated the environment state and agent position assignments to use the new constants, improving maintainability.
2025-08-01 08:41:54 +02:00
Adil Zouitine f2b79656eb refactor(pipeline): Transition from tuple to dictionary format for EnvTransition
- Updated the EnvTransition structure to use a dictionary format instead of a tuple, enhancing readability and maintainability.
- Replaced instances of TransitionIndex with TransitionKey for accessing transition components.
- Adjusted related processing functions and tests to accommodate the new dictionary format, ensuring consistent handling of transitions across the codebase.
2025-08-01 08:41:53 +02:00
pre-commit-ci[bot] 14c2ece004 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-01 08:41:53 +02:00
Adil Zouitine 35612c61e1 refactor(pipeline): Introduce ComplementaryDataProcessor for handling complementary data in transitions 2025-08-01 08:41:53 +02:00
pre-commit-ci[bot] f7bb3e2d90 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-01 08:41:53 +02:00
Adil Zouitine 1e0d667a22 Apply suggestions from code review
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Signed-off-by: Adil Zouitine <adilzouitinegm@gmail.com>
2025-08-01 08:41:53 +02:00
Adil Zouitine 33969a0337 refactor(pipeline): Simplify observation and padding data handling in batch transitions 2025-08-01 08:41:53 +02:00
Adil Zouitine fa26290e8c feat(pipeline): Enhance step_through method to support both tuple and dict inputs 2025-08-01 08:41:53 +02:00
Adil Zouitine e9f7f5127b chore(learner): nit comment from copilot 2025-08-01 08:41:53 +02:00
Adil Zouitine 097842c70f chore(normalization): addressing comments from copilot 2025-08-01 08:41:53 +02:00
Adil Zouitine 3b8a3a32a0 feat (overrides): Implement support for loading processors with parameter overrides
- Added the ability to provide non-serializable objects when loading processors from saved configurations using the `overrides` parameter.
- Enhanced error handling for invalid override keys and instantiation errors.
- Updated documentation and examples to illustrate the usage of overrides for both registered and unregistered steps.
- Added comprehensive tests to validate the new functionality and ensure backward compatibility.
2025-08-01 08:41:53 +02:00
Adil Zouitine 1c56779dd9 chore (type): add typing for multiprocess envs 2025-08-01 08:41:53 +02:00
Adil Zouitine 83a4338f8b chore (output format): improves output format 2025-08-01 08:41:53 +02:00
Adil Zouitine 730c7b2f35 fix(test): linting issue 2025-08-01 08:41:53 +02:00
pre-commit-ci[bot] 116059a43e [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-01 08:41:53 +02:00
Adil Zouitine b08149a113 chore (batch handling): Enhance processing components with batch conversion utilities 2025-08-01 08:41:53 +02:00
Adil Zouitine c227107f60 feat (device processor): Implement device processor 2025-08-01 08:41:53 +02:00
Adil Zouitine 01dc289f3d chore (docstrin):Improve docstring for NormalizerProcessor 2025-08-01 08:41:53 +02:00
Adil Zouitine 6830ca7645 Refactor normalization components and update tests
- Renamed `ObservationNormalizer` to `NormalizerProcessor` and `ActionUnnormalizer` to `UnnormalizerProcessor` for clarity.
- Consolidated normalization logic for both observations and actions into `NormalizerProcessor` and `UnnormalizerProcessor`.
- Updated tests to reflect the new class names and ensure proper functionality of normalization and unnormalization processes.
- Enhanced handling of missing statistics in normalization processes.
2025-08-01 08:41:52 +02:00
Adil Zouitine ed42c71fc3 fix(test): import issue 2025-08-01 08:41:52 +02:00
Adil Zouitine e0139065bd chore(test): add suggestion made by copilot regarding numpy test 2025-08-01 08:41:52 +02:00
Adil Zouitine e509f255af Update tests/processor/test_observation_processor.py
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Adil Zouitine <adilzouitinegm@gmail.com>
2025-08-01 08:41:52 +02:00
Adil Zouitine e2fcd140b0 fix(test): policies 2025-08-01 08:41:52 +02:00
Adil Zouitine 2a7a0e6129 fix (test): test factory 2025-08-01 08:41:52 +02:00
Adil Zouitine 9f33791b19 chore (docs): add docstring for processor 2025-08-01 08:41:52 +02:00
Adil Zouitine 453e0a995f Enhance processing architecture with new components
- Added `RenameProcessor` to facilitate key renaming in observations, improving data handling flexibility.
- Updated `__init__.py` to include `RenameProcessor` in module exports.
- Refactored `NormalizationProcessor` and `ObservationNormalizer` to use `rsplit` for better key handling.
- Introduced comprehensive tests for `NormalizationProcessor` and `RenameProcessor` to ensure functionality and robustness.
2025-08-01 08:41:52 +02:00
pre-commit-ci[bot] 8ebf79c494 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-01 08:41:52 +02:00
Adil Zouitine 8774aec304 Add normalization processor and related components
- Introduced `NormalizationProcessor` to handle both observation normalization and action unnormalization.
- Added `ObservationNormalizer` and `ActionUnnormalizer` classes for specific normalization tasks.
- Updated `__init__.py` to include the new `NormalizationProcessor` in the module exports.
- Enhanced `ObservationProcessor` with registration in the `ProcessorStepRegistry` for better modularity.
- Created `RenameProcessor` for renaming keys in observations, improving flexibility in data processing.
2025-08-01 08:41:52 +02:00
pre-commit-ci[bot] ac742c9f0d [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-01 08:41:52 +02:00
Adil Zouitine cd13f1ecfd Add RobotProcessor tutorial to documentation
- Introduced a new tutorial on using RobotProcessor for preprocessing robot data.
- Added a section in the table of contents for easy navigation to the new tutorial.
- The tutorial covers key concepts, real-world scenarios, and practical examples for effective use of the RobotProcessor pipeline.
2025-08-01 08:41:52 +02:00
Adil Zouitine 9aa632968f Refactor processing architecture to use RobotProcessor
- Replaced instances of RobotPipeline with RobotProcessor across the codebase for improved modularity and clarity.
- Introduced ProcessorStepRegistry for better management of processing steps.
- Updated relevant documentation and tests to reflect the new processing structure.
- Enhanced the save/load functionality to support the new processor design.
- Added a model card template for RobotProcessor to facilitate sharing and documentation.
2025-08-01 08:41:52 +02:00
Adil Zouitine 62caaf07b0 Remove redundant tests for None observation and serialization methods in test_observation_processor.py to streamline the test suite and improve maintainability. 2025-08-01 08:41:52 +02:00
Adil Zouitine 3355f04ca6 Refactor observation processing and improve modularity
- Updated `ObservationProcessor` to enhance the modular design for processing observations.
- Cleaned up imports and improved code readability by removing unnecessary lines and comments.
- Ensured backward compatibility while integrating new processing components.
- Added tests to validate the functionality of the updated processing architecture.
2025-08-01 08:41:52 +02:00
pre-commit-ci[bot] 769f531603 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-01 08:41:51 +02:00
Adil Zouitine f6c7287ae7 Refactor observation preprocessing to use a modular pipeline system
- Introduced `RobotPipeline` and `ObservationProcessor` for handling observation transformations.
- Updated `preprocess_observation` to maintain backward compatibility while leveraging the new pipeline.
- Added tests for the new processing components and ensured they match the original functionality.
- Removed hardcoded logic in favor of a more flexible, composable architecture.
2025-08-01 08:41:51 +02:00
54 changed files with 4453 additions and 4189 deletions
+7 -45
View File
@@ -19,11 +19,6 @@ 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:
@@ -55,7 +50,6 @@ 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 }}
@@ -92,29 +86,13 @@ jobs:
env: env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# zizmor: ignore[template-injection] # zizmor: ignore[template-injection]
run: | run: gh release create ${{ github.ref_name }} --release-name "Release ${{ github.ref_name }}" --generate-notes ./dist/*
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 TestPyPI for pre-releases
# 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]
with:
repository-url: https://test.pypi.org/legacy/
verbose: true
print-hash: true
- name: Publish to PyPI - name: Publish to PyPI
if: startsWith(github.ref, 'refs/tags/v') && !contains(github.ref, '-') if: startsWith(github.ref, 'refs/tags/v')
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:
verbose: true password: ${{ secrets.PYPI_API_TOKEN }}
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:
@@ -141,31 +119,15 @@ 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
# zizmor: ignore[template-injection] run: uv run pip install lerobot==${{ needs.build-and-publish.outputs.version }} # 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 python -c "import lerobot; print(lerobot.__version__)" run: uv run 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 weekly # TODO(Steven): Publish draft/pre-release and to test pypi
# 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
+167 -74
View File
@@ -1,21 +1,25 @@
<p align="center"> <p align="center">
<img alt="LeRobot, Hugging Face Robotics Library" src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/lerobot-logo-thumbnail.png" width="100%"> <picture>
<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.yml/badge.svg?branch=main)](https://github.com/huggingface/lerobot/actions/workflows/nighty.yml?query=branch%3Amain) [![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)
[![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/)
[![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-v2.1-ff69b4.svg)](https://github.com/huggingface/lerobot/blob/main/CODE_OF_CONDUCT.md) [![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%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">
@@ -25,10 +29,10 @@
<div align="center"> <div align="center">
<img <img
src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/hope_jr/hopejr.png" src="media/hope_jr/hopejr.png?raw=true"
alt="HopeJR robot" alt="HopeJR robot"
title="HopeJR robot" title="HopeJR robot"
width="60%" style="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>
@@ -47,12 +51,20 @@
</h2> </h2>
<div align="center"> <div align="center">
<table> <div style="display: flex; gap: 1rem; justify-content: center; align-items: center;" >
<tr> <img
<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> src="media/so101/so101.webp?raw=true"
<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> alt="SO-101 follower arm"
</tr> title="SO-101 follower arm"
</table> style="width: 40%;"
/>
<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>
@@ -64,7 +76,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="https://raw.githubusercontent.com/huggingface/lerobot/main/media/lekiwi/kiwi.webp" alt="LeKiwi mobile robot" title="LeKiwi mobile robot" width="50%"> <img src="media/lekiwi/kiwi.webp?raw=true" alt="LeKiwi mobile robot" title="LeKiwi mobile robot" width="50%">
</div> </div>
<br/> <br/>
@@ -87,9 +99,9 @@
<table> <table>
<tr> <tr>
<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/aloha_act.gif" width="100%" alt="ACT policy on ALOHA 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/simxarm_tdmpc.gif" width="100%" alt="TDMPC policy on SimXArm 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> <td><img src="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>
@@ -98,11 +110,23 @@
</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
LeRobot works with Python 3.10+ and PyTorch 2.2+. Download our source code:
### Environment Setup ```bash
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):
@@ -127,18 +151,7 @@ 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 .
@@ -159,34 +172,6 @@ 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
@@ -197,7 +182,7 @@ wandb login
### Visualize datasets ### Visualize datasets
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. 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.
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:
@@ -227,7 +212,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](https://github.com/huggingface/lerobot/blob/main/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](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.
@@ -271,7 +256,7 @@ Dataset can be uploaded/downloaded from the HuggingFace hub seamlessly. To work
### Evaluate a pretrained policy ### Evaluate a pretrained policy
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. 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.
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):
@@ -295,13 +280,13 @@ See `python -m lerobot.scripts.eval --help` for more instructions.
### Train your own policy ### Train your own policy
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. 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.
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](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. 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.
\<img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/wandb.png" alt="WandB logs example"\> ![](media/wandb.png)
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 `python -m lerobot.scripts.eval --help` for more instructions.
@@ -320,6 +305,26 @@ 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)).
@@ -336,16 +341,34 @@ 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/src/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/lerobot/scripts/eval.py) for an example of how other people may use your policy.
### Acknowledgment ### Improve your code with profiling
- The LeRobot team 🤗 for building SmolVLA [Paper](https://arxiv.org/abs/2506.01844), [Blog](https://huggingface.co/blog/smolvla). An example of a code snippet to profile the evaluation of a policy:
- 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). <!-- prettier-ignore-start -->
- 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). ```python
- Thanks to Antonio Loquercio and Ashish Kumar for their early support. from torch.profiler import profile, record_function, ProfilerActivity
- 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
@@ -353,13 +376,83 @@ 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 Pascal, 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 Pascale, 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)
-3
View File
@@ -1,3 +0,0 @@
# docs-requirements.txt
hf-doc-builder @ git+https://github.com/huggingface/doc-builder.git@main
watchdog>=6.0.0
+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 . -r docs-requirements.txt pip install -e ".[docs]"
``` ```
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)
+2
View File
@@ -13,6 +13,8 @@
title: Cameras title: Cameras
- local: integrate_hardware - local: integrate_hardware
title: Bring Your Own Hardware title: Bring Your Own Hardware
- local: processor_tutorial
title: RobotProcessor Pipeline
- local: hilserl - local: hilserl
title: Train a Robot with RL title: Train a Robot with RL
- local: hilserl_sim - local: hilserl_sim
+125 -55
View File
@@ -56,27 +56,41 @@ pip install -e ".[hilserl]"
### Understanding Configuration ### Understanding Configuration
The training process begins with proper configuration for the HILSerl environment. The configuration class of interest is `HILSerlRobotEnvConfig` in `lerobot/envs/configs.py`. Which is defined as: The training process begins with proper configuration for the HILSerl environment. The main configuration class is `GymManipulatorConfig` in `lerobot/scripts/rl/gym_manipulator.py`, which contains nested `HILSerlRobotEnvConfig` and `DatasetConfig`. The configuration is organized into focused, nested sub-configs:
<!-- prettier-ignore-start --> <!-- prettier-ignore-start -->
```python ```python
class GymManipulatorConfig:
env: HILSerlRobotEnvConfig # Environment configuration (nested)
dataset: DatasetConfig # Dataset recording/replay configuration (nested)
mode: str | None = None # "record", "replay", or None (for training)
class HILSerlRobotEnvConfig(EnvConfig): class HILSerlRobotEnvConfig(EnvConfig):
robot: RobotConfig | None = None # Main robot agent (defined in `lerobot/robots`) robot: RobotConfig | None = None # Main robot agent (defined in `lerobot/robots`)
teleop: TeleoperatorConfig | None = None # Teleoperator agent, e.g., gamepad or leader arm, (defined in `lerobot/teleoperators`) teleop: TeleoperatorConfig | None = None # Teleoperator agent, e.g., gamepad or leader arm
wrapper: EnvTransformConfig | None = None # Environment wrapper settings; check `lerobot/scripts/server/gym_manipulator.py` processor: HILSerlProcessorConfig # Processing pipeline configuration (nested)
fps: int = 10 # Control frequency
name: str = "real_robot" # Environment name name: str = "real_robot" # Environment name
mode: str = None # "record", "replay", or None (for training)
repo_id: str | None = None # LeRobot dataset repository ID
dataset_root: str | None = None # Local dataset root (optional)
task: str = "" # Task identifier
num_episodes: int = 10 # Number of episodes for recording
episode: int = 0 # episode index for replay
device: str = "cuda" # Compute device device: str = "cuda" # Compute device
push_to_hub: bool = True # Whether to push the recorded datasets to Hub fps: int = 30 # Control frequency
pretrained_policy_name_or_path: str | None = None # For policy loading
reward_classifier_pretrained_path: str | None = None # For reward model # Nested processor configuration
number_of_steps_after_success: int = 0 # For reward classifier, collect more positive examples after a success to train a classifier class HILSerlProcessorConfig:
control_mode: str = "gamepad" # Control mode
observation: ObservationConfig # Observation processing settings
image_preprocessing: ImagePreprocessingConfig # Image crop/resize settings
gripper: GripperConfig # Gripper control and penalty settings
reset: ResetConfig # Environment reset and timing settings
inverse_kinematics: InverseKinematicsConfig # IK processing settings
reward_classifier: RewardClassifierConfig # Reward classifier settings
# Dataset configuration
class DatasetConfig:
repo_id: str # LeRobot dataset repository ID
dataset_root: str | None = None # Local dataset root (optional)
task: str # Task identifier
num_episodes: int # Number of episodes for recording
episode: int # Episode index for replay
push_to_hub: bool # Whether to push datasets to Hub
``` ```
<!-- prettier-ignore-end --> <!-- prettier-ignore-end -->
@@ -130,22 +144,31 @@ With the bounds defined, you can safely collect demonstrations for training. Tra
Create a configuration file for recording demonstrations (or edit an existing one like [env_config_so100.json](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json)): Create a configuration file for recording demonstrations (or edit an existing one like [env_config_so100.json](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json)):
1. Set `mode` to `"record"` 1. Set `mode` to `"record"` at the root level
2. Specify a unique `repo_id` for your dataset (e.g., "username/task_name") 2. Specify a unique `repo_id` for your dataset in the `dataset` section (e.g., "username/task_name")
3. Set `num_episodes` to the number of demonstrations you want to collect 3. Set `num_episodes` in the `dataset` section to the number of demonstrations you want to collect
4. Set `crop_params_dict` to `null` initially (we'll determine crops later) 4. Set `env.processor.image_preprocessing.crop_params_dict` to `{}` initially (we'll determine crops later)
5. Configure `robot`, `cameras`, and other hardware settings 5. Configure `env.robot`, `env.teleop`, and other hardware settings in the `env` section
Example configuration section: Example configuration section:
```json ```json
"mode": "record", {
"repo_id": "username/pick_lift_cube", "env": {
"dataset_root": null, "type": "gym_manipulator",
"task": "pick_and_lift", "fps": 10
"num_episodes": 15, // ... robot, teleop, processor configs ...
"episode": 0, },
"push_to_hub": true "dataset": {
"repo_id": "username/pick_lift_cube",
"dataset_root": null,
"task": "pick_and_lift",
"num_episodes": 15,
"episode": 0,
"push_to_hub": true
},
"mode": "record"
}
``` ```
### Using a Teleoperation Device ### Using a Teleoperation Device
@@ -191,10 +214,17 @@ The gamepad provides a very convenient way to control the robot and the episode
To setup the gamepad, you need to set the `control_mode` to `"gamepad"` and define the `teleop` section in the configuration file. To setup the gamepad, you need to set the `control_mode` to `"gamepad"` and define the `teleop` section in the configuration file.
```json ```json
{
"env": {
"teleop": { "teleop": {
"type": "gamepad", "type": "gamepad",
"use_gripper": true "use_gripper": true
}, },
"processor": {
"control_mode": "gamepad"
}
}
}
``` ```
<p align="center"> <p align="center">
@@ -216,11 +246,18 @@ The SO101 leader arm has reduced gears that allows it to move and track the foll
To setup the SO101 leader, you need to set the `control_mode` to `"leader"` and define the `teleop` section in the configuration file. To setup the SO101 leader, you need to set the `control_mode` to `"leader"` and define the `teleop` section in the configuration file.
```json ```json
{
"env": {
"teleop": { "teleop": {
"type": "so101_leader", "type": "so101_leader",
"port": "/dev/tty.usbmodem585A0077921", # check your port number "port": "/dev/tty.usbmodem585A0077921",
"use_degrees": true "use_degrees": true
}, },
"processor": {
"control_mode": "leader"
}
}
}
``` ```
In order to annotate the success/failure of the episode, **you will need** to use a keyboard to press `s` for success, `esc` for failure. In order to annotate the success/failure of the episode, **you will need** to use a keyboard to press `s` for success, `esc` for failure.
@@ -251,7 +288,7 @@ python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/e
During recording: During recording:
1. The robot will reset to the initial position defined in the configuration file `fixed_reset_joint_positions` 1. The robot will reset to the initial position defined in the configuration file `env.processor.reset.fixed_reset_joint_positions`
2. Complete the task successfully 2. Complete the task successfully
3. The episode ends with a reward of 1 when you press the "success" button 3. The episode ends with a reward of 1 when you press the "success" button
4. If the time limit is reached, or the fail button is pressed, the episode ends with a reward of 0 4. If the time limit is reached, or the fail button is pressed, the episode ends with a reward of 0
@@ -310,11 +347,19 @@ observation.images.front: [180, 250, 120, 150]
Add these crop parameters to your training configuration: Add these crop parameters to your training configuration:
```json ```json
"crop_params_dict": { {
"observation.images.side": [180, 207, 180, 200], "env": {
"observation.images.front": [180, 250, 120, 150] "processor": {
}, "image_preprocessing": {
"resize_size": [128, 128] "crop_params_dict": {
"observation.images.side": [180, 207, 180, 200],
"observation.images.front": [180, 250, 120, 150]
},
"resize_size": [128, 128]
}
}
}
}
``` ```
**Recommended image resolution** **Recommended image resolution**
@@ -343,26 +388,35 @@ python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/r
**Key Parameters for Data Collection** **Key Parameters for Data Collection**
- **mode**: set it to `"record"` to collect a dataset - **mode**: set it to `"record"` to collect a dataset (at root level)
- **repo_id**: `"hf_username/dataset_name"`, name of the dataset and repo on the hub - **dataset.repo_id**: `"hf_username/dataset_name"`, name of the dataset and repo on the hub
- **num_episodes**: Number of episodes to record - **dataset.num_episodes**: Number of episodes to record
- **number_of_steps_after_success**: Number of additional frames to record after a success (reward=1) is detected - **env.processor.reset.number_of_steps_after_success**: Number of additional frames to record after a success (reward=1) is detected
- **fps**: Number of frames per second to record - **env.fps**: Number of frames per second to record
- **push_to_hub**: Whether to push the dataset to the hub - **dataset.push_to_hub**: Whether to push the dataset to the hub
The `number_of_steps_after_success` parameter is crucial as it allows you to collect more positive examples. When a success is detected, the system will continue recording for the specified number of steps while maintaining the reward=1 label. Otherwise, there won't be enough states in the dataset labeled to 1 to train a good classifier. The `env.processor.reset.number_of_steps_after_success` parameter is crucial as it allows you to collect more positive examples. When a success is detected, the system will continue recording for the specified number of steps while maintaining the reward=1 label. Otherwise, there won't be enough states in the dataset labeled to 1 to train a good classifier.
Example configuration section for data collection: Example configuration section for data collection:
```json ```json
{ {
"mode": "record", "env": {
"repo_id": "hf_username/dataset_name", "type": "gym_manipulator",
"dataset_root": "data/your_dataset", "fps": 10,
"num_episodes": 20, "processor": {
"push_to_hub": true, "reset": {
"fps": 10, "number_of_steps_after_success": 15
"number_of_steps_after_success": 15 }
}
},
"dataset": {
"repo_id": "hf_username/dataset_name",
"dataset_root": "data/your_dataset",
"num_episodes": 20,
"push_to_hub": true
},
"mode": "record"
} }
``` ```
@@ -421,9 +475,17 @@ To use your trained reward classifier, configure the `HILSerlRobotEnvConfig` to
<!-- prettier-ignore-start --> <!-- prettier-ignore-start -->
```python ```python
env_config = HILSerlRobotEnvConfig( config = GymManipulatorConfig(
reward_classifier_pretrained_path="path_to_your_pretrained_trained_model", env=HILSerlRobotEnvConfig(
# Other environment parameters processor=HILSerlProcessorConfig(
reward_classifier=RewardClassifierConfig(
pretrained_path="path_to_your_pretrained_trained_model"
)
),
# Other environment parameters
),
dataset=DatasetConfig(...),
mode=None # For training
) )
``` ```
<!-- prettier-ignore-end --> <!-- prettier-ignore-end -->
@@ -432,7 +494,15 @@ or set the argument in the json config file.
```json ```json
{ {
"reward_classifier_pretrained_path": "path_to_your_pretrained_model" "env": {
"processor": {
"reward_classifier": {
"pretrained_path": "path_to_your_pretrained_model",
"success_threshold": 0.7,
"success_reward": 1.0
}
}
}
} }
``` ```
+3 -3
View File
@@ -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/${HF_USER}/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/cadene/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
@@ -428,7 +428,7 @@ 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`](https://github.com/huggingface/lerobot/blob/main/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 [`python -m lerobot.scripts.train`](../src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
```bash ```bash
python -m lerobot.scripts.train \ python -m lerobot.scripts.train \
@@ -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`](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. 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.
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`.
+2 -2
View File
@@ -96,7 +96,7 @@ 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`](https://github.com/huggingface/lerobot/blob/main/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 [`python -m lerobot.scripts.train`](../src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
```bash ```bash
python -m lerobot.scripts.train \ python -m lerobot.scripts.train \
@@ -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`](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. 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.
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`.
+11 -39
View File
@@ -1,6 +1,15 @@
# Installation # Installation
## Environment Setup ## Install LeRobot
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)
@@ -31,49 +40,12 @@ 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`.
-14
View File
@@ -1,14 +0,0 @@
## 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}
}
```
-14
View File
@@ -1,14 +0,0 @@
## 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},
}
```
-14
View File
@@ -1,14 +0,0 @@
## 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}
}
```
-14
View File
@@ -1,14 +0,0 @@
## 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}
}
```
-14
View File
@@ -1,14 +0,0 @@
## 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}
}
```
File diff suppressed because it is too large Load Diff
+6 -5
View File
@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
[project] [project]
name = "lerobot" name = "lerobot"
version = "0.3.3" version = "0.2.0"
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,16 +68,15 @@ 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
@@ -126,6 +125,7 @@ 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,6 +147,7 @@ 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]",
-625
View File
@@ -1,625 +0,0 @@
# 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
-650
View File
@@ -1,650 +0,0 @@
#
# 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
-9
View File
@@ -1,9 +0,0 @@
# 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]
+1 -5
View File
@@ -82,9 +82,5 @@ def calibrate(cfg: CalibrateConfig):
device.disconnect() device.disconnect()
def main():
calibrate()
if __name__ == "__main__": if __name__ == "__main__":
main() calibrate()
+4 -5
View File
@@ -27,7 +27,6 @@ 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
@@ -120,8 +119,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_name, ft in self.input_features.items(): for _, ft in self.input_features.items():
if ft.type is FeatureType.STATE and ft_name == OBS_STATE: if ft.type is FeatureType.STATE:
return ft return ft
return None return None
@@ -138,8 +137,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
@property @property
def action_feature(self) -> PolicyFeature | None: def action_feature(self) -> PolicyFeature | None:
for ft_name, ft in self.output_features.items(): for _, ft in self.output_features.items():
if ft.type is FeatureType.ACTION and ft_name == ACTION: if ft.type is FeatureType.ACTION:
return ft return ft
return None return None
+58 -85
View File
@@ -161,35 +161,74 @@ class XarmEnv(EnvConfig):
@dataclass @dataclass
class VideoRecordConfig: class ImagePreprocessingConfig:
"""Configuration for video recording in ManiSkill environments.""" crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
resize_size: tuple[int, int] | None = None
enabled: bool = False
record_dir: str = "videos"
trajectory_name: str = "trajectory"
@dataclass @dataclass
class EnvTransformConfig: class RewardClassifierConfig:
"""Configuration for environment wrappers.""" """Configuration for reward classification."""
pretrained_path: str | None = None
success_threshold: float = 0.5
success_reward: float = 1.0
@dataclass
class InverseKinematicsConfig:
"""Configuration for inverse kinematics processing."""
urdf_path: str | None = None
target_frame_name: str | None = None
end_effector_bounds: dict[str, list[float]] | None = None
end_effector_step_sizes: dict[str, float] | None = None
max_gripper_pos: float | None = None
@dataclass
class ObservationConfig:
"""Configuration for observation processing."""
# ee_action_space_params: EEActionSpaceConfig = field(default_factory=EEActionSpaceConfig)
control_mode: str = "gamepad"
display_cameras: bool = False
add_joint_velocity_to_observation: bool = False add_joint_velocity_to_observation: bool = False
add_current_to_observation: bool = False add_current_to_observation: bool = False
add_ee_pose_to_observation: bool = False add_ee_pose_to_observation: bool = False
crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None display_cameras: bool = False
resize_size: tuple[int, int] | None = None
control_time_s: float = 20.0
fixed_reset_joint_positions: Any | None = None @dataclass
reset_time_s: float = 5.0 class GripperConfig:
"""Configuration for gripper control and penalties."""
use_gripper: bool = True use_gripper: bool = True
gripper_quantization_threshold: float | None = 0.8
gripper_penalty: float = 0.0 gripper_penalty: float = 0.0
gripper_penalty_in_reward: bool = False gripper_penalty_in_reward: bool = False
@dataclass
class ResetConfig:
"""Configuration for environment reset behavior."""
fixed_reset_joint_positions: Any | None = None
reset_time_s: float = 5.0
control_time_s: float = 20.0
terminate_on_success: bool = True
number_of_steps_after_success: int = 0
@dataclass
class HILSerlProcessorConfig:
"""Configuration for environment processing pipeline."""
control_mode: str = "gamepad"
observation: ObservationConfig = field(default_factory=ObservationConfig)
image_preprocessing: ImagePreprocessingConfig = field(default_factory=ImagePreprocessingConfig)
gripper: GripperConfig = field(default_factory=GripperConfig)
reset: ResetConfig = field(default_factory=ResetConfig)
inverse_kinematics: InverseKinematicsConfig = field(default_factory=InverseKinematicsConfig)
reward_classifier: RewardClassifierConfig = field(default_factory=RewardClassifierConfig)
@EnvConfig.register_subclass(name="gym_manipulator") @EnvConfig.register_subclass(name="gym_manipulator")
@dataclass @dataclass
class HILSerlRobotEnvConfig(EnvConfig): class HILSerlRobotEnvConfig(EnvConfig):
@@ -197,77 +236,11 @@ class HILSerlRobotEnvConfig(EnvConfig):
robot: RobotConfig | None = None robot: RobotConfig | None = None
teleop: TeleoperatorConfig | None = None teleop: TeleoperatorConfig | None = None
wrapper: EnvTransformConfig | None = None processor: HILSerlProcessorConfig = field(default_factory=HILSerlProcessorConfig)
fps: int = 10
name: str = "real_robot" name: str = "real_robot"
mode: str | None = None # Either "record", "replay", None
repo_id: str | None = None
dataset_root: str | None = None
task: str | None = ""
num_episodes: int = 10 # only for record mode
episode: int = 0
device: str = "cuda" device: str = "cuda"
push_to_hub: bool = True
pretrained_policy_name_or_path: str | None = None
reward_classifier_pretrained_path: str | None = None
# For the reward classifier, to record more positive examples after a success
number_of_steps_after_success: int = 0
@property @property
def gym_kwargs(self) -> dict: def gym_kwargs(self) -> dict:
return {} return {}
@EnvConfig.register_subclass("hil")
@dataclass
class HILEnvConfig(EnvConfig):
"""Configuration for the HIL environment."""
name: str = "PandaPickCube"
task: str | None = "PandaPickCubeKeyboard-v0"
use_viewer: bool = True
gripper_penalty: float = 0.0
use_gamepad: bool = True
state_dim: int = 18
action_dim: int = 4
fps: int = 100
episode_length: int = 100
video_record: VideoRecordConfig = field(default_factory=VideoRecordConfig)
features: dict[str, PolicyFeature] = field(
default_factory=lambda: {
"action": PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
"observation.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128)),
"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(18,)),
}
)
features_map: dict[str, str] = field(
default_factory=lambda: {
"action": ACTION,
"observation.image": OBS_IMAGE,
"observation.state": OBS_STATE,
}
)
################# args from hilserlrobotenv
reward_classifier_pretrained_path: str | None = None
robot_config: RobotConfig | None = None
teleop_config: TeleoperatorConfig | None = None
wrapper: EnvTransformConfig | None = None
mode: str | None = None # Either "record", "replay", None
repo_id: str | None = None
dataset_root: str | None = None
num_episodes: int = 10 # only for record mode
episode: int = 0
device: str = "cuda"
push_to_hub: bool = True
pretrained_policy_name_or_path: str | None = None
# For the reward classifier, to record more positive examples after a success
number_of_steps_after_success: int = 0
############################
@property
def gym_kwargs(self) -> dict:
return {
"use_viewer": self.use_viewer,
"use_gamepad": self.use_gamepad,
"gripper_penalty": self.gripper_penalty,
}
+21 -45
View File
@@ -16,10 +16,8 @@
import warnings import warnings
from typing import Any from typing import Any
import einops
import gymnasium as gym import gymnasium as gym
import numpy as np import numpy as np
import torch
from torch import Tensor from torch import Tensor
from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.configs.types import FeatureType, PolicyFeature
@@ -28,62 +26,40 @@ from lerobot.utils.utils import get_channel_first_image_shape
def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Tensor]: def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Tensor]:
# TODO(aliberts, rcadene): refactor this to use features from the environment (no hardcoding)
"""Convert environment observation to LeRobot format observation. """Convert environment observation to LeRobot format observation.
This function uses the new pipeline system internally but maintains
backward compatibility with the original interface.
Args: Args:
observation: Dictionary of observation batches from a Gym vector environment. observation: Dictionary of observation batches from a Gym vector environment.
Returns: Returns:
Dictionary of observation batches with keys renamed to LeRobot format and values as tensors. Dictionary of observation batches with keys renamed to LeRobot format and values as tensors.
""" """
# map to expected inputs for the policy from lerobot.processor import RobotProcessor, TransitionKey, VanillaObservationProcessor
return_observations = {}
if "pixels" in observations:
if isinstance(observations["pixels"], dict):
imgs = {f"observation.images.{key}": img for key, img in observations["pixels"].items()}
else:
imgs = {"observation.image": observations["pixels"]}
for imgkey, img in imgs.items(): # Create processor with observation processor
# TODO(aliberts, rcadene): use transforms.ToTensor()? processor = RobotProcessor([VanillaObservationProcessor()])
img = torch.from_numpy(img)
# When preprocessing observations in a non-vectorized environment, we need to add a batch dimension. # Create transition dictionary and process
# This is the case for human-in-the-loop RL where there is only one environment. transition = {
if img.ndim == 3: TransitionKey.OBSERVATION: observations,
img = img.unsqueeze(0) TransitionKey.ACTION: None,
# sanity check that images are channel last TransitionKey.REWARD: None,
_, h, w, c = img.shape TransitionKey.DONE: None,
assert c < h and c < w, f"expect channel last images, but instead got {img.shape=}" TransitionKey.TRUNCATED: None,
TransitionKey.INFO: None,
TransitionKey.COMPLEMENTARY_DATA: None,
}
result = processor(transition)
# sanity check that images are uint8 # Extract and return the processed observation
assert img.dtype == torch.uint8, f"expect torch.uint8, but instead {img.dtype=}" return result[TransitionKey.OBSERVATION]
# convert to channel first of type float32 in range [0,1]
img = einops.rearrange(img, "b h w c -> b c h w").contiguous()
img = img.type(torch.float32)
img /= 255
return_observations[imgkey] = img
if "environment_state" in observations:
env_state = torch.from_numpy(observations["environment_state"]).float()
if env_state.dim() == 1:
env_state = env_state.unsqueeze(0)
return_observations["observation.environment_state"] = env_state
# TODO(rcadene): enable pixels only baseline with `obs_type="pixels"` in environment by removing
agent_pos = torch.from_numpy(observations["agent_pos"]).float()
if agent_pos.dim() == 1:
agent_pos = agent_pos.unsqueeze(0)
return_observations["observation.state"] = agent_pos
return return_observations
def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]: def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
# TODO(aliberts, rcadene): remove this hardcoding of keys and just use the nested keys as is # TODO(aliberts, rcadene): remove this hardcoding of keys and just use the nested keys as is
# (need to also refactor preprocess_observation and externalize normalization from policies) # (need to externalize normalization from policies)
policy_features = {} policy_features = {}
for key, ft in env_cfg.features.items(): for key, ft in env_cfg.features.items():
if ft.type is FeatureType.VISUAL: if ft.type is FeatureType.VISUAL:
+1 -5
View File
@@ -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}")
def main(): if __name__ == "__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,7 +313,3 @@ def 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()
+1 -5
View File
@@ -61,9 +61,5 @@ 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}).")
def main():
find_port()
if __name__ == "__main__": if __name__ == "__main__":
main() find_port()
-1
View File
@@ -1 +0,0 @@
../../../../docs/source/policy_act_README.md
-1
View File
@@ -1 +0,0 @@
../../../../docs/source/policy_diffusion_README.md
-1
View File
@@ -1 +0,0 @@
../../../../docs/source/policy_smolvla_README.md
-1
View File
@@ -1 +0,0 @@
../../../../docs/source/policy_tdmpc_README.md
-1
View File
@@ -1 +0,0 @@
../../../../docs/source/policy_vqbet_README.md
+25 -1
View File
@@ -15,8 +15,18 @@
# limitations under the License. # limitations under the License.
from .device_processor import DeviceProcessor from .device_processor import DeviceProcessor
from .hil_processor import (
GripperPenaltyProcessor,
ImageCropResizeProcessor,
InterventionActionProcessor,
TimeLimitProcessor,
)
from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor
from .observation_processor import VanillaObservationProcessor from .observation_processor import (
ImageProcessor,
StateProcessor,
VanillaObservationProcessor,
)
from .pipeline import ( from .pipeline import (
ActionProcessor, ActionProcessor,
DoneProcessor, DoneProcessor,
@@ -32,14 +42,26 @@ from .pipeline import (
TruncatedProcessor, TruncatedProcessor,
) )
from .rename_processor import RenameProcessor from .rename_processor import RenameProcessor
from .robot_processor import (
InverseKinematicsProcessor,
JointVelocityProcessor,
MotorCurrentProcessor,
)
__all__ = [ __all__ = [
"ActionProcessor", "ActionProcessor",
"DeviceProcessor", "DeviceProcessor",
"DoneProcessor", "DoneProcessor",
"EnvTransition", "EnvTransition",
"GripperPenaltyProcessor",
"IdentityProcessor", "IdentityProcessor",
"ImageCropResizeProcessor",
"ImageProcessor",
"InfoProcessor", "InfoProcessor",
"InterventionActionProcessor",
"InverseKinematicsProcessor",
"JointVelocityProcessor",
"MotorCurrentProcessor",
"NormalizerProcessor", "NormalizerProcessor",
"UnnormalizerProcessor", "UnnormalizerProcessor",
"ObservationProcessor", "ObservationProcessor",
@@ -48,6 +70,8 @@ __all__ = [
"RenameProcessor", "RenameProcessor",
"RewardProcessor", "RewardProcessor",
"RobotProcessor", "RobotProcessor",
"StateProcessor",
"TimeLimitProcessor",
"TransitionKey", "TransitionKey",
"TruncatedProcessor", "TruncatedProcessor",
"VanillaObservationProcessor", "VanillaObservationProcessor",
+2 -4
View File
@@ -20,7 +20,6 @@ import torch
from lerobot.configs.types import PolicyFeature from lerobot.configs.types import PolicyFeature
from lerobot.processor.pipeline import EnvTransition, TransitionKey from lerobot.processor.pipeline import EnvTransition, TransitionKey
from lerobot.utils.utils import get_safe_torch_device
@dataclass @dataclass
@@ -31,11 +30,10 @@ class DeviceProcessor:
specified device (CPU or GPU) before they are returned. specified device (CPU or GPU) before they are returned.
""" """
device: torch.device = "cpu" device: str = "cpu"
def __post_init__(self): def __post_init__(self):
self.device = get_safe_torch_device(self.device) self.non_blocking = "cuda" in self.device
self.non_blocking = "cuda" in str(self.device)
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
# Create a copy of the transition # Create a copy of the transition
+331
View File
@@ -0,0 +1,331 @@
import time
from dataclasses import dataclass
from typing import Any
import torch
import torchvision.transforms.functional as F # noqa: N812
from lerobot.configs.types import PolicyFeature
from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
@dataclass
@ProcessorStepRegistry.register("image_crop_resize_processor")
class ImageCropResizeProcessor:
"""Crop and resize image observations."""
crop_params_dict: dict[str, tuple[int, int, int, int]]
resize_size: tuple[int, int] = (128, 128)
def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition.get(TransitionKey.OBSERVATION)
if observation is None:
return transition
if self.resize_size is None and not self.crop_params_dict:
return transition
new_observation = dict(observation)
# Process all image keys in the observation
for key in observation:
if "image" not in key:
continue
image = observation[key]
device = image.device
if device.type == "mps":
image = image.cpu()
# Crop if crop params are provided for this key
if key in self.crop_params_dict:
crop_params = self.crop_params_dict[key]
image = F.crop(image, *crop_params)
# Always resize
image = F.resize(image, self.resize_size)
image = image.clamp(0.0, 1.0)
new_observation[key] = image.to(device)
new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = new_observation
return new_transition
def get_config(self) -> dict[str, Any]:
return {
"crop_params_dict": self.crop_params_dict,
"resize_size": self.resize_size,
}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
@dataclass
@ProcessorStepRegistry.register("time_limit_processor")
class TimeLimitProcessor:
"""Track episode time and enforce time limits."""
max_episode_steps: int
current_step: int = 0
def __call__(self, transition: EnvTransition) -> EnvTransition:
truncated = transition.get(TransitionKey.TRUNCATED)
if truncated is None:
return transition
self.current_step += 1
if self.current_step >= self.max_episode_steps:
truncated = True
new_transition = transition.copy()
new_transition[TransitionKey.TRUNCATED] = truncated
return new_transition
def get_config(self) -> dict[str, Any]:
return {
"max_episode_steps": self.max_episode_steps,
}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
self.current_step = 0
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
@dataclass
@ProcessorStepRegistry.register("gripper_penalty_processor")
class GripperPenaltyProcessor:
penalty: float = -0.01
max_gripper_pos: float = 30.0
def __call__(self, transition: EnvTransition) -> EnvTransition:
"""Calculate gripper penalty and add to complementary data."""
action = transition.get(TransitionKey.ACTION)
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
if complementary_data is None or action is None:
return transition
current_gripper_pos = complementary_data.get("raw_joint_positions", None)[-1]
if current_gripper_pos is None:
return transition
gripper_action = action[-1].item()
gripper_action_normalized = gripper_action / self.max_gripper_pos
# Normalize gripper state and action
gripper_state_normalized = current_gripper_pos / self.max_gripper_pos
gripper_action_normalized = gripper_action - 1.0
# Calculate penalty boolean as in original
gripper_penalty_bool = (gripper_state_normalized < 0.5 and gripper_action_normalized > 0.5) or (
gripper_state_normalized > 0.75 and gripper_action_normalized < 0.5
)
gripper_penalty = self.penalty * int(gripper_penalty_bool)
# Add penalty information to complementary data
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
# Create new complementary data with penalty info
new_complementary_data = dict(complementary_data)
new_complementary_data["discrete_penalty"] = gripper_penalty
# Create new transition with updated complementary data
new_transition = transition.copy()
new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
return new_transition
def get_config(self) -> dict[str, Any]:
return {
"penalty": self.penalty,
"max_gripper_pos": self.max_gripper_pos,
}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
"""Reset the processor state."""
self.last_gripper_state = None
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
@dataclass
@ProcessorStepRegistry.register("intervention_action_processor")
class InterventionActionProcessor:
"""Handle action intervention based on signals in the transition.
This processor checks for intervention signals in the transition's complementary data
and overrides agent actions when intervention is active.
"""
use_gripper: bool = False
def __call__(self, transition: EnvTransition) -> EnvTransition:
action = transition.get(TransitionKey.ACTION)
if action is None:
return transition
# Get intervention signals from complementary data
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
teleop_action = complementary_data.get("teleop_action", {})
is_intervention = complementary_data.get("is_intervention", False)
terminate_episode = complementary_data.get("terminate_episode", False)
success = complementary_data.get("success", False)
rerecord_episode = complementary_data.get("rerecord_episode", False)
new_transition = transition.copy()
# Override action if intervention is active
if is_intervention and teleop_action:
# Convert teleop_action dict to tensor format
action_list = [
teleop_action.get("delta_x", 0.0),
teleop_action.get("delta_y", 0.0),
teleop_action.get("delta_z", 0.0),
]
if self.use_gripper:
action_list.append(teleop_action.get("gripper", 1.0))
teleop_action_tensor = torch.tensor(action_list, dtype=action.dtype, device=action.device)
new_transition[TransitionKey.ACTION] = teleop_action_tensor
# Handle episode termination
new_transition[TransitionKey.DONE] = bool(terminate_episode)
new_transition[TransitionKey.REWARD] = float(success)
# Update info with intervention metadata
info = new_transition.get(TransitionKey.INFO, {})
info["is_intervention"] = is_intervention
info["rerecord_episode"] = rerecord_episode
info["next.success"] = success if terminate_episode else info.get("next.success", False)
new_transition[TransitionKey.INFO] = info
new_transition[TransitionKey.COMPLEMENTARY_DATA]["teleop_action"] = new_transition[
TransitionKey.ACTION
]
return new_transition
def get_config(self) -> dict[str, Any]:
return {
"use_gripper": self.use_gripper,
}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
@dataclass
@ProcessorStepRegistry.register("reward_classifier_processor")
class RewardClassifierProcessor:
"""Apply reward classification to image observations.
This processor runs a trained reward classifier on image observations
to predict rewards and success states, potentially terminating episodes
when success is achieved.
"""
pretrained_path: str = None
device: str = "cpu"
success_threshold: float = 0.5
success_reward: float = 1.0
terminate_on_success: bool = True
reward_classifier: Any = None
def __post_init__(self):
"""Initialize the reward classifier after dataclass initialization."""
if self.pretrained_path is not None:
from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
self.reward_classifier = Classifier.from_pretrained(self.pretrained_path)
self.reward_classifier.to(self.device)
self.reward_classifier.eval()
def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition.get(TransitionKey.OBSERVATION)
if observation is None or self.reward_classifier is None:
return transition
# Extract images from observation
images = {key: value for key, value in observation.items() if "image" in key}
if not images:
return transition
# Run reward classifier
start_time = time.perf_counter()
with torch.inference_mode():
success = self.reward_classifier.predict_reward(images, threshold=self.success_threshold)
classifier_frequency = 1 / (time.perf_counter() - start_time)
# Calculate reward and termination
reward = transition.get(TransitionKey.REWARD, 0.0)
terminated = transition.get(TransitionKey.DONE, False)
if success == 1.0:
reward = self.success_reward
if self.terminate_on_success:
terminated = True
# Update transition
new_transition = transition.copy()
new_transition[TransitionKey.REWARD] = reward
new_transition[TransitionKey.DONE] = terminated
# Update info with classifier frequency
info = new_transition.get(TransitionKey.INFO, {})
info["reward_classifier_frequency"] = classifier_frequency
new_transition[TransitionKey.INFO] = info
return new_transition
def get_config(self) -> dict[str, Any]:
return {
"device": self.device,
"success_threshold": self.success_threshold,
"success_reward": self.success_reward,
"terminate_on_success": self.terminate_on_success,
}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
+5 -1
View File
@@ -220,6 +220,7 @@ class UnnormalizerProcessor:
features: dict[str, PolicyFeature] features: dict[str, PolicyFeature]
norm_map: dict[FeatureType, NormalizationMode] norm_map: dict[FeatureType, NormalizationMode]
stats: dict[str, dict[str, Any]] | None = None stats: dict[str, dict[str, Any]] | None = None
eps: float = 1e-8
_tensor_stats: dict[str, dict[str, Tensor]] = field(default_factory=dict, init=False, repr=False) _tensor_stats: dict[str, dict[str, Tensor]] = field(default_factory=dict, init=False, repr=False)
@@ -229,8 +230,10 @@ class UnnormalizerProcessor:
dataset: LeRobotDataset, dataset: LeRobotDataset,
features: dict[str, PolicyFeature], features: dict[str, PolicyFeature],
norm_map: dict[FeatureType, NormalizationMode], norm_map: dict[FeatureType, NormalizationMode],
*,
eps: float = 1e-8,
) -> UnnormalizerProcessor: ) -> UnnormalizerProcessor:
return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats) return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats, eps=eps)
def __post_init__(self): def __post_init__(self):
# Handle deserialization from JSON config # Handle deserialization from JSON config
@@ -305,6 +308,7 @@ class UnnormalizerProcessor:
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
return { return {
"eps": self.eps,
"features": { "features": {
key: {"type": ft.type.value, "shape": ft.shape} for key, ft in self.features.items() key: {"type": ft.type.value, "shape": ft.shape} for key, ft in self.features.items()
}, },
+194 -84
View File
@@ -13,7 +13,8 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 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.
from dataclasses import dataclass from dataclasses import dataclass, field
from typing import Any
import einops import einops
import numpy as np import numpy as np
@@ -22,28 +23,53 @@ from torch import Tensor
from lerobot.configs.types import PolicyFeature from lerobot.configs.types import PolicyFeature
from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
from lerobot.processor.pipeline import ObservationProcessor, ProcessorStepRegistry from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
@dataclass @dataclass
@ProcessorStepRegistry.register(name="observation_processor") class ImageProcessor:
class VanillaObservationProcessor(ObservationProcessor): """Process image observations from environment format to policy format.
"""
Processes environment observations into the LeRobot format by handling both images and states.
Image processing: Converts images from:
- Converts channel-last (H, W, C) images to channel-first (C, H, W) - Channel-last (H, W, C) to channel-first (C, H, W)
- Normalizes uint8 images ([0, 255]) to float32 ([0, 1]) - uint8 [0, 255] to float32 [0, 1]
- Adds a batch dimension if missing - Adds batch dimension if needed
- Supports single images and image dictionaries - Handles both single images and dictionaries of images
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 __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition.get(TransitionKey.OBSERVATION)
if observation is None:
return transition
processed_obs = {}
# Copy all observations first
for key, value in observation.items():
processed_obs[key] = value
# Handle pixels key if present
pixels = observation.get("pixels")
if pixels is not None:
# Remove pixels from processed_obs since we'll replace it with processed images
processed_obs.pop("pixels", None)
# Determine image mapping
if isinstance(pixels, dict):
imgs = {f"{OBS_IMAGES}.{key}": img for key, img in pixels.items()}
else:
imgs = {OBS_IMAGE: pixels}
# Process each image
for imgkey, img in imgs.items():
processed_img = self._process_single_image(img)
processed_obs[imgkey] = processed_img
# Return new transition with processed observation
new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_obs
return new_transition
def _process_single_image(self, img: np.ndarray) -> Tensor: def _process_single_image(self, img: np.ndarray) -> Tensor:
"""Process a single image array.""" """Process a single image array."""
# Convert to tensor # Convert to tensor
@@ -69,89 +95,173 @@ class VanillaObservationProcessor(ObservationProcessor):
return img_tensor return img_tensor
def _process_observation(self, observation): def get_config(self) -> dict[str, Any]:
""" """Return configuration for serialization."""
Processes both image and state observations. return {}
def state_dict(self) -> dict[str, torch.Tensor]:
"""Return state dictionary (empty for this processor)."""
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
"""Load state dictionary (no-op for this processor)."""
pass
def reset(self) -> None:
"""Reset processor state (no-op for this processor)."""
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
"""Transforms:
pixels -> OBS_IMAGE,
observation.pixels -> OBS_IMAGE,
pixels.<cam> -> OBS_IMAGES.<cam>,
observation.pixels.<cam> -> OBS_IMAGES.<cam>
""" """
if "pixels" in features:
features[OBS_IMAGE] = features.pop("pixels")
if "observation.pixels" in features:
features[OBS_IMAGE] = features.pop("observation.pixels")
processed_obs = observation.copy() prefixes = ("pixels.", "observation.pixels.")
for key in list(features.keys()):
for p in prefixes:
if key.startswith(p):
suffix = key[len(p) :]
features[f"{OBS_IMAGES}.{suffix}"] = features.pop(key)
break
return features
if "pixels" in processed_obs:
pixels = processed_obs.pop("pixels")
if isinstance(pixels, dict): @dataclass
imgs = {f"{OBS_IMAGES}.{key}": img for key, img in pixels.items()} class StateProcessor:
else: """Process state observations from environment format to policy format.
imgs = {OBS_IMAGE: pixels}
for imgkey, img in imgs.items(): Handles:
processed_obs[imgkey] = self._process_single_image(img) - environment_state -> observation.environment_state
- agent_pos -> observation.state
- Converts numpy arrays to tensors
- Adds batch dimension if needed
"""
if "environment_state" in processed_obs: def __call__(self, transition: EnvTransition) -> EnvTransition:
env_state_np = processed_obs.pop("environment_state") observation = transition.get(TransitionKey.OBSERVATION)
env_state = torch.from_numpy(env_state_np).float()
if observation is None:
return transition
processed_obs = dict(observation) # Copy existing observations
# Process environment_state
if "environment_state" in observation:
env_state = torch.from_numpy(observation["environment_state"]).float()
if env_state.dim() == 1: if env_state.dim() == 1:
env_state = env_state.unsqueeze(0) env_state = env_state.unsqueeze(0)
processed_obs[OBS_ENV_STATE] = env_state processed_obs[OBS_ENV_STATE] = env_state
# Remove original key
del processed_obs["environment_state"]
if "agent_pos" in processed_obs: # Process agent_pos
agent_pos_np = processed_obs.pop("agent_pos") if "agent_pos" in observation:
agent_pos = torch.from_numpy(agent_pos_np).float() agent_pos = torch.from_numpy(observation["agent_pos"]).float()
if agent_pos.dim() == 1: if agent_pos.dim() == 1:
agent_pos = agent_pos.unsqueeze(0) agent_pos = agent_pos.unsqueeze(0)
processed_obs[OBS_STATE] = agent_pos processed_obs[OBS_STATE] = agent_pos
# Remove original key
del processed_obs["agent_pos"]
return processed_obs # Return new transition with processed observation
new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_obs
return new_transition
def observation(self, observation): def get_config(self) -> dict[str, Any]:
return self._process_observation(observation) """Return configuration for serialization."""
return {}
def state_dict(self) -> dict[str, torch.Tensor]:
"""Return state dictionary (empty for this processor)."""
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
"""Load state dictionary (no-op for this processor)."""
pass
def reset(self) -> None:
"""Reset processor state (no-op for this processor)."""
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]: def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
"""Transforms feature keys to a standardized contract. """Transforms:
environment_state -> OBS_ENV_STATE,
This method handles several renaming patterns: agent_pos -> OBS_STATE,
- Exact matches (e.g., 'pixels' -> 'OBS_IMAGE'). observation.environment_state -> OBS_ENV_STATE,
- Prefixed exact matches (e.g., 'observation.pixels' -> 'OBS_IMAGE'). observation.agent_pos -> OBS_STATE
- 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 = { pairs = (
"pixels": OBS_IMAGE, ("environment_state", OBS_ENV_STATE),
"environment_state": OBS_ENV_STATE, ("agent_pos", OBS_STATE),
"agent_pos": OBS_STATE, )
} for old, new in pairs:
if old in features:
prefix_pairs = { features[new] = features.pop(old)
"pixels.": f"{OBS_IMAGES}.", prefixed = f"observation.{old}"
} if prefixed in features:
features[new] = features.pop(prefixed)
for key in list(features.keys()): return features
matched_prefix = False
for old_prefix, new_prefix in prefix_pairs.items():
prefixed_old = f"observation.{old_prefix}" @dataclass
if key.startswith(prefixed_old): @ProcessorStepRegistry.register(name="observation_processor")
suffix = key[len(prefixed_old) :] class VanillaObservationProcessor:
features[f"{new_prefix}{suffix}"] = features.pop(key) """Complete observation processor that combines image and state processing.
matched_prefix = True
break This processor replicates the functionality of the original preprocess_observation
function but in a modular, composable way that fits into the pipeline architecture.
if key.startswith(old_prefix): """
suffix = key[len(old_prefix) :]
features[f"{new_prefix}{suffix}"] = features.pop(key) image_processor: ImageProcessor = field(default_factory=ImageProcessor)
matched_prefix = True state_processor: StateProcessor = field(default_factory=StateProcessor)
break
def __call__(self, transition: EnvTransition) -> EnvTransition:
if matched_prefix: # First process images
continue transition = self.image_processor(transition)
# Then process state
for old, new in exact_pairs.items(): transition = self.state_processor(transition)
if key == old or key == f"observation.{old}": return transition
if key in features:
features[new] = features.pop(key) def get_config(self) -> dict[str, Any]:
break """Return configuration for serialization."""
return {
"image_processor": self.image_processor.get_config(),
"state_processor": self.state_processor.get_config(),
}
def state_dict(self) -> dict[str, torch.Tensor]:
"""Return state dictionary."""
state = {}
state.update({f"image_processor.{k}": v for k, v in self.image_processor.state_dict().items()})
state.update({f"state_processor.{k}": v for k, v in self.state_processor.state_dict().items()})
return state
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
"""Load state dictionary."""
image_state = {
k.replace("image_processor.", ""): v for k, v in state.items() if k.startswith("image_processor.")
}
state_state = {
k.replace("state_processor.", ""): v for k, v in state.items() if k.startswith("state_processor.")
}
self.image_processor.load_state_dict(image_state)
self.state_processor.load_state_dict(state_state)
def reset(self) -> None:
"""Reset processor state."""
self.image_processor.reset()
self.state_processor.reset()
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
features = self.image_processor.feature_contract(features)
features = self.state_processor.feature_contract(features)
return features return features
+169 -247
View File
@@ -31,12 +31,12 @@ from huggingface_hub.errors import HfHubHTTPError
from safetensors.torch import load_file, save_file from safetensors.torch import load_file, save_file
from lerobot.configs.types import PolicyFeature from lerobot.configs.types import PolicyFeature
from lerobot.utils.utils import get_safe_torch_device
class TransitionKey(str, Enum): class TransitionKey(str, Enum):
"""Keys for accessing EnvTransition dictionary components.""" """Keys for accessing EnvTransition dictionary components."""
# TODO(Steven): Use consts
OBSERVATION = "observation" OBSERVATION = "observation"
ACTION = "action" ACTION = "action"
REWARD = "reward" REWARD = "reward"
@@ -46,18 +46,19 @@ class TransitionKey(str, Enum):
COMPLEMENTARY_DATA = "complementary_data" COMPLEMENTARY_DATA = "complementary_data"
EnvTransition = TypedDict( class EnvTransition(TypedDict, total=False):
"EnvTransition", """Environment transition data structure.
{
TransitionKey.OBSERVATION.value: dict[str, Any] | None, All fields are optional (total=False) to allow flexible usage.
TransitionKey.ACTION.value: Any | torch.Tensor | None, """
TransitionKey.REWARD.value: float | torch.Tensor | None,
TransitionKey.DONE.value: bool | torch.Tensor | None, observation: dict[str, Any] | None
TransitionKey.TRUNCATED.value: bool | torch.Tensor | None, action: Any | torch.Tensor | None
TransitionKey.INFO.value: dict[str, Any] | None, reward: float | torch.Tensor | None
TransitionKey.COMPLEMENTARY_DATA.value: dict[str, Any] | None, done: bool | torch.Tensor | None
}, truncated: bool | torch.Tensor | None
) info: dict[str, Any] | None
complementary_data: dict[str, Any] | None
class ProcessorStepRegistry: class ProcessorStepRegistry:
@@ -135,8 +136,8 @@ class ProcessorStepRegistry:
class ProcessorStep(Protocol): class ProcessorStep(Protocol):
"""Structural typing interface for a single processor step. """Structural typing interface for a single processor step.
A step is any callable accepting a full `EnvTransition` dict and A step is any callable accepting a full `EnvTransition` tuple and
returning a (possibly modified) dict of the same structure. Implementers returning a (possibly modified) tuple of the same structure. Implementers
are encouragedbut not requiredto expose the optional helper methods are encouragedbut not requiredto expose the optional helper methods
listed below. When present, these hooks let `RobotProcessor` listed below. When present, these hooks let `RobotProcessor`
automatically serialise the step's configuration and learnable state using automatically serialise the step's configuration and learnable state using
@@ -254,22 +255,24 @@ class RobotProcessor(ModelHubMixin):
Composable, debuggable post-processing processor for robot transitions. Composable, debuggable post-processing processor for robot transitions.
The class orchestrates an ordered collection of small, functional transformsstepsexecuted The class orchestrates an ordered collection of small, functional transformsstepsexecuted
left-to-right on each incoming `EnvTransition`. It can process both `EnvTransition` dicts left-to-right on each incoming `EnvTransition`. It can process both `EnvTransition` tuples
and batch dictionaries, automatically converting between formats as needed. and batch dictionaries, automatically converting between formats as needed.
Args: Args:
steps: Ordered list of processing steps executed on every call. Defaults to empty list. steps: Ordered list of processing steps executed on every call. Defaults to empty list.
name: Human-readable identifier that is persisted inside the JSON config. name: Human-readable identifier that is persisted inside the JSON config.
Defaults to "RobotProcessor". Defaults to "RobotProcessor".
to_transition: Function to convert batch dict to EnvTransition dict. seed: Global seed forwarded to steps that choose to consume it. Defaults to None.
to_transition: Function to convert batch dict to EnvTransition tuple.
Defaults to _default_batch_to_transition. Defaults to _default_batch_to_transition.
to_output: Function to convert EnvTransition dict to the desired output format. to_output: Function to convert EnvTransition tuple to the desired output format.
Usually it is a batch dict or EnvTransition dict. Usually it is a batch dict or EnvTransition tuple.
Defaults to _default_transition_to_batch. Defaults to _default_transition_to_batch.
before_step_hooks: List of hooks called before each step. Each hook receives the step before_step_hooks: List of hooks called before each step. Each hook receives the step
index and transition, and can optionally return a modified transition. index and transition, and can optionally return a modified transition.
after_step_hooks: List of hooks called after each step. Each hook receives the step after_step_hooks: List of hooks called after each step. Each hook receives the step
index and transition, and can optionally return a modified transition. index and transition, and can optionally return a modified transition.
reset_hooks: List of hooks called during processor reset.
Hook Semantics: Hook Semantics:
- Hooks are executed sequentially in the order they were registered. There is no way to - Hooks are executed sequentially in the order they were registered. There is no way to
@@ -281,13 +284,11 @@ class RobotProcessor(ModelHubMixin):
- Hooks should generally be stateless to maintain predictable behavior. If you need stateful - Hooks should generally be stateless to maintain predictable behavior. If you need stateful
processing, consider implementing a proper ProcessorStep instead. processing, consider implementing a proper ProcessorStep instead.
- To remove hooks, use the unregister methods. To remove steps, you must create a new pipeline. - To remove hooks, use the unregister methods. To remove steps, you must create a new pipeline.
- Hooks ALWAYS receive transitions in EnvTransition format, regardless of the input format
passed to __call__. This ensures consistent hook behavior whether processing batch dicts
or EnvTransition objects.
""" """
steps: Sequence[ProcessorStep] = field(default_factory=list) steps: Sequence[ProcessorStep] = field(default_factory=list)
name: str = "RobotProcessor" name: str = "RobotProcessor"
seed: int | None = None
to_transition: Callable[[dict[str, Any]], EnvTransition] = field( to_transition: Callable[[dict[str, Any]], EnvTransition] = field(
default_factory=lambda: _default_batch_to_transition, repr=False default_factory=lambda: _default_batch_to_transition, repr=False
@@ -300,6 +301,7 @@ class RobotProcessor(ModelHubMixin):
# Hooks do not modify transitions - they are called for logging, debugging, or monitoring purposes # Hooks do not modify transitions - they are called for logging, debugging, or monitoring purposes
before_step_hooks: list[Callable[[int, EnvTransition], None]] = field(default_factory=list, repr=False) before_step_hooks: list[Callable[[int, EnvTransition], None]] = field(default_factory=list, repr=False)
after_step_hooks: list[Callable[[int, EnvTransition], None]] = field(default_factory=list, repr=False) after_step_hooks: list[Callable[[int, EnvTransition], None]] = field(default_factory=list, repr=False)
reset_hooks: list[Callable[[], None]] = field(default_factory=list, repr=False)
def __call__(self, data: EnvTransition | dict[str, Any]): def __call__(self, data: EnvTransition | dict[str, Any]):
"""Process data through all steps. """Process data through all steps.
@@ -319,30 +321,22 @@ class RobotProcessor(ModelHubMixin):
Raises: Raises:
ValueError: If the transition is not a valid EnvTransition format. ValueError: If the transition is not a valid EnvTransition format.
""" """
# Check if we need to convert back to batch format at the end iterator = self.step_through(data)
_, called_with_batch = self._prepare_transition(data) current_result = next(iterator) # Get initial state
# Use step_through to get the iterator # Process through all steps with hooks
step_iterator = self.step_through(data) for idx, step_result in enumerate(iterator):
# Apply before hooks
# Get initial state (before any steps)
current_transition = next(step_iterator)
# Process each step with hooks
for idx, next_transition in enumerate(step_iterator):
# Apply before hooks with current state (before step execution)
for hook in self.before_step_hooks: for hook in self.before_step_hooks:
hook(idx, current_transition) _ = hook(idx, step_result)
# Move to next state (after step execution) # Apply after hooks
current_transition = next_transition
# Apply after hooks with updated state
for hook in self.after_step_hooks: for hook in self.after_step_hooks:
hook(idx, current_transition) _ = hook(idx, step_result)
# Convert back to original format if needed current_result = step_result
return self.to_output(current_transition) if called_with_batch else current_transition
return current_result
def _prepare_transition(self, data: EnvTransition | dict[str, Any]) -> tuple[EnvTransition, bool]: def _prepare_transition(self, data: EnvTransition | dict[str, Any]) -> tuple[EnvTransition, bool]:
"""Prepare and validate transition data for processing. """Prepare and validate transition data for processing.
@@ -372,48 +366,64 @@ class RobotProcessor(ModelHubMixin):
return transition, called_with_batch return transition, called_with_batch
def step_through(self, data: EnvTransition | dict[str, Any]) -> Iterable[EnvTransition]: def step_through(self, data: EnvTransition | dict[str, Any]) -> Iterable[EnvTransition | dict[str, Any]]:
"""Yield the intermediate results after each processor step. """Yield the intermediate results after each processor step.
This is a low-level method that does NOT apply hooks. It simply executes each step This is a low-level method that does NOT apply hooks. It simply executes each step
and yields the intermediate results. This allows users to debug the pipeline or and yields the intermediate results. This allows users to debug the pipeline or
apply custom logic between steps if needed. apply custom logic between steps if needed.
Note: This method always yields EnvTransition objects regardless of input format. Like __call__, this method accepts either EnvTransition dicts or batch dictionaries
If you need the results in the original input format, you'll need to convert them and preserves the input format in the yielded results.
using `to_output()`.
Args: Args:
data: Either an EnvTransition dict or a batch dictionary to process. data: Either an EnvTransition dict or a batch dictionary to process.
Yields: Yields:
The intermediate EnvTransition results after each step. The intermediate results after each step, in the same format as the input.
""" """
transition, _ = self._prepare_transition(data) transition, called_with_batch = self._prepare_transition(data)
# Yield initial state # Yield initial state
yield transition yield self.to_output(transition) if called_with_batch else transition
# Process each step WITHOUT hooks (low-level method) # Process each step WITHOUT hooks (low-level method)
for processor_step in self.steps: for processor_step in self.steps:
transition = processor_step(transition) transition = processor_step(transition)
yield transition yield self.to_output(transition) if called_with_batch else transition
def _save_pretrained(self, save_directory: Path, **kwargs): def _save_pretrained(self, destination_path: str, **kwargs):
"""Internal save method for ModelHubMixin compatibility.""" """Internal save method for ModelHubMixin compatibility."""
# Extract config_filename from kwargs if provided # Extract config_filename from kwargs if provided
config_filename = kwargs.pop("config_filename", None) config_filename = kwargs.pop("config_filename", None)
self.save_pretrained(save_directory, config_filename=config_filename) self.save_pretrained(destination_path, config_filename=config_filename)
def save_pretrained(self, save_directory: str | Path, config_filename: str | None = None, **kwargs): def _generate_model_card(self, destination_path: str) -> None:
"""Serialize the processor definition and parameters to *save_directory*. """Generate README.md from the RobotProcessor model card template."""
# Read the template
template_path = Path(__file__).parent.parent / "templates" / "robotprocessor_modelcard_template.md"
if not template_path.exists():
# Fallback: if template doesn't exist, skip model card generation
return
with open(template_path) as f:
model_card_content = f.read()
# Write the README.md
readme_path = os.path.join(destination_path, "README.md")
with open(readme_path, "w") as f:
f.write(model_card_content)
def save_pretrained(self, destination_path: str, config_filename: str | None = None, **kwargs):
"""Serialize the processor definition and parameters to *destination_path*.
Args: Args:
save_directory: Directory where the processor will be saved. destination_path: Directory where the processor will be saved.
config_filename: Optional custom config filename. If not provided, defaults to config_filename: Optional custom config filename. If not provided, defaults to
"{self.name}.json" where self.name is sanitized for filesystem compatibility. "{self.name}.json" where self.name is sanitized for filesystem compatibility.
""" """
os.makedirs(str(save_directory), exist_ok=True) os.makedirs(destination_path, exist_ok=True)
# Sanitize processor name for use in filenames # Sanitize processor name for use in filenames
import re import re
@@ -427,6 +437,7 @@ class RobotProcessor(ModelHubMixin):
config: dict[str, Any] = { config: dict[str, Any] = {
"name": self.name, "name": self.name,
"seed": self.seed,
"steps": [], "steps": [],
} }
@@ -434,15 +445,16 @@ class RobotProcessor(ModelHubMixin):
# Check if step was registered # Check if step was registered
registry_name = getattr(processor_step.__class__, "_registry_name", None) registry_name = getattr(processor_step.__class__, "_registry_name", None)
step_entry: dict[str, Any] = {}
if registry_name: if registry_name:
# Use registry name for registered steps # Use registry name for registered steps
step_entry["registry_name"] = registry_name step_entry: dict[str, Any] = {
"registry_name": registry_name,
}
else: else:
# Fall back to full module path for unregistered steps # Fall back to full module path for unregistered steps
step_entry["class"] = ( step_entry: dict[str, Any] = {
f"{processor_step.__class__.__module__}.{processor_step.__class__.__name__}" "class": f"{processor_step.__class__.__module__}.{processor_step.__class__.__name__}",
) }
if hasattr(processor_step, "get_config"): if hasattr(processor_step, "get_config"):
step_entry["config"] = processor_step.get_config() step_entry["config"] = processor_step.get_config()
@@ -469,34 +481,43 @@ class RobotProcessor(ModelHubMixin):
else: else:
state_filename = f"{sanitized_name}_step_{step_index}.safetensors" state_filename = f"{sanitized_name}_step_{step_index}.safetensors"
save_file(cloned_state, os.path.join(str(save_directory), state_filename)) save_file(cloned_state, os.path.join(destination_path, state_filename))
step_entry["state_file"] = state_filename step_entry["state_file"] = state_filename
config["steps"].append(step_entry) config["steps"].append(step_entry)
with open(os.path.join(str(save_directory), config_filename), "w") as file_pointer: with open(os.path.join(destination_path, config_filename), "w") as file_pointer:
json.dump(config, file_pointer, indent=2) json.dump(config, file_pointer, indent=2)
# Generate README.md from template
self._generate_model_card(destination_path)
def to(self, device: str | torch.device):
"""Move all tensor states inside each step to device and return self.
Uses a generic mechanism: fetch each step's state dict, move every tensor
to the target device, and reload it. Only works for steps that implement
both state_dict() and load_state_dict() methods.
"""
device = get_safe_torch_device(device)
for step in self.steps:
if hasattr(step, "state_dict") and hasattr(step, "load_state_dict"):
state = step.state_dict()
if state: # Only process if there's actual state
moved_state = {k: v.to(device) for k, v in state.items()}
step.load_state_dict(moved_state)
return self
@classmethod @classmethod
def from_pretrained( def from_pretrained(
cls, cls, source: str, *, config_filename: str | None = None, overrides: dict[str, Any] | None = None
pretrained_model_name_or_path: str | Path,
*,
force_download: bool = False,
resume_download: bool | None = None,
proxies: dict[str, str] | None = None,
token: str | bool | None = None,
cache_dir: str | Path | None = None,
local_files_only: bool = False,
revision: str | None = None,
config_filename: str | None = None,
overrides: dict[str, Any] | None = None,
**kwargs,
) -> RobotProcessor: ) -> RobotProcessor:
"""Load a serialized processor from source (local path or Hugging Face Hub identifier). """Load a serialized processor from source (local path or Hugging Face Hub identifier).
Args: Args:
pretrained_model_name_or_path: Local path to a saved processor directory or Hugging Face Hub identifier source: Local path to a saved processor directory or Hugging Face Hub identifier
(e.g., "username/processor-name"). (e.g., "username/processor-name").
config_filename: Optional specific config filename to load. If not provided, will: config_filename: Optional specific config filename to load. If not provided, will:
- For local paths: look for any .json file in the directory (error if multiple found) - For local paths: look for any .json file in the directory (error if multiple found)
@@ -549,9 +570,6 @@ class RobotProcessor(ModelHubMixin):
) )
``` ```
""" """
# Use the local variable name 'source' for clarity
source = str(pretrained_model_name_or_path)
if Path(source).is_dir(): if Path(source).is_dir():
# Local path - use it directly # Local path - use it directly
base_path = Path(source) base_path = Path(source)
@@ -569,7 +587,7 @@ class RobotProcessor(ModelHubMixin):
config_filename = json_files[0].name config_filename = json_files[0].name
with open(base_path / config_filename) as file_pointer: with open(base_path / config_filename) as file_pointer:
loaded_config: dict[str, Any] = json.load(file_pointer) config: dict[str, Any] = json.load(file_pointer)
else: else:
# Hugging Face Hub - download all required files # Hugging Face Hub - download all required files
if config_filename is None: if config_filename is None:
@@ -583,18 +601,7 @@ class RobotProcessor(ModelHubMixin):
config_path = None config_path = None
for name in common_names: for name in common_names:
try: try:
config_path = hf_hub_download( config_path = hf_hub_download(source, name, repo_type="model")
source,
name,
repo_type="model",
force_download=force_download,
resume_download=resume_download,
proxies=proxies,
token=token,
cache_dir=cache_dir,
local_files_only=local_files_only,
revision=revision,
)
config_filename = name config_filename = name
break break
except (FileNotFoundError, OSError, HfHubHTTPError): except (FileNotFoundError, OSError, HfHubHTTPError):
@@ -610,21 +617,10 @@ class RobotProcessor(ModelHubMixin):
) )
else: else:
# Download specific config file # Download specific config file
config_path = hf_hub_download( config_path = hf_hub_download(source, config_filename, repo_type="model")
source,
config_filename,
repo_type="model",
force_download=force_download,
resume_download=resume_download,
proxies=proxies,
token=token,
cache_dir=cache_dir,
local_files_only=local_files_only,
revision=revision,
)
with open(config_path) as file_pointer: with open(config_path) as file_pointer:
loaded_config = json.load(file_pointer) config: dict[str, Any] = json.load(file_pointer)
# Store downloaded files in the same directory as the config # Store downloaded files in the same directory as the config
base_path = Path(config_path).parent base_path = Path(config_path).parent
@@ -637,7 +633,7 @@ class RobotProcessor(ModelHubMixin):
override_keys = set(overrides.keys()) override_keys = set(overrides.keys())
steps: list[ProcessorStep] = [] steps: list[ProcessorStep] = []
for step_entry in loaded_config["steps"]: for step_entry in config["steps"]:
# Check if step uses registry name or module path # Check if step uses registry name or module path
if "registry_name" in step_entry: if "registry_name" in step_entry:
# Load from registry # Load from registry
@@ -689,18 +685,7 @@ class RobotProcessor(ModelHubMixin):
state_path = str(base_path / step_entry["state_file"]) state_path = str(base_path / step_entry["state_file"])
else: else:
# Hugging Face Hub - download the state file # Hugging Face Hub - download the state file
state_path = hf_hub_download( state_path = hf_hub_download(source, step_entry["state_file"], repo_type="model")
source,
step_entry["state_file"],
repo_type="model",
force_download=force_download,
resume_download=resume_download,
proxies=proxies,
token=token,
cache_dir=cache_dir,
local_files_only=local_files_only,
revision=revision,
)
step_instance.load_state_dict(load_file(state_path)) step_instance.load_state_dict(load_file(state_path))
@@ -709,7 +694,7 @@ class RobotProcessor(ModelHubMixin):
# Check for unused override keys # Check for unused override keys
if override_keys: if override_keys:
available_keys = [] available_keys = []
for step_entry in loaded_config["steps"]: for step_entry in config["steps"]:
if "registry_name" in step_entry: if "registry_name" in step_entry:
available_keys.append(step_entry["registry_name"]) available_keys.append(step_entry["registry_name"])
else: else:
@@ -723,7 +708,7 @@ class RobotProcessor(ModelHubMixin):
f"Make sure override keys match exact step class names or registry names." f"Make sure override keys match exact step class names or registry names."
) )
return cls(steps, loaded_config.get("name", "RobotProcessor")) return cls(steps, config.get("name", "RobotProcessor"), config.get("seed"))
def __len__(self) -> int: def __len__(self) -> int:
"""Return the number of steps in the processor.""" """Return the number of steps in the processor."""
@@ -735,7 +720,7 @@ class RobotProcessor(ModelHubMixin):
* ``slice`` returns a new RobotProcessor with the sliced steps. * ``slice`` returns a new RobotProcessor with the sliced steps.
""" """
if isinstance(idx, slice): if isinstance(idx, slice):
return RobotProcessor(self.steps[idx], self.name) return RobotProcessor(self.steps[idx], self.name, self.seed)
return self.steps[idx] return self.steps[idx]
def register_before_step_hook(self, fn: Callable[[int, EnvTransition], None]): def register_before_step_hook(self, fn: Callable[[int, EnvTransition], None]):
@@ -778,11 +763,71 @@ class RobotProcessor(ModelHubMixin):
f"Hook {fn} not found in after_step_hooks. Make sure to pass the exact same function reference." f"Hook {fn} not found in after_step_hooks. Make sure to pass the exact same function reference."
) from None ) from None
def register_reset_hook(self, fn: Callable[[], None]):
"""Attach fn to be executed when reset is called."""
self.reset_hooks.append(fn)
def unregister_reset_hook(self, fn: Callable[[], None]):
"""Remove a previously registered reset hook.
Args:
fn: The exact function reference that was registered. Must be the same object.
Raises:
ValueError: If the hook is not found in the registered hooks.
"""
try:
self.reset_hooks.remove(fn)
except ValueError:
raise ValueError(
f"Hook {fn} not found in reset_hooks. Make sure to pass the exact same function reference."
) from None
def reset(self): def reset(self):
"""Clear state in every step that implements ``reset()`` and fire registered hooks.""" """Clear state in every step that implements ``reset()`` and fire registered hooks."""
for step in self.steps: for step in self.steps:
if hasattr(step, "reset"): if hasattr(step, "reset"):
step.reset() # type: ignore[attr-defined] step.reset() # type: ignore[attr-defined]
for fn in self.reset_hooks:
fn()
def profile_steps(
self, transition: EnvTransition, num_runs: int = 100, warmup_runs: int = 5
) -> dict[str, float]:
"""Profile the execution time of each step for performance optimization."""
import copy
import time
profile_results = {}
# Make a copy to avoid altering the original transition
transition_copy = copy.deepcopy(transition)
# Get intermediate transitions for each step using step_through
intermediate_transitions = list(self.step_through(transition_copy))
for idx, processor_step in enumerate(self.steps):
step_name = f"step_{idx}_{processor_step.__class__.__name__}"
# Use the appropriate input transition for this step
input_transition = intermediate_transitions[idx]
# Warm up - copy transition for each run to ensure consistent conditions
for _ in range(warmup_runs):
transition_copy = copy.deepcopy(input_transition)
_ = processor_step(transition_copy)
# Time the step - copy transition for each run to ensure consistent conditions
start_time = time.perf_counter()
for _ in range(num_runs):
transition_copy = copy.deepcopy(input_transition)
_ = processor_step(transition_copy)
end_time = time.perf_counter()
avg_time = (end_time - start_time) / num_runs * 1000 # Convert to milliseconds
profile_results[step_name] = avg_time
return profile_results
def __repr__(self) -> str: def __repr__(self) -> str:
"""Return a readable string representation of the processor.""" """Return a readable string representation of the processor."""
@@ -799,6 +844,9 @@ class RobotProcessor(ModelHubMixin):
parts = [f"name='{self.name}'", steps_repr] parts = [f"name='{self.name}'", steps_repr]
if self.seed is not None:
parts.append(f"seed={self.seed}")
return f"RobotProcessor({', '.join(parts)})" return f"RobotProcessor({', '.join(parts)})"
def __post_init__(self): def __post_init__(self):
@@ -863,30 +911,12 @@ class ObservationProcessor:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition.get(TransitionKey.OBSERVATION) observation = transition.get(TransitionKey.OBSERVATION)
if observation is None:
return transition
processed_observation = self.observation(observation) processed_observation = self.observation(observation)
# Create a new transition dict with the processed observation # Create a new transition dict with the processed observation
new_transition = transition.copy() new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_observation new_transition[TransitionKey.OBSERVATION] = processed_observation
return new_transition return new_transition
def get_config(self) -> dict[str, Any]:
return {}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
class ActionProcessor: class ActionProcessor:
"""Base class for processors that modify only the action component of a transition. """Base class for processors that modify only the action component of a transition.
@@ -923,30 +953,12 @@ class ActionProcessor:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
action = transition.get(TransitionKey.ACTION) action = transition.get(TransitionKey.ACTION)
if action is None:
return transition
processed_action = self.action(action) processed_action = self.action(action)
# Create a new transition dict with the processed action # Create a new transition dict with the processed action
new_transition = transition.copy() new_transition = transition.copy()
new_transition[TransitionKey.ACTION] = processed_action new_transition[TransitionKey.ACTION] = processed_action
return new_transition return new_transition
def get_config(self) -> dict[str, Any]:
return {}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
class RewardProcessor: class RewardProcessor:
"""Base class for processors that modify only the reward component of a transition. """Base class for processors that modify only the reward component of a transition.
@@ -982,30 +994,12 @@ class RewardProcessor:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
reward = transition.get(TransitionKey.REWARD) reward = transition.get(TransitionKey.REWARD)
if reward is None:
return transition
processed_reward = self.reward(reward) processed_reward = self.reward(reward)
# Create a new transition dict with the processed reward # Create a new transition dict with the processed reward
new_transition = transition.copy() new_transition = transition.copy()
new_transition[TransitionKey.REWARD] = processed_reward new_transition[TransitionKey.REWARD] = processed_reward
return new_transition return new_transition
def get_config(self) -> dict[str, Any]:
return {}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
class DoneProcessor: class DoneProcessor:
"""Base class for processors that modify only the done flag of a transition. """Base class for processors that modify only the done flag of a transition.
@@ -1046,30 +1040,12 @@ class DoneProcessor:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
done = transition.get(TransitionKey.DONE) done = transition.get(TransitionKey.DONE)
if done is None:
return transition
processed_done = self.done(done) processed_done = self.done(done)
# Create a new transition dict with the processed done flag # Create a new transition dict with the processed done flag
new_transition = transition.copy() new_transition = transition.copy()
new_transition[TransitionKey.DONE] = processed_done new_transition[TransitionKey.DONE] = processed_done
return new_transition return new_transition
def get_config(self) -> dict[str, Any]:
return {}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
class TruncatedProcessor: class TruncatedProcessor:
"""Base class for processors that modify only the truncated flag of a transition. """Base class for processors that modify only the truncated flag of a transition.
@@ -1106,30 +1082,12 @@ class TruncatedProcessor:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
truncated = transition.get(TransitionKey.TRUNCATED) truncated = transition.get(TransitionKey.TRUNCATED)
if truncated is None:
return transition
processed_truncated = self.truncated(truncated) processed_truncated = self.truncated(truncated)
# Create a new transition dict with the processed truncated flag # Create a new transition dict with the processed truncated flag
new_transition = transition.copy() new_transition = transition.copy()
new_transition[TransitionKey.TRUNCATED] = processed_truncated new_transition[TransitionKey.TRUNCATED] = processed_truncated
return new_transition return new_transition
def get_config(self) -> dict[str, Any]:
return {}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
class InfoProcessor: class InfoProcessor:
"""Base class for processors that modify only the info dictionary of a transition. """Base class for processors that modify only the info dictionary of a transition.
@@ -1171,30 +1129,12 @@ class InfoProcessor:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
info = transition.get(TransitionKey.INFO) info = transition.get(TransitionKey.INFO)
if info is None:
return transition
processed_info = self.info(info) processed_info = self.info(info)
# Create a new transition dict with the processed info # Create a new transition dict with the processed info
new_transition = transition.copy() new_transition = transition.copy()
new_transition[TransitionKey.INFO] = processed_info new_transition[TransitionKey.INFO] = processed_info
return new_transition return new_transition
def get_config(self) -> dict[str, Any]:
return {}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
class ComplementaryDataProcessor: class ComplementaryDataProcessor:
"""Base class for processors that modify only the complementary data of a transition. """Base class for processors that modify only the complementary data of a transition.
@@ -1217,30 +1157,12 @@ class ComplementaryDataProcessor:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA) complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
if complementary_data is None:
return transition
processed_complementary_data = self.complementary_data(complementary_data) processed_complementary_data = self.complementary_data(complementary_data)
# Create a new transition dict with the processed complementary data # Create a new transition dict with the processed complementary data
new_transition = transition.copy() new_transition = transition.copy()
new_transition[TransitionKey.COMPLEMENTARY_DATA] = processed_complementary_data new_transition[TransitionKey.COMPLEMENTARY_DATA] = processed_complementary_data
return new_transition return new_transition
def get_config(self) -> dict[str, Any]:
return {}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
class IdentityProcessor: class IdentityProcessor:
"""Identity processor that does nothing.""" """Identity processor that does nothing."""
+19 -7
View File
@@ -16,21 +16,24 @@
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any from typing import Any
import torch
from lerobot.configs.types import PolicyFeature from lerobot.configs.types import PolicyFeature
from lerobot.processor.pipeline import ( from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
ObservationProcessor,
ProcessorStepRegistry,
)
@dataclass @dataclass
@ProcessorStepRegistry.register(name="rename_processor") @ProcessorStepRegistry.register(name="rename_processor")
class RenameProcessor(ObservationProcessor): class RenameProcessor:
"""Rename processor that renames keys in the observation.""" """Rename processor that renames keys in the observation."""
rename_map: dict[str, str] = field(default_factory=dict) rename_map: dict[str, str] = field(default_factory=dict)
def observation(self, observation): def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition.get(TransitionKey.OBSERVATION)
if observation is None:
return transition
processed_obs = {} processed_obs = {}
for key, value in observation.items(): for key, value in observation.items():
if key in self.rename_map: if key in self.rename_map:
@@ -38,11 +41,20 @@ class RenameProcessor(ObservationProcessor):
else: else:
processed_obs[key] = value processed_obs[key] = value
return processed_obs # Create a new transition with the renamed observation
new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_obs
return new_transition
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
return {"rename_map": self.rename_map} return {"rename_map": self.rename_map}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]: def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
"""Transforms: """Transforms:
- Each key in the observation that appears in `rename_map` is renamed to its value. - Each key in the observation that appears in `rename_map` is renamed to its value.
+245
View File
@@ -0,0 +1,245 @@
from dataclasses import dataclass, field
from typing import Any
import gymnasium as gym
import numpy as np
import torch
from lerobot.configs.types import PolicyFeature
from lerobot.model.kinematics import RobotKinematics
from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
@dataclass
@ProcessorStepRegistry.register("joint_velocity_processor")
class JointVelocityProcessor:
"""Add joint velocity information to observations.
Computes joint velocities by tracking changes in joint positions over time.
"""
joint_velocity_limits: float = 100.0
dt: float = 1.0 / 10
last_joint_positions: torch.Tensor | None = None
def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition.get(TransitionKey.OBSERVATION)
if observation is None:
return transition
# Get current joint positions (assuming they're in observation.state)
current_positions = observation.get("observation.state")
if current_positions is None:
return transition
# Initialize last joint positions if not already set
if self.last_joint_positions is None:
self.last_joint_positions = current_positions.clone()
# Compute velocities
joint_velocities = (current_positions - self.last_joint_positions) / self.dt
self.last_joint_positions = current_positions.clone()
# Extend observation with velocities
extended_state = torch.cat([current_positions, joint_velocities], dim=-1)
# Create new observation dict
new_observation = dict(observation)
new_observation["observation.state"] = extended_state
# Return new transition
new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = new_observation
return new_transition
def get_config(self) -> dict[str, Any]:
return {
"joint_velocity_limits": self.joint_velocity_limits,
"dt": self.dt,
}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
self.last_joint_positions = None
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
@dataclass
@ProcessorStepRegistry.register("current_processor")
class MotorCurrentProcessor:
"""Add motor current information to observations."""
env: gym.Env = None
def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition.get(TransitionKey.OBSERVATION)
if observation is None:
return transition
# Get current values from complementary_data (where robot state would be stored)
present_current_dict = self.env.unwrapped.robot.bus.sync_read("Present_Current")
motor_currents = torch.tensor(
[present_current_dict[name] for name in self.env.unwrapped.robot.bus.motors],
dtype=torch.float32,
).unsqueeze(0)
current_state = observation.get("observation.state")
if current_state is None:
return transition
extended_state = torch.cat([current_state, motor_currents], dim=-1)
# Create new observation dict
new_observation = dict(observation)
new_observation["observation.state"] = extended_state
# Return new transition
new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = new_observation
return new_transition
def get_config(self) -> dict[str, Any]:
return {}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
pass
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
@dataclass
@ProcessorStepRegistry.register("inverse_kinematics_processor")
class InverseKinematicsProcessor:
"""Convert end-effector space actions to joint space using inverse kinematics.
This processor transforms delta commands in end-effector space (delta_x, delta_y, delta_z)
to joint space commands using forward and inverse kinematics. It maintains the current
end-effector pose and joint positions to compute the transformations.
"""
urdf_path: str
target_frame_name: str = "gripper_link"
end_effector_step_sizes: dict[str, float] = field(default_factory=lambda: {"x": 1.0, "y": 1.0, "z": 1.0})
end_effector_bounds: dict[str, list[float]] | None = None
max_gripper_pos: float = 30.0
# State tracking
current_ee_pos: np.ndarray | None = field(default=None, init=False, repr=False)
current_joint_pos: np.ndarray | None = field(default=None, init=False, repr=False)
kinematics: RobotKinematics | None = field(default=None, init=False, repr=False)
def __post_init__(self):
"""Initialize the kinematics module after dataclass initialization."""
if self.urdf_path:
self.kinematics = RobotKinematics(
urdf_path=self.urdf_path,
target_frame_name=self.target_frame_name,
)
def __call__(self, transition: EnvTransition) -> EnvTransition:
action = transition.get(TransitionKey.ACTION)
if action is None:
return transition
action_np = action.detach().cpu().numpy().squeeze()
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
raw_joint_positions = complementary_data.get("raw_joint_positions")
current_gripper_pos = raw_joint_positions[-1]
if self.current_joint_pos is None:
self.current_joint_pos = raw_joint_positions
# Initialize end-effector position if not available
if self.current_joint_pos is None:
return transition # Cannot proceed without joint positions
# Calculate current end-effector position using forward kinematics
if self.current_ee_pos is None:
self.current_ee_pos = self.kinematics.forward_kinematics(self.current_joint_pos)
# Scale deltas by step sizes
delta_ee = np.array(
[
action_np[0] * self.end_effector_step_sizes["x"],
action_np[1] * self.end_effector_step_sizes["y"],
action_np[2] * self.end_effector_step_sizes["z"],
],
dtype=np.float32,
)
# Set desired end-effector position by adding delta
desired_ee_pos = np.eye(4)
desired_ee_pos[:3, :3] = self.current_ee_pos[:3, :3] # Keep orientation
# Add delta to position and clip to bounds
desired_ee_pos[:3, 3] = self.current_ee_pos[:3, 3] + delta_ee
if self.end_effector_bounds is not None:
desired_ee_pos[:3, 3] = np.clip(
desired_ee_pos[:3, 3],
self.end_effector_bounds["min"],
self.end_effector_bounds["max"],
)
# Compute inverse kinematics to get joint positions
target_joint_values = self.kinematics.inverse_kinematics(self.current_joint_pos, desired_ee_pos)
# Update current state
self.current_ee_pos = desired_ee_pos.copy()
self.current_joint_pos = target_joint_values.copy()
# Create new action with joint space commands
gripper_action = current_gripper_pos
if len(action_np) > 3:
# Handle gripper command separately
gripper_command = action_np[3]
# Process gripper command (convert from [0,2] to delta) and discretize
gripper_delta = np.round(gripper_command - 1.0).astype(int) * self.max_gripper_pos
gripper_action = np.clip(current_gripper_pos + gripper_delta, 0, self.max_gripper_pos)
# Combine joint positions and gripper
target_joint_values[-1] = gripper_action
converted_action = torch.from_numpy(target_joint_values).to(action.device).to(action.dtype)
new_transition = transition.copy()
new_transition[TransitionKey.ACTION] = converted_action
return new_transition
def get_config(self) -> dict[str, Any]:
return {
"urdf_path": self.urdf_path,
"target_frame_name": self.target_frame_name,
"end_effector_step_sizes": self.end_effector_step_sizes,
"end_effector_bounds": self.end_effector_bounds,
"max_gripper_pos": self.max_gripper_pos,
}
def state_dict(self) -> dict[str, torch.Tensor]:
return {}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
def reset(self) -> None:
"""Reset the processor state."""
self.current_ee_pos = None
self.current_joint_pos = None
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
return features
+1 -5
View File
@@ -393,9 +393,5 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
return dataset return dataset
def main():
record()
if __name__ == "__main__": if __name__ == "__main__":
main() record()
+1 -5
View File
@@ -112,9 +112,5 @@ def replay(cfg: ReplayConfig):
robot.disconnect() robot.disconnect()
def main():
replay()
if __name__ == "__main__": if __name__ == "__main__":
main() replay()
+20 -8
View File
@@ -68,10 +68,11 @@ from tqdm import trange
from lerobot.configs import parser from lerobot.configs import parser
from lerobot.configs.eval import EvalPipelineConfig from lerobot.configs.eval import EvalPipelineConfig
from lerobot.envs.factory import make_env from lerobot.envs.factory import make_env
from lerobot.envs.utils import add_envs_task, check_env_attributes_and_types, preprocess_observation from lerobot.envs.utils import add_envs_task, check_env_attributes_and_types
from lerobot.policies.factory import make_policy from lerobot.policies.factory import make_policy
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.utils import get_device_from_parameters from lerobot.policies.utils import get_device_from_parameters
from lerobot.processor import RobotProcessor, TransitionKey, VanillaObservationProcessor
from lerobot.utils.io_utils import write_video from lerobot.utils.io_utils import write_video
from lerobot.utils.random_utils import set_seed from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import ( from lerobot.utils.utils import (
@@ -128,6 +129,16 @@ def rollout(
if render_callback is not None: if render_callback is not None:
render_callback(env) render_callback(env)
# Create observation processing processor
# NOTE: During environment interaction, we skip batch dictionary conversion
# since that format is only needed for loss computation during training.
# Using identity functions to avoid unnecessary format transformations.
obs_processor = RobotProcessor(
[VanillaObservationProcessor()],
to_transition=lambda x: x,
to_output=lambda x: x,
)
all_observations = [] all_observations = []
all_actions = [] all_actions = []
all_rewards = [] all_rewards = []
@@ -147,10 +158,13 @@ def rollout(
check_env_attributes_and_types(env) check_env_attributes_and_types(env)
while not np.all(done): while not np.all(done):
# Numpy array to tensor and changing dictionary keys to LeRobot policy format. # Numpy array to tensor and changing dictionary keys to LeRobot policy format.
observation = preprocess_observation(observation) transition = (observation, None, None, None, None, None, None)
processed_transition = obs_processor(transition)
observation = processed_transition[TransitionKey.OBSERVATION]
if return_observations: if return_observations:
all_observations.append(deepcopy(observation)) all_observations.append(deepcopy(observation))
# TODO(azouitine): Move this in processor side
observation = { observation = {
key: observation[key].to(device, non_blocking=device.type == "cuda") for key in observation key: observation[key].to(device, non_blocking=device.type == "cuda") for key in observation
} }
@@ -195,7 +209,9 @@ def rollout(
# Track the final observation. # Track the final observation.
if return_observations: if return_observations:
observation = preprocess_observation(observation) transition = (observation, None, None, None, None, None, None)
processed_transition = obs_processor(transition)
observation = processed_transition[TransitionKey.OBSERVATION]
all_observations.append(deepcopy(observation)) all_observations.append(deepcopy(observation))
# Stack the sequence along the first dimension so that we have (batch, sequence, *) tensors. # Stack the sequence along the first dimension so that we have (batch, sequence, *) tensors.
@@ -501,10 +517,6 @@ def eval_main(cfg: EvalPipelineConfig):
logging.info("End of eval") logging.info("End of eval")
def main(): if __name__ == "__main__":
init_logging() init_logging()
eval_main() eval_main()
if __name__ == "__main__":
main()
+61 -26
View File
@@ -62,8 +62,14 @@ from lerobot.configs import parser
from lerobot.configs.train import TrainRLServerPipelineConfig from lerobot.configs.train import TrainRLServerPipelineConfig
from lerobot.policies.factory import make_policy from lerobot.policies.factory import make_policy
from lerobot.policies.sac.modeling_sac import SACPolicy from lerobot.policies.sac.modeling_sac import SACPolicy
from lerobot.processor.pipeline import TransitionKey
from lerobot.robots import so100_follower # noqa: F401 from lerobot.robots import so100_follower # noqa: F401
from lerobot.scripts.rl.gym_manipulator import make_robot_env from lerobot.scripts.rl.gym_manipulator import (
create_transition,
make_processors,
make_robot_env,
step_env_and_process_transition,
)
from lerobot.teleoperators import gamepad, so101_leader # noqa: F401 from lerobot.teleoperators import gamepad, so101_leader # noqa: F401
from lerobot.transport import services_pb2, services_pb2_grpc from lerobot.transport import services_pb2, services_pb2_grpc
from lerobot.transport.utils import ( from lerobot.transport.utils import (
@@ -236,7 +242,8 @@ def act_with_policy(
logging.info("make_env online") logging.info("make_env online")
online_env = make_robot_env(cfg=cfg.env) online_env, teleop_device = make_robot_env(cfg=cfg.env)
env_processor, action_processor = make_processors(online_env, cfg.env)
set_seed(cfg.seed) set_seed(cfg.seed)
device = get_safe_torch_device(cfg.policy.device, log=True) device = get_safe_torch_device(cfg.policy.device, log=True)
@@ -257,6 +264,13 @@ def act_with_policy(
assert isinstance(policy, nn.Module) assert isinstance(policy, nn.Module)
obs, info = online_env.reset() obs, info = online_env.reset()
complementary_data = {"raw_joint_positions": info.pop("raw_joint_positions")}
env_processor.reset()
action_processor.reset()
# Process initial observation
transition = create_transition(observation=obs, info=info, complementary_data=complementary_data)
transition = env_processor(transition)
# NOTE: For the moment we will solely handle the case of a single environment # NOTE: For the moment we will solely handle the case of a single environment
sum_reward_episode = 0 sum_reward_episode = 0
@@ -274,45 +288,57 @@ def act_with_policy(
logging.info("[ACTOR] Shutting down act_with_policy") logging.info("[ACTOR] Shutting down act_with_policy")
return return
if interaction_step >= cfg.policy.online_step_before_learning: observation = transition[TransitionKey.OBSERVATION]
# Time policy inference and check if it meets FPS requirement
with policy_timer:
action = policy.select_action(batch=obs)
policy_fps = policy_timer.fps_last
log_policy_frequency_issue(policy_fps=policy_fps, cfg=cfg, interaction_step=interaction_step) # Time policy inference and check if it meets FPS requirement
with policy_timer:
# Extract observation from transition for policy
action = policy.select_action(batch=observation)
policy_fps = policy_timer.fps_last
else: log_policy_frequency_issue(policy_fps=policy_fps, cfg=cfg, interaction_step=interaction_step)
action = online_env.action_space.sample()
next_obs, reward, done, truncated, info = online_env.step(action) # Use the new step function
new_transition, terminate_episode = step_env_and_process_transition(
env=online_env,
transition=transition,
action=action,
teleop_device=teleop_device,
env_processor=env_processor,
action_processor=action_processor,
)
# Extract values from processed transition
next_observation = new_transition[TransitionKey.OBSERVATION]
executed_action = new_transition[TransitionKey.COMPLEMENTARY_DATA]["teleop_action"]
reward = new_transition[TransitionKey.REWARD]
done = new_transition.get(TransitionKey.DONE, False)
truncated = new_transition.get(TransitionKey.TRUNCATED, False)
sum_reward_episode += float(reward) sum_reward_episode += float(reward)
# Increment total steps counter for intervention rate
episode_total_steps += 1 episode_total_steps += 1
# NOTE: We override the action if the intervention is True, because the action applied is the intervention action # Check for intervention from transition info
if "is_intervention" in info and info["is_intervention"]: intervention_info = new_transition[TransitionKey.INFO]
# NOTE: The action space for demonstration before hand is with the full action space if intervention_info.get("is_intervention", False):
# but sometimes for example we want to deactivate the gripper
action = info["action_intervention"]
episode_intervention = True episode_intervention = True
# Increment intervention steps counter
episode_intervention_steps += 1 episode_intervention_steps += 1
# Create transition for learner (convert to old format)
list_transition_to_send_to_learner.append( list_transition_to_send_to_learner.append(
Transition( Transition(
state=obs, state=observation,
action=action, action=executed_action,
reward=reward, reward=reward,
next_state=next_obs, next_state=next_observation,
done=done, done=done,
truncated=truncated, # TODO: (azouitine) Handle truncation properly truncated=truncated,
complementary_info=info, complementary_info={}, # new_transition[TransitionKey.COMPLEMENTARY_DATA],
) )
) )
# assign obs to the next obs and continue the rollout
obs = next_obs # Update transition for next iteration
transition = new_transition
if done or truncated: if done or truncated:
logging.info(f"[ACTOR] Global step {interaction_step}: Episode reward: {sum_reward_episode}") logging.info(f"[ACTOR] Global step {interaction_step}: Episode reward: {sum_reward_episode}")
@@ -347,12 +373,21 @@ def act_with_policy(
) )
) )
# Reset intervention counters # Reset intervention counters and environment
sum_reward_episode = 0.0 sum_reward_episode = 0.0
episode_intervention = False episode_intervention = False
episode_intervention_steps = 0 episode_intervention_steps = 0
episode_total_steps = 0 episode_total_steps = 0
# Reset environment and processors
obs, info = online_env.reset() obs, info = online_env.reset()
complementary_data = {"raw_joint_positions": info.pop("raw_joint_positions")}
env_processor.reset()
action_processor.reset()
# Process initial observation
transition = create_transition(observation=obs, info=info, complementary_data=complementary_data)
transition = env_processor(transition)
if cfg.env.fps is not None: if cfg.env.fps is not None:
dt_time = time.perf_counter() - start_time dt_time = time.perf_counter() - start_time
File diff suppressed because it is too large Load Diff
+1 -5
View File
@@ -286,10 +286,6 @@ def train(cfg: TrainPipelineConfig):
policy.push_model_to_hub(cfg) policy.push_model_to_hub(cfg)
def main(): if __name__ == "__main__":
init_logging() init_logging()
train() train()
if __name__ == "__main__":
main()
+1 -5
View File
@@ -80,9 +80,5 @@ def setup_motors(cfg: SetupConfig):
device.setup_motors() device.setup_motors()
def main():
setup_motors()
if __name__ == "__main__": if __name__ == "__main__":
main() setup_motors()
+1 -5
View File
@@ -153,9 +153,5 @@ def teleoperate(cfg: TeleoperateConfig):
robot.disconnect() robot.disconnect()
def main():
teleoperate()
if __name__ == "__main__": if __name__ == "__main__":
main() teleoperate()
@@ -107,6 +107,45 @@ class GamepadTeleop(Teleoperator):
return action_dict return action_dict
def get_teleop_events(self) -> dict[str, Any]:
"""
Get extra control events from the gamepad such as intervention status,
episode termination, success indicators, etc.
Returns:
Dictionary containing:
- is_intervention: bool - Whether human is currently intervening
- terminate_episode: bool - Whether to terminate the current episode
- success: bool - Whether the episode was successful
- rerecord_episode: bool - Whether to rerecord the episode
"""
if self.gamepad is None:
return {
"is_intervention": False,
"terminate_episode": False,
"success": False,
"rerecord_episode": False,
}
# Update gamepad state to get fresh inputs
self.gamepad.update()
# Check if intervention is active
is_intervention = self.gamepad.should_intervene()
# Get episode end status
episode_end_status = self.gamepad.get_episode_end_status()
terminate_episode = episode_end_status is not None
success = episode_end_status == "success"
rerecord_episode = episode_end_status == "rerecord_episode"
return {
"is_intervention": is_intervention,
"terminate_episode": terminate_episode,
"success": success,
"rerecord_episode": rerecord_episode,
}
def disconnect(self) -> None: def disconnect(self) -> None:
"""Disconnect from the gamepad.""" """Disconnect from the gamepad."""
if self.gamepad is not None: if self.gamepad is not None:
@@ -235,3 +235,67 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop):
action_dict["gripper"] = gripper_action action_dict["gripper"] = gripper_action
return action_dict return action_dict
def get_teleop_events(self) -> dict[str, Any]:
"""
Get extra control events from the keyboard such as intervention status,
episode termination, success indicators, etc.
Keyboard mappings:
- Any movement keys pressed = intervention active
- 's' key = success (terminate episode successfully)
- 'r' key = rerecord episode (terminate and rerecord)
- 'q' key = quit episode (terminate without success)
Returns:
Dictionary containing:
- is_intervention: bool - Whether human is currently intervening
- terminate_episode: bool - Whether to terminate the current episode
- success: bool - Whether the episode was successful
- rerecord_episode: bool - Whether to rerecord the episode
"""
if not self.is_connected:
return {
"is_intervention": False,
"terminate_episode": False,
"success": False,
"rerecord_episode": False,
}
# Check if any movement keys are currently pressed (indicates intervention)
movement_keys = [
keyboard.Key.up,
keyboard.Key.down,
keyboard.Key.left,
keyboard.Key.right,
keyboard.Key.shift,
keyboard.Key.shift_r,
keyboard.Key.ctrl_r,
keyboard.Key.ctrl_l,
]
is_intervention = any(self.current_pressed.get(key, False) for key in movement_keys)
# Check for episode control commands from misc_keys_queue
terminate_episode = False
success = False
rerecord_episode = False
# Process any pending misc keys
while not self.misc_keys_queue.empty():
key = self.misc_keys_queue.get_nowait()
if key == "s":
terminate_episode = True
success = True
elif key == "r":
terminate_episode = True
rerecord_episode = True
elif key == "q":
terminate_episode = True
success = False
return {
"is_intervention": is_intervention,
"terminate_episode": terminate_episode,
"success": success,
"rerecord_episode": rerecord_episode,
}
+12
View File
@@ -160,6 +160,18 @@ class Teleoperator(abc.ABC):
""" """
pass pass
@abc.abstractmethod
def get_teleop_events(self) -> dict[str, Any]:
"""
Get extra control events from the teleoperator such as intervention status,
episode termination, success indicators, etc.
Check the implementation of the gamepad for an example.
Returns:
dict[str, Any]: A dictionary containing control events with keys and values that are specific to the setup.
"""
pass
@abc.abstractmethod @abc.abstractmethod
def send_feedback(self, feedback: dict[str, Any]) -> None: def send_feedback(self, feedback: dict[str, Any]) -> None:
""" """
@@ -0,0 +1,195 @@
---
library_name: lerobot
tags:
- robotics
- lerobot
- safetensors
pipeline_tag: robotics
---
# RobotProcessor
## Overview
RobotProcessor is a composable, debuggable post-processing pipeline for robot transitions in the LeRobot framework. It orchestrates an ordered collection of small, functional transforms (steps) that are executed left-to-right on each incoming `EnvTransition`.
## Architecture
The RobotProcessor provides a modular architecture for processing robot environment transitions through a sequence of composable steps. Each step is a callable that accepts a full `EnvTransition` tuple and returns a potentially modified tuple of the same structure.
### EnvTransition Structure
An `EnvTransition` is a 7-tuple containing:
1. **observation**: Current state observation
2. **action**: Action taken (can be None)
3. **reward**: Reward received (float or None)
4. **done**: Episode termination flag (bool or None)
5. **truncated**: Episode truncation flag (bool or None)
6. **info**: Additional information dictionary
7. **complementary_data**: Extra data dictionary
## Key Features
- **Composable Pipeline**: Chain multiple processing steps in a specific order
- **State Persistence**: Save and load processor state using SafeTensors format
- **Hugging Face Hub Integration**: Easy sharing and loading via `save_pretrained()` and `from_pretrained()`
- **Debugging Support**: Step-through functionality to inspect intermediate transformations
- **Hook System**: Before/after step hooks for additional processing or monitoring
- **Device Support**: Move tensor states to different devices (CPU/GPU)
- **Performance Profiling**: Built-in profiling to identify bottlenecks
## Installation
Follow the [installation instructions](https://huggingface.co/docs/lerobot/installation) to install the package.
## Usage
### Basic Example
```python
from lerobot.processor.pipeline import RobotProcessor
from your_steps import ObservationNormalizer, VelocityCalculator
# Create a processor with multiple steps
processor = RobotProcessor(
steps=[
ObservationNormalizer(mean=0, std=1),
VelocityCalculator(window_size=5),
],
name="my_robot_processor",
seed=42
)
# Process a transition
obs, info = env.reset()
transition = (obs, None, 0.0, False, False, info, {})
processed_transition = processor(transition)
# Extract processed observation
processed_obs = processed_transition[0]
```
### Saving and Loading
```python
# Save locally
processor.save_pretrained("./my_processor")
# Push to Hugging Face Hub
processor.push_to_hub("username/my-robot-processor")
# Load from Hub
loaded_processor = RobotProcessor.from_pretrained("username/my-robot-processor")
```
### Debugging with Step-Through
```python
# Inspect intermediate results
for idx, intermediate_transition in enumerate(processor.step_through(transition)):
print(f"After step {idx}: {intermediate_transition[0]}") # Print observation
```
### Using Hooks
```python
# Add monitoring hook
def log_observation(step_idx, transition):
print(f"Step {step_idx}: obs shape = {transition[0].shape}")
return None # Don't modify transition
processor.register_before_step_hook(log_observation)
```
## Creating Custom Steps
To create a custom processor step, implement the `ProcessorStep` protocol:
```python
from lerobot.processor.pipeline import ProcessorStepRegistry, EnvTransition
@ProcessorStepRegistry.register("my_custom_step")
class MyCustomStep:
def __init__(self, param1=1.0):
self.param1 = param1
self.buffer = []
def __call__(self, transition: EnvTransition) -> EnvTransition:
obs, action, reward, done, truncated, info, comp_data = transition
# Process observation
processed_obs = obs * self.param1
return (processed_obs, action, reward, done, truncated, info, comp_data)
def get_config(self) -> dict:
return {"param1": self.param1}
def state_dict(self) -> dict:
# Return only torch.Tensor state
return {}
def load_state_dict(self, state: dict) -> None:
# Load tensor state
pass
def reset(self) -> None:
# Clear buffers at episode boundaries
self.buffer.clear()
```
## Advanced Features
### Device Management
```python
# Move all tensor states to GPU
processor = processor.to("cuda")
# Move to specific device
processor = processor.to(torch.device("cuda:1"))
```
### Performance Profiling
```python
# Profile step execution times
profile_results = processor.profile_steps(transition, num_runs=100)
for step_name, time_ms in profile_results.items():
print(f"{step_name}: {time_ms:.3f} ms")
```
### Processor Slicing
```python
# Get a single step
first_step = processor[0]
# Create a sub-processor with steps 1-3
sub_processor = processor[1:4]
```
## Model Card Specifications
- **Pipeline Tag**: robotics
- **Library**: lerobot
- **Format**: safetensors
- **License**: Apache 2.0
## Limitations
- Steps must maintain the 7-tuple structure of EnvTransition
- All tensor state must be separated from configuration for proper serialization
- Steps are executed sequentially (no parallel processing within a single transition)
## Citation
If you use RobotProcessor in your research, please cite:
```bibtex
@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},
title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch},
howpublished = "\url{https://github.com/huggingface/lerobot}",
year = {2024}
}
```
+7 -2
View File
@@ -22,7 +22,7 @@ from gymnasium.utils.env_checker import check_env
import lerobot import lerobot
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.processor import RobotProcessor, TransitionKey, VanillaObservationProcessor
from tests.utils import require_env from tests.utils import require_env
OBS_TYPES = ["state", "pixels", "pixels_agent_pos"] OBS_TYPES = ["state", "pixels", "pixels_agent_pos"]
@@ -48,7 +48,12 @@ def test_factory(env_name):
cfg = make_env_config(env_name) cfg = make_env_config(env_name)
env = make_env(cfg, n_envs=1) env = make_env(cfg, n_envs=1)
obs, _ = env.reset() obs, _ = env.reset()
obs = preprocess_observation(obs)
# Process observation using processor
obs_processor = RobotProcessor([VanillaObservationProcessor()])
transition = (obs, None, None, None, None, None, None)
processed_transition = obs_processor(transition)
obs = processed_transition[TransitionKey.OBSERVATION]
# test image keys are float32 in range [0,1] # test image keys are float32 in range [0,1]
for key in obs: for key in obs:
+5 -52
View File
@@ -27,13 +27,10 @@ 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.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,
@@ -42,6 +39,7 @@ from lerobot.policies.factory import (
) )
from lerobot.policies.normalize import Normalize, Unnormalize from lerobot.policies.normalize import Normalize, Unnormalize
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.processor import RobotProcessor, TransitionKey, VanillaObservationProcessor
from lerobot.utils.random_utils import seeded_context from lerobot.utils.random_utils import seeded_context
from tests.artifacts.policies.save_policy_to_safetensors import get_policy_stats from tests.artifacts.policies.save_policy_to_safetensors import get_policy_stats
from tests.utils import DEVICE, require_cpu, require_env, require_x86_64_kernel from tests.utils import DEVICE, require_cpu, require_env, require_x86_64_kernel
@@ -187,7 +185,10 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
observation, _ = env.reset(seed=train_cfg.seed) observation, _ = env.reset(seed=train_cfg.seed)
# apply transform to normalize the observations # apply transform to normalize the observations
observation = preprocess_observation(observation) obs_processor = RobotProcessor([VanillaObservationProcessor()])
transition = (observation, None, None, None, None, None, None)
processed_transition = obs_processor(transition)
observation = processed_transition[TransitionKey.OBSERVATION]
# send observation to device/gpu # send observation to device/gpu
observation = {key: observation[key].to(DEVICE, non_blocking=True) for key in observation} observation = {key: observation[key].to(DEVICE, non_blocking=True) for key in observation}
@@ -365,54 +366,6 @@ 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",
[ [
+34 -19
View File
@@ -20,7 +20,11 @@ import torch
from lerobot.configs.types import FeatureType from lerobot.configs.types import FeatureType
from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
from lerobot.processor import VanillaObservationProcessor from lerobot.processor import (
ImageProcessor,
StateProcessor,
VanillaObservationProcessor,
)
from lerobot.processor.pipeline import TransitionKey from lerobot.processor.pipeline import TransitionKey
from tests.conftest import assert_contract_is_typed from tests.conftest import assert_contract_is_typed
@@ -42,7 +46,7 @@ def create_transition(
def test_process_single_image(): def test_process_single_image():
"""Test processing a single image.""" """Test processing a single image."""
processor = VanillaObservationProcessor() processor = ImageProcessor()
# Create a mock image (H, W, C) format, uint8 # Create a mock image (H, W, C) format, uint8
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
@@ -68,7 +72,7 @@ def test_process_single_image():
def test_process_image_dict(): def test_process_image_dict():
"""Test processing multiple images in a dictionary.""" """Test processing multiple images in a dictionary."""
processor = VanillaObservationProcessor() processor = ImageProcessor()
# Create mock images # Create mock images
image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8) image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8)
@@ -91,7 +95,7 @@ def test_process_image_dict():
def test_process_batched_image(): def test_process_batched_image():
"""Test processing already batched images.""" """Test processing already batched images."""
processor = VanillaObservationProcessor() processor = ImageProcessor()
# Create a batched image (B, H, W, C) # Create a batched image (B, H, W, C)
image = np.random.randint(0, 256, size=(2, 64, 64, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(2, 64, 64, 3), dtype=np.uint8)
@@ -108,7 +112,7 @@ def test_process_batched_image():
def test_invalid_image_format(): def test_invalid_image_format():
"""Test error handling for invalid image formats.""" """Test error handling for invalid image formats."""
processor = VanillaObservationProcessor() processor = ImageProcessor()
# Test wrong channel order (channels first) # Test wrong channel order (channels first)
image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8) image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8)
@@ -121,7 +125,7 @@ def test_invalid_image_format():
def test_invalid_image_dtype(): def test_invalid_image_dtype():
"""Test error handling for invalid image dtype.""" """Test error handling for invalid image dtype."""
processor = VanillaObservationProcessor() processor = ImageProcessor()
# Test wrong dtype # Test wrong dtype
image = np.random.rand(64, 64, 3).astype(np.float32) image = np.random.rand(64, 64, 3).astype(np.float32)
@@ -134,7 +138,7 @@ def test_invalid_image_dtype():
def test_no_pixels_in_observation(): def test_no_pixels_in_observation():
"""Test processor when no pixels are in observation.""" """Test processor when no pixels are in observation."""
processor = VanillaObservationProcessor() processor = ImageProcessor()
observation = {"other_data": np.array([1, 2, 3])} observation = {"other_data": np.array([1, 2, 3])}
transition = create_transition(observation=observation) transition = create_transition(observation=observation)
@@ -149,7 +153,7 @@ def test_no_pixels_in_observation():
def test_none_observation(): def test_none_observation():
"""Test processor with None observation.""" """Test processor with None observation."""
processor = VanillaObservationProcessor() processor = ImageProcessor()
transition = create_transition() transition = create_transition()
result = processor(transition) result = processor(transition)
@@ -159,7 +163,7 @@ def test_none_observation():
def test_serialization_methods(): def test_serialization_methods():
"""Test serialization methods.""" """Test serialization methods."""
processor = VanillaObservationProcessor() processor = ImageProcessor()
# Test get_config # Test get_config
config = processor.get_config() config = processor.get_config()
@@ -178,7 +182,7 @@ def test_serialization_methods():
def test_process_environment_state(): def test_process_environment_state():
"""Test processing environment_state.""" """Test processing environment_state."""
processor = VanillaObservationProcessor() processor = StateProcessor()
env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32) env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
observation = {"environment_state": env_state} observation = {"environment_state": env_state}
@@ -199,7 +203,7 @@ def test_process_environment_state():
def test_process_agent_pos(): def test_process_agent_pos():
"""Test processing agent_pos.""" """Test processing agent_pos."""
processor = VanillaObservationProcessor() processor = StateProcessor()
agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32) agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
observation = {"agent_pos": agent_pos} observation = {"agent_pos": agent_pos}
@@ -220,7 +224,7 @@ def test_process_agent_pos():
def test_process_batched_states(): def test_process_batched_states():
"""Test processing already batched states.""" """Test processing already batched states."""
processor = VanillaObservationProcessor() processor = StateProcessor()
env_state = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) 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) agent_pos = np.array([[0.5, -0.5], [1.0, -1.0]], dtype=np.float32)
@@ -238,7 +242,7 @@ def test_process_batched_states():
def test_process_both_states(): def test_process_both_states():
"""Test processing both environment_state and agent_pos.""" """Test processing both environment_state and agent_pos."""
processor = VanillaObservationProcessor() processor = StateProcessor()
env_state = np.array([1.0, 2.0], dtype=np.float32) env_state = np.array([1.0, 2.0], dtype=np.float32)
agent_pos = np.array([0.5, -0.5], dtype=np.float32) agent_pos = np.array([0.5, -0.5], dtype=np.float32)
@@ -263,7 +267,7 @@ def test_process_both_states():
def test_no_states_in_observation(): def test_no_states_in_observation():
"""Test processor when no states are in observation.""" """Test processor when no states are in observation."""
processor = VanillaObservationProcessor() processor = StateProcessor()
observation = {"other_data": np.array([1, 2, 3])} observation = {"other_data": np.array([1, 2, 3])}
transition = create_transition(observation=observation) transition = create_transition(observation=observation)
@@ -355,6 +359,17 @@ def test_empty_observation():
assert processed_obs == {} assert processed_obs == {}
def test_custom_sub_processors():
"""Test ObservationProcessor with custom sub-processors."""
image_proc = ImageProcessor()
state_proc = StateProcessor()
processor = VanillaObservationProcessor(image_processor=image_proc, state_processor=state_proc)
# Should use the provided processors
assert processor.image_processor is image_proc
assert processor.state_processor is state_proc
def test_equivalent_to_original_function(): def test_equivalent_to_original_function():
"""Test that ObservationProcessor produces equivalent results to preprocess_observation.""" """Test that ObservationProcessor produces equivalent results to preprocess_observation."""
# Import the original function for comparison # Import the original function for comparison
@@ -411,7 +426,7 @@ def test_equivalent_with_image_dict():
def test_image_processor_feature_contract_pixels_to_image(policy_feature_factory): def test_image_processor_feature_contract_pixels_to_image(policy_feature_factory):
processor = VanillaObservationProcessor() processor = ImageProcessor()
features = { features = {
"pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"keep": policy_feature_factory(FeatureType.ENV, (1,)), "keep": policy_feature_factory(FeatureType.ENV, (1,)),
@@ -425,7 +440,7 @@ def test_image_processor_feature_contract_pixels_to_image(policy_feature_factory
def test_image_processor_feature_contract_observation_pixels_to_image(policy_feature_factory): def test_image_processor_feature_contract_observation_pixels_to_image(policy_feature_factory):
processor = VanillaObservationProcessor() processor = ImageProcessor()
features = { features = {
"observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"keep": policy_feature_factory(FeatureType.ENV, (1,)), "keep": policy_feature_factory(FeatureType.ENV, (1,)),
@@ -439,7 +454,7 @@ def test_image_processor_feature_contract_observation_pixels_to_image(policy_fea
def test_image_processor_feature_contract_multi_camera_and_prefixed(policy_feature_factory): def test_image_processor_feature_contract_multi_camera_and_prefixed(policy_feature_factory):
processor = VanillaObservationProcessor() processor = ImageProcessor()
features = { features = {
"pixels.front": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "pixels.front": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
"pixels.wrist": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "pixels.wrist": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
@@ -457,7 +472,7 @@ def test_image_processor_feature_contract_multi_camera_and_prefixed(policy_featu
def test_state_processor_feature_contract_environment_and_agent_pos(policy_feature_factory): def test_state_processor_feature_contract_environment_and_agent_pos(policy_feature_factory):
processor = VanillaObservationProcessor() processor = StateProcessor()
features = { features = {
"environment_state": policy_feature_factory(FeatureType.STATE, (3,)), "environment_state": policy_feature_factory(FeatureType.STATE, (3,)),
"agent_pos": policy_feature_factory(FeatureType.STATE, (7,)), "agent_pos": policy_feature_factory(FeatureType.STATE, (7,)),
@@ -473,7 +488,7 @@ def test_state_processor_feature_contract_environment_and_agent_pos(policy_featu
def test_state_processor_feature_contract_prefixed_inputs(policy_feature_factory): def test_state_processor_feature_contract_prefixed_inputs(policy_feature_factory):
proc = VanillaObservationProcessor() proc = StateProcessor()
features = { features = {
"observation.environment_state": policy_feature_factory(FeatureType.STATE, (2,)), "observation.environment_state": policy_feature_factory(FeatureType.STATE, (2,)),
"observation.agent_pos": policy_feature_factory(FeatureType.STATE, (4,)), "observation.agent_pos": policy_feature_factory(FeatureType.STATE, (4,)),
+467 -24
View File
@@ -21,6 +21,7 @@ from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
import numpy as np
import pytest import pytest
import torch import torch
import torch.nn as nn import torch.nn as nn
@@ -267,25 +268,14 @@ def test_step_through_with_dict():
assert len(results) == 3 # Original + 2 steps assert len(results) == 3 # Original + 2 steps
# Ensure all results are EnvTransition dicts (regardless of input format) # Ensure all results are dicts (same format as input)
for result in results: for result in results:
assert isinstance(result, dict) assert isinstance(result, dict)
# Check that keys are TransitionKey enums or at least valid transition keys
for key in result:
assert key in [
TransitionKey.OBSERVATION,
TransitionKey.ACTION,
TransitionKey.REWARD,
TransitionKey.DONE,
TransitionKey.TRUNCATED,
TransitionKey.INFO,
TransitionKey.COMPLEMENTARY_DATA,
]
# Check that the processing worked - verify step counters in complementary_data # Check that the processing worked - the complementary data from steps
assert results[1].get(TransitionKey.COMPLEMENTARY_DATA, {}).get("step1_counter") == 0 # should show up in the info or complementary_data fields when converted back to dict
assert results[2].get(TransitionKey.COMPLEMENTARY_DATA, {}).get("step1_counter") == 0 # Note: This depends on how _default_transition_to_batch handles complementary_data
assert results[2].get(TransitionKey.COMPLEMENTARY_DATA, {}).get("step2_counter") == 0 # For now, just check that we get dict outputs
def test_step_through_no_hooks(): def test_step_through_no_hooks():
@@ -363,6 +353,32 @@ def test_hooks():
assert after_calls == [0] assert after_calls == [0]
def test_reset():
"""Test pipeline reset functionality."""
step = MockStep("test_step")
pipeline = RobotProcessor([step])
reset_called = []
def reset_hook():
reset_called.append(True)
pipeline.register_reset_hook(reset_hook)
# Make some calls to increment counter
transition = create_transition()
pipeline(transition)
pipeline(transition)
assert step.counter == 2
# Reset should reset step and call hook
pipeline.reset()
assert step.counter == 0
assert len(reset_called) == 1
def test_unregister_hooks(): def test_unregister_hooks():
"""Test unregistering hooks from the pipeline.""" """Test unregistering hooks from the pipeline."""
step = MockStep("test_step") step = MockStep("test_step")
@@ -402,6 +418,21 @@ def test_unregister_hooks():
pipeline(transition) pipeline(transition)
assert len(after_calls) == 0 assert len(after_calls) == 0
# Test reset_hook
reset_calls = []
def reset_hook():
reset_calls.append(True)
pipeline.register_reset_hook(reset_hook)
pipeline.reset()
assert len(reset_calls) == 1
pipeline.unregister_reset_hook(reset_hook)
reset_calls.clear()
pipeline.reset()
assert len(reset_calls) == 0
def test_unregister_nonexistent_hook(): def test_unregister_nonexistent_hook():
"""Test error handling when unregistering hooks that don't exist.""" """Test error handling when unregistering hooks that don't exist."""
@@ -420,6 +451,9 @@ def test_unregister_nonexistent_hook():
with pytest.raises(ValueError, match="not found in after_step_hooks"): with pytest.raises(ValueError, match="not found in after_step_hooks"):
pipeline.unregister_after_step_hook(some_hook) pipeline.unregister_after_step_hook(some_hook)
with pytest.raises(ValueError, match="not found in reset_hooks"):
pipeline.unregister_reset_hook(reset_hook)
def test_multiple_hooks_and_selective_unregister(): def test_multiple_hooks_and_selective_unregister():
"""Test registering multiple hooks and selectively unregistering them.""" """Test registering multiple hooks and selectively unregistering them."""
@@ -508,6 +542,22 @@ def test_hook_execution_order_documentation():
assert execution_order == ["A", "C", "B"] # B is now last assert execution_order == ["A", "C", "B"] # B is now last
def test_profile_steps():
"""Test step profiling functionality."""
step1 = MockStep("step1")
step2 = MockStep("step2")
pipeline = RobotProcessor([step1, step2])
transition = create_transition()
profile_results = pipeline.profile_steps(transition, num_runs=10)
assert len(profile_results) == 2
assert "step_0_MockStep" in profile_results
assert "step_1_MockStep" in profile_results
assert all(isinstance(time, float) and time >= 0 for time in profile_results.values())
def test_save_and_load_pretrained(): def test_save_and_load_pretrained():
"""Test saving and loading pipeline. """Test saving and loading pipeline.
@@ -521,7 +571,7 @@ def test_save_and_load_pretrained():
step1.counter = 5 step1.counter = 5
step2.counter = 10 step2.counter = 10
pipeline = RobotProcessor([step1, step2], name="TestPipeline") pipeline = RobotProcessor([step1, step2], name="TestPipeline", seed=42)
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
# Save pipeline # Save pipeline
@@ -536,6 +586,7 @@ def test_save_and_load_pretrained():
config = json.load(f) config = json.load(f)
assert config["name"] == "TestPipeline" assert config["name"] == "TestPipeline"
assert config["seed"] == 42
assert len(config["steps"]) == 2 assert len(config["steps"]) == 2
# Verify counters are saved in config, not in separate state files # Verify counters are saved in config, not in separate state files
@@ -546,6 +597,7 @@ def test_save_and_load_pretrained():
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir)
assert loaded_pipeline.name == "TestPipeline" assert loaded_pipeline.name == "TestPipeline"
assert loaded_pipeline.seed == 42
assert len(loaded_pipeline) == 2 assert len(loaded_pipeline) == 2
# Check that counter was restored from config # Check that counter was restored from config
@@ -667,6 +719,182 @@ class MockModuleStep(nn.Module):
return features return features
def test_to_device_with_state_dict():
"""Test moving pipeline to device for steps with state_dict."""
step = MockStepWithTensorState(name="device_test", window_size=5)
pipeline = RobotProcessor([step])
# Process some transitions to populate state
for i in range(10):
transition = create_transition(reward=float(i))
pipeline(transition)
# Check initial device (should be CPU)
assert step.running_mean.device.type == "cpu"
assert step.running_count.device.type == "cpu"
# Move to same device (CPU)
result = pipeline.to("cpu")
assert result is pipeline # Check it returns self
assert step.running_mean.device.type == "cpu"
assert step.running_count.device.type == "cpu"
# Test with torch.device object
result = pipeline.to(torch.device("cpu"))
assert result is pipeline
assert step.running_mean.device.type == "cpu"
# If CUDA is available, test GPU transfer
if torch.cuda.is_available():
result = pipeline.to("cuda")
assert result is pipeline
assert step.running_mean.device.type == "cuda"
assert step.running_count.device.type == "cuda"
# Move back to CPU
pipeline.to("cpu")
assert step.running_mean.device.type == "cpu"
assert step.running_count.device.type == "cpu"
def test_to_device_with_module():
"""Test moving pipeline to device for steps that inherit from nn.Module.
Even though the step inherits from nn.Module, the pipeline will use the
state_dict/load_state_dict approach to move tensors to the device.
"""
module_step = MockModuleStep(input_dim=5, hidden_dim=3)
pipeline = RobotProcessor([module_step])
# Process some data
obs = torch.randn(2, 5)
transition = create_transition(observation=obs, reward=1.0)
pipeline(transition)
# Check initial device
assert module_step.linear.weight.device.type == "cpu"
assert module_step.running_mean.device.type == "cpu"
# Move to same device
result = pipeline.to("cpu")
assert result is pipeline
assert module_step.linear.weight.device.type == "cpu"
assert module_step.running_mean.device.type == "cpu"
# If CUDA is available, test GPU transfer
if torch.cuda.is_available():
result = pipeline.to("cuda:0")
assert result is pipeline
assert module_step.linear.weight.device.type == "cuda"
assert module_step.linear.weight.device.index == 0
assert module_step.running_mean.device.type == "cuda"
assert module_step.running_mean.device.index == 0
# Verify the module still works after transfer
obs_cuda = torch.randn(2, 5, device="cuda:0")
transition = create_transition(observation=obs_cuda, reward=1.0)
pipeline(transition) # Should not raise an error
def test_to_device_mixed_steps():
"""Test moving pipeline with various types of steps, all using state_dict approach."""
module_step = MockModuleStep()
state_dict_step = MockStepWithTensorState()
simple_step = MockStepWithoutOptionalMethods() # No tensor state
pipeline = RobotProcessor([module_step, state_dict_step, simple_step])
# Process some data
for i in range(5):
transition = create_transition(observation=torch.randn(2, 10), reward=float(i))
pipeline(transition)
# Check initial state
assert module_step.linear.weight.device.type == "cpu"
assert state_dict_step.running_mean.device.type == "cpu"
# Move to device
result = pipeline.to("cpu")
assert result is pipeline
if torch.cuda.is_available():
pipeline.to("cuda")
assert module_step.linear.weight.device.type == "cuda"
assert module_step.running_mean.device.type == "cuda"
assert state_dict_step.running_mean.device.type == "cuda"
assert state_dict_step.running_count.device.type == "cuda"
def test_to_device_empty_state():
"""Test moving pipeline with steps that have empty state_dict."""
step = MockStep("empty_state") # This step has empty state_dict
pipeline = RobotProcessor([step])
# Should not raise an error even with empty state
result = pipeline.to("cpu")
assert result is pipeline
if torch.cuda.is_available():
result = pipeline.to("cuda")
assert result is pipeline
def test_to_device_preserves_functionality():
"""Test that pipeline functionality is preserved after device transfer."""
step = MockStepWithTensorState(window_size=3)
pipeline = RobotProcessor([step])
# Process initial data
rewards = [1.0, 2.0, 3.0]
for r in rewards:
transition = create_transition(reward=r)
pipeline(transition)
# Check state before transfer
initial_mean = step.running_mean.clone()
initial_count = step.running_count.clone()
# Move to device (CPU to CPU in this case, but tests the mechanism)
pipeline.to("cpu")
# Verify state is preserved
assert torch.allclose(step.running_mean, initial_mean)
assert step.running_count == initial_count
# Process more data to ensure functionality
transition = create_transition(reward=4.0)
_ = pipeline(transition)
assert step.running_count == 4
assert step.running_mean[0] == 4.0 # First slot should have been overwritten with 4.0
def test_to_device_invalid_device():
"""Test error handling for invalid devices."""
pipeline = RobotProcessor([MockStep()])
# Invalid device names should raise an error from PyTorch
with pytest.raises(RuntimeError):
pipeline.to("invalid_device")
def test_to_device_chaining():
"""Test that to() returns self for method chaining."""
step1 = MockStepWithTensorState()
step2 = MockModuleStep()
pipeline = RobotProcessor([step1, step2])
# Test chaining
result = pipeline.to("cpu").reset()
assert result is None # reset() returns None
# Can chain multiple to() calls
result1 = pipeline.to("cpu")
result2 = result1.to("cpu")
assert result1 is pipeline
assert result2 is pipeline
class MockNonModuleStepWithState: class MockNonModuleStepWithState:
"""Mock step that explicitly does NOT inherit from nn.Module but has tensor state. """Mock step that explicitly does NOT inherit from nn.Module but has tensor state.
@@ -749,6 +977,129 @@ class MockNonModuleStepWithState:
return features return features
def test_to_device_non_module_class():
"""Test moving pipeline to device for regular classes (non nn.Module) with tensor state.
This ensures the state_dict/load_state_dict approach works for classes that
don't inherit from nn.Module but still have tensor state to manage.
"""
# Create a non-module step with tensor state
non_module_step = MockNonModuleStepWithState(name="device_test", feature_dim=5)
pipeline = RobotProcessor([non_module_step])
# Process some data to populate state
for i in range(3):
obs = torch.randn(2, 5)
transition = create_transition(observation=obs, reward=float(i))
result = pipeline(transition)
comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
assert f"{non_module_step.name}_steps" in comp_data
# Verify all tensors are on CPU initially
assert non_module_step.weights.device.type == "cpu"
assert non_module_step.bias.device.type == "cpu"
assert non_module_step.running_stats.device.type == "cpu"
assert non_module_step.step_count.device.type == "cpu"
# Verify step count
assert non_module_step.step_count.item() == 3
# Store initial values for comparison
initial_weights = non_module_step.weights.clone()
initial_bias = non_module_step.bias.clone()
initial_stats = non_module_step.running_stats.clone()
# Move to same device (CPU)
result = pipeline.to("cpu")
assert result is pipeline
# Verify tensors are still on CPU and values unchanged
assert non_module_step.weights.device.type == "cpu"
assert torch.allclose(non_module_step.weights, initial_weights)
assert torch.allclose(non_module_step.bias, initial_bias)
assert torch.allclose(non_module_step.running_stats, initial_stats)
# If CUDA is available, test GPU transfer
if torch.cuda.is_available():
# Move to GPU
pipeline.to("cuda")
# Verify all tensors moved to GPU
assert non_module_step.weights.device.type == "cuda"
assert non_module_step.bias.device.type == "cuda"
assert non_module_step.running_stats.device.type == "cuda"
assert non_module_step.step_count.device.type == "cuda"
# Verify values are preserved
assert torch.allclose(non_module_step.weights.cpu(), initial_weights)
assert torch.allclose(non_module_step.bias.cpu(), initial_bias)
assert torch.allclose(non_module_step.running_stats.cpu(), initial_stats)
assert non_module_step.step_count.item() == 3
# Test that step still works on GPU
obs_gpu = torch.randn(2, 5, device="cuda")
transition = create_transition(observation=obs_gpu, reward=1.0)
result = pipeline(transition)
comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
# Verify processing worked
assert comp_data[f"{non_module_step.name}_steps"] == 4
# Move back to CPU
pipeline.to("cpu")
assert non_module_step.weights.device.type == "cpu"
assert non_module_step.step_count.item() == 4
def test_to_device_module_vs_non_module():
"""Test that both nn.Module and non-Module steps work with the same state_dict approach."""
# Create both types of steps
module_step = MockModuleStep(input_dim=5, hidden_dim=3)
non_module_step = MockNonModuleStepWithState(name="non_module", feature_dim=5)
# Create pipeline with both
pipeline = RobotProcessor([module_step, non_module_step])
# Process some data
obs = torch.randn(2, 5)
transition = create_transition(observation=obs, reward=1.0)
_ = pipeline(transition)
# Check initial devices
assert module_step.linear.weight.device.type == "cpu"
assert module_step.running_mean.device.type == "cpu"
assert non_module_step.weights.device.type == "cpu"
assert non_module_step.running_stats.device.type == "cpu"
# Both should have been called
assert module_step.counter == 1
assert non_module_step.step_count.item() == 1
if torch.cuda.is_available():
# Move to GPU
pipeline.to("cuda")
# Verify both types of steps moved correctly
assert module_step.linear.weight.device.type == "cuda"
assert module_step.running_mean.device.type == "cuda"
assert non_module_step.weights.device.type == "cuda"
assert non_module_step.running_stats.device.type == "cuda"
# Process data on GPU
obs_gpu = torch.randn(2, 5, device="cuda")
transition = create_transition(observation=obs_gpu, reward=2.0)
_ = pipeline(transition)
# Verify both steps processed the data
assert module_step.counter == 2
assert non_module_step.step_count.item() == 2
# Move back to CPU and verify
pipeline.to("cpu")
assert module_step.linear.weight.device.type == "cpu"
assert non_module_step.weights.device.type == "cpu"
# Tests for overrides functionality # Tests for overrides functionality
@dataclass @dataclass
class MockStepWithNonSerializableParam: class MockStepWithNonSerializableParam:
@@ -1127,6 +1478,96 @@ def test_from_pretrained_override_error_messages():
assert "registered_mock_step" in error_msg assert "registered_mock_step" in error_msg
class MockStepWithMixedState:
"""Mock step demonstrating proper separation of tensor and non-tensor state.
Non-tensor state should go in get_config(), only tensors in state_dict().
"""
def __init__(self, name: str = "mixed_state"):
self.name = name
self.tensor_data = torch.randn(5)
self.numpy_data = np.array([1, 2, 3, 4, 5]) # Goes in config
self.scalar_value = 42 # Goes in config
self.list_value = [1, 2, 3] # Goes in config
def __call__(self, transition: EnvTransition) -> EnvTransition:
# Simple pass-through
return transition
def state_dict(self) -> dict[str, torch.Tensor]:
"""Return ONLY tensor state as per the type contract."""
return {
"tensor_data": self.tensor_data,
}
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
"""Load tensor state only."""
self.tensor_data = state["tensor_data"]
def get_config(self) -> dict[str, Any]:
"""Non-tensor state goes here."""
return {
"name": self.name,
"numpy_data": self.numpy_data.tolist(), # Convert to list for JSON serialization
"scalar_value": self.scalar_value,
"list_value": self.list_value,
}
def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
# We do not test feature_contract here
return features
def test_to_device_with_mixed_state_types():
"""Test that to() only moves tensor state, while non-tensor state remains in config."""
step = MockStepWithMixedState()
pipeline = RobotProcessor([step])
# Store initial values
initial_numpy = step.numpy_data.copy()
initial_scalar = step.scalar_value
initial_list = step.list_value.copy()
# Check initial state
assert step.tensor_data.device.type == "cpu"
assert isinstance(step.numpy_data, np.ndarray)
assert isinstance(step.scalar_value, int)
assert isinstance(step.list_value, list)
# Verify state_dict only contains tensors
state = step.state_dict()
assert all(isinstance(v, torch.Tensor) for v in state.values())
assert "tensor_data" in state
assert "numpy_data" not in state
# Move to same device
pipeline.to("cpu")
# Verify tensor moved and non-tensor attributes unchanged
assert step.tensor_data.device.type == "cpu"
assert np.array_equal(step.numpy_data, initial_numpy)
assert step.scalar_value == initial_scalar
assert step.list_value == initial_list
if torch.cuda.is_available():
# Move to GPU
pipeline.to("cuda")
# Only tensor should move to GPU
assert step.tensor_data.device.type == "cuda"
# Non-tensor values should remain unchanged
assert isinstance(step.numpy_data, np.ndarray)
assert np.array_equal(step.numpy_data, initial_numpy)
assert step.scalar_value == initial_scalar
assert step.list_value == initial_list
# Move back to CPU
pipeline.to("cpu")
assert step.tensor_data.device.type == "cpu"
def test_repr_empty_processor(): def test_repr_empty_processor():
"""Test __repr__ with empty processor.""" """Test __repr__ with empty processor."""
pipeline = RobotProcessor() pipeline = RobotProcessor()
@@ -1193,10 +1634,10 @@ def test_repr_with_custom_name():
def test_repr_with_seed(): def test_repr_with_seed():
"""Test __repr__ with seed parameter.""" """Test __repr__ with seed parameter."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step]) pipeline = RobotProcessor([step], seed=42)
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='RobotProcessor', steps=1: [MockStep])" expected = "RobotProcessor(name='RobotProcessor', steps=1: [MockStep], seed=42)"
assert repr_str == expected assert repr_str == expected
@@ -1204,17 +1645,19 @@ def test_repr_with_custom_name_and_seed():
"""Test __repr__ with both custom name and seed.""" """Test __repr__ with both custom name and seed."""
step1 = MockStep("step1") step1 = MockStep("step1")
step2 = MockStepWithoutOptionalMethods() step2 = MockStepWithoutOptionalMethods()
pipeline = RobotProcessor([step1, step2], name="MyProcessor") pipeline = RobotProcessor([step1, step2], name="MyProcessor", seed=123)
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='MyProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods])" expected = (
"RobotProcessor(name='MyProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods], seed=123)"
)
assert repr_str == expected assert repr_str == expected
def test_repr_without_seed(): def test_repr_without_seed():
"""Test __repr__ when seed is explicitly None (should not show seed).""" """Test __repr__ when seed is explicitly None (should not show seed)."""
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step], name="TestProcessor") pipeline = RobotProcessor([step], name="TestProcessor", seed=None)
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='TestProcessor', steps=1: [MockStep])" expected = "RobotProcessor(name='TestProcessor', steps=1: [MockStep])"
@@ -1242,10 +1685,10 @@ def test_repr_edge_case_long_names():
step3 = MockStepWithTensorState() step3 = MockStepWithTensorState()
step4 = MockNonModuleStepWithState() step4 = MockNonModuleStepWithState()
pipeline = RobotProcessor([step1, step2, step3, step4], name="LongNames") pipeline = RobotProcessor([step1, step2, step3, step4], name="LongNames", seed=999)
repr_str = repr(pipeline) repr_str = repr(pipeline)
expected = "RobotProcessor(name='LongNames', steps=4: [MockStepWithNonSerializableParam, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])" expected = "RobotProcessor(name='LongNames', steps=4: [MockStepWithNonSerializableParam, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState], seed=999)"
assert repr_str == expected assert repr_str == expected