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