Derive streaming decoder cap from episode pool

docs/source/training_dataset_streaming.mdx

@@ -58,17 +58,17 @@ Subset sidecars cannot be published.
 
 The production defaults are:
 
-| Option                              | Default | Meaning                                             |
+| Option                              |        Default | Meaning                                             |
-| ----------------------------------- | ------: | --------------------------------------------------- |
+| ----------------------------------- | -------------: | --------------------------------------------------- |
-| `streaming_episode_pool_size`       |      32 | Maximum complete episodes mixed by each rank        |
+| `streaming_episode_pool_size`       |             32 | Maximum complete episodes mixed by each rank        |
-| `streaming_prefetch_episodes`       |       8 | Episodes fetched ahead of the active pool           |
+| `streaming_prefetch_episodes`       |              8 | Episodes fetched ahead of the active pool           |
-| `streaming_byte_budget_gb`          |       8 | Maximum synthesized MP4 bytes per rank              |
+| `streaming_byte_budget_gb`          |              8 | Maximum synthesized MP4 bytes per rank              |
-| `streaming_decode_threads`          |       2 | Parallel sample assembly and video decode workers   |
+| `streaming_decode_threads`          |              2 | Parallel sample assembly and video decode workers   |
-| `streaming_decoded_queue_size`      |       8 | Decoded samples buffered ahead, in planner order    |
+| `streaming_decoded_queue_size`      |              8 | Decoded samples buffered ahead, in planner order    |
-| `streaming_max_open_decoders`       |      64 | Independent open-decoder LRU cap per rank           |
+| `streaming_max_open_decoders`       | pool × cameras | Independent open-decoder LRU cap per rank           |
-| `streaming_native_http_connections` |   unset | Native HTTP connection cap per rank                 |
+| `streaming_native_http_connections` |          unset | Native HTTP connection cap per rank                 |
-| `streaming_native_http_subranges`   |       1 | Concurrent subranges per native HTTP range read     |
+| `streaming_native_http_subranges`   |              1 | Concurrent subranges per native HTTP range read     |
-| `streaming_data_root`               |   unset | Optional local, Hub, bucket, or fsspec payload root |
+| `streaming_data_root`               |          unset | Optional local, Hub, bucket, or fsspec payload root |
 
 The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes
 computed from the sidecar. An episode larger than the complete rank budget fails before training
@@ -84,7 +84,6 @@ lerobot-train \
   --dataset.streaming_byte_budget_gb=4 \
   --dataset.streaming_decode_threads=2 \
   --dataset.streaming_decoded_queue_size=8 \
-  --dataset.streaming_max_open_decoders=64 \
   --num_workers=4 \
   --policy.type=act \
   --output_dir=outputs/train/act_streaming

scripts/bench_streaming_dataset.py

@@ -52,7 +52,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--byte-budget-gb", type=float, default=8.0)
     parser.add_argument("--decode-threads", type=int, default=2)
     parser.add_argument("--decoded-queue-size", type=int, default=8)
-    parser.add_argument("--max-open-decoders", type=int, default=64)
+    parser.add_argument("--max-open-decoders", type=int, default=None)
     parser.add_argument("--native-http-connections", type=int, default=None)
     parser.add_argument("--native-http-subranges", type=int, default=1)
     parser.add_argument("--warmup-batches", type=int, default=8)
@@ -196,7 +196,8 @@ def main() -> None:
             "byte_budget_gb": args.byte_budget_gb,
             "decode_threads": args.decode_threads,
             "decoded_queue_size": args.decoded_queue_size,
-            "max_open_decoders": args.max_open_decoders,
+            "max_open_decoders": dataset.max_open_decoders,
+            "requested_max_open_decoders": args.max_open_decoders,
             "native_http_connections": args.native_http_connections,
             "native_http_subranges": args.native_http_subranges,
             "warmup_batches": args.warmup_batches,

src/lerobot/configs/default.py

@@ -55,8 +55,8 @@ class DatasetConfig:
     # Parallel sample assembly/decode workers and their bounded in-order result queue.
     streaming_decode_threads: int = 2
     streaming_decoded_queue_size: int = 8
-    # Independent decoder-state cap.
+    # Independent decoder-state cap. None covers every camera in the configured episode pool.
-    streaming_max_open_decoders: int = 64
+    streaming_max_open_decoders: int | None = None
     # Per-rank native HTTP limits. None preserves the fetcher's worker-derived default.
     streaming_native_http_connections: int | None = None
     streaming_native_http_subranges: int = 1
@@ -80,7 +80,7 @@ class DatasetConfig:
             raise ValueError("streaming_decode_threads must be positive")
         if self.streaming_decoded_queue_size <= 0:
             raise ValueError("streaming_decoded_queue_size must be positive")
-        if self.streaming_max_open_decoders <= 0:
+        if self.streaming_max_open_decoders is not None and self.streaming_max_open_decoders <= 0:
             raise ValueError("streaming_max_open_decoders must be positive")
         if self.streaming_native_http_connections is not None and self.streaming_native_http_connections <= 0:
             raise ValueError("streaming_native_http_connections must be positive")

src/lerobot/datasets/streaming_dataset.py

@@ -125,7 +125,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         token: str | bool | None = None,
         decode_threads: int = 2,
         decoded_queue_size: int = 8,
-        max_open_decoders: int = 64,
+        max_open_decoders: int | None = None,
         native_http_connections: int | None = None,
         native_http_subranges: int = 1,
     ):
@@ -161,7 +161,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             decode_threads (int, optional): Parallel sample assembly and video decode workers.
             decoded_queue_size (int, optional): Maximum number of decoded samples produced ahead
                 of the consumer. Results are yielded in exact planner order.
-            max_open_decoders (int, optional): Maximum number of open video decoders per rank.
+            max_open_decoders (int | None, optional): Maximum number of open video decoders per
+                rank. By default every camera in the configured episode pool can remain open.
             native_http_connections (int | None, optional): Native HTTP connection limit per rank.
                 ``None`` preserves the fetcher's worker-derived default.
             native_http_subranges (int, optional): Concurrent HTTP subranges per video range read.
@@ -213,7 +214,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             raise ValueError("decode_threads must be positive")
         if decoded_queue_size <= 0:
             raise ValueError("decoded_queue_size must be positive")
-        if max_open_decoders <= 0:
+        if max_open_decoders is not None and max_open_decoders <= 0:
             raise ValueError("max_open_decoders must be positive")
         if native_http_connections is not None and native_http_connections <= 0:
             raise ValueError("native_http_connections must be positive")
@@ -224,7 +225,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         self.byte_budget = int(byte_budget_gb * 1024**3)
         self.decode_threads = decode_threads
         self.decoded_queue_size = decoded_queue_size
-        self.max_open_decoders = max_open_decoders
         self.native_http_connections = native_http_connections
         self.native_http_subranges = native_http_subranges
         self.repeat = repeat
@@ -250,6 +250,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         self.meta.rescale_depth_stats(self._depth_output_unit)
         # Check version
         check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION)
+        self.max_open_decoders = (
+            max_open_decoders
+            if max_open_decoders is not None
+            else max(1, self.episode_pool_size * len(self.meta.video_keys))
+        )
 
         self._depth_encoder_configs: dict[str, DepthEncoderConfig] = {
             vid_key: DepthEncoderConfig.from_video_info(self.meta.features[vid_key].get("info"))

tests/configs/test_default.py

@@ -38,6 +38,10 @@ def test_dataset_config_empty_episodes_ok():
     DatasetConfig(repo_id="user/repo", episodes=[])
 
 
+def test_dataset_config_derives_streaming_decoder_limit_by_default():
+    assert DatasetConfig(repo_id="user/repo").streaming_max_open_decoders is None
+
+
 @pytest.mark.parametrize(
     ("field", "value", "message"),
     [

tests/datasets/test_streaming_production.py

@@ -118,6 +118,35 @@ def test_parallel_decode_queue_preserves_planner_order(
     assert 1 < max_active <= parallel.decode_threads
 
 
+def test_default_decoder_limit_covers_the_configured_episode_pool(
+    tmp_path: Path,
+    lerobot_dataset_factory,
+) -> None:
+    root = tmp_path / "dataset"
+    lerobot_dataset_factory(
+        root=root,
+        repo_id=DUMMY_REPO_ID,
+        total_episodes=2,
+        total_frames=20,
+    )
+
+    streaming = StreamingLeRobotDataset(
+        DUMMY_REPO_ID,
+        root=root,
+        episode_pool_size=7,
+    )
+
+    assert streaming.max_open_decoders == 7 * len(streaming.meta.video_keys)
+
+    overridden = StreamingLeRobotDataset(
+        DUMMY_REPO_ID,
+        root=root,
+        episode_pool_size=7,
+        max_open_decoders=5,
+    )
+    assert overridden.max_open_decoders == 5
+
+
 @pytest.mark.parametrize("video_backend", ["torchcodec", "pyav"])
 def test_streaming_rgb_video_matches_map_style(
     tmp_path: Path,