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