adding metric comparisons and updating (note: pre-debug) init
This commit is contained in:
Arec
@@ -1,249 +0,0 @@
|
||||
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
|
||||
|
||||
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]:
|
||||
"""
|
||||
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
|
||||
Returns:
|
||||
policy (Optional[BaseAlgorithm]): the policy to evaluate
|
||||
"""
|
||||
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 == '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}
|
||||
|
||||
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 (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]
|
||||
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]]:
|
||||
"""
|
||||
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
|
||||
|
||||
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__))
|
||||
|
||||
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)
|
||||
policy_metrics[i] = evaluator.evaluate(policy)
|
||||
|
||||
return policy_metrics
|
||||
|
||||
|
||||
|
||||
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
|
||||
"""
|
||||
pass
|
||||
|
||||
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()
|
||||
Reference in New Issue
Block a user