TIR_PROJ/herding/world/diffdrive.py

"""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 by default.  Pass ``slip_std`` and an ``rng`` to
:func:`kinematics_step` / :func:`mecanum_kinematics_step` to add
per-wheel Gaussian speed noise for domain randomisation.
"""

from __future__ import annotations

import math
from typing import Optional

import numpy as np


def kinematics_step(x, y, h, w_left, w_right, wheel_radius, wheel_base, dt,
                    slip_std: float = 0.0,
                    rng: Optional[np.random.Generator] = None):
    """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)
    slip_std : optional Gaussian std (rad/s) added to each wheel speed
    rng  : numpy Generator for slip noise; required when slip_std > 0

    Returns (new_x, new_y, new_h).
    """
    if slip_std > 0.0 and rng is not None:
        noise = rng.normal(0.0, slip_std, size=2)
        w_left = w_left + noise[0]
        w_right = w_right + noise[1]
    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]² to wheel speeds.

    Forward speed scales by ``cos(err)`` (clamped to ±90°); 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 in wheel rad/s, target as a heading angle."""
    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


# ---------------------------------------------------------------------------
# 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,
                             slip_std: float = 0.0,
                             rng: Optional[np.random.Generator] = None,
                             strafe_efficiency: float = 1.0,
                             strafe_to_forward_bleed: float = 0.0):
    """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)
    slip_std : optional Gaussian std (rad/s) added to each wheel speed
    rng  : numpy Generator for slip noise; required when slip_std > 0
    strafe_efficiency : scales the realised lateral (vy_body) velocity.
        ``1.0`` (default) = perfect mecanum (textbook X-pattern). Set to
        the value that matches deployed-platform calibration to train
        a policy that compensates for under-actuated strafing — Webots
        with the roller-hinge mecanum proto currently calibrates to
        ~0.4 of textbook on strafe.
    strafe_to_forward_bleed : fraction of |vy_body_ideal| added to
        vx_body to simulate the consistent body-x bleed-through that
        accompanies strafing in Webots' physical-roller mecanum. Use a
        *negative* value (Webots calibrates to ≈ -0.28) to model the
        backward bleed seen on strafe; positive would model forward
        bleed. The bleed magnitude is symmetric in strafe sign — both
        +y and -y commands produce the same x-direction error.

    Returns (new_x, new_y, new_h).
    """
    if slip_std > 0.0 and rng is not None:
        noise = rng.normal(0.0, slip_std, size=4)
        w_fl, w_fr = w_fl + noise[0], w_fr + noise[1]
        w_rl, w_rr = w_rl + noise[2], w_rr + noise[3]
    r = wheel_radius
    vx_body = (w_fl + w_fr + w_rl + w_rr) * r / 4.0
    vy_body_ideal = (-w_fl + w_fr + w_rl - w_rr) * r / 4.0
    vy_body = vy_body_ideal * strafe_efficiency
    if strafe_to_forward_bleed != 0.0:
        # Bleed-through is asymmetric — forward in body frame, matching
        # Webots behaviour where strafe commands push the dog forward
        # regardless of strafe sign (the rollers slip the same way
        # symmetrically across the body's longitudinal axis).
        vx_body += strafe_to_forward_bleed * abs(vy_body_ideal)
    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,
    )