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.Matrix;
008import edu.wpi.first.math.Nat;
009import edu.wpi.first.math.Num;
010import edu.wpi.first.math.numbers.N1;
011import edu.wpi.first.math.system.NumericalJacobian;
012import java.util.function.BiFunction;
013import java.util.function.Function;
014
015/**
016 * Constructs a control-affine plant inversion model-based feedforward from given model dynamics.
017 *
018 * <p>If given the vector valued function as f(x, u) where x is the state vector and u is the input
019 * vector, the B matrix(continuous input matrix) is calculated through a {@link
020 * edu.wpi.first.math.system.NumericalJacobian}. In this case f has to be control-affine (of the
021 * form f(x) + Bu).
022 *
023 * <p>The feedforward is calculated as <strong> u_ff = B<sup>+</sup> (rDot - f(x))</strong>, where
024 * <strong> B<sup>+</sup> </strong> is the pseudoinverse of B.
025 *
026 * <p>This feedforward does not account for a dynamic B matrix, B is either determined or supplied
027 * when the feedforward is created and remains constant.
028 *
029 * <p>For more on the underlying math, read
030 * https://file.tavsys.net/control/controls-engineering-in-frc.pdf.
031 */
032@SuppressWarnings({"ParameterName", "LocalVariableName", "MemberName", "ClassTypeParameterName"})
033public class ControlAffinePlantInversionFeedforward<States extends Num, Inputs extends Num> {
034  /** The current reference state. */
035  @SuppressWarnings("MemberName")
036  private Matrix<States, N1> m_r;
037
038  /** The computed feedforward. */
039  private Matrix<Inputs, N1> m_uff;
040
041  @SuppressWarnings("MemberName")
042  private final Matrix<States, Inputs> m_B;
043
044  private final Nat<Inputs> m_inputs;
045
046  private final double m_dt;
047
048  /** The model dynamics. */
049  private final BiFunction<Matrix<States, N1>, Matrix<Inputs, N1>, Matrix<States, N1>> m_f;
050
051  /**
052   * Constructs a feedforward with given model dynamics as a function of state and input.
053   *
054   * @param states A {@link Nat} representing the number of states.
055   * @param inputs A {@link Nat} representing the number of inputs.
056   * @param f A vector-valued function of x, the state, and u, the input, that returns the
057   *     derivative of the state vector. HAS to be control-affine (of the form f(x) + Bu).
058   * @param dtSeconds The timestep between calls of calculate().
059   */
060  public ControlAffinePlantInversionFeedforward(
061      Nat<States> states,
062      Nat<Inputs> inputs,
063      BiFunction<Matrix<States, N1>, Matrix<Inputs, N1>, Matrix<States, N1>> f,
064      double dtSeconds) {
065    this.m_dt = dtSeconds;
066    this.m_f = f;
067    this.m_inputs = inputs;
068
069    this.m_B =
070        NumericalJacobian.numericalJacobianU(
071            states, inputs, m_f, new Matrix<>(states, Nat.N1()), new Matrix<>(inputs, Nat.N1()));
072
073    m_r = new Matrix<>(states, Nat.N1());
074    m_uff = new Matrix<>(inputs, Nat.N1());
075
076    reset();
077  }
078
079  /**
080   * Constructs a feedforward with given model dynamics as a function of state, and the plant's
081   * B(continuous input matrix) matrix.
082   *
083   * @param states A {@link Nat} representing the number of states.
084   * @param inputs A {@link Nat} representing the number of inputs.
085   * @param f A vector-valued function of x, the state, that returns the derivative of the state
086   *     vector.
087   * @param B Continuous input matrix of the plant being controlled.
088   * @param dtSeconds The timestep between calls of calculate().
089   */
090  public ControlAffinePlantInversionFeedforward(
091      Nat<States> states,
092      Nat<Inputs> inputs,
093      Function<Matrix<States, N1>, Matrix<States, N1>> f,
094      Matrix<States, Inputs> B,
095      double dtSeconds) {
096    this.m_dt = dtSeconds;
097    this.m_inputs = inputs;
098
099    this.m_f = (x, u) -> f.apply(x);
100    this.m_B = B;
101
102    m_r = new Matrix<>(states, Nat.N1());
103    m_uff = new Matrix<>(inputs, Nat.N1());
104
105    reset();
106  }
107
108  /**
109   * Returns the previously calculated feedforward as an input vector.
110   *
111   * @return The calculated feedforward.
112   */
113  public Matrix<Inputs, N1> getUff() {
114    return m_uff;
115  }
116
117  /**
118   * Returns an element of the previously calculated feedforward.
119   *
120   * @param row Row of uff.
121   * @return The row of the calculated feedforward.
122   */
123  public double getUff(int row) {
124    return m_uff.get(row, 0);
125  }
126
127  /**
128   * Returns the current reference vector r.
129   *
130   * @return The current reference vector.
131   */
132  public Matrix<States, N1> getR() {
133    return m_r;
134  }
135
136  /**
137   * Returns an element of the current reference vector r.
138   *
139   * @param row Row of r.
140   * @return The row of the current reference vector.
141   */
142  public double getR(int row) {
143    return m_r.get(row, 0);
144  }
145
146  /**
147   * Resets the feedforward with a specified initial state vector.
148   *
149   * @param initialState The initial state vector.
150   */
151  public void reset(Matrix<States, N1> initialState) {
152    m_r = initialState;
153    m_uff.fill(0.0);
154  }
155
156  /** Resets the feedforward with a zero initial state vector. */
157  public void reset() {
158    m_r.fill(0.0);
159    m_uff.fill(0.0);
160  }
161
162  /**
163   * Calculate the feedforward with only the desired future reference. This uses the internally
164   * stored "current" reference.
165   *
166   * <p>If this method is used the initial state of the system is the one set using {@link
167   * LinearPlantInversionFeedforward#reset(Matrix)}. If the initial state is not set it defaults to
168   * a zero vector.
169   *
170   * @param nextR The reference state of the future timestep (k + dt).
171   * @return The calculated feedforward.
172   */
173  public Matrix<Inputs, N1> calculate(Matrix<States, N1> nextR) {
174    return calculate(m_r, nextR);
175  }
176
177  /**
178   * Calculate the feedforward with current and future reference vectors.
179   *
180   * @param r The reference state of the current timestep (k).
181   * @param nextR The reference state of the future timestep (k + dt).
182   * @return The calculated feedforward.
183   */
184  @SuppressWarnings({"ParameterName", "LocalVariableName"})
185  public Matrix<Inputs, N1> calculate(Matrix<States, N1> r, Matrix<States, N1> nextR) {
186    var rDot = (nextR.minus(r)).div(m_dt);
187
188    m_uff = m_B.solve(rDot.minus(m_f.apply(r, new Matrix<>(m_inputs, Nat.N1()))));
189
190    m_r = nextR;
191    return m_uff;
192  }
193}