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