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.

baselines/results_plotter.py

@@ -10,6 +10,8 @@
 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__':