[wpimath] Fix ResetTranslation and ResetRotation in PoseEstimator and PoseEstimator3d causing the robot to teleport

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

@@ -48,7 +48,9 @@ public class PoseEstimator<T> {
       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
+  // been no vision measurements after the last reset. May contain one entry while
+  // m_odometryPoseBuffer is empty to correct for translation/rotation after a call to
+  // ResetRotation/ResetTranslation.
   private final NavigableMap<Double, VisionUpdate> m_visionUpdates = new TreeMap<>();
 
   private Pose2d m_poseEstimate;
@@ -145,9 +147,22 @@ public void resetPose(Pose2d pose) {
    */
   public void resetTranslation(Translation2d translation) {
     m_odometry.resetTranslation(translation);
+
+    final var latestVisionUpdate = m_visionUpdates.lastEntry();
     m_odometryPoseBuffer.clear();
     m_visionUpdates.clear();
-    m_poseEstimate = m_odometry.getPoseMeters();
+
+    if (latestVisionUpdate != null) {
+      // apply vision compensation to the pose rotation
+      final var visionUpdate =
+          new VisionUpdate(
+              new Pose2d(translation, latestVisionUpdate.getValue().visionPose.getRotation()),
+              new Pose2d(translation, latestVisionUpdate.getValue().odometryPose.getRotation()));
+      m_visionUpdates.put(latestVisionUpdate.getKey(), visionUpdate);
+      m_poseEstimate = visionUpdate.compensate(m_odometry.getPoseMeters());
+    } else {
+      m_poseEstimate = m_odometry.getPoseMeters();
+    }
   }
 
   /**
@@ -157,9 +172,22 @@ public void resetTranslation(Translation2d translation) {
    */
   public void resetRotation(Rotation2d rotation) {
     m_odometry.resetRotation(rotation);
+
+    final var latestVisionUpdate = m_visionUpdates.lastEntry();
     m_odometryPoseBuffer.clear();
     m_visionUpdates.clear();
-    m_poseEstimate = m_odometry.getPoseMeters();
+
+    if (latestVisionUpdate != null) {
+      // apply vision compensation to the pose translation
+      final var visionUpdate =
+          new VisionUpdate(
+              new Pose2d(latestVisionUpdate.getValue().visionPose.getTranslation(), rotation),
+              new Pose2d(latestVisionUpdate.getValue().odometryPose.getTranslation(), rotation));
+      m_visionUpdates.put(latestVisionUpdate.getKey(), visionUpdate);
+      m_poseEstimate = visionUpdate.compensate(m_odometry.getPoseMeters());
+    } else {
+      m_poseEstimate = m_odometry.getPoseMeters();
+    }
   }
 
   /**

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

@@ -56,7 +56,9 @@ public class PoseEstimator3d<T> {
       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
+  // been no vision measurements after the last reset. May contain one entry while
+  // m_odometryPoseBuffer is empty to correct for translation/rotation after a call to
+  // ResetRotation/ResetTranslation.
   private final NavigableMap<Double, VisionUpdate> m_visionUpdates = new TreeMap<>();
 
   private Pose3d m_poseEstimate;
@@ -157,9 +159,22 @@ public void resetPose(Pose3d pose) {
    */
   public void resetTranslation(Translation3d translation) {
     m_odometry.resetTranslation(translation);
+
+    final var latestVisionUpdate = m_visionUpdates.lastEntry();
     m_odometryPoseBuffer.clear();
     m_visionUpdates.clear();
-    m_poseEstimate = m_odometry.getPoseMeters();
+
+    if (latestVisionUpdate != null) {
+      // apply vision compensation to the pose rotation
+      final var visionUpdate =
+          new VisionUpdate(
+              new Pose3d(translation, latestVisionUpdate.getValue().visionPose.getRotation()),
+              new Pose3d(translation, latestVisionUpdate.getValue().odometryPose.getRotation()));
+      m_visionUpdates.put(latestVisionUpdate.getKey(), visionUpdate);
+      m_poseEstimate = visionUpdate.compensate(m_odometry.getPoseMeters());
+    } else {
+      m_poseEstimate = m_odometry.getPoseMeters();
+    }
   }
 
   /**
@@ -169,9 +184,22 @@ public void resetTranslation(Translation3d translation) {
    */
   public void resetRotation(Rotation3d rotation) {
     m_odometry.resetRotation(rotation);
+
+    final var latestVisionUpdate = m_visionUpdates.lastEntry();
     m_odometryPoseBuffer.clear();
     m_visionUpdates.clear();
-    m_poseEstimate = m_odometry.getPoseMeters();
+
+    if (latestVisionUpdate != null) {
+      // apply vision compensation to the pose translation
+      final var visionUpdate =
+          new VisionUpdate(
+              new Pose3d(latestVisionUpdate.getValue().visionPose.getTranslation(), rotation),
+              new Pose3d(latestVisionUpdate.getValue().odometryPose.getTranslation(), rotation));
+      m_visionUpdates.put(latestVisionUpdate.getKey(), visionUpdate);
+      m_poseEstimate = visionUpdate.compensate(m_odometry.getPoseMeters());
+    } else {
+      m_poseEstimate = m_odometry.getPoseMeters();
+    }
   }
 
   /**

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

@@ -136,24 +136,70 @@ class WPILIB_DLLEXPORT PoseEstimator {
    *
    * @param translation The pose to translation to.
    */
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
+#endif  // defined(__GNUC__) && !defined(__clang__)
   void ResetTranslation(const Translation2d& translation) {
     m_odometry.ResetTranslation(translation);
+
+    const std::optional<std::pair<units::second_t, VisionUpdate>>
+        latestVisionUpdate =
+            m_visionUpdates.empty() ? std::nullopt
+                                    : std::optional{*m_visionUpdates.crbegin()};
     m_odometryPoseBuffer.Clear();
     m_visionUpdates.clear();
-    m_poseEstimate = m_odometry.GetPose();
+
+    if (latestVisionUpdate) {
+      // apply vision compensation to the pose rotation
+      const VisionUpdate visionUpdate{
+          Pose2d{translation, latestVisionUpdate->second.visionPose.Rotation()},
+          Pose2d{translation,
+                 latestVisionUpdate->second.odometryPose.Rotation()}};
+      m_visionUpdates[latestVisionUpdate->first] = visionUpdate;
+      m_poseEstimate = visionUpdate.Compensate(m_odometry.GetPose());
+    } else {
+      m_poseEstimate = m_odometry.GetPose();
+    }
   }
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic pop
+#endif  // defined(__GNUC__) && !defined(__clang__)
 
   /**
    * Resets the robot's rotation.
    *
    * @param rotation The rotation to reset to.
    */
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
+#endif  // defined(__GNUC__) && !defined(__clang__)
   void ResetRotation(const Rotation2d& rotation) {
     m_odometry.ResetRotation(rotation);
+
+    const std::optional<std::pair<units::second_t, VisionUpdate>>
+        latestVisionUpdate =
+            m_visionUpdates.empty() ? std::nullopt
+                                    : std::optional{*m_visionUpdates.crbegin()};
     m_odometryPoseBuffer.Clear();
     m_visionUpdates.clear();
-    m_poseEstimate = m_odometry.GetPose();
+
+    if (latestVisionUpdate) {
+      // apply vision compensation to the pose translation
+      const VisionUpdate visionUpdate{
+          Pose2d{latestVisionUpdate->second.visionPose.Translation(), rotation},
+          Pose2d{latestVisionUpdate->second.odometryPose.Translation(),
+                 rotation}};
+      m_visionUpdates[latestVisionUpdate->first] = visionUpdate;
+      m_poseEstimate = visionUpdate.Compensate(m_odometry.GetPose());
+    } else {
+      m_poseEstimate = m_odometry.GetPose();
+    }
   }
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic pop
+#endif  // defined(__GNUC__) && !defined(__clang__)
 
   /**
    * Gets the estimated robot pose.
@@ -431,7 +477,9 @@ class WPILIB_DLLEXPORT PoseEstimator {
   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
+  // unless there have been no vision measurements after the last reset. May
+  // contain one entry while m_odometryPoseBuffer is empty to correct for
+  // translation/rotation after a call to ResetRotation/ResetTranslation.
   std::map<units::second_t, VisionUpdate> m_visionUpdates;
 
   Pose2d m_poseEstimate;

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

@@ -145,24 +145,70 @@ class WPILIB_DLLEXPORT PoseEstimator3d {
    *
    * @param translation The pose to translation to.
    */
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
+#endif  // defined(__GNUC__) && !defined(__clang__)
   void ResetTranslation(const Translation3d& translation) {
     m_odometry.ResetTranslation(translation);
+
+    const std::optional<std::pair<units::second_t, VisionUpdate>>
+        latestVisionUpdate =
+            m_visionUpdates.empty() ? std::nullopt
+                                    : std::optional{*m_visionUpdates.crbegin()};
     m_odometryPoseBuffer.Clear();
     m_visionUpdates.clear();
-    m_poseEstimate = m_odometry.GetPose();
+
+    if (latestVisionUpdate) {
+      // apply vision compensation to the pose rotation
+      const VisionUpdate visionUpdate{
+          Pose3d{translation, latestVisionUpdate->second.visionPose.Rotation()},
+          Pose3d{translation,
+                 latestVisionUpdate->second.odometryPose.Rotation()}};
+      m_visionUpdates[latestVisionUpdate->first] = visionUpdate;
+      m_poseEstimate = visionUpdate.Compensate(m_odometry.GetPose());
+    } else {
+      m_poseEstimate = m_odometry.GetPose();
+    }
   }
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic pop
+#endif  // defined(__GNUC__) && !defined(__clang__)
 
   /**
    * Resets the robot's rotation.
    *
    * @param rotation The rotation to reset to.
    */
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
+#endif  // defined(__GNUC__) && !defined(__clang__)
   void ResetRotation(const Rotation3d& rotation) {
     m_odometry.ResetRotation(rotation);
+
+    const std::optional<std::pair<units::second_t, VisionUpdate>>
+        latestVisionUpdate =
+            m_visionUpdates.empty() ? std::nullopt
+                                    : std::optional{*m_visionUpdates.crbegin()};
     m_odometryPoseBuffer.Clear();
     m_visionUpdates.clear();
-    m_poseEstimate = m_odometry.GetPose();
+
+    if (latestVisionUpdate) {
+      // apply vision compensation to the pose translation
+      const VisionUpdate visionUpdate{
+          Pose3d{latestVisionUpdate->second.visionPose.Translation(), rotation},
+          Pose3d{latestVisionUpdate->second.odometryPose.Translation(),
+                 rotation}};
+      m_visionUpdates[latestVisionUpdate->first] = visionUpdate;
+      m_poseEstimate = visionUpdate.Compensate(m_odometry.GetPose());
+    } else {
+      m_poseEstimate = m_odometry.GetPose();
+    }
   }
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic pop
+#endif  // defined(__GNUC__) && !defined(__clang__)
 
   /**
    * Gets the estimated robot pose.
@@ -442,7 +488,9 @@ class WPILIB_DLLEXPORT PoseEstimator3d {
   TimeInterpolatableBuffer<Pose3d> 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
+  // unless there have been no vision measurements after the last reset. May
+  // contain one entry while m_odometryPoseBuffer is empty to correct for
+  // translation/rotation after a call to ResetRotation/ResetTranslation.
   std::map<units::second_t, VisionUpdate> m_visionUpdates;
 
   Pose3d m_poseEstimate;

wpimath/src/test/java/edu/wpi/first/math/estimator/SwerveDrivePoseEstimator3dTest.java

@@ -8,6 +8,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import edu.wpi.first.math.MathSharedStore;
 import edu.wpi.first.math.VecBuilder;
 import edu.wpi.first.math.geometry.Pose2d;
 import edu.wpi.first.math.geometry.Pose3d;
@@ -509,11 +510,23 @@ void testReset() {
         () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getY(), kEpsilon),
         () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getZ(), kEpsilon));
 
+    // Add a vision measurement with a different translation
+    estimator.addVisionMeasurement(
+        new Pose3d(3, 0, 0, Rotation3d.kZero), MathSharedStore.getTimestamp());
+
+    assertAll(
+        () -> assertEquals(2.5, estimator.getEstimatedPosition().getX(), kEpsilon),
+        () -> assertEquals(0, estimator.getEstimatedPosition().getY(), kEpsilon),
+        () -> assertEquals(0, estimator.getEstimatedPosition().getZ(), kEpsilon),
+        () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getX(), kEpsilon),
+        () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getY(), kEpsilon),
+        () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getZ(), kEpsilon));
+
     // Test reset rotation
     estimator.resetRotation(new Rotation3d(Rotation2d.kCCW_Pi_2));
 
     assertAll(
-        () -> assertEquals(2, estimator.getEstimatedPosition().getX(), kEpsilon),
+        () -> assertEquals(2.5, estimator.getEstimatedPosition().getX(), kEpsilon),
         () -> assertEquals(0, estimator.getEstimatedPosition().getY(), kEpsilon),
         () -> assertEquals(0, estimator.getEstimatedPosition().getZ(), kEpsilon),
         () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getX(), kEpsilon),
@@ -533,7 +546,7 @@ void testReset() {
     }
 
     assertAll(
-        () -> assertEquals(2, estimator.getEstimatedPosition().getX(), kEpsilon),
+        () -> assertEquals(2.5, estimator.getEstimatedPosition().getX(), kEpsilon),
         () -> assertEquals(1, estimator.getEstimatedPosition().getY(), kEpsilon),
         () -> assertEquals(0, estimator.getEstimatedPosition().getZ(), kEpsilon),
         () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getX(), kEpsilon),
@@ -542,6 +555,22 @@ void testReset() {
             assertEquals(
                 Math.PI / 2, estimator.getEstimatedPosition().getRotation().getZ(), kEpsilon));
 
+    // Add a vision measurement with a different rotation
+    estimator.addVisionMeasurement(
+        new Pose3d(2.5, 1, 0, new Rotation3d(Rotation2d.kPi)), MathSharedStore.getTimestamp());
+
+    assertAll(
+        () -> assertEquals(2.5, estimator.getEstimatedPosition().getX(), kEpsilon),
+        () -> assertEquals(1, estimator.getEstimatedPosition().getY(), kEpsilon),
+        () -> assertEquals(0, estimator.getEstimatedPosition().getZ(), kEpsilon),
+        () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getX(), kEpsilon),
+        () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getY(), kEpsilon),
+        () ->
+            assertEquals(
+                Math.PI * 3.0 / 4,
+                estimator.getEstimatedPosition().getRotation().getZ(),
+                kEpsilon));
+
     // Test reset translation
     estimator.resetTranslation(new Translation3d(-1, -1, -1));
 
@@ -553,7 +582,9 @@ void testReset() {
         () -> assertEquals(0, estimator.getEstimatedPosition().getRotation().getY(), kEpsilon),
         () ->
             assertEquals(
-                Math.PI / 2, estimator.getEstimatedPosition().getRotation().getZ(), kEpsilon));
+                Math.PI * 3.0 / 4,
+                estimator.getEstimatedPosition().getRotation().getZ(),
+                kEpsilon));
 
     // Test reset pose
     estimator.resetPose(Pose3d.kZero);