adding tool for visualizing acceleration distributions, and making nframes an arg

src/bc/bc.py

@@ -39,7 +39,7 @@ def bc_config(ray_config):
             'lr':ray_config['lr'],
             'weight_decay':ray_config['weight_decay']
         },
-        'train_epochs': 20,
+        'train_epochs': 40,
         'train_batch_size': ray_config['train_batch_size'],
         'loss': ray_config['loss'],
 

src/main.py

@@ -47,7 +47,7 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
 
         # make policy, train and test datasets, and send to 
         policy = policy_class(config)
-        train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0])#,1,2])
+        train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0,1,2])
         cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3])
         train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
     
@@ -59,7 +59,7 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
 
         # simulate policy
         track = 4
-        simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=500)
+        simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=kwargs['nframes'])
 
         # run test metrics
         test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[track])
@@ -101,4 +101,4 @@ def simulate_policy(policy, loc=0, track=0, filestr='', nframes=float('inf')):
 
             pbar.update()
 
-    env.close(filestr=filestr)
+    env.close(filestr=filestr+'_sim')

src/metrics.py

@@ -1,7 +1,8 @@
 import torch
 import pickle
 import numpy as np
-
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader
 import intersim.collisions
 
 def metrics(filestr: str, test_dataset, policy):
@@ -26,9 +27,38 @@ def metrics(filestr: str, test_dataset, policy):
     # calculate divergence between velocity distributions
     
     # calcuate divergence between acceleration distributions
-     
-    pass
+    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)
+    import pdb
+    pdb.set_trace()
+    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(x):
     """
 
@@ -40,7 +70,7 @@ def divergence(p, q, type='kl'):
     Calculate a divergence between p and q
     Args:
         p (torch.tensor): (n) samples from p
-        q (torch.tensor): (n) samples from q
+        q (torch.tensor): (m) samples from q
     Returns:
         d (float): approximate divergence
     """