Skip to content

Commit d7df94c

Browse files
megarubbervirtuald
megarubber
and
virtuald
authored
Add SwerveBot example (#98)
Issue #49 Add the [SwerveBot](https://github.com/wpilibsuite/allwpilib/tree/main/wpilibjExamples/src/main/java/edu/wpi/first/wpilibj/examples/swervebot) example. I added a constants.py file to remove redundant imports, since the original Java example doesn't have a Constants.java. --------- Co-authored-by: Dustin Spicuzza <[email protected]>
1 parent 13bb632 commit d7df94c

File tree

4 files changed

+307
-0
lines changed

4 files changed

+307
-0
lines changed

‎SwerveBot/drivetrain.py

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
#
2+
# Copyright (c) FIRST and other WPILib contributors.
3+
# Open Source Software; you can modify and/or share it under the terms of
4+
# the WPILib BSD license file in the root directory of this project.
5+
#
6+
7+
import math
8+
import wpilib
9+
import wpimath.geometry
10+
import wpimath.kinematics
11+
import swervemodule
12+
13+
kMaxSpeed = 3.0 # 3 meters per second
14+
kMaxAngularSpeed = math.pi # 1/2 rotation per second
15+
16+
17+
class Drivetrain:
18+
"""
19+
Represents a swerve drive style drivetrain.
20+
"""
21+
22+
def __init__(self) -> None:
23+
self.frontLeftLocation = wpimath.geometry.Translation2d(0.381, 0.381)
24+
self.frontRightLocation = wpimath.geometry.Translation2d(0.381, -0.381)
25+
self.backLeftLocation = wpimath.geometry.Translation2d(-0.381, 0.381)
26+
self.backRightLocation = wpimath.geometry.Translation2d(-0.381, -0.381)
27+
28+
self.frontLeft = swervemodule.SwerveModule(1, 2, 0, 1, 2, 3)
29+
self.frontRight = swervemodule.SwerveModule(3, 4, 4, 5, 6, 7)
30+
self.backLeft = swervemodule.SwerveModule(5, 6, 8, 9, 10, 11)
31+
self.backRight = swervemodule.SwerveModule(7, 8, 12, 13, 14, 15)
32+
33+
self.gyro = wpilib.AnalogGyro(0)
34+
35+
self.kinematics = wpimath.kinematics.SwerveDrive4Kinematics(
36+
self.frontLeftLocation,
37+
self.frontRightLocation,
38+
self.backLeftLocation,
39+
self.backRightLocation,
40+
)
41+
42+
self.odometry = wpimath.kinematics.SwerveDrive4Odometry(
43+
self.kinematics,
44+
self.gyro.getRotation2d(),
45+
(
46+
self.frontLeft.getPosition(),
47+
self.frontRight.getPosition(),
48+
self.backLeft.getPosition(),
49+
self.backRight.getPosition(),
50+
),
51+
)
52+
53+
self.gyro.reset()
54+
55+
def drive(
56+
self,
57+
xSpeed: float,
58+
ySpeed: float,
59+
rot: float,
60+
fieldRelative: bool,
61+
periodSeconds: float,
62+
) -> None:
63+
"""
64+
Method to drive the robot using joystick info.
65+
:param xSpeed: Speed of the robot in the x direction (forward).
66+
:param ySpeed: Speed of the robot in the y direction (sideways).
67+
:param rot: Angular rate of the robot.
68+
:param fieldRelative: Whether the provided x and y speeds are relative to the field.
69+
:param periodSeconds: Time
70+
"""
71+
swerveModuleStates = self.kinematics.toSwerveModuleStates(
72+
wpimath.kinematics.ChassisSpeeds.discretize(
73+
wpimath.kinematics.ChassisSpeeds.fromFieldRelativeSpeeds(
74+
xSpeed, ySpeed, rot, self.gyro.getRotation2d()
75+
)
76+
if fieldRelative
77+
else wpimath.kinematics.ChassisSpeeds(xSpeed, ySpeed, rot),
78+
periodSeconds,
79+
)
80+
)
81+
wpimath.kinematics.SwerveDrive4Kinematics.desaturateWheelSpeeds(
82+
swerveModuleStates, kMaxSpeed
83+
)
84+
self.frontLeft.setDesiredState(swerveModuleStates[0])
85+
self.frontRight.setDesiredState(swerveModuleStates[1])
86+
self.backLeft.setDesiredState(swerveModuleStates[2])
87+
self.backRight.setDesiredState(swerveModuleStates[3])
88+
89+
def updateOdometry(self) -> None:
90+
"""Updates the field relative position of the robot."""
91+
self.odometry.update(
92+
self.gyro.getRotation2d(),
93+
(
94+
self.frontLeft.getPosition(),
95+
self.frontRight.getPosition(),
96+
self.backLeft.getPosition(),
97+
self.backRight.getPosition(),
98+
),
99+
)

‎SwerveBot/robot.py

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
#!/usr/bin/env python3
2+
#
3+
# Copyright (c) FIRST and other WPILib contributors.
4+
# Open Source Software; you can modify and/or share it under the terms of
5+
# the WPILib BSD license file in the root directory of this project.
6+
#
7+
8+
import wpilib
9+
import wpimath
10+
import wpilib.drive
11+
import wpimath.filter
12+
import wpimath.controller
13+
import drivetrain
14+
15+
16+
class MyRobot(wpilib.TimedRobot):
17+
def robotInit(self) -> None:
18+
"""Robot initialization function"""
19+
self.controller = wpilib.XboxController(0)
20+
self.swerve = drivetrain.Drivetrain()
21+
22+
# Slew rate limiters to make joystick inputs more gentle; 1/3 sec from 0 to 1.
23+
self.xspeedLimiter = wpimath.filter.SlewRateLimiter(3)
24+
self.yspeedLimiter = wpimath.filter.SlewRateLimiter(3)
25+
self.rotLimiter = wpimath.filter.SlewRateLimiter(3)
26+
27+
def autonomousPeriodic(self) -> None:
28+
self.driveWithJoystick(False)
29+
self.swerve.updateOdometry()
30+
31+
def teleopPeriodic(self) -> None:
32+
self.driveWithJoystick(True)
33+
34+
def driveWithJoystick(self, fieldRelative: bool) -> None:
35+
# Get the x speed. We are inverting this because Xbox controllers return
36+
# negative values when we push forward.
37+
xSpeed = (
38+
-self.xspeedLimiter.calculate(
39+
wpimath.applyDeadband(self.controller.getLeftY(), 0.02)
40+
)
41+
* drivetrain.kMaxSpeed
42+
)
43+
44+
# Get the y speed or sideways/strafe speed. We are inverting this because
45+
# we want a positive value when we pull to the left. Xbox controllers
46+
# return positive values when you pull to the right by default.
47+
ySpeed = (
48+
-self.yspeedLimiter.calculate(
49+
wpimath.applyDeadband(self.controller.getLeftX(), 0.02)
50+
)
51+
* drivetrain.kMaxSpeed
52+
)
53+
54+
# Get the rate of angular rotation. We are inverting this because we want a
55+
# positive value when we pull to the left (remember, CCW is positive in
56+
# mathematics). Xbox controllers return positive values when you pull to
57+
# the right by default.
58+
rot = (
59+
-self.rotLimiter.calculate(
60+
wpimath.applyDeadband(self.controller.getRightX(), 0.02)
61+
)
62+
* drivetrain.kMaxSpeed
63+
)
64+
65+
self.swerve.drive(xSpeed, ySpeed, rot, fieldRelative, self.getPeriod())
66+
67+
68+
if __name__ == "__main__":
69+
wpilib.run(MyRobot)

‎SwerveBot/swervemodule.py

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
#
2+
# Copyright (c) FIRST and other WPILib contributors.
3+
# Open Source Software; you can modify and/or share it under the terms of
4+
# the WPILib BSD license file in the root directory of this project.
5+
#
6+
7+
import math
8+
import wpilib
9+
import wpimath.kinematics
10+
import wpimath.geometry
11+
import wpimath.controller
12+
import wpimath.trajectory
13+
14+
kWheelRadius = 0.0508
15+
kEncoderResolution = 4096
16+
kModuleMaxAngularVelocity = math.pi
17+
kModuleMaxAngularAcceleration = math.tau
18+
19+
20+
class SwerveModule:
21+
def __init__(
22+
self,
23+
driveMotorChannel: int,
24+
turningMotorChannel: int,
25+
driveEncoderChannelA: int,
26+
driveEncoderChannelB: int,
27+
turningEncoderChannelA: int,
28+
turningEncoderChannelB: int,
29+
) -> None:
30+
"""Constructs a SwerveModule with a drive motor, turning motor, drive encoder and turning encoder.
31+
32+
:param driveMotorChannel: PWM output for the drive motor.
33+
:param turningMotorChannel: PWM output for the turning motor.
34+
:param driveEncoderChannelA: DIO input for the drive encoder channel A
35+
:param driveEncoderChannelB: DIO input for the drive encoder channel B
36+
:param turningEncoderChannelA: DIO input for the turning encoder channel A
37+
:param turningEncoderChannelB: DIO input for the turning encoder channel B
38+
"""
39+
self.driveMotor = wpilib.PWMSparkMax(driveMotorChannel)
40+
self.turningMotor = wpilib.PWMSparkMax(turningMotorChannel)
41+
42+
self.driveEncoder = wpilib.Encoder(driveEncoderChannelA, driveEncoderChannelB)
43+
self.turningEncoder = wpilib.Encoder(
44+
turningEncoderChannelA, turningEncoderChannelB
45+
)
46+
47+
# Gains are for example purposes only - must be determined for your own robot!
48+
self.drivePIDController = wpimath.controller.PIDController(1, 0, 0)
49+
50+
# Gains are for example purposes only - must be determined for your own robot!
51+
self.turningPIDController = wpimath.controller.ProfiledPIDController(
52+
1,
53+
0,
54+
0,
55+
wpimath.trajectory.TrapezoidProfile.Constraints(
56+
kModuleMaxAngularVelocity,
57+
kModuleMaxAngularAcceleration,
58+
),
59+
)
60+
61+
# Gains are for example purposes only - must be determined for your own robot!
62+
self.driveFeedforward = wpimath.controller.SimpleMotorFeedforwardMeters(1, 3)
63+
self.turnFeedforward = wpimath.controller.SimpleMotorFeedforwardMeters(1, 0.5)
64+
65+
# Set the distance per pulse for the drive encoder. We can simply use the
66+
# distance traveled for one rotation of the wheel divided by the encoder
67+
# resolution.
68+
self.driveEncoder.setDistancePerPulse(
69+
math.tau * kWheelRadius / kEncoderResolution
70+
)
71+
72+
# Set the distance (in this case, angle) in radians per pulse for the turning encoder.
73+
# This is the the angle through an entire rotation (2 * pi) divided by the
74+
# encoder resolution.
75+
self.turningEncoder.setDistancePerPulse(math.tau / kEncoderResolution)
76+
77+
# Limit the PID Controller's input range between -pi and pi and set the input
78+
# to be continuous.
79+
self.turningPIDController.enableContinuousInput(-math.pi, math.pi)
80+
81+
def getState(self) -> wpimath.kinematics.SwerveModuleState:
82+
"""Returns the current state of the module.
83+
84+
:returns: The current state of the module.
85+
"""
86+
return wpimath.kinematics.SwerveModuleState(
87+
self.driveEncoder.getRate(),
88+
wpimath.geometry.Rotation2d(self.turningEncoder.getDistance()),
89+
)
90+
91+
def getPosition(self) -> wpimath.kinematics.SwerveModulePosition:
92+
"""Returns the current position of the module.
93+
94+
:returns: The current position of the module.
95+
"""
96+
return wpimath.kinematics.SwerveModulePosition(
97+
self.driveEncoder.getRate(),
98+
wpimath.geometry.Rotation2d(self.turningEncoder.getDistance()),
99+
)
100+
101+
def setDesiredState(
102+
self, desiredState: wpimath.kinematics.SwerveModuleState
103+
) -> None:
104+
"""Sets the desired state for the module.
105+
106+
:param desiredState: Desired state with speed and angle.
107+
"""
108+
109+
encoderRotation = wpimath.geometry.Rotation2d(self.turningEncoder.getDistance())
110+
111+
# Optimize the reference state to avoid spinning further than 90 degrees
112+
state = wpimath.kinematics.SwerveModuleState.optimize(
113+
desiredState, encoderRotation
114+
)
115+
116+
# Scale speed by cosine of angle error. This scales down movement perpendicular to the desired
117+
# direction of travel that can occur when modules change directions. This results in smoother
118+
# driving.
119+
state.speed *= (state.angle - encoderRotation).cos()
120+
121+
# Calculate the drive output from the drive PID controller.
122+
driveOutput = self.drivePIDController.calculate(
123+
self.driveEncoder.getRate(), state.speed
124+
)
125+
126+
driveFeedforward = self.driveFeedforward.calculate(state.speed)
127+
128+
# Calculate the turning motor output from the turning PID controller.
129+
turnOutput = self.turningPIDController.calculate(
130+
self.turningEncoder.getDistance(), state.angle.radians()
131+
)
132+
133+
turnFeedforward = self.turnFeedforward.calculate(
134+
self.turningPIDController.getSetpoint().velocity
135+
)
136+
137+
self.driveMotor.setVoltage(driveOutput + driveFeedforward)
138+
self.turningMotor.setVoltage(turnOutput + turnFeedforward)

‎run_tests.sh

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ BASE_TESTS="
4545
Solenoid
4646
StatefulAutonomous
4747
StateSpaceFlywheel
48+
SwerveBot
4849
TankDrive
4950
TankDriveXboxController
5051
Timed/src

0 commit comments

Comments
 (0)