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