Johnny Fernandes
201 lines
7.3 KiB
Python
201 lines
7.3 KiB
Python
"""Universal shepherd teacher — Strömbom core + mecanum omega + straggler recovery.
|
|
|
|
The core collect/drive logic is **identical** to :mod:`strombom` (same
|
|
``F_FACTOR``, ``DELTA_COLLECT``, ``DELTA_DRIVE`` thresholds and target
|
|
computation) so it inherits the proven ~100 % success rate at n ≤ 8.
|
|
Two additions make it useful as a universal teacher:
|
|
|
|
1. **Omega for mecanum.** When ``drive_mode="mecanum"``, the teacher
|
|
outputs a non-zero ``omega`` channel so the dog **faces the
|
|
direction of travel**. During collect the dog faces the target
|
|
sheep; during drive it faces the pen. This gives the BC student a
|
|
real rotation signal to learn from.
|
|
|
|
2. **Last-straggler recovery.** When exactly one sheep remains active
|
|
and it is near the gate, the dog positions itself behind that
|
|
straggler (opposite the gate) and pushes it straight through. This
|
|
handles the edge case where the last sheep circles the gate posts.
|
|
|
|
Call signature::
|
|
|
|
vx, vy, omega, mode = compute_action(
|
|
dog_xy, dog_heading, sheep_positions, pen_target,
|
|
drive_mode="differential",
|
|
)
|
|
|
|
For differential drive ``omega`` is always 0.0 and can be ignored.
|
|
"""
|
|
|
|
import math
|
|
|
|
from herding.world.geometry import (
|
|
PEN_ENTRY, GATE_X, GATE_Y, in_pen,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tuning constants — match Strömbom exactly for proven success rates.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
F_FACTOR = 4.0 # collect/drive threshold scaled by √n
|
|
DELTA_COLLECT = 1.5 # standoff behind the furthest sheep
|
|
DELTA_DRIVE = 2.0 # standoff behind flock CoM
|
|
|
|
# Omega gain for mecanum (how strongly the dog turns to face target)
|
|
OMEGA_GAIN = 0.6
|
|
|
|
# Recovery: push small flocks (≤ RECOVERY_MAX_N) through the gate one
|
|
# sheep at a time. n=1 alone is not enough — at n=2..3 on the round
|
|
# field the flock is too small to self-cohere through the 3 m gate but
|
|
# the standard collect/drive standoff just orbits them. Push the sheep
|
|
# nearest the gate first; once it pens, the rule re-applies to the next.
|
|
RECOVERY_MAX_N = 3
|
|
RECOVERY_GATE_DIST = 8.0 # only when target sheep is this close to gate
|
|
RECOVERY_PUSH_DIST = 1.2 # stand-off behind sheep, away from gate
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
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 _angle_diff(a, b):
|
|
"""Signed shortest angular difference a - b, in [-π, π]."""
|
|
return math.atan2(math.sin(a - b), math.cos(a - b))
|
|
|
|
|
|
def _gate_center():
|
|
"""Centre of the gate opening."""
|
|
return (0.5 * (GATE_X[0] + GATE_X[1]), GATE_Y)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Core teacher
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def compute_action(dog_xy, dog_heading, sheep_positions,
|
|
pen_target=PEN_ENTRY, drive_mode="differential"):
|
|
"""Return ``(vx, vy, omega, mode)``.
|
|
|
|
Parameters
|
|
----------
|
|
dog_xy : (float, float)
|
|
Dog position in world frame.
|
|
dog_heading : float
|
|
Dog heading in world frame (rad), 0 = +x axis.
|
|
sheep_positions : dict[str, (float, float)]
|
|
Visible sheep positions.
|
|
pen_target : (float, float)
|
|
Centre of the pen gate (defaults to geometry.PEN_ENTRY).
|
|
drive_mode : str
|
|
``"differential"`` or ``"mecanum"``.
|
|
|
|
Returns
|
|
-------
|
|
vx, vy : float
|
|
Velocity intent in [-1, 1].
|
|
omega : float
|
|
Yaw intent in [-1, 1] (0 for differential).
|
|
mode : str
|
|
Phase label: ``"idle"``, ``"collect"``, ``"drive"``, ``"recovery"``.
|
|
"""
|
|
active = [(x, y) for (x, y) in sheep_positions.values()
|
|
if _is_active(x, y)]
|
|
if not active:
|
|
return 0.0, 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)
|
|
|
|
# ---- Small-flock recovery (push sheep through the gate one by one) ----
|
|
# Triggers when the active flock is small (≤ RECOVERY_MAX_N) and the
|
|
# sheep nearest the gate is close enough that direct pushing works.
|
|
# For larger flocks the standard collect/drive logic handles them.
|
|
gc = _gate_center()
|
|
if n <= RECOVERY_MAX_N:
|
|
# Pick the sheep closest to the gate as the recovery target —
|
|
# finishing that one first reduces the active count and lets the
|
|
# remaining sheep get their own recovery turn.
|
|
gate_dists = [math.hypot(p[0] - gc[0], p[1] - gc[1]) for p in active]
|
|
target_idx = min(range(n), key=lambda i: gate_dists[i])
|
|
sx, sy = active[target_idx]
|
|
d_to_gate = gate_dists[target_idx]
|
|
if d_to_gate < RECOVERY_GATE_DIST:
|
|
dx_g = sx - gc[0]
|
|
dy_g = sy - gc[1]
|
|
d_g = math.hypot(dx_g, dy_g)
|
|
if d_g > 0.3:
|
|
ux, uy = dx_g / d_g, dy_g / d_g
|
|
else:
|
|
ux, uy = 0.0, 1.0
|
|
tx = sx + RECOVERY_PUSH_DIST * ux
|
|
ty = sy + RECOVERY_PUSH_DIST * uy
|
|
ax, ay = _unit(tx - dog_xy[0], ty - dog_xy[1])
|
|
mode = "recovery"
|
|
face_target = (sx, sy)
|
|
omega = 0.0
|
|
if drive_mode == "mecanum":
|
|
desired = math.atan2(
|
|
face_target[1] - dog_xy[1],
|
|
face_target[0] - dog_xy[0],
|
|
)
|
|
err = _angle_diff(desired, dog_heading)
|
|
omega = max(-1.0, min(1.0, OMEGA_GAIN * err / math.pi))
|
|
return ax, ay, omega, mode
|
|
|
|
# ---- Standard Strömbom collect/drive (proven core) ----
|
|
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"
|
|
face_target = (sx, sy)
|
|
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"
|
|
face_target = pen_target
|
|
|
|
ax, ay = _unit(tx - dog_xy[0], ty - dog_xy[1])
|
|
|
|
# ---- Omega (mecanum only) ----
|
|
omega = 0.0
|
|
if drive_mode == "mecanum" and mode != "idle":
|
|
desired_heading = math.atan2(
|
|
face_target[1] - dog_xy[1],
|
|
face_target[0] - dog_xy[0],
|
|
)
|
|
err = _angle_diff(desired_heading, dog_heading)
|
|
omega = max(-1.0, min(1.0, OMEGA_GAIN * err / math.pi))
|
|
|
|
return ax, ay, omega, mode
|
|
|
|
|
|
def compute_action_diff(dog_xy, dog_heading, sheep_positions,
|
|
pen_target=PEN_ENTRY):
|
|
"""Compatibility wrapper returning ``(vx, vy, mode)`` — same as Strömbom.
|
|
|
|
Use this when plugging into existing differential-drive code that
|
|
doesn't expect omega.
|
|
"""
|
|
vx, vy, _omega, mode = compute_action(
|
|
dog_xy, dog_heading, sheep_positions, pen_target,
|
|
drive_mode="differential",
|
|
)
|
|
return vx, vy, mode
|