From f9e058a7d956eb9d7e492785d802b52e9b60b63b Mon Sep 17 00:00:00 2001
From: ebuehrle <43623224+ebuehrle@users.noreply.github.com>
Date: Thu, 28 Oct 2021 12:35:32 +0200
Subject: [PATCH] Vectorize available actions computation

---
 src/gail/options.py    |  9 +++++++--
 src/util/collisions.py | 24 +++++++++++++-----------
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/src/gail/options.py b/src/gail/options.py
index 21dd7af..5cd7099 100644
--- a/src/gail/options.py
+++ b/src/gail/options.py
@@ -48,7 +48,7 @@ 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):
+            while not self.done and self.plan and feasible(self.env, self.plan, self.ch.to('cpu')):
                 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)
@@ -129,7 +129,12 @@ class RenderOptions(LLOptions):
 
 def available_actions(env, options):
     """Return mask of available actions given current `env` state."""
-    valid = np.array([feasible(env, generate_plan(env, i, options), i) for i in range(len(options))])
+    plan_indices = list(range(len(options)))
+    plans = [generate_plan(env, i, options) for i in plan_indices]
+    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)
     return valid
 
 def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
diff --git a/src/util/collisions.py b/src/util/collisions.py
index f73bf5d..33668a4 100644
--- a/src/util/collisions.py
+++ b/src/util/collisions.py
@@ -4,23 +4,25 @@ from intersim.collisions import state_to_polygon
 
 def feasible(env, plan, ch, method='exact'):
     """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
-    if ch == 0:
-        return True
+    # 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
 
-    # zero pad plan - Take (T,) np plan and convert it to (T, nv, 1) torch.Tensor
-    full_plan = torch.zeros(len(plan), env._env._nv, 1)
-    full_plan[:, env._agent, 0] = torch.tensor(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]) 
+        valid = check_future_collisions_fast(env, full_plan) 
     elif method=='ncircles':
-        valid = check_future_collisions_ncircles(env, [full_plan])
+        valid = check_future_collisions_ncircles(env, full_plan)
     elif method=='exact':
-        valid = check_future_collisions_exact(env, [full_plan])
+        valid = check_future_collisions_exact(env, full_plan)
     else:
         raise NotImplementedError('Invalid collision-checking method')
-    return valid.item()
+
+    return valid | (ch == 0)
 
 def check_future_collisions_ncircles(env, actions, n_circles:int=2):
     """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
@@ -36,7 +38,7 @@ def check_future_collisions_ncircles(env, actions, n_circles:int=2):
     assert n_circles >= 2
     B, (T, nv, _) = len(actions), actions[0].shape
 
-    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
+    states = env._env.propagate_action_profile_vectorized(actions)
     assert states.shape == (B, T, nv, 5)
     centers = states[:, :, :, :2]
     psi = states[:, :, :, 3]
@@ -81,7 +83,7 @@ def check_future_collisions_circle(env, actions):
     """
     B, (T, nv, _) = len(actions), actions[0].shape
 
-    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
+    states = env._env.propagate_action_profile_vectorized(actions)
     assert states.shape == (B, T, nv, 5)
 
     distance = ((states[:, :, :, :2] - states[:, :, env._agent:env._agent+1, :2])**2).sum(-1).sqrt()