updating evaluation wrapper to only store relevant variables during execution

src/evaluation/evaluation.py

@@ -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:
+        Args:
-    def __init__(self, filestr, eval_env, expert_data, n_eval_episodes=10):
+            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.n_episodes = eval_env.nv
-        self.expert_data = expert_data
+
-        self.compute_expert_features(expert_data)
+        # 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 = []
+        Reset metrics prior to evaluation
-        self._episode_done = True
+        """
-        self._accelerations = []
+        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 compute_expert_features(self, expert_data):
+    def save(self, filestr):
-        # expert velocities
+        """
-        extract_state = lambda info: info['projected_state'][info['agent']]
+        Save metrics to filestr
-        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])
 
-    def evaluate(self, epoch, generator, discriminator):
+        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}')
+
+        # 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(
+        evaluate_policy(
-            generator, 
+            policy, 
             self.env,
             n_eval_episodes=self.n_eval_episodes,
             callback=self.evaluate_policy_callback, 
-            return_episode_rewards=True
+            return_episode_rewards=False
         )
-
+        self.post_proc()
-        collision_rate = self._n_collisions / self.n_eval_episodes
+        self.save(filestr)
-        metrics['collision_rate'] = collision_rate
+        return self._metrics
-
-        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
 
     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)
     
-        # if last episode is done, start new trajectory
+    def post_proc(self):
-        # this is currently not necessary, only if velocity is to be averaged over individual trajectories first
+        """
-        # and then averaging over all trajectories
+        Postprocess and metrics after simulation episodes
-        if self._episode_done:
+        """
-            self._trajectories.append([])
+        # self.metric_keys_all = ['v_all', 'a_all', 'col_all']
-        self._trajectories[-1].append(info['projected_state'][_agent])
+        # self.metric_keys_single = ['j_all', 'v_avg','a_avg', 'col','brake', 't']
-        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