From eecf32e77a73c3e4ad243021719a41adaa3b0f82 Mon Sep 17 00:00:00 2001
From: Haskely <Haskely@live.com>
Date: Wed, 26 Feb 2025 00:27:36 +0800
Subject: [PATCH] feat: Add root directory option for dataset configuration
 (#765)

Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
---
 lerobot/common/datasets/factory.py | 5 ++++-
 lerobot/configs/default.py         | 2 ++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/lerobot/common/datasets/factory.py b/lerobot/common/datasets/factory.py
index fb1fe6d6f..38c01b42f 100644
--- a/lerobot/common/datasets/factory.py
+++ b/lerobot/common/datasets/factory.py
@@ -83,10 +83,13 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
     )
 
     if isinstance(cfg.dataset.repo_id, str):
-        ds_meta = LeRobotDatasetMetadata(cfg.dataset.repo_id, revision=cfg.dataset.revision)
+        ds_meta = LeRobotDatasetMetadata(
+            cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
+        )
         delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
         dataset = LeRobotDataset(
             cfg.dataset.repo_id,
+            root=cfg.dataset.root,
             episodes=cfg.dataset.episodes,
             delta_timestamps=delta_timestamps,
             image_transforms=image_transforms,
diff --git a/lerobot/configs/default.py b/lerobot/configs/default.py
index a50134315..0dc7f4c73 100644
--- a/lerobot/configs/default.py
+++ b/lerobot/configs/default.py
@@ -29,6 +29,8 @@ class DatasetConfig:
     # "dataset_index" into the returned item. The index mapping is made according to the order in which the
     # datsets are provided.
     repo_id: str
+    # Root directory where the dataset will be stored (e.g. 'dataset/path').
+    root: str | None = None
     episodes: list[int] | None = None
     image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
     revision: str | None = None