Add generic drivetrain EKF

This will be followed by a commit with specific utilities for handling
localization in the 2019 game with cameras and targets.

Change-Id: Ic15084e8f7c8b21eb7ac65a493d04da5d2dd941b
diff --git a/frc971/control_loops/drivetrain/BUILD b/frc971/control_loops/drivetrain/BUILD
index ee05fc5..6725987 100644
--- a/frc971/control_loops/drivetrain/BUILD
+++ b/frc971/control_loops/drivetrain/BUILD
@@ -41,6 +41,31 @@
 )
 
 cc_library(
+    name = "hybrid_ekf",
+    hdrs = ["hybrid_ekf.h"],
+    deps = [
+        ":drivetrain_config",
+        "//aos/containers:priority_queue",
+        "//frc971/control_loops:c2d",
+        "//frc971/control_loops:runge_kutta",
+        "//third_party/eigen",
+    ],
+)
+
+cc_test(
+    name = "hybrid_ekf_test",
+    srcs = ["hybrid_ekf_test.cc"],
+    deps = [
+        ":drivetrain_test_lib",
+        ":hybrid_ekf",
+        ":trajectory",
+        "//aos/testing:googletest",
+        "//aos/testing:random_seed",
+        "//aos/testing:test_shm",
+    ],
+)
+
+cc_library(
     name = "gear",
     hdrs = [
         "gear.h",
@@ -57,12 +82,12 @@
         "splinedrivetrain.h",
     ],
     deps = [
+        ":distance_spline",
         ":drivetrain_config",
         ":drivetrain_queue",
         ":spline",
-        ":distance_spline",
         ":trajectory",
-    ]
+    ],
 )
 
 cc_library(
@@ -186,8 +211,8 @@
         ":drivetrain_queue",
         ":gear",
         ":polydrivetrain",
-        ":ssdrivetrain",
         ":splinedrivetrain",
+        ":ssdrivetrain",
         "//aos/controls:control_loop",
         "//aos/logging:matrix_logging",
         "//aos/logging:queue_logging",
@@ -304,6 +329,7 @@
 cc_binary(
     name = "spline.so",
     srcs = ["libspline.cc"],
+    linkshared = True,
     deps = [
         ":distance_spline",
         ":spline",
@@ -313,7 +339,6 @@
         "//third_party/eigen",
         "//y2019/control_loops/drivetrain:drivetrain_base",
     ],
-    linkshared=True,
 )
 
 cc_test(
diff --git a/frc971/control_loops/drivetrain/hybrid_ekf.h b/frc971/control_loops/drivetrain/hybrid_ekf.h
new file mode 100644
index 0000000..245e772
--- /dev/null
+++ b/frc971/control_loops/drivetrain/hybrid_ekf.h
@@ -0,0 +1,452 @@
+#ifndef FRC971_CONTROL_LOOPS_DRIVETRAIN_HYBRID_EKF_H_
+#define FRC971_CONTROL_LOOPS_DRIVETRAIN_HYBRID_EKF_H_
+
+#include <chrono>
+
+#include "aos/containers/priority_queue.h"
+#include "frc971/control_loops/c2d.h"
+#include "frc971/control_loops/runge_kutta.h"
+#include "Eigen/Dense"
+#include "frc971/control_loops/drivetrain/drivetrain_config.h"
+
+namespace frc971 {
+namespace control_loops {
+namespace drivetrain {
+
+namespace testing {
+class HybridEkfTest;
+}
+
+// HybridEkf is an EKF for use in robot localization. It is currently
+// coded for use with drivetrains in particular, and so the states and inputs
+// are chosen as such.
+// The "Hybrid" part of the name refers to the fact that it can take in
+// measurements with variable time-steps.
+// measurements can also have been taken in the past and we maintain a buffer
+// so that we can replay the kalman filter whenever we get an old measurement.
+// Currently, this class provides the necessary utilities for doing
+// measurement updates with an encoder/gyro as well as a more generic
+// update function that can be used for arbitrary nonlinear updates (presumably
+// a camera update).
+template <typename Scalar = double>
+class HybridEkf {
+ public:
+  // An enum specifying what each index in the state vector is for.
+  enum StateIdx {
+    kX = 0,
+    kY = 1,
+    kTheta = 2,
+    kLeftEncoder = 3,
+    kLeftVelocity = 4,
+    kRightEncoder = 5,
+    kRightVelocity = 6,
+    kLeftVoltageError = 7,
+    kRightVoltageError = 8 ,
+    kAngularError = 9,
+  };
+  static constexpr int kNStates = 10;
+  static constexpr int kNInputs = 2;
+  // Number of previous samples to save.
+  static constexpr int kSaveSamples = 50;
+  // Assume that all correction steps will have kNOutputs
+  // dimensions.
+  // TODO(james): Relax this assumption; relaxing it requires
+  // figuring out how to deal with storing variable size
+  // observation matrices, though.
+  static constexpr int kNOutputs = 3;
+  // Inputs are [left_volts, right_volts]
+  typedef Eigen::Matrix<Scalar, kNInputs, 1> Input;
+  // Outputs are either:
+  // [left_encoder, right_encoder, gyro_vel]; or [heading, distance, skew] to
+  // some target. This makes it so we don't have to figure out how we store
+  // variable-size measurement updates.
+  typedef Eigen::Matrix<Scalar, kNOutputs, 1> Output;
+  typedef Eigen::Matrix<Scalar, kNStates, kNStates> StateSquare;
+  // State is [x_position, y_position, theta, Kalman States], where
+  // Kalman States are the states from the standard drivetrain Kalman Filter,
+  // which is: [left encoder, left ground vel, right encoder, right ground vel,
+  // left voltage error, right voltage error, angular_error], where:
+  // left/right encoder should correspond directly to encoder readings
+  // left/right velocities are the velocity of the left/right sides over the
+  //   ground (i.e., corrected for angular_error).
+  // voltage errors are the difference between commanded and effective voltage,
+  //   used to estimate consistent modelling errors (e.g., friction).
+  // angular error is the difference between the angular velocity as estimated
+  //   by the encoders vs. estimated by the gyro, such as might be caused by
+  //   wheels on one side of the drivetrain being too small or one side's
+  //   wheels slipping more than the other.
+  typedef Eigen::Matrix<Scalar, kNStates, 1> State;
+
+  // Constructs a HybridEkf for a particular drivetrain.
+  // Currently, we use the drivetrain config for modelling constants
+  // (continuous time A and B matrices) and for the noise matrices for the
+  // encoders/gyro.
+  HybridEkf(const DrivetrainConfig<Scalar> &dt_config)
+      : dt_config_(dt_config),
+        velocity_drivetrain_coefficients_(
+            dt_config.make_hybrid_drivetrain_velocity_loop()
+                .plant()
+                .coefficients()) {
+    InitializeMatrices();
+  }
+
+  // Set the initial guess of the state. Can only be called once, and before
+  // any measurement updates have occured.
+  // TODO(james): We may want to actually re-initialize and reset things on
+  // the field. Create some sort of Reset() function.
+  void ResetInitialState(::aos::monotonic_clock::time_point t,
+                       const State &state) {
+    observations_.clear();
+    X_hat_ = state;
+    observations_.PushFromBottom(
+        {t,
+         t,
+         X_hat_,
+         P_,
+         Input::Zero(),
+         Output::Zero(),
+         {},
+         [](const State &, const Input &) { return Output::Zero(); },
+         [](const State &) {
+           return Eigen::Matrix<Scalar, kNOutputs, kNStates>::Zero();
+         },
+         Eigen::Matrix<Scalar, kNOutputs, kNOutputs>::Identity()});
+  }
+
+  // Correct with:
+  // A measurement z at time t with z = h(X_hat, U) + v where v has noise
+  // covariance R.
+  // Input U is applied from the previous timestep until time t.
+  // If t is later than any previous measurements, then U must be provided.
+  // If the measurement falls between two previous measurements, then U
+  // can be provided or not; if U is not provided, then it is filled in based
+  // on an assumption that the voltage was held constant between the time steps.
+  // TODO(james): Is it necessary to explicitly to provide a version with H as a
+  // matrix for linear cases?
+  void Correct(
+      const Output &z, const Input *U,
+      ::std::function<
+          void(const State &, const StateSquare &,
+               ::std::function<Output(const State &, const Input &)> *,
+               ::std::function<Eigen::Matrix<Scalar, kNOutputs, kNStates>(
+                   const State &)> *)> make_h,
+      ::std::function<Output(const State &, const Input &)> h,
+      ::std::function<Eigen::Matrix<Scalar, kNOutputs, kNStates>(const State &)>
+          dhdx, const Eigen::Matrix<Scalar, kNOutputs, kNOutputs> &R,
+      aos::monotonic_clock::time_point t);
+
+  // A utility function for specifically updating with encoder and gyro
+  // measurements.
+  void UpdateEncodersAndGyro(const Scalar left_encoder,
+                             const Scalar right_encoder, const Scalar gyro_rate,
+                             const Input &U,
+                             ::aos::monotonic_clock::time_point t) {
+    Output z(left_encoder, right_encoder, gyro_rate);
+    Eigen::Matrix<Scalar, kNOutputs, kNOutputs> R;
+    R.setZero();
+    R.diagonal() << encoder_noise_, encoder_noise_, gyro_noise_;
+    Correct(z, &U, {}, [this](const State &X, const Input &) {
+                         return H_encoders_and_gyro_ * X;
+                       },
+            [this](const State &) { return H_encoders_and_gyro_; }, R, t);
+  }
+
+  // Sundry accessor:
+  State X_hat() const { return X_hat_; }
+  Scalar X_hat(long i) const { return X_hat_(i, 0); }
+  StateSquare P() const { return P_; }
+  ::aos::monotonic_clock::time_point latest_t() const {
+    return observations_.top().t;
+  }
+
+ private:
+  struct Observation {
+    // Time when the observation was taken.
+    aos::monotonic_clock::time_point t;
+    // Time that the previous observation was taken:
+    aos::monotonic_clock::time_point prev_t;
+    // Estimate of state at previous observation time t, after accounting for
+    // the previous observation.
+    State X_hat;
+    // Noise matrix corresponding to X_hat_.
+    StateSquare P;
+    // The input applied from previous observation until time t.
+    Input U;
+    // Measurement taken at that time.
+    Output z;
+    // A function to create h and dhdx from a given position/covariance
+    // estimate. This is used by the camera to make it so that we only have to
+    // match targets once.
+    // Only called if h and dhdx are empty.
+    ::std::function<
+        void(const State &, const StateSquare &,
+             ::std::function<Output(const State &, const Input &)> *,
+             ::std::function<Eigen::Matrix<Scalar, kNOutputs, kNStates>(
+                 const State &)> *)> make_h;
+    // A function to calculate the expected output at a given state/input.
+    // TODO(james): For encoders/gyro, it is linear and the function call may
+    // be expensive. Potential source of optimization.
+    ::std::function<Output(const State &, const Input &)> h;
+    // The Jacobian of h with respect to x.
+    // We assume that U has no impact on the Jacobian.
+    // TODO(james): Currently, none of the users of this actually make use of
+    // the ability to have dynamic dhdx (technically, the camera code should
+    // recalculate it to be strictly correct, but I was both too lazy to do
+    // so and it seemed unnecessary). This is a potential source for future
+    // optimizations if function calls are being expensive.
+    ::std::function<
+        Eigen::Matrix<Scalar, kNOutputs, kNStates>(const State &)> dhdx;
+    // The measurement noise matrix.
+    Eigen::Matrix<Scalar, kNOutputs, kNOutputs> R;
+
+    // In order to sort the observations in the PriorityQueue object, we
+    // need a comparison function.
+    friend bool operator <(const Observation &l, const Observation &r) {
+      return l.t < r.t;
+    }
+  };
+
+  void InitializeMatrices();
+
+  StateSquare AForState(const State &X) const {
+    StateSquare A_continuous = A_continuous_;
+    const Scalar theta = X(kTheta, 0);
+    const Scalar linear_vel =
+        (X(kLeftVelocity, 0) + X(kRightVelocity, 0)) / 2.0;
+    const Scalar stheta = ::std::sin(theta);
+    const Scalar ctheta = ::std::cos(theta);
+    // X and Y derivatives
+    A_continuous(kX, kTheta) = -stheta * linear_vel;
+    A_continuous(kX, kLeftVelocity) = ctheta / 2.0;
+    A_continuous(kX, kRightVelocity) = ctheta / 2.0;
+    A_continuous(kY, kTheta) = ctheta * linear_vel;
+    A_continuous(kY, kLeftVelocity) = stheta / 2.0;
+    A_continuous(kY, kRightVelocity) = stheta / 2.0;
+    return A_continuous;
+  }
+
+  State DiffEq(const State &X, const Input &U) const {
+    State Xdot = A_continuous_ * X + B_continuous_ * U;
+    // And then we need to add on the terms for the x/y change:
+    const Scalar theta = X(kTheta, 0);
+    const Scalar linear_vel =
+        (X(kLeftVelocity, 0) + X(kRightVelocity, 0)) / 2.0;
+    const Scalar stheta = ::std::sin(theta);
+    const Scalar ctheta = ::std::cos(theta);
+    Xdot(kX, 0) = ctheta * linear_vel;
+    Xdot(kY, 0) = stheta * linear_vel;
+    return Xdot;
+  }
+
+  void PredictImpl(const Input &U, std::chrono::nanoseconds dt, State *state,
+                   StateSquare *P) {
+    StateSquare A_c = AForState(*state);
+    StateSquare A_d;
+    Eigen::Matrix<Scalar, kNStates, 0> dummy_B;
+    controls::C2D<Scalar, kNStates, 0>(
+        A_c, Eigen::Matrix<Scalar, kNStates, 0>::Zero(), dt, &A_d, &dummy_B);
+    StateSquare Q_d;
+    controls::DiscretizeQ(Q_continuous_, A_c, dt, &Q_d);
+    *P = A_d * *P * A_d.transpose() + Q_d;
+
+    *state = RungeKuttaU(
+        [this](const State &X,
+               const Input &U) { return DiffEq(X, U); },
+        *state, U,
+        ::std::chrono::duration_cast<::std::chrono::duration<double>>(dt)
+            .count());
+  }
+
+  void CorrectImpl(const Eigen::Matrix<Scalar, kNOutputs, kNOutputs> &R,
+                   const Output &Z, const Output &expected_Z,
+                   const Eigen::Matrix<Scalar, kNOutputs, kNStates> &H,
+                   State *state, StateSquare *P) {
+    Output err = Z - expected_Z;
+    Eigen::Matrix<Scalar, kNStates, kNOutputs> PH = *P * H.transpose();
+    Eigen::Matrix<Scalar, kNOutputs, kNOutputs> S = H * PH + R;
+    Eigen::Matrix<Scalar, kNStates, kNOutputs> K = PH * S.inverse();
+    *state = *state + K * err;
+    *P = (StateSquare::Identity() - K * H) * *P;
+  }
+
+  void ProcessObservation(Observation *obs, const std::chrono::nanoseconds dt,
+                          State *state, StateSquare *P) {
+    *state = obs->X_hat;
+    *P = obs->P;
+    if (dt.count() != 0) {
+      PredictImpl(obs->U, dt, state, P);
+    }
+    if (!(obs->h && obs->dhdx)) {
+      CHECK(obs->make_h);
+      obs->make_h(*state, *P, &obs->h, &obs->dhdx);
+    }
+    CorrectImpl(obs->R, obs->z, obs->h(*state, obs->U), obs->dhdx(*state),
+                state, P);
+  }
+
+  DrivetrainConfig<Scalar> dt_config_;
+  State X_hat_;
+  StateFeedbackHybridPlantCoefficients<2, 2, 2, Scalar>
+      velocity_drivetrain_coefficients_;
+  StateSquare A_continuous_;
+  StateSquare Q_continuous_;
+  StateSquare P_;
+  Eigen::Matrix<Scalar, kNOutputs, kNStates> H_encoders_and_gyro_;
+  Scalar encoder_noise_, gyro_noise_;
+  Eigen::Matrix<Scalar, kNStates, kNInputs> B_continuous_;
+
+  aos::PriorityQueue<Observation, kSaveSamples, ::std::less<Observation>>
+      observations_;
+
+  friend class testing::HybridEkfTest;
+};  // class HybridEkf
+
+template <typename Scalar>
+void HybridEkf<Scalar>::Correct(
+    const Output &z, const Input *U,
+    ::std::function<
+        void(const State &, const StateSquare &,
+             ::std::function<Output(const State &, const Input &)> *,
+             ::std::function<Eigen::Matrix<Scalar, kNOutputs, kNStates>(
+                 const State &)> *)> make_h,
+    ::std::function<Output(const State &, const Input &)> h,
+    ::std::function<Eigen::Matrix<Scalar, kNOutputs, kNStates>(const State &)>
+        dhdx, const Eigen::Matrix<Scalar, kNOutputs, kNOutputs> &R,
+    aos::monotonic_clock::time_point t) {
+  CHECK(!observations_.empty());
+  if (!observations_.full() && t < observations_.begin()->t) {
+    LOG(ERROR,
+        "Dropped an observation that was received before we "
+        "initialized.\n");
+    return;
+  }
+  auto cur_it =
+      observations_.PushFromBottom({t, t, State::Zero(), StateSquare::Zero(),
+                                    Input::Zero(), z, make_h, h, dhdx, R});
+  if (cur_it == observations_.end()) {
+    LOG(DEBUG,
+        "Camera dropped off of end with time of %fs; earliest observation in "
+        "queue has time of %fs.\n",
+        ::std::chrono::duration_cast<::std::chrono::duration<double>>(
+            t.time_since_epoch()).count(),
+        ::std::chrono::duration_cast<::std::chrono::duration<double>>(
+            observations_.begin()->t.time_since_epoch()).count());
+    return;
+  }
+
+  // Now we populate any state information that depends on where the
+  // observation was inserted into the queue. X_hat and P must be populated
+  // from the values present in the observation *following* this one in
+  // the queue (note that the X_hat and P that we store in each observation
+  // is the values that they held after accounting for the previous
+  // measurement and before accounting for the time between the previous and
+  // current measurement). If we appended to the end of the queue, then
+  // we need to pull from X_hat_ and P_ specifically.
+  // Furthermore, for U:
+  // -If the observation was inserted at the end, then the user must've
+  //  provided U and we use it.
+  // -Otherwise, only grab U if necessary.
+  auto next_it = cur_it;
+  ++next_it;
+  if (next_it == observations_.end()) {
+    cur_it->X_hat = X_hat_;
+    cur_it->P = P_;
+    // Note that if next_it == observations_.end(), then because we already
+    // checked for !observations_.empty(), we are guaranteed to have
+    // valid prev_it.
+    auto prev_it = cur_it;
+    --prev_it;
+    cur_it->prev_t = prev_it->t;
+    // TODO(james): Figure out a saner way of handling this.
+    CHECK(U != nullptr);
+    cur_it->U = *U;
+  } else {
+    cur_it->X_hat = next_it->X_hat;
+    cur_it->P = next_it->P;
+    cur_it->prev_t = next_it->prev_t;
+    next_it->prev_t = cur_it->t;
+    cur_it->U = (U == nullptr) ? next_it->U : *U;
+  }
+  // Now we need to rerun the predict step from the previous to the new
+  // observation as well as every following correct/predict up to the current
+  // time.
+  while (true) {
+    // We use X_hat_ and P_ to store the intermediate states, and then
+    // once we reach the end they will all be up-to-date.
+    ProcessObservation(&*cur_it, cur_it->t - cur_it->prev_t, &X_hat_, &P_);
+    CHECK(X_hat_.allFinite());
+    if (next_it != observations_.end()) {
+      next_it->X_hat = X_hat_;
+      next_it->P = P_;
+    } else {
+      break;
+    }
+    ++cur_it;
+    ++next_it;
+  }
+}
+
+template <typename Scalar>
+void HybridEkf<Scalar>::InitializeMatrices() {
+  A_continuous_.setZero();
+  const Scalar diameter = 2.0 * dt_config_.robot_radius;
+  // Theta derivative
+  A_continuous_(kTheta, kLeftVelocity) = -1.0 / diameter;
+  A_continuous_(kTheta, kRightVelocity) = 1.0 / diameter;
+
+  // Encoder derivatives
+  A_continuous_(kLeftEncoder, kLeftVelocity) = 1.0;
+  A_continuous_(kLeftEncoder, kAngularError) = 1.0;
+  A_continuous_(kRightEncoder, kRightVelocity) = 1.0;
+  A_continuous_(kRightEncoder, kAngularError) = -1.0;
+
+  // Pull velocity derivatives from velocity matrices.
+  // Note that this looks really awkward (doesn't use
+  // Eigen blocks) because someone decided that the full
+  // drivetrain Kalman Filter should half a weird convention.
+  // TODO(james): Support shifting drivetrains with changing A_continuous
+  const auto &vel_coefs = velocity_drivetrain_coefficients_;
+  A_continuous_(kLeftVelocity, kLeftVelocity) = vel_coefs.A_continuous(0, 0);
+  A_continuous_(kLeftVelocity, kRightVelocity) = vel_coefs.A_continuous(0, 1);
+  A_continuous_(kRightVelocity, kLeftVelocity) = vel_coefs.A_continuous(1, 0);
+  A_continuous_(kRightVelocity, kRightVelocity) = vel_coefs.A_continuous(1, 1);
+
+  // Provide for voltage error terms:
+  B_continuous_.setZero();
+  B_continuous_.row(kLeftVelocity) = vel_coefs.B_continuous.row(0);
+  B_continuous_.row(kRightVelocity) = vel_coefs.B_continuous.row(1);
+  A_continuous_.template block<kNStates, kNInputs>(0, 7) = B_continuous_;
+
+  Q_continuous_.setZero();
+  // TODO(james): Improve estimates of process noise--e.g., X/Y noise can
+  // probably be reduced because we rarely randomly jump in space, but maybe
+  // wheel speed noise should increase due to likelihood of slippage.
+  Q_continuous_(kX, kX) = 0.005;
+  Q_continuous_(kY, kY) = 0.005;
+  Q_continuous_(kTheta, kTheta) = 0.001;
+  Q_continuous_.template block<7, 7>(3, 3) =
+      dt_config_.make_kf_drivetrain_loop().observer().coefficients().Q;
+
+  P_.setZero();
+  P_.diagonal() << 1, 1, 1, 0.5, 0.1, 0.5, 0.1, 1, 1, 1;
+
+  H_encoders_and_gyro_.setZero();
+  // Encoders are stored directly in the state matrix, so are a minor
+  // transform away.
+  H_encoders_and_gyro_(0, kLeftEncoder) = 1.0;
+  H_encoders_and_gyro_(1, kRightEncoder) = 1.0;
+  // Gyro rate is just the difference between right/left side speeds:
+  H_encoders_and_gyro_(2, kLeftVelocity) = -1.0 / diameter;
+  H_encoders_and_gyro_(2, kRightVelocity) = 1.0 / diameter;
+
+  const Eigen::Matrix<Scalar, 4, 4> R_kf_drivetrain =
+      dt_config_.make_kf_drivetrain_loop().observer().coefficients().R;
+  encoder_noise_ = R_kf_drivetrain(0, 0);
+  gyro_noise_ = R_kf_drivetrain(2, 2);
+}
+
+}  // namespace drivetrain
+}  // namespace control_loops
+}  // namespace frc971
+
+#endif  // FRC971_CONTROL_LOOPS_DRIVETRAIN_HYBRID_EKF_H_
diff --git a/frc971/control_loops/drivetrain/hybrid_ekf_test.cc b/frc971/control_loops/drivetrain/hybrid_ekf_test.cc
new file mode 100644
index 0000000..9bdb30b
--- /dev/null
+++ b/frc971/control_loops/drivetrain/hybrid_ekf_test.cc
@@ -0,0 +1,427 @@
+#include "frc971/control_loops/drivetrain/hybrid_ekf.h"
+
+#include <random>
+
+#include "aos/testing/random_seed.h"
+#include "aos/testing/test_shm.h"
+#include "frc971/control_loops/drivetrain/trajectory.h"
+#include "frc971/control_loops/drivetrain/drivetrain_test_lib.h"
+#include "gtest/gtest.h"
+
+namespace frc971 {
+namespace control_loops {
+namespace drivetrain {
+namespace testing {
+
+typedef HybridEkf<>::StateIdx StateIdx;
+
+
+class HybridEkfTest : public ::testing::Test {
+ public:
+  typedef HybridEkf<>::State State;
+  typedef HybridEkf<>::Input Input;
+  ::aos::testing::TestSharedMemory shm_;
+  HybridEkfTest()
+      : dt_config_(GetTestDrivetrainConfig()),
+        ekf_(dt_config_),
+        t0_(::std::chrono::seconds(0)),
+        velocity_plant_coefs_(dt_config_.make_hybrid_drivetrain_velocity_loop()
+                                  .plant()
+                                  .coefficients()) {
+    ekf_.ResetInitialState(t0_, HybridEkf<>::State::Zero());
+  }
+
+ protected:
+  const State Update(const State &X, const Input &U) {
+    return RungeKuttaU(
+        ::std::bind(&HybridEkfTest::DiffEq, this, ::std::placeholders::_1,
+                    ::std::placeholders::_2),
+        X, U, ::std::chrono::duration_cast<::std::chrono::duration<double>>(
+                  dt_config_.dt).count());
+  }
+  void CheckDiffEq(const State &X, const Input &U) {
+    // Re-implement dynamics as a sanity check:
+    const double diameter = 2.0 * dt_config_.robot_radius;
+    const double theta = X(StateIdx::kTheta, 0);
+    const double stheta = ::std::sin(theta);
+    const double ctheta = ::std::cos(theta);
+    const double left_vel = X(StateIdx::kLeftVelocity, 0);
+    const double right_vel = X(StateIdx::kRightVelocity, 0);
+    const State Xdot_ekf = DiffEq(X, U);
+    EXPECT_EQ(Xdot_ekf(StateIdx::kX, 0), ctheta * (left_vel + right_vel) / 2.0);
+    EXPECT_EQ(Xdot_ekf(StateIdx::kY, 0), stheta * (left_vel + right_vel) / 2.0);
+    EXPECT_EQ(Xdot_ekf(StateIdx::kTheta, 0), (right_vel - left_vel) / diameter);
+    EXPECT_EQ(Xdot_ekf(StateIdx::kLeftEncoder, 0),
+              left_vel + X(StateIdx::kAngularError, 0));
+    EXPECT_EQ(Xdot_ekf(StateIdx::kRightEncoder, 0),
+              right_vel - X(StateIdx::kAngularError, 0));
+
+    Eigen::Matrix<double, 2, 1> vel_x(X(StateIdx::kLeftVelocity, 0),
+                                      X(StateIdx::kRightVelocity, 0));
+    Eigen::Matrix<double, 2, 1> expected_vel_X =
+        velocity_plant_coefs_.A_continuous * vel_x +
+        velocity_plant_coefs_.B_continuous *
+            (U + X.middleRows<2>(StateIdx::kLeftVoltageError));
+    EXPECT_EQ(Xdot_ekf(StateIdx::kLeftVelocity, 0), expected_vel_X(0, 0));
+    EXPECT_EQ(Xdot_ekf(StateIdx::kRightVelocity, 0), expected_vel_X(1, 0));
+
+    // Dynamics don't expect error terms to change:
+    EXPECT_EQ(0.0, Xdot_ekf.bottomRows<3>().squaredNorm());
+  }
+  State DiffEq(const State &X, const Input &U) {
+    return ekf_.DiffEq(X, U);
+  }
+
+  // Returns a random value sampled from a normal distribution with a standard
+  // deviation of std and a mean of zero.
+  double Normal(double std) {
+    return normal_(gen_) * std;
+  }
+
+  DrivetrainConfig<double> dt_config_;
+  HybridEkf<> ekf_;
+  ::aos::monotonic_clock::time_point t0_;
+  StateFeedbackHybridPlantCoefficients<2, 2, 2> velocity_plant_coefs_;
+
+  ::std::mt19937 gen_{static_cast<uint32_t>(::aos::testing::RandomSeed())};
+  ::std::normal_distribution<> normal_;
+};
+
+// Tests that the dynamics from DiffEq behave as expected:
+TEST_F(HybridEkfTest, CheckDynamics) {
+  CheckDiffEq(State::Zero(), Input::Zero());
+  CheckDiffEq(State::Zero(), {-5.0, 5.0});
+  CheckDiffEq(State::Zero(), {12.0, -3.0});
+  CheckDiffEq((State() << 100.0, 200.0, M_PI, 1.234, 0.5, 1.2, 0.6, 3.0, -4.0,
+               0.3).finished(),
+              {5.0, 6.0});
+  CheckDiffEq((State() << 100.0, 200.0, 2.0, 1.234, 0.5, 1.2, 0.6, 3.0, -4.0,
+               0.3).finished(),
+              {5.0, 6.0});
+  CheckDiffEq((State() << 100.0, 200.0, -2.0, 1.234, 0.5, 1.2, 0.6, 3.0, -4.0,
+               0.3).finished(),
+              {5.0, 6.0});
+  // And check that a theta outisde of [-M_PI, M_PI] works.
+  CheckDiffEq((State() << 100.0, 200.0, 200.0, 1.234, 0.5, 1.2, 0.6, 3.0, -4.0,
+               0.3).finished(),
+              {5.0, 6.0});
+}
+
+// Tests that if we provide a bunch of observations of the angular position
+// with zero change in time, the state should approach the estimation.
+TEST_F(HybridEkfTest, ZeroTimeCorrect) {
+  HybridEkf<>::Output Z(1, 0, 0);
+  Eigen::Matrix<double, 3, 10> H;
+  H.setZero();
+  H(0, 2) = 1;
+  auto h = [H](const State &X, const Input &) { return H * X; };
+  auto dhdx = [H](const State &) { return H; };
+  Eigen::Matrix<double, 3, 3> R;
+  R.setZero();
+  R.diagonal() << 0.01, 0.01, 0.01;
+  Input U(0, 0);
+  EXPECT_EQ(0.0, ekf_.X_hat(StateIdx::kTheta));
+  const double starting_p_norm = ekf_.P().norm();
+  for (int ii = 0; ii < 100; ++ii) {
+    ekf_.Correct(Z, &U, {}, h, dhdx, R, t0_);
+  }
+  EXPECT_NEAR(Z(0, 0), ekf_.X_hat(StateIdx::kTheta), 1e-3);
+  const double ending_p_norm = ekf_.P().norm();
+  // Due to corrections, noise should've decreased.
+  EXPECT_LT(ending_p_norm, starting_p_norm * 0.95);
+}
+
+// Tests that prediction steps alone, with no corrections, results in sane
+// results. In order to implement the "no corrections" part of that, we just set
+// H to zero.
+TEST_F(HybridEkfTest, PredictionsAreSane) {
+  HybridEkf<>::Output Z(0, 0, 0);
+  // Use true_X to track what we think the true robot state is.
+  State true_X = ekf_.X_hat();
+  Eigen::Matrix<double, 3, 10> H;
+  H.setZero();
+  auto h = [H](const State &X, const Input &) { return H * X; };
+  auto dhdx = [H](const State &) { return H; };
+  // Provide constant input voltage.
+  Eigen::Matrix<double, 2, 1> U;
+  U << 12.0, 10.0;
+  // The exact value of the noise matrix ix unimportant so long as it is
+  // non-zero.
+  Eigen::Matrix<double, 3, 3> R;
+  R.setIdentity();
+  EXPECT_EQ(0.0, ekf_.X_hat().norm());
+  const double starting_p_norm = ekf_.P().norm();
+  for (int ii = 0; ii < 100; ++ii) {
+    ekf_.Correct(Z, &U, {}, h, dhdx, R, t0_ + dt_config_.dt * (ii + 1));
+    true_X = Update(true_X, U);
+    EXPECT_EQ(true_X, ekf_.X_hat());
+  }
+  // We don't care about precise results, just that they are generally sane:
+  // robot should've travelled forwards and slightly to the right.
+  EXPECT_NEAR(0.9, ekf_.X_hat(StateIdx::kX), 0.1);
+  EXPECT_NEAR(-0.15, ekf_.X_hat(StateIdx::kY), 0.05);
+  EXPECT_NEAR(-0.3, ekf_.X_hat(StateIdx::kTheta), 0.05);
+  EXPECT_NEAR(1.0, ekf_.X_hat(StateIdx::kLeftEncoder), 0.1);
+  EXPECT_NEAR(ekf_.X_hat(StateIdx::kLeftEncoder) * 0.8,
+              ekf_.X_hat(StateIdx::kRightEncoder),
+              ekf_.X_hat(StateIdx::kLeftEncoder) * 0.1);
+  EXPECT_NEAR(2.4, ekf_.X_hat(StateIdx::kLeftVelocity), 0.1);
+  EXPECT_NEAR(ekf_.X_hat(StateIdx::kLeftVelocity) * 0.8,
+              ekf_.X_hat(StateIdx::kRightVelocity),
+              ekf_.X_hat(StateIdx::kLeftVelocity) * 0.1);
+  EXPECT_EQ(0.0, ekf_.X_hat(StateIdx::kLeftVoltageError));
+  EXPECT_EQ(0.0, ekf_.X_hat(StateIdx::kRightVoltageError));
+  EXPECT_EQ(0.0, ekf_.X_hat(StateIdx::kAngularError));
+  const double ending_p_norm = ekf_.P().norm();
+  // Due to lack of corrections, noise should've increased.
+  EXPECT_GT(ending_p_norm, starting_p_norm * 1.10);
+}
+
+// Parameterize HybridEkf off of the number of prediction-only updates to
+// provide before doing measurements. This is so that we make sure to exercise
+// the corner case when we apply measurements that appear at precisely the
+// oldest time-step (during development, this corner case caused some issues).
+class HybridEkfOldCorrectionsTest
+    : public HybridEkfTest,
+      public ::testing::WithParamInterface<size_t> {};
+
+// Tests that creating an old measurement works, by basically running the
+// previous two tests in reverse (running a series of predictions, and then
+// inserting observation(s) at the start to change everything).
+TEST_P(HybridEkfOldCorrectionsTest, CreateOldCorrection) {
+  HybridEkf<>::Output Z;
+  Z.setZero();
+  Eigen::Matrix<double, 3, 10> H;
+  H.setZero();
+  auto h_zero = [H](const State &X, const Input &) { return H * X; };
+  auto dhdx_zero = [H](const State &) { return H; };
+  Eigen::Matrix<double, 2, 1> U;
+  U << 12.0, 12.0;
+  Eigen::Matrix<double, 3, 3> R;
+  R = R.Identity();
+  EXPECT_EQ(0.0, ekf_.X_hat().norm());
+  // We fill up the buffer to be as full as demanded by the user.
+  const size_t n_predictions = GetParam();
+  for (size_t ii = 0; ii < n_predictions; ++ii) {
+    ekf_.Correct(Z, &U, {}, h_zero, dhdx_zero, R,
+                 t0_ + dt_config_.dt * (ii + 1));
+  }
+
+  // Store state and covariance after prediction steps.
+  const State modeled_X_hat = ekf_.X_hat();
+  const double modeled_p_norm = ekf_.P().norm();
+
+  Z << 1, 1, M_PI / 2.0;
+  H.setZero();
+  H(0, 0) = 1;
+  H(1, 1) = 1;
+  H(2, 2) = 1;
+  auto h = [H](const State &X, const Input &) { return H * X; };
+  auto dhdx = [H](const State &) { return H; };
+  R.setZero();
+  R.diagonal() << 1e-5, 1e-5, 1e-5;
+  U.setZero();
+  for (int ii = 0; ii < 20; ++ii) {
+    ekf_.Correct(Z, &U, {}, h, dhdx, R, t0_);
+  }
+  const double corrected_p_norm = ekf_.P().norm();
+  State expected_X_hat = modeled_X_hat;
+  expected_X_hat(0, 0) = Z(0, 0);
+  expected_X_hat(1, 0) = Z(1, 0) + modeled_X_hat(0, 0);
+  expected_X_hat(2, 0) = Z(2, 0);
+  EXPECT_LT((expected_X_hat.topRows<7>() - ekf_.X_hat().topRows<7>()).norm(),
+           1e-3)
+      << "X_hat: " << ekf_.X_hat() << " expected " << expected_X_hat;
+  // The covariance after the predictions but before the corrections should
+  // be higher than after the corrections are made.
+  EXPECT_GT(modeled_p_norm, corrected_p_norm);
+}
+
+// Ensure that we check kSaveSamples - 1, for potential corner cases.
+INSTANTIATE_TEST_CASE_P(OldCorrectionTest, HybridEkfOldCorrectionsTest,
+                        ::testing::Values(0, 1, 10,
+                                          HybridEkf<>::kSaveSamples - 1));
+
+// Tests that creating a correction that is too old results in the correction
+// being dropped and ignored.
+TEST_F(HybridEkfTest, DiscardTooOldCorrection) {
+  HybridEkf<>::Output Z;
+  Z.setZero();
+  Eigen::Matrix<double, 3, 10> H;
+  H.setZero();
+  auto h_zero = [H](const State &X, const Input &) { return H * X; };
+  auto dhdx_zero = [H](const State &) { return H; };
+  Eigen::Matrix<double, 2, 1> U(12.0, 12.0);
+  Eigen::Matrix<double, 3, 3> R;
+  R.setIdentity();
+
+  EXPECT_EQ(0.0, ekf_.X_hat().norm());
+  for (int ii = 0; ii < HybridEkf<>::kSaveSamples; ++ii) {
+    ekf_.Correct(Z, &U, {}, h_zero, dhdx_zero, R,
+                 t0_ + dt_config_.dt * (ii + 1));
+  }
+  const State modeled_X_hat = ekf_.X_hat();
+  const HybridEkf<>::StateSquare modeled_P = ekf_.P();
+
+  Z << 1, 1, M_PI / 2.0;
+  H.setZero();
+  H(0, 0) = 1;
+  H(1, 1) = 1;
+  H(2, 2) = 1;
+  auto h = [H](const State &X, const Input &) { return H * X; };
+  auto dhdx = [H](const State &) { return H; };
+  R.setIdentity();
+  R *= 1e-5;
+  U.setZero();
+  ekf_.Correct(Z, &U, {}, h, dhdx, R, t0_);
+  EXPECT_EQ(ekf_.X_hat(), modeled_X_hat)
+      << "Expected too-old correction to have no effect; X_hat: "
+      << ekf_.X_hat() << " expected " << modeled_X_hat;
+  EXPECT_EQ(ekf_.P(), modeled_P)
+      << "Expected too-old correction to have no effect; P: " << ekf_.P()
+      << " expected " << modeled_P;
+}
+
+// Tests The UpdateEncodersAndGyro function works when perfect measurements are
+// provided:
+TEST_F(HybridEkfTest, PerfectEncoderUpdate) {
+  State true_X = ekf_.X_hat();
+  Input U(-1.0, 5.0);
+  for (int ii = 0; ii < 100; ++ii) {
+    true_X = Update(true_X, U);
+    ekf_.UpdateEncodersAndGyro(true_X(StateIdx::kLeftEncoder, 0),
+                               true_X(StateIdx::kRightEncoder, 0),
+                               (true_X(StateIdx::kRightVelocity, 0) -
+                                true_X(StateIdx::kLeftVelocity, 0)) /
+                                   dt_config_.robot_radius / 2.0,
+                               U, t0_ + (ii + 1) * dt_config_.dt);
+    EXPECT_NEAR((true_X - ekf_.X_hat()).squaredNorm(), 0.0, 1e-25)
+        << "Expected only floating point precision errors in update step. "
+           "Estimated X_hat:\n" << ekf_.X_hat() << "\ntrue X:\n" << true_X;
+  }
+}
+
+// Tests that encoder updates cause everything to converge properly in the
+// presence of voltage error.
+TEST_F(HybridEkfTest, PerfectEncoderUpdatesWithVoltageError) {
+  State true_X = ekf_.X_hat();
+  true_X(StateIdx::kLeftVoltageError, 0) = 2.0;
+  true_X(StateIdx::kRightVoltageError, 0) = 2.0;
+  Input U(5.0, 5.0);
+  for (int ii = 0; ii < 1000; ++ii) {
+    true_X = Update(true_X, U);
+    ekf_.UpdateEncodersAndGyro(true_X(StateIdx::kLeftEncoder, 0),
+                               true_X(StateIdx::kRightEncoder, 0),
+                               (true_X(StateIdx::kRightVelocity, 0) -
+                                true_X(StateIdx::kLeftVelocity, 0)) /
+                                   dt_config_.robot_radius / 2.0,
+                               U, t0_ + (ii + 1) * dt_config_.dt);
+  }
+  EXPECT_NEAR((true_X - ekf_.X_hat()).norm(), 0.0, 5e-3)
+      << "Expected non-x/y estimates to converge to correct. "
+         "Estimated X_hat:\n"
+      << ekf_.X_hat() << "\ntrue X:\n"
+      << true_X;
+}
+
+// Tests encoder/gyro updates when we have some errors in our estimate.
+TEST_F(HybridEkfTest, PerfectEncoderUpdateConverges) {
+  // In order to simulate modelling errors, we add an angular_error and start
+  // the encoder values slightly off.
+  State true_X = ekf_.X_hat();
+  true_X(StateIdx::kAngularError, 0) = 1.0;
+  true_X(StateIdx::kLeftEncoder, 0) += 2.0;
+  true_X(StateIdx::kRightEncoder, 0) -= 2.0;
+  // After enough time, everything should converge to near-perfect (if there
+  // were any errors in the original absolute state (x/y/theta) state, then we
+  // can't correct for those).
+  // Note: Because we don't have any absolute measurements used for corrections,
+  // we will get slightly off on the absolute x/y/theta, but the errors are so
+  // small they are negligible.
+  Input U(10.0, 5.0);
+  for (int ii = 0; ii < 100; ++ii) {
+    true_X = Update(true_X, U);
+    ekf_.UpdateEncodersAndGyro(true_X(StateIdx::kLeftEncoder, 0),
+                               true_X(StateIdx::kRightEncoder, 0),
+                               (true_X(StateIdx::kRightVelocity, 0) -
+                                true_X(StateIdx::kLeftVelocity, 0)) /
+                                   dt_config_.robot_radius / 2.0,
+                               U, t0_ + (ii + 1) * dt_config_.dt);
+  }
+  EXPECT_NEAR((true_X - ekf_.X_hat()).norm(), 0.0, 1e-5)
+      << "Expected non-x/y estimates to converge to correct. "
+         "Estimated X_hat:\n"
+      << ekf_.X_hat() << "\ntrue X:\n"
+      << true_X;
+}
+
+// Tests encoder/gyro updates in a realistic-ish scenario with noise:
+TEST_F(HybridEkfTest, RealisticEncoderUpdateConverges) {
+  // In order to simulate modelling errors, we add an angular_error and start
+  // the encoder values slightly off.
+  State true_X = ekf_.X_hat();
+  true_X(StateIdx::kAngularError, 0) = 1.0;
+  true_X(StateIdx::kLeftEncoder, 0) += 2.0;
+  true_X(StateIdx::kRightEncoder, 0) -= 2.0;
+  Input U(10.0, 5.0);
+  for (int ii = 0; ii < 100; ++ii) {
+    true_X = Update(true_X, U);
+    ekf_.UpdateEncodersAndGyro(
+        true_X(StateIdx::kLeftEncoder, 0) + Normal(1e-3),
+        true_X(StateIdx::kRightEncoder, 0) + Normal(1e-3),
+        (true_X(StateIdx::kRightVelocity, 0) -
+         true_X(StateIdx::kLeftVelocity, 0)) /
+            dt_config_.robot_radius / 2.0 + Normal(1e-4),
+        U, t0_ + (ii + 1) * dt_config_.dt);
+  }
+  EXPECT_NEAR(
+      (true_X.bottomRows<9>() - ekf_.X_hat().bottomRows<9>()).squaredNorm(),
+      0.0, 2e-3)
+      << "Expected non-x/y estimates to converge to correct. "
+         "Estimated X_hat:\n" << ekf_.X_hat() << "\ntrue X:\n" << true_X;
+}
+
+class HybridEkfDeathTest : public HybridEkfTest {};
+
+TEST_F(HybridEkfDeathTest, DieIfUninitialized) {
+  HybridEkf<> ekf(dt_config_);
+  // Expect death if we fail to initialize before starting to provide updates.
+  EXPECT_DEATH(ekf.UpdateEncodersAndGyro(0.0, 1.0, 2.0, {3.0, 3.0}, t0_),
+               "observations_.empty()");
+}
+
+TEST_F(HybridEkfDeathTest, DieOnNoU) {
+  // Expect death if the user does not provide U when creating a fresh
+  // measurement.
+  EXPECT_DEATH(ekf_.Correct({1, 2, 3}, nullptr, {}, {}, {}, {},
+                            t0_ + ::std::chrono::seconds(1)),
+               "U != nullptr");
+}
+
+// Because the user can choose to provide only one of make_h or (h, dhdx), check
+// that we die when an improper combination is provided.
+TEST_F(HybridEkfDeathTest, DieOnNoH) {
+  // Check that we die when no h-related functions are provided:
+  Input U(1, 2);
+  EXPECT_DEATH(ekf_.Correct({1, 2, 3}, &U, {}, {}, {}, {},
+                            t0_ + ::std::chrono::seconds(1)),
+               "make_h");
+  // Check that we die when only one of h and dhdx are provided:
+  EXPECT_DEATH(ekf_.Correct({1, 2, 3}, &U, {}, {},
+                            [](const State &) {
+                              return Eigen::Matrix<double, 3, 10>::Zero();
+                            },
+                            {}, t0_ + ::std::chrono::seconds(1)),
+               "make_h");
+  EXPECT_DEATH(ekf_.Correct({1, 2, 3}, &U, {},
+                            [](const State &, const Input &) {
+                              return Eigen::Matrix<double, 3, 1>::Zero();
+                            },
+                            {}, {}, t0_ + ::std::chrono::seconds(1)),
+               "make_h");
+}
+
+}  // namespace testing
+}  // namespace drivetrain
+}  // namespace control_loops
+}  // namespace frc971