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.wpilibj.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@SuppressWarnings("MemberName")
033public class SwerveControllerCommand extends CommandBase {
034  private final Timer m_timer = new Timer();
035  private final Trajectory m_trajectory;
036  private final Supplier<Pose2d> m_pose;
037  private final SwerveDriveKinematics m_kinematics;
038  private final HolonomicDriveController m_controller;
039  private final Consumer<SwerveModuleState[]> m_outputModuleStates;
040  private final Supplier<Rotation2d> m_desiredRotation;
041
042  /**
043   * Constructs a new SwerveControllerCommand that when executed will follow the provided
044   * trajectory. This command will not return output voltages but rather raw module states from the
045   * position controllers which need to be put into a velocity PID.
046   *
047   * <p>Note: The controllers will *not* set the outputVolts to zero upon completion of the path-
048   * this is left to the user, since it is not appropriate for paths with nonstationary endstates.
049   *
050   * @param trajectory The trajectory to follow.
051   * @param pose A function that supplies the robot pose - use one of the odometry classes to
052   *     provide this.
053   * @param kinematics The kinematics for the robot drivetrain.
054   * @param xController The Trajectory Tracker PID controller for the robot's x position.
055   * @param yController The Trajectory Tracker PID controller for the robot's y position.
056   * @param thetaController The Trajectory Tracker PID controller for angle for the robot.
057   * @param desiredRotation The angle that the drivetrain should be facing. This is sampled at each
058   *     time step.
059   * @param outputModuleStates The raw output module states from the position controllers.
060   * @param requirements The subsystems to require.
061   */
062  @SuppressWarnings("ParameterName")
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, "xController", "SwerveControllerCommand"),
081            requireNonNullParam(thetaController, "thetaController", "SwerveControllerCommand"));
082
083    m_outputModuleStates =
084        requireNonNullParam(outputModuleStates, "frontLeftOutput", "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  @SuppressWarnings("ParameterName")
116  public SwerveControllerCommand(
117      Trajectory trajectory,
118      Supplier<Pose2d> pose,
119      SwerveDriveKinematics kinematics,
120      PIDController xController,
121      PIDController yController,
122      ProfiledPIDController thetaController,
123      Consumer<SwerveModuleState[]> outputModuleStates,
124      Subsystem... requirements) {
125    this(
126        trajectory,
127        pose,
128        kinematics,
129        xController,
130        yController,
131        thetaController,
132        () ->
133            trajectory.getStates().get(trajectory.getStates().size() - 1).poseMeters.getRotation(),
134        outputModuleStates,
135        requirements);
136  }
137
138  @Override
139  public void initialize() {
140    m_timer.reset();
141    m_timer.start();
142  }
143
144  @Override
145  @SuppressWarnings("LocalVariableName")
146  public void execute() {
147    double curTime = m_timer.get();
148    var desiredState = m_trajectory.sample(curTime);
149
150    var targetChassisSpeeds =
151        m_controller.calculate(m_pose.get(), desiredState, m_desiredRotation.get());
152    var targetModuleStates = m_kinematics.toSwerveModuleStates(targetChassisSpeeds);
153
154    m_outputModuleStates.accept(targetModuleStates);
155  }
156
157  @Override
158  public void end(boolean interrupted) {
159    m_timer.stop();
160  }
161
162  @Override
163  public boolean isFinished() {
164    return m_timer.hasElapsed(m_trajectory.getTotalTimeSeconds());
165  }
166}