253 lines
9.0 KiB
Python
253 lines
9.0 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
|
|
Returns:
|
|
info (dict): metrics in a dictionary
|
|
"""
|
|
info = {}
|
|
|
|
# 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)
|
|
info['n_collisions'] = n_collisions
|
|
|
|
# calculate average velocity
|
|
avg_v = average_velocity(states)
|
|
info['average_velocity'] = avg_v
|
|
|
|
# convert policy dtype between float32 and float64
|
|
policy.policy = policy.policy.type(test_dataset[0]['state'].dtype)
|
|
|
|
# generate actions in test dataset
|
|
true_actions, pred_actions = [], []
|
|
true_velocities = []
|
|
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_velocities.append(batch['state'][:,2])
|
|
|
|
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')
|
|
|
|
# calculate divergence between acceleration distributions
|
|
acceleration_kl = divergence(pred_actions, true_actions, type='kl', n_components=-1)
|
|
info['acceleration_kl'] = acceleration_kl
|
|
|
|
# calculate divergence between velocity distributions
|
|
sim_velocities = states[:,:,2]
|
|
sim_velocities = sim_velocities[~torch.isnan(sim_velocities)].flatten()
|
|
true_velocities = torch.cat(true_velocities, dim=0)
|
|
velocity_kl = divergence(sim_velocities, true_velocities, type='kl', n_components=-1)
|
|
info['velocity_kl'] = velocity_kl
|
|
|
|
return info
|
|
|
|
|
|
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 where T is the number of time steps and nv the number of vehicles
|
|
Returns
|
|
avg_v (float): average velocity
|
|
"""
|
|
velocities = states[:,:,2]
|
|
# average velocity per vehicle
|
|
vehicle_avg_v = nanmean(velocities, dim=0)
|
|
arg_v = nanmean(vehicle_avg_v)
|
|
return arg_v
|
|
|
|
def divergence(p, q, type='js', 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
|
|
type (str): divergence to use
|
|
'kl': Kullback-Leibler divergence KL(p||q)
|
|
'js': Jensen-Shannon divergence (symmetric KLD)
|
|
n_components (int): method to use to compute kl divergence
|
|
n_components < 0: approximate samples with histogram density
|
|
n_components == 0: approximate samples by Gaussian distributions and compute analytically
|
|
n_components > 0: approximate samples as Gaussian mixture models with n_components components
|
|
Returns:
|
|
d (float): approximate divergence
|
|
"""
|
|
if type == 'js':
|
|
# Use histogram binning to discretize sampled distributions
|
|
p_hist = np.histogram(p, bins='auto', density=True)
|
|
q_hist = np.histogram(q, bins='auto', density=True)
|
|
m = torch.cat([p, q], dim=0)
|
|
m_weights = torch.cat([torch.full_like(p, 1./len(p)), torch.full_like(q, 1./len(q))], dim=0)
|
|
m_bins = np.sort(np.concatenate([p_hist[1], q_hist[1]]))
|
|
m_hist = np.histogram(m, bins=m_bins, density=True, weights=m_weights)
|
|
d = .5 * kl_histogram(p, p_hist, m_hist) + .5 * kl_histogram(q, q_hist, m_hist)
|
|
return d
|
|
|
|
elif type == 'kl':
|
|
if n_components < 0:
|
|
# Use histogram binning to discretize sampled distributions
|
|
p_hist = np.histogram(p, bins='auto', density=True)
|
|
q_hist = np.histogram(q, bins='auto', density=True)
|
|
d = kl_histogram(p, p_hist, q_hist)
|
|
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_histogram(p_sample, p_hist, q_hist):
|
|
"""
|
|
Calculate the kl divergence between p and q based on a histogram representation
|
|
Args:
|
|
p_sample (torch.tensor): (n) samples from p
|
|
p_hist (tuple): result of np.histogram(density=True) for samples from p
|
|
q_hist (tuple): result of np.histogram(density=True) for samples from q
|
|
Returns:
|
|
d (float): approximate KL divergence
|
|
"""
|
|
p_density, p_edges = p_hist
|
|
q_density, q_edges = q_hist
|
|
px = evaluate_histogram(p_sample, p_density, p_edges)
|
|
qx = evaluate_histogram(p_sample, q_density, 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.inf
|
|
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
|
|
|
|
|
|
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
|
|
|