From 5a05e6769f8ec463ac92e75a9eba4c906c2e2cc1 Mon Sep 17 00:00:00 2001 From: Hongrui Zheng Date: Sat, 30 Dec 2023 14:26:56 -0500 Subject: [PATCH] initial impl, separated scan and pose update, used jax (numba doesn't work), slower than just using scan, fix example. --- examples/waypoint_follow.py | 5 +- gym/f110_gym/envs/base_classes.py | 109 +++++++++++++++++++------- gym/f110_gym/envs/collision_models.py | 56 ++++++++++++- gym/f110_gym/envs/f110_env.py | 24 ++++-- 4 files changed, 158 insertions(+), 36 deletions(-) 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/base_classes.py b/gym/f110_gym/envs/base_classes.py index c9754d50..c825dd22 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,6 +357,7 @@ 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 @@ -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,7 +441,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 +450,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 +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): """ diff --git a/gym/f110_gym/envs/collision_models.py b/gym/f110_gym/envs/collision_models.py index 87547ab6..29a4614e 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) @@ -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 """ 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 )