pomdp
This commit is contained in:
@ -1,6 +1,5 @@
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
from collections import defaultdict
|
||||
|
||||
from stable_baselines3.common.callbacks import BaseCallback
|
||||
|
||||
@ -9,51 +8,6 @@ from environments.logging.plotting import prepare_plot
|
||||
import pandas as pd
|
||||
|
||||
|
||||
class FactoryMonitor:
|
||||
|
||||
def __init__(self, env):
|
||||
self._env = env
|
||||
self._monitor = defaultdict(lambda: defaultdict(lambda: 0))
|
||||
self._last_vals = defaultdict(lambda: 0)
|
||||
|
||||
def __iter__(self):
|
||||
for key, value in self._monitor.items():
|
||||
yield key, dict(value)
|
||||
|
||||
def add(self, key, value, step=None):
|
||||
assert step is None or step >= 1 # Is this good practice?
|
||||
step = step or self._env.steps
|
||||
self._last_vals[key] = self._last_vals[key] + value
|
||||
self._monitor[key][step] = self._last_vals[key]
|
||||
return self._last_vals[key]
|
||||
|
||||
def set(self, key, value, step=None):
|
||||
assert step is None or step >= 1 # Is this good practice?
|
||||
step = step or self._env.steps
|
||||
self._last_vals[key] = value
|
||||
self._monitor[key][step] = self._last_vals[key]
|
||||
return self._last_vals[key]
|
||||
|
||||
def remove(self, key, value, step=None):
|
||||
assert step is None or step >= 1 # Is this good practice?
|
||||
step = step or self._env.steps
|
||||
self._last_vals[key] = self._last_vals[key] - value
|
||||
self._monitor[key][step] = self._last_vals[key]
|
||||
return self._last_vals[key]
|
||||
|
||||
def to_dict(self):
|
||||
return dict(self)
|
||||
|
||||
def to_pd_dataframe(self):
|
||||
import pandas as pd
|
||||
df = pd.DataFrame.from_dict(self.to_dict())
|
||||
df.fillna(0)
|
||||
return df
|
||||
|
||||
def reset(self):
|
||||
raise RuntimeError("DO NOT DO THIS! Always initalize a new Monitor per Env-Run.")
|
||||
|
||||
|
||||
class MonitorCallback(BaseCallback):
|
||||
|
||||
ext = 'png'
|
||||
@ -62,6 +16,7 @@ class MonitorCallback(BaseCallback):
|
||||
super(MonitorCallback, self).__init__()
|
||||
self.filepath = Path(filepath)
|
||||
self._monitor_df = pd.DataFrame()
|
||||
self._monitor_dict = dict()
|
||||
self.env = env
|
||||
self.plotting = plotting
|
||||
self.started = False
|
||||
@ -113,12 +68,17 @@ class MonitorCallback(BaseCallback):
|
||||
self.closed = True
|
||||
|
||||
def _on_step(self) -> bool:
|
||||
for _, info in enumerate(self.locals.get('infos', [])):
|
||||
self._monitor_dict[self.num_timesteps] = {key: val for key, val in info.items()
|
||||
if key not in ['terminal_observation', 'episode']}
|
||||
|
||||
for env_idx, done in enumerate(self.locals.get('dones', [])):
|
||||
if done:
|
||||
env_monitor_df = self.locals['infos'][env_idx]['monitor'].to_pd_dataframe()
|
||||
env_monitor_df = pd.DataFrame.from_dict(self._monitor_dict, orient='index')
|
||||
self._monitor_dict = 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 'amount' in col or 'count' in col else 'sum' for col in columns}
|
||||
{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])
|
||||
|
Reference in New Issue
Block a user