[wpimath] Use Odometry for internal state in Pose Estimation (#4668)

This effectively replaces the Unscented Kalman Filter used for Pose Estimation with the Odometry model, and uses a recalculable Kalman gain matrix to update pose estimations whenever a vision measurement is added.

Notably, this change removes the need for the confusing generics used in Java, and the C++ implementation got quite a bit less complex as well.

Co-authored-by: Tyler Veness <calcmogul@gmail.com>

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

@@ -6,35 +6,37 @@ package edu.wpi.first.math.estimator;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
-import edu.wpi.first.math.MatBuilder;
-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.Transform2d;
 import edu.wpi.first.math.geometry.Translation2d;
 import edu.wpi.first.math.kinematics.ChassisSpeeds;
 import edu.wpi.first.math.kinematics.DifferentialDriveKinematics;
-import edu.wpi.first.math.kinematics.DifferentialDriveWheelSpeeds;
 import edu.wpi.first.math.trajectory.Trajectory;
 import edu.wpi.first.math.trajectory.TrajectoryConfig;
 import edu.wpi.first.math.trajectory.TrajectoryGenerator;
 import java.util.List;
 import java.util.Random;
+import java.util.TreeMap;
+import java.util.function.Function;
 import org.junit.jupiter.api.Test;
 
 class DifferentialDrivePoseEstimatorTest {
   @Test
   void testAccuracy() {
+    var kinematics = new DifferentialDriveKinematics(1);
+
     var estimator =
         new DifferentialDrivePoseEstimator(
+            kinematics,
             new Rotation2d(),
             0,
             0,
             new Pose2d(),
-            new MatBuilder<>(Nat.N5(), Nat.N1()).fill(0.02, 0.02, 0.01, 0.02, 0.02),
-            new MatBuilder<>(Nat.N3(), Nat.N1()).fill(0.01, 0.01, 0.001),
-            new MatBuilder<>(Nat.N3(), Nat.N1()).fill(0.1, 0.1, 0.01));
-
-    var traj =
+            VecBuilder.fill(0.02, 0.02, 0.01),
+            VecBuilder.fill(0.1, 0.1, 0.1));
+    var trajectory =
         TrajectoryGenerator.generateTrajectory(
             List.of(
                 new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
@@ -42,67 +44,165 @@ class DifferentialDrivePoseEstimatorTest {
                 new Pose2d(0, 0, Rotation2d.fromDegrees(135)),
                 new Pose2d(-3, 0, Rotation2d.fromDegrees(-90)),
                 new Pose2d(0, 0, Rotation2d.fromDegrees(45))),
-            new TrajectoryConfig(10, 5));
+            new TrajectoryConfig(2, 2));
 
+    testFollowTrajectory(
+        kinematics,
+        estimator,
+        trajectory,
+        state ->
+            new ChassisSpeeds(
+                state.velocityMetersPerSecond,
+                0,
+                state.velocityMetersPerSecond * state.curvatureRadPerMeter),
+        state -> state.poseMeters,
+        trajectory.getInitialPose(),
+        new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
+        0.02,
+        0.1,
+        0.25,
+        true);
+  }
+
+  @Test
+  void testBadInitialPose() {
     var kinematics = new DifferentialDriveKinematics(1);
-    var rand = new Random(4915);
 
-    final double dt = 0.02;
+    var estimator =
+        new DifferentialDrivePoseEstimator(
+            kinematics,
+            new Rotation2d(),
+            0,
+            0,
+            new Pose2d(),
+            VecBuilder.fill(0.02, 0.02, 0.01),
+            VecBuilder.fill(0.1, 0.1, 0.1));
+
+    var trajectory =
+        TrajectoryGenerator.generateTrajectory(
+            List.of(
+                new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
+                new Pose2d(3, 0, Rotation2d.fromDegrees(-90)),
+                new Pose2d(0, 0, Rotation2d.fromDegrees(135)),
+                new Pose2d(-3, 0, Rotation2d.fromDegrees(-90)),
+                new Pose2d(0, 0, Rotation2d.fromDegrees(45))),
+            new TrajectoryConfig(2, 2));
+
+    for (int offset_direction_degs = 0; offset_direction_degs < 360; offset_direction_degs += 45) {
+      for (int offset_heading_degs = 0; offset_heading_degs < 360; offset_heading_degs += 45) {
+        var pose_offset = Rotation2d.fromDegrees(offset_direction_degs);
+        var heading_offset = Rotation2d.fromDegrees(offset_heading_degs);
+
+        var initial_pose =
+            trajectory
+                .getInitialPose()
+                .plus(
+                    new Transform2d(
+                        new Translation2d(pose_offset.getCos(), pose_offset.getSin()),
+                        heading_offset));
+
+        testFollowTrajectory(
+            kinematics,
+            estimator,
+            trajectory,
+            state ->
+                new ChassisSpeeds(
+                    state.velocityMetersPerSecond,
+                    0,
+                    state.velocityMetersPerSecond * state.curvatureRadPerMeter),
+            state -> state.poseMeters,
+            initial_pose,
+            new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
+            0.02,
+            0.1,
+            0.25,
+            false);
+      }
+    }
+  }
+
+  void testFollowTrajectory(
+      final DifferentialDriveKinematics kinematics,
+      final DifferentialDrivePoseEstimator estimator,
+      final Trajectory trajectory,
+      final Function<Trajectory.State, ChassisSpeeds> chassisSpeedsGenerator,
+      final Function<Trajectory.State, Pose2d> visionMeasurementGenerator,
+      final Pose2d startingPose,
+      final Pose2d endingPose,
+      final double dt,
+      final double visionUpdateRate,
+      final double visionUpdateDelay,
+      final boolean checkError) {
+    double leftDistanceMeters = 0;
+    double rightDistanceMeters = 0;
+
+    estimator.resetPosition(
+        new Rotation2d(), leftDistanceMeters, rightDistanceMeters, startingPose);
+
+    var rand = new Random(3538);
+
     double t = 0.0;
 
-    final double visionUpdateRate = 0.1;
-    Pose2d lastVisionPose = null;
-    double lastVisionUpdateTime = Double.NEGATIVE_INFINITY;
+    System.out.print("time, est_x, est_y, est_theta, true_x, true_y, true_theta\n");
 
-    double distanceLeft = 0.0;
-    double distanceRight = 0.0;
+    final TreeMap<Double, Pose2d> visionUpdateQueue = new TreeMap<>();
 
     double maxError = Double.NEGATIVE_INFINITY;
     double errorSum = 0;
-    Trajectory.State groundtruthState;
-    DifferentialDriveWheelSpeeds input;
-    while (t <= traj.getTotalTimeSeconds()) {
-      groundtruthState = traj.sample(t);
-      input =
-          kinematics.toWheelSpeeds(
-              new ChassisSpeeds(
-                  groundtruthState.velocityMetersPerSecond,
-                  0.0,
-                  // ds/dt * dtheta/ds = dtheta/dt
-                  groundtruthState.velocityMetersPerSecond
-                      * groundtruthState.curvatureRadPerMeter));
+    while (t <= trajectory.getTotalTimeSeconds()) {
+      var groundTruthState = trajectory.sample(t);
 
-      if (lastVisionUpdateTime + visionUpdateRate + rand.nextGaussian() * 0.4 < t) {
-        if (lastVisionPose != null) {
-          estimator.addVisionMeasurement(lastVisionPose, lastVisionUpdateTime);
-        }
-        var groundPose = groundtruthState.poseMeters;
-        lastVisionPose =
-            new Pose2d(
-                new Translation2d(
-                    groundPose.getTranslation().getX() + rand.nextGaussian() * 0.1,
-                    groundPose.getTranslation().getY() + rand.nextGaussian() * 0.1),
-                new Rotation2d(rand.nextGaussian() * 0.01).plus(groundPose.getRotation()));
-        lastVisionUpdateTime = t;
+      // We are due for a new vision measurement if it's been `visionUpdateRate` seconds since the
+      // last vision measurement
+      if (visionUpdateQueue.isEmpty() || visionUpdateQueue.lastKey() + visionUpdateRate < t) {
+        Pose2d newVisionPose =
+            visionMeasurementGenerator
+                .apply(groundTruthState)
+                .plus(
+                    new Transform2d(
+                        new Translation2d(rand.nextGaussian() * 0.1, rand.nextGaussian() * 0.1),
+                        new Rotation2d(rand.nextGaussian() * 0.05)));
+
+        visionUpdateQueue.put(t, newVisionPose);
       }
 
-      input.leftMetersPerSecond += rand.nextGaussian() * 0.01;
-      input.rightMetersPerSecond += rand.nextGaussian() * 0.01;
+      // We should apply the oldest vision measurement if it has been `visionUpdateDelay` seconds
+      // since it was measured
+      if (!visionUpdateQueue.isEmpty() && visionUpdateQueue.firstKey() + visionUpdateDelay < t) {
+        var visionEntry = visionUpdateQueue.pollFirstEntry();
+        estimator.addVisionMeasurement(visionEntry.getValue(), visionEntry.getKey());
+      }
 
-      distanceLeft += input.leftMetersPerSecond * dt;
-      distanceRight += input.rightMetersPerSecond * dt;
+      var chassisSpeeds = chassisSpeedsGenerator.apply(groundTruthState);
+
+      var wheelSpeeds = kinematics.toWheelSpeeds(chassisSpeeds);
+
+      leftDistanceMeters += wheelSpeeds.leftMetersPerSecond * dt;
+      rightDistanceMeters += wheelSpeeds.rightMetersPerSecond * dt;
 
-      var rotNoise = new Rotation2d(rand.nextGaussian() * 0.001);
       var xHat =
           estimator.updateWithTime(
               t,
-              groundtruthState.poseMeters.getRotation().plus(rotNoise),
-              input,
-              distanceLeft,
-              distanceRight);
+              groundTruthState
+                  .poseMeters
+                  .getRotation()
+                  .plus(new Rotation2d(rand.nextGaussian() * 0.05))
+                  .minus(trajectory.getInitialPose().getRotation()),
+              leftDistanceMeters,
+              rightDistanceMeters);
+
+      System.out.printf(
+          "%f, %f, %f, %f, %f, %f, %f\n",
+          t,
+          xHat.getX(),
+          xHat.getY(),
+          xHat.getRotation().getRadians(),
+          groundTruthState.poseMeters.getX(),
+          groundTruthState.poseMeters.getY(),
+          groundTruthState.poseMeters.getRotation().getRadians());
 
       double error =
-          groundtruthState.poseMeters.getTranslation().getDistance(xHat.getTranslation());
+          groundTruthState.poseMeters.getTranslation().getDistance(xHat.getTranslation());
       if (error > maxError) {
         maxError = error;
       }
@@ -111,7 +211,20 @@ class DifferentialDrivePoseEstimatorTest {
       t += dt;
     }
 
-    assertEquals(0.0, errorSum / (traj.getTotalTimeSeconds() / dt), 0.05, "Incorrect mean error");
-    assertEquals(0.0, maxError, 0.125, "Incorrect max error");
+    assertEquals(
+        endingPose.getX(), estimator.getEstimatedPosition().getX(), 0.08, "Incorrect Final X");
+    assertEquals(
+        endingPose.getY(), estimator.getEstimatedPosition().getY(), 0.08, "Incorrect Final Y");
+    assertEquals(
+        endingPose.getRotation().getRadians(),
+        estimator.getEstimatedPosition().getRotation().getRadians(),
+        0.15,
+        "Incorrect Final Theta");
+
+    if (checkError) {
+      assertEquals(
+          0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.07, "Incorrect mean error");
+      assertEquals(0.0, maxError, 0.2, "Incorrect max error");
+    }
   }
 }

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

@@ -9,14 +9,18 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 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.Transform2d;
 import edu.wpi.first.math.geometry.Translation2d;
 import edu.wpi.first.math.kinematics.ChassisSpeeds;
 import edu.wpi.first.math.kinematics.MecanumDriveKinematics;
 import edu.wpi.first.math.kinematics.MecanumDriveWheelPositions;
+import edu.wpi.first.math.trajectory.Trajectory;
 import edu.wpi.first.math.trajectory.TrajectoryConfig;
 import edu.wpi.first.math.trajectory.TrajectoryGenerator;
 import java.util.List;
 import java.util.Random;
+import java.util.TreeMap;
+import java.util.function.Function;
 import org.junit.jupiter.api.Test;
 
 class MecanumDrivePoseEstimatorTest {
@@ -31,68 +35,155 @@ class MecanumDrivePoseEstimatorTest {
 
     var estimator =
         new MecanumDrivePoseEstimator(
+            kinematics,
             new Rotation2d(),
             wheelPositions,
             new Pose2d(),
-            kinematics,
-            VecBuilder.fill(0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1),
-            VecBuilder.fill(0.05, 0.05, 0.05, 0.05, 0.05),
-            VecBuilder.fill(0.1, 0.1, 0.1));
+            VecBuilder.fill(0.1, 0.1, 0.1),
+            VecBuilder.fill(0.45, 0.45, 0.1));
 
     var trajectory =
         TrajectoryGenerator.generateTrajectory(
             List.of(
-                new Pose2d(),
-                new Pose2d(20, 20, Rotation2d.fromDegrees(0)),
-                new Pose2d(10, 10, Rotation2d.fromDegrees(180)),
-                new Pose2d(30, 30, Rotation2d.fromDegrees(0)),
-                new Pose2d(20, 20, Rotation2d.fromDegrees(180)),
-                new Pose2d(10, 10, Rotation2d.fromDegrees(0))),
-            new TrajectoryConfig(0.5, 2));
+                new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
+                new Pose2d(3, 0, Rotation2d.fromDegrees(-90)),
+                new Pose2d(0, 0, Rotation2d.fromDegrees(135)),
+                new Pose2d(-3, 0, Rotation2d.fromDegrees(-90)),
+                new Pose2d(0, 0, Rotation2d.fromDegrees(45))),
+            new TrajectoryConfig(2, 2));
 
-    var rand = new Random(5190);
+    testFollowTrajectory(
+        kinematics,
+        estimator,
+        trajectory,
+        state ->
+            new ChassisSpeeds(
+                state.velocityMetersPerSecond,
+                0,
+                state.velocityMetersPerSecond * state.curvatureRadPerMeter),
+        state -> state.poseMeters,
+        trajectory.getInitialPose(),
+        new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
+        0.02,
+        0.1,
+        0.25,
+        true);
+  }
+
+  @Test
+  void testBadInitialPose() {
+    var kinematics =
+        new MecanumDriveKinematics(
+            new Translation2d(1, 1), new Translation2d(1, -1),
+            new Translation2d(-1, -1), new Translation2d(-1, 1));
+
+    var wheelPositions = new MecanumDriveWheelPositions();
+
+    var estimator =
+        new MecanumDrivePoseEstimator(
+            kinematics,
+            new Rotation2d(),
+            wheelPositions,
+            new Pose2d(),
+            VecBuilder.fill(0.1, 0.1, 0.1),
+            VecBuilder.fill(0.45, 0.45, 0.1));
+
+    var trajectory =
+        TrajectoryGenerator.generateTrajectory(
+            List.of(
+                new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
+                new Pose2d(3, 0, Rotation2d.fromDegrees(-90)),
+                new Pose2d(0, 0, Rotation2d.fromDegrees(135)),
+                new Pose2d(-3, 0, Rotation2d.fromDegrees(-90)),
+                new Pose2d(0, 0, Rotation2d.fromDegrees(45))),
+            new TrajectoryConfig(2, 2));
+
+    for (int offset_direction_degs = 0; offset_direction_degs < 360; offset_direction_degs += 45) {
+      for (int offset_heading_degs = 0; offset_heading_degs < 360; offset_heading_degs += 45) {
+        var pose_offset = Rotation2d.fromDegrees(offset_direction_degs);
+        var heading_offset = Rotation2d.fromDegrees(offset_heading_degs);
+
+        var initial_pose =
+            trajectory
+                .getInitialPose()
+                .plus(
+                    new Transform2d(
+                        new Translation2d(pose_offset.getCos(), pose_offset.getSin()),
+                        heading_offset));
+
+        testFollowTrajectory(
+            kinematics,
+            estimator,
+            trajectory,
+            state ->
+                new ChassisSpeeds(
+                    state.velocityMetersPerSecond,
+                    0,
+                    state.velocityMetersPerSecond * state.curvatureRadPerMeter),
+            state -> state.poseMeters,
+            initial_pose,
+            new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
+            0.02,
+            0.1,
+            0.25,
+            false);
+      }
+    }
+  }
+
+  void testFollowTrajectory(
+      final MecanumDriveKinematics kinematics,
+      final MecanumDrivePoseEstimator estimator,
+      final Trajectory trajectory,
+      final Function<Trajectory.State, ChassisSpeeds> chassisSpeedsGenerator,
+      final Function<Trajectory.State, Pose2d> visionMeasurementGenerator,
+      final Pose2d startingPose,
+      final Pose2d endingPose,
+      final double dt,
+      final double visionUpdateRate,
+      final double visionUpdateDelay,
+      final boolean checkError) {
+    var wheelPositions = new MecanumDriveWheelPositions();
+
+    estimator.resetPosition(new Rotation2d(), wheelPositions, startingPose);
+
+    var rand = new Random(3538);
 
-    final double dt = 0.02;
     double t = 0.0;
 
-    final double visionUpdateRate = 0.1;
-    Pose2d lastVisionPose = null;
-    double lastVisionUpdateTime = Double.NEGATIVE_INFINITY;
+    System.out.print("time, est_x, est_y, est_theta, true_x, true_y, true_theta\n");
+
+    final TreeMap<Double, Pose2d> visionUpdateQueue = new TreeMap<>();
 
     double maxError = Double.NEGATIVE_INFINITY;
     double errorSum = 0;
     while (t <= trajectory.getTotalTimeSeconds()) {
       var groundTruthState = trajectory.sample(t);
 
-      if (lastVisionUpdateTime + visionUpdateRate < t) {
-        if (lastVisionPose != null) {
-          estimator.addVisionMeasurement(lastVisionPose, lastVisionUpdateTime);
-        }
+      // We are due for a new vision measurement if it's been `visionUpdateRate` seconds since the
+      // last vision measurement
+      if (visionUpdateQueue.isEmpty() || visionUpdateQueue.lastKey() + visionUpdateRate < t) {
+        Pose2d newVisionPose =
+            visionMeasurementGenerator
+                .apply(groundTruthState)
+                .plus(
+                    new Transform2d(
+                        new Translation2d(rand.nextGaussian() * 0.1, rand.nextGaussian() * 0.1),
+                        new Rotation2d(rand.nextGaussian() * 0.05)));
 
-        lastVisionPose =
-            new Pose2d(
-                new Translation2d(
-                    groundTruthState.poseMeters.getTranslation().getX() + rand.nextGaussian() * 0.1,
-                    groundTruthState.poseMeters.getTranslation().getY()
-                        + rand.nextGaussian() * 0.1),
-                new Rotation2d(rand.nextGaussian() * 0.1)
-                    .plus(groundTruthState.poseMeters.getRotation()));
-
-        lastVisionUpdateTime = t;
+        visionUpdateQueue.put(t, newVisionPose);
       }
 
-      var wheelSpeeds =
-          kinematics.toWheelSpeeds(
-              new ChassisSpeeds(
-                  groundTruthState.velocityMetersPerSecond,
-                  0,
-                  groundTruthState.velocityMetersPerSecond
-                      * groundTruthState.curvatureRadPerMeter));
+      // We should apply the oldest vision measurement if it has been `visionUpdateDelay` seconds
+      // since it was measured
+      if (!visionUpdateQueue.isEmpty() && visionUpdateQueue.firstKey() + visionUpdateDelay < t) {
+        var visionEntry = visionUpdateQueue.pollFirstEntry();
+        estimator.addVisionMeasurement(visionEntry.getValue(), visionEntry.getKey());
+      }
 
-      wheelSpeeds.frontLeftMetersPerSecond += rand.nextGaussian() * 0.1;
-      wheelSpeeds.frontRightMetersPerSecond += rand.nextGaussian() * 0.1;
-      wheelSpeeds.rearLeftMetersPerSecond += rand.nextGaussian() * 0.1;
-      wheelSpeeds.rearRightMetersPerSecond += rand.nextGaussian() * 0.1;
+      var chassisSpeeds = chassisSpeedsGenerator.apply(groundTruthState);
+
+      var wheelSpeeds = kinematics.toWheelSpeeds(chassisSpeeds);
 
       wheelPositions.frontLeftMeters += wheelSpeeds.frontLeftMetersPerSecond * dt;
       wheelPositions.frontRightMeters += wheelSpeeds.frontRightMetersPerSecond * dt;
@@ -105,108 +196,19 @@ class MecanumDrivePoseEstimatorTest {
               groundTruthState
                   .poseMeters
                   .getRotation()
-                  .plus(new Rotation2d(rand.nextGaussian() * 0.05)),
-              wheelSpeeds,
+                  .plus(new Rotation2d(rand.nextGaussian() * 0.05))
+                  .minus(trajectory.getInitialPose().getRotation()),
               wheelPositions);
 
-      double error =
-          groundTruthState.poseMeters.getTranslation().getDistance(xHat.getTranslation());
-      if (error > maxError) {
-        maxError = error;
-      }
-      errorSum += error;
-
-      t += dt;
-    }
-
-    assertEquals(
-        0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.05, "Incorrect mean error");
-    assertEquals(0.0, maxError, 0.125, "Incorrect max error");
-  }
-
-  @Test
-  void testAccuracyFacingXAxis() {
-    var kinematics =
-        new MecanumDriveKinematics(
-            new Translation2d(1, 1), new Translation2d(1, -1),
-            new Translation2d(-1, -1), new Translation2d(-1, 1));
-
-    var wheelPositions = new MecanumDriveWheelPositions();
-
-    var estimator =
-        new MecanumDrivePoseEstimator(
-            new Rotation2d(),
-            wheelPositions,
-            new Pose2d(),
-            kinematics,
-            VecBuilder.fill(0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1),
-            VecBuilder.fill(0.05, 0.05, 0.05, 0.05, 0.05),
-            VecBuilder.fill(0.1, 0.1, 0.1));
-
-    var trajectory =
-        TrajectoryGenerator.generateTrajectory(
-            List.of(
-                new Pose2d(),
-                new Pose2d(20, 20, Rotation2d.fromDegrees(0)),
-                new Pose2d(10, 10, Rotation2d.fromDegrees(180)),
-                new Pose2d(30, 30, Rotation2d.fromDegrees(0)),
-                new Pose2d(20, 20, Rotation2d.fromDegrees(180)),
-                new Pose2d(10, 10, Rotation2d.fromDegrees(0))),
-            new TrajectoryConfig(0.5, 2));
-
-    var rand = new Random(5190);
-
-    final double dt = 0.02;
-    double t = 0.0;
-
-    final double visionUpdateRate = 0.1;
-    Pose2d lastVisionPose = null;
-    double lastVisionUpdateTime = Double.NEGATIVE_INFINITY;
-
-    double maxError = Double.NEGATIVE_INFINITY;
-    double errorSum = 0;
-    while (t <= trajectory.getTotalTimeSeconds()) {
-      var groundTruthState = trajectory.sample(t);
-
-      if (lastVisionUpdateTime + visionUpdateRate < t) {
-        if (lastVisionPose != null) {
-          estimator.addVisionMeasurement(lastVisionPose, lastVisionUpdateTime);
-        }
-
-        lastVisionPose =
-            new Pose2d(
-                new Translation2d(
-                    groundTruthState.poseMeters.getTranslation().getX() + rand.nextGaussian() * 0.1,
-                    groundTruthState.poseMeters.getTranslation().getY()
-                        + rand.nextGaussian() * 0.1),
-                new Rotation2d(rand.nextGaussian() * 0.1)
-                    .plus(groundTruthState.poseMeters.getRotation()));
-
-        lastVisionUpdateTime = t;
-      }
-
-      var wheelSpeeds =
-          kinematics.toWheelSpeeds(
-              new ChassisSpeeds(
-                  groundTruthState.velocityMetersPerSecond
-                      * groundTruthState.poseMeters.getRotation().getCos(),
-                  groundTruthState.velocityMetersPerSecond
-                      * groundTruthState.poseMeters.getRotation().getSin(),
-                  0));
-
-      wheelSpeeds.frontLeftMetersPerSecond += rand.nextGaussian() * 0.1;
-      wheelSpeeds.frontRightMetersPerSecond += rand.nextGaussian() * 0.1;
-      wheelSpeeds.rearLeftMetersPerSecond += rand.nextGaussian() * 0.1;
-      wheelSpeeds.rearRightMetersPerSecond += rand.nextGaussian() * 0.1;
-
-      wheelPositions.frontLeftMeters += wheelSpeeds.frontLeftMetersPerSecond * dt;
-      wheelPositions.frontRightMeters += wheelSpeeds.frontRightMetersPerSecond * dt;
-      wheelPositions.rearLeftMeters += wheelSpeeds.rearLeftMetersPerSecond * dt;
-      wheelPositions.rearRightMeters += wheelSpeeds.rearRightMetersPerSecond * dt;
-
-      var xHat =
-          estimator.updateWithTime(
-              t, new Rotation2d(rand.nextGaussian() * 0.05), wheelSpeeds, wheelPositions);
+      System.out.printf(
+          "%f, %f, %f, %f, %f, %f, %f\n",
+          t,
+          xHat.getX(),
+          xHat.getY(),
+          xHat.getRotation().getRadians(),
+          groundTruthState.poseMeters.getX(),
+          groundTruthState.poseMeters.getY(),
+          groundTruthState.poseMeters.getRotation().getRadians());
 
       double error =
           groundTruthState.poseMeters.getTranslation().getDistance(xHat.getTranslation());
@@ -219,7 +221,19 @@ class MecanumDrivePoseEstimatorTest {
     }
 
     assertEquals(
-        0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.05, "Incorrect mean error");
-    assertEquals(0.0, maxError, 0.125, "Incorrect max error");
+        endingPose.getX(), estimator.getEstimatedPosition().getX(), 0.08, "Incorrect Final X");
+    assertEquals(
+        endingPose.getY(), estimator.getEstimatedPosition().getY(), 0.08, "Incorrect Final Y");
+    assertEquals(
+        endingPose.getRotation().getRadians(),
+        estimator.getEstimatedPosition().getRotation().getRadians(),
+        0.15,
+        "Incorrect Final Theta");
+
+    if (checkError) {
+      assertEquals(
+          0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.07, "Incorrect mean error");
+      assertEquals(0.0, maxError, 0.2, "Incorrect max error");
+    }
   }
 }

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

@@ -6,20 +6,21 @@ package edu.wpi.first.math.estimator;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
-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.Transform2d;
 import edu.wpi.first.math.geometry.Translation2d;
 import edu.wpi.first.math.kinematics.ChassisSpeeds;
 import edu.wpi.first.math.kinematics.SwerveDriveKinematics;
 import edu.wpi.first.math.kinematics.SwerveModulePosition;
-import edu.wpi.first.math.numbers.N5;
-import edu.wpi.first.math.numbers.N7;
+import edu.wpi.first.math.trajectory.Trajectory;
 import edu.wpi.first.math.trajectory.TrajectoryConfig;
 import edu.wpi.first.math.trajectory.TrajectoryGenerator;
 import java.util.List;
 import java.util.Random;
+import java.util.TreeMap;
+import java.util.function.Function;
 import org.junit.jupiter.api.Test;
 
 class SwerveDrivePoseEstimatorTest {
@@ -38,17 +39,13 @@ class SwerveDrivePoseEstimatorTest {
     var br = new SwerveModulePosition();
 
     var estimator =
-        new SwerveDrivePoseEstimator<N7, N7, N5>(
-            Nat.N7(),
-            Nat.N7(),
-            Nat.N5(),
+        new SwerveDrivePoseEstimator(
+            kinematics,
             new Rotation2d(),
             new SwerveModulePosition[] {fl, fr, bl, br},
             new Pose2d(),
-            kinematics,
-            VecBuilder.fill(0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1),
-            VecBuilder.fill(0.005, 0.005, 0.005, 0.005, 0.005),
-            VecBuilder.fill(0.1, 0.1, 0.1));
+            VecBuilder.fill(0.1, 0.1, 0.1),
+            VecBuilder.fill(0.5, 0.5, 0.5));
 
     var trajectory =
         TrajectoryGenerator.generateTrajectory(
@@ -58,88 +55,28 @@ class SwerveDrivePoseEstimatorTest {
                 new Pose2d(0, 0, Rotation2d.fromDegrees(135)),
                 new Pose2d(-3, 0, Rotation2d.fromDegrees(-90)),
                 new Pose2d(0, 0, Rotation2d.fromDegrees(45))),
-            new TrajectoryConfig(0.5, 2));
+            new TrajectoryConfig(2, 2));
 
-    var rand = new Random(4915);
-
-    final double dt = 0.02;
-    double t = 0.0;
-
-    final double visionUpdateRate = 0.1;
-    Pose2d lastVisionPose = null;
-    double lastVisionUpdateTime = Double.NEGATIVE_INFINITY;
-
-    double maxError = Double.NEGATIVE_INFINITY;
-    double errorSum = 0;
-    while (t <= trajectory.getTotalTimeSeconds()) {
-      var groundTruthState = trajectory.sample(t);
-
-      if (lastVisionUpdateTime + visionUpdateRate < t) {
-        if (lastVisionPose != null) {
-          estimator.addVisionMeasurement(lastVisionPose, lastVisionUpdateTime);
-        }
-
-        lastVisionPose =
-            new Pose2d(
-                new Translation2d(
-                    groundTruthState.poseMeters.getTranslation().getX() + rand.nextGaussian() * 0.1,
-                    groundTruthState.poseMeters.getTranslation().getY()
-                        + rand.nextGaussian() * 0.1),
-                new Rotation2d(rand.nextGaussian() * 0.1)
-                    .plus(groundTruthState.poseMeters.getRotation()));
-
-        lastVisionUpdateTime = t;
-      }
-
-      var moduleStates =
-          kinematics.toSwerveModuleStates(
-              new ChassisSpeeds(
-                  groundTruthState.velocityMetersPerSecond,
-                  0.0,
-                  groundTruthState.velocityMetersPerSecond
-                      * groundTruthState.curvatureRadPerMeter));
-      for (var moduleState : moduleStates) {
-        moduleState.angle = moduleState.angle.plus(new Rotation2d(rand.nextGaussian() * 0.005));
-        moduleState.speedMetersPerSecond += rand.nextGaussian() * 0.1;
-      }
-
-      fl.distanceMeters += moduleStates[0].speedMetersPerSecond * dt;
-      fr.distanceMeters += moduleStates[1].speedMetersPerSecond * dt;
-      bl.distanceMeters += moduleStates[2].speedMetersPerSecond * dt;
-      br.distanceMeters += moduleStates[3].speedMetersPerSecond * dt;
-
-      fl.angle = moduleStates[0].angle;
-      fr.angle = moduleStates[1].angle;
-      bl.angle = moduleStates[2].angle;
-      br.angle = moduleStates[3].angle;
-
-      var xHat =
-          estimator.updateWithTime(
-              t,
-              groundTruthState
-                  .poseMeters
-                  .getRotation()
-                  .plus(new Rotation2d(rand.nextGaussian() * 0.05)),
-              moduleStates,
-              new SwerveModulePosition[] {fl, fr, bl, br});
-
-      double error =
-          groundTruthState.poseMeters.getTranslation().getDistance(xHat.getTranslation());
-      if (error > maxError) {
-        maxError = error;
-      }
-      errorSum += error;
-
-      t += dt;
-    }
-
-    assertEquals(
-        0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.05, "Incorrect mean error");
-    assertEquals(0.0, maxError, 0.125, "Incorrect max error");
+    testFollowTrajectory(
+        kinematics,
+        estimator,
+        trajectory,
+        state ->
+            new ChassisSpeeds(
+                state.velocityMetersPerSecond,
+                0,
+                state.velocityMetersPerSecond * state.curvatureRadPerMeter),
+        state -> state.poseMeters,
+        trajectory.getInitialPose(),
+        new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
+        0.02,
+        0.1,
+        0.25,
+        true);
   }
 
   @Test
-  void testAccuracyFacingXAxis() {
+  void testBadInitialPose() {
     var kinematics =
         new SwerveDriveKinematics(
             new Translation2d(1, 1),
@@ -153,18 +90,13 @@ class SwerveDrivePoseEstimatorTest {
     var br = new SwerveModulePosition();
 
     var estimator =
-        new SwerveDrivePoseEstimator<N7, N7, N5>(
-            Nat.N7(),
-            Nat.N7(),
-            Nat.N5(),
+        new SwerveDrivePoseEstimator(
+            kinematics,
             new Rotation2d(),
             new SwerveModulePosition[] {fl, fr, bl, br},
-            new Pose2d(),
-            kinematics,
-            VecBuilder.fill(0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1),
-            VecBuilder.fill(0.005, 0.005, 0.005, 0.005, 0.005),
-            VecBuilder.fill(0.1, 0.1, 0.1));
-
+            new Pose2d(-1, -1, Rotation2d.fromRadians(-1)),
+            VecBuilder.fill(0.1, 0.1, 0.1),
+            VecBuilder.fill(0.9, 0.9, 0.9));
     var trajectory =
         TrajectoryGenerator.generateTrajectory(
             List.of(
@@ -173,76 +105,137 @@ class SwerveDrivePoseEstimatorTest {
                 new Pose2d(0, 0, Rotation2d.fromDegrees(135)),
                 new Pose2d(-3, 0, Rotation2d.fromDegrees(-90)),
                 new Pose2d(0, 0, Rotation2d.fromDegrees(45))),
-            new TrajectoryConfig(0.5, 2));
+            new TrajectoryConfig(2, 2));
 
-    var rand = new Random(4915);
+    for (int offset_direction_degs = 0; offset_direction_degs < 360; offset_direction_degs += 45) {
+      for (int offset_heading_degs = 0; offset_heading_degs < 360; offset_heading_degs += 45) {
+        var pose_offset = Rotation2d.fromDegrees(offset_direction_degs);
+        var heading_offset = Rotation2d.fromDegrees(offset_heading_degs);
+
+        var initial_pose =
+            trajectory
+                .getInitialPose()
+                .plus(
+                    new Transform2d(
+                        new Translation2d(pose_offset.getCos(), pose_offset.getSin()),
+                        heading_offset));
+
+        testFollowTrajectory(
+            kinematics,
+            estimator,
+            trajectory,
+            state ->
+                new ChassisSpeeds(
+                    state.velocityMetersPerSecond,
+                    0,
+                    state.velocityMetersPerSecond * state.curvatureRadPerMeter),
+            state -> state.poseMeters,
+            initial_pose,
+            new Pose2d(0, 0, Rotation2d.fromDegrees(45)),
+            0.02,
+            0.1,
+            1.0,
+            false);
+      }
+    }
+  }
+
+  void testFollowTrajectory(
+      final SwerveDriveKinematics kinematics,
+      final SwerveDrivePoseEstimator estimator,
+      final Trajectory trajectory,
+      final Function<Trajectory.State, ChassisSpeeds> chassisSpeedsGenerator,
+      final Function<Trajectory.State, Pose2d> visionMeasurementGenerator,
+      final Pose2d startingPose,
+      final Pose2d endingPose,
+      final double dt,
+      final double visionUpdateRate,
+      final double visionUpdateDelay,
+      final boolean checkError) {
+    SwerveModulePosition[] positions = {
+      new SwerveModulePosition(),
+      new SwerveModulePosition(),
+      new SwerveModulePosition(),
+      new SwerveModulePosition()
+    };
+
+    estimator.resetPosition(new Rotation2d(), positions, startingPose);
+
+    var rand = new Random(3538);
 
-    final double dt = 0.02;
     double t = 0.0;
 
-    final double visionUpdateRate = 0.1;
-    Pose2d lastVisionPose = null;
-    double lastVisionUpdateTime = Double.NEGATIVE_INFINITY;
+    System.out.print(
+        "time, est_x, est_y, est_theta, true_x, true_y, true_theta, "
+            + "distance_1, distance_2, distance_3, distance_4, "
+            + "angle_1, angle_2, angle_3, angle_4\n");
+
+    final TreeMap<Double, Pose2d> visionUpdateQueue = new TreeMap<>();
 
     double maxError = Double.NEGATIVE_INFINITY;
     double errorSum = 0;
     while (t <= trajectory.getTotalTimeSeconds()) {
       var groundTruthState = trajectory.sample(t);
 
-      if (lastVisionUpdateTime + visionUpdateRate < t) {
-        if (lastVisionPose != null) {
-          estimator.addVisionMeasurement(lastVisionPose, lastVisionUpdateTime);
-        }
+      // We are due for a new vision measurement if it's been `visionUpdateRate` seconds since the
+      // last vision measurement
+      if (visionUpdateQueue.isEmpty() || visionUpdateQueue.lastKey() + visionUpdateRate < t) {
+        Pose2d newVisionPose =
+            visionMeasurementGenerator
+                .apply(groundTruthState)
+                .plus(
+                    new Transform2d(
+                        new Translation2d(rand.nextGaussian() * 0.1, rand.nextGaussian() * 0.1),
+                        new Rotation2d(rand.nextGaussian() * 0.05)));
 
-        lastVisionPose =
-            new Pose2d(
-                new Translation2d(
-                    groundTruthState.poseMeters.getTranslation().getX() + rand.nextGaussian() * 0.1,
-                    groundTruthState.poseMeters.getTranslation().getY()
-                        + rand.nextGaussian() * 0.1),
-                new Rotation2d(rand.nextGaussian() * 0.1)
-                    .plus(groundTruthState.poseMeters.getRotation()));
-
-        lastVisionUpdateTime = t;
+        visionUpdateQueue.put(t, newVisionPose);
       }
 
-      var moduleStates =
-          kinematics.toSwerveModuleStates(
-              new ChassisSpeeds(
-                  groundTruthState.velocityMetersPerSecond
-                      * groundTruthState.poseMeters.getRotation().getCos(),
-                  groundTruthState.velocityMetersPerSecond
-                      * groundTruthState.poseMeters.getRotation().getSin(),
-                  0.0));
-      for (var moduleState : moduleStates) {
-        moduleState.angle = moduleState.angle.plus(new Rotation2d(rand.nextGaussian() * 0.005));
-        moduleState.speedMetersPerSecond += rand.nextGaussian() * 0.1;
+      // We should apply the oldest vision measurement if it has been `visionUpdateDelay` seconds
+      // since it was measured
+      if (!visionUpdateQueue.isEmpty() && visionUpdateQueue.firstKey() + visionUpdateDelay < t) {
+        var visionEntry = visionUpdateQueue.pollFirstEntry();
+        estimator.addVisionMeasurement(visionEntry.getValue(), visionEntry.getKey());
       }
 
-      fl.distanceMeters +=
-          groundTruthState.velocityMetersPerSecond * dt
-              + 0.5 * groundTruthState.accelerationMetersPerSecondSq * dt * dt;
-      fr.distanceMeters +=
-          groundTruthState.velocityMetersPerSecond * dt
-              + 0.5 * groundTruthState.accelerationMetersPerSecondSq * dt * dt;
-      bl.distanceMeters +=
-          groundTruthState.velocityMetersPerSecond * dt
-              + 0.5 * groundTruthState.accelerationMetersPerSecondSq * dt * dt;
-      br.distanceMeters +=
-          groundTruthState.velocityMetersPerSecond * dt
-              + 0.5 * groundTruthState.accelerationMetersPerSecondSq * dt * dt;
+      var chassisSpeeds = chassisSpeedsGenerator.apply(groundTruthState);
 
-      fl.angle = groundTruthState.poseMeters.getRotation();
-      fr.angle = groundTruthState.poseMeters.getRotation();
-      bl.angle = groundTruthState.poseMeters.getRotation();
-      br.angle = groundTruthState.poseMeters.getRotation();
+      var moduleStates = kinematics.toSwerveModuleStates(chassisSpeeds);
+
+      for (int i = 0; i < moduleStates.length; i++) {
+        positions[i].distanceMeters +=
+            moduleStates[i].speedMetersPerSecond * (1 - rand.nextGaussian() * 0.05) * dt;
+        positions[i].angle =
+            moduleStates[i].angle.plus(new Rotation2d(rand.nextGaussian() * 0.005));
+      }
 
       var xHat =
           estimator.updateWithTime(
               t,
-              new Rotation2d(rand.nextGaussian() * 0.05),
-              moduleStates,
-              new SwerveModulePosition[] {fl, fr, bl, br});
+              groundTruthState
+                  .poseMeters
+                  .getRotation()
+                  .plus(new Rotation2d(rand.nextGaussian() * 0.05))
+                  .minus(trajectory.getInitialPose().getRotation()),
+              positions);
+
+      System.out.printf(
+          "%f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f\n",
+          t,
+          xHat.getX(),
+          xHat.getY(),
+          xHat.getRotation().getRadians(),
+          groundTruthState.poseMeters.getX(),
+          groundTruthState.poseMeters.getY(),
+          groundTruthState.poseMeters.getRotation().getRadians(),
+          positions[0].distanceMeters,
+          positions[1].distanceMeters,
+          positions[2].distanceMeters,
+          positions[3].distanceMeters,
+          positions[0].angle.getRadians(),
+          positions[1].angle.getRadians(),
+          positions[2].angle.getRadians(),
+          positions[3].angle.getRadians());
 
       double error =
           groundTruthState.poseMeters.getTranslation().getDistance(xHat.getTranslation());
@@ -255,7 +248,19 @@ class SwerveDrivePoseEstimatorTest {
     }
 
     assertEquals(
-        0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.05, "Incorrect mean error");
-    assertEquals(0.0, maxError, 0.125, "Incorrect max error");
+        endingPose.getX(), estimator.getEstimatedPosition().getX(), 0.08, "Incorrect Final X");
+    assertEquals(
+        endingPose.getY(), estimator.getEstimatedPosition().getY(), 0.08, "Incorrect Final Y");
+    assertEquals(
+        endingPose.getRotation().getRadians(),
+        estimator.getEstimatedPosition().getRotation().getRadians(),
+        0.15,
+        "Incorrect Final Theta");
+
+    if (checkError) {
+      assertEquals(
+          0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.07, "Incorrect mean error");
+      assertEquals(0.0, maxError, 0.2, "Incorrect max error");
+    }
   }
 }

wpimath/src/test/java/edu/wpi/first/math/kinematics/SwerveDriveOdometryTest.java

@@ -178,6 +178,14 @@ class SwerveDriveOdometryTest {
       t += dt;
     }
 
+    assertEquals(0.0, odometry.getPoseMeters().getX(), 1e-1, "Incorrect Final X");
+    assertEquals(0.0, odometry.getPoseMeters().getY(), 1e-1, "Incorrect Final Y");
+    assertEquals(
+        Math.PI / 4,
+        odometry.getPoseMeters().getRotation().getRadians(),
+        10 * Math.PI / 180,
+        "Incorrect Final Theta");
+
     assertEquals(
         0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.05, "Incorrect mean error");
     assertEquals(0.0, maxError, 0.125, "Incorrect max error");
@@ -253,6 +261,14 @@ class SwerveDriveOdometryTest {
       t += dt;
     }
 
+    assertEquals(0.0, odometry.getPoseMeters().getX(), 1e-1, "Incorrect Final X");
+    assertEquals(0.0, odometry.getPoseMeters().getY(), 1e-1, "Incorrect Final Y");
+    assertEquals(
+        0.0,
+        odometry.getPoseMeters().getRotation().getRadians(),
+        10 * Math.PI / 180,
+        "Incorrect Final Theta");
+
     assertEquals(
         0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.06, "Incorrect mean error");
     assertEquals(0.0, maxError, 0.125, "Incorrect max error");