Merge branch 'test' into main

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

scratch/arec/intersimple/test_model.py

@@ -1,148 +0,0 @@
-from tqdm import tqdm
-from copy import deepcopy
-import stable_baselines3 as sb3
-import intersim
-
-ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
-
-def load_model(model_path:str, method:str):
-    """
-    Load a model given a path and the method
-    
-    Args:
-        model_path (str): the path to the model
-        method (str): the method for the model
-    Returns:
-        model: the action model
-        is_heir (bool): whether the method is heirarchial
-    """
-    model = None
-    is_heir = False
-    if method == 'expert':
-        raise NotImplementedError
-    elif method == 'bc':
-        raise NotImplementedError
-    elif method == 'gail':
-        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
-        raise NotImplementedError
-    else:
-        raise NotImplementedError
-    return model, is_heir
-
-def load_expert_states(roundabout, track):
-    """
-    Load expert states from roundabout/track info
-    Args:
-        roundabout (str): roundabout name
-        track (str): 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
-    import pdb
-    pdb.set_trace()
-    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):
-    """
-    Test a particular model at different locations/tracks
-
-    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
-    """
-
-    # load policy
-    policy, is_heir = load_model(model_name, method)
-
-    # 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)
-        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()
-
-        # 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):
-    """
-    Print and save stats
-    """
-    pass
-
-def load_compare():
-    pass
-
-if __name__=='__main__':
-    import fire
-    fire.Fire()

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/__init__.py

@@ -0,0 +1 @@
+from src.baselines.rule_policies import IDMRulePolicy, PControllerPolicy

src/baselines/rule_policies.py

@@ -0,0 +1,220 @@
+from stable_baselines3.common.base_class import BaseAlgorithm
+from intersim.envs.intersimple import Intersimple, NormalizedActionSpace
+from typing import Tuple, Optional, List
+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
+    
+    # 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
+
+        (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: 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 (List[float]): list of future time at which to compare closest vehicle
+            half_angle (float): half angle to look inside for closest vehicle
+        """
+
+        self._env = env
+        self.t_future = t_future
+        self.half_angle = half_angle
+
+        # 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
+
+        # 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, None]:
+        """
+        Predict 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 (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
+        """
+        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[:,3:4] # (nv, 1)
+
+        d, r, i = self.get_ego_dr(agent, xy, v, psi)
+
+        # propagate environment forward at constant velocity
+        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 = 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)
+        
+        # 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[float, float, Optional[int]]:
+        """
+        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 (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
+
+        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 = int(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

src/eval_main.py

@@ -0,0 +1,330 @@
+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.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
+
+#(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
+    
+    Args:
+        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 method == 'idm':
+        policy = IDMRulePolicy(env, **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 == 'sgail':
+        policy = sb3.PPO.load(policy_file)
+        raise NotImplementedError
+    else:
+        raise NotImplementedError
+    return policy
+
+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}
+    
+    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 (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
+    """
+    rname = intersim.LOCATIONS[roundabout]
+    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(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
+
+    Args:
+        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
+        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 = dict(intersim.envs.intersimple.__dict__)
+    envs_dict.update(dict(options_envs.__dict__))
+    policy_metrics = [None]* len(locations)
+
+    # iterate through vehicles
+    for i, location in tqdm(enumerate(locations)):
+        
+        # add roundabout and track to environent 
+        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)
+
+        # initialize environment
+        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]:
+    """
+    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']
+    summary_metrics = {}
+
+    # average average-velocity
+    all_vavgs =  sum([d['v_avg'] for d in metrics],[]) # aggregate to single list
+    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]
+    summary_metrics['mean positive acceleration'] = np.mean(pos_accels)
+
+    # average deceleration
+    decels = all_aalls[all_aalls<0]
+    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|
+    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
+    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
+    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
+    summary_metrics['mean episode length'] = sum(all_ts)/len(all_ts)
+
+    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
+
+    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
+
+    Returns:   
+        comparison_metrics (Dict[str,float]): dict mapping comparison metric description to value
+    """
+    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)
+    comparison_metrics['mean shortfall velocity'] = np.mean(expert_vavg - policy_vavg)
+    
+    # velocity JSD
+    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 = 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 = 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')
+
+    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='', 
+    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)
+    """
+    print(f'Evaluating {method} on {env}')
+
+    # set seed
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    # load expert metrics
+    expert_metrics = generate_expert_metrics(locations)
+
+    # 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(locations, env, env_kwargs, method, policy_file, policy_kwargs)
+        summary_metrics(policy_metrics)
+        comparison_metrics(policy_metrics, expert_metrics)
+
+if __name__=='__main__':
+    import fire
+    fire.Fire(eval_main)

src/evaluation/__init__.py

@@ -0,0 +1,2 @@
+from src.evaluation.evaluation import IntersimpleEvaluation
+from src.evaluation.metrics import divergence, visualize_distribution

src/evaluation/evaluation.py

@@ -1,76 +1,110 @@
-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 intersim.envs.intersimple import Intersimple, IncrementingAgent
+from typing import Callable, Dict, Optional
 import os
+import pickle
+from tqdm import tqdm
 
+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:IncrementingAgent, use_pbar:bool=True):
+        """
+        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 (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)
-        self.filestr = filestr
+  
         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
+        self.use_pbar = use_pbar
+
+        # 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: [[] 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 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])
+    def save(self, filestr):
+        """
+        Save metrics to filestr
 
-    def evaluate(self, epoch, generator, discriminator):
+        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}'
+
+        # 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: 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 or None
+        """
         self.reset()
-        metrics = {}
+        if self.use_pbar:
+            self.pbar = tqdm(total=self.n_episodes)
 
-        episode_rewards, episode_lengths = evaluate_policy(
-            generator, 
+        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=True
+            return_episode_rewards=False
         )
+        if self.use_pbar:
+            self.pbar.close()
 
-        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()
+        if filestr:
+            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 +113,56 @@ class Evaluation:
         assert isinstance(env, Intersimple)
 
         # Increase collision counter if episode terminated with a collision
-        if info['collision']:
-            assert done
-            self._n_collisions += 1
+        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 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
+        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):
+        """
+        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']
+
+        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.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  

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):
     """