diff --git a/commands2/__init__.py b/commands2/__init__.py index 4d9c2e8c..52efd58b 100644 --- a/commands2/__init__.py +++ b/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 @@ -75,6 +76,7 @@ "SequentialCommandGroup", "StartEndCommand", "Subsystem", + "SwerveControllerCommand", "TimedCommandRobot", "TrapezoidProfileSubsystem", "WaitCommand", diff --git a/commands2/swervecontrollercommand.py b/commands2/swervecontrollercommand.py new file mode 100644 index 00000000..fac15df7 --- /dev/null +++ b/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()) diff --git a/setup.py b/setup.py index b1c069cb..efa8c01c 100644 --- a/setup.py +++ b/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, diff --git a/tests/test_swervecontrollercommand.py b/tests/test_swervecontrollercommand.py new file mode 100644 index 00000000..317461d4 --- /dev/null +++ b/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, + )