From 95ae23b0e7bdbbe9419d80b581de5854c69d2fa6 Mon Sep 17 00:00:00 2001
From: Tyler Veness <calcmogul@gmail.com>
Date: Sat, 19 Mar 2022 20:41:28 -0700
Subject: [PATCH] [wpimath] Improve EKF numerical stability (#4093)

The Joseph form of the error covariance update equation is more
numerically stable when the Kalman gain isn't optimal. Numerical
instability and filter divergence can occur if the user goes long time
periods between updates and the error covariance becomes ill-conditioned
(the ratio between the largest and smallest eigenvalue gets too large).
---
 .../wpi/first/math/estimator/ExtendedKalmanFilter.java | 10 ++++++++--
 .../include/frc/estimator/ExtendedKalmanFilter.h       | 10 +++++++---
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/wpimath/src/main/java/edu/wpi/first/math/estimator/ExtendedKalmanFilter.java b/wpimath/src/main/java/edu/wpi/first/math/estimator/ExtendedKalmanFilter.java
index dc268f92c2..10d6a548f3 100644
--- a/wpimath/src/main/java/edu/wpi/first/math/estimator/ExtendedKalmanFilter.java
+++ b/wpimath/src/main/java/edu/wpi/first/math/estimator/ExtendedKalmanFilter.java
@@ -366,7 +366,13 @@ public class ExtendedKalmanFilter<States extends Num, Inputs extends Num, Output
     // x̂ₖ₊₁⁺ = x̂ₖ₊₁⁻ + K(y − h(x̂ₖ₊₁⁻, uₖ₊₁))
     m_xHat = addFuncX.apply(m_xHat, K.times(residualFuncY.apply(y, h.apply(m_xHat, u))));
 
-    // Pₖ₊₁⁺ = (I − KC)Pₖ₊₁⁻
-    m_P = Matrix.eye(m_states).minus(K.times(C)).times(m_P);
+    // Pₖ₊₁⁺ = (I−Kₖ₊₁C)Pₖ₊₁⁻(I−Kₖ₊₁C)ᵀ + Kₖ₊₁RKₖ₊₁ᵀ
+    // Use Joseph form for numerical stability
+    m_P =
+        Matrix.eye(m_states)
+            .minus(K.times(C))
+            .times(m_P)
+            .times(Matrix.eye(m_states).minus(K.times(C)).transpose())
+            .plus(K.times(discR).times(K.transpose()));
   }
 }
diff --git a/wpimath/src/main/native/include/frc/estimator/ExtendedKalmanFilter.h b/wpimath/src/main/native/include/frc/estimator/ExtendedKalmanFilter.h
index 3e5edb8d0b..a7e0f3df94 100644
--- a/wpimath/src/main/native/include/frc/estimator/ExtendedKalmanFilter.h
+++ b/wpimath/src/main/native/include/frc/estimator/ExtendedKalmanFilter.h
@@ -327,11 +327,15 @@ class ExtendedKalmanFilter {
     Eigen::Matrix<double, States, Rows> K =
         S.transpose().ldlt().solve(C * m_P.transpose()).transpose();
 
-    // x̂ₖ₊₁⁺ = x̂ₖ₊₁⁻ + K(y − h(x̂ₖ₊₁⁻, uₖ₊₁))
+    // x̂ₖ₊₁⁺ = x̂ₖ₊₁⁻ + Kₖ₊₁(y − h(x̂ₖ₊₁⁻, uₖ₊₁))
     m_xHat = addFuncX(m_xHat, K * residualFuncY(y, h(m_xHat, u)));
 
-    // Pₖ₊₁⁺ = (I − KC)Pₖ₊₁⁻
-    m_P = (Eigen::Matrix<double, States, States>::Identity() - K * C) * m_P;
+    // Pₖ₊₁⁺ = (I−Kₖ₊₁C)Pₖ₊₁⁻(I−Kₖ₊₁C)ᵀ + Kₖ₊₁RKₖ₊₁ᵀ
+    // Use Joseph form for numerical stability
+    m_P = (Eigen::Matrix<double, States, States>::Identity() - K * C) * m_P *
+              (Eigen::Matrix<double, States, States>::Identity() - K * C)
+                  .transpose() +
+          K * discR * K.transpose();
   }
 
  private: