From 7ffcc0b4b8bb8601502edab3617e79ef56bab259 Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Thu, 28 Oct 2021 13:29:00 +0200 Subject: [PATCH] Extract evaluation code to separate file --- scratch/johannes/gail_options_image.py | 95 +------------------------- src/evaluation/evaluation.py | 92 +++++++++++++++++++++++++ 2 files changed, 95 insertions(+), 92 deletions(-) create mode 100644 src/evaluation/evaluation.py diff --git a/scratch/johannes/gail_options_image.py b/scratch/johannes/gail_options_image.py index f405db1..ac1676a 100644 --- a/scratch/johannes/gail_options_image.py +++ b/scratch/johannes/gail_options_image.py @@ -8,7 +8,7 @@ import stable_baselines3 from stable_baselines3.common.evaluation import evaluate_policy import torch.utils.data import numpy as np -from intersim.envs.intersimple import Intersimple, NormalizedActionSpace, NRasterized, NRasterizedInfo, NRasterizedIncrementingAgent, NRasterizedRandomAgent, speed_reward +from intersim.envs.intersimple import Intersimple, NRasterized, NRasterizedInfo, NRasterizedIncrementingAgent, NRasterizedRandomAgent, speed_reward import itertools import functools from torch.distributions import Categorical @@ -24,10 +24,9 @@ from tqdm import tqdm from src.policies.options import OptionsCnnPolicy from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions from src.gail.train import train_discriminator, train_generator -from src.metrics import nanmean, divergence, visualize_distribution model_name = 'gail_options_image' -Env = NRasterized +Env = NRasterizedRandomAgent env_settings = {'width': 36, 'height': 36, 'm_per_px': 2} ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5]] # option 0 is safe fallback @@ -41,7 +40,7 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=32, disc logger.configure(tempdir_path / "GAIL/") print(f"All Tensorboards and logging are being written inside {tempdir_path}/.") - venv = make_vec_env(NRasterized, n_envs=1, env_kwargs=env_settings) + venv = make_vec_env(Env, n_envs=1, env_kwargs=env_settings) discriminator = adversarial.GAIL( expert_data=expert_data, expert_batch_size=expert_batch_size, @@ -75,94 +74,6 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=32, disc return generator -from stable_baselines3.common.vec_env import VecEnv -class Evaluation: - def __init__(self, eval_env, n_eval_episodes=10): - # 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.env = eval_env - self.n_eval_episodes = n_eval_episodes - self.reset() - - def reset(self): - self._n_collisions = 0 - self._trajectories = [] - self._episode_done = True - self._accelerations = [] - - def evaluate(self, epoch, generator, discriminator, expert_data): - self.reset() - metrics = {} - - episode_rewards, episode_lengths = evaluate_policy( - generator, - self.env, - n_eval_episodes=self.n_eval_episodes, - callback=self.evaluate_policy_callback, - return_episode_rewards=True - ) - - collision_rate = self._n_collisions / self.n_eval_episodes - metrics['collision_rate'] = collision_rate - - assert len(self._trajectories) >= self.n_eval_episodes - - # average velocity of each episode - # this first averages velocity over single trajectories and then averages over trajectories - # avg_velocities = [nanmean(torch.stack(t)[:,2]) for t in self._trajectories] - # avg_velocity = np.mean(avg_velocities) - - # 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)] - - # 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] - expert_velocities = expert_velocities[~torch.isnan(expert_velocities)] - - metrics['avg_velocity_loss'] = (expert_velocities.mean() - policy_velocities.mean()).item() - metrics['velocity_divergence'] = divergence(policy_velocities, expert_velocities, type='js') - - - # accelerations produced by generator - policy_accelerations = torch.tensor(self._accelerations) - # expert accelerations - extract_accel = lambda info: info['action_taken'][info['agent']] - expert_accelerations = torch.cat([extract_accel(info) for info in expert_data.infos]) - - metrics['acceleration_divergence'] = divergence(policy_accelerations, expert_accelerations, type='js') - visualize_distribution(expert_accelerations, policy_accelerations, 'output/_action_viz{:02}'.format(epoch)) - - print(metrics) - return metrics - - def evaluate_policy_callback(self, local_vars, global_vars): - venv_i = local_vars['i'] - info = local_vars['info'] - done = local_vars['done'] - _agent = info['agent'] - env = local_vars['env'].envs[venv_i] - assert isinstance(env, Intersimple) - - # Increase collision counter if episode terminated with a collision - if info['collision']: - assert done - self._n_collisions += 1 - - # if last episode is done, start new trajectory - if self._episode_done: - self._trajectories.append([]) - agent_state = info['projected_state'][_agent] - self._trajectories[-1].append(agent_state) - # this does not work since agent_action are high level options - # agent_action = local_vars['actions'][venv_i] # normalized intersimple action - # acceleration = env._unnormalize(agent_action) if isinstance(env, NormalizedActionSpace) else agent_action - self._accelerations.append(info['action_taken'][_agent]) - self._episode_done = done - diff --git a/src/evaluation/evaluation.py b/src/evaluation/evaluation.py new file mode 100644 index 0000000..882069a --- /dev/null +++ b/src/evaluation/evaluation.py @@ -0,0 +1,92 @@ +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 + +class Evaluation: + def __init__(self, eval_env, n_eval_episodes=10): + # 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.env = eval_env + self.n_eval_episodes = n_eval_episodes + self.reset() + + def reset(self): + self._n_collisions = 0 + self._trajectories = [] + self._episode_done = True + self._accelerations = [] + + def evaluate(self, epoch, generator, discriminator, expert_data): + self.reset() + metrics = {} + + episode_rewards, episode_lengths = evaluate_policy( + generator, + self.env, + n_eval_episodes=self.n_eval_episodes, + callback=self.evaluate_policy_callback, + return_episode_rewards=True + ) + + collision_rate = self._n_collisions / self.n_eval_episodes + metrics['collision_rate'] = collision_rate + + assert len(self._trajectories) >= self.n_eval_episodes + + # average velocity of each episode + # this first averages velocity over single trajectories and then averages over trajectories + # avg_velocities = [nanmean(torch.stack(t)[:,2]) for t in self._trajectories] + # avg_velocity = np.mean(avg_velocities) + + # 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)] + + # 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] + expert_velocities = expert_velocities[~torch.isnan(expert_velocities)] + + metrics['avg_velocity_loss'] = (expert_velocities.mean() - policy_velocities.mean()).item() + metrics['velocity_divergence'] = divergence(policy_velocities, expert_velocities, type='js') + + + # accelerations produced by generator + policy_accelerations = torch.tensor(self._accelerations) + # expert accelerations + extract_accel = lambda info: info['action_taken'][info['agent']] + expert_accelerations = torch.cat([extract_accel(info) for info in expert_data.infos]) + + metrics['acceleration_divergence'] = divergence(policy_accelerations, expert_accelerations, type='js') + visualize_distribution(expert_accelerations, policy_accelerations, 'output/_action_viz{:02}'.format(epoch)) + + print(metrics) + return metrics + + def evaluate_policy_callback(self, local_vars, global_vars): + venv_i = local_vars['i'] + info = local_vars['info'] + done = local_vars['done'] + _agent = info['agent'] + env = local_vars['env'].envs[venv_i] + assert isinstance(env, Intersimple) + + # Increase collision counter if episode terminated with a collision + if info['collision']: + assert done + self._n_collisions += 1 + + # if last episode is done, start new trajectory + if self._episode_done: + self._trajectories.append([]) + agent_state = info['projected_state'][_agent] + self._trajectories[-1].append(agent_state) + # this does not work since agent_action are high level options + # agent_action = local_vars['actions'][venv_i] # normalized intersimple action + # acceleration = env._unnormalize(agent_action) if isinstance(env, NormalizedActionSpace) else agent_action + self._accelerations.append(info['action_taken'][_agent]) + self._episode_done = done