linear_quadratic_regulator.h

#pragma once
 
#include "core/utils/math/eigen_interface.h"
#include "../vendor/eigen/unsupported/Eigen/MatrixFunctions"
 
#include "core/utils/math/systems/dare_solver.h"
#include "core/utils/math/systems/discretization.h"
#include "core/utils/math/systems/linear_system.h"

template <int DIM> EMat<DIM, DIM> cost_matrix(const EVec<DIM> &tolerances) {
    EMat<DIM, DIM> Q = EMat<DIM, DIM>::Zero();
    for (int i = 0; i < DIM; i++) {
        Q(i, i) = 1.0 / (tolerances(i) * tolerances(i));
    }
    return Q;
}

template <int STATES, int INPUTS> class LinearQuadraticRegulator {
  public:
    // Definitions to shorten some lines.
    using MatrixA = EMat<STATES, STATES>;
    using MatrixB = EMat<STATES, INPUTS>;
    using VectorX = EVec<STATES>;
    using VectorU = EVec<INPUTS>;

    template <int OUTPUTS>
    LinearQuadraticRegulator(
      LinearSystem<STATES, INPUTS, OUTPUTS> &plant, const VectorX &Qtolerances, const VectorU &Rtolerances,
      const double &dt
    )
        : LinearQuadraticRegulator(plant.A(), plant.B(), Qtolerances, Rtolerances, dt) {}

    LinearQuadraticRegulator(
      const MatrixA &A, const MatrixB &B, const VectorX &Qtolerances, const VectorU &Rtolerances, const double &dt
    )
        : LinearQuadraticRegulator(A, B, cost_matrix(Qtolerances), cost_matrix(Rtolerances), dt) {}

    LinearQuadraticRegulator(
      const MatrixA &A, const MatrixB &B, const EMat<STATES, STATES> &Q, const EMat<INPUTS, INPUTS> &R, const double &dt
    ) {
        auto [Ad, Bd] = discretize_AB(A, B, dt);
 
        MatrixA S = DARE<STATES, INPUTS>(Ad, Bd, Q, R);
 
        // (BᵀSD + R) \ (BᵀSA)
        K_ = (Bd.transpose() * S * Bd + R).llt().solve(Bd.transpose() * S * Ad);
    }

    VectorU calculate(const VectorX &x, const VectorX &r) { return K_ * (r - x); }

    template <int OUTPUTS>
    void latency_compensate(LinearSystem<STATES, INPUTS, OUTPUTS> &plant, const double &dt, const double &input_delay) {
        auto [Ad, Bd] = discretize_AB(plant.A(), plant.B(), dt);
 
        // Kdelay = K(A - BK)^(delay / dt)
        K_ = K_ * (Ad - Bd * K_).pow(input_delay / dt);
    }
 
  private:
    EMat<INPUTS, STATES> K_;
};