Merge branch 'master' of https://github.com/vacancy/Jacinle

* 'master' of https://github.com/vacancy/Jacinle:
  [master] more functionality to gym recorder

Author: vacancy

jaclearn/rl/engines/mujoco/gym_recorder.py

@@ -31,6 +31,7 @@ def __init__(self, gym_environ):
 
         self.object_names = dict()
         self.object_poses = list()
+        self.logs = dict()
 
         self.reset()
 
@@ -43,8 +44,15 @@ def reset(self):
         self._init_object_names()
         self.gym_states = list()
         self.object_poses = list()
+        self.logs = dict()
         self.step()
 
+    def log(self, key, message):
+        log_index = len(self.gym_states) - 1
+        if log_index not in self.logs:
+            self.logs[log_index] = dict()
+        self.logs[log_index][key] = message
+
     def step(self):
         self.gym_states.append(self.sim.get_state())
         new_object_poses = dict()
@@ -66,6 +74,7 @@ def dump(self, save_dir):
             names=self.object_names,
             poses=self.object_poses
         ))
+        io.dump(osp.join(save_dir, 'logs.json'), self.logs)
 
     def hook(self):
         old_reset = self.gym_environ.reset

jactorch/data/dataset.py

@@ -87,6 +87,10 @@ def __len__(self):
         return len(self.list)
 
 
+class JacDataset(Dataset, IterableDatasetMixin):
+    pass
+
+
 class FilterableDatasetUnwrapped(Dataset, IterableDatasetMixin):
     """
     A filterable dataset. User can call various `filter_*` operations to obtain a subset of the dataset.