Files
2018-05-30 11:37:13 -06:00

111 lines
4.1 KiB
Python

# Adapted from https://github.com/openai/baselines/blob/master/baselines/results_plotter.py
from ..component.bench import load_monitor_log
import numpy as np
import os
import re
class Plotter:
COLORS = ['blue', 'green', 'red', 'cyan', 'magenta', 'yellow', 'black', 'purple', 'pink',
'brown', 'orange', 'teal', 'coral', 'lightblue', 'lime', 'lavender', 'turquoise',
'darkgreen', 'tan', 'salmon', 'gold', 'lightpurple', 'darkred', 'darkblue']
X_TIMESTEPS = 'timesteps'
X_EPISODES = 'episodes'
X_WALLTIME = 'walltime_hrs'
def __init__(self):
pass
def rolling_window(self, a, window):
shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
strides = a.strides + (a.strides[-1],)
return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
def window_func(self, x, y, window, func):
yw = self.rolling_window(y, window)
yw_func = func(yw, axis=-1)
return x[window - 1:], yw_func
def ts2xy(self, ts, xaxis):
if xaxis == Plotter.X_TIMESTEPS:
x = np.cumsum(ts.l.values)
y = ts.r.values
elif xaxis == Plotter.X_EPISODES:
x = np.arange(len(ts))
y = ts.r.values
elif xaxis == Plotter.X_WALLTIME:
x = ts.t.values / 3600.
y = ts.r.values
else:
raise NotImplementedError
return x, y
def load_results(self, dirs, max_timesteps=1e8, x_axis=X_TIMESTEPS, episode_window=100):
tslist = []
for dir in dirs:
ts = load_monitor_log(dir)
ts = ts[ts.l.cumsum() <= max_timesteps]
tslist.append(ts)
xy_list = [self.ts2xy(ts, x_axis) for ts in tslist]
if episode_window:
xy_list = [self.window_func(x, y, episode_window, np.mean) for x, y in xy_list]
return xy_list
def load_evaluation_episodes_results(self,
dirs,
evaluation_episodes_interval,
evaluation_episodes,
max_timesteps=1e8):
raw_data = self.load_results(dirs, max_timesteps, episode_window=0)
ys = []
for x, y in raw_data:
y = np.reshape(np.asarray(y), (-1, evaluation_episodes)).mean(-1)
x = np.arange(y.shape[0]) * evaluation_episodes_interval
ys.append(y)
return x, np.stack(ys)
def average(self, xy_list, bin, max_timesteps, top_k=0):
if top_k:
perf = [np.max(y) for _, y in xy_list]
top_k_runs = np.argsort(perf)[-top_k:]
new_xy_list = []
for r, (x, y) in enumerate(xy_list):
if r in top_k_runs:
new_xy_list.append((x, y))
xy_list = new_xy_list
new_x = np.arange(0, max_timesteps, bin)
new_y = []
for x, y in xy_list:
new_y.append(np.interp(new_x, x, y))
return new_x, np.asarray(new_y)
def plot_results(self, dirs, max_timesteps=1e8, x_axis=X_TIMESTEPS, episode_window=100, title=None):
import matplotlib.pyplot as plt
plt.ticklabel_format(axis='x', style='sci', scilimits=(1, 1))
xy_list = self.load_results(dirs, max_timesteps, x_axis, episode_window)
for (i, (x, y)) in enumerate(xy_list):
color = Plotter.COLORS[i]
plt.plot(x, y, color=color)
plt.xlabel(x_axis)
plt.ylabel("Episode Rewards")
if title is not None:
plt.title(title)
def load_log_dirs(self, pattern, negative_pattern=' ', root='./log', **kwargs):
dirs = [item[0] for item in os.walk(root)]
leaf_dirs = []
for i in range(len(dirs)):
if i + 1 < len(dirs) and dirs[i + 1].startswith(dirs[i]):
continue
leaf_dirs.append(dirs[i])
names = []
p = re.compile(pattern)
np = re.compile(negative_pattern)
for dir in leaf_dirs:
if p.match(dir) and not np.match(dir):
names.append(dir)
print(dir)
return sorted(names)