diff --git a/evaluate_models.sh b/evaluate_models.sh new file mode 100755 index 0000000..15ff738 --- /dev/null +++ b/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 + + diff --git a/scratch/arec/intersimple/test_model.py b/scratch/arec/intersimple/test_model.py deleted file mode 100644 index 8ae5f2d..0000000 --- a/scratch/arec/intersimple/test_model.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py index aebb694..08e852e 100644 --- a/src/__init__.py +++ b/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 \ No newline at end of file +from src.data.data_utils import InteractionDatasetSingleAgent \ No newline at end of file diff --git a/src/baselines/__init__.py b/src/baselines/__init__.py new file mode 100644 index 0000000..ea73cdf --- /dev/null +++ b/src/baselines/__init__.py @@ -0,0 +1 @@ +from src.baselines.rule_policies import IDMRulePolicy, PControllerPolicy \ No newline at end of file diff --git a/src/baselines/rule_policies.py b/src/baselines/rule_policies.py new file mode 100644 index 0000000..1d5948b --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/eval_main.py b/src/eval_main.py new file mode 100644 index 0000000..edef877 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/src/evaluation/__init__.py b/src/evaluation/__init__.py index e69de29..acfa054 100644 --- a/src/evaluation/__init__.py +++ b/src/evaluation/__init__.py @@ -0,0 +1,2 @@ +from src.evaluation.evaluation import IntersimpleEvaluation +from src.evaluation.metrics import divergence, visualize_distribution \ No newline at end of file diff --git a/src/evaluation/evaluation.py b/src/evaluation/evaluation.py index 0cefc22..3374c71 100644 --- a/src/evaluation/evaluation.py +++ b/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 save(self, filestr): + """ + Save metrics to filestr - 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]) + 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}' - def evaluate(self, epoch, generator, discriminator): + # 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 = {} - - episode_rewards, episode_lengths = evaluate_policy( - generator, + if self.use_pbar: + self.pbar = tqdm(total=self.n_episodes) + + 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 \ No newline at end of file diff --git a/src/evaluation/metrics.py b/src/evaluation/metrics.py index 3b1028d..5ef5d30 100644 --- a/src/evaluation/metrics.py +++ b/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): """