One graph for both average and max

graph_reward.py

 import numpy as np
 import pickle as pl
 import matplotlib.pyplot as plt
-from pathlib import Path
 import glob
 import re
-import sys
-import os
 
-if sys.argv.__contains__("max"):
-    pops = []
-    rewards = []
-    all_paths = []
-    all_numbers = []
-    best_high = []
-    for path in glob.glob("ai_training/models/*/*.pth"):
-        if not path.__contains__("cnn.pth"):
-            all_paths.append(path)
-    for i in range(len(all_paths)):
-        scores = re.findall(r"[^a-zA-Z]+\d.pth", all_paths[i])
-        pop =  all_paths[i].split("pop")[1].split("-")[0]
-        if len(scores) == 1:
-            score = ''.join(scores[0])
-            all_numbers.append((int(pop), float(score[:-4])))
-    for pop, number in all_numbers:
-        best_high.append((pop,  number))
 
-    best_high = sorted(best_high, reverse=False, key=lambda x: x[0])
+popsMax = []
+rewardsMax = []
+all_pathsMax = []
+all_numbersMax = []
+best_highMax = []
+for path in glob.glob("ai_training/new_4-1_on_1-1_stuck/*/*.pth"):
+    if not path.__contains__("cnn.pth"):
+        all_pathsMax.append(path)
+for i in range(len(all_pathsMax)):
+    scores = re.findall(r"[^a-zA-Z]+\d.pth", all_pathsMax[i])
+    pop =  all_pathsMax[i].split("pop")[1].split("-")[0]
+    if len(scores) == 1:
+        score = ''.join(scores[0])
+        all_numbersMax.append((int(pop), float(score[:-4])))
+for pop, number in all_numbersMax:
+    best_highMax.append((pop,  number))
 
-    for pop, reward in best_high:
-        pops.append(pop)
-        rewards.append(reward)
-    plt.xlim(0, max(pops))
-    plt.ylim(0, max(rewards))
-    # naming the x axis
-    plt.xlabel('Population')
-    # naming the y axis
-    plt.ylabel('Max Reward')
-    plt.plot(pops, rewards, scalex=False, scaley=False)
+best_high = sorted(best_highMax, reverse=False, key=lambda x: x[0])
 
+for pop, reward in best_high:
+    popsMax.append(pop)
+    rewardsMax.append(reward)
+plt.xlim(0, max(popsMax))
+plt.ylim(0, max(rewardsMax))
+plt.plot(popsMax, rewardsMax, 'b', scalex=False, scaley=False)
 
-    # giving a title to my graph
-    plt.title('The relationship between population and max rewards')
-    plt.show()
 
-if sys.argv.__contains__("average"):
-    pops = []
-    rewards = []
-    all_paths = []
-    data = []
-    for path in glob.glob("ai_training/models/*/*.pickle"):
-        all_paths.append(path)
-    for i in range(len(all_paths)):
-        infile = open(all_paths[i], 'rb')
-        R = pl.load(infile)
-        infile.close()
-        avg_reward = (sum(R) / len(R))
-        pop =  all_paths[i].split("pop")[1].split("-")[0]
-        data.append((int(pop), avg_reward))
-    
-    data = sorted(data, reverse=False, key=lambda x: x[0])
-
-    for pop, reward in data:
-        pops.append(pop)
-        rewards.append(reward)
-    plt.xlim(0, len(pops))
-    plt.ylim(0, max(rewards))
-    # naming the x axis
-    plt.xlabel('Population')
-    # naming the y axis
-    plt.ylabel('Average rewards')
-    plt.plot(pops, rewards, scalex=False, scaley=False)
-
-
-    # giving a title to my graph
-    plt.title('The relationship between population and average rewards')
-    plt.show()
 
+#if sys.argv.__contains__("average"):
+pops = []
+rewards = []
+all_paths = []
+data = []
+for path in glob.glob("ai_training/new_4-1_on_1-1_stuck/*/*.pickle"):
+    all_paths.append(path)
+for i in range(len(all_paths)):
+    infile = open(all_paths[i], 'rb')
+    R = pl.load(infile)
+    infile.close()
+    avg_reward = (sum(R) / len(R))
+    pop =  all_paths[i].split("pop")[1].split("-")[0]
+    data.append((int(pop), avg_reward))
 
+data = sorted(data, reverse=False, key=lambda x: x[0])
 
+for pop, reward in data:
+    pops.append(pop)
+    rewards.append(reward)
+plt.xlim(0, len(pops))
+plt.ylim(0, max(rewards))
+# naming the x axis
+plt.xlabel('Population')
+# naming the y axis
+plt.ylabel('Average rewards')
+plt.plot(pops, rewards, 'r', scalex=False, scaley=False)
 
 
+# giving a title to my graph
+plt.title('The relationship between population and average rewards')
+plt.figtext(.83, .99, "Red is average\nBlue is max", verticalalignment='top',  bbox=dict(facecolor='yellow', alpha=0.1))
+plt.show()