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.system;
006
007import edu.wpi.first.math.Matrix;
008import edu.wpi.first.math.Num;
009import edu.wpi.first.math.StateSpaceUtil;
010import edu.wpi.first.math.controller.LinearPlantInversionFeedforward;
011import edu.wpi.first.math.controller.LinearQuadraticRegulator;
012import edu.wpi.first.math.estimator.KalmanFilter;
013import edu.wpi.first.math.numbers.N1;
014import java.util.function.Function;
015import org.ejml.MatrixDimensionException;
016import org.ejml.simple.SimpleMatrix;
017
018/**
019 * Combines a controller, feedforward, and observer for controlling a mechanism with full state
020 * feedback.
021 *
022 * <p>For everything in this file, "inputs" and "outputs" are defined from the perspective of the
023 * plant. This means U is an input and Y is an output (because you give the plant U (powers) and it
024 * gives you back a Y (sensor values)). This is the opposite of what they mean from the perspective
025 * of the controller (U is an output because that's what goes to the motors and Y is an input
026 * because that's what comes back from the sensors).
027 *
028 * <p>For more on the underlying math, read
029 * https://file.tavsys.net/control/controls-engineering-in-frc.pdf.
030 */
031public class LinearSystemLoop<States extends Num, Inputs extends Num, Outputs extends Num> {
032  private final LinearQuadraticRegulator<States, Inputs, Outputs> m_controller;
033  private final LinearPlantInversionFeedforward<States, Inputs, Outputs> m_feedforward;
034  private final KalmanFilter<States, Inputs, Outputs> m_observer;
035  private Matrix<States, N1> m_nextR;
036  private Function<Matrix<Inputs, N1>, Matrix<Inputs, N1>> m_clampFunction;
037
038  /**
039   * Constructs a state-space loop with the given plant, controller, and observer. By default, the
040   * initial reference is all zeros. Users should call reset with the initial system state before
041   * enabling the loop. This constructor assumes that the input(s) to this system are voltage.
042   *
043   * @param plant State-space plant.
044   * @param controller State-space controller.
045   * @param observer State-space observer.
046   * @param maxVoltageVolts The maximum voltage that can be applied. Commonly 12.
047   * @param dtSeconds The nominal timestep.
048   */
049  public LinearSystemLoop(
050      LinearSystem<States, Inputs, Outputs> plant,
051      LinearQuadraticRegulator<States, Inputs, Outputs> controller,
052      KalmanFilter<States, Inputs, Outputs> observer,
053      double maxVoltageVolts,
054      double dtSeconds) {
055    this(
056        controller,
057        new LinearPlantInversionFeedforward<>(plant, dtSeconds),
058        observer,
059        u -> StateSpaceUtil.desaturateInputVector(u, maxVoltageVolts));
060  }
061
062  /**
063   * Constructs a state-space loop with the given plant, controller, and observer. By default, the
064   * initial reference is all zeros. Users should call reset with the initial system state before
065   * enabling the loop.
066   *
067   * @param plant State-space plant.
068   * @param controller State-space controller.
069   * @param observer State-space observer.
070   * @param clampFunction The function used to clamp the input U.
071   * @param dtSeconds The nominal timestep.
072   */
073  public LinearSystemLoop(
074      LinearSystem<States, Inputs, Outputs> plant,
075      LinearQuadraticRegulator<States, Inputs, Outputs> controller,
076      KalmanFilter<States, Inputs, Outputs> observer,
077      Function<Matrix<Inputs, N1>, Matrix<Inputs, N1>> clampFunction,
078      double dtSeconds) {
079    this(
080        controller,
081        new LinearPlantInversionFeedforward<>(plant, dtSeconds),
082        observer,
083        clampFunction);
084  }
085
086  /**
087   * Constructs a state-space loop with the given controller, feedforward and observer. By default,
088   * the initial reference is all zeros. Users should call reset with the initial system state
089   * before enabling the loop.
090   *
091   * @param controller State-space controller.
092   * @param feedforward Plant inversion feedforward.
093   * @param observer State-space observer.
094   * @param maxVoltageVolts The maximum voltage that can be applied. Assumes that the inputs are
095   *     voltages.
096   */
097  public LinearSystemLoop(
098      LinearQuadraticRegulator<States, Inputs, Outputs> controller,
099      LinearPlantInversionFeedforward<States, Inputs, Outputs> feedforward,
100      KalmanFilter<States, Inputs, Outputs> observer,
101      double maxVoltageVolts) {
102    this(
103        controller,
104        feedforward,
105        observer,
106        u -> StateSpaceUtil.desaturateInputVector(u, maxVoltageVolts));
107  }
108
109  /**
110   * Constructs a state-space loop with the given controller, feedforward, and observer. By default,
111   * the initial reference is all zeros. Users should call reset with the initial system state
112   * before enabling the loop.
113   *
114   * @param controller State-space controller.
115   * @param feedforward Plant inversion feedforward.
116   * @param observer State-space observer.
117   * @param clampFunction The function used to clamp the input U.
118   */
119  public LinearSystemLoop(
120      LinearQuadraticRegulator<States, Inputs, Outputs> controller,
121      LinearPlantInversionFeedforward<States, Inputs, Outputs> feedforward,
122      KalmanFilter<States, Inputs, Outputs> observer,
123      Function<Matrix<Inputs, N1>, Matrix<Inputs, N1>> clampFunction) {
124    this.m_controller = controller;
125    this.m_feedforward = feedforward;
126    this.m_observer = observer;
127    this.m_clampFunction = clampFunction;
128
129    m_nextR = new Matrix<>(new SimpleMatrix(controller.getK().getNumCols(), 1));
130    reset(m_nextR);
131  }
132
133  /**
134   * Returns the observer's state estimate x-hat.
135   *
136   * @return the observer's state estimate x-hat.
137   */
138  public Matrix<States, N1> getXHat() {
139    return getObserver().getXhat();
140  }
141
142  /**
143   * Returns an element of the observer's state estimate x-hat.
144   *
145   * @param row Row of x-hat.
146   * @return the i-th element of the observer's state estimate x-hat.
147   */
148  public double getXHat(int row) {
149    return getObserver().getXhat(row);
150  }
151
152  /**
153   * Set the initial state estimate x-hat.
154   *
155   * @param xhat The initial state estimate x-hat.
156   */
157  public void setXHat(Matrix<States, N1> xhat) {
158    getObserver().setXhat(xhat);
159  }
160
161  /**
162   * Set an element of the initial state estimate x-hat.
163   *
164   * @param row Row of x-hat.
165   * @param value Value for element of x-hat.
166   */
167  public void setXHat(int row, double value) {
168    getObserver().setXhat(row, value);
169  }
170
171  /**
172   * Returns an element of the controller's next reference r.
173   *
174   * @param row Row of r.
175   * @return the element i of the controller's next reference r.
176   */
177  public double getNextR(int row) {
178    return getNextR().get(row, 0);
179  }
180
181  /**
182   * Returns the controller's next reference r.
183   *
184   * @return the controller's next reference r.
185   */
186  public Matrix<States, N1> getNextR() {
187    return m_nextR;
188  }
189
190  /**
191   * Set the next reference r.
192   *
193   * @param nextR Next reference.
194   */
195  public void setNextR(Matrix<States, N1> nextR) {
196    m_nextR = nextR;
197  }
198
199  /**
200   * Set the next reference r.
201   *
202   * @param nextR Next reference.
203   */
204  public void setNextR(double... nextR) {
205    if (nextR.length != m_nextR.getNumRows()) {
206      throw new MatrixDimensionException(
207          String.format(
208              "The next reference does not have the "
209                  + "correct number of entries! Expected %s, but got %s.",
210              m_nextR.getNumRows(), nextR.length));
211    }
212    m_nextR = new Matrix<>(new SimpleMatrix(m_nextR.getNumRows(), 1, true, nextR));
213  }
214
215  /**
216   * Returns the controller's calculated control input u plus the calculated feedforward u_ff.
217   *
218   * @return the calculated control input u.
219   */
220  public Matrix<Inputs, N1> getU() {
221    return clampInput(m_controller.getU().plus(m_feedforward.getUff()));
222  }
223
224  /**
225   * Returns an element of the controller's calculated control input u.
226   *
227   * @param row Row of u.
228   * @return the calculated control input u at the row i.
229   */
230  public double getU(int row) {
231    return getU().get(row, 0);
232  }
233
234  /**
235   * Return the controller used internally.
236   *
237   * @return the controller used internally.
238   */
239  public LinearQuadraticRegulator<States, Inputs, Outputs> getController() {
240    return m_controller;
241  }
242
243  /**
244   * Return the feedforward used internally.
245   *
246   * @return the feedforward used internally.
247   */
248  public LinearPlantInversionFeedforward<States, Inputs, Outputs> getFeedforward() {
249    return m_feedforward;
250  }
251
252  /**
253   * Return the observer used internally.
254   *
255   * @return the observer used internally.
256   */
257  public KalmanFilter<States, Inputs, Outputs> getObserver() {
258    return m_observer;
259  }
260
261  /**
262   * Zeroes reference r and controller output u. The previous reference of the
263   * PlantInversionFeedforward and the initial state estimate of the KalmanFilter are set to the
264   * initial state provided.
265   *
266   * @param initialState The initial state.
267   */
268  public void reset(Matrix<States, N1> initialState) {
269    m_nextR.fill(0.0);
270    m_controller.reset();
271    m_feedforward.reset(initialState);
272    m_observer.setXhat(initialState);
273  }
274
275  /**
276   * Returns difference between reference r and current state x-hat.
277   *
278   * @return The state error matrix.
279   */
280  public Matrix<States, N1> getError() {
281    return getController().getR().minus(m_observer.getXhat());
282  }
283
284  /**
285   * Returns difference between reference r and current state x-hat.
286   *
287   * @param index The index of the error matrix to return.
288   * @return The error at that index.
289   */
290  public double getError(int index) {
291    return (getController().getR().minus(m_observer.getXhat())).get(index, 0);
292  }
293
294  /**
295   * Get the function used to clamp the input u.
296   *
297   * @return The clamping function.
298   */
299  public Function<Matrix<Inputs, N1>, Matrix<Inputs, N1>> getClampFunction() {
300    return m_clampFunction;
301  }
302
303  /**
304   * Set the clamping function used to clamp inputs.
305   *
306   * @param clampFunction The clamping function.
307   */
308  public void setClampFunction(Function<Matrix<Inputs, N1>, Matrix<Inputs, N1>> clampFunction) {
309    this.m_clampFunction = clampFunction;
310  }
311
312  /**
313   * Correct the state estimate x-hat using the measurements in y.
314   *
315   * @param y Measurement vector.
316   */
317  public void correct(Matrix<Outputs, N1> y) {
318    getObserver().correct(getU(), y);
319  }
320
321  /**
322   * Sets new controller output, projects model forward, and runs observer prediction.
323   *
324   * <p>After calling this, the user should send the elements of u to the actuators.
325   *
326   * @param dtSeconds Timestep for model update.
327   */
328  public void predict(double dtSeconds) {
329    var u =
330        clampInput(
331            m_controller
332                .calculate(getObserver().getXhat(), m_nextR)
333                .plus(m_feedforward.calculate(m_nextR)));
334    getObserver().predict(u, dtSeconds);
335  }
336
337  /**
338   * Clamp the input u to the min and max.
339   *
340   * @param unclampedU The input to clamp.
341   * @return The clamped input.
342   */
343  public Matrix<Inputs, N1> clampInput(Matrix<Inputs, N1> unclampedU) {
344    return m_clampFunction.apply(unclampedU);
345  }
346}