chore(dependecies): untangle dependecies across internal modules (#3149)

examples/phone_to_so100/evaluate.py

@@ -23,8 +23,6 @@ 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 (
-    RobotAction,
-    RobotObservation,
     RobotProcessorPipeline,
     make_default_teleop_action_processor,
 )
@@ -40,6 +38,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.scripts.lerobot_record import record_loop
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
 from lerobot.utils.visualization_utils import init_rerun

examples/phone_to_so100/record.py

@@ -19,7 +19,7 @@ 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 RobotAction, RobotObservation, RobotProcessorPipeline
+from lerobot.processor import RobotProcessorPipeline
 from lerobot.processor.converters import (
     observation_to_transition,
     robot_action_observation_to_transition,
@@ -38,6 +38,7 @@ from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
 from lerobot.teleoperators.phone.teleop_phone import Phone
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
 from lerobot.utils.visualization_utils import init_rerun

examples/phone_to_so100/replay.py

@@ -18,7 +18,7 @@ import time
 
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
+from lerobot.processor import RobotProcessorPipeline
 from lerobot.processor.converters import (
     robot_action_observation_to_transition,
     transition_to_robot_action,
@@ -27,6 +27,7 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
 from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say

examples/phone_to_so100/teleoperate.py

@@ -16,7 +16,7 @@
 import time
 
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
+from lerobot.processor import RobotProcessorPipeline
 from lerobot.processor.converters import (
     robot_action_observation_to_transition,
     transition_to_robot_action,
@@ -31,6 +31,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
 from lerobot.teleoperators.phone.teleop_phone import Phone
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 

examples/so100_to_so100_EE/evaluate.py

@@ -23,8 +23,6 @@ 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 (
-    RobotAction,
-    RobotObservation,
     RobotProcessorPipeline,
     make_default_teleop_action_processor,
 )
@@ -40,6 +38,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.scripts.lerobot_record import record_loop
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
 from lerobot.utils.visualization_utils import init_rerun

examples/so100_to_so100_EE/record.py

@@ -20,7 +20,7 @@ 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 RobotAction, RobotObservation, RobotProcessorPipeline
+from lerobot.processor import RobotProcessorPipeline
 from lerobot.processor.converters import (
     observation_to_transition,
     robot_action_observation_to_transition,
@@ -35,6 +35,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 )
 from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
 from lerobot.utils.visualization_utils import init_rerun

examples/so100_to_so100_EE/replay.py

@@ -19,7 +19,7 @@ import time
 
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
+from lerobot.processor import RobotProcessorPipeline
 from lerobot.processor.converters import (
     robot_action_observation_to_transition,
     transition_to_robot_action,
@@ -28,6 +28,7 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
 from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say

examples/so100_to_so100_EE/teleoperate.py

@@ -17,7 +17,7 @@
 import time
 
 from lerobot.model.kinematics import RobotKinematics
-from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
+from lerobot.processor import RobotProcessorPipeline
 from lerobot.processor.converters import (
     robot_action_observation_to_transition,
     robot_action_to_transition,
@@ -30,6 +30,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 

src/lerobot/async_inference/policy_server.py

@@ -39,15 +39,13 @@ import grpc
 import torch
 
 from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.processor import (
-    PolicyAction,
-    PolicyProcessorPipeline,
-)
+from lerobot.processor import PolicyProcessorPipeline
 from lerobot.transport import (
     services_pb2,  # type: ignore
     services_pb2_grpc,  # type: ignore
 )
 from lerobot.transport.utils import receive_bytes_in_chunks
+from lerobot.types import PolicyAction
 
 from .configs import PolicyServerConfig
 from .constants import SUPPORTED_POLICIES

src/lerobot/configs/policies.py

@@ -30,8 +30,8 @@ from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.optim.optimizers import OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
 from lerobot.utils.constants import ACTION, OBS_STATE
+from lerobot.utils.device_utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 from lerobot.utils.hub import HubMixin
-from lerobot.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 
 T = TypeVar("T", bound="PreTrainedConfig")
 logger = getLogger(__name__)

src/lerobot/datasets/pipeline_features.py

@@ -18,7 +18,8 @@ from typing import Any
 
 from lerobot.configs.types import PipelineFeatureType
 from lerobot.datasets.utils import hw_to_dataset_features
-from lerobot.processor import DataProcessorPipeline, RobotAction, RobotObservation
+from lerobot.processor import DataProcessorPipeline
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
 
 

src/lerobot/envs/libero.py

@@ -29,7 +29,7 @@ from gymnasium import spaces
 from libero.libero import benchmark, get_libero_path
 from libero.libero.envs import OffScreenRenderEnv
 
-from lerobot.processor import RobotObservation
+from lerobot.types import RobotObservation
 
 
 def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:

src/lerobot/envs/metaworld.py

@@ -25,7 +25,7 @@ import metaworld.policies as policies
 import numpy as np
 from gymnasium import spaces
 
-from lerobot.processor import RobotObservation
+from lerobot.types import RobotObservation
 
 # ---- Load configuration data from the external JSON file ----
 CONFIG_PATH = Path(__file__).parent / "metaworld_config.json"

src/lerobot/envs/utils.py

@@ -29,7 +29,7 @@ from torch import Tensor
 
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.envs.configs import EnvConfig
-from lerobot.processor import RobotObservation
+from lerobot.types import RobotObservation
 from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR
 from lerobot.utils.utils import get_channel_first_image_shape
 

src/lerobot/policies/factory.py

@@ -43,13 +43,14 @@ from lerobot.policies.utils import validate_visual_features_consistency
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.policies.wall_x.configuration_wall_x import WallXConfig
 from lerobot.policies.xvla.configuration_xvla import XVLAConfig
-from lerobot.processor import PolicyAction, PolicyProcessorPipeline
+from lerobot.processor import PolicyProcessorPipeline
 from lerobot.processor.converters import (
     batch_to_transition,
     policy_action_to_transition,
     transition_to_batch,
     transition_to_policy_action,
 )
+from lerobot.types import PolicyAction
 from lerobot.utils.constants import (
     ACTION,
     POLICY_POSTPROCESSOR_DEFAULT_NAME,

src/lerobot/policies/groot/processor_groot.py

@@ -49,7 +49,7 @@ from lerobot.processor.converters import (
     policy_action_to_transition,
     transition_to_policy_action,
 )
-from lerobot.processor.core import EnvTransition, TransitionKey
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     ACTION,
     HF_LEROBOT_HOME,

src/lerobot/policies/pi05/processor_pi05.py

@@ -36,7 +36,7 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.processor.core import EnvTransition, TransitionKey
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     OBS_STATE,
     POLICY_POSTPROCESSOR_DEFAULT_NAME,

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -37,7 +37,7 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.processor.core import EnvTransition, TransitionKey
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     OBS_STATE,
     POLICY_POSTPROCESSOR_DEFAULT_NAME,

src/lerobot/policies/sarm/processor_sarm.py

@@ -48,8 +48,8 @@ from lerobot.processor.converters import (
     policy_action_to_transition,
     transition_to_policy_action,
 )
-from lerobot.processor.core import EnvTransition, TransitionKey
 from lerobot.processor.pipeline import PipelineFeatureType
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -68,7 +68,7 @@ from lerobot.policies.utils import (
     populate_queues,
 )
 from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
-from lerobot.utils.utils import get_safe_dtype
+from lerobot.utils.device_utils import get_safe_dtype
 
 
 class ActionSelectKwargs(TypedDict, total=False):

src/lerobot/policies/utils.py

@@ -24,7 +24,7 @@ from torch import nn
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.datasets.utils import build_dataset_frame
-from lerobot.processor import PolicyAction, RobotAction, RobotObservation
+from lerobot.types import PolicyAction, RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
 
 

src/lerobot/policies/xvla/processor_xvla.py

@@ -38,7 +38,7 @@ from lerobot.processor import (
     UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.processor.core import EnvTransition, TransitionKey
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     OBS_IMAGES,
     OBS_PREFIX,

src/lerobot/processor/__init__.py

@@ -14,13 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .batch_processor import AddBatchDimensionProcessorStep
-from .converters import (
-    batch_to_transition,
-    create_transition,
-    transition_to_batch,
-)
-from .core import (
+from lerobot.types import (
     EnvAction,
     EnvTransition,
     PolicyAction,
@@ -28,6 +22,13 @@ from .core import (
     RobotObservation,
     TransitionKey,
 )
+
+from .batch_processor import AddBatchDimensionProcessorStep
+from .converters import (
+    batch_to_transition,
+    create_transition,
+    transition_to_batch,
+)
 from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
 from .device_processor import DeviceProcessorStep
 from .factory import (

src/lerobot/processor/batch_processor.py

@@ -25,9 +25,9 @@ from dataclasses import dataclass, field
 from torch import Tensor
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.types import EnvTransition, PolicyAction
 from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 
-from .core import EnvTransition, PolicyAction
 from .pipeline import (
     ComplementaryDataProcessorStep,
     ObservationProcessorStep,

src/lerobot/processor/converters.py

@@ -23,10 +23,9 @@ from typing import Any
 import numpy as np
 import torch
 
+from lerobot.types import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
 from lerobot.utils.constants import ACTION, DONE, INFO, OBS_PREFIX, REWARD, TRUNCATED
 
-from .core import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
-
 
 @singledispatch
 def to_tensor(

src/lerobot/processor/delta_action_processor.py

@@ -17,8 +17,8 @@
 from dataclasses import dataclass
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.types import PolicyAction, RobotAction
 
-from .core import PolicyAction, RobotAction
 from .pipeline import ActionProcessorStep, ProcessorStepRegistry, RobotActionProcessorStep
 
 

src/lerobot/processor/device_processor.py

@@ -25,9 +25,9 @@ from typing import Any
 import torch
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.utils.utils import get_safe_torch_device
+from lerobot.types import EnvTransition, PolicyAction, TransitionKey
+from lerobot.utils.device_utils import get_safe_torch_device
 
-from .core import EnvTransition, PolicyAction, TransitionKey
 from .pipeline import ProcessorStep, ProcessorStepRegistry
 
 

src/lerobot/processor/factory.py

@@ -14,13 +14,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from lerobot.types import RobotAction, RobotObservation
+
 from .converters import (
     observation_to_transition,
     robot_action_observation_to_transition,
     transition_to_observation,
     transition_to_robot_action,
 )
-from .core import RobotAction, RobotObservation
 from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
 
 

src/lerobot/processor/gym_action_processor.py

@@ -17,9 +17,9 @@
 from dataclasses import dataclass
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.types import EnvAction, EnvTransition, PolicyAction
 
 from .converters import to_tensor
-from .core import EnvAction, EnvTransition, PolicyAction
 from .hil_processor import TELEOP_ACTION_KEY
 from .pipeline import ActionProcessorStep, ProcessorStep, ProcessorStepRegistry
 
@@ -75,7 +75,7 @@ class Numpy2TorchActionProcessorStep(ProcessorStep):
 
     def __call__(self, transition: EnvTransition) -> EnvTransition:
         """Converts numpy action to torch tensor if action exists, otherwise passes through."""
-        from .core import TransitionKey
+        from lerobot.types import TransitionKey
 
         self._current_transition = transition.copy()
         new_transition = self._current_transition

src/lerobot/processor/hil_processor.py

@@ -30,7 +30,8 @@ from lerobot.teleoperators.utils import TeleopEvents
 if TYPE_CHECKING:
     from lerobot.teleoperators.teleoperator import Teleoperator
 
-from .core import EnvTransition, PolicyAction, TransitionKey
+from lerobot.types import EnvTransition, PolicyAction, TransitionKey
+
 from .pipeline import (
     ComplementaryDataProcessorStep,
     InfoProcessorStep,

src/lerobot/processor/normalize_processor.py

@@ -26,10 +26,10 @@ from torch import Tensor
 
 from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.types import EnvTransition, PolicyAction, TransitionKey
 from lerobot.utils.constants import ACTION
 
 from .converters import from_tensor_to_numpy, to_tensor
-from .core import EnvTransition, PolicyAction, TransitionKey
 from .pipeline import PolicyProcessorPipeline, ProcessorStep, ProcessorStepRegistry, RobotObservation
 
 

src/lerobot/processor/pipeline.py

@@ -46,10 +46,10 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file, save_file
 
 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.types import EnvAction, EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
 from lerobot.utils.hub import HubMixin
 
 from .converters import batch_to_transition, create_transition, transition_to_batch
-from .core import EnvAction, EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
 
 # Generic type variables for pipeline input and output.
 TInput = TypeVar("TInput")

src/lerobot/processor/tokenizer_processor.py

@@ -30,6 +30,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
+from lerobot.types import EnvTransition, RobotObservation, TransitionKey
 from lerobot.utils.constants import (
     ACTION_TOKEN_MASK,
     ACTION_TOKENS,
@@ -40,7 +41,6 @@ from lerobot.utils.constants import (
 )
 from lerobot.utils.import_utils import _transformers_available
 
-from .core import EnvTransition, RobotObservation, TransitionKey
 from .pipeline import ActionProcessorStep, ObservationProcessorStep, ProcessorStepRegistry
 
 # Conditional import for type checking and lazy loading

src/lerobot/rl/actor.py

@@ -62,7 +62,6 @@ 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 import TransitionKey
 from lerobot.rl.process import ProcessSignalHandler
 from lerobot.rl.queue import get_last_item_from_queue
 from lerobot.robots import so_follower  # noqa: F401
@@ -77,6 +76,8 @@ from lerobot.transport.utils import (
     send_bytes_in_chunks,
     transitions_to_bytes,
 )
+from lerobot.types import TransitionKey
+from lerobot.utils.device_utils import get_safe_torch_device
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.transition import (
@@ -86,7 +87,6 @@ from lerobot.utils.transition import (
 )
 from lerobot.utils.utils import (
     TimerManager,
-    get_safe_torch_device,
     init_logging,
 )
 

src/lerobot/rl/learner.py

@@ -86,6 +86,7 @@ from lerobot.utils.constants import (
     PRETRAINED_MODEL_DIR,
     TRAINING_STATE_DIR,
 )
+from lerobot.utils.device_utils import get_safe_torch_device
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.train_utils import (
     get_step_checkpoint_dir,
@@ -96,7 +97,6 @@ from lerobot.utils.train_utils import (
 from lerobot.utils.transition import move_state_dict_to_device, move_transition_to_device
 from lerobot.utils.utils import (
     format_big_number,
-    get_safe_torch_device,
     init_logging,
 )
 

src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py

@@ -17,8 +17,8 @@
 import logging
 from functools import cached_property
 
-from lerobot.processor import RobotAction, RobotObservation
 from lerobot.robots.openarm_follower import OpenArmFollower, OpenArmFollowerConfig
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/bi_so_follower/bi_so_follower.py

@@ -17,8 +17,8 @@
 import logging
 from functools import cached_property
 
-from lerobot.processor import RobotAction, RobotObservation
 from lerobot.robots.so_follower import SOFollower, SOFollowerRobotConfig
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py

@@ -23,7 +23,7 @@ import cv2
 import numpy as np
 import requests
 
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 from lerobot.utils.errors import DeviceNotConnectedError
 

src/lerobot/robots/hope_jr/hope_jr_arm.py

@@ -24,7 +24,7 @@ from lerobot.motors.calibration_gui import RangeFinderGUI
 from lerobot.motors.feetech import (
     FeetechMotorsBus,
 )
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/hope_jr/hope_jr_hand.py

@@ -24,7 +24,7 @@ from lerobot.motors.calibration_gui import RangeFinderGUI
 from lerobot.motors.feetech import (
     FeetechMotorsBus,
 )
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/koch_follower/koch_follower.py

@@ -24,7 +24,7 @@ from lerobot.motors.dynamixel import (
     DynamixelMotorsBus,
     OperatingMode,
 )
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/lekiwi/lekiwi.py

@@ -28,7 +28,7 @@ from lerobot.motors.feetech import (
     FeetechMotorsBus,
     OperatingMode,
 )
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/lekiwi/lekiwi_client.py

@@ -22,7 +22,7 @@ from functools import cached_property
 import cv2
 import numpy as np
 
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STATE
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 from lerobot.utils.errors import DeviceNotConnectedError

src/lerobot/robots/omx_follower/omx_follower.py

@@ -25,7 +25,7 @@ from lerobot.motors.dynamixel import (
     DynamixelMotorsBus,
     OperatingMode,
 )
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/openarm_follower/openarm_follower.py

@@ -22,7 +22,7 @@ from typing import Any
 from lerobot.cameras.utils import make_cameras_from_configs
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.damiao import DamiaoMotorsBus
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/reachy2/robot_reachy2.py

@@ -19,7 +19,7 @@ import time
 from typing import TYPE_CHECKING, Any
 
 from lerobot.cameras.utils import make_cameras_from_configs
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.import_utils import _reachy2_sdk_available
 
 from ..robot import Robot

src/lerobot/robots/robot.py

@@ -19,7 +19,7 @@ from pathlib import Path
 import draccus
 
 from lerobot.motors import MotorCalibration
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, ROBOTS
 
 from .config import RobotConfig

src/lerobot/robots/so_follower/so_follower.py

@@ -24,7 +24,7 @@ from lerobot.motors.feetech import (
     FeetechMotorsBus,
     OperatingMode,
 )
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..robot import Robot

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -26,8 +26,6 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable
 import numpy as np
 
 from lerobot.cameras.utils import make_cameras_from_configs
-from lerobot.envs.factory import make_env
-from lerobot.processor import RobotAction, RobotObservation
 from lerobot.robots.unitree_g1.g1_kinematics import G1_29_ArmIK
 from lerobot.robots.unitree_g1.g1_utils import (
     REMOTE_AXES,
@@ -37,6 +35,7 @@ from lerobot.robots.unitree_g1.g1_utils import (
     default_remote_input,
     make_locomotion_controller,
 )
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.import_utils import _unitree_sdk_available
 
 from ..robot import Robot
@@ -291,6 +290,8 @@ class UnitreeG1(Robot):
     def connect(self, calibrate: bool = True) -> None:  # connect to DDS
         # Initialize DDS channel and simulation environment
         if self.config.is_simulation:
+            from lerobot.envs.factory import make_env
+
             self._ChannelFactoryInitialize(0, "lo")
             self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
             # Extract the actual gym env from the dict structure

src/lerobot/scripts/lerobot_eval.py

@@ -80,13 +80,14 @@ from lerobot.envs.utils import (
 )
 from lerobot.policies.factory import make_policy, make_pre_post_processors
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.processor import PolicyAction, PolicyProcessorPipeline
+from lerobot.processor import PolicyProcessorPipeline
+from lerobot.types import PolicyAction
 from lerobot.utils.constants import ACTION, DONE, OBS_STR, REWARD
+from lerobot.utils.device_utils import get_safe_torch_device
 from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.io_utils import write_video
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.utils import (
-    get_safe_torch_device,
     init_logging,
     inside_slurm,
 )

src/lerobot/scripts/lerobot_record.py

@@ -139,10 +139,10 @@ from lerobot.utils.control_utils import (
     sanity_check_dataset_name,
     sanity_check_dataset_robot_compatibility,
 )
+from lerobot.utils.device_utils import get_safe_torch_device
 from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import (
-    get_safe_torch_device,
     init_logging,
     log_say,
 )

src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py

@@ -17,8 +17,8 @@
 import logging
 from functools import cached_property
 
-from lerobot.processor import RobotAction
 from lerobot.teleoperators.openarm_leader import OpenArmLeaderConfig
+from lerobot.types import RobotAction
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..openarm_leader import OpenArmLeader

src/lerobot/teleoperators/gamepad/teleop_gamepad.py

@@ -20,7 +20,7 @@ from typing import Any
 
 import numpy as np
 
-from lerobot.processor import RobotAction
+from lerobot.types import RobotAction
 from lerobot.utils.decorators import check_if_not_connected
 
 from ..teleoperator import Teleoperator

src/lerobot/teleoperators/keyboard/teleop_keyboard.py

@@ -21,7 +21,7 @@ import time
 from queue import Queue
 from typing import Any
 
-from lerobot.processor import RobotAction
+from lerobot.types import RobotAction
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..teleoperator import Teleoperator

src/lerobot/teleoperators/openarm_leader/openarm_leader.py

@@ -20,7 +20,7 @@ from typing import Any
 
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.damiao import DamiaoMotorsBus
-from lerobot.processor import RobotAction
+from lerobot.types import RobotAction
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..teleoperator import Teleoperator

src/lerobot/teleoperators/openarm_mini/openarm_mini.py

@@ -23,7 +23,7 @@ from lerobot.motors.feetech import (
     FeetechMotorsBus,
     OperatingMode,
 )
-from lerobot.processor import RobotAction
+from lerobot.types import RobotAction
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 from ..teleoperator import Teleoperator

src/lerobot/teleoperators/phone/phone_processor.py

@@ -17,8 +17,9 @@
 from dataclasses import dataclass, field
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.processor import ProcessorStepRegistry, RobotAction, RobotActionProcessorStep
+from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep
 from lerobot.teleoperators.phone.config_phone import PhoneOS
+from lerobot.types import RobotAction
 
 
 @ProcessorStepRegistry.register("map_phone_action_to_robot_action")

src/lerobot/teleoperators/teleoperator.py

@@ -20,7 +20,7 @@ from typing import Any
 import draccus
 
 from lerobot.motors.motors_bus import MotorCalibration
-from lerobot.processor import RobotAction
+from lerobot.types import RobotAction
 from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, TELEOPERATORS
 
 from .config import TeleoperatorConfig

src/lerobot/utils/control_utils.py

@@ -32,8 +32,9 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import DEFAULT_FEATURES
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import prepare_observation_for_inference
-from lerobot.processor import PolicyAction, PolicyProcessorPipeline
+from lerobot.processor import PolicyProcessorPipeline
 from lerobot.robots import Robot
+from lerobot.types import PolicyAction
 
 
 @cache

src/lerobot/utils/device_utils.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import torch
+
+
+def auto_select_torch_device() -> torch.device:
+    """Tries to select automatically a torch device."""
+    if torch.cuda.is_available():
+        logging.info("Cuda backend detected, using cuda.")
+        return torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        logging.info("Metal backend detected, using mps.")
+        return torch.device("mps")
+    elif torch.xpu.is_available():
+        logging.info("Intel XPU backend detected, using xpu.")
+        return torch.device("xpu")
+    else:
+        logging.warning("No accelerated backend detected. Using default cpu, this will be slow.")
+        return torch.device("cpu")
+
+
+# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
+def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
+    """Given a string, return a torch.device with checks on whether the device is available."""
+    try_device = str(try_device)
+    if try_device.startswith("cuda"):
+        assert torch.cuda.is_available()
+        device = torch.device(try_device)
+    elif try_device == "mps":
+        assert torch.backends.mps.is_available()
+        device = torch.device("mps")
+    elif try_device == "xpu":
+        assert torch.xpu.is_available()
+        device = torch.device("xpu")
+    elif try_device == "cpu":
+        device = torch.device("cpu")
+        if log:
+            logging.warning("Using CPU, this will be slow.")
+    else:
+        device = torch.device(try_device)
+        if log:
+            logging.warning(f"Using custom {try_device} device.")
+    return device
+
+
+def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
+    """
+    mps is currently not compatible with float64
+    """
+    if isinstance(device, torch.device):
+        device = device.type
+    if device == "mps" and dtype == torch.float64:
+        return torch.float32
+    if device == "xpu" and dtype == torch.float64:
+        if hasattr(torch.xpu, "get_device_capability"):
+            device_capability = torch.xpu.get_device_capability()
+            # NOTE: Some Intel XPU devices do not support double precision (FP64).
+            # The `has_fp64` flag is returned by `torch.xpu.get_device_capability()`
+            # when available; if False, we fall back to float32 for compatibility.
+            if not device_capability.get("has_fp64", False):
+                logging.warning(f"Device {device} does not support float64, using float32 instead.")
+                return torch.float32
+        else:
+            logging.warning(
+                f"Device {device} capability check failed. Assuming no support for float64, using float32 instead."
+            )
+            return torch.float32
+        return dtype
+    else:
+        return dtype
+
+
+def is_torch_device_available(try_device: str) -> bool:
+    try_device = str(try_device)  # Ensure try_device is a string
+    if try_device.startswith("cuda"):
+        return torch.cuda.is_available()
+    elif try_device == "mps":
+        return torch.backends.mps.is_available()
+    elif try_device == "xpu":
+        return torch.xpu.is_available()
+    elif try_device == "cpu":
+        return True
+    else:
+        raise ValueError(f"Unknown device {try_device}. Supported devices are: cuda, mps, xpu or cpu.")
+
+
+def is_amp_available(device: str):
+    if device in ["cuda", "xpu", "cpu"]:
+        return True
+    elif device == "mps":
+        return False
+    else:
+        raise ValueError(f"Unknown device '{device}.")

src/lerobot/utils/utils.py

@@ -13,6 +13,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
 import logging
 import os
 import platform
@@ -24,11 +26,12 @@ from copy import copy, deepcopy
 from datetime import datetime
 from pathlib import Path
 from statistics import mean
+from typing import TYPE_CHECKING
 
 import numpy as np
-import torch
+
+if TYPE_CHECKING:
     from accelerate import Accelerator
-from datasets.utils.logging import disable_progress_bar, enable_progress_bar
 
 
 def inside_slurm():
@@ -37,96 +40,6 @@ def inside_slurm():
     return "SLURM_JOB_ID" in os.environ
 
 
-def auto_select_torch_device() -> torch.device:
-    """Tries to select automatically a torch device."""
-    if torch.cuda.is_available():
-        logging.info("Cuda backend detected, using cuda.")
-        return torch.device("cuda")
-    elif torch.backends.mps.is_available():
-        logging.info("Metal backend detected, using mps.")
-        return torch.device("mps")
-    elif torch.xpu.is_available():
-        logging.info("Intel XPU backend detected, using xpu.")
-        return torch.device("xpu")
-    else:
-        logging.warning("No accelerated backend detected. Using default cpu, this will be slow.")
-        return torch.device("cpu")
-
-
-# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
-def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
-    """Given a string, return a torch.device with checks on whether the device is available."""
-    try_device = str(try_device)
-    if try_device.startswith("cuda"):
-        assert torch.cuda.is_available()
-        device = torch.device(try_device)
-    elif try_device == "mps":
-        assert torch.backends.mps.is_available()
-        device = torch.device("mps")
-    elif try_device == "xpu":
-        assert torch.xpu.is_available()
-        device = torch.device("xpu")
-    elif try_device == "cpu":
-        device = torch.device("cpu")
-        if log:
-            logging.warning("Using CPU, this will be slow.")
-    else:
-        device = torch.device(try_device)
-        if log:
-            logging.warning(f"Using custom {try_device} device.")
-    return device
-
-
-def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
-    """
-    mps is currently not compatible with float64
-    """
-    if isinstance(device, torch.device):
-        device = device.type
-    if device == "mps" and dtype == torch.float64:
-        return torch.float32
-    if device == "xpu" and dtype == torch.float64:
-        if hasattr(torch.xpu, "get_device_capability"):
-            device_capability = torch.xpu.get_device_capability()
-            # NOTE: Some Intel XPU devices do not support double precision (FP64).
-            # The `has_fp64` flag is returned by `torch.xpu.get_device_capability()`
-            # when available; if False, we fall back to float32 for compatibility.
-            if not device_capability.get("has_fp64", False):
-                logging.warning(f"Device {device} does not support float64, using float32 instead.")
-                return torch.float32
-        else:
-            logging.warning(
-                f"Device {device} capability check failed. Assuming no support for float64, using float32 instead."
-            )
-            return torch.float32
-        return dtype
-    else:
-        return dtype
-
-
-def is_torch_device_available(try_device: str) -> bool:
-    try_device = str(try_device)  # Ensure try_device is a string
-    if try_device.startswith("cuda"):
-        return torch.cuda.is_available()
-    elif try_device == "mps":
-        return torch.backends.mps.is_available()
-    elif try_device == "xpu":
-        return torch.xpu.is_available()
-    elif try_device == "cpu":
-        return True
-    else:
-        raise ValueError(f"Unknown device {try_device}. Supported devices are: cuda, mps, xpu or cpu.")
-
-
-def is_amp_available(device: str):
-    if device in ["cuda", "xpu", "cpu"]:
-        return True
-    elif device == "mps":
-        return False
-    else:
-        raise ValueError(f"Unknown device '{device}.")
-
-
 def init_logging(
     log_file: Path | None = None,
     display_pid: bool = False,
@@ -297,9 +210,13 @@ class SuppressProgressBars:
     """
 
     def __enter__(self):
+        from datasets.utils.logging import disable_progress_bar
+
         disable_progress_bar()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        from datasets.utils.logging import enable_progress_bar
+
         enable_progress_bar()
 
 

src/lerobot/utils/visualization_utils.py

@@ -18,7 +18,7 @@ import os
 import numpy as np
 import rerun as rr
 
-from lerobot.processor import RobotAction, RobotObservation
+from lerobot.types import RobotAction, RobotObservation
 
 from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR
 

tests/mocks/mock_robot.py

@@ -20,8 +20,8 @@ from functools import cached_property
 
 from lerobot.cameras import CameraConfig, make_cameras_from_configs
 from lerobot.motors.motors_bus import Motor, MotorNormMode
-from lerobot.processor import RobotAction, RobotObservation
 from lerobot.robots import Robot, RobotConfig
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 from tests.mocks.mock_motors_bus import MockMotorsBus
 

tests/mocks/mock_teleop.py

@@ -19,8 +19,8 @@ from dataclasses import dataclass
 from functools import cached_property
 from typing import Any
 
-from lerobot.processor import RobotAction
 from lerobot.teleoperators import Teleoperator, TeleoperatorConfig
+from lerobot.types import RobotAction
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
 
 

tests/policies/groot/test_groot_lerobot.py

@@ -28,8 +28,9 @@ import torch
 from lerobot.policies.groot.configuration_groot import GrootConfig
 from lerobot.policies.groot.modeling_groot import GrootPolicy
 from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
-from lerobot.processor import PolicyAction, PolicyProcessorPipeline
-from lerobot.utils.utils import auto_select_torch_device
+from lerobot.processor import PolicyProcessorPipeline
+from lerobot.types import PolicyAction
+from lerobot.utils.device_utils import auto_select_torch_device
 from tests.utils import require_cuda  # noqa: E402
 
 pytest.importorskip("transformers")

tests/policies/groot/test_groot_vs_original.py

@@ -28,7 +28,8 @@ import torch
 from lerobot.policies.groot.configuration_groot import GrootConfig
 from lerobot.policies.groot.modeling_groot import GrootPolicy
 from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
-from lerobot.processor import PolicyAction, PolicyProcessorPipeline
+from lerobot.processor import PolicyProcessorPipeline
+from lerobot.types import PolicyAction
 
 pytest.importorskip("gr00t")
 pytest.importorskip("transformers")

tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py

@@ -31,7 +31,8 @@ pytest.importorskip("scipy")
 from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
 from lerobot.policies.pi0_fast.modeling_pi0_fast import PI0FastPolicy
 from lerobot.policies.pi0_fast.processor_pi0_fast import make_pi0_fast_pre_post_processors
-from lerobot.processor import PolicyAction, PolicyProcessorPipeline  # noqa: E402
+from lerobot.processor import PolicyProcessorPipeline  # noqa: E402
+from lerobot.types import PolicyAction  # noqa: E402
 from lerobot.utils.constants import (
     ACTION_TOKEN_MASK,
     ACTION_TOKENS,

tests/policies/pi0_pi05/test_pi05_original_vs_lerobot.py

@@ -42,7 +42,8 @@ from transformers import AutoTokenizer  # noqa: E402
 
 from lerobot.policies.pi05 import PI05Config, PI05Policy  # noqa: E402
 from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors  # noqa: E402
-from lerobot.processor import PolicyAction, PolicyProcessorPipeline  # noqa: E402
+from lerobot.processor import PolicyProcessorPipeline  # noqa: E402
+from lerobot.types import PolicyAction  # noqa: E402
 
 # TODO: ADDING DEFAULT IMAGES_FEATURES TO CONFIG
 DUMMY_ACTION_DIM = 32

tests/policies/pi0_pi05/test_pi0_original_vs_lerobot.py

@@ -41,7 +41,8 @@ from transformers import AutoTokenizer  # noqa: E402
 
 from lerobot.policies.pi0 import PI0Config, PI0Policy  # noqa: E402
 from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors  # noqa: E402
-from lerobot.processor import PolicyAction, PolicyProcessorPipeline  # noqa: E402
+from lerobot.processor import PolicyProcessorPipeline  # noqa: E402
+from lerobot.types import PolicyAction  # noqa: E402
 
 # TODO: ADDING DEFAULT IMAGES_FEATURES TO CONFIG
 DUMMY_ACTION_DIM = 32

tests/policies/test_sarm_processor.py

@@ -25,7 +25,7 @@ import pandas as pd
 import pytest
 import torch
 
-from lerobot.processor.core import TransitionKey
+from lerobot.types import TransitionKey
 
 
 class MockDatasetMeta:

tests/policies/xvla/test_xvla_original_vs_lerobot.py

@@ -30,7 +30,8 @@ pytest.importorskip("transformers")
 from lerobot.policies.xvla.configuration_xvla import XVLAConfig
 from lerobot.policies.xvla.modeling_xvla import XVLAPolicy
 from lerobot.policies.xvla.processor_xvla import make_xvla_pre_post_processors
-from lerobot.processor import PolicyAction, PolicyProcessorPipeline  # noqa: E402
+from lerobot.processor import PolicyProcessorPipeline  # noqa: E402
+from lerobot.types import PolicyAction  # noqa: E402
 from lerobot.utils.constants import OBS_IMAGES, OBS_STATE  # noqa: E402
 from tests.utils import require_cuda  # noqa: E402
 

tests/processor/test_batch_conversion.py

@@ -16,8 +16,9 @@
 
 import torch
 
-from lerobot.processor import DataProcessorPipeline, TransitionKey
+from lerobot.processor import DataProcessorPipeline
 from lerobot.processor.converters import batch_to_transition, transition_to_batch
+from lerobot.types import TransitionKey
 from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_PREFIX, OBS_STATE, REWARD, TRUNCATED
 
 

tests/processor/test_converters.py

@@ -18,13 +18,13 @@ import numpy as np
 import pytest
 import torch
 
-from lerobot.processor import TransitionKey
 from lerobot.processor.converters import (
     batch_to_transition,
     create_transition,
     to_tensor,
     transition_to_batch,
 )
+from lerobot.types import TransitionKey
 from lerobot.utils.constants import ACTION, DONE, OBS_STATE, OBS_STR, REWARD
 
 

tests/processor/test_device_processor.py

@@ -19,8 +19,9 @@ import pytest
 import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.processor import DataProcessorPipeline, DeviceProcessorStep, TransitionKey
+from lerobot.processor import DataProcessorPipeline, DeviceProcessorStep
 from lerobot.processor.converters import create_transition, identity_transition
+from lerobot.types import TransitionKey
 from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
 
 

tests/processor/test_normalize_processor.py

@@ -30,7 +30,7 @@ from lerobot.processor import (
 )
 from lerobot.processor.converters import create_transition, identity_transition, to_tensor
 from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE, OBS_STR
-from lerobot.utils.utils import auto_select_torch_device
+from lerobot.utils.device_utils import auto_select_torch_device
 
 
 def test_numpy_conversion():

tests/processor/test_observation_processor.py

@@ -19,8 +19,9 @@ import pytest
 import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType
-from lerobot.processor import TransitionKey, VanillaObservationProcessorStep
+from lerobot.processor import VanillaObservationProcessorStep
 from lerobot.processor.converters import create_transition
+from lerobot.types import TransitionKey
 from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from tests.conftest import assert_contract_is_typed
 

tests/processor/test_tokenizer_processor.py

@@ -25,8 +25,9 @@ import pytest
 import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep, TransitionKey
+from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep
 from lerobot.processor.converters import create_transition, identity_transition
+from lerobot.types import TransitionKey
 from lerobot.utils.constants import (
     ACTION,
     OBS_IMAGE,

tests/training/test_visual_validation.py

@@ -37,7 +37,7 @@ from lerobot.configs.train import TrainPipelineConfig
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.policies.factory import make_policy_config
 from lerobot.scripts.lerobot_train import train
-from lerobot.utils.utils import auto_select_torch_device
+from lerobot.utils.device_utils import auto_select_torch_device
 
 pytest.importorskip("transformers")
 

tests/utils.py

@@ -21,8 +21,8 @@ import pytest
 import torch
 
 from lerobot import available_cameras, available_motors, available_robots
+from lerobot.utils.device_utils import auto_select_torch_device
 from lerobot.utils.import_utils import is_package_available
-from lerobot.utils.utils import auto_select_torch_device
 
 DEVICE = os.environ.get("LEROBOT_TEST_DEVICE", str(auto_select_torch_device()))
 

tests/utils/test_visualization_utils.py

@@ -21,7 +21,7 @@ from types import SimpleNamespace
 import numpy as np
 import pytest
 
-from lerobot.processor import TransitionKey
+from lerobot.types import TransitionKey
 from lerobot.utils.constants import OBS_STATE