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

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