diff --git a/examples/isaac_teleop_to_so101/common.py b/examples/isaac_teleop_to_so101/common.py index 80f56bae4..76a89be24 100644 --- a/examples/isaac_teleop_to_so101/common.py +++ b/examples/isaac_teleop_to_so101/common.py @@ -44,6 +44,7 @@ from typing import Protocol import numpy as np +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( RobotProcessorPipeline, @@ -56,7 +57,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( EEBoundsAndSafety, InverseKinematicsEEToJoints, ) -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, HF_LEROBOT_HOME, TELEOPERATORS from lerobot.utils.robot_utils import precise_sleep diff --git a/examples/isaac_teleop_to_so101/isaac_teleop/teleop_so101_leader_arm.py b/examples/isaac_teleop_to_so101/isaac_teleop/teleop_so101_leader_arm.py index f23c4e9b7..b610c6e7a 100644 --- a/examples/isaac_teleop_to_so101/isaac_teleop/teleop_so101_leader_arm.py +++ b/examples/isaac_teleop_to_so101/isaac_teleop/teleop_so101_leader_arm.py @@ -38,7 +38,7 @@ from typing import TYPE_CHECKING import numpy as np -from lerobot.types import RobotAction +from lerobot.lerobot_types import RobotAction from .base import _GRIPPER_MOTOR_SCALE, IsaacTeleopTeleoperator, _isaacteleop_available from .config_isaac_teleop import SO101LeaderArmConfig diff --git a/examples/isaac_teleop_to_so101/isaac_teleop/teleop_xr_controller.py b/examples/isaac_teleop_to_so101/isaac_teleop/teleop_xr_controller.py index 25949876b..ac298afff 100644 --- a/examples/isaac_teleop_to_so101/isaac_teleop/teleop_xr_controller.py +++ b/examples/isaac_teleop_to_so101/isaac_teleop/teleop_xr_controller.py @@ -32,7 +32,7 @@ from typing import TYPE_CHECKING, Any import numpy as np -from lerobot.types import RobotAction +from lerobot.lerobot_types import RobotAction from .base import IsaacTeleopTeleoperator, _isaacteleop_available from .config_isaac_teleop import XRControllerConfig diff --git a/examples/isaac_teleop_to_so101/isaac_teleop/xr_controller_processor.py b/examples/isaac_teleop_to_so101/isaac_teleop/xr_controller_processor.py index b0eaa8280..99165ebfb 100644 --- a/examples/isaac_teleop_to_so101/isaac_teleop/xr_controller_processor.py +++ b/examples/isaac_teleop_to_so101/isaac_teleop/xr_controller_processor.py @@ -26,8 +26,8 @@ from __future__ import annotations from dataclasses import dataclass from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import RobotAction from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep -from lerobot.types import RobotAction from lerobot.utils.rotation import Rotation from .base import _GRIPPER_MOTOR_SCALE diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index d1fb4de67..03925b1e1 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -21,6 +21,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.common.control_utils import predict_action from lerobot.configs import FeatureType, PolicyFeature from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.policies import make_pre_post_processors from lerobot.policies.act import ACTPolicy @@ -38,7 +39,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( ForwardKinematicsJointsToEE, InverseKinematicsEEToJoints, ) -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts from lerobot.utils.keyboard_input import init_keyboard_listener diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index 612e94ab9..826bee1a1 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -16,6 +16,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( RobotProcessorPipeline, @@ -36,7 +37,6 @@ from lerobot.scripts.lerobot_record import record_loop from lerobot.teleoperators.phone import Phone, PhoneConfig from lerobot.teleoperators.phone.config_phone import PhoneOS from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.feature_utils import combine_feature_dicts from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say diff --git a/examples/phone_to_so100/replay.py b/examples/phone_to_so100/replay.py index 7a2cb340f..bd9633180 100644 --- a/examples/phone_to_so100/replay.py +++ b/examples/phone_to_so100/replay.py @@ -17,6 +17,7 @@ import time from lerobot.datasets import LeRobotDataset +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( RobotProcessorPipeline, @@ -27,7 +28,6 @@ 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 diff --git a/examples/phone_to_so100/rollout.py b/examples/phone_to_so100/rollout.py index ca6706c52..c3f014d63 100644 --- a/examples/phone_to_so100/rollout.py +++ b/examples/phone_to_so100/rollout.py @@ -27,6 +27,7 @@ Highlight, or DAgger via ``lerobot-rollout --strategy.type=...``. from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.configs import PreTrainedConfig +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( RobotProcessorPipeline, @@ -43,7 +44,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context from lerobot.rollout.inference import SyncInferenceConfig from lerobot.rollout.strategies import BaseStrategy -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.process import ProcessSignalHandler from lerobot.utils.utils import init_logging diff --git a/examples/phone_to_so100/teleoperate.py b/examples/phone_to_so100/teleoperate.py index 9065de4bb..d28d3da5c 100644 --- a/examples/phone_to_so100/teleoperate.py +++ b/examples/phone_to_so100/teleoperate.py @@ -15,6 +15,7 @@ import time +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( RobotProcessorPipeline, @@ -31,7 +32,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.teleoperators.phone import Phone, PhoneConfig from lerobot.teleoperators.phone.config_phone import PhoneOS from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction -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 diff --git a/examples/so100_to_so100_EE/evaluate.py b/examples/so100_to_so100_EE/evaluate.py index 2a2022623..48ae5bf11 100644 --- a/examples/so100_to_so100_EE/evaluate.py +++ b/examples/so100_to_so100_EE/evaluate.py @@ -21,6 +21,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.common.control_utils import predict_action from lerobot.configs import FeatureType, PolicyFeature from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.policies import make_pre_post_processors from lerobot.policies.act import ACTPolicy @@ -38,7 +39,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( ForwardKinematicsJointsToEE, InverseKinematicsEEToJoints, ) -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts from lerobot.utils.keyboard_input import init_keyboard_listener diff --git a/examples/so100_to_so100_EE/record.py b/examples/so100_to_so100_EE/record.py index 3706ee4f5..a8e49bdaf 100644 --- a/examples/so100_to_so100_EE/record.py +++ b/examples/so100_to_so100_EE/record.py @@ -17,6 +17,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( RobotProcessorPipeline, @@ -33,7 +34,6 @@ 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.feature_utils import combine_feature_dicts from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say diff --git a/examples/so100_to_so100_EE/replay.py b/examples/so100_to_so100_EE/replay.py index 24d6f26fb..0d98de708 100644 --- a/examples/so100_to_so100_EE/replay.py +++ b/examples/so100_to_so100_EE/replay.py @@ -18,6 +18,7 @@ import time from lerobot.datasets import LeRobotDataset +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( RobotProcessorPipeline, @@ -28,7 +29,6 @@ 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 diff --git a/examples/so100_to_so100_EE/rollout.py b/examples/so100_to_so100_EE/rollout.py index d608bfab2..a55db1aad 100644 --- a/examples/so100_to_so100_EE/rollout.py +++ b/examples/so100_to_so100_EE/rollout.py @@ -25,6 +25,7 @@ forward/inverse kinematics. from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.configs import PreTrainedConfig +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( RobotProcessorPipeline, @@ -41,7 +42,6 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context from lerobot.rollout.inference import SyncInferenceConfig from lerobot.rollout.strategies import BaseStrategy -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.process import ProcessSignalHandler from lerobot.utils.utils import init_logging diff --git a/examples/so100_to_so100_EE/teleoperate.py b/examples/so100_to_so100_EE/teleoperate.py index c0ff7482b..fcad09621 100644 --- a/examples/so100_to_so100_EE/teleoperate.py +++ b/examples/so100_to_so100_EE/teleoperate.py @@ -16,6 +16,7 @@ import time +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( RobotProcessorPipeline, @@ -30,7 +31,6 @@ 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 diff --git a/src/lerobot/async_inference/policy_server.py b/src/lerobot/async_inference/policy_server.py index 787d39abf..c43aeed80 100644 --- a/src/lerobot/async_inference/policy_server.py +++ b/src/lerobot/async_inference/policy_server.py @@ -38,6 +38,7 @@ import draccus import grpc import torch +from lerobot.lerobot_types import PolicyAction from lerobot.policies import get_policy_class, make_pre_post_processors from lerobot.processor import PolicyProcessorPipeline from lerobot.transport import ( @@ -45,7 +46,6 @@ from lerobot.transport import ( 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 diff --git a/src/lerobot/common/control_utils.py b/src/lerobot/common/control_utils.py index e3130643d..6a5e8bded 100644 --- a/src/lerobot/common/control_utils.py +++ b/src/lerobot/common/control_utils.py @@ -35,9 +35,9 @@ else: if TYPE_CHECKING: from lerobot.datasets import LeRobotDataset +from lerobot.lerobot_types import PolicyAction from lerobot.processor import PolicyProcessorPipeline from lerobot.robots import Robot -from lerobot.types import PolicyAction def predict_action( diff --git a/src/lerobot/datasets/pipeline_features.py b/src/lerobot/datasets/pipeline_features.py index 91feee2cd..f35744097 100644 --- a/src/lerobot/datasets/pipeline_features.py +++ b/src/lerobot/datasets/pipeline_features.py @@ -17,8 +17,8 @@ from collections.abc import Sequence from typing import Any from lerobot.configs import PipelineFeatureType +from lerobot.lerobot_types import 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 from lerobot.utils.feature_utils import hw_to_dataset_features diff --git a/src/lerobot/envs/libero.py b/src/lerobot/envs/libero.py index c0cc03e51..b959ca8fb 100644 --- a/src/lerobot/envs/libero.py +++ b/src/lerobot/envs/libero.py @@ -30,7 +30,7 @@ from gymnasium import spaces from libero.libero import benchmark, get_libero_path from libero.libero.envs import OffScreenRenderEnv -from lerobot.types import RobotObservation +from lerobot.lerobot_types import RobotObservation from .utils import _LazyAsyncVectorEnv, parse_camera_names diff --git a/src/lerobot/envs/metaworld.py b/src/lerobot/envs/metaworld.py index a32acbff1..fabe428ef 100644 --- a/src/lerobot/envs/metaworld.py +++ b/src/lerobot/envs/metaworld.py @@ -25,7 +25,7 @@ import metaworld.policies as policies import numpy as np from gymnasium import spaces -from lerobot.types import RobotObservation +from lerobot.lerobot_types import RobotObservation from .utils import _LazyAsyncVectorEnv diff --git a/src/lerobot/envs/robocasa.py b/src/lerobot/envs/robocasa.py index a84a7c766..35caeee62 100644 --- a/src/lerobot/envs/robocasa.py +++ b/src/lerobot/envs/robocasa.py @@ -25,7 +25,7 @@ import gymnasium as gym import numpy as np from gymnasium import spaces -from lerobot.types import RobotObservation +from lerobot.lerobot_types import RobotObservation from .utils import _LazyAsyncVectorEnv, parse_camera_names diff --git a/src/lerobot/envs/robotwin.py b/src/lerobot/envs/robotwin.py index cd0cef28f..65d0186c7 100644 --- a/src/lerobot/envs/robotwin.py +++ b/src/lerobot/envs/robotwin.py @@ -28,7 +28,7 @@ import numpy as np import torch from gymnasium import spaces -from lerobot.types import RobotObservation +from lerobot.lerobot_types import RobotObservation from lerobot.utils.import_utils import _scipy_available from .utils import _LazyAsyncVectorEnv diff --git a/src/lerobot/envs/vlabench.py b/src/lerobot/envs/vlabench.py index 02d6bdcc5..9fd7a777e 100644 --- a/src/lerobot/envs/vlabench.py +++ b/src/lerobot/envs/vlabench.py @@ -37,7 +37,7 @@ import numpy as np from gymnasium import spaces from scipy.spatial.transform import Rotation -from lerobot.types import RobotObservation +from lerobot.lerobot_types import RobotObservation from .utils import _LazyAsyncVectorEnv diff --git a/src/lerobot/types.py b/src/lerobot/lerobot_types.py similarity index 97% rename from src/lerobot/types.py rename to src/lerobot/lerobot_types.py index 9de504870..a05b27712 100644 --- a/src/lerobot/types.py +++ b/src/lerobot/lerobot_types.py @@ -23,7 +23,7 @@ import numpy as np import torch -class TransitionKey(str, Enum): +class TransitionKey(str, Enum): # noqa: UP042 """Keys for accessing EnvTransition dictionary components.""" # TODO(Steven): Use consts diff --git a/src/lerobot/policies/eo1/processor_eo1.py b/src/lerobot/policies/eo1/processor_eo1.py index 1d3dc0515..18859339d 100644 --- a/src/lerobot/policies/eo1/processor_eo1.py +++ b/src/lerobot/policies/eo1/processor_eo1.py @@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Any import torch from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import TransitionKey from lerobot.processor import ( ComplementaryDataProcessorStep, PolicyAction, @@ -31,7 +32,6 @@ from lerobot.processor import ( make_default_policy_processor_steps, make_policy_processor_pipelines, ) -from lerobot.types import TransitionKey from lerobot.utils.constants import OBS_STATE from lerobot.utils.import_utils import _transformers_available, require_package diff --git a/src/lerobot/policies/evo1/processor_evo1.py b/src/lerobot/policies/evo1/processor_evo1.py index ba818f9db..d0ddf22b5 100644 --- a/src/lerobot/policies/evo1/processor_evo1.py +++ b/src/lerobot/policies/evo1/processor_evo1.py @@ -21,6 +21,7 @@ from typing import Any import torch from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, @@ -40,7 +41,6 @@ from lerobot.processor.converters import ( policy_action_to_transition, transition_to_policy_action, ) -from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import ( ACTION, DONE, diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 20c45f6c7..3527a1821 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -28,6 +28,7 @@ if TYPE_CHECKING: from lerobot.configs import FeatureType, PreTrainedConfig from lerobot.envs import EnvConfig, env_to_policy_features +from lerobot.lerobot_types import PolicyAction from lerobot.processor import ( AbsoluteActionsProcessorStep, PolicyProcessorPipeline, @@ -37,7 +38,6 @@ from lerobot.processor import ( transition_to_batch, transition_to_policy_action, ) -from lerobot.types import PolicyAction from lerobot.utils.constants import ( ACTION, POLICY_POSTPROCESSOR_DEFAULT_NAME, diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 0bd976a85..3cdd25b18 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -50,6 +50,7 @@ if TYPE_CHECKING or _datasets_available: else: LeRobotDataset = None +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.processor import ( AbsoluteActionsProcessorStep, AddBatchDimensionProcessorStep, @@ -66,7 +67,6 @@ from lerobot.processor import ( transition_to_batch, transition_to_policy_action, ) -from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import ( ACTION, OBS_IMAGE, diff --git a/src/lerobot/policies/molmoact2/processor_molmoact2.py b/src/lerobot/policies/molmoact2/processor_molmoact2.py index d2db817ef..701092225 100644 --- a/src/lerobot/policies/molmoact2/processor_molmoact2.py +++ b/src/lerobot/policies/molmoact2/processor_molmoact2.py @@ -36,6 +36,7 @@ import torch from torch import Tensor from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, @@ -49,7 +50,6 @@ from lerobot.processor import ( policy_action_to_transition, transition_to_policy_action, ) -from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import ( ACTION, OBS_IMAGES, diff --git a/src/lerobot/policies/pi05/processor_pi05.py b/src/lerobot/policies/pi05/processor_pi05.py index ac16269a0..a0b5a9f0e 100644 --- a/src/lerobot/policies/pi05/processor_pi05.py +++ b/src/lerobot/policies/pi05/processor_pi05.py @@ -22,6 +22,7 @@ import numpy as np import torch from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.processor import ( AbsoluteActionsProcessorStep, PolicyAction, @@ -33,7 +34,6 @@ from lerobot.processor import ( make_default_policy_processor_steps, make_policy_processor_pipelines, ) -from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import OBS_STATE from .configuration_pi05 import PI05Config diff --git a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py b/src/lerobot/policies/pi0_fast/processor_pi0_fast.py index fdf759e29..864dafcfb 100644 --- a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py +++ b/src/lerobot/policies/pi0_fast/processor_pi0_fast.py @@ -22,6 +22,7 @@ import numpy as np import torch from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.processor import ( AbsoluteActionsProcessorStep, ActionTokenizerProcessorStep, @@ -34,7 +35,6 @@ from lerobot.processor import ( make_default_policy_processor_steps, make_policy_processor_pipelines, ) -from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import OBS_STATE from .configuration_pi0_fast import PI0FastConfig diff --git a/src/lerobot/policies/utils.py b/src/lerobot/policies/utils.py index f465fcff8..3f7192f31 100644 --- a/src/lerobot/policies/utils.py +++ b/src/lerobot/policies/utils.py @@ -22,7 +22,7 @@ import torch from torch import nn from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig -from lerobot.types import PolicyAction, RobotAction, RobotObservation +from lerobot.lerobot_types import PolicyAction, RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame diff --git a/src/lerobot/policies/xvla/processor_xvla.py b/src/lerobot/policies/xvla/processor_xvla.py index 02c54f9c9..4ee8fe793 100644 --- a/src/lerobot/policies/xvla/processor_xvla.py +++ b/src/lerobot/policies/xvla/processor_xvla.py @@ -21,6 +21,7 @@ import numpy as np import torch from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.processor import ( ObservationProcessorStep, PolicyAction, @@ -31,7 +32,6 @@ from lerobot.processor import ( make_default_policy_processor_steps, make_policy_processor_pipelines, ) -from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import ( IMAGENET_STATS, OBS_IMAGES, diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 79bb44c85..6ed762b6a 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from lerobot.types import ( +from lerobot.lerobot_types import ( EnvAction, EnvTransition, PolicyAction, diff --git a/src/lerobot/processor/batch_processor.py b/src/lerobot/processor/batch_processor.py index 669c68a0a..27742f3b2 100644 --- a/src/lerobot/processor/batch_processor.py +++ b/src/lerobot/processor/batch_processor.py @@ -25,7 +25,7 @@ from dataclasses import dataclass, field from torch import Tensor from lerobot.configs import PipelineFeatureType, PolicyFeature -from lerobot.types import EnvTransition, PolicyAction +from lerobot.lerobot_types import EnvTransition, PolicyAction from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE from .pipeline import ( diff --git a/src/lerobot/processor/converters.py b/src/lerobot/processor/converters.py index faa4d5cd9..3cd9debf0 100644 --- a/src/lerobot/processor/converters.py +++ b/src/lerobot/processor/converters.py @@ -23,7 +23,7 @@ from typing import Any import numpy as np import torch -from lerobot.types import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey +from lerobot.lerobot_types import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey from lerobot.utils.constants import ACTION, DONE, INFO, OBS_PREFIX, REWARD, TRUNCATED diff --git a/src/lerobot/processor/delta_action_processor.py b/src/lerobot/processor/delta_action_processor.py index 57d353c5c..6af0d8c36 100644 --- a/src/lerobot/processor/delta_action_processor.py +++ b/src/lerobot/processor/delta_action_processor.py @@ -17,7 +17,7 @@ from dataclasses import dataclass from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature -from lerobot.types import PolicyAction, RobotAction +from lerobot.lerobot_types import PolicyAction, RobotAction from .pipeline import ActionProcessorStep, ProcessorStepRegistry, RobotActionProcessorStep diff --git a/src/lerobot/processor/device_processor.py b/src/lerobot/processor/device_processor.py index 1171c7e78..2064b37fd 100644 --- a/src/lerobot/processor/device_processor.py +++ b/src/lerobot/processor/device_processor.py @@ -25,7 +25,7 @@ from typing import Any import torch from lerobot.configs import PipelineFeatureType, PolicyFeature -from lerobot.types import EnvTransition, PolicyAction, TransitionKey +from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey from lerobot.utils.device_utils import get_safe_torch_device from .pipeline import ProcessorStep, ProcessorStepRegistry diff --git a/src/lerobot/processor/factory.py b/src/lerobot/processor/factory.py index ec929f365..bb038b108 100644 --- a/src/lerobot/processor/factory.py +++ b/src/lerobot/processor/factory.py @@ -20,7 +20,7 @@ from typing import Any import torch from lerobot.configs.policies import PreTrainedConfig -from lerobot.types import PolicyAction, RobotAction, RobotObservation +from lerobot.lerobot_types import PolicyAction, RobotAction, RobotObservation from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .batch_processor import AddBatchDimensionProcessorStep diff --git a/src/lerobot/processor/gym_action_processor.py b/src/lerobot/processor/gym_action_processor.py index 2ec5f6e64..549e2b1e0 100644 --- a/src/lerobot/processor/gym_action_processor.py +++ b/src/lerobot/processor/gym_action_processor.py @@ -17,7 +17,7 @@ from dataclasses import dataclass from lerobot.configs import PipelineFeatureType, PolicyFeature -from lerobot.types import EnvAction, EnvTransition, PolicyAction, TransitionKey +from lerobot.lerobot_types import EnvAction, EnvTransition, PolicyAction, TransitionKey from .converters import to_tensor from .hil_processor import TELEOP_ACTION_KEY diff --git a/src/lerobot/processor/hil_processor.py b/src/lerobot/processor/hil_processor.py index e7351827b..e0c9fe733 100644 --- a/src/lerobot/processor/hil_processor.py +++ b/src/lerobot/processor/hil_processor.py @@ -29,7 +29,7 @@ from lerobot.teleoperators.utils import TeleopEvents if TYPE_CHECKING: from lerobot.teleoperators.teleoperator import Teleoperator -from lerobot.types import EnvTransition, PolicyAction, TransitionKey +from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey from .pipeline import ( ComplementaryDataProcessorStep, diff --git a/src/lerobot/processor/normalize_processor.py b/src/lerobot/processor/normalize_processor.py index 1649b4b31..73c8922b6 100644 --- a/src/lerobot/processor/normalize_processor.py +++ b/src/lerobot/processor/normalize_processor.py @@ -25,7 +25,7 @@ import torch from torch import Tensor from lerobot.configs import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature -from lerobot.types import EnvTransition, PolicyAction, TransitionKey +from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey if TYPE_CHECKING: from lerobot.datasets import LeRobotDataset diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index e40a7c479..2db1f395b 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -45,7 +45,14 @@ from huggingface_hub import hf_hub_download from safetensors.torch import load_file, save_file from lerobot.configs import PipelineFeatureType, PolicyFeature -from lerobot.types import EnvAction, EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey +from lerobot.lerobot_types import ( + EnvAction, + EnvTransition, + PolicyAction, + RobotAction, + RobotObservation, + TransitionKey, +) from lerobot.utils.constants import HF_LEROBOT_HOME from lerobot.utils.hub import HubMixin diff --git a/src/lerobot/processor/policy_robot_bridge.py b/src/lerobot/processor/policy_robot_bridge.py index 25d622dc2..bd43c0e81 100644 --- a/src/lerobot/processor/policy_robot_bridge.py +++ b/src/lerobot/processor/policy_robot_bridge.py @@ -20,7 +20,7 @@ from typing import Any import torch from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature -from lerobot.types import PolicyAction, RobotAction +from lerobot.lerobot_types import PolicyAction, RobotAction from lerobot.utils.constants import ACTION from .pipeline import ActionProcessorStep, ProcessorStepRegistry diff --git a/src/lerobot/processor/relative_action_processor.py b/src/lerobot/processor/relative_action_processor.py index 5b1039291..4fe007f79 100644 --- a/src/lerobot/processor/relative_action_processor.py +++ b/src/lerobot/processor/relative_action_processor.py @@ -20,7 +20,7 @@ import torch from torch import Tensor from lerobot.configs import PipelineFeatureType, PolicyFeature -from lerobot.types import EnvTransition, TransitionKey +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.utils.constants import OBS_STATE from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep diff --git a/src/lerobot/processor/render_messages_processor.py b/src/lerobot/processor/render_messages_processor.py index 140592f0e..92ba9b8d9 100644 --- a/src/lerobot/processor/render_messages_processor.py +++ b/src/lerobot/processor/render_messages_processor.py @@ -23,7 +23,7 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.configs.recipe import TrainingRecipe from lerobot.datasets.language import LANGUAGE_EVENTS, LANGUAGE_PERSISTENT from lerobot.datasets.language_render import render_sample -from lerobot.types import EnvTransition, TransitionKey +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.utils.utils import unwrap_scalar from .pipeline import ProcessorStep, ProcessorStepRegistry diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py index a808e6127..967159144 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -30,7 +30,7 @@ from typing import TYPE_CHECKING, Any import torch from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature -from lerobot.types import EnvTransition, RobotObservation, TransitionKey +from lerobot.lerobot_types import EnvTransition, RobotObservation, TransitionKey from lerobot.utils.constants import ( ACTION_TOKEN_MASK, ACTION_TOKENS, diff --git a/src/lerobot/rewards/robometer/compute_rabc_weights.py b/src/lerobot/rewards/robometer/compute_rabc_weights.py index 8e2b016fb..a4f368df0 100644 --- a/src/lerobot/rewards/robometer/compute_rabc_weights.py +++ b/src/lerobot/rewards/robometer/compute_rabc_weights.py @@ -57,10 +57,10 @@ import torch from tqdm import tqdm from lerobot.datasets import LeRobotDataset +from lerobot.lerobot_types import TransitionKey from lerobot.rewards.robometer.configuration_robometer import RobometerConfig from lerobot.rewards.robometer.modeling_robometer import RobometerRewardModel from lerobot.rewards.robometer.processor_robometer import RobometerEncoderProcessorStep -from lerobot.types import TransitionKey DEFAULT_OUTPUT_FILENAME = "robometer_progress.parquet" diff --git a/src/lerobot/rewards/robometer/processor_robometer.py b/src/lerobot/rewards/robometer/processor_robometer.py index d98f8b9aa..764833202 100644 --- a/src/lerobot/rewards/robometer/processor_robometer.py +++ b/src/lerobot/rewards/robometer/processor_robometer.py @@ -25,6 +25,7 @@ from PIL import Image from torch import Tensor from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, @@ -39,7 +40,6 @@ from lerobot.rewards.robometer.configuration_robometer import ( RobometerConfig, ) from lerobot.rewards.robometer.modeling_robometer import ROBOMETER_FEATURE_PREFIX -from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import ( OBS_IMAGES, POLICY_POSTPROCESSOR_DEFAULT_NAME, diff --git a/src/lerobot/rewards/sarm/processor_sarm.py b/src/lerobot/rewards/sarm/processor_sarm.py index 37db374d4..bb3a5c69a 100644 --- a/src/lerobot/rewards/sarm/processor_sarm.py +++ b/src/lerobot/rewards/sarm/processor_sarm.py @@ -47,6 +47,7 @@ else: Faker = None # type: ignore[assignment, misc] from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import EnvTransition, PolicyAction, TransitionKey from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, @@ -58,7 +59,6 @@ from lerobot.processor import ( policy_action_to_transition, transition_to_policy_action, ) -from lerobot.types import EnvTransition, PolicyAction, TransitionKey from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME from .configuration_sarm import SARMConfig diff --git a/src/lerobot/rewards/topreward/compute_rabc_weights.py b/src/lerobot/rewards/topreward/compute_rabc_weights.py index a448654e5..9789c5b63 100644 --- a/src/lerobot/rewards/topreward/compute_rabc_weights.py +++ b/src/lerobot/rewards/topreward/compute_rabc_weights.py @@ -48,10 +48,10 @@ import torch from tqdm import tqdm from lerobot.datasets import LeRobotDataset +from lerobot.lerobot_types import TransitionKey from lerobot.rewards.topreward.configuration_topreward import TOPRewardConfig from lerobot.rewards.topreward.modeling_topreward import TOPRewardModel from lerobot.rewards.topreward.processor_topreward import TOPRewardEncoderProcessorStep -from lerobot.types import TransitionKey DEFAULT_OUTPUT_FILENAME = "topreward_progress.parquet" diff --git a/src/lerobot/rewards/topreward/processor_topreward.py b/src/lerobot/rewards/topreward/processor_topreward.py index ff0646e49..75c5fd02a 100644 --- a/src/lerobot/rewards/topreward/processor_topreward.py +++ b/src/lerobot/rewards/topreward/processor_topreward.py @@ -23,6 +23,7 @@ import torch from torch import Tensor from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import EnvTransition, TransitionKey from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, @@ -37,7 +38,6 @@ from lerobot.rewards.topreward.configuration_topreward import ( DEFAULT_PROMPT_SUFFIX_TEMPLATE, TOPRewardConfig, ) -from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import ( OBS_IMAGES, OBS_PREFIX, diff --git a/src/lerobot/rl/algorithms/base.py b/src/lerobot/rl/algorithms/base.py index 01c34584b..2c03bfc6e 100644 --- a/src/lerobot/rl/algorithms/base.py +++ b/src/lerobot/rl/algorithms/base.py @@ -28,7 +28,7 @@ from huggingface_hub.errors import HfHubHTTPError from safetensors.torch import load_file as load_safetensors, save_file as save_safetensors from torch.optim import Optimizer -from lerobot.types import BatchType +from lerobot.lerobot_types import BatchType from lerobot.utils.hub import HubMixin from .configs import RLAlgorithmConfig, TrainingStats diff --git a/src/lerobot/rl/algorithms/sac/sac_algorithm.py b/src/lerobot/rl/algorithms/sac/sac_algorithm.py index 81c44068f..31042596b 100644 --- a/src/lerobot/rl/algorithms/sac/sac_algorithm.py +++ b/src/lerobot/rl/algorithms/sac/sac_algorithm.py @@ -26,6 +26,7 @@ import torch.nn.functional as F # noqa: N812 from torch import Tensor from torch.optim import Optimizer +from lerobot.lerobot_types import BatchType from lerobot.policies.gaussian_actor.modeling_gaussian_actor import ( DISCRETE_DIMENSION_INDEX, MLP, @@ -35,7 +36,6 @@ from lerobot.policies.gaussian_actor.modeling_gaussian_actor import ( orthogonal_init, ) from lerobot.policies.utils import get_device_from_parameters -from lerobot.types import BatchType from lerobot.utils.constants import ACTION from lerobot.utils.transition import move_state_dict_to_device diff --git a/src/lerobot/rl/data_sources/__init__.py b/src/lerobot/rl/data_sources/__init__.py index 97cfe5001..0cd49cb2f 100644 --- a/src/lerobot/rl/data_sources/__init__.py +++ b/src/lerobot/rl/data_sources/__init__.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from lerobot.types import BatchType +from lerobot.lerobot_types import BatchType from .data_mixer import DataMixer, OnlineOfflineMixer diff --git a/src/lerobot/rl/data_sources/data_mixer.py b/src/lerobot/rl/data_sources/data_mixer.py index 57a2d86be..04d9691cc 100644 --- a/src/lerobot/rl/data_sources/data_mixer.py +++ b/src/lerobot/rl/data_sources/data_mixer.py @@ -16,7 +16,7 @@ from __future__ import annotations import abc -from lerobot.types import BatchType +from lerobot.lerobot_types import BatchType from ..buffer import ReplayBuffer, concatenate_batch_transitions diff --git a/src/lerobot/rl/trainer.py b/src/lerobot/rl/trainer.py index 65f00568e..11847c842 100644 --- a/src/lerobot/rl/trainer.py +++ b/src/lerobot/rl/trainer.py @@ -17,7 +17,7 @@ from __future__ import annotations from collections.abc import Iterator from typing import Any -from lerobot.types import BatchType +from lerobot.lerobot_types import BatchType from .algorithms.base import RLAlgorithm from .algorithms.configs import TrainingStats diff --git a/src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py b/src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py index 1613fa177..b92632627 100644 --- a/src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py +++ b/src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py @@ -17,7 +17,7 @@ import logging from functools import cached_property -from lerobot.types import RobotAction, RobotObservation +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.utils.bimanual import BimanualMixin from lerobot.utils.decorators import check_if_not_connected diff --git a/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py b/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py index 7780ac0fe..704fd9382 100644 --- a/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py +++ b/src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py @@ -17,7 +17,7 @@ import logging from functools import cached_property -from lerobot.types import RobotAction, RobotObservation +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.utils.bimanual import BimanualMixin from lerobot.utils.decorators import check_if_not_connected diff --git a/src/lerobot/robots/bi_so_follower/bi_so_follower.py b/src/lerobot/robots/bi_so_follower/bi_so_follower.py index 66eec3c01..dcc64224d 100644 --- a/src/lerobot/robots/bi_so_follower/bi_so_follower.py +++ b/src/lerobot/robots/bi_so_follower/bi_so_follower.py @@ -17,7 +17,7 @@ import logging from functools import cached_property -from lerobot.types import RobotAction, RobotObservation +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.utils.bimanual import BimanualMixin from lerobot.utils.decorators import check_if_not_connected diff --git a/src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py b/src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py index 76707a80c..250e28d43 100644 --- a/src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py +++ b/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.types import RobotAction, RobotObservation +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.errors import DeviceNotConnectedError diff --git a/src/lerobot/robots/hope_jr/hope_jr_arm.py b/src/lerobot/robots/hope_jr/hope_jr_arm.py index b606a4fe7..735e190af 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_arm.py +++ b/src/lerobot/robots/hope_jr/hope_jr_arm.py @@ -19,12 +19,12 @@ import time from functools import cached_property from lerobot.cameras import make_cameras_from_configs +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors import Motor, MotorNormMode from lerobot.motors.calibration_gui import RangeFinderGUI from lerobot.motors.feetech import ( FeetechMotorsBus, ) -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..robot import Robot diff --git a/src/lerobot/robots/hope_jr/hope_jr_hand.py b/src/lerobot/robots/hope_jr/hope_jr_hand.py index ce70e7e13..9f2de6d20 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_hand.py +++ b/src/lerobot/robots/hope_jr/hope_jr_hand.py @@ -19,12 +19,12 @@ import time from functools import cached_property from lerobot.cameras import make_cameras_from_configs +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors import Motor, MotorNormMode from lerobot.motors.calibration_gui import RangeFinderGUI from lerobot.motors.feetech import ( FeetechMotorsBus, ) -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..robot import Robot diff --git a/src/lerobot/robots/koch_follower/koch_follower.py b/src/lerobot/robots/koch_follower/koch_follower.py index de6f9c4a3..f72521d87 100644 --- a/src/lerobot/robots/koch_follower/koch_follower.py +++ b/src/lerobot/robots/koch_follower/koch_follower.py @@ -19,12 +19,12 @@ import time from functools import cached_property from lerobot.cameras import make_cameras_from_configs +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.dynamixel import ( DynamixelMotorsBus, OperatingMode, ) -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..robot import Robot diff --git a/src/lerobot/robots/lekiwi/lekiwi.py b/src/lerobot/robots/lekiwi/lekiwi.py index 3712a64d3..7d7f984ab 100644 --- a/src/lerobot/robots/lekiwi/lekiwi.py +++ b/src/lerobot/robots/lekiwi/lekiwi.py @@ -23,12 +23,12 @@ from typing import Any import numpy as np from lerobot.cameras import make_cameras_from_configs +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.feetech import ( FeetechMotorsBus, OperatingMode, ) -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..robot import Robot diff --git a/src/lerobot/robots/lekiwi/lekiwi_client.py b/src/lerobot/robots/lekiwi/lekiwi_client.py index 2850842a2..898730953 100644 --- a/src/lerobot/robots/lekiwi/lekiwi_client.py +++ b/src/lerobot/robots/lekiwi/lekiwi_client.py @@ -22,7 +22,7 @@ from functools import cached_property import cv2 import numpy as np -from lerobot.types import RobotAction, RobotObservation +from lerobot.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 diff --git a/src/lerobot/robots/omx_follower/omx_follower.py b/src/lerobot/robots/omx_follower/omx_follower.py index b2cfb52e9..e6e363235 100644 --- a/src/lerobot/robots/omx_follower/omx_follower.py +++ b/src/lerobot/robots/omx_follower/omx_follower.py @@ -19,13 +19,13 @@ import time from functools import cached_property from lerobot.cameras import make_cameras_from_configs +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.dynamixel import ( DriveMode, DynamixelMotorsBus, OperatingMode, ) -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..robot import Robot diff --git a/src/lerobot/robots/openarm_follower/openarm_follower.py b/src/lerobot/robots/openarm_follower/openarm_follower.py index 5f2286eb4..00adf41ab 100644 --- a/src/lerobot/robots/openarm_follower/openarm_follower.py +++ b/src/lerobot/robots/openarm_follower/openarm_follower.py @@ -20,9 +20,9 @@ from functools import cached_property from typing import Any from lerobot.cameras import make_cameras_from_configs +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.damiao import DamiaoMotorsBus -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..robot import Robot diff --git a/src/lerobot/robots/reachy2/robot_reachy2.py b/src/lerobot/robots/reachy2/robot_reachy2.py index ac5c9ef2f..a8bb5e462 100644 --- a/src/lerobot/robots/reachy2/robot_reachy2.py +++ b/src/lerobot/robots/reachy2/robot_reachy2.py @@ -19,7 +19,7 @@ import time from typing import TYPE_CHECKING, Any from lerobot.cameras import make_cameras_from_configs -from lerobot.types import RobotAction, RobotObservation +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.utils.import_utils import _reachy2_sdk_available, require_package from ..robot import Robot diff --git a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py index bff0b86c1..fc44799f9 100644 --- a/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py +++ b/src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py @@ -21,8 +21,8 @@ from functools import cached_property from typing import TYPE_CHECKING from lerobot.cameras import make_cameras_from_configs +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors import MotorCalibration -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.import_utils import _motorbridge_available, require_package diff --git a/src/lerobot/robots/robot.py b/src/lerobot/robots/robot.py index 1b556f963..07d3fd2b8 100644 --- a/src/lerobot/robots/robot.py +++ b/src/lerobot/robots/robot.py @@ -18,8 +18,8 @@ from pathlib import Path import draccus +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors import MotorCalibration -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, ROBOTS from .config import RobotConfig diff --git a/src/lerobot/robots/so_follower/so_follower.py b/src/lerobot/robots/so_follower/so_follower.py index 4b477e75c..85dc43dc0 100644 --- a/src/lerobot/robots/so_follower/so_follower.py +++ b/src/lerobot/robots/so_follower/so_follower.py @@ -19,12 +19,12 @@ import time from functools import cached_property from lerobot.cameras import make_cameras_from_configs +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.feetech import ( FeetechMotorsBus, OperatingMode, ) -from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..robot import Robot diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 5b8be0941..efbfbbe10 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -26,7 +26,7 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable import numpy as np from lerobot.cameras import make_cameras_from_configs -from lerobot.types import RobotAction, RobotObservation +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.utils.import_utils import _unitree_sdk_available, require_package from ..robot import Robot diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index de4936653..fa6ee30a2 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -82,9 +82,9 @@ from lerobot.envs import ( make_env_pre_post_processors, preprocess_observation, ) +from lerobot.lerobot_types import PolicyAction from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.processor import PolicyProcessorPipeline -from lerobot.types import PolicyAction from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package diff --git a/src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py b/src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py index 640c45a57..09b43fb92 100644 --- a/src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py +++ b/src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py @@ -17,7 +17,7 @@ import logging from functools import cached_property -from lerobot.types import RobotAction +from lerobot.lerobot_types import RobotAction from lerobot.utils.bimanual import BimanualMixin from lerobot.utils.decorators import check_if_not_connected diff --git a/src/lerobot/teleoperators/bi_openarm_mini/bi_openarm_mini.py b/src/lerobot/teleoperators/bi_openarm_mini/bi_openarm_mini.py index 41ebdba0d..4e6f9fb90 100644 --- a/src/lerobot/teleoperators/bi_openarm_mini/bi_openarm_mini.py +++ b/src/lerobot/teleoperators/bi_openarm_mini/bi_openarm_mini.py @@ -17,7 +17,7 @@ import logging from functools import cached_property -from lerobot.types import RobotAction +from lerobot.lerobot_types import RobotAction from lerobot.utils.bimanual import BimanualMixin from lerobot.utils.decorators import check_if_not_connected diff --git a/src/lerobot/teleoperators/bi_rebot_102_leader/bi_rebot_102_leader.py b/src/lerobot/teleoperators/bi_rebot_102_leader/bi_rebot_102_leader.py index 9da866b43..ebaf6e9c2 100644 --- a/src/lerobot/teleoperators/bi_rebot_102_leader/bi_rebot_102_leader.py +++ b/src/lerobot/teleoperators/bi_rebot_102_leader/bi_rebot_102_leader.py @@ -17,7 +17,7 @@ import logging from functools import cached_property -from lerobot.types import RobotAction +from lerobot.lerobot_types import RobotAction from lerobot.utils.bimanual import BimanualMixin from lerobot.utils.decorators import check_if_not_connected diff --git a/src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py b/src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py index 690ae7aa8..47532c317 100644 --- a/src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py +++ b/src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py @@ -17,7 +17,7 @@ import logging from functools import cached_property -from lerobot.types import RobotAction +from lerobot.lerobot_types import RobotAction from lerobot.utils.bimanual import BimanualMixin from lerobot.utils.decorators import check_if_not_connected diff --git a/src/lerobot/teleoperators/gamepad/teleop_gamepad.py b/src/lerobot/teleoperators/gamepad/teleop_gamepad.py index d86d4f486..7202f044d 100644 --- a/src/lerobot/teleoperators/gamepad/teleop_gamepad.py +++ b/src/lerobot/teleoperators/gamepad/teleop_gamepad.py @@ -21,7 +21,7 @@ from typing import Any import numpy as np -from lerobot.types import RobotAction +from lerobot.lerobot_types import RobotAction from lerobot.utils.decorators import check_if_not_connected from ..teleoperator import Teleoperator diff --git a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py index 872cc7a26..8c60b6599 100644 --- a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py +++ b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py @@ -19,7 +19,7 @@ import time from queue import Queue from typing import Any -from lerobot.types import RobotAction +from lerobot.lerobot_types import RobotAction from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.import_utils import _pynput_available, require_package from lerobot.utils.keyboard_input import pynput_can_capture diff --git a/src/lerobot/teleoperators/openarm_leader/openarm_leader.py b/src/lerobot/teleoperators/openarm_leader/openarm_leader.py index a82ad26af..0a4140257 100644 --- a/src/lerobot/teleoperators/openarm_leader/openarm_leader.py +++ b/src/lerobot/teleoperators/openarm_leader/openarm_leader.py @@ -18,9 +18,9 @@ import logging import time from typing import Any +from lerobot.lerobot_types import RobotAction from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.damiao import DamiaoMotorsBus -from lerobot.types import RobotAction from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..teleoperator import Teleoperator diff --git a/src/lerobot/teleoperators/openarm_mini/openarm_mini.py b/src/lerobot/teleoperators/openarm_mini/openarm_mini.py index 9793d0b6f..6b95a63f8 100644 --- a/src/lerobot/teleoperators/openarm_mini/openarm_mini.py +++ b/src/lerobot/teleoperators/openarm_mini/openarm_mini.py @@ -18,12 +18,12 @@ import logging import time from typing import Any +from lerobot.lerobot_types import RobotAction from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.feetech import ( FeetechMotorsBus, OperatingMode, ) -from lerobot.types import RobotAction from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..teleoperator import Teleoperator diff --git a/src/lerobot/teleoperators/phone/phone_processor.py b/src/lerobot/teleoperators/phone/phone_processor.py index 3d57a5a71..c54b6dffe 100644 --- a/src/lerobot/teleoperators/phone/phone_processor.py +++ b/src/lerobot/teleoperators/phone/phone_processor.py @@ -17,8 +17,8 @@ from dataclasses import dataclass, field from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import RobotAction from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep -from lerobot.types import RobotAction from .config_phone import PhoneOS diff --git a/src/lerobot/teleoperators/rebot_102_leader/rebot_102_leader.py b/src/lerobot/teleoperators/rebot_102_leader/rebot_102_leader.py index f9f10ed69..13dbbb2b6 100644 --- a/src/lerobot/teleoperators/rebot_102_leader/rebot_102_leader.py +++ b/src/lerobot/teleoperators/rebot_102_leader/rebot_102_leader.py @@ -18,8 +18,8 @@ import logging import time from typing import TYPE_CHECKING +from lerobot.lerobot_types import RobotAction from lerobot.motors import MotorCalibration -from lerobot.types import RobotAction from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.import_utils import _motorbridge_smart_servo_available, require_package diff --git a/src/lerobot/teleoperators/teleoperator.py b/src/lerobot/teleoperators/teleoperator.py index f47904423..9368c97da 100644 --- a/src/lerobot/teleoperators/teleoperator.py +++ b/src/lerobot/teleoperators/teleoperator.py @@ -19,8 +19,8 @@ from typing import Any import draccus +from lerobot.lerobot_types import RobotAction from lerobot.motors.motors_bus import MotorCalibration -from lerobot.types import RobotAction from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, TELEOPERATORS from .config import TeleoperatorConfig diff --git a/src/lerobot/utils/foxglove_visualization.py b/src/lerobot/utils/foxglove_visualization.py index fc4136e12..90b5019b1 100644 --- a/src/lerobot/utils/foxglove_visualization.py +++ b/src/lerobot/utils/foxglove_visualization.py @@ -27,7 +27,7 @@ import time import cv2 import numpy as np -from lerobot.types import RobotAction, RobotObservation +from lerobot.lerobot_types import RobotAction, RobotObservation from .constants import ( ACTION, diff --git a/src/lerobot/utils/rerun_visualization.py b/src/lerobot/utils/rerun_visualization.py index 46f2c0b4b..b7f0cceef 100644 --- a/src/lerobot/utils/rerun_visualization.py +++ b/src/lerobot/utils/rerun_visualization.py @@ -25,7 +25,7 @@ import os import numpy as np from lerobot.configs import DEPTH_MILLIMETER_UNIT, infer_depth_unit -from lerobot.types import RobotAction, RobotObservation +from lerobot.lerobot_types import RobotAction, RobotObservation from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR from .import_utils import require_package diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index 09a89b20a..482342a03 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -21,7 +21,7 @@ this module does not import ``rerun`` or ``foxglove`` (each backend imports its ``require_package`` guard). """ -from lerobot.types import RobotAction, RobotObservation +from lerobot.lerobot_types import RobotAction, RobotObservation from .foxglove_visualization import init_foxglove, log_foxglove_data, shutdown_foxglove from .rerun_visualization import init_rerun, log_rerun_data, shutdown_rerun diff --git a/tests/mocks/mock_robot.py b/tests/mocks/mock_robot.py index 5504b30bf..ad6bbecaf 100644 --- a/tests/mocks/mock_robot.py +++ b/tests/mocks/mock_robot.py @@ -19,9 +19,9 @@ from dataclasses import dataclass, field from functools import cached_property from lerobot.cameras import CameraConfig, make_cameras_from_configs +from lerobot.lerobot_types import RobotAction, RobotObservation from lerobot.motors.motors_bus import Motor, MotorNormMode 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 diff --git a/tests/mocks/mock_teleop.py b/tests/mocks/mock_teleop.py index b84b2b891..9f1201efa 100644 --- a/tests/mocks/mock_teleop.py +++ b/tests/mocks/mock_teleop.py @@ -19,8 +19,8 @@ from dataclasses import dataclass from functools import cached_property from typing import Any +from lerobot.lerobot_types 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 diff --git a/tests/policies/groot/test_groot_lerobot.py b/tests/policies/groot/test_groot_lerobot.py index 59dd6cede..f1c22facc 100644 --- a/tests/policies/groot/test_groot_lerobot.py +++ b/tests/policies/groot/test_groot_lerobot.py @@ -27,11 +27,11 @@ import torch pytest.importorskip("transformers", reason="groot requires the `groot` extra (transformers)") +from lerobot.lerobot_types import PolicyAction 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 PolicyProcessorPipeline -from lerobot.types import PolicyAction from lerobot.utils.device_utils import auto_select_torch_device from tests.utils import require_cuda diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index 8b74e4664..08979cad4 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -27,6 +27,7 @@ from safetensors.torch import load_file from torch import nn from lerobot.configs import FeatureType, PolicyFeature +from lerobot.lerobot_types import TransitionKey from lerobot.policies.factory import make_policy_config, make_pre_post_processors from lerobot.policies.groot.configuration_groot import ( GROOT_ACTION_DECODE_TRANSFORM_LIBERO, @@ -54,7 +55,6 @@ from lerobot.processor import ( PolicyProcessorPipeline, RelativeActionsProcessorStep, ) -from lerobot.types import TransitionKey from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE diff --git a/tests/policies/groot/test_groot_n1_7_oss_parity.py b/tests/policies/groot/test_groot_n1_7_oss_parity.py index 3fced5909..bdeada369 100644 --- a/tests/policies/groot/test_groot_n1_7_oss_parity.py +++ b/tests/policies/groot/test_groot_n1_7_oss_parity.py @@ -22,6 +22,7 @@ import numpy as np import pytest import torch +from lerobot.lerobot_types import TransitionKey from lerobot.policies.groot.action_head.cross_attention_dit import AlternateVLDiT from lerobot.policies.groot.groot_n1_7 import GR00TN17 from lerobot.policies.groot.processor_groot import ( @@ -30,7 +31,6 @@ from lerobot.policies.groot.processor_groot import ( GrootN17VLMEncodeStep, _transform_n1_7_image_for_vlm_albumentations, ) -from lerobot.types import TransitionKey from lerobot.utils.constants import OBS_STATE OSS_REFERENCE_COMMIT = "ab88b50c718f6528e1df9dcbaf75865d1b604760" diff --git a/tests/policies/groot/test_groot_state_dropout.py b/tests/policies/groot/test_groot_state_dropout.py index fcfebcb8f..c6a031d69 100644 --- a/tests/policies/groot/test_groot_state_dropout.py +++ b/tests/policies/groot/test_groot_state_dropout.py @@ -24,8 +24,8 @@ regularization never activated. These tests pin the train/eval split. import torch +from lerobot.lerobot_types import TransitionKey from lerobot.policies.groot.processor_groot import GrootN17PackInputsStep -from lerobot.types import TransitionKey from lerobot.utils.constants import OBS_STATE diff --git a/tests/policies/molmoact2/test_molmoact2.py b/tests/policies/molmoact2/test_molmoact2.py index 095b73180..41c0f235f 100644 --- a/tests/policies/molmoact2/test_molmoact2.py +++ b/tests/policies/molmoact2/test_molmoact2.py @@ -31,6 +31,7 @@ pytest.importorskip("transformers") pytest.importorskip("scipy") from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature +from lerobot.lerobot_types import TransitionKey from lerobot.policies import get_policy_class, make_policy_config from lerobot.policies.molmoact2 import ( modeling_molmoact2 as molmoact2_modeling, @@ -57,7 +58,6 @@ from lerobot.policies.molmoact2.processor_molmoact2 import ( make_molmoact2_pre_post_processors, ) from lerobot.policies.rtc.configuration_rtc import RTCConfig -from lerobot.types import TransitionKey from lerobot.utils.constants import ACTION, OBS_STATE diff --git a/tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py b/tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py index b757d5a94..db485deed 100644 --- a/tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py +++ b/tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py @@ -28,11 +28,11 @@ import torch pytest.importorskip("transformers") pytest.importorskip("scipy") +from lerobot.lerobot_types import PolicyAction # noqa: E402 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 PolicyProcessorPipeline # noqa: E402 -from lerobot.types import PolicyAction # noqa: E402 from lerobot.utils.constants import ( ACTION_TOKEN_MASK, ACTION_TOKENS, diff --git a/tests/policies/xvla/test_xvla_original_vs_lerobot.py b/tests/policies/xvla/test_xvla_original_vs_lerobot.py index 3cea11329..4caeb923c 100644 --- a/tests/policies/xvla/test_xvla_original_vs_lerobot.py +++ b/tests/policies/xvla/test_xvla_original_vs_lerobot.py @@ -27,11 +27,11 @@ import torch pytest.importorskip("transformers") +from lerobot.lerobot_types import PolicyAction # noqa: E402 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 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 diff --git a/tests/processor/test_batch_conversion.py b/tests/processor/test_batch_conversion.py index d589b6c5e..76c2ace45 100644 --- a/tests/processor/test_batch_conversion.py +++ b/tests/processor/test_batch_conversion.py @@ -16,9 +16,9 @@ import torch +from lerobot.lerobot_types import 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 diff --git a/tests/processor/test_converters.py b/tests/processor/test_converters.py index 91afdd0e5..dc7beb169 100644 --- a/tests/processor/test_converters.py +++ b/tests/processor/test_converters.py @@ -18,13 +18,13 @@ import numpy as np import pytest import torch +from lerobot.lerobot_types 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 diff --git a/tests/processor/test_device_processor.py b/tests/processor/test_device_processor.py index 57b923076..2db9cab74 100644 --- a/tests/processor/test_device_processor.py +++ b/tests/processor/test_device_processor.py @@ -19,9 +19,9 @@ import pytest import torch from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import 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 diff --git a/tests/processor/test_observation_processor.py b/tests/processor/test_observation_processor.py index 923059210..2cb72aedb 100644 --- a/tests/processor/test_observation_processor.py +++ b/tests/processor/test_observation_processor.py @@ -19,9 +19,9 @@ import pytest import torch from lerobot.configs.types import FeatureType, PipelineFeatureType +from lerobot.lerobot_types import TransitionKey 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 diff --git a/tests/processor/test_pipeline_from_pretrained_helpers.py b/tests/processor/test_pipeline_from_pretrained_helpers.py index 36f9a8fad..ce61b7c29 100644 --- a/tests/processor/test_pipeline_from_pretrained_helpers.py +++ b/tests/processor/test_pipeline_from_pretrained_helpers.py @@ -30,13 +30,13 @@ import torch from safetensors.torch import save_file from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import EnvTransition from lerobot.processor.pipeline import ( DataProcessorPipeline, ProcessorMigrationError, ProcessorStep, ProcessorStepRegistry, ) -from lerobot.types import EnvTransition # Simplified Config Loading Tests diff --git a/tests/processor/test_render_messages_processor.py b/tests/processor/test_render_messages_processor.py index f96e3c0ab..62c02454a 100644 --- a/tests/processor/test_render_messages_processor.py +++ b/tests/processor/test_render_messages_processor.py @@ -7,9 +7,9 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da import torch # noqa: E402 from lerobot.configs.recipe import MessageTurn, TrainingRecipe # noqa: E402 +from lerobot.lerobot_types import TransitionKey # noqa: E402 from lerobot.processor.converters import create_transition # noqa: E402 from lerobot.processor.render_messages_processor import RenderMessagesStep # noqa: E402 -from lerobot.types import TransitionKey # noqa: E402 def test_render_messages_step_noops_without_language_columns(): diff --git a/tests/processor/test_tokenizer_processor.py b/tests/processor/test_tokenizer_processor.py index 5708e6e81..52277c78f 100644 --- a/tests/processor/test_tokenizer_processor.py +++ b/tests/processor/test_tokenizer_processor.py @@ -25,9 +25,9 @@ import pytest import torch from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import 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, diff --git a/tests/rewards/test_sarm_processor.py b/tests/rewards/test_sarm_processor.py index 65f70d396..e377ee879 100644 --- a/tests/rewards/test_sarm_processor.py +++ b/tests/rewards/test_sarm_processor.py @@ -25,7 +25,7 @@ import pandas as pd import pytest import torch -from lerobot.types import TransitionKey +from lerobot.lerobot_types import TransitionKey class MockDatasetMeta: diff --git a/tests/rewards/test_topreward_processor.py b/tests/rewards/test_topreward_processor.py index df379276e..5bc18a578 100644 --- a/tests/rewards/test_topreward_processor.py +++ b/tests/rewards/test_topreward_processor.py @@ -20,13 +20,13 @@ import pytest import torch from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.lerobot_types import TransitionKey from lerobot.rewards.topreward.processor_topreward import ( TOPREWARD_FEATURE_PREFIX, TOPREWARD_INPUT_KEYS, _expand_tasks, _prepare_video_batch, ) -from lerobot.types import TransitionKey from tests.utils import skip_if_package_missing # --------------------------------------------------------------------------- diff --git a/tests/test_package_imports.py b/tests/test_package_imports.py new file mode 100644 index 000000000..43e50399f --- /dev/null +++ b/tests/test_package_imports.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python + +# Copyright 2026 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 subprocess +import sys +from pathlib import Path + + +def test_lerobot_package_directory_does_not_shadow_stdlib_types(): + package_dir = Path(__file__).parents[1] / "src" / "lerobot" + command = ( + "from types import DynamicClassAttribute, MappingProxyType; " + "from lerobot.lerobot_types import RobotAction" + ) + + result = subprocess.run( + [sys.executable, "-c", command], + cwd=package_dir, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr diff --git a/tests/utils/test_rerun_visualization.py b/tests/utils/test_rerun_visualization.py index d4c3e6999..f93ae562b 100644 --- a/tests/utils/test_rerun_visualization.py +++ b/tests/utils/test_rerun_visualization.py @@ -23,7 +23,7 @@ import pytest pytest.importorskip("rerun", reason="rerun-sdk is required (install lerobot[viz])") -from lerobot.types import TransitionKey +from lerobot.lerobot_types import TransitionKey from lerobot.utils.constants import OBS_STATE