From d626964119177f400d7d427d0b1fda4f2df38e68 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Sat, 11 Apr 2026 15:03:24 +0200 Subject: [PATCH] big imports refactor --- pyproject.toml | 14 ++- src/lerobot/async_inference/__init__.py | 6 + src/lerobot/async_inference/helpers.py | 2 +- src/lerobot/async_inference/policy_server.py | 2 +- src/lerobot/async_inference/robot_client.py | 4 +- src/lerobot/cameras/__init__.py | 6 + src/lerobot/cameras/zmq/image_server.py | 4 +- src/lerobot/configs/__init__.py | 47 ++++++++ src/lerobot/configs/eval.py | 5 +- src/lerobot/configs/policies.py | 6 +- src/lerobot/configs/train.py | 8 +- .../sarm_annotations/subtask_annotation.py | 2 +- src/lerobot/datasets/__init__.py | 44 +++++-- src/lerobot/datasets/aggregate.py | 12 +- src/lerobot/datasets/compute_stats.py | 6 +- src/lerobot/datasets/dataset_metadata.py | 15 +-- src/lerobot/datasets/dataset_reader.py | 8 +- src/lerobot/datasets/dataset_tools.py | 26 ++-- src/lerobot/datasets/dataset_writer.py | 14 +-- src/lerobot/datasets/factory.py | 11 +- src/lerobot/datasets/feature_utils.py | 7 +- src/lerobot/datasets/io_utils.py | 7 +- src/lerobot/datasets/lerobot_dataset.py | 13 +- src/lerobot/datasets/multi_dataset.py | 9 +- src/lerobot/datasets/pipeline_features.py | 2 +- src/lerobot/datasets/streaming_dataset.py | 13 +- src/lerobot/datasets/video_utils.py | 2 +- src/lerobot/envs/__init__.py | 19 ++- src/lerobot/envs/configs.py | 15 +-- src/lerobot/envs/factory.py | 4 +- src/lerobot/envs/libero.py | 3 +- src/lerobot/envs/metaworld.py | 3 +- src/lerobot/envs/utils.py | 5 +- src/lerobot/model/__init__.py | 5 + src/lerobot/motors/__init__.py | 2 + src/lerobot/motors/dynamixel/__init__.py | 4 + src/lerobot/motors/feetech/__init__.py | 4 + src/lerobot/optim/__init__.py | 43 ++++++- src/lerobot/optim/factory.py | 2 +- src/lerobot/optim/schedulers.py | 3 + src/lerobot/policies/__init__.py | 34 +++++- src/lerobot/policies/act/configuration_act.py | 5 +- src/lerobot/policies/act/modeling_act.py | 5 +- src/lerobot/policies/act/processor_act.py | 6 +- .../diffusion/configuration_diffusion.py | 6 +- .../policies/diffusion/modeling_diffusion.py | 11 +- .../policies/diffusion/processor_diffusion.py | 6 +- src/lerobot/policies/factory.py | 113 +++++++++--------- src/lerobot/policies/groot/__init__.py | 4 + .../groot/action_head/cross_attention_dit.py | 2 +- .../action_head/flow_matching_action_head.py | 3 +- .../policies/groot/configuration_groot.py | 6 +- src/lerobot/policies/groot/groot_n1.py | 7 +- src/lerobot/policies/groot/modeling_groot.py | 9 +- src/lerobot/policies/groot/processor_groot.py | 7 +- .../configuration_multi_task_dit.py | 6 +- .../multi_task_dit/modeling_multi_task_dit.py | 10 +- .../processor_multi_task_dit.py | 6 +- src/lerobot/policies/pi0/configuration_pi0.py | 9 +- src/lerobot/policies/pi0/modeling_pi0.py | 11 +- src/lerobot/policies/pi0/processor_pi0.py | 8 +- .../policies/pi05/configuration_pi05.py | 9 +- src/lerobot/policies/pi05/modeling_pi05.py | 11 +- src/lerobot/policies/pi05/processor_pi05.py | 8 +- .../pi0_fast/configuration_pi0_fast.py | 9 +- .../policies/pi0_fast/modeling_pi0_fast.py | 11 +- .../policies/pi0_fast/processor_pi0_fast.py | 8 +- src/lerobot/policies/pretrained.py | 5 +- src/lerobot/policies/rtc/__init__.py | 10 +- src/lerobot/policies/rtc/action_queue.py | 2 +- src/lerobot/policies/rtc/configuration_rtc.py | 2 +- src/lerobot/policies/rtc/modeling_rtc.py | 7 +- src/lerobot/policies/sac/configuration_sac.py | 5 +- src/lerobot/policies/sac/modeling_sac.py | 7 +- src/lerobot/policies/sac/processor_sac.py | 6 +- .../reward_model/configuration_classifier.py | 6 +- .../sac/reward_model/modeling_classifier.py | 5 +- .../sac/reward_model/processor_classifier.py | 6 +- src/lerobot/policies/sarm/__init__.py | 3 + .../policies/sarm/compute_rabc_weights.py | 9 +- .../policies/sarm/configuration_sarm.py | 6 +- src/lerobot/policies/sarm/modeling_sarm.py | 9 +- src/lerobot/policies/sarm/processor_sarm.py | 42 +++---- .../policies/smolvla/configuration_smolvla.py | 11 +- .../policies/smolvla/modeling_smolvla.py | 15 +-- .../policies/smolvla/processor_smolvla.py | 8 +- .../policies/smolvla/smolvlm_with_expert.py | 5 + .../policies/tdmpc/configuration_tdmpc.py | 5 +- src/lerobot/policies/tdmpc/modeling_tdmpc.py | 7 +- src/lerobot/policies/tdmpc/processor_tdmpc.py | 6 +- src/lerobot/policies/utils.py | 6 +- .../policies/vqbet/configuration_vqbet.py | 6 +- src/lerobot/policies/vqbet/modeling_vqbet.py | 9 +- src/lerobot/policies/vqbet/processor_vqbet.py | 6 +- src/lerobot/policies/vqbet/vqbet_utils.py | 2 +- src/lerobot/policies/wall_x/__init__.py | 2 +- .../policies/wall_x/configuration_wall_x.py | 6 +- .../policies/wall_x/modeling_wall_x.py | 25 ++-- .../policies/wall_x/processor_wall_x.py | 8 +- src/lerobot/policies/wall_x/utils.py | 5 +- src/lerobot/policies/xvla/__init__.py | 2 +- .../policies/xvla/configuration_xvla.py | 6 +- src/lerobot/policies/xvla/modeling_xvla.py | 10 +- src/lerobot/policies/xvla/processor_xvla.py | 10 +- src/lerobot/processor/__init__.py | 14 ++- src/lerobot/processor/batch_processor.py | 2 +- .../processor/delta_action_processor.py | 2 +- src/lerobot/processor/device_processor.py | 2 +- src/lerobot/processor/env_processor.py | 2 +- src/lerobot/processor/gym_action_processor.py | 6 +- src/lerobot/processor/hil_processor.py | 2 +- .../processor/migrate_policy_normalization.py | 4 +- src/lerobot/processor/normalize_processor.py | 4 +- .../processor/observation_processor.py | 2 +- src/lerobot/processor/pipeline.py | 5 +- src/lerobot/processor/policy_robot_bridge.py | 2 +- .../processor/relative_action_processor.py | 2 +- src/lerobot/processor/rename_processor.py | 2 +- src/lerobot/processor/tokenizer_processor.py | 2 +- src/lerobot/rl/__init__.py | 6 + src/lerobot/rl/actor.py | 6 +- src/lerobot/rl/buffer.py | 2 +- src/lerobot/rl/crop_dataset_roi.py | 2 +- src/lerobot/rl/eval_policy.py | 4 +- src/lerobot/rl/gym_manipulator.py | 8 +- .../rl/joint_observations_processor.py | 4 +- src/lerobot/rl/learner.py | 10 +- src/lerobot/rl/learner_service.py | 3 +- src/lerobot/robots/__init__.py | 2 + .../bi_openarm_follower.py | 2 +- .../config_bi_openarm_follower.py | 2 +- .../robots/bi_so_follower/bi_so_follower.py | 2 +- .../bi_so_follower/config_bi_so_follower.py | 3 +- src/lerobot/robots/hope_jr/hope_jr_arm.py | 2 +- src/lerobot/robots/hope_jr/hope_jr_hand.py | 2 +- .../robots/koch_follower/koch_follower.py | 2 +- src/lerobot/robots/lekiwi/__init__.py | 4 + src/lerobot/robots/lekiwi/config_lekiwi.py | 4 +- src/lerobot/robots/lekiwi/lekiwi.py | 2 +- .../robots/omx_follower/omx_follower.py | 2 +- .../openarm_follower/openarm_follower.py | 2 +- .../robots/reachy2/configuration_reachy2.py | 3 +- src/lerobot/robots/reachy2/robot_reachy2.py | 2 +- .../so_follower/robot_kinematic_processor.py | 4 +- src/lerobot/robots/so_follower/so_follower.py | 2 +- .../robots/unitree_g1/gr00t_locomotion.py | 2 +- .../robots/unitree_g1/holosoma_locomotion.py | 2 +- src/lerobot/robots/unitree_g1/unitree_g1.py | 20 ++-- .../robots/unitree_g1/unitree_sdk2_socket.py | 2 +- .../scripts/augment_dataset_quantile_stats.py | 11 +- .../scripts/convert_dataset_v21_to_v30.py | 4 +- src/lerobot/scripts/lerobot_calibrate.py | 4 +- src/lerobot/scripts/lerobot_edit_dataset.py | 4 +- src/lerobot/scripts/lerobot_eval.py | 8 +- src/lerobot/scripts/lerobot_find_cameras.py | 6 +- .../scripts/lerobot_find_joint_limits.py | 2 +- .../scripts/lerobot_imgtransform_viz.py | 2 +- src/lerobot/scripts/lerobot_record.py | 28 ++--- src/lerobot/scripts/lerobot_teleoperate.py | 6 +- src/lerobot/scripts/lerobot_train.py | 15 +-- .../scripts/lerobot_train_tokenizer.py | 3 +- src/lerobot/teleoperators/__init__.py | 2 + .../bi_openarm_leader/bi_openarm_leader.py | 3 +- .../config_bi_openarm_leader.py | 3 +- .../bi_so_leader/bi_so_leader.py | 3 +- .../bi_so_leader/config_bi_so_leader.py | 3 +- .../homunculus/homunculus_glove.py | 2 +- .../teleoperators/phone/phone_processor.py | 5 +- .../teleoperators/phone/teleop_phone.py | 5 +- src/lerobot/transforms/__init__.py | 2 +- src/lerobot/transport/__init__.py | 2 + src/lerobot/utils/__init__.py | 66 ++++++++++ src/lerobot/utils/control_utils.py | 9 +- src/lerobot/utils/decorators.py | 2 +- src/lerobot/utils/feature_utils.py | 5 +- src/lerobot/utils/io_utils.py | 2 +- src/lerobot/utils/logging_utils.py | 2 +- src/lerobot/utils/random_utils.py | 4 +- src/lerobot/utils/train_utils.py | 17 ++- src/lerobot/utils/transition.py | 2 +- src/lerobot/utils/visualization_utils.py | 2 +- tests/fixtures/dataset_factories.py | 4 +- uv.lock | 19 ++- 183 files changed, 892 insertions(+), 558 deletions(-) create mode 100644 src/lerobot/async_inference/__init__.py create mode 100644 src/lerobot/configs/__init__.py create mode 100644 src/lerobot/model/__init__.py create mode 100644 src/lerobot/policies/sarm/__init__.py create mode 100644 src/lerobot/rl/__init__.py create mode 100644 src/lerobot/transport/__init__.py create mode 100644 src/lerobot/utils/__init__.py diff --git a/pyproject.toml b/pyproject.toml index 3d036afa8..213ba05ce 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/lerobot/async_inference/__init__.py b/src/lerobot/async_inference/__init__.py new file mode 100644 index 000000000..0da31b6a8 --- /dev/null +++ b/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") diff --git a/src/lerobot/async_inference/helpers.py b/src/lerobot/async_inference/helpers.py index e0e96c95f..4931c68c5 100644 --- a/src/lerobot/async_inference/helpers.py +++ b/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 diff --git a/src/lerobot/async_inference/policy_server.py b/src/lerobot/async_inference/policy_server.py index 3f63929df..787d39abf 100644 --- a/src/lerobot/async_inference/policy_server.py +++ b/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 diff --git a/src/lerobot/async_inference/robot_client.py b/src/lerobot/async_inference/robot_client.py index 0ee70a0e6..a250a08fb 100644 --- a/src/lerobot/async_inference/robot_client.py +++ b/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.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 +from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, diff --git a/src/lerobot/cameras/__init__.py b/src/lerobot/cameras/__init__.py index cbf1f11bf..3598d58aa 100644 --- a/src/lerobot/cameras/__init__.py +++ b/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"] diff --git a/src/lerobot/cameras/zmq/image_server.py b/src/lerobot/cameras/zmq/image_server.py index 8222b9fee..b8b6f8e74 100644 --- a/src/lerobot/cameras/zmq/image_server.py +++ b/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 lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig +from ..configs import ColorMode +from ..opencv import OpenCVCamera, OpenCVCameraConfig logger = logging.getLogger(__name__) diff --git a/src/lerobot/configs/__init__.py b/src/lerobot/configs/__init__.py new file mode 100644 index 000000000..3ddaec1af --- /dev/null +++ b/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", +] diff --git a/src/lerobot/configs/eval.py b/src/lerobot/configs/eval.py index da8bee6b2..d1cebd27f 100644 --- a/src/lerobot/configs/eval.py +++ b/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__) diff --git a/src/lerobot/configs/policies.py b/src/lerobot/configs/policies.py index ce567b8f5..91701af6d 100644 --- a/src/lerobot/configs/policies.py +++ b/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.optimizers import OptimizerConfig -from lerobot.optim.schedulers import LRSchedulerConfig +from lerobot.optim import LRSchedulerConfig, OptimizerConfig 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__) diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index 8b8aedb26..d754a0847 100644 --- a/src/lerobot/configs/train.py +++ b/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.configs.policies import PreTrainedConfig -from lerobot.optim import OptimizerConfig -from lerobot.optim.schedulers import LRSchedulerConfig +from lerobot.optim import LRSchedulerConfig, OptimizerConfig from lerobot.utils.hub import HubMixin +from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig +from .policies import PreTrainedConfig + TRAIN_CONFIG_NAME = "train_config.json" diff --git a/src/lerobot/data_processing/sarm_annotations/subtask_annotation.py b/src/lerobot/data_processing/sarm_annotations/subtask_annotation.py index 6865a6d1e..169485513 100644 --- a/src/lerobot/data_processing/sarm_annotations/subtask_annotation.py +++ b/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 diff --git a/src/lerobot/datasets/__init__.py b/src/lerobot/datasets/__init__.py index ec7b9f45f..c24ed4d9a 100644 --- a/src/lerobot/datasets/__init__.py +++ b/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 lerobot.datasets.factory import make_dataset -from lerobot.datasets.image_writer import safe_stop_image_writer -from lerobot.datasets.io_utils import load_episodes, write_stats -from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.datasets.multi_dataset import MultiLeRobotDataset -from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features -from lerobot.datasets.sampler import EpisodeAwareSampler -from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset -from lerobot.datasets.video_utils import VideoEncodingManager +from .compute_stats import aggregate_stats, get_feature_stats +from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata +from .dataset_tools import ( + convert_image_to_video_dataset, + delete_episodes, + merge_datasets, + modify_tasks, + recompute_stats, + remove_feature, + 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", ] diff --git a/src/lerobot/datasets/aggregate.py b/src/lerobot/datasets/aggregate.py index 66f055f04..0da1da964 100644 --- a/src/lerobot/datasets/aggregate.py +++ b/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 lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata -from lerobot.datasets.feature_utils import get_hf_features_from_features -from lerobot.datasets.io_utils import ( +from .compute_stats import aggregate_stats +from .dataset_metadata import LeRobotDatasetMetadata +from .feature_utils import get_hf_features_from_features +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]): diff --git a/src/lerobot/datasets/compute_stats.py b/src/lerobot/datasets/compute_stats.py index 03eefe40e..f489c84a7 100644 --- a/src/lerobot/datasets/compute_stats.py +++ b/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 = [] diff --git a/src/lerobot/datasets/dataset_metadata.py b/src/lerobot/datasets/dataset_metadata.py index cf150c6dd..4b2d42838 100644 --- a/src/lerobot/datasets/dataset_metadata.py +++ b/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.datasets.feature_utils import create_empty_dataset_info -from lerobot.datasets.io_utils import ( +from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE +from lerobot.utils.feature_utils import _validate_feature_names + +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 lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE -from lerobot.utils.feature_utils import _validate_feature_names +from .video_utils import get_video_info CODEBASE_VERSION = "v3.0" diff --git a/src/lerobot/datasets/dataset_reader.py b/src/lerobot/datasets/dataset_reader.py index 3720a5084..fc7ce36ce 100644 --- a/src/lerobot/datasets/dataset_reader.py +++ b/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 lerobot.datasets.feature_utils import ( +from .dataset_metadata import LeRobotDatasetMetadata +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: diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index fa91ef412..cbf4e5c49 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/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.datasets.compute_stats import ( +from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE +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 lerobot.datasets.io_utils import ( +from .dataset_metadata import LeRobotDatasetMetadata +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.datasets.utils import ( +from .lerobot_dataset import LeRobotDataset +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 lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE +from .video_utils import encode_video_frames, get_video_info 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 lerobot.datasets.io_utils import embed_images + from .feature_utils import get_hf_features_from_features + 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")) diff --git a/src/lerobot/datasets/dataset_writer.py b/src/lerobot/datasets/dataset_writer.py index 787ecd337..60ec9e348 100644 --- a/src/lerobot/datasets/dataset_writer.py +++ b/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 lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata -from lerobot.datasets.feature_utils import ( +from .compute_stats import compute_episode_stats +from .dataset_metadata import LeRobotDatasetMetadata +from .feature_utils import ( get_hf_features_from_features, validate_episode_buffer, validate_frame, ) -from lerobot.datasets.image_writer import AsyncImageWriter, write_image -from lerobot.datasets.io_utils import ( +from .image_writer import AsyncImageWriter, write_image +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, diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py index 1baac5e1e..040cba5cb 100644 --- a/src/lerobot/datasets/factory.py +++ b/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 diff --git a/src/lerobot/datasets/feature_utils.py b/src/lerobot/datasets/feature_utils.py index 137c9f58a..b05dbf2cc 100644 --- a/src/lerobot/datasets/feature_utils.py +++ b/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: diff --git a/src/lerobot/datasets/io_utils.py b/src/lerobot/datasets/io_utils.py index 446d6ca34..c8da4fb94 100644 --- a/src/lerobot/datasets/io_utils.py +++ b/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: diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index 2f0154cda..7cda5d677 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/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.datasets.dataset_reader import DatasetReader -from lerobot.datasets.dataset_writer import DatasetWriter -from lerobot.datasets.utils import ( +from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE + +from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata +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__) diff --git a/src/lerobot/datasets/multi_dataset.py b/src/lerobot/datasets/multi_dataset.py index 092443077..b4b7a941d 100644 --- a/src/lerobot/datasets/multi_dataset.py +++ b/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__) diff --git a/src/lerobot/datasets/pipeline_features.py b/src/lerobot/datasets/pipeline_features.py index 76e2a6511..cf02a52ac 100644 --- a/src/lerobot/datasets/pipeline_features.py +++ b/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 diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index 1767cc79d..f47d71367 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/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.datasets.feature_utils import get_delta_indices -from lerobot.datasets.io_utils import item_to_torch -from lerobot.datasets.utils import ( +from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE + +from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata +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): diff --git a/src/lerobot/datasets/video_utils.py b/src/lerobot/datasets/video_utils.py index 141c2eb64..1060213f5 100644 --- a/src/lerobot/datasets/video_utils.py +++ b/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 diff --git a/src/lerobot/envs/__init__.py b/src/lerobot/envs/__init__.py index 183c12325..adb40490c 100644 --- a/src/lerobot/envs/__init__.py +++ b/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", +] diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index af5bda33f..2a7c52d45 100644 --- a/src/lerobot/envs/configs.py +++ b/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: diff --git a/src/lerobot/envs/factory.py b/src/lerobot/envs/factory.py index 40d5425cc..317cf2e6f 100644 --- a/src/lerobot/envs/factory.py +++ b/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 lerobot.envs.utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result +from .configs import EnvConfig, HubEnvConfig +from .utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result def make_env_config(env_type: str, **kwargs) -> EnvConfig: diff --git a/src/lerobot/envs/libero.py b/src/lerobot/envs/libero.py index 1b814db52..ec90d0ffd 100644 --- a/src/lerobot/envs/libero.py +++ b/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.""" diff --git a/src/lerobot/envs/metaworld.py b/src/lerobot/envs/metaworld.py index 49c775957..1dc513a68 100644 --- a/src/lerobot/envs/metaworld.py +++ b/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: diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index ff5f53735..b0d834a05 100644 --- a/src/lerobot/envs/utils.py +++ b/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.envs.configs import EnvConfig +from lerobot.configs import FeatureType, PolicyFeature 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 = {} diff --git a/src/lerobot/model/__init__.py b/src/lerobot/model/__init__.py new file mode 100644 index 000000000..53e1ed852 --- /dev/null +++ b/src/lerobot/model/__init__.py @@ -0,0 +1,5 @@ +# Kinematics utilities for robot modeling. + +from .kinematics import RobotKinematics as RobotKinematics + +__all__ = ["RobotKinematics"] diff --git a/src/lerobot/motors/__init__.py b/src/lerobot/motors/__init__.py index 5df80d5ba..63ac14a38 100644 --- a/src/lerobot/motors/__init__.py +++ b/src/lerobot/motors/__init__.py @@ -19,3 +19,5 @@ from .motors_bus import ( MotorCalibration, MotorNormMode, ) + +__all__ = ["Motor", "MotorCalibration", "MotorNormMode"] diff --git a/src/lerobot/motors/dynamixel/__init__.py b/src/lerobot/motors/dynamixel/__init__.py index 425f8538a..19f92d7fe 100644 --- a/src/lerobot/motors/dynamixel/__init__.py +++ b/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 * diff --git a/src/lerobot/motors/feetech/__init__.py b/src/lerobot/motors/feetech/__init__.py index 75da2d221..798a79a9a 100644 --- a/src/lerobot/motors/feetech/__init__.py +++ b/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 * diff --git a/src/lerobot/optim/__init__.py b/src/lerobot/optim/__init__.py index de2c4c996..46676027b 100644 --- a/src/lerobot/optim/__init__.py +++ b/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", +] diff --git a/src/lerobot/optim/factory.py b/src/lerobot/optim/factory.py index 699289993..ce519e0b2 100644 --- a/src/lerobot/optim/factory.py +++ b/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( diff --git a/src/lerobot/optim/schedulers.py b/src/lerobot/optim/schedulers.py index f39db6efd..914edd2db 100644 --- a/src/lerobot/optim/schedulers.py +++ b/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} diff --git a/src/lerobot/policies/__init__.py b/src/lerobot/policies/__init__.py index 55ce09cf9..896cc8617 100644 --- a/src/lerobot/policies/__init__.py +++ b/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", ] diff --git a/src/lerobot/policies/act/configuration_act.py b/src/lerobot/policies/act/configuration_act.py index bd89185fd..b5c3d68f1 100644 --- a/src/lerobot/policies/act/configuration_act.py +++ b/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.types import NormalizationMode -from lerobot.optim.optimizers import AdamWConfig +from lerobot.configs import NormalizationMode, PreTrainedConfig +from lerobot.optim import AdamWConfig @PreTrainedConfig.register_subclass("act") diff --git a/src/lerobot/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index a5c48eb3d..0120258ee 100644 --- a/src/lerobot/policies/act/modeling_act.py +++ b/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): """ diff --git a/src/lerobot/policies/act/processor_act.py b/src/lerobot/policies/act/processor_act.py index 727b18cef..d87ade900 100644 --- a/src/lerobot/policies/act/processor_act.py +++ b/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, diff --git a/src/lerobot/policies/diffusion/configuration_diffusion.py b/src/lerobot/policies/diffusion/configuration_diffusion.py index 91b3df214..8e3d4bf19 100644 --- a/src/lerobot/policies/diffusion/configuration_diffusion.py +++ b/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.types import NormalizationMode -from lerobot.optim.optimizers import AdamConfig -from lerobot.optim.schedulers import DiffuserSchedulerConfig +from lerobot.configs import NormalizationMode, PreTrainedConfig +from lerobot.optim import AdamConfig, DiffuserSchedulerConfig @PreTrainedConfig.register_subclass("diffusion") diff --git a/src/lerobot/policies/diffusion/modeling_diffusion.py b/src/lerobot/policies/diffusion/modeling_diffusion.py index 0334f0e70..5b3b97571 100644 --- a/src/lerobot/policies/diffusion/modeling_diffusion.py +++ b/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.policies.pretrained import PreTrainedPolicy -from lerobot.policies.utils import ( +from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE + +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 diff --git a/src/lerobot/policies/diffusion/processor_diffusion.py b/src/lerobot/policies/diffusion/processor_diffusion.py index a7799be64..c4bc17680 100644 --- a/src/lerobot/policies/diffusion/processor_diffusion.py +++ b/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, diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 051a70bcc..611a6e9bc 100644 --- a/src/lerobot/policies/factory.py +++ b/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.types import FeatureType -from lerobot.envs.configs import EnvConfig -from lerobot.envs.utils import env_to_policy_features -from lerobot.policies.act.configuration_act import ACTConfig -from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig -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 ( +from lerobot.configs import FeatureType, PreTrainedConfig +from lerobot.envs import EnvConfig, env_to_policy_features +from lerobot.processor import ( + AbsoluteActionsProcessorStep, + PolicyProcessorPipeline, + RelativeActionsProcessorStep, 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: diff --git a/src/lerobot/policies/groot/__init__.py b/src/lerobot/policies/groot/__init__.py index c8933ff56..4f86d871e 100644 --- a/src/lerobot/policies/groot/__init__.py +++ b/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 diff --git a/src/lerobot/policies/groot/action_head/cross_attention_dit.py b/src/lerobot/policies/groot/action_head/cross_attention_dit.py index 424c475a6..ca40bbe78 100755 --- a/src/lerobot/policies/groot/action_head/cross_attention_dit.py +++ b/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 diff --git a/src/lerobot/policies/groot/action_head/flow_matching_action_head.py b/src/lerobot/policies/groot/action_head/flow_matching_action_head.py index bfc456ba0..4fda21ca5 100644 --- a/src/lerobot/policies/groot/action_head/flow_matching_action_head.py +++ b/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 diff --git a/src/lerobot/policies/groot/configuration_groot.py b/src/lerobot/policies/groot/configuration_groot.py index 4f3d78222..17cb631d7 100644 --- a/src/lerobot/policies/groot/configuration_groot.py +++ b/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.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.optim.optimizers import AdamWConfig -from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig +from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig from lerobot.utils.constants import ACTION, OBS_STATE diff --git a/src/lerobot/policies/groot/groot_n1.py b/src/lerobot/policies/groot/groot_n1.py index 06ff5a04d..fc753839a 100644 --- a/src/lerobot/policies/groot/groot_n1.py +++ b/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 lerobot.utils.constants import ACTION, HF_LEROBOT_HOME +from .utils import ensure_eagle_cache_ready DEFAULT_VENDOR_EAGLE_PATH = str((Path(__file__).resolve().parent / "eagle2_hg_model").resolve()) DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5" diff --git a/src/lerobot/policies/groot/modeling_groot.py b/src/lerobot/policies/groot/modeling_groot.py index 9a479b8f9..4b612bca4 100644 --- a/src/lerobot/policies/groot/modeling_groot.py +++ b/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.policies.groot.configuration_groot import GrootConfig -from lerobot.policies.groot.groot_n1 import GR00TN15 -from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.configs import FeatureType, PolicyFeature 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") diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index 8bf9dabca..3367de711 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/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" diff --git a/src/lerobot/policies/multi_task_dit/configuration_multi_task_dit.py b/src/lerobot/policies/multi_task_dit/configuration_multi_task_dit.py index 061230687..33be3113f 100644 --- a/src/lerobot/policies/multi_task_dit/configuration_multi_task_dit.py +++ b/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.types import NormalizationMode -from lerobot.optim.optimizers import AdamConfig -from lerobot.optim.schedulers import DiffuserSchedulerConfig +from lerobot.configs import NormalizationMode, PreTrainedConfig +from lerobot.optim import AdamConfig, DiffuserSchedulerConfig @PreTrainedConfig.register_subclass("multi_task_dit") diff --git a/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py b/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py index 45d017b68..8e5d1e3cb 100644 --- a/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py +++ b/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 diff --git a/src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py b/src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py index fc94599c2..5f5b9994e 100644 --- a/src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py +++ b/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, diff --git a/src/lerobot/policies/pi0/configuration_pi0.py b/src/lerobot/policies/pi0/configuration_pi0.py index cf4b636a3..a06315f07 100644 --- a/src/lerobot/policies/pi0/configuration_pi0.py +++ b/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.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.optim.optimizers import AdamWConfig -from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig -from lerobot.policies.rtc.configuration_rtc import RTCConfig +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig +from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE +from ..rtc.configuration_rtc import RTCConfig + DEFAULT_IMAGE_SIZE = 224 diff --git a/src/lerobot/policies/pi0/modeling_pi0.py b/src/lerobot/policies/pi0/modeling_pi0.py index aebf32964..22e4e6a26 100644 --- a/src/lerobot/policies/pi0/modeling_pi0.py +++ b/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.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.configs import PreTrainedConfig 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 diff --git a/src/lerobot/policies/pi0/processor_pi0.py b/src/lerobot/policies/pi0/processor_pi0.py index 0302876a1..ad861f85b 100644 --- a/src/lerobot/policies/pi0/processor_pi0.py +++ b/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.policies.pi0.configuration_pi0 import PI0Config +from lerobot.configs import PipelineFeatureType, PolicyFeature 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): diff --git a/src/lerobot/policies/pi05/configuration_pi05.py b/src/lerobot/policies/pi05/configuration_pi05.py index 6760be0a2..124e85cc9 100644 --- a/src/lerobot/policies/pi05/configuration_pi05.py +++ b/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.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.optim.optimizers import AdamWConfig -from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig -from lerobot.policies.rtc.configuration_rtc import RTCConfig +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig +from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE +from ..rtc.configuration_rtc import RTCConfig + DEFAULT_IMAGE_SIZE = 224 diff --git a/src/lerobot/policies/pi05/modeling_pi05.py b/src/lerobot/policies/pi05/modeling_pi05.py index 96c4002f2..a44817a74 100644 --- a/src/lerobot/policies/pi05/modeling_pi05.py +++ b/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.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.configs import PreTrainedConfig 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 diff --git a/src/lerobot/policies/pi05/processor_pi05.py b/src/lerobot/policies/pi05/processor_pi05.py index cb616af87..2d015b24f 100644 --- a/src/lerobot/policies/pi05/processor_pi05.py +++ b/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.policies.pi05.configuration_pi05 import PI05Config +from lerobot.configs import PipelineFeatureType, PolicyFeature 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 diff --git a/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py b/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py index 6a645fae1..e5c6851f4 100644 --- a/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py +++ b/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.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.optim.optimizers import AdamWConfig -from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig -from lerobot.policies.rtc.configuration_rtc import RTCConfig +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig +from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE +from ..rtc.configuration_rtc import RTCConfig + DEFAULT_IMAGE_SIZE = 224 diff --git a/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py b/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py index 1bcf9794c..e86b8ad27 100644 --- a/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py +++ b/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.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.configs import PreTrainedConfig 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 diff --git a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py b/src/lerobot/policies/pi0_fast/processor_pi0_fast.py index c4a510615..60a519786 100644 --- a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py +++ b/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.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig +from lerobot.configs import PipelineFeatureType, PolicyFeature 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 diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 70efeba6f..724f920f3 100644 --- a/src/lerobot/policies/pretrained.py +++ b/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") diff --git a/src/lerobot/policies/rtc/__init__.py b/src/lerobot/policies/rtc/__init__.py index ac7b72ef7..7a29dcac0 100644 --- a/src/lerobot/policies/rtc/__init__.py +++ b/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 lerobot.policies.rtc.action_queue import ActionQueue -from lerobot.policies.rtc.configuration_rtc import RTCConfig -from lerobot.policies.rtc.latency_tracker import LatencyTracker -from lerobot.policies.rtc.modeling_rtc import RTCProcessor +from .action_interpolator import ActionInterpolator +from .action_queue import ActionQueue +from .configuration_rtc import RTCConfig +from .latency_tracker import LatencyTracker +from .modeling_rtc import RTCProcessor __all__ = [ "ActionInterpolator", diff --git a/src/lerobot/policies/rtc/action_queue.py b/src/lerobot/policies/rtc/action_queue.py index 3c20d6d21..dbbdc41df 100644 --- a/src/lerobot/policies/rtc/action_queue.py +++ b/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__) diff --git a/src/lerobot/policies/rtc/configuration_rtc.py b/src/lerobot/policies/rtc/configuration_rtc.py index 70a8dfb09..c70fe3de0 100644 --- a/src/lerobot/policies/rtc/configuration_rtc.py +++ b/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 diff --git a/src/lerobot/policies/rtc/modeling_rtc.py b/src/lerobot/policies/rtc/modeling_rtc.py index 280905adf..c1aeed328 100644 --- a/src/lerobot/policies/rtc/modeling_rtc.py +++ b/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.policies.rtc.configuration_rtc import RTCConfig -from lerobot.policies.rtc.debug_tracker import Tracker +from lerobot.configs import RTCAttentionSchedule + +from .configuration_rtc import RTCConfig +from .debug_tracker import Tracker logger = logging.getLogger(__name__) diff --git a/src/lerobot/policies/sac/configuration_sac.py b/src/lerobot/policies/sac/configuration_sac.py index ada12330c..db0a77672 100644 --- a/src/lerobot/policies/sac/configuration_sac.py +++ b/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.types import NormalizationMode -from lerobot.optim.optimizers import MultiAdamConfig +from lerobot.configs import NormalizationMode, PreTrainedConfig +from lerobot.optim import MultiAdamConfig from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE diff --git a/src/lerobot/policies/sac/modeling_sac.py b/src/lerobot/policies/sac/modeling_sac.py index d5dd71a48..cc7030ce2 100644 --- a/src/lerobot/policies/sac/modeling_sac.py +++ b/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 diff --git a/src/lerobot/policies/sac/processor_sac.py b/src/lerobot/policies/sac/processor_sac.py index cf90e3cb4..3409307c2 100644 --- a/src/lerobot/policies/sac/processor_sac.py +++ b/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, diff --git a/src/lerobot/policies/sac/reward_model/configuration_classifier.py b/src/lerobot/policies/sac/reward_model/configuration_classifier.py index 879e3c1af..3a5bfa424 100644 --- a/src/lerobot/policies/sac/reward_model/configuration_classifier.py +++ b/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.types import NormalizationMode -from lerobot.optim.optimizers import AdamWConfig, OptimizerConfig -from lerobot.optim.schedulers import LRSchedulerConfig +from lerobot.configs import NormalizationMode, PreTrainedConfig +from lerobot.optim import AdamWConfig, LRSchedulerConfig, OptimizerConfig from lerobot.utils.constants import OBS_IMAGE diff --git a/src/lerobot/policies/sac/reward_model/modeling_classifier.py b/src/lerobot/policies/sac/reward_model/modeling_classifier.py index dba6a174b..c8b7efe58 100644 --- a/src/lerobot/policies/sac/reward_model/modeling_classifier.py +++ b/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.""" diff --git a/src/lerobot/policies/sac/reward_model/processor_classifier.py b/src/lerobot/policies/sac/reward_model/processor_classifier.py index c2a34eab2..1f7a66e58 100644 --- a/src/lerobot/policies/sac/reward_model/processor_classifier.py +++ b/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( diff --git a/src/lerobot/policies/sarm/__init__.py b/src/lerobot/policies/sarm/__init__.py new file mode 100644 index 000000000..3d52651d9 --- /dev/null +++ b/src/lerobot/policies/sarm/__init__.py @@ -0,0 +1,3 @@ +from .configuration_sarm import SARMConfig + +__all__ = ["SARMConfig"] diff --git a/src/lerobot/policies/sarm/compute_rabc_weights.py b/src/lerobot/policies/sarm/compute_rabc_weights.py index 485c1096b..07d0780b5 100644 --- a/src/lerobot/policies/sarm/compute_rabc_weights.py +++ b/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.policies.sarm.modeling_sarm import SARMRewardModel -from lerobot.policies.sarm.processor_sarm import make_sarm_pre_post_processors -from lerobot.policies.sarm.sarm_utils import normalize_stage_tau +from lerobot.datasets import LeRobotDataset + +from .modeling_sarm import SARMRewardModel +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: diff --git a/src/lerobot/policies/sarm/configuration_sarm.py b/src/lerobot/policies/sarm/configuration_sarm.py index 673422fe2..fc8daa055 100644 --- a/src/lerobot/policies/sarm/configuration_sarm.py +++ b/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.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.optim.optimizers import AdamWConfig -from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig +from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig from lerobot.utils.constants import OBS_IMAGES, OBS_STATE diff --git a/src/lerobot/policies/sarm/modeling_sarm.py b/src/lerobot/policies/sarm/modeling_sarm.py index 6051d90f8..710554e4b 100644 --- a/src/lerobot/policies/sarm/modeling_sarm.py +++ b/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.policies.sarm.configuration_sarm import SARMConfig -from lerobot.policies.sarm.sarm_utils import ( +from lerobot.utils.constants import OBS_STR + +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): diff --git a/src/lerobot/policies/sarm/processor_sarm.py b/src/lerobot/policies/sarm/processor_sarm.py index f377a7ffa..57fb7c8af 100644 --- a/src/lerobot/policies/sarm/processor_sarm.py +++ b/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.policies.sarm.configuration_sarm import SARMConfig -from lerobot.policies.sarm.sarm_utils import ( +from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.processor 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): diff --git a/src/lerobot/policies/smolvla/configuration_smolvla.py b/src/lerobot/policies/smolvla/configuration_smolvla.py index 5007abbb4..6d5288db3 100644 --- a/src/lerobot/policies/smolvla/configuration_smolvla.py +++ b/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.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.optim.optimizers import AdamWConfig -from lerobot.optim.schedulers import ( - CosineDecayWithWarmupSchedulerConfig, -) -from lerobot.policies.rtc.configuration_rtc import RTCConfig +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig +from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig from lerobot.utils.constants import OBS_IMAGES +from ..rtc.configuration_rtc import RTCConfig + @PreTrainedConfig.register_subclass("smolvla") @dataclass diff --git a/src/lerobot/policies/smolvla/modeling_smolvla.py b/src/lerobot/policies/smolvla/modeling_smolvla.py index 7110ba7d2..ee3ff4db9 100644 --- a/src/lerobot/policies/smolvla/modeling_smolvla.py +++ b/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 diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 3fc130aa1..0aa2beb19 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/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.policies.smolvla.configuration_smolvla import SmolVLAConfig +from lerobot.configs import PipelineFeatureType, PolicyFeature 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, diff --git a/src/lerobot/policies/smolvla/smolvlm_with_expert.py b/src/lerobot/policies/smolvla/smolvlm_with_expert.py index caca41dab..b1d93d72d 100644 --- a/src/lerobot/policies/smolvla/smolvlm_with_expert.py +++ b/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, diff --git a/src/lerobot/policies/tdmpc/configuration_tdmpc.py b/src/lerobot/policies/tdmpc/configuration_tdmpc.py index 3ec493472..bb8a2cf96 100644 --- a/src/lerobot/policies/tdmpc/configuration_tdmpc.py +++ b/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.types import NormalizationMode -from lerobot.optim.optimizers import AdamConfig +from lerobot.configs import NormalizationMode, PreTrainedConfig +from lerobot.optim import AdamConfig @PreTrainedConfig.register_subclass("tdmpc") diff --git a/src/lerobot/policies/tdmpc/modeling_tdmpc.py b/src/lerobot/policies/tdmpc/modeling_tdmpc.py index f83c82e21..a50bb9670 100644 --- a/src/lerobot/policies/tdmpc/modeling_tdmpc.py +++ b/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. diff --git a/src/lerobot/policies/tdmpc/processor_tdmpc.py b/src/lerobot/policies/tdmpc/processor_tdmpc.py index 9b6f97e50..7afe956dc 100644 --- a/src/lerobot/policies/tdmpc/processor_tdmpc.py +++ b/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, diff --git a/src/lerobot/policies/utils.py b/src/lerobot/policies/utils.py index 8c57bcc07..c37127813 100644 --- a/src/lerobot/policies/utils.py +++ b/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.types import FeatureType, PolicyFeature +from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig 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) diff --git a/src/lerobot/policies/vqbet/configuration_vqbet.py b/src/lerobot/policies/vqbet/configuration_vqbet.py index 32906e528..d02745321 100644 --- a/src/lerobot/policies/vqbet/configuration_vqbet.py +++ b/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.types import NormalizationMode -from lerobot.optim.optimizers import AdamConfig -from lerobot.optim.schedulers import VQBeTSchedulerConfig +from lerobot.configs import NormalizationMode, PreTrainedConfig +from lerobot.optim import AdamConfig, VQBeTSchedulerConfig @PreTrainedConfig.register_subclass("vqbet") diff --git a/src/lerobot/policies/vqbet/modeling_vqbet.py b/src/lerobot/policies/vqbet/modeling_vqbet.py index 6d3976b79..153f7fe3c 100644 --- a/src/lerobot/policies/vqbet/modeling_vqbet.py +++ b/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 diff --git a/src/lerobot/policies/vqbet/processor_vqbet.py b/src/lerobot/policies/vqbet/processor_vqbet.py index 1e19ff779..f7b6a061e 100644 --- a/src/lerobot/policies/vqbet/processor_vqbet.py +++ b/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, diff --git a/src/lerobot/policies/vqbet/vqbet_utils.py b/src/lerobot/policies/vqbet/vqbet_utils.py index 7b13577f6..f8bfcb06a 100644 --- a/src/lerobot/policies/vqbet/vqbet_utils.py +++ b/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 diff --git a/src/lerobot/policies/wall_x/__init__.py b/src/lerobot/policies/wall_x/__init__.py index d80c27bda..be4542fde 100644 --- a/src/lerobot/policies/wall_x/__init__.py +++ b/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"] diff --git a/src/lerobot/policies/wall_x/configuration_wall_x.py b/src/lerobot/policies/wall_x/configuration_wall_x.py index 5269c4e10..70576a46b 100644 --- a/src/lerobot/policies/wall_x/configuration_wall_x.py +++ b/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.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.optim.optimizers import AdamWConfig -from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig +from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig from lerobot.utils.constants import ACTION, OBS_STATE diff --git a/src/lerobot/policies/wall_x/modeling_wall_x.py b/src/lerobot/policies/wall_x/modeling_wall_x.py index 84ee05743..a95b42e9d 100644 --- a/src/lerobot/policies/wall_x/modeling_wall_x.py +++ b/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.policies.utils import populate_queues -from lerobot.policies.wall_x.configuration_wall_x import WallXConfig -from lerobot.policies.wall_x.constant import ( +from lerobot.utils.constants import ACTION, OBS_STATE + +from ..pretrained import PreTrainedPolicy +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 lerobot.policies.wall_x.qwen_model.qwen2_5_vl_moe import ( +from .qwen_model.configuration_qwen2_5_vl import Qwen2_5_VLConfig +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__) diff --git a/src/lerobot/policies/wall_x/processor_wall_x.py b/src/lerobot/policies/wall_x/processor_wall_x.py index e4e281541..069cef5d6 100644 --- a/src/lerobot/policies/wall_x/processor_wall_x.py +++ b/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.policies.wall_x.configuration_wall_x import WallXConfig +from lerobot.configs import PipelineFeatureType, PolicyFeature 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, diff --git a/src/lerobot/policies/wall_x/utils.py b/src/lerobot/policies/wall_x/utils.py index e08ef69d5..39a20afce 100644 --- a/src/lerobot/policies/wall_x/utils.py +++ b/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 diff --git a/src/lerobot/policies/xvla/__init__.py b/src/lerobot/policies/xvla/__init__.py index 71b04e76f..84f237ef5 100644 --- a/src/lerobot/policies/xvla/__init__.py +++ b/src/lerobot/policies/xvla/__init__.py @@ -1,5 +1,5 @@ # register the processor steps -from lerobot.policies.xvla.processor_xvla import ( +from .processor_xvla import ( XVLAAddDomainIdProcessorStep, XVLAImageNetNormalizeProcessorStep, XVLAImageToFloatProcessorStep, diff --git a/src/lerobot/policies/xvla/configuration_xvla.py b/src/lerobot/policies/xvla/configuration_xvla.py index 30700b042..614c9a944 100644 --- a/src/lerobot/policies/xvla/configuration_xvla.py +++ b/src/lerobot/policies/xvla/configuration_xvla.py @@ -21,10 +21,8 @@ from __future__ import annotations from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any -from lerobot.configs.policies import PreTrainedConfig -from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.optim.optimizers import XVLAAdamWConfig -from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig +from lerobot.optim import CosineDecayWithWarmupSchedulerConfig, XVLAAdamWConfig from lerobot.utils.constants import OBS_IMAGES # Conditional import for type checking and lazy loading diff --git a/src/lerobot/policies/xvla/modeling_xvla.py b/src/lerobot/policies/xvla/modeling_xvla.py index 0436ae527..0f8bdd036 100644 --- a/src/lerobot/policies/xvla/modeling_xvla.py +++ b/src/lerobot/policies/xvla/modeling_xvla.py @@ -24,15 +24,19 @@ import os from collections import deque from pathlib import Path +from lerobot.utils.import_utils import require_package + +require_package("transformers", extra="xvla") + import torch import torch.nn.functional as F # noqa: N812 from torch import Tensor, nn -from lerobot.configs.policies import PreTrainedConfig -from lerobot.policies.pretrained import PreTrainedPolicy, T -from lerobot.policies.utils import populate_queues +from lerobot.configs import PreTrainedConfig from lerobot.utils.constants import ACTION, OBS_LANGUAGE_TOKENS, OBS_STATE +from ..pretrained import PreTrainedPolicy, T +from ..utils import populate_queues from .action_hub import build_action_space from .configuration_florence2 import Florence2Config from .configuration_xvla import XVLAConfig diff --git a/src/lerobot/policies/xvla/processor_xvla.py b/src/lerobot/policies/xvla/processor_xvla.py index 6fb68587a..0336ec722 100644 --- a/src/lerobot/policies/xvla/processor_xvla.py +++ b/src/lerobot/policies/xvla/processor_xvla.py @@ -20,9 +20,7 @@ from typing import Any import numpy as np import torch -from lerobot.configs.types import PipelineFeatureType, PolicyFeature -from lerobot.policies.xvla.configuration_xvla import XVLAConfig -from lerobot.policies.xvla.utils import rotate6d_to_axis_angle +from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, @@ -35,8 +33,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 ( IMAGENET_STATS, @@ -47,6 +46,9 @@ from lerobot.utils.constants import ( POLICY_PREPROCESSOR_DEFAULT_NAME, ) +from .configuration_xvla import XVLAConfig +from .utils import rotate6d_to_axis_angle + def make_xvla_pre_post_processors( config: XVLAConfig, diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 122b3533c..213bd93f2 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -27,10 +27,15 @@ from .batch_processor import AddBatchDimensionProcessorStep from .converters import ( batch_to_transition, create_transition, + from_tensor_to_numpy, + identity_transition, + policy_action_to_transition, transition_to_batch, + transition_to_policy_action, ) from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep from .device_processor import DeviceProcessorStep +from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep from .factory import ( make_default_processors, make_default_robot_action_processor, @@ -81,7 +86,7 @@ from .relative_action_processor import ( to_absolute_actions, to_relative_actions, ) -from .rename_processor import RenameObservationsProcessorStep +from .rename_processor import RenameObservationsProcessorStep, rename_stats from .tokenizer_processor import ActionTokenizerProcessorStep, TokenizerProcessorStep __all__ = [ @@ -91,6 +96,10 @@ __all__ = [ "ComplementaryDataProcessorStep", "batch_to_transition", "create_transition", + "from_tensor_to_numpy", + "identity_transition", + "policy_action_to_transition", + "transition_to_policy_action", "DeviceProcessorStep", "DoneProcessorStep", "EnvAction", @@ -122,10 +131,13 @@ __all__ = [ "RobotAction", "RobotActionProcessorStep", "RobotObservation", + "rename_stats", "RenameObservationsProcessorStep", "RewardClassifierProcessorStep", "RewardProcessorStep", "DataProcessorPipeline", + "IsaaclabArenaProcessorStep", + "LiberoProcessorStep", "TimeLimitProcessorStep", "AddBatchDimensionProcessorStep", "RobotProcessorPipeline", diff --git a/src/lerobot/processor/batch_processor.py b/src/lerobot/processor/batch_processor.py index c904acf84..eb7db255a 100644 --- a/src/lerobot/processor/batch_processor.py +++ b/src/lerobot/processor/batch_processor.py @@ -24,7 +24,7 @@ from dataclasses import dataclass, field from torch import Tensor -from lerobot.configs.types import PipelineFeatureType, PolicyFeature +from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.types import EnvTransition, PolicyAction from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE diff --git a/src/lerobot/processor/delta_action_processor.py b/src/lerobot/processor/delta_action_processor.py index f7f5676ac..86b2feec1 100644 --- a/src/lerobot/processor/delta_action_processor.py +++ b/src/lerobot/processor/delta_action_processor.py @@ -16,7 +16,7 @@ from dataclasses import dataclass -from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.types import PolicyAction, RobotAction from .pipeline import ActionProcessorStep, ProcessorStepRegistry, RobotActionProcessorStep diff --git a/src/lerobot/processor/device_processor.py b/src/lerobot/processor/device_processor.py index 36c80e58e..1171c7e78 100644 --- a/src/lerobot/processor/device_processor.py +++ b/src/lerobot/processor/device_processor.py @@ -24,7 +24,7 @@ from typing import Any import torch -from lerobot.configs.types import PipelineFeatureType, PolicyFeature +from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.types import EnvTransition, PolicyAction, TransitionKey from lerobot.utils.device_utils import get_safe_torch_device diff --git a/src/lerobot/processor/env_processor.py b/src/lerobot/processor/env_processor.py index a77e066cf..75cbb79de 100644 --- a/src/lerobot/processor/env_processor.py +++ b/src/lerobot/processor/env_processor.py @@ -17,7 +17,7 @@ from dataclasses import dataclass import torch -from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.utils.constants import OBS_IMAGES, OBS_PREFIX, OBS_STATE, OBS_STR from .pipeline import ObservationProcessorStep, ProcessorStepRegistry diff --git a/src/lerobot/processor/gym_action_processor.py b/src/lerobot/processor/gym_action_processor.py index e756ded7f..2ec5f6e64 100644 --- a/src/lerobot/processor/gym_action_processor.py +++ b/src/lerobot/processor/gym_action_processor.py @@ -16,8 +16,8 @@ from dataclasses import dataclass -from lerobot.configs.types import PipelineFeatureType, PolicyFeature -from lerobot.types import EnvAction, EnvTransition, PolicyAction +from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.types import EnvAction, EnvTransition, PolicyAction, TransitionKey from .converters import to_tensor from .hil_processor import TELEOP_ACTION_KEY @@ -75,8 +75,6 @@ class Numpy2TorchActionProcessorStep(ProcessorStep): def __call__(self, transition: EnvTransition) -> EnvTransition: """Converts numpy action to torch tensor if action exists, otherwise passes through.""" - from lerobot.types import TransitionKey - self._current_transition = transition.copy() new_transition = self._current_transition diff --git a/src/lerobot/processor/hil_processor.py b/src/lerobot/processor/hil_processor.py index 0b8521c2b..c6f98c689 100644 --- a/src/lerobot/processor/hil_processor.py +++ b/src/lerobot/processor/hil_processor.py @@ -24,7 +24,7 @@ import numpy as np import torch import torchvision.transforms.functional as F # noqa: N812 -from lerobot.configs.types import PipelineFeatureType, PolicyFeature +from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.teleoperators.utils import TeleopEvents if TYPE_CHECKING: diff --git a/src/lerobot/processor/migrate_policy_normalization.py b/src/lerobot/processor/migrate_policy_normalization.py index 525b7431c..37df4be41 100644 --- a/src/lerobot/processor/migrate_policy_normalization.py +++ b/src/lerobot/processor/migrate_policy_normalization.py @@ -57,8 +57,8 @@ import torch from huggingface_hub import HfApi, hf_hub_download from safetensors.torch import load_file as load_safetensors -from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature -from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature +from lerobot.policies import get_policy_class, make_policy_config, make_pre_post_processors from lerobot.utils.constants import ACTION diff --git a/src/lerobot/processor/normalize_processor.py b/src/lerobot/processor/normalize_processor.py index 9099538d2..7516c7b47 100644 --- a/src/lerobot/processor/normalize_processor.py +++ b/src/lerobot/processor/normalize_processor.py @@ -24,11 +24,11 @@ from typing import TYPE_CHECKING, Any import torch from torch import Tensor -from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature +from lerobot.configs import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature from lerobot.types import EnvTransition, PolicyAction, TransitionKey if TYPE_CHECKING: - from lerobot.datasets.lerobot_dataset import LeRobotDataset + from lerobot.datasets import LeRobotDataset from lerobot.utils.constants import ACTION diff --git a/src/lerobot/processor/observation_processor.py b/src/lerobot/processor/observation_processor.py index d22d8fb96..12d1f82a2 100644 --- a/src/lerobot/processor/observation_processor.py +++ b/src/lerobot/processor/observation_processor.py @@ -20,7 +20,7 @@ import numpy as np import torch from torch import Tensor -from lerobot.configs.types import PipelineFeatureType, PolicyFeature +from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR from .pipeline import ObservationProcessorStep, ProcessorStepRegistry diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index abfb31421..2b949d5cb 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -45,8 +45,9 @@ import torch from huggingface_hub import hf_hub_download from safetensors.torch import load_file, save_file -from lerobot.configs.types import PipelineFeatureType, PolicyFeature +from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.types import EnvAction, EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey +from lerobot.utils.constants import HF_LEROBOT_HOME from lerobot.utils.hub import HubMixin from .converters import batch_to_transition, create_transition, transition_to_batch @@ -422,8 +423,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): """ if save_directory is None: # Use default directory in HF_LEROBOT_HOME - from lerobot.utils.constants import HF_LEROBOT_HOME - sanitized_name = re.sub(r"[^a-zA-Z0-9_]", "_", self.name.lower()) save_directory = HF_LEROBOT_HOME / "processors" / sanitized_name diff --git a/src/lerobot/processor/policy_robot_bridge.py b/src/lerobot/processor/policy_robot_bridge.py index 25887d414..1c5327bb5 100644 --- a/src/lerobot/processor/policy_robot_bridge.py +++ b/src/lerobot/processor/policy_robot_bridge.py @@ -19,7 +19,7 @@ from typing import Any import torch -from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.processor import ActionProcessorStep, PolicyAction, ProcessorStepRegistry, RobotAction from lerobot.utils.constants import ACTION diff --git a/src/lerobot/processor/relative_action_processor.py b/src/lerobot/processor/relative_action_processor.py index e00d26e98..d9f97f2c6 100644 --- a/src/lerobot/processor/relative_action_processor.py +++ b/src/lerobot/processor/relative_action_processor.py @@ -19,7 +19,7 @@ from typing import Any import torch from torch import Tensor -from lerobot.configs.types import PipelineFeatureType, PolicyFeature +from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.types import EnvTransition, TransitionKey from lerobot.utils.constants import OBS_STATE diff --git a/src/lerobot/processor/rename_processor.py b/src/lerobot/processor/rename_processor.py index 6cae5921f..5ffec6868 100644 --- a/src/lerobot/processor/rename_processor.py +++ b/src/lerobot/processor/rename_processor.py @@ -17,7 +17,7 @@ from copy import deepcopy from dataclasses import dataclass, field from typing import Any -from lerobot.configs.types import PipelineFeatureType, PolicyFeature +from lerobot.configs import PipelineFeatureType, PolicyFeature from .pipeline import ObservationProcessorStep, ProcessorStepRegistry diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py index 0b5305dcf..a808e6127 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -29,7 +29,7 @@ from typing import TYPE_CHECKING, Any import torch -from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.types import EnvTransition, RobotObservation, TransitionKey from lerobot.utils.constants import ( ACTION_TOKEN_MASK, diff --git a/src/lerobot/rl/__init__.py b/src/lerobot/rl/__init__.py new file mode 100644 index 000000000..41037b98d --- /dev/null +++ b/src/lerobot/rl/__init__.py @@ -0,0 +1,6 @@ +# Reinforcement learning modules. +# Requires: lerobot[hilserl] + +from lerobot.utils.import_utils import require_package + +require_package("grpcio", extra="hilserl", import_name="grpc") diff --git a/src/lerobot/rl/actor.py b/src/lerobot/rl/actor.py index 18c0ca1ea..0d785bde3 100644 --- a/src/lerobot/rl/actor.py +++ b/src/lerobot/rl/actor.py @@ -60,10 +60,8 @@ from torch.multiprocessing import Event, Queue from lerobot.cameras import opencv # noqa: F401 from lerobot.configs import parser from lerobot.configs.train import TrainRLServerPipelineConfig -from lerobot.policies.factory import make_policy +from lerobot.policies import make_policy from lerobot.policies.sac.modeling_sac import SACPolicy -from lerobot.rl.process import ProcessSignalHandler -from lerobot.rl.queue import get_last_item_from_queue from lerobot.robots import so_follower # noqa: F401 from lerobot.teleoperators import gamepad, so_leader # noqa: F401 from lerobot.teleoperators.utils import TeleopEvents @@ -96,6 +94,8 @@ from .gym_manipulator import ( make_robot_env, step_env_and_process_transition, ) +from .process import ProcessSignalHandler +from .queue import get_last_item_from_queue # Main entry point diff --git a/src/lerobot/rl/buffer.py b/src/lerobot/rl/buffer.py index 68954162d..97aaa9caa 100644 --- a/src/lerobot/rl/buffer.py +++ b/src/lerobot/rl/buffer.py @@ -23,7 +23,7 @@ import torch import torch.nn.functional as F # noqa: N812 from tqdm import tqdm -from lerobot.datasets.lerobot_dataset import LeRobotDataset +from lerobot.datasets import LeRobotDataset from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, REWARD from lerobot.utils.transition import Transition diff --git a/src/lerobot/rl/crop_dataset_roi.py b/src/lerobot/rl/crop_dataset_roi.py index 4345fed3c..b6bde2273 100644 --- a/src/lerobot/rl/crop_dataset_roi.py +++ b/src/lerobot/rl/crop_dataset_roi.py @@ -24,7 +24,7 @@ import torch import torchvision.transforms.functional as F # type: ignore # noqa: N812 from tqdm import tqdm # type: ignore -from lerobot.datasets.lerobot_dataset import LeRobotDataset +from lerobot.datasets import LeRobotDataset from lerobot.utils.constants import DONE, REWARD diff --git a/src/lerobot/rl/eval_policy.py b/src/lerobot/rl/eval_policy.py index fb2504f2a..4398351c5 100644 --- a/src/lerobot/rl/eval_policy.py +++ b/src/lerobot/rl/eval_policy.py @@ -18,8 +18,8 @@ import logging from lerobot.cameras import opencv # noqa: F401 from lerobot.configs import parser from lerobot.configs.train import TrainRLServerPipelineConfig -from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.policies.factory import make_policy +from lerobot.datasets import LeRobotDataset +from lerobot.policies import make_policy from lerobot.robots import ( # noqa: F401 RobotConfig, make_robot_from_config, diff --git a/src/lerobot/rl/gym_manipulator.py b/src/lerobot/rl/gym_manipulator.py index bd64d205f..b6ff7155a 100644 --- a/src/lerobot/rl/gym_manipulator.py +++ b/src/lerobot/rl/gym_manipulator.py @@ -25,9 +25,9 @@ import torch from lerobot.cameras import opencv # noqa: F401 from lerobot.configs import parser -from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.envs.configs import HILSerlRobotEnvConfig -from lerobot.model.kinematics import RobotKinematics +from lerobot.datasets import LeRobotDataset +from lerobot.envs import HILSerlRobotEnvConfig +from lerobot.model import RobotKinematics from lerobot.processor import ( AddBatchDimensionProcessorStep, AddTeleopActionAsComplimentaryDataStep, @@ -50,8 +50,8 @@ from lerobot.processor import ( TransitionKey, VanillaObservationProcessorStep, create_transition, + identity_transition, ) -from lerobot.processor.converters import identity_transition from lerobot.robots import ( # noqa: F401 RobotConfig, make_robot_from_config, diff --git a/src/lerobot/rl/joint_observations_processor.py b/src/lerobot/rl/joint_observations_processor.py index 2fbcc7c46..dc677e26c 100644 --- a/src/lerobot/rl/joint_observations_processor.py +++ b/src/lerobot/rl/joint_observations_processor.py @@ -19,8 +19,8 @@ from typing import Any import torch -from lerobot.configs.types import PipelineFeatureType, PolicyFeature -from lerobot.processor.pipeline import ( +from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.processor import ( ObservationProcessorStep, ProcessorStepRegistry, ) diff --git a/src/lerobot/rl/learner.py b/src/lerobot/rl/learner.py index 7034b11d8..396d6341e 100644 --- a/src/lerobot/rl/learner.py +++ b/src/lerobot/rl/learner.py @@ -63,11 +63,8 @@ from lerobot.cameras import opencv # noqa: F401 from lerobot.configs import parser from lerobot.configs.train import TrainRLServerPipelineConfig from lerobot.datasets import LeRobotDataset, make_dataset -from lerobot.policies.factory import make_policy +from lerobot.policies import make_policy from lerobot.policies.sac.modeling_sac import SACPolicy -from lerobot.rl.buffer import ReplayBuffer, concatenate_batch_transitions -from lerobot.rl.process import ProcessSignalHandler -from lerobot.rl.wandb_utils import WandBLogger from lerobot.robots import so_follower # noqa: F401 from lerobot.teleoperators import gamepad, so_leader # noqa: F401 from lerobot.teleoperators.utils import TeleopEvents @@ -99,7 +96,10 @@ from lerobot.utils.utils import ( init_logging, ) +from .buffer import ReplayBuffer, concatenate_batch_transitions from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService +from .process import ProcessSignalHandler +from .wandb_utils import WandBLogger @parser.wrap() @@ -151,7 +151,7 @@ def train(cfg: TrainRLServerPipelineConfig, job_name: str | None = None): # Setup WandB logging if enabled if cfg.wandb.enable and cfg.wandb.project: - from lerobot.rl.wandb_utils import WandBLogger + from .wandb_utils import WandBLogger wandb_logger = WandBLogger(cfg) else: diff --git a/src/lerobot/rl/learner_service.py b/src/lerobot/rl/learner_service.py index 7ef38119b..4128cdf55 100644 --- a/src/lerobot/rl/learner_service.py +++ b/src/lerobot/rl/learner_service.py @@ -19,10 +19,11 @@ import logging import time from multiprocessing import Event, Queue -from lerobot.rl.queue import get_last_item_from_queue from lerobot.transport import services_pb2, services_pb2_grpc from lerobot.transport.utils import receive_bytes_in_chunks, send_bytes_in_chunks +from .queue import get_last_item_from_queue + MAX_WORKERS = 3 # Stream parameters, send transitions and interactions SHUTDOWN_TIMEOUT = 10 diff --git a/src/lerobot/robots/__init__.py b/src/lerobot/robots/__init__.py index 1dba0f1b0..eb8b06fb8 100644 --- a/src/lerobot/robots/__init__.py +++ b/src/lerobot/robots/__init__.py @@ -17,3 +17,5 @@ from .config import RobotConfig from .robot import Robot from .utils import make_robot_from_config + +__all__ = ["Robot", "RobotConfig", "make_robot_from_config"] diff --git a/src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py b/src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py index c48ac5934..c27398278 100644 --- a/src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py +++ b/src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py @@ -17,10 +17,10 @@ import logging from functools import cached_property -from lerobot.robots.openarm_follower import OpenArmFollower, OpenArmFollowerConfig from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected +from ..openarm_follower import OpenArmFollower, OpenArmFollowerConfig from ..robot import Robot from .config_bi_openarm_follower import BiOpenArmFollowerConfig diff --git a/src/lerobot/robots/bi_openarm_follower/config_bi_openarm_follower.py b/src/lerobot/robots/bi_openarm_follower/config_bi_openarm_follower.py index ef5d70cab..9ed56aeac 100644 --- a/src/lerobot/robots/bi_openarm_follower/config_bi_openarm_follower.py +++ b/src/lerobot/robots/bi_openarm_follower/config_bi_openarm_follower.py @@ -17,9 +17,9 @@ from dataclasses import dataclass, field from lerobot.cameras import CameraConfig -from lerobot.robots.openarm_follower import OpenArmFollowerConfigBase from ..config import RobotConfig +from ..openarm_follower import OpenArmFollowerConfigBase @RobotConfig.register_subclass("bi_openarm_follower") diff --git a/src/lerobot/robots/bi_so_follower/bi_so_follower.py b/src/lerobot/robots/bi_so_follower/bi_so_follower.py index ba1826e29..f592150a6 100644 --- a/src/lerobot/robots/bi_so_follower/bi_so_follower.py +++ b/src/lerobot/robots/bi_so_follower/bi_so_follower.py @@ -17,11 +17,11 @@ import logging from functools import cached_property -from lerobot.robots.so_follower import SOFollower, SOFollowerRobotConfig from lerobot.types import RobotAction, RobotObservation from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from ..robot import Robot +from ..so_follower import SOFollower, SOFollowerRobotConfig from .config_bi_so_follower import BiSOFollowerConfig logger = logging.getLogger(__name__) diff --git a/src/lerobot/robots/bi_so_follower/config_bi_so_follower.py b/src/lerobot/robots/bi_so_follower/config_bi_so_follower.py index dca74fa2d..97afbab4f 100644 --- a/src/lerobot/robots/bi_so_follower/config_bi_so_follower.py +++ b/src/lerobot/robots/bi_so_follower/config_bi_so_follower.py @@ -16,9 +16,8 @@ from dataclasses import dataclass -from lerobot.robots.so_follower import SOFollowerConfig - from ..config import RobotConfig +from ..so_follower import SOFollowerConfig @RobotConfig.register_subclass("bi_so_follower") diff --git a/src/lerobot/robots/hope_jr/hope_jr_arm.py b/src/lerobot/robots/hope_jr/hope_jr_arm.py index 7f6492ef0..4918bcae3 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_arm.py +++ b/src/lerobot/robots/hope_jr/hope_jr_arm.py @@ -18,7 +18,7 @@ import logging import time from functools import cached_property -from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.cameras import make_cameras_from_configs from lerobot.motors import Motor, MotorNormMode from lerobot.motors.calibration_gui import RangeFinderGUI from lerobot.motors.feetech import ( diff --git a/src/lerobot/robots/hope_jr/hope_jr_hand.py b/src/lerobot/robots/hope_jr/hope_jr_hand.py index 784804836..566628724 100644 --- a/src/lerobot/robots/hope_jr/hope_jr_hand.py +++ b/src/lerobot/robots/hope_jr/hope_jr_hand.py @@ -18,7 +18,7 @@ import logging import time from functools import cached_property -from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.cameras import make_cameras_from_configs from lerobot.motors import Motor, MotorNormMode from lerobot.motors.calibration_gui import RangeFinderGUI from lerobot.motors.feetech import ( diff --git a/src/lerobot/robots/koch_follower/koch_follower.py b/src/lerobot/robots/koch_follower/koch_follower.py index 44e83f6a3..3f40ac738 100644 --- a/src/lerobot/robots/koch_follower/koch_follower.py +++ b/src/lerobot/robots/koch_follower/koch_follower.py @@ -18,7 +18,7 @@ import logging import time from functools import cached_property -from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.cameras import make_cameras_from_configs from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.dynamixel import ( DynamixelMotorsBus, diff --git a/src/lerobot/robots/lekiwi/__init__.py b/src/lerobot/robots/lekiwi/__init__.py index ada2ff368..d98bde590 100644 --- a/src/lerobot/robots/lekiwi/__init__.py +++ b/src/lerobot/robots/lekiwi/__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("pyzmq", extra="lekiwi", import_name="zmq") + from .config_lekiwi import LeKiwiClientConfig, LeKiwiConfig from .lekiwi import LeKiwi from .lekiwi_client import LeKiwiClient diff --git a/src/lerobot/robots/lekiwi/config_lekiwi.py b/src/lerobot/robots/lekiwi/config_lekiwi.py index acaf5f0ec..51fa8f03f 100644 --- a/src/lerobot/robots/lekiwi/config_lekiwi.py +++ b/src/lerobot/robots/lekiwi/config_lekiwi.py @@ -14,8 +14,8 @@ from dataclasses import dataclass, field -from lerobot.cameras.configs import CameraConfig, Cv2Rotation -from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig +from lerobot.cameras import CameraConfig, Cv2Rotation +from lerobot.cameras.opencv import OpenCVCameraConfig from ..config import RobotConfig diff --git a/src/lerobot/robots/lekiwi/lekiwi.py b/src/lerobot/robots/lekiwi/lekiwi.py index 60fac89e5..b73ebeab9 100644 --- a/src/lerobot/robots/lekiwi/lekiwi.py +++ b/src/lerobot/robots/lekiwi/lekiwi.py @@ -22,7 +22,7 @@ from typing import Any import numpy as np -from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.cameras import make_cameras_from_configs from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.feetech import ( FeetechMotorsBus, diff --git a/src/lerobot/robots/omx_follower/omx_follower.py b/src/lerobot/robots/omx_follower/omx_follower.py index 5d161daa2..c30eec97a 100644 --- a/src/lerobot/robots/omx_follower/omx_follower.py +++ b/src/lerobot/robots/omx_follower/omx_follower.py @@ -18,7 +18,7 @@ import logging import time from functools import cached_property -from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.cameras import make_cameras_from_configs from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.dynamixel import ( DriveMode, diff --git a/src/lerobot/robots/openarm_follower/openarm_follower.py b/src/lerobot/robots/openarm_follower/openarm_follower.py index 99e8b920b..4d1765f07 100644 --- a/src/lerobot/robots/openarm_follower/openarm_follower.py +++ b/src/lerobot/robots/openarm_follower/openarm_follower.py @@ -19,7 +19,7 @@ import time from functools import cached_property from typing import Any -from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.cameras import make_cameras_from_configs from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.damiao import DamiaoMotorsBus from lerobot.types import RobotAction, RobotObservation diff --git a/src/lerobot/robots/reachy2/configuration_reachy2.py b/src/lerobot/robots/reachy2/configuration_reachy2.py index 63293e675..8cb67a495 100644 --- a/src/lerobot/robots/reachy2/configuration_reachy2.py +++ b/src/lerobot/robots/reachy2/configuration_reachy2.py @@ -14,8 +14,7 @@ from dataclasses import dataclass, field -from lerobot.cameras import CameraConfig -from lerobot.cameras.configs import ColorMode +from lerobot.cameras import CameraConfig, ColorMode from lerobot.cameras.reachy2_camera import Reachy2CameraConfig from ..config import RobotConfig diff --git a/src/lerobot/robots/reachy2/robot_reachy2.py b/src/lerobot/robots/reachy2/robot_reachy2.py index 5227a096a..ef55f71b9 100644 --- a/src/lerobot/robots/reachy2/robot_reachy2.py +++ b/src/lerobot/robots/reachy2/robot_reachy2.py @@ -18,7 +18,7 @@ from __future__ import annotations import time from typing import TYPE_CHECKING, Any -from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.cameras import make_cameras_from_configs from lerobot.types import RobotAction, RobotObservation from lerobot.utils.import_utils import _reachy2_sdk_available diff --git a/src/lerobot/robots/so_follower/robot_kinematic_processor.py b/src/lerobot/robots/so_follower/robot_kinematic_processor.py index 2aa60e12a..8114fdc2c 100644 --- a/src/lerobot/robots/so_follower/robot_kinematic_processor.py +++ b/src/lerobot/robots/so_follower/robot_kinematic_processor.py @@ -19,8 +19,8 @@ from typing import Any import numpy as np -from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature -from lerobot.model.kinematics import RobotKinematics +from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.model import RobotKinematics from lerobot.processor import ( EnvTransition, ObservationProcessorStep, diff --git a/src/lerobot/robots/so_follower/so_follower.py b/src/lerobot/robots/so_follower/so_follower.py index ca132d102..0651f566c 100644 --- a/src/lerobot/robots/so_follower/so_follower.py +++ b/src/lerobot/robots/so_follower/so_follower.py @@ -18,7 +18,7 @@ import logging import time from functools import cached_property -from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.cameras import make_cameras_from_configs from lerobot.motors import Motor, MotorCalibration, MotorNormMode from lerobot.motors.feetech import ( FeetechMotorsBus, diff --git a/src/lerobot/robots/unitree_g1/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/gr00t_locomotion.py index 31166e123..12fe26073 100644 --- a/src/lerobot/robots/unitree_g1/gr00t_locomotion.py +++ b/src/lerobot/robots/unitree_g1/gr00t_locomotion.py @@ -21,7 +21,7 @@ import numpy as np import onnxruntime as ort from huggingface_hub import hf_hub_download -from lerobot.robots.unitree_g1.g1_utils import ( +from .g1_utils import ( REMOTE_AXES, REMOTE_BUTTONS, G1_29_JointIndex, diff --git a/src/lerobot/robots/unitree_g1/holosoma_locomotion.py b/src/lerobot/robots/unitree_g1/holosoma_locomotion.py index 857bb97bc..3d3bccbdc 100644 --- a/src/lerobot/robots/unitree_g1/holosoma_locomotion.py +++ b/src/lerobot/robots/unitree_g1/holosoma_locomotion.py @@ -22,7 +22,7 @@ import onnx import onnxruntime as ort from huggingface_hub import hf_hub_download -from lerobot.robots.unitree_g1.g1_utils import ( +from .g1_utils import ( REMOTE_AXES, G1_29_JointArmIndex, G1_29_JointIndex, diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 9e373c05f..785861a5a 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -25,9 +25,14 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable import numpy as np -from lerobot.cameras.utils import make_cameras_from_configs -from lerobot.robots.unitree_g1.g1_kinematics import G1_29_ArmIK -from lerobot.robots.unitree_g1.g1_utils import ( +from lerobot.cameras import make_cameras_from_configs +from lerobot.types import RobotAction, RobotObservation +from lerobot.utils.import_utils import _unitree_sdk_available + +from ..robot import Robot +from .config_unitree_g1 import UnitreeG1Config +from .g1_kinematics import G1_29_ArmIK +from .g1_utils import ( REMOTE_AXES, REMOTE_KEYS, G1_29_JointArmIndex, @@ -35,11 +40,6 @@ from lerobot.robots.unitree_g1.g1_utils import ( default_remote_input, make_locomotion_controller, ) -from lerobot.types import RobotAction, RobotObservation -from lerobot.utils.import_utils import _unitree_sdk_available - -from ..robot import Robot -from .config_unitree_g1 import UnitreeG1Config if TYPE_CHECKING or _unitree_sdk_available: from unitree_sdk2py.core.channel import ( @@ -127,7 +127,7 @@ class UnitreeG1(Robot): self._ChannelPublisher = _SDKChannelPublisher self._ChannelSubscriber = _SDKChannelSubscriber else: - from lerobot.robots.unitree_g1.unitree_sdk2_socket import ( + from .unitree_sdk2_socket import ( ChannelFactoryInitialize, ChannelPublisher, ChannelSubscriber, @@ -290,7 +290,7 @@ class UnitreeG1(Robot): def connect(self, calibrate: bool = True) -> None: # connect to DDS # Initialize DDS channel and simulation environment if self.config.is_simulation: - from lerobot.envs.factory import make_env + from lerobot.envs import make_env self._ChannelFactoryInitialize(0, "lo") self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True) diff --git a/src/lerobot/robots/unitree_g1/unitree_sdk2_socket.py b/src/lerobot/robots/unitree_g1/unitree_sdk2_socket.py index 0f1f8f8d6..4f0b787aa 100644 --- a/src/lerobot/robots/unitree_g1/unitree_sdk2_socket.py +++ b/src/lerobot/robots/unitree_g1/unitree_sdk2_socket.py @@ -20,7 +20,7 @@ from typing import Any import zmq -from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config +from .config_unitree_g1 import UnitreeG1Config # Module-level ZMQ state mirrors the Unitree SDK's global ChannelFactory Singleton. # Only one robot connection per process is supported. diff --git a/src/lerobot/scripts/augment_dataset_quantile_stats.py b/src/lerobot/scripts/augment_dataset_quantile_stats.py index 7ec066b0a..a5b9edf63 100644 --- a/src/lerobot/scripts/augment_dataset_quantile_stats.py +++ b/src/lerobot/scripts/augment_dataset_quantile_stats.py @@ -44,9 +44,14 @@ from huggingface_hub import HfApi from requests import HTTPError from tqdm import tqdm -from lerobot.datasets import LeRobotDataset, write_stats -from lerobot.datasets.compute_stats import DEFAULT_QUANTILES, aggregate_stats, get_feature_stats -from lerobot.datasets.dataset_metadata import CODEBASE_VERSION +from lerobot.datasets import ( + CODEBASE_VERSION, + LeRobotDataset, + aggregate_stats, + get_feature_stats, + write_stats, +) +from lerobot.datasets.compute_stats import DEFAULT_QUANTILES from lerobot.utils.utils import init_logging diff --git a/src/lerobot/scripts/convert_dataset_v21_to_v30.py b/src/lerobot/scripts/convert_dataset_v21_to_v30.py index 3f76726fd..34a4d73b7 100644 --- a/src/lerobot/scripts/convert_dataset_v21_to_v30.py +++ b/src/lerobot/scripts/convert_dataset_v21_to_v30.py @@ -59,9 +59,7 @@ from datasets import Dataset, Features, Image from huggingface_hub import HfApi, snapshot_download from requests import HTTPError -from lerobot.datasets import LeRobotDataset -from lerobot.datasets.compute_stats import aggregate_stats -from lerobot.datasets.dataset_metadata import CODEBASE_VERSION +from lerobot.datasets import CODEBASE_VERSION, LeRobotDataset, aggregate_stats from lerobot.datasets.io_utils import ( cast_stats_to_numpy, get_file_size_in_mb, diff --git a/src/lerobot/scripts/lerobot_calibrate.py b/src/lerobot/scripts/lerobot_calibrate.py index 242067978..192daa69f 100644 --- a/src/lerobot/scripts/lerobot_calibrate.py +++ b/src/lerobot/scripts/lerobot_calibrate.py @@ -31,8 +31,8 @@ from pprint import pformat import draccus -from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 -from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 +from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py index 4d0a63f00..162bf7e41 100644 --- a/src/lerobot/scripts/lerobot_edit_dataset.py +++ b/src/lerobot/scripts/lerobot_edit_dataset.py @@ -178,8 +178,8 @@ from pathlib import Path import draccus from lerobot.configs import parser -from lerobot.datasets import LeRobotDataset -from lerobot.datasets.dataset_tools import ( +from lerobot.datasets import ( + LeRobotDataset, convert_image_to_video_dataset, delete_episodes, merge_datasets, diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index cd912280f..9f13bcc1e 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -71,14 +71,14 @@ from tqdm import trange from lerobot.configs import parser from lerobot.configs.eval import EvalPipelineConfig -from lerobot.envs.factory import make_env, make_env_pre_post_processors -from lerobot.envs.utils import ( +from lerobot.envs import ( check_env_attributes_and_types, close_envs, + make_env, + make_env_pre_post_processors, preprocess_observation, ) -from lerobot.policies.factory import make_policy, make_pre_post_processors -from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.processor import PolicyProcessorPipeline from lerobot.types import PolicyAction from lerobot.utils.constants import ACTION, DONE, OBS_STR, REWARD diff --git a/src/lerobot/scripts/lerobot_find_cameras.py b/src/lerobot/scripts/lerobot_find_cameras.py index 0248a2768..ab999f7ae 100644 --- a/src/lerobot/scripts/lerobot_find_cameras.py +++ b/src/lerobot/scripts/lerobot_find_cameras.py @@ -37,11 +37,11 @@ from typing import Any import numpy as np from PIL import Image -from lerobot.cameras.configs import ColorMode +from lerobot.cameras import ColorMode +from lerobot.cameras.opencv import OpenCVCameraConfig from lerobot.cameras.opencv.camera_opencv import OpenCVCamera -from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig +from lerobot.cameras.realsense import RealSenseCameraConfig from lerobot.cameras.realsense.camera_realsense import RealSenseCamera -from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig logger = logging.getLogger(__name__) diff --git a/src/lerobot/scripts/lerobot_find_joint_limits.py b/src/lerobot/scripts/lerobot_find_joint_limits.py index bcb93ba12..c4f867631 100644 --- a/src/lerobot/scripts/lerobot_find_joint_limits.py +++ b/src/lerobot/scripts/lerobot_find_joint_limits.py @@ -41,7 +41,7 @@ from dataclasses import dataclass import draccus import numpy as np -from lerobot.model.kinematics import RobotKinematics +from lerobot.model import RobotKinematics from lerobot.robots import ( # noqa: F401 RobotConfig, bi_openarm_follower, diff --git a/src/lerobot/scripts/lerobot_imgtransform_viz.py b/src/lerobot/scripts/lerobot_imgtransform_viz.py index 0b9d697af..7cd4c782d 100644 --- a/src/lerobot/scripts/lerobot_imgtransform_viz.py +++ b/src/lerobot/scripts/lerobot_imgtransform_viz.py @@ -35,7 +35,7 @@ from pathlib import Path import draccus from torchvision.transforms import ToPILImage -from lerobot.configs.default import DatasetConfig +from lerobot.configs import DatasetConfig from lerobot.datasets import LeRobotDataset from lerobot.transforms import ( ImageTransforms, diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 063489238..739a3a1ff 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -76,15 +76,12 @@ from typing import Any import torch -from lerobot.cameras import ( # noqa: F401 - CameraConfig, # noqa: F401 -) -from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 -from lerobot.cameras.reachy2_camera.configuration_reachy2_camera import Reachy2CameraConfig # noqa: F401 -from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 -from lerobot.cameras.zmq.configuration_zmq import ZMQCameraConfig # noqa: F401 -from lerobot.configs import parser -from lerobot.configs.policies import PreTrainedConfig +from lerobot.cameras import CameraConfig # noqa: F401 +from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 +from lerobot.cameras.reachy2_camera import Reachy2CameraConfig # noqa: F401 +from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.cameras.zmq import ZMQCameraConfig # noqa: F401 +from lerobot.configs import PreTrainedConfig, parser from lerobot.datasets import ( LeRobotDataset, VideoEncodingManager, @@ -92,10 +89,13 @@ from lerobot.datasets import ( create_initial_features, safe_stop_image_writer, ) -from lerobot.policies.factory import make_policy, make_pre_post_processors -from lerobot.policies.pretrained import PreTrainedPolicy -from lerobot.policies.rtc import ActionInterpolator -from lerobot.policies.utils import make_robot_action +from lerobot.policies import ( + ActionInterpolator, + PreTrainedPolicy, + make_policy, + make_pre_post_processors, + make_robot_action, +) from lerobot.processor import ( PolicyAction, PolicyProcessorPipeline, @@ -103,8 +103,8 @@ from lerobot.processor import ( RobotObservation, RobotProcessorPipeline, make_default_processors, + rename_stats, ) -from lerobot.processor.rename_processor import rename_stats from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 8f465739c..36d842af4 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -56,9 +56,9 @@ import time from dataclasses import asdict, dataclass from pprint import pformat -from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401 -from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401 -from lerobot.cameras.zmq.configuration_zmq import ZMQCameraConfig # noqa: F401 +from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 +from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 +from lerobot.cameras.zmq import ZMQCameraConfig # noqa: F401 from lerobot.configs import parser from lerobot.processor import ( RobotAction, diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 9572c4e1b..c2332b9aa 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -13,8 +13,6 @@ # 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. -from __future__ import annotations - import dataclasses import logging import time @@ -33,13 +31,10 @@ from tqdm import tqdm from lerobot.configs import parser from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets import EpisodeAwareSampler, make_dataset -from lerobot.envs.factory import make_env, make_env_pre_post_processors -from lerobot.envs.utils import close_envs +from lerobot.envs import close_envs, make_env, make_env_pre_post_processors from lerobot.optim.factory import make_optimizer_and_scheduler -from lerobot.policies.factory import make_policy, make_pre_post_processors -from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.rl.wandb_utils import WandBLogger -from lerobot.scripts.lerobot_eval import eval_policy_all from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.logging_utils import AverageMeter, MetricsTracker from lerobot.utils.random_utils import set_seed @@ -58,6 +53,8 @@ from lerobot.utils.utils import ( inside_slurm, ) +from .lerobot_eval import eval_policy_all + def update_policy( train_metrics: MetricsTracker, @@ -65,7 +62,7 @@ def update_policy( batch: Any, optimizer: Optimizer, grad_clip_norm: float, - accelerator: Accelerator, + accelerator: "Accelerator", lr_scheduler=None, lock=None, rabc_weights_provider=None, @@ -154,7 +151,7 @@ def update_policy( @parser.wrap() -def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None): +def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): """ Main function to train a policy. diff --git a/src/lerobot/scripts/lerobot_train_tokenizer.py b/src/lerobot/scripts/lerobot_train_tokenizer.py index 9afd9fb9d..c821a4d54 100644 --- a/src/lerobot/scripts/lerobot_train_tokenizer.py +++ b/src/lerobot/scripts/lerobot_train_tokenizer.py @@ -60,8 +60,7 @@ if TYPE_CHECKING or _transformers_available: else: AutoProcessor = None -from lerobot.configs import parser -from lerobot.configs.types import NormalizationMode +from lerobot.configs import NormalizationMode, parser from lerobot.datasets import LeRobotDataset from lerobot.utils.constants import ACTION, OBS_STATE diff --git a/src/lerobot/teleoperators/__init__.py b/src/lerobot/teleoperators/__init__.py index ee508dddb..d66e4b67d 100644 --- a/src/lerobot/teleoperators/__init__.py +++ b/src/lerobot/teleoperators/__init__.py @@ -17,3 +17,5 @@ from .config import TeleoperatorConfig from .teleoperator import Teleoperator from .utils import TeleopEvents, make_teleoperator_from_config + +__all__ = ["Teleoperator", "TeleoperatorConfig", "TeleopEvents", "make_teleoperator_from_config"] diff --git a/src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py b/src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py index b44f1fbea..624729c02 100644 --- a/src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py +++ b/src/lerobot/teleoperators/bi_openarm_leader/bi_openarm_leader.py @@ -17,11 +17,10 @@ import logging from functools import cached_property -from lerobot.teleoperators.openarm_leader import OpenArmLeaderConfig from lerobot.types import RobotAction from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected -from ..openarm_leader import OpenArmLeader +from ..openarm_leader import OpenArmLeader, OpenArmLeaderConfig from ..teleoperator import Teleoperator from .config_bi_openarm_leader import BiOpenArmLeaderConfig diff --git a/src/lerobot/teleoperators/bi_openarm_leader/config_bi_openarm_leader.py b/src/lerobot/teleoperators/bi_openarm_leader/config_bi_openarm_leader.py index 39fc90add..f7ec929ed 100644 --- a/src/lerobot/teleoperators/bi_openarm_leader/config_bi_openarm_leader.py +++ b/src/lerobot/teleoperators/bi_openarm_leader/config_bi_openarm_leader.py @@ -16,9 +16,8 @@ from dataclasses import dataclass -from lerobot.teleoperators.openarm_leader import OpenArmLeaderConfigBase - from ..config import TeleoperatorConfig +from ..openarm_leader import OpenArmLeaderConfigBase @TeleoperatorConfig.register_subclass("bi_openarm_leader") diff --git a/src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py b/src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py index e84ac6f50..f2e88d20a 100644 --- a/src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py +++ b/src/lerobot/teleoperators/bi_so_leader/bi_so_leader.py @@ -17,10 +17,9 @@ import logging from functools import cached_property -from lerobot.teleoperators.so_leader import SOLeaderTeleopConfig from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected -from ..so_leader import SOLeader +from ..so_leader import SOLeader, SOLeaderTeleopConfig from ..teleoperator import Teleoperator from .config_bi_so_leader import BiSOLeaderConfig diff --git a/src/lerobot/teleoperators/bi_so_leader/config_bi_so_leader.py b/src/lerobot/teleoperators/bi_so_leader/config_bi_so_leader.py index c2f23c617..f477d0f26 100644 --- a/src/lerobot/teleoperators/bi_so_leader/config_bi_so_leader.py +++ b/src/lerobot/teleoperators/bi_so_leader/config_bi_so_leader.py @@ -16,9 +16,8 @@ from dataclasses import dataclass -from lerobot.teleoperators.so_leader import SOLeaderConfig - from ..config import TeleoperatorConfig +from ..so_leader import SOLeaderConfig @TeleoperatorConfig.register_subclass("bi_so_leader") diff --git a/src/lerobot/teleoperators/homunculus/homunculus_glove.py b/src/lerobot/teleoperators/homunculus/homunculus_glove.py index c4393d660..a58e35025 100644 --- a/src/lerobot/teleoperators/homunculus/homunculus_glove.py +++ b/src/lerobot/teleoperators/homunculus/homunculus_glove.py @@ -23,12 +23,12 @@ import serial from lerobot.motors import MotorCalibration from lerobot.motors.motors_bus import MotorNormMode -from lerobot.teleoperators.homunculus.joints_translation import homunculus_glove_to_hope_jr_hand from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.utils import enter_pressed, move_cursor_up from ..teleoperator import Teleoperator from .config_homunculus import HomunculusGloveConfig +from .joints_translation import homunculus_glove_to_hope_jr_hand logger = logging.getLogger(__name__) diff --git a/src/lerobot/teleoperators/phone/phone_processor.py b/src/lerobot/teleoperators/phone/phone_processor.py index c498bed7d..3d57a5a71 100644 --- a/src/lerobot/teleoperators/phone/phone_processor.py +++ b/src/lerobot/teleoperators/phone/phone_processor.py @@ -16,11 +16,12 @@ from dataclasses import dataclass, field -from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep -from lerobot.teleoperators.phone.config_phone import PhoneOS from lerobot.types import RobotAction +from .config_phone import PhoneOS + @ProcessorStepRegistry.register("map_phone_action_to_robot_action") @dataclass diff --git a/src/lerobot/teleoperators/phone/teleop_phone.py b/src/lerobot/teleoperators/phone/teleop_phone.py index 221ee8083..f68843194 100644 --- a/src/lerobot/teleoperators/phone/teleop_phone.py +++ b/src/lerobot/teleoperators/phone/teleop_phone.py @@ -26,11 +26,12 @@ import hebi import numpy as np from teleop import Teleop -from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS -from lerobot.teleoperators.teleoperator import Teleoperator from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.rotation import Rotation +from ..teleoperator import Teleoperator +from .config_phone import PhoneConfig, PhoneOS + logger = logging.getLogger(__name__) diff --git a/src/lerobot/transforms/__init__.py b/src/lerobot/transforms/__init__.py index f5437fa46..6cf9699d0 100644 --- a/src/lerobot/transforms/__init__.py +++ b/src/lerobot/transforms/__init__.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from lerobot.transforms.transforms import ( +from .transforms import ( ImageTransformConfig, ImageTransforms, ImageTransformsConfig, diff --git a/src/lerobot/transport/__init__.py b/src/lerobot/transport/__init__.py new file mode 100644 index 000000000..7ea343691 --- /dev/null +++ b/src/lerobot/transport/__init__.py @@ -0,0 +1,2 @@ +# gRPC transport layer for async inference. +# Requires: lerobot[grpcio-dep] diff --git a/src/lerobot/utils/__init__.py b/src/lerobot/utils/__init__.py new file mode 100644 index 000000000..4da9306e1 --- /dev/null +++ b/src/lerobot/utils/__init__.py @@ -0,0 +1,66 @@ +# 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 lightweight, base-dependency-only utilities. + +Heavy utility modules (train_utils, control_utils, visualization_utils) +are intentionally NOT re-exported here to avoid pulling in optional +dependencies. Import them directly, e.g.: +``from lerobot.utils.train_utils import save_checkpoint`` +""" + +from .constants import ( + ACTION, + DEFAULT_FEATURES, + DONE, + IMAGENET_STATS, + OBS_ENV_STATE, + OBS_IMAGE, + OBS_IMAGES, + OBS_STATE, + OBS_STR, + REWARD, +) +from .decorators import check_if_already_connected, check_if_not_connected +from .device_utils import auto_select_torch_device, get_safe_torch_device, is_torch_device_available +from .errors import DeviceAlreadyConnectedError, DeviceNotConnectedError +from .import_utils import is_package_available, require_package + +__all__ = [ + # Constants + "ACTION", + "DEFAULT_FEATURES", + "DONE", + "IMAGENET_STATS", + "OBS_ENV_STATE", + "OBS_IMAGE", + "OBS_IMAGES", + "OBS_STATE", + "OBS_STR", + "REWARD", + # Device utilities + "auto_select_torch_device", + "get_safe_torch_device", + "is_torch_device_available", + # Import guards + "is_package_available", + "require_package", + # Decorators + "check_if_already_connected", + "check_if_not_connected", + # Errors + "DeviceAlreadyConnectedError", + "DeviceNotConnectedError", +] diff --git a/src/lerobot/utils/control_utils.py b/src/lerobot/utils/control_utils.py index 68aa653f6..44b0aef0a 100644 --- a/src/lerobot/utils/control_utils.py +++ b/src/lerobot/utils/control_utils.py @@ -27,11 +27,10 @@ from typing import TYPE_CHECKING, Any import numpy as np import torch -from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.policies import PreTrainedPolicy, prepare_observation_for_inference if TYPE_CHECKING: - from lerobot.datasets.lerobot_dataset import LeRobotDataset -from lerobot.policies.utils import prepare_observation_for_inference + from lerobot.datasets import LeRobotDataset from lerobot.processor import PolicyProcessorPipeline from lerobot.robots import Robot from lerobot.types import PolicyAction @@ -218,12 +217,12 @@ def sanity_check_dataset_robot_compatibility( Raises: ValueError: If any of the checked metadata fields do not match. """ - from lerobot.utils.import_utils import require_package + from .import_utils import require_package require_package("deepdiff", extra="hardware") from deepdiff import DeepDiff - from lerobot.utils.constants import DEFAULT_FEATURES + from .constants import DEFAULT_FEATURES fields = [ ("robot_type", dataset.meta.robot_type, robot.robot_type), diff --git a/src/lerobot/utils/decorators.py b/src/lerobot/utils/decorators.py index 8fc2f9a07..75171f637 100644 --- a/src/lerobot/utils/decorators.py +++ b/src/lerobot/utils/decorators.py @@ -16,7 +16,7 @@ from functools import wraps -from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError +from .errors import DeviceAlreadyConnectedError, DeviceNotConnectedError def check_if_not_connected(func): diff --git a/src/lerobot/utils/feature_utils.py b/src/lerobot/utils/feature_utils.py index 45ba80800..2a4886234 100644 --- a/src/lerobot/utils/feature_utils.py +++ b/src/lerobot/utils/feature_utils.py @@ -25,8 +25,9 @@ from typing import Any import numpy as np -from lerobot.configs.types import FeatureType, PolicyFeature -from lerobot.utils.constants import ACTION, DEFAULT_FEATURES, OBS_ENV_STATE, OBS_STR +from lerobot.configs import FeatureType, PolicyFeature + +from .constants import ACTION, DEFAULT_FEATURES, OBS_ENV_STATE, OBS_STR def _validate_feature_names(features: dict[str, dict]) -> None: diff --git a/src/lerobot/utils/io_utils.py b/src/lerobot/utils/io_utils.py index 227c0265b..e037b412c 100644 --- a/src/lerobot/utils/io_utils.py +++ b/src/lerobot/utils/io_utils.py @@ -58,7 +58,7 @@ def write_video(video_path: str | Path, stacked_frames: list, fps: int) -> None: stacked_frames: List of HWC uint8 numpy arrays (RGB). fps: Frames per second for the output video. """ - from lerobot.utils.import_utils import require_package + from .import_utils import require_package require_package("av", extra="av-dep") import av diff --git a/src/lerobot/utils/logging_utils.py b/src/lerobot/utils/logging_utils.py index 1497c0585..0ce596f55 100644 --- a/src/lerobot/utils/logging_utils.py +++ b/src/lerobot/utils/logging_utils.py @@ -16,7 +16,7 @@ from collections.abc import Callable from typing import Any -from lerobot.utils.utils import format_big_number +from .utils import format_big_number class AverageMeter: diff --git a/src/lerobot/utils/random_utils.py b/src/lerobot/utils/random_utils.py index 7bd55cd20..e280fc342 100644 --- a/src/lerobot/utils/random_utils.py +++ b/src/lerobot/utils/random_utils.py @@ -23,8 +23,8 @@ import numpy as np import torch from safetensors.torch import load_file, save_file -from lerobot.utils.constants import RNG_STATE -from lerobot.utils.utils import flatten_dict, unflatten_dict +from .constants import RNG_STATE +from .utils import flatten_dict, unflatten_dict def serialize_python_rng_state() -> dict[str, torch.Tensor]: diff --git a/src/lerobot/utils/train_utils.py b/src/lerobot/utils/train_utils.py index 2ace347ee..eb2c364d1 100644 --- a/src/lerobot/utils/train_utils.py +++ b/src/lerobot/utils/train_utils.py @@ -19,19 +19,24 @@ from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler from lerobot.configs.train import TrainPipelineConfig -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.optim import ( + load_optimizer_state, + load_scheduler_state, + save_optimizer_state, + save_scheduler_state, +) +from lerobot.policies import PreTrainedPolicy from lerobot.processor import PolicyProcessorPipeline -from lerobot.utils.constants import ( + +from .constants import ( CHECKPOINTS_DIR, LAST_CHECKPOINT_LINK, PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR, TRAINING_STEP, ) -from lerobot.utils.io_utils import load_json, write_json -from lerobot.utils.random_utils import load_rng_state, save_rng_state +from .io_utils import load_json, write_json +from .random_utils import load_rng_state, save_rng_state def get_step_identifier(step: int, total_steps: int) -> str: diff --git a/src/lerobot/utils/transition.py b/src/lerobot/utils/transition.py index fe3620861..a79b95151 100644 --- a/src/lerobot/utils/transition.py +++ b/src/lerobot/utils/transition.py @@ -18,7 +18,7 @@ from typing import TypedDict import torch -from lerobot.utils.constants import ACTION +from .constants import ACTION class Transition(TypedDict): diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index 2fe61ff9f..d9d5bf6b5 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -18,9 +18,9 @@ import os import numpy as np from lerobot.types import RobotAction, RobotObservation -from lerobot.utils.import_utils import require_package from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR +from .import_utils import require_package def init_rerun( diff --git a/tests/fixtures/dataset_factories.py b/tests/fixtures/dataset_factories.py index 5ecb52145..e068484b0 100644 --- a/tests/fixtures/dataset_factories.py +++ b/tests/fixtures/dataset_factories.py @@ -34,12 +34,12 @@ from lerobot.datasets.utils import ( DEFAULT_CHUNK_SIZE, DEFAULT_DATA_FILE_SIZE_IN_MB, DEFAULT_DATA_PATH, - DEFAULT_FEATURES, DEFAULT_VIDEO_FILE_SIZE_IN_MB, DEFAULT_VIDEO_PATH, - flatten_dict, ) from lerobot.datasets.video_utils import encode_video_frames +from lerobot.utils.constants import DEFAULT_FEATURES +from lerobot.utils.utils import flatten_dict from tests.fixtures.constants import ( DEFAULT_FPS, DUMMY_CAMERA_FEATURES, diff --git a/uv.lock b/uv.lock index 7c7bb3518..7d6e7b43e 100644 --- a/uv.lock +++ b/uv.lock @@ -2308,7 +2308,6 @@ dev = [ { name = "datasets" }, { name = "debugpy" }, { name = "deepdiff" }, - { name = "diffusers" }, { name = "grpcio" }, { name = "grpcio-tools" }, { name = "jsonlines" }, @@ -2321,6 +2320,12 @@ dev = [ { name = "torchcodec", marker = "(platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l' and sys_platform == 'linux') or (platform_machine != 'x86_64' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, { name = "wandb" }, ] +diffusers-dep = [ + { name = "diffusers" }, +] +diffusion = [ + { name = "diffusers" }, +] dynamixel = [ { name = "dynamixel-sdk" }, ] @@ -2336,6 +2341,7 @@ gamepad = [ ] groot = [ { name = "decord", marker = "platform_machine == 'AMD64' or platform_machine == 'x86_64'" }, + { name = "diffusers" }, { name = "dm-tree", version = "0.1.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, { name = "dm-tree", version = "0.1.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, { name = "flash-attn", marker = "sys_platform != 'darwin'" }, @@ -2398,6 +2404,7 @@ metaworld = [ { name = "torchcodec", marker = "(platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l' and sys_platform == 'linux') or (platform_machine != 'x86_64' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, ] multi-task-dit = [ + { name = "diffusers" }, { name = "transformers" }, ] openarms = [ @@ -2473,7 +2480,6 @@ test = [ { name = "av" }, { name = "datasets" }, { name = "deepdiff" }, - { name = "diffusers" }, { name = "jsonlines" }, { name = "mock-serial", marker = "sys_platform != 'win32'" }, { name = "pynput" }, @@ -2489,7 +2495,6 @@ training = [ { name = "accelerate" }, { name = "av" }, { name = "datasets" }, - { name = "diffusers" }, { name = "jsonlines" }, { name = "torchcodec", marker = "(platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l' and sys_platform == 'linux') or (platform_machine != 'x86_64' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, { name = "wandb" }, @@ -2535,7 +2540,7 @@ requires-dist = [ { name = "debugpy", marker = "extra == 'dev'", specifier = ">=1.8.1,<1.9.0" }, { name = "decord", marker = "(platform_machine == 'AMD64' and extra == 'groot') or (platform_machine == 'x86_64' and extra == 'groot')", specifier = ">=0.6.0,<1.0.0" }, { name = "deepdiff", marker = "extra == 'hardware'", specifier = ">=7.0.1,<9.0.0" }, - { name = "diffusers", marker = "extra == 'training'", specifier = ">=0.27.2,<0.36.0" }, + { name = "diffusers", marker = "extra == 'diffusers-dep'", specifier = ">=0.27.2,<0.36.0" }, { name = "dm-tree", marker = "extra == 'groot'", specifier = ">=0.1.8,<1.0.0" }, { name = "draccus", specifier = "==0.10.0" }, { name = "dynamixel-sdk", marker = "extra == 'dynamixel'", specifier = ">=3.7.31,<3.9.0" }, @@ -2574,6 +2579,10 @@ requires-dist = [ { name = "lerobot", extras = ["dataset"], marker = "extra == 'test'" }, { name = "lerobot", extras = ["dataset"], marker = "extra == 'training'" }, { name = "lerobot", extras = ["dev"], marker = "extra == 'all'" }, + { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'diffusion'" }, + { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'groot'" }, + { name = "lerobot", extras = ["diffusers-dep"], marker = "extra == 'multi-task-dit'" }, + { name = "lerobot", extras = ["diffusion"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["dynamixel"], marker = "extra == 'all'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'hopejr'" }, { name = "lerobot", extras = ["feetech"], marker = "extra == 'lekiwi'" }, @@ -2689,7 +2698,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'transformers-dep'", specifier = "==5.3.0" }, { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.25.0" }, ] -provides-extras = ["dataset", "training", "hardware", "viz", "build", "robot", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "can-dep", "peft-dep", "scipy-dep", "qwen-vl-utils-dep", "matplotlib-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "kinematics", "intelrealsense", "phone", "wallx", "pi", "smolvla", "multi-task-dit", "groot", "sarm", "xvla", "hilserl", "async", "peft", "dev", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] +provides-extras = ["dataset", "training", "hardware", "viz", "build", "robot", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "smolvla", "multi-task-dit", "groot", "sarm", "xvla", "hilserl", "async", "peft", "dev", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"] [[package]] name = "librt"