lazy import molmoact2 scipy

src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py

@@ -22,8 +22,6 @@ from pathlib import Path
 from typing import ClassVar
 
 import numpy as np
-from scipy.fft import dct
-from scipy.fft import idct
 from tokenizers import ByteLevelBPETokenizer
 from tokenizers.trainers import BpeTrainer
 from huggingface_hub import snapshot_download
@@ -86,6 +84,8 @@ class UniversalActionProcessor(ProcessorMixin):
         self.bpe_tokenizer = self.tokenizer
 
     def __call__(self, action_chunk: np.array) -> np.array:
+        from scipy.fft import dct
+
         assert action_chunk.ndim <= 3, "Only 3 dimensions supported: [batch, timesteps, action_dim]"
         if action_chunk.ndim == 2:
             action_chunk = action_chunk[None, ...]
@@ -109,6 +109,8 @@ class UniversalActionProcessor(ProcessorMixin):
         time_horizon: int | None = None,
         action_dim: int | None = None,
     ) -> np.array:
+        from scipy.fft import idct
+
         self.time_horizon = time_horizon or self.time_horizon or self.called_time_horizon
         self.action_dim = action_dim or self.action_dim or self.called_action_dim
 
@@ -150,6 +152,8 @@ class UniversalActionProcessor(ProcessorMixin):
         time_horizon: int | None = None,
         action_dim: int | None = None,
     ) -> "UniversalActionProcessor":
+        from scipy.fft import dct
+
         # Run DCT over all inputs
         dct_tokens = [dct(a, axis=0, norm="ortho").flatten() for a in action_data]
 

src/lerobot/policies/molmoact2/modeling_molmoact2.py

@@ -29,24 +29,30 @@ import torch.nn.functional as F  # noqa: N812
 from safetensors.torch import load_file as load_safetensors_file
 from torch import Tensor
 from torch.distributions import Beta
-from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME
 
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.utils.constants import ACTION
-from lerobot.utils.import_utils import _transformers_available, require_package
+from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
 
 from ..rtc.modeling_rtc import RTCProcessor
 from .configuration_molmoact2 import MolmoAct2Config, _hf_token, _resolve_checkpoint_location
 
 if TYPE_CHECKING or _transformers_available:
-    from .hf_model.action_tokenizer import UniversalActionProcessor
+    from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME
+
     from .hf_model.configuration_molmoact2 import MolmoAct2Config as HFMolmoAct2Config
     from .hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration
 else:
-    UniversalActionProcessor = None
+    SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json"
+    SAFE_WEIGHTS_NAME = "model.safetensors"
     HFMolmoAct2Config = None
     MolmoAct2ForConditionalGeneration = None
 
+if TYPE_CHECKING or (_transformers_available and _scipy_available):
+    from .hf_model.action_tokenizer import UniversalActionProcessor
+else:
+    UniversalActionProcessor = None
+
 _MODEL_INPUT_KEYS = {
     "input_ids",
     "pixel_values",
@@ -483,9 +489,10 @@ class MolmoAct2Policy(PreTrainedPolicy):
     def _load_discrete_action_tokenizer(self) -> Any:
         if self.action_tokenizer is None:
             require_package("transformers", extra="molmoact2")
+            require_package("scipy", extra="molmoact2")
 
             if UniversalActionProcessor is None:
-                raise RuntimeError("transformers is required to load MolmoAct2 action tokenizer.")
+                raise RuntimeError("transformers and scipy are required to load MolmoAct2 action tokenizer.")
             self.action_tokenizer = UniversalActionProcessor.from_pretrained_local(
                 self.config.discrete_action_tokenizer,
             )

src/lerobot/policies/molmoact2/processor_molmoact2.py

@@ -52,24 +52,27 @@ from lerobot.utils.constants import (
     POLICY_POSTPROCESSOR_DEFAULT_NAME,
     POLICY_PREPROCESSOR_DEFAULT_NAME,
 )
-from lerobot.utils.import_utils import _transformers_available, require_package
+from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
 
 from .configuration_molmoact2 import MolmoAct2Config, infer_molmoact2_max_sequence_length
 
 if TYPE_CHECKING or _transformers_available:
     from transformers import Qwen2Tokenizer
 
-    from .hf_model.action_tokenizer import UniversalActionProcessor
     from .hf_model.image_processing_molmoact2 import MolmoAct2ImageProcessor
     from .hf_model.processing_molmoact2 import MolmoAct2Processor
     from .hf_model.video_processing_molmoact2 import MolmoAct2VideoProcessor
 else:
     Qwen2Tokenizer = None
-    UniversalActionProcessor = None
     MolmoAct2ImageProcessor = None
     MolmoAct2Processor = None
     MolmoAct2VideoProcessor = None
 
+if TYPE_CHECKING or (_transformers_available and _scipy_available):
+    from .hf_model.action_tokenizer import UniversalActionProcessor
+else:
+    UniversalActionProcessor = None
+
 ACTION_OUTPUT_TOKEN = "<action_output>"  # nosec B105
 ACTION_START_TOKEN = "<action_start>"  # nosec B105
 ACTION_END_TOKEN = "<action_end>"  # nosec B105
@@ -586,8 +589,9 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep):
         self.processor = _load_local_molmoact2_processor(checkpoint_location)
         self.action_processor = None
         if self.action_mode in {"discrete", "both"}:
+            require_package("scipy", extra="molmoact2")
             if UniversalActionProcessor is None:
-                raise RuntimeError("transformers is required to load MolmoAct2 action tokenizer.")
+                raise RuntimeError("transformers and scipy are required to load MolmoAct2 action tokenizer.")
             self.action_processor = UniversalActionProcessor.from_pretrained_local(
                 self.discrete_action_tokenizer,
             )