Plots for peg insertion project

plots.py

import seaborn as sns; sns.set()
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

def plot(run_IDs, run_labels, factors, titles, 
         colors=['steelblue', 'orange', 'mediumseagreen', 'red', 'purple'],
         ecolors=['lightblue', 'navajowhite', 'lightgreen', 'pink', 'plum'],
         dpi=150):
    csv_names = ['run_%d/progress.csv' % i for i in run_IDs]
    dfs = [pd.read_csv(f) for f in csv_names]
    
    for f_idx, factor in enumerate(factors):
        factor_dfdata = []
        for i, df in enumerate(dfs):
            dfdict = {
                'iter': df['Iteration'],
                'avg': df['%sAverage' % factor],
                'std': df['%sStd' % factor]
            }
            factor_dfdata.append((i, dfdict))
        
        plt.figure(dpi=dpi)
        plt.title(titles[f_idx])
        plt.legend()
        for j, data in factor_dfdata:
            ax = sns.lineplot(x=data['iter'], y=data['avg'], label=run_labels[j], color=colors[j])
            
            # Clip error for length factor
            yerr_upper = data['avg']+data['std']
            if factor == 'Length':
                yerr_upper = yerr_upper.clip(upper=2000)
            
            ax.fill_between(data['iter'], data['avg']-data['std'], yerr_upper, color=ecolors[j], alpha=0.8)
        plt.tight_layout()
        
plot([37, 38], ['1 F/T reading', '8 F/T readings'], ['DiscountedReturn'], ['Avg. Discounted Return with varying num. of F/T readings, 0.03 noise'])