wrapping all environments in timelimit to stop runs longer than 100s, since some others were erroring

src/eval_main.py

@@ -14,6 +14,7 @@ from src.core.reparam_module import ReparamPolicy, ReparamSafePolicy
 from src.options import envs as options_envs2
 from src.safe_options.policy import SetMaskedDiscretePolicy
 from src.safe_options import options as options_envs3
+from src.util.wrappers import IntersimpleTimeLimit
 
 from typing import Optional, List, Dict, Tuple
 import torch
@@ -201,7 +202,7 @@ def evaluate_policy(locations:List[Tuple[int,int]],
 
     # iterate through vehicles
     for i, location in tqdm(enumerate(locations)):
-        
+
         # add roundabout and track to environent 
         iround, track = location
         rname = intersim.LOCATIONS[iround]
@@ -214,7 +215,14 @@ def evaluate_policy(locations:List[Tuple[int,int]],
 
         # initialize environment
         Env = envs_dict[env_class]
-        eval_env = Env(**it_env_kwargs)
+
+        # wrap in TimeLimit
+        if 'max_episode_steps' in it_env_kwargs.keys():
+            steps = it_env_kwargs.pop('max_episode_steps')
+            eval_env = IntersimpleTimeLimit(Env(**it_env_kwargs), 
+                max_episode_steps=steps)
+        else:
+            eval_env = Env(**it_env_kwargs)
         evaluator = IntersimpleEvaluation(eval_env)
 
         # load policy
@@ -364,7 +372,8 @@ def eval_main(
     np.random.seed(seed)
     torch.manual_seed(seed)
     pfilename = policy_file.split('/')[-1].split('.')[0]
-    outbase = f'out/{method}/{pfilename}_seed{seed}'
+    locstr = 'loc_'+'_'.join([f'r{ro}t{tr}' for (ro,tr) in locations])
+    outbase = f'out/{method}/{locstr}/{pfilename}_seed{seed}'
 
     # load expert metrics
     expert_metrics = generate_expert_metrics(locations)

src/evaluation/evaluation.py

@@ -6,8 +6,9 @@ from typing import Callable, Dict, Optional
 import os
 import pickle
 from tqdm import tqdm
+from src.util.wrappers import IntersimpleTimeLimit
 from src.options.envs import OptionsEnv
-from src.util.wrappers import OptionsTimeLimit
+from src.safe_options.options import SafeOptionsEnv
 
 class IntersimpleEvaluation:
     """
@@ -36,7 +37,10 @@ class IntersimpleEvaluation:
         self.env = eval_env
         self.n_episodes = eval_env.nv
         self.use_pbar = use_pbar
-        self.is_options_env = isinstance(self.env, (OptionsEnv, OptionsTimeLimit))
+        if isinstance(self.env, IntersimpleTimeLimit):
+            self.is_options_env = isinstance(self.env.env, (OptionsEnv, SafeOptionsEnv))
+        else:
+            self.is_options_env = isinstance(self.env, (OptionsEnv, SafeOptionsEnv))
 
         # metrics present on every step of every episode
         self.metric_keys_all = ['x_all', 'y_all', 'v_all', 'a_all', 'col_all']

src/evaluation/metrics.py

@@ -24,7 +24,7 @@ def rwse(expert:List[np.ndarray], policy:List[np.ndarray], dt:float=0.1) -> Dict
     assert len(expert) == len(policy)
 
     # calculate rwse
-    times = [1,2,5,10,15,20]
+    times = [1,2,5,10,15,20,25,30]
     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}

src/evaluation/utils.py

@@ -2,6 +2,7 @@ import pickle
 import os
 import numpy as np
 from typing import List,Dict
+
 def save_metrics(metrics:dict, filestr:str):
     """
     Save metric dict to filestr

src/safe_options/options.py

@@ -14,7 +14,7 @@ from src.options.envs import OptionsEnv
 from src.safe_options.collisions import feasible
 
 from intersim.envs import IntersimpleLidarFlatIncrementingAgent
-from src.util.wrappers import OptionsTimeLimit, Setobs, TransformObservation
+from src.util.wrappers import Setobs, TransformObservation
 
 @dataclass
 class Buffer:
@@ -251,10 +251,10 @@ obs_max = np.array([
     [50, np.pi, 20, 20, np.pi, 1e-1],
 ]).reshape(-1)
 
-def NormalizedSafeOptionsEvalEnv(max_episode_steps=float('inf'), safe_actions_collision_method=None, abort_unsafe_collision_method=None, **kwargs):
-    return OptionsTimeLimit(SafeOptionsEnv(Setobs(
+def NormalizedSafeOptionsEvalEnv(safe_actions_collision_method=None, abort_unsafe_collision_method=None, **kwargs):
+    return SafeOptionsEnv(Setobs(
         TransformObservation(IntersimpleLidarFlatIncrementingAgent(
             n_rays=5,
             **kwargs,
         ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
-    ), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5), (10, 5)], safe_actions_collision_method=safe_actions_collision_method, abort_unsafe_collision_method=abort_unsafe_collision_method), max_episode_steps=max_episode_steps)
+    ), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5), (10, 5)], safe_actions_collision_method=safe_actions_collision_method, abort_unsafe_collision_method=abort_unsafe_collision_method)

src/util/wrappers.py

@@ -9,7 +9,7 @@ class TransformObservation(gym.wrappers.TransformObservation):
     def __getattr__(self, name):
         return getattr(self.env, name)
 
-class OptionsTimeLimit(gym.wrappers.TimeLimit):
+class IntersimpleTimeLimit(gym.wrappers.TimeLimit):
     def __getattr__(self, name):
         return getattr(self.env, name)