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>

tests/processor/test_converters.py

@@ -0,0 +1,196 @@
+import numpy as np
+import pytest
+import torch
+
+from lerobot.processor.converters import (
+    to_dataset_frame,
+    to_output_robot_action,
+    to_transition_robot_observation,
+    to_transition_teleop_action,
+)
+from lerobot.processor.pipeline import TransitionKey
+
+
+def test_to_transition_teleop_action_prefix_and_tensor_conversion():
+    # Scalars, arrays, and "image-like" uint8 arrays are supported
+    img = np.zeros((8, 12, 3), dtype=np.uint8)
+    act = {
+        "ee.x": 0.5,  # scalar to torch tensor
+        "delta": np.array([1.0, 2.0]),  # ndarray to torch tensor
+        "raw_img": img,  # uint8 HWC to passthrough ndarray
+    }
+
+    tr = to_transition_teleop_action(act)
+
+    # Should be an EnvTransition-like dict with ACTION populated
+    assert isinstance(tr, dict)
+    assert TransitionKey.ACTION in tr
+    assert "action.ee.x" in tr[TransitionKey.ACTION]
+    assert "action.delta" in tr[TransitionKey.ACTION]
+    assert "action.raw_img" in tr[TransitionKey.ACTION]
+
+    # Types: scalars/arrays -> torch tensor; images to np.ndarray
+    assert isinstance(tr[TransitionKey.ACTION]["action.ee.x"], torch.Tensor)
+    assert tr[TransitionKey.ACTION]["action.ee.x"].item() == pytest.approx(0.5)
+
+    assert isinstance(tr[TransitionKey.ACTION]["action.delta"], torch.Tensor)
+    assert tr[TransitionKey.ACTION]["action.delta"].shape == (2,)
+    assert torch.allclose(tr[TransitionKey.ACTION]["action.delta"], torch.tensor([1.0, 2.0]))
+
+    assert isinstance(tr[TransitionKey.ACTION]["action.raw_img"], np.ndarray)
+    assert tr[TransitionKey.ACTION]["action.raw_img"].dtype == np.uint8
+    assert tr[TransitionKey.ACTION]["action.raw_img"].shape == (8, 12, 3)
+
+    # Observation is created as empty dict by make_transition
+    assert TransitionKey.OBSERVATION in tr
+    assert isinstance(tr[TransitionKey.OBSERVATION], dict)
+    assert tr[TransitionKey.OBSERVATION] == {}
+
+
+def test_to_transition_robot_observation_state_vs_images_split():
+    # Create an observation with mixed content
+    img = np.full((10, 20, 3), 255, dtype=np.uint8)  # image (uint8 HWC)
+    obs = {
+        "j1.pos": 10.0,  # scalar to state to torch tensor
+        "j2.pos": np.float32(20.0),  # scalar np to state to torch tensor
+        "image_front": img,  # to images passthrough
+        "flag": np.int32(7),  # scalar to state to torch tensor
+        "arr": np.array([1.5, 2.5]),  # vector to state to torch tensor
+    }
+
+    tr = to_transition_robot_observation(obs)
+    assert isinstance(tr, dict)
+    assert TransitionKey.OBSERVATION in tr
+
+    out = tr[TransitionKey.OBSERVATION]
+    # Check state keys are present and converted to tensors
+    for k in ("j1.pos", "j2.pos", "flag", "arr"):
+        key = f"observation.state.{k}"
+        assert key in out
+        v = out[key]
+        if k != "arr":
+            assert isinstance(v, torch.Tensor) and v.ndim == 0
+        else:
+            assert isinstance(v, torch.Tensor) and v.ndim == 1 and v.shape == (2,)
+
+    # Check image present as is
+    assert "observation.images.image_front" in out
+    assert isinstance(out["observation.images.image_front"], np.ndarray)
+    assert out["observation.images.image_front"].dtype == np.uint8
+    assert out["observation.images.image_front"].shape == (10, 20, 3)
+
+    # ACTION should be empty dict by make_transition
+    assert TransitionKey.ACTION in tr
+    assert isinstance(tr[TransitionKey.ACTION], dict)
+    assert tr[TransitionKey.ACTION] == {}
+
+
+def test_to_output_robot_action_strips_prefix_and_filters_pos_keys_only():
+    # Build a transition with mixed action keys
+    tr = {
+        TransitionKey.ACTION: {
+            "action.j1.pos": 11.0,  # keep "j1.pos"
+            "action.gripper.pos": torch.tensor(33.0),  # keep: tensor accepted
+            "action.ee.x": 0.5,  # ignore (doesn't end with .pos)
+            "misc": "ignore_me",  # ignore (no 'action.' prefix)
+        }
+    }
+
+    out = to_output_robot_action(tr)
+    # Only ".pos" keys with "action." prefix are retained and stripped to base names
+    assert set(out.keys()) == {"j1.pos", "gripper.pos"}
+    # Values converted to float
+    assert isinstance(out["j1.pos"], float)
+    assert isinstance(out["gripper.pos"], float)
+    assert out["j1.pos"] == pytest.approx(11.0)
+    assert out["gripper.pos"] == pytest.approx(33.0)
+
+
+def test_to_dataset_frame_merge_and_pack_vectors_and_metadata():
+    # Fabricate dataset features (as stored in dataset.meta["features"])
+    features = {
+        # Action vector: 3 elements in specific order
+        "action": {
+            "dtype": "float32",
+            "shape": (3,),
+            "names": ["j1.pos", "j2.pos", "gripper.pos"],
+        },
+        # Observation state vector: 2 elements
+        "observation.state": {
+            "dtype": "float32",
+            "shape": (2,),
+            "names": ["j1.pos", "j2.pos"],
+        },
+        # Image spec (video/image dtype acceptable)
+        "observation.images.front": {
+            "dtype": "image",
+            "shape": (480, 640, 3),
+            "names": ["h", "w", "c"],
+        },
+    }
+
+    # Build two transitions to be merged: teleop (action) and robot obs (state/images)
+    img = np.random.randint(0, 255, size=(480, 640, 3), dtype=np.uint8)
+
+    teleop_transition = {
+        TransitionKey.OBSERVATION: {},
+        TransitionKey.ACTION: {
+            "action.j1.pos": torch.tensor(1.1),
+            "action.j2.pos": torch.tensor(2.2),
+            # gripper.pos missing → defaults to 0.0
+            "action.ee.x": 0.5,  # ignored, not in features["action"]["names"]
+        },
+        TransitionKey.COMPLEMENTARY_DATA: {
+            "frame_is_pad": True,
+            "task": "Pick cube",
+        },
+    }
+
+    robot_transition = {
+        TransitionKey.OBSERVATION: {
+            "observation.state.j1.pos": torch.tensor(10.0),
+            "observation.state.j2.pos": torch.tensor(20.0),
+            "observation.images.front": img,
+        },
+        TransitionKey.REWARD: torch.tensor(5.0),
+        TransitionKey.DONE: True,
+        TransitionKey.TRUNCATED: False,
+        TransitionKey.INFO: {"note": "ok"},
+    }
+
+    # Directly call the refactored function
+    batch = to_dataset_frame([teleop_transition, robot_transition], features)
+
+    # Images passthrough
+    assert "observation.images.front" in batch
+    assert batch["observation.images.front"].shape == img.shape
+    assert batch["observation.images.front"].dtype == np.uint8
+    assert np.shares_memory(batch["observation.images.front"], img) or np.array_equal(
+        batch["observation.images.front"], img
+    )
+
+    # Observation.state vector
+    assert "observation.state" in batch
+    obs_vec = batch["observation.state"]
+    assert isinstance(obs_vec, np.ndarray) and obs_vec.dtype == np.float32
+    assert obs_vec.shape == (2,)
+    assert obs_vec[0] == pytest.approx(10.0)
+    assert obs_vec[1] == pytest.approx(20.0)
+
+    # Action vector
+    assert "action" in batch
+    act_vec = batch["action"]
+    assert isinstance(act_vec, np.ndarray) and act_vec.dtype == np.float32
+    assert act_vec.shape == (3,)
+    assert act_vec[0] == pytest.approx(1.1)
+    assert act_vec[1] == pytest.approx(2.2)
+    assert act_vec[2] == pytest.approx(0.0)  # default for missing gripper.pos
+
+    # Next.* metadata
+    assert batch["next.reward"] == pytest.approx(5.0)
+    assert batch["next.done"] is True
+    assert batch["next.truncated"] is False
+
+    # Complementary data
+    assert batch["frame_is_pad"] is True
+    assert batch["task"] == "Pick cube"

tests/processor/test_device_processor.py

@@ -288,8 +288,8 @@ def test_serialization_methods():
     assert processor.device == device
 
 
-def test_feature_contract():
-    """Test that feature_contract returns features unchanged."""
+def test_features():
+    """Test that features returns features unchanged."""
     processor = DeviceProcessor(device="cpu")
 
     features = {
@@ -297,7 +297,7 @@ def test_feature_contract():
         "action": PolicyFeature(type=FeatureType.ACTION, shape=(5,)),
     }
 
-    result = processor.feature_contract(features)
+    result = processor.transform_features(features)
     assert result == features
     assert result is features  # Should return the same object
 

tests/processor/test_normalize_processor.py

@@ -621,10 +621,19 @@ def test_serialization_roundtrip(full_stats):
     assert torch.allclose(result1[TransitionKey.ACTION], result2[TransitionKey.ACTION])
 
     # Verify features and norm_map are correctly reconstructed
-    assert new_processor.features.keys() == original_processor.features.keys()
-    for key in new_processor.features:
-        assert new_processor.features[key].type == original_processor.features[key].type
-        assert new_processor.features[key].shape == original_processor.features[key].shape
+    assert (
+        new_processor.transform_features(features).keys()
+        == original_processor.transform_features(features).keys()
+    )
+    for key in new_processor.transform_features(features):
+        assert (
+            new_processor.transform_features(features)[key].type
+            == original_processor.transform_features(features)[key].type
+        )
+        assert (
+            new_processor.transform_features(features)[key].shape
+            == original_processor.transform_features(features)[key].shape
+        )
 
     assert new_processor.norm_map == original_processor.norm_map
 

tests/processor/test_observation_processor.py

@@ -410,13 +410,13 @@ def test_equivalent_with_image_dict():
         torch.testing.assert_close(original_result[key], processor_result[key])
 
 
-def test_image_processor_feature_contract_pixels_to_image(policy_feature_factory):
+def test_image_processor_features_pixels_to_image(policy_feature_factory):
     processor = VanillaObservationProcessor()
     features = {
         "pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
         "keep": policy_feature_factory(FeatureType.ENV, (1,)),
     }
-    out = processor.feature_contract(features.copy())
+    out = processor.transform_features(features.copy())
 
     assert OBS_IMAGE in out and out[OBS_IMAGE] == features["pixels"]
     assert "pixels" not in out
@@ -424,13 +424,13 @@ def test_image_processor_feature_contract_pixels_to_image(policy_feature_factory
     assert_contract_is_typed(out)
 
 
-def test_image_processor_feature_contract_observation_pixels_to_image(policy_feature_factory):
+def test_image_processor_features_observation_pixels_to_image(policy_feature_factory):
     processor = VanillaObservationProcessor()
     features = {
         "observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
         "keep": policy_feature_factory(FeatureType.ENV, (1,)),
     }
-    out = processor.feature_contract(features.copy())
+    out = processor.transform_features(features.copy())
 
     assert OBS_IMAGE in out and out[OBS_IMAGE] == features["observation.pixels"]
     assert "observation.pixels" not in out
@@ -438,7 +438,7 @@ def test_image_processor_feature_contract_observation_pixels_to_image(policy_fea
     assert_contract_is_typed(out)
 
 
-def test_image_processor_feature_contract_multi_camera_and_prefixed(policy_feature_factory):
+def test_image_processor_features_multi_camera_and_prefixed(policy_feature_factory):
     processor = VanillaObservationProcessor()
     features = {
         "pixels.front": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
@@ -446,7 +446,7 @@ def test_image_processor_feature_contract_multi_camera_and_prefixed(policy_featu
         "observation.pixels.rear": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
         "keep": policy_feature_factory(FeatureType.ENV, (7,)),
     }
-    out = processor.feature_contract(features.copy())
+    out = processor.transform_features(features.copy())
 
     assert f"{OBS_IMAGES}.front" in out and out[f"{OBS_IMAGES}.front"] == features["pixels.front"]
     assert f"{OBS_IMAGES}.wrist" in out and out[f"{OBS_IMAGES}.wrist"] == features["pixels.wrist"]
@@ -456,14 +456,14 @@ def test_image_processor_feature_contract_multi_camera_and_prefixed(policy_featu
     assert_contract_is_typed(out)
 
 
-def test_state_processor_feature_contract_environment_and_agent_pos(policy_feature_factory):
+def test_state_processor_features_environment_and_agent_pos(policy_feature_factory):
     processor = VanillaObservationProcessor()
     features = {
         "environment_state": policy_feature_factory(FeatureType.STATE, (3,)),
         "agent_pos": policy_feature_factory(FeatureType.STATE, (7,)),
         "keep": policy_feature_factory(FeatureType.ENV, (1,)),
     }
-    out = processor.feature_contract(features.copy())
+    out = processor.transform_features(features.copy())
 
     assert OBS_ENV_STATE in out and out[OBS_ENV_STATE] == features["environment_state"]
     assert OBS_STATE in out and out[OBS_STATE] == features["agent_pos"]
@@ -472,13 +472,13 @@ def test_state_processor_feature_contract_environment_and_agent_pos(policy_featu
     assert_contract_is_typed(out)
 
 
-def test_state_processor_feature_contract_prefixed_inputs(policy_feature_factory):
+def test_state_processor_features_prefixed_inputs(policy_feature_factory):
     proc = VanillaObservationProcessor()
     features = {
         "observation.environment_state": policy_feature_factory(FeatureType.STATE, (2,)),
         "observation.agent_pos": policy_feature_factory(FeatureType.STATE, (4,)),
     }
-    out = proc.feature_contract(features.copy())
+    out = proc.transform_features(features.copy())
 
     assert OBS_ENV_STATE in out and out[OBS_ENV_STATE] == features["observation.environment_state"]
     assert OBS_STATE in out and out[OBS_STATE] == features["observation.agent_pos"]

tests/processor/test_pipeline.py

@@ -26,6 +26,7 @@ import torch
 import torch.nn as nn
 
 from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
 from lerobot.processor import EnvTransition, ProcessorStepRegistry, RobotProcessor
 from lerobot.processor.pipeline import TransitionKey
 from tests.conftest import assert_contract_is_typed
@@ -90,8 +91,8 @@ class MockStep:
     def reset(self) -> None:
         self.counter = 0
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # We do not test feature_contract here
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        # We do not test features here
         return features
 
 
@@ -112,8 +113,8 @@ class MockStepWithoutOptionalMethods:
 
         return transition
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # We do not test feature_contract here
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        # We do not test features here
         return features
 
 
@@ -168,8 +169,8 @@ class MockStepWithTensorState:
         self.running_mean.zero_()
         self.running_count.zero_()
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # We do not test feature_contract here
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        # We do not test features here
         return features
 
 
@@ -662,8 +663,8 @@ class MockModuleStep(nn.Module):
         self.running_mean.zero_()
         self.counter = 0
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # We do not test feature_contract here
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        # We do not test features here
         return features
 
 
@@ -744,8 +745,8 @@ class MockNonModuleStepWithState:
         self.step_count.zero_()
         self.history.clear()
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # We do not test feature_contract here
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        # We do not test features here
         return features
 
 
@@ -799,8 +800,8 @@ class MockStepWithNonSerializableParam:
     def reset(self) -> None:
         pass
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # We do not test feature_contract here
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        # We do not test features here
         return features
 
 
@@ -838,8 +839,8 @@ class RegisteredMockStep:
     def reset(self) -> None:
         pass
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # We do not test feature_contract here
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        # We do not test features here
         return features
 
 
@@ -1382,8 +1383,8 @@ def test_state_file_naming_with_registry():
         def load_state_dict(self, state):
             self.state_tensor = state["state_tensor"]
 
-        def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-            # We do not test feature_contract here
+        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+            # We do not test features here
             return features
 
     try:
@@ -1439,8 +1440,8 @@ def test_override_with_nested_config():
         def get_config(self):
             return {"name": self.name, "simple_param": self.simple_param, "nested_config": self.nested_config}
 
-        def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-            # We do not test feature_contract here
+        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+            # We do not test features here
             return features
 
     try:
@@ -1531,8 +1532,8 @@ def test_override_with_callables():
         def get_config(self):
             return {"name": self.name}
 
-        def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-            # We do not test feature_contract here
+        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+            # We do not test features here
             return features
 
     try:
@@ -1766,8 +1767,8 @@ def test_override_with_device_strings():
         def load_state_dict(self, state):
             self.buffer = state["buffer"]
 
-        def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-            # We do not test feature_contract here
+        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+            # We do not test features here
             return features
 
     try:
@@ -1860,21 +1861,16 @@ def test_save_load_with_custom_converter_functions():
 
 
 class NonCompliantStep:
-    """Intentionally non-compliant: missing feature_contract."""
+    """Intentionally non-compliant: missing features."""
 
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
 
-def test_construction_rejects_step_without_feature_contract():
-    with pytest.raises(TypeError, match=r"must define feature_contract\(features\) -> dict\[str, Any\]"):
-        RobotProcessor([NonCompliantStep()])
-
-
 class NonCallableStep:
     """Intentionally non-compliant: missing __call__."""
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return features
 
 
@@ -1893,7 +1889,7 @@ class FeatureContractAddStep:
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         features[self.key] = self.value
         return features
 
@@ -1908,7 +1904,7 @@ class FeatureContractMutateStep:
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         features[self.key] = self.fn(features.get(self.key))
         return features
 
@@ -1920,7 +1916,7 @@ class FeatureContractBadReturnStep:
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         return ["not-a-dict"]
 
 
@@ -1933,12 +1929,12 @@ class FeatureContractRemoveStep:
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
-    def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
         features.pop(self.key, None)
         return features
 
 
-def test_feature_contract_orders_and_merges(policy_feature_factory):
+def test_features_orders_and_merges(policy_feature_factory):
     p = RobotProcessor(
         [
             FeatureContractAddStep("a", policy_feature_factory(FeatureType.STATE, (1,))),
@@ -1946,14 +1942,14 @@ def test_feature_contract_orders_and_merges(policy_feature_factory):
             FeatureContractAddStep("b", policy_feature_factory(FeatureType.ENV, (2,))),
         ]
     )
-    out = p.feature_contract({})
+    out = p.transform_features({})
 
     assert out["a"].type == FeatureType.STATE and out["a"].shape == (3,)
     assert out["b"].type == FeatureType.ENV and out["b"].shape == (2,)
     assert_contract_is_typed(out)
 
 
-def test_feature_contract_respects_initial_without_mutation(policy_feature_factory):
+def test_features_respects_initial_without_mutation(policy_feature_factory):
     initial = {
         "seed": policy_feature_factory(FeatureType.STATE, (7,)),
         "nested": policy_feature_factory(FeatureType.ENV, (0,)),
@@ -1966,7 +1962,7 @@ def test_feature_contract_respects_initial_without_mutation(policy_feature_facto
             ),
         ]
     )
-    out = p.feature_contract(initial_features=initial)
+    out = p.transform_features(initial_features=initial)
 
     assert out["seed"].shape == (8,)
     assert out["nested"].shape == (5,)
@@ -1977,13 +1973,7 @@ def test_feature_contract_respects_initial_without_mutation(policy_feature_facto
     assert_contract_is_typed(out)
 
 
-def test_feature_contract_type_error_on_bad_step():
-    p = RobotProcessor([FeatureContractAddStep(), FeatureContractBadReturnStep()])
-    with pytest.raises(TypeError, match=r"\w+\.feature_contract must return dict\[str, Any\]"):
-        _ = p.feature_contract({})
-
-
-def test_feature_contract_execution_order_tracking():
+def test_features_execution_order_tracking():
     class Track:
         def __init__(self, label):
             self.label = label
@@ -1991,32 +1981,186 @@ def test_feature_contract_execution_order_tracking():
         def __call__(self, transition: EnvTransition) -> EnvTransition:
             return transition
 
-        def feature_contract(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
             code = {"A": 1, "B": 2, "C": 3}[self.label]
             pf = features.get("order", PolicyFeature(type=FeatureType.ENV, shape=()))
             features["order"] = PolicyFeature(type=pf.type, shape=pf.shape + (code,))
             return features
 
-    out = RobotProcessor([Track("A"), Track("B"), Track("C")]).feature_contract({})
+    out = RobotProcessor([Track("A"), Track("B"), Track("C")]).transform_features({})
     assert out["order"].shape == (1, 2, 3)
 
 
-def test_feature_contract_remove_key(policy_feature_factory):
+def test_features_remove_key(policy_feature_factory):
     p = RobotProcessor(
         [
             FeatureContractAddStep("a", policy_feature_factory(FeatureType.STATE, (1,))),
             FeatureContractRemoveStep("a"),
         ]
     )
-    out = p.feature_contract({})
+    out = p.transform_features({})
     assert "a" not in out
 
 
-def test_feature_contract_remove_from_initial(policy_feature_factory):
+def test_features_remove_from_initial(policy_feature_factory):
     initial = {
         "keep": policy_feature_factory(FeatureType.STATE, (1,)),
         "drop": policy_feature_factory(FeatureType.STATE, (1,)),
     }
     p = RobotProcessor([FeatureContractRemoveStep("drop")])
-    out = p.feature_contract(initial_features=initial)
+    out = p.transform_features(initial_features=initial)
     assert "drop" not in out and out["keep"] == initial["keep"]
+
+
+@dataclass
+class AddActionEEAndJointFeatures:
+    """Adds both EE and JOINT action features."""
+
+    def __call__(self, tr):
+        return tr
+
+    def transform_features(self, features: dict) -> dict:
+        # EE features
+        features["action.ee.x"] = float
+        features["action.ee.y"] = float
+        # JOINT features
+        features["action.j1.pos"] = float
+        features["action.j2.pos"] = float
+        return features
+
+
+@dataclass
+class AddObservationStateFeatures:
+    """Adds state features (and optionally an image spec to test precedence)."""
+
+    add_front_image: bool = False
+    front_image_shape: tuple = (240, 320, 3)
+
+    def __call__(self, tr):
+        return tr
+
+    def transform_features(self, features: dict) -> dict:
+        # State features (mix EE and a joint state)
+        features["observation.state.ee.x"] = float
+        features["observation.state.j1.pos"] = float
+        if self.add_front_image:
+            features["observation.images.front"] = self.front_image_shape
+        return features
+
+
+def test_aggregate_joint_action_only():
+    rp = RobotProcessor([AddActionEEAndJointFeatures()])
+    initial = {"front": (480, 640, 3)}
+
+    out = aggregate_pipeline_dataset_features(
+        pipeline=rp,
+        initial_features=initial,
+        use_videos=True,
+        patterns=["action.j1.pos", "action.j2.pos"],
+    )
+
+    # Expect only "action" with joint names
+    assert "action" in out and "observation.state" not in out
+    assert out["action"]["dtype"] == "float32"
+    assert set(out["action"]["names"]) == {"j1.pos", "j2.pos"}
+    assert out["action"]["shape"] == (len(out["action"]["names"]),)
+
+
+def test_aggregate_ee_action_and_observation_with_videos():
+    rp = RobotProcessor([AddActionEEAndJointFeatures(), AddObservationStateFeatures()])
+    initial = {"front": (480, 640, 3), "side": (720, 1280, 3)}
+
+    out = aggregate_pipeline_dataset_features(
+        pipeline=rp,
+        initial_features=initial,
+        use_videos=True,
+        patterns=["action.ee", "observation.state"],
+    )
+
+    # Action should pack only EE names
+    assert "action" in out
+    assert set(out["action"]["names"]) == {"ee.x", "ee.y"}
+    assert out["action"]["dtype"] == "float32"
+
+    # Observation state should pack both ee.x and j1.pos as a vector
+    assert "observation.state" in out
+    assert set(out["observation.state"]["names"]) == {"ee.x", "j1.pos"}
+    assert out["observation.state"]["dtype"] == "float32"
+
+    # Cameras from initial_features appear as videos
+    for cam in ("front", "side"):
+        key = f"observation.images.{cam}"
+        assert key in out
+        assert out[key]["dtype"] == "video"
+        assert out[key]["shape"] == initial[cam]
+        assert out[key]["names"] == ["height", "width", "channels"]
+
+
+def test_aggregate_both_action_types():
+    rp = RobotProcessor([AddActionEEAndJointFeatures()])
+    out = aggregate_pipeline_dataset_features(
+        pipeline=rp,
+        initial_features={},
+        use_videos=True,
+        patterns=["action.ee", "action.j1", "action.j2.pos"],
+    )
+
+    assert "action" in out
+    expected = {"ee.x", "ee.y", "j1.pos", "j2.pos"}
+    assert set(out["action"]["names"]) == expected
+    assert out["action"]["shape"] == (len(expected),)
+
+
+def test_aggregate_images_when_use_videos_false():
+    rp = RobotProcessor([AddObservationStateFeatures(add_front_image=True)])
+    initial = {"back": (480, 640, 3)}
+
+    out = aggregate_pipeline_dataset_features(
+        pipeline=rp,
+        initial_features=initial,
+        use_videos=False,  # expect "image" dtype
+        patterns=None,
+    )
+
+    key = "observation.images.back"
+    key_front = "observation.images.front"
+    assert key not in out
+    assert key_front not in out
+
+
+def test_aggregate_images_when_use_videos_true():
+    rp = RobotProcessor([AddObservationStateFeatures(add_front_image=True)])
+    initial = {"back": (480, 640, 3)}
+
+    out = aggregate_pipeline_dataset_features(
+        pipeline=rp,
+        initial_features=initial,
+        use_videos=True,
+        patterns=None,
+    )
+
+    key = "observation.images.front"
+    key_back = "observation.images.back"
+    assert key in out
+    assert key_back in out
+    assert out[key]["dtype"] == "video"
+    assert out[key_back]["dtype"] == "video"
+    assert out[key_back]["shape"] == initial["back"]
+
+
+def test_initial_camera_not_overridden_by_step_image():
+    # Step explicitly sets a different front image shape; initial has another shape.
+    # aggregate_pipeline_dataset_features should keep the step's value (setdefault behavior on initial cams).
+    rp = RobotProcessor([AddObservationStateFeatures(add_front_image=True, front_image_shape=(240, 320, 3))])
+    initial = {"front": (480, 640, 3)}  # should NOT override the step-provided (240, 320, 3)
+
+    out = aggregate_pipeline_dataset_features(
+        pipeline=rp,
+        initial_features=initial,
+        use_videos=True,
+        patterns=["observation.images.front"],
+    )
+
+    key = "observation.images.front"
+    assert key in out
+    assert out[key]["shape"] == (240, 320, 3)  # from the step, not from initial

tests/processor/test_rename_processor.py

@@ -410,7 +410,7 @@ def test_value_types_preserved():
     assert processed_obs["old_list"] == [1, 2, 3]
 
 
-def test_feature_contract_basic_renaming(policy_feature_factory):
+def test_features_basic_renaming(policy_feature_factory):
     processor = RenameProcessor(rename_map={"a": "x", "b": "y"})
     features = {
         "a": policy_feature_factory(FeatureType.STATE, (2,)),
@@ -418,7 +418,7 @@ def test_feature_contract_basic_renaming(policy_feature_factory):
         "c": policy_feature_factory(FeatureType.ENV, (1,)),
     }
 
-    out = processor.feature_contract(features.copy())
+    out = processor.transform_features(features.copy())
 
     # Values preserved and typed
     assert out["x"] == features["a"]
@@ -430,14 +430,14 @@ def test_feature_contract_basic_renaming(policy_feature_factory):
     assert set(features) == {"a", "b", "c"}
 
 
-def test_feature_contract_overlapping_keys(policy_feature_factory):
+def test_features_overlapping_keys(policy_feature_factory):
     # Overlapping renames: both 'a' and 'b' exist. 'a'->'b', 'b'->'c'
     processor = RenameProcessor(rename_map={"a": "b", "b": "c"})
     features = {
         "a": policy_feature_factory(FeatureType.STATE, (1,)),
         "b": policy_feature_factory(FeatureType.STATE, (2,)),
     }
-    out = processor.feature_contract(features)
+    out = processor.transform_features(features)
 
     assert set(out) == {"b", "c"}
     assert out["b"] == features["a"]  # 'a' renamed to'b'
@@ -445,7 +445,7 @@ def test_feature_contract_overlapping_keys(policy_feature_factory):
     assert_contract_is_typed(out)
 
 
-def test_feature_contract_chained_processors(policy_feature_factory):
+def test_features_chained_processors(policy_feature_factory):
     # Chain two rename processors at the contract level
     processor1 = RenameProcessor(rename_map={"pos": "agent_position", "img": "camera_image"})
     processor2 = RenameProcessor(
@@ -458,7 +458,7 @@ def test_feature_contract_chained_processors(policy_feature_factory):
         "img": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)),
         "extra": policy_feature_factory(FeatureType.ENV, (1,)),
     }
-    out = pipeline.feature_contract(initial_features=spec)
+    out = pipeline.transform_features(initial_features=spec)
 
     assert set(out) == {"observation.state", "observation.image", "extra"}
     assert out["observation.state"] == spec["pos"]

tests/processor/test_tokenizer_processor.py

@@ -470,7 +470,7 @@ def test_registry_functionality():
 
 
 @require_package("transformers")
-def test_feature_contract_basic():
+def test_features_basic():
     """Test basic feature contract functionality."""
     mock_tokenizer = MockTokenizer(vocab_size=100)
     processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=128)
@@ -480,7 +480,7 @@ def test_feature_contract_basic():
         "action": PolicyFeature(type=FeatureType.ACTION, shape=(5,)),
     }
 
-    output_features = processor.feature_contract(input_features)
+    output_features = processor.transform_features(input_features)
 
     # Check that original features are preserved
     assert "observation.state" in output_features
@@ -501,13 +501,13 @@ def test_feature_contract_basic():
 
 
 @require_package("transformers")
-def test_feature_contract_with_custom_max_length():
+def test_features_with_custom_max_length():
     """Test feature contract with custom max_length."""
     mock_tokenizer = MockTokenizer(vocab_size=100)
     processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=64)
 
     input_features = {}
-    output_features = processor.feature_contract(input_features)
+    output_features = processor.transform_features(input_features)
 
     # Check that features use correct max_length
     assert f"{OBS_LANGUAGE}.tokens" in output_features
@@ -521,7 +521,7 @@ def test_feature_contract_with_custom_max_length():
 
 
 @require_package("transformers")
-def test_feature_contract_existing_features():
+def test_features_existing_features():
     """Test feature contract when tokenized features already exist."""
     mock_tokenizer = MockTokenizer(vocab_size=100)
     processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=256)
@@ -531,7 +531,7 @@ def test_feature_contract_existing_features():
         f"{OBS_LANGUAGE}.attention_mask": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)),
     }
 
-    output_features = processor.feature_contract(input_features)
+    output_features = processor.transform_features(input_features)
 
     # Should not overwrite existing features
     assert output_features[f"{OBS_LANGUAGE}.tokens"].shape == (100,)  # Original shape preserved