mirror of
https://github.com/wpilibsuite/allwpilib
synced 2026-06-19 00:41:43 +00:00
dca59147e1
This also updates the bazel scripts to behave more like the C++ and Java examples, and updates the copybara scripts to be able to sync up `mostrobotpy`
130 lines
5.3 KiB
Python
Executable File
130 lines
5.3 KiB
Python
Executable File
#
|
|
# 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
|
|
|
|
|
|
class Drivetrain:
|
|
"""Represents a mecanum drive style drivetrain."""
|
|
|
|
MAX_VELOCITY = 3.0 # 3 meters per second
|
|
MAX_ANGULAR_VELOCITY = math.pi # 1/2 rotation per second
|
|
|
|
def __init__(self) -> None:
|
|
self.frontLeftMotor = wpilib.PWMSparkMax(1)
|
|
self.frontRightMotor = wpilib.PWMSparkMax(2)
|
|
self.backLeftMotor = wpilib.PWMSparkMax(3)
|
|
self.backRightMotor = wpilib.PWMSparkMax(4)
|
|
|
|
self.frontLeftEncoder = wpilib.Encoder(0, 1)
|
|
self.frontRightEncoder = wpilib.Encoder(2, 3)
|
|
self.backLeftEncoder = wpilib.Encoder(4, 5)
|
|
self.backRightEncoder = wpilib.Encoder(6, 7)
|
|
|
|
self.frontLeftLocation = wpimath.Translation2d(0.381, 0.381)
|
|
self.frontRightLocation = wpimath.Translation2d(0.381, -0.381)
|
|
self.backLeftLocation = wpimath.Translation2d(-0.381, 0.381)
|
|
self.backRightLocation = wpimath.Translation2d(-0.381, -0.381)
|
|
|
|
self.frontLeftPIDController = wpimath.PIDController(1, 0, 0)
|
|
self.frontRightPIDController = wpimath.PIDController(1, 0, 0)
|
|
self.backLeftPIDController = wpimath.PIDController(1, 0, 0)
|
|
self.backRightPIDController = wpimath.PIDController(1, 0, 0)
|
|
|
|
self.imu = wpilib.OnboardIMU(wpilib.OnboardIMU.MountOrientation.FLAT)
|
|
|
|
self.kinematics = wpimath.MecanumDriveKinematics(
|
|
self.frontLeftLocation,
|
|
self.frontRightLocation,
|
|
self.backLeftLocation,
|
|
self.backRightLocation,
|
|
)
|
|
|
|
self.odometry = wpimath.MecanumDriveOdometry(
|
|
self.kinematics, self.imu.getRotation2d(), self.getCurrentDistances()
|
|
)
|
|
|
|
# Gains are for example purposes only - must be determined for your own robot!
|
|
self.feedforward = wpimath.SimpleMotorFeedforwardMeters(1, 3)
|
|
|
|
self.imu.resetYaw()
|
|
|
|
# We need to invert one side of the drivetrain so that positive voltages
|
|
# result in both sides moving forward. Depending on how your robot's
|
|
# gearbox is constructed, you might have to invert the left side instead.
|
|
self.frontRightMotor.setInverted(True)
|
|
self.backRightMotor.setInverted(True)
|
|
|
|
def getCurrentState(self) -> wpimath.MecanumDriveWheelVelocities:
|
|
"""Returns the current state of the drivetrain."""
|
|
return wpimath.MecanumDriveWheelVelocities(
|
|
self.frontLeftEncoder.getRate(),
|
|
self.frontRightEncoder.getRate(),
|
|
self.backLeftEncoder.getRate(),
|
|
self.backRightEncoder.getRate(),
|
|
)
|
|
|
|
def getCurrentDistances(self) -> wpimath.MecanumDriveWheelPositions:
|
|
"""Returns the current distances measured by the drivetrain."""
|
|
positions = wpimath.MecanumDriveWheelPositions()
|
|
positions.frontLeft = self.frontLeftEncoder.getDistance()
|
|
positions.frontRight = self.frontRightEncoder.getDistance()
|
|
positions.rearLeft = self.backLeftEncoder.getDistance()
|
|
positions.rearRight = self.backRightEncoder.getDistance()
|
|
return positions
|
|
|
|
def setVelocities(self, velocities: wpimath.MecanumDriveWheelVelocities) -> None:
|
|
"""Sets the desired velocities for each wheel."""
|
|
frontLeftFeedforward = self.feedforward.calculate(velocities.frontLeft)
|
|
frontRightFeedforward = self.feedforward.calculate(velocities.frontRight)
|
|
backLeftFeedforward = self.feedforward.calculate(velocities.rearLeft)
|
|
backRightFeedforward = self.feedforward.calculate(velocities.rearRight)
|
|
|
|
frontLeftOutput = self.frontLeftPIDController.calculate(
|
|
self.frontLeftEncoder.getRate(), velocities.frontLeft
|
|
)
|
|
frontRightOutput = self.frontRightPIDController.calculate(
|
|
self.frontRightEncoder.getRate(), velocities.frontRight
|
|
)
|
|
backLeftOutput = self.frontLeftPIDController.calculate(
|
|
self.backLeftEncoder.getRate(), velocities.rearLeft
|
|
)
|
|
backRightOutput = self.frontRightPIDController.calculate(
|
|
self.backRightEncoder.getRate(), velocities.rearRight
|
|
)
|
|
|
|
self.frontLeftMotor.setVoltage(frontLeftOutput + frontLeftFeedforward)
|
|
self.frontRightMotor.setVoltage(frontRightOutput + frontRightFeedforward)
|
|
self.backLeftMotor.setVoltage(backLeftOutput + backLeftFeedforward)
|
|
self.backRightMotor.setVoltage(backRightOutput + backRightFeedforward)
|
|
|
|
def drive(
|
|
self,
|
|
xVelocity: float,
|
|
yVelocity: float,
|
|
rot: float,
|
|
fieldRelative: bool,
|
|
periodSeconds: float,
|
|
) -> None:
|
|
"""Method to drive the robot using joystick info."""
|
|
chassisVelocities = wpimath.ChassisVelocities(xVelocity, yVelocity, rot)
|
|
if fieldRelative:
|
|
chassisVelocities = chassisVelocities.toRobotRelative(
|
|
self.imu.getRotation2d()
|
|
)
|
|
|
|
self.setVelocities(
|
|
self.kinematics.toWheelVelocities(
|
|
chassisVelocities.discretize(periodSeconds)
|
|
).desaturate(self.MAX_VELOCITY)
|
|
)
|
|
|
|
def updateOdometry(self) -> None:
|
|
"""Updates the field-relative position."""
|
|
self.odometry.update(self.imu.getRotation2d(), self.getCurrentDistances())
|