A `SwerveSubsystem` shell for GRT 2022-2023.

FollowPathCommand.java

package frc.robot.commands.swerve;

import java.util.List;

import edu.wpi.first.math.controller.PIDController;
import edu.wpi.first.math.controller.ProfiledPIDController;
import edu.wpi.first.math.geometry.Pose2d;
import edu.wpi.first.math.geometry.Rotation2d;
import edu.wpi.first.math.geometry.Translation2d;
import edu.wpi.first.math.trajectory.Trajectory;
import edu.wpi.first.math.trajectory.TrajectoryConfig;
import edu.wpi.first.math.trajectory.TrajectoryGenerator;
import edu.wpi.first.math.trajectory.TrapezoidProfile;
import edu.wpi.first.math.trajectory.constraint.SwerveDriveKinematicsConstraint;
import edu.wpi.first.wpilibj2.command.SwerveControllerCommand;

import frc.robot.subsystems.swerve.SwerveSubsystem;

public class FollowPathCommand extends SwerveControllerCommand {
    // TODO: tune / measure
    private static final double Kp = 3.3283;
    private static final double Ki = 0;
    private static final double Kd = 0;

    /**
     * Creates a FollowPathCommand from a given trajectory.
     * 
     * @param swerveSubsystem The swerve subsystem.
     * @param trajectory The trajectory to follow.
     */
    public FollowPathCommand(SwerveSubsystem swerveSubsystem, Trajectory trajectory) {
        super(
            trajectory,
            swerveSubsystem::getRobotPosition,
            swerveSubsystem.getKinematics(),
            new PIDController(Kp, Ki, Kd),
            new PIDController(Kp, Ki, Kd),
            new ProfiledPIDController(
                Kp, Ki, Kd, 
                new TrapezoidProfile.Constraints(SwerveSubsystem.MAX_VEL, SwerveSubsystem.MAX_ACCEL)
            ),
            () -> new Rotation2d(), // TODO: this can control the angle of swerve at every timestep; hub locking?
            swerveSubsystem::setSwerveModuleStates,
            swerveSubsystem
        );
    }

    /**
     * Creates a FollowPathCommand from a given start point, list of waypoints, end point, and boolean representing whether
     * the path should be reversed (if the robot should drive backwards through the trajectory).
     * 
     * @param swerveSubsystem The swerve subsystem.
     * @param start The start point of the trajectory as a Pose2d.
     * @param waypoints A list of waypoints the robot must pass through as a List<Translation2d>.
     * @param end The end point of the trajectory as a Pose2d.
     * @param reversed Whether the trajectory is reversed.
     */
    public FollowPathCommand(SwerveSubsystem swerveSubsystem, Pose2d start, List<Translation2d> waypoints, Pose2d end, boolean reversed) {
        this(
            swerveSubsystem,
            // Target trajectory
            TrajectoryGenerator.generateTrajectory(
                start, waypoints, end, 
                new TrajectoryConfig(SwerveSubsystem.MAX_VEL, SwerveSubsystem.MAX_ACCEL)
                    .setReversed(reversed)
                    .setKinematics(swerveSubsystem.getKinematics())
                    .addConstraint(
                        new SwerveDriveKinematicsConstraint(
                            swerveSubsystem.getKinematics(), 
                            SwerveSubsystem.MAX_VEL
                        )
                    )
            )
        );
    }

    /**
     * Creates a FollowPathCommand from a given start point, list of waypoints, and end point.
     * 
     * @param swerveSubsystem The swerve subsystem.
     * @param start The start point of the trajectory as a Pose2d.
     * @param waypoints A list of waypoints the robot must pass through as a List<Translation2d>.
     * @param end The end point of the trajectory as a Pose2d.
     */
    public FollowPathCommand(SwerveSubsystem swerveSubsystem, Pose2d start, List<Translation2d> waypoints, Pose2d end) {
        this(swerveSubsystem, start, waypoints, end, false);
    }
}

SwerveModule.java

package frc.robot.subsystems.swerve;

import com.revrobotics.CANSparkMax;
import com.revrobotics.RelativeEncoder;
import com.revrobotics.SparkMaxPIDController;
import com.revrobotics.CANSparkMax.ControlType;
import com.revrobotics.CANSparkMax.IdleMode;
import com.revrobotics.CANSparkMaxLowLevel.MotorType;

import edu.wpi.first.math.geometry.Rotation2d;
import edu.wpi.first.math.kinematics.SwerveModuleState;

public class SwerveModule {
    private final CANSparkMax driveMotor;
    private final RelativeEncoder driveEncoder;
    private final SparkMaxPIDController drivePidController;

    private final CANSparkMax steerMotor;
    private final RelativeEncoder steerEncoder;
    private final SparkMaxPIDController steerPidController;

    private static final double DRIVE_ROTATIONS_TO_METERS = 0;
    private static final double STEER_ROTATIONS_TO_RADIANS = 0;

    private static final double driveP = 0;
    private static final double driveI = 0;
    private static final double driveD = 0;
    private static final double driveFF = 0;

    private static final double steerP = 0;
    private static final double steerI = 0;
    private static final double steerD = 0;
    private static final double steerFF = 0;

    public SwerveModule(int drivePort, int rotationPort) {
        driveMotor = new CANSparkMax(drivePort, MotorType.kBrushless);
        driveMotor.restoreFactoryDefaults();
        driveMotor.setIdleMode(IdleMode.kBrake);

        driveEncoder = driveMotor.getEncoder();
        driveEncoder.setVelocityConversionFactor(DRIVE_ROTATIONS_TO_METERS / 60); // RPM -> m/s

        drivePidController = driveMotor.getPIDController();
        drivePidController.setP(driveP);
        drivePidController.setI(driveI);
        drivePidController.setD(driveD);
        drivePidController.setFF(driveFF);

        steerMotor = new CANSparkMax(rotationPort, MotorType.kBrushless);
        steerMotor.restoreFactoryDefaults();
        steerMotor.setIdleMode(IdleMode.kBrake);

        steerEncoder = steerMotor.getAlternateEncoder(4096);
        steerEncoder.setPositionConversionFactor(STEER_ROTATIONS_TO_RADIANS);

        steerPidController = steerMotor.getPIDController();
        steerPidController.setP(steerP);
        steerPidController.setI(steerI);
        steerPidController.setD(steerD);
        steerPidController.setFF(steerFF);
    }

    /**
     * Gets the current state of the module as a `SwerveModuleState`.
     * @return The state of the module.
     */
    public SwerveModuleState getState() {
        return new SwerveModuleState(
            driveEncoder.getVelocity(), 
            getAngle()
        );
    }

    /**
     * Sets the desired state of the module.
     * @param state The desired state of the module as a `SwerveModuleState`.
     */
    public void setDesiredState(SwerveModuleState state) {
        SwerveModuleState optimized = SwerveModuleState.optimize(state, getAngle());
        drivePidController.setReference(optimized.speedMetersPerSecond, ControlType.kVelocity);
        steerPidController.setReference(optimized.angle.getRadians(), ControlType.kSmartMotion);
    }

    /**
     * Returns the current angle of the module.
     * @return The current angle of the module, as a `Rotation2d`.
     */
    private Rotation2d getAngle() {
        return new Rotation2d(steerEncoder.getPosition());
    }
}

SwerveSubsystem.java

package frc.robot.subsystems.swerve;

import com.kauailabs.navx.frc.AHRS;

import edu.wpi.first.math.MatBuilder;
import edu.wpi.first.math.Nat;
import edu.wpi.first.math.estimator.SwerveDrivePoseEstimator;
import edu.wpi.first.math.geometry.Pose2d;
import edu.wpi.first.math.geometry.Rotation2d;
import edu.wpi.first.math.geometry.Translation2d;
import edu.wpi.first.math.kinematics.ChassisSpeeds;
import edu.wpi.first.math.kinematics.SwerveDriveKinematics;
import edu.wpi.first.math.kinematics.SwerveModuleState;
import edu.wpi.first.wpilibj.SPI;
import edu.wpi.first.wpilibj2.command.SubsystemBase;

public class SwerveSubsystem extends SubsystemBase {
    private final SwerveModule topLeftModule;
    private final SwerveModule topRightModule;
    private final SwerveModule bottomLeftModule;
    private final SwerveModule bottomRightModule;

    private final SwerveDrivePoseEstimator poseEstimator;
    private final AHRS ahrs;

    private final SwerveDriveKinematics kinematics;

    public static final double MAX_VEL = 3; // Max tangential velocity, in m/s
    public static final double MAX_ACCEL = 3; // Max tangential acceleration, in m/s^2
    public static final double MAX_OMEGA = Math.toRadians(30); // Max angular velocity, in rads/s

    public SwerveSubsystem() {
        // Initialize swerve modules
        // TODO: CAN IDs
        topLeftModule = new SwerveModule(0, 1);
        topRightModule = new SwerveModule(2, 3);
        bottomLeftModule = new SwerveModule(4, 5);
        bottomRightModule = new SwerveModule(6, 7);

        // Initialize system kinematics with top left, top right, bottom left, and bottom right swerve
        // module positions
        // TODO: positions
        kinematics = new SwerveDriveKinematics(
            new Translation2d(),
            new Translation2d(),
            new Translation2d(),
            new Translation2d()
        );

        // Initialize NaxX and pose estimator
        ahrs = new AHRS(SPI.Port.kMXP);
        poseEstimator = new SwerveDrivePoseEstimator(
            getGyroHeading(), 
            new Pose2d(), 
            kinematics, 
            // State measurement standard deviations: [X, Y, theta]
            new MatBuilder<>(Nat.N3(), Nat.N1()).fill(0.02, 0.02, 0.01), 
            // Odometry measurement standard deviations: [gyro]
            new MatBuilder<>(Nat.N1(), Nat.N1()).fill(0.01), 
            // Vision measurement standard deviations: [X, Y, theta]
            new MatBuilder<>(Nat.N3(), Nat.N1()).fill(0.1, 0.1, 0.01)
        );
    }

    /**
     * Set the field-centric swerve drive powers of the subsystem.
     * @param xPower The power [-1.0, 1.0] in the x (forward) direction.
     * @param yPower The power [-1.0, 1.0] in the y (left) direction.
     * @param angularPower The angular (rotational) power [-1.0, 1.0].
     */
    public void setSwerveDrivePowers(double xPower, double yPower, double angularPower) {
        /*
        swerveSubsystem.setDefaultCommand(new RunCommand(() -> {
            double xPower = -driveController.getLeftY();
            double yPower = -driveController.getLeftX();
            double angularPower = -driveController.getRightX();
            swerveSubsystem.setSwerveDrivePowers(xPower, yPower, angularPower);
        }, swerveSubsystem));
        */

        // Scale [-1.0, 1.0] powers to desired velocity, turning field-relative powers
        // into robot relative chassis speeds.
        ChassisSpeeds speeds = ChassisSpeeds.fromFieldRelativeSpeeds(
            xPower * MAX_VEL, 
            yPower * MAX_VEL, 
            angularPower * MAX_OMEGA,
            getGyroHeading()
        );

        // Calculate swerve module states from desired chassis speeds, desaturating them to
        // ensure all velocities are under MAX_VEL after kinematics.
        SwerveModuleState[] states = kinematics.toSwerveModuleStates(speeds);
        SwerveDriveKinematics.desaturateWheelSpeeds(states, MAX_VEL);

        // Set module states
        setSwerveModuleStates(states);
    }

    /**
     * Sets the swerve module states of this subsystem. Module states are assumed to be passed
     * in a tuple of [top left, top right, bottom left, bottom right].
     * @param states The swerve module states to set.
     */
    public void setSwerveModuleStates(SwerveModuleState... states) {
        topLeftModule.setDesiredState(states[0]);
        topRightModule.setDesiredState(states[1]);
        bottomLeftModule.setDesiredState(states[2]);
        bottomRightModule.setDesiredState(states[3]);
    }

    /**
     * Gets the subsystems `SwerveDriveKinematics` instance.
     * @return The SwerveDriveKinematics representing this system's kinematics.
     */
    public SwerveDriveKinematics getKinematics() {
        return kinematics;
    }

    @Override
    public void periodic() {
        // Update pose estimator from swerve module states
        Rotation2d gyroAngle = getGyroHeading();
        poseEstimator.update(
            gyroAngle,
            topLeftModule.getState(),
            topRightModule.getState(),
            bottomLeftModule.getState(),
            bottomRightModule.getState()
        );
    }

    /**
     * Gets the estimated current position of the robot.
     * @return The estimated position of the robot as a Pose2d.
     */
    public Pose2d getRobotPosition() {
        return poseEstimator.getEstimatedPosition();
    }

    /**
     * Reset the robot's position to a given Pose2d.
     * @param position The position to reset the pose estimator to.
     */
    public void resetPosition(Pose2d position) {
        Rotation2d gyroAngle = getGyroHeading();
        poseEstimator.resetPosition(position, gyroAngle);
    }

    /**
     * Zeros the robot's position.
     * This method zeros both the robot's translation *and* rotation.
     */
    public void resetPosition() {
        resetPosition(new Pose2d());
    }

    /**
     * Gets the gyro angle given by the NavX AHRS, inverted to be counterclockwise positive.
     * @return The robot heading as a Rotation2d.
     */
    private Rotation2d getGyroHeading() {
        return Rotation2d.fromDegrees(-ahrs.getAngle());
    }
}