Implement the C++ drivetrain trajectory optimizer
This implements the curvature, forwards, and backwards passes, and adds
a test which makes sure the feed forwards gets us close enough to the
end. Also adds a plotting tool (trajectory_plot) to simulate everything
and tune.
Change-Id: I9f8d6088893cc0b7263b3ff0d79667c027604700
diff --git a/frc971/control_loops/drivetrain/trajectory.cc b/frc971/control_loops/drivetrain/trajectory.cc
new file mode 100644
index 0000000..2c07992
--- /dev/null
+++ b/frc971/control_loops/drivetrain/trajectory.cc
@@ -0,0 +1,378 @@
+#include "frc971/control_loops/drivetrain/trajectory.h"
+
+#include <chrono>
+
+#include "Eigen/Dense"
+#include "aos/logging/matrix_logging.h"
+#include "frc971/control_loops/dlqr.h"
+#include "frc971/control_loops/c2d.h"
+#include "frc971/control_loops/drivetrain/distance_spline.h"
+#include "frc971/control_loops/drivetrain/drivetrain_config.h"
+#include "frc971/control_loops/hybrid_state_feedback_loop.h"
+#include "frc971/control_loops/state_feedback_loop.h"
+
+namespace frc971 {
+namespace control_loops {
+namespace drivetrain {
+
+Trajectory::Trajectory(const DistanceSpline *spline,
+ 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()))),
+ robot_radius_l_(config.robot_radius),
+ robot_radius_r_(config.robot_radius),
+ longitudal_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()),
+ Tla_to_lr_(Tlr_to_la_.inverse()),
+ plan_(num_distance, vmax) {}
+
+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));
+ }
+}
+
+// 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) {
+ ::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();
+ 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);
+ }
+ }
+
+ // Then, assume an acceleration oval and stay inside it.
+ const double lateral_acceleration = v * v * spline_->DDXY(x).norm();
+ const double squared =
+ 1.0 - ::std::pow(lateral_acceleration / lateral_acceleration_, 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) {
+ LOG(ERROR, "Imaginary %f, d %f\n", squared, x);
+ return 0.0;
+ }
+ const double longitudal_acceleration =
+ ::std::sqrt(squared) * longitudal_acceleration_;
+ return ::std::min(longitudal_acceleration, maxa);
+}
+
+void Trajectory::ForwardPass() {
+ plan_[0] = 0.0;
+ const double delta_distance = Distance(1) - Distance(0);
+ for (size_t i = 0; i < plan_.size() - 1; ++i) {
+ const double distance = Distance(i);
+
+ // Integrate our acceleration forward one step.
+ plan_[i + 1] = ::std::min(
+ plan_[i + 1],
+ IntegrateAccelForDistance(
+ [this](double x, double v) { return ForwardAcceleration(x, v); },
+ plan_[i], distance, delta_distance));
+ }
+}
+
+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 squared =
+ 1.0 - ::std::pow(lateral_acceleration / lateral_acceleration_, 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) {
+ LOG(ERROR, "Imaginary %f, d %f\n", squared, x);
+ return 0.0;
+ }
+ const double longitudal_acceleration =
+ -::std::sqrt(squared) * longitudal_acceleration_;
+ return ::std::max(longitudal_acceleration, mina);
+}
+
+void Trajectory::BackwardPass() {
+ const double delta_distance = Distance(0) - Distance(1);
+ plan_.back() = 0.0;
+ for (size_t i = plan_.size() - 1; i > 0; --i) {
+ const double distance = Distance(i);
+
+ // Integrate our deceleration back one step.
+ plan_[i - 1] = ::std::min(
+ plan_[i - 1],
+ IntegrateAccelForDistance(
+ [this](double x, double v) { return BackwardAcceleration(x, v); },
+ plan_[i], distance, delta_distance));
+ }
+}
+
+::Eigen::Matrix<double, 3, 1> Trajectory::FFAcceleration(double distance) {
+ size_t before_index;
+ size_t after_index;
+ if (distance < Distance(1)) {
+ // Within the first step.
+ after_index = 1;
+ // Make sure we don't end up off the beginning of the curve.
+ if (distance < 0.0) {
+ distance = 0.0;
+ }
+ } else if (distance > Distance(plan_.size() - 2)) {
+ // Within the last step.
+ after_index = plan_.size() - 1;
+ // Make sure we don't end up off the end of the curve.
+ if (distance > length()) {
+ distance = length();
+ }
+ } else {
+ // Otherwise do the calculation normally.
+ after_index = static_cast<size_t>(
+ ::std::ceil(distance / length() * (plan_.size() - 1)));
+ }
+ before_index = after_index - 1;
+ const double before_distance = Distance(before_index);
+ const double after_distance = Distance(after_index);
+
+ // Now, compute the velocity that we could have if we accelerated from the
+ // previous step and decelerated from the next step. The min will tell us
+ // which is in effect.
+ const double velocity_forwards = IntegrateAccelForDistance(
+ [this](double x, double v) { return ForwardAcceleration(x, v); },
+ plan_[before_index], before_distance, distance - before_distance);
+ const double velocity_backward = IntegrateAccelForDistance(
+ [this](double x, double v) { return BackwardAcceleration(x, v); },
+ plan_[after_index], after_distance, distance - after_distance);
+
+ // And then also make sure we aren't curvature limited.
+ const double vcurvature = LateralVelocityCurvature(distance);
+
+ double acceleration;
+ double velocity;
+ if (vcurvature < velocity_forwards && vcurvature < velocity_backward) {
+ // If we are curvature limited, we can't accelerate.
+ velocity = vcurvature;
+ acceleration = 0.0;
+ } else if (velocity_forwards < velocity_backward) {
+ // Otherwise, pick the acceleration and velocity from the forward pass if it
+ // was the predominate factor in this step.
+ velocity = velocity_forwards;
+ acceleration = ForwardAcceleration(distance, velocity);
+ } else {
+ // Otherwise, pick the acceleration and velocity from the backward pass if
+ // it was the predominate factor in this step.
+ velocity = velocity_backward;
+ acceleration = BackwardAcceleration(distance, velocity);
+ }
+ return (::Eigen::Matrix<double, 3, 1>() << distance, velocity, acceleration)
+ .finished();
+}
+
+::Eigen::Matrix<double, 2, 1> Trajectory::FFVoltage(double distance) {
+ const Eigen::Matrix<double, 3, 1> xva = FFAcceleration(distance);
+ const double velocity = xva(1);
+ const double acceleration = xva(2);
+ const double current_ddtheta = spline_->DDTheta(distance);
+ const double current_dtheta = spline_->DTheta(distance);
+ // We've now got the equation:
+ // K2 * d^x/dt^2 + K1 (dx/dt)^2 = A * K2 * dx/dt + B * U
+ const ::Eigen::Matrix<double, 2, 1> my_K2 = K2(current_dtheta);
+
+ const ::Eigen::Matrix<double, 2, 2> B_inverse =
+ velocity_drivetrain_->plant().coefficients().B_continuous.inverse();
+
+ // Now, rephrase it as K5 a + K3 v^2 + K4 v = U
+ const ::Eigen::Matrix<double, 2, 1> K5 = B_inverse * my_K2;
+ const ::Eigen::Matrix<double, 2, 1> K3 = B_inverse * K1(current_ddtheta);
+ const ::Eigen::Matrix<double, 2, 1> K4 =
+ -B_inverse * velocity_drivetrain_->plant().coefficients().A_continuous *
+ my_K2;
+
+ return K5 * acceleration + K3 * velocity * velocity + K4 * velocity;
+}
+
+const ::std::vector<double> Trajectory::Distances() const {
+ ::std::vector<double> d;
+ d.reserve(plan_.size());
+ for (size_t i = 0; i < plan_.size(); ++i) {
+ d.push_back(Distance(i));
+ }
+ return d;
+}
+
+::Eigen::Matrix<double, 5, 5> Trajectory::ALinearizedContinuous(
+ const ::Eigen::Matrix<double, 5, 1> &state) const {
+
+ const double sintheta = ::std::sin(state(2));
+ const double costheta = ::std::cos(state(2));
+ const ::Eigen::Matrix<double, 2, 1> linear_angular =
+ Tlr_to_la_ * state.block<2, 1>(3, 0);
+
+ // When stopped, just roll with a min velocity.
+ double linear_velocity = 0.0;
+ constexpr double kMinVelocity = 0.1;
+ if (::std::abs(linear_angular(0)) < kMinVelocity / 100.0) {
+ linear_velocity = 0.1;
+ } else if (::std::abs(linear_angular(0)) > kMinVelocity) {
+ linear_velocity = linear_angular(0);
+ } else if (linear_angular(0) > 0) {
+ linear_velocity = kMinVelocity;
+ } else if (linear_angular(0) < 0) {
+ linear_velocity = -kMinVelocity;
+ }
+
+ ::Eigen::Matrix<double, 5, 5> result = ::Eigen::Matrix<double, 5, 5>::Zero();
+ result(0, 2) = -sintheta * linear_velocity;
+ result(0, 3) = 0.5 * costheta;
+ result(0, 4) = 0.5 * costheta;
+
+ result(1, 2) = costheta * linear_velocity;
+ result(1, 3) = 0.5 * sintheta;
+ result(1, 4) = 0.5 * sintheta;
+
+ result(2, 3) = Tlr_to_la_(1, 0);
+ result(2, 4) = Tlr_to_la_(1, 1);
+
+ result.block<2, 2>(3, 3) =
+ velocity_drivetrain_->plant().coefficients().A_continuous;
+ return result;
+}
+
+::Eigen::Matrix<double, 5, 2> Trajectory::BLinearizedContinuous() const {
+ ::Eigen::Matrix<double, 5, 2> result = ::Eigen::Matrix<double, 5, 2>::Zero();
+ result.block<2, 2>(3, 0) =
+ velocity_drivetrain_->plant().coefficients().B_continuous;
+ return result;
+}
+
+void Trajectory::AB(const ::Eigen::Matrix<double, 5, 1> &state,
+ ::std::chrono::nanoseconds dt,
+ ::Eigen::Matrix<double, 5, 5> *A,
+ ::Eigen::Matrix<double, 5, 2> *B) const {
+ ::Eigen::Matrix<double, 5, 5> A_linearized_continuous =
+ ALinearizedContinuous(state);
+ ::Eigen::Matrix<double, 5, 2> B_linearized_continuous =
+ BLinearizedContinuous();
+
+ // Now, convert it to discrete.
+ controls::C2D(A_linearized_continuous, B_linearized_continuous,
+ dt, A, B);
+}
+
+::Eigen::Matrix<double, 2, 5> Trajectory::KForState(
+ const ::Eigen::Matrix<double, 5, 1> &state, ::std::chrono::nanoseconds dt,
+ const ::Eigen::DiagonalMatrix<double, 5> &Q,
+ const ::Eigen::DiagonalMatrix<double, 2> &R) const {
+ ::Eigen::Matrix<double, 5, 5> A;
+ ::Eigen::Matrix<double, 5, 2> B;
+ AB(state, dt, &A, &B);
+
+ ::Eigen::Matrix<double, 5, 5> S = ::Eigen::Matrix<double, 5, 5>::Zero();
+ ::Eigen::Matrix<double, 2, 5> K = ::Eigen::Matrix<double, 2, 5>::Zero();
+
+ int info = ::frc971::controls::dlqr<5, 2>(A, B, Q, R, &K, &S);
+ if (info == 0) {
+ LOG_MATRIX(INFO, "K", K);
+ } else {
+ LOG(ERROR, "Failed to solve %d, controllability: %d\n", info,
+ controls::Controllability(A, B));
+ // TODO(austin): Can we be more clever here? Use the last one? We should
+ // collect more info about when this breaks down from logs.
+ K = ::Eigen::Matrix<double, 2, 5>::Zero();
+ }
+ ::Eigen::EigenSolver<::Eigen::Matrix<double, 5, 5>> eigensolver(A - B * K);
+ const auto eigenvalues = eigensolver.eigenvalues();
+ LOG(DEBUG,
+ "Eigenvalues: (%f + %fj), (%f + %fj), (%f + %fj), (%f + %fj), (%f + "
+ "%fj)\n",
+ eigenvalues(0).real(), eigenvalues(0).imag(), eigenvalues(1).real(),
+ eigenvalues(1).imag(), eigenvalues(2).real(), eigenvalues(2).imag(),
+ eigenvalues(3).real(), eigenvalues(3).imag(), eigenvalues(4).real(),
+ eigenvalues(4).imag());
+ return K;
+}
+
+const ::Eigen::Matrix<double, 5, 1> Trajectory::GoalState(double distance,
+ double velocity) {
+ ::Eigen::Matrix<double, 5, 1> result;
+ result.block<2, 1>(0, 0) = spline_->XY(distance);
+ result(2, 0) = spline_->Theta(distance);
+
+ result.block<2, 1>(3, 0) = Tla_to_lr_ *
+ (::Eigen::Matrix<double, 2, 1>() << velocity,
+ spline_->DThetaDt(distance, velocity))
+ .finished();
+ return result;
+}
+
+::std::vector<::Eigen::Matrix<double, 3, 1>> Trajectory::PlanXVA(
+ ::std::chrono::nanoseconds dt) {
+ double dt_float =
+ ::std::chrono::duration_cast<::std::chrono::duration<double>>(dt).count();
+ double t = 0.0;
+ ::Eigen::Matrix<double, 2, 1> state = ::Eigen::Matrix<double, 2, 1>::Zero();
+
+ ::std::vector<::Eigen::Matrix<double, 3, 1>> result;
+ result.emplace_back(FFAcceleration(0));
+ result.back()(1) = 0.0;
+
+ while (state(0) < length() - 1e-4) {
+ t += dt_float;
+
+ // TODO(austin): This feels like something that should be pulled out into
+ // a library for re-use.
+ state = RungeKutta(
+ [this](const ::Eigen::Matrix<double, 2, 1> x) {
+ ::Eigen::Matrix<double, 3, 1> xva = FFAcceleration(x(0));
+ return (::Eigen::Matrix<double, 2, 1>() << x(1), xva(2)).finished();
+ },
+ state, dt_float);
+
+ result.emplace_back(FFAcceleration(state(0)));
+ state(1) = result.back()(1);
+ }
+ return result;
+}
+
+} // namespace drivetrain
+} // namespace control_loops
+} // namespace frc971