From 9c01c96691785cc9b0c16ae94d781ca723413954 Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Tue, 30 Jun 2026 18:23:25 +0200 Subject: [PATCH 1/2] Fix random track generation --- lsy_drone_racing/envs/randomize.py | 200 +++++++++++++---------------- tests/unit/envs/test_randomize.py | 70 ++++++++++ 2 files changed, 156 insertions(+), 114 deletions(-) create mode 100644 tests/unit/envs/test_randomize.py diff --git a/lsy_drone_racing/envs/randomize.py b/lsy_drone_racing/envs/randomize.py index 7dcceb4e7..c20abec6d 100644 --- a/lsy_drone_racing/envs/randomize.py +++ b/lsy_drone_racing/envs/randomize.py @@ -125,39 +125,42 @@ def build_random_track_fn( pos_limit_high: Array, *, border_margin: float = 0.5, - start_excl_r: float = 1.0, + drone_excl_r: float = 1.0, gate_excl_r: float = 1.0, - obstacle_excl_r: float = 1.0, - gate_corridor_width: float = 0.4, - obstacle_corridor_width: float = 0.4, + obstacle_excl_r: float = 0.8, + obstacle_obstacle_excl_r: float = 0.3, + gate_corridor_width: float = 0.5, + obstacle_corridor_width: float = 0.5, yaw_range: float = 0.75, - grid_h: int = 40, - grid_w: int = 40, -) -> Callable[[Array], tuple[Array, Array, Array]]: + grid_res: int = 40, +) -> Callable[[Array, Array], tuple[Array, Array, Array]]: """Build a JIT- and vmap-compatible function that generates a complete random track layout. - Gates and obstacles are placed on a 2-D grid using iterative exclusion zones and corridor - masks. Gate z-heights and obstacle z-heights are fixed at the values provided. + The track is built around a given drone start position. Gates and obstacles are placed one after + another on a static 2-D grid. Each object type keeps a running per-cell clearance field (signed + distance to the nearest placed object minus its radius) and samples a cell that still keeps + every exclusion radius. If the grid saturates, the least-violating cell is used. Gate and + obstacle heights are fixed at the values provided. Args: - n_objects: Number of gates (= number of obstacles). gates_z: Z-height for each gate, shape ``(n_objects,)``. obstacles_z: Z-height for each obstacle, shape ``(n_objects,)``. pos_limit_low: XY lower bounds of the arena ``[xmin, ymin]``. pos_limit_high: XY upper bounds of the arena ``[xmax, ymax]``. border_margin: Min distance [m] of all objects from the arena boundary. - start_excl_r: Exclusion radius [m] around the drone start position. + drone_excl_r: Min distance [m] of gates and obstacles from the drone start position. gate_excl_r: Min distance [m] between consecutive gates. - obstacle_excl_r: Min distance [m] from gates to obstacles and between obstacles. + obstacle_excl_r: Min distance [m] from gates to obstacles. + obstacle_obstacle_excl_r: Min distance [m] between obstacles. gate_corridor_width: Half-width [m] of the flight corridor masked out for gate placement. obstacle_corridor_width: Half-width [m] of the corridor used for obstacle placement. yaw_range: Maximum yaw offset [rad] from the travel direction for gate orientation. - grid_h: Grid height (number of rows). - grid_w: Grid width (number of columns). + grid_res: Number of grid nodes along the x-axis. The y-axis count is derived from the arena + aspect ratio so the grid spacing is equal in both axes (square cells). Returns: - ``generate(key) -> (gates_pos, gates_quat, obstacles_pos)`` — a pure JAX function that - produces one random track per call. Shapes: ``(N, 3)``, ``(N, 4)`` xyzw, ``(N, 3)``. + ``sample_track(drone_pos, key) -> (gates_pos, gates_quat, obstacles_pos)``, a pure JAX + function that produces one random track per call. """ gates_z = jp.array(gates_z, dtype=jp.float32) obstacles_z = jp.array(obstacles_z, dtype=jp.float32) @@ -167,122 +170,94 @@ def build_random_track_fn( xmin, ymin = jp.array(pos_limit_low[:2], dtype=jp.float32) + border_margin xmax, ymax = jp.array(pos_limit_high[:2], dtype=jp.float32) - border_margin - # Precompute the placement grid (static). + width = float(pos_limit_high[0] - pos_limit_low[0]) - 2 * border_margin + height = float(pos_limit_high[1] - pos_limit_low[1]) - 2 * border_margin + grid_w = grid_res + grid_h = max(1, round(grid_res * height / width)) + xs = jp.linspace(xmin, xmax, grid_w) ys = jp.linspace(ymin, ymax, grid_h) grid = jp.stack(jp.meshgrid(xs, ys), axis=-1) # (H, W, 2) grid_flat = grid.reshape(-1, 2) - cell_dxy = jp.array([(xmax - xmin) / grid_w, (ymax - ymin) / grid_h]) - def _sample(weight: Array, key: Array) -> Array: - """Weighted sample from the placement grid with sub-cell jitter.""" - flat = weight.reshape(-1) - total = flat.sum() - p = jp.where(total > 0, flat / total, jp.ones_like(flat) / flat.size) - k_choice, k_jitter = jax.random.split(key) - pos = grid_flat[jax.random.choice(k_choice, flat.shape[0], p=p)] - return pos + (jax.random.uniform(k_jitter, (2,)) - 0.5) * cell_dxy + def _sample(clearance: Array, preference: Array, key: Array) -> Array: + """Sample a grid cell. - def _excl_circle(center: Array, radius: float) -> Array: - """Float mask: 1 where grid point is farther than `radius` from `center`.""" - return (jp.sum((grid - center) ** 2, axis=-1) > radius**2).astype(jp.float32) + Args: + clearance: The signed margin to every placed object. + preference: A soft 0/1 mask marking where the cell would ideally lie. + key: JAX PRNG key. + """ + valid = (clearance > 0).reshape(-1).astype(jp.float32) + preferred = valid * preference.reshape(-1) + # Case 1: cells that keep every exclusion radius *and* lie in the soft preference region. + # Case 2: preference leaves nothing -> any cell that still keeps every exclusion radius. + weight = jp.where(preferred.sum() > 0, preferred, valid) + # Case 3: grid saturated, no cell keeps all radii -> the single least-violating cell. + best = (jp.arange(valid.size) == jp.argmax(clearance)).astype(jp.float32) + weight = jp.where(weight.sum() > 0, weight, best) + return grid_flat[jax.random.choice(key, weight.shape[0], p=weight / weight.sum())] + + def _clearance(center: Array, radius: float) -> Array: + """Calculate the signed clearance per cell. + + Args: + center: The XY position of the object. + radius: The exclusion radius of the object. + """ + return jp.sqrt(jp.sum((grid - center) ** 2, axis=-1)) - radius def _corridor(from_xy: Array, to_xy: Array, width: float) -> Array: - """Float mask: 1 for grid points inside the corridor of given `width`.""" + """Create a grid mask for cells within `width` of the segment `from_xy -> to_xy`.""" v = to_xy - from_xy n = jp.linalg.norm(v) + 1e-8 - u = v / n - to_cell = grid - from_xy - proj = jp.sum(to_cell * u, axis=-1) - closest = from_xy + proj[..., None] * u - perp = jp.linalg.norm(grid - closest, axis=-1) + proj = jp.sum((grid - from_xy) * (v / n), axis=-1) + perp = jp.linalg.norm(grid - (from_xy + proj[..., None] * v / n), axis=-1) return ((perp < width) & (proj >= 0) & (proj <= n)).astype(jp.float32) - def generate(key: Array) -> tuple[Array, Array, Array]: - """Generate one random track. + def sample_track(drone_pos: Array, key: Array) -> tuple[Array, Array, Array]: + """Sample one random track around the given drone start position. Args: + drone_pos: Drone start position ``(3,)``. key: JAX PRNG key. Returns: - ``(gates_pos, gates_quat, obstacles_pos)`` with shapes ``(N, 3)``, ``(N, 4)`` (xyzw), - ``(N, 3)``. + ``(gates_pos (N, 3), gates_quat (N, 4), obstacles_pos (N, 3))``. """ - k_start, *sub_keys = jax.random.split(key, 1 + 3 * N) - k_gates = jp.array(sub_keys[:N]) - k_yaws = jp.array(sub_keys[N : 2 * N]) - k_obs = jp.array(sub_keys[2 * N :]) - - start_xy = jax.random.uniform( - k_start, - (2,), - minval=jp.array([xmin - border_margin, ymin - border_margin]), - maxval=jp.array([xmax + border_margin, ymax + border_margin]), - ) - start_excl = _excl_circle(start_xy, start_excl_r) - - ones = jp.ones((grid_h, grid_w), jp.float32) - init = ( - start_excl, # gate placement weight - ones, # cumulative gate exclusion (gate-to-gate) - ones, # cumulative gate exclusion (gate-to-obstacle) - ones, # cumulative obstacle exclusion - jp.zeros((N, 3), jp.float32), # placed gates: [x, y, yaw] - jp.zeros((N, 2), jp.float32), # placed obstacles: [x, y] - ) - - def place_one( - carry: tuple[Array, Array, Array, Array, Array, Array], i: int - ) -> tuple[tuple[Array, Array, Array, Array, Array, Array], None]: - gate_w, gate_excl, gate_excl_obs, obs_excl, gates, obstacles = carry - - # Place gate - gate_xy = _sample(gate_w * gate_excl, k_gates[i]) - prev_xy = jax.lax.cond( - i == 0, lambda _: start_xy, lambda _: gates[i - 1, :2], operand=None - ) - travel_dir = gate_xy - prev_xy - yaw_offset = jax.random.uniform(k_yaws[i], minval=-yaw_range, maxval=yaw_range) - yaw = (yaw_offset + jp.arctan2(travel_dir[1], travel_dir[0])) % (2 * jp.pi) - gates = gates.at[i].set(jp.array([gate_xy[0], gate_xy[1], yaw])) - - # Place obstacle inside the travel corridor - in_corridor = _corridor(prev_xy, gate_xy, obstacle_corridor_width) - obs_weight = in_corridor * gate_excl_obs * obs_excl * start_excl - obs_xy = _sample(obs_weight, k_obs[i]) - obstacles = obstacles.at[i].set(obs_xy) - - # Update exclusion zones - gate_excl_new = gate_excl * _excl_circle(gate_xy, gate_excl_r) - gate_excl_obs_new = gate_excl_obs * _excl_circle(gate_xy, obstacle_excl_r) - obs_excl_new = obs_excl * _excl_circle(obs_xy, obstacle_excl_r) - gate_corr = _corridor(prev_xy, gate_xy, gate_corridor_width) - gate_w_new = gate_w * gate_excl_new * (1.0 - gate_corr) - - return ( - gate_w_new, - gate_excl_new, - gate_excl_obs_new, - obs_excl_new, - gates, - obstacles, - ), None - - (_, _, _, _, gates, obstacles), _ = jax.lax.scan(place_one, init, jp.arange(N)) - - # Assemble output arrays with correct z-heights. + keys = jax.random.split(key, 3 * N) + k_gates, k_yaws, k_obs = keys[:N], keys[N : 2 * N], keys[2 * N :] + prev_xy = drone_pos[:2] + gate_clear = obs_clear = _clearance(prev_xy, drone_excl_r) + gate_soft = jp.ones((grid_h, grid_w), jp.float32) + + gates, obstacles = [], [] + for i in range(N): # N is usually small, so this unrolled loop instead of scan is fine. + gate_xy = _sample(gate_clear, gate_soft, k_gates[i]) + travel = gate_xy - prev_xy + yaw = jax.random.uniform(k_yaws[i], minval=-yaw_range, maxval=yaw_range) + yaw = (yaw + jp.arctan2(travel[1], travel[0])) % (2 * jp.pi) + gates.append(jp.array([gate_xy[0], gate_xy[1], yaw])) + gate_clear = jp.minimum(gate_clear, _clearance(gate_xy, gate_excl_r)) + obs_clear = jp.minimum(obs_clear, _clearance(gate_xy, obstacle_excl_r)) + gate_soft = gate_soft * (1.0 - _corridor(prev_xy, gate_xy, gate_corridor_width)) + + corridor = _corridor(prev_xy, gate_xy, obstacle_corridor_width) + obs_xy = _sample(obs_clear, corridor, k_obs[i]) + obstacles.append(obs_xy) + gate_clear = jp.minimum(gate_clear, _clearance(obs_xy, obstacle_excl_r)) + obs_clear = jp.minimum(obs_clear, _clearance(obs_xy, obstacle_obstacle_excl_r)) + prev_xy = gate_xy + + gates, obstacles = jp.stack(gates), jp.stack(obstacles) gates_pos = jp.concatenate([gates[:, :2], gates_z[:, None]], axis=-1) half_yaw = gates[:, 2] / 2.0 - # Pure-yaw quaternion (xyzw): roll=pitch=0 → [0, 0, sin(yaw/2), cos(yaw/2)] - gates_quat = jp.stack( - [jp.zeros_like(half_yaw), jp.zeros_like(half_yaw), jp.sin(half_yaw), jp.cos(half_yaw)], - axis=-1, - ) + zeros = jp.zeros_like(half_yaw) + gates_quat = jp.stack([zeros, zeros, jp.sin(half_yaw), jp.cos(half_yaw)], axis=-1) obstacles_pos = jp.concatenate([obstacles, obstacles_z[:, None]], axis=-1) - return gates_pos, gates_quat, obstacles_pos - return generate + return sample_track def build_full_track_randomization_fn( @@ -290,11 +265,7 @@ def build_full_track_randomization_fn( ) -> Callable[[EnvData, Array, Array], EnvData]: """Build a track randomization function that fully regenerates the track per world. - Unlike the perturbation-based approach, this generates an entirely new gate and obstacle layout - for every environment world that is being reset. Z-heights are fixed at the provided values. - Args: - n_objects: Number of gates (= number of obstacles). gates_z: Z-height for each gate, shape ``(n_objects,)``. obstacles_z: Z-height for each obstacle, shape ``(n_objects,)``. pos_limit_low: XY lower bounds of the arena ``[xmin, ymin]``. @@ -310,7 +281,8 @@ def build_full_track_randomization_fn( def randomize_track(data: EnvData, mask: Array, key: Array) -> EnvData: n_envs = data.gates_pos.shape[0] keys = jax.random.split(key, n_envs) - gates_pos, gates_quat, obstacles_pos = batched_generate(keys) + drones_pos = data.sim_data.states.pos[:, 0] # build each track around the first drone + gates_pos, gates_quat, obstacles_pos = batched_generate(drones_pos, keys) return leaf_replace( data, mask, diff --git a/tests/unit/envs/test_randomize.py b/tests/unit/envs/test_randomize.py new file mode 100644 index 000000000..fd3410a58 --- /dev/null +++ b/tests/unit/envs/test_randomize.py @@ -0,0 +1,70 @@ +"""Integration test: level3 produces collision-free tracks through the real reset pipeline.""" + +import os + +os.environ["SCIPY_ARRAY_API"] = "1" + +from pathlib import Path + +import gymnasium +import numpy as np +from jax import Array +from ml_collections import ConfigDict + +import lsy_drone_racing # noqa: F401 (registers the gymnasium environments) +from lsy_drone_racing.utils import load_config + +CONFIG_PATH = Path(__file__).parents[3] / "config" +N_WORLDS = 1000 + + +def min_distance_between(a: Array, b: Array) -> float: + """Smallest XY distance between sets `a` and `b`, minimized over all worlds.""" + return np.linalg.norm(a[:, :, None] - b[:, None], axis=-1).min() + + +def min_distance_within(points: Array) -> float: + """Smallest XY distance between two distinct points of the same set, over all worlds.""" + distances = np.linalg.norm(points[:, :, None] - points[:, None], axis=-1) + self_pairs = np.arange(points.shape[1]) + distances[:, self_pairs, self_pairs] = np.inf + return distances.min() + + +def max_xy_shift(randomization: ConfigDict) -> float: + """Largest XY displacement a uniform position randomization can apply to an object.""" + kwargs = randomization.kwargs + return float(np.hypot(*np.maximum(np.abs(kwargs.minval[:2]), np.abs(kwargs.maxval[:2])))) + + +def test_level3_tracks_are_collision_free(): + config = load_config(CONFIG_PATH / "level3.toml") + env = gymnasium.make_vec( + "DroneRacing-v0", + num_envs=N_WORLDS, + freq=config.env.freq, + sim_config=config.sim, + sensor_range=config.env.sensor_range, + track=config.env.track, + disturbances=config.env.get("disturbances"), + randomizations=config.env.get("randomizations"), + seed=config.env.seed, + ) + env = gymnasium.wrappers.vector.JaxToNumpy(env) + env.reset() + data = env.unwrapped.data # ground-truth state, not the sensor-masked observation + env.close() + + gates = np.asarray(data.gates_pos)[..., :2] + obstacles = np.asarray(data.obstacles_pos)[..., :2] + drone = np.asarray(data.sim_data.states.pos)[:, 0, :2][:, None, :] + + # Gates and obstacles are perturbed after track generation. Relax each exclusion radius by the + # maximum displacement those randomizations can apply. + gate_displacement = max_xy_shift(config.env.randomizations.gate_pos) + obstacle_displacement = max_xy_shift(config.env.randomizations.obstacle_pos) + assert min_distance_within(gates) >= 1.0 - 2 * gate_displacement + assert min_distance_within(obstacles) >= 0.3 - 2 * obstacle_displacement + assert min_distance_between(gates, obstacles) >= 0.8 - gate_displacement - obstacle_displacement + assert min_distance_between(drone, gates) >= 1.0 - gate_displacement + assert min_distance_between(drone, obstacles) >= 1.0 - obstacle_displacement From 2dcc1e5edd93e326e73f309008e2e193e52780f6 Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Tue, 30 Jun 2026 18:41:41 +0200 Subject: [PATCH 2/2] Remove gate corridors --- lsy_drone_racing/envs/randomize.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/lsy_drone_racing/envs/randomize.py b/lsy_drone_racing/envs/randomize.py index c20abec6d..209a651fd 100644 --- a/lsy_drone_racing/envs/randomize.py +++ b/lsy_drone_racing/envs/randomize.py @@ -129,7 +129,6 @@ def build_random_track_fn( gate_excl_r: float = 1.0, obstacle_excl_r: float = 0.8, obstacle_obstacle_excl_r: float = 0.3, - gate_corridor_width: float = 0.5, obstacle_corridor_width: float = 0.5, yaw_range: float = 0.75, grid_res: int = 40, @@ -152,7 +151,6 @@ def build_random_track_fn( gate_excl_r: Min distance [m] between consecutive gates. obstacle_excl_r: Min distance [m] from gates to obstacles. obstacle_obstacle_excl_r: Min distance [m] between obstacles. - gate_corridor_width: Half-width [m] of the flight corridor masked out for gate placement. obstacle_corridor_width: Half-width [m] of the corridor used for obstacle placement. yaw_range: Maximum yaw offset [rad] from the travel direction for gate orientation. grid_res: Number of grid nodes along the x-axis. The y-axis count is derived from the arena @@ -229,18 +227,17 @@ def sample_track(drone_pos: Array, key: Array) -> tuple[Array, Array, Array]: k_gates, k_yaws, k_obs = keys[:N], keys[N : 2 * N], keys[2 * N :] prev_xy = drone_pos[:2] gate_clear = obs_clear = _clearance(prev_xy, drone_excl_r) - gate_soft = jp.ones((grid_h, grid_w), jp.float32) + no_preference = jp.ones((grid_h, grid_w), jp.float32) gates, obstacles = [], [] for i in range(N): # N is usually small, so this unrolled loop instead of scan is fine. - gate_xy = _sample(gate_clear, gate_soft, k_gates[i]) + gate_xy = _sample(gate_clear, no_preference, k_gates[i]) travel = gate_xy - prev_xy yaw = jax.random.uniform(k_yaws[i], minval=-yaw_range, maxval=yaw_range) yaw = (yaw + jp.arctan2(travel[1], travel[0])) % (2 * jp.pi) gates.append(jp.array([gate_xy[0], gate_xy[1], yaw])) gate_clear = jp.minimum(gate_clear, _clearance(gate_xy, gate_excl_r)) obs_clear = jp.minimum(obs_clear, _clearance(gate_xy, obstacle_excl_r)) - gate_soft = gate_soft * (1.0 - _corridor(prev_xy, gate_xy, gate_corridor_width)) corridor = _corridor(prev_xy, gate_xy, obstacle_corridor_width) obs_xy = _sample(obs_clear, corridor, k_obs[i])