Refactor DataCollectorV0 and HDF5 dependencies isolation

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ repos:
     hooks:
       - id: flake8
         args:
-          - '--per-file-ignores=*/__init__.py:F401 test/all_parameter_combs_test.py:F405 pettingzoo/classic/go/go.py:W605'
+          - '--per-file-ignores=*/__init__.py:F401'
           - --extend-ignore=E203
           - --max-complexity=205
           - --max-line-length=300
@@ -48,7 +48,7 @@ repos:
           - --count
           # TODO: Remove ignoring rules D101, D102, D103, D105
           - --add-ignore=D100,D107,D101,D102,D103,D105
-        exclude: "__init__.py$|^pettingzoo.test|^docs"
+        exclude: "__init__.py$|^docs"
         additional_dependencies: ["toml"]
   - repo: local
     hooks:

minari/data_collector/callbacks/episode_metadata.py

@@ -1,29 +1,30 @@
-import h5py
+from typing import Dict
+
 import numpy as np
 
 
 class EpisodeMetadataCallback:
     """Callback to full episode after saving to hdf5 file as a group.
 
     This callback can be overridden to add extra metadata attributes or statistics to
-    each HDF5 episode group in the Minari dataset. The custom callback can then be
+    each episode in the Minari dataset. The custom callback can then be
     passed to the DataCollectorV0 wrapper to the `episode_metadata_callback` argument.
 
     TODO: add more default statistics to episode datasets
     """
 
-    def __call__(self, eps_group: h5py.Group):
+    def __call__(self, episode: Dict):
         """Callback method.
 
         Override this method to add custom attribute metadata to the episode group.
 
         Args:
-            eps_group (h5py.Group): the HDF5 group that contains an episode's datasets
+            episode (dict): the dict that contains an episode's data
         """
-        eps_group["rewards"].attrs["sum"] = np.sum(eps_group["rewards"])
-        eps_group["rewards"].attrs["mean"] = np.mean(eps_group["rewards"])
-        eps_group["rewards"].attrs["std"] = np.std(eps_group["rewards"])
-        eps_group["rewards"].attrs["max"] = np.max(eps_group["rewards"])
-        eps_group["rewards"].attrs["min"] = np.min(eps_group["rewards"])
-
-        eps_group.attrs["total_steps"] = eps_group["rewards"].shape[0]
+        return {
+            "rewards_sum": np.sum(episode["rewards"]),
+            "rewards_mean": np.mean(episode["rewards"]),
+            "rewards_std": np.std(episode["rewards"]),
+            "rewards_max": np.max(episode["rewards"]),
+            "rewards_min": np.min(episode["rewards"]),
+        }

minari/data_collector/callbacks/step_data.py

@@ -19,6 +19,7 @@ class StepData(TypedDict):
     "rewards",
     "truncations",
     "terminations",
+    "infos",
 }
 
 
@@ -58,8 +59,8 @@ def __call__(self, env, **kwargs):
 
                     return step_data
 
-        The episode groups in the HDF5 file of this Minari dataset will contain a subgroup called `environment_states` with dataset `velocity` and another subgroup called `pose`
-        with datasets `position` and `orientation`
+        The Minari dataset will contain a dictionary called `environment_states` with `velocity` value and another dictionary `pose`
+        with `position` and `orientation`
 
         Args:
             env (gym.Env): current Gymnasium environment.