81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
import pickle
|
|
from collections import defaultdict
|
|
from os import PathLike
|
|
from pathlib import Path
|
|
from typing import List, Dict, Union
|
|
|
|
from stable_baselines3.common.callbacks import BaseCallback
|
|
|
|
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, 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)
|
|
|
|
def __getattr__(self, item):
|
|
return getattr(self.unwrapped, item)
|
|
|
|
def step(self, action):
|
|
obs, reward, done, info = self.unwrapped.step(action)
|
|
self._read_info(0, info)
|
|
self._read_done(0, done)
|
|
return obs, reward, done, info
|
|
|
|
def reset(self):
|
|
return self.unwrapped.reset()
|
|
|
|
def _on_training_start(self) -> None:
|
|
pass
|
|
|
|
def _on_training_end(self) -> None:
|
|
pass
|
|
|
|
def _on_step(self, alt_infos: List[Dict] = None, alt_dones: List[bool] = None) -> bool:
|
|
for env_idx, info in enumerate(self.locals.get('infos', [])):
|
|
self._read_info(env_idx, info)
|
|
|
|
for env_idx, done in list(
|
|
enumerate(self.locals.get('dones', []))) + list(enumerate(self.locals.get('done', []))):
|
|
self._read_done(env_idx, done)
|
|
return True
|
|
|
|
def _read_info(self, env_idx, info: dict):
|
|
self._monitor_dicts[env_idx][len(self._monitor_dicts[env_idx])] = {
|
|
key: val for key, val in info.items() if
|
|
key not in ['terminal_observation', 'episode'] and not key.startswith('rec_')}
|
|
return
|
|
|
|
def _read_done(self, env_idx, done):
|
|
if done:
|
|
env_monitor_df = pd.DataFrame.from_dict(self._monitor_dicts[env_idx], orient='index')
|
|
self._monitor_dicts[env_idx] = dict()
|
|
columns = [col for col in env_monitor_df.columns if col not in IGNORED_DF_COLUMNS]
|
|
env_monitor_df = env_monitor_df.aggregate(
|
|
{col: 'mean' if col.endswith('ount') else 'sum' for col in columns}
|
|
)
|
|
env_monitor_df['episode'] = len(self._monitor_df)
|
|
self._monitor_df = self._monitor_df.append([env_monitor_df])
|
|
else:
|
|
pass
|
|
return
|
|
|
|
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)
|
|
if auto_plotting_keys:
|
|
plot_single_run(filepath, column_keys=auto_plotting_keys)
|