201 lines
6.7 KiB
Python
201 lines
6.7 KiB
Python
import torch
|
|
import pickle
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from torch.utils.data import DataLoader
|
|
from intersim import collisions
|
|
|
|
def metrics(filestr: str, test_dataset, policy):
|
|
"""
|
|
Calculate metrics using a) base filestring to a simulation, and b) the test dataset and learned policy
|
|
Args:
|
|
filestr (str): base string to outputs of a simulation
|
|
test_dataset: a dataset held for testing
|
|
policy: policy
|
|
"""
|
|
|
|
# compute metrics using either
|
|
# a) simulation files that were saved under the trained policy with prefix 'policy'
|
|
# b) applying the policy to observations in the test dataset
|
|
|
|
# load simulated trajectory
|
|
states = torch.load(filestr + '_sim_states.pt').detach()
|
|
lengths = torch.load(filestr + '_sim_lengths.pt').detach()
|
|
widths = torch.load(filestr + '_sim_widths.pt').detach()
|
|
xpoly = torch.load(filestr + '_sim_xpoly.pt').detach()
|
|
ypoly = torch.load(filestr + '_sim_ypoly.pt').detach()
|
|
|
|
# count collisions (from function in intersim.collisions)
|
|
n_collisions = collisions.count_collisions_trajectory(states, lengths, widths)
|
|
|
|
# calculate average velocity
|
|
avg_v = average_velocity(states)
|
|
|
|
# calculate divergence between velocity distributions
|
|
|
|
# calcuate divergence between acceleration distributions
|
|
policy.policy = policy.policy.type(test_dataset[0]['state'].dtype)
|
|
|
|
# generate actions in test dataset
|
|
true_actions, pred_actions = [], []
|
|
test_loader = DataLoader(test_dataset, batch_size=1024)
|
|
with torch.no_grad():
|
|
for (batch_idx, batch) in enumerate(test_loader):
|
|
pred_actions.append(policy(batch))
|
|
true_actions.append(batch['action'])
|
|
|
|
true_actions, pred_actions = torch.cat(true_actions,dim=0), torch.cat(pred_actions,dim=0)
|
|
visualize_distribution(true_actions[:,0], pred_actions[:,0], filestr+'_action_viz')
|
|
|
|
|
|
def visualize_distribution(true, pred, filestr):
|
|
"""
|
|
Visualize two distributions
|
|
Args:
|
|
true (torch.tensor): (n,)-sized true distribution
|
|
pred (torch.tensor): (m,)-sized pred distribution
|
|
filestr (str): string to save figure to
|
|
"""
|
|
nni1 = ~torch.isnan(true)
|
|
nni2 = ~torch.isnan(pred)
|
|
plt.figure()
|
|
plt.hist(true[nni1].numpy(), density=True, bins=20)
|
|
plt.hist(pred[nni2].numpy(), density=True, bins=20)
|
|
plt.legend(['True', 'Predicted'])
|
|
plt.savefig(filestr+'.png')
|
|
|
|
def average_velocity(states):
|
|
"""
|
|
Compute average of average velocity over all vehicles.
|
|
Args:
|
|
states (torch.tensor): (T,nv,5) vehicle states
|
|
Returns
|
|
avg_v (float): average velocity
|
|
"""
|
|
velocities = states[:,:,2]
|
|
vehicle_avg_v = nanmean(velocities, dim=0)
|
|
arg_v = nanmean(vehicle_avg_v)
|
|
return arg_v
|
|
|
|
def divergence(p, q, type='kl', n_components=0):
|
|
"""
|
|
Calculate a divergence between p and q
|
|
Args:
|
|
p (torch.tensor): (n) samples from p
|
|
q (torch.tensor): (m) samples from q
|
|
Returns:
|
|
d (float): approximate divergence
|
|
"""
|
|
if type == 'kl':
|
|
if n_components < 0:
|
|
# Use histogram binning to discretize sampled distributions
|
|
p_hist, p_edges = np.histogram(p.unsqueeze(-1), bins='auto', density=True)
|
|
q_hist, q_edges = np.histogram(q.unsqueeze(-1), bins='auto', density=True)
|
|
px = evaluate_histogram(p, p_hist, p_edges)
|
|
qx = evaluate_histogram(p, q_hist, q_edges)
|
|
p_supp = ~np.isclose(px, 0.0)
|
|
q_supp = ~np.isclose(qx, 0.0)
|
|
if np.any(np.logical_and(p_supp, ~q_supp)):
|
|
# if not support(p) subset support(q)
|
|
return np.nan
|
|
elif ~np.any(p_supp):
|
|
# if p is zero everywhere
|
|
return 0.
|
|
d = np.mean(np.log(px[p_supp] / qx[p_supp]))
|
|
return d
|
|
elif n_components == 0:
|
|
# Assume p and q to be Gaussian
|
|
pm = torch.mean(p)
|
|
qm = torch.mean(q)
|
|
pv = torch.var(p)
|
|
qv = torch.var(q)
|
|
d = kl_normal(pm, pv, qm, qv).item()
|
|
return d
|
|
else:
|
|
from sklearn.mixture import GaussianMixture
|
|
p = p.unsqueeze(-1)
|
|
q = q.unsqueeze(-1)
|
|
p_gmm = GaussianMixture(n_components=n_components).fit(p)
|
|
q_gmm = GaussianMixture(n_components=n_components).fit(q)
|
|
px = p_gmm.score_samples(p)
|
|
qx = q_gmm.score_samples(p)
|
|
d = np.mean(px - qx).item()
|
|
return d
|
|
else:
|
|
raise NotImplementedError("Please implement divergence for type '{}'".format(type))
|
|
|
|
|
|
def kl_normal(pm, pv, qm, qv):
|
|
"""
|
|
Computes the elem-wise KL divergence between two normal distributions KL(p || q) and
|
|
sum over the last dimension
|
|
|
|
Args:
|
|
pm: tensor: (batch, dim): p mean
|
|
pv: tensor: (batch, dim): p variance
|
|
qm: tensor: (batch, dim): q mean
|
|
qv: tensor: (batch, dim): q variance
|
|
|
|
Return:
|
|
kl: tensor: (batch,): kl between each sample
|
|
"""
|
|
element_wise = 0.5 * (torch.log(qv) - torch.log(pv) + pv / qv + (pm - qm).pow(2) / qv - 1)
|
|
kl = element_wise.sum(-1)
|
|
return kl
|
|
|
|
|
|
def kl_cat(q, log_q, log_p):
|
|
"""
|
|
Computes the KL divergence between two categorical distributions
|
|
|
|
Args:
|
|
q: tensor: (batch, dim): Categorical distribution parameters
|
|
log_q: tensor: (batch, dim): Log of q
|
|
log_p: tensor: (batch, dim): Log of p
|
|
|
|
Return:
|
|
kl: tensor: (batch,) kl between each sample
|
|
"""
|
|
element_wise = (q * (log_q - log_p))
|
|
kl = element_wise.sum(-1)
|
|
return kl
|
|
|
|
|
|
def nanmean(v, *args, inplace=False, **kwargs):
|
|
"""
|
|
Calculate mean over not nan entries
|
|
|
|
To be added to torch as torch.nanmean in the next release
|
|
https://github.com/pytorch/pytorch/issues/61474, https://github.com/pytorch/pytorch/issues/21987
|
|
|
|
Args:
|
|
v (torch.tensor): arbitrary tensor
|
|
Returns:
|
|
result (torch.tensor): mean over non nan elements
|
|
"""
|
|
if not inplace:
|
|
v = v.clone()
|
|
is_nan = torch.isnan(v)
|
|
v[is_nan] = 0
|
|
result = v.sum(*args, **kwargs) / (~is_nan).float().sum(*args, **kwargs)
|
|
return result
|
|
|
|
|
|
def evaluate_histogram(x, hist, bin_edges):
|
|
"""
|
|
Evaluate a histogram
|
|
Args:
|
|
x (array) : points at which to evaluate the histogram
|
|
hist (array): histogram values in terms of number of occurrences or probability
|
|
bin_edges (array): edges of histogram bins
|
|
e.g. from hist, bin_edges = np.histogram(p, bins='auto', density=True)
|
|
Return:
|
|
r: tensor: (batch,) kl between each sample
|
|
"""
|
|
idx = np.digitize(x, bin_edges)
|
|
mask = np.logical_and(np.less(0, idx), np.less(idx, len(bin_edges)))
|
|
r = np.zeros_like(x)
|
|
r[mask] = hist[idx[mask] - 1]
|
|
return r
|
|
|