Restructuring and Testing Done

This commit is contained in:
steffen-illium
2021-07-13 11:12:03 +02:00
parent eee4760e72
commit 35f5bdeed4
14 changed files with 1160 additions and 842 deletions

View File

@ -1,5 +1,6 @@
import pickle
from pathlib import Path
from typing import List, Dict
from stable_baselines3.common.callbacks import BaseCallback
@ -66,13 +67,15 @@ class MonitorCallback(BaseCallback):
print('Plotting done.')
self.closed = True
def _on_step(self) -> bool:
for _, info in enumerate(self.locals.get('infos', [])):
def _on_step(self, alt_infos: List[Dict] = None, alt_dones: List[bool] = None) -> bool:
infos = alt_infos or self.locals.get('infos', [])
dones = alt_dones or self.locals.get('dones', None) or self.locals.get('done', [None])
for _, info in enumerate(infos):
self._monitor_dict[self.num_timesteps] = {key: val for key, val in info.items()
if key not in ['terminal_observation', 'episode']}
if key not in ['terminal_observation', 'episode']
and not key.startswith('rec_')}
for env_idx, done in list(enumerate(self.locals.get('dones', []))) + \
list(enumerate(self.locals.get('done', []))):
for env_idx, done in enumerate(dones):
if done:
env_monitor_df = pd.DataFrame.from_dict(self._monitor_dict, orient='index')
self._monitor_dict = dict()