From e867359d0915ac448a53860d8c16b7a02ed0cbca Mon Sep 17 00:00:00 2001
From: Steven Palma <imstevenpmwork@ieee.org>
Date: Mon, 3 Aug 2026 16:00:48 +0200
Subject: [PATCH] feat(Train): enable buckets with streaming dataset (#4312)

---
 docs/source/lerobot-dataset-v3.mdx | 16 ++++++++++++++++
 src/lerobot/configs/default.py     | 13 +++++++++++++
 src/lerobot/datasets/factory.py    | 10 +++++++++-
 tests/configs/test_default.py      | 19 +++++++++++++++++++
 4 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/docs/source/lerobot-dataset-v3.mdx b/docs/source/lerobot-dataset-v3.mdx
index 0647af0b0..9ae7588db 100644
--- a/docs/source/lerobot-dataset-v3.mdx
+++ b/docs/source/lerobot-dataset-v3.mdx
@@ -142,6 +142,22 @@ repo_id = "yaak-ai/L2D-v3"
 dataset = StreamingLeRobotDataset(repo_id)  # streams directly from the Hub
 ```
 
+Datasets stored in an [HF Storage Bucket](https://huggingface.co/docs/hub/storage-buckets) (`hf://buckets/`) can be streamed the same way by passing `repo_type="bucket"`:
+
+```python
+dataset = StreamingLeRobotDataset("my-org/my-bucket", repo_type="bucket")
+```
+
+Both options are available in `lerobot-train` through `--dataset.streaming=true`, and `--dataset.repo_type=bucket` to stream from a bucket instead of a Hub dataset repo:
+
+```bash
+lerobot-train \
+  --dataset.repo_id=my-org/my-bucket \
+  --dataset.repo_type=bucket \
+  --dataset.streaming=true \
+  ...
+```
+
 <div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
   <figure style="margin:0; text-align:center;">
     <img
diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py
index 38991a665..0734d0bc2 100644
--- a/src/lerobot/configs/default.py
+++ b/src/lerobot/configs/default.py
@@ -29,6 +29,9 @@ class DatasetConfig:
     # "dataset_index" into the returned item. The index mapping is made according to the order in which the
     # datasets are provided.
     repo_id: str
+    # Hub repository type: "dataset" (default) or "bucket" for an HF Storage Bucket streamed over
+    # hf://buckets/. Buckets are streaming-only, so "bucket" requires streaming=true.
+    repo_type: str = "dataset"
     # Root directory for a concrete local dataset tree (e.g. 'dataset/path'). If None, local datasets are
     # looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
     root: str | None = None
@@ -48,6 +51,16 @@ class DatasetConfig:
     eval_split: float = 0.0
 
     def __post_init__(self) -> None:
+        if self.repo_type not in ("dataset", "bucket"):
+            raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {self.repo_type!r}")
+        if self.repo_type == "bucket" and not self.streaming:
+            raise ValueError(
+                "repo_type='bucket' is streaming-only: set streaming=true to train from an HF Storage Bucket."
+            )
+        if self.repo_type == "bucket" and self.eval_split != 0.0:
+            raise ValueError(
+                "eval_split requires map-style datasets and is not supported with repo_type='bucket'."
+            )
         if self.depth_output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT):
             raise ValueError(
                 f"depth_output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {self.depth_output_unit!r}"
diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py
index da7b4365a..a727bf924 100644
--- a/src/lerobot/datasets/factory.py
+++ b/src/lerobot/datasets/factory.py
@@ -84,10 +84,17 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
 
     if isinstance(cfg.dataset.repo_id, str):
         ds_meta = LeRobotDatasetMetadata(
-            cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
+            cfg.dataset.repo_id,
+            root=cfg.dataset.root,
+            revision=cfg.dataset.revision,
+            repo_type=cfg.dataset.repo_type,
         )
         delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
         if not cfg.dataset.streaming:
+            if cfg.dataset.repo_type == "bucket":
+                raise ValueError(
+                    "repo_type='bucket' is streaming-only: set dataset.streaming=true to train from an HF Storage Bucket."
+                )
             dataset = LeRobotDataset(
                 cfg.dataset.repo_id,
                 root=cfg.dataset.root,
@@ -111,6 +118,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
                 max_num_shards=cfg.num_workers,
                 tolerance_s=cfg.tolerance_s,
                 return_uint8=True,
+                repo_type=cfg.dataset.repo_type,
             )
     else:
         raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
diff --git a/tests/configs/test_default.py b/tests/configs/test_default.py
index 238b8bacd..979509ec3 100644
--- a/tests/configs/test_default.py
+++ b/tests/configs/test_default.py
@@ -36,3 +36,22 @@ def test_dataset_config_none_episodes_ok():
 
 def test_dataset_config_empty_episodes_ok():
     DatasetConfig(repo_id="user/repo", episodes=[])
+
+
+def test_dataset_config_bucket_streaming_ok():
+    DatasetConfig(repo_id="user/repo", repo_type="bucket", streaming=True)
+
+
+def test_dataset_config_invalid_repo_type():
+    with pytest.raises(ValueError, match="repo_type"):
+        DatasetConfig(repo_id="user/repo", repo_type="model")
+
+
+def test_dataset_config_bucket_requires_streaming():
+    with pytest.raises(ValueError, match="streaming-only"):
+        DatasetConfig(repo_id="user/repo", repo_type="bucket")
+
+
+def test_dataset_config_bucket_rejects_eval_split():
+    with pytest.raises(ValueError, match="eval_split"):
+        DatasetConfig(repo_id="user/repo", repo_type="bucket", streaming=True, eval_split=0.1)