From aeb70812c14f3517afd5ac091d659d63fd4186c3 Mon Sep 17 00:00:00 2001 From: AdilZouitine Date: Thu, 11 Sep 2025 14:24:36 +0200 Subject: [PATCH] refactor(processor): unify action imports and enhance type clarity across multiple files - Updated imports in various files to include RobotAction and PolicyAction directly from the processor module, improving clarity and consistency. - Removed redundant imports from core, streamlining the codebase and enhancing maintainability. - Adjusted type annotations and references in the RobotProcessorPipeline and related components to align with the new import structure, ensuring better type safety and readability. --- examples/phone_to_so100/evaluate.py | 3 +-- examples/phone_to_so100/record.py | 3 +-- examples/phone_to_so100/replay.py | 3 +-- examples/phone_to_so100/teleoperate.py | 3 +-- src/lerobot/policies/act/processor_act.py | 2 +- src/lerobot/policies/diffusion/processor_diffusion.py | 2 +- src/lerobot/policies/factory.py | 3 +-- src/lerobot/policies/pi0/processor_pi0.py | 2 +- src/lerobot/policies/pi0fast/processor_pi0fast.py | 2 +- src/lerobot/policies/sac/processor_sac.py | 2 +- .../policies/sac/reward_model/processor_classifier.py | 2 +- src/lerobot/policies/smolvla/processor_smolvla.py | 2 +- src/lerobot/policies/tdmpc/processor_tdmpc.py | 2 +- src/lerobot/policies/vqbet/processor_vqbet.py | 2 +- src/lerobot/processor/__init__.py | 3 +++ src/lerobot/record.py | 3 ++- src/lerobot/replay.py | 3 +-- .../robots/so100_follower/robot_kinematic_processor.py | 2 +- src/lerobot/scripts/eval.py | 3 +-- src/lerobot/teleoperate.py | 9 +++++++-- src/lerobot/utils/control_utils.py | 3 +-- tests/artifacts/policies/save_policy_to_safetensors.py | 3 +-- 22 files changed, 31 insertions(+), 31 deletions(-) diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index bd7272b0b..1715412b2 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -23,13 +23,12 @@ from lerobot.datasets.utils import combine_feature_dicts from lerobot.model.kinematics import RobotKinematics from lerobot.policies.act.modeling_act import ACTPolicy from lerobot.policies.factory import make_pre_post_processors -from lerobot.processor import RobotProcessorPipeline +from lerobot.processor import EnvTransition, RobotAction, RobotProcessorPipeline from lerobot.processor.converters import ( identity_transition, observation_to_transition, transition_to_robot_action, ) -from lerobot.processor.core import EnvTransition, RobotAction 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 ( diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index f310a12d1..f0aa0e5ad 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -21,14 +21,13 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features from lerobot.datasets.utils import combine_feature_dicts from lerobot.model.kinematics import RobotKinematics -from lerobot.processor import RobotProcessorPipeline +from lerobot.processor import EnvTransition, RobotAction, RobotProcessorPipeline from lerobot.processor.converters import ( identity_transition, observation_to_transition, robot_action_to_transition, transition_to_robot_action, ) -from lerobot.processor.core import EnvTransition, RobotAction 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 ( diff --git a/examples/phone_to_so100/replay.py b/examples/phone_to_so100/replay.py index 4f4dcc62f..51fdc71b0 100644 --- a/examples/phone_to_so100/replay.py +++ b/examples/phone_to_so100/replay.py @@ -19,9 +19,8 @@ import time from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.model.kinematics import RobotKinematics -from lerobot.processor import RobotProcessorPipeline +from lerobot.processor import RobotAction, RobotProcessorPipeline from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action -from lerobot.processor.core import RobotAction from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.robot_kinematic_processor import ( AddRobotObservationAsComplimentaryData, diff --git a/examples/phone_to_so100/teleoperate.py b/examples/phone_to_so100/teleoperate.py index b91842d5a..b0ba95dee 100644 --- a/examples/phone_to_so100/teleoperate.py +++ b/examples/phone_to_so100/teleoperate.py @@ -16,9 +16,8 @@ import time from lerobot.model.kinematics import RobotKinematics -from lerobot.processor import RobotProcessorPipeline +from lerobot.processor import RobotAction, RobotProcessorPipeline from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action -from lerobot.processor.core import RobotAction from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.robot_kinematic_processor import ( AddRobotObservationAsComplimentaryData, diff --git a/src/lerobot/policies/act/processor_act.py b/src/lerobot/policies/act/processor_act.py index 129d2997f..8888191c2 100644 --- a/src/lerobot/policies/act/processor_act.py +++ b/src/lerobot/policies/act/processor_act.py @@ -23,12 +23,12 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, + PolicyAction, PolicyProcessorPipeline, RenameObservationsProcessorStep, UnnormalizerProcessorStep, ) from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.processor.core import PolicyAction def make_act_pre_post_processors( diff --git a/src/lerobot/policies/diffusion/processor_diffusion.py b/src/lerobot/policies/diffusion/processor_diffusion.py index e9ada252b..f89e48aa7 100644 --- a/src/lerobot/policies/diffusion/processor_diffusion.py +++ b/src/lerobot/policies/diffusion/processor_diffusion.py @@ -24,12 +24,12 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, + PolicyAction, PolicyProcessorPipeline, RenameObservationsProcessorStep, UnnormalizerProcessorStep, ) from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.processor.core import PolicyAction def make_diffusion_pre_post_processors( diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 8d94f5837..06c0c4ba5 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -39,14 +39,13 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardCla from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig -from lerobot.processor import PolicyProcessorPipeline +from lerobot.processor import PolicyAction, PolicyProcessorPipeline from lerobot.processor.converters import ( batch_to_transition, policy_action_to_transition, transition_to_batch, transition_to_policy_action, ) -from lerobot.processor.core import PolicyAction def get_policy_class(name: str) -> type[PreTrainedPolicy]: diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index f6470d8cf..ebf2b03ed 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -26,6 +26,7 @@ from lerobot.processor import ( ComplementaryDataProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, + PolicyAction, PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, @@ -34,7 +35,6 @@ from lerobot.processor import ( UnnormalizerProcessorStep, ) from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.processor.core import PolicyAction @ProcessorStepRegistry.register(name="pi0_new_line_processor") diff --git a/src/lerobot/policies/pi0fast/processor_pi0fast.py b/src/lerobot/policies/pi0fast/processor_pi0fast.py index de4443413..b585589b3 100644 --- a/src/lerobot/policies/pi0fast/processor_pi0fast.py +++ b/src/lerobot/policies/pi0fast/processor_pi0fast.py @@ -24,12 +24,12 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, + PolicyAction, PolicyProcessorPipeline, RenameObservationsProcessorStep, UnnormalizerProcessorStep, ) from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.processor.core import PolicyAction def make_pi0fast_pre_post_processors( diff --git a/src/lerobot/policies/sac/processor_sac.py b/src/lerobot/policies/sac/processor_sac.py index c0cd8f751..e6c2bc139 100644 --- a/src/lerobot/policies/sac/processor_sac.py +++ b/src/lerobot/policies/sac/processor_sac.py @@ -25,12 +25,12 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, + PolicyAction, PolicyProcessorPipeline, RenameObservationsProcessorStep, UnnormalizerProcessorStep, ) from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.processor.core import PolicyAction def make_sac_pre_post_processors( diff --git a/src/lerobot/policies/sac/reward_model/processor_classifier.py b/src/lerobot/policies/sac/reward_model/processor_classifier.py index a0c6bd503..c2a34eab2 100644 --- a/src/lerobot/policies/sac/reward_model/processor_classifier.py +++ b/src/lerobot/policies/sac/reward_model/processor_classifier.py @@ -23,10 +23,10 @@ from lerobot.processor import ( DeviceProcessorStep, IdentityProcessorStep, NormalizerProcessorStep, + PolicyAction, PolicyProcessorPipeline, ) from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.processor.core import PolicyAction def make_classifier_processor( diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 90ac0fa9a..a5f132173 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -26,6 +26,7 @@ from lerobot.processor import ( ComplementaryDataProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, + PolicyAction, PolicyProcessorPipeline, ProcessorStepRegistry, RenameObservationsProcessorStep, @@ -33,7 +34,6 @@ from lerobot.processor import ( UnnormalizerProcessorStep, ) from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.processor.core import PolicyAction def make_smolvla_pre_post_processors( diff --git a/src/lerobot/policies/tdmpc/processor_tdmpc.py b/src/lerobot/policies/tdmpc/processor_tdmpc.py index 51f97bd6e..dd20b7843 100644 --- a/src/lerobot/policies/tdmpc/processor_tdmpc.py +++ b/src/lerobot/policies/tdmpc/processor_tdmpc.py @@ -24,12 +24,12 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, + PolicyAction, PolicyProcessorPipeline, RenameObservationsProcessorStep, UnnormalizerProcessorStep, ) from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.processor.core import PolicyAction def make_tdmpc_pre_post_processors( diff --git a/src/lerobot/policies/vqbet/processor_vqbet.py b/src/lerobot/policies/vqbet/processor_vqbet.py index bd9f7102a..db4e88233 100644 --- a/src/lerobot/policies/vqbet/processor_vqbet.py +++ b/src/lerobot/policies/vqbet/processor_vqbet.py @@ -25,12 +25,12 @@ from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, + PolicyAction, PolicyProcessorPipeline, RenameObservationsProcessorStep, UnnormalizerProcessorStep, ) from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -from lerobot.processor.core import PolicyAction def make_vqbet_pre_post_processors( diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index a4505900e..38d8ca514 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -68,6 +68,7 @@ __all__ = [ "create_transition", "DeviceProcessorStep", "DoneProcessorStep", + "EnvAction", "EnvTransition", "GripperPenaltyProcessorStep", "hotswap_stats", @@ -83,11 +84,13 @@ __all__ = [ "NormalizerProcessorStep", "Numpy2TorchActionProcessorStep", "ObservationProcessorStep", + "PolicyAction", "PolicyActionProcessorStep", "PolicyProcessorPipeline", "ProcessorKwargs", "ProcessorStep", "ProcessorStepRegistry", + "RobotAction", "RobotActionProcessorStep", "RenameObservationsProcessorStep", "RewardClassifierProcessorStep", diff --git a/src/lerobot/record.py b/src/lerobot/record.py index e82afc21c..71de5dd5f 100644 --- a/src/lerobot/record.py +++ b/src/lerobot/record.py @@ -80,7 +80,9 @@ from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.processor import ( EnvTransition, IdentityProcessorStep, + PolicyAction, PolicyProcessorPipeline, + RobotAction, RobotProcessorPipeline, TransitionKey, ) @@ -91,7 +93,6 @@ from lerobot.processor.converters import ( transition_to_dataset_frame, transition_to_robot_action, ) -from lerobot.processor.core import PolicyAction, RobotAction from lerobot.processor.rename_processor import rename_stats from lerobot.robots import ( # noqa: F401 Robot, diff --git a/src/lerobot/replay.py b/src/lerobot/replay.py index 3c184703e..fee54bc64 100644 --- a/src/lerobot/replay.py +++ b/src/lerobot/replay.py @@ -47,9 +47,8 @@ from pprint import pformat from lerobot.configs import parser from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.processor import IdentityProcessorStep, RobotProcessorPipeline +from lerobot.processor import IdentityProcessorStep, RobotAction, RobotProcessorPipeline from lerobot.processor.converters import robot_action_to_transition, transition_to_robot_action -from lerobot.processor.core import RobotAction from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, diff --git a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py index 7c0c1eed7..09796910c 100644 --- a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py +++ b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py @@ -27,10 +27,10 @@ from lerobot.processor import ( ObservationProcessorStep, ProcessorStep, ProcessorStepRegistry, + RobotAction, RobotActionProcessorStep, TransitionKey, ) -from lerobot.processor.core import RobotAction from lerobot.robots.robot import Robot from lerobot.utils.rotation import Rotation diff --git a/src/lerobot/scripts/eval.py b/src/lerobot/scripts/eval.py index 7a0bafc68..0ccd4e417 100644 --- a/src/lerobot/scripts/eval.py +++ b/src/lerobot/scripts/eval.py @@ -72,8 +72,7 @@ from lerobot.envs.factory import make_env from lerobot.envs.utils import add_envs_task, check_env_attributes_and_types, preprocess_observation from lerobot.policies.factory import make_policy, make_pre_post_processors from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor.core import PolicyAction -from lerobot.processor.pipeline import PolicyProcessorPipeline +from lerobot.processor import PolicyAction, PolicyProcessorPipeline from lerobot.utils.io_utils import write_video from lerobot.utils.random_utils import set_seed from lerobot.utils.utils import ( diff --git a/src/lerobot/teleoperate.py b/src/lerobot/teleoperate.py index e2a5e10f0..1f28e3794 100644 --- a/src/lerobot/teleoperate.py +++ b/src/lerobot/teleoperate.py @@ -62,14 +62,19 @@ import rerun as rr from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 from lerobot.configs import parser -from lerobot.processor import EnvTransition, IdentityProcessorStep, RobotProcessorPipeline, TransitionKey +from lerobot.processor import ( + EnvTransition, + IdentityProcessorStep, + RobotAction, + RobotProcessorPipeline, + TransitionKey, +) from lerobot.processor.converters import ( identity_transition, observation_to_transition, robot_action_to_transition, transition_to_robot_action, ) -from lerobot.processor.core import RobotAction from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, diff --git a/src/lerobot/utils/control_utils.py b/src/lerobot/utils/control_utils.py index febe7070e..47beb5746 100644 --- a/src/lerobot/utils/control_utils.py +++ b/src/lerobot/utils/control_utils.py @@ -32,8 +32,7 @@ from termcolor import colored from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.utils import DEFAULT_FEATURES from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor import PolicyProcessorPipeline -from lerobot.processor.core import PolicyAction +from lerobot.processor import PolicyAction, PolicyProcessorPipeline from lerobot.robots import Robot diff --git a/tests/artifacts/policies/save_policy_to_safetensors.py b/tests/artifacts/policies/save_policy_to_safetensors.py index d91c31dbd..b0ffa9a31 100644 --- a/tests/artifacts/policies/save_policy_to_safetensors.py +++ b/tests/artifacts/policies/save_policy_to_safetensors.py @@ -24,7 +24,6 @@ from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets.factory import make_dataset from lerobot.optim.factory import make_optimizer_and_scheduler from lerobot.policies.factory import make_policy, make_policy_config, make_pre_post_processors -from lerobot.processor import TransitionKey from lerobot.utils.random_utils import set_seed @@ -104,7 +103,7 @@ def get_policy_stats(ds_repo_id: str, policy_name: str, policy_kwargs: dict): actions = {} for i in range(actions_queue): unnormalized_action = policy.select_action(obs).contiguous() - action_robot = postprocessor({TransitionKey.ACTION: unnormalized_action}).get(TransitionKey.ACTION) + action_robot = postprocessor(unnormalized_action) actions[str(i)] = action_robot return output_dict, grad_stats, param_stats, actions