feat(Train): enable buckets with streaming dataset (#4312)

This commit is contained in:
Steven Palma
2026-08-03 16:00:48 +02:00
committed by GitHub
parent f1efa588b8
commit e867359d09
4 changed files with 57 additions and 1 deletions
+16
View File
@@ -142,6 +142,22 @@ repo_id = "yaak-ai/L2D-v3"
dataset = StreamingLeRobotDataset(repo_id) # streams directly from the Hub 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;"> <div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
<figure style="margin:0; text-align:center;"> <figure style="margin:0; text-align:center;">
<img <img
+13
View File
@@ -29,6 +29,9 @@ class DatasetConfig:
# "dataset_index" into the returned item. The index mapping is made according to the order in which the # "dataset_index" into the returned item. The index mapping is made according to the order in which the
# datasets are provided. # datasets are provided.
repo_id: str 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 # 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. # 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 root: str | None = None
@@ -48,6 +51,16 @@ class DatasetConfig:
eval_split: float = 0.0 eval_split: float = 0.0
def __post_init__(self) -> None: 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): if self.depth_output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT):
raise ValueError( raise ValueError(
f"depth_output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {self.depth_output_unit!r}" f"depth_output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {self.depth_output_unit!r}"
+9 -1
View File
@@ -84,10 +84,17 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
if isinstance(cfg.dataset.repo_id, str): if isinstance(cfg.dataset.repo_id, str):
ds_meta = LeRobotDatasetMetadata( 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) delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
if not cfg.dataset.streaming: 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( dataset = LeRobotDataset(
cfg.dataset.repo_id, cfg.dataset.repo_id,
root=cfg.dataset.root, root=cfg.dataset.root,
@@ -111,6 +118,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
max_num_shards=cfg.num_workers, max_num_shards=cfg.num_workers,
tolerance_s=cfg.tolerance_s, tolerance_s=cfg.tolerance_s,
return_uint8=True, return_uint8=True,
repo_type=cfg.dataset.repo_type,
) )
else: else:
raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.") raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
+19
View File
@@ -36,3 +36,22 @@ def test_dataset_config_none_episodes_ok():
def test_dataset_config_empty_episodes_ok(): def test_dataset_config_empty_episodes_ok():
DatasetConfig(repo_id="user/repo", episodes=[]) 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)