For plotting some A2C results from OpenAI baselines.

Gist for Plotting A2C

"""
Quick script I ran to visualuze some A2C results. Just run `python plot.py` in the same directory 
as where the stuff in `files` is located.
"""
import argparse, csv, os, pickle, sys, matplotlib
from os.path import join
matplotlib.use('Agg')
import matplotlib.pyplot as plt
plt.style.use('seaborn-darkgrid')
import numpy as np


def smoothed(x, w):
    """Smooth x by averaging over sliding windows of w.
    """
    assert len(x) > w
    smooth = []
    for i in range(1, w):
        smooth.append( np.mean(x[0:i]) )
    for i in range(w, len(x)+1):
        smooth.append( np.mean(x[i-w:i]) )
    assert len(x) == len(smooth)
    return np.array(smooth)


# For each item here: (file path, plot title, plot row index).
files = [
    ('openai-02cpu-2018-09-26-17-52-19-054836', 'A2C, Pong, env=2',  0), 
    ('openai-04cpu-2018-09-27-08-23-19-939494', 'A2C, Pong, env=4',  1),
    ('openai-08cpu-2018-09-26-21-04-18-045069', 'A2C, Pong, env=8',  2),
    ('openai-16cpu-2018-09-26-17-45-39-553564', 'A2C, Pong, env=16', 3),
]

titlesize = 35
xsize = 30
nrows, ncols = 4, 3
fig, ax = plt.subplots(nrows, ncols, squeeze=False, figsize=(11*ncols,8*nrows))


for f in files:
    print("\nCurrently on: {}".format(f))
    fname, title, row = f 
    monitors = sorted([xx for xx in os.listdir(fname) if 'monitor.csv' in xx and '.swp' not in xx])

    # Put all the monitor curves. Might need to smooth them.
    for env_idx,mon in enumerate(monitors):
        scores = []
        steps  = []
        times  = []
        csv_path = join(fname,mon)
        with open(csv_path) as csv_file:
            csv_reader = csv.reader(csv_file, delimiter=',')
            line_count = 0
            for csv_row in csv_reader:
                if line_count == 0 or line_count == 1:
                    line_count += 1
                    continue
                scores.append(float(csv_row[0]))
                steps.append(int(csv_row[1]))
                times.append(float(csv_row[2]))
                line_count += 1
        print("finished: {}".format(csv_path))
        
        # Now process to see as a function of episodes and training steps, etc.
        # In the csv file, _steps_ are NOT cumulative, but _time_ IS.
        num_episodes = len(scores)
        tr_episodes = np.arange(num_episodes)
        tr_steps = np.cumsum(steps)
        tr_times = np.array(times) / 60.0 # get it in minutes

        # Plot!! For scores, use the raw ones.
        envlabel = 'env {}'.format(env_idx)
        ax[row,0].plot(tr_episodes, scores, label=envlabel+'; {:.1f}'.format(scores[-1]))

        # Then for here, use the smoothed scores across window `w`.
        sm_10 = smoothed(scores, w=10)
        ax[row,1].plot(tr_steps, sm_10, label=envlabel+'; {:.1f}'.format(sm_10[-1]))

        # And another way of smoothing.
        sm_100 = smoothed(scores, w=100)
        ax[row,2].plot(tr_times, sm_100, label=envlabel+'; {:.1f}'.format(sm_100[-1]))

    # Bells and whistles
    for col in range(ncols):
        ax[row,col].set_ylabel("Scores", fontsize=30)
        ax[row,col].tick_params(axis='x', labelsize=25)
        ax[row,col].tick_params(axis='y', labelsize=25)
        leg = ax[row,col].legend(loc="best", ncol=2, prop={'size':25})
        for legobj in leg.legendHandles:
            legobj.set_linewidth(5.0)
    ax[row,0].set_title(title+', Raw Scores', fontsize=titlesize)
    ax[row,0].set_xlabel("Training Episodes", fontsize=xsize)
    ax[row,1].set_title(title+', Smoothed (w=10)', fontsize=titlesize)
    ax[row,1].set_xlabel("Training Steps", fontsize=xsize)
    ax[row,2].set_title(title+', Smoothed (w=100)', fontsize=titlesize)
    ax[row,2].set_xlabel("Training Time (Minutes)", fontsize=xsize)

plt.tight_layout()
figname = 'a2c_tests.png'
plt.savefig(figname)
print("\nJust saved: {}".format(figname))