Check emergency braking in available actions

src/gail/options.py

@@ -5,7 +5,7 @@ import numpy as np
 
 class OptionsEnv(gym.Wrapper):
     
-    def __init__(self, env, options=[(v,t) for v in [0,2,4,6,8] for t in [5]], *args, **kwargs):
+    def __init__(self, env, options=[(0, 5), (5, 5), (10, 5)], *args, **kwargs):
         """option 0 is treated as safe fallback"""
 
         super().__init__(env, *args, **kwargs)
@@ -36,9 +36,13 @@ class OptionsEnv(gym.Wrapper):
                 self.episode_start = True
 
             self.m = available_actions(self.env, self.options)
+            if not self.m.any():
+                # action 0 is considered safe fallback
+                self.m[0] = True
+            
             self.ch, self.value, self.log_prob = generator.policy.predict({
                 'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
-                'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
+                'mask': self.m.unsqueeze(0).to(generator.policy.device),
             })
             self.plan = list(map(float, generate_plan(self.env, self.ch, self.options)))
 
@@ -48,7 +52,9 @@ class OptionsEnv(gym.Wrapper):
             assert self.plan
             #assert feasible(self.env, self.plan, self.ch)
 
-            while not self.done and self.plan and feasible(self.env, self.plan, self.ch.to('cpu')):
+            while not self.done and self.plan and \
+                (feasible(self.env, safety_plan(self.env, self.plan)) or self.m.sum() == 1):
+                
                 self.a, self.plan = self.plan[0], self.plan[1:]
                 self.a = self.env._normalize(self.a)
                 self.nexts, _, self.done, _ = self.env.step(self.a)
@@ -127,14 +133,20 @@ class RenderOptions(LLOptions):
     def close(self, *args, **kwargs):
         self.env.close(*args, **kwargs)
 
+def safety_plan(env, plan):
+    return np.concatenate((plan, np.array(5 * [env._env._min_acc])), axis=0)
+
 def available_actions(env, options):
     """Return mask of available actions given current `env` state."""
-    plan_indices = list(range(len(options)))
+    plans = [generate_plan(env, i, options) for i, _ in enumerate(options)]
-    plans = [generate_plan(env, i, options) for i in plan_indices]
+    # is emergency braking still possible?
+    plans = list(map(lambda p: safety_plan(env, p), plans))
+
     T = max(len(p) for p in plans)
     plans = [np.pad(p, ((0, T-len(p)),), constant_values=np.nan) for p in plans]
     plans = np.stack(plans, axis=0)
-    valid = feasible(env, plans, plan_indices)
+    
+    valid = feasible(env, plans)
     return valid
 
 def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):

src/util/collisions.py

@@ -2,16 +2,14 @@ import torch
 import numpy as np
 from intersim.collisions import state_to_polygon
 
-def feasible(env, plan, ch, method='exact'):
+def feasible(env, plan, method='exact'):
-    """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
+    """Check if input profile is feasible given current `env` state."""
     # zero pad plan - Take (B, T) or (T,) np plan and convert it to (B, T, nv, 1) torch.Tensor
     plan = torch.tensor(plan)
     plan = plan.reshape(-1, plan.shape[-1])
     full_plan = torch.zeros(*plan.shape, env._env._nv, 1)
     full_plan[:, :, env._agent, 0] = plan
 
-    ch = torch.tensor(ch)
-
     # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
     if method=='circle':
         valid = check_future_collisions_fast(env, full_plan) 
@@ -22,7 +20,7 @@ def feasible(env, plan, ch, method='exact'):
     else:
         raise NotImplementedError('Invalid collision-checking method')
 
-    return valid | (ch == 0)
+    return valid
 
 def check_future_collisions_ncircles(env, actions, n_circles:int=2):
     """Checks whether `env._agent` would collide with other agents assuming `actions` as input.