From 407a8c1d7dd990cb3eace4ecb2e33cb66949f45b Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Wed, 8 Jul 2026 12:16:01 +0200 Subject: [PATCH] feat(annotate): support video datasets in VF advantage scoring --- .../steerable_pipeline/modules/advantage.py | 68 +++++++++++++++---- .../annotations/steerable_pipeline/reader.py | 18 +++-- src/lerobot/scripts/lerobot_annotate.py | 39 ++++++++++- 3 files changed, 106 insertions(+), 19 deletions(-) diff --git a/src/lerobot/annotations/steerable_pipeline/modules/advantage.py b/src/lerobot/annotations/steerable_pipeline/modules/advantage.py index 2b34bea99..b010b1cc0 100644 --- a/src/lerobot/annotations/steerable_pipeline/modules/advantage.py +++ b/src/lerobot/annotations/steerable_pipeline/modules/advantage.py @@ -33,6 +33,7 @@ import numpy as np import torch from ..config import AdvantageConfig +from ..frames import VideoFrameProvider, null_provider from ..reader import EpisodeRecord from ..staging import EpisodeStaging @@ -124,6 +125,9 @@ class AdvantageModule: def _compute_values(self, record: EpisodeRecord, skip_mask: np.ndarray | None = None) -> np.ndarray: """Run frozen VF over all frames to get V(s_t) predictions. + Supports both image datasets (columns in parquet) and video datasets + (frames decoded from .mp4 via the shared VideoFrameProvider). + Args: record: Episode data. skip_mask: Optional boolean mask [num_frames]. Frames where True are @@ -133,17 +137,23 @@ class AdvantageModule: num_frames = len(df) values = np.zeros(num_frames, dtype=np.float32) - image_key = self._resolve_image_key(df) - if image_key is None: - logger.warning("No image key found for episode %d; returning zero values.", record.episode_index) - return values - # Determine which frame indices actually need inference infer_indices = np.where(~skip_mask)[0] if skip_mask is not None else np.arange(num_frames) - if len(infer_indices) == 0: return values + # Try parquet image columns first, fall back to video decoding + image_key = self._resolve_image_key(df) + video_frames = None + + if image_key is None: + image_key, video_frames = self._decode_video_frames(record, infer_indices) + if image_key is None: + logger.warning( + "No image/video key found for episode %d; returning zero values.", record.episode_index + ) + return values + task_text = record.episode_task for batch_start in range(0, len(infer_indices), self.config.batch_size): @@ -151,14 +161,18 @@ class AdvantageModule: batch_indices = infer_indices[batch_start:batch_end] batch_images = [] - for idx in batch_indices: - img_val = df.iloc[idx][image_key] - if isinstance(img_val, np.ndarray): - img_tensor = torch.from_numpy(img_val).float() - elif isinstance(img_val, torch.Tensor): - img_tensor = img_val.float() + for local_i in range(len(batch_indices)): + if video_frames is not None: + img_tensor = video_frames[batch_start + local_i].float() else: - img_tensor = torch.zeros(3, 224, 224) + idx = batch_indices[local_i] + img_val = df.iloc[idx][image_key] + if isinstance(img_val, np.ndarray): + img_tensor = torch.from_numpy(img_val).float() + elif isinstance(img_val, torch.Tensor): + img_tensor = img_val.float() + else: + img_tensor = torch.zeros(3, 224, 224) batch_images.append(img_tensor) batch_images_tensor = torch.stack(batch_images) @@ -178,6 +192,34 @@ class AdvantageModule: return values + def _decode_video_frames( + self, record: EpisodeRecord, infer_indices: np.ndarray + ) -> tuple[str | None, torch.Tensor | None]: + """Decode video frames using the existing VideoFrameProvider infrastructure. + + Returns (image_key, decoded_frames_tensor) or (None, None) on failure. + """ + dataset_root = record.data_path.parent.parent.parent + + if not hasattr(self, "_frame_provider") or self._frame_provider is None: + try: + self._frame_provider = VideoFrameProvider(root=dataset_root) + except Exception: + self._frame_provider = null_provider() + + if not self._frame_provider.camera_keys: + return None, None + + camera_key = self._frame_provider.camera_keys[0] + timestamps = [float(record.frame_timestamps[i]) for i in infer_indices] + + frames = self._frame_provider.frames_at(record, timestamps, camera_key=camera_key) + if not frames: + return None, None + + frames_tensor = torch.stack(frames) + return camera_key, frames_tensor + def _compute_n_step_advantages( self, mc_returns: np.ndarray, values: np.ndarray, record: EpisodeRecord, n: int ) -> np.ndarray: diff --git a/src/lerobot/annotations/steerable_pipeline/reader.py b/src/lerobot/annotations/steerable_pipeline/reader.py index 22fe4ac26..1cf1c679b 100644 --- a/src/lerobot/annotations/steerable_pipeline/reader.py +++ b/src/lerobot/annotations/steerable_pipeline/reader.py @@ -31,6 +31,7 @@ rows into memory at once. from __future__ import annotations +import functools from collections.abc import Iterator, Sequence from dataclasses import dataclass, field from pathlib import Path @@ -42,6 +43,18 @@ from lerobot.datasets.io_utils import load_tasks from lerobot.datasets.utils import DEFAULT_TASKS_PATH +@functools.lru_cache(maxsize=8) +def _read_parquet_as_pandas(path: Path): # type: ignore[no-untyped-def] + """Read a parquet shard once and cache the pandas DataFrame. + + Multiple EpisodeRecords from the same shard share this single read. + The LRU cache (keyed by path) avoids re-reading the same file + across 100+ episodes that all live in one chunk. + """ + + return pq.read_table(path).to_pandas() + + @dataclass class EpisodeRecord: """Per-episode record yielded by the reader.""" @@ -61,10 +74,7 @@ class EpisodeRecord: def frames_df(self): # type: ignore[no-untyped-def] """Lazy-load the pandas slice for this episode (memoized).""" if self._frames_df_cache is None: - import pandas as pd # noqa: PLC0415 - deferred for optional dataset extra - - table = pq.read_table(self.data_path) - df: pd.DataFrame = table.to_pandas() + df = _read_parquet_as_pandas(self.data_path) self._frames_df_cache = df.iloc[self.row_offset : self.row_offset + self.row_count].reset_index( drop=True ) diff --git a/src/lerobot/scripts/lerobot_annotate.py b/src/lerobot/scripts/lerobot_annotate.py index 6bda51978..0911288ca 100644 --- a/src/lerobot/scripts/lerobot_annotate.py +++ b/src/lerobot/scripts/lerobot_annotate.py @@ -53,6 +53,37 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path: if cfg.repo_id is not None: from huggingface_hub import snapshot_download + needs_vlm = cfg.plan.enabled or cfg.interjections.enabled or cfg.vqa.enabled + advantage_only = cfg.advantage.enabled and not needs_vlm + + if advantage_only: + # Download only metadata + parquet first to resolve the camera key, + # then fetch only the single camera's videos the advantage module needs. + root = Path( + snapshot_download( + repo_id=cfg.repo_id, + repo_type="dataset", + allow_patterns=["meta/**", "data/**"], + ) + ) + camera_key = cfg.vlm.camera_key + if camera_key is None: + from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata # noqa: PLC0415 + + meta = LeRobotDatasetMetadata(repo_id="local", root=root) + depth_keys = set(meta.depth_keys) + video_keys = [k for k in meta.video_keys if k not in depth_keys] + camera_key = video_keys[0] if video_keys else None + + if camera_key: + logger.info("advantage-only mode: downloading only camera %s", camera_key) + snapshot_download( + repo_id=cfg.repo_id, + repo_type="dataset", + allow_patterns=[f"videos/{camera_key}/**"], + ) + return root + return Path(snapshot_download(repo_id=cfg.repo_id, repo_type="dataset")) raise ValueError("Either --root or --repo_id must be provided.") @@ -65,10 +96,11 @@ def annotate(cfg: AnnotationPipelineConfig) -> None: logger.info("annotate: root=%s", root) needs_vlm = cfg.plan.enabled or cfg.interjections.enabled or cfg.vqa.enabled + needs_video = needs_vlm or cfg.advantage.enabled vlm = make_vlm_client(cfg.vlm) if needs_vlm else None frame_provider = ( make_frame_provider(root, camera_key=cfg.vlm.camera_key, video_backend=cfg.video_backend) - if needs_vlm + if needs_video else None ) # Surface the resolved cameras up front so a silent vqa-module no-op @@ -105,7 +137,10 @@ def annotate(cfg: AnnotationPipelineConfig) -> None: if needs_vlm else None ) - advantage = AdvantageModule(config=cfg.advantage) + advantage = AdvantageModule( + config=cfg.advantage, + **({"frame_provider": frame_provider} if frame_provider is not None else {}), + ) writer = LanguageColumnsWriter() validator = StagingValidator( dataset_camera_keys=tuple(cam_keys) or None,