Fix random track generation for config level 3

lsy_drone_racing/envs/randomize.py

@@ -125,39 +125,40 @@ 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,
+    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.
-        gate_corridor_width: Half-width [m] of the flight corridor masked out for gate placement.
+        obstacle_excl_r: Min distance [m] from gates to obstacles.
+        obstacle_obstacle_excl_r: Min distance [m] between obstacles.
         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,134 +168,101 @@ 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)
+        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, 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))
+
+            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(
     gates_z: Array, obstacles_z: Array, pos_limit_low: Array, pos_limit_high: Array
 ) -> 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 +278,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,

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