plotting

environments/logging/monitor.py

@@ -4,6 +4,8 @@ from collections import defaultdict
 
 from stable_baselines3.common.callbacks import BaseCallback
 
+from environments.logging.plotting import prepare_plot
+
 
 class FactoryMonitor:
 
@@ -58,11 +60,12 @@ class MonitorCallback(BaseCallback):
 
     ext = 'png'
 
-    def __init__(self, env, filepath=Path('debug_out/monitor.pick')):
+    def __init__(self, env, filepath=Path('debug_out/monitor.pick'), plotting=True):
         super(MonitorCallback, self).__init__()
         self.filepath = Path(filepath)
         self._monitor_list = list()
         self.env = env
+        self.plotting = plotting
         self.started = False
         self.closed = False
 
@@ -91,7 +94,18 @@ class MonitorCallback(BaseCallback):
             # self.out_file.unlink(missing_ok=True)
             with self.filepath.open('wb') as f:
                 pickle.dump(self.monitor_as_df_list, f, protocol=pickle.HIGHEST_PROTOCOL)
-            self.prepare_plot()
+            if self.plotting:
+                print('Monitor files were dumped to disk, now plotting....')
+                # %% Imports
+                import pandas as pd
+                # %% Load MonitorList from Disk
+                with self.filepath.open('rb') as f:
+                    monitor_list = pickle.load(f)
+
+                result = pd.concat(monitor_list, sort=False)
+                # result.tail()
+                prepare_plot(filepath=self.filepath, results_df=result, tag='monitor')
+                print('Plotting done.')
             self.closed = True
 
     def _on_step(self) -> bool:
@@ -99,50 +113,5 @@ class MonitorCallback(BaseCallback):
             self._monitor_list.append(self.env.monitor)
         else:
             pass
+        return True
 
-    def plot(self, **kwargs):
-        from matplotlib import pyplot as plt
-        plt.rcParams.update(kwargs)
-
-        plt.tight_layout()
-        figure = plt.gcf()
-        plt.show()
-        figure.savefig(str(self.filepath.parent / f'{self.filepath.stem}_monitor_measures.{self.ext}'), format=self.ext)
-
-    def prepare_plot(self):
-        # %% Imports
-        import pandas as pd
-        import seaborn as sns
-
-        # %% Load MonitorList from Disk
-        with self.filepath.open('rb') as f:
-            monitor_list = pickle.load(f)
-
-        result = pd.concat(monitor_list, sort=False)
-        # result.tail()
-
-        # %%
-        lineplot = sns.lineplot(data=result)
-        lineplot.title.title = f'Lineplot Summary of {len(monitor_list)} Episodes'
-
-        # %%
-        sns.set_theme(palette='husl', style='whitegrid')
-        font_size = 16
-        tex_fonts = {
-            # Use LaTeX to write all text
-            "text.usetex": True,
-            "font.family": "serif",
-            # Use 10pt font in plots, to match 10pt font in document
-            "axes.labelsize": font_size,
-            "font.size": font_size,
-            # Make the legend/label fonts a little smaller
-            "legend.fontsize": font_size - 2,
-            "xtick.labelsize": font_size - 2,
-            "ytick.labelsize": font_size - 2
-        }
-
-        try:
-            self.plot(**tex_fonts)
-        except FileNotFoundError:
-            tex_fonts['text.usetex'] = False
-            self.plot(**tex_fonts)