Checkpoint 8

herding/world/diffdrive.py

@@ -1,4 +1,5 @@
-"""Differential-drive kinematics, shared by the env and Webots controllers.
+"""Differential-drive and mecanum kinematics, shared by the env and Webots
+controllers.
 
 First-order rigid-body model — no slip, wheel-accel limits, or contact
 forces. Webots' ODE physics handles those at inference; the env stays
@@ -59,3 +60,131 @@ def heading_speed_to_wheels(heading, speed_motor, h, max_wheel_omega,
     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
+
+
+# ---------------------------------------------------------------------------
+# Mecanum (4-wheel omnidirectional) kinematics
+# ---------------------------------------------------------------------------
+
+def mecanum_kinematics_step(x, y, h, w_fl, w_fr, w_rl, w_rr,
+                             wheel_radius, lx, ly, dt):
+    """Integrate one step of mecanum forward kinematics.
+
+    Parameters
+    ----------
+    x, y : robot position (m)
+    h    : robot heading (rad), 0 = +x axis
+    w_fl, w_fr, w_rl, w_rr : wheel angular velocities (rad/s)
+    wheel_radius : wheel radius (m)
+    lx   : half the front-to-back axle distance (m)
+    ly   : half the left-to-right axle distance (m)
+    dt   : timestep (s)
+
+    Returns (new_x, new_y, new_h).
+    """
+    r = wheel_radius
+    vx_body = (w_fl + w_fr + w_rl + w_rr) * r / 4.0
+    vy_body = (-w_fl + w_fr + w_rl - w_rr) * r / 4.0
+    omega = (-w_fl + w_fr - w_rl + w_rr) * r / (4.0 * (lx + ly))
+
+    cos_h = math.cos(h)
+    sin_h = math.sin(h)
+    vx_world = vx_body * cos_h - vy_body * sin_h
+    vy_world = vx_body * sin_h + vy_body * cos_h
+
+    new_x = x + vx_world * dt
+    new_y = y + vy_world * dt
+    new_h = math.atan2(math.sin(h + omega * dt), math.cos(h + omega * dt))
+    return new_x, new_y, new_h
+
+
+def mecanum_inverse(vx_body, vy_body, omega, wheel_radius, lx, ly,
+                    max_wheel_omega):
+    """Mecanum inverse kinematics: body-frame velocities to 4 wheel speeds.
+
+    Parameters
+    ----------
+    vx_body, vy_body : desired body-frame linear velocities (m/s)
+    omega            : desired yaw rate (rad/s)
+    wheel_radius     : wheel radius (m)
+    lx               : half front-to-back axle distance (m)
+    ly               : half left-to-right axle distance (m)
+    max_wheel_omega  : wheel angular velocity clamp (rad/s)
+
+    Returns (w_fl, w_fr, w_rl, w_rr).
+    """
+    r = wheel_radius
+    k = lx + ly
+    w_fl = (vx_body - vy_body - k * omega) / r
+    w_fr = (vx_body + vy_body + k * omega) / r
+    w_rl = (vx_body + vy_body - k * omega) / r
+    w_rr = (vx_body - vy_body + k * omega) / r
+
+    scale = max(abs(w_fl), abs(w_fr), abs(w_rl), abs(w_rr), 1e-9)
+    if scale > max_wheel_omega:
+        ratio = max_wheel_omega / scale
+        w_fl *= ratio
+        w_fr *= ratio
+        w_rl *= ratio
+        w_rr *= ratio
+
+    return w_fl, w_fr, w_rl, w_rr
+
+
+def velocity_to_mecanum_wheels(vx, vy, omega, h, max_linear, wheel_radius,
+                               lx, ly, max_wheel_omega,
+                               k_turn=4.0, wheel_base=0.28):
+    """Convert world-frame (vx, vy, omega) action in [-1, 1]^3 to 4 wheel speeds.
+
+    Truly holonomic interpretation: (vx, vy) is the desired *world-frame*
+    velocity (magnitude up to ``max_linear`` m/s) and ``omega`` is the
+    desired yaw rate (independent of motion direction). The dog can
+    crab-walk and rotate at the same time.
+
+    This matches the universal teacher's signal: drive toward a standoff
+    point while facing the sheep / pen separately. With the older
+    non-holonomic version, ``omega`` from the teacher fought against the
+    forward-only kinematics and dropped success rates instead of helping.
+
+    Parameters
+    ----------
+    vx, vy : desired world-frame velocity intent in [-1, 1] (clamped on
+             magnitude to ≤ 1)
+    omega  : desired yaw rate intent in [-1, 1]
+    h      : current heading (rad), 0 = +x
+    max_linear     : max linear speed (m/s)
+    wheel_radius   : wheel radius (m)
+    lx, ly         : half axle distances (m)
+    max_wheel_omega : wheel angular velocity clamp (rad/s)
+    k_turn         : unused (kept for signature compatibility)
+    wheel_base     : unused (kept for signature compatibility)
+
+    Returns (w_fl, w_fr, w_rl, w_rr).
+    """
+    # Clamp the action magnitude in the (vx, vy) unit disk.
+    norm = math.hypot(vx, vy)
+    if norm > 1.0:
+        vx /= norm
+        vy /= norm
+
+    # World-frame velocity → body-frame velocity (rotate by -h).
+    vx_world = vx * max_linear
+    vy_world = vy * max_linear
+    cos_h = math.cos(h)
+    sin_h = math.sin(h)
+    vx_body =  cos_h * vx_world + sin_h * vy_world
+    vy_body = -sin_h * vx_world + cos_h * vy_world
+
+    # Yaw rate: omega ∈ [-1, 1] maps to ± max_linear / (lx + ly) — same
+    # peak yaw as the old "omega_extra" channel, but used directly
+    # rather than added to a heading-tracker.
+    yaw_max = max_linear / max(lx + ly, 1e-6)
+    omega_rad = omega * yaw_max
+
+    if abs(vx_body) < 1e-3 and abs(vy_body) < 1e-3 and abs(omega_rad) < 1e-3:
+        return 0.0, 0.0, 0.0, 0.0
+
+    return mecanum_inverse(
+        vx_body, vy_body, omega_rad,
+        wheel_radius, lx, ly, max_wheel_omega,
+    )