diff --git a/docs/source/hilserl.mdx b/docs/source/hilserl.mdx index dfdf431ac..7a2d6790b 100644 --- a/docs/source/hilserl.mdx +++ b/docs/source/hilserl.mdx @@ -143,27 +143,27 @@ HIL-SERL uses a modular processor pipeline architecture that processes robot obs The environment processor (`env_processor`) handles incoming observations and environment state: -1. **VanillaObservationProcessor**: Converts raw robot observations into standardized format -2. **JointVelocityProcessor** (optional): Adds joint velocity information to observations -3. **MotorCurrentProcessor** (optional): Adds motor current readings to observations +1. **VanillaObservationProcessorStep**: Converts raw robot observations into standardized format +2. **JointVelocityProcessorStep** (optional): Adds joint velocity information to observations +3. **MotorCurrentProcessorStep** (optional): Adds motor current readings to observations 4. **ForwardKinematicsJointsToEE** (optional): Computes end-effector pose from joint positions -5. **ImageCropResizeProcessor** (optional): Crops and resizes camera images -6. **TimeLimitProcessor** (optional): Enforces episode time limits -7. **GripperPenaltyProcessor** (optional): Applies penalties for inappropriate gripper usage -8. **RewardClassifierProcessor** (optional): Automated reward detection using vision models -9. **ToBatchProcessor**: Converts data to batch format for neural network processing -10. **DeviceProcessor**: Moves data to the specified compute device (CPU/GPU) +5. **ImageCropResizeProcessorStep** (optional): Crops and resizes camera images +6. **TimeLimitProcessorStep** (optional): Enforces episode time limits +7. **GripperPenaltyProcessorStep** (optional): Applies penalties for inappropriate gripper usage +8. **RewardClassifierProcessorStep** (optional): Automated reward detection using vision models +9. **AddBatchDimensionProcessorStep**: Converts data to batch format for neural network processing +10. **DeviceProcessorStep**: Moves data to the specified compute device (CPU/GPU) #### Action Processor Pipeline The action processor (`action_processor`) handles outgoing actions and human interventions: -1. **AddTeleopActionAsComplimentaryData**: Captures teleoperator actions for logging -2. **AddTeleopEventsAsInfo**: Records intervention events and episode control signals +1. **AddTeleopActionAsComplimentaryDataStep**: Captures teleoperator actions for logging +2. **AddTeleopEventsAsInfoStep**: Records intervention events and episode control signals 3. **AddRobotObservationAsComplimentaryData**: Stores raw robot state for processing -4. **InterventionActionProcessor**: Handles human interventions and episode termination +4. **InterventionActionProcessorStep**: Handles human interventions and episode termination 5. **Inverse Kinematics Pipeline** (when enabled): - - **MapDeltaActionToRobotAction**: Converts delta actions to robot action format + - **MapDeltaActionToRobotActionStep**: Converts delta actions to robot action format - **EEReferenceAndDelta**: Computes end-effector reference and delta movements - **EEBoundsAndSafety**: Enforces workspace safety bounds - **InverseKinematicsEEToJoints**: Converts end-effector actions to joint targets diff --git a/src/lerobot/policies/act/processor_act.py b/src/lerobot/policies/act/processor_act.py index 0d7e466f3..f33aeac96 100644 --- a/src/lerobot/policies/act/processor_act.py +++ b/src/lerobot/policies/act/processor_act.py @@ -18,13 +18,13 @@ import torch from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.act.configuration_act import ACTConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, ProcessorKwargs, - RenameProcessor, - ToBatchProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -40,18 +40,18 @@ def make_act_pre_post_processors( postprocessor_kwargs = {} input_steps = [ - RenameProcessor(rename_map={}), - NormalizerProcessor( + RenameProcessorStep(rename_map={}), + NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), - DeviceProcessor(device=config.device), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor( + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] diff --git a/src/lerobot/policies/diffusion/processor_diffusion.py b/src/lerobot/policies/diffusion/processor_diffusion.py index ba57f6dba..ef47ec8b3 100644 --- a/src/lerobot/policies/diffusion/processor_diffusion.py +++ b/src/lerobot/policies/diffusion/processor_diffusion.py @@ -19,13 +19,13 @@ import torch from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, ProcessorKwargs, - RenameProcessor, - ToBatchProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -41,18 +41,18 @@ def make_diffusion_pre_post_processors( postprocessor_kwargs = {} input_steps = [ - RenameProcessor(rename_map={}), - NormalizerProcessor( + RenameProcessorStep(rename_map={}), + NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), - DeviceProcessor(device=config.device), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor( + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index cc8e762cf..ddffef28f 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -21,17 +21,17 @@ from lerobot.configs.types import PolicyFeature from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( + AddBatchDimensionProcessorStep, ComplementaryDataProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, ProcessorKwargs, ProcessorStep, ProcessorStepRegistry, - RenameProcessor, - ToBatchProcessor, - TokenizerProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + TokenizerProcessorStep, + UnnormalizerProcessorStep, ) @@ -80,26 +80,26 @@ def make_pi0_pre_post_processors( # Add remaining processors input_steps: list[ProcessorStep] = [ - RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one - NormalizerProcessor( + RenameProcessorStep(rename_map={}), # To mimic the same processor as pretrained one + NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma - TokenizerProcessor( + TokenizerProcessorStep( tokenizer_name="google/paligemma-3b-pt-224", max_length=config.tokenizer_max_length, padding_side="right", padding="max_length", ), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps: list[ProcessorStep] = [ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor( + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] diff --git a/src/lerobot/policies/pi0fast/processor_pi0fast.py b/src/lerobot/policies/pi0fast/processor_pi0fast.py index a4c971626..fab103280 100644 --- a/src/lerobot/policies/pi0fast/processor_pi0fast.py +++ b/src/lerobot/policies/pi0fast/processor_pi0fast.py @@ -19,13 +19,13 @@ import torch from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, ProcessorKwargs, - RenameProcessor, - ToBatchProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -41,18 +41,18 @@ def make_pi0fast_pre_post_processors( postprocessor_kwargs = {} input_steps = [ - RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one - NormalizerProcessor( + RenameProcessorStep(rename_map={}), # To mimic the same processor as pretrained one + NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), - DeviceProcessor(device=config.device), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor( + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] diff --git a/src/lerobot/policies/sac/processor_sac.py b/src/lerobot/policies/sac/processor_sac.py index 0632693ba..7c90cc08f 100644 --- a/src/lerobot/policies/sac/processor_sac.py +++ b/src/lerobot/policies/sac/processor_sac.py @@ -20,13 +20,13 @@ import torch from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, ProcessorKwargs, - RenameProcessor, - ToBatchProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -42,18 +42,18 @@ def make_sac_pre_post_processors( postprocessor_kwargs = {} input_steps = [ - RenameProcessor(rename_map={}), - NormalizerProcessor( + RenameProcessorStep(rename_map={}), + NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), - DeviceProcessor(device=config.device), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor( + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] diff --git a/src/lerobot/policies/sac/reward_model/processor_classifier.py b/src/lerobot/policies/sac/reward_model/processor_classifier.py index 0108748a2..c5ffae621 100644 --- a/src/lerobot/policies/sac/reward_model/processor_classifier.py +++ b/src/lerobot/policies/sac/reward_model/processor_classifier.py @@ -18,9 +18,9 @@ import torch from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.processor import ( DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, IdentityProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, ProcessorKwargs, ) @@ -37,15 +37,15 @@ def make_classifier_processor( postprocessor_kwargs = {} input_steps = [ - NormalizerProcessor( + NormalizerProcessorStep( features=config.input_features, norm_map=config.normalization_mapping, stats=dataset_stats ), - NormalizerProcessor( + NormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] - output_steps = [DeviceProcessor(device="cpu"), IdentityProcessorStep()] + output_steps = [DeviceProcessorStep(device="cpu"), IdentityProcessorStep()] return ( DataProcessorPipeline( diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 644593760..3d2be30b2 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -20,16 +20,16 @@ from lerobot.configs.types import PolicyFeature from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, ComplementaryDataProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, ProcessorKwargs, ProcessorStepRegistry, - RenameProcessor, - ToBatchProcessor, - TokenizerProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + TokenizerProcessorStep, + UnnormalizerProcessorStep, ) @@ -45,25 +45,25 @@ def make_smolvla_pre_post_processors( postprocessor_kwargs = {} input_steps = [ - RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one - NormalizerProcessor( + RenameProcessorStep(rename_map={}), # To mimic the same processor as pretrained one + NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), SmolVLANewLineProcessor(), - TokenizerProcessor( + TokenizerProcessorStep( tokenizer_name=config.vlm_model_name, padding=config.pad_language_to, padding_side="right", max_length=config.tokenizer_max_length, ), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor( + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] diff --git a/src/lerobot/policies/tdmpc/processor_tdmpc.py b/src/lerobot/policies/tdmpc/processor_tdmpc.py index 660c66b72..fecddef0e 100644 --- a/src/lerobot/policies/tdmpc/processor_tdmpc.py +++ b/src/lerobot/policies/tdmpc/processor_tdmpc.py @@ -19,13 +19,13 @@ import torch from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, ProcessorKwargs, - RenameProcessor, - ToBatchProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -41,18 +41,18 @@ def make_tdmpc_pre_post_processors( postprocessor_kwargs = {} input_steps = [ - RenameProcessor(rename_map={}), - NormalizerProcessor( + RenameProcessorStep(rename_map={}), + NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), - DeviceProcessor(device=config.device), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor( + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] diff --git a/src/lerobot/policies/vqbet/processor_vqbet.py b/src/lerobot/policies/vqbet/processor_vqbet.py index f92a4e88e..37be43150 100644 --- a/src/lerobot/policies/vqbet/processor_vqbet.py +++ b/src/lerobot/policies/vqbet/processor_vqbet.py @@ -20,13 +20,13 @@ import torch from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, ProcessorKwargs, - RenameProcessor, - ToBatchProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -42,18 +42,18 @@ def make_vqbet_pre_post_processors( postprocessor_kwargs = {} input_steps = [ - RenameProcessor(rename_map={}), # Let the possibility to the user to rename the keys - NormalizerProcessor( + RenameProcessorStep(rename_map={}), # Let the possibility to the user to rename the keys + NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), - DeviceProcessor(device=config.device), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor( + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 4e6c0767b..4832d8a4e 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .batch_processor import ToBatchProcessor +from .batch_processor import AddBatchDimensionProcessorStep from .converters import ( batch_to_transition, create_transition, @@ -23,21 +23,21 @@ from .converters import ( transition_to_dataset_frame, ) from .core import EnvTransition, TransitionKey -from .delta_action_processor import MapDeltaActionToRobotAction, MapTensorToDeltaActionDict -from .device_processor import DeviceProcessor -from .gym_action_processor import Numpy2TorchActionProcessor, Torch2NumpyActionProcessor +from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep +from .device_processor import DeviceProcessorStep +from .gym_action_processor import Numpy2TorchActionProcessorStep, Torch2NumpyActionProcessorStep from .hil_processor import ( - AddTeleopActionAsComplimentaryData, - AddTeleopEventsAsInfo, - GripperPenaltyProcessor, - ImageCropResizeProcessor, - InterventionActionProcessor, - RewardClassifierProcessor, - TimeLimitProcessor, + AddTeleopActionAsComplimentaryDataStep, + AddTeleopEventsAsInfoStep, + GripperPenaltyProcessorStep, + ImageCropResizeProcessorStep, + InterventionActionProcessorStep, + RewardClassifierProcessorStep, + TimeLimitProcessorStep, ) -from .joint_observations_processor import JointVelocityProcessor, MotorCurrentProcessor -from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor, hotswap_stats -from .observation_processor import VanillaObservationProcessor +from .joint_observations_processor import JointVelocityProcessorStep, MotorCurrentProcessorStep +from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep, hotswap_stats +from .observation_processor import VanillaObservationProcessorStep from .pipeline import ( ActionProcessorStep, ComplementaryDataProcessorStep, @@ -52,48 +52,48 @@ from .pipeline import ( RewardProcessorStep, TruncatedProcessorStep, ) -from .rename_processor import RenameProcessor -from .tokenizer_processor import TokenizerProcessor +from .rename_processor import RenameProcessorStep +from .tokenizer_processor import TokenizerProcessorStep __all__ = [ "ActionProcessorStep", - "AddTeleopActionAsComplimentaryData", - "AddTeleopEventsAsInfo", + "AddTeleopActionAsComplimentaryDataStep", + "AddTeleopEventsAsInfoStep", "ComplementaryDataProcessorStep", "batch_to_transition", "create_transition", - "DeviceProcessor", + "DeviceProcessorStep", "DoneProcessorStep", "EnvTransition", - "GripperPenaltyProcessor", + "GripperPenaltyProcessorStep", "hotswap_stats", "IdentityProcessorStep", - "ImageCropResizeProcessor", + "ImageCropResizeProcessorStep", "InfoProcessorStep", - "InterventionActionProcessor", - "JointVelocityProcessor", - "MapDeltaActionToRobotAction", - "MapTensorToDeltaActionDict", + "InterventionActionProcessorStep", + "JointVelocityProcessorStep", + "MapDeltaActionToRobotActionStep", + "MapTensorToDeltaActionDictStep", "merge_transitions", - "MotorCurrentProcessor", - "NormalizerProcessor", - "Numpy2TorchActionProcessor", + "MotorCurrentProcessorStep", + "NormalizerProcessorStep", + "Numpy2TorchActionProcessorStep", "ObservationProcessorStep", "ProcessorKwargs", "ProcessorStep", "ProcessorStepRegistry", - "RenameProcessor", - "RewardClassifierProcessor", + "RenameProcessorStep", + "RewardClassifierProcessorStep", "RewardProcessorStep", "DataProcessorPipeline", - "TimeLimitProcessor", - "ToBatchProcessor", - "TokenizerProcessor", - "Torch2NumpyActionProcessor", + "TimeLimitProcessorStep", + "AddBatchDimensionProcessorStep", + "TokenizerProcessorStep", + "Torch2NumpyActionProcessorStep", "transition_to_batch", "transition_to_dataset_frame", "TransitionKey", "TruncatedProcessorStep", - "UnnormalizerProcessor", - "VanillaObservationProcessor", + "UnnormalizerProcessorStep", + "VanillaObservationProcessorStep", ] diff --git a/src/lerobot/processor/batch_processor.py b/src/lerobot/processor/batch_processor.py index f0858d2eb..02d1dd7e7 100644 --- a/src/lerobot/processor/batch_processor.py +++ b/src/lerobot/processor/batch_processor.py @@ -30,7 +30,7 @@ from .pipeline import ( @dataclass @ProcessorStepRegistry.register(name="to_batch_processor_action") -class ToBatchProcessorAction(ActionProcessorStep): +class AddBatchDimensionActionStep(ActionProcessorStep): """Process action component in-place, adding batch dimension if needed.""" def action(self, action): @@ -45,7 +45,7 @@ class ToBatchProcessorAction(ActionProcessorStep): @dataclass @ProcessorStepRegistry.register(name="to_batch_processor_observation") -class ToBatchProcessorObservation(ObservationProcessorStep): +class AddBatchDimensionObservationStep(ObservationProcessorStep): """Process observation component in-place, adding batch dimensions where needed.""" def observation(self, observation): @@ -74,7 +74,7 @@ class ToBatchProcessorObservation(ObservationProcessorStep): @dataclass @ProcessorStepRegistry.register(name="to_batch_processor_complementary_data") -class ToBatchProcessorComplementaryData(ComplementaryDataProcessorStep): +class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep): """Process complementary data in-place, handling task field batching.""" def complementary_data(self, complementary_data): @@ -103,7 +103,7 @@ class ToBatchProcessorComplementaryData(ComplementaryDataProcessorStep): @dataclass @ProcessorStepRegistry.register(name="to_batch_processor") -class ToBatchProcessor(ProcessorStep): +class AddBatchDimensionProcessorStep(ProcessorStep): """Processor that adds batch dimensions to observations and actions when needed. This processor ensures that observations and actions have proper batch dimensions for model processing: @@ -138,12 +138,14 @@ class ToBatchProcessor(ProcessorStep): ``` """ - to_batch_action_processor: ToBatchProcessorAction = field(default_factory=ToBatchProcessorAction) - to_batch_observation_processor: ToBatchProcessorObservation = field( - default_factory=ToBatchProcessorObservation + to_batch_action_processor: AddBatchDimensionActionStep = field( + default_factory=AddBatchDimensionActionStep ) - to_batch_complementary_data_processor: ToBatchProcessorComplementaryData = field( - default_factory=ToBatchProcessorComplementaryData + to_batch_observation_processor: AddBatchDimensionObservationStep = field( + default_factory=AddBatchDimensionObservationStep + ) + to_batch_complementary_data_processor: AddBatchDimensionComplementaryDataStep = field( + default_factory=AddBatchDimensionComplementaryDataStep ) def __call__(self, transition: EnvTransition) -> EnvTransition: diff --git a/src/lerobot/processor/delta_action_processor.py b/src/lerobot/processor/delta_action_processor.py index 229757b6e..1aeadba2d 100644 --- a/src/lerobot/processor/delta_action_processor.py +++ b/src/lerobot/processor/delta_action_processor.py @@ -26,7 +26,7 @@ from .pipeline import ActionProcessorStep, ProcessorStepRegistry @ProcessorStepRegistry.register("map_tensor_to_delta_action_dict") @dataclass -class MapTensorToDeltaActionDict(ActionProcessorStep): +class MapTensorToDeltaActionDictStep(ActionProcessorStep): """ Map a tensor to a delta action dictionary. """ @@ -58,7 +58,7 @@ class MapTensorToDeltaActionDict(ActionProcessorStep): @ProcessorStepRegistry.register("map_delta_action_to_robot_action") @dataclass -class MapDeltaActionToRobotAction(ActionProcessorStep): +class MapDeltaActionToRobotActionStep(ActionProcessorStep): """ Map delta actions from teleoperators (gamepad, keyboard) to robot target actions for use with inverse kinematics processors. diff --git a/src/lerobot/processor/device_processor.py b/src/lerobot/processor/device_processor.py index 649b63a44..4e2a16979 100644 --- a/src/lerobot/processor/device_processor.py +++ b/src/lerobot/processor/device_processor.py @@ -27,7 +27,7 @@ from .pipeline import ProcessorStep, ProcessorStepRegistry @ProcessorStepRegistry.register("device_processor") @dataclass -class DeviceProcessor(ProcessorStep): +class DeviceProcessorStep(ProcessorStep): """Processes transitions by moving tensors to the specified device and optionally converting float dtypes. This processor ensures that all tensors in the transition are moved to the diff --git a/src/lerobot/processor/gym_action_processor.py b/src/lerobot/processor/gym_action_processor.py index c5962646d..a5e7b1777 100644 --- a/src/lerobot/processor/gym_action_processor.py +++ b/src/lerobot/processor/gym_action_processor.py @@ -24,7 +24,7 @@ from .pipeline import ActionProcessorStep, ProcessorStepRegistry @ProcessorStepRegistry.register("torch2numpy_action_processor") @dataclass -class Torch2NumpyActionProcessor(ActionProcessorStep): +class Torch2NumpyActionProcessorStep(ActionProcessorStep): """Convert PyTorch tensor actions to NumPy arrays.""" squeeze_batch_dim: bool = True @@ -56,7 +56,7 @@ class Torch2NumpyActionProcessor(ActionProcessorStep): @ProcessorStepRegistry.register("numpy2torch_action_processor") @dataclass -class Numpy2TorchActionProcessor(ActionProcessorStep): +class Numpy2TorchActionProcessorStep(ActionProcessorStep): """Convert NumPy array action to PyTorch tensor.""" def action(self, action: np.ndarray) -> torch.Tensor: diff --git a/src/lerobot/processor/hil_processor.py b/src/lerobot/processor/hil_processor.py index ac1b5e62d..3c1570eaf 100644 --- a/src/lerobot/processor/hil_processor.py +++ b/src/lerobot/processor/hil_processor.py @@ -29,7 +29,7 @@ TELEOP_ACTION_KEY = "teleop_action" @ProcessorStepRegistry.register("add_teleop_action_as_complementary_data") @dataclass -class AddTeleopActionAsComplimentaryData(ComplementaryDataProcessorStep): +class AddTeleopActionAsComplimentaryDataStep(ComplementaryDataProcessorStep): """Add teleoperator action to transition complementary data.""" teleop_device: Teleoperator @@ -45,7 +45,7 @@ class AddTeleopActionAsComplimentaryData(ComplementaryDataProcessorStep): @ProcessorStepRegistry.register("add_teleop_action_as_info") @dataclass -class AddTeleopEventsAsInfo(InfoProcessorStep): +class AddTeleopEventsAsInfoStep(InfoProcessorStep): """Add teleoperator control events to transition info.""" teleop_device: Teleoperator @@ -62,7 +62,7 @@ class AddTeleopEventsAsInfo(InfoProcessorStep): @ProcessorStepRegistry.register("image_crop_resize_processor") @dataclass -class ImageCropResizeProcessor(ObservationProcessorStep): +class ImageCropResizeProcessorStep(ObservationProcessorStep): """Crop and resize image observations.""" crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None @@ -112,7 +112,7 @@ class ImageCropResizeProcessor(ObservationProcessorStep): @dataclass @ProcessorStepRegistry.register("time_limit_processor") -class TimeLimitProcessor(TruncatedProcessorStep): +class TimeLimitProcessorStep(TruncatedProcessorStep): """Track episode steps and enforce time limits.""" max_episode_steps: int @@ -139,7 +139,7 @@ class TimeLimitProcessor(TruncatedProcessorStep): @dataclass @ProcessorStepRegistry.register("gripper_penalty_processor") -class GripperPenaltyProcessor(ComplementaryDataProcessorStep): +class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep): """Apply penalty for inappropriate gripper usage.""" penalty: float = -0.01 @@ -188,7 +188,7 @@ class GripperPenaltyProcessor(ComplementaryDataProcessorStep): @dataclass @ProcessorStepRegistry.register("intervention_action_processor") -class InterventionActionProcessor(ProcessorStep): +class InterventionActionProcessorStep(ProcessorStep): """Handle human intervention actions and episode termination.""" use_gripper: bool = False @@ -261,7 +261,7 @@ class InterventionActionProcessor(ProcessorStep): @dataclass @ProcessorStepRegistry.register("reward_classifier_processor") -class RewardClassifierProcessor(ProcessorStep): +class RewardClassifierProcessorStep(ProcessorStep): """Apply reward classification to image observations.""" pretrained_path: str | None = None diff --git a/src/lerobot/processor/joint_observations_processor.py b/src/lerobot/processor/joint_observations_processor.py index 8f652b6f8..dc91d725e 100644 --- a/src/lerobot/processor/joint_observations_processor.py +++ b/src/lerobot/processor/joint_observations_processor.py @@ -14,7 +14,7 @@ from lerobot.robots import Robot @dataclass @ProcessorStepRegistry.register("joint_velocity_processor") -class JointVelocityProcessor(ObservationProcessorStep): +class JointVelocityProcessorStep(ObservationProcessorStep): """Add joint velocity information to observations.""" dt: float = 0.1 @@ -67,7 +67,7 @@ class JointVelocityProcessor(ObservationProcessorStep): @dataclass @ProcessorStepRegistry.register("current_processor") -class MotorCurrentProcessor(ObservationProcessorStep): +class MotorCurrentProcessorStep(ObservationProcessorStep): """Add motor current information to observations.""" robot: Robot | None = None diff --git a/src/lerobot/processor/migrate_policy_normalization.py b/src/lerobot/processor/migrate_policy_normalization.py index 1530f0aaa..84eac12af 100644 --- a/src/lerobot/processor/migrate_policy_normalization.py +++ b/src/lerobot/processor/migrate_policy_normalization.py @@ -47,11 +47,11 @@ from safetensors.torch import load_file as load_safetensors from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature -from .batch_processor import ToBatchProcessor -from .device_processor import DeviceProcessor -from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor +from .batch_processor import AddBatchDimensionProcessorStep +from .device_processor import DeviceProcessorStep +from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep from .pipeline import DataProcessorPipeline -from .rename_processor import RenameProcessor +from .rename_processor import RenameProcessorStep # Policy type to class mapping POLICY_CLASSES = { @@ -404,8 +404,8 @@ def main(): # Now create preprocessor and postprocessor with cleaned_config available print("Creating preprocessor and postprocessor...") # The pattern from existing processor factories: - # - Preprocessor has two NormalizerProcessors: one for input_features, one for output_features - # - Postprocessor has one UnnormalizerProcessor for output_features only + # - Preprocessor has two NormalizerProcessorSteps: one for input_features, one for output_features + # - Postprocessor has one UnnormalizerProcessorStep for output_features only # Get features from cleaned_config (now they're PolicyFeature objects) input_features = cleaned_config.get("input_features", {}) @@ -413,21 +413,21 @@ def main(): # Create preprocessor with two normalizers (following the pattern from processor factories) preprocessor_steps = [ - RenameProcessor(rename_map={}), - NormalizerProcessor( + RenameProcessorStep(rename_map={}), + NormalizerProcessorStep( features={**input_features, **output_features}, norm_map=norm_map, stats=stats, ), - ToBatchProcessor(), - DeviceProcessor(device=policy_config.device), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=policy_config.device), ] preprocessor = DataProcessorPipeline(steps=preprocessor_steps, name="robot_preprocessor") # Create postprocessor with unnormalizer for outputs only postprocessor_steps = [ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor(features=output_features, norm_map=norm_map, stats=stats), + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep(features=output_features, norm_map=norm_map, stats=stats), ] postprocessor = DataProcessorPipeline(steps=postprocessor_steps, name="robot_postprocessor") diff --git a/src/lerobot/processor/normalize_processor.py b/src/lerobot/processor/normalize_processor.py index f962ca18f..205efe83a 100644 --- a/src/lerobot/processor/normalize_processor.py +++ b/src/lerobot/processor/normalize_processor.py @@ -161,7 +161,7 @@ class _NormalizationMixin: @dataclass @ProcessorStepRegistry.register(name="normalizer_processor") -class NormalizerProcessor(_NormalizationMixin, ProcessorStep): +class NormalizerProcessorStep(_NormalizationMixin, ProcessorStep): """ A processor that applies normalization to observations and actions in a transition. @@ -180,7 +180,7 @@ class NormalizerProcessor(_NormalizationMixin, ProcessorStep): normalize_observation_keys: set[str] | None = None, eps: float = 1e-8, device: torch.device | str | None = None, - ) -> NormalizerProcessor: + ) -> NormalizerProcessorStep: return cls( features=features, norm_map=norm_map, @@ -213,7 +213,7 @@ class NormalizerProcessor(_NormalizationMixin, ProcessorStep): @dataclass @ProcessorStepRegistry.register(name="unnormalizer_processor") -class UnnormalizerProcessor(_NormalizationMixin, ProcessorStep): +class UnnormalizerProcessorStep(_NormalizationMixin, ProcessorStep): """ A processor that applies unnormalization (the inverse of normalization) to observations and actions in a transition. @@ -230,7 +230,7 @@ class UnnormalizerProcessor(_NormalizationMixin, ProcessorStep): norm_map: dict[FeatureType, NormalizationMode], *, device: torch.device | str | None = None, - ) -> UnnormalizerProcessor: + ) -> UnnormalizerProcessorStep: return cls(features=features, norm_map=norm_map, stats=dataset.meta.stats, device=device) def __call__(self, transition: EnvTransition) -> EnvTransition: @@ -259,7 +259,7 @@ def hotswap_stats( Replaces normalization statistics in a RobotProcessor pipeline. This function creates a deep copy of the provided `RobotProcessor` and updates the - statistics of any `NormalizerProcessor` or `UnnormalizerProcessor` steps within it. + statistics of any `NormalizerProcessorStep` or `UnnormalizerProcessorStep` steps within it. It's useful for adapting a trained policy to a new environment or dataset with different data distributions. """ diff --git a/src/lerobot/processor/observation_processor.py b/src/lerobot/processor/observation_processor.py index 902377304..b5502f769 100644 --- a/src/lerobot/processor/observation_processor.py +++ b/src/lerobot/processor/observation_processor.py @@ -28,7 +28,7 @@ from .pipeline import ObservationProcessorStep, ProcessorStepRegistry @dataclass @ProcessorStepRegistry.register(name="observation_processor") -class VanillaObservationProcessor(ObservationProcessorStep): +class VanillaObservationProcessorStep(ObservationProcessorStep): """ Processes environment observations into the LeRobot format by handling both images and states. diff --git a/src/lerobot/processor/rename_processor.py b/src/lerobot/processor/rename_processor.py index 8c80826e9..215365803 100644 --- a/src/lerobot/processor/rename_processor.py +++ b/src/lerobot/processor/rename_processor.py @@ -24,7 +24,7 @@ from .pipeline import ObservationProcessorStep, ProcessorStepRegistry @dataclass @ProcessorStepRegistry.register(name="rename_processor") -class RenameProcessor(ObservationProcessorStep): +class RenameProcessorStep(ObservationProcessorStep): """Rename processor that renames keys in the observation.""" rename_map: dict[str, str] = field(default_factory=dict) diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py index 0d1dfa761..0f54644b4 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -24,7 +24,7 @@ else: @dataclass @ProcessorStepRegistry.register(name="tokenizer_processor") -class TokenizerProcessor(ObservationProcessorStep): +class TokenizerProcessorStep(ObservationProcessorStep): """Tokenizes text tasks in complementary data using a huggingface tokenizer. This processor handles tokenization of task strings found in the complementary_data @@ -48,7 +48,7 @@ class TokenizerProcessor(ObservationProcessorStep): Examples: Using tokenizer name (auto-loaded): ```python - processor = TokenizerProcessor(tokenizer_name="bert-base-uncased", max_length=128) + processor = TokenizerProcessorStep(tokenizer_name="bert-base-uncased", max_length=128) ``` Using custom tokenizer object: @@ -56,7 +56,7 @@ class TokenizerProcessor(ObservationProcessorStep): from transformers import AutoTokenizer custom_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") - processor = TokenizerProcessor(tokenizer=custom_tokenizer, max_length=128) + processor = TokenizerProcessorStep(tokenizer=custom_tokenizer, max_length=128) ``` """ @@ -76,7 +76,7 @@ class TokenizerProcessor(ObservationProcessorStep): if not _transformers_available: raise ImportError( "The 'transformers' library is not installed. " - "Please install it with `pip install 'lerobot[transformers-dep]'` to use TokenizerProcessor." + "Please install it with `pip install 'lerobot[transformers-dep]'` to use TokenizerProcessorStep." ) if self.tokenizer is not None: diff --git a/src/lerobot/scripts/rl/gym_manipulator.py b/src/lerobot/scripts/rl/gym_manipulator.py index ccc9be2de..5a2111a7b 100644 --- a/src/lerobot/scripts/rl/gym_manipulator.py +++ b/src/lerobot/scripts/rl/gym_manipulator.py @@ -29,25 +29,25 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.envs.configs import HILSerlRobotEnvConfig from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( - AddTeleopActionAsComplimentaryData, - AddTeleopEventsAsInfo, + AddBatchDimensionProcessorStep, + AddTeleopActionAsComplimentaryDataStep, + AddTeleopEventsAsInfoStep, DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, EnvTransition, - GripperPenaltyProcessor, - ImageCropResizeProcessor, - InterventionActionProcessor, - JointVelocityProcessor, - MapDeltaActionToRobotAction, - MapTensorToDeltaActionDict, - MotorCurrentProcessor, - Numpy2TorchActionProcessor, - RewardClassifierProcessor, - TimeLimitProcessor, - ToBatchProcessor, - Torch2NumpyActionProcessor, + GripperPenaltyProcessorStep, + ImageCropResizeProcessorStep, + InterventionActionProcessorStep, + JointVelocityProcessorStep, + MapDeltaActionToRobotActionStep, + MapTensorToDeltaActionDictStep, + MotorCurrentProcessorStep, + Numpy2TorchActionProcessorStep, + RewardClassifierProcessorStep, + TimeLimitProcessorStep, + Torch2NumpyActionProcessorStep, TransitionKey, - VanillaObservationProcessor, + VanillaObservationProcessorStep, create_transition, ) from lerobot.robots import ( # noqa: F401 @@ -362,16 +362,16 @@ def make_processors( if cfg.name == "gym_hil": action_pipeline_steps = [ - InterventionActionProcessor(terminate_on_success=terminate_on_success), - Torch2NumpyActionProcessor(), + InterventionActionProcessorStep(terminate_on_success=terminate_on_success), + Torch2NumpyActionProcessorStep(), ] # Minimal processor pipeline for GymHIL simulation env_pipeline_steps = [ - Numpy2TorchActionProcessor(), - VanillaObservationProcessor(), - ToBatchProcessor(), - DeviceProcessor(device=device), + Numpy2TorchActionProcessorStep(), + VanillaObservationProcessorStep(), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=device), ] return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline( @@ -391,13 +391,13 @@ def make_processors( joint_names=motor_names, ) - env_pipeline_steps = [VanillaObservationProcessor()] + env_pipeline_steps = [VanillaObservationProcessorStep()] if cfg.processor.observation is not None: if cfg.processor.observation.add_joint_velocity_to_observation: - env_pipeline_steps.append(JointVelocityProcessor(dt=1.0 / cfg.fps)) + env_pipeline_steps.append(JointVelocityProcessorStep(dt=1.0 / cfg.fps)) if cfg.processor.observation.add_current_to_observation: - env_pipeline_steps.append(MotorCurrentProcessor(robot=env.robot)) + env_pipeline_steps.append(MotorCurrentProcessorStep(robot=env.robot)) if kinematics_solver is not None: env_pipeline_steps.append( @@ -409,7 +409,7 @@ def make_processors( if cfg.processor.image_preprocessing is not None: env_pipeline_steps.append( - ImageCropResizeProcessor( + ImageCropResizeProcessorStep( crop_params_dict=cfg.processor.image_preprocessing.crop_params_dict, resize_size=cfg.processor.image_preprocessing.resize_size, ) @@ -418,13 +418,13 @@ def make_processors( # Add time limit processor if reset config exists if cfg.processor.reset is not None: env_pipeline_steps.append( - TimeLimitProcessor(max_episode_steps=int(cfg.processor.reset.control_time_s * cfg.fps)) + TimeLimitProcessorStep(max_episode_steps=int(cfg.processor.reset.control_time_s * cfg.fps)) ) # Add gripper penalty processor if gripper config exists and enabled if cfg.processor.gripper is not None and cfg.processor.gripper.use_gripper: env_pipeline_steps.append( - GripperPenaltyProcessor( + GripperPenaltyProcessorStep( penalty=cfg.processor.gripper.gripper_penalty, max_gripper_pos=cfg.processor.max_gripper_pos, ) @@ -435,7 +435,7 @@ def make_processors( and cfg.processor.reward_classifier.pretrained_path is not None ): env_pipeline_steps.append( - RewardClassifierProcessor( + RewardClassifierProcessorStep( pretrained_path=cfg.processor.reward_classifier.pretrained_path, device=device, success_threshold=cfg.processor.reward_classifier.success_threshold, @@ -444,14 +444,14 @@ def make_processors( ) ) - env_pipeline_steps.append(ToBatchProcessor()) - env_pipeline_steps.append(DeviceProcessor(device=device)) + env_pipeline_steps.append(AddBatchDimensionProcessorStep()) + env_pipeline_steps.append(DeviceProcessorStep(device=device)) action_pipeline_steps = [ - AddTeleopActionAsComplimentaryData(teleop_device=teleop_device), - AddTeleopEventsAsInfo(teleop_device=teleop_device), + AddTeleopActionAsComplimentaryDataStep(teleop_device=teleop_device), + AddTeleopEventsAsInfoStep(teleop_device=teleop_device), AddRobotObservationAsComplimentaryData(robot=env.robot), - InterventionActionProcessor( + InterventionActionProcessorStep( use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False, terminate_on_success=terminate_on_success, ), @@ -461,10 +461,10 @@ def make_processors( if cfg.processor.inverse_kinematics is not None and kinematics_solver is not None: # Add EE bounds and safety processor inverse_kinematics_steps = [ - MapTensorToDeltaActionDict( + MapTensorToDeltaActionDictStep( use_gripper=cfg.processor.gripper.use_gripper if cfg.processor.gripper is not None else False ), - MapDeltaActionToRobotAction(), + MapDeltaActionToRobotActionStep(), EEReferenceAndDelta( kinematics=kinematics_solver, end_effector_step_sizes=cfg.processor.inverse_kinematics.end_effector_step_sizes, diff --git a/tests/processor/test_act_processor.py b/tests/processor/test_act_processor.py index 56cc1b9b4..b102aef80 100644 --- a/tests/processor/test_act_processor.py +++ b/tests/processor/test_act_processor.py @@ -25,13 +25,13 @@ from lerobot.constants import ACTION, OBS_STATE from lerobot.policies.act.configuration_act import ACTConfig from lerobot.policies.act.processor_act import make_act_pre_post_processors from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, - RenameProcessor, - ToBatchProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, + RenameProcessorStep, TransitionKey, - UnnormalizerProcessor, + UnnormalizerProcessorStep, ) @@ -86,15 +86,15 @@ def test_make_act_processor_basic(): # Check steps in preprocessor assert len(preprocessor.steps) == 4 - assert isinstance(preprocessor.steps[0], RenameProcessor) - assert isinstance(preprocessor.steps[1], NormalizerProcessor) - assert isinstance(preprocessor.steps[2], ToBatchProcessor) - assert isinstance(preprocessor.steps[3], DeviceProcessor) + assert isinstance(preprocessor.steps[0], RenameProcessorStep) + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) + assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) + assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor assert len(postprocessor.steps) == 2 - assert isinstance(postprocessor.steps[0], DeviceProcessor) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[0], DeviceProcessorStep) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_act_processor_normalization(): @@ -303,11 +303,11 @@ def test_act_processor_mixed_precision(): postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x}, ) - # Replace DeviceProcessor with one that uses float16 + # Replace DeviceProcessorStep with one that uses float16 modified_steps = [] for step in preprocessor.steps: - if isinstance(step, DeviceProcessor): - modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16")) + if isinstance(step, DeviceProcessorStep): + modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16")) else: modified_steps.append(step) preprocessor.steps = modified_steps diff --git a/tests/processor/test_batch_processor.py b/tests/processor/test_batch_processor.py index 9b3a7fa1e..568bfa5c4 100644 --- a/tests/processor/test_batch_processor.py +++ b/tests/processor/test_batch_processor.py @@ -22,7 +22,12 @@ import pytest import torch from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE -from lerobot.processor import DataProcessorPipeline, ProcessorStepRegistry, ToBatchProcessor, TransitionKey +from lerobot.processor import ( + AddBatchDimensionProcessorStep, + DataProcessorPipeline, + ProcessorStepRegistry, + TransitionKey, +) def create_transition( @@ -42,7 +47,7 @@ def create_transition( def test_state_1d_to_2d(): """Test that 1D state tensors get unsqueezed to 2D.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Test observation.state state_1d = torch.randn(7) @@ -58,7 +63,7 @@ def test_state_1d_to_2d(): def test_env_state_1d_to_2d(): """Test that 1D environment state tensors get unsqueezed to 2D.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Test observation.environment_state env_state_1d = torch.randn(10) @@ -74,7 +79,7 @@ def test_env_state_1d_to_2d(): def test_image_3d_to_4d(): """Test that 3D image tensors get unsqueezed to 4D.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Test observation.image image_3d = torch.randn(224, 224, 3) @@ -90,7 +95,7 @@ def test_image_3d_to_4d(): def test_multiple_images_3d_to_4d(): """Test that 3D image tensors in observation.images.* get unsqueezed to 4D.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Test observation.images.camera1 and observation.images.camera2 image1_3d = torch.randn(64, 64, 3) @@ -115,7 +120,7 @@ def test_multiple_images_3d_to_4d(): def test_already_batched_tensors_unchanged(): """Test that already batched tensors remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create already batched tensors state_2d = torch.randn(1, 7) @@ -141,7 +146,7 @@ def test_already_batched_tensors_unchanged(): def test_higher_dimensional_tensors_unchanged(): """Test that tensors with more dimensions than expected remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create tensors with more dimensions state_3d = torch.randn(2, 7, 5) # More than 1D @@ -164,7 +169,7 @@ def test_higher_dimensional_tensors_unchanged(): def test_non_tensor_values_unchanged(): """Test that non-tensor values in observations remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() observation = { OBS_STATE: [1, 2, 3], # List, not tensor @@ -187,7 +192,7 @@ def test_non_tensor_values_unchanged(): def test_none_observation(): """Test processor handles None observation gracefully.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() transition = create_transition(observation=None) result = processor(transition) @@ -197,7 +202,7 @@ def test_none_observation(): def test_empty_observation(): """Test processor handles empty observation dict.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() observation = {} transition = create_transition(observation=observation) @@ -209,7 +214,7 @@ def test_empty_observation(): def test_mixed_observation(): """Test processor with mixed observation containing various types and dimensions.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() state_1d = torch.randn(5) env_state_2d = torch.randn(1, 8) # Already batched @@ -241,8 +246,8 @@ def test_mixed_observation(): def test_integration_with_robot_processor(): - """Test ToBatchProcessor integration with RobotProcessor.""" - to_batch_processor = ToBatchProcessor() + """Test AddBatchDimensionProcessorStep integration with RobotProcessor.""" + to_batch_processor = AddBatchDimensionProcessorStep() pipeline = DataProcessorPipeline([to_batch_processor], to_transition=lambda x: x, to_output=lambda x: x) # Create unbatched observation @@ -261,7 +266,7 @@ def test_integration_with_robot_processor(): def test_serialization_methods(): """Test get_config, state_dict, load_state_dict, and reset methods.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Test get_config config = processor.get_config() @@ -281,8 +286,8 @@ def test_serialization_methods(): def test_save_and_load_pretrained(): - """Test saving and loading ToBatchProcessor with RobotProcessor.""" - processor = ToBatchProcessor() + """Test saving and loading AddBatchDimensionProcessorStep with RobotProcessor.""" + processor = AddBatchDimensionProcessorStep() pipeline = DataProcessorPipeline( [processor], name="BatchPipeline", to_transition=lambda x: x, to_output=lambda x: x ) @@ -302,7 +307,7 @@ def test_save_and_load_pretrained(): assert loaded_pipeline.name == "BatchPipeline" assert len(loaded_pipeline) == 1 - assert isinstance(loaded_pipeline.steps[0], ToBatchProcessor) + assert isinstance(loaded_pipeline.steps[0], AddBatchDimensionProcessorStep) # Test functionality of loaded processor observation = {OBS_STATE: torch.randn(5)} @@ -313,10 +318,10 @@ def test_save_and_load_pretrained(): def test_registry_functionality(): - """Test that ToBatchProcessor is properly registered.""" + """Test that AddBatchDimensionProcessorStep is properly registered.""" # Check that the processor is registered registered_class = ProcessorStepRegistry.get("to_batch_processor") - assert registered_class is ToBatchProcessor + assert registered_class is AddBatchDimensionProcessorStep # Check that it's in the list of registered processors assert "to_batch_processor" in ProcessorStepRegistry.list() @@ -324,7 +329,7 @@ def test_registry_functionality(): def test_registry_based_save_load(): """Test saving and loading using registry name.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() pipeline = DataProcessorPipeline([processor], to_transition=lambda x: x, to_output=lambda x: x) with tempfile.TemporaryDirectory() as tmp_dir: @@ -350,7 +355,7 @@ def test_registry_based_save_load(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_device_compatibility(): """Test processor works with tensors on different devices.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create tensors on GPU state_1d = torch.randn(7, device="cuda") @@ -374,7 +379,7 @@ def test_device_compatibility(): def test_processor_preserves_other_transition_keys(): """Test that processor only modifies observation and preserves other transition keys.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() action = torch.randn(5) reward = 1.5 @@ -411,7 +416,7 @@ def test_processor_preserves_other_transition_keys(): def test_edge_case_zero_dimensional_tensors(): """Test processor handles 0D tensors (scalars) correctly.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # 0D tensors should not be modified scalar_tensor = torch.tensor(42.0) @@ -433,7 +438,7 @@ def test_edge_case_zero_dimensional_tensors(): # Action-specific tests def test_action_1d_to_2d(): """Test that 1D action tensors get batch dimension added.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create 1D action tensor action_1d = torch.randn(4) @@ -448,7 +453,7 @@ def test_action_1d_to_2d(): def test_action_already_batched(): """Test that already batched action tensors remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Test various batch sizes action_batched_1 = torch.randn(1, 4) @@ -467,7 +472,7 @@ def test_action_already_batched(): def test_action_higher_dimensional(): """Test that higher dimensional action tensors remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # 3D action tensor (e.g., sequence of actions) action_3d = torch.randn(2, 4, 3) @@ -484,7 +489,7 @@ def test_action_higher_dimensional(): def test_action_scalar_tensor(): """Test that scalar (0D) action tensors remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() action_scalar = torch.tensor(1.5) transition = create_transition(action=action_scalar) @@ -497,7 +502,7 @@ def test_action_scalar_tensor(): def test_action_non_tensor(): """Test that non-tensor actions remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # List action action_list = [0.1, 0.2, 0.3, 0.4] @@ -526,7 +531,7 @@ def test_action_non_tensor(): def test_action_none(): """Test that None action is handled correctly.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() transition = create_transition(action=None) result = processor(transition) @@ -535,7 +540,7 @@ def test_action_none(): def test_action_with_observation(): """Test action processing together with observation processing.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Both need batching observation = { @@ -555,7 +560,7 @@ def test_action_with_observation(): def test_action_different_sizes(): """Test action processing with various action dimensions.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Different action sizes (robot with different DOF) action_sizes = [1, 2, 4, 7, 10, 20] @@ -572,7 +577,7 @@ def test_action_different_sizes(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_action_device_compatibility(): """Test action processing on different devices.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # CUDA action action_cuda = torch.randn(4, device="cuda") @@ -593,7 +598,7 @@ def test_action_device_compatibility(): def test_action_dtype_preservation(): """Test that action dtype is preserved during processing.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Different dtypes dtypes = [torch.float32, torch.float64, torch.int32, torch.int64] @@ -609,7 +614,7 @@ def test_action_dtype_preservation(): def test_empty_action_tensor(): """Test handling of empty action tensors.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Empty 1D tensor action_empty = torch.tensor([]) @@ -631,7 +636,7 @@ def test_empty_action_tensor(): # Task-specific tests def test_task_string_to_list(): """Test that string tasks get wrapped in lists to add batch dimension.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create complementary data with string task complementary_data = {"task": "pick_cube"} @@ -648,7 +653,7 @@ def test_task_string_to_list(): def test_task_string_validation(): """Test that only string and list of strings are valid task values.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Valid string task - should be converted to list complementary_data = {"task": "valid_task"} @@ -667,7 +672,7 @@ def test_task_string_validation(): def test_task_list_of_strings(): """Test that lists of strings remain unchanged (already batched).""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Test various list of strings test_lists = [ @@ -693,7 +698,7 @@ def test_task_list_of_strings(): def test_complementary_data_none(): """Test processor handles None complementary_data gracefully.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() transition = create_transition(complementary_data=None) result = processor(transition) @@ -703,7 +708,7 @@ def test_complementary_data_none(): def test_complementary_data_empty(): """Test processor handles empty complementary_data dict.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() complementary_data = {} transition = create_transition(complementary_data=complementary_data) @@ -715,7 +720,7 @@ def test_complementary_data_empty(): def test_complementary_data_no_task(): """Test processor handles complementary_data without task field.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() complementary_data = { "episode_id": 123, @@ -733,7 +738,7 @@ def test_complementary_data_no_task(): def test_complementary_data_mixed(): """Test processor with mixed complementary_data containing task and other fields.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() complementary_data = { "task": "stack_blocks", @@ -758,7 +763,7 @@ def test_complementary_data_mixed(): def test_task_with_observation_and_action(): """Test task processing together with observation and action processing.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # All components need batching observation = { @@ -783,7 +788,7 @@ def test_task_with_observation_and_action(): def test_task_comprehensive_string_cases(): """Test task processing with comprehensive string cases and edge cases.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Test various string formats string_tasks = [ @@ -841,7 +846,7 @@ def test_task_comprehensive_string_cases(): def test_task_preserves_other_keys(): """Test that task processing preserves other keys in complementary_data.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() complementary_data = { "task": "clean_table", @@ -869,7 +874,7 @@ def test_task_preserves_other_keys(): # Index and task_index specific tests def test_index_scalar_to_1d(): """Test that 0D index tensor gets unsqueezed to 1D.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create 0D index tensor (scalar) index_0d = torch.tensor(42, dtype=torch.int64) @@ -886,7 +891,7 @@ def test_index_scalar_to_1d(): def test_task_index_scalar_to_1d(): """Test that 0D task_index tensor gets unsqueezed to 1D.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create 0D task_index tensor (scalar) task_index_0d = torch.tensor(7, dtype=torch.int64) @@ -903,7 +908,7 @@ def test_task_index_scalar_to_1d(): def test_index_and_task_index_together(): """Test processing both index and task_index together.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create 0D tensors for both index_0d = torch.tensor(100, dtype=torch.int64) @@ -933,7 +938,7 @@ def test_index_and_task_index_together(): def test_index_already_batched(): """Test that already batched index tensors remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create already batched tensors index_1d = torch.tensor([42], dtype=torch.int64) @@ -954,7 +959,7 @@ def test_index_already_batched(): def test_task_index_already_batched(): """Test that already batched task_index tensors remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create already batched tensors task_index_1d = torch.tensor([7], dtype=torch.int64) @@ -975,7 +980,7 @@ def test_task_index_already_batched(): def test_index_non_tensor_unchanged(): """Test that non-tensor index values remain unchanged.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() complementary_data = { "index": 42, # Plain int, not tensor @@ -992,7 +997,7 @@ def test_index_non_tensor_unchanged(): def test_index_dtype_preservation(): """Test that index and task_index dtype is preserved during processing.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Test different dtypes dtypes = [torch.int32, torch.int64, torch.long] @@ -1015,7 +1020,7 @@ def test_index_dtype_preservation(): def test_index_with_full_transition(): """Test index/task_index processing with full transition data.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create full transition with all components observation = { @@ -1057,7 +1062,7 @@ def test_index_with_full_transition(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_index_device_compatibility(): """Test processor works with index/task_index tensors on different devices.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Create tensors on GPU index_0d = torch.tensor(42, dtype=torch.int64, device="cuda") @@ -1081,7 +1086,7 @@ def test_index_device_compatibility(): def test_empty_index_tensor(): """Test handling of empty index tensors.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() # Empty 0D tensor doesn't make sense, but test empty 1D index_empty = torch.tensor([], dtype=torch.int64) @@ -1096,7 +1101,7 @@ def test_empty_index_tensor(): def test_action_processing_creates_new_transition(): """Test that the processor creates a new transition object with correctly processed action.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() action = torch.randn(4) transition = create_transition(action=action) @@ -1118,7 +1123,7 @@ def test_action_processing_creates_new_transition(): def test_task_processing_creates_new_transition(): """Test that the processor creates a new transition object with correctly processed task.""" - processor = ToBatchProcessor() + processor = AddBatchDimensionProcessorStep() complementary_data = {"task": "sort_objects"} transition = create_transition(complementary_data=complementary_data) diff --git a/tests/processor/test_classifier_processor.py b/tests/processor/test_classifier_processor.py index 66d750086..1c9118bd1 100644 --- a/tests/processor/test_classifier_processor.py +++ b/tests/processor/test_classifier_processor.py @@ -26,9 +26,9 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardCla from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor from lerobot.processor import ( DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, IdentityProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, TransitionKey, ) @@ -87,13 +87,13 @@ def test_make_classifier_processor_basic(): # Check steps in preprocessor assert len(preprocessor.steps) == 3 - assert isinstance(preprocessor.steps[0], NormalizerProcessor) # For input features - assert isinstance(preprocessor.steps[1], NormalizerProcessor) # For output features - assert isinstance(preprocessor.steps[2], DeviceProcessor) + assert isinstance(preprocessor.steps[0], NormalizerProcessorStep) # For input features + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) # For output features + assert isinstance(preprocessor.steps[2], DeviceProcessorStep) # Check steps in postprocessor assert len(postprocessor.steps) == 2 - assert isinstance(postprocessor.steps[0], DeviceProcessor) + assert isinstance(postprocessor.steps[0], DeviceProcessorStep) assert isinstance(postprocessor.steps[1], IdentityProcessorStep) @@ -291,11 +291,11 @@ def test_classifier_processor_mixed_precision(): # Get the steps from the factory function factory_preprocessor, factory_postprocessor = make_classifier_processor(config, stats) - # Replace DeviceProcessor with one that uses float16 + # Replace DeviceProcessorStep with one that uses float16 modified_steps = [] for step in factory_preprocessor.steps: - if isinstance(step, DeviceProcessor): - modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16")) + if isinstance(step, DeviceProcessorStep): + modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16")) else: modified_steps.append(step) diff --git a/tests/processor/test_device_processor.py b/tests/processor/test_device_processor.py index b9906b720..88bd184fb 100644 --- a/tests/processor/test_device_processor.py +++ b/tests/processor/test_device_processor.py @@ -19,7 +19,7 @@ import pytest import torch from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.processor import DataProcessorPipeline, DeviceProcessor, TransitionKey +from lerobot.processor import DataProcessorPipeline, DeviceProcessorStep, TransitionKey def create_transition( @@ -46,7 +46,7 @@ def create_transition( def test_basic_functionality(): """Test basic device processor functionality on CPU.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") # Create a transition with CPU tensors observation = {"observation.state": torch.randn(10), "observation.image": torch.randn(3, 224, 224)} @@ -73,7 +73,7 @@ def test_basic_functionality(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_cuda_functionality(): """Test device processor functionality on CUDA.""" - processor = DeviceProcessor(device="cuda") + processor = DeviceProcessorStep(device="cuda") # Create a transition with CPU tensors observation = {"observation.state": torch.randn(10), "observation.image": torch.randn(3, 224, 224)} @@ -100,7 +100,7 @@ def test_cuda_functionality(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_specific_cuda_device(): """Test device processor with specific CUDA device.""" - processor = DeviceProcessor(device="cuda:0") + processor = DeviceProcessorStep(device="cuda:0") observation = {"observation.state": torch.randn(10)} action = torch.randn(5) @@ -116,7 +116,7 @@ def test_specific_cuda_device(): def test_non_tensor_values(): """Test that non-tensor values are preserved.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") observation = { "observation.state": torch.randn(10), @@ -142,7 +142,7 @@ def test_non_tensor_values(): def test_none_values(): """Test handling of None values.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") # Test with None observation transition = create_transition(observation=None, action=torch.randn(5)) @@ -159,7 +159,7 @@ def test_none_values(): def test_empty_observation(): """Test handling of empty observation dictionary.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") transition = create_transition(observation={}, action=torch.randn(5)) result = processor(transition) @@ -170,7 +170,7 @@ def test_empty_observation(): def test_scalar_tensors(): """Test handling of scalar tensors.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") observation = {"observation.scalar": torch.tensor(1.5)} action = torch.tensor(2.0) @@ -187,7 +187,7 @@ def test_scalar_tensors(): def test_dtype_preservation(): """Test that tensor dtypes are preserved.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") observation = { "observation.float32": torch.randn(5, dtype=torch.float32), @@ -209,7 +209,7 @@ def test_dtype_preservation(): def test_shape_preservation(): """Test that tensor shapes are preserved.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") observation = { "observation.1d": torch.randn(10), @@ -232,7 +232,7 @@ def test_shape_preservation(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_mixed_devices(): """Test handling of tensors already on different devices.""" - processor = DeviceProcessor(device="cuda") + processor = DeviceProcessorStep(device="cuda") # Create tensors on different devices observation = { @@ -253,22 +253,22 @@ def test_mixed_devices(): def test_non_blocking_flag(): """Test that non_blocking flag is set correctly.""" # CPU processor should have non_blocking=False - cpu_processor = DeviceProcessor(device="cpu") + cpu_processor = DeviceProcessorStep(device="cpu") assert cpu_processor.non_blocking is False if torch.cuda.is_available(): # CUDA processor should have non_blocking=True - cuda_processor = DeviceProcessor(device="cuda") + cuda_processor = DeviceProcessorStep(device="cuda") assert cuda_processor.non_blocking is True - cuda_0_processor = DeviceProcessor(device="cuda:0") + cuda_0_processor = DeviceProcessorStep(device="cuda:0") assert cuda_0_processor.non_blocking is True def test_serialization_methods(): """Test get_config, state_dict, and load_state_dict methods.""" device = "cuda" if torch.cuda.is_available() else "cpu" - processor = DeviceProcessor(device=device) + processor = DeviceProcessorStep(device=device) # Test get_config config = processor.get_config() @@ -289,7 +289,7 @@ def test_serialization_methods(): def test_features(): """Test that features returns features unchanged.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") features = { "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)), @@ -304,11 +304,11 @@ def test_features(): def test_integration_with_robot_processor(): """Test integration with RobotProcessor.""" from lerobot.constants import OBS_STATE - from lerobot.processor import ToBatchProcessor + from lerobot.processor import AddBatchDimensionProcessorStep - # Create a pipeline with DeviceProcessor - device_processor = DeviceProcessor(device="cpu") - batch_processor = ToBatchProcessor() + # Create a pipeline with DeviceProcessorStep + device_processor = DeviceProcessorStep(device="cpu") + batch_processor = AddBatchDimensionProcessorStep() processor = DataProcessorPipeline( steps=[batch_processor, device_processor], @@ -333,9 +333,9 @@ def test_integration_with_robot_processor(): def test_save_and_load_pretrained(): - """Test saving and loading processor with DeviceProcessor.""" + """Test saving and loading processor with DeviceProcessorStep.""" device = "cuda:0" if torch.cuda.is_available() else "cpu" - processor = DeviceProcessor(device=device, float_dtype="float16") + processor = DeviceProcessorStep(device=device, float_dtype="float16") robot_processor = DataProcessorPipeline(steps=[processor], name="device_test_processor") with tempfile.TemporaryDirectory() as tmpdir: @@ -347,7 +347,7 @@ def test_save_and_load_pretrained(): assert len(loaded_processor.steps) == 1 loaded_device_processor = loaded_processor.steps[0] - assert isinstance(loaded_device_processor, DeviceProcessor) + assert isinstance(loaded_device_processor, DeviceProcessorStep) # Use getattr to access attributes safely assert ( getattr(loaded_device_processor, "device", None) == device.split(":")[0] @@ -356,18 +356,18 @@ def test_save_and_load_pretrained(): def test_registry_functionality(): - """Test that DeviceProcessor is properly registered.""" + """Test that DeviceProcessorStep is properly registered.""" from lerobot.processor import ProcessorStepRegistry - # Check that DeviceProcessor is registered + # Check that DeviceProcessorStep is registered registered_class = ProcessorStepRegistry.get("device_processor") - assert registered_class is DeviceProcessor + assert registered_class is DeviceProcessorStep @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_performance_with_large_tensors(): """Test performance with large tensors and non_blocking flag.""" - processor = DeviceProcessor(device="cuda") + processor = DeviceProcessorStep(device="cuda") # Create large tensors observation = { @@ -389,7 +389,7 @@ def test_performance_with_large_tensors(): def test_reward_done_truncated_types(): """Test handling of different types for reward, done, and truncated.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") # Test with scalar values (not tensors) transition = create_transition( @@ -429,7 +429,7 @@ def test_reward_done_truncated_types(): def test_complementary_data_preserved(): """Test that complementary_data is preserved unchanged.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") complementary_data = { "task": "pick_object", @@ -449,13 +449,13 @@ def test_complementary_data_preserved(): assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == "pick_object" assert result[TransitionKey.COMPLEMENTARY_DATA]["episode_id"] == 42 assert result[TransitionKey.COMPLEMENTARY_DATA]["metadata"] == {"sensor": "camera_1"} - # Note: Currently DeviceProcessor doesn't process tensors in complementary_data + # Note: Currently DeviceProcessorStep doesn't process tensors in complementary_data # This is intentional as complementary_data is typically metadata def test_float_dtype_conversion(): """Test float dtype conversion functionality.""" - processor = DeviceProcessor(device="cpu", float_dtype="float16") + processor = DeviceProcessorStep(device="cpu", float_dtype="float16") # Create tensors of different types observation = { @@ -485,7 +485,7 @@ def test_float_dtype_conversion(): def test_float_dtype_none(): """Test that when float_dtype is None, no dtype conversion occurs.""" - processor = DeviceProcessor(device="cpu", float_dtype=None) + processor = DeviceProcessorStep(device="cpu", float_dtype=None) observation = { "observation.float32": torch.randn(5, dtype=torch.float32), @@ -506,7 +506,7 @@ def test_float_dtype_none(): def test_float_dtype_bfloat16(): """Test conversion to bfloat16.""" - processor = DeviceProcessor(device="cpu", float_dtype="bfloat16") + processor = DeviceProcessorStep(device="cpu", float_dtype="bfloat16") observation = {"observation.state": torch.randn(5, dtype=torch.float32)} action = torch.randn(3, dtype=torch.float64) @@ -520,7 +520,7 @@ def test_float_dtype_bfloat16(): def test_float_dtype_float64(): """Test conversion to float64.""" - processor = DeviceProcessor(device="cpu", float_dtype="float64") + processor = DeviceProcessorStep(device="cpu", float_dtype="float64") observation = {"observation.state": torch.randn(5, dtype=torch.float16)} action = torch.randn(3, dtype=torch.float32) @@ -535,27 +535,27 @@ def test_float_dtype_float64(): def test_float_dtype_invalid(): """Test that invalid float_dtype raises ValueError.""" with pytest.raises(ValueError, match="Invalid float_dtype 'invalid_dtype'"): - DeviceProcessor(device="cpu", float_dtype="invalid_dtype") + DeviceProcessorStep(device="cpu", float_dtype="invalid_dtype") def test_float_dtype_aliases(): """Test that dtype aliases work correctly.""" # Test 'half' alias for float16 - processor_half = DeviceProcessor(device="cpu", float_dtype="half") + processor_half = DeviceProcessorStep(device="cpu", float_dtype="half") assert processor_half._target_float_dtype == torch.float16 # Test 'float' alias for float32 - processor_float = DeviceProcessor(device="cpu", float_dtype="float") + processor_float = DeviceProcessorStep(device="cpu", float_dtype="float") assert processor_float._target_float_dtype == torch.float32 # Test 'double' alias for float64 - processor_double = DeviceProcessor(device="cpu", float_dtype="double") + processor_double = DeviceProcessorStep(device="cpu", float_dtype="double") assert processor_double._target_float_dtype == torch.float64 def test_float_dtype_with_mixed_tensors(): """Test float dtype conversion with mixed tensor types.""" - processor = DeviceProcessor(device="cpu", float_dtype="float32") + processor = DeviceProcessorStep(device="cpu", float_dtype="float32") observation = { "observation.image": torch.randint(0, 255, (3, 64, 64), dtype=torch.uint8), # Should not convert @@ -579,13 +579,13 @@ def test_float_dtype_with_mixed_tensors(): def test_float_dtype_serialization(): """Test that float_dtype is properly serialized in get_config.""" device = "cuda" if torch.cuda.is_available() else "cpu" - processor = DeviceProcessor(device=device, float_dtype="float16") + processor = DeviceProcessorStep(device=device, float_dtype="float16") config = processor.get_config() assert config == {"device": device, "float_dtype": "float16"} # Test with None float_dtype - processor_none = DeviceProcessor(device="cpu", float_dtype=None) + processor_none = DeviceProcessorStep(device="cpu", float_dtype=None) config_none = processor_none.get_config() assert config_none == {"device": "cpu", "float_dtype": None} @@ -594,7 +594,7 @@ def test_float_dtype_serialization(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_float_dtype_with_cuda(): """Test float dtype conversion combined with CUDA device.""" - processor = DeviceProcessor(device="cuda", float_dtype="float16") + processor = DeviceProcessorStep(device="cuda", float_dtype="float16") # Create tensors on CPU with different dtypes observation = { @@ -619,7 +619,7 @@ def test_float_dtype_with_cuda(): def test_complementary_data_index_fields(): """Test processing of index and task_index fields in complementary_data.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") # Create transition with index and task_index in complementary_data complementary_data = { @@ -657,7 +657,7 @@ def test_complementary_data_index_fields(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_complementary_data_index_fields_cuda(): """Test moving index and task_index fields to CUDA.""" - processor = DeviceProcessor(device="cuda:0") + processor = DeviceProcessorStep(device="cuda:0") # Create CPU tensors complementary_data = { @@ -679,7 +679,7 @@ def test_complementary_data_index_fields_cuda(): def test_complementary_data_without_index_fields(): """Test that complementary_data without index/task_index fields works correctly.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") complementary_data = { "task": ["navigate"], @@ -697,7 +697,7 @@ def test_complementary_data_without_index_fields(): def test_complementary_data_mixed_tensors(): """Test complementary_data with mix of tensors and non-tensors.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") complementary_data = { "task": ["pick_and_place"], @@ -726,7 +726,7 @@ def test_complementary_data_mixed_tensors(): def test_complementary_data_float_dtype_conversion(): """Test that float dtype conversion doesn't affect int tensors in complementary_data.""" - processor = DeviceProcessor(device="cpu", float_dtype="float16") + processor = DeviceProcessorStep(device="cpu", float_dtype="float16") complementary_data = { "index": torch.tensor([42], dtype=torch.int64), @@ -750,7 +750,7 @@ def test_complementary_data_float_dtype_conversion(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_complementary_data_full_pipeline_cuda(): """Test full transition with complementary_data on CUDA.""" - processor = DeviceProcessor(device="cuda:0", float_dtype="float16") + processor = DeviceProcessorStep(device="cuda:0", float_dtype="float16") # Create full transition with mixed CPU tensors observation = {"observation.state": torch.randn(1, 7, dtype=torch.float32)} @@ -796,7 +796,7 @@ def test_complementary_data_full_pipeline_cuda(): def test_complementary_data_empty(): """Test empty complementary_data handling.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") transition = create_transition( observation={"observation.state": torch.randn(1, 7)}, @@ -811,7 +811,7 @@ def test_complementary_data_empty(): def test_complementary_data_none(): """Test None complementary_data handling.""" - processor = DeviceProcessor(device="cpu") + processor = DeviceProcessorStep(device="cpu") transition = create_transition( observation={"observation.state": torch.randn(1, 7)}, @@ -826,8 +826,8 @@ def test_complementary_data_none(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_preserves_gpu_placement(): - """Test that DeviceProcessor preserves GPU placement when tensor is already on GPU.""" - processor = DeviceProcessor(device="cuda:0") + """Test that DeviceProcessorStep preserves GPU placement when tensor is already on GPU.""" + processor = DeviceProcessorStep(device="cuda:0") # Create tensors already on GPU observation = { @@ -852,9 +852,9 @@ def test_preserves_gpu_placement(): @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") def test_multi_gpu_preservation(): - """Test that DeviceProcessor preserves placement on different GPUs in multi-GPU setup.""" + """Test that DeviceProcessorStep preserves placement on different GPUs in multi-GPU setup.""" # Test 1: GPU-to-GPU preservation (cuda:0 config, cuda:1 input) - processor_gpu = DeviceProcessor(device="cuda:0") + processor_gpu = DeviceProcessorStep(device="cuda:0") # Create tensors on cuda:1 (simulating Accelerate placement) cuda1_device = torch.device("cuda:1") @@ -873,7 +873,7 @@ def test_multi_gpu_preservation(): assert result[TransitionKey.ACTION].device == cuda1_device # Test 2: GPU-to-CPU should move to CPU (not preserve GPU) - processor_cpu = DeviceProcessor(device="cpu") + processor_cpu = DeviceProcessorStep(device="cpu") transition_gpu = create_transition( observation={"observation.state": torch.randn(10).cuda()}, action=torch.randn(5).cuda() @@ -889,7 +889,7 @@ def test_multi_gpu_preservation(): def test_multi_gpu_with_cpu_tensors(): """Test that CPU tensors are moved to configured device even in multi-GPU context.""" # Processor configured for cuda:1 - processor = DeviceProcessor(device="cuda:1") + processor = DeviceProcessorStep(device="cuda:1") # Mix of CPU and GPU tensors observation = { @@ -916,7 +916,7 @@ def test_multi_gpu_with_cpu_tensors(): @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") def test_multi_gpu_with_float_dtype(): """Test float dtype conversion works correctly with multi-GPU preservation.""" - processor = DeviceProcessor(device="cuda:0", float_dtype="float16") + processor = DeviceProcessorStep(device="cuda:0", float_dtype="float16") # Create float tensors on different GPUs observation = { @@ -946,7 +946,7 @@ def test_simulated_accelerate_scenario(): for gpu_id in range(min(torch.cuda.device_count(), 2)): # Each "process" has a processor configured for cuda:0 # but data comes in already placed on the process's GPU - processor = DeviceProcessor(device="cuda:0") + processor = DeviceProcessorStep(device="cuda:0") # Simulate data already placed by Accelerate device = torch.device(f"cuda:{gpu_id}") @@ -966,7 +966,11 @@ def test_policy_processor_integration(): """Test integration with policy processors - input on GPU, output on CPU.""" from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.constants import ACTION, OBS_STATE - from lerobot.processor import NormalizerProcessor, ToBatchProcessor, UnnormalizerProcessor + from lerobot.processor import ( + AddBatchDimensionProcessorStep, + NormalizerProcessorStep, + UnnormalizerProcessorStep, + ) # Create features and stats features = { @@ -984,9 +988,9 @@ def test_policy_processor_integration(): # Create input processor (preprocessor) that moves to GPU input_processor = DataProcessorPipeline( steps=[ - NormalizerProcessor(features=features, norm_map=norm_map, stats=stats), - ToBatchProcessor(), - DeviceProcessor(device="cuda"), + NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device="cuda"), ], name="test_preprocessor", to_transition=lambda x: x, @@ -996,8 +1000,8 @@ def test_policy_processor_integration(): # Create output processor (postprocessor) that moves to CPU output_processor = DataProcessorPipeline( steps=[ - DeviceProcessor(device="cpu"), - UnnormalizerProcessor(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats), + DeviceProcessorStep(device="cpu"), + UnnormalizerProcessorStep(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats), ], name="test_postprocessor", to_transition=lambda x: x, diff --git a/tests/processor/test_diffusion_processor.py b/tests/processor/test_diffusion_processor.py index 0ebca478a..7fcb92e3a 100644 --- a/tests/processor/test_diffusion_processor.py +++ b/tests/processor/test_diffusion_processor.py @@ -25,13 +25,13 @@ from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, - RenameProcessor, - ToBatchProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, + RenameProcessorStep, TransitionKey, - UnnormalizerProcessor, + UnnormalizerProcessorStep, ) @@ -89,15 +89,15 @@ def test_make_diffusion_processor_basic(): # Check steps in preprocessor assert len(preprocessor.steps) == 4 - assert isinstance(preprocessor.steps[0], RenameProcessor) - assert isinstance(preprocessor.steps[1], NormalizerProcessor) - assert isinstance(preprocessor.steps[2], ToBatchProcessor) - assert isinstance(preprocessor.steps[3], DeviceProcessor) + assert isinstance(preprocessor.steps[0], RenameProcessorStep) + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) + assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) + assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor assert len(postprocessor.steps) == 2 - assert isinstance(postprocessor.steps[0], DeviceProcessor) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[0], DeviceProcessorStep) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_diffusion_processor_with_images(): @@ -294,11 +294,11 @@ def test_diffusion_processor_mixed_precision(): # Get the steps from the factory function factory_preprocessor, factory_postprocessor = make_diffusion_pre_post_processors(config, stats) - # Replace DeviceProcessor with one that uses float16 + # Replace DeviceProcessorStep with one that uses float16 modified_steps = [] for step in factory_preprocessor.steps: - if isinstance(step, DeviceProcessor): - modified_steps.append(DeviceProcessor(device=config.device, float_dtype="float16")) + if isinstance(step, DeviceProcessorStep): + modified_steps.append(DeviceProcessorStep(device=config.device, float_dtype="float16")) else: modified_steps.append(step) diff --git a/tests/processor/test_normalize_processor.py b/tests/processor/test_normalize_processor.py index 134445441..f8bf4660e 100644 --- a/tests/processor/test_normalize_processor.py +++ b/tests/processor/test_normalize_processor.py @@ -23,9 +23,9 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.processor import ( DataProcessorPipeline, IdentityProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, TransitionKey, - UnnormalizerProcessor, + UnnormalizerProcessorStep, hotswap_stats, ) from lerobot.processor.converters import to_tensor @@ -125,7 +125,7 @@ def _create_observation_norm_map(): } -# Fixtures for observation normalisation tests using NormalizerProcessor +# Fixtures for observation normalisation tests using NormalizerProcessorStep @pytest.fixture def observation_stats(): return { @@ -142,10 +142,10 @@ def observation_stats(): @pytest.fixture def observation_normalizer(observation_stats): - """Return a NormalizerProcessor that only has observation stats (no action).""" + """Return a NormalizerProcessorStep that only has observation stats (no action).""" features = _create_observation_features() norm_map = _create_observation_norm_map() - return NormalizerProcessor(features=features, norm_map=norm_map, stats=observation_stats) + return NormalizerProcessorStep(features=features, norm_map=norm_map, stats=observation_stats) def test_mean_std_normalization(observation_normalizer): @@ -182,7 +182,7 @@ def test_min_max_normalization(observation_normalizer): def test_selective_normalization(observation_stats): features = _create_observation_features() norm_map = _create_observation_norm_map() - normalizer = NormalizerProcessor( + normalizer = NormalizerProcessorStep( features=features, norm_map=norm_map, stats=observation_stats, @@ -208,7 +208,7 @@ def test_selective_normalization(observation_stats): def test_device_compatibility(observation_stats): features = _create_observation_features() norm_map = _create_observation_norm_map() - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=observation_stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=observation_stats) observation = { "observation.image": torch.tensor([0.7, 0.5, 0.3]).cuda(), } @@ -237,7 +237,7 @@ def test_from_lerobot_dataset(): FeatureType.ACTION: NormalizationMode.MEAN_STD, } - normalizer = NormalizerProcessor.from_lerobot_dataset(mock_dataset, features, norm_map) + normalizer = NormalizerProcessorStep.from_lerobot_dataset(mock_dataset, features, norm_map) # Both observation and action statistics should be present in tensor stats assert "observation.image" in normalizer._tensor_stats @@ -252,7 +252,7 @@ def test_state_dict_save_load(observation_normalizer): # Create new normalizer and load state features = _create_observation_features() norm_map = _create_observation_norm_map() - new_normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats={}) + new_normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats={}) new_normalizer.load_state_dict(state_dict) # Test that it works the same @@ -303,7 +303,7 @@ def _create_action_norm_map_min_max(): def test_mean_std_unnormalization(action_stats_mean_std): features = _create_action_features() norm_map = _create_action_norm_map_mean_std() - unnormalizer = UnnormalizerProcessor( + unnormalizer = UnnormalizerProcessorStep( features=features, norm_map=norm_map, stats={"action": action_stats_mean_std} ) @@ -321,7 +321,7 @@ def test_mean_std_unnormalization(action_stats_mean_std): def test_min_max_unnormalization(action_stats_min_max): features = _create_action_features() norm_map = _create_action_norm_map_min_max() - unnormalizer = UnnormalizerProcessor( + unnormalizer = UnnormalizerProcessorStep( features=features, norm_map=norm_map, stats={"action": action_stats_min_max} ) @@ -347,7 +347,7 @@ def test_min_max_unnormalization(action_stats_min_max): def test_numpy_action_input(action_stats_mean_std): features = _create_action_features() norm_map = _create_action_norm_map_mean_std() - unnormalizer = UnnormalizerProcessor( + unnormalizer = UnnormalizerProcessorStep( features=features, norm_map=norm_map, stats={"action": action_stats_mean_std} ) @@ -365,7 +365,7 @@ def test_numpy_action_input(action_stats_mean_std): def test_none_action(action_stats_mean_std): features = _create_action_features() norm_map = _create_action_norm_map_mean_std() - unnormalizer = UnnormalizerProcessor( + unnormalizer = UnnormalizerProcessorStep( features=features, norm_map=norm_map, stats={"action": action_stats_mean_std} ) @@ -381,11 +381,11 @@ def test_action_from_lerobot_dataset(): mock_dataset.meta.stats = {"action": {"mean": [0.0], "std": [1.0]}} features = {"action": PolicyFeature(FeatureType.ACTION, (1,))} norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD} - unnormalizer = UnnormalizerProcessor.from_lerobot_dataset(mock_dataset, features, norm_map) + unnormalizer = UnnormalizerProcessorStep.from_lerobot_dataset(mock_dataset, features, norm_map) assert "mean" in unnormalizer._tensor_stats["action"] -# Fixtures for NormalizerProcessor tests +# Fixtures for NormalizerProcessorStep tests @pytest.fixture def full_stats(): return { @@ -424,7 +424,7 @@ def _create_full_norm_map(): def normalizer_processor(full_stats): features = _create_full_features() norm_map = _create_full_norm_map() - return NormalizerProcessor(features=features, norm_map=norm_map, stats=full_stats) + return NormalizerProcessorStep(features=features, norm_map=norm_map, stats=full_stats) def test_combined_normalization(normalizer_processor): @@ -468,7 +468,7 @@ def test_processor_from_lerobot_dataset(full_stats): features = _create_full_features() norm_map = _create_full_norm_map() - processor = NormalizerProcessor.from_lerobot_dataset( + processor = NormalizerProcessorStep.from_lerobot_dataset( mock_dataset, features, norm_map, normalize_observation_keys={"observation.image"} ) @@ -480,7 +480,7 @@ def test_processor_from_lerobot_dataset(full_stats): def test_get_config(full_stats): features = _create_full_features() norm_map = _create_full_norm_map() - processor = NormalizerProcessor( + processor = NormalizerProcessorStep( features=features, norm_map=norm_map, stats=full_stats, @@ -539,7 +539,7 @@ def test_empty_observation(): stats = {"observation.image": {"mean": [0.5], "std": [0.2]}} features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) transition = create_transition() result = normalizer(transition) @@ -550,7 +550,7 @@ def test_empty_observation(): def test_empty_stats(): features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats={}) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats={}) observation = {"observation.image": torch.tensor([0.5])} transition = create_transition(observation=observation) @@ -566,7 +566,7 @@ def test_partial_stats(): stats = {"observation.image": {"mean": [0.5]}} # Missing std / (min,max) features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) observation = {"observation.image": torch.tensor([0.7])} transition = create_transition(observation=observation) @@ -581,7 +581,7 @@ def test_missing_action_stats_no_error(): features = {"observation.image": PolicyFeature(FeatureType.VISUAL, (3, 96, 96))} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} - processor = UnnormalizerProcessor.from_lerobot_dataset(mock_dataset, features, norm_map) + processor = UnnormalizerProcessorStep.from_lerobot_dataset(mock_dataset, features, norm_map) # The tensor stats should not contain the 'action' key assert "action" not in processor._tensor_stats @@ -590,7 +590,7 @@ def test_serialization_roundtrip(full_stats): """Test that features and norm_map can be serialized and deserialized correctly.""" features = _create_full_features() norm_map = _create_full_norm_map() - original_processor = NormalizerProcessor( + original_processor = NormalizerProcessorStep( features=features, norm_map=norm_map, stats=full_stats, @@ -602,7 +602,7 @@ def test_serialization_roundtrip(full_stats): config = original_processor.get_config() # Create a new processor from the config (deserialization) - new_processor = NormalizerProcessor( + new_processor = NormalizerProcessorStep( features=config["features"], norm_map=config["norm_map"], stats=full_stats, @@ -670,7 +670,7 @@ def test_identity_normalization_observations(): "observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}, } - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) observation = { "observation.image": torch.tensor([0.7, 0.5, 0.3]), @@ -695,7 +695,7 @@ def test_identity_normalization_actions(): norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY} stats = {"action": {"mean": [0.0, 0.0], "std": [1.0, 2.0]}} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) action = torch.tensor([1.0, -0.5]) transition = create_transition(action=action) @@ -721,7 +721,7 @@ def test_identity_unnormalization_observations(): "observation.state": {"min": [-1.0, -1.0], "max": [1.0, 1.0]}, } - unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) observation = { "observation.image": torch.tensor([0.7, 0.5, 0.3]), @@ -748,7 +748,7 @@ def test_identity_unnormalization_actions(): norm_map = {FeatureType.ACTION: NormalizationMode.IDENTITY} stats = {"action": {"min": [-1.0, -2.0], "max": [1.0, 2.0]}} - unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) action = torch.tensor([0.5, -0.8]) # Normalized values transition = create_transition(action=action) @@ -771,8 +771,8 @@ def test_identity_with_missing_stats(): } stats = {} # No stats provided - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) - unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) + unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])} action = torch.tensor([1.0, -0.5]) @@ -812,7 +812,7 @@ def test_identity_mixed_with_other_modes(): "action": {"min": [-1.0, -1.0], "max": [1.0, 1.0]}, } - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) observation = { "observation.image": torch.tensor([0.7, 0.5, 0.3]), @@ -854,7 +854,7 @@ def test_identity_defaults_when_not_in_norm_map(): "observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}, } - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) observation = { "observation.image": torch.tensor([0.7, 0.5, 0.3]), @@ -888,8 +888,8 @@ def test_identity_roundtrip(): "action": {"min": [-1.0, -1.0], "max": [1.0, 1.0]}, } - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) - unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) + unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) original_observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])} original_action = torch.tensor([0.5, -0.2]) @@ -921,7 +921,7 @@ def test_identity_config_serialization(): "action": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}, } - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) # Get config config = normalizer.get_config() @@ -931,7 +931,7 @@ def test_identity_config_serialization(): assert config["norm_map"]["ACTION"] == "MEAN_STD" # Create new processor from config (simulating load) - new_normalizer = NormalizerProcessor( + new_normalizer = NormalizerProcessorStep( features=config["features"], norm_map=config["norm_map"], stats=stats, @@ -969,7 +969,7 @@ def test_identity_config_serialization(): # norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD} # stats = {"observation.state": {"mean": [0.0, 0.0], "std": [1.0, 1.0]}} -# normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) +# normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) # # Manually inject an invalid mode to test error handling # normalizer.norm_map[FeatureType.STATE] = "INVALID_MODE" @@ -1006,8 +1006,8 @@ def test_hotswap_stats_basic_functionality(): } # Create processors - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) - unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) + unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) identity = IdentityProcessorStep() # Create robot processor @@ -1047,7 +1047,7 @@ def test_hotswap_stats_deep_copy(): } norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) original_processor = DataProcessorPipeline(steps=[normalizer]) # Store reference to original stats @@ -1072,7 +1072,7 @@ def test_hotswap_stats_deep_copy(): def test_hotswap_stats_only_affects_normalizer_steps(): - """Test that hotswap_stats only modifies NormalizerProcessor and UnnormalizerProcessor steps.""" + """Test that hotswap_stats only modifies NormalizerProcessorStep and UnnormalizerProcessorStep steps.""" stats = { "observation.image": {"mean": np.array([0.5]), "std": np.array([0.2])}, } @@ -1087,8 +1087,8 @@ def test_hotswap_stats_only_affects_normalizer_steps(): norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} # Create mixed steps - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) - unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) + unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) identity = IdentityProcessorStep() robot_processor = DataProcessorPipeline(steps=[normalizer, identity, unnormalizer]) @@ -1117,7 +1117,7 @@ def test_hotswap_stats_empty_stats(): } norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) robot_processor = DataProcessorPipeline(steps=[normalizer]) # Hotswap with empty stats @@ -1167,7 +1167,7 @@ def test_hotswap_stats_preserves_other_attributes(): normalize_observation_keys = {"observation.image"} eps = 1e-6 - normalizer = NormalizerProcessor( + normalizer = NormalizerProcessorStep( features=features, norm_map=norm_map, stats=initial_stats, @@ -1212,10 +1212,10 @@ def test_hotswap_stats_multiple_normalizer_types(): } # Create multiple normalizers and unnormalizers - normalizer1 = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) - normalizer2 = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) - unnormalizer1 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) - unnormalizer2 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) + normalizer1 = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) + normalizer2 = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) + unnormalizer1 = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) + unnormalizer2 = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) robot_processor = DataProcessorPipeline(steps=[normalizer1, unnormalizer1, normalizer2, unnormalizer2]) @@ -1264,7 +1264,7 @@ def test_hotswap_stats_with_different_data_types(): FeatureType.ACTION: NormalizationMode.MEAN_STD, } - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) robot_processor = DataProcessorPipeline(steps=[normalizer]) # Hotswap stats @@ -1320,7 +1320,7 @@ def test_hotswap_stats_functional_test(): } # Create original processor - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=initial_stats) original_processor = DataProcessorPipeline( steps=[normalizer], to_transition=lambda x: x, to_output=lambda x: x ) @@ -1366,7 +1366,7 @@ def test_zero_std_uses_eps(): features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))} norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD} stats = {"observation.state": {"mean": np.array([0.5]), "std": np.array([0.0])}} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats, eps=1e-6) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats, eps=1e-6) observation = {"observation.state": torch.tensor([0.5])} # equals mean out = normalizer(create_transition(observation=observation)) @@ -1378,7 +1378,7 @@ def test_min_equals_max_maps_to_minus_one(): features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))} norm_map = {FeatureType.STATE: NormalizationMode.MIN_MAX} stats = {"observation.state": {"min": np.array([2.0]), "max": np.array([2.0])}} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats, eps=1e-6) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats, eps=1e-6) observation = {"observation.state": torch.tensor([2.0])} out = normalizer(create_transition(observation=observation)) @@ -1393,7 +1393,7 @@ def test_action_normalized_despite_normalize_observation_keys(): } norm_map = {FeatureType.STATE: NormalizationMode.IDENTITY, FeatureType.ACTION: NormalizationMode.MEAN_STD} stats = {"action": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}} - normalizer = NormalizerProcessor( + normalizer = NormalizerProcessorStep( features=features, norm_map=norm_map, stats=stats, normalize_observation_keys={"observation.state"} ) @@ -1411,12 +1411,12 @@ def test_unnormalize_observations_mean_std_and_min_max(): "observation.mm": PolicyFeature(FeatureType.STATE, (2,)), } # Build two processors: one mean/std and one min/max - unnorm_ms = UnnormalizerProcessor( + unnorm_ms = UnnormalizerProcessorStep( features={"observation.ms": features["observation.ms"]}, norm_map={FeatureType.STATE: NormalizationMode.MEAN_STD}, stats={"observation.ms": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}}, ) - unnorm_mm = UnnormalizerProcessor( + unnorm_mm = UnnormalizerProcessorStep( features={"observation.mm": features["observation.mm"]}, norm_map={FeatureType.STATE: NormalizationMode.MIN_MAX}, stats={"observation.mm": {"min": np.array([0.0, -2.0]), "max": np.array([2.0, 2.0])}}, @@ -1438,7 +1438,7 @@ def test_unknown_observation_keys_ignored(): features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))} norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD} stats = {"observation.state": {"mean": np.array([0.0]), "std": np.array([1.0])}} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) obs = {"observation.state": torch.tensor([1.0]), "observation.unknown": torch.tensor([5.0])} tr = create_transition(observation=obs) @@ -1452,7 +1452,7 @@ def test_batched_action_normalization(): features = {"action": PolicyFeature(FeatureType.ACTION, (2,))} norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD} stats = {"action": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) actions = torch.tensor([[1.0, -1.0], [3.0, 3.0]]) # first equals mean → zeros; second → [1, 1] out = normalizer(create_transition(action=actions))[TransitionKey.ACTION] @@ -1464,7 +1464,7 @@ def test_complementary_data_preservation(): features = {"observation.state": PolicyFeature(FeatureType.STATE, (1,))} norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD} stats = {"observation.state": {"mean": np.array([0.0]), "std": np.array([1.0])}} - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) comp = {"existing": 123} tr = create_transition(observation={"observation.state": torch.tensor([1.0])}, complementary_data=comp) @@ -1483,8 +1483,8 @@ def test_roundtrip_normalize_unnormalize_non_identity(): "observation.state": {"mean": np.array([1.0, -1.0]), "std": np.array([2.0, 4.0])}, "action": {"min": np.array([-2.0, 0.0]), "max": np.array([2.0, 4.0])}, } - normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) - unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) + normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) + unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats) # Add a time dimension in action for broadcasting check (B,T,D) obs = {"observation.state": torch.tensor([[3.0, 3.0], [1.0, -1.0]])} diff --git a/tests/processor/test_observation_processor.py b/tests/processor/test_observation_processor.py index 0102820d5..ddd705442 100644 --- a/tests/processor/test_observation_processor.py +++ b/tests/processor/test_observation_processor.py @@ -20,7 +20,7 @@ import torch from lerobot.configs.types import FeatureType from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE -from lerobot.processor import TransitionKey, VanillaObservationProcessor +from lerobot.processor import TransitionKey, VanillaObservationProcessorStep from tests.conftest import assert_contract_is_typed @@ -41,7 +41,7 @@ def create_transition( def test_process_single_image(): """Test processing a single image.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Create a mock image (H, W, C) format, uint8 image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) @@ -67,7 +67,7 @@ def test_process_single_image(): def test_process_image_dict(): """Test processing multiple images in a dictionary.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Create mock images image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8) @@ -90,7 +90,7 @@ def test_process_image_dict(): def test_process_batched_image(): """Test processing already batched images.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Create a batched image (B, H, W, C) image = np.random.randint(0, 256, size=(2, 64, 64, 3), dtype=np.uint8) @@ -107,7 +107,7 @@ def test_process_batched_image(): def test_invalid_image_format(): """Test error handling for invalid image formats.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Test wrong channel order (channels first) image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8) @@ -120,7 +120,7 @@ def test_invalid_image_format(): def test_invalid_image_dtype(): """Test error handling for invalid image dtype.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Test wrong dtype image = np.random.rand(64, 64, 3).astype(np.float32) @@ -133,7 +133,7 @@ def test_invalid_image_dtype(): def test_no_pixels_in_observation(): """Test processor when no pixels are in observation.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() observation = {"other_data": np.array([1, 2, 3])} transition = create_transition(observation=observation) @@ -148,7 +148,7 @@ def test_no_pixels_in_observation(): def test_none_observation(): """Test processor with None observation.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() transition = create_transition() result = processor(transition) @@ -158,7 +158,7 @@ def test_none_observation(): def test_serialization_methods(): """Test serialization methods.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Test get_config config = processor.get_config() @@ -177,7 +177,7 @@ def test_serialization_methods(): def test_process_environment_state(): """Test processing environment_state.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32) observation = {"environment_state": env_state} @@ -198,7 +198,7 @@ def test_process_environment_state(): def test_process_agent_pos(): """Test processing agent_pos.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32) observation = {"agent_pos": agent_pos} @@ -219,7 +219,7 @@ def test_process_agent_pos(): def test_process_batched_states(): """Test processing already batched states.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() env_state = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) agent_pos = np.array([[0.5, -0.5], [1.0, -1.0]], dtype=np.float32) @@ -237,7 +237,7 @@ def test_process_batched_states(): def test_process_both_states(): """Test processing both environment_state and agent_pos.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() env_state = np.array([1.0, 2.0], dtype=np.float32) agent_pos = np.array([0.5, -0.5], dtype=np.float32) @@ -262,7 +262,7 @@ def test_process_both_states(): def test_no_states_in_observation(): """Test processor when no states are in observation.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() observation = {"other_data": np.array([1, 2, 3])} transition = create_transition(observation=observation) @@ -276,7 +276,7 @@ def test_no_states_in_observation(): def test_complete_observation_processing(): """Test processing a complete observation with both images and states.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Create mock data image = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8) @@ -313,7 +313,7 @@ def test_complete_observation_processing(): def test_image_only_processing(): """Test processing observation with only images.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) observation = {"pixels": image} @@ -328,7 +328,7 @@ def test_image_only_processing(): def test_state_only_processing(): """Test processing observation with only states.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() agent_pos = np.array([1.0, 2.0], dtype=np.float32) observation = {"agent_pos": agent_pos} @@ -343,7 +343,7 @@ def test_state_only_processing(): def test_empty_observation(): """Test processing empty observation.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() observation = {} transition = create_transition(observation=observation) @@ -359,7 +359,7 @@ def test_equivalent_to_original_function(): # Import the original function for comparison from lerobot.envs.utils import preprocess_observation - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Create test data similar to what the original function expects image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) @@ -386,7 +386,7 @@ def test_equivalent_with_image_dict(): """Test equivalence with dictionary of images.""" from lerobot.envs.utils import preprocess_observation - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Create test data with multiple cameras image1 = np.random.randint(0, 256, size=(32, 32, 3), dtype=np.uint8) @@ -410,7 +410,7 @@ def test_equivalent_with_image_dict(): def test_image_processor_features_pixels_to_image(policy_feature_factory): - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() features = { "pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "keep": policy_feature_factory(FeatureType.ENV, (1,)), @@ -424,7 +424,7 @@ def test_image_processor_features_pixels_to_image(policy_feature_factory): def test_image_processor_features_observation_pixels_to_image(policy_feature_factory): - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() features = { "observation.pixels": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "keep": policy_feature_factory(FeatureType.ENV, (1,)), @@ -438,7 +438,7 @@ def test_image_processor_features_observation_pixels_to_image(policy_feature_fac def test_image_processor_features_multi_camera_and_prefixed(policy_feature_factory): - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() features = { "pixels.front": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), "pixels.wrist": policy_feature_factory(FeatureType.VISUAL, (3, 64, 64)), @@ -456,7 +456,7 @@ def test_image_processor_features_multi_camera_and_prefixed(policy_feature_facto def test_state_processor_features_environment_and_agent_pos(policy_feature_factory): - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() features = { "environment_state": policy_feature_factory(FeatureType.STATE, (3,)), "agent_pos": policy_feature_factory(FeatureType.STATE, (7,)), @@ -472,7 +472,7 @@ def test_state_processor_features_environment_and_agent_pos(policy_feature_facto def test_state_processor_features_prefixed_inputs(policy_feature_factory): - proc = VanillaObservationProcessor() + proc = VanillaObservationProcessorStep() features = { "observation.environment_state": policy_feature_factory(FeatureType.STATE, (2,)), "observation.agent_pos": policy_feature_factory(FeatureType.STATE, (4,)), diff --git a/tests/processor/test_pi0_processor.py b/tests/processor/test_pi0_processor.py index 759704014..589bf209f 100644 --- a/tests/processor/test_pi0_processor.py +++ b/tests/processor/test_pi0_processor.py @@ -25,12 +25,12 @@ from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.policies.pi0.processor_pi0 import Pi0NewLineProcessor, make_pi0_pre_post_processors from lerobot.processor import ( - DeviceProcessor, - NormalizerProcessor, - RenameProcessor, - ToBatchProcessor, + AddBatchDimensionProcessorStep, + DeviceProcessorStep, + NormalizerProcessorStep, + RenameProcessorStep, TransitionKey, - UnnormalizerProcessor, + UnnormalizerProcessorStep, ) @@ -83,7 +83,7 @@ def test_make_pi0_processor_basic(): config = create_default_config() stats = create_default_stats() - with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor"): + with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep"): preprocessor, postprocessor = make_pi0_pre_post_processors( config, stats, @@ -97,17 +97,17 @@ def test_make_pi0_processor_basic(): # Check steps in preprocessor assert len(preprocessor.steps) == 6 - assert isinstance(preprocessor.steps[0], RenameProcessor) - assert isinstance(preprocessor.steps[1], NormalizerProcessor) - assert isinstance(preprocessor.steps[2], ToBatchProcessor) + assert isinstance(preprocessor.steps[0], RenameProcessorStep) + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) + assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], Pi0NewLineProcessor) - # Step 4 would be TokenizerProcessor but it's mocked - assert isinstance(preprocessor.steps[5], DeviceProcessor) + # Step 4 would be TokenizerProcessorStep but it's mocked + assert isinstance(preprocessor.steps[5], DeviceProcessorStep) # Check steps in postprocessor assert len(postprocessor.steps) == 2 - assert isinstance(postprocessor.steps[0], DeviceProcessor) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[0], DeviceProcessorStep) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_pi0_newline_processor_single_task(): @@ -165,7 +165,7 @@ def test_pi0_processor_cuda(): stats = create_default_stats() # Mock the tokenizer processor to act as pass-through - class MockTokenizerProcessor: + class MockTokenizerProcessorStep: def __init__(self, *args, **kwargs): pass @@ -187,7 +187,7 @@ def test_pi0_processor_cuda(): def transform_features(self, features): return features - with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor): + with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep", MockTokenizerProcessorStep): preprocessor, postprocessor = make_pi0_pre_post_processors( config, stats, @@ -220,7 +220,7 @@ def test_pi0_processor_accelerate_scenario(): stats = create_default_stats() # Mock the tokenizer processor to act as pass-through - class MockTokenizerProcessor: + class MockTokenizerProcessorStep: def __init__(self, *args, **kwargs): pass @@ -242,7 +242,7 @@ def test_pi0_processor_accelerate_scenario(): def transform_features(self, features): return features - with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor): + with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep", MockTokenizerProcessorStep): preprocessor, postprocessor = make_pi0_pre_post_processors( config, stats, @@ -276,7 +276,7 @@ def test_pi0_processor_multi_gpu(): stats = create_default_stats() # Mock the tokenizer processor to act as pass-through - class MockTokenizerProcessor: + class MockTokenizerProcessorStep: def __init__(self, *args, **kwargs): pass @@ -298,7 +298,7 @@ def test_pi0_processor_multi_gpu(): def transform_features(self, features): return features - with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor): + with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep", MockTokenizerProcessorStep): preprocessor, postprocessor = make_pi0_pre_post_processors( config, stats, @@ -329,7 +329,7 @@ def test_pi0_processor_without_stats(): config = create_default_config() # Mock the tokenizer processor - with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor"): + with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep"): preprocessor, postprocessor = make_pi0_pre_post_processors( config, dataset_stats=None, diff --git a/tests/processor/test_rename_processor.py b/tests/processor/test_rename_processor.py index a8c29e973..4d345daba 100644 --- a/tests/processor/test_rename_processor.py +++ b/tests/processor/test_rename_processor.py @@ -23,7 +23,7 @@ from lerobot.configs.types import FeatureType from lerobot.processor import ( DataProcessorPipeline, ProcessorStepRegistry, - RenameProcessor, + RenameProcessorStep, TransitionKey, ) from lerobot.processor.rename_processor import rename_stats @@ -51,7 +51,7 @@ def test_basic_renaming(): "old_key1": "new_key1", "old_key2": "new_key2", } - processor = RenameProcessor(rename_map=rename_map) + processor = RenameProcessorStep(rename_map=rename_map) observation = { "old_key1": torch.tensor([1.0, 2.0]), @@ -79,7 +79,7 @@ def test_basic_renaming(): def test_empty_rename_map(): """Test processor with empty rename map (should pass through unchanged).""" - processor = RenameProcessor(rename_map={}) + processor = RenameProcessorStep(rename_map={}) observation = { "key1": torch.tensor([1.0]), @@ -98,7 +98,7 @@ def test_empty_rename_map(): def test_none_observation(): """Test processor with None observation.""" - processor = RenameProcessor(rename_map={"old": "new"}) + processor = RenameProcessorStep(rename_map={"old": "new"}) transition = create_transition() result = processor(transition) @@ -113,7 +113,7 @@ def test_overlapping_rename(): "a": "b", "b": "c", # This creates a potential conflict } - processor = RenameProcessor(rename_map=rename_map) + processor = RenameProcessorStep(rename_map=rename_map) observation = { "a": 1, @@ -138,7 +138,7 @@ def test_partial_rename(): "observation.state": "observation.proprio_state", "pixels": "observation.image", } - processor = RenameProcessor(rename_map=rename_map) + processor = RenameProcessorStep(rename_map=rename_map) observation = { "observation.state": torch.randn(10), @@ -168,15 +168,15 @@ def test_get_config(): "old1": "new1", "old2": "new2", } - processor = RenameProcessor(rename_map=rename_map) + processor = RenameProcessorStep(rename_map=rename_map) config = processor.get_config() assert config == {"rename_map": rename_map} def test_state_dict(): - """Test state dict (should be empty for RenameProcessor).""" - processor = RenameProcessor(rename_map={"old": "new"}) + """Test state dict (should be empty for RenameProcessorStep).""" + processor = RenameProcessorStep(rename_map={"old": "new"}) state = processor.state_dict() assert state == {} @@ -191,7 +191,7 @@ def test_integration_with_robot_processor(): "agent_pos": "observation.state", "pixels": "observation.image", } - rename_processor = RenameProcessor(rename_map=rename_map) + rename_processor = RenameProcessorStep(rename_map=rename_map) pipeline = DataProcessorPipeline([rename_processor], to_transition=lambda x: x, to_output=lambda x: x) @@ -225,18 +225,20 @@ def test_save_and_load_pretrained(): "old_state": "observation.state", "old_image": "observation.image", } - processor = RenameProcessor(rename_map=rename_map) - pipeline = DataProcessorPipeline([processor], name="TestRenameProcessor") + processor = RenameProcessorStep(rename_map=rename_map) + pipeline = DataProcessorPipeline([processor], name="TestRenameProcessorStep") with tempfile.TemporaryDirectory() as tmp_dir: # Save pipeline pipeline.save_pretrained(tmp_dir) # Check files were created - config_path = Path(tmp_dir) / "testrenameprocessor.json" # Based on name="TestRenameProcessor" + config_path = ( + Path(tmp_dir) / "testrenameprocessorstep.json" + ) # Based on name="TestRenameProcessorStep" assert config_path.exists() - # No state files should be created for RenameProcessor + # No state files should be created for RenameProcessorStep state_files = list(Path(tmp_dir).glob("*.safetensors")) assert len(state_files) == 0 @@ -245,12 +247,12 @@ def test_save_and_load_pretrained(): tmp_dir, to_transition=lambda x: x, to_output=lambda x: x ) - assert loaded_pipeline.name == "TestRenameProcessor" + assert loaded_pipeline.name == "TestRenameProcessorStep" assert len(loaded_pipeline) == 1 # Check that loaded processor works correctly loaded_processor = loaded_pipeline.steps[0] - assert isinstance(loaded_processor, RenameProcessor) + assert isinstance(loaded_processor, RenameProcessorStep) assert loaded_processor.rename_map == rename_map # Test functionality after loading @@ -267,23 +269,23 @@ def test_save_and_load_pretrained(): def test_registry_functionality(): - """Test that RenameProcessor is properly registered.""" + """Test that RenameProcessorStep is properly registered.""" # Check that it's registered assert "rename_processor" in ProcessorStepRegistry.list() # Get from registry retrieved_class = ProcessorStepRegistry.get("rename_processor") - assert retrieved_class is RenameProcessor + assert retrieved_class is RenameProcessorStep # Create instance from registry instance = retrieved_class(rename_map={"old": "new"}) - assert isinstance(instance, RenameProcessor) + assert isinstance(instance, RenameProcessorStep) assert instance.rename_map == {"old": "new"} def test_registry_based_save_load(): """Test save/load using registry name instead of module path.""" - processor = RenameProcessor(rename_map={"key1": "renamed_key1"}) + processor = RenameProcessorStep(rename_map={"key1": "renamed_key1"}) pipeline = DataProcessorPipeline([processor], to_transition=lambda x: x, to_output=lambda x: x) with tempfile.TemporaryDirectory() as tmp_dir: @@ -303,14 +305,14 @@ def test_registry_based_save_load(): # Load should work loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) loaded_processor = loaded_pipeline.steps[0] - assert isinstance(loaded_processor, RenameProcessor) + assert isinstance(loaded_processor, RenameProcessorStep) assert loaded_processor.rename_map == {"key1": "renamed_key1"} def test_chained_rename_processors(): - """Test multiple RenameProcessors in a pipeline.""" + """Test multiple RenameProcessorSteps in a pipeline.""" # First processor: rename raw keys to intermediate format - processor1 = RenameProcessor( + processor1 = RenameProcessorStep( rename_map={ "pos": "agent_position", "img": "camera_image", @@ -318,7 +320,7 @@ def test_chained_rename_processors(): ) # Second processor: rename to final format - processor2 = RenameProcessor( + processor2 = RenameProcessorStep( rename_map={ "agent_position": "observation.state", "camera_image": "observation.image", @@ -363,7 +365,7 @@ def test_nested_observation_rename(): "observation.images.right": "observation.camera.right_view", "observation.proprio": "observation.proprioception", } - processor = RenameProcessor(rename_map=rename_map) + processor = RenameProcessorStep(rename_map=rename_map) observation = { "observation.images.left": torch.randn(3, 64, 64), @@ -393,7 +395,7 @@ def test_nested_observation_rename(): def test_value_types_preserved(): """Test that various value types are preserved during renaming.""" rename_map = {"old_tensor": "new_tensor", "old_array": "new_array", "old_scalar": "new_scalar"} - processor = RenameProcessor(rename_map=rename_map) + processor = RenameProcessorStep(rename_map=rename_map) tensor_value = torch.randn(3, 3) array_value = np.random.rand(2, 2) @@ -421,7 +423,7 @@ def test_value_types_preserved(): def test_features_basic_renaming(policy_feature_factory): - processor = RenameProcessor(rename_map={"a": "x", "b": "y"}) + processor = RenameProcessorStep(rename_map={"a": "x", "b": "y"}) features = { "a": policy_feature_factory(FeatureType.STATE, (2,)), "b": policy_feature_factory(FeatureType.ACTION, (3,)), @@ -442,7 +444,7 @@ def test_features_basic_renaming(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"}) + processor = RenameProcessorStep(rename_map={"a": "b", "b": "c"}) features = { "a": policy_feature_factory(FeatureType.STATE, (1,)), "b": policy_feature_factory(FeatureType.STATE, (2,)), @@ -457,8 +459,8 @@ def test_features_overlapping_keys(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( + processor1 = RenameProcessorStep(rename_map={"pos": "agent_position", "img": "camera_image"}) + processor2 = RenameProcessorStep( rename_map={"agent_position": "observation.state", "camera_image": "observation.image"} ) pipeline = DataProcessorPipeline([processor1, processor2]) diff --git a/tests/processor/test_sac_processor.py b/tests/processor/test_sac_processor.py index ca15bd55d..d60bb42a6 100644 --- a/tests/processor/test_sac_processor.py +++ b/tests/processor/test_sac_processor.py @@ -25,13 +25,13 @@ from lerobot.constants import ACTION, OBS_STATE from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, - RenameProcessor, - ToBatchProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, + RenameProcessorStep, TransitionKey, - UnnormalizerProcessor, + UnnormalizerProcessorStep, ) @@ -91,15 +91,15 @@ def test_make_sac_processor_basic(): # Check steps in preprocessor assert len(preprocessor.steps) == 4 - assert isinstance(preprocessor.steps[0], RenameProcessor) - assert isinstance(preprocessor.steps[1], NormalizerProcessor) - assert isinstance(preprocessor.steps[2], ToBatchProcessor) - assert isinstance(preprocessor.steps[3], DeviceProcessor) + assert isinstance(preprocessor.steps[0], RenameProcessorStep) + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) + assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) + assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor assert len(postprocessor.steps) == 2 - assert isinstance(postprocessor.steps[0], DeviceProcessor) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[0], DeviceProcessorStep) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_sac_processor_normalization_modes(): @@ -306,10 +306,10 @@ def test_sac_processor_mixed_precision(): postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x}, ) - # Replace DeviceProcessor with one that uses float16 + # Replace DeviceProcessorStep with one that uses float16 for i, step in enumerate(preprocessor.steps): - if isinstance(step, DeviceProcessor): - preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16") + if isinstance(step, DeviceProcessorStep): + preprocessor.steps[i] = DeviceProcessorStep(device=config.device, float_dtype="float16") # Create test data observation = {OBS_STATE: torch.randn(10, dtype=torch.float32)} diff --git a/tests/processor/test_smolvla_processor.py b/tests/processor/test_smolvla_processor.py index e4d673e62..dc1b1a83f 100644 --- a/tests/processor/test_smolvla_processor.py +++ b/tests/processor/test_smolvla_processor.py @@ -28,12 +28,12 @@ from lerobot.policies.smolvla.processor_smolvla import ( make_smolvla_pre_post_processors, ) from lerobot.processor import ( - DeviceProcessor, - NormalizerProcessor, - RenameProcessor, - ToBatchProcessor, + AddBatchDimensionProcessorStep, + DeviceProcessorStep, + NormalizerProcessorStep, + RenameProcessorStep, TransitionKey, - UnnormalizerProcessor, + UnnormalizerProcessorStep, ) @@ -88,7 +88,7 @@ def test_make_smolvla_processor_basic(): config = create_default_config() stats = create_default_stats() - with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor"): + with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep"): preprocessor, postprocessor = make_smolvla_pre_post_processors( config, stats, @@ -102,17 +102,17 @@ def test_make_smolvla_processor_basic(): # Check steps in preprocessor assert len(preprocessor.steps) == 6 - assert isinstance(preprocessor.steps[0], RenameProcessor) - assert isinstance(preprocessor.steps[1], NormalizerProcessor) - assert isinstance(preprocessor.steps[2], ToBatchProcessor) + assert isinstance(preprocessor.steps[0], RenameProcessorStep) + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) + assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], SmolVLANewLineProcessor) - # Step 4 would be TokenizerProcessor but it's mocked - assert isinstance(preprocessor.steps[5], DeviceProcessor) + # Step 4 would be TokenizerProcessorStep but it's mocked + assert isinstance(preprocessor.steps[5], DeviceProcessorStep) # Check steps in postprocessor assert len(postprocessor.steps) == 2 - assert isinstance(postprocessor.steps[0], DeviceProcessor) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[0], DeviceProcessorStep) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_smolvla_newline_processor_single_task(): @@ -170,7 +170,7 @@ def test_smolvla_processor_cuda(): stats = create_default_stats() # Mock the tokenizer processor to act as pass-through - class MockTokenizerProcessor: + class MockTokenizerProcessorStep: def __init__(self, *args, **kwargs): pass @@ -192,7 +192,9 @@ def test_smolvla_processor_cuda(): def transform_features(self, features): return features - with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor): + with patch( + "lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep + ): preprocessor, postprocessor = make_smolvla_pre_post_processors( config, stats, @@ -225,7 +227,7 @@ def test_smolvla_processor_accelerate_scenario(): stats = create_default_stats() # Mock the tokenizer processor to act as pass-through - class MockTokenizerProcessor: + class MockTokenizerProcessorStep: def __init__(self, *args, **kwargs): pass @@ -247,7 +249,9 @@ def test_smolvla_processor_accelerate_scenario(): def transform_features(self, features): return features - with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor): + with patch( + "lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep + ): preprocessor, postprocessor = make_smolvla_pre_post_processors( config, stats, @@ -281,7 +285,7 @@ def test_smolvla_processor_multi_gpu(): stats = create_default_stats() # Mock the tokenizer processor to act as pass-through - class MockTokenizerProcessor: + class MockTokenizerProcessorStep: def __init__(self, *args, **kwargs): pass @@ -303,7 +307,9 @@ def test_smolvla_processor_multi_gpu(): def transform_features(self, features): return features - with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor): + with patch( + "lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessorStep + ): preprocessor, postprocessor = make_smolvla_pre_post_processors( config, stats, @@ -334,7 +340,7 @@ def test_smolvla_processor_without_stats(): config = create_default_config() # Mock the tokenizer processor - with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor"): + with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep"): preprocessor, postprocessor = make_smolvla_pre_post_processors( config, dataset_stats=None, diff --git a/tests/processor/test_tdmpc_processor.py b/tests/processor/test_tdmpc_processor.py index b9be43ec4..14858b8b1 100644 --- a/tests/processor/test_tdmpc_processor.py +++ b/tests/processor/test_tdmpc_processor.py @@ -25,13 +25,13 @@ from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, - RenameProcessor, - ToBatchProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, + RenameProcessorStep, TransitionKey, - UnnormalizerProcessor, + UnnormalizerProcessorStep, ) @@ -94,15 +94,15 @@ def test_make_tdmpc_processor_basic(): # Check steps in preprocessor assert len(preprocessor.steps) == 4 - assert isinstance(preprocessor.steps[0], RenameProcessor) - assert isinstance(preprocessor.steps[1], NormalizerProcessor) - assert isinstance(preprocessor.steps[2], ToBatchProcessor) - assert isinstance(preprocessor.steps[3], DeviceProcessor) + assert isinstance(preprocessor.steps[0], RenameProcessorStep) + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) + assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) + assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor assert len(postprocessor.steps) == 2 - assert isinstance(postprocessor.steps[0], DeviceProcessor) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[0], DeviceProcessorStep) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_tdmpc_processor_normalization(): @@ -330,10 +330,10 @@ def test_tdmpc_processor_mixed_precision(): postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x}, ) - # Replace DeviceProcessor with one that uses float16 + # Replace DeviceProcessorStep with one that uses float16 for i, step in enumerate(preprocessor.steps): - if isinstance(step, DeviceProcessor): - preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16") + if isinstance(step, DeviceProcessorStep): + preprocessor.steps[i] = DeviceProcessorStep(device=config.device, float_dtype="float16") # Create test data observation = { diff --git a/tests/processor/test_tokenizer_processor.py b/tests/processor/test_tokenizer_processor.py index 6a518a250..5a3c97240 100644 --- a/tests/processor/test_tokenizer_processor.py +++ b/tests/processor/test_tokenizer_processor.py @@ -1,5 +1,5 @@ """ -Tests for the TokenizerProcessor class. +Tests for the TokenizerProcessorStep class. """ import tempfile @@ -10,7 +10,7 @@ import torch from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.constants import OBS_LANGUAGE -from lerobot.processor import DataProcessorPipeline, TokenizerProcessor, TransitionKey +from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep, TransitionKey from tests.utils import require_package @@ -95,7 +95,7 @@ def test_basic_tokenization(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10) + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=10) transition = create_transition( observation={"state": torch.tensor([1.0, 2.0])}, @@ -127,7 +127,7 @@ def test_basic_tokenization_with_tokenizer_object(): """Test basic string tokenization functionality using tokenizer object directly.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10) transition = create_transition( observation={"state": torch.tensor([1.0, 2.0])}, @@ -161,7 +161,7 @@ def test_list_of_strings_tokenization(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=8) + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=8) transition = create_transition( observation={"state": torch.tensor([1.0, 2.0])}, @@ -189,7 +189,7 @@ def test_custom_keys(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer", task_key="instruction", max_length=5) + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", task_key="instruction", max_length=5) transition = create_transition( observation={"state": torch.tensor([1.0, 2.0])}, @@ -215,7 +215,7 @@ def test_none_complementary_data(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer") + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer") transition = create_transition(complementary_data=None) @@ -230,7 +230,7 @@ def test_missing_task_key(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer") + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer") transition = create_transition(complementary_data={"other_field": "some value"}) @@ -245,7 +245,7 @@ def test_none_task_value(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer") + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer") transition = create_transition(complementary_data={"task": None}) @@ -260,7 +260,7 @@ def test_unsupported_task_type(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer") + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer") # Test with integer task transition = create_transition(complementary_data={"task": 123}) @@ -279,7 +279,7 @@ def test_unsupported_task_type(mock_auto_tokenizer): def test_no_tokenizer_error(): """Test that ValueError is raised when neither tokenizer nor tokenizer_name is provided.""" with pytest.raises(ValueError, match="Either 'tokenizer' or 'tokenizer_name' must be provided"): - TokenizerProcessor() + TokenizerProcessorStep() @require_package("transformers") @@ -290,7 +290,7 @@ def test_invalid_tokenizer_name_error(): mock_auto_tokenizer.from_pretrained.side_effect = Exception("Model not found") with pytest.raises(Exception, match="Model not found"): - TokenizerProcessor(tokenizer_name="invalid-tokenizer") + TokenizerProcessorStep(tokenizer_name="invalid-tokenizer") @require_package("transformers") @@ -300,7 +300,7 @@ def test_get_config_with_tokenizer_name(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor( + processor = TokenizerProcessorStep( tokenizer_name="test-tokenizer", max_length=256, task_key="instruction", @@ -327,7 +327,7 @@ def test_get_config_with_tokenizer_object(): """Test configuration serialization when using tokenizer object.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor( + processor = TokenizerProcessorStep( tokenizer=mock_tokenizer, max_length=256, task_key="instruction", @@ -357,7 +357,7 @@ def test_state_dict_methods(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer") + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer") # Should return empty dict state = processor.state_dict() @@ -374,7 +374,7 @@ def test_reset_method(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer") + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer") # Should not raise error processor.reset() @@ -387,7 +387,7 @@ def test_integration_with_robot_processor(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6) + tokenizer_processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=6) robot_processor = DataProcessorPipeline( [tokenizer_processor], to_transition=lambda x: x, to_output=lambda x: x ) @@ -424,7 +424,7 @@ def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - original_processor = TokenizerProcessor( + original_processor = TokenizerProcessorStep( tokenizer_name="test-tokenizer", max_length=32, task_key="instruction" ) @@ -459,7 +459,9 @@ def test_save_and_load_pretrained_with_tokenizer_object(): """Test saving and loading processor with tokenizer object using overrides.""" mock_tokenizer = MockTokenizer(vocab_size=100) - original_processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=32, task_key="instruction") + original_processor = TokenizerProcessorStep( + tokenizer=mock_tokenizer, max_length=32, task_key="instruction" + ) robot_processor = DataProcessorPipeline( [original_processor], to_transition=lambda x: x, to_output=lambda x: x @@ -500,14 +502,14 @@ def test_registry_functionality(): # Check that we can retrieve it retrieved_class = ProcessorStepRegistry.get("tokenizer_processor") - assert retrieved_class is TokenizerProcessor + assert retrieved_class is TokenizerProcessorStep @require_package("transformers") def test_features_basic(): """Test basic feature contract functionality.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=128) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=128) input_features = { "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,)), @@ -538,7 +540,7 @@ def test_features_basic(): 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) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=64) input_features = {} output_features = processor.transform_features(input_features) @@ -558,7 +560,7 @@ def test_features_with_custom_max_length(): 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) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=256) input_features = { f"{OBS_LANGUAGE}.tokens": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)), @@ -595,7 +597,7 @@ def test_tokenization_parameters(mock_auto_tokenizer): tracking_tokenizer = TrackingMockTokenizer() mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer - processor = TokenizerProcessor( + processor = TokenizerProcessorStep( tokenizer_name="test-tokenizer", max_length=16, padding="longest", @@ -627,7 +629,7 @@ def test_preserves_other_complementary_data(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer") + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer") transition = create_transition( observation={"state": torch.tensor([1.0, 2.0])}, @@ -662,7 +664,7 @@ def test_deterministic_tokenization(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10) + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=10) transition = create_transition( observation={"state": torch.tensor([1.0, 2.0])}, @@ -690,7 +692,7 @@ def test_empty_string_task(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=8) + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=8) transition = create_transition( observation={"state": torch.tensor([1.0, 2.0])}, @@ -714,7 +716,7 @@ def test_very_long_task(mock_auto_tokenizer): mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=5, truncation=True) + processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=5, truncation=True) long_task = " ".join(["word"] * 100) # Very long task transition = create_transition( @@ -764,7 +766,9 @@ def test_custom_padding_side(mock_auto_tokenizer): mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer # Test left padding - processor_left = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10, padding_side="left") + processor_left = TokenizerProcessorStep( + tokenizer_name="test-tokenizer", max_length=10, padding_side="left" + ) transition = create_transition( observation={"state": torch.tensor([1.0, 2.0])}, @@ -776,7 +780,9 @@ def test_custom_padding_side(mock_auto_tokenizer): assert tracking_tokenizer.padding_side_calls[-1] == "left" # Test right padding (default) - processor_right = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=10, padding_side="right") + processor_right = TokenizerProcessorStep( + tokenizer_name="test-tokenizer", max_length=10, padding_side="right" + ) processor_right(transition) @@ -787,7 +793,7 @@ def test_custom_padding_side(mock_auto_tokenizer): def test_device_detection_cpu(): """Test that tokenized tensors stay on CPU when other tensors are on CPU.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10) # Create transition with CPU tensors observation = {"observation.state": torch.randn(10)} # CPU tensor @@ -811,7 +817,7 @@ def test_device_detection_cpu(): def test_device_detection_cuda(): """Test that tokenized tensors are moved to CUDA when other tensors are on CUDA.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10) # Create transition with CUDA tensors observation = {"observation.state": torch.randn(10).cuda()} # CUDA tensor @@ -836,7 +842,7 @@ def test_device_detection_cuda(): def test_device_detection_multi_gpu(): """Test that tokenized tensors match device in multi-GPU setup.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10) # Test with tensors on cuda:1 device = torch.device("cuda:1") @@ -860,7 +866,7 @@ def test_device_detection_multi_gpu(): def test_device_detection_no_tensors(): """Test that tokenized tensors stay on CPU when no other tensors exist.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10) # Create transition with no tensors transition = create_transition( @@ -882,7 +888,7 @@ def test_device_detection_no_tensors(): def test_device_detection_mixed_devices(): """Test device detection when tensors are on different devices (uses first found).""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10) if torch.cuda.is_available(): # Create transition with mixed devices @@ -910,7 +916,7 @@ def test_device_detection_mixed_devices(): def test_device_detection_from_action(): """Test that device is detected from action tensor when no observation tensors exist.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10) # Create transition with action on CUDA but no observation tensors observation = {"metadata": {"key": "value"}} # No tensors in observation @@ -933,7 +939,7 @@ def test_device_detection_from_action(): def test_device_detection_preserves_dtype(): """Test that device detection doesn't affect dtype of tokenized tensors.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10) # Create transition with float tensor (to test dtype isn't affected) observation = {"observation.state": torch.randn(10, dtype=torch.float16)} @@ -953,15 +959,15 @@ def test_device_detection_preserves_dtype(): @require_package("transformers") @patch("lerobot.processor.tokenizer_processor.AutoTokenizer") def test_integration_with_device_processor(mock_auto_tokenizer): - """Test that TokenizerProcessor works correctly with DeviceProcessor in pipeline.""" - from lerobot.processor import DeviceProcessor + """Test that TokenizerProcessorStep works correctly with DeviceProcessorStep in pipeline.""" + from lerobot.processor import DeviceProcessorStep mock_tokenizer = MockTokenizer(vocab_size=100) mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer - # Create pipeline with TokenizerProcessor then DeviceProcessor - tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6) - device_processor = DeviceProcessor(device="cuda:0") + # Create pipeline with TokenizerProcessorStep then DeviceProcessorStep + tokenizer_processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=6) + device_processor = DeviceProcessorStep(device="cuda:0") robot_processor = DataProcessorPipeline( [tokenizer_processor, device_processor], to_transition=lambda x: x, to_output=lambda x: x ) @@ -975,7 +981,7 @@ def test_integration_with_device_processor(mock_auto_tokenizer): result = robot_processor(transition) - # All tensors should end up on CUDA (moved by DeviceProcessor) + # All tensors should end up on CUDA (moved by DeviceProcessorStep) assert result[TransitionKey.OBSERVATION]["observation.state"].device.type == "cuda" assert result[TransitionKey.ACTION].device.type == "cuda" @@ -991,7 +997,7 @@ def test_integration_with_device_processor(mock_auto_tokenizer): def test_simulated_accelerate_scenario(): """Test scenario simulating Accelerate with data already on GPU.""" mock_tokenizer = MockTokenizer(vocab_size=100) - processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=10) + processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10) # Simulate Accelerate scenario: batch already on GPU device = torch.device("cuda:0") diff --git a/tests/processor/test_vqbet_processor.py b/tests/processor/test_vqbet_processor.py index 6f7e21b41..891e88056 100644 --- a/tests/processor/test_vqbet_processor.py +++ b/tests/processor/test_vqbet_processor.py @@ -25,13 +25,13 @@ from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors from lerobot.processor import ( + AddBatchDimensionProcessorStep, DataProcessorPipeline, - DeviceProcessor, - NormalizerProcessor, - RenameProcessor, - ToBatchProcessor, + DeviceProcessorStep, + NormalizerProcessorStep, + RenameProcessorStep, TransitionKey, - UnnormalizerProcessor, + UnnormalizerProcessorStep, ) @@ -94,15 +94,15 @@ def test_make_vqbet_processor_basic(): # Check steps in preprocessor assert len(preprocessor.steps) == 4 - assert isinstance(preprocessor.steps[0], RenameProcessor) - assert isinstance(preprocessor.steps[1], NormalizerProcessor) - assert isinstance(preprocessor.steps[2], ToBatchProcessor) - assert isinstance(preprocessor.steps[3], DeviceProcessor) + assert isinstance(preprocessor.steps[0], RenameProcessorStep) + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) + assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) + assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor assert len(postprocessor.steps) == 2 - assert isinstance(postprocessor.steps[0], DeviceProcessor) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[0], DeviceProcessorStep) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_vqbet_processor_with_images(): @@ -323,10 +323,10 @@ def test_vqbet_processor_mixed_precision(): postprocessor_kwargs={"to_transition": lambda x: x, "to_output": lambda x: x}, ) - # Replace DeviceProcessor with one that uses float16 + # Replace DeviceProcessorStep with one that uses float16 for i, step in enumerate(preprocessor.steps): - if isinstance(step, DeviceProcessor): - preprocessor.steps[i] = DeviceProcessor(device=config.device, float_dtype="float16") + if isinstance(step, DeviceProcessorStep): + preprocessor.steps[i] = DeviceProcessorStep(device=config.device, float_dtype="float16") # Create test data observation = {