From af9307b1465ccd1f8b569438306edf066f004dcd Mon Sep 17 00:00:00 2001 From: Hong-Yu Chiu Date: Thu, 2 Jul 2026 06:49:57 +0000 Subject: [PATCH] Lazily import lerobot in datasets. Signed-off-by: Hong-Yu Chiu --- .../action/datasets/agibotworld_beta_lerobot_dataset.py | 4 +++- .../vfm/action/datasets/bridge_orig_lerobot_dataset.py | 4 +++- .../data/vfm/action/datasets/droid_lerobot_dataset.py | 4 +++- .../data/vfm/action/datasets/fractal_lerobot_dataset.py | 4 +++- .../data/vfm/action/datasets/libero_lerobot_dataset.py | 4 +++- .../data/vfm/action/datasets/robomind_franka_dataset.py | 7 ++++++- .../data/vfm/action/datasets/robomind_ur_dataset.py | 4 +++- .../data/vfm/action/datasets/umi_lerobot_dataset.py | 4 +++- 8 files changed, 27 insertions(+), 8 deletions(-) diff --git a/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py index ce22773c..81a9a364 100644 --- a/cosmos_framework/data/vfm/action/datasets/agibotworld_beta_lerobot_dataset.py +++ b/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), diff --git a/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py index 25e167b0..e6ea9d2b 100644 --- a/cosmos_framework/data/vfm/action/datasets/bridge_orig_lerobot_dataset.py +++ b/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), diff --git a/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py index f3faa91f..5d31a999 100644 --- a/cosmos_framework/data/vfm/action/datasets/droid_lerobot_dataset.py +++ b/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( diff --git a/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py index 10628db1..b9dd23d7 100644 --- a/cosmos_framework/data/vfm/action/datasets/fractal_lerobot_dataset.py +++ b/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), diff --git a/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py index fb8d5e8c..fe31f9d5 100644 --- a/cosmos_framework/data/vfm/action/datasets/libero_lerobot_dataset.py +++ b/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)) diff --git a/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py b/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py index 53b0a2c6..7bda3053 100644 --- a/cosmos_framework/data/vfm/action/datasets/robomind_franka_dataset.py +++ b/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( diff --git a/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py b/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py index 317525f7..ab7bb433 100644 --- a/cosmos_framework/data/vfm/action/datasets/robomind_ur_dataset.py +++ b/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), diff --git a/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py b/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py index 59395961..61a31497 100644 --- a/cosmos_framework/data/vfm/action/datasets/umi_lerobot_dataset.py +++ b/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),