updating rwse to work at different times, updating correct testing environment from roundabout, removing the assertion that a collision implies done in the evaluator, using nanmean and nanstd in averaging

src/eval_main.py

@@ -214,7 +214,7 @@ def evaluate_policy(locations:List[Tuple[int,int]],
 
         # initialize environment
         Env = envs_dict[env_class]
-        eval_env = Env(**env_kwargs)
+        eval_env = Env(**it_env_kwargs)
         evaluator = IntersimpleEvaluation(eval_env)
 
         # load policy
@@ -303,7 +303,7 @@ def comparison_metrics(policy_metrics:List[Dict[str,list]],
             expert_traj.append(np.vstack((expert_metrics[iR]['x_all'][iTraj], expert_metrics[iR]['y_all'][iTraj]))) 
             policy_traj.append(np.vstack((policy_metrics[iR]['x_all'][iTraj], policy_metrics[iR]['y_all'][iTraj]))) 
     assert len(expert_traj)==len(policy_traj)
-    comparison_metrics['rwse'] = rwse(expert_traj, policy_traj)
+    comparison_metrics.update(rwse(expert_traj, policy_traj))
 
     # average velocity shortfall
     expert_vavg = np.array(sum([d['v_avg'] for d in expert_metrics],[]))
@@ -356,7 +356,9 @@ def eval_main(
         env (str): environment class
         method (str): method (expert, bc, gail, rail, hgail, hrail)
     """
-    print(f'Evaluating {method} on {env}')
+    print(f'#############################################################################')
+    print(f'Evaluating {method} from file {policy_file} on {env} at locations {locations}')
+    print(f'#############################################################################')
 
     # set seed
     np.random.seed(seed)

src/evaluation/evaluation.py

@@ -135,8 +135,8 @@ class IntersimpleEvaluation:
         self._metrics['a_all'][_agent].append(info['action_taken'][_agent,0].item())
         col = info['collision']
 
-        if col:
-            assert done
+        # if col: # commenting out if we dont want to end on collision
+        #     assert done
         self._metrics['col_all'][_agent].append(col)
         
         if done and self.use_pbar:

src/evaluation/metrics.py

@@ -4,10 +4,58 @@ import numpy as np
 import matplotlib.pyplot as plt
 from torch.utils.data import DataLoader
 from intersim import collisions
-from typing import List
+from typing import List, Dict
 # import tikzplotlib 
 
-def rwse(expert:List[np.ndarray], policy:List[np.ndarray], dt:float=0.1) -> float:
+def rwse(expert:List[np.ndarray], policy:List[np.ndarray], dt:float=0.1) -> Dict[str,float]:
+    """
+    Calculate average mean squared displacement error
+    
+    Args:
+        expert (List[np.ndarray]): all position trajectories for all expert rollouts
+        policy (List[np.ndarray]): all position trajectories for all policy rollouts
+
+    each trajectory in the list should have shape (2, T). however expert[i] might have a 
+    different T than policy[i]
+    
+    Returns
+        rwse_dict (Dict[str,float]): dict of different RWSEs
+    """
+    assert len(expert) == len(policy)
+
+    # calculate rwse
+    times = [1,2,5,10,15,20]
+    time_indices = [int(t/dt) for t in times]
+    rwse_dict_keys = [f'rwse_{t}s' for t in times]+['rwse_end']
+    se_dict = {key:[] for key in rwse_dict_keys}
+    for expert_trajectory, policy_trajectory in zip(expert, policy):
+        _, T1 = expert_trajectory.shape
+        _, T2 = policy_trajectory.shape
+        minT = min(T1, T2)
+
+        crop_expert_trajectory = expert_trajectory[:, :minT]
+        crop_policy_trajectory = policy_trajectory[:, :minT]
+        
+        # square error along every time
+        se = ((crop_policy_trajectory - crop_expert_trajectory)**2).sum(0)
+        
+        # add to dict with appropriate indexing
+        for time, idx in zip(times, time_indices):
+            if minT >= idx:
+                se_dict[f'rwse_{time}s'].append(se[idx-1])
+        se_dict['rwse_end'].append(se[-1])
+
+    assert len(se_dict['rwse_end']) == len(expert)
+
+    # print how many trajectories of each time:
+    for key in rwse_dict_keys:
+        print('%s has %i elements'%(key, len(se_dict[key])))
+
+    rwse_dict = {key:np.mean(np.array(se_dict[key]))**0.5 for key in rwse_dict_keys}
+
+    return rwse_dict
+
+def rwse_basic(expert:List[np.ndarray], policy:List[np.ndarray], dt:float=0.1) -> float:
     """
     Calculate average mean squared displacement error
     
@@ -42,8 +90,6 @@ def rwse(expert:List[np.ndarray], policy:List[np.ndarray], dt:float=0.1) -> floa
 
     return avg_rwse
 
-
-
 def visualize_distribution(expert, policy, filestr):
     """
     Visualize two distributions

src/evaluation/utils.py

@@ -47,8 +47,8 @@ def average_metrics(metric_list:List[Dict[str,float]]):
         for i in range(N):
             master_dict[key].append(metric_list[i][key])
         master_dict[key] = np.array(master_dict[key])
-        mu = np.mean(master_dict[key])
-        std2 = np.std(master_dict[key])*2
+        mu = np.nanmean(master_dict[key])
+        std2 = np.nanstd(master_dict[key])*2
         print(f'{key}: {mu} \pm {std2}')