|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +import wpilib.drive |
| 4 | +import wpimath.controller |
| 5 | +import wpimath.kinematics |
| 6 | +import wpimath.units |
| 7 | + |
| 8 | +import math |
| 9 | + |
| 10 | + |
| 11 | +class Drivetrain: |
| 12 | + """Represents a differential drive style drivetrain.""" |
| 13 | + |
| 14 | + MAX_SPEED = 3.0 # meters per second |
| 15 | + MAX_ANGULAR_SPEED = 2 * math.pi # one rotation per second |
| 16 | + |
| 17 | + TRACK_WIDTH = 0.381 * 2 # meters |
| 18 | + WHEEL_RADIUS = 0.0508 # meters |
| 19 | + ENCODER_RESOLUTION = 4096 # counts per revolution |
| 20 | + |
| 21 | + def __init__(self): |
| 22 | + leftLeader = wpilib.PWMSparkMax(1) |
| 23 | + leftFollower = wpilib.PWMSparkMax(2) |
| 24 | + rightLeader = wpilib.PWMSparkMax(3) |
| 25 | + rightFollower = wpilib.PWMSparkMax(4) |
| 26 | + |
| 27 | + self.leftEncoder = wpilib.Encoder(0, 1) |
| 28 | + self.rightEncoder = wpilib.Encoder(2, 3) |
| 29 | + |
| 30 | + self.leftGroup = wpilib.MotorControllerGroup(leftLeader, leftFollower) |
| 31 | + self.rightGroup = wpilib.MotorControllerGroup(rightLeader, rightFollower) |
| 32 | + |
| 33 | + self.gyro = wpilib.AnalogGyro(0) |
| 34 | + |
| 35 | + self.leftPIDController = wpimath.controller.PIDController(1.0, 0.0, 0.0) |
| 36 | + self.rightPIDController = wpimath.controller.PIDController(1.0, 0.0, 0.0) |
| 37 | + |
| 38 | + self.kinematics = wpimath.kinematics.DifferentialDriveKinematics( |
| 39 | + self.TRACK_WIDTH |
| 40 | + ) |
| 41 | + |
| 42 | + # Gains are for example purposes only - must be determined for your own robot! |
| 43 | + self.feedforward = wpimath.controller.SimpleMotorFeedforwardMeters(1, 3) |
| 44 | + |
| 45 | + self.gyro.reset() |
| 46 | + |
| 47 | + # We need to invert one side of the drivetrain so that positive voltages |
| 48 | + # result in both sides moving forward. Depending on how your robot's |
| 49 | + # gearbox is constructed, you might have to invert the left side instead. |
| 50 | + self.rightGroup.setInverted(True) |
| 51 | + |
| 52 | + # Set the distance per pulse for the drive encoders. We can simply use the |
| 53 | + # distance traveled for one rotation of the wheel divided by the encoder |
| 54 | + # resolution. |
| 55 | + self.leftEncoder.setDistancePerPulse( |
| 56 | + 2 * math.pi * self.WHEEL_RADIUS / self.ENCODER_RESOLUTION |
| 57 | + ) |
| 58 | + self.rightEncoder.setDistancePerPulse( |
| 59 | + 2 * math.pi * self.WHEEL_RADIUS / self.ENCODER_RESOLUTION |
| 60 | + ) |
| 61 | + |
| 62 | + self.leftEncoder.reset() |
| 63 | + self.rightEncoder.reset() |
| 64 | + |
| 65 | + self.odometry = wpimath.kinematics.DifferentialDriveOdometry( |
| 66 | + self.gyro.getRotation2d(), |
| 67 | + self.leftEncoder.getDistance(), |
| 68 | + self.rightEncoder.getDistance(), |
| 69 | + ) |
| 70 | + |
| 71 | + def setSpeeds(self, speeds: wpimath.kinematics.DifferentialDriveWheelSpeeds): |
| 72 | + """Sets the desired wheel speeds.""" |
| 73 | + leftFeedforward = self.feedforward.calculate(speeds.left) |
| 74 | + rightFeedforward = self.feedforward.calculate(speeds.right) |
| 75 | + |
| 76 | + leftOutput = self.leftPIDController.calculate( |
| 77 | + self.leftEncoder.getRate(), speeds.left |
| 78 | + ) |
| 79 | + rightOutput = self.rightPIDController.calculate( |
| 80 | + self.rightEncoder.getRate(), speeds.right |
| 81 | + ) |
| 82 | + |
| 83 | + self.leftGroup.setVoltage(leftOutput + leftFeedforward) |
| 84 | + self.rightGroup.setVoltage(rightOutput + rightFeedforward) |
| 85 | + |
| 86 | + def drive(self, xSpeed, rot): |
| 87 | + """Drives the robot with the given linear velocity and angular velocity.""" |
| 88 | + wheelSpeeds = self.kinematics.toWheelSpeeds( |
| 89 | + wpimath.kinematics.ChassisSpeeds(xSpeed, 0, rot) |
| 90 | + ) |
| 91 | + self.setSpeeds(wheelSpeeds) |
| 92 | + |
| 93 | + def updateOdometry(self): |
| 94 | + """Updates the field-relative position.""" |
| 95 | + self.odometry.update( |
| 96 | + self.gyro.getRotation2d(), |
| 97 | + self.leftEncoder.getDistance(), |
| 98 | + self.rightEncoder.getDistance(), |
| 99 | + ) |
0 commit comments