Add generic linearized LQR class
This provides a convenient class for taking a system with non-linear
dynamics and generating an LQR controller for it. It still needs some
optimizations to be able to run effectively on a roborio.
Change-Id: I57eaec75738b05713bef272b677aeaa439980354
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index b88fd62..637e42c 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -400,6 +400,34 @@
)
cc_library(
+ name = "linearization_utils",
+ hdrs = ["linearization_utils.h"],
+)
+
+cc_library(
+ name = "linearized_controller",
+ hdrs = ["linearized_controller.h"],
+ deps = [
+ ":eigen_dynamics",
+ ":linearization_utils",
+ "//frc971/control_loops:c2d",
+ "//frc971/control_loops:dlqr",
+ "//frc971/control_loops:jacobian",
+ "@com_google_absl//absl/log",
+ "@com_google_absl//absl/log:check",
+ ],
+)
+
+cc_test(
+ name = "linearized_controller_test",
+ srcs = ["linearized_controller_test.cc"],
+ deps = [
+ ":linearized_controller",
+ "//aos/testing:googletest",
+ ],
+)
+
+cc_library(
name = "auto_diff_jacobian",
hdrs = ["auto_diff_jacobian.h"],
target_compatible_with = ["@platforms//os:linux"],
diff --git a/frc971/control_loops/swerve/linearization_utils.h b/frc971/control_loops/swerve/linearization_utils.h
new file mode 100644
index 0000000..ee1275e
--- /dev/null
+++ b/frc971/control_loops/swerve/linearization_utils.h
@@ -0,0 +1,13 @@
+#ifndef FRC971_CONTROL_LOOPS_SWERVE_LINEARIZATION_UTILS_H_
+#define FRC971_CONTROL_LOOPS_SWERVE_LINEARIZATION_UTILS_H_
+
+namespace frc971::control_loops::swerve {
+template <typename State, typename Input>
+struct DynamicsInterface {
+ virtual ~DynamicsInterface() {}
+ // To be overridden by the implementation; returns the derivative of the state
+ // at the given state with the provided control input.
+ virtual State operator()(const State &X, const Input &U) const = 0;
+};
+} // namespace frc971::control_loops::swerve
+#endif // FRC971_CONTROL_LOOPS_SWERVE_LINEARIZATION_UTILS_H_
diff --git a/frc971/control_loops/swerve/linearized_controller.h b/frc971/control_loops/swerve/linearized_controller.h
new file mode 100644
index 0000000..1df9640
--- /dev/null
+++ b/frc971/control_loops/swerve/linearized_controller.h
@@ -0,0 +1,126 @@
+#include <memory>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include <Eigen/Dense>
+
+#include "frc971/control_loops/c2d.h"
+#include "frc971/control_loops/dlqr.h"
+#include "frc971/control_loops/jacobian.h"
+#include "frc971/control_loops/swerve/dynamics.h"
+#include "frc971/control_loops/swerve/linearization_utils.h"
+
+namespace frc971::control_loops::swerve {
+
+// Provides a simple LQR controller that takes a non-linear system, linearizes
+// the dynamics at each timepoint, recalculates the LQR gains for those
+// dynamics, and calculates the relevant feedback inputs to provide.
+template <int NStates, typename Scalar = double>
+class LinearizedController {
+ public:
+ typedef Eigen::Matrix<Scalar, NStates, 1> State;
+ typedef Eigen::Matrix<Scalar, NStates, NStates> StateSquare;
+ typedef Eigen::Matrix<Scalar, kNumInputs, 1> Input;
+ typedef Eigen::Matrix<Scalar, kNumInputs, kNumInputs> InputSquare;
+ typedef Eigen::Matrix<Scalar, NStates, kNumInputs> BMatrix;
+ typedef DynamicsInterface<State, Input> Dynamics;
+
+ struct Parameters {
+ // State cost matrix.
+ StateSquare Q;
+ // Input cost matrix.
+ InputSquare R;
+ // period at which the controller is called.
+ std::chrono::nanoseconds dt;
+ // The dynamics to use.
+ // TODO(james): I wrote this before creating the auto-differentiation
+ // functions; we should swap to the auto-differentiation, since the
+ // numerical linearization is one of the bigger timesinks in this controller
+ // right now.
+ std::unique_ptr<Dynamics> dynamics;
+ };
+
+ // Represents the linearized dynamics of the system.
+ struct LinearDynamics {
+ StateSquare A;
+ BMatrix B;
+ };
+
+ // Debug information for a given cycle of the controller.
+ struct ControllerDebug {
+ // Feedforward input which we provided.
+ Input U_ff;
+ // Calculated feedback input to provide.
+ Input U_feedback;
+ Eigen::Matrix<Scalar, kNumInputs, NStates> feedback_contributions;
+ };
+
+ struct ControllerResult {
+ // Control input to provide to the robot.
+ Input U;
+ ControllerDebug debug;
+ };
+
+ LinearizedController(Parameters params) : params_(std::move(params)) {}
+
+ // Runs the controller for a given iteration, relinearizing the dynamics about
+ // the provided current state X, attempting to control the robot to the
+ // desired goal state.
+ // The U_ff input will be added into the returned control input.
+ ControllerResult RunController(const State &X, const State &goal,
+ Input U_ff) {
+ auto start_time = aos::monotonic_clock::now();
+ // TODO(james): Swap this to the auto-diff methods; this is currently about
+ // a third of the total time spent in this method when run on the roborio.
+ const struct LinearDynamics continuous_dynamics =
+ LinearizeDynamics(X, U_ff);
+ auto linearization_time = aos::monotonic_clock::now();
+ struct LinearDynamics discrete_dynamics;
+ frc971::controls::C2D(continuous_dynamics.A, continuous_dynamics.B,
+ params_.dt, &discrete_dynamics.A,
+ &discrete_dynamics.B);
+ auto c2d_time = aos::monotonic_clock::now();
+ VLOG(2) << "Controllability of dynamics (ideally should be " << NStates
+ << "): "
+ << frc971::controls::Controllability(discrete_dynamics.A,
+ discrete_dynamics.B);
+ Eigen::Matrix<Scalar, kNumInputs, NStates> K;
+ Eigen::Matrix<Scalar, NStates, NStates> S;
+ // TODO(james): Swap this to a cheaper DARE solver; we should probably just
+ // do something like we do in Trajectory::CalculatePathGains for the tank
+ // spline controller where we approximate the infinite-horizon DARE solution
+ // by doing a finite-horizon LQR.
+ // Currently the dlqr call represents ~60% of the time spent in the
+ // RunController() method.
+ frc971::controls::dlqr(discrete_dynamics.A, discrete_dynamics.B, params_.Q,
+ params_.R, &K, &S);
+ auto dlqr_time = aos::monotonic_clock::now();
+ const Input U_feedback = K * (goal - X);
+ const Input U = U_ff + U_feedback;
+ Eigen::Matrix<Scalar, kNumInputs, NStates> feedback_contributions;
+ for (int state_idx = 0; state_idx < NStates; ++state_idx) {
+ feedback_contributions.col(state_idx) =
+ K.col(state_idx) * (goal - X)(state_idx);
+ }
+ VLOG(2) << "linearization time "
+ << aos::time::DurationInSeconds(linearization_time - start_time)
+ << " c2d time "
+ << aos::time::DurationInSeconds(c2d_time - linearization_time)
+ << " dlqr time "
+ << aos::time::DurationInSeconds(dlqr_time - c2d_time);
+ return {.U = U,
+ .debug = {.U_ff = U_ff,
+ .U_feedback = U_feedback,
+ .feedback_contributions = feedback_contributions}};
+ }
+
+ LinearDynamics LinearizeDynamics(const State &X, const Input &U) {
+ return {.A = NumericalJacobianX(*params_.dynamics, X, U),
+ .B = NumericalJacobianU(*params_.dynamics, X, U)};
+ }
+
+ private:
+ const Parameters params_;
+};
+
+} // namespace frc971::control_loops::swerve
diff --git a/frc971/control_loops/swerve/linearized_controller_test.cc b/frc971/control_loops/swerve/linearized_controller_test.cc
new file mode 100644
index 0000000..d5e7aee
--- /dev/null
+++ b/frc971/control_loops/swerve/linearized_controller_test.cc
@@ -0,0 +1,87 @@
+#include "frc971/control_loops/swerve/linearized_controller.h"
+
+#include "gtest/gtest.h"
+
+namespace frc971::control_loops::swerve::test {
+
+class LinearizedControllerTest : public ::testing::Test {
+ protected:
+ typedef LinearizedController<2> Controller;
+ typedef Controller::State State;
+ typedef Controller::Input Input;
+ struct LinearDynamics : public Controller::Dynamics {
+ Eigen::Vector2d operator()(
+ const Eigen::Vector2d &X,
+ const Eigen::Matrix<double, 8, 1> &U) const override {
+ return Eigen::Matrix2d{{0.0, 1.0}, {0.0, -0.01}} * X +
+ Eigen::Matrix<double, 2, 8>{
+ {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0},
+ {1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}} *
+ U;
+ }
+ };
+ static Eigen::Matrix<double, kNumInputs, kNumInputs> MakeR() {
+ Eigen::Matrix<double, kNumInputs, kNumInputs> R;
+ R.setIdentity();
+ return R;
+ }
+ LinearizedControllerTest()
+ : controller_({.Q = Eigen::Matrix<double, 2, 2>{{1.0, 0.0}, {0.0, 1.0}},
+ .R = MakeR(),
+ .dt = std::chrono::milliseconds(10),
+ .dynamics = std::make_unique<LinearDynamics>()}) {}
+
+ Controller controller_;
+};
+
+// Sanity check that the dynamics linearization is working correctly.
+TEST_F(LinearizedControllerTest, LinearizedDynamics) {
+ auto dynamics = controller_.LinearizeDynamics(State::Zero(), Input::Zero());
+ EXPECT_EQ(0.0, dynamics.A(0, 0));
+ EXPECT_EQ(0.0, dynamics.A(1, 0));
+ EXPECT_EQ(1.0, dynamics.A(0, 1));
+ EXPECT_EQ(-0.01, dynamics.A(1, 1));
+ // All elements of B except for (1, 0) should be exactly 0.
+ EXPECT_EQ(1.0, dynamics.B(1, 0));
+ EXPECT_EQ(1.0, dynamics.B.norm());
+}
+
+// Confirm that the generated LQR controller is able to generate correct
+// inputs when state and goal are at zero.
+TEST_F(LinearizedControllerTest, ControllerResultAtZero) {
+ auto result =
+ controller_.RunController(State::Zero(), State::Zero(), Input::Zero());
+ EXPECT_EQ(0.0, result.U.norm());
+ EXPECT_EQ(0.0, result.debug.U_ff.norm());
+ EXPECT_EQ(0.0, result.debug.U_feedback.norm());
+}
+
+// Confirm that the generated LQR controller is able to generate correct
+// inputs when state is zero and the goal is non-zero.
+TEST_F(LinearizedControllerTest, ControlToZero) {
+ auto result = controller_.RunController(State::Zero(), State{{1.0}, {0.0}},
+ Input::Zero());
+ EXPECT_LT(0.0, result.U(0, 0));
+ // All other U inputs should be zero.
+ EXPECT_EQ(0.0, result.U.bottomRows<7>().norm());
+ EXPECT_EQ(0.0, result.debug.U_ff.norm());
+ EXPECT_EQ(0.0,
+ (result.U - (result.debug.U_ff + result.debug.U_feedback)).norm());
+}
+
+// Confirm that the generated LQR controller is able to pass through the
+// feedforwards when we have no difference between the goal and the current
+// state.
+TEST_F(LinearizedControllerTest, ControlToNonzeroState) {
+ const State state{{1.0}, {1.0}};
+ auto result = controller_.RunController(
+ state, state,
+ Input{{1.0}, {0.0}, {0.0}, {0.0}, {0.0}, {0.0}, {0.0}, {0.0}});
+ EXPECT_EQ(1.0, result.U(0, 0));
+ // All other U inputs should be zero.
+ EXPECT_EQ(0.0, result.U.bottomRows<7>().norm());
+ EXPECT_EQ(0.0, result.debug.U_feedback.norm());
+ EXPECT_EQ(0.0,
+ (result.U - (result.debug.U_ff + result.debug.U_feedback)).norm());
+}
+} // namespace frc971::control_loops::swerve::test