Checkpoint 8

tests/test_control.py

@@ -14,6 +14,7 @@ from herding.control.sequential import compute_action as sequential_action
 from herding.control.strombom import (
     DELTA_DRIVE, F_FACTOR, compute_action as strombom_action,
 )
+from herding.control.universal import compute_action as universal_action
 from herding.world.geometry import PEN_ENTRY
 
 
@@ -119,8 +120,10 @@ def test_sequential_targets_closest_to_pen():
 def test_active_scan_initial_phase_rotates():
     teacher = ActiveScanTeacher(strombom_action)
     # First call → opening rotation regardless of input.
-    vx, vy, mode = teacher((0.0, 0.0), 0.0, {"s0": (5.0, 0.0)}, PEN_ENTRY)
+    vx, vy, omega, mode = teacher(
+        (0.0, 0.0), 0.0, {"s0": (5.0, 0.0)}, PEN_ENTRY)
     assert mode == "scan_initial"
+    assert omega == 0.0
     assert math.isclose(math.hypot(vx, vy), 1.0, abs_tol=1e-6)
 
 
@@ -129,7 +132,8 @@ def test_active_scan_hands_off_to_base_after_opener():
     # Burn through the opener.
     for _ in range(2):
         teacher((0.0, 0.0), 0.0, {"s0": (0.0, 8.0)}, PEN_ENTRY)
-    _vx, _vy, mode = teacher((0.0, 0.0), 0.0, {"s0": (0.0, 8.0)}, PEN_ENTRY)
+    _vx, _vy, _omega, mode = teacher(
+        (0.0, 0.0), 0.0, {"s0": (0.0, 8.0)}, PEN_ENTRY)
     # Either drive (Strömbom mode label) or collect; not scan_initial.
     assert "scan" not in mode
 
@@ -141,7 +145,7 @@ def test_active_scan_holds_last_action_on_brief_empty():
     teacher((0.0, 0.0), 0.0, {"s0": (0.0, 8.0)}, PEN_ENTRY)
     last = teacher.last_action
     # Now a single empty frame → hold.
-    vx, vy, mode = teacher((0.0, 0.0), 0.0, {}, PEN_ENTRY)
+    vx, vy, _omega, mode = teacher((0.0, 0.0), 0.0, {}, PEN_ENTRY)
     assert mode == "hold"
     assert (vx, vy) == last
 
@@ -150,10 +154,30 @@ def test_active_scan_explores_after_sustained_empty():
     teacher = ActiveScanTeacher(strombom_action, initial_scan_steps=1)
     teacher((0.0, 0.0), 0.0, {}, PEN_ENTRY)  # opener
     for _ in range(EMPTY_DEBOUNCE_STEPS):
-        last_vx, last_vy, mode = teacher((5.0, 5.0), 0.0, {}, PEN_ENTRY)
+        last_vx, last_vy, _omega, mode = teacher(
+            (5.0, 5.0), 0.0, {}, PEN_ENTRY)
     assert mode in ("explore", "scan_at_centre")
 
 
+def test_active_scan_preserves_mecanum_omega():
+    """Regression: ActiveScanTeacher must propagate omega from a mecanum
+    base teacher, not silently drop it. Without this, BC mecanum demos
+    have omega=0 everywhere and the policy never learns to rotate.
+    """
+    teacher = ActiveScanTeacher(universal_action, initial_scan_steps=1)
+    # Burn the opener so we exit phase 1.
+    teacher((0.0, 0.0), 0.0, {"s0": (8.0, 8.0)}, PEN_ENTRY,
+            drive_mode="mecanum")
+    # Place a sheep off to the side so the dog needs to face it.
+    # Dog at origin facing +x (heading=0); target at (0, 8) → desired
+    # heading +π/2, so omega should be positive.
+    vx, vy, omega, mode = teacher(
+        (0.0, 0.0), 0.0, {"s0": (0.0, 8.0)}, PEN_ENTRY,
+        drive_mode="mecanum")
+    assert mode in ("collect", "drive", "recovery")
+    assert abs(omega) > 0.05, f"omega should be non-zero on mecanum, got {omega}"
+
+
 def test_active_scan_reset_clears_state():
     teacher = ActiveScanTeacher(strombom_action, initial_scan_steps=5)
     for _ in range(3):