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

@@ -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 (

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 (

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,

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,

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(

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(

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]:

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

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(

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(

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(

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(

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(

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(

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",

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,

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,

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
 

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 (

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,

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
 
 

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