Improve robustness of spline planning

Previously, there were corner cases in the spline math that we handled
poorly due to laziness in how we handled certain constraints.

This patch should make it so that we can correctly plan any spline that
we actually produce--although due to the fact that it is now more
properly following the actual physical constraints of the splines, it
actually forced me to increase slightly the tolerances in the tests.

This has yet to be tested on a real robot.

Change-Id: I63f68eede9d0fbe6d41bf3caea8aca19ece9fa1f
diff --git a/frc971/control_loops/drivetrain/trajectory.cc b/frc971/control_loops/drivetrain/trajectory.cc
index ba4c0c2..97662d4 100644
--- a/frc971/control_loops/drivetrain/trajectory.cc
+++ b/frc971/control_loops/drivetrain/trajectory.cc
@@ -18,22 +18,20 @@
                        const DrivetrainConfig<double> &config, double vmax,
                        int num_distance)
     : spline_(spline),
-      velocity_drivetrain_(
-          ::std::unique_ptr<StateFeedbackLoop<2, 2, 2, double,
-                                              StateFeedbackHybridPlant<2, 2, 2>,
-                                              HybridKalman<2, 2, 2>>>(
-              new StateFeedbackLoop<2, 2, 2, double,
-                                    StateFeedbackHybridPlant<2, 2, 2>,
-                                    HybridKalman<2, 2, 2>>(
-                  config.make_hybrid_drivetrain_velocity_loop()))),
+      velocity_drivetrain_(::std::unique_ptr<
+          StateFeedbackLoop<2, 2, 2, double, StateFeedbackHybridPlant<2, 2, 2>,
+                            HybridKalman<2, 2, 2>>>(
+          new StateFeedbackLoop<2, 2, 2, double,
+                                StateFeedbackHybridPlant<2, 2, 2>,
+                                HybridKalman<2, 2, 2>>(
+              config.make_hybrid_drivetrain_velocity_loop()))),
       robot_radius_l_(config.robot_radius),
       robot_radius_r_(config.robot_radius),
-      longitudal_acceleration_(3.0),
+      longitudinal_acceleration_(3.0),
       lateral_acceleration_(2.0),
       Tlr_to_la_((::Eigen::Matrix<double, 2, 2>() << 0.5, 0.5,
                   -1.0 / (robot_radius_l_ + robot_radius_r_),
-                  1.0 / (robot_radius_l_ + robot_radius_r_))
-                     .finished()),
+                  1.0 / (robot_radius_l_ + robot_radius_r_)).finished()),
       Tla_to_lr_(Tlr_to_la_.inverse()),
       plan_(num_distance == 0
                 ? ::std::max(100, static_cast<int>(spline_->length() / 0.0025))
@@ -44,56 +42,293 @@
 void Trajectory::LateralAccelPass() {
   for (size_t i = 0; i < plan_.size(); ++i) {
     const double distance = Distance(i);
-    plan_[i] = ::std::min(plan_[i], LateralVelocityCurvature(distance));
+    const double velocity_limit =  LateralVelocityCurvature(distance);
+    if (velocity_limit < plan_[i]) {
+      plan_[i] = velocity_limit;
+      plan_segment_type_[i] = CURVATURE_LIMITED;
+    }
   }
 }
 
-// TODO(austin): Deduplicate this potentially with the backward accel function.
-// Need to sort out how the max velocity limit is going to work since the
-// velocity and acceleration need to match at all points.
-// TODO(austin): Accel check the wheels instead of the center of mass.
-double Trajectory::ForwardAcceleration(const double x, const double v) {
+void Trajectory::VoltageFeasibilityPass(VoltageLimit limit_type) {
+  for (size_t i = 0; i < plan_.size(); ++i) {
+    const double distance = Distance(i);
+    const double velocity_limit = VoltageVelocityLimit(distance, limit_type);
+    if (velocity_limit < plan_[i]) {
+      plan_[i] = velocity_limit;
+      plan_segment_type_[i] = VOLTAGE_LIMITED;
+    }
+  }
+}
+
+double Trajectory::BestAcceleration(double x, double v, bool backwards) {
   ::Eigen::Matrix<double, 2, 1> K3;
   ::Eigen::Matrix<double, 2, 1> K4;
   ::Eigen::Matrix<double, 2, 1> K5;
   K345(x, &K3, &K4, &K5);
 
-  const ::Eigen::Matrix<double, 2, 1> C = K3 * v * v + K4 * v;
   // Now, solve for all a's and find the best one which meets our criteria.
-  double maxa = -::std::numeric_limits<double>::infinity();
+  const ::Eigen::Matrix<double, 2, 1> C = K3 * v * v + K4 * v;
+  double min_voltage_accel = ::std::numeric_limits<double>::infinity();
+  double max_voltage_accel = -min_voltage_accel;
   for (const double a : {(voltage_limit_ - C(0, 0)) / K5(0, 0),
                          (voltage_limit_ - C(1, 0)) / K5(1, 0),
                          (-voltage_limit_ - C(0, 0)) / K5(0, 0),
                          (-voltage_limit_ - C(1, 0)) / K5(1, 0)}) {
     const ::Eigen::Matrix<double, 2, 1> U = K5 * a + K3 * v * v + K4 * v;
     if ((U.array().abs() < voltage_limit_ + 1e-6).all()) {
-      maxa = ::std::max(maxa, a);
+      min_voltage_accel = ::std::min(a, min_voltage_accel);
+      max_voltage_accel = ::std::max(a, max_voltage_accel);
     }
   }
+  double best_accel = backwards ? min_voltage_accel : max_voltage_accel;
 
-  // Then, assume an acceleration oval and stay inside it.
-  const double lateral_acceleration = v * v * spline_->DDXY(x).norm();
-  const double ellipse_down_shift = longitudal_acceleration_ * 1.0;
-  const double ellipse_width_stretch = ::std::sqrt(
-      1.0 / (1.0 - ::std::pow(ellipse_down_shift / (longitudal_acceleration_ +
-                                                    ellipse_down_shift),
-                              2.0)));
-  const double squared =
-      1.0 - ::std::pow(lateral_acceleration / lateral_acceleration_ /
-                           ellipse_width_stretch,
-                       2.0);
-  // If we would end up with an imaginary number, cap us at 0 acceleration.
-  // TODO(austin): Investigate when this happens, why, and fix it.
-  if (squared < 0.0) {
-    AOS_LOG(ERROR, "Imaginary %f, maxa: %f, fa(%f, %f) -> 0.0\n", squared, maxa,
-            x, v);
-    return 0.0;
+  double min_friction_accel, max_friction_accel;
+  FrictionLngAccelLimits(x, v, &min_friction_accel, &max_friction_accel);
+  if (backwards) {
+    best_accel = ::std::max(best_accel, min_friction_accel);
+  } else {
+    best_accel = ::std::min(best_accel, max_friction_accel);
   }
-  const double longitudal_acceleration =
-      ::std::sqrt(::std::abs(squared)) *
-          (longitudal_acceleration_ + ellipse_down_shift) -
-      ellipse_down_shift;
-  return ::std::min(longitudal_acceleration, maxa);
+
+  if (max_friction_accel < min_friction_accel - 0.1 ||
+      best_accel < min_voltage_accel || best_accel > max_voltage_accel) {
+    AOS_LOG(WARNING,
+        "Viable friction limits and viable voltage limits do not overlap (x: "
+        "%f, v: %f, backwards: %d) best_accel = %f, min voltage %f, max "
+        "voltage %f min friction %f max friction %f.\n",
+        x, v, backwards, best_accel, min_voltage_accel, max_voltage_accel,
+        min_friction_accel, max_friction_accel);
+    // Don't actually do anything--this will just result in attempting to drive
+    // higher voltages thatn we have available. In practice, that'll probably
+    // work out fine.
+  }
+
+  return best_accel;
+}
+
+double Trajectory::LateralVelocityCurvature(double distance) const {
+  // To calculate these constraints, we first note that:
+  // wheel accels = K2 * v_robot' + K1 * v_robot^2
+  // All that this logic does is solve for v_robot, leaving v_robot' free,
+  // assuming that the wheels are at their limits.
+  // To do this, we:
+  //
+  // 1) Determine what the wheel accels will be at the limit--since we have
+  // two free variables (v_robot, v_robot'), both wheels will be at their
+  // limits--if in a sufficiently tight turn (such that the signs of the
+  // coefficients of K2 are different), then the wheels will be accelerating
+  // in opposite directions; otherwise, they accelerate in the same direction.
+  // The magnitude of these per-wheel accelerations is a function of velocity,
+  // so it must also be solved for.
+  //
+  // 2) Eliminate that v_robot' term (since we don't care
+  // about it) by multiplying be a "K2prime" term (where K2prime * K2 = 0) on
+  // both sides of the equation.
+  //
+  // 3) Solving the relatively tractable remaining equation, which is
+  // basically just grouping all the terms together in one spot and taking the
+  // 4th root of everything.
+  const double dtheta = spline_->DTheta(distance);
+  const ::Eigen::Matrix<double, 1, 2> K2prime =
+      K2(dtheta).transpose() *
+      (::Eigen::Matrix<double, 2, 2>() << 0, 1, -1, 0).finished();
+  // Calculate whether the wheels are spinning in opposite directions.
+  const bool opposites = K2prime(0) * K2prime(1) < 0;
+  const ::Eigen::Matrix<double, 2, 1> K1calc = K1(spline_->DDTheta(distance));
+  const double lat_accel_squared =
+      ::std::pow(dtheta / lateral_acceleration_, 2);
+  const double curvature_change_term =
+      (K2prime * K1calc).value() /
+      (K2prime *
+       (::Eigen::Matrix<double, 2, 1>() << 1.0, (opposites ? -1.0 : 1.0))
+           .finished() *
+       longitudinal_acceleration_)
+          .value();
+  const double vel_inv = ::std::sqrt(
+      ::std::sqrt(::std::pow(curvature_change_term, 2) + lat_accel_squared));
+  if (vel_inv == 0.0) {
+    return ::std::numeric_limits<double>::infinity();
+  }
+  return 1.0 / vel_inv;
+}
+
+void Trajectory::FrictionLngAccelLimits(double x, double v, double *min_accel,
+                                        double *max_accel) const {
+  // First, calculate the max longitudinal acceleration that can be achieved
+  // by either wheel given the friction elliipse that we have.
+  const double lateral_acceleration = v * v * spline_->DDXY(x).norm();
+  const double max_wheel_lng_accel_squared =
+      1.0 - ::std::pow(lateral_acceleration / lateral_acceleration_, 2.0);
+  if (max_wheel_lng_accel_squared < 0.0) {
+    AOS_LOG(DEBUG,
+        "Something (probably Runge-Kutta) queried invalid velocity %f at "
+        "distance %f\n",
+        v, x);
+    // If we encounter this, it means that the Runge-Kutta has attempted to
+    // sample points a bit past the edge of the friction boundary. If so, we
+    // gradually ramp the min/max accels to be more and more incorrect (note
+    // how min_accel > max_accel if we reach this case) to avoid causing any
+    // numerical issues.
+    *min_accel =
+        ::std::sqrt(-max_wheel_lng_accel_squared) * longitudinal_acceleration_;
+    *max_accel = -*min_accel;
+    return;
+  }
+  *min_accel = -::std::numeric_limits<double>::infinity();
+  *max_accel = ::std::numeric_limits<double>::infinity();
+
+  // Calculate max/min accelerations by calculating what the robots overall
+  // longitudinal acceleration would be if each wheel were running at the max
+  // forwards/backwards longitudinal acceleration.
+  const double max_wheel_lng_accel =
+      longitudinal_acceleration_ * ::std::sqrt(max_wheel_lng_accel_squared);
+  const ::Eigen::Matrix<double, 2, 1> K1v2 = K1(spline_->DDTheta(x)) * v * v;
+  const ::Eigen::Matrix<double, 2, 1> K2inv =
+      K2(spline_->DTheta(x)).cwiseInverse();
+  // Store the accelerations of the robot corresponding to each wheel being at
+  // the max/min acceleration. The first coefficient in each vector
+  // corresponds to the left wheel, the second to the right wheel.
+  const ::Eigen::Matrix<double, 2, 1> accels1 =
+      K2inv.array() * (-K1v2.array() + max_wheel_lng_accel);
+  const ::Eigen::Matrix<double, 2, 1> accels2 =
+      K2inv.array() * (-K1v2.array() - max_wheel_lng_accel);
+
+  // If either term is non-finite, that suggests that a term of K2 is zero
+  // (which is physically possible when turning such that one wheel is
+  // stationary), so just ignore that side of the drivetrain.
+  if (::std::isfinite(accels1(0))) {
+    // The inner max/min in this case determines which of the two cases (+ or
+    // - acceleration on the left wheel) we care about--in a sufficiently
+    // tight turning radius, the left hweel may be accelerating backwards when
+    // the robot as a whole accelerates forwards. We then use that
+    // acceleration to bound the min/max accel.
+    *min_accel = ::std::max(*min_accel, ::std::min(accels1(0), accels2(0)));
+    *max_accel = ::std::min(*max_accel, ::std::max(accels1(0), accels2(0)));
+  }
+  // Same logic as previous if-statement, but for the right wheel.
+  if (::std::isfinite(accels1(1))) {
+    *min_accel = ::std::max(*min_accel, ::std::min(accels1(1), accels2(1)));
+    *max_accel = ::std::min(*max_accel, ::std::max(accels1(1), accels2(1)));
+  }
+}
+
+double Trajectory::VoltageVelocityLimit(
+    double distance, VoltageLimit limit_type,
+    Eigen::Matrix<double, 2, 1> *constraint_voltages) const {
+  // To sketch an outline of the math going on here, we start with the basic
+  // dynamics of the robot along the spline:
+  // K2 * v_robot' + K1 * v_robot^2 = A * K2 * v_robot + B * U
+  // We need to determine the maximum v_robot given constrained U and free
+  // v_robot'.
+  // Similarly to the friction constraints, we accomplish this by first
+  // multiplying by a K2prime term to eliminate the v_robot' term.
+  // As with the friction constraints, we also know that the limits will occur
+  // when both sides of the drivetrain are driven at their max magnitude
+  // voltages, although they may be driven at different signs.
+  // Once we determine whether the voltages match signs, we still have to
+  // consider both possible pairings (technically we could probably
+  // predetermine which pairing, e.g. +/- or -/+, we acre about, but we don't
+  // need to).
+  //
+  // For each pairing, we then get to solve a quadratic formula for the robot
+  // velocity at those voltages. This gives us up to 4 solutions, of which
+  // up to 3 will give us positive velocities; each solution velocity
+  // corresponds to a transition from feasibility to infeasibility, where a
+  // velocity of zero is always feasible, and there will always be 0, 1, or 3
+  // positive solutions. Among the positive solutions, we take both the min
+  // and the max--the min will be the highest velocity such that all
+  // velocities between zero and that velocity are valid; the max will be
+  // the highest feasible velocity. Which we return depends on what the
+  // limit_type is.
+  //
+  // Sketching the actual math:
+  // K2 * v_robot' + K1 * v_robot^2 = A * K2 * v_robot +/- B * U_max
+  // K2prime * K1 * v_robot^2 = K2prime * (A * K2 * v_robot +/- B * U_max)
+  // a v_robot^2 + b v_robot +/- c = 0
+  const ::Eigen::Matrix<double, 2, 2> B =
+      velocity_drivetrain_->plant().coefficients().B_continuous;
+  const double dtheta = spline_->DTheta(distance);
+  const ::Eigen::Matrix<double, 2, 1> BinvK2 = B.inverse() * K2(dtheta);
+  // Because voltages can actually impact *both* wheels, in order to determine
+  // whether the voltages will have opposite signs, we need to use B^-1 * K2.
+  const bool opposite_voltages = BinvK2(0) * BinvK2(1) > 0.0;
+  const ::Eigen::Matrix<double, 1, 2> K2prime =
+      K2(dtheta).transpose() *
+      (::Eigen::Matrix<double, 2, 2>() << 0, 1, -1, 0).finished();
+  const double a = K2prime * K1(spline_->DDTheta(distance));
+  const double b = -K2prime *
+                   velocity_drivetrain_->plant().coefficients().A_continuous *
+                   K2(dtheta);
+  const ::Eigen::Matrix<double, 1, 2> c_coeff = -K2prime * B;
+  // Calculate the "positive" version of the voltage limits we will use.
+  const ::Eigen::Matrix<double, 2, 1> abs_volts =
+      voltage_limit_ *
+      (::Eigen::Matrix<double, 2, 1>() << 1.0, (opposite_voltages ? -1.0 : 1.0))
+          .finished();
+
+  double min_valid_vel = ::std::numeric_limits<double>::infinity();
+  if (limit_type == VoltageLimit::kAggressive) {
+    min_valid_vel = 0.0;
+  }
+  // Iterate over both possibilites for +/- voltage, and solve the quadratic
+  // formula. For every positive solution, adjust the velocity limit
+  // appropriately.
+  for (const double sign : {1.0, -1.0}) {
+    const ::Eigen::Matrix<double, 2, 1> U = sign * abs_volts;
+    const double prev_vel = min_valid_vel;
+    const double c = c_coeff * U;
+    const double determinant = b * b - 4 * a * c;
+    if (a == 0) {
+      // If a == 0, that implies we are on a constant curvature path, in which
+      // case we just have b * v + c = 0.
+      // Note that if -b * c > 0.0, then vel will be greater than zero and b
+      // will be non-zero.
+      if (-b * c > 0.0) {
+        const double vel = -c / b;
+        if (limit_type == VoltageLimit::kConservative) {
+          min_valid_vel = ::std::min(min_valid_vel, vel);
+        } else {
+          min_valid_vel = ::std::max(min_valid_vel, vel);
+        }
+      } else if (b == 0) {
+        // If a and b are zero, then we are travelling in a straight line and
+        // have no voltage-based velocity constraints.
+        min_valid_vel = ::std::numeric_limits<double>::infinity();
+      }
+    } else if (determinant > 0) {
+      const double sqrt_determinant = ::std::sqrt(determinant);
+      const double high_vel = (-b + sqrt_determinant) / (2.0 * a);
+      const double low_vel = (-b - sqrt_determinant) / (2.0 * a);
+      if (low_vel > 0) {
+        if (limit_type == VoltageLimit::kConservative) {
+          min_valid_vel = ::std::min(min_valid_vel, low_vel);
+        } else {
+          min_valid_vel = ::std::max(min_valid_vel, low_vel);
+        }
+      }
+      if (high_vel > 0) {
+        if (limit_type == VoltageLimit::kConservative) {
+          min_valid_vel = ::std::min(min_valid_vel, high_vel);
+        } else {
+          min_valid_vel = ::std::max(min_valid_vel, high_vel);
+        }
+      }
+    } else if (determinant == 0 && -b * a > 0) {
+      const double vel = -b / (2.0 * a);
+      if (vel > 0.0) {
+        if (limit_type == VoltageLimit::kConservative) {
+          min_valid_vel = ::std::min(min_valid_vel, vel);
+        } else {
+          min_valid_vel = ::std::max(min_valid_vel, vel);
+        }
+      }
+    }
+    if (constraint_voltages != nullptr && prev_vel != min_valid_vel) {
+      *constraint_voltages = U;
+    }
+  }
+  return min_valid_vel;
 }
 
 void Trajectory::ForwardPass() {
@@ -107,57 +342,13 @@
         [this](double x, double v) { return ForwardAcceleration(x, v); },
         plan_[i], distance, delta_distance);
 
-    if (new_plan_velocity < plan_[i + 1]) {
+    if (new_plan_velocity <= plan_[i + 1]) {
       plan_[i + 1] = new_plan_velocity;
       plan_segment_type_[i] = SegmentType::ACCELERATION_LIMITED;
     }
   }
 }
 
-double Trajectory::BackwardAcceleration(double x, double v) {
-  ::Eigen::Matrix<double, 2, 1> K3;
-  ::Eigen::Matrix<double, 2, 1> K4;
-  ::Eigen::Matrix<double, 2, 1> K5;
-  K345(x, &K3, &K4, &K5);
-
-  // Now, solve for all a's and find the best one which meets our criteria.
-  const ::Eigen::Matrix<double, 2, 1> C = K3 * v * v + K4 * v;
-  double mina = ::std::numeric_limits<double>::infinity();
-  for (const double a : {(voltage_limit_ - C(0, 0)) / K5(0, 0),
-                         (voltage_limit_ - C(1, 0)) / K5(1, 0),
-                         (-voltage_limit_ - C(0, 0)) / K5(0, 0),
-                         (-voltage_limit_ - C(1, 0)) / K5(1, 0)}) {
-    const ::Eigen::Matrix<double, 2, 1> U = K5 * a + K3 * v * v + K4 * v;
-    if ((U.array().abs() < voltage_limit_ + 1e-6).all()) {
-      mina = ::std::min(mina, a);
-    }
-  }
-
-  // Then, assume an acceleration oval and stay inside it.
-  const double lateral_acceleration = v * v * spline_->DDXY(x).norm();
-  const double ellipse_down_shift = longitudal_acceleration_ * 1.0;
-  const double ellipse_width_stretch = ::std::sqrt(
-      1.0 / (1.0 - ::std::pow(ellipse_down_shift / (longitudal_acceleration_ +
-                                                    ellipse_down_shift),
-                              2.0)));
-  const double squared =
-      1.0 - ::std::pow(lateral_acceleration / lateral_acceleration_ /
-                           ellipse_width_stretch,
-                       2.0);
-  // If we would end up with an imaginary number, cap us at 0 acceleration.
-  // TODO(austin): Investigate when this happens, why, and fix it.
-  if (squared < 0.0) {
-    AOS_LOG(ERROR, "Imaginary %f, mina: %f, fa(%f, %f) -> 0.0\n", squared, mina,
-            x, v);
-    return 0.0;
-  }
-  const double longitudal_acceleration =
-      -::std::sqrt(::std::abs(squared)) *
-          (longitudal_acceleration_ + ellipse_down_shift) +
-      ellipse_down_shift;
-  return ::std::max(longitudal_acceleration, mina);
-}
-
 void Trajectory::BackwardPass() {
   const double delta_distance = Distance(0) - Distance(1);
   plan_.back() = 0.0;
@@ -169,7 +360,7 @@
         [this](double x, double v) { return BackwardAcceleration(x, v); },
         plan_[i], distance, delta_distance);
 
-    if (new_plan_velocity < plan_[i - 1]) {
+    if (new_plan_velocity <= plan_[i - 1]) {
       plan_[i - 1] = new_plan_velocity;
       plan_segment_type_[i - 1] = SegmentType::DECELERATION_LIMITED;
     }
@@ -195,6 +386,9 @@
 
   double acceleration;
   double velocity;
+  // TODO(james): While technically correct for sufficiently small segment
+  // steps, this method of switching between limits has a tendency to produce
+  // sudden jumps in acceelrations, which is undesirable.
   switch (plan_segment_type_[DistanceToSegment(distance)]) {
     case SegmentType::VELOCITY_LIMITED:
       acceleration = 0.0;
@@ -205,9 +399,20 @@
       break;
     case SegmentType::CURVATURE_LIMITED:
       velocity = vcurvature;
-      acceleration = 0.0;
+      FrictionLngAccelLimits(distance, velocity, &acceleration, &acceleration);
+      break;
+    case SegmentType::VOLTAGE_LIMITED:
+      // Normally, we expect that voltage limited plans will all get dominated
+      // by the acceleration/deceleration limits. This may not always be true;
+      // if we ever encounter this error, we just need to back out what the
+      // accelerations would be in this case.
+      LOG(FATAL) <<  "Unexpectedly got VOLTAGE_LIMITED plan.";
       break;
     case SegmentType::ACCELERATION_LIMITED:
+      // TODO(james): The integration done here and in the DECELERATION_LIMITED
+      // can technically cause us to violate friction constraints. We currently
+      // don't do anything about it to avoid causing sudden jumps in voltage,
+      // but we probably *should* at some point.
       velocity = IntegrateAccelForDistance(
           [this](double x, double v) { return ForwardAcceleration(x, v); },
           plan_[before_index], before_distance, distance - before_distance);
@@ -226,11 +431,6 @@
       break;
   }
 
-  if (vcurvature < velocity) {
-    velocity = vcurvature;
-    acceleration = 0.0;
-    AOS_LOG(ERROR, "Curvature won\n");
-  }
   return (::Eigen::Matrix<double, 3, 1>() << distance, velocity, acceleration)
       .finished();
 }