411 lines
16 KiB
Python
411 lines
16 KiB
Python
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, rwse
|
|
from src.evaluation.utils import save_metrics
|
|
from src.core.policy import SetPolicy, SetDiscretePolicy
|
|
from src.core.reparam_module import ReparamPolicy, ReparamSafePolicy
|
|
from src.options import envs as options_envs2
|
|
from src.safe_options.policy import SetMaskedDiscretePolicy
|
|
from src.safe_options import options as options_envs3
|
|
from src.util.wrappers import IntersimpleTimeLimit
|
|
|
|
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:
|
|
method (str): the method for the model
|
|
policy_file (str): the path to the model
|
|
policy_kwargs (str): the path to the model
|
|
env (Intersimple): Intersimple environment for evaluation (necessary for IDM policy)
|
|
Returns:
|
|
policy (Optional[BaseAlgorithm]): the policy to evaluate
|
|
"""
|
|
ml = torch.device('cpu') if not torch.cuda.is_available() else None
|
|
if method == 'idm':
|
|
policy = IDMRulePolicy(env, **policy_kwargs)
|
|
elif method == 'bc':
|
|
policy = SetPolicy(env.action_space.shape[-1])
|
|
policy.load_state_dict(torch.load(policy_file, map_location=ml))
|
|
policy.eval()
|
|
elif method == 'gail':
|
|
policy = SetPolicy(env.action_space.shape[-1])
|
|
policy(torch.zeros(env.observation_space.shape))
|
|
policy = ReparamPolicy(policy)
|
|
policy.load_state_dict(torch.load(policy_file, map_location=ml))
|
|
policy.eval()
|
|
elif method == 'gail-ppo':
|
|
policy = SetPolicy(env.action_space.shape[-1])
|
|
policy.load_state_dict(torch.load(policy_file, map_location=ml))
|
|
policy.eval()
|
|
elif method == 'rail':
|
|
raise NotImplementedError
|
|
elif method == 'ogail':
|
|
policy = SetDiscretePolicy(env.action_space.n)
|
|
policy(torch.zeros(env.observation_space.shape))
|
|
policy = ReparamPolicy(policy)
|
|
policy.load_state_dict(torch.load(policy_file, map_location=ml))
|
|
policy.eval()
|
|
elif method == 'ogail-ppo':
|
|
policy = SetDiscretePolicy(env.action_space.n)
|
|
policy.load_state_dict(torch.load(policy_file, map_location=ml))
|
|
policy.eval()
|
|
elif method == 'sgail':
|
|
policy = SetMaskedDiscretePolicy(env.action_space.n)
|
|
policy(
|
|
torch.zeros(env.observation_space['observation'].shape),
|
|
torch.zeros(env.observation_space['safe_actions'].shape)
|
|
)
|
|
policy = ReparamSafePolicy(policy)
|
|
policy.load_state_dict(torch.load(policy_file, map_location=ml))
|
|
policy.eval()
|
|
elif method == 'sgail-ppo':
|
|
policy = SetMaskedDiscretePolicy(env.action_space.n)
|
|
policy.load_state_dict(torch.load(policy_file, map_location=ml))
|
|
policy.eval()
|
|
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','x_all','y_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['x_all'][i] = states[nni,i,0].numpy()
|
|
metrics['y_all'][i] = states[nni,i,1].numpy()
|
|
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__))
|
|
envs_dict.update(dict(options_envs2.__dict__))
|
|
envs_dict.update(dict(options_envs3.__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]
|
|
|
|
# wrap in TimeLimit
|
|
if 'max_episode_steps' in it_env_kwargs.keys():
|
|
steps = it_env_kwargs.pop('max_episode_steps')
|
|
eval_env = IntersimpleTimeLimit(Env(**it_env_kwargs),
|
|
max_episode_steps=steps)
|
|
else:
|
|
eval_env = Env(**it_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 = {}
|
|
|
|
# mean travel distance
|
|
dt = 0.1
|
|
travel_ds = []
|
|
for iRound in range(len(metrics)):
|
|
for iTraj in range(len(metrics[iRound]['v_all'])):
|
|
travel_ds.append(dt*sum(metrics[iRound]['v_all'][iTraj]))
|
|
summary_metrics['mean travel distance'] = sum(travel_ds)/len(travel_ds)
|
|
|
|
# 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 acceleration'] = 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)
|
|
summary_metrics['success rate'] = 1 - summary_metrics['collision rate']
|
|
|
|
# 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)
|
|
summary_metrics['mean episode time'] = summary_metrics['mean episode length'] * dt
|
|
|
|
|
|
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]], outbase:str='' ) -> 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
|
|
outbase (str): path to save output figs to
|
|
|
|
Returns:
|
|
comparison_metrics (Dict[str,float]): dict mapping comparison metric description to value
|
|
"""
|
|
comparison_metrics = {}
|
|
|
|
# rwse
|
|
expert_traj, policy_traj = [], []
|
|
for iR in range(len(policy_metrics)):
|
|
for iTraj in range(len(policy_metrics[iR]['x_all'])):
|
|
expert_traj.append(np.vstack((expert_metrics[iR]['x_all'][iTraj], expert_metrics[iR]['y_all'][iTraj])))
|
|
policy_traj.append(np.vstack((policy_metrics[iR]['x_all'][iTraj], policy_metrics[iR]['y_all'][iTraj])))
|
|
assert len(expert_traj)==len(policy_traj)
|
|
comparison_metrics.update(rwse(expert_traj, policy_traj))
|
|
|
|
# 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)
|
|
|
|
# Average |Delta V average|
|
|
comparison_metrics['average absolute average velocity'] = np.mean(np.abs(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, outbase+'_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, outbase+'_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, outbase+'_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'#############################################################################')
|
|
print(f'Evaluating {method} from file {policy_file} on {env} at locations {locations}')
|
|
print(f'#############################################################################')
|
|
|
|
# set seed
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
pfilename = policy_file.split('/')[-1].split('.')[0]
|
|
locstr = 'loc_'+'_'.join([f'r{ro}t{tr}' for (ro,tr) in locations])
|
|
outbase = f'out/{method}/{locstr}/{pfilename}_seed{seed}'
|
|
|
|
# load expert metrics
|
|
expert_metrics = generate_expert_metrics(locations)
|
|
|
|
# no comparison for expert
|
|
if method=='expert':
|
|
smetrics = summary_metrics(expert_metrics)
|
|
save_metrics(smetrics, outbase+'_summary.pkl')
|
|
|
|
# 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)
|
|
smetrics = summary_metrics(policy_metrics)
|
|
save_metrics(smetrics, outbase+'_summary.pkl')
|
|
cmetrics = comparison_metrics(policy_metrics, expert_metrics, outbase=outbase)
|
|
save_metrics(cmetrics, outbase+'_comparison.pkl')
|
|
|
|
if __name__=='__main__':
|
|
import fire
|
|
fire.Fire(eval_main)
|