diff --git a/commands2/__init__.py b/commands2/__init__.py index cc8e6a4c..a45bd3d0 100644 --- a/commands2/__init__.py +++ b/commands2/__init__.py @@ -25,6 +25,7 @@ from .conditionalcommand import ConditionalCommand from .functionalcommand import FunctionalCommand from .instantcommand import InstantCommand +from .mecanumcontrollercommand import MecanumControllerCommand from .notifiercommand import NotifierCommand from .parallelcommandgroup import ParallelCommandGroup from .paralleldeadlinegroup import ParallelDeadlineGroup @@ -56,6 +57,7 @@ "IllegalCommandUse", "InstantCommand", "InterruptionBehavior", + "MecanumControllerCommand", "NotifierCommand", "ParallelCommandGroup", "ParallelDeadlineGroup", diff --git a/commands2/mecanumcontrollercommand.py b/commands2/mecanumcontrollercommand.py new file mode 100644 index 00000000..c2a721f7 --- /dev/null +++ b/commands2/mecanumcontrollercommand.py @@ -0,0 +1,222 @@ +# 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. +from __future__ import annotations +from typing import Set, Callable, Union, List + +from wpimath.controller import ( + HolonomicDriveController, + PIDController, + ProfiledPIDControllerRadians, + SimpleMotorFeedforwardMeters, +) +from wpimath.geometry import Pose2d, Rotation2d +from wpimath.kinematics import ( + ChassisSpeeds, + MecanumDriveKinematics, + MecanumDriveWheelSpeeds, +) +from wpimath.trajectory import Trajectory +from wpilib import Timer + +from .subsystem import Subsystem +from .command import Command + + +class MecanumControllerCommand(Command): + """ + A command that uses two PID controllers (PIDController) and a ProfiledPIDController + (ProfiledPIDController) to follow a trajectory Trajectory with a mecanum drive. + + The command handles trajectory-following, Velocity PID calculations, and feedforwards + internally. This is intended to be a more-or-less "complete solution" that can be used by teams + without a great deal of controls expertise. + + Advanced teams seeking more flexibility (for example, those who wish to use the onboard PID + functionality of a "smart" motor controller) may use the secondary constructor that omits the PID + and feedforward functionality, returning only the raw wheel speeds from the PID controllers. + + The robot angle controller does not follow the angle given by the trajectory but rather goes + to the angle given in the final state of the trajectory. + """ + + __FRONT_LEFT_INDEX = 0 + __REAR_LEFT_INDEX = 1 + __FRONT_RIGHT_INDEX = 2 + __REAR_RIGHT_INDEX = 3 + + def __init__( + self, + trajectory: Trajectory, + pose: Callable[[], Pose2d], + kinematics: MecanumDriveKinematics, + xController: PIDController, + yController: PIDController, + thetaController: ProfiledPIDControllerRadians, + # rotationSupplier: Callable[[], Rotation2d], + maxWheelVelocityMetersPerSecond: float, + outputConsumer: Callable[[Union[List[float], MecanumDriveWheelSpeeds]], None], + *requirements: Subsystem, + feedforward: SimpleMotorFeedforwardMeters | None = None, + chassisWheelPIDControllers: List[PIDController] | None = None, + currentWheelSpeedsSupplier: Callable[[], MecanumDriveWheelSpeeds] | None = None, + ): + super().__init__() + + self.trajectory: Trajectory = trajectory + self.pose = pose + self.kinematics = kinematics + self.controller = HolonomicDriveController( + xController, yController, thetaController + ) + # self.rotationSupplier = rotationSupplier + self.maxWheelVelocityMetersPerSecond = maxWheelVelocityMetersPerSecond + self.outputConsumer = outputConsumer + + # Member to track time of last loop execution + self.prevTime = 0 + self.usePID = False + + # Optional requirements checks. If any one of the feedforward, pidcontrollers, or wheelspeedsuppliers + # are not None, then all should be non-null. Raise RuntimeError if they are. The List of PID controllers + # should be equal to each corner of the robot being commanded. + if ( + feedforward + or chassisWheelPIDControllers + or currentWheelSpeedsSupplier is not None + ): + if ( + feedforward + or chassisWheelPIDControllers + or currentWheelSpeedsSupplier is None + ): + raise RuntimeError( + f"Failed to instantiate MecanumControllerCommand, one of the arguments passed in was None: \n \ + \t Feedforward: {feedforward} - chassisWheelPIDControllers: {chassisWheelPIDControllers} - wheelSpeedSupplier: {currentWheelSpeedsSupplier} " + ) + + # check the length of the PID controllers + if len(chassisWheelPIDControllers != 4): + raise RuntimeError( + f"Failed to instantiate MecanumControllerCommand, not enough PID controllers provided.\n \ + \t Required: 4 - Provided: {len(chassisWheelPIDControllers)}" + ) + + self.frontLeftController = chassisWheelPIDControllers[ + self.__FRONT_LEFT_INDEX + ] + self.rearLeftController = chassisWheelPIDControllers[self.__REAR_LEFT_INDEX] + self.frontRightController = chassisWheelPIDControllers[ + self.__FRONT_RIGHT_INDEX + ] + self.rearRightController = chassisWheelPIDControllers[ + self.__REAR_RIGHT_INDEX + ] + + if currentWheelSpeedsSupplier is None: + raise RuntimeError( + f"Failed to instantiate MecanumControllerCommand, no WheelSpeedSupplierProvided." + ) + + self.currentWheelSpeeds = currentWheelSpeedsSupplier + + if feedforward is None: + raise RuntimeError( + f"Failed to instantiate MecanumControllerCommand, no SimpleMotorFeedforwardMeters supplied." + ) + + self.feedforward: SimpleMotorFeedforwardMeters = feedforward + + # All the optional requirements verify, set usePid flag to True + self.usePID = True + + # Set the requirements for the command + self.addRequirements(*requirements) + + self.timer = Timer() + + def initialize(self): + initialState: Trajectory.State = self.trajectory.sample(0) + initialXVelocity = initialState.velocity * initialState.pose.rotation().cos() + initialYVelocity = initialState.velocity * initialState.pose.rotation().sin() + self.m_prevSpeeds = self.kinematics.toWheelSpeeds( + ChassisSpeeds(initialXVelocity, initialYVelocity, 0.0) + ) + self.timer.restart() + self.prevTime = 0 + + def execute(self): + curTime = self.timer.get() + dt = curTime - self.prevTime + desiredState: Trajectory.State = self.trajectory.sample(curTime) + targetChassisSpeeds = self.controller.calculate( + self.pose(), + desiredState, + desiredState.pose.rotation() + # self.pose.get(), desiredState, self.rotationSupplier.get() + ) + targetWheelSpeeds = self.kinematics.toWheelSpeeds(targetChassisSpeeds) + targetWheelSpeeds.desaturate(self.maxWheelVelocityMetersPerSecond) + frontLeftSpeedSetpoint = targetWheelSpeeds.frontLeft + rearLeftSpeedSetpoint = targetWheelSpeeds.rearLeft + frontRightSpeedSetpoint = targetWheelSpeeds.frontRight + rearRightSpeedSetpoint = targetWheelSpeeds.rearRight + + if not self.usePID: + self.outputConsumer( + MecanumDriveWheelSpeeds( + frontLeftSpeedSetpoint, + frontRightSpeedSetpoint, + rearLeftSpeedSetpoint, + rearRightSpeedSetpoint, + ) + ) + else: + frontLeftFeedforward = self.feedforward.calculate( + frontLeftSpeedSetpoint, + (frontLeftSpeedSetpoint - self.m_prevSpeeds.frontLeft) / dt, + ) + rearLeftFeedforward = self.feedforward.calculate( + rearLeftSpeedSetpoint, + (rearLeftSpeedSetpoint - self.m_prevSpeeds.rearLeft) / dt, + ) + frontRightFeedforward = self.feedforward.calculate( + frontRightSpeedSetpoint, + (frontRightSpeedSetpoint - self.m_prevSpeeds.frontRight) / dt, + ) + rearRightFeedforward = self.feedforward.calculate( + rearRightSpeedSetpoint, + (rearRightSpeedSetpoint - self.m_prevSpeeds.rearRight) / dt, + ) + frontLeftOutput = frontLeftFeedforward + self.frontLeftController.calculate( + self.currentWheelSpeeds().frontLeft, + frontLeftSpeedSetpoint, + ) + rearLeftOutput = rearLeftFeedforward + self.rearLeftController.calculate( + self.currentWheelSpeeds().rearLeft, + rearLeftSpeedSetpoint, + ) + frontRightOutput = ( + frontRightFeedforward + + self.frontRightController.calculate( + self.currentWheelSpeeds().frontRight, + frontRightSpeedSetpoint, + ) + ) + rearRightOutput = rearRightFeedforward + self.rearRightController.calculate( + self.currentWheelSpeeds().rearRight, + rearRightSpeedSetpoint, + ) + self.outputConsumer( + frontLeftOutput, frontRightOutput, rearLeftOutput, rearRightOutput + ) + + # Store current call information for next call + self.prevTime = curTime + self.m_prevSpeeds = targetWheelSpeeds + + def end(self, interrupted): + self.timer.stop() + + def isFinished(self): + return self.timer.hasElapsed(self.trajectory.totalTime()) diff --git a/tests/test_mecanumcontrollercommand.py b/tests/test_mecanumcontrollercommand.py new file mode 100644 index 00000000..d1eaefdb --- /dev/null +++ b/tests/test_mecanumcontrollercommand.py @@ -0,0 +1,156 @@ +# 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. + +from typing import TYPE_CHECKING, List +import math + +import wpimath.controller as controller +import wpimath.trajectory as trajectory +import wpimath.geometry as geometry +import wpimath.kinematics as kinematics +from wpilib import Timer + +from util import * # type: ignore + +if TYPE_CHECKING: + from .util import * + +import pytest + +import commands2 + + +class MecanumControllerCommandTestDataFixtures: + def __init__(self): + self.timer = Timer() + self.angle: geometry.Rotation2d = geometry.Rotation2d(0) + + # Track speeds and distances + self.frontLeftSpeed = 0 + self.frontLeftDistance = 0 + self.rearLeftSpeed = 0 + self.rearLeftDistance = 0 + self.frontRightSpeed = 0 + self.frontRightDistance = 0 + self.rearRightSpeed = 0 + self.rearRightDistance = 0 + + # Profile Controller and constraints + trapProfileConstraints: trajectory.TrapezoidProfileRadians.Constraints = ( + trajectory.TrapezoidProfileRadians.Constraints(3 * math.pi, math.pi) + ) + self.rotationController: controller.ProfiledPIDControllerRadians = ( + controller.ProfiledPIDControllerRadians( + 1.0, 0.0, 0.0, trapProfileConstraints + ) + ) + + # Chassis/Drivetrain constants + self.kxTolerance = 1 / 12.0 + self.kyTolerance = 1 / 12.0 + self.kAngularTolerance = 1 / 12.0 + self.kWheelBase = 0.5 + self.kTrackWidth = 0.5 + + self.command_kinematics: kinematics.MecanumDriveKinematics = ( + kinematics.MecanumDriveKinematics( + geometry.Translation2d(self.kWheelBase / 2, self.kTrackWidth / 2), + geometry.Translation2d(self.kWheelBase / 2, -self.kTrackWidth / 2), + geometry.Translation2d(self.kWheelBase / 2, self.kTrackWidth / 2), + geometry.Translation2d(self.kWheelBase / 2, -self.kTrackWidth / 2), + ) + ) + + self.command_odometry: kinematics.MecanumDriveOdometry = ( + kinematics.MecanumDriveOdometry( + self.command_kinematics, + geometry.Rotation2d(0), + kinematics.MecanumDriveWheelPositions(), + geometry.Pose2d(0, 0, geometry.Rotation2d(0)), + ) + ) + + def setWheelSpeeds(self, wheelSpeeds: kinematics.MecanumDriveWheelSpeeds) -> None: + self.frontLeftSpeed = wheelSpeeds.frontLeft + self.rearLeftSpeed = wheelSpeeds.rearLeft + self.frontRightSpeed = wheelSpeeds.frontRight + self.rearRightSpeed = wheelSpeeds.rearRight + + def getCurrentWheelDistances(self) -> kinematics.MecanumDriveWheelPositions: + positions = kinematics.MecanumDriveWheelPositions() + positions.frontLeft = self.frontLeftDistance + positions.frontRight = self.frontRightDistance + positions.rearLeft = self.rearLeftDistance + positions.rearRight = self.rearRightDistance + + return positions + + def getRobotPose(self) -> geometry.Pose2d: + self.command_odometry.update(self.angle, self.getCurrentWheelDistances()) + return self.command_odometry.getPose() + + +@pytest.fixture() +def get_mec_controller_data() -> MecanumControllerCommandTestDataFixtures: + return MecanumControllerCommandTestDataFixtures() + + +def test_mecanumControllerCommand( + scheduler: commands2.CommandScheduler, get_mec_controller_data +): + with ManualSimTime() as sim: + subsystem = commands2.Subsystem() + waypoints: List[geometry.Pose2d] = [] + waypoints.append(geometry.Pose2d(0, 0, geometry.Rotation2d(0))) + waypoints.append(geometry.Pose2d(1, 5, geometry.Rotation2d(3))) + traj_config: trajectory.TrajectoryConfig = trajectory.TrajectoryConfig(8.8, 0.1) + new_trajectory: trajectory.Trajectory = ( + trajectory.TrajectoryGenerator.generateTrajectory(waypoints, traj_config) + ) + end_state = new_trajectory.sample(new_trajectory.totalTime()) + + fixture_data = get_mec_controller_data + + mecContCommand = commands2.MecanumControllerCommand( + new_trajectory, + fixture_data.getRobotPose, + fixture_data.command_kinematics, + controller.PIDController(0.6, 0, 0), + controller.PIDController(0.6, 0, 0), + fixture_data.rotationController, + 8.8, + fixture_data.setWheelSpeeds, + subsystem, + ) + + fixture_data.timer.restart() + + mecContCommand.initialize() + + while not mecContCommand.isFinished(): + mecContCommand.execute() + fixture_data.angle = new_trajectory.sample( + fixture_data.timer.get() + ).pose.rotation() + + fixture_data.frontLeftDistance += fixture_data.frontLeftSpeed * 0.005 + fixture_data.rearLeftDistance += fixture_data.rearLeftSpeed * 0.005 + fixture_data.frontRightDistance += fixture_data.frontRightSpeed * 0.005 + fixture_data.rearRightDistance += fixture_data.rearRightSpeed * 0.005 + + sim.step(0.005) + + fixture_data.timer.stop() + mecContCommand.end(True) + + assert end_state.pose.X() == pytest.approx( + fixture_data.getRobotPose().X(), fixture_data.kxTolerance + ) + assert end_state.pose.Y() == pytest.approx( + fixture_data.getRobotPose().Y(), fixture_data.kyTolerance + ) + assert end_state.pose.rotation().radians() == pytest.approx( + fixture_data.getRobotPose().rotation().radians(), + fixture_data.kAngularTolerance, + )