Adds basic SwerveControllerCommand (#45)

* Adds Swerve2/3/4/6ControllerCommand to Commands2

* Fixes for mypy errors

* Fixes broken test run for incomplete pid

* Experimenting with mypy assignment error

* Ignoring mypy type error with explicit annotation

* Re-formatted with black

* Added Typing.overload to constructor definitions, forced kwargs for optional contructor arguments.

* Removed superfluous comment

* Updated implementation to use Overtake and Beartype.

* Fixed test fixtures to supply tuple to odometry update

* Fixed tests for odometry update calls to send Tuples

* Fix argument

* Bug fix: Added addRequirements call to contstructors

* Pinned versions of experimental libraries

* Add Simple SwerveControllerCommands

* black

---------

Co-authored-by: NewtonCrosby <lospugs@>
Co-authored-by: lospugs <[email protected]>
Co-authored-by: Dustin Spicuzza <[email protected]>

commands2/__init__.py

@@ -42,6 +42,7 @@
 from .sequentialcommandgroup import SequentialCommandGroup
 from .startendcommand import StartEndCommand
 from .subsystem import Subsystem
+from .swervecontrollercommand import SwerveControllerCommand
 from .timedcommandrobot import TimedCommandRobot
 from .trapezoidprofilesubsystem import TrapezoidProfileSubsystem
 from .waitcommand import WaitCommand
@@ -77,6 +78,7 @@
     "SequentialCommandGroup",
     "StartEndCommand",
     "Subsystem",
+    "SwerveControllerCommand",
     "TimedCommandRobot",
     "TrapezoidProfileSubsystem",
     "WaitCommand",

commands2/swervecontrollercommand.py

@@ -0,0 +1,109 @@
+# 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 Callable, Optional, Union, Tuple, Sequence
+from typing_extensions import overload
+from wpimath.controller import (
+    HolonomicDriveController,
+)
+from wpimath.geometry import Pose2d, Rotation2d
+from wpimath.kinematics import (
+    SwerveDrive2Kinematics,
+    SwerveDrive3Kinematics,
+    SwerveDrive4Kinematics,
+    SwerveDrive6Kinematics,
+    SwerveModuleState,
+)
+from wpimath.trajectory import Trajectory
+from wpilib import Timer
+
+from .command import Command
+from .subsystem import Subsystem
+
+
+class SwerveControllerCommand(Command):
+    """
+    A command that uses two PID controllers (PIDController) and a ProfiledPIDController
+    (ProfiledPIDController) to follow a trajectory (Trajectory) with a swerve drive.
+
+    This command outputs the raw desired Swerve Module States (SwerveModuleState) in an
+    array. The desired wheel and module rotation velocities should be taken from those and used in
+    velocity PIDs.
+
+    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.
+    """
+
+    def __init__(
+        self,
+        trajectory: Trajectory,
+        pose: Callable[[], Pose2d],
+        kinematics: Union[
+            SwerveDrive2Kinematics,
+            SwerveDrive3Kinematics,
+            SwerveDrive4Kinematics,
+            SwerveDrive6Kinematics,
+        ],
+        controller: HolonomicDriveController,
+        outputModuleStates: Callable[[Sequence[SwerveModuleState]], None],
+        requirements: Tuple[Subsystem],
+        desiredRotation: Optional[Callable[[], Rotation2d]] = None,
+    ) -> None:
+        """
+        Constructs a new SwerveControllerCommand that when executed will follow the
+        provided trajectory. This command will not return output voltages but
+        rather raw module states from the position controllers which need to be put
+        into a velocity PID.
+
+        Note: The controllers will *not* set the outputVolts to zero upon
+        completion of the path- this is left to the user, since it is not
+        appropriate for paths with nonstationary endstates.
+
+        :param trajectory:         The trajectory to follow.
+        :param pose:               A function that supplies the robot pose - use one of the odometry classes to
+                                   provide this.
+        :param kinematics:         The kinematics for the robot drivetrain. Can be kinematics for 2/3/4/6
+                                   SwerveKinematics.
+        :param controller:         The HolonomicDriveController for the drivetrain.
+                                   If you have x, y, and theta controllers, pass them into
+                                   HolonomicPIDController.
+        :param outputModuleStates: The raw output module states from the position controllers.
+        :param requirements:       The subsystems to require.
+        :param desiredRotation:    (optional) The angle that the drivetrain should be
+                                   facing. This is sampled at each time step. If not specified, that rotation of
+                                   the final pose in the trajectory is used.
+        """
+        super().__init__()
+        self._trajectory = trajectory
+        self._pose = pose
+        self._kinematics = kinematics
+        self._outputModuleStates = outputModuleStates
+        self._controller = controller
+        if desiredRotation is None:
+            self._desiredRotation = trajectory.states()[-1].pose.rotation
+        else:
+            self._desiredRotation = desiredRotation
+
+        self._timer = Timer()
+        self.addRequirements(*requirements)
+
+    def initialize(self):
+        self._timer.restart()
+
+    def execute(self):
+        curTime = self._timer.get()
+        desiredState = self._trajectory.sample(curTime)
+
+        targetChassisSpeeds = self._controller.calculate(
+            self._pose(), desiredState, self._desiredRotation()
+        )
+        targetModuleStates = self._kinematics.toSwerveModuleStates(targetChassisSpeeds)
+
+        self._outputModuleStates(targetModuleStates)
+
+    def end(self, interrupted):
+        self._timer.stop()
+
+    def isFinished(self):
+        return self._timer.hasElapsed(self._trajectory.totalTime())

setup.py

@@ -11,7 +11,10 @@
     description="WPILib command framework v2",
     url="https://github.com/robotpy/robotpy-commands-v2",
     packages=["commands2"],
-    install_requires=["wpilib<2025,>=2024.0.0b2", "typing_extensions>=4.1.0,<5"],
+    install_requires=[
+        "wpilib<2025,>=2024.0.0b2",
+        "typing_extensions>=4.1.0,<5",
+    ],
     license="BSD-3-Clause",
     python_requires=">=3.8",
     include_package_data=True,

tests/test_swervecontrollercommand.py

@@ -0,0 +1,142 @@
+# 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, Tuple
+import math
+
+from wpilib import Timer
+
+from wpimath.geometry import Pose2d, Rotation2d, Translation2d
+from wpimath.kinematics import (
+    SwerveModuleState,
+    SwerveModulePosition,
+    SwerveDrive4Kinematics,
+    SwerveDrive4Odometry,
+)
+from wpimath.controller import (
+    ProfiledPIDControllerRadians,
+    PIDController,
+    HolonomicDriveController,
+)
+from wpimath.trajectory import (
+    TrapezoidProfileRadians,
+    Trajectory,
+    TrajectoryConfig,
+    TrajectoryGenerator,
+)
+
+
+from util import *  # type: ignore
+
+if TYPE_CHECKING:
+    from .util import *
+
+import pytest
+
+import commands2
+
+
+def test_swervecontrollercommand():
+    timer = Timer()
+    angle = Rotation2d(0)
+
+    swerve_module_states = (
+        SwerveModuleState(0, Rotation2d(0)),
+        SwerveModuleState(0, Rotation2d(0)),
+        SwerveModuleState(0, Rotation2d(0)),
+        SwerveModuleState(0, Rotation2d(0)),
+    )
+
+    swerve_module_positions = (
+        SwerveModulePosition(0, Rotation2d(0)),
+        SwerveModulePosition(0, Rotation2d(0)),
+        SwerveModulePosition(0, Rotation2d(0)),
+        SwerveModulePosition(0, Rotation2d(0)),
+    )
+
+    rot_controller = ProfiledPIDControllerRadians(
+        1,
+        0,
+        0,
+        TrapezoidProfileRadians.Constraints(3 * math.pi, math.pi),
+    )
+
+    x_tolerance = 1 / 12.0
+    y_tolerance = 1 / 12.0
+    angular_tolerance = 1 / 12.0
+
+    wheel_base = 0.5
+    track_width = 0.5
+
+    kinematics = SwerveDrive4Kinematics(
+        Translation2d(wheel_base / 2, track_width / 2),
+        Translation2d(wheel_base / 2, -track_width / 2),
+        Translation2d(-wheel_base / 2, track_width / 2),
+        Translation2d(-wheel_base / 2, -track_width / 2),
+    )
+
+    odometry = SwerveDrive4Odometry(
+        kinematics,
+        Rotation2d(0),
+        swerve_module_positions,
+        Pose2d(0, 0, Rotation2d(0)),
+    )
+
+    def set_module_states(states):
+        nonlocal swerve_module_states
+        swerve_module_states = states
+
+    def get_robot_pose() -> Pose2d:
+        odometry.update(angle, swerve_module_positions)
+        return odometry.getPose()
+
+    with ManualSimTime() as sim:
+        subsystem = commands2.Subsystem()
+        waypoints: List[Pose2d] = []
+        waypoints.append(Pose2d(0, 0, Rotation2d(0)))
+        waypoints.append(Pose2d(1, 5, Rotation2d(3)))
+        config = TrajectoryConfig(8.8, 0.1)
+        trajectory = TrajectoryGenerator.generateTrajectory(waypoints, config)
+
+        end_state = trajectory.sample(trajectory.totalTime())
+
+        command = commands2.SwerveControllerCommand(
+            trajectory=trajectory,
+            pose=get_robot_pose,
+            kinematics=kinematics,
+            controller=HolonomicDriveController(
+                PIDController(0.6, 0, 0),
+                PIDController(0.6, 0, 0),
+                rot_controller,
+            ),
+            outputModuleStates=set_module_states,
+            requirements=(subsystem,),
+        )
+
+        timer.restart()
+
+        command.initialize()
+        while not command.isFinished():
+            command.execute()
+            angle = trajectory.sample(timer.get()).pose.rotation()
+
+            for i in range(0, len(swerve_module_positions)):
+                swerve_module_positions[i].distance += (
+                    swerve_module_states[i].speed * 0.005
+                )
+                swerve_module_positions[i].angle = swerve_module_states[i].angle
+
+            sim.step(0.005)
+
+        timer.stop()
+        command.end(True)
+
+        assert end_state.pose.X() == pytest.approx(get_robot_pose().X(), x_tolerance)
+
+        assert end_state.pose.Y() == pytest.approx(get_robot_pose().Y(), y_tolerance)
+
+        assert end_state.pose.rotation().radians() == pytest.approx(
+            get_robot_pose().rotation().radians(),
+            angular_tolerance,
+        )