Johnny Fernandes
99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
"""Sequential single-target shepherd dog algorithm.
|
||
|
||
Strömbom drives the flock's centre of mass; with N sheep and a narrow
|
||
3 m gate, this fails because the flock is wider than the gate and CoM
|
||
driving abandons stragglers. Real sheepdogs solve this differently:
|
||
they pick *one* sheep at a time, drive it through, return for the next.
|
||
|
||
This module implements that "pin-and-push" approach.
|
||
|
||
Algorithm (one step):
|
||
1. Active sheep = those still in the field (not yet penned).
|
||
2. Target = the active sheep currently closest to the pen entry.
|
||
3. Drive position = ``target + Δ · unit(target − pen_entry)`` —
|
||
directly behind the target relative to the goal.
|
||
4. Output unit vector pointing the dog at the drive position.
|
||
|
||
Once the target crosses the gate it latches as penned and is removed
|
||
from the active set; the next-closest unpenned sheep becomes the
|
||
target. The algorithm naturally "queues" sheep through the gate.
|
||
|
||
Empirically (with our flocking dynamics) this scales linearly with
|
||
flock size and works up to at least n=10 within a 15 000-step budget.
|
||
"""
|
||
|
||
import math
|
||
|
||
from herding.geometry import GATE_Y, PEN_ENTRY, in_pen
|
||
|
||
|
||
DELTA_DRIVE = 1.5 # standoff behind the target sheep
|
||
APPROACH_GAIN = 1.0 # action magnitude scale (1 = full speed)
|
||
|
||
|
||
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:
|
||
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)`` where mode encodes the current target.
|
||
|
||
Compatible with the Strömbom call signature so it can be drop-in
|
||
swapped in the dog controller and the env's imitation reward.
|
||
"""
|
||
active = [(name, x, y) for name, (x, y) in sheep_positions.items()
|
||
if _is_active(x, y)]
|
||
if not active:
|
||
return 0.0, 0.0, "idle"
|
||
|
||
# Pick target = sheep closest to pen entry. Stable choice: as one
|
||
# sheep approaches and crosses the gate it stays the target until
|
||
# latched; then the next-closest takes over.
|
||
name, sx, sy = min(
|
||
active,
|
||
key=lambda s: math.hypot(s[1] - pen_target[0], s[2] - pen_target[1]),
|
||
)
|
||
|
||
# Drive position behind the target along the (target → pen) line.
|
||
ux, uy = _unit(sx - pen_target[0], sy - pen_target[1])
|
||
tx = sx + DELTA_DRIVE * ux
|
||
ty = sy + DELTA_DRIVE * uy
|
||
|
||
ax, ay = _unit(tx - dog_xy[0], ty - dog_xy[1])
|
||
return APPROACH_GAIN * ax, APPROACH_GAIN * ay, f"drive:{name}"
|
||
|
||
|
||
def compute_action_debug(dog_xy, sheep_positions, pen_target=PEN_ENTRY):
|
||
"""Debug variant returning ``(vx, vy, mode, debug_dict)``."""
|
||
active = [(name, x, y) for name, (x, y) in sheep_positions.items()
|
||
if _is_active(x, y)]
|
||
if not active:
|
||
return 0.0, 0.0, "idle", {
|
||
"n_active": 0, "target_name": "",
|
||
"target_x": 0.0, "target_y": 0.0,
|
||
"drive_x": dog_xy[0], "drive_y": dog_xy[1],
|
||
}
|
||
|
||
name, sx, sy = min(
|
||
active,
|
||
key=lambda s: math.hypot(s[1] - pen_target[0], s[2] - pen_target[1]),
|
||
)
|
||
|
||
ux, uy = _unit(sx - pen_target[0], sy - pen_target[1])
|
||
tx = sx + DELTA_DRIVE * ux
|
||
ty = sy + DELTA_DRIVE * uy
|
||
ax, ay = _unit(tx - dog_xy[0], ty - dog_xy[1])
|
||
|
||
return APPROACH_GAIN * ax, APPROACH_GAIN * ay, f"drive:{name}", {
|
||
"n_active": len(active), "target_name": name,
|
||
"target_x": sx, "target_y": sy,
|
||
"drive_x": tx, "drive_y": ty,
|
||
}
|