Adjustments and Documentation, recording and new environments, refactoring

environments/logging/envmonitor.py

@@ -71,8 +71,8 @@ class EnvMonitor(BaseCallback):
             pass
         return
 
-    def save_run(self, filepath: Union[Path, str], auto_plotting_keys=None):
-        filepath = Path(filepath)
+    def save_run(self, filepath: Union[Path, str, None] = None, auto_plotting_keys=None):
+        filepath = Path(filepath or self._filepath)
         filepath.parent.mkdir(exist_ok=True, parents=True)
         with filepath.open('wb') as f:
             pickle.dump(self._monitor_df.reset_index(), f, protocol=pickle.HIGHEST_PROTOCOL)

environments/logging/recorder.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+from os import PathLike
 from pathlib import Path
 from typing import Union
 
@@ -12,11 +13,14 @@ from environments.factory.base.base_factory import REC_TAC
 
 class EnvRecorder(BaseCallback):
 
-    def __init__(self, env, entities='all'):
+    def __init__(self, env, entities: str = 'all', filepath: Union[str, PathLike] = None, freq: int = 0):
         super(EnvRecorder, self).__init__()
+        self.filepath = filepath
         self.unwrapped = env
+        self.freq = freq
         self._recorder_dict = defaultdict(list)
         self._recorder_out_list = list()
+        self._episode_counter = 1
         if isinstance(entities, str):
             if entities.lower() == 'all':
                 self._entities = None
@@ -29,46 +33,47 @@ class EnvRecorder(BaseCallback):
         return getattr(self.unwrapped, item)
 
     def reset(self):
-        self.unwrapped.start_recording()
+        self._on_training_start()
         return self.unwrapped.reset()
 
     def _on_training_start(self) -> None:
-        self.unwrapped._record_episodes = True
-        pass
+        assert self.start_recording()
 
     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}
-
-            info_dict.update(episode=(self.num_timesteps + env_idx))
             self._recorder_dict[env_idx].append(info_dict)
         else:
             pass
-        return
+        return True
 
     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)})
+                                            'episode': self._episode_counter})
             self._recorder_dict[env_idx] = list()
         else:
             pass
 
     def step(self, actions):
         step_result = self.unwrapped.step(actions)
-        # 0, 1,     2    ,     3    =    idx
-        # _, _, done_bool, info_obj = step_result
-        self._read_info(0, step_result[3])
-        self._read_done(0, step_result[2])
+        self._on_step()
         return step_result
 
-    def save_records(self, filepath: Union[Path, str], save_occupation_map=False, save_trajectory_map=False):
-        filepath = Path(filepath)
+    def finalize(self):
+        self._on_training_end()
+        return True
+
+    def save_records(self, filepath: Union[Path, str, None] = None, save_occupation_map=False, save_trajectory_map=False):
+        filepath = Path(filepath or self.filepath)
         filepath.parent.mkdir(exist_ok=True, parents=True)
         # cls.out_file.unlink(missing_ok=True)
         with filepath.open('w') as f:
-            out_dict = {'episodes': self._recorder_out_list, 'header': self.unwrapped.params}
+            out_dict = {'n_episodes': self._episode_counter,
+                        'header': self.unwrapped.params,
+                        'episodes': self._recorder_out_list
+                        }
             try:
                 simplejson.dump(out_dict, f, indent=4)
             except TypeError:
@@ -76,6 +81,7 @@ class EnvRecorder(BaseCallback):
 
         if save_occupation_map:
             a = np.zeros((15, 15))
+            # noinspection PyTypeChecker
             for episode in out_dict['episodes']:
                 df = pd.DataFrame([y for x in episode['steps'] for y in x['Agents']])
 
@@ -94,15 +100,22 @@ class EnvRecorder(BaseCallback):
             raise NotImplementedError('This has not yet been implemented.')
 
     def _on_step(self) -> bool:
+        do_record = self.freq == -1 or self._episode_counter % self.freq == 0
         for env_idx, info in enumerate(self.locals.get('infos', [])):
-            self._read_info(env_idx, info)
-
+            if do_record:
+                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)
-
+            if do_record:
+                self._read_done(env_idx, done)
+            if done:
+                self._episode_counter += 1
         return True
 
     def _on_training_end(self) -> None:
+        for env_idx in range(len(self._recorder_dict)):
+            if self._recorder_dict[env_idx]:
+                self._recorder_out_list.append({'steps': self._recorder_dict[env_idx],
+                                                'episode': self._episode_counter})
         pass