Checkpoint 8

herding/control/active_scan.py

@@ -33,7 +33,11 @@ class ActiveScanTeacher:
 
     Call signature::
 
-        vx, vy, mode = teacher(dog_xy, dog_heading, sheep_positions, pen_target)
+        vx, vy, omega, mode = teacher(dog_xy, dog_heading, sheep_positions,
+                                      pen_target, drive_mode="differential")
+
+    ``omega`` is the yaw-rate intent (mecanum only); 0.0 for differential
+    drive and during blind exploration phases.
     """
 
     def __init__(self, base_action_fn, initial_scan_steps: int = INITIAL_SCAN_STEPS):
@@ -62,7 +66,8 @@ class ActiveScanTeacher:
             return 0.0, 0.0
         return EXPLORE_SPEED * dx / d, EXPLORE_SPEED * dy / d
 
-    def __call__(self, dog_xy, dog_heading, sheep_positions, pen_target):
+    def __call__(self, dog_xy, dog_heading, sheep_positions, pen_target,
+                 drive_mode="differential"):
         self.step += 1
         n_visible = len(sheep_positions)
 
@@ -75,7 +80,7 @@ class ActiveScanTeacher:
         if self.step <= self.initial_scan:
             vx, vy = self._scan_action(dog_heading)
             self.last_action = (vx, vy)
-            return vx, vy, "scan_initial"
+            return vx, vy, 0.0, "scan_initial"
 
         # Phase 2: walk-to-centre after a sustained empty tracker.
         if self.empty_streak >= EMPTY_DEBOUNCE_STEPS:
@@ -87,16 +92,31 @@ class ActiveScanTeacher:
                 vx, vy = ex, ey
                 mode = "explore"
             self.last_action = (vx, vy)
-            return vx, vy, mode
+            return vx, vy, 0.0, mode
 
         # Phase 2b: brief tracker blink — hold the previous action.
         if n_visible == 0:
             vx, vy = self.last_action
-            return vx, vy, "hold"
+            return vx, vy, 0.0, "hold"
 
         # Phase 3: hand off to the underlying analytic teacher, then
         # apply the shared near-sheep speed modulation.
-        vx, vy, mode = self.base(dog_xy, sheep_positions, pen_target)
+        # Handle both old-style (dog_xy, sheep, pen) and new-style
+        # (dog_xy, heading, sheep, pen, drive_mode) teachers.
+        try:
+            result = self.base(dog_xy, dog_heading, sheep_positions,
+                               pen_target, drive_mode)
+        except TypeError:
+            try:
+                result = self.base(dog_xy, dog_heading, sheep_positions,
+                                   pen_target)
+            except TypeError:
+                result = self.base(dog_xy, sheep_positions, pen_target)
+        if len(result) == 4:
+            vx, vy, omega, mode = result
+        else:
+            vx, vy, mode = result
+            omega = 0.0
         vx, vy = modulate_speed_near_sheep(vx, vy, dog_xy, sheep_positions)
         self.last_action = (vx, vy)
-        return vx, vy, mode
+        return vx, vy, omega, mode