BUG FIXES: moving around when policy is loaded, adding BaseAlgorithm abstract classes, correcting metrics, normalizng actions if idm environment is a normalized action one, manually updating environment graph when using idm, implementing idm forward class

src/evaluation/evaluation.py

@@ -1,7 +1,7 @@
 import numpy as np
 from stable_baselines3.common.vec_env import VecEnv
 from stable_baselines3.common.evaluation import evaluate_policy
-from intersim.envs.intersimple import Intersimple
+from intersim.envs.intersimple import Intersimple, IncrementingAgent
 from typing import Callable, Dict, Optional
 import os
 import pickle
@@ -19,21 +19,18 @@ class IntersimpleEvaluation:
         - whether there was a collision
         - whether there was a hard brake
     """
-    def __init__(self, eval_env, use_pbar:bool=True):
+    def __init__(self, eval_env:IncrementingAgent, use_pbar:bool=True):
         """
         Initialize evaluation environment with an Intersimple IncrementingAgent environment
 
         Args:
-            eval_env (Intersimple.IncrementingAgent): evaluation environment that increments agent upon reset
+            eval_env (IncrementingAgent): evaluation environment that increments agent upon reset
             use_pbar (bool): whether to use a progress bar
         """
         # if env is a VecEnv, the code needs to be adapted, since the callback will be called after each step, 
         # so transitions of different envs will be mixed and the total number of episodes could be larger than n_eval_episodes!
         assert not isinstance(eval_env, VecEnv)
-        
-        # make sure we have specified an environment which increments the agent number on reset
-        assert isinstance(eval_env, Intersimple.IncrementingAgent)
-
+  
         self.env = eval_env
         self.n_episodes = eval_env.nv
         self.use_pbar = use_pbar
@@ -54,7 +51,7 @@ class IntersimpleEvaluation:
         """
         Reset metrics prior to evaluation
         """
-        self._metrics = {key: [[]]*self.n_episodes for key in self.metric_keys_all}
+        self._metrics = {key: [[] for _ in range(self.n_episodes)] for key in self.metric_keys_all}
         self._metrics.update({key: [None]*self.n_episodes for key in self.metric_keys_single})
     
     def save(self, filestr):
@@ -66,8 +63,8 @@ class IntersimpleEvaluation:
         """
         # assert metrics all have correct length
         for key in self.metric_keys:
-            assert(len(self._metrics[key])==self.n_episodes, 
-                f'_metrics[{key}] does not have length {self.n_episodes}')
+            assert len(self._metrics[key])==self.n_episodes, \
+                f'_metrics[{key}] does not have length {self.n_episodes}'
 
         # make filepath
         os.makedirs(os.path.dirname(filestr))
@@ -91,13 +88,13 @@ class IntersimpleEvaluation:
         evaluate_policy(
             policy, 
             self.env,
-            n_eval_episodes=self.n_eval_episodes,
+            n_eval_episodes=self.n_episodes,
             callback=self.evaluate_policy_callback, 
             return_episode_rewards=False
         )
         if self.use_pbar:
             self.pbar.close()
-            
+
         self.post_proc()
         if filestr:
             self.save(filestr)
@@ -119,6 +116,7 @@ class IntersimpleEvaluation:
         self._metrics['v_all'][_agent].append(info['prev_state'][_agent,2].item())
         self._metrics['a_all'][_agent].append(info['action_taken'][_agent,0].item())
         col = info['collision']
+
         if col:
             assert done
         self._metrics['col_all'][_agent].append(col)
@@ -138,7 +136,7 @@ class IntersimpleEvaluation:
             self._metrics['a_all'][i] = np.array(self._metrics['a_all'][i])
 
             # jerk
-            self._metrics['j_all'][i] = np.diff(self._metrics['a_all'][i]) / self.eval_env._env._dt
+            self._metrics['j_all'][i] = np.diff(self._metrics['a_all'][i]) / self.env._env._dt
 
             # average velocity and acceleration
             self._metrics['v_avg'][i] = np.mean(self._metrics['v_all'][i])

src/evaluation/metrics.py

@@ -4,83 +4,24 @@ import numpy as np
 import matplotlib.pyplot as plt
 from torch.utils.data import DataLoader
 from intersim import collisions
+# import tikzplotlib 
 
-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
-    Returns:
-        info (dict): metrics in a dictionary
-    """
-    info = {}
-
-    # 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 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()
-    xpoly = torch.load(filestr + '_sim_xpoly.pt').detach()
-    ypoly = torch.load(filestr + '_sim_ypoly.pt').detach()
-
-    # count collisions (from function in intersim.collisions)
-    n_collisions = collisions.count_collisions_trajectory(states, lengths, widths)
-    info['n_collisions'] = n_collisions
-
-    # calculate average velocity
-    avg_v = average_velocity(states)
-    info['average_velocity'] = avg_v
-
-    # convert policy dtype between float32 and float64
-    policy.policy = policy.policy.type(test_dataset[0]['state']['ego_state'].dtype)
-    
-    # generate actions in test dataset
-    true_actions, pred_actions = [], []
-    true_velocities = []
-    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['state']))
-            true_actions.append(batch['action'])
-            true_velocities.append(batch['state']['ego_state'][:,2])
-
-    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') 
-
-    # calculate divergence between acceleration distributions
-    acceleration_divergence = divergence(pred_actions, true_actions, type='js')
-    info['acceleration_divergence'] = acceleration_divergence
-
-    # calculate divergence between velocity distributions
-    sim_velocities = states[:,:,2]
-    sim_velocities = sim_velocities[~torch.isnan(sim_velocities)].flatten()
-    true_velocities = torch.cat(true_velocities, dim=0)
-    velocity_divergence = divergence(sim_velocities, true_velocities, type='js')
-    info['velocity_divergence'] = velocity_divergence
-
-    return info
-
-
-def visualize_distribution(true, pred, filestr):
+def visualize_distribution(expert, policy, filestr):
     """
     Visualize two distributions
     Args:
-        true (torch.tensor): (n,)-sized true distribution
-        pred (torch.tensor): (m,)-sized pred distribution
+        expert (torch.tensor): (n,)-sized true distribution
+        generated (torch.tensor): (m,)-sized pred distribution
         filestr (str): string to save figure to
     """
-    nni1 = ~torch.isnan(true)
-    nni2 = ~torch.isnan(pred)
+    nni1 = ~torch.isnan(expert)
+    nni2 = ~torch.isnan(policy)
     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.hist(expert[nni1].numpy(), density=True, bins=20)
+    plt.hist(policy[nni2].numpy(), density=True, bins=20)
+    plt.legend(['Expert', 'Predicted'])
     plt.savefig(filestr+'.png')
+    # tikzplotlib.save(filestr+'.tex')
     
 def average_velocity(states):
     """