big imports refactor

pyproject.toml

@@ -98,7 +98,6 @@ training = [
     "lerobot[dataset]",
     "accelerate>=1.10.0,<2.0.0",
     "wandb>=0.24.0,<0.25.0",
-    "diffusers>=0.27.2,<0.36.0",
 ]
 hardware = [
     "pynput>=1.7.8,<1.9.0",
@@ -116,7 +115,7 @@ build = [
 # ── User-facing composite extras (map to CLI scripts) ─────
 # lerobot-record, lerobot-replay, lerobot-calibrate, lerobot-teleoperate, etc.
 robot = ["lerobot[dataset]", "lerobot[hardware]", "lerobot[viz]"]
-# lerobot-eval
+# lerobot-eval -- base evaluation framework. You also need the policy's extra (e.g., lerobot[pi]).
 evaluation = ["lerobot[av-dep]"]
 # lerobot-dataset-viz, lerobot-imgtransform-viz
 dataset_viz = ["lerobot[dataset]", "lerobot[viz]"]
@@ -130,6 +129,7 @@ grpcio-dep = ["grpcio==1.73.1", "protobuf>=6.31.1,<6.32.0"]
 can-dep = ["python-can>=4.2.0,<5.0.0"]
 peft-dep = ["peft>=0.18.0,<1.0.0"]
 scipy-dep = ["scipy>=1.14.0,<2.0.0"]
+diffusers-dep = ["diffusers>=0.27.2,<0.36.0"]
 qwen-vl-utils-dep = ["qwen-vl-utils>=0.0.11,<0.1.0"]
 matplotlib-dep = ["matplotlib>=3.10.3,<4.0.0", "contourpy>=1.3.0,<2.0.0"] # NOTE: Explicitly listing contourpy helps the resolver converge faster.
 
@@ -162,6 +162,7 @@ intelrealsense = [
 phone = ["hebi-py>=2.8.0,<2.12.0", "teleop>=0.1.0,<0.2.0", "fastapi<1.0", "lerobot[scipy-dep]"]
 
 # Policies
+diffusion = ["lerobot[diffusers-dep]"]
 wallx = [
     "lerobot[transformers-dep]",
     "lerobot[peft]",
@@ -171,10 +172,11 @@ wallx = [
 ]
 pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]"]
 smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "accelerate>=1.7.0,<2.0.0"]
-multi_task_dit = ["lerobot[transformers-dep]"]
+multi_task_dit = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]"]
 groot = [
     "lerobot[transformers-dep]",
     "lerobot[peft]",
+    "lerobot[diffusers-dep]",
     "dm-tree>=0.1.8,<1.0.0",
     "timm>=1.0.0,<1.1.0",
     "Pillow>=10.0.0,<13.0.0",
@@ -222,6 +224,7 @@ all = [
     "lerobot[reachy2]",
     "lerobot[kinematics]",
     "lerobot[intelrealsense]",
+    "lerobot[diffusion]",
     "lerobot[wallx]",
     "lerobot[pi]",
     "lerobot[smolvla]",
@@ -299,6 +302,11 @@ ignore = [
 
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401", "F403", "E402"]
+# E402: require_package() guards must precede the imports they protect
+"src/lerobot/policies/wall_x/modeling_wall_x.py" = ["E402"]
+"src/lerobot/policies/smolvla/smolvlm_with_expert.py" = ["E402"]
+"src/lerobot/policies/xvla/modeling_xvla.py" = ["E402"]
+"src/lerobot/policies/sarm/processor_sarm.py" = ["E402"]
 "src/lerobot/policies/wall_x/**" = ["N801", "N812", "SIM102", "SIM108", "SIM210", "SIM211", "B006", "B007", "SIM118"] # Supprese these as they are coming from original Qwen2_5_vl code TODO(pepijn): refactor original
 
 [tool.ruff.lint.isort]

src/lerobot/async_inference/__init__.py

@@ -0,0 +1,6 @@
+# Async inference server/client.
+# Requires: lerobot[async]
+
+from lerobot.utils.import_utils import require_package
+
+require_package("grpcio", extra="async", import_name="grpc")

src/lerobot/async_inference/helpers.py

@@ -22,7 +22,7 @@ from typing import Any
 
 import torch
 
-from lerobot.configs.types import PolicyFeature
+from lerobot.configs import PolicyFeature
 
 # NOTE: Configs need to be loaded for the client to be able to instantiate the policy config
 from lerobot.policies import (  # noqa: F401

src/lerobot/async_inference/policy_server.py

@@ -38,7 +38,7 @@ import draccus
 import grpc
 import torch
 
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
+from lerobot.policies import get_policy_class, make_pre_post_processors
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.transport import (
     services_pb2,  # type: ignore

src/lerobot/async_inference/robot_client.py

@@ -47,8 +47,8 @@ import draccus
 import grpc
 import torch
 
-from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
+from lerobot.cameras.opencv import OpenCVCameraConfig  # noqa: F401
-from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
+from lerobot.cameras.realsense import RealSenseCameraConfig  # noqa: F401
 from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,

src/lerobot/cameras/__init__.py

@@ -15,3 +15,9 @@
 from .camera import Camera
 from .configs import CameraConfig, ColorMode, Cv2Backends, Cv2Rotation
 from .utils import make_cameras_from_configs
+
+# NOTE: Camera submodule configs and implementations (OpenCVCameraConfig, RealSenseCamera, etc.)
+# are intentionally NOT re-exported here to avoid pulling backend-specific dependencies.
+# Import from submodules: ``from lerobot.cameras.opencv import OpenCVCameraConfig``
+
+__all__ = ["Camera", "CameraConfig", "ColorMode", "Cv2Backends", "Cv2Rotation", "make_cameras_from_configs"]

src/lerobot/cameras/zmq/image_server.py

@@ -31,8 +31,8 @@ import cv2
 import numpy as np
 import zmq
 
-from lerobot.cameras.configs import ColorMode
+from ..configs import ColorMode
-from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
+from ..opencv import OpenCVCamera, OpenCVCameraConfig
 
 logger = logging.getLogger(__name__)
 

src/lerobot/configs/__init__.py

@@ -0,0 +1,47 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Public API for lerobot configuration types and base config classes.
+
+NOTE: TrainPipelineConfig, EvalPipelineConfig, and TrainRLServerPipelineConfig
+are intentionally NOT re-exported here to avoid circular dependencies
+(they import lerobot.envs and lerobot.policies at module level).
+Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
+"""
+
+from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
+from .policies import PreTrainedConfig
+from .types import (
+    FeatureType,
+    NormalizationMode,
+    PipelineFeatureType,
+    PolicyFeature,
+    RTCAttentionSchedule,
+)
+
+__all__ = [
+    # Types
+    "FeatureType",
+    "NormalizationMode",
+    "PipelineFeatureType",
+    "PolicyFeature",
+    "RTCAttentionSchedule",
+    # Config classes
+    "DatasetConfig",
+    "EvalConfig",
+    "PeftConfig",
+    "PreTrainedConfig",
+    "WandBConfig",
+]

src/lerobot/configs/eval.py

@@ -19,8 +19,9 @@ from pathlib import Path
 
 from lerobot import envs, policies  # noqa: F401
 from lerobot.configs import parser
-from lerobot.configs.default import EvalConfig
+
-from lerobot.configs.policies import PreTrainedConfig
+from .default import EvalConfig
+from .policies import PreTrainedConfig
 
 logger = getLogger(__name__)
 

src/lerobot/configs/policies.py

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

src/lerobot/configs/train.py

@@ -24,12 +24,12 @@ from huggingface_hub.errors import HfHubHTTPError
 
 from lerobot import envs
 from lerobot.configs import parser
-from lerobot.configs.default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
+from lerobot.optim import LRSchedulerConfig, OptimizerConfig
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.optim import OptimizerConfig
-from lerobot.optim.schedulers import LRSchedulerConfig
 from lerobot.utils.hub import HubMixin
 
+from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
+from .policies import PreTrainedConfig
+
 TRAIN_CONFIG_NAME = "train_config.json"
 
 

src/lerobot/data_processing/sarm_annotations/subtask_annotation.py

@@ -76,7 +76,7 @@ import torch
 from pydantic import BaseModel, Field
 from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
 
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets import LeRobotDataset
 
 
 # Pydantic Models for SARM Subtask Annotation

src/lerobot/datasets/__init__.py

@@ -19,18 +19,33 @@ from lerobot.utils.import_utils import require_package
 
 require_package("datasets", extra="dataset")
 
-from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+from .compute_stats import aggregate_stats, get_feature_stats
-from lerobot.datasets.factory import make_dataset
+from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
-from lerobot.datasets.image_writer import safe_stop_image_writer
+from .dataset_tools import (
-from lerobot.datasets.io_utils import load_episodes, write_stats
+    convert_image_to_video_dataset,
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
+    delete_episodes,
-from lerobot.datasets.multi_dataset import MultiLeRobotDataset
+    merge_datasets,
-from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
+    modify_tasks,
-from lerobot.datasets.sampler import EpisodeAwareSampler
+    recompute_stats,
-from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
+    remove_feature,
-from lerobot.datasets.video_utils import VideoEncodingManager
+    split_dataset,
+)
+from .factory import make_dataset
+from .image_writer import safe_stop_image_writer
+from .io_utils import load_episodes, write_stats
+from .lerobot_dataset import LeRobotDataset
+from .multi_dataset import MultiLeRobotDataset
+from .pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
+from .sampler import EpisodeAwareSampler
+from .streaming_dataset import StreamingLeRobotDataset
+from .video_utils import VideoEncodingManager
+
+# NOTE: Low-level I/O functions (cast_stats_to_numpy, get_parquet_file_size_in_mb, etc.)
+# and legacy migration constants are intentionally NOT re-exported here.
+# Import directly: ``from lerobot.datasets.io_utils import ...``
 
 __all__ = [
+    "CODEBASE_VERSION",
     "EpisodeAwareSampler",
     "LeRobotDataset",
     "LeRobotDatasetMetadata",
@@ -38,9 +53,18 @@ __all__ = [
     "StreamingLeRobotDataset",
     "VideoEncodingManager",
     "aggregate_pipeline_dataset_features",
+    "aggregate_stats",
+    "convert_image_to_video_dataset",
     "create_initial_features",
+    "delete_episodes",
+    "get_feature_stats",
     "load_episodes",
     "make_dataset",
+    "merge_datasets",
+    "modify_tasks",
+    "recompute_stats",
+    "remove_feature",
     "safe_stop_image_writer",
+    "split_dataset",
     "write_stats",
 ]

src/lerobot/datasets/aggregate.py

@@ -23,10 +23,10 @@ import datasets
 import pandas as pd
 import tqdm
 
-from lerobot.datasets.compute_stats import aggregate_stats
+from .compute_stats import aggregate_stats
-from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+from .dataset_metadata import LeRobotDatasetMetadata
-from lerobot.datasets.feature_utils import get_hf_features_from_features
+from .feature_utils import get_hf_features_from_features
-from lerobot.datasets.io_utils import (
+from .io_utils import (
     get_file_size_in_mb,
     get_parquet_file_size_in_mb,
     to_parquet_with_hf_images,
@@ -34,7 +34,7 @@ from lerobot.datasets.io_utils import (
     write_stats,
     write_tasks,
 )
-from lerobot.datasets.utils import (
+from .utils import (
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_DATA_PATH,
@@ -43,7 +43,7 @@ from lerobot.datasets.utils import (
     DEFAULT_VIDEO_PATH,
     update_chunk_file_indices,
 )
-from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
+from .video_utils import concatenate_video_files, get_video_duration_in_s
 
 
 def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):

src/lerobot/datasets/compute_stats.py

@@ -19,9 +19,11 @@ import logging
 
 import numpy as np
 
-from lerobot.datasets.io_utils import load_image_as_numpy
+from lerobot.processor import RelativeActionsProcessorStep
 from lerobot.utils.constants import ACTION, OBS_STATE
 
+from .io_utils import load_image_as_numpy
+
 DEFAULT_QUANTILES = [0.01, 0.10, 0.50, 0.90, 0.99]
 
 
@@ -696,8 +698,6 @@ def compute_relative_action_stats(
         ValueError: If the dataset has fewer frames than ``chunk_size``.
         RuntimeError: If no valid (single-episode) chunks are found.
     """
-    from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep
-
     if exclude_joints is None:
         exclude_joints = []
 

src/lerobot/datasets/dataset_metadata.py

@@ -23,9 +23,12 @@ import pyarrow as pa
 import pyarrow.parquet as pq
 from huggingface_hub import snapshot_download
 
-from lerobot.datasets.compute_stats import aggregate_stats
+from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE
-from lerobot.datasets.feature_utils import create_empty_dataset_info
+from lerobot.utils.feature_utils import _validate_feature_names
-from lerobot.datasets.io_utils import (
+
+from .compute_stats import aggregate_stats
+from .feature_utils import create_empty_dataset_info
+from .io_utils import (
     get_file_size_in_mb,
     load_episodes,
     load_info,
@@ -37,7 +40,7 @@ from lerobot.datasets.io_utils import (
     write_stats,
     write_tasks,
 )
-from lerobot.datasets.utils import (
+from .utils import (
     DEFAULT_EPISODES_PATH,
     INFO_PATH,
     check_version_compatibility,
@@ -47,9 +50,7 @@ from lerobot.datasets.utils import (
     is_valid_version,
     update_chunk_file_indices,
 )
-from lerobot.datasets.video_utils import get_video_info
+from .video_utils import get_video_info
-from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE
-from lerobot.utils.feature_utils import _validate_feature_names
 
 CODEBASE_VERSION = "v3.0"
 

src/lerobot/datasets/dataset_reader.py

@@ -21,17 +21,17 @@ from pathlib import Path
 import datasets
 import torch
 
-from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+from .dataset_metadata import LeRobotDatasetMetadata
-from lerobot.datasets.feature_utils import (
+from .feature_utils import (
     check_delta_timestamps,
     get_delta_indices,
     get_hf_features_from_features,
 )
-from lerobot.datasets.io_utils import (
+from .io_utils import (
     hf_transform_to_torch,
     load_nested_dataset,
 )
-from lerobot.datasets.video_utils import decode_video_frames
+from .video_utils import decode_video_frames
 
 
 class DatasetReader:

src/lerobot/datasets/dataset_tools.py

@@ -36,22 +36,25 @@ import pyarrow.parquet as pq
 import torch
 from tqdm import tqdm
 
-from lerobot.datasets.aggregate import aggregate_datasets
+from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE
-from lerobot.datasets.compute_stats import (
+from lerobot.utils.utils import flatten_dict
+
+from .aggregate import aggregate_datasets
+from .compute_stats import (
     aggregate_stats,
     compute_episode_stats,
     compute_relative_action_stats,
 )
-from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+from .dataset_metadata import LeRobotDatasetMetadata
-from lerobot.datasets.io_utils import (
+from .io_utils import (
     get_parquet_file_size_in_mb,
     load_episodes,
     write_info,
     write_stats,
     write_tasks,
 )
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from .lerobot_dataset import LeRobotDataset
-from lerobot.datasets.utils import (
+from .utils import (
     DATA_DIR,
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
@@ -59,8 +62,7 @@ from lerobot.datasets.utils import (
     DEFAULT_EPISODES_PATH,
     update_chunk_file_indices,
 )
-from lerobot.datasets.video_utils import encode_video_frames, get_video_info
+from .video_utils import encode_video_frames, get_video_info
-from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE
 
 
 def _load_episode_with_stats(src_dataset: LeRobotDataset, episode_idx: int) -> dict:
@@ -829,8 +831,6 @@ def _copy_and_reindex_episodes_metadata(
         data_metadata: Dict mapping new episode index to its data file metadata
         video_metadata: Optional dict mapping new episode index to its video metadata
     """
-    from lerobot.utils.utils import flatten_dict
-
     if src_dataset.meta.episodes is None:
         src_dataset.meta.episodes = load_episodes(src_dataset.meta.root)
 
@@ -922,8 +922,8 @@ def _write_parquet(df: pd.DataFrame, path: Path, meta: LeRobotDatasetMetadata) -
 
     This ensures images are properly embedded and the file can be loaded correctly by HF datasets.
     """
-    from lerobot.datasets.feature_utils import get_hf_features_from_features
+    from .feature_utils import get_hf_features_from_features
-    from lerobot.datasets.io_utils import embed_images
+    from .io_utils import embed_images
 
     hf_features = get_hf_features_from_features(meta.features)
     ep_dataset = datasets.Dataset.from_dict(df.to_dict(orient="list"), features=hf_features, split="train")
@@ -1367,7 +1367,7 @@ def _copy_data_without_images(
         episode_indices: Episodes to include
         img_keys: Image keys to remove
     """
-    from lerobot.datasets.utils import DATA_DIR
+    from .utils import DATA_DIR
 
     data_dir = src_dataset.root / DATA_DIR
     parquet_files = sorted(data_dir.glob("*/*.parquet"))

src/lerobot/datasets/dataset_writer.py

@@ -31,26 +31,26 @@ import PIL.Image
 import pyarrow.parquet as pq
 import torch
 
-from lerobot.datasets.compute_stats import compute_episode_stats
+from .compute_stats import compute_episode_stats
-from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+from .dataset_metadata import LeRobotDatasetMetadata
-from lerobot.datasets.feature_utils import (
+from .feature_utils import (
     get_hf_features_from_features,
     validate_episode_buffer,
     validate_frame,
 )
-from lerobot.datasets.image_writer import AsyncImageWriter, write_image
+from .image_writer import AsyncImageWriter, write_image
-from lerobot.datasets.io_utils import (
+from .io_utils import (
     embed_images,
     get_file_size_in_mb,
     load_episodes,
     write_info,
 )
-from lerobot.datasets.utils import (
+from .utils import (
     DEFAULT_EPISODES_PATH,
     DEFAULT_IMAGE_PATH,
     update_chunk_file_indices,
 )
-from lerobot.datasets.video_utils import (
+from .video_utils import (
     StreamingVideoEncoder,
     concatenate_video_files,
     encode_video_frames,

src/lerobot/datasets/factory.py

@@ -18,15 +18,16 @@ from pprint import pformat
 
 import torch
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
-from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.multi_dataset import MultiLeRobotDataset
-from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
 from lerobot.transforms import ImageTransforms
 from lerobot.utils.constants import ACTION, IMAGENET_STATS, OBS_PREFIX, REWARD
 
+from .dataset_metadata import LeRobotDatasetMetadata
+from .lerobot_dataset import LeRobotDataset
+from .multi_dataset import MultiLeRobotDataset
+from .streaming_dataset import StreamingLeRobotDataset
+
 
 def resolve_delta_timestamps(
     cfg: PreTrainedConfig, ds_meta: LeRobotDatasetMetadata

src/lerobot/datasets/feature_utils.py

@@ -19,15 +19,16 @@ import datasets
 import numpy as np
 from PIL import Image as PILImage
 
-from lerobot.datasets.utils import (
+from lerobot.utils.constants import DEFAULT_FEATURES
+from lerobot.utils.utils import is_valid_numpy_dtype_string
+
+from .utils import (
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_DATA_PATH,
     DEFAULT_VIDEO_FILE_SIZE_IN_MB,
     DEFAULT_VIDEO_PATH,
 )
-from lerobot.utils.constants import DEFAULT_FEATURES
-from lerobot.utils.utils import is_valid_numpy_dtype_string
 
 
 def get_hf_features_from_features(features: dict) -> datasets.Features:

src/lerobot/datasets/io_utils.py

@@ -28,7 +28,10 @@ from datasets.table import embed_table_storage
 from PIL import Image as PILImage
 from torchvision import transforms
 
-from lerobot.datasets.utils import (
+from lerobot.utils.io_utils import load_json, write_json
+from lerobot.utils.utils import SuppressProgressBars
+
+from .utils import (
     DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_EPISODES_PATH,
     DEFAULT_SUBTASKS_PATH,
@@ -40,8 +43,6 @@ from lerobot.datasets.utils import (
     serialize_dict,
     unflatten_dict,
 )
-from lerobot.utils.io_utils import load_json, write_json
-from lerobot.utils.utils import SuppressProgressBars
 
 
 def get_parquet_file_size_in_mb(parquet_path: str | Path) -> float:

src/lerobot/datasets/lerobot_dataset.py

@@ -24,20 +24,21 @@ import torch.utils
 from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.errors import RevisionNotFoundError
 
-from lerobot.datasets.dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
+from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE
-from lerobot.datasets.dataset_reader import DatasetReader
+
-from lerobot.datasets.dataset_writer import DatasetWriter
+from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
-from lerobot.datasets.utils import (
+from .dataset_reader import DatasetReader
+from .dataset_writer import DatasetWriter
+from .utils import (
     create_lerobot_dataset_card,
     get_safe_version,
     is_valid_version,
 )
-from lerobot.datasets.video_utils import (
+from .video_utils import (
     StreamingVideoEncoder,
     get_safe_default_codec,
     resolve_vcodec,
 )
-from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE
 
 logger = logging.getLogger(__name__)
 

src/lerobot/datasets/multi_dataset.py

@@ -21,12 +21,13 @@ import datasets
 import torch
 import torch.utils
 
-from lerobot.datasets.compute_stats import aggregate_stats
-from lerobot.datasets.feature_utils import get_hf_features_from_features
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.datasets.video_utils import VideoFrame
 from lerobot.utils.constants import HF_LEROBOT_HOME
 
+from .compute_stats import aggregate_stats
+from .feature_utils import get_hf_features_from_features
+from .lerobot_dataset import LeRobotDataset
+from .video_utils import VideoFrame
+
 logger = logging.getLogger(__name__)
 
 

src/lerobot/datasets/pipeline_features.py

@@ -16,7 +16,7 @@ import re
 from collections.abc import Sequence
 from typing import Any
 
-from lerobot.configs.types import PipelineFeatureType
+from lerobot.configs import PipelineFeatureType
 from lerobot.processor import DataProcessorPipeline
 from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR

src/lerobot/datasets/streaming_dataset.py

@@ -22,20 +22,21 @@ import numpy as np
 import torch
 from datasets import load_dataset
 
-from lerobot.datasets.dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
+from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
-from lerobot.datasets.feature_utils import get_delta_indices
+
-from lerobot.datasets.io_utils import item_to_torch
+from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
-from lerobot.datasets.utils import (
+from .feature_utils import get_delta_indices
+from .io_utils import item_to_torch
+from .utils import (
     check_version_compatibility,
     find_float_index,
     is_float_in_list,
     safe_shard,
 )
-from lerobot.datasets.video_utils import (
+from .video_utils import (
     VideoDecoderCache,
     decode_video_frames_torchcodec,
 )
-from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
 
 
 class LookBackError(Exception):

src/lerobot/datasets/video_utils.py

@@ -598,7 +598,7 @@ class _CameraEncoderThread(threading.Thread):
         self.encoder_threads = encoder_threads
 
     def run(self) -> None:
-        from lerobot.datasets.compute_stats import RunningQuantileStats, auto_downsample_height_width
+        from .compute_stats import RunningQuantileStats, auto_downsample_height_width
 
         container = None
         output_stream = None

src/lerobot/envs/__init__.py

@@ -12,4 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configs import AlohaEnv, EnvConfig, HubEnvConfig, PushtEnv  # noqa: F401
+from .configs import AlohaEnv, EnvConfig, HILSerlRobotEnvConfig, HubEnvConfig, PushtEnv
+from .factory import make_env, make_env_config, make_env_pre_post_processors
+from .utils import check_env_attributes_and_types, close_envs, env_to_policy_features, preprocess_observation
+
+__all__ = [
+    "AlohaEnv",
+    "EnvConfig",
+    "HILSerlRobotEnvConfig",
+    "HubEnvConfig",
+    "PushtEnv",
+    "check_env_attributes_and_types",
+    "close_envs",
+    "env_to_policy_features",
+    "make_env",
+    "make_env_config",
+    "make_env_pre_post_processors",
+    "preprocess_observation",
+]

src/lerobot/envs/configs.py

@@ -23,7 +23,8 @@ import draccus
 import gymnasium as gym
 from gymnasium.envs.registration import registry as gym_registry
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs import FeatureType, PolicyFeature
+from lerobot.processor import IsaaclabArenaProcessorStep, LiberoProcessorStep, PolicyProcessorPipeline
 from lerobot.robots import RobotConfig
 from lerobot.teleoperators.config import TeleoperatorConfig
 from lerobot.utils.constants import (
@@ -124,8 +125,6 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
 
     def get_env_processors(self):
         """Return (preprocessor, postprocessor) for this env. Default: identity."""
-        from lerobot.processor.pipeline import PolicyProcessorPipeline
-
         return PolicyProcessorPipeline(steps=[]), PolicyProcessorPipeline(steps=[])
 
 
@@ -418,7 +417,7 @@ class LiberoEnv(EnvConfig):
         return kwargs
 
     def create_envs(self, n_envs: int, use_async_envs: bool = False):
-        from lerobot.envs.libero import create_libero_envs
+        from .libero import create_libero_envs
 
         if self.task is None:
             raise ValueError("LiberoEnv requires a task to be specified")
@@ -436,9 +435,6 @@ class LiberoEnv(EnvConfig):
         )
 
     def get_env_processors(self):
-        from lerobot.processor.env_processor import LiberoProcessorStep
-        from lerobot.processor.pipeline import PolicyProcessorPipeline
-
         return (
             PolicyProcessorPipeline(steps=[LiberoProcessorStep()]),
             PolicyProcessorPipeline(steps=[]),
@@ -487,7 +483,7 @@ class MetaworldEnv(EnvConfig):
         }
 
     def create_envs(self, n_envs: int, use_async_envs: bool = False):
-        from lerobot.envs.metaworld import create_metaworld_envs
+        from .metaworld import create_metaworld_envs
 
         if self.task is None:
             raise ValueError("MetaWorld requires a task to be specified")
@@ -568,9 +564,6 @@ class IsaaclabArenaEnv(HubEnvConfig):
         return {}
 
     def get_env_processors(self):
-        from lerobot.processor.env_processor import IsaaclabArenaProcessorStep
-        from lerobot.processor.pipeline import PolicyProcessorPipeline
-
         state_keys = tuple(k.strip() for k in (self.state_keys or "").split(",") if k.strip())
         camera_keys = tuple(k.strip() for k in (self.camera_keys or "").split(",") if k.strip())
         if not state_keys and not camera_keys:

src/lerobot/envs/factory.py

@@ -19,8 +19,8 @@ from typing import Any
 
 import gymnasium as gym
 
-from lerobot.envs.configs import EnvConfig, HubEnvConfig
+from .configs import EnvConfig, HubEnvConfig
-from lerobot.envs.utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result
+from .utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result
 
 
 def make_env_config(env_type: str, **kwargs) -> EnvConfig:

src/lerobot/envs/libero.py

@@ -29,9 +29,10 @@ from gymnasium import spaces
 from libero.libero import benchmark, get_libero_path
 from libero.libero.envs import OffScreenRenderEnv
 
-from lerobot.envs.utils import _LazyAsyncVectorEnv
 from lerobot.types import RobotObservation
 
+from .utils import _LazyAsyncVectorEnv
+
 
 def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
     """Normalize camera_name into a non-empty list of strings."""

src/lerobot/envs/metaworld.py

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

src/lerobot/envs/utils.py

@@ -27,11 +27,12 @@ import torch
 from huggingface_hub import hf_hub_download, snapshot_download
 from torch import Tensor
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs import FeatureType, PolicyFeature
-from lerobot.envs.configs import EnvConfig
 from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR
 from lerobot.utils.utils import get_channel_first_image_shape
 
+from .configs import EnvConfig
+
 
 def _convert_nested_dict(d):
     result = {}

src/lerobot/model/__init__.py

@@ -0,0 +1,5 @@
+# Kinematics utilities for robot modeling.
+
+from .kinematics import RobotKinematics as RobotKinematics
+
+__all__ = ["RobotKinematics"]

src/lerobot/motors/__init__.py

@@ -19,3 +19,5 @@ from .motors_bus import (
     MotorCalibration,
     MotorNormMode,
 )
+
+__all__ = ["Motor", "MotorCalibration", "MotorNormMode"]

src/lerobot/motors/dynamixel/__init__.py

@@ -14,5 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from lerobot.utils.import_utils import require_package
+
+require_package("dynamixel-sdk", extra="dynamixel", import_name="dynamixel_sdk")
+
 from .dynamixel import DriveMode, DynamixelMotorsBus, OperatingMode, TorqueMode
 from .tables import *

src/lerobot/motors/feetech/__init__.py

@@ -14,5 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from lerobot.utils.import_utils import require_package
+
+require_package("feetech-servo-sdk", extra="feetech", import_name="scservo_sdk")
+
 from .feetech import DriveMode, FeetechMotorsBus, OperatingMode, TorqueMode
 from .tables import *

src/lerobot/optim/__init__.py

@@ -12,4 +12,45 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .optimizers import OptimizerConfig as OptimizerConfig
+from .optimizers import (
+    AdamConfig as AdamConfig,
+    AdamWConfig as AdamWConfig,
+    MultiAdamConfig as MultiAdamConfig,
+    OptimizerConfig as OptimizerConfig,
+    SGDConfig as SGDConfig,
+    XVLAAdamWConfig as XVLAAdamWConfig,
+    load_optimizer_state,
+    save_optimizer_state,
+)
+from .schedulers import (
+    CosineDecayWithWarmupSchedulerConfig as CosineDecayWithWarmupSchedulerConfig,
+    DiffuserSchedulerConfig as DiffuserSchedulerConfig,
+    LRSchedulerConfig as LRSchedulerConfig,
+    VQBeTSchedulerConfig as VQBeTSchedulerConfig,
+    load_scheduler_state,
+    save_scheduler_state,
+)
+
+# NOTE: make_optimizer_and_scheduler is intentionally NOT re-exported here
+# to avoid circular dependencies (it imports lerobot.configs.train and lerobot.policies).
+# Import directly: ``from lerobot.optim.factory import make_optimizer_and_scheduler``
+
+__all__ = [
+    # Optimizer configs
+    "AdamConfig",
+    "AdamWConfig",
+    "MultiAdamConfig",
+    "OptimizerConfig",
+    "SGDConfig",
+    "XVLAAdamWConfig",
+    # Scheduler configs
+    "CosineDecayWithWarmupSchedulerConfig",
+    "DiffuserSchedulerConfig",
+    "LRSchedulerConfig",
+    "VQBeTSchedulerConfig",
+    # State management
+    "load_optimizer_state",
+    "load_scheduler_state",
+    "save_optimizer_state",
+    "save_scheduler_state",
+]

src/lerobot/optim/factory.py

@@ -19,7 +19,7 @@ from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
 from lerobot.configs.train import TrainPipelineConfig
-from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies import PreTrainedPolicy
 
 
 def make_optimizer_and_scheduler(

src/lerobot/optim/schedulers.py

@@ -47,6 +47,9 @@ class DiffuserSchedulerConfig(LRSchedulerConfig):
     num_warmup_steps: int | None = None
 
     def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
+        from lerobot.utils.import_utils import require_package
+
+        require_package("diffusers", extra="diffusion")
         from diffusers.optimization import get_scheduler
 
         kwargs = {**asdict(self), "num_training_steps": num_training_steps, "optimizer": optimizer}

src/lerobot/policies/__init__.py

@@ -14,30 +14,54 @@
 
 from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
+from .factory import get_policy_class, make_policy, make_policy_config, make_pre_post_processors
 from .groot.configuration_groot import GrootConfig as GrootConfig
 from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig as MultiTaskDiTConfig
 from .pi0.configuration_pi0 import PI0Config as PI0Config
 from .pi0_fast.configuration_pi0_fast import PI0FastConfig as PI0FastConfig
 from .pi05.configuration_pi05 import PI05Config as PI05Config
+from .pretrained import PreTrainedPolicy as PreTrainedPolicy
+from .rtc import ActionInterpolator as ActionInterpolator
+from .sarm.configuration_sarm import SARMConfig as SARMConfig
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
-from .smolvla.processor_smolvla import SmolVLANewLineProcessor
+from .smolvla.processor_smolvla import (
+    SmolVLANewLineProcessor as SmolVLANewLineProcessor,  # noqa: F401 - registers with ProcessorStepRegistry
+)
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
+from .utils import make_robot_action, prepare_observation_for_inference
 from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
 from .wall_x.configuration_wall_x import WallXConfig as WallXConfig
 from .xvla.configuration_xvla import XVLAConfig as XVLAConfig
 
+# NOTE: Policy modeling classes (e.g., SACPolicy) are intentionally NOT re-exported here.
+# They have heavy optional dependencies and are loaded lazily via get_policy_class().
+# Import directly: ``from lerobot.policies.sac.modeling_sac import SACPolicy``
+
 __all__ = [
+    # Configuration classes
     "ACTConfig",
     "DiffusionConfig",
+    "GrootConfig",
     "MultiTaskDiTConfig",
     "PI0Config",
-    "PI05Config",
     "PI0FastConfig",
-    "SmolVLAConfig",
+    "PI05Config",
     "SARMConfig",
+    "SmolVLAConfig",
     "TDMPCConfig",
     "VQBeTConfig",
-    "GrootConfig",
-    "XVLAConfig",
     "WallXConfig",
+    "XVLAConfig",
+    # Base class
+    "PreTrainedPolicy",
+    # RTC utilities
+    "ActionInterpolator",
+    # Utility functions
+    "make_robot_action",
+    "prepare_observation_for_inference",
+    # Factory functions
+    "get_policy_class",
+    "make_policy",
+    "make_policy_config",
+    "make_pre_post_processors",
 ]

src/lerobot/policies/act/configuration_act.py

@@ -15,9 +15,8 @@
 # limitations under the License.
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import NormalizationMode, PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
+from lerobot.optim import AdamWConfig
-from lerobot.optim.optimizers import AdamWConfig
 
 
 @PreTrainedConfig.register_subclass("act")

src/lerobot/policies/act/modeling_act.py

@@ -33,10 +33,11 @@ from torch import Tensor, nn
 from torchvision.models._utils import IntermediateLayerGetter
 from torchvision.ops.misc import FrozenBatchNorm2d
 
-from lerobot.policies.act.configuration_act import ACTConfig
-from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 
+from ..pretrained import PreTrainedPolicy
+from .configuration_act import ACTConfig
+
 
 class ACTPolicy(PreTrainedPolicy):
     """

src/lerobot/policies/act/processor_act.py

@@ -17,7 +17,6 @@ from typing import Any
 
 import torch
 
-from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -26,10 +25,13 @@ from lerobot.processor import (
     PolicyProcessorPipeline,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
+from .configuration_act import ACTConfig
+
 
 def make_act_pre_post_processors(
     config: ACTConfig,

src/lerobot/policies/diffusion/configuration_diffusion.py

@@ -16,10 +16,8 @@
 # limitations under the License.
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import NormalizationMode, PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
+from lerobot.optim import AdamConfig, DiffuserSchedulerConfig
-from lerobot.optim.optimizers import AdamConfig
-from lerobot.optim.schedulers import DiffuserSchedulerConfig
 
 
 @PreTrainedConfig.register_subclass("diffusion")

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -31,15 +31,16 @@ import torch.nn.functional as F  # noqa: N812
 import torchvision
 from torch import Tensor, nn
 
-from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
+from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
-from lerobot.policies.pretrained import PreTrainedPolicy
+
-from lerobot.policies.utils import (
+from ..pretrained import PreTrainedPolicy
+from ..utils import (
     get_device_from_parameters,
     get_dtype_from_parameters,
     get_output_shape,
     populate_queues,
 )
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
+from .configuration_diffusion import DiffusionConfig
 
 
 class DiffusionPolicy(PreTrainedPolicy):
@@ -156,7 +157,7 @@ def _make_noise_scheduler(name: str, **kwargs: dict):
     """
     from lerobot.utils.import_utils import require_package
 
-    require_package("diffusers", extra="training")
+    require_package("diffusers", extra="diffusion")
     from diffusers.schedulers.scheduling_ddim import DDIMScheduler
     from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
 

src/lerobot/policies/diffusion/processor_diffusion.py

@@ -18,7 +18,6 @@ from typing import Any
 
 import torch
 
-from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -27,10 +26,13 @@ from lerobot.processor import (
     PolicyProcessorPipeline,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
+from .configuration_diffusion import DiffusionConfig
+
 
 def make_diffusion_pre_post_processors(
     config: DiffusionConfig,

src/lerobot/policies/factory.py

@@ -23,30 +23,14 @@ from typing import TYPE_CHECKING, Any, TypedDict, Unpack
 import torch
 
 if TYPE_CHECKING:
-    from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+    from lerobot.datasets import LeRobotDatasetMetadata
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import FeatureType, PreTrainedConfig
-from lerobot.configs.types import FeatureType
+from lerobot.envs import EnvConfig, env_to_policy_features
-from lerobot.envs.configs import EnvConfig
+from lerobot.processor import (
-from lerobot.envs.utils import env_to_policy_features
+    AbsoluteActionsProcessorStep,
-from lerobot.policies.act.configuration_act import ACTConfig
+    PolicyProcessorPipeline,
-from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
+    RelativeActionsProcessorStep,
-from lerobot.policies.groot.configuration_groot import GrootConfig
-from lerobot.policies.multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
-from lerobot.policies.pi0.configuration_pi0 import PI0Config
-from lerobot.policies.pi05.configuration_pi05 import PI05Config
-from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.sac.configuration_sac import SACConfig
-from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
-from lerobot.policies.sarm.configuration_sarm import SARMConfig
-from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
-from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.policies.utils import validate_visual_features_consistency
-from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
-from lerobot.policies.wall_x.configuration_wall_x import WallXConfig
-from lerobot.policies.xvla.configuration_xvla import XVLAConfig
-from lerobot.processor import PolicyProcessorPipeline
-from lerobot.processor.converters import (
     batch_to_transition,
     policy_action_to_transition,
     transition_to_batch,
@@ -58,6 +42,24 @@ from lerobot.utils.constants import (
     POLICY_POSTPROCESSOR_DEFAULT_NAME,
     POLICY_PREPROCESSOR_DEFAULT_NAME,
 )
+from lerobot.utils.feature_utils import dataset_to_policy_features
+
+from .act.configuration_act import ACTConfig
+from .diffusion.configuration_diffusion import DiffusionConfig
+from .groot.configuration_groot import GrootConfig
+from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
+from .pi0.configuration_pi0 import PI0Config
+from .pi05.configuration_pi05 import PI05Config
+from .pretrained import PreTrainedPolicy
+from .sac.configuration_sac import SACConfig
+from .sac.reward_model.configuration_classifier import RewardClassifierConfig
+from .sarm.configuration_sarm import SARMConfig
+from .smolvla.configuration_smolvla import SmolVLAConfig
+from .tdmpc.configuration_tdmpc import TDMPCConfig
+from .utils import validate_visual_features_consistency
+from .vqbet.configuration_vqbet import VQBeTConfig
+from .wall_x.configuration_wall_x import WallXConfig
+from .xvla.configuration_xvla import XVLAConfig
 
 
 def _reconnect_relative_absolute_steps(
@@ -70,11 +72,6 @@ def _reconnect_relative_absolute_steps(
     the RelativeActionsProcessorStep so it can read the cached state at inference time.
     That reference is not serializable, so we re-establish it here after loading.
     """
-    from lerobot.processor.relative_action_processor import (
-        AbsoluteActionsProcessorStep,
-        RelativeActionsProcessorStep,
-    )
-
     relative_step = next((s for s in preprocessor.steps if isinstance(s, RelativeActionsProcessorStep)), None)
     if relative_step is None:
         return
@@ -100,63 +97,63 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         NotImplementedError: If the policy name is not recognized.
     """
     if name == "tdmpc":
-        from lerobot.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
+        from .tdmpc.modeling_tdmpc import TDMPCPolicy
 
         return TDMPCPolicy
     elif name == "diffusion":
-        from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
+        from .diffusion.modeling_diffusion import DiffusionPolicy
 
         return DiffusionPolicy
     elif name == "act":
-        from lerobot.policies.act.modeling_act import ACTPolicy
+        from .act.modeling_act import ACTPolicy
 
         return ACTPolicy
     elif name == "multi_task_dit":
-        from lerobot.policies.multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
+        from .multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
 
         return MultiTaskDiTPolicy
     elif name == "vqbet":
-        from lerobot.policies.vqbet.modeling_vqbet import VQBeTPolicy
+        from .vqbet.modeling_vqbet import VQBeTPolicy
 
         return VQBeTPolicy
     elif name == "pi0":
-        from lerobot.policies.pi0.modeling_pi0 import PI0Policy
+        from .pi0.modeling_pi0 import PI0Policy
 
         return PI0Policy
     elif name == "pi0_fast":
-        from lerobot.policies.pi0_fast.modeling_pi0_fast import PI0FastPolicy
+        from .pi0_fast.modeling_pi0_fast import PI0FastPolicy
 
         return PI0FastPolicy
     elif name == "pi05":
-        from lerobot.policies.pi05.modeling_pi05 import PI05Policy
+        from .pi05.modeling_pi05 import PI05Policy
 
         return PI05Policy
     elif name == "sac":
-        from lerobot.policies.sac.modeling_sac import SACPolicy
+        from .sac.modeling_sac import SACPolicy
 
         return SACPolicy
     elif name == "reward_classifier":
-        from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
+        from .sac.reward_model.modeling_classifier import Classifier
 
         return Classifier
     elif name == "smolvla":
-        from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
+        from .smolvla.modeling_smolvla import SmolVLAPolicy
 
         return SmolVLAPolicy
     elif name == "sarm":
-        from lerobot.policies.sarm.modeling_sarm import SARMRewardModel
+        from .sarm.modeling_sarm import SARMRewardModel
 
         return SARMRewardModel
     elif name == "groot":
-        from lerobot.policies.groot.modeling_groot import GrootPolicy
+        from .groot.modeling_groot import GrootPolicy
 
         return GrootPolicy
     elif name == "xvla":
-        from lerobot.policies.xvla.modeling_xvla import XVLAPolicy
+        from .xvla.modeling_xvla import XVLAPolicy
 
         return XVLAPolicy
     elif name == "wall_x":
-        from lerobot.policies.wall_x.modeling_wall_x import WallXPolicy
+        from .wall_x.modeling_wall_x import WallXPolicy
 
         return WallXPolicy
     else:
@@ -316,7 +313,7 @@ def make_pre_post_processors(
 
     # Create a new processor based on policy type
     if isinstance(policy_cfg, TDMPCConfig):
-        from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
+        from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
 
         processors = make_tdmpc_pre_post_processors(
             config=policy_cfg,
@@ -324,7 +321,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, DiffusionConfig):
-        from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors
+        from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
 
         processors = make_diffusion_pre_post_processors(
             config=policy_cfg,
@@ -332,7 +329,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, ACTConfig):
-        from lerobot.policies.act.processor_act import make_act_pre_post_processors
+        from .act.processor_act import make_act_pre_post_processors
 
         processors = make_act_pre_post_processors(
             config=policy_cfg,
@@ -340,7 +337,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, MultiTaskDiTConfig):
-        from lerobot.policies.multi_task_dit.processor_multi_task_dit import (
+        from .multi_task_dit.processor_multi_task_dit import (
             make_multi_task_dit_pre_post_processors,
         )
 
@@ -350,7 +347,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, VQBeTConfig):
-        from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors
+        from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
 
         processors = make_vqbet_pre_post_processors(
             config=policy_cfg,
@@ -358,7 +355,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, PI0Config):
-        from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
+        from .pi0.processor_pi0 import make_pi0_pre_post_processors
 
         processors = make_pi0_pre_post_processors(
             config=policy_cfg,
@@ -366,7 +363,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, PI05Config):
-        from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors
+        from .pi05.processor_pi05 import make_pi05_pre_post_processors
 
         processors = make_pi05_pre_post_processors(
             config=policy_cfg,
@@ -374,7 +371,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, SACConfig):
-        from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors
+        from .sac.processor_sac import make_sac_pre_post_processors
 
         processors = make_sac_pre_post_processors(
             config=policy_cfg,
@@ -382,7 +379,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, RewardClassifierConfig):
-        from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor
+        from .sac.reward_model.processor_classifier import make_classifier_processor
 
         processors = make_classifier_processor(
             config=policy_cfg,
@@ -390,7 +387,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, SmolVLAConfig):
-        from lerobot.policies.smolvla.processor_smolvla import make_smolvla_pre_post_processors
+        from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
 
         processors = make_smolvla_pre_post_processors(
             config=policy_cfg,
@@ -398,7 +395,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, SARMConfig):
-        from lerobot.policies.sarm.processor_sarm import make_sarm_pre_post_processors
+        from .sarm.processor_sarm import make_sarm_pre_post_processors
 
         processors = make_sarm_pre_post_processors(
             config=policy_cfg,
@@ -406,7 +403,7 @@ def make_pre_post_processors(
             dataset_meta=kwargs.get("dataset_meta"),
         )
     elif isinstance(policy_cfg, GrootConfig):
-        from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
+        from .groot.processor_groot import make_groot_pre_post_processors
 
         processors = make_groot_pre_post_processors(
             config=policy_cfg,
@@ -414,7 +411,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, XVLAConfig):
-        from lerobot.policies.xvla.processor_xvla import (
+        from .xvla.processor_xvla import (
             make_xvla_pre_post_processors,
         )
 
@@ -424,7 +421,7 @@ def make_pre_post_processors(
         )
 
     elif isinstance(policy_cfg, WallXConfig):
-        from lerobot.policies.wall_x.processor_wall_x import make_wall_x_pre_post_processors
+        from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
 
         processors = make_wall_x_pre_post_processors(
             config=policy_cfg,
@@ -495,8 +492,6 @@ def make_policy(
 
     kwargs = {}
     if ds_meta is not None:
-        from lerobot.utils.feature_utils import dataset_to_policy_features
-
         features = dataset_to_policy_features(ds_meta.features)
     else:
         if not cfg.pretrained_path:

src/lerobot/policies/groot/__init__.py

@@ -14,6 +14,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from lerobot.utils.import_utils import require_package
+
+require_package("transformers", extra="groot")
+
 from .configuration_groot import GrootConfig
 from .modeling_groot import GrootPolicy
 from .processor_groot import make_groot_pre_post_processors

src/lerobot/policies/groot/action_head/cross_attention_dit.py

@@ -19,7 +19,7 @@ import torch.nn.functional as F  # noqa: N812
 
 from lerobot.utils.import_utils import require_package
 
-require_package("diffusers", extra="training")
+require_package("diffusers", extra="groot")
 
 from diffusers import ConfigMixin, ModelMixin  # noqa: E402
 from diffusers.configuration_utils import register_to_config  # noqa: E402

src/lerobot/policies/groot/action_head/flow_matching_action_head.py

@@ -31,11 +31,10 @@ else:
     PretrainedConfig = object
     BatchFeature = None
 
-from lerobot.policies.groot.action_head.action_encoder import (
+from .action_encoder import (
     SinusoidalPositionalEncoding,
     swish,
 )
-
 from .cross_attention_dit import DiT, SelfAttentionTransformer
 
 

src/lerobot/policies/groot/configuration_groot.py

@@ -16,10 +16,8 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
-from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
 from lerobot.utils.constants import ACTION, OBS_STATE
 
 

src/lerobot/policies/groot/groot_n1.py

@@ -41,12 +41,13 @@ try:
 except ImportError:
     tree = None
 
-from lerobot.policies.groot.action_head.flow_matching_action_head import (
+from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME
+
+from .action_head.flow_matching_action_head import (
     FlowmatchingActionHead,
     FlowmatchingActionHeadConfig,
 )
-from lerobot.policies.groot.utils import ensure_eagle_cache_ready
+from .utils import ensure_eagle_cache_ready
-from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME
 
 DEFAULT_VENDOR_EAGLE_PATH = str((Path(__file__).resolve().parent / "eagle2_hg_model").resolve())
 DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5"

src/lerobot/policies/groot/modeling_groot.py

@@ -41,12 +41,13 @@ from typing import TypeVar
 import torch
 from torch import Tensor
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs import FeatureType, PolicyFeature
-from lerobot.policies.groot.configuration_groot import GrootConfig
-from lerobot.policies.groot.groot_n1 import GR00TN15
-from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.utils.constants import ACTION, OBS_IMAGES
 
+from ..pretrained import PreTrainedPolicy
+from .configuration_groot import GrootConfig
+from .groot_n1 import GR00TN15
+
 T = TypeVar("T", bound="GrootPolicy")
 
 

src/lerobot/policies/groot/processor_groot.py

@@ -30,12 +30,11 @@ else:
     AutoProcessor = None
     ProcessorMixin = object
 
-from lerobot.configs.types import (
+from lerobot.configs import (
     FeatureType,
     NormalizationMode,
     PolicyFeature,
 )
-from lerobot.policies.groot.configuration_groot import GrootConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -44,8 +43,6 @@ from lerobot.processor import (
     ProcessorStep,
     ProcessorStepRegistry,
     RenameObservationsProcessorStep,
-)
-from lerobot.processor.converters import (
     policy_action_to_transition,
     transition_to_policy_action,
 )
@@ -60,6 +57,8 @@ from lerobot.utils.constants import (
     POLICY_PREPROCESSOR_DEFAULT_NAME,
 )
 
+from .configuration_groot import GrootConfig
+
 # Defaults for Eagle processor locations
 DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5"
 

src/lerobot/policies/multi_task_dit/configuration_multi_task_dit.py

@@ -17,10 +17,8 @@
 import logging
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import NormalizationMode, PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
+from lerobot.optim import AdamConfig, DiffuserSchedulerConfig
-from lerobot.optim.optimizers import AdamConfig
-from lerobot.optim.schedulers import DiffuserSchedulerConfig
 
 
 @PreTrainedConfig.register_subclass("multi_task_dit")

src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py

@@ -36,17 +36,16 @@ import torch.nn.functional as F  # noqa: N812
 import torchvision
 from torch import Tensor
 
-from lerobot.policies.multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
 from lerobot.utils.import_utils import _transformers_available
 
+from .configuration_multi_task_dit import MultiTaskDiTConfig
+
 # Conditional import for type checking and lazy loading
 if TYPE_CHECKING or _transformers_available:
     from transformers import CLIPTextModel, CLIPVisionModel
 else:
     CLIPTextModel = None
     CLIPVisionModel = None
-from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.utils import populate_queues
 from lerobot.utils.constants import (
     ACTION,
     OBS_IMAGES,
@@ -55,6 +54,9 @@ from lerobot.utils.constants import (
     OBS_STATE,
 )
 
+from ..pretrained import PreTrainedPolicy
+from ..utils import populate_queues
+
 # -- Policy --
 
 
@@ -643,7 +645,7 @@ class DiffusionObjective(nn.Module):
 
         from lerobot.utils.import_utils import require_package
 
-        require_package("diffusers", extra="training")
+        require_package("diffusers", extra="multi_task_dit")
         from diffusers.schedulers.scheduling_ddim import DDIMScheduler
         from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
 

src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py

@@ -18,7 +18,6 @@ from typing import Any
 
 import torch
 
-from lerobot.policies.multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -28,10 +27,13 @@ from lerobot.processor import (
     RenameObservationsProcessorStep,
     TokenizerProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
+from .configuration_multi_task_dit import MultiTaskDiTConfig
+
 
 def make_multi_task_dit_pre_post_processors(
     config: MultiTaskDiTConfig,

src/lerobot/policies/pi0/configuration_pi0.py

@@ -16,13 +16,12 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
-from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
+from ..rtc.configuration_rtc import RTCConfig
+
 DEFAULT_IMAGE_SIZE = 224
 
 

src/lerobot/policies/pi0/modeling_pi0.py

@@ -33,7 +33,7 @@ if TYPE_CHECKING or _transformers_available:
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
 
-    from lerobot.policies.pi_gemma import (
+    from ..pi_gemma import (
         PaliGemmaForConditionalGenerationWithPiGemma,
         PiGemmaForCausalLM,
         _gated_residual,
@@ -48,10 +48,7 @@ else:
     PaliGemmaForConditionalGenerationWithPiGemma = None
 
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import PreTrainedConfig
-from lerobot.policies.pi0.configuration_pi0 import DEFAULT_IMAGE_SIZE, PI0Config
-from lerobot.policies.pretrained import PreTrainedPolicy, T
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
 from lerobot.utils.constants import (
     ACTION,
     OBS_LANGUAGE_ATTENTION_MASK,
@@ -60,6 +57,10 @@ from lerobot.utils.constants import (
     OPENPI_ATTENTION_MASK_VALUE,
 )
 
+from ..pretrained import PreTrainedPolicy, T
+from ..rtc.modeling_rtc import RTCProcessor
+from .configuration_pi0 import DEFAULT_IMAGE_SIZE, PI0Config
+
 
 class ActionSelectKwargs(TypedDict, total=False):
     inference_delay: int | None

src/lerobot/policies/pi0/processor_pi0.py

@@ -18,8 +18,7 @@ from typing import Any
 
 import torch
 
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     AddBatchDimensionProcessorStep,
@@ -34,10 +33,13 @@ from lerobot.processor import (
     RenameObservationsProcessorStep,
     TokenizerProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
+from .configuration_pi0 import PI0Config
+
 
 @ProcessorStepRegistry.register(name="pi0_new_line_processor")
 class Pi0NewLineProcessor(ComplementaryDataProcessorStep):

src/lerobot/policies/pi05/configuration_pi05.py

@@ -16,13 +16,12 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
-from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
+from ..rtc.configuration_rtc import RTCConfig
+
 DEFAULT_IMAGE_SIZE = 224
 
 

src/lerobot/policies/pi05/modeling_pi05.py

@@ -33,7 +33,7 @@ if TYPE_CHECKING or _transformers_available:
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
 
-    from lerobot.policies.pi_gemma import (
+    from ..pi_gemma import (
         PaliGemmaForConditionalGenerationWithPiGemma,
         PiGemmaForCausalLM,
         _gated_residual,
@@ -46,10 +46,7 @@ else:
     _gated_residual = None
     layernorm_forward = None
     PaliGemmaForConditionalGenerationWithPiGemma = None
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import PreTrainedConfig
-from lerobot.policies.pi05.configuration_pi05 import DEFAULT_IMAGE_SIZE, PI05Config
-from lerobot.policies.pretrained import PreTrainedPolicy, T
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
 from lerobot.utils.constants import (
     ACTION,
     OBS_LANGUAGE_ATTENTION_MASK,
@@ -57,6 +54,10 @@ from lerobot.utils.constants import (
     OPENPI_ATTENTION_MASK_VALUE,
 )
 
+from ..pretrained import PreTrainedPolicy, T
+from ..rtc.modeling_rtc import RTCProcessor
+from .configuration_pi05 import DEFAULT_IMAGE_SIZE, PI05Config
+
 
 class ActionSelectKwargs(TypedDict, total=False):
     inference_delay: int | None

src/lerobot/policies/pi05/processor_pi05.py

@@ -21,8 +21,7 @@ from typing import Any
 import numpy as np
 import torch
 
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.policies.pi05.configuration_pi05 import PI05Config
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     AddBatchDimensionProcessorStep,
@@ -36,8 +35,9 @@ from lerobot.processor import (
     RenameObservationsProcessorStep,
     TokenizerProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     OBS_STATE,
@@ -45,6 +45,8 @@ from lerobot.utils.constants import (
     POLICY_PREPROCESSOR_DEFAULT_NAME,
 )
 
+from .configuration_pi05 import PI05Config
+
 
 @ProcessorStepRegistry.register(name="pi05_prepare_state_tokenizer_processor_step")
 @dataclass

src/lerobot/policies/pi0_fast/configuration_pi0_fast.py

@@ -16,13 +16,12 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
-from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
+from ..rtc.configuration_rtc import RTCConfig
+
 DEFAULT_IMAGE_SIZE = 224
 
 

src/lerobot/policies/pi0_fast/modeling_pi0_fast.py

@@ -38,7 +38,7 @@ if TYPE_CHECKING or _transformers_available:
     from transformers import AutoTokenizer
     from transformers.models.auto import CONFIG_MAPPING
 
-    from lerobot.policies.pi_gemma import (
+    from ..pi_gemma import (
         PaliGemmaForConditionalGenerationWithPiGemma,
         PiGemmaModel,
     )
@@ -48,10 +48,7 @@ else:
     PiGemmaModel = None
     PaliGemmaForConditionalGenerationWithPiGemma = None
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import PreTrainedConfig
-from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
-from lerobot.policies.pretrained import PreTrainedPolicy, T
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
 from lerobot.utils.constants import (
     ACTION,
     ACTION_TOKEN_MASK,
@@ -61,6 +58,10 @@ from lerobot.utils.constants import (
     OPENPI_ATTENTION_MASK_VALUE,
 )
 
+from ..pretrained import PreTrainedPolicy, T
+from ..rtc.modeling_rtc import RTCProcessor
+from .configuration_pi0_fast import PI0FastConfig
+
 
 class ActionSelectKwargs(TypedDict, total=False):
     temperature: float | None

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -21,8 +21,7 @@ from typing import Any
 import numpy as np
 import torch
 
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     ActionTokenizerProcessorStep,
@@ -37,8 +36,9 @@ from lerobot.processor import (
     RenameObservationsProcessorStep,
     TokenizerProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     OBS_STATE,
@@ -46,6 +46,8 @@ from lerobot.utils.constants import (
     POLICY_PREPROCESSOR_DEFAULT_NAME,
 )
 
+from .configuration_pi0_fast import PI0FastConfig
+
 
 @ProcessorStepRegistry.register(name="pi0_fast_prepare_state_tokenizer_processor_step")
 @dataclass

src/lerobot/policies/pretrained.py

@@ -29,11 +29,12 @@ from huggingface_hub.errors import HfHubHTTPError
 from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
 from torch import Tensor, nn
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
-from lerobot.policies.utils import log_model_loading_keys
 from lerobot.utils.hub import HubMixin
 
+from .utils import log_model_loading_keys
+
 T = TypeVar("T", bound="PreTrainedPolicy")
 
 

src/lerobot/policies/rtc/__init__.py

@@ -14,11 +14,11 @@
 
 """Real-Time Chunking (RTC) utilities for action-chunking policies."""
 
-from lerobot.policies.rtc.action_interpolator import ActionInterpolator
+from .action_interpolator import ActionInterpolator
-from lerobot.policies.rtc.action_queue import ActionQueue
+from .action_queue import ActionQueue
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
+from .configuration_rtc import RTCConfig
-from lerobot.policies.rtc.latency_tracker import LatencyTracker
+from .latency_tracker import LatencyTracker
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
+from .modeling_rtc import RTCProcessor
 
 __all__ = [
     "ActionInterpolator",

src/lerobot/policies/rtc/action_queue.py

@@ -27,7 +27,7 @@ from threading import Lock
 import torch
 from torch import Tensor
 
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
+from .configuration_rtc import RTCConfig
 
 logger = logging.getLogger(__name__)
 

src/lerobot/policies/rtc/configuration_rtc.py

@@ -23,7 +23,7 @@ Based on:
 
 from dataclasses import dataclass
 
-from lerobot.configs.types import RTCAttentionSchedule
+from lerobot.configs import RTCAttentionSchedule
 
 
 @dataclass

src/lerobot/policies/rtc/modeling_rtc.py

@@ -27,9 +27,10 @@ import math
 import torch
 from torch import Tensor
 
-from lerobot.configs.types import RTCAttentionSchedule
+from lerobot.configs import RTCAttentionSchedule
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
+
-from lerobot.policies.rtc.debug_tracker import Tracker
+from .configuration_rtc import RTCConfig
+from .debug_tracker import Tracker
 
 logger = logging.getLogger(__name__)
 

src/lerobot/policies/sac/configuration_sac.py

@@ -17,9 +17,8 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import NormalizationMode, PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
+from lerobot.optim import MultiAdamConfig
-from lerobot.optim.optimizers import MultiAdamConfig
 from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
 
 

src/lerobot/policies/sac/modeling_sac.py

@@ -28,11 +28,12 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 from torch.distributions import MultivariateNormal, TanhTransform, Transform, TransformedDistribution
 
-from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.sac.configuration_sac import SACConfig, is_image_feature
-from lerobot.policies.utils import get_device_from_parameters
 from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STATE
 
+from ..pretrained import PreTrainedPolicy
+from ..utils import get_device_from_parameters
+from .configuration_sac import SACConfig, is_image_feature
+
 DISCRETE_DIMENSION_INDEX = -1  # Gripper is always the last dimension
 
 

src/lerobot/policies/sac/processor_sac.py

@@ -19,7 +19,6 @@ from typing import Any
 
 import torch
 
-from lerobot.policies.sac.configuration_sac import SACConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -28,10 +27,13 @@ from lerobot.processor import (
     PolicyProcessorPipeline,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
+from .configuration_sac import SACConfig
+
 
 def make_sac_pre_post_processors(
     config: SACConfig,

src/lerobot/policies/sac/reward_model/configuration_classifier.py

@@ -15,10 +15,8 @@
 # limitations under the License.
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import NormalizationMode, PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
+from lerobot.optim import AdamWConfig, LRSchedulerConfig, OptimizerConfig
-from lerobot.optim.optimizers import AdamWConfig, OptimizerConfig
-from lerobot.optim.schedulers import LRSchedulerConfig
 from lerobot.utils.constants import OBS_IMAGE
 
 

src/lerobot/policies/sac/reward_model/modeling_classifier.py

@@ -19,10 +19,11 @@ import logging
 import torch
 from torch import Tensor, nn
 
-from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 from lerobot.utils.constants import OBS_IMAGE, REWARD
 
+from ...pretrained import PreTrainedPolicy
+from .configuration_classifier import RewardClassifierConfig
+
 
 class ClassifierOutput:
     """Wrapper for classifier outputs with additional metadata."""

src/lerobot/policies/sac/reward_model/processor_classifier.py

@@ -18,15 +18,17 @@ from typing import Any
 
 import torch
 
-from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 from lerobot.processor import (
     DeviceProcessorStep,
     IdentityProcessorStep,
     NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+
+from .configuration_classifier import RewardClassifierConfig
 
 
 def make_classifier_processor(

src/lerobot/policies/sarm/__init__.py

@@ -0,0 +1,3 @@
+from .configuration_sarm import SARMConfig
+
+__all__ = ["SARMConfig"]

src/lerobot/policies/sarm/compute_rabc_weights.py

@@ -57,10 +57,11 @@ import pyarrow.parquet as pq
 import torch
 from tqdm import tqdm
 
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets import LeRobotDataset
-from lerobot.policies.sarm.modeling_sarm import SARMRewardModel
+
-from lerobot.policies.sarm.processor_sarm import make_sarm_pre_post_processors
+from .modeling_sarm import SARMRewardModel
-from lerobot.policies.sarm.sarm_utils import normalize_stage_tau
+from .processor_sarm import make_sarm_pre_post_processors
+from .sarm_utils import normalize_stage_tau
 
 
 def get_reward_model_path_from_parquet(parquet_path: Path) -> str | None:

src/lerobot/policies/sarm/configuration_sarm.py

@@ -22,10 +22,8 @@ Paper: https://arxiv.org/abs/2509.25358
 
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
-from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
 from lerobot.utils.constants import OBS_IMAGES, OBS_STATE
 
 

src/lerobot/policies/sarm/modeling_sarm.py

@@ -34,13 +34,14 @@ import torch.nn as nn
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 
-from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.utils.constants import OBS_STR
-from lerobot.policies.sarm.configuration_sarm import SARMConfig
+
-from lerobot.policies.sarm.sarm_utils import (
+from ..pretrained import PreTrainedPolicy
+from .configuration_sarm import SARMConfig
+from .sarm_utils import (
     normalize_stage_tau,
     pad_state_to_max_dim,
 )
-from lerobot.utils.constants import OBS_STR
 
 
 class StageTransformer(nn.Module):

src/lerobot/policies/sarm/processor_sarm.py

@@ -24,33 +24,35 @@ import pandas as pd
 import torch
 from faker import Faker
 from PIL import Image
+
+from lerobot.utils.import_utils import require_package
+
+require_package("transformers", extra="sarm")
+
 from transformers import CLIPModel, CLIPProcessor
 
-from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.policies.sarm.configuration_sarm import SARMConfig
+from lerobot.processor import (
-from lerobot.policies.sarm.sarm_utils import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    RenameObservationsProcessorStep,
+    from_tensor_to_numpy,
+    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
+from .sarm_utils import (
     apply_rewind_augmentation,
     compute_absolute_indices,
     find_stage_and_tau,
     pad_state_to_max_dim,
 )
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    ProcessorStep,
-    RenameObservationsProcessorStep,
-)
-from lerobot.processor.converters import (
-    from_tensor_to_numpy,
-    policy_action_to_transition,
-    transition_to_policy_action,
-)
-from lerobot.processor.pipeline import PipelineFeatureType
-from lerobot.types import EnvTransition, TransitionKey
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 class SARMEncodingProcessorStep(ProcessorStep):

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -14,15 +14,12 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
-from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import (
-    CosineDecayWithWarmupSchedulerConfig,
-)
-from lerobot.policies.rtc.configuration_rtc import RTCConfig
 from lerobot.utils.constants import OBS_IMAGES
 
+from ..rtc.configuration_rtc import RTCConfig
+
 
 @PreTrainedConfig.register_subclass("smolvla")
 @dataclass

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -60,16 +60,17 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 
-from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
-from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
-from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
-from lerobot.policies.utils import (
-    populate_queues,
-)
 from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.utils.device_utils import get_safe_dtype
 
+from ..pretrained import PreTrainedPolicy
+from ..rtc.modeling_rtc import RTCProcessor
+from ..utils import (
+    populate_queues,
+)
+from .configuration_smolvla import SmolVLAConfig
+from .smolvlm_with_expert import SmolVLMWithExpertModel
+
 
 class ActionSelectKwargs(TypedDict, total=False):
     inference_delay: int | None

src/lerobot/policies/smolvla/processor_smolvla.py

@@ -18,8 +18,7 @@ from typing import Any
 
 import torch
 
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     ComplementaryDataProcessorStep,
@@ -31,10 +30,13 @@ from lerobot.processor import (
     RenameObservationsProcessorStep,
     TokenizerProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
+from .configuration_smolvla import SmolVLAConfig
+
 
 def make_smolvla_pre_post_processors(
     config: SmolVLAConfig,

src/lerobot/policies/smolvla/smolvlm_with_expert.py

@@ -16,6 +16,11 @@ import copy
 
 import torch
 from torch import nn
+
+from lerobot.utils.import_utils import require_package
+
+require_package("transformers", extra="smolvla")
+
 from transformers import (
     AutoConfig,
     AutoModel,

src/lerobot/policies/tdmpc/configuration_tdmpc.py

@@ -16,9 +16,8 @@
 # limitations under the License.
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import NormalizationMode, PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
+from lerobot.optim import AdamConfig
-from lerobot.optim.optimizers import AdamConfig
 
 
 @PreTrainedConfig.register_subclass("tdmpc")

src/lerobot/policies/tdmpc/modeling_tdmpc.py

@@ -35,11 +35,12 @@ import torch.nn as nn
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 
-from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
 from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_PREFIX, OBS_STATE, OBS_STR, REWARD
 
+from ..pretrained import PreTrainedPolicy
+from ..utils import get_device_from_parameters, get_output_shape, populate_queues
+from .configuration_tdmpc import TDMPCConfig
+
 
 class TDMPCPolicy(PreTrainedPolicy):
     """Implementation of TD-MPC learning + inference.

src/lerobot/policies/tdmpc/processor_tdmpc.py

@@ -18,7 +18,6 @@ from typing import Any
 
 import torch
 
-from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -27,10 +26,13 @@ from lerobot.processor import (
     PolicyProcessorPipeline,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
+from .configuration_tdmpc import TDMPCConfig
+
 
 def make_tdmpc_pre_post_processors(
     config: TDMPCConfig,

src/lerobot/policies/utils.py

@@ -21,10 +21,10 @@ import numpy as np
 import torch
 from torch import nn
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig
-from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.types import PolicyAction, RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
+from lerobot.utils.feature_utils import build_dataset_frame
 
 
 def populate_queues(
@@ -162,8 +162,6 @@ def build_inference_frame(
     Returns:
         A dictionary of preprocessed tensors ready for model inference.
     """
-    from lerobot.utils.feature_utils import build_dataset_frame
-
     # Extracts the correct keys from the incoming raw observation
     observation = build_dataset_frame(ds_features, observation, prefix=OBS_STR)
 

src/lerobot/policies/vqbet/configuration_vqbet.py

@@ -18,10 +18,8 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import NormalizationMode, PreTrainedConfig
-from lerobot.configs.types import NormalizationMode
+from lerobot.optim import AdamConfig, VQBeTSchedulerConfig
-from lerobot.optim.optimizers import AdamConfig
-from lerobot.optim.schedulers import VQBeTSchedulerConfig
 
 
 @PreTrainedConfig.register_subclass("vqbet")

src/lerobot/policies/vqbet/modeling_vqbet.py

@@ -27,12 +27,13 @@ import torch.nn.functional as F  # noqa: N812
 import torchvision
 from torch import Tensor, nn
 
-from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
-from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
-from lerobot.policies.vqbet.vqbet_utils import GPT, ResidualVQ
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
 
+from ..pretrained import PreTrainedPolicy
+from ..utils import get_device_from_parameters, get_output_shape, populate_queues
+from .configuration_vqbet import VQBeTConfig
+from .vqbet_utils import GPT, ResidualVQ
+
 # ruff: noqa: N806
 
 

src/lerobot/policies/vqbet/processor_vqbet.py

@@ -19,7 +19,6 @@ from typing import Any
 
 import torch
 
-from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -28,10 +27,13 @@ from lerobot.processor import (
     PolicyProcessorPipeline,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
+from .configuration_vqbet import VQBeTConfig
+
 
 def make_vqbet_pre_post_processors(
     config: VQBeTConfig,

src/lerobot/policies/vqbet/vqbet_utils.py

@@ -30,7 +30,7 @@ from torch import einsum, nn
 from torch.cuda.amp import autocast
 from torch.optim import Optimizer
 
-from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
+from .configuration_vqbet import VQBeTConfig
 
 # ruff: noqa: N806
 

src/lerobot/policies/wall_x/__init__.py

@@ -16,4 +16,4 @@
 
 from .configuration_wall_x import WallXConfig
 
-__all__ = ["WallXConfig", "WallXPolicy", "make_wall_x_pre_post_processors"]
+__all__ = ["WallXConfig"]

src/lerobot/policies/wall_x/configuration_wall_x.py

@@ -14,10 +14,8 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
-from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
 from lerobot.utils.constants import ACTION, OBS_STATE
 
 

src/lerobot/policies/wall_x/modeling_wall_x.py

@@ -43,6 +43,14 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+
+from lerobot.utils.import_utils import require_package
+
+require_package("transformers", extra="wallx")
+require_package("peft", extra="wallx")
+require_package("torchdiffeq", extra="wallx")
+require_package("qwen-vl-utils", extra="wallx", import_name="qwen_vl_utils")
+
 from peft import LoraConfig, get_peft_model
 from PIL import Image
 from qwen_vl_utils.vision_process import smart_resize
@@ -59,10 +67,12 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
 )
 from transformers.utils import is_torchdynamo_compiling, logging
 
-from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.utils.constants import ACTION, OBS_STATE
-from lerobot.policies.utils import populate_queues
+
-from lerobot.policies.wall_x.configuration_wall_x import WallXConfig
+from ..pretrained import PreTrainedPolicy
-from lerobot.policies.wall_x.constant import (
+from ..utils import populate_queues
+from .configuration_wall_x import WallXConfig
+from .constant import (
     GENERATE_SUBTASK_RATIO,
     IMAGE_FACTOR,
     MAX_PIXELS,
@@ -72,19 +82,18 @@ from lerobot.policies.wall_x.constant import (
     RESOLUTION,
     TOKENIZER_MAX_LENGTH,
 )
-from lerobot.policies.wall_x.qwen_model.configuration_qwen2_5_vl import Qwen2_5_VLConfig
+from .qwen_model.configuration_qwen2_5_vl import Qwen2_5_VLConfig
-from lerobot.policies.wall_x.qwen_model.qwen2_5_vl_moe import (
+from .qwen_model.qwen2_5_vl_moe import (
     Qwen2_5_VisionTransformerPretrainedModel,
     Qwen2_5_VLACausalLMOutputWithPast,
     Qwen2_5_VLMoEModel,
 )
-from lerobot.policies.wall_x.utils import (
+from .utils import (
     get_wallx_normal_text,
     preprocesser_call,
     process_grounding_points,
     replace_action_token,
 )
-from lerobot.utils.constants import ACTION, OBS_STATE
 
 logger = logging.get_logger(__name__)
 

src/lerobot/policies/wall_x/processor_wall_x.py

@@ -18,8 +18,7 @@ from typing import Any
 
 import torch
 
-from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.policies.wall_x.configuration_wall_x import WallXConfig
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     ComplementaryDataProcessorStep,
@@ -30,10 +29,13 @@ from lerobot.processor import (
     ProcessorStepRegistry,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
+from .configuration_wall_x import WallXConfig
+
 
 def make_wall_x_pre_post_processors(
     config: WallXConfig,

src/lerobot/policies/wall_x/utils.py

@@ -30,10 +30,11 @@ from typing import Any
 import torch
 from transformers import BatchFeature
 
-from lerobot.policies.wall_x.constant import (
+from lerobot.utils.constants import OBS_IMAGES
+
+from .constant import (
     CAMERA_NAME_MAPPING,
 )
-from lerobot.utils.constants import OBS_IMAGES
 
 
 @dataclass