diff --git a/y2022/control_loops/drivetrain/localizer.cc b/y2022/control_loops/drivetrain/localizer.cc
index 79ced2d..0aa4456 100644
--- a/y2022/control_loops/drivetrain/localizer.cc
+++ b/y2022/control_loops/drivetrain/localizer.cc
@@ -11,6 +11,7 @@
     : event_loop_(event_loop),
       dt_config_(dt_config),
       ekf_(dt_config),
+      observations_(&ekf_),
       localizer_output_fetcher_(
           event_loop_->MakeFetcher<frc971::controls::LocalizerOutput>(
               "/localizer")),
@@ -48,7 +49,7 @@
       joystick_state_fetcher_->autonomous()) {
     // TODO(james): This is an inelegant way to avoid having the localizer mess
     // up splines. Do better.
-    //return;
+    // return;
   }
   if (localizer_output_fetcher_.Fetch()) {
     clock_offset_fetcher_.Fetch();
@@ -89,11 +90,6 @@
       }
     }
 
-    Eigen::Matrix<float, HybridEkf::kNOutputs, HybridEkf::kNStates> H;
-    H.setZero();
-    H(0, StateIdx::kX) = 1;
-    H(1, StateIdx::kY) = 1;
-    H(2, StateIdx::kTheta) = 1;
     const Eigen::Vector3f Z{
         static_cast<float>(localizer_output_fetcher_->x()),
         static_cast<float>(localizer_output_fetcher_->y()),
@@ -101,15 +97,8 @@
     Eigen::Matrix3f R = Eigen::Matrix3f::Zero();
     R.diagonal() << 0.01, 0.01, 1e-4;
     const Input U_correct = ekf_.MostRecentInput();
-    ekf_.Correct(
-        Eigen::Vector3f::Zero(), &U_correct, {},
-        [H, state_at_capture, Z](const State &,
-                                 const Input &) -> Eigen::Vector3f {
-          Eigen::Vector3f error = H * state_at_capture.value() - Z;
-          error(2) = aos::math::NormalizeAngle(error(2));
-          return error;
-        },
-        [H](const State &) { return H; }, R, now);
+    observations_.CorrectKnownH(Eigen::Vector3f::Zero(), &U_correct,
+                                Corrector(state_at_capture.value(), Z), R, now);
   }
 }
 
diff --git a/y2022/control_loops/drivetrain/localizer.h b/y2022/control_loops/drivetrain/localizer.h
index 0d2673b..77b29eb 100644
--- a/y2022/control_loops/drivetrain/localizer.h
+++ b/y2022/control_loops/drivetrain/localizer.h
@@ -4,11 +4,11 @@
 #include <string_view>
 
 #include "aos/events/event_loop.h"
+#include "aos/network/message_bridge_server_generated.h"
 #include "frc971/control_loops/drivetrain/hybrid_ekf.h"
 #include "frc971/control_loops/drivetrain/localizer.h"
-#include "y2022/localizer/localizer_output_generated.h"
-#include "aos/network/message_bridge_server_generated.h"
 #include "frc971/input/joystick_state_generated.h"
+#include "y2022/localizer/localizer_output_generated.h"
 
 namespace y2022 {
 namespace control_loops {
@@ -63,9 +63,34 @@
   }
 
  private:
+  class Corrector : public HybridEkf::ExpectedObservationFunctor {
+   public:
+    Corrector(const State &state_at_capture, const Eigen::Vector3f &Z)
+        : state_at_capture_(state_at_capture), Z_(Z) {
+      H_.setZero();
+      H_(0, StateIdx::kX) = 1;
+      H_(1, StateIdx::kY) = 1;
+      H_(2, StateIdx::kTheta) = 1;
+    }
+    Output H(const State &, const Input &) final {
+      Eigen::Vector3f error = H_ * state_at_capture_ - Z_;
+      error(2) = aos::math::NormalizeAngle(error(2));
+      return error;
+    }
+    Eigen::Matrix<float, HybridEkf::kNOutputs, HybridEkf::kNStates> DHDX(
+        const State &) final {
+      return H_;
+    }
+
+   private:
+    Eigen::Matrix<float, HybridEkf::kNOutputs, HybridEkf::kNStates> H_;
+    State state_at_capture_;
+    Eigen::Vector3f Z_;
+  };
   aos::EventLoop *const event_loop_;
   const frc971::control_loops::drivetrain::DrivetrainConfig<double> dt_config_;
   HybridEkf ekf_;
+  HybridEkf::ExpectedObservationAllocator<Corrector> observations_;
 
   aos::Fetcher<frc971::controls::LocalizerOutput> localizer_output_fetcher_;
   aos::Fetcher<aos::JoystickState> joystick_state_fetcher_;
diff --git a/y2022/control_loops/drivetrain/localizer_test.cc b/y2022/control_loops/drivetrain/localizer_test.cc
index 1e33826..77f3988 100644
--- a/y2022/control_loops/drivetrain/localizer_test.cc
+++ b/y2022/control_loops/drivetrain/localizer_test.cc
@@ -10,12 +10,13 @@
 #include "frc971/control_loops/drivetrain/drivetrain.h"
 #include "frc971/control_loops/drivetrain/drivetrain_test_lib.h"
 #include "frc971/control_loops/team_number_test_environment.h"
-#include "y2022/localizer/localizer_output_generated.h"
 #include "gtest/gtest.h"
 #include "y2022/control_loops/drivetrain/drivetrain_base.h"
+#include "y2022/localizer/localizer_output_generated.h"
 
 DEFINE_string(output_folder, "",
               "If set, logs all channels to the provided logfile.");
+DECLARE_bool(die_on_malloc);
 
 namespace y2022 {
 namespace control_loops {
@@ -31,7 +32,7 @@
   DrivetrainConfig<double> config = GetDrivetrainConfig();
   return config;
 }
-}
+}  // namespace
 
 namespace chrono = std::chrono;
 using aos::monotonic_clock;
@@ -74,6 +75,7 @@
         drivetrain_plant_(drivetrain_plant_event_loop_.get(),
                           drivetrain_plant_imu_event_loop_.get(), dt_config_,
                           std::chrono::microseconds(500)) {
+    FLAGS_die_on_malloc = true;
     set_team_id(frc971::control_loops::testing::kTeamNumber);
     set_battery_voltage(12.0);
 
@@ -95,6 +97,7 @@
           output_builder.add_x(drivetrain_plant_.state()(0));
           output_builder.add_y(drivetrain_plant_.state()(1));
           output_builder.add_theta(drivetrain_plant_.state()(2));
+          builder.CheckOk(builder.Send(output_builder.Finish()));
         })
         ->Setup(imu_test_event_loop_->monotonic_now(),
                 std::chrono::milliseconds(5));