Clean up down estimator to work with current IMU

Change-Id: Icac944b3b687d9ff27af218b5740bcee88177989
diff --git a/frc971/control_loops/drivetrain/drivetrain.cc b/frc971/control_loops/drivetrain/drivetrain.cc
index 1afc32d..7dfd0e6 100644
--- a/frc971/control_loops/drivetrain/drivetrain.cc
+++ b/frc971/control_loops/drivetrain/drivetrain.cc
@@ -58,6 +58,9 @@
     // present.
     imu_values_fetcher_.Fetch();
   });
+  if (dt_config.is_simulated) {
+    down_estimator_.assume_perfect_gravity();
+  }
 }
 
 int DrivetrainLoop::ControllerIndexFromGears() {
@@ -316,7 +319,10 @@
         dt_openloop_.PopulateStatus(status->fbb());
 
     const flatbuffers::Offset<DownEstimatorState> down_estimator_state_offset =
-        down_estimator_.PopulateStatus(status->fbb());
+        down_estimator_.PopulateStatus(status->fbb(), monotonic_now);
+
+    const flatbuffers::Offset<ImuZeroerState> zeroer_offset =
+        imu_zeroer_.PopulateStatus(status->fbb());
 
     flatbuffers::Offset<LineFollowLogging> line_follow_logging_offset =
         dt_line_follow_.PopulateStatus(status);
@@ -357,6 +363,7 @@
     builder.add_line_follow_logging(line_follow_logging_offset);
     builder.add_trajectory_logging(trajectory_logging_offset);
     builder.add_down_estimator(down_estimator_state_offset);
+    builder.add_zeroing(zeroer_offset);
     status->Send(builder.Finish());
   }
 
diff --git a/frc971/control_loops/drivetrain/drivetrain_config.h b/frc971/control_loops/drivetrain/drivetrain_config.h
index 685516f..5a29008 100644
--- a/frc971/control_loops/drivetrain/drivetrain_config.h
+++ b/frc971/control_loops/drivetrain/drivetrain_config.h
@@ -96,6 +96,16 @@
   // Whether the shift button on the pistol grip enables line following mode.
   bool pistol_grip_shift_enables_line_follow = false;
 
+  // Rotation matrix from the IMU's coordinate frame to the robot's coordinate
+  // frame.
+  // I.e., imu_transform * imu_readings will give the imu readings in the
+  // robot frame.
+  Eigen::Matrix<double, 3, 3> imu_transform =
+      Eigen::Matrix<double, 3, 3>::Identity();
+
+  // True if we are running a simulated drivetrain.
+  bool is_simulated = false;
+
   // Converts the robot state to a linear distance position, velocity.
   static Eigen::Matrix<Scalar, 2, 1> LeftRightToLinear(
       const Eigen::Matrix<Scalar, 7, 1> &left_right) {
diff --git a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
index f606495..1de5d40 100644
--- a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
+++ b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
@@ -72,6 +72,8 @@
 
     // Run for enough time to allow the gyro/imu zeroing code to run.
     RunFor(std::chrono::seconds(10));
+    drivetrain_status_fetcher_.Fetch();
+    EXPECT_TRUE(CHECK_NOTNULL(drivetrain_status_fetcher_->zeroing())->zeroed());
   }
   virtual ~DrivetrainTest() {}
 
diff --git a/frc971/control_loops/drivetrain/drivetrain_test_lib.cc b/frc971/control_loops/drivetrain/drivetrain_test_lib.cc
index 8b4aab1..336025a 100644
--- a/frc971/control_loops/drivetrain/drivetrain_test_lib.cc
+++ b/frc971/control_loops/drivetrain/drivetrain_test_lib.cc
@@ -75,7 +75,10 @@
 
       0.25,
       1.00,
-      1.00};
+      1.00,
+      false,
+      Eigen::Matrix3d::Identity(),
+      /*is_simulated=*/true};
 
   return kDrivetrainConfig;
 };
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator.cc b/frc971/control_loops/drivetrain/improved_down_estimator.cc
index 8b2ff4c..cdb1587 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator.cc
+++ b/frc971/control_loops/drivetrain/improved_down_estimator.cc
@@ -136,11 +136,20 @@
   return corrected_X.block<3, 1>(0, 0) * scalar;
 }
 
-// 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 aos::monotonic_clock::duration dt) {
+  const Eigen::Matrix<double, 3, 1> calibrated_measurement =
+      imu_transform_ * measurement;
+  const Eigen::Matrix<double, 3, 1> calibrated_U = imu_transform_ * U;
+  DoPredict(calibrated_U, calibrated_measurement, dt);
+  IterationCleanup(calibrated_measurement, calibrated_U);
+}
+
+// States are X_hat_bar (position estimate) and P (Covariance)
+void QuaternionUkf::DoPredict(const Eigen::Matrix<double, 3, 1> &U,
+                            const Eigen::Matrix<double, 3, 1> &measurement,
+                            const aos::monotonic_clock::duration 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,
@@ -160,7 +169,7 @@
   // somewhere in this system, we'll have to revisit this assumption.
 
   // Now, compute the actual sigma points using the columns of S as the
-  // pertubation vectors. The last point is the original mean.
+  // perturbation vectors. The last point is the original mean.
   const Eigen::Matrix<double, 4, 3 * 2 + 1> X =
       GenerateSigmaPoints(X_hat_, P_ + Q_);
 
@@ -183,13 +192,29 @@
 
   // If the only obvious acceleration is that due to gravity, then accept the
   // measurement.
-  // 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) {
+  const double kUseAccelThreshold = assume_perfect_gravity_ ? 1e-10 : 0.025;
+  const double accel_norm = measurement.norm();
+  if (std::abs(accel_norm - 1.0) > kUseAccelThreshold) {
     P_ = P_prior;
+    consecutive_still_ = 0;
     return;
   }
+  // Whenever we seem to have been still for a while, zero the integrated
+  // velocity. Because we just use this for debugging, only set it once per time
+  // duration when we are paused--this lets us observe how far things drift
+  // while sitting still.
+  if (consecutive_still_ == 1000) {
+    pos_vel_.block<3, 1>(3, 0).setZero();
+  }
+  // Don't do accelerometer updates unless we have been roughly still for a
+  // decent number of iterations.
+  if (++consecutive_still_ < 50) {
+    return;
+  }
+
+  // Update the gravity_magnitude_ using a semi-arbitrary time-constant that
+  // seems to provide decent results.
+  gravity_magnitude_ += 0.001 * (accel_norm - gravity_magnitude_);
 
   // TODO(austin): Maybe re-calculate the sigma points here before transforming
   // them?  Otherwise we can't cleanly decouple the model and measurement
@@ -200,7 +225,7 @@
   Eigen::Matrix<double, kNumMeasurements, 3 * 2 + 1> Z;
   Z_hat_.setZero();
   for (int i = 0; i < Z.cols(); ++i) {
-    Z.col(i) = H(Y.col(i));
+    Z.col(i) = H(Y.col(i)) * accel_norm;
 
     // Compute the mean in addition.
     Z_hat_ += Z.col(i) / Z.cols();
@@ -242,6 +267,47 @@
   P_ = P_prior - K * P_vv * K.transpose();
 }
 
+void QuaternionUkf::IterationCleanup(const Eigen::Vector3d &accel,
+                                     const Eigen::Vector3d &gyro) {
+  // To make certain debugging simpler, resolve the ambiguity in the quaternion
+  // space by making the w coefficient always be positive.
+  if (X_hat_.coeffs().w() < 0) {
+    X_hat_.coeffs() *= -1.0;
+  }
+  const Eigen::Vector3d robot_x_in_global_frame =
+      X_hat() * Eigen::Vector3d::UnitX();
+  const double yaw =
+      std::atan2(robot_x_in_global_frame.y(), robot_x_in_global_frame.x());
+  // The down estimator UKF does a poor job of estimating yaw, so back out and
+  // remove the yaw estimate and let downstream estimators take care of it.
+  X_hat_ = Eigen::AngleAxis<double>(-yaw, Eigen::Vector3d::UnitZ()) * X_hat();
+  last_accels_[buffer_index_] =
+      (X_hat_ * accel - Eigen::Vector3d::UnitZ() * gravity_magnitude_) *
+      9.80665;
+  last_yaw_rates_[buffer_index_] = gyro.z();
+  buffer_index_ = (buffer_index_ + 1) % last_accels_.size();
+}
+
+Eigen::Matrix<double, 6, 1> QuaternionUkf::PosVelDerivative(
+    const Eigen::Matrix<double, 6, 1> &pos_vel,
+    const Eigen::Matrix<double, 3, 1> &accel) {
+  Eigen::Matrix<double, 6, 1> derivative;
+  derivative.block<3, 1>(0, 0) = pos_vel.block<3, 1>(3, 0);
+  derivative.block<3, 1>(3, 0) = accel;
+  return derivative;
+}
+
+void QuaternionUkf::UpdatePosition(aos::monotonic_clock::duration dt,
+                                   const Eigen::Vector3d &accel) {
+  const double dt_sec = aos::time::DurationInSeconds(dt);
+  yaw_ += avg_recent_yaw_rates() * dt_sec;
+  pos_vel_ = RungeKutta(
+      std::bind(
+          &QuaternionUkf::PosVelDerivative, this, std::placeholders::_1,
+          Eigen::AngleAxis<double>(yaw_, Eigen::Vector3d::UnitZ()) * accel),
+      pos_vel_, dt_sec);
+}
+
 Eigen::Matrix<double, 3, 3> ComputeQuaternionCovariance(
     const Eigen::Quaternion<double> &mean,
     const Eigen::Matrix<double, 4, 7> &points,
@@ -286,24 +352,38 @@
   return X;
 }
 
+void DrivetrainUkf::UpdateIntegratedPositions(
+    aos::monotonic_clock::time_point now) {
+  if (last_pos_vel_update_ != aos::monotonic_clock::min_time) {
+    UpdatePosition(now - last_pos_vel_update_, avg_recent_accel());
+  }
+  last_pos_vel_update_ = now;
+}
+
 flatbuffers::Offset<DownEstimatorState> DrivetrainUkf::PopulateStatus(
-    flatbuffers::FlatBufferBuilder *fbb) const {
+    flatbuffers::FlatBufferBuilder *fbb, aos::monotonic_clock::time_point now) {
+  // On everything but the first iteration, integrate the position/velocity
+  // estimates.
+  // TODO(james): Consider optionally disabling this to avoid excess CPU usage.
+  UpdateIntegratedPositions(now);
+
   DownEstimatorState::Builder builder(*fbb);
   builder.add_quaternion_x(X_hat().x());
   builder.add_quaternion_y(X_hat().y());
   builder.add_quaternion_z(X_hat().z());
   builder.add_quaternion_w(X_hat().w());
 
+  builder.add_yaw(yaw());
+
+  // Calculate the current pitch numbers to provide a more human-readable
+  // debugging output.
   {
-    // Note that this algorithm is not very numerically stable near pitches of
-    // +/- pi / 2.
     const Eigen::Vector3d robot_x_in_global_frame =
         X_hat() * Eigen::Vector3d::UnitX();
-    builder.add_yaw(
-        std::atan2(robot_x_in_global_frame.y(), robot_x_in_global_frame.x()));
     const double xy_norm = robot_x_in_global_frame.block<2, 1>(0, 0).norm();
-    builder.add_longitudinal_pitch(
-        std::atan2(-robot_x_in_global_frame.z(), xy_norm));
+    const double pitch = std::atan2(-robot_x_in_global_frame.z(), xy_norm);
+
+    builder.add_longitudinal_pitch(pitch);
   }
   {
     const Eigen::Vector3d robot_y_in_global_frame =
@@ -313,6 +393,31 @@
         std::atan2(robot_y_in_global_frame.z(), xy_norm));
   }
 
+  builder.add_position_x(pos_vel()(0, 0));
+  builder.add_position_y(pos_vel()(1, 0));
+  builder.add_position_z(pos_vel()(2, 0));
+  builder.add_velocity_x(pos_vel()(3, 0));
+  builder.add_velocity_y(pos_vel()(4, 0));
+  builder.add_velocity_z(pos_vel()(5, 0));
+
+  {
+    const Eigen::Vector3d last_accel_avg = avg_recent_accel();
+    builder.add_accel_x(last_accel_avg.x());
+    builder.add_accel_y(last_accel_avg.y());
+    builder.add_accel_z(last_accel_avg.z());
+  }
+
+  {
+    const Eigen::Vector3d expected_accel = H(X_hat().coeffs());
+    builder.add_expected_accel_x(expected_accel.x());
+    builder.add_expected_accel_y(expected_accel.y());
+    builder.add_expected_accel_z(expected_accel.z());
+  }
+
+  builder.add_gravity_magnitude(gravity_magnitude());
+
+  builder.add_consecutive_still(consecutive_still());
+
   return builder.Finish();
 }
 
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator.h b/frc971/control_loops/drivetrain/improved_down_estimator.h
index aa45e4e..27de867 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator.h
+++ b/frc971/control_loops/drivetrain/improved_down_estimator.h
@@ -112,21 +112,30 @@
   // Measurements to use for correcting the estimated system state. These
   // correspond to (x, y, z) measurements from the accelerometer.
   constexpr static int kNumMeasurements = 3;
-  QuaternionUkf() {
-    // TODO(james): Tune the process/measurement noises.
+  QuaternionUkf(const Eigen::Matrix<double, 3, 3> &imu_transform =
+                    Eigen::Matrix<double, 3, 3>::Identity())
+      : imu_transform_(imu_transform) {
+    // The various noise matrices are tuned to provide values that seem
+    // reasonable given our current setup (using the ADIS16470 for
+    // measurements).
     R_.setIdentity();
-    R_ /= 100.0;
+    R_ /= std::pow(1000.0, 2);
 
     Q_.setIdentity();
-    Q_ /= 10000.0;
+    Q_ /= std::pow(2000.0 * 500.0, 2);
 
     // Assume that the robot starts flat on the ground pointed straight forward.
     X_hat_ = Eigen::Quaternion<double>(1.0, 0.0, 0.0, 0.0);
 
-    // TODO(james): Determine an appropriate starting noise estimate. Probably
-    // not too critical.
     P_.setIdentity();
     P_ /= 1000.0;
+
+    pos_vel_.setZero();
+    for (auto &last_accel : last_accels_) {
+      last_accel.setZero();
+    }
+
+    last_yaw_rates_.fill(0);
   }
 
   // Handles updating the state of the UKF, given the gyro and accelerometer
@@ -135,6 +144,8 @@
   // 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.
+  // Accelerometer measurements should be in g's, and gyro measurements in
+  // radians / sec.
   void Predict(const Eigen::Matrix<double, kNumInputs, 1> &U,
                const Eigen::Matrix<double, kNumMeasurements, 1> &measurement,
                const aos::monotonic_clock::duration dt);
@@ -159,7 +170,57 @@
 
   Eigen::Matrix<double, kNumMeasurements, 1> Z_hat() const { return Z_hat_; };
 
+  Eigen::Matrix<double, 6, 1> pos_vel() const { return pos_vel_; }
+  double yaw() const { return yaw_; }
+  double avg_recent_yaw_rates() const {
+    double avg = 0.0;
+    for (const auto &yaw_rate : last_yaw_rates_) {
+      avg += yaw_rate;
+    };
+    avg /= last_yaw_rates_.size();
+    return avg;
+  }
+  Eigen::Matrix<double, 3, 1> avg_recent_accel() const {
+    Eigen::Vector3d avg;
+    avg.setZero();
+    for (const auto &accel : last_accels_) {
+      avg += accel;
+    };
+    avg /= last_accels_.size();
+    return avg;
+  }
+
+  double gravity_magnitude() const { return gravity_magnitude_; }
+  int consecutive_still() const { return consecutive_still_; }
+
+  // Causes the down estimator to assume that gravity is exactly one g and that
+  // the accelerometer readings have no meaningful noise. This is used for
+  // dealing with tests where the IMU readings are actually nearly perfect.
+  void assume_perfect_gravity() { assume_perfect_gravity_ = true; }
+
+ protected:
+  // Updates the position/velocity integration.
+  void UpdatePosition(aos::monotonic_clock::duration dt,
+                      const Eigen::Vector3d &accel);
+
+  aos::monotonic_clock::time_point last_pos_vel_update_ =
+      aos::monotonic_clock::min_time;
+
  private:
+  // Length of buffer to maintain for averaging recent acceleration/gyro values.
+  static constexpr size_t kBufferSize = 10;
+
+  // Does all the heavy lifting from the Predict() call.
+  void DoPredict(const Eigen::Vector3d &accel, const Eigen::Vector3d &gyro,
+                 const aos::monotonic_clock::duration dt);
+  // Runs some cleanup that needs to happen at the end of any iteration.
+  void IterationCleanup(const Eigen::Vector3d &accel,
+                        const Eigen::Vector3d &gyro);
+
+  // Takes in the current pos_vel_ vector as well as an acceleration in the
+  // world-frame and returns the derivative. All numbers in m, m/s, or m/s/s.
+  Eigen::Matrix<double, 6, 1> PosVelDerivative(
+      const Eigen::Matrix<double, 6, 1> &pos_vel, const Eigen::Vector3d &accel);
   // Measurement Noise (Uncertainty)
   Eigen::Matrix<double, kNumInputs, kNumInputs> R_;
   // Model noise. Note that both this and P are 3 x 3 matrices, despite the
@@ -173,8 +234,46 @@
 
   // Current expected accelerometer measurement.
   Eigen::Matrix<double, kNumMeasurements, 1> Z_hat_;
+
+  // Current position and velocity vector in format:
+  // {pos_x, pos_y, pos_z, vel_x, vel_y, vel_z}, in meters and meters / sec.
+  // This is just used for cosmetic purposes, as it only accounts for IMU
+  // measurements and so is prone to drift.
+  Eigen::Matrix<double, 6, 1> pos_vel_;
+  // Current yaw estimate, obtained purely by integrating the gyro measurements.
+  double yaw_ = 0;
+  // Circular buffer in which to store the most recent acceleration
+  // measurements. These accelerations are transformed into the robot's yaw
+  // frame and have gravity removed (i.e., the users of last_accels_ should not
+  // have to worry about pitch/roll or the gravitational component of the
+  // acceleration).
+  // As such, x is the robot's longitudinal acceleration, y is the lateral
+  // acceleration, and z is the up/down acceleration. All are in m/s/s.
+  int buffer_index_ = 0;
+  std::array<Eigen::Matrix<double, 3, 1>, kBufferSize> last_accels_;
+  // Array of the most recent yaw rates, in rad/sec.
+  std::array<double, kBufferSize> last_yaw_rates_;
+
+  // Number of consecutive iterations in which we think that the robot has been
+  // in a zero-acceleration state. We only accept accelerometer corrections to
+  // the down estimator when we've been static for a sufficiently long time.
+  int consecutive_still_ = 0;
+  // This variable tracks the current estimated acceleration due to gravity, in
+  // g's. This helps to compensate for both local variations in gravity as well
+  // as for calibration errors on individual accelerometer axes.
+  double gravity_magnitude_ = 1.0;
+
+  // The transformation from the IMU's frame to the robot frame.
+  Eigen::Matrix<double, 3, 3> imu_transform_;
+
+  bool assume_perfect_gravity_ = false;
 };
 
+// TODO(james): The lines between DrivetrainUkf and QuaternionUkf have blurred
+// to the point where there is minimal distinction. Either remove the
+// unnecessary abstraction or figure out what we actually care about abstracting
+// (e.g., we do eventually need to add some ability to provide a custom
+// accelerometer calibration).
 class DrivetrainUkf : public QuaternionUkf {
  public:
   // UKF for http://kodlab.seas.upenn.edu/uploads/Arun/UKFpaper.pdf
@@ -228,12 +327,16 @@
     // determine that calibration routines would be unnecessary).
     Eigen::Quaternion<double> Xquat(X);
     Eigen::Matrix<double, 3, 1> gprime =
-        Xquat.conjugate() * Eigen::Matrix<double, 3, 1>(0.0, 0.0, 1.0);
+        Xquat.conjugate() * Eigen::Matrix<double, 3, 1>(0.0, 0.0, 1.0) *
+        1.0;
     return gprime;
   }
 
+  void UpdateIntegratedPositions(aos::monotonic_clock::time_point now);
+
   flatbuffers::Offset<DownEstimatorState> PopulateStatus(
-      flatbuffers::FlatBufferBuilder *fbb) const;
+      flatbuffers::FlatBufferBuilder *fbb,
+      aos::monotonic_clock::time_point now);
 };
 
 }  // namespace drivetrain
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator_test.cc b/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
index a6d02f6..86ef5e9 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
+++ b/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
@@ -27,7 +27,7 @@
 
 // Do a known transformation to see if quaternion integration is working
 // correctly.
-TEST(RungeKuttaTest, QuaternionIntegral) {
+TEST(DownEstimatorTest, QuaternionIntegral) {
   Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
   Eigen::Vector3d uy = Eigen::Vector3d::UnitY();
   Eigen::Vector3d uz = Eigen::Vector3d::UnitZ();
@@ -90,7 +90,7 @@
   EXPECT_NEAR(0.0, (uz - integral3 * uz).norm(), 5e-2);
 }
 
-TEST(RungeKuttaTest, UkfConstantRotation) {
+TEST(DownEstimatorTest, UkfConstantRotation) {
   drivetrain::DrivetrainUkf dtukf;
   const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
   EXPECT_EQ(0.0,
@@ -111,42 +111,53 @@
 }
 
 // Tests that the euler angles in the status message are correct.
-TEST(RungeKuttaTest, UkfEulerStatus) {
+TEST(DownEstimatorTest, UkfEulerStatus) {
   drivetrain::DrivetrainUkf dtukf;
   const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
   const Eigen::Vector3d uy = Eigen::Vector3d::UnitY();
   const Eigen::Vector3d uz = Eigen::Vector3d::UnitZ();
   // First, rotate 3 radians in the yaw axis, then 0.5 radians in the pitch
   // axis, and then 0.1 radians about the roll axis.
+  // The down estimator should ignore any of the pitch movement.
   constexpr double kYaw = 3.0;
   constexpr double kPitch = 0.5;
   constexpr double kRoll = 0.1;
   Eigen::Matrix<double, 3, 1> measurement;
   measurement.setZero();
+  aos::monotonic_clock::time_point now = aos::monotonic_clock::epoch();
+  const std::chrono::milliseconds dt(5);
+  // Run a bunch of one-second rotations at the appropriate rate to cause the
+  // total pitch/roll/yaw to be kPitch/kRoll/kYaw.
   for (int ii = 0; ii < 200; ++ii) {
-    dtukf.Predict(uz * kYaw, measurement, std::chrono::milliseconds(5));
+    dtukf.UpdateIntegratedPositions(now);
+    now += dt;
+    dtukf.Predict(uz * kYaw, measurement, dt);
   }
   for (int ii = 0; ii < 200; ++ii) {
-    dtukf.Predict(uy * kPitch, measurement, std::chrono::milliseconds(5));
+    dtukf.UpdateIntegratedPositions(now);
+    now += dt;
+    dtukf.Predict(uy * kPitch, measurement, dt);
   }
+  EXPECT_FLOAT_EQ(kYaw, dtukf.yaw());
   for (int ii = 0; ii < 200; ++ii) {
-    dtukf.Predict(ux * kRoll, measurement, std::chrono::milliseconds(5));
+    dtukf.UpdateIntegratedPositions(now);
+    now += dt;
+    dtukf.Predict(ux * kRoll, measurement, dt);
   }
-  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(kYaw, uz) *
-                                    Eigen::AngleAxis<double>(kPitch, uy) *
+  EXPECT_FLOAT_EQ(kYaw, dtukf.yaw());
+  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(kPitch, uy) *
                                     Eigen::AngleAxis<double>(kRoll, ux));
   flatbuffers::FlatBufferBuilder fbb;
   fbb.ForceDefaults(true);
-  fbb.Finish(dtukf.PopulateStatus(&fbb));
+  fbb.Finish(dtukf.PopulateStatus(&fbb, now));
 
   aos::FlatbufferDetachedBuffer<drivetrain::DownEstimatorState> state(
       fbb.Release());
   EXPECT_EQ(kPitch, state.message().longitudinal_pitch());
-  EXPECT_EQ(kYaw, state.message().yaw());
   // The longitudinal pitch is not actually the same number as the roll, so we
   // don't check it here.
 
-  EXPECT_TRUE(QuaternionEqual(expected, dtukf.X_hat(), 0.01))
+  EXPECT_TRUE(QuaternionEqual(expected, dtukf.X_hat(), 0.0001))
       << "Expected: " << expected.coeffs()
       << " Got: " << dtukf.X_hat().coeffs();
 }
@@ -155,7 +166,7 @@
 // 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) {
+TEST(DownEstimatorTest, UkfAccelCorrectsBias) {
   drivetrain::DrivetrainUkf dtukf;
   const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
   Eigen::Matrix<double, 3, 1> measurement;
@@ -180,31 +191,33 @@
 // 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) {
+TEST(DownEstimatorTest, UkfIgnoreBadAccel) {
   drivetrain::DrivetrainUkf dtukf;
-  const Eigen::Vector3d uz = Eigen::Vector3d::UnitZ();
+  const Eigen::Vector3d uy = Eigen::Vector3d::UnitY();
   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.
+  // are only rotating about the Y (pitch) axis.
   measurement << 0.3, 1.0, 0.0;
   for (int ii = 0; ii < 200; ++ii) {
-    dtukf.Predict({0.0, 0.0, 1.0}, measurement, std::chrono::milliseconds(5));
+    dtukf.Predict({0.0, M_PI_2, 0.0}, measurement,
+                  std::chrono::milliseconds(5));
   }
-  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(1.0, uz));
+  const Eigen::Quaterniond expected(Eigen::AngleAxis<double>(M_PI_2, uy));
   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(),
+      (Eigen::Vector3d(-1.0, 0.0, 0.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
 // return the original quaternion.
-TEST(RungeKuttaTest, QuaternionMean) {
+TEST(DownEstimatorTest, QuaternionMean) {
   Eigen::Matrix<double, 4, 7> vectors;
   vectors.col(0) << 0, 0, 0, 1;
   for (int i = 0; i < 3; ++i) {
@@ -229,7 +242,7 @@
 
 // Tests that computing sigma points, and then computing the mean and covariance
 // returns the original answer.
-TEST(RungeKuttaTest, SigmaPoints) {
+TEST(DownEstimatorTest, SigmaPoints) {
   const Eigen::Quaternion<double> mean(
       Eigen::AngleAxis<double>(M_PI / 2.0, Eigen::Vector3d::UnitX()));
 
@@ -260,7 +273,7 @@
 
 // Tests that computing sigma points with a large covariance that will precisely
 // wrap, that we do clip the perturbations.
-TEST(RungeKuttaTest, ClippedSigmaPoints) {
+TEST(DownEstimatorTest, ClippedSigmaPoints) {
   const Eigen::Quaternion<double> mean(
       Eigen::AngleAxis<double>(M_PI / 2.0, Eigen::Vector3d::UnitX()));
 
@@ -291,7 +304,7 @@
 }
 
 // Tests that ToRotationVectorFromQuaternion works for a 0 rotation.
-TEST(RungeKuttaTest, ToRotationVectorFromQuaternionAtZero) {
+TEST(DownEstimatorTest, ToRotationVectorFromQuaternionAtZero) {
   Eigen::Matrix<double, 3, 1> vector =
       drivetrain::ToRotationVectorFromQuaternion(
           Eigen::Quaternion<double>(
@@ -302,7 +315,7 @@
 }
 
 // Tests that ToRotationVectorFromQuaternion works for a real rotation.
-TEST(RungeKuttaTest, ToRotationVectorFromQuaternion) {
+TEST(DownEstimatorTest, ToRotationVectorFromQuaternion) {
   Eigen::Matrix<double, 3, 1> vector =
       drivetrain::ToRotationVectorFromQuaternion(
           Eigen::Quaternion<double>(
@@ -315,7 +328,7 @@
 
 // Tests that ToRotationVectorFromQuaternion works for a solution with negative
 // coefficients.
-TEST(RungeKuttaTest, ToRotationVectorFromQuaternionNegative) {
+TEST(DownEstimatorTest, ToRotationVectorFromQuaternionNegative) {
   Eigen::Matrix<double, 3, 1> vector =
       drivetrain::ToRotationVectorFromQuaternion(
           Eigen::Quaternion<double>(
@@ -330,7 +343,7 @@
 }
 
 // Tests that ToQuaternionFromRotationVector works for a 0 rotation.
-TEST(RungeKuttaTest, ToQuaternionFromRotationVectorAtZero) {
+TEST(DownEstimatorTest, ToQuaternionFromRotationVectorAtZero) {
   Eigen::Matrix<double, 4, 1> quaternion =
       drivetrain::ToQuaternionFromRotationVector(Eigen::Vector3d::Zero());
 
@@ -342,7 +355,7 @@
 }
 
 // Tests that ToQuaternionFromRotationVector works for a real rotation.
-TEST(RungeKuttaTest, ToQuaternionFromRotationVector) {
+TEST(DownEstimatorTest, ToQuaternionFromRotationVector) {
   Eigen::Matrix<double, 4, 1> quaternion =
       drivetrain::ToQuaternionFromRotationVector(Eigen::Vector3d::UnitX() *
                                                  M_PI * 0.5);
@@ -358,7 +371,7 @@
 
 // Tests that ToQuaternionFromRotationVector correctly clips a rotation vector
 // that is too large in magnitude.
-TEST(RungeKuttaTest, ToQuaternionFromLargeRotationVector) {
+TEST(DownEstimatorTest, ToQuaternionFromLargeRotationVector) {
   const double kMaxAngle = 2.0;
   const Eigen::Vector3d rotation_vector =
       Eigen::Vector3d::UnitX() * kMaxAngle * 2.0;
@@ -372,7 +385,7 @@
 
 // Tests that ToQuaternionFromRotationVector and ToRotationVectorFromQuaternion
 // works for random rotations.
-TEST(RungeKuttaTest, RandomQuaternions) {
+TEST(DownEstimatorTest, RandomQuaternions) {
   std::mt19937 generator(aos::testing::RandomSeed());
   std::uniform_real_distribution<double> random_scalar(-1.0, 1.0);
 
diff --git a/y2020/control_loops/drivetrain/drivetrain_base.cc b/y2020/control_loops/drivetrain/drivetrain_base.cc
index c3597f5..62679bc 100644
--- a/y2020/control_loops/drivetrain/drivetrain_base.cc
+++ b/y2020/control_loops/drivetrain/drivetrain_base.cc
@@ -53,6 +53,9 @@
       1.2 /* quickturn_wheel_multiplier */,
       1.2 /* wheel_multiplier */,
       true /*pistol_grip_shift_enables_line_follow*/,
+      (Eigen::Matrix<double, 3, 3>() << 0.0, -1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0,
+       1.0)
+          .finished() /*imu_transform*/,
   };
 
   return kDrivetrainConfig;