Update plot utility

main.py

@@ -338,23 +338,16 @@ def ddpg_continuous():
 def plot():
     import matplotlib.pyplot as plt
     plotter = Plotter()
-    # name = 'log/ppo_continuous-180408-002056'
-    # plotter.plot_results([name])
-    # plt.show()
-    names = [
-            # 'a2c_pixel_atari-180407-92711',
-            # 'categorical_dqn_pixel_atari-180407-094006',
-            # 'dqn_pixel_atari-180407-01414',
-            # 'quantile_regression_dqn_pixel_atari-180407-01604',
-            #  'n_step_dqn_pixel_atari-180408-001104',
-            #  'ppo_continuous-180408-002056',
-            #  'ddpg_continuous-180407-234141'
-            'ppo_pixel_atari-180410-235529',
-             ]
-    for name in names:
-        plotter.plot_results(['to_plot/%s' % (name)], title='BreakoutNoFrameskip-v4')
-        plt.savefig('images/%s.png' % (name))
-        plt.close()
+    names = plotter.load_log_dirs('')
+    data = plotter.load_results(names)
+
+    for i, name in enumerate(names):
+        x, y = data[i]
+        plt.plot(x, y, color=Plotter.COLORS[i], label=name)
+    plt.legend()
+    plt.xlabel('timesteps')
+    plt.ylabel('episode return')
+    plt.show()
 
 def action_conditional_video_prediction():
     game = 'PongNoFrameskip-v4'

utils/plot.py

@@ -2,6 +2,8 @@
 
 import numpy as np
 import component
+import os
+import re
 
 class Plotter:
     COLORS = ['blue', 'green', 'red', 'cyan', 'magenta', 'yellow', 'black', 'purple', 'pink',
@@ -61,3 +63,20 @@ class Plotter:
         plt.ylabel("Episode Rewards")
         if title is not None:
             plt.title(title)
+
+    def load_log_dirs(self, pattern, negative_pattern=' ', root='./log'):
+        dirs = [item[0] for item in os.walk(root)]
+        leaf_dirs = []
+        for i in range(len(dirs)):
+            if i + 1 < len(dirs) and dirs[i + 1].startswith(dirs[i]):
+                continue
+            leaf_dirs.append(dirs[i])
+        names = []
+        p = re.compile(pattern)
+        np = re.compile(negative_pattern)
+        for dir in leaf_dirs:
+            if p.match(dir) and not np.match(dir):
+                names.append(dir)
+                print(dir)
+
+        return sorted(names)