diff --git a/main.py b/main.py index 89b1686..e3bc505 100644 --- a/main.py +++ b/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' diff --git a/utils/plot.py b/utils/plot.py index 62c75d3..bac9541 100644 --- a/utils/plot.py +++ b/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)