Replace iterative solver in SwerveSetpointGenerator with direct solutions

pathplannerlib-python/pathplannerlib/util/swerve.py

@@ -25,8 +25,6 @@ class SwerveSetpointGenerator:
     forces acting on a module's wheel from exceeding the force of friction.
     """
     _k_epsilon = 1e-8
-    _MAX_STEER_ITERATIONS = 8
-    _MAX_DRIVE_ITERATIONS = 10
 
     def __init__(self, config: RobotConfig, max_steer_velocity_rads_per_sec: float) -> None:
         """
@@ -288,15 +286,7 @@ def generateSetpoint(self, prev_setpoint: SwerveSetpoint, desired_state_robot_re
                 desired_vy[m] if min_s == 1.0 else (desired_vy[m] - prev_vy[m]) * min_s + prev_vy[m]
             # Find the max s for this drive wheel. Search on the interval between 0 and min_s, because we
             # already know we can't go faster than that.
-            s = self._findDriveMaxS(
-                prev_vx[m],
-                prev_vy[m],
-                math.hypot(prev_vx[m], prev_vy[m]),
-                vx_min_s,
-                vy_min_s,
-                math.hypot(vx_min_s, vy_min_s),
-                max_vel_step
-            )
+            s = self._findDriveMaxS(prev_vx[m], prev_vy[m], vx_min_s, vy_min_s, max_vel_step)
             min_s = min(min_s, s)
 
         ret_speeds = ChassisSpeeds(
@@ -383,97 +373,108 @@ def _unwrapAngle(ref: float, angle: float) -> float:
         else:
             return angle
 
-    @classmethod
-    def _findRoot(
-            cls,
-            func: Callable[[float, float], float],
-            x_0: float,
-            y_0: float,
-            f_0: float,
-            x_1: float,
-            y_1: float,
-            f_1: float,
-            iterations_left: int
-    ) -> float:
-        """
-        Find the root of the generic 2D parametric function 'func' using the regula falsi technique.
-        This is a pretty naive way to do root finding, but it's usually faster than simple bisection
-        while being robust in ways that e.g. the Newton-Raphson method isn't.
-        :param func: The Function2d to take the root of.
-        :param x_0: x value of the lower bracket.
-        :param y_0: y value of the lower bracket.
-        :param f_0: value of 'func' at x_0, y_0 (passed in by caller to save a call to 'func' during
-        recursion)
-        :param x_1: x value of the upper bracket.
-        :param y_1: y value of the upper bracket.
-        :param f_1: value of 'func' at x_1, y_1 (passed in by caller to save a call to 'func' during
-        recursion)
-        :param iterations_left: Number of iterations of root finding left.
-        :return: The parameter value 's' that interpolating between 0 and 1 that corresponds to the
-        (approximate) root.
-        """
-        s_guess = max(0.0, min(1.0, -f_0 / (f_1 - f_0)))
-
-        if iterations_left < 0 or cls._epsilonEquals(f_0, f_1):
-            return s_guess
-
-        x_guess = (x_1 - x_0) * s_guess + x_0
-        y_guess = (y_1 - y_0) * s_guess + y_0
-        f_guess = func(x_guess, y_guess)
-        if cls._signum(f_0) == cls._signum(f_guess):
-            # 0 and guess on same side of root, so use upper bracket.
-            return s_guess \
-                + (1.0 - s_guess) \
-                * cls._findRoot(func, x_guess, y_guess, f_guess, x_1, y_1, f_1, iterations_left - 1)
-        else:
-            # Use lower bracket.
-            return s_guess \
-                * cls._findRoot(func, x_0, y_0, f_0, x_guess, y_guess, f_guess, iterations_left - 1)
-
     @classmethod
     def _findSteeringMaxS(
             cls,
             x_0: float,
             y_0: float,
-            f_0: float,
+            theta_0: float,
             x_1: float,
             y_1: float,
-            f_1: float,
+            theta_1: float,
             max_deviation: float
     ) -> float:
-        f_1 = cls._unwrapAngle(f_0, f_1)
-        diff = f_1 - f_0
+        theta_1 = cls._unwrapAngle(theta_0, theta_1)
+        diff = theta_1 - theta_0
         if math.fabs(diff) <= max_deviation:
             # Can go all the way to s=1
             return 1.0
-        offset = f_0 + cls._signum(diff) * max_deviation
 
-        def func(x: float, y: float):
-            return cls._unwrapAngle(f_0, math.atan2(y, x)) - offset
-
-        return cls._findRoot(func, x_0, y_0, f_0 - offset, x_1, y_1, f_1 - offset, cls._MAX_STEER_ITERATIONS)
+        target = theta_0 + math.copysign(max_deviation, diff)
+
+        # Rotate the velocity vectors such that the target angle becomes the +X
+        # axis. We only need find the Y components, h_0 and h_1, since they are
+        # proportional to the distances from the two points to the solution
+        # point (x_0 + (x_1 - x_0)s, y_0 + (y_1 - y_0)s).
+        sin = math.sin(-target)
+        cos = math.cos(-target)
+        h_0 = sin * x_0 + cos * y_0
+        h_1 = sin * x_1 + cos * y_1
+
+        # Undo linear interpolation from h_0 to h_1:
+        # 0 = h_0 + (h_1 - h_0) * s
+        # -h_0 = (h_1 - h_0) * s
+        # -h_0 / (h_1 - h_0) = s
+        # h_0 / (h_0 - h_1) = s
+        # Guaranteed to not divide by zero, since if h_0 was equal to h_1, theta_0
+        # would be equal to theta_1, which is caught by the difference check.
+        return h_0 / (h_0 - h_1)
 
     @classmethod
     def _findDriveMaxS(
             cls,
             x_0: float,
             y_0: float,
-            f_0: float,
             x_1: float,
             y_1: float,
-            f_1: float,
             max_vel_step: float
     ) -> float:
-        diff = f_1 - f_0
+        # Derivation:
+        # Want to find point P(s) between (x_0, y_0) and (x_1, y_1) where the
+        # length of P(s) is the target T. P(s) is linearly interpolated between the
+        # points, so P(s) = (x_0 + (x_1 - x_0) * s, y_0 + (y_1 - y_0) * s).
+        # Then,
+        #     T = sqrt(P(s).x^2 + P(s).y^2)
+        #   T^2 = (x_0 + (x_1 - x_0) * s)^2 + (y_0 + (y_1 - y_0) * s)^2
+        #   T^2 = x_0^2 + 2x_0(x_1-x_0)s + (x_1-x_0)^2*s^2
+        #       + y_0^2 + 2y_0(y_1-y_0)s + (y_1-y_0)^2*s^2
+        #   T^2 = x_0^2 + 2x_0x_1s - 2x_0^2*s + x_1^2*s^2 - 2x_0x_1s^2 + x_0^2*s^2
+        #       + y_0^2 + 2y_0y_1s - 2y_0^2*s + y_1^2*s^2 - 2y_0y_1s^2 + y_0^2*s^2
+        #     0 = (x_0^2 + y_0^2 + x_1^2 + y_1^2 - 2x_0x_1 - 2y_0y_1)s^2
+        #       + (2x_0x_1 + 2y_0y_1 - 2x_0^2 - 2y_0^2)s
+        #       + (x_0^2 + y_0^2 - T^2).
+        #
+        # To simplify, we can factor out some common parts:
+        # Let l_0 = x_0^2 + y_0^2, l_1 = x_1^2 + y_1^2, and
+        # p = x_0 * x_1 + y_0 * y_1.
+        # Then we have
+        #   0 = (l_0 + l_1 - 2p)s^2 + 2(p - l_0)s + (l_0 - T^2),
+        # with which we can solve for s using the quadratic formula.
+
+        l_0 = x_0 * x_0 + y_0 * y_0
+        l_1 = x_1 * x_1 + y_1 * y_1
+        sqrt_l_0 = math.sqrt(l_0)
+        diff = math.sqrt(l_1) - sqrt_l_0
         if math.fabs(diff) <= max_vel_step:
             # Can go all the way to s=1.
             return 1.0
-        offset = f_0 + cls._signum(diff) * max_vel_step
-
-        def func(x: float, y: float):
-            return math.hypot(x, y) - offset
 
-        return cls._findRoot(func, x_0, y_0, f_0 - offset, x_1, y_1, f_1 - offset, cls._MAX_DRIVE_ITERATIONS)
+        target = sqrt_l_0 + math.copysign(max_vel_step, diff)
+        p = x_0 * x_1 + y_0 * y_1
+
+        # Quadratic of s
+        a = l_0 + l_1 - 2 * p
+        b = 2 * (p - l_0)
+        c = l_0 - target * target
+        root = math.sqrt(b * b - 4 * a * c)
+
+        def is_valid_s(s):
+            return math.isfinite(s) and s >= 0 and s <= 1
+
+        # Check if either of the solutions are valid
+        # Won't divide by zero because it is only possible for a to be zero if the
+        # target velocity is exactly the same or the reverse of the current
+        # velocity, which would be caught by the difference check.
+        s_1 = (-b + root) / (2 * a)
+        if is_valid_s(s_1):
+            return s_1
+        s_2 = (-b - root) / (2 * a)
+        if is_valid_s(s_2):
+            return s_2
+
+        # Since we passed the initial max_vel_step check, a solution should exist,
+        # but if no solution was found anyway, just don't limit movement
+        return 1.0
 
     @classmethod
     def _epsilonEquals(

pathplannerlib/src/main/java/com/pathplanner/lib/util/swerve/SwerveSetpointGenerator.java

@@ -26,8 +26,6 @@
  */
 public class SwerveSetpointGenerator {
   private static final double kEpsilon = 1E-8;
-  private static final int MAX_STEER_ITERATIONS = 8;
-  private static final int MAX_DRIVE_ITERATIONS = 10;
 
   private final RobotConfig config;
   private final double maxSteerVelocityRadsPerSec;
@@ -326,15 +324,7 @@ public SwerveSetpoint generateSetpoint(
           min_s == 1.0 ? desired_vy[m] : (desired_vy[m] - prev_vy[m]) * min_s + prev_vy[m];
       // Find the max s for this drive wheel. Search on the interval between 0 and min_s, because we
       // already know we can't go faster than that.
-      double s =
-          findDriveMaxS(
-              prev_vx[m],
-              prev_vy[m],
-              Math.hypot(prev_vx[m], prev_vy[m]),
-              vx_min_s,
-              vy_min_s,
-              Math.hypot(vx_min_s, vy_min_s),
-              maxVelStep);
+      double s = findDriveMaxS(prev_vx[m], prev_vy[m], vx_min_s, vy_min_s, maxVelStep);
       min_s = Math.min(min_s, s);
     }
 
@@ -501,88 +491,104 @@ private static double unwrapAngle(double ref, double angle) {
     }
   }
 
-  @FunctionalInterface
-  private interface Function2d {
-    double f(double x, double y);
-  }
-
-  /**
-   * Find the root of the generic 2D parametric function 'func' using the regula falsi technique.
-   * This is a pretty naive way to do root finding, but it's usually faster than simple bisection
-   * while being robust in ways that e.g. the Newton-Raphson method isn't.
-   *
-   * @param func The Function2d to take the root of.
-   * @param x_0 x value of the lower bracket.
-   * @param y_0 y value of the lower bracket.
-   * @param f_0 value of 'func' at x_0, y_0 (passed in by caller to save a call to 'func' during
-   *     recursion)
-   * @param x_1 x value of the upper bracket.
-   * @param y_1 y value of the upper bracket.
-   * @param f_1 value of 'func' at x_1, y_1 (passed in by caller to save a call to 'func' during
-   *     recursion)
-   * @param iterations_left Number of iterations of root finding left.
-   * @return The parameter value 's' that interpolating between 0 and 1 that corresponds to the
-   *     (approximate) root.
-   */
-  private static double findRoot(
-      Function2d func,
-      double x_0,
-      double y_0,
-      double f_0,
-      double x_1,
-      double y_1,
-      double f_1,
-      int iterations_left) {
-    var s_guess = Math.max(0.0, Math.min(1.0, -f_0 / (f_1 - f_0)));
-
-    if (iterations_left < 0 || epsilonEquals(f_0, f_1)) {
-      return s_guess;
-    }
-
-    var x_guess = (x_1 - x_0) * s_guess + x_0;
-    var y_guess = (y_1 - y_0) * s_guess + y_0;
-    var f_guess = func.f(x_guess, y_guess);
-    if (Math.signum(f_0) == Math.signum(f_guess)) {
-      // 0 and guess on same side of root, so use upper bracket.
-      return s_guess
-          + (1.0 - s_guess)
-              * findRoot(func, x_guess, y_guess, f_guess, x_1, y_1, f_1, iterations_left - 1);
-    } else {
-      // Use lower bracket.
-      return s_guess
-          * findRoot(func, x_0, y_0, f_0, x_guess, y_guess, f_guess, iterations_left - 1);
-    }
-  }
-
   private static double findSteeringMaxS(
       double x_0,
       double y_0,
-      double f_0,
+      double theta_0,
       double x_1,
       double y_1,
-      double f_1,
+      double theta_1,
       double max_deviation) {
-    f_1 = unwrapAngle(f_0, f_1);
-    double diff = f_1 - f_0;
+    theta_1 = unwrapAngle(theta_0, theta_1);
+    double diff = theta_1 - theta_0;
     if (Math.abs(diff) <= max_deviation) {
       // Can go all the way to s=1.
       return 1.0;
     }
-    double offset = f_0 + Math.signum(diff) * max_deviation;
-    Function2d func = (x, y) -> unwrapAngle(f_0, Math.atan2(y, x)) - offset;
-    return findRoot(func, x_0, y_0, f_0 - offset, x_1, y_1, f_1 - offset, MAX_STEER_ITERATIONS);
+
+    double target = theta_0 + Math.copySign(max_deviation, diff);
+
+    // Rotate the velocity vectors such that the target angle becomes the +X
+    // axis. We only need find the Y components, h_0 and h_1, since they are
+    // proportional to the distances from the two points to the solution
+    // point (x_0 + (x_1 - x_0)s, y_0 + (y_1 - y_0)s).
+    double sin = Math.sin(-target);
+    double cos = Math.cos(-target);
+    double h_0 = sin * x_0 + cos * y_0;
+    double h_1 = sin * x_1 + cos * y_1;
+
+    // Undo linear interpolation from h_0 to h_1:
+    // 0 = h_0 + (h_1 - h_0) * s
+    // -h_0 = (h_1 - h_0) * s
+    // -h_0 / (h_1 - h_0) = s
+    // h_0 / (h_0 - h_1) = s
+    // Guaranteed to not divide by zero, since if h_0 was equal to h_1, theta_0
+    // would be equal to theta_1, which is caught by the difference check.
+    return h_0 / (h_0 - h_1);
+  }
+
+  private static boolean isValidS(double s) {
+    return Double.isFinite(s) && s >= 0 && s <= 1;
   }
 
   private static double findDriveMaxS(
-      double x_0, double y_0, double f_0, double x_1, double y_1, double f_1, double max_vel_step) {
-    double diff = f_1 - f_0;
+      double x_0, double y_0, double x_1, double y_1, double max_vel_step) {
+    // Derivation:
+    // Want to find point P(s) between (x_0, y_0) and (x_1, y_1) where the
+    // length of P(s) is the target T. P(s) is linearly interpolated between the
+    // points, so P(s) = (x_0 + (x_1 - x_0) * s, y_0 + (y_1 - y_0) * s).
+    // Then,
+    //     T = sqrt(P(s).x^2 + P(s).y^2)
+    //   T^2 = (x_0 + (x_1 - x_0) * s)^2 + (y_0 + (y_1 - y_0) * s)^2
+    //   T^2 = x_0^2 + 2x_0(x_1-x_0)s + (x_1-x_0)^2*s^2
+    //       + y_0^2 + 2y_0(y_1-y_0)s + (y_1-y_0)^2*s^2
+    //   T^2 = x_0^2 + 2x_0x_1s - 2x_0^2*s + x_1^2*s^2 - 2x_0x_1s^2 + x_0^2*s^2
+    //       + y_0^2 + 2y_0y_1s - 2y_0^2*s + y_1^2*s^2 - 2y_0y_1s^2 + y_0^2*s^2
+    //     0 = (x_0^2 + y_0^2 + x_1^2 + y_1^2 - 2x_0x_1 - 2y_0y_1)s^2
+    //       + (2x_0x_1 + 2y_0y_1 - 2x_0^2 - 2y_0^2)s
+    //       + (x_0^2 + y_0^2 - T^2).
+    //
+    // To simplify, we can factor out some common parts:
+    // Let l_0 = x_0^2 + y_0^2, l_1 = x_1^2 + y_1^2, and
+    // p = x_0 * x_1 + y_0 * y_1.
+    // Then we have
+    //   0 = (l_0 + l_1 - 2p)s^2 + 2(p - l_0)s + (l_0 - T^2),
+    // with which we can solve for s using the quadratic formula.
+
+    double l_0 = x_0 * x_0 + y_0 * y_0;
+    double l_1 = x_1 * x_1 + y_1 * y_1;
+    double sqrt_l_0 = Math.sqrt(l_0);
+    double diff = Math.sqrt(l_1) - sqrt_l_0;
     if (Math.abs(diff) <= max_vel_step) {
       // Can go all the way to s=1.
       return 1.0;
     }
-    double offset = f_0 + Math.signum(diff) * max_vel_step;
-    Function2d func = (x, y) -> Math.hypot(x, y) - offset;
-    return findRoot(func, x_0, y_0, f_0 - offset, x_1, y_1, f_1 - offset, MAX_DRIVE_ITERATIONS);
+
+    double target = sqrt_l_0 + Math.copySign(max_vel_step, diff);
+    double p = x_0 * x_1 + y_0 * y_1;
+
+    // Quadratic of s
+    double a = l_0 + l_1 - 2 * p;
+    double b = 2 * (p - l_0);
+    double c = l_0 - target * target;
+    double root = Math.sqrt(b * b - 4 * a * c);
+
+    // Check if either of the solutions are valid
+    // Won't divide by zero because it is only possible for a to be zero if the
+    // target velocity is exactly the same or the reverse of the current
+    // velocity, which would be caught by the difference check.
+    double s_1 = (-b + root) / (2 * a);
+    if (isValidS(s_1)) {
+      return s_1;
+    }
+    double s_2 = (-b - root) / (2 * a);
+    if (isValidS(s_2)) {
+      return s_2;
+    }
+
+    // Since we passed the initial max_vel_step check, a solution should exist,
+    // but if no solution was found anyway, just don't limit movement
+    return 1.0;
   }
 
   private static boolean epsilonEquals(double a, double b, double epsilon) {