Separate scan and pose update, non-scan dependent collision checking

examples/waypoint_follow.py

@@ -298,19 +298,20 @@ def main():
         "vgain": 1,
     }
 
-    num_agents = 3
+    num_agents = 1
     env = gym.make(
         "f110_gym:f110-v0",
         config={
             "map": "Spielberg",
             "num_agents": num_agents,
             "timestep": 0.01,
             "integrator": "rk4",
+            "scan": False,
             "control_input": ["speed", "steering_angle"],
             "model": "st",
             "observation_config": {"type": "kinematic_state"},
             "params": {"mu": 1.0},
-            "reset_config": {"type": "random_static"},
+            # "reset_config": {"type": "random_static"},
         },
         render_mode="human",
     )

gym/f110_gym/envs/action.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from abc import abstractmethod
 from enum import Enum
 from typing import Any, Dict, Tuple

gym/f110_gym/envs/base_classes.py

@@ -30,7 +30,12 @@
 import numpy as np
 from f110_gym.envs.dynamic_models import DynamicModel
 from f110_gym.envs.action import CarAction
-from f110_gym.envs.collision_models import collision_multiple, get_vertices
+from f110_gym.envs.collision_models import (
+    collision_multiple,
+    get_vertices,
+    collision_multiple_map,
+)
+from f110_gym.envs.track import Track
 from f110_gym.envs.integrator import EulerIntegrator, IntegratorType
 from f110_gym.envs.laser_models import ScanSimulator2D, check_ttc_jit, ray_cast
 
@@ -314,13 +319,6 @@ def update_pose(self, raw_steer, vel):
         elif self.state[4] < 0:
             self.state[4] = self.state[4] + 2 * np.pi
 
-        # update scan
-        current_scan = RaceCar.scan_simulator.scan(
-            np.append(self.state[0:2], self.state[4]), self.scan_rng
-        )
-
-        return current_scan
-
     def update_opp_poses(self, opp_poses):
         """
         Updates the vehicle's information on other vehicles
@@ -346,16 +344,24 @@ def update_scan(self, agent_scans, agent_index):
             None
         """
 
-        current_scan = agent_scans[agent_index]
+        # update scan with current pose
+        current_scan = RaceCar.scan_simulator.scan(
+            np.append(self.state[0:2], self.state[4]), self.scan_rng
+        )
+
+        # current_scan = agent_scans[agent_index]
 
         # check ttc
         self.check_ttc(current_scan)
 
         # ray cast other agents to modify scan
         new_scan = self.ray_cast_agents(current_scan)
 
+        # update attribute
         agent_scans[agent_index] = new_scan
 
+        return new_scan
+
 
 class Simulator(object):
     """
@@ -373,6 +379,7 @@ class Simulator(object):
         integrator (Integrator): integrator to use for vehicle dynamics
         model (Model): model to use for vehicle dynamics
         action_type (Action): action type to use for vehicle dynamics
+        track (Track): track of the current simulator
     """
 
     def __init__(
@@ -381,10 +388,12 @@ def __init__(
         num_agents,
         seed,
         action_type: CarAction,
+        track: Track,
         integrator=IntegratorType.RK4,
         model=DynamicModel.ST,
         time_step=0.01,
         ego_idx=0,
+        gen_scan=True,
     ):
         """
         Init function
@@ -412,6 +421,8 @@ def __init__(
         self.collisions = np.zeros((self.num_agents,))
         self.collision_idx = -1 * np.ones((self.num_agents,))
         self.model = model
+        self.track = track
+        self.gen_scan = gen_scan
 
         # initializing agents
         for i in range(self.num_agents):
@@ -432,16 +443,32 @@ def __init__(
 
     def set_map(self, map_name):
         """
-        Sets the map of the environment and sets the map for scan simulator of each agent
+        Sets the map of the environment, initialize map pixel world coordinates,
+        Sets the map for scan simulator of each agent if scan enabled
 
         Args:
             map_name (str): name of the map
 
         Returns:
             None
         """
-        for agent in self.agents:
-            agent.set_map(map_name)
+        # generate map pixels in world frame (pixel center)
+        map_img = self.track.occupancy_map
+        h, w = map_img.shape
+        reso = self.track.spec.resolution
+        ox = self.track.spec.origin[0]
+        oy = self.track.spec.origin[1]
+        x_ind, y_ind = np.meshgrid(np.arange(w), np.arange(h))
+        pcx = (x_ind * reso + ox + reso / 2).flatten()
+        pcy = (y_ind * reso + oy + reso / 2).flatten()
+        self.pixel_centers = np.vstack((pcx, pcy)).T
+        map_mask = (map_img == 0.0).flatten()
+        self.pixel_centers = self.pixel_centers[map_mask, :]
+
+        # set scan simulator map if scan enabled
+        if self.gen_scan:
+            for agent in self.agents:
+                agent.set_map(map_name)
 
     def update_params(self, params, agent_idx=-1):
         """
@@ -486,6 +513,30 @@ def check_collision(self):
             )
         self.collisions, self.collision_idx = collision_multiple(all_vertices)
 
+    def check_map_collision(self):
+        """
+        Checks for collision between agents and map
+
+        Args:
+            None
+
+        Returns:
+            None
+        """
+        # get vertices of all agents
+        all_vertices = np.empty((self.num_agents, 4, 2))
+        for i in range(self.num_agents):
+            all_vertices[i, :, :] = get_vertices(
+                np.append(self.agents[i].state[0:2], self.agents[i].state[4]),
+                self.params["length"],
+                self.params["width"],
+            )
+
+        self.collisions = collision_multiple_map(
+            all_vertices,
+            self.pixel_centers,
+        )
+
     def step(self, control_inputs):
         """
         Steps the simulation environment
@@ -497,11 +548,11 @@ def step(self, control_inputs):
             observations (dict): dictionary for observations: poses of agents, current laser scan of each agent, collision indicators, etc.
         """
 
+        # pose updates
         # looping over agents
         for i, agent in enumerate(self.agents):
             # update each agent's pose
-            current_scan = agent.update_pose(control_inputs[i, 0], control_inputs[i, 1])
-            self.agent_scans[i, :] = current_scan
+            agent.update_pose(control_inputs[i, 0], control_inputs[i, 1])
 
             # update sim's information of agent poses
             self.agent_poses[i, :] = np.append(agent.state[0:2], agent.state[4])
@@ -510,19 +561,25 @@ def step(self, control_inputs):
         # check collisions between all agents
         self.check_collision()
 
-        for i, agent in enumerate(self.agents):
-            # update agent's information on other agents
-            opp_poses = np.concatenate(
-                (self.agent_poses[0:i, :], self.agent_poses[i + 1 :, :]), axis=0
-            )
-            agent.update_opp_poses(opp_poses)
-
-            # update each agent's current scan based on other agents
-            agent.update_scan(self.agent_scans, i)
-
-            # update agent collision with environment
-            if agent.in_collision:
-                self.collisions[i] = 1.0
+        # update scan if requested, with other agent's poses updated
+        if self.gen_scan:
+            for i, agent in enumerate(self.agents):
+                # update agent's information on other agents
+                opp_poses = np.concatenate(
+                    (self.agent_poses[0:i, :], self.agent_poses[i + 1 :, :]), axis=0
+                )
+                agent.update_opp_poses(opp_poses)
+                # generate scan
+                current_scan = agent.update_scan(self.agent_scans, i)
+                self.agent_scans[i, :] = current_scan
+
+                # update agent collision with environment using ittc
+                if agent.in_collision:
+                    self.collisions[i] = 1.0
+
+        # check map collision if not updating scan
+        else:
+            self.check_map_collision()
 
     def reset(self, poses):
         """

gym/f110_gym/envs/collision_models.py

@@ -29,6 +29,10 @@
 
 import numpy as np
 from numba import njit
+from numba.np.extensions import cross2d
+import jax.numpy as jnp
+from jax import jit
+import jax
 
 
 @njit(cache=True)
@@ -212,6 +216,85 @@ def collision_multiple(vertices):
 
     return collisions, collision_idx
 
+@jit
+def collision_multiple_map(vertices, pixel_centers):
+    """
+    Check vertices collision with map occupancy
+    Rasters car polygon to map occupancy
+    vmap across number of cars, and number of occupied pixels
+
+    Args:
+        vertices (np.ndarray (num_bodies, 4, 2)): agent rectangle vertices, ccw winding order
+        pixel_centers (np.ndarray (HxW, 2)): x, y position of pixel centers of map image
+
+    Returns:
+        collisions (np.ndarray (num_bodies, )): whether each body is in collision with map
+    """
+    edges = jnp.roll(vertices, 1, axis=1) - vertices
+    center_p = pixel_centers[:, None, None] - edges
+    cross_prods = jnp.cross(center_p, edges)
+    left_of = jnp.where(cross_prods <= 0, 1.0, 0.0)
+    all_left_of = jnp.sum(left_of, axis=-1)
+    collisions = jnp.where(jnp.sum(jnp.where(all_left_of == 4.0, 1.0, 0.0), axis=0) > 0.0, 1.0, 0.0)
+    return collisions
+
+
+
+@jit
+def collision_multiple_map_jaxloop(vertices, pixel_centers):
+    """
+    Check vertices collision with map occupancy
+    Rasters car polygon to map occupancy
+    JAX impl is about twice faster than the Numba impl
+
+    Args:
+        vertices (np.ndarray (num_bodies, 4, 2)): agent rectangle vertices, ccw winding order
+        pixel_centers (np.ndarray (HxW, 2)): x, y position of pixel centers of map image
+
+    Returns:
+        collisions (np.ndarray (num_bodies, )): whether each body is in collision with map
+    """
+    collisions = jnp.zeros((vertices.shape[0],))
+    # check if center of pixel to the LEFT of all 4 edges
+    # loop because vectorizing is way slower
+    for car_ind in range(vertices.shape[0]):
+        left_of = jnp.empty((4, pixel_centers.shape[0]))
+        for v_ind in range(-1, 3):
+            edge = vertices[car_ind, v_ind + 1] - vertices[car_ind, v_ind]
+            center_p = pixel_centers - vertices[car_ind, v_ind]
+            left_of = left_of.at[v_ind + 1, :].set((jnp.cross(center_p, edge) <= 0))
+        ls = jnp.any((jnp.sum(left_of, axis=0) == 4.0))
+        collisions = collisions.at[car_ind].set(jnp.where(ls, 1.0, 0.0))
+
+    return collisions
+
+
+@njit(cache=True)
+def collision_multiple_map_nb(vertices, pixel_centers):
+    """
+    Check vertices collision with map occupancy
+    Rasters car polygon to map occupancy
+
+    Args:
+        vertices (np.ndarray (num_bodies, 4, 2)): agent rectangle vertices, ccw winding order
+        pixel_centers (np.ndarray (HxW, 2)): x, y position of pixel centers of map image
+
+    Returns:
+        collisions (np.ndarray (num_bodies, )): whether each body is in collision with map
+    """
+    collisions = np.zeros((vertices.shape[0],))
+    # check if center of pixel to the LEFT of all 4 edges
+    # loop because vectorizing is way slower
+    for car_ind in range(vertices.shape[0]):
+        left_of = np.empty((4, pixel_centers.shape[0]))
+        for v_ind in range(-1, 3):
+            edge = vertices[car_ind, v_ind + 1] - vertices[car_ind, v_ind]
+            center_p = pixel_centers - vertices[car_ind, v_ind]
+            left_of[v_ind + 1, :] = cross2d(center_p, edge) <= 0
+        ls = np.any((np.sum(left_of, axis=0) == 4.0))
+        collisions[car_ind] = np.where(ls, 1.0, 0.0)
+    return collisions
+
 
 """
 Utility functions for getting vertices by pose and shape