huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

updating evaluation wrapper to only store relevant variables during execution

Browse Source
This commit is contained in:
Arec
2022-01-31 16:28:13 -08:00
parent 3b60c14319
commit 3991306da0
1 changed files with 124 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

182
src/evaluation/evaluation.py
View File

@@ -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