From 3991306da0d9b6684f540eda260b54a283241d93 Mon Sep 17 00:00:00 2001 From: Arec Date: Mon, 31 Jan 2022 16:28:13 -0800 Subject: [PATCH] updating evaluation wrapper to only store relevant variables during execution --- src/evaluation/evaluation.py | 182 ++++++++++++++++++++++++----------- 1 file changed, 124 insertions(+), 58 deletions(-) diff --git a/src/evaluation/evaluation.py b/src/evaluation/evaluation.py index 0cefc22..de1306c 100644 --- a/src/evaluation/evaluation.py +++ b/src/evaluation/evaluation.py @@ -1,76 +1,105 @@ -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 typing import Callable, Dict import os +import pickle +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): + """ + 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 (Intersimple.IncrementingAgent): evaluation environment that increments agent upon reset + """ # 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 + + # make sure we have specified an environment which increments the agent number on reset + assert isinstance(eval_env, Intersimple.IncrementingAgent) + 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 + + # 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: [[]]*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: str) -> 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 + """ self.reset() metrics = {} - episode_rewards, episode_lengths = evaluate_policy( - generator, + evaluate_policy( + policy, self.env, n_eval_episodes=self.n_eval_episodes, callback=self.evaluate_policy_callback, - return_episode_rewards=True + return_episode_rewards=False ) - - 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() + 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 +108,52 @@ class Evaluation: assert isinstance(env, Intersimple) # Increase collision counter if episode terminated with a collision - if info['collision']: + 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 col: assert done - self._n_collisions += 1 + self._metrics['col_all'][_agent].append(col) + + 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'] - # 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 + 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.eval_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