Add exact two-stage collision checking method

scratch/etienne/intersimple/gail_options_image.py

@@ -5,6 +5,7 @@ import stable_baselines3
 import torch.utils.data
 import numpy as np
 from intersim.envs.intersimple import NRasterized
+from intersim.collisions import state_to_polygon
 import itertools
 from torch.distributions import Categorical
 import gym
@@ -179,7 +180,14 @@ def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
     return a*np.ones((t,))
 
 def generate_plan(env, i):
-    """Generate input profile for high-level action `i`."""
+    """Generate input profile for high-level action `i`.
+
+    Args:
+        env (gym.Env): current environment state
+        i (int): high-level action `i`
+    Returns:
+        plan (np.array): length T array of acceleration values
+    """
     assert i < len(ALL_OPTIONS), "Invalid option index {i}"
     target_v, t = ALL_OPTIONS[i]
     current_v = env._env.state[env._agent, 1].item() # extract from env
@@ -187,16 +195,16 @@ def generate_plan(env, i):
     assert len(plan) == t, "incorrect plan length"
     return plan
 
-def check_future_collisions_fast(env, actions):
-    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
-    
-    Vehicles are (over-)approximated by single circles.
+def check_future_collisions_circle(env, actions):
+    """Compute collision information for circular vehicle approximations
 
     Args:
         env (gym.Env): current environment state
         actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
     Returns:
-        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+        states (torch.Tensor): tensor of shape (B, T, nv, 5) of future states based on the action profiles
+        collision_tensor (torch.Tensor): tensor of shape (B, T, nv) of bools indicating which plan collides with which vehicles in which time frame
+            false: colliding, true: not colliding
     """
     B, (T, nv, _) = len(actions), actions[0].shape
 
@@ -213,7 +221,64 @@ def check_future_collisions_fast(env, actions):
     min_distance = min_distance.unsqueeze(0).unsqueeze(0)
     assert min_distance.shape == (1, 1, nv)
 
-    return (distance > min_distance).all(-1).all(-1)
+    collision_tensor = distance > min_distance
+    assert collision_tensor.shape == (B, T, nv)
+    return states, collision_tensor
+
+def check_future_collisions_fast(env, actions):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by single circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    _, collision_tensor = check_future_collisions_circle(env, actions)
+    return collision_tensor.all(-1).all(-1)
+
+def check_future_collisions_exact(env, actions):
+    """
+        Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    # First check with simple circle collision check
+    states, collision_tensor = check_future_coqllisions_circle(env, actions)
+    (B, T, nv, _) = states.shape
+    # For those that have colliding circles, check exactly
+    colliding_mask = ~collision_tensor
+    
+    ego_states = states[:, :, env._agent:env._agent+1, :].expand(states.shape)
+    assert ego_states.shape == states.shape
+
+    # get dimensions
+    lengths = env._env._lengths.expand(states.shape[:3])
+    widths = env._env._widths.expand(states.shape[:3])
+    ego_lengths = lengths[:, :, env._agent:env._agent+1].expand(lengths.shape)
+    ego_widths = widths[:, :, env._agent:env._agent+1].expand(widths.shape)
+    assert lengths.shape == widths.shape == ego_lengths.shape == ego_widths.shape == (B, T, nv)
+    
+    # For every collision instance between ego and other vehicle, check whether rectangles intersect
+    exact_collisions = torch.zeros_like(collision_tensor[colliding_mask])
+    for i, (ego_state, ego_length, ego_width, other_state, other_length, other_width) in enumerate(zip(
+        ego_states[colliding_mask], ego_lengths[colliding_mask], ego_widths[colliding_mask],
+        states[colliding_mask], lengths[colliding_mask], widths[colliding_mask]
+    )):
+        assert ego_state.shape == other_state.shape == (5,)
+        assert ego_length.shape == ego_width.shape == other_length.shape == other_width.shape == ()
+        p_ego = state_to_polygon(ego_state, ego_length, ego_width)
+        p_other = state_to_polygon(other_state, other_length, other_width)
+        exact_collisions[i] = p_ego.intersects(p_other)
+
+    collision_tensor[colliding_mask] = ~exact_collisions
+    return collision_tensor.all(-1).all(-1) 
 
 def feasible(env, plan, ch):
     """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""