diff --git a/setup.py b/setup.py index e598108..e7602bf 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,9 @@ setup(name='torchy_baselines', 'gym[classic_control]>=0.10.9', 'numpy', 'torch>=1.2.0', - 'cloudpickle' + 'cloudpickle', + # For reading logs + 'pandas' ], extras_require={ 'tests': [ @@ -32,8 +34,6 @@ setup(name='torchy_baselines', 'extra': [ # For render 'opencv-python', - # For reading logs - 'pandas' ] }, description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.', diff --git a/torchy_baselines/common/monitor.py b/torchy_baselines/common/monitor.py index 53bea70..8d460d2 100644 --- a/torchy_baselines/common/monitor.py +++ b/torchy_baselines/common/monitor.py @@ -1,29 +1,37 @@ -""" -Taken from stable-baselines -""" +__all__ = ['Monitor', 'get_monitor_files', 'load_results'] + import csv import json import os import time +from glob import glob +from typing import Tuple, Dict, Any, List, Optional -from gym.core import Wrapper +import gym +import pandas +import numpy as np -class Monitor(Wrapper): +class Monitor(gym.Wrapper): EXT = "monitor.csv" file_handler = None - def __init__(self, env, filename=None, allow_early_resets=True, reset_keywords=(), info_keywords=()): + def __init__(self, + env: gym.Env, + filename: Optional[str], + allow_early_resets: bool = True, + reset_keywords=(), + info_keywords=()): """ A monitor wrapper for Gym environments, it is used to know the episode reward, length, time and other data. - :param env: (Gym environment) The environment - :param filename: (str) the location to save a log file, can be None for no log + :param env: (gym.Env) The environment + :param filename: (Optional[str]) the location to save a log file, can be None for no log :param allow_early_resets: (bool) allows the reset of the environment before it is done :param reset_keywords: (tuple) extra keywords for the reset call, if extra parameters are needed at reset :param info_keywords: (tuple) extra information to log, from the information return of environment.step """ - Wrapper.__init__(self, env=env) + super(Monitor, self).__init__(env=env) self.t_start = time.time() if filename is None: self.file_handler = None @@ -52,12 +60,12 @@ class Monitor(Wrapper): self.total_steps = 0 self.current_reset_info = {} # extra info about the current episode, that was passed in during reset() - def reset(self, **kwargs): + def reset(self, **kwargs) -> np.ndarray: """ Calls the Gym environment reset. Can only be called if the environment is over, or if allow_early_resets is True :param kwargs: Extra keywords saved for the next episode. only if defined by reset_keywords - :return: ([int] or [float]) the first observation of the environment + :return: (np.ndarray) the first observation of the environment """ if not self.allow_early_resets and not self.needs_reset: raise RuntimeError("Tried to reset an environment before done. If you want to allow early resets, " @@ -67,16 +75,16 @@ class Monitor(Wrapper): for key in self.reset_keywords: value = kwargs.get(key) if value is None: - raise ValueError('Expected you to pass kwarg %s into reset' % key) + raise ValueError('Expected you to pass kwarg {} into reset'.format(key)) self.current_reset_info[key] = value return self.env.reset(**kwargs) - def step(self, action): + def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, Dict[Any, Any]]: """ Step the environment with the given action - :param action: ([int] or [float]) the action - :return: ([int] or [float], [float], [bool], dict) observation, reward, done, information + :param action: (np.ndarray) the action + :return: (Tuple[np.ndarray, float, bool, Dict[Any, Any]]) observation, reward, done, information """ if self.needs_reset: raise RuntimeError("Tried to step environment that needs reset") @@ -104,10 +112,11 @@ class Monitor(Wrapper): """ Closes the environment """ + super(Monitor, self).close() if self.file_handler is not None: self.file_handler.close() - def get_total_steps(self): + def get_total_steps(self) -> int: """ Returns the total number of timesteps @@ -115,7 +124,7 @@ class Monitor(Wrapper): """ return self.total_steps - def get_episode_rewards(self): + def get_episode_rewards(self) -> List[float]: """ Returns the rewards of all the episodes @@ -123,7 +132,7 @@ class Monitor(Wrapper): """ return self.episode_rewards - def get_episode_lengths(self): + def get_episode_lengths(self) -> List[int]: """ Returns the number of timesteps of all the episodes @@ -131,10 +140,69 @@ class Monitor(Wrapper): """ return self.episode_lengths - def get_episode_times(self): + def get_episode_times(self) -> List[float]: """ Returns the runtime in seconds of all the episodes :return: ([float]) """ return self.episode_times + + +class LoadMonitorResultsError(Exception): + """ + Raised when loading the monitor log fails. + """ + pass + + +def get_monitor_files(path: str) -> List[str]: + """ + get all the monitor files in the given path + + :param path: (str) the logging folder + :return: ([str]) the log files + """ + return glob(os.path.join(path, "*" + Monitor.EXT)) + + +def load_results(path: str) -> pandas.DataFrame: + """ + Load all Monitor logs from a given directory path matching ``*monitor.csv`` and ``*monitor.json`` + + :param path: (str) the directory path containing the log file(s) + :return: (pandas.DataFrame) the logged data + """ + # get both csv and (old) json files + monitor_files = (glob(os.path.join(path, "*monitor.json")) + get_monitor_files(path)) + if not monitor_files: + raise LoadMonitorResultsError("no monitor files of the form *%s found in %s" % (Monitor.EXT, path)) + data_frames = [] + headers = [] + for file_name in monitor_files: + with open(file_name, 'rt') as file_handler: + if file_name.endswith('csv'): + first_line = file_handler.readline() + assert first_line[0] == '#' + header = json.loads(first_line[1:]) + data_frame = pandas.read_csv(file_handler, index_col=None) + headers.append(header) + elif file_name.endswith('json'): # Deprecated json format + episodes = [] + lines = file_handler.readlines() + header = json.loads(lines[0]) + headers.append(header) + for line in lines[1:]: + episode = json.loads(line) + episodes.append(episode) + data_frame = pandas.DataFrame(episodes) + else: + assert 0, 'unreachable' + data_frame['t'] += header['t_start'] + data_frames.append(data_frame) + data_frame = pandas.concat(data_frames) + data_frame.sort_values('t', inplace=True) + data_frame.reset_index(inplace=True) + data_frame['t'] -= min(header['t_start'] for header in headers) + # data_frame.headers = headers # HACK to preserve backwards compatibility + return data_frame