Loads episode_data_index and stats during dataset __init__ (#85)

Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Co-authored-by: Alexander Soare <alexander.soare159@gmail.com>

tests/data/aloha_sim_insertion_human/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 50
+}

tests/data/aloha_sim_insertion_human/train/dataset_info.json

@@ -21,11 +21,11 @@
       "length": 14,
       "_type": "Sequence"
     },
-    "episode_id": {
+    "episode_index": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "frame_id": {
+    "frame_index": {
       "dtype": "int64",
       "_type": "Value"
     },
@@ -37,14 +37,6 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_index_from": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
-    "episode_data_index_to": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
     "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": "d79cf82ffc86f110",
+  "_fingerprint": "22eeca7a3f4725ee",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/data/aloha_sim_insertion_scripted/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 50
+}

tests/data/aloha_sim_insertion_scripted/train/dataset_info.json

@@ -21,11 +21,11 @@
       "length": 14,
       "_type": "Sequence"
     },
-    "episode_id": {
+    "episode_index": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "frame_id": {
+    "frame_index": {
       "dtype": "int64",
       "_type": "Value"
     },
@@ -37,14 +37,6 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_index_from": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
-    "episode_data_index_to": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
     "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": "d8e4a817b5449498",
+  "_fingerprint": "97c28d4ad1536e4c",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/data/aloha_sim_transfer_cube_human/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 50
+}

tests/data/aloha_sim_transfer_cube_human/train/dataset_info.json

@@ -21,11 +21,11 @@
       "length": 14,
       "_type": "Sequence"
     },
-    "episode_id": {
+    "episode_index": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "frame_id": {
+    "frame_index": {
       "dtype": "int64",
       "_type": "Value"
     },
@@ -37,14 +37,6 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_index_from": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
-    "episode_data_index_to": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
     "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": "f03482befa767127",
+  "_fingerprint": "cb9349b5c92951e8",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/data/aloha_sim_transfer_cube_scripted/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 50
+}

tests/data/aloha_sim_transfer_cube_scripted/train/dataset_info.json

@@ -21,11 +21,11 @@
       "length": 14,
       "_type": "Sequence"
     },
-    "episode_id": {
+    "episode_index": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "frame_id": {
+    "frame_index": {
       "dtype": "int64",
       "_type": "Value"
     },
@@ -37,14 +37,6 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_index_from": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
-    "episode_data_index_to": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
     "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": "93e03c6320c7d56e",
+  "_fingerprint": "e4d7ad2b360db1af",
   "_format_columns": null,
   "_format_kwargs": {},
   "_format_type": "torch",

tests/data/pusht/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 10
+}

tests/data/pusht/train/dataset_info.json

@@ -21,11 +21,11 @@
       "length": 2,
       "_type": "Sequence"
     },
-    "episode_id": {
+    "episode_index": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "frame_id": {
+    "frame_index": {
       "dtype": "int64",
       "_type": "Value"
     },
@@ -45,14 +45,6 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_index_from": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
-    "episode_data_index_to": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
     "index": {
       "dtype": "int64",
       "_type": "Value"

tests/data/pusht/train/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 10
+}

tests/data/pusht/train/state.json

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

tests/data/xarm_lift_medium/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 15
+}

tests/data/xarm_lift_medium/train/dataset_info.json

@@ -21,11 +21,11 @@
       "length": 4,
       "_type": "Sequence"
     },
-    "episode_id": {
+    "episode_index": {
       "dtype": "int64",
       "_type": "Value"
     },
-    "frame_id": {
+    "frame_index": {
       "dtype": "int64",
       "_type": "Value"
     },
@@ -41,14 +41,6 @@
       "dtype": "bool",
       "_type": "Value"
     },
-    "episode_data_index_from": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
-    "episode_data_index_to": {
-      "dtype": "int64",
-      "_type": "Value"
-    },
     "index": {
       "dtype": "int64",
       "_type": "Value"

tests/data/xarm_lift_medium/train/state.json

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

tests/data/xarm_lift_medium_replay/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 15
+}

tests/data/xarm_lift_medium_replay/train/dataset_info.json

@@ -0,0 +1,51 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.image": {
+      "_type": "Image"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 4,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 4,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.reward": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/xarm_lift_medium_replay/train/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "9f8e1a8c1845df55",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": "torch",
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/xarm_push_medium/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 15
+}

tests/data/xarm_push_medium/train/dataset_info.json

@@ -0,0 +1,51 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.image": {
+      "_type": "Image"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 4,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 3,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.reward": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/xarm_push_medium/train/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "c900258061dd0b3f",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": "torch",
+  "_output_all_columns": false,
+  "_split": null
+}

tests/data/xarm_push_medium_replay/meta_data/info.json

@@ -0,0 +1,3 @@
+{
+    "fps": 15
+}

tests/data/xarm_push_medium_replay/train/dataset_info.json

@@ -0,0 +1,51 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "observation.image": {
+      "_type": "Image"
+    },
+    "observation.state": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 4,
+      "_type": "Sequence"
+    },
+    "action": {
+      "feature": {
+        "dtype": "float32",
+        "_type": "Value"
+      },
+      "length": 3,
+      "_type": "Sequence"
+    },
+    "episode_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "frame_index": {
+      "dtype": "int64",
+      "_type": "Value"
+    },
+    "timestamp": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.reward": {
+      "dtype": "float32",
+      "_type": "Value"
+    },
+    "next.done": {
+      "dtype": "bool",
+      "_type": "Value"
+    },
+    "index": {
+      "dtype": "int64",
+      "_type": "Value"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

tests/data/xarm_push_medium_replay/train/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "e51c80a33c7688c0",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": "torch",
+  "_output_all_columns": false,
+  "_split": null
+}