317 lines
14 KiB
Python
317 lines
14 KiB
Python
"""
|
||
components/drive.py – DriveSubsystem ported to MagicBot component.
|
||
|
||
Owns four SwerveModule instances. All swerve math uses standard wpimath.
|
||
Integrates with PathPlannerLib and Limelight vision.
|
||
"""
|
||
|
||
import math
|
||
import os
|
||
|
||
import magicbot
|
||
from pathplannerlib.auto import AutoBuilder
|
||
from pathplannerlib.config import PIDConstants, RobotConfig
|
||
from pathplannerlib.controller import PPHolonomicDriveController
|
||
from pathplannerlib.logging import PathPlannerLogging
|
||
from phoenix6.hardware import Pigeon2
|
||
from wpilib import DriverStation, RobotState, SendableChooser, SmartDashboard, getDeployDirectory
|
||
from wpimath.estimator import SwerveDrive4PoseEstimator
|
||
from wpimath.geometry import Pose2d, Pose3d
|
||
from wpimath.kinematics import (
|
||
ChassisSpeeds,
|
||
SwerveDrive4Kinematics,
|
||
)
|
||
|
||
import constants
|
||
import LimelightHelpers
|
||
from components.swerve_module import SwerveModule
|
||
|
||
|
||
class DriveComponent:
|
||
"""
|
||
MagicBot component that drives four swerve modules and fuses vision
|
||
measurements via a SwerveDrive4PoseEstimator.
|
||
"""
|
||
|
||
def setup(self) -> None:
|
||
# ── Swerve modules ─────────────────────────────────────────────────
|
||
self._front_left = SwerveModule(
|
||
name="Front Left",
|
||
drive_id=constants.SwerveDriveIDs.FRONT_LEFT_DRIVE_ID,
|
||
steer_id=constants.SwerveDriveIDs.FRONT_LEFT_STEER_ID,
|
||
encoder_id=constants.SwerveDriveIDs.FRONT_LEFT_ENCODER_ID,
|
||
steer_offset_rotations=constants.SwerveDriveOffsets.FRONT_LEFT_MOTOR_OFFSET,
|
||
drive_inverted=False,
|
||
canbus="Drive",
|
||
)
|
||
self._front_right = SwerveModule(
|
||
name="Front Right",
|
||
drive_id=constants.SwerveDriveIDs.FRONT_RIGHT_DRIVE_ID,
|
||
steer_id=constants.SwerveDriveIDs.FRONT_RIGHT_STEER_ID,
|
||
encoder_id=constants.SwerveDriveIDs.FRONT_RIGHT_ENCODER_ID,
|
||
steer_offset_rotations=constants.SwerveDriveOffsets.FRONT_RIGHT_MOTOR_OFFSET,
|
||
drive_inverted=True, # Front Right is inverted
|
||
canbus="Drive",
|
||
)
|
||
self._back_left = SwerveModule(
|
||
name="Back Left",
|
||
drive_id=constants.SwerveDriveIDs.BACK_LEFT_DRIVE_ID,
|
||
steer_id=constants.SwerveDriveIDs.BACK_LEFT_STEER_ID,
|
||
encoder_id=constants.SwerveDriveIDs.BACK_LEFT_ENCODER_ID,
|
||
steer_offset_rotations=constants.SwerveDriveOffsets.BACK_LEFT_MOTOR_OFFSET,
|
||
drive_inverted=False,
|
||
canbus="Drive",
|
||
)
|
||
self._back_right = SwerveModule(
|
||
name="Back Right",
|
||
drive_id=constants.SwerveDriveIDs.BACK_RIGHT_DRIVE_ID,
|
||
steer_id=constants.SwerveDriveIDs.BACK_RIGHT_STEER_ID,
|
||
encoder_id=constants.SwerveDriveIDs.BACK_RIGHT_ENCODER_ID,
|
||
steer_offset_rotations=constants.SwerveDriveOffsets.BACK_RIGHT_MOTOR_OFFSET,
|
||
drive_inverted=True, # Back Right is inverted
|
||
canbus="Drive",
|
||
)
|
||
self._modules = [
|
||
self._front_left,
|
||
self._front_right,
|
||
self._back_left,
|
||
self._back_right,
|
||
]
|
||
|
||
# ── Kinematics ─────────────────────────────────────────────────────
|
||
self._kinematics = SwerveDrive4Kinematics(
|
||
constants.SwerveDriveModuleLocations.FRONT_LEFT,
|
||
constants.SwerveDriveModuleLocations.FRONT_RIGHT,
|
||
constants.SwerveDriveModuleLocations.BACK_LEFT,
|
||
constants.SwerveDriveModuleLocations.BACK_RIGHT,
|
||
)
|
||
|
||
# ── Gyro ───────────────────────────────────────────────────────────
|
||
self._pigeon = Pigeon2(constants.ExtraIDEntities.PIGEON_ID, "Drive")
|
||
|
||
# ── Pose estimator ─────────────────────────────────────────────────
|
||
self._pose_estimator = SwerveDrive4PoseEstimator(
|
||
self._kinematics,
|
||
self._pigeon.getRotation2d(),
|
||
self._module_positions(),
|
||
Pose2d(),
|
||
)
|
||
# 3 std-dev values: (x_std, y_std, heading_std)
|
||
self._pose_estimator.setVisionMeasurementStdDevs((0.7, 0.7, 0.1))
|
||
|
||
# ── Desired chassis speeds (reset each loop) ────────────────────────
|
||
self._desired_speeds = ChassisSpeeds()
|
||
|
||
# ── PathPlanner target pose (for @feedback) ────────────────────────
|
||
self._pp_target_pose = Pose2d()
|
||
|
||
# ── PathPlanner / AutoBuilder ──────────────────────────────────────
|
||
# _DriveSubsystemAdapter is imported from robot.py at runtime to
|
||
# avoid circular imports; it is passed in via inject_drive_adapter().
|
||
self._drive_adapter = None
|
||
self._auto_chooser: SendableChooser | None = None
|
||
|
||
def inject_drive_adapter(self, adapter, robot_config: RobotConfig) -> None:
|
||
"""
|
||
Called by robot.py after all objects are constructed.
|
||
Finalises AutoBuilder configuration and builds the auto chooser.
|
||
"""
|
||
AutoBuilder.configure(
|
||
self.get_robot_pose_2d, # pose supplier
|
||
self.set_pose, # pose resetter
|
||
self.get_robot_relative_speeds, # speeds supplier
|
||
lambda speeds, ff: self.set_chassis_speeds(speeds), # consumer (no negation)
|
||
PPHolonomicDriveController(
|
||
PIDConstants(5, 0, 0),
|
||
PIDConstants(math.pi, 0, 0),
|
||
),
|
||
robot_config,
|
||
self._is_red_alliance, # flip-path predicate
|
||
adapter, # dummy Subsystem
|
||
)
|
||
|
||
# Build auto chooser manually so a bad path file skips that auto instead
|
||
# of crashing the entire chooser (PathPlanner throws ZeroDivisionError on
|
||
# paths with consecutive rotation targets at the same point).
|
||
autos_dir = os.path.join(getDeployDirectory(), "pathplanner", "autos")
|
||
self._auto_chooser = SendableChooser()
|
||
self._auto_chooser.setDefaultOption("Do Nothing", None)
|
||
if os.path.isdir(autos_dir):
|
||
for filename in os.listdir(autos_dir):
|
||
auto_name = filename.removesuffix(".auto")
|
||
try:
|
||
self._auto_chooser.addOption(auto_name, AutoBuilder.buildAuto(auto_name))
|
||
except Exception as build_err:
|
||
SmartDashboard.putString(f"Auto/{auto_name} error", str(build_err))
|
||
SmartDashboard.putData("Auto Chooser", self._auto_chooser)
|
||
|
||
PathPlannerLogging.setLogTargetPoseCallback(
|
||
lambda pose: setattr(self, "_pp_target_pose", pose)
|
||
)
|
||
|
||
# ── Public drive API ───────────────────────────────────────────────────────
|
||
|
||
def set_chassis_speeds_field_relative(
|
||
self, vx: float, vy: float, omega: float
|
||
) -> None:
|
||
"""
|
||
Convert joystick inputs (with quadratic scaling) to robot-relative
|
||
ChassisSpeeds and store them to be applied in execute().
|
||
"""
|
||
scaled_vx = vx * abs(vx) * constants.RobotChassisLimits.MAX_ROBOT_LINEAR_VELOCITY
|
||
scaled_vy = vy * abs(vy) * constants.RobotChassisLimits.MAX_ROBOT_LINEAR_VELOCITY
|
||
scaled_omega = omega * abs(omega) * constants.RobotChassisLimits.MAX_ROBOT_ANGULAR_VELOCITY
|
||
|
||
self._desired_speeds = ChassisSpeeds.fromFieldRelativeSpeeds(
|
||
scaled_vx,
|
||
scaled_vy,
|
||
scaled_omega,
|
||
self._pigeon.getRotation2d(),
|
||
)
|
||
|
||
def set_chassis_speeds(self, speeds: ChassisSpeeds) -> None:
|
||
"""Direct robot-relative speed input (used by PathPlanner consumer)."""
|
||
self._desired_speeds = speeds
|
||
|
||
def stop_drive(self) -> None:
|
||
self._desired_speeds = ChassisSpeeds()
|
||
|
||
def rotate_chassis(self, rot_speed: float) -> None:
|
||
self._desired_speeds = ChassisSpeeds.fromRobotRelativeSpeeds(
|
||
0, 0, rot_speed, self._pigeon.getRotation2d()
|
||
)
|
||
|
||
def get_bot_pose(self) -> Pose3d:
|
||
return Pose3d(self._pose_estimator.getEstimatedPosition())
|
||
|
||
def get_robot_pose_2d(self) -> Pose2d:
|
||
return self._pose_estimator.getEstimatedPosition()
|
||
|
||
def get_robot_relative_speeds(self) -> ChassisSpeeds:
|
||
states = self._module_states()
|
||
return self._kinematics.toChassisSpeeds(states)
|
||
|
||
def set_pose(self, pose: Pose2d) -> None:
|
||
self._pose_estimator.resetPosition(
|
||
self._pigeon.getRotation2d(),
|
||
self._module_positions(),
|
||
pose,
|
||
)
|
||
|
||
def set_brake_mode(self, brake: bool) -> None:
|
||
from phoenix6.configs import MotorOutputConfigs
|
||
from phoenix6.signals import NeutralModeValue
|
||
mode = NeutralModeValue.BRAKE if brake else NeutralModeValue.COAST
|
||
for module in self._modules:
|
||
cfg = MotorOutputConfigs().with_neutral_mode(mode)
|
||
module._drive_motor.configurator.apply(cfg)
|
||
module._steer_motor.configurator.apply(cfg)
|
||
|
||
def get_auto_command(self):
|
||
if self._auto_chooser is not None:
|
||
return self._auto_chooser.getSelected()
|
||
return None
|
||
|
||
# ── @feedback properties ───────────────────────────────────────────────────
|
||
|
||
@magicbot.feedback
|
||
def robot_x(self) -> float:
|
||
return self._pose_estimator.getEstimatedPosition().X()
|
||
|
||
@magicbot.feedback
|
||
def robot_y(self) -> float:
|
||
return self._pose_estimator.getEstimatedPosition().Y()
|
||
|
||
@magicbot.feedback
|
||
def robot_heading_degrees(self) -> float:
|
||
return self._pose_estimator.getEstimatedPosition().rotation().degrees()
|
||
|
||
@magicbot.feedback
|
||
def pathplanner_target_x(self) -> float:
|
||
return self._pp_target_pose.X()
|
||
|
||
@magicbot.feedback
|
||
def pathplanner_target_y(self) -> float:
|
||
return self._pp_target_pose.Y()
|
||
|
||
# ── MagicBot lifecycle ─────────────────────────────────────────────────────
|
||
|
||
def execute(self) -> None:
|
||
# Apply desired speeds through kinematics
|
||
module_states = self._kinematics.toSwerveModuleStates(self._desired_speeds)
|
||
|
||
# Desaturate wheel speeds to enforce linear velocity limit
|
||
SwerveDrive4Kinematics.desaturateWheelSpeeds(
|
||
module_states,
|
||
constants.RobotChassisLimits.MAX_ROBOT_LINEAR_VELOCITY,
|
||
)
|
||
|
||
# Send states to modules
|
||
for module, state in zip(self._modules, module_states):
|
||
module.set_desired_state(state)
|
||
|
||
# Reset desired speeds so the robot stops if nothing sets them this loop
|
||
self._desired_speeds = ChassisSpeeds()
|
||
|
||
# Update pose estimator with gyro + module positions
|
||
self._pose_estimator.update(
|
||
self._pigeon.getRotation2d(),
|
||
self._module_positions(),
|
||
)
|
||
|
||
heading_deg = self.get_robot_pose_2d().rotation().degrees()
|
||
|
||
# Update Limelight orientation so MegaTag2 knows where the robot is facing
|
||
LimelightHelpers.SetRobotOrientation(
|
||
constants.CameraConstants.FRONT_CAMERA_NAME, heading_deg, 0, 0, 0, 0, 0
|
||
)
|
||
LimelightHelpers.SetRobotOrientation(
|
||
constants.CameraConstants.BACK_CAMERA_NAME, heading_deg, 0, 0, 0, 0, 0
|
||
)
|
||
|
||
# IMU mode: 1 when disabled (external), 4 when enabled (fused)
|
||
imu_mode = 1 if RobotState.isDisabled() else 4
|
||
LimelightHelpers.SetIMUMode(constants.CameraConstants.FRONT_CAMERA_NAME, imu_mode)
|
||
LimelightHelpers.SetIMUMode(constants.CameraConstants.BACK_CAMERA_NAME, imu_mode)
|
||
|
||
# MegaTag2 vision fusion
|
||
front_estimate = LimelightHelpers.getBotPoseEstimate_wpiBlue_MegaTag2(
|
||
constants.CameraConstants.FRONT_CAMERA_NAME
|
||
)
|
||
back_estimate = LimelightHelpers.getBotPoseEstimate_wpiBlue_MegaTag2(
|
||
constants.CameraConstants.BACK_CAMERA_NAME
|
||
)
|
||
|
||
if front_estimate is not None and front_estimate.tagCount >= 2:
|
||
self._pose_estimator.addVisionMeasurement(
|
||
front_estimate.pose,
|
||
front_estimate.timestampSeconds,
|
||
)
|
||
|
||
if back_estimate is not None and back_estimate.tagCount >= 2:
|
||
self._pose_estimator.addVisionMeasurement(
|
||
back_estimate.pose,
|
||
back_estimate.timestampSeconds,
|
||
)
|
||
|
||
# Publish per-module telemetry
|
||
for module in self._modules:
|
||
module.publish_telemetry()
|
||
|
||
# ── Helpers ────────────────────────────────────────────────────────────────
|
||
|
||
def _module_positions(self):
|
||
fl, fr, bl, br = self._modules
|
||
return (fl.get_position(), fr.get_position(), bl.get_position(), br.get_position())
|
||
|
||
def _module_states(self):
|
||
fl, fr, bl, br = self._modules
|
||
return (fl.get_state(), fr.get_state(), bl.get_state(), br.get_state())
|
||
|
||
@staticmethod
|
||
def _is_red_alliance() -> bool:
|
||
alliance = DriverStation.getAlliance()
|
||
if alliance:
|
||
return alliance == DriverStation.Alliance.kRed
|
||
return False
|