Johnny Fernandes
54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
"""Shared low-level control helpers used by every dog mode.
|
|
|
|
Centralised here so the BC student, Strömbom, Sequential, and the DAgger
|
|
teacher all apply identical post-processing to their action outputs.
|
|
The downstream wheel-velocity layer (``herding.diffdrive``) is unchanged.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
|
|
# Speed-modulation: scale action magnitude down when close to the
|
|
# nearest sheep. Stops the dog from charging in at full speed and
|
|
# scattering the flock. Action norm linearly ramps from MIN_SPEED at
|
|
# distance 0 to 1.0 at SLOW_NEAR_SHEEP.
|
|
SLOW_NEAR_SHEEP = 2.5
|
|
MIN_SPEED = 0.30
|
|
|
|
|
|
def modulate_speed_near_sheep(
|
|
vx: float, vy: float,
|
|
dog_xy: tuple[float, float],
|
|
sheep_positions,
|
|
slow_dist: float = SLOW_NEAR_SHEEP,
|
|
min_scale: float = MIN_SPEED,
|
|
) -> tuple[float, float]:
|
|
"""Scale (vx, vy) magnitude down when close to the nearest sheep.
|
|
|
|
``sheep_positions`` accepts either a ``{name: (x, y)}`` dict
|
|
(matching what the trackers emit) or an iterable of ``(x, y)``
|
|
tuples. Empty input → action returned unchanged.
|
|
|
|
The intent direction is preserved; only magnitude is reduced. With
|
|
``slow_dist=2.5`` and ``min_scale=0.3``, an action that started at
|
|
norm 1 is multiplied by 0.3 right next to a sheep, by 0.65 at 1 m
|
|
away, and by 1.0 once the nearest sheep is ≥ 2.5 m off.
|
|
"""
|
|
if not sheep_positions:
|
|
return vx, vy
|
|
if hasattr(sheep_positions, "values"):
|
|
positions = sheep_positions.values()
|
|
else:
|
|
positions = sheep_positions
|
|
nearest = float("inf")
|
|
for sx, sy in positions:
|
|
d = math.hypot(sx - dog_xy[0], sy - dog_xy[1])
|
|
if d < nearest:
|
|
nearest = d
|
|
if nearest >= slow_dist or nearest == float("inf"):
|
|
return vx, vy
|
|
scale = min_scale + (1.0 - min_scale) * (nearest / slow_dist)
|
|
return vx * scale, vy * scale
|