Checkpoint 6

herding/control/modulation.py

@@ -0,0 +1,53 @@
+"""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