From 3b4ef6ffb5bf708dd51935f7b4006a0da8265b15 Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Thu, 29 Jul 2021 07:35:21 -0700
Subject: [PATCH] adding tool for visualizing acceleration distributions, and
 making nframes an arg

---
 experiments/experiment.py |  7 ++++++-
 src/bc/bc.py              |  2 +-
 src/main.py               |  6 +++---
 src/metrics.py            | 38 ++++++++++++++++++++++++++++++++++----
 4 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/experiments/experiment.py b/experiments/experiment.py
index 593f6b0..bed727e 100644
--- a/experiments/experiment.py
+++ b/experiments/experiment.py
@@ -33,6 +33,8 @@ def parse_args():
         default=None, type=str)
     parser.add_argument('--seed', default=0, type=int,
         help='seed')    
+    parser.add_argument('--nframes', default=500, type=int,
+        help='frames for test animation') 
     args = parser.parse_args()
     kwargs = {
         'train':args.train, 
@@ -41,7 +43,8 @@ def parse_args():
         'loc':args.loc,
         'config_path':args.config,
         'seed':args.seed,
-        'ray':args.ray
+        'ray':args.ray,
+        'nframes':args.nframes,
     }
     return kwargs
 
@@ -100,6 +103,8 @@ if __name__ == '__main__':
         if not os.path.isdir(outdir):
             os.makedirs(outdir)   
         filestr = opj(outdir, basestr(**kwargs)) 
+        if kwargs['ray']:
+            filestr = kwargs['config_path'].replace('_config.json','')
         main(config, filestr=filestr, **kwargs)
 
     elif kwargs['ray'] and kwargs['train']:
diff --git a/src/bc/bc.py b/src/bc/bc.py
index 418df82..687e476 100644
--- a/src/bc/bc.py
+++ b/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'],
 
diff --git a/src/main.py b/src/main.py
index c924d28..04282a8 100644
--- a/src/main.py
+++ b/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')
diff --git a/src/metrics.py b/src/metrics.py
index b75e0e7..00d56b3 100644
--- a/src/metrics.py
+++ b/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
     """