diff --git a/src/lerobot/async_inference/__init__.py b/src/lerobot/async_inference/__init__.py index 0da31b6a8..f1edccb87 100644 --- a/src/lerobot/async_inference/__init__.py +++ b/src/lerobot/async_inference/__init__.py @@ -1,6 +1,16 @@ -# Async inference server/client. -# Requires: lerobot[async] +""" +Async inference server/client. + +Requires: ``pip install 'lerobot[async]'`` + +Available modules (import directly):: + + from lerobot.async_inference.policy_server import ... + from lerobot.async_inference.robot_client import ... +""" from lerobot.utils.import_utils import require_package require_package("grpcio", extra="async", import_name="grpc") + +__all__: list[str] = [] diff --git a/src/lerobot/common/__init__.py b/src/lerobot/common/__init__.py index 2f9bf10d9..dfd97496d 100644 --- a/src/lerobot/common/__init__.py +++ b/src/lerobot/common/__init__.py @@ -5,4 +5,12 @@ Unlike ``lerobot.utils`` (which must remain dependency-free), modules here are allowed to import from ``lerobot.policies``, ``lerobot.processor``, ``lerobot.configs``, etc. They are deliberately NOT re-exported from the top-level ``lerobot`` package. + +Available modules (import directly):: + + from lerobot.common.control_utils import predict_action, ... + from lerobot.common.train_utils import save_checkpoint, ... + from lerobot.common.wandb_utils import WandBLogger, ... """ + +__all__: list[str] = [] diff --git a/src/lerobot/rl/wandb_utils.py b/src/lerobot/common/wandb_utils.py similarity index 100% rename from src/lerobot/rl/wandb_utils.py rename to src/lerobot/common/wandb_utils.py diff --git a/src/lerobot/data_processing/__init__.py b/src/lerobot/data_processing/__init__.py index 2f76d5676..cd55d46fc 100644 --- a/src/lerobot/data_processing/__init__.py +++ b/src/lerobot/data_processing/__init__.py @@ -11,3 +11,13 @@ # 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. + +""" +Data processing utilities (annotation tools, dataset transformations). + +Available sub-modules (import directly):: + + from lerobot.data_processing.sarm_annotations import ... +""" + +__all__: list[str] = [] diff --git a/src/lerobot/data_processing/sarm_annotations/__init__.py b/src/lerobot/data_processing/sarm_annotations/__init__.py index 2f76d5676..cd4c38f33 100644 --- a/src/lerobot/data_processing/sarm_annotations/__init__.py +++ b/src/lerobot/data_processing/sarm_annotations/__init__.py @@ -11,3 +11,13 @@ # 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. + +""" +SARM subtask annotation tools. + +Available modules (import directly):: + + from lerobot.data_processing.sarm_annotations.subtask_annotation import ... +""" + +__all__: list[str] = [] diff --git a/src/lerobot/envs/__init__.py b/src/lerobot/envs/__init__.py index adb40490c..730207d6c 100644 --- a/src/lerobot/envs/__init__.py +++ b/src/lerobot/envs/__init__.py @@ -12,6 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +# NOTE: gymnasium is currently a core dependency but is a candidate for moving to an +# optional extra in the future. This guard is here to ensure a clear error message +# if/when that transition happens. +from lerobot.utils.import_utils import require_package + +require_package("gymnasium", extra="evaluation", import_name="gymnasium") + 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 diff --git a/src/lerobot/policies/__init__.py b/src/lerobot/policies/__init__.py index 896cc8617..e138a84d9 100644 --- a/src/lerobot/policies/__init__.py +++ b/src/lerobot/policies/__init__.py @@ -22,11 +22,10 @@ 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 .sac.configuration_sac import SACConfig as SACConfig +from .sac.reward_model.configuration_classifier import RewardClassifierConfig as RewardClassifierConfig from .sarm.configuration_sarm import SARMConfig as SARMConfig from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig -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 @@ -46,6 +45,8 @@ __all__ = [ "PI0Config", "PI0FastConfig", "PI05Config", + "RewardClassifierConfig", + "SACConfig", "SARMConfig", "SmolVLAConfig", "TDMPCConfig", diff --git a/src/lerobot/policies/sarm/__init__.py b/src/lerobot/policies/sarm/__init__.py index 3d52651d9..85b24c59e 100644 --- a/src/lerobot/policies/sarm/__init__.py +++ b/src/lerobot/policies/sarm/__init__.py @@ -1,3 +1,4 @@ from .configuration_sarm import SARMConfig +from .modeling_sarm import SARMRewardModel -__all__ = ["SARMConfig"] +__all__ = ["SARMConfig", "SARMRewardModel"] diff --git a/src/lerobot/policies/smolvla/__init__.py b/src/lerobot/policies/smolvla/__init__.py index d610dcf46..516faaaf9 100644 --- a/src/lerobot/policies/smolvla/__init__.py +++ b/src/lerobot/policies/smolvla/__init__.py @@ -1,7 +1,5 @@ from .configuration_smolvla import SmolVLAConfig -from .processor_smolvla import ( - SmolVLANewLineProcessor, - make_smolvla_pre_post_processors, -) +from .modeling_smolvla import SmolVLAPolicy +from .processor_smolvla import make_smolvla_pre_post_processors -__all__ = ["SmolVLAConfig", "SmolVLANewLineProcessor", "make_smolvla_pre_post_processors"] +__all__ = ["SmolVLAConfig", "SmolVLAPolicy", "make_smolvla_pre_post_processors"] diff --git a/src/lerobot/policies/smolvla/processor_smolvla.py b/src/lerobot/policies/smolvla/processor_smolvla.py index 0aa2beb19..8d6c8aca4 100644 --- a/src/lerobot/policies/smolvla/processor_smolvla.py +++ b/src/lerobot/policies/smolvla/processor_smolvla.py @@ -18,15 +18,13 @@ from typing import Any import torch -from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.processor import ( AddBatchDimensionProcessorStep, - ComplementaryDataProcessorStep, DeviceProcessorStep, + NewLineTaskProcessorStep, NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, - ProcessorStepRegistry, RenameObservationsProcessorStep, TokenizerProcessorStep, UnnormalizerProcessorStep, @@ -71,7 +69,7 @@ def make_smolvla_pre_post_processors( input_steps = [ RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one AddBatchDimensionProcessorStep(), - SmolVLANewLineProcessor(), + NewLineTaskProcessorStep(), TokenizerProcessorStep( tokenizer_name=config.vlm_model_name, padding=config.pad_language_to, @@ -103,41 +101,3 @@ def make_smolvla_pre_post_processors( to_output=transition_to_policy_action, ), ) - - -@ProcessorStepRegistry.register(name="smolvla_new_line_processor") -class SmolVLANewLineProcessor(ComplementaryDataProcessorStep): - """ - A processor step that ensures the 'task' description ends with a newline character. - - This step is necessary for certain tokenizers (e.g., PaliGemma) that expect a - newline at the end of the prompt. It handles both single string tasks and lists - of string tasks. - """ - - def complementary_data(self, complementary_data): - if "task" not in complementary_data: - return complementary_data - - task = complementary_data["task"] - if task is None: - return complementary_data - - new_complementary_data = dict(complementary_data) - - # Handle both string and list of strings - if isinstance(task, str): - # Single string: add newline if not present - if not task.endswith("\n"): - new_complementary_data["task"] = f"{task}\n" - elif isinstance(task, list) and all(isinstance(t, str) for t in task): - # List of strings: add newline to each if not present - new_complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task] - # If task is neither string nor list of strings, leave unchanged - - return new_complementary_data - - def transform_features( - self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] - ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: - return features diff --git a/src/lerobot/policies/smolvla/smolvlm_with_expert.py b/src/lerobot/policies/smolvla/smolvlm_with_expert.py index b1d93d72d..ea806f185 100644 --- a/src/lerobot/policies/smolvla/smolvlm_with_expert.py +++ b/src/lerobot/policies/smolvla/smolvlm_with_expert.py @@ -13,21 +13,27 @@ # limitations under the License. import copy +from typing import TYPE_CHECKING import torch from torch import nn -from lerobot.utils.import_utils import require_package +from lerobot.utils.import_utils import _transformers_available, require_package -require_package("transformers", extra="smolvla") - -from transformers import ( - AutoConfig, - AutoModel, - AutoModelForImageTextToText, - AutoProcessor, - SmolVLMForConditionalGeneration, -) +if TYPE_CHECKING or _transformers_available: + from transformers import ( + AutoConfig, + AutoModel, + AutoModelForImageTextToText, + AutoProcessor, + SmolVLMForConditionalGeneration, + ) +else: + AutoConfig = None + AutoModel = None + AutoModelForImageTextToText = None + AutoProcessor = None + SmolVLMForConditionalGeneration = None def apply_rope(x, positions, max_wavelength=10_000): @@ -78,6 +84,7 @@ class SmolVLMWithExpertModel(nn.Module): device: str = "auto", ): super().__init__() + require_package("transformers", extra="smolvla") if load_vlm_weights: print(f"Loading {model_id} weights ...") self.vlm = AutoModelForImageTextToText.from_pretrained( diff --git a/src/lerobot/policies/wall_x/__init__.py b/src/lerobot/policies/wall_x/__init__.py index 7c10b8b2a..16fd2c8ab 100644 --- a/src/lerobot/policies/wall_x/__init__.py +++ b/src/lerobot/policies/wall_x/__init__.py @@ -15,6 +15,7 @@ # limitations under the License. from .configuration_wall_x import WallXConfig +from .modeling_wall_x import WallXPolicy from .processor_wall_x import make_wall_x_pre_post_processors -__all__ = ["WallXConfig", "make_wall_x_pre_post_processors"] +__all__ = ["WallXConfig", "WallXPolicy", "make_wall_x_pre_post_processors"] diff --git a/src/lerobot/policies/wall_x/modeling_wall_x.py b/src/lerobot/policies/wall_x/modeling_wall_x.py index a95b42e9d..b90b7e564 100644 --- a/src/lerobot/policies/wall_x/modeling_wall_x.py +++ b/src/lerobot/policies/wall_x/modeling_wall_x.py @@ -34,40 +34,28 @@ lerobot-train \ ``` """ +import logging import math from collections import deque from os import PathLike -from typing import Any +from typing import TYPE_CHECKING, Any 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 from torch import Tensor from torch.distributions import Beta from torch.nn import CrossEntropyLoss -from torchdiffeq import odeint -from transformers import AutoProcessor, BatchFeature -from transformers.cache_utils import ( - StaticCache, -) -from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( - Qwen2_5_VLForConditionalGeneration, -) -from transformers.utils import is_torchdynamo_compiling, logging from lerobot.utils.constants import ACTION, OBS_STATE +from lerobot.utils.import_utils import ( + _peft_available, + _transformers_available, + is_package_available, + require_package, +) from ..pretrained import PreTrainedPolicy from ..utils import populate_queues @@ -82,12 +70,45 @@ from .constant import ( RESOLUTION, TOKENIZER_MAX_LENGTH, ) -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, + +_torchdiffeq_available = is_package_available("torchdiffeq") +_qwen_vl_utils_available = is_package_available("qwen-vl-utils", import_name="qwen_vl_utils") +_wallx_deps_available = ( + _transformers_available and _peft_available and _torchdiffeq_available and _qwen_vl_utils_available ) + +if TYPE_CHECKING or _wallx_deps_available: + from peft import LoraConfig, get_peft_model + from qwen_vl_utils.vision_process import smart_resize + from torchdiffeq import odeint + from transformers import AutoProcessor, BatchFeature + from transformers.cache_utils import StaticCache + from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VLForConditionalGeneration, + ) + from transformers.utils import is_torchdynamo_compiling + + 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, + ) +else: + LoraConfig = None + get_peft_model = None + smart_resize = None + odeint = None + AutoProcessor = None + BatchFeature = None + StaticCache = None + Qwen2_5_VLForConditionalGeneration = None + is_torchdynamo_compiling = None + Qwen2_5_VLConfig = None + Qwen2_5_VisionTransformerPretrainedModel = None + Qwen2_5_VLACausalLMOutputWithPast = None + Qwen2_5_VLMoEModel = None + from .utils import ( get_wallx_normal_text, preprocesser_call, @@ -95,7 +116,7 @@ from .utils import ( replace_action_token, ) -logger = logging.get_logger(__name__) +logger = logging.getLogger(__name__) class SinusoidalPosEmb(nn.Module): @@ -262,7 +283,13 @@ class ActionHead(nn.Module): return self.propri_proj(proprioception) -class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration): +# Conditional base: when transformers is unavailable the class still parses +# (inheriting from nn.Module) but cannot be instantiated—require_package in +# WallXPolicy.__init__ gives the user a clear error before that happens. +_Qwen2_5_VLForAction_Base = Qwen2_5_VLForConditionalGeneration if _wallx_deps_available else nn.Module + + +class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base): """ Qwen2.5 Vision-Language Mixture of Experts model for action processing. @@ -1717,6 +1744,10 @@ class WallXPolicy(PreTrainedPolicy): name = "wall_x" def __init__(self, config: WallXConfig, **kwargs): + 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") super().__init__(config) config.validate_features() self.config = config diff --git a/src/lerobot/policies/wall_x/utils.py b/src/lerobot/policies/wall_x/utils.py index 39a20afce..d38a2d509 100644 --- a/src/lerobot/policies/wall_x/utils.py +++ b/src/lerobot/policies/wall_x/utils.py @@ -25,10 +25,16 @@ import random import re from collections import OrderedDict from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any import torch -from transformers import BatchFeature + +from lerobot.utils.import_utils import _transformers_available + +if TYPE_CHECKING or _transformers_available: + from transformers import BatchFeature +else: + BatchFeature = None from lerobot.utils.constants import OBS_IMAGES diff --git a/src/lerobot/policies/xvla/__init__.py b/src/lerobot/policies/xvla/__init__.py index 7a8cce8bb..58609e91c 100644 --- a/src/lerobot/policies/xvla/__init__.py +++ b/src/lerobot/policies/xvla/__init__.py @@ -1,4 +1,5 @@ from .configuration_xvla import XVLAConfig +from .modeling_xvla import XVLAPolicy from .processor_xvla import ( XVLAAddDomainIdProcessorStep, XVLAImageNetNormalizeProcessorStep, @@ -7,6 +8,7 @@ from .processor_xvla import ( __all__ = [ "XVLAConfig", + "XVLAPolicy", "XVLAAddDomainIdProcessorStep", "XVLAImageNetNormalizeProcessorStep", "XVLAImageToFloatProcessorStep", diff --git a/src/lerobot/policies/xvla/modeling_xvla.py b/src/lerobot/policies/xvla/modeling_xvla.py index 0f8bdd036..04e923fdd 100644 --- a/src/lerobot/policies/xvla/modeling_xvla.py +++ b/src/lerobot/policies/xvla/modeling_xvla.py @@ -23,10 +23,7 @@ import logging import os from collections import deque from pathlib import Path - -from lerobot.utils.import_utils import require_package - -require_package("transformers", extra="xvla") +from typing import TYPE_CHECKING import torch import torch.nn.functional as F # noqa: N812 @@ -34,15 +31,22 @@ from torch import Tensor, nn from lerobot.configs import PreTrainedConfig from lerobot.utils.constants import ACTION, OBS_LANGUAGE_TOKENS, OBS_STATE +from lerobot.utils.import_utils import _transformers_available, require_package 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 -from .modeling_florence2 import Florence2ForConditionalGeneration from .soft_transformer import SoftPromptedTransformer +# Florence2 config and modeling depend on transformers +if TYPE_CHECKING or _transformers_available: + from .configuration_florence2 import Florence2Config + from .modeling_florence2 import Florence2ForConditionalGeneration +else: + Florence2Config = None + Florence2ForConditionalGeneration = None + class XVLAModel(nn.Module): """ @@ -278,6 +282,7 @@ class XVLAPolicy(PreTrainedPolicy): name = "xvla" def __init__(self, config: XVLAConfig, **kwargs): + require_package("transformers", extra="xvla") super().__init__(config) config.validate_features() florence_config = config.get_florence_config() diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 213bd93f2..20cdf310c 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -56,6 +56,7 @@ from .hil_processor import ( RewardClassifierProcessorStep, TimeLimitProcessorStep, ) +from .newline_task_processor import NewLineTaskProcessorStep from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep, hotswap_stats from .observation_processor import VanillaObservationProcessorStep from .pipeline import ( @@ -119,6 +120,7 @@ __all__ = [ "RelativeActionsProcessorStep", "MapDeltaActionToRobotActionStep", "MapTensorToDeltaActionDictStep", + "NewLineTaskProcessorStep", "NormalizerProcessorStep", "Numpy2TorchActionProcessorStep", "ObservationProcessorStep", diff --git a/src/lerobot/processor/newline_task_processor.py b/src/lerobot/processor/newline_task_processor.py new file mode 100644 index 000000000..ea61bdd71 --- /dev/null +++ b/src/lerobot/processor/newline_task_processor.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python + +# Copyright 2025 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. + +from lerobot.configs import PipelineFeatureType, PolicyFeature + +from .pipeline import ComplementaryDataProcessorStep, ProcessorStepRegistry + + +# NOTE: The registry name "smolvla_new_line_processor" is kept for backward compatibility +# with serialized processor configs that reference this name. +@ProcessorStepRegistry.register(name="smolvla_new_line_processor") +class NewLineTaskProcessorStep(ComplementaryDataProcessorStep): + """ + A processor step that ensures the 'task' description ends with a newline character. + + This step is necessary for certain tokenizers (e.g., PaliGemma) that expect a + newline at the end of the prompt. It handles both single string tasks and lists + of string tasks. + """ + + def complementary_data(self, complementary_data): + if "task" not in complementary_data: + return complementary_data + + task = complementary_data["task"] + if task is None: + return complementary_data + + new_complementary_data = dict(complementary_data) + + # Handle both string and list of strings + if isinstance(task, str): + # Single string: add newline if not present + if not task.endswith("\n"): + new_complementary_data["task"] = f"{task}\n" + elif isinstance(task, list) and all(isinstance(t, str) for t in task): + # List of strings: add newline to each if not present + new_complementary_data["task"] = [t if t.endswith("\n") else f"{t}\n" for t in task] + # If task is neither string nor list of strings, leave unchanged + + return new_complementary_data + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features diff --git a/src/lerobot/processor/policy_robot_bridge.py b/src/lerobot/processor/policy_robot_bridge.py index 1c5327bb5..25d622dc2 100644 --- a/src/lerobot/processor/policy_robot_bridge.py +++ b/src/lerobot/processor/policy_robot_bridge.py @@ -20,9 +20,11 @@ from typing import Any import torch from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature -from lerobot.processor import ActionProcessorStep, PolicyAction, ProcessorStepRegistry, RobotAction +from lerobot.types import PolicyAction, RobotAction from lerobot.utils.constants import ACTION +from .pipeline import ActionProcessorStep, ProcessorStepRegistry + @dataclass @ProcessorStepRegistry.register("robot_action_to_policy_action_processor") diff --git a/src/lerobot/rl/__init__.py b/src/lerobot/rl/__init__.py index 41037b98d..f6c5df0db 100644 --- a/src/lerobot/rl/__init__.py +++ b/src/lerobot/rl/__init__.py @@ -1,6 +1,20 @@ -# Reinforcement learning modules. -# Requires: lerobot[hilserl] +""" +Reinforcement learning modules. + +Requires: ``pip install 'lerobot[hilserl]'`` + +Available modules (import directly):: + + from lerobot.rl.actor import ... + from lerobot.rl.learner import ... + from lerobot.rl.learner_service import ... + from lerobot.rl.buffer import ... + from lerobot.rl.eval_policy import ... + from lerobot.rl.gym_manipulator import ... +""" from lerobot.utils.import_utils import require_package require_package("grpcio", extra="hilserl", import_name="grpc") + +__all__: list[str] = [] diff --git a/src/lerobot/rl/learner.py b/src/lerobot/rl/learner.py index 8f54f8c1b..073d9a65f 100644 --- a/src/lerobot/rl/learner.py +++ b/src/lerobot/rl/learner.py @@ -66,6 +66,7 @@ from lerobot.common.train_utils import ( save_checkpoint, update_last_checkpoint, ) +from lerobot.common.wandb_utils import WandBLogger from lerobot.configs import parser from lerobot.configs.train import TrainRLServerPipelineConfig from lerobot.datasets import LeRobotDataset, make_dataset @@ -99,7 +100,6 @@ from lerobot.utils.utils import ( 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 .wandb_utils import WandBLogger + from lerobot.common.wandb_utils import WandBLogger wandb_logger = WandBLogger(cfg) else: diff --git a/src/lerobot/scripts/lerobot_find_cameras.py b/src/lerobot/scripts/lerobot_find_cameras.py index ab999f7ae..72f4096da 100644 --- a/src/lerobot/scripts/lerobot_find_cameras.py +++ b/src/lerobot/scripts/lerobot_find_cameras.py @@ -38,10 +38,8 @@ import numpy as np from PIL import Image from lerobot.cameras import ColorMode -from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.cameras.opencv.camera_opencv import OpenCVCamera -from lerobot.cameras.realsense import RealSenseCameraConfig -from lerobot.cameras.realsense.camera_realsense import RealSenseCamera +from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig +from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig logger = logging.getLogger(__name__) diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index f6787e149..8f1d27602 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -142,7 +142,7 @@ from lerobot.teleoperators import ( # noqa: F401 so_leader, unitree_g1, ) -from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop +from lerobot.teleoperators.keyboard import KeyboardTeleop from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 0e27b45bd..c943767e3 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -35,13 +35,13 @@ from lerobot.common.train_utils import ( save_checkpoint, update_last_checkpoint, ) +from lerobot.common.wandb_utils import WandBLogger from lerobot.configs import parser from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets import EpisodeAwareSampler, make_dataset 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 import PreTrainedPolicy, make_policy, make_pre_post_processors -from lerobot.rl.wandb_utils import WandBLogger 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 diff --git a/src/lerobot/transport/__init__.py b/src/lerobot/transport/__init__.py index 7ea343691..355ad4cb9 100644 --- a/src/lerobot/transport/__init__.py +++ b/src/lerobot/transport/__init__.py @@ -1,2 +1,11 @@ -# gRPC transport layer for async inference. -# Requires: lerobot[grpcio-dep] +""" +gRPC transport layer for async inference. + +Requires: ``pip install 'lerobot[grpcio-dep]'`` + +Available modules (import directly):: + + from lerobot.transport.utils import ... +""" + +__all__: list[str] = [] diff --git a/tests/processor/test_smolvla_processor.py b/tests/processor/test_smolvla_processor.py index 227b1dc35..2aa7d4bdf 100644 --- a/tests/processor/test_smolvla_processor.py +++ b/tests/processor/test_smolvla_processor.py @@ -22,14 +22,12 @@ import torch from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig -from lerobot.policies.smolvla.processor_smolvla import ( - SmolVLANewLineProcessor, - make_smolvla_pre_post_processors, -) +from lerobot.policies.smolvla.processor_smolvla import make_smolvla_pre_post_processors from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, EnvTransition, + NewLineTaskProcessorStep, NormalizerProcessorStep, ProcessorStep, RenameObservationsProcessorStep, @@ -108,7 +106,7 @@ def test_make_smolvla_processor_basic(): assert len(preprocessor.steps) == 6 assert isinstance(preprocessor.steps[0], RenameObservationsProcessorStep) assert isinstance(preprocessor.steps[1], AddBatchDimensionProcessorStep) - assert isinstance(preprocessor.steps[2], SmolVLANewLineProcessor) + assert isinstance(preprocessor.steps[2], NewLineTaskProcessorStep) # Step 3 would be TokenizerProcessorStep but it's mocked assert isinstance(preprocessor.steps[4], DeviceProcessorStep) assert isinstance(preprocessor.steps[5], NormalizerProcessorStep) @@ -120,8 +118,8 @@ def test_make_smolvla_processor_basic(): def test_smolvla_newline_processor_single_task(): - """Test SmolVLANewLineProcessor with single task string.""" - processor = SmolVLANewLineProcessor() + """Test NewLineTaskProcessorStep with single task string.""" + processor = NewLineTaskProcessorStep() # Test with task that doesn't have newline transition = create_transition(complementary_data={"task": "test task"}) @@ -135,8 +133,8 @@ def test_smolvla_newline_processor_single_task(): def test_smolvla_newline_processor_list_of_tasks(): - """Test SmolVLANewLineProcessor with list of task strings.""" - processor = SmolVLANewLineProcessor() + """Test NewLineTaskProcessorStep with list of task strings.""" + processor = NewLineTaskProcessorStep() # Test with list of tasks tasks = ["task1", "task2\n", "task3"] @@ -147,8 +145,8 @@ def test_smolvla_newline_processor_list_of_tasks(): def test_smolvla_newline_processor_empty_transition(): - """Test SmolVLANewLineProcessor with empty transition.""" - processor = SmolVLANewLineProcessor() + """Test NewLineTaskProcessorStep with empty transition.""" + processor = NewLineTaskProcessorStep() # Test with no complementary_data transition = create_transition() @@ -361,8 +359,8 @@ def test_smolvla_processor_without_stats(): def test_smolvla_newline_processor_state_dict(): - """Test SmolVLANewLineProcessor state dict methods.""" - processor = SmolVLANewLineProcessor() + """Test NewLineTaskProcessorStep state dict methods.""" + processor = NewLineTaskProcessorStep() # Test state_dict (should be empty) state = processor.state_dict() @@ -380,8 +378,8 @@ def test_smolvla_newline_processor_state_dict(): def test_smolvla_newline_processor_transform_features(): - """Test SmolVLANewLineProcessor transform_features method.""" - processor = SmolVLANewLineProcessor() + """Test NewLineTaskProcessorStep transform_features method.""" + processor = NewLineTaskProcessorStep() # Test transform_features features = {