Checkpoint 6

herding/world/diffdrive.py

@@ -0,0 +1,70 @@
+"""Differential-drive kinematics matching the Webots robot specs.
+
+The Webots controllers and the training env both use these helpers so the
+sim and the real (Webots) physics agree to first order. They do not model
+slip, wheel acceleration limits, or contact forces — Webots does that for
+us at inference time. The training env has to be close enough that a
+policy trained against this kinematic model still works when handed off
+to ODE physics.
+"""
+
+import math
+
+
+def kinematics_step(x, y, h, w_left, w_right, wheel_radius, wheel_base, dt):
+    """Integrate one step of differential-drive forward kinematics.
+
+    Inputs
+    ------
+    x, y : robot position (m)
+    h    : robot heading (rad), 0 = +x axis
+    w_left, w_right : wheel angular velocities (rad/s)
+    wheel_radius, wheel_base : robot dimensions (m)
+    dt   : timestep (s)
+
+    Returns (new_x, new_y, new_h).
+    """
+    v = (w_right + w_left) * wheel_radius * 0.5
+    omega = (w_right - w_left) * wheel_radius / wheel_base
+    new_x = x + v * math.cos(h) * dt
+    new_y = y + v * math.sin(h) * dt
+    new_h = math.atan2(math.sin(h + omega * dt), math.cos(h + omega * dt))
+    return new_x, new_y, new_h
+
+
+def velocity_to_wheels(vx, vy, h, max_linear, wheel_radius, max_wheel_omega,
+                       k_turn=4.0):
+    """Convert a desired (vx, vy) intent in [-1, 1]^2 to wheel speeds.
+
+    Mirrors ``drive_action`` in controllers/shepherd_dog/shepherd_dog.py:
+    forward speed scales by ``cos(err)`` (clamped to ±90°), and a P
+    controller on heading error contributes the wheel-rate differential.
+    """
+    speed_ms = math.hypot(vx, vy) * max_linear
+    if speed_ms < 1e-3:
+        return 0.0, 0.0
+    target_h = math.atan2(vy, vx)
+    err = math.atan2(math.sin(target_h - h), math.cos(target_h - h))
+    clamped_err = max(-math.pi / 2, min(math.pi / 2, err))
+    fwd_ms = speed_ms * math.cos(clamped_err)
+    fwd_rad = fwd_ms / wheel_radius
+    turn = k_turn * err
+    left = max(-max_wheel_omega, min(max_wheel_omega, fwd_rad - turn))
+    right = max(-max_wheel_omega, min(max_wheel_omega, fwd_rad + turn))
+    return left, right
+
+
+def heading_speed_to_wheels(heading, speed_motor, h, max_wheel_omega,
+                            k_turn=4.0):
+    """Sheep variant: speed already expressed in motor (wheel rad/s) units.
+
+    Matches the existing sheep controller (``controllers/sheep/sheep.py``)
+    where ``speed = max(WANDER_SPEED, min(FLEE_SPEED, mag * 3.0))`` and
+    these constants are wheel angular velocities, not linear m/s.
+    """
+    err = math.atan2(math.sin(heading - h), math.cos(heading - h))
+    fwd = max(0.0, math.cos(err)) * speed_motor
+    turn = k_turn * err
+    left = max(-max_wheel_omega, min(max_wheel_omega, fwd - turn))
+    right = max(-max_wheel_omega, min(max_wheel_omega, fwd + turn))
+    return left, right