Johnny Fernandes
115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
"""Strömbom collect/drive heuristic for the shepherd dog.
|
|
|
|
Adapted from the original ``controllers/shepherd_dog/strombom.py`` and
|
|
updated for the external pen layout. Used as a baseline controller and
|
|
as the fallback when the RL policy isn't available.
|
|
|
|
Reference: Strömbom et al. 2014, "Solving the shepherding problem".
|
|
"""
|
|
|
|
import math
|
|
|
|
from herding.geometry import PEN_ENTRY, GATE_Y, in_pen
|
|
|
|
# Algorithm parameters. DELTA_DRIVE / DELTA_COLLECT were tightened from
|
|
# the original (4.0 / 2.5) because the new external pen sits ~26 m from
|
|
# typical sheep spawn locations — at the old 4 m standoff, the flee force
|
|
# (quadratic ramp, 3.7 at 4 m vs ~10 at 2 m) couldn't move sheep through
|
|
# the path inside the 3000-step episode budget.
|
|
#
|
|
# F_FACTOR was 2.0 in the original Strömbom paper; raised to 4.0 here so
|
|
# the dog stays in *drive* mode much longer. With our tighter cohesion
|
|
# (flocking_sim.py), partially-collected flocks consolidate naturally
|
|
# during a drive, and we don't waste 80% of the time budget on a slow
|
|
# "collect" pre-phase.
|
|
F_FACTOR = 4.0
|
|
DELTA_COLLECT = 1.5
|
|
DELTA_DRIVE = 2.0
|
|
|
|
|
|
def _unit(x, y):
|
|
d = math.hypot(x, y)
|
|
if d < 1e-6:
|
|
return 0.0, 0.0
|
|
return x / d, y / d
|
|
|
|
|
|
def _is_active(x, y) -> bool:
|
|
"""A sheep is "active" if it's still in the field — not in or below
|
|
the gate plane (we treat anything south of the gate as committed to
|
|
the pen and stop trying to herd it)."""
|
|
return (not in_pen(x, y)) and y > GATE_Y
|
|
|
|
|
|
def compute_action(dog_xy, sheep_positions, pen_target=PEN_ENTRY):
|
|
"""Return ``(vx, vy, mode)`` — mode in {idle, collect, drive}.
|
|
|
|
``sheep_positions`` is a ``{name: (x, y)}`` mapping (matches the
|
|
Webots controller's representation).
|
|
"""
|
|
active = [(x, y) for (x, y) in sheep_positions.values() if _is_active(x, y)]
|
|
if not active:
|
|
return 0.0, 0.0, "idle"
|
|
|
|
n = len(active)
|
|
com_x = sum(p[0] for p in active) / n
|
|
com_y = sum(p[1] for p in active) / n
|
|
dists = [math.hypot(p[0] - com_x, p[1] - com_y) for p in active]
|
|
radius = max(dists)
|
|
|
|
if radius > F_FACTOR * math.sqrt(n):
|
|
# Collect: aim at a point behind the furthest sheep, opposite the CoM.
|
|
idx = max(range(n), key=lambda i: dists[i])
|
|
sx, sy = active[idx]
|
|
ux, uy = _unit(sx - com_x, sy - com_y)
|
|
tx, ty = sx + DELTA_COLLECT * ux, sy + DELTA_COLLECT * uy
|
|
mode = "collect"
|
|
else:
|
|
# Drive: aim at a point behind the flock CoM relative to the goal.
|
|
ux, uy = _unit(com_x - pen_target[0], com_y - pen_target[1])
|
|
tx, ty = com_x + DELTA_DRIVE * ux, com_y + DELTA_DRIVE * uy
|
|
mode = "drive"
|
|
|
|
ax, ay = _unit(tx - dog_xy[0], ty - dog_xy[1])
|
|
return ax, ay, mode
|
|
|
|
|
|
def compute_action_debug(dog_xy, sheep_positions, pen_target=PEN_ENTRY):
|
|
"""Variant of compute_action that also returns a small debug dict.
|
|
|
|
Kept for parity with the legacy controller's CSV logger.
|
|
"""
|
|
active = [(x, y) for (x, y) in sheep_positions.values() if _is_active(x, y)]
|
|
if not active:
|
|
return 0.0, 0.0, "idle", {
|
|
"n_active": 0, "radius": 0.0, "threshold": 0.0,
|
|
"com_x": 0.0, "com_y": 0.0,
|
|
"target_x": dog_xy[0], "target_y": dog_xy[1],
|
|
}
|
|
|
|
n = len(active)
|
|
com_x = sum(p[0] for p in active) / n
|
|
com_y = sum(p[1] for p in active) / n
|
|
dists = [math.hypot(p[0] - com_x, p[1] - com_y) for p in active]
|
|
radius = max(dists)
|
|
threshold = F_FACTOR * math.sqrt(n)
|
|
|
|
if radius > threshold:
|
|
idx = max(range(n), key=lambda i: dists[i])
|
|
sx, sy = active[idx]
|
|
ux, uy = _unit(sx - com_x, sy - com_y)
|
|
tx, ty = sx + DELTA_COLLECT * ux, sy + DELTA_COLLECT * uy
|
|
mode = "collect"
|
|
else:
|
|
ux, uy = _unit(com_x - pen_target[0], com_y - pen_target[1])
|
|
tx, ty = com_x + DELTA_DRIVE * ux, com_y + DELTA_DRIVE * uy
|
|
mode = "drive"
|
|
|
|
ax, ay = _unit(tx - dog_xy[0], ty - dog_xy[1])
|
|
dbg = {
|
|
"n_active": n, "radius": radius, "threshold": threshold,
|
|
"com_x": com_x, "com_y": com_y,
|
|
"target_x": tx, "target_y": ty,
|
|
}
|
|
return ax, ay, mode, dbg
|