adding pbar to evaluator and making metric save optional, adding typing to baselines

src/baselines/__init__.py

@@ -0,0 +1 @@
+from src.baselines.rule_policies import IDMRulePolicy, PControllerPolicy

src/baselines/rule_policies.py

@@ -1,5 +1,6 @@
 from stable_baselines3.common.base_class import BaseAlgorithm
 from intersim.envs.intersimple import Intersimple
+from typing import Tuple, Optional
 import numpy as np
 
 class PControllerPolicy(BaseAlgorithm):
@@ -14,7 +15,7 @@ class PControllerPolicy(BaseAlgorithm):
         self.target_v = 8.94 # m/s
         self.attn_weight = 20
     
-    def predict(self, observation, *args, **kwargs):
+    def predict(self, observation: np.ndarray, *args, **kwargs):
         """
         Generate action, state from observation
 
@@ -58,15 +59,16 @@ class IDMRulePolicy(BaseAlgorithm):
     
     """
 
-    def __init__(self, env, target_speed: float= 8.94, t_future=0):
+    def __init__(self, env: Intersimple, target_speed: float= 8.94, t_future:float=0):
         """
         Initialize policy with pointer to environment it will run on and target speed
 
         Args:
             env (Intersimple): intersimple environment which IDM runs on
             target_speed (float): target speed in roundabout (default: 8.94=20 mph)
+            t_future (float): future time at which to compare closest
         """
-        assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
+        # assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
         self._env = env
 
         
@@ -82,8 +84,10 @@ class IDMRulePolicy(BaseAlgorithm):
         self.tau = 0.5 # desired time headway
         self.b_pref = 2.5 # preferred deceleration
         self.d_min = 1 #minimum spacing
+        super().__init__()
 
-    def predict(self, observation, *args, **kwargs):
+    def predict(self, observation:np.ndarray, 
+        *args, **kwargs) -> Tuple[np.ndarray,Optional[np.ndarray]]:
         """
         Generate action, state from observation
 
@@ -96,6 +100,8 @@ class IDMRulePolicy(BaseAlgorithm):
             action (np.ndarray): action for controlled agent to take
             state (np.ndarray): the index of the chosen vehicle for IDM
         """
+        import pdb
+        pdb.set_trace()
         agent = self._env._agent
         full_state = self._env._env.projected_state.numpy() #(nv, 5)
         ego_state = full_state[agent] # (5,)
@@ -125,7 +131,8 @@ class IDMRulePolicy(BaseAlgorithm):
         assert(action.shape==(1,))
         return action, i
     
-    def get_ego_dr(self, agent:int, xy: np.ndarray, v: np.ndarray, psi: np.ndarray):
+    def get_ego_dr(self, agent:int, xy: np.ndarray, 
+        v: np.ndarray, psi: np.ndarray) -> Tuple[Optional[np.ndarray], float, float]:
         """
         Return distance and relative speed of closest car within half angle from heading
         

src/evaluation/__init__.py

@@ -0,0 +1 @@
+from src.evaluation.evaluation import IntersimpleEvaluation

src/evaluation/evaluation.py

@@ -2,9 +2,10 @@ 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 typing import Callable, Dict
+from typing import Callable, Dict, Optional
 import os
 import pickle
+from tqdm import tqdm
 
 class IntersimpleEvaluation:
     """
@@ -18,12 +19,13 @@ class IntersimpleEvaluation:
         - whether there was a collision
         - whether there was a hard brake
     """
-    def __init__(self, eval_env):
+    def __init__(self, eval_env, use_pbar:bool=True):
         """
         Initialize evaluation environment with an Intersimple IncrementingAgent environment
 
         Args:
             eval_env (Intersimple.IncrementingAgent): evaluation environment that increments agent upon reset
+            use_pbar (bool): whether to use a progress bar
         """
         # 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!
@@ -34,6 +36,7 @@ class IntersimpleEvaluation:
 
         self.env = eval_env
         self.n_episodes = eval_env.nv
+        self.use_pbar = use_pbar
 
         # metrics present on every step of every episode
         self.metric_keys_all = ['v_all', 'a_all', 'col_all']
@@ -42,7 +45,7 @@ class IntersimpleEvaluation:
         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'
+        self.hard_brake = -3. # acceleration for 'hard brake'
 
         # reset metrics
         self.reset()
@@ -73,17 +76,18 @@ class IntersimpleEvaluation:
         with open(filestr, 'wb') as f:
             pickle.dump(self._metrics, f)     
 
-    def evaluate(self, policy, filestr: str) -> Dict[str, list]:
+    def evaluate(self, policy, filestr: Optional[str] = None) -> 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
+            filestr (str): path-like string to dump metrics to or None
         """
         self.reset()
-        metrics = {}
-        
+        if self.use_pbar:
+            self.pbar = tqdm(total=self.n_episodes)
+
         evaluate_policy(
             policy, 
             self.env,
@@ -91,8 +95,12 @@ class IntersimpleEvaluation:
             callback=self.evaluate_policy_callback, 
             return_episode_rewards=False
         )
+        if self.use_pbar:
+            self.pbar.close()
+            
         self.post_proc()
-        self.save(filestr)
+        if filestr:
+            self.save(filestr)
         return self._metrics
 
     def evaluate_policy_callback(self, local_vars, global_vars):
@@ -114,6 +122,9 @@ class IntersimpleEvaluation:
         if col:
             assert done
         self._metrics['col_all'][_agent].append(col)
+        
+        if done and self.use_pbar:
+            self.pbar.update(1)
     
     def post_proc(self):
         """