refactor(processors): unify import statements by consolidating pipeline imports into the main processor module (#1845)

examples/phone_to_so100/evaluate.py

@@ -21,11 +21,11 @@ from lerobot.datasets.utils import merge_features
 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 RobotProcessor
 from lerobot.processor.converters import (
     to_output_robot_action,
     to_transition_robot_observation,
 )
-from lerobot.processor.pipeline import RobotProcessor
 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

@@ -20,12 +20,12 @@ 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.model.kinematics import RobotKinematics
+from lerobot.processor import RobotProcessor
 from lerobot.processor.converters import (
     to_output_robot_action,
     to_transition_robot_observation,
     to_transition_teleop_action,
 )
-from lerobot.processor.pipeline import RobotProcessor
 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,8 +19,8 @@ import time
 
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.model.kinematics import RobotKinematics
+from lerobot.processor import RobotProcessor
 from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action
-from lerobot.processor.pipeline import RobotProcessor
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
     AddRobotObservationAsComplimentaryData,

src/lerobot/datasets/pipeline_features.py

@@ -17,7 +17,7 @@ from typing import Any
 
 from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
 from lerobot.datasets.utils import hw_to_dataset_features
-from lerobot.processor.pipeline import RobotProcessor
+from lerobot.processor import RobotProcessor
 
 
 def aggregate_pipeline_dataset_features(

src/lerobot/policies/factory.py

@@ -38,7 +38,7 @@ 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.pipeline import ProcessorKwargs, RobotProcessor
+from lerobot.processor import ProcessorKwargs, RobotProcessor
 
 
 def get_policy_class(name: str) -> type[PreTrainedPolicy]:

src/lerobot/policies/pi0/processor_pi0.py

@@ -21,20 +21,18 @@ from lerobot.configs.types import PolicyFeature
 from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.processor import (
+    ComplementaryDataProcessor,
     DeviceProcessor,
     NormalizerProcessor,
     ProcessorKwargs,
+    ProcessorStep,
+    ProcessorStepRegistry,
+    RenameProcessor,
     RobotProcessor,
     ToBatchProcessor,
     TokenizerProcessor,
     UnnormalizerProcessor,
 )
-from lerobot.processor.pipeline import (
-    ComplementaryDataProcessor,
-    ProcessorStep,
-    ProcessorStepRegistry,
-)
-from lerobot.processor.rename_processor import RenameProcessor
 
 
 @ProcessorStepRegistry.register(name="pi0_new_line_processor")

src/lerobot/policies/smolvla/processor_smolvla.py

@@ -20,19 +20,17 @@ from lerobot.configs.types import PolicyFeature
 from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.processor import (
+    ComplementaryDataProcessor,
     DeviceProcessor,
     NormalizerProcessor,
     ProcessorKwargs,
+    ProcessorStepRegistry,
     RenameProcessor,
     RobotProcessor,
     ToBatchProcessor,
     TokenizerProcessor,
     UnnormalizerProcessor,
 )
-from lerobot.processor.pipeline import (
-    ComplementaryDataProcessor,
-    ProcessorStepRegistry,
-)
 
 
 def make_smolvla_pre_post_processors(

src/lerobot/processor/__init__.py

@@ -40,6 +40,7 @@ from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor, hot
 from .observation_processor import VanillaObservationProcessor
 from .pipeline import (
     ActionProcessor,
+    ComplementaryDataProcessor,
     DoneProcessor,
     IdentityProcessor,
     InfoProcessor,
@@ -58,6 +59,7 @@ __all__ = [
     "ActionProcessor",
     "AddTeleopActionAsComplimentaryData",
     "AddTeleopEventsAsInfo",
+    "ComplementaryDataProcessor",
     "batch_to_transition",
     "create_transition",
     "DeviceProcessor",

src/lerobot/processor/batch_processor.py

@@ -17,10 +17,11 @@ from torch import Tensor
 
 from lerobot.configs.types import PolicyFeature
 from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
-from lerobot.processor.pipeline import (
+
+from .core import EnvTransition
+from .pipeline import (
     ActionProcessor,
     ComplementaryDataProcessor,
-    EnvTransition,
     ObservationProcessor,
     ProcessorStep,
     ProcessorStepRegistry,

src/lerobot/processor/delta_action_processor.py

@@ -19,7 +19,8 @@ from dataclasses import dataclass
 from torch import Tensor
 
 from lerobot.configs.types import FeatureType, PolicyFeature
-from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry
+
+from .pipeline import ActionProcessor, ProcessorStepRegistry
 
 
 @ProcessorStepRegistry.register("map_tensor_to_delta_action_dict")

src/lerobot/processor/device_processor.py

@@ -19,10 +19,11 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PolicyFeature
-from lerobot.processor.core import EnvTransition, TransitionKey
-from lerobot.processor.pipeline import ProcessorStep, ProcessorStepRegistry
 from lerobot.utils.utils import get_safe_torch_device
 
+from .core import EnvTransition, TransitionKey
+from .pipeline import ProcessorStep, ProcessorStepRegistry
+
 
 @ProcessorStepRegistry.register("device_processor")
 @dataclass

src/lerobot/processor/gym_action_processor.py

@@ -17,8 +17,9 @@ import numpy as np
 import torch
 
 from lerobot.configs.types import PolicyFeature
-from lerobot.processor.converters import to_tensor
-from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry
+
+from .converters import to_tensor
+from .pipeline import ActionProcessor, ProcessorStepRegistry
 
 
 @ProcessorStepRegistry.register("torch2numpy_action_processor")

src/lerobot/processor/hil_processor.py

@@ -9,18 +9,18 @@ import torchvision.transforms.functional as F  # noqa: N812
 
 from lerobot.configs.types import PolicyFeature
 from lerobot.constants import ACTION
-from lerobot.processor.pipeline import (
+from lerobot.teleoperators.teleoperator import Teleoperator
+from lerobot.teleoperators.utils import TeleopEvents
+
+from .core import EnvTransition, TransitionKey
+from .pipeline import (
     ComplementaryDataProcessor,
-    EnvTransition,
     InfoProcessor,
     ObservationProcessor,
     ProcessorStep,
     ProcessorStepRegistry,
-    TransitionKey,
     TruncatedProcessor,
 )
-from lerobot.teleoperators.teleoperator import Teleoperator
-from lerobot.teleoperators.utils import TeleopEvents
 
 GRIPPER_KEY = "gripper"
 DISCRETE_PENALTY_KEY = "discrete_penalty"

src/lerobot/processor/joint_observations_processor.py

@@ -4,12 +4,10 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PolicyFeature
-from lerobot.processor.pipeline import (
-    ObservationProcessor,
-    ProcessorStepRegistry,
-)
 from lerobot.robots import Robot
 
+from .pipeline import ObservationProcessor, ProcessorStepRegistry
+
 
 @dataclass
 @ProcessorStepRegistry.register("joint_velocity_processor")

src/lerobot/processor/migrate_policy_normalization.py

@@ -46,11 +46,12 @@ from huggingface_hub import hf_hub_download
 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.normalize_processor import NormalizerProcessor, UnnormalizerProcessor
-from lerobot.processor.pipeline import RobotProcessor
-from lerobot.processor.rename_processor import RenameProcessor
+
+from .batch_processor import ToBatchProcessor
+from .device_processor import DeviceProcessor
+from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor
+from .pipeline import RobotProcessor
+from .rename_processor import RenameProcessor
 
 # Policy type to class mapping
 POLICY_CLASSES = {

src/lerobot/processor/normalize_processor.py

@@ -9,14 +9,10 @@ from torch import Tensor
 
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.processor.converters import to_tensor
-from lerobot.processor.pipeline import (
-    EnvTransition,
-    ProcessorStep,
-    ProcessorStepRegistry,
-    RobotProcessor,
-    TransitionKey,
-)
+
+from .converters import to_tensor
+from .core import EnvTransition, TransitionKey
+from .pipeline import ProcessorStep, ProcessorStepRegistry, RobotProcessor
 
 
 @dataclass

src/lerobot/processor/observation_processor.py

@@ -22,7 +22,8 @@ from torch import Tensor
 
 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 .pipeline import ObservationProcessor, ProcessorStepRegistry
 
 
 @dataclass

src/lerobot/processor/rename_processor.py

@@ -18,10 +18,8 @@ from dataclasses import dataclass, field
 from typing import Any
 
 from lerobot.configs.types import PolicyFeature
-from lerobot.processor.pipeline import (
-    ObservationProcessor,
-    ProcessorStepRegistry,
-)
+
+from .pipeline import ObservationProcessor, ProcessorStepRegistry
 
 
 @dataclass

src/lerobot/processor/tokenizer_processor.py

@@ -11,14 +11,11 @@ import torch
 
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
-from lerobot.processor.pipeline import (
-    EnvTransition,
-    ObservationProcessor,
-    ProcessorStepRegistry,
-    TransitionKey,
-)
 from lerobot.utils.import_utils import _transformers_available
 
+from .core import EnvTransition, TransitionKey
+from .pipeline import ObservationProcessor, ProcessorStepRegistry
+
 if TYPE_CHECKING or _transformers_available:
     from transformers import AutoTokenizer
 else:

src/lerobot/record.py

@@ -76,14 +76,13 @@ from lerobot.datasets.utils import hw_to_dataset_features
 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 IdentityProcessor, RobotProcessor, TransitionKey
 from lerobot.processor.converters import (
     to_output_robot_action,
     to_transition_robot_observation,
     to_transition_teleop_action,
     transition_to_dataset_frame,
 )
-from lerobot.processor.pipeline import IdentityProcessor, TransitionKey
 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 RobotProcessor
+from lerobot.processor import IdentityProcessor, RobotProcessor
 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
     Robot,
     RobotConfig,

src/lerobot/robots/so100_follower/robot_kinematic_processor.py

@@ -22,7 +22,7 @@ from scipy.spatial.transform import Rotation
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.constants import ACTION, OBS_STATE
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor.pipeline import (
+from lerobot.processor import (
     ActionProcessor,
     ComplementaryDataProcessor,
     EnvTransition,

src/lerobot/scripts/rl/actor.py

@@ -62,7 +62,7 @@ from lerobot.configs import parser
 from lerobot.configs.train import TrainRLServerPipelineConfig
 from lerobot.policies.factory import make_policy
 from lerobot.policies.sac.modeling_sac import SACPolicy
-from lerobot.processor.pipeline import TransitionKey
+from lerobot.processor import TransitionKey
 from lerobot.robots import so100_follower  # noqa: F401
 from lerobot.scripts.rl.gym_manipulator import (
     create_transition,

src/lerobot/scripts/rl/gym_manipulator.py

@@ -32,6 +32,7 @@ from lerobot.processor import (
     AddTeleopActionAsComplimentaryData,
     AddTeleopEventsAsInfo,
     DeviceProcessor,
+    EnvTransition,
     GripperPenaltyProcessor,
     ImageCropResizeProcessor,
     InterventionActionProcessor,
@@ -45,10 +46,10 @@ from lerobot.processor import (
     TimeLimitProcessor,
     ToBatchProcessor,
     Torch2NumpyActionProcessor,
+    TransitionKey,
     VanillaObservationProcessor,
     create_transition,
 )
-from lerobot.processor.pipeline import EnvTransition, TransitionKey
 from lerobot.robots import (  # noqa: F401
     RobotConfig,
     make_robot_from_config,

src/lerobot/teleoperate.py

@@ -61,13 +61,12 @@ 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 RobotProcessor
+from lerobot.processor import IdentityProcessor, RobotProcessor
 from lerobot.processor.converters import (
     to_output_robot_action,
     to_transition_robot_observation,
     to_transition_teleop_action,
 )
-from lerobot.processor.pipeline import IdentityProcessor
 from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,

src/lerobot/teleoperators/phone/phone_processor.py

@@ -17,7 +17,7 @@
 from dataclasses import dataclass, field
 
 from lerobot.configs.types import FeatureType, PolicyFeature
-from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry
+from lerobot.processor import ActionProcessor, ProcessorStepRegistry
 from lerobot.teleoperators.phone.config_phone import PhoneOS
 
 

src/lerobot/utils/train_utils.py

@@ -32,7 +32,7 @@ from lerobot.datasets.utils import load_json, write_json
 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 import RobotProcessor
 from lerobot.utils.random_utils import load_rng_state, save_rng_state
 
 

src/lerobot/utils/visualization_utils.py

@@ -19,7 +19,7 @@ from typing import Any
 import numpy as np
 import rerun as rr
 
-from lerobot.processor.pipeline import EnvTransition, TransitionKey
+from lerobot.processor import EnvTransition, TransitionKey
 
 
 def _init_rerun(session_name: str = "lerobot_control_loop") -> None:

tests/processor/test_act_processor.py

@@ -30,9 +30,9 @@ from lerobot.processor import (
     RenameProcessor,
     RobotProcessor,
     ToBatchProcessor,
+    TransitionKey,
     UnnormalizerProcessor,
 )
-from lerobot.processor.pipeline import TransitionKey
 
 
 def create_transition(observation=None, action=None, **kwargs):

tests/processor/test_batch_conversion.py

@@ -1,7 +1,7 @@
 import torch
 
+from lerobot.processor import RobotProcessor, TransitionKey
 from lerobot.processor.converters import batch_to_transition, transition_to_batch
-from lerobot.processor.pipeline import RobotProcessor, TransitionKey
 
 
 def _dummy_batch():

tests/processor/test_batch_processor.py

@@ -22,9 +22,7 @@ import pytest
 import torch
 
 from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
-from lerobot.processor import ProcessorStepRegistry, RobotProcessor
-from lerobot.processor.batch_processor import ToBatchProcessor
-from lerobot.processor.pipeline import TransitionKey
+from lerobot.processor import ProcessorStepRegistry, RobotProcessor, ToBatchProcessor, TransitionKey
 
 
 def create_transition(

tests/processor/test_classifier_processor.py

@@ -24,8 +24,13 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 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.pipeline import TransitionKey
+from lerobot.processor import (
+    DeviceProcessor,
+    IdentityProcessor,
+    NormalizerProcessor,
+    RobotProcessor,
+    TransitionKey,
+)
 
 
 def create_transition(observation=None, action=None, **kwargs):

tests/processor/test_converters.py

@@ -2,6 +2,7 @@ import numpy as np
 import pytest
 import torch
 
+from lerobot.processor import TransitionKey
 from lerobot.processor.converters import (
     batch_to_transition,
     to_output_robot_action,
@@ -11,7 +12,6 @@ from lerobot.processor.converters import (
     transition_to_batch,
     transition_to_dataset_frame,
 )
-from lerobot.processor.pipeline import TransitionKey
 
 
 def test_to_transition_teleop_action_prefix_and_tensor_conversion():

tests/processor/test_device_processor.py

@@ -19,8 +19,7 @@ import pytest
 import torch
 
 from lerobot.configs.types import FeatureType, PolicyFeature
-from lerobot.processor import DeviceProcessor, RobotProcessor
-from lerobot.processor.pipeline import TransitionKey
+from lerobot.processor import DeviceProcessor, RobotProcessor, TransitionKey
 
 
 def create_transition(
@@ -358,7 +357,7 @@ def test_save_and_load_pretrained():
 
 def test_registry_functionality():
     """Test that DeviceProcessor is properly registered."""
-    from lerobot.processor.pipeline import ProcessorStepRegistry
+    from lerobot.processor import ProcessorStepRegistry
 
     # Check that DeviceProcessor is registered
     registered_class = ProcessorStepRegistry.get("device_processor")

tests/processor/test_diffusion_processor.py

@@ -30,9 +30,9 @@ from lerobot.processor import (
     RenameProcessor,
     RobotProcessor,
     ToBatchProcessor,
+    TransitionKey,
     UnnormalizerProcessor,
 )
-from lerobot.processor.pipeline import TransitionKey
 
 
 def create_transition(observation=None, action=None, **kwargs):

tests/processor/test_normalize_processor.py

@@ -20,13 +20,15 @@ import pytest
 import torch
 
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
-from lerobot.processor.converters import to_tensor
-from lerobot.processor.normalize_processor import (
+from lerobot.processor import (
+    IdentityProcessor,
     NormalizerProcessor,
+    RobotProcessor,
+    TransitionKey,
     UnnormalizerProcessor,
     hotswap_stats,
 )
-from lerobot.processor.pipeline import IdentityProcessor, RobotProcessor, TransitionKey
+from lerobot.processor.converters import to_tensor
 
 
 def create_transition(

tests/processor/test_observation_processor.py

@@ -20,8 +20,7 @@ import torch
 
 from lerobot.configs.types import FeatureType
 from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
-from lerobot.processor import VanillaObservationProcessor
-from lerobot.processor.pipeline import TransitionKey
+from lerobot.processor import TransitionKey, VanillaObservationProcessor
 from tests.conftest import assert_contract_is_typed
 
 

tests/processor/test_pi0_processor.py

@@ -29,9 +29,9 @@ from lerobot.processor import (
     NormalizerProcessor,
     RenameProcessor,
     ToBatchProcessor,
+    TransitionKey,
     UnnormalizerProcessor,
 )
-from lerobot.processor.pipeline import TransitionKey
 
 
 def create_transition(observation=None, action=None, **kwargs):

tests/processor/test_pipeline.py

@@ -27,8 +27,7 @@ import torch.nn as nn
 
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features
-from lerobot.processor import EnvTransition, ProcessorStepRegistry, RobotProcessor
-from lerobot.processor.pipeline import TransitionKey
+from lerobot.processor import EnvTransition, ProcessorStepRegistry, RobotProcessor, TransitionKey
 from tests.conftest import assert_contract_is_typed
 
 

tests/processor/test_rename_processor.py

@@ -20,7 +20,12 @@ import numpy as np
 import torch
 
 from lerobot.configs.types import FeatureType
-from lerobot.processor import ProcessorStepRegistry, RenameProcessor, RobotProcessor, TransitionKey
+from lerobot.processor import (
+    ProcessorStepRegistry,
+    RenameProcessor,
+    RobotProcessor,
+    TransitionKey,
+)
 from lerobot.processor.rename_processor import rename_stats
 from tests.conftest import assert_contract_is_typed
 

tests/processor/test_sac_processor.py

@@ -30,9 +30,9 @@ from lerobot.processor import (
     RenameProcessor,
     RobotProcessor,
     ToBatchProcessor,
+    TransitionKey,
     UnnormalizerProcessor,
 )
-from lerobot.processor.pipeline import TransitionKey
 
 
 def create_transition(observation=None, action=None, **kwargs):

tests/processor/test_smolvla_processor.py

@@ -32,9 +32,9 @@ from lerobot.processor import (
     NormalizerProcessor,
     RenameProcessor,
     ToBatchProcessor,
+    TransitionKey,
     UnnormalizerProcessor,
 )
-from lerobot.processor.pipeline import TransitionKey
 
 
 def create_transition(observation=None, action=None, **kwargs):

tests/processor/test_tdmpc_processor.py

@@ -30,9 +30,9 @@ from lerobot.processor import (
     RenameProcessor,
     RobotProcessor,
     ToBatchProcessor,
+    TransitionKey,
     UnnormalizerProcessor,
 )
-from lerobot.processor.pipeline import TransitionKey
 
 
 def create_transition(observation=None, action=None, **kwargs):

tests/processor/test_tokenizer_processor.py

@@ -10,8 +10,7 @@ import torch
 
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.constants import OBS_LANGUAGE
-from lerobot.processor.pipeline import RobotProcessor, TransitionKey
-from lerobot.processor.tokenizer_processor import TokenizerProcessor
+from lerobot.processor import RobotProcessor, TokenizerProcessor, TransitionKey
 from tests.utils import require_package
 
 
@@ -488,7 +487,7 @@ def test_save_and_load_pretrained_with_tokenizer_object():
 @require_package("transformers")
 def test_registry_functionality():
     """Test that the processor is properly registered."""
-    from lerobot.processor.pipeline import ProcessorStepRegistry
+    from lerobot.processor import ProcessorStepRegistry
 
     # Check that the processor is registered
     assert "tokenizer_processor" in ProcessorStepRegistry.list()

tests/processor/test_vqbet_processor.py

@@ -30,9 +30,9 @@ from lerobot.processor import (
     RenameProcessor,
     RobotProcessor,
     ToBatchProcessor,
+    TransitionKey,
     UnnormalizerProcessor,
 )
-from lerobot.processor.pipeline import TransitionKey
 
 
 def create_transition(observation=None, action=None, **kwargs):

tests/utils/test_visualization_utils.py

@@ -5,7 +5,7 @@ from types import SimpleNamespace
 import numpy as np
 import pytest
 
-from lerobot.processor.pipeline import TransitionKey
+from lerobot.processor import TransitionKey
 
 
 @pytest.fixture