/
make_plots.py
248 lines (220 loc) · 8.32 KB
/
make_plots.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
import argparse
from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib
import numpy as np
import os
import sys
import pandas as pd
matplotlib.rcParams['font.size'] = 12
def assign_color(path, params):
assignments = {'joint':'tab:gray', 'disagreement':'tab:green',
'concat':'tab:purple', 'visual':'#ff7f0e', 'fft':'#1f77b4'}
for method in assignments.keys():
if method in path or method in params:
if 'noise' in params:
alpha = 0.5
else:
alpha = 1.0
return assignments[method], alpha
return None, 1.0
def assign_label(path, params):
if 'b-' in path:
label = 'Visual prediction curiosity'
elif 'concat' in path:
label = 'Audio-visual prediction'
else:
label = 'Audio-visual association (ours)'
if 'noise' in params:
if '(ours)' in label:
label = label[:-7]
label += ' ' + params.split('noise')[1] + 'with noise'
if 'joint' in path:
label = 'Combined curiosities (ours)'
if 'disagreement' in params:
label = 'Disagreement'
return label
def plot_run(paths, x='tcount', ys='eprew', label='', assign_colors=True, params=''):
for y in ys.split(','):
all_runs = []
nframes = None
for path in paths:
f = open(path + '/progress.csv', 'r')
if 'concat' in path and not args.concat:
continue
try:
df = pd.read_csv(f)
all_runs.append(df[y])
nframes = df[x]
except:
print("Exception while reading file", sys.exc_info()[0], path)
continue
if len(all_runs) == 0:
return
color = None
alpha = 0.75
if assign_colors:
color, alpha = assign_color(path, params)
if label != '':
if params != '':
label = assign_label(path, params)
if len(ys.split(',')) > 1:
label += ' ' + y
min_length = min([len(run) for run in all_runs])
all_runs = np.asarray([run[:min_length] for run in all_runs])
mean_run = np.mean(all_runs, axis=0)
nframes = nframes[:min_length] * 4
ax.plot(nframes, mean_run, '-', label=label, color=color, alpha=alpha)
if all_runs.shape[0] > 1:
error = np.std(all_runs, axis=0)
alpha = 0.2
if 'noise' in params:
alpha = 0.1
ax.fill_between(nframes, mean_run-error, mean_run+error,
alpha=alpha, linewidth=0.0, color=color)
def get_paths(paths_str):
paths = []
for path in paths_str.split(','):
if '*' in path:
try:
paths.extend(os.popen('ls -d logs/*' + path[1:] + '*/ 2> /dev/null').read().split())
except:
print("Exception: paths not found")
continue
else:
paths.append(path)
return paths
def create_params(path):
if ('_s-' in path) != args.sticky:
return None
if 'noise' in path and not args.noise:
return None
if 'joint' in path:
return None
if args.mean:
params = '_'.join(path.split('_')[-6:])
if 'noise' in path:
noise_amt = float(path.split('noise')[1].split('-')[0])
if noise_amt > 0.1:
return None
if noise_amt != 0:
params = params[:-1]
params += 'noise'
if 'concat' in path:
concat_type_real = ' ' + path.split('concat')[0].split('breakout')[-1][1:-1]
concat_type = path.split('concat')[1].split('_')[0][:-1]
if len(concat_type) == 0 or concat_type == '-both':
params += '-concat'
else:
params += concat_type
if '_s-' in path:
params += '_sticky'
if 'disagreement' in path or 'openai' in path:
params = get_disagreement_params(path)
if 'joint' in path:
params += '_joint'
return params
def create_figure():
if args.all:
if args.noise:
envs = ['Asterix', 'MsPacman', 'SpaceInvaders']
fig, axs = plt.subplots(1, 3, sharex=True, figsize=(18, 6))
else:
envs = ['AirRaid', 'Alien', 'Assault', 'Asterix', 'Asteroids', 'Atlantis', \
'BeamRider', 'Breakout', 'MsPacman', 'Qbert', 'Seaquest', 'SpaceInvaders']
if args.sticky:
fig, axs = plt.subplots(2, 6, sharex=True, figsize=(18, 6))
else:
fig, axs = plt.subplots(3, 4, sharex=True, figsize=(20, 15))
fig.add_subplot(111, frameon=False)
# hide tick and tick label of the big axes
plt.tick_params(axis='both', which='both', bottom=False, top=False,
left=False, right=False, labelcolor='none')
plt.grid(False)
else:
fig, ax = plt.subplots(1, 1)
axs = np.array(ax)
envs = ['']
fsize = 32
if args.noise:
fsize = 20
plt.xlabel("\nFrames (millions)\n", fontsize=fsize)
plt.ylabel("Extrinsic Reward Per Episode\n", fontsize=fsize)
return fig, axs, envs
def finish_plot():
if args.xlim:
plt.xlim((-5*1e6, args.xlim * 1e6))
if args.ylim:
plt.ylim((-0.0 * args.ylim, args.ylim))
mean_str = '_mean' if args.mean else ''
if args.all:
sticky_str = '_sticky' if args.sticky else ''
noise_str = '_noise' if args.noise else ''
save_name = 'all' + sticky_str + mean_str + noise_str + "_" + args.y + '.png'
else:
save_name = path.split('/')[-1] + '_' + str(len(ps)) + \
'_' + str(args.y) + mean_str + '.png'
plt.savefig(save_name)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--paths', type=str)
parser.add_argument('-x', type=str, default='tcount')
parser.add_argument('-xlim', type=int, default=200)
parser.add_argument('-ylim', type=float)
parser.add_argument('-y', type=str, default='eprew')
parser.add_argument('--mean', type=bool, default=False)
parser.add_argument('--all', type=bool, default=False)
parser.add_argument('--sticky', type=bool, default=False)
parser.add_argument('--noise', type=bool, default=False)
parser.add_argument('--concat', type=bool, default=False)
parser.add_argument('--assign_colors', type=bool, default=False)
args = parser.parse_args()
fig, axs, envs = create_figure()
env_index = 0
for ax in axs.flat:
runs_by_params = defaultdict(list)
if args.all:
if env_index >= len(envs):
break
ps = get_paths('*' + envs[env_index])
else:
ps = get_paths(args.paths)
if len(ps) == 0:
print("No paths found")
sys.exit()
for path in ps:
# Remove trailing slash from pathname
if path[-1] == '/':
path = path[:-1]
if args.mean:
params = create_params(path)
if params is not None:
runs_by_params[params].append(path)
else:
run_name = path.split('/')[-1]
plot_run([path], args.x, args.y, label=run_name, assign_colors=args.assign_colors)
if args.mean:
for params, runs in runs_by_params.items():
# Only add labels for legend in last plot
label = params
if args.all and env_index != len(envs) - 1:
label = ''
plot_run(runs, args.x, args.y, label=label, assign_colors=True, params=params)
ax.ticklabel_format(style='sci', axis='x', scilimits=(6, 6))
if not args.all or env_index == len(envs) - 1:
ncol = 3
fsize = 25
loc = (0.05, 0.01)
if args.noise:
ncol = 2
fsize = 18
loc = (0.25, 0.01)
leg = fig.legend(prop={'size': fsize}, ncol=ncol, loc=loc)
for line in leg.get_lines():
line.set_linewidth(4.0)
ax.tick_params(axis='both', labelsize=14)
ax.set_title(envs[env_index], fontsize=16)
ax.set_xlim((-5 * 1e6, args.xlim * 1e6))
env_index += 1
fig.tight_layout(pad=0.5)
finish_plot()