001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.math.controller;
006
007import edu.wpi.first.math.DARE;
008import edu.wpi.first.math.MathSharedStore;
009import edu.wpi.first.math.Matrix;
010import edu.wpi.first.math.Num;
011import edu.wpi.first.math.StateSpaceUtil;
012import edu.wpi.first.math.Vector;
013import edu.wpi.first.math.numbers.N1;
014import edu.wpi.first.math.system.Discretization;
015import edu.wpi.first.math.system.LinearSystem;
016import org.ejml.simple.SimpleMatrix;
017
018/**
019 * Contains the controller coefficients and logic for a linear-quadratic regulator (LQR). LQRs use
020 * the control law u = K(r - x).
021 *
022 * <p>For more on the underlying math, read <a
023 * href="https://file.tavsys.net/control/controls-engineering-in-frc.pdf">https://file.tavsys.net/control/controls-engineering-in-frc.pdf</a>.
024 */
025public class LinearQuadraticRegulator<States extends Num, Inputs extends Num, Outputs extends Num> {
026  /** The current reference state. */
027  private Matrix<States, N1> m_r;
028
029  /** The computed and capped controller output. */
030  private Matrix<Inputs, N1> m_u;
031
032  // Controller gain.
033  private Matrix<Inputs, States> m_K;
034
035  /**
036   * Constructs a controller with the given coefficients and plant. Rho is defaulted to 1.
037   *
038   * <p>See <a
039   * href="https://docs.wpilib.org/en/stable/docs/software/advanced-controls/state-space/state-space-intro.html#lqr-tuning">https://docs.wpilib.org/en/stable/docs/software/advanced-controls/state-space/state-space-intro.html#lqr-tuning</a>
040   * for how to select the tolerances.
041   *
042   * @param plant The plant being controlled.
043   * @param qelms The maximum desired error tolerance for each state.
044   * @param relms The maximum desired control effort for each input.
045   * @param dtSeconds Discretization timestep.
046   * @throws IllegalArgumentException If the system is uncontrollable.
047   */
048  public LinearQuadraticRegulator(
049      LinearSystem<States, Inputs, Outputs> plant,
050      Vector<States> qelms,
051      Vector<Inputs> relms,
052      double dtSeconds) {
053    this(
054        plant.getA(),
055        plant.getB(),
056        StateSpaceUtil.makeCostMatrix(qelms),
057        StateSpaceUtil.makeCostMatrix(relms),
058        dtSeconds);
059  }
060
061  /**
062   * Constructs a controller with the given coefficients and plant.
063   *
064   * <p>See <a
065   * href="https://docs.wpilib.org/en/stable/docs/software/advanced-controls/state-space/state-space-intro.html#lqr-tuning">https://docs.wpilib.org/en/stable/docs/software/advanced-controls/state-space/state-space-intro.html#lqr-tuning</a>
066   * for how to select the tolerances.
067   *
068   * @param A Continuous system matrix of the plant being controlled.
069   * @param B Continuous input matrix of the plant being controlled.
070   * @param qelms The maximum desired error tolerance for each state.
071   * @param relms The maximum desired control effort for each input.
072   * @param dtSeconds Discretization timestep.
073   * @throws IllegalArgumentException If the system is uncontrollable.
074   */
075  public LinearQuadraticRegulator(
076      Matrix<States, States> A,
077      Matrix<States, Inputs> B,
078      Vector<States> qelms,
079      Vector<Inputs> relms,
080      double dtSeconds) {
081    this(
082        A,
083        B,
084        StateSpaceUtil.makeCostMatrix(qelms),
085        StateSpaceUtil.makeCostMatrix(relms),
086        dtSeconds);
087  }
088
089  /**
090   * Constructs a controller with the given coefficients and plant.
091   *
092   * @param A Continuous system matrix of the plant being controlled.
093   * @param B Continuous input matrix of the plant being controlled.
094   * @param Q The state cost matrix.
095   * @param R The input cost matrix.
096   * @param dtSeconds Discretization timestep.
097   * @throws IllegalArgumentException If the system is uncontrollable.
098   */
099  public LinearQuadraticRegulator(
100      Matrix<States, States> A,
101      Matrix<States, Inputs> B,
102      Matrix<States, States> Q,
103      Matrix<Inputs, Inputs> R,
104      double dtSeconds) {
105    var discABPair = Discretization.discretizeAB(A, B, dtSeconds);
106    var discA = discABPair.getFirst();
107    var discB = discABPair.getSecond();
108
109    if (!StateSpaceUtil.isStabilizable(discA, discB)) {
110      var msg =
111          "The system passed to the LQR is uncontrollable!\n\nA =\n"
112              + discA.getStorage().toString()
113              + "\nB =\n"
114              + discB.getStorage().toString()
115              + '\n';
116      MathSharedStore.reportError(msg, Thread.currentThread().getStackTrace());
117      throw new IllegalArgumentException(msg);
118    }
119
120    var S = DARE.dare(discA, discB, Q, R);
121
122    // K = (BᵀSB + R)⁻¹BᵀSA
123    m_K =
124        discB
125            .transpose()
126            .times(S)
127            .times(discB)
128            .plus(R)
129            .solve(discB.transpose().times(S).times(discA));
130
131    m_r = new Matrix<>(new SimpleMatrix(B.getNumRows(), 1));
132    m_u = new Matrix<>(new SimpleMatrix(B.getNumCols(), 1));
133
134    reset();
135  }
136
137  /**
138   * Constructs a controller with the given coefficients and plant.
139   *
140   * @param A Continuous system matrix of the plant being controlled.
141   * @param B Continuous input matrix of the plant being controlled.
142   * @param Q The state cost matrix.
143   * @param R The input cost matrix.
144   * @param N The state-input cross-term cost matrix.
145   * @param dtSeconds Discretization timestep.
146   * @throws IllegalArgumentException If the system is uncontrollable.
147   */
148  public LinearQuadraticRegulator(
149      Matrix<States, States> A,
150      Matrix<States, Inputs> B,
151      Matrix<States, States> Q,
152      Matrix<Inputs, Inputs> R,
153      Matrix<States, Inputs> N,
154      double dtSeconds) {
155    var discABPair = Discretization.discretizeAB(A, B, dtSeconds);
156    var discA = discABPair.getFirst();
157    var discB = discABPair.getSecond();
158
159    var S = DARE.dare(discA, discB, Q, R, N);
160
161    // K = (BᵀSB + R)⁻¹(BᵀSA + Nᵀ)
162    m_K =
163        discB
164            .transpose()
165            .times(S)
166            .times(discB)
167            .plus(R)
168            .solve(discB.transpose().times(S).times(discA).plus(N.transpose()));
169
170    m_r = new Matrix<>(new SimpleMatrix(B.getNumRows(), 1));
171    m_u = new Matrix<>(new SimpleMatrix(B.getNumCols(), 1));
172
173    reset();
174  }
175
176  /**
177   * Returns the control input vector u.
178   *
179   * @return The control input.
180   */
181  public Matrix<Inputs, N1> getU() {
182    return m_u;
183  }
184
185  /**
186   * Returns an element of the control input vector u.
187   *
188   * @param row Row of u.
189   * @return The row of the control input vector.
190   */
191  public double getU(int row) {
192    return m_u.get(row, 0);
193  }
194
195  /**
196   * Returns the reference vector r.
197   *
198   * @return The reference vector.
199   */
200  public Matrix<States, N1> getR() {
201    return m_r;
202  }
203
204  /**
205   * Returns an element of the reference vector r.
206   *
207   * @param row Row of r.
208   * @return The row of the reference vector.
209   */
210  public double getR(int row) {
211    return m_r.get(row, 0);
212  }
213
214  /**
215   * Returns the controller matrix K.
216   *
217   * @return the controller matrix K.
218   */
219  public Matrix<Inputs, States> getK() {
220    return m_K;
221  }
222
223  /** Resets the controller. */
224  public final void reset() {
225    m_r.fill(0.0);
226    m_u.fill(0.0);
227  }
228
229  /**
230   * Returns the next output of the controller.
231   *
232   * @param x The current state x.
233   * @return The next controller output.
234   */
235  public Matrix<Inputs, N1> calculate(Matrix<States, N1> x) {
236    m_u = m_K.times(m_r.minus(x));
237    return m_u;
238  }
239
240  /**
241   * Returns the next output of the controller.
242   *
243   * @param x The current state x.
244   * @param nextR the next reference vector r.
245   * @return The next controller output.
246   */
247  public Matrix<Inputs, N1> calculate(Matrix<States, N1> x, Matrix<States, N1> nextR) {
248    m_r = nextR;
249    return calculate(x);
250  }
251
252  /**
253   * Adjusts LQR controller gain to compensate for a pure time delay in the input.
254   *
255   * <p>Linear-Quadratic regulator controller gains tend to be aggressive. If sensor measurements
256   * are time-delayed too long, the LQR may be unstable. However, if we know the amount of delay, we
257   * can compute the control based on where the system will be after the time delay.
258   *
259   * <p>See <a
260   * href="https://file.tavsys.net/control/controls-engineering-in-frc.pdf">https://file.tavsys.net/control/controls-engineering-in-frc.pdf</a>
261   * appendix C.4 for a derivation.
262   *
263   * @param plant The plant being controlled.
264   * @param dtSeconds Discretization timestep in seconds.
265   * @param inputDelaySeconds Input time delay in seconds.
266   */
267  public void latencyCompensate(
268      LinearSystem<States, Inputs, Outputs> plant, double dtSeconds, double inputDelaySeconds) {
269    var discABPair = Discretization.discretizeAB(plant.getA(), plant.getB(), dtSeconds);
270    var discA = discABPair.getFirst();
271    var discB = discABPair.getSecond();
272
273    m_K = m_K.times((discA.minus(discB.times(m_K))).pow(inputDelaySeconds / dtSeconds));
274  }
275}