mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-06 08:29:46 +00:00
feat(Train): enable buckets with streaming dataset (#4312)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user