EquivariantFilter.h

Go to the documentation of this file.

 
#pragma once
 
#include <gtsam/base/GroupAction.h>
#include <gtsam/base/Matrix.h>
#include <gtsam/base/Vector.h>
#include <gtsam/navigation/ManifoldEKF.h>
#include <gtsam/nonlinear/Values.h>
 
namespace gtsam {

template <typename M, typename Symmetry>
class EquivariantFilter : public ManifoldEKF<M> {
 public:
  using Base = ManifoldEKF<M>;
 
  // Manifold traits
  static constexpr int DimM = Base::Dim;
  using TangentM = typename Base::TangentVector;
  using MatrixM = typename Base::Jacobian;
  using CovarianceM = typename Base::Covariance;
 
  // Group traits
  using G = typename Symmetry::Group;
  static constexpr int DimG = traits<G>::dimension;
  using TangentG = typename traits<G>::TangentVector;
 
  // Cross-dimension helpers
  using MatrixMG = Eigen::Matrix<double, DimM, DimG>;
  using MatrixGM = Eigen::Matrix<double, DimG, DimM>;
 
 private:
  M xi_ref_;  // Origin (reference) state on the manifold
  typename Symmetry::Orbit act_on_ref_;  // Orbit of the reference state
  MatrixMG Dphi0_;           // Differential of state action at identity
  MatrixGM InnovationLift_;  // Innovation lift matrix ((Dphi0)^+)
 
  G g_;  // Group element estimate
 
 protected:


  void resetReferenceAndGroup(const M& xi_ref, const CovarianceM& P,
                              const G& g) {
    xi_ref_ = xi_ref;
    g_ = g;
 
    // Recompute Dphi0 and innovation lift matrix from current reference state
    act_on_ref_ = typename Symmetry::Orbit(xi_ref_);
    const G identity_at_g = traits<G>::Compose(g_.inverse(), g_);
    act_on_ref_(identity_at_g, &Dphi0_);
    InnovationLift_ = Dphi0_.completeOrthogonalDecomposition().pseudoInverse();
 
    if constexpr (DimM == Eigen::Dynamic) {
      this->n_ = traits<M>::GetDimension(xi_ref_);
      this->I_ = MatrixM::Identity(this->n_, this->n_);
    }
    this->P_ = P;
    this->X_ = act_on_ref_(g_);
  }

  const M& referenceState() const { return xi_ref_; }

  template <typename Lift, typename InputOrbit>
  TangentG liftAtOrigin(const InputOrbit& psi_u,
                        OptionalJacobian<DimG, DimM> D_lift = {}) const {
    Lift lift_u_origin(psi_u(g_.inverse()));
    return lift_u_origin(xi_ref_, D_lift);
  }

  void propagate(const TangentG& increment, const MatrixM& Phi,
                 const CovarianceM& Qd) {
    const G step = traits<G>::Expmap(increment);
    const G g_next = Symmetry::type == ActionType::Left
                         ? traits<G>::Compose(step, g_)
                         : traits<G>::Compose(g_, step);
    Base::predict(act_on_ref_(g_next), Phi, Qd);
    g_ = g_next;
  }

  void applyCorrection(const TangentG& delta_x) {
    const G step = traits<G>::Expmap(delta_x);
    g_ = Symmetry::type == ActionType::Left
             ? traits<G>::Compose(g_, step)
             : traits<G>::Compose(step, g_);
    this->X_ = act_on_ref_(g_);
  }
 
 public:
  EquivariantFilter(const M& xi_ref, const CovarianceM& Sigma,
                    const G& X0 = traits<G>::Identity())
      : Base(xi_ref, Sigma), xi_ref_(xi_ref), act_on_ref_(xi_ref), g_(X0) {
    // Compute differential of action phi at identity (Dphi0)
    const G identity_at_g = traits<G>::Compose(g_.inverse(), g_);
    act_on_ref_(identity_at_g, &Dphi0_);
 
    // Precompute the Innovation Lift matrix (pseudo-inverse of Dphi0)
    InnovationLift_ = Dphi0_.completeOrthogonalDecomposition().pseudoInverse();
    this->X_ = act_on_ref_(g_);
  }

  using Base::state;

  const typename Base::Covariance& errorCovariance() const { return this->P_; }

  MatrixM actionDifferential() const {
    MatrixM J;
    if constexpr (MatrixM::RowsAtCompileTime == Eigen::Dynamic) {
      J.resize(this->n_, this->n_);
    }
    const typename Symmetry::Diffeomorphism action_at_g(g_);
    action_at_g(xi_ref_, &J);
    return J;
  }

  CovarianceM covariance() const {
    const MatrixM J = actionDifferential();
    return J * this->P_ * J.transpose();
  }

  const G& groupEstimate() const { return g_; }

  template <typename Lift, typename InputOrbit>
  MatrixM computeErrorDynamicsMatrix(const InputOrbit& psi_u) const {
    MatrixGM D_lift;
    liftAtOrigin<Lift>(psi_u, &D_lift);
    return Dphi0_ * D_lift;
  }

  template <size_t K = 1>
  MatrixM transitionMatrix(const MatrixM& A, double dt) const {
    if constexpr (K == 1) {
      return this->I_ + A * dt;
    } else {
      return MatrixM(expm(A * dt, K));
    }
  }

  template <size_t K = 1, typename Lift, typename InputOrbit>
  void predict(const Lift& lift_u, const InputOrbit& psi_u, const MatrixM& Qc,
               double dt) {
    const MatrixM A = computeErrorDynamicsMatrix<Lift>(psi_u);
    predictWithJacobian<K>(lift_u, A, Qc, dt);
  }

  template <size_t K = 1, typename Lift>
  void predictWithJacobian(const Lift& lift_u, const MatrixM& A,
                           const MatrixM& Qc, double dt) {
    const TangentG Lambda = lift_u(this->state());
    propagate(Lambda * dt, transitionMatrix<K>(A, dt), CovarianceM(Qc * dt));
  }

  template <typename Lift>
  void predictWithTransition(const Lift& lift_u, const MatrixM& Phi,
                             const CovarianceM& Qd, double dt) {
    propagate(lift_u(this->state()) * dt, Phi, Qd);
  }

  template <typename Measurement>
  void update(
      const Measurement& prediction,
      const Eigen::Matrix<double, traits<Measurement>::dimension, DimM>& H,
      const Measurement& z,
      const Eigen::Matrix<double, traits<Measurement>::dimension,
                          traits<Measurement>::dimension>& R) {
    static constexpr int MeasDim = traits<Measurement>::dimension;
 
    // Innovation: y = h(x_pred) - z. In tangent space: local(z, h(x_pred))
    // NOTE: we use the `z_hat - z` sign convention, NOT `z - z_hat`.
    typename traits<Measurement>::TangentVector innovation =
        traits<Measurement>::Local(z, prediction);
 
    // Kalman Gain: K = P H^T S^-1
    // K will be Eigen::Matrix<double, Dim, MeasDim>
    Eigen::Matrix<double, DimM, MeasDim> K = this->KalmanGain(H, R);
 
    // Correction in Manifold tangent space
    // K matches dimensions with innovation, so result is TangentM
    TangentM delta_xi = -K * innovation;
 
    // Lift correction to Group tangent space
    TangentG delta_x = InnovationLift_ * delta_xi;
    applyCorrection(delta_x);
 
    // Update covariance on Manifold using Joseph form
    this->JosephUpdate(K, H, R);
  }

  template <typename Z, typename Func>
  void update(Func&& h, const Z& z,
              const Eigen::Matrix<double, traits<Z>::dimension,
                                  traits<Z>::dimension>& R) {
    static_assert(IsManifold<Z>::value,
                  "Template parameter Z must be a GTSAM Manifold.");
 
    Matrix H(traits<Z>::GetDimension(z), this->n_);
    Z prediction = h(this->X_, H);
    update<Z>(prediction, H, z, R);
  }

  void updateWithVector(const gtsam::Vector& prediction, const Matrix& H,
                        const gtsam::Vector& z, const Matrix& R) {
    this->validateInputs(prediction, H, z, R);
    update<Vector>(prediction, H, z, R);
  }

  template <typename InnovationLiftFn>
  void updateWithVector(const gtsam::Vector& prediction, const Matrix& H,
                        const gtsam::Vector& z, const Matrix& R,
                        InnovationLiftFn&& innovationLift) {
    this->validateInputs(prediction, H, z, R);
 
    const gtsam::Vector innovation = traits<Vector>::Local(z, prediction);
    const Matrix K = this->KalmanGain(H, R);
    const TangentM delta_xi = -K * innovation;
    const TangentG delta_x = innovationLift(delta_xi);
 
    applyCorrection(delta_x);
    this->JosephUpdate(K, H, R);
  }
};
 
}  // namespace gtsam