Add aloha dataset

2024-03-06 10:15:11 +00:00
parent 49c0955f97
commit d782b029e1
3 changed files with 216 additions and 0 deletions
--- a/lerobot/common/datasets/factory.py
+++ b/lerobot/common/datasets/factory.py
@@ -66,6 +66,12 @@ def make_offline_buffer(

        clsfunc = PushtExperienceReplay
        dataset_id = "pusht"
+
+    elif cfg.env.name == "aloha":
+        from lerobot.common.datasets.aloha import AlohaExperienceReplay
+
+        clsfunc = AlohaExperienceReplay
+        dataset_id = f"aloha_{cfg.env.task}"
    else:
        raise ValueError(cfg.env.name)