Address comments

tests/data/aloha_sim_insertion_human/train/dataset_info.json

@@ -2,6 +2,9 @@
   "citation": "",
   "description": "",
   "features": {
+    "observation.images.top": {
+      "_type": "Image"
+    },
     "observation.state": {
       "feature": {
         "dtype": "float32",
@@ -34,17 +37,14 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_id_from": {
+    "episode_data_index_from": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "episode_data_id_to": {
+    "episode_data_index_to": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "observation.images.top": {
-      "_type": "Image"
-    },
     "index": {
       "dtype": "int64",
       "_type": "Value"

tests/data/aloha_sim_insertion_human/train/state.json

@@ -4,7 +4,7 @@
       "filename": "data-00000-of-00001.arrow"
     }
   ],
-  "_fingerprint": "05980bca35112ebd",
+  "_fingerprint": "d79cf82ffc86f110",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/data/aloha_sim_insertion_scripted/train/dataset_info.json

@@ -2,6 +2,9 @@
   "citation": "",
   "description": "",
   "features": {
+    "observation.images.top": {
+      "_type": "Image"
+    },
     "observation.state": {
       "feature": {
         "dtype": "float32",
@@ -34,17 +37,14 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_id_from": {
+    "episode_data_index_from": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "episode_data_id_to": {
+    "episode_data_index_to": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "observation.images.top": {
-      "_type": "Image"
-    },
     "index": {
       "dtype": "int64",
       "_type": "Value"

tests/data/aloha_sim_insertion_scripted/train/state.json

@@ -4,7 +4,7 @@
       "filename": "data-00000-of-00001.arrow"
     }
   ],
-  "_fingerprint": "f3330a7e1d8bc55b",
+  "_fingerprint": "d8e4a817b5449498",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/data/aloha_sim_transfer_cube_human/train/dataset_info.json

@@ -2,6 +2,9 @@
   "citation": "",
   "description": "",
   "features": {
+    "observation.images.top": {
+      "_type": "Image"
+    },
     "observation.state": {
       "feature": {
         "dtype": "float32",
@@ -34,17 +37,14 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_id_from": {
+    "episode_data_index_from": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "episode_data_id_to": {
+    "episode_data_index_to": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "observation.images.top": {
-      "_type": "Image"
-    },
     "index": {
       "dtype": "int64",
       "_type": "Value"

tests/data/aloha_sim_transfer_cube_human/train/state.json

@@ -4,7 +4,7 @@
       "filename": "data-00000-of-00001.arrow"
     }
   ],
-  "_fingerprint": "42aa77ffb6863924",
+  "_fingerprint": "f03482befa767127",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/data/aloha_sim_transfer_cube_scripted/train/dataset_info.json

@@ -2,6 +2,9 @@
   "citation": "",
   "description": "",
   "features": {
+    "observation.images.top": {
+      "_type": "Image"
+    },
     "observation.state": {
       "feature": {
         "dtype": "float32",
@@ -34,17 +37,14 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_id_from": {
+    "episode_data_index_from": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "episode_data_id_to": {
+    "episode_data_index_to": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "observation.images.top": {
-      "_type": "Image"
-    },
     "index": {
       "dtype": "int64",
       "_type": "Value"

tests/data/aloha_sim_transfer_cube_scripted/train/state.json

@@ -4,7 +4,7 @@
       "filename": "data-00000-of-00001.arrow"
     }
   ],
-  "_fingerprint": "43f176a3740fe622",
+  "_fingerprint": "93e03c6320c7d56e",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/data/pusht/train/dataset_info.json

@@ -45,11 +45,11 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_id_from": {
+    "episode_data_index_from": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "episode_data_id_to": {
+    "episode_data_index_to": {
       "dtype": "int64",
       "_type": "Value"
     },

tests/data/pusht/train/state.json

@@ -4,7 +4,7 @@
       "filename": "data-00000-of-00001.arrow"
     }
   ],
-  "_fingerprint": "f7ed966ae18000ae",
+  "_fingerprint": "21bb9a76ed78a475",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/data/xarm_lift_medium/train/dataset_info.json

@@ -41,11 +41,11 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_id_from": {
+    "episode_data_index_from": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "episode_data_id_to": {
+    "episode_data_index_to": {
       "dtype": "int64",
       "_type": "Value"
     },

tests/data/xarm_lift_medium/train/state.json

@@ -4,7 +4,7 @@
       "filename": "data-00000-of-00001.arrow"
     }
   ],
-  "_fingerprint": "7dcd82fc3815bba6",
+  "_fingerprint": "a95cbec45e3bb9d6",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/test_datasets.py

@@ -95,12 +95,14 @@ def test_compute_stats():
     """
     from lerobot.common.datasets.xarm import XarmDataset
 
+    DATA_DIR = Path(os.environ["DATA_DIR"]) if "DATA_DIR" in os.environ else None
 
     # get transform to convert images from uint8 [0,255] to float32 [0,1]
     transform = Prod(in_keys=XarmDataset.image_keys, prod=1 / 255.0)
 
     dataset = XarmDataset(
         dataset_id="xarm_lift_medium",
+        root=DATA_DIR,
         transform=transform,
     )
 
@@ -115,11 +117,11 @@ def test_compute_stats():
     # get all frames from the dataset in the same dtype and range as during compute_stats
     dataloader = torch.utils.data.DataLoader(
         dataset,
-        num_workers=16,
+        num_workers=8,
         batch_size=len(dataset),
         shuffle=False,
     )
-    data_dict = next(iter(dataloader))  # takes 23 seconds
+    data_dict = next(iter(dataloader))
 
     # compute stats based on all frames from the dataset without any batching
     expected_stats = {}
@@ -154,8 +156,8 @@ def test_load_previous_and_future_frames_within_tolerance():
     data_dict = Dataset.from_dict({
         "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5],
         "index": [0, 1, 2, 3, 4],
-        "episode_data_id_from": [0, 0, 0, 0, 0],
-        "episode_data_id_to": [4, 4, 4, 4, 4],
+        "episode_data_index_from": [0, 0, 0, 0, 0],
+        "episode_data_index_to": [4, 4, 4, 4, 4],
     })
     data_dict = data_dict.with_format("torch")
     item = data_dict[2]
@@ -170,8 +172,8 @@ def test_load_previous_and_future_frames_outside_tolerance_inside_episode_range(
     data_dict = Dataset.from_dict({
         "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5],
         "index": [0, 1, 2, 3, 4],
-        "episode_data_id_from": [0, 0, 0, 0, 0],
-        "episode_data_id_to": [4, 4, 4, 4, 4],
+        "episode_data_index_from": [0, 0, 0, 0, 0],
+        "episode_data_index_to": [4, 4, 4, 4, 4],
     })
     data_dict = data_dict.with_format("torch")
     item = data_dict[2]
@@ -184,8 +186,8 @@ def test_load_previous_and_future_frames_outside_tolerance_outside_episode_range
     data_dict = Dataset.from_dict({
         "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5],
         "index": [0, 1, 2, 3, 4],
-        "episode_data_id_from": [0, 0, 0, 0, 0],
-        "episode_data_id_to": [4, 4, 4, 4, 4],
+        "episode_data_index_from": [0, 0, 0, 0, 0],
+        "episode_data_index_to": [4, 4, 4, 4, 4],
     })
     data_dict = data_dict.with_format("torch")
     item = data_dict[2]