diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index d1190b363..48fcb7407 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -21,11 +21,11 @@ from lerobot.datasets.utils import merge_features from lerobot.model.kinematics import RobotKinematics from lerobot.policies.act.modeling_act import ACTPolicy from lerobot.policies.factory import make_pre_post_processors +from lerobot.processor import RobotProcessor from lerobot.processor.converters import ( to_output_robot_action, to_transition_robot_observation, ) -from lerobot.processor.pipeline import RobotProcessor from lerobot.record import record_loop from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.robot_kinematic_processor import ( diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index e9d22ef80..32723c280 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -20,12 +20,12 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features from lerobot.datasets.utils import merge_features from lerobot.model.kinematics import RobotKinematics +from lerobot.processor import RobotProcessor from lerobot.processor.converters import ( to_output_robot_action, to_transition_robot_observation, to_transition_teleop_action, ) -from lerobot.processor.pipeline import RobotProcessor from lerobot.record import record_loop from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.robot_kinematic_processor import ( diff --git a/examples/phone_to_so100/replay.py b/examples/phone_to_so100/replay.py index e39f482c2..83938d7ca 100644 --- a/examples/phone_to_so100/replay.py +++ b/examples/phone_to_so100/replay.py @@ -19,8 +19,8 @@ import time from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.model.kinematics import RobotKinematics +from lerobot.processor import RobotProcessor from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action -from lerobot.processor.pipeline import RobotProcessor from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.robot_kinematic_processor import ( AddRobotObservationAsComplimentaryData, diff --git a/src/lerobot/datasets/pipeline_features.py b/src/lerobot/datasets/pipeline_features.py index c22c6264f..98e936a24 100644 --- a/src/lerobot/datasets/pipeline_features.py +++ b/src/lerobot/datasets/pipeline_features.py @@ -17,7 +17,7 @@ from typing import Any from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE from lerobot.datasets.utils import hw_to_dataset_features -from lerobot.processor.pipeline import RobotProcessor +from lerobot.processor import RobotProcessor def aggregate_pipeline_dataset_features( diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 0bda4ad27..0bb91eeac 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -38,7 +38,7 @@ from lerobot.policies.sac.reward_model.configuration_classifier import RewardCla from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig -from lerobot.processor.pipeline import ProcessorKwargs, RobotProcessor +from lerobot.processor import ProcessorKwargs, RobotProcessor def get_policy_class(name: str) -> type[PreTrainedPolicy]: diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index 3c2b18635..56854273c 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/src/lerobot/policies/pi0/processor_pi0.py @@ -21,20 +21,18 @@ from lerobot.configs.types import PolicyFeature from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.processor import ( + ComplementaryDataProcessor, DeviceProcessor, NormalizerProcessor, ProcessorKwargs, + ProcessorStep, + ProcessorStepRegistry, + RenameProcessor, RobotProcessor, ToBatchProcessor, TokenizerProcessor, UnnormalizerProcessor, ) -from lerobot.processor.pipeline import ( - ComplementaryDataProcessor, - ProcessorStep, - ProcessorStepRegistry, -) -from lerobot.processor.rename_processor import RenameProcessor @ProcessorStepRegistry.register(name="pi0_new_line_processor") diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 0b535e238..f95908d22 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -20,19 +20,17 @@ from lerobot.configs.types import PolicyFeature from lerobot.constants import POSTPROCESSOR_DEFAULT_NAME, PREPROCESSOR_DEFAULT_NAME from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig from lerobot.processor import ( + ComplementaryDataProcessor, DeviceProcessor, NormalizerProcessor, ProcessorKwargs, + ProcessorStepRegistry, RenameProcessor, RobotProcessor, ToBatchProcessor, TokenizerProcessor, UnnormalizerProcessor, ) -from lerobot.processor.pipeline import ( - ComplementaryDataProcessor, - ProcessorStepRegistry, -) def make_smolvla_pre_post_processors( diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 9589ef695..cbb5de664 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -40,6 +40,7 @@ from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor, hot from .observation_processor import VanillaObservationProcessor from .pipeline import ( ActionProcessor, + ComplementaryDataProcessor, DoneProcessor, IdentityProcessor, InfoProcessor, @@ -58,6 +59,7 @@ __all__ = [ "ActionProcessor", "AddTeleopActionAsComplimentaryData", "AddTeleopEventsAsInfo", + "ComplementaryDataProcessor", "batch_to_transition", "create_transition", "DeviceProcessor", diff --git a/src/lerobot/processor/batch_processor.py b/src/lerobot/processor/batch_processor.py index cfa57bb26..d696bb092 100644 --- a/src/lerobot/processor/batch_processor.py +++ b/src/lerobot/processor/batch_processor.py @@ -17,10 +17,11 @@ from torch import Tensor from lerobot.configs.types import PolicyFeature from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE -from lerobot.processor.pipeline import ( + +from .core import EnvTransition +from .pipeline import ( ActionProcessor, ComplementaryDataProcessor, - EnvTransition, ObservationProcessor, ProcessorStep, ProcessorStepRegistry, diff --git a/src/lerobot/processor/delta_action_processor.py b/src/lerobot/processor/delta_action_processor.py index 38964f1fa..5b8847cb8 100644 --- a/src/lerobot/processor/delta_action_processor.py +++ b/src/lerobot/processor/delta_action_processor.py @@ -19,7 +19,8 @@ from dataclasses import dataclass from torch import Tensor from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry + +from .pipeline import ActionProcessor, ProcessorStepRegistry @ProcessorStepRegistry.register("map_tensor_to_delta_action_dict") diff --git a/src/lerobot/processor/device_processor.py b/src/lerobot/processor/device_processor.py index 3f68d9d03..649b63a44 100644 --- a/src/lerobot/processor/device_processor.py +++ b/src/lerobot/processor/device_processor.py @@ -19,10 +19,11 @@ from typing import Any import torch from lerobot.configs.types import PolicyFeature -from lerobot.processor.core import EnvTransition, TransitionKey -from lerobot.processor.pipeline import ProcessorStep, ProcessorStepRegistry from lerobot.utils.utils import get_safe_torch_device +from .core import EnvTransition, TransitionKey +from .pipeline import ProcessorStep, ProcessorStepRegistry + @ProcessorStepRegistry.register("device_processor") @dataclass diff --git a/src/lerobot/processor/gym_action_processor.py b/src/lerobot/processor/gym_action_processor.py index 54142ca9f..cdadd36d7 100644 --- a/src/lerobot/processor/gym_action_processor.py +++ b/src/lerobot/processor/gym_action_processor.py @@ -17,8 +17,9 @@ import numpy as np import torch from lerobot.configs.types import PolicyFeature -from lerobot.processor.converters import to_tensor -from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry + +from .converters import to_tensor +from .pipeline import ActionProcessor, ProcessorStepRegistry @ProcessorStepRegistry.register("torch2numpy_action_processor") diff --git a/src/lerobot/processor/hil_processor.py b/src/lerobot/processor/hil_processor.py index 82bec9f9d..740dfb68e 100644 --- a/src/lerobot/processor/hil_processor.py +++ b/src/lerobot/processor/hil_processor.py @@ -9,18 +9,18 @@ import torchvision.transforms.functional as F # noqa: N812 from lerobot.configs.types import PolicyFeature from lerobot.constants import ACTION -from lerobot.processor.pipeline import ( +from lerobot.teleoperators.teleoperator import Teleoperator +from lerobot.teleoperators.utils import TeleopEvents + +from .core import EnvTransition, TransitionKey +from .pipeline import ( ComplementaryDataProcessor, - EnvTransition, InfoProcessor, ObservationProcessor, ProcessorStep, ProcessorStepRegistry, - TransitionKey, TruncatedProcessor, ) -from lerobot.teleoperators.teleoperator import Teleoperator -from lerobot.teleoperators.utils import TeleopEvents GRIPPER_KEY = "gripper" DISCRETE_PENALTY_KEY = "discrete_penalty" diff --git a/src/lerobot/processor/joint_observations_processor.py b/src/lerobot/processor/joint_observations_processor.py index 40ec4ea1e..d98b4b6c1 100644 --- a/src/lerobot/processor/joint_observations_processor.py +++ b/src/lerobot/processor/joint_observations_processor.py @@ -4,12 +4,10 @@ from typing import Any import torch from lerobot.configs.types import PolicyFeature -from lerobot.processor.pipeline import ( - ObservationProcessor, - ProcessorStepRegistry, -) from lerobot.robots import Robot +from .pipeline import ObservationProcessor, ProcessorStepRegistry + @dataclass @ProcessorStepRegistry.register("joint_velocity_processor") diff --git a/src/lerobot/processor/migrate_policy_normalization.py b/src/lerobot/processor/migrate_policy_normalization.py index ae4b1fb23..d5b186bab 100644 --- a/src/lerobot/processor/migrate_policy_normalization.py +++ b/src/lerobot/processor/migrate_policy_normalization.py @@ -46,11 +46,12 @@ from huggingface_hub import hf_hub_download from safetensors.torch import load_file as load_safetensors from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.processor.batch_processor import ToBatchProcessor -from lerobot.processor.device_processor import DeviceProcessor -from lerobot.processor.normalize_processor import NormalizerProcessor, UnnormalizerProcessor -from lerobot.processor.pipeline import RobotProcessor -from lerobot.processor.rename_processor import RenameProcessor + +from .batch_processor import ToBatchProcessor +from .device_processor import DeviceProcessor +from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor +from .pipeline import RobotProcessor +from .rename_processor import RenameProcessor # Policy type to class mapping POLICY_CLASSES = { diff --git a/src/lerobot/processor/normalize_processor.py b/src/lerobot/processor/normalize_processor.py index 4b6949b2d..af2a810f5 100644 --- a/src/lerobot/processor/normalize_processor.py +++ b/src/lerobot/processor/normalize_processor.py @@ -9,14 +9,10 @@ from torch import Tensor from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.processor.converters import to_tensor -from lerobot.processor.pipeline import ( - EnvTransition, - ProcessorStep, - ProcessorStepRegistry, - RobotProcessor, - TransitionKey, -) + +from .converters import to_tensor +from .core import EnvTransition, TransitionKey +from .pipeline import ProcessorStep, ProcessorStepRegistry, RobotProcessor @dataclass diff --git a/src/lerobot/processor/observation_processor.py b/src/lerobot/processor/observation_processor.py index 40273548e..eeb72bb4e 100644 --- a/src/lerobot/processor/observation_processor.py +++ b/src/lerobot/processor/observation_processor.py @@ -22,7 +22,8 @@ from torch import Tensor from lerobot.configs.types import PolicyFeature from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE -from lerobot.processor.pipeline import ObservationProcessor, ProcessorStepRegistry + +from .pipeline import ObservationProcessor, ProcessorStepRegistry @dataclass diff --git a/src/lerobot/processor/rename_processor.py b/src/lerobot/processor/rename_processor.py index ebc867cac..f22a916be 100644 --- a/src/lerobot/processor/rename_processor.py +++ b/src/lerobot/processor/rename_processor.py @@ -18,10 +18,8 @@ from dataclasses import dataclass, field from typing import Any from lerobot.configs.types import PolicyFeature -from lerobot.processor.pipeline import ( - ObservationProcessor, - ProcessorStepRegistry, -) + +from .pipeline import ObservationProcessor, ProcessorStepRegistry @dataclass diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py index 6a6698a38..b743a2a3c 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -11,14 +11,11 @@ import torch from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS -from lerobot.processor.pipeline import ( - EnvTransition, - ObservationProcessor, - ProcessorStepRegistry, - TransitionKey, -) from lerobot.utils.import_utils import _transformers_available +from .core import EnvTransition, TransitionKey +from .pipeline import ObservationProcessor, ProcessorStepRegistry + if TYPE_CHECKING or _transformers_available: from transformers import AutoTokenizer else: diff --git a/src/lerobot/record.py b/src/lerobot/record.py index e83aa8e52..33a6bb98c 100644 --- a/src/lerobot/record.py +++ b/src/lerobot/record.py @@ -76,14 +76,13 @@ from lerobot.datasets.utils import hw_to_dataset_features from lerobot.datasets.video_utils import VideoEncodingManager from lerobot.policies.factory import make_policy, make_pre_post_processors from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor import RobotProcessor +from lerobot.processor import IdentityProcessor, RobotProcessor, TransitionKey from lerobot.processor.converters import ( to_output_robot_action, to_transition_robot_observation, to_transition_teleop_action, transition_to_dataset_frame, ) -from lerobot.processor.pipeline import IdentityProcessor, TransitionKey from lerobot.processor.rename_processor import rename_stats from lerobot.robots import ( # noqa: F401 Robot, diff --git a/src/lerobot/replay.py b/src/lerobot/replay.py index f6628ac6b..88148f7a7 100644 --- a/src/lerobot/replay.py +++ b/src/lerobot/replay.py @@ -47,9 +47,8 @@ from pprint import pformat from lerobot.configs import parser from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.processor import RobotProcessor +from lerobot.processor import IdentityProcessor, RobotProcessor from lerobot.processor.converters import to_output_robot_action, to_transition_teleop_action -from lerobot.processor.pipeline import IdentityProcessor from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, diff --git a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py index 1a349d106..470b56c31 100644 --- a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py +++ b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py @@ -22,7 +22,7 @@ from scipy.spatial.transform import Rotation from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.constants import ACTION, OBS_STATE from lerobot.model.kinematics import RobotKinematics -from lerobot.processor.pipeline import ( +from lerobot.processor import ( ActionProcessor, ComplementaryDataProcessor, EnvTransition, diff --git a/src/lerobot/scripts/rl/actor.py b/src/lerobot/scripts/rl/actor.py index 997bc620b..ee2f6c6b1 100644 --- a/src/lerobot/scripts/rl/actor.py +++ b/src/lerobot/scripts/rl/actor.py @@ -62,7 +62,7 @@ from lerobot.configs import parser from lerobot.configs.train import TrainRLServerPipelineConfig from lerobot.policies.factory import make_policy from lerobot.policies.sac.modeling_sac import SACPolicy -from lerobot.processor.pipeline import TransitionKey +from lerobot.processor import TransitionKey from lerobot.robots import so100_follower # noqa: F401 from lerobot.scripts.rl.gym_manipulator import ( create_transition, diff --git a/src/lerobot/scripts/rl/gym_manipulator.py b/src/lerobot/scripts/rl/gym_manipulator.py index 206fc3a05..5fb4dc697 100644 --- a/src/lerobot/scripts/rl/gym_manipulator.py +++ b/src/lerobot/scripts/rl/gym_manipulator.py @@ -32,6 +32,7 @@ from lerobot.processor import ( AddTeleopActionAsComplimentaryData, AddTeleopEventsAsInfo, DeviceProcessor, + EnvTransition, GripperPenaltyProcessor, ImageCropResizeProcessor, InterventionActionProcessor, @@ -45,10 +46,10 @@ from lerobot.processor import ( TimeLimitProcessor, ToBatchProcessor, Torch2NumpyActionProcessor, + TransitionKey, VanillaObservationProcessor, create_transition, ) -from lerobot.processor.pipeline import EnvTransition, TransitionKey from lerobot.robots import ( # noqa: F401 RobotConfig, make_robot_from_config, diff --git a/src/lerobot/teleoperate.py b/src/lerobot/teleoperate.py index d8101f0b3..97b34b310 100644 --- a/src/lerobot/teleoperate.py +++ b/src/lerobot/teleoperate.py @@ -61,13 +61,12 @@ import rerun as rr from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 from lerobot.configs import parser -from lerobot.processor import RobotProcessor +from lerobot.processor import IdentityProcessor, RobotProcessor from lerobot.processor.converters import ( to_output_robot_action, to_transition_robot_observation, to_transition_teleop_action, ) -from lerobot.processor.pipeline import IdentityProcessor from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, diff --git a/src/lerobot/teleoperators/phone/phone_processor.py b/src/lerobot/teleoperators/phone/phone_processor.py index 36880e0c8..4b8c1809a 100644 --- a/src/lerobot/teleoperators/phone/phone_processor.py +++ b/src/lerobot/teleoperators/phone/phone_processor.py @@ -17,7 +17,7 @@ from dataclasses import dataclass, field from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry +from lerobot.processor import ActionProcessor, ProcessorStepRegistry from lerobot.teleoperators.phone.config_phone import PhoneOS diff --git a/src/lerobot/utils/train_utils.py b/src/lerobot/utils/train_utils.py index 1067fa619..25cdeb23b 100644 --- a/src/lerobot/utils/train_utils.py +++ b/src/lerobot/utils/train_utils.py @@ -32,7 +32,7 @@ from lerobot.datasets.utils import load_json, write_json from lerobot.optim.optimizers import load_optimizer_state, save_optimizer_state from lerobot.optim.schedulers import load_scheduler_state, save_scheduler_state from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.processor.pipeline import RobotProcessor +from lerobot.processor import RobotProcessor from lerobot.utils.random_utils import load_rng_state, save_rng_state diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index 8a4f65a03..1bb5e65cd 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -19,7 +19,7 @@ from typing import Any import numpy as np import rerun as rr -from lerobot.processor.pipeline import EnvTransition, TransitionKey +from lerobot.processor import EnvTransition, TransitionKey def _init_rerun(session_name: str = "lerobot_control_loop") -> None: diff --git a/tests/processor/test_act_processor.py b/tests/processor/test_act_processor.py index 6550b007e..f405da419 100644 --- a/tests/processor/test_act_processor.py +++ b/tests/processor/test_act_processor.py @@ -30,9 +30,9 @@ from lerobot.processor import ( RenameProcessor, RobotProcessor, ToBatchProcessor, + TransitionKey, UnnormalizerProcessor, ) -from lerobot.processor.pipeline import TransitionKey def create_transition(observation=None, action=None, **kwargs): diff --git a/tests/processor/test_batch_conversion.py b/tests/processor/test_batch_conversion.py index eb09383a8..9be80821a 100644 --- a/tests/processor/test_batch_conversion.py +++ b/tests/processor/test_batch_conversion.py @@ -1,7 +1,7 @@ import torch +from lerobot.processor import RobotProcessor, TransitionKey from lerobot.processor.converters import batch_to_transition, transition_to_batch -from lerobot.processor.pipeline import RobotProcessor, TransitionKey def _dummy_batch(): diff --git a/tests/processor/test_batch_processor.py b/tests/processor/test_batch_processor.py index 881e05e0f..681ee9a60 100644 --- a/tests/processor/test_batch_processor.py +++ b/tests/processor/test_batch_processor.py @@ -22,9 +22,7 @@ import pytest import torch from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE -from lerobot.processor import ProcessorStepRegistry, RobotProcessor -from lerobot.processor.batch_processor import ToBatchProcessor -from lerobot.processor.pipeline import TransitionKey +from lerobot.processor import ProcessorStepRegistry, RobotProcessor, ToBatchProcessor, TransitionKey def create_transition( diff --git a/tests/processor/test_classifier_processor.py b/tests/processor/test_classifier_processor.py index 3b7417538..1ec62eda4 100644 --- a/tests/processor/test_classifier_processor.py +++ b/tests/processor/test_classifier_processor.py @@ -24,8 +24,13 @@ from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.constants import OBS_IMAGE, OBS_STATE from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor -from lerobot.processor import DeviceProcessor, IdentityProcessor, NormalizerProcessor, RobotProcessor -from lerobot.processor.pipeline import TransitionKey +from lerobot.processor import ( + DeviceProcessor, + IdentityProcessor, + NormalizerProcessor, + RobotProcessor, + TransitionKey, +) def create_transition(observation=None, action=None, **kwargs): diff --git a/tests/processor/test_converters.py b/tests/processor/test_converters.py index a7b5a7b8f..23acaf11e 100644 --- a/tests/processor/test_converters.py +++ b/tests/processor/test_converters.py @@ -2,6 +2,7 @@ import numpy as np import pytest import torch +from lerobot.processor import TransitionKey from lerobot.processor.converters import ( batch_to_transition, to_output_robot_action, @@ -11,7 +12,6 @@ from lerobot.processor.converters import ( transition_to_batch, transition_to_dataset_frame, ) -from lerobot.processor.pipeline import TransitionKey def test_to_transition_teleop_action_prefix_and_tensor_conversion(): diff --git a/tests/processor/test_device_processor.py b/tests/processor/test_device_processor.py index 8449881cf..0482b4916 100644 --- a/tests/processor/test_device_processor.py +++ b/tests/processor/test_device_processor.py @@ -19,8 +19,7 @@ import pytest import torch from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.processor import DeviceProcessor, RobotProcessor -from lerobot.processor.pipeline import TransitionKey +from lerobot.processor import DeviceProcessor, RobotProcessor, TransitionKey def create_transition( @@ -358,7 +357,7 @@ def test_save_and_load_pretrained(): def test_registry_functionality(): """Test that DeviceProcessor is properly registered.""" - from lerobot.processor.pipeline import ProcessorStepRegistry + from lerobot.processor import ProcessorStepRegistry # Check that DeviceProcessor is registered registered_class = ProcessorStepRegistry.get("device_processor") diff --git a/tests/processor/test_diffusion_processor.py b/tests/processor/test_diffusion_processor.py index b075a8d87..0efe0e5e1 100644 --- a/tests/processor/test_diffusion_processor.py +++ b/tests/processor/test_diffusion_processor.py @@ -30,9 +30,9 @@ from lerobot.processor import ( RenameProcessor, RobotProcessor, ToBatchProcessor, + TransitionKey, UnnormalizerProcessor, ) -from lerobot.processor.pipeline import TransitionKey def create_transition(observation=None, action=None, **kwargs): diff --git a/tests/processor/test_normalize_processor.py b/tests/processor/test_normalize_processor.py index 13bf14192..816a9f36a 100644 --- a/tests/processor/test_normalize_processor.py +++ b/tests/processor/test_normalize_processor.py @@ -20,13 +20,15 @@ import pytest import torch from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.processor.converters import to_tensor -from lerobot.processor.normalize_processor import ( +from lerobot.processor import ( + IdentityProcessor, NormalizerProcessor, + RobotProcessor, + TransitionKey, UnnormalizerProcessor, hotswap_stats, ) -from lerobot.processor.pipeline import IdentityProcessor, RobotProcessor, TransitionKey +from lerobot.processor.converters import to_tensor def create_transition( diff --git a/tests/processor/test_observation_processor.py b/tests/processor/test_observation_processor.py index 4e6efdb6c..0102820d5 100644 --- a/tests/processor/test_observation_processor.py +++ b/tests/processor/test_observation_processor.py @@ -20,8 +20,7 @@ import torch from lerobot.configs.types import FeatureType from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE -from lerobot.processor import VanillaObservationProcessor -from lerobot.processor.pipeline import TransitionKey +from lerobot.processor import TransitionKey, VanillaObservationProcessor from tests.conftest import assert_contract_is_typed diff --git a/tests/processor/test_pi0_processor.py b/tests/processor/test_pi0_processor.py index 6f8564277..759704014 100644 --- a/tests/processor/test_pi0_processor.py +++ b/tests/processor/test_pi0_processor.py @@ -29,9 +29,9 @@ from lerobot.processor import ( NormalizerProcessor, RenameProcessor, ToBatchProcessor, + TransitionKey, UnnormalizerProcessor, ) -from lerobot.processor.pipeline import TransitionKey def create_transition(observation=None, action=None, **kwargs): diff --git a/tests/processor/test_pipeline.py b/tests/processor/test_pipeline.py index 4f5813024..e1b61e0df 100644 --- a/tests/processor/test_pipeline.py +++ b/tests/processor/test_pipeline.py @@ -27,8 +27,7 @@ import torch.nn as nn from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features -from lerobot.processor import EnvTransition, ProcessorStepRegistry, RobotProcessor -from lerobot.processor.pipeline import TransitionKey +from lerobot.processor import EnvTransition, ProcessorStepRegistry, RobotProcessor, TransitionKey from tests.conftest import assert_contract_is_typed diff --git a/tests/processor/test_rename_processor.py b/tests/processor/test_rename_processor.py index ae8cc123e..6f8dfd3c5 100644 --- a/tests/processor/test_rename_processor.py +++ b/tests/processor/test_rename_processor.py @@ -20,7 +20,12 @@ import numpy as np import torch from lerobot.configs.types import FeatureType -from lerobot.processor import ProcessorStepRegistry, RenameProcessor, RobotProcessor, TransitionKey +from lerobot.processor import ( + ProcessorStepRegistry, + RenameProcessor, + RobotProcessor, + TransitionKey, +) from lerobot.processor.rename_processor import rename_stats from tests.conftest import assert_contract_is_typed diff --git a/tests/processor/test_sac_processor.py b/tests/processor/test_sac_processor.py index ece7e4b18..c2e50f6c5 100644 --- a/tests/processor/test_sac_processor.py +++ b/tests/processor/test_sac_processor.py @@ -30,9 +30,9 @@ from lerobot.processor import ( RenameProcessor, RobotProcessor, ToBatchProcessor, + TransitionKey, UnnormalizerProcessor, ) -from lerobot.processor.pipeline import TransitionKey def create_transition(observation=None, action=None, **kwargs): diff --git a/tests/processor/test_smolvla_processor.py b/tests/processor/test_smolvla_processor.py index fb88df2ce..e4d673e62 100644 --- a/tests/processor/test_smolvla_processor.py +++ b/tests/processor/test_smolvla_processor.py @@ -32,9 +32,9 @@ from lerobot.processor import ( NormalizerProcessor, RenameProcessor, ToBatchProcessor, + TransitionKey, UnnormalizerProcessor, ) -from lerobot.processor.pipeline import TransitionKey def create_transition(observation=None, action=None, **kwargs): diff --git a/tests/processor/test_tdmpc_processor.py b/tests/processor/test_tdmpc_processor.py index 035e3c93a..6e2259e83 100644 --- a/tests/processor/test_tdmpc_processor.py +++ b/tests/processor/test_tdmpc_processor.py @@ -30,9 +30,9 @@ from lerobot.processor import ( RenameProcessor, RobotProcessor, ToBatchProcessor, + TransitionKey, UnnormalizerProcessor, ) -from lerobot.processor.pipeline import TransitionKey def create_transition(observation=None, action=None, **kwargs): diff --git a/tests/processor/test_tokenizer_processor.py b/tests/processor/test_tokenizer_processor.py index 14b5c6ee6..022767de3 100644 --- a/tests/processor/test_tokenizer_processor.py +++ b/tests/processor/test_tokenizer_processor.py @@ -10,8 +10,7 @@ import torch from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.constants import OBS_LANGUAGE -from lerobot.processor.pipeline import RobotProcessor, TransitionKey -from lerobot.processor.tokenizer_processor import TokenizerProcessor +from lerobot.processor import RobotProcessor, TokenizerProcessor, TransitionKey from tests.utils import require_package @@ -488,7 +487,7 @@ def test_save_and_load_pretrained_with_tokenizer_object(): @require_package("transformers") def test_registry_functionality(): """Test that the processor is properly registered.""" - from lerobot.processor.pipeline import ProcessorStepRegistry + from lerobot.processor import ProcessorStepRegistry # Check that the processor is registered assert "tokenizer_processor" in ProcessorStepRegistry.list() diff --git a/tests/processor/test_vqbet_processor.py b/tests/processor/test_vqbet_processor.py index d46b9f9a4..bb57d8cd0 100644 --- a/tests/processor/test_vqbet_processor.py +++ b/tests/processor/test_vqbet_processor.py @@ -30,9 +30,9 @@ from lerobot.processor import ( RenameProcessor, RobotProcessor, ToBatchProcessor, + TransitionKey, UnnormalizerProcessor, ) -from lerobot.processor.pipeline import TransitionKey def create_transition(observation=None, action=None, **kwargs): diff --git a/tests/utils/test_visualization_utils.py b/tests/utils/test_visualization_utils.py index 5e1eb4bab..bcdfff5e2 100644 --- a/tests/utils/test_visualization_utils.py +++ b/tests/utils/test_visualization_utils.py @@ -5,7 +5,7 @@ from types import SimpleNamespace import numpy as np import pytest -from lerobot.processor.pipeline import TransitionKey +from lerobot.processor import TransitionKey @pytest.fixture