plot_script.py

import matplotlib.pyplot as plt
import matplotlib.colors as colors

import numpy as np
import sys
import glob

colors_map = {
    'semi-commnet': '#dc8d6d',
    'commnet': '#5785c1',
    'semi-mlp': '#78b38a',
    'mlp': '#ba71af'
}

def read_file(vec, file_name, scalar, term):
    print(file_name)
    with open(file_name, 'r') as f:
        lines = f.readlines()
        if len(lines) < 2:
            return vec

        mean_reward = False
        for idx, line in enumerate(lines):
            if term not in line:
                continue
            epoch_idx = idx
            epoch_line = line
            while 'Epoch' not in epoch_line:
                epoch_idx -= 1
                epoch_line = lines[epoch_idx]

            epoch = int(epoch_line.split(' ')[1].split('\t')[0])
            if not scalar:
                floats = line.split('\t')[1]
                left_bracket = floats.find('[')
                right_bracket = floats.find(']')
                floats = np.fromstring(floats[left_bracket + 2:right_bracket], dtype=float, sep=' ')

                if epoch > len(vec):
                    vec.append([floats.mean()])
                else:
                    vec[epoch - 1].append(floats.mean())
            else:
                floats = line.split('\t')[0]
                if epoch > len(vec):
                    vec.append([float(floats.split(' ')[-1].strip())])
                else:
                    vec[epoch - 1].append(float(floats.split(' ')[-1].strip()))

                # if term == 'Mean':
                #     vec.append(float(floats.split(' ')[2]))
                # elif term == 'Success':
                #     vec.append(float(floats.split(' ')[1].strip()))
                # else:
                #     raise RuntimeError("Unkown term used as line parser.")

    return vec

number = 1

def parse_plot(files, scalar=False, term='Epoch'):
    global number
    fig = plt.subplot(6, 1, number)
    number += 2
    coll = dict()
    for fname in files:
        f = fname.split('.')
        if 'semi' in fname and 'commnet' in fname:
            label = 'semi-commnet'
        elif 'commnet' in fname:
            label = 'commnet'
        elif 'semi' in fname and 'mlp' in fname:
            label = 'semi-mlp'
        else:
            label = 'mlp'

        if label not in coll:
            coll[label] = []

        # if len(f) == 2:
        #     label = f[0]
        # else:
        #     label = '-'.join(f[0].split('_')[1:])
        coll[label] = read_file(coll[label], fname, scalar, term)

    for label in coll.keys():
        coll[label] = coll[label][:1000]

        mean_values = []
        max_values = []
        min_values = []

        for val in coll[label]:
            mean = sum(val) / len(val)

            if term == 'Success':
                mean *= 100
            mean_values.append(mean)
            variance = np.var(val)

            if term == 'Success':
                variance *= 100
            variance = variance if variance < 20 else 20
            max_values.append(mean + variance)
            min_values.append(mean - variance)

        fig.plot(np.arange(len(coll[label])), mean_values, linewidth=1.5, label=label, color=colors_map[label])
        fig.fill_between(np.arange(len(coll[label])), min_values, max_values, color=colors.to_rgba(colors_map[label], alpha=0.2))

    term = 'Rewards' if term == 'Epoch' else term

    fig.set_xlabel('Epochs')
    fig.set_ylabel(term)
    fig.legend()
    fig.grid()
    fig.set_title('StarCraft {} {}'.format(sys.argv[2], term))

files = glob.glob(sys.argv[1] + "*")
files = list(filter(lambda x: x.find(".pt") == -1, files))

# rewards
parse_plot(files, False, 'Epoch')

# success
parse_plot(files, True, 'Success')

# steps
parse_plot(files, True, 'Steps taken')
plt.show()