Merge branch 'main' of https://github.com/sisl/InteractionImitation into main

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')

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
+    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
+

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)