Adds MecanumControllerCommand to Commands2.

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",

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())

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,
+        )