forked from Kaixhin/imitation-learning
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
24 lines (19 loc) · 747 Bytes
/
utils.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
from matplotlib import pyplot as plt
import numpy as np
import seaborn as sns
import torch
sns.set(style='white')
# Flattens a list of dicts with torch Tensors
def flatten_list_dicts(list_dicts):
return {k: torch.cat([d[k] for d in list_dicts], dim=0) for k in list_dicts[-1].keys()}
# Makes a lineplot with scalar x and statistics of vector y
def lineplot(x, y, filename, xaxis='Steps', yaxis='Returns'):
y = np.array(y)
y_mean, y_std = y.mean(axis=1), y.std(axis=1)
sns.lineplot(x=x, y=y_mean, color='coral')
plt.fill_between(x, y_mean - y_std, y_mean + y_std, color='coral', alpha=0.3)
plt.xlim(left=0, right=x[-1])
plt.xlabel(xaxis)
plt.ylabel(yaxis)
plt.savefig(f'{filename}.png')
plt.close()