Merge "Make skew in Ekf match skew from cameras"
diff --git a/frc971/control_loops/drivetrain/BUILD b/frc971/control_loops/drivetrain/BUILD
index c99053f..a4d2ca6 100644
--- a/frc971/control_loops/drivetrain/BUILD
+++ b/frc971/control_loops/drivetrain/BUILD
@@ -65,6 +65,13 @@
     ],
 )
 
+queue_library(
+    name = "localizer_queue",
+    srcs = [
+        "localizer.q",
+    ],
+)
+
 cc_library(
     name = "localizer",
     hdrs = ["localizer.h"],
@@ -262,6 +269,7 @@
         ":gear",
         ":line_follow_drivetrain",
         ":localizer",
+        ":localizer_queue",
         ":polydrivetrain",
         ":splinedrivetrain",
         ":ssdrivetrain",
@@ -308,6 +316,7 @@
     deps = [
         ":drivetrain_config",
         ":drivetrain_lib",
+        ":localizer_queue",
         ":drivetrain_queue",
         ":drivetrain_test_lib",
         "//aos:queues",
diff --git a/frc971/control_loops/drivetrain/drivetrain.cc b/frc971/control_loops/drivetrain/drivetrain.cc
index cab8043..1b37612 100644
--- a/frc971/control_loops/drivetrain/drivetrain.cc
+++ b/frc971/control_loops/drivetrain/drivetrain.cc
@@ -36,6 +36,8 @@
     : aos::controls::ControlLoop<::frc971::control_loops::DrivetrainQueue>(
           event_loop, name),
       dt_config_(dt_config),
+      localizer_control_fetcher_(event_loop->MakeFetcher<LocalizerControl>(
+          ".frc971.control_loops.drivetrain.localizer_control")),
       localizer_(localizer),
       kf_(dt_config_.make_kf_drivetrain_loop()),
       dt_openloop_(dt_config_, &kf_),
@@ -232,6 +234,15 @@
     Y << position->left_encoder, position->right_encoder, last_gyro_rate_,
         last_accel_;
     kf_.Correct(Y);
+    // If we get a new message setting the absolute position, then reset the
+    // localizer.
+    // TODO(james): Use a watcher (instead of a fetcher) once we support it in
+    // simulation.
+    if (localizer_control_fetcher_.Fetch()) {
+      localizer_->ResetPosition(localizer_control_fetcher_->x,
+                                localizer_control_fetcher_->y,
+                                localizer_control_fetcher_->theta);
+    }
     localizer_->Update({last_last_left_voltage_, last_last_right_voltage_},
                        monotonic_now, position->left_encoder,
                        position->right_encoder, last_gyro_rate_, last_accel_);
diff --git a/frc971/control_loops/drivetrain/drivetrain.h b/frc971/control_loops/drivetrain/drivetrain.h
index 800199b..f07d840 100644
--- a/frc971/control_loops/drivetrain/drivetrain.h
+++ b/frc971/control_loops/drivetrain/drivetrain.h
@@ -13,6 +13,7 @@
 #include "frc971/control_loops/drivetrain/polydrivetrain.h"
 #include "frc971/control_loops/drivetrain/ssdrivetrain.h"
 #include "frc971/control_loops/drivetrain/line_follow_drivetrain.h"
+#include "frc971/control_loops/drivetrain/localizer.q.h"
 #include "frc971/control_loops/drivetrain/splinedrivetrain.h"
 
 namespace frc971 {
@@ -45,6 +46,7 @@
 
   const DrivetrainConfig<double> dt_config_;
 
+  ::aos::Fetcher<LocalizerControl> localizer_control_fetcher_;
   LocalizerInterface *localizer_;
 
   StateFeedbackLoop<7, 2, 4> kf_;
diff --git a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
index 97aaba5..5216916 100644
--- a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
+++ b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
@@ -17,6 +17,7 @@
 #include "frc971/control_loops/drivetrain/drivetrain.q.h"
 #include "frc971/control_loops/drivetrain/drivetrain_config.h"
 #include "frc971/control_loops/drivetrain/drivetrain_test_lib.h"
+#include "frc971/control_loops/drivetrain/localizer.q.h"
 #include "frc971/queues/gyro.q.h"
 
 DEFINE_bool(plot, false, "If true, plot");
@@ -755,6 +756,22 @@
   EXPECT_NEAR(0.0, drivetrain_motor_plant_.GetPosition().y(), 1e-4);
 }
 
+// Tests that we can reset the localizer to a new position.
+TEST_F(DrivetrainTest, ResetLocalizer) {
+  EXPECT_EQ(0.0, localizer_.x());
+  EXPECT_EQ(0.0, localizer_.y());
+  EXPECT_EQ(0.0, localizer_.theta());
+  ::aos::Queue<LocalizerControl> localizer_queue(
+      ".frc971.control_loops.drivetrain.localizer_control");
+  ASSERT_TRUE(
+      localizer_queue.MakeWithBuilder().x(9.0).y(7.0).theta(1.0).Send());
+  RunIteration();
+
+  EXPECT_EQ(9.0, localizer_.x());
+  EXPECT_EQ(7.0, localizer_.y());
+  EXPECT_EQ(1.0, localizer_.theta());
+}
+
 ::aos::controls::HVPolytope<2, 4, 4> MakeBox(double x1_min, double x1_max,
                                              double x2_min, double x2_max) {
   Eigen::Matrix<double, 4, 2> box_H;
diff --git a/frc971/control_loops/drivetrain/localizer.h b/frc971/control_loops/drivetrain/localizer.h
index 5d5b58d..af07089 100644
--- a/frc971/control_loops/drivetrain/localizer.h
+++ b/frc971/control_loops/drivetrain/localizer.h
@@ -40,6 +40,8 @@
                       ::aos::monotonic_clock::time_point now,
                       double left_encoder, double right_encoder,
                       double gyro_rate, double longitudinal_accelerometer) = 0;
+  // Reset the absolute position of the estimator.
+  virtual void ResetPosition(double x, double y, double theta) = 0;
   // There are several subtly different norms floating around for state
   // matrices. In order to avoid that mess, we jus tprovide direct accessors for
   // the values that most people care about.
@@ -92,6 +94,13 @@
     ekf_.UpdateEncodersAndGyro(left_encoder, right_encoder, gyro_rate, U, now);
   }
 
+  void ResetPosition(double x, double y, double theta) override {
+    ekf_.ResetInitialState(
+        ::aos::monotonic_clock::now(),
+        (Ekf::State() << x, y, theta, 0, 0, 0, 0, 0, 0, 0).finished(),
+        ekf_.P());
+  };
+
   double x() const override { return ekf_.X_hat(StateIdx::kX); }
   double y() const override { return ekf_.X_hat(StateIdx::kY); }
   double theta() const override { return ekf_.X_hat(StateIdx::kTheta); }
diff --git a/frc971/control_loops/drivetrain/localizer.q b/frc971/control_loops/drivetrain/localizer.q
new file mode 100644
index 0000000..1ae0adc
--- /dev/null
+++ b/frc971/control_loops/drivetrain/localizer.q
@@ -0,0 +1,11 @@
+package frc971.control_loops.drivetrain;
+
+// Allows you to reset the state of the localizer to a specific position on the
+// field.
+message LocalizerControl {
+  float x;      // X position, meters
+  float y;      // Y position, meters
+  float theta;  // heading, radians
+};
+
+queue LocalizerControl localizer_control;
diff --git a/y2019/control_loops/drivetrain/event_loop_localizer.h b/y2019/control_loops/drivetrain/event_loop_localizer.h
index eb03627..773410f 100644
--- a/y2019/control_loops/drivetrain/event_loop_localizer.h
+++ b/y2019/control_loops/drivetrain/event_loop_localizer.h
@@ -33,6 +33,9 @@
       ::aos::EventLoop *event_loop);
 
   void Reset(const Localizer::State &state);
+  void ResetPosition(double x, double y, double theta) override {
+    Reset((Localizer::State() << x, y, theta, 0, 0, 0, 0, 0, 0, 0).finished());
+  }
 
   void Update(const ::Eigen::Matrix<double, 2, 1> &U,
               ::aos::monotonic_clock::time_point now, double left_encoder,
diff --git a/y2019/control_loops/drivetrain/localized_drivetrain_test.cc b/y2019/control_loops/drivetrain/localized_drivetrain_test.cc
index 8c699c5..3dde98d 100644
--- a/y2019/control_loops/drivetrain/localized_drivetrain_test.cc
+++ b/y2019/control_loops/drivetrain/localized_drivetrain_test.cc
@@ -238,6 +238,28 @@
   EXPECT_NEAR(true_state(4, 0), localizer_.right_velocity(), 1e-5);
 }
 
+// Tests that when we reset the position of the localizer via the queue to
+// compensate for the disturbance, the estimator behaves perfectly.
+TEST_F(LocalizedDrivetrainTest, ResetLocalizer) {
+  set_enable_cameras(false);
+  (*drivetrain_motor_plant_.mutable_state())(0, 0) += 0.05;
+  my_drivetrain_queue_.goal.MakeWithBuilder()
+      .controller_type(1)
+      .left_goal(-1.0)
+      .right_goal(1.0)
+      .Send();
+  ::aos::Queue<::frc971::control_loops::drivetrain::LocalizerControl>
+      localizer_queue(".frc971.control_loops.drivetrain.localizer_control");
+  ASSERT_TRUE(localizer_queue.MakeWithBuilder()
+                  .x(drivetrain_motor_plant_.state().x())
+                  .y(drivetrain_motor_plant_.state().y())
+                  .theta(drivetrain_motor_plant_.state()(2, 0))
+                  .Send());
+  RunForTime(chrono::seconds(3));
+  VerifyNearGoal();
+  VerifyEstimatorAccurate(1e-5);
+}
+
 // Tests that, when a small error in X is introduced, the camera corrections do
 // correct for it.
 TEST_F(LocalizedDrivetrainTest, CameraUpdate) {