Respect WasReset() indicator in the drivetrain

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

Change-Id: I9cf8720c7ad9b2bd1f16b8e65acf824b10f3c7ed
diff --git a/aos/events/simulated_event_loop.cc b/aos/events/simulated_event_loop.cc
index c429eef..2eaffa0 100644
--- a/aos/events/simulated_event_loop.cc
+++ b/aos/events/simulated_event_loop.cc
@@ -346,7 +346,9 @@
       return std::make_pair(false, monotonic_clock::min_time);
     }
 
-    CHECK(!fell_behind_) << ": Got behind";
+    CHECK(!fell_behind_) << ": Got behind on "
+                         << configuration::StrippedChannelToString(
+                                simulated_channel_->channel());
 
     SetMsg(msgs_.front());
     msgs_.pop_front();
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 {
diff --git a/y2019/control_loops/drivetrain/event_loop_localizer.cc b/y2019/control_loops/drivetrain/event_loop_localizer.cc
index 16cd723..634cedb 100644
--- a/y2019/control_loops/drivetrain/event_loop_localizer.cc
+++ b/y2019/control_loops/drivetrain/event_loop_localizer.cc
@@ -49,6 +49,11 @@
 }
 
 void EventLoopLocalizer::Reset(::aos::monotonic_clock::time_point now,
+                               const Localizer::State &state) {
+  localizer_.ResetInitialState(now, state, localizer_.P());
+}
+
+void EventLoopLocalizer::Reset(::aos::monotonic_clock::time_point now,
                                const Localizer::State &state,
                                double theta_uncertainty) {
   Localizer::StateSquare newP = localizer_.P();
diff --git a/y2019/control_loops/drivetrain/event_loop_localizer.h b/y2019/control_loops/drivetrain/event_loop_localizer.h
index 5305b8d..a78f3a2 100644
--- a/y2019/control_loops/drivetrain/event_loop_localizer.h
+++ b/y2019/control_loops/drivetrain/event_loop_localizer.h
@@ -34,6 +34,8 @@
           &dt_config);
 
   void Reset(::aos::monotonic_clock::time_point t,
+             const Localizer::State &state) override;
+  void Reset(::aos::monotonic_clock::time_point t,
              const Localizer::State &state, double theta_uncertainty);
   void ResetPosition(::aos::monotonic_clock::time_point t, double x, double y,
                      double theta, double theta_uncertainty,
diff --git a/y2020/control_loops/drivetrain/localizer.cc b/y2020/control_loops/drivetrain/localizer.cc
index 7d7c83f..79035df 100644
--- a/y2020/control_loops/drivetrain/localizer.cc
+++ b/y2020/control_loops/drivetrain/localizer.cc
@@ -95,6 +95,16 @@
   target_selector_.set_has_target(false);
 }
 
+void Localizer::Reset(
+    aos::monotonic_clock::time_point t,
+    const frc971::control_loops::drivetrain::HybridEkf<double>::State &state) {
+  // Go through and clear out all of the fetchers so that we don't get behind.
+  for (auto &fetcher : image_fetchers_) {
+    fetcher.Fetch();
+  }
+  ekf_.ResetInitialState(t, state.cast<float>(), ekf_.P());
+}
+
 void Localizer::HandleSuperstructureStatus(
     const y2020::control_loops::superstructure::Status &status) {
   CHECK(status.has_turret());
@@ -162,7 +172,10 @@
     LOG(WARNING) << "Got camera frame from the future.";
     return;
   }
-  CHECK(result.has_camera_calibration());
+  if (!result.has_camera_calibration()) {
+    LOG(WARNING) << "Got camera frame without calibration data.";
+    return;
+  }
   // Per the ImageMatchResult specification, we can actually determine whether
   // the camera is the turret camera just from the presence of the
   // turret_extrinsics member.
diff --git a/y2020/control_loops/drivetrain/localizer.h b/y2020/control_loops/drivetrain/localizer.h
index 341f304..5dbad02 100644
--- a/y2020/control_loops/drivetrain/localizer.h
+++ b/y2020/control_loops/drivetrain/localizer.h
@@ -51,9 +51,9 @@
               double right_encoder, double gyro_rate,
               const Eigen::Vector3d &accel) override;
 
-  void Reset(aos::monotonic_clock::time_point t, const State &state) {
-    ekf_.ResetInitialState(t, state, ekf_.P());
-  }
+  void Reset(aos::monotonic_clock::time_point t,
+             const frc971::control_loops::drivetrain::HybridEkf<double>::State
+                 &state) override;
 
   void ResetPosition(aos::monotonic_clock::time_point t, double x, double y,
                      double theta, double /*theta_override*/,
diff --git a/y2020/control_loops/drivetrain/localizer_test.cc b/y2020/control_loops/drivetrain/localizer_test.cc
index dbb4f2f..da3ef17 100644
--- a/y2020/control_loops/drivetrain/localizer_test.cc
+++ b/y2020/control_loops/drivetrain/localizer_test.cc
@@ -206,9 +206,9 @@
   void SetStartingPosition(const Eigen::Matrix<double, 3, 1> &xytheta) {
     *drivetrain_plant_.mutable_state() << xytheta.x(), xytheta.y(),
         xytheta(2, 0), 0.0, 0.0;
-    Eigen::Matrix<float, Localizer::HybridEkf::kNStates, 1> localizer_state;
+    Eigen::Matrix<double, Localizer::HybridEkf::kNStates, 1> localizer_state;
     localizer_state.setZero();
-    localizer_state.block<3, 1>(0, 0) = xytheta.cast<float>();
+    localizer_state.block<3, 1>(0, 0) = xytheta;
     localizer_.Reset(monotonic_now(), localizer_state);
   }
 
@@ -547,6 +547,33 @@
   EXPECT_TRUE(VerifyEstimatorAccurate(1e-2));
 }
 
+// Tests that we don't blow up if we stop getting updates for an extended period
+// of time and fall behind on fetching fron the cameras.
+TEST_F(LocalizedDrivetrainTest, FetchersHandleTimeGap) {
+  set_enable_cameras(true);
+  set_send_delay(std::chrono::seconds(0));
+  event_loop_factory()->set_network_delay(std::chrono::nanoseconds(1));
+  test_event_loop_
+      ->AddTimer([this]() { drivetrain_plant_.set_send_messages(false); })
+      ->Setup(test_event_loop_->monotonic_now());
+  test_event_loop_->AddPhasedLoop(
+      [this](int) {
+        auto builder = camera_sender_.MakeBuilder();
+        ImageMatchResultT image;
+        ASSERT_TRUE(
+            builder.Send(ImageMatchResult::Pack(*builder.fbb(), &image)));
+      },
+      std::chrono::milliseconds(20));
+  test_event_loop_
+      ->AddTimer([this]() {
+        drivetrain_plant_.set_send_messages(true);
+        SimulateSensorReset();
+      })
+      ->Setup(test_event_loop_->monotonic_now() + std::chrono::seconds(10));
+
+  RunFor(chrono::seconds(20));
+}
+
 }  // namespace testing
 }  // namespace drivetrain
 }  // namespace control_loops