Fix H() function in DrivetrainUkf and improve tests

We had a convention backwards that led to the calculation of the
expected accelerometer reading being incorrect.

Change-Id: I32a3c3fb6ac0698c575b4a9d1f28f9d57123379a
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator.cc b/frc971/control_loops/drivetrain/improved_down_estimator.cc
index 40dc8e3..31de9fe 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator.cc
+++ b/frc971/control_loops/drivetrain/improved_down_estimator.cc
@@ -139,16 +139,17 @@
 // States are X_hat_bar (position estimate) and P (Covariance)
 
 void QuaternionUkf::Predict(const Eigen::Matrix<double, 3, 1> &U,
-                            const Eigen::Matrix<double, 3, 1> &measurement) {
+                            const Eigen::Matrix<double, 3, 1> &measurement,
+                            const double dt) {
   // Compute the sigma points.
   // Our system is pretty linear. The traditional way of dealing with process
   // noise is to augment your state vector with the mean of the process noise,
   // and augment your covariance matrix with the covariance of your process
-  // noise. Sigma points are then computed. These points are then propegated
+  // noise. Sigma points are then computed. These points are then propagated
   // through the model. This ends up effectively meaning that perturbations
-  // from the unaugmented state with covariance P are propegated through the
+  // from the unaugmented state with covariance P are propagated through the
   // model, and points which are at the mean but with perturbations to simulated
-  // process noise are propegated through the system. The covariance is then
+  // process noise are propagated through the system. The covariance is then
   // calculated from this set of points, and works out to have a covariance of
   // essentially P + Q.
   //
@@ -163,12 +164,12 @@
   const Eigen::Matrix<double, 4, 3 * 2 + 1> X =
       GenerateSigmaPoints(X_hat_, P_ + Q_);
 
-  // Now, compute Y, the sigma points which have been propegated forwards by the
+  // Now, compute Y, the sigma points which have been propagated forwards by the
   // model.
   Eigen::Matrix<double, 4, 3 * 2 + 1> Y;
   for (int i = 0; i < Y.cols(); ++i) {
     // Y = Transformed sigma points
-    Y.col(i) = A(X.col(i), U);
+    Y.col(i) = A(X.col(i), U, dt);
   }
 
   // We now have the sigma points after the model update.
@@ -182,14 +183,17 @@
 
   // If the only obvious acceleration is that due to gravity, then accept the
   // measurement.
-  constexpr double kUseAccelThreshold = 0.1;
-  if (std::abs(measurement.squaredNorm() - 1.0) < kUseAccelThreshold) {
+  // TODO(james): Calibrate this on a real robot. This may require some sort of
+  // calibration routine.
+  constexpr double kUseAccelThreshold = 0.02;
+  if (std::abs(measurement.squaredNorm() - 1.0) > kUseAccelThreshold) {
     P_ = P_prior;
     return;
   }
 
   // TODO(austin): Maybe re-calculate the sigma points here before transforming
-  // them?  Otherwise we can't cleanly decouple the model and measurement updates.
+  // them?  Otherwise we can't cleanly decouple the model and measurement
+  // updates.
 
   // Apply the measurement transform to all the sigma points to get a
   // representation of the distribution of the measurement.
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator.h b/frc971/control_loops/drivetrain/improved_down_estimator.h
index 2b47f9e..75accfa 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator.h
+++ b/frc971/control_loops/drivetrain/improved_down_estimator.h
@@ -94,6 +94,14 @@
   // position flat on the ground with a heading of zero (which, in our normal
   // field coordinates, means pointed straight away from our driver's station
   // wall).
+  // The X axis is pointed straight out from the driver's station, the Z axis
+  // straight up and the Y axis straight to the left (i.e., a right-handed
+  // coordinate system).
+  // The quaternion itself represents the transformation from the body frame to
+  // global frame. E.g., for the gravity vector, the acceleration due to gravity
+  // in the global frame is equal to X_hat_ * gravity_in_robot_frame. Note that
+  // this convention does seem to be the inverse of that used in the paper
+  // referenced on Quaternion UKFs.
   constexpr static int kNumStates = 4;
   // Inputs to the system--we use the (x, y, z)  gyro measurements as the inputs
   // to the system.
@@ -119,20 +127,25 @@
   }
 
   // Handles updating the state of the UKF, given the gyro and accelerometer
-  // measurements.
+  // measurements. Given the design of the filter, U is the x/y/z gyro
+  // measurements and measurement is the accelerometer x/y/z measurements.
+  // dt is the length of the current timestep.
+  // U specifically corresponds with the U in the paper, which corresponds with
+  // the input to the system used by the filter.
   void Predict(const Eigen::Matrix<double, kNumInputs, 1> &U,
-               const Eigen::Matrix<double, 3, 1> &measurement);
+               const Eigen::Matrix<double, kNumMeasurements, 1> &measurement,
+               const double dt);
 
   // Returns the updated state for X after one time step, given the current
   // state and gyro measurements.
-  virtual Eigen::Matrix<double, 4, 1> A(
-      const Eigen::Matrix<double, 4, 1> &X,
-      const Eigen::Matrix<double, kNumInputs, 1> &U) const = 0;
+  virtual Eigen::Matrix<double, kNumStates, 1> A(
+      const Eigen::Matrix<double, kNumStates, 1> &X,
+      const Eigen::Matrix<double, kNumInputs, 1> &U, const double dt) const = 0;
 
   // Returns the current expected accelerometer measurements given the current
   // state.
-  virtual Eigen::Matrix<double, 3, 1> H(
-      const Eigen::Matrix<double, 4, 1> &X) const = 0;
+  virtual Eigen::Matrix<double, kNumMeasurements, 1> H(
+      const Eigen::Matrix<double, kNumStates, 1> &X) const = 0;
 
   // Returns the current estimate of the robot's orientation. Note that this
   // filter does not have anything other than the gyro with which to estimate
@@ -140,7 +153,7 @@
   // filters.
   const Eigen::Quaternion<double> &X_hat() const { return X_hat_; }
 
-  Eigen::Matrix<double, 3, 1> Z_hat() const { return Z_hat_; };
+  Eigen::Matrix<double, kNumMeasurements, 1> Z_hat() const { return Z_hat_; };
 
  private:
   // Measurement Noise (Uncertainty)
@@ -155,13 +168,11 @@
   Eigen::Quaternion<double> X_hat_;
 
   // Current expected accelerometer measurement.
-  Eigen::Matrix<double, 3, 1> Z_hat_;
+  Eigen::Matrix<double, kNumMeasurements, 1> Z_hat_;
 };
 
 class DrivetrainUkf : public QuaternionUkf {
  public:
-  constexpr static double kDt = 0.00505;
-
   // UKF for http://kodlab.seas.upenn.edu/uploads/Arun/UKFpaper.pdf
   // Reference in case the link is dead:
   // Kraft, Edgar. "A quaternion-based unscented Kalman filter for orientation
@@ -195,19 +206,24 @@
   // Moves the robot by the provided rotation vector (U).
   Eigen::Matrix<double, kNumStates, 1> A(
       const Eigen::Matrix<double, kNumStates, 1> &X,
-      const Eigen::Matrix<double, kNumInputs, 1> &U) const override {
+      const Eigen::Matrix<double, kNumInputs, 1> &U,
+      const double dt) const override {
     return RungeKutta(
-        std::bind(&QuaternionDerivative, U, std::placeholders::_1), X, kDt);
+        std::bind(&QuaternionDerivative, U, std::placeholders::_1), X, dt);
   }
 
   // Returns the expected accelerometer measurement (which is just going to be
   // 1g downwards).
   Eigen::Matrix<double, kNumMeasurements, 1> H(
       const Eigen::Matrix<double, kNumStates, 1> &X) const override {
-    // TODO(austin): Figure out how to compute what the sensors *should* read.
+    // Assume that we expect to see a reading of (0, 0, 1) when flat on the
+    // ground.
+    // TODO(james): Figure out a calibration routine for managing the fact that
+    // the accelerometer will not be perfectly oriented within the robot (or
+    // determine that calibration routines would be unnecessary).
     Eigen::Quaternion<double> Xquat(X);
     Eigen::Matrix<double, 3, 1> gprime =
-        Xquat * Eigen::Matrix<double, 3, 1>(0.0, 0.0, -1.0);
+        Xquat.conjugate() * Eigen::Matrix<double, 3, 1>(0.0, 0.0, 1.0);
     return gprime;
   }
 };
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator_test.cc b/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
index db7eeee..31adaf2 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
+++ b/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
@@ -12,6 +12,19 @@
 namespace control_loops {
 namespace testing {
 
+namespace {
+// Check if two quaternions are logically equal, to within some reasonable
+// tolerance. This is needed because a single rotation can be represented by two
+// quaternions.
+bool QuaternionEqual(const Eigen::Quaterniond &a, const Eigen::Quaterniond &b,
+                     double tolerance) {
+  // If a == b, then a.inverse() * b will be the identity. The identity
+  // quaternion is the only time where the vector portion of the quaternion is
+  // zero.
+  return (a.inverse() * b).vec().norm() <= tolerance;
+}
+}  // namespace
+
 // Do a known transformation to see if quaternion integration is working
 // correctly.
 TEST(RungeKuttaTest, QuaternionIntegral) {
@@ -31,7 +44,7 @@
   VLOG(1) << "ux is " << ux;
   VLOG(1) << "qux is " << qux;
 
-  // Start by rotating around the X world vector for pi/2
+  // Start by rotating around the X body vector for pi/2
   Eigen::Quaternion<double> integral1(
       RungeKutta(std::bind(&drivetrain::DrivetrainUkf::QuaternionDerivative, ux,
                            std::placeholders::_1),
@@ -39,7 +52,7 @@
 
   VLOG(1) << "integral1 * uz => " << integral1 * uz;
 
-  // Then rotate around the Y world vector for pi/2
+  // Then rotate around the Y body vector for pi/2
   Eigen::Quaternion<double> integral2(
       RungeKutta(std::bind(&drivetrain::DrivetrainUkf::QuaternionDerivative, uy,
                            std::placeholders::_1),
@@ -47,12 +60,14 @@
 
   VLOG(1) << "integral2 * uz => " << integral2 * uz;
 
-  // Then rotate around the X world vector for -pi/2
+  // Then rotate around the X body vector for -pi/2
   Eigen::Quaternion<double> integral3(
       RungeKutta(std::bind(&drivetrain::DrivetrainUkf::QuaternionDerivative,
                            -ux, std::placeholders::_1),
                  integral2.normalized().coeffs(), 0.5 * M_PI));
 
+  integral1.normalize();
+  integral2.normalize();
   integral3.normalize();
 
   VLOG(1) << "Integral is w: " << integral1.w() << " vec: " << integral1.vec()
@@ -62,15 +77,87 @@
           << " norm " << integral3.norm();
 
   VLOG(1) << "ux => " << integral3 * ux;
+  EXPECT_NEAR(0.0, (ux - integral1 * ux).norm(), 5e-2);
+  EXPECT_NEAR(0.0, (uz - integral1 * uy).norm(), 5e-2);
+  EXPECT_NEAR(0.0, (-uy - integral1 * uz).norm(), 5e-2);
+
+  EXPECT_NEAR(0.0, (uy - integral2 * ux).norm(), 5e-2);
+  EXPECT_NEAR(0.0, (uz - integral2 * uy).norm(), 5e-2);
+  EXPECT_NEAR(0.0, (ux - integral2 * uz).norm(), 5e-2);
+
   EXPECT_NEAR(0.0, (uy - integral3 * ux).norm(), 5e-2);
+  EXPECT_NEAR(0.0, (-ux - integral3 * uy).norm(), 5e-2);
+  EXPECT_NEAR(0.0, (uz - integral3 * uz).norm(), 5e-2);
 }
 
-TEST(RungeKuttaTest, Ukf) {
+TEST(RungeKuttaTest, UkfConstantRotation) {
   drivetrain::DrivetrainUkf dtukf;
-  Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
+  const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
+  EXPECT_EQ(0.0,
+            (Eigen::Vector3d(0.0, 0.0, 1.0) - dtukf.H(dtukf.X_hat().coeffs()))
+                .norm());
   Eigen::Matrix<double, 3, 1> measurement;
   measurement.setZero();
-  dtukf.Predict(ux, measurement);
+  for (int ii = 0; ii < 200; ++ii) {
+    dtukf.Predict(ux * M_PI_2, measurement, 0.005);
+  }
+  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(M_PI_2, ux));
+  EXPECT_TRUE(QuaternionEqual(expected, dtukf.X_hat(), 0.01))
+      << "Expected: " << expected.coeffs()
+      << " Got: " << dtukf.X_hat().coeffs();
+  EXPECT_NEAR(0.0,
+            (Eigen::Vector3d(0.0, 1.0, 0.0) - dtukf.H(dtukf.X_hat().coeffs()))
+                .norm(), 1e-10);
+}
+
+// Tests that if the gyro indicates no movement but that the accelerometer shows
+// that we are slightly rotated, that we eventually adjust our estimate to be
+// correct.
+TEST(RungeKuttaTest, UkfAccelCorrectsBias) {
+  drivetrain::DrivetrainUkf dtukf;
+  const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
+  Eigen::Matrix<double, 3, 1> measurement;
+  // Supply the accelerometer with a slightly off reading to ensure that we
+  // don't require exactly 1g to work.
+  measurement << 0.01, 0.99, 0.0;
+  EXPECT_TRUE(
+      QuaternionEqual(Eigen::Quaterniond::Identity(), dtukf.X_hat(), 0.0))
+      << "X_hat: " << dtukf.X_hat().coeffs();
+  EXPECT_EQ(0.0,
+            (Eigen::Vector3d(0.0, 0.0, 1.0) - dtukf.H(dtukf.X_hat().coeffs()))
+                .norm());
+  for (int ii = 0; ii < 200; ++ii) {
+    dtukf.Predict({0.0, 0.0, 0.0}, measurement, 0.005);
+  }
+  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(M_PI_2, ux));
+  EXPECT_TRUE(QuaternionEqual(expected, dtukf.X_hat(), 0.01))
+      << "Expected: " << expected.coeffs()
+      << " Got: " << dtukf.X_hat().coeffs();
+}
+
+// Tests that if the accelerometer is reading values with a magnitude that isn't ~1g,
+// that we are slightly rotated, that we eventually adjust our estimate to be
+// correct.
+TEST(RungeKuttaTest, UkfIgnoreBadAccel) {
+  drivetrain::DrivetrainUkf dtukf;
+  const Eigen::Vector3d uz = Eigen::Vector3d::UnitZ();
+  Eigen::Matrix<double, 3, 1> measurement;
+  // Set up a scenario where, if we naively took the accelerometer readings, we
+  // would think that we were rotated. But the gyro readings indicate that we
+  // are only rotating about the Z (yaw) axis.
+  measurement << 0.3, 1.0, 0.0;
+  for (int ii = 0; ii < 200; ++ii) {
+    dtukf.Predict({0.0, 0.0, 1.0}, measurement, 0.005);
+  }
+  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(1.0, uz));
+  EXPECT_TRUE(QuaternionEqual(expected, dtukf.X_hat(), 1e-1))
+      << "Expected: " << expected.coeffs()
+      << " Got: " << dtukf.X_hat().coeffs();
+  EXPECT_NEAR(
+      0.0,
+      (Eigen::Vector3d(0.0, 0.0, 1.0) - dtukf.H(dtukf.X_hat().coeffs())).norm(),
+      1e-10)
+      << dtukf.H(dtukf.X_hat().coeffs());
 }
 
 // Tests that small perturbations around a couple quaternions averaged out