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.wpilibj2.command;
006
007import static edu.wpi.first.util.ErrorMessages.requireNonNullParam;
008
009import edu.wpi.first.math.controller.HolonomicDriveController;
010import edu.wpi.first.math.controller.PIDController;
011import edu.wpi.first.math.controller.ProfiledPIDController;
012import edu.wpi.first.math.geometry.Pose2d;
013import edu.wpi.first.math.geometry.Rotation2d;
014import edu.wpi.first.math.kinematics.SwerveDriveKinematics;
015import edu.wpi.first.math.kinematics.SwerveModuleState;
016import edu.wpi.first.math.trajectory.Trajectory;
017import edu.wpi.first.wpilibj.Timer;
018import java.util.function.Consumer;
019import java.util.function.Supplier;
020
021/**
022 * A command that uses two PID controllers ({@link PIDController}) and a ProfiledPIDController
023 * ({@link ProfiledPIDController}) to follow a trajectory {@link Trajectory} with a swerve drive.
024 *
025 * <p>This command outputs the raw desired Swerve Module States ({@link SwerveModuleState}) in an
026 * array. The desired wheel and module rotation velocities should be taken from those and used in
027 * velocity PIDs.
028 *
029 * <p>The robot angle controller does not follow the angle given by the trajectory but rather goes
030 * to the angle given in the final state of the trajectory.
031 *
032 * <p>This class is provided by the NewCommands VendorDep
033 */
034public class SwerveControllerCommand extends Command {
035  private final Timer m_timer = new Timer();
036  private final Trajectory m_trajectory;
037  private final Supplier<Pose2d> m_pose;
038  private final SwerveDriveKinematics m_kinematics;
039  private final HolonomicDriveController m_controller;
040  private final Consumer<SwerveModuleState[]> m_outputModuleStates;
041  private final Supplier<Rotation2d> m_desiredRotation;
042
043  /**
044   * Constructs a new SwerveControllerCommand that when executed will follow the provided
045   * trajectory. This command will not return output voltages but rather raw module states from the
046   * position controllers which need to be put into a velocity PID.
047   *
048   * <p>Note: The controllers will *not* set the outputVolts to zero upon completion of the path.
049   * This is left to the user to do since it is not appropriate for paths with nonstationary
050   * endstates.
051   *
052   * @param trajectory The trajectory to follow.
053   * @param pose A function that supplies the robot pose - use one of the odometry classes to
054   *     provide this.
055   * @param kinematics The kinematics for the robot drivetrain.
056   * @param xController The Trajectory Tracker PID controller for the robot's x position.
057   * @param yController The Trajectory Tracker PID controller for the robot's y position.
058   * @param thetaController The Trajectory Tracker PID controller for angle for the robot.
059   * @param desiredRotation The angle that the drivetrain should be facing. This is sampled at each
060   *     time step.
061   * @param outputModuleStates The raw output module states from the position controllers.
062   * @param requirements The subsystems to require.
063   */
064  public SwerveControllerCommand(
065      Trajectory trajectory,
066      Supplier<Pose2d> pose,
067      SwerveDriveKinematics kinematics,
068      PIDController xController,
069      PIDController yController,
070      ProfiledPIDController thetaController,
071      Supplier<Rotation2d> desiredRotation,
072      Consumer<SwerveModuleState[]> outputModuleStates,
073      Subsystem... requirements) {
074    this(
075        trajectory,
076        pose,
077        kinematics,
078        new HolonomicDriveController(
079            requireNonNullParam(xController, "xController", "SwerveControllerCommand"),
080            requireNonNullParam(yController, "yController", "SwerveControllerCommand"),
081            requireNonNullParam(thetaController, "thetaController", "SwerveControllerCommand")),
082        desiredRotation,
083        outputModuleStates,
084        requirements);
085  }
086
087  /**
088   * Constructs a new SwerveControllerCommand that when executed will follow the provided
089   * trajectory. This command will not return output voltages but rather raw module states from the
090   * position controllers which need to be put into a velocity PID.
091   *
092   * <p>Note: The controllers will *not* set the outputVolts to zero upon completion of the path.
093   * This is left to the user since it is not appropriate for paths with nonstationary endstates.
094   *
095   * <p>Note 2: The final rotation of the robot will be set to the rotation of the final pose in the
096   * trajectory. The robot will not follow the rotations from the poses at each timestep. If
097   * alternate rotation behavior is desired, the other constructor with a supplier for rotation
098   * should be used.
099   *
100   * @param trajectory The trajectory to follow.
101   * @param pose A function that supplies the robot pose - use one of the odometry classes to
102   *     provide this.
103   * @param kinematics The kinematics for the robot drivetrain.
104   * @param xController The Trajectory Tracker PID controller for the robot's x position.
105   * @param yController The Trajectory Tracker PID controller for the robot's y position.
106   * @param thetaController The Trajectory Tracker PID controller for angle for the robot.
107   * @param outputModuleStates The raw output module states from the position controllers.
108   * @param requirements The subsystems to require.
109   */
110  public SwerveControllerCommand(
111      Trajectory trajectory,
112      Supplier<Pose2d> pose,
113      SwerveDriveKinematics kinematics,
114      PIDController xController,
115      PIDController yController,
116      ProfiledPIDController thetaController,
117      Consumer<SwerveModuleState[]> outputModuleStates,
118      Subsystem... requirements) {
119    this(
120        trajectory,
121        pose,
122        kinematics,
123        xController,
124        yController,
125        thetaController,
126        () -> trajectory.getStates().get(trajectory.getStates().size() - 1).pose.getRotation(),
127        outputModuleStates,
128        requirements);
129  }
130
131  /**
132   * Constructs a new SwerveControllerCommand that when executed will follow the provided
133   * trajectory. This command will not return output voltages but rather raw module states from the
134   * position controllers which need to be put into a velocity PID.
135   *
136   * <p>Note: The controllers will *not* set the outputVolts to zero upon completion of the path-
137   * this is left to the user, since it is not appropriate for paths with nonstationary endstates.
138   *
139   * <p>Note 2: The final rotation of the robot will be set to the rotation of the final pose in the
140   * trajectory. The robot will not follow the rotations from the poses at each timestep. If
141   * alternate rotation behavior is desired, the other constructor with a supplier for rotation
142   * should be used.
143   *
144   * @param trajectory The trajectory to follow.
145   * @param pose A function that supplies the robot pose - use one of the odometry classes to
146   *     provide this.
147   * @param kinematics The kinematics for the robot drivetrain.
148   * @param controller The HolonomicDriveController for the drivetrain.
149   * @param outputModuleStates The raw output module states from the position controllers.
150   * @param requirements The subsystems to require.
151   */
152  public SwerveControllerCommand(
153      Trajectory trajectory,
154      Supplier<Pose2d> pose,
155      SwerveDriveKinematics kinematics,
156      HolonomicDriveController controller,
157      Consumer<SwerveModuleState[]> outputModuleStates,
158      Subsystem... requirements) {
159    this(
160        trajectory,
161        pose,
162        kinematics,
163        controller,
164        () -> trajectory.getStates().get(trajectory.getStates().size() - 1).pose.getRotation(),
165        outputModuleStates,
166        requirements);
167  }
168
169  /**
170   * Constructs a new SwerveControllerCommand that when executed will follow the provided
171   * trajectory. This command will not return output voltages but rather raw module states from the
172   * position controllers which need to be put into a velocity PID.
173   *
174   * <p>Note: The controllers will *not* set the outputVolts to zero upon completion of the path-
175   * this is left to the user, since it is not appropriate for paths with nonstationary endstates.
176   *
177   * @param trajectory The trajectory to follow.
178   * @param pose A function that supplies the robot pose - use one of the odometry classes to
179   *     provide this.
180   * @param kinematics The kinematics for the robot drivetrain.
181   * @param controller The HolonomicDriveController for the drivetrain.
182   * @param desiredRotation The angle that the drivetrain should be facing. This is sampled at each
183   *     time step.
184   * @param outputModuleStates The raw output module states from the position controllers.
185   * @param requirements The subsystems to require.
186   */
187  @SuppressWarnings("this-escape")
188  public SwerveControllerCommand(
189      Trajectory trajectory,
190      Supplier<Pose2d> pose,
191      SwerveDriveKinematics kinematics,
192      HolonomicDriveController controller,
193      Supplier<Rotation2d> desiredRotation,
194      Consumer<SwerveModuleState[]> outputModuleStates,
195      Subsystem... requirements) {
196    m_trajectory = requireNonNullParam(trajectory, "trajectory", "SwerveControllerCommand");
197    m_pose = requireNonNullParam(pose, "pose", "SwerveControllerCommand");
198    m_kinematics = requireNonNullParam(kinematics, "kinematics", "SwerveControllerCommand");
199    m_controller = requireNonNullParam(controller, "controller", "SwerveControllerCommand");
200
201    m_desiredRotation =
202        requireNonNullParam(desiredRotation, "desiredRotation", "SwerveControllerCommand");
203
204    m_outputModuleStates =
205        requireNonNullParam(outputModuleStates, "outputModuleStates", "SwerveControllerCommand");
206
207    addRequirements(requirements);
208  }
209
210  @Override
211  public void initialize() {
212    m_timer.restart();
213  }
214
215  @Override
216  public void execute() {
217    double curTime = m_timer.get();
218    var desiredState = m_trajectory.sample(curTime);
219
220    var targetChassisSpeeds =
221        m_controller.calculate(m_pose.get(), desiredState, m_desiredRotation.get());
222    var targetModuleStates = m_kinematics.toSwerveModuleStates(targetChassisSpeeds);
223
224    m_outputModuleStates.accept(targetModuleStates);
225  }
226
227  @Override
228  public void end(boolean interrupted) {
229    m_timer.stop();
230  }
231
232  @Override
233  public boolean isFinished() {
234    return m_timer.hasElapsed(m_trajectory.getTotalTime());
235  }
236}