[wpimath] Replace UKF implementation with square root form (#4168)

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

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

@@ -110,7 +110,7 @@ class DifferentialDrivePoseEstimatorTest {
       t += dt;
     }
 
-    assertEquals(0.0, errorSum / (traj.getTotalTimeSeconds() / dt), 0.035, "Incorrect mean error");
-    assertEquals(0.0, maxError, 0.055, "Incorrect max error");
+    assertEquals(0.0, errorSum / (traj.getTotalTimeSeconds() / dt), 0.05, "Incorrect mean error");
+    assertEquals(0.0, maxError, 0.1, "Incorrect max error");
   }
 }

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

@@ -111,7 +111,7 @@ class MecanumDrivePoseEstimatorTest {
     }
 
     assertEquals(
-        0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.25, "Incorrect mean error");
-    assertEquals(0.0, maxError, 0.42, "Incorrect max error");
+        0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.05, "Incorrect mean error");
+    assertEquals(0.0, maxError, 0.1, "Incorrect max error");
   }
 }

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

@@ -16,7 +16,7 @@ class MerweScaledSigmaPointsTest {
   void testZeroMeanPoints() {
     var merweScaledSigmaPoints = new MerweScaledSigmaPoints<>(Nat.N2());
     var points =
-        merweScaledSigmaPoints.sigmaPoints(
+        merweScaledSigmaPoints.squareRootSigmaPoints(
             VecBuilder.fill(0, 0), Matrix.mat(Nat.N2(), Nat.N2()).fill(1, 0, 0, 1));
 
     assertTrue(
@@ -31,8 +31,8 @@ class MerweScaledSigmaPointsTest {
   void testNonzeroMeanPoints() {
     var merweScaledSigmaPoints = new MerweScaledSigmaPoints<>(Nat.N2());
     var points =
-        merweScaledSigmaPoints.sigmaPoints(
-            VecBuilder.fill(1, 2), Matrix.mat(Nat.N2(), Nat.N2()).fill(1, 0, 0, 10));
+        merweScaledSigmaPoints.squareRootSigmaPoints(
+            VecBuilder.fill(1, 2), Matrix.mat(Nat.N2(), Nat.N2()).fill(1, 0, 0, Math.sqrt(10)));
 
     assertTrue(
         points.isEqual(

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

@@ -110,7 +110,7 @@ class SwerveDrivePoseEstimatorTest {
     }
 
     assertEquals(
-        0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.25, "Incorrect mean error");
-    assertEquals(0.0, maxError, 0.42, "Incorrect max error");
+        0.0, errorSum / (trajectory.getTotalTimeSeconds() / dt), 0.05, "Incorrect mean error");
+    assertEquals(0.0, maxError, 0.1, "Incorrect max error");
   }
 }

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

@@ -16,8 +16,8 @@ import edu.wpi.first.math.geometry.Pose2d;
 import edu.wpi.first.math.geometry.Rotation2d;
 import edu.wpi.first.math.numbers.N1;
 import edu.wpi.first.math.numbers.N2;
-import edu.wpi.first.math.numbers.N4;
-import edu.wpi.first.math.numbers.N6;
+import edu.wpi.first.math.numbers.N3;
+import edu.wpi.first.math.numbers.N5;
 import edu.wpi.first.math.system.Discretization;
 import edu.wpi.first.math.system.NumericalIntegration;
 import edu.wpi.first.math.system.NumericalJacobian;
@@ -31,91 +31,111 @@ import org.junit.jupiter.api.Test;
 
 class UnscentedKalmanFilterTest {
   @SuppressWarnings({"LocalVariableName", "ParameterName"})
-  private static Matrix<N6, N1> getDynamics(Matrix<N6, N1> x, Matrix<N2, N1> u) {
+  private static Matrix<N5, N1> getDynamics(Matrix<N5, N1> x, Matrix<N2, N1> u) {
     var motors = DCMotor.getCIM(2);
 
-    var gHigh = 7.08;
-    var rb = 0.8382 / 2.0;
-    var r = 0.0746125;
-    var m = 63.503;
-    var J = 5.6;
+    // var gLow = 15.32;    // Low gear ratio
+    var gHigh = 7.08; // High gear ratio
+    var rb = 0.8382 / 2.0; // Robot radius
+    var r = 0.0746125; // Wheel radius
+    var m = 63.503; // Robot mass
+    var J = 5.6; // Robot moment of inertia
 
     var C1 =
         -Math.pow(gHigh, 2)
             * motors.KtNMPerAmp
             / (motors.KvRadPerSecPerVolt * motors.rOhms * r * r);
     var C2 = gHigh * motors.KtNMPerAmp / (motors.rOhms * r);
+    var k1 = 1.0 / m + Math.pow(rb, 2) / J;
+    var k2 = 1.0 / m - Math.pow(rb, 2) / J;
 
-    var c = x.get(2, 0);
-    var s = x.get(3, 0);
-    var vl = x.get(4, 0);
-    var vr = x.get(5, 0);
-
+    var vl = x.get(3, 0);
+    var vr = x.get(4, 0);
     var Vl = u.get(0, 0);
     var Vr = u.get(1, 0);
 
-    var k1 = 1.0 / m + rb * rb / J;
-    var k2 = 1.0 / m - rb * rb / J;
-
-    var xvel = (vl + vr) / 2;
-    var w = (vr - vl) / (2.0 * rb);
-
+    var v = 0.5 * (vl + vr);
     return VecBuilder.fill(
-        xvel * c,
-        xvel * s,
-        -s * w,
-        c * w,
-        k1 * ((C1 * vl) + (C2 * Vl)) + k2 * ((C1 * vr) + (C2 * Vr)),
-        k2 * ((C1 * vl) + (C2 * Vl)) + k1 * ((C1 * vr) + (C2 * Vr)));
+        v * Math.cos(x.get(2, 0)),
+        v * Math.sin(x.get(2, 0)),
+        (vr - vl) / (2.0 * rb),
+        k1 * (C1 * vl + C2 * Vl) + k2 * (C1 * vr + C2 * Vr),
+        k2 * (C1 * vl + C2 * Vl) + k1 * (C1 * vr + C2 * Vr));
   }
 
   @SuppressWarnings({"PMD.UnusedFormalParameter", "ParameterName"})
-  private static Matrix<N4, N1> getLocalMeasurementModel(Matrix<N6, N1> x, Matrix<N2, N1> u) {
-    return VecBuilder.fill(x.get(2, 0), x.get(3, 0), x.get(4, 0), x.get(5, 0));
+  private static Matrix<N3, N1> getLocalMeasurementModel(Matrix<N5, N1> x, Matrix<N2, N1> u) {
+    return VecBuilder.fill(x.get(2, 0), x.get(3, 0), x.get(4, 0));
   }
 
   @SuppressWarnings({"PMD.UnusedFormalParameter", "ParameterName"})
-  private static Matrix<N6, N1> getGlobalMeasurementModel(Matrix<N6, N1> x, Matrix<N2, N1> u) {
+  private static Matrix<N5, N1> getGlobalMeasurementModel(Matrix<N5, N1> x, Matrix<N2, N1> u) {
     return x.copy();
   }
 
   @Test
   @SuppressWarnings("LocalVariableName")
   void testInit() {
+    var dtSeconds = 0.005;
     assertDoesNotThrow(
         () -> {
-          UnscentedKalmanFilter<N6, N2, N4> observer =
+          UnscentedKalmanFilter<N5, N2, N3> observer =
               new UnscentedKalmanFilter<>(
-                  Nat.N6(),
-                  Nat.N4(),
+                  Nat.N5(),
+                  Nat.N3(),
                   UnscentedKalmanFilterTest::getDynamics,
                   UnscentedKalmanFilterTest::getLocalMeasurementModel,
-                  VecBuilder.fill(0.5, 0.5, 0.7, 0.7, 1.0, 1.0),
-                  VecBuilder.fill(0.001, 0.001, 0.5, 0.5),
-                  0.00505);
+                  VecBuilder.fill(0.5, 0.5, 10.0, 1.0, 1.0),
+                  VecBuilder.fill(0.0001, 0.01, 0.01),
+                  AngleStatistics.angleMean(2),
+                  AngleStatistics.angleMean(0),
+                  AngleStatistics.angleResidual(2),
+                  AngleStatistics.angleResidual(0),
+                  AngleStatistics.angleAdd(2),
+                  dtSeconds);
 
           var u = VecBuilder.fill(12.0, 12.0);
-          observer.predict(u, 0.00505);
+          observer.predict(u, dtSeconds);
 
           var localY = getLocalMeasurementModel(observer.getXhat(), u);
           observer.correct(u, localY);
+
+          var globalY = getGlobalMeasurementModel(observer.getXhat(), u);
+          var R =
+              StateSpaceUtil.makeCovarianceMatrix(
+                  Nat.N5(), VecBuilder.fill(0.01, 0.01, 0.0001, 0.01, 0.01));
+          observer.correct(
+              Nat.N5(),
+              u,
+              globalY,
+              UnscentedKalmanFilterTest::getGlobalMeasurementModel,
+              R,
+              AngleStatistics.angleMean(2),
+              AngleStatistics.angleResidual(2),
+              AngleStatistics.angleResidual(2),
+              AngleStatistics.angleAdd(2));
         });
   }
 
   @SuppressWarnings("LocalVariableName")
   @Test
   void testConvergence() {
-    double dtSeconds = 0.00505;
+    double dtSeconds = 0.005;
     double rbMeters = 0.8382 / 2.0; // Robot radius
 
-    UnscentedKalmanFilter<N6, N2, N4> observer =
+    UnscentedKalmanFilter<N5, N2, N3> observer =
         new UnscentedKalmanFilter<>(
-            Nat.N6(),
-            Nat.N4(),
+            Nat.N5(),
+            Nat.N3(),
             UnscentedKalmanFilterTest::getDynamics,
             UnscentedKalmanFilterTest::getLocalMeasurementModel,
-            VecBuilder.fill(0.5, 0.5, 0.7, 0.7, 1.0, 1.0),
-            VecBuilder.fill(0.001, 0.001, 0.5, 0.5),
+            VecBuilder.fill(0.5, 0.5, 10.0, 1.0, 1.0),
+            VecBuilder.fill(0.0001, 0.5, 0.5),
+            AngleStatistics.angleMean(2),
+            AngleStatistics.angleMean(0),
+            AngleStatistics.angleResidual(2),
+            AngleStatistics.angleResidual(0),
+            AngleStatistics.angleAdd(2),
             dtSeconds);
 
     List<Pose2d> waypoints =
@@ -125,56 +145,52 @@ class UnscentedKalmanFilterTest {
     var trajectory =
         TrajectoryGenerator.generateTrajectory(waypoints, new TrajectoryConfig(8.8, 0.1));
 
-    Matrix<N6, N1> nextR;
+    Matrix<N5, N1> r = new Matrix<>(Nat.N5(), Nat.N1());
     Matrix<N2, N1> u = new Matrix<>(Nat.N2(), Nat.N1());
 
     var B =
         NumericalJacobian.numericalJacobianU(
-            Nat.N6(),
+            Nat.N5(),
             Nat.N2(),
             UnscentedKalmanFilterTest::getDynamics,
-            new Matrix<>(Nat.N6(), Nat.N1()),
-            u);
+            new Matrix<>(Nat.N5(), Nat.N1()),
+            new Matrix<>(Nat.N2(), Nat.N1()));
 
-    observer.setXhat(VecBuilder.fill(2.75, 22.521, 1.0, 0.0, 0.0, 0.0)); // TODO not hard code this
-
-    var ref = trajectory.sample(0.0);
-
-    Matrix<N6, N1> r =
+    observer.setXhat(
         VecBuilder.fill(
-            ref.poseMeters.getTranslation().getX(),
-            ref.poseMeters.getTranslation().getY(),
-            ref.poseMeters.getRotation().getCos(),
-            ref.poseMeters.getRotation().getSin(),
-            ref.velocityMetersPerSecond * (1 - (ref.curvatureRadPerMeter * rbMeters)),
-            ref.velocityMetersPerSecond * (1 + (ref.curvatureRadPerMeter * rbMeters)));
-    nextR = r.copy();
+            trajectory.getInitialPose().getTranslation().getX(),
+            trajectory.getInitialPose().getTranslation().getY(),
+            trajectory.getInitialPose().getRotation().getRadians(),
+            0.0,
+            0.0));
 
     var trueXhat = observer.getXhat();
 
     double totalTime = trajectory.getTotalTimeSeconds();
     for (int i = 0; i < (totalTime / dtSeconds); i++) {
-      ref = trajectory.sample(dtSeconds * i);
+      var ref = trajectory.sample(dtSeconds * i);
       double vl = ref.velocityMetersPerSecond * (1 - (ref.curvatureRadPerMeter * rbMeters));
       double vr = ref.velocityMetersPerSecond * (1 + (ref.curvatureRadPerMeter * rbMeters));
 
-      nextR.set(0, 0, ref.poseMeters.getTranslation().getX());
-      nextR.set(1, 0, ref.poseMeters.getTranslation().getY());
-      nextR.set(2, 0, ref.poseMeters.getRotation().getCos());
-      nextR.set(3, 0, ref.poseMeters.getRotation().getSin());
-      nextR.set(4, 0, vl);
-      nextR.set(5, 0, vr);
+      var nextR =
+          VecBuilder.fill(
+              ref.poseMeters.getTranslation().getX(),
+              ref.poseMeters.getTranslation().getY(),
+              ref.poseMeters.getRotation().getRadians(),
+              vl,
+              vr);
 
-      Matrix<N4, N1> localY = getLocalMeasurementModel(trueXhat, new Matrix<>(Nat.N2(), Nat.N1()));
-      var noiseStdDev = VecBuilder.fill(0.001, 0.001, 0.5, 0.5);
+      Matrix<N3, N1> localY = getLocalMeasurementModel(trueXhat, new Matrix<>(Nat.N2(), Nat.N1()));
+      var noiseStdDev = VecBuilder.fill(0.0001, 0.5, 0.5);
 
       observer.correct(u, localY.plus(StateSpaceUtil.makeWhiteNoiseVector(noiseStdDev)));
 
       var rdot = nextR.minus(r).div(dtSeconds);
       u = new Matrix<>(B.solve(rdot.minus(getDynamics(r, new Matrix<>(Nat.N2(), Nat.N1())))));
 
-      r = nextR;
       observer.predict(u, dtSeconds);
+
+      r = nextR;
       trueXhat =
           NumericalIntegration.rk4(UnscentedKalmanFilterTest::getDynamics, trueXhat, u, dtSeconds);
     }
@@ -183,25 +199,28 @@ class UnscentedKalmanFilterTest {
     observer.correct(u, localY);
 
     var globalY = getGlobalMeasurementModel(trueXhat, u);
-    var R = StateSpaceUtil.makeCostMatrix(VecBuilder.fill(0.01, 0.01, 0.0001, 0.0001, 0.5, 0.5));
+    var R =
+        StateSpaceUtil.makeCovarianceMatrix(
+            Nat.N5(), VecBuilder.fill(0.01, 0.01, 0.0001, 0.5, 0.5));
     observer.correct(
-        Nat.N6(),
+        Nat.N5(),
         u,
         globalY,
         UnscentedKalmanFilterTest::getGlobalMeasurementModel,
         R,
-        (sigmas, weights) -> sigmas.times(Matrix.changeBoundsUnchecked(weights)),
-        Matrix::minus,
-        Matrix::minus,
-        Matrix::plus);
+        AngleStatistics.angleMean(2),
+        AngleStatistics.angleResidual(2),
+        AngleStatistics.angleResidual(2),
+        AngleStatistics.angleAdd(2));
 
     final var finalPosition = trajectory.sample(trajectory.getTotalTimeSeconds());
 
-    assertEquals(finalPosition.poseMeters.getTranslation().getX(), observer.getXhat(0), 0.25);
-    assertEquals(finalPosition.poseMeters.getTranslation().getY(), observer.getXhat(1), 0.25);
-    assertEquals(finalPosition.poseMeters.getRotation().getRadians(), observer.getXhat(2), 1.0);
-    assertEquals(0.0, observer.getXhat(3), 1.0);
-    assertEquals(0.0, observer.getXhat(4), 1.0);
+    assertEquals(finalPosition.poseMeters.getTranslation().getX(), observer.getXhat(0), 0.055);
+    assertEquals(finalPosition.poseMeters.getTranslation().getY(), observer.getXhat(1), 0.15);
+    assertEquals(
+        finalPosition.poseMeters.getRotation().getRadians(), observer.getXhat(2), 0.000005);
+    assertEquals(0.0, observer.getXhat(3), 0.1);
+    assertEquals(0.0, observer.getXhat(4), 0.1);
   }
 
   @Test
@@ -235,95 +254,24 @@ class UnscentedKalmanFilterTest {
     assertEquals(ref.get(0, 0), observer.getXhat(0), 5);
   }
 
+  @SuppressWarnings("LocalVariableName")
   @Test
-  void testUnscentedTransform() {
-    // From FilterPy
-    var ret =
-        UnscentedKalmanFilter.unscentedTransform(
-            Nat.N4(),
-            Nat.N4(),
-            Matrix.mat(Nat.N4(), Nat.N9())
-                .fill(
-                    -0.9,
-                    -0.822540333075852,
-                    -0.8922540333075852,
-                    -0.9,
-                    -0.9,
-                    -0.9774596669241481,
-                    -0.9077459666924148,
-                    -0.9,
-                    -0.9,
-                    1.0,
-                    1.0,
-                    1.077459666924148,
-                    1.0,
-                    1.0,
-                    1.0,
-                    0.9225403330758519,
-                    1.0,
-                    1.0,
-                    -0.9,
-                    -0.9,
-                    -0.9,
-                    -0.822540333075852,
-                    -0.8922540333075852,
-                    -0.9,
-                    -0.9,
-                    -0.9774596669241481,
-                    -0.9077459666924148,
-                    1.0,
-                    1.0,
-                    1.0,
-                    1.0,
-                    1.077459666924148,
-                    1.0,
-                    1.0,
-                    1.0,
-                    0.9225403330758519),
-            VecBuilder.fill(
-                -132.33333333,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667),
-            VecBuilder.fill(
-                -129.34333333,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667,
-                16.66666667),
-            (sigmas, weights) -> sigmas.times(Matrix.changeBoundsUnchecked(weights)),
-            Matrix::minus);
+  void testRoundTripP() {
+    var dtSeconds = 0.005;
 
-    assertTrue(VecBuilder.fill(-0.9, 1, -0.9, 1).isEqual(ret.getFirst(), 1E-5));
+    var observer =
+        new UnscentedKalmanFilter<>(
+            Nat.N2(),
+            Nat.N2(),
+            (x, u) -> x,
+            (x, u) -> x,
+            VecBuilder.fill(0.0, 0.0),
+            VecBuilder.fill(0.0, 0.0),
+            dtSeconds);
 
-    assertTrue(
-        Matrix.mat(Nat.N4(), Nat.N4())
-            .fill(
-                2.02000002e-01,
-                2.00000500e-02,
-                -2.69044710e-29,
-                -4.59511477e-29,
-                2.00000500e-02,
-                2.00001000e-01,
-                -2.98781068e-29,
-                -5.12759588e-29,
-                -2.73372625e-29,
-                -3.09882635e-29,
-                2.02000002e-01,
-                2.00000500e-02,
-                -4.67065917e-29,
-                -5.10705197e-29,
-                2.00000500e-02,
-                2.00001000e-01)
-            .isEqual(ret.getSecond(), 1E-5));
+    var P = Matrix.mat(Nat.N2(), Nat.N2()).fill(2.0, 1.0, 1.0, 2.0);
+    observer.setP(P);
+
+    assertTrue(observer.getP().isEqual(P, 1e-9));
   }
 }

wpimath/src/test/native/cpp/estimator/DifferentialDrivePoseEstimatorTest.cpp

@@ -93,6 +93,6 @@ TEST(DifferentialDrivePoseEstimatorTest, Accuracy) {
   }
 
   EXPECT_NEAR(0.0, errorSum / (trajectory.TotalTime().value() / dt.value()),
-              0.2);
-  EXPECT_NEAR(0.0, maxError, 0.4);
+              0.05);
+  EXPECT_NEAR(0.0, maxError, 0.1);
 }

wpimath/src/test/native/cpp/estimator/MecanumDrivePoseEstimatorTest.cpp

@@ -84,6 +84,6 @@ TEST(MecanumDrivePoseEstimatorTest, Accuracy) {
     t += dt;
   }
 
-  EXPECT_LT(errorSum / (trajectory.TotalTime().value() / dt.value()), 0.2);
-  EXPECT_LT(maxError, 0.4);
+  EXPECT_LT(errorSum / (trajectory.TotalTime().value() / dt.value()), 0.05);
+  EXPECT_LT(maxError, 0.1);
 }

wpimath/src/test/native/cpp/estimator/MerweScaledSigmaPointsTest.cpp

@@ -10,7 +10,7 @@ namespace drake::math {
 namespace {
 TEST(MerweScaledSigmaPointsTest, ZeroMean) {
   frc::MerweScaledSigmaPoints<2> sigmaPoints;
-  auto points = sigmaPoints.SigmaPoints(
+  auto points = sigmaPoints.SquareRootSigmaPoints(
       frc::Vectord<2>{0.0, 0.0}, frc::Matrixd<2, 2>{{1.0, 0.0}, {0.0, 1.0}});
 
   EXPECT_TRUE(
@@ -21,8 +21,9 @@ TEST(MerweScaledSigmaPointsTest, ZeroMean) {
 
 TEST(MerweScaledSigmaPointsTest, NonzeroMean) {
   frc::MerweScaledSigmaPoints<2> sigmaPoints;
-  auto points = sigmaPoints.SigmaPoints(
-      frc::Vectord<2>{1.0, 2.0}, frc::Matrixd<2, 2>{{1.0, 0.0}, {0.0, 10.0}});
+  auto points = sigmaPoints.SquareRootSigmaPoints(
+      frc::Vectord<2>{1.0, 2.0},
+      frc::Matrixd<2, 2>{{1.0, 0.0}, {0.0, std::sqrt(10.0)}});
 
   EXPECT_TRUE(
       (points - frc::Matrixd<2, 5>{{1.0, 1.00173205, 1.0, 0.998268, 1.0},

wpimath/src/test/native/cpp/estimator/SwerveDrivePoseEstimatorTest.cpp

@@ -84,6 +84,6 @@ TEST(SwerveDrivePoseEstimatorTest, Accuracy) {
     t += dt;
   }
 
-  EXPECT_LT(errorSum / (trajectory.TotalTime().value() / dt.value()), 0.2);
-  EXPECT_LT(maxError, 0.4);
+  EXPECT_LT(errorSum / (trajectory.TotalTime().value() / dt.value()), 0.05);
+  EXPECT_LT(maxError, 0.1);
 }

wpimath/src/test/native/cpp/estimator/UnscentedKalmanFilterTest.cpp

@@ -23,7 +23,7 @@ namespace {
 frc::Vectord<5> Dynamics(const frc::Vectord<5>& x, const frc::Vectord<2>& u) {
   auto motors = frc::DCMotor::CIM(2);
 
-  // constexpr double Glow = 15.32;       // Low gear ratio
+  // constexpr double Glow = 15.32;    // Low gear ratio
   constexpr double Ghigh = 7.08;       // High gear ratio
   constexpr auto rb = 0.8382_m / 2.0;  // Robot radius
   constexpr auto r = 0.0746125_m;      // Wheel radius
@@ -71,6 +71,11 @@ TEST(UnscentedKalmanFilterTest, Init) {
                                                LocalMeasurementModel,
                                                {0.5, 0.5, 10.0, 1.0, 1.0},
                                                {0.0001, 0.01, 0.01},
+                                               frc::AngleMean<5, 5>(2),
+                                               frc::AngleMean<3, 5>(0),
+                                               frc::AngleResidual<5>(2),
+                                               frc::AngleResidual<3>(0),
+                                               frc::AngleAdd<5>(2),
                                                dt};
   frc::Vectord<2> u{12.0, 12.0};
   observer.Predict(u, dt);
@@ -93,6 +98,11 @@ TEST(UnscentedKalmanFilterTest, Convergence) {
                                                LocalMeasurementModel,
                                                {0.5, 0.5, 10.0, 1.0, 1.0},
                                                {0.0001, 0.5, 0.5},
+                                               frc::AngleMean<5, 5>(2),
+                                               frc::AngleMean<3, 5>(0),
+                                               frc::AngleResidual<5>(2),
+                                               frc::AngleResidual<3>(0),
+                                               frc::AngleAdd<5>(2),
                                                dt};
 
   auto waypoints =
@@ -150,12 +160,28 @@ TEST(UnscentedKalmanFilterTest, Convergence) {
   );
 
   auto finalPosition = trajectory.Sample(trajectory.TotalTime());
-  ASSERT_NEAR(finalPosition.pose.Translation().X().value(), observer.Xhat(0),
-              1.0);
-  ASSERT_NEAR(finalPosition.pose.Translation().Y().value(), observer.Xhat(1),
-              1.0);
-  ASSERT_NEAR(finalPosition.pose.Rotation().Radians().value(), observer.Xhat(2),
-              1.0);
-  ASSERT_NEAR(0.0, observer.Xhat(3), 1.0);
-  ASSERT_NEAR(0.0, observer.Xhat(4), 1.0);
+  EXPECT_NEAR(finalPosition.pose.Translation().X().value(), observer.Xhat(0),
+              0.055);
+  EXPECT_NEAR(finalPosition.pose.Translation().Y().value(), observer.Xhat(1),
+              0.15);
+  EXPECT_NEAR(finalPosition.pose.Rotation().Radians().value(), observer.Xhat(2),
+              0.000005);
+  EXPECT_NEAR(0.0, observer.Xhat(3), 0.1);
+  EXPECT_NEAR(0.0, observer.Xhat(4), 0.1);
+}
+
+TEST(UnscentedKalmanFilterTest, RoundTripP) {
+  constexpr auto dt = 5_ms;
+
+  frc::UnscentedKalmanFilter<2, 2, 2> observer{
+      [](const frc::Vectord<2>& x, const frc::Vectord<2>& u) { return x; },
+      [](const frc::Vectord<2>& x, const frc::Vectord<2>& u) { return x; },
+      {0.0, 0.0},
+      {0.0, 0.0},
+      dt};
+
+  frc::Matrixd<2, 2> P({{2, 1}, {1, 2}});
+  observer.SetP(P);
+
+  ASSERT_TRUE(observer.P().isApprox(P));
 }