From d1f9e3d7c435e070be192f6225e0b2eb9d3f2b37 Mon Sep 17 00:00:00 2001 From: Arec Date: Wed, 2 Feb 2022 22:29:05 -0800 Subject: [PATCH] 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