diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 1848a6ffd..48f915ce8 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -44,12 +44,19 @@ from lerobot.utils.constants import ( POLICY_PREPROCESSOR_DEFAULT_NAME, ) from lerobot.utils.feature_utils import dataset_to_policy_features +from lerobot.utils.import_utils import _peft_available, require_package from .evo1.configuration_evo1 import Evo1Config from .groot.configuration_groot import GrootConfig from .pretrained import PreTrainedPolicy from .utils import validate_visual_features_consistency +if TYPE_CHECKING or _peft_available: + from peft import PeftConfig, PeftModel +else: + PeftConfig = None + PeftModel = None + def _reconnect_relative_absolute_steps( preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline @@ -334,7 +341,7 @@ def make_policy( # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo # of the adapter and the adapter's config contains the path to the base policy. So we need the # adapter config first, then load the correct policy and then apply PEFT. - from peft import PeftConfig, PeftModel + require_package("peft", extra="peft") logging.info("Loading policy's PEFT adapter.") diff --git a/src/lerobot/policies/molmoact2/modeling_molmoact2.py b/src/lerobot/policies/molmoact2/modeling_molmoact2.py index 2cc85ab02..14d9e19e2 100644 --- a/src/lerobot/policies/molmoact2/modeling_molmoact2.py +++ b/src/lerobot/policies/molmoact2/modeling_molmoact2.py @@ -43,11 +43,22 @@ from torch.distributions import Beta from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.utils.constants import ACTION -from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package +from lerobot.utils.import_utils import ( + _peft_available, + _scipy_available, + _transformers_available, + require_package, +) from ..rtc.modeling_rtc import RTCProcessor from .configuration_molmoact2 import MolmoAct2Config +if TYPE_CHECKING or _peft_available: + from peft import LoraConfig, get_peft_model +else: + LoraConfig = None + get_peft_model = None + logger = logging.getLogger(__name__) @@ -1731,13 +1742,11 @@ class MolmoAct2Policy(PreTrainedPolicy): def _build_inner_lora_config(self): require_package("peft", extra="molmoact2") - from peft import LoraConfig return LoraConfig(**self._get_inner_peft_targets()) def _apply_lora_adapters(self) -> None: require_package("peft", extra="molmoact2") - from peft import get_peft_model peft_config = self._build_inner_lora_config() self._validate_peft_config(peft_config) diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 90d894e30..0283f389a 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -34,14 +34,22 @@ from lerobot.configs import PreTrainedConfig from lerobot.configs.train import TrainPipelineConfig from lerobot.utils.device_utils import resolve_safetensors_device from lerobot.utils.hub import HubMixin +from lerobot.utils.import_utils import _peft_available, require_package from .utils import log_model_loading_keys -T = TypeVar("T", bound="PreTrainedPolicy") +if TYPE_CHECKING or _peft_available: + from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model +else: + PEFT_TYPE_TO_CONFIG_MAPPING = None + PeftType = None + get_peft_model = None if TYPE_CHECKING: from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata +T = TypeVar("T", bound="PreTrainedPolicy") + def _build_card_context( cfg: TrainPipelineConfig | None, @@ -384,7 +392,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.) These are merged with policy defaults to build the final config. """ - from peft import get_peft_model + require_package("peft", extra="peft") # If user provided a complete config, use it directly (with overrides) if peft_config is not None: @@ -455,7 +463,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): Returns: Preprocessed dict with renamed keys and init_type mapped to method-specific key. """ - from peft import PeftType + require_package("peft", extra="peft") cli_overrides = cli_overrides.copy() @@ -480,7 +488,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): def _build_peft_config(self, cli_overrides: dict): """Build a PEFT config from policy defaults and CLI overrides.""" - from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType + require_package("peft", extra="peft") # Determine PEFT method type (default to LORA) method_type_str = cli_overrides.get("method_type") or "lora" @@ -507,7 +515,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict): """Apply CLI overrides to an existing PEFT config.""" - from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType + require_package("peft", extra="peft") # Get method type from existing config or CLI override method_type_str = cli_overrides.get("method_type") diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index 2b39ab871..cf69d57b5 100644 --- a/src/lerobot/rollout/context.py +++ b/src/lerobot/rollout/context.py @@ -24,6 +24,7 @@ from __future__ import annotations import logging from dataclasses import dataclass, field from threading import Event +from typing import TYPE_CHECKING import torch @@ -47,6 +48,7 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor from lerobot.robots import make_robot_from_config from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features +from lerobot.utils.import_utils import _peft_available, require_package from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig from .inference import ( @@ -57,6 +59,12 @@ from .inference import ( ) from .robot_wrapper import ThreadSafeRobot +if TYPE_CHECKING or _peft_available: + from peft import PeftConfig, PeftModel +else: + PeftConfig = None + PeftModel = None + logger = logging.getLogger(__name__) @@ -171,7 +179,7 @@ def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy revision=pretrained_revision, ) - from peft import PeftConfig, PeftModel + require_package("peft", extra="peft") peft_path = policy_config.pretrained_path peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision) diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index c4ed35145..7ee8063a1 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -62,7 +62,7 @@ from dataclasses import asdict from functools import partial from pathlib import Path from pprint import pformat -from typing import Any, TypedDict +from typing import TYPE_CHECKING, Any, TypedDict import einops import gymnasium as gym @@ -87,7 +87,7 @@ from lerobot.processor import PolicyProcessorPipeline from lerobot.types import PolicyAction from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD from lerobot.utils.device_utils import get_safe_torch_device -from lerobot.utils.import_utils import register_third_party_plugins +from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package from lerobot.utils.io_utils import write_video from lerobot.utils.random_utils import set_seed from lerobot.utils.utils import ( @@ -95,6 +95,11 @@ from lerobot.utils.utils import ( inside_slurm, ) +if TYPE_CHECKING or _peft_available: + from peft import PeftModel +else: + PeftModel = None + def _env_features_to_dataset_features(env_features: dict) -> dict: """Convert EnvConfig.features to the dict format expected by LeRobotDataset.create().""" @@ -444,13 +449,11 @@ def eval_policy( exc = ValueError( f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided." ) - try: - from peft import PeftModel - - if not isinstance(policy, PeftModel): - raise exc - except ImportError: - raise exc from None + if not _peft_available: + raise exc + require_package("peft", extra="peft") + if not isinstance(policy, PeftModel): + raise exc start = time.time() # Preserve the mode for direct callers. eval_policy_all scopes the mode diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 8bfa16a98..8d0e8f593 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -57,7 +57,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.rewards import make_reward_pre_post_processors from lerobot.utils.collate import lerobot_collate_fn -from lerobot.utils.import_utils import register_third_party_plugins +from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package from lerobot.utils.logging_utils import AverageMeter, MetricsTracker from lerobot.utils.random_utils import set_seed from lerobot.utils.utils import ( @@ -68,6 +68,11 @@ from lerobot.utils.utils import ( inside_slurm, ) +if TYPE_CHECKING or _peft_available: + from peft import PeftModel +else: + PeftModel = None + from .lerobot_eval import eval_policy_all @@ -207,8 +212,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if cfg.job.is_remote: return submit_to_hf(cfg) - from lerobot.utils.import_utils import require_package - require_package("accelerate", extra="training") from accelerate import Accelerator from accelerate.utils import DistributedDataParallelKwargs, DistributedType @@ -312,7 +315,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if cfg.peft is not None: if cfg.is_reward_model_training: raise ValueError("PEFT is only supported for policy training. ") - from peft import PeftModel + require_package("peft", extra="peft") if isinstance(policy, PeftModel): logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.") diff --git a/tests/test_rollout.py b/tests/test_rollout.py index c247ff3c1..e50949654 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -185,18 +185,25 @@ def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(m peft_config_from_pretrained = MagicMock(return_value=peft_config) adapted_policy = MagicMock() peft_model_from_pretrained = MagicMock(return_value=adapted_policy) - monkeypatch.setitem( - sys.modules, - "peft", - SimpleNamespace( - PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained), - PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained), - ), + require_package = MagicMock() + monkeypatch.setattr(rollout_context, "require_package", require_package) + monkeypatch.setattr( + rollout_context, + "PeftConfig", + SimpleNamespace(from_pretrained=peft_config_from_pretrained), + raising=False, + ) + monkeypatch.setattr( + rollout_context, + "PeftModel", + SimpleNamespace(from_pretrained=peft_model_from_pretrained), + raising=False, ) policy = rollout_context._load_pretrained_policy(policy_config) assert policy is adapted_policy + require_package.assert_called_once_with("peft", extra="peft") peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha") policy_class.from_pretrained.assert_called_once_with( pretrained_name_or_path="user/base-policy",