adding main test sequence. must debug and add summary and comparison metric generators tomorrow

scratch/arec/intersimple/test_model.py

@@ -2,87 +2,165 @@ 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
+        policy (Optional[BaseAlgorithm]): the policy to evaluate
-        is_heir (bool): whether the method is heirarchial
     """
-    model = None
+    if skip_load:
-    is_heir = False
+        return None
-    if method == 'expert':
+    elif method == 'idm':
-        raise NotImplementedError
+        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':
+    elif method == 'sgail':
-        is_heir = True
+        policy = sb3.PPO.load(policy_file)
-        model = sb3.PPO.load(model_path)
-    elif method == 'hrail':
-        is_heir = True
         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
+        roundabout (int): roundabout index
-        track (str): track id
+        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
+    rname = intersim.LOCATIONS[roundabout]
-    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()
+    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(
+def evaluate_policy(policy:BaseAlgorithm, locations:List[Tuple[int,int]], 
-    locations=[(0,0)], 
+    env_class:str, env_kwargs:dict) -> List[Dict[str,list]]:
-    model_name='gail_image_multiagent_nocollision', 
-    env='NRasterizedRouteIncrementingAgent', 
-    method='expert',
-    options_list=ALL_OPTIONS, 
-    **env_kwargs):
     """
-    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
+        policy (BaseAlgorithm): policy to evaluate
-        model_name (str): name of model to test
+        locations (list of tuples): list of locations to evaluate policy
-        env (str): environment class
+        env_class (str): name of environment to evaluate policy with
-        method (str): method (expert, bc, gail, rail, hgail, hrail)
+        env_kwargs (dict): key word arguments to initialize environment with
-        options_list (list): list of options
-    """
     
-    # load policy
+    Returns:
-    policy, is_heir = load_model(model_name, method)
+        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__))
+    
+    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, track = location
-        iround = intersim.LOCATIONS.index(roundabout)
+        rname = intersim.LOCATIONS[iround]
         it_env_kwargs = deepcopy(env_kwargs)
         loc_kwargs = {
             'loc':iround,
@@ -90,59 +168,82 @@ def test_model(
         }
         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 = envs_dict[env_class]
-            Env = src.options.envs.__dict__[env]
+        eval_env = Env(**env_kwargs)
-        else: 
+        evaluator = IntersimpleEvaluation(eval_env)
-            Env = intersim.envs.intersimple.__dict__[env]
+        policy_metrics[i] = evaluator.evaluate(policy)
-        env = Env(**env_kwargs)
-        s = env.reset()
     
-        # Iterate through every vehicle and time
+    return policy_metrics
-        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()