Add a catapult MPC solver in C++ which runs fast enough
Change-Id: Ieccddb951dce287fbdd74d75916896e97360c50b
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/y2022/control_loops/superstructure/catapult/BUILD b/y2022/control_loops/superstructure/catapult/BUILD
new file mode 100644
index 0000000..289dfde
--- /dev/null
+++ b/y2022/control_loops/superstructure/catapult/BUILD
@@ -0,0 +1,68 @@
+genrule(
+ name = "genrule_catapult",
+ outs = [
+ "catapult_plant.h",
+ "catapult_plant.cc",
+ "integral_catapult_plant.h",
+ "integral_catapult_plant.cc",
+ ],
+ cmd = "$(location //y2022/control_loops/python:catapult) $(OUTS)",
+ target_compatible_with = ["@platforms//os:linux"],
+ tools = [
+ "//y2022/control_loops/python:catapult",
+ ],
+)
+
+cc_library(
+ name = "catapult_plants",
+ srcs = [
+ "catapult_plant.cc",
+ "integral_catapult_plant.cc",
+ ],
+ hdrs = [
+ "catapult_plant.h",
+ "integral_catapult_plant.h",
+ ],
+ visibility = ["//visibility:public"],
+ deps = [
+ "//frc971/control_loops:state_feedback_loop",
+ ],
+)
+
+cc_library(
+ name = "catapult",
+ srcs = [
+ "catapult.cc",
+ ],
+ hdrs = [
+ "catapult.h",
+ ],
+ visibility = ["//visibility:public"],
+ deps = [
+ ":catapult_plants",
+ "//aos:realtime",
+ "//third_party/osqp-cpp",
+ ],
+)
+
+cc_test(
+ name = "catapult_test",
+ srcs = [
+ "catapult_test.cc",
+ ],
+ deps = [
+ ":catapult",
+ "//aos/testing:googletest",
+ ],
+)
+
+cc_binary(
+ name = "catapult_main",
+ srcs = [
+ "catapult_main.cc",
+ ],
+ deps = [
+ ":catapult",
+ "//aos:init",
+ ],
+)
diff --git a/y2022/control_loops/superstructure/catapult/catapult.cc b/y2022/control_loops/superstructure/catapult/catapult.cc
new file mode 100644
index 0000000..243cb2d
--- /dev/null
+++ b/y2022/control_loops/superstructure/catapult/catapult.cc
@@ -0,0 +1,260 @@
+#include "y2022/control_loops/superstructure/catapult/catapult.h"
+
+#include "Eigen/Dense"
+#include "Eigen/Sparse"
+#include "aos/realtime.h"
+#include "aos/time/time.h"
+#include "glog/logging.h"
+#include "osqp++.h"
+#include "osqp.h"
+#include "y2022/control_loops/superstructure/catapult/catapult_plant.h"
+
+namespace y2022 {
+namespace control_loops {
+namespace superstructure {
+namespace catapult {
+namespace chrono = std::chrono;
+
+namespace {
+osqp::OsqpInstance MakeInstance(
+ size_t horizon, Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> P) {
+ osqp::OsqpInstance instance;
+ instance.objective_matrix = P.cast<c_float>().sparseView();
+
+ instance.constraint_matrix =
+ Eigen::SparseMatrix<c_float, Eigen::ColMajor, osqp::c_int>(horizon,
+ horizon);
+ instance.constraint_matrix.setIdentity();
+
+ instance.lower_bounds =
+ Eigen::Matrix<c_float, Eigen::Dynamic, 1>::Zero(horizon, 1);
+ instance.upper_bounds =
+ Eigen::Matrix<c_float, Eigen::Dynamic, 1>::Ones(horizon, 1) * 12.0;
+ return instance;
+}
+} // namespace
+
+MPCProblem::MPCProblem(size_t horizon,
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> P,
+ Eigen::Matrix<double, Eigen::Dynamic, 1> accel_q,
+ Eigen::Matrix<double, 2, 2> Af,
+ Eigen::Matrix<double, Eigen::Dynamic, 2> final_q)
+ : horizon_(horizon),
+ accel_q_(std::move(accel_q.cast<c_float>())),
+ Af_(std::move(Af.cast<c_float>())),
+ final_q_(std::move(final_q.cast<c_float>())),
+ instance_(MakeInstance(horizon, std::move(P))) {
+ instance_.objective_vector =
+ Eigen::Matrix<c_float, Eigen::Dynamic, 1>(horizon, 1);
+ settings_.max_iter = 25;
+ settings_.check_termination = 25;
+ settings_.warm_start = 0;
+ auto status = solver_.Init(instance_, settings_);
+ CHECK(status.ok()) << status;
+}
+
+void MPCProblem::SetState(Eigen::Matrix<c_float, 2, 1> X_initial,
+ Eigen::Matrix<c_float, 2, 1> X_final) {
+ X_initial_ = X_initial;
+ X_final_ = X_final;
+ instance_.objective_vector =
+ X_initial(1, 0) * accel_q_ + final_q_ * (Af_ * X_initial - X_final);
+
+ auto status = solver_.SetObjectiveVector(instance_.objective_vector);
+ CHECK(status.ok()) << status;
+}
+
+bool MPCProblem::Solve() {
+ const aos::monotonic_clock::time_point start_time =
+ aos::monotonic_clock::now();
+ osqp::OsqpExitCode exit_code = solver_.Solve();
+ const aos::monotonic_clock::time_point end_time = aos::monotonic_clock::now();
+ VLOG(1) << "OSQP solved in "
+ << std::chrono::duration<double>(end_time - start_time).count();
+ solve_time_ = std::chrono::duration<double>(end_time - start_time).count();
+ // TODO(austin): Dump the exit codes out as an enum for logging.
+ //
+ // TODO(austin): The dual problem doesn't appear to be converging on all
+ // problems. Are we phrasing something wrong?
+
+ // TODO(austin): Set a time limit so we can't run forever, and signal back
+ // when we hit our limit.
+ return exit_code == osqp::OsqpExitCode::kOptimal;
+}
+
+void MPCProblem::WarmStart(const MPCProblem &p) {
+ CHECK_GE(p.horizon(), horizon())
+ << ": Can only copy a bigger problem's solution into a smaller problem.";
+ auto status = solver_.SetPrimalWarmStart(p.solver_.primal_solution().block(
+ p.horizon() - horizon(), 0, horizon(), 1));
+ CHECK(status.ok()) << status;
+ status = solver_.SetDualWarmStart(p.solver_.dual_solution().block(
+ p.horizon() - horizon(), 0, horizon(), 1));
+ CHECK(status.ok()) << status;
+}
+
+CatapultProblemGenerator::CatapultProblemGenerator(size_t horizon)
+ : plant_(MakeCatapultPlant()),
+ horizon_(horizon),
+ Q_final_(
+ (Eigen::DiagonalMatrix<double, 2>().diagonal() << 10000.0, 10000.0)
+ .finished()),
+ As_(MakeAs()),
+ Bs_(MakeBs()),
+ m_(Makem()),
+ M_(MakeM()),
+ W_(MakeW()),
+ w_(Makew()),
+ Pi_(MakePi()),
+ WM_(W_ * M_),
+ Wmpw_(W_ * m_ + w_) {}
+
+std::unique_ptr<MPCProblem> CatapultProblemGenerator::MakeProblem(
+ size_t horizon) {
+ return std::make_unique<MPCProblem>(
+ horizon, P(horizon), accel_q(horizon), Af(horizon),
+ (2.0 * Q_final_ * Bf(horizon)).transpose());
+}
+
+const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
+CatapultProblemGenerator::P(size_t horizon) {
+ CHECK_GT(horizon, 0u);
+ CHECK_LE(horizon, horizon_);
+ return 2.0 * (WM_.block(0, 0, horizon, horizon).transpose() * Pi(horizon) *
+ WM_.block(0, 0, horizon, horizon) +
+ Bf(horizon).transpose() * Q_final_ * Bf(horizon));
+}
+
+const Eigen::Matrix<double, Eigen::Dynamic, 1> CatapultProblemGenerator::q(
+ size_t horizon, Eigen::Matrix<c_float, 2, 1> X_initial,
+ Eigen::Matrix<c_float, 2, 1> X_final) {
+ CHECK_GT(horizon, 0u);
+ CHECK_LE(horizon, horizon_);
+ return 2.0 * X_initial(1, 0) * accel_q(horizon) +
+ 2.0 * ((Af(horizon) * X_initial - X_final).transpose() * Q_final_ *
+ Bf(horizon))
+ .transpose();
+}
+
+const Eigen::Matrix<double, Eigen::Dynamic, 1>
+CatapultProblemGenerator::accel_q(size_t horizon) {
+ return 2.0 * ((Wmpw_.block(0, 0, horizon, 1)).transpose() * Pi(horizon) *
+ WM_.block(0, 0, horizon, horizon))
+ .transpose();
+}
+
+const Eigen::Matrix<double, 2, 2> CatapultProblemGenerator::Af(size_t horizon) {
+ CHECK_GT(horizon, 0u);
+ CHECK_LE(horizon, horizon_);
+ return As_.block<2, 2>(2 * (horizon - 1), 0);
+}
+
+const Eigen::Matrix<double, 2, Eigen::Dynamic> CatapultProblemGenerator::Bf(
+ size_t horizon) {
+ CHECK_GT(horizon, 0u);
+ CHECK_LE(horizon, horizon_);
+ return Bs_.block(2 * (horizon - 1), 0, 2, horizon);
+}
+
+const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
+CatapultProblemGenerator::Pi(size_t horizon) {
+ CHECK_GT(horizon, 0u);
+ CHECK_LE(horizon, horizon_);
+ return Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>(Pi_).block(
+ horizon_ - horizon, horizon_ - horizon, horizon, horizon);
+}
+
+Eigen::Matrix<double, Eigen::Dynamic, 2> CatapultProblemGenerator::MakeAs() {
+ Eigen::Matrix<double, Eigen::Dynamic, 2> As =
+ Eigen::Matrix<double, Eigen::Dynamic, 2>::Zero(horizon_ * 2, 2);
+ for (size_t i = 0; i < horizon_; ++i) {
+ if (i == 0) {
+ As.block<2, 2>(0, 0) = plant_.A();
+ } else {
+ As.block<2, 2>(i * 2, 0) = plant_.A() * As.block<2, 2>((i - 1) * 2, 0);
+ }
+ }
+ return As;
+}
+
+Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
+CatapultProblemGenerator::MakeBs() {
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Bs =
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>::Zero(horizon_ * 2,
+ horizon_);
+ for (size_t i = 0; i < horizon_; ++i) {
+ for (size_t j = 0; j < i + 1; ++j) {
+ if (i == j) {
+ Bs.block<2, 1>(i * 2, j) = plant_.B();
+ } else {
+ Bs.block<2, 1>(i * 2, j) =
+ As_.block<2, 2>((i - j - 1) * 2, 0) * plant_.B();
+ }
+ }
+ }
+ return Bs;
+}
+
+Eigen::Matrix<double, Eigen::Dynamic, 1> CatapultProblemGenerator::Makem() {
+ Eigen::Matrix<double, Eigen::Dynamic, 1> m =
+ Eigen::Matrix<double, Eigen::Dynamic, 1>::Zero(horizon_, 1);
+ for (size_t i = 0; i < horizon_; ++i) {
+ m(i, 0) = As_(1 + 2 * i, 1);
+ }
+ return m;
+}
+
+Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
+CatapultProblemGenerator::MakeM() {
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> M =
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>::Zero(horizon_,
+ horizon_);
+ for (size_t i = 0; i < horizon_; ++i) {
+ for (size_t j = 0; j < horizon_; ++j) {
+ M(i, j) = Bs_(2 * i + 1, j);
+ }
+ }
+ return M;
+}
+
+Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
+CatapultProblemGenerator::MakeW() {
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> W =
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>::Identity(horizon_,
+ horizon_);
+ for (size_t i = 0; i < horizon_ - 1; ++i) {
+ W(i + 1, i) = -1.0;
+ }
+ W /= std::chrono::duration<double>(plant_.dt()).count();
+ return W;
+}
+
+Eigen::Matrix<double, Eigen::Dynamic, 1> CatapultProblemGenerator::Makew() {
+ Eigen::Matrix<double, Eigen::Dynamic, 1> w =
+ Eigen::Matrix<double, Eigen::Dynamic, 1>::Zero(horizon_, 1);
+ w(0, 0) = -1.0 / std::chrono::duration<double>(plant_.dt()).count();
+ return w;
+}
+
+Eigen::DiagonalMatrix<double, Eigen::Dynamic>
+CatapultProblemGenerator::MakePi() {
+ Eigen::DiagonalMatrix<double, Eigen::Dynamic> Pi(horizon_);
+ for (size_t i = 0; i < horizon_; ++i) {
+ Pi.diagonal()(i) =
+ std::pow(0.01, 2.0) +
+ std::pow(0.02 * std::max(0.0, (20 - ((int)horizon_ - (int)i)) / 20.),
+ 2.0);
+ }
+ return Pi;
+}
+
+Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
+CatapultProblemGenerator::MakeP() {
+ return 2.0 * (M_.transpose() * W_.transpose() * Pi_ * W_ * M_ +
+ Bf(horizon_).transpose() * Q_final_ * Bf(horizon_));
+}
+
+} // namespace catapult
+} // namespace superstructure
+} // namespace control_loops
+} // namespace y2022
diff --git a/y2022/control_loops/superstructure/catapult/catapult.h b/y2022/control_loops/superstructure/catapult/catapult.h
new file mode 100644
index 0000000..85e9242
--- /dev/null
+++ b/y2022/control_loops/superstructure/catapult/catapult.h
@@ -0,0 +1,132 @@
+#ifndef Y2022_CONTROL_LOOPS_SUPERSTRUCTURE_CATAPULT_CATAPULT_H_
+#define Y2022_CONTROL_LOOPS_SUPERSTRUCTURE_CATAPULT_CATAPULT_H_
+
+#include "Eigen/Dense"
+#include "frc971/control_loops/state_feedback_loop.h"
+#include "glog/logging.h"
+#include "osqp++.h"
+
+namespace y2022 {
+namespace control_loops {
+namespace superstructure {
+namespace catapult {
+
+// MPC problem for a specified horizon. This contains all the state for the
+// solver, setters to modify the current and target state, and a way to fetch
+// the solution.
+class MPCProblem {
+ public:
+ MPCProblem(size_t horizon,
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> P,
+ Eigen::Matrix<double, Eigen::Dynamic, 1> accel_q,
+ Eigen::Matrix<double, 2, 2> Af,
+ Eigen::Matrix<double, Eigen::Dynamic, 2> final_q);
+
+ MPCProblem(MPCProblem const &) = delete;
+ void operator=(MPCProblem const &x) = delete;
+
+ // Sets the current and final state. This keeps the problem in tact and
+ // doesn't recreate it, so it will be fast.
+ void SetState(Eigen::Matrix<double, 2, 1> X_initial,
+ Eigen::Matrix<double, 2, 1> X_final);
+
+ // Solves our problem.
+ bool Solve();
+
+ double solve_time() const { return solve_time_; }
+
+ // Returns the solution that the solver found when Solve was last called.
+ double U(size_t i) const { return solver_.primal_solution()(i); }
+
+ // Returns the number of U's to be solved.
+ size_t horizon() const { return horizon_; }
+
+ // Warm starts the optimizer with the provided solution to make it solve
+ // faster.
+ void WarmStart(const MPCProblem &p);
+
+ private:
+ // The number of u's to solve for.
+ const size_t horizon_;
+
+ // The problem statement variables needed by SetState to update q.
+ const Eigen::Matrix<double, Eigen::Dynamic, 1> accel_q_;
+ const Eigen::Matrix<double, 2, 2> Af_;
+ const Eigen::Matrix<double, Eigen::Dynamic, 2> final_q_;
+
+ Eigen::Matrix<double, 2, 1> X_initial_;
+ Eigen::Matrix<double, 2, 1> X_final_;
+
+ // Solver state.
+ osqp::OsqpInstance instance_;
+ osqp::OsqpSolver solver_;
+ osqp::OsqpSettings settings_;
+
+ double solve_time_ = 0;
+};
+
+// Decently efficient problem generator for multiple horizons given a max
+// horizon to solve for.
+//
+// The math is documented in mpc.tex
+class CatapultProblemGenerator {
+ public:
+ // Builds a problem generator for the specified max horizon and caches a lot
+ // of the state.
+ CatapultProblemGenerator(size_t horizon);
+
+ // Returns the maximum horizon.
+ size_t horizon() const { return horizon_; }
+
+ // Makes a problem for the specificed horizon.
+ std::unique_ptr<MPCProblem> MakeProblem(size_t horizon);
+
+ // Returns the P and Q matrices for the problem statement.
+ // cost = 0.5 X.T P X + q.T X
+ const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> P(size_t horizon);
+ const Eigen::Matrix<double, Eigen::Dynamic, 1> q(
+ size_t horizon, Eigen::Matrix<double, 2, 1> X_initial,
+ Eigen::Matrix<double, 2, 1> X_final);
+
+ private:
+ const Eigen::Matrix<double, Eigen::Dynamic, 1> accel_q(size_t horizon);
+
+ const Eigen::Matrix<double, 2, 2> Af(size_t horizon);
+ const Eigen::Matrix<double, 2, Eigen::Dynamic> Bf(size_t horizon);
+ const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Pi(
+ size_t horizon);
+
+ // These functions are used in the constructor to build up the matrices below.
+ Eigen::Matrix<double, Eigen::Dynamic, 2> MakeAs();
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> MakeBs();
+ Eigen::Matrix<double, Eigen::Dynamic, 1> Makem();
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> MakeM();
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> MakeW();
+ Eigen::Matrix<double, Eigen::Dynamic, 1> Makew();
+ Eigen::DiagonalMatrix<double, Eigen::Dynamic> MakePi();
+ Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> MakeP();
+
+ const StateFeedbackPlant<2, 1, 1> plant_;
+ const size_t horizon_;
+
+ const Eigen::DiagonalMatrix<double, 2> Q_final_;
+
+ const Eigen::Matrix<double, Eigen::Dynamic, 2> As_;
+ const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Bs_;
+ const Eigen::Matrix<double, Eigen::Dynamic, 1> m_;
+ const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> M_;
+
+ const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> W_;
+ const Eigen::Matrix<double, Eigen::Dynamic, 1> w_;
+ const Eigen::DiagonalMatrix<double, Eigen::Dynamic> Pi_;
+
+ const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> WM_;
+ const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Wmpw_;
+};
+
+} // namespace catapult
+} // namespace superstructure
+} // namespace control_loops
+} // namespace y2022
+
+#endif // Y2022_CONTROL_LOOPS_SUPERSTRUCTURE_CATAPULT_CATAPULT_H_
diff --git a/y2022/control_loops/superstructure/catapult/catapult_main.cc b/y2022/control_loops/superstructure/catapult/catapult_main.cc
new file mode 100644
index 0000000..53d357a
--- /dev/null
+++ b/y2022/control_loops/superstructure/catapult/catapult_main.cc
@@ -0,0 +1,126 @@
+#include "aos/init.h"
+#include "aos/realtime.h"
+#include "aos/time/time.h"
+#include "y2022/control_loops/superstructure/catapult/catapult.h"
+#include "y2022/control_loops/superstructure/catapult/catapult_plant.h"
+
+namespace y2022 {
+namespace control_loops {
+namespace superstructure {
+namespace catapult {
+namespace chrono = std::chrono;
+
+void OSQPSolve() {
+ Eigen::Matrix<double, 2, 1> X_initial(0.0, 0.0);
+ Eigen::Matrix<double, 2, 1> X_final(2.0, 25.0);
+
+ LOG(INFO) << "Starting a dynamic OSQP solve";
+ CatapultProblemGenerator g(35);
+ const StateFeedbackPlant<2, 1, 1> plant = MakeCatapultPlant();
+
+ constexpr int kHorizon = 35;
+
+ // TODO(austin): This is a good unit test! Make sure computing the problem
+ // different ways comes out the same.
+ {
+ CatapultProblemGenerator g2(10);
+ constexpr int kTestHorizon = 10;
+ CHECK(g2.P(kTestHorizon) == g.P(kTestHorizon))
+ << g2.P(kTestHorizon) - g.P(kTestHorizon);
+ CHECK(g2.q(kTestHorizon, X_initial, X_final) ==
+ g.q(kTestHorizon, X_initial, X_final))
+ << g2.q(kTestHorizon, X_initial, X_final) -
+ g.q(kTestHorizon, X_initial, X_final);
+ }
+
+ osqp::OsqpInstance instance;
+ instance.objective_matrix = g.P(kHorizon).sparseView();
+
+ instance.objective_vector = g.q(kHorizon, X_initial, X_final);
+
+ instance.constraint_matrix =
+ Eigen::SparseMatrix<double, Eigen::ColMajor, osqp::c_int>(kHorizon,
+ kHorizon);
+ instance.constraint_matrix.setIdentity();
+
+ instance.lower_bounds =
+ Eigen::Matrix<double, Eigen::Dynamic, 1>::Zero(kHorizon, 1);
+ instance.upper_bounds =
+ Eigen::Matrix<double, Eigen::Dynamic, 1>::Ones(kHorizon, 1) * 12.0;
+
+ osqp::OsqpSolver solver;
+ osqp::OsqpSettings settings;
+ // Edit settings if appropriate.
+ auto status = solver.Init(instance, settings);
+ CHECK(status.ok()) << status;
+
+ aos::LockAllMemory();
+ aos::ExpandStackSize();
+ aos::SetCurrentThreadRealtimePriority(60);
+
+ for (int i = 0; i < 10; ++i) {
+ const aos::monotonic_clock::time_point start_time =
+ aos::monotonic_clock::now();
+ osqp::OsqpExitCode exit_code = solver.Solve();
+ const aos::monotonic_clock::time_point end_time =
+ aos::monotonic_clock::now();
+ LOG(INFO) << "OSQP solved in "
+ << chrono::duration<double>(end_time - start_time).count();
+ CHECK(exit_code == osqp::OsqpExitCode::kOptimal);
+ }
+
+ double optimal_objective = solver.objective_value();
+ Eigen::Matrix<double, Eigen::Dynamic, 1> optimal_solution =
+ solver.primal_solution();
+
+ LOG(INFO) << "Cost: " << optimal_objective;
+ LOG(INFO) << "U: " << optimal_solution;
+
+ std::vector<std::unique_ptr<MPCProblem>> problems;
+ for (size_t i = g.horizon(); i > 0; --i) {
+ LOG(INFO) << "Made problem " << i;
+ problems.emplace_back(g.MakeProblem(i));
+ }
+
+ std::unique_ptr<MPCProblem> p = g.MakeProblem(kHorizon);
+
+ p->SetState(X_initial, X_final);
+ p->Solve();
+ p->Solve();
+ p->Solve();
+ p->Solve();
+
+ Eigen::Vector2d X = X_initial;
+ for (size_t i = 0; i < g.horizon(); ++i) {
+ problems[i]->SetState(X, X_final);
+ if (i != 0) {
+ problems[i]->WarmStart(*problems[i - 1]);
+ }
+
+ problems[i]->Solve();
+ X = plant.A() * X + plant.B() * problems[i]->U(0);
+
+ LOG(INFO) << "Dynamic u(" << i << "): " << problems[i]->U(0) << " -> "
+ << X.transpose();
+ }
+
+ aos::UnsetCurrentThreadRealtimePriority();
+}
+
+int Main(int /*argc*/, char ** /*argv*/) {
+ OSQPSolve();
+
+ gflags::ShutDownCommandLineFlags();
+ return 0;
+}
+
+} // namespace catapult
+} // namespace superstructure
+} // namespace control_loops
+} // namespace y2022
+
+int main(int argc, char **argv) {
+ ::aos::InitGoogle(&argc, &argv);
+
+ return y2022::control_loops::superstructure::catapult::Main(argc, argv);
+}
diff --git a/y2022/control_loops/superstructure/catapult/catapult_test.cc b/y2022/control_loops/superstructure/catapult/catapult_test.cc
new file mode 100644
index 0000000..b37c3e0
--- /dev/null
+++ b/y2022/control_loops/superstructure/catapult/catapult_test.cc
@@ -0,0 +1,62 @@
+#include "y2022/control_loops/superstructure/catapult/catapult.h"
+
+#include "glog/logging.h"
+#include "gtest/gtest.h"
+#include "y2022/control_loops/superstructure/catapult/catapult_plant.h"
+
+namespace y2022 {
+namespace control_loops {
+namespace superstructure {
+namespace catapult {
+namespace testing {
+
+// Tests that computing P and q with 2 different horizons comes out the same.
+TEST(MPCTest, HorizonTest) {
+ Eigen::Matrix<double, 2, 1> X_initial(0.0, 0.0);
+ Eigen::Matrix<double, 2, 1> X_final(2.0, 25.0);
+
+ CatapultProblemGenerator g(35);
+
+ CatapultProblemGenerator g2(10);
+ constexpr int kTestHorizon = 10;
+ EXPECT_TRUE(g2.P(kTestHorizon) == g.P(kTestHorizon))
+ << g2.P(kTestHorizon) - g.P(kTestHorizon);
+ EXPECT_TRUE(g2.q(kTestHorizon, X_initial, X_final) ==
+ g.q(kTestHorizon, X_initial, X_final))
+ << g2.q(kTestHorizon, X_initial, X_final) -
+ g.q(kTestHorizon, X_initial, X_final);
+}
+
+// Tests that we can get to the destination.
+TEST(MPCTest, NearGoal) {
+ Eigen::Matrix<double, 2, 1> X_initial(0.0, 0.0);
+ Eigen::Matrix<double, 2, 1> X_final(2.0, 25.0);
+
+ CatapultProblemGenerator g(35);
+ const StateFeedbackPlant<2, 1, 1> plant = MakeCatapultPlant();
+
+ std::vector<std::unique_ptr<MPCProblem>> problems;
+ for (size_t i = g.horizon(); i > 0; --i) {
+ problems.emplace_back(g.MakeProblem(i));
+ }
+
+ Eigen::Vector2d X = X_initial;
+ for (size_t i = 0; i < g.horizon(); ++i) {
+ problems[i]->SetState(X, X_final);
+ if (i != 0) {
+ problems[i]->WarmStart(*problems[i - 1]);
+ }
+
+ problems[i]->Solve();
+ X = plant.A() * X + plant.B() * problems[i]->U(0);
+ }
+
+ EXPECT_NEAR(X_final.x(), X.x(), 1e-2);
+ EXPECT_NEAR(X_final.y(), X.y(), 1e-2);
+}
+
+} // namespace testing
+} // namespace catapult
+} // namespace superstructure
+} // namespace control_loops
+} // namespace y2022