making general-purpose metric function

src/main.py

@@ -4,7 +4,7 @@ import intersim
 import numpy as np
 import os
 opj = os.path.join
-from src import InteractionDatasetSingleAgent
+from src import InteractionDatasetSingleAgent, metrics
 
 def basestr(**kwargs):
     """
@@ -31,15 +31,13 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs):
         os.mkdir(outdir)
     filestr = opj(outdir, basestr(**kwargs))
    
-    # define transforms
-    transforms={}
+
     
     # method-based training
     if method=='bc':
         from src import bc
         policy_class = bc.BehaviorCloningPolicy
         load_policy_fn = bc.load_policy
-        metrics_fn = bc.metrics
         train_fn = bc.train
     else:
         raise NotImplementedError
@@ -60,8 +58,8 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs):
         simulate_policy(policy, loc=loc, track=track, filestr=filestr)
 
         # run test metrics
-        # test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[4])
-        # metrics_fn(test_dataset, policy)
+        test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[4])
+        metrics(filestr, test_dataset, policy)
 
 
 def simulate_policy(policy, loc=0, track=0, filestr=''):

src/metrics.py

@@ -0,0 +1,47 @@
+import torch
+import pickle
+import numpy as np
+
+import intersim.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
+    """
+
+    # 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 trajectory
+
+    # count collisions (from function in intersim.collisions)
+
+    # calculate average velocity
+
+    # calculate divergence between velocity distributions
+    
+    # calcuate divergence between acceleration distributions
+     
+    pass
+
+def average_velocity(x):
+    """
+
+    """
+    pass
+
+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
+    Returns:
+        d (float): approximate divergence
+    """
+    pass