Add SwerveBot example

SwerveBot/drivetrain.py

@@ -0,0 +1,99 @@
+#
+# Copyright (c) FIRST and other WPILib contributors.
+# Open Source Software; you can modify and/or share it under the terms of
+# the WPILib BSD license file in the root directory of this project.
+#
+
+import math
+import wpilib
+import wpimath.geometry
+import wpimath.kinematics
+import swervemodule
+
+kMaxSpeed = 3.0 # 3 meters per second
+kMaxAngularSpeed = math.pi # 1/2 rotation per second
+
+
+class Drivetrain:
+ """
+ Represents a swerve drive style drivetrain.
+ """
+
+ def __init__(self) -> None:
+ self.frontLeftLocation = wpimath.geometry.Translation2d(0.381, 0.381)
+ self.frontRightLocation = wpimath.geometry.Translation2d(0.381, -0.381)
+ self.backLeftLocation = wpimath.geometry.Translation2d(-0.381, 0.381)
+ self.backRightLocation = wpimath.geometry.Translation2d(-0.381, -0.381)
+
+ self.frontLeft = swervemodule.SwerveModule(1, 2, 0, 1, 2, 3)
+ self.frontRight = swervemodule.SwerveModule(3, 4, 4, 5, 6, 7)
+ self.backLeft = swervemodule.SwerveModule(5, 6, 8, 9, 10, 11)
+ self.backRight = swervemodule.SwerveModule(7, 8, 12, 13, 14, 15)
+
+ self.gyro = wpilib.AnalogGyro(0)
+
+ self.kinematics = wpimath.kinematics.SwerveDrive4Kinematics(
+ self.frontLeftLocation,
+ self.frontRightLocation,
+ self.backLeftLocation,
+ self.backRightLocation,
+ )
+
+ self.odometry = wpimath.kinematics.SwerveDrive4Odometry(
+ self.kinematics,
+ self.gyro.getRotation2d(),
+ (
+ self.frontLeft.getPosition(),
+ self.frontRight.getPosition(),
+ self.backLeft.getPosition(),
+ self.backRight.getPosition(),
+ ),
+ )
+
+ self.gyro.reset()
+
+ def drive(
+ self,
+ xSpeed: float,
+ ySpeed: float,
+ rot: float,
+ fieldRelative: bool,
+ periodSeconds: float,
+ ) -> None:
+ """
+ Method to drive the robot using joystick info.
+ :param xSpeed: Speed of the robot in the x direction (forward).
+ :param ySpeed: Speed of the robot in the y direction (sideways).
+ :param rot: Angular rate of the robot.
+ :param fieldRelative: Whether the provided x and y speeds are relative to the field.
+ :param periodSeconds: Time
+ """
+ swerveModuleStates = self.kinematics.toSwerveModuleStates(
+ wpimath.kinematics.ChassisSpeeds.discretize(
+ wpimath.kinematics.ChassisSpeeds.fromFieldRelativeSpeeds(
+ xSpeed, ySpeed, rot, self.gyro.getRotation2d()
+ )
+ if fieldRelative
+ else wpimath.kinematics.ChassisSpeeds(xSpeed, ySpeed, rot),
+ periodSeconds,
+ )
+ )
+ wpimath.kinematics.SwerveDrive4Kinematics.desaturateWheelSpeeds(
+ swerveModuleStates, kMaxSpeed
+ )
+ self.frontLeft.setDesiredState(swerveModuleStates[0])
+ self.frontRight.setDesiredState(swerveModuleStates[1])
+ self.backLeft.setDesiredState(swerveModuleStates[2])
+ self.backRight.setDesiredState(swerveModuleStates[3])
+
+ def updateOdometry(self) -> None:
+ """Updates the field relative position of the robot."""
+ self.odometry.update(
+ self.gyro.getRotation2d(),
+ (
+ self.frontLeft.getPosition(),
+ self.frontRight.getPosition(),
+ self.backLeft.getPosition(),
+ self.backRight.getPosition(),
+ ),
+ )

SwerveBot/robot.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) FIRST and other WPILib contributors.
+# Open Source Software; you can modify and/or share it under the terms of
+# the WPILib BSD license file in the root directory of this project.
+#
+
+import wpilib
+import wpimath
+import wpilib.drive
+import wpimath.filter
+import wpimath.controller
+import drivetrain
+
+
+class MyRobot(wpilib.TimedRobot):
+ def robotInit(self) -> None:
+ """Robot initialization function"""
+ self.controller = wpilib.XboxController(0)
+ self.swerve = drivetrain.Drivetrain()
+
+ # Slew rate limiters to make joystick inputs more gentle; 1/3 sec from 0 to 1.
+ self.xspeedLimiter = wpimath.filter.SlewRateLimiter(3)
+ self.yspeedLimiter = wpimath.filter.SlewRateLimiter(3)
+ self.rotLimiter = wpimath.filter.SlewRateLimiter(3)
+
+ def autonomousPeriodic(self) -> None:
+ self.driveWithJoystick(False)
+ self.swerve.updateOdometry()
+
+ def teleopPeriodic(self) -> None:
+ self.driveWithJoystick(True)
+
+ def driveWithJoystick(self, fieldRelative: bool) -> None:
+ # Get the x speed. We are inverting this because Xbox controllers return
+ # negative values when we push forward.
+ xSpeed = (
+ -self.xspeedLimiter.calculate(
+ wpimath.applyDeadband(self.controller.getLeftY(), 0.02)
+ )
+ * drivetrain.kMaxSpeed
+ )
+
+ # Get the y speed or sideways/strafe speed. We are inverting this because
+ # we want a positive value when we pull to the left. Xbox controllers
+ # return positive values when you pull to the right by default.
+ ySpeed = (
+ -self.yspeedLimiter.calculate(
+ wpimath.applyDeadband(self.controller.getLeftX(), 0.02)
+ )
+ * drivetrain.kMaxSpeed
+ )
+
+ # Get the rate of angular rotation. We are inverting this because we want a
+ # positive value when we pull to the left (remember, CCW is positive in
+ # mathematics). Xbox controllers return positive values when you pull to
+ # the right by default.
+ rot = (
+ -self.rotLimiter.calculate(
+ wpimath.applyDeadband(self.controller.getRightX(), 0.02)
+ )
+ * drivetrain.kMaxSpeed
+ )
+
+ self.swerve.drive(xSpeed, ySpeed, rot, fieldRelative, self.getPeriod())
+
+
+if __name__ == "__main__":
+ wpilib.run(MyRobot)

SwerveBot/swervemodule.py

@@ -0,0 +1,138 @@
+#
+# Copyright (c) FIRST and other WPILib contributors.
+# Open Source Software; you can modify and/or share it under the terms of
+# the WPILib BSD license file in the root directory of this project.
+#
+
+import math
+import wpilib
+import wpimath.kinematics
+import wpimath.geometry
+import wpimath.controller
+import wpimath.trajectory
+
+kWheelRadius = 0.0508
+kEncoderResolution = 4096
+kModuleMaxAngularVelocity = math.pi
+kModuleMaxAngularAcceleration = math.tau
+
+
+class SwerveModule:
+ def __init__(
+ self,
+ driveMotorChannel: int,
+ turningMotorChannel: int,
+ driveEncoderChannelA: int,
+ driveEncoderChannelB: int,
+ turningEncoderChannelA: int,
+ turningEncoderChannelB: int,
+ ) -> None:
+ """Constructs a SwerveModule with a drive motor, turning motor, drive encoder and turning encoder.
+
+ :param driveMotorChannel: PWM output for the drive motor.
+ :param turningMotorChannel: PWM output for the turning motor.
+ :param driveEncoderChannelA: DIO input for the drive encoder channel A
+ :param driveEncoderChannelB: DIO input for the drive encoder channel B
+ :param turningEncoderChannelA: DIO input for the turning encoder channel A
+ :param turningEncoderChannelB: DIO input for the turning encoder channel B
+ """
+ self.driveMotor = wpilib.PWMSparkMax(driveMotorChannel)
+ self.turningMotor = wpilib.PWMSparkMax(turningMotorChannel)
+
+ self.driveEncoder = wpilib.Encoder(driveEncoderChannelA, driveEncoderChannelB)
+ self.turningEncoder = wpilib.Encoder(
+ turningEncoderChannelA, turningEncoderChannelB
+ )
+
+ # Gains are for example purposes only - must be determined for your own robot!
+ self.drivePIDController = wpimath.controller.PIDController(1, 0, 0)
+
+ # Gains are for example purposes only - must be determined for your own robot!
+ self.turningPIDController = wpimath.controller.ProfiledPIDController(
+ 1,
+ 0,
+ 0,
+ wpimath.trajectory.TrapezoidProfile.Constraints(
+ kModuleMaxAngularVelocity,
+ kModuleMaxAngularAcceleration,
+ ),
+ )
+
+ # Gains are for example purposes only - must be determined for your own robot!
+ self.driveFeedforward = wpimath.controller.SimpleMotorFeedforwardMeters(1, 3)
+ self.turnFeedforward = wpimath.controller.SimpleMotorFeedforwardMeters(1, 0.5)
+
+ # Set the distance per pulse for the drive encoder. We can simply use the
+ # distance traveled for one rotation of the wheel divided by the encoder
+ # resolution.
+ self.driveEncoder.setDistancePerPulse(
+ math.tau * kWheelRadius / kEncoderResolution
+ )
+
+ # Set the distance (in this case, angle) in radians per pulse for the turning encoder.
+ # This is the the angle through an entire rotation (2 * pi) divided by the
+ # encoder resolution.
+ self.turningEncoder.setDistancePerPulse(math.tau / kEncoderResolution)
+
+ # Limit the PID Controller's input range between -pi and pi and set the input
+ # to be continuous.
+ self.turningPIDController.enableContinuousInput(-math.pi, math.pi)
+
+ def getState(self) -> wpimath.kinematics.SwerveModuleState:
+ """Returns the current state of the module.
+
+ :returns: The current state of the module.
+ """
+ return wpimath.kinematics.SwerveModuleState(
+ self.driveEncoder.getRate(),
+ wpimath.geometry.Rotation2d(self.turningEncoder.getDistance()),
+ )
+
+ def getPosition(self) -> wpimath.kinematics.SwerveModulePosition:
+ """Returns the current position of the module.
+
+ :returns: The current position of the module.
+ """
+ return wpimath.kinematics.SwerveModulePosition(
+ self.driveEncoder.getRate(),
+ wpimath.geometry.Rotation2d(self.turningEncoder.getDistance()),
+ )
+
+ def setDesiredState(
+ self, desiredState: wpimath.kinematics.SwerveModuleState
+ ) -> None:
+ """Sets the desired state for the module.
+
+ :param desiredState: Desired state with speed and angle.
+ """
+
+ encoderRotation = wpimath.geometry.Rotation2d(self.turningEncoder.getDistance())
+
+ # Optimize the reference state to avoid spinning further than 90 degrees
+ state = wpimath.kinematics.SwerveModuleState.optimize(
+ desiredState, encoderRotation
+ )
+
+ # Scale speed by cosine of angle error. This scales down movement perpendicular to the desired
+ # direction of travel that can occur when modules change directions. This results in smoother
+ # driving.
+ state.speed *= (state.angle - encoderRotation).cos()
+
+ # Calculate the drive output from the drive PID controller.
+ driveOutput = self.drivePIDController.calculate(
+ self.driveEncoder.getRate(), state.speed
+ )
+
+ driveFeedforward = self.driveFeedforward.calculate(state.speed)
+
+ # Calculate the turning motor output from the turning PID controller.
+ turnOutput = self.turningPIDController.calculate(
+ self.turningEncoder.getDistance(), state.angle.radians()
+ )
+
+ turnFeedforward = self.turnFeedforward.calculate(
+ self.turningPIDController.getSetpoint().velocity
+ )
+
+ self.driveMotor.setVoltage(driveOutput + driveFeedforward)
+ self.turningMotor.setVoltage(turnOutput + turnFeedforward)

run_tests.sh

@@ -45,6 +45,7 @@ BASE_TESTS="
  Solenoid
  StatefulAutonomous
  StateSpaceFlywheel
+ SwerveBot
  TankDrive
  TankDriveXboxController
  Timed/src