From 1a74fa5237f88479ca64faef8351a9542301e91a Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Tue, 20 Jul 2021 07:05:09 -0700
Subject: [PATCH] making general-purpose metric function

---
 src/main.py    | 10 ++++------
 src/metrics.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 6 deletions(-)
 create mode 100644 src/metrics.py

diff --git a/src/main.py b/src/main.py
index c071dc4..b7068e5 100644
--- a/src/main.py
+++ b/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=''):
diff --git a/src/metrics.py b/src/metrics.py
new file mode 100644
index 0000000..cfe0568
--- /dev/null
+++ b/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
\ No newline at end of file