Add possibility of plotting timesteps vs episodes (#578)
* Add possibility of plotting timesteps vs episodes * Remove leftover from personal project patch * Auto plt.tight_layout() on resize window event Calls `plt.tight_layout()` if a `resize_event` is issued. This means that the plot will look good even after the user has resized the plotting window.
This commit is contained in:
committed by
pzhokhov
parent
115b59d28b
commit
85be74500d
@@ -10,6 +10,8 @@ from baselines.bench.monitor import load_results
|
||||
X_TIMESTEPS = 'timesteps'
|
||||
X_EPISODES = 'episodes'
|
||||
X_WALLTIME = 'walltime_hrs'
|
||||
Y_REWARD = 'reward'
|
||||
Y_TIMESTEPS = 'timesteps'
|
||||
POSSIBLE_X_AXES = [X_TIMESTEPS, X_EPISODES, X_WALLTIME]
|
||||
EPISODES_WINDOW = 100
|
||||
COLORS = ['blue', 'green', 'red', 'cyan', 'magenta', 'yellow', 'black', 'purple', 'pink',
|
||||
@@ -26,22 +28,25 @@ def window_func(x, y, window, func):
|
||||
yw_func = func(yw, axis=-1)
|
||||
return x[window-1:], yw_func
|
||||
|
||||
def ts2xy(ts, xaxis):
|
||||
def ts2xy(ts, xaxis, yaxis):
|
||||
if xaxis == X_TIMESTEPS:
|
||||
x = np.cumsum(ts.l.values)
|
||||
y = ts.r.values
|
||||
elif xaxis == X_EPISODES:
|
||||
x = np.arange(len(ts))
|
||||
y = ts.r.values
|
||||
elif xaxis == X_WALLTIME:
|
||||
x = ts.t.values / 3600.
|
||||
else:
|
||||
raise NotImplementedError
|
||||
if yaxis == Y_REWARD:
|
||||
y = ts.r.values
|
||||
elif yaxis == Y_TIMESTEPS:
|
||||
y = ts.l.values
|
||||
else:
|
||||
raise NotImplementedError
|
||||
return x, y
|
||||
|
||||
def plot_curves(xy_list, xaxis, title):
|
||||
plt.figure(figsize=(8,2))
|
||||
def plot_curves(xy_list, xaxis, yaxis, title):
|
||||
fig = plt.figure(figsize=(8,2))
|
||||
maxx = max(xy[0][-1] for xy in xy_list)
|
||||
minx = 0
|
||||
for (i, (x, y)) in enumerate(xy_list):
|
||||
@@ -52,17 +57,19 @@ def plot_curves(xy_list, xaxis, title):
|
||||
plt.xlim(minx, maxx)
|
||||
plt.title(title)
|
||||
plt.xlabel(xaxis)
|
||||
plt.ylabel("Episode Rewards")
|
||||
plt.ylabel(yaxis)
|
||||
plt.tight_layout()
|
||||
fig.canvas.mpl_connect('resize_event', lambda event: plt.tight_layout())
|
||||
plt.grid(True)
|
||||
|
||||
def plot_results(dirs, num_timesteps, xaxis, task_name):
|
||||
def plot_results(dirs, num_timesteps, xaxis, yaxis, task_name):
|
||||
tslist = []
|
||||
for dir in dirs:
|
||||
ts = load_results(dir)
|
||||
ts = ts[ts.l.cumsum() <= num_timesteps]
|
||||
tslist.append(ts)
|
||||
xy_list = [ts2xy(ts, xaxis) for ts in tslist]
|
||||
plot_curves(xy_list, xaxis, task_name)
|
||||
xy_list = [ts2xy(ts, xaxis, yaxis) for ts in tslist]
|
||||
plot_curves(xy_list, xaxis, yaxis, task_name)
|
||||
|
||||
# Example usage in jupyter-notebook
|
||||
# from baselines import log_viewer
|
||||
@@ -77,10 +84,11 @@ def main():
|
||||
parser.add_argument('--dirs', help='List of log directories', nargs = '*', default=['./log'])
|
||||
parser.add_argument('--num_timesteps', type=int, default=int(10e6))
|
||||
parser.add_argument('--xaxis', help = 'Varible on X-axis', default = X_TIMESTEPS)
|
||||
parser.add_argument('--yaxis', help = 'Varible on Y-axis', default = Y_REWARD)
|
||||
parser.add_argument('--task_name', help = 'Title of plot', default = 'Breakout')
|
||||
args = parser.parse_args()
|
||||
args.dirs = [os.path.abspath(dir) for dir in args.dirs]
|
||||
plot_results(args.dirs, args.num_timesteps, args.xaxis, args.task_name)
|
||||
plot_results(args.dirs, args.num_timesteps, args.xaxis, args.yaxis, args.task_name)
|
||||
plt.show()
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
Reference in New Issue
Block a user