feat(datasets): support episode selection and filtering in streaming

StreamingLeRobotDataset previously accepted an `episodes` argument but
never applied it, so the whole stream was returned regardless. Resolve
the selection against the dataset range, apply it as a lazy filter over
the streamed shards, and add an `episode_filter` predicate (mirroring
LeRobotDataset). Empty selections and non-matching filters now raise
instead of silently streaming everything, and num_frames/num_episodes
reflect the selection.

src/lerobot/datasets/streaming_dataset.py

@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
 from collections import deque
 from collections.abc import Callable, Generator, Iterable, Iterator
 from pathlib import Path
@@ -34,6 +35,7 @@ from .utils import (
     check_version_compatibility,
     find_float_index,
     is_float_in_list,
+    resolve_episode_indices,
     safe_shard,
 )
 from .video_utils import (
@@ -42,6 +44,8 @@ from .video_utils import (
     decode_video_frames_torchcodec,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class LookBackError(Exception):
     """
@@ -248,6 +252,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         repo_id: str,
         root: str | Path | None = None,
         episodes: list[int] | None = None,
+        episode_filter: Callable[[dict], bool] | None = None,
         image_transforms: Callable | None = None,
         delta_timestamps: dict[list[float]] | None = None,
         tolerance_s: float = 1e-4,
@@ -274,6 +279,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
                 When omitted, Hub metadata is resolved through the cache under ``$HF_LEROBOT_HOME/hub``.
             episodes (list[int] | None, optional): If specified, this will only load episodes specified by
                 their episode_index in this list.
+            episode_filter (Callable[[dict], bool] | None, optional): Predicate over per-episode
+                metadata rows used to select episodes (e.g. ``lambda ep: ep["length"] >= 100``).
+                Intersected with ``episodes`` when both are set. Defaults to None.
             image_transforms (Callable | None, optional): Transform to apply to image data.
             tolerance_s (float, optional): Tolerance in seconds for timestamp matching.
             revision (str, optional): Git revision id (branch name, tag, or commit hash).
@@ -305,7 +313,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         self.streaming_from_local = root is not None and self.repo_type == "dataset"
 
         self.image_transforms = image_transforms
-        self.episodes = episodes
         self.tolerance_s = tolerance_s
         self.revision = revision if revision else CODEBASE_VERSION
         self.seed = seed
@@ -338,6 +345,21 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         # Check version
         check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION)
 
+        self.episodes = resolve_episode_indices(episodes, self.meta.total_episodes)
+        if self.episodes is not None and not self.episodes:
+            raise ValueError(
+                "No valid episodes: the requested episode selection is empty after resolving "
+                f"against the dataset range [0, {self.meta.total_episodes})."
+            )
+        if episode_filter is not None:
+            resolved = self.meta.filter_episodes(episode_filter, candidates=self.episodes)
+            if not resolved:
+                raise ValueError(
+                    "The episode filter did not match any episode. Make sure the filter and episodes list are valid and compatible."
+                )
+            logger.info(f"The episode filter matched {len(resolved)} episode(s).")
+            self.episodes = resolved
+
         self._depth_encoder_configs: dict[str, DepthEncoderConfig] = {
             vid_key: DepthEncoderConfig.from_video_info(self.meta.features[vid_key].get("info"))
             for vid_key in self.meta.depth_keys
@@ -379,15 +401,22 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
                 **token_kwargs,
             )
 
+        # Streaming loads all shards, so episode selection is applied as a lazy stream filter.
+        if self.episodes is not None:
+            selected = set(self.episodes)
+            self.hf_dataset = self.hf_dataset.filter(lambda x: x["episode_index"] in selected)
+
         self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
 
     @property
     def num_frames(self):
-        return self.meta.total_frames
+        if self.episodes is None:
+            return self.meta.total_frames
+        return sum(self.meta.episodes[ep]["length"] for ep in self.episodes)
 
     @property
     def num_episodes(self):
-        return self.meta.total_episodes
+        return len(self.episodes) if self.episodes is not None else self.meta.total_episodes
 
     @property
     def fps(self):

tests/datasets/test_streaming.py

@@ -84,6 +84,7 @@ def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, mon
         root=requested_root or tmp_path / "snapshot",
         revision=streaming_dataset_module.CODEBASE_VERSION,
         _version=streaming_dataset_module.CODEBASE_VERSION,
+        total_episodes=10,
         features={},
         depth_keys=[],
         image_keys=[],
@@ -155,6 +156,81 @@ def test_single_frame_consistency(tmp_path, lerobot_dataset_factory):
         )
 
 
+def test_streaming_episode_selection(tmp_path, lerobot_dataset_factory):
+    """episodes=[...] restricts both the streamed frames and the count properties to the selection."""
+    ds_num_frames = 200
+    ds_num_episodes = 10
+    selected = [2, 5, 7]
+
+    local_path = tmp_path / "test"
+    repo_id = DUMMY_REPO_ID
+
+    ds = lerobot_dataset_factory(
+        root=local_path,
+        repo_id=repo_id,
+        total_episodes=ds_num_episodes,
+        total_frames=ds_num_frames,
+    )
+
+    streaming_ds = StreamingLeRobotDataset(
+        repo_id=repo_id, root=local_path, episodes=selected, buffer_size=50, shuffle=False
+    )
+
+    assert streaming_ds.num_episodes == len(selected)
+    assert streaming_ds.num_frames == sum(ds.meta.episodes[ep]["length"] for ep in selected)
+
+    episode_indices = [int(frame["episode_index"]) for frame in streaming_ds]
+    assert set(episode_indices) == set(selected)
+    assert len(episode_indices) == streaming_ds.num_frames
+
+
+def test_streaming_episode_filter(tmp_path, lerobot_dataset_factory):
+    """episode_filter restricts the stream to episodes whose metadata matches the predicate."""
+    ds_num_episodes = 6
+    ds_num_frames = 120  # 20 frames per episode
+
+    local_path = tmp_path / "test"
+    repo_id = DUMMY_REPO_ID
+
+    ds = lerobot_dataset_factory(
+        root=local_path,
+        repo_id=repo_id,
+        total_episodes=ds_num_episodes,
+        total_frames=ds_num_frames,
+    )
+
+    keep = {1, 4}
+    streaming_ds = StreamingLeRobotDataset(
+        repo_id=repo_id,
+        root=local_path,
+        episode_filter=lambda ep: ep["episode_index"] in keep,
+        buffer_size=50,
+        shuffle=False,
+    )
+
+    assert set(streaming_ds.episodes) == keep
+    assert streaming_ds.num_episodes == len(keep)
+    assert {int(frame["episode_index"]) for frame in streaming_ds} == keep
+
+
+@pytest.mark.parametrize(
+    "kwargs, match",
+    [
+        ({"episodes": [99]}, "No valid episodes"),
+        ({"episode_filter": lambda ep: False}, "episode filter did not match"),
+    ],
+)
+def test_streaming_empty_selection_raises(tmp_path, lerobot_dataset_factory, kwargs, match):
+    """An empty selection (out-of-range episodes or a non-matching filter) must fail loudly."""
+    local_path = tmp_path / "test"
+    repo_id = DUMMY_REPO_ID
+
+    lerobot_dataset_factory(root=local_path, repo_id=repo_id, total_episodes=4, total_frames=40)
+
+    with pytest.raises(ValueError, match=match):
+        StreamingLeRobotDataset(repo_id=repo_id, root=local_path, **kwargs)
+
+
 @pytest.mark.parametrize(
     "shuffle",
     [False, True],
@@ -494,6 +570,7 @@ def _fake_meta(*args, **kwargs):
     meta.root = root or "/tmp/_streaming_meta"
     meta.revision = revision or "v0"
     meta._version = "v3.0"
+    meta.total_episodes = 10
     meta.depth_keys = []
     meta.image_keys = []
     meta.rescale_depth_stats = lambda *_a, **_k: None
@@ -616,7 +693,7 @@ def test_repo_type_is_keyword_only_and_preserves_positional_episodes():
         patch("lerobot.datasets.streaming_dataset.check_version_compatibility"),
         patch(
             "lerobot.datasets.streaming_dataset.load_dataset",
-            return_value=SimpleNamespace(num_shards=1),
+            return_value=SimpleNamespace(num_shards=1, filter=lambda *a, **k: SimpleNamespace(num_shards=1)),
         ),
     ):
         dataset = StreamingLeRobotDataset(DUMMY_REPO_ID, None, episodes)