Merge branch 'main' into dev

src/__init__.py

@@ -1,3 +1,3 @@
 from src.data.expert_data import generate_expert_data, load_expert_data
 from src.data.data_utils import InteractionDatasetSingleAgent
-from src.metrics import metrics
+from src.evaluation.metrics import metrics

src/evaluation/evaluation.py

@@ -0,0 +1,87 @@
+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
+
+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
+
+        # 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
+        # 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

src/gail/options.py

@@ -40,7 +40,7 @@ class OptionsEnv(gym.Wrapper):
                 # action 0 is considered safe fallback
                 self.m[0] = True
             
-            self.ch, self.value, self.log_prob = generator.policy.predict({
+            self.ch, self.value, self.log_prob = generator.policy.forward({
                 'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
                 'mask': self.m.unsqueeze(0).to(generator.policy.device),
             })

src/gail/train.py

@@ -9,10 +9,11 @@ def flatten_transitions(transitions):
         'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
     }
 
-def train_discriminator(env, generator, discriminator, num_samples):
+def train_discriminator(env, generator, discriminator, num_samples, n_updates=1):
     transitions = list(itertools.islice(env.sample_ll(generator), num_samples))
     generator_samples = flatten_transitions(transitions)
-    discriminator.train_disc(gen_samples=generator_samples)
+    for _ in range(n_updates):
+        discriminator.train_disc(gen_samples=generator_samples)
 
 def train_generator(env, generator, discriminator, num_samples):
     generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1))

src/policies/options.py

@@ -22,7 +22,7 @@ class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
         values = self.value_net(latent_vf)
         return values, distribution.distribution
 
-    def predict(self, obs, eps=1e-6):
+    def forward(self, obs, eps=1e-6):
         """
         Will mask invalid states before making action selections
         Args: