From 3991306da0d9b6684f540eda260b54a283241d93 Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Mon, 31 Jan 2022 16:28:13 -0800
Subject: [PATCH 1/6] updating evaluation wrapper to only store relevant
 variables during execution

---
 src/evaluation/evaluation.py | 182 ++++++++++++++++++++++++-----------
 1 file changed, 124 insertions(+), 58 deletions(-)

diff --git a/src/evaluation/evaluation.py b/src/evaluation/evaluation.py
index 0cefc22..de1306c 100644
--- a/src/evaluation/evaluation.py
+++ b/src/evaluation/evaluation.py
@@ -1,76 +1,105 @@
-import torch
 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 src.evaluation.metrics import nanmean, divergence, visualize_distribution
+from typing import Callable, Dict
 import os
+import pickle
 
+class IntersimpleEvaluation:
+    """
+    Class to evaluate a policy on n_agents in an intersimple environment and store metrics for each single agent:
+        - all velocities
+        - all accelerations
+        - all jerks
+        - average velocity
+        - average acceleration
+        - existence time
+        - whether there was a collision
+        - whether there was a hard brake
+    """
+    def __init__(self, eval_env):
+        """
+        Initialize evaluation environment with an Intersimple IncrementingAgent environment
 
-class Evaluation:
-    def __init__(self, filestr, eval_env, expert_data, n_eval_episodes=10):
+        Args:
+            eval_env (Intersimple.IncrementingAgent): evaluation environment that increments agent upon reset
+        """
         # 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)
-        self.filestr = filestr
+        
+        # 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_eval_episodes = n_eval_episodes
-        self.expert_data = expert_data
-        self.compute_expert_features(expert_data)
+        self.n_episodes = eval_env.nv
+
+        # metrics present on every step of every episode
+        self.metric_keys_all = ['v_all', 'a_all', 'col_all']
+
+        # metrics calculated after the fact, with one per episode
+        self.metric_keys_single = ['j_all', 'v_avg','a_avg', 'col','brake', 't']
+        
+        # numbers for calculating metrics
+        self.hard_brake = -3 # acceleration for 'hard brake'
+
+        # reset metrics
         self.reset()
 
     def reset(self):
-        self._n_collisions = 0
-        self._trajectories = []
-        self._episode_done = True
-        self._accelerations = []
+        """
+        Reset metrics prior to evaluation
+        """
+        self._metrics = {key: [[]]*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):
+        """
+        Save metrics to filestr
 
-    def compute_expert_features(self, expert_data):
-        # expert velocities
-        extract_state = lambda info: info['projected_state'][info['agent']]
-        expert_velocities = torch.stack([extract_state(info) for info in expert_data.infos])[:,2]
-        self.expert_velocities = expert_velocities[~torch.isnan(expert_velocities)]
-        # expert accelerations
-        extract_accel = lambda info: info['action_taken'][info['agent']]
-        self.expert_accelerations = torch.cat([extract_accel(info) for info in expert_data.infos])
+        Args:
+            filestr (str): path-like string to dump metrics to
+        """
+        # 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}')
 
-    def evaluate(self, epoch, generator, discriminator):
+        # make filepath
+        os.makedirs(os.path.dirname(filestr))
+
+        # pickle dump
+        with open(filestr, 'wb') as f:
+            pickle.dump(self._metrics, f)     
+
+    def evaluate(self, policy, filestr: str) -> Dict[str, list]:
+        """
+        Evaluate a policy on the incrementing agent evaluation environment
+        
+        Args:
+            policy (BaseClass.BaseAlgorithm): policy in which policy.predict(observation)[0] returns an action
+            filestr (str): path-like string to dump metrics to
+        """
         self.reset()
         metrics = {}
         
-        episode_rewards, episode_lengths = evaluate_policy(
-            generator, 
+        evaluate_policy(
+            policy, 
             self.env,
             n_eval_episodes=self.n_eval_episodes,
             callback=self.evaluate_policy_callback, 
-            return_episode_rewards=True
+            return_episode_rewards=False
         )
-
-        collision_rate = self._n_collisions / self.n_eval_episodes
-        metrics['collision_rate'] = collision_rate
-
-        assert len(self._trajectories) >= self.n_eval_episodes
-
-        # velocities produced by generator
-        policy_velocities = torch.cat([torch.stack(t)[:,2] for t in self._trajectories])
-        # if episodes terminate without collisions, then the state is fully nan
-        policy_velocities = policy_velocities[~torch.isnan(policy_velocities)]
-
-        metrics['avg_velocity_loss'] = (self.expert_velocities.mean() - policy_velocities.mean()).item()
-        metrics['velocity_divergence'] = divergence(policy_velocities, self.expert_velocities, type='js')
-        
-
-        # accelerations produced by generator
-        policy_accelerations = torch.tensor(self._accelerations)
-
-        metrics['acceleration_divergence'] = divergence(policy_accelerations, self.expert_accelerations, type='js')
-        visualize_distribution(self.expert_accelerations, policy_accelerations, os.path.join(self.filestr, '_action_viz{:02}'.format(epoch))) 
-
-        print(metrics)
-        return metrics
+        self.post_proc()
+        self.save(filestr)
+        return self._metrics
 
     def evaluate_policy_callback(self, local_vars, global_vars):
+        """
+        Callback run in evaluate_policy after taking an action and receiving an observation
+        
+        """
         venv_i = local_vars['i']
         info = local_vars['info']
         done = local_vars['done']
@@ -79,15 +108,52 @@ class Evaluation:
         assert isinstance(env, Intersimple)
 
         # Increase collision counter if episode terminated with a collision
-        if info['collision']:
+        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._n_collisions += 1
+        self._metrics['col_all'][_agent].append(col)
+    
+    def post_proc(self):
+        """
+        Postprocess and metrics after simulation episodes
+        """
+        # self.metric_keys_all = ['v_all', 'a_all', 'col_all']
+        # self.metric_keys_single = ['j_all', 'v_avg','a_avg', 'col','brake', 't']
 
-        # if last episode is done, start new trajectory
-        # this is currently not necessary, only if velocity is to be averaged over individual trajectories first
-        # and then averaging over all trajectories
-        if self._episode_done:
-            self._trajectories.append([])
-        self._trajectories[-1].append(info['projected_state'][_agent])
-        self._accelerations.append(info['action_taken'][_agent])
-        self._episode_done = done
+        for i in range(self.n_episodes):
+            self._metrics['v_all'][i] = np.array(self._metrics['v_all'][i])
+            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
+
+            # average velocity and acceleration
+            self._metrics['v_avg'][i] = np.mean(self._metrics['v_all'][i])
+            self._metrics['a_avg'][i] = np.mean(self._metrics['a_all'][i])
+
+            # collision?
+            self._metrics['col'][i] = any(self._metrics['col_all'][i]) 
+
+            # brake?
+            self._metrics['brake'][i] = any(self._metrics['a_all'][i] < self.hard_brake)
+
+            # time length 
+            self._metrics['t'][i] =len(self._metrics['v_all'][i])
+
+
+def load_metrics(filestr:str) -> Dict[str,list]:
+    """
+    Load metrics to filestr
+
+    Args:
+        filestr (str): path-like string to dump metrics to
+
+    Returns
+        metrics (Dict[str, list])
+    """
+    # pickle load
+    with open(filestr, 'rb') as f:
+        metrics = pickle.load(f)
+    return metrics  
\ No newline at end of file

From 3ce86b31f7f36d0e1c1076e1c35054c35b843434 Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Wed, 2 Feb 2022 15:43:23 -0800
Subject: [PATCH 2/6] adding Prop controller and IDMRulePolicy

---
 src/baselines/__init__.py      |   0
 src/baselines/rule_policies.py | 181 +++++++++++++++++++++++++++++++++
 2 files changed, 181 insertions(+)
 create mode 100644 src/baselines/__init__.py
 create mode 100644 src/baselines/rule_policies.py

diff --git a/src/baselines/__init__.py b/src/baselines/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/baselines/rule_policies.py b/src/baselines/rule_policies.py
new file mode 100644
index 0000000..2f395d7
--- /dev/null
+++ b/src/baselines/rule_policies.py
@@ -0,0 +1,181 @@
+from stable_baselines3.common.base_class import BaseAlgorithm
+from intersim.envs.intersimple import Intersimple
+import numpy as np
+
+class PControllerPolicy(BaseAlgorithm):
+
+
+    def __init__(self, env):
+        """
+        Initialize policy with pointer to environment it will run on
+        """
+        assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
+        self._env = env
+        self.target_v = 8.94 # m/s
+        self.attn_weight = 20
+    
+    def predict(self, observation, *args, **kwargs):
+        """
+        Generate action, state from observation
+
+        (But actually generate next action from underlying environment state)
+        
+        Args:
+            observation (np.ndarray): instantaneous observation from environment
+
+        Returns
+            action (np.ndarray): action for controlled agent to take
+            state (np.ndarray): hidden state for use in next prediction (null)
+        """
+        agent = self._env._agent
+        ego_state = self._env._env.projected_state[agent].numpy() # (5,) tensor
+        
+        
+        # relative_state = np.delete(self._env._env.relative_state[agent].numpy(), agent, axis=0) #(nv-1, 6) tensor
+        
+        # calculate front and left distances from ego
+
+        # calculate relative speed in direction of position difference vector
+        
+        # calculate angle alpha and distance d of vehicle i from ego heading
+        
+        # attn[i] ~= exp( -(alpha[i])^2 - .01 * d[i] - .1 * vrel[i]
+
+
+        # Proportional controller
+        # action = (self.target_v  - self.attn_weight * attn.sum()) - ego_state[2]
+        action = self.target_v  - ego_state[2]
+        return action, None
+
+class IDMRulePolicy(BaseAlgorithm):
+    """
+    IDMRulePolicy returns action predictions based on an IDM policy.
+
+    The front car is chosen as the closer of:
+        - closest car within a 45 degree half angle cone of the ego's heading
+        - ''' after propagating the environment forward by `t_future' seconds with 
+            current headings and velocities
+    
+    """
+
+    def __init__(self, env, target_speed: float= 8.94, t_future=0):
+        """
+        Initialize policy with pointer to environment it will run on and target speed
+
+        Args:
+            env (Intersimple): intersimple environment which IDM runs on
+            target_speed (float): target speed in roundabout (default: 8.94=20 mph)
+        """
+        assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
+        self._env = env
+
+        
+        assert(t_future >=0, 'negative target speed')
+        self.t_future = t_future
+
+        self.half_angle = 45 # degrees for finding car to follow
+
+        # Default IDM parameters
+        assert(target_speed>0, 'negative target speed')
+        self.s_max = target_speed
+        self.a_max = np.array([3.]) # nominal acceleration
+        self.tau = 0.5 # desired time headway
+        self.b_pref = 2.5 # preferred deceleration
+        self.d_min = 1 #minimum spacing
+
+    def predict(self, observation, *args, **kwargs):
+        """
+        Generate action, state from observation
+
+        (But actually generate next action from underlying environment state)
+        
+        Args:
+            observation (np.ndarray): instantaneous observation from environment
+
+        Returns
+            action (np.ndarray): action for controlled agent to take
+            state (np.ndarray): the index of the chosen vehicle for IDM
+        """
+        agent = self._env._agent
+        full_state = self._env._env.projected_state.numpy() #(nv, 5)
+        ego_state = full_state[agent] # (5,)
+        s = ego_state[2]
+        xy = full_state[:,0:2] # (nv, 2)
+        v = full_state[:,2:3] # (nv, 1)
+        psi = full_state[:,2:3] # (nv, 1)
+
+
+        d, r, i = self.get_ego_dr(agent, xy, v, psi)
+
+        # propagate environment forward at constant velocity
+        if self.t_future > 0:
+            xy2 = xy + self.t_future * v * np.vstack((np.cos(psi[:,0]), np.sin(psi[:,0])))
+            d2, r2, i2 = self.get_ego_dr(agent, xy2, v, psi)
+            
+            # choose closer vehicle (now vs imagined)
+            if d2 < d:
+                d, r, i = d2, r2, i2   
+        
+        if d == np.inf:
+            d_des = 0
+        else:
+            d_des = self.d_min + self.tau * s + s * r / (2* (self.a_max*self.b_pref)**0.5 )
+        action = self.a_max*(1 - (s/self.s_max)**4 - (d_des/d)**2)
+
+        assert(action.shape==(1,))
+        return action, i
+    
+    def get_ego_dr(self, agent:int, xy: np.ndarray, v: np.ndarray, psi: np.ndarray):
+        """
+        Return distance and relative speed of closest car within half angle from heading
+        
+        Args:
+            agent (int): agent index
+            xy (np.ndarray): (nv, 2) x and y positions
+            v (np.ndarray): (nv, 1) velocity
+            psi (np.ndarray): (nv, 1) heading angle
+
+        Returns:
+            d (float): distance to closest vehicle in cone
+            r (float): relative speed between the two vehicles
+            i (Union[None,np.ndarray]): (1,) array of closest vehicle index, or None 
+        """
+        nv, nxy = xy.shape
+        nv2, nvel = v.shape
+        nv3, npsi = psi.shape
+        assert(nv==nv2==nv3)
+        assert(nxy==2)
+        assert(nvel==npsi==1)
+
+        dxys = xy - xy[agent] # (nv, 2)
+        ds = np.linalg.norm(dxys,axis=1) # (nv,)
+        df = (dxys*np.hstack((np.cos(psi),np.sin(psi)))).sum(-1) # (nv, )
+        dl = (dxys*np.hstack((-np.sin(psi), np.cos(psi)))).sum(-1) # (nv, )
+        alpha = to_circle(np.arctan2(dl, df))
+
+        val_idx = np.arange(nv)[(np.abs(alpha) < self.half_angle*np.pi/180 & np.arange(nv) != agent)]
+        
+        if len(val_idx)==0:
+            i = None
+            d = float('inf')
+            r = float('inf')
+        else:
+            idx = np.argmin(ds[val_idx]) # closest car which meets requirements
+            i = val_idx[idx]
+            d = ds[i]
+            r = v[i,0]-v[agent,0] 
+            
+        return d, r, i
+    
+def to_circle(x: np.ndarray) -> np.ndarray:
+    """
+    Casts x (in rad) to [-pi, pi)
+
+    Args:
+        x (np.ndarray): (*) input angle (radians)
+
+    Returns:
+        y (np.ndarray): (*) x cast to [-pi, pi)
+    """
+    y = np.remainder(x + np.pi, 2*np.pi) - np.pi
+    return y
\ No newline at end of file

From 31912416f1fe1bacab5eaaa6da7d86ef335e5331 Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Wed, 2 Feb 2022 22:19:40 -0800
Subject: [PATCH 3/6] adding pbar to evaluator and making metric save optional,
 adding typing to baselines

---
 src/baselines/__init__.py      |  1 +
 src/baselines/rule_policies.py | 17 ++++++++++++-----
 src/evaluation/__init__.py     |  1 +
 src/evaluation/evaluation.py   | 27 +++++++++++++++++++--------
 4 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/src/baselines/__init__.py b/src/baselines/__init__.py
index e69de29..ea73cdf 100644
--- a/src/baselines/__init__.py
+++ b/src/baselines/__init__.py
@@ -0,0 +1 @@
+from src.baselines.rule_policies import IDMRulePolicy, PControllerPolicy
\ No newline at end of file
diff --git a/src/baselines/rule_policies.py b/src/baselines/rule_policies.py
index 2f395d7..fb2b271 100644
--- a/src/baselines/rule_policies.py
+++ b/src/baselines/rule_policies.py
@@ -1,5 +1,6 @@
 from stable_baselines3.common.base_class import BaseAlgorithm
 from intersim.envs.intersimple import Intersimple
+from typing import Tuple, Optional
 import numpy as np
 
 class PControllerPolicy(BaseAlgorithm):
@@ -14,7 +15,7 @@ class PControllerPolicy(BaseAlgorithm):
         self.target_v = 8.94 # m/s
         self.attn_weight = 20
     
-    def predict(self, observation, *args, **kwargs):
+    def predict(self, observation: np.ndarray, *args, **kwargs):
         """
         Generate action, state from observation
 
@@ -58,15 +59,16 @@ class IDMRulePolicy(BaseAlgorithm):
     
     """
 
-    def __init__(self, env, target_speed: float= 8.94, t_future=0):
+    def __init__(self, env: Intersimple, target_speed: float= 8.94, t_future:float=0):
         """
         Initialize policy with pointer to environment it will run on and target speed
 
         Args:
             env (Intersimple): intersimple environment which IDM runs on
             target_speed (float): target speed in roundabout (default: 8.94=20 mph)
+            t_future (float): future time at which to compare closest
         """
-        assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
+        # assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
         self._env = env
 
         
@@ -82,8 +84,10 @@ class IDMRulePolicy(BaseAlgorithm):
         self.tau = 0.5 # desired time headway
         self.b_pref = 2.5 # preferred deceleration
         self.d_min = 1 #minimum spacing
+        super().__init__()
 
-    def predict(self, observation, *args, **kwargs):
+    def predict(self, observation:np.ndarray, 
+        *args, **kwargs) -> Tuple[np.ndarray,Optional[np.ndarray]]:
         """
         Generate action, state from observation
 
@@ -96,6 +100,8 @@ class IDMRulePolicy(BaseAlgorithm):
             action (np.ndarray): action for controlled agent to take
             state (np.ndarray): the index of the chosen vehicle for IDM
         """
+        import pdb
+        pdb.set_trace()
         agent = self._env._agent
         full_state = self._env._env.projected_state.numpy() #(nv, 5)
         ego_state = full_state[agent] # (5,)
@@ -125,7 +131,8 @@ class IDMRulePolicy(BaseAlgorithm):
         assert(action.shape==(1,))
         return action, i
     
-    def get_ego_dr(self, agent:int, xy: np.ndarray, v: np.ndarray, psi: np.ndarray):
+    def get_ego_dr(self, agent:int, xy: np.ndarray, 
+        v: np.ndarray, psi: np.ndarray) -> Tuple[Optional[np.ndarray], float, float]:
         """
         Return distance and relative speed of closest car within half angle from heading
         
diff --git a/src/evaluation/__init__.py b/src/evaluation/__init__.py
index e69de29..3410d72 100644
--- a/src/evaluation/__init__.py
+++ b/src/evaluation/__init__.py
@@ -0,0 +1 @@
+from src.evaluation.evaluation import IntersimpleEvaluation
diff --git a/src/evaluation/evaluation.py b/src/evaluation/evaluation.py
index de1306c..e047afe 100644
--- a/src/evaluation/evaluation.py
+++ b/src/evaluation/evaluation.py
@@ -2,9 +2,10 @@ 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 typing import Callable, Dict
+from typing import Callable, Dict, Optional
 import os
 import pickle
+from tqdm import tqdm
 
 class IntersimpleEvaluation:
     """
@@ -18,12 +19,13 @@ class IntersimpleEvaluation:
         - whether there was a collision
         - whether there was a hard brake
     """
-    def __init__(self, eval_env):
+    def __init__(self, eval_env, use_pbar:bool=True):
         """
         Initialize evaluation environment with an Intersimple IncrementingAgent environment
 
         Args:
             eval_env (Intersimple.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!
@@ -34,6 +36,7 @@ class IntersimpleEvaluation:
 
         self.env = eval_env
         self.n_episodes = eval_env.nv
+        self.use_pbar = use_pbar
 
         # metrics present on every step of every episode
         self.metric_keys_all = ['v_all', 'a_all', 'col_all']
@@ -42,7 +45,7 @@ class IntersimpleEvaluation:
         self.metric_keys_single = ['j_all', 'v_avg','a_avg', 'col','brake', 't']
         
         # numbers for calculating metrics
-        self.hard_brake = -3 # acceleration for 'hard brake'
+        self.hard_brake = -3. # acceleration for 'hard brake'
 
         # reset metrics
         self.reset()
@@ -73,17 +76,18 @@ class IntersimpleEvaluation:
         with open(filestr, 'wb') as f:
             pickle.dump(self._metrics, f)     
 
-    def evaluate(self, policy, filestr: str) -> Dict[str, list]:
+    def evaluate(self, policy, filestr: Optional[str] = None) -> Dict[str, list]:
         """
         Evaluate a policy on the incrementing agent evaluation environment
         
         Args:
             policy (BaseClass.BaseAlgorithm): policy in which policy.predict(observation)[0] returns an action
-            filestr (str): path-like string to dump metrics to
+            filestr (str): path-like string to dump metrics to or None
         """
         self.reset()
-        metrics = {}
-        
+        if self.use_pbar:
+            self.pbar = tqdm(total=self.n_episodes)
+
         evaluate_policy(
             policy, 
             self.env,
@@ -91,8 +95,12 @@ class IntersimpleEvaluation:
             callback=self.evaluate_policy_callback, 
             return_episode_rewards=False
         )
+        if self.use_pbar:
+            self.pbar.close()
+            
         self.post_proc()
-        self.save(filestr)
+        if filestr:
+            self.save(filestr)
         return self._metrics
 
     def evaluate_policy_callback(self, local_vars, global_vars):
@@ -114,6 +122,9 @@ class IntersimpleEvaluation:
         if col:
             assert done
         self._metrics['col_all'][_agent].append(col)
+        
+        if done and self.use_pbar:
+            self.pbar.update(1)
     
     def post_proc(self):
         """

From d1f9e3d7c435e070be192f6225e0b2eb9d3f2b37 Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Wed, 2 Feb 2022 22:29:05 -0800
Subject: [PATCH 4/6] adding main test sequence. must debug and add summary and
 comparison metric generators tomorrow

---
 scratch/arec/intersimple/test_model.py | 273 +++++++++++++++++--------
 1 file changed, 187 insertions(+), 86 deletions(-)

diff --git a/scratch/arec/intersimple/test_model.py b/scratch/arec/intersimple/test_model.py
index 8ae5f2d..9934c4c 100644
--- a/scratch/arec/intersimple/test_model.py
+++ b/scratch/arec/intersimple/test_model.py
@@ -2,147 +2,248 @@ from tqdm import tqdm
 from copy import deepcopy
 import stable_baselines3 as sb3
 import intersim
+from stable_baselines3.common.base_class import BaseAlgorithm
+from src.baselines import IDMRulePolicy
+from src.evaluation import IntersimpleEvaluation
+import src.options.envs as options_envs
 
-ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
+from typing import Optional, List, Dict, Tuple
+import torch
+import numpy as np
 
-def load_model(model_path:str, method:str):
+def load_policy(policy_file:str, method:str, 
+    policy_kwargs:dict, skip_load:bool=False) -> Optional[BaseAlgorithm]:
     """
     Load a model given a path and the method
     
     Args:
-        model_path (str): the path to the model
+        load_policy (str): the path to the model
         method (str): the method for the model
+        skip_load (bool): whether to skip loading
     Returns:
-        model: the action model
-        is_heir (bool): whether the method is heirarchial
+        policy (Optional[BaseAlgorithm]): the policy to evaluate
     """
-    model = None
-    is_heir = False
-    if method == 'expert':
-        raise NotImplementedError
+    if skip_load:
+        return None
+    elif method == 'idm':
+        policy = IDMRulePolicy(policy_kwargs)
     elif method == 'bc':
         raise NotImplementedError
     elif method == 'gail':
+        policy = sb3.PPO.load(policy_file)
         raise NotImplementedError
     elif method == 'rail':
         raise NotImplementedError
-    elif method == 'hgail':
-        is_heir = True
-        model = sb3.PPO.load(model_path)
-    elif method == 'hrail':
-        is_heir = True
+    elif method == 'sgail':
+        policy = sb3.PPO.load(policy_file)
         raise NotImplementedError
     else:
         raise NotImplementedError
-    return model, is_heir
+    return policy
 
-def load_expert_states(roundabout, track):
+def form_expert_metrics(states:torch.Tensor, actions:torch.Tensor) ->Dict[str, list]:
+    """
+    Given experts of tensor states and actions, form a dictionary of metrics
+
+    Args:
+        states (torch.tensor): (T+1, nv, 5) expert states for track file
+        actions (torch.tensor): (T, nv, 1) expert actions for track file
+    
+    Returns:
+        metrics (dict): dictionary maping strings to lists
+    """
+
+    T1, nv, _ = states.shape
+    T, nv2, _ = actions.shape
+    assert(nv==nv2)
+    assert(T1==T+1)
+    states = states[:T]
+
+    # make sure metric keys and calculations match that in src.evaluation.IntersimpleEvaluation
+    
+    hard_brake = -3.
+    timestep = 0.1
+    keys = ['col_all','v_all', 'a_all','j_all', 'v_avg', 'a_avg', 'col', 'brake', 't']
+    metrics = {key:[None]*nv for key in keys}
+    
+    import pdb
+    pdb.set_trace()
+    
+    for i in range(nv):
+        
+        nni = ~torch.isnan(states[:,i,0])
+
+        metrics['col_all'][i] = [False] * sum(nni) 
+        metrics['v_all'][i] = states[nni,i,2].numpy()
+        metrics['a_all'][i] = actions[nni,i,0].numpy()
+        
+        # jerk
+        metrics['j_all'][i] = np.diff(metrics['a_all'][i]) / timestep
+
+        # average velocity and acceleration
+        metrics['v_avg'][i] = np.mean(metrics['v_all'][i])
+        metrics['a_avg'][i] = np.mean(metrics['a_all'][i])
+
+        # collision?
+        metrics['col'][i] = any(metrics['col_all'][i]) 
+
+        # brake?
+        metrics['brake'][i] = any(metrics['a_all'][i] < hard_brake)
+
+        # time length 
+        metrics['t'][i] = sum(nni) 
+
+    for key in metrics.keys():
+        assert(len(metrics[key])==nv)
+    return metrics
+
+def generate_expert_metrics(locations: List[Tuple[int,int]]) -> List[Dict[str, list]]:
+    """"
+    Given a list of locations, for and return a list of metrics for each location
+
+    Args:
+        locations (list): list of (roundabout, track) ints
+
+    Returns:
+        expert_metrics (list of dicts): expert_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith expert roundabout trackfile
+    """
+    expert_metrics = []
+    for (roundabout, track) in locations:
+        states, actions = load_expert_states(roundabout, track)
+        expert_metrics.append(form_expert_metrics(states, actions))
+    return expert_metrics
+
+
+def load_expert_states(roundabout:int, track:int):
     """
     Load expert states from roundabout/track info
     Args:
-        roundabout (str): roundabout name
-        track (str): track id
+        roundabout (int): roundabout index
+        track (int): track id
     Returns:
         states (torch.tensor): (T+1, nv, 5) expert states for track file
         actions (torch.tensor): (T, nv, 1) expert actions for track file
     """
-    state_path = '../../../expert_data/%s/track%04i/joint_expert_states.pt'%(roundabout, track)] #FIXME when moving
-    action_path = '../../../expert_data/%s/track%04i/joint_expert_actions.pt'%(roundabout, track)] #FIXME when moving
-    states = torch.load(path)
-    actions = torch.load(path)
-    # nanify actions where vehicle's don't exist
+    rname = intersim.LOCATIONS[roundabout]
     import pdb
     pdb.set_trace()
+    state_path = 'expert_data/%s/track%04i/joint_expert_states.pt'%(rname, track)
+    action_path = 'expert_data/%s/track%04i/joint_expert_actions.pt'%(rname, track)
+    states = torch.load(state_path)
+    actions = torch.load(action_path)
     return states, actions
 
-def test_model(
-    locations=[(0,0)], 
-    model_name='gail_image_multiagent_nocollision', 
-    env='NRasterizedRouteIncrementingAgent', 
-    method='expert',
-    options_list=ALL_OPTIONS, 
-    **env_kwargs):
+def evaluate_policy(policy:BaseAlgorithm, locations:List[Tuple[int,int]], 
+    env_class:str, env_kwargs:dict) -> List[Dict[str,list]]:
     """
-    Test a particular model at different locations/tracks
+    Evaluate policy on an incrementing agent environment at all locations. 
+    Return metrics for that policy
 
     Args:
-        locations (list of tuples): list of (roundabout, track) integer pairs
-        model_name (str): name of model to test
-        env (str): environment class
-        method (str): method (expert, bc, gail, rail, hgail, hrail)
-        options_list (list): list of options
+        policy (BaseAlgorithm): policy to evaluate
+        locations (list of tuples): list of locations to evaluate policy
+        env_class (str): name of environment to evaluate policy with
+        env_kwargs (dict): key word arguments to initialize environment with
+    
+    Returns:
+        policy_metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
     """
-
-    # load policy
-    policy, is_heir = load_model(model_name, method)
+    envs_dict = deepcopy(intersim.envs.intersimple.__dict__)
+    envs_dict.update(deepcopy(options_envs.__dict__))
+    
+    policy_metrics = [None]* len(locations)
 
     # iterate through vehicles
-    all_vehicle_infos = []
     for i, location in tqdm(enumerate(locations)):
         
         # add roundabout and track to environent 
-        roundabout, track = location
-        iround = intersim.LOCATIONS.index(roundabout)
+        iround, track = location
+        rname = intersim.LOCATIONS[iround]
         it_env_kwargs = deepcopy(env_kwargs)
         loc_kwargs = {
             'loc':iround,
             'track':track
         }
         it_env_kwargs.update(loc_kwargs)
-        
-        # load expert states and get average velocities
-        expert_states, expert_actions = load_expert_states(roundabout, track)
-        expert_vavg = torch.nanmean(expert_states[:,:,3], dim=-1)
 
         # initialize environment
-        if not is_heir:
-            Env = src.options.envs.__dict__[env]
-        else: 
-            Env = intersim.envs.intersimple.__dict__[env]
-        env = Env(**env_kwargs)
-        s = env.reset()
+        Env = envs_dict[env_class]
+        eval_env = Env(**env_kwargs)
+        evaluator = IntersimpleEvaluation(eval_env)
+        policy_metrics[i] = evaluator.evaluate(policy)
+    
+    return policy_metrics
+        
 
-        # Iterate through every vehicle and time
-        vehicle_infos, done = [], False
-        for iv in range(env.nv):
-            v_number = env.agent
-            i_vehicle_infos = {'s':[], 'a':[], 'it':[]}
-            while not done:
-                a = policy(s)
-                sp, r, done, info   = env.step(a)
-                i_vehicle_infos['s'].append(env._env.state) # FIX
-                i_vehicle_infos['a'].append(a) 
-                i_vehicle_infos['it'].append(env._env.it) # FIX
-            i_vehicle_info.update({
-                'vehicle_id': env.agent,
-                'n_steps': len(i_vehicle_infos['a']),
-                'T': len(i_vehicle_infos['a'])*env._env.dt, # FIX
-                'n_collisions': collision.check(i_vehicle_infos['s'], env._env.lengths. env._env.widths), # FIX
-                'expert_vavg': expert_vavg[env.agent]
-            })
-            vehicle_infos.append(i_vehicle_info)
-            env.reset()
 
-        all_vehicle_infos.append({
-            'loc': location,
-            'track': track,
-            'stats': vehicle_infos
-            })
-        env.close()
-
-    # print and save model-specific metrics
-    outfolder = 'test_metrics'
-    print_and_save(all_vehicle_infos, method, model, outfolder)
-
-def print_and_save(stats, method, model, outfolder):
+def summary_metrics(metrics:List[Dict[str,list]]):
     """
-    Print and save stats
+    Summarize and print metrics averaged over vehicles and roundabouts
+
+    Args:
+        metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
     """
     pass
 
-def load_compare():
+def comparison_metrics(policy_metrics:List[Dict[str,list]], expert_metrics:List[Dict[str,list]]):
+    """
+    Provide distributional comparison between different sets of metrics
+
+    Args:
+        policy_metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
+        expert_metrics (list of dicts): expert_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith expert roundabout trackfile
+    
+    """
     pass
 
+def test_model(
+    locations: List[Tuple[int,int]]= [(0,0)],
+    method: str='expert', 
+    policy_file: str='', 
+    policy_kwargs: dict={},
+    env: str='NRasterizedRouteIncrementingAgent', 
+    env_kwargs: dict={},
+    seed: int=0):
+    """
+    Test a particular model at different testing locations/tracks and compute average metrics 
+    over all files. 
+
+    Args:
+        locations (list of tuples): list of (roundabout, track) integer pair testing locations
+        method (str): method string
+        policy_file (str): path to saved policy
+        env (str): environment class
+        method (str): method (expert, bc, gail, rail, hgail, hrail)
+    """
+    # set seed
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    # load policy
+    policy = load_policy(policy_file, method, policy_kwargs, skip_load=(method=='expert'))
+    
+    # load expert metrics
+    expert_metrics = generate_expert_metrics(locations)
+
+    # if we have a policy
+    if policy:
+
+        # evaluate it on the given roundabouts
+        policy_metrics = evaluate_policy(policy, locations, env, env_kwargs)
+
+        # 
+        summary_metrics(policy_metrics)
+        comparison_metrics(policy_metrics, expert_metrics)
+    
+    else:
+        # if no policy, only generate summary metrics for the expert
+        summary_metrics(expert_metrics)
+
 if __name__=='__main__':
     import fire
     fire.Fire()
\ No newline at end of file

From 795e1c08b602468d07880eaf4cfcd8970b96a3f2 Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Fri, 4 Feb 2022 15:51:17 -0800
Subject: [PATCH 5/6] adding metric comparisons and updating (note: pre-debug)
 init

---
 .../test_model.py => src/eval_main.py         | 73 ++++++++++++++++++-
 src/evaluation/__init__.py                    |  1 +
 2 files changed, 72 insertions(+), 2 deletions(-)
 rename scratch/arec/intersimple/test_model.py => src/eval_main.py (72%)

diff --git a/scratch/arec/intersimple/test_model.py b/src/eval_main.py
similarity index 72%
rename from scratch/arec/intersimple/test_model.py
rename to src/eval_main.py
index 9934c4c..d1872c5 100644
--- a/scratch/arec/intersimple/test_model.py
+++ b/src/eval_main.py
@@ -6,6 +6,7 @@ from stable_baselines3.common.base_class import BaseAlgorithm
 from src.baselines import IDMRulePolicy
 from src.evaluation import IntersimpleEvaluation
 import src.options.envs as options_envs
+from src.evaluation.metrics import divergence, visualize_distribution
 
 from typing import Optional, List, Dict, Tuple
 import torch
@@ -186,7 +187,46 @@ def summary_metrics(metrics:List[Dict[str,list]]):
         metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j' 
             evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
     """
-    pass
+    # keys = ['col_all','v_all', 'a_all','j_all', 'v_avg', 'a_avg', 'col', 'brake', 't']
+    import pdb
+    pdb.set_trace()
+
+    # average average-velocity
+    all_vavgs =  sum([d['v_avg'] for d in metrics],[]) # aggregate to single list
+    mean_vavg = sum(all_vavgs)/len(all_vavgs)
+    print(f'Mean Average Velocity: {mean_vavg}')
+
+    
+    all_aalls = np.concatenate([np.concatenate(d['a_all']) for d in metrics])
+    # average +acceleration
+    pos_accels = all_aalls[all_aalls>0]
+    mean_pos_accels = np.mean(pos_accels)
+    print(f'Mean positive acceleration: {mean_pos_accels}')
+
+    # average deceleration
+    decels = all_aalls[all_aalls<0]
+    mean_decels = np.mean(decels)
+    print(f'Mean positive acceleration: {mean_decels}')
+
+    # average |jerk|
+    all_jerks = np.concatenate([np.concatenate(d['j_all']) for d in metrics])
+    mean_abs_jerk = np.mean(np.abs(all_jerks))
+    print(f'Mean |Jerk|: {mean_abs_jerk}')
+
+    # collision rate
+    all_collisions =  sum([d['col'] for d in metrics],[]) # aggregate to single list
+    collision_rate = sum(all_collisions)/len(all_collisions)
+    print(f'Collision Rate: {collision_rate}')
+
+    # hard brake rate
+    all_hard_brakes =  sum([d['brake'] for d in metrics],[]) # aggregate to single list
+    hard_brake_rate = sum(all_hard_brakes)/len(all_hard_brakes)
+    print(f'Hard Brake Rate: {hard_brake_rate}')
+
+    # average number of timesteps
+    all_ts = sum([d['t'] for d in metrics],[]) # aggregate to single list
+    mean_t = sum(all_ts)/len(all_ts)
+    print(f'Mean episode length: {mean_t}')
 
 def comparison_metrics(policy_metrics:List[Dict[str,list]], expert_metrics:List[Dict[str,list]]):
     """
@@ -199,7 +239,36 @@ def comparison_metrics(policy_metrics:List[Dict[str,list]], expert_metrics:List[
             evaluated on the kth episode (car) of the ith expert roundabout trackfile
     
     """
-    pass
+    import pdb
+    pdb.set_trace()
+
+    # average velocity shortfall
+    expert_vavg = np.array(sum([d['v_avg'] for d in expert_metrics],[]))
+    policy_vavg = np.array(sum([d['v_avg'] for d in policy_metrics],[]))
+    assert(len(expert_vavg)==len(policy_vavg))
+    mean_shortfall = np.mean(expert_vavg - policy_vavg)
+    print(f'Mean shortfall velocity: {mean_shortfall}')
+    
+    # velocity JSD
+    expert_vs = np.concatenate([np.concatenate(d['v_all']) for d in expert_metrics])
+    policy_vs = np.concatenate([np.concatenate(d['v_all']) for d in policy_metrics])
+    vel_div = divergence(expert_vs, policy_vs)
+    print(f'Velocity distribution divergence: {vel_div}')
+    visualize_distribution(expert_vs, policy_vs, 'velocity_jsd.png')
+    
+    # acceleration JSD
+    expert_as = np.concatenate([np.concatenate(d['a_all']) for d in expert_metrics])
+    policy_as = np.concatenate([np.concatenate(d['a_all']) for d in policy_metrics])
+    accel_div = divergence(expert_as, policy_as)
+    print(f'Velocity distribution divergence: {accel_div}')
+    visualize_distribution(expert_as, policy_as, 'accel_jsd.png')
+    
+    # jerk JSD
+    expert_jerks = np.concatenate([np.concatenate(d['j_all']) for d in expert_metrics])
+    policy_jerks = np.concatenate([np.concatenate(d['j_all']) for d in policy_metrics])
+    jerk_div = divergence(expert_jerks, policy_jerks)
+    print(f'Jerk distribution divergence: {jerk_div}')
+    visualize_distribution(expert_jerks, policy_jerks, 'jerk_jsd.png')
 
 def test_model(
     locations: List[Tuple[int,int]]= [(0,0)],
diff --git a/src/evaluation/__init__.py b/src/evaluation/__init__.py
index 3410d72..acfa054 100644
--- a/src/evaluation/__init__.py
+++ b/src/evaluation/__init__.py
@@ -1 +1,2 @@
 from src.evaluation.evaluation import IntersimpleEvaluation
+from src.evaluation.metrics import divergence, visualize_distribution
\ No newline at end of file

From 3e6fce42ee295859e33f7a9ab4f3b40a3bd1205c Mon Sep 17 00:00:00 2001
From: Arec <jamgochian95@gmail.com>
Date: Sat, 5 Feb 2022 21:48:56 -0800
Subject: [PATCH 6/6] 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

---
 evaluate_models.sh             |  16 ++++
 src/__init__.py                |   3 +-
 src/baselines/rule_policies.py | 108 ++++++++++++++--------
 src/eval_main.py               | 162 ++++++++++++++++++---------------
 src/evaluation/evaluation.py   |  24 +++--
 src/evaluation/metrics.py      |  79 ++--------------
 6 files changed, 195 insertions(+), 197 deletions(-)
 create mode 100755 evaluate_models.sh

diff --git a/evaluate_models.sh b/evaluate_models.sh
new file mode 100755
index 0000000..15ff738
--- /dev/null
+++ b/evaluate_models.sh
@@ -0,0 +1,16 @@
+# eval_main inputs
+# locations: List[Tuple[int,int]]= [(0,0)],
+# method: str='expert', 
+# policy_file: str='', 
+# policy_kwargs: dict={},
+# env: str='NRasterizedRouteIncrementingAgent', 
+# env_kwargs: dict={},
+# seed: int=0
+
+# expert
+python -m src.eval_main
+
+# idm
+python -m src.eval_main --method=idm
+
+
diff --git a/src/__init__.py b/src/__init__.py
index aebb694..08e852e 100644
--- a/src/__init__.py
+++ b/src/__init__.py
@@ -1,3 +1,2 @@
 from src.data.expert_data import generate_expert_data, load_expert_data
-from src.data.data_utils import InteractionDatasetSingleAgent
-from src.evaluation.metrics import metrics
\ No newline at end of file
+from src.data.data_utils import InteractionDatasetSingleAgent
\ No newline at end of file
diff --git a/src/baselines/rule_policies.py b/src/baselines/rule_policies.py
index fb2b271..1d5948b 100644
--- a/src/baselines/rule_policies.py
+++ b/src/baselines/rule_policies.py
@@ -1,6 +1,6 @@
 from stable_baselines3.common.base_class import BaseAlgorithm
-from intersim.envs.intersimple import Intersimple
-from typing import Tuple, Optional
+from intersim.envs.intersimple import Intersimple, NormalizedActionSpace
+from typing import Tuple, Optional, List
 import numpy as np
 
 class PControllerPolicy(BaseAlgorithm):
@@ -10,11 +10,17 @@ class PControllerPolicy(BaseAlgorithm):
         """
         Initialize policy with pointer to environment it will run on
         """
-        assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
+        assert isinstance(env, Intersimple), 'Environment is not an intersimple environment'
         self._env = env
         self.target_v = 8.94 # m/s
         self.attn_weight = 20
     
+    # BaseAlgorithm abstract methods
+    def _setup_model(self): 
+        return None
+    def learn(self, *args, **kwargs):
+        return self
+
     def predict(self, observation: np.ndarray, *args, **kwargs):
         """
         Generate action, state from observation
@@ -59,37 +65,45 @@ class IDMRulePolicy(BaseAlgorithm):
     
     """
 
-    def __init__(self, env: Intersimple, target_speed: float= 8.94, t_future:float=0):
+    def __init__(self, env: Intersimple, 
+        target_speed:float= 8.94, 
+        t_future:List[float]=[0., 1., 2., 3.], 
+        half_angle:float=60.):
         """
         Initialize policy with pointer to environment it will run on and target speed
 
         Args:
             env (Intersimple): intersimple environment which IDM runs on
             target_speed (float): target speed in roundabout (default: 8.94=20 mph)
-            t_future (float): future time at which to compare closest
+            t_future (List[float]): list of future time at which to compare closest vehicle
+            half_angle (float): half angle to look inside for closest vehicle
         """
-        # assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
+
         self._env = env
-
-        
-        assert(t_future >=0, 'negative target speed')
         self.t_future = t_future
-
-        self.half_angle = 45 # degrees for finding car to follow
+        self.half_angle = half_angle
 
         # Default IDM parameters
-        assert(target_speed>0, 'negative target speed')
+        assert target_speed>0, 'negative target speed'
         self.s_max = target_speed
         self.a_max = np.array([3.]) # nominal acceleration
         self.tau = 0.5 # desired time headway
         self.b_pref = 2.5 # preferred deceleration
         self.d_min = 1 #minimum spacing
-        super().__init__()
 
+        # for np.remainder nan warnings
+        np.seterr(invalid='ignore')
+
+    # BaseAlgorithm abstract methods
+    def _setup_model(self): 
+        return None
+    def learn(self, *args, **kwargs):
+        return self
+    
     def predict(self, observation:np.ndarray, 
-        *args, **kwargs) -> Tuple[np.ndarray,Optional[np.ndarray]]:
+        *args, **kwargs) -> Tuple[np.ndarray, None]:
         """
-        Generate action, state from observation
+        Predict action, state from observation
 
         (But actually generate next action from underlying environment state)
         
@@ -98,41 +112,59 @@ class IDMRulePolicy(BaseAlgorithm):
 
         Returns
             action (np.ndarray): action for controlled agent to take
-            state (np.ndarray): the index of the chosen vehicle for IDM
+            state (None): None (hidden state for a recurrent policy)
+        """
+        return self.forward(observation, *args, **kwargs), None
+
+    def forward(self, *args, **kwargs) -> np.ndarray:
+        """
+        Generate action from underlying environment
+
+        Returns
+            action (np.ndarray): action for controlled agent to take
         """
-        import pdb
-        pdb.set_trace()
         agent = self._env._agent
         full_state = self._env._env.projected_state.numpy() #(nv, 5)
         ego_state = full_state[agent] # (5,)
         s = ego_state[2]
         xy = full_state[:,0:2] # (nv, 2)
         v = full_state[:,2:3] # (nv, 1)
-        psi = full_state[:,2:3] # (nv, 1)
-
+        psi = full_state[:,3:4] # (nv, 1)
 
         d, r, i = self.get_ego_dr(agent, xy, v, psi)
 
         # propagate environment forward at constant velocity
-        if self.t_future > 0:
-            xy2 = xy + self.t_future * v * np.vstack((np.cos(psi[:,0]), np.sin(psi[:,0])))
-            d2, r2, i2 = self.get_ego_dr(agent, xy2, v, psi)
-            
-            # choose closer vehicle (now vs imagined)
-            if d2 < d:
-                d, r, i = d2, r2, i2   
+        for t in self.t_future:
+            if t > 0:
+                xy2 = xy + t * v * np.vstack((np.cos(psi[:,0]), np.sin(psi[:,0]))).T
+                d2, r2, i2 = self.get_ego_dr(agent, xy2, v, psi)
+                
+                # choose closer vehicle (now vs imagined)
+                if d2 < d:
+                    d, r, i = d2, r2, i2
         
+        # Update environment interaction graph with i
+        if i:
+            self._env._env._graph._neighbor_dict={agent:[i]}
+
         if d == np.inf:
-            d_des = 0
+            d_des = self.d_min
         else:
             d_des = self.d_min + self.tau * s + s * r / (2* (self.a_max*self.b_pref)**0.5 )
-        action = self.a_max*(1 - (s/self.s_max)**4 - (d_des/d)**2)
+            d_des = max(d_des, self.d_min)
 
-        assert(action.shape==(1,))
-        return action, i
+        assert (d_des>= self.d_min)
+        action = self.a_max*(1 - (s/self.s_max)**4 - (d_des/d)**2)
+        
+        # normalize action to range if env is a NormalizedActionSpace
+        if isinstance(self._env, NormalizedActionSpace):
+            action = self._env._normalize(action)
+
+        assert action.shape==(1,)
+        return action
     
     def get_ego_dr(self, agent:int, xy: np.ndarray, 
-        v: np.ndarray, psi: np.ndarray) -> Tuple[Optional[np.ndarray], float, float]:
+        v: np.ndarray, psi: np.ndarray) -> Tuple[float, float, Optional[int]]:
         """
         Return distance and relative speed of closest car within half angle from heading
         
@@ -145,14 +177,14 @@ class IDMRulePolicy(BaseAlgorithm):
         Returns:
             d (float): distance to closest vehicle in cone
             r (float): relative speed between the two vehicles
-            i (Union[None,np.ndarray]): (1,) array of closest vehicle index, or None 
+            i (Optional[int]): index of closest vehicle, or None 
         """
         nv, nxy = xy.shape
         nv2, nvel = v.shape
         nv3, npsi = psi.shape
-        assert(nv==nv2==nv3)
-        assert(nxy==2)
-        assert(nvel==npsi==1)
+        assert nv==nv2==nv3
+        assert nxy==2
+        assert nvel==npsi==1
 
         dxys = xy - xy[agent] # (nv, 2)
         ds = np.linalg.norm(dxys,axis=1) # (nv,)
@@ -160,7 +192,7 @@ class IDMRulePolicy(BaseAlgorithm):
         dl = (dxys*np.hstack((-np.sin(psi), np.cos(psi)))).sum(-1) # (nv, )
         alpha = to_circle(np.arctan2(dl, df))
 
-        val_idx = np.arange(nv)[(np.abs(alpha) < self.half_angle*np.pi/180 & np.arange(nv) != agent)]
+        val_idx = np.arange(nv)[(np.abs(alpha) < self.half_angle*np.pi/180) & (np.arange(nv) != agent)]
         
         if len(val_idx)==0:
             i = None
@@ -168,7 +200,7 @@ class IDMRulePolicy(BaseAlgorithm):
             r = float('inf')
         else:
             idx = np.argmin(ds[val_idx]) # closest car which meets requirements
-            i = val_idx[idx]
+            i = int(val_idx[idx])
             d = ds[i]
             r = v[i,0]-v[agent,0] 
             
diff --git a/src/eval_main.py b/src/eval_main.py
index d1872c5..edef877 100644
--- a/src/eval_main.py
+++ b/src/eval_main.py
@@ -2,18 +2,23 @@ from tqdm import tqdm
 from copy import deepcopy
 import stable_baselines3 as sb3
 import intersim
+from intersim.envs import Intersimple
 from stable_baselines3.common.base_class import BaseAlgorithm
 from src.baselines import IDMRulePolicy
 from src.evaluation import IntersimpleEvaluation
-import src.options.envs as options_envs
+import src.gail.options as options_envs
 from src.evaluation.metrics import divergence, visualize_distribution
 
 from typing import Optional, List, Dict, Tuple
 import torch
 import numpy as np
 
-def load_policy(policy_file:str, method:str, 
-    policy_kwargs:dict, skip_load:bool=False) -> Optional[BaseAlgorithm]:
+#(method, policy_file, policy_kwargs, eval_env)
+
+def load_policy(method:str, 
+    policy_file:str, 
+    policy_kwargs:dict, 
+    env: Intersimple) -> BaseAlgorithm:
     """
     Load a model given a path and the method
     
@@ -21,13 +26,12 @@ def load_policy(policy_file:str, method:str,
         load_policy (str): the path to the model
         method (str): the method for the model
         skip_load (bool): whether to skip loading
+        env (Intersimple): Intersimple environment for evaluation (necessary for IDM policy)
     Returns:
         policy (Optional[BaseAlgorithm]): the policy to evaluate
     """
-    if skip_load:
-        return None
-    elif method == 'idm':
-        policy = IDMRulePolicy(policy_kwargs)
+    if method == 'idm':
+        policy = IDMRulePolicy(env, **policy_kwargs)
     elif method == 'bc':
         raise NotImplementedError
     elif method == 'gail':
@@ -67,9 +71,6 @@ def form_expert_metrics(states:torch.Tensor, actions:torch.Tensor) ->Dict[str, l
     keys = ['col_all','v_all', 'a_all','j_all', 'v_avg', 'a_avg', 'col', 'brake', 't']
     metrics = {key:[None]*nv for key in keys}
     
-    import pdb
-    pdb.set_trace()
-    
     for i in range(nv):
         
         nni = ~torch.isnan(states[:,i,0])
@@ -115,7 +116,6 @@ def generate_expert_metrics(locations: List[Tuple[int,int]]) -> List[Dict[str, l
         expert_metrics.append(form_expert_metrics(states, actions))
     return expert_metrics
 
-
 def load_expert_states(roundabout:int, track:int):
     """
     Load expert states from roundabout/track info
@@ -127,16 +127,18 @@ def load_expert_states(roundabout:int, track:int):
         actions (torch.tensor): (T, nv, 1) expert actions for track file
     """
     rname = intersim.LOCATIONS[roundabout]
-    import pdb
-    pdb.set_trace()
     state_path = 'expert_data/%s/track%04i/joint_expert_states.pt'%(rname, track)
     action_path = 'expert_data/%s/track%04i/joint_expert_actions.pt'%(rname, track)
     states = torch.load(state_path)
     actions = torch.load(action_path)
     return states, actions
 
-def evaluate_policy(policy:BaseAlgorithm, locations:List[Tuple[int,int]], 
-    env_class:str, env_kwargs:dict) -> List[Dict[str,list]]:
+def evaluate_policy(locations:List[Tuple[int,int]], 
+    env_class:str, 
+    env_kwargs:dict,
+    method: str,
+    policy_file: str,
+    policy_kwargs:dict) -> List[Dict[str,list]]:
     """
     Evaluate policy on an incrementing agent environment at all locations. 
     Return metrics for that policy
@@ -146,14 +148,16 @@ def evaluate_policy(policy:BaseAlgorithm, locations:List[Tuple[int,int]],
         locations (list of tuples): list of locations to evaluate policy
         env_class (str): name of environment to evaluate policy with
         env_kwargs (dict): key word arguments to initialize environment with
+        method (str): policy method
+        policy_file (str): policy file path
+        policy_kwargs (dict): policy kwargs
     
     Returns:
         policy_metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j' 
             evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
     """
-    envs_dict = deepcopy(intersim.envs.intersimple.__dict__)
-    envs_dict.update(deepcopy(options_envs.__dict__))
-    
+    envs_dict = dict(intersim.envs.intersimple.__dict__)
+    envs_dict.update(dict(options_envs.__dict__))
     policy_metrics = [None]* len(locations)
 
     # iterate through vehicles
@@ -173,62 +177,71 @@ def evaluate_policy(policy:BaseAlgorithm, locations:List[Tuple[int,int]],
         Env = envs_dict[env_class]
         eval_env = Env(**env_kwargs)
         evaluator = IntersimpleEvaluation(eval_env)
+
+        # load policy
+        policy = load_policy(method, policy_file, policy_kwargs, eval_env)
+
+        # run policy on environment
         policy_metrics[i] = evaluator.evaluate(policy)
     
     return policy_metrics
         
-
-
-def summary_metrics(metrics:List[Dict[str,list]]):
+def summary_metrics(metrics:List[Dict[str,list]]) -> Dict[str,float]:
     """
     Summarize and print metrics averaged over vehicles and roundabouts
 
     Args:
         metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j' 
             evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
+    
+    Returns:
+        summary_metrics (Dict[str,float]): maps summary metric descriptions to values
     """
     # keys = ['col_all','v_all', 'a_all','j_all', 'v_avg', 'a_avg', 'col', 'brake', 't']
-    import pdb
-    pdb.set_trace()
+    summary_metrics = {}
 
     # average average-velocity
     all_vavgs =  sum([d['v_avg'] for d in metrics],[]) # aggregate to single list
-    mean_vavg = sum(all_vavgs)/len(all_vavgs)
-    print(f'Mean Average Velocity: {mean_vavg}')
-
+    summary_metrics['mean average velocity'] = sum(all_vavgs)/len(all_vavgs)
     
+    # average acceleration
     all_aalls = np.concatenate([np.concatenate(d['a_all']) for d in metrics])
+    summary_metrics['mean acceleartion'] = np.mean(all_aalls)
+
     # average +acceleration
     pos_accels = all_aalls[all_aalls>0]
-    mean_pos_accels = np.mean(pos_accels)
-    print(f'Mean positive acceleration: {mean_pos_accels}')
+    summary_metrics['mean positive acceleration'] = np.mean(pos_accels)
 
     # average deceleration
     decels = all_aalls[all_aalls<0]
-    mean_decels = np.mean(decels)
-    print(f'Mean positive acceleration: {mean_decels}')
+    summary_metrics['mean deceleration'] = np.mean(decels)
+
+    # average jerk
+    all_jerks = np.concatenate([np.concatenate(d['j_all']) for d in metrics])
+    summary_metrics['mean jerk'] = np.mean(all_jerks)
 
     # average |jerk|
-    all_jerks = np.concatenate([np.concatenate(d['j_all']) for d in metrics])
-    mean_abs_jerk = np.mean(np.abs(all_jerks))
-    print(f'Mean |Jerk|: {mean_abs_jerk}')
+    summary_metrics['mean |jerk|'] = np.mean(np.abs(all_jerks))
 
     # collision rate
     all_collisions =  sum([d['col'] for d in metrics],[]) # aggregate to single list
-    collision_rate = sum(all_collisions)/len(all_collisions)
-    print(f'Collision Rate: {collision_rate}')
+    summary_metrics['collision rate'] = sum(all_collisions)/len(all_collisions)
 
     # hard brake rate
     all_hard_brakes =  sum([d['brake'] for d in metrics],[]) # aggregate to single list
-    hard_brake_rate = sum(all_hard_brakes)/len(all_hard_brakes)
-    print(f'Hard Brake Rate: {hard_brake_rate}')
+    summary_metrics['hard brake rate'] = sum(all_hard_brakes)/len(all_hard_brakes)
 
     # average number of timesteps
     all_ts = sum([d['t'] for d in metrics],[]) # aggregate to single list
-    mean_t = sum(all_ts)/len(all_ts)
-    print(f'Mean episode length: {mean_t}')
+    summary_metrics['mean episode length'] = sum(all_ts)/len(all_ts)
 
-def comparison_metrics(policy_metrics:List[Dict[str,list]], expert_metrics:List[Dict[str,list]]):
+    for key in summary_metrics.keys():
+        print(f'{key}: {summary_metrics[key]}')
+
+    return summary_metrics
+
+def comparison_metrics(policy_metrics:List[Dict[str,list]], 
+    expert_metrics:List[Dict[str,list]]) -> Dict[str,float]:
     """
     Provide distributional comparison between different sets of metrics
 
@@ -237,40 +250,42 @@ def comparison_metrics(policy_metrics:List[Dict[str,list]], expert_metrics:List[
             evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
         expert_metrics (list of dicts): expert_metrics[i][j][k] returns the value of metric 'j' 
             evaluated on the kth episode (car) of the ith expert roundabout trackfile
-    
+
+    Returns:   
+        comparison_metrics (Dict[str,float]): dict mapping comparison metric description to value
     """
-    import pdb
-    pdb.set_trace()
+    comparison_metrics = {}
 
     # average velocity shortfall
     expert_vavg = np.array(sum([d['v_avg'] for d in expert_metrics],[]))
     policy_vavg = np.array(sum([d['v_avg'] for d in policy_metrics],[]))
-    assert(len(expert_vavg)==len(policy_vavg))
-    mean_shortfall = np.mean(expert_vavg - policy_vavg)
-    print(f'Mean shortfall velocity: {mean_shortfall}')
+    assert len(expert_vavg)==len(policy_vavg)
+    comparison_metrics['mean shortfall velocity'] = np.mean(expert_vavg - policy_vavg)
     
     # velocity JSD
-    expert_vs = np.concatenate([np.concatenate(d['v_all']) for d in expert_metrics])
-    policy_vs = np.concatenate([np.concatenate(d['v_all']) for d in policy_metrics])
-    vel_div = divergence(expert_vs, policy_vs)
-    print(f'Velocity distribution divergence: {vel_div}')
-    visualize_distribution(expert_vs, policy_vs, 'velocity_jsd.png')
+    expert_vs = torch.tensor(np.concatenate([np.concatenate(d['v_all']) for d in expert_metrics]))
+    policy_vs = torch.tensor(np.concatenate([np.concatenate(d['v_all']) for d in policy_metrics]))
+    comparison_metrics['velocity distribution divergence'] = divergence(expert_vs, policy_vs)
+    visualize_distribution(expert_vs, policy_vs, 'velocity_jsd')
     
     # acceleration JSD
-    expert_as = np.concatenate([np.concatenate(d['a_all']) for d in expert_metrics])
-    policy_as = np.concatenate([np.concatenate(d['a_all']) for d in policy_metrics])
-    accel_div = divergence(expert_as, policy_as)
-    print(f'Velocity distribution divergence: {accel_div}')
-    visualize_distribution(expert_as, policy_as, 'accel_jsd.png')
+    expert_as = torch.tensor(np.concatenate([np.concatenate(d['a_all']) for d in expert_metrics]))
+    policy_as = torch.tensor(np.concatenate([np.concatenate(d['a_all']) for d in policy_metrics]))
+    comparison_metrics['acceleration distribution divergence'] = divergence(expert_as, policy_as)
+    visualize_distribution(expert_as, policy_as, 'accel_jsd')
     
     # jerk JSD
-    expert_jerks = np.concatenate([np.concatenate(d['j_all']) for d in expert_metrics])
-    policy_jerks = np.concatenate([np.concatenate(d['j_all']) for d in policy_metrics])
-    jerk_div = divergence(expert_jerks, policy_jerks)
-    print(f'Jerk distribution divergence: {jerk_div}')
-    visualize_distribution(expert_jerks, policy_jerks, 'jerk_jsd.png')
+    expert_jerks = torch.tensor(np.concatenate([np.concatenate(d['j_all']) for d in expert_metrics]))
+    policy_jerks = torch.tensor(np.concatenate([np.concatenate(d['j_all']) for d in policy_metrics]))
+    comparison_metrics['jerk distribution divergence'] = divergence(expert_jerks, policy_jerks)
+    visualize_distribution(expert_jerks, policy_jerks, 'jerk_jsd')
 
-def test_model(
+    for key in comparison_metrics.keys():
+        print(f'{key}: {comparison_metrics[key]}')
+        
+    return comparison_metrics
+
+def eval_main(
     locations: List[Tuple[int,int]]= [(0,0)],
     method: str='expert', 
     policy_file: str='', 
@@ -289,30 +304,27 @@ def test_model(
         env (str): environment class
         method (str): method (expert, bc, gail, rail, hgail, hrail)
     """
+    print(f'Evaluating {method} on {env}')
+
     # set seed
     np.random.seed(seed)
     torch.manual_seed(seed)
 
-    # load policy
-    policy = load_policy(policy_file, method, policy_kwargs, skip_load=(method=='expert'))
-    
     # load expert metrics
     expert_metrics = generate_expert_metrics(locations)
 
-    # if we have a policy
-    if policy:
+    # no comparison for expert
+    if method=='expert':
+        summary_metrics(expert_metrics)
+    
+    # otherwise evaluate policy on roundabouts and generate metrics
+    else:
 
         # evaluate it on the given roundabouts
-        policy_metrics = evaluate_policy(policy, locations, env, env_kwargs)
-
-        # 
+        policy_metrics = evaluate_policy(locations, env, env_kwargs, method, policy_file, policy_kwargs)
         summary_metrics(policy_metrics)
         comparison_metrics(policy_metrics, expert_metrics)
-    
-    else:
-        # if no policy, only generate summary metrics for the expert
-        summary_metrics(expert_metrics)
 
 if __name__=='__main__':
     import fire
-    fire.Fire()
\ No newline at end of file
+    fire.Fire(eval_main)
\ No newline at end of file
diff --git a/src/evaluation/evaluation.py b/src/evaluation/evaluation.py
index e047afe..3374c71 100644
--- a/src/evaluation/evaluation.py
+++ b/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])
diff --git a/src/evaluation/metrics.py b/src/evaluation/metrics.py
index 3b1028d..5ef5d30 100644
--- a/src/evaluation/metrics.py
+++ b/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):
     """