# %% # import sys # sys.path.append('../../../') from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction from imitation.algorithms import adversarial 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, NRasterized, NRasterizedIncrementingAgent, NRasterizedRandomAgent, speed_reward import itertools import functools from torch.distributions import Categorical import gym import torch import pickle import imitation.data.rollout as rollout import tempfile import pathlib from imitation.util import logger from stable_baselines3.common.env_util import make_vec_env 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_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 def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=32, discount=0.99): env = NRasterizedRandomAgent(**env_settings) env.discount = discount tempdir = tempfile.TemporaryDirectory(prefix="quickstart") tempdir_path = pathlib.Path(tempdir.name) 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) discriminator = adversarial.GAIL( expert_data=expert_data, expert_batch_size=expert_batch_size, discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)}, #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)}, venv=venv, # unused gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused ) generator = stable_baselines3.PPO( OptionsCnnPolicy, OptionsEnv(env, options=ALL_OPTIONS), verbose=1, n_steps=generator_steps, ) # PPO.train requires logger as set up in # PPO._setup_learn (called by PPO.learn) generator._logger = stable_baselines3.common.utils.configure_logger( generator.verbose, generator.tensorboard_log, ) for epoch in tqdm(range(epochs)): train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size) train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps) eval_env = NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_settings) ev = Evaluation(eval_env, n_eval_episodes=10) ev.evaluate(epoch, generator, discriminator, expert_data) 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 def evaluate(self, epoch, generator, discriminator, expert_data): self.reset() info = {} 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 info['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]) 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)] info['avg_velocity_loss'] = expert_velocities.mean() - policy_velocities.mean() # divergence (true, policy) # same for acceleration print(info) return info def evaluate_policy_callback(self, local_vars, global_vars): info = local_vars['info'] done = local_vars['done'] env = local_vars['env'].envs[local_vars['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([]) self._trajectories[-1].append(info['projected_state'][env._agent]) self._episode_done = done # class CAPolicy: # def __init__(self, a): # self.a = torch.tensor([a]) # def predict(self, obs, state=None, deterministic=False): # return self.a, state # %% if __name__ == '__main__': # %% # env = NRasterizedIncrementingAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_settings) # generator = CAPolicy(0.0) # ev = Evaluation(env, 10) # ev.evaluate(1, generator, None, None) # exit() # %% with open("scratch/etienne/intersimple/data/NormalizedIntersimpleExpertMu.001N200_NRasterizedInfoAgent51w36h36mppx2.pkl", "rb") as f: trajectories = pickle.load(f) transitions = rollout.flatten_trajectories(trajectories) generator = train(transitions) generator.save(model_name) # %% model = stable_baselines3.PPO.load(model_name) env = RenderOptions(NRasterizedRandomAgent(**env_settings), options=ALL_OPTIONS) for s in env.sample_ll(model): if s['dones']: break env.close(filestr='render/'+model_name)