mirror of
https://github.com/illiumst/marl-factory-grid.git
synced 2025-05-22 14:56:43 +02:00
74 lines
2.6 KiB
Python
74 lines
2.6 KiB
Python
import pickle
|
|
from os import PathLike
|
|
from pathlib import Path
|
|
from typing import Union
|
|
|
|
from gymnasium import Wrapper
|
|
|
|
from marl_factory_grid.utils.helpers import IGNORED_DF_COLUMNS
|
|
|
|
import pandas as pd
|
|
|
|
from marl_factory_grid.utils.plotting.plot_single_runs import plot_single_run
|
|
|
|
|
|
class EnvMonitor(Wrapper):
|
|
|
|
ext = 'png'
|
|
|
|
def __init__(self, env, filepath: Union[str, PathLike] = None):
|
|
"""
|
|
EnvMonitor is a wrapper for Gymnasium environments that monitors and logs key information during interactions.
|
|
"""
|
|
super(EnvMonitor, self).__init__(env)
|
|
self._filepath = filepath
|
|
self._monitor_df = pd.DataFrame()
|
|
self._monitor_dict = dict()
|
|
|
|
def step(self, action):
|
|
obs_type, obs, reward, done, info = self.env.step(action)
|
|
self._read_info(info)
|
|
self._read_done(done)
|
|
return obs_type, obs, reward, done, info
|
|
|
|
def reset(self):
|
|
return self.env.reset()
|
|
|
|
def _read_info(self, info: dict):
|
|
self._monitor_dict[len(self._monitor_dict)] = {
|
|
key: val for key, val in info.items() if
|
|
key not in ['terminal_observation', 'episode']}
|
|
return
|
|
|
|
def _read_done(self, done):
|
|
if done:
|
|
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 col.endswith('ount') else 'sum' for col in columns}
|
|
)
|
|
env_monitor_df['episode'] = len(self._monitor_df)
|
|
self._monitor_df = pd.concat([self._monitor_df, pd.DataFrame([env_monitor_df])], ignore_index=True)
|
|
else:
|
|
pass
|
|
return
|
|
|
|
def save_monitor(self, filepath: Union[Path, str, None] = None, auto_plotting_keys=None):
|
|
"""
|
|
Saves the monitoring data to a file and optionally generates plots.
|
|
|
|
:param filepath: The path to save the monitoring data file.
|
|
:type filepath: Union[Path, str, None]
|
|
:param auto_plotting_keys: Keys to use for automatic plot generation.
|
|
:type auto_plotting_keys: Any
|
|
"""
|
|
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)
|
|
|
|
def report_possible_colum_keys(self):
|
|
print(self._monitor_df.columns) |