Gain-schedule localizer corrections for auto vs not

Change-Id: I2adc1f3e1413efc6e2723cb3b6ed532987f4772a
Signed-off-by: James Kuszmaul <jabukuszmaul@gmail.com>
diff --git a/y2022/localizer/BUILD b/y2022/localizer/BUILD
index edc3791..243259a 100644
--- a/y2022/localizer/BUILD
+++ b/y2022/localizer/BUILD
@@ -105,6 +105,7 @@
         "//frc971/control_loops/drivetrain:drivetrain_status_fbs",
         "//frc971/control_loops/drivetrain:improved_down_estimator",
         "//frc971/control_loops/drivetrain:localizer_fbs",
+        "//frc971/input:joystick_state_fbs",
         "//frc971/wpilib:imu_batch_fbs",
         "//frc971/wpilib:imu_fbs",
         "//frc971/zeroing:imu_zeroer",
diff --git a/y2022/localizer/localizer.cc b/y2022/localizer/localizer.cc
index 783d368..8ca5395 100644
--- a/y2022/localizer/localizer.cc
+++ b/y2022/localizer/localizer.cc
@@ -642,7 +642,11 @@
   H_model(1, kY) = 1.0;
   H_accel(0, kX) = 1.0;
   H_accel(1, kY) = 1.0;
-  R.diagonal() << 1e-2, 1e-2;
+  if (aggressive_corrections_) {
+    R.diagonal() << 1e-2, 1e-2;
+  } else {
+    R.diagonal() << 1e-0, 1e-0;
+  }
 
   const Eigen::Matrix<double, kNModelStates, 2> K_model =
       P_model_ * H_model.transpose() *
@@ -886,6 +890,8 @@
           event_loop_
               ->MakeFetcher<y2022::control_loops::superstructure::Status>(
                   "/superstructure")),
+      joystick_state_fetcher_(
+          event_loop_->MakeFetcher<aos::JoystickState>("/roborio/aos")),
       zeroer_(zeroing::ImuZeroer::FaultBehavior::kTemporary),
       left_encoder_(-DrivetrainWrapPeriod() / 2.0, DrivetrainWrapPeriod()),
       right_encoder_(-DrivetrainWrapPeriod() / 2.0, DrivetrainWrapPeriod()) {
@@ -926,39 +932,44 @@
             absl::StrCat("/", kPisToUse[camera_index], "/camera"));
   }
   aos::TimerHandler *estimate_timer = event_loop_->AddTimer([this]() {
-    for (size_t camera_index = 0; camera_index < kPisToUse.size();
-         ++camera_index) {
-      if (model_based_.NumQueuedImageDebugs() ==
-              ModelBasedLocalizer::kDebugBufferSize ||
-          (last_visualization_send_ + kMinVisualizationPeriod <
-           event_loop_->monotonic_now())) {
-        auto builder = visualization_sender_.MakeBuilder();
-        visualization_sender_.CheckOk(
-            builder.Send(model_based_.PopulateVisualization(builder.fbb())));
-      }
-      if (target_estimate_fetchers_[camera_index].Fetch()) {
-        const std::optional<aos::monotonic_clock::duration> monotonic_offset =
-            ClockOffset(kPisToUse[camera_index]);
-        if (!monotonic_offset.has_value()) {
-          continue;
+      joystick_state_fetcher_.Fetch();
+      const bool maybe_in_auto = (joystick_state_fetcher_.get() != nullptr)
+                                     ? joystick_state_fetcher_->autonomous()
+                                     : true;
+      model_based_.set_use_aggressive_image_corrections(!maybe_in_auto);
+      for (size_t camera_index = 0; camera_index < kPisToUse.size();
+           ++camera_index) {
+        if (model_based_.NumQueuedImageDebugs() ==
+                ModelBasedLocalizer::kDebugBufferSize ||
+            (last_visualization_send_ + kMinVisualizationPeriod <
+             event_loop_->monotonic_now())) {
+          auto builder = visualization_sender_.MakeBuilder();
+          visualization_sender_.CheckOk(
+              builder.Send(model_based_.PopulateVisualization(builder.fbb())));
         }
-        // TODO(james): Get timestamp from message contents.
-        aos::monotonic_clock::time_point capture_time(
-            target_estimate_fetchers_[camera_index]
-                .context()
-                .monotonic_remote_time -
-            monotonic_offset.value());
-        if (capture_time > target_estimate_fetchers_[camera_index]
-                               .context()
-                               .monotonic_event_time) {
-          model_based_.TallyRejection(RejectionReason::IMAGE_FROM_FUTURE);
-          continue;
+        if (target_estimate_fetchers_[camera_index].Fetch()) {
+          const std::optional<aos::monotonic_clock::duration> monotonic_offset =
+              ClockOffset(kPisToUse[camera_index]);
+          if (!monotonic_offset.has_value()) {
+            continue;
+          }
+          // TODO(james): Get timestamp from message contents.
+          aos::monotonic_clock::time_point capture_time(
+              target_estimate_fetchers_[camera_index]
+                  .context()
+                  .monotonic_remote_time -
+              monotonic_offset.value());
+          if (capture_time > target_estimate_fetchers_[camera_index]
+                                 .context()
+                                 .monotonic_event_time) {
+            model_based_.TallyRejection(RejectionReason::IMAGE_FROM_FUTURE);
+            continue;
+          }
+          capture_time -= pico_offset_error_;
+          model_based_.HandleImageMatch(
+              capture_time, target_estimate_fetchers_[camera_index].get(),
+              camera_index);
         }
-        capture_time -= pico_offset_error_;
-        model_based_.HandleImageMatch(
-            capture_time, target_estimate_fetchers_[camera_index].get(),
-            camera_index);
-      }
     }
   });
   event_loop_->OnRun([this, estimate_timer]() {
diff --git a/y2022/localizer/localizer.h b/y2022/localizer/localizer.h
index f8205d7..5620a4d 100644
--- a/y2022/localizer/localizer.h
+++ b/y2022/localizer/localizer.h
@@ -9,6 +9,7 @@
 #include "aos/network/message_bridge_server_generated.h"
 #include "aos/time/time.h"
 #include "frc971/control_loops/drivetrain/drivetrain_output_generated.h"
+#include "frc971/input/joystick_state_generated.h"
 #include "frc971/control_loops/drivetrain/improved_down_estimator.h"
 #include "frc971/control_loops/drivetrain/localizer_generated.h"
 #include "frc971/zeroing/imu_zeroer.h"
@@ -136,6 +137,9 @@
   AccelState accel_state() const { return current_state_.accel_state; };
 
   void set_longitudinal_offset(double offset) { long_offset_ = offset; }
+  void set_use_aggressive_image_corrections(bool aggressive) {
+    aggressive_corrections_ = aggressive;
+  }
 
   void TallyRejection(const RejectionReason reason);
 
@@ -267,6 +271,10 @@
   // center, negative = behind center.
   double long_offset_ = -0.15;
 
+  // Whether to use more aggressive corrections on the localizer. Only do this
+  // in teleop, since it can make spline control really jumpy.
+  bool aggressive_corrections_ = false;
+
   double last_residual_ = 0.0;
   double filtered_residual_ = 0.0;
   Eigen::Vector2d filtered_residual_accel_ = Eigen::Vector2d::Zero();
@@ -334,6 +342,7 @@
       target_estimate_fetchers_;
   aos::Fetcher<y2022::control_loops::superstructure::Status>
       superstructure_fetcher_;
+  aos::Fetcher<aos::JoystickState> joystick_state_fetcher_;
   zeroing::ImuZeroer zeroer_;
   aos::monotonic_clock::time_point last_output_send_ =
       aos::monotonic_clock::min_time;