Checkpoint 6
This commit is contained in:
Johnny Fernandes
@@ -0,0 +1,193 @@
|
||||
"""Fast 2D LiDAR simulator for the Gymnasium env.
|
||||
|
||||
Raycasts against:
|
||||
* **Sheep** — discs of radius ``SHEEP_RADIUS``.
|
||||
* **Static world geometry** — axis-aligned wall segments and gate
|
||||
posts taken from ``worlds/field.wbt``. Without these, demos
|
||||
collected in-env would never include the false-positive clusters
|
||||
Webots produces from the stone walls and gate-post boxes, and the
|
||||
BC student trained on those demos collapses on deployment.
|
||||
|
||||
Returns a range array matching the Webots Lidar device on the dog
|
||||
(see ``protos/ShepherdDog.proto``: 180 rays, 140° FOV centred on
|
||||
forward, 12 m max range, 5 mm noise).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
# Match protos/ShepherdDog.proto Lidar device.
|
||||
LIDAR_N_RAYS = 180
|
||||
LIDAR_FOV = 2.44 # rad ≈ 140°
|
||||
LIDAR_MAX_RANGE = 12.0
|
||||
LIDAR_NOISE = 0.005 # m, gaussian std
|
||||
|
||||
# Sheep modelled as a vertical cylinder; this is the horizontal-section
|
||||
# radius the LiDAR plane intersects. Tuned to the proto sheep (~0.45 m
|
||||
# body length). The exact value is not load-bearing — the perception
|
||||
# clusterer is range-tolerant.
|
||||
SHEEP_RADIUS = 0.30
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Static world geometry — must match worlds/field.wbt
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Vertical walls: (x, y_min, y_max). Field east/west walls and the two
|
||||
# pen side walls are visible through the open gate.
|
||||
_VERTICAL_WALLS = (
|
||||
( 15.0, -15.0, 15.0), # field east
|
||||
(-15.0, -15.0, 15.0), # field west
|
||||
( 10.0, -22.0, -15.0), # pen west
|
||||
( 13.0, -22.0, -15.0), # pen east
|
||||
)
|
||||
|
||||
# Horizontal walls: (y, x_min, x_max). South wall is split by the 3 m
|
||||
# gate at x ∈ [10, 13]; the pen south wall closes the back of the pen.
|
||||
_HORIZONTAL_WALLS = (
|
||||
( 15.0, -15.0, 15.0), # field north
|
||||
(-15.0, -15.0, 10.0), # field south-west of gate
|
||||
(-15.0, 13.0, 15.0), # field south-east of gate
|
||||
(-22.0, 10.0, 13.0), # pen south
|
||||
)
|
||||
|
||||
# Gate posts and field corner pillars treated as vertical cylinders at
|
||||
# LiDAR height. Radius 0.25 m comes from the 0.44 × 0.44 m boxes in the
|
||||
# wbt — close enough to a circular cross-section for this purpose.
|
||||
_POSTS_XY = np.array([
|
||||
( 10.0, -15.0), # west gate post
|
||||
( 13.0, -15.0), # east gate post
|
||||
( 15.0, 15.0), # NE field corner
|
||||
( 15.0, -15.0), # SE field corner
|
||||
(-15.0, 15.0), # NW field corner
|
||||
(-15.0, -15.0), # SW field corner
|
||||
], dtype=np.float64)
|
||||
POST_RADIUS = 0.25
|
||||
|
||||
|
||||
def ray_angles(n: int = LIDAR_N_RAYS, fov: float = LIDAR_FOV) -> np.ndarray:
|
||||
"""Local-frame ray angles, sweeping from +fov/2 to -fov/2.
|
||||
|
||||
Convention: angle is measured CCW from the dog's forward axis. Ray 0
|
||||
points to the dog's left, last ray to the right. Webots' default
|
||||
Lidar sweep matches this.
|
||||
"""
|
||||
return np.linspace(fov / 2.0, -fov / 2.0, n, dtype=np.float64)
|
||||
|
||||
|
||||
# Cached so we don't rebuild every step.
|
||||
_ANGLES = ray_angles()
|
||||
_COS = np.cos(_ANGLES)
|
||||
_SIN = np.sin(_ANGLES)
|
||||
|
||||
|
||||
def _raycast_static(
|
||||
ox: float, oy: float, cos_w: np.ndarray, sin_w: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
"""Per-ray distance to nearest wall or post hit (∞ if none).
|
||||
|
||||
Walls are axis-aligned line segments; for each ray we compute t at
|
||||
which it crosses the wall's constant-coord plane and check the
|
||||
other coord lies in the segment. Posts are circles; same disc
|
||||
intersection as for sheep.
|
||||
"""
|
||||
n_rays = cos_w.shape[0]
|
||||
best = np.full(n_rays, np.inf, dtype=np.float64)
|
||||
|
||||
EPS = 1e-3
|
||||
safe_cos = np.where(np.abs(cos_w) < 1e-9, 1e-9, cos_w)
|
||||
safe_sin = np.where(np.abs(sin_w) < 1e-9, 1e-9, sin_w)
|
||||
|
||||
# Vertical walls (x = const)
|
||||
for wx, ymin, ymax in _VERTICAL_WALLS:
|
||||
t = (wx - ox) / safe_cos
|
||||
y_at = oy + t * sin_w
|
||||
valid = (t > EPS) & (y_at >= ymin - EPS) & (y_at <= ymax + EPS)
|
||||
cand = np.where(valid, t, np.inf)
|
||||
np.minimum(best, cand, out=best)
|
||||
|
||||
# Horizontal walls (y = const)
|
||||
for wy, xmin, xmax in _HORIZONTAL_WALLS:
|
||||
t = (wy - oy) / safe_sin
|
||||
x_at = ox + t * cos_w
|
||||
valid = (t > EPS) & (x_at >= xmin - EPS) & (x_at <= xmax + EPS)
|
||||
cand = np.where(valid, t, np.inf)
|
||||
np.minimum(best, cand, out=best)
|
||||
|
||||
# Posts (treat as discs)
|
||||
if _POSTS_XY.size:
|
||||
px = _POSTS_XY[:, 0] - ox
|
||||
py = _POSTS_XY[:, 1] - oy
|
||||
t_post = np.outer(px, cos_w) + np.outer(py, sin_w) # (P, N)
|
||||
d2 = (px ** 2 + py ** 2)[:, None] # (P, 1)
|
||||
perp2 = d2 - t_post ** 2
|
||||
R2 = POST_RADIUS ** 2
|
||||
hit = (perp2 < R2) & (t_post > 0.0)
|
||||
half = np.sqrt(np.clip(R2 - perp2, 0.0, None))
|
||||
cand = np.where(hit, t_post - half, np.inf)
|
||||
nearest = cand.min(axis=0)
|
||||
np.minimum(best, nearest, out=best)
|
||||
|
||||
return best
|
||||
|
||||
|
||||
def simulate_scan(
|
||||
dog_x: float, dog_y: float, dog_heading: float,
|
||||
sheep_xy: list[tuple[float, float]],
|
||||
noise: float = LIDAR_NOISE,
|
||||
max_range: float = LIDAR_MAX_RANGE,
|
||||
rng: np.random.Generator | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Return a (N,) float32 range array. No-hit entries equal ``max_range``.
|
||||
|
||||
``sheep_xy`` is the list of (x, y) world positions of every sheep in
|
||||
the scene (penned and active). Static world geometry (walls and
|
||||
posts) is also raycast so demos contain the same false-positive
|
||||
clusters Webots produces.
|
||||
"""
|
||||
n_rays = _ANGLES.shape[0]
|
||||
|
||||
ch, sh = math.cos(dog_heading), math.sin(dog_heading)
|
||||
cos_w = ch * _COS - sh * _SIN
|
||||
sin_w = sh * _COS + ch * _SIN
|
||||
|
||||
# Walls + posts
|
||||
best = _raycast_static(dog_x, dog_y, cos_w, sin_w)
|
||||
|
||||
# Sheep discs
|
||||
if sheep_xy:
|
||||
sx = np.asarray([p[0] for p in sheep_xy], dtype=np.float64) - dog_x
|
||||
sy = np.asarray([p[1] for p in sheep_xy], dtype=np.float64) - dog_y
|
||||
t = np.outer(sx, cos_w) + np.outer(sy, sin_w)
|
||||
s_dist2 = (sx ** 2 + sy ** 2)[:, None]
|
||||
perp2 = s_dist2 - t ** 2
|
||||
R2 = SHEEP_RADIUS ** 2
|
||||
hit = (perp2 < R2) & (t > 0.0)
|
||||
half = np.sqrt(np.clip(R2 - perp2, 0.0, None))
|
||||
candidate = np.where(hit, t - half, np.inf)
|
||||
nearest = candidate.min(axis=0)
|
||||
np.minimum(best, nearest, out=best)
|
||||
|
||||
# Clip to LIDAR_MAX_RANGE; entries that never got a hit stay at inf
|
||||
# → clipped down to max_range like the real Webots device.
|
||||
ranges = np.minimum(best, max_range).astype(np.float32)
|
||||
return _add_noise(ranges, noise, rng, max_range)
|
||||
|
||||
|
||||
def _add_noise(ranges: np.ndarray, sigma: float,
|
||||
rng: np.random.Generator | None, max_range: float) -> np.ndarray:
|
||||
if sigma <= 0.0:
|
||||
return ranges
|
||||
if rng is None:
|
||||
rng = np.random.default_rng()
|
||||
hit_mask = ranges < max_range - 1e-3
|
||||
n_hit = int(hit_mask.sum())
|
||||
if n_hit:
|
||||
ranges = ranges.copy()
|
||||
ranges[hit_mask] += rng.normal(0.0, sigma, size=n_hit).astype(np.float32)
|
||||
np.clip(ranges, 0.0, max_range, out=ranges)
|
||||
return ranges
|
||||
Reference in New Issue
Block a user