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

@@ -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
+
+

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

src/baselines/rule_policies.py

@@ -1,6 +1,6 @@
 from stable_baselines3.common.base_class import BaseAlgorithm
-from intersim.envs.intersimple import Intersimple
+from intersim.envs.intersimple import Intersimple, NormalizedActionSpace
-from typing import Tuple, Optional
+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 = half_angle
-        self.half_angle = 45 # degrees for finding car to follow
 
         # 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:
+        for t in self.t_future:
-            xy2 = xy + self.t_future * v * np.vstack((np.cos(psi[:,0]), np.sin(psi[:,0])))
+            if t > 0:
-            d2, r2, i2 = self.get_ego_dr(agent, xy2, v, psi)
+                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)
+                # choose closer vehicle (now vs imagined)
-            if d2 < d:
+                if d2 < d:
-                d, r, i = d2, r2, i2   
+                    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 )
+            d_des = max(d_des, self.d_min)
+
+        assert (d_des>= self.d_min)
         action = self.a_max*(1 - (s/self.s_max)**4 - (d_des/d)**2)
         
-        assert(action.shape==(1,))
+        # normalize action to range if env is a NormalizedActionSpace
-        return action, i
+        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 nv==nv2==nv3
-        assert(nxy==2)
+        assert nxy==2
-        assert(nvel==npsi==1)
+        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] 
             

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, 
+#(method, policy_file, policy_kwargs, eval_env)
-    policy_kwargs:dict, skip_load:bool=False) -> Optional[BaseAlgorithm]:
+
+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:
+    if method == 'idm':
-        return None
+        policy = IDMRulePolicy(env, **policy_kwargs)
-    elif method == 'idm':
-        policy = IDMRulePolicy(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]], 
+def evaluate_policy(locations:List[Tuple[int,int]], 
-    env_class:str, env_kwargs:dict) -> List[Dict[str,list]]:
+    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 = dict(intersim.envs.intersimple.__dict__)
-    envs_dict.update(deepcopy(options_envs.__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]]) -> Dict[str,float]:
-
-def summary_metrics(metrics:List[Dict[str,list]]):
     """
     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
+    summary_metrics = {}
-    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)
+    summary_metrics['mean average velocity'] = sum(all_vavgs)/len(all_vavgs)
-    print(f'Mean Average Velocity: {mean_vavg}')
-
     
+    # 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)
+    summary_metrics['mean positive acceleration'] = 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)
+    summary_metrics['mean deceleration'] = 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])
+    summary_metrics['mean jerk'] = np.mean(all_jerks)
 
     # average |jerk|
-    all_jerks = np.concatenate([np.concatenate(d['j_all']) for d in metrics])
+    summary_metrics['mean |jerk|'] = np.mean(np.abs(all_jerks))
-    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)
+    summary_metrics['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)
+    summary_metrics['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)
+    summary_metrics['mean episode length'] = 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]]):
+    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
 
@@ -238,39 +251,41 @@ def comparison_metrics(policy_metrics:List[Dict[str,list]], expert_metrics:List[
         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
+    comparison_metrics = {}
-    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))
+    assert len(expert_vavg)==len(policy_vavg)
-    mean_shortfall = np.mean(expert_vavg - policy_vavg)
+    comparison_metrics['mean shortfall velocity'] = 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])
+    expert_vs = torch.tensor(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])
+    policy_vs = torch.tensor(np.concatenate([np.concatenate(d['v_all']) for d in policy_metrics]))
-    vel_div = divergence(expert_vs, policy_vs)
+    comparison_metrics['velocity distribution divergence'] = divergence(expert_vs, policy_vs)
-    print(f'Velocity distribution divergence: {vel_div}')
+    visualize_distribution(expert_vs, policy_vs, 'velocity_jsd')
-    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])
+    expert_as = torch.tensor(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])
+    policy_as = torch.tensor(np.concatenate([np.concatenate(d['a_all']) for d in policy_metrics]))
-    accel_div = divergence(expert_as, policy_as)
+    comparison_metrics['acceleration distribution divergence'] = divergence(expert_as, policy_as)
-    print(f'Velocity distribution divergence: {accel_div}')
+    visualize_distribution(expert_as, policy_as, 'accel_jsd')
-    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])
+    expert_jerks = torch.tensor(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])
+    policy_jerks = torch.tensor(np.concatenate([np.concatenate(d['j_all']) for d in policy_metrics]))
-    jerk_div = divergence(expert_jerks, policy_jerks)
+    comparison_metrics['jerk distribution divergence'] = divergence(expert_jerks, policy_jerks)
-    print(f'Jerk distribution divergence: {jerk_div}')
+    visualize_distribution(expert_jerks, policy_jerks, 'jerk_jsd')
-    visualize_distribution(expert_jerks, policy_jerks, 'jerk_jsd.png')
 
-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
+    # no comparison for expert
-    if policy:
+    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()
+    fire.Fire(eval_main)

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, 
+            assert len(self._metrics[key])==self.n_episodes, \
-                f'_metrics[{key}] does not have length {self.n_episodes}')
+                f'_metrics[{key}] does not have length {self.n_episodes}'
 
         # make filepath
         os.makedirs(os.path.dirname(filestr))
@@ -91,7 +88,7 @@ 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
         )
@@ -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):
+def visualize_distribution(expert, policy, filestr):
-    """
-    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):
     """
     Visualize two distributions
     Args:
-        true (torch.tensor): (n,)-sized true distribution
+        expert (torch.tensor): (n,)-sized true distribution
-        pred (torch.tensor): (m,)-sized pred distribution
+        generated (torch.tensor): (m,)-sized pred distribution
         filestr (str): string to save figure to
     """
-    nni1 = ~torch.isnan(true)
+    nni1 = ~torch.isnan(expert)
-    nni2 = ~torch.isnan(pred)
+    nni2 = ~torch.isnan(policy)
     plt.figure()
-    plt.hist(true[nni1].numpy(), density=True, bins=20)
+    plt.hist(expert[nni1].numpy(), density=True, bins=20)
-    plt.hist(pred[nni2].numpy(), density=True, bins=20)
+    plt.hist(policy[nni2].numpy(), density=True, bins=20)
-    plt.legend(['True', 'Predicted'])
+    plt.legend(['Expert', 'Predicted'])
     plt.savefig(filestr+'.png')
+    # tikzplotlib.save(filestr+'.tex')
     
 def average_velocity(states):
     """