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 CommandBase {
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, since it is not appropriate for paths with nonstationary endstates.
050   *
051   * @param trajectory The trajectory to follow.
052   * @param pose A function that supplies the robot pose - use one of the odometry classes to
053   *     provide this.
054   * @param kinematics The kinematics for the robot drivetrain.
055   * @param xController The Trajectory Tracker PID controller for the robot's x position.
056   * @param yController The Trajectory Tracker PID controller for the robot's y position.
057   * @param thetaController The Trajectory Tracker PID controller for angle for the robot.
058   * @param desiredRotation The angle that the drivetrain should be facing. This is sampled at each
059   *     time step.
060   * @param outputModuleStates The raw output module states from the position controllers.
061   * @param requirements The subsystems to require.
062   */
063  public SwerveControllerCommand(
064      Trajectory trajectory,
065      Supplier<Pose2d> pose,
066      SwerveDriveKinematics kinematics,
067      PIDController xController,
068      PIDController yController,
069      ProfiledPIDController thetaController,
070      Supplier<Rotation2d> desiredRotation,
071      Consumer<SwerveModuleState[]> outputModuleStates,
072      Subsystem... requirements) {
073    m_trajectory = requireNonNullParam(trajectory, "trajectory", "SwerveControllerCommand");
074    m_pose = requireNonNullParam(pose, "pose", "SwerveControllerCommand");
075    m_kinematics = requireNonNullParam(kinematics, "kinematics", "SwerveControllerCommand");
076
077    m_controller =
078        new HolonomicDriveController(
079            requireNonNullParam(xController, "xController", "SwerveControllerCommand"),
080            requireNonNullParam(yController, "yController", "SwerveControllerCommand"),
081            requireNonNullParam(thetaController, "thetaController", "SwerveControllerCommand"));
082
083    m_outputModuleStates =
084        requireNonNullParam(outputModuleStates, "outputModuleStates", "SwerveControllerCommand");
085
086    m_desiredRotation =
087        requireNonNullParam(desiredRotation, "desiredRotation", "SwerveControllerCommand");
088
089    addRequirements(requirements);
090  }
091
092  /**
093   * Constructs a new SwerveControllerCommand that when executed will follow the provided
094   * trajectory. This command will not return output voltages but rather raw module states from the
095   * position controllers which need to be put into a velocity PID.
096   *
097   * <p>Note: The controllers will *not* set the outputVolts to zero upon completion of the path-
098   * this is left to the user, since it is not appropriate for paths with nonstationary endstates.
099   *
100   * <p>Note 2: The final rotation of the robot will be set to the rotation of the final pose in the
101   * trajectory. The robot will not follow the rotations from the poses at each timestep. If
102   * alternate rotation behavior is desired, the other constructor with a supplier for rotation
103   * should be used.
104   *
105   * @param trajectory The trajectory to follow.
106   * @param pose A function that supplies the robot pose - use one of the odometry classes to
107   *     provide this.
108   * @param kinematics The kinematics for the robot drivetrain.
109   * @param xController The Trajectory Tracker PID controller for the robot's x position.
110   * @param yController The Trajectory Tracker PID controller for the robot's y position.
111   * @param thetaController The Trajectory Tracker PID controller for angle for the robot.
112   * @param outputModuleStates The raw output module states from the position controllers.
113   * @param requirements The subsystems to require.
114   */
115  public SwerveControllerCommand(
116      Trajectory trajectory,
117      Supplier<Pose2d> pose,
118      SwerveDriveKinematics kinematics,
119      PIDController xController,
120      PIDController yController,
121      ProfiledPIDController thetaController,
122      Consumer<SwerveModuleState[]> outputModuleStates,
123      Subsystem... requirements) {
124    this(
125        trajectory,
126        pose,
127        kinematics,
128        xController,
129        yController,
130        thetaController,
131        () ->
132            trajectory.getStates().get(trajectory.getStates().size() - 1).poseMeters.getRotation(),
133        outputModuleStates,
134        requirements);
135  }
136
137  @Override
138  public void initialize() {
139    m_timer.reset();
140    m_timer.start();
141  }
142
143  @Override
144  public void execute() {
145    double curTime = m_timer.get();
146    var desiredState = m_trajectory.sample(curTime);
147
148    var targetChassisSpeeds =
149        m_controller.calculate(m_pose.get(), desiredState, m_desiredRotation.get());
150    var targetModuleStates = m_kinematics.toSwerveModuleStates(targetChassisSpeeds);
151
152    m_outputModuleStates.accept(targetModuleStates);
153  }
154
155  @Override
156  public void end(boolean interrupted) {
157    m_timer.stop();
158  }
159
160  @Override
161  public boolean isFinished() {
162    return m_timer.hasElapsed(m_trajectory.getTotalTimeSeconds());
163  }
164}