Johnny Fernandes
96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
"""Strömbom (2014) collect/drive heuristic for the shepherd dog.
|
|
|
|
When the flock is scattered (max radius > F_FACTOR · √n) the dog moves
|
|
to a point behind the furthest sheep and pushes it back toward the
|
|
flock CoM. Otherwise it drives, parking behind the CoM relative to
|
|
the pen target. Returns a unit-vector intent ``(vx, vy, mode)``.
|
|
|
|
Reference: Strömbom et al. 2014, "Solving the shepherding problem."
|
|
"""
|
|
|
|
import math
|
|
|
|
from herding.world.geometry import PEN_ENTRY, GATE_Y, in_pen
|
|
|
|
F_FACTOR = 4.0 # collect/drive threshold scaled by √n
|
|
DELTA_COLLECT = 1.5 # drive-position offset behind the furthest sheep
|
|
DELTA_DRIVE = 2.0 # drive-position offset behind the flock CoM
|
|
|
|
|
|
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 still in the field counts; one south of the gate doesn't."""
|
|
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}."""
|
|
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 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 behind the CoM, opposite the pen.
|
|
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):
|
|
"""``compute_action`` plus a small debug dict (CoM, target, radius)."""
|
|
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
|