diff --git a/examples/waypoint_follow.py b/examples/waypoint_follow.py index d37add9a..178dbc69 100644 --- a/examples/waypoint_follow.py +++ b/examples/waypoint_follow.py @@ -298,7 +298,7 @@ def main(): "vgain": 1, } - num_agents = 3 + num_agents = 1 env = gym.make( "f110_gym:f110-v0", config={ @@ -306,11 +306,12 @@ def main(): "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", ) diff --git a/gym/f110_gym/envs/action.py b/gym/f110_gym/envs/action.py index 93658462..fcdbc817 100644 --- a/gym/f110_gym/envs/action.py +++ b/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 diff --git a/gym/f110_gym/envs/base_classes.py b/gym/f110_gym/envs/base_classes.py index c9754d50..608d1e68 100644 --- a/gym/f110_gym/envs/base_classes.py +++ b/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,7 +344,12 @@ 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) @@ -354,8 +357,11 @@ def update_scan(self, agent_scans, agent_index): # 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,7 +443,8 @@ 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 @@ -440,8 +452,23 @@ def set_map(self, map_name): 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): """ diff --git a/gym/f110_gym/envs/collision_models.py b/gym/f110_gym/envs/collision_models.py index 87547ab6..ee74213a 100644 --- a/gym/f110_gym/envs/collision_models.py +++ b/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 diff --git a/gym/f110_gym/envs/f110_env.py b/gym/f110_gym/envs/f110_env.py index 5cb1ad41..0852e20b 100644 --- a/gym/f110_gym/envs/f110_env.py +++ b/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,11 +146,12 @@ 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)] @@ -154,6 +159,10 @@ def __init__(self, config: dict = None, render_mode=None, **kwargs): 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 ) diff --git a/setup.py b/setup.py index 8d5d5b50..56b6cd3d 100644 --- a/setup.py +++ b/setup.py @@ -20,6 +20,7 @@ "requests", "shapely", "opencv-python", + "jax[cpu]", ], extras_require={ "dev": [ diff --git a/tests/test_f110_env.py b/tests/test_f110_env.py index 545b9ea3..001f82b8 100644 --- a/tests/test_f110_env.py +++ b/tests/test_f110_env.py @@ -81,7 +81,7 @@ def test_configure_method(self): for k in obs0: self.assertTrue( np.allclose(obs0[k], obs1[k]), - f"Observations {k} should be the same", + f"Observations {k} should be the same, got {obs0[k]} and {obs1[k]}", ) self.assertTrue(done0 == done1, "Done should be the same") t += 1