Add baselines' plotter

component/__init__.py

@@ -2,4 +2,5 @@ from .atari_wrapper import *
 from .policy import *
 from .replay import *
 from .task import *
-from .random_process import *
+from .random_process import *
+from .bench import *

component/bench.py

@@ -0,0 +1,138 @@
+# from https://raw.githubusercontent.com/openai/baselines/master/baselines/bench/monitor.py
+
+__all__ = ['Monitor', 'get_monitor_files', 'load_results']
+
+import gym
+from gym.core import Wrapper
+import time
+from glob import glob
+import csv
+import os.path as osp
+import json
+
+class Monitor(Wrapper):
+    EXT = "monitor.csv"
+    f = None
+
+    def __init__(self, env, filename, allow_early_resets=True, reset_keywords=(), info_keywords=()):
+        Wrapper.__init__(self, env=env)
+        self.tstart = time.time()
+        if filename is None:
+            self.f = None
+            self.logger = None
+        else:
+            if not filename.endswith(Monitor.EXT):
+                if osp.isdir(filename):
+                    filename = osp.join(filename, Monitor.EXT)
+                else:
+                    filename = filename + "." + Monitor.EXT
+            self.f = open(filename, "wt")
+            self.f.write('#%s\n'%json.dumps({"t_start": self.tstart, 'env_id' : env.spec and env.spec.id}))
+            self.logger = csv.DictWriter(self.f, fieldnames=('r', 'l', 't')+reset_keywords+info_keywords)
+            self.logger.writeheader()
+            self.f.flush()
+
+        self.reset_keywords = reset_keywords
+        self.info_keywords = info_keywords
+        self.allow_early_resets = allow_early_resets
+        self.rewards = None
+        self.needs_reset = True
+        self.episode_rewards = []
+        self.episode_lengths = []
+        self.episode_times = []
+        self.total_steps = 0
+        self.current_reset_info = {} # extra info about the current episode, that was passed in during reset()
+
+    def reset(self, **kwargs):
+        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, wrap your env with Monitor(env, path, allow_early_resets=True)")
+        self.rewards = []
+        self.needs_reset = False
+        for k in self.reset_keywords:
+            v = kwargs.get(k)
+            if v is None:
+                raise ValueError('Expected you to pass kwarg %s into reset'%k)
+            self.current_reset_info[k] = v
+        return self.env.reset(**kwargs)
+
+    def step(self, action):
+        if self.needs_reset:
+            raise RuntimeError("Tried to step environment that needs reset")
+        ob, rew, done, info = self.env.step(action)
+        self.rewards.append(rew)
+        if done:
+            self.needs_reset = True
+            eprew = sum(self.rewards)
+            eplen = len(self.rewards)
+            epinfo = {"r": round(eprew, 6), "l": eplen, "t": round(time.time() - self.tstart, 6)}
+            for k in self.info_keywords:
+                epinfo[k] = info[k]
+            self.episode_rewards.append(eprew)
+            self.episode_lengths.append(eplen)
+            self.episode_times.append(time.time() - self.tstart)
+            epinfo.update(self.current_reset_info)
+            if self.logger:
+                self.logger.writerow(epinfo)
+                self.f.flush()
+            info['episode'] = epinfo
+        self.total_steps += 1
+        return (ob, rew, done, info)
+
+    def close(self):
+        if self.f is not None:
+            self.f.close()
+
+    def get_total_steps(self):
+        return self.total_steps
+
+    def get_episode_rewards(self):
+        return self.episode_rewards
+
+    def get_episode_lengths(self):
+        return self.episode_lengths
+
+    def get_episode_times(self):
+        return self.episode_times
+
+class LoadMonitorResultsError(Exception):
+    pass
+
+def get_monitor_files(dir):
+    return glob(osp.join(dir, "*" + Monitor.EXT))
+
+def load_results(dir):
+    import pandas
+    monitor_files = (
+        glob(osp.join(dir, "*monitor.json")) +
+        glob(osp.join(dir, "*monitor.csv"))) # get both csv and (old) json files
+    if not monitor_files:
+        raise LoadMonitorResultsError("no monitor files of the form *%s found in %s" % (Monitor.EXT, dir))
+    dfs = []
+    headers = []
+    for fname in monitor_files:
+        with open(fname, 'rt') as fh:
+            if fname.endswith('csv'):
+                firstline = fh.readline()
+                assert firstline[0] == '#'
+                header = json.loads(firstline[1:])
+                df = pandas.read_csv(fh, index_col=None)
+                headers.append(header)
+            elif fname.endswith('json'): # Deprecated json format
+                episodes = []
+                lines = fh.readlines()
+                header = json.loads(lines[0])
+                headers.append(header)
+                for line in lines[1:]:
+                    episode = json.loads(line)
+                    episodes.append(episode)
+                df = pandas.DataFrame(episodes)
+            else:
+                assert 0, 'unreachable'
+            df['t'] += header['t_start']
+        dfs.append(df)
+    df = pandas.concat(dfs)
+    df.sort_values('t', inplace=True)
+    df.reset_index(inplace=True)
+    df['t'] -= min(header['t_start'] for header in headers)
+    # df.headers = headers # HACK to preserve backwards compatibility
+    return df

component/task.py

@@ -9,6 +9,8 @@ import numpy as np
 from .atari_wrapper import *
 import multiprocessing as mp
 import sys
+from .bench import Monitor
+from utils import *
 
 class BasicTask:
     def __init__(self, max_steps=sys.maxsize):
@@ -32,6 +34,9 @@ class BasicTask:
     def random_action(self):
         return self.env.action_space.sample()
 
+    def set_monitor(self, filename):
+        self.env = Monitor(self.env, filename)
+
 class ClassicalControl(BasicTask):
     def __init__(self, name='CartPole-v0', max_steps=200):
         BasicTask.__init__(self, max_steps)
@@ -138,9 +143,11 @@ class Roboschool(BasicTask):
     def step(self, action):
         return BasicTask.step(self, np.clip(action, -1, 1))
 
-def sub_task(parent_pipe, pipe, task_fn):
+def sub_task(parent_pipe, pipe, task_fn, filename=None):
     parent_pipe.close()
     task = task_fn()
+    if filename is not None:
+        task.set_monitor(filename)
     task.env.seed(np.random.randint(0, sys.maxsize))
     while True:
         op, data = pipe.recv()
@@ -155,12 +162,17 @@ def sub_task(parent_pipe, pipe, task_fn):
             assert False, 'Unknown Operation'
 
 class ParallelizedTask:
-    def __init__(self, task_fn, num_workers):
+    def __init__(self, task_fn, num_workers, tag='vanilla'):
         self.task_fn = task_fn
         self.task = task_fn()
         self.name = self.task.name
+        # date = datetime.datetime.now().strftime("%I:%M%p-on-%B-%d-%Y")
+        mkdir('./log/%s-%s' % (self.name, tag))
+        filenames = ['./log/%s-%s/worker-%d' % (self.name, tag, i)
+                     for i in range(num_workers)]
         self.pipes, worker_pipes = zip(*[mp.Pipe() for _ in range(num_workers)])
-        args = [(p, wp, task_fn) for p, wp in zip(self.pipes, worker_pipes)]
+        args = [(p, wp, task_fn, filename)
+                for p, wp, filename in zip(self.pipes, worker_pipes, filenames)]
         self.workers = [mp.Process(target=sub_task, args=arg) for arg in args]
         for p in self.workers: p.start()
         for p in worker_pipes: p.close()