Extend graph plotting to multiple seeds

Author: hedrox

aggregate_runs.py

@@ -1,25 +1,28 @@
-import numpy as np
-import pandas as pd
+import sys
 import argparse
+import pandas as pd
 
 from tqdm import tqdm
 
 
-parser = argparse.ArgumentParser()
-parser.add_argument("model1_path", type=str, help="Path to the progress.csv of the first")
-parser.add_argument("model2_path", type=str, help="Path to the progress.csv of the second")
+parser = argparse.ArgumentParser(description="Aggregate multiple runs of a model")
+parser.add_argument("path_list", type=str, nargs="+", default=[],
+                    help="Paths to the progress.csv files")
 args = parser.parse_args()
 
-data_csv1 = pd.read_csv(args.model1_path)
-data_csv2 = pd.read_csv(args.model2_path)
 
-res = data_csv1.copy()
+if len(args.path_list) < 2:
+    print("You need to give at least two paths to progress.csv files,
+          so that we have something to aggregate")
+    sys.exit(1)
 
-base_reward = data_csv1.rewtotal.values[-1]
-base_tcount = data_csv1.tcount.values[-1]
-visited_rooms = set()
-highest_reward = data_csv1.best_ret.values[-1]
+data_csv = []
+for csv_path in args.path_list:
+    data_csv.append(pd.read_csv(csv_path))
 
+res = data_csv[0].copy()
+
+visited_rooms = set()
 def format_rooms(rooms):
     # Input rooms are strings: '[1;2;3]'
     rooms = rooms.strip("][")
@@ -28,31 +31,34 @@ def format_rooms(rooms):
         return rooms
     return []
 
-rooms = format_rooms(data_csv1.rooms.values[-1])
+rooms = format_rooms(data_csv[0].rooms.values[-1])
 for r in rooms:
     visited_rooms.add(r)
 
-for row in tqdm(data_csv2.values):
-    reward = row[43]
-    if pd.isna(reward):
-        continue
-
-    tcount = row[32]
-    rooms = format_rooms(row[31])
-    for r in rooms:
-        visited_rooms.add(r)
-    highest_reward = max(row[2], highest_reward)
-
-    total_reward = base_reward + reward
-    total_tcount = base_tcount + tcount
-    frame = list(row)
-    frame[2] = highest_reward
-    frame[8] = len(visited_rooms)    
-    frame[43] = total_reward
-    frame[32] = total_tcount
-    df = pd.DataFrame([frame], columns=list(data_csv1.columns))
-    res = res.append(df, ignore_index=True)
+for data_cv in tqdm(data_csv[1:], desc="files"):
+    base_reward = res.rewtotal.values[-1]
+    base_tcount = res.tcount.values[-1]
+    highest_reward = res.best_ret.values[-1]
 
-res.to_csv("aggregated_progress.csv")
+    for row in tqdm(data_cv.values, desc="rows", leave=False):
+        reward = row[43]
+        if pd.isna(reward):
+            continue
+
+        tcount = row[32]
+        rooms = format_rooms(row[31])
+        for r in rooms:
+            visited_rooms.add(r)
+        highest_reward = max(row[2], highest_reward)
 
-    
+        total_reward = base_reward + reward
+        total_tcount = base_tcount + tcount
+        frame = list(row)
+        frame[2] = highest_reward
+        frame[8] = len(visited_rooms)
+        frame[43] = total_reward
+        frame[32] = total_tcount
+        df = pd.DataFrame([frame], columns=list(data_csv[0].columns))
+        res = res.append(df, ignore_index=True)
+
+res.to_csv("aggregated_progress.csv")

plot_graphs.py

@@ -1,74 +1,131 @@
+import os
+import sys
+import argparse
 import numpy as np
 import pandas as pd
-import os, sys, argparse
 
 import matplotlib.pyplot as plt
 
+
 parser = argparse.ArgumentParser()
-parser.add_argument("--rnd_path", type=str, help="Path to the progress.csv of the RND run")
-parser.add_argument("--aarnd_path", type=str,
-                    help="Path to the progress.csv of the AA RND run")
-parser.add_argument("--egornd_path", type=str,
-                    help="Path to the progress.csv of the Ego RND run")
+parser.add_argument("--base_path", type=str,
+                    help="The common directory path to all runs")
+parser.add_argument("--rnd_paths", type=str, nargs="+", default=[],
+                    help="Paths to the RND directory where the progress.csv file exists")
+parser.add_argument("--aarnd_paths", type=str, nargs="+", default=[],
+                    help="Paths to the AA RND directory where the progress.csv file exists")
+parser.add_argument("--egornd_paths", type=str, nargs="+", default=[],
+                    help="Paths to the Ego RND directory where the progress.csv file exists")
 args = parser.parse_args()
 
-data1 = pd.read_csv(args.rnd_path)
-data2 = pd.read_csv(args.aarnd_path)
-data3 = pd.read_csv(args.egornd_path)
-
-data = [data1, data2, data3]
-for idx in range(len(data)):
-    twomillion = 200000000
-    while True:
-        res = data[idx]['tcount'][data[idx]['tcount'] == twomillion]
-        if not res.empty:
-            data[idx] = data[idx][:res.index[0]]
-            break
-        else:
-            twomillion += 1
-
-data1, data2, data3 = data
 
-#fig, axes = plt.subplots(figsize=(19.20, 10.80), nrows=2, ncols=2)
+def read_from_csv(paths):
+    result = []
+    for path in paths:
+        expanded_path = os.path.join(args.base_path, path, "progress.csv")
+        if not os.path.exists(expanded_path):
+            raise Exception("Path: {} does not exist".format(expanded_path))
+        result.append(pd.read_csv(expanded_path))
+    return result
+
+
+def clean_data(data):
+    for idx in range(len(data)):
+        if len(data[idx].columns) == 45:
+            data[idx] = data[idx].drop(columns=['Unnamed: 0'])
+
+    for idx in range(len(data)):
+        twomillion = 200000000
+        while True:
+            res = data[idx]['tcount'][data[idx]['tcount'] == twomillion]
+            if not res.empty:
+                data[idx] = data[idx][:res.index[0]]
+                break
+            else:
+                twomillion += 1
+    return data
+
+
+def equalize_rows(data):
+    for dta in data:
+        min_len = len(dta[0])
+        for idx in range(len(dta)):
+            if len(dta[idx]) < min_len:
+                min_len = len(dta[idx])
+        for idx in range(len(dta)):
+            dta[idx] = dta[idx][:min_len]
+    return data
+
+
+def calculate_mu_sigma(data, column):
+    all_data = []
+    for dta in data:
+        dta_val = dta[column].values
+        dta_val[np.isnan(dta_val)] = 0
+        dta_val[np.isinf(dta_val)] = 0
+        all_data.append(dta_val)
+
+    total = np.stack(all_data)
+    mu = total.mean(axis=0)
+    sigma = total.std(axis=0)
+    ci = sigma
+    tcount = data[0]['tcount'].values
+    return mu, ci, tcount
+
+
+def plot_fill(data, column, ax, ylabel):
+    rnd_data = data[0]
+    aarnd_data = data[1]
+    egornd_data = data[2]
+
+    mu, ci, tcount = calculate_mu_sigma(rnd_data, column)
+    ax.plot(tcount, mu, lw=1, color='red', label='RND')
+    ax.fill_between(tcount, (mu-ci), (mu+ci), facecolor='red', alpha=0.25)
+
+    mu, ci, tcount = calculate_mu_sigma(aarnd_data, column)
+    ax.plot(tcount, mu, lw=1, color='green', label='AA RND')
+    ax.fill_between(tcount, (mu-ci), (mu+ci), facecolor='green', alpha=0.25)
+
+    mu, ci, tcount = calculate_mu_sigma(egornd_data, column)
+    ax.plot(tcount, mu, lw=1, color='blue', label='Ego RND')
+    ax.fill_between(tcount, (mu-ci), (mu+ci), facecolor='blue', alpha=0.25)
+
+    ax.set_xlabel('Frames')
+    ax.set_ylabel(ylabel)
+    ax.legend()
+
+
+if not args.base_path or not args.rnd_paths or not args.aarnd_paths or not args.egornd_paths:
+    print("Command line arguments were not set properly")
+    sys.exit(1)
+
+assert len(args.rnd_paths) == len(args.aarnd_paths) == len(args.egornd_paths), \
+    "Number of path not equal for all methods"
+
+data_rnd = np.array(read_from_csv(args.rnd_paths))
+data_aarnd = np.array(read_from_csv(args.aarnd_paths))
+data_egornd = np.array(read_from_csv(args.egornd_paths))
+
+data = np.concatenate((data_rnd, data_aarnd, data_egornd))
+data = clean_data(data)
+
+# fig, axes = plt.subplots(figsize=(19.20, 10.80), nrows=2, ncols=2)
 fig, axes = plt.subplots(nrows=2, ncols=2)
+# fig.suptitle("Montezuma's Revenge Ego vs AA-RND", fontsize=10,y=0.9,x=0.51)
 
 """
 retextmean, retextstd, retintmean, retintstd, rewintmean_norm, rewintmean_unnorm,
 vpredextmean, vpredintmean are interesting metrics
 """
 
-#fig.suptitle("Montezuma's Revenge Ego vs AA-RND", fontsize=10,y=0.9,x=0.51)
-data1.plot(x='tcount', y='rewtotal', ax=axes[0,0], color='red', label='RND')
-data2.plot(x='tcount', y='rewtotal', ax=axes[0,0], color='green', label='AA RND')
-data3.plot(x='tcount', y='rewtotal', ax=axes[0,0], color='blue', label='Ego RND')
-axes[0,0].set_xlabel('Frames')
-axes[0,0].set_ylabel('Total Rewards')
-
-
-data1.plot(x='tcount', y='n_rooms', ax=axes[0,1], color='red', label='RND')
-data2.plot(x='tcount', y='n_rooms', ax=axes[0,1], color='green', label='AA RND')
-data3.plot(x='tcount', y='n_rooms', ax=axes[0,1], color='blue', label='Ego RND')
-axes[0,1].set_xlabel('Frames')
-axes[0,1].set_ylabel('No Rooms')
+data = data.reshape(3, -1)
+data = equalize_rows(data)
 
-
-data1.plot(x='tcount', y='eprew', ax=axes[1,0], color='red', label='RND')
-data2.plot(x='tcount', y='eprew', ax=axes[1,0], color='green', label='AA RND')
-data3.plot(x='tcount', y='eprew', ax=axes[1,0], color='blue', label='Ego RND')
-axes[1,0].set_xlabel('Frames')
-axes[1,0].set_ylabel('Episodic Rewards')
-
-
-data1.plot(x='tcount', y='best_ret', ax=axes[1,1], color='red', label='RND')
-data2.plot(x='tcount', y='best_ret', ax=axes[1,1], color='green', label='AA RND')
-data3.plot(x='tcount', y='best_ret', ax=axes[1,1], color='blue', label='Ego RND')
-axes[1,1].set_xlabel('Frames')
-axes[1,1].set_ylabel('Best Rewards')
-
-#fig.show()
-#plt.show()
+plot_fill(data, 'rewtotal', axes[0, 0], 'Total Rewards')
+plot_fill(data, 'n_rooms', axes[0, 1], 'No Rooms')
+plot_fill(data, 'eprew', axes[1, 0], 'Episodic Rewards')
+plot_fill(data, 'best_ret', axes[1, 1], 'Best Rewards')
 
 plot_name = 'montezuma-all-three'
 plt.tight_layout()
-plt.savefig(f'{plot_name}.eps')
-
+plt.savefig(f'{plot_name}.png')

run_atari.py

@@ -2,6 +2,7 @@
 import functools
 import os
 
+from copy import copy
 from baselines import logger
 from mpi4py import MPI
 import mpi_util
@@ -22,8 +23,11 @@ def train(*, env_id, num_env, hps, num_timesteps, seed):
         # for the ego experiment we needed a higher intrinsic coefficient
         hps['int_coeff'] = 3.0
 
+    hyperparams = copy(hps)
+    hyperparams.update({'seed': seed})
     logger.info("Hyperparameters:")
-    logger.info(hps)
+    logger.info(hyperparams)
+
     venv = VecFrameStack(
         make_atari_env(env_id, num_env, seed, wrapper_kwargs={},
                        start_index=num_env * MPI.COMM_WORLD.Get_rank(),