diff --git a/experiments/experiment.py b/experiments/experiment.py index 43d50eb..8f26d9d 100644 --- a/experiments/experiment.py +++ b/experiments/experiment.py @@ -166,7 +166,7 @@ if __name__ == '__main__': config_path = filestr+'_config.json' with open(config_path, 'r') as cfg: config = json5.load(cfg) - print(filepath) + print("Best ray experiment:", filepath) main(config, filestr=filestr, **kwargs) else: raise Exception('No valid config found') diff --git a/src/metrics.py b/src/metrics.py index 98a677e..5e4997f 100644 --- a/src/metrics.py +++ b/src/metrics.py @@ -18,7 +18,7 @@ def metrics(filestr: str, test_dataset, policy): # a) simulation files that were saved under the trained policy with prefix 'policy' # b) applying the policy to observations in the test dataset - # load trajectory + # 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() @@ -77,16 +77,94 @@ def average_velocity(states): arg_v = nanmean(vehicle_avg_v) return arg_v -def divergence(p, q, type='kl'): +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 + type (str): divergence to use + 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 """ - pass + 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): """ @@ -105,4 +183,23 @@ def nanmean(v, *args, inplace=False, **kwargs): is_nan = torch.isnan(v) v[is_nan] = 0 result = v.sum(*args, **kwargs) / (~is_nan).float().sum(*args, **kwargs) - return result \ No newline at end of file + 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 + diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..bef990d --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,62 @@ +import torch +import numpy as np + +from src import metrics +from src.metrics import divergence, evaluate_histogram +from sklearn.model_selection import GridSearchCV +from sklearn.neighbors import KernelDensity + +def test_kl_divergence(): + p = torch.randn(1000000) + q = 1.0 + 2.0 * torch.randn(1000000) + + d1 = divergence(p, q, type='kl', n_components=0) + print(d1) + assert isinstance(d1, float) + d2 = divergence(p, q, type='kl', n_components=1) + print(d2) + assert isinstance(d2, float) + assert np.isclose(d1, d2, atol=1e-5) + d3 = divergence(p, q, type='kl', n_components=3) + print(d3) + assert isinstance(d3, float) + d4 = divergence(p, q, type='kl', n_components=-1) + assert isinstance(d4, float) + d5 = divergence(q, p, type='kl', n_components=-1) + assert isinstance(d5, float) + + p = 1000000 * torch.randn(1000) + q = torch.randn(1000) + d6 = divergence(p, q, type='kl', n_components=-1) + assert np.isnan(d6) + +def test_evaluate_histogram(): + N = 10000 + p = torch.randn(N) + q = 1.0 + 2.0 * torch.randn(2*N) + + 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) + assert px.shape == p.shape + qx = evaluate_histogram(q, p_hist, p_edges) + assert qx.shape == q.shape + px = evaluate_histogram(p, q_hist, q_edges) + assert px.shape == p.shape + qx = evaluate_histogram(q, q_hist, q_edges) + assert qx.shape == q.shape + + +if __name__ == '__main__': + p = torch.randn(10) + q = 1.0 + 2.0 * torch.randn(10) + + p = p.unsqueeze(-1) + q = q.unsqueeze(-1) + 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 = p_hist[np.digitize(p, p_edges) - 1] + + # qx = q_hist[np.digitize(p, q_edges) - 1] + px = evaluate_histogram(p, p_hist, p_edges) + qx = evaluate_histogram(q, p_hist, p_edges)