Monitor and Recorder are Wrappers.

environments/logging/recorder.py

@@ -1,4 +1,3 @@
-import json
 from collections import defaultdict
 from pathlib import Path
 from typing import Union
@@ -11,22 +10,13 @@ from stable_baselines3.common.callbacks import BaseCallback
 from environments.factory.base.base_factory import REC_TAC
 
 
-# noinspection PyAttributeOutsideInit
-from environments.helpers import Constants as c
+class EnvRecorder(BaseCallback):
 
-
-class RecorderCallback(BaseCallback):
-
-    def __init__(self, filepath: Union[str, Path], occupation_map: bool = False, trajectory_map: bool = False,
-                 entities='all'):
-        super(RecorderCallback, self).__init__()
-        self.trajectory_map = trajectory_map
-        self.occupation_map = occupation_map
-        self.filepath = Path(filepath)
+    def __init__(self, env, entities='all'):
+        super(EnvRecorder, self).__init__()
+        self.unwrapped = env
         self._recorder_dict = defaultdict(list)
         self._recorder_out_list = list()
-        self._env_params = None
-        self.do_record: bool
         if isinstance(entities, str):
             if entities.lower() == 'all':
                 self._entities = None
@@ -37,10 +27,18 @@ class RecorderCallback(BaseCallback):
         self.started = False
         self.closed = False
 
-    def read_params(self, params):
-        self._env_params = params
+    def __getattr__(self, item):
+        return getattr(self.unwrapped, item)
 
-    def read_info(self, env_idx, info: dict):
+    def reset(self):
+        self.unwrapped._record_episodes = True
+        return self.unwrapped.reset()
+
+    def _on_training_start(self) -> None:
+        self.unwrapped._record_episodes = True
+        pass
+
+    def _read_info(self, env_idx, info: dict):
         if info_dict := {key.replace(REC_TAC, ''): val for key, val in info.items() if key.startswith(f'{REC_TAC}')}:
             if self._entities:
                 info_dict = {k: v for k, v in info_dict.items() if k in self._entities}
@@ -51,7 +49,7 @@ class RecorderCallback(BaseCallback):
             pass
         return
 
-    def read_done(self, env_idx, done):
+    def _read_done(self, env_idx, done):
         if done:
             self._recorder_out_list.append({'steps': self._recorder_dict[env_idx],
                                             'episode': len(self._recorder_out_list)})
@@ -59,77 +57,46 @@ class RecorderCallback(BaseCallback):
         else:
             pass
 
-    def start(self, force=False):
-        if (hasattr(self.training_env, 'record_episodes') and self.training_env.record_episodes) or force:
-            self.do_record = True
-            self.filepath.parent.mkdir(exist_ok=True, parents=True)
-            self.started = True
-        else:
-            self.do_record = False
+    def save_records(self, filepath: Union[Path, str], save_occupation_map=False, save_trajectory_map=False):
+        filepath = Path(filepath)
+        filepath.parent.mkdir(exist_ok=True, parents=True)
+        # self.out_file.unlink(missing_ok=True)
+        with filepath.open('w') as f:
+            out_dict = {'episodes': self._recorder_out_list, 'header': self.unwrapped.params}
+            try:
+                simplejson.dump(out_dict, f, indent=4)
+            except TypeError:
+                print('Shit')
 
-    def stop(self):
-        if self.do_record and self.started:
-            # self.out_file.unlink(missing_ok=True)
-            with self.filepath.open('w') as f:
-                out_dict = {'episodes': self._recorder_out_list, 'header': self._env_params}
-                try:
-                    simplejson.dump(out_dict, f, indent=4)
-                except TypeError:
-                    print('Shit')
+        if save_occupation_map:
+            a = np.zeros((15, 15))
+            for episode in out_dict['episodes']:
+                df = pd.DataFrame([y for x in episode['steps'] for y in x['Agents']])
 
-            if self.occupation_map:
-                a = np.zeros((15, 15))
-                for episode in  out_dict['episodes']:
-                    df = pd.DataFrame([y for x in episode['steps'] for y in x['Agents']])
+                b = list(df[['x', 'y']].to_records(index=False))
 
-                    b = list(df[['x', 'y']].to_records(index=False))
+                np.add.at(a, tuple(zip(*b)), 1)
 
-                    np.add.at(a, tuple(zip(*b)), 1)
+            # a = np.rot90(a)
+            import seaborn as sns
+            from matplotlib import pyplot as plt
+            hm = sns.heatmap(data=a)
+            hm.set_title('Very Nice Heatmap')
+            plt.show()
 
-                # a = np.rot90(a)
-                import seaborn as sns
-                from matplotlib import pyplot as plt
-                hm = sns.heatmap(data=a)
-                hm.set_title('Very Nice Heatmap')
-                plt.show()
-
-            if self.trajectory_map:
-                print('Recorder files were dumped to disk, now plotting the occupation map...')
-
-            self.closed = True
-            self.started = False
-        else:
-            pass
+        if save_trajectory_map:
+            raise NotImplementedError('This has not yet been implemented.')
 
     def _on_step(self) -> bool:
-        if self.do_record and self.started:
-            for env_idx, info in enumerate(self.locals.get('infos', [])):
-                self.read_info(env_idx, info)
+        for env_idx, info in enumerate(self.locals.get('infos', [])):
+            self._read_info(env_idx, info)
+
+        dones = list(enumerate(self.locals.get('dones', [])))
+        dones.extend(list(enumerate(self.locals.get('done', []))))
+        for env_idx, done in dones:
+            self._read_done(env_idx, done)
 
-            for env_idx, done in list(
-                    enumerate(self.locals.get('dones', []))) + list(
-                enumerate(self.locals.get('done', []))):
-                self.read_done(env_idx, done)
-        else:
-            pass
         return True
 
-    def __enter__(self):
-        self.start(force=True)
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.stop()
-
-    def _on_training_start(self) -> None:
-        if self.started:
-            pass
-        else:
-            self.start()
-        pass
-
     def _on_training_end(self) -> None:
-        if self.closed:
-            pass
-        else:
-            self.stop()
+        pass