Compare commits

..

201 Commits

Author SHA1 Message Date
Jade Choghari aa517b5780 precommit 2025-09-17 11:20:16 +02:00
Jade Choghari 2c17433f4d make it work 2025-09-17 11:16:42 +02:00
Jade Choghari d31283cc5d refactor with pipeline 2025-09-17 09:49:19 +02:00
Steven Palma d8feb22f93 test(processor): fix isinstance and cuda test 2025-09-16 18:31:41 +02:00
Steven Palma 9073d64050 Merge branch 'main' into user/azouitine/2025-7-4-convert-codebase-with-pipeline 2025-09-16 18:05:45 +02:00
Steven Palma d6ce7bd330 chore(processor): update comments in record.py 2025-09-16 18:04:44 +02:00
Steven Palma 43eb0e375f fix(processor): enforce signatures 2025-09-16 17:13:07 +02:00
AdilZouitine fa8be1c4fe test(processor): update tests to handle missing or invalid task keys
- Modified tests to assert that the processor raises appropriate exceptions when the task key is missing or has an invalid value in the complementary data.
- Ensured that the tests cover cases for None, integer, and mixed list task values, improving robustness against invalid inputs.
2025-09-16 16:52:10 +02:00
Adil Zouitine a7d1179aab fix(processor): Preserve stats overrides in normalizer load_state_dict and fix training resumption (#1958)
* feat(processor): enhance normalization handling and state management

- Added support for additional normalization modes including IDENTITY.
- Introduced a new function `clean_state_dict` to remove specific substrings from state dict keys.
- Implemented preservation of explicitly provided normalization statistics during state loading.
- Updated training script to conditionally provide dataset statistics based on resume state.
- Expanded tests to verify the correct behavior of stats override preservation and loading.

* fix(train): remove redundant comment regarding state loading

- Removed a comment that noted the preprocessor and postprocessor state is already loaded when resuming training, as it was deemed unnecessary for clarity.
2025-09-16 16:45:13 +02:00
Steven Palma 772da63a8e test(async): fix feature manipulation (#1957)
* test(async): fix feature manipulation

* chore(processor): remove unused functions
2025-09-16 15:49:32 +02:00
Steven Palma 27a229ea64 chore(examples): homogenize style across example files (#1955)
* chore(examples): homogenize style across example files

* chore(examples): homogenize style across example files eval + replay

* chore(examples): homogenize headers
2025-09-16 14:56:36 +02:00
Jade Choghari 0d1e57f032 improve annotation and info dict
Signed-off-by: Jade Choghari <chogharijade@gmail.com>
2025-09-16 14:21:11 +02:00
Pepijn e2eff72ec0 feat(ee): add so100_to_so100_EE replay and evaluate examples 2025-09-16 13:56:44 +02:00
Jade Choghari d0d9036304 fix test failing
Signed-off-by: Jade Choghari <chogharijade@gmail.com>
2025-09-16 13:51:05 +02:00
Steven Palma 5d79869934 chore(processor): remove unused transition_features dict 2025-09-16 13:13:09 +02:00
Jade Choghari 7403060ad6 add cmake
Signed-off-by: Jade Choghari <chogharijade@gmail.com>
2025-09-16 13:12:31 +02:00
Jade Choghari 89aaaf1556 add cmake
Signed-off-by: Jade Choghari <chogharijade@gmail.com>
2025-09-16 13:12:04 +02:00
Jade Choghari c7e976d812 Change max_parallel_tasks to 1
Reduce the maximum number of parallel tasks from 5 to 1.

Signed-off-by: Jade Choghari <chogharijade@gmail.com>
2025-09-16 12:59:43 +02:00
Jade Choghari 1a4a47e804 Change max_parallel_tasks from 5 to 1
Signed-off-by: Jade Choghari <chogharijade@gmail.com>
2025-09-16 12:58:56 +02:00
Jade Choghari 4fe5c3ab70 Add libero (#1950)
* add libero

* backup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add factory

* Add LIBERO as a submodule

* add changes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add multitask

* remove photos

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* bug remove

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix video paths and train.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix renaming issues with cams

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update .gitignore

* final refactor/fix

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add safethread support

* ad

* Add dep (#4)

* Add 'libero' dependencies to pyproject.toml

* Add Git dependencies for egl_probe and LIBERO

* Update libero-requirements.txt

* add future dep

* update bash

* quick fix

* remove step1

* cleanup (#5)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* clean

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cleanup 2

* improve install

* Delete libero-requirements.txt

* iterate on review

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add docs for eval

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* doc

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update doc

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove brkpt

* fix docs

* update docs/script

* update doc

* skip test warning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* hotfix: flip actions

* add train

* new things

* More things

* add new changes

* iterate on review

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove files

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove unces

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* new changes

* iterate on review

* remove files

* factor

* update installation

* doc title

* make it reproducible

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* iterate on review

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove files

* update tests

* update doc

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove sh files

* add gym

---------

Signed-off-by: Jade Choghari <chogharijade@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jade Choghari (jchoghar) <chogharijade@gmai.com>
Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
2025-09-16 12:05:32 +02:00
Steven Palma cf7946e602 chore(processors): tokenizers raises and remove tensor conversion (#1949) 2025-09-16 11:44:02 +02:00
AdilZouitine b12a386334 docs(debug): enhance debugging guide for processor pipelines
- Streamlined the introduction to clarify the challenges of debugging complex processor pipelines.
- Expanded the section on hooks, detailing their purpose and implementation for runtime monitoring.
- Introduced step-by-step debugging techniques, emphasizing the use of the `step_through()` method for inspecting intermediate states.
- Added examples of feature validation to ensure data structure contracts are met.
- Consolidated best practices for debugging, highlighting the synergy between hooks, step-through debugging, and feature validation.
2025-09-16 10:26:05 +02:00
AdilZouitine cee5a3fec5 docs(processor): enhance tutorial on implementing custom processors
- Updated the tutorial to use `NormalizerProcessorStep` as the primary example, clarifying its role in normalizing observations and actions.
- Improved explanations of the need for custom processors, emphasizing data compatibility and processing requirements.
- Added code snippets demonstrating the normalization process and the configuration of processor pipelines.
- Enhanced the introduction to processors, detailing their function as translators between raw robot data and model inputs.
- Included examples of real-world processor configurations for both training and inference scenarios.
2025-09-16 09:13:05 +02:00
Steven Palma 8fb18109ef Merge branch 'main' into user/azouitine/2025-7-4-convert-codebase-with-pipeline 2025-09-15 23:22:31 +02:00
Steven Palma 170d8be7c2 feat(example): Add SO100 EE pipeline control (teleop+record) (#1943)
* feat(examples): add ee so100 processors teleop & record

* refactor(processor): improve FK processor for better use compatability
2025-09-15 23:21:47 +02:00
Steven Palma 8063cd5ed3 test(processor): fix batch expectation 2025-09-15 22:22:17 +02:00
Steven Palma 03891f66da chore(processor): update input output of main 3 processors for better semantics (#1942)
* chore(processor): update input output of main 3 processors for better semantics

* refactor(processor): replace Any with RobotObservation for improved type safety in processors

* fix(processors): no PolicyObservation

* chore(processor): update with RobotObservation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: AdilZouitine <adilzouitinegm@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2025-09-15 20:16:43 +02:00
Pepijn 99213daa3e fix cmd record, eval 2025-09-15 17:35:12 +02:00
Steven Palma 42cffd6f2e Merge pull request #1941 from huggingface/chore/merge_main_to_pipeline
chore: merge main to pipeline
2025-09-15 15:28:37 +02:00
Steven Palma 7b1b37b696 Merge branch 'main' into chore/merge_main_to_pipeline 2025-09-15 15:17:24 +02:00
Pepijn 4382742681 fix teleop, record and eval (#1940) 2025-09-15 14:39:05 +02:00
AdilZouitine e8d79b5191 refactor(docs): streamline monitoring hooks and enhance performance reporting
- Removed the log_shapes and measure_performance hooks, simplifying the monitoring process to focus on NaN checks.
- Updated performance reporting to include maximum processing times alongside average times for better insights.
- Clarified documentation regarding the processing pipeline and feature transformations.
2025-09-15 14:01:04 +02:00
Adil Zouitine 066308ceb8 refactor(processor): replace ModelHubMixin with HubMixin and enhance save_pretrained method (#1937)
- Updated DataProcessorPipeline to use HubMixin instead of ModelHubMixin for improved functionality.
- Refactored save_pretrained method to handle saving
2025-09-15 13:13:35 +02:00
Steven Palma 40e9ddd1ed fix(processors): make sure nested dict are also shallow copied (#1939) 2025-09-15 13:10:10 +02:00
Steven Palma c69f23723e refactor(processor): transform_features loop + EAFP (#1932) 2025-09-14 16:07:32 +02:00
Steven Palma 50293bb17b refactor(processors): several additions (#1926)
* chore(processor): remove merge_transitions functions (#1925)

* refactor(processors): move processors out of configs (#1927)

* chore(processor): streamline combine_features_dict (#1928)

* chore(policies): use new constants (#1929)

* fix(deps): right version transformers (#1930)

* fix(tests): add none + disable async tests for now (#1931)
2025-09-13 23:53:20 +02:00
Pepijn 839ac5f2aa fix(processor): phone examples (#1921)
* fix(processor): phone examples

* chore(processor): simplify gripper in phone example kinematic chain

---------

Co-authored-by: Steven Palma <steven.palma@huggingface.co>
2025-09-13 21:25:19 +02:00
AdilZouitine 0479eb8f69 docs: Add new section for debugging processor pipelines
- Introduced a new documentation entry for debugging processor pipelines, enhancing the existing guide on processors.
- This addition aims to provide users with insights and best practices for troubleshooting and optimizing their processor workflows.
2025-09-12 18:09:23 +02:00
Adil Zouitine a877c596ba chore(docs): Processor doc (#1685)
* chore(docs): initialize doc

* Added script for the second part of the processor doc

* precommit style nit

* improved part 2 of processor guide

* Add comprehensive documentation for processors in robotics

- Introduced a detailed guide on processors, covering their role in transforming raw robot data into model-ready inputs and vice versa.
- Explained core concepts such as EnvTransition, ProcessorStep, and RobotProcessor, along with their functionalities.
- Included examples of common processor steps like normalization, device management, batch processing, and text tokenization.
- Provided insights on building complete pipelines, integrating processors into training loops, and saving/loading configurations.
- Emphasized best practices and advanced features for effective usage of processors in robotics applications.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat(docs): Enhance introduction to processors with additional converter functions

- Updated the introduction to processors documentation to include default batch-to-transition and transition-to-batch converters.
- Added detailed descriptions and examples for new specialized converter functions: `to_transition_teleop_action`, `to_transition_robot_observation`, `to_output_robot_action`, and `to_dataset_frame`.
- Improved clarity on how these converters facilitate integration with existing robotics applications.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Improved doc implement_your_own_pipeline
- Use normalization processor as default example
- Add section on transform features
- Add section on overrides.

* Add phone docs and use pipeline for robots/teleop docs

* Fix typo in documentation for adapters in robots/teleop section

* Enhance documentation for processors with detailed explanations and examples

- Updated the introduction to processors, clarifying the role of `EnvTransition` and `ProcessorStep`.
- Introduced `DataProcessorPipeline` as a generic orchestrator for chaining processor steps.
- Added comprehensive descriptions of new converter functions and their applications.
- Improved clarity on type safety and the differences between `RobotProcessorPipeline` and `PolicyProcessorPipeline`.
- Included examples for various processing scenarios, emphasizing best practices for data handling in robotics.

* Enhance documentation for processor migration and debugging

- Added detailed sections on the migration of models to the new `PolicyProcessorPipeline` system, including breaking changes and migration scripts.
- Introduced a comprehensive guide for debugging processor pipelines, covering common issues, step-by-step inspection, and runtime monitoring techniques.
- Updated examples to reflect new usage patterns and best practices for processor implementation and error handling.
- Clarified the role of various processor steps and their configurations in the context of robotics applications.

---------

Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Pepijn <pepijn@huggingface.co>
2025-09-12 18:00:37 +02:00
Steven Palma 1ccdf365d2 docs(processor): update docstrings pipeline (#1920) 2025-09-12 17:54:27 +02:00
Pepijn 6bdcd460e0 use observation instead of obs 2025-09-12 12:18:00 +02:00
Pepijn 2005a28a00 add empty obs and act in create_initial_features 2025-09-12 12:17:46 +02:00
Pepijn 58b91dc886 fixes for rotation matrix 2025-09-12 11:57:59 +02:00
Pepijn 8e0f5cd052 fixes for processors used in phone teleop 2025-09-12 11:57:48 +02:00
AdilZouitine f51272362c refactor(processor): update migration script for policy normalization and hub integration
- Modified the migration script to include a branch argument for pushing to the hub, enhancing flexibility in version control.
- Improved error handling by ensuring the policy type is extracted from the configuration, promoting robustness.
- Streamlined the process of saving and pushing model components to the hub, allowing for a single commit with optional PR creation.
- Updated the commit message and description for better clarity on the migration changes and benefits, ensuring users are informed of the new architecture and usage.
2025-09-11 21:06:25 +02:00
Steven Palma cd0098a5f7 debug(scripts): simplify record with processors (#1918)
Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
2025-09-11 18:20:39 +02:00
AdilZouitine efde42d4a9 refactor(processor): migrate policy normalization to use factory functions
- Updated the migration script to utilize `make_pre_post_processors` and `make_policy_config` from `lerobot.policies.factory`, enhancing consistency with the current codebase.
- Improved normalization statistics extraction and processor pipeline creation, ensuring compatibility with the new `PolicyProcessorPipeline` architecture.
- Cleaned up configuration handling by removing unnecessary fields and adding normalization mapping directly to the config.
- Enhanced type safety and readability by refining feature type and normalization mode handling.
2025-09-11 18:14:13 +02:00
AdilZouitine aeb70812c1 refactor(processor): unify action imports and enhance type clarity across multiple files
- Updated imports in various files to include RobotAction and PolicyAction directly from the processor module, improving clarity and consistency.
- Removed redundant imports from core, streamlining the codebase and enhancing maintainability.
- Adjusted type annotations and references in the RobotProcessorPipeline and related components to align with the new import structure, ensuring better type safety and readability.
2025-09-11 14:24:36 +02:00
Adil Zouitine 376a6457cf feat(processor): enhance type safety with generic DataProcessorPipeline for policy and robot pipelines (#1915)
* refactor(processor): enhance type annotations for processors in record, replay, teleoperate, and control utils

- Updated type annotations for preprocessor and postprocessor parameters in record_loop and predict_action functions to specify the expected dictionary types.
- Adjusted robot_action_processor type in ReplayConfig and TeleoperateConfig to improve clarity and maintainability.
- Ensured consistency in type definitions across multiple files, enhancing overall code readability.

* refactor(processor): enhance type annotations for RobotProcessorPipeline in various files

- Updated type annotations for RobotProcessorPipeline instances in evaluate.py, record.py, replay.py, teleoperate.py, and other related files to specify input and output types more clearly.
- Introduced new type conversions for PolicyAction and EnvTransition to improve type safety and maintainability across the processing pipelines.
- Ensured consistency in type definitions, enhancing overall code readability and reducing potential runtime errors.

* refactor(processor): update transition handling in processors to use transition_to_batch

- Replaced direct transition handling with transition_to_batch in various processor tests and implementations to ensure consistent batching of input data.
- Updated assertions in tests to reflect changes in data structure, enhancing clarity and maintainability.
- Improved overall code readability by standardizing the way transitions are processed across different processor types.

* refactor(tests): standardize transition key usage in processor tests

- Updated assertions in processor test files to utilize the TransitionKey for action references, enhancing consistency across tests.
- Replaced direct string references with TransitionKey constants for improved readability and maintainability.
- Ensured that all relevant tests reflect these changes, contributing to a more uniform approach in handling transitions.
2025-09-11 13:36:04 +02:00
Steven Palma a2489ab0da fix(robots): remove action prefix hard-coded in teleop keyboard and gamepad 2025-09-11 12:37:09 +02:00
Steven Palma 014486999e fix(processor): use subprocessors in AddBatchDimensionProcessorStep only if we have the ingredients 2025-09-10 23:52:58 +02:00
Steven Palma cda44e5a52 refactor(processor): phone processor is now an RobotActionProcessorStep 2025-09-10 23:16:49 +02:00
Adil Zouitine 9183083e75 refactor(processor): clarify action types, distinguish PolicyAction, RobotAction, and EnvAction (#1908)
* refactor(processor): split action from policy, robots and environment

- Updated function names to robot_action_to_transition and robot_transition_to_action across multiple files to better reflect their purpose in processing robot actions.
- Adjusted references in the RobotProcessorPipeline and related components to ensure compatibility with the new naming convention.
- Enhanced type annotations for action parameters to improve code readability and maintainability.

* refactor(converters): rename robot_transition_to_action to transition_to_robot_action

- Updated function names across multiple files to improve clarity and consistency in processing robot actions.
- Adjusted references in RobotProcessorPipeline and related components to align with the new naming convention.
- Simplified action handling in the AddBatchDimensionProcessorStep by removing unnecessary checks for action presence.

* refactor(converters): update references to transition_to_robot_action

- Renamed all instances of robot_transition_to_action to transition_to_robot_action across multiple files for consistency and clarity in the processing of robot actions.
- Adjusted the RobotProcessorPipeline configurations to reflect the new naming convention, enhancing code readability.

* refactor(processor): update Torch2NumpyActionProcessorStep to extend ActionProcessorStep

- Changed the base class of Torch2NumpyActionProcessorStep from PolicyActionProcessorStep to ActionProcessorStep, aligning it with the current architecture of action processing.
- This modification enhances the clarity of the class's role in the processing pipeline.

* fix(processor): main action processor can take also EnvAction

---------

Co-authored-by: Steven Palma <steven.palma@huggingface.co>
2025-09-10 22:40:37 +02:00
Steven Palma 6745958362 fix(processor): specialized processors respect contract by raising if none (#1909)
* fix(processor): specialized processor now raise

* test(processor): fix tests for now raise specialized processors

* test(processor): use identity in newly introduced pipeline
2025-09-10 18:45:47 +02:00
Steven Palma 51588f741b test(processor): all processors use now the same create_transition (#1906)
* test(processor): all processors use now the same create_transition

* test(processor): use identity instead of lambda for transition in pipelines
2025-09-10 18:39:06 +02:00
Steven Palma df4292f6ed chore(processor): remove action prefixes (#1905) 2025-09-10 18:32:08 +02:00
Adil Zouitine 7e30090e97 refactor(eval): specify type parameters for preprocessor and postprocessor in eval_policy function (#1904) 2025-09-10 10:31:05 +02:00
Steven Palma e881fb6678 refactor(pipeline): feature contract now categorizes between OBS or Action (#1867)
* refactor(processor): signature of transform_features

* refactor(processor): remove prefixes + processor respect new transform_features signature + update test accordingly

* refactor(processor): rename now is only for visual

* refactor(processor): update normalize processor

* refactor(processor): update vanilla processor features

* refactor(processor): feature contract now uses its own enum

* chore(processor): rename renameprocessor

* chore(processor): minor changes

* refactor(processor): add create & change aggregate

* refactor(processor): update aggregate

* refactor(processor): simplify to functions, fix features contracts and rename function

* test(processor): remove to converter tests as now they are very simple

* chore(docs): recover docs joint observations processor

* fix(processor): update RKP

* fix(tests): recv diff test_pipeline

* chore(tests): add docs to test

* chore(processor): leave obs language constant untouched

* fix(processor): correct new shape of feature in crop image processor
2025-09-09 18:27:30 +02:00
Adil Zouitine acf0ba7fb3 refactor(converters): rename _from_tensor to from_tensor_to_numpy for clarity (#1902)
- Updated the function name from _from_tensor to from_tensor_to_numpy to better reflect its purpose of converting PyTorch tensors to numpy arrays or scalars.
- Adjusted all references to the renamed function throughout the codebase to maintain consistency.
- Enhanced the _NormalizationMixin class to reconstruct the stats dictionary from tensor stats using the new function, ensuring compatibility after loading state dicts.
- Added tests to verify the correct reconstruction of stats and functionality of methods dependent on self.stats after loading.
2025-09-09 17:51:47 +02:00
Adil Zouitine a74b90edd1 refactor(eval): integrate preprocessor and postprocessor into rollout and eval_policy functions (#1900)
* refactor(eval): integrate preprocessor and postprocessor into rollout and eval_policy functions

- Updated the `rollout` and `eval_policy` functions to accept preprocessor and postprocessor parameters, enhancing the flexibility of the evaluation pipeline.
- Adjusted the implementation to apply preprocessing and postprocessing steps during policy evaluation, improving the overall data handling and processing flow.

* refactor(eval): remove redundant observation device conversion in rollout function

- Eliminated unnecessary device conversion for the observation dictionary within the `rollout` function, streamlining the code and enhancing readability.
- This change simplifies the observation handling process, aligning with the preference for clearer solutions.

* debug

* refactor(utils): enhance task handling in add_envs_task function

- Improved the `add_envs_task` function to validate the output of `task_description` and `task` calls, ensuring they return lists of strings.
- Removed the use of `else` statement for environments without language instructions, simplifying the logic and enhancing readability.
- Streamlined the observation dictionary handling by ensuring consistent data types for task attributes.
2025-09-09 17:00:34 +02:00
Steven Palma 846677f9cc Merge branch 'main' into user/azouitine/2025-7-4-convert-codebase-with-pipeline 2025-09-08 22:35:13 +02:00
Steven Palma af9ddcf9a2 chore(docs): update doctrines pipeline files (#1872)
* docs(processor): update docstrings batch_processor

* docs(processor): update docstrings device_processor

* docs(processor): update docstrings tokenizer_processor

* update docstrings processor_act

* update docstrings for pipeline_features

* update docstrings for utils

* update docstring for processor_diffusion

* update docstrings factory

* add docstrings to pi0 processor

* add docstring to pi0fast processor

* add docstring classifier processor

* add docstring to sac processor

* add docstring smolvla processor

* add docstring to tdmpc processor

* add docstring to vqbet processor

* add docstrings to converters

* add docstrings for delta_action_processor

* add docstring to gym action processor

* update hil processor

* add docstring to joint obs processor

* add docstring to migrate_normalize_processor

* update docstrings normalize processor

* update docstring normalize processor

* update docstrings observation processor

* update docstrings rename_processor

* add docstrings robot_kinematic_processor

* cleanup rl comments

* add docstring to train.py

* add docstring to teleoperate.py

* add docstrings to phone_processor.py

* add docstrings to teleop_phone.py

* add docstrings to control_utils.py

* add docstrings to visualization_utils.py

---------

Co-authored-by: Pepijn <pepijn@huggingface.co>
2025-09-08 18:44:15 +02:00
Adil Zouitine d32006440c refactor(processors): Improve Normalization Processor Performance and Device/Dtype Adaptability (#1880)
* refactor(processors): reorder processor steps for consistency across implementations

- Updated the order of processor steps in multiple files to ensure consistency, placing AddBatchDimensionProcessorStep and DeviceProcessorStep before NormalizerProcessorStep.
- Adjusted related test assertions to reflect the new order of steps in the preprocessor, enhancing clarity and maintainability.

* refactor(normalization): remove dtype specification in tensor conversion for adaptation logic

- Updated tensor conversion in the _NormalizationMixin class to remove explicit dtype specification, allowing for automatic adaptation of tensor types.
- Adjusted related tests to ensure proper functionality with the new tensor conversion logic, verifying that normalizers adapt correctly to input types.
2025-09-08 10:46:35 +02:00
Steven Palma f1cfdfced9 fix(processor): recover type inference for use of processors (#1873) 2025-09-05 11:31:30 +02:00
Adil Zouitine 888a5b6249 refactor(utils): simplify log_rerun_data function (#1864)
* refactor(logging): enhance log_rerun_data to handle observation and action separately

- Updated the `log_rerun_data` function to accept and log observation and action data more clearly, improving readability and maintainability.
- Refactored the `record_loop` and `teleop_loop` functions to extract and pass observation and action data to `log_rerun_data`, ensuring consistent logging format.

* refactor(tests): update test_log_rerun_data to align with log_rerun_data changes

- Modified test cases in `test_visualization_utils.py` to extract and pass observation and action data separately to `log_rerun_data`, improving clarity and consistency with recent function updates.
- Ensured that the tests reflect the new structure of `log_rerun_data` for better maintainability.

* refactor(processors): simplify calls to log_rerun + replace lambda functions with identity_transition

---------

Co-authored-by: Steven Palma <steven.palma@huggingface.co>
2025-09-04 19:25:51 +02:00
Adil Zouitine f247aa0701 refactor(tests): update processor test assertions to reflect new preprocessor and postprocessor names (#1869)
- Changed assertions in multiple processor test files to verify the updated names from "robot_preprocessor" and "robot_postprocessor" to "policy_preprocessor" and "policy_postprocessor" for consistency with recent refactoring.
2025-09-04 17:34:06 +02:00
Adil Zouitine 1ac6a6d3fe refactor(constants): rename preprocessor and postprocessor constants for clarity (#1868)
- Updated constant names from PREPROCESSOR_DEFAULT_NAME and POSTPROCESSOR_DEFAULT_NAME to POLICY_PREPROCESSOR_DEFAULT_NAME and POLICY_POSTPROCESSOR_DEFAULT_NAME for better context.
- Adjusted references across multiple files to use the new constant names, ensuring consistency in the codebase.
2025-09-04 17:01:53 +02:00
Steven Palma e698c709d8 fix(deps): use in-house rotation utils over scipy throughout the codebase 2025-09-04 16:44:18 +02:00
Adil Zouitine a988da4789 feat(teleoperation): introduce HasTeleopEvents protocol and enhance teleop event handling (#1866)
- Added the HasTeleopEvents protocol to define a standard for teleoperators that provide control events.
- Implemented a runtime check to ensure teleoperators implement the get_teleop_events() method.
- Updated AddTeleopEventsAsInfoStep to utilize the new protocol, enhancing compatibility with custom teleoperators.
- Improved documentation for clarity on teleoperation event extraction and compatibility with built-in teleoperators.
2025-09-04 16:28:49 +02:00
Adil Zouitine 99963b6968 refactor(dependencies): remove scipy dependency and introduce custom rotation utilities (#1863)
- Removed the scipy dependency from the project to streamline requirements.
- Added a new `rotation.py` module containing a custom `Rotation` class that replicates essential functionalities of `scipy.spatial.transform.Rotation`, allowing for rotation vector, matrix, and quaternion conversions without external dependencies.
- Updated the `robot_kinematic_processor.py` to utilize the new custom rotation utilities.
2025-09-04 16:26:28 +02:00
Adil Zouitine 332ca4ccc5 refactor(pipeline): enforce ProcessorStep inheritance for pipeline steps (#1862)
- Updated the DataProcessorPipeline to require that all steps inherit from ProcessorStep, enhancing type safety and clarity.
- Adjusted tests to utilize a MockTokenizerProcessorStep that adheres to the ProcessorStep interface, ensuring consistent behavior across tests.
- Refactored various mock step classes in tests to inherit from ProcessorStep for improved consistency and maintainability.
2025-09-04 16:22:03 +02:00
Adil Zouitine fc43246942 feat(record): add transition features to dataset and handle scalar vs array formatting in converters (#1861)
- Introduced new transition features (`next.reward`, `next.done`, `next.truncated`) in the dataset during recording.
- Updated the `transition_to_dataset_frame` function to handle scalar values correctly, ensuring compatibility with expected array formats for reward, done, and truncated features.
2025-09-04 16:17:31 +02:00
Adil Zouitine 793ad86fc9 refactor(processor): enforce config_filename requirement for HF Hub loading (#1860)
- Updated the DataProcessorPipeline to require a specific config_filename when loading from Hugging Face Hub, enhancing clarity and preventing errors.
- Simplified local path checks and improved error handling for invalid paths.
- Adjusted tests to reflect the new requirement and ensure proper error handling for various loading scenarios.
2025-09-04 10:31:18 +02:00
Adil Zouitine a6dbb65917 chore(processor): add type alias RobotProcessorPipeline and PolicyProcessorPipeline (#1859)
* feat(processor): introduce PolicyProcessorPipeline and RobotProcessorPipeline as type aliases for DataProcessorPipeline

- Added PolicyProcessorPipeline and RobotProcessorPipeline type aliases to enhance clarity and maintainability in the processor module.
- Updated the __all__ list to include the new pipelines for better module export consistency.

* refactor(processor): replace DataProcessorPipeline with PolicyProcessorPipeline across multiple modules

- Updated all instances of DataProcessorPipeline to PolicyProcessorPipeline in various processor files for consistency and clarity.
- Adjusted function signatures to reflect the new pipeline type, enhancing maintainability and readability.

* refactor(processor): update hotswap_stats function to use PolicyProcessorPipeline

- Changed the parameter name from robot_processor to policy_processor for clarity.
- Ensured consistency with recent updates to the processor module by reflecting the new pipeline type in the function signature.

* refactor(processor): replace DataProcessorPipeline with PolicyProcessorPipeline in migrate_policy_normalization.py

- Updated the preprocessor and postprocessor to use PolicyProcessorPipeline for consistency with recent changes in the processor module.
- Enhanced clarity and maintainability by aligning with the new pipeline structure.

* refactor(processor): update hotswap_stats to use PolicyProcessorPipeline

- Changed the parameter type in hotswap_stats from DataProcessorPipeline to PolicyProcessorPipeline for consistency with recent updates.
- Enhanced clarity by updating the function documentation to reflect the new pipeline type.

* refactor(processor): replace DataProcessorPipeline with RobotProcessorPipeline across multiple files

- Updated instances of DataProcessorPipeline to RobotProcessorPipeline in evaluate.py, record.py, replay.py, teleoperate.py, and other relevant files for consistency and clarity.
- Adjusted function signatures and variable types to reflect the new pipeline structure, enhancing maintainability and readability.
2025-09-03 19:01:28 +02:00
Steven Palma 6c7169c4af chore(processor): rename teleop_phone variable names (#1858) 2025-09-03 18:42:13 +02:00
Adil Zouitine f125d5e3bf refactor(processor): rename internal device variable for clarity (#1857)
- Changed the internal device variable from `_device` to `tensor_device` for improved readability and consistency.
- Updated references throughout the class to reflect the new variable name.
2025-09-03 18:39:06 +02:00
Steven Palma 75dcfd4886 chore(processor): rename merge_features -> combine_feature_dicts (#1856) 2025-09-03 18:20:35 +02:00
Adil Zouitine ff3cbaa872 refactor(processor): rename internal tokenizer variable for clarity (#1855)
- Changed the internal tokenizer variable name from `_tokenizer` to `input_tokenizer` for improved readability and consistency.
- Updated references throughout the class to reflect the new variable name.
2025-09-03 18:20:12 +02:00
Adil Zouitine ce793cde64 chore(processor): add Step suffix to all processors (#1854)
* refactor(processor): rename MapDeltaActionToRobotAction and MapTensorToDeltaActionDict for consistency

* refactor(processor): rename DeviceProcessor to DeviceProcessorStep for consistency across modules

* refactor(processor): rename Torch2NumpyActionProcessor to Torch2NumpyActionProcessorStep for consistency

* refactor(processor): rename Numpy2TorchActionProcessor to Numpy2TorchActionProcessorStep for consistency

* refactor(processor): rename AddTeleopActionAsComplimentaryData to AddTeleopActionAsComplimentaryDataStep for consistency

* refactor(processor): rename ImageCropResizeProcessor and AddTeleopEventsAsInfo for consistency

* refactor(processor): rename TimeLimitProcessor to TimeLimitProcessorStep for consistency

* refactor(processor): rename GripperPenaltyProcessor to GripperPenaltyProcessorStep for consistency

* refactor(processor): rename InterventionActionProcessor to InterventionActionProcessorStep for consistency

* refactor(processor): rename RewardClassifierProcessor to RewardClassifierProcessorStep for consistency

* refactor(processor): rename JointVelocityProcessor to JointVelocityProcessorStep for consistency

* refactor(processor): rename MotorCurrentProcessor to MotorCurrentProcessorStep for consistency

* refactor(processor): rename NormalizerProcessor and UnnormalizerProcessor to NormalizerProcessorStep and UnnormalizerProcessorStep for consistency

* refactor(processor): rename VanillaObservationProcessor to VanillaObservationProcessorStep for consistency

* refactor(processor): rename RenameProcessor to RenameProcessorStep for consistency

* refactor(processor): rename TokenizerProcessor to TokenizerProcessorStep for consistency

* refactor(processor): rename ToBatchProcessor to AddBatchDimensionProcessorStep for consistency

* refactor(processor): update config file name in test for RenameProcessorStep consistency
2025-09-03 18:12:11 +02:00
Steven Palma 029c4a9a76 chore(processor): rename converters function names (#1853)
* chore(processor): rename to_transition_teleop_action -> action_to_transition

* chore(processor): rename to_transition_robot_observation -> observation_to_transition

* chore(processor): rename to_output_robot_action -> transition_to_robot_action
2025-09-03 18:08:54 +02:00
Steven Palma d893bf1e30 chore(processor): rename specialized processor -> XYZProcessorStep (#1852) 2025-09-03 17:30:47 +02:00
Steven Palma 8c796b39f5 chore(processor): rename RobotProcessor -> DataProcessorPipeline (#1850) 2025-09-03 17:13:16 +02:00
Adil Zouitine 4ebe482a7e refactor(processors): enhance transform_features method across multiple processors (#1849)
* refactor(processors): enhance transform_features method across multiple processors

- Updated the transform_features method in various processors to utilize a copy of the features dictionary, ensuring immutability of the original features.
- Added handling for new feature keys and removed obsolete ones in the MapTensorToDeltaActionDict, JointVelocityProcessor, and others.
- Improved readability and maintainability by following consistent patterns in feature transformation.

* refactor(processors): standardize action and observation keys in delta_action_processor and joint_observations_processor

- Updated action and observation keys to use constants for improved readability and maintainability.
- Refactored the transform_features method in multiple processors to ensure consistent handling of feature keys.
- Enhanced error handling by raising exceptions for missing required components in action and observation processing.
- Removed obsolete code and improved overall structure for better clarity.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor(processors): remove unused import in joint_observations_processor

* refactor(processors): simplify transform_features method in delta_action_processor

* refactor(processors): streamline transform_features method in ImageCropResizeProcessor

* refactor(processors): improve error handling and streamline transform_features method in phone_processor

- Raised a ValueError for missing position and rotation in action to enhance error handling.

* refactor(processors): enhance error handling in JointVelocityProcessor

- Added a ValueError raise for missing current joint positions in the observation method to improve error handling and ensure the integrity of the transform_features method.

* refactor(processors): simplify transform_features method in robot kinematic processors

* refactor(processors): standardize action keys in phone_processor

* fix(processor): RKP feature obs -> act

---------

Signed-off-by: Adil Zouitine <adilzouitinegm@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Steven Palma <steven.palma@huggingface.co>
2025-09-03 16:54:41 +02:00
Steven Palma 2fcc358e98 refactor(processors): add extended api for specialized pipelines (#1848) 2025-09-03 12:28:40 +02:00
Steven Palma b052843f08 refactor(processors): unify import statements by consolidating pipeline imports into the main processor module (#1845) 2025-09-02 18:26:59 +02:00
Steven Palma ebb464c255 refactor(processors): update transition handling in RewardClassifierProcessor and InverseKinematicsEEToJoints (#1844) 2025-09-02 17:57:49 +02:00
Steven Palma 2914ae2a96 refactor(processors): add transform_features method to various processors (#1843) 2025-09-02 17:15:01 +02:00
Adil Zouitine 645c87e3a9 refactor(converters): gather converters and refactor the logic (#1833)
* refactor(converters): move batch transition functions to converters module

- Moved `_default_batch_to_transition` and `_default_transition_to_batch` functions from `pipeline.py` to `converters.py` for better organization and separation of concerns.
- Updated references in `RobotProcessor` to use the new location of these functions.
- Added tests to ensure correct functionality of the transition functions, including handling of index and task_index fields.
- Removed redundant tests from `pipeline.py` to streamline the test suite.

* refactor(processor): reorganize EnvTransition and TransitionKey definitions

- Moved `EnvTransition` and `TransitionKey` classes from `pipeline.py` to a new `core.py` module for better structure and maintainability.
- Updated import statements across relevant modules to reflect the new location of these definitions, ensuring consistent access throughout the codebase.

* refactor(converters): rename and update dataset frame conversion functions

- Replaced `to_dataset_frame` with `transition_to_dataset_frame` for clarity and consistency in naming.
- Updated references in `record.py`, `pipeline.py`, and tests to use the new function name.
- Introduced `merge_transitions` to streamline the merging of transitions, enhancing readability and maintainability.
- Adjusted related tests to ensure correct functionality with the new naming conventions.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(processor): solve conflict artefacts

* refactor(converters): remove unused identity function and update type hints for merge_transitions

* refactor(processor): remove unused identity import and clean up gym_manipulator.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Steven Palma <steven.palma@huggingface.co>
2025-09-02 15:33:38 +02:00
Steven Palma 2c802ac134 refactor(converters): implement unified tensor conversion function (#1841)
- Introduced `to_tensor` function using `singledispatch` to handle various input types, including scalars, arrays, and dictionaries, converting them to PyTorch tensors.
- Replaced previous tensor conversion logic in `gym_action_processor`, `normalize_processor`, and `test_converters` with the new `to_tensor` function for improved readability and maintainability.
- Updated tests to cover new functionality and ensure correct tensor conversion behavior.

Co-authored-by: AdilZouitine <adilzouitinegm@gmail.com>
2025-09-02 13:47:04 +02:00
Steven Palma 15ffc01fb3 Revert "refactor(converters): implement unified tensor conversion function (#…" (#1840)
This reverts commit a837685bf8.
2025-09-02 13:43:35 +02:00
Adil Zouitine a837685bf8 refactor(converters): implement unified tensor conversion function (#1830)
- Introduced `to_tensor` function using `singledispatch` to handle various input types, including scalars, arrays, and dictionaries, converting them to PyTorch tensors.
- Replaced previous tensor conversion logic in `gym_action_processor`, `normalize_processor`, and `test_converters` with the new `to_tensor` function for improved readability and maintainability.
- Updated tests to cover new functionality and ensure correct tensor conversion behavior.
2025-09-02 13:28:26 +02:00
Adil Zouitine d32b76cc66 refactor(processor): improve processor pipeline typing with generic type (#1810)
* refactor(processor): introduce generic type for to_output

- Always return `TOutput`
- Remove `_prepare_transition`, so `__call__` now always returns `TOutput`
- Update tests accordingly
- This refactor paves the way for adding settings for `to_transition` and `to_output` in `make_processor` and the post-processor

* refactor(processor): consolidate ProcessorKwargs usage across policies

- Removed the ProcessorTypes module and integrated ProcessorKwargs directly into the processor pipeline.
- Updated multiple policy files to utilize the new ProcessorKwargs structure for preprocessor and postprocessor arguments.
- Simplified the handling of processor kwargs by initializing them to empty dictionaries when not provided.
2025-09-02 12:57:14 +02:00
Adil Zouitine 08fb310eaa refactor(constants, processor): standardize action and observation keys across multiple files (#1808)
- Added new constants for truncated and done states in constants.py.
- Updated references to action and observation keys in pipeline_features.py, converters.py, hil_processor.py, tokenizer_processor.py, and robot_kinematic_processor.py to use the new constants for improved readability and maintainability.
2025-08-31 22:53:13 +02:00
Steven Palma 574a708950 Merge branch 'main' into user/azouitine/2025-7-4-convert-codebase-with-pipeline 2025-08-31 20:46:59 +02:00
Steven Palma ce665160ae feat(processor): multiple improvements to the pipeline porting (#1749)
* [Port codebase pipeline] General fixes for RL and scripts (#1748)

* Refactor dataset configuration in documentation and codebase

- Updated dataset configuration keys from `dataset_root` to `root` and `num_episodes` to `num_episodes_to_record` for consistency.
- Adjusted replay episode handling by renaming `episode` to `replay_episode`.
- Enhanced documentation
- added specific processor to transform from policy actions to delta actions

* Added Robot action to tensor processor
Added new processor script for dealing with gym specific action processing

* removed RobotAction2Tensor processor; imrpoved choosing observations in actor

* nit in delta action

* added missing reset functions to kinematics

* Adapt teleoperate and replay to pipeline similar to record

* refactor(processors): move to inheritance (#1750)

* fix(teleoperator): improvements phone implementation (#1752)

* fix(teleoperator): protect shared state in phone implementation

* refactor(teleop): separate classes in phone

* fix: solve breaking changes (#1753)

* refactor(policies): multiple improvements (#1754)

* refactor(processor): simpler logic in device processor (#1755)

* refactor(processor): euclidean distance in delta action processor (#1757)

* refactor(processor): improvements to joint observations processor migration (#1758)

* refactor(processor): improvements to tokenizer migration (#1759)

* refactor(processor): improvements to tokenizer migration

* fix(tests): tokenizer tests regression from #1750

* fix(processors): fix float comparison and config in hil processors (#1760)

* chore(teleop): remove unnecessary callbacks in KeyboardEndEffectorTeleop (#1761)

* refactor(processor): improvements normalize pipeline migration (#1756)

* refactor(processor): several improvements normalize processor step

* refactor(processor): more improvements normalize processor

* refactor(processor): more changes to normalizer

* refactor(processor): take a different approach to DRY

* refactor(processor): final design

* chore(record): revert comment and continue deleted (#1764)

* refactor(examples): pipeline phone examples (#1769)

* refactor(examples): phone teleop + teleop script

* refactor(examples): phone replay + replay

* chore(examples): rename phone example files & folders

* feat(processor): fix improvements to the pipeline porting (#1796)

* refactor(processor): enhance tensor device handling in normalization process (#1795)

* refactor(tests): remove unsupported device detection test for complementary data (#1797)

* chore(tests): update ToBatchProcessor test (#1798)

* refactor(tests): remove in-place mutation tests for actions and complementary data in batch processor

* test(tests): add tests for action and task processing in batch processor

* add names for android and ios phone (#1799)

* use _tensor_stats in normalize processor (#1800)

* fix(normalize_processor): correct device reference for tensor epsilon handling (#1801)

* add point 5 add missing feature contracts (#1806)

* Fix PR comments 1452 (#1807)

* use key to determine image

* Address rest of PR comments

* use PolicyFeatures in transform_features

---------

Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>

---------

Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
2025-08-31 20:38:52 +02:00
AdilZouitine 35c5d43255 chore(processor): Add default names for preprocessor and postprocessor in constants
- Introduced `PREPROCESSOR_DEFAULT_NAME` and `POSTPROCESSOR_DEFAULT_NAME` constants for consistent naming across various processor implementations.
- Updated processor creation in multiple policy files to utilize these constants, enhancing code readability and maintainability.
- Modified the training script to load and save the preprocessor and postprocessor using the new constants.
2025-08-11 18:00:25 +02:00
Steven Palma 95c1e32aa5 Merge branch 'main' into user/azouitine/2025-7-4-convert-codebase-with-pipeline 2025-08-11 13:56:03 +02:00
Michel Aractingi e4db65a127 Remove HILEnvConfig references 2025-08-11 11:14:57 +02:00
Michel Aractingi 0053defa2e Refactorgym_manipulator.py using the universal pipeline (#1650)
* Migrate gym_manipulator to use the pipeline
Added get_teleop_events function to capture relevant events from teleop devices unrelated to actions

* Added the capability to record a dataset

* Added the replay functionality with the pipeline

* Refactored `actor.py` to use the pipeline

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* RL works at this commit - fixed actor.py and bugs in gym_manipulator

* change folder structure to reduce the size of gym_manip

* Refactored hilserl config

* Remove dataset and mode from HilSerlEnvConfig to a GymManipulatorConfig to reduce verbose of configs during training

* format docs

* removed get_teleop_events from abc

* Refactor environment configuration and processing pipeline for GymHIL support. Removed device attribute from HILSerlRobotEnvConfig, added DummyTeleopDevice for simulation, and updated processor creation to accommodate GymHIL environments.

* Improved typing for HILRobotEnv config and GymManipulator config

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Migrated `gym_manipulator` to use a more modular structure similar to phone teleop

* Refactor gripper handling and transition processing in HIL and robot kinematic processors

- Updated gripper position handling to use a consistent key format across processors
- Improved the EEReferenceAndDelta class to handle reference joint positions.
- Added support for discrete gripper actions in the GripperVelocityToJoint processor.
- Refactored the gym manipulator to improve modularity and clarity in processing steps.

* Added delta_action_processor mapping wrapper

* Added missing file delta_action_processor and improved imports in `gym_manipulator`

* nit

* Added missing file joint_observation_processor

* Enhance processing architecture with new teleoperation processors

- Introduced `AddTeleopActionAsComplimentaryData` and `AddTeleopEventsAsInfo` for integrating teleoperator actions and events into transitions.
- Added `Torch2NumpyActionProcessor` and `Numpy2TorchActionProcessor` for seamless conversion between PyTorch tensors and NumPy arrays.
- Updated `__init__.py` to include new processors in module exports, improving modularity and clarity in the processing pipeline.
- GymHIL is now fully supported with HIL using the pipeline

* Refactor configuration structure for gym_hil integration

- Renamed sections for better readability, such as changing "Gym Wrappers Configuration" to "Processor Configuration."
- Enhanced documentation with clear examples for dataset collection and policy evaluation configurations.

* Enhance reset configuration and teleoperation event handling

- Added `terminate_on_success` parameter to `ResetConfig` and `InterventionActionProcessor` for controlling episode termination behavior upon success detection.
- Updated documentation to clarify the impact of `terminate_on_success` on data collection for reward classifier training.
- Refactored teleoperation event handling to use `TeleopEvents` constants for improved readability and maintainability across various modules.

* fix(keyboard teleop), delta action keys

* Added transform features and feature contract

* Added transform features for image crop

* Enum for TeleopEvents

* Update tranform_features delta action proc

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2025-08-11 11:07:55 +02:00
AdilZouitine fd5d8b3d5f refactor(train): Remove unnecessary tensor device handling in training loop 2025-08-08 19:35:15 +02:00
AdilZouitine 5bf82f8229 feat(tests): Add comprehensive tests for various policy processors
- Introduced new test files for ACT, Classifier, Diffusion, PI0, SAC, SmolVLA, TDMPC, and VQBeT policy processors.
- Each test file includes unit tests to validate functionality, including handling of batch sizes, device management, and data type conversions.
- Enhanced test coverage to ensure robustness and reliability of processor implementations across different scenarios.
2025-08-08 19:34:50 +02:00
AdilZouitine 5ca3920611 feat(DeviceProcessor): Enhance tensor processing with device detection and float dtype conversion
- Improved the _process_tensor method to preserve GPU placement for tensors already on a GPU, facilitating multi-GPU training scenarios.
- Introduced a new _detect_device method in TokenizerProcessor to ensure tokenized tensors match the device of existing tensors in transitions.
- Added comprehensive unit tests to validate the functionality of device detection and float dtype conversion across various scenarios.
2025-08-08 19:33:24 +02:00
AdilZouitine 8bde9d0ab7 refactor(factory): streamline processor loading by removing unused comments
- Removed commented-out code related to loading pretrained processors in the make_processor function.
- This change enhances code clarity and maintains focus on the current implementation.
2025-08-08 13:23:26 +02:00
AdilZouitine abcbc16126 refactor(normalization): remove Normalize and Unnormalize classes
- Deleted the Normalize and Unnormalize classes from the normalization module to streamline the codebase.
- Updated tests to ensure compatibility with the removal of these classes, focusing on the new NormalizerProcessor and UnnormalizerProcessor implementations.
- Enhanced the handling of normalization statistics and improved overall code clarity.
2025-08-08 13:23:10 +02:00
AdilZouitine e4fd30a8d4 feat(policies): convert save_policy_to_safetensors with pipeline 2025-08-08 13:21:50 +02:00
Adil Zouitine 5f759b1637 feat(dependencies): Add scipy as a required dependency
- Included `scipy>=1.15.2` in the project dependencies to enhance functionality and support for scientific computing tasks.
2025-08-07 18:09:49 +02:00
Adil Zouitine 6a75b4761a refactor(TokenizerProcessor): improve dependency handling and observation management
- Updated TokenizerProcessor to conditionally import AutoTokenizer based on the availability of the transformers library, enhancing flexibility.
- Modified tokenizer attribute type to Any to accommodate scenarios where transformers may not be installed.
- Improved observation handling by using a more concise approach to manage the transition dictionary, ensuring compatibility with existing data structures.
- Added error handling for missing transformers library, providing clear guidance for users on installation requirements.
2025-08-07 17:07:20 +02:00
Pepijn e5ade5565d Integrate pipeline and add phone teleop (#1681)
* 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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* 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.

* chore (docs): add docstring for processor

* fix (test): test factory

* fix(test): policies

* 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>

* chore(test): add suggestion made by copilot regarding numpy test

* fix(test): import issue

* 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.

* chore (docstrin):Improve docstring for NormalizerProcessor

* feat (device processor): Implement device processor

* chore (batch handling): Enhance processing components with batch conversion utilities

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(test): linting issue

* chore (output format): improves output format

* chore (type): add typing for multiprocess envs

* 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.

* chore(normalization): addressing comments from copilot

* chore(learner): nit comment from copilot

* feat(pipeline): Enhance step_through method to support both tuple and dict inputs

* refactor(pipeline): Simplify observation and padding data handling in batch transitions

* Apply suggestions from code review

Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Signed-off-by: Adil Zouitine <adilzouitinegm@gmail.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor(pipeline): Introduce ComplementaryDataProcessor for handling complementary data in transitions

* fix(ci): temporary fix on dataset deps version

* feat(processors): Introduce processors for various policy types

- Added `make_processor` function to create processor instances for different policy types, including `tdmpc`, `diffusion`, `act`, `vqbet`, `pi0`, `pi0fast`, `sac`, and `reward_classifier`.
- Implemented corresponding processor files for each policy type, encapsulating normalization and unnormalization steps.
- Updated existing policies to remove direct normalization dependencies, enhancing modularity and clarity.
- Enhanced test coverage to validate the integration of new processors with existing policy configurations.

* refactor(learner): Remove normalization from cached image features retrieval

- Simplified the retrieval of observation features by removing the normalization step from the `get_cached_image_features` method calls.
- This change enhances clarity and aligns with the recent updates to policy processors.

* refactor(policies): Remove unnormalization step from action predictions

- Eliminated the unnormalization of actions in both `TDMPCPolicy` and `VQBeTPolicy` classes to streamline action prediction.
- This change improves code clarity and aligns with recent updates to policy processors.

* feat(train): Integrate preprocessor into training pipeline

* refactor(train): Update preprocessor initialization to include dataset statistics

* refactor(policies): Enhance processor creation and add NaN detection hook

* refactor(train): Update memory pinning logic for mps compatibility

* feat: initial commit phone teleop

* ugly delta control

* use quaternion

* 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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* 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.

* Remove redundant tests for None observation and serialization methods in `test_observation_processor.py` to streamline the test suite and improve maintainability.

* 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.

* 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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* 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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* 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.

* chore (docs): add docstring for processor

* fix (test): test factory

* fix(test): policies

* 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>

* chore(test): add suggestion made by copilot regarding numpy test

* fix(test): import issue

* 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.

* chore (docstrin):Improve docstring for NormalizerProcessor

* feat (device processor): Implement device processor

* chore (batch handling): Enhance processing components with batch conversion utilities

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(test): linting issue

* chore (output format): improves output format

* chore (type): add typing for multiprocess envs

* 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.

* chore(normalization): addressing comments from copilot

* chore(learner): nit comment from copilot

* feat(pipeline): Enhance step_through method to support both tuple and dict inputs

* refactor(pipeline): Simplify observation and padding data handling in batch transitions

* Apply suggestions from code review

Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Signed-off-by: Adil Zouitine <adilzouitinegm@gmail.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor(pipeline): Introduce ComplementaryDataProcessor for handling complementary data in transitions

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* 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.

* 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.

* 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.

* 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.

* 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.

* chore(pipeline): Move _CFG_NAME along other class member

* 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.

* 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.

* chore(doc): address pip install commant lerobot that not exist yet

* 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.

* 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.

* docs(pipeline): Add clarification for repo name sanitization process

* feat(processors): Introduce processors for various policy types

- Added `make_processor` function to create processor instances for different policy types, including `tdmpc`, `diffusion`, `act`, `vqbet`, `pi0`, `pi0fast`, `sac`, and `reward_classifier`.
- Implemented corresponding processor files for each policy type, encapsulating normalization and unnormalization steps.
- Updated existing policies to remove direct normalization dependencies, enhancing modularity and clarity.
- Enhanced test coverage to validate the integration of new processors with existing policy configurations.

* refactor(learner): Remove normalization from cached image features retrieval

- Simplified the retrieval of observation features by removing the normalization step from the `get_cached_image_features` method calls.
- This change enhances clarity and aligns with the recent updates to policy processors.

* refactor(policies): Remove unnormalization step from action predictions

- Eliminated the unnormalization of actions in both `TDMPCPolicy` and `VQBeTPolicy` classes to streamline action prediction.
- This change improves code clarity and aligns with recent updates to policy processors.

* feat(train): Integrate preprocessor into training pipeline

* refactor(train): Update preprocessor initialization to include dataset statistics

* refactor(policies): Enhance processor creation and add NaN detection hook

* feat(record): Integrate RobotProcessor into recording loop and update policy handling

- Added support for RobotProcessor in the record_loop function to enhance data processing capabilities.
- Updated the logic to reset both policy and processor when provided, ensuring proper state management.
- Modified action prediction to utilize the processor, improving the overall functionality of the recording process.
- Adjusted the save_checkpoint function to include preprocessor state saving, enhancing checkpointing capabilities.

* feat(migration): Add script for migrating policy models with normalization layers

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat(migrate): Enhance migration script to create preprocessor and postprocessor for policy models

- Updated the migration script to generate both a preprocessor and a postprocessor, improving the handling of normalization for training and inference.
- Added functionality to convert features to PolicyFeature objects, ensuring compatibility with the new processor architecture.
- Refined the extraction and removal of normalization statistics and layers, streamlining the migration process.
- Improved error handling for missing mandatory configuration fields during model instantiation.

* feat(migrate): Add model card generation and saving to migration script

- Implemented functionality to generate and save a model card for the migrated model, including metadata such as dataset repository ID, license, and tags.
- Enhanced the script to push the model card to the hub if requested, improving model documentation and accessibility.
- Refactored the saving process to ensure the model card is saved locally and uploaded correctly when pushing to the hub.

* feat(processor): Introduce ToBatchProcessor for handling observation batching

- Added ToBatchProcessor to ensure observations have proper batch dimensions for model processing.
- Implemented functionality to add batch dimensions to state and image observations as needed.
- Created comprehensive unit tests to validate the processor's behavior with various tensor dimensions and types.
- Ensured compatibility with existing transition keys and maintained the integrity of non-observation data.

* feat(processors): Add ToBatchProcessor to multiple policy processors

- Integrated ToBatchProcessor into various policy processors to handle observation batching.
- Updated make functions for act, diffusion, pi0, pi0fast, sac, smolvla, tdmpc, and vqbet processors to include the new batching functionality.
- Ensured consistency across all processor implementations for improved data handling.

* refactor(factory): Remove unused imports and NaN detection hook from processor creation

* feat(batch_processor): Enhance ToBatchProcessor to handle action batching

- Updated ToBatchProcessor to add batch dimensions to actions in addition to observations.
- Implemented separate methods for processing observations and actions, improving code readability.
- Added comprehensive unit tests to validate action batching functionality across various tensor dimensions and types.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat(factory): Enhance make_processor to support preprocessor and postprocessor configuration

- Introduced ProcessorConfigKwargs TypedDict for better type safety in processor configuration.
- Updated make_processor to accept preprocessor and postprocessor configuration filenames, improving flexibility in processor instantiation.
- Refactored the loading of pretrained processors to utilize the new configuration options.

* refactor(factory): Clean up imports in factory.py

- Removed unused import of IdentityProcessor to streamline the code.

* feat(migrate): Extend load_model_from_hub to include train configuration

- Updated load_model_from_hub to return the train configuration alongside the model state_dict and config.
- Modified main function to handle the additional train configuration when loading models from both the hub and local paths.
- Adjusted dataset_repo_id extraction to utilize the train configuration for improved accuracy.

* refactor(record): Rename processor parameters and update processing logic

- Renamed `processor` to `preprocessor` and added `postprocessor` parameter for clarity.
- Updated the `record_loop` and `predict_action` functions to utilize the new preprocessor and postprocessor, enhancing the processing flow.
- Ensured compatibility with existing functionality while improving code readability.

* feat(batch_processor): Add task field processing to ToBatchProcessor

- Enhanced ToBatchProcessor to wrap string tasks in a list, adding batch dimensions for compatibility with model inference.
- Implemented a new method for processing complementary data, ensuring that task values are correctly handled as either strings or lists of strings.
- Added comprehensive unit tests to validate task processing, including edge cases and in-place mutation of complementary data.

* feat(normalization): Implement IDENTITY mode for normalization and unnormalization

- Enhanced NormalizerProcessor and UnnormalizerProcessor to support IDENTITY mode, allowing features to bypass normalization when specified.
- Updated processing logic to check normalization modes and handle missing statistics gracefully.
- Added comprehensive unit tests to validate IDENTITY mode functionality for both observations and actions, ensuring correct behavior across various scenarios.
- Improved error handling for unsupported normalization modes.

* fix(rebase): remove residual normalization layer:

* refactor(diffusion): remove normalization layer from input processing

* Add debug + calib

* cleanup

* Add pipeline

* fix int

* Add record example

* nit

* Add feature contract to pipelinestep and pipeline

* Add tests

* Add processor tests

* PR feedback

* encorperate pr feedback

* type in doc

* oops

* cleaned up steps and integrated pipeline with feature_contract

* refactor steps and robot to pipeline

* cleanup pipeline

* cleanup code further

* make it run

* feat(processors): Introduce processors for various policy types

- Added `make_processor` function to create processor instances for different policy types, including `tdmpc`, `diffusion`, `act`, `vqbet`, `pi0`, `pi0fast`, `sac`, and `reward_classifier`.
- Implemented corresponding processor files for each policy type, encapsulating normalization and unnormalization steps.
- Updated existing policies to remove direct normalization dependencies, enhancing modularity and clarity.
- Enhanced test coverage to validate the integration of new processors with existing policy configurations.

* refactor(learner): Remove normalization from cached image features retrieval

- Simplified the retrieval of observation features by removing the normalization step from the `get_cached_image_features` method calls.
- This change enhances clarity and aligns with the recent updates to policy processors.

* refactor(policies): Remove unnormalization step from action predictions

- Eliminated the unnormalization of actions in both `TDMPCPolicy` and `VQBeTPolicy` classes to streamline action prediction.
- This change improves code clarity and aligns with recent updates to policy processors.

* feat(train): Integrate preprocessor into training pipeline

* refactor(train): Update preprocessor initialization to include dataset statistics

* refactor(policies): Enhance processor creation and add NaN detection hook

* feat(record): Integrate RobotProcessor into recording loop and update policy handling

- Added support for RobotProcessor in the record_loop function to enhance data processing capabilities.
- Updated the logic to reset both policy and processor when provided, ensuring proper state management.
- Modified action prediction to utilize the processor, improving the overall functionality of the recording process.
- Adjusted the save_checkpoint function to include preprocessor state saving, enhancing checkpointing capabilities.

* feat(migration): Add script for migrating policy models with normalization layers

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat(migrate): Enhance migration script to create preprocessor and postprocessor for policy models

- Updated the migration script to generate both a preprocessor and a postprocessor, improving the handling of normalization for training and inference.
- Added functionality to convert features to PolicyFeature objects, ensuring compatibility with the new processor architecture.
- Refined the extraction and removal of normalization statistics and layers, streamlining the migration process.
- Improved error handling for missing mandatory configuration fields during model instantiation.

* feat(migrate): Add model card generation and saving to migration script

- Implemented functionality to generate and save a model card for the migrated model, including metadata such as dataset repository ID, license, and tags.
- Enhanced the script to push the model card to the hub if requested, improving model documentation and accessibility.
- Refactored the saving process to ensure the model card is saved locally and uploaded correctly when pushing to the hub.

* feat(processor): Introduce ToBatchProcessor for handling observation batching

- Added ToBatchProcessor to ensure observations have proper batch dimensions for model processing.
- Implemented functionality to add batch dimensions to state and image observations as needed.
- Created comprehensive unit tests to validate the processor's behavior with various tensor dimensions and types.
- Ensured compatibility with existing transition keys and maintained the integrity of non-observation data.

* feat(processors): Add ToBatchProcessor to multiple policy processors

- Integrated ToBatchProcessor into various policy processors to handle observation batching.
- Updated make functions for act, diffusion, pi0, pi0fast, sac, smolvla, tdmpc, and vqbet processors to include the new batching functionality.
- Ensured consistency across all processor implementations for improved data handling.

* refactor(factory): Remove unused imports and NaN detection hook from processor creation

* feat(batch_processor): Enhance ToBatchProcessor to handle action batching

- Updated ToBatchProcessor to add batch dimensions to actions in addition to observations.
- Implemented separate methods for processing observations and actions, improving code readability.
- Added comprehensive unit tests to validate action batching functionality across various tensor dimensions and types.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat(factory): Enhance make_processor to support preprocessor and postprocessor configuration

- Introduced ProcessorConfigKwargs TypedDict for better type safety in processor configuration.
- Updated make_processor to accept preprocessor and postprocessor configuration filenames, improving flexibility in processor instantiation.
- Refactored the loading of pretrained processors to utilize the new configuration options.

* refactor(factory): Clean up imports in factory.py

- Removed unused import of IdentityProcessor to streamline the code.

* feat(migrate): Extend load_model_from_hub to include train configuration

- Updated load_model_from_hub to return the train configuration alongside the model state_dict and config.
- Modified main function to handle the additional train configuration when loading models from both the hub and local paths.
- Adjusted dataset_repo_id extraction to utilize the train configuration for improved accuracy.

* refactor(record): Rename processor parameters and update processing logic

- Renamed `processor` to `preprocessor` and added `postprocessor` parameter for clarity.
- Updated the `record_loop` and `predict_action` functions to utilize the new preprocessor and postprocessor, enhancing the processing flow.
- Ensured compatibility with existing functionality while improving code readability.

* feat(batch_processor): Add task field processing to ToBatchProcessor

- Enhanced ToBatchProcessor to wrap string tasks in a list, adding batch dimensions for compatibility with model inference.
- Implemented a new method for processing complementary data, ensuring that task values are correctly handled as either strings or lists of strings.
- Added comprehensive unit tests to validate task processing, including edge cases and in-place mutation of complementary data.

* feat(normalization): Implement IDENTITY mode for normalization and unnormalization

- Enhanced NormalizerProcessor and UnnormalizerProcessor to support IDENTITY mode, allowing features to bypass normalization when specified.
- Updated processing logic to check normalization modes and handle missing statistics gracefully.
- Added comprehensive unit tests to validate IDENTITY mode functionality for both observations and actions, ensuring correct behavior across various scenarios.
- Improved error handling for unsupported normalization modes.

* fix(rebase): remove residual normalization layer:

* refactor(diffusion): remove normalization layer from input processing

* refactor(normalization): Remove unused state dict transformation methods and streamline imports

- Eliminated the _transform_state_dict_keys and _load_as_safetensor methods from PI0Policy, simplifying the model loading process.
- Cleaned up imports in modeling_pi0.py by removing log_model_loading_keys and init_logging.
- Updated TDMPCPolicy and VQBeTPolicy to handle action removal from batches during offline evaluation.
- Introduced hotswap_stats function in normalize_processor.py to update normalization statistics dynamically, with corresponding tests to ensure functionality.

* refactor(normalization): Clean up imports in normalize_processor.py

* feat(batch_processor): Add feature_contract method to ToBatchProcessor

- Introduced feature_contract method that returns features without modification, maintaining the no-op behavior of the processor.
- This addition enhances the flexibility of the ToBatchProcessor for future feature processing needs.

* fix(dependencies): Update transformers dependency constraint to allow only versions up to 4.52.0

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat(tokenizer): Introduce TokenizerProcessor for text tokenization

- Added TokenizerProcessor class to handle tokenization of task strings using Hugging Face's AutoTokenizer.
- Supports both string and list inputs, with customizable parameters for task key, output key, and tokenization settings.
- Implemented comprehensive unit tests to validate functionality, including handling of various input scenarios and integration with RobotProcessor.
- Updated types.py to include LANGUAGE feature type and modified __init__.py to register the new processor.

* feat(language): Enhance language processing in TokenizerProcessor

- Added OBS_LANGUAGE constant to define the observation language key.
- Updated TokenizerProcessor to store tokenized task data in the observation dictionary, ensuring compatibility with the new language feature.
- Introduced Pi0NewLineProcessor to append newlines to tasks for proper tokenization.
- Modified tests to validate the integration of language tokens and attention masks in the observation structure.

* feat(tokenizer): Add padding configuration to TokenizerProcessor

- Introduced `padding_side` parameter to the TokenizerProcessor for customizable padding direction.
- Updated the `make_pi0_processor` function to include the new padding configuration.
- Enhanced unit tests to validate the functionality of the `padding_side` parameter in various scenarios.

* feat(processor): Add state management methods to Pi0NewLineProcessor

* feat(normalization): Track normalization and unnormalization info in complementary data

- Updated NormalizerProcessor and UnnormalizerProcessor to accept additional parameters for tracking normalization modes.
- Enhanced the __call__ methods to store normalization and unnormalization information in the complementary data of transitions.
- Added unit tests to verify the correct tracking of normalization info, including scenarios with missing stats and selective normalization keys.

* feat(factory): Add preprocessor and postprocessor overrides to ProcessorConfigKwargs

- Updated ProcessorConfigKwargs to include optional overrides for preprocessor and postprocessor configurations.
- Enhanced the make_processor function to utilize the new overrides, allowing for more flexible processor initialization.

* feat(processors): Integrate RenameProcessor into various processor configurations

- Added RenameProcessor to the input steps of multiple processor functions, including make_act_processor, make_diffusion_processor, make_pi0_processor, make_sac_processor, make_tdmpc_processor, make_vqbet_processor, and make_smolvla_processor.
- Consolidated normalization features from input and output into a single NormalizerProcessor for improved efficiency.
- Updated the input steps to ensure compatibility with the new RenameProcessor integration.

* Do some todos and cleanup

* change feature_contract to dataset_features

* use one method for conversion pipeline output to add_frame dict and use base processors where possible

* Add back in and use record_loop

* update todo

* rename to_dataset_frame

* feat(smolvla): Refactor language processing and introduce new line processor (#1658)

- Removed the prepare_language method and directly accessed language tokens and masks from the batch using the OBS_LANGUAGE constant.
- Added SmolVLANewLineProcessor to ensure tasks end with a newline, enhancing tokenization compatibility.
- Updated the make_smolvla_processor function to include the new line processor and tokenizer processor for improved input handling.

* feat(processors): Integrate DeviceProcessor into multiple processor configurations

- Added DeviceProcessor to the input and output steps of various processor functions, including make_act_processor, make_diffusion_processor, make_pi0_processor, make_pi0fast_processor, make_sac_processor, make_tdmpc_processor, make_vqbet_processor, and make_smolvla_processor.
- Enhanced the DeviceProcessor class with state management methods and ensured compatibility with existing processor pipelines.
- Introduced unit tests for DeviceProcessor to validate functionality across different scenarios, including CPU and CUDA operations.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix

* fix reference frame

* refactor(pipeline): Remove to() method for device management

- Eliminated the to() method from RobotProcessor, which was responsible for moving tensor states to specified devices.
- Removed associated unit tests that validated the functionality of the to() method across various scenarios.
- Streamlined the pipeline code by focusing on other device management strategies.

* feat(processor): Enhance DeviceProcessor with float dtype conversion

- Added support for optional float dtype conversion in DeviceProcessor, allowing tensors to be converted to specified floating-point types while preserving non-float types.
- Implemented validation for float dtype input and updated the processor's configuration methods to include float dtype.
- Refactored tensor processing logic to streamline device movement and dtype conversion.
- Introduced comprehensive unit tests to validate the new float dtype functionality across various scenarios.

* update data visualization

* update teleop example

* fix record bugs

* Add replay

* Not code

* feature(pipeline): port tokenizer pipeline for VLA (#1645)

* feat(tokenizer): Introduce TokenizerProcessor for text tokenization

- Added TokenizerProcessor class to handle tokenization of task strings using Hugging Face's AutoTokenizer.
- Supports both string and list inputs, with customizable parameters for task key, output key, and tokenization settings.
- Implemented comprehensive unit tests to validate functionality, including handling of various input scenarios and integration with RobotProcessor.
- Updated types.py to include LANGUAGE feature type and modified __init__.py to register the new processor.

* feat(language): Enhance language processing in TokenizerProcessor

- Added OBS_LANGUAGE constant to define the observation language key.
- Updated TokenizerProcessor to store tokenized task data in the observation dictionary, ensuring compatibility with the new language feature.
- Introduced Pi0NewLineProcessor to append newlines to tasks for proper tokenization.
- Modified tests to validate the integration of language tokens and attention masks in the observation structure.

* feat(tokenizer): Add padding configuration to TokenizerProcessor

- Introduced `padding_side` parameter to the TokenizerProcessor for customizable padding direction.
- Updated the `make_pi0_processor` function to include the new padding configuration.
- Enhanced unit tests to validate the functionality of the `padding_side` parameter in various scenarios.

* feat(processor): Add state management methods to Pi0NewLineProcessor

* feat(normalization): Track normalization and unnormalization info in complementary data

- Updated NormalizerProcessor and UnnormalizerProcessor to accept additional parameters for tracking normalization modes.
- Enhanced the __call__ methods to store normalization and unnormalization information in the complementary data of transitions.
- Added unit tests to verify the correct tracking of normalization info, including scenarios with missing stats and selective normalization keys.

* feat(factory): Add preprocessor and postprocessor overrides to ProcessorConfigKwargs

- Updated ProcessorConfigKwargs to include optional overrides for preprocessor and postprocessor configurations.
- Enhanced the make_processor function to utilize the new overrides, allowing for more flexible processor initialization.

* feat(processors): Integrate RenameProcessor into various processor configurations

- Added RenameProcessor to the input steps of multiple processor functions, including make_act_processor, make_diffusion_processor, make_pi0_processor, make_sac_processor, make_tdmpc_processor, make_vqbet_processor, and make_smolvla_processor.
- Consolidated normalization features from input and output into a single NormalizerProcessor for improved efficiency.
- Updated the input steps to ensure compatibility with the new RenameProcessor integration.

* feat(smolvla): Refactor language processing and introduce new line processor (#1658)

- Removed the prepare_language method and directly accessed language tokens and masks from the batch using the OBS_LANGUAGE constant.
- Added SmolVLANewLineProcessor to ensure tasks end with a newline, enhancing tokenization compatibility.
- Updated the make_smolvla_processor function to include the new line processor and tokenizer processor for improved input handling.

* feture(policies): add device processor (#1659)

* feat(processors): Integrate DeviceProcessor into multiple processor configurations

- Added DeviceProcessor to the input and output steps of various processor functions, including make_act_processor, make_diffusion_processor, make_pi0_processor, make_pi0fast_processor, make_sac_processor, make_tdmpc_processor, make_vqbet_processor, and make_smolvla_processor.
- Enhanced the DeviceProcessor class with state management methods and ensured compatibility with existing processor pipelines.
- Introduced unit tests for DeviceProcessor to validate functionality across different scenarios, including CPU and CUDA operations.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor(pipeline): Remove to() method for device management

- Eliminated the to() method from RobotProcessor, which was responsible for moving tensor states to specified devices.
- Removed associated unit tests that validated the functionality of the to() method across various scenarios.
- Streamlined the pipeline code by focusing on other device management strategies.

* feat(processor): Enhance DeviceProcessor with float dtype conversion

- Added support for optional float dtype conversion in DeviceProcessor, allowing tensors to be converted to specified floating-point types while preserving non-float types.
- Implemented validation for float dtype input and updated the processor's configuration methods to include float dtype.
- Refactored tensor processing logic to streamline device movement and dtype conversion.
- Introduced comprehensive unit tests to validate the new float dtype functionality across various scenarios.

* feat(policies): Add new line processors and update module exports

* feat(processor): Enhance batch and device processors to handle index and task_index fields

- Added logic to ToBatchProcessor for unsqueezing 0D tensors for index and task_index fields, ensuring they are processed as 1D tensors.
- Updated DeviceProcessor to process index and task_index fields in complementary data, preserving their tensor types and ensuring non-tensor fields remain unchanged.
- Enhanced unit tests to validate the correct handling of index and task_index fields across various scenarios, including device compatibility and dtype preservation.

* Add eval script

* fix `q_curr` in InverseKinematicsEEToJoints to the IK solution

* feat(processors): Introduce processors for various policy types

- Added `make_processor` function to create processor instances for different policy types, including `tdmpc`, `diffusion`, `act`, `vqbet`, `pi0`, `pi0fast`, `sac`, and `reward_classifier`.
- Implemented corresponding processor files for each policy type, encapsulating normalization and unnormalization steps.
- Updated existing policies to remove direct normalization dependencies, enhancing modularity and clarity.
- Enhanced test coverage to validate the integration of new processors with existing policy configurations.

* refactor(learner): Remove normalization from cached image features retrieval

- Simplified the retrieval of observation features by removing the normalization step from the `get_cached_image_features` method calls.
- This change enhances clarity and aligns with the recent updates to policy processors.

* refactor(policies): Remove unnormalization step from action predictions

- Eliminated the unnormalization of actions in both `TDMPCPolicy` and `VQBeTPolicy` classes to streamline action prediction.
- This change improves code clarity and aligns with recent updates to policy processors.

* feat(train): Integrate preprocessor into training pipeline

* refactor(train): Update preprocessor initialization to include dataset statistics

* refactor(policies): Enhance processor creation and add NaN detection hook

* feat(record): Integrate RobotProcessor into recording loop and update policy handling

- Added support for RobotProcessor in the record_loop function to enhance data processing capabilities.
- Updated the logic to reset both policy and processor when provided, ensuring proper state management.
- Modified action prediction to utilize the processor, improving the overall functionality of the recording process.
- Adjusted the save_checkpoint function to include preprocessor state saving, enhancing checkpointing capabilities.

* feat(migration): Add script for migrating policy models with normalization layers

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat(migrate): Enhance migration script to create preprocessor and postprocessor for policy models

- Updated the migration script to generate both a preprocessor and a postprocessor, improving the handling of normalization for training and inference.
- Added functionality to convert features to PolicyFeature objects, ensuring compatibility with the new processor architecture.
- Refined the extraction and removal of normalization statistics and layers, streamlining the migration process.
- Improved error handling for missing mandatory configuration fields during model instantiation.

* feat(migrate): Add model card generation and saving to migration script

- Implemented functionality to generate and save a model card for the migrated model, including metadata such as dataset repository ID, license, and tags.
- Enhanced the script to push the model card to the hub if requested, improving model documentation and accessibility.
- Refactored the saving process to ensure the model card is saved locally and uploaded correctly when pushing to the hub.

* feat(processor): Introduce ToBatchProcessor for handling observation batching

- Added ToBatchProcessor to ensure observations have proper batch dimensions for model processing.
- Implemented functionality to add batch dimensions to state and image observations as needed.
- Created comprehensive unit tests to validate the processor's behavior with various tensor dimensions and types.
- Ensured compatibility with existing transition keys and maintained the integrity of non-observation data.

* feat(processors): Add ToBatchProcessor to multiple policy processors

- Integrated ToBatchProcessor into various policy processors to handle observation batching.
- Updated make functions for act, diffusion, pi0, pi0fast, sac, smolvla, tdmpc, and vqbet processors to include the new batching functionality.
- Ensured consistency across all processor implementations for improved data handling.

* refactor(factory): Remove unused imports and NaN detection hook from processor creation

* feat(batch_processor): Enhance ToBatchProcessor to handle action batching

- Updated ToBatchProcessor to add batch dimensions to actions in addition to observations.
- Implemented separate methods for processing observations and actions, improving code readability.
- Added comprehensive unit tests to validate action batching functionality across various tensor dimensions and types.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat(factory): Enhance make_processor to support preprocessor and postprocessor configuration

- Introduced ProcessorConfigKwargs TypedDict for better type safety in processor configuration.
- Updated make_processor to accept preprocessor and postprocessor configuration filenames, improving flexibility in processor instantiation.
- Refactored the loading of pretrained processors to utilize the new configuration options.

* refactor(factory): Clean up imports in factory.py

- Removed unused import of IdentityProcessor to streamline the code.

* feat(migrate): Extend load_model_from_hub to include train configuration

- Updated load_model_from_hub to return the train configuration alongside the model state_dict and config.
- Modified main function to handle the additional train configuration when loading models from both the hub and local paths.
- Adjusted dataset_repo_id extraction to utilize the train configuration for improved accuracy.

* refactor(record): Rename processor parameters and update processing logic

- Renamed `processor` to `preprocessor` and added `postprocessor` parameter for clarity.
- Updated the `record_loop` and `predict_action` functions to utilize the new preprocessor and postprocessor, enhancing the processing flow.
- Ensured compatibility with existing functionality while improving code readability.

* feat(batch_processor): Add task field processing to ToBatchProcessor

- Enhanced ToBatchProcessor to wrap string tasks in a list, adding batch dimensions for compatibility with model inference.
- Implemented a new method for processing complementary data, ensuring that task values are correctly handled as either strings or lists of strings.
- Added comprehensive unit tests to validate task processing, including edge cases and in-place mutation of complementary data.

* feat(normalization): Implement IDENTITY mode for normalization and unnormalization

- Enhanced NormalizerProcessor and UnnormalizerProcessor to support IDENTITY mode, allowing features to bypass normalization when specified.
- Updated processing logic to check normalization modes and handle missing statistics gracefully.
- Added comprehensive unit tests to validate IDENTITY mode functionality for both observations and actions, ensuring correct behavior across various scenarios.
- Improved error handling for unsupported normalization modes.

* fix(rebase): remove residual normalization layer:

* refactor(diffusion): remove normalization layer from input processing

* refactor(normalization): Remove unused state dict transformation methods and streamline imports

- Eliminated the _transform_state_dict_keys and _load_as_safetensor methods from PI0Policy, simplifying the model loading process.
- Cleaned up imports in modeling_pi0.py by removing log_model_loading_keys and init_logging.
- Updated TDMPCPolicy and VQBeTPolicy to handle action removal from batches during offline evaluation.
- Introduced hotswap_stats function in normalize_processor.py to update normalization statistics dynamically, with corresponding tests to ensure functionality.

* refactor(normalization): Clean up imports in normalize_processor.py

* feat(batch_processor): Add feature_contract method to ToBatchProcessor

- Introduced feature_contract method that returns features without modification, maintaining the no-op behavior of the processor.
- This addition enhances the flexibility of the ToBatchProcessor for future feature processing needs.

* fix(dependencies): Update transformers dependency constraint to allow only versions up to 4.52.0

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feature(pipeline): port tokenizer pipeline for VLA (#1645)

* feat(tokenizer): Introduce TokenizerProcessor for text tokenization

- Added TokenizerProcessor class to handle tokenization of task strings using Hugging Face's AutoTokenizer.
- Supports both string and list inputs, with customizable parameters for task key, output key, and tokenization settings.
- Implemented comprehensive unit tests to validate functionality, including handling of various input scenarios and integration with RobotProcessor.
- Updated types.py to include LANGUAGE feature type and modified __init__.py to register the new processor.

* feat(language): Enhance language processing in TokenizerProcessor

- Added OBS_LANGUAGE constant to define the observation language key.
- Updated TokenizerProcessor to store tokenized task data in the observation dictionary, ensuring compatibility with the new language feature.
- Introduced Pi0NewLineProcessor to append newlines to tasks for proper tokenization.
- Modified tests to validate the integration of language tokens and attention masks in the observation structure.

* feat(tokenizer): Add padding configuration to TokenizerProcessor

- Introduced `padding_side` parameter to the TokenizerProcessor for customizable padding direction.
- Updated the `make_pi0_processor` function to include the new padding configuration.
- Enhanced unit tests to validate the functionality of the `padding_side` parameter in various scenarios.

* feat(processor): Add state management methods to Pi0NewLineProcessor

* feat(normalization): Track normalization and unnormalization info in complementary data

- Updated NormalizerProcessor and UnnormalizerProcessor to accept additional parameters for tracking normalization modes.
- Enhanced the __call__ methods to store normalization and unnormalization information in the complementary data of transitions.
- Added unit tests to verify the correct tracking of normalization info, including scenarios with missing stats and selective normalization keys.

* feat(factory): Add preprocessor and postprocessor overrides to ProcessorConfigKwargs

- Updated ProcessorConfigKwargs to include optional overrides for preprocessor and postprocessor configurations.
- Enhanced the make_processor function to utilize the new overrides, allowing for more flexible processor initialization.

* feat(processors): Integrate RenameProcessor into various processor configurations

- Added RenameProcessor to the input steps of multiple processor functions, including make_act_processor, make_diffusion_processor, make_pi0_processor, make_sac_processor, make_tdmpc_processor, make_vqbet_processor, and make_smolvla_processor.
- Consolidated normalization features from input and output into a single NormalizerProcessor for improved efficiency.
- Updated the input steps to ensure compatibility with the new RenameProcessor integration.

* feat(smolvla): Refactor language processing and introduce new line processor (#1658)

- Removed the prepare_language method and directly accessed language tokens and masks from the batch using the OBS_LANGUAGE constant.
- Added SmolVLANewLineProcessor to ensure tasks end with a newline, enhancing tokenization compatibility.
- Updated the make_smolvla_processor function to include the new line processor and tokenizer processor for improved input handling.

* feture(policies): add device processor (#1659)

* feat(processors): Integrate DeviceProcessor into multiple processor configurations

- Added DeviceProcessor to the input and output steps of various processor functions, including make_act_processor, make_diffusion_processor, make_pi0_processor, make_pi0fast_processor, make_sac_processor, make_tdmpc_processor, make_vqbet_processor, and make_smolvla_processor.
- Enhanced the DeviceProcessor class with state management methods and ensured compatibility with existing processor pipelines.
- Introduced unit tests for DeviceProcessor to validate functionality across different scenarios, including CPU and CUDA operations.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor(pipeline): Remove to() method for device management

- Eliminated the to() method from RobotProcessor, which was responsible for moving tensor states to specified devices.
- Removed associated unit tests that validated the functionality of the to() method across various scenarios.
- Streamlined the pipeline code by focusing on other device management strategies.

* feat(processor): Enhance DeviceProcessor with float dtype conversion

- Added support for optional float dtype conversion in DeviceProcessor, allowing tensors to be converted to specified floating-point types while preserving non-float types.
- Implemented validation for float dtype input and updated the processor's configuration methods to include float dtype.
- Refactored tensor processing logic to streamline device movement and dtype conversion.
- Introduced comprehensive unit tests to validate the new float dtype functionality across various scenarios.

* feat(policies): Add new line processors and update module exports

* feat(processor): Enhance batch and device processors to handle index and task_index fields

- Added logic to ToBatchProcessor for unsqueezing 0D tensors for index and task_index fields, ensuring they are processed as 1D tensors.
- Updated DeviceProcessor to process index and task_index fields in complementary data, preserving their tensor types and ensuring non-tensor fields remain unchanged.
- Enhanced unit tests to validate the correct handling of index and task_index fields across various scenarios, including device compatibility and dtype preservation.

* refactor(processors): Standardize processor naming conventions

- Updated processor names across various files to use a consistent "robot_preprocessor" and "robot_postprocessor" format.
- Modified the make_processor functions in factory, act, diffusion, pi0, pi0fast, sac, smolvla, tdmpc, and vqbet to reflect the new naming scheme.
- Enhanced the pipeline configuration to align with the updated processor names, improving clarity and maintainability.

* refactor(factory): Update processor configuration and type hints

- Changed return type of get_policy_class to type[PreTrainedPolicy] for improved type safety.
- Enhanced make_processor function to utilize dataset_stats in processor creation for better flexibility.
- Updated ProcessorConfigKwargs to include dataset_stats, allowing for more comprehensive processor configurations.
- Streamlined processor initialization by removing unnecessary kwargs and ensuring clarity in processor type handling.

* Fix eval and android gripper

* add some tests

* refactor(factory, pi0fast): Update processor function names and parameters

- Renamed make_pi0_processor to make_pi0fast_processor for clarity and consistency.
- Updated parameter names in the factory's make_processor function to use pretrained_model_name_or_path instead of source, enhancing readability and alignment with naming conventions.

* fix(train.py) push postprocessor with preprocessor
- Add preprocesser policy overrides for device and rename_map
- Add rename_map to DatasetRecordConfig (record.py)

* Cleanup pr

* fix more git diff pr issues

* add path as type in save_pretrained

* small nit

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* rename test file

* fix: make dataset_features/feature_contract is optional

* fix tests

* Encorperate pr feedback

* clean up record.py

* add ascii art, fix normal record

* remove merge issues

* fix merge

* remove features

* Add feedback PR

* fix last 4 tests

* remove features check

* rename to transform_features

* add transform_features

* fix lekiwi eval and update eval api example

---------

Signed-off-by: Adil Zouitine <adilzouitinegm@gmail.com>
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
2025-08-07 16:13:34 +02:00
Adil Zouitine 0524551f52 refactor(migrate_policy_normalization): Enhance preprocessor and postprocessor structure
- Introduced RenameProcessor in the preprocessor to handle renaming features.
- Combined input and output features in a single NormalizerProcessor for improved efficiency.
- Updated RobotProcessor initialization to clarify step naming for preprocessor and postprocessor.
- Added DeviceProcessor to both preprocessor and postprocessor for better device management.
2025-08-07 11:04:15 +02:00
Steven Palma 862bc7ef85 Merge branch 'main' into user/azouitine/2025-7-4-convert-codebase-with-pipeline 2025-08-06 21:08:32 +02:00
Adil Zouitine d38792d6e5 test(tokenizer_processor): Add require_package decorator for transformers
- Introduced @require_package("transformers") decorator in multiple test functions to ensure the transformers package is available before running tests.
- This change enhances test reliability by preventing failures due to missing dependencies.
2025-08-06 19:22:23 +02:00
pre-commit-ci[bot] db3cf0158c [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-06 16:08:39 +00:00
Adil Zouitine 0535f2a59a refactor(device_processor): Update device handling and improve type hints
- Changed device attribute type from torch.device to str for better clarity.
- Introduced a private _device attribute to store the actual torch.device instance.
- Updated tests to conditionally check for CUDA availability, ensuring compatibility across different environments.
- Refactored device-related assertions in tests to use a consistent approach for device type verification.
2025-08-06 18:08:15 +02:00
Michel Aractingi 2805ae347c fix(train.py) push postprocessor with preprocessor
- Add preprocesser policy overrides for device and rename_map
- Add rename_map to DatasetRecordConfig (record.py)
2025-08-06 17:21:17 +02:00
Adil Zouitine 28ef6fcd14 refactor(factory, pi0fast): Update processor function names and parameters
- Renamed make_pi0_processor to make_pi0fast_processor for clarity and consistency.
- Updated parameter names in the factory's make_processor function to use pretrained_model_name_or_path instead of source, enhancing readability and alignment with naming conventions.
2025-08-06 17:21:16 +02:00
Adil Zouitine 7fc7ec75bb refactor(factory): Update processor configuration and type hints
- Changed return type of get_policy_class to type[PreTrainedPolicy] for improved type safety.
- Enhanced make_processor function to utilize dataset_stats in processor creation for better flexibility.
- Updated ProcessorConfigKwargs to include dataset_stats, allowing for more comprehensive processor configurations.
- Streamlined processor initialization by removing unnecessary kwargs and ensuring clarity in processor type handling.
2025-08-06 17:21:15 +02:00
Adil Zouitine 87890cbf38 refactor(processors): Standardize processor naming conventions
- Updated processor names across various files to use a consistent "robot_preprocessor" and "robot_postprocessor" format.
- Modified the make_processor functions in factory, act, diffusion, pi0, pi0fast, sac, smolvla, tdmpc, and vqbet to reflect the new naming scheme.
- Enhanced the pipeline configuration to align with the updated processor names, improving clarity and maintainability.
2025-08-06 17:21:14 +02:00
Adil Zouitine 5326ffe77e feature(pipeline): port tokenizer pipeline for VLA (#1645)
* feat(tokenizer): Introduce TokenizerProcessor for text tokenization

- Added TokenizerProcessor class to handle tokenization of task strings using Hugging Face's AutoTokenizer.
- Supports both string and list inputs, with customizable parameters for task key, output key, and tokenization settings.
- Implemented comprehensive unit tests to validate functionality, including handling of various input scenarios and integration with RobotProcessor.
- Updated types.py to include LANGUAGE feature type and modified __init__.py to register the new processor.

* feat(language): Enhance language processing in TokenizerProcessor

- Added OBS_LANGUAGE constant to define the observation language key.
- Updated TokenizerProcessor to store tokenized task data in the observation dictionary, ensuring compatibility with the new language feature.
- Introduced Pi0NewLineProcessor to append newlines to tasks for proper tokenization.
- Modified tests to validate the integration of language tokens and attention masks in the observation structure.

* feat(tokenizer): Add padding configuration to TokenizerProcessor

- Introduced `padding_side` parameter to the TokenizerProcessor for customizable padding direction.
- Updated the `make_pi0_processor` function to include the new padding configuration.
- Enhanced unit tests to validate the functionality of the `padding_side` parameter in various scenarios.

* feat(processor): Add state management methods to Pi0NewLineProcessor

* feat(normalization): Track normalization and unnormalization info in complementary data

- Updated NormalizerProcessor and UnnormalizerProcessor to accept additional parameters for tracking normalization modes.
- Enhanced the __call__ methods to store normalization and unnormalization information in the complementary data of transitions.
- Added unit tests to verify the correct tracking of normalization info, including scenarios with missing stats and selective normalization keys.

* feat(factory): Add preprocessor and postprocessor overrides to ProcessorConfigKwargs

- Updated ProcessorConfigKwargs to include optional overrides for preprocessor and postprocessor configurations.
- Enhanced the make_processor function to utilize the new overrides, allowing for more flexible processor initialization.

* feat(processors): Integrate RenameProcessor into various processor configurations

- Added RenameProcessor to the input steps of multiple processor functions, including make_act_processor, make_diffusion_processor, make_pi0_processor, make_sac_processor, make_tdmpc_processor, make_vqbet_processor, and make_smolvla_processor.
- Consolidated normalization features from input and output into a single NormalizerProcessor for improved efficiency.
- Updated the input steps to ensure compatibility with the new RenameProcessor integration.

* feat(smolvla): Refactor language processing and introduce new line processor (#1658)

- Removed the prepare_language method and directly accessed language tokens and masks from the batch using the OBS_LANGUAGE constant.
- Added SmolVLANewLineProcessor to ensure tasks end with a newline, enhancing tokenization compatibility.
- Updated the make_smolvla_processor function to include the new line processor and tokenizer processor for improved input handling.

* feture(policies): add device processor (#1659)

* feat(processors): Integrate DeviceProcessor into multiple processor configurations

- Added DeviceProcessor to the input and output steps of various processor functions, including make_act_processor, make_diffusion_processor, make_pi0_processor, make_pi0fast_processor, make_sac_processor, make_tdmpc_processor, make_vqbet_processor, and make_smolvla_processor.
- Enhanced the DeviceProcessor class with state management methods and ensured compatibility with existing processor pipelines.
- Introduced unit tests for DeviceProcessor to validate functionality across different scenarios, including CPU and CUDA operations.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor(pipeline): Remove to() method for device management

- Eliminated the to() method from RobotProcessor, which was responsible for moving tensor states to specified devices.
- Removed associated unit tests that validated the functionality of the to() method across various scenarios.
- Streamlined the pipeline code by focusing on other device management strategies.

* feat(processor): Enhance DeviceProcessor with float dtype conversion

- Added support for optional float dtype conversion in DeviceProcessor, allowing tensors to be converted to specified floating-point types while preserving non-float types.
- Implemented validation for float dtype input and updated the processor's configuration methods to include float dtype.
- Refactored tensor processing logic to streamline device movement and dtype conversion.
- Introduced comprehensive unit tests to validate the new float dtype functionality across various scenarios.

* feat(policies): Add new line processors and update module exports

* feat(processor): Enhance batch and device processors to handle index and task_index fields

- Added logic to ToBatchProcessor for unsqueezing 0D tensors for index and task_index fields, ensuring they are processed as 1D tensors.
- Updated DeviceProcessor to process index and task_index fields in complementary data, preserving their tensor types and ensuring non-tensor fields remain unchanged.
- Enhanced unit tests to validate the correct handling of index and task_index fields across various scenarios, including device compatibility and dtype preservation.
2025-08-06 17:21:13 +02:00
pre-commit-ci[bot] a1734cf575 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-06 17:21:12 +02:00
Adil Zouitine 82f300e880 fix(dependencies): Update transformers dependency constraint to allow only versions up to 4.52.0 2025-08-06 17:21:11 +02:00
Adil Zouitine 3e7c9d7afc feat(batch_processor): Add feature_contract method to ToBatchProcessor
- Introduced feature_contract method that returns features without modification, maintaining the no-op behavior of the processor.
- This addition enhances the flexibility of the ToBatchProcessor for future feature processing needs.
2025-08-06 17:21:09 +02:00
Adil Zouitine e9cb779eab refactor(normalization): Clean up imports in normalize_processor.py 2025-08-06 17:21:08 +02:00
Adil Zouitine 8ff95be04c refactor(normalization): Remove unused state dict transformation methods and streamline imports
- Eliminated the _transform_state_dict_keys and _load_as_safetensor methods from PI0Policy, simplifying the model loading process.
- Cleaned up imports in modeling_pi0.py by removing log_model_loading_keys and init_logging.
- Updated TDMPCPolicy and VQBeTPolicy to handle action removal from batches during offline evaluation.
- Introduced hotswap_stats function in normalize_processor.py to update normalization statistics dynamically, with corresponding tests to ensure functionality.
2025-08-06 17:21:07 +02:00
Adil Zouitine f02ce69df0 refactor(diffusion): remove normalization layer from input processing 2025-08-06 17:21:07 +02:00
Adil Zouitine 1feb7b5d88 fix(rebase): remove residual normalization layer: 2025-08-06 17:21:06 +02:00
Adil Zouitine fbe9009db2 feat(normalization): Implement IDENTITY mode for normalization and unnormalization
- Enhanced NormalizerProcessor and UnnormalizerProcessor to support IDENTITY mode, allowing features to bypass normalization when specified.
- Updated processing logic to check normalization modes and handle missing statistics gracefully.
- Added comprehensive unit tests to validate IDENTITY mode functionality for both observations and actions, ensuring correct behavior across various scenarios.
- Improved error handling for unsupported normalization modes.
2025-08-06 17:21:05 +02:00
Adil Zouitine c0013b130b feat(batch_processor): Add task field processing to ToBatchProcessor
- Enhanced ToBatchProcessor to wrap string tasks in a list, adding batch dimensions for compatibility with model inference.
- Implemented a new method for processing complementary data, ensuring that task values are correctly handled as either strings or lists of strings.
- Added comprehensive unit tests to validate task processing, including edge cases and in-place mutation of complementary data.
2025-08-06 17:21:04 +02:00
Adil Zouitine c4763f61a1 refactor(record): Rename processor parameters and update processing logic
- Renamed `processor` to `preprocessor` and added `postprocessor` parameter for clarity.
- Updated the `record_loop` and `predict_action` functions to utilize the new preprocessor and postprocessor, enhancing the processing flow.
- Ensured compatibility with existing functionality while improving code readability.
2025-08-06 17:21:03 +02:00
Adil Zouitine b95c219d96 feat(migrate): Extend load_model_from_hub to include train configuration
- Updated load_model_from_hub to return the train configuration alongside the model state_dict and config.
- Modified main function to handle the additional train configuration when loading models from both the hub and local paths.
- Adjusted dataset_repo_id extraction to utilize the train configuration for improved accuracy.
2025-08-06 17:21:02 +02:00
Adil Zouitine 9b1138171e refactor(factory): Clean up imports in factory.py
- Removed unused import of IdentityProcessor to streamline the code.
2025-08-06 17:21:02 +02:00
Adil Zouitine 023b8f3466 feat(factory): Enhance make_processor to support preprocessor and postprocessor configuration
- Introduced ProcessorConfigKwargs TypedDict for better type safety in processor configuration.
- Updated make_processor to accept preprocessor and postprocessor configuration filenames, improving flexibility in processor instantiation.
- Refactored the loading of pretrained processors to utilize the new configuration options.
2025-08-06 17:21:00 +02:00
pre-commit-ci[bot] 1cad87ebd2 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-06 17:21:00 +02:00
Adil Zouitine 99de7567e6 feat(batch_processor): Enhance ToBatchProcessor to handle action batching
- Updated ToBatchProcessor to add batch dimensions to actions in addition to observations.
- Implemented separate methods for processing observations and actions, improving code readability.
- Added comprehensive unit tests to validate action batching functionality across various tensor dimensions and types.
2025-08-06 17:20:58 +02:00
Adil Zouitine 21baa8fa02 refactor(factory): Remove unused imports and NaN detection hook from processor creation 2025-08-06 17:20:53 +02:00
Adil Zouitine 8b4a5368b3 feat(processors): Add ToBatchProcessor to multiple policy processors
- Integrated ToBatchProcessor into various policy processors to handle observation batching.
- Updated make functions for act, diffusion, pi0, pi0fast, sac, smolvla, tdmpc, and vqbet processors to include the new batching functionality.
- Ensured consistency across all processor implementations for improved data handling.
2025-08-06 17:20:52 +02:00
Adil Zouitine f5c6b03b61 feat(processor): Introduce ToBatchProcessor for handling observation batching
- Added ToBatchProcessor to ensure observations have proper batch dimensions for model processing.
- Implemented functionality to add batch dimensions to state and image observations as needed.
- Created comprehensive unit tests to validate the processor's behavior with various tensor dimensions and types.
- Ensured compatibility with existing transition keys and maintained the integrity of non-observation data.
2025-08-06 17:20:51 +02:00
Adil Zouitine e7be2fd113 feat(migrate): Add model card generation and saving to migration script
- Implemented functionality to generate and save a model card for the migrated model, including metadata such as dataset repository ID, license, and tags.
- Enhanced the script to push the model card to the hub if requested, improving model documentation and accessibility.
- Refactored the saving process to ensure the model card is saved locally and uploaded correctly when pushing to the hub.
2025-08-06 17:20:50 +02:00
Adil Zouitine b632490b4b feat(migrate): Enhance migration script to create preprocessor and postprocessor for policy models
- Updated the migration script to generate both a preprocessor and a postprocessor, improving the handling of normalization for training and inference.
- Added functionality to convert features to PolicyFeature objects, ensuring compatibility with the new processor architecture.
- Refined the extraction and removal of normalization statistics and layers, streamlining the migration process.
- Improved error handling for missing mandatory configuration fields during model instantiation.
2025-08-06 17:20:50 +02:00
pre-commit-ci[bot] 9a9c7208d2 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-06 17:20:49 +02:00
pre-commit-ci[bot] 427b97d198 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2025-08-06 17:20:48 +02:00
AdilZouitine 2c2bb1e8bf feat(migration): Add script for migrating policy models with normalization layers 2025-08-06 17:20:47 +02:00
AdilZouitine 4b24f94225 feat(record): Integrate RobotProcessor into recording loop and update policy handling
- Added support for RobotProcessor in the record_loop function to enhance data processing capabilities.
- Updated the logic to reset both policy and processor when provided, ensuring proper state management.
- Modified action prediction to utilize the processor, improving the overall functionality of the recording process.
- Adjusted the save_checkpoint function to include preprocessor state saving, enhancing checkpointing capabilities.
2025-08-06 17:20:46 +02:00
AdilZouitine 670a278cbc refactor(policies): Enhance processor creation and add NaN detection hook 2025-08-06 17:20:45 +02:00
AdilZouitine fc74001202 refactor(train): Update preprocessor initialization to include dataset statistics 2025-08-06 17:20:45 +02:00
Adil Zouitine f14ac5d486 feat(train): Integrate preprocessor into training pipeline 2025-08-06 17:20:44 +02:00
Adil Zouitine 7bd0d62ce5 refactor(policies): Remove unnormalization step from action predictions
- Eliminated the unnormalization of actions in both `TDMPCPolicy` and `VQBeTPolicy` classes to streamline action prediction.
- This change improves code clarity and aligns with recent updates to policy processors.
2025-08-06 17:20:43 +02:00
Adil Zouitine 7eccefe235 refactor(learner): Remove normalization from cached image features retrieval
- Simplified the retrieval of observation features by removing the normalization step from the `get_cached_image_features` method calls.
- This change enhances clarity and aligns with the recent updates to policy processors.
2025-08-06 17:20:42 +02:00
Adil Zouitine b72274066e feat(processors): Introduce processors for various policy types
- Added `make_processor` function to create processor instances for different policy types, including `tdmpc`, `diffusion`, `act`, `vqbet`, `pi0`, `pi0fast`, `sac`, and `reward_classifier`.
- Implemented corresponding processor files for each policy type, encapsulating normalization and unnormalization steps.
- Updated existing policies to remove direct normalization dependencies, enhancing modularity and clarity.
- Enhanced test coverage to validate the integration of new processors with existing policy configurations.
2025-08-06 17:20:41 +02:00
Steven Palma 20f2910b63 Merge branch 'main' into user/azouitine/2025-7-2-implement-pipeline 2025-08-06 17:20:39 +02:00
Steven Palma fd4ae3466b refactor(pipeline): minor improvements (#1684)
* chore(pipeline): remove unused features + device torch + envtransition keys

* refactor(pipeline): ImageProcessor & StateProcessor are both implemented directly in VanillaObservationPRocessor

* refactor(pipeline): RenameProcessor now inherits from ObservationProcessor + remove unused code

* test(pipeline): fix broken test after refactors

* docs(pipeline): update docstrings VanillaObservationProcessor

* chore(pipeline): move None check to base pipeline classes
2025-08-06 14:00:13 +02:00
Adil Zouitine 7beb040e8e refactor(pipeline): Rename parameters for clarity and enhance save/load functionality
- Updated parameter names in the save_pretrained and from_pretrained methods for improved readability, changing destination_path to save_directory and source to pretrained_model_name_or_path.
- Enhanced the save_pretrained method to ensure directory creation and file handling is consistent with the new parameter names.
- Streamlined the loading process in from_pretrained to utilize loaded_config for better clarity and maintainability.
2025-08-05 17:44:21 +02:00
Adil Zouitine 05bd18f453 refactor(observation): Streamline observation preprocessing and remove unused processor methods
- Updated the `preprocess_observation` function to enhance image handling and ensure proper tensor formatting.
- Removed the `RobotProcessor` and associated transition handling from the `rollout` function, simplifying the observation processing flow.
- Integrated direct calls to `preprocess_observation` for improved clarity and efficiency in the evaluation script.
2025-08-05 10:32:56 +02:00
Adil Zouitine 8077456c00 refactor(pipeline): Remove model card generation and streamline processor methods
- Eliminated the _generate_model_card method from RobotProcessor, which was responsible for generating README.md files from a template.
- Updated save_pretrained method to remove model card generation, focusing on serialization of processor definitions and parameters.
- Added default implementations for get_config, state_dict, load_state_dict, reset, and feature_contract methods in various processor classes to enhance consistency and usability.
2025-08-05 10:31:09 +02:00
AdilZouitine 5595887fd0 refactor(pipeline): Remove to() method for device management
- Eliminated the to() method from RobotProcessor, which was responsible for moving tensor states to specified devices.
- Removed associated unit tests that validated the functionality of the to() method across various scenarios.
- Streamlined the pipeline code by focusing on other device management strategies.
2025-08-05 10:27:25 +02:00
Adil Zouitine 41959389b6 docs(pipeline): Clarify transition handling and hook behavior
- Updated documentation to specify that hooks always receive transitions in EnvTransition format, ensuring consistent behavior across input formats.
- Refactored the step_through method to yield only EnvTransition objects, regardless of the input format, and updated related tests to reflect this change.
- Enhanced test assertions to verify the structure of results and the correctness of processing steps.
2025-08-02 14:51:52 +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
99 changed files with 1867 additions and 3953 deletions
+1 -1
View File
@@ -25,7 +25,7 @@ body:
id: system-info
attributes:
label: System Info
description: Please share your LeRobot configuration by running `lerobot-info` (if installed) or `python -m lerobot.scripts.display_sys_info` (if not installed) and pasting the output below.
description: If needed, you can share your lerobot configuration with us by running `python -m lerobot.scripts.display_sys_info` and copy-pasting its outputs below
render: Shell
placeholder: lerobot version, OS, python version, numpy version, torch version, and lerobot's configuration
validations:
-68
View File
@@ -1,68 +0,0 @@
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This workflow handles closing stale issues and PRs.
name: Stale
on:
# Allows running this workflow manually from the Actions tab
workflow_dispatch:
# Runs at 02:00
schedule:
- cron: "0 2 * * *"
env:
CLOSE_ISSUE_MESSAGE: >
This issue was closed because it has been stalled for 14 days with no activity.
Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
CLOSE_PR_MESSAGE: >
This PR was closed because it has been stalled for 14 days with no activity.
Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
WARN_ISSUE_MESSAGE: >
This issue has been automatically marked as stale because it has not had
recent activity (1 year). It will be closed if no further activity occurs.
Thank you for your contributions.
WARN_PR_MESSAGE: >
This PR has been automatically marked as stale because it has not had
recent activity (1 year). It will be closed if no further activity occurs.
Thank you for your contributions.
jobs:
# This job runs the actions/stale action to close stale issues and PRs.
stale:
name: Close Stale Issues and PRs
runs-on: ubuntu-latest
permissions:
actions: write
contents: write # only for delete-branch option
issues: write
pull-requests: write
steps:
- uses: actions/stale@v10
with:
repo-token: ${{ secrets.GITHUB_TOKEN }}
stale-issue-label: stale
stale-pr-label: stale
exempt-issue-labels: never-stale
exempt-pr-labels: never-stale
days-before-issue-stale: 180 # TODO(Steven): Will modify this to 90 after initial cleanup
days-before-issue-close: 14
days-before-pr-stale: 180
days-before-pr-close: 14
delete-branch: true
close-issue-message: ${{ env.CLOSE_ISSUE_MESSAGE }}
close-pr-message: ${{ env.CLOSE_PR_MESSAGE }}
stale-issue-message: ${{ env.WARN_ISSUE_MESSAGE }}
stale-pr-message: ${{ env.WARN_PR_MESSAGE }}
operations-per-run: 500
+42 -10
View File
@@ -202,7 +202,7 @@ Check out [example 1](https://github.com/huggingface/lerobot/blob/main/examples/
You can also locally visualize episodes from a dataset on the hub by executing our script from the command line:
```bash
lerobot-dataset-viz \
python -m lerobot.scripts.visualize_dataset \
--repo-id lerobot/pusht \
--episode-index 0
```
@@ -210,7 +210,7 @@ lerobot-dataset-viz \
or from a dataset in a local folder with the `root` option and the `--local-files-only` (in the following case the dataset will be searched for in `./my_local_data_dir/lerobot/pusht`)
```bash
lerobot-dataset-viz \
python -m lerobot.scripts.visualize_dataset \
--repo-id lerobot/pusht \
--root ./my_local_data_dir \
--local-files-only 1 \
@@ -221,19 +221,19 @@ It will open `rerun.io` and display the camera streams, robot states and actions
https://github-production-user-asset-6210df.s3.amazonaws.com/4681518/328035972-fd46b787-b532-47e2-bb6f-fd536a55a7ed.mov?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240505%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240505T172924Z&X-Amz-Expires=300&X-Amz-Signature=d680b26c532eeaf80740f08af3320d22ad0b8a4e4da1bcc4f33142c15b509eda&X-Amz-SignedHeaders=host&actor_id=24889239&key_id=0&repo_id=748713144
Our script can also visualize datasets stored on a distant server. See `lerobot-dataset-viz --help` for more instructions.
Our script can also visualize datasets stored on a distant server. See `python -m lerobot.scripts.visualize_dataset --help` for more instructions.
### The `LeRobotDataset` format
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/dataset/load_lerobot_dataset.py) for more details on `delta_timestamps`.
A specificity of `LeRobotDataset` is that, rather than retrieving a single frame by its index, we can retrieve several frames based on their temporal relationship with the indexed frame, by setting `delta_timestamps` to a list of relative times with respect to the indexed frame. For example, with `delta_timestamps = {"observation.image": [-1, -0.5, -0.2, 0]}` one can retrieve, for a given index, 4 frames: 3 "previous" frames 1 second, 0.5 seconds, and 0.2 seconds before the indexed frame, and the indexed frame itself (corresponding to the 0 entry). See example [1_load_lerobot_dataset.py](https://github.com/huggingface/lerobot/blob/main/examples/1_load_lerobot_dataset.py) for more details on `delta_timestamps`.
Under the hood, the `LeRobotDataset` format makes use of several ways to serialize data which can be useful to understand if you plan to work more closely with this format. We tried to make a flexible yet simple dataset format that would cover most type of features and specificities present in reinforcement learning and robotics, in simulation and in real-world, with a focus on cameras and robot states but easily extended to other types of sensory inputs as long as they can be represented by a tensor.
Here are the important details and internal structure organization of a typical `LeRobotDataset` instantiated with `dataset = LeRobotDataset("lerobot/aloha_static_coffee")`. The exact features will change from dataset to dataset but not the main aspects:
```
````
dataset attributes:
├ hf_dataset: a Hugging Face dataset (backed by Arrow/parquet). Typical features example:
│ ├ observation.images.cam_high (VideoFrame):
@@ -269,7 +269,7 @@ dataset attributes:
├ root (Path): local directory where the dataset is stored
├ image_transforms (Callable): optional image transformations to apply to visual modalities
└ delta_timestamps (dict): optional delta timestamps for temporal queries
```
decoding videos (e.g., 'pyav', 'torchcodec')
A `LeRobotDataset` is serialised using several widespread file formats for each of its parts, namely:
@@ -279,6 +279,42 @@ A `LeRobotDataset` is serialised using several widespread file formats for each
Dataset can be uploaded/downloaded from the HuggingFace hub seamlessly. To work on a local dataset, you can specify its location with the `root` argument if it's not in the default `~/.cache/huggingface/lerobot` location.
### 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.
We also provide a more capable script to parallelize the evaluation over multiple environments during the same rollout. Here is an example with a pretrained model hosted on [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht):
```bash
lerobot-eval \
--policy.path=lerobot/diffusion_pusht \
--env.type=pusht \
--eval.batch_size=10 \
--eval.n_episodes=10 \
--policy.use_amp=false \
--policy.device=cuda
````
Note: After training your own policy, you can re-evaluate the checkpoints with:
```bash
lerobot-eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
```
See `lerobot-eval --help` for more instructions.
### 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.
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.
\<img src="https://raw.githubusercontent.com/huggingface/lerobot/main/media/wandb.png" alt="WandB logs example"\>
Note: For efficiency, during training every checkpoint is evaluated on a low number of episodes. You may use `--eval.n_episodes=500` to evaluate on more episodes than the default. Or, after training, you may want to re-evaluate your best checkpoints on more episodes or change the evaluation settings. See `lerobot-eval --help` for more instructions.
#### Reproduce state-of-the-art (SOTA)
We provide some pretrained policies on our [hub page](https://huggingface.co/lerobot) that can achieve state-of-the-art performances.
@@ -337,7 +373,3 @@ If you want, you can cite this work with:
## Star History
[![Star History Chart](https://api.star-history.com/svg?repos=huggingface/lerobot&type=Timeline)](https://star-history.com/#huggingface/lerobot&Timeline)
```
```
-378
View File
@@ -1,378 +0,0 @@
"""
Benchmark memory footprint and inference latency of a policy on arbitrary devices.
This script loads a pretrained policy directly (similar to the async inference server)
and generates dummy input data based on the policy's input_features to perform
accurate benchmarking without requiring datasets.
"""
import argparse
import os
import signal
import statistics
from contextlib import contextmanager
from datetime import datetime
from pathlib import Path
import psutil
import torch
from tqdm import tqdm
from lerobot.configs.types import FeatureType
from lerobot.policies.factory import get_policy_class
from lerobot.policies.pretrained import PreTrainedPolicy
class TimeoutException:
pass
@contextmanager
def timeout(seconds):
def signal_handler(signum, frame):
raise TimeoutException(f"Timed out after {seconds} seconds")
# On Windows, signal is not available, so we can't use this timeout mechanism
if not hasattr(signal, "SIGALRM"):
yield
return
old_handler = signal.signal(signal.SIGALRM, signal_handler)
try:
# signal.alarm expects integer seconds
# for float seconds, we can use setitimer
signal.setitimer(signal.ITIMER_REAL, seconds)
yield
finally:
signal.setitimer(signal.ITIMER_REAL, 0)
signal.signal(signal.SIGALRM, old_handler)
def bytes_to_human(n: int) -> str:
for unit in ["B", "KB", "MB", "GB", "TB"]:
if n < 1024:
return f"{n:.2f} {unit}"
n /= 1024
return f"{n:.2f} PB"
def percentile(values: list[float], p: float) -> float:
if not values:
return float("nan")
k = (len(values) - 1) * (p / 100.0)
f = int(k)
c = min(f + 1, len(values) - 1)
if f == c:
return values[f]
return values[f] + (values[c] - values[f]) * (k - f)
def generate_dummy_observation(input_features: dict, device: str = "cpu") -> dict:
"""Generate dummy observation data based on policy input features."""
dummy_obs = {}
for key, feature in input_features.items():
shape = feature.shape
if feature.type == FeatureType.VISUAL:
# Images: random values in [0, 1] range (already normalized)
dummy_obs[key] = torch.rand(shape, dtype=torch.float32, device=device)
elif feature.type in [FeatureType.STATE, FeatureType.ACTION, FeatureType.ENV]:
# State/action/env: random normal distribution
dummy_obs[key] = torch.randn(shape, dtype=torch.float32, device=device)
else:
# Default: random normal for unknown types
dummy_obs[key] = torch.randn(shape, dtype=torch.float32, device=device)
# Add batch dimension
for key in dummy_obs:
dummy_obs[key] = dummy_obs[key].unsqueeze(0)
# Add task string for language-conditioned policies
dummy_obs["task"] = ""
return dummy_obs
def main():
parser = argparse.ArgumentParser(description="Policy inference benchmark")
parser.add_argument(
"--policy-id", type=str, required=True, help="Model ID or local path to pretrained policy"
)
parser.add_argument(
"--policy-type", type=str, required=True, help="Type of policy (smolvla, act, diffusion, etc.)"
)
parser.add_argument(
"--device", type=str, default="mps", choices=["cuda", "cpu", "mps"], help="Device to run on"
)
parser.add_argument("--seed", type=int, default=42, help="Random seed")
parser.add_argument(
"--num-samples", type=int, default=100, help="Number of inference samples to benchmark"
)
parser.add_argument("--warmup", type=int, default=10, help="Number of warmup samples (not timed)")
parser.add_argument(
"--output-dir", type=str, default="outputs/benchmarks", help="Directory to save benchmark results"
)
parser.add_argument(
"--timeout",
type=float,
default=0.3,
help="Timeout for each inference pass in seconds (default: 0.3s = 300ms)",
)
args = parser.parse_args()
# Seed & deterministic-ish setup
torch.manual_seed(args.seed)
if args.device == "cuda":
torch.cuda.manual_seed_all(args.seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = False # leave False to avoid perf cliffs
# Resolve device availability
device = args.device.lower()
if device == "cuda" and not torch.cuda.is_available():
print("[!] CUDA requested but unavailable. Falling back to CPU.")
device = "cpu"
elif device == "mps" and not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()):
print("[!] MPS requested but unavailable. Falling back to CPU.")
device = "cpu"
use_cuda = device == "cuda"
# Create output directory and log file
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
policy_name = args.policy_id.replace("/", "_").replace("\\", "_")
log_file = output_dir / f"benchmark_{args.policy_type}_{policy_name}_{device}_{timestamp}.txt"
# Load policy directly from pretrained (similar to async inference server)
print(f"Loading policy {args.policy_type} from {args.policy_id}...")
policy_class = get_policy_class(args.policy_type)
policy: PreTrainedPolicy = policy_class.from_pretrained(args.policy_id)
policy.eval()
policy.to(device)
print(f"Policy loaded on {device}")
print(f"Input features: {list(policy.config.input_features.keys())}")
print(f"Output features: {list(policy.config.output_features.keys())}")
# Generate dummy observation based on policy input features
dummy_observation = generate_dummy_observation(policy.config.input_features, device)
dummy_observation["task"] = ""
# Helper to sync for fair timings
def _sync(dev_=device):
if dev_ == "cuda" and torch.cuda.is_available():
torch.cuda.synchronize()
elif dev_ == "mps" and hasattr(torch, "mps"):
try:
torch.mps.synchronize()
except AttributeError:
pass # MPS sync not available in this PyTorch version
# Warmup (to stabilize kernels/caches)
print("Warming up...")
with torch.no_grad():
policy.reset()
for _ in range(args.warmup):
_ = policy.select_action(dummy_observation)
_sync()
# Memory footprint before timing
process = psutil.Process(os.getpid())
rss_before = process.memory_info().rss
if use_cuda:
torch.cuda.reset_peak_memory_stats()
# PyTorch timing with Event objects for more accurate GPU timing
print(f"Running benchmark: {args.num_samples} samples...")
if use_cuda:
# Use CUDA Events for precise GPU timing
start_events = []
end_events = []
timeout_count = 0
with torch.no_grad():
for forward in tqdm(range(args.num_samples), desc="Trials"):
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
try:
with timeout(args.timeout):
start_event.record()
_ = policy.select_action(dummy_observation)
end_event.record()
start_events.append(start_event)
end_events.append(end_event)
except TimeoutException:
timeout_count += 1
# Add placeholder for timeout
start_events.append(None)
end_events.append(None)
print(f"\n[!] Timeout on forward {forward + 1}")
continue
# Synchronize and collect timing results
torch.cuda.synchronize()
per_forward_ms = []
for start_event, end_event in zip(start_events, end_events, strict=True):
if start_event is None:
per_forward_ms.append(args.timeout * 1000)
else:
per_forward_ms.append(start_event.elapsed_time(end_event))
if timeout_count > 0:
print(f"[!] {timeout_count} inference passes timed out (>{args.timeout * 1000:.1f}ms)")
else:
# Use simple time.perf_counter for CPU/MPS timing with timeout
import time
per_forward_ms = []
timeout_count = 0
with torch.no_grad():
for sample in tqdm(range(args.num_samples), desc="Samples"):
try:
with timeout(args.timeout):
start_time = time.perf_counter()
_ = policy.select_action(dummy_observation)
end_time = time.perf_counter()
per_forward_ms.append((end_time - start_time) * 1000) # Convert to ms
except TimeoutException:
timeout_count += 1
per_forward_ms.append(args.timeout * 1000)
print(f"\n[!] Timeout on sample {sample + 1}")
continue
if timeout_count > 0:
print(f"[!] {timeout_count} inference passes timed out (>{args.timeout * 1000:.1f}ms)")
# Memory footprint after timing
rss_after = process.memory_info().rss
rss_delta = rss_after - rss_before
cuda_peak = torch.cuda.max_memory_allocated() if use_cuda else 0
# Sort timing results for percentile calculations
per_forward_ms_sorted = sorted(per_forward_ms)
mean_ms = statistics.fmean(per_forward_ms) if per_forward_ms else float("nan")
std_ms = statistics.pstdev(per_forward_ms) if len(per_forward_ms) > 1 else 0.0
min_ms = per_forward_ms_sorted[0] if per_forward_ms_sorted else float("nan")
max_ms = per_forward_ms_sorted[-1] if per_forward_ms_sorted else float("nan")
p50_ms = percentile(per_forward_ms_sorted, 50)
p95_ms = percentile(per_forward_ms_sorted, 95)
# Model size
num_params = sum(p.numel() for p in policy.parameters())
# Prepare results for logging
results = {
"timestamp": datetime.now().isoformat(),
"policy_type": args.policy_type,
"policy_id": args.policy_id,
"device": device,
"num_trials": args.num_samples,
"forwards_per_trial": 1,
"warmup": args.warmup,
"timeout_ms": args.timeout * 1000,
"seed": args.seed,
"num_params": num_params,
"timeout_count": timeout_count,
"latency_mean_ms": mean_ms,
"latency_std_ms": std_ms,
"latency_min_ms": min_ms,
"latency_max_ms": max_ms,
"latency_p50_ms": p50_ms,
"latency_p95_ms": p95_ms,
"cpu_rss_before": rss_before,
"cpu_rss_after": rss_after,
"cpu_rss_delta": rss_delta,
"cuda_peak_alloc": cuda_peak,
"input_features": list(policy.config.input_features.keys()),
"output_features": list(policy.config.output_features.keys()),
}
# Format and write results to log file
log_content = f"""
=== LeRobot Policy Inference Benchmark ===
Timestamp: {results["timestamp"]}
Policy: {results["policy_type"]} ({results["policy_id"]})
Device: {results["device"]}
Seed: {results["seed"]}
=== Model Information ===
Parameters: {results["num_params"]:,}
Input Features: {", ".join(results["input_features"])}
Output Features: {", ".join(results["output_features"])}
=== Benchmark Configuration ===
Samples: {results["num_trials"]}
Warmup: {results["warmup"]}
Total Measurements: {len(per_forward_ms)}
Timeout: {results["timeout_ms"]:.1f}ms
Timeouts: {results["timeout_count"]} / {results["num_trials"]}
=== Latency Results (ms) ===
Mean: {results["latency_mean_ms"]:.3f}
Std Dev: {results["latency_std_ms"]:.3f}
Min: {results["latency_min_ms"]:.3f}
Max: {results["latency_max_ms"]:.3f}
P50: {results["latency_p50_ms"]:.3f}
P95: {results["latency_p95_ms"]:.3f}
=== Memory Footprint ===
CPU RSS Before: {bytes_to_human(results["cpu_rss_before"])}
CPU RSS After: {bytes_to_human(results["cpu_rss_after"])}{bytes_to_human(results["cpu_rss_delta"])})
"""
if use_cuda:
log_content += f"CUDA Peak: {bytes_to_human(results['cuda_peak_alloc'])} (reset before timing)\n"
log_content += f"""
=== Raw Timing Data (first 20 measurements, ms) ===
{", ".join(f"{t:.3f}" for t in per_forward_ms[:20])}
{"..." if len(per_forward_ms) > 20 else ""}
=== Summary Statistics ===
Timing Method: {"CUDA Events" if use_cuda else "torch.utils.benchmark.Timer"}
Device Available: {torch.cuda.is_available() if device == "cuda" else torch.backends.mps.is_available() if device == "mps" else True}
PyTorch Version: {torch.__version__}
Benchmark completed successfully at {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
"""
# Write to log file
with open(log_file, "w") as f:
f.write(log_content)
# Print to console (shorter version)
print("\n=== Inference Benchmark Results ===")
print(f"Policy: {args.policy_type} ({args.policy_id})")
print(f"Device: {device}")
print(f"Samples: {args.num_samples} | Warmup: {args.warmup}")
print(f"Model params: {num_params:,}")
print("\nLatency per forward (ms):")
print(f" mean: {mean_ms:.3f} std: {std_ms:.3f}")
print(f" min: {min_ms:.3f} max: {max_ms:.3f}")
print(f" p50: {p50_ms:.3f} p95: {p95_ms:.3f}")
print("\nMemory footprint:")
print(f" CPU RSS before: {bytes_to_human(rss_before)}")
print(f" CPU RSS after : {bytes_to_human(rss_after)}{bytes_to_human(rss_delta)})")
if use_cuda:
print(
f" CUDA peak allocated: {bytes_to_human(cuda_peak)} "
f"(reset by reset_peak_memory_stats before timing)"
)
print(f"\nResults saved to: {log_file}")
print("Benchmark completed successfully!")
if __name__ == "__main__":
main()
+2 -2
View File
@@ -23,14 +23,14 @@
- sections:
- local: lerobot-dataset-v3
title: Using LeRobotDataset
- local: libero
title: Using Libero
- local: porting_datasets_v3
title: Porting Large Datasets
title: "Datasets"
- sections:
- local: smolvla
title: Finetune SmolVLA
- local: libero
title: Using Libero
title: "Policies"
- sections:
+15 -14
View File
@@ -62,7 +62,7 @@ pip install -e ".[hilserl]"
### Understanding Configuration
The training process begins with proper configuration for the HILSerl environment. The main configuration class is `GymManipulatorConfig` in `lerobot/rl/gym_manipulator.py`, which contains nested `HILSerlRobotEnvConfig` and `DatasetConfig`. The configuration is organized into focused, nested sub-configs:
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 -->
```python
@@ -160,8 +160,9 @@ The action processor (`action_processor`) handles outgoing actions and human int
1. **AddTeleopActionAsComplimentaryDataStep**: Captures teleoperator actions for logging
2. **AddTeleopEventsAsInfoStep**: Records intervention events and episode control signals
3. **InterventionActionProcessorStep**: Handles human interventions and episode termination
4. **Inverse Kinematics Pipeline** (when enabled):
3. **AddRobotObservationAsComplimentaryData**: Stores raw robot state for processing
4. **InterventionActionProcessorStep**: Handles human interventions and episode termination
5. **Inverse Kinematics Pipeline** (when enabled):
- **MapDeltaActionToRobotActionStep**: Converts delta actions to robot action format
- **EEReferenceAndDelta**: Computes end-effector reference and delta movements
- **EEBoundsAndSafety**: Enforces workspace safety bounds
@@ -346,7 +347,7 @@ With the bounds defined, you can safely collect demonstrations for training. Tra
**Setting Up Record Mode**
Create a configuration file for recording demonstrations (or edit an existing one like [env_config.json](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/env_config.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"` at the root level
2. Specify a unique `repo_id` for your dataset in the `dataset` section (e.g., "username/task_name")
@@ -518,7 +519,7 @@ During the online training, press `space` to take over the policy and `space` ag
Start the recording process, an example of the config file can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json):
```bash
python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config_so100.json
python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config_so100.json
```
During recording:
@@ -549,7 +550,7 @@ Note: If you already know the crop parameters, you can skip this step and just s
Use the `crop_dataset_roi.py` script to interactively select regions of interest in your camera images:
```bash
python -m lerobot.rl.crop_dataset_roi --repo-id username/pick_lift_cube
python -m lerobot.scripts.rl.crop_dataset_roi --repo-id username/pick_lift_cube
```
1. For each camera view, the script will display the first frame
@@ -618,7 +619,7 @@ Before training, you need to collect a dataset with labeled examples. The `recor
To collect a dataset, you need to modify some parameters in the environment configuration based on HILSerlRobotEnvConfig.
```bash
python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/reward_classifier_train_config.json
python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/reward_classifier_train_config.json
```
**Key Parameters for Data Collection**
@@ -764,7 +765,7 @@ or set the argument in the json config file.
Run `gym_manipulator.py` to test the model.
```bash
python -m lerobot.rl.gym_manipulator --config_path path/to/env_config.json
python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config.json
```
The reward classifier will automatically provide rewards based on the visual input from the robot's cameras.
@@ -772,12 +773,12 @@ The reward classifier will automatically provide rewards based on the visual inp
**Example Workflow for training the reward classifier**
1. **Create the configuration files**:
Create the necessary json configuration files for the reward classifier and the environment. Check the examples [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/reward_classifier/config.json).
Create the necessary json configuration files for the reward classifier and the environment. Check the examples [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/tree/main).
2. **Collect a dataset**:
```bash
python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
```
3. **Train the classifier**:
@@ -788,7 +789,7 @@ The reward classifier will automatically provide rewards based on the visual inp
4. **Test the classifier**:
```bash
python -m lerobot.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
```
### Training with Actor-Learner
@@ -797,7 +798,7 @@ The LeRobot system uses a distributed actor-learner architecture for training. T
**Configuration Setup**
Create a training configuration file (example available [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/train_config.json)). The training config is based on the main `TrainRLServerPipelineConfig` class in `lerobot/configs/train.py`.
Create a training configuration file (example available [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/train_config_hilserl_so100.json)). The training config is based on the main `TrainRLServerPipelineConfig` class in `lerobot/configs/train.py`.
1. Configure the policy settings (`type="sac"`, `device`, etc.)
2. Set `dataset` to your cropped dataset
@@ -810,7 +811,7 @@ Create a training configuration file (example available [here](https://huggingfa
First, start the learner server process:
```bash
python -m lerobot.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
python -m lerobot.scripts.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
```
The learner:
@@ -825,7 +826,7 @@ The learner:
In a separate terminal, start the actor process with the same configuration:
```bash
python -m lerobot.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
python -m lerobot.scripts.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
```
The actor:
+6 -6
View File
@@ -26,7 +26,7 @@ pip install -e ".[hilserl]"
## Configuration
To use `gym_hil` with LeRobot, you need to create a configuration file. An example is provided [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/gym_hil/env_config.json). Key configuration sections include:
To use `gym_hil` with LeRobot, you need to create a configuration file. An example is provided [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/gym_hil_env.json). Key configuration sections include:
### Environment Type and Task
@@ -91,7 +91,7 @@ Important parameters:
To run the environment, set mode to null:
```bash
python -m lerobot.rl.gym_manipulator --config_path path/to/gym_hil_env.json
python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
```
### Recording a Dataset
@@ -118,21 +118,21 @@ To collect a dataset, set the mode to `record` whilst defining the repo_id and n
```
```bash
python -m lerobot.rl.gym_manipulator --config_path path/to/gym_hil_env.json
python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
```
### Training a Policy
To train a policy, checkout the configuration example available [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/rl/gym_hil/train_config.json) and run the actor and learner servers:
To train a policy, checkout the configuration example available [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/train_gym_hil_env.json) and run the actor and learner servers:
```bash
python -m lerobot.rl.actor --config_path path/to/train_gym_hil_env.json
python -m lerobot.scripts.rl.actor --config_path path/to/train_gym_hil_env.json
```
In a different terminal, run the learner server:
```bash
python -m lerobot.rl.learner --config_path path/to/train_gym_hil_env.json
python -m lerobot.scripts.rl.learner --config_path path/to/train_gym_hil_env.json
```
The simulation environment provides a safe and repeatable way to develop and test your Human-In-the-Loop reinforcement learning components before deploying to real robots.
+6 -6
View File
@@ -22,7 +22,7 @@ pip install -e ".[hilserl]"
## Teleoperate and Record a Dataset
To use `gym_hil` with LeRobot, you need to use a configuration file. An example config file can be found [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/sim_il/env_config.json).
To use `gym_hil` with LeRobot, you need to use a configuration file. An example config file can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_gym_hil_il.json).
To teleoperate and collect a dataset, we need to modify this config file. Here's an example configuration for imitation learning data collection:
@@ -61,14 +61,14 @@ Then we can run this command to start:
<hfoption id="Linux">
```bash
python -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
```
</hfoption>
<hfoption id="MacOS">
```bash
mjpython -m lerobot.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
mjpython -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
```
</hfoption>
@@ -165,7 +165,7 @@ huggingface-cli upload ${HF_USER}/il_sim_test${CKPT} \
## Evaluate your policy in Sim
To evaluate your policy we have to use a configuration file. An example can be found [here](https://huggingface.co/datasets/lerobot/config_examples/resolve/main/sim_il/eval_config.json).
To evaluate your policy we have to use a configuration file. An example can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/eval_config_gym_hil.json).
Here's an example evaluation configuration:
@@ -198,14 +198,14 @@ Then you can run this command to visualize your trained policy
<hfoption id="Linux">
```bash
python -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
python -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
```
</hfoption>
<hfoption id="MacOS">
```bash
mjpython -m lerobot.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
mjpython -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
```
</hfoption>
+4 -4
View File
@@ -14,7 +14,7 @@ Training data from one robot setup needs adaptation for deployment on different
**That's where processors come in.** They serve as universal translators that bridge these gaps, ensuring seamless data flow from sensors to models to actuators.
Processors handle all the preprocessing and postprocessing steps needed to convert raw environment data into model-ready inputs and vice versa.
This means that your favorite policy can be used like this:
Now your favorite policy can be used like this:
```python
import torch
@@ -40,7 +40,7 @@ postprocessed_action = postprocessor(action)
## What are Processors?
In robotics, data comes in many forms: images from cameras, joint positions from sensors, text instructions from users, and more. Each type of data requires specific transformations before a model can use it effectively. Models need this data to be:
In robotics, data comes in many forms - images from cameras, joint positions from sensors, text instructions from users, and more. Each type of data requires specific transformations before a model can use it effectively. Models need this data to be:
- **Normalized**: Scaled to appropriate ranges for neural network processing
- **Batched**: Organized with proper dimensions for batch processing
@@ -181,8 +181,8 @@ training_preprocessor = PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
RenameObservationsProcessorStep(rename_map={}), # Standardize keys
AddBatchDimensionProcessorStep(), # Add batch dims
TokenizerProcessorStep(tokenizer_name="...", ...), # Tokenize language
DeviceProcessorStep(device="cuda"), # Move to GPU first
NormalizerProcessorStep(features=..., stats=...), # Normalize on GPU
DeviceProcessorStep(device="cuda"), # Move to GPU first
NormalizerProcessorStep(features=..., stats=...), # Normalize on GPU
]
)
-112
View File
@@ -8,7 +8,6 @@ This docs will guide you to:
- Record a dataset and push it to the Hub
- Load datasets for training with `LeRobotDataset`
- Stream datasets without downloading using `StreamingLeRobotDataset`
- Apply image transforms for data augmentation during training
- Migrate existing `v2.1` datasets to `v3.0`
## Whats new in `v3`
@@ -151,117 +150,6 @@ dataset = StreamingLeRobotDataset(repo_id) # streams directly from the Hub
</figure>
</div>
## Image transforms
Image transforms are data augmentations applied to camera frames during training to improve model robustness and generalization. LeRobot supports various transforms including brightness, contrast, saturation, hue, and sharpness adjustments.
### Using transforms during dataset creation/recording
Currently, transforms are applied during **training time only**, not during recording. When you create or record a dataset, the raw images are stored without transforms. This allows you to experiment with different augmentations later without re-recording data.
### Adding transforms to existing datasets (API)
Use the `image_transforms` parameter when loading a dataset for training:
```python
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.transforms import ImageTransforms, ImageTransformsConfig, ImageTransformConfig
# Option 1: Use default transform configuration (disabled by default)
transforms_config = ImageTransformsConfig(
enable=True, # Enable transforms
max_num_transforms=3, # Apply up to 3 transforms per frame
random_order=False, # Apply in standard order
)
transforms = ImageTransforms(transforms_config)
dataset = LeRobotDataset(
repo_id="your-username/your-dataset",
image_transforms=transforms
)
# Option 2: Create custom transform configuration
custom_transforms_config = ImageTransformsConfig(
enable=True,
max_num_transforms=2,
random_order=True,
tfs={
"brightness": ImageTransformConfig(
weight=1.0,
type="ColorJitter",
kwargs={"brightness": (0.7, 1.3)} # Adjust brightness range
),
"contrast": ImageTransformConfig(
weight=2.0, # Higher weight = more likely to be selected
type="ColorJitter",
kwargs={"contrast": (0.8, 1.2)}
),
"sharpness": ImageTransformConfig(
weight=0.5, # Lower weight = less likely to be selected
type="SharpnessJitter",
kwargs={"sharpness": (0.3, 2.0)}
),
}
)
dataset = LeRobotDataset(
repo_id="your-username/your-dataset",
image_transforms=ImageTransforms(custom_transforms_config)
)
# Option 3: Use pure torchvision transforms
from torchvision.transforms import v2
torchvision_transforms = v2.Compose([
v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
v2.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
])
dataset = LeRobotDataset(
repo_id="your-username/your-dataset",
image_transforms=torchvision_transforms
)
```
### Available transform types
LeRobot provides several transform types:
- **`ColorJitter`**: Adjusts brightness, contrast, saturation, and hue
- **`SharpnessJitter`**: Randomly adjusts image sharpness
- **`Identity`**: No transformation (useful for testing)
You can also use any `torchvision.transforms.v2` transform by passing it directly to the `image_transforms` parameter.
### Configuration options
- **`enable`**: Enable/disable transforms (default: `False`)
- **`max_num_transforms`**: Maximum number of transforms applied per frame (default: `3`)
- **`random_order`**: Apply transforms in random order vs. standard order (default: `False`)
- **`weight`**: Sampling probability for each transform (higher = more likely, if sum of weights is not 1, they will be normalized)
- **`kwargs`**: Transform-specific parameters (e.g., brightness range)
### Visualizing transforms
Use the visualization script to preview how transforms affect your data:
```bash
lerobot-imgtransform-viz \
--repo-id=your-username/your-dataset \
--output-dir=./transform_examples \
--n-examples=5
```
This saves example images showing the effect of each transform, helping you tune parameters.
### Best practices
- **Start conservative**: Begin with small ranges (e.g., brightness 0.9-1.1) and increase gradually
- **Test first**: Use the visualization script to ensure transforms look reasonable
- **Monitor training**: Strong augmentations can hurt performance if too aggressive
- **Match your domain**: If your robot operates in varying lighting, use brightness/contrast transforms
- **Combine wisely**: Using too many transforms simultaneously can make training unstable
## Migrate `v2.1` → `v3.0`
A converter aggregates perepisode files into larger shards and writes episode offsets/metadata. Convert your dataset using the instructions below.
+41 -38
View File
@@ -28,7 +28,7 @@ Links:
### Phone orientation and controls
- Orientation: hold the phone with the screen facing up and the top edge pointing in the same direction as the robot gripper. This ensures calibration aligns the phones frame with the robot frame so motion feels natural, see the image below for reference.
- Orientation: hold the phone with the screen facing up and the top edge pointing in the same direction as the robot gripper. This ensures calibration aligns the phones frame with the robot frame so motion feels natural.
- Enable/disable:
- iOS: Hold `B1` to enable teleoperation, release to stop. The first press captures a reference pose.
- Android: Press and hold the `Move` button, release to stop. The first press captures a reference pose.
@@ -36,8 +36,6 @@ Links:
- iOS: Analog input `A3` controls the gripper as velocity input.
- Android: Buttons `A` and `B` act like increment/decrement (A opens, B closes). You can tune velocity in the `GripperVelocityToJoint` step.
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/phone_teleop.webp" alt="Phone teleop orientation" title="Phone teleop orientation" width="40%">
### Step 1: Choose the platform
Modify the examples to use `PhoneOS.IOS` or `PhoneOS.ANDROID` in `PhoneConfig`. The API is identical across platforms, only the input source differs. All examples are under `examples/` and have `phone_so100_*.py` variants.
@@ -66,79 +64,80 @@ Run on of the examples scripts to teleoperate, record a dataset, replay a datase
All scripts assume you configured your robot (e.g., SO-100 follower) and set the correct serial port.
Additionally you need to **copy the urdf of the robot to the examples folder**. For the examples in this tutorial (Using SO100/SO101) it is highly recommended to use the urdf in the [SO-ARM100 repo](https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf)
- Android: after starting the script, open the printed local URL on your phone, tap Start, then press and hold Move.
- iOS: open HEBI Mobile I/O first; B1 enables motion. A3 controls the gripper.
You can customize mapping or safety limits by editing the processor steps shown in the examples.
You can also remap inputs (e.g., use a different analog input) or adapt the pipeline to other robots (e.g., LeKiwi) by modifying the input and kinematics steps. More about this in the [Processors for Robots and Teleoperators](./processors_robots_teleop.mdx) guide.
- Run this example to teleoperate:
```bash
python examples/phone_to_so100/teleoperate.py
python examples/phone_so100_teleop.py
```
After running the example:
- Android: after starting the script, open the printed local URL on your phone, tap Start, then press and hold Move.
- iOS: open HEBI Mobile I/O first; B1 enables motion. A3 controls the gripper.
Additionally you can customize mapping or safety limits by editing the processor steps shown in the examples. You can also remap inputs (e.g., use a different analog input) or adapt the pipeline to other robots (e.g., LeKiwi) by modifying the input and kinematics steps. More about this in the [Processors for Robots and Teleoperators](./processors_robots_teleop.mdx) guide.
- Run this example to record a dataset, which saves absolute end effector observations and actions:
```bash
python examples/phone_to_so100/record.py
python examples/phone_so100_record.py
```
- Run this example to replay recorded episodes:
```bash
python examples/phone_to_so100/replay.py
python examples/phone_so100_replay.py
```
- Run this example to evaluate a pretrained policy:
```bash
python examples/phone_to_so100/evaluate.py
python examples/phone_so100_eval.py
```
### Important pipeline steps and options
- Kinematics are used in multiple steps. We use [Placo](https://github.com/Rhoban/placo) which is a wrapper around Pinocchio for handling our kinematics. We construct the kinematics object by passing the robot's URDF and target frame. We set `target_frame_name` to the gripper frame.
```examples/phone_to_so100/teleoperate.py
kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()),
```44:49:examples/phone_so100_teleop.py
RobotKinematics(
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()),
)
```
- The `MapPhoneActionToRobotAction` step converts the calibrated phone pose and inputs into target deltas and gripper commands, below is shown what the step outputs.
```src/lerobot/teleoperators/phone/phone_processor.py
action["enabled"] = enabled
action["target_x"] = -pos[1] if enabled else 0.0
action["target_y"] = pos[0] if enabled else 0.0
action["target_z"] = pos[2] if enabled else 0.0
action["target_wx"] = rotvec[1] if enabled else 0.0
action["target_wy"] = rotvec[0] if enabled else 0.0
action["target_wz"] = -rotvec[2] if enabled else 0.0
action["gripper_vel"] = gripper_vel # Still send gripper action when disabled
```72:83:src/lerobot/teleoperators/phone/phone_processor.py
# Map calibrated phone pose to robot targets (enabled gates the motion)
act.update(
{
"action.enabled": enabled,
"action.target_x": -pos[1] if enabled else 0.0,
"action.target_y": pos[0] if enabled else 0.0,
"action.target_z": pos[2] if enabled else 0.0,
"action.target_wx": rotvec[1] if enabled else 0.0,
"action.target_wy": rotvec[0] if enabled else 0.0,
"action.target_wz": -rotvec[2] if enabled else 0.0,
"action.gripper": gripper,
}
)
```
- The `EEReferenceAndDelta` step converts target deltas to an absolute desired EE pose, storing a reference on enable, the `end_effector_step_sizes` are the step sizes for the EE pose and can be modified to change the motion speed.
```examples/phone_to_so100/teleoperate.py
```56:65:examples/phone_so100_teleop.py
EEReferenceAndDelta(
kinematics=kinematics_solver,
end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
motor_names=list(robot.bus.motors.keys()),
use_latched_reference=True,
),
)
```
- The `EEBoundsAndSafety` step clamps EE motion to a workspace and checks for large ee step jumps to ensure safety. The `end_effector_bounds` are the bounds for the EE pose and can be modified to change the workspace. The `max_ee_step_m` and `max_ee_twist_step_rad` are the step limits for the EE pose and can be modified to change the safety limits.
```examples/phone_to_so100/teleoperate.py
```61:66:examples/phone_so100_teleop.py
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
max_ee_step_m=0.10,
@@ -148,8 +147,11 @@ Additionally you can customize mapping or safety limits by editing the processor
- The `GripperVelocityToJoint` step turns a velocitylike gripper input into absolute gripper position using the current measured state. The `speed_factor` is the factor by which the velocity is multiplied.
```examples/phone_to_so100/teleoperate.py
GripperVelocityToJoint(speed_factor=20.0)
```78:81:examples/phone_so100_teleop.py
GripperVelocityToJoint(
motor_names=list(robot.bus.motors.keys()),
speed_factor=20.0,
)
```
#### Different IK initial guesses
@@ -158,7 +160,7 @@ We use different IK initial guesses in the kinematic steps. As initial guess eit
- Closed loop (used in record/eval): sets `initial_guess_current_joints=True` so IK starts from the measured joints each frame.
```examples/phone_to_so100/record.py
```71:76:examples/phone_so100_eval.py
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
motor_names=list(robot.bus.motors.keys()),
@@ -168,7 +170,7 @@ We use different IK initial guesses in the kinematic steps. As initial guess eit
- Open loop (used in replay): sets `initial_guess_current_joints=False` so IK continues from the previous IK solution rather than the measured state. This preserves action stability when we replay without feedback.
```examples/phone_to_so100/replay.py
```80:86:examples/phone_so100_replay.py
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
motor_names=list(robot.bus.motors.keys()),
@@ -179,6 +181,7 @@ We use different IK initial guesses in the kinematic steps. As initial guess eit
### Pipeline steps explained
- MapPhoneActionToRobotAction: converts calibrated phone pose and inputs into target deltas and a gripper command. Motion is gated by an enable signal (B1 on iOS, Move on Android).
- AddRobotObservationAsComplimentaryData: reads current robot joints and inserts them under `complementary_data.raw_joint_positions` for FK/IK steps to use.
- EEReferenceAndDelta: latches a reference EE pose on enable and combines it with target deltas to produce an absolute desired EE pose each frame. When disabled, it keeps sending the last commanded pose.
- EEBoundsAndSafety: clamps the EE pose to a workspace and ratelimits jumps for safety. Also declares `action.ee.*` features.
- InverseKinematicsEEToJoints: turns an EE pose into joint positions with IK. `initial_guess_current_joints=True` is recommended for closedloop control; set `False` for openloop replay for stability.
+68 -71
View File
@@ -17,7 +17,7 @@ We use the Phone to SO100 follower examples for concreteness, but the same pa
The examples in this guide use absolute end effector (EE) poses because they are easy to reason about. In practice, relative EE deltas or joint position are often preferred as learning features.
With processors, you choose the learning features you want to use for your policy. This could be joints positions/velocities, absolute EE, or relative EE positions. You can also choose to store other features, such as joint torques, motor currents, etc.
You can choose what you save and learn from the teleop and robot action spaces, joints, absolute EE, or relative EE by using/implementing the right steps (and `transform_features()`) in your pipelines.
## Three pipelines
@@ -31,102 +31,99 @@ Each of these pipelines handle different conversions between different action an
Below is an example of the three pipelines that we use in the phone to SO-100 follower examples:
```69:90:examples/phone_so100_record.py
phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotAction]( # teleop -> dataset action
steps=[
MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
EEReferenceAndDelta(
kinematics=kinematics_solver, end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5}, motor_names=list(robot.bus.motors.keys()),
),
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, max_ee_step_m=0.20, max_ee_twist_step_rad=0.50,
),
GripperVelocityToJoint(),
],
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
phone_to_robot_ee_pose = RobotProcessor( # teleop -> dataset action
steps=[MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
AddRobotObservationAsComplimentaryData(robot=robot),
EEReferenceAndDelta(kinematics=kinematics_solver,
end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
motor_names=list(robot.bus.motors.keys())),
EEBoundsAndSafety(end_effector_bounds={"min": [-1, -1, -1], "max": [1, 1, 1]},
max_ee_step_m=0.20, max_ee_twist_step_rad=0.50)],
to_transition=to_transition_teleop_action,
to_output=lambda tr: tr,
)
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction]( # dataset action -> robot
steps=[
InverseKinematicsEEToJoints(
kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()), initial_guess_current_joints=True,
),
],
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
robot_ee_to_joints = RobotProcessor( # dataset action -> robot
steps=[InverseKinematicsEEToJoints(kinematics=kinematics_solver,
motor_names=list(robot.bus.motors.keys()),
initial_guess_current_joints=True),
GripperVelocityToJoint(motor_names=list(robot.bus.motors.keys()), speed_factor=20.0)],
to_transition=lambda tr: tr,
to_output=to_output_robot_action,
)
robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation]( # robot obs -> dataset obs
steps=[
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys()))
],
to_transition=observation_to_transition,
to_output=transition_to_observation,
robot_joints_to_ee_pose = RobotProcessor( # robot obs -> dataset obs
steps=[ForwardKinematicsJointsToEE(kinematics=kinematics_solver,
motor_names=list(robot.bus.motors.keys()))],
to_transition=to_transition_robot_observation,
to_output=lambda tr: tr,
)
```
## Why to_transition / to_output
To convert from robot/teleoperator to pipeline and back, we use the `to_transition` and `to_output` pipeline adapters.
They standardize conversions to reduce boilerplate code, and form the bridge between the robot and teleoperators raw dictionaries and the pipelines `EnvTransition` format.
They standardize conversions to reduce boilerplate code, and form the bridge between the robot and teleoperators raw dicts and the pipelines `EnvTransition` format.
In the phone to SO-100 follower examples we use the following adapters:
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
- `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
- `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
- `transition_to_observation`: transforms the pipeline transition to a observation dict.
- `to_transition_teleop_action`: transforms the teleop action dict to a pipeline transition (puts keys under `action.*`, converts scalars/arrays to tensors, keeps objects like `Rotation` intact)
- `to_output_robot_action`: transforms the pipeline transition to a robot action dict (extracts keys ending with `.pos`/`.vel` and strips `action.` prefix)
- `to_transition_robot_observation`: transforms the robot observation dict to a pipeline transition (splits state vs images; stores state under `observation.state.*` and images under `observation.images.*`)
Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
See `src/lerobot/processor/converters.py` for more details.
## Dataset feature contracts
Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
Dataset features are the keys saved in the dataset. Each step can declare what its dataset features are via `transform_features(...)`. We can then aggregate features per pipeline with `aggregate_pipeline_dataset_features()` and merge multiple groups with `merge_features(...)`.
Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
```src/lerobot/robots/so100_follower/robot_kinematic_processor.py
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
# We only use the ee pose in the dataset, so we don't need the joint positions
for n in self.motor_names:
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
# We specify the dataset features of this step that we want to be stored in the dataset
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
type=FeatureType.STATE, shape=(1,)
)
return features
```203:211:src/lerobot/robots/so100_follower/robot_kinematic_processor.py
def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
# Because this is last step we specify the dataset features of this step that we want to be stored in the dataset
features["action.ee.x"] = float
features["action.ee.y"] = float
features["action.ee.z"] = float
features["action.ee.wx"] = float
features["action.ee.wy"] = float
features["action.ee.wz"] = float
return features
```
Here we declare what PolicyFeatures we modify in this step, so we know what features we can expect when we run the processor. These features can then be aggregated and used to create the dataset features.
Tip: declare features at the last step that produces them (e.g., `EEBoundsAndSafety` declares `action.ee.*`, `ForwardKinematicsJointsToEE` declares `observation.state.ee.*`).
Below is an example of how we aggregate and merge features in the phone to SO-100 record example:
Below is an example of how we aggregate and merge features in the phone to SO-100 follower examples:
```121:145:examples/phone_so100_record.py
features=combine_feature_dicts(
# Run the feature contract of the pipelines
# This tells you how the features would look like after the pipeline steps
aggregate_pipeline_dataset_features(
pipeline=phone_to_robot_ee_pose_processor,
initial_features=create_initial_features(action=phone.action_features), # <- Action features we can expect, these come from our teleop device (phone) and action processor
use_videos=True,
),
aggregate_pipeline_dataset_features(
pipeline=robot_joints_to_ee_pose,
initial_features=create_initial_features(observation=robot.observation_features), # <- Observation features we can expect, these come from our robot and observation processor
use_videos=True,
patterns=["observation.state.ee"], # <- Here you could optionally filter the features we want to store in the dataset, with a specific pattern
action_ee = aggregate_pipeline_dataset_features(
pipeline=phone_to_robot_ee_pose,
initial_features=phone.action_features,
use_videos=True,
patterns=["action.ee"],
)
),
),
gripper = aggregate_pipeline_dataset_features(
pipeline=robot_ee_to_joints,
initial_features={},
use_videos=True,
patterns=["action.gripper.pos", "observation.state.gripper.pos"],
)
observation_ee = aggregate_pipeline_dataset_features(
pipeline=robot_joints_to_ee_pose,
initial_features=robot.observation_features,
use_videos=True,
patterns=["observation.state.ee"],
)
dataset_features = merge_features(action_ee, gripper, observation_ee)
```
How it works:
- `aggregate_pipeline_dataset_features(...)`: applies `transform_features` across the pipeline and filters by patterns (images included when `use_videos=True`, and state features included when `patterns` is specified).
- `combine_feature_dicts(...)`: combine multiple feature dicts.
- Recording with `record_loop(...)` uses `build_dataset_frame(...)` to build frames consistent with `dataset.features` before we call `add_frame(...)` to add the frame to the dataset.
- `aggregate_pipeline_dataset_features(...)`: applies `transform_features` across the pipeline and filters by patterns (images included when `use_videos=True`).
- `merge_features(...)`: combine multiple feature dicts.
- Recording uses `to_dataset_frame(...)` to build frames consistent with `dataset.features` before we call `add_frame(...)` to add the frame to the dataset.
## Guidance when customizing robot pipelines
@@ -135,17 +132,17 @@ You can store any of the following features as your action/observation space:
- Joint positions
- Absolute EE poses
- Relative EE deltas
- Other features: joint velocity, torques, etc.
- Other features: joint velocity, etc.
Pick what you want to use for your policy action and observation space and configure/modify the pipelines and steps accordingly.
### Different robots
- You can easily reuse pipelines, for example to use another robot with phone teleop, modify the examples and swap the robot `RobotKinematics` (URDF) and `motor_names` to use your own robot with Phone teleop. Additionally you should ensure `target_frame_name` points to your gripper/wrist.
- Swap `RobotKinematics` URDF and `motor_names`. Ensure `target_frame_name` points to your gripper/wrist.
### Safety first
- When changing pipelines, start with tight bounds, implement safety steps when working with real robots.
- Its advised to start with simulation first and then move to real robots.
Thats it! We hope this guide helps you get started with customizing your robot pipelines, If you run into any issues at any point, jump into our [Discord community](https://discord.com/invite/s3KuuzsPFb) for support.
Hope this guide helps you get started with customizing your robot pipelines, If you run into any issues at any point, jump into our [Discord community](https://discord.com/invite/s3KuuzsPFb) for support.
@@ -136,7 +136,7 @@ print(f"{dataset[0]['action'].shape=}\n") # (64, c)
# PyTorch datasets.
dataloader = torch.utils.data.DataLoader(
dataset,
num_workers=4,
num_workers=0,
batch_size=32,
shuffle=True,
)
+139
View File
@@ -0,0 +1,139 @@
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This script demonstrates how to evaluate a pretrained policy from the HuggingFace Hub or from your local
training outputs directory. In the latter case, you might want to run examples/3_train_policy.py first.
It requires the installation of the 'gym_pusht' simulation environment. Install it by running:
```bash
pip install -e ".[pusht]"
```
"""
from pathlib import Path
import gym_pusht # noqa: F401
import gymnasium as gym
import imageio
import numpy
import torch
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
# Create a directory to store the video of the evaluation
output_directory = Path("outputs/eval/example_pusht_diffusion")
output_directory.mkdir(parents=True, exist_ok=True)
# Select your device
device = "cuda"
# Provide the [hugging face repo id](https://huggingface.co/lerobot/diffusion_pusht):
pretrained_policy_path = "lerobot/diffusion_pusht"
# OR a path to a local outputs/train folder.
# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")
policy = DiffusionPolicy.from_pretrained(pretrained_policy_path)
# Initialize evaluation environment to render two observation types:
# an image of the scene and state/position of the agent. The environment
# also automatically stops running after 300 interactions/steps.
env = gym.make(
"gym_pusht/PushT-v0",
obs_type="pixels_agent_pos",
max_episode_steps=300,
)
# We can verify that the shapes of the features expected by the policy match the ones from the observations
# produced by the environment
print(policy.config.input_features)
print(env.observation_space)
# Similarly, we can check that the actions produced by the policy will match the actions expected by the
# environment
print(policy.config.output_features)
print(env.action_space)
# Reset the policy and environments to prepare for rollout
policy.reset()
numpy_observation, info = env.reset(seed=42)
# Prepare to collect every rewards and all the frames of the episode,
# from initial state to final state.
rewards = []
frames = []
# Render frame of the initial state
frames.append(env.render())
step = 0
done = False
while not done:
# Prepare observation for the policy running in Pytorch
state = torch.from_numpy(numpy_observation["agent_pos"])
image = torch.from_numpy(numpy_observation["pixels"])
# Convert to float32 with image from channel first in [0,255]
# to channel last in [0,1]
state = state.to(torch.float32)
image = image.to(torch.float32) / 255
image = image.permute(2, 0, 1)
# Send data tensors from CPU to GPU
state = state.to(device, non_blocking=True)
image = image.to(device, non_blocking=True)
# Add extra (empty) batch dimension, required to forward the policy
state = state.unsqueeze(0)
image = image.unsqueeze(0)
# Create the policy input dictionary
observation = {
"observation.state": state,
"observation.image": image,
}
# Predict the next action with respect to the current observation
with torch.inference_mode():
action = policy.select_action(observation)
# Prepare the action for the environment
numpy_action = action.squeeze(0).to("cpu").numpy()
# Step through the environment and receive a new observation
numpy_observation, reward, terminated, truncated, info = env.step(numpy_action)
print(f"{step=} {reward=} {terminated=}")
# Keep track of all the rewards and frames
rewards.append(reward)
frames.append(env.render())
# The rollout is considered done when the success state is reached (i.e. terminated is True),
# or the maximum number of iterations is reached (i.e. truncated is True)
done = terminated | truncated | done
step += 1
if terminated:
print("Success!")
else:
print("Failure!")
# Get the speed of environment (i.e. its number of frames per second).
fps = env.metadata["render_fps"]
# Encode all frames into a mp4 video.
video_path = output_directory / "rollout.mp4"
imageio.mimsave(str(video_path), numpy.stack(frames), fps=fps)
print(f"Video of the evaluation is available in '{video_path}'.")
@@ -12,7 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""This script demonstrates how to train Diffusion Policy on the PushT environment."""
"""This script demonstrates how to train Diffusion Policy on the PushT environment.
Once you have trained a model with this script, you can try to evaluate it on
examples/2_evaluate_pretrained_policy.py
"""
from pathlib import Path
@@ -23,7 +27,6 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetad
from lerobot.datasets.utils import dataset_to_policy_features
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
from lerobot.policies.factory import make_pre_post_processors
def main():
@@ -53,10 +56,9 @@ def main():
cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
# We can now instantiate our policy with this config and the dataset stats.
policy = DiffusionPolicy(cfg)
policy = DiffusionPolicy(cfg, dataset_stats=dataset_metadata.stats)
policy.train()
policy.to(device)
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
# Another policy-dataset interaction is with the delta_timestamps. Each policy expects a given number frames
# which can differ for inputs, outputs and rewards (if there are some).
@@ -97,7 +99,7 @@ def main():
done = False
while not done:
for batch in dataloader:
batch = preprocessor(batch)
batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
loss, _ = policy.forward(batch)
loss.backward()
optimizer.step()
@@ -112,8 +114,6 @@ def main():
# Save a policy checkpoint.
policy.save_pretrained(output_directory)
preprocessor.save_pretrained(output_directory)
postprocessor.save_pretrained(output_directory)
if __name__ == "__main__":
+311
View File
@@ -0,0 +1,311 @@
This tutorial will explain the training script, how to use it, and particularly how to configure everything needed for the training run.
> **Note:** The following assumes you're running these commands on a machine equipped with a cuda GPU. If you don't have one (or if you're using a Mac), you can add `--policy.device=cpu` (`--policy.device=mps` respectively). However, be advised that the code executes much slower on cpu.
## The training script
LeRobot offers a training script at [`lerobot/scripts/train.py`](../src/lerobot/scripts/train.py). At a high level it does the following:
- Initialize/load a configuration for the following steps using.
- Instantiates a dataset.
- (Optional) Instantiates a simulation environment corresponding to that dataset.
- Instantiates a policy.
- Runs a standard training loop with forward pass, backward pass, optimization step, and occasional logging, evaluation (of the policy on the environment), and checkpointing.
## Overview of the configuration system
In the training script, the main function `train` expects a `TrainPipelineConfig` object:
<!-- prettier-ignore-start -->
```python
# train.py
@parser.wrap()
def train(cfg: TrainPipelineConfig):
```
<!-- prettier-ignore-end -->
You can inspect the `TrainPipelineConfig` defined in [`lerobot/configs/train.py`](../src/lerobot/configs/train.py) (which is heavily commented and meant to be a reference to understand any option)
When running the script, inputs for the command line are parsed thanks to the `@parser.wrap()` decorator and an instance of this class is automatically generated. Under the hood, this is done with [Draccus](https://github.com/dlwh/draccus) which is a tool dedicated to this purpose. If you're familiar with Hydra, Draccus can similarly load configurations from config files (.json, .yaml) and also override their values through command line inputs. Unlike Hydra, these configurations are pre-defined in the code through dataclasses rather than being defined entirely in config files. This allows for more rigorous serialization/deserialization, typing, and to manipulate configuration as objects directly in the code and not as dictionaries or namespaces (which enables nice features in an IDE such as autocomplete, jump-to-def, etc.)
Let's have a look at a simplified example. Amongst other attributes, the training config has the following attributes:
<!-- prettier-ignore-start -->
```python
@dataclass
class TrainPipelineConfig:
dataset: DatasetConfig
env: envs.EnvConfig | None = None
policy: PreTrainedConfig | None = None
```
<!-- prettier-ignore-end -->
in which `DatasetConfig` for example is defined as such:
<!-- prettier-ignore-start -->
```python
@dataclass
class DatasetConfig:
repo_id: str
episodes: list[int] | None = None
video_backend: str = "pyav"
```
<!-- prettier-ignore-end -->
This creates a hierarchical relationship where, for example assuming we have a `cfg` instance of `TrainPipelineConfig`, we can access the `repo_id` value with `cfg.dataset.repo_id`.
From the command line, we can specify this value by using a very similar syntax `--dataset.repo_id=repo/id`.
By default, every field takes its default value specified in the dataclass. If a field doesn't have a default value, it needs to be specified either from the command line or from a config file which path is also given in the command line (more in this below). In the example above, the `dataset` field doesn't have a default value which means it must be specified.
## Specifying values from the CLI
Let's say that we want to train [Diffusion Policy](../src/lerobot/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
```bash
lerobot-train \
--dataset.repo_id=lerobot/pusht \
--policy.type=diffusion \
--env.type=pusht
```
Let's break this down:
- To specify the dataset, we just need to specify its `repo_id` on the hub which is the only required argument in the `DatasetConfig`. The rest of the fields have default values and in this case we are fine with those so we can just add the option `--dataset.repo_id=lerobot/pusht`.
- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/policies](../src/lerobot/policies)
- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/envs/configs.py`](../src/lerobot/envs/configs.py)
Let's see another example. Let's say you've been training [ACT](../src/lerobot/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
```bash
lerobot-train \
--policy.type=act \
--dataset.repo_id=lerobot/aloha_sim_insertion_human \
--env.type=aloha \
--output_dir=outputs/train/act_aloha_insertion
```
> Notice we added `--output_dir` to explicitly tell where to write outputs from this run (checkpoints, training state, configs etc.). This is not mandatory and if you don't specify it, a default directory will be created from the current date and time, env.type and policy.type. This will typically look like `outputs/train/2025-01-24/16-10-05_aloha_act`.
We now want to train a different policy for aloha on another task. We'll change the dataset and use [lerobot/aloha_sim_transfer_cube_human](https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_human) instead. Of course, we also need to change the task of the environment as well to match this other task.
Looking at the [`AlohaEnv`](../src/lerobot/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
```bash
lerobot-train \
--policy.type=act \
--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
--env.type=aloha \
--env.task=AlohaTransferCube-v0 \
--output_dir=outputs/train/act_aloha_transfer
```
## Loading from a config file
Now, let's assume that we want to reproduce the run just above. That run has produced a `train_config.json` file in its checkpoints, which serializes the `TrainPipelineConfig` instance it used:
```json
{
"dataset": {
"repo_id": "lerobot/aloha_sim_transfer_cube_human",
"episodes": null,
...
},
"env": {
"type": "aloha",
"task": "AlohaTransferCube-v0",
"fps": 50,
...
},
"policy": {
"type": "act",
"n_obs_steps": 1,
...
},
...
}
```
We can then simply load the config values from this file using:
```bash
lerobot-train \
--config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
--output_dir=outputs/train/act_aloha_transfer_2
```
`--config_path` is also a special argument which allows to initialize the config from a local config file. It can point to a directory that contains `train_config.json` or to the config file itself directly.
Similarly to Hydra, we can still override some parameters in the CLI if we want to, e.g.:
```bash
lerobot-train \
--config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
--output_dir=outputs/train/act_aloha_transfer_2
--policy.n_action_steps=80
```
> Note: While `--output_dir` is not required in general, in this case we need to specify it since it will otherwise take the value from the `train_config.json` (which is `outputs/train/act_aloha_transfer`). In order to prevent accidental deletion of previous run checkpoints, we raise an error if you're trying to write in an existing directory. This is not the case when resuming a run, which is what you'll learn next.
`--config_path` can also accept the repo_id of a repo on the hub that contains a `train_config.json` file, e.g. running:
```bash
lerobot-train --config_path=lerobot/diffusion_pusht
```
will start a training run with the same configuration used for training [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht)
## Resume training
Being able to resume a training run is important in case it crashed or aborted for any reason. We'll demonstrate how to do that here.
Let's reuse the command from the previous run and add a few more options:
```bash
lerobot-train \
--policy.type=act \
--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
--env.type=aloha \
--env.task=AlohaTransferCube-v0 \
--log_freq=25 \
--save_freq=100 \
--output_dir=outputs/train/run_resumption
```
Here we've taken care to set up the log frequency and checkpointing frequency to low numbers so we can showcase resumption. You should be able to see some logging and have a first checkpoint within 1 minute (depending on hardware). Wait for the first checkpoint to happen, you should see a line that looks like this in your terminal:
```
INFO 2025-01-24 16:10:56 ts/train.py:263 Checkpoint policy after step 100
```
Now let's simulate a crash by killing the process (hit `ctrl`+`c`). We can then simply resume this run from the last checkpoint available with:
```bash
lerobot-train \
--config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
--resume=true
```
You should see from the logging that your training picks up from where it left off.
Another reason for which you might want to resume a run is simply to extend training and add more training steps. The number of training steps is set by the option `--steps`, which is 100 000 by default.
You could double the number of steps of the previous run with:
```bash
lerobot-train \
--config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
--resume=true \
--steps=200000
```
## Outputs of a run
In the output directory, there will be a folder called `checkpoints` with the following structure:
```bash
outputs/train/run_resumption/checkpoints
├── 000100 # checkpoint_dir for training step 100
│ ├── pretrained_model/
│ │ ├── config.json # policy config
│ │ ├── model.safetensors # policy weights
│ │ └── train_config.json # train config
│ └── training_state/
│ ├── optimizer_param_groups.json # optimizer param groups
│ ├── optimizer_state.safetensors # optimizer state
│ ├── rng_state.safetensors # rng states
│ ├── scheduler_state.json # scheduler state
│ └── training_step.json # training step
├── 000200
└── last -> 000200 # symlink to the last available checkpoint
```
## Fine-tuning a pre-trained policy
In addition to the features currently in Draccus, we've added a special `.path` argument for the policy, which allows to load a policy as you would with `PreTrainedPolicy.from_pretrained()`. In that case, `path` can be a local directory that contains a checkpoint or a repo_id pointing to a pretrained policy on the hub.
For example, we could fine-tune a [policy pre-trained on the aloha transfer task](https://huggingface.co/lerobot/act_aloha_sim_transfer_cube_human) on the aloha insertion task. We can achieve this with:
```bash
lerobot-train \
--policy.path=lerobot/act_aloha_sim_transfer_cube_human \
--dataset.repo_id=lerobot/aloha_sim_insertion_human \
--env.type=aloha \
--env.task=AlohaInsertion-v0
```
When doing so, keep in mind that the features of the fine-tuning dataset would have to match the input/output features of the pretrained policy.
## Typical logs and metrics
When you start the training process, you will first see your full configuration being printed in the terminal. You can check it to make sure that you configured your run correctly. The final configuration will also be saved with the checkpoint.
After that, you will see training log like this one:
```
INFO 2024-08-14 13:35:12 ts/train.py:192 step:0 smpl:64 ep:1 epch:0.00 loss:1.112 grdn:15.387 lr:2.0e-07 updt_s:1.738 data_s:4.774
```
or evaluation log:
```
INFO 2024-08-14 13:38:45 ts/train.py:226 step:100 smpl:6K ep:52 epch:0.25 ∑rwrd:20.693 success:0.0% eval_s:120.266
```
These logs will also be saved in wandb if `wandb.enable` is set to `true`. Here are the meaning of some abbreviations:
- `smpl`: number of samples seen during training.
- `ep`: number of episodes seen during training. An episode contains multiple samples in a complete manipulation task.
- `epch`: number of time all unique samples are seen (epoch).
- `grdn`: gradient norm.
- `∑rwrd`: compute the sum of rewards in every evaluation episode and then take an average of them.
- `success`: average success rate of eval episodes. Reward and success are usually different except for the sparsing reward setting, where reward=1 only when the task is completed successfully.
- `eval_s`: time to evaluate the policy in the environment, in second.
- `updt_s`: time to update the network parameters, in second.
- `data_s`: time to load a batch of data, in second.
Some metrics are useful for initial performance profiling. For example, if you find the current GPU utilization is low via the `nvidia-smi` command and `data_s` sometimes is too high, you may need to modify batch size or number of dataloading workers to accelerate dataloading. We also recommend [pytorch profiler](https://github.com/huggingface/lerobot?tab=readme-ov-file#improve-your-code-with-profiling) for detailed performance probing.
## In short
We'll summarize here the main use cases to remember from this tutorial.
#### Train a policy from scratch CLI
```bash
lerobot-train \
--policy.type=act \ # <- select 'act' policy
--env.type=pusht \ # <- select 'pusht' environment
--dataset.repo_id=lerobot/pusht # <- train on this dataset
```
#### Train a policy from scratch - config file + CLI
```bash
lerobot-train \
--config_path=path/to/pretrained_model \ # <- can also be a repo_id
--policy.n_action_steps=80 # <- you may still override values
```
#### Resume/continue a training run
```bash
lerobot-train \
--config_path=checkpoint/pretrained_model/ \
--resume=true \
--steps=200000 # <- you can change some training parameters
```
#### Fine-tuning
```bash
lerobot-train \
--policy.path=lerobot/act_aloha_sim_transfer_cube_human \ # <- can also be a local path to a checkpoint
--dataset.repo_id=lerobot/aloha_sim_insertion_human \
--env.type=aloha \
--env.task=AlohaInsertion-v0
```
---
Now that you know the basics of how to train a policy, you might want to know how to apply this knowledge to actual robots, or how to record your own datasets and train policies on your specific task?
If that's the case, head over to the next tutorial [`7_get_started_with_real_robot.md`](./7_get_started_with_real_robot.md).
Or in the meantime, happy training! 🤗
@@ -13,7 +13,11 @@
# limitations under the License.
"""This script demonstrates how to train a Diffusion Policy on the PushT environment,
using a dataset processed in streaming mode."""
using a dataset processed in streaming mode.
Once you have trained a model with this script, you can try to evaluate it on
examples/2_evaluate_pretrained_policy.py
"""
from pathlib import Path
@@ -26,7 +30,6 @@ from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.datasets.utils import dataset_to_policy_features
from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.policies.act.modeling_act import ACTPolicy
from lerobot.policies.factory import make_pre_post_processors
def main():
@@ -47,7 +50,9 @@ def main():
training_steps = 10
log_freq = 1
dataset_id = "lerobot/droid_1.0.1" # 26M frames! Would require 4TB of disk space if installed locally (:
dataset_id = (
"aractingi/droid_1.0.1" # 26M frames! Would require 4TB of disk space if installed locally (:
)
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
features = dataset_to_policy_features(dataset_metadata.features)
output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
@@ -55,10 +60,9 @@ def main():
# We can now instantiate our policy with this config and the dataset stats.
cfg = ACTConfig(input_features=input_features, output_features=output_features)
policy = ACTPolicy(cfg)
policy = ACTPolicy(cfg, dataset_stats=dataset_metadata.stats)
policy.train()
policy.to(device)
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
# Delta timestamps are used to (1) augment frames used during training and (2) supervise the policy.
# Here, we use delta-timestamps to only provide ground truth actions for supervision
@@ -85,7 +89,13 @@ def main():
done = False
while not done:
for batch in dataloader:
batch = preprocessor(batch)
batch = {
k: (v.type(torch.float32) if isinstance(v, torch.Tensor) and v.dtype != torch.bool else v)
for k, v in batch.items()
}
batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
# batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
loss, _ = policy.forward(batch)
loss.backward()
optimizer.step()
@@ -100,8 +110,6 @@ def main():
# Save a policy checkpoint.
policy.save_pretrained(output_directory)
preprocessor.save_pretrained(output_directory)
postprocessor.save_pretrained(output_directory)
if __name__ == "__main__":
@@ -1,177 +0,0 @@
#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This example demonstrates how to use image transforms with LeRobot datasets for data augmentation during training.
Image transforms are applied to camera frames to improve model robustness and generalization. They are applied
at training time only, not during dataset recording, allowing you to experiment with different augmentations
without re-recording data.
"""
import torch
from torchvision.transforms import v2
from torchvision.transforms.functional import to_pil_image
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.transforms import ImageTransformConfig, ImageTransforms, ImageTransformsConfig
def save_image(tensor, filename):
"""Helper function to save a tensor as an image file."""
if tensor.dim() == 3: # [C, H, W]
if tensor.max() > 1.0:
tensor = tensor / 255.0
tensor = torch.clamp(tensor, 0.0, 1.0)
pil_image = to_pil_image(tensor)
pil_image.save(filename)
print(f"Saved: {filename}")
else:
print(f"Skipped {filename}: unexpected tensor shape {tensor.shape}")
def example_1_default_transforms():
"""Example 1: Use default transform configuration and save original vs transformed images"""
print("\n Example 1: Default Transform Configuration with Image Saving")
repo_id = "pepijn223/record_main_0" # Example dataset
try:
# Load dataset without transforms (original)
dataset_original = LeRobotDataset(repo_id=repo_id)
# Load dataset with transforms enabled
transforms_config = ImageTransformsConfig(
enable=True, # Enable transforms (disabled by default)
max_num_transforms=2, # Apply up to 2 transforms per frame
random_order=False, # Apply in standard order
)
dataset_with_transforms = LeRobotDataset(
repo_id=repo_id, image_transforms=ImageTransforms(transforms_config)
)
# Save original and transformed images for comparison
if len(dataset_original) > 0:
frame_idx = 0 # Use first frame
original_sample = dataset_original[frame_idx]
transformed_sample = dataset_with_transforms[frame_idx]
print(f"Saving comparison images (frame {frame_idx}):")
for cam_key in dataset_original.meta.camera_keys:
if cam_key in original_sample and cam_key in transformed_sample:
cam_name = cam_key.replace(".", "_").replace("/", "_")
# Save original and transformed images
save_image(original_sample[cam_key], f"{cam_name}_original.png")
save_image(transformed_sample[cam_key], f"{cam_name}_transformed.png")
except Exception as e:
print(f"Could not load dataset '{repo_id}': {e}")
def example_2_custom_transforms():
"""Example 2: Create custom transform configuration and save examples"""
print("\n Example 2: Custom Transform Configuration")
repo_id = "pepijn223/record_main_0" # Example dataset
try:
# Create custom transform configuration with strong effects
custom_transforms_config = ImageTransformsConfig(
enable=True,
max_num_transforms=2, # Apply up to 2 transforms per frame
random_order=True, # Apply transforms in random order
tfs={
"brightness": ImageTransformConfig(
weight=1.0,
type="ColorJitter",
kwargs={"brightness": (0.5, 1.5)}, # Strong brightness range
),
"contrast": ImageTransformConfig(
weight=1.0, # Higher weight = more likely to be selected
type="ColorJitter",
kwargs={"contrast": (0.6, 1.4)}, # Strong contrast
),
"sharpness": ImageTransformConfig(
weight=0.5, # Lower weight = less likely to be selected
type="SharpnessJitter",
kwargs={"sharpness": (0.2, 2.0)}, # Strong sharpness variation
),
},
)
dataset_with_custom_transforms = LeRobotDataset(
repo_id=repo_id, image_transforms=ImageTransforms(custom_transforms_config)
)
# Save examples with strong transforms
if len(dataset_with_custom_transforms) > 0:
sample = dataset_with_custom_transforms[0]
print("Saving custom transform examples:")
for cam_key in dataset_with_custom_transforms.meta.camera_keys:
if cam_key in sample:
cam_name = cam_key.replace(".", "_").replace("/", "_")
save_image(sample[cam_key], f"{cam_name}_custom_transforms.png")
except Exception as e:
print(f"Could not load dataset '{repo_id}': {e}")
def example_3_torchvision_transforms():
"""Example 3: Use pure torchvision transforms and save examples"""
print("\n Example 3: Pure Torchvision Transforms")
repo_id = "pepijn223/record_main_0" # Example dataset
try:
# Create torchvision transform pipeline
torchvision_transforms = v2.Compose(
[
v2.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1),
v2.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
v2.RandomRotation(degrees=10), # Small rotation
]
)
dataset_with_torchvision = LeRobotDataset(repo_id=repo_id, image_transforms=torchvision_transforms)
# Save examples with torchvision transforms
if len(dataset_with_torchvision) > 0:
sample = dataset_with_torchvision[0]
print("Saving torchvision transform examples:")
for cam_key in dataset_with_torchvision.meta.camera_keys:
if cam_key in sample:
cam_name = cam_key.replace(".", "_").replace("/", "_")
save_image(sample[cam_key], f"{cam_name}_torchvision.png")
except Exception as e:
print(f"Could not load dataset '{repo_id}': {e}")
def main():
"""Run all examples"""
print("LeRobot Dataset Image Transforms Examples")
example_1_default_transforms()
example_2_custom_transforms()
example_3_torchvision_transforms()
if __name__ == "__main__":
main()
-2
View File
@@ -60,8 +60,6 @@ preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy,
pretrained_path=HF_MODEL_ID,
dataset_stats=dataset.meta.stats,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
)
# Connect the robot
+1 -1
View File
@@ -26,7 +26,7 @@ from lerobot.utils.control_utils import init_keyboard_listener
from lerobot.utils.utils import log_say
from lerobot.utils.visualization_utils import _init_rerun
NUM_EPISODES = 2
NUM_EPISODES = 3
FPS = 30
EPISODE_TIME_SEC = 30
RESET_TIME_SEC = 10
+2 -4
View File
@@ -32,9 +32,7 @@ robot = LeKiwiClient(robot_config)
# Fetch the dataset to replay
dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns("action")
actions = dataset.hf_dataset.select_columns("action")
# Connect to the robot
robot.connect()
@@ -44,7 +42,7 @@ if not robot.is_connected:
print("Starting replay loop...")
log_say(f"Replaying episode {EPISODE_IDX}")
for idx in range(len(episode_frames)):
for idx in range(dataset.num_frames):
t0 = time.perf_counter()
# Get recorded action from dataset
+6 -6
View File
@@ -30,13 +30,14 @@ from lerobot.processor import (
)
from lerobot.processor.converters import (
observation_to_transition,
robot_action_observation_to_transition,
robot_action_to_transition,
transition_to_observation,
transition_to_robot_action,
)
from lerobot.record import record_loop
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import (
AddRobotObservationAsComplimentaryData,
ForwardKinematicsJointsToEE,
InverseKinematicsEEToJoints,
)
@@ -68,21 +69,22 @@ policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()),
)
# Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[
AddRobotObservationAsComplimentaryData(robot=robot),
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
motor_names=list(robot.bus.motors.keys()),
initial_guess_current_joints=True,
),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
@@ -128,8 +130,6 @@ preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy,
pretrained_path=HF_MODEL_ID,
dataset_stats=dataset.meta.stats,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
)
# Connect the robot
+11 -10
View File
@@ -22,13 +22,14 @@ from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
from lerobot.processor.converters import (
observation_to_transition,
robot_action_observation_to_transition,
robot_action_to_transition,
transition_to_observation,
transition_to_robot_action,
)
from lerobot.record import record_loop
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import (
AddRobotObservationAsComplimentaryData,
EEBoundsAndSafety,
EEReferenceAndDelta,
ForwardKinematicsJointsToEE,
@@ -43,7 +44,7 @@ from lerobot.utils.control_utils import init_keyboard_listener
from lerobot.utils.utils import log_say
from lerobot.utils.visualization_utils import _init_rerun
NUM_EPISODES = 2
NUM_EPISODES = 10
FPS = 30
EPISODE_TIME_SEC = 60
RESET_TIME_SEC = 30
@@ -53,7 +54,7 @@ HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Create the robot and teleoperator configurations
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411",
port="/dev/tty.usbmodem58760434471",
id="my_awesome_follower_arm",
cameras=camera_config,
use_degrees=True,
@@ -66,34 +67,34 @@ phone = Phone(teleop_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()),
)
# Build pipeline to convert phone action to EE action
phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[
MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
AddRobotObservationAsComplimentaryData(robot=robot),
EEReferenceAndDelta(
kinematics=kinematics_solver,
end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
motor_names=list(robot.bus.motors.keys()),
use_latched_reference=True,
),
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
max_ee_step_m=0.20,
max_ee_twist_step_rad=0.50,
),
GripperVelocityToJoint(speed_factor=20.0),
GripperVelocityToJoint(),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
# Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
@@ -101,7 +102,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
initial_guess_current_joints=True,
),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
+11 -17
View File
@@ -18,13 +18,11 @@ import time
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
from lerobot.processor.converters import (
robot_action_observation_to_transition,
transition_to_robot_action,
)
from lerobot.processor import RobotAction, RobotProcessorPipeline
from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import (
AddRobotObservationAsComplimentaryData,
InverseKinematicsEEToJoints,
)
from lerobot.robots.so100_follower.so100_follower import SO100Follower
@@ -36,7 +34,7 @@ HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Initialize the robot config
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
)
# Initialize the robot
@@ -44,29 +42,28 @@ robot = SO100Follower(robot_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()),
)
# Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[
AddRobotObservationAsComplimentaryData(robot=robot),
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
motor_names=list(robot.bus.motors.keys()),
initial_guess_current_joints=False, # Because replay is open loop
),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
# Fetch the dataset to replay
dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns("action")
actions = dataset.hf_dataset.select_columns("action")
# Connect to the robot
robot.connect()
@@ -76,7 +73,7 @@ if not robot.is_connected:
print("Starting replay loop...")
log_say(f"Replaying episode {EPISODE_IDX}")
for idx in range(len(episode_frames)):
for idx in range(dataset.num_frames):
t0 = time.perf_counter()
# Get recorded action from dataset
@@ -84,11 +81,8 @@ for idx in range(len(episode_frames)):
name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
}
# Get robot observation
robot_obs = robot.get_observation()
# Dataset EE -> robot joints
joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
joint_action = robot_ee_to_joints_processor(ee_action)
# Send action to robot
_ = robot.send_action(joint_action)
+9 -15
View File
@@ -16,13 +16,11 @@
import time
from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
from lerobot.processor.converters import (
robot_action_observation_to_transition,
transition_to_robot_action,
)
from lerobot.processor import RobotAction, RobotProcessorPipeline
from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import (
AddRobotObservationAsComplimentaryData,
EEBoundsAndSafety,
EEReferenceAndDelta,
GripperVelocityToJoint,
@@ -39,7 +37,7 @@ FPS = 30
# Initialize the robot and teleoperator
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", use_degrees=True
)
teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID
@@ -49,20 +47,20 @@ teleop_device = Phone(teleop_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()),
)
# Build pipeline to convert phone action to ee pose action to joint action
phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
phone_to_robot_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[
MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
AddRobotObservationAsComplimentaryData(robot=robot),
EEReferenceAndDelta(
kinematics=kinematics_solver,
end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5},
motor_names=list(robot.bus.motors.keys()),
use_latched_reference=True,
),
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
@@ -75,10 +73,9 @@ phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, Robo
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
motor_names=list(robot.bus.motors.keys()),
initial_guess_current_joints=True,
),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
@@ -96,14 +93,11 @@ print("Starting teleop loop. Move your phone to teleoperate the robot...")
while True:
t0 = time.perf_counter()
# Get robot observation
robot_obs = robot.get_observation()
# Get teleop action
phone_obs = teleop_device.get_action()
# Phone -> EE pose -> Joints transition
joint_action = phone_to_robot_joints_processor((phone_obs, robot_obs))
joint_action = phone_to_robot_joints_processor(phone_obs)
# Send action to robot
_ = robot.send_action(joint_action)
+6 -6
View File
@@ -30,13 +30,14 @@ from lerobot.processor import (
)
from lerobot.processor.converters import (
observation_to_transition,
robot_action_observation_to_transition,
robot_action_to_transition,
transition_to_observation,
transition_to_robot_action,
)
from lerobot.record import record_loop
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import (
AddRobotObservationAsComplimentaryData,
ForwardKinematicsJointsToEE,
InverseKinematicsEEToJoints,
)
@@ -68,21 +69,22 @@ policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()),
)
# Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[
AddRobotObservationAsComplimentaryData(robot=robot),
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
motor_names=list(robot.bus.motors.keys()),
initial_guess_current_joints=True,
),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
@@ -129,8 +131,6 @@ preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy,
pretrained_path=HF_MODEL_ID,
dataset_stats=dataset.meta.stats,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
)
# Connect the robot and teleoperator
+10 -9
View File
@@ -23,13 +23,14 @@ from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
from lerobot.processor.converters import (
observation_to_transition,
robot_action_observation_to_transition,
robot_action_to_transition,
transition_to_observation,
transition_to_robot_action,
)
from lerobot.record import record_loop
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import (
AddRobotObservationAsComplimentaryData,
EEBoundsAndSafety,
ForwardKinematicsJointsToEE,
InverseKinematicsEEToJoints,
@@ -41,7 +42,7 @@ from lerobot.utils.control_utils import init_keyboard_listener
from lerobot.utils.utils import log_say
from lerobot.utils.visualization_utils import _init_rerun
NUM_EPISODES = 2
NUM_EPISODES = 10
FPS = 30
EPISODE_TIME_SEC = 60
RESET_TIME_SEC = 30
@@ -61,14 +62,14 @@ leader = SO100Leader(leader_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
follower_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(follower.bus.motors.keys()),
)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
leader_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(leader.bus.motors.keys()),
)
@@ -85,19 +86,20 @@ follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservatio
)
# Build pipeline to convert leader joints to EE action
leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
leader_joints_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[
ForwardKinematicsJointsToEE(
kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
# Build pipeline to convert EE action to follower joints
ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction](
[
AddRobotObservationAsComplimentaryData(robot=follower),
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
max_ee_step_m=0.10,
@@ -106,10 +108,9 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
InverseKinematicsEEToJoints(
kinematics=follower_kinematics_solver,
motor_names=list(follower.bus.motors.keys()),
initial_guess_current_joints=True,
),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
+10 -16
View File
@@ -19,13 +19,11 @@ import time
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
from lerobot.processor.converters import (
robot_action_observation_to_transition,
transition_to_robot_action,
)
from lerobot.processor import RobotAction, RobotProcessorPipeline
from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import (
AddRobotObservationAsComplimentaryData,
InverseKinematicsEEToJoints,
)
from lerobot.robots.so100_follower.so100_follower import SO100Follower
@@ -45,29 +43,28 @@ robot = SO100Follower(robot_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./src/lerobot/teleoperators/sim/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()),
)
# Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
robot_ee_to_joints_processor = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[
AddRobotObservationAsComplimentaryData(robot=robot),
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
motor_names=list(robot.bus.motors.keys()),
initial_guess_current_joints=False, # Because replay is open loop
),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
# Fetch the dataset to replay
dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns("action")
actions = dataset.hf_dataset.select_columns("action")
# Connect to the robot
robot.connect()
@@ -77,7 +74,7 @@ if not robot.is_connected:
print("Starting replay loop...")
log_say(f"Replaying episode {EPISODE_IDX}")
for idx in range(len(episode_frames)):
for idx in range(dataset.num_frames):
t0 = time.perf_counter()
# Get recorded action from dataset
@@ -85,11 +82,8 @@ for idx in range(len(episode_frames)):
name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
}
# Get robot observation
robot_obs = robot.get_observation()
# Dataset EE -> robot joints
joint_action = robot_ee_to_joints_processor((ee_action, robot_obs))
joint_action = robot_ee_to_joints_processor(ee_action)
# Send action to robot
_ = robot.send_action(joint_action)
+8 -11
View File
@@ -17,14 +17,14 @@
import time
from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
from lerobot.processor import RobotAction, RobotProcessorPipeline
from lerobot.processor.converters import (
robot_action_observation_to_transition,
robot_action_to_transition,
transition_to_robot_action,
)
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.robot_kinematic_processor import (
AddRobotObservationAsComplimentaryData,
EEBoundsAndSafety,
ForwardKinematicsJointsToEE,
InverseKinematicsEEToJoints,
@@ -49,14 +49,14 @@ leader = SO100Leader(leader_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
follower_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(follower.bus.motors.keys()),
)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
leader_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf",
urdf_path="./examples/phone_to_so100/SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link",
joint_names=list(leader.bus.motors.keys()),
)
@@ -73,8 +73,9 @@ leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
)
# build pipeline to convert EE action to robot joints
ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
ee_to_follower_joints = RobotProcessorPipeline[RobotAction, RobotAction](
[
AddRobotObservationAsComplimentaryData(robot=follower),
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
max_ee_step_m=0.10,
@@ -83,10 +84,9 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
InverseKinematicsEEToJoints(
kinematics=follower_kinematics_solver,
motor_names=list(follower.bus.motors.keys()),
initial_guess_current_joints=False,
),
],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
@@ -101,9 +101,6 @@ print("Starting teleop loop...")
while True:
t0 = time.perf_counter()
# Get robot observation
robot_obs = follower.get_observation()
# Get teleop observation
leader_joints_obs = leader.get_action()
@@ -111,7 +108,7 @@ while True:
leader_ee_act = leader_to_ee(leader_joints_obs)
# teleop EE -> robot joints
follower_joints_act = ee_to_follower_joints((leader_ee_act, robot_obs))
follower_joints_act = ee_to_follower_joints(leader_ee_act)
# Send action to robot
_ = follower.send_action(follower_joints_act)
+58
View File
@@ -0,0 +1,58 @@
#!/bin/bash
# storage / caches
RAID=/raid/jade
export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
export HF_HOME=$RAID/.cache/huggingface
export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
export WANDB_CACHE_DIR=$RAID/.cache/wandb
export TMPDIR=$RAID/.cache/tmp
mkdir -p $TMPDIR
export WANDB_MODE=offline
export HF_DATASETS_OFFLINE=1
export HF_HUB_OFFLINE=1
export TOKENIZERS_PARALLELISM=false
export MUJOCO_GL=egl
export CUDA_VISIBLE_DEVICES=2
# CONFIGURATION
POLICY_PATH="/raid/jade/logs/lerobot/lerobot_2_HuggingFaceVLA_libero_smolvla_lr1e-4bs32steps100000/checkpoints/100000/pretrained_model"
POLICY_PATH="/raid/jade/models/smolvla_pipe"
TASK=libero_spatial
ENV_TYPE="libero"
BATCH_SIZE=1
N_EPISODES=1
# storage / caches
RAID=/raid/jade
N_ACTION_STEPS=1
export TRANSFORMERS_CACHE=$RAID/.cache/huggingface/transformers
export HF_HOME=$RAID/.cache/huggingface
export HF_DATASETS_CACHE=$RAID/.cache/huggingface/datasets
export HF_LEROBOT_HOME=$RAID/.cache/huggingface/lerobot
export WANDB_CACHE_DIR=$RAID/.cache/wandb
export TMPDIR=$RAID/.cache/tmp
mkdir -p $TMPDIR
export WANDB_MODE=offline
# export HF_DATASETS_OFFLINE=1
# export HF_HUB_OFFLINE=1
export TOKENIZERS_PARALLELISM=false
export MUJOCO_GL=egl
export MUJOCO_GL=egl
unset HF_HUB_OFFLINE
# RUN EVALUATION
python src/lerobot/scripts/eval.py \
--policy.path="$POLICY_PATH" \
--env.type="$ENV_TYPE" \
--eval.batch_size="$BATCH_SIZE" \
--eval.n_episodes="$N_EPISODES" \
--env.task=$TASK \
--env.max_parallel_tasks=10 \
# python examples/evaluate_libero.py \
# --policy_path "$POLICY_PATH" \
# --task_suite_name "$TASK" \
# --num_steps_wait 10 \
# --num_trials_per_task 10 \
# --video_out_path "data/libero/videos" \
# --device "cuda" \
# --seed 7
+18 -9
View File
@@ -59,7 +59,7 @@ keywords = ["lerobot", "huggingface", "robotics", "machine learning", "artifici
dependencies = [
# Hugging Face dependencies
"datasets>=4.0.0",
"datasets>=2.19.0,<=3.6.0", # TODO: Bumb dependency
"diffusers>=0.27.2",
"huggingface-hub[hf-transfer,cli]>=0.34.2",
@@ -121,7 +121,7 @@ phone = ["hebi-py>=2.8.0", "teleop>=0.1.0"]
# Policies
pi0 = ["lerobot[transformers-dep]"]
smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14", "accelerate>=1.7.0", "safetensors>=0.4.3"]
hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.11", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.9", "lerobot[grpcio-dep]", "lerobot[placo-dep]"]
# Features
async = ["lerobot[grpcio-dep]", "matplotlib>=3.10.3"]
@@ -135,9 +135,21 @@ video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"]
aloha = ["gym-aloha>=0.1.1"]
pusht = ["gym-pusht>=0.1.5", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
xarm = ["gym-xarm>=0.1.1"]
libero = ["lerobot[transformers-dep]", "libero @ git+https://github.com/huggingface/lerobot-libero.git@main#egg=libero"]
libero = [
"hydra-core>=1.2,<1.4",
"easydict>=1.9",
"lerobot[transformers-dep]",
"robomimic==0.2.0",
"thop>=0.1.0.post2206102148",
"robosuite==1.4.0",
"bddl==1.0.1",
"matplotlib>=3.5.3",
"cloudpickle>=2.0.0",
"gym>=0.25,<0.27",
"future>=0.18.3",
"egl_probe @ git+https://github.com/jadechoghari/egl_probe.git#egg=egl_probe",
"libero @ git+https://github.com/jadechoghari/LIBERO.git@main#egg=libero",
]
# All
all = [
"lerobot[dynamixel]",
@@ -158,7 +170,7 @@ all = [
"lerobot[pusht]",
"lerobot[xarm]",
"lerobot[phone]",
"lerobot[libero]",
"lerobot[libero]"
]
[project.scripts]
@@ -171,9 +183,6 @@ lerobot-setup-motors="lerobot.setup_motors:main"
lerobot-teleoperate="lerobot.teleoperate:main"
lerobot-eval="lerobot.scripts.eval:main"
lerobot-train="lerobot.scripts.train:main"
lerobot-dataset-viz="lerobot.scripts.lerobot_dataset_viz:main"
lerobot-info="lerobot.scripts.lerobot_info:main"
lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"
# ---------------- Tool Configurations ----------------
[tool.setuptools.packages.find]
+5 -4
View File
@@ -196,10 +196,11 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
config = json.load(f)
config.pop("type")
with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
with tempfile.NamedTemporaryFile("w+") as f:
json.dump(config, f)
config_file = f.name
f.flush()
cli_overrides = policy_kwargs.pop("cli_overrides", [])
with draccus.config_type("json"):
return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
cli_overrides = policy_kwargs.pop("cli_overrides", [])
with draccus.config_type("json"):
return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
@@ -404,7 +404,7 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
info["video_files_size_in_mb"] = video_file_size_in_mb
info["data_path"] = DEFAULT_DATA_PATH
info["video_path"] = DEFAULT_VIDEO_PATH
info["fps"] = int(info["fps"])
info["fps"] = float(info["fps"])
for key in info["features"]:
if info["features"][key]["dtype"] == "video":
# already has fps in video_info
+87 -59
View File
@@ -31,7 +31,6 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
features: dict[str, PolicyFeature] = field(default_factory=dict)
features_map: dict[str, str] = field(default_factory=dict)
max_parallel_tasks: int = 1
disable_env_checker: bool = True
@property
def type(self) -> str:
@@ -163,71 +162,33 @@ class XarmEnv(EnvConfig):
@dataclass
class ImagePreprocessingConfig:
crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
resize_size: tuple[int, int] | None = None
class VideoRecordConfig:
"""Configuration for video recording in ManiSkill environments."""
enabled: bool = False
record_dir: str = "videos"
trajectory_name: str = "trajectory"
@dataclass
class RewardClassifierConfig:
"""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
@dataclass
class ObservationConfig:
"""Configuration for observation processing."""
class EnvTransformConfig:
"""Configuration for environment wrappers."""
# 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_current_to_observation: bool = False
add_ee_pose_to_observation: bool = False
display_cameras: bool = False
@dataclass
class GripperConfig:
"""Configuration for gripper control and penalties."""
use_gripper: bool = True
gripper_penalty: float = 0.0
gripper_penalty_in_reward: bool = False
@dataclass
class ResetConfig:
"""Configuration for environment reset behavior."""
crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
resize_size: tuple[int, int] | None = None
control_time_s: float = 20.0
fixed_reset_joint_positions: Any | None = None
reset_time_s: float = 5.0
control_time_s: float = 20.0
terminate_on_success: bool = True
@dataclass
class HILSerlProcessorConfig:
"""Configuration for environment processing pipeline."""
control_mode: str = "gamepad"
observation: ObservationConfig | None = None
image_preprocessing: ImagePreprocessingConfig | None = None
gripper: GripperConfig | None = None
reset: ResetConfig | None = None
inverse_kinematics: InverseKinematicsConfig | None = None
reward_classifier: RewardClassifierConfig | None = None
max_gripper_pos: float | None = 100.0
use_gripper: bool = True
gripper_quantization_threshold: float | None = 0.8
gripper_penalty: float = 0.0
gripper_penalty_in_reward: bool = False
@EnvConfig.register_subclass(name="gym_manipulator")
@@ -237,15 +198,82 @@ class HILSerlRobotEnvConfig(EnvConfig):
robot: RobotConfig | None = None
teleop: TeleoperatorConfig | None = None
processor: HILSerlProcessorConfig = field(default_factory=HILSerlProcessorConfig)
wrapper: EnvTransformConfig | None = None
fps: int = 10
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"
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
def gym_kwargs(self) -> dict:
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,
}
@EnvConfig.register_subclass("libero")
@dataclass
class LiberoEnv(EnvConfig):
+7 -4
View File
@@ -17,7 +17,7 @@ import importlib
import gymnasium as gym
from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv, XarmEnv
from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, LiberoEnv, PushtEnv, XarmEnv
def make_env_config(env_type: str, **kwargs) -> EnvConfig:
@@ -27,6 +27,8 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
return PushtEnv(**kwargs)
elif env_type == "xarm":
return XarmEnv(**kwargs)
elif env_type == "hil":
return HILEnvConfig(**kwargs)
elif env_type == "libero":
return LiberoEnv(**kwargs)
else:
@@ -76,13 +78,14 @@ def make_env(
try:
importlib.import_module(package_name)
except ModuleNotFoundError as e:
print(f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.type}]'`")
raise e
raise ModuleNotFoundError(
f'{package_name} is not installed. Install with: pip install "lerobot[{cfg.type}]"'
) from e
gym_handle = f"{package_name}/{cfg.task}"
def _make_one():
return gym.make(gym_handle, disable_env_checker=cfg.disable_env_checker, **(cfg.gym_kwargs or {}))
return gym.make(gym_handle, disable_env_checker=True, **(cfg.gym_kwargs or {}))
vec = env_cls([_make_one for _ in range(n_envs)])
+36 -14
View File
@@ -15,6 +15,7 @@
# limitations under the License.
from __future__ import annotations
import math
import os
from collections import defaultdict
from collections.abc import Callable, Iterable, Mapping, Sequence
@@ -28,7 +29,6 @@ import torch
from gymnasium import spaces
from libero.libero import benchmark, get_libero_path
from libero.libero.envs import OffScreenRenderEnv
from robosuite.utils.transform_utils import quat2axisangle
def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
@@ -44,7 +44,7 @@ def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
return cams
def _get_suite(name: str) -> benchmark.Benchmark:
def _get_suite(name: str) -> Any:
"""Instantiate a LIBERO suite by name with clear validation."""
bench = benchmark.get_benchmark_dict()
if name not in bench:
@@ -66,6 +66,33 @@ def _select_task_ids(total_tasks: int, task_ids: Iterable[int] | None) -> list[i
return ids
def quat2axisangle(quat: np.ndarray) -> np.ndarray:
"""
Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
Converts quaternion to axis-angle format.
Returns a unit vector direction scaled by its angle in radians.
Args:
quat (np.array): (x,y,z,w) vec4 float angles
Returns:
np.array: (ax,ay,az) axis-angle exponential coordinates
"""
# clip quaternion
if quat[3] > 1.0:
quat[3] = 1.0
elif quat[3] < -1.0:
quat[3] = -1.0
den = np.sqrt(1.0 - quat[3] * quat[3])
if math.isclose(den, 0.0):
# This is (close to) a zero degree rotation, immediately return
return np.zeros(3)
return (quat[:3] * 2.0 * math.acos(quat[3])) / den
def get_task_init_states(task_suite: Any, i: int) -> np.ndarray:
init_states_path = (
Path(get_libero_path("init_states"))
@@ -83,10 +110,6 @@ def get_libero_dummy_action():
OBS_STATE_DIM = 8
ACTION_DIM = 7
AGENT_POS_LOW = -1000.0
AGENT_POS_HIGH = 1000.0
ACTION_LOW = -1.0
ACTION_HIGH = 1.0
TASK_SUITE_MAX_STEPS: dict[str, int] = {
"libero_spatial": 280, # longest training demo has 193 steps
"libero_object": 280, # longest training demo has 254 steps
@@ -125,8 +148,8 @@ class LiberoEnv(gym.Env):
self.visualization_width = visualization_width
self.visualization_height = visualization_height
self.init_states = init_states
self.camera_name = _parse_camera_names(
camera_name
self.camera_name = camera_name.split(
","
) # agentview_image (main) or robot0_eye_in_hand_image (wrist)
# Map raw camera names to "image1" and "image2".
@@ -176,17 +199,15 @@ class LiberoEnv(gym.Env):
{
"pixels": spaces.Dict(images),
"agent_pos": spaces.Box(
low=AGENT_POS_LOW,
high=AGENT_POS_HIGH,
low=-1000.0,
high=1000.0,
shape=(OBS_STATE_DIM,),
dtype=np.float64,
),
}
)
self.action_space = spaces.Box(
low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32
)
self.action_space = spaces.Box(low=-1, high=1, shape=(ACTION_DIM,), dtype=np.float32)
def render(self):
raw_obs = self._env.env._get_observations()
@@ -291,6 +312,7 @@ def _make_env_fns(
gym_kwargs: Mapping[str, Any],
) -> list[Callable[[], LiberoEnv]]:
"""Build n_envs factory callables for a single (suite, task_id)."""
joined_cams = ",".join(camera_names) # keep backward-compat: downstream expects a string
def _make_env(episode_index: int, **kwargs) -> LiberoEnv:
local_kwargs = dict(kwargs)
@@ -298,7 +320,7 @@ def _make_env_fns(
task_suite=suite,
task_id=task_id,
task_suite_name=suite_name,
camera_name=camera_names,
camera_name=joined_cams,
init_states=init_states,
episode_index=episode_index,
**local_kwargs,
-1
View File
@@ -99,7 +99,6 @@ def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
policy_key = env_cfg.features_map[key]
policy_features[policy_key] = feature
return policy_features
+1 -1
View File
@@ -65,10 +65,10 @@ def make_act_pre_post_processors(
),
]
output_steps = [
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
@@ -73,10 +73,10 @@ def make_diffusion_pre_post_processors(
),
]
output_steps = [
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-4
View File
@@ -359,7 +359,6 @@ def make_policy(
if env_cfg is None:
raise ValueError("env_cfg cannot be None when ds_meta is not provided")
features = env_to_policy_features(env_cfg)
cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
kwargs["config"] = cfg
@@ -372,10 +371,7 @@ def make_policy(
else:
# Make a fresh policy.
policy = policy_cls(**kwargs)
policy.to(cfg.device)
assert isinstance(policy, torch.nn.Module)
# policy = torch.compile(policy, mode="reduce-overhead")
return policy
+1 -1
View File
@@ -146,10 +146,10 @@ def make_pi0_pre_post_processors(
]
output_steps: list[ProcessorStep] = [
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
@@ -73,10 +73,10 @@ def make_pi0fast_pre_post_processors(
),
]
output_steps = [
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+1 -3
View File
@@ -246,9 +246,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
base_model=base_model,
)
template_card = (
files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text(encoding="utf-8")
)
template_card = files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text()
card = ModelCard.from_template(card_data, template_str=template_card)
card.validate()
return card
+1 -1
View File
@@ -73,10 +73,10 @@ def make_sac_pre_post_processors(
),
]
output_steps = [
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
@@ -84,10 +84,10 @@ def make_smolvla_pre_post_processors(
),
]
output_steps = [
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
@@ -71,10 +71,10 @@ def make_tdmpc_pre_post_processors(
),
]
output_steps = [
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
@@ -72,10 +72,10 @@ def make_vqbet_pre_post_processors(
),
]
output_steps = [
DeviceProcessorStep(device="cpu"),
UnnormalizerProcessorStep(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
DeviceProcessorStep(device="cpu"),
]
return (
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+1 -10
View File
@@ -36,10 +36,7 @@ from .factory import (
make_default_robot_observation_processor,
make_default_teleop_action_processor,
)
from .gym_action_processor import (
Numpy2TorchActionProcessorStep,
Torch2NumpyActionProcessorStep,
)
from .gym_action_processor import Numpy2TorchActionProcessorStep, Torch2NumpyActionProcessorStep
from .hil_processor import (
AddTeleopActionAsComplimentaryDataStep,
AddTeleopEventsAsInfoStep,
@@ -70,10 +67,6 @@ from .pipeline import (
RobotProcessorPipeline,
TruncatedProcessorStep,
)
from .policy_robot_bridge import (
PolicyActionToRobotActionProcessorStep,
RobotActionToPolicyActionProcessorStep,
)
from .rename_processor import RenameObservationsProcessorStep
from .tokenizer_processor import TokenizerProcessorStep
@@ -123,8 +116,6 @@ __all__ = [
"RobotProcessorPipeline",
"TokenizerProcessorStep",
"Torch2NumpyActionProcessorStep",
"RobotActionToPolicyActionProcessorStep",
"PolicyActionToRobotActionProcessorStep",
"transition_to_batch",
"TransitionKey",
"TruncatedProcessorStep",
+8 -27
View File
@@ -207,37 +207,14 @@ def create_transition(
}
def robot_action_observation_to_transition(
action_observation: tuple[RobotAction, RobotObservation],
) -> EnvTransition:
"""
Convert a raw robot action and observation dictionary into a standardized `EnvTransition`.
Args:
action: The raw action dictionary from a teleoperation device or controller.
observation: The raw observation dictionary from the environment.
Returns:
An `EnvTransition` containing the formatted observation.
"""
if not isinstance(action_observation, tuple):
raise ValueError("action_observation should be a tuple type with an action and observation")
action, observation = action_observation
if action is not None and not isinstance(action, dict):
raise ValueError(f"Action should be a RobotAction type got {type(action)}")
if observation is not None and not isinstance(observation, dict):
raise ValueError(f"Observation should be a RobotObservation type got {type(observation)}")
return create_transition(action=action, observation=observation)
def robot_action_to_transition(action: RobotAction) -> EnvTransition:
"""
Convert a raw robot action dictionary into a standardized `EnvTransition`.
The keys in the action dictionary are prefixed with "action." and stored under
the `ACTION` key in the transition. Values are converted to tensors, except for
special types like `Rotation`.
Args:
action: The raw action dictionary from a teleoperation device or controller.
@@ -253,6 +230,10 @@ def observation_to_transition(observation: RobotObservation) -> EnvTransition:
"""
Convert a raw robot observation dictionary into a standardized `EnvTransition`.
The observation is split into state and image components. State keys are prefixed
with "observation.state." and image keys with "observation.images.". The result is
stored under the `OBSERVATION` key in the transition.
Args:
observation: The raw observation dictionary from the environment.
@@ -48,12 +48,12 @@ class MapTensorToDeltaActionDictStep(ActionProcessorStep):
# TODO (maractingi): add rotation
delta_action = {
"delta_x": action[0].item(),
"delta_y": action[1].item(),
"delta_z": action[2].item(),
"delta_x": action[0],
"delta_y": action[1],
"delta_z": action[2],
}
if self.use_gripper:
delta_action["gripper"] = action[3].item()
delta_action["gripper"] = action[3]
return delta_action
def transform_features(
@@ -126,7 +126,7 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
"target_wx": target_wx,
"target_wy": target_wy,
"target_wz": target_wz,
"gripper_vel": float(gripper),
"gripper": float(gripper),
}
return action
@@ -105,10 +105,6 @@ class DeviceProcessorStep(ProcessorStep):
# In both cases, use the configured device.
target_device = self.tensor_device
# MPS workaround: Convert float64 to float32 since MPS doesn't support float64
if target_device.type == "mps" and tensor.dtype == torch.float64:
tensor = tensor.to(dtype=torch.float32)
# Only move if necessary
if tensor.device != target_device:
tensor = tensor.to(target_device, non_blocking=self.non_blocking)
+7 -11
View File
@@ -16,7 +16,7 @@
from .converters import (
observation_to_transition,
robot_action_observation_to_transition,
robot_action_to_transition,
transition_to_observation,
transition_to_robot_action,
)
@@ -24,23 +24,19 @@ from .core import RobotAction, RobotObservation
from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
def make_default_teleop_action_processor() -> RobotProcessorPipeline[
tuple[RobotAction, RobotObservation], RobotAction
]:
teleop_action_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
def make_default_teleop_action_processor() -> RobotProcessorPipeline[RobotAction, RobotAction]:
teleop_action_processor = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[IdentityProcessorStep()],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
return teleop_action_processor
def make_default_robot_action_processor() -> RobotProcessorPipeline[
tuple[RobotAction, RobotObservation], RobotAction
]:
robot_action_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
def make_default_robot_action_processor() -> RobotProcessorPipeline[RobotAction, RobotAction]:
robot_action_processor = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[IdentityProcessorStep()],
to_transition=robot_action_observation_to_transition,
to_transition=robot_action_to_transition,
to_output=transition_to_robot_action,
)
return robot_action_processor
+17 -21
View File
@@ -19,8 +19,8 @@ from dataclasses import dataclass
from lerobot.configs.types import PipelineFeatureType, PolicyFeature
from .converters import to_tensor
from .core import EnvAction, EnvTransition, PolicyAction
from .pipeline import ActionProcessorStep, ProcessorStep, ProcessorStepRegistry
from .core import EnvAction, PolicyAction
from .pipeline import ActionProcessorStep, ProcessorStepRegistry
@ProcessorStepRegistry.register("torch2numpy_action_processor")
@@ -69,27 +69,23 @@ class Torch2NumpyActionProcessorStep(ActionProcessorStep):
@ProcessorStepRegistry.register("numpy2torch_action_processor")
@dataclass
class Numpy2TorchActionProcessorStep(ProcessorStep):
"""Converts a NumPy array action to a PyTorch tensor when action is present."""
class Numpy2TorchActionProcessorStep(ActionProcessorStep):
"""
Converts a NumPy array action to a PyTorch tensor.
def __call__(self, transition: EnvTransition) -> EnvTransition:
"""Converts numpy action to torch tensor if action exists, otherwise passes through."""
from .core import TransitionKey
This step is useful for converting actions from environments or hardware,
which are often NumPy arrays, into PyTorch tensors that can be processed
by a policy or model.
"""
self._current_transition = transition.copy()
new_transition = self._current_transition
action = new_transition.get(TransitionKey.ACTION)
if action is not None:
if not isinstance(action, EnvAction):
raise TypeError(
f"Expected np.ndarray or None, got {type(action).__name__}. "
"Use appropriate processor for non-tensor actions."
)
torch_action = to_tensor(action, dtype=None) # Preserve original dtype
new_transition[TransitionKey.ACTION] = torch_action
return new_transition
def action(self, action: EnvAction) -> PolicyAction:
if not isinstance(action, EnvAction):
raise TypeError(
f"Expected np.ndarray or None, got {type(action).__name__}. "
"Use appropriate processor for non-tensor actions."
)
torch_action = to_tensor(action, dtype=None) # Preserve original dtype
return torch_action
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+2 -7
View File
@@ -340,16 +340,11 @@ class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep):
"""
action = self.transition.get(TransitionKey.ACTION)
raw_joint_positions = complementary_data.get("raw_joint_positions", None)
if raw_joint_positions is None:
return complementary_data
current_gripper_pos = raw_joint_positions.get(GRIPPER_KEY, None)
current_gripper_pos = complementary_data.get("raw_joint_positions", None).get(GRIPPER_KEY, None)
if current_gripper_pos is None:
return complementary_data
# Gripper action is a PolicyAction at this stage
gripper_action = action[-1].item()
gripper_action = action[f"{GRIPPER_KEY}.pos"]
gripper_action_normalized = gripper_action / self.max_gripper_pos
# Normalize gripper state and action
+145 -693
View File
@@ -237,19 +237,6 @@ class ProcessorKwargs(TypedDict, total=False):
after_step_hooks: list[Callable[[int, EnvTransition], None]] | None
class ProcessorMigrationError(Exception):
"""Raised when a model needs migration to the processor format"""
def __init__(self, model_path: str | Path, migration_command: str, original_error: str):
self.model_path = model_path
self.migration_command = migration_command
self.original_error = original_error
super().__init__(
f"Model '{model_path}' requires migration to processor format. "
f"Run: {migration_command}\n\nOriginal error: {original_error}"
)
@dataclass
class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
"""A sequential pipeline for processing data, integrated with the Hugging Face Hub.
@@ -452,7 +439,6 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
def from_pretrained(
cls,
pretrained_model_name_or_path: str | Path,
config_filename: str,
*,
force_download: bool = False,
resume_download: bool | None = None,
@@ -461,74 +447,24 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
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,
to_transition: Callable[[TInput], EnvTransition] | None = None,
to_output: Callable[[EnvTransition], TOutput] | None = None,
**kwargs,
) -> DataProcessorPipeline[TInput, TOutput]:
"""Loads a pipeline from a local directory, single file, or Hugging Face Hub repository.
"""Loads a pipeline from a local directory or a Hugging Face Hub repository.
This method implements a simplified loading pipeline with intelligent migration detection:
**Simplified Loading Strategy**:
1. **Config Loading** (_load_config):
- **Directory**: Load specified config_filename from directory
- **Single file**: Load file directly (config_filename ignored)
- **Hub repository**: Download specified config_filename from Hub
2. **Config Validation** (_validate_loaded_config):
- Format validation: Ensure config is valid processor format
- Migration detection: Guide users to migrate old LeRobot models
- Clear errors: Provide actionable error messages
3. **Step Construction** (_build_steps_with_overrides):
- Class resolution: Registry lookup or dynamic imports
- Override merging: User parameters override saved config
- State loading: Load .safetensors files for stateful steps
4. **Override Validation** (_validate_overrides_used):
- Ensure all user overrides were applied (catch typos)
- Provide helpful error messages with available keys
**Migration Detection**:
- **Smart detection**: Analyzes JSON files to detect old LeRobot models
- **Precise targeting**: Avoids false positives on other HuggingFace models
- **Clear guidance**: Provides exact migration command to run
- **Error mode**: Always raises ProcessorMigrationError for clear user action
**Loading Examples**:
```python
# Directory loading
pipeline = DataProcessorPipeline.from_pretrained("/models/my_model", config_filename="processor.json")
# Single file loading
pipeline = DataProcessorPipeline.from_pretrained(
"/models/my_model/processor.json", config_filename="processor.json"
)
# Hub loading
pipeline = DataProcessorPipeline.from_pretrained("user/repo", config_filename="processor.json")
# Multiple configs (preprocessor/postprocessor)
preprocessor = DataProcessorPipeline.from_pretrained(
"model", config_filename="policy_preprocessor.json"
)
postprocessor = DataProcessorPipeline.from_pretrained(
"model", config_filename="policy_postprocessor.json"
)
```
**Override System**:
- **Key matching**: Use registry names or class names as override keys
- **Config merging**: User overrides take precedence over saved config
- **Validation**: Ensure all override keys match actual steps (catch typos)
- **Example**: overrides={"NormalizeStep": {"device": "cuda"}}
This method reconstructs a `DataProcessorPipeline` by:
1. Loading the main JSON configuration file.
2. Iterating through the steps defined in the config.
3. Dynamically importing or looking up each step's class.
4. Instantiating each step with its saved configuration, potentially with overrides.
5. Loading the step's state from its `.safetensors` file, if it exists.
Args:
pretrained_model_name_or_path: The identifier of the repository on the Hugging Face Hub,
a path to a local directory, or a path to a single config file.
config_filename: The name of the pipeline's JSON configuration file. Always required
to prevent ambiguity when multiple configs exist (e.g., preprocessor vs postprocessor).
pretrained_model_name_or_path: The identifier of the repository on the Hugging Face Hub
or a path to a local directory.
force_download: Whether to force (re)downloading the files.
resume_download: Whether to resume a previously interrupted download.
proxies: A dictionary of proxy servers to use.
@@ -536,6 +472,9 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
cache_dir: The path to a specific cache folder to store downloaded files.
local_files_only: If True, avoid downloading files from the Hub.
revision: The specific model version to use (e.g., a branch name, tag name, or commit id).
config_filename: The name of the pipeline's JSON configuration file. If not provided,
it's auto-detected in local directories (if only one .json file exists). This parameter
is mandatory when loading from Hugging Face Hub repositories.
overrides: A dictionary to override the configuration of specific steps. Keys should
match the step's class name or registry name.
to_transition: A custom function to convert input data to `EnvTransition`.
@@ -550,667 +489,180 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
ValueError: If configuration is ambiguous or instantiation fails.
ImportError: If a step's class cannot be imported.
KeyError: If an override key doesn't match any step in the pipeline.
ProcessorMigrationError: If the model requires migration to processor format.
"""
model_id = str(pretrained_model_name_or_path)
hub_download_kwargs = {
"force_download": force_download,
"resume_download": resume_download,
"proxies": proxies,
"token": token,
"cache_dir": cache_dir,
"local_files_only": local_files_only,
"revision": revision,
}
loaded_config: dict[str, Any] | None = None
base_path: Path | None = None
# 1. Load configuration using simplified 3-way logic
loaded_config, base_path = cls._load_config(model_id, config_filename, hub_download_kwargs)
# Standard pattern: try local directory first
if Path(model_id).is_dir():
base_path = Path(model_id)
# 2. Validate configuration and handle migration
cls._validate_loaded_config(model_id, loaded_config, config_filename)
# Handle config filename
if config_filename is None:
json_files = list(base_path.glob("*.json"))
if len(json_files) == 0:
# No config files found locally, will try Hub next
pass
elif len(json_files) == 1:
config_filename = json_files[0].name
else:
raise ValueError(
f"Multiple .json files found in {model_id}: {[f.name for f in json_files]}. "
f"Please specify which one to load using the config_filename parameter."
)
# 3. Build steps with overrides
steps, validated_overrides = cls._build_steps_with_overrides(
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs
)
# Try to load config from local directory
if config_filename and (base_path / config_filename).exists():
with open(base_path / config_filename) as f:
loaded_config = json.load(f)
# 4. Validate that all overrides were used
cls._validate_overrides_used(validated_overrides, loaded_config)
# If not found locally, try Hub
if loaded_config is None:
# Check if this looks like a local path that doesn't exist
# Hub repo IDs have format "user/repo" with exactly one slash
# Local paths typically have multiple slashes, backslashes, or start with ./ or ../
looks_like_local_path = (
model_id.count("/") > 1 # Multiple slashes suggest local path
or "\\" in model_id # Backslashes are only in local paths
or Path(model_id).is_absolute() # Absolute paths are local
or model_id.startswith("./")
or model_id.startswith("../") # Relative path indicators
)
# 5. Construct and return the final pipeline instance
return cls(
steps=steps,
name=loaded_config.get("name", "DataProcessorPipeline"),
to_transition=to_transition or cast(Callable[[TInput], EnvTransition], batch_to_transition),
to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch),
)
@classmethod
def _load_config(
cls,
model_id: str,
config_filename: str,
hub_download_kwargs: dict[str, Any],
) -> tuple[dict[str, Any], Path]:
"""Load configuration from local file or Hugging Face Hub.
This method implements a super-simplified 3-way loading strategy:
1. **Local directory**: Load config_filename from directory
- Example: model_id="/models/my_model", config_filename="processor.json"
- Loads: "/models/my_model/processor.json"
2. **Single file**: Load file directly (ignore config_filename)
- Example: model_id="/models/my_model/processor.json"
- Loads: "/models/my_model/processor.json" (config_filename ignored)
3. **Hub repository**: Download config_filename from Hub
- Example: model_id="user/repo", config_filename="processor.json"
- Downloads and loads: config_filename from Hub repo
**Benefits of Explicit config_filename**:
- No auto-detection complexity or edge cases
- No risk of loading wrong config (preprocessor vs postprocessor)
- Consistent behavior across local and Hub usage
- Clear, predictable errors
Args:
model_id: The model identifier (Hub repo ID, local directory, or file path)
config_filename: The explicit config filename to load (always required)
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
Returns:
Tuple of (loaded_config, base_path)
- loaded_config: Parsed JSON config dict (always loaded, never None)
- base_path: Directory containing config file (for state file resolution)
Raises:
FileNotFoundError: If config file cannot be found locally or on Hub
"""
model_path = Path(model_id)
if model_path.is_dir():
# Directory: load specified config from directory
config_path = model_path / config_filename
if not config_path.exists():
# Check for migration before giving clear error
if cls._should_suggest_migration(model_path):
cls._suggest_processor_migration(model_id, f"Config file '{config_filename}' not found")
raise FileNotFoundError(
f"Config file '{config_filename}' not found in directory '{model_id}'"
if looks_like_local_path:
# This appears to be a local path that doesn't exist
raise FileNotFoundError(f"Local path '{model_id}' does not exist")
# For Hub repositories, config_filename is mandatory
if config_filename is None:
raise ValueError(
f"When loading from Hugging Face Hub, 'config_filename' must be specified. "
f"Example: DataProcessorPipeline.from_pretrained('{model_id}', config_filename='processor.json')"
)
with open(config_path) as f:
return json.load(f), model_path
elif model_path.is_file():
# File: load file directly (config_filename is ignored for single files)
with open(model_path) as f:
return json.load(f), model_path.parent
else:
# Hub: download specified config
try:
# Download the configuration file from the Hub
config_path = hf_hub_download(
repo_id=model_id,
filename=config_filename,
repo_type="model",
**hub_download_kwargs,
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 f:
return json.load(f), Path(config_path).parent
loaded_config = json.load(f)
# The base path for other files (like state tensors) is the directory of the config file
base_path = Path(config_path).parent
except Exception as e:
raise FileNotFoundError(
f"Could not find '{config_filename}' on the HuggingFace Hub at '{model_id}'"
f"Could not find {config_filename} on the HuggingFace Hub at {model_id}"
) from e
@classmethod
def _validate_loaded_config(
cls, model_id: str, loaded_config: dict[str, Any], config_filename: str
) -> None:
"""Validate that a config was loaded and is a valid processor config.
# At this point, loaded_config must be loaded successfully
if loaded_config is None:
raise RuntimeError("Failed to load configuration from local directory or Hub")
This method validates processor config format with intelligent migration detection:
if overrides is None:
overrides = {}
**Config Format Validation**:
- Use _is_processor_config() to validate structure
- Must have "steps" field with list of step configurations
- Each step needs "class" or "registry_name"
- If validation fails AND local directory: Check for migration need
- If migration needed: Raise ProcessorMigrationError with command
- If no migration: Raise ValueError with helpful error message
**Migration Detection Logic**:
- Only triggered for local directories (not Hub repos)
- Analyzes all JSON files in directory to detect old LeRobot models
- Provides exact migration command with model path
Args:
model_id: The model identifier (used for migration detection)
loaded_config: The loaded config dictionary (guaranteed non-None)
config_filename: The config filename that was loaded (for error messages)
Raises:
ValueError: If config format is invalid
ProcessorMigrationError: If model needs migration to processor format
"""
# Validate that this is actually a processor config
if not cls._is_processor_config(loaded_config):
if Path(model_id).is_dir() and cls._should_suggest_migration(Path(model_id)):
cls._suggest_processor_migration(
model_id,
f"Config file '{config_filename}' is not a valid processor configuration",
)
raise ValueError(
f"Config file '{config_filename}' is not a valid processor configuration. "
f"Expected a config with 'steps' field, but got: {list(loaded_config.keys())}"
)
@classmethod
def _build_steps_with_overrides(
cls,
loaded_config: dict[str, Any],
overrides: dict[str, Any],
model_id: str,
base_path: Path | None,
hub_download_kwargs: dict[str, Any],
) -> tuple[list[ProcessorStep], set[str]]:
"""Build all processor steps with overrides and state loading.
This method orchestrates the complete step construction pipeline:
**For each step in loaded_config["steps"]**:
1. **Class Resolution** (via _resolve_step_class):
- **If "registry_name" exists**: Look up in ProcessorStepRegistry
Example: {"registry_name": "normalize_step"} -> Get registered class
- **Else use "class" field**: Dynamic import from full module path
Example: {"class": "lerobot.processor.normalize.NormalizeStep"}
- **Result**: (step_class, step_key) where step_key is used for overrides
2. **Step Instantiation** (via _instantiate_step):
- **Merge configs**: saved_config + user_overrides
- **Override priority**: User overrides take precedence over saved config
- **Example**: saved={"mean": 0.0}, override={"mean": 1.0} -> final={"mean": 1.0}
- **Result**: Instantiated ProcessorStep object
3. **State Loading** (via _load_step_state):
- **If step has "state_file"**: Load tensor state from .safetensors
- **Local first**: Check base_path/state_file.safetensors
- **Hub fallback**: Download state file if not found locally
- **Optional**: Only load if step has load_state_dict method
4. **Override Tracking**:
- **Track used overrides**: Remove step_key from remaining set
- **Purpose**: Validate all user overrides were applied (detect typos)
**Error Handling**:
- Class resolution errors -> ImportError with helpful message
- Instantiation errors -> ValueError with config details
- State loading errors -> Propagated from load_state_dict
Args:
loaded_config: The loaded processor configuration (must have "steps" field)
overrides: User-provided parameter overrides (keyed by class/registry name)
model_id: The model identifier (needed for Hub state file downloads)
base_path: Local directory path for finding state files
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
Returns:
Tuple of (instantiated_steps_list, unused_override_keys)
- instantiated_steps_list: List of ready-to-use ProcessorStep instances
- unused_override_keys: Override keys that didn't match any step (for validation)
Raises:
ImportError: If a step class cannot be imported or found in registry
ValueError: If a step cannot be instantiated with its configuration
"""
steps: list[ProcessorStep] = []
override_keys = set(overrides.keys())
steps: list[ProcessorStep] = []
for step_entry in loaded_config["steps"]:
# 1. Get step class and key
step_class, step_key = cls._resolve_step_class(step_entry)
# Determine the step class, prioritizing the registry.
if "registry_name" in step_entry:
try:
step_class = ProcessorStepRegistry.get(step_entry["registry_name"])
step_key = step_entry["registry_name"]
except KeyError as e:
raise ImportError(f"Failed to load processor step from registry. {str(e)}") from e
else:
# Fallback to dynamic import using the full class path.
full_class_path = step_entry["class"]
module_path, class_name = full_class_path.rsplit(".", 1)
# 2. Instantiate step with overrides
step_instance = cls._instantiate_step(step_entry, step_class, step_key, overrides)
try:
module = importlib.import_module(module_path)
step_class = getattr(module, class_name)
step_key = class_name
except (ImportError, AttributeError) as e:
raise ImportError(
f"Failed to load processor step '{full_class_path}'. "
f"Make sure the module '{module_path}' is installed and contains class '{class_name}'. "
f"Consider registering the step using @ProcessorStepRegistry.register() for better portability. "
f"Error: {str(e)}"
) from e
# 3. Load step state if available
cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
# 4. Track used overrides
if step_key in override_keys:
override_keys.discard(step_key)
steps.append(step_instance)
return steps, override_keys
@classmethod
def _resolve_step_class(cls, step_entry: dict[str, Any]) -> tuple[type[ProcessorStep], str]:
"""Resolve step class from registry or import path.
This method implements a two-tier resolution strategy:
**Tier 1: Registry-based resolution** (preferred):
- **If "registry_name" in step_entry**: Look up in ProcessorStepRegistry
- **Advantage**: Faster, no imports needed, guaranteed compatibility
- **Example**: {"registry_name": "normalize_step"} -> Get pre-registered class
- **Error**: KeyError if registry_name not found -> Convert to ImportError
**Tier 2: Dynamic import fallback**:
- **Else use "class" field**: Full module.ClassName import path
- **Process**: Split "module.path.ClassName" into module + class parts
- **Import**: Use importlib.import_module() + getattr()
- **Example**: "lerobot.processor.normalize.NormalizeStep"
a. Import module: "lerobot.processor.normalize"
b. Get class: getattr(module, "NormalizeStep")
- **step_key**: Use class_name ("NormalizeStep") for overrides
**Override Key Strategy**:
- Registry steps: Use registry_name ("normalize_step")
- Import steps: Use class_name ("NormalizeStep")
- This allows users to override with: {"normalize_step": {...}} or {"NormalizeStep": {...}}
**Error Handling**:
- Registry KeyError -> ImportError with registry context
- Import/Attribute errors -> ImportError with helpful suggestions
- All errors include troubleshooting guidance
Args:
step_entry: The step configuration dictionary (must have "registry_name" or "class")
Returns:
Tuple of (step_class, step_key)
- step_class: The resolved ProcessorStep class (ready for instantiation)
- step_key: The key used for user overrides (registry_name or class_name)
Raises:
ImportError: If step class cannot be loaded from registry or import path
"""
if "registry_name" in step_entry:
# Instantiate the step, merging saved config with user-provided overrides.
try:
step_class = ProcessorStepRegistry.get(step_entry["registry_name"])
return step_class, step_entry["registry_name"]
except KeyError as e:
raise ImportError(f"Failed to load processor step from registry. {str(e)}") from e
else:
# Fallback to dynamic import using the full class path
full_class_path = step_entry["class"]
module_path, class_name = full_class_path.rsplit(".", 1)
saved_cfg = step_entry.get("config", {})
step_overrides = overrides.get(step_key, {})
merged_cfg = {**saved_cfg, **step_overrides}
step_instance: ProcessorStep = step_class(**merged_cfg)
try:
module = importlib.import_module(module_path)
step_class = getattr(module, class_name)
return step_class, class_name
except (ImportError, AttributeError) as e:
raise ImportError(
f"Failed to load processor step '{full_class_path}'. "
f"Make sure the module '{module_path}' is installed and contains class '{class_name}'. "
f"Consider registering the step using @ProcessorStepRegistry.register() for better portability. "
if step_key in override_keys:
override_keys.discard(step_key)
except Exception as e:
step_name = step_entry.get("registry_name", step_entry.get("class", "Unknown"))
raise ValueError(
f"Failed to instantiate processor step '{step_name}' with config: {step_entry.get('config', {})}. "
f"Error: {str(e)}"
) from e
@classmethod
def _instantiate_step(
cls,
step_entry: dict[str, Any],
step_class: type[ProcessorStep],
step_key: str,
overrides: dict[str, Any],
) -> ProcessorStep:
"""Instantiate a single processor step with config overrides.
# Load the step's state if a state file is specified.
if "state_file" in step_entry and hasattr(step_instance, "load_state_dict"):
# Check if state file exists locally first
if base_path and (base_path / step_entry["state_file"]).exists():
state_path = str(base_path / step_entry["state_file"])
else:
# Download the state file from the Hub.
state_path = hf_hub_download(
repo_id=model_id,
filename=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,
)
This method handles the configuration merging and instantiation logic:
step_instance.load_state_dict(load_file(state_path))
**Configuration Merging Strategy**:
1. **Extract saved config**: Get step_entry.get("config", {}) from saved pipeline
- Example: {"config": {"mean": 0.0, "std": 1.0}}
2. **Extract user overrides**: Get overrides.get(step_key, {}) for this step
- Example: overrides = {"NormalizeStep": {"mean": 2.0, "device": "cuda"}}
3. **Merge with priority**: {**saved_cfg, **step_overrides}
- **Override priority**: User values override saved values
- **Result**: {"mean": 2.0, "std": 1.0, "device": "cuda"}
steps.append(step_instance)
**Instantiation Process**:
- **Call constructor**: step_class(**merged_cfg)
- **Example**: NormalizeStep(mean=2.0, std=1.0, device="cuda")
# Check for any unused override keys, which likely indicates a typo by the user.
if override_keys:
available_keys = [
step.get("registry_name") or step["class"].rsplit(".", 1)[1]
for step in loaded_config["steps"]
]
**Error Handling**:
- **Any exception during instantiation**: Convert to ValueError
- **Include context**: step name, attempted config, original error
- **Purpose**: Help users debug configuration issues
- **Common causes**:
a. Invalid parameter types (str instead of float)
b. Missing required parameters
c. Incompatible parameter combinations
Args:
step_entry: The step configuration from saved config (contains "config" dict)
step_class: The step class to instantiate (already resolved)
step_key: The key used for overrides ("registry_name" or class name)
overrides: User-provided parameter overrides (keyed by step_key)
Returns:
The instantiated processor step (ready for use)
Raises:
ValueError: If step cannot be instantiated, with detailed error context
"""
try:
saved_cfg = step_entry.get("config", {})
step_overrides = overrides.get(step_key, {})
merged_cfg = {**saved_cfg, **step_overrides}
return step_class(**merged_cfg)
except Exception as e:
step_name = step_entry.get("registry_name", step_entry.get("class", "Unknown"))
raise ValueError(
f"Failed to instantiate processor step '{step_name}' with config: {step_entry.get('config', {})}. "
f"Error: {str(e)}"
) from e
@classmethod
def _load_step_state(
cls,
step_instance: ProcessorStep,
step_entry: dict[str, Any],
model_id: str,
base_path: Path | None,
hub_download_kwargs: dict[str, Any],
) -> None:
"""Load state dictionary for a processor step if available.
This method implements conditional state loading with local/Hub fallback:
**Precondition Checks** (early return if not met):
1. **"state_file" in step_entry**: Step config specifies a state file
- **If missing**: Step has no saved state (e.g., stateless transforms)
2. **hasattr(step_instance, "load_state_dict")**: Step supports state loading
- **If missing**: Step doesn't implement state loading (rare)
**State File Resolution Strategy**:
1. **Local file priority**: Check base_path/state_filename exists
- **Advantage**: Faster, no network calls
- **Example**: "/models/my_model/normalize_step_0.safetensors"
- **Use case**: Loading from local saved model directory
2. **Hub download fallback**: Download state file from repository
- **When triggered**: Local file not found or base_path is None
- **Process**: Use hf_hub_download with same parameters as config
- **Example**: Download "normalize_step_0.safetensors" from "user/repo"
- **Result**: Downloaded to local cache, path returned
**State Loading Process**:
- **Load tensors**: Use safetensors.torch.load_file()
- **Apply to step**: Call step_instance.load_state_dict(tensor_dict)
- **In-place modification**: Updates step's internal tensor state
**Common state file examples**:
- "normalize_step_0.safetensors" - normalization statistics
- "custom_step_1.safetensors" - learned parameters
- "tokenizer_step_2.safetensors" - vocabulary embeddings
Args:
step_instance: The step instance to load state into (must have load_state_dict)
step_entry: The step configuration dictionary (may contain "state_file")
model_id: The model identifier (used for Hub downloads if needed)
base_path: Local directory path for finding state files (None for Hub-only)
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
Note:
This method modifies step_instance in-place and returns None.
If state loading fails, exceptions from load_state_dict propagate.
"""
if "state_file" not in step_entry or not hasattr(step_instance, "load_state_dict"):
return
state_filename = step_entry["state_file"]
# Try local file first
if base_path and (base_path / state_filename).exists():
state_path = str(base_path / state_filename)
else:
# Download from Hub
state_path = hf_hub_download(
repo_id=model_id,
filename=state_filename,
repo_type="model",
**hub_download_kwargs,
raise KeyError(
f"Override keys {list(override_keys)} do not match any step in the saved configuration. "
f"Available step keys: {available_keys}. "
f"Make sure override keys match exact step class names or registry names."
)
step_instance.load_state_dict(load_file(state_path))
@classmethod
def _validate_overrides_used(
cls, remaining_override_keys: set[str], loaded_config: dict[str, Any]
) -> None:
"""Validate that all provided overrides were used.
This method ensures user overrides are valid to catch typos and configuration errors:
**Validation Logic**:
1. **If remaining_override_keys is empty**: All overrides were used -> Success
- **Early return**: No validation needed
- **Normal case**: User provided correct override keys
2. **If remaining_override_keys has entries**: Some overrides unused -> Error
- **Root cause**: User provided keys that don't match any step
- **Common issues**:
a. Typos in step names ("NormalizStep" vs "NormalizeStep")
b. Using wrong key type (class name vs registry name)
c. Step doesn't exist in saved pipeline
**Helpful Error Generation**:
- **Extract available keys**: Build list of valid override keys from config
a. **Registry steps**: Use "registry_name" directly
b. **Import steps**: Extract class name from "class" field
- Example: "lerobot.processor.normalize.NormalizeStep" -> "NormalizeStep"
- **Error message includes**:
a. Invalid keys provided by user
b. List of valid keys they can use
c. Guidance about registry vs class names
**Override Key Resolution Rules**:
- Steps with "registry_name": Use registry_name for overrides
- Steps with "class": Use final class name for overrides
- Users must match these exact keys in their overrides dict
Args:
remaining_override_keys: Override keys that weren't matched to any step
loaded_config: The loaded processor configuration (contains "steps" list)
Raises:
KeyError: If any override keys were not used, with helpful error message
"""
if not remaining_override_keys:
return
available_keys = [
step.get("registry_name") or step["class"].rsplit(".", 1)[1] for step in loaded_config["steps"]
]
raise KeyError(
f"Override keys {list(remaining_override_keys)} do not match any step in the saved configuration. "
f"Available step keys: {available_keys}. "
f"Make sure override keys match exact step class names or registry names."
# Construct and return the final pipeline instance.
return cls(
steps=steps,
name=loaded_config.get("name", "DataProcessorPipeline"),
to_transition=to_transition or batch_to_transition,
to_output=to_output or cast(Callable[[EnvTransition], TOutput], transition_to_batch),
)
@classmethod
def _should_suggest_migration(cls, model_path: Path) -> bool:
"""Check if directory has JSON files but no processor configs.
This method implements smart migration detection to avoid false positives:
**Decision Logic**:
1. **No JSON files found**: Return False
- **Reason**: Empty directory or only non-config files
- **Example**: Directory with only .safetensors, .md files
- **Action**: No migration needed
2. **JSON files exist**: Analyze each file
- **Goal**: Determine if ANY file is a valid processor config
- **Process**:
a. Try to parse each .json file
b. Skip files with JSON parse errors (malformed)
c. Check if parsed config passes _is_processor_config()
- **If ANY valid processor found**: Return False (no migration)
- **If NO valid processors found**: Return True (migration needed)
**Examples**:
- **No migration**: ["processor.json", "config.json"] where processor.json is valid
- **Migration needed**: ["config.json", "train.json"] where both are model configs
- **No migration**: [] (empty directory)
- **Migration needed**: ["old_model_config.json"] with old LeRobot format
**Why this works**:
- **Precise detection**: Only suggests migration for actual old LeRobot models
- **Avoids false positives**: Won't trigger on other HuggingFace model types
- **Graceful handling**: Ignores malformed JSON files
Args:
model_path: Path to local directory to analyze
Returns:
True if directory has JSON configs but none are processor configs (migration needed)
False if no JSON files or at least one valid processor config exists
"""
json_files = list(model_path.glob("*.json"))
if len(json_files) == 0:
return False
# Check if any JSON file is a processor config
for json_file in json_files:
try:
with open(json_file) as f:
config = json.load(f)
if cls._is_processor_config(config):
return False # Found at least one processor config, no migration needed
except (json.JSONDecodeError, OSError):
# Skip files that can't be parsed as JSON
continue
# Have JSON files but no processor configs - suggest migration
return True
@classmethod
def _is_processor_config(cls, config: dict) -> bool:
"""Check if config follows DataProcessorPipeline format.
This method validates the processor configuration structure:
**Required Structure Validation**:
1. **"steps" field existence**: Must have top-level "steps" key
- **If missing**: Not a processor config (e.g., model config, train config)
- **Example invalid**: {"type": "act", "hidden_dim": 256}
2. **"steps" field type**: Must be a list, not other types
- **If not list**: Invalid format
- **Example invalid**: {"steps": "some_string"} or {"steps": {"key": "value"}}
3. **Empty steps validation**: Empty list is valid
- **If len(steps) == 0**: Return True immediately
- **Use case**: Empty processor pipeline (no-op)
- **Example valid**: {"name": "EmptyProcessor", "steps": []}
**Individual Step Validation** (for non-empty steps):
For each step in the steps list:
1. **Step type**: Must be a dictionary
- **If not dict**: Invalid step format
- **Example invalid**: ["string_step", 123, true]
2. **Step identifier**: Must have either "class" OR "registry_name"
- **"registry_name"**: Registered step (preferred)
Example: {"registry_name": "normalize_step", "config": {...}}
- **"class"**: Full import path
Example: {"class": "lerobot.processor.normalize.NormalizeStep"}
- **If neither**: Invalid step (can't resolve class)
- **If both**: Also valid (registry_name takes precedence)
**Valid Processor Config Examples**:
- {"steps": []} - Empty processor
- {"steps": [{"registry_name": "normalize"}]} - Registry step
- {"steps": [{"class": "my.module.Step"}]} - Import step
- {"name": "MyProcessor", "steps": [...]} - With name
**Invalid Config Examples**:
- {"type": "act"} - Missing "steps"
- {"steps": "normalize"} - Steps not a list
- {"steps": [{}]} - Step missing class/registry_name
- {"steps": ["string"]} - Step not a dict
Args:
config: The configuration dictionary to validate
Returns:
True if config follows valid DataProcessorPipeline format, False otherwise
"""
# Must have a "steps" field with a list of step configurations
if not isinstance(config.get("steps"), list):
return False
steps = config["steps"]
if len(steps) == 0:
return True # Empty processor is valid
# Each step must be a dict with either "class" or "registry_name"
for step in steps:
if not isinstance(step, dict):
return False
if not ("class" in step or "registry_name" in step):
return False
return True
@classmethod
def _suggest_processor_migration(cls, model_path: str | Path, original_error: str) -> None:
"""Raise migration error when we detect JSON files but no processor configs.
This method is called when migration detection determines that a model
directory contains configuration files but none are valid processor configs.
This typically indicates an old LeRobot model that needs migration.
**When this is called**:
- User tries to load DataProcessorPipeline from local directory
- Directory contains JSON configuration files
- None of the JSON files follow processor config format
- _should_suggest_migration() returned True
**Migration Command Generation**:
- Constructs exact command user needs to run
- Uses the migration script: migrate_policy_normalization.py
- Includes the model path automatically
- Example: "python src/lerobot/processor/migrate_policy_normalization.py --pretrained-path /models/old_model"
**Error Structure**:
- **Always raises**: ProcessorMigrationError (never returns)
- **Includes**: model_path, migration_command, original_error
- **Purpose**: Force user attention to migration need
- **User experience**: Clear actionable error with exact command to run
**Migration Process**:
The suggested command will:
1. Extract normalization stats from old model
2. Create new processor configs (preprocessor + postprocessor)
3. Remove normalization layers from model
4. Save migrated model with processor pipeline
Args:
model_path: Path to the model directory needing migration
original_error: The error that triggered migration detection (for context)
Raises:
ProcessorMigrationError: Always raised (this method never returns normally)
"""
migration_command = (
f"python src/lerobot/processor/migrate_policy_normalization.py --pretrained-path {model_path}"
)
raise ProcessorMigrationError(model_path, migration_command, original_error)
def __len__(self) -> int:
"""Returns the number of steps in the pipeline."""
return len(self.steps)
@@ -1,52 +0,0 @@
from dataclasses import asdict, dataclass
from typing import Any
import torch
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.processor import ActionProcessorStep, PolicyAction, ProcessorStepRegistry, RobotAction
@dataclass
@ProcessorStepRegistry.register("robot_action_to_policy_action_processor")
class RobotActionToPolicyActionProcessorStep(ActionProcessorStep):
"""Processor step to map a dictionary to a tensor action."""
motor_names: list[str]
def action(self, action: RobotAction) -> PolicyAction:
if len(self.motor_names) != len(action):
raise ValueError(f"Action must have {len(self.motor_names)} elements, got {len(action)}")
return torch.tensor([action[f"{name}.pos"] for name in self.motor_names])
def get_config(self) -> dict[str, Any]:
return asdict(self)
def transform_features(self, features):
features[PipelineFeatureType.ACTION]["action"] = PolicyFeature(
type=FeatureType.ACTION, shape=(len(self.motor_names),)
)
return features
@dataclass
@ProcessorStepRegistry.register("policy_action_to_robot_action_processor")
class PolicyActionToRobotActionProcessorStep(ActionProcessorStep):
"""Processor step to map a policy action to a robot action."""
motor_names: list[str]
def action(self, action: PolicyAction) -> RobotAction:
if len(self.motor_names) != len(action):
raise ValueError(f"Action must have {len(self.motor_names)} elements, got {len(action)}")
return {f"{name}.pos": action[i] for i, name in enumerate(self.motor_names)}
def get_config(self) -> dict[str, Any]:
return asdict(self)
def transform_features(self, features):
for name in self.motor_names:
features[PipelineFeatureType.ACTION][f"{name}.pos"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
)
return features
+8 -13
View File
@@ -21,12 +21,11 @@ Example:
lerobot-record \
--robot.type=so100_follower \
--robot.port=/dev/tty.usbmodem58760431541 \
--robot.cameras="{laptop: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
--robot.cameras="{laptop: {type: opencv, camera_index: 0, width: 640, height: 480}}" \
--robot.id=black \
--dataset.repo_id=<my_username>/<my_dataset_name> \
--dataset.repo_id=aliberts/record-test \
--dataset.num_episodes=2 \
--dataset.single_task="Grab the cube" \
--display_data=true
# <- Teleop optional if you want to teleoperate to record or in between episodes with a policy \
# --teleop.type=so100_leader \
# --teleop.port=/dev/tty.usbmodem58760431551 \
@@ -236,12 +235,8 @@ def record_loop(
robot: Robot,
events: dict,
fps: int,
teleop_action_processor: RobotProcessorPipeline[
tuple[RobotAction, RobotObservation], RobotAction
], # runs after teleop
robot_action_processor: RobotProcessorPipeline[
tuple[RobotAction, RobotObservation], RobotAction
], # runs before robot
teleop_action_processor: RobotProcessorPipeline[RobotAction, RobotAction], # runs after teleop
robot_action_processor: RobotProcessorPipeline[RobotAction, RobotAction], # runs before robot
robot_observation_processor: RobotProcessorPipeline[
RobotObservation, RobotObservation
], # runs after robot
@@ -327,7 +322,7 @@ def record_loop(
act = teleop.get_action()
# Applies a pipeline to the raw teleop action, default is IdentityProcessor
act_processed_teleop = teleop_action_processor((act, obs))
act_processed_teleop = teleop_action_processor(act)
elif policy is None and isinstance(teleop, list):
arm_action = teleop_arm.get_action()
@@ -335,7 +330,7 @@ def record_loop(
keyboard_action = teleop_keyboard.get_action()
base_action = robot._from_keyboard_to_base_action(keyboard_action)
act = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
act_processed_teleop = teleop_action_processor((act, obs))
act_processed_teleop = teleop_action_processor(act)
else:
logging.info(
"No policy or teleoperator provided, skipping action generation."
@@ -347,10 +342,10 @@ def record_loop(
# Applies a pipeline to the action, default is IdentityProcessor
if policy is not None and act_processed_policy is not None:
action_values = act_processed_policy
robot_action_to_send = robot_action_processor((act_processed_policy, obs))
robot_action_to_send = robot_action_processor(act_processed_policy)
else:
action_values = act_processed_teleop
robot_action_to_send = robot_action_processor((act_processed_teleop, obs))
robot_action_to_send = robot_action_processor(act_processed_teleop)
# Send action to robot
# Action can eventually be clipped using `max_relative_target`,
+3 -4
View File
@@ -23,7 +23,7 @@ lerobot-replay \
--robot.port=/dev/tty.usbmodem58760431541 \
--robot.id=black \
--dataset.repo_id=aliberts/record-test \
--dataset.episode=0
--dataset.episode=2
```
Example replay with bimanual so100:
@@ -57,6 +57,7 @@ from lerobot.robots import ( # noqa: F401
hope_jr,
koch_follower,
make_robot_from_config,
reachy2,
so100_follower,
so101_follower,
)
@@ -112,9 +113,7 @@ def replay(cfg: ReplayConfig):
for i, name in enumerate(dataset.features["action"]["names"]):
action[name] = action_array[i]
robot_obs = robot.get_observation()
processed_action = robot_action_processor((action, robot_obs))
processed_action = robot_action_processor(action)
_ = robot.send_action(processed_action)
@@ -22,6 +22,7 @@ import numpy as np
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.model.kinematics import RobotKinematics
from lerobot.processor import (
ComplementaryDataProcessorStep,
EnvTransition,
ObservationProcessorStep,
ProcessorStep,
@@ -30,6 +31,7 @@ from lerobot.processor import (
RobotActionProcessorStep,
TransitionKey,
)
from lerobot.robots.robot import Robot
from lerobot.utils.rotation import Rotation
@@ -66,46 +68,38 @@ class EEReferenceAndDelta(RobotActionProcessorStep):
use_latched_reference: bool = (
True # If True, latch reference on enable; if False, always use current pose
)
use_ik_solution: bool = False
reference_ee_pose: np.ndarray | None = field(default=None, init=False, repr=False)
_prev_enabled: bool = field(default=False, init=False, repr=False)
_command_when_disabled: np.ndarray | None = field(default=None, init=False, repr=False)
def action(self, action: RobotAction) -> RobotAction:
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
new_action = action.copy()
comp = self.transition.get(TransitionKey.COMPLEMENTARY_DATA)
if observation is None:
raise ValueError("Joints observation is require for computing robot kinematics")
if self.use_ik_solution and "IK_solution" in self.transition.get(TransitionKey.COMPLEMENTARY_DATA):
q_raw = self.transition.get(TransitionKey.COMPLEMENTARY_DATA)["IK_solution"]
else:
q_raw = np.array(
[
float(v)
for k, v in observation.items()
if isinstance(k, str)
and k.endswith(".pos")
and k.removesuffix(".pos") in self.motor_names
],
dtype=float,
# Get joint positions from complimentary data
raw = comp["raw_joint_positions"]
if raw is None:
raise ValueError(
"raw_joint_positions is not in complementary data and is required for EEReferenceAndDelta"
)
if q_raw is None:
raise ValueError("Joints observation is require for computing robot kinematics")
if "reference_joint_positions" in comp:
q = comp["reference_joint_positions"]
else:
q = np.array([float(raw[n]) for n in self.motor_names], dtype=float)
# Current pose from FK on measured joints
t_curr = self.kinematics.forward_kinematics(q_raw)
t_curr = self.kinematics.forward_kinematics(q)
enabled = bool(action.pop("enabled"))
tx = float(action.pop("target_x"))
ty = float(action.pop("target_y"))
tz = float(action.pop("target_z"))
wx = float(action.pop("target_wx"))
wy = float(action.pop("target_wy"))
wz = float(action.pop("target_wz"))
gripper_vel = float(action.pop("gripper_vel"))
enabled = bool(new_action.pop("enabled"))
tx = float(new_action.pop("target_x"))
ty = float(new_action.pop("target_y"))
tz = float(new_action.pop("target_z"))
wx = float(new_action.pop("target_wx"))
wy = float(new_action.pop("target_wy"))
wz = float(new_action.pop("target_wz"))
gripper_vel = float(new_action.pop("gripper_vel"))
desired = None
@@ -141,16 +135,16 @@ class EEReferenceAndDelta(RobotActionProcessorStep):
# Write action fields
pos = desired[:3, 3]
tw = Rotation.from_matrix(desired[:3, :3]).as_rotvec()
action["ee.x"] = float(pos[0])
action["ee.y"] = float(pos[1])
action["ee.z"] = float(pos[2])
action["ee.wx"] = float(tw[0])
action["ee.wy"] = float(tw[1])
action["ee.wz"] = float(tw[2])
action["ee.gripper_vel"] = gripper_vel
new_action["ee.x"] = float(pos[0])
new_action["ee.y"] = float(pos[1])
new_action["ee.z"] = float(pos[2])
new_action["ee.wx"] = float(tw[0])
new_action["ee.wy"] = float(tw[1])
new_action["ee.wz"] = float(tw[2])
new_action["ee.gripper_vel"] = gripper_vel
self._prev_enabled = enabled
return action
return new_action
def reset(self):
"""Resets the internal state of the processor."""
@@ -256,7 +250,7 @@ class EEBoundsAndSafety(RobotActionProcessorStep):
@ProcessorStepRegistry.register("inverse_kinematics_ee_to_joints")
@dataclass
class InverseKinematicsEEToJoints(RobotActionProcessorStep):
class InverseKinematicsEEToJoints(ProcessorStep):
"""
Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
@@ -276,36 +270,40 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
q_curr: np.ndarray | None = field(default=None, init=False, repr=False)
initial_guess_current_joints: bool = True
def action(self, action: RobotAction) -> RobotAction:
x = action.pop("ee.x")
y = action.pop("ee.y")
z = action.pop("ee.z")
wx = action.pop("ee.wx")
wy = action.pop("ee.wy")
wz = action.pop("ee.wz")
gripper_pos = action.pop("ee.gripper_pos")
def __call__(self, transition: EnvTransition) -> EnvTransition:
new_transition = transition.copy()
act = new_transition.get(TransitionKey.ACTION).copy()
if not isinstance(act, dict):
raise ValueError(f"Action should be a RobotAction type got {type(act)}")
comp = new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
x = act.pop("ee.x")
y = act.pop("ee.y")
z = act.pop("ee.z")
wx = act.pop("ee.wx")
wy = act.pop("ee.wy")
wz = act.pop("ee.wz")
gripper_pos = act.pop("ee.gripper_pos")
if None in (x, y, z, wx, wy, wz, gripper_pos):
raise ValueError(
"Missing required end-effector pose components: ee.x, ee.y, ee.z, ee.wx, ee.wy, ee.wz, ee.gripper_pos must all be present in action"
)
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
if observation is None:
raise ValueError("Joints observation is require for computing robot kinematics")
q_raw = np.array(
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
dtype=float,
)
if q_raw is None:
raise ValueError("Joints observation is require for computing robot kinematics")
# Get joint positions from complimentary data
raw = comp["raw_joint_positions"]
if raw is None:
raise ValueError(
"raw_joint_positions is not in complementary data and is required for EEReferenceAndDelta"
)
if self.initial_guess_current_joints: # Use current joints as initial guess
self.q_curr = q_raw
self.q_curr = np.array([float(raw[n]) for n in self.motor_names], dtype=float)
else: # Use previous ik solution as initial guess
if self.q_curr is None:
self.q_curr = q_raw
self.q_curr = np.array([float(raw[n]) for n in self.motor_names], dtype=float)
# Build desired 4x4 transform from pos + rotvec (twist)
t_des = np.eye(4, dtype=float)
@@ -316,14 +314,17 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
q_target = self.kinematics.inverse_kinematics(self.q_curr, t_des)
self.q_curr = q_target
new_act = dict(act)
# TODO: This is sentitive to order of motor_names = q_target mapping
for i, name in enumerate(self.motor_names):
if name != "gripper":
action[f"{name}.pos"] = float(q_target[i])
new_act[f"{name}.pos"] = float(q_target[i])
else:
action["gripper.pos"] = float(gripper_pos)
return action
new_act["gripper.pos"] = float(gripper_pos)
new_transition[TransitionKey.ACTION] = new_act
if not self.initial_guess_current_joints:
new_transition[TransitionKey.COMPLEMENTARY_DATA]["reference_joint_positions"] = q_target
return new_transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
@@ -367,30 +368,29 @@ class GripperVelocityToJoint(RobotActionProcessorStep):
discrete_gripper: bool = False
def action(self, action: RobotAction) -> RobotAction:
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
complementary_data = self.transition.get(TransitionKey.COMPLEMENTARY_DATA)
gripper_vel = action.pop("ee.gripper_vel")
if observation is None:
raise ValueError("Joints observation is require for computing robot kinematics")
if "raw_joint_positions" not in complementary_data:
raise ValueError(
"raw_joint_positions is not in complementary data and is required for GripperVelocityToJoint"
)
q_raw = np.array(
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
dtype=float,
)
if q_raw is None:
raise ValueError("Joints observation is require for computing robot kinematics")
curr_gripper_pos = complementary_data["raw_joint_positions"]["gripper"]
if self.discrete_gripper:
# Discrete gripper actions are in [0, 1, 2]
# 0: open, 1: close, 2: stay
# We need to shift them to [-1, 0, 1] and then scale them to clip_max
gripper_vel = (gripper_vel - 1) * self.clip_max
# TODO(Michel,Adil): Fix this logic
# if self.discrete_gripper:
# # Discrete gripper actions are in [0, 1, 2]
# # 0: open, 1: close, 2: stay
# # We need to shift them to [-1, 0, 1] and then scale them to clip_max
# gripper_action = gripper_vel
# gripper_action *= self.clip_max
# action["ee.gripper_pos"] = gripper_action
# Compute desired gripper position
delta = gripper_vel * float(self.speed_factor)
# TODO: This assumes gripper is the last specified joint in the robot
gripper_pos = float(np.clip(q_raw[-1] + delta, self.clip_min, self.clip_max))
gripper_pos = float(np.clip(curr_gripper_pos + delta, self.clip_min, self.clip_max))
action["ee.gripper_pos"] = gripper_pos
return action
@@ -494,6 +494,41 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
return features
@ProcessorStepRegistry.register("add_robot_observation")
@dataclass
class AddRobotObservationAsComplimentaryData(ComplementaryDataProcessorStep):
"""
Reads the robot's current observation and adds it to the transition's complementary data.
This step acts as a bridge to the physical robot, injecting its real-time sensor readings
(like raw joint positions) into the data processing pipeline. This data is then available
for other processing steps.
Attributes:
robot: An instance of a `Robot` class used to get observations from hardware.
"""
robot: Robot
def complementary_data(self, comp: dict | None) -> dict:
new_comp = dict(comp)
obs = (
self.robot.get_observation()
) # todo(steven): why not self.trtansition.get(TransitionKey.OBSERVATION)?
new_comp["raw_joint_positions"] = {
k.removesuffix(".pos"): float(v)
for k, v in obs.items()
if isinstance(k, str) and k.endswith(".pos")
}
return new_comp
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
@ProcessorStepRegistry.register(name="forward_kinematics_joints_to_ee")
@dataclass
class ForwardKinematicsJointsToEE(ProcessorStep):
@@ -523,94 +558,3 @@ class ForwardKinematicsJointsToEE(ProcessorStep):
if features[PipelineFeatureType.OBSERVATION] is not None:
features = self.joints_to_ee_observation_processor.transform_features(features)
return features
@ProcessorStepRegistry.register("inverse_kinematics_rl_step")
@dataclass
class InverseKinematicsRLStep(ProcessorStep):
"""
Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
This is modified from the InverseKinematicsEEToJoints step to be used in the RL pipeline.
"""
kinematics: RobotKinematics
motor_names: list[str]
q_curr: np.ndarray | None = field(default=None, init=False, repr=False)
initial_guess_current_joints: bool = True
def __call__(self, transition: EnvTransition) -> EnvTransition:
new_transition = dict(transition)
action = new_transition.get(TransitionKey.ACTION)
if action is None:
raise ValueError("Action is required for InverseKinematicsEEToJoints")
action = dict(action)
x = action.pop("ee.x")
y = action.pop("ee.y")
z = action.pop("ee.z")
wx = action.pop("ee.wx")
wy = action.pop("ee.wy")
wz = action.pop("ee.wz")
gripper_pos = action.pop("ee.gripper_pos")
if None in (x, y, z, wx, wy, wz, gripper_pos):
raise ValueError(
"Missing required end-effector pose components: ee.x, ee.y, ee.z, ee.wx, ee.wy, ee.wz, ee.gripper_pos must all be present in action"
)
observation = new_transition.get(TransitionKey.OBSERVATION).copy()
if observation is None:
raise ValueError("Joints observation is require for computing robot kinematics")
q_raw = np.array(
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
dtype=float,
)
if q_raw is None:
raise ValueError("Joints observation is require for computing robot kinematics")
if self.initial_guess_current_joints: # Use current joints as initial guess
self.q_curr = q_raw
else: # Use previous ik solution as initial guess
if self.q_curr is None:
self.q_curr = q_raw
# Build desired 4x4 transform from pos + rotvec (twist)
t_des = np.eye(4, dtype=float)
t_des[:3, :3] = Rotation.from_rotvec([wx, wy, wz]).as_matrix()
t_des[:3, 3] = [x, y, z]
# Compute inverse kinematics
q_target = self.kinematics.inverse_kinematics(self.q_curr, t_des)
self.q_curr = q_target
# TODO: This is sentitive to order of motor_names = q_target mapping
for i, name in enumerate(self.motor_names):
if name != "gripper":
action[f"{name}.pos"] = float(q_target[i])
else:
action["gripper.pos"] = float(gripper_pos)
new_transition[TransitionKey.ACTION] = action
complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
complementary_data["IK_solution"] = q_target
new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
return new_transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
for feat in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
features[PipelineFeatureType.ACTION].pop(f"ee.{feat}", None)
for name in self.motor_names:
features[PipelineFeatureType.ACTION][f"{name}.pos"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
)
return features
def reset(self):
"""Resets the initial guess for the IK solver."""
self.q_curr = None
+4
View File
@@ -170,4 +170,8 @@ python lerobot/scripts/control_robot.py \
--control.episode=0
```
Follow [previous tutorial](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md#4-train-a-policy-on-your-data) to train a policy on your data and run inference on your robot. You will need to adapt the code for Stretch.
> TODO(rcadene, aliberts): Add already setup environment and policy yaml configuration files
If you need help, please reach out on Discord in the channel `#stretch3-mobile-arm`.
+3 -1
View File
@@ -118,7 +118,7 @@ echo ${HF_USER}/aloha_test
If you didn't upload with `--control.push_to_hub=false`, you can also visualize it locally with [Rerun](https://github.com/rerun-io/rerun):
```bash
lerobot-dataset-viz \
python -m lerobot.scripts.visualize_dataset \
--repo-id ${HF_USER}/aloha_test --episode 0
```
@@ -193,4 +193,6 @@ As you can see, it's almost the same command as previously used to record your t
## More
Follow this [previous tutorial](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md#4-train-a-policy-on-your-data) for a more in-depth explanation.
If you have any question or need help, please reach out on Discord in the channel `#aloha-arm`.
+90
View File
@@ -0,0 +1,90 @@
#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Use this script to get a quick summary of your system config.
It should be able to run without any of LeRobot's dependencies or LeRobot itself installed.
"""
import platform
HAS_HF_HUB = True
HAS_HF_DATASETS = True
HAS_NP = True
HAS_TORCH = True
HAS_LEROBOT = True
try:
import huggingface_hub
except ImportError:
HAS_HF_HUB = False
try:
import datasets
except ImportError:
HAS_HF_DATASETS = False
try:
import numpy as np
except ImportError:
HAS_NP = False
try:
import torch
except ImportError:
HAS_TORCH = False
try:
import lerobot
except ImportError:
HAS_LEROBOT = False
lerobot_version = lerobot.__version__ if HAS_LEROBOT else "N/A"
hf_hub_version = huggingface_hub.__version__ if HAS_HF_HUB else "N/A"
hf_datasets_version = datasets.__version__ if HAS_HF_DATASETS else "N/A"
np_version = np.__version__ if HAS_NP else "N/A"
torch_version = torch.__version__ if HAS_TORCH else "N/A"
torch_cuda_available = torch.cuda.is_available() if HAS_TORCH else "N/A"
cuda_version = torch._C._cuda_getCompiledVersion() if HAS_TORCH and torch.version.cuda is not None else "N/A"
# TODO(aliberts): refactor into an actual command `lerobot env`
def display_sys_info() -> dict:
"""Run this to get basic system info to help for tracking issues & bugs."""
info = {
"`lerobot` version": lerobot_version,
"Platform": platform.platform(),
"Python version": platform.python_version(),
"Huggingface_hub version": hf_hub_version,
"Dataset version": hf_datasets_version,
"Numpy version": np_version,
"PyTorch version (GPU?)": f"{torch_version} ({torch_cuda_available})",
"Cuda version": cuda_version,
"Using GPU in script?": "<fill in>",
# "Using distributed or parallel set-up in script?": "<fill in>",
}
print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the last point.\n")
print(format_dict(info))
return info
def format_dict(d: dict) -> str:
return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"
if __name__ == "__main__":
display_sys_info()
+197 -190
View File
@@ -52,11 +52,10 @@ import logging
import threading
import time
from collections import defaultdict
from collections.abc import Callable
from collections.abc import Callable, Iterator
from contextlib import nullcontext
from copy import deepcopy
from dataclasses import asdict
from functools import partial
from pathlib import Path
from pprint import pformat
from typing import Any, TypedDict
@@ -164,11 +163,11 @@ def rollout(
# Infer "task" from attributes of environments.
# TODO: works with SyncVectorEnv but not AsyncVectorEnv
observation = add_envs_task(env, observation)
observation = preprocessor(observation)
with torch.inference_mode():
action = policy.select_action(observation)
action = postprocessor(action)
# Convert to CPU / numpy.
action_numpy: np.ndarray = action.to("cpu").numpy()
assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)"
@@ -232,9 +231,9 @@ def rollout(
def eval_policy(
env: gym.vector.VectorEnv,
policy: PreTrainedPolicy,
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
n_episodes: int,
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]] | None = None,
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] | None = None,
max_episodes_rendered: int = 0,
videos_dir: Path | None = None,
return_episode_data: bool = False,
@@ -418,7 +417,6 @@ def eval_policy(
"eval_ep_s": (time.time() - start) / n_episodes,
},
}
if return_episode_data:
info["episodes"] = episode_data
@@ -479,7 +477,6 @@ def eval_main(cfg: EvalPipelineConfig):
# Check device is available
device = get_safe_torch_device(cfg.policy.device, log=True)
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
set_seed(cfg.seed)
@@ -490,7 +487,6 @@ def eval_main(cfg: EvalPipelineConfig):
envs = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
logging.info("Making policy.")
policy = make_policy(
cfg=cfg.policy,
env_cfg=cfg.env,
@@ -498,10 +494,7 @@ def eval_main(cfg: EvalPipelineConfig):
policy.eval()
preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=cfg.policy,
pretrained_path=cfg.policy.pretrained_path,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path
)
with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
info = eval_policy_all(
@@ -514,14 +507,17 @@ def eval_main(cfg: EvalPipelineConfig):
videos_dir=Path(cfg.output_dir) / "videos",
start_seed=cfg.seed,
max_parallel_tasks=cfg.env.max_parallel_tasks,
verbose=False,
)
print("Overall Aggregated Metrics:")
print(info["overall"])
print(info["overall"]["aggregated"])
# Print per-suite stats
for task_group, task_group_info in info.items():
if task_group == "overall":
continue # Skip the overall stats since we already printed it
print(f"\nAggregated Metrics for {task_group}:")
print(task_group_info)
print(task_group_info["aggregated"])
# Close all vec envs
close_envs(envs)
@@ -543,205 +539,216 @@ class TaskMetrics(TypedDict):
ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths")
def eval_one(
env: gym.vector.VectorEnv,
*,
def eval_policy_all(
envs: dict[str, dict[int, gym.vector.VectorEnv]],
policy: PreTrainedPolicy,
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
n_episodes: int,
max_episodes_rendered: int,
videos_dir: Path | None,
return_episode_data: bool,
start_seed: int | None,
) -> TaskMetrics:
"""Evaluates one task_id of one suite using the provided vec env."""
task_videos_dir = videos_dir
task_result = eval_policy(
env=env,
policy=policy,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=n_episodes,
max_episodes_rendered=max_episodes_rendered,
videos_dir=task_videos_dir,
return_episode_data=return_episode_data,
start_seed=start_seed,
)
per_episode = task_result["per_episode"]
return TaskMetrics(
sum_rewards=[ep["sum_reward"] for ep in per_episode],
max_rewards=[ep["max_reward"] for ep in per_episode],
successes=[ep["success"] for ep in per_episode],
video_paths=task_result.get("video_paths", []),
)
def run_one(
task_group: str,
task_id: int,
env,
*,
policy,
preprocessor,
postprocessor,
n_episodes: int,
max_episodes_rendered: int,
videos_dir: Path | None,
return_episode_data: bool,
start_seed: int | None,
):
"""
Run eval_one for a single (task_group, task_id, env).
Returns (task_group, task_id, task_metrics_dict).
This function is intentionally module-level to make it easy to test.
"""
task_videos_dir = None
if videos_dir is not None:
task_videos_dir = videos_dir / f"{task_group}_{task_id}"
task_videos_dir.mkdir(parents=True, exist_ok=True)
# Call the existing eval_one (assumed to return TaskMetrics-like dict)
metrics = eval_one(
env,
policy=policy,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=n_episodes,
max_episodes_rendered=max_episodes_rendered,
videos_dir=task_videos_dir,
return_episode_data=return_episode_data,
start_seed=start_seed,
)
# ensure we always provide video_paths key to simplify accumulation
if max_episodes_rendered > 0:
metrics.setdefault("video_paths", [])
return task_group, task_id, metrics
def eval_policy_all(
envs: dict[str, dict[int, gym.vector.VectorEnv]],
policy,
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
n_episodes: int,
*,
max_episodes_rendered: int = 0,
videos_dir: Path | None = None,
return_episode_data: bool = False,
start_seed: int | None = None,
max_parallel_tasks: int = 1,
verbose: bool = True,
) -> dict:
"""
Evaluate a nested `envs` dict: {task_group: {task_id: vec_env}}.
This implementation flattens tasks, runs them sequentially or via ThreadPoolExecutor,
accumulates per-group and overall statistics, and returns the same aggregate metrics
schema as the single-env evaluator (avg_sum_reward / avg_max_reward / pc_success / timings)
plus per-task infos.
Evaluate a policy over a dict-of-dicts of vectorized envs:
envs[suite_name][task_id] -> gym.vector.VectorEnv
Returns a dict with per-suite aggregates and an 'overall' block.
"""
start_t = time.time()
global_start = time.time()
# Flatten envs into list of (task_group, task_id, env)
tasks = [(tg, tid, vec) for tg, group in envs.items() for tid, vec in group.items()]
# inner: evaluate a single (suite, task)
def eval_one(
task_group: str,
task_id: int,
env: gym.vector.VectorEnv,
*,
policy: PreTrainedPolicy,
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
n_episodes: int,
max_episodes_rendered: int,
videos_dir: Path | None,
return_episode_data: bool,
start_seed: int | None,
) -> TaskMetrics:
"""Evaluates one task_id of one suite using the provided vec env."""
if verbose:
print(f"Evaluating: task_group={task_group}, task_id={task_id} ...")
# accumulators: track metrics at both per-group level and across all groups
task_videos_dir = None
if videos_dir is not None:
task_videos_dir = videos_dir / f"{task_group}_{task_id}"
task_videos_dir.mkdir(parents=True, exist_ok=True)
task_result = eval_policy(
env=env,
policy=policy,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=n_episodes,
max_episodes_rendered=max_episodes_rendered,
videos_dir=task_videos_dir,
return_episode_data=return_episode_data,
start_seed=start_seed,
)
per_episode = task_result["per_episode"]
return TaskMetrics(
sum_rewards=[ep["sum_reward"] for ep in per_episode],
max_rewards=[ep["max_reward"] for ep in per_episode],
successes=[ep["success"] for ep in per_episode],
video_paths=task_result.get("video_paths", []),
)
def _eval_monotask(
envs, policy, n_episodes, max_episodes_rendered, videos_dir, return_episode_data, start_seed
):
for task_group, tasks in envs.items():
for task_id, vec in tasks.items():
yield (
task_group,
task_id,
eval_one(
task_group,
task_id,
vec,
policy=policy,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=n_episodes,
max_episodes_rendered=max_episodes_rendered,
videos_dir=videos_dir,
return_episode_data=return_episode_data,
start_seed=start_seed,
),
)
def _eval_parallel(
envs,
policy,
preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
n_episodes,
max_episodes_rendered,
videos_dir,
return_episode_data,
start_seed,
max_parallel_tasks,
):
with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor:
fut2key: dict[cf.Future, tuple[str, int]] = {}
for task_group, tasks in envs.items():
for task_id, vec in tasks.items():
fut = executor.submit(
eval_one,
task_group,
task_id,
vec,
policy=policy,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=n_episodes,
max_episodes_rendered=max_episodes_rendered,
videos_dir=videos_dir,
return_episode_data=return_episode_data,
start_seed=start_seed,
)
fut2key[fut] = (task_group, task_id)
for fut in cf.as_completed(fut2key):
task_group, task_id = fut2key[fut]
yield task_group, task_id, fut.result()
# result producer: sequential or threaded, same consumer
def iter_task_results() -> Iterator[tuple[str, int, TaskMetrics]]:
"""
Yield evaluation results for each (task_group, task_id).
Depending on `max_parallel_tasks`, runs sequentially or in parallel,
but always returns a generator of tuples:
(task_group, task_id, TaskMetrics).
"""
if max_parallel_tasks == 1:
yield from _eval_monotask(
envs,
policy,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=n_episodes,
max_episodes_rendered=max_episodes_rendered,
videos_dir=videos_dir,
return_episode_data=return_episode_data,
start_seed=start_seed,
)
else:
yield from _eval_parallel(
envs,
policy=policy,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=n_episodes,
max_episodes_rendered=max_episodes_rendered,
videos_dir=videos_dir,
return_episode_data=return_episode_data,
start_seed=start_seed,
max_parallel_tasks=max_parallel_tasks,
)
# single accumulator path on the main thread
group_acc: dict[str, dict[str, list]] = defaultdict(lambda: {k: [] for k in ACC_KEYS})
overall: dict[str, list] = {k: [] for k in ACC_KEYS}
per_task_infos: list[dict] = []
# small inline helper to accumulate one task's metrics into accumulators
def _accumulate_to(group: str, metrics: dict):
# metrics expected to contain 'sum_rewards', 'max_rewards', 'successes', optionally 'video_paths'
# but eval_one may store per-episode lists; we assume metrics uses scalars averaged per task as before.
# To be robust, accept scalars or lists.
def _append(key, value):
if value is None:
return
if isinstance(value, list):
group_acc[group][key].extend(value)
overall[key].extend(value)
else:
group_acc[group][key].append(value)
overall[key].append(value)
for task_group, _task_id, metrics in iter_task_results():
acc = group_acc[task_group]
for k in ACC_KEYS:
acc[k].extend(metrics[k])
overall[k].extend(metrics[k])
_append("sum_rewards", metrics.get("sum_rewards"))
_append("max_rewards", metrics.get("max_rewards"))
_append("successes", metrics.get("successes"))
# video_paths is list-like
paths = metrics.get("video_paths", [])
if paths:
group_acc[group]["video_paths"].extend(paths)
overall["video_paths"].extend(paths)
# build outputs
results: dict[str, dict] = {}
for task_group, data in group_acc.items():
suite_rewards = data["sum_rewards"]
suite_max = data["max_rewards"]
suite_succ = data["successes"]
suite_vids = data["video_paths"]
# Choose runner (sequential vs threaded)
task_runner = partial(
run_one,
policy=policy,
preprocessor=preprocessor,
postprocessor=postprocessor,
n_episodes=n_episodes,
max_episodes_rendered=max_episodes_rendered,
videos_dir=videos_dir,
return_episode_data=return_episode_data,
start_seed=start_seed,
)
suite_eval_s = time.time() - global_start
suite_eval_ep_s = suite_eval_s / max(1, len(suite_rewards))
if max_parallel_tasks <= 1:
# sequential path (single accumulator path on the main thread)
# NOTE: keeping a single-threaded accumulator avoids concurrent list appends or locks
for task_group, task_id, env in tasks:
tg, tid, metrics = task_runner(task_group, task_id, env)
_accumulate_to(tg, metrics)
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
else:
# threaded path: submit all tasks, consume completions on main thread and accumulate there
with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor:
fut2meta = {}
for task_group, task_id, env in tasks:
fut = executor.submit(task_runner, task_group, task_id, env)
fut2meta[fut] = (task_group, task_id)
for fut in cf.as_completed(fut2meta):
tg, tid, metrics = fut.result()
_accumulate_to(tg, metrics)
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
# compute aggregated metrics helper (robust to lists/scalars)
def _agg_from_list(xs):
if not xs:
return float("nan")
arr = np.array(xs, dtype=float)
return float(np.nanmean(arr))
# compute per-group aggregates
groups_aggregated = {}
for group, acc in group_acc.items():
groups_aggregated[group] = {
"avg_sum_reward": _agg_from_list(acc["sum_rewards"]),
"avg_max_reward": _agg_from_list(acc["max_rewards"]),
"pc_success": _agg_from_list(acc["successes"]) * 100 if acc["successes"] else float("nan"),
"n_episodes": len(acc["sum_rewards"]),
"video_paths": list(acc["video_paths"]),
results[task_group] = {
"aggregated": {
"avg_sum_reward": float(np.nanmean(suite_rewards)) if suite_rewards else float("nan"),
"avg_max_reward": float(np.nanmean(suite_max)) if suite_max else float("nan"),
"pc_success": float(np.nanmean(suite_succ) * 100) if suite_succ else float("nan"),
"eval_s": suite_eval_s,
"eval_ep_s": suite_eval_ep_s,
},
"video_paths": suite_vids,
"episodes": None,
}
# overall aggregates
overall_agg = {
"avg_sum_reward": _agg_from_list(overall["sum_rewards"]),
"avg_max_reward": _agg_from_list(overall["max_rewards"]),
"pc_success": _agg_from_list(overall["successes"]) * 100 if overall["successes"] else float("nan"),
"n_episodes": len(overall["sum_rewards"]),
"eval_s": time.time() - start_t,
"eval_ep_s": (time.time() - start_t) / max(1, len(overall["sum_rewards"])),
"video_paths": list(overall["video_paths"]),
}
return {
"per_task": per_task_infos,
"per_group": groups_aggregated,
"overall": overall_agg,
global_eval_s = time.time() - global_start
global_eval_ep_s = global_eval_s / max(1, len(overall["sum_rewards"]))
results["overall"] = {
"aggregated": {
"avg_sum_reward": float(np.nanmean(overall["sum_rewards"]))
if overall["sum_rewards"]
else float("nan"),
"avg_max_reward": float(np.nanmean(overall["max_rewards"]))
if overall["max_rewards"]
else float("nan"),
"pc_success": float(np.nanmean(overall["successes"]) * 100)
if overall["successes"]
else float("nan"),
"eval_s": global_eval_s,
"eval_ep_s": global_eval_ep_s,
},
"video_paths": overall["video_paths"],
"episodes": None,
}
return results
def main():
-96
View File
@@ -1,96 +0,0 @@
#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Use this script to get a quick summary of your system config.
It should be able to run without any of LeRobot's dependencies or LeRobot itself installed.
Example:
```shell
lerobot-info
```
"""
import importlib
import platform
def get_package_version(package_name: str) -> str:
"""Get the version of a package if it exists, otherwise return 'N/A'."""
try:
module = importlib.import_module(package_name)
return getattr(module, "__version__", "Installed (version not found)")
except ImportError:
return "N/A"
def get_sys_info() -> dict:
"""Run this to get basic system info to help for tracking issues & bugs."""
# General package versions
info = {
"lerobot version": get_package_version("lerobot"),
"Platform": platform.platform(),
"Python version": platform.python_version(),
"Huggingface Hub version": get_package_version("huggingface_hub"),
"Datasets version": get_package_version("datasets"),
"Numpy version": get_package_version("numpy"),
}
# PyTorch and GPU specific information
torch_version = "N/A"
torch_cuda_available = "N/A"
cuda_version = "N/A"
gpu_model = "N/A"
try:
import torch
torch_version = torch.__version__
torch_cuda_available = torch.cuda.is_available()
if torch_cuda_available:
cuda_version = torch.version.cuda
# Gets the name of the first available GPU
gpu_model = torch.cuda.get_device_name(0)
except ImportError:
# If torch is not installed, the default "N/A" values will be used.
pass
info.update(
{
"PyTorch version": torch_version,
"Is PyTorch built with CUDA support?": torch_cuda_available,
"Cuda version": cuda_version,
"GPU model": gpu_model,
"Using GPU in script?": "<fill in>",
}
)
return info
def format_dict_for_markdown(d: dict) -> str:
"""Formats a dictionary into a markdown-friendly bulleted list."""
return "\n".join([f"- {prop}: {val}" for prop, val in d.items()])
def main():
system_info = get_sys_info()
print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the last point.\n")
print(format_dict_for_markdown(system_info))
if __name__ == "__main__":
main()
@@ -24,7 +24,7 @@ Examples of usage:
- Start an actor server for real robot training with human-in-the-loop intervention:
```bash
python -m lerobot.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
python -m lerobot.scripts.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
```
**NOTE**: The actor server requires a running learner server to connect to. Ensure the learner
@@ -64,6 +64,12 @@ from lerobot.policies.factory import make_policy
from lerobot.policies.sac.modeling_sac import SACPolicy
from lerobot.processor import TransitionKey
from lerobot.robots import so100_follower # noqa: F401
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.utils import TeleopEvents
from lerobot.transport import services_pb2, services_pb2_grpc
@@ -90,13 +96,6 @@ from lerobot.utils.utils import (
init_logging,
)
from .gym_manipulator import (
create_transition,
make_processors,
make_robot_env,
step_env_and_process_transition,
)
ACTOR_SHUTDOWN_TIMEOUT = 30
# Main entry point
@@ -25,13 +25,12 @@ from lerobot.robots import ( # noqa: F401
make_robot_from_config,
so100_follower,
)
from lerobot.scripts.rl.gym_manipulator import make_robot_env
from lerobot.teleoperators import (
gamepad, # noqa: F401
so101_leader, # noqa: F401
)
from .gym_manipulator import make_robot_env
logging.basicConfig(level=logging.INFO)
@@ -44,14 +44,12 @@ from lerobot.processor import (
MotorCurrentProcessorStep,
Numpy2TorchActionProcessorStep,
RewardClassifierProcessorStep,
RobotActionToPolicyActionProcessorStep,
TimeLimitProcessorStep,
Torch2NumpyActionProcessorStep,
TransitionKey,
VanillaObservationProcessorStep,
create_transition,
)
from lerobot.processor.converters import identity_transition
from lerobot.robots import ( # noqa: F401
RobotConfig,
make_robot_from_config,
@@ -59,11 +57,12 @@ from lerobot.robots import ( # noqa: F401
)
from lerobot.robots.robot import Robot
from lerobot.robots.so100_follower.robot_kinematic_processor import (
AddRobotObservationAsComplimentaryData,
EEBoundsAndSafety,
EEReferenceAndDelta,
ForwardKinematicsJointsToEEObservation,
ForwardKinematicsJointsToEE,
GripperVelocityToJoint,
InverseKinematicsRLStep,
InverseKinematicsEEToJoints,
)
from lerobot.teleoperators import (
gamepad, # noqa: F401
@@ -157,20 +156,15 @@ class RobotEnv(gym.Env):
self.use_gripper = use_gripper
self._joint_names = list(self.robot.bus.motors.keys())
self._raw_joint_positions = None
self._setup_spaces()
def _get_observation(self) -> dict[str, Any]:
"""Get current robot observation including joint positions and camera images."""
obs_dict = self.robot.get_observation()
raw_joint_joint_position = {f"{name}.pos": obs_dict[f"{name}.pos"] for name in self._joint_names}
joint_positions = np.array([raw_joint_joint_position[f"{name}.pos"] for name in self._joint_names])
joint_positions = np.array([obs_dict[name] for name in self._joint_names])
images = {key: obs_dict[key] for key in self._image_keys}
return {"agent_pos": joint_positions, "pixels": images, **raw_joint_joint_position}
return {"agent_pos": joint_positions, "pixels": images}
def _setup_spaces(self) -> None:
"""Configure observation and action spaces based on robot capabilities."""
@@ -245,19 +239,21 @@ class RobotEnv(gym.Env):
self.current_step = 0
self.episode_data = None
obs = self._get_observation()
self._raw_joint_positions = {f"{key}.pos": obs[f"{key}.pos"] for key in self._joint_names}
return obs, {TeleopEvents.IS_INTERVENTION: False}
return obs, {
TeleopEvents.IS_INTERVENTION: False,
"raw_joint_positions": obs["agent_pos"],
}
def step(self, action) -> tuple[dict[str, np.ndarray], float, bool, bool, dict[str, Any]]:
"""Execute one environment step with given action."""
joint_targets_dict = {f"{key}.pos": action[i] for i, key in enumerate(self.robot.bus.motors.keys())}
joint_targets_dict = {
f"{key}.pos": action[f"action.{key}.pos"] for key in self.robot.bus.motors.keys()
}
self.robot.send_action(joint_targets_dict)
obs = self._get_observation()
self._raw_joint_positions = {f"{key}.pos": obs[f"{key}.pos"] for key in self._joint_names}
if self.display_cameras:
self.render()
@@ -272,7 +268,7 @@ class RobotEnv(gym.Env):
reward,
terminated,
truncated,
{TeleopEvents.IS_INTERVENTION: False},
{TeleopEvents.IS_INTERVENTION: False, "raw_joint_positions": obs["agent_pos"]},
)
def render(self) -> None:
@@ -292,10 +288,6 @@ class RobotEnv(gym.Env):
if self.robot.is_connected:
self.robot.disconnect()
def get_raw_joint_positions(self) -> dict[str, float]:
"""Get raw joint positions."""
return self._raw_joint_positions
def make_robot_env(cfg: HILSerlRobotEnvConfig) -> tuple[gym.Env, Any]:
"""Create robot environment from configuration.
@@ -352,9 +344,7 @@ def make_robot_env(cfg: HILSerlRobotEnvConfig) -> tuple[gym.Env, Any]:
def make_processors(
env: gym.Env, teleop_device: Teleoperator | None, cfg: HILSerlRobotEnvConfig, device: str = "cpu"
) -> tuple[
DataProcessorPipeline[EnvTransition, EnvTransition], DataProcessorPipeline[EnvTransition, EnvTransition]
]:
) -> tuple[Any, Any]:
"""Create environment and action processors.
Args:
@@ -376,6 +366,7 @@ def make_processors(
Torch2NumpyActionProcessorStep(),
]
# Minimal processor pipeline for GymHIL simulation
env_pipeline_steps = [
Numpy2TorchActionProcessorStep(),
VanillaObservationProcessorStep(),
@@ -383,10 +374,8 @@ def make_processors(
DeviceProcessorStep(device=device),
]
return DataProcessorPipeline(
steps=env_pipeline_steps, to_transition=identity_transition, to_output=identity_transition
), DataProcessorPipeline(
steps=action_pipeline_steps, to_transition=identity_transition, to_output=identity_transition
return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline(
steps=action_pipeline_steps
)
# Full processor pipeline for real robot environment
@@ -412,7 +401,7 @@ def make_processors(
if kinematics_solver is not None:
env_pipeline_steps.append(
ForwardKinematicsJointsToEEObservation(
ForwardKinematicsJointsToEE(
kinematics=kinematics_solver,
motor_names=motor_names,
)
@@ -461,6 +450,7 @@ def make_processors(
action_pipeline_steps = [
AddTeleopActionAsComplimentaryDataStep(teleop_device=teleop_device),
AddTeleopEventsAsInfoStep(teleop_device=teleop_device),
AddRobotObservationAsComplimentaryData(robot=env.robot),
InterventionActionProcessorStep(
use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False,
terminate_on_success=terminate_on_success,
@@ -480,37 +470,34 @@ def make_processors(
end_effector_step_sizes=cfg.processor.inverse_kinematics.end_effector_step_sizes,
motor_names=motor_names,
use_latched_reference=False,
use_ik_solution=True,
),
EEBoundsAndSafety(
end_effector_bounds=cfg.processor.inverse_kinematics.end_effector_bounds,
),
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
motor_names=motor_names,
initial_guess_current_joints=False,
),
GripperVelocityToJoint(
motor_names=motor_names,
clip_max=cfg.processor.max_gripper_pos,
speed_factor=1.0,
discrete_gripper=True,
),
InverseKinematicsRLStep(
kinematics=kinematics_solver, motor_names=motor_names, initial_guess_current_joints=False
),
]
action_pipeline_steps.extend(inverse_kinematics_steps)
action_pipeline_steps.append(RobotActionToPolicyActionProcessorStep(motor_names=motor_names))
return DataProcessorPipeline(
steps=env_pipeline_steps, to_transition=identity_transition, to_output=identity_transition
), DataProcessorPipeline(
steps=action_pipeline_steps, to_transition=identity_transition, to_output=identity_transition
)
return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline(steps=action_pipeline_steps)
def step_env_and_process_transition(
env: gym.Env,
transition: EnvTransition,
action: torch.Tensor,
env_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
action_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
) -> EnvTransition:
env_processor: DataProcessorPipeline,
action_processor: DataProcessorPipeline,
):
"""
Execute one step with processor pipeline.
@@ -527,14 +514,14 @@ def step_env_and_process_transition(
# Create action transition
transition[TransitionKey.ACTION] = action
transition[TransitionKey.OBSERVATION] = (
env.get_raw_joint_positions() if hasattr(env, "get_raw_joint_positions") else {}
)
processed_action_transition = action_processor(transition)
processed_action = processed_action_transition[TransitionKey.ACTION]
# Step environment with processed action
obs, reward, terminated, truncated, info = env.step(processed_action)
# Combine rewards from environment and action processor
reward = reward + processed_action_transition[TransitionKey.REWARD]
terminated = terminated or processed_action_transition[TransitionKey.DONE]
truncated = truncated or processed_action_transition[TransitionKey.TRUNCATED]
@@ -558,8 +545,8 @@ def step_env_and_process_transition(
def control_loop(
env: gym.Env,
env_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
action_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
env_processor: DataProcessorPipeline,
action_processor: DataProcessorPipeline,
teleop_device: Teleoperator,
cfg: GymManipulatorConfig,
) -> None:
@@ -591,7 +578,7 @@ def control_loop(
# Process initial observation
transition = create_transition(observation=obs, info=info, complementary_data=complementary_data)
transition = env_processor(data=transition)
transition = env_processor(transition)
# Determine if gripper is used
use_gripper = cfg.env.processor.gripper.use_gripper if cfg.env.processor.gripper is not None else True
@@ -660,11 +647,7 @@ def control_loop(
truncated = transition.get(TransitionKey.TRUNCATED, False)
if cfg.mode == "record":
observations = {
k: v.squeeze(0).cpu()
for k, v in transition[TransitionKey.OBSERVATION].items()
if isinstance(v, torch.Tensor)
}
observations = {k: v.squeeze(0).cpu() for k, v in transition[TransitionKey.OBSERVATION].items()}
# Use teleop_action if available, otherwise use the action from the transition
action_to_record = transition[TransitionKey.COMPLEMENTARY_DATA].get(
"teleop_action", transition[TransitionKey.ACTION]
@@ -678,10 +661,8 @@ def control_loop(
if use_gripper:
discrete_penalty = transition[TransitionKey.COMPLEMENTARY_DATA].get("discrete_penalty", 0.0)
frame["complementary_info.discrete_penalty"] = np.array([discrete_penalty], dtype=np.float32)
if dataset is not None:
frame["task"] = cfg.dataset.task
dataset.add_frame(frame)
dataset.add_frame(frame, task=cfg.dataset.task)
episode_step += 1
@@ -731,19 +712,17 @@ def replay_trajectory(
episodes=[cfg.dataset.replay_episode],
download_videos=False,
)
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == cfg.dataset.replay_episode)
actions = episode_frames.select_columns("action")
dataset_actions = dataset.hf_dataset.select_columns(["action"])
_, info = env.reset()
for action_data in actions:
for action_data in dataset_actions:
start_time = time.perf_counter()
transition = create_transition(
observation=env.get_raw_joint_positions() if hasattr(env, "get_raw_joint_positions") else {},
action=action_data["action"],
complementary_data={"raw_joint_positions": info["raw_joint_positions"]},
)
transition = action_processor(transition)
env.step(transition[TransitionKey.ACTION])
_, _, _, _, info = env.step(transition[TransitionKey.ACTION])
busy_wait(1 / cfg.env.fps - (time.perf_counter() - start_time))
@@ -25,7 +25,7 @@ Examples of usage:
- Start a learner server for training:
```bash
python -m lerobot.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
python -m lerobot.scripts.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
```
**NOTE**: Start the learner server before launching the actor server. The learner opens a gRPC server
@@ -73,6 +73,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.policies.factory import make_policy
from lerobot.policies.sac.modeling_sac import SACPolicy
from lerobot.robots import so100_follower # noqa: F401
from lerobot.scripts.rl import learner_service
from lerobot.teleoperators import gamepad, so101_leader # noqa: F401
from lerobot.teleoperators.utils import TeleopEvents
from lerobot.transport import services_pb2_grpc
@@ -99,8 +100,6 @@ from lerobot.utils.utils import (
)
from lerobot.utils.wandb_utils import WandBLogger
from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService
LOG_PREFIX = "[LEARNER]"
@@ -640,7 +639,7 @@ def start_learner(
# TODO: Check if its useful
_ = ProcessSignalHandler(False, display_pid=True)
service = LearnerService(
service = learner_service.LearnerService(
shutdown_event=shutdown_event,
parameters_queue=parameters_queue,
seconds_between_pushes=cfg.policy.actor_learner_config.policy_parameters_push_frequency,
@@ -650,7 +649,7 @@ def start_learner(
)
server = grpc.server(
ThreadPoolExecutor(max_workers=MAX_WORKERS),
ThreadPoolExecutor(max_workers=learner_service.MAX_WORKERS),
options=[
("grpc.max_receive_message_length", MAX_MESSAGE_SIZE),
("grpc.max_send_message_length", MAX_MESSAGE_SIZE),
@@ -671,7 +670,7 @@ def start_learner(
shutdown_event.wait()
logging.info("[LEARNER] Stopping gRPC server...")
server.stop(SHUTDOWN_TIMEOUT)
server.stop(learner_service.SHUTDOWN_TIMEOUT)
logging.info("[LEARNER] gRPC server stopped")
+13 -19
View File
@@ -67,10 +67,8 @@ def update_policy(
) -> tuple[MetricsTracker, dict]:
"""
Performs a single training step to update the policy's weights.
This function executes the forward and backward passes, clips gradients, and steps the optimizer and
learning rate scheduler. It also handles mixed-precision training via a GradScaler.
Args:
train_metrics: A MetricsTracker instance to record training statistics.
policy: The policy model to be trained.
@@ -81,7 +79,6 @@ def update_policy(
lr_scheduler: An optional learning rate scheduler.
use_amp: A boolean indicating whether to use automatic mixed precision.
lock: An optional lock for thread-safe optimizer updates.
Returns:
A tuple containing:
- The updated MetricsTracker with new statistics for this step.
@@ -132,7 +129,6 @@ def update_policy(
def train(cfg: TrainPipelineConfig):
"""
Main function to train a policy.
This function orchestrates the entire training pipeline, including:
- Setting up logging, seeding, and device configuration.
- Creating the dataset, evaluation environment (if applicable), policy, and optimizer.
@@ -140,7 +136,6 @@ def train(cfg: TrainPipelineConfig):
- Running the main training loop, which involves fetching data batches and calling `update_policy`.
- Periodically logging metrics, saving model checkpoints, and evaluating the policy.
- Pushing the final trained model to the Hugging Face Hub if configured.
Args:
cfg: A `TrainPipelineConfig` object containing all training configurations.
"""
@@ -163,7 +158,6 @@ def train(cfg: TrainPipelineConfig):
logging.info("Creating dataset")
dataset = make_dataset(cfg)
# Create environment used for evaluating checkpoints during training on simulation data.
# On real-world data, no need to create an environment as evaluations are done outside train.py,
# using the eval.py instead, with gym_dora environment and dora-rs.
@@ -173,20 +167,17 @@ def train(cfg: TrainPipelineConfig):
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
logging.info("Creating policy")
policy = make_policy(
cfg=cfg.policy,
ds_meta=dataset.meta,
)
# Create processors - only provide dataset_stats if not resuming from saved processors
processor_kwargs = {}
if not (cfg.resume and cfg.policy.pretrained_path):
# Only provide dataset_stats when not resuming from saved processor state
processor_kwargs["dataset_stats"] = dataset.meta.stats
if cfg.policy.pretrained_path is not None:
processor_kwargs["preprocessor_overrides"] = {"device_processor": {"device": device.type}}
preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path, **processor_kwargs
)
@@ -238,7 +229,6 @@ def train(cfg: TrainPipelineConfig):
dl_iter = cycle(dataloader)
policy.train()
train_metrics = {
"loss": AverageMeter("loss", ":.3f"),
"grad_norm": AverageMeter("grdn", ":.3f"),
@@ -255,8 +245,8 @@ def train(cfg: TrainPipelineConfig):
for _ in range(step, cfg.steps):
start_time = time.perf_counter()
batch = next(dl_iter)
batch = preprocessor(batch)
train_tracker.dataloading_s = time.perf_counter() - start_time
batch = preprocessor(batch)
train_tracker, output_dict = update_policy(
train_tracker,
@@ -304,7 +294,7 @@ def train(cfg: TrainPipelineConfig):
torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
):
eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env
env=eval_env, # dict[suite][task_id] -> vec_env
policy=policy,
preprocessor=preprocessor,
postprocessor=postprocessor,
@@ -313,13 +303,17 @@ def train(cfg: TrainPipelineConfig):
max_episodes_rendered=4,
start_seed=cfg.seed,
max_parallel_tasks=cfg.env.max_parallel_tasks,
verbose=False,
)
# overall metrics (suite-agnostic)
aggregated = eval_info["overall"]
aggregated = eval_info["overall"]["aggregated"]
# optional: per-suite logging
for suite, suite_info in eval_info.items():
logging.info("Suite %s aggregated: %s", suite, suite_info)
if suite == "overall":
continue
logging.info("Suite %s aggregated: %s", suite, suite_info["aggregated"])
# meters/tracker
eval_metrics = {
@@ -330,13 +324,13 @@ def train(cfg: TrainPipelineConfig):
eval_tracker = MetricsTracker(
cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step
)
eval_tracker.eval_s = aggregated.pop("eval_s")
eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward")
eval_tracker.pc_success = aggregated.pop("pc_success")
eval_tracker.eval_s = aggregated.get("eval_s", 0.0)
eval_tracker.avg_sum_reward = aggregated.get("avg_sum_reward", float("nan"))
eval_tracker.pc_success = aggregated.get("pc_success", float("nan"))
if wandb_logger:
wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
wandb_logger.log_video(eval_info["overall"]["video_paths"][0], step, mode="eval")
wandb_logger.log_video(eval_info["video_paths"][0], step, mode="eval")
if eval_env:
close_envs(eval_env)
@@ -29,14 +29,14 @@ Examples:
- Visualize data stored on a local machine:
```
local$ lerobot-dataset-viz \
local$ python -m lerobot.scripts.visualize_dataset \
--repo-id lerobot/pusht \
--episode-index 0
```
- Visualize data stored on a distant machine with a local viewer:
```
distant$ lerobot-dataset-viz \
distant$ python -m lerobot.scripts.visualize_dataset \
--repo-id lerobot/pusht \
--episode-index 0 \
--save 1 \
@@ -50,7 +50,7 @@ local$ rerun lerobot_pusht_episode_0.rrd
(You need to forward the websocket port to the distant machine, with
`ssh -L 9087:localhost:9087 username@remote-host`)
```
distant$ lerobot-dataset-viz \
distant$ python -m lerobot.scripts.visualize_dataset \
--repo-id lerobot/pusht \
--episode-index 0 \
--mode distant \
@@ -20,10 +20,10 @@ Additionally, each individual transform can be visualized separately as well as
Example:
```bash
lerobot-imgtransform-viz \
--repo_id=lerobot/pusht \
--episodes='[0]' \
--image_transforms.enable=True
python -m lerobot.scripts.visualize_image_transforms \
--repo_id=lerobot/pusht \
--episodes='[0]' \
--image_transforms.enable=True
```
"""
@@ -126,9 +126,5 @@ def visualize_image_transforms(cfg: DatasetConfig, output_dir: Path = OUTPUT_DIR
save_each_transform(cfg.image_transforms, original_frame, output_dir, n_examples)
def main():
visualize_image_transforms()
if __name__ == "__main__":
main()
visualize_image_transforms()
+6 -10
View File
@@ -109,8 +109,8 @@ def teleop_loop(
teleop: Teleoperator,
robot: Robot,
fps: int,
teleop_action_processor: RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction],
robot_action_processor: RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction],
teleop_action_processor: RobotProcessorPipeline[RobotAction, RobotAction],
robot_action_processor: RobotProcessorPipeline[RobotAction, RobotAction],
robot_observation_processor: RobotProcessorPipeline[RobotObservation, RobotObservation],
display_data: bool = False,
duration: float | None = None,
@@ -137,25 +137,21 @@ def teleop_loop(
while True:
loop_start = time.perf_counter()
# Get robot observation
# Not really needed for now other than for visualization
# teleop_action_processor can take None as an observation
# given that it is the identity processor as default
obs = robot.get_observation()
# Get teleop action
raw_action = teleop.get_action()
# Process teleop action through pipeline
teleop_action = teleop_action_processor((raw_action, obs))
teleop_action = teleop_action_processor(raw_action)
# Process action for robot through pipeline
robot_action_to_send = robot_action_processor((teleop_action, obs))
robot_action_to_send = robot_action_processor(teleop_action)
# Send processed action to robot (robot_action_processor.to_output should return dict[str, Any])
_ = robot.send_action(robot_action_to_send)
if display_data:
# Get robot observation
obs = robot.get_observation()
# Process robot observation through pipeline
obs_transition = robot_observation_processor(obs)
@@ -297,8 +297,8 @@ class GamepadController(InputController):
try:
# Read joystick axes
# Left stick X and Y (typically axes 0 and 1)
y_input = self.joystick.get_axis(0) # Up/Down (often inverted)
x_input = self.joystick.get_axis(1) # Left/Right
x_input = self.joystick.get_axis(0) # Left/Right
y_input = self.joystick.get_axis(1) # Up/Down (often inverted)
# Right stick Y (typically axis 3 or 4)
z_input = self.joystick.get_axis(3) # Up/Down for Z
@@ -310,7 +310,7 @@ class GamepadController(InputController):
# Calculate deltas (note: may need to invert axes depending on controller)
delta_x = -x_input * self.x_step_size # Forward/backward
delta_y = -y_input * self.y_step_size # Left/right
delta_y = y_input * self.y_step_size # Left/right
delta_z = -z_input * self.z_step_size # Up/down
return delta_x, delta_y, delta_z
+1 -1
View File
@@ -29,7 +29,7 @@ from lerobot.datasets.transforms import (
SharpnessJitter,
make_transform_from_config,
)
from lerobot.scripts.lerobot_imgtransform_viz import (
from lerobot.scripts.visualize_image_transforms import (
save_all_transforms,
save_each_transform,
)
+1 -1
View File
@@ -15,7 +15,7 @@
# limitations under the License.
import pytest
from lerobot.scripts.lerobot_dataset_viz import visualize_dataset
from lerobot.scripts.visualize_dataset import visualize_dataset
@pytest.mark.skip("TODO: add dummy videos")
+147
View File
@@ -0,0 +1,147 @@
#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import io
import subprocess
import sys
from pathlib import Path
import pytest
from tests.fixtures.constants import DUMMY_REPO_ID
from tests.utils import require_package
def _find_and_replace(text: str, finds_and_replaces: list[tuple[str, str]]) -> str:
for f, r in finds_and_replaces:
assert f in text
text = text.replace(f, r)
return text
# TODO(aliberts): Remove usage of subprocess calls and patch code with fixtures
def _run_script(path):
subprocess.run([sys.executable, path], check=True)
def _read_file(path):
with open(path) as file:
return file.read()
@pytest.mark.skip("TODO Fix and remove subprocess / excec calls")
def test_example_1(tmp_path, lerobot_dataset_factory):
_ = lerobot_dataset_factory(root=tmp_path, repo_id=DUMMY_REPO_ID)
path = "examples/1_load_lerobot_dataset.py"
file_contents = _read_file(path)
file_contents = _find_and_replace(
file_contents,
[
('repo_id = "lerobot/pusht"', f'repo_id = "{DUMMY_REPO_ID}"'),
(
"LeRobotDataset(repo_id",
f"LeRobotDataset(repo_id, root='{str(tmp_path)}'",
),
],
)
exec(file_contents, {})
assert Path("outputs/examples/1_load_lerobot_dataset/episode_0.mp4").exists()
@pytest.mark.skip("TODO Fix and remove subprocess / excec calls")
@require_package("gym_pusht")
def test_examples_basic2_basic3_advanced1():
"""
Train a model with example 3, check the outputs.
Evaluate the trained model with example 2, check the outputs.
Calculate the validation loss with advanced example 1, check the outputs.
"""
### Test example 3
file_contents = _read_file("examples/3_train_policy.py")
# Do fewer steps, use smaller batch, use CPU, and don't complicate things with dataloader workers.
file_contents = _find_and_replace(
file_contents,
[
("training_steps = 5000", "training_steps = 1"),
("num_workers=4", "num_workers=0"),
('device = torch.device("cuda")', 'device = torch.device("cpu")'),
("batch_size=64", "batch_size=1"),
],
)
# Pass empty globals to allow dictionary comprehension https://stackoverflow.com/a/32897127/4391249.
exec(file_contents, {})
for file_name in ["model.safetensors", "config.json"]:
assert Path(f"outputs/train/example_pusht_diffusion/{file_name}").exists()
### Test example 2
file_contents = _read_file("examples/2_evaluate_pretrained_policy.py")
# Do fewer evals, use CPU, and use the local model.
file_contents = _find_and_replace(
file_contents,
[
(
'pretrained_policy_path = Path(snapshot_download("lerobot/diffusion_pusht"))',
"",
),
(
'# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
'pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
),
('device = torch.device("cuda")', 'device = torch.device("cpu")'),
("step += 1", "break"),
],
)
exec(file_contents, {})
assert Path("outputs/eval/example_pusht_diffusion/rollout.mp4").exists()
## Test example 4
file_contents = _read_file("examples/advanced/2_calculate_validation_loss.py")
# Run on a single example from the last episode, use CPU, and use the local model.
file_contents = _find_and_replace(
file_contents,
[
(
'pretrained_policy_path = Path(snapshot_download("lerobot/diffusion_pusht"))',
"",
),
(
'# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
'pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")',
),
("train_episodes = episodes[:num_train_episodes]", "train_episodes = [0]"),
("val_episodes = episodes[num_train_episodes:]", "val_episodes = [1]"),
("num_workers=4", "num_workers=0"),
('device = torch.device("cuda")', 'device = torch.device("cpu")'),
("batch_size=64", "batch_size=1"),
],
)
# Capture the output of the script
output_buffer = io.StringIO()
sys.stdout = output_buffer
exec(file_contents, {})
printed_output = output_buffer.getvalue()
# Restore stdout to its original state
sys.stdout = sys.__stdout__
assert "Average loss on validation set" in printed_output
+1
View File
@@ -194,6 +194,7 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
suite_name = next(iter(envs))
task_id = next(iter(envs[suite_name]))
env = envs[suite_name][task_id]
observation, _ = env.reset(seed=train_cfg.seed)
# apply transform to normalize the observations
+3 -5
View File
@@ -81,8 +81,8 @@ def test_make_act_processor_basic():
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_act_processor_normalization():
@@ -239,9 +239,7 @@ def test_act_processor_save_and_load():
preprocessor.save_pretrained(tmpdir)
# Load preprocessor
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, config_filename="policy_preprocessor.json"
)
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
# Test that loaded processor works
observation = {OBS_STATE: torch.randn(7)}
+2 -8
View File
@@ -290,10 +290,7 @@ def test_save_and_load_pretrained():
# Load pipeline
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="batchpipeline.json",
to_transition=identity_transition,
to_output=identity_transition,
tmp_dir, to_transition=identity_transition, to_output=identity_transition
)
assert loaded_pipeline.name == "BatchPipeline"
@@ -328,10 +325,7 @@ def test_registry_based_save_load():
with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir)
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
to_transition=identity_transition,
to_output=identity_transition,
tmp_dir, to_transition=identity_transition, to_output=identity_transition
)
# Verify the loaded processor works
+1 -3
View File
@@ -250,9 +250,7 @@ def test_classifier_processor_save_and_load():
preprocessor.save_pretrained(tmpdir)
# Load preprocessor
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, config_filename="classifier_preprocessor.json"
)
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
# Test that loaded processor works
observation = {
+1 -147
View File
@@ -324,9 +324,7 @@ def test_save_and_load_pretrained():
robot_processor.save_pretrained(tmpdir)
# Load
loaded_processor = DataProcessorPipeline.from_pretrained(
tmpdir, config_filename="device_test_processor.json"
)
loaded_processor = DataProcessorPipeline.from_pretrained(tmpdir)
assert len(loaded_processor.steps) == 1
loaded_device_processor = loaded_processor.steps[0]
@@ -1015,147 +1013,3 @@ def test_policy_processor_integration():
# Verify action is back on CPU and unnormalized
assert output_result[TransitionKey.ACTION].device.type == "cpu"
assert output_result[TransitionKey.ACTION].shape == (1, 5)
@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
def test_mps_float64_compatibility():
"""Test MPS device compatibility with float64 tensors (automatic conversion to float32)."""
processor = DeviceProcessorStep(device="mps")
# Create tensors with different dtypes, including float64 which MPS doesn't support
observation = {
"observation.float64": torch.randn(5, dtype=torch.float64), # Should be converted to float32
"observation.float32": torch.randn(5, dtype=torch.float32), # Should remain float32
"observation.float16": torch.randn(5, dtype=torch.float16), # Should remain float16
"observation.int64": torch.randint(0, 10, (5,), dtype=torch.int64), # Should remain int64
"observation.bool": torch.tensor([True, False, True], dtype=torch.bool), # Should remain bool
}
action = torch.randn(3, dtype=torch.float64) # Should be converted to float32
reward = torch.tensor(1.0, dtype=torch.float64) # Should be converted to float32
done = torch.tensor(False, dtype=torch.bool) # Should remain bool
truncated = torch.tensor(True, dtype=torch.bool) # Should remain bool
transition = create_transition(
observation=observation, action=action, reward=reward, done=done, truncated=truncated
)
result = processor(transition)
# Check that all tensors are on MPS device
assert result[TransitionKey.OBSERVATION]["observation.float64"].device.type == "mps"
assert result[TransitionKey.OBSERVATION]["observation.float32"].device.type == "mps"
assert result[TransitionKey.OBSERVATION]["observation.float16"].device.type == "mps"
assert result[TransitionKey.OBSERVATION]["observation.int64"].device.type == "mps"
assert result[TransitionKey.OBSERVATION]["observation.bool"].device.type == "mps"
assert result[TransitionKey.ACTION].device.type == "mps"
assert result[TransitionKey.REWARD].device.type == "mps"
assert result[TransitionKey.DONE].device.type == "mps"
assert result[TransitionKey.TRUNCATED].device.type == "mps"
# Check that float64 tensors were automatically converted to float32
assert result[TransitionKey.OBSERVATION]["observation.float64"].dtype == torch.float32
assert result[TransitionKey.ACTION].dtype == torch.float32
assert result[TransitionKey.REWARD].dtype == torch.float32
# Check that other dtypes were preserved
assert result[TransitionKey.OBSERVATION]["observation.float32"].dtype == torch.float32
assert result[TransitionKey.OBSERVATION]["observation.float16"].dtype == torch.float16
assert result[TransitionKey.OBSERVATION]["observation.int64"].dtype == torch.int64
assert result[TransitionKey.OBSERVATION]["observation.bool"].dtype == torch.bool
assert result[TransitionKey.DONE].dtype == torch.bool
assert result[TransitionKey.TRUNCATED].dtype == torch.bool
@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
def test_mps_float64_with_complementary_data():
"""Test MPS float64 conversion with complementary_data tensors."""
processor = DeviceProcessorStep(device="mps")
# Create complementary_data with float64 tensors
complementary_data = {
"task": ["pick_object"],
"index": torch.tensor([42], dtype=torch.int64), # Should remain int64
"task_index": torch.tensor([3], dtype=torch.int64), # Should remain int64
"float64_tensor": torch.tensor([1.5, 2.5], dtype=torch.float64), # Should convert to float32
"float32_tensor": torch.tensor([3.5], dtype=torch.float32), # Should remain float32
}
transition = create_transition(
observation={"observation.state": torch.randn(5, dtype=torch.float64)},
action=torch.randn(3, dtype=torch.float64),
complementary_data=complementary_data,
)
result = processor(transition)
# Check that all tensors are on MPS device
assert result[TransitionKey.OBSERVATION]["observation.state"].device.type == "mps"
assert result[TransitionKey.ACTION].device.type == "mps"
processed_comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
assert processed_comp_data["index"].device.type == "mps"
assert processed_comp_data["task_index"].device.type == "mps"
assert processed_comp_data["float64_tensor"].device.type == "mps"
assert processed_comp_data["float32_tensor"].device.type == "mps"
# Check dtype conversions
assert result[TransitionKey.OBSERVATION]["observation.state"].dtype == torch.float32 # Converted
assert result[TransitionKey.ACTION].dtype == torch.float32 # Converted
assert processed_comp_data["float64_tensor"].dtype == torch.float32 # Converted
assert processed_comp_data["float32_tensor"].dtype == torch.float32 # Unchanged
assert processed_comp_data["index"].dtype == torch.int64 # Unchanged
assert processed_comp_data["task_index"].dtype == torch.int64 # Unchanged
# Check non-tensor data preserved
assert processed_comp_data["task"] == ["pick_object"]
@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
def test_mps_with_explicit_float_dtype():
"""Test MPS device with explicit float_dtype setting."""
# Test that explicit float_dtype still works on MPS
processor = DeviceProcessorStep(device="mps", float_dtype="float16")
observation = {
"observation.float64": torch.randn(
5, dtype=torch.float64
), # First converted to float32, then to float16
"observation.float32": torch.randn(5, dtype=torch.float32), # Converted to float16
"observation.int32": torch.randint(0, 10, (5,), dtype=torch.int32), # Should remain int32
}
action = torch.randn(3, dtype=torch.float64)
transition = create_transition(observation=observation, action=action)
result = processor(transition)
# Check device placement
assert result[TransitionKey.OBSERVATION]["observation.float64"].device.type == "mps"
assert result[TransitionKey.OBSERVATION]["observation.float32"].device.type == "mps"
assert result[TransitionKey.OBSERVATION]["observation.int32"].device.type == "mps"
assert result[TransitionKey.ACTION].device.type == "mps"
# Check that all float tensors end up as float16 (the target dtype)
assert result[TransitionKey.OBSERVATION]["observation.float64"].dtype == torch.float16
assert result[TransitionKey.OBSERVATION]["observation.float32"].dtype == torch.float16
assert result[TransitionKey.ACTION].dtype == torch.float16
# Check that non-float tensors are preserved
assert result[TransitionKey.OBSERVATION]["observation.int32"].dtype == torch.int32
@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
def test_mps_serialization():
"""Test that MPS device processor can be serialized and loaded correctly."""
processor = DeviceProcessorStep(device="mps", float_dtype="float32")
# Test get_config
config = processor.get_config()
assert config == {"device": "mps", "float_dtype": "float32"}
# Test state_dict (should be empty)
state = processor.state_dict()
assert state == {}
# Test load_state_dict (should be no-op)
processor.load_state_dict({})
assert processor.device == "mps"
+3 -5
View File
@@ -84,8 +84,8 @@ def test_make_diffusion_processor_basic():
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_diffusion_processor_with_images():
@@ -258,9 +258,7 @@ def test_diffusion_processor_save_and_load():
preprocessor.save_pretrained(tmpdir)
# Load preprocessor
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, config_filename="policy_preprocessor.json"
)
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
# Test that loaded processor works
observation = {
-341
View File
@@ -1,341 +0,0 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Tests for processor migration detection functionality.
"""
import json
import tempfile
from pathlib import Path
import pytest
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
def test_is_processor_config_valid_configs():
"""Test processor config detection with valid configurations."""
valid_configs = [
{"steps": []}, # Empty steps
{"steps": [{"class": "MyClass"}]}, # Class-based step
{"steps": [{"registry_name": "my_step"}]}, # Registry-based step
{"steps": [{"class": "A"}, {"registry_name": "B"}]}, # Mixed
{"name": "Test", "steps": [{"class": "MyClass"}]}, # With name
]
for i, config in enumerate(valid_configs):
assert DataProcessorPipeline._is_processor_config(config), (
f"Valid config {i} should be detected as processor config: {config}"
)
def test_is_processor_config_invalid_configs():
"""Test processor config detection with invalid configurations."""
invalid_configs = [
{}, # No steps field
{"steps": "not a list"}, # Steps is not a list
{"steps": [{}]}, # Step without class or registry_name
{"steps": ["not a dict"]}, # Step is not a dict
{"steps": [{"other_field": "value"}]}, # Step with wrong fields
{"other_field": "value"}, # Completely different structure
]
for i, config in enumerate(invalid_configs):
assert not DataProcessorPipeline._is_processor_config(config), (
f"Invalid config {i} should not be detected as processor config: {config}"
)
def test_should_suggest_migration_with_processor_config():
"""Test that migration is NOT suggested when processor config exists."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create a valid processor config
processor_config = {
"name": "TestProcessor",
"steps": [
{
"class": "lerobot.processor.normalize.NormalizeStep",
"config": {"mean": 0.0, "std": 1.0},
}
],
}
with open(tmp_path / "processor.json", "w") as f:
json.dump(processor_config, f)
# Should NOT suggest migration (processor config exists)
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
assert not result
def test_should_suggest_migration_with_empty_processor_config():
"""Test that migration is NOT suggested when empty processor config exists."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create an empty processor config
empty_processor_config = {
"name": "EmptyProcessor",
"steps": [], # Empty steps is valid
}
with open(tmp_path / "empty_processor.json", "w") as f:
json.dump(empty_processor_config, f)
# Should NOT suggest migration (processor config exists, even if empty)
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
assert not result
def test_should_suggest_migration_with_model_config_only():
"""Test that migration IS suggested when only model config exists."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create a model config (like old LeRobot format)
model_config = {
"type": "act",
"input_features": {"observation.state": {"shape": [7]}},
"output_features": {"action": {"shape": [7]}},
"hidden_dim": 256,
"n_obs_steps": 1,
"n_action_steps": 1,
}
with open(tmp_path / "config.json", "w") as f:
json.dump(model_config, f)
# SHOULD suggest migration (model config exists but no processor)
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
assert result
def test_should_suggest_migration_no_json_files():
"""Test that migration is NOT suggested when no JSON files exist."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create some non-JSON files
with open(tmp_path / "model.safetensors", "w") as f:
f.write("fake model data")
with open(tmp_path / "README.md", "w") as f:
f.write("# Model README")
# Should NOT suggest migration (no JSON files)
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
assert not result
def test_should_suggest_migration_random_json_files():
"""Test that migration IS suggested when JSON files exist but none are processor configs."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create some random JSON file (not a processor config)
random_config = {"some_field": "some_value", "another_field": 123}
with open(tmp_path / "random.json", "w") as f:
json.dump(random_config, f)
# SHOULD suggest migration (JSON files exist but none are processor configs)
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
assert result
def test_should_suggest_migration_mixed_configs():
"""Test that migration is NOT suggested when processor config exists alongside other configs."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create both a processor config and a model config
processor_config = {"name": "TestProcessor", "steps": [{"registry_name": "normalize_step"}]}
model_config = {"type": "diffusion", "hidden_dim": 512}
with open(tmp_path / "processor.json", "w") as f:
json.dump(processor_config, f)
with open(tmp_path / "config.json", "w") as f:
json.dump(model_config, f)
# Should NOT suggest migration (processor config exists)
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
assert not result
def test_should_suggest_migration_invalid_json():
"""Test that invalid JSON is handled gracefully."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create an invalid JSON file
with open(tmp_path / "invalid.json", "w") as f:
f.write("{ invalid json")
# Create a valid non-processor config
model_config = {"type": "act"}
with open(tmp_path / "model.json", "w") as f:
json.dump(model_config, f)
# SHOULD suggest migration (invalid JSON is ignored, but we have non-processor JSON)
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
assert result
def test_from_pretrained_multiple_json_files_migration_error():
"""Test that multiple JSON files trigger ProcessorMigrationError."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create multiple non-processor configs
model_config = {"type": "act", "hidden_dim": 128}
train_config = {"batch_size": 32, "lr": 0.001}
with open(tmp_path / "config.json", "w") as f:
json.dump(model_config, f)
with open(tmp_path / "train_config.json", "w") as f:
json.dump(train_config, f)
# Should raise ProcessorMigrationError
with pytest.raises(ProcessorMigrationError) as exc_info:
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
# Check the error details
error = exc_info.value
assert str(tmp_path) in str(error.model_path)
assert "migrate_policy_normalization.py" in error.migration_command
assert "not a valid processor configuration" in error.original_error
def test_from_pretrained_no_processor_config_migration_error():
"""Test that missing processor config triggers ProcessorMigrationError."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create a model config but no processor
model_config = {"type": "diffusion", "hidden_dim": 256}
with open(tmp_path / "config.json", "w") as f:
json.dump(model_config, f)
# Should raise ProcessorMigrationError
with pytest.raises(ProcessorMigrationError) as exc_info:
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
# Check the error details
error = exc_info.value
assert str(tmp_path) in str(error.model_path)
assert "migrate_policy_normalization.py" in error.migration_command
assert "not a valid processor configuration" in error.original_error
def test_from_pretrained_valid_processor_no_migration_error():
"""Test that valid processor config does NOT trigger migration error."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create a valid processor config
processor_config = {
"name": "TestProcessor",
"steps": [], # Empty is valid
}
with open(tmp_path / "processor.json", "w") as f:
json.dump(processor_config, f)
# Should succeed and create pipeline
pipeline = DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
assert pipeline is not None
assert pipeline.name == "TestProcessor"
assert len(pipeline) == 0
def test_from_pretrained_no_json_files_no_migration_error():
"""Test that directories with no JSON files don't trigger migration errors."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create some non-JSON files
with open(tmp_path / "model.safetensors", "w") as f:
f.write("fake model data")
# Should raise FileNotFoundError (config file not found)
with pytest.raises(FileNotFoundError, match="not found in directory"):
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
def test_processor_migration_error_creation():
"""Test that ProcessorMigrationError is created correctly."""
model_path = "/path/to/model"
migration_command = "python migrate.py --path /path/to/model"
original_error = "Config not found"
error = ProcessorMigrationError(model_path, migration_command, original_error)
assert error.model_path == model_path
assert error.migration_command == migration_command
assert error.original_error == original_error
assert model_path in str(error)
assert migration_command in str(error)
assert original_error in str(error)
def test_processor_migration_error_attributes():
"""Test that ProcessorMigrationError has correct attributes."""
model_path = Path("/test/path")
migration_command = "python test.py"
original_error = "Test error"
error = ProcessorMigrationError(model_path, migration_command, original_error)
# Test that attributes are accessible
assert hasattr(error, "model_path")
assert hasattr(error, "migration_command")
assert hasattr(error, "original_error")
# Test that it's still an Exception
assert isinstance(error, Exception)
def test_migration_suggestion_raises_error():
"""Test that migration suggestion always raises ProcessorMigrationError."""
with pytest.raises(ProcessorMigrationError) as exc_info:
DataProcessorPipeline._suggest_processor_migration("/test/path", "Test error")
error = exc_info.value
assert "/test/path" in str(error.model_path)
assert "Test error" in error.original_error
assert "migrate_policy_normalization.py" in error.migration_command
def test_migration_error_always_raised_for_invalid_configs():
"""Test that ProcessorMigrationError is always raised for invalid configs."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create a model config
model_config = {"type": "test", "param": "value"}
with open(tmp_path / "config.json", "w") as f:
json.dump(model_config, f)
# Should always raise ProcessorMigrationError
with pytest.raises(ProcessorMigrationError):
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
+1 -3
View File
@@ -1714,9 +1714,7 @@ def test_pipeline_from_pretrained_with_stats_overrides():
# Load the pipeline with stat overrides
overrides = {"normalizer_processor": {"stats": override_stats}}
loaded_pipeline = DataProcessorPipeline.from_pretrained(
temp_dir, config_filename="test_pipeline.json", overrides=overrides
)
loaded_pipeline = DataProcessorPipeline.from_pretrained(temp_dir, overrides=overrides)
# The critical test: the loaded pipeline should use override stats, not original stats
loaded_normalizer = loaded_pipeline.steps[0]
+2 -2
View File
@@ -108,8 +108,8 @@ def test_make_pi0_processor_basic():
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_pi0_newline_processor_single_task():
+38 -81
View File
@@ -543,7 +543,7 @@ def test_save_and_load_pretrained():
assert config["steps"][1]["config"]["counter"] == 10
# Load pipeline
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="testpipeline.json")
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
assert loaded_pipeline.name == "TestPipeline"
assert len(loaded_pipeline) == 2
@@ -571,9 +571,7 @@ def test_step_without_optional_methods():
# Save/load should work even without optional methods
with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir)
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, config_filename="dataprocessorpipeline.json"
)
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
assert len(loaded_pipeline) == 1
@@ -602,9 +600,7 @@ def test_mixed_json_and_tensor_state():
assert state_path.exists()
# Load and verify
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, config_filename="dataprocessorpipeline.json"
)
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
loaded_step = loaded_pipeline.steps[0]
# Check JSON attributes were restored
@@ -892,11 +888,7 @@ def test_from_pretrained_with_overrides():
}
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="testoverrides.json",
overrides=overrides,
to_transition=identity_transition,
to_output=identity_transition,
tmp_dir, overrides=overrides, to_transition=identity_transition, to_output=identity_transition
)
# Verify the pipeline was loaded correctly
@@ -934,11 +926,7 @@ def test_from_pretrained_with_partial_overrides():
# The current implementation applies overrides to ALL steps with the same class name
# Both steps will get the override
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
overrides=overrides,
to_transition=identity_transition,
to_output=identity_transition,
tmp_dir, overrides=overrides, to_transition=identity_transition, to_output=identity_transition
)
transition = create_transition(reward=1.0)
@@ -962,9 +950,7 @@ def test_from_pretrained_invalid_override_key():
overrides = {"NonExistentStep": {"param": "value"}}
with pytest.raises(KeyError, match="Override keys.*do not match any step"):
DataProcessorPipeline.from_pretrained(
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
)
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
def test_from_pretrained_multiple_invalid_override_keys():
@@ -979,9 +965,7 @@ def test_from_pretrained_multiple_invalid_override_keys():
overrides = {"NonExistentStep1": {"param": "value1"}, "NonExistentStep2": {"param": "value2"}}
with pytest.raises(KeyError) as exc_info:
DataProcessorPipeline.from_pretrained(
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
)
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
error_msg = str(exc_info.value)
assert "NonExistentStep1" in error_msg
@@ -1001,11 +985,7 @@ def test_from_pretrained_registered_step_override():
overrides = {"registered_mock_step": {"value": 999, "device": "cuda"}}
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
overrides=overrides,
to_transition=identity_transition,
to_output=identity_transition,
tmp_dir, overrides=overrides, to_transition=identity_transition, to_output=identity_transition
)
# Test that overrides were applied
@@ -1035,11 +1015,7 @@ def test_from_pretrained_mixed_registered_and_unregistered():
}
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
overrides=overrides,
to_transition=identity_transition,
to_output=identity_transition,
tmp_dir, overrides=overrides, to_transition=identity_transition, to_output=identity_transition
)
# Test both steps
@@ -1062,10 +1038,7 @@ def test_from_pretrained_no_overrides():
# Load without overrides
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
to_transition=identity_transition,
to_output=identity_transition,
tmp_dir, to_transition=identity_transition, to_output=identity_transition
)
assert len(loaded_pipeline) == 1
@@ -1087,11 +1060,7 @@ def test_from_pretrained_empty_overrides():
# Load with empty overrides
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
overrides={},
to_transition=identity_transition,
to_output=identity_transition,
tmp_dir, overrides={}, to_transition=identity_transition, to_output=identity_transition
)
assert len(loaded_pipeline) == 1
@@ -1119,9 +1088,7 @@ def test_from_pretrained_override_instantiation_error():
}
with pytest.raises(ValueError, match="Failed to instantiate processor step"):
DataProcessorPipeline.from_pretrained(
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
)
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
def test_from_pretrained_with_state_and_overrides():
@@ -1145,9 +1112,7 @@ def test_from_pretrained_with_state_and_overrides():
}
}
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
)
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
loaded_step = loaded_pipeline.steps[0]
# Check that config overrides were applied
@@ -1175,9 +1140,7 @@ def test_from_pretrained_override_error_messages():
overrides = {"WrongStepName": {"param": "value"}}
with pytest.raises(KeyError) as exc_info:
DataProcessorPipeline.from_pretrained(
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
)
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
error_msg = str(exc_info.value)
assert "WrongStepName" in error_msg
@@ -1362,21 +1325,21 @@ def test_multiple_processors_same_directory():
assert len(loaded_post) == 1
def test_explicit_config_filename_loading():
"""Test explicit config filename loading (no more auto-detection)."""
def test_auto_detect_single_config():
"""Test automatic config detection when there's only one JSON file."""
step = MockStepWithTensorState()
pipeline = DataProcessorPipeline([step], name="SingleConfig")
with tempfile.TemporaryDirectory() as tmp_dir:
pipeline.save_pretrained(tmp_dir)
# Load with explicit config_filename (now required)
loaded = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="singleconfig.json")
# Load without specifying config_filename
loaded = DataProcessorPipeline.from_pretrained(tmp_dir)
assert loaded.name == "SingleConfig"
def test_explicit_config_selection_with_multiple_configs():
"""Test explicit config selection when multiple configs exist."""
def test_error_multiple_configs_no_filename():
"""Test error when multiple configs exist and no filename specified."""
proc1 = DataProcessorPipeline([MockStep()], name="processor1")
proc2 = DataProcessorPipeline([MockStep()], name="processor2")
@@ -1384,12 +1347,9 @@ def test_explicit_config_selection_with_multiple_configs():
proc1.save_pretrained(tmp_dir)
proc2.save_pretrained(tmp_dir)
# Can load specific configs explicitly
loaded1 = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="processor1.json")
loaded2 = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="processor2.json")
assert loaded1.name == "processor1"
assert loaded2.name == "processor2"
# Should raise error
with pytest.raises(ValueError, match="Multiple .json files found"):
DataProcessorPipeline.from_pretrained(tmp_dir)
def test_state_file_naming_with_indices():
@@ -1523,7 +1483,6 @@ def test_override_with_nested_config():
# Load with nested override
loaded = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
overrides={"complex_config_step": {"nested_config": {"level1": {"level2": "overridden"}}}},
to_transition=identity_transition,
to_output=identity_transition,
@@ -1548,7 +1507,6 @@ def test_override_preserves_defaults():
# Override only one parameter
loaded = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
overrides={
"MockStepWithNonSerializableParam": {
"multiplier": 5.0 # Only override multiplier
@@ -1578,9 +1536,7 @@ def test_override_type_validation():
}
with pytest.raises(ValueError, match="Failed to instantiate"):
DataProcessorPipeline.from_pretrained(
tmp_dir, config_filename="dataprocessorpipeline.json", overrides=overrides
)
DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides)
def test_override_with_callables():
@@ -1631,7 +1587,6 @@ def test_override_with_callables():
# Load with callable override
loaded = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
overrides={"callable_step": {"transform_fn": double_values}},
to_transition=identity_transition,
to_output=identity_transition,
@@ -1656,9 +1611,7 @@ def test_override_multiple_same_class_warning():
# Override affects all instances of the class
loaded = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
overrides={"MockStepWithNonSerializableParam": {"multiplier": 10.0}},
tmp_dir, overrides={"MockStepWithNonSerializableParam": {"multiplier": 10.0}}
)
# Both steps get the same override
@@ -1763,9 +1716,7 @@ def test_override_with_device_strings():
# Override device
if torch.cuda.is_available():
loaded = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="dataprocessorpipeline.json",
overrides={"device_aware_step": {"device": "cuda:0"}},
tmp_dir, overrides={"device_aware_step": {"device": "cuda:0"}}
)
loaded_step = loaded.steps[0]
@@ -1783,13 +1734,19 @@ def test_from_pretrained_nonexistent_path():
# Test with an invalid local path - should raise FileNotFoundError
with pytest.raises(FileNotFoundError):
DataProcessorPipeline.from_pretrained("/path/that/does/not/exist", config_filename="processor.json")
DataProcessorPipeline.from_pretrained("/path/that/does/not/exist")
# Test with a path that doesn't exist as a directory
with pytest.raises(FileNotFoundError):
DataProcessorPipeline.from_pretrained("user/repo/extra/path", config_filename="processor.json")
DataProcessorPipeline.from_pretrained("user/repo/extra/path")
# Test with a non-existent Hub repo
# Test with a Hub repo without specifying config_filename (should raise ValueError)
with pytest.raises(
ValueError, match="When loading from Hugging Face Hub, 'config_filename' must be specified"
):
DataProcessorPipeline.from_pretrained("nonexistent-user/nonexistent-repo")
# Test with a non-existent Hub repo when config_filename is provided
with pytest.raises((FileNotFoundError, HfHubHTTPError)):
DataProcessorPipeline.from_pretrained(
"nonexistent-user/nonexistent-repo", config_filename="processor.json"
@@ -1797,9 +1754,9 @@ def test_from_pretrained_nonexistent_path():
# Test with a local directory that exists but has no config files
with tempfile.TemporaryDirectory() as tmp_dir:
# Since the directory exists but has no config, it will raise FileNotFoundError
# Since the directory exists but has no config, it will try Hub and fail
with pytest.raises(FileNotFoundError):
DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="processor.json")
DataProcessorPipeline.from_pretrained(tmp_dir)
def test_save_load_with_custom_converter_functions():
@@ -1836,7 +1793,7 @@ def test_save_load_with_custom_converter_functions():
pipeline.save_pretrained(tmp_dir)
# Load - should use default converters
loaded = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="dataprocessorpipeline.json")
loaded = DataProcessorPipeline.from_pretrained(tmp_dir)
# Verify it uses default converters by checking with standard batch format
batch = {
@@ -1,259 +0,0 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Tests for DataProcessorPipeline.from_pretrained helper methods.
These tests focus on the individual private methods that were extracted from
the main from_pretrained method to improve modularity and testability.
"""
import json
import tempfile
from pathlib import Path
import pytest
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
# Simplified Config Loading Tests
def test_load_config_directory():
"""Test loading config from directory."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create a config file
config_file = tmp_path / "processor.json"
test_config = {"name": "TestProcessor", "steps": []}
config_file.write_text(json.dumps(test_config))
# Load from directory
loaded_config, base_path = DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
assert loaded_config == test_config
assert base_path == tmp_path
def test_load_config_single_file():
"""Test loading config from a single file path."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create a config file
config_file = tmp_path / "processor.json"
test_config = {"name": "TestProcessor", "steps": []}
config_file.write_text(json.dumps(test_config))
# Load using file path directly
loaded_config, base_path = DataProcessorPipeline._load_config(
str(config_file), "any_filename_ignored", {}
)
assert loaded_config == test_config
assert base_path == tmp_path
def test_load_config_directory_file_not_found():
"""Test directory loading when config file doesn't exist."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Directory exists but no processor.json
with pytest.raises(FileNotFoundError, match="not found in directory"):
DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
def test_load_config_directory_with_migration_detection():
"""Test that missing config triggers migration detection."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create old-style config to trigger migration
(tmp_path / "config.json").write_text(json.dumps({"type": "act"}))
# Try to load processor.json (doesn't exist), should trigger migration
with pytest.raises(ProcessorMigrationError):
DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
def test_load_config_nonexistent_path_tries_hub():
"""Test that nonexistent paths try Hub (simplified logic)."""
# This path doesn't exist locally, should try Hub
with pytest.raises(FileNotFoundError, match="on the HuggingFace Hub"):
DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
# Config Validation Tests
def test_validate_loaded_config_valid_config():
"""Test validation with valid processor config."""
valid_config = {"name": "TestProcessor", "steps": []}
# Should not raise any exception
DataProcessorPipeline._validate_loaded_config("any-path", valid_config, "processor.json")
def test_validate_loaded_config_invalid_config():
"""Test validation with invalid processor config."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Create non-processor config to trigger migration
(tmp_path / "config.json").write_text(json.dumps({"type": "act"}))
invalid_config = {"type": "act", "hidden_dim": 256}
with pytest.raises(ProcessorMigrationError):
DataProcessorPipeline._validate_loaded_config(str(tmp_path), invalid_config, "config.json")
def test_validate_loaded_config_invalid_config_no_migration():
"""Test validation with invalid config when no migration is detected."""
# Non-directory path (Hub repo) - no migration detection
invalid_config = {"type": "act", "hidden_dim": 256}
with pytest.raises(ValueError, match="not a valid processor configuration"):
DataProcessorPipeline._validate_loaded_config("user/repo", invalid_config, "config.json")
# Step Class Resolution Tests
def test_resolve_step_class_registry_name():
"""Test resolution using registry name."""
from lerobot.processor.pipeline import ProcessorStep, ProcessorStepRegistry
# Register a test step
@ProcessorStepRegistry.register("test_step")
class TestStep(ProcessorStep):
def __call__(self, transition):
return transition
def transform_features(self, features):
return features
try:
step_entry = {"registry_name": "test_step"}
step_class, step_key = DataProcessorPipeline._resolve_step_class(step_entry)
assert step_class is TestStep
assert step_key == "test_step"
finally:
ProcessorStepRegistry.unregister("test_step")
def test_resolve_step_class_registry_name_not_found():
"""Test resolution with non-existent registry name."""
step_entry = {"registry_name": "nonexistent_step"}
with pytest.raises(ImportError, match="Failed to load processor step from registry"):
DataProcessorPipeline._resolve_step_class(step_entry)
def test_resolve_step_class_import_path():
"""Test resolution using full import path."""
# Use a valid existing class (this should work)
step_entry = {"class": "lerobot.processor.pipeline.ProcessorStep"}
# This should succeed - ProcessorStep can be imported, just not instantiated
step_class, step_key = DataProcessorPipeline._resolve_step_class(step_entry)
from lerobot.processor.pipeline import ProcessorStep
assert step_class is ProcessorStep
assert step_key == "ProcessorStep"
def test_resolve_step_class_invalid_import_path():
"""Test resolution with invalid import path."""
step_entry = {"class": "nonexistent.module.ClassName"}
with pytest.raises(ImportError, match="Failed to load processor step"):
DataProcessorPipeline._resolve_step_class(step_entry)
# Override Validation Tests
def test_validate_overrides_used_all_used():
"""Test validation when all overrides are used."""
# Empty set means all overrides were used
remaining_overrides = set()
config = {"steps": [{"class": "SomeStep"}]}
# Should not raise
DataProcessorPipeline._validate_overrides_used(remaining_overrides, config)
def test_validate_overrides_used_some_unused():
"""Test validation when some overrides are unused."""
remaining_overrides = {"NonExistentStep", "AnotherMissingStep"}
config = {
"steps": [
{"registry_name": "normalize_step"},
{"class": "some.module.TransformStep"},
]
}
with pytest.raises(KeyError, match="Override keys.*do not match any step"):
DataProcessorPipeline._validate_overrides_used(remaining_overrides, config)
def test_validate_overrides_used_helpful_error_message():
"""Test that error message includes available step keys."""
remaining_overrides = {"WrongStep"}
config = {
"steps": [
{"registry_name": "correct_step"},
{"class": "module.path.CorrectClass"},
]
}
with pytest.raises(KeyError) as exc_info:
DataProcessorPipeline._validate_overrides_used(remaining_overrides, config)
error_msg = str(exc_info.value)
assert "Available step keys" in error_msg
assert "correct_step" in error_msg
assert "CorrectClass" in error_msg
# Integration Tests for Simplified Logic
def test_simplified_three_way_loading():
"""Test that the simplified 3-way loading logic works correctly."""
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
# Test 1: Directory loading
config_file = tmp_path / "processor.json"
test_config = {"name": "DirectoryTest", "steps": []}
config_file.write_text(json.dumps(test_config))
loaded_config, base_path = DataProcessorPipeline._load_config(str(tmp_path), "processor.json", {})
assert loaded_config["name"] == "DirectoryTest"
assert base_path == tmp_path
# Test 2: Single file loading
loaded_config, base_path = DataProcessorPipeline._load_config(
str(config_file), "ignored_filename", {}
)
assert loaded_config["name"] == "DirectoryTest"
assert base_path == tmp_path
-525
View File
@@ -1,525 +0,0 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import tempfile
from pathlib import Path
import pytest
import torch
from lerobot.configs.types import FeatureType, PipelineFeatureType
from lerobot.processor import (
DataProcessorPipeline,
PolicyActionToRobotActionProcessorStep,
ProcessorStepRegistry,
RobotActionToPolicyActionProcessorStep,
)
from lerobot.processor.converters import identity_transition
from tests.conftest import assert_contract_is_typed
def test_robot_to_policy_basic_action_conversion():
"""Test basic robot action to policy action conversion."""
motor_names = ["joint1", "joint2", "joint3"]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
robot_action = {
"joint1.pos": 1.0,
"joint2.pos": 2.0,
"joint3.pos": 3.0,
}
policy_action = processor.action(robot_action)
assert isinstance(policy_action, torch.Tensor)
assert policy_action.shape == (3,)
torch.testing.assert_close(policy_action, torch.tensor([1.0, 2.0, 3.0]))
def test_robot_to_policy_action_conversion_preserves_order():
"""Test that motor names order is preserved in conversion."""
motor_names = ["gripper", "arm", "wrist"]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
robot_action = {
"arm.pos": 10.0,
"gripper.pos": 5.0,
"wrist.pos": 15.0,
}
policy_action = processor.action(robot_action)
expected = torch.tensor([5.0, 10.0, 15.0])
torch.testing.assert_close(policy_action, expected)
def test_robot_to_policy_action_conversion_with_floats_and_tensors():
"""Test conversion with mixed float and tensor values."""
motor_names = ["joint1", "joint2"]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
robot_action = {
"joint1.pos": torch.tensor(1.5),
"joint2.pos": 2.5, # Regular float
}
policy_action = processor.action(robot_action)
assert isinstance(policy_action, torch.Tensor)
torch.testing.assert_close(policy_action, torch.tensor([1.5, 2.5]))
def test_robot_to_policy_action_length_mismatch_error():
"""Test error when robot action length doesn't match motor names."""
motor_names = ["joint1", "joint2", "joint3"]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
# Too few actions
robot_action = {"joint1.pos": 1.0, "joint2.pos": 2.0}
with pytest.raises(ValueError, match="Action must have 3 elements, got 2"):
processor.action(robot_action)
robot_action = {
"joint1.pos": 1.0,
"joint2.pos": 2.0,
"joint3.pos": 3.0,
"extra.pos": 4.0,
}
with pytest.raises(ValueError, match="Action must have 3 elements, got 4"):
processor.action(robot_action)
def test_robot_to_policy_missing_motor_key_error():
"""Test error when robot action is missing expected motor keys."""
motor_names = ["joint1", "joint2"]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
robot_action = {
"joint1.pos": 1.0,
"wrong_key.pos": 2.0,
}
with pytest.raises(KeyError):
processor.action(robot_action)
def test_robot_to_policy_transform_features():
"""Test feature transformation for robot to policy action processor."""
motor_names = ["joint1", "joint2", "joint3"]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
features = {
PipelineFeatureType.ACTION: {
"joint1.pos": {"type": FeatureType.ACTION, "shape": (1,)},
"joint2.pos": {"type": FeatureType.ACTION, "shape": (1,)},
"joint3.pos": {"type": FeatureType.ACTION, "shape": (1,)},
"other_data": {"type": FeatureType.ENV, "shape": (1,)},
}
}
transformed = processor.transform_features(features)
assert "action" in transformed[PipelineFeatureType.ACTION]
action_feature = transformed[PipelineFeatureType.ACTION]["action"]
assert action_feature.type == FeatureType.ACTION
assert action_feature.shape == (3,)
assert "joint1.pos" in transformed[PipelineFeatureType.ACTION]
assert "joint2.pos" in transformed[PipelineFeatureType.ACTION]
assert "joint3.pos" in transformed[PipelineFeatureType.ACTION]
assert "other_data" in transformed[PipelineFeatureType.ACTION]
def test_robot_to_policy_get_config():
"""Test configuration serialization."""
motor_names = ["motor1", "motor2"]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
config = processor.get_config()
assert config == {"motor_names": motor_names}
def test_robot_to_policy_state_dict():
"""Test state dict operations."""
processor = RobotActionToPolicyActionProcessorStep(motor_names=["joint1"])
state = processor.state_dict()
assert state == {}
processor.load_state_dict({})
def test_robot_to_policy_single_motor():
"""Test with single motor."""
processor = RobotActionToPolicyActionProcessorStep(motor_names=["single_joint"])
robot_action = {"single_joint.pos": 42.0}
policy_action = processor.action(robot_action)
assert policy_action.shape == (1,)
torch.testing.assert_close(policy_action, torch.tensor([42.0]))
def test_policy_to_robot_basic_action_conversion():
"""Test basic policy action to robot action conversion."""
motor_names = ["joint1", "joint2", "joint3"]
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
policy_action = torch.tensor([1.0, 2.0, 3.0])
robot_action = processor.action(policy_action)
assert isinstance(robot_action, dict)
assert len(robot_action) == 3
expected = {
"joint1.pos": 1.0,
"joint2.pos": 2.0,
"joint3.pos": 3.0,
}
for key, expected_value in expected.items():
assert key in robot_action
actual_value = robot_action[key]
if isinstance(actual_value, torch.Tensor):
actual_value = actual_value.item()
assert actual_value == pytest.approx(expected_value)
def test_policy_to_robot_action_conversion_preserves_order():
"""Test that motor names order corresponds to tensor indices."""
motor_names = ["gripper", "arm", "wrist"]
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
policy_action = torch.tensor([5.0, 10.0, 15.0])
robot_action = processor.action(policy_action)
assert robot_action["gripper.pos"] == pytest.approx(5.0)
assert robot_action["arm.pos"] == pytest.approx(10.0)
assert robot_action["wrist.pos"] == pytest.approx(15.0)
def test_policy_to_robot_action_conversion_with_numpy_input():
"""Test conversion with numpy array input."""
import numpy as np
motor_names = ["joint1", "joint2"]
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
policy_action = np.array([1.5, 2.5])
robot_action = processor.action(policy_action)
assert robot_action["joint1.pos"] == pytest.approx(1.5)
assert robot_action["joint2.pos"] == pytest.approx(2.5)
def test_policy_to_robot_action_length_mismatch_error():
"""Test error when policy action length doesn't match motor names."""
motor_names = ["joint1", "joint2", "joint3"]
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
policy_action = torch.tensor([1.0, 2.0])
with pytest.raises(ValueError, match="Action must have 3 elements, got 2"):
processor.action(policy_action)
policy_action = torch.tensor([1.0, 2.0, 3.0, 4.0])
with pytest.raises(ValueError, match="Action must have 3 elements, got 4"):
processor.action(policy_action)
def test_policy_to_robot_transform_features():
"""Test feature transformation for policy to robot action processor."""
motor_names = ["joint1", "joint2"]
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
features = {
PipelineFeatureType.ACTION: {
"action": {"type": FeatureType.ACTION, "shape": (2,)},
"other_data": {"type": FeatureType.ENV, "shape": (1,)},
}
}
transformed = processor.transform_features(features)
assert "joint1.pos" in transformed[PipelineFeatureType.ACTION]
assert "joint2.pos" in transformed[PipelineFeatureType.ACTION]
for motor in motor_names:
motor_feature = transformed[PipelineFeatureType.ACTION][f"{motor}.pos"]
assert motor_feature.type == FeatureType.ACTION
assert motor_feature.shape == (1,)
assert "action" in transformed[PipelineFeatureType.ACTION]
assert "other_data" in transformed[PipelineFeatureType.ACTION]
def test_policy_to_robot_get_config():
"""Test configuration serialization."""
motor_names = ["motor1", "motor2"]
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
config = processor.get_config()
assert config == {"motor_names": motor_names}
def test_policy_to_robot_state_dict():
"""Test state dict operations."""
processor = PolicyActionToRobotActionProcessorStep(motor_names=["joint1"])
state = processor.state_dict()
assert state == {}
processor.load_state_dict({})
def test_policy_to_robot_single_motor():
"""Test with single motor."""
processor = PolicyActionToRobotActionProcessorStep(motor_names=["single_joint"])
policy_action = torch.tensor([42.0])
robot_action = processor.action(policy_action)
assert len(robot_action) == 1
assert robot_action["single_joint.pos"] == pytest.approx(42.0)
def test_robot_to_policy_registry():
"""Test RobotActionToPolicyActionProcessorStep registry."""
assert "robot_action_to_policy_action_processor" in ProcessorStepRegistry.list()
retrieved_class = ProcessorStepRegistry.get("robot_action_to_policy_action_processor")
assert retrieved_class is RobotActionToPolicyActionProcessorStep
instance = retrieved_class(motor_names=["test"])
assert isinstance(instance, RobotActionToPolicyActionProcessorStep)
assert instance.motor_names == ["test"]
def test_policy_to_robot_registry():
"""Test PolicyActionToRobotActionProcessorStep registry."""
assert "policy_action_to_robot_action_processor" in ProcessorStepRegistry.list()
retrieved_class = ProcessorStepRegistry.get("policy_action_to_robot_action_processor")
assert retrieved_class is PolicyActionToRobotActionProcessorStep
instance = retrieved_class(motor_names=["test"])
assert isinstance(instance, PolicyActionToRobotActionProcessorStep)
assert instance.motor_names == ["test"]
def test_save_and_load_robot_to_policy():
"""Test saving and loading RobotActionToPolicyActionProcessorStep."""
motor_names = ["joint1", "joint2", "joint3"]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
pipeline = DataProcessorPipeline([processor], name="TestRobotToPolicy")
with tempfile.TemporaryDirectory() as tmp_dir:
# Save pipeline
pipeline.save_pretrained(tmp_dir)
# Check config file exists
config_path = Path(tmp_dir) / "testrobottopolicy.json"
assert config_path.exists()
# Load pipeline
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
"testrobottopolicy.json",
to_transition=identity_transition,
to_output=identity_transition,
)
assert loaded_pipeline.name == "TestRobotToPolicy"
assert len(loaded_pipeline) == 1
# Check loaded processor
loaded_processor = loaded_pipeline.steps[0]
assert isinstance(loaded_processor, RobotActionToPolicyActionProcessorStep)
assert loaded_processor.motor_names == motor_names
# Test functionality after loading
robot_action = {"joint1.pos": 1.0, "joint2.pos": 2.0, "joint3.pos": 3.0}
policy_action = loaded_processor.action(robot_action)
torch.testing.assert_close(policy_action, torch.tensor([1.0, 2.0, 3.0]))
def test_save_and_load_policy_to_robot():
"""Test saving and loading PolicyActionToRobotActionProcessorStep."""
motor_names = ["motor_a", "motor_b"]
processor = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
pipeline = DataProcessorPipeline([processor], name="TestPolicyToRobot")
with tempfile.TemporaryDirectory() as tmp_dir:
# Save pipeline
pipeline.save_pretrained(tmp_dir)
# Load pipeline
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
"testpolicytorobot.json",
to_transition=identity_transition,
to_output=identity_transition,
)
loaded_processor = loaded_pipeline.steps[0]
assert isinstance(loaded_processor, PolicyActionToRobotActionProcessorStep)
assert loaded_processor.motor_names == motor_names
policy_action = torch.tensor([10.0, 20.0])
robot_action = loaded_processor.action(policy_action)
assert robot_action["motor_a.pos"] == pytest.approx(10.0)
assert robot_action["motor_b.pos"] == pytest.approx(20.0)
# Integration and chaining tests
def test_round_trip_conversion():
"""Test that robot->policy->robot conversion preserves values."""
motor_names = ["joint1", "joint2", "joint3"]
robot_to_policy = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
policy_to_robot = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
original_robot_action = {
"joint1.pos": 1.5,
"joint2.pos": -2.3,
"joint3.pos": 0.7,
}
policy_action = robot_to_policy.action(original_robot_action)
final_robot_action = policy_to_robot.action(policy_action)
for key in original_robot_action:
original_val = original_robot_action[key]
final_val = final_robot_action[key]
if isinstance(final_val, torch.Tensor):
final_val = final_val.item()
assert final_val == pytest.approx(original_val, abs=1e-6)
def test_chained_processors_in_pipeline():
"""Test both processors chained in a pipeline."""
motor_names = ["joint1", "joint2"]
robot_to_policy = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
policy_to_robot = PolicyActionToRobotActionProcessorStep(motor_names=motor_names)
pipeline = DataProcessorPipeline(
[robot_to_policy, policy_to_robot],
to_transition=identity_transition,
to_output=identity_transition,
)
assert len(pipeline.steps) == 2
assert isinstance(pipeline.steps[0], RobotActionToPolicyActionProcessorStep)
assert isinstance(pipeline.steps[1], PolicyActionToRobotActionProcessorStep)
def test_robot_to_policy_features_contract(policy_feature_factory):
"""Test feature transformation maintains proper typing contract."""
processor = RobotActionToPolicyActionProcessorStep(motor_names=["j1", "j2"])
features = {
PipelineFeatureType.ACTION: {
"j1.pos": policy_feature_factory(FeatureType.ACTION, (1,)),
"j2.pos": policy_feature_factory(FeatureType.ACTION, (1,)),
"other": policy_feature_factory(FeatureType.ENV, (3,)),
}
}
out = processor.transform_features(features.copy())
assert_contract_is_typed(out)
assert "action" in out[PipelineFeatureType.ACTION]
action_feature = out[PipelineFeatureType.ACTION]["action"]
assert action_feature.type == FeatureType.ACTION
assert action_feature.shape == (2,)
def test_policy_to_robot_features_contract(policy_feature_factory):
"""Test feature transformation maintains proper typing contract."""
processor = PolicyActionToRobotActionProcessorStep(motor_names=["m1", "m2", "m3"])
features = {
PipelineFeatureType.ACTION: {
"action": policy_feature_factory(FeatureType.ACTION, (3,)),
"other": policy_feature_factory(FeatureType.ENV, (1,)),
}
}
out = processor.transform_features(features.copy())
assert_contract_is_typed(out)
for motor in ["m1", "m2", "m3"]:
key = f"{motor}.pos"
assert key in out[PipelineFeatureType.ACTION]
motor_feature = out[PipelineFeatureType.ACTION][key]
assert motor_feature.type == FeatureType.ACTION
assert motor_feature.shape == (1,)
def test_empty_motor_names_list():
"""Test behavior with empty motor names list."""
processor = RobotActionToPolicyActionProcessorStep(motor_names=[])
robot_action = {}
policy_action = processor.action(robot_action)
assert isinstance(policy_action, torch.Tensor)
assert policy_action.shape == (0,)
def test_empty_motor_names_list_policy_to_robot():
"""Test PolicyActionToRobotActionProcessorStep with empty motor names."""
processor = PolicyActionToRobotActionProcessorStep(motor_names=[])
policy_action = torch.tensor([])
robot_action = processor.action(policy_action)
assert isinstance(robot_action, dict)
assert len(robot_action) == 0
def test_very_long_motor_names():
"""Test with many motor names."""
motor_names = [f"joint_{i}" for i in range(100)]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
robot_action = {f"joint_{i}.pos": float(i) for i in range(100)}
policy_action = processor.action(robot_action)
assert policy_action.shape == (100,)
expected = torch.tensor([float(i) for i in range(100)])
torch.testing.assert_close(policy_action, expected)
def test_special_characters_in_motor_names():
"""Test with special characters in motor names."""
motor_names = ["motor-1", "motor_2", "motor.3"]
processor = RobotActionToPolicyActionProcessorStep(motor_names=motor_names)
robot_action = {
"motor-1.pos": 1.0,
"motor_2.pos": 2.0,
"motor.3.pos": 3.0,
}
policy_action = processor.action(robot_action)
torch.testing.assert_close(policy_action, torch.tensor([1.0, 2.0, 3.0]))
+2 -7
View File
@@ -232,10 +232,7 @@ def test_save_and_load_pretrained():
# Load pipeline
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir,
config_filename="testrenameprocessorstep.json",
to_transition=identity_transition,
to_output=identity_transition,
tmp_dir, to_transition=identity_transition, to_output=identity_transition
)
assert loaded_pipeline.name == "TestRenameProcessorStep"
@@ -296,9 +293,7 @@ def test_registry_based_save_load():
assert "class" not in config["steps"][0] # Should use registry, not module path
# Load should work
loaded_pipeline = DataProcessorPipeline.from_pretrained(
tmp_dir, config_filename="dataprocessorpipeline.json"
)
loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir)
loaded_processor = loaded_pipeline.steps[0]
assert isinstance(loaded_processor, RenameObservationsProcessorStep)
assert loaded_processor.rename_map == {"key1": "renamed_key1"}
+3 -5
View File
@@ -84,8 +84,8 @@ def test_make_sac_processor_basic():
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_sac_processor_normalization_modes():
@@ -241,9 +241,7 @@ def test_sac_processor_save_and_load():
preprocessor.save_pretrained(tmpdir)
# Load preprocessor
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, config_filename="policy_preprocessor.json"
)
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
# Test that loaded processor works
observation = {OBS_STATE: torch.randn(10)}
+2 -2
View File
@@ -115,8 +115,8 @@ def test_make_smolvla_processor_basic():
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_smolvla_newline_processor_single_task():
+3 -5
View File
@@ -87,8 +87,8 @@ def test_make_tdmpc_processor_basic():
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_tdmpc_processor_normalization():
@@ -269,9 +269,7 @@ def test_tdmpc_processor_save_and_load():
preprocessor.save_pretrained(tmpdir)
# Load preprocessor
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, config_filename="policy_preprocessor.json"
)
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
# Test that loaded processor works
observation = {
+1 -5
View File
@@ -425,10 +425,7 @@ def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer):
# Load processor - tokenizer will be recreated from saved config
loaded_processor = DataProcessorPipeline.from_pretrained(
temp_dir,
config_filename="dataprocessorpipeline.json",
to_transition=identity_transition,
to_output=identity_transition,
temp_dir, to_transition=identity_transition, to_output=identity_transition
)
# Test that loaded processor works
@@ -464,7 +461,6 @@ def test_save_and_load_pretrained_with_tokenizer_object():
# Load processor with tokenizer override (since tokenizer object wasn't saved)
loaded_processor = DataProcessorPipeline.from_pretrained(
temp_dir,
config_filename="dataprocessorpipeline.json",
overrides={"tokenizer_processor": {"tokenizer": mock_tokenizer}},
to_transition=identity_transition,
to_output=identity_transition,
+3 -5
View File
@@ -87,8 +87,8 @@ def test_make_vqbet_processor_basic():
# Check steps in postprocessor
assert len(postprocessor.steps) == 2
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
assert isinstance(postprocessor.steps[1], DeviceProcessorStep)
assert isinstance(postprocessor.steps[0], DeviceProcessorStep)
assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep)
def test_vqbet_processor_with_images():
@@ -264,9 +264,7 @@ def test_vqbet_processor_save_and_load():
preprocessor.save_pretrained(tmpdir)
# Load preprocessor
loaded_preprocessor = DataProcessorPipeline.from_pretrained(
tmpdir, config_filename="policy_preprocessor.json"
)
loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir)
# Test that loaded processor works
observation = {
+5 -5
View File
@@ -65,7 +65,7 @@ def close_service_stub(channel, server):
@require_package("grpc")
def test_establish_learner_connection_success():
from lerobot.rl.actor import establish_learner_connection
from lerobot.scripts.rl.actor import establish_learner_connection
"""Test successful connection establishment."""
stub, _servicer, channel, server = create_learner_service_stub()
@@ -82,7 +82,7 @@ def test_establish_learner_connection_success():
@require_package("grpc")
def test_establish_learner_connection_failure():
from lerobot.rl.actor import establish_learner_connection
from lerobot.scripts.rl.actor import establish_learner_connection
"""Test connection failure."""
stub, servicer, channel, server = create_learner_service_stub()
@@ -101,7 +101,7 @@ def test_establish_learner_connection_failure():
@require_package("grpc")
def test_push_transitions_to_transport_queue():
from lerobot.rl.actor import push_transitions_to_transport_queue
from lerobot.scripts.rl.actor import push_transitions_to_transport_queue
from lerobot.transport.utils import bytes_to_transitions
from tests.transport.test_transport_utils import assert_transitions_equal
@@ -137,7 +137,7 @@ def test_push_transitions_to_transport_queue():
@require_package("grpc")
@pytest.mark.timeout(3) # force cross-platform watchdog
def test_transitions_stream():
from lerobot.rl.actor import transitions_stream
from lerobot.scripts.rl.actor import transitions_stream
"""Test transitions stream functionality."""
shutdown_event = Event()
@@ -169,7 +169,7 @@ def test_transitions_stream():
@require_package("grpc")
@pytest.mark.timeout(3) # force cross-platform watchdog
def test_interactions_stream():
from lerobot.rl.actor import interactions_stream
from lerobot.scripts.rl.actor import interactions_stream
from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
"""Test interactions stream functionality."""
+6 -6
View File
@@ -90,13 +90,13 @@ def cfg():
@require_package("grpc")
@pytest.mark.timeout(10) # force cross-platform watchdog
def test_end_to_end_transitions_flow(cfg):
from lerobot.rl.actor import (
from lerobot.scripts.rl.actor import (
establish_learner_connection,
learner_service_client,
push_transitions_to_transport_queue,
send_transitions,
)
from lerobot.rl.learner import start_learner
from lerobot.scripts.rl.learner import start_learner
from lerobot.transport.utils import bytes_to_transitions
from tests.transport.test_transport_utils import assert_transitions_equal
@@ -152,12 +152,12 @@ def test_end_to_end_transitions_flow(cfg):
@require_package("grpc")
@pytest.mark.timeout(10)
def test_end_to_end_interactions_flow(cfg):
from lerobot.rl.actor import (
from lerobot.scripts.rl.actor import (
establish_learner_connection,
learner_service_client,
send_interactions,
)
from lerobot.rl.learner import start_learner
from lerobot.scripts.rl.learner import start_learner
from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
"""Test complete interactions flow from actor to learner."""
@@ -226,8 +226,8 @@ def test_end_to_end_interactions_flow(cfg):
@pytest.mark.parametrize("data_size", ["small", "large"])
@pytest.mark.timeout(10)
def test_end_to_end_parameters_flow(cfg, data_size):
from lerobot.rl.actor import establish_learner_connection, learner_service_client, receive_policy
from lerobot.rl.learner import start_learner
from lerobot.scripts.rl.actor import establish_learner_connection, learner_service_client, receive_policy
from lerobot.scripts.rl.learner import start_learner
from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
"""Test complete parameter flow from learner to actor, with small and large data."""
+1 -1
View File
@@ -50,7 +50,7 @@ def create_learner_service_stub(
):
import grpc
from lerobot.rl.learner_service import LearnerService
from lerobot.scripts.rl.learner_service import LearnerService
from lerobot.transport import services_pb2_grpc # generated from .proto
"""Fixture to start a LearnerService gRPC server and provide a connected stub."""