[wpimath] Simplify pose estimator (#6705)

wpimath/src/main/java/edu/wpi/first/math/estimator/PoseEstimator.java

@@ -5,22 +5,21 @@
 package edu.wpi.first.math.estimator;
 
 import edu.wpi.first.math.MathSharedStore;
+import edu.wpi.first.math.MathUtil;
 import edu.wpi.first.math.Matrix;
 import edu.wpi.first.math.Nat;
 import edu.wpi.first.math.VecBuilder;
 import edu.wpi.first.math.geometry.Pose2d;
 import edu.wpi.first.math.geometry.Rotation2d;
 import edu.wpi.first.math.geometry.Twist2d;
-import edu.wpi.first.math.interpolation.Interpolatable;
 import edu.wpi.first.math.interpolation.TimeInterpolatableBuffer;
 import edu.wpi.first.math.kinematics.Kinematics;
 import edu.wpi.first.math.kinematics.Odometry;
 import edu.wpi.first.math.numbers.N1;
 import edu.wpi.first.math.numbers.N3;
-import java.util.Map;
-import java.util.NoSuchElementException;
-import java.util.Objects;
+import java.util.NavigableMap;
 import java.util.Optional;
+import java.util.TreeMap;
 
 /**
  * This class wraps {@link Odometry} to fuse latency-compensated vision measurements with encoder
@@ -38,14 +37,20 @@ import java.util.Optional;
  * @param <T> Wheel positions type.
  */
 public class PoseEstimator<T> {
-  private final Kinematics<?, T> m_kinematics;
   private final Odometry<T> m_odometry;
   private final Matrix<N3, N1> m_q = new Matrix<>(Nat.N3(), Nat.N1());
   private final Matrix<N3, N3> m_visionK = new Matrix<>(Nat.N3(), Nat.N3());
 
   private static final double kBufferDuration = 1.5;
-  private final TimeInterpolatableBuffer<InterpolationRecord> m_poseBuffer =
+  // Maps timestamps to odometry-only pose estimates
+  private final TimeInterpolatableBuffer<Pose2d> m_odometryPoseBuffer =
       TimeInterpolatableBuffer.createBuffer(kBufferDuration);
+  // Maps timestamps to vision updates
+  // Always contains one entry before the oldest entry in m_odometryPoseBuffer, unless there have
+  // been no vision measurements after the last reset
+  private final NavigableMap<Double, VisionUpdate> m_visionUpdates = new TreeMap<>();
+
+  private Pose2d m_poseEstimate;
 
   /**
    * Constructs a PoseEstimator.
@@ -59,14 +64,16 @@ public class PoseEstimator<T> {
    *     in meters, y position in meters, and heading in radians). Increase these numbers to trust
    *     the vision pose measurement less.
    */
+  @SuppressWarnings("PMD.UnusedFormalParameter")
   public PoseEstimator(
       Kinematics<?, T> kinematics,
       Odometry<T> odometry,
       Matrix<N3, N1> stateStdDevs,
       Matrix<N3, N1> visionMeasurementStdDevs) {
-    m_kinematics = kinematics;
     m_odometry = odometry;
 
+    m_poseEstimate = m_odometry.getPoseMeters();
+
     for (int i = 0; i < 3; ++i) {
       m_q.set(i, 0, stateStdDevs.get(i, 0) * stateStdDevs.get(i, 0));
     }
@@ -113,7 +120,9 @@ public class PoseEstimator<T> {
   public void resetPosition(Rotation2d gyroAngle, T wheelPositions, Pose2d poseMeters) {
     // Reset state estimate and error covariance
     m_odometry.resetPosition(gyroAngle, wheelPositions, poseMeters);
-    m_poseBuffer.clear();
+    m_odometryPoseBuffer.clear();
+    m_visionUpdates.clear();
+    m_poseEstimate = m_odometry.getPoseMeters();
   }
 
   /**
@@ -122,7 +131,7 @@ public class PoseEstimator<T> {
    * @return The estimated robot pose in meters.
    */
   public Pose2d getEstimatedPosition() {
-    return m_odometry.getPoseMeters();
+    return m_poseEstimate;
   }
 
   /**
@@ -132,7 +141,54 @@ public class PoseEstimator<T> {
    * @return The pose at the given timestamp (or Optional.empty() if the buffer is empty).
    */
   public Optional<Pose2d> sampleAt(double timestampSeconds) {
-    return m_poseBuffer.getSample(timestampSeconds).map(record -> record.poseMeters);
+    // Step 0: If there are no odometry updates to sample, skip.
+    if (m_odometryPoseBuffer.getInternalBuffer().isEmpty()) {
+      return Optional.empty();
+    }
+
+    // Step 1: Make sure timestamp matches the sample from the odometry pose buffer. (When sampling,
+    // the buffer will always use a timestamp between the first and last timestamps)
+    double oldestOdometryTimestamp = m_odometryPoseBuffer.getInternalBuffer().firstKey();
+    double newestOdometryTimestamp = m_odometryPoseBuffer.getInternalBuffer().lastKey();
+    timestampSeconds =
+        MathUtil.clamp(timestampSeconds, oldestOdometryTimestamp, newestOdometryTimestamp);
+
+    // Step 2: If there are no applicable vision updates, use the odometry-only information.
+    if (m_visionUpdates.isEmpty() || timestampSeconds < m_visionUpdates.firstKey()) {
+      return m_odometryPoseBuffer.getSample(timestampSeconds);
+    }
+
+    // Step 3: Get the latest vision update from before or at the timestamp to sample at.
+    double floorTimestamp = m_visionUpdates.floorKey(timestampSeconds);
+    var visionUpdate = m_visionUpdates.get(floorTimestamp);
+
+    // Step 4: Get the pose measured by odometry at the time of the sample.
+    var odometryEstimate = m_odometryPoseBuffer.getSample(timestampSeconds);
+
+    // Step 5: Apply the vision compensation to the odometry pose.
+    return odometryEstimate.map(odometryPose -> visionUpdate.compensate(odometryPose));
+  }
+
+  /** Removes stale vision updates that won't affect sampling. */
+  private void cleanUpVisionUpdates() {
+    // Step 0: If there are no odometry samples, skip.
+    if (m_odometryPoseBuffer.getInternalBuffer().isEmpty()) {
+      return;
+    }
+
+    // Step 1: Find the oldest timestamp that needs a vision update.
+    double oldestOdometryTimestamp = m_odometryPoseBuffer.getInternalBuffer().firstKey();
+
+    // Step 2: If there are no vision updates before that timestamp, skip.
+    if (m_visionUpdates.isEmpty() || oldestOdometryTimestamp < m_visionUpdates.firstKey()) {
+      return;
+    }
+
+    // Step 3: Find the newest vision update timestamp before or at the oldest timestamp.
+    double newestNeededVisionUpdateTimestamp = m_visionUpdates.floorKey(oldestOdometryTimestamp);
+
+    // Step 4: Remove all entries strictly before the newest timestamp we need.
+    m_visionUpdates.headMap(newestNeededVisionUpdateTimestamp, false).clear();
   }
 
   /**
@@ -156,50 +212,51 @@ public class PoseEstimator<T> {
    */
   public void addVisionMeasurement(Pose2d visionRobotPoseMeters, double timestampSeconds) {
     // Step 0: If this measurement is old enough to be outside the pose buffer's timespan, skip.
-    try {
-      if (m_poseBuffer.getInternalBuffer().lastKey() - kBufferDuration > timestampSeconds) {
-        return;
-      }
-    } catch (NoSuchElementException ex) {
+    if (m_odometryPoseBuffer.getInternalBuffer().isEmpty()
+        || m_odometryPoseBuffer.getInternalBuffer().lastKey() - kBufferDuration
+            > timestampSeconds) {
       return;
     }
 
-    // Step 1: Get the pose odometry measured at the moment the vision measurement was made.
-    var sample = m_poseBuffer.getSample(timestampSeconds);
+    // Step 1: Clean up any old entries
+    cleanUpVisionUpdates();
 
-    if (sample.isEmpty()) {
+    // Step 2: Get the pose measured by odometry at the moment the vision measurement was made.
+    var odometrySample = m_odometryPoseBuffer.getSample(timestampSeconds);
+
+    if (odometrySample.isEmpty()) {
       return;
     }
 
-    // Step 2: Measure the twist between the odometry pose and the vision pose.
-    var twist = sample.get().poseMeters.log(visionRobotPoseMeters);
+    // Step 3: Get the vision-compensated pose estimate at the moment the vision measurement was
+    // made.
+    var visionSample = sampleAt(timestampSeconds);
 
-    // Step 3: We should not trust the twist entirely, so instead we scale this twist by a Kalman
+    if (visionSample.isEmpty()) {
+      return;
+    }
+
+    // Step 4: Measure the twist between the old pose estimate and the vision pose.
+    var twist = visionSample.get().log(visionRobotPoseMeters);
+
+    // Step 5: We should not trust the twist entirely, so instead we scale this twist by a Kalman
     // gain matrix representing how much we trust vision measurements compared to our current pose.
     var k_times_twist = m_visionK.times(VecBuilder.fill(twist.dx, twist.dy, twist.dtheta));
 
-    // Step 4: Convert back to Twist2d.
+    // Step 6: Convert back to Twist2d.
     var scaledTwist =
         new Twist2d(k_times_twist.get(0, 0), k_times_twist.get(1, 0), k_times_twist.get(2, 0));
 
-    // Step 5: Reset Odometry to state at sample with vision adjustment.
-    m_odometry.resetPosition(
-        sample.get().gyroAngle,
-        sample.get().wheelPositions,
-        sample.get().poseMeters.exp(scaledTwist));
+    // Step 7: Calculate and record the vision update.
+    var visionUpdate = new VisionUpdate(visionSample.get().exp(scaledTwist), odometrySample.get());
+    m_visionUpdates.put(timestampSeconds, visionUpdate);
 
-    // Step 6: Record the current pose to allow multiple measurements from the same timestamp
-    m_poseBuffer.addSample(
-        timestampSeconds,
-        new InterpolationRecord(
-            getEstimatedPosition(), sample.get().gyroAngle, sample.get().wheelPositions));
+    // Step 8: Remove later vision measurements. (Matches previous behavior)
+    m_visionUpdates.tailMap(timestampSeconds, false).entrySet().clear();
 
-    // Step 7: Replay odometry inputs between sample time and latest recorded sample to update the
-    // pose buffer and correct odometry.
-    for (Map.Entry<Double, InterpolationRecord> entry :
-        m_poseBuffer.getInternalBuffer().tailMap(timestampSeconds).entrySet()) {
-      updateWithTime(entry.getKey(), entry.getValue().gyroAngle, entry.getValue().wheelPositions);
-    }
+    // Step 9: Update latest pose estimate. Since we cleared all updates after this vision update,
+    // it's guaranteed to be the latest vision update.
+    m_poseEstimate = visionUpdate.compensate(m_odometry.getPoseMeters());
   }
 
   /**
@@ -258,83 +315,52 @@ public class PoseEstimator<T> {
    * @return The estimated pose of the robot in meters.
    */
   public Pose2d updateWithTime(double currentTimeSeconds, Rotation2d gyroAngle, T wheelPositions) {
-    m_odometry.update(gyroAngle, wheelPositions);
-    m_poseBuffer.addSample(
-        currentTimeSeconds,
-        new InterpolationRecord(
-            getEstimatedPosition(), gyroAngle, m_kinematics.copy(wheelPositions)));
+    var odometryEstimate = m_odometry.update(gyroAngle, wheelPositions);
+
+    m_odometryPoseBuffer.addSample(currentTimeSeconds, odometryEstimate);
+
+    if (m_visionUpdates.isEmpty()) {
+      m_poseEstimate = odometryEstimate;
+    } else {
+      var visionUpdate = m_visionUpdates.get(m_visionUpdates.lastKey());
+      m_poseEstimate = visionUpdate.compensate(odometryEstimate);
+    }
 
     return getEstimatedPosition();
   }
 
   /**
-   * Represents an odometry record. The record contains the inputs provided as well as the pose that
-   * was observed based on these inputs, as well as the previous record and its inputs.
+   * Represents a vision update record. The record contains the vision-compensated pose estimate as
+   * well as the corresponding odometry pose estimate.
    */
-  private final class InterpolationRecord implements Interpolatable<InterpolationRecord> {
-    // The pose observed given the current sensor inputs and the previous pose.
-    private final Pose2d poseMeters;
+  private static final class VisionUpdate {
+    // The vision-compensated pose estimate.
+    private final Pose2d visionPose;
 
-    // The current gyro angle.
-    private final Rotation2d gyroAngle;
-
-    // The current encoder readings.
-    private final T wheelPositions;
+    // The pose estimated based solely on odometry.
+    private final Pose2d odometryPose;
 
     /**
-     * Constructs an Interpolation Record with the specified parameters.
+     * Constructs a vision update record with the specified parameters.
      *
-     * @param poseMeters The pose observed given the current sensor inputs and the previous pose.
-     * @param gyro The current gyro angle.
-     * @param wheelPositions The current encoder readings.
+     * @param visionPose The vision-compensated pose estimate.
+     * @param odometryPose The pose estimate based solely on odometry.
      */
-    private InterpolationRecord(Pose2d poseMeters, Rotation2d gyro, T wheelPositions) {
-      this.poseMeters = poseMeters;
-      this.gyroAngle = gyro;
-      this.wheelPositions = wheelPositions;
+    private VisionUpdate(Pose2d visionPose, Pose2d odometryPose) {
+      this.visionPose = visionPose;
+      this.odometryPose = odometryPose;
     }
 
     /**
-     * Return the interpolated record. This object is assumed to be the starting position, or lower
-     * bound.
+     * Returns the vision-compensated version of the pose. Specifically, changes the pose from being
+     * relative to this record's odometry pose to being relative to this record's vision pose.
      *
-     * @param endValue The upper bound, or end.
-     * @param t How far between the lower and upper bound we are. This should be bounded in [0, 1].
-     * @return The interpolated value.
+     * @param pose The pose to compensate.
+     * @return The compensated pose.
      */
-    @Override
-    public InterpolationRecord interpolate(InterpolationRecord endValue, double t) {
-      if (t < 0) {
-        return this;
-      } else if (t >= 1) {
-        return endValue;
-      } else {
-        // Find the new wheel distances.
-        var wheelLerp = m_kinematics.interpolate(wheelPositions, endValue.wheelPositions, t);
-
-        // Find the new gyro angle.
-        var gyroLerp = gyroAngle.interpolate(endValue.gyroAngle, t);
-
-        // Create a twist to represent the change based on the interpolated sensor inputs.
-        Twist2d twist = m_kinematics.toTwist2d(wheelPositions, wheelLerp);
-        twist.dtheta = gyroLerp.minus(gyroAngle).getRadians();
-
-        return new InterpolationRecord(poseMeters.exp(twist), gyroLerp, wheelLerp);
-      }
-    }
-
-    @Override
-    public boolean equals(Object obj) {
-      return this == obj
-          || obj instanceof PoseEstimator<?>.InterpolationRecord record
-              && Objects.equals(gyroAngle, record.gyroAngle)
-              && Objects.equals(wheelPositions, record.wheelPositions)
-              && Objects.equals(poseMeters, record.poseMeters);
-    }
-
-    @Override
-    public int hashCode() {
-      return Objects.hash(gyroAngle, wheelPositions, poseMeters);
+    public Pose2d compensate(Pose2d pose) {
+      var delta = pose.minus(this.odometryPose);
+      return this.visionPose.plus(delta);
     }
   }
 }

wpimath/src/main/native/include/frc/estimator/PoseEstimator.h

@@ -4,7 +4,10 @@
 
 #pragma once
 
+#include <map>
 #include <optional>
+#include <utility>
+#include <vector>
 
 #include <Eigen/Core>
 #include <wpi/SymbolExports.h>
@@ -185,57 +188,46 @@ class WPILIB_DLLEXPORT PoseEstimator {
                         const WheelPositions& wheelPositions);
 
  private:
-  struct InterpolationRecord {
-    // The pose observed given the current sensor inputs and the previous pose.
-    Pose2d pose;
+  /**
+   * Removes stale vision updates that won't affect sampling.
+   */
+  void CleanUpVisionUpdates();
 
-    // The current gyroscope angle.
-    Rotation2d gyroAngle;
+  struct VisionUpdate {
+    // The vision-compensated pose estimate
+    Pose2d visionPose;
 
-    // The distances traveled by the wheels.
-    WheelPositions wheelPositions;
+    // The pose estimated based solely on odometry
+    Pose2d odometryPose;
 
     /**
-     * Checks equality between this InterpolationRecord and another object.
+     * Returns the vision-compensated version of the pose. Specifically, changes
+     * the pose from being relative to this record's odometry pose to being
+     * relative to this record's vision pose.
      *
-     * @param other The other object.
-     * @return Whether the two objects are equal.
+     * @param pose The pose to compensate.
+     * @return The compensated pose.
      */
-    bool operator==(const InterpolationRecord& other) const = default;
-
-    /**
-     * Checks inequality between this InterpolationRecord and another object.
-     *
-     * @param other The other object.
-     * @return Whether the two objects are not equal.
-     */
-    bool operator!=(const InterpolationRecord& other) const = default;
-
-    /**
-     * Interpolates between two InterpolationRecords.
-     *
-     * @param endValue The end value for the interpolation.
-     * @param i The interpolant (fraction).
-     *
-     * @return The interpolated state.
-     */
-    InterpolationRecord Interpolate(
-        Kinematics<WheelSpeeds, WheelPositions>& kinematics,
-        InterpolationRecord endValue, double i) const;
+    Pose2d Compensate(const Pose2d& pose) const {
+      auto delta = pose - odometryPose;
+      return visionPose + delta;
+    }
   };
 
   static constexpr units::second_t kBufferDuration = 1.5_s;
 
-  Kinematics<WheelSpeeds, WheelPositions>& m_kinematics;
   Odometry<WheelSpeeds, WheelPositions>& m_odometry;
   wpi::array<double, 3> m_q{wpi::empty_array};
   Eigen::Matrix3d m_visionK = Eigen::Matrix3d::Zero();
 
-  TimeInterpolatableBuffer<InterpolationRecord> m_poseBuffer{
-      kBufferDuration, [this](const InterpolationRecord& start,
-                              const InterpolationRecord& end, double t) {
-        return start.Interpolate(this->m_kinematics, end, t);
-      }};
+  // Maps timestamps to odometry-only pose estimates
+  TimeInterpolatableBuffer<Pose2d> m_odometryPoseBuffer{kBufferDuration};
+  // Maps timestamps to vision updates
+  // Always contains one entry before the oldest entry in m_odometryPoseBuffer,
+  // unless there have been no vision measurements after the last reset
+  std::map<units::second_t, VisionUpdate> m_visionUpdates;
+
+  Pose2d m_poseEstimate;
 };
 }  // namespace frc
 

wpimath/src/main/native/include/frc/estimator/PoseEstimator.inc

@@ -14,7 +14,7 @@ PoseEstimator<WheelSpeeds, WheelPositions>::PoseEstimator(
     Odometry<WheelSpeeds, WheelPositions>& odometry,
     const wpi::array<double, 3>& stateStdDevs,
     const wpi::array<double, 3>& visionMeasurementStdDevs)
-    : m_kinematics(kinematics), m_odometry(odometry) {
+    : m_odometry(odometry), m_poseEstimate(m_odometry.GetPose()) {
   for (size_t i = 0; i < 3; ++i) {
     m_q[i] = stateStdDevs[i] * stateStdDevs[i];
   }
@@ -48,26 +48,93 @@ void PoseEstimator<WheelSpeeds, WheelPositions>::ResetPosition(
     const Pose2d& pose) {
   // Reset state estimate and error covariance
   m_odometry.ResetPosition(gyroAngle, wheelPositions, pose);
-  m_poseBuffer.Clear();
+  m_odometryPoseBuffer.Clear();
+  m_visionUpdates.clear();
+  m_poseEstimate = m_odometry.GetPose();
 }
 
 template <typename WheelSpeeds, typename WheelPositions>
 Pose2d PoseEstimator<WheelSpeeds, WheelPositions>::GetEstimatedPosition()
     const {
-  return m_odometry.GetPose();
+  return m_poseEstimate;
+  if (m_visionUpdates.empty()) {
+    return m_odometry.GetPose();
+  }
+  auto visionUpdate = m_visionUpdates.rbegin()->second;
+  return visionUpdate.Compensate(m_odometry.GetPose());
 }
 
 template <typename WheelSpeeds, typename WheelPositions>
 std::optional<Pose2d> PoseEstimator<WheelSpeeds, WheelPositions>::SampleAt(
     units::second_t timestamp) const {
-  // TODO Replace with std::optional::transform() in C++23
-  std::optional<PoseEstimator<WheelSpeeds, WheelPositions>::InterpolationRecord>
-      sample = m_poseBuffer.Sample(timestamp);
-  if (sample) {
-    return sample->pose;
-  } else {
+  // Step 0: If there are no odometry updates to sample, skip.
+  if (m_odometryPoseBuffer.GetInternalBuffer().empty()) {
     return std::nullopt;
   }
+
+  // Step 1: Make sure timestamp matches the sample from the odometry pose
+  // buffer. (When sampling, the buffer will always use a timestamp
+  // between the first and last timestamps)
+  units::second_t oldestOdometryTimestamp =
+      m_odometryPoseBuffer.GetInternalBuffer().front().first;
+  units::second_t newestOdometryTimestamp =
+      m_odometryPoseBuffer.GetInternalBuffer().back().first;
+  timestamp =
+      std::clamp(timestamp, oldestOdometryTimestamp, newestOdometryTimestamp);
+
+  // Step 2: If there are no applicable vision updates, use the odometry-only
+  // information.
+  if (m_visionUpdates.empty() || timestamp < m_visionUpdates.begin()->first) {
+    return m_odometryPoseBuffer.Sample(timestamp);
+  }
+
+  // Step 3: Get the latest vision update from before or at the timestamp to
+  // sample at.
+  // First, find the iterator past the sample timestamp, then go back one. Note
+  // that upper_bound() won't return begin() because we check begin() earlier.
+  auto floorIter = m_visionUpdates.upper_bound(timestamp);
+  --floorIter;
+  auto visionUpdate = floorIter->second;
+
+  // Step 4: Get the pose measured by odometry at the time of the sample.
+  auto odometryEstimate = m_odometryPoseBuffer.Sample(timestamp);
+
+  // Step 5: Apply the vision compensation to the odometry pose.
+  // TODO Replace with std::optional::transform() in C++23
+  if (odometryEstimate) {
+    return visionUpdate.Compensate(*odometryEstimate);
+  }
+  return std::nullopt;
+}
+
+template <typename WheelSpeeds, typename WheelPositions>
+void PoseEstimator<WheelSpeeds, WheelPositions>::CleanUpVisionUpdates() {
+  // Step 0: If there are no odometry samples, skip.
+  if (m_odometryPoseBuffer.GetInternalBuffer().empty()) {
+    return;
+  }
+
+  // Step 1: Find the oldest timestamp that needs a vision update.
+  units::second_t oldestOdometryTimestamp =
+      m_odometryPoseBuffer.GetInternalBuffer().front().first;
+
+  // Step 2: If there are no vision updates before that timestamp, skip.
+  if (m_visionUpdates.empty() ||
+      oldestOdometryTimestamp < m_visionUpdates.begin()->first) {
+    return;
+  }
+
+  // Step 3: Find the newest vision update timestamp before or at the oldest
+  // timestamp.
+  // First, find the iterator past the oldest odometry timestamp, then go
+  // back one. Note that upper_bound() won't return begin() because we check
+  // begin() earlier.
+  auto newestNeededVisionUpdate =
+      m_visionUpdates.upper_bound(oldestOdometryTimestamp);
+  --newestNeededVisionUpdate;
+
+  // Step 4: Remove all entries strictly before the newest timestamp we need.
+  m_visionUpdates.erase(m_visionUpdates.begin(), newestNeededVisionUpdate);
 }
 
 template <typename WheelSpeeds, typename WheelPositions>
@@ -75,57 +142,58 @@ void PoseEstimator<WheelSpeeds, WheelPositions>::AddVisionMeasurement(
     const Pose2d& visionRobotPose, units::second_t timestamp) {
   // Step 0: If this measurement is old enough to be outside the pose buffer's
   // timespan, skip.
-  if (!m_poseBuffer.GetInternalBuffer().empty() &&
-      m_poseBuffer.GetInternalBuffer().front().first - kBufferDuration >
+  if (m_odometryPoseBuffer.GetInternalBuffer().empty() ||
+      m_odometryPoseBuffer.GetInternalBuffer().front().first - kBufferDuration >
           timestamp) {
     return;
   }
 
-  // Step 1: Get the estimated pose from when the vision measurement was made.
-  auto sample = m_poseBuffer.Sample(timestamp);
+  // Step 1: Clean up any old entries
+  CleanUpVisionUpdates();
 
-  if (!sample.has_value()) {
+  // Step 2: Get the pose measured by odometry at the moment the vision
+  // measurement was made.
+  auto odometrySample = m_odometryPoseBuffer.Sample(timestamp);
+
+  if (!odometrySample) {
     return;
   }
 
-  // Step 2: Measure the twist between the odometry pose and the vision pose
-  auto twist = sample.value().pose.Log(visionRobotPose);
+  // Step 3: Get the vision-compensated pose estimate at the moment the vision
+  // measurement was made.
+  auto visionSample = SampleAt(timestamp);
 
-  // Step 3: We should not trust the twist entirely, so instead we scale this
+  if (!visionSample) {
+    return;
+  }
+
+  // Step 4: Measure the twist between the old pose estimate and the vision
+  // pose.
+  auto twist = visionSample.value().Log(visionRobotPose);
+
+  // Step 5: We should not trust the twist entirely, so instead we scale this
   // twist by a Kalman gain matrix representing how much we trust vision
   // measurements compared to our current pose.
   Eigen::Vector3d k_times_twist =
       m_visionK *
       Eigen::Vector3d{twist.dx.value(), twist.dy.value(), twist.dtheta.value()};
 
-  // Step 4: Convert back to Twist2d
+  // Step 6: Convert back to Twist2d.
   Twist2d scaledTwist{units::meter_t{k_times_twist(0)},
                       units::meter_t{k_times_twist(1)},
                       units::radian_t{k_times_twist(2)}};
 
-  // Step 5: Reset Odometry to state at sample with vision adjustment.
-  m_odometry.ResetPosition(sample.value().gyroAngle,
-                           sample.value().wheelPositions,
-                           sample.value().pose.Exp(scaledTwist));
+  // Step 7: Calculate and record the vision update.
+  VisionUpdate visionUpdate{visionSample->Exp(scaledTwist), *odometrySample};
+  m_visionUpdates[timestamp] = visionUpdate;
 
-  // Step 6: Record the current pose to allow multiple measurements from the
-  // same timestamp
-  m_poseBuffer.AddSample(timestamp,
-                         {GetEstimatedPosition(), sample.value().gyroAngle,
-                          sample.value().wheelPositions});
+  // Step 8: Remove later vision measurements. (Matches previous behavior)
+  auto firstAfter = m_visionUpdates.upper_bound(timestamp);
+  m_visionUpdates.erase(firstAfter, m_visionUpdates.end());
 
-  // Step 7: Replay odometry inputs between sample time and latest recorded
-  // sample to update the pose buffer and correct odometry.
-  auto internal_buf = m_poseBuffer.GetInternalBuffer();
-
-  auto upper_bound =
-      std::lower_bound(internal_buf.begin(), internal_buf.end(), timestamp,
-                       [](const auto& pair, auto t) { return t > pair.first; });
-
-  for (auto entry = upper_bound; entry != internal_buf.end(); entry++) {
-    UpdateWithTime(entry->first, entry->second.gyroAngle,
-                   entry->second.wheelPositions);
-  }
+  // Step 9: Update latest pose estimate. Since we cleared all updates after
+  // this vision update, it's guaranteed to be the latest vision update.
+  m_poseEstimate = visionUpdate.Compensate(m_odometry.GetPose());
 }
 
 template <typename WheelSpeeds, typename WheelPositions>
@@ -139,40 +207,18 @@ template <typename WheelSpeeds, typename WheelPositions>
 Pose2d PoseEstimator<WheelSpeeds, WheelPositions>::UpdateWithTime(
     units::second_t currentTime, const Rotation2d& gyroAngle,
     const WheelPositions& wheelPositions) {
-  m_odometry.Update(gyroAngle, wheelPositions);
+  auto odometryEstimate = m_odometry.Update(gyroAngle, wheelPositions);
 
-  WheelPositions internalWheelPositions = wheelPositions;
+  m_odometryPoseBuffer.AddSample(currentTime, odometryEstimate);
 
-  m_poseBuffer.AddSample(
-      currentTime, {GetEstimatedPosition(), gyroAngle, internalWheelPositions});
+  if (m_visionUpdates.empty()) {
+    m_poseEstimate = odometryEstimate;
+  } else {
+    auto visionUpdate = m_visionUpdates.rbegin()->second;
+    m_poseEstimate = visionUpdate.Compensate(odometryEstimate);
+  }
 
   return GetEstimatedPosition();
 }
 
-template <typename WheelSpeeds, typename WheelPositions>
-typename PoseEstimator<WheelSpeeds, WheelPositions>::InterpolationRecord
-PoseEstimator<WheelSpeeds, WheelPositions>::InterpolationRecord::Interpolate(
-    Kinematics<WheelSpeeds, WheelPositions>& kinematics,
-    InterpolationRecord endValue, double i) const {
-  if (i < 0) {
-    return *this;
-  } else if (i > 1) {
-    return endValue;
-  } else {
-    // Find the new wheel distance measurements.
-    WheelPositions wheels_lerp = kinematics.Interpolate(
-        this->wheelPositions, endValue.wheelPositions, i);
-
-    // Find the new gyro angle.
-    auto gyro = wpi::Lerp(this->gyroAngle, endValue.gyroAngle, i);
-
-    // Create a twist to represent the change based on the interpolated
-    // sensor inputs.
-    auto twist = kinematics.ToTwist2d(this->wheelPositions, wheels_lerp);
-    twist.dtheta = (gyro - gyroAngle).Radians();
-
-    return {pose.Exp(twist), gyro, wheels_lerp};
-  }
-}
-
 }  // namespace frc

wpimath/src/main/native/include/frc/interpolation/TimeInterpolatableBuffer.h

@@ -144,6 +144,13 @@ class TimeInterpolatableBuffer {
     return m_pastSnapshots;
   }
 
+  /**
+   * Grant access to the internal sample buffer.
+   */
+  const std::vector<std::pair<units::second_t, T>>& GetInternalBuffer() const {
+    return m_pastSnapshots;
+  }
+
  private:
   units::second_t m_historySize;
   std::vector<std::pair<units::second_t, T>> m_pastSnapshots;