[wpimath] Fix UnscentedKalmanFilter and improve math docs (#7850)

Throughout the code the state sqrt covariance S and innovation covariance Sy are maintained as upper triangular cholesky factors of those covariance matrices. The original paper defines P=S*S', so S should be lower triangular. The functions in the paper reflect this. In the code implementation, the sqrt covariance matrices are upper triangular, but the algorithm expects them to be lower triangular.

This bug was likely missed because the incorrect version of the filter is able to converge for some systems where all the states are observed, and the test case is set up such that all states are observed.

To fix the bug, a couple things needed to be changed:
all instances of rankUpdate() needed to be changed to use the lower triangular cholesky factor,
In the unscented transform, when S is found via QR decomposition, we need to take the transpose because R is upper triangular,
P() and SetP() functions need to be modified to be P=S*S' instead of P=S'*S, and P.llt().matrixL() instead of P.llt().matrixU() respectively.

Each part of the algorithm has also had the comments changed to clarify exactly which equation from the paper it implements.

Co-authored-by: Tyler Veness <[email protected]>

Author: PascalSkylake

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

@@ -69,8 +69,8 @@ public int getNumSigmas() {
   }
 
   /**
-   * Computes the sigma points for an unscented Kalman filter given the mean (x) and covariance(P)
-   * of the filter.
+   * Computes the sigma points for an unscented Kalman filter given the mean (x) and square-root
+   * covariance (s) of the filter.
    *
    * @param x An array of the means.
    * @param s Square-root covariance of the filter.
@@ -86,6 +86,8 @@ public int getNumSigmas() {
     // 2 * states + 1 by states
     Matrix<S, ?> sigmas =
         new Matrix<>(new SimpleMatrix(m_states.getNum(), 2 * m_states.getNum() + 1));
+
+    // equation (17)
     sigmas.setColumn(0, x);
     for (int k = 0; k < m_states.getNum(); k++) {
       var xPlusU = x.plus(U.extractColumnVector(k));

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

@@ -35,9 +35,9 @@
  * href="https://file.tavsys.net/control/controls-engineering-in-frc.pdf">https://file.tavsys.net/control/controls-engineering-in-frc.pdf</a>
  * chapter 9 "Stochastic control theory".
  *
- * <p>This class implements a square-root-form unscented Kalman filter (SR-UKF). For more
- * information about the SR-UKF, see <a
- * href="https://www.researchgate.net/publication/3908304">https://www.researchgate.net/publication/3908304</a>.
+ * <p>This class implements a square-root-form unscented Kalman filter (SR-UKF). The main reason for
+ * this is to guarantee that the covariance matrix remains positive definite. For more information
+ * about the SR-UKF, see https://www.researchgate.net/publication/3908304.
  *
  * @param <States> Number of states.
  * @param <Inputs> Number of inputs.
@@ -105,7 +105,7 @@ public UnscentedKalmanFilter(
   }
 
   /**
-   * Constructs an unscented Kalman filter with custom mean, residual, and addition functions. Using
+   * Constructs an Unscented Kalman filter with custom mean, residual, and addition functions. Using
    * custom functions for arithmetic can be useful if you have angles in the state or measurements,
    * because they allow you to correctly account for the modular nature of angle arithmetic.
    *
@@ -193,12 +193,21 @@ Pair<Matrix<C, N1>, Matrix<C, C>> squareRootUnscentedTransform(
           "Wc must be 2 * states + 1 by 1! Got " + Wc.getNumRows() + " by " + Wc.getNumCols());
     }
 
-    // New mean is usually just the sum of the sigmas * weight:
-    //       n
-    // dot = Σ W[k] Xᵢ[k]
-    //      k=1
+    // New mean is usually just the sum of the sigmas * weights:
+    //
+    //      2n
+    //   x̂ = Σ Wᵢ⁽ᵐ⁾𝒳ᵢ
+    //      i=0
+    //
+    // equations (19) and (23) in the paper show this,
+    // but we allow a custom function, usually for angle wrapping
     Matrix<C, N1> x = meanFunc.apply(sigmas, Wm);
 
+    // Form an intermediate matrix S⁻ as:
+    //
+    //   [√{W₁⁽ᶜ⁾}(𝒳_{1:2L} - x̂) √{Rᵛ}]
+    //
+    // the part of equations (20) and (24) within the "qr{}"
     Matrix<C, ?> Sbar = new Matrix<>(new SimpleMatrix(dim.getNum(), 2 * s.getNum() + dim.getNum()));
     for (int i = 0; i < 2 * s.getNum(); i++) {
       Sbar.setColumn(
@@ -214,8 +223,24 @@ Pair<Matrix<C, N1>, Matrix<C, C>> squareRootUnscentedTransform(
       throw new RuntimeException("QR decomposition failed! Input matrix:\n" + qrStorage);
     }
 
-    Matrix<C, C> newS = new Matrix<>(new SimpleMatrix(qr.getR(null, true)));
-    newS.rankUpdate(residualFunc.apply(sigmas.extractColumnVector(0), x), Wc.get(0, 0), false);
+    // Compute the square-root covariance of the sigma points
+    //
+    // We transpose S⁻ first because we formed it by horizontally
+    // concatenating each part; it should be vertical so we can take
+    // the QR decomposition as defined in the "QR Decomposition" passage
+    // of section 3. "EFFICIENT SQUARE-ROOT IMPLEMENTATION"
+    //
+    // The resulting matrix R is the square-root covariance S, but it
+    // is upper triangular, so we need to transpose it.
+    //
+    // equations (20) and (24)
+    Matrix<C, C> newS = new Matrix<>(new SimpleMatrix(qr.getR(null, true)).transpose());
+
+    // Update or downdate the square-root covariance with (𝒳₀-x̂)
+    // depending on whether its weight (W₀⁽ᶜ⁾) is positive or negative.
+    //
+    // equations (21) and (25)
+    newS.rankUpdate(residualFunc.apply(sigmas.extractColumnVector(0), x), Wc.get(0, 0), true);
 
     return new Pair<>(x, newS);
   }
@@ -256,7 +281,7 @@ public void setS(Matrix<States, States> newS) {
    */
   @Override
   public Matrix<States, States> getP() {
-    return m_S.transpose().times(m_S);
+    return m_S.times(m_S.transpose());
   }
 
   /**
@@ -280,7 +305,7 @@ public double getP(int row, int col) {
    */
   @Override
   public void setP(Matrix<States, States> newP) {
-    m_S = newP.lltDecompose(false);
+    m_S = newP.lltDecompose(true);
   }
 
   /**
@@ -347,14 +372,28 @@ public void predict(Matrix<Inputs, N1> u, double dtSeconds) {
     var discQ = Discretization.discretizeAQ(contA, m_contQ, dtSeconds).getSecond();
     var squareRootDiscQ = discQ.lltDecompose(true);
 
+    // Generate sigma points around the state mean
+    //
+    // equation (17)
     var sigmas = m_pts.squareRootSigmaPoints(m_xHat, m_S);
 
+    // Project each sigma point forward in time according to the
+    // dynamics f(x, u)
+    //
+    //   sigmas  = 𝒳ₖ₋₁
+    //   sigmasF = 𝒳ₖ,ₖ₋₁ or just 𝒳 for readability
+    //
+    // equation (18)
     for (int i = 0; i < m_pts.getNumSigmas(); ++i) {
       Matrix<States, N1> x = sigmas.extractColumnVector(i);
 
       m_sigmasF.setColumn(i, NumericalIntegration.rk4(m_f, x, u, dtSeconds));
     }
 
+    // Pass the predicted sigmas (𝒳) through the Unscented Transform
+    // to compute the prior state mean and covariance
+    //
+    // equations (18) (19) and (20)
     var ret =
         squareRootUnscentedTransform(
             m_states,
@@ -459,15 +498,27 @@ public <R extends Num> void correct(
     final var discR = Discretization.discretizeR(R, m_dtSeconds);
     final var squareRootDiscR = discR.lltDecompose(true);
 
-    // Transform sigma points into measurement space
+    // Generate new sigma points from the prior mean and covariance
+    // and transform them into measurement space using h(x, u)
+    //
+    //   sigmas  = 𝒳
+    //   sigmasH = 𝒴
+    //
+    // This differs from equation (22) which uses
+    // the prior sigma points, regenerating them allows
+    // multiple measurement updates per time update
     Matrix<R, ?> sigmasH = new Matrix<>(new SimpleMatrix(rows.getNum(), 2 * m_states.getNum() + 1));
     var sigmas = m_pts.squareRootSigmaPoints(m_xHat, m_S);
     for (int i = 0; i < m_pts.getNumSigmas(); i++) {
       Matrix<R, N1> hRet = h.apply(sigmas.extractColumnVector(i), u);
       sigmasH.setColumn(i, hRet);
     }
 
-    // Mean and covariance of prediction passed through unscented transform
+    // Pass the predicted measurement sigmas through the Unscented Transform
+    // to compute the mean predicted measurement and square-root innovation
+    // covariance.
+    //
+    // equations (23) (24) and (25)
     var transRet =
         squareRootUnscentedTransform(
             m_states,
@@ -481,30 +532,54 @@ public <R extends Num> void correct(
     var yHat = transRet.getFirst();
     var Sy = transRet.getSecond();
 
-    // Compute cross covariance of the state and the measurements
+    // Compute cross covariance of the predicted state and measurement sigma
+    // points given as:
+    //
+    //           2n
+    //   P_{xy} = Σ Wᵢ⁽ᶜ⁾[𝒳ᵢ - x̂][𝒴ᵢ - ŷ⁻]ᵀ
+    //           i=0
+    //
+    // equation (26)
     Matrix<States, R> Pxy = new Matrix<>(m_states, rows);
     for (int i = 0; i < m_pts.getNumSigmas(); i++) {
-      // Pxy += (sigmas_f[:, i] - x̂)(sigmas_h[:, i] - ŷ)ᵀ W_c[i]
       var dx = residualFuncX.apply(m_sigmasF.extractColumnVector(i), m_xHat);
       var dy = residualFuncY.apply(sigmasH.extractColumnVector(i), yHat).transpose();
 
       Pxy = Pxy.plus(dx.times(dy).times(m_pts.getWc(i)));
     }
 
-    // K = (P_{xy} / S_yᵀ) / S_y
-    // K = (S_y \ P_{xy}ᵀ)ᵀ / S_y
-    // K = (S_yᵀ \ (S_y \ P_{xy}ᵀ))ᵀ
+    // Compute the Kalman gain. We use Eigen's QR decomposition to solve. This
+    // is equivalent to MATLAB's \ operator, so we need to rearrange to use
+    // that.
+    //
+    //   K = (P_{xy} / S_{y}ᵀ) / S_{y}
+    //   K = (S_{y} \ P_{xy})ᵀ / S_{y}
+    //   K = (S_{y}ᵀ \ (S_{y} \ P_{xy}ᵀ))ᵀ
+    //
+    // equation (27)
     Matrix<States, R> K =
         Sy.transpose()
             .solveFullPivHouseholderQr(Sy.solveFullPivHouseholderQr(Pxy.transpose()))
             .transpose();
 
-    // x̂ₖ₊₁⁺ = x̂ₖ₊₁⁻ + K(y − ŷ)
+    // Compute the posterior state mean
+    //
+    // x̂ = x̂⁻ + K(y − ŷ⁻)
+    //
+    // second part of equation (27)
     m_xHat = addFuncX.apply(m_xHat, K.times(residualFuncY.apply(y, yHat)));
 
+    // Compute the intermediate matrix U for downdating
+    // the square-root covariance
+    //
+    // equation (28)
     Matrix<States, R> U = K.times(Sy);
+
+    // Downdate the posterior square-root state covariance
+    //
+    // equation (29)
     for (int i = 0; i < rows.getNum(); i++) {
-      m_S.rankUpdate(U.extractColumnVector(i), -1, false);
+      m_S.rankUpdate(U.extractColumnVector(i), -1, true);
     }
   }
 }

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

@@ -51,7 +51,7 @@ class MerweScaledSigmaPoints {
 
   /**
    * Computes the sigma points for an unscented Kalman filter given the mean
-   * (x) and square-root covariance(S) of the filter.
+   * (x) and square-root covariance (S) of the filter.
    *
    * @param x An array of the means.
    * @param S Square-root covariance of the filter.
@@ -68,6 +68,8 @@ class MerweScaledSigmaPoints {
     Matrixd<States, States> U = eta * S;
 
     Matrixd<States, 2 * States + 1> sigmas;
+
+    // equation (17)
     sigmas.template block<States, 1>(0, 0) = x;
     for (int k = 0; k < States; ++k) {
       sigmas.template block<States, 1>(0, k + 1) =