chore(comments): simplifying comments

Merge the separate episode_filter no-match and unknown-key tests into a
single parametrized test, and add coverage for an out-of-range episodes
list raising rather than silently resolving to empty.

src/lerobot/datasets/dataset_reader.py

@@ -39,7 +39,6 @@ from .io_utils import (
     hf_transform_to_torch,
     load_nested_dataset,
 )
-from .utils import resolve_episode_indices
 from .video_utils import decode_video_frames
 
 
@@ -69,8 +68,9 @@ class DatasetReader:
         Args:
             meta: Dataset metadata instance.
             root: Local dataset root directory.
-            episodes: Optional list of episode indices to select. ``None``
-                means all episodes.
+            episodes: Optional list of episode indices to select, assumed
+                already validated by the caller. ``None`` means
+                all episodes.
             tolerance_s: Timestamp synchronization tolerance in seconds.
             video_backend: Video decoding backend identifier.
             delta_timestamps: Optional dict mapping feature keys to lists of
@@ -84,7 +84,7 @@ class DatasetReader:
         """
         self._meta = meta
         self.root = root
-        self.episodes = resolve_episode_indices(episodes, meta.total_episodes)
+        self.episodes = episodes
         self._tolerance_s = tolerance_s
         self._video_backend = video_backend
         if image_transforms is not None and not callable(image_transforms):
@@ -152,9 +152,9 @@ class DatasetReader:
     @property
     def num_frames(self) -> int:
         """Number of frames in selected episodes."""
-        if self.episodes is not None and self.hf_dataset is not None:
-            return len(self.hf_dataset)
-        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) -> int:

src/lerobot/datasets/factory.py

@@ -91,6 +91,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             repo_type=cfg.dataset.repo_type,
         )
         delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
+        # Resolved here because dataset reader does not apply exclude_episodes.
         episodes = resolve_episode_indices(
             cfg.dataset.episodes, ds_meta.total_episodes, cfg.dataset.exclude_episodes
         )

src/lerobot/datasets/lerobot_dataset.py

@@ -34,6 +34,7 @@ from .utils import (
     create_lerobot_dataset_card,
     get_safe_version,
     is_valid_version,
+    resolve_episode_indices,
 )
 from .video_utils import (
     StreamingVideoEncoder,
@@ -237,13 +238,12 @@ class LeRobotDataset(torch.utils.data.Dataset):
         self.revision = self.meta.revision
         self.meta.rescale_depth_stats(self._depth_output_unit)
 
-        if episodes is not None and any(
-            episode >= self.meta.total_episodes or episode < 0 for episode in episodes
-        ):
-            logger.warning(
-                f"Some episodes in the provided episodes list are out of range for this dataset ({self.meta.total_episodes})."
+        episodes = resolve_episode_indices(episodes, self.meta.total_episodes)
+        if episodes is not None and not 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=episodes)
             if not resolved:

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_datasets.py

@@ -1828,25 +1828,21 @@ def test_episode_filter_intersects_with_episodes(tmp_path, lerobot_dataset_facto
     assert seen_eps == set(expected_eps)
 
 
-def test_episode_filter_no_match_raises(tmp_path, lerobot_dataset_factory):
-    """An empty match in LeRobotDataset's episode_filter raises a ValueError rather than silently returning an empty dataset."""
+@pytest.mark.parametrize(
+    "kwargs, exc, match",
+    [
+        (
+            {"episode_filter": lambda ep: ep["length"] < 0},
+            ValueError,
+            r"The episode filter did not match any episode",
+        ),
+        ({"episode_filter": lambda ep: ep["not_a_real_field"] > 0}, KeyError, "not_a_real_field"),
+        ({"episodes": [99]}, ValueError, "No valid episodes"),
+    ],
+)
+def test_episode_selection_invalid_raises(tmp_path, lerobot_dataset_factory, kwargs, exc, match):
+    """Invalid selections fail loudly: non-matching filter, unknown filter key, or out-of-range episodes."""
     dataset = lerobot_dataset_factory(root=tmp_path / "test", total_episodes=4, total_frames=100)
 
-    with pytest.raises(ValueError, match=r"The episode filter did not match any episode"):
-        LeRobotDataset(
-            dataset.repo_id,
-            root=dataset.root,
-            episode_filter=lambda ep: ep["length"] < 0,
-        )
-
-
-def test_episode_filter_unknown_key_raises(tmp_path, lerobot_dataset_factory):
-    """A predicate referencing a column absent from meta.episodes surfaces a clear KeyError."""
-    dataset = lerobot_dataset_factory(root=tmp_path / "test", total_episodes=4, total_frames=100)
-
-    with pytest.raises(KeyError, match="not_a_real_field"):
-        LeRobotDataset(
-            dataset.repo_id,
-            root=dataset.root,
-            episode_filter=lambda ep: ep["not_a_real_field"] > 0,
-        )
+    with pytest.raises(exc, match=match):
+        LeRobotDataset(dataset.repo_id, root=dataset.root, **kwargs)

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
+
+    local_path = tmp_path / "test"
+    repo_id = DUMMY_REPO_ID
+
+    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)