"""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 close enough to first order that a policy trained here transfers. """ 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]² 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): """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, )