Add first metrics

scratch/johannes/gail_options_image.py

@@ -0,0 +1,190 @@
+# %%
+# 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)