Merge branch 'main' into refactor/reward-models

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,