feat(processors): use pipelines across the codebase (#1452)

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

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

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

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

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

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

* 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

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

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

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

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

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

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

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

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

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

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

* feat(policies): convert save_policy_to_safetensors with pipeline

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

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

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

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

* refactor(train): Remove unnecessary tensor device handling in training loop

* Refactor`gym_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>

* Remove HILEnvConfig references

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

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

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

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

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

* Revert "refactor(converters): implement unified tensor conversion function (#…" (#1840)

This reverts commit a837685bf870919fc07ada287a71711cebabb1ea.

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

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

* refactor(processors): add transform_features method to various processors (#1843)

* refactor(processors): update transition handling in RewardClassifierProcessor and InverseKinematicsEEToJoints (#1844)

* refactor(processors): unify import statements by consolidating pipeline imports into the main processor module (#1845)

* refactor(processors): add extended api for specialized pipelines (#1848)

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

* chore(processor): rename RobotProcessor -> DataProcessorPipeline (#1850)

* chore(processor): rename specialized processor -> XYZProcessorStep (#1852)

* 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

* 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

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

* chore(processor): rename merge_features -> combine_feature_dicts (#1856)

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

* chore(processor): rename teleop_phone variable names (#1858)

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

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

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

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

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

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

* fix(deps): use in-house rotation utils over scipy throughout the codebase

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

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

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

* fix(processor): recover type inference for use of processors (#1873)

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

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

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

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

* 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

* refactor(eval): specify type parameters for preprocessor and postprocessor in eval_policy function (#1904)

* chore(processor): remove action prefixes (#1905)

* 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

* 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

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

* refactor(processor): phone processor is now an RobotActionProcessorStep

* fix(processor): use subprocessors in AddBatchDimensionProcessorStep only if we have the ingredients

* fix(robots): remove action prefix hard-coded in teleop keyboard and gamepad

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

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

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

* debug(scripts): simplify record with processors (#1918)

Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>

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

* fixes for processors used in phone teleop

* fixes for rotation matrix

* add empty obs and act in create_initial_features

* use observation instead of obs

* docs(processor): update docstrings pipeline (#1920)

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

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

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

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

* refactor(processor): transform_features loop + EAFP (#1932)

* fix(processors): make sure nested dict are also shallow copied (#1939)

* 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

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

* fix teleop, record and eval (#1940)

* fix cmd record, eval

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

* test(processor): fix batch expectation

* 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

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

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

* chore(processors): tokenizers raises and remove tensor conversion (#1949)

* chore(processor): remove unused transition_features dict

* feat(ee): add so100_to_so100_EE replay and evaluate examples

* 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

* test(async): fix feature manipulation (#1957)

* test(async): fix feature manipulation

* chore(processor): remove unused functions

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

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

* fix(processor): enforce signatures

* chore(processor): update comments in record.py

* test(processor): fix isinstance and cuda test

* modify phone docs

* fix(processor): reorder output steps to ensure correct processing sequence (#1961)

- Moved DeviceProcessorStep to the end of the output steps in multiple processor files to maintain the intended processing order.
- Updated corresponding tests to reflect the change in step order.

* fix(processors): assumptions for robot_action_processor & teleop_action_processor (#1964)

* fix(processors): new assumptions pipeline

* fix(processors): ee jj phone teleop replay record working

* chore(processors): update comments and default vars

* chore(processor): remove unnecessary copy

* chore(processor): added todo assumption gripper

* fix(processors): eval using detected device

* finish phone docs

* fix correct image link

* feat(processor): implement migration detection and error handling for  processor configurations (#1968)

* feat(processor): implement migration detection and error handling for processor configurations

- Added ProcessorMigrationError to handle migration requirements for old model formats.
- Enhanced DataProcessorPipeline.from_pretrained to include robust migration detection logic.
- Implemented methods for resolving configuration sources, validating loaded configs, and checking for valid processor configurations.
- Introduced comprehensive tests for migration detection and configuration validation to ensure correct behavior.

* refactor(processor): simplify loading logic and enhance migration detection

- Refactored DataProcessorPipeline to implement a simplified three-way loading strategy for configuration files.
- Introduced explicit config_filename parameter to avoid ambiguity during loading.
- Updated ProcessorMigrationError to provide clearer error messages for migration requirements.
- Enhanced tests to cover new loading logic and ensure proper migration detection.
- Removed deprecated methods related to config source resolution.

* fix(processor) RL (#1953)

* fix(gym_manipulator) general fixes to make it compitable

* fix for dataset v3.0

* fix for gym_manipulator

* add map policy action to robot action wrappers in a seperate scripts

* added unittest for policy to robot bridge

* fixes for gripper penalty

* fix style

* fix gamepad controller

* fixes for sim teleop

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

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

* modify numpy2torch to a regular processor as a quick fix

* missing imports?!

* - Removed the use of `AddRobotObservationAsComplimentaryData` from `gym_manipulator` and thus the codebase
- Added get_raw_joint_positions functions to RobotEnv
- Pass raw_joint_positions as input to the action_pipeline in `gym_manipulator`
- Add `InverseKinematicsRLStep` to be tailored towards the need of RL which requires the use of the IK solution as the main reference point of the control loop
- Added the option `use_ik_solution` in `EEReferenceDelta` step to rely on the ik solution rather than the joint values

* -Updated links to all the config files to place them in the new repo with configs compatible with the pipeline

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

* fix(tests): update test cases for loading pipelines with specific config filenames

- Modified test cases to include explicit configuration filenames when loading pipelines in `test_policy_robot_bridge.py`.
- Ensured that the tests reflect the correct loading behavior for both robot-to-policy and policy-to-robot transitions.

* fix(examples): train mps processor (#1970)

* fix(examples): train mps processor

* fix(processor): add MPS compatibility for float64 tensors

- Implemented a workaround to convert float64 tensors to float32 when using the MPS device, as MPS does not support float64.
- Added unit tests to verify the automatic conversion of float64 tensors to float32 and ensure compatibility with various tensor types on the MPS device.

---------

Co-authored-by: AdilZouitine <adilzouitinegm@gmail.com>

---------

Signed-off-by: Adil Zouitine <adilzouitinegm@gmail.com>
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.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: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
Co-authored-by: Steven Palma <steven.palma@huggingface.co>
Co-authored-by: Pepijn <pepijn@huggingface.co>

src/lerobot/configs/policies.py

@@ -26,7 +26,7 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError
 
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.constants import ACTION, OBS_STATE
 from lerobot.optim.optimizers import OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
@@ -53,7 +53,6 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     """
 
     n_obs_steps: int = 1
-    normalization_mapping: dict[str, NormalizationMode] = field(default_factory=dict)
 
     input_features: dict[str, PolicyFeature] = field(default_factory=dict)
     output_features: dict[str, PolicyFeature] = field(default_factory=dict)

src/lerobot/configs/types.py

@@ -24,6 +24,12 @@ class FeatureType(str, Enum):
     ENV = "ENV"
     ACTION = "ACTION"
     REWARD = "REWARD"
+    LANGUAGE = "LANGUAGE"
+
+
+class PipelineFeatureType(str, Enum):
+    ACTION = "ACTION"
+    OBSERVATION = "OBSERVATION"
 
 
 class NormalizationMode(str, Enum):

src/lerobot/constants.py

@@ -21,8 +21,14 @@ OBS_ENV_STATE = "observation.environment_state"
 OBS_STATE = "observation.state"
 OBS_IMAGE = "observation.image"
 OBS_IMAGES = "observation.images"
+OBS_LANGUAGE = "observation.language"
 ACTION = "action"
 REWARD = "next.reward"
+TRUNCATED = "next.truncated"
+DONE = "next.done"
+
+OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
+OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
 
 ROBOTS = "robots"
 ROBOT_TYPE = "robot_type"
@@ -39,6 +45,9 @@ OPTIMIZER_STATE = "optimizer_state.safetensors"
 OPTIMIZER_PARAM_GROUPS = "optimizer_param_groups.json"
 SCHEDULER_STATE = "scheduler_state.json"
 
+POLICY_PREPROCESSOR_DEFAULT_NAME = "policy_preprocessor"
+POLICY_POSTPROCESSOR_DEFAULT_NAME = "policy_postprocessor"
+
 if "LEROBOT_HOME" in os.environ:
     raise ValueError(
         f"You have a 'LEROBOT_HOME' environment variable set to '{os.getenv('LEROBOT_HOME')}'.\n"

src/lerobot/datasets/pipeline_features.py

@@ -0,0 +1,141 @@
+# 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 re
+from collections.abc import Sequence
+from typing import Any
+
+from lerobot.configs.types import PipelineFeatureType
+from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
+from lerobot.datasets.utils import hw_to_dataset_features
+from lerobot.processor import DataProcessorPipeline
+
+
+def create_initial_features(
+    action: dict[str, Any] | None = None, observation: dict[str, Any] | None = None
+) -> dict[PipelineFeatureType, dict[str, Any]]:
+    """
+    Creates the initial features dict for the dataset from action and observation specs.
+
+    Args:
+        action: A dictionary of action feature names to their types/shapes.
+        observation: A dictionary of observation feature names to their types/shapes.
+
+    Returns:
+        The initial features dictionary structured by PipelineFeatureType.
+    """
+    features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
+    if action:
+        features[PipelineFeatureType.ACTION] = action
+    if observation:
+        features[PipelineFeatureType.OBSERVATION] = observation
+    return features
+
+
+# Helper to filter state/action keys based on regex patterns.
+def should_keep(key: str, patterns: tuple[str]) -> bool:
+    if patterns is None:
+        return True
+    return any(re.search(pat, key) for pat in patterns)
+
+
+def strip_prefix(key: str, prefixes_to_strip: tuple[str]) -> str:
+    for prefix in prefixes_to_strip:
+        if key.startswith(prefix):
+            return key[len(prefix) :]
+    return key
+
+
+# Define prefixes to strip from feature keys for clean names.
+# Handles both fully qualified (e.g., "action.state") and short (e.g., "state") forms.
+PREFIXES_TO_STRIP = tuple(
+    f"{token}." for const in (ACTION, OBS_STATE, OBS_IMAGES) for token in (const, const.split(".")[-1])
+)
+
+
+def aggregate_pipeline_dataset_features(
+    pipeline: DataProcessorPipeline,
+    initial_features: dict[PipelineFeatureType, dict[str, Any]],
+    *,
+    use_videos: bool = True,
+    patterns: Sequence[str] | None = None,
+) -> dict[str, dict]:
+    """
+    Aggregates and filters pipeline features to create a dataset-ready features dictionary.
+
+    This function transforms initial features using the pipeline, categorizes them as action or observations
+    (image or state), filters them based on `use_videos` and `patterns`, and finally
+    formats them for use with a Hugging Face LeRobot Dataset.
+
+    Args:
+        pipeline: The DataProcessorPipeline to apply.
+        initial_features: A dictionary of raw feature specs for actions and observations.
+        use_videos: If False, image features are excluded.
+        patterns: A sequence of regex patterns to filter action and state features.
+                  Image features are not affected by this filter.
+
+    Returns:
+        A dictionary of features formatted for a Hugging Face LeRobot Dataset.
+    """
+    all_features = pipeline.transform_features(initial_features)
+
+    # Intermediate storage for categorized and filtered features.
+    processed_features: dict[str, dict[str, Any]] = {
+        "action": {},
+        "observation": {},
+    }
+    images_token = OBS_IMAGES.split(".")[-1]
+
+    # Iterate through all features transformed by the pipeline.
+    for ptype, feats in all_features.items():
+        if ptype not in [PipelineFeatureType.ACTION, PipelineFeatureType.OBSERVATION]:
+            continue
+
+        for key, value in feats.items():
+            # 1. Categorize the feature.
+            is_action = ptype == PipelineFeatureType.ACTION
+            # Observations are classified as images if their key matches image-related tokens or if the shape of the feature is 3.
+            # All other observations are treated as state.
+            is_image = not is_action and (
+                (isinstance(value, tuple) and len(value) == 3)
+                or (
+                    key.startswith(f"{OBS_IMAGES}.")
+                    or key.startswith(f"{images_token}.")
+                    or f".{images_token}." in key
+                )
+            )
+
+            # 2. Apply filtering rules.
+            if is_image and not use_videos:
+                continue
+            if not is_image and not should_keep(key, patterns):
+                continue
+
+            # 3. Add the feature to the appropriate group with a clean name.
+            name = strip_prefix(key, PREFIXES_TO_STRIP)
+            if is_action:
+                processed_features["action"][name] = value
+            else:
+                processed_features["observation"][name] = value
+
+    # Convert the processed features into the final dataset format.
+    dataset_features = {}
+    if processed_features["action"]:
+        dataset_features.update(hw_to_dataset_features(processed_features["action"], ACTION, use_videos))
+    if processed_features["observation"]:
+        dataset_features.update(
+            hw_to_dataset_features(processed_features["observation"], "observation", use_videos)
+        )
+
+    return dataset_features

src/lerobot/datasets/utils.py

@@ -150,14 +150,20 @@ def get_video_size_in_mb(mp4_path: Path) -> float:
 
 
 def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
-    """Flatten a nested dictionary structure by collapsing nested keys into one key with a separator.
+    """Flatten a nested dictionary by joining keys with a separator.
 
-    For example:
-    ```
-    >>> dct = {"a": {"b": 1, "c": {"d": 2}}, "e": 3}`
-    >>> print(flatten_dict(dct))
-    {"a/b": 1, "a/c/d": 2, "e": 3}
-    ```
+    Example:
+        >>> dct = {"a": {"b": 1, "c": {"d": 2}}, "e": 3}
+        >>> print(flatten_dict(dct))
+        {'a/b': 1, 'a/c/d': 2, 'e': 3}
+
+    Args:
+        d (dict): The dictionary to flatten.
+        parent_key (str): The base key to prepend to the keys in this level.
+        sep (str): The separator to use between keys.
+
+    Returns:
+        dict: A flattened dictionary.
     """
     items = []
     for k, v in d.items():
@@ -170,6 +176,20 @@ def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
 
 
 def unflatten_dict(d: dict, sep: str = "/") -> dict:
+    """Unflatten a dictionary with delimited keys into a nested dictionary.
+
+    Example:
+        >>> flat_dct = {"a/b": 1, "a/c/d": 2, "e": 3}
+        >>> print(unflatten_dict(flat_dct))
+        {'a': {'b': 1, 'c': {'d': 2}}, 'e': 3}
+
+    Args:
+        d (dict): A dictionary with flattened keys.
+        sep (str): The separator used in the keys.
+
+    Returns:
+        dict: A nested dictionary.
+    """
     outdict = {}
     for key, value in d.items():
         parts = key.split(sep)
@@ -183,6 +203,19 @@ def unflatten_dict(d: dict, sep: str = "/") -> dict:
 
 
 def serialize_dict(stats: dict[str, torch.Tensor | np.ndarray | dict]) -> dict:
+    """Serialize a dictionary containing tensors or numpy arrays to be JSON-compatible.
+
+    Converts torch.Tensor, np.ndarray, and np.generic types to lists or native Python types.
+
+    Args:
+        stats (dict): A dictionary that may contain non-serializable numeric types.
+
+    Returns:
+        dict: A dictionary with all values converted to JSON-serializable types.
+
+    Raises:
+        NotImplementedError: If a value has an unsupported type.
+    """
     serialized_dict = {}
     for key, value in flatten_dict(stats).items():
         if isinstance(value, (torch.Tensor, np.ndarray)):
@@ -199,6 +232,17 @@ def serialize_dict(stats: dict[str, torch.Tensor | np.ndarray | dict]) -> dict:
 
 
 def embed_images(dataset: datasets.Dataset) -> datasets.Dataset:
+    """Embed image bytes into the dataset table before saving to Parquet.
+
+    This function prepares a Hugging Face dataset for serialization by converting
+    image objects into an embedded format that can be stored in Arrow/Parquet.
+
+    Args:
+        dataset (datasets.Dataset): The input dataset, possibly containing image features.
+
+    Returns:
+        datasets.Dataset: The dataset with images embedded in the table storage.
+    """
     # Embed image bytes into the table before saving to parquet
     format = dataset.format
     dataset = dataset.with_format("arrow")
@@ -208,11 +252,27 @@ def embed_images(dataset: datasets.Dataset) -> datasets.Dataset:
 
 
 def load_json(fpath: Path) -> Any:
+    """Load data from a JSON file.
+
+    Args:
+        fpath (Path): Path to the JSON file.
+
+    Returns:
+        Any: The data loaded from the JSON file.
+    """
     with open(fpath) as f:
         return json.load(f)
 
 
 def write_json(data: dict, fpath: Path) -> None:
+    """Write data to a JSON file.
+
+    Creates parent directories if they don't exist.
+
+    Args:
+        data (dict): The dictionary to write.
+        fpath (Path): The path to the output JSON file.
+    """
     fpath.parent.mkdir(exist_ok=True, parents=True)
     with open(fpath, "w") as f:
         json.dump(data, f, indent=4, ensure_ascii=False)
@@ -223,6 +283,16 @@ def write_info(info: dict, local_dir: Path) -> None:
 
 
 def load_info(local_dir: Path) -> dict:
+    """Load dataset info metadata from its standard file path.
+
+    Also converts shape lists to tuples for consistency.
+
+    Args:
+        local_dir (Path): The root directory of the dataset.
+
+    Returns:
+        dict: The dataset information dictionary.
+    """
     info = load_json(local_dir / INFO_PATH)
     for ft in info["features"].values():
         ft["shape"] = tuple(ft["shape"])
@@ -230,16 +300,40 @@ def load_info(local_dir: Path) -> dict:
 
 
 def write_stats(stats: dict, local_dir: Path) -> None:
+    """Serialize and write dataset statistics to their standard file path.
+
+    Args:
+        stats (dict): The statistics dictionary (can contain tensors/numpy arrays).
+        local_dir (Path): The root directory of the dataset.
+    """
     serialized_stats = serialize_dict(stats)
     write_json(serialized_stats, local_dir / STATS_PATH)
 
 
 def cast_stats_to_numpy(stats: dict) -> dict[str, dict[str, np.ndarray]]:
+    """Recursively cast numerical values in a stats dictionary to numpy arrays.
+
+    Args:
+        stats (dict): The statistics dictionary.
+
+    Returns:
+        dict: The statistics dictionary with values cast to numpy arrays.
+    """
     stats = {key: np.array(value) for key, value in flatten_dict(stats).items()}
     return unflatten_dict(stats)
 
 
 def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]] | None:
+    """Load dataset statistics and cast numerical values to numpy arrays.
+
+    Returns None if the stats file doesn't exist.
+
+    Args:
+        local_dir (Path): The root directory of the dataset.
+
+    Returns:
+        A dictionary of statistics or None if the file is not found.
+    """
     if not (local_dir / STATS_PATH).exists():
         return None
     stats = load_json(local_dir / STATS_PATH)
@@ -297,6 +391,18 @@ def backward_compatible_episodes_stats(
 def load_image_as_numpy(
     fpath: str | Path, dtype: np.dtype = np.float32, channel_first: bool = True
 ) -> np.ndarray:
+    """Load an image from a file into a numpy array.
+
+    Args:
+        fpath (str | Path): Path to the image file.
+        dtype (np.dtype): The desired data type of the output array. If floating,
+            pixels are scaled to [0, 1].
+        channel_first (bool): If True, converts the image to (C, H, W) format.
+            Otherwise, it remains in (H, W, C) format.
+
+    Returns:
+        np.ndarray: The image as a numpy array.
+    """
     img = PILImage.open(fpath).convert("RGB")
     img_array = np.array(img, dtype=dtype)
     if channel_first:  # (H, W, C) -> (C, H, W)
@@ -307,10 +413,19 @@ def load_image_as_numpy(
 
 
 def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[torch.Tensor | str]]:
-    """Get a transform function that convert items from Hugging Face dataset (pyarrow)
-    to torch tensors. Importantly, images are converted from PIL, which corresponds to
-    a channel last representation (h w c) of uint8 type, to a torch image representation
-    with channel first (c h w) of float32 type in range [0,1].
+    """Convert a batch from a Hugging Face dataset to torch tensors.
+
+    This transform function converts items from Hugging Face dataset format (pyarrow)
+    to torch tensors. Importantly, images are converted from PIL objects (H, W, C, uint8)
+    to a torch image representation (C, H, W, float32) in the range [0, 1]. Other
+    types are converted to torch.tensor.
+
+    Args:
+        items_dict (dict): A dictionary representing a batch of data from a
+            Hugging Face dataset.
+
+    Returns:
+        dict: The batch with items converted to torch tensors.
     """
     for key in items_dict:
         first_item = items_dict[key][0]
@@ -325,6 +440,14 @@ def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[to
 
 
 def is_valid_version(version: str) -> bool:
+    """Check if a string is a valid PEP 440 version.
+
+    Args:
+        version (str): The version string to check.
+
+    Returns:
+        bool: True if the version string is valid, False otherwise.
+    """
     try:
         packaging.version.parse(version)
         return True
@@ -338,6 +461,18 @@ def check_version_compatibility(
     current_version: str | packaging.version.Version,
     enforce_breaking_major: bool = True,
 ) -> None:
+    """Check for version compatibility between a dataset and the current codebase.
+
+    Args:
+        repo_id (str): The repository ID for logging purposes.
+        version_to_check (str | packaging.version.Version): The version of the dataset.
+        current_version (str | packaging.version.Version): The current version of the codebase.
+        enforce_breaking_major (bool): If True, raise an error on major version mismatch.
+
+    Raises:
+        BackwardCompatibilityError: If the dataset version is from a newer, incompatible
+            major version of the codebase.
+    """
     v_check = (
         packaging.version.parse(version_to_check)
         if not isinstance(version_to_check, packaging.version.Version)
@@ -355,7 +490,14 @@ def check_version_compatibility(
 
 
 def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
-    """Returns available valid versions (branches and tags) on given repo."""
+    """Return available valid versions (branches and tags) on a given Hub repo.
+
+    Args:
+        repo_id (str): The repository ID on the Hugging Face Hub.
+
+    Returns:
+        list[packaging.version.Version]: A list of valid versions found.
+    """
     api = HfApi()
     repo_refs = api.list_repo_refs(repo_id, repo_type="dataset")
     repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags]
@@ -368,9 +510,22 @@ def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
 
 
 def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> str:
-    """
-    Returns the version if available on repo or the latest compatible one.
-    Otherwise, will throw a `CompatibilityError`.
+    """Return the specified version if available on repo, or the latest compatible one.
+
+    If the exact version is not found, it looks for the latest version with the
+    same major version number that is less than or equal to the target minor version.
+
+    Args:
+        repo_id (str): The repository ID on the Hugging Face Hub.
+        version (str | packaging.version.Version): The target version.
+
+    Returns:
+        str: The safe version string (e.g., "v1.2.3") to use as a revision.
+
+    Raises:
+        RevisionNotFoundError: If the repo has no version tags.
+        BackwardCompatibilityError: If only older major versions are available.
+        ForwardCompatibilityError: If only newer major versions are available.
     """
     target_version = (
         packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version
@@ -412,6 +567,17 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
 
 
 def get_hf_features_from_features(features: dict) -> datasets.Features:
+    """Convert a LeRobot features dictionary to a `datasets.Features` object.
+
+    Args:
+        features (dict): A LeRobot-style feature dictionary.
+
+    Returns:
+        datasets.Features: The corresponding Hugging Face `datasets.Features` object.
+
+    Raises:
+        ValueError: If a feature has an unsupported shape.
+    """
     hf_features = {}
     for key, ft in features.items():
         if ft["dtype"] == "video":
@@ -439,6 +605,14 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
 
 
 def _validate_feature_names(features: dict[str, dict]) -> None:
+    """Validate that feature names do not contain invalid characters.
+
+    Args:
+        features (dict): The LeRobot features dictionary.
+
+    Raises:
+        ValueError: If any feature name contains '/'.
+    """
     invalid_features = {name: ft for name, ft in features.items() if "/" in name}
     if invalid_features:
         raise ValueError(f"Feature names should not contain '/'. Found '/' in '{invalid_features}'.")
@@ -447,8 +621,28 @@ def _validate_feature_names(features: dict[str, dict]) -> None:
 def hw_to_dataset_features(
     hw_features: dict[str, type | tuple], prefix: str, use_video: bool = True
 ) -> dict[str, dict]:
+    """Convert hardware-specific features to a LeRobot dataset feature dictionary.
+
+    This function takes a dictionary describing hardware outputs (like joint states
+    or camera image shapes) and formats it into the standard LeRobot feature
+    specification.
+
+    Args:
+        hw_features (dict): Dictionary mapping feature names to their type (float for
+            joints) or shape (tuple for images).
+        prefix (str): The prefix to add to the feature keys (e.g., "observation"
+            or "action").
+        use_video (bool): If True, image features are marked as "video", otherwise "image".
+
+    Returns:
+        dict: A LeRobot features dictionary.
+    """
     features = {}
-    joint_fts = {key: ftype for key, ftype in hw_features.items() if ftype is float}
+    joint_fts = {
+        key: ftype
+        for key, ftype in hw_features.items()
+        if ftype is float or (isinstance(ftype, PolicyFeature) and ftype.type != FeatureType.VISUAL)
+    }
     cam_fts = {key: shape for key, shape in hw_features.items() if isinstance(shape, tuple)}
 
     if joint_fts and prefix == "action":
@@ -479,6 +673,20 @@ def hw_to_dataset_features(
 def build_dataset_frame(
     ds_features: dict[str, dict], values: dict[str, Any], prefix: str
 ) -> dict[str, np.ndarray]:
+    """Construct a single data frame from raw values based on dataset features.
+
+    A "frame" is a dictionary containing all the data for a single timestep,
+    formatted as numpy arrays according to the feature specification.
+
+    Args:
+        ds_features (dict): The LeRobot dataset features dictionary.
+        values (dict): A dictionary of raw values from the hardware/environment.
+        prefix (str): The prefix to filter features by (e.g., "observation"
+            or "action").
+
+    Returns:
+        dict: A dictionary representing a single frame of data.
+    """
     frame = {}
     for key, ft in ds_features.items():
         if key in DEFAULT_FEATURES or not key.startswith(prefix):
@@ -492,6 +700,21 @@ def build_dataset_frame(
 
 
 def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFeature]:
+    """Convert dataset features to policy features.
+
+    This function transforms the dataset's feature specification into a format
+    that a policy can use, classifying features by type (e.g., visual, state,
+    action) and ensuring correct shapes (e.g., channel-first for images).
+
+    Args:
+        features (dict): The LeRobot dataset features dictionary.
+
+    Returns:
+        dict: A dictionary mapping feature keys to `PolicyFeature` objects.
+
+    Raises:
+        ValueError: If an image feature does not have a 3D shape.
+    """
     # TODO(aliberts): Implement "type" in dataset features and simplify this
     policy_features = {}
     for key, ft in features.items():
@@ -522,6 +745,58 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
     return policy_features
 
 
+def combine_feature_dicts(*dicts: dict) -> dict:
+    """Merge LeRobot grouped feature dicts.
+
+    - For 1D numeric specs (dtype not image/video/string) with "names": we merge the names and recompute the shape.
+    - For others (e.g. `observation.images.*`), the last one wins (if they are identical).
+
+    Args:
+        *dicts: A variable number of LeRobot feature dictionaries to merge.
+
+    Returns:
+        dict: A single merged feature dictionary.
+
+    Raises:
+        ValueError: If there's a dtype mismatch for a feature being merged.
+    """
+    out: dict = {}
+    for d in dicts:
+        for key, value in d.items():
+            if not isinstance(value, dict):
+                out[key] = value
+                continue
+
+            dtype = value.get("dtype")
+            shape = value.get("shape")
+            is_vector = (
+                dtype not in ("image", "video", "string")
+                and isinstance(shape, tuple)
+                and len(shape) == 1
+                and "names" in value
+            )
+
+            if is_vector:
+                # Initialize or retrieve the accumulating dict for this feature key
+                target = out.setdefault(key, {"dtype": dtype, "names": [], "shape": (0,)})
+                # Ensure consistent data types across merged entries
+                if "dtype" in target and dtype != target["dtype"]:
+                    raise ValueError(f"dtype mismatch for '{key}': {target['dtype']} vs {dtype}")
+
+                # Merge feature names: append only new ones to preserve order without duplicates
+                seen = set(target["names"])
+                for n in value["names"]:
+                    if n not in seen:
+                        target["names"].append(n)
+                        seen.add(n)
+                # Recompute the shape to reflect the updated number of features
+                target["shape"] = (len(target["names"]),)
+            else:
+                # For images/videos and non-1D entries: override with the latest definition
+                out[key] = value
+    return out
+
+
 def create_empty_dataset_info(
     codebase_version: str,
     fps: int,
@@ -532,6 +807,18 @@ def create_empty_dataset_info(
     data_files_size_in_mb: int | None = None,
     video_files_size_in_mb: int | None = None,
 ) -> dict:
+    """Create a template dictionary for a new dataset's `info.json`.
+
+    Args:
+        codebase_version (str): The version of the LeRobot codebase.
+        fps (int): The frames per second of the data.
+        features (dict): The LeRobot features dictionary for the dataset.
+        use_videos (bool): Whether the dataset will store videos.
+        robot_type (str | None): The type of robot used, if any.
+
+    Returns:
+        dict: A dictionary with the initial dataset metadata.
+    """
     return {
         "codebase_version": codebase_version,
         "robot_type": robot_type,
@@ -552,9 +839,23 @@ def create_empty_dataset_info(
 def check_delta_timestamps(
     delta_timestamps: dict[str, list[float]], fps: int, tolerance_s: float, raise_value_error: bool = True
 ) -> bool:
-    """This will check if all the values in delta_timestamps are multiples of 1/fps +/- tolerance.
-    This is to ensure that these delta_timestamps added to any timestamp from a dataset will themselves be
-    actual timestamps from the dataset.
+    """Check if delta timestamps are multiples of 1/fps +/- tolerance.
+
+    This ensures that adding these delta timestamps to any existing timestamp in
+    the dataset will result in a value that aligns with the dataset's frame rate.
+
+    Args:
+        delta_timestamps (dict): A dictionary where values are lists of time
+            deltas in seconds.
+        fps (int): The frames per second of the dataset.
+        tolerance_s (float): The allowed tolerance in seconds.
+        raise_value_error (bool): If True, raises an error on failure.
+
+    Returns:
+        bool: True if all deltas are valid, False otherwise.
+
+    Raises:
+        ValueError: If any delta is outside the tolerance and `raise_value_error` is True.
     """
     outside_tolerance = {}
     for key, delta_ts in delta_timestamps.items():
@@ -580,6 +881,15 @@ def check_delta_timestamps(
 
 
 def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dict[str, list[int]]:
+    """Convert delta timestamps in seconds to delta indices in frames.
+
+    Args:
+        delta_timestamps (dict): A dictionary of time deltas in seconds.
+        fps (int): The frames per second of the dataset.
+
+    Returns:
+        dict: A dictionary of frame delta indices.
+    """
     delta_indices = {}
     for key, delta_ts in delta_timestamps.items():
         delta_indices[key] = [round(d * fps) for d in delta_ts]
@@ -588,9 +898,17 @@ def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dic
 
 
 def cycle(iterable: Any) -> Iterator[Any]:
-    """The equivalent of itertools.cycle, but safe for Pytorch dataloaders.
+    """Create a dataloader-safe cyclical iterator.
 
-    See https://github.com/pytorch/pytorch/issues/23900 for information on why itertools.cycle is not safe.
+    This is an equivalent of `itertools.cycle` but is safe for use with
+    PyTorch DataLoaders with multiple workers.
+    See https://github.com/pytorch/pytorch/issues/23900 for details.
+
+    Args:
+        iterable: The iterable to cycle over.
+
+    Yields:
+        Items from the iterable, restarting from the beginning when exhausted.
     """
     iterator = iter(iterable)
     while True:
@@ -601,8 +919,14 @@ def cycle(iterable: Any) -> Iterator[Any]:
 
 
 def create_branch(repo_id: str, *, branch: str, repo_type: str | None = None) -> None:
-    """Create a branch on a existing Hugging Face repo. Delete the branch if it already
-    exists before creating it.
+    """Create a branch on an existing Hugging Face repo.
+
+    Deletes the branch if it already exists before creating it.
+
+    Args:
+        repo_id (str): The ID of the repository.
+        branch (str): The name of the branch to create.
+        repo_type (str | None): The type of the repository (e.g., "dataset").
     """
     api = HfApi()
 
@@ -620,9 +944,20 @@ def create_lerobot_dataset_card(
     dataset_info: dict | None = None,
     **kwargs,
 ) -> DatasetCard:
-    """
-    Keyword arguments will be used to replace values in src/lerobot/datasets/card_template.md.
-    Note: If specified, license must be one of https://huggingface.co/docs/hub/repositories-licenses.
+    """Create a `DatasetCard` for a LeRobot dataset.
+
+    Keyword arguments are used to replace values in the card template.
+    Note: If specified, `license` must be a valid license identifier from
+    https://huggingface.co/docs/hub/repositories-licenses.
+
+    Args:
+        tags (list | None): A list of tags to add to the dataset card.
+        dataset_info (dict | None): The dataset's info dictionary, which will
+            be displayed on the card.
+        **kwargs: Additional keyword arguments to populate the card template.
+
+    Returns:
+        DatasetCard: The generated dataset card object.
     """
     card_tags = ["LeRobot"]
 
@@ -675,6 +1010,15 @@ def validate_frame(frame: dict, features: dict) -> None:
 
 
 def validate_features_presence(actual_features: set[str], expected_features: set[str]) -> str:
+    """Check for missing or extra features in a frame.
+
+    Args:
+        actual_features (set[str]): The set of feature names present in the frame.
+        expected_features (set[str]): The set of feature names expected in the frame.
+
+    Returns:
+        str: An error message string if there's a mismatch, otherwise an empty string.
+    """
     error_message = ""
     missing_features = expected_features - actual_features
     extra_features = actual_features - expected_features
@@ -692,6 +1036,19 @@ def validate_features_presence(actual_features: set[str], expected_features: set
 def validate_feature_dtype_and_shape(
     name: str, feature: dict, value: np.ndarray | PILImage.Image | str
 ) -> str:
+    """Validate the dtype and shape of a single feature's value.
+
+    Args:
+        name (str): The name of the feature.
+        feature (dict): The feature specification from the LeRobot features dictionary.
+        value: The value of the feature to validate.
+
+    Returns:
+        str: An error message if validation fails, otherwise an empty string.
+
+    Raises:
+        NotImplementedError: If the feature dtype is not supported for validation.
+    """
     expected_dtype = feature["dtype"]
     expected_shape = feature["shape"]
     if is_valid_numpy_dtype_string(expected_dtype):
@@ -707,6 +1064,17 @@ def validate_feature_dtype_and_shape(
 def validate_feature_numpy_array(
     name: str, expected_dtype: str, expected_shape: list[int], value: np.ndarray
 ) -> str:
+    """Validate a feature that is expected to be a numpy array.
+
+    Args:
+        name (str): The name of the feature.
+        expected_dtype (str): The expected numpy dtype as a string.
+        expected_shape (list[int]): The expected shape.
+        value (np.ndarray): The numpy array to validate.
+
+    Returns:
+        str: An error message if validation fails, otherwise an empty string.
+    """
     error_message = ""
     if isinstance(value, np.ndarray):
         actual_dtype = value.dtype
@@ -726,6 +1094,18 @@ def validate_feature_numpy_array(
 def validate_feature_image_or_video(
     name: str, expected_shape: list[str], value: np.ndarray | PILImage.Image
 ) -> str:
+    """Validate a feature that is expected to be an image or video frame.
+
+    Accepts `np.ndarray` (channel-first or channel-last) or `PIL.Image.Image`.
+
+    Args:
+        name (str): The name of the feature.
+        expected_shape (list[str]): The expected shape (C, H, W).
+        value: The image data to validate.
+
+    Returns:
+        str: An error message if validation fails, otherwise an empty string.
+    """
     # Note: The check of pixels range ([0,1] for float and [0,255] for uint8) is done by the image writer threads.
     error_message = ""
     if isinstance(value, np.ndarray):
@@ -742,12 +1122,35 @@ def validate_feature_image_or_video(
 
 
 def validate_feature_string(name: str, value: str) -> str:
+    """Validate a feature that is expected to be a string.
+
+    Args:
+        name (str): The name of the feature.
+        value (str): The value to validate.
+
+    Returns:
+        str: An error message if validation fails, otherwise an empty string.
+    """
     if not isinstance(value, str):
         return f"The feature '{name}' is expected to be of type 'str', but type '{type(value)}' provided instead.\n"
     return ""
 
 
 def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features: dict) -> None:
+    """Validate the episode buffer before it's written to disk.
+
+    Ensures the buffer has the required keys, contains at least one frame, and
+    has features consistent with the dataset's specification.
+
+    Args:
+        episode_buffer (dict): The buffer containing data for a single episode.
+        total_episodes (int): The current total number of episodes in the dataset.
+        features (dict): The LeRobot features dictionary for the dataset.
+
+    Raises:
+        ValueError: If the buffer is invalid.
+        NotImplementedError: If the episode index is manually set and doesn't match.
+    """
     if "size" not in episode_buffer:
         raise ValueError("size key not found in episode_buffer")
 

src/lerobot/envs/configs.py

@@ -161,35 +161,73 @@ class XarmEnv(EnvConfig):
 
 
 @dataclass
-class VideoRecordConfig:
-    """Configuration for video recording in ManiSkill environments."""
-
-    enabled: bool = False
-    record_dir: str = "videos"
-    trajectory_name: str = "trajectory"
+class ImagePreprocessingConfig:
+    crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
+    resize_size: tuple[int, int] | None = None
 
 
 @dataclass
-class EnvTransformConfig:
-    """Configuration for environment wrappers."""
+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."""
 
-    # 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
-    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
+    display_cameras: bool = False
+
+
+@dataclass
+class GripperConfig:
+    """Configuration for gripper control and penalties."""
+
     use_gripper: bool = True
-    gripper_quantization_threshold: float | None = 0.8
     gripper_penalty: float = 0.0
     gripper_penalty_in_reward: bool = False
 
 
+@dataclass
+class ResetConfig:
+    """Configuration for environment reset behavior."""
+
+    fixed_reset_joint_positions: Any | None = None
+    reset_time_s: float = 5.0
+    control_time_s: float = 20.0
+    terminate_on_success: bool = True
+
+
+@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
+
+
 @EnvConfig.register_subclass(name="gym_manipulator")
 @dataclass
 class HILSerlRobotEnvConfig(EnvConfig):
@@ -197,77 +235,10 @@ class HILSerlRobotEnvConfig(EnvConfig):
 
     robot: RobotConfig | None = None
     teleop: TeleoperatorConfig | None = None
-    wrapper: EnvTransformConfig | None = None
-    fps: int = 10
+    processor: HILSerlProcessorConfig = field(default_factory=HILSerlProcessorConfig)
+
     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,
-        }

src/lerobot/envs/factory.py

@@ -17,7 +17,7 @@ import importlib
 
 import gymnasium as gym
 
-from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, PushtEnv, XarmEnv
+from lerobot.envs.configs import AlohaEnv, EnvConfig, PushtEnv, XarmEnv
 
 
 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
@@ -27,8 +27,6 @@ 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)
     else:
         raise ValueError(f"Policy type '{env_type}' is not available.")
 

src/lerobot/envs/utils.py

@@ -127,9 +127,29 @@ def check_env_attributes_and_types(env: gym.vector.VectorEnv) -> None:
 def add_envs_task(env: gym.vector.VectorEnv, observation: dict[str, Any]) -> dict[str, Any]:
     """Adds task feature to the observation dict with respect to the first environment attribute."""
     if hasattr(env.envs[0], "task_description"):
-        observation["task"] = env.call("task_description")
+        task_result = env.call("task_description")
+
+        if isinstance(task_result, tuple):
+            task_result = list(task_result)
+
+        if not isinstance(task_result, list):
+            raise TypeError(f"Expected task_description to return a list, got {type(task_result)}")
+        if not all(isinstance(item, str) for item in task_result):
+            raise TypeError("All items in task_description result must be strings")
+
+        observation["task"] = task_result
     elif hasattr(env.envs[0], "task"):
-        observation["task"] = env.call("task")
+        task_result = env.call("task")
+
+        if isinstance(task_result, tuple):
+            task_result = list(task_result)
+
+        if not isinstance(task_result, list):
+            raise TypeError(f"Expected task to return a list, got {type(task_result)}")
+        if not all(isinstance(item, str) for item in task_result):
+            raise TypeError("All items in task result must be strings")
+
+        observation["task"] = task_result
     else:  #  For envs without language instructions, e.g. aloha transfer cube and etc.
         num_envs = observation[list(observation.keys())[0]].shape[0]
         observation["task"] = ["" for _ in range(num_envs)]

src/lerobot/policies/__init__.py

@@ -15,6 +15,17 @@
 from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
 from .pi0.configuration_pi0 import PI0Config as PI0Config
+from .pi0.processor_pi0 import Pi0NewLineProcessor
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
+from .smolvla.processor_smolvla import SmolVLANewLineProcessor
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
 from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
+
+__all__ = [
+    "ACTConfig",
+    "DiffusionConfig",
+    "PI0Config",
+    "SmolVLAConfig",
+    "TDMPCConfig",
+    "VQBeTConfig",
+]

src/lerobot/policies/act/modeling_act.py

@@ -35,7 +35,6 @@ from torchvision.ops.misc import FrozenBatchNorm2d
 
 from lerobot.constants import ACTION, OBS_IMAGES
 from lerobot.policies.act.configuration_act import ACTConfig
-from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 
 
@@ -51,27 +50,16 @@ class ACTPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: ACTConfig,
-        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                     the configuration class is used.
-            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
-                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
         super().__init__(config)
         config.validate_features()
         self.config = config
 
-        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
-        self.normalize_targets = Normalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-        self.unnormalize_outputs = Unnormalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-
         self.model = ACT(config)
 
         if config.temporal_ensemble_coeff is not None:
@@ -137,23 +125,19 @@ class ACTPolicy(PreTrainedPolicy):
         """Predict a chunk of actions given environment observations."""
         self.eval()
 
-        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = [batch[key] for key in self.config.image_features]
 
         actions = self.model(batch)[0]
-        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
         return actions
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Run the batch through the model and compute the loss for training or validation."""
-        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = [batch[key] for key in self.config.image_features]
 
-        batch = self.normalize_targets(batch)
         actions_hat, (mu_hat, log_sigma_x2_hat) = self.model(batch)
 
         l1_loss = (

src/lerobot/policies/act/processor_act.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+# Copyright 2024 Tony Z. Zhao and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any
+
+import torch
+
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+from lerobot.policies.act.configuration_act import ACTConfig
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+
+def make_act_pre_post_processors(
+    config: ACTConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """Creates the pre- and post-processing pipelines for the ACT policy.
+
+    The pre-processing pipeline handles normalization, batching, and device placement for the model inputs.
+    The post-processing pipeline handles unnormalization and moves the model outputs back to the CPU.
+
+    Args:
+        config (ACTConfig): The ACT policy configuration object.
+        dataset_stats (dict[str, dict[str, torch.Tensor]] | None): A dictionary containing dataset
+            statistics (e.g., mean and std) used for normalization. Defaults to None.
+
+    Returns:
+        tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the
+        pre-processor pipeline and the post-processor pipeline.
+    """
+
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+            device=config.device,
+        ),
+    ]
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -35,7 +35,6 @@ from torch import Tensor, nn
 
 from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
-from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import (
     get_device_from_parameters,
@@ -57,7 +56,6 @@ class DiffusionPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: DiffusionConfig,
-        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
@@ -70,14 +68,6 @@ class DiffusionPolicy(PreTrainedPolicy):
         config.validate_features()
         self.config = config
 
-        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
-        self.normalize_targets = Normalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-        self.unnormalize_outputs = Unnormalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-
         # queues are populated during rollout of the policy, they contain the n latest observations and actions
         self._queues = None
 
@@ -106,9 +96,6 @@ class DiffusionPolicy(PreTrainedPolicy):
         batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
         actions = self.diffusion.generate_actions(batch)
 
-        # TODO(rcadene): make above methods return output dictionary?
-        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
-
         return actions
 
     @torch.no_grad()
@@ -137,7 +124,6 @@ class DiffusionPolicy(PreTrainedPolicy):
         if ACTION in batch:
             batch.pop(ACTION)
 
-        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
@@ -153,11 +139,9 @@ class DiffusionPolicy(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, None]:
         """Run the batch through the model and compute the loss for training or validation."""
-        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
-        batch = self.normalize_targets(batch)
         loss = self.diffusion.compute_loss(batch)
         # no output_dict so returning None
         return loss, None

src/lerobot/policies/diffusion/processor_diffusion.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python
+
+# Copyright 2024 Columbia Artificial Intelligence, Robotics Lab,
+# and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any
+
+import torch
+
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+
+def make_diffusion_pre_post_processors(
+    config: DiffusionConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Constructs pre-processor and post-processor pipelines for a diffusion policy.
+
+    The pre-processing pipeline prepares the input data for the model by:
+    1. Renaming features.
+    2. Normalizing the input and output features based on dataset statistics.
+    3. Adding a batch dimension.
+    4. Moving the data to the specified device.
+
+    The post-processing pipeline handles the model's output by:
+    1. Moving the data to the CPU.
+    2. Unnormalizing the output features to their original scale.
+
+    Args:
+        config: The configuration object for the diffusion policy,
+            containing feature definitions, normalization mappings, and device information.
+        dataset_stats: A dictionary of statistics used for normalization.
+            Defaults to None.
+
+    Returns:
+        A tuple containing the configured pre-processor and post-processor pipelines.
+    """
+
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/factory.py

@@ -14,12 +14,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
+from __future__ import annotations
 
-from torch import nn
+import logging
+from typing import Any, TypedDict
+
+import torch
+from typing_extensions import Unpack
 
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.datasets.utils import dataset_to_policy_features
 from lerobot.envs.configs import EnvConfig
@@ -34,10 +39,32 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardCla
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
+from lerobot.processor import PolicyAction, PolicyProcessorPipeline
+from lerobot.processor.converters import (
+    batch_to_transition,
+    policy_action_to_transition,
+    transition_to_batch,
+    transition_to_policy_action,
+)
 
 
-def get_policy_class(name: str) -> PreTrainedPolicy:
-    """Get the policy's class and config class given a name (matching the policy class' `name` attribute)."""
+def get_policy_class(name: str) -> type[PreTrainedPolicy]:
+    """
+    Retrieves a policy class by its registered name.
+
+    This function uses dynamic imports to avoid loading all policy classes into memory
+    at once, improving startup time and reducing dependencies.
+
+    Args:
+        name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
+              "vqbet", "pi0", "pi0fast", "sac", "reward_classifier", "smolvla".
+
+    Returns:
+        The policy class corresponding to the given name.
+
+    Raises:
+        NotImplementedError: If the policy name is not recognized.
+    """
     if name == "tdmpc":
         from lerobot.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
 
@@ -79,6 +106,24 @@ def get_policy_class(name: str) -> PreTrainedPolicy:
 
 
 def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
+    """
+    Instantiates a policy configuration object based on the policy type.
+
+    This factory function simplifies the creation of policy configuration objects by
+    mapping a string identifier to the corresponding config class.
+
+    Args:
+        policy_type: The type of the policy. Supported types include "tdmpc",
+                     "diffusion", "act", "vqbet", "pi0", "pi0fast", "sac", "smolvla",
+                     "reward_classifier".
+        **kwargs: Keyword arguments to be passed to the configuration class constructor.
+
+    Returns:
+        An instance of a `PreTrainedConfig` subclass.
+
+    Raises:
+        ValueError: If the `policy_type` is not recognized.
+    """
     if policy_type == "tdmpc":
         return TDMPCConfig(**kwargs)
     elif policy_type == "diffusion":
@@ -101,30 +146,187 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         raise ValueError(f"Policy type '{policy_type}' is not available.")
 
 
+class ProcessorConfigKwargs(TypedDict, total=False):
+    """
+    A TypedDict defining the keyword arguments for processor configuration.
+
+    This provides type hints for the optional arguments passed to `make_pre_post_processors`,
+    improving code clarity and enabling static analysis.
+
+    Attributes:
+        preprocessor_config_filename: The filename for the preprocessor configuration.
+        postprocessor_config_filename: The filename for the postprocessor configuration.
+        preprocessor_overrides: A dictionary of overrides for the preprocessor configuration.
+        postprocessor_overrides: A dictionary of overrides for the postprocessor configuration.
+        dataset_stats: Dataset statistics for normalization.
+    """
+
+    preprocessor_config_filename: str | None
+    postprocessor_config_filename: str | None
+    preprocessor_overrides: dict[str, Any] | None
+    postprocessor_overrides: dict[str, Any] | None
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None
+
+
+def make_pre_post_processors(
+    policy_cfg: PreTrainedConfig,
+    pretrained_path: str | None = None,
+    **kwargs: Unpack[ProcessorConfigKwargs],
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Create or load pre- and post-processor pipelines for a given policy.
+
+    This function acts as a factory. It can either load existing processor pipelines
+    from a pretrained path or create new ones from scratch based on the policy
+    configuration. Each policy type has a dedicated factory function for its
+    processors (e.g., `make_tdmpc_pre_post_processors`).
+
+    Args:
+        policy_cfg: The configuration of the policy for which to create processors.
+        pretrained_path: An optional path to load pretrained processor pipelines from.
+            If provided, pipelines are loaded from this path.
+        **kwargs: Keyword arguments for processor configuration, as defined in
+            `ProcessorConfigKwargs`.
+
+    Returns:
+        A tuple containing the input (pre-processor) and output (post-processor) pipelines.
+
+    Raises:
+        NotImplementedError: If a processor factory is not implemented for the given
+            policy configuration type.
+    """
+    if pretrained_path:
+        return (
+            PolicyProcessorPipeline.from_pretrained(
+                pretrained_model_name_or_path=pretrained_path,
+                config_filename=kwargs.get(
+                    "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
+                ),
+                overrides=kwargs.get("preprocessor_overrides", {}),
+                to_transition=batch_to_transition,
+                to_output=transition_to_batch,
+            ),
+            PolicyProcessorPipeline.from_pretrained(
+                pretrained_model_name_or_path=pretrained_path,
+                config_filename=kwargs.get(
+                    "postprocessor_config_filename", f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json"
+                ),
+                overrides=kwargs.get("postprocessor_overrides", {}),
+                to_transition=policy_action_to_transition,
+                to_output=transition_to_policy_action,
+            ),
+        )
+
+    # Create a new processor based on policy type
+    if isinstance(policy_cfg, TDMPCConfig):
+        from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
+
+        processors = make_tdmpc_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, DiffusionConfig):
+        from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors
+
+        processors = make_diffusion_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, ACTConfig):
+        from lerobot.policies.act.processor_act import make_act_pre_post_processors
+
+        processors = make_act_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, VQBeTConfig):
+        from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors
+
+        processors = make_vqbet_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, PI0Config):
+        from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
+
+        processors = make_pi0_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, PI0FASTConfig):
+        from lerobot.policies.pi0fast.processor_pi0fast import make_pi0fast_pre_post_processors
+
+        processors = make_pi0fast_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, SACConfig):
+        from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors
+
+        processors = make_sac_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, RewardClassifierConfig):
+        from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor
+
+        processors = make_classifier_processor(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, SmolVLAConfig):
+        from lerobot.policies.smolvla.processor_smolvla import make_smolvla_pre_post_processors
+
+        processors = make_smolvla_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    else:
+        raise NotImplementedError(f"Processor for policy type '{policy_cfg.type}' is not implemented.")
+
+    return processors
+
+
 def make_policy(
     cfg: PreTrainedConfig,
     ds_meta: LeRobotDatasetMetadata | None = None,
     env_cfg: EnvConfig | None = None,
 ) -> PreTrainedPolicy:
-    """Make an instance of a policy class.
+    """
+    Instantiate a policy model.
 
-    This function exists because (for now) we need to parse features from either a dataset or an environment
-    in order to properly dimension and instantiate a policy for that dataset or environment.
+    This factory function handles the logic of creating a policy, which requires
+    determining the input and output feature shapes. These shapes can be derived
+    either from a `LeRobotDatasetMetadata` object or an `EnvConfig` object. The function
+    can either initialize a new policy from scratch or load a pretrained one.
 
     Args:
-        cfg (PreTrainedConfig): The config of the policy to make. If `pretrained_path` is set, the policy will
-            be loaded with the weights from that path.
-        ds_meta (LeRobotDatasetMetadata | None, optional): Dataset metadata to take input/output shapes and
-            statistics to use for (un)normalization of inputs/outputs in the policy. Defaults to None.
-        env_cfg (EnvConfig | None, optional): The config of a gym environment to parse features from. Must be
-            provided if ds_meta is not. Defaults to None.
-
-    Raises:
-        ValueError: Either ds_meta or env and env_cfg must be provided.
-        NotImplementedError: if the policy.type is 'vqbet' and the policy device 'mps' (due to an incompatibility)
+        cfg: The configuration for the policy to be created. If `cfg.pretrained_path` is
+             set, the policy will be loaded with weights from that path.
+        ds_meta: Dataset metadata used to infer feature shapes and types. Also provides
+                 statistics for normalization layers.
+        env_cfg: Environment configuration used to infer feature shapes and types.
+                 One of `ds_meta` or `env_cfg` must be provided.
 
     Returns:
-        PreTrainedPolicy: _description_
+        An instantiated and device-placed policy model.
+
+    Raises:
+        ValueError: If both or neither of `ds_meta` and `env_cfg` are provided.
+        NotImplementedError: If attempting to use an unsupported policy-backend
+                             combination (e.g., VQBeT with 'mps').
     """
     if bool(ds_meta) == bool(env_cfg):
         raise ValueError("Either one of a dataset metadata or a sim env must be provided.")
@@ -147,7 +349,6 @@ def make_policy(
     kwargs = {}
     if ds_meta is not None:
         features = dataset_to_policy_features(ds_meta.features)
-        kwargs["dataset_stats"] = ds_meta.stats
     else:
         if not cfg.pretrained_path:
             logging.warning(
@@ -155,6 +356,8 @@ def make_policy(
                 "rather than a dataset. Normalization modules inside the policy will have infinite values "
                 "by default without stats from a dataset."
             )
+        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}
@@ -171,7 +374,7 @@ def make_policy(
         policy = policy_cls(**kwargs)
 
     policy.to(cfg.device)
-    assert isinstance(policy, nn.Module)
+    assert isinstance(policy, torch.nn.Module)
 
     # policy = torch.compile(policy, mode="reduce-overhead")
 

src/lerobot/policies/normalize.py

@@ -1,420 +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.
-import numpy as np
-import torch
-from torch import Tensor, nn
-
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
-
-
-def create_stats_buffers(
-    features: dict[str, PolicyFeature],
-    norm_map: dict[str, NormalizationMode],
-    stats: dict[str, dict[str, Tensor]] | None = None,
-) -> dict[str, dict[str, nn.ParameterDict]]:
-    """
-    Create buffers per modality (e.g. "observation.image", "action") containing their mean, std, min, max
-    statistics.
-
-    Args: (see Normalize and Unnormalize)
-
-    Returns:
-        dict: A dictionary where keys are modalities and values are `nn.ParameterDict` containing
-            `nn.Parameters` set to `requires_grad=False`, suitable to not be updated during backpropagation.
-    """
-    stats_buffers = {}
-
-    for key, ft in features.items():
-        norm_mode = norm_map.get(ft.type, NormalizationMode.IDENTITY)
-        if norm_mode is NormalizationMode.IDENTITY:
-            continue
-
-        assert isinstance(norm_mode, NormalizationMode)
-
-        shape = tuple(ft.shape)
-
-        if ft.type is FeatureType.VISUAL:
-            # sanity checks
-            assert len(shape) == 3, f"number of dimensions of {key} != 3 ({shape=}"
-            c, h, w = shape
-            assert c < h and c < w, f"{key} is not channel first ({shape=})"
-            # override image shape to be invariant to height and width
-            shape = (c, 1, 1)
-
-        # Note: we initialize mean, std, min, max to infinity. They should be overwritten
-        # downstream by `stats` or `policy.load_state_dict`, as expected. During forward,
-        # we assert they are not infinity anymore.
-
-        buffer = {}
-        if norm_mode is NormalizationMode.MEAN_STD:
-            mean = torch.ones(shape, dtype=torch.float32) * torch.inf
-            std = torch.ones(shape, dtype=torch.float32) * torch.inf
-            buffer = nn.ParameterDict(
-                {
-                    "mean": nn.Parameter(mean, requires_grad=False),
-                    "std": nn.Parameter(std, requires_grad=False),
-                }
-            )
-        elif norm_mode is NormalizationMode.MIN_MAX:
-            min = torch.ones(shape, dtype=torch.float32) * torch.inf
-            max = torch.ones(shape, dtype=torch.float32) * torch.inf
-            buffer = nn.ParameterDict(
-                {
-                    "min": nn.Parameter(min, requires_grad=False),
-                    "max": nn.Parameter(max, requires_grad=False),
-                }
-            )
-
-        # TODO(aliberts, rcadene): harmonize this to only use one framework (np or torch)
-        if stats:
-            if isinstance(stats[key]["mean"], np.ndarray):
-                if norm_mode is NormalizationMode.MEAN_STD:
-                    buffer["mean"].data = torch.from_numpy(stats[key]["mean"]).to(dtype=torch.float32)
-                    buffer["std"].data = torch.from_numpy(stats[key]["std"]).to(dtype=torch.float32)
-                elif norm_mode is NormalizationMode.MIN_MAX:
-                    buffer["min"].data = torch.from_numpy(stats[key]["min"]).to(dtype=torch.float32)
-                    buffer["max"].data = torch.from_numpy(stats[key]["max"]).to(dtype=torch.float32)
-            elif isinstance(stats[key]["mean"], torch.Tensor):
-                # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
-                # tensors anywhere (for example, when we use the same stats for normalization and
-                # unnormalization). See the logic here
-                # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
-                if norm_mode is NormalizationMode.MEAN_STD:
-                    buffer["mean"].data = stats[key]["mean"].clone().to(dtype=torch.float32)
-                    buffer["std"].data = stats[key]["std"].clone().to(dtype=torch.float32)
-                elif norm_mode is NormalizationMode.MIN_MAX:
-                    buffer["min"].data = stats[key]["min"].clone().to(dtype=torch.float32)
-                    buffer["max"].data = stats[key]["max"].clone().to(dtype=torch.float32)
-            else:
-                type_ = type(stats[key]["mean"])
-                raise ValueError(f"np.ndarray or torch.Tensor expected, but type is '{type_}' instead.")
-
-        stats_buffers[key] = buffer
-    return stats_buffers
-
-
-def _no_stats_error_str(name: str) -> str:
-    return (
-        f"`{name}` is infinity. You should either initialize with `stats` as an argument, or use a "
-        "pretrained model."
-    )
-
-
-class Normalize(nn.Module):
-    """Normalizes data (e.g. "observation.image") for more stable and faster convergence during training."""
-
-    def __init__(
-        self,
-        features: dict[str, PolicyFeature],
-        norm_map: dict[str, NormalizationMode],
-        stats: dict[str, dict[str, Tensor]] | None = None,
-    ):
-        """
-        Args:
-            shapes (dict): A dictionary where keys are input modalities (e.g. "observation.image") and values
-            are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing
-            mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape
-            is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.
-            modes (dict): A dictionary where keys are output modalities (e.g. "observation.image") and values
-                are their normalization modes among:
-                    - "mean_std": subtract the mean and divide by standard deviation.
-                    - "min_max": map to [-1, 1] range.
-            stats (dict, optional): A dictionary where keys are output modalities (e.g. "observation.image")
-                and values are dictionaries of statistic types and their values (e.g.
-                `{"mean": torch.randn(3,1,1)}, "std": torch.randn(3,1,1)}`). If provided, as expected for
-                training the model for the first time, these statistics will overwrite the default buffers. If
-                not provided, as expected for finetuning or evaluation, the default buffers should to be
-                overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the
-                dataset is not needed to get the stats, since they are already in the policy state_dict.
-        """
-        super().__init__()
-        self.features = features
-        self.norm_map = norm_map
-        self.stats = stats
-        stats_buffers = create_stats_buffers(features, norm_map, stats)
-        for key, buffer in stats_buffers.items():
-            setattr(self, "buffer_" + key.replace(".", "_"), buffer)
-
-    # TODO(rcadene): should we remove torch.no_grad?
-    @torch.no_grad()
-    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
-        # TODO: Remove this shallow copy
-        batch = dict(batch)  # shallow copy avoids mutating the input batch
-        for key, ft in self.features.items():
-            if key not in batch:
-                # FIXME(aliberts, rcadene): This might lead to silent fail!
-                continue
-
-            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
-            if norm_mode is NormalizationMode.IDENTITY:
-                continue
-
-            buffer = getattr(self, "buffer_" + key.replace(".", "_"))
-
-            if norm_mode is NormalizationMode.MEAN_STD:
-                mean = buffer["mean"]
-                std = buffer["std"]
-                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
-                assert not torch.isinf(std).any(), _no_stats_error_str("std")
-                batch[key] = (batch[key] - mean) / (std + 1e-8)
-            elif norm_mode is NormalizationMode.MIN_MAX:
-                min = buffer["min"]
-                max = buffer["max"]
-                assert not torch.isinf(min).any(), _no_stats_error_str("min")
-                assert not torch.isinf(max).any(), _no_stats_error_str("max")
-                # normalize to [0,1]
-                batch[key] = (batch[key] - min) / (max - min + 1e-8)
-                # normalize to [-1, 1]
-                batch[key] = batch[key] * 2 - 1
-            else:
-                raise ValueError(norm_mode)
-        return batch
-
-
-class Unnormalize(nn.Module):
-    """
-    Similar to `Normalize` but unnormalizes output data (e.g. `{"action": torch.randn(b,c)}`) in their
-    original range used by the environment.
-    """
-
-    def __init__(
-        self,
-        features: dict[str, PolicyFeature],
-        norm_map: dict[str, NormalizationMode],
-        stats: dict[str, dict[str, Tensor]] | None = None,
-    ):
-        """
-        Args:
-            shapes (dict): A dictionary where keys are input modalities (e.g. "observation.image") and values
-            are their shapes (e.g. `[3,96,96]`]). These shapes are used to create the tensor buffer containing
-            mean, std, min, max statistics. If the provided `shapes` contain keys related to images, the shape
-            is adjusted to be invariant to height and width, assuming a channel-first (c, h, w) format.
-            modes (dict): A dictionary where keys are output modalities (e.g. "observation.image") and values
-                are their normalization modes among:
-                    - "mean_std": subtract the mean and divide by standard deviation.
-                    - "min_max": map to [-1, 1] range.
-            stats (dict, optional): A dictionary where keys are output modalities (e.g. "observation.image")
-                and values are dictionaries of statistic types and their values (e.g.
-                `{"mean": torch.randn(3,1,1)}, "std": torch.randn(3,1,1)}`). If provided, as expected for
-                training the model for the first time, these statistics will overwrite the default buffers. If
-                not provided, as expected for finetuning or evaluation, the default buffers should to be
-                overwritten by a call to `policy.load_state_dict(state_dict)`. That way, initializing the
-                dataset is not needed to get the stats, since they are already in the policy state_dict.
-        """
-        super().__init__()
-        self.features = features
-        self.norm_map = norm_map
-        self.stats = stats
-        # `self.buffer_observation_state["mean"]` contains `torch.tensor(state_dim)`
-        stats_buffers = create_stats_buffers(features, norm_map, stats)
-        for key, buffer in stats_buffers.items():
-            setattr(self, "buffer_" + key.replace(".", "_"), buffer)
-
-    # TODO(rcadene): should we remove torch.no_grad?
-    @torch.no_grad()
-    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
-        batch = dict(batch)  # shallow copy avoids mutating the input batch
-        for key, ft in self.features.items():
-            if key not in batch:
-                continue
-
-            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
-            if norm_mode is NormalizationMode.IDENTITY:
-                continue
-
-            buffer = getattr(self, "buffer_" + key.replace(".", "_"))
-
-            if norm_mode is NormalizationMode.MEAN_STD:
-                mean = buffer["mean"]
-                std = buffer["std"]
-                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
-                assert not torch.isinf(std).any(), _no_stats_error_str("std")
-                batch[key] = batch[key] * std + mean
-            elif norm_mode is NormalizationMode.MIN_MAX:
-                min = buffer["min"]
-                max = buffer["max"]
-                assert not torch.isinf(min).any(), _no_stats_error_str("min")
-                assert not torch.isinf(max).any(), _no_stats_error_str("max")
-                batch[key] = (batch[key] + 1) / 2
-                batch[key] = batch[key] * (max - min) + min
-            else:
-                raise ValueError(norm_mode)
-        return batch
-
-
-# TODO (azouitine): We should replace all normalization on the policies with register_buffer normalization
-#       and remove the `Normalize` and `Unnormalize` classes.
-def _initialize_stats_buffers(
-    module: nn.Module,
-    features: dict[str, PolicyFeature],
-    norm_map: dict[str, NormalizationMode],
-    stats: dict[str, dict[str, Tensor]] | None = None,
-) -> None:
-    """Register statistics buffers (mean/std or min/max) on the given *module*.
-
-    The logic matches the previous constructors of `NormalizeBuffer` and `UnnormalizeBuffer`,
-    but is factored out so it can be reused by both classes and stay in sync.
-    """
-    for key, ft in features.items():
-        norm_mode = norm_map.get(ft.type, NormalizationMode.IDENTITY)
-        if norm_mode is NormalizationMode.IDENTITY:
-            continue
-
-        shape: tuple[int, ...] = tuple(ft.shape)
-        if ft.type is FeatureType.VISUAL:
-            # reduce spatial dimensions, keep channel dimension only
-            c, *_ = shape
-            shape = (c, 1, 1)
-
-        prefix = key.replace(".", "_")
-
-        if norm_mode is NormalizationMode.MEAN_STD:
-            mean = torch.full(shape, torch.inf, dtype=torch.float32)
-            std = torch.full(shape, torch.inf, dtype=torch.float32)
-
-            if stats and key in stats and "mean" in stats[key] and "std" in stats[key]:
-                mean_data = stats[key]["mean"]
-                std_data = stats[key]["std"]
-                if isinstance(mean_data, torch.Tensor):
-                    # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
-                    # tensors anywhere (for example, when we use the same stats for normalization and
-                    # unnormalization). See the logic here
-                    # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
-                    mean = mean_data.clone().to(dtype=torch.float32)
-                    std = std_data.clone().to(dtype=torch.float32)
-                else:
-                    raise ValueError(f"Unsupported stats type for key '{key}' (expected ndarray or Tensor).")
-
-            module.register_buffer(f"{prefix}_mean", mean)
-            module.register_buffer(f"{prefix}_std", std)
-            continue
-
-        if norm_mode is NormalizationMode.MIN_MAX:
-            min_val = torch.full(shape, torch.inf, dtype=torch.float32)
-            max_val = torch.full(shape, torch.inf, dtype=torch.float32)
-
-            if stats and key in stats and "min" in stats[key] and "max" in stats[key]:
-                min_data = stats[key]["min"]
-                max_data = stats[key]["max"]
-                if isinstance(min_data, torch.Tensor):
-                    min_val = min_data.clone().to(dtype=torch.float32)
-                    max_val = max_data.clone().to(dtype=torch.float32)
-                else:
-                    raise ValueError(f"Unsupported stats type for key '{key}' (expected ndarray or Tensor).")
-
-            module.register_buffer(f"{prefix}_min", min_val)
-            module.register_buffer(f"{prefix}_max", max_val)
-            continue
-
-        raise ValueError(norm_mode)
-
-
-class NormalizeBuffer(nn.Module):
-    """Same as `Normalize` but statistics are stored as registered buffers rather than parameters."""
-
-    def __init__(
-        self,
-        features: dict[str, PolicyFeature],
-        norm_map: dict[str, NormalizationMode],
-        stats: dict[str, dict[str, Tensor]] | None = None,
-    ):
-        super().__init__()
-        self.features = features
-        self.norm_map = norm_map
-
-        _initialize_stats_buffers(self, features, norm_map, stats)
-
-    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
-        batch = dict(batch)
-        for key, ft in self.features.items():
-            if key not in batch:
-                continue
-
-            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
-            if norm_mode is NormalizationMode.IDENTITY:
-                continue
-
-            prefix = key.replace(".", "_")
-
-            if norm_mode is NormalizationMode.MEAN_STD:
-                mean = getattr(self, f"{prefix}_mean")
-                std = getattr(self, f"{prefix}_std")
-                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
-                assert not torch.isinf(std).any(), _no_stats_error_str("std")
-                batch[key] = (batch[key] - mean) / (std + 1e-8)
-                continue
-
-            if norm_mode is NormalizationMode.MIN_MAX:
-                min_val = getattr(self, f"{prefix}_min")
-                max_val = getattr(self, f"{prefix}_max")
-                assert not torch.isinf(min_val).any(), _no_stats_error_str("min")
-                assert not torch.isinf(max_val).any(), _no_stats_error_str("max")
-                batch[key] = (batch[key] - min_val) / (max_val - min_val + 1e-8)
-                batch[key] = batch[key] * 2 - 1
-                continue
-
-            raise ValueError(norm_mode)
-
-        return batch
-
-
-class UnnormalizeBuffer(nn.Module):
-    """Inverse operation of `NormalizeBuffer`. Uses registered buffers for statistics."""
-
-    def __init__(
-        self,
-        features: dict[str, PolicyFeature],
-        norm_map: dict[str, NormalizationMode],
-        stats: dict[str, dict[str, Tensor]] | None = None,
-    ):
-        super().__init__()
-        self.features = features
-        self.norm_map = norm_map
-
-        _initialize_stats_buffers(self, features, norm_map, stats)
-
-    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
-        # batch = dict(batch)
-        for key, ft in self.features.items():
-            if key not in batch:
-                continue
-
-            norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
-            if norm_mode is NormalizationMode.IDENTITY:
-                continue
-
-            prefix = key.replace(".", "_")
-
-            if norm_mode is NormalizationMode.MEAN_STD:
-                mean = getattr(self, f"{prefix}_mean")
-                std = getattr(self, f"{prefix}_std")
-                assert not torch.isinf(mean).any(), _no_stats_error_str("mean")
-                assert not torch.isinf(std).any(), _no_stats_error_str("std")
-                batch[key] = batch[key] * std + mean
-                continue
-
-            if norm_mode is NormalizationMode.MIN_MAX:
-                min_val = getattr(self, f"{prefix}_min")
-                max_val = getattr(self, f"{prefix}_max")
-                assert not torch.isinf(min_val).any(), _no_stats_error_str("min")
-                assert not torch.isinf(max_val).any(), _no_stats_error_str("max")
-                batch[key] = (batch[key] + 1) / 2
-                batch[key] = batch[key] * (max_val - min_val) + min_val
-                continue
-
-            raise ValueError(norm_mode)
-
-        return batch

src/lerobot/policies/pi0/modeling_pi0.py

@@ -56,18 +56,15 @@ from collections import deque
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
-from transformers import AutoTokenizer
 
-from lerobot.constants import ACTION, OBS_STATE
-from lerobot.policies.normalize import Normalize, Unnormalize
+from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.policies.pi0.paligemma_with_expert import (
     PaliGemmaWithExpertConfig,
     PaliGemmaWithExpertModel,
 )
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.utils import log_model_loading_keys
-from lerobot.utils.utils import get_safe_dtype, init_logging
+from lerobot.utils.utils import get_safe_dtype
 
 
 def create_sinusoidal_pos_embedding(
@@ -223,28 +220,17 @@ class PI0Policy(PreTrainedPolicy):
     def __init__(
         self,
         config: PI0Config,
-        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                     the configuration class is used.
-            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
-                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
 
         super().__init__(config)
         config.validate_features()
         self.config = config
-        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
-        self.normalize_targets = Normalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-        self.unnormalize_outputs = Unnormalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
 
-        self.language_tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
         self.model = PI0FlowMatching(config)
 
         self.reset()
@@ -253,99 +239,6 @@ class PI0Policy(PreTrainedPolicy):
         """This should be called whenever the environment is reset."""
         self._action_queue = deque([], maxlen=self.config.n_action_steps)
 
-    @classmethod
-    def _transform_state_dict_keys(cls, state_dict: dict) -> dict:
-        """
-        Transform state dict keys to match expected model structure.
-
-        Transformations:
-        - model.paligemma_with_expert.paligemma.language_model.lm_head ->
-          model.paligemma_with_expert.paligemma.lm_head
-        - model.paligemma_with_expert.paligemma.language_model.model ->
-          model.paligemma_with_expert.paligemma.model.language_model
-        - model.paligemma_with_expert.paligemma.vision_tower ->
-          model.paligemma_with_expert.paligemma.model.vision_tower
-        - model.paligemma_with_expert.paligemma.multi_modal_projector ->
-          model.paligemma_with_expert.paligemma.model.multi_modal_projector
-
-        Also handles tied weights between lm_head.weight and
-        embed_tokens.weight.
-        """
-        import re
-
-        transformed_dict = {}
-
-        transformations = [
-            (
-                re.compile(r"\.paligemma_with_expert\.paligemma\.language_model\.lm_head"),
-                ".paligemma_with_expert.paligemma.lm_head",
-            ),
-            (
-                re.compile(r"\.paligemma_with_expert\.paligemma\.language_model\.model"),
-                ".paligemma_with_expert.paligemma.model.language_model",
-            ),
-            (
-                re.compile(r"\.paligemma_with_expert\.paligemma\.vision_tower"),
-                ".paligemma_with_expert.paligemma.model.vision_tower",
-            ),
-            (
-                re.compile(r"\.paligemma_with_expert\.paligemma\.multi_modal_projector"),
-                ".paligemma_with_expert.paligemma.model.multi_modal_projector",
-            ),
-        ]
-
-        for key, value in state_dict.items():
-            new_key = key
-            for pattern, replacement in transformations:
-                new_key = pattern.sub(replacement, new_key)
-            transformed_dict[new_key] = value
-
-        # Handle tied weights: lm_head.weight and embed_tokens.weight share memory
-        lm_head_key = None
-        embed_tokens_key = None
-
-        for key in transformed_dict:
-            if key.endswith(".paligemma_with_expert.paligemma.lm_head.weight"):
-                lm_head_key = key
-            elif key.endswith(".paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"):
-                embed_tokens_key = key
-            if lm_head_key and embed_tokens_key:
-                break
-
-        if lm_head_key and not embed_tokens_key:
-            embed_tokens_key = lm_head_key.replace(
-                ".lm_head.weight", ".model.language_model.embed_tokens.weight"
-            )
-            transformed_dict[embed_tokens_key] = transformed_dict[lm_head_key]
-        elif embed_tokens_key and not lm_head_key:
-            lm_head_key = embed_tokens_key.replace(
-                ".model.language_model.embed_tokens.weight", ".lm_head.weight"
-            )
-            transformed_dict[lm_head_key] = transformed_dict[embed_tokens_key]
-
-        return transformed_dict
-
-    @classmethod
-    def _load_as_safetensor(
-        cls, model: "PI0Policy", model_file: str, map_location: str, strict: bool
-    ) -> "PI0Policy":
-        """Override to apply key transformations before loading."""
-        from safetensors.torch import load_file
-
-        init_logging()
-        # Load the state dict from file safely
-        state_dict = load_file(model_file, device=map_location)
-
-        # Apply key transformations
-        transformed_state_dict = cls._transform_state_dict_keys(state_dict)
-
-        # Load the transformed state dict
-        msg = model.load_state_dict(transformed_state_dict, strict=strict)
-
-        # Log message
-        log_model_loading_keys(msg.missing_keys, msg.unexpected_keys)
-        return model
-
     def get_optim_params(self) -> dict:
         return self.parameters()
 
@@ -377,14 +270,13 @@ class PI0Policy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
-        batch = self.normalize_inputs(batch)
-
         # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
         # querying the policy.
         if len(self._action_queue) == 0:
             images, img_masks = self.prepare_images(batch)
             state = self.prepare_state(batch)
-            lang_tokens, lang_masks = self.prepare_language(batch)
+            lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
+            lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
 
             actions = self.model.sample_actions(
                 images, img_masks, lang_tokens, lang_masks, state, noise=noise
@@ -394,8 +286,6 @@ class PI0Policy(PreTrainedPolicy):
             original_action_dim = self.config.action_feature.shape[0]
             actions = actions[:, :, :original_action_dim]
 
-            actions = self.unnormalize_outputs({"action": actions})["action"]
-
             if self.config.adapt_to_pi_aloha:
                 actions = self._pi_aloha_encode_actions(actions)
 
@@ -410,12 +300,10 @@ class PI0Policy(PreTrainedPolicy):
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
             batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
 
-        batch = self.normalize_inputs(batch)
-        batch = self.normalize_targets(batch)
-
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
-        lang_tokens, lang_masks = self.prepare_language(batch)
+        lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
+        lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
         actions = self.prepare_action(batch)
         actions_is_pad = batch.get("action_is_pad")
 
@@ -482,26 +370,6 @@ class PI0Policy(PreTrainedPolicy):
 
         return images, img_masks
 
-    def prepare_language(self, batch) -> tuple[Tensor, Tensor]:
-        """Tokenize the text input"""
-        device = batch[OBS_STATE].device
-        tasks = batch["task"]
-
-        # PaliGemma prompt has to end with a new line
-        tasks = [task if task.endswith("\n") else f"{task}\n" for task in tasks]
-
-        tokenized_prompt = self.language_tokenizer.__call__(
-            tasks,
-            padding="max_length",
-            padding_side="right",
-            max_length=self.config.tokenizer_max_length,
-            return_tensors="pt",
-        )
-        lang_tokens = tokenized_prompt["input_ids"].to(device=device)
-        lang_masks = tokenized_prompt["attention_mask"].to(device=device, dtype=torch.bool)
-
-        return lang_tokens, lang_masks
-
     def _pi_aloha_decode_state(self, state):
         # Flip the joints.
         for motor_idx in [1, 2, 8, 9]:
@@ -567,7 +435,7 @@ class PI0FlowMatching(nn.Module):
     └──────────────────────────────┘
     """
 
-    def __init__(self, config):
+    def __init__(self, config: PI0Config):
         super().__init__()
         self.config = config
 

src/lerobot/policies/pi0/processor_pi0.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python
+
+# Copyright 2025 Physical Intelligence and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+from lerobot.policies.pi0.configuration_pi0 import PI0Config
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    ComplementaryDataProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    RenameObservationsProcessorStep,
+    TokenizerProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+
+@ProcessorStepRegistry.register(name="pi0_new_line_processor")
+class Pi0NewLineProcessor(ComplementaryDataProcessorStep):
+    """
+    Ensures that the task description string ends with a newline character.
+
+    This processing step is required for compatibility with the PaliGemma tokenizer,
+    which expects a newline at the end of the text prompt. It handles both single
+    strings and lists of strings for the 'task' key in complementary data.
+    """
+
+    def complementary_data(self, complementary_data):
+        """
+        Adds a newline to the 'task' field if it doesn't already have one.
+
+        Args:
+            complementary_data: A dictionary that may contain a 'task' key with a
+                                string or list of strings.
+
+        Returns:
+            A new dictionary with the modified 'task' field.
+        """
+        if "task" not in complementary_data:
+            return complementary_data
+
+        task = complementary_data["task"]
+        if task is None:
+            return complementary_data
+
+        new_complementary_data = dict(complementary_data)
+
+        # Handle both string and list of strings
+        if isinstance(task, str):
+            # Single string: add newline if not present
+            if not task.endswith("\n"):
+                new_complementary_data["task"] = f"{task}\n"
+        elif isinstance(task, list) and all(isinstance(t, str) for t in task):
+            # List of strings: add newline to each if not present
+            new_complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task]
+        # If task is neither string nor list of strings, leave unchanged
+
+        return new_complementary_data
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        This step does not alter the feature definitions.
+
+        Args:
+            features: The input feature dictionary.
+
+        Returns:
+            The unchanged feature dictionary.
+        """
+        return features
+
+
+def make_pi0_pre_post_processors(
+    config: PI0Config,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Constructs pre-processor and post-processor pipelines for the PI0 policy.
+
+    The pre-processing pipeline prepares input data for the model by:
+    1. Renaming features to match pretrained configurations.
+    2. Normalizing input and output features based on dataset statistics.
+    3. Adding a batch dimension.
+    4. Appending a newline character to the task description for tokenizer compatibility.
+    5. Tokenizing the text prompt using the PaliGemma tokenizer.
+    6. Moving all data to the specified device.
+
+    The post-processing pipeline handles the model's output by:
+    1. Moving data to the CPU.
+    2. Unnormalizing the output features to their original scale.
+
+    Args:
+        config: The configuration object for the PI0 policy.
+        dataset_stats: A dictionary of statistics for normalization.
+        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
+        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
+
+    Returns:
+        A tuple containing the configured pre-processor and post-processor pipelines.
+    """
+
+    # Add remaining processors
+    input_steps: list[ProcessorStep] = [
+        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
+        AddBatchDimensionProcessorStep(),
+        Pi0NewLineProcessor(),  # Add newlines before tokenization for PaliGemma
+        TokenizerProcessorStep(
+            tokenizer_name="google/paligemma-3b-pt-224",
+            max_length=config.tokenizer_max_length,
+            padding_side="right",
+            padding="max_length",
+        ),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+
+    output_steps: list[ProcessorStep] = [
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/pi0fast/modeling_pi0fast.py

@@ -58,7 +58,6 @@ from transformers.cache_utils import HybridCache, StaticCache
 from transformers.models.auto import CONFIG_MAPPING
 
 from lerobot.constants import ACTION, OBS_STATE
-from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
 
@@ -146,14 +145,6 @@ class PI0FASTPolicy(PreTrainedPolicy):
         config.validate_features()
         self.config = config
 
-        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
-        self.normalize_targets = Normalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-        self.unnormalize_outputs = Unnormalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-
         self.language_tokenizer = AutoProcessor.from_pretrained("google/paligemma-3b-pt-224")
         self.model = PI0FAST(config)
 
@@ -221,8 +212,6 @@ class PI0FASTPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
-        batch = self.normalize_inputs(batch)
-
         # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
         # querying the policy.
         if len(self._action_queue) == 0:
@@ -235,8 +224,6 @@ class PI0FASTPolicy(PreTrainedPolicy):
             ]  # self.config.max_action_dim  # self.config.action_feature.shape[0]
             actions = actions[:, :, :original_action_dim]
 
-            actions = self.unnormalize_outputs({"action": actions})["action"]
-
             if self.config.adapt_to_pi_aloha:
                 actions = self._pi_aloha_encode_actions(actions)
 
@@ -249,8 +236,6 @@ class PI0FASTPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
             batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
-        batch = self.normalize_inputs(batch)
-        batch = self.normalize_targets(batch)
         loss_dict = self.model.forward(batch)
         return loss_dict["loss"], loss_dict
 

src/lerobot/policies/pi0fast/processor_pi0fast.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python
+
+# Copyright 2025 Physical Intelligence and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+
+def make_pi0fast_pre_post_processors(
+    config: PI0FASTConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Constructs pre-processor and post-processor pipelines for the PI0Fast policy.
+
+    The pre-processing pipeline prepares input data for the model by:
+    1. Renaming features to match pretrained configurations.
+    2. Normalizing input and output features based on dataset statistics.
+    3. Adding a batch dimension.
+    4. Moving all data to the specified device.
+
+    The post-processing pipeline handles the model's output by:
+    1. Moving data to the CPU.
+    2. Unnormalizing the output features to their original scale.
+
+    Args:
+        config: The configuration object for the PI0Fast policy.
+        dataset_stats: A dictionary of statistics for normalization.
+        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
+        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
+
+    Returns:
+        A tuple containing the configured pre-processor and post-processor pipelines.
+    """
+
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
+        AddBatchDimensionProcessorStep(),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/sac/modeling_sac.py

@@ -28,7 +28,6 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 from torch.distributions import MultivariateNormal, TanhTransform, Transform, TransformedDistribution
 
-from lerobot.policies.normalize import NormalizeBuffer
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.configuration_sac import SACConfig, is_image_feature
 from lerobot.policies.utils import get_device_from_parameters
@@ -45,7 +44,6 @@ class SACPolicy(
     def __init__(
         self,
         config: SACConfig | None = None,
-        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         super().__init__(config)
         config.validate_features()
@@ -53,7 +51,6 @@ class SACPolicy(
 
         # Determine action dimension and initialize all components
         continuous_action_dim = config.output_features["action"].shape[0]
-        self._init_normalization(dataset_stats)
         self._init_encoders()
         self._init_critics(continuous_action_dim)
         self._init_actor(continuous_action_dim)
@@ -88,8 +85,7 @@ class SACPolicy(
 
         observations_features = None
         if self.shared_encoder and self.actor.encoder.has_images:
-            # Cache and normalize image features
-            observations_features = self.actor.encoder.get_cached_image_features(batch, normalize=True)
+            observations_features = self.actor.encoder.get_cached_image_features(batch)
 
         actions, _, _ = self.actor(batch, observations_features)
 
@@ -391,28 +387,12 @@ class SACPolicy(
         actor_loss = ((self.temperature * log_probs) - min_q_preds).mean()
         return actor_loss
 
-    def _init_normalization(self, dataset_stats):
-        """Initialize input/output normalization modules."""
-        self.normalize_inputs = nn.Identity()
-        self.normalize_targets = nn.Identity()
-        if self.config.dataset_stats is not None:
-            params = _convert_normalization_params_to_tensor(self.config.dataset_stats)
-            self.normalize_inputs = NormalizeBuffer(
-                self.config.input_features, self.config.normalization_mapping, params
-            )
-            stats = dataset_stats or params
-            self.normalize_targets = NormalizeBuffer(
-                self.config.output_features, self.config.normalization_mapping, stats
-            )
-
     def _init_encoders(self):
         """Initialize shared or separate encoders for actor and critic."""
         self.shared_encoder = self.config.shared_encoder
-        self.encoder_critic = SACObservationEncoder(self.config, self.normalize_inputs)
+        self.encoder_critic = SACObservationEncoder(self.config)
         self.encoder_actor = (
-            self.encoder_critic
-            if self.shared_encoder
-            else SACObservationEncoder(self.config, self.normalize_inputs)
+            self.encoder_critic if self.shared_encoder else SACObservationEncoder(self.config)
         )
 
     def _init_critics(self, continuous_action_dim):
@@ -424,9 +404,7 @@ class SACPolicy(
             )
             for _ in range(self.config.num_critics)
         ]
-        self.critic_ensemble = CriticEnsemble(
-            encoder=self.encoder_critic, ensemble=heads, output_normalization=self.normalize_targets
-        )
+        self.critic_ensemble = CriticEnsemble(encoder=self.encoder_critic, ensemble=heads)
         target_heads = [
             CriticHead(
                 input_dim=self.encoder_critic.output_dim + continuous_action_dim,
@@ -434,9 +412,7 @@ class SACPolicy(
             )
             for _ in range(self.config.num_critics)
         ]
-        self.critic_target = CriticEnsemble(
-            encoder=self.encoder_critic, ensemble=target_heads, output_normalization=self.normalize_targets
-        )
+        self.critic_target = CriticEnsemble(encoder=self.encoder_critic, ensemble=target_heads)
         self.critic_target.load_state_dict(self.critic_ensemble.state_dict())
 
         if self.config.use_torch_compile:
@@ -490,10 +466,9 @@ class SACPolicy(
 class SACObservationEncoder(nn.Module):
     """Encode image and/or state vector observations."""
 
-    def __init__(self, config: SACConfig, input_normalizer: nn.Module) -> None:
+    def __init__(self, config: SACConfig) -> None:
         super().__init__()
         self.config = config
-        self.input_normalization = input_normalizer
         self._init_image_layers()
         self._init_state_layers()
         self._compute_output_dim()
@@ -568,11 +543,10 @@ class SACObservationEncoder(nn.Module):
     def forward(
         self, obs: dict[str, Tensor], cache: dict[str, Tensor] | None = None, detach: bool = False
     ) -> Tensor:
-        obs = self.input_normalization(obs)
         parts = []
         if self.has_images:
             if cache is None:
-                cache = self.get_cached_image_features(obs, normalize=False)
+                cache = self.get_cached_image_features(obs)
             parts.append(self._encode_images(cache, detach))
         if self.has_env:
             parts.append(self.env_encoder(obs["observation.environment_state"]))
@@ -585,7 +559,7 @@ class SACObservationEncoder(nn.Module):
             "No parts to concatenate, you should have at least one image or environment state or state"
         )
 
-    def get_cached_image_features(self, obs: dict[str, Tensor], normalize: bool = False) -> dict[str, Tensor]:
+    def get_cached_image_features(self, obs: dict[str, Tensor]) -> dict[str, Tensor]:
         """Extract and optionally cache image features from observations.
 
         This function processes image observations through the vision encoder once and returns
@@ -597,26 +571,17 @@ class SACObservationEncoder(nn.Module):
         - The vision encoder forward pass is typically the main computational bottleneck during training and inference
         - Caching these features can provide 2-4x speedup in training and inference
 
-        Normalization behavior:
-        - When called from inside forward(): set normalize=False since inputs are already normalized
-        - When called from outside forward(): set normalize=True to ensure proper input normalization
-
         Usage patterns:
-        - Called in select_action() with normalize=True
+        - Called in select_action()
         - Called in learner.py's get_observation_features() to pre-compute features for all policy components
-        - Called internally by forward() with normalize=False
+        - Called internally by forward()
 
         Args:
             obs: Dictionary of observation tensors containing image keys
-            normalize: Whether to normalize observations before encoding
-                      Set to True when calling directly from outside the encoder's forward method
-                      Set to False when calling from within forward() where inputs are already normalized
 
         Returns:
             Dictionary mapping image keys to their corresponding encoded features
         """
-        if normalize:
-            obs = self.input_normalization(obs)
         batched = torch.cat([obs[k] for k in self.image_keys], dim=0)
         out = self.image_encoder(batched)
         chunks = torch.chunk(out, len(self.image_keys), dim=0)
@@ -747,7 +712,6 @@ class CriticEnsemble(nn.Module):
     Args:
         encoder (SACObservationEncoder): encoder for observations.
         ensemble (List[CriticHead]): list of critic heads.
-        output_normalization (nn.Module): normalization layer for actions.
         init_final (float | None): optional initializer scale for final layers.
 
     Forward returns a tensor of shape (num_critics, batch_size) containing Q-values.
@@ -757,13 +721,11 @@ class CriticEnsemble(nn.Module):
         self,
         encoder: SACObservationEncoder,
         ensemble: list[CriticHead],
-        output_normalization: nn.Module,
         init_final: float | None = None,
     ):
         super().__init__()
         self.encoder = encoder
         self.init_final = init_final
-        self.output_normalization = output_normalization
         self.critics = nn.ModuleList(ensemble)
 
     def forward(
@@ -775,11 +737,6 @@ class CriticEnsemble(nn.Module):
         device = get_device_from_parameters(self)
         # Move each tensor in observations to device
         observations = {k: v.to(device) for k, v in observations.items()}
-        # NOTE: We normalize actions it helps for sample efficiency
-        actions: dict[str, torch.tensor] = {"action": actions}
-        # NOTE: Normalization layer took dict in input and outputs a dict that why
-        actions = self.output_normalization(actions)["action"]
-        actions = actions.to(device)
 
         obs_enc = self.encoder(observations, cache=observation_features)
 

src/lerobot/policies/sac/processor_sac.py

@@ -0,0 +1,92 @@
+#!/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.
+
+from typing import Any
+
+import torch
+
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+from lerobot.policies.sac.configuration_sac import SACConfig
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+
+def make_sac_pre_post_processors(
+    config: SACConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Constructs pre-processor and post-processor pipelines for the SAC policy.
+
+    The pre-processing pipeline prepares input data for the model by:
+    1. Renaming features to match pretrained configurations.
+    2. Normalizing input and output features based on dataset statistics.
+    3. Adding a batch dimension.
+    4. Moving all data to the specified device.
+
+    The post-processing pipeline handles the model's output by:
+    1. Moving data to the CPU.
+    2. Unnormalizing the output features to their original scale.
+
+    Args:
+        config: The configuration object for the SAC policy.
+        dataset_stats: A dictionary of statistics for normalization.
+
+    Returns:
+        A tuple containing the configured pre-processor and post-processor pipelines.
+    """
+
+    # Add remaining processors
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/sac/reward_model/modeling_classifier.py

@@ -20,7 +20,6 @@ import torch
 from torch import Tensor, nn
 
 from lerobot.constants import OBS_IMAGE, REWARD
-from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 
@@ -108,22 +107,12 @@ class Classifier(PreTrainedPolicy):
     def __init__(
         self,
         config: RewardClassifierConfig,
-        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         from transformers import AutoModel
 
         super().__init__(config)
         self.config = config
 
-        # Initialize normalization (standardized with the policy framework)
-        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
-        self.normalize_targets = Normalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-        self.unnormalize_outputs = Unnormalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-
         # Set up encoder
         encoder = AutoModel.from_pretrained(self.config.model_name, trust_remote_code=True)
         # Extract vision model if we're given a multimodal model
@@ -247,10 +236,6 @@ class Classifier(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Tensor]]:
         """Standard forward pass for training compatible with train.py."""
-        # Normalize inputs if needed
-        batch = self.normalize_inputs(batch)
-        batch = self.normalize_targets(batch)
-
         # Extract images and labels
         images, labels = self.extract_images_and_labels(batch)
 

src/lerobot/policies/sac/reward_model/processor_classifier.py

@@ -0,0 +1,82 @@
+# !/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
+from lerobot.processor import (
+    DeviceProcessorStep,
+    IdentityProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+
+def make_classifier_processor(
+    config: RewardClassifierConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Constructs pre-processor and post-processor pipelines for the reward classifier.
+
+    The pre-processing pipeline prepares input data for the classifier by:
+    1. Normalizing both input and output features based on dataset statistics.
+    2. Moving the data to the specified device.
+
+    The post-processing pipeline handles the classifier's output by:
+    1. Moving the data to the CPU.
+    2. Applying an identity step, as no unnormalization is needed for the output logits.
+
+    Args:
+        config: The configuration object for the RewardClassifier.
+        dataset_stats: A dictionary of statistics for normalization.
+        preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
+        postprocessor_kwargs: Additional arguments for the post-processor pipeline.
+
+    Returns:
+        A tuple containing the configured pre-processor and post-processor pipelines.
+    """
+
+    input_steps = [
+        NormalizerProcessorStep(
+            features=config.input_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        NormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device=config.device),
+    ]
+    output_steps = [DeviceProcessorStep(device="cpu"), IdentityProcessorStep()]
+
+    return (
+        PolicyProcessorPipeline(
+            steps=input_steps,
+            name="classifier_preprocessor",
+        ),
+        PolicyProcessorPipeline(
+            steps=output_steps,
+            name="classifier_postprocessor",
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -53,21 +53,13 @@ policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
 """
 
 import math
-import os
-import re
 from collections import deque
 
-import safetensors
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
-from transformers import AutoProcessor
 
-from lerobot.constants import ACTION, OBS_STATE
-from lerobot.policies.normalize import (
-    Normalize,
-    Unnormalize,
-)
+from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
@@ -76,102 +68,6 @@ from lerobot.policies.utils import (
 )
 from lerobot.utils.utils import get_safe_dtype
 
-# Matches ".soNNN", optionally followed by "-something", up to the "_buffer_" marker
-_VARIANT_RE = re.compile(r"\.so\d+(?:-[\w]+)?_buffer_")
-
-
-def canonicalise(k: str) -> str:
-    """
-    Remove dataset-variant markers like '.so100-blue_' or '.so100_' from a
-    normalisation-buffer key.
-    """
-    return _VARIANT_RE.sub(".buffer_", k)
-
-
-def standardise_state_dict(
-    checkpoint: dict[str, torch.Tensor], ref_keys: set[str], *, verbose: bool = True
-) -> tuple[dict[str, torch.Tensor], list[str]]:
-    """
-    • Re-keys `checkpoint ` so that every entry matches the *reference* key set.
-    • If several variant keys collapse to the same canonical name we keep the
-      first one and log the collision.
-    • Returns the new dict + a list of entries that could not be matched.
-    """
-    out, collisions, unmatched = {}, {}, []
-
-    for k, v in checkpoint.items():
-        canon = canonicalise(k)
-        if canon in ref_keys:
-            if canon in out:  # duplicate after collapsing
-                collisions.setdefault(canon, []).append(k)
-            else:
-                out[canon] = v
-        else:
-            unmatched.append(k)
-
-    if verbose:
-        for canon, variants in collisions.items():
-            print(f"[standardise_state_dict] '{canon}'  ←  {variants}")
-        if unmatched:
-            print(f"[standardise_state_dict] kept {len(unmatched)} unmatched keys")
-
-    out.update({k: checkpoint[k] for k in unmatched})
-    return out, unmatched
-
-
-def rename_checkpoint_keys(checkpoint: dict, rename_str: str):
-    """
-    Renames keys in a checkpoint dictionary based on the given rename string.
-
-    Args:
-        checkpoint (dict): The checkpoint dictionary.
-        rename_str (str): A string specifying key mappings in the format "old1//new1,old2//new2".
-
-    Returns:
-        dict: The modified checkpoint with renamed keys.
-    """
-
-    rename_dict = dict(pair.split("//") for pair in rename_str.split(","))
-
-    new_checkpoint = {}
-    for k, v in checkpoint.items():
-        for old_key, new_key in rename_dict.items():
-            if old_key in k:
-                k = k.replace(old_key, new_key)
-        new_checkpoint[k] = v
-    return new_checkpoint
-
-
-def load_smolvla(
-    model: torch.nn.Module,
-    filename: str | os.PathLike,
-    *,
-    device: str = "cpu",
-    checkpoint_keys_mapping: str = "",
-) -> torch.nn.Module:
-    state_dict = safetensors.torch.load_file(filename, device=device)
-
-    # Optional user-supplied renames (e.g. "model._orig_mod.//model.")
-    if checkpoint_keys_mapping and "//" in checkpoint_keys_mapping:
-        state_dict = rename_checkpoint_keys(state_dict, checkpoint_keys_mapping)
-
-    state_dict, _ = standardise_state_dict(state_dict, set(model.state_dict().keys()))
-
-    # HACK(aliberts): to not overwrite normalization parameters as they should come from the dataset
-    norm_keys = ("normalize_inputs", "normalize_targets", "unnormalize_outputs")
-    state_dict = {k: v for k, v in state_dict.items() if not k.startswith(norm_keys)}
-
-    missing, unexpected = model.load_state_dict(state_dict, strict=False)
-
-    if not all(key.startswith(norm_keys) for key in missing) or unexpected:
-        raise RuntimeError(
-            "SmolVLA %d missing / %d unexpected keys",
-            len(missing),
-            len(unexpected),
-        )
-
-    return model
-
 
 def create_sinusoidal_pos_embedding(
     time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu"
@@ -326,28 +222,17 @@ class SmolVLAPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: SmolVLAConfig,
-        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                     the configuration class is used.
-            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
-                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
 
         super().__init__(config)
         config.validate_features()
         self.config = config
-        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
-        self.normalize_targets = Normalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-        self.unnormalize_outputs = Unnormalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
 
-        self.language_tokenizer = AutoProcessor.from_pretrained(self.config.vlm_model_name).tokenizer
         self.model = VLAFlowMatching(config)
         self.reset()
 
@@ -357,23 +242,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
             ACTION: deque(maxlen=self.config.n_action_steps),
         }
 
-    # HACK(aliberts, danaaubakirova): we overwrite this classmethod here to fix smolVLA-specific issues
-    @classmethod
-    def _load_as_safetensor(
-        cls,
-        model: "SmolVLAPolicy",
-        model_file: str,
-        map_location: str,
-        strict: bool,
-    ):
-        safetensors.torch.load_model(model, model_file, strict=strict, device=map_location)
-        return load_smolvla(
-            model,
-            model_file,
-            device=map_location,
-            checkpoint_keys_mapping="model._orig_mod.//model.",
-        )
-
     def get_optim_params(self) -> dict:
         return self.parameters()
 
@@ -389,7 +257,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
 
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
-        lang_tokens, lang_masks = self.prepare_language(batch)
+        lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
+        lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
 
         actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state, noise=noise)
 
@@ -397,8 +266,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
         original_action_dim = self.config.action_feature.shape[0]
         actions = actions[:, :, :original_action_dim]
 
-        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
-
         if self.config.adapt_to_pi_aloha:
             actions = self._pi_aloha_encode_actions(actions)
 
@@ -408,8 +275,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
 
-        batch = self.normalize_inputs(batch)
-
         return batch
 
     @torch.no_grad()
@@ -450,11 +315,11 @@ class SmolVLAPolicy(PreTrainedPolicy):
         if self.config.adapt_to_pi_aloha:
             batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
             batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION])
-        batch = self.normalize_inputs(batch)
-        batch = self.normalize_targets(batch)
+
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
-        lang_tokens, lang_masks = self.prepare_language(batch)
+        lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
+        lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
         actions = self.prepare_action(batch)
         actions_is_pad = batch.get("actions_id_pad")
         loss_dict = {}
@@ -518,30 +383,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
             img_masks.append(mask)
         return images, img_masks
 
-    def prepare_language(self, batch) -> tuple[Tensor, Tensor]:
-        """Tokenize the text input"""
-        device = batch[OBS_STATE].device
-        tasks = batch["task"]
-        if isinstance(tasks, str):
-            tasks = [tasks]
-
-        if len(tasks) == 1:
-            tasks = [tasks[0] for _ in range(batch[OBS_STATE].shape[0])]
-
-        tasks = [task if task.endswith("\n") else f"{task}\n" for task in tasks]
-
-        tokenized_prompt = self.language_tokenizer.__call__(
-            tasks,
-            padding=self.config.pad_language_to,
-            padding_side="right",
-            max_length=self.config.tokenizer_max_length,
-            return_tensors="pt",
-        )
-        lang_tokens = tokenized_prompt["input_ids"].to(device=device)
-        lang_masks = tokenized_prompt["attention_mask"].to(device=device, dtype=torch.bool)
-
-        return lang_tokens, lang_masks
-
     def _pi_aloha_decode_state(self, state):
         # Flip the joints.
         for motor_idx in [1, 2, 8, 9]:

src/lerobot/policies/smolvla/processor_smolvla.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python
+
+# Copyright 2025 HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    ComplementaryDataProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    ProcessorStepRegistry,
+    RenameObservationsProcessorStep,
+    TokenizerProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+
+def make_smolvla_pre_post_processors(
+    config: SmolVLAConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Constructs pre-processor and post-processor pipelines for the SmolVLA policy.
+
+    The pre-processing pipeline prepares input data for the model by:
+    1.  Renaming features to match pretrained configurations.
+    2.  Normalizing input and output features based on dataset statistics.
+    3.  Adding a batch dimension.
+    4.  Ensuring the language task description ends with a newline character.
+    5.  Tokenizing the language task description.
+    6.  Moving all data to the specified device.
+
+    The post-processing pipeline handles the model's output by:
+    1.  Moving data to the CPU.
+    2.  Unnormalizing the output actions to their original scale.
+
+    Args:
+        config: The configuration object for the SmolVLA policy.
+        dataset_stats: A dictionary of statistics for normalization.
+
+    Returns:
+        A tuple containing the configured pre-processor and post-processor pipelines.
+    """
+
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
+        AddBatchDimensionProcessorStep(),
+        SmolVLANewLineProcessor(),
+        TokenizerProcessorStep(
+            tokenizer_name=config.vlm_model_name,
+            padding=config.pad_language_to,
+            padding_side="right",
+            max_length=config.tokenizer_max_length,
+        ),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )
+
+
+@ProcessorStepRegistry.register(name="smolvla_new_line_processor")
+class SmolVLANewLineProcessor(ComplementaryDataProcessorStep):
+    """
+    A processor step that ensures the 'task' description ends with a newline character.
+
+    This step is necessary for certain tokenizers (e.g., PaliGemma) that expect a
+    newline at the end of the prompt. It handles both single string tasks and lists
+    of string tasks.
+    """
+
+    def complementary_data(self, complementary_data):
+        if "task" not in complementary_data:
+            return complementary_data
+
+        task = complementary_data["task"]
+        if task is None:
+            return complementary_data
+
+        new_complementary_data = dict(complementary_data)
+
+        # Handle both string and list of strings
+        if isinstance(task, str):
+            # Single string: add newline if not present
+            if not task.endswith("\n"):
+                new_complementary_data["task"] = f"{task}\n"
+        elif isinstance(task, list) and all(isinstance(t, str) for t in task):
+            # List of strings: add newline to each if not present
+            new_complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task]
+        # If task is neither string nor list of strings, leave unchanged
+
+        return new_complementary_data
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features

src/lerobot/policies/tdmpc/modeling_tdmpc.py

@@ -36,7 +36,6 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 
 from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_STATE, REWARD
-from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
@@ -63,26 +62,19 @@ class TDMPCPolicy(PreTrainedPolicy):
     config_class = TDMPCConfig
     name = "tdmpc"
 
-    def __init__(self, config: TDMPCConfig, dataset_stats: dict[str, dict[str, Tensor]] | None = None):
+    def __init__(
+        self,
+        config: TDMPCConfig,
+    ):
         """
         Args:
             config: Policy configuration class instance or None, in which case the default instantiation of
                 the configuration class is used.
-            dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected
-                that they will be passed with a call to `load_state_dict` before the policy is used.
         """
         super().__init__(config)
         config.validate_features()
         self.config = config
 
-        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
-        self.normalize_targets = Normalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-        self.unnormalize_outputs = Unnormalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-
         self.model = TDMPCTOLD(config)
         self.model_target = deepcopy(self.model)
         for param in self.model_target.parameters():
@@ -137,7 +129,6 @@ class TDMPCPolicy(PreTrainedPolicy):
 
         actions = torch.clamp(actions, -1, +1)
 
-        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
         return actions
 
     @torch.no_grad()
@@ -147,11 +138,12 @@ class TDMPCPolicy(PreTrainedPolicy):
         if ACTION in batch:
             batch.pop(ACTION)
 
-        batch = self.normalize_inputs(batch)
-
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGE] = batch[next(iter(self.config.image_features))]
+        # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
+        if ACTION in batch:
+            batch.pop(ACTION)
 
         self._queues = populate_queues(self._queues, batch)
 
@@ -320,11 +312,9 @@ class TDMPCPolicy(PreTrainedPolicy):
         """
         device = get_device_from_parameters(self)
 
-        batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
             batch[OBS_IMAGE] = batch[next(iter(self.config.image_features))]
-        batch = self.normalize_targets(batch)
 
         info = {}
 

src/lerobot/policies/tdmpc/processor_tdmpc.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python
+
+# Copyright 2024 Nicklas Hansen, Xiaolong Wang, Hao Su,
+# and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any
+
+import torch
+
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+
+def make_tdmpc_pre_post_processors(
+    config: TDMPCConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Constructs pre-processor and post-processor pipelines for the TDMPC policy.
+
+    The pre-processing pipeline prepares input data for the model by:
+    1. Renaming features to match pretrained configurations.
+    2. Normalizing input and output features based on dataset statistics.
+    3. Adding a batch dimension.
+    4. Moving all data to the specified device.
+
+    The post-processing pipeline handles the model's output by:
+    1. Moving data to the CPU.
+    2. Unnormalizing the output features to their original scale.
+
+    Args:
+        config: The configuration object for the TDMPC policy.
+        dataset_stats: A dictionary of statistics for normalization.
+
+    Returns:
+        A tuple containing the configured pre-processor and post-processor pipelines.
+    """
+
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/vqbet/modeling_vqbet.py

@@ -28,7 +28,6 @@ import torchvision
 from torch import Tensor, nn
 
 from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
-from lerobot.policies.normalize import Normalize, Unnormalize
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
@@ -48,7 +47,6 @@ class VQBeTPolicy(PreTrainedPolicy):
     def __init__(
         self,
         config: VQBeTConfig | None = None,
-        dataset_stats: dict[str, dict[str, Tensor]] | None = None,
     ):
         """
         Args:
@@ -61,14 +59,6 @@ class VQBeTPolicy(PreTrainedPolicy):
         config.validate_features()
         self.config = config
 
-        self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
-        self.normalize_targets = Normalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-        self.unnormalize_outputs = Unnormalize(
-            config.output_features, config.normalization_mapping, dataset_stats
-        )
-
         self.vqbet = VQBeTModel(config)
 
         self.reset()
@@ -128,7 +118,6 @@ class VQBeTPolicy(PreTrainedPolicy):
     def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
         batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
         actions = self.vqbet(batch, rollout=True)[:, : self.config.action_chunk_size]
-        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
         return actions
 
     @torch.no_grad()
@@ -142,10 +131,12 @@ class VQBeTPolicy(PreTrainedPolicy):
         # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
         if ACTION in batch:
             batch.pop(ACTION)
-        batch = self.normalize_inputs(batch)
         batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
         # NOTE: It's important that this happens after stacking the images into a single key.
         batch["observation.images"] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
+        if ACTION in batch:
+            batch.pop(ACTION)
 
         self._queues = populate_queues(self._queues, batch)
 
@@ -165,10 +156,8 @@ class VQBeTPolicy(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Run the batch through the model and compute the loss for training or validation."""
-        batch = self.normalize_inputs(batch)
         batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
         batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
-        batch = self.normalize_targets(batch)
         # VQ-BeT discretizes action using VQ-VAE before training BeT (please refer to section 3.2 in the VQ-BeT paper https://huggingface.co/papers/2403.03181)
         if not self.vqbet.action_head.vqvae_model.discretized.item():
             # loss: total loss of training RVQ

src/lerobot/policies/vqbet/processor_vqbet.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python
+
+# Copyright 2024 Seungjae Lee and Yibin Wang and Haritheja Etukuru
+# and H. Jin Kim and Nur Muhammad Mahi Shafiullah and Lerrel Pinto
+# and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any
+
+import torch
+
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+
+def make_vqbet_pre_post_processors(
+    config: VQBeTConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """
+    Constructs pre-processor and post-processor pipelines for the VQ-BeT policy.
+
+    The pre-processing pipeline prepares input data for the model by:
+    1. Renaming features, allowing customization to match pretrained configurations.
+    2. Normalizing input and output features based on dataset statistics.
+    3. Adding a batch dimension.
+    4. Moving all data to the specified device.
+
+    The post-processing pipeline handles the model's output by:
+    1. Moving data to the CPU.
+    2. Unnormalizing the output features to their original scale.
+
+    Args:
+        config: The configuration object for the VQ-BeT policy.
+        dataset_stats: A dictionary of statistics for normalization.
+
+    Returns:
+        A tuple containing the configured pre-processor and post-processor pipelines.
+    """
+
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),  # Let the possibility to the user to rename the keys
+        AddBatchDimensionProcessorStep(),
+        DeviceProcessorStep(device=config.device),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+    ]
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/processor/__init__.py

@@ -14,41 +14,120 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .device_processor import DeviceProcessor
-from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor
-from .observation_processor import VanillaObservationProcessor
-from .pipeline import (
-    ActionProcessor,
-    DoneProcessor,
+from .batch_processor import AddBatchDimensionProcessorStep
+from .converters import (
+    batch_to_transition,
+    create_transition,
+    transition_to_batch,
+)
+from .core import (
+    EnvAction,
     EnvTransition,
-    IdentityProcessor,
-    InfoProcessor,
-    ObservationProcessor,
+    PolicyAction,
+    RobotAction,
+    RobotObservation,
+    TransitionKey,
+)
+from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
+from .device_processor import DeviceProcessorStep
+from .factory import (
+    make_default_processors,
+    make_default_robot_action_processor,
+    make_default_robot_observation_processor,
+    make_default_teleop_action_processor,
+)
+from .gym_action_processor import (
+    Numpy2TorchActionProcessorStep,
+    Torch2NumpyActionProcessorStep,
+)
+from .hil_processor import (
+    AddTeleopActionAsComplimentaryDataStep,
+    AddTeleopEventsAsInfoStep,
+    GripperPenaltyProcessorStep,
+    ImageCropResizeProcessorStep,
+    InterventionActionProcessorStep,
+    RewardClassifierProcessorStep,
+    TimeLimitProcessorStep,
+)
+from .joint_observations_processor import JointVelocityProcessorStep, MotorCurrentProcessorStep
+from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep, hotswap_stats
+from .observation_processor import VanillaObservationProcessorStep
+from .pipeline import (
+    ActionProcessorStep,
+    ComplementaryDataProcessorStep,
+    DataProcessorPipeline,
+    DoneProcessorStep,
+    IdentityProcessorStep,
+    InfoProcessorStep,
+    ObservationProcessorStep,
+    PolicyActionProcessorStep,
+    PolicyProcessorPipeline,
+    ProcessorKwargs,
     ProcessorStep,
     ProcessorStepRegistry,
-    RewardProcessor,
-    RobotProcessor,
-    TransitionKey,
-    TruncatedProcessor,
+    RewardProcessorStep,
+    RobotActionProcessorStep,
+    RobotProcessorPipeline,
+    TruncatedProcessorStep,
 )
-from .rename_processor import RenameProcessor
+from .policy_robot_bridge import (
+    PolicyActionToRobotActionProcessorStep,
+    RobotActionToPolicyActionProcessorStep,
+)
+from .rename_processor import RenameObservationsProcessorStep
+from .tokenizer_processor import TokenizerProcessorStep
 
 __all__ = [
-    "ActionProcessor",
-    "DeviceProcessor",
-    "DoneProcessor",
+    "ActionProcessorStep",
+    "AddTeleopActionAsComplimentaryDataStep",
+    "AddTeleopEventsAsInfoStep",
+    "ComplementaryDataProcessorStep",
+    "batch_to_transition",
+    "create_transition",
+    "DeviceProcessorStep",
+    "DoneProcessorStep",
+    "EnvAction",
     "EnvTransition",
-    "IdentityProcessor",
-    "InfoProcessor",
-    "NormalizerProcessor",
-    "UnnormalizerProcessor",
-    "ObservationProcessor",
+    "GripperPenaltyProcessorStep",
+    "hotswap_stats",
+    "IdentityProcessorStep",
+    "ImageCropResizeProcessorStep",
+    "InfoProcessorStep",
+    "InterventionActionProcessorStep",
+    "JointVelocityProcessorStep",
+    "make_default_processors",
+    "make_default_teleop_action_processor",
+    "make_default_robot_action_processor",
+    "make_default_robot_observation_processor",
+    "MapDeltaActionToRobotActionStep",
+    "MapTensorToDeltaActionDictStep",
+    "MotorCurrentProcessorStep",
+    "NormalizerProcessorStep",
+    "Numpy2TorchActionProcessorStep",
+    "ObservationProcessorStep",
+    "PolicyAction",
+    "PolicyActionProcessorStep",
+    "PolicyProcessorPipeline",
+    "ProcessorKwargs",
     "ProcessorStep",
     "ProcessorStepRegistry",
-    "RenameProcessor",
-    "RewardProcessor",
-    "RobotProcessor",
+    "RobotAction",
+    "RobotActionProcessorStep",
+    "RobotObservation",
+    "RenameObservationsProcessorStep",
+    "RewardClassifierProcessorStep",
+    "RewardProcessorStep",
+    "DataProcessorPipeline",
+    "TimeLimitProcessorStep",
+    "AddBatchDimensionProcessorStep",
+    "RobotProcessorPipeline",
+    "TokenizerProcessorStep",
+    "Torch2NumpyActionProcessorStep",
+    "RobotActionToPolicyActionProcessorStep",
+    "PolicyActionToRobotActionProcessorStep",
+    "transition_to_batch",
     "TransitionKey",
-    "TruncatedProcessor",
-    "VanillaObservationProcessor",
+    "TruncatedProcessorStep",
+    "UnnormalizerProcessorStep",
+    "VanillaObservationProcessorStep",
 ]

src/lerobot/processor/batch_processor.py

@@ -0,0 +1,254 @@
+#!/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.
+
+"""
+This script defines processor steps for adding a batch dimension to various components of an environment transition.
+
+These steps are designed to process actions, observations, and complementary data, making them suitable for batch processing by adding a leading dimension. This is a common requirement before feeding data into a neural network model.
+"""
+
+from dataclasses import dataclass, field
+
+from torch import Tensor
+
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
+
+from .core import EnvTransition, PolicyAction
+from .pipeline import (
+    ComplementaryDataProcessorStep,
+    ObservationProcessorStep,
+    PolicyActionProcessorStep,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    TransitionKey,
+)
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="to_batch_processor_action")
+class AddBatchDimensionActionStep(PolicyActionProcessorStep):
+    """
+    Processor step to add a batch dimension to a 1D tensor action.
+
+    This is useful for creating a batch of size 1 from a single action sample.
+    """
+
+    def action(self, action: PolicyAction) -> PolicyAction:
+        """
+        Adds a batch dimension to the action if it's a 1D tensor.
+
+        Args:
+            action: The action tensor.
+
+        Returns:
+            The action tensor with an added batch dimension.
+        """
+        if action.dim() != 1:
+            return action
+        return action.unsqueeze(0)
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        Returns the input features unchanged.
+
+        Adding a batch dimension does not alter the feature definition.
+
+        Args:
+            features: A dictionary of policy features.
+
+        Returns:
+            The original dictionary of policy features.
+        """
+        return features
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="to_batch_processor_observation")
+class AddBatchDimensionObservationStep(ObservationProcessorStep):
+    """
+    Processor step to add a batch dimension to observations.
+
+    It handles different types of observations:
+    - State vectors (1D tensors).
+    - Single images (3D tensors).
+    - Dictionaries of multiple images (3D tensors).
+    """
+
+    def observation(self, observation: dict[str, Tensor]) -> dict[str, Tensor]:
+        """
+        Adds a batch dimension to tensor-based observations in the observation dictionary.
+
+        Args:
+            observation: The observation dictionary.
+
+        Returns:
+            The observation dictionary with batch dimensions added to tensors.
+        """
+        # Process state observations - add batch dim if 1D
+        for state_key in [OBS_STATE, OBS_ENV_STATE]:
+            if state_key in observation:
+                state_value = observation[state_key]
+                if isinstance(state_value, Tensor) and state_value.dim() == 1:
+                    observation[state_key] = state_value.unsqueeze(0)
+
+        # Process single image observation - add batch dim if 3D
+        if OBS_IMAGE in observation:
+            image_value = observation[OBS_IMAGE]
+            if isinstance(image_value, Tensor) and image_value.dim() == 3:
+                observation[OBS_IMAGE] = image_value.unsqueeze(0)
+
+        # Process multiple image observations - add batch dim if 3D
+        for key, value in observation.items():
+            if key.startswith(f"{OBS_IMAGES}.") and isinstance(value, Tensor) and value.dim() == 3:
+                observation[key] = value.unsqueeze(0)
+        return observation
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        Returns the input features unchanged.
+
+        Adding a batch dimension does not alter the feature definition.
+
+        Args:
+            features: A dictionary of policy features.
+
+        Returns:
+            The original dictionary of policy features.
+        """
+        return features
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="to_batch_processor_complementary_data")
+class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
+    """
+    Processor step to add a batch dimension to complementary data fields.
+
+    Handles specific keys like 'task', 'index', and 'task_index' to make them batched.
+    - 'task' (str) is wrapped in a list.
+    - 'index' and 'task_index' (0D tensors) get a batch dimension.
+    """
+
+    def complementary_data(self, complementary_data: dict) -> dict:
+        """
+        Adds a batch dimension to specific fields in the complementary data dictionary.
+
+        Args:
+            complementary_data: The complementary data dictionary.
+
+        Returns:
+            The complementary data dictionary with batch dimensions added.
+        """
+        # Process task field - wrap string in list to add batch dimension
+        if "task" in complementary_data:
+            task_value = complementary_data["task"]
+            if isinstance(task_value, str):
+                complementary_data["task"] = [task_value]
+
+        # Process index field - add batch dim if 0D
+        if "index" in complementary_data:
+            index_value = complementary_data["index"]
+            if isinstance(index_value, Tensor) and index_value.dim() == 0:
+                complementary_data["index"] = index_value.unsqueeze(0)
+
+        # Process task_index field - add batch dim if 0D
+        if "task_index" in complementary_data:
+            task_index_value = complementary_data["task_index"]
+            if isinstance(task_index_value, Tensor) and task_index_value.dim() == 0:
+                complementary_data["task_index"] = task_index_value.unsqueeze(0)
+        return complementary_data
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        Returns the input features unchanged.
+
+        Adding a batch dimension does not alter the feature definition.
+
+        Args:
+            features: A dictionary of policy features.
+
+        Returns:
+            The original dictionary of policy features.
+        """
+        return features
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="to_batch_processor")
+class AddBatchDimensionProcessorStep(ProcessorStep):
+    """
+    A composite processor step that adds a batch dimension to the entire environment transition.
+
+    This step combines individual processors for actions, observations, and complementary data
+    to create a batched transition (batch size 1) from a single-instance transition.
+
+    Attributes:
+        to_batch_action_processor: Processor for the action component.
+        to_batch_observation_processor: Processor for the observation component.
+        to_batch_complementary_data_processor: Processor for the complementary data component.
+    """
+
+    to_batch_action_processor: AddBatchDimensionActionStep = field(
+        default_factory=AddBatchDimensionActionStep
+    )
+    to_batch_observation_processor: AddBatchDimensionObservationStep = field(
+        default_factory=AddBatchDimensionObservationStep
+    )
+    to_batch_complementary_data_processor: AddBatchDimensionComplementaryDataStep = field(
+        default_factory=AddBatchDimensionComplementaryDataStep
+    )
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        """
+        Applies the batching process to all relevant parts of an environment transition.
+
+        Args:
+            transition: The environment transition to process.
+
+        Returns:
+            The environment transition with a batch dimension added.
+        """
+        if transition[TransitionKey.ACTION] is not None:
+            transition = self.to_batch_action_processor(transition)
+        if transition[TransitionKey.OBSERVATION] is not None:
+            transition = self.to_batch_observation_processor(transition)
+        if transition[TransitionKey.COMPLEMENTARY_DATA] is not None:
+            transition = self.to_batch_complementary_data_processor(transition)
+        return transition
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        Returns the input features unchanged.
+
+        Adding a batch dimension does not alter the feature definition.
+
+        Args:
+            features: A dictionary of policy features.
+
+        Returns:
+            The original dictionary of policy features.
+        """
+        # NOTE: We ignore the batch dimension when transforming features
+        return features

src/lerobot/processor/converters.py

@@ -0,0 +1,412 @@
+# !/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import singledispatch
+from typing import Any
+
+import numpy as np
+import torch
+
+from .core import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
+
+
+@singledispatch
+def to_tensor(
+    value: Any,
+    *,
+    dtype: torch.dtype | None = torch.float32,
+    device: torch.device | str | None = None,
+) -> torch.Tensor:
+    """
+    Convert various data types to PyTorch tensors with configurable options.
+
+    This is a unified tensor conversion function using single dispatch to handle
+    different input types appropriately.
+
+    Args:
+        value: Input value to convert (tensor, array, scalar, sequence, etc.).
+        dtype: Target tensor dtype. If None, preserves original dtype.
+        device: Target device for the tensor.
+
+    Returns:
+        A PyTorch tensor.
+
+    Raises:
+        TypeError: If the input type is not supported.
+    """
+    raise TypeError(f"Unsupported type for tensor conversion: {type(value)}")
+
+
+@to_tensor.register(torch.Tensor)
+def _(value: torch.Tensor, *, dtype=torch.float32, device=None, **kwargs) -> torch.Tensor:
+    """Handle conversion for existing PyTorch tensors."""
+    if dtype is not None:
+        value = value.to(dtype=dtype)
+    if device is not None:
+        value = value.to(device=device)
+    return value
+
+
+@to_tensor.register(np.ndarray)
+def _(
+    value: np.ndarray,
+    *,
+    dtype=torch.float32,
+    device=None,
+    **kwargs,
+) -> torch.Tensor:
+    """Handle conversion for numpy arrays."""
+    # Check for numpy scalars (0-dimensional arrays) and treat them as scalars.
+    if value.ndim == 0:
+        # Numpy scalars should be converted to 0-dimensional tensors.
+        scalar_value = value.item()
+        return torch.tensor(scalar_value, dtype=dtype, device=device)
+
+    # Create tensor from numpy array.
+    tensor = torch.from_numpy(value)
+
+    # Apply dtype and device conversion if specified.
+    if dtype is not None:
+        tensor = tensor.to(dtype=dtype)
+    if device is not None:
+        tensor = tensor.to(device=device)
+
+    return tensor
+
+
+@to_tensor.register(int)
+@to_tensor.register(float)
+@to_tensor.register(np.integer)
+@to_tensor.register(np.floating)
+def _(value, *, dtype=torch.float32, device=None, **kwargs) -> torch.Tensor:
+    """Handle conversion for scalar values including numpy scalars."""
+    return torch.tensor(value, dtype=dtype, device=device)
+
+
+@to_tensor.register(list)
+@to_tensor.register(tuple)
+def _(value: Sequence, *, dtype=torch.float32, device=None, **kwargs) -> torch.Tensor:
+    """Handle conversion for sequences (lists, tuples)."""
+    return torch.tensor(value, dtype=dtype, device=device)
+
+
+@to_tensor.register(dict)
+def _(value: dict, *, device=None, **kwargs) -> dict:
+    """Handle conversion for dictionaries by recursively converting their values to tensors."""
+    if not value:
+        return {}
+
+    result = {}
+    for key, sub_value in value.items():
+        if sub_value is None:
+            continue
+
+        if isinstance(sub_value, dict):
+            # Recursively process nested dictionaries.
+            result[key] = to_tensor(
+                sub_value,
+                device=device,
+                **kwargs,
+            )
+            continue
+
+        # Convert individual values to tensors.
+        result[key] = to_tensor(
+            sub_value,
+            device=device,
+            **kwargs,
+        )
+    return result
+
+
+def from_tensor_to_numpy(x: torch.Tensor | Any) -> np.ndarray | float | int | Any:
+    """
+    Convert a PyTorch tensor to a numpy array or scalar if applicable.
+
+    If the input is not a tensor, it is returned unchanged.
+
+    Args:
+        x: The input, which can be a tensor or any other type.
+
+    Returns:
+        A numpy array, a scalar, or the original input.
+    """
+    if isinstance(x, torch.Tensor):
+        return x.item() if x.numel() == 1 else x.detach().cpu().numpy()
+    return x
+
+
+def _extract_complementary_data(batch: dict[str, Any]) -> dict[str, Any]:
+    """
+    Extract complementary data from a batch dictionary.
+
+    This includes padding flags, task description, and indices.
+
+    Args:
+        batch: The batch dictionary.
+
+    Returns:
+        A dictionary with the extracted complementary data.
+    """
+    pad_keys = {k: v for k, v in batch.items() if "_is_pad" in k}
+    task_key = {"task": batch["task"]} if "task" in batch else {}
+    index_key = {"index": batch["index"]} if "index" in batch else {}
+    task_index_key = {"task_index": batch["task_index"]} if "task_index" in batch else {}
+
+    return {**pad_keys, **task_key, **index_key, **task_index_key}
+
+
+def create_transition(
+    observation: dict[str, Any] | None = None,
+    action: PolicyAction | RobotAction | None = None,
+    reward: float = 0.0,
+    done: bool = False,
+    truncated: bool = False,
+    info: dict[str, Any] | None = None,
+    complementary_data: dict[str, Any] | None = None,
+) -> EnvTransition:
+    """
+    Create an `EnvTransition` dictionary with sensible defaults.
+
+    Args:
+        observation: Observation dictionary.
+        action: Action dictionary.
+        reward: Scalar reward value.
+        done: Episode termination flag.
+        truncated: Episode truncation flag.
+        info: Additional info dictionary.
+        complementary_data: Complementary data dictionary.
+
+    Returns:
+        A complete `EnvTransition` dictionary.
+    """
+    return {
+        TransitionKey.OBSERVATION: observation,
+        TransitionKey.ACTION: action,
+        TransitionKey.REWARD: reward,
+        TransitionKey.DONE: done,
+        TransitionKey.TRUNCATED: truncated,
+        TransitionKey.INFO: info if info is not None else {},
+        TransitionKey.COMPLEMENTARY_DATA: complementary_data if complementary_data is not None else {},
+    }
+
+
+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`.
+
+    Args:
+        action: The raw action dictionary from a teleoperation device or controller.
+
+    Returns:
+        An `EnvTransition` containing the formatted action.
+    """
+    if not isinstance(action, dict):
+        raise ValueError(f"Action should be a RobotAction type got {type(action)}")
+    return create_transition(action=action)
+
+
+def observation_to_transition(observation: RobotObservation) -> EnvTransition:
+    """
+    Convert a raw robot observation dictionary into a standardized `EnvTransition`.
+
+    Args:
+        observation: The raw observation dictionary from the environment.
+
+    Returns:
+        An `EnvTransition` containing the formatted observation.
+    """
+    if not isinstance(observation, dict):
+        raise ValueError(f"Observation should be a RobotObservation type got {type(observation)}")
+    return create_transition(observation=observation)
+
+
+def transition_to_robot_action(transition: EnvTransition) -> RobotAction:
+    """
+    Extract a raw robot action dictionary for a robot from an `EnvTransition`.
+
+    This function searches for keys in the format "action.*.pos" or "action.*.vel"
+    and converts them into a flat dictionary suitable for sending to a robot controller.
+
+    Args:
+        transition: The `EnvTransition` containing the action.
+
+    Returns:
+        A dictionary representing the raw robot action.
+    """
+    if not isinstance(transition, dict):
+        raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
+
+    action = transition.get(TransitionKey.ACTION)
+    if not isinstance(action, dict):
+        raise ValueError(f"Action should be a RobotAction type (dict) got {type(action)}")
+    return transition.get(TransitionKey.ACTION)
+
+
+def transition_to_policy_action(transition: EnvTransition) -> PolicyAction:
+    """
+    Convert an `EnvTransition` to a `PolicyAction`.
+    """
+    if not isinstance(transition, dict):
+        raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
+
+    action = transition.get(TransitionKey.ACTION)
+    if not isinstance(action, PolicyAction):
+        raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
+    return action
+
+
+def transition_to_observation(transition: EnvTransition) -> RobotObservation:
+    """
+    Convert an `EnvTransition` to a `RobotObservation`.
+    """
+    if not isinstance(transition, dict):
+        raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
+
+    observation = transition.get(TransitionKey.OBSERVATION)
+    if not isinstance(observation, dict):
+        raise ValueError(f"Observation should be a RobotObservation (dict) type got {type(observation)}")
+    return observation
+
+
+def policy_action_to_transition(action: PolicyAction) -> EnvTransition:
+    """
+    Convert a `PolicyAction` to an `EnvTransition`.
+    """
+    if not isinstance(action, PolicyAction):
+        raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
+    return create_transition(action=action)
+
+
+def batch_to_transition(batch: dict[str, Any]) -> EnvTransition:
+    """
+    Convert a batch dictionary from a dataset/dataloader into an `EnvTransition`.
+
+    This function maps recognized keys from a batch to the `EnvTransition` structure,
+    filling in missing keys with sensible defaults.
+
+    Args:
+        batch: A batch dictionary.
+
+    Returns:
+        An `EnvTransition` dictionary.
+
+    Raises:
+        ValueError: If the input is not a dictionary.
+    """
+
+    # Validate input type.
+    if not isinstance(batch, dict):
+        raise ValueError(f"EnvTransition must be a dictionary. Got {type(batch).__name__}")
+
+    action = batch.get("action")
+    if action is not None and not isinstance(action, PolicyAction):
+        raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
+
+    # Extract observation and complementary data keys.
+    observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
+    complementary_data = _extract_complementary_data(batch)
+
+    return create_transition(
+        observation=observation_keys if observation_keys else None,
+        action=batch.get("action"),
+        reward=batch.get("next.reward", 0.0),
+        done=batch.get("next.done", False),
+        truncated=batch.get("next.truncated", False),
+        info=batch.get("info", {}),
+        complementary_data=complementary_data if complementary_data else None,
+    )
+
+
+def transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
+    """
+    Convert an `EnvTransition` back to the canonical batch format used in LeRobot.
+
+    This is the inverse of `batch_to_transition`.
+
+    Args:
+        transition: The `EnvTransition` to convert.
+
+    Returns:
+        A batch dictionary with canonical LeRobot field names.
+    """
+    if not isinstance(transition, dict):
+        raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")
+
+    batch = {
+        "action": transition.get(TransitionKey.ACTION),
+        "next.reward": transition.get(TransitionKey.REWARD, 0.0),
+        "next.done": transition.get(TransitionKey.DONE, False),
+        "next.truncated": transition.get(TransitionKey.TRUNCATED, False),
+        "info": transition.get(TransitionKey.INFO, {}),
+    }
+
+    # Add complementary data.
+    comp_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
+    if comp_data:
+        batch.update(comp_data)
+
+    # Flatten observation dictionary.
+    observation = transition.get(TransitionKey.OBSERVATION)
+    if isinstance(observation, dict):
+        batch.update(observation)
+
+    return batch
+
+
+def identity_transition(transition: EnvTransition) -> EnvTransition:
+    """
+    An identity function for transitions, returning the input unchanged.
+
+    Useful as a default or placeholder in processing pipelines.
+
+    Args:
+        tr: An `EnvTransition`.
+
+    Returns:
+        The same `EnvTransition`.
+    """
+    return transition

src/lerobot/processor/core.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any, TypeAlias, TypedDict
+
+import numpy as np
+import torch
+
+
+class TransitionKey(str, Enum):
+    """Keys for accessing EnvTransition dictionary components."""
+
+    # TODO(Steven): Use consts
+    OBSERVATION = "observation"
+    ACTION = "action"
+    REWARD = "reward"
+    DONE = "done"
+    TRUNCATED = "truncated"
+    INFO = "info"
+    COMPLEMENTARY_DATA = "complementary_data"
+
+
+PolicyAction: TypeAlias = torch.Tensor
+RobotAction: TypeAlias = dict[str, Any]
+EnvAction: TypeAlias = np.ndarray
+RobotObservation: TypeAlias = dict[str, Any]
+
+
+EnvTransition = TypedDict(
+    "EnvTransition",
+    {
+        TransitionKey.OBSERVATION.value: dict[str, Any] | None,
+        TransitionKey.ACTION.value: PolicyAction | RobotAction | EnvAction | None,
+        TransitionKey.REWARD.value: float | torch.Tensor | None,
+        TransitionKey.DONE.value: bool | torch.Tensor | None,
+        TransitionKey.TRUNCATED.value: bool | torch.Tensor | None,
+        TransitionKey.INFO.value: dict[str, Any] | None,
+        TransitionKey.COMPLEMENTARY_DATA.value: dict[str, Any] | None,
+    },
+)

src/lerobot/processor/delta_action_processor.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+
+from .core import PolicyAction, RobotAction
+from .pipeline import ActionProcessorStep, ProcessorStepRegistry, RobotActionProcessorStep
+
+
+@ProcessorStepRegistry.register("map_tensor_to_delta_action_dict")
+@dataclass
+class MapTensorToDeltaActionDictStep(ActionProcessorStep):
+    """
+    Maps a flat action tensor from a policy to a structured delta action dictionary.
+
+    This step is typically used after a policy outputs a continuous action vector.
+    It decomposes the vector into named components for delta movements of the
+    end-effector (x, y, z) and optionally the gripper.
+
+    Attributes:
+        use_gripper: If True, assumes the 4th element of the tensor is the
+                     gripper action.
+    """
+
+    use_gripper: bool = True
+
+    def action(self, action: PolicyAction) -> RobotAction:
+        if not isinstance(action, PolicyAction):
+            raise ValueError("Only PolicyAction is supported for this processor")
+
+        if action.dim() > 1:
+            action = action.squeeze(0)
+
+        # TODO (maractingi): add rotation
+        delta_action = {
+            "delta_x": action[0].item(),
+            "delta_y": action[1].item(),
+            "delta_z": action[2].item(),
+        }
+        if self.use_gripper:
+            delta_action["gripper"] = action[3].item()
+        return delta_action
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        for axis in ["x", "y", "z"]:
+            features[PipelineFeatureType.ACTION][f"delta_{axis}"] = PolicyFeature(
+                type=FeatureType.ACTION, shape=(1,)
+            )
+
+        if self.use_gripper:
+            features[PipelineFeatureType.ACTION]["gripper"] = PolicyFeature(
+                type=FeatureType.ACTION, shape=(1,)
+            )
+        return features
+
+
+@ProcessorStepRegistry.register("map_delta_action_to_robot_action")
+@dataclass
+class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
+    """
+    Maps delta actions from teleoperators to robot target actions for inverse kinematics.
+
+    This step converts a dictionary of delta movements (e.g., from a gamepad)
+    into a target action format that includes an "enabled" flag and target
+    end-effector positions. It also handles scaling and noise filtering.
+
+    Attributes:
+        position_scale: A factor to scale the delta position inputs.
+        rotation_scale: A factor to scale the delta rotation inputs (currently unused).
+        noise_threshold: The magnitude below which delta inputs are considered noise
+                         and do not trigger an "enabled" state.
+    """
+
+    # Scale factors for delta movements
+    position_scale: float = 1.0
+    rotation_scale: float = 0.0  # No rotation deltas for gamepad/keyboard
+    noise_threshold: float = 1e-3  # 1 mm threshold to filter out noise
+
+    def action(self, action: RobotAction) -> RobotAction:
+        # NOTE (maractingi): Action can be a dict from the teleop_devices or a tensor from the policy
+        # TODO (maractingi): changing this target_xyz naming convention from the teleop_devices
+        delta_x = action.pop("delta_x")
+        delta_y = action.pop("delta_y")
+        delta_z = action.pop("delta_z")
+        gripper = action.pop("gripper")
+
+        # Determine if the teleoperator is actively providing input
+        # Consider enabled if any significant movement delta is detected
+        position_magnitude = (delta_x**2 + delta_y**2 + delta_z**2) ** 0.5  # Use Euclidean norm for position
+        enabled = position_magnitude > self.noise_threshold  # Small threshold to avoid noise
+
+        # Scale the deltas appropriately
+        scaled_delta_x = delta_x * self.position_scale
+        scaled_delta_y = delta_y * self.position_scale
+        scaled_delta_z = delta_z * self.position_scale
+
+        # For gamepad/keyboard, we don't have rotation input, so set to 0
+        # These could be extended in the future for more sophisticated teleoperators
+        target_wx = 0.0
+        target_wy = 0.0
+        target_wz = 0.0
+
+        # Update action with robot target format
+        action = {
+            "enabled": enabled,
+            "target_x": scaled_delta_x,
+            "target_y": scaled_delta_y,
+            "target_z": scaled_delta_z,
+            "target_wx": target_wx,
+            "target_wy": target_wy,
+            "target_wz": target_wz,
+            "gripper_vel": float(gripper),
+        }
+
+        return action
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        for axis in ["x", "y", "z", "gripper"]:
+            features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
+
+        for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
+            features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
+                type=FeatureType.ACTION, shape=(1,)
+            )
+
+        return features

src/lerobot/processor/device_processor.py

@@ -13,70 +13,182 @@
 # 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 defines a processor step for moving environment transition data to a specific torch device and casting
+its floating-point precision.
+"""
+
 from dataclasses import dataclass
 from typing import Any
 
 import torch
 
-from lerobot.configs.types import PolicyFeature
-from lerobot.processor.pipeline import EnvTransition, TransitionKey
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
 from lerobot.utils.utils import get_safe_torch_device
 
+from .core import EnvTransition, PolicyAction, TransitionKey
+from .pipeline import ProcessorStep, ProcessorStepRegistry
 
+
+@ProcessorStepRegistry.register("device_processor")
 @dataclass
-class DeviceProcessor:
-    """Processes transitions by moving tensors to the specified device.
+class DeviceProcessorStep(ProcessorStep):
+    """
+    Processor step to move all tensors within an `EnvTransition` to a specified device and optionally cast their
+    floating-point data type.
 
-    This processor ensures that all tensors in the transition are moved to the
-    specified device (CPU or GPU) before they are returned.
+    This is crucial for preparing data for model training or inference on hardware like GPUs.
+
+    Attributes:
+        device: The target device for tensors (e.g., "cpu", "cuda", "cuda:0").
+        float_dtype: The target floating-point dtype as a string (e.g., "float32", "float16", "bfloat16").
+                     If None, the dtype is not changed.
     """
 
-    device: torch.device = "cpu"
+    device: str = "cpu"
+    float_dtype: str | None = None
+
+    DTYPE_MAPPING = {
+        "float16": torch.float16,
+        "float32": torch.float32,
+        "float64": torch.float64,
+        "bfloat16": torch.bfloat16,
+        "half": torch.float16,
+        "float": torch.float32,
+        "double": torch.float64,
+    }
 
     def __post_init__(self):
-        self.device = get_safe_torch_device(self.device)
+        """
+        Initializes the processor by converting string configurations to torch objects.
+
+        This method sets up the `torch.device`, determines if transfers can be non-blocking, and validates the
+        `float_dtype` string, converting it to a `torch.dtype` object.
+        """
+        self.tensor_device: torch.device = get_safe_torch_device(self.device)
+        # Update device string in case a specific GPU was selected (e.g., "cuda" -> "cuda:0")
+        self.device = self.tensor_device.type
         self.non_blocking = "cuda" in str(self.device)
 
+        # Validate and convert float_dtype string to torch dtype
+        if self.float_dtype is not None:
+            if self.float_dtype not in self.DTYPE_MAPPING:
+                raise ValueError(
+                    f"Invalid float_dtype '{self.float_dtype}'. Available options: {list(self.DTYPE_MAPPING.keys())}"
+                )
+            self._target_float_dtype = self.DTYPE_MAPPING[self.float_dtype]
+        else:
+            self._target_float_dtype = None
+
+    def _process_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
+        """
+        Moves a single tensor to the target device and casts its dtype.
+
+        Handles multi-GPU scenarios by not moving a tensor if it's already on a different CUDA device than
+        the target, which is useful when using frameworks like Accelerate.
+
+        Args:
+            tensor: The input torch.Tensor.
+
+        Returns:
+            The processed tensor on the correct device and with the correct dtype.
+        """
+        # Determine target device
+        if tensor.is_cuda and self.tensor_device.type == "cuda":
+            # Both tensor and target are on GPU - preserve tensor's GPU placement.
+            # This handles multi-GPU scenarios where Accelerate has already placed
+            # tensors on the correct GPU for each process.
+            target_device = tensor.device
+        else:
+            # Either tensor is on CPU, or we're configured for CPU.
+            # 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)
+
+        # Convert float dtype if specified and tensor is floating point
+        if self._target_float_dtype is not None and tensor.is_floating_point():
+            tensor = tensor.to(dtype=self._target_float_dtype)
+
+        return tensor
+
     def __call__(self, transition: EnvTransition) -> EnvTransition:
-        # Create a copy of the transition
+        """
+        Applies device and dtype conversion to all tensors in an environment transition.
+
+        It iterates through the transition, finds all `torch.Tensor` objects (including those nested in
+        dictionaries like `observation`), and processes them.
+
+        Args:
+            transition: The input `EnvTransition` object.
+
+        Returns:
+            A new `EnvTransition` object with all tensors moved to the target device and dtype.
+        """
         new_transition = transition.copy()
+        action = new_transition.get(TransitionKey.ACTION)
 
-        # Process observation tensors
-        observation = transition.get(TransitionKey.OBSERVATION)
-        if observation is not None:
-            new_observation = {
-                k: v.to(self.device, non_blocking=self.non_blocking) if isinstance(v, torch.Tensor) else v
-                for k, v in observation.items()
-            }
-            new_transition[TransitionKey.OBSERVATION] = new_observation
+        if action is not None and not isinstance(action, PolicyAction):
+            raise ValueError(f"If action is not None should be a PolicyAction type got {type(action)}")
 
-        # Process action tensor
-        action = transition.get(TransitionKey.ACTION)
-        if action is not None and isinstance(action, torch.Tensor):
-            new_transition[TransitionKey.ACTION] = action.to(self.device, non_blocking=self.non_blocking)
+        simple_tensor_keys = [
+            TransitionKey.ACTION,
+            TransitionKey.REWARD,
+            TransitionKey.DONE,
+            TransitionKey.TRUNCATED,
+        ]
 
-        # Process reward tensor
-        reward = transition.get(TransitionKey.REWARD)
-        if reward is not None and isinstance(reward, torch.Tensor):
-            new_transition[TransitionKey.REWARD] = reward.to(self.device, non_blocking=self.non_blocking)
+        dict_tensor_keys = [
+            TransitionKey.OBSERVATION,
+            TransitionKey.COMPLEMENTARY_DATA,
+        ]
 
-        # Process done tensor
-        done = transition.get(TransitionKey.DONE)
-        if done is not None and isinstance(done, torch.Tensor):
-            new_transition[TransitionKey.DONE] = done.to(self.device, non_blocking=self.non_blocking)
+        # Process simple, top-level tensors
+        for key in simple_tensor_keys:
+            value = transition.get(key)
+            if isinstance(value, torch.Tensor):
+                new_transition[key] = self._process_tensor(value)
 
-        # Process truncated tensor
-        truncated = transition.get(TransitionKey.TRUNCATED)
-        if truncated is not None and isinstance(truncated, torch.Tensor):
-            new_transition[TransitionKey.TRUNCATED] = truncated.to(
-                self.device, non_blocking=self.non_blocking
-            )
+        # Process tensors nested within dictionaries
+        for key in dict_tensor_keys:
+            data_dict = transition.get(key)
+            if data_dict is not None:
+                new_data_dict = {
+                    k: self._process_tensor(v) if isinstance(v, torch.Tensor) else v
+                    for k, v in data_dict.items()
+                }
+                new_transition[key] = new_data_dict
 
         return new_transition
 
     def get_config(self) -> dict[str, Any]:
-        """Return configuration for serialization."""
-        return {"device": self.device}
+        """
+        Returns the serializable configuration of the processor.
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        Returns:
+            A dictionary containing the device and float_dtype settings.
+        """
+        return {"device": self.device, "float_dtype": self.float_dtype}
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        Returns the input features unchanged.
+
+        Device and dtype transformations do not alter the fundamental definition of the features (e.g., shape).
+
+        Args:
+            features: A dictionary of policy features.
+
+        Returns:
+            The original dictionary of policy features.
+        """
         return features

src/lerobot/processor/factory.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .converters import (
+    observation_to_transition,
+    robot_action_observation_to_transition,
+    transition_to_observation,
+    transition_to_robot_action,
+)
+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](
+        steps=[IdentityProcessorStep()],
+        to_transition=robot_action_observation_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](
+        steps=[IdentityProcessorStep()],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
+    return robot_action_processor
+
+
+def make_default_robot_observation_processor() -> RobotProcessorPipeline[RobotObservation, RobotObservation]:
+    robot_observation_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
+        steps=[IdentityProcessorStep()],
+        to_transition=observation_to_transition,
+        to_output=transition_to_observation,
+    )
+    return robot_observation_processor
+
+
+def make_default_processors():
+    teleop_action_processor = make_default_teleop_action_processor()
+    robot_action_processor = make_default_robot_action_processor()
+    robot_observation_processor = make_default_robot_observation_processor()
+    return (teleop_action_processor, robot_action_processor, robot_observation_processor)

src/lerobot/processor/gym_action_processor.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+
+from .converters import to_tensor
+from .core import EnvAction, EnvTransition, PolicyAction
+from .pipeline import ActionProcessorStep, ProcessorStep, ProcessorStepRegistry
+
+
+@ProcessorStepRegistry.register("torch2numpy_action_processor")
+@dataclass
+class Torch2NumpyActionProcessorStep(ActionProcessorStep):
+    """
+    Converts a PyTorch tensor action to a NumPy array.
+
+    This step is useful when the output of a policy (typically a torch.Tensor)
+    needs to be passed to an environment or component that expects a NumPy array.
+
+    Attributes:
+        squeeze_batch_dim: If True, removes the first dimension of the array
+                           if it is of size 1. This is useful for converting a
+                           batched action of size (1, D) to a single action of size (D,).
+    """
+
+    squeeze_batch_dim: bool = True
+
+    def action(self, action: PolicyAction) -> EnvAction:
+        if not isinstance(action, PolicyAction):
+            raise TypeError(
+                f"Expected PolicyAction or None, got {type(action).__name__}. "
+                "Use appropriate processor for non-tensor actions."
+            )
+
+        numpy_action = action.detach().cpu().numpy()
+
+        # Remove batch dimensions but preserve action dimensions.
+        # Only squeeze if there's a batch dimension (first dim == 1).
+        if (
+            self.squeeze_batch_dim
+            and numpy_action.shape
+            and len(numpy_action.shape) > 1
+            and numpy_action.shape[0] == 1
+        ):
+            numpy_action = numpy_action.squeeze(0)
+
+        return numpy_action
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@ProcessorStepRegistry.register("numpy2torch_action_processor")
+@dataclass
+class Numpy2TorchActionProcessorStep(ProcessorStep):
+    """Converts a NumPy array action to a PyTorch tensor when action is present."""
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        """Converts numpy action to torch tensor if action exists, otherwise passes through."""
+        from .core import TransitionKey
+
+        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 transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features

src/lerobot/processor/hil_processor.py

@@ -0,0 +1,596 @@
+#!/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 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 math
+import time
+from dataclasses import dataclass
+from typing import Any, Protocol, TypeVar, runtime_checkable
+
+import numpy as np
+import torch
+import torchvision.transforms.functional as F  # noqa: N812
+
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.teleoperators.teleoperator import Teleoperator
+from lerobot.teleoperators.utils import TeleopEvents
+
+from .core import EnvTransition, PolicyAction, TransitionKey
+from .pipeline import (
+    ComplementaryDataProcessorStep,
+    InfoProcessorStep,
+    ObservationProcessorStep,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    TruncatedProcessorStep,
+)
+
+GRIPPER_KEY = "gripper"
+DISCRETE_PENALTY_KEY = "discrete_penalty"
+TELEOP_ACTION_KEY = "teleop_action"
+
+
+@runtime_checkable
+class HasTeleopEvents(Protocol):
+    """
+    Minimal protocol for objects that provide teleoperation events.
+
+    This protocol defines the `get_teleop_events()` method, allowing processor
+    steps to interact with teleoperators that support event-based controls
+    (like episode termination or success flagging) without needing to know the
+    teleoperator's specific class.
+    """
+
+    def get_teleop_events(self) -> dict[str, Any]:
+        """
+        Get extra control events from the teleoperator.
+
+        Returns:
+            A dictionary containing control events such as:
+            - `is_intervention`: bool - Whether the human is currently intervening.
+            - `terminate_episode`: bool - Whether to terminate the current episode.
+            - `success`: bool - Whether the episode was successful.
+            - `rerecord_episode`: bool - Whether to rerecord the episode.
+        """
+        ...
+
+
+# Type variable constrained to Teleoperator subclasses that also implement events
+TeleopWithEvents = TypeVar("TeleopWithEvents", bound=Teleoperator)
+
+
+def _check_teleop_with_events(teleop: Teleoperator) -> None:
+    """
+    Runtime check that a teleoperator implements the `HasTeleopEvents` protocol.
+
+    Args:
+        teleop: The teleoperator instance to check.
+
+    Raises:
+        TypeError: If the teleoperator does not have a `get_teleop_events` method.
+    """
+    if not isinstance(teleop, HasTeleopEvents):
+        raise TypeError(
+            f"Teleoperator {type(teleop).__name__} must implement get_teleop_events() method. "
+            f"Compatible teleoperators: GamepadTeleop, KeyboardEndEffectorTeleop"
+        )
+
+
+@ProcessorStepRegistry.register("add_teleop_action_as_complementary_data")
+@dataclass
+class AddTeleopActionAsComplimentaryDataStep(ComplementaryDataProcessorStep):
+    """
+    Adds the raw action from a teleoperator to the transition's complementary data.
+
+    This is useful for human-in-the-loop scenarios where the human's input needs to
+    be available to downstream processors, for example, to override a policy's action
+    during an intervention.
+
+    Attributes:
+        teleop_device: The teleoperator instance to get the action from.
+    """
+
+    teleop_device: Teleoperator
+
+    def complementary_data(self, complementary_data: dict) -> dict:
+        """
+        Retrieves the teleoperator's action and adds it to the complementary data.
+
+        Args:
+            complementary_data: The incoming complementary data dictionary.
+
+        Returns:
+            A new dictionary with the teleoperator action added under the
+            `teleop_action` key.
+        """
+        new_complementary_data = dict(complementary_data)
+        new_complementary_data[TELEOP_ACTION_KEY] = self.teleop_device.get_action()
+        return new_complementary_data
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@ProcessorStepRegistry.register("add_teleop_action_as_info")
+@dataclass
+class AddTeleopEventsAsInfoStep(InfoProcessorStep):
+    """
+    Adds teleoperator control events (e.g., terminate, success) to the transition's info.
+
+    This step extracts control events from teleoperators that support event-based
+    interaction, making these signals available to other parts of the system.
+
+    Attributes:
+        teleop_device: An instance of a teleoperator that implements the
+                       `HasTeleopEvents` protocol.
+    """
+
+    teleop_device: TeleopWithEvents
+
+    def __post_init__(self):
+        """Validates that the provided teleoperator supports events after initialization."""
+        _check_teleop_with_events(self.teleop_device)
+
+    def info(self, info: dict) -> dict:
+        """
+        Retrieves teleoperator events and updates the info dictionary.
+
+        Args:
+            info: The incoming info dictionary.
+
+        Returns:
+            A new dictionary including the teleoperator events.
+        """
+        new_info = dict(info)
+
+        teleop_events = self.teleop_device.get_teleop_events()
+        new_info.update(teleop_events)
+        return new_info
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@ProcessorStepRegistry.register("image_crop_resize_processor")
+@dataclass
+class ImageCropResizeProcessorStep(ObservationProcessorStep):
+    """
+    Crops and/or resizes image observations.
+
+    This step iterates through all image keys in an observation dictionary and applies
+    the specified transformations. It handles device placement, moving tensors to the
+    CPU if necessary for operations not supported on certain accelerators like MPS.
+
+    Attributes:
+        crop_params_dict: A dictionary mapping image keys to cropping parameters
+                          (top, left, height, width).
+        resize_size: A tuple (height, width) to resize all images to.
+    """
+
+    crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
+    resize_size: tuple[int, int] | None = None
+
+    def observation(self, observation: dict) -> dict:
+        """
+        Applies cropping and resizing to all images in the observation dictionary.
+
+        Args:
+            observation: The observation dictionary, potentially containing image tensors.
+
+        Returns:
+            A new observation dictionary with transformed images.
+        """
+        if self.resize_size is None and not self.crop_params_dict:
+            return observation
+
+        new_observation = dict(observation)
+
+        # Process all image keys in the observation
+        for key in observation:
+            if "image" not in key:
+                continue
+
+            image = observation[key]
+            device = image.device
+            # NOTE (maractingi): No mps kernel for crop and resize, so we need to move to cpu
+            if device.type == "mps":
+                image = image.cpu()
+            # Crop if crop params are provided for this key
+            if self.crop_params_dict is not None and key in self.crop_params_dict:
+                crop_params = self.crop_params_dict[key]
+                image = F.crop(image, *crop_params)
+            if self.resize_size is not None:
+                image = F.resize(image, self.resize_size)
+                image = image.clamp(0.0, 1.0)
+            new_observation[key] = image.to(device)
+
+        return new_observation
+
+    def get_config(self) -> dict[str, Any]:
+        """
+        Returns the configuration of the step for serialization.
+
+        Returns:
+            A dictionary with the crop parameters and resize dimensions.
+        """
+        return {
+            "crop_params_dict": self.crop_params_dict,
+            "resize_size": self.resize_size,
+        }
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        Updates the image feature shapes in the policy features dictionary if resizing is applied.
+
+        Args:
+            features: The policy features dictionary.
+
+        Returns:
+            The updated policy features dictionary with new image shapes.
+        """
+        if self.resize_size is None:
+            return features
+        for key in features[PipelineFeatureType.OBSERVATION]:
+            if "image" in key:
+                nb_channel = features[PipelineFeatureType.OBSERVATION][key].shape[0]
+                features[PipelineFeatureType.OBSERVATION][key] = PolicyFeature(
+                    type=features[PipelineFeatureType.OBSERVATION][key].type,
+                    shape=(nb_channel, *self.resize_size),
+                )
+        return features
+
+
+@dataclass
+@ProcessorStepRegistry.register("time_limit_processor")
+class TimeLimitProcessorStep(TruncatedProcessorStep):
+    """
+    Tracks episode steps and enforces a time limit by truncating the episode.
+
+    Attributes:
+        max_episode_steps: The maximum number of steps allowed per episode.
+        current_step: The current step count for the active episode.
+    """
+
+    max_episode_steps: int
+    current_step: int = 0
+
+    def truncated(self, truncated: bool) -> bool:
+        """
+        Increments the step counter and sets the truncated flag if the time limit is reached.
+
+        Args:
+            truncated: The incoming truncated flag.
+
+        Returns:
+            True if the episode step limit is reached, otherwise the incoming value.
+        """
+        self.current_step += 1
+        if self.current_step >= self.max_episode_steps:
+            truncated = True
+        # TODO (steven): missing an else truncated = False?
+        return truncated
+
+    def get_config(self) -> dict[str, Any]:
+        """
+        Returns the configuration of the step for serialization.
+
+        Returns:
+            A dictionary containing the `max_episode_steps`.
+        """
+        return {
+            "max_episode_steps": self.max_episode_steps,
+        }
+
+    def reset(self) -> None:
+        """Resets the step counter, typically called at the start of a new episode."""
+        self.current_step = 0
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@dataclass
+@ProcessorStepRegistry.register("gripper_penalty_processor")
+class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep):
+    """
+    Applies a penalty for inefficient gripper usage.
+
+    This step penalizes actions that attempt to close an already closed gripper or
+    open an already open one, based on position thresholds.
+
+    Attributes:
+        penalty: The negative reward value to apply.
+        max_gripper_pos: The maximum position value for the gripper, used for normalization.
+    """
+
+    penalty: float = -0.01
+    max_gripper_pos: float = 30.0
+
+    def complementary_data(self, complementary_data: dict) -> dict:
+        """
+        Calculates the gripper penalty and adds it to the complementary data.
+
+        Args:
+            complementary_data: The incoming complementary data, which should contain
+                                raw joint positions.
+
+        Returns:
+            A new complementary data dictionary with the `discrete_penalty` key added.
+        """
+        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)
+        if current_gripper_pos is None:
+            return complementary_data
+
+        # Gripper action is a PolicyAction at this stage
+        gripper_action = action[-1].item()
+        gripper_action_normalized = gripper_action / self.max_gripper_pos
+
+        # Normalize gripper state and action
+        gripper_state_normalized = current_gripper_pos / self.max_gripper_pos
+
+        # Calculate penalty boolean as in original
+        gripper_penalty_bool = (gripper_state_normalized < 0.5 and gripper_action_normalized > 0.5) or (
+            gripper_state_normalized > 0.75 and gripper_action_normalized < 0.5
+        )
+
+        gripper_penalty = self.penalty * int(gripper_penalty_bool)
+
+        # Create new complementary data with penalty info
+        new_complementary_data = dict(complementary_data)
+        new_complementary_data[DISCRETE_PENALTY_KEY] = gripper_penalty
+
+        return new_complementary_data
+
+    def get_config(self) -> dict[str, Any]:
+        """
+        Returns the configuration of the step for serialization.
+
+        Returns:
+            A dictionary containing the penalty value and max gripper position.
+        """
+        return {
+            "penalty": self.penalty,
+            "max_gripper_pos": self.max_gripper_pos,
+        }
+
+    def reset(self) -> None:
+        """Resets the processor's internal state."""
+        pass
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@dataclass
+@ProcessorStepRegistry.register("intervention_action_processor")
+class InterventionActionProcessorStep(ProcessorStep):
+    """
+    Handles human intervention, overriding policy actions and managing episode termination.
+
+    When an intervention is detected (via teleoperator events in the `info` dict),
+    this step replaces the policy's action with the human's teleoperated action.
+    It also processes signals to terminate the episode or flag success.
+
+    Attributes:
+        use_gripper: Whether to include the gripper in the teleoperated action.
+        terminate_on_success: If True, automatically sets the `done` flag when a
+                              `success` event is received.
+    """
+
+    use_gripper: bool = False
+    terminate_on_success: bool = True
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        """
+        Processes the transition to handle interventions.
+
+        Args:
+            transition: The incoming environment transition.
+
+        Returns:
+            The modified transition, potentially with an overridden action, updated
+            reward, and termination status.
+        """
+        action = transition.get(TransitionKey.ACTION)
+        if not isinstance(action, PolicyAction):
+            raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
+
+        # Get intervention signals from complementary data
+        info = transition.get(TransitionKey.INFO, {})
+        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
+        teleop_action = complementary_data.get(TELEOP_ACTION_KEY, {})
+        is_intervention = info.get(TeleopEvents.IS_INTERVENTION, False)
+        terminate_episode = info.get(TeleopEvents.TERMINATE_EPISODE, False)
+        success = info.get(TeleopEvents.SUCCESS, False)
+        rerecord_episode = info.get(TeleopEvents.RERECORD_EPISODE, False)
+
+        new_transition = transition.copy()
+
+        # Override action if intervention is active
+        if is_intervention and teleop_action is not None:
+            if isinstance(teleop_action, dict):
+                # Convert teleop_action dict to tensor format
+                action_list = [
+                    teleop_action.get("delta_x", 0.0),
+                    teleop_action.get("delta_y", 0.0),
+                    teleop_action.get("delta_z", 0.0),
+                ]
+                if self.use_gripper:
+                    action_list.append(teleop_action.get(GRIPPER_KEY, 1.0))
+            elif isinstance(teleop_action, np.ndarray):
+                action_list = teleop_action.tolist()
+            else:
+                action_list = teleop_action
+
+            teleop_action_tensor = torch.tensor(action_list, dtype=action.dtype, device=action.device)
+            new_transition[TransitionKey.ACTION] = teleop_action_tensor
+
+        # Handle episode termination
+        new_transition[TransitionKey.DONE] = bool(terminate_episode) or (
+            self.terminate_on_success and success
+        )
+        new_transition[TransitionKey.REWARD] = float(success)
+
+        # Update info with intervention metadata
+        info = new_transition.get(TransitionKey.INFO, {})
+        info[TeleopEvents.IS_INTERVENTION] = is_intervention
+        info[TeleopEvents.RERECORD_EPISODE] = rerecord_episode
+        info[TeleopEvents.SUCCESS] = success
+        new_transition[TransitionKey.INFO] = info
+
+        # Update complementary data with teleop action
+        complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
+        complementary_data[TELEOP_ACTION_KEY] = new_transition.get(TransitionKey.ACTION)
+        new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
+
+        return new_transition
+
+    def get_config(self) -> dict[str, Any]:
+        """
+        Returns the configuration of the step for serialization.
+
+        Returns:
+            A dictionary containing the step's configuration attributes.
+        """
+        return {
+            "use_gripper": self.use_gripper,
+            "terminate_on_success": self.terminate_on_success,
+        }
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@dataclass
+@ProcessorStepRegistry.register("reward_classifier_processor")
+class RewardClassifierProcessorStep(ProcessorStep):
+    """
+    Applies a pretrained reward classifier to image observations to predict success.
+
+    This step uses a model to determine if the current state is successful, updating
+    the reward and potentially terminating the episode.
+
+    Attributes:
+        pretrained_path: Path to the pretrained reward classifier model.
+        device: The device to run the classifier on.
+        success_threshold: The probability threshold to consider a prediction as successful.
+        success_reward: The reward value to assign on success.
+        terminate_on_success: If True, terminates the episode upon successful classification.
+        reward_classifier: The loaded classifier model instance.
+    """
+
+    pretrained_path: str | None = None
+    device: str = "cpu"
+    success_threshold: float = 0.5
+    success_reward: float = 1.0
+    terminate_on_success: bool = True
+
+    reward_classifier: Any = None
+
+    def __post_init__(self):
+        """Initializes the reward classifier model after the dataclass is created."""
+        if self.pretrained_path is not None:
+            from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
+
+            self.reward_classifier = Classifier.from_pretrained(self.pretrained_path)
+            self.reward_classifier.to(self.device)
+            self.reward_classifier.eval()
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        """
+        Processes a transition, applying the reward classifier to its image observations.
+
+        Args:
+            transition: The incoming environment transition.
+
+        Returns:
+            The modified transition with an updated reward and done flag based on the
+            classifier's prediction.
+        """
+        new_transition = transition.copy()
+        observation = new_transition.get(TransitionKey.OBSERVATION)
+        if observation is None or self.reward_classifier is None:
+            return new_transition
+
+        # Extract images from observation
+        images = {key: value for key, value in observation.items() if "image" in key}
+
+        if not images:
+            return new_transition
+
+        # Run reward classifier
+        start_time = time.perf_counter()
+        with torch.inference_mode():
+            success = self.reward_classifier.predict_reward(images, threshold=self.success_threshold)
+
+        classifier_frequency = 1 / (time.perf_counter() - start_time)
+
+        # Calculate reward and termination
+        reward = new_transition.get(TransitionKey.REWARD, 0.0)
+        terminated = new_transition.get(TransitionKey.DONE, False)
+
+        if math.isclose(success, 1, abs_tol=1e-2):
+            reward = self.success_reward
+            if self.terminate_on_success:
+                terminated = True
+
+        # Update transition
+        new_transition[TransitionKey.REWARD] = reward
+        new_transition[TransitionKey.DONE] = terminated
+
+        # Update info with classifier frequency
+        info = new_transition.get(TransitionKey.INFO, {})
+        info["reward_classifier_frequency"] = classifier_frequency
+        new_transition[TransitionKey.INFO] = info
+
+        return new_transition
+
+    def get_config(self) -> dict[str, Any]:
+        """
+        Returns the configuration of the step for serialization.
+
+        Returns:
+            A dictionary containing the step's configuration attributes.
+        """
+        return {
+            "device": self.device,
+            "success_threshold": self.success_threshold,
+            "success_reward": self.success_reward,
+            "terminate_on_success": self.terminate_on_success,
+        }
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features

src/lerobot/processor/joint_observations_processor.py

@@ -0,0 +1,211 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import OBS_STATE
+from lerobot.processor.pipeline import (
+    ObservationProcessorStep,
+    ProcessorStepRegistry,
+)
+from lerobot.robots import Robot
+
+
+@dataclass
+@ProcessorStepRegistry.register("joint_velocity_processor")
+class JointVelocityProcessorStep(ObservationProcessorStep):
+    """
+    Calculates and appends joint velocity information to the observation state.
+
+    This step computes the velocity of each joint by calculating the finite
+    difference between the current and the last observed joint positions. The
+    resulting velocity vector is then concatenated to the original state vector.
+
+    Attributes:
+        dt: The time step (delta time) in seconds between observations, used for
+            calculating velocity.
+        last_joint_positions: Stores the joint positions from the previous step
+                              to enable velocity calculation.
+    """
+
+    dt: float = 0.1
+
+    last_joint_positions: torch.Tensor | None = None
+
+    def observation(self, observation: dict) -> dict:
+        """
+        Computes joint velocities and adds them to the observation state.
+
+        Args:
+            observation: The input observation dictionary, expected to contain
+                         an `observation.state` key with joint positions.
+
+        Returns:
+            A new observation dictionary with the `observation.state` tensor
+            extended to include joint velocities.
+
+        Raises:
+            ValueError: If `observation.state` is not found in the observation.
+        """
+        # Get current joint positions (assuming they're in observation.state)
+        current_positions = observation.get(OBS_STATE)
+        if current_positions is None:
+            raise ValueError(f"{OBS_STATE} is not in observation")
+
+        # Initialize last joint positions if not already set
+        if self.last_joint_positions is None:
+            self.last_joint_positions = current_positions.clone()
+            joint_velocities = torch.zeros_like(current_positions)
+        else:
+            # Compute velocities
+            joint_velocities = (current_positions - self.last_joint_positions) / self.dt
+
+        self.last_joint_positions = current_positions.clone()
+
+        # Extend observation with velocities
+        extended_state = torch.cat([current_positions, joint_velocities], dim=-1)
+
+        # Create new observation dict
+        new_observation = dict(observation)
+        new_observation[OBS_STATE] = extended_state
+
+        return new_observation
+
+    def get_config(self) -> dict[str, Any]:
+        """
+        Returns the configuration of the step for serialization.
+
+        Returns:
+            A dictionary containing the time step `dt`.
+        """
+        return {
+            "dt": self.dt,
+        }
+
+    def reset(self) -> None:
+        """Resets the internal state, clearing the last known joint positions."""
+        self.last_joint_positions = None
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        Updates the `observation.state` feature to reflect the added velocities.
+
+        This method doubles the size of the first dimension of the `observation.state`
+        shape to account for the concatenation of position and velocity vectors.
+
+        Args:
+            features: The policy features dictionary.
+
+        Returns:
+            The updated policy features dictionary.
+        """
+        if OBS_STATE in features[PipelineFeatureType.OBSERVATION]:
+            original_feature = features[PipelineFeatureType.OBSERVATION][OBS_STATE]
+            # Double the shape to account for positions + velocities
+            new_shape = (original_feature.shape[0] * 2,) + original_feature.shape[1:]
+
+            features[PipelineFeatureType.OBSERVATION][OBS_STATE] = PolicyFeature(
+                type=original_feature.type, shape=new_shape
+            )
+        return features
+
+
+@dataclass
+@ProcessorStepRegistry.register("current_processor")
+class MotorCurrentProcessorStep(ObservationProcessorStep):
+    """
+    Reads motor currents from a robot and appends them to the observation state.
+
+    This step queries the robot's hardware interface to get the present current
+    for each motor and concatenates this information to the existing state vector.
+
+    Attributes:
+        robot: An instance of a `lerobot` Robot class that provides access to
+               the hardware bus.
+    """
+
+    robot: Robot | None = None
+
+    def observation(self, observation: dict) -> dict:
+        """
+        Fetches motor currents and adds them to the observation state.
+
+        Args:
+            observation: The input observation dictionary.
+
+        Returns:
+            A new observation dictionary with the `observation.state` tensor
+            extended to include motor currents.
+
+        Raises:
+            ValueError: If the `robot` attribute has not been set.
+        """
+        # Get current values from robot state
+        if self.robot is None:
+            raise ValueError("Robot is not set")
+
+        present_current_dict = self.robot.bus.sync_read("Present_Current")  # type: ignore[attr-defined]
+        motor_currents = torch.tensor(
+            [present_current_dict[name] for name in self.robot.bus.motors],  # type: ignore[attr-defined]
+            dtype=torch.float32,
+        ).unsqueeze(0)
+
+        current_state = observation.get(OBS_STATE)
+        if current_state is None:
+            return observation
+
+        extended_state = torch.cat([current_state, motor_currents], dim=-1)
+
+        # Create new observation dict
+        new_observation = dict(observation)
+        new_observation[OBS_STATE] = extended_state
+
+        return new_observation
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        Updates the `observation.state` feature to reflect the added motor currents.
+
+        This method increases the size of the first dimension of the `observation.state`
+        shape by the number of motors in the robot.
+
+        Args:
+            features: The policy features dictionary.
+
+        Returns:
+            The updated policy features dictionary.
+        """
+        if OBS_STATE in features[PipelineFeatureType.OBSERVATION] and self.robot is not None:
+            original_feature = features[PipelineFeatureType.OBSERVATION][OBS_STATE]
+            # Add motor current dimensions to the original state shape
+            num_motors = 0
+            if hasattr(self.robot, "bus") and hasattr(self.robot.bus, "motors"):  # type: ignore[attr-defined]
+                num_motors = len(self.robot.bus.motors)  # type: ignore[attr-defined]
+
+            if num_motors > 0:
+                new_shape = (original_feature.shape[0] + num_motors,) + original_feature.shape[1:]
+                features[PipelineFeatureType.OBSERVATION][OBS_STATE] = PolicyFeature(
+                    type=original_feature.type, shape=new_shape
+                )
+        return features

src/lerobot/processor/migrate_policy_normalization.py

@@ -0,0 +1,646 @@
+#!/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.
+
+"""
+A generic script to migrate LeRobot policies with built-in normalization layers to the new
+pipeline-based processor system.
+
+This script performs the following steps:
+1.  Loads a pretrained policy model and its configuration from a local path or the
+    Hugging Face Hub.
+2.  Scans the model's state dictionary to extract normalization statistics (e.g., mean,
+    std, min, max) for all features.
+3.  Creates two new processor pipelines:
+    - A preprocessor that normalizes inputs (observations) and outputs (actions).
+    - A postprocessor that unnormalizes outputs (actions) for inference.
+4.  Removes the original normalization layers from the model's state dictionary,
+    creating a "clean" model.
+5.  Saves the new clean model, the preprocessor, the postprocessor, and a generated
+    model card to a new directory.
+6.  Optionally pushes all the new artifacts to the Hugging Face Hub.
+
+Usage:
+    python src/lerobot/processor/migrate_policy_normalization.py \
+        --pretrained-path lerobot/act_aloha_sim_transfer_cube_human \
+        --push-to-hub \
+        --branch main
+
+Note: This script now uses the modern `make_pre_post_processors` and `make_policy_config`
+factory functions from `lerobot.policies.factory` to create processors and configurations,
+ensuring consistency with the current codebase.
+
+The script extracts normalization statistics from the old model's state_dict, creates clean
+processor pipelines using the factory functions, and saves a migrated model that is compatible
+with the new PolicyProcessorPipeline architecture.
+"""
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import torch
+from huggingface_hub import HfApi, hf_hub_download
+from safetensors.torch import load_file as load_safetensors
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors
+
+
+def extract_normalization_stats(state_dict: dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
+    """
+    Scans a model's state_dict to find and extract normalization statistics.
+
+    This function identifies keys corresponding to normalization layers (e.g., those
+    for mean, std, min, max) based on a set of predefined patterns and organizes
+    them into a nested dictionary.
+
+    Args:
+        state_dict: The state dictionary of a pretrained policy model.
+
+    Returns:
+        A nested dictionary where outer keys are feature names (e.g.,
+        'observation.state') and inner keys are statistic types ('mean', 'std'),
+        mapping to their corresponding tensor values.
+    """
+    stats = {}
+
+    # Define patterns to match and their prefixes to remove
+    normalization_patterns = [
+        "normalize_inputs.buffer_",
+        "unnormalize_outputs.buffer_",
+        "normalize_targets.buffer_",
+        "normalize.",  # Must come after normalize_* patterns
+        "unnormalize.",  # Must come after unnormalize_* patterns
+        "input_normalizer.",
+        "output_normalizer.",
+        "normalalize_inputs.",
+        "unnormalize_outputs.",
+        "normalize_targets.",
+        "unnormalize_targets.",
+    ]
+
+    # Process each key in state_dict
+    for key, tensor in state_dict.items():
+        # Try each pattern
+        for pattern in normalization_patterns:
+            if key.startswith(pattern):
+                # Extract the remaining part after the pattern
+                remaining = key[len(pattern) :]
+                parts = remaining.split(".")
+
+                # Need at least feature name and stat type
+                if len(parts) >= 2:
+                    # Last part is the stat type (mean, std, min, max, etc.)
+                    stat_type = parts[-1]
+                    # Everything else is the feature name
+                    feature_name = ".".join(parts[:-1]).replace("_", ".")
+
+                    # Add to stats
+                    if feature_name not in stats:
+                        stats[feature_name] = {}
+                    stats[feature_name][stat_type] = tensor.clone()
+
+                # Only process the first matching pattern
+                break
+
+    return stats
+
+
+def detect_features_and_norm_modes(
+    config: dict[str, Any], stats: dict[str, dict[str, torch.Tensor]]
+) -> tuple[dict[str, PolicyFeature], dict[FeatureType, NormalizationMode]]:
+    """
+    Infers policy features and normalization modes from the model config and stats.
+
+    This function first attempts to find feature definitions and normalization
+    mappings directly from the policy's configuration file. If this information is
+    not present, it infers it from the extracted normalization statistics, using
+    tensor shapes to determine feature shapes and the presence of specific stat
+    keys (e.g., 'mean'/'std' vs 'min'/'max') to determine the normalization mode.
+    It applies sensible defaults if inference is not possible.
+
+    Args:
+        config: The policy's configuration dictionary from `config.json`.
+        stats: The normalization statistics extracted from the model's state_dict.
+
+    Returns:
+        A tuple containing:
+        - A dictionary mapping feature names to `PolicyFeature` objects.
+        - A dictionary mapping `FeatureType` enums to `NormalizationMode` enums.
+    """
+    features = {}
+    norm_modes = {}
+
+    # First, check if there's a normalization_mapping in the config
+    if "normalization_mapping" in config:
+        print(f"Found normalization_mapping in config: {config['normalization_mapping']}")
+        # Extract normalization modes from config
+        for feature_type_str, mode_str in config["normalization_mapping"].items():
+            # Convert string to FeatureType enum
+            try:
+                if feature_type_str == "VISUAL":
+                    feature_type = FeatureType.VISUAL
+                elif feature_type_str == "STATE":
+                    feature_type = FeatureType.STATE
+                elif feature_type_str == "ACTION":
+                    feature_type = FeatureType.ACTION
+                else:
+                    print(f"Warning: Unknown feature type '{feature_type_str}', skipping")
+                    continue
+            except (AttributeError, ValueError):
+                print(f"Warning: Could not parse feature type '{feature_type_str}', skipping")
+                continue
+
+            # Convert string to NormalizationMode enum
+            try:
+                if mode_str == "MEAN_STD":
+                    mode = NormalizationMode.MEAN_STD
+                elif mode_str == "MIN_MAX":
+                    mode = NormalizationMode.MIN_MAX
+                elif mode_str == "IDENTITY":
+                    mode = NormalizationMode.IDENTITY
+                else:
+                    print(
+                        f"Warning: Unknown normalization mode '{mode_str}' for feature type '{feature_type_str}'"
+                    )
+                    continue
+            except (AttributeError, ValueError):
+                print(f"Warning: Could not parse normalization mode '{mode_str}', skipping")
+                continue
+
+            norm_modes[feature_type] = mode
+
+    # Try to extract from config
+    if "features" in config:
+        for key, feature_config in config["features"].items():
+            shape = feature_config.get("shape", feature_config.get("dim"))
+            shape = (shape,) if isinstance(shape, int) else tuple(shape)
+
+            # Determine feature type
+            if "image" in key or "visual" in key:
+                feature_type = FeatureType.VISUAL
+            elif "state" in key:
+                feature_type = FeatureType.STATE
+            elif "action" in key:
+                feature_type = FeatureType.ACTION
+            else:
+                feature_type = FeatureType.STATE  # Default
+
+            features[key] = PolicyFeature(feature_type, shape)
+
+    # If no features in config, infer from stats
+    if not features:
+        for key, stat_dict in stats.items():
+            # Get shape from any stat tensor
+            tensor = next(iter(stat_dict.values()))
+            shape = tuple(tensor.shape)
+
+            # Determine feature type based on key
+            if "image" in key or "visual" in key or "pixels" in key:
+                feature_type = FeatureType.VISUAL
+            elif "state" in key or "joint" in key or "position" in key:
+                feature_type = FeatureType.STATE
+            elif "action" in key:
+                feature_type = FeatureType.ACTION
+            else:
+                feature_type = FeatureType.STATE
+
+            features[key] = PolicyFeature(feature_type, shape)
+
+    # If normalization modes weren't in config, determine based on available stats
+    if not norm_modes:
+        for key, stat_dict in stats.items():
+            if key in features:
+                if "mean" in stat_dict and "std" in stat_dict:
+                    feature_type = features[key].type
+                    if feature_type not in norm_modes:
+                        norm_modes[feature_type] = NormalizationMode.MEAN_STD
+                elif "min" in stat_dict and "max" in stat_dict:
+                    feature_type = features[key].type
+                    if feature_type not in norm_modes:
+                        norm_modes[feature_type] = NormalizationMode.MIN_MAX
+
+    # Default normalization modes if not detected
+    if FeatureType.VISUAL not in norm_modes:
+        norm_modes[FeatureType.VISUAL] = NormalizationMode.MEAN_STD
+    if FeatureType.STATE not in norm_modes:
+        norm_modes[FeatureType.STATE] = NormalizationMode.MIN_MAX
+    if FeatureType.ACTION not in norm_modes:
+        norm_modes[FeatureType.ACTION] = NormalizationMode.MEAN_STD
+
+    return features, norm_modes
+
+
+def remove_normalization_layers(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """
+    Creates a new state_dict with all normalization-related layers removed.
+
+    This function filters the original state dictionary, excluding any keys that
+    match a set of predefined patterns associated with normalization modules.
+
+    Args:
+        state_dict: The original model state dictionary.
+
+    Returns:
+        A new state dictionary containing only the core model weights, without
+        any normalization parameters.
+    """
+    new_state_dict = {}
+
+    # Patterns to remove
+    remove_patterns = [
+        "normalize_inputs.",
+        "unnormalize_outputs.",
+        "normalize_targets.",  # Added pattern for target normalization
+        "normalize.",
+        "unnormalize.",
+        "input_normalizer.",
+        "output_normalizer.",
+        "normalizer.",
+    ]
+
+    for key, tensor in state_dict.items():
+        should_remove = any(pattern in key for pattern in remove_patterns)
+        if not should_remove:
+            new_state_dict[key] = tensor
+
+    return new_state_dict
+
+
+def clean_state_dict(
+    state_dict: dict[str, torch.Tensor], remove_str: str = "._orig_mod"
+) -> dict[str, torch.Tensor]:
+    """
+    Remove a substring (e.g. '._orig_mod') from all keys in a state dict.
+
+    Args:
+        state_dict (dict): The original state dict.
+        remove_str (str): The substring to remove from the keys.
+
+    Returns:
+        dict: A new state dict with cleaned keys.
+    """
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        new_k = k.replace(remove_str, "")
+        new_state_dict[new_k] = v
+    return new_state_dict
+
+
+def convert_features_to_policy_features(features_dict: dict[str, dict]) -> dict[str, PolicyFeature]:
+    """
+    Converts a feature dictionary from the old config format to the new `PolicyFeature` format.
+
+    Args:
+        features_dict: The feature dictionary in the old format, where values are
+                       simple dictionaries (e.g., `{"shape": [7]}`).
+
+    Returns:
+        A dictionary mapping feature names to `PolicyFeature` dataclass objects.
+    """
+    converted_features = {}
+
+    for key, feature_dict in features_dict.items():
+        # Determine feature type based on key
+        if "image" in key or "visual" in key:
+            feature_type = FeatureType.VISUAL
+        elif "state" in key:
+            feature_type = FeatureType.STATE
+        elif "action" in key:
+            feature_type = FeatureType.ACTION
+        else:
+            feature_type = FeatureType.STATE
+
+        # Get shape from feature dict
+        shape = feature_dict.get("shape", feature_dict.get("dim"))
+        shape = (shape,) if isinstance(shape, int) else tuple(shape) if shape is not None else ()
+
+        converted_features[key] = PolicyFeature(feature_type, shape)
+
+    return converted_features
+
+
+def load_model_from_hub(
+    repo_id: str, revision: str | None = None
+) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, Any]]:
+    """
+    Downloads and loads a model's state_dict and configs from the Hugging Face Hub.
+
+    Args:
+        repo_id: The repository ID on the Hub (e.g., 'lerobot/aloha').
+        revision: The specific git revision (branch, tag, or commit hash) to use.
+
+    Returns:
+        A tuple containing the model's state dictionary, the policy configuration,
+        and the training configuration.
+    """
+    # Download files.
+    safetensors_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", revision=revision)
+
+    config_path = hf_hub_download(repo_id=repo_id, filename="config.json", revision=revision)
+    train_config_path = hf_hub_download(repo_id=repo_id, filename="train_config.json", revision=revision)
+
+    # Load state_dict
+    state_dict = load_safetensors(safetensors_path)
+
+    # Load config
+    with open(config_path) as f:
+        config = json.load(f)
+
+    with open(train_config_path) as f:
+        train_config = json.load(f)
+
+    return state_dict, config, train_config
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Migrate policy models with normalization layers to new pipeline system"
+    )
+    parser.add_argument(
+        "--pretrained-path",
+        type=str,
+        required=True,
+        help="Path to pretrained model (hub repo or local directory)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=None,
+        help="Output directory for migrated model (default: same as pretrained-path)",
+    )
+    parser.add_argument("--push-to-hub", action="store_true", help="Push migrated model to hub")
+    parser.add_argument(
+        "--hub-repo-id",
+        type=str,
+        default=None,
+        help="Hub repository ID for pushing (default: same as pretrained-path)",
+    )
+    parser.add_argument("--revision", type=str, default=None, help="Revision of the model to load")
+    parser.add_argument("--private", action="store_true", help="Make the hub repository private")
+    parser.add_argument(
+        "--branch",
+        type=str,
+        default=None,
+        help="Git branch to use when pushing to hub. If specified, a PR will be created automatically (default: push directly to main)",
+    )
+
+    args = parser.parse_args()
+
+    # Load model and config
+    print(f"Loading model from {args.pretrained_path}...")
+    if os.path.isdir(args.pretrained_path):
+        # Local directory
+        state_dict = load_safetensors(os.path.join(args.pretrained_path, "model.safetensors"))
+        with open(os.path.join(args.pretrained_path, "config.json")) as f:
+            config = json.load(f)
+        with open(os.path.join(args.pretrained_path, "train_config.json")) as f:
+            train_config = json.load(f)
+    else:
+        # Hub repository
+        state_dict, config, train_config = load_model_from_hub(args.pretrained_path, args.revision)
+
+    # Extract normalization statistics
+    print("Extracting normalization statistics...")
+    stats = extract_normalization_stats(state_dict)
+
+    print(f"Found normalization statistics for: {list(stats.keys())}")
+
+    # Detect input features and normalization modes
+    print("Detecting features and normalization modes...")
+    features, norm_map = detect_features_and_norm_modes(config, stats)
+
+    print(f"Detected features: {list(features.keys())}")
+    print(f"Normalization modes: {norm_map}")
+
+    # Remove normalization layers from state_dict
+    print("Removing normalization layers from model...")
+    new_state_dict = remove_normalization_layers(state_dict)
+    new_state_dict = clean_state_dict(new_state_dict, remove_str="._orig_mod")
+
+    removed_keys = set(state_dict.keys()) - set(new_state_dict.keys())
+    if removed_keys:
+        print(f"Removed {len(removed_keys)} normalization layer keys")
+
+    # Determine output path
+    if args.output_dir:
+        output_dir = Path(args.output_dir)
+    else:
+        if os.path.isdir(args.pretrained_path):
+            output_dir = Path(args.pretrained_path).parent / f"{Path(args.pretrained_path).name}_migrated"
+        else:
+            output_dir = Path(f"./{args.pretrained_path.replace('/', '_')}_migrated")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Extract policy type from config
+    if "type" not in config:
+        raise ValueError("Policy type not found in config.json. The config must contain a 'type' field.")
+
+    policy_type = config["type"]
+    print(f"Detected policy type: {policy_type}")
+
+    # Clean up config - remove fields that shouldn't be passed to config constructor
+    cleaned_config = dict(config)
+
+    # Remove fields that are not part of the config class constructors
+    fields_to_remove = ["normalization_mapping", "type"]
+    for field in fields_to_remove:
+        if field in cleaned_config:
+            print(f"Removing '{field}' field from config")
+            del cleaned_config[field]
+
+    # Convert input_features and output_features to PolicyFeature objects if they exist
+    if "input_features" in cleaned_config:
+        cleaned_config["input_features"] = convert_features_to_policy_features(
+            cleaned_config["input_features"]
+        )
+    if "output_features" in cleaned_config:
+        cleaned_config["output_features"] = convert_features_to_policy_features(
+            cleaned_config["output_features"]
+        )
+
+    # Add normalization mapping to config
+    cleaned_config["normalization_mapping"] = norm_map
+
+    # Create policy configuration using the factory
+    print(f"Creating {policy_type} policy configuration...")
+    policy_config = make_policy_config(policy_type, **cleaned_config)
+
+    # Create policy instance using the factory
+    print(f"Instantiating {policy_type} policy...")
+    policy_class = get_policy_class(policy_type)
+    policy = policy_class(policy_config)
+
+    # Load the cleaned state dict
+    policy.load_state_dict(new_state_dict, strict=True)
+    print("Successfully loaded cleaned state dict into policy model")
+
+    # Create preprocessor and postprocessor using the factory
+    print("Creating preprocessor and postprocessor using make_pre_post_processors...")
+    preprocessor, postprocessor = make_pre_post_processors(policy_cfg=policy_config, dataset_stats=stats)
+
+    # Determine hub repo ID if pushing to hub
+    hub_repo_id = None
+    if args.push_to_hub:
+        if args.hub_repo_id:
+            hub_repo_id = args.hub_repo_id
+        else:
+            if not os.path.isdir(args.pretrained_path):
+                # Use same repo with "_migrated" suffix
+                hub_repo_id = f"{args.pretrained_path}_migrated"
+            else:
+                raise ValueError("--hub-repo-id must be specified when pushing local model to hub")
+
+    # Save all components to local directory first
+    print(f"Saving preprocessor to {output_dir}...")
+    preprocessor.save_pretrained(output_dir)
+
+    print(f"Saving postprocessor to {output_dir}...")
+    postprocessor.save_pretrained(output_dir)
+
+    print(f"Saving model to {output_dir}...")
+    policy.save_pretrained(output_dir)
+
+    # Generate and save model card
+    print("Generating model card...")
+    # Get metadata from original config
+    dataset_repo_id = train_config.get("repo_id", "unknown")
+    license = config.get("license", "apache-2.0")
+
+    tags = config.get("tags", ["robotics", "lerobot", policy_type]) or ["robotics", "lerobot", policy_type]
+    tags = set(tags).union({"robotics", "lerobot", policy_type})
+    tags = list(tags)
+
+    # Generate model card
+    card = policy.generate_model_card(
+        dataset_repo_id=dataset_repo_id, model_type=policy_type, license=license, tags=tags
+    )
+
+    # Save model card locally
+    card.save(str(output_dir / "README.md"))
+    print(f"Model card saved to {output_dir / 'README.md'}")
+    # Push all files to hub in a single operation if requested
+    if args.push_to_hub and hub_repo_id:
+        api = HfApi()
+
+        # Determine if we should create a PR (automatically if branch is specified)
+        create_pr = args.branch is not None
+        target_location = f"branch '{args.branch}'" if args.branch else "main branch"
+
+        print(f"Pushing all migrated files to {hub_repo_id} on {target_location}...")
+
+        # Upload all files in a single commit with automatic PR creation if branch specified
+        commit_message = "Migrate policy to PolicyProcessorPipeline system"
+        commit_description = None
+
+        if create_pr:
+            # Separate commit description for PR body
+            commit_description = """🤖 **Automated Policy Migration to PolicyProcessorPipeline**
+
+This PR migrates your model to the new LeRobot policy format using the modern PolicyProcessorPipeline architecture.
+
+## What Changed
+
+### ✨ **New Architecture - PolicyProcessorPipeline**
+Your model now uses external PolicyProcessorPipeline components for data processing instead of built-in normalization layers. This provides:
+- **Modularity**: Separate preprocessing and postprocessing pipelines
+- **Flexibility**: Easy to swap, configure, and debug processing steps
+- **Compatibility**: Works with the latest LeRobot ecosystem
+
+### 🔧 **Normalization Extraction**
+We've extracted normalization statistics from your model's state_dict and removed the built-in normalization layers:
+- **Extracted patterns**: `normalize_inputs.*`, `unnormalize_outputs.*`, `normalize.*`, `unnormalize.*`, `input_normalizer.*`, `output_normalizer.*`
+- **Statistics preserved**: Mean, std, min, max values for all features
+- **Clean model**: State dict now contains only core model weights
+
+### 📦 **Files Added**
+- **preprocessor_config.json**: Configuration for input preprocessing pipeline
+- **postprocessor_config.json**: Configuration for output postprocessing pipeline
+- **model.safetensors**: Clean model weights without normalization layers
+- **config.json**: Updated model configuration
+- **train_config.json**: Training configuration
+- **README.md**: Updated model card with migration information
+
+### 🚀 **Benefits**
+- **Backward Compatible**: Your model behavior remains identical
+- **Future Ready**: Compatible with latest LeRobot features and updates
+- **Debuggable**: Easy to inspect and modify processing steps
+- **Portable**: Processors can be shared and reused across models
+
+### 💻 **Usage**
+```python
+# Load your migrated model
+from lerobot.policies import get_policy_class
+from lerobot.processor import PolicyProcessorPipeline
+
+# The preprocessor and postprocessor are now external
+preprocessor = PolicyProcessorPipeline.from_pretrained("your-model-repo", config_filename="preprocessor_config.json")
+postprocessor = PolicyProcessorPipeline.from_pretrained("your-model-repo", config_filename="postprocessor_config.json")
+policy = get_policy_class("your-policy-type").from_pretrained("your-model-repo")
+
+# Process data through the pipeline
+processed_batch = preprocessor(raw_batch)
+action = policy(processed_batch)
+final_action = postprocessor(action)
+```
+
+*Generated automatically by the LeRobot policy migration script*"""
+
+        upload_kwargs = {
+            "repo_id": hub_repo_id,
+            "folder_path": output_dir,
+            "repo_type": "model",
+            "commit_message": commit_message,
+            "revision": args.branch,
+            "create_pr": create_pr,
+            "allow_patterns": ["*.json", "*.safetensors", "*.md"],
+            "ignore_patterns": ["*.tmp", "*.log"],
+        }
+
+        # Add commit_description for PR body if creating PR
+        if create_pr and commit_description:
+            upload_kwargs["commit_description"] = commit_description
+
+        api.upload_folder(**upload_kwargs)
+
+        if create_pr:
+            print("All files pushed and pull request created successfully!")
+        else:
+            print("All files pushed to main branch successfully!")
+
+    print("\nMigration complete!")
+    print(f"Migrated model saved to: {output_dir}")
+    if args.push_to_hub and hub_repo_id:
+        if args.branch:
+            print(
+                f"Successfully pushed all files to branch '{args.branch}' and created PR on https://huggingface.co/{hub_repo_id}"
+            )
+        else:
+            print(f"Successfully pushed to https://huggingface.co/{hub_repo_id}")
+        if args.branch:
+            print(f"\nView the branch at: https://huggingface.co/{hub_repo_id}/tree/{args.branch}")
+            print(
+                f"View the PR at: https://huggingface.co/{hub_repo_id}/discussions (look for the most recent PR)"
+            )
+        else:
+            print(f"\nView the changes at: https://huggingface.co/{hub_repo_id}")
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/processor/normalize_processor.py

@@ -1,67 +1,353 @@
+#!/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 not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from __future__ import annotations
 
-from collections.abc import Mapping
+from copy import deepcopy
 from dataclasses import dataclass, field
 from typing import Any
 
-import numpy as np
 import torch
 from torch import Tensor
 
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
+
+from .converters import from_tensor_to_numpy, to_tensor
+from .core import EnvTransition, PolicyAction, TransitionKey
+from .pipeline import PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry
 
 
-def _convert_stats_to_tensors(stats: dict[str, dict[str, Any]]) -> dict[str, dict[str, Tensor]]:
-    """Convert numpy arrays and other types to torch tensors."""
-    tensor_stats: dict[str, dict[str, Tensor]] = {}
-    for key, sub in stats.items():
-        tensor_stats[key] = {}
-        for stat_name, value in sub.items():
-            if isinstance(value, np.ndarray):
-                tensor_val = torch.from_numpy(value.astype(np.float32))
-            elif isinstance(value, torch.Tensor):
-                tensor_val = value.to(dtype=torch.float32)
-            elif isinstance(value, (int, float, list, tuple)):
-                tensor_val = torch.tensor(value, dtype=torch.float32)
-            else:
-                raise TypeError(f"Unsupported type for stats['{key}']['{stat_name}']: {type(value)}")
-            tensor_stats[key][stat_name] = tensor_val
-    return tensor_stats
+@dataclass
+class _NormalizationMixin:
+    """
+    A mixin class providing core functionality for normalization and unnormalization.
+
+    This class manages normalization statistics (`stats`), converts them to tensors for
+    efficient computation, handles device placement, and implements the logic for
+    applying normalization transformations (mean/std and min/max). It is designed to
+    be inherited by concrete `ProcessorStep` implementations and should not be used
+    directly.
+
+    **Stats Override Preservation:**
+    When stats are explicitly provided during construction (e.g., via overrides in
+    `DataProcessorPipeline.from_pretrained()`), they are preserved even when
+    `load_state_dict()` is called. This allows users to override normalization
+    statistics from saved models while keeping the rest of the model state intact.
+
+    Examples:
+        ```python
+        # Common use case: Override with dataset stats
+        from lerobot.datasets import LeRobotDataset
+
+        dataset = LeRobotDataset("my_dataset")
+        pipeline = DataProcessorPipeline.from_pretrained(
+            "model_path", overrides={"normalizer_processor": {"stats": dataset.meta.stats}}
+        )
+        # dataset.meta.stats will be used, not the stats from the saved model
+
+        # Custom stats override
+        custom_stats = {"action": {"mean": [0.0], "std": [1.0]}}
+        pipeline = DataProcessorPipeline.from_pretrained(
+            "model_path", overrides={"normalizer_processor": {"stats": custom_stats}}
+        )
+        ```
+
+    Attributes:
+        features: A dictionary mapping feature names to `PolicyFeature` objects, defining
+            the data structure to be processed.
+        norm_map: A dictionary mapping `FeatureType` to `NormalizationMode`, specifying
+            which normalization method to use for each type of feature.
+        stats: A dictionary containing the normalization statistics (e.g., mean, std,
+            min, max) for each feature.
+        device: The PyTorch device on which to store and perform tensor operations.
+        eps: A small epsilon value to prevent division by zero in normalization
+            calculations.
+        normalize_observation_keys: An optional set of keys to selectively apply
+            normalization to specific observation features.
+        _tensor_stats: An internal dictionary holding the normalization statistics as
+            PyTorch tensors.
+        _stats_explicitly_provided: Internal flag tracking whether stats were explicitly
+            provided during construction (used for override preservation).
+    """
+
+    features: dict[str, PolicyFeature]
+    norm_map: dict[FeatureType, NormalizationMode]
+    stats: dict[str, dict[str, Any]] | None = None
+    device: torch.device | str | None = None
+    dtype: torch.dtype | None = None
+    eps: float = 1e-8
+    normalize_observation_keys: set[str] | None = None
+
+    _tensor_stats: dict[str, dict[str, Tensor]] = field(default_factory=dict, init=False, repr=False)
+    _stats_explicitly_provided: bool = field(default=False, init=False, repr=False)
+
+    def __post_init__(self):
+        """
+        Initializes the mixin after dataclass construction.
+
+        This method handles the robust deserialization of `features` and `norm_map`
+        from JSON-compatible formats (where enums become strings and tuples become
+        lists) and converts the provided `stats` dictionary into a dictionary of
+        tensors (`_tensor_stats`) on the specified device.
+        """
+        # Track if stats were explicitly provided (not None and not empty)
+        self._stats_explicitly_provided = self.stats is not None and bool(self.stats)
+        # Robust JSON deserialization handling (guard empty maps).
+        if self.features:
+            first_val = next(iter(self.features.values()))
+            if isinstance(first_val, dict):
+                reconstructed = {}
+                for key, ft_dict in self.features.items():
+                    reconstructed[key] = PolicyFeature(
+                        type=FeatureType(ft_dict["type"]), shape=tuple(ft_dict["shape"])
+                    )
+                self.features = reconstructed
+
+        if self.norm_map:
+            # if keys are strings (JSON), rebuild enum map
+            if all(isinstance(k, str) for k in self.norm_map.keys()):
+                reconstructed = {}
+                for ft_type_str, norm_mode_str in self.norm_map.items():
+                    reconstructed[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
+                self.norm_map = reconstructed
+
+        # Convert stats to tensors and move to the target device once during initialization.
+        self.stats = self.stats or {}
+        if self.dtype is None:
+            self.dtype = torch.float32
+        self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=self.dtype)
+
+    def to(
+        self, device: torch.device | str | None = None, dtype: torch.dtype | None = None
+    ) -> _NormalizationMixin:
+        """
+        Moves the processor's normalization stats to the specified device.
+
+        Args:
+            device: The target PyTorch device.
+
+        Returns:
+            The instance of the class, allowing for method chaining.
+        """
+        if device is not None:
+            self.device = device
+        if dtype is not None:
+            self.dtype = dtype
+        self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=self.dtype)
+        return self
+
+    def state_dict(self) -> dict[str, Tensor]:
+        """
+        Returns the normalization statistics as a flat state dictionary.
+
+        All tensors are moved to the CPU before being returned, which is standard practice
+        for saving state dictionaries.
+
+        Returns:
+            A flat dictionary mapping from `'feature_name.stat_name'` to the
+            corresponding statistics tensor on the CPU.
+        """
+        flat: dict[str, Tensor] = {}
+        for key, sub in self._tensor_stats.items():
+            for stat_name, tensor in sub.items():
+                flat[f"{key}.{stat_name}"] = tensor.cpu()  # Always save to CPU
+        return flat
+
+    def load_state_dict(self, state: dict[str, Tensor]) -> None:
+        """
+        Loads normalization statistics from a state dictionary.
+
+        The loaded tensors are moved to the processor's configured device.
+
+        **Stats Override Preservation:**
+        If stats were explicitly provided during construction (e.g., via overrides in
+        `DataProcessorPipeline.from_pretrained()`), they are preserved and the state
+        dictionary is ignored. This allows users to override normalization statistics
+        while still loading the rest of the model state.
+
+        This behavior is crucial for scenarios where users want to adapt a pretrained
+        model to a new dataset with different statistics without retraining the entire
+        model.
+
+        Args:
+            state: A flat state dictionary with keys in the format
+                   `'feature_name.stat_name'`.
+
+        Note:
+            When stats are preserved due to explicit provision, only the tensor
+            representation is updated to ensure consistency with the current device
+            and dtype settings.
+        """
+        # If stats were explicitly provided during construction, preserve them
+        if self._stats_explicitly_provided and self.stats is not None:
+            # Don't load from state_dict, keep the explicitly provided stats
+            # But ensure _tensor_stats is properly initialized
+            self._tensor_stats = to_tensor(self.stats, device=self.device, dtype=self.dtype)  # type: ignore[assignment]
+            return
+
+        # Normal behavior: load stats from state_dict
+        self._tensor_stats.clear()
+        for flat_key, tensor in state.items():
+            key, stat_name = flat_key.rsplit(".", 1)
+            # Load to the processor's configured device.
+            self._tensor_stats.setdefault(key, {})[stat_name] = tensor.to(
+                dtype=torch.float32, device=self.device
+            )
+
+        # Reconstruct the original stats dict from tensor stats for compatibility with to() method
+        # and other functions that rely on self.stats
+        self.stats = {}
+        for key, tensor_dict in self._tensor_stats.items():
+            self.stats[key] = {}
+            for stat_name, tensor in tensor_dict.items():
+                # Convert tensor back to python/numpy format
+                self.stats[key][stat_name] = from_tensor_to_numpy(tensor)
+
+    def get_config(self) -> dict[str, Any]:
+        """
+        Returns a serializable dictionary of the processor's configuration.
+
+        This method is used when saving the processor to disk, ensuring that its
+        configuration can be reconstructed later.
+
+        Returns:
+            A JSON-serializable dictionary containing the configuration.
+        """
+        config = {
+            "eps": self.eps,
+            "features": {
+                key: {"type": ft.type.value, "shape": ft.shape} for key, ft in self.features.items()
+            },
+            "norm_map": {ft_type.value: norm_mode.value for ft_type, norm_mode in self.norm_map.items()},
+        }
+        if self.normalize_observation_keys is not None:
+            config["normalize_observation_keys"] = sorted(self.normalize_observation_keys)
+        return config
+
+    def _normalize_observation(self, observation: dict[str, Any], inverse: bool) -> dict[str, Tensor]:
+        """
+        Applies (un)normalization to all relevant features in an observation dictionary.
+
+        Args:
+            observation: The observation dictionary to process.
+            inverse: If `True`, applies unnormalization; otherwise, applies normalization.
+
+        Returns:
+            A new observation dictionary with the transformed tensor values.
+        """
+        new_observation = dict(observation)
+        for key, feature in self.features.items():
+            if self.normalize_observation_keys is not None and key not in self.normalize_observation_keys:
+                continue
+            if feature.type != FeatureType.ACTION and key in new_observation:
+                # Convert to tensor but preserve original dtype for adaptation logic
+                tensor = torch.as_tensor(new_observation[key])
+                new_observation[key] = self._apply_transform(tensor, key, feature.type, inverse=inverse)
+        return new_observation
+
+    def _normalize_action(self, action: Tensor, inverse: bool) -> Tensor:
+        # Convert to tensor but preserve original dtype for adaptation logic
+        """
+        Applies (un)normalization to an action tensor.
+
+        Args:
+            action: The action tensor to process.
+            inverse: If `True`, applies unnormalization; otherwise, applies normalization.
+
+        Returns:
+            The transformed action tensor.
+        """
+        processed_action = self._apply_transform(action, "action", FeatureType.ACTION, inverse=inverse)
+        return processed_action
+
+    def _apply_transform(
+        self, tensor: Tensor, key: str, feature_type: FeatureType, *, inverse: bool = False
+    ) -> Tensor:
+        """
+        Core logic to apply a normalization or unnormalization transformation to a tensor.
+
+        This method selects the appropriate normalization mode (e.g., mean/std, min/max)
+        based on the feature type and applies the corresponding mathematical operation.
+
+        Args:
+            tensor: The input tensor to transform.
+            key: The feature key corresponding to the tensor.
+            feature_type: The `FeatureType` of the tensor.
+            inverse: If `True`, applies the inverse transformation (unnormalization).
+
+        Returns:
+            The transformed tensor.
+
+        Raises:
+            ValueError: If an unsupported normalization mode is encountered.
+        """
+        norm_mode = self.norm_map.get(feature_type, NormalizationMode.IDENTITY)
+        if norm_mode == NormalizationMode.IDENTITY or key not in self._tensor_stats:
+            return tensor
+
+        if norm_mode not in (NormalizationMode.MEAN_STD, NormalizationMode.MIN_MAX):
+            raise ValueError(f"Unsupported normalization mode: {norm_mode}")
+
+        # For Accelerate compatibility: Ensure stats are on the same device and dtype as the input tensor
+        if self._tensor_stats and key in self._tensor_stats:
+            first_stat = next(iter(self._tensor_stats[key].values()))
+            if first_stat.device != tensor.device or first_stat.dtype != tensor.dtype:
+                self.to(device=tensor.device, dtype=tensor.dtype)
+
+        stats = self._tensor_stats[key]
+
+        if norm_mode == NormalizationMode.MEAN_STD and "mean" in stats and "std" in stats:
+            mean, std = stats["mean"], stats["std"]
+            # Avoid division by zero by adding a small epsilon.
+            denom = std + self.eps
+            if inverse:
+                return tensor * std + mean
+            return (tensor - mean) / denom
+
+        if norm_mode == NormalizationMode.MIN_MAX and "min" in stats and "max" in stats:
+            min_val, max_val = stats["min"], stats["max"]
+            denom = max_val - min_val
+            # When min_val == max_val, substitute the denominator with a small epsilon
+            # to prevent division by zero. This consistently maps an input equal to
+            # min_val to -1, ensuring a stable transformation.
+            denom = torch.where(
+                denom == 0, torch.tensor(self.eps, device=tensor.device, dtype=tensor.dtype), denom
+            )
+            if inverse:
+                # Map from [-1, 1] back to [min, max]
+                return (tensor + 1) / 2 * denom + min_val
+            # Map from [min, max] to [-1, 1]
+            return 2 * (tensor - min_val) / denom - 1
+
+        # If necessary stats are missing, return input unchanged.
+        return tensor
 
 
 @dataclass
 @ProcessorStepRegistry.register(name="normalizer_processor")
-class NormalizerProcessor:
-    """Normalizes observations and actions in a single processor step.
-
-    This processor handles normalization of both observation and action tensors
-    using either mean/std normalization or min/max scaling to a [-1, 1] range.
-
-    For each tensor key in the stats dictionary, the processor will:
-    - Use mean/std normalization if those statistics are provided: (x - mean) / std
-    - Use min/max scaling if those statistics are provided: 2 * (x - min) / (max - min) - 1
-
-    The processor can be configured to normalize only specific keys by setting
-    the normalize_keys parameter.
+class NormalizerProcessorStep(_NormalizationMixin, ProcessorStep):
     """
+    A processor step that applies normalization to observations and actions in a transition.
 
-    # Features and normalisation map are mandatory to match the design of normalize.py
-    features: dict[str, PolicyFeature]
-    norm_map: dict[FeatureType, NormalizationMode]
-
-    # Pre-computed statistics coming from dataset.meta.stats for instance.
-    stats: dict[str, dict[str, Any]] | None = None
-
-    # Explicit subset of keys to normalise. If ``None`` every key (except
-    # "action") found in ``stats`` will be normalised. Using a ``set`` makes
-    # membership checks O(1).
-    normalize_keys: set[str] | None = None
-
-    eps: float = 1e-8
-
-    _tensor_stats: dict[str, dict[str, Tensor]] = field(default_factory=dict, init=False, repr=False)
+    This class uses the logic from `_NormalizationMixin` to perform forward normalization
+    (e.g., scaling data to have zero mean and unit variance, or to the range [-1, 1]).
+    It is typically used in the pre-processing pipeline before feeding data to a policy.
+    """
 
     @classmethod
     def from_lerobot_dataset(
@@ -70,158 +356,73 @@ class NormalizerProcessor:
         features: dict[str, PolicyFeature],
         norm_map: dict[FeatureType, NormalizationMode],
         *,
-        normalize_keys: set[str] | None = None,
+        normalize_observation_keys: set[str] | None = None,
         eps: float = 1e-8,
-    ) -> NormalizerProcessor:
-        """Factory helper that pulls statistics from a :class:`LeRobotDataset`.
-
-        The features and norm_map parameters are mandatory to match the design
-        pattern used in normalize.py.
+        device: torch.device | str | None = None,
+    ) -> NormalizerProcessorStep:
         """
+        Creates a `NormalizerProcessorStep` instance using statistics from a `LeRobotDataset`.
 
+        Args:
+            dataset: The dataset from which to extract normalization statistics.
+            features: The feature definition for the processor.
+            norm_map: The mapping from feature types to normalization modes.
+            normalize_observation_keys: An optional set of observation keys to normalize.
+            eps: A small epsilon value for numerical stability.
+            device: The target device for the processor.
+
+        Returns:
+            A new instance of `NormalizerProcessorStep`.
+        """
         return cls(
             features=features,
             norm_map=norm_map,
             stats=dataset.meta.stats,
-            normalize_keys=normalize_keys,
+            normalize_observation_keys=normalize_observation_keys,
             eps=eps,
+            device=device,
         )
 
-    def __post_init__(self):
-        # Handle deserialization from JSON config
-        if self.features and isinstance(list(self.features.values())[0], dict):
-            # Features came from JSON - need to reconstruct PolicyFeature objects
-            reconstructed_features = {}
-            for key, ft_dict in self.features.items():
-                reconstructed_features[key] = PolicyFeature(
-                    type=FeatureType(ft_dict["type"]), shape=tuple(ft_dict["shape"])
-                )
-            self.features = reconstructed_features
-
-        if self.norm_map and isinstance(list(self.norm_map.keys())[0], str):
-            # norm_map came from JSON - need to reconstruct enum keys and values
-            reconstructed_norm_map = {}
-            for ft_type_str, norm_mode_str in self.norm_map.items():
-                reconstructed_norm_map[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
-            self.norm_map = reconstructed_norm_map
-
-        # Convert statistics once so we avoid repeated numpy→Tensor conversions
-        # during runtime.
-        self.stats = self.stats or {}
-        self._tensor_stats = _convert_stats_to_tensors(self.stats)
-
-        # Ensure *normalize_keys* is a set for fast look-ups and compare by
-        # value later when returning the configuration.
-        if self.normalize_keys is not None and not isinstance(self.normalize_keys, set):
-            self.normalize_keys = set(self.normalize_keys)
-
-    def _normalize_obs(self, observation):
-        if observation is None:
-            return None
-
-        # Decide which keys should be normalised for this call.
-        if self.normalize_keys is not None:
-            keys_to_norm = self.normalize_keys
-        else:
-            # Use feature map to skip action keys.
-            keys_to_norm = {k for k, ft in self.features.items() if ft.type is not FeatureType.ACTION}
-
-        processed = dict(observation)
-        for key in keys_to_norm:
-            if key not in processed or key not in self._tensor_stats:
-                continue
-
-            orig_val = processed[key]
-            tensor = (
-                orig_val.to(dtype=torch.float32)
-                if isinstance(orig_val, torch.Tensor)
-                else torch.as_tensor(orig_val, dtype=torch.float32)
-            )
-            stats = {k: v.to(tensor.device) for k, v in self._tensor_stats[key].items()}
-
-            if "mean" in stats and "std" in stats:
-                mean, std = stats["mean"], stats["std"]
-                processed[key] = (tensor - mean) / (std + self.eps)
-            elif "min" in stats and "max" in stats:
-                min_val, max_val = stats["min"], stats["max"]
-                processed[key] = 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
-        return processed
-
-    def _normalize_action(self, action):
-        if action is None or "action" not in self._tensor_stats:
-            return action
-
-        tensor = (
-            action.to(dtype=torch.float32)
-            if isinstance(action, torch.Tensor)
-            else torch.as_tensor(action, dtype=torch.float32)
-        )
-        stats = {k: v.to(tensor.device) for k, v in self._tensor_stats["action"].items()}
-        if "mean" in stats and "std" in stats:
-            mean, std = stats["mean"], stats["std"]
-            return (tensor - mean) / (std + self.eps)
-        if "min" in stats and "max" in stats:
-            min_val, max_val = stats["min"], stats["max"]
-            return 2 * (tensor - min_val) / (max_val - min_val + self.eps) - 1
-        raise ValueError("Action stats must contain either ('mean','std') or ('min','max')")
-
     def __call__(self, transition: EnvTransition) -> EnvTransition:
-        observation = self._normalize_obs(transition.get(TransitionKey.OBSERVATION))
-        action = self._normalize_action(transition.get(TransitionKey.ACTION))
-
-        # Create a new transition with normalized values
         new_transition = transition.copy()
-        new_transition[TransitionKey.OBSERVATION] = observation
-        new_transition[TransitionKey.ACTION] = action
+
+        # Handle observation normalization.
+        observation = new_transition.get(TransitionKey.OBSERVATION)
+        if observation is not None:
+            new_transition[TransitionKey.OBSERVATION] = self._normalize_observation(
+                observation, inverse=False
+            )
+
+        # Handle action normalization.
+        action = new_transition.get(TransitionKey.ACTION)
+
+        if action is None:
+            return new_transition
+
+        if not isinstance(action, PolicyAction):
+            raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
+
+        new_transition[TransitionKey.ACTION] = self._normalize_action(action, inverse=False)
+
         return new_transition
 
-    def get_config(self) -> dict[str, Any]:
-        config = {
-            "eps": self.eps,
-            "features": {
-                key: {"type": ft.type.value, "shape": ft.shape} for key, ft in self.features.items()
-            },
-            "norm_map": {ft_type.value: norm_mode.value for ft_type, norm_mode in self.norm_map.items()},
-        }
-        if self.normalize_keys is not None:
-            # Serialise as a list for YAML / JSON friendliness
-            config["normalize_keys"] = sorted(self.normalize_keys)
-        return config
-
-    def state_dict(self) -> dict[str, Tensor]:
-        flat = {}
-        for key, sub in self._tensor_stats.items():
-            for stat_name, tensor in sub.items():
-                flat[f"{key}.{stat_name}"] = tensor
-        return flat
-
-    def load_state_dict(self, state: Mapping[str, Tensor]) -> None:
-        self._tensor_stats.clear()
-        for flat_key, tensor in state.items():
-            key, stat_name = flat_key.rsplit(".", 1)
-            self._tensor_stats.setdefault(key, {})[stat_name] = tensor
-
-    def reset(self):
-        pass
-
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         return features
 
 
 @dataclass
 @ProcessorStepRegistry.register(name="unnormalizer_processor")
-class UnnormalizerProcessor:
-    """Inverse normalisation for observations and actions.
-
-    Exactly mirrors :class:`NormalizerProcessor` but applies the inverse
-    transform.
+class UnnormalizerProcessorStep(_NormalizationMixin, ProcessorStep):
     """
+    A processor step that applies unnormalization to observations and actions.
 
-    features: dict[str, PolicyFeature]
-    norm_map: dict[FeatureType, NormalizationMode]
-    stats: dict[str, dict[str, Any]] | None = None
-
-    _tensor_stats: dict[str, dict[str, Tensor]] = field(default_factory=dict, init=False, repr=False)
+    This class inverts the normalization process, scaling data back to its original
+    range. It is typically used in the post-processing pipeline to convert a policy's
+    normalized action output into a format that can be executed by a robot or
+    environment.
+    """
 
     @classmethod
     def from_lerobot_dataset(
@@ -229,103 +430,72 @@ class UnnormalizerProcessor:
         dataset: LeRobotDataset,
         features: dict[str, PolicyFeature],
         norm_map: dict[FeatureType, NormalizationMode],
-    ) -> UnnormalizerProcessor:
-        return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats)
+        *,
+        device: torch.device | str | None = None,
+    ) -> UnnormalizerProcessorStep:
+        """
+        Creates an `UnnormalizerProcessorStep` using statistics from a `LeRobotDataset`.
 
-    def __post_init__(self):
-        # Handle deserialization from JSON config
-        if self.features and isinstance(list(self.features.values())[0], dict):
-            # Features came from JSON - need to reconstruct PolicyFeature objects
-            reconstructed_features = {}
-            for key, ft_dict in self.features.items():
-                reconstructed_features[key] = PolicyFeature(
-                    type=FeatureType(ft_dict["type"]), shape=tuple(ft_dict["shape"])
-                )
-            self.features = reconstructed_features
+        Args:
+            dataset: The dataset from which to extract normalization statistics.
+            features: The feature definition for the processor.
+            norm_map: The mapping from feature types to normalization modes.
+            device: The target device for the processor.
 
-        if self.norm_map and isinstance(list(self.norm_map.keys())[0], str):
-            # norm_map came from JSON - need to reconstruct enum keys and values
-            reconstructed_norm_map = {}
-            for ft_type_str, norm_mode_str in self.norm_map.items():
-                reconstructed_norm_map[FeatureType(ft_type_str)] = NormalizationMode(norm_mode_str)
-            self.norm_map = reconstructed_norm_map
-
-        self.stats = self.stats or {}
-        self._tensor_stats = _convert_stats_to_tensors(self.stats)
-
-    def _unnormalize_obs(self, observation):
-        if observation is None:
-            return None
-        keys = [k for k, ft in self.features.items() if ft.type is not FeatureType.ACTION]
-        processed = dict(observation)
-        for key in keys:
-            if key not in processed or key not in self._tensor_stats:
-                continue
-            orig_val = processed[key]
-            tensor = (
-                orig_val.to(dtype=torch.float32)
-                if isinstance(orig_val, torch.Tensor)
-                else torch.as_tensor(orig_val, dtype=torch.float32)
-            )
-            stats = {k: v.to(tensor.device) for k, v in self._tensor_stats[key].items()}
-            if "mean" in stats and "std" in stats:
-                mean, std = stats["mean"], stats["std"]
-                processed[key] = tensor * std + mean
-            elif "min" in stats and "max" in stats:
-                min_val, max_val = stats["min"], stats["max"]
-                processed[key] = (tensor + 1) / 2 * (max_val - min_val) + min_val
-        return processed
-
-    def _unnormalize_action(self, action):
-        if action is None or "action" not in self._tensor_stats:
-            return action
-        tensor = (
-            action.to(dtype=torch.float32)
-            if isinstance(action, torch.Tensor)
-            else torch.as_tensor(action, dtype=torch.float32)
-        )
-        stats = {k: v.to(tensor.device) for k, v in self._tensor_stats["action"].items()}
-        if "mean" in stats and "std" in stats:
-            mean, std = stats["mean"], stats["std"]
-            return tensor * std + mean
-        if "min" in stats and "max" in stats:
-            min_val, max_val = stats["min"], stats["max"]
-            return (tensor + 1) / 2 * (max_val - min_val) + min_val
-        raise ValueError("Action stats must contain either ('mean','std') or ('min','max')")
+        Returns:
+            A new instance of `UnnormalizerProcessorStep`.
+        """
+        return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats, device=device)
 
     def __call__(self, transition: EnvTransition) -> EnvTransition:
-        observation = self._unnormalize_obs(transition.get(TransitionKey.OBSERVATION))
-        action = self._unnormalize_action(transition.get(TransitionKey.ACTION))
-
-        # Create a new transition with unnormalized values
         new_transition = transition.copy()
-        new_transition[TransitionKey.OBSERVATION] = observation
-        new_transition[TransitionKey.ACTION] = action
+
+        # Handle observation unnormalization.
+        observation = new_transition.get(TransitionKey.OBSERVATION)
+        if observation is not None:
+            new_transition[TransitionKey.OBSERVATION] = self._normalize_observation(observation, inverse=True)
+
+        # Handle action unnormalization.
+        action = new_transition.get(TransitionKey.ACTION)
+
+        if action is None:
+            return new_transition
+        if not isinstance(action, PolicyAction):
+            raise ValueError(f"Action should be a PolicyAction type got {type(action)}")
+
+        new_transition[TransitionKey.ACTION] = self._normalize_action(action, inverse=True)
+
         return new_transition
 
-    def get_config(self) -> dict[str, Any]:
-        return {
-            "features": {
-                key: {"type": ft.type.value, "shape": ft.shape} for key, ft in self.features.items()
-            },
-            "norm_map": {ft_type.value: norm_mode.value for ft_type, norm_mode in self.norm_map.items()},
-        }
-
-    def state_dict(self) -> dict[str, Tensor]:
-        flat = {}
-        for key, sub in self._tensor_stats.items():
-            for stat_name, tensor in sub.items():
-                flat[f"{key}.{stat_name}"] = tensor
-        return flat
-
-    def load_state_dict(self, state: Mapping[str, Tensor]) -> None:
-        self._tensor_stats.clear()
-        for flat_key, tensor in state.items():
-            key, stat_name = flat_key.rsplit(".", 1)
-            self._tensor_stats.setdefault(key, {})[stat_name] = tensor
-
-    def reset(self):
-        pass
-
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         return features
+
+
+def hotswap_stats(
+    policy_processor: PolicyProcessorPipeline, stats: dict[str, dict[str, Any]]
+) -> PolicyProcessorPipeline:
+    """
+    Replaces normalization statistics in an existing `PolicyProcessorPipeline` instance.
+
+    This function creates a deep copy of the provided pipeline and updates the
+    statistics of any `NormalizerProcessorStep` or `UnnormalizerProcessorStep` it
+    contains. This is useful for adapting a trained policy to a new environment or
+    dataset with different data distributions without having to reconstruct the entire
+    pipeline.
+
+    Args:
+        policy_processor: The policy processor pipeline to modify.
+        stats: The new dictionary of normalization statistics to apply.
+
+    Returns:
+        A new `PolicyProcessorPipeline` instance with the updated statistics.
+    """
+    rp = deepcopy(policy_processor)
+    for step in rp.steps:
+        if isinstance(step, _NormalizationMixin):
+            step.stats = stats
+            # Re-initialize tensor_stats on the correct device.
+            step._tensor_stats = to_tensor(stats, device=step.device, dtype=step.dtype)  # type: ignore[assignment]
+    return rp

src/lerobot/processor/observation_processor.py

@@ -20,32 +20,54 @@ import numpy as np
 import torch
 from torch import Tensor
 
-from lerobot.configs.types import PolicyFeature
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
 from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
-from lerobot.processor.pipeline import ObservationProcessor, ProcessorStepRegistry
+
+from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
 
 
 @dataclass
 @ProcessorStepRegistry.register(name="observation_processor")
-class VanillaObservationProcessor(ObservationProcessor):
+class VanillaObservationProcessorStep(ObservationProcessorStep):
     """
-    Processes environment observations into the LeRobot format by handling both images and states.
+    Processes standard Gymnasium observations into the LeRobot format.
 
-    Image processing:
-        - Converts channel-last (H, W, C) images to channel-first (C, H, W)
-        - Normalizes uint8 images ([0, 255]) to float32 ([0, 1])
-        - Adds a batch dimension if missing
-        - Supports single images and image dictionaries
+    This step handles both image and state data from a typical observation dictionary,
+    preparing it for use in a LeRobot policy.
 
-    State processing:
-        - Maps 'environment_state' to observation.environment_state
-        - Maps 'agent_pos' to observation.state
-        - Converts numpy arrays to tensors
-        - Adds a batch dimension if missing
+    **Image Processing:**
+    -   Converts channel-last (H, W, C), `uint8` images to channel-first (C, H, W),
+        `float32` tensors.
+    -   Normalizes pixel values from the [0, 255] range to [0, 1].
+    -   Adds a batch dimension if one is not already present.
+    -   Recognizes a single image under the key `"pixels"` and maps it to
+        `"observation.image"`.
+    -   Recognizes a dictionary of images under the key `"pixels"` and maps them
+        to `"observation.images.{camera_name}"`.
+
+    **State Processing:**
+    -   Maps the `"environment_state"` key to `"observation.environment_state"`.
+    -   Maps the `"agent_pos"` key to `"observation.state"`.
+    -   Converts NumPy arrays to PyTorch tensors.
+    -   Adds a batch dimension if one is not already present.
     """
 
     def _process_single_image(self, img: np.ndarray) -> Tensor:
-        """Process a single image array."""
+        """
+        Processes a single NumPy image array into a channel-first, normalized tensor.
+
+        Args:
+            img: A NumPy array representing the image, expected to be in channel-last
+                 (H, W, C) format with a `uint8` dtype.
+
+        Returns:
+            A `float32` PyTorch tensor in channel-first (B, C, H, W) format, with
+            pixel values normalized to the [0, 1] range.
+
+        Raises:
+            ValueError: If the input image does not appear to be in channel-last
+                        format or is not of `uint8` dtype.
+        """
         # Convert to tensor
         img_tensor = torch.from_numpy(img)
 
@@ -106,19 +128,32 @@ class VanillaObservationProcessor(ObservationProcessor):
     def observation(self, observation):
         return self._process_observation(observation)
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        """Transforms feature keys to a standardized contract.
-
-        This method handles several renaming patterns:
-        - Exact matches (e.g., 'pixels' -> 'OBS_IMAGE').
-        - Prefixed exact matches (e.g., 'observation.pixels' -> 'OBS_IMAGE').
-        - Prefix matches (e.g., 'pixels.cam1' -> 'OBS_IMAGES.cam1').
-        - Prefixed prefix matches (e.g., 'observation.pixels.cam1' -> 'OBS_IMAGES.cam1').
-        - environment_state -> OBS_ENV_STATE,
-        - agent_pos -> OBS_STATE,
-        - observation.environment_state -> OBS_ENV_STATE,
-        - observation.agent_pos -> OBS_STATE
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         """
+        Transforms feature keys from the Gym standard to the LeRobot standard.
+
+        This method standardizes the feature dictionary by renaming keys according
+        to LeRobot's conventions, ensuring that policies can be constructed correctly.
+        It handles various raw key formats, including those with an "observation." prefix.
+
+        **Renaming Rules:**
+        - `pixels` or `observation.pixels` -> `observation.image`
+        - `pixels.{cam}` or `observation.pixels.{cam}` -> `observation.images.{cam}`
+        - `environment_state` or `observation.environment_state` -> `observation.environment_state`
+        - `agent_pos` or `observation.agent_pos` -> `observation.state`
+
+        Args:
+            features: The policy features dictionary with Gym-style keys.
+
+        Returns:
+            The policy features dictionary with standardized LeRobot keys.
+        """
+        # Build a new features mapping keyed by the same FeatureType buckets
+        # We assume callers already placed features in the correct FeatureType.
+        new_features: dict[PipelineFeatureType, dict[str, PolicyFeature]] = {ft: {} for ft in features.keys()}
+
         exact_pairs = {
             "pixels": OBS_IMAGE,
             "environment_state": OBS_ENV_STATE,
@@ -129,29 +164,43 @@ class VanillaObservationProcessor(ObservationProcessor):
             "pixels.": f"{OBS_IMAGES}.",
         }
 
-        for key in list(features.keys()):
-            matched_prefix = False
-            for old_prefix, new_prefix in prefix_pairs.items():
-                prefixed_old = f"observation.{old_prefix}"
-                if key.startswith(prefixed_old):
-                    suffix = key[len(prefixed_old) :]
-                    features[f"{new_prefix}{suffix}"] = features.pop(key)
-                    matched_prefix = True
-                    break
+        # Iterate over all incoming feature buckets and normalize/move each entry
+        for src_ft, bucket in features.items():
+            for key, feat in list(bucket.items()):
+                handled = False
 
-                if key.startswith(old_prefix):
-                    suffix = key[len(old_prefix) :]
-                    features[f"{new_prefix}{suffix}"] = features.pop(key)
-                    matched_prefix = True
-                    break
-
-            if matched_prefix:
-                continue
-
-            for old, new in exact_pairs.items():
-                if key == old or key == f"observation.{old}":
-                    if key in features:
-                        features[new] = features.pop(key)
+                # Prefix-based rules (e.g. pixels.cam1 -> OBS_IMAGES.cam1)
+                for old_prefix, new_prefix in prefix_pairs.items():
+                    prefixed_old = f"observation.{old_prefix}"
+                    if key.startswith(prefixed_old):
+                        suffix = key[len(prefixed_old) :]
+                        new_key = f"{new_prefix}{suffix}"
+                        new_features[src_ft][new_key] = feat
+                        handled = True
                         break
 
-        return features
+                    if key.startswith(old_prefix):
+                        suffix = key[len(old_prefix) :]
+                        new_key = f"{new_prefix}{suffix}"
+                        new_features[src_ft][new_key] = feat
+                        handled = True
+                        break
+
+                if handled:
+                    continue
+
+                # Exact-name rules (pixels, environment_state, agent_pos)
+                for old, new in exact_pairs.items():
+                    if key == old or key == f"observation.{old}":
+                        new_key = new
+                        new_features[src_ft][new_key] = feat
+                        handled = True
+                        break
+
+                if handled:
+                    continue
+
+                # Default: keep key in the same source FeatureType bucket
+                new_features[src_ft][key] = feat
+
+        return new_features

src/lerobot/processor/policy_robot_bridge.py

@@ -0,0 +1,52 @@
+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

src/lerobot/processor/rename_processor.py

@@ -13,20 +13,30 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from copy import deepcopy
 from dataclasses import dataclass, field
 from typing import Any
 
-from lerobot.configs.types import PolicyFeature
-from lerobot.processor.pipeline import (
-    ObservationProcessor,
-    ProcessorStepRegistry,
-)
+from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+
+from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
 
 
 @dataclass
-@ProcessorStepRegistry.register(name="rename_processor")
-class RenameProcessor(ObservationProcessor):
-    """Rename processor that renames keys in the observation."""
+@ProcessorStepRegistry.register(name="rename_observations_processor")
+class RenameObservationsProcessorStep(ObservationProcessorStep):
+    """
+    A processor step that renames keys in an observation dictionary.
+
+    This step is useful for creating a standardized data interface by mapping keys
+    from an environment's format to the format expected by a LeRobot policy or
+    other downstream components.
+
+    Attributes:
+        rename_map: A dictionary mapping from old key names to new key names.
+                    Keys present in an observation that are not in this map will
+                    be kept with their original names.
+    """
 
     rename_map: dict[str, str] = field(default_factory=dict)
 
@@ -43,9 +53,41 @@ class RenameProcessor(ObservationProcessor):
     def get_config(self) -> dict[str, Any]:
         return {"rename_map": self.rename_map}
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         """Transforms:
         - Each key in the observation that appears in `rename_map` is renamed to its value.
         - Keys not in `rename_map` remain unchanged.
         """
-        return {self.rename_map.get(k, k): v for k, v in features.items()}
+        new_features: dict[PipelineFeatureType, dict[str, PolicyFeature]] = features.copy()
+        new_features[PipelineFeatureType.OBSERVATION] = {
+            self.rename_map.get(k, k): v for k, v in features[PipelineFeatureType.OBSERVATION].items()
+        }
+        return new_features
+
+
+def rename_stats(stats: dict[str, dict[str, Any]], rename_map: dict[str, str]) -> dict[str, dict[str, Any]]:
+    """
+    Renames the top-level keys in a statistics dictionary using a provided mapping.
+
+    This is a helper function typically used to keep normalization statistics
+    consistent with renamed observation or action features. It performs a defensive
+    deep copy to avoid modifying the original `stats` dictionary.
+
+    Args:
+        stats: A nested dictionary of statistics, where top-level keys are
+               feature names (e.g., `{"observation.state": {"mean": 0.5}}`).
+        rename_map: A dictionary mapping old feature names to new feature names.
+
+    Returns:
+        A new statistics dictionary with its top-level keys renamed. Returns an
+        empty dictionary if the input `stats` is empty.
+    """
+    if not stats:
+        return {}
+    renamed: dict[str, dict[str, Any]] = {}
+    for old_key, sub_stats in stats.items():
+        new_key = rename_map.get(old_key, old_key)
+        renamed[new_key] = deepcopy(sub_stats) if sub_stats is not None else {}
+    return renamed

src/lerobot/processor/tokenizer_processor.py

@@ -0,0 +1,270 @@
+#!/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.
+
+"""
+This script defines a processor for tokenizing natural language instructions from an environment transition.
+
+It uses a tokenizer from the Hugging Face `transformers` library to convert task descriptions (text) into
+token IDs and attention masks, which are then added to the observation dictionary.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
+from lerobot.utils.import_utils import _transformers_available
+
+from .core import EnvTransition, TransitionKey
+from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
+
+# Conditional import for type checking and lazy loading
+if TYPE_CHECKING or _transformers_available:
+    from transformers import AutoTokenizer
+else:
+    AutoTokenizer = None
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="tokenizer_processor")
+class TokenizerProcessorStep(ObservationProcessorStep):
+    """
+    Processor step to tokenize a natural language task description.
+
+    This step extracts a task string from the `complementary_data` of an `EnvTransition`,
+    tokenizes it using a Hugging Face `transformers` tokenizer, and adds the resulting
+    token IDs and attention mask to the `observation` dictionary.
+
+    Requires the `transformers` library to be installed.
+
+    Attributes:
+        tokenizer_name: The name of a pretrained tokenizer from the Hugging Face Hub (e.g., "bert-base-uncased").
+        tokenizer: A pre-initialized tokenizer object. If provided, `tokenizer_name` is ignored.
+        max_length: The maximum length to pad or truncate sequences to.
+        task_key: The key in `complementary_data` where the task string is stored.
+        padding_side: The side to pad on ('left' or 'right').
+        padding: The padding strategy ('max_length', 'longest', etc.).
+        truncation: Whether to truncate sequences longer than `max_length`.
+        input_tokenizer: The internal tokenizer instance, loaded during initialization.
+    """
+
+    tokenizer_name: str | None = None
+    tokenizer: Any | None = None  # Use `Any` for compatibility without a hard dependency
+    max_length: int = 512
+    task_key: str = "task"
+    padding_side: str = "right"
+    padding: str = "max_length"
+    truncation: bool = True
+
+    # Internal tokenizer instance (not part of the config)
+    input_tokenizer: Any = field(default=None, init=False, repr=False)
+
+    def __post_init__(self):
+        """
+        Initializes the tokenizer after the dataclass is created.
+
+        It checks for the availability of the `transformers` library and loads the tokenizer
+        either from a provided object or by name from the Hugging Face Hub.
+
+        Raises:
+            ImportError: If the `transformers` library is not installed.
+            ValueError: If neither `tokenizer` nor `tokenizer_name` is provided.
+        """
+        if not _transformers_available:
+            raise ImportError(
+                "The 'transformers' library is not installed. "
+                "Please install it with `pip install 'lerobot[transformers-dep]'` to use TokenizerProcessorStep."
+            )
+
+        if self.tokenizer is not None:
+            # Use provided tokenizer object directly
+            self.input_tokenizer = self.tokenizer
+        elif self.tokenizer_name is not None:
+            if AutoTokenizer is None:
+                raise ImportError("AutoTokenizer is not available")
+            self.input_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
+        else:
+            raise ValueError(
+                "Either 'tokenizer' or 'tokenizer_name' must be provided. "
+                "Pass a tokenizer object directly or a tokenizer name to auto-load."
+            )
+
+    def get_task(self, transition: EnvTransition) -> list[str] | None:
+        """
+        Extracts the task description(s) from the transition's complementary data.
+
+        Args:
+            transition: The environment transition.
+
+        Returns:
+            A list of task strings, or None if the task key is not found or the value is None.
+        """
+        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
+        if complementary_data is None:
+            raise ValueError("Complementary data is None so no task can be extracted from it")
+
+        task = complementary_data[self.task_key]
+        if task is None:
+            raise ValueError("Task extracted from Complementary data is None")
+
+        # Standardize to a list of strings for the tokenizer
+        if isinstance(task, str):
+            return [task]
+        elif isinstance(task, list) and all(isinstance(t, str) for t in task):
+            return task
+
+        return None
+
+    def observation(self, observation: dict[str, Any]) -> dict[str, Any]:
+        """
+        Tokenizes the task description and adds it to the observation dictionary.
+
+        This method retrieves the task, tokenizes it, moves the resulting tensors to the
+        same device as other data in the transition, and updates the observation.
+
+        Args:
+            observation: The original observation dictionary.
+
+        Returns:
+            The updated observation dictionary including token IDs and an attention mask.
+        """
+        task = self.get_task(self.transition)
+        if task is None:
+            raise ValueError("Task cannot be None")
+
+        # Tokenize the task (this will create CPU tensors)
+        tokenized_prompt = self._tokenize_text(task)
+
+        # Detect the device from existing tensors in the transition to ensure consistency
+        target_device = self._detect_device(self.transition)
+
+        # Move new tokenized tensors to the detected device
+        if target_device is not None:
+            tokenized_prompt = {
+                k: v.to(target_device) if isinstance(v, torch.Tensor) else v
+                for k, v in tokenized_prompt.items()
+            }
+
+        # Create a new observation dict to avoid modifying the original in place
+        new_observation = dict(observation)
+
+        # Add tokenized data to the observation
+        new_observation[OBS_LANGUAGE_TOKENS] = tokenized_prompt["input_ids"]
+        new_observation[OBS_LANGUAGE_ATTENTION_MASK] = tokenized_prompt["attention_mask"].to(dtype=torch.bool)
+
+        return new_observation
+
+    def _detect_device(self, transition: EnvTransition) -> torch.device | None:
+        """
+        Detects the torch.device from existing tensors in the transition.
+
+        It checks tensors in the observation dictionary first, then the action tensor.
+
+        Args:
+            transition: The environment transition.
+
+        Returns:
+            The detected `torch.device`, or None if no tensors are found.
+        """
+        # Check observation tensors first (most likely place to find tensors)
+        observation = transition.get(TransitionKey.OBSERVATION)
+        if observation:
+            for value in observation.values():
+                if isinstance(value, torch.Tensor):
+                    return value.device
+
+        # Fallback to checking the action tensor
+        action = transition.get(TransitionKey.ACTION)
+        if isinstance(action, torch.Tensor):
+            return action.device
+
+        return None  # No tensors found, default will be CPU
+
+    def _tokenize_text(self, text: str | list[str]) -> dict[str, torch.Tensor]:
+        """
+        A wrapper around the tokenizer call.
+
+        Args:
+            text: A string or list of strings to tokenize.
+
+        Returns:
+            A dictionary containing tokenized 'input_ids' and 'attention_mask' as PyTorch tensors.
+        """
+        return self.input_tokenizer(
+            text,
+            max_length=self.max_length,
+            truncation=self.truncation,
+            padding=self.padding,
+            padding_side=self.padding_side,
+            return_tensors="pt",
+        )
+
+    def get_config(self) -> dict[str, Any]:
+        """
+        Returns the serializable configuration of the processor.
+
+        Note: The tokenizer object itself is not serialized. If the processor was initialized
+        with a tokenizer name, that name will be included in the config.
+
+        Returns:
+            A dictionary with the processor's configuration parameters.
+        """
+        config = {
+            "max_length": self.max_length,
+            "task_key": self.task_key,
+            "padding_side": self.padding_side,
+            "padding": self.padding,
+            "truncation": self.truncation,
+        }
+
+        # Only save tokenizer_name if it was used to create the tokenizer
+        if self.tokenizer_name is not None and self.tokenizer is None:
+            config["tokenizer_name"] = self.tokenizer_name
+
+        return config
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """
+        Adds feature definitions for the language tokens and attention mask.
+
+        This updates the policy features dictionary to include the new data added to the
+        observation, ensuring downstream components are aware of their shape and type.
+
+        Args:
+            features: The dictionary of existing policy features.
+
+        Returns:
+            The updated dictionary of policy features.
+        """
+        # Add a feature for the token IDs if it doesn't already exist
+        if OBS_LANGUAGE_TOKENS not in features[PipelineFeatureType.OBSERVATION]:
+            features[PipelineFeatureType.OBSERVATION][OBS_LANGUAGE_TOKENS] = PolicyFeature(
+                type=FeatureType.LANGUAGE, shape=(self.max_length,)
+            )
+
+        # Add a feature for the attention mask if it doesn't already exist
+        if OBS_LANGUAGE_ATTENTION_MASK not in features[PipelineFeatureType.OBSERVATION]:
+            features[PipelineFeatureType.OBSERVATION][OBS_LANGUAGE_ATTENTION_MASK] = PolicyFeature(
+                type=FeatureType.LANGUAGE, shape=(self.max_length,)
+            )
+
+        return features

src/lerobot/record.py

@@ -21,11 +21,12 @@ Example:
 lerobot-record \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58760431541 \
-    --robot.cameras="{laptop: {type: opencv, camera_index: 0, width: 640, height: 480}}" \
+    --robot.cameras="{laptop: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
     --robot.id=black \
-    --dataset.repo_id=aliberts/record-test \
+    --dataset.repo_id=<my_username>/<my_dataset_name> \
     --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 \
@@ -59,9 +60,10 @@ lerobot-record \
 
 import logging
 import time
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 from pathlib import Path
 from pprint import pformat
+from typing import Any
 
 from lerobot.cameras import (  # noqa: F401
     CameraConfig,  # noqa: F401
@@ -72,10 +74,20 @@ from lerobot.configs import parser
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.datasets.image_writer import safe_stop_image_writer
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
+from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.datasets.utils import build_dataset_frame, combine_feature_dicts
 from lerobot.datasets.video_utils import VideoEncodingManager
-from lerobot.policies.factory import make_policy
+from lerobot.policies.factory import make_policy, make_pre_post_processors
 from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.processor import (
+    PolicyAction,
+    PolicyProcessorPipeline,
+    RobotAction,
+    RobotObservation,
+    RobotProcessorPipeline,
+    make_default_processors,
+)
+from lerobot.processor.rename_processor import rename_stats
 from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
@@ -149,6 +161,8 @@ class DatasetRecordConfig:
     # Number of episodes to record before batch encoding videos
     # Set to 1 for immediate encoding (default behavior), or higher for batched encoding
     video_encoding_batch_size: int = 1
+    # Rename map for the observation to override the image and state keys
+    rename_map: dict[str, str] = field(default_factory=dict)
 
     def __post_init__(self):
         if self.single_task is None:
@@ -187,14 +201,55 @@ class RecordConfig:
         return ["policy"]
 
 
+""" --------------- record_loop() data flow --------------------------
+       [ Robot ]
+           V
+     [ robot.get_observation() ] ---> raw_obs
+           V
+     [ robot_observation_processor ] ---> processed_obs
+           V
+     .-----( ACTION LOGIC )------------------.
+     V                                       V
+     [ From Teleoperator ]                   [ From Policy ]
+     |                                       |
+     |  [teleop.get_action] -> raw_action    |   [predict_action]
+     |          |                            |          |
+     |          V                            |          V
+     | [teleop_action_processor]             |          |
+     |          |                            |          |
+     '---> processed_teleop_action           '---> processed_policy_action
+     |                                       |
+     '-------------------------.-------------'
+                               V
+                  [ robot_action_processor ] --> robot_action_to_send
+                               V
+                    [ robot.send_action() ] -- (Robot Executes)
+                               V
+                    ( Save to Dataset )
+                               V
+                  ( Rerun Log / Loop Wait )
+"""
+
+
 @safe_stop_image_writer
 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
+    robot_observation_processor: RobotProcessorPipeline[
+        RobotObservation, RobotObservation
+    ],  # runs after robot
     dataset: LeRobotDataset | None = None,
     teleop: Teleoperator | list[Teleoperator] | None = None,
     policy: PreTrainedPolicy | None = None,
+    preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]] | None = None,
+    postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] | None = None,
     control_time_s: int | None = None,
     single_task: str | None = None,
     display_data: bool = False,
@@ -226,9 +281,11 @@ def record_loop(
                 "For multi-teleop, the list must contain exactly one KeyboardTeleop and one arm teleoperator. Currently only supported for LeKiwi robot."
             )
 
-    # if policy is given it needs cleaning up
-    if policy is not None:
+    # Reset policy and processor if they are provided
+    if policy is not None and preprocessor is not None and postprocessor is not None:
         policy.reset()
+        preprocessor.reset()
+        postprocessor.reset()
 
     timestamp = 0
     start_episode_t = time.perf_counter()
@@ -239,32 +296,46 @@ def record_loop(
             events["exit_early"] = False
             break
 
-        observation = robot.get_observation()
+        # Get robot observation
+        obs = robot.get_observation()
+
+        # Applies a pipeline to the raw robot observation, default is IdentityProcessor
+        obs_processed = robot_observation_processor(obs)
 
         if policy is not None or dataset is not None:
-            observation_frame = build_dataset_frame(dataset.features, observation, prefix="observation")
+            observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix="observation")
 
-        if policy is not None:
+        # Get action from either policy or teleop
+        if policy is not None and preprocessor is not None and postprocessor is not None:
             action_values = predict_action(
-                observation_frame,
-                policy,
-                get_safe_torch_device(policy.config.device),
-                policy.config.use_amp,
+                observation=observation_frame,
+                policy=policy,
+                device=get_safe_torch_device(policy.config.device),
+                preprocessor=preprocessor,
+                postprocessor=postprocessor,
+                use_amp=policy.config.use_amp,
                 task=single_task,
                 robot_type=robot.robot_type,
             )
-            action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
+
+            action_names = dataset.features["action"]["names"]
+            act_processed_policy: RobotAction = {
+                f"{name}": float(action_values[i]) for i, name in enumerate(action_names)
+            }
+
         elif policy is None and isinstance(teleop, Teleoperator):
-            action = teleop.get_action()
+            act = teleop.get_action()
+
+            # Applies a pipeline to the raw teleop action, default is IdentityProcessor
+            act_processed_teleop = teleop_action_processor((act, obs))
+
         elif policy is None and isinstance(teleop, list):
-            # TODO(pepijn, steven): clean the record loop for use of multiple robots (possibly with pipeline)
             arm_action = teleop_arm.get_action()
             arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
-
             keyboard_action = teleop_keyboard.get_action()
             base_action = robot._from_keyboard_to_base_action(keyboard_action)
-
-            action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
+            act = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
+            act_processed_teleop = teleop_action_processor((act, obs))
         else:
             logging.info(
                 "No policy or teleoperator provided, skipping action generation."
@@ -273,17 +344,28 @@ def record_loop(
             )
             continue
 
-        # Action can eventually be clipped using `max_relative_target`,
-        # so action actually sent is saved in the dataset.
-        sent_action = robot.send_action(action)
+        # 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))
+        else:
+            action_values = act_processed_teleop
+            robot_action_to_send = robot_action_processor((act_processed_teleop, obs))
 
+        # Send action to robot
+        # Action can eventually be clipped using `max_relative_target`,
+        # so action actually sent is saved in the dataset. action = postprocessor.process(action)
+        # TODO(steven, pepijn, adil): we should use a pipeline step to clip the action, so the sent action is the action that we input to the robot.
+        _sent_action = robot.send_action(robot_action_to_send)
+
+        # Write to dataset
         if dataset is not None:
-            action_frame = build_dataset_frame(dataset.features, sent_action, prefix="action")
+            action_frame = build_dataset_frame(dataset.features, action_values, prefix="action")
             frame = {**observation_frame, **action_frame, "task": single_task}
             dataset.add_frame(frame)
 
         if display_data:
-            log_rerun_data(observation, action)
+            log_rerun_data(observation=obs_processed, action=action_values)
 
         dt_s = time.perf_counter() - start_loop_t
         busy_wait(1 / fps - dt_s)
@@ -301,9 +383,22 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
     robot = make_robot_from_config(cfg.robot)
     teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None
 
-    action_features = hw_to_dataset_features(robot.action_features, "action", cfg.dataset.video)
-    obs_features = hw_to_dataset_features(robot.observation_features, "observation", cfg.dataset.video)
-    dataset_features = {**action_features, **obs_features}
+    teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
+
+    dataset_features = combine_feature_dicts(
+        aggregate_pipeline_dataset_features(
+            pipeline=teleop_action_processor,
+            initial_features=create_initial_features(
+                action=robot.action_features
+            ),  # TODO(steven, pepijn): in future this should be come from teleop or policy
+            use_videos=cfg.dataset.video,
+        ),
+        aggregate_pipeline_dataset_features(
+            pipeline=robot_observation_processor,
+            initial_features=create_initial_features(observation=robot.observation_features),
+            use_videos=cfg.dataset.video,
+        ),
+    )
 
     if cfg.resume:
         dataset = LeRobotDataset(
@@ -335,6 +430,18 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
 
     # Load pretrained policy
     policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)
+    preprocessor = None
+    postprocessor = None
+    if cfg.policy is not None:
+        preprocessor, postprocessor = make_pre_post_processors(
+            policy_cfg=cfg.policy,
+            pretrained_path=cfg.policy.pretrained_path,
+            dataset_stats=rename_stats(dataset.meta.stats, cfg.dataset.rename_map),
+            preprocessor_overrides={
+                "device_processor": {"device": cfg.policy.device},
+                "rename_observations_processor": {"rename_map": cfg.dataset.rename_map},
+            },
+        )
 
     robot.connect()
     if teleop is not None:
@@ -350,8 +457,13 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
                 robot=robot,
                 events=events,
                 fps=cfg.dataset.fps,
+                teleop_action_processor=teleop_action_processor,
+                robot_action_processor=robot_action_processor,
+                robot_observation_processor=robot_observation_processor,
                 teleop=teleop,
                 policy=policy,
+                preprocessor=preprocessor,
+                postprocessor=postprocessor,
                 dataset=dataset,
                 control_time_s=cfg.dataset.episode_time_s,
                 single_task=cfg.dataset.single_task,
@@ -368,6 +480,9 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
                     robot=robot,
                     events=events,
                     fps=cfg.dataset.fps,
+                    teleop_action_processor=teleop_action_processor,
+                    robot_action_processor=robot_action_processor,
+                    robot_observation_processor=robot_observation_processor,
                     teleop=teleop,
                     control_time_s=cfg.dataset.reset_time_s,
                     single_task=cfg.dataset.single_task,

src/lerobot/replay.py

@@ -23,7 +23,7 @@ lerobot-replay \
     --robot.port=/dev/tty.usbmodem58760431541 \
     --robot.id=black \
     --dataset.repo_id=aliberts/record-test \
-    --dataset.episode=2
+    --dataset.episode=0
 ```
 
 Example replay with bimanual so100:
@@ -45,9 +45,11 @@ from dataclasses import asdict, dataclass
 from pathlib import Path
 from pprint import pformat
 
-import draccus
-
+from lerobot.configs import parser
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.processor import (
+    make_default_robot_action_processor,
+)
 from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
@@ -55,7 +57,6 @@ from lerobot.robots import (  # noqa: F401
     hope_jr,
     koch_follower,
     make_robot_from_config,
-    reachy2,
     so100_follower,
     so101_follower,
 )
@@ -86,11 +87,13 @@ class ReplayConfig:
     play_sounds: bool = True
 
 
-@draccus.wrap()
+@parser.wrap()
 def replay(cfg: ReplayConfig):
     init_logging()
     logging.info(pformat(asdict(cfg)))
 
+    robot_action_processor = make_default_robot_action_processor()
+
     robot = make_robot_from_config(cfg.robot)
     dataset = LeRobotDataset(cfg.dataset.repo_id, root=cfg.dataset.root, episodes=[cfg.dataset.episode])
 
@@ -109,7 +112,11 @@ def replay(cfg: ReplayConfig):
         for i, name in enumerate(dataset.features["action"]["names"]):
             action[name] = action_array[i]
 
-        robot.send_action(action)
+        robot_obs = robot.get_observation()
+
+        processed_action = robot_action_processor((action, robot_obs))
+
+        _ = robot.send_action(processed_action)
 
         dt_s = time.perf_counter() - start_episode_t
         busy_wait(1 / dataset.fps - dt_s)

src/lerobot/robots/so100_follower/__init__.py

@@ -14,6 +14,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .config_so100_follower import SO100FollowerConfig, SO100FollowerEndEffectorConfig
+from .config_so100_follower import SO100FollowerConfig
 from .so100_follower import SO100Follower
-from .so100_follower_end_effector import SO100FollowerEndEffector

src/lerobot/robots/so100_follower/config_so100_follower.py

@@ -39,35 +39,3 @@ class SO100FollowerConfig(RobotConfig):
 
     # Set to `True` for backward compatibility with previous policies/dataset
     use_degrees: bool = False
-
-
-@RobotConfig.register_subclass("so100_follower_end_effector")
-@dataclass
-class SO100FollowerEndEffectorConfig(SO100FollowerConfig):
-    """Configuration for the SO100FollowerEndEffector robot."""
-
-    # Path to URDF file for kinematics
-    # 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
-    urdf_path: str | None = None
-
-    # End-effector frame name in the URDF
-    target_frame_name: str = "gripper_frame_link"
-
-    # Default bounds for the end-effector position (in meters)
-    end_effector_bounds: dict[str, list[float]] = field(
-        default_factory=lambda: {
-            "min": [-1.0, -1.0, -1.0],  # min x, y, z
-            "max": [1.0, 1.0, 1.0],  # max x, y, z
-        }
-    )
-
-    max_gripper_pos: float = 50
-
-    end_effector_step_sizes: dict[str, float] = field(
-        default_factory=lambda: {
-            "x": 0.02,
-            "y": 0.02,
-            "z": 0.02,
-        }
-    )

src/lerobot/robots/so100_follower/robot_kinematic_processor.py

@@ -0,0 +1,616 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+
+from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.model.kinematics import RobotKinematics
+from lerobot.processor import (
+    EnvTransition,
+    ObservationProcessorStep,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    RobotAction,
+    RobotActionProcessorStep,
+    TransitionKey,
+)
+from lerobot.utils.rotation import Rotation
+
+
+@ProcessorStepRegistry.register("ee_reference_and_delta")
+@dataclass
+class EEReferenceAndDelta(RobotActionProcessorStep):
+    """
+    Computes a target end-effector pose from a relative delta command.
+
+    This step takes a desired change in position and orientation (`target_*`) and applies it to a
+    reference end-effector pose to calculate an absolute target pose. The reference pose is derived
+    from the current robot joint positions using forward kinematics.
+
+    The processor can operate in two modes:
+    1.  `use_latched_reference=True`: The reference pose is "latched" or saved at the moment the action
+        is first enabled. Subsequent commands are relative to this fixed reference.
+    2.  `use_latched_reference=False`: The reference pose is updated to the robot's current pose at
+        every step.
+
+    Attributes:
+        kinematics: The robot's kinematic model for forward kinematics.
+        end_effector_step_sizes: A dictionary scaling the input delta commands.
+        motor_names: A list of motor names required for forward kinematics.
+        use_latched_reference: If True, latch the reference pose on enable; otherwise, always use the
+            current pose as the reference.
+        reference_ee_pose: Internal state storing the latched reference pose.
+        _prev_enabled: Internal state to detect the rising edge of the enable signal.
+        _command_when_disabled: Internal state to hold the last command while disabled.
+    """
+
+    kinematics: RobotKinematics
+    end_effector_step_sizes: dict
+    motor_names: list[str]
+    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()
+
+        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,
+            )
+
+        if q_raw is None:
+            raise ValueError("Joints observation is require for computing robot kinematics")
+
+        # Current pose from FK on measured joints
+        t_curr = self.kinematics.forward_kinematics(q_raw)
+
+        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"))
+
+        desired = None
+
+        if enabled:
+            ref = t_curr
+            if self.use_latched_reference:
+                # Latched reference mode: latch reference at the rising edge
+                if not self._prev_enabled or self.reference_ee_pose is None:
+                    self.reference_ee_pose = t_curr.copy()
+                ref = self.reference_ee_pose if self.reference_ee_pose is not None else t_curr
+
+            delta_p = np.array(
+                [
+                    tx * self.end_effector_step_sizes["x"],
+                    ty * self.end_effector_step_sizes["y"],
+                    tz * self.end_effector_step_sizes["z"],
+                ],
+                dtype=float,
+            )
+            r_abs = Rotation.from_rotvec([wx, wy, wz]).as_matrix()
+            desired = np.eye(4, dtype=float)
+            desired[:3, :3] = ref[:3, :3] @ r_abs
+            desired[:3, 3] = ref[:3, 3] + delta_p
+
+            self._command_when_disabled = desired.copy()
+        else:
+            # While disabled, keep sending the same command to avoid drift.
+            if self._command_when_disabled is None:
+                # If we've never had an enabled command yet, freeze current FK pose once.
+                self._command_when_disabled = t_curr.copy()
+            desired = self._command_when_disabled.copy()
+
+        # 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
+
+        self._prev_enabled = enabled
+        return action
+
+    def reset(self):
+        """Resets the internal state of the processor."""
+        self._prev_enabled = False
+        self.reference_ee_pose = None
+        self._command_when_disabled = None
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        for feat in [
+            "enabled",
+            "target_x",
+            "target_y",
+            "target_z",
+            "target_wx",
+            "target_wy",
+            "target_wz",
+            "gripper_vel",
+        ]:
+            features[PipelineFeatureType.ACTION].pop(f"{feat}", None)
+
+        for feat in ["x", "y", "z", "wx", "wy", "wz", "gripper_vel"]:
+            features[PipelineFeatureType.ACTION][f"ee.{feat}"] = PolicyFeature(
+                type=FeatureType.ACTION, shape=(1,)
+            )
+
+        return features
+
+
+@ProcessorStepRegistry.register("ee_bounds_and_safety")
+@dataclass
+class EEBoundsAndSafety(RobotActionProcessorStep):
+    """
+    Clips the end-effector pose to predefined bounds and checks for unsafe jumps.
+
+    This step ensures that the target end-effector pose remains within a safe operational workspace.
+    It also moderates the command to prevent large, sudden movements between consecutive steps.
+
+    Attributes:
+        end_effector_bounds: A dictionary with "min" and "max" keys for position clipping.
+        max_ee_step_m: The maximum allowed change in position (in meters) between steps.
+        max_ee_twist_step_rad: The maximum allowed change in orientation (in radians) between steps.
+        _last_pos: Internal state storing the last commanded position.
+        _last_twist: Internal state storing the last commanded orientation.
+    """
+
+    end_effector_bounds: dict
+    max_ee_step_m: float = 0.05
+    max_ee_twist_step_rad: float = 0.20
+    _last_pos: np.ndarray | None = field(default=None, init=False, repr=False)
+    _last_twist: np.ndarray | None = field(default=None, init=False, repr=False)
+
+    def action(self, action: RobotAction) -> RobotAction:
+        x = action["ee.x"]
+        y = action["ee.y"]
+        z = action["ee.z"]
+        wx = action["ee.wx"]
+        wy = action["ee.wy"]
+        wz = action["ee.wz"]
+        # TODO(Steven): ee.gripper_vel does not need to be bounded
+
+        if None in (x, y, z, wx, wy, wz):
+            raise ValueError(
+                "Missing required end-effector pose components: x, y, z, wx, wy, wz must all be present in action"
+            )
+
+        pos = np.array([x, y, z], dtype=float)
+        twist = np.array([wx, wy, wz], dtype=float)
+
+        # Clip position
+        pos = np.clip(pos, self.end_effector_bounds["min"], self.end_effector_bounds["max"])
+
+        # Check for jumps in position
+        if self._last_pos is not None:
+            dpos = pos - self._last_pos
+            n = float(np.linalg.norm(dpos))
+            if n > self.max_ee_step_m and n > 0:
+                pos = self._last_pos + dpos * (self.max_ee_step_m / n)
+                raise ValueError(f"EE jump {n:.3f}m > {self.max_ee_step_m}m")
+
+        self._last_pos = pos
+        self._last_twist = twist
+
+        action["ee.x"] = float(pos[0])
+        action["ee.y"] = float(pos[1])
+        action["ee.z"] = float(pos[2])
+        action["ee.wx"] = float(twist[0])
+        action["ee.wy"] = float(twist[1])
+        action["ee.wz"] = float(twist[2])
+        return action
+
+    def reset(self):
+        """Resets the last known position and orientation."""
+        self._last_pos = None
+        self._last_twist = None
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@ProcessorStepRegistry.register("inverse_kinematics_ee_to_joints")
+@dataclass
+class InverseKinematicsEEToJoints(RobotActionProcessorStep):
+    """
+    Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
+
+    This step translates a Cartesian command (position and orientation of the end-effector) into
+    the corresponding joint-space commands for each motor.
+
+    Attributes:
+        kinematics: The robot's kinematic model for inverse kinematics.
+        motor_names: A list of motor names for which to compute joint positions.
+        q_curr: Internal state storing the last joint positions, used as an initial guess for the IK solver.
+        initial_guess_current_joints: If True, use the robot's current joint state as the IK guess.
+            If False, use the solution from the previous step.
+    """
+
+    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 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")
+
+        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")
+
+        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)
+
+        return action
+
+    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
+
+
+@ProcessorStepRegistry.register("gripper_velocity_to_joint")
+@dataclass
+class GripperVelocityToJoint(RobotActionProcessorStep):
+    """
+    Converts a gripper velocity command into a target gripper joint position.
+
+    This step integrates a normalized velocity command over time to produce a position command,
+    taking the current gripper position as a starting point. It also supports a discrete mode
+    where integer actions map to open, close, or no-op.
+
+    Attributes:
+        motor_names: A list of motor names, which must include 'gripper'.
+        speed_factor: A scaling factor to convert the normalized velocity command to a position change.
+        clip_min: The minimum allowed gripper joint position.
+        clip_max: The maximum allowed gripper joint position.
+        discrete_gripper: If True, treat the input action as discrete (0: open, 1: close, 2: stay).
+    """
+
+    speed_factor: float = 20.0
+    clip_min: float = 0.0
+    clip_max: float = 100.0
+    discrete_gripper: bool = False
+
+    def action(self, action: RobotAction) -> RobotAction:
+        observation = self.transition.get(TransitionKey.OBSERVATION).copy()
+
+        gripper_vel = action.pop("ee.gripper_vel")
+
+        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.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
+
+        # 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))
+        action["ee.gripper_pos"] = gripper_pos
+
+        return action
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        features[PipelineFeatureType.ACTION].pop("ee.gripper_vel", None)
+        features[PipelineFeatureType.ACTION]["ee.gripper_pos"] = PolicyFeature(
+            type=FeatureType.ACTION, shape=(1,)
+        )
+
+        return features
+
+
+def compute_forward_kinematics_joints_to_ee(
+    joints: dict[str, Any], kinematics: RobotKinematics, motor_names: list[str]
+) -> dict[str, Any]:
+    motor_joint_values = [joints[f"{n}.pos"] for n in motor_names]
+
+    q = np.array(motor_joint_values, dtype=float)
+    t = kinematics.forward_kinematics(q)
+    pos = t[:3, 3]
+    tw = Rotation.from_matrix(t[:3, :3]).as_rotvec()
+    gripper_pos = joints["gripper.pos"]
+    for n in motor_names:
+        joints.pop(f"{n}.pos")
+    joints["ee.x"] = float(pos[0])
+    joints["ee.y"] = float(pos[1])
+    joints["ee.z"] = float(pos[2])
+    joints["ee.wx"] = float(tw[0])
+    joints["ee.wy"] = float(tw[1])
+    joints["ee.wz"] = float(tw[2])
+    joints["ee.gripper_pos"] = float(gripper_pos)
+    return joints
+
+
+@ProcessorStepRegistry.register("forward_kinematics_joints_to_ee_observation")
+@dataclass
+class ForwardKinematicsJointsToEEObservation(ObservationProcessorStep):
+    """
+    Computes the end-effector pose from joint positions using forward kinematics (FK).
+
+    This step is typically used to add the robot's Cartesian pose to the observation space,
+    which can be useful for visualization or as an input to a policy.
+
+    Attributes:
+        kinematics: The robot's kinematic model.
+    """
+
+    kinematics: RobotKinematics
+    motor_names: list[str]
+
+    def observation(self, observation: dict[str, Any]) -> dict[str, Any]:
+        return compute_forward_kinematics_joints_to_ee(observation, self.kinematics, self.motor_names)
+
+    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.OBSERVATION].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.OBSERVATION][f"ee.{k}"] = PolicyFeature(
+                type=FeatureType.STATE, shape=(1,)
+            )
+        return features
+
+
+@ProcessorStepRegistry.register("forward_kinematics_joints_to_ee_action")
+@dataclass
+class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
+    """
+    Computes the end-effector pose from joint positions using forward kinematics (FK).
+
+    This step is typically used to add the robot's Cartesian pose to the observation space,
+    which can be useful for visualization or as an input to a policy.
+
+    Attributes:
+        kinematics: The robot's kinematic model.
+    """
+
+    kinematics: RobotKinematics
+    motor_names: list[str]
+
+    def action(self, action: RobotAction) -> RobotAction:
+        return compute_forward_kinematics_joints_to_ee(action, self.kinematics, self.motor_names)
+
+    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
+
+
+@ProcessorStepRegistry.register(name="forward_kinematics_joints_to_ee")
+@dataclass
+class ForwardKinematicsJointsToEE(ProcessorStep):
+    kinematics: RobotKinematics
+    motor_names: list[str]
+
+    def __post_init__(self):
+        self.joints_to_ee_action_processor = ForwardKinematicsJointsToEEAction(
+            kinematics=self.kinematics, motor_names=self.motor_names
+        )
+        self.joints_to_ee_observation_processor = ForwardKinematicsJointsToEEObservation(
+            kinematics=self.kinematics, motor_names=self.motor_names
+        )
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        if transition.get(TransitionKey.ACTION) is not None:
+            transition = self.joints_to_ee_action_processor(transition)
+        if transition.get(TransitionKey.OBSERVATION) is not None:
+            transition = self.joints_to_ee_observation_processor(transition)
+        return transition
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        if features[PipelineFeatureType.ACTION] is not None:
+            features = self.joints_to_ee_action_processor.transform_features(features)
+        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

src/lerobot/robots/so100_follower/so100_follower_end_effector.py

@@ -1,200 +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 logging
-import time
-from typing import Any
-
-import numpy as np
-
-from lerobot.cameras import make_cameras_from_configs
-from lerobot.errors import DeviceNotConnectedError
-from lerobot.model.kinematics import RobotKinematics
-from lerobot.motors import Motor, MotorNormMode
-from lerobot.motors.feetech import FeetechMotorsBus
-
-from . import SO100Follower
-from .config_so100_follower import SO100FollowerEndEffectorConfig
-
-logger = logging.getLogger(__name__)
-
-
-class SO100FollowerEndEffector(SO100Follower):
-    """
-    SO100Follower robot with end-effector space control.
-
-    This robot inherits from SO100Follower but transforms actions from
-    end-effector space to joint space before sending them to the motors.
-    """
-
-    config_class = SO100FollowerEndEffectorConfig
-    name = "so100_follower_end_effector"
-
-    def __init__(self, config: SO100FollowerEndEffectorConfig):
-        super().__init__(config)
-        self.bus = FeetechMotorsBus(
-            port=self.config.port,
-            motors={
-                "shoulder_pan": Motor(1, "sts3215", MotorNormMode.DEGREES),
-                "shoulder_lift": Motor(2, "sts3215", MotorNormMode.DEGREES),
-                "elbow_flex": Motor(3, "sts3215", MotorNormMode.DEGREES),
-                "wrist_flex": Motor(4, "sts3215", MotorNormMode.DEGREES),
-                "wrist_roll": Motor(5, "sts3215", MotorNormMode.DEGREES),
-                "gripper": Motor(6, "sts3215", MotorNormMode.RANGE_0_100),
-            },
-            calibration=self.calibration,
-        )
-
-        self.cameras = make_cameras_from_configs(config.cameras)
-
-        self.config = config
-
-        # Initialize the kinematics module for the so100 robot
-        if self.config.urdf_path is None:
-            raise ValueError(
-                "urdf_path must be provided in the configuration for end-effector control. "
-                "Please set urdf_path in your SO100FollowerEndEffectorConfig."
-            )
-
-        self.kinematics = RobotKinematics(
-            urdf_path=self.config.urdf_path,
-            target_frame_name=self.config.target_frame_name,
-        )
-
-        # Store the bounds for end-effector position
-        self.end_effector_bounds = self.config.end_effector_bounds
-
-        self.current_ee_pos = None
-        self.current_joint_pos = None
-
-    @property
-    def action_features(self) -> dict[str, Any]:
-        """
-        Define action features for end-effector control.
-        Returns dictionary with dtype, shape, and names.
-        """
-        return {
-            "dtype": "float32",
-            "shape": (4,),
-            "names": {"delta_x": 0, "delta_y": 1, "delta_z": 2, "gripper": 3},
-        }
-
-    def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
-        """
-        Transform action from end-effector space to joint space and send to motors.
-
-        Args:
-            action: Dictionary with keys 'delta_x', 'delta_y', 'delta_z' for end-effector control
-                   or a numpy array with [delta_x, delta_y, delta_z]
-
-        Returns:
-            The joint-space action that was sent to the motors
-        """
-
-        if not self.is_connected:
-            raise DeviceNotConnectedError(f"{self} is not connected.")
-
-        # Convert action to numpy array if not already
-        if isinstance(action, dict):
-            if all(k in action for k in ["delta_x", "delta_y", "delta_z"]):
-                delta_ee = np.array(
-                    [
-                        action["delta_x"] * self.config.end_effector_step_sizes["x"],
-                        action["delta_y"] * self.config.end_effector_step_sizes["y"],
-                        action["delta_z"] * self.config.end_effector_step_sizes["z"],
-                    ],
-                    dtype=np.float32,
-                )
-                if "gripper" not in action:
-                    action["gripper"] = [1.0]
-                action = np.append(delta_ee, action["gripper"])
-            else:
-                logger.warning(
-                    f"Expected action keys 'delta_x', 'delta_y', 'delta_z', got {list(action.keys())}"
-                )
-                action = np.zeros(4, dtype=np.float32)
-
-        if self.current_joint_pos is None:
-            # Read current joint positions
-            current_joint_pos = self.bus.sync_read("Present_Position")
-            self.current_joint_pos = np.array([current_joint_pos[name] for name in self.bus.motors])
-
-        # Calculate current end-effector position using forward kinematics
-        if self.current_ee_pos is None:
-            self.current_ee_pos = self.kinematics.forward_kinematics(self.current_joint_pos)
-
-        # Set desired end-effector position by adding delta
-        desired_ee_pos = np.eye(4)
-        desired_ee_pos[:3, :3] = self.current_ee_pos[:3, :3]  # Keep orientation
-
-        # Add delta to position and clip to bounds
-        desired_ee_pos[:3, 3] = self.current_ee_pos[:3, 3] + action[:3]
-        if self.end_effector_bounds is not None:
-            desired_ee_pos[:3, 3] = np.clip(
-                desired_ee_pos[:3, 3],
-                self.end_effector_bounds["min"],
-                self.end_effector_bounds["max"],
-            )
-
-        # Compute inverse kinematics to get joint positions
-        target_joint_values_in_degrees = self.kinematics.inverse_kinematics(
-            self.current_joint_pos, desired_ee_pos
-        )
-
-        # Create joint space action dictionary
-        joint_action = {
-            f"{key}.pos": target_joint_values_in_degrees[i] for i, key in enumerate(self.bus.motors.keys())
-        }
-
-        # Handle gripper separately if included in action
-        # Gripper delta action is in the range 0 - 2,
-        # We need to shift the action to the range -1, 1 so that we can expand it to -Max_gripper_pos, Max_gripper_pos
-        joint_action["gripper.pos"] = np.clip(
-            self.current_joint_pos[-1] + (action[-1] - 1) * self.config.max_gripper_pos,
-            5,
-            self.config.max_gripper_pos,
-        )
-
-        self.current_ee_pos = desired_ee_pos.copy()
-        self.current_joint_pos = target_joint_values_in_degrees.copy()
-        self.current_joint_pos[-1] = joint_action["gripper.pos"]
-
-        # Send joint space action to parent class
-        return super().send_action(joint_action)
-
-    def get_observation(self) -> dict[str, Any]:
-        if not self.is_connected:
-            raise DeviceNotConnectedError(f"{self} is not connected.")
-
-        # Read arm position
-        start = time.perf_counter()
-        obs_dict = self.bus.sync_read("Present_Position")
-        obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
-        dt_ms = (time.perf_counter() - start) * 1e3
-        logger.debug(f"{self} read state: {dt_ms:.1f}ms")
-
-        # Capture images from cameras
-        for cam_key, cam in self.cameras.items():
-            start = time.perf_counter()
-            obs_dict[cam_key] = cam.async_read()
-            dt_ms = (time.perf_counter() - start) * 1e3
-            logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms")
-
-        return obs_dict
-
-    def reset(self):
-        self.current_ee_pos = None
-        self.current_joint_pos = None

src/lerobot/robots/utils.py

@@ -29,10 +29,6 @@ def make_robot_from_config(config: RobotConfig) -> Robot:
         from .so100_follower import SO100Follower
 
         return SO100Follower(config)
-    elif config.type == "so100_follower_end_effector":
-        from .so100_follower import SO100FollowerEndEffector
-
-        return SO100FollowerEndEffector(config)
     elif config.type == "so101_follower":
         from .so101_follower import SO101Follower
 
@@ -73,6 +69,7 @@ def make_robot_from_config(config: RobotConfig) -> Robot:
         raise ValueError(config.type)
 
 
+# TODO(pepijn): Move to pipeline step to make sure we don't have to do this in the robot code and send action to robot is clean for use in dataset
 def ensure_safe_goal_position(
     goal_present_pos: dict[str, tuple[float, float]], max_relative_target: float | dict[str, float]
 ) -> dict[str, float]:

src/lerobot/scripts/eval.py

@@ -56,6 +56,7 @@ from copy import deepcopy
 from dataclasses import asdict
 from pathlib import Path
 from pprint import pformat
+from typing import Any
 
 import einops
 import gymnasium as gym
@@ -69,9 +70,9 @@ from lerobot.configs import parser
 from lerobot.configs.eval import EvalPipelineConfig
 from lerobot.envs.factory import make_env
 from lerobot.envs.utils import add_envs_task, check_env_attributes_and_types, preprocess_observation
-from lerobot.policies.factory import make_policy
+from lerobot.policies.factory import make_policy, make_pre_post_processors
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.utils import get_device_from_parameters
+from lerobot.processor import PolicyAction, PolicyProcessorPipeline
 from lerobot.utils.io_utils import write_video
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.utils import (
@@ -84,6 +85,8 @@ from lerobot.utils.utils import (
 def rollout(
     env: gym.vector.VectorEnv,
     policy: PreTrainedPolicy,
+    preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
     seeds: list[int] | None = None,
     return_observations: bool = False,
     render_callback: Callable[[gym.vector.VectorEnv], None] | None = None,
@@ -120,7 +123,6 @@ def rollout(
         The dictionary described above.
     """
     assert isinstance(policy, nn.Module), "Policy must be a PyTorch nn module."
-    device = get_device_from_parameters(policy)
 
     # Reset the policy and environments.
     policy.reset()
@@ -151,23 +153,20 @@ def rollout(
         if return_observations:
             all_observations.append(deepcopy(observation))
 
-        observation = {
-            key: observation[key].to(device, non_blocking=device.type == "cuda") for key in observation
-        }
-
         # 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 = action.to("cpu").numpy()
-        assert action.ndim == 2, "Action dimensions should be (batch, action_dim)"
+        action_numpy: np.ndarray = action.to("cpu").numpy()
+        assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)"
 
         # Apply the next action.
-        observation, reward, terminated, truncated, info = env.step(action)
+        observation, reward, terminated, truncated, info = env.step(action_numpy)
         if render_callback is not None:
             render_callback(env)
 
@@ -181,7 +180,7 @@ def rollout(
         # Keep track of which environments are done so far.
         done = terminated | truncated | done
 
-        all_actions.append(torch.from_numpy(action))
+        all_actions.append(torch.from_numpy(action_numpy))
         all_rewards.append(torch.from_numpy(reward))
         all_dones.append(torch.from_numpy(done))
         all_successes.append(torch.tensor(successes))
@@ -220,6 +219,8 @@ 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,
     max_episodes_rendered: int = 0,
     videos_dir: Path | None = None,
@@ -296,8 +297,10 @@ def eval_policy(
                 start_seed + (batch_ix * env.num_envs), start_seed + ((batch_ix + 1) * env.num_envs)
             )
         rollout_data = rollout(
-            env,
-            policy,
+            env=env,
+            policy=policy,
+            preprocessor=preprocessor,
+            postprocessor=postprocessor,
             seeds=list(seeds) if seeds else None,
             return_observations=return_episode_data,
             render_callback=render_frame if max_episodes_rendered > 0 else None,
@@ -479,13 +482,22 @@ def eval_main(cfg: EvalPipelineConfig):
         cfg=cfg.policy,
         env_cfg=cfg.env,
     )
+
     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)}},
+    )
 
     with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
         info = eval_policy(
-            env,
-            policy,
-            cfg.eval.n_episodes,
+            env=env,
+            policy=policy,
+            preprocessor=preprocessor,
+            postprocessor=postprocessor,
+            n_episodes=cfg.eval.n_episodes,
             max_episodes_rendered=10,
             videos_dir=Path(cfg.output_dir) / "videos",
             start_seed=cfg.seed,

src/lerobot/scripts/rl/actor.py

@@ -62,9 +62,16 @@ from lerobot.configs import parser
 from lerobot.configs.train import TrainRLServerPipelineConfig
 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 make_robot_env
+from lerobot.scripts.rl.gym_manipulator import (
+    create_transition,
+    make_processors,
+    make_robot_env,
+    step_env_and_process_transition,
+)
 from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
+from lerobot.teleoperators.utils import TeleopEvents
 from lerobot.transport import services_pb2, services_pb2_grpc
 from lerobot.transport.utils import (
     bytes_to_state_dict,
@@ -91,10 +98,7 @@ from lerobot.utils.utils import (
 
 ACTOR_SHUTDOWN_TIMEOUT = 30
 
-
-#################################################
-# Main entry point #
-#################################################
+# Main entry point
 
 
 @parser.wrap()
@@ -201,9 +205,7 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
     logging.info("[ACTOR] queues closed")
 
 
-#################################################
-# Core algorithm functions #
-#################################################
+# Core algorithm functions
 
 
 def act_with_policy(
@@ -236,7 +238,8 @@ def act_with_policy(
 
     logging.info("make_env online")
 
-    online_env = make_robot_env(cfg=cfg.env)
+    online_env, teleop_device = make_robot_env(cfg=cfg.env)
+    env_processor, action_processor = make_processors(online_env, teleop_device, cfg.env, cfg.policy.device)
 
     set_seed(cfg.seed)
     device = get_safe_torch_device(cfg.policy.device, log=True)
@@ -257,6 +260,12 @@ def act_with_policy(
     assert isinstance(policy, nn.Module)
 
     obs, info = online_env.reset()
+    env_processor.reset()
+    action_processor.reset()
+
+    # Process initial observation
+    transition = create_transition(observation=obs, info=info)
+    transition = env_processor(transition)
 
     # NOTE: For the moment we will solely handle the case of a single environment
     sum_reward_episode = 0
@@ -274,45 +283,71 @@ def act_with_policy(
             logging.info("[ACTOR] Shutting down act_with_policy")
             return
 
-        if interaction_step >= cfg.policy.online_step_before_learning:
-            # Time policy inference and check if it meets FPS requirement
-            with policy_timer:
-                action = policy.select_action(batch=obs)
-            policy_fps = policy_timer.fps_last
+        observation = {
+            k: v for k, v in transition[TransitionKey.OBSERVATION].items() if k in cfg.policy.input_features
+        }
 
-            log_policy_frequency_issue(policy_fps=policy_fps, cfg=cfg, interaction_step=interaction_step)
+        # Time policy inference and check if it meets FPS requirement
+        with policy_timer:
+            # Extract observation from transition for policy
+            action = policy.select_action(batch=observation)
+        policy_fps = policy_timer.fps_last
 
-        else:
-            action = online_env.action_space.sample()
+        log_policy_frequency_issue(policy_fps=policy_fps, cfg=cfg, interaction_step=interaction_step)
 
-        next_obs, reward, done, truncated, info = online_env.step(action)
+        # Use the new step function
+        new_transition = step_env_and_process_transition(
+            env=online_env,
+            transition=transition,
+            action=action,
+            env_processor=env_processor,
+            action_processor=action_processor,
+        )
+
+        # Extract values from processed transition
+        next_observation = {
+            k: v
+            for k, v in new_transition[TransitionKey.OBSERVATION].items()
+            if k in cfg.policy.input_features
+        }
+
+        # Teleop action is the action that was executed in the environment
+        # It is either the action from the teleop device or the action from the policy
+        executed_action = new_transition[TransitionKey.COMPLEMENTARY_DATA]["teleop_action"]
+
+        reward = new_transition[TransitionKey.REWARD]
+        done = new_transition.get(TransitionKey.DONE, False)
+        truncated = new_transition.get(TransitionKey.TRUNCATED, False)
 
         sum_reward_episode += float(reward)
-        # Increment total steps counter for intervention rate
         episode_total_steps += 1
 
-        # NOTE: We override the action if the intervention is True, because the action applied is the intervention action
-        if "is_intervention" in info and info["is_intervention"]:
-            # NOTE: The action space for demonstration before hand is with the full action space
-            # but sometimes for example we want to deactivate the gripper
-            action = info["action_intervention"]
+        # Check for intervention from transition info
+        intervention_info = new_transition[TransitionKey.INFO]
+        if intervention_info.get(TeleopEvents.IS_INTERVENTION, False):
             episode_intervention = True
-            # Increment intervention steps counter
             episode_intervention_steps += 1
 
+        complementary_info = {
+            "discrete_penalty": torch.tensor(
+                [new_transition[TransitionKey.COMPLEMENTARY_DATA].get("discrete_penalty", 0.0)]
+            ),
+        }
+        # Create transition for learner (convert to old format)
         list_transition_to_send_to_learner.append(
             Transition(
-                state=obs,
-                action=action,
+                state=observation,
+                action=executed_action,
                 reward=reward,
-                next_state=next_obs,
+                next_state=next_observation,
                 done=done,
-                truncated=truncated,  # TODO: (azouitine) Handle truncation properly
-                complementary_info=info,
+                truncated=truncated,
+                complementary_info=complementary_info,
             )
         )
-        # assign obs to the next obs and continue the rollout
-        obs = next_obs
+
+        # Update transition for next iteration
+        transition = new_transition
 
         if done or truncated:
             logging.info(f"[ACTOR] Global step {interaction_step}: Episode reward: {sum_reward_episode}")
@@ -347,21 +382,27 @@ def act_with_policy(
                 )
             )
 
-            # Reset intervention counters
+            # Reset intervention counters and environment
             sum_reward_episode = 0.0
             episode_intervention = False
             episode_intervention_steps = 0
             episode_total_steps = 0
+
+            # Reset environment and processors
             obs, info = online_env.reset()
+            env_processor.reset()
+            action_processor.reset()
+
+            # Process initial observation
+            transition = create_transition(observation=obs, info=info)
+            transition = env_processor(transition)
 
         if cfg.env.fps is not None:
             dt_time = time.perf_counter() - start_time
             busy_wait(1 / cfg.env.fps - dt_time)
 
 
-#################################################
-#  Communication Functions - Group all gRPC/messaging functions  #
-#################################################
+#  Communication Functions - Group all gRPC/messaging functions
 
 
 def establish_learner_connection(
@@ -606,9 +647,7 @@ def interactions_stream(
     return services_pb2.Empty()
 
 
-#################################################
-#  Policy functions #
-#################################################
+#  Policy functions
 
 
 def update_policy_parameters(policy: SACPolicy, parameters_queue: Queue, device):
@@ -640,9 +679,7 @@ def update_policy_parameters(policy: SACPolicy, parameters_queue: Queue, device)
             logging.info("[ACTOR] Loaded discrete critic parameters from Learner.")
 
 
-#################################################
-#  Utilities functions #
-#################################################
+#  Utilities functions
 
 
 def push_transitions_to_transport_queue(transitions: list, transitions_queue):

src/lerobot/scripts/rl/learner.py

@@ -75,6 +75,7 @@ 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
 from lerobot.transport.utils import (
     MAX_MESSAGE_SIZE,
@@ -102,11 +103,6 @@ from lerobot.utils.wandb_utils import WandBLogger
 LOG_PREFIX = "[LEARNER]"
 
 
-#################################################
-# MAIN ENTRY POINTS AND CORE ALGORITHM FUNCTIONS #
-#################################################
-
-
 @parser.wrap()
 def train_cli(cfg: TrainRLServerPipelineConfig):
     if not use_threads(cfg):
@@ -249,9 +245,7 @@ def start_learner_threads(
     logging.info("[LEARNER] queues closed")
 
 
-#################################################
-# Core algorithm functions #
-#################################################
+# Core algorithm functions
 
 
 def add_actor_information_and_train(
@@ -819,9 +813,7 @@ def make_optimizers_and_scheduler(cfg: TrainRLServerPipelineConfig, policy: nn.M
     return optimizers, lr_scheduler
 
 
-#################################################
-# Training setup functions #
-#################################################
+# Training setup functions
 
 
 def handle_resume_logic(cfg: TrainRLServerPipelineConfig) -> TrainRLServerPipelineConfig:
@@ -1022,9 +1014,7 @@ def initialize_offline_replay_buffer(
     return offline_replay_buffer
 
 
-#################################################
-# Utilities/Helpers functions #
-#################################################
+# Utilities/Helpers functions
 
 
 def get_observation_features(
@@ -1048,10 +1038,8 @@ def get_observation_features(
         return None, None
 
     with torch.no_grad():
-        observation_features = policy.actor.encoder.get_cached_image_features(observations, normalize=True)
-        next_observation_features = policy.actor.encoder.get_cached_image_features(
-            next_observations, normalize=True
-        )
+        observation_features = policy.actor.encoder.get_cached_image_features(observations)
+        next_observation_features = policy.actor.encoder.get_cached_image_features(next_observations)
 
     return observation_features, next_observation_features
 
@@ -1176,7 +1164,7 @@ def process_transitions(
 
             # Add to offline buffer if it's an intervention
             if dataset_repo_id is not None and transition.get("complementary_info", {}).get(
-                "is_intervention"
+                TeleopEvents.IS_INTERVENTION
             ):
                 offline_replay_buffer.add(**transition)
 

src/lerobot/scripts/train.py

@@ -31,7 +31,7 @@ from lerobot.datasets.sampler import EpisodeAwareSampler
 from lerobot.datasets.utils import cycle
 from lerobot.envs.factory import make_env
 from lerobot.optim.factory import make_optimizer_and_scheduler
-from lerobot.policies.factory import make_policy
+from lerobot.policies.factory import make_policy, make_pre_post_processors
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import get_device_from_parameters
 from lerobot.scripts.eval import eval_policy
@@ -64,6 +64,28 @@ def update_policy(
     use_amp: bool = False,
     lock=None,
 ) -> 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.
+        batch: A batch of training data.
+        optimizer: The optimizer used to update the policy's parameters.
+        grad_clip_norm: The maximum norm for gradient clipping.
+        grad_scaler: The GradScaler for automatic mixed-precision training.
+        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.
+        - A dictionary of outputs from the policy's forward pass, for logging purposes.
+    """
     start_time = time.perf_counter()
     device = get_device_from_parameters(policy)
     policy.train()
@@ -107,6 +129,20 @@ def update_policy(
 
 @parser.wrap()
 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.
+    - Handling resumption from a checkpoint.
+    - 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.
+    """
     cfg.validate()
     logging.info(pformat(cfg.to_dict()))
 
@@ -141,6 +177,16 @@ def train(cfg: TrainPipelineConfig):
         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
+
+    preprocessor, postprocessor = make_pre_post_processors(
+        policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path, **processor_kwargs
+    )
+
     logging.info("Creating optimizer and scheduler")
     optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
     grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp)
@@ -205,15 +251,9 @@ 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
 
-        for key in batch:
-            if isinstance(batch[key], torch.Tensor):
-                if batch[key].dtype != torch.bool:
-                    batch[key] = batch[key].type(torch.float32) if device.type == "mps" else batch[key]
-
-                batch[key] = batch[key].to(device, non_blocking=device.type == "cuda")
-
         train_tracker, output_dict = update_policy(
             train_tracker,
             policy,
@@ -245,7 +285,9 @@ def train(cfg: TrainPipelineConfig):
         if cfg.save_checkpoint and is_saving_step:
             logging.info(f"Checkpoint policy after step {step}")
             checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
-            save_checkpoint(checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler)
+            save_checkpoint(
+                checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler, preprocessor, postprocessor
+            )
             update_last_checkpoint(checkpoint_dir)
             if wandb_logger:
                 wandb_logger.log_policy(checkpoint_dir)
@@ -258,9 +300,11 @@ def train(cfg: TrainPipelineConfig):
                 torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
             ):
                 eval_info = eval_policy(
-                    eval_env,
-                    policy,
-                    cfg.eval.n_episodes,
+                    env=eval_env,
+                    policy=policy,
+                    preprocessor=preprocessor,
+                    postprocessor=postprocessor,
+                    n_episodes=cfg.eval.n_episodes,
                     videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
                     max_episodes_rendered=4,
                     start_seed=cfg.seed,
@@ -289,6 +333,8 @@ def train(cfg: TrainPipelineConfig):
 
     if cfg.policy.push_to_hub:
         policy.push_model_to_hub(cfg)
+        preprocessor.push_to_hub(cfg.policy.repo_id)
+        postprocessor.push_to_hub(cfg.policy.repo_id)
 
 
 def main():

src/lerobot/teleoperate.py

@@ -56,11 +56,17 @@ import time
 from dataclasses import asdict, dataclass
 from pprint import pformat
 
-import draccus
 import rerun as rr
 
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
+from lerobot.configs import parser
+from lerobot.processor import (
+    RobotAction,
+    RobotObservation,
+    RobotProcessorPipeline,
+    make_default_processors,
+)
 from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
@@ -100,36 +106,81 @@ class TeleoperateConfig:
 
 
 def teleop_loop(
-    teleop: Teleoperator, robot: Robot, fps: int, display_data: bool = False, duration: float | None = None
+    teleop: Teleoperator,
+    robot: Robot,
+    fps: int,
+    teleop_action_processor: RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction],
+    robot_action_processor: RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction],
+    robot_observation_processor: RobotProcessorPipeline[RobotObservation, RobotObservation],
+    display_data: bool = False,
+    duration: float | None = None,
 ):
+    """
+    This function continuously reads actions from a teleoperation device, processes them through optional
+    pipelines, sends them to a robot, and optionally displays the robot's state. The loop runs at a
+    specified frequency until a set duration is reached or it is manually interrupted.
+
+    Args:
+        teleop: The teleoperator device instance providing control actions.
+        robot: The robot instance being controlled.
+        fps: The target frequency for the control loop in frames per second.
+        display_data: If True, fetches robot observations and displays them in the console and Rerun.
+        duration: The maximum duration of the teleoperation loop in seconds. If None, the loop runs indefinitely.
+        teleop_action_processor: An optional pipeline to process raw actions from the teleoperator.
+        robot_action_processor: An optional pipeline to process actions before they are sent to the robot.
+        robot_observation_processor: An optional pipeline to process raw observations from the robot.
+    """
+
     display_len = max(len(key) for key in robot.action_features)
     start = time.perf_counter()
+
     while True:
         loop_start = time.perf_counter()
-        action = teleop.get_action()
-        if display_data:
-            observation = robot.get_observation()
-            log_rerun_data(observation, action)
 
-        robot.send_action(action)
+        # 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))
+
+        # Process action for robot through pipeline
+        robot_action_to_send = robot_action_processor((teleop_action, obs))
+
+        # 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:
+            # Process robot observation through pipeline
+            obs_transition = robot_observation_processor(obs)
+
+            log_rerun_data(
+                observation=obs_transition,
+                action=teleop_action,
+            )
+
+            print("\n" + "-" * (display_len + 10))
+            print(f"{'NAME':<{display_len}} | {'NORM':>7}")
+            # Display the final robot action that was sent
+            for motor, value in robot_action_to_send.items():
+                print(f"{motor:<{display_len}} | {value:>7.2f}")
+            move_cursor_up(len(robot_action_to_send) + 5)
+
         dt_s = time.perf_counter() - loop_start
         busy_wait(1 / fps - dt_s)
-
         loop_s = time.perf_counter() - loop_start
-
-        print("\n" + "-" * (display_len + 10))
-        print(f"{'NAME':<{display_len}} | {'NORM':>7}")
-        for motor, value in action.items():
-            print(f"{motor:<{display_len}} | {value:>7.2f}")
         print(f"\ntime: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)")
 
         if duration is not None and time.perf_counter() - start >= duration:
             return
 
-        move_cursor_up(len(action) + 5)
 
-
-@draccus.wrap()
+@parser.wrap()
 def teleoperate(cfg: TeleoperateConfig):
     init_logging()
     logging.info(pformat(asdict(cfg)))
@@ -138,12 +189,22 @@ def teleoperate(cfg: TeleoperateConfig):
 
     teleop = make_teleoperator_from_config(cfg.teleop)
     robot = make_robot_from_config(cfg.robot)
+    teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
 
     teleop.connect()
     robot.connect()
 
     try:
-        teleop_loop(teleop, robot, cfg.fps, display_data=cfg.display_data, duration=cfg.teleop_time_s)
+        teleop_loop(
+            teleop=teleop,
+            robot=robot,
+            fps=cfg.fps,
+            display_data=cfg.display_data,
+            duration=cfg.teleop_time_s,
+            teleop_action_processor=teleop_action_processor,
+            robot_action_processor=robot_action_processor,
+            robot_observation_processor=robot_observation_processor,
+        )
     except KeyboardInterrupt:
         pass
     finally:

src/lerobot/teleoperators/__init__.py

@@ -16,4 +16,4 @@
 
 from .config import TeleoperatorConfig
 from .teleoperator import Teleoperator
-from .utils import make_teleoperator_from_config
+from .utils import TeleopEvents, make_teleoperator_from_config

src/lerobot/teleoperators/gamepad/gamepad_utils.py

@@ -16,6 +16,8 @@
 
 import logging
 
+from ..utils import TeleopEvents
+
 
 class InputController:
     """Base class for input controllers that generate motion deltas."""
@@ -134,10 +136,10 @@ class KeyboardController(InputController):
                     return False
                 elif key == keyboard.Key.enter:
                     self.key_states["success"] = True
-                    self.episode_end_status = "success"
+                    self.episode_end_status = TeleopEvents.SUCCESS
                 elif key == keyboard.Key.backspace:
                     self.key_states["failure"] = True
-                    self.episode_end_status = "failure"
+                    self.episode_end_status = TeleopEvents.FAILURE
             except AttributeError:
                 pass
 
@@ -255,13 +257,13 @@ class GamepadController(InputController):
         for event in pygame.event.get():
             if event.type == pygame.JOYBUTTONDOWN:
                 if event.button == 3:
-                    self.episode_end_status = "success"
+                    self.episode_end_status = TeleopEvents.SUCCESS
                 # A button (1) for failure
                 elif event.button == 1:
-                    self.episode_end_status = "failure"
+                    self.episode_end_status = TeleopEvents.FAILURE
                 # X button (0) for rerecord
                 elif event.button == 0:
-                    self.episode_end_status = "rerecord_episode"
+                    self.episode_end_status = TeleopEvents.RERECORD_EPISODE
 
                 # RB button (6) for closing gripper
                 elif event.button == 6:
@@ -295,8 +297,8 @@ class GamepadController(InputController):
         try:
             # Read joystick axes
             # Left stick X and Y (typically axes 0 and 1)
-            x_input = self.joystick.get_axis(0)  # Left/Right
-            y_input = self.joystick.get_axis(1)  # Up/Down (often inverted)
+            y_input = self.joystick.get_axis(0)  # Up/Down (often inverted)
+            x_input = self.joystick.get_axis(1)  # Left/Right
 
             # Right stick Y (typically axis 3 or 4)
             z_input = self.joystick.get_axis(3)  # Up/Down for Z
@@ -308,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
@@ -451,11 +453,11 @@ class GamepadControllerHID(InputController):
                 # Check if X/Square button (bit 5) is pressed for failure
                 # Check if A/Cross button (bit 4) is pressed for rerecording
                 if buttons & 1 << 7:
-                    self.episode_end_status = "success"
+                    self.episode_end_status = TeleopEvents.SUCCESS
                 elif buttons & 1 << 5:
-                    self.episode_end_status = "failure"
+                    self.episode_end_status = TeleopEvents.FAILURE
                 elif buttons & 1 << 4:
-                    self.episode_end_status = "rerecord_episode"
+                    self.episode_end_status = TeleopEvents.RERECORD_EPISODE
                 else:
                     self.episode_end_status = None
 

src/lerobot/teleoperators/gamepad/teleop_gamepad.py

@@ -21,6 +21,7 @@ from typing import Any
 import numpy as np
 
 from ..teleoperator import Teleoperator
+from ..utils import TeleopEvents
 from .configuration_gamepad import GamepadTeleopConfig
 
 
@@ -107,6 +108,48 @@ class GamepadTeleop(Teleoperator):
 
         return action_dict
 
+    def get_teleop_events(self) -> dict[str, Any]:
+        """
+        Get extra control events from the gamepad such as intervention status,
+        episode termination, success indicators, etc.
+
+        Returns:
+            Dictionary containing:
+                - is_intervention: bool - Whether human is currently intervening
+                - terminate_episode: bool - Whether to terminate the current episode
+                - success: bool - Whether the episode was successful
+                - rerecord_episode: bool - Whether to rerecord the episode
+        """
+        if self.gamepad is None:
+            return {
+                TeleopEvents.IS_INTERVENTION: False,
+                TeleopEvents.TERMINATE_EPISODE: False,
+                TeleopEvents.SUCCESS: False,
+                TeleopEvents.RERECORD_EPISODE: False,
+            }
+
+        # Update gamepad state to get fresh inputs
+        self.gamepad.update()
+
+        # Check if intervention is active
+        is_intervention = self.gamepad.should_intervene()
+
+        # Get episode end status
+        episode_end_status = self.gamepad.get_episode_end_status()
+        terminate_episode = episode_end_status in [
+            TeleopEvents.RERECORD_EPISODE,
+            TeleopEvents.FAILURE,
+        ]
+        success = episode_end_status == TeleopEvents.SUCCESS
+        rerecord_episode = episode_end_status == TeleopEvents.RERECORD_EPISODE
+
+        return {
+            TeleopEvents.IS_INTERVENTION: is_intervention,
+            TeleopEvents.TERMINATE_EPISODE: terminate_episode,
+            TeleopEvents.SUCCESS: success,
+            TeleopEvents.RERECORD_EPISODE: rerecord_episode,
+        }
+
     def disconnect(self) -> None:
         """Disconnect from the gamepad."""
         if self.gamepad is not None:

src/lerobot/teleoperators/keyboard/teleop_keyboard.py

@@ -24,6 +24,7 @@ from typing import Any
 from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..teleoperator import Teleoperator
+from ..utils import TeleopEvents
 from .configuration_keyboard import KeyboardEndEffectorTeleopConfig, KeyboardTeleopConfig
 
 PYNPUT_AVAILABLE = True
@@ -176,16 +177,6 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop):
                 "names": {"delta_x": 0, "delta_y": 1, "delta_z": 2},
             }
 
-    def _on_press(self, key):
-        if hasattr(key, "char"):
-            key = key.char
-        self.event_queue.put((key, True))
-
-    def _on_release(self, key):
-        if hasattr(key, "char"):
-            key = key.char
-        self.event_queue.put((key, False))
-
     def get_action(self) -> dict[str, Any]:
         if not self.is_connected:
             raise DeviceNotConnectedError(
@@ -235,3 +226,66 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop):
             action_dict["gripper"] = gripper_action
 
         return action_dict
+
+    def get_teleop_events(self) -> dict[str, Any]:
+        """
+        Get extra control events from the keyboard such as intervention status,
+        episode termination, success indicators, etc.
+
+        Keyboard mappings:
+        - Any movement keys pressed = intervention active
+        - 's' key = success (terminate episode successfully)
+        - 'r' key = rerecord episode (terminate and rerecord)
+        - 'q' key = quit episode (terminate without success)
+
+        Returns:
+            Dictionary containing:
+                - is_intervention: bool - Whether human is currently intervening
+                - terminate_episode: bool - Whether to terminate the current episode
+                - success: bool - Whether the episode was successful
+                - rerecord_episode: bool - Whether to rerecord the episode
+        """
+        if not self.is_connected:
+            return {
+                TeleopEvents.IS_INTERVENTION: False,
+                TeleopEvents.TERMINATE_EPISODE: False,
+                TeleopEvents.SUCCESS: False,
+                TeleopEvents.RERECORD_EPISODE: False,
+            }
+
+        # Check if any movement keys are currently pressed (indicates intervention)
+        movement_keys = [
+            keyboard.Key.up,
+            keyboard.Key.down,
+            keyboard.Key.left,
+            keyboard.Key.right,
+            keyboard.Key.shift,
+            keyboard.Key.shift_r,
+            keyboard.Key.ctrl_r,
+            keyboard.Key.ctrl_l,
+        ]
+        is_intervention = any(self.current_pressed.get(key, False) for key in movement_keys)
+
+        # Check for episode control commands from misc_keys_queue
+        terminate_episode = False
+        success = False
+        rerecord_episode = False
+
+        # Process any pending misc keys
+        while not self.misc_keys_queue.empty():
+            key = self.misc_keys_queue.get_nowait()
+            if key == "s":
+                success = True
+            elif key == "r":
+                terminate_episode = True
+                rerecord_episode = True
+            elif key == "q":
+                terminate_episode = True
+                success = False
+
+        return {
+            TeleopEvents.IS_INTERVENTION: is_intervention,
+            TeleopEvents.TERMINATE_EPISODE: terminate_episode,
+            TeleopEvents.SUCCESS: success,
+            TeleopEvents.RERECORD_EPISODE: rerecord_episode,
+        }

src/lerobot/teleoperators/phone/__init__.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .config_phone import PhoneConfig
+from .teleop_phone import Phone

src/lerobot/teleoperators/phone/config_phone.py

@@ -0,0 +1,36 @@
+#!/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.
+
+from dataclasses import dataclass
+from enum import Enum
+
+import numpy as np
+
+from ..config import TeleoperatorConfig
+
+
+class PhoneOS(Enum):
+    ANDROID = "android"
+    IOS = "ios"
+
+
+@TeleoperatorConfig.register_subclass("phone")
+@dataclass
+class PhoneConfig(TeleoperatorConfig):
+    phone_os: PhoneOS = PhoneOS.IOS
+    camera_offset = np.array(
+        [0.0, -0.02, 0.04]
+    )  # iPhone 14 Pro camera is 2cm off center and 4cm above center

src/lerobot/teleoperators/phone/phone_processor.py

@@ -0,0 +1,110 @@
+# !/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+
+from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.processor import ProcessorStepRegistry, RobotAction, RobotActionProcessorStep
+from lerobot.teleoperators.phone.config_phone import PhoneOS
+
+
+@ProcessorStepRegistry.register("map_phone_action_to_robot_action")
+@dataclass
+class MapPhoneActionToRobotAction(RobotActionProcessorStep):
+    """
+    Maps calibrated phone pose actions to standardized robot action inputs.
+
+    This processor step acts as a bridge between the phone teleoperator's output
+    and the robot's expected action format. It remaps the phone's 6-DoF pose
+    (position and rotation) to the robot's target end-effector pose, applying
+    necessary axis inversions and swaps. It also interprets platform-specific
+    button presses to generate a gripper command.
+
+    Attributes:
+        platform: The operating system of the phone (iOS or Android), used
+            to determine the correct button mappings for the gripper.
+    """
+
+    # TODO(Steven): Gripper vel could be output of phone_teleop directly
+    platform: PhoneOS
+    _enabled_prev: bool = field(default=False, init=False, repr=False)
+
+    def action(self, action: RobotAction) -> RobotAction:
+        """
+        Processes the phone action dictionary to create a robot action dictionary.
+
+        Args:
+            act: The input action dictionary from the phone teleoperator.
+
+        Returns:
+            A new action dictionary formatted for the robot controller.
+
+        Raises:
+            ValueError: If 'pos' or 'rot' keys are missing from the input action.
+        """
+        # Pop them from the action
+        enabled = bool(action.pop("phone.enabled"))
+        pos = action.pop("phone.pos")
+        rot = action.pop("phone.rot")
+        inputs = action.pop("phone.raw_inputs")
+
+        if pos is None or rot is None:
+            raise ValueError("pos and rot must be present in action")
+
+        rotvec = rot.as_rotvec()  # Absolute orientation as rotvec
+
+        # Map certain inputs to certain actions
+        if self.platform == PhoneOS.IOS:
+            gripper_vel = float(inputs.get("a3", 0.0))
+        else:
+            a = float(inputs.get("reservedButtonA", 0.0))
+            b = float(inputs.get("reservedButtonB", 0.0))
+            gripper_vel = (
+                a - b
+            )  # Positive if a is pressed, negative if b is pressed, 0 if both or neither are pressed
+
+        # For some actions we need to invert the axis
+        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
+        return action
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        for feat in ["enabled", "pos", "rot", "raw_inputs"]:
+            features[PipelineFeatureType.ACTION].pop(f"phone.{feat}", None)
+
+        for feat in [
+            "enabled",
+            "target_x",
+            "target_y",
+            "target_z",
+            "target_wx",
+            "target_wy",
+            "target_wz",
+            "gripper_vel",
+        ]:
+            features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
+                type=FeatureType.ACTION, shape=(1,)
+            )
+
+        return features

src/lerobot/teleoperators/phone/teleop_phone.py

@@ -0,0 +1,421 @@
+#!/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.
+
+# Docs:
+# hebi: https://docs.hebi.us/tools.html#mobile-io
+# teleop: https://github.com/SpesRobotics/teleop
+
+import logging
+import threading
+import time
+
+import hebi
+import numpy as np
+from teleop import Teleop
+
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
+from lerobot.teleoperators.teleoperator import Teleoperator
+from lerobot.utils.rotation import Rotation
+
+logger = logging.getLogger(__name__)
+
+
+class BasePhone:
+    _enabled: bool = False
+    _calib_pos: np.ndarray | None = None
+    _calib_rot_inv: Rotation | None = None
+
+    def _reapply_position_calibration(self, pos: np.ndarray) -> None:
+        self._calib_pos = pos.copy()
+
+    @property
+    def is_calibrated(self) -> bool:
+        return (self._calib_pos is not None) and (self._calib_rot_inv is not None)
+
+    @property
+    def action_features(self) -> dict[str, type]:
+        return {
+            "phone.pos": np.ndarray,  # shape (3,)
+            "phone.rot": Rotation,  # scipy.spatial.transform.Rotation
+            "phone.raw_inputs": dict,  # analogs/buttons or webXR meta
+            "phone.enabled": bool,
+        }
+
+    @property
+    def feedback_features(self) -> dict[str, type]:
+        # No haptic or other feedback implemented yet
+        pass
+
+    def configure(self) -> None:
+        # No additional configuration required for phone teleop
+        pass
+
+    def send_feedback(self, feedback: dict[str, float]) -> None:
+        # We could add haptic feedback (vibrations) here, but it's not implemented yet
+        raise NotImplementedError
+
+
+class IOSPhone(BasePhone, Teleoperator):
+    name = "ios_phone"
+
+    def __init__(self, config: PhoneConfig):
+        super().__init__(config)
+        self.config = config
+        self._group = None
+
+    @property
+    def is_connected(self) -> bool:
+        return self._group is not None
+
+    def connect(self) -> None:
+        if self.is_connected:
+            raise DeviceAlreadyConnectedError(f"{self} already connected")
+
+        logger.info("Connecting to IPhone, make sure to open the HEBI Mobile I/O app.")
+        lookup = hebi.Lookup()
+        time.sleep(2.0)
+        group = lookup.get_group_from_names(["HEBI"], ["mobileIO"])
+        if group is None:
+            raise RuntimeError("Mobile I/O not found — check name/family settings in the app.")
+        self._group = group
+        logger.info(f"{self} connected to HEBI group with {group.size} module(s).")
+
+        self.calibrate()
+
+    def calibrate(self) -> None:
+        print(
+            "Hold the phone so that: top edge points forward in same direction as the robot (robot +x) and screen points up (robot +z)"
+        )
+        print("Press and hold B1 in the HEBI Mobile I/O app to capture this pose...\n")
+        position, rotation = self._wait_for_capture_trigger()
+        self._calib_pos = position.copy()
+        self._calib_rot_inv = rotation.inv()
+        self._enabled = False
+        print("Calibration done\n")
+
+    def _wait_for_capture_trigger(self) -> tuple[np.ndarray, Rotation]:
+        """
+        Blocks execution until the calibration trigger is detected from the iOS device.
+
+        This method enters a loop, continuously reading the phone's state. It waits for the user to press
+        and hold the 'B1' button in the HEBI Mobile I/O app. Once B1 is pressed, the loop breaks and
+        returns the phone's pose at that exact moment.
+
+        Returns:
+            A tuple containing the position (np.ndarray) and rotation (Rotation) of the phone at the
+            moment the trigger was activated.
+        """
+        while True:
+            has_pose, position, rotation, fb_pose = self._read_current_pose()
+            if not has_pose:
+                time.sleep(0.01)
+                continue
+
+            io = getattr(fb_pose, "io", None)
+            button_b = getattr(io, "b", None) if io is not None else None
+            button_b1_pressed = False
+            if button_b is not None:
+                button_b1_pressed = bool(button_b.get_int(1))
+            if button_b1_pressed:
+                return position, rotation
+
+            time.sleep(0.01)
+
+    def _read_current_pose(self) -> tuple[bool, np.ndarray | None, Rotation | None, object | None]:
+        """
+        Reads the instantaneous 6-DoF pose from the connected iOS device via the HEBI SDK.
+
+        This method fetches the latest feedback packet from the HEBI group, extracts the ARKit
+        position and orientation, and converts them into a standard format. It also applies a
+        configured camera offset to adjust the pose from the camera's frame to the phone's
+        physical frame.
+
+        Returns:
+            A tuple containing:
+            - A boolean indicating if a valid pose was successfully read.
+            - The 3D position as a NumPy array, or None if not available.
+            - The orientation as a `Rotation` object, or None if not available.
+            - The raw HEBI feedback object for accessing other data like button presses.
+        """
+        fbk = self._group.get_next_feedback()
+        pose = fbk[0]
+        ar_pos = getattr(pose, "ar_position", None)
+        ar_quat = getattr(pose, "ar_orientation", None)
+        if ar_pos is None or ar_quat is None:
+            return False, None, None, None
+        # HEBI provides orientation in w, x, y, z format.
+        # Scipy's Rotation expects x, y, z, w.
+        quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]]))  # wxyz to xyzw
+        rot = Rotation.from_quat(quat_xyzw)
+        pos = ar_pos - rot.apply(self.config.camera_offset)
+        return True, pos, rot, pose
+
+    def get_action(self) -> dict:
+        has_pose, raw_position, raw_rotation, fb_pose = self._read_current_pose()
+        if not has_pose or not self.is_calibrated:
+            return {}
+
+        # Collect raw inputs (B1 / analogs on iOS, move/scale on Android)
+        raw_inputs: dict[str, float | int | bool] = {}
+        io = getattr(fb_pose, "io", None)
+        if io is not None:
+            bank_a, bank_b = io.a, io.b
+            if bank_a:
+                for ch in range(1, 9):
+                    if bank_a.has_float(ch):
+                        raw_inputs[f"a{ch}"] = float(bank_a.get_float(ch))
+            if bank_b:
+                for ch in range(1, 9):
+                    if bank_b.has_int(ch):
+                        raw_inputs[f"b{ch}"] = int(bank_b.get_int(ch))
+                    elif hasattr(bank_b, "has_bool") and bank_b.has_bool(ch):
+                        raw_inputs[f"b{ch}"] = int(bank_b.get_bool(ch))
+
+        enable = bool(raw_inputs.get("b1", 0))
+
+        # Rising edge then re-capture calibration immediately from current raw pose
+        if enable and not self._enabled:
+            self._reapply_position_calibration(raw_position)
+
+        # Apply calibration
+        pos_cal = self._calib_rot_inv.apply(raw_position - self._calib_pos)
+        rot_cal = self._calib_rot_inv * raw_rotation
+
+        self._enabled = enable
+
+        return {
+            "phone.pos": pos_cal,
+            "phone.rot": rot_cal,
+            "phone.raw_inputs": raw_inputs,
+            "phone.enabled": self._enabled,
+        }
+
+    def disconnect(self) -> None:
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+
+        self._group = None
+
+
+class AndroidPhone(BasePhone, Teleoperator):
+    name = "android_phone"
+
+    def __init__(self, config: PhoneConfig):
+        super().__init__(config)
+        self.config = config
+        self._teleop = None
+        self._teleop_thread = None
+        self._latest_pose = None
+        self._latest_message = None
+        self._android_lock = threading.Lock()
+
+    @property
+    def is_connected(self) -> bool:
+        return self._teleop is not None
+
+    def connect(self) -> None:
+        if self.is_connected:
+            raise DeviceAlreadyConnectedError(f"{self} already connected")
+
+        logger.info("Starting teleop stream for Android...")
+        self._teleop = Teleop()
+        self._teleop.subscribe(self._android_callback)
+        self._teleop_thread = threading.Thread(target=self._teleop.run, daemon=True)
+        self._teleop_thread.start()
+        logger.info(f"{self} connected, teleop stream started.")
+
+        self.calibrate()
+
+    def calibrate(self) -> None:
+        print(
+            "Hold the phone so that: top edge points forward in same direction as the robot (robot +x) and screen points up (robot +z)"
+        )
+        print("Touch and move on the WebXR page to capture this pose...\n")
+
+        pos, rot = self._wait_for_capture_trigger()
+        self._calib_pos = pos.copy()
+        self._calib_rot_inv = rot.inv()
+        self._enabled = False
+        print("Calibration done\n")
+
+    def _wait_for_capture_trigger(self) -> tuple[np.ndarray, Rotation]:
+        """
+        Blocks execution until the calibration trigger is detected from the Android device.
+
+        This method enters a loop, continuously checking the latest message received from the WebXR
+        session. It waits for the user to touch and move their finger on the screen, which generates
+        a `move` event. Once this event is detected, the loop breaks and returns the phone's current
+        pose.
+
+        Returns:
+            A tuple containing the position (np.ndarray) and rotation (Rotation) of the phone at the
+            moment the trigger was activated.
+        """
+        while True:
+            with self._android_lock:
+                msg = self._latest_message or {}
+
+            if bool(msg.get("move", False)):
+                ok, pos, rot, _pose = self._read_current_pose()
+                if ok:
+                    return pos, rot
+
+            time.sleep(0.01)
+
+    def _read_current_pose(self) -> tuple[bool, np.ndarray | None, Rotation | None, object | None]:
+        """
+        Reads the latest 6-DoF pose received from the Android device's WebXR session.
+
+        This method accesses the most recent pose data stored by the `_android_callback`. It uses a
+        thread lock to safely read the shared `_latest_pose` variable. The pose, a 4x4 matrix, is
+        then decomposed into position and rotation, and the configured camera offset is applied.
+
+        Returns:
+            A tuple containing:
+            - A boolean indicating if a valid pose was available.
+            - The 3D position as a NumPy array, or None if no pose has been received yet.
+            - The orientation as a `Rotation` object, or None if no pose has been received.
+            - The raw 4x4 pose matrix as received from the teleop stream.
+        """
+        with self._android_lock:
+            if self._latest_pose is None:
+                return False, None, None, None
+            p = self._latest_pose.copy()
+            pose = self._latest_pose
+        rot = Rotation.from_matrix(p[:3, :3])
+        pos = p[:3, 3] - rot.apply(self.config.camera_offset)
+        return True, pos, rot, pose
+
+    def _android_callback(self, pose: np.ndarray, message: dict) -> None:
+        """
+        Callback function to handle incoming data from the Android teleop stream.
+
+        This method is executed by the `teleop` package's subscriber thread whenever a new
+        pose and message are received from the WebXR session on the Android phone. It updates
+        the internal state (`_latest_pose` and `_latest_message`) with the new data.
+        A thread lock is used to ensure that these shared variables are updated atomically,
+        preventing race conditions with the main thread that reads them.
+
+        Args:
+            pose: A 4x4 NumPy array representing the phone's transformation matrix.
+            message: A dictionary containing additional data, such as button presses or touch events.
+        """
+        with self._android_lock:
+            self._latest_pose = pose
+            self._latest_message = message
+
+    def get_action(self) -> dict:
+        ok, raw_pos, raw_rot, pose = self._read_current_pose()
+        if not ok or not self.is_calibrated:
+            return {}
+
+        # Collect raw inputs (B1 / analogs on iOS, move/scale on Android)
+        raw_inputs: dict[str, float | int | bool] = {}
+        msg = self._latest_message or {}
+        raw_inputs["move"] = bool(msg.get("move", False))
+        raw_inputs["scale"] = float(msg.get("scale", 1.0))
+        raw_inputs["reservedButtonA"] = bool(msg.get("reservedButtonA", False))
+        raw_inputs["reservedButtonB"] = bool(msg.get("reservedButtonB", False))
+
+        enable = bool(raw_inputs.get("move", False))
+
+        # Rising edge then re-capture calibration immediately from current raw pose
+        if enable and not self._enabled:
+            self._reapply_position_calibration(raw_pos)
+
+        # Apply calibration
+        pos_cal = self._calib_rot_inv.apply(raw_pos - self._calib_pos)
+        rot_cal = self._calib_rot_inv * raw_rot
+
+        self._enabled = enable
+
+        return {
+            "phone.pos": pos_cal,
+            "phone.rot": rot_cal,
+            "phone.raw_inputs": raw_inputs,
+            "phone.enabled": self._enabled,
+        }
+
+    def disconnect(self) -> None:
+        if not self.is_connected:
+            raise DeviceNotConnectedError(f"{self} is not connected.")
+
+        self._teleop = None
+        if self._teleop_thread and self._teleop_thread.is_alive():
+            self._teleop_thread.join(timeout=1.0)
+            self._teleop_thread = None
+            self._latest_pose = None
+
+
+class Phone(Teleoperator):
+    """
+    Phone-based teleoperator using ARKit (iOS via HEBI Mobile I/O App) or the teleop Python package (Android via WebXR API).
+    For HEBI Mobile I/O we also expose 8 analog (a1-a8) and 8 digital (b1-b8) inputs.
+
+    Press and hold **B1** to enable teleoperation. While enabled, the first B1 press
+    captures a reference pose and rotation, when disabled and pressed again the position is reapplied.
+    """
+
+    config_class = PhoneConfig
+    name = "phone"
+
+    def __init__(self, config: PhoneConfig):
+        super().__init__(config)
+        self.config = config
+
+        self._phone_impl: Teleoperator
+
+        if self.config.phone_os == PhoneOS.IOS:
+            self._phone_impl = IOSPhone(config)
+        elif self.config.phone_os == PhoneOS.ANDROID:
+            self._phone_impl = AndroidPhone(config)
+        else:
+            raise ValueError(f"Invalid config phone_os: {self.config.phone_os}")
+
+    @property
+    def is_connected(self) -> bool:
+        return self._phone_impl.is_connected
+
+    def connect(self) -> None:
+        return self._phone_impl.connect()
+
+    def calibrate(self) -> None:
+        return self._phone_impl.calibrate()
+
+    @property
+    def is_calibrated(self) -> bool:
+        return self._phone_impl.is_calibrated
+
+    @property
+    def action_features(self) -> dict[str, type]:
+        return self._phone_impl.action_features
+
+    @property
+    def feedback_features(self) -> dict[str, type]:
+        return self._phone_impl.feedback_features
+
+    def configure(self) -> None:
+        return self._phone_impl.configure()
+
+    def get_action(self) -> dict:
+        return self._phone_impl.get_action()
+
+    def send_feedback(self, feedback: dict[str, float]) -> None:
+        return self._phone_impl.send_feedback(feedback)
+
+    def disconnect(self) -> None:
+        return self._phone_impl.disconnect()

src/lerobot/teleoperators/utils.py

@@ -12,10 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from enum import Enum
+
 from .config import TeleoperatorConfig
 from .teleoperator import Teleoperator
 
 
+class TeleopEvents(Enum):
+    """Shared constants for teleoperator events across teleoperators."""
+
+    SUCCESS = "success"
+    FAILURE = "failure"
+    RERECORD_EPISODE = "rerecord_episode"
+    IS_INTERVENTION = "is_intervention"
+    TERMINATE_EPISODE = "terminate_episode"
+
+
 def make_teleoperator_from_config(config: TeleoperatorConfig) -> Teleoperator:
     if config.type == "keyboard":
         from .keyboard import KeyboardTeleop

src/lerobot/utils/control_utils.py

@@ -22,6 +22,7 @@ import traceback
 from contextlib import nullcontext
 from copy import copy
 from functools import cache
+from typing import Any
 
 import numpy as np
 import torch
@@ -31,10 +32,25 @@ from termcolor import colored
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import DEFAULT_FEATURES
 from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.processor import PolicyAction, PolicyProcessorPipeline
 from lerobot.robots import Robot
 
 
 def log_control_info(robot: Robot, dt_s, episode_index=None, frame_index=None, fps=None):
+    """
+    Logs performance metrics for a single step of the robot control loop.
+
+    This function formats and prints a single line of log information, including episode/frame counters,
+    total loop time (dt), and detailed timings for various robot and camera operations. It can also
+    highlight performance drops in yellow if the actual FPS is lower than the target FPS.
+
+    Args:
+        robot: The `Robot` instance, used to access its internal logs for detailed timings.
+        dt_s: The total duration of the control loop step in seconds.
+        episode_index: The index of the current episode.
+        frame_index: The index of the current frame within the episode.
+        fps: The target frames per second, used to check for performance degradation.
+    """
     log_items = []
     if episode_index is not None:
         log_items.append(f"ep:{episode_index}")
@@ -80,7 +96,16 @@ def log_control_info(robot: Robot, dt_s, episode_index=None, frame_index=None, f
 
 @cache
 def is_headless():
-    """Detects if python is running without a monitor."""
+    """
+    Detects if the Python script is running in a headless environment (e.g., without a display).
+
+    This function attempts to import `pynput`, a library that requires a graphical environment.
+    If the import fails, it assumes the environment is headless. The result is cached to avoid
+    re-running the check.
+
+    Returns:
+        True if the environment is determined to be headless, False otherwise.
+    """
     try:
         import pynput  # noqa
 
@@ -101,10 +126,35 @@ def predict_action(
     observation: dict[str, np.ndarray],
     policy: PreTrainedPolicy,
     device: torch.device,
+    preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction],
     use_amp: bool,
     task: str | None = None,
     robot_type: str | None = None,
 ):
+    """
+    Performs a single-step inference to predict a robot action from an observation.
+
+    This function encapsulates the full inference pipeline:
+    1. Prepares the observation by converting it to PyTorch tensors and adding a batch dimension.
+    2. Runs the preprocessor pipeline on the observation.
+    3. Feeds the processed observation to the policy to get a raw action.
+    4. Runs the postprocessor pipeline on the raw action.
+    5. Formats the final action by removing the batch dimension and moving it to the CPU.
+
+    Args:
+        observation: A dictionary of NumPy arrays representing the robot's current observation.
+        policy: The `PreTrainedPolicy` model to use for action prediction.
+        device: The `torch.device` (e.g., 'cuda' or 'cpu') to run inference on.
+        preprocessor: The `PolicyProcessorPipeline` for preprocessing observations.
+        postprocessor: The `PolicyProcessorPipeline` for postprocessing actions.
+        use_amp: A boolean to enable/disable Automatic Mixed Precision for CUDA inference.
+        task: An optional string identifier for the task.
+        robot_type: An optional string identifier for the robot type.
+
+    Returns:
+        A `torch.Tensor` containing the predicted action, ready for the robot.
+    """
     observation = copy(observation)
     with (
         torch.inference_mode(),
@@ -122,10 +172,14 @@ def predict_action(
         observation["task"] = task if task else ""
         observation["robot_type"] = robot_type if robot_type else ""
 
+        observation = preprocessor(observation)
+
         # Compute the next action with the policy
         # based on the current observation
         action = policy.select_action(observation)
 
+        action = postprocessor(action)
+
         # Remove batch dimension
         action = action.squeeze(0)
 
@@ -136,6 +190,18 @@ def predict_action(
 
 
 def init_keyboard_listener():
+    """
+    Initializes a non-blocking keyboard listener for real-time user interaction.
+
+    This function sets up a listener for specific keys (right arrow, left arrow, escape) to control
+    the program flow during execution, such as stopping recording or exiting loops. It gracefully
+    handles headless environments where keyboard listening is not possible.
+
+    Returns:
+        A tuple containing:
+        - The `pynput.keyboard.Listener` instance, or `None` if in a headless environment.
+        - A dictionary of event flags (e.g., `exit_early`) that are set by key presses.
+    """
     # Allow to exit early while recording an episode or resetting the environment,
     # by tapping the right arrow key '->'. This might require a sudo permission
     # to allow your terminal to monitor keyboard events.
@@ -177,6 +243,19 @@ def init_keyboard_listener():
 
 
 def sanity_check_dataset_name(repo_id, policy_cfg):
+    """
+    Validates the dataset repository name against the presence of a policy configuration.
+
+    This function enforces a naming convention: a dataset repository ID should start with "eval_"
+    if and only if a policy configuration is provided for evaluation purposes.
+
+    Args:
+        repo_id: The Hugging Face Hub repository ID of the dataset.
+        policy_cfg: The configuration object for the policy, or `None`.
+
+    Raises:
+        ValueError: If the naming convention is violated.
+    """
     _, dataset_name = repo_id.split("/")
     # either repo_id doesnt start with "eval_" and there is no policy
     # or repo_id starts with "eval_" and there is a policy
@@ -197,6 +276,21 @@ def sanity_check_dataset_name(repo_id, policy_cfg):
 def sanity_check_dataset_robot_compatibility(
     dataset: LeRobotDataset, robot: Robot, fps: int, features: dict
 ) -> None:
+    """
+    Checks if a dataset's metadata is compatible with the current robot and recording setup.
+
+    This function compares key metadata fields (`robot_type`, `fps`, and `features`) from the
+    dataset against the current configuration to ensure that appended data will be consistent.
+
+    Args:
+        dataset: The `LeRobotDataset` instance to check.
+        robot: The `Robot` instance representing the current hardware setup.
+        fps: The current recording frequency (frames per second).
+        features: The dictionary of features for the current recording session.
+
+    Raises:
+        ValueError: If any of the checked metadata fields do not match.
+    """
     fields = [
         ("robot_type", dataset.meta.robot_type, robot.robot_type),
         ("fps", dataset.fps, fps),

src/lerobot/utils/import_utils.py

@@ -58,6 +58,7 @@ def is_package_available(pkg_name: str, return_version: bool = False) -> tuple[b
 
 
 _torch_available, _torch_version = is_package_available("torch", return_version=True)
+_transformers_available = is_package_available("transformers")
 _gym_xarm_available = is_package_available("gym_xarm")
 _gym_aloha_available = is_package_available("gym_aloha")
 _gym_pusht_available = is_package_available("gym_pusht")

src/lerobot/utils/rotation.py

@@ -0,0 +1,270 @@
+#!/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.
+
+"""Custom rotation utilities to replace scipy.spatial.transform.Rotation."""
+
+import numpy as np
+
+
+class Rotation:
+    """
+    Custom rotation class that provides a subset of scipy.spatial.transform.Rotation functionality.
+
+    Supports conversions between rotation vectors, rotation matrices, and quaternions.
+    """
+
+    def __init__(self, quat: np.ndarray) -> None:
+        """Initialize rotation from quaternion [x, y, z, w]."""
+        self._quat = np.asarray(quat, dtype=float)
+        # Normalize quaternion
+        norm = np.linalg.norm(self._quat)
+        if norm > 0:
+            self._quat = self._quat / norm
+
+    @classmethod
+    def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
+        """
+        Create rotation from rotation vector using Rodrigues' formula.
+
+        Args:
+            rotvec: Rotation vector [x, y, z] where magnitude is angle in radians
+
+        Returns:
+            Rotation instance
+        """
+        rotvec = np.asarray(rotvec, dtype=float)
+        angle = np.linalg.norm(rotvec)
+
+        if angle < 1e-8:
+            # For very small angles, use identity quaternion
+            quat = np.array([0.0, 0.0, 0.0, 1.0])
+        else:
+            axis = rotvec / angle
+            half_angle = angle / 2.0
+            sin_half = np.sin(half_angle)
+            cos_half = np.cos(half_angle)
+
+            # Quaternion [x, y, z, w]
+            quat = np.array([axis[0] * sin_half, axis[1] * sin_half, axis[2] * sin_half, cos_half])
+
+        return cls(quat)
+
+    @classmethod
+    def from_matrix(cls, matrix: np.ndarray) -> "Rotation":
+        """
+        Create rotation from 3x3 rotation matrix.
+
+        Args:
+            matrix: 3x3 rotation matrix
+
+        Returns:
+            Rotation instance
+        """
+        matrix = np.asarray(matrix, dtype=float)
+
+        # Shepherd's method for converting rotation matrix to quaternion
+        trace = np.trace(matrix)
+
+        if trace > 0:
+            s = np.sqrt(trace + 1.0) * 2  # s = 4 * qw
+            qw = 0.25 * s
+            qx = (matrix[2, 1] - matrix[1, 2]) / s
+            qy = (matrix[0, 2] - matrix[2, 0]) / s
+            qz = (matrix[1, 0] - matrix[0, 1]) / s
+        elif matrix[0, 0] > matrix[1, 1] and matrix[0, 0] > matrix[2, 2]:
+            s = np.sqrt(1.0 + matrix[0, 0] - matrix[1, 1] - matrix[2, 2]) * 2  # s = 4 * qx
+            qw = (matrix[2, 1] - matrix[1, 2]) / s
+            qx = 0.25 * s
+            qy = (matrix[0, 1] + matrix[1, 0]) / s
+            qz = (matrix[0, 2] + matrix[2, 0]) / s
+        elif matrix[1, 1] > matrix[2, 2]:
+            s = np.sqrt(1.0 + matrix[1, 1] - matrix[0, 0] - matrix[2, 2]) * 2  # s = 4 * qy
+            qw = (matrix[0, 2] - matrix[2, 0]) / s
+            qx = (matrix[0, 1] + matrix[1, 0]) / s
+            qy = 0.25 * s
+            qz = (matrix[1, 2] + matrix[2, 1]) / s
+        else:
+            s = np.sqrt(1.0 + matrix[2, 2] - matrix[0, 0] - matrix[1, 1]) * 2  # s = 4 * qz
+            qw = (matrix[1, 0] - matrix[0, 1]) / s
+            qx = (matrix[0, 2] + matrix[2, 0]) / s
+            qy = (matrix[1, 2] + matrix[2, 1]) / s
+            qz = 0.25 * s
+
+        quat = np.array([qx, qy, qz, qw])
+        return cls(quat)
+
+    @classmethod
+    def from_quat(cls, quat: np.ndarray) -> "Rotation":
+        """
+        Create rotation from quaternion.
+
+        Args:
+            quat: Quaternion [x, y, z, w] or [w, x, y, z] (specify convention in docstring)
+                  This implementation expects [x, y, z, w] format
+
+        Returns:
+            Rotation instance
+        """
+        return cls(quat)
+
+    def as_matrix(self) -> np.ndarray:
+        """
+        Convert rotation to 3x3 rotation matrix.
+
+        Returns:
+            3x3 rotation matrix
+        """
+        qx, qy, qz, qw = self._quat
+
+        # Compute rotation matrix from quaternion
+        return np.array(
+            [
+                [1 - 2 * (qy * qy + qz * qz), 2 * (qx * qy - qz * qw), 2 * (qx * qz + qy * qw)],
+                [2 * (qx * qy + qz * qw), 1 - 2 * (qx * qx + qz * qz), 2 * (qy * qz - qx * qw)],
+                [2 * (qx * qz - qy * qw), 2 * (qy * qz + qx * qw), 1 - 2 * (qx * qx + qy * qy)],
+            ],
+            dtype=float,
+        )
+
+    def as_rotvec(self) -> np.ndarray:
+        """
+        Convert rotation to rotation vector.
+
+        Returns:
+            Rotation vector [x, y, z] where magnitude is angle in radians
+        """
+        qx, qy, qz, qw = self._quat
+
+        # Ensure qw is positive for unique representation
+        if qw < 0:
+            qx, qy, qz, qw = -qx, -qy, -qz, -qw
+
+        # Compute angle and axis
+        angle = 2.0 * np.arccos(np.clip(abs(qw), 0.0, 1.0))
+        sin_half_angle = np.sqrt(1.0 - qw * qw)
+
+        if sin_half_angle < 1e-8:
+            # For very small angles, use linearization: rotvec ≈ 2 * [qx, qy, qz]
+            return 2.0 * np.array([qx, qy, qz])
+
+        # Extract axis and scale by angle
+        axis = np.array([qx, qy, qz]) / sin_half_angle
+        return angle * axis
+
+    def as_quat(self) -> np.ndarray:
+        """
+        Get quaternion representation.
+
+        Returns:
+            Quaternion [x, y, z, w]
+        """
+        return self._quat.copy()
+
+    def apply(self, vectors: np.ndarray, inverse: bool = False) -> np.ndarray:
+        """
+        Apply this rotation to a set of vectors.
+
+        This is equivalent to applying the rotation matrix to the vectors:
+        self.as_matrix() @ vectors (or self.as_matrix().T @ vectors if inverse=True).
+
+        Args:
+            vectors: Array of shape (3,) or (N, 3) representing vectors in 3D space
+            inverse: If True, apply the inverse of the rotation. Default is False.
+
+        Returns:
+            Rotated vectors with shape:
+            - (3,) if input was single vector with shape (3,)
+            - (N, 3) in all other cases
+        """
+        vectors = np.asarray(vectors, dtype=float)
+        original_shape = vectors.shape
+
+        # Handle single vector case - ensure it's 2D for matrix multiplication
+        if vectors.ndim == 1:
+            if len(vectors) != 3:
+                raise ValueError("Single vector must have length 3")
+            vectors = vectors.reshape(1, 3)
+            single_vector = True
+        elif vectors.ndim == 2:
+            if vectors.shape[1] != 3:
+                raise ValueError("Vectors must have shape (N, 3)")
+            single_vector = False
+        else:
+            raise ValueError("Vectors must be 1D or 2D array")
+
+        # Get rotation matrix
+        rotation_matrix = self.as_matrix()
+
+        # Apply inverse if requested (transpose for orthogonal rotation matrices)
+        if inverse:
+            rotation_matrix = rotation_matrix.T
+
+        # Apply rotation: (N, 3) @ (3, 3).T -> (N, 3)
+        rotated_vectors = vectors @ rotation_matrix.T
+
+        # Return original shape for single vector case
+        if single_vector and original_shape == (3,):
+            return rotated_vectors.flatten()
+
+        return rotated_vectors
+
+    def inv(self) -> "Rotation":
+        """
+        Invert this rotation.
+
+        Composition of a rotation with its inverse results in an identity transformation.
+
+        Returns:
+            Rotation instance containing the inverse of this rotation
+        """
+        qx, qy, qz, qw = self._quat
+
+        # For a unit quaternion, the inverse is the conjugate: [-x, -y, -z, w]
+        inverse_quat = np.array([-qx, -qy, -qz, qw])
+
+        return Rotation(inverse_quat)
+
+    def __mul__(self, other: "Rotation") -> "Rotation":
+        """
+        Compose this rotation with another rotation using the * operator.
+
+        The composition `r2 * r1` means "apply r1 first, then r2".
+        This is equivalent to applying rotation matrices: r2.as_matrix() @ r1.as_matrix()
+
+        Args:
+            other: Another Rotation instance to compose with
+
+        Returns:
+            Rotation instance representing the composition of rotations
+        """
+        if not isinstance(other, Rotation):
+            return NotImplemented
+
+        # Get quaternions [x, y, z, w]
+        x1, y1, z1, w1 = other._quat  # Apply first
+        x2, y2, z2, w2 = self._quat  # Apply second
+
+        # Quaternion multiplication: q2 * q1 (apply q1 first, then q2)
+        composed_quat = np.array(
+            [
+                w2 * x1 + x2 * w1 + y2 * z1 - z2 * y1,  # x component
+                w2 * y1 - x2 * z1 + y2 * w1 + z2 * x1,  # y component
+                w2 * z1 + x2 * y1 - y2 * x1 + z2 * w1,  # z component
+                w2 * w1 - x2 * x1 - y2 * y1 - z2 * z1,  # w component
+            ]
+        )
+
+        return Rotation(composed_quat)

src/lerobot/utils/train_utils.py

@@ -32,6 +32,7 @@ from lerobot.datasets.utils import load_json, write_json
 from lerobot.optim.optimizers import load_optimizer_state, save_optimizer_state
 from lerobot.optim.schedulers import load_scheduler_state, save_scheduler_state
 from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.processor import PolicyProcessorPipeline
 from lerobot.utils.random_utils import load_rng_state, save_rng_state
 
 
@@ -74,6 +75,8 @@ def save_checkpoint(
     policy: PreTrainedPolicy,
     optimizer: Optimizer,
     scheduler: LRScheduler | None = None,
+    preprocessor: PolicyProcessorPipeline | None = None,
+    postprocessor: PolicyProcessorPipeline | None = None,
 ) -> None:
     """This function creates the following directory structure:
 
@@ -81,7 +84,9 @@ def save_checkpoint(
     ├── pretrained_model/
     │   ├── config.json  # policy config
     │   ├── model.safetensors  # policy weights
-    │   └── train_config.json  # train config
+    │   ├── train_config.json  # train config
+    │   ├── processor.json  # processor config (if preprocessor provided)
+    │   └── step_*.safetensors  # processor state files (if any)
     └── training_state/
         ├── optimizer_param_groups.json  #  optimizer param groups
         ├── optimizer_state.safetensors  # optimizer state
@@ -95,10 +100,15 @@ def save_checkpoint(
         policy (PreTrainedPolicy): The policy to save.
         optimizer (Optimizer | None, optional): The optimizer to save the state from. Defaults to None.
         scheduler (LRScheduler | None, optional): The scheduler to save the state from. Defaults to None.
+        preprocessor: The preprocessor/pipeline to save. Defaults to None.
     """
     pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR
     policy.save_pretrained(pretrained_dir)
     cfg.save_pretrained(pretrained_dir)
+    if preprocessor is not None:
+        preprocessor.save_pretrained(pretrained_dir)
+    if postprocessor is not None:
+        postprocessor.save_pretrained(pretrained_dir)
     save_training_state(checkpoint_dir, step, optimizer, scheduler)
 
 

src/lerobot/utils/visualization_utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numbers
 import os
 from typing import Any
 
@@ -28,19 +29,69 @@ def _init_rerun(session_name: str = "lerobot_control_loop") -> None:
     rr.spawn(memory_limit=memory_limit)
 
 
-def log_rerun_data(observation: dict[str | Any], action: dict[str | Any]):
-    for obs, val in observation.items():
-        if isinstance(val, float):
-            rr.log(f"observation.{obs}", rr.Scalar(val))
-        elif isinstance(val, np.ndarray):
-            if val.ndim == 1:
-                for i, v in enumerate(val):
-                    rr.log(f"observation.{obs}_{i}", rr.Scalar(float(v)))
-            else:
-                rr.log(f"observation.{obs}", rr.Image(val), static=True)
-    for act, val in action.items():
-        if isinstance(val, float):
-            rr.log(f"action.{act}", rr.Scalar(val))
-        elif isinstance(val, np.ndarray):
-            for i, v in enumerate(val):
-                rr.log(f"action.{act}_{i}", rr.Scalar(float(v)))
+def _is_scalar(x):
+    return (
+        isinstance(x, float)
+        or isinstance(x, numbers.Real)
+        or isinstance(x, (np.integer, np.floating))
+        or (isinstance(x, np.ndarray) and x.ndim == 0)
+    )
+
+
+def log_rerun_data(
+    observation: dict[str, Any] | None = None,
+    action: dict[str, Any] | None = None,
+) -> None:
+    """
+    Logs observation and action data to Rerun for real-time visualization.
+
+    This function iterates through the provided observation and action dictionaries and sends their contents
+    to the Rerun viewer. It handles different data types appropriately:
+    - Scalar values (floats, ints) are logged as `rr.Scalar`.
+    - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed
+      from CHW to HWC format and logged as `rr.Image`.
+    - 1D NumPy arrays are logged as a series of individual scalars, with each element indexed.
+    - Other multi-dimensional arrays are flattened and logged as individual scalars.
+
+    Keys are automatically namespaced with "observation." or "action." if not already present.
+
+    Args:
+        observation: An optional dictionary containing observation data to log.
+        action: An optional dictionary containing action data to log.
+    """
+    if observation:
+        for k, v in observation.items():
+            if v is None:
+                continue
+            key = k if str(k).startswith("observation.") else f"observation.{k}"
+
+            if _is_scalar(v):
+                rr.log(key, rr.Scalar(float(v)))
+            elif isinstance(v, np.ndarray):
+                arr = v
+                # Convert CHW -> HWC when needed
+                if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
+                    arr = np.transpose(arr, (1, 2, 0))
+                if arr.ndim == 1:
+                    for i, vi in enumerate(arr):
+                        rr.log(f"{key}_{i}", rr.Scalar(float(vi)))
+                else:
+                    rr.log(key, rr.Image(arr), static=True)
+
+    if action:
+        for k, v in action.items():
+            if v is None:
+                continue
+            key = k if str(k).startswith("action.") else f"action.{k}"
+
+            if _is_scalar(v):
+                rr.log(key, rr.Scalar(float(v)))
+            elif isinstance(v, np.ndarray):
+                if v.ndim == 1:
+                    for i, vi in enumerate(v):
+                        rr.log(f"{key}_{i}", rr.Scalar(float(vi)))
+                else:
+                    # Fall back to flattening higher-dimensional arrays
+                    flat = v.flatten()
+                    for i, vi in enumerate(flat):
+                        rr.log(f"{key}_{i}", rr.Scalar(float(vi)))