initial impl, separated scan and pose update, used jax (numba doesn't…

… work), slower than just using scan, fix example.

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/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,14 +344,20 @@ 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
 
 
@@ -373,6 +377,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 +386,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 +419,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 +441,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 +511,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 +546,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 +559,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)
@@ -182,7 +186,7 @@ def collision(vertices1, vertices2):
     return False
 
 
-@njit(cache=True)
+# @njit(cache=True)
 def collision_multiple(vertices):
     """
     Check pair-wise collisions for all provided vertices
@@ -213,6 +217,56 @@ def collision_multiple(vertices):
     return collisions, collision_idx
 
 
+# @njit(cache=True)
+@jit
+def collision_multiple_map(vertices, pixel_centers):
+    """
+    Check pair-wise collisions for all provided vertices
+
+    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
+        map_occupied (np.ndarray (HxW, )): occupancy indicator of map_image
+
+    Returns:
+        collisions (np.ndarray (num_vertices, )): whether each body is in collision with map
+    """
+    collisions = jnp.zeros((vertices.shape[0],))
+
+    # TODO: reduce check size, only pixel_centers occupied around current poses
+    # NOTE: this doesn't work with jit since pixel_center sizes changes
+    # centers = jnp.mean(vertices, axis=1)
+    # dist = jnp.abs(all_pixel_centers[:, None] - centers)
+    # effective_ind = jnp.where(jnp.logical_and(dist[:, :, 0] < 0.5, dist[:, :, 1] < 0.5))
+    # pixel_centers = all_pixel_centers[:, None][effective_ind]
+
+    # 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))
+
+    # center_vecs = (
+    #     pixel_centers[:, None, None] - vertices
+    # )  # result is (HxW, num_bodies, 4, 2)
+    # cross_prod = np.cross(edges, center_vecs, -1)  # cross_prod is (HxW, num_bodeis, 4)
+    # (pix_ind, body_ind, edge_ind) = np.where((cross_prod < 0))
+
+    # check if enclosed pixels occupied, inflate?
+    # inside_occupied = map_occupied[pix_ind]
+    # which_collision = np.unique(body_ind[np.where(inside_occupied)[0]])
+    # which_collision = np.unique(body_ind)
+
+    # figure out which car collided
+    # collisions[which_collision] = 1
+    return collisions
+
+
 """
 Utility functions for getting vertices by pose and shape
 """

gym/f110_gym/envs/f110_env.py

@@ -27,8 +27,7 @@
 # gym imports
 import gymnasium as gym
 
-from f110_gym.envs.action import (CarAction,
-                                  from_single_to_multi_action_space)
+from f110_gym.envs.action import CarAction, from_single_to_multi_action_space
 from f110_gym.envs.integrator import IntegratorType
 from f110_gym.envs.rendering import make_renderer
 
@@ -104,6 +103,7 @@ def __init__(self, config: dict = None, render_mode=None, **kwargs):
         self.model = DynamicModel.from_string(self.config["model"])
         self.observation_config = self.config["observation_config"]
         self.action_type = CarAction(self.config["control_input"], params=self.params)
+        self.gen_scan = self.config["scan"]
 
         # radius to consider done
         self.start_thresh = 0.5  # 10cm
@@ -133,6 +133,10 @@ def __init__(self, config: dict = None, render_mode=None, **kwargs):
         self.start_thetas = np.zeros((self.num_agents,))
         self.start_rot = np.eye(2)
 
+        self.track = Track.from_track_name(
+            self.map_name
+        )  # load track in gym env for convenience
+
         # initiate stuff
         self.sim = Simulator(
             self.params,
@@ -142,18 +146,23 @@ def __init__(self, config: dict = None, render_mode=None, **kwargs):
             integrator=self.integrator,
             model=self.model,
             action_type=self.action_type,
+            track=self.track,
+            gen_scan=self.gen_scan,
         )
+
+        # TODO: from_track_name is basically called twice here, necessary?
         self.sim.set_map(self.map_name)
-        self.track = Track.from_track_name(
-            self.map_name
-        )  # load track in gym env for convenience
 
         # observations
         self.agent_ids = [f"agent_{i}" for i in range(self.num_agents)]
 
         assert (
             "type" in self.observation_config
         ), "observation_config must contain 'type' key"
+
+        # TODO: how to handle obs when not generating scan?
+        #       right now it's just returning empty array
+
         self.observation_type = observation_factory(env=self, **self.observation_config)
         self.observation_space = self.observation_type.space()
 
@@ -223,6 +232,7 @@ def default_config(cls) -> dict:
             "num_agents": 2,
             "timestep": 0.01,
             "ego_idx": 0,
+            "scan": True,
             "integrator": "rk4",
             "model": "st",
             "control_input": ["speed", "steering_angle"],
@@ -240,7 +250,9 @@ def configure(self, config: dict) -> None:
 
             if hasattr(self, "action_space"):
                 # if some parameters changed, recompute action space
-                self.action_type = CarAction(self.config["control_input"], params=self.params)
+                self.action_type = CarAction(
+                    self.config["control_input"], params=self.params
+                )
                 self.action_space = from_single_to_multi_action_space(
                     self.action_type.space, self.num_agents
                 )