Experiments look good

environments/logging/envmonitor.py

@@ -1,5 +1,6 @@
 import pickle
 from collections import defaultdict
+from os import PathLike
 from pathlib import Path
 from typing import List, Dict, Union
 
@@ -9,14 +10,17 @@ from environments.helpers import IGNORED_DF_COLUMNS
 
 import pandas as pd
 
+from plotting.compare_runs import plot_single_run
+
 
 class EnvMonitor(BaseCallback):
 
     ext = 'png'
 
-    def __init__(self, env):
+    def __init__(self, env, filepath: Union[str, PathLike] = None):
         super(EnvMonitor, self).__init__()
         self.unwrapped = env
+        self._filepath = filepath
         self._monitor_df = pd.DataFrame()
         self._monitor_dicts = defaultdict(dict)
 
@@ -67,8 +71,10 @@ class EnvMonitor(BaseCallback):
             pass
         return
 
-    def save_run(self, filepath: Union[Path, str]):
+    def save_run(self, filepath: Union[Path, str], auto_plotting_keys=None):
         filepath = Path(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)
+        if auto_plotting_keys:
+            plot_single_run(filepath, column_keys=auto_plotting_keys)

environments/logging/recorder.py

@@ -24,14 +24,12 @@ class EnvRecorder(BaseCallback):
                 self._entities = [entities]
         else:
             self._entities = entities
-        self.started = False
-        self.closed = False
 
     def __getattr__(self, item):
         return getattr(self.unwrapped, item)
 
     def reset(self):
-        self.unwrapped._record_episodes = True
+        self.unwrapped.start_recording()
         return self.unwrapped.reset()
 
     def _on_training_start(self) -> None:
@@ -57,6 +55,14 @@ class EnvRecorder(BaseCallback):
         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])
+        return step_result
+
     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)