feat(annotate): support video datasets in VF advantage scoring

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:
+            for local_i in range(len(batch_indices)):
-                img_val = df.iloc[idx][image_key]
+                if video_frames is not None:
-                if isinstance(img_val, np.ndarray):
+                    img_tensor = video_frames[batch_start + local_i].float()
-                    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)
+                    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:

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
+            df = _read_parquet_as_pandas(self.data_path)
-
-            table = pq.read_table(self.data_path)
-            df: pd.DataFrame = table.to_pandas()
             self._frames_df_cache = df.iloc[self.row_offset : self.row_offset + self.row_count].reset_index(
                 drop=True
             )

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,