Lazily import lerobot in datasets.

cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py

@@ -11,7 +11,6 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
-from lerobot.datasets.video_utils import decode_video_frames
 
 from cosmos_framework.data.vfm.action.agibot_fk import (
     AGIBOT_WORLD_GRIPPER_TO_OPENCV_BY_WRIST,
@@ -232,6 +231,9 @@ def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str,
         return self._compose_multi_view(top, left, right)
 
     def _load_video_key(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]], key: str) -> torch.Tensor:
+        # lerobot is a heavy, optional ("train" extra) dependency; import lazily.
+        from lerobot.datasets.video_utils import decode_video_frames
+
         timestamps = [float(row["timestamp"]) for row in observation_rows]
         return decode_video_frames(
             self._video_path(episode, key),

cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py

@@ -11,7 +11,6 @@
 
 import numpy as np
 import torch
-from lerobot.datasets.video_utils import decode_video_frames
 
 from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec
 from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset
@@ -126,6 +125,9 @@ def __getitem__(self, idx: int) -> dict[str, Any]:
         )
 
     def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor:
+        # lerobot is a heavy, optional ("train" extra) dependency; import lazily.
+        from lerobot.datasets.video_utils import decode_video_frames
+
         timestamps = [float(row["timestamp"]) for row in observation_rows]
         return decode_video_frames(
             self._video_path(episode, _IMAGE_FEATURE),

cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py

@@ -15,7 +15,6 @@
 import torch
 import torch.nn.functional as F
 import torchvision.transforms as T
-from lerobot.datasets.video_utils import decode_video_frames
 
 from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Joint, Pos, Rot, build_action_spec
 from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset
@@ -290,6 +289,9 @@ def _load_concat_video(
         episode: dict[str, Any],
         observation_rows: list[dict[str, Any]],
     ) -> torch.Tensor:
+        # lerobot is a heavy, optional ("train" extra) dependency; import lazily.
+        from lerobot.datasets.video_utils import decode_video_frames
+
         timestamps = [float(row["timestamp"]) for row in observation_rows]
         frames_by_view = {
             name: decode_video_frames(

cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py

@@ -18,7 +18,6 @@
 
 import numpy as np
 import torch
-from lerobot.datasets.video_utils import decode_video_frames
 
 from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec
 from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset
@@ -155,6 +154,9 @@ def __getitem__(self, idx: int) -> dict[str, Any]:
         )
 
     def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor:
+        # lerobot is a heavy, optional ("train" extra) dependency; import lazily.
+        from lerobot.datasets.video_utils import decode_video_frames
+
         timestamps = [float(row["timestamp"]) for row in observation_rows]
         return decode_video_frames(
             self._video_path(episode, _IMAGE_FEATURE),

cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py

@@ -31,7 +31,6 @@
 import pyarrow.parquet as pq
 import torch
 import torch.nn.functional as F
-from lerobot.datasets.video_utils import decode_video_frames
 
 from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec
 from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset
@@ -307,6 +306,9 @@ def _build_frame_wise_action(self, raw: np.ndarray) -> torch.Tensor:
         return torch.cat([translation, rotation, gripper], dim=-1)  # [chunk, action_dim]
 
     def _load_video(self, episode: dict[str, Any], timestamps: list[float]) -> torch.Tensor:
+        # lerobot is a heavy, optional ("train" extra) dependency; import lazily.
+        from lerobot.datasets.video_utils import decode_video_frames
+
         frames_by_view = {}
         for key in self._video_keys:
             from_ts = float(episode.get(f"videos/{key}/from_timestamp", 0.0))

cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py

@@ -12,7 +12,6 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
-from lerobot.datasets.video_utils import decode_video_frames
 
 from cosmos_framework.data.vfm.action.action_normalization import load_action_stats
 from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec
@@ -181,6 +180,9 @@ def _load_single_video(
         episode: dict[str, Any],
         observation_rows: list[dict[str, Any]],
     ) -> torch.Tensor:
+        # lerobot is a heavy, optional ("train" extra) dependency; import lazily.
+        from lerobot.datasets.video_utils import decode_video_frames
+
         video_key = self._image_features["top"]
         timestamps = [float(row["timestamp"]) for row in observation_rows]
         return decode_video_frames(
@@ -194,6 +196,9 @@ def _load_concat_video(
         episode: dict[str, Any],
         observation_rows: list[dict[str, Any]],
     ) -> torch.Tensor:
+        # lerobot is a heavy, optional ("train" extra) dependency; import lazily.
+        from lerobot.datasets.video_utils import decode_video_frames
+
         timestamps = [float(row["timestamp"]) for row in observation_rows]
         frames_by_view = {
             name: decode_video_frames(

cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py

@@ -11,7 +11,6 @@
 
 import numpy as np
 import torch
-from lerobot.datasets.video_utils import decode_video_frames
 
 from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec
 from cosmos_framework.data.vfm.action.datasets.base_dataset import ActionBaseDataset
@@ -168,6 +167,9 @@ def __getitem__(self, idx: int) -> dict[str, Any]:
         )
 
     def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor:
+        # lerobot is a heavy, optional ("train" extra) dependency; import lazily.
+        from lerobot.datasets.video_utils import decode_video_frames
+
         timestamps = [float(row["timestamp"]) for row in observation_rows]
         return decode_video_frames(
             self._video_path(episode, _IMAGE_FEATURE),

cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py

@@ -11,7 +11,6 @@
 
 import numpy as np
 import torch
-from lerobot.datasets.video_utils import decode_video_frames
 
 from cosmos_framework.data.vfm.action.action_normalization import load_action_stats
 from cosmos_framework.data.vfm.action.action_spec import ActionSpec, Gripper, Pos, Rot, build_action_spec
@@ -140,6 +139,9 @@ def __getitem__(self, idx: int) -> dict[str, Any]:
         )
 
     def _load_video(self, episode: dict[str, Any], observation_rows: list[dict[str, Any]]) -> torch.Tensor:
+        # lerobot is a heavy, optional ("train" extra) dependency; import lazily.
+        from lerobot.datasets.video_utils import decode_video_frames
+
         timestamps = [float(row["timestamp"]) for row in observation_rows]
         return decode_video_frames(
             self._video_path(episode, self._image_key),