From 4279229c0d13b643ff85da90f3dffa8bb928d396 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Thu, 28 Aug 2025 17:43:27 +0200 Subject: [PATCH 01/12] Chore/rename RobotProcessor to DataProcessorPipeline (#1802) * refactor(processor): rename RobotProcessor to DataProcessorPipeline for clarity and consistency * refactor(pipeline): replace RobotProcessor with DataProcessorPipeline across multiple files for consistency --- examples/phone_to_so100/evaluate.py | 6 +- examples/phone_to_so100/record.py | 8 +- examples/phone_to_so100/replay.py | 4 +- examples/phone_to_so100/teleoperate.py | 4 +- src/lerobot/datasets/pipeline_features.py | 4 +- src/lerobot/policies/act/processor_act.py | 6 +- .../policies/diffusion/processor_diffusion.py | 6 +- src/lerobot/policies/factory.py | 8 +- src/lerobot/policies/pi0/processor_pi0.py | 6 +- .../policies/pi0fast/processor_pi0fast.py | 6 +- src/lerobot/policies/sac/processor_sac.py | 6 +- .../sac/reward_model/processor_classifier.py | 6 +- .../policies/smolvla/processor_smolvla.py | 6 +- src/lerobot/policies/tdmpc/processor_tdmpc.py | 6 +- src/lerobot/policies/vqbet/processor_vqbet.py | 6 +- src/lerobot/processor/__init__.py | 4 +- .../processor/migrate_policy_normalization.py | 6 +- src/lerobot/processor/normalize_processor.py | 10 +- src/lerobot/processor/pipeline.py | 36 +-- src/lerobot/record.py | 18 +- src/lerobot/replay.py | 6 +- src/lerobot/scripts/rl/gym_manipulator.py | 20 +- src/lerobot/teleoperate.py | 20 +- src/lerobot/utils/control_utils.py | 6 +- src/lerobot/utils/train_utils.py | 6 +- tests/processor/test_act_processor.py | 4 +- tests/processor/test_batch_conversion.py | 6 +- tests/processor/test_batch_processor.py | 16 +- tests/processor/test_classifier_processor.py | 4 +- tests/processor/test_device_processor.py | 14 +- tests/processor/test_diffusion_processor.py | 4 +- tests/processor/test_normalize_processor.py | 24 +- tests/processor/test_pipeline.py | 246 +++++++++--------- tests/processor/test_rename_processor.py | 24 +- tests/processor/test_sac_processor.py | 4 +- tests/processor/test_tdmpc_processor.py | 4 +- tests/processor/test_tokenizer_processor.py | 16 +- tests/processor/test_vqbet_processor.py | 4 +- 38 files changed, 303 insertions(+), 287 deletions(-) diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index d1190b363d..d67fc125bb 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -25,7 +25,7 @@ to_output_robot_action, to_transition_robot_observation, ) -from lerobot.processor.pipeline import RobotProcessor +from lerobot.processor.pipeline import DataProcessorPipeline from lerobot.record import record_loop from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.robot_kinematic_processor import ( @@ -65,7 +65,7 @@ ) # Build pipeline to convert ee pose action to joint action -robot_ee_to_joints = RobotProcessor( +robot_ee_to_joints = DataProcessorPipeline( steps=[ AddRobotObservationAsComplimentaryData(robot=robot), InverseKinematicsEEToJoints( @@ -79,7 +79,7 @@ ) # Build pipeline to convert joint observation to ee pose observation -robot_joints_to_ee_pose = RobotProcessor( +robot_joints_to_ee_pose = DataProcessorPipeline( steps=[ ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ], diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index e9d22ef80d..575199e66c 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -25,7 +25,7 @@ to_transition_robot_observation, to_transition_teleop_action, ) -from lerobot.processor.pipeline import RobotProcessor +from lerobot.processor.pipeline import DataProcessorPipeline from lerobot.record import record_loop from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.robot_kinematic_processor import ( @@ -73,7 +73,7 @@ ) # Build pipeline to convert phone action to ee pose action -phone_to_robot_ee_pose = RobotProcessor( +phone_to_robot_ee_pose = DataProcessorPipeline( steps=[ MapPhoneActionToRobotAction(platform=teleop_config.phone_os), AddRobotObservationAsComplimentaryData(robot=robot), @@ -93,7 +93,7 @@ ) # Build pipeline to convert ee pose action to joint action -robot_ee_to_joints = RobotProcessor( +robot_ee_to_joints = DataProcessorPipeline( steps=[ InverseKinematicsEEToJoints( kinematics=kinematics_solver, @@ -110,7 +110,7 @@ ) # Build pipeline to convert joint observation to ee pose observation -robot_joints_to_ee_pose = RobotProcessor( +robot_joints_to_ee_pose = DataProcessorPipeline( steps=[ ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ], diff --git a/examples/phone_to_so100/replay.py b/examples/phone_to_so100/replay.py index e39f482c27..a91f29d6f5 100644 --- a/examples/phone_to_so100/replay.py +++ b/examples/phone_to_so100/replay.py @@ -20,7 +20,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.model.kinematics import RobotKinematics from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action -from lerobot.processor.pipeline import RobotProcessor +from lerobot.processor.pipeline import DataProcessorPipeline from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.robot_kinematic_processor import ( AddRobotObservationAsComplimentaryData, @@ -50,7 +50,7 @@ ) # Build pipeline to convert ee pose action to joint action -robot_ee_to_joints = RobotProcessor( +robot_ee_to_joints = DataProcessorPipeline( steps=[ AddRobotObservationAsComplimentaryData(robot=robot), InverseKinematicsEEToJoints( diff --git a/examples/phone_to_so100/teleoperate.py b/examples/phone_to_so100/teleoperate.py index 1eef0f8ae1..85bfe4020e 100644 --- a/examples/phone_to_so100/teleoperate.py +++ b/examples/phone_to_so100/teleoperate.py @@ -16,7 +16,7 @@ import time from lerobot.model.kinematics import RobotKinematics -from lerobot.processor import RobotProcessor +from lerobot.processor import DataProcessorPipeline from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.robot_kinematic_processor import ( @@ -49,7 +49,7 @@ ) # Build pipeline to convert phone action to ee pose action to joint action -phone_to_robot_joints = RobotProcessor( +phone_to_robot_joints = DataProcessorPipeline( steps=[ MapPhoneActionToRobotAction(platform=teleop_config.phone_os), AddRobotObservationAsComplimentaryData(robot=robot), diff --git a/src/lerobot/datasets/pipeline_features.py b/src/lerobot/datasets/pipeline_features.py index fef75b4079..fd2cf3363b 100644 --- a/src/lerobot/datasets/pipeline_features.py +++ b/src/lerobot/datasets/pipeline_features.py @@ -16,11 +16,11 @@ from typing import Any from lerobot.datasets.utils import hw_to_dataset_features -from lerobot.processor.pipeline import RobotProcessor +from lerobot.processor.pipeline import DataProcessorPipeline def aggregate_pipeline_dataset_features( - pipeline: RobotProcessor, + pipeline: DataProcessorPipeline, initial_features: dict[str, Any], *, use_videos: bool = True, diff --git a/src/lerobot/policies/act/processor_act.py b/src/lerobot/policies/act/processor_act.py index db6d4d85a3..7b1f1693db 100644 --- a/src/lerobot/policies/act/processor_act.py +++ b/src/lerobot/policies/act/processor_act.py @@ -18,10 +18,10 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.act.configuration_act import ACTConfig from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -29,7 +29,7 @@ def make_act_pre_post_processors( config: ACTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), NormalizerProcessor( @@ -46,6 +46,6 @@ def make_act_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return RobotProcessor(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), RobotProcessor( + return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME ) diff --git a/src/lerobot/policies/diffusion/processor_diffusion.py b/src/lerobot/policies/diffusion/processor_diffusion.py index 2ecf5d700e..08d36a1f83 100644 --- a/src/lerobot/policies/diffusion/processor_diffusion.py +++ b/src/lerobot/policies/diffusion/processor_diffusion.py @@ -19,10 +19,10 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -30,7 +30,7 @@ def make_diffusion_pre_post_processors( config: DiffusionConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), NormalizerProcessor( @@ -47,6 +47,6 @@ def make_diffusion_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return RobotProcessor(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), RobotProcessor( + return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME ) diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 67dd928eca..fe1d7bfbca 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -38,7 +38,7 @@ from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig -from lerobot.processor.pipeline import RobotProcessor +from lerobot.processor.pipeline import DataProcessorPipeline def get_policy_class(name: str) -> type[PreTrainedPolicy]: @@ -120,7 +120,7 @@ def make_pre_post_processors( policy_cfg: PreTrainedConfig, pretrained_path: str | None = None, **kwargs: Unpack[ProcessorConfigKwargs], -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: """Make a processor instance for a given policy type. This function creates the appropriate processor configuration based on the policy type. @@ -140,12 +140,12 @@ def make_pre_post_processors( """ if pretrained_path: return ( - RobotProcessor.from_pretrained( + DataProcessorPipeline.from_pretrained( pretrained_model_name_or_path=pretrained_path, config_filename=kwargs.get("preprocessor_config_filename", "robot_preprocessor.json"), overrides=kwargs.get("preprocessor_overrides", {}), ), - RobotProcessor.from_pretrained( + DataProcessorPipeline.from_pretrained( pretrained_model_name_or_path=pretrained_path, config_filename=kwargs.get("postprocessor_config_filename", "robot_postprocessor.json"), overrides=kwargs.get("postprocessor_overrides", {}), diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index 8b1fc83013..def619add3 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -20,9 +20,9 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, - RobotProcessor, ToBatchProcessor, TokenizerProcessor, UnnormalizerProcessor, @@ -66,7 +66,7 @@ def complementary_data(self, complementary_data): def make_pi0_pre_post_processors( config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: # Add remaining processors input_steps: list[ProcessorStep] = [ RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one @@ -93,6 +93,6 @@ def make_pi0_pre_post_processors( ), ] - return RobotProcessor(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), RobotProcessor( + return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME ) diff --git a/src/lerobot/policies/pi0fast/processor_pi0fast.py b/src/lerobot/policies/pi0fast/processor_pi0fast.py index 9cfd24c762..28642e4b2c 100644 --- a/src/lerobot/policies/pi0fast/processor_pi0fast.py +++ b/src/lerobot/policies/pi0fast/processor_pi0fast.py @@ -19,10 +19,10 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -30,7 +30,7 @@ def make_pi0fast_pre_post_processors( config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one NormalizerProcessor( @@ -47,6 +47,6 @@ def make_pi0fast_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return RobotProcessor(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), RobotProcessor( + return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME ) diff --git a/src/lerobot/policies/sac/processor_sac.py b/src/lerobot/policies/sac/processor_sac.py index d96d1e003a..badaffa394 100644 --- a/src/lerobot/policies/sac/processor_sac.py +++ b/src/lerobot/policies/sac/processor_sac.py @@ -20,10 +20,10 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -31,7 +31,7 @@ def make_sac_pre_post_processors( config: SACConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), NormalizerProcessor( @@ -48,6 +48,6 @@ def make_sac_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return RobotProcessor(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), RobotProcessor( + return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME ) diff --git a/src/lerobot/policies/sac/reward_model/processor_classifier.py b/src/lerobot/policies/sac/reward_model/processor_classifier.py index 084634a732..bf794ac02a 100644 --- a/src/lerobot/policies/sac/reward_model/processor_classifier.py +++ b/src/lerobot/policies/sac/reward_model/processor_classifier.py @@ -17,16 +17,16 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, IdentityProcessor, NormalizerProcessor, - RobotProcessor, ) def make_classifier_processor( config: RewardClassifierConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: input_steps = [ NormalizerProcessor( features=config.input_features, norm_map=config.normalization_mapping, stats=dataset_stats @@ -37,6 +37,6 @@ def make_classifier_processor( DeviceProcessor(device=config.device), ] output_steps = [DeviceProcessor(device="cpu"), IdentityProcessor()] - return RobotProcessor(steps=input_steps, name="classifier_preprocessor"), RobotProcessor( + return DataProcessorPipeline(steps=input_steps, name="classifier_preprocessor"), DataProcessorPipeline( steps=output_steps, name="classifier_postprocessor" ) diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index f9d6a45947..2bf26a407c 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -19,10 +19,10 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, TokenizerProcessor, UnnormalizerProcessor, @@ -35,7 +35,7 @@ def make_smolvla_pre_post_processors( config: SmolVLAConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one NormalizerProcessor( @@ -59,7 +59,7 @@ def make_smolvla_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return RobotProcessor(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), RobotProcessor( + return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME ) diff --git a/src/lerobot/policies/tdmpc/processor_tdmpc.py b/src/lerobot/policies/tdmpc/processor_tdmpc.py index 9d13ef0e11..01d7e3cf45 100644 --- a/src/lerobot/policies/tdmpc/processor_tdmpc.py +++ b/src/lerobot/policies/tdmpc/processor_tdmpc.py @@ -19,10 +19,10 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -30,7 +30,7 @@ def make_tdmpc_pre_post_processors( config: TDMPCConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), NormalizerProcessor( @@ -47,6 +47,6 @@ def make_tdmpc_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return RobotProcessor(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), RobotProcessor( + return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME ) diff --git a/src/lerobot/policies/vqbet/processor_vqbet.py b/src/lerobot/policies/vqbet/processor_vqbet.py index 242e6c3b09..23b600b61f 100644 --- a/src/lerobot/policies/vqbet/processor_vqbet.py +++ b/src/lerobot/policies/vqbet/processor_vqbet.py @@ -20,10 +20,10 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -31,7 +31,7 @@ def make_vqbet_pre_post_processors( config: VQBeTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[RobotProcessor, RobotProcessor]: +) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), # Let the possibility to the user to rename the keys NormalizerProcessor( @@ -48,6 +48,6 @@ def make_vqbet_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return RobotProcessor(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), RobotProcessor( + return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME ) diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 813756490f..4e25ea1e6a 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -32,6 +32,7 @@ from .observation_processor import VanillaObservationProcessor from .pipeline import ( ActionProcessor, + DataProcessorPipeline, DoneProcessor, EnvTransition, IdentityProcessor, @@ -40,7 +41,6 @@ ProcessorStep, ProcessorStepRegistry, RewardProcessor, - RobotProcessor, TransitionKey, TruncatedProcessor, ) @@ -73,7 +73,7 @@ "RenameProcessor", "RewardClassifierProcessor", "RewardProcessor", - "RobotProcessor", + "DataProcessorPipeline", "ToBatchProcessor", "TokenizerProcessor", "TimeLimitProcessor", diff --git a/src/lerobot/processor/migrate_policy_normalization.py b/src/lerobot/processor/migrate_policy_normalization.py index ae4b1fb235..348a18d58a 100644 --- a/src/lerobot/processor/migrate_policy_normalization.py +++ b/src/lerobot/processor/migrate_policy_normalization.py @@ -49,7 +49,7 @@ from lerobot.processor.batch_processor import ToBatchProcessor from lerobot.processor.device_processor import DeviceProcessor from lerobot.processor.normalize_processor import NormalizerProcessor, UnnormalizerProcessor -from lerobot.processor.pipeline import RobotProcessor +from lerobot.processor.pipeline import DataProcessorPipeline from lerobot.processor.rename_processor import RenameProcessor # Policy type to class mapping @@ -421,14 +421,14 @@ def main(): ToBatchProcessor(), DeviceProcessor(device=policy_config.device), ] - preprocessor = RobotProcessor(steps=preprocessor_steps, name="robot_preprocessor") + 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), ] - postprocessor = RobotProcessor(steps=postprocessor_steps, name="robot_postprocessor") + postprocessor = DataProcessorPipeline(steps=postprocessor_steps, name="robot_postprocessor") # Determine hub repo ID if pushing to hub if args.push_to_hub: diff --git a/src/lerobot/processor/normalize_processor.py b/src/lerobot/processor/normalize_processor.py index fa635414ca..a75e82097d 100644 --- a/src/lerobot/processor/normalize_processor.py +++ b/src/lerobot/processor/normalize_processor.py @@ -11,10 +11,10 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.processor.pipeline import ( + DataProcessorPipeline, EnvTransition, ProcessorStep, ProcessorStepRegistry, - RobotProcessor, TransitionKey, ) @@ -281,11 +281,13 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: return new_transition -def hotswap_stats(robot_processor: RobotProcessor, stats: dict[str, dict[str, Any]]) -> RobotProcessor: +def hotswap_stats( + robot_processor: DataProcessorPipeline, stats: dict[str, dict[str, Any]] +) -> DataProcessorPipeline: """ - Replaces normalization statistics in a RobotProcessor pipeline. + Replaces normalization statistics in a DataProcessorPipeline pipeline. - This function creates a deep copy of the provided `RobotProcessor` and updates the + This function creates a deep copy of the provided `DataProcessorPipeline` and updates the statistics of any `NormalizerProcessor` or `UnnormalizerProcessor` 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/pipeline.py b/src/lerobot/processor/pipeline.py index 8ffe490d6c..018b66b6b1 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -139,7 +139,7 @@ class ProcessorStep(ABC): A step is any callable accepting a full `EnvTransition` dict and returning a (possibly modified) dict of the same structure. Implementers are encouraged—but not required—to expose the optional helper methods - listed below. When present, these hooks let `RobotProcessor` + listed below. When present, these hooks let `DataProcessorPipeline` automatically serialise the step's configuration and learnable state using a safe-to-share JSON + SafeTensors format. @@ -280,7 +280,7 @@ def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]: @dataclass -class RobotProcessor(ModelHubMixin): +class DataProcessorPipeline(ModelHubMixin): """ Composable, debuggable post-processing processor for robot transitions. @@ -291,7 +291,7 @@ class RobotProcessor(ModelHubMixin): Args: steps: Ordered list of processing steps executed on every call. Defaults to empty list. name: Human-readable identifier that is persisted inside the JSON config. - Defaults to "RobotProcessor". + Defaults to "DataProcessorPipeline". to_transition: Function to convert batch dict to EnvTransition dict. Defaults to _default_batch_to_transition. to_output: Function to convert EnvTransition dict to the desired output format. @@ -318,7 +318,7 @@ class RobotProcessor(ModelHubMixin): """ steps: Sequence[ProcessorStep] = field(default_factory=list) - name: str = "RobotProcessor" + name: str = "DataProcessorPipeline" to_transition: Callable[[dict[str, Any]], EnvTransition] = field( default_factory=lambda: _default_batch_to_transition, repr=False @@ -526,7 +526,7 @@ def from_pretrained( config_filename: str | None = None, overrides: dict[str, Any] | None = None, **kwargs, - ) -> RobotProcessor: + ) -> DataProcessorPipeline: """Load a serialized processor from source (local path or Hugging Face Hub identifier). Args: @@ -542,7 +542,7 @@ def from_pretrained( non-serializable objects like environment instances. Returns: - A RobotProcessor instance loaded from the saved configuration. + A DataProcessorPipeline instance loaded from the saved configuration. Raises: ImportError: If a processor step class cannot be loaded or imported. @@ -552,12 +552,12 @@ def from_pretrained( Examples: Basic loading: ```python - processor = RobotProcessor.from_pretrained("path/to/processor") + processor = DataProcessorPipeline.from_pretrained("path/to/processor") ``` Loading specific config file: ```python - processor = RobotProcessor.from_pretrained( + processor = DataProcessorPipeline.from_pretrained( "username/multi-processor-repo", config_filename="preprocessor.json" ) ``` @@ -567,14 +567,14 @@ def from_pretrained( import gym env = gym.make("CartPole-v1") - processor = RobotProcessor.from_pretrained( + processor = DataProcessorPipeline.from_pretrained( "username/cartpole-processor", overrides={"ActionRepeatStep": {"env": env}} ) ``` Multiple overrides: ```python - processor = RobotProcessor.from_pretrained( + processor = DataProcessorPipeline.from_pretrained( "path/to/processor", overrides={ "CustomStep": {"param1": "new_value"}, @@ -609,9 +609,9 @@ def from_pretrained( if config_filename is None: # Try common config names common_names = [ - "robot_processor.json", - "robot_preprocessor.json", - "robot_postprocessor.json", + "data_processor.json", + "data_preprocessor.json", + "data_postprocessor.json", ] config_path = None for name in common_names: @@ -756,19 +756,19 @@ def from_pretrained( f"Make sure override keys match exact step class names or registry names." ) - return cls(steps, loaded_config.get("name", "RobotProcessor")) + return cls(steps, loaded_config.get("name", "DataProcessorPipeline")) def __len__(self) -> int: """Return the number of steps in the processor.""" return len(self.steps) - def __getitem__(self, idx: int | slice) -> ProcessorStep | RobotProcessor: + def __getitem__(self, idx: int | slice) -> ProcessorStep | DataProcessorPipeline: """Indexing helper exposing underlying steps. * ``int`` – returns the idx-th ProcessorStep. - * ``slice`` – returns a new RobotProcessor with the sliced steps. + * ``slice`` – returns a new DataProcessorPipeline with the sliced steps. """ if isinstance(idx, slice): - return RobotProcessor(self.steps[idx], self.name) + return DataProcessorPipeline(self.steps[idx], self.name) return self.steps[idx] def register_before_step_hook(self, fn: Callable[[int, EnvTransition], None]): @@ -832,7 +832,7 @@ def __repr__(self) -> str: parts = [f"name='{self.name}'", steps_repr] - return f"RobotProcessor({', '.join(parts)})" + return f"DataProcessorPipeline({', '.join(parts)})" def __post_init__(self): for i, step in enumerate(self.steps): diff --git a/src/lerobot/record.py b/src/lerobot/record.py index 0ebe23501e..441d7e2c08 100644 --- a/src/lerobot/record.py +++ b/src/lerobot/record.py @@ -76,7 +76,7 @@ from lerobot.datasets.video_utils import VideoEncodingManager from lerobot.policies.factory import make_policy, make_pre_post_processors from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor import RobotProcessor +from lerobot.processor import DataProcessorPipeline from lerobot.processor.converters import ( to_dataset_frame, to_output_robot_action, @@ -236,22 +236,22 @@ def record_loop( dataset: LeRobotDataset | None = None, teleop: Teleoperator | list[Teleoperator] | None = None, policy: PreTrainedPolicy | None = None, - preprocessor: RobotProcessor | None = None, - postprocessor: RobotProcessor | None = None, + preprocessor: DataProcessorPipeline | None = None, + postprocessor: DataProcessorPipeline | None = None, control_time_s: int | None = None, - teleop_action_processor: RobotProcessor | None = None, # runs after teleop - robot_action_processor: RobotProcessor | None = None, # runs before robot - robot_observation_processor: RobotProcessor | None = None, # runs after robot + teleop_action_processor: DataProcessorPipeline | None = None, # runs after teleop + robot_action_processor: DataProcessorPipeline | None = None, # runs before robot + robot_observation_processor: DataProcessorPipeline | None = None, # runs after robot single_task: str | None = None, display_data: bool = False, ): - teleop_action_processor = teleop_action_processor or RobotProcessor( + teleop_action_processor = teleop_action_processor or DataProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr ) - robot_action_processor = robot_action_processor or RobotProcessor( + robot_action_processor = robot_action_processor or DataProcessorPipeline( steps=[IdentityProcessor()], to_transition=lambda tr: tr, to_output=to_output_robot_action ) - robot_observation_processor = robot_observation_processor or RobotProcessor( + robot_observation_processor = robot_observation_processor or DataProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr ) diff --git a/src/lerobot/replay.py b/src/lerobot/replay.py index f6628ac6b5..c09554d744 100644 --- a/src/lerobot/replay.py +++ b/src/lerobot/replay.py @@ -47,7 +47,7 @@ from lerobot.configs import parser from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.processor import RobotProcessor +from lerobot.processor import DataProcessorPipeline from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action from lerobot.processor.pipeline import IdentityProcessor from lerobot.robots import ( # noqa: F401 @@ -86,7 +86,7 @@ class ReplayConfig: # Use vocal synthesis to read events. play_sounds: bool = True # Optional processor for actions before sending to robot - robot_action_processor: RobotProcessor | None = None + robot_action_processor: DataProcessorPipeline | None = None @parser.wrap() @@ -95,7 +95,7 @@ def replay(cfg: ReplayConfig): logging.info(pformat(asdict(cfg))) # Initialize robot action processor with default if not provided - robot_action_processor = cfg.robot_action_processor or RobotProcessor( + robot_action_processor = cfg.robot_action_processor or DataProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=to_output_robot_action, # type: ignore[arg-type] diff --git a/src/lerobot/scripts/rl/gym_manipulator.py b/src/lerobot/scripts/rl/gym_manipulator.py index 835b85190c..1a2324a412 100644 --- a/src/lerobot/scripts/rl/gym_manipulator.py +++ b/src/lerobot/scripts/rl/gym_manipulator.py @@ -31,6 +31,7 @@ from lerobot.processor import ( AddTeleopActionAsComplimentaryData, AddTeleopEventsAsInfo, + DataProcessorPipeline, DeviceProcessor, GripperPenaltyProcessor, ImageCropResizeProcessor, @@ -41,7 +42,6 @@ MotorCurrentProcessor, Numpy2TorchActionProcessor, RewardClassifierProcessor, - RobotProcessor, TimeLimitProcessor, ToBatchProcessor, Torch2NumpyActionProcessor, @@ -387,7 +387,9 @@ def make_processors( DeviceProcessor(device=device), ] - return RobotProcessor(steps=env_pipeline_steps), RobotProcessor(steps=action_pipeline_steps) + return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline( + steps=action_pipeline_steps + ) # Full processor pipeline for real robot environment # Get robot and motor information for kinematics @@ -497,15 +499,15 @@ def make_processors( ] action_pipeline_steps.extend(inverse_kinematics_steps) - return RobotProcessor(steps=env_pipeline_steps), RobotProcessor(steps=action_pipeline_steps) + return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline(steps=action_pipeline_steps) def step_env_and_process_transition( env: gym.Env, transition: EnvTransition, action: torch.Tensor, - env_processor: RobotProcessor, - action_processor: RobotProcessor, + env_processor: DataProcessorPipeline, + action_processor: DataProcessorPipeline, ): """ Execute one step with processor pipeline. @@ -554,8 +556,8 @@ def step_env_and_process_transition( def control_loop( env: gym.Env, - env_processor: RobotProcessor, - action_processor: RobotProcessor, + env_processor: DataProcessorPipeline, + action_processor: DataProcessorPipeline, teleop_device: Teleoperator, cfg: GymManipulatorConfig, ) -> None: @@ -709,7 +711,9 @@ def control_loop( dataset.push_to_hub() -def replay_trajectory(env: gym.Env, action_processor: RobotProcessor, cfg: GymManipulatorConfig) -> None: +def replay_trajectory( + env: gym.Env, action_processor: DataProcessorPipeline, cfg: GymManipulatorConfig +) -> None: """Replay recorded trajectory on robot environment.""" assert cfg.dataset.replay_episode is not None, "Replay episode must be provided for replay" diff --git a/src/lerobot/teleoperate.py b/src/lerobot/teleoperate.py index d8101f0b35..33aafad04e 100644 --- a/src/lerobot/teleoperate.py +++ b/src/lerobot/teleoperate.py @@ -61,7 +61,7 @@ from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 from lerobot.configs import parser -from lerobot.processor import RobotProcessor +from lerobot.processor import DataProcessorPipeline from lerobot.processor.converters import ( to_output_robot_action, to_transition_robot_observation, @@ -105,9 +105,9 @@ class TeleoperateConfig: # Display all cameras on screen display_data: bool = False # Optional processors for data transformation - teleop_action_processor: RobotProcessor | None = None # runs after teleop - robot_action_processor: RobotProcessor | None = None # runs before robot - robot_observation_processor: RobotProcessor | None = None # runs after robot + teleop_action_processor: DataProcessorPipeline | None = None # runs after teleop + robot_action_processor: DataProcessorPipeline | None = None # runs before robot + robot_observation_processor: DataProcessorPipeline | None = None # runs after robot def teleop_loop( @@ -116,20 +116,20 @@ def teleop_loop( fps: int, display_data: bool = False, duration: float | None = None, - teleop_action_processor: RobotProcessor | None = None, - robot_action_processor: RobotProcessor | None = None, - robot_observation_processor: RobotProcessor | None = None, + teleop_action_processor: DataProcessorPipeline | None = None, + robot_action_processor: DataProcessorPipeline | None = None, + robot_observation_processor: DataProcessorPipeline | None = None, ): # Initialize processors with defaults if not provided - teleop_action_processor = teleop_action_processor or RobotProcessor( + teleop_action_processor = teleop_action_processor or DataProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr ) - robot_action_processor = robot_action_processor or RobotProcessor( + robot_action_processor = robot_action_processor or DataProcessorPipeline( steps=[IdentityProcessor()], to_transition=lambda tr: tr, to_output=to_output_robot_action, # type: ignore[arg-type] ) - robot_observation_processor = robot_observation_processor or RobotProcessor( + robot_observation_processor = robot_observation_processor or DataProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr ) diff --git a/src/lerobot/utils/control_utils.py b/src/lerobot/utils/control_utils.py index d8c7c9d570..5553d56f23 100644 --- a/src/lerobot/utils/control_utils.py +++ b/src/lerobot/utils/control_utils.py @@ -31,7 +31,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.utils import DEFAULT_FEATURES from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor import RobotProcessor, TransitionKey +from lerobot.processor import DataProcessorPipeline, TransitionKey from lerobot.robots import Robot @@ -102,8 +102,8 @@ def predict_action( observation: dict[str, np.ndarray], policy: PreTrainedPolicy, device: torch.device, - preprocessor: RobotProcessor, - postprocessor: RobotProcessor, + preprocessor: DataProcessorPipeline, + postprocessor: DataProcessorPipeline, use_amp: bool, task: str | None = None, robot_type: str | None = None, diff --git a/src/lerobot/utils/train_utils.py b/src/lerobot/utils/train_utils.py index 1067fa619e..948674c3d3 100644 --- a/src/lerobot/utils/train_utils.py +++ b/src/lerobot/utils/train_utils.py @@ -32,7 +32,7 @@ from lerobot.optim.optimizers import load_optimizer_state, save_optimizer_state from lerobot.optim.schedulers import load_scheduler_state, save_scheduler_state from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor.pipeline import RobotProcessor +from lerobot.processor.pipeline import DataProcessorPipeline from lerobot.utils.random_utils import load_rng_state, save_rng_state @@ -75,8 +75,8 @@ def save_checkpoint( policy: PreTrainedPolicy, optimizer: Optimizer, scheduler: LRScheduler | None = None, - preprocessor: RobotProcessor | None = None, - postprocessor: RobotProcessor | None = None, + preprocessor: DataProcessorPipeline | None = None, + postprocessor: DataProcessorPipeline | None = None, ) -> None: """This function creates the following directory structure: diff --git a/tests/processor/test_act_processor.py b/tests/processor/test_act_processor.py index afba2aa268..6ee926da2c 100644 --- a/tests/processor/test_act_processor.py +++ b/tests/processor/test_act_processor.py @@ -25,10 +25,10 @@ from lerobot.policies.act.configuration_act import ACTConfig from lerobot.policies.act.processor_act import make_act_pre_post_processors from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -230,7 +230,7 @@ def test_act_processor_save_and_load(): preprocessor.save_pretrained(tmpdir) # Load preprocessor - loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir) + loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir) # Test that loaded processor works observation = {OBS_STATE: torch.randn(7)} diff --git a/tests/processor/test_batch_conversion.py b/tests/processor/test_batch_conversion.py index 63894025d3..1a47a8adc5 100644 --- a/tests/processor/test_batch_conversion.py +++ b/tests/processor/test_batch_conversion.py @@ -1,7 +1,7 @@ import torch from lerobot.processor.pipeline import ( - RobotProcessor, + DataProcessorPipeline, TransitionKey, _default_batch_to_transition, _default_transition_to_batch, @@ -24,7 +24,7 @@ def _dummy_batch(): def test_observation_grouping_roundtrip(): """Test that observation.* keys are properly grouped and ungrouped.""" - proc = RobotProcessor([]) + proc = DataProcessorPipeline([]) batch_in = _dummy_batch() batch_out = proc(batch_in) @@ -265,7 +265,7 @@ def to_batch(tr): batch = _default_transition_to_batch(tr) return batch - processor = RobotProcessor(steps=[], to_transition=to_tr, to_output=to_batch) + processor = DataProcessorPipeline(steps=[], to_transition=to_tr, to_output=to_batch) batch = { "observation.state": torch.randn(1, 4), diff --git a/tests/processor/test_batch_processor.py b/tests/processor/test_batch_processor.py index 0bf050e20c..ecb694143e 100644 --- a/tests/processor/test_batch_processor.py +++ b/tests/processor/test_batch_processor.py @@ -22,7 +22,7 @@ import torch from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE -from lerobot.processor import ProcessorStepRegistry, RobotProcessor +from lerobot.processor import DataProcessorPipeline, ProcessorStepRegistry from lerobot.processor.batch_processor import ToBatchProcessor from lerobot.processor.pipeline import TransitionKey @@ -243,9 +243,9 @@ def test_mixed_observation(): def test_integration_with_robot_processor(): - """Test ToBatchProcessor integration with RobotProcessor.""" + """Test ToBatchProcessor integration with DataProcessorPipeline.""" to_batch_processor = ToBatchProcessor() - pipeline = RobotProcessor([to_batch_processor]) + pipeline = DataProcessorPipeline([to_batch_processor]) # Create unbatched observation observation = { @@ -283,9 +283,9 @@ def test_serialization_methods(): def test_save_and_load_pretrained(): - """Test saving and loading ToBatchProcessor with RobotProcessor.""" + """Test saving and loading ToBatchProcessor with DataProcessorPipeline.""" processor = ToBatchProcessor() - pipeline = RobotProcessor([processor], name="BatchPipeline") + pipeline = DataProcessorPipeline([processor], name="BatchPipeline") with tempfile.TemporaryDirectory() as tmp_dir: # Save pipeline @@ -296,7 +296,7 @@ def test_save_and_load_pretrained(): assert config_path.exists() # Load pipeline - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) assert loaded_pipeline.name == "BatchPipeline" assert len(loaded_pipeline) == 1 @@ -323,11 +323,11 @@ def test_registry_functionality(): def test_registry_based_save_load(): """Test saving and loading using registry name.""" processor = ToBatchProcessor() - pipeline = RobotProcessor([processor]) + pipeline = DataProcessorPipeline([processor]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) # Verify the loaded processor works observation = { diff --git a/tests/processor/test_classifier_processor.py b/tests/processor/test_classifier_processor.py index ddc8739849..07abccfbb5 100644 --- a/tests/processor/test_classifier_processor.py +++ b/tests/processor/test_classifier_processor.py @@ -24,7 +24,7 @@ from lerobot.constants import OBS_IMAGE, OBS_STATE from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor -from lerobot.processor import DeviceProcessor, IdentityProcessor, NormalizerProcessor, RobotProcessor +from lerobot.processor import DataProcessorPipeline, DeviceProcessor, IdentityProcessor, NormalizerProcessor from lerobot.processor.pipeline import TransitionKey @@ -237,7 +237,7 @@ def test_classifier_processor_save_and_load(): preprocessor.save_pretrained(tmpdir) # Load preprocessor - loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir) + loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir) # Test that loaded processor works observation = { diff --git a/tests/processor/test_device_processor.py b/tests/processor/test_device_processor.py index 354eb22ec8..3f86d06fb8 100644 --- a/tests/processor/test_device_processor.py +++ b/tests/processor/test_device_processor.py @@ -19,7 +19,7 @@ import torch from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.processor import DeviceProcessor, RobotProcessor +from lerobot.processor import DataProcessorPipeline, DeviceProcessor from lerobot.processor.pipeline import TransitionKey @@ -303,7 +303,7 @@ def test_features(): def test_integration_with_robot_processor(): - """Test integration with RobotProcessor.""" + """Test integration with DataProcessorPipeline.""" from lerobot.constants import OBS_STATE from lerobot.processor import ToBatchProcessor @@ -311,7 +311,7 @@ def test_integration_with_robot_processor(): device_processor = DeviceProcessor(device="cpu") batch_processor = ToBatchProcessor() - processor = RobotProcessor(steps=[batch_processor, device_processor], name="test_pipeline") + processor = DataProcessorPipeline(steps=[batch_processor, device_processor], name="test_pipeline") # Create test data observation = {OBS_STATE: torch.randn(10)} @@ -332,14 +332,14 @@ def test_save_and_load_pretrained(): """Test saving and loading processor with DeviceProcessor.""" device = "cuda:0" if torch.cuda.is_available() else "cpu" processor = DeviceProcessor(device=device, float_dtype="float16") - robot_processor = RobotProcessor(steps=[processor], name="device_test_processor") + robot_processor = DataProcessorPipeline(steps=[processor], name="device_test_processor") with tempfile.TemporaryDirectory() as tmpdir: # Save robot_processor.save_pretrained(tmpdir) # Load - loaded_processor = RobotProcessor.from_pretrained(tmpdir) + loaded_processor = DataProcessorPipeline.from_pretrained(tmpdir) assert len(loaded_processor.steps) == 1 loaded_device_processor = loaded_processor.steps[0] @@ -978,7 +978,7 @@ def test_policy_processor_integration(): norm_map = {FeatureType.STATE: NormalizationMode.MEAN_STD, FeatureType.ACTION: NormalizationMode.MEAN_STD} # Create input processor (preprocessor) that moves to GPU - input_processor = RobotProcessor( + input_processor = DataProcessorPipeline( steps=[ NormalizerProcessor(features=features, norm_map=norm_map, stats=stats), ToBatchProcessor(), @@ -988,7 +988,7 @@ def test_policy_processor_integration(): ) # Create output processor (postprocessor) that moves to CPU - output_processor = RobotProcessor( + output_processor = DataProcessorPipeline( steps=[ DeviceProcessor(device="cpu"), UnnormalizerProcessor(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats), diff --git a/tests/processor/test_diffusion_processor.py b/tests/processor/test_diffusion_processor.py index e6d3ea5908..b4f62e2037 100644 --- a/tests/processor/test_diffusion_processor.py +++ b/tests/processor/test_diffusion_processor.py @@ -25,10 +25,10 @@ from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -245,7 +245,7 @@ def test_diffusion_processor_save_and_load(): preprocessor.save_pretrained(tmpdir) # Load preprocessor - loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir) + loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir) # Test that loaded processor works observation = { diff --git a/tests/processor/test_normalize_processor.py b/tests/processor/test_normalize_processor.py index 6b904eee7c..261153a837 100644 --- a/tests/processor/test_normalize_processor.py +++ b/tests/processor/test_normalize_processor.py @@ -26,7 +26,7 @@ _convert_stats_to_tensors, hotswap_stats, ) -from lerobot.processor.pipeline import IdentityProcessor, RobotProcessor, TransitionKey +from lerobot.processor.pipeline import DataProcessorPipeline, IdentityProcessor, TransitionKey def create_transition( @@ -505,8 +505,8 @@ def test_get_config(full_stats): def test_integration_with_robot_processor(normalizer_processor): - """Test integration with RobotProcessor pipeline""" - robot_processor = RobotProcessor([normalizer_processor]) + """Test integration with DataProcessorPipeline pipeline""" + robot_processor = DataProcessorPipeline([normalizer_processor]) observation = { "observation.image": torch.tensor([0.7, 0.5, 0.3]), @@ -1007,7 +1007,7 @@ def test_hotswap_stats_basic_functionality(): identity = IdentityProcessor() # Create robot processor - robot_processor = RobotProcessor(steps=[normalizer, unnormalizer, identity]) + robot_processor = DataProcessorPipeline(steps=[normalizer, unnormalizer, identity]) # Hotswap stats new_processor = hotswap_stats(robot_processor, new_stats) @@ -1044,7 +1044,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) - original_processor = RobotProcessor(steps=[normalizer]) + original_processor = DataProcessorPipeline(steps=[normalizer]) # Store reference to original stats original_stats_reference = original_processor.steps[0].stats @@ -1087,7 +1087,7 @@ def test_hotswap_stats_only_affects_normalizer_steps(): unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) identity = IdentityProcessor() - robot_processor = RobotProcessor(steps=[normalizer, identity, unnormalizer]) + robot_processor = DataProcessorPipeline(steps=[normalizer, identity, unnormalizer]) # Hotswap stats new_processor = hotswap_stats(robot_processor, new_stats) @@ -1114,7 +1114,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) - robot_processor = RobotProcessor(steps=[normalizer]) + robot_processor = DataProcessorPipeline(steps=[normalizer]) # Hotswap with empty stats new_processor = hotswap_stats(robot_processor, empty_stats) @@ -1131,7 +1131,7 @@ def test_hotswap_stats_no_normalizer_steps(): } # Create processor with only identity steps - robot_processor = RobotProcessor(steps=[IdentityProcessor(), IdentityProcessor()]) + robot_processor = DataProcessorPipeline(steps=[IdentityProcessor(), IdentityProcessor()]) # Hotswap stats - should work without error new_processor = hotswap_stats(robot_processor, stats) @@ -1170,7 +1170,7 @@ def test_hotswap_stats_preserves_other_attributes(): normalize_observation_keys=normalize_observation_keys, eps=eps, ) - robot_processor = RobotProcessor(steps=[normalizer]) + robot_processor = DataProcessorPipeline(steps=[normalizer]) # Hotswap stats new_processor = hotswap_stats(robot_processor, new_stats) @@ -1213,7 +1213,7 @@ def test_hotswap_stats_multiple_normalizer_types(): unnormalizer1 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) unnormalizer2 = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) - robot_processor = RobotProcessor(steps=[normalizer1, unnormalizer1, normalizer2, unnormalizer2]) + robot_processor = DataProcessorPipeline(steps=[normalizer1, unnormalizer1, normalizer2, unnormalizer2]) # Hotswap stats new_processor = hotswap_stats(robot_processor, new_stats) @@ -1261,7 +1261,7 @@ def test_hotswap_stats_with_different_data_types(): } normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) - robot_processor = RobotProcessor(steps=[normalizer]) + robot_processor = DataProcessorPipeline(steps=[normalizer]) # Hotswap stats new_processor = hotswap_stats(robot_processor, new_stats) @@ -1317,7 +1317,7 @@ def test_hotswap_stats_functional_test(): # Create original processor normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=initial_stats) - original_processor = RobotProcessor(steps=[normalizer]) + original_processor = DataProcessorPipeline(steps=[normalizer]) # Process with original stats original_result = original_processor(transition) diff --git a/tests/processor/test_pipeline.py b/tests/processor/test_pipeline.py index 42a8eb538f..df1f018f02 100644 --- a/tests/processor/test_pipeline.py +++ b/tests/processor/test_pipeline.py @@ -27,7 +27,7 @@ from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features -from lerobot.processor import EnvTransition, ProcessorStepRegistry, RobotProcessor +from lerobot.processor import DataProcessorPipeline, EnvTransition, ProcessorStepRegistry from lerobot.processor.pipeline import TransitionKey from tests.conftest import assert_contract_is_typed @@ -176,7 +176,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po def test_empty_pipeline(): """Test pipeline with no steps.""" - pipeline = RobotProcessor() + pipeline = DataProcessorPipeline() transition = create_transition() result = pipeline(transition) @@ -188,7 +188,7 @@ def test_empty_pipeline(): def test_single_step_pipeline(): """Test pipeline with a single step.""" step = MockStep("test_step") - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) transition = create_transition() result = pipeline(transition) @@ -205,7 +205,7 @@ def test_multiple_steps_pipeline(): """Test pipeline with multiple steps.""" step1 = MockStep("step1") step2 = MockStep("step2") - pipeline = RobotProcessor([step1, step2]) + pipeline = DataProcessorPipeline([step1, step2]) transition = create_transition() result = pipeline(transition) @@ -217,7 +217,7 @@ def test_multiple_steps_pipeline(): def test_invalid_transition_format(): """Test pipeline with invalid transition format.""" - pipeline = RobotProcessor([MockStep()]) + pipeline = DataProcessorPipeline([MockStep()]) # Test with wrong type (tuple instead of dict) with pytest.raises(ValueError, match="EnvTransition must be a dictionary"): @@ -232,7 +232,7 @@ def test_step_through(): """Test step_through method with dict input.""" step1 = MockStep("step1") step2 = MockStep("step2") - pipeline = RobotProcessor([step1, step2]) + pipeline = DataProcessorPipeline([step1, step2]) transition = create_transition() @@ -253,7 +253,7 @@ def test_step_through_with_dict(): """Test step_through method with dict input.""" step1 = MockStep("step1") step2 = MockStep("step2") - pipeline = RobotProcessor([step1, step2]) + pipeline = DataProcessorPipeline([step1, step2]) batch = { "observation.image": None, @@ -292,7 +292,7 @@ def test_step_through_with_dict(): def test_step_through_no_hooks(): """Test that step_through doesn't execute hooks.""" step = MockStep("test_step") - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) hook_calls = [] @@ -327,7 +327,7 @@ def test_indexing(): """Test pipeline indexing.""" step1 = MockStep("step1") step2 = MockStep("step2") - pipeline = RobotProcessor([step1, step2]) + pipeline = DataProcessorPipeline([step1, step2]) # Test integer indexing assert pipeline[0] is step1 @@ -335,7 +335,7 @@ def test_indexing(): # Test slice indexing sub_pipeline = pipeline[0:1] - assert isinstance(sub_pipeline, RobotProcessor) + assert isinstance(sub_pipeline, DataProcessorPipeline) assert len(sub_pipeline) == 1 assert sub_pipeline[0] is step1 @@ -343,7 +343,7 @@ def test_indexing(): def test_hooks(): """Test before/after step hooks.""" step = MockStep("test_step") - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) before_calls = [] after_calls = [] @@ -367,7 +367,7 @@ def after_hook(idx: int, transition: EnvTransition): def test_unregister_hooks(): """Test unregistering hooks from the pipeline.""" step = MockStep("test_step") - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) # Test before_step_hook before_calls = [] @@ -406,7 +406,7 @@ def after_hook(idx: int, transition: EnvTransition): def test_unregister_nonexistent_hook(): """Test error handling when unregistering hooks that don't exist.""" - pipeline = RobotProcessor([MockStep()]) + pipeline = DataProcessorPipeline([MockStep()]) def some_hook(idx: int, transition: EnvTransition): pass @@ -424,7 +424,7 @@ def reset_hook(): def test_multiple_hooks_and_selective_unregister(): """Test registering multiple hooks and selectively unregistering them.""" - pipeline = RobotProcessor([MockStep("step1"), MockStep("step2")]) + pipeline = DataProcessorPipeline([MockStep("step1"), MockStep("step2")]) calls_1 = [] calls_2 = [] @@ -470,7 +470,7 @@ def hook3(idx: int, transition: EnvTransition): def test_hook_execution_order_documentation(): """Test and document that hooks are executed sequentially in registration order.""" - pipeline = RobotProcessor([MockStep("step")]) + pipeline = DataProcessorPipeline([MockStep("step")]) execution_order = [] @@ -522,7 +522,7 @@ def test_save_and_load_pretrained(): step1.counter = 5 step2.counter = 10 - pipeline = RobotProcessor([step1, step2], name="TestPipeline") + pipeline = DataProcessorPipeline([step1, step2], name="TestPipeline") with tempfile.TemporaryDirectory() as tmp_dir: # Save pipeline @@ -544,7 +544,7 @@ def test_save_and_load_pretrained(): assert config["steps"][1]["config"]["counter"] == 10 # Load pipeline - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) assert loaded_pipeline.name == "TestPipeline" assert len(loaded_pipeline) == 2 @@ -557,7 +557,7 @@ def test_save_and_load_pretrained(): def test_step_without_optional_methods(): """Test pipeline with steps that don't implement optional methods.""" step = MockStepWithoutOptionalMethods(multiplier=3.0) - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) transition = create_transition(reward=2.0) result = pipeline(transition) @@ -570,14 +570,14 @@ def test_step_without_optional_methods(): # Save/load should work even without optional methods with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) assert len(loaded_pipeline) == 1 def test_mixed_json_and_tensor_state(): """Test step with both JSON attributes and tensor state.""" step = MockStepWithTensorState(name="stats", learning_rate=0.05, window_size=5) - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) # Process some transitions with rewards for i in range(10): @@ -593,13 +593,13 @@ def test_mixed_json_and_tensor_state(): pipeline.save_pretrained(tmp_dir) # Check that both config and state files were created - config_path = Path(tmp_dir) / "robotprocessor.json" # Default name is "RobotProcessor" - state_path = Path(tmp_dir) / "robotprocessor_step_0.safetensors" + config_path = Path(tmp_dir) / "dataprocessorpipeline.json" # Default name is "DataProcessorPipeline" + state_path = Path(tmp_dir) / "dataprocessorpipeline_step_0.safetensors" assert config_path.exists() assert state_path.exists() # Load and verify - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) loaded_step = loaded_pipeline.steps[0] # Check JSON attributes were restored @@ -860,7 +860,7 @@ def test_from_pretrained_with_overrides(): env_step = MockStepWithNonSerializableParam(name="env_step", multiplier=2.0) registered_step = RegisteredMockStep(value=100, device="cpu") - pipeline = RobotProcessor([env_step, registered_step], name="TestOverrides") + pipeline = DataProcessorPipeline([env_step, registered_step], name="TestOverrides") with tempfile.TemporaryDirectory() as tmp_dir: # Save the pipeline @@ -878,7 +878,7 @@ def test_from_pretrained_with_overrides(): "registered_mock_step": {"device": "cuda", "value": 200}, } - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) # Verify the pipeline was loaded correctly assert len(loaded_pipeline) == 2 @@ -904,7 +904,7 @@ def test_from_pretrained_with_partial_overrides(): step1 = MockStepWithNonSerializableParam(name="step1", multiplier=1.0) step2 = MockStepWithNonSerializableParam(name="step2", multiplier=2.0) - pipeline = RobotProcessor([step1, step2]) + pipeline = DataProcessorPipeline([step1, step2]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -914,7 +914,7 @@ def test_from_pretrained_with_partial_overrides(): # The current implementation applies overrides to ALL steps with the same class name # Both steps will get the override - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) transition = create_transition(reward=1.0) result = loaded_pipeline(transition) @@ -928,7 +928,7 @@ def test_from_pretrained_with_partial_overrides(): def test_from_pretrained_invalid_override_key(): """Test that invalid override keys raise KeyError.""" step = MockStepWithNonSerializableParam() - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -937,13 +937,13 @@ def test_from_pretrained_invalid_override_key(): overrides = {"NonExistentStep": {"param": "value"}} with pytest.raises(KeyError, match="Override keys.*do not match any step"): - RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) def test_from_pretrained_multiple_invalid_override_keys(): """Test that multiple invalid override keys are reported.""" step = MockStepWithNonSerializableParam() - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -952,7 +952,7 @@ def test_from_pretrained_multiple_invalid_override_keys(): overrides = {"NonExistentStep1": {"param": "value1"}, "NonExistentStep2": {"param": "value2"}} with pytest.raises(KeyError) as exc_info: - RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) error_msg = str(exc_info.value) assert "NonExistentStep1" in error_msg @@ -963,7 +963,7 @@ def test_from_pretrained_multiple_invalid_override_keys(): def test_from_pretrained_registered_step_override(): """Test overriding registered steps using registry names.""" registered_step = RegisteredMockStep(value=50, device="cpu") - pipeline = RobotProcessor([registered_step]) + pipeline = DataProcessorPipeline([registered_step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -971,7 +971,7 @@ def test_from_pretrained_registered_step_override(): # Override using registry name overrides = {"registered_mock_step": {"value": 999, "device": "cuda"}} - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) # Test that overrides were applied transition = create_transition() @@ -987,7 +987,7 @@ def test_from_pretrained_mixed_registered_and_unregistered(): unregistered_step = MockStepWithNonSerializableParam(name="unregistered", multiplier=1.0) registered_step = RegisteredMockStep(value=10, device="cpu") - pipeline = RobotProcessor([unregistered_step, registered_step]) + pipeline = DataProcessorPipeline([unregistered_step, registered_step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -999,7 +999,7 @@ def test_from_pretrained_mixed_registered_and_unregistered(): "registered_mock_step": {"value": 777}, } - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) # Test both steps transition = create_transition(reward=2.0) @@ -1014,13 +1014,13 @@ def test_from_pretrained_mixed_registered_and_unregistered(): def test_from_pretrained_no_overrides(): """Test that from_pretrained works without overrides (backward compatibility).""" step = MockStepWithNonSerializableParam(name="no_override", multiplier=3.0) - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) # Load without overrides - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) assert len(loaded_pipeline) == 1 @@ -1034,13 +1034,13 @@ def test_from_pretrained_no_overrides(): def test_from_pretrained_empty_overrides(): """Test that from_pretrained works with empty overrides dict.""" step = MockStepWithNonSerializableParam(multiplier=2.0) - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) # Load with empty overrides - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides={}) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, overrides={}) assert len(loaded_pipeline) == 1 @@ -1054,7 +1054,7 @@ def test_from_pretrained_empty_overrides(): def test_from_pretrained_override_instantiation_error(): """Test that instantiation errors with overrides are properly reported.""" step = MockStepWithNonSerializableParam(multiplier=1.0) - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -1067,13 +1067,13 @@ def test_from_pretrained_override_instantiation_error(): } with pytest.raises(ValueError, match="Failed to instantiate processor step"): - RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) def test_from_pretrained_with_state_and_overrides(): """Test that overrides work correctly with steps that have tensor state.""" step = MockStepWithTensorState(name="tensor_step", learning_rate=0.01, window_size=5) - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) # Process some data to create state for i in range(10): @@ -1091,7 +1091,7 @@ def test_from_pretrained_with_state_and_overrides(): } } - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) loaded_step = loaded_pipeline.steps[0] # Check that config overrides were applied @@ -1110,7 +1110,7 @@ def test_from_pretrained_override_error_messages(): """Test that error messages for override failures are helpful.""" step1 = MockStepWithNonSerializableParam(name="step1") step2 = RegisteredMockStep() - pipeline = RobotProcessor([step1, step2]) + pipeline = DataProcessorPipeline([step1, step2]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -1119,7 +1119,7 @@ def test_from_pretrained_override_error_messages(): overrides = {"WrongStepName": {"param": "value"}} with pytest.raises(KeyError) as exc_info: - RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) error_msg = str(exc_info.value) assert "WrongStepName" in error_msg @@ -1130,20 +1130,20 @@ def test_from_pretrained_override_error_messages(): def test_repr_empty_processor(): """Test __repr__ with empty processor.""" - pipeline = RobotProcessor() + pipeline = DataProcessorPipeline() repr_str = repr(pipeline) - expected = "RobotProcessor(name='RobotProcessor', steps=0: [])" + expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=0: [])" assert repr_str == expected def test_repr_single_step(): """Test __repr__ with single step.""" step = MockStep("test_step") - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) repr_str = repr(pipeline) - expected = "RobotProcessor(name='RobotProcessor', steps=1: [MockStep])" + expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=1: [MockStep])" assert repr_str == expected @@ -1151,18 +1151,18 @@ def test_repr_multiple_steps_under_limit(): """Test __repr__ with 2-3 steps (all shown).""" step1 = MockStep("step1") step2 = MockStepWithoutOptionalMethods() - pipeline = RobotProcessor([step1, step2]) + pipeline = DataProcessorPipeline([step1, step2]) repr_str = repr(pipeline) - expected = "RobotProcessor(name='RobotProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods])" + expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=2: [MockStep, MockStepWithoutOptionalMethods])" assert repr_str == expected # Test with 3 steps (boundary case) step3 = MockStepWithTensorState() - pipeline = RobotProcessor([step1, step2, step3]) + pipeline = DataProcessorPipeline([step1, step2, step3]) repr_str = repr(pipeline) - expected = "RobotProcessor(name='RobotProcessor', steps=3: [MockStep, MockStepWithoutOptionalMethods, MockStepWithTensorState])" + expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=3: [MockStep, MockStepWithoutOptionalMethods, MockStepWithTensorState])" assert repr_str == expected @@ -1174,30 +1174,30 @@ def test_repr_many_steps_truncated(): step4 = MockModuleStep() step5 = MockNonModuleStepWithState() - pipeline = RobotProcessor([step1, step2, step3, step4, step5]) + pipeline = DataProcessorPipeline([step1, step2, step3, step4, step5]) repr_str = repr(pipeline) - expected = "RobotProcessor(name='RobotProcessor', steps=5: [MockStep, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])" + expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=5: [MockStep, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])" assert repr_str == expected def test_repr_with_custom_name(): """Test __repr__ with custom processor name.""" step = MockStep("test_step") - pipeline = RobotProcessor([step], name="CustomProcessor") + pipeline = DataProcessorPipeline([step], name="CustomProcessor") repr_str = repr(pipeline) - expected = "RobotProcessor(name='CustomProcessor', steps=1: [MockStep])" + expected = "DataProcessorPipeline(name='CustomProcessor', steps=1: [MockStep])" assert repr_str == expected def test_repr_with_seed(): """Test __repr__ with seed parameter.""" step = MockStep("test_step") - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) repr_str = repr(pipeline) - expected = "RobotProcessor(name='RobotProcessor', steps=1: [MockStep])" + expected = "DataProcessorPipeline(name='DataProcessorPipeline', steps=1: [MockStep])" assert repr_str == expected @@ -1205,20 +1205,22 @@ def test_repr_with_custom_name_and_seed(): """Test __repr__ with both custom name and seed.""" step1 = MockStep("step1") step2 = MockStepWithoutOptionalMethods() - pipeline = RobotProcessor([step1, step2], name="MyProcessor") + pipeline = DataProcessorPipeline([step1, step2], name="MyProcessor") repr_str = repr(pipeline) - expected = "RobotProcessor(name='MyProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods])" + expected = ( + "DataProcessorPipeline(name='MyProcessor', steps=2: [MockStep, MockStepWithoutOptionalMethods])" + ) assert repr_str == expected def test_repr_without_seed(): """Test __repr__ when seed is explicitly None (should not show seed).""" step = MockStep("test_step") - pipeline = RobotProcessor([step], name="TestProcessor") + pipeline = DataProcessorPipeline([step], name="TestProcessor") repr_str = repr(pipeline) - expected = "RobotProcessor(name='TestProcessor', steps=1: [MockStep])" + expected = "DataProcessorPipeline(name='TestProcessor', steps=1: [MockStep])" assert repr_str == expected @@ -1229,10 +1231,10 @@ def test_repr_various_step_types(): step3 = MockModuleStep() step4 = MockNonModuleStepWithState() - pipeline = RobotProcessor([step1, step2, step3, step4], name="MixedSteps") + pipeline = DataProcessorPipeline([step1, step2, step3, step4], name="MixedSteps") repr_str = repr(pipeline) - expected = "RobotProcessor(name='MixedSteps', steps=4: [MockStep, MockStepWithTensorState, ..., MockNonModuleStepWithState])" + expected = "DataProcessorPipeline(name='MixedSteps', steps=4: [MockStep, MockStepWithTensorState, ..., MockNonModuleStepWithState])" assert repr_str == expected @@ -1243,10 +1245,10 @@ def test_repr_edge_case_long_names(): step3 = MockStepWithTensorState() step4 = MockNonModuleStepWithState() - pipeline = RobotProcessor([step1, step2, step3, step4], name="LongNames") + pipeline = DataProcessorPipeline([step1, step2, step3, step4], name="LongNames") repr_str = repr(pipeline) - expected = "RobotProcessor(name='LongNames', steps=4: [MockStepWithNonSerializableParam, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])" + expected = "DataProcessorPipeline(name='LongNames', steps=4: [MockStepWithNonSerializableParam, MockStepWithoutOptionalMethods, ..., MockNonModuleStepWithState])" assert repr_str == expected @@ -1254,7 +1256,7 @@ def test_repr_edge_case_long_names(): def test_save_with_custom_config_filename(): """Test saving processor with custom config filename.""" step = MockStep("test") - pipeline = RobotProcessor([step], name="TestProcessor") + pipeline = DataProcessorPipeline([step], name="TestProcessor") with tempfile.TemporaryDirectory() as tmp_dir: # Save with custom filename @@ -1270,16 +1272,18 @@ def test_save_with_custom_config_filename(): assert config["name"] == "TestProcessor" # Load with specific filename - loaded = RobotProcessor.from_pretrained(tmp_dir, config_filename="my_custom_config.json") + loaded = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="my_custom_config.json") assert loaded.name == "TestProcessor" def test_multiple_processors_same_directory(): """Test saving multiple processors to the same directory with different config files.""" # Create different processors - preprocessor = RobotProcessor([MockStep("pre1"), MockStep("pre2")], name="preprocessor") + preprocessor = DataProcessorPipeline([MockStep("pre1"), MockStep("pre2")], name="preprocessor") - postprocessor = RobotProcessor([MockStepWithoutOptionalMethods(multiplier=0.5)], name="postprocessor") + postprocessor = DataProcessorPipeline( + [MockStepWithoutOptionalMethods(multiplier=0.5)], name="postprocessor" + ) with tempfile.TemporaryDirectory() as tmp_dir: # Save both to same directory @@ -1291,8 +1295,8 @@ def test_multiple_processors_same_directory(): assert (Path(tmp_dir) / "postprocessor.json").exists() # Load them back - loaded_pre = RobotProcessor.from_pretrained(tmp_dir, config_filename="preprocessor.json") - loaded_post = RobotProcessor.from_pretrained(tmp_dir, config_filename="postprocessor.json") + loaded_pre = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="preprocessor.json") + loaded_post = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="postprocessor.json") assert loaded_pre.name == "preprocessor" assert loaded_post.name == "postprocessor" @@ -1303,20 +1307,20 @@ def test_multiple_processors_same_directory(): def test_auto_detect_single_config(): """Test automatic config detection when there's only one JSON file.""" step = MockStepWithTensorState() - pipeline = RobotProcessor([step], name="SingleConfig") + pipeline = DataProcessorPipeline([step], name="SingleConfig") with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) # Load without specifying config_filename - loaded = RobotProcessor.from_pretrained(tmp_dir) + loaded = DataProcessorPipeline.from_pretrained(tmp_dir) assert loaded.name == "SingleConfig" def test_error_multiple_configs_no_filename(): """Test error when multiple configs exist and no filename specified.""" - proc1 = RobotProcessor([MockStep()], name="processor1") - proc2 = RobotProcessor([MockStep()], name="processor2") + proc1 = DataProcessorPipeline([MockStep()], name="processor1") + proc2 = DataProcessorPipeline([MockStep()], name="processor2") with tempfile.TemporaryDirectory() as tmp_dir: proc1.save_pretrained(tmp_dir) @@ -1324,7 +1328,7 @@ def test_error_multiple_configs_no_filename(): # Should raise error with pytest.raises(ValueError, match="Multiple .json files found"): - RobotProcessor.from_pretrained(tmp_dir) + DataProcessorPipeline.from_pretrained(tmp_dir) def test_state_file_naming_with_indices(): @@ -1334,7 +1338,7 @@ def test_state_file_naming_with_indices(): step2 = MockStepWithTensorState(name="norm2", window_size=10) step3 = MockModuleStep(input_dim=5) - pipeline = RobotProcessor([step1, step2, step3]) + pipeline = DataProcessorPipeline([step1, step2, step3]) # Process some data to create state for i in range(5): @@ -1350,9 +1354,9 @@ def test_state_file_naming_with_indices(): # Files should be named with pipeline name prefix and indices expected_names = [ - "robotprocessor_step_0.safetensors", - "robotprocessor_step_1.safetensors", - "robotprocessor_step_2.safetensors", + "dataprocessorpipeline_step_0.safetensors", + "dataprocessorpipeline_step_1.safetensors", + "dataprocessorpipeline_step_2.safetensors", ] actual_names = [f.name for f in state_files] assert actual_names == expected_names @@ -1391,7 +1395,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po # Create pipeline with registered steps step1 = TestStatefulStep(1) step2 = TestStatefulStep(2) - pipeline = RobotProcessor([step1, step2]) + pipeline = DataProcessorPipeline([step1, step2]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -1402,8 +1406,8 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po # Should include pipeline name, index and registry name expected_names = [ - "robotprocessor_step_0_test_stateful_step.safetensors", - "robotprocessor_step_1_test_stateful_step.safetensors", + "dataprocessorpipeline_step_0_test_stateful_step.safetensors", + "dataprocessorpipeline_step_1_test_stateful_step.safetensors", ] actual_names = [f.name for f in state_files] assert actual_names == expected_names @@ -1446,13 +1450,13 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po try: step = ComplexConfigStep() - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) # Load with nested override - loaded = RobotProcessor.from_pretrained( + loaded = DataProcessorPipeline.from_pretrained( tmp_dir, overrides={"complex_config_step": {"nested_config": {"level1": {"level2": "overridden"}}}}, ) @@ -1468,13 +1472,13 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po def test_override_preserves_defaults(): """Test that overrides only affect specified parameters.""" step = MockStepWithNonSerializableParam(name="test", multiplier=2.0) - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) # Override only one parameter - loaded = RobotProcessor.from_pretrained( + loaded = DataProcessorPipeline.from_pretrained( tmp_dir, overrides={ "MockStepWithNonSerializableParam": { @@ -1492,7 +1496,7 @@ def test_override_preserves_defaults(): def test_override_type_validation(): """Test that type errors in overrides are caught properly.""" step = MockStepWithTensorState(learning_rate=0.01) - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -1505,7 +1509,7 @@ def test_override_type_validation(): } with pytest.raises(ValueError, match="Failed to instantiate"): - RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) + DataProcessorPipeline.from_pretrained(tmp_dir, overrides=overrides) def test_override_with_callables(): @@ -1538,7 +1542,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po try: step = CallableStep() - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -1552,7 +1556,7 @@ def double_values(x): return x # Load with callable override - loaded = RobotProcessor.from_pretrained( + loaded = DataProcessorPipeline.from_pretrained( tmp_dir, overrides={"callable_step": {"transform_fn": double_values}} ) @@ -1568,13 +1572,13 @@ def test_override_multiple_same_class_warning(): """Test behavior when multiple steps of same class exist.""" step1 = MockStepWithNonSerializableParam(name="step1", multiplier=1.0) step2 = MockStepWithNonSerializableParam(name="step2", multiplier=2.0) - pipeline = RobotProcessor([step1, step2]) + pipeline = DataProcessorPipeline([step1, step2]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) # Override affects all instances of the class - loaded = RobotProcessor.from_pretrained( + loaded = DataProcessorPipeline.from_pretrained( tmp_dir, overrides={"MockStepWithNonSerializableParam": {"multiplier": 10.0}} ) @@ -1590,7 +1594,7 @@ def test_override_multiple_same_class_warning(): def test_config_filename_special_characters(): """Test config filenames with special characters are sanitized.""" # Processor name with special characters - pipeline = RobotProcessor([MockStep()], name="My/Processor\\With:Special*Chars") + pipeline = DataProcessorPipeline([MockStep()], name="My/Processor\\With:Special*Chars") with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) @@ -1608,10 +1612,10 @@ def test_state_file_naming_with_multiple_processors(): """Test that state files are properly prefixed with pipeline names to avoid conflicts.""" # Create two processors with state step1 = MockStepWithTensorState(name="norm", window_size=5) - preprocessor = RobotProcessor([step1], name="PreProcessor") + preprocessor = DataProcessorPipeline([step1], name="PreProcessor") step2 = MockStepWithTensorState(name="norm", window_size=10) - postprocessor = RobotProcessor([step2], name="PostProcessor") + postprocessor = DataProcessorPipeline([step2], name="PostProcessor") # Process some data to create state for i in range(3): @@ -1631,8 +1635,8 @@ def test_state_file_naming_with_multiple_processors(): assert (Path(tmp_dir) / "postprocessor_step_0.safetensors").exists() # Load both back and verify they work correctly - loaded_pre = RobotProcessor.from_pretrained(tmp_dir, config_filename="preprocessor.json") - loaded_post = RobotProcessor.from_pretrained(tmp_dir, config_filename="postprocessor.json") + loaded_pre = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="preprocessor.json") + loaded_post = DataProcessorPipeline.from_pretrained(tmp_dir, config_filename="postprocessor.json") assert loaded_pre.name == "PreProcessor" assert loaded_post.name == "PostProcessor" @@ -1773,14 +1777,14 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po try: step = DeviceAwareStep(device="cpu") - pipeline = RobotProcessor([step]) + pipeline = DataProcessorPipeline([step]) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) # Override device if torch.cuda.is_available(): - loaded = RobotProcessor.from_pretrained( + loaded = DataProcessorPipeline.from_pretrained( tmp_dir, overrides={"device_aware_step": {"device": "cuda:0"}} ) @@ -1799,16 +1803,16 @@ def test_from_pretrained_nonexistent_path(): # Test with an invalid repo ID (too many slashes) - caught by HF validation with pytest.raises(HFValidationError): - RobotProcessor.from_pretrained("/path/that/does/not/exist") + DataProcessorPipeline.from_pretrained("/path/that/does/not/exist") # Test with a non-existent but valid Hub repo format with pytest.raises((FileNotFoundError, HfHubHTTPError)): - RobotProcessor.from_pretrained("nonexistent-user/nonexistent-repo") + DataProcessorPipeline.from_pretrained("nonexistent-user/nonexistent-repo") # Test with a local directory that exists but has no config files with tempfile.TemporaryDirectory() as tmp_dir: with pytest.raises(FileNotFoundError, match="No .json configuration files found"): - RobotProcessor.from_pretrained(tmp_dir) + DataProcessorPipeline.from_pretrained(tmp_dir) def test_save_load_with_custom_converter_functions(): @@ -1837,13 +1841,15 @@ def custom_to_output(transition): } # Create processor with custom converters - pipeline = RobotProcessor([MockStep()], to_transition=custom_to_transition, to_output=custom_to_output) + pipeline = DataProcessorPipeline( + [MockStep()], to_transition=custom_to_transition, to_output=custom_to_output + ) with tempfile.TemporaryDirectory() as tmp_dir: pipeline.save_pretrained(tmp_dir) # Load - should use default converters - loaded = RobotProcessor.from_pretrained(tmp_dir) + loaded = DataProcessorPipeline.from_pretrained(tmp_dir) # Verify it uses default converters by checking with standard batch format batch = { @@ -1876,7 +1882,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po def test_construction_rejects_step_without_call(): with pytest.raises(TypeError, match=r"must define __call__"): - RobotProcessor([NonCallableStep()]) + DataProcessorPipeline([NonCallableStep()]) @dataclass @@ -1935,7 +1941,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po def test_features_orders_and_merges(policy_feature_factory): - p = RobotProcessor( + p = DataProcessorPipeline( [ FeatureContractAddStep("a", policy_feature_factory(FeatureType.STATE, (1,))), FeatureContractMutateStep("a", lambda v: PolicyFeature(type=v.type, shape=(3,))), @@ -1954,7 +1960,7 @@ def test_features_respects_initial_without_mutation(policy_feature_factory): "seed": policy_feature_factory(FeatureType.STATE, (7,)), "nested": policy_feature_factory(FeatureType.ENV, (0,)), } - p = RobotProcessor( + p = DataProcessorPipeline( [ FeatureContractMutateStep("seed", lambda v: PolicyFeature(type=v.type, shape=(v.shape[0] + 1,))), FeatureContractMutateStep( @@ -1987,12 +1993,12 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po features["order"] = PolicyFeature(type=pf.type, shape=pf.shape + (code,)) return features - out = RobotProcessor([Track("A"), Track("B"), Track("C")]).transform_features({}) + out = DataProcessorPipeline([Track("A"), Track("B"), Track("C")]).transform_features({}) assert out["order"].shape == (1, 2, 3) def test_features_remove_key(policy_feature_factory): - p = RobotProcessor( + p = DataProcessorPipeline( [ FeatureContractAddStep("a", policy_feature_factory(FeatureType.STATE, (1,))), FeatureContractRemoveStep("a"), @@ -2007,7 +2013,7 @@ def test_features_remove_from_initial(policy_feature_factory): "keep": policy_feature_factory(FeatureType.STATE, (1,)), "drop": policy_feature_factory(FeatureType.STATE, (1,)), } - p = RobotProcessor([FeatureContractRemoveStep("drop")]) + p = DataProcessorPipeline([FeatureContractRemoveStep("drop")]) out = p.transform_features(initial_features=initial) assert "drop" not in out and out["keep"] == initial["keep"] @@ -2049,7 +2055,7 @@ def transform_features(self, features: dict) -> dict: def test_aggregate_joint_action_only(): - rp = RobotProcessor([AddActionEEAndJointFeatures()]) + rp = DataProcessorPipeline([AddActionEEAndJointFeatures()]) initial = {"front": (480, 640, 3)} out = aggregate_pipeline_dataset_features( @@ -2067,7 +2073,7 @@ def test_aggregate_joint_action_only(): def test_aggregate_ee_action_and_observation_with_videos(): - rp = RobotProcessor([AddActionEEAndJointFeatures(), AddObservationStateFeatures()]) + rp = DataProcessorPipeline([AddActionEEAndJointFeatures(), AddObservationStateFeatures()]) initial = {"front": (480, 640, 3), "side": (720, 1280, 3)} out = aggregate_pipeline_dataset_features( @@ -2097,7 +2103,7 @@ def test_aggregate_ee_action_and_observation_with_videos(): def test_aggregate_both_action_types(): - rp = RobotProcessor([AddActionEEAndJointFeatures()]) + rp = DataProcessorPipeline([AddActionEEAndJointFeatures()]) out = aggregate_pipeline_dataset_features( pipeline=rp, initial_features={}, @@ -2112,7 +2118,7 @@ def test_aggregate_both_action_types(): def test_aggregate_images_when_use_videos_false(): - rp = RobotProcessor([AddObservationStateFeatures(add_front_image=True)]) + rp = DataProcessorPipeline([AddObservationStateFeatures(add_front_image=True)]) initial = {"back": (480, 640, 3)} out = aggregate_pipeline_dataset_features( @@ -2129,7 +2135,7 @@ def test_aggregate_images_when_use_videos_false(): def test_aggregate_images_when_use_videos_true(): - rp = RobotProcessor([AddObservationStateFeatures(add_front_image=True)]) + rp = DataProcessorPipeline([AddObservationStateFeatures(add_front_image=True)]) initial = {"back": (480, 640, 3)} out = aggregate_pipeline_dataset_features( @@ -2151,7 +2157,9 @@ def test_aggregate_images_when_use_videos_true(): def test_initial_camera_not_overridden_by_step_image(): # Step explicitly sets a different front image shape; initial has another shape. # aggregate_pipeline_dataset_features should keep the step's value (setdefault behavior on initial cams). - rp = RobotProcessor([AddObservationStateFeatures(add_front_image=True, front_image_shape=(240, 320, 3))]) + rp = DataProcessorPipeline( + [AddObservationStateFeatures(add_front_image=True, front_image_shape=(240, 320, 3))] + ) initial = {"front": (480, 640, 3)} # should NOT override the step-provided (240, 320, 3) out = aggregate_pipeline_dataset_features( diff --git a/tests/processor/test_rename_processor.py b/tests/processor/test_rename_processor.py index 4efb249dd7..e31d706262 100644 --- a/tests/processor/test_rename_processor.py +++ b/tests/processor/test_rename_processor.py @@ -20,7 +20,7 @@ import torch from lerobot.configs.types import FeatureType -from lerobot.processor import ProcessorStepRegistry, RenameProcessor, RobotProcessor, TransitionKey +from lerobot.processor import DataProcessorPipeline, ProcessorStepRegistry, RenameProcessor, TransitionKey from lerobot.processor.rename_processor import rename_stats from tests.conftest import assert_contract_is_typed @@ -181,14 +181,14 @@ def test_state_dict(): def test_integration_with_robot_processor(): - """Test integration with RobotProcessor pipeline.""" + """Test integration with DataProcessorPipeline pipeline.""" rename_map = { "agent_pos": "observation.state", "pixels": "observation.image", } rename_processor = RenameProcessor(rename_map=rename_map) - pipeline = RobotProcessor([rename_processor]) + pipeline = DataProcessorPipeline([rename_processor]) observation = { "agent_pos": np.array([1.0, 2.0, 3.0]), @@ -215,13 +215,13 @@ def test_integration_with_robot_processor(): def test_save_and_load_pretrained(): - """Test saving and loading processor with RobotProcessor.""" + """Test saving and loading processor with DataProcessorPipeline.""" rename_map = { "old_state": "observation.state", "old_image": "observation.image", } processor = RenameProcessor(rename_map=rename_map) - pipeline = RobotProcessor([processor], name="TestRenameProcessor") + pipeline = DataProcessorPipeline([processor], name="TestRenameProcessor") with tempfile.TemporaryDirectory() as tmp_dir: # Save pipeline @@ -236,7 +236,7 @@ def test_save_and_load_pretrained(): assert len(state_files) == 0 # Load pipeline - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) assert loaded_pipeline.name == "TestRenameProcessor" assert len(loaded_pipeline) == 1 @@ -277,7 +277,7 @@ def test_registry_functionality(): def test_registry_based_save_load(): """Test save/load using registry name instead of module path.""" processor = RenameProcessor(rename_map={"key1": "renamed_key1"}) - pipeline = RobotProcessor([processor]) + pipeline = DataProcessorPipeline([processor]) with tempfile.TemporaryDirectory() as tmp_dir: # Save and load @@ -286,7 +286,9 @@ def test_registry_based_save_load(): # Verify config uses registry name import json - with open(Path(tmp_dir) / "robotprocessor.json") as f: # Default name is "RobotProcessor" + with open( + Path(tmp_dir) / "dataprocessorpipeline.json" + ) as f: # Default name is "DataProcessorPipeline" config = json.load(f) assert "registry_name" in config["steps"][0] @@ -294,7 +296,7 @@ def test_registry_based_save_load(): assert "class" not in config["steps"][0] # Should use registry, not module path # Load should work - loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir) + loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) loaded_processor = loaded_pipeline.steps[0] assert isinstance(loaded_processor, RenameProcessor) assert loaded_processor.rename_map == {"key1": "renamed_key1"} @@ -318,7 +320,7 @@ def test_chained_rename_processors(): } ) - pipeline = RobotProcessor([processor1, processor2]) + pipeline = DataProcessorPipeline([processor1, processor2]) observation = { "pos": np.array([1.0, 2.0]), @@ -452,7 +454,7 @@ def test_features_chained_processors(policy_feature_factory): processor2 = RenameProcessor( rename_map={"agent_position": "observation.state", "camera_image": "observation.image"} ) - pipeline = RobotProcessor([processor1, processor2]) + pipeline = DataProcessorPipeline([processor1, processor2]) spec = { "pos": policy_feature_factory(FeatureType.STATE, (7,)), diff --git a/tests/processor/test_sac_processor.py b/tests/processor/test_sac_processor.py index 33f3330dcd..23e859fbe6 100644 --- a/tests/processor/test_sac_processor.py +++ b/tests/processor/test_sac_processor.py @@ -25,10 +25,10 @@ from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -232,7 +232,7 @@ def test_sac_processor_save_and_load(): preprocessor.save_pretrained(tmpdir) # Load preprocessor - loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir) + loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir) # Test that loaded processor works observation = {OBS_STATE: torch.randn(10)} diff --git a/tests/processor/test_tdmpc_processor.py b/tests/processor/test_tdmpc_processor.py index 473fc93dbe..ac95791d5a 100644 --- a/tests/processor/test_tdmpc_processor.py +++ b/tests/processor/test_tdmpc_processor.py @@ -25,10 +25,10 @@ from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -252,7 +252,7 @@ def test_tdmpc_processor_save_and_load(): preprocessor.save_pretrained(tmpdir) # Load preprocessor - loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir) + loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir) # Test that loaded processor works observation = { diff --git a/tests/processor/test_tokenizer_processor.py b/tests/processor/test_tokenizer_processor.py index 300191d863..b25560400d 100644 --- a/tests/processor/test_tokenizer_processor.py +++ b/tests/processor/test_tokenizer_processor.py @@ -10,7 +10,7 @@ from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.constants import OBS_LANGUAGE -from lerobot.processor.pipeline import RobotProcessor, TransitionKey +from lerobot.processor.pipeline import DataProcessorPipeline, TransitionKey from lerobot.processor.tokenizer_processor import TokenizerProcessor from tests.utils import require_package @@ -384,12 +384,12 @@ def test_reset_method(mock_auto_tokenizer): @require_package("transformers") @patch("lerobot.processor.tokenizer_processor.AutoTokenizer") def test_integration_with_robot_processor(mock_auto_tokenizer): - """Test integration with RobotProcessor.""" + """Test integration with DataProcessorPipeline.""" 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) - robot_processor = RobotProcessor([tokenizer_processor]) + robot_processor = DataProcessorPipeline([tokenizer_processor]) transition = create_transition( observation={"state": torch.tensor([1.0, 2.0])}, @@ -427,14 +427,14 @@ def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer): tokenizer_name="test-tokenizer", max_length=32, task_key="instruction" ) - robot_processor = RobotProcessor([original_processor]) + robot_processor = DataProcessorPipeline([original_processor]) with tempfile.TemporaryDirectory() as temp_dir: # Save processor robot_processor.save_pretrained(temp_dir) # Load processor - tokenizer will be recreated from saved config - loaded_processor = RobotProcessor.from_pretrained(temp_dir) + loaded_processor = DataProcessorPipeline.from_pretrained(temp_dir) # Test that loaded processor works transition = create_transition( @@ -456,14 +456,14 @@ def test_save_and_load_pretrained_with_tokenizer_object(): original_processor = TokenizerProcessor(tokenizer=mock_tokenizer, max_length=32, task_key="instruction") - robot_processor = RobotProcessor([original_processor]) + robot_processor = DataProcessorPipeline([original_processor]) with tempfile.TemporaryDirectory() as temp_dir: # Save processor robot_processor.save_pretrained(temp_dir) # Load processor with tokenizer override (since tokenizer object wasn't saved) - loaded_processor = RobotProcessor.from_pretrained( + loaded_processor = DataProcessorPipeline.from_pretrained( temp_dir, overrides={"tokenizer_processor": {"tokenizer": mock_tokenizer}} ) @@ -952,7 +952,7 @@ def test_integration_with_device_processor(mock_auto_tokenizer): # Create pipeline with TokenizerProcessor then DeviceProcessor tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6) device_processor = DeviceProcessor(device="cuda:0") - robot_processor = RobotProcessor([tokenizer_processor, device_processor]) + robot_processor = DataProcessorPipeline([tokenizer_processor, device_processor]) # Start with CPU tensors transition = create_transition( diff --git a/tests/processor/test_vqbet_processor.py b/tests/processor/test_vqbet_processor.py index 6df59c99d5..a0c56bdf4c 100644 --- a/tests/processor/test_vqbet_processor.py +++ b/tests/processor/test_vqbet_processor.py @@ -25,10 +25,10 @@ from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors from lerobot.processor import ( + DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, RenameProcessor, - RobotProcessor, ToBatchProcessor, UnnormalizerProcessor, ) @@ -245,7 +245,7 @@ def test_vqbet_processor_save_and_load(): preprocessor.save_pretrained(tmpdir) # Load preprocessor - loaded_preprocessor = RobotProcessor.from_pretrained(tmpdir) + loaded_preprocessor = DataProcessorPipeline.from_pretrained(tmpdir) # Test that loaded processor works observation = { From 0ed900714c81fd3b54f84eabcbe4db970cc3feda Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Fri, 29 Aug 2025 16:36:10 +0200 Subject: [PATCH 02/12] refactor(pipeline): introduce RobotProcessorPipeline and PolicyProcessorPipeline as type aliases for DataProcessorPipeline (#1804) * refactor(pipeline): introduce RobotProcessorPipeline and PolicyProcessorPipeline as type aliases for DataProcessorPipeline * refactor(pipeline): replace DataProcessorPipeline with PolicyProcessorPipeline alias across multiple policy files for consistency * refactor(pipeline): update references to DataProcessorPipeline with RobotProcessorPipeline in record, replay, and teleoperate modules for improved consistency --- src/lerobot/policies/act/processor_act.py | 10 +++++----- .../policies/diffusion/processor_diffusion.py | 10 +++++----- src/lerobot/policies/factory.py | 8 ++++---- src/lerobot/policies/pi0/processor_pi0.py | 10 +++++----- .../policies/pi0fast/processor_pi0fast.py | 10 +++++----- src/lerobot/policies/sac/processor_sac.py | 10 +++++----- .../sac/reward_model/processor_classifier.py | 10 +++++----- .../policies/smolvla/processor_smolvla.py | 10 +++++----- src/lerobot/policies/tdmpc/processor_tdmpc.py | 10 +++++----- src/lerobot/policies/vqbet/processor_vqbet.py | 10 +++++----- src/lerobot/processor/__init__.py | 4 ++++ src/lerobot/processor/pipeline.py | 7 ++++++- src/lerobot/record.py | 18 ++++++++--------- src/lerobot/replay.py | 6 +++--- src/lerobot/teleoperate.py | 20 +++++++++---------- src/lerobot/utils/train_utils.py | 6 +++--- 16 files changed, 84 insertions(+), 75 deletions(-) diff --git a/src/lerobot/policies/act/processor_act.py b/src/lerobot/policies/act/processor_act.py index 7b1f1693db..9e7bf3fb81 100644 --- a/src/lerobot/policies/act/processor_act.py +++ b/src/lerobot/policies/act/processor_act.py @@ -18,9 +18,9 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.act.configuration_act import ACTConfig from lerobot.processor import ( - DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, + PolicyProcessorPipeline, RenameProcessor, ToBatchProcessor, UnnormalizerProcessor, @@ -29,7 +29,7 @@ def make_act_pre_post_processors( config: ACTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), NormalizerProcessor( @@ -46,6 +46,6 @@ def make_act_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( - steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME - ) + return PolicyProcessorPipeline( + steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME + ), PolicyProcessorPipeline(steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME) diff --git a/src/lerobot/policies/diffusion/processor_diffusion.py b/src/lerobot/policies/diffusion/processor_diffusion.py index 08d36a1f83..03a27bd84c 100644 --- a/src/lerobot/policies/diffusion/processor_diffusion.py +++ b/src/lerobot/policies/diffusion/processor_diffusion.py @@ -19,9 +19,9 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.processor import ( - DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, + PolicyProcessorPipeline, RenameProcessor, ToBatchProcessor, UnnormalizerProcessor, @@ -30,7 +30,7 @@ def make_diffusion_pre_post_processors( config: DiffusionConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), NormalizerProcessor( @@ -47,6 +47,6 @@ def make_diffusion_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( - steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME - ) + return PolicyProcessorPipeline( + steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME + ), PolicyProcessorPipeline(steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME) diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index fe1d7bfbca..1d37d4e63b 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -38,7 +38,7 @@ from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig -from lerobot.processor.pipeline import DataProcessorPipeline +from lerobot.processor.pipeline import PolicyProcessorPipeline def get_policy_class(name: str) -> type[PreTrainedPolicy]: @@ -120,7 +120,7 @@ def make_pre_post_processors( policy_cfg: PreTrainedConfig, pretrained_path: str | None = None, **kwargs: Unpack[ProcessorConfigKwargs], -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: """Make a processor instance for a given policy type. This function creates the appropriate processor configuration based on the policy type. @@ -140,12 +140,12 @@ def make_pre_post_processors( """ if pretrained_path: return ( - DataProcessorPipeline.from_pretrained( + PolicyProcessorPipeline.from_pretrained( pretrained_model_name_or_path=pretrained_path, config_filename=kwargs.get("preprocessor_config_filename", "robot_preprocessor.json"), overrides=kwargs.get("preprocessor_overrides", {}), ), - DataProcessorPipeline.from_pretrained( + PolicyProcessorPipeline.from_pretrained( pretrained_model_name_or_path=pretrained_path, config_filename=kwargs.get("postprocessor_config_filename", "robot_postprocessor.json"), overrides=kwargs.get("postprocessor_overrides", {}), diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index def619add3..bdf95bdcc9 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -20,9 +20,9 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( - DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, + PolicyProcessorPipeline, ToBatchProcessor, TokenizerProcessor, UnnormalizerProcessor, @@ -66,7 +66,7 @@ def complementary_data(self, complementary_data): def make_pi0_pre_post_processors( config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: # Add remaining processors input_steps: list[ProcessorStep] = [ RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one @@ -93,6 +93,6 @@ def make_pi0_pre_post_processors( ), ] - return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( - steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME - ) + return PolicyProcessorPipeline( + steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME + ), PolicyProcessorPipeline(steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME) diff --git a/src/lerobot/policies/pi0fast/processor_pi0fast.py b/src/lerobot/policies/pi0fast/processor_pi0fast.py index 28642e4b2c..1dc48a6277 100644 --- a/src/lerobot/policies/pi0fast/processor_pi0fast.py +++ b/src/lerobot/policies/pi0fast/processor_pi0fast.py @@ -19,9 +19,9 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( - DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, + PolicyProcessorPipeline, RenameProcessor, ToBatchProcessor, UnnormalizerProcessor, @@ -30,7 +30,7 @@ def make_pi0fast_pre_post_processors( config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one NormalizerProcessor( @@ -47,6 +47,6 @@ def make_pi0fast_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( - steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME - ) + return PolicyProcessorPipeline( + steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME + ), PolicyProcessorPipeline(steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME) diff --git a/src/lerobot/policies/sac/processor_sac.py b/src/lerobot/policies/sac/processor_sac.py index badaffa394..16f667c80e 100644 --- a/src/lerobot/policies/sac/processor_sac.py +++ b/src/lerobot/policies/sac/processor_sac.py @@ -20,9 +20,9 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.processor import ( - DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, + PolicyProcessorPipeline, RenameProcessor, ToBatchProcessor, UnnormalizerProcessor, @@ -31,7 +31,7 @@ def make_sac_pre_post_processors( config: SACConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), NormalizerProcessor( @@ -48,6 +48,6 @@ def make_sac_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( - steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME - ) + return PolicyProcessorPipeline( + steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME + ), PolicyProcessorPipeline(steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME) diff --git a/src/lerobot/policies/sac/reward_model/processor_classifier.py b/src/lerobot/policies/sac/reward_model/processor_classifier.py index bf794ac02a..0fe955b2b4 100644 --- a/src/lerobot/policies/sac/reward_model/processor_classifier.py +++ b/src/lerobot/policies/sac/reward_model/processor_classifier.py @@ -17,16 +17,16 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.processor import ( - DataProcessorPipeline, DeviceProcessor, IdentityProcessor, NormalizerProcessor, + PolicyProcessorPipeline, ) def make_classifier_processor( config: RewardClassifierConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: input_steps = [ NormalizerProcessor( features=config.input_features, norm_map=config.normalization_mapping, stats=dataset_stats @@ -37,6 +37,6 @@ def make_classifier_processor( DeviceProcessor(device=config.device), ] output_steps = [DeviceProcessor(device="cpu"), IdentityProcessor()] - return DataProcessorPipeline(steps=input_steps, name="classifier_preprocessor"), DataProcessorPipeline( - steps=output_steps, name="classifier_postprocessor" - ) + return PolicyProcessorPipeline( + steps=input_steps, name="classifier_preprocessor" + ), PolicyProcessorPipeline(steps=output_steps, name="classifier_postprocessor") diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 2bf26a407c..dc7b71c495 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -19,9 +19,9 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.processor import ( - DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, + PolicyProcessorPipeline, RenameProcessor, ToBatchProcessor, TokenizerProcessor, @@ -35,7 +35,7 @@ def make_smolvla_pre_post_processors( config: SmolVLAConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), # To mimic the same processor as pretrained one NormalizerProcessor( @@ -59,9 +59,9 @@ def make_smolvla_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( - steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME - ) + return PolicyProcessorPipeline( + steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME + ), PolicyProcessorPipeline(steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME) @ProcessorStepRegistry.register(name="smolvla_new_line_processor") diff --git a/src/lerobot/policies/tdmpc/processor_tdmpc.py b/src/lerobot/policies/tdmpc/processor_tdmpc.py index 01d7e3cf45..b3d3acd2e6 100644 --- a/src/lerobot/policies/tdmpc/processor_tdmpc.py +++ b/src/lerobot/policies/tdmpc/processor_tdmpc.py @@ -19,9 +19,9 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.processor import ( - DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, + PolicyProcessorPipeline, RenameProcessor, ToBatchProcessor, UnnormalizerProcessor, @@ -30,7 +30,7 @@ def make_tdmpc_pre_post_processors( config: TDMPCConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), NormalizerProcessor( @@ -47,6 +47,6 @@ def make_tdmpc_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( - steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME - ) + return PolicyProcessorPipeline( + steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME + ), PolicyProcessorPipeline(steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME) diff --git a/src/lerobot/policies/vqbet/processor_vqbet.py b/src/lerobot/policies/vqbet/processor_vqbet.py index 23b600b61f..62f16216ef 100644 --- a/src/lerobot/policies/vqbet/processor_vqbet.py +++ b/src/lerobot/policies/vqbet/processor_vqbet.py @@ -20,9 +20,9 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.processor import ( - DataProcessorPipeline, DeviceProcessor, NormalizerProcessor, + PolicyProcessorPipeline, RenameProcessor, ToBatchProcessor, UnnormalizerProcessor, @@ -31,7 +31,7 @@ def make_vqbet_pre_post_processors( config: VQBeTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None -) -> tuple[DataProcessorPipeline, DataProcessorPipeline]: +) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: input_steps = [ RenameProcessor(rename_map={}), # Let the possibility to the user to rename the keys NormalizerProcessor( @@ -48,6 +48,6 @@ def make_vqbet_pre_post_processors( features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats ), ] - return DataProcessorPipeline(steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME), DataProcessorPipeline( - steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME - ) + return PolicyProcessorPipeline( + steps=input_steps, name=PREPROCESSOR_DEFAULT_NAME + ), PolicyProcessorPipeline(steps=output_steps, name=POSTPROCESSOR_DEFAULT_NAME) diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 4e25ea1e6a..2c13892252 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -38,9 +38,11 @@ IdentityProcessor, InfoProcessor, ObservationProcessor, + PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, RewardProcessor, + RobotProcessorPipeline, TransitionKey, TruncatedProcessor, ) @@ -74,6 +76,8 @@ "RewardClassifierProcessor", "RewardProcessor", "DataProcessorPipeline", + "RobotProcessorPipeline", + "PolicyProcessorPipeline", "ToBatchProcessor", "TokenizerProcessor", "TimeLimitProcessor", diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index 018b66b6b1..a43960fd90 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -24,7 +24,7 @@ from dataclasses import dataclass, field from enum import Enum from pathlib import Path -from typing import Any, TypedDict +from typing import Any, TypeAlias, TypedDict import torch from huggingface_hub import ModelHubMixin, hf_hub_download @@ -855,6 +855,11 @@ def transform_features(self, initial_features: dict[str, PolicyFeature]) -> dict return features +# TODO(Adil): when we upgrade to python 3.12, we can use TypeAlias instead of this +RobotProcessorPipeline: TypeAlias = DataProcessorPipeline +PolicyProcessorPipeline: TypeAlias = DataProcessorPipeline + + class ObservationProcessor(ProcessorStep, ABC): """Base class for processors that modify only the observation component of a transition. diff --git a/src/lerobot/record.py b/src/lerobot/record.py index 441d7e2c08..b196779a1b 100644 --- a/src/lerobot/record.py +++ b/src/lerobot/record.py @@ -76,7 +76,7 @@ from lerobot.datasets.video_utils import VideoEncodingManager from lerobot.policies.factory import make_policy, make_pre_post_processors from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor import DataProcessorPipeline +from lerobot.processor import PolicyProcessorPipeline, RobotProcessorPipeline from lerobot.processor.converters import ( to_dataset_frame, to_output_robot_action, @@ -236,22 +236,22 @@ def record_loop( dataset: LeRobotDataset | None = None, teleop: Teleoperator | list[Teleoperator] | None = None, policy: PreTrainedPolicy | None = None, - preprocessor: DataProcessorPipeline | None = None, - postprocessor: DataProcessorPipeline | None = None, + preprocessor: PolicyProcessorPipeline | None = None, + postprocessor: PolicyProcessorPipeline | None = None, control_time_s: int | None = None, - teleop_action_processor: DataProcessorPipeline | None = None, # runs after teleop - robot_action_processor: DataProcessorPipeline | None = None, # runs before robot - robot_observation_processor: DataProcessorPipeline | None = None, # runs after robot + teleop_action_processor: RobotProcessorPipeline | None = None, # runs after teleop + robot_action_processor: RobotProcessorPipeline | None = None, # runs before robot + robot_observation_processor: RobotProcessorPipeline | None = None, # runs after robot single_task: str | None = None, display_data: bool = False, ): - teleop_action_processor = teleop_action_processor or DataProcessorPipeline( + teleop_action_processor = teleop_action_processor or RobotProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr ) - robot_action_processor = robot_action_processor or DataProcessorPipeline( + robot_action_processor = robot_action_processor or RobotProcessorPipeline( steps=[IdentityProcessor()], to_transition=lambda tr: tr, to_output=to_output_robot_action ) - robot_observation_processor = robot_observation_processor or DataProcessorPipeline( + robot_observation_processor = robot_observation_processor or RobotProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr ) diff --git a/src/lerobot/replay.py b/src/lerobot/replay.py index c09554d744..2abcf9e68c 100644 --- a/src/lerobot/replay.py +++ b/src/lerobot/replay.py @@ -47,7 +47,7 @@ from lerobot.configs import parser from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.processor import DataProcessorPipeline +from lerobot.processor import RobotProcessorPipeline from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action from lerobot.processor.pipeline import IdentityProcessor from lerobot.robots import ( # noqa: F401 @@ -86,7 +86,7 @@ class ReplayConfig: # Use vocal synthesis to read events. play_sounds: bool = True # Optional processor for actions before sending to robot - robot_action_processor: DataProcessorPipeline | None = None + robot_action_processor: RobotProcessorPipeline | None = None @parser.wrap() @@ -95,7 +95,7 @@ def replay(cfg: ReplayConfig): logging.info(pformat(asdict(cfg))) # Initialize robot action processor with default if not provided - robot_action_processor = cfg.robot_action_processor or DataProcessorPipeline( + robot_action_processor = cfg.robot_action_processor or RobotProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=to_output_robot_action, # type: ignore[arg-type] diff --git a/src/lerobot/teleoperate.py b/src/lerobot/teleoperate.py index 33aafad04e..b6816c29c9 100644 --- a/src/lerobot/teleoperate.py +++ b/src/lerobot/teleoperate.py @@ -61,7 +61,7 @@ from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 from lerobot.configs import parser -from lerobot.processor import DataProcessorPipeline +from lerobot.processor import RobotProcessorPipeline from lerobot.processor.converters import ( to_output_robot_action, to_transition_robot_observation, @@ -105,9 +105,9 @@ class TeleoperateConfig: # Display all cameras on screen display_data: bool = False # Optional processors for data transformation - teleop_action_processor: DataProcessorPipeline | None = None # runs after teleop - robot_action_processor: DataProcessorPipeline | None = None # runs before robot - robot_observation_processor: DataProcessorPipeline | None = None # runs after robot + teleop_action_processor: RobotProcessorPipeline | None = None # runs after teleop + robot_action_processor: RobotProcessorPipeline | None = None # runs before robot + robot_observation_processor: RobotProcessorPipeline | None = None # runs after robot def teleop_loop( @@ -116,20 +116,20 @@ def teleop_loop( fps: int, display_data: bool = False, duration: float | None = None, - teleop_action_processor: DataProcessorPipeline | None = None, - robot_action_processor: DataProcessorPipeline | None = None, - robot_observation_processor: DataProcessorPipeline | None = None, + teleop_action_processor: RobotProcessorPipeline | None = None, + robot_action_processor: RobotProcessorPipeline | None = None, + robot_observation_processor: RobotProcessorPipeline | None = None, ): # Initialize processors with defaults if not provided - teleop_action_processor = teleop_action_processor or DataProcessorPipeline( + teleop_action_processor = teleop_action_processor or RobotProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr ) - robot_action_processor = robot_action_processor or DataProcessorPipeline( + robot_action_processor = robot_action_processor or RobotProcessorPipeline( steps=[IdentityProcessor()], to_transition=lambda tr: tr, to_output=to_output_robot_action, # type: ignore[arg-type] ) - robot_observation_processor = robot_observation_processor or DataProcessorPipeline( + robot_observation_processor = robot_observation_processor or RobotProcessorPipeline( steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr ) diff --git a/src/lerobot/utils/train_utils.py b/src/lerobot/utils/train_utils.py index 948674c3d3..1bec46b3d3 100644 --- a/src/lerobot/utils/train_utils.py +++ b/src/lerobot/utils/train_utils.py @@ -32,7 +32,7 @@ from lerobot.optim.optimizers import load_optimizer_state, save_optimizer_state from lerobot.optim.schedulers import load_scheduler_state, save_scheduler_state from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor.pipeline import DataProcessorPipeline +from lerobot.processor.pipeline import PolicyProcessorPipeline from lerobot.utils.random_utils import load_rng_state, save_rng_state @@ -75,8 +75,8 @@ def save_checkpoint( policy: PreTrainedPolicy, optimizer: Optimizer, scheduler: LRScheduler | None = None, - preprocessor: DataProcessorPipeline | None = None, - postprocessor: DataProcessorPipeline | None = None, + preprocessor: PolicyProcessorPipeline | None = None, + postprocessor: PolicyProcessorPipeline | None = None, ) -> None: """This function creates the following directory structure: From 690a80f9f0098d33ab5cc212d41fb01b7972a7e9 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Fri, 29 Aug 2025 16:53:58 +0200 Subject: [PATCH 03/12] refactor(processor): Add Step suffix on base processor (#1813) --- src/lerobot/policies/pi0/processor_pi0.py | 4 +-- .../sac/reward_model/processor_classifier.py | 4 +-- .../policies/smolvla/processor_smolvla.py | 4 +-- src/lerobot/processor/__init__.py | 28 +++++++++---------- src/lerobot/processor/batch_processor.py | 12 ++++---- .../processor/delta_action_processor.py | 6 ++-- src/lerobot/processor/gym_action_processor.py | 6 ++-- src/lerobot/processor/hil_processor.py | 18 ++++++------ .../processor/joint_observations_processor.py | 6 ++-- .../processor/observation_processor.py | 4 +-- src/lerobot/processor/pipeline.py | 28 +++++++++---------- src/lerobot/processor/rename_processor.py | 4 +-- src/lerobot/processor/tokenizer_processor.py | 4 +-- src/lerobot/record.py | 16 ++++++----- src/lerobot/replay.py | 4 +-- .../robot_kinematic_processor.py | 14 +++++----- src/lerobot/teleoperate.py | 10 ++++--- .../teleoperators/phone/phone_processor.py | 4 +-- src/lerobot/utils/control_utils.py | 6 ++-- tests/processor/test_classifier_processor.py | 13 ++++++--- tests/processor/test_normalize_processor.py | 8 +++--- tests/processor/test_observation_processor.py | 2 +- 22 files changed, 107 insertions(+), 98 deletions(-) diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index bdf95bdcc9..896a3f8b20 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -28,7 +28,7 @@ UnnormalizerProcessor, ) from lerobot.processor.pipeline import ( - ComplementaryDataProcessor, + ComplementaryDataProcessorStep, ProcessorStep, ProcessorStepRegistry, ) @@ -36,7 +36,7 @@ @ProcessorStepRegistry.register(name="pi0_new_line_processor") -class Pi0NewLineProcessor(ComplementaryDataProcessor): +class Pi0NewLineProcessor(ComplementaryDataProcessorStep): """Add a new line to the end of the task if it doesn't have one. This is required for the PaliGemma tokenizer. """ diff --git a/src/lerobot/policies/sac/reward_model/processor_classifier.py b/src/lerobot/policies/sac/reward_model/processor_classifier.py index 0fe955b2b4..59c383df78 100644 --- a/src/lerobot/policies/sac/reward_model/processor_classifier.py +++ b/src/lerobot/policies/sac/reward_model/processor_classifier.py @@ -18,7 +18,7 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.processor import ( DeviceProcessor, - IdentityProcessor, + IdentityProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, ) @@ -36,7 +36,7 @@ def make_classifier_processor( ), DeviceProcessor(device=config.device), ] - output_steps = [DeviceProcessor(device="cpu"), IdentityProcessor()] + output_steps = [DeviceProcessor(device="cpu"), IdentityProcessorStep()] return PolicyProcessorPipeline( steps=input_steps, name="classifier_preprocessor" ), PolicyProcessorPipeline(steps=output_steps, name="classifier_postprocessor") diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index dc7b71c495..32f5de0f59 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -28,7 +28,7 @@ UnnormalizerProcessor, ) from lerobot.processor.pipeline import ( - ComplementaryDataProcessor, + ComplementaryDataProcessorStep, ProcessorStepRegistry, ) @@ -65,7 +65,7 @@ def make_smolvla_pre_post_processors( @ProcessorStepRegistry.register(name="smolvla_new_line_processor") -class SmolVLANewLineProcessor(ComplementaryDataProcessor): +class SmolVLANewLineProcessor(ComplementaryDataProcessorStep): """Add a new line to the end of the task if it doesn't have one.""" def complementary_data(self, complementary_data): diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 2c13892252..d3bdcf5902 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -31,37 +31,37 @@ from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor, hotswap_stats from .observation_processor import VanillaObservationProcessor from .pipeline import ( - ActionProcessor, + ActionProcessorStep, DataProcessorPipeline, - DoneProcessor, + DoneProcessorStep, EnvTransition, - IdentityProcessor, - InfoProcessor, - ObservationProcessor, + IdentityProcessorStep, + InfoProcessorStep, + ObservationProcessorStep, PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, - RewardProcessor, + RewardProcessorStep, RobotProcessorPipeline, TransitionKey, - TruncatedProcessor, + TruncatedProcessorStep, ) from .rename_processor import RenameProcessor from .tokenizer_processor import TokenizerProcessor __all__ = [ - "ActionProcessor", + "ActionProcessorStep", "AddTeleopActionAsComplimentaryData", "AddTeleopEventsAsInfo", "DeviceProcessor", - "DoneProcessor", + "RewardProcessorStep", "MapDeltaActionToRobotAction", "MapTensorToDeltaActionDict", "EnvTransition", "GripperPenaltyProcessor", - "IdentityProcessor", + "IdentityProcessorStep", "ImageCropResizeProcessor", - "InfoProcessor", + "InfoProcessorStep", "InterventionActionProcessor", "JointVelocityProcessor", "MapDeltaActionToRobotAction", @@ -69,12 +69,12 @@ "NormalizerProcessor", "UnnormalizerProcessor", "hotswap_stats", - "ObservationProcessor", + "ObservationProcessorStep", "ProcessorStep", "ProcessorStepRegistry", "RenameProcessor", "RewardClassifierProcessor", - "RewardProcessor", + "DoneProcessorStep", "DataProcessorPipeline", "RobotProcessorPipeline", "PolicyProcessorPipeline", @@ -84,6 +84,6 @@ "Numpy2TorchActionProcessor", "Torch2NumpyActionProcessor", "TransitionKey", - "TruncatedProcessor", + "TruncatedProcessorStep", "VanillaObservationProcessor", ] diff --git a/src/lerobot/processor/batch_processor.py b/src/lerobot/processor/batch_processor.py index aab575ef71..3f42be8ee5 100644 --- a/src/lerobot/processor/batch_processor.py +++ b/src/lerobot/processor/batch_processor.py @@ -17,10 +17,10 @@ from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE from lerobot.processor.pipeline import ( - ActionProcessor, - ComplementaryDataProcessor, + ActionProcessorStep, + ComplementaryDataProcessorStep, EnvTransition, - ObservationProcessor, + ObservationProcessorStep, ProcessorStep, ProcessorStepRegistry, ) @@ -28,7 +28,7 @@ @dataclass @ProcessorStepRegistry.register(name="to_batch_processor_action") -class ToBatchProcessorAction(ActionProcessor): +class ToBatchProcessorAction(ActionProcessorStep): """Process action component in-place, adding batch dimension if needed.""" def action(self, action): @@ -40,7 +40,7 @@ def action(self, action): @dataclass @ProcessorStepRegistry.register(name="to_batch_processor_observation") -class ToBatchProcessorObservation(ObservationProcessor): +class ToBatchProcessorObservation(ObservationProcessorStep): """Process observation component in-place, adding batch dimensions where needed.""" def observation(self, observation): @@ -66,7 +66,7 @@ def observation(self, observation): @dataclass @ProcessorStepRegistry.register(name="to_batch_processor_complementary_data") -class ToBatchProcessorComplementaryData(ComplementaryDataProcessor): +class ToBatchProcessorComplementaryData(ComplementaryDataProcessorStep): """Process complementary data in-place, handling task field batching.""" def complementary_data(self, complementary_data): diff --git a/src/lerobot/processor/delta_action_processor.py b/src/lerobot/processor/delta_action_processor.py index 63eff9aad1..343cb3b388 100644 --- a/src/lerobot/processor/delta_action_processor.py +++ b/src/lerobot/processor/delta_action_processor.py @@ -19,12 +19,12 @@ from torch import Tensor from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry +from lerobot.processor.pipeline import ActionProcessorStep, ProcessorStepRegistry @ProcessorStepRegistry.register("map_tensor_to_delta_action_dict") @dataclass -class MapTensorToDeltaActionDict(ActionProcessor): +class MapTensorToDeltaActionDict(ActionProcessorStep): """ Map a tensor to a delta action dictionary. """ @@ -48,7 +48,7 @@ def action(self, action: Tensor) -> dict: @ProcessorStepRegistry.register("map_delta_action_to_robot_action") @dataclass -class MapDeltaActionToRobotAction(ActionProcessor): +class MapDeltaActionToRobotAction(ActionProcessorStep): """ Map delta actions from teleoperators (gamepad, keyboard) to robot target actions for use with inverse kinematics processors. diff --git a/src/lerobot/processor/gym_action_processor.py b/src/lerobot/processor/gym_action_processor.py index 70c54cbded..4ab6813ac4 100644 --- a/src/lerobot/processor/gym_action_processor.py +++ b/src/lerobot/processor/gym_action_processor.py @@ -16,12 +16,12 @@ import numpy as np import torch -from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry +from lerobot.processor.pipeline import ActionProcessorStep, ProcessorStepRegistry @ProcessorStepRegistry.register("torch2numpy_action_processor") @dataclass -class Torch2NumpyActionProcessor(ActionProcessor): +class Torch2NumpyActionProcessor(ActionProcessorStep): """Convert PyTorch tensor actions to NumPy arrays.""" squeeze_batch_dim: bool = True @@ -50,7 +50,7 @@ def action(self, action: torch.Tensor) -> np.ndarray: @ProcessorStepRegistry.register("numpy2torch_action_processor") @dataclass -class Numpy2TorchActionProcessor(ActionProcessor): +class Numpy2TorchActionProcessor(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 c1f4569eda..0bd9de8bce 100644 --- a/src/lerobot/processor/hil_processor.py +++ b/src/lerobot/processor/hil_processor.py @@ -9,14 +9,14 @@ from lerobot.configs.types import PolicyFeature from lerobot.processor.pipeline import ( - ComplementaryDataProcessor, + ComplementaryDataProcessorStep, EnvTransition, - InfoProcessor, - ObservationProcessor, + InfoProcessorStep, + ObservationProcessorStep, ProcessorStep, ProcessorStepRegistry, TransitionKey, - TruncatedProcessor, + TruncatedProcessorStep, ) from lerobot.teleoperators.teleoperator import Teleoperator from lerobot.teleoperators.utils import TeleopEvents @@ -26,7 +26,7 @@ @ProcessorStepRegistry.register("add_teleop_action_as_complementary_data") @dataclass -class AddTeleopActionAsComplimentaryData(ComplementaryDataProcessor): +class AddTeleopActionAsComplimentaryData(ComplementaryDataProcessorStep): """Add teleoperator action to transition complementary data.""" teleop_device: Teleoperator @@ -39,7 +39,7 @@ def complementary_data(self, complementary_data: dict) -> dict: @ProcessorStepRegistry.register("add_teleop_action_as_info") @dataclass -class AddTeleopEventsAsInfo(InfoProcessor): +class AddTeleopEventsAsInfo(InfoProcessorStep): """Add teleoperator control events to transition info.""" teleop_device: Teleoperator @@ -53,7 +53,7 @@ def info(self, info: dict) -> dict: @ProcessorStepRegistry.register("image_crop_resize_processor") @dataclass -class ImageCropResizeProcessor(ObservationProcessor): +class ImageCropResizeProcessor(ObservationProcessorStep): """Crop and resize image observations.""" crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None @@ -103,7 +103,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po @dataclass @ProcessorStepRegistry.register("time_limit_processor") -class TimeLimitProcessor(TruncatedProcessor): +class TimeLimitProcessor(TruncatedProcessorStep): """Track episode steps and enforce time limits.""" max_episode_steps: int @@ -127,7 +127,7 @@ def reset(self) -> None: @dataclass @ProcessorStepRegistry.register("gripper_penalty_processor") -class GripperPenaltyProcessor(ComplementaryDataProcessor): +class GripperPenaltyProcessor(ComplementaryDataProcessorStep): """Apply penalty for inappropriate gripper usage.""" penalty: float = -0.01 diff --git a/src/lerobot/processor/joint_observations_processor.py b/src/lerobot/processor/joint_observations_processor.py index 40ec4ea1ec..aae0921cc7 100644 --- a/src/lerobot/processor/joint_observations_processor.py +++ b/src/lerobot/processor/joint_observations_processor.py @@ -5,7 +5,7 @@ from lerobot.configs.types import PolicyFeature from lerobot.processor.pipeline import ( - ObservationProcessor, + ObservationProcessorStep, ProcessorStepRegistry, ) from lerobot.robots import Robot @@ -13,7 +13,7 @@ @dataclass @ProcessorStepRegistry.register("joint_velocity_processor") -class JointVelocityProcessor(ObservationProcessor): +class JointVelocityProcessor(ObservationProcessorStep): """Add joint velocity information to observations.""" dt: float = 0.1 @@ -66,7 +66,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po @dataclass @ProcessorStepRegistry.register("current_processor") -class MotorCurrentProcessor(ObservationProcessor): +class MotorCurrentProcessor(ObservationProcessorStep): """Add motor current information to observations.""" robot: Robot | None = None diff --git a/src/lerobot/processor/observation_processor.py b/src/lerobot/processor/observation_processor.py index 40273548ef..1648a97981 100644 --- a/src/lerobot/processor/observation_processor.py +++ b/src/lerobot/processor/observation_processor.py @@ -22,12 +22,12 @@ from lerobot.configs.types import PolicyFeature from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE -from lerobot.processor.pipeline import ObservationProcessor, ProcessorStepRegistry +from lerobot.processor.pipeline import ObservationProcessorStep, ProcessorStepRegistry @dataclass @ProcessorStepRegistry.register(name="observation_processor") -class VanillaObservationProcessor(ObservationProcessor): +class VanillaObservationProcessor(ObservationProcessorStep): """ Processes environment observations into the LeRobot format by handling both images and states. diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index a43960fd90..bf8fd0ef0a 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -860,7 +860,7 @@ def transform_features(self, initial_features: dict[str, PolicyFeature]) -> dict PolicyProcessorPipeline: TypeAlias = DataProcessorPipeline -class ObservationProcessor(ProcessorStep, ABC): +class ObservationProcessorStep(ProcessorStep, ABC): """Base class for processors that modify only the observation component of a transition. Subclasses should override the `observation` method to implement custom observation processing. @@ -869,7 +869,7 @@ class ObservationProcessor(ProcessorStep, ABC): Example: ```python - class MyObservationScaler(ObservationProcessor): + class MyObservationScaler(ObservationProcessorStep): def __init__(self, scale_factor): self.scale_factor = scale_factor @@ -906,7 +906,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: return new_transition -class ActionProcessor(ProcessorStep, ABC): +class ActionProcessorStep(ProcessorStep, ABC): """Base class for processors that modify only the action component of a transition. Subclasses should override the `action` method to implement custom action processing. @@ -915,7 +915,7 @@ class ActionProcessor(ProcessorStep, ABC): Example: ```python - class ActionClipping(ActionProcessor): + class ActionClipping(ActionProcessorStep): def __init__(self, min_val, max_val): self.min_val = min_val self.max_val = max_val @@ -953,7 +953,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: return new_transition -class RewardProcessor(ProcessorStep, ABC): +class RewardProcessorStep(ProcessorStep, ABC): """Base class for processors that modify only the reward component of a transition. Subclasses should override the `reward` method to implement custom reward processing. @@ -962,7 +962,7 @@ class RewardProcessor(ProcessorStep, ABC): Example: ```python - class RewardScaler(RewardProcessor): + class RewardScaler(RewardProcessorStep): def __init__(self, scale_factor): self.scale_factor = scale_factor @@ -999,7 +999,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: return new_transition -class DoneProcessor(ProcessorStep, ABC): +class DoneProcessorStep(ProcessorStep, ABC): """Base class for processors that modify only the done flag of a transition. Subclasses should override the `done` method to implement custom done flag processing. @@ -1008,7 +1008,7 @@ class DoneProcessor(ProcessorStep, ABC): Example: ```python - class TimeoutDone(DoneProcessor): + class TimeoutDone(DoneProcessorStep): def __init__(self, max_steps): self.steps = 0 self.max_steps = max_steps @@ -1050,7 +1050,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: return new_transition -class TruncatedProcessor(ProcessorStep, ABC): +class TruncatedProcessorStep(ProcessorStep, ABC): """Base class for processors that modify only the truncated flag of a transition. Subclasses should override the `truncated` method to implement custom truncated flag processing. @@ -1059,7 +1059,7 @@ class TruncatedProcessor(ProcessorStep, ABC): Example: ```python - class EarlyTruncation(TruncatedProcessor): + class EarlyTruncation(TruncatedProcessorStep): def __init__(self, threshold): self.threshold = threshold @@ -1097,7 +1097,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: return new_transition -class InfoProcessor(ProcessorStep, ABC): +class InfoProcessorStep(ProcessorStep, ABC): """Base class for processors that modify only the info dictionary of a transition. Subclasses should override the `info` method to implement custom info processing. @@ -1106,7 +1106,7 @@ class InfoProcessor(ProcessorStep, ABC): Example: ```python - class InfoAugmenter(InfoProcessor): + class InfoAugmenter(InfoProcessorStep): def __init__(self): self.step_count = 0 @@ -1149,7 +1149,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: return new_transition -class ComplementaryDataProcessor(ProcessorStep, ABC): +class ComplementaryDataProcessorStep(ProcessorStep, ABC): """Base class for processors that modify only the complementary data of a transition. Subclasses should override the `complementary_data` method to implement custom complementary data processing. @@ -1182,7 +1182,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: return new_transition -class IdentityProcessor(ProcessorStep): +class IdentityProcessorStep(ProcessorStep): """Identity processor that does nothing.""" def __call__(self, transition: EnvTransition) -> EnvTransition: diff --git a/src/lerobot/processor/rename_processor.py b/src/lerobot/processor/rename_processor.py index ebc867cacc..63dff1a3de 100644 --- a/src/lerobot/processor/rename_processor.py +++ b/src/lerobot/processor/rename_processor.py @@ -19,14 +19,14 @@ from lerobot.configs.types import PolicyFeature from lerobot.processor.pipeline import ( - ObservationProcessor, + ObservationProcessorStep, ProcessorStepRegistry, ) @dataclass @ProcessorStepRegistry.register(name="rename_processor") -class RenameProcessor(ObservationProcessor): +class RenameProcessor(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 d2c04e44cd..1d444e3363 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -13,7 +13,7 @@ from lerobot.constants import OBS_LANGUAGE from lerobot.processor.pipeline import ( EnvTransition, - ObservationProcessor, + ObservationProcessorStep, ProcessorStepRegistry, TransitionKey, ) @@ -27,7 +27,7 @@ @dataclass @ProcessorStepRegistry.register(name="tokenizer_processor") -class TokenizerProcessor(ObservationProcessor): +class TokenizerProcessor(ObservationProcessorStep): """Tokenizes text tasks in complementary data using a huggingface tokenizer. This processor handles tokenization of task strings found in the complementary_data diff --git a/src/lerobot/record.py b/src/lerobot/record.py index b196779a1b..e467c93f8e 100644 --- a/src/lerobot/record.py +++ b/src/lerobot/record.py @@ -83,7 +83,7 @@ to_transition_robot_observation, to_transition_teleop_action, ) -from lerobot.processor.pipeline import IdentityProcessor, TransitionKey +from lerobot.processor.pipeline import IdentityProcessorStep, TransitionKey from lerobot.processor.rename_processor import rename_stats from lerobot.robots import ( # noqa: F401 Robot, @@ -246,13 +246,15 @@ def record_loop( display_data: bool = False, ): teleop_action_processor = teleop_action_processor or RobotProcessorPipeline( - steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr + steps=[IdentityProcessorStep()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr ) robot_action_processor = robot_action_processor or RobotProcessorPipeline( - steps=[IdentityProcessor()], to_transition=lambda tr: tr, to_output=to_output_robot_action + steps=[IdentityProcessorStep()], to_transition=lambda tr: tr, to_output=to_output_robot_action ) robot_observation_processor = robot_observation_processor or RobotProcessorPipeline( - steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr + steps=[IdentityProcessorStep()], + to_transition=to_transition_robot_observation, + to_output=lambda tr: tr, ) if dataset is not None and dataset.fps != fps: @@ -302,7 +304,7 @@ def record_loop( # Get robot observation obs = robot.get_observation() - # Applies a pipeline to the raw robot observation, default is IdentityProcessor + # Applies a pipeline to the raw robot observation, default is IdentityProcessorStep obs_transition = robot_observation_processor(obs) # Get action from either policy or teleop @@ -333,7 +335,7 @@ def record_loop( elif isinstance(teleop, Teleoperator): act = teleop.get_action() - # Applies a pipeline to the raw teleop action, default is IdentityProcessor + # Applies a pipeline to the raw teleop action, default is IdentityProcessorStep teleop_transition = teleop_action_processor(act) elif isinstance(teleop, list): @@ -351,7 +353,7 @@ def record_loop( ) continue - # Applies a pipeline to the action, default is IdentityProcessor + # Applies a pipeline to the action, default is IdentityProcessorStep # IMPORTANT: action_pipeline.to_output must return a dict suitable for robot.send_action() if policy is not None and policy_transition is not None: robot_action_to_send = robot_action_processor(policy_transition) diff --git a/src/lerobot/replay.py b/src/lerobot/replay.py index 2abcf9e68c..cc21be8f5a 100644 --- a/src/lerobot/replay.py +++ b/src/lerobot/replay.py @@ -49,7 +49,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.processor import RobotProcessorPipeline from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action -from lerobot.processor.pipeline import IdentityProcessor +from lerobot.processor.pipeline import IdentityProcessorStep from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, @@ -96,7 +96,7 @@ def replay(cfg: ReplayConfig): # Initialize robot action processor with default if not provided robot_action_processor = cfg.robot_action_processor or RobotProcessorPipeline( - steps=[IdentityProcessor()], + steps=[IdentityProcessorStep()], to_transition=to_transition_teleop_action, to_output=to_output_robot_action, # type: ignore[arg-type] ) diff --git a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py index 7c6c73a4d9..2b086d2951 100644 --- a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py +++ b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py @@ -22,10 +22,10 @@ from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.model.kinematics import RobotKinematics from lerobot.processor.pipeline import ( - ActionProcessor, - ComplementaryDataProcessor, + ActionProcessorStep, + ComplementaryDataProcessorStep, EnvTransition, - ObservationProcessor, + ObservationProcessorStep, ProcessorStep, ProcessorStepRegistry, TransitionKey, @@ -35,7 +35,7 @@ @ProcessorStepRegistry.register("ee_reference_and_delta") @dataclass -class EEReferenceAndDelta(ActionProcessor): +class EEReferenceAndDelta(ActionProcessorStep): """ Compute the desired end-effector pose from the target pose and the current pose. @@ -158,7 +158,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po @ProcessorStepRegistry.register("ee_bounds_and_safety") @dataclass -class EEBoundsAndSafety(ActionProcessor): +class EEBoundsAndSafety(ActionProcessorStep): """ Clip the end-effector pose to the bounds and check for jumps. @@ -377,7 +377,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po @ProcessorStepRegistry.register("forward_kinematics_joints_to_ee") @dataclass -class ForwardKinematicsJointsToEE(ObservationProcessor): +class ForwardKinematicsJointsToEE(ObservationProcessorStep): """ Compute the end-effector pose from the joint positions. @@ -421,7 +421,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po @ProcessorStepRegistry.register("add_robot_observation") @dataclass -class AddRobotObservationAsComplimentaryData(ComplementaryDataProcessor): +class AddRobotObservationAsComplimentaryData(ComplementaryDataProcessorStep): """ Read the robot's current observation and insert it into the transition as complementary data. diff --git a/src/lerobot/teleoperate.py b/src/lerobot/teleoperate.py index b6816c29c9..404e8a5bf1 100644 --- a/src/lerobot/teleoperate.py +++ b/src/lerobot/teleoperate.py @@ -67,7 +67,7 @@ to_transition_robot_observation, to_transition_teleop_action, ) -from lerobot.processor.pipeline import IdentityProcessor +from lerobot.processor.pipeline import IdentityProcessorStep from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, @@ -122,15 +122,17 @@ def teleop_loop( ): # Initialize processors with defaults if not provided teleop_action_processor = teleop_action_processor or RobotProcessorPipeline( - steps=[IdentityProcessor()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr + steps=[IdentityProcessorStep()], to_transition=to_transition_teleop_action, to_output=lambda tr: tr ) robot_action_processor = robot_action_processor or RobotProcessorPipeline( - steps=[IdentityProcessor()], + steps=[IdentityProcessorStep()], to_transition=lambda tr: tr, to_output=to_output_robot_action, # type: ignore[arg-type] ) robot_observation_processor = robot_observation_processor or RobotProcessorPipeline( - steps=[IdentityProcessor()], to_transition=to_transition_robot_observation, to_output=lambda tr: tr + steps=[IdentityProcessorStep()], + to_transition=to_transition_robot_observation, + to_output=lambda tr: tr, ) # Reset processors diff --git a/src/lerobot/teleoperators/phone/phone_processor.py b/src/lerobot/teleoperators/phone/phone_processor.py index 36880e0c85..63f11fb98d 100644 --- a/src/lerobot/teleoperators/phone/phone_processor.py +++ b/src/lerobot/teleoperators/phone/phone_processor.py @@ -17,13 +17,13 @@ from dataclasses import dataclass, field from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry +from lerobot.processor.pipeline import ActionProcessorStep, ProcessorStepRegistry from lerobot.teleoperators.phone.config_phone import PhoneOS @ProcessorStepRegistry.register("map_phone_action_to_robot_action") @dataclass -class MapPhoneActionToRobotAction(ActionProcessor): +class MapPhoneActionToRobotAction(ActionProcessorStep): """ Map calibrated phone pose (actions) to the inputs for robot actions diff --git a/src/lerobot/utils/control_utils.py b/src/lerobot/utils/control_utils.py index 5553d56f23..fb883d2c52 100644 --- a/src/lerobot/utils/control_utils.py +++ b/src/lerobot/utils/control_utils.py @@ -31,7 +31,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.utils import DEFAULT_FEATURES from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor import DataProcessorPipeline, TransitionKey +from lerobot.processor import PolicyProcessorPipeline, TransitionKey from lerobot.robots import Robot @@ -102,8 +102,8 @@ def predict_action( observation: dict[str, np.ndarray], policy: PreTrainedPolicy, device: torch.device, - preprocessor: DataProcessorPipeline, - postprocessor: DataProcessorPipeline, + preprocessor: PolicyProcessorPipeline, + postprocessor: PolicyProcessorPipeline, use_amp: bool, task: str | None = None, robot_type: str | None = None, diff --git a/tests/processor/test_classifier_processor.py b/tests/processor/test_classifier_processor.py index 07abccfbb5..8c060cb0b0 100644 --- a/tests/processor/test_classifier_processor.py +++ b/tests/processor/test_classifier_processor.py @@ -24,7 +24,12 @@ from lerobot.constants import OBS_IMAGE, OBS_STATE from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor -from lerobot.processor import DataProcessorPipeline, DeviceProcessor, IdentityProcessor, NormalizerProcessor +from lerobot.processor import ( + DataProcessorPipeline, + DeviceProcessor, + IdentityProcessorStep, + NormalizerProcessor, +) from lerobot.processor.pipeline import TransitionKey @@ -89,7 +94,7 @@ def test_make_classifier_processor_basic(): # Check steps in postprocessor assert len(postprocessor.steps) == 2 assert isinstance(postprocessor.steps[0], DeviceProcessor) - assert isinstance(postprocessor.steps[1], IdentityProcessor) + assert isinstance(postprocessor.steps[1], IdentityProcessorStep) def test_classifier_processor_normalization(): @@ -311,7 +316,7 @@ def test_classifier_processor_batch_data(): def test_classifier_processor_postprocessor_identity(): - """Test that Classifier postprocessor uses IdentityProcessor correctly.""" + """Test that Classifier postprocessor uses IdentityProcessorStep correctly.""" config = create_default_config() stats = create_default_stats() @@ -324,6 +329,6 @@ def test_classifier_processor_postprocessor_identity(): # Process through postprocessor processed = postprocessor(transition) - # IdentityProcessor should leave values unchanged (except device) + # IdentityProcessorStep should leave values unchanged (except device) assert torch.allclose(processed[TransitionKey.ACTION].cpu(), reward.cpu()) assert processed[TransitionKey.ACTION].device.type == "cpu" diff --git a/tests/processor/test_normalize_processor.py b/tests/processor/test_normalize_processor.py index 261153a837..32240ee461 100644 --- a/tests/processor/test_normalize_processor.py +++ b/tests/processor/test_normalize_processor.py @@ -26,7 +26,7 @@ _convert_stats_to_tensors, hotswap_stats, ) -from lerobot.processor.pipeline import DataProcessorPipeline, IdentityProcessor, TransitionKey +from lerobot.processor.pipeline import DataProcessorPipeline, IdentityProcessorStep, TransitionKey def create_transition( @@ -1004,7 +1004,7 @@ 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) - identity = IdentityProcessor() + identity = IdentityProcessorStep() # Create robot processor robot_processor = DataProcessorPipeline(steps=[normalizer, unnormalizer, identity]) @@ -1085,7 +1085,7 @@ def test_hotswap_stats_only_affects_normalizer_steps(): # Create mixed steps normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) unnormalizer = UnnormalizerProcessor(features=features, norm_map=norm_map, stats=stats) - identity = IdentityProcessor() + identity = IdentityProcessorStep() robot_processor = DataProcessorPipeline(steps=[normalizer, identity, unnormalizer]) @@ -1131,7 +1131,7 @@ def test_hotswap_stats_no_normalizer_steps(): } # Create processor with only identity steps - robot_processor = DataProcessorPipeline(steps=[IdentityProcessor(), IdentityProcessor()]) + robot_processor = DataProcessorPipeline(steps=[IdentityProcessorStep(), IdentityProcessorStep()]) # Hotswap stats - should work without error new_processor = hotswap_stats(robot_processor, stats) diff --git a/tests/processor/test_observation_processor.py b/tests/processor/test_observation_processor.py index 4e6efdb6c5..3196416e32 100644 --- a/tests/processor/test_observation_processor.py +++ b/tests/processor/test_observation_processor.py @@ -356,7 +356,7 @@ def test_empty_observation(): def test_equivalent_to_original_function(): - """Test that ObservationProcessor produces equivalent results to preprocess_observation.""" + """Test that ObservationProcessorStep produces equivalent results to preprocess_observation.""" # Import the original function for comparison from lerobot.envs.utils import preprocess_observation From 8382208e13d344e8f68e0ee0ad68b646da2b5ad7 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Fri, 29 Aug 2025 17:15:52 +0200 Subject: [PATCH 04/12] Chore(processor): Rename DeviceProcessor to DeviceProcessorStep and rename _device field (#1814) * refactor(processor): rename DeviceProcessor to DeviceProcessorStep across multiple files for consistency * refactor(processor): update device reference in DeviceProcessorStep for clarity --- src/lerobot/policies/act/processor_act.py | 6 +- .../policies/diffusion/processor_diffusion.py | 6 +- src/lerobot/policies/pi0/processor_pi0.py | 6 +- .../policies/pi0fast/processor_pi0fast.py | 6 +- src/lerobot/policies/sac/processor_sac.py | 6 +- .../sac/reward_model/processor_classifier.py | 6 +- .../policies/smolvla/processor_smolvla.py | 6 +- src/lerobot/policies/tdmpc/processor_tdmpc.py | 6 +- src/lerobot/policies/vqbet/processor_vqbet.py | 6 +- src/lerobot/processor/__init__.py | 4 +- src/lerobot/processor/device_processor.py | 10 +- tests/processor/test_act_processor.py | 12 +- tests/processor/test_classifier_processor.py | 12 +- tests/processor/test_device_processor.py | 116 +++++++++--------- tests/processor/test_diffusion_processor.py | 12 +- tests/processor/test_pi0_processor.py | 6 +- tests/processor/test_sac_processor.py | 12 +- tests/processor/test_smolvla_processor.py | 6 +- tests/processor/test_tdmpc_processor.py | 12 +- tests/processor/test_tokenizer_processor.py | 10 +- tests/processor/test_vqbet_processor.py | 12 +- 21 files changed, 139 insertions(+), 139 deletions(-) diff --git a/src/lerobot/policies/act/processor_act.py b/src/lerobot/policies/act/processor_act.py index 9e7bf3fb81..e9f3fa264e 100644 --- a/src/lerobot/policies/act/processor_act.py +++ b/src/lerobot/policies/act/processor_act.py @@ -18,7 +18,7 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.act.configuration_act import ACTConfig from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, @@ -38,10 +38,10 @@ def make_act_pre_post_processors( stats=dataset_stats, ), ToBatchProcessor(), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor( 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 03a27bd84c..15f0b5cc42 100644 --- a/src/lerobot/policies/diffusion/processor_diffusion.py +++ b/src/lerobot/policies/diffusion/processor_diffusion.py @@ -19,7 +19,7 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, @@ -39,10 +39,10 @@ def make_diffusion_pre_post_processors( stats=dataset_stats, ), ToBatchProcessor(), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor( 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 896a3f8b20..b538a4abc8 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -20,7 +20,7 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, ToBatchProcessor, @@ -83,11 +83,11 @@ def make_pi0_pre_post_processors( padding_side="right", padding="max_length", ), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps: list[ProcessorStep] = [ - DeviceProcessor(device="cpu"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor( 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 1dc48a6277..11fc05c838 100644 --- a/src/lerobot/policies/pi0fast/processor_pi0fast.py +++ b/src/lerobot/policies/pi0fast/processor_pi0fast.py @@ -19,7 +19,7 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, @@ -39,10 +39,10 @@ def make_pi0fast_pre_post_processors( stats=dataset_stats, ), ToBatchProcessor(), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor( 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 16f667c80e..2cbdb7fd3c 100644 --- a/src/lerobot/policies/sac/processor_sac.py +++ b/src/lerobot/policies/sac/processor_sac.py @@ -20,7 +20,7 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, @@ -40,10 +40,10 @@ def make_sac_pre_post_processors( stats=dataset_stats, ), ToBatchProcessor(), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor( 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 59c383df78..857ae450de 100644 --- a/src/lerobot/policies/sac/reward_model/processor_classifier.py +++ b/src/lerobot/policies/sac/reward_model/processor_classifier.py @@ -17,7 +17,7 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, IdentityProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, @@ -34,9 +34,9 @@ def make_classifier_processor( NormalizerProcessor( 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 PolicyProcessorPipeline( steps=input_steps, name="classifier_preprocessor" ), PolicyProcessorPipeline(steps=output_steps, name="classifier_postprocessor") diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 32f5de0f59..16a6fabcfe 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -19,7 +19,7 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, @@ -51,10 +51,10 @@ def make_smolvla_pre_post_processors( padding_side="right", max_length=config.tokenizer_max_length, ), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor( 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 b3d3acd2e6..faf73e0039 100644 --- a/src/lerobot/policies/tdmpc/processor_tdmpc.py +++ b/src/lerobot/policies/tdmpc/processor_tdmpc.py @@ -19,7 +19,7 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, @@ -39,10 +39,10 @@ def make_tdmpc_pre_post_processors( stats=dataset_stats, ), ToBatchProcessor(), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor( 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 62f16216ef..ef65433322 100644 --- a/src/lerobot/policies/vqbet/processor_vqbet.py +++ b/src/lerobot/policies/vqbet/processor_vqbet.py @@ -20,7 +20,7 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, @@ -40,10 +40,10 @@ def make_vqbet_pre_post_processors( stats=dataset_stats, ), ToBatchProcessor(), - DeviceProcessor(device=config.device), + DeviceProcessorStep(device=config.device), ] output_steps = [ - DeviceProcessor(device="cpu"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor( 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 d3bdcf5902..cc390c588d 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -16,7 +16,7 @@ from .batch_processor import ToBatchProcessor from .delta_action_processor import MapDeltaActionToRobotAction, MapTensorToDeltaActionDict -from .device_processor import DeviceProcessor +from .device_processor import DeviceProcessorStep from .gym_action_processor import Numpy2TorchActionProcessor, Torch2NumpyActionProcessor from .hil_processor import ( AddTeleopActionAsComplimentaryData, @@ -53,7 +53,7 @@ "ActionProcessorStep", "AddTeleopActionAsComplimentaryData", "AddTeleopEventsAsInfo", - "DeviceProcessor", + "DeviceProcessorStep", "RewardProcessorStep", "MapDeltaActionToRobotAction", "MapTensorToDeltaActionDict", diff --git a/src/lerobot/processor/device_processor.py b/src/lerobot/processor/device_processor.py index 78a3ad7972..8352b18c06 100644 --- a/src/lerobot/processor/device_processor.py +++ b/src/lerobot/processor/device_processor.py @@ -24,7 +24,7 @@ @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 @@ -47,8 +47,8 @@ class DeviceProcessor(ProcessorStep): } def __post_init__(self): - self._device: torch.device = get_safe_torch_device(self.device) - self.device = self._device.type # cuda might have changed to cuda:1 + self.torch_device: torch.device = get_safe_torch_device(self.device) + self.device = self.torch_device.type # cuda might have changed to cuda:1 self.non_blocking = "cuda" in str(self.device) # Validate and convert float_dtype string to torch dtype @@ -70,7 +70,7 @@ def _process_tensor(self, tensor: torch.Tensor) -> torch.Tensor: Otherwise, it moves to the configured device. """ # Determine target device - if tensor.is_cuda and self._device.type == "cuda": + if tensor.is_cuda and self.torch_device.type == "cuda": # Both tensor and target are on GPU - preserve tensor's GPU placement # This handles multi-GPU scenarios where Accelerate has already placed # tensors on the correct GPU for each process @@ -78,7 +78,7 @@ def _process_tensor(self, tensor: torch.Tensor) -> torch.Tensor: else: # Either tensor is on CPU, or we're configured for CPU # In both cases, use the configured device - target_device = self._device + target_device = self.torch_device # Only move if necessary if tensor.device != target_device: diff --git a/tests/processor/test_act_processor.py b/tests/processor/test_act_processor.py index 6ee926da2c..50df9597ad 100644 --- a/tests/processor/test_act_processor.py +++ b/tests/processor/test_act_processor.py @@ -26,7 +26,7 @@ from lerobot.policies.act.processor_act import make_act_pre_post_processors from lerobot.processor import ( DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, RenameProcessor, ToBatchProcessor, @@ -89,11 +89,11 @@ def test_make_act_processor_basic(): 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[3], 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], UnnormalizerProcessor) @@ -271,10 +271,10 @@ def test_act_processor_mixed_precision(): # Modify the device processor to use float16 preprocessor, postprocessor = make_act_pre_post_processors(config, stats) - # 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(7, dtype=torch.float32)} diff --git a/tests/processor/test_classifier_processor.py b/tests/processor/test_classifier_processor.py index 8c060cb0b0..546cd54041 100644 --- a/tests/processor/test_classifier_processor.py +++ b/tests/processor/test_classifier_processor.py @@ -26,7 +26,7 @@ from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor from lerobot.processor import ( DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, IdentityProcessorStep, NormalizerProcessor, ) @@ -89,11 +89,11 @@ def test_make_classifier_processor_basic(): 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[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) @@ -268,10 +268,10 @@ def test_classifier_processor_mixed_precision(): # Create processor preprocessor, postprocessor = make_classifier_processor(config, stats) - # 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_device_processor.py b/tests/processor/test_device_processor.py index 3f86d06fb8..5c693857ed 100644 --- a/tests/processor/test_device_processor.py +++ b/tests/processor/test_device_processor.py @@ -19,7 +19,7 @@ import torch from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.processor import DataProcessorPipeline, DeviceProcessor +from lerobot.processor import DataProcessorPipeline, DeviceProcessorStep from lerobot.processor.pipeline import TransitionKey @@ -47,7 +47,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)} @@ -74,7 +74,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)} @@ -101,7 +101,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) @@ -117,7 +117,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), @@ -143,7 +143,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)) @@ -160,7 +160,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) @@ -171,7 +171,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) @@ -188,7 +188,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), @@ -210,7 +210,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), @@ -233,7 +233,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 = { @@ -254,22 +254,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() @@ -290,7 +290,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,)), @@ -307,8 +307,8 @@ def test_integration_with_robot_processor(): from lerobot.constants import OBS_STATE from lerobot.processor import ToBatchProcessor - # Create a pipeline with DeviceProcessor - device_processor = DeviceProcessor(device="cpu") + # Create a pipeline with DeviceProcessorStep + device_processor = DeviceProcessorStep(device="cpu") batch_processor = ToBatchProcessor() processor = DataProcessorPipeline(steps=[batch_processor, device_processor], name="test_pipeline") @@ -329,9 +329,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: @@ -343,7 +343,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] @@ -352,18 +352,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.pipeline 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 = { @@ -385,7 +385,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( @@ -425,7 +425,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", @@ -445,13 +445,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 = { @@ -481,7 +481,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), @@ -502,7 +502,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) @@ -516,7 +516,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) @@ -531,27 +531,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 @@ -575,13 +575,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} @@ -590,7 +590,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 = { @@ -615,7 +615,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 = { @@ -653,7 +653,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 = { @@ -675,7 +675,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"], @@ -693,7 +693,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"], @@ -722,7 +722,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), @@ -746,7 +746,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)} @@ -792,7 +792,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)}, @@ -807,7 +807,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)}, @@ -822,8 +822,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 = { @@ -848,9 +848,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") @@ -869,7 +869,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() @@ -885,7 +885,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 = { @@ -912,7 +912,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 = { @@ -942,7 +942,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}") @@ -982,7 +982,7 @@ def test_policy_processor_integration(): steps=[ NormalizerProcessor(features=features, norm_map=norm_map, stats=stats), ToBatchProcessor(), - DeviceProcessor(device="cuda"), + DeviceProcessorStep(device="cuda"), ], name="test_preprocessor", ) @@ -990,7 +990,7 @@ def test_policy_processor_integration(): # Create output processor (postprocessor) that moves to CPU output_processor = DataProcessorPipeline( steps=[ - DeviceProcessor(device="cpu"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats), ], name="test_postprocessor", diff --git a/tests/processor/test_diffusion_processor.py b/tests/processor/test_diffusion_processor.py index b4f62e2037..ba39e7206d 100644 --- a/tests/processor/test_diffusion_processor.py +++ b/tests/processor/test_diffusion_processor.py @@ -26,7 +26,7 @@ from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors from lerobot.processor import ( DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, RenameProcessor, ToBatchProcessor, @@ -92,11 +92,11 @@ def test_make_diffusion_processor_basic(): 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[3], 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], UnnormalizerProcessor) @@ -271,10 +271,10 @@ def test_diffusion_processor_mixed_precision(): # Create processor preprocessor, postprocessor = make_diffusion_pre_post_processors(config, stats) - # 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_pi0_processor.py b/tests/processor/test_pi0_processor.py index 056814f37b..d2611a4535 100644 --- a/tests/processor/test_pi0_processor.py +++ b/tests/processor/test_pi0_processor.py @@ -25,7 +25,7 @@ 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, + DeviceProcessorStep, NormalizerProcessor, RenameProcessor, ToBatchProcessor, @@ -97,11 +97,11 @@ def test_make_pi0_processor_basic(): assert isinstance(preprocessor.steps[2], ToBatchProcessor) assert isinstance(preprocessor.steps[3], Pi0NewLineProcessor) # Step 4 would be TokenizerProcessor but it's mocked - assert isinstance(preprocessor.steps[5], DeviceProcessor) + 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[0], DeviceProcessorStep) assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) diff --git a/tests/processor/test_sac_processor.py b/tests/processor/test_sac_processor.py index 23e859fbe6..2b82de18b9 100644 --- a/tests/processor/test_sac_processor.py +++ b/tests/processor/test_sac_processor.py @@ -26,7 +26,7 @@ from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors from lerobot.processor import ( DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, RenameProcessor, ToBatchProcessor, @@ -89,11 +89,11 @@ def test_make_sac_processor_basic(): 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[3], 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], UnnormalizerProcessor) @@ -254,10 +254,10 @@ def test_sac_processor_mixed_precision(): # Create processor preprocessor, postprocessor = make_sac_pre_post_processors(config, stats) - # 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 d327317397..f05b13a75c 100644 --- a/tests/processor/test_smolvla_processor.py +++ b/tests/processor/test_smolvla_processor.py @@ -28,7 +28,7 @@ make_smolvla_pre_post_processors, ) from lerobot.processor import ( - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, RenameProcessor, ToBatchProcessor, @@ -102,11 +102,11 @@ def test_make_smolvla_processor_basic(): assert isinstance(preprocessor.steps[2], ToBatchProcessor) assert isinstance(preprocessor.steps[3], SmolVLANewLineProcessor) # Step 4 would be TokenizerProcessor but it's mocked - assert isinstance(preprocessor.steps[5], DeviceProcessor) + 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[0], DeviceProcessorStep) assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) diff --git a/tests/processor/test_tdmpc_processor.py b/tests/processor/test_tdmpc_processor.py index ac95791d5a..78cd796596 100644 --- a/tests/processor/test_tdmpc_processor.py +++ b/tests/processor/test_tdmpc_processor.py @@ -26,7 +26,7 @@ from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors from lerobot.processor import ( DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, RenameProcessor, ToBatchProcessor, @@ -92,11 +92,11 @@ def test_make_tdmpc_processor_basic(): 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[3], 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], UnnormalizerProcessor) @@ -278,10 +278,10 @@ def test_tdmpc_processor_mixed_precision(): # Create processor preprocessor, postprocessor = make_tdmpc_pre_post_processors(config, stats) - # 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 b25560400d..eeb36db1c7 100644 --- a/tests/processor/test_tokenizer_processor.py +++ b/tests/processor/test_tokenizer_processor.py @@ -943,15 +943,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 TokenizerProcessor 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 + # Create pipeline with TokenizerProcessor then DeviceProcessorStep tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6) - device_processor = DeviceProcessor(device="cuda:0") + device_processor = DeviceProcessorStep(device="cuda:0") robot_processor = DataProcessorPipeline([tokenizer_processor, device_processor]) # Start with CPU tensors @@ -963,7 +963,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" diff --git a/tests/processor/test_vqbet_processor.py b/tests/processor/test_vqbet_processor.py index a0c56bdf4c..a6d91e3704 100644 --- a/tests/processor/test_vqbet_processor.py +++ b/tests/processor/test_vqbet_processor.py @@ -26,7 +26,7 @@ from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors from lerobot.processor import ( DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, NormalizerProcessor, RenameProcessor, ToBatchProcessor, @@ -92,11 +92,11 @@ def test_make_vqbet_processor_basic(): 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[3], 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], UnnormalizerProcessor) @@ -271,10 +271,10 @@ def test_vqbet_processor_mixed_precision(): # Create processor preprocessor, postprocessor = make_vqbet_pre_post_processors(config, stats) - # 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 = { From d18732239eaf9d373a520a0815cecda2e469a495 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Fri, 29 Aug 2025 18:21:54 +0200 Subject: [PATCH 05/12] refactor(processor): replace ToBatchProcessor with AddBatchDimensionProcessorStep across multiple files for consistency (#1815) --- src/lerobot/policies/act/processor_act.py | 4 +- .../policies/diffusion/processor_diffusion.py | 4 +- src/lerobot/policies/pi0/processor_pi0.py | 4 +- .../policies/pi0fast/processor_pi0fast.py | 4 +- src/lerobot/policies/sac/processor_sac.py | 4 +- .../policies/smolvla/processor_smolvla.py | 4 +- src/lerobot/policies/tdmpc/processor_tdmpc.py | 4 +- src/lerobot/policies/vqbet/processor_vqbet.py | 4 +- src/lerobot/processor/__init__.py | 4 +- src/lerobot/processor/batch_processor.py | 2 +- .../processor/migrate_policy_normalization.py | 10 +- src/lerobot/scripts/rl/gym_manipulator.py | 12 +- tests/processor/test_act_processor.py | 4 +- tests/processor/test_batch_processor.py | 112 +++++++++--------- tests/processor/test_device_processor.py | 8 +- tests/processor/test_diffusion_processor.py | 4 +- tests/processor/test_pi0_processor.py | 4 +- tests/processor/test_sac_processor.py | 4 +- tests/processor/test_smolvla_processor.py | 4 +- tests/processor/test_tdmpc_processor.py | 4 +- tests/processor/test_vqbet_processor.py | 4 +- 21 files changed, 104 insertions(+), 104 deletions(-) diff --git a/src/lerobot/policies/act/processor_act.py b/src/lerobot/policies/act/processor_act.py index e9f3fa264e..7628192ef3 100644 --- a/src/lerobot/policies/act/processor_act.py +++ b/src/lerobot/policies/act/processor_act.py @@ -18,11 +18,11 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.act.configuration_act import ACTConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) @@ -37,7 +37,7 @@ def make_act_pre_post_processors( norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), DeviceProcessorStep(device=config.device), ] output_steps = [ diff --git a/src/lerobot/policies/diffusion/processor_diffusion.py b/src/lerobot/policies/diffusion/processor_diffusion.py index 15f0b5cc42..44444373c1 100644 --- a/src/lerobot/policies/diffusion/processor_diffusion.py +++ b/src/lerobot/policies/diffusion/processor_diffusion.py @@ -19,11 +19,11 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) @@ -38,7 +38,7 @@ def make_diffusion_pre_post_processors( norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), DeviceProcessorStep(device=config.device), ] output_steps = [ diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index b538a4abc8..3426c13ffd 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -20,10 +20,10 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, - ToBatchProcessor, TokenizerProcessor, UnnormalizerProcessor, ) @@ -75,7 +75,7 @@ def make_pi0_pre_post_processors( norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma TokenizerProcessor( tokenizer_name="google/paligemma-3b-pt-224", diff --git a/src/lerobot/policies/pi0fast/processor_pi0fast.py b/src/lerobot/policies/pi0fast/processor_pi0fast.py index 11fc05c838..4e8d3ea8ef 100644 --- a/src/lerobot/policies/pi0fast/processor_pi0fast.py +++ b/src/lerobot/policies/pi0fast/processor_pi0fast.py @@ -19,11 +19,11 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) @@ -38,7 +38,7 @@ def make_pi0fast_pre_post_processors( norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), DeviceProcessorStep(device=config.device), ] output_steps = [ diff --git a/src/lerobot/policies/sac/processor_sac.py b/src/lerobot/policies/sac/processor_sac.py index 2cbdb7fd3c..539682bfff 100644 --- a/src/lerobot/policies/sac/processor_sac.py +++ b/src/lerobot/policies/sac/processor_sac.py @@ -20,11 +20,11 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) @@ -39,7 +39,7 @@ def make_sac_pre_post_processors( norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), DeviceProcessorStep(device=config.device), ] output_steps = [ diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 16a6fabcfe..b4eb590c12 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -19,11 +19,11 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, - ToBatchProcessor, TokenizerProcessor, UnnormalizerProcessor, ) @@ -43,7 +43,7 @@ def make_smolvla_pre_post_processors( norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), SmolVLANewLineProcessor(), TokenizerProcessor( tokenizer_name=config.vlm_model_name, diff --git a/src/lerobot/policies/tdmpc/processor_tdmpc.py b/src/lerobot/policies/tdmpc/processor_tdmpc.py index faf73e0039..f63bd63663 100644 --- a/src/lerobot/policies/tdmpc/processor_tdmpc.py +++ b/src/lerobot/policies/tdmpc/processor_tdmpc.py @@ -19,11 +19,11 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) @@ -38,7 +38,7 @@ def make_tdmpc_pre_post_processors( norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), DeviceProcessorStep(device=config.device), ] output_steps = [ diff --git a/src/lerobot/policies/vqbet/processor_vqbet.py b/src/lerobot/policies/vqbet/processor_vqbet.py index ef65433322..2e47c8f7b5 100644 --- a/src/lerobot/policies/vqbet/processor_vqbet.py +++ b/src/lerobot/policies/vqbet/processor_vqbet.py @@ -20,11 +20,11 @@ from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig from lerobot.processor import ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) @@ -39,7 +39,7 @@ def make_vqbet_pre_post_processors( norm_map=config.normalization_mapping, stats=dataset_stats, ), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), DeviceProcessorStep(device=config.device), ] output_steps = [ diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index cc390c588d..c09301a96f 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 .delta_action_processor import MapDeltaActionToRobotAction, MapTensorToDeltaActionDict from .device_processor import DeviceProcessorStep from .gym_action_processor import Numpy2TorchActionProcessor, Torch2NumpyActionProcessor @@ -78,7 +78,7 @@ "DataProcessorPipeline", "RobotProcessorPipeline", "PolicyProcessorPipeline", - "ToBatchProcessor", + "AddBatchDimensionProcessorStep", "TokenizerProcessor", "TimeLimitProcessor", "Numpy2TorchActionProcessor", diff --git a/src/lerobot/processor/batch_processor.py b/src/lerobot/processor/batch_processor.py index 3f42be8ee5..93af7d6381 100644 --- a/src/lerobot/processor/batch_processor.py +++ b/src/lerobot/processor/batch_processor.py @@ -92,7 +92,7 @@ def complementary_data(self, complementary_data): @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: diff --git a/src/lerobot/processor/migrate_policy_normalization.py b/src/lerobot/processor/migrate_policy_normalization.py index 348a18d58a..103cc548f0 100644 --- a/src/lerobot/processor/migrate_policy_normalization.py +++ b/src/lerobot/processor/migrate_policy_normalization.py @@ -46,8 +46,8 @@ from safetensors.torch import load_file as load_safetensors from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.processor.batch_processor import ToBatchProcessor -from lerobot.processor.device_processor import DeviceProcessor +from lerobot.processor.batch_processor import AddBatchDimensionProcessorStep +from lerobot.processor.device_processor import DeviceProcessorStep from lerobot.processor.normalize_processor import NormalizerProcessor, UnnormalizerProcessor from lerobot.processor.pipeline import DataProcessorPipeline from lerobot.processor.rename_processor import RenameProcessor @@ -418,14 +418,14 @@ def main(): 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"), + DeviceProcessorStep(device="cpu"), UnnormalizerProcessor(features=output_features, norm_map=norm_map, stats=stats), ] postprocessor = DataProcessorPipeline(steps=postprocessor_steps, name="robot_postprocessor") diff --git a/src/lerobot/scripts/rl/gym_manipulator.py b/src/lerobot/scripts/rl/gym_manipulator.py index 1a2324a412..005d52aa44 100644 --- a/src/lerobot/scripts/rl/gym_manipulator.py +++ b/src/lerobot/scripts/rl/gym_manipulator.py @@ -29,10 +29,11 @@ from lerobot.envs.configs import HILSerlRobotEnvConfig from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( + AddBatchDimensionProcessorStep, AddTeleopActionAsComplimentaryData, AddTeleopEventsAsInfo, DataProcessorPipeline, - DeviceProcessor, + DeviceProcessorStep, GripperPenaltyProcessor, ImageCropResizeProcessor, InterventionActionProcessor, @@ -43,7 +44,6 @@ Numpy2TorchActionProcessor, RewardClassifierProcessor, TimeLimitProcessor, - ToBatchProcessor, Torch2NumpyActionProcessor, VanillaObservationProcessor, ) @@ -383,8 +383,8 @@ def make_processors( env_pipeline_steps = [ Numpy2TorchActionProcessor(), VanillaObservationProcessor(), - ToBatchProcessor(), - DeviceProcessor(device=device), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=device), ] return DataProcessorPipeline(steps=env_pipeline_steps), DataProcessorPipeline( @@ -457,8 +457,8 @@ 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), diff --git a/tests/processor/test_act_processor.py b/tests/processor/test_act_processor.py index 50df9597ad..d5c4fc5d8b 100644 --- a/tests/processor/test_act_processor.py +++ b/tests/processor/test_act_processor.py @@ -25,11 +25,11 @@ 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, DeviceProcessorStep, NormalizerProcessor, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) from lerobot.processor.pipeline import TransitionKey @@ -88,7 +88,7 @@ def test_make_act_processor_basic(): 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[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor diff --git a/tests/processor/test_batch_processor.py b/tests/processor/test_batch_processor.py index ecb694143e..f07379bf83 100644 --- a/tests/processor/test_batch_processor.py +++ b/tests/processor/test_batch_processor.py @@ -23,7 +23,7 @@ from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE from lerobot.processor import DataProcessorPipeline, ProcessorStepRegistry -from lerobot.processor.batch_processor import ToBatchProcessor +from lerobot.processor.batch_processor import AddBatchDimensionProcessorStep from lerobot.processor.pipeline import TransitionKey @@ -44,7 +44,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) @@ -60,7 +60,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) @@ -76,7 +76,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) @@ -92,7 +92,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) @@ -117,7 +117,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) @@ -143,7 +143,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 @@ -166,7 +166,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 @@ -189,7 +189,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) @@ -199,7 +199,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) @@ -211,7 +211,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 @@ -243,8 +243,8 @@ def test_mixed_observation(): def test_integration_with_robot_processor(): - """Test ToBatchProcessor integration with DataProcessorPipeline.""" - to_batch_processor = ToBatchProcessor() + """Test AddBatchDimensionProcessorStep integration with DataProcessorPipeline.""" + to_batch_processor = AddBatchDimensionProcessorStep() pipeline = DataProcessorPipeline([to_batch_processor]) # Create unbatched observation @@ -263,7 +263,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() @@ -283,8 +283,8 @@ def test_serialization_methods(): def test_save_and_load_pretrained(): - """Test saving and loading ToBatchProcessor with DataProcessorPipeline.""" - processor = ToBatchProcessor() + """Test saving and loading AddBatchDimensionProcessorStep with DataProcessorPipeline.""" + processor = AddBatchDimensionProcessorStep() pipeline = DataProcessorPipeline([processor], name="BatchPipeline") with tempfile.TemporaryDirectory() as tmp_dir: @@ -300,7 +300,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)} @@ -311,10 +311,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() @@ -322,7 +322,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]) with tempfile.TemporaryDirectory() as tmp_dir: @@ -346,7 +346,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") @@ -370,7 +370,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 @@ -407,7 +407,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) @@ -429,7 +429,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) @@ -444,7 +444,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) @@ -463,7 +463,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) @@ -480,7 +480,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) @@ -493,7 +493,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] @@ -522,7 +522,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) @@ -531,7 +531,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 = { @@ -551,7 +551,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] @@ -568,7 +568,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") @@ -589,7 +589,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] @@ -605,7 +605,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([]) @@ -627,7 +627,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"} @@ -644,7 +644,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"} @@ -663,7 +663,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 = [ @@ -689,7 +689,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) @@ -699,7 +699,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) @@ -711,7 +711,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, @@ -729,7 +729,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", @@ -754,7 +754,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 = { @@ -779,7 +779,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 = [ @@ -837,7 +837,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", @@ -865,7 +865,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) @@ -882,7 +882,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) @@ -899,7 +899,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) @@ -929,7 +929,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) @@ -950,7 +950,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) @@ -971,7 +971,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 @@ -988,7 +988,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] @@ -1011,7 +1011,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 = { @@ -1053,7 +1053,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") @@ -1077,7 +1077,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) @@ -1092,7 +1092,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) @@ -1114,7 +1114,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_device_processor.py b/tests/processor/test_device_processor.py index 5c693857ed..3b4ebe4bc3 100644 --- a/tests/processor/test_device_processor.py +++ b/tests/processor/test_device_processor.py @@ -305,11 +305,11 @@ def test_features(): def test_integration_with_robot_processor(): """Test integration with DataProcessorPipeline.""" from lerobot.constants import OBS_STATE - from lerobot.processor import ToBatchProcessor + from lerobot.processor import AddBatchDimensionProcessorStep # Create a pipeline with DeviceProcessorStep device_processor = DeviceProcessorStep(device="cpu") - batch_processor = ToBatchProcessor() + batch_processor = AddBatchDimensionProcessorStep() processor = DataProcessorPipeline(steps=[batch_processor, device_processor], name="test_pipeline") @@ -962,7 +962,7 @@ 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, NormalizerProcessor, UnnormalizerProcessor # Create features and stats features = { @@ -981,7 +981,7 @@ def test_policy_processor_integration(): input_processor = DataProcessorPipeline( steps=[ NormalizerProcessor(features=features, norm_map=norm_map, stats=stats), - ToBatchProcessor(), + AddBatchDimensionProcessorStep(), DeviceProcessorStep(device="cuda"), ], name="test_preprocessor", diff --git a/tests/processor/test_diffusion_processor.py b/tests/processor/test_diffusion_processor.py index ba39e7206d..d36d8b88fc 100644 --- a/tests/processor/test_diffusion_processor.py +++ b/tests/processor/test_diffusion_processor.py @@ -25,11 +25,11 @@ 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, DeviceProcessorStep, NormalizerProcessor, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) from lerobot.processor.pipeline import TransitionKey @@ -91,7 +91,7 @@ def test_make_diffusion_processor_basic(): 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[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor diff --git a/tests/processor/test_pi0_processor.py b/tests/processor/test_pi0_processor.py index d2611a4535..fd594d7ea7 100644 --- a/tests/processor/test_pi0_processor.py +++ b/tests/processor/test_pi0_processor.py @@ -25,10 +25,10 @@ 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 ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) from lerobot.processor.pipeline import TransitionKey @@ -94,7 +94,7 @@ def test_make_pi0_processor_basic(): 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[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], Pi0NewLineProcessor) # Step 4 would be TokenizerProcessor but it's mocked assert isinstance(preprocessor.steps[5], DeviceProcessorStep) diff --git a/tests/processor/test_sac_processor.py b/tests/processor/test_sac_processor.py index 2b82de18b9..7a19c44c5f 100644 --- a/tests/processor/test_sac_processor.py +++ b/tests/processor/test_sac_processor.py @@ -25,11 +25,11 @@ 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, DeviceProcessorStep, NormalizerProcessor, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) from lerobot.processor.pipeline import TransitionKey @@ -88,7 +88,7 @@ def test_make_sac_processor_basic(): 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[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor diff --git a/tests/processor/test_smolvla_processor.py b/tests/processor/test_smolvla_processor.py index f05b13a75c..f0dbb93a92 100644 --- a/tests/processor/test_smolvla_processor.py +++ b/tests/processor/test_smolvla_processor.py @@ -28,10 +28,10 @@ make_smolvla_pre_post_processors, ) from lerobot.processor import ( + AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessor, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) from lerobot.processor.pipeline import TransitionKey @@ -99,7 +99,7 @@ def test_make_smolvla_processor_basic(): 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[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], SmolVLANewLineProcessor) # Step 4 would be TokenizerProcessor but it's mocked assert isinstance(preprocessor.steps[5], DeviceProcessorStep) diff --git a/tests/processor/test_tdmpc_processor.py b/tests/processor/test_tdmpc_processor.py index 78cd796596..9f9e0de505 100644 --- a/tests/processor/test_tdmpc_processor.py +++ b/tests/processor/test_tdmpc_processor.py @@ -25,11 +25,11 @@ 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, DeviceProcessorStep, NormalizerProcessor, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) from lerobot.processor.pipeline import TransitionKey @@ -91,7 +91,7 @@ def test_make_tdmpc_processor_basic(): 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[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor diff --git a/tests/processor/test_vqbet_processor.py b/tests/processor/test_vqbet_processor.py index a6d91e3704..e878dc802b 100644 --- a/tests/processor/test_vqbet_processor.py +++ b/tests/processor/test_vqbet_processor.py @@ -25,11 +25,11 @@ 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, DeviceProcessorStep, NormalizerProcessor, RenameProcessor, - ToBatchProcessor, UnnormalizerProcessor, ) from lerobot.processor.pipeline import TransitionKey @@ -91,7 +91,7 @@ def test_make_vqbet_processor_basic(): 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[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], DeviceProcessorStep) # Check steps in postprocessor From 80987ce298f7a7abac9658d19835d034e638a2bf Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Fri, 29 Aug 2025 18:46:17 +0200 Subject: [PATCH 06/12] chore(processor): rename tokenizer step and tokenizer field (#1817) * refactor(processor): rename internal tokenizer field for clarity in TokenizerProcessor * refactor(processor): rename TokenizerProcessor to TokenizerProcessorStep for consistency across multiple files --- src/lerobot/policies/pi0/processor_pi0.py | 4 +- .../policies/smolvla/processor_smolvla.py | 4 +- src/lerobot/processor/__init__.py | 4 +- src/lerobot/processor/tokenizer_processor.py | 18 ++-- tests/processor/test_pi0_processor.py | 12 +-- tests/processor/test_smolvla_processor.py | 12 +-- tests/processor/test_tokenizer_processor.py | 88 ++++++++++--------- 7 files changed, 74 insertions(+), 68 deletions(-) diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index 3426c13ffd..f78474895f 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -24,7 +24,7 @@ DeviceProcessorStep, NormalizerProcessor, PolicyProcessorPipeline, - TokenizerProcessor, + TokenizerProcessorStep, UnnormalizerProcessor, ) from lerobot.processor.pipeline import ( @@ -77,7 +77,7 @@ def make_pi0_pre_post_processors( ), 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", diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index b4eb590c12..999621e27e 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -24,7 +24,7 @@ NormalizerProcessor, PolicyProcessorPipeline, RenameProcessor, - TokenizerProcessor, + TokenizerProcessorStep, UnnormalizerProcessor, ) from lerobot.processor.pipeline import ( @@ -45,7 +45,7 @@ def make_smolvla_pre_post_processors( ), AddBatchDimensionProcessorStep(), SmolVLANewLineProcessor(), - TokenizerProcessor( + TokenizerProcessorStep( tokenizer_name=config.vlm_model_name, padding=config.pad_language_to, padding_side="right", diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index c09301a96f..b78d57bb7b 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -47,7 +47,7 @@ TruncatedProcessorStep, ) from .rename_processor import RenameProcessor -from .tokenizer_processor import TokenizerProcessor +from .tokenizer_processor import TokenizerProcessorStep __all__ = [ "ActionProcessorStep", @@ -79,7 +79,7 @@ "RobotProcessorPipeline", "PolicyProcessorPipeline", "AddBatchDimensionProcessorStep", - "TokenizerProcessor", + "TokenizerProcessorStep", "TimeLimitProcessor", "Numpy2TorchActionProcessor", "Torch2NumpyActionProcessor", diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py index 1d444e3363..8332fe5191 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -27,7 +27,7 @@ @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 @@ -51,7 +51,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: @@ -59,7 +59,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) ``` """ @@ -72,23 +72,23 @@ class TokenizerProcessor(ObservationProcessorStep): truncation: bool = True # Internal tokenizer instance (not serialized) - _tokenizer: Any = field(default=None, init=False, repr=False) + input_tokenizer: Any = field(default=None, init=False, repr=False) def __post_init__(self): """Initialize the tokenizer from the provided tokenizer or tokenizer name.""" 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: # Use provided tokenizer object directly - self._tokenizer = self.tokenizer + self.input_tokenizer = self.tokenizer elif self.tokenizer_name is not None: if AutoTokenizer is None: raise ImportError("AutoTokenizer is not available") - self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) + self.input_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) else: raise ValueError( "Either 'tokenizer' or 'tokenizer_name' must be provided. " @@ -198,7 +198,7 @@ def _tokenize_text(self, text: str | list[str]) -> dict[str, torch.Tensor]: Returns: Dictionary containing tokenized output with keys like 'input_ids', 'attention_mask'. """ - return self._tokenizer( + return self.input_tokenizer( text, max_length=self.max_length, truncation=self.truncation, @@ -222,7 +222,7 @@ def get_config(self) -> dict[str, Any]: } # Only include tokenizer_name if it was used (not when tokenizer object was provided) - # TODO(steven): Consider saving the name of the _tokenizer if it was loaded + # TODO(steven): Consider saving the name of the input_tokenizer if it was loaded if self.tokenizer_name is not None and self.tokenizer is None: config["tokenizer_name"] = self.tokenizer_name diff --git a/tests/processor/test_pi0_processor.py b/tests/processor/test_pi0_processor.py index fd594d7ea7..aca811ca82 100644 --- a/tests/processor/test_pi0_processor.py +++ b/tests/processor/test_pi0_processor.py @@ -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) # Check processor names @@ -96,7 +96,7 @@ def test_make_pi0_processor_basic(): assert isinstance(preprocessor.steps[1], NormalizerProcessor) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], Pi0NewLineProcessor) - # Step 4 would be TokenizerProcessor but it's mocked + # Step 4 would be TokenizerProcessorStep but it's mocked assert isinstance(preprocessor.steps[5], DeviceProcessorStep) # Check steps in postprocessor @@ -182,7 +182,7 @@ def get_config(self): def transform_features(self, features): return features - with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor): + with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep", MockTokenizerProcessor): preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats) # Create CPU data @@ -232,7 +232,7 @@ def get_config(self): def transform_features(self, features): return features - with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor): + with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep", MockTokenizerProcessor): preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats) # Simulate Accelerate: data already on GPU and batched @@ -283,7 +283,7 @@ def get_config(self): def transform_features(self, features): return features - with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessor", MockTokenizerProcessor): + with patch("lerobot.policies.pi0.processor_pi0.TokenizerProcessorStep", MockTokenizerProcessor): preprocessor, postprocessor = make_pi0_pre_post_processors(config, stats) # Simulate data on different GPU @@ -309,7 +309,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) # Should still create processors diff --git a/tests/processor/test_smolvla_processor.py b/tests/processor/test_smolvla_processor.py index f0dbb93a92..127dc87ada 100644 --- a/tests/processor/test_smolvla_processor.py +++ b/tests/processor/test_smolvla_processor.py @@ -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) # Check processor names @@ -101,7 +101,7 @@ def test_make_smolvla_processor_basic(): assert isinstance(preprocessor.steps[1], NormalizerProcessor) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], SmolVLANewLineProcessor) - # Step 4 would be TokenizerProcessor but it's mocked + # Step 4 would be TokenizerProcessorStep but it's mocked assert isinstance(preprocessor.steps[5], DeviceProcessorStep) # Check steps in postprocessor @@ -187,7 +187,7 @@ def get_config(self): def transform_features(self, features): return features - with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor): + with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessor): preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats) # Create CPU data @@ -237,7 +237,7 @@ def get_config(self): def transform_features(self, features): return features - with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor): + with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessor): preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats) # Simulate Accelerate: data already on GPU and batched @@ -288,7 +288,7 @@ def get_config(self): def transform_features(self, features): return features - with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessor", MockTokenizerProcessor): + with patch("lerobot.policies.smolvla.processor_smolvla.TokenizerProcessorStep", MockTokenizerProcessor): preprocessor, postprocessor = make_smolvla_pre_post_processors(config, stats) # Simulate data on different GPU @@ -314,7 +314,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) # Should still create processors diff --git a/tests/processor/test_tokenizer_processor.py b/tests/processor/test_tokenizer_processor.py index eeb36db1c7..10f44087f8 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 @@ -11,7 +11,7 @@ from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.constants import OBS_LANGUAGE from lerobot.processor.pipeline import DataProcessorPipeline, TransitionKey -from lerobot.processor.tokenizer_processor import TokenizerProcessor +from lerobot.processor.tokenizer_processor import TokenizerProcessorStep from tests.utils import require_package @@ -96,7 +96,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])}, @@ -128,7 +128,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])}, @@ -162,7 +162,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])}, @@ -190,7 +190,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])}, @@ -216,7 +216,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) @@ -231,7 +231,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"}) @@ -246,7 +246,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}) @@ -261,7 +261,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}) @@ -280,7 +280,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") @@ -291,7 +291,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") @@ -301,7 +301,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", @@ -328,7 +328,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", @@ -358,7 +358,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() @@ -375,7 +375,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() @@ -388,7 +388,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]) transition = create_transition( @@ -423,7 +423,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" ) @@ -454,7 +454,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]) @@ -490,14 +492,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,)), @@ -528,7 +530,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) @@ -548,7 +550,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,)), @@ -585,7 +587,7 @@ def __call__(self, *args, **kwargs): 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", @@ -617,7 +619,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])}, @@ -652,7 +654,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])}, @@ -680,7 +682,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])}, @@ -704,7 +706,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( @@ -754,7 +756,9 @@ def __call__( 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])}, @@ -766,7 +770,9 @@ def __call__( 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) @@ -777,7 +783,7 @@ def __call__( 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 @@ -801,7 +807,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 @@ -826,7 +832,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") @@ -850,7 +856,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( @@ -872,7 +878,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 @@ -900,7 +906,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 @@ -923,7 +929,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)} @@ -943,14 +949,14 @@ 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 DeviceProcessorStep in pipeline.""" + """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 DeviceProcessorStep - tokenizer_processor = TokenizerProcessor(tokenizer_name="test-tokenizer", max_length=6) + # 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]) @@ -979,7 +985,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") From d428bb5161109aba73e406f24fabf57023feacb0 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Fri, 29 Aug 2025 19:32:01 +0200 Subject: [PATCH 07/12] refactor(processor): update action keys from target to displacement format across multiple files (#1816) --- .../processor/delta_action_processor.py | 44 +++++++++---------- .../robot_kinematic_processor.py | 27 ++++++------ .../teleoperators/phone/phone_processor.py | 26 +++++------ 3 files changed, 49 insertions(+), 48 deletions(-) diff --git a/src/lerobot/processor/delta_action_processor.py b/src/lerobot/processor/delta_action_processor.py index 343cb3b388..a3bd51de38 100644 --- a/src/lerobot/processor/delta_action_processor.py +++ b/src/lerobot/processor/delta_action_processor.py @@ -64,12 +64,12 @@ class MapDeltaActionToRobotAction(ActionProcessorStep): Output ACTION keys: { "action.enabled": bool, - "action.target_x": float, - "action.target_y": float, - "action.target_z": float, - "action.target_wx": float, - "action.target_wy": float, - "action.target_wz": float, + "action.displacement_x": float, + "action.displacement_y": float, + "action.displacement_z": float, + "action.displacement_wx": float, + "action.displacement_wy": float, + "action.displacement_wz": float, "action.gripper": float, } """ @@ -99,19 +99,19 @@ def action(self, action: dict) -> dict: # For gamepad/keyboard, we don't have rotation input, so set to 0 # These could be extended in the future for more sophisticated teleoperators - target_wx = 0.0 - target_wy = 0.0 - target_wz = 0.0 + displacement_wx = 0.0 + displacement_wy = 0.0 + displacement_wz = 0.0 - # Update action with robot target format + # Update action with robot displacement format action = { "action.enabled": enabled, - "action.target_x": scaled_delta_x, - "action.target_y": scaled_delta_y, - "action.target_z": scaled_delta_z, - "action.target_wx": target_wx, - "action.target_wy": target_wy, - "action.target_wz": target_wz, + "action.displacement_x": scaled_delta_x, + "action.displacement_y": scaled_delta_y, + "action.displacement_z": scaled_delta_z, + "action.displacement_wx": displacement_wx, + "action.displacement_wy": displacement_wy, + "action.displacement_wz": displacement_wz, "action.gripper": float(gripper), } @@ -123,12 +123,12 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po features.update( { "action.enabled": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), - "action.target_x": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), - "action.target_y": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), - "action.target_z": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), - "action.target_wx": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), - "action.target_wy": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), - "action.target_wz": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), + "action.displacement_x": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), + "action.displacement_y": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), + "action.displacement_z": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), + "action.displacement_wx": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), + "action.displacement_wy": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), + "action.displacement_wz": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), "action.gripper": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), } ) diff --git a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py index 2b086d2951..fae3f1f592 100644 --- a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py +++ b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py @@ -41,7 +41,8 @@ class EEReferenceAndDelta(ActionProcessorStep): Input ACTION keys: { - "action.ee.{x,y,z,wx,wy,wz}" : float + "action.enabled": bool, + "action.displacement_{x,y,z,wx,wy,wz}" : float "complementary_data.raw_joint_positions": dict, } @@ -82,12 +83,12 @@ def action(self, action): t_curr = self.kinematics.forward_kinematics(q) enabled = bool(new_action.pop("action.enabled", 0)) - tx = float(new_action.pop("action.target_x", 0.0)) - ty = float(new_action.pop("action.target_y", 0.0)) - tz = float(new_action.pop("action.target_z", 0.0)) - wx = float(new_action.pop("action.target_wx", 0.0)) - wy = float(new_action.pop("action.target_wy", 0.0)) - wz = float(new_action.pop("action.target_wz", 0.0)) + tx = float(new_action.pop("action.displacement_x", 0.0)) + ty = float(new_action.pop("action.displacement_y", 0.0)) + tz = float(new_action.pop("action.displacement_z", 0.0)) + wx = float(new_action.pop("action.displacement_wx", 0.0)) + wy = float(new_action.pop("action.displacement_wy", 0.0)) + wz = float(new_action.pop("action.displacement_wz", 0.0)) desired = None @@ -140,12 +141,12 @@ def reset(self): def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]: features.pop("action.enabled", None) - features.pop("action.target_x", None) - features.pop("action.target_y", None) - features.pop("action.target_z", None) - features.pop("action.target_wx", None) - features.pop("action.target_wy", None) - features.pop("action.target_wz", None) + features.pop("action.displacement_x", None) + features.pop("action.displacement_y", None) + features.pop("action.displacement_z", None) + features.pop("action.displacement_wx", None) + features.pop("action.displacement_wy", None) + features.pop("action.displacement_wz", None) features["action.ee.x"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) features["action.ee.y"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) diff --git a/src/lerobot/teleoperators/phone/phone_processor.py b/src/lerobot/teleoperators/phone/phone_processor.py index 63f11fb98d..849cf203cf 100644 --- a/src/lerobot/teleoperators/phone/phone_processor.py +++ b/src/lerobot/teleoperators/phone/phone_processor.py @@ -38,7 +38,7 @@ class MapPhoneActionToRobotAction(ActionProcessorStep): Output ACTION keys: { "action.enabled": bool, - "action.ee.{x,y,z,wx,wy,wz}" : float + "action.displacement_{x,y,z,wx,wy,wz}" : float "action.gripper": float, } """ @@ -70,12 +70,12 @@ def action(self, act: dict) -> dict: # For some actions we need to invert the axis act["action.enabled"] = enabled - act["action.target_x"] = -pos[1] if enabled else 0.0 - act["action.target_y"] = pos[0] if enabled else 0.0 - act["action.target_z"] = pos[2] if enabled else 0.0 - act["action.target_wx"] = rotvec[1] if enabled else 0.0 - act["action.target_wy"] = rotvec[0] if enabled else 0.0 - act["action.target_wz"] = -rotvec[2] if enabled else 0.0 + act["action.displacement_x"] = -pos[1] if enabled else 0.0 + act["action.displacement_y"] = pos[0] if enabled else 0.0 + act["action.displacement_z"] = pos[2] if enabled else 0.0 + act["action.displacement_wx"] = rotvec[1] if enabled else 0.0 + act["action.displacement_wy"] = rotvec[0] if enabled else 0.0 + act["action.displacement_wz"] = -rotvec[2] if enabled else 0.0 act["action.gripper"] = gripper # Still send gripper action when disabled return act @@ -86,11 +86,11 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po features.pop("action.phone.raw_inputs", None) features["action.enabled"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) - features["action.target_x"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) - features["action.target_y"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) - features["action.target_z"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) - features["action.target_wx"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) - features["action.target_wy"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) - features["action.target_wz"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) + features["action.displacement_x"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) + features["action.displacement_y"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) + features["action.displacement_z"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) + features["action.displacement_wx"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) + features["action.displacement_wy"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) + features["action.displacement_wz"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) features["action.gripper"] = (PolicyFeature(type=FeatureType.ACTION, shape=(1,)),) return features From 780a5e5bf9afccf2eead79bfb20870893570086f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 1 Sep 2025 07:54:32 +0000 Subject: [PATCH 08/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/lerobot/robots/so100_follower/robot_kinematic_processor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py index d18d400035..b7b0a7178b 100644 --- a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py +++ b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py @@ -82,7 +82,7 @@ def action(self, action): # Current pose from FK on measured joints t_curr = self.kinematics.forward_kinematics(q) - + enabled = bool(new_action.pop(f"{ACTION}.enabled", 0)) tx = float(new_action.pop(f"{ACTION}.target_x", 0.0)) ty = float(new_action.pop(f"{ACTION}.target_y", 0.0)) From 3d0997d8f0db62a8eb199c245b1e1fc60613e06a Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Mon, 1 Sep 2025 10:18:52 +0200 Subject: [PATCH 09/12] refactor(converters): rename to_transition_robot_observation to robot_observation_to_transition for clarity (#1825) - Updated references across evaluate.py, record.py, teleoperate.py, and other files to reflect the new naming convention. - Improved readability and maintainability of the codebase. --- examples/phone_to_so100/evaluate.py | 4 ++-- examples/phone_to_so100/record.py | 4 ++-- src/lerobot/processor/converters.py | 2 +- src/lerobot/record.py | 4 ++-- src/lerobot/teleoperate.py | 4 ++-- tests/processor/test_converters.py | 6 +++--- 6 files changed, 12 insertions(+), 12 deletions(-) diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index d67fc125bb..038be22510 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -22,8 +22,8 @@ from lerobot.policies.act.modeling_act import ACTPolicy from lerobot.policies.factory import make_pre_post_processors from lerobot.processor.converters import ( + robot_observation_to_transition, to_output_robot_action, - to_transition_robot_observation, ) from lerobot.processor.pipeline import DataProcessorPipeline from lerobot.record import record_loop @@ -83,7 +83,7 @@ steps=[ ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ], - to_transition=to_transition_robot_observation, + to_transition=robot_observation_to_transition, to_output=lambda tr: tr, ) diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index 575199e66c..74eb9e50a5 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -21,8 +21,8 @@ from lerobot.datasets.utils import merge_features from lerobot.model.kinematics import RobotKinematics from lerobot.processor.converters import ( + robot_observation_to_transition, to_output_robot_action, - to_transition_robot_observation, to_transition_teleop_action, ) from lerobot.processor.pipeline import DataProcessorPipeline @@ -114,7 +114,7 @@ steps=[ ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ], - to_transition=to_transition_robot_observation, + to_transition=robot_observation_to_transition, to_output=lambda tr: tr, ) diff --git a/src/lerobot/processor/converters.py b/src/lerobot/processor/converters.py index c9218a6505..fd3dfef7d8 100644 --- a/src/lerobot/processor/converters.py +++ b/src/lerobot/processor/converters.py @@ -94,7 +94,7 @@ def to_transition_teleop_action(action: dict[str, Any]) -> EnvTransition: # TODO(Adil, Pepijn): Overtime we can maybe add these converters to pipeline.py itself -def to_transition_robot_observation(observation: dict[str, Any]) -> EnvTransition: +def robot_observation_to_transition(observation: dict[str, Any]) -> EnvTransition: """ Convert a raw robot observation dict into an EnvTransition under the OBSERVATION TransitionKey. """ diff --git a/src/lerobot/record.py b/src/lerobot/record.py index e467c93f8e..7d4ea5b4db 100644 --- a/src/lerobot/record.py +++ b/src/lerobot/record.py @@ -78,9 +78,9 @@ from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.processor import PolicyProcessorPipeline, RobotProcessorPipeline from lerobot.processor.converters import ( + robot_observation_to_transition, to_dataset_frame, to_output_robot_action, - to_transition_robot_observation, to_transition_teleop_action, ) from lerobot.processor.pipeline import IdentityProcessorStep, TransitionKey @@ -253,7 +253,7 @@ def record_loop( ) robot_observation_processor = robot_observation_processor or RobotProcessorPipeline( steps=[IdentityProcessorStep()], - to_transition=to_transition_robot_observation, + to_transition=robot_observation_to_transition, to_output=lambda tr: tr, ) diff --git a/src/lerobot/teleoperate.py b/src/lerobot/teleoperate.py index 404e8a5bf1..e197ba8a09 100644 --- a/src/lerobot/teleoperate.py +++ b/src/lerobot/teleoperate.py @@ -63,8 +63,8 @@ from lerobot.configs import parser from lerobot.processor import RobotProcessorPipeline from lerobot.processor.converters import ( + robot_observation_to_transition, to_output_robot_action, - to_transition_robot_observation, to_transition_teleop_action, ) from lerobot.processor.pipeline import IdentityProcessorStep @@ -131,7 +131,7 @@ def teleop_loop( ) robot_observation_processor = robot_observation_processor or RobotProcessorPipeline( steps=[IdentityProcessorStep()], - to_transition=to_transition_robot_observation, + to_transition=robot_observation_to_transition, to_output=lambda tr: tr, ) diff --git a/tests/processor/test_converters.py b/tests/processor/test_converters.py index 590f6a8927..51225bc123 100644 --- a/tests/processor/test_converters.py +++ b/tests/processor/test_converters.py @@ -3,9 +3,9 @@ import torch from lerobot.processor.converters import ( + robot_observation_to_transition, to_dataset_frame, to_output_robot_action, - to_transition_robot_observation, to_transition_teleop_action, ) from lerobot.processor.pipeline import TransitionKey @@ -47,7 +47,7 @@ def test_to_transition_teleop_action_prefix_and_tensor_conversion(): assert tr[TransitionKey.OBSERVATION] == {} -def test_to_transition_robot_observation_state_vs_images_split(): +def test_robot_observation_to_transition_state_vs_images_split(): # Create an observation with mixed content img = np.full((10, 20, 3), 255, dtype=np.uint8) # image (uint8 HWC) obs = { @@ -58,7 +58,7 @@ def test_to_transition_robot_observation_state_vs_images_split(): "arr": np.array([1.5, 2.5]), # vector to state to torch tensor } - tr = to_transition_robot_observation(obs) + tr = robot_observation_to_transition(obs) assert isinstance(tr, dict) assert TransitionKey.OBSERVATION in tr From 5b42b95cd09e91d618e279f1f216d05aa66f33f4 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Mon, 1 Sep 2025 10:31:12 +0200 Subject: [PATCH 10/12] refactor(datasets): rename merge_features to merge_dataset_features for clarity (#1826) - Updated references in evaluate.py, record.py, and test_dataset_utils.py to reflect the new function name. - Improved code readability and maintainability by standardizing the naming convention. --- examples/phone_to_so100/evaluate.py | 4 ++-- examples/phone_to_so100/record.py | 4 ++-- src/lerobot/datasets/utils.py | 2 +- tests/datasets/test_dataset_utils.py | 12 ++++++------ 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index 038be22510..b2ad48878e 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -17,7 +17,7 @@ from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features -from lerobot.datasets.utils import merge_features +from lerobot.datasets.utils import merge_dataset_features from lerobot.model.kinematics import RobotKinematics from lerobot.policies.act.modeling_act import ACTPolicy from lerobot.policies.factory import make_pre_post_processors @@ -103,7 +103,7 @@ patterns=["observation.state.ee"], ) # Get all ee observation features -dataset_features = merge_features(obs_ee, action_ee_and_gripper) +dataset_features = merge_dataset_features(obs_ee, action_ee_and_gripper) print("All dataset features: ", dataset_features) diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index 74eb9e50a5..e5b6865601 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -18,7 +18,7 @@ from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features -from lerobot.datasets.utils import merge_features +from lerobot.datasets.utils import merge_dataset_features from lerobot.model.kinematics import RobotKinematics from lerobot.processor.converters import ( robot_observation_to_transition, @@ -142,7 +142,7 @@ patterns=["observation.state.ee"], ) -dataset_features = merge_features(action_ee, gripper, observation_ee) +dataset_features = merge_dataset_features(action_ee, gripper, observation_ee) print("All dataset features: ", dataset_features) diff --git a/src/lerobot/datasets/utils.py b/src/lerobot/datasets/utils.py index db60e63b30..ce60b38d1e 100644 --- a/src/lerobot/datasets/utils.py +++ b/src/lerobot/datasets/utils.py @@ -470,7 +470,7 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea return policy_features -def merge_features(*dicts: dict) -> dict: +def merge_dataset_features(*dicts: dict) -> dict: """ Merge LeRobot grouped feature dicts. diff --git a/tests/datasets/test_dataset_utils.py b/tests/datasets/test_dataset_utils.py index ae09fb2624..95a8dd2889 100644 --- a/tests/datasets/test_dataset_utils.py +++ b/tests/datasets/test_dataset_utils.py @@ -20,7 +20,7 @@ from huggingface_hub import DatasetCard from lerobot.datasets.push_dataset_to_hub.utils import calculate_episode_data_index -from lerobot.datasets.utils import create_lerobot_dataset_card, hf_transform_to_torch, merge_features +from lerobot.datasets.utils import create_lerobot_dataset_card, hf_transform_to_torch, merge_dataset_features def test_default_parameters(): @@ -72,7 +72,7 @@ def test_merge_simple_vectors(): } } - out = merge_features(g1, g2) + out = merge_dataset_features(g1, g2) assert "action" in out assert out["action"]["dtype"] == "float32" @@ -87,7 +87,7 @@ def test_merge_multiple_groups_order_and_dedup(): g2 = {"action": {"dtype": "float32", "shape": (2,), "names": ["b", "c"]}} g3 = {"action": {"dtype": "float32", "shape": (3,), "names": ["a", "c", "d"]}} - out = merge_features(g1, g2, g3) + out = merge_dataset_features(g1, g2, g3) assert out["action"]["names"] == ["a", "b", "c", "d"] assert out["action"]["shape"] == (4,) @@ -110,7 +110,7 @@ def test_non_vector_last_wins_for_images(): } } - out = merge_features(g1, g2) + out = merge_dataset_features(g1, g2) assert out["observation.images.front"]["shape"] == (3, 720, 1280) assert out["observation.images.front"]["dtype"] == "image" @@ -120,13 +120,13 @@ def test_dtype_mismatch_raises(): g2 = {"action": {"dtype": "float64", "shape": (1,), "names": ["b"]}} with pytest.raises(ValueError, match="dtype mismatch for 'action'"): - _ = merge_features(g1, g2) + _ = merge_dataset_features(g1, g2) def test_non_dict_passthrough_last_wins(): g1 = {"misc": 123} g2 = {"misc": 456} - out = merge_features(g1, g2) + out = merge_dataset_features(g1, g2) # For non-dict entries the last one wins assert out["misc"] == 456 From 88adca1d966b9a960c52ef84e530c454ebb65478 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Mon, 1 Sep 2025 10:39:31 +0200 Subject: [PATCH 11/12] refactor(teleoperators): rename _wait_for_capture_trigger to _capture_calibration_pose_on_trigger for clarity (#1827) - Updated method names in IOSPhone and AndroidPhone classes to improve readability and maintainability. - Ensured consistent naming convention across teleoperation modules. --- src/lerobot/teleoperators/phone/teleop_phone.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/lerobot/teleoperators/phone/teleop_phone.py b/src/lerobot/teleoperators/phone/teleop_phone.py index ed985c0815..2ce89933f9 100644 --- a/src/lerobot/teleoperators/phone/teleop_phone.py +++ b/src/lerobot/teleoperators/phone/teleop_phone.py @@ -102,13 +102,13 @@ def calibrate(self) -> None: ) print("Press and hold B1 in the HEBI Mobile I/O app to capture this pose...\n") - pos, rot = self._wait_for_capture_trigger() + pos, rot = self._capture_calibration_pose_on_trigger() self._calib_pos = pos.copy() self._calib_rot_inv = rot.inv() self._enabled = False print("Calibration done\n") - def _wait_for_capture_trigger(self) -> tuple[np.ndarray, Rotation]: + def _capture_calibration_pose_on_trigger(self) -> tuple[np.ndarray, Rotation]: """Wait trigger for calibration: iOS: B1. Android: 'move'.""" while True: ok, pos, rot, pose = self._read_current_pose() @@ -222,13 +222,13 @@ def calibrate(self) -> None: ) print("Touch and move on the WebXR page to capture this pose...\n") - pos, rot = self._wait_for_capture_trigger() + pos, rot = self._capture_calibration_pose_on_trigger() self._calib_pos = pos.copy() self._calib_rot_inv = rot.inv() self._enabled = False print("Calibration done\n") - def _wait_for_capture_trigger(self) -> tuple[np.ndarray, Rotation]: + def _capture_calibration_pose_on_trigger(self) -> tuple[np.ndarray, Rotation]: """Wait trigger for calibration: iOS: B1. Android: 'move'.""" while True: with self._android_lock: From f7515837bed053d3e7dea067156d66b448df29fe Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Mon, 1 Sep 2025 11:12:59 +0200 Subject: [PATCH 12/12] refactor(processor): standardize processor class naming with 'Step' suffix for consistency (#1829) * refactor(processor): rename MapTensorToDeltaActionDict to MapTensorToDeltaActionDictStep for consistency - Updated references in init.py and gym_manipulator.py to reflect the new class name. - Improved clarity and maintainability of the codebase by standardizing naming conventions. * refactor(processor): rename MapDeltaActionToRobotAction to MapDeltaActionToRobotActionStep for consistency - Updated references in init.py, gym_manipulator.py, and delta_action_processor.py to reflect the new class name. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions. * refactor(processor): rename Torch2NumpyActionProcessor to Torch2NumpyActionProcessorStep for consistency - Updated references in init.py, gym_action_processor.py, and gym_manipulator.py to reflect the new class name. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions. * refactor(processor): rename Numpy2TorchActionProcessor to Numpy2TorchActionProcessorStep for consistency - Updated references in init.py, gym_action_processor.py, and gym_manipulator.py to reflect the new class name. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions. * refactor(processor): rename processors for consistency and clarity - Updated class names in hil_processor.py, init.py, and gym_manipulator.py to include 'Step' suffix for better consistency. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions across processor classes. * refactor(processor): rename JointVelocityProcessor and MotorCurrentProcessor for consistency - Updated class names in joint_observations_processor.py, init.py, and gym_manipulator.py to include 'Step' suffix for better consistency. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions across processor classes. * refactor(processor): rename NormalizerProcessor and UnnormalizerProcessor to include 'Step' suffix for consistency - Updated class names in normalize_processor.py and various processor files to reflect the new naming convention. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions across processor classes. * refactor(processor): rename VanillaObservationProcessor to VanillaObservationProcessorStep for consistency - Updated references in observation_processor.py, init.py, gym_manipulator.py, and test_observation_processor.py to reflect the new class name. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions across processor classes. * refactor(processor): rename NormalizerProcessor and UnnormalizerProcessor to NormalizerProcessorStep and UnnormalizerProcessorStep for consistency - Updated references in migrate_policy_normalization.py to reflect the new class names. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions across processor classes. * refactor(processor): rename RenameProcessor to RenameProcessorStep for consistency - Updated references across multiple processor files and tests to reflect the new class name. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions across processor classes. * refactor(processor): rename ToBatchProcessor classes to include 'Step' suffix for consistency - Updated class names in batch_processor.py to reflect the new naming convention. - Enhanced clarity and maintainability of the codebase by standardizing naming conventions across processor classes. * refactor(tests): update config path in test_save_and_load_pretrained to reflect RenameProcessorStep naming - Adjusted the configuration file path in the test to match the updated class name "TestRenameProcessorStep". - Ensured consistency in naming conventions across test files. --- docs/source/hilserl.mdx | 24 ++-- src/lerobot/policies/act/processor_act.py | 12 +- .../policies/diffusion/processor_diffusion.py | 12 +- src/lerobot/policies/pi0/processor_pi0.py | 12 +- .../policies/pi0fast/processor_pi0fast.py | 12 +- src/lerobot/policies/sac/processor_sac.py | 12 +- .../sac/reward_model/processor_classifier.py | 6 +- .../policies/smolvla/processor_smolvla.py | 12 +- src/lerobot/policies/tdmpc/processor_tdmpc.py | 12 +- src/lerobot/policies/vqbet/processor_vqbet.py | 12 +- src/lerobot/processor/__init__.py | 62 ++++----- src/lerobot/processor/batch_processor.py | 16 +-- .../processor/delta_action_processor.py | 4 +- src/lerobot/processor/gym_action_processor.py | 4 +- src/lerobot/processor/hil_processor.py | 14 +- .../processor/joint_observations_processor.py | 4 +- .../processor/migrate_policy_normalization.py | 14 +- src/lerobot/processor/normalize_processor.py | 10 +- .../processor/observation_processor.py | 2 +- src/lerobot/processor/rename_processor.py | 2 +- src/lerobot/scripts/rl/gym_manipulator.py | 60 ++++----- tests/processor/test_act_processor.py | 12 +- tests/processor/test_classifier_processor.py | 6 +- tests/processor/test_device_processor.py | 10 +- tests/processor/test_diffusion_processor.py | 12 +- tests/processor/test_normalize_processor.py | 122 +++++++++--------- tests/processor/test_observation_processor.py | 50 +++---- tests/processor/test_pi0_processor.py | 12 +- tests/processor/test_rename_processor.py | 62 ++++----- tests/processor/test_sac_processor.py | 12 +- tests/processor/test_smolvla_processor.py | 12 +- tests/processor/test_tdmpc_processor.py | 12 +- tests/processor/test_vqbet_processor.py | 12 +- 33 files changed, 329 insertions(+), 323 deletions(-) diff --git a/docs/source/hilserl.mdx b/docs/source/hilserl.mdx index dfdf431ac1..cb235776e7 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 +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. **DeviceProcessor**: 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 7628192ef3..5fa1690736 100644 --- a/src/lerobot/policies/act/processor_act.py +++ b/src/lerobot/policies/act/processor_act.py @@ -20,10 +20,10 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, PolicyProcessorPipeline, - RenameProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -31,8 +31,8 @@ def make_act_pre_post_processors( config: ACTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: 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, @@ -42,7 +42,7 @@ def make_act_pre_post_processors( ] output_steps = [ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor( + 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 44444373c1..b2078e7a44 100644 --- a/src/lerobot/policies/diffusion/processor_diffusion.py +++ b/src/lerobot/policies/diffusion/processor_diffusion.py @@ -21,10 +21,10 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, PolicyProcessorPipeline, - RenameProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -32,8 +32,8 @@ def make_diffusion_pre_post_processors( config: DiffusionConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: 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, @@ -43,7 +43,7 @@ def make_diffusion_pre_post_processors( ] output_steps = [ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor( + 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 f78474895f..4c1237fd1a 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -22,17 +22,17 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, PolicyProcessorPipeline, TokenizerProcessorStep, - UnnormalizerProcessor, + UnnormalizerProcessorStep, ) from lerobot.processor.pipeline import ( ComplementaryDataProcessorStep, ProcessorStep, ProcessorStepRegistry, ) -from lerobot.processor.rename_processor import RenameProcessor +from lerobot.processor.rename_processor import RenameProcessorStep @ProcessorStepRegistry.register(name="pi0_new_line_processor") @@ -69,8 +69,8 @@ def make_pi0_pre_post_processors( ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: # 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, @@ -88,7 +88,7 @@ def make_pi0_pre_post_processors( output_steps: list[ProcessorStep] = [ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor( + 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 4e8d3ea8ef..70d0328d08 100644 --- a/src/lerobot/policies/pi0fast/processor_pi0fast.py +++ b/src/lerobot/policies/pi0fast/processor_pi0fast.py @@ -21,10 +21,10 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, PolicyProcessorPipeline, - RenameProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -32,8 +32,8 @@ def make_pi0fast_pre_post_processors( config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: 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, @@ -43,7 +43,7 @@ def make_pi0fast_pre_post_processors( ] output_steps = [ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor( + 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 539682bfff..92c90d1b04 100644 --- a/src/lerobot/policies/sac/processor_sac.py +++ b/src/lerobot/policies/sac/processor_sac.py @@ -22,10 +22,10 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, PolicyProcessorPipeline, - RenameProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -33,8 +33,8 @@ def make_sac_pre_post_processors( config: SACConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: 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, @@ -44,7 +44,7 @@ def make_sac_pre_post_processors( ] output_steps = [ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor( + 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 857ae450de..8dc86df727 100644 --- a/src/lerobot/policies/sac/reward_model/processor_classifier.py +++ b/src/lerobot/policies/sac/reward_model/processor_classifier.py @@ -19,7 +19,7 @@ from lerobot.processor import ( DeviceProcessorStep, IdentityProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, PolicyProcessorPipeline, ) @@ -28,10 +28,10 @@ def make_classifier_processor( config: RewardClassifierConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: 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 ), DeviceProcessorStep(device=config.device), diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 999621e27e..9afbe2f977 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -21,11 +21,11 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, PolicyProcessorPipeline, - RenameProcessor, + RenameProcessorStep, TokenizerProcessorStep, - UnnormalizerProcessor, + UnnormalizerProcessorStep, ) from lerobot.processor.pipeline import ( ComplementaryDataProcessorStep, @@ -37,8 +37,8 @@ def make_smolvla_pre_post_processors( config: SmolVLAConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: 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, @@ -55,7 +55,7 @@ def make_smolvla_pre_post_processors( ] output_steps = [ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor( + 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 f63bd63663..5b9f50ff55 100644 --- a/src/lerobot/policies/tdmpc/processor_tdmpc.py +++ b/src/lerobot/policies/tdmpc/processor_tdmpc.py @@ -21,10 +21,10 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, PolicyProcessorPipeline, - RenameProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -32,8 +32,8 @@ def make_tdmpc_pre_post_processors( config: TDMPCConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: 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, @@ -43,7 +43,7 @@ def make_tdmpc_pre_post_processors( ] output_steps = [ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor( + 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 2e47c8f7b5..6ec538ed23 100644 --- a/src/lerobot/policies/vqbet/processor_vqbet.py +++ b/src/lerobot/policies/vqbet/processor_vqbet.py @@ -22,10 +22,10 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, PolicyProcessorPipeline, - RenameProcessor, - UnnormalizerProcessor, + RenameProcessorStep, + UnnormalizerProcessorStep, ) @@ -33,8 +33,8 @@ def make_vqbet_pre_post_processors( config: VQBeTConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: 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, @@ -44,7 +44,7 @@ def make_vqbet_pre_post_processors( ] output_steps = [ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor( + 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 b78d57bb7b..d7a29b5206 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -15,21 +15,21 @@ # limitations under the License. from .batch_processor import AddBatchDimensionProcessorStep -from .delta_action_processor import MapDeltaActionToRobotAction, MapTensorToDeltaActionDict +from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep from .device_processor import DeviceProcessorStep -from .gym_action_processor import Numpy2TorchActionProcessor, Torch2NumpyActionProcessor +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, DataProcessorPipeline, @@ -46,44 +46,44 @@ TransitionKey, TruncatedProcessorStep, ) -from .rename_processor import RenameProcessor +from .rename_processor import RenameProcessorStep from .tokenizer_processor import TokenizerProcessorStep __all__ = [ "ActionProcessorStep", - "AddTeleopActionAsComplimentaryData", - "AddTeleopEventsAsInfo", + "AddTeleopActionAsComplimentaryDataStep", + "AddTeleopEventsAsInfoStep", "DeviceProcessorStep", "RewardProcessorStep", - "MapDeltaActionToRobotAction", - "MapTensorToDeltaActionDict", + "MapDeltaActionToRobotActionStep", + "MapTensorToDeltaActionDictStep", "EnvTransition", - "GripperPenaltyProcessor", + "GripperPenaltyProcessorStep", "IdentityProcessorStep", - "ImageCropResizeProcessor", + "ImageCropResizeProcessorStep", "InfoProcessorStep", - "InterventionActionProcessor", - "JointVelocityProcessor", - "MapDeltaActionToRobotAction", - "MotorCurrentProcessor", - "NormalizerProcessor", - "UnnormalizerProcessor", + "InterventionActionProcessorStep", + "JointVelocityProcessorStep", + "MapDeltaActionToRobotActionStep", + "MotorCurrentProcessorStep", + "NormalizerProcessorStep", + "UnnormalizerProcessorStep", "hotswap_stats", "ObservationProcessorStep", "ProcessorStep", "ProcessorStepRegistry", - "RenameProcessor", - "RewardClassifierProcessor", + "RenameProcessorStep", + "RewardClassifierProcessorStep", "DoneProcessorStep", "DataProcessorPipeline", "RobotProcessorPipeline", "PolicyProcessorPipeline", "AddBatchDimensionProcessorStep", "TokenizerProcessorStep", - "TimeLimitProcessor", - "Numpy2TorchActionProcessor", - "Torch2NumpyActionProcessor", + "TimeLimitProcessorStep", + "Numpy2TorchActionProcessorStep", + "Torch2NumpyActionProcessorStep", "TransitionKey", "TruncatedProcessorStep", - "VanillaObservationProcessor", + "VanillaObservationProcessorStep", ] diff --git a/src/lerobot/processor/batch_processor.py b/src/lerobot/processor/batch_processor.py index 93af7d6381..b3577b8602 100644 --- a/src/lerobot/processor/batch_processor.py +++ b/src/lerobot/processor/batch_processor.py @@ -28,7 +28,7 @@ @dataclass @ProcessorStepRegistry.register(name="to_batch_processor_action") -class ToBatchProcessorAction(ActionProcessorStep): +class ToBatchProcessorActionStep(ActionProcessorStep): """Process action component in-place, adding batch dimension if needed.""" def action(self, action): @@ -40,7 +40,7 @@ def action(self, action): @dataclass @ProcessorStepRegistry.register(name="to_batch_processor_observation") -class ToBatchProcessorObservation(ObservationProcessorStep): +class ToBatchProcessorObservationStep(ObservationProcessorStep): """Process observation component in-place, adding batch dimensions where needed.""" def observation(self, observation): @@ -66,7 +66,7 @@ def observation(self, observation): @dataclass @ProcessorStepRegistry.register(name="to_batch_processor_complementary_data") -class ToBatchProcessorComplementaryData(ComplementaryDataProcessorStep): +class ToBatchProcessorComplementaryDataStep(ComplementaryDataProcessorStep): """Process complementary data in-place, handling task field batching.""" def complementary_data(self, complementary_data): @@ -127,12 +127,12 @@ class AddBatchDimensionProcessorStep(ProcessorStep): ``` """ - to_batch_action_processor: ToBatchProcessorAction = field(default_factory=ToBatchProcessorAction) - to_batch_observation_processor: ToBatchProcessorObservation = field( - default_factory=ToBatchProcessorObservation + to_batch_action_processor: ToBatchProcessorActionStep = field(default_factory=ToBatchProcessorActionStep) + to_batch_observation_processor: ToBatchProcessorObservationStep = field( + default_factory=ToBatchProcessorObservationStep ) - to_batch_complementary_data_processor: ToBatchProcessorComplementaryData = field( - default_factory=ToBatchProcessorComplementaryData + to_batch_complementary_data_processor: ToBatchProcessorComplementaryDataStep = field( + default_factory=ToBatchProcessorComplementaryDataStep ) 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 a3bd51de38..a150f9d689 100644 --- a/src/lerobot/processor/delta_action_processor.py +++ b/src/lerobot/processor/delta_action_processor.py @@ -24,7 +24,7 @@ @ProcessorStepRegistry.register("map_tensor_to_delta_action_dict") @dataclass -class MapTensorToDeltaActionDict(ActionProcessorStep): +class MapTensorToDeltaActionDictStep(ActionProcessorStep): """ Map a tensor to a delta action dictionary. """ @@ -48,7 +48,7 @@ def action(self, action: Tensor) -> dict: @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/gym_action_processor.py b/src/lerobot/processor/gym_action_processor.py index 4ab6813ac4..37bfa4b0c4 100644 --- a/src/lerobot/processor/gym_action_processor.py +++ b/src/lerobot/processor/gym_action_processor.py @@ -21,7 +21,7 @@ @ProcessorStepRegistry.register("torch2numpy_action_processor") @dataclass -class Torch2NumpyActionProcessor(ActionProcessorStep): +class Torch2NumpyActionProcessorStep(ActionProcessorStep): """Convert PyTorch tensor actions to NumPy arrays.""" squeeze_batch_dim: bool = True @@ -50,7 +50,7 @@ def action(self, action: torch.Tensor) -> np.ndarray: @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 02caa663ef..bbb7c12b7b 100644 --- a/src/lerobot/processor/hil_processor.py +++ b/src/lerobot/processor/hil_processor.py @@ -29,7 +29,7 @@ @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 @@ -42,7 +42,7 @@ def complementary_data(self, complementary_data: dict) -> dict: @ProcessorStepRegistry.register("add_teleop_action_as_info") @dataclass -class AddTeleopEventsAsInfo(InfoProcessorStep): +class AddTeleopEventsAsInfoStep(InfoProcessorStep): """Add teleoperator control events to transition info.""" teleop_device: Teleoperator @@ -56,7 +56,7 @@ def info(self, info: dict) -> dict: @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 @@ -106,7 +106,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po @dataclass @ProcessorStepRegistry.register("time_limit_processor") -class TimeLimitProcessor(TruncatedProcessorStep): +class TimeLimitProcessorStep(TruncatedProcessorStep): """Track episode steps and enforce time limits.""" max_episode_steps: int @@ -130,7 +130,7 @@ def reset(self) -> None: @dataclass @ProcessorStepRegistry.register("gripper_penalty_processor") -class GripperPenaltyProcessor(ComplementaryDataProcessorStep): +class GripperPenaltyProcessorStep(ComplementaryDataProcessorStep): """Apply penalty for inappropriate gripper usage.""" penalty: float = -0.01 @@ -176,7 +176,7 @@ def reset(self) -> None: @dataclass @ProcessorStepRegistry.register("intervention_action_processor") -class InterventionActionProcessor(ProcessorStep): +class InterventionActionProcessorStep(ProcessorStep): """Handle human intervention actions and episode termination.""" use_gripper: bool = False @@ -246,7 +246,7 @@ def get_config(self) -> dict[str, Any]: @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 aae0921cc7..baea44362d 100644 --- a/src/lerobot/processor/joint_observations_processor.py +++ b/src/lerobot/processor/joint_observations_processor.py @@ -13,7 +13,7 @@ @dataclass @ProcessorStepRegistry.register("joint_velocity_processor") -class JointVelocityProcessor(ObservationProcessorStep): +class JointVelocityProcessorStep(ObservationProcessorStep): """Add joint velocity information to observations.""" dt: float = 0.1 @@ -66,7 +66,7 @@ def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, Po @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 103cc548f0..48d4d82771 100644 --- a/src/lerobot/processor/migrate_policy_normalization.py +++ b/src/lerobot/processor/migrate_policy_normalization.py @@ -48,9 +48,9 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.processor.batch_processor import AddBatchDimensionProcessorStep from lerobot.processor.device_processor import DeviceProcessorStep -from lerobot.processor.normalize_processor import NormalizerProcessor, UnnormalizerProcessor +from lerobot.processor.normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep from lerobot.processor.pipeline import DataProcessorPipeline -from lerobot.processor.rename_processor import RenameProcessor +from lerobot.processor.rename_processor import RenameProcessorStep # Policy type to class mapping POLICY_CLASSES = { @@ -403,8 +403,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", {}) @@ -412,8 +412,8 @@ 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, @@ -426,7 +426,7 @@ def main(): # Create postprocessor with unnormalizer for outputs only postprocessor_steps = [ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor(features=output_features, norm_map=norm_map, stats=stats), + 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 a75e82097d..2a9c052c01 100644 --- a/src/lerobot/processor/normalize_processor.py +++ b/src/lerobot/processor/normalize_processor.py @@ -196,7 +196,7 @@ def _apply_transform( @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. @@ -215,7 +215,7 @@ def from_lerobot_dataset( 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, @@ -245,7 +245,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition: @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. @@ -262,7 +262,7 @@ def from_lerobot_dataset( 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: @@ -288,7 +288,7 @@ def hotswap_stats( Replaces normalization statistics in a DataProcessorPipeline pipeline. This function creates a deep copy of the provided `DataProcessorPipeline` 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 1648a97981..7998eb0e23 100644 --- a/src/lerobot/processor/observation_processor.py +++ b/src/lerobot/processor/observation_processor.py @@ -27,7 +27,7 @@ @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 63dff1a3de..d927ff5659 100644 --- a/src/lerobot/processor/rename_processor.py +++ b/src/lerobot/processor/rename_processor.py @@ -26,7 +26,7 @@ @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/scripts/rl/gym_manipulator.py b/src/lerobot/scripts/rl/gym_manipulator.py index 005d52aa44..5e269b683b 100644 --- a/src/lerobot/scripts/rl/gym_manipulator.py +++ b/src/lerobot/scripts/rl/gym_manipulator.py @@ -30,22 +30,22 @@ from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( AddBatchDimensionProcessorStep, - AddTeleopActionAsComplimentaryData, - AddTeleopEventsAsInfo, + AddTeleopActionAsComplimentaryDataStep, + AddTeleopEventsAsInfoStep, DataProcessorPipeline, DeviceProcessorStep, - GripperPenaltyProcessor, - ImageCropResizeProcessor, - InterventionActionProcessor, - JointVelocityProcessor, - MapDeltaActionToRobotAction, - MapTensorToDeltaActionDict, - MotorCurrentProcessor, - Numpy2TorchActionProcessor, - RewardClassifierProcessor, - TimeLimitProcessor, - Torch2NumpyActionProcessor, - VanillaObservationProcessor, + GripperPenaltyProcessorStep, + ImageCropResizeProcessorStep, + InterventionActionProcessorStep, + JointVelocityProcessorStep, + MapDeltaActionToRobotActionStep, + MapTensorToDeltaActionDictStep, + MotorCurrentProcessorStep, + Numpy2TorchActionProcessorStep, + RewardClassifierProcessorStep, + TimeLimitProcessorStep, + Torch2NumpyActionProcessorStep, + VanillaObservationProcessorStep, ) from lerobot.processor.pipeline import EnvTransition, TransitionKey from lerobot.robots import ( # noqa: F401 @@ -375,14 +375,14 @@ 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(), + Numpy2TorchActionProcessorStep(), + VanillaObservationProcessorStep(), AddBatchDimensionProcessorStep(), DeviceProcessorStep(device=device), ] @@ -404,13 +404,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( @@ -422,7 +422,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, ) @@ -431,13 +431,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, ) @@ -448,7 +448,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, @@ -461,10 +461,10 @@ def make_processors( 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, ), @@ -474,8 +474,8 @@ 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(), - MapDeltaActionToRobotAction(), + MapTensorToDeltaActionDictStep(), + 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 d5c4fc5d8b..899f22e4dd 100644 --- a/tests/processor/test_act_processor.py +++ b/tests/processor/test_act_processor.py @@ -28,9 +28,9 @@ AddBatchDimensionProcessorStep, DataProcessorPipeline, DeviceProcessorStep, - NormalizerProcessor, - RenameProcessor, - UnnormalizerProcessor, + NormalizerProcessorStep, + RenameProcessorStep, + UnnormalizerProcessorStep, ) from lerobot.processor.pipeline import TransitionKey @@ -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[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], DeviceProcessorStep) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_act_processor_normalization(): diff --git a/tests/processor/test_classifier_processor.py b/tests/processor/test_classifier_processor.py index 546cd54041..0ab9eb79b4 100644 --- a/tests/processor/test_classifier_processor.py +++ b/tests/processor/test_classifier_processor.py @@ -28,7 +28,7 @@ DataProcessorPipeline, DeviceProcessorStep, IdentityProcessorStep, - NormalizerProcessor, + NormalizerProcessorStep, ) from lerobot.processor.pipeline import TransitionKey @@ -87,8 +87,8 @@ 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[0], NormalizerProcessorStep) # For input features + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) # For output features assert isinstance(preprocessor.steps[2], DeviceProcessorStep) # Check steps in postprocessor diff --git a/tests/processor/test_device_processor.py b/tests/processor/test_device_processor.py index 3b4ebe4bc3..d78daa62a2 100644 --- a/tests/processor/test_device_processor.py +++ b/tests/processor/test_device_processor.py @@ -962,7 +962,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 AddBatchDimensionProcessorStep, NormalizerProcessor, UnnormalizerProcessor + from lerobot.processor import ( + AddBatchDimensionProcessorStep, + NormalizerProcessorStep, + UnnormalizerProcessorStep, + ) # Create features and stats features = { @@ -980,7 +984,7 @@ 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), + NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats), AddBatchDimensionProcessorStep(), DeviceProcessorStep(device="cuda"), ], @@ -991,7 +995,7 @@ def test_policy_processor_integration(): output_processor = DataProcessorPipeline( steps=[ DeviceProcessorStep(device="cpu"), - UnnormalizerProcessor(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats), + UnnormalizerProcessorStep(features={ACTION: features[ACTION]}, norm_map=norm_map, stats=stats), ], name="test_postprocessor", ) diff --git a/tests/processor/test_diffusion_processor.py b/tests/processor/test_diffusion_processor.py index d36d8b88fc..93173d6daf 100644 --- a/tests/processor/test_diffusion_processor.py +++ b/tests/processor/test_diffusion_processor.py @@ -28,9 +28,9 @@ AddBatchDimensionProcessorStep, DataProcessorPipeline, DeviceProcessorStep, - NormalizerProcessor, - RenameProcessor, - UnnormalizerProcessor, + NormalizerProcessorStep, + RenameProcessorStep, + UnnormalizerProcessorStep, ) from lerobot.processor.pipeline import TransitionKey @@ -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[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], DeviceProcessorStep) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_diffusion_processor_with_images(): diff --git a/tests/processor/test_normalize_processor.py b/tests/processor/test_normalize_processor.py index 32240ee461..1f321cb64c 100644 --- a/tests/processor/test_normalize_processor.py +++ b/tests/processor/test_normalize_processor.py @@ -21,8 +21,8 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.processor.normalize_processor import ( - NormalizerProcessor, - UnnormalizerProcessor, + NormalizerProcessorStep, + UnnormalizerProcessorStep, _convert_stats_to_tensors, hotswap_stats, ) @@ -123,7 +123,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 { @@ -140,10 +140,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): @@ -180,7 +180,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, @@ -206,7 +206,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(), } @@ -235,7 +235,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 @@ -250,7 +250,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 @@ -301,7 +301,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} ) @@ -319,7 +319,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} ) @@ -345,7 +345,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} ) @@ -363,7 +363,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} ) @@ -379,11 +379,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 { @@ -422,7 +422,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): @@ -466,7 +466,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"} ) @@ -478,7 +478,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, @@ -535,7 +535,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) @@ -546,7 +546,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) @@ -562,7 +562,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) @@ -577,7 +577,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 @@ -586,7 +586,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, @@ -598,7 +598,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, @@ -666,7 +666,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]), @@ -691,7 +691,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) @@ -717,7 +717,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]), @@ -744,7 +744,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) @@ -767,8 +767,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]) @@ -808,7 +808,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]), @@ -850,7 +850,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]), @@ -884,8 +884,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]) @@ -917,7 +917,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() @@ -927,7 +927,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, @@ -965,7 +965,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" @@ -1002,8 +1002,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 @@ -1043,7 +1043,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 @@ -1068,7 +1068,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])}, } @@ -1083,8 +1083,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]) @@ -1113,7 +1113,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 @@ -1163,7 +1163,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, @@ -1208,10 +1208,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]) @@ -1260,7 +1260,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 @@ -1316,7 +1316,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]) # Process with original stats @@ -1360,7 +1360,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)) @@ -1372,7 +1372,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)) @@ -1387,7 +1387,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"} ) @@ -1405,12 +1405,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])}}, @@ -1432,7 +1432,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) @@ -1446,7 +1446,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] @@ -1458,7 +1458,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) @@ -1477,8 +1477,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 3196416e32..3a5c1a60c1 100644 --- a/tests/processor/test_observation_processor.py +++ b/tests/processor/test_observation_processor.py @@ -20,7 +20,7 @@ from lerobot.configs.types import FeatureType from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE -from lerobot.processor import VanillaObservationProcessor +from lerobot.processor import VanillaObservationProcessorStep from lerobot.processor.pipeline import TransitionKey from tests.conftest import assert_contract_is_typed @@ -42,7 +42,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) @@ -68,7 +68,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) @@ -91,7 +91,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) @@ -108,7 +108,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) @@ -121,7 +121,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) @@ -134,7 +134,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) @@ -149,7 +149,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) @@ -159,7 +159,7 @@ def test_none_observation(): def test_serialization_methods(): """Test serialization methods.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() # Test get_config config = processor.get_config() @@ -178,7 +178,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} @@ -199,7 +199,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} @@ -220,7 +220,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) @@ -238,7 +238,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) @@ -263,7 +263,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) @@ -277,7 +277,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) @@ -314,7 +314,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} @@ -329,7 +329,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} @@ -344,7 +344,7 @@ def test_state_only_processing(): def test_empty_observation(): """Test processing empty observation.""" - processor = VanillaObservationProcessor() + processor = VanillaObservationProcessorStep() observation = {} transition = create_transition(observation=observation) @@ -360,7 +360,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) @@ -387,7 +387,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) @@ -411,7 +411,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,)), @@ -425,7 +425,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,)), @@ -439,7 +439,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)), @@ -457,7 +457,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,)), @@ -473,7 +473,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 aca811ca82..2e1af8ecbb 100644 --- a/tests/processor/test_pi0_processor.py +++ b/tests/processor/test_pi0_processor.py @@ -27,9 +27,9 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, - RenameProcessor, - UnnormalizerProcessor, + NormalizerProcessorStep, + RenameProcessorStep, + UnnormalizerProcessorStep, ) from lerobot.processor.pipeline import TransitionKey @@ -92,8 +92,8 @@ 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[0], RenameProcessorStep) + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], Pi0NewLineProcessor) # Step 4 would be TokenizerProcessorStep but it's mocked @@ -102,7 +102,7 @@ def test_make_pi0_processor_basic(): # Check steps in postprocessor assert len(postprocessor.steps) == 2 assert isinstance(postprocessor.steps[0], DeviceProcessorStep) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_pi0_newline_processor_single_task(): diff --git a/tests/processor/test_rename_processor.py b/tests/processor/test_rename_processor.py index e31d706262..928f7f1d49 100644 --- a/tests/processor/test_rename_processor.py +++ b/tests/processor/test_rename_processor.py @@ -20,7 +20,7 @@ import torch from lerobot.configs.types import FeatureType -from lerobot.processor import DataProcessorPipeline, ProcessorStepRegistry, RenameProcessor, TransitionKey +from lerobot.processor import DataProcessorPipeline, ProcessorStepRegistry, RenameProcessorStep, TransitionKey from lerobot.processor.rename_processor import rename_stats from tests.conftest import assert_contract_is_typed @@ -46,7 +46,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]), @@ -74,7 +74,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]), @@ -93,7 +93,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) @@ -108,7 +108,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, @@ -133,7 +133,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), @@ -163,15 +163,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 == {} @@ -186,7 +186,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]) @@ -220,30 +220,32 @@ 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 # Load pipeline loaded_pipeline = DataProcessorPipeline.from_pretrained(tmp_dir) - 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 @@ -260,23 +262,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]) with tempfile.TemporaryDirectory() as tmp_dir: @@ -298,14 +300,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", @@ -313,7 +315,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", @@ -356,7 +358,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), @@ -386,7 +388,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) @@ -414,7 +416,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,)), @@ -435,7 +437,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,)), @@ -450,8 +452,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 7a19c44c5f..6e1811bf94 100644 --- a/tests/processor/test_sac_processor.py +++ b/tests/processor/test_sac_processor.py @@ -28,9 +28,9 @@ AddBatchDimensionProcessorStep, DataProcessorPipeline, DeviceProcessorStep, - NormalizerProcessor, - RenameProcessor, - UnnormalizerProcessor, + NormalizerProcessorStep, + RenameProcessorStep, + UnnormalizerProcessorStep, ) from lerobot.processor.pipeline import TransitionKey @@ -86,15 +86,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[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], DeviceProcessorStep) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_sac_processor_normalization_modes(): diff --git a/tests/processor/test_smolvla_processor.py b/tests/processor/test_smolvla_processor.py index 127dc87ada..3bc2eda55f 100644 --- a/tests/processor/test_smolvla_processor.py +++ b/tests/processor/test_smolvla_processor.py @@ -30,9 +30,9 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, - NormalizerProcessor, - RenameProcessor, - UnnormalizerProcessor, + NormalizerProcessorStep, + RenameProcessorStep, + UnnormalizerProcessorStep, ) from lerobot.processor.pipeline import TransitionKey @@ -97,8 +97,8 @@ 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[0], RenameProcessorStep) + assert isinstance(preprocessor.steps[1], NormalizerProcessorStep) assert isinstance(preprocessor.steps[2], AddBatchDimensionProcessorStep) assert isinstance(preprocessor.steps[3], SmolVLANewLineProcessor) # Step 4 would be TokenizerProcessorStep but it's mocked @@ -107,7 +107,7 @@ def test_make_smolvla_processor_basic(): # Check steps in postprocessor assert len(postprocessor.steps) == 2 assert isinstance(postprocessor.steps[0], DeviceProcessorStep) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_smolvla_newline_processor_single_task(): diff --git a/tests/processor/test_tdmpc_processor.py b/tests/processor/test_tdmpc_processor.py index 9f9e0de505..b9cd0decf0 100644 --- a/tests/processor/test_tdmpc_processor.py +++ b/tests/processor/test_tdmpc_processor.py @@ -28,9 +28,9 @@ AddBatchDimensionProcessorStep, DataProcessorPipeline, DeviceProcessorStep, - NormalizerProcessor, - RenameProcessor, - UnnormalizerProcessor, + NormalizerProcessorStep, + RenameProcessorStep, + UnnormalizerProcessorStep, ) from lerobot.processor.pipeline import TransitionKey @@ -89,15 +89,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[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], DeviceProcessorStep) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_tdmpc_processor_normalization(): diff --git a/tests/processor/test_vqbet_processor.py b/tests/processor/test_vqbet_processor.py index e878dc802b..a384ffe1d2 100644 --- a/tests/processor/test_vqbet_processor.py +++ b/tests/processor/test_vqbet_processor.py @@ -28,9 +28,9 @@ AddBatchDimensionProcessorStep, DataProcessorPipeline, DeviceProcessorStep, - NormalizerProcessor, - RenameProcessor, - UnnormalizerProcessor, + NormalizerProcessorStep, + RenameProcessorStep, + UnnormalizerProcessorStep, ) from lerobot.processor.pipeline import TransitionKey @@ -89,15 +89,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[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], DeviceProcessorStep) - assert isinstance(postprocessor.steps[1], UnnormalizerProcessor) + assert isinstance(postprocessor.steps[1], UnnormalizerProcessorStep) def test_vqbet_processor_with_images():