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