Respect WasReset() indicator in the drivetrain

At least, pay attention to it for the localizer/down estimator.

Change-Id: I9cf8720c7ad9b2bd1f16b8e65acf824b10f3c7ed
diff --git a/frc971/control_loops/drivetrain/drivetrain.cc b/frc971/control_loops/drivetrain/drivetrain.cc
index a822ac4..f517e83 100644
--- a/frc971/control_loops/drivetrain/drivetrain.cc
+++ b/frc971/control_loops/drivetrain/drivetrain.cc
@@ -107,6 +107,19 @@
     has_been_enabled_ = true;
   }
 
+  if (WasReset()) {
+    // If all the sensors got reset (e.g., due to wpilib_interface restarting),
+    // reset the localizer and down estimator to avoid weird jumps in the
+    // filters.
+    down_estimator_.Reset();
+    // Just reset the localizer to the current state, except for the encoders.
+    LocalizerInterface::Ekf::State X_hat = localizer_->Xhat();
+    X_hat(LocalizerInterface::StateIdx::kLeftEncoder) = position->left_encoder();
+    X_hat(LocalizerInterface::StateIdx::kRightEncoder) =
+        position->right_encoder();
+    localizer_->Reset(monotonic_now, X_hat);
+  }
+
   // TODO(austin): Put gear detection logic here.
   switch (dt_config_.shifter_type) {
     case ShifterType::SIMPLE_SHIFTER:
diff --git a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
index faed4af..9ce3a48 100644
--- a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
+++ b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
@@ -1463,6 +1463,47 @@
   EXPECT_EQ(1.0, localizer_.theta());
 }
 
+// Tests that if wpilib_interface restarts, the drivetrain handles it.
+TEST_F(DrivetrainTest, ResetDrivetrain) {
+  SetEnabled(true);
+  EXPECT_EQ(0.0, localizer_.x());
+  EXPECT_EQ(0.0, localizer_.y());
+  EXPECT_EQ(0.0, localizer_.theta());
+
+  {
+    auto builder = drivetrain_goal_sender_.MakeBuilder();
+    Goal::Builder goal_builder = builder.MakeBuilder<Goal>();
+    goal_builder.add_controller_type(ControllerType::MOTION_PROFILE);
+    goal_builder.add_left_goal(4.0);
+    goal_builder.add_right_goal(4.0);
+    EXPECT_TRUE(builder.Send(goal_builder.Finish()));
+  }
+
+  RunFor(chrono::seconds(2));
+
+  const double x_pos = localizer_.x();
+  // The dead-reckoned X position in these sorts of situations tends to be
+  // relatively poor with the current localizer, since it ignores voltage inputs
+  // and models the robot's speed as always decaying towards zero, hence the
+  // large tolerances on the x position.
+  EXPECT_NEAR(4.0, x_pos, 2.0);
+  EXPECT_NEAR(0.0, localizer_.y(), 1e-5);
+  EXPECT_NEAR(0.0, localizer_.theta(), 1e-5);
+  EXPECT_NEAR(4.0, localizer_.left_encoder(), 1e-3);
+  EXPECT_NEAR(4.0, localizer_.right_encoder(), 1e-3);
+
+  SimulateSensorReset();
+  drivetrain_plant_.Reinitialize();
+
+  RunFor(dt());
+
+  EXPECT_EQ(x_pos, localizer_.x());
+  EXPECT_NEAR(0.0, localizer_.y(), 1e-5);
+  EXPECT_NEAR(0.0, localizer_.theta(), 1e-5);
+  EXPECT_NEAR(0.0, localizer_.left_encoder(), 1e-5);
+  EXPECT_NEAR(0.0, localizer_.right_encoder(), 1e-5);
+}
+
 // TODO(austin): Make sure the profile reset code when we disable works.
 
 }  // namespace testing
diff --git a/frc971/control_loops/drivetrain/drivetrain_test_lib.cc b/frc971/control_loops/drivetrain/drivetrain_test_lib.cc
index 3ed8d00..c7ef9fd 100644
--- a/frc971/control_loops/drivetrain/drivetrain_test_lib.cc
+++ b/frc971/control_loops/drivetrain/drivetrain_test_lib.cc
@@ -174,7 +174,7 @@
   const double left_encoder = GetLeftPosition();
   const double right_encoder = GetRightPosition();
 
-  {
+  if (send_messages_) {
     ::aos::Sender<::frc971::control_loops::drivetrain::Position>::Builder
         builder = drivetrain_position_sender_.MakeBuilder();
     frc971::control_loops::drivetrain::Position::Builder position_builder =
@@ -188,6 +188,9 @@
 }
 
 void DrivetrainSimulation::SendImuMessage() {
+  if (!send_messages_) {
+    return;
+  }
   auto builder = imu_sender_.MakeBuilder();
   frc971::IMUValues::Builder imu_builder =
       builder.MakeBuilder<frc971::IMUValues>();
@@ -228,10 +231,16 @@
 void DrivetrainSimulation::Simulate() {
   last_left_position_ = drivetrain_plant_.Y(0, 0);
   last_right_position_ = drivetrain_plant_.Y(1, 0);
-  EXPECT_TRUE(drivetrain_output_fetcher_.Fetch());
   ::Eigen::Matrix<double, 2, 1> U = last_U_;
-  last_U_ << drivetrain_output_fetcher_->left_voltage(),
-      drivetrain_output_fetcher_->right_voltage();
+  if (send_messages_) {
+    EXPECT_TRUE(drivetrain_output_fetcher_.Fetch());
+    last_U_ << drivetrain_output_fetcher_->left_voltage(),
+        drivetrain_output_fetcher_->right_voltage();
+    left_gear_high_ = drivetrain_output_fetcher_->left_high();
+    right_gear_high_ = drivetrain_output_fetcher_->right_high();
+  } else {
+    U = U.Zero();
+  }
   {
     robot_state_fetcher_.Fetch();
     const double scalar = robot_state_fetcher_.get()
@@ -239,8 +248,6 @@
                               : 1.0;
     last_U_ *= scalar;
   }
-  left_gear_high_ = drivetrain_output_fetcher_->left_high();
-  right_gear_high_ = drivetrain_output_fetcher_->right_high();
 
   if (left_gear_high_) {
     if (right_gear_high_) {
diff --git a/frc971/control_loops/drivetrain/drivetrain_test_lib.h b/frc971/control_loops/drivetrain/drivetrain_test_lib.h
index f749626..800b758 100644
--- a/frc971/control_loops/drivetrain/drivetrain_test_lib.h
+++ b/frc971/control_loops/drivetrain/drivetrain_test_lib.h
@@ -73,6 +73,12 @@
 
   void MaybePlot();
 
+  // Set whether we should send out the drivetrain Position and IMU messages
+  // (this will keep sending the "truth" message).
+  void set_send_messages(const bool send_messages) {
+    send_messages_ = send_messages;
+  }
+
  private:
   // Sends out the position queue messages.
   void SendPositionMessage();
@@ -123,6 +129,8 @@
   ::std::vector<double> actual_y_;
   ::std::vector<double> trajectory_x_;
   ::std::vector<double> trajectory_y_;
+
+  bool send_messages_ = true;
 };
 
 }  // namespace testing
diff --git a/frc971/control_loops/drivetrain/hybrid_ekf.h b/frc971/control_loops/drivetrain/hybrid_ekf.h
index 4fd2d43..61408aa 100644
--- a/frc971/control_loops/drivetrain/hybrid_ekf.h
+++ b/frc971/control_loops/drivetrain/hybrid_ekf.h
@@ -703,7 +703,7 @@
   // 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.
+  // drivetrain Kalman Filter should have 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);
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator.cc b/frc971/control_loops/drivetrain/improved_down_estimator.cc
index 665bf32..8d2592b 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator.cc
+++ b/frc971/control_loops/drivetrain/improved_down_estimator.cc
@@ -9,6 +9,30 @@
 namespace control_loops {
 namespace drivetrain {
 
+void QuaternionUkf::Reset() {
+  // The various noise matrices are tuned to provide values that seem
+  // reasonable given our current setup (using the ADIS16470 for
+  // measurements).
+  R_.setIdentity();
+  R_ /= std::pow(1000.0, 2);
+
+  Q_.setIdentity();
+  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);
+
+  P_.setIdentity();
+  P_ /= 1000.0;
+
+  pos_vel_.setZero();
+  for (auto &last_accel : last_accels_) {
+    last_accel.setZero();
+  }
+
+  last_yaw_rates_.fill(0);
+}
+
 void QuaternionUkf::Predict(const Eigen::Matrix<double, 3, 1> &U,
                             const Eigen::Matrix<double, 3, 1> &measurement,
                             const aos::monotonic_clock::duration dt) {
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator.h b/frc971/control_loops/drivetrain/improved_down_estimator.h
index 457bbef..79e9b7f 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator.h
+++ b/frc971/control_loops/drivetrain/improved_down_estimator.h
@@ -59,29 +59,11 @@
   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_ /= std::pow(1000.0, 2);
-
-    Q_.setIdentity();
-    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);
-
-    P_.setIdentity();
-    P_ /= 1000.0;
-
-    pos_vel_.setZero();
-    for (auto &last_accel : last_accels_) {
-      last_accel.setZero();
-    }
-
-    last_yaw_rates_.fill(0);
+    Reset();
   }
 
+  void Reset();
+
   // Handles updating the state of the UKF, given the gyro and accelerometer
   // measurements. Given the design of the filter, U is the x/y/z gyro
   // measurements and measurement is the accelerometer x/y/z measurements.
diff --git a/frc971/control_loops/drivetrain/localizer.h b/frc971/control_loops/drivetrain/localizer.h
index 6677145..67d681e 100644
--- a/frc971/control_loops/drivetrain/localizer.h
+++ b/frc971/control_loops/drivetrain/localizer.h
@@ -55,6 +55,8 @@
                       ::aos::monotonic_clock::time_point now,
                       double left_encoder, double right_encoder,
                       double gyro_rate, const Eigen::Vector3d &accel) = 0;
+  virtual void Reset(aos::monotonic_clock::time_point t,
+                     const Ekf::State &state) = 0;
   // Reset the absolute position of the estimator.
   virtual void ResetPosition(::aos::monotonic_clock::time_point t, double x,
                              double y, double theta, double theta_uncertainty,
@@ -148,6 +150,11 @@
                                now);
   }
 
+  void Reset(aos::monotonic_clock::time_point t,
+             const HybridEkf<double>::State &state) override {
+    ekf_.ResetInitialState(t, state, ekf_.P());
+  }
+
   void ResetPosition(::aos::monotonic_clock::time_point t, double x, double y,
                      double theta, double /*theta_override*/,
                      bool /*reset_theta*/) override {