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",