001/************************* PROJECT RON *************************/ 002/* Copyright (c) 2026 StuyPulse Robotics. All rights reserved. */ 003/* Use of this source code is governed by an MIT-style license */ 004/* that can be found in the repository LICENSE file. */ 005/***************************************************************/ 006package com.stuypulse.robot.util.simulation; 007 008// Copyright 2021-2025 Iron Maple 5516 009// Original Source: 010// https://github.com/Shenzhen-Robotics-Alliance/maple-sim/blob/main/templates/CTRE%20Swerve%20with%20maple-sim/src/main/java/frc/robot/utils/simulation/MapleSimSwerveDrivetrain.java 011// 012// This code is licensed under MIT license (see https://mit-license.org/) 013import static edu.wpi.first.units.Units.*; 014 015import com.ctre.phoenix6.configs.CANcoderConfiguration; 016import com.ctre.phoenix6.configs.Slot0Configs; 017import com.ctre.phoenix6.configs.TalonFXConfiguration; 018import com.ctre.phoenix6.hardware.CANcoder; 019import com.ctre.phoenix6.hardware.Pigeon2; 020import com.ctre.phoenix6.hardware.TalonFX; 021import com.ctre.phoenix6.signals.StaticFeedforwardSignValue; 022import com.ctre.phoenix6.sim.CANcoderSimState; 023import com.ctre.phoenix6.sim.Pigeon2SimState; 024import com.ctre.phoenix6.sim.TalonFXSimState; 025import com.ctre.phoenix6.swerve.SwerveDrivetrain; 026import com.ctre.phoenix6.swerve.SwerveModule; 027import com.ctre.phoenix6.swerve.SwerveModuleConstants; 028import edu.wpi.first.math.geometry.Pose2d; 029import edu.wpi.first.math.geometry.Translation2d; 030import edu.wpi.first.math.system.plant.DCMotor; 031import edu.wpi.first.units.measure.*; 032import edu.wpi.first.wpilibj.RobotBase; 033import org.ironmaple.simulation.SimulatedArena; 034import org.ironmaple.simulation.drivesims.COTS; 035import org.ironmaple.simulation.drivesims.SwerveDriveSimulation; 036import org.ironmaple.simulation.drivesims.SwerveModuleSimulation; 037import org.ironmaple.simulation.drivesims.configs.DriveTrainSimulationConfig; 038import org.ironmaple.simulation.drivesims.configs.SwerveModuleSimulationConfig; 039import org.ironmaple.simulation.motorsims.SimulatedBattery; 040import org.ironmaple.simulation.motorsims.SimulatedMotorController; 041 042/** 043 * 044 * 045 * <h2>Injects Maple-Sim simulation data into a CTRE swerve drivetrain.</h2> 046 * 047 * <p> 048 * This class retrieves simulation data from Maple-Sim and injects it into the 049 * CTRE {@link 050 * com.ctre.phoenix6.swerve.SwerveDrivetrain} instance. 051 * 052 * <p> 053 * It replaces the {@link com.ctre.phoenix6.swerve.SimSwerveDrivetrain} class. 054 */ 055public class MapleSimSwerveDrivetrain { 056 057 private final Pigeon2SimState pigeonSim; 058 059 private final SimSwerveModule[] simModules; 060 061 public final SwerveDriveSimulation mapleSimDrive; 062 063 /** 064 * 065 * 066 * <h4>Constructs a drivetrain simulation using the specified parameters.</h4> 067 * 068 * @param simPeriod the time period of the simulation 069 * @param robotMassWithBumpers the total mass of the robot, including bumpers 070 * @param bumperLengthX the length of the bumper along the X-axis 071 * (influences the collision space 072 * of the robot) 073 * @param bumperWidthY the width of the bumper along the Y-axis 074 * (influences the collision space of 075 * the robot) 076 * @param driveMotorModel the {@link DCMotor} model for the drive motor, 077 * typically <code> 078 * DCMotor.getKrakenX60Foc() 079 * </code> 080 * @param steerMotorModel the {@link DCMotor} model for the steer motor, 081 * typically <code> 082 * DCMotor.getKrakenX60Foc() 083 * </code> 084 * @param wheelCOF the coefficient of friction of the drive wheels 085 * @param moduleLocations the locations of the swerve modules on the robot, 086 * in the order <code> 087 * FL, FR, BL, BR</code> 088 * @param pigeon the {@link Pigeon2} IMU used in the drivetrain 089 * @param modules the {@link SwerveModule}s, typically obtained via 090 * {@link 091 * SwerveDrivetrain#getModules()} 092 * @param moduleConstants the constants for the swerve modules 093 */ 094 public MapleSimSwerveDrivetrain( 095 Time simPeriod, 096 Mass robotMassWithBumpers, 097 Distance bumperLengthX, 098 Distance bumperWidthY, 099 DCMotor driveMotorModel, 100 DCMotor steerMotorModel, 101 double wheelCOF, 102 Translation2d[] moduleLocations, 103 Pigeon2 pigeon, 104 SwerveModule<TalonFX, TalonFX, CANcoder>[] modules, 105 SwerveModuleConstants<TalonFXConfiguration, TalonFXConfiguration, CANcoderConfiguration>... moduleConstants) { 106 this.pigeonSim = pigeon.getSimState(); 107 simModules = new SimSwerveModule[moduleConstants.length]; 108 DriveTrainSimulationConfig simulationConfig = DriveTrainSimulationConfig.Default() 109 .withRobotMass(robotMassWithBumpers) 110 .withBumperSize(bumperLengthX, bumperWidthY) 111 .withGyro(COTS.ofPigeon2()) 112 .withCustomModuleTranslations(moduleLocations) 113 .withSwerveModule( 114 new SwerveModuleSimulationConfig( 115 driveMotorModel, 116 steerMotorModel, 117 moduleConstants[0].DriveMotorGearRatio, 118 moduleConstants[0].SteerMotorGearRatio, 119 Volts.of(moduleConstants[0].DriveFrictionVoltage), 120 Volts.of(moduleConstants[0].SteerFrictionVoltage), 121 Meters.of(moduleConstants[0].WheelRadius), 122 KilogramSquareMeters 123 .of(moduleConstants[0].SteerInertia), 124 wheelCOF)); 125 mapleSimDrive = new SwerveDriveSimulation(simulationConfig, new Pose2d()); 126 SwerveModuleSimulation[] moduleSimulations = mapleSimDrive.getModules(); 127 for (int i = 0; i < this.simModules.length; i++) 128 simModules[i] = new SimSwerveModule(moduleConstants[0], moduleSimulations[i], modules[i]); 129 SimulatedArena.overrideSimulationTimings(simPeriod, 1); 130 } 131 132 /** 133 * 134 * 135 * <h4>Update the simulation.</h4> 136 * 137 * <p> 138 * Updates the Maple-Sim simulation and injects the results into the simulated 139 * CTRE devices, 140 * including motors and the IMU. 141 */ 142 public void update() { 143 SimulatedArena.getInstance().simulationPeriodic(); 144 pigeonSim.setRawYaw(mapleSimDrive.getSimulatedDriveTrainPose().getRotation().getMeasure()); 145 pigeonSim.setAngularVelocityZ( 146 RadiansPerSecond.of( 147 mapleSimDrive.getDriveTrainSimulatedChassisSpeedsRobotRelative().omegaRadiansPerSecond)); 148 } 149 150 /** 151 * 152 * 153 * <h2>Represents the simulation of a single {@link SwerveModule}.</h2> 154 */ 155 protected static class SimSwerveModule { 156 157 public final SwerveModuleConstants<TalonFXConfiguration, TalonFXConfiguration, CANcoderConfiguration> moduleConstant; 158 159 public final SwerveModuleSimulation moduleSimulation; 160 161 public SimSwerveModule( 162 SwerveModuleConstants<TalonFXConfiguration, TalonFXConfiguration, CANcoderConfiguration> moduleConstant, 163 SwerveModuleSimulation moduleSimulation, 164 SwerveModule<TalonFX, TalonFX, CANcoder> module) { 165 this.moduleConstant = moduleConstant; 166 this.moduleSimulation = moduleSimulation; 167 moduleSimulation.useDriveMotorController( 168 new TalonFXMotorControllerSim(module.getDriveMotor())); 169 moduleSimulation.useSteerMotorController( 170 new TalonFXMotorControllerWithRemoteCanCoderSim( 171 module.getSteerMotor(), module.getEncoder())); 172 } 173 } 174 175 // Static utils classes 176 public static class TalonFXMotorControllerSim implements SimulatedMotorController { 177 178 public final int id; 179 180 private final TalonFXSimState talonFXSimState; 181 182 public TalonFXMotorControllerSim(TalonFX talonFX) { 183 this.id = talonFX.getDeviceID(); 184 this.talonFXSimState = talonFX.getSimState(); 185 } 186 187 @Override 188 public Voltage updateControlSignal( 189 Angle mechanismAngle, 190 AngularVelocity mechanismVelocity, 191 Angle encoderAngle, 192 AngularVelocity encoderVelocity) { 193 talonFXSimState.setRawRotorPosition(encoderAngle); 194 talonFXSimState.setRotorVelocity(encoderVelocity); 195 talonFXSimState.setSupplyVoltage(SimulatedBattery.getBatteryVoltage()); 196 return talonFXSimState.getMotorVoltageMeasure(); 197 } 198 } 199 200 public static class TalonFXMotorControllerWithRemoteCanCoderSim 201 extends TalonFXMotorControllerSim { 202 203 private final int encoderId; 204 205 private final CANcoderSimState remoteCancoderSimState; 206 207 public TalonFXMotorControllerWithRemoteCanCoderSim(TalonFX talonFX, CANcoder cancoder) { 208 super(talonFX); 209 this.remoteCancoderSimState = cancoder.getSimState(); 210 this.encoderId = cancoder.getDeviceID(); 211 } 212 213 @Override 214 public Voltage updateControlSignal( 215 Angle mechanismAngle, 216 AngularVelocity mechanismVelocity, 217 Angle encoderAngle, 218 AngularVelocity encoderVelocity) { 219 remoteCancoderSimState.setSupplyVoltage(SimulatedBattery.getBatteryVoltage()); 220 remoteCancoderSimState.setRawPosition(mechanismAngle); 221 remoteCancoderSimState.setVelocity(mechanismVelocity); 222 return super.updateControlSignal( 223 mechanismAngle, mechanismVelocity, encoderAngle, encoderVelocity); 224 } 225 } 226 227 /** 228 * 229 * 230 * <h4>Regulates all {@link SwerveModuleConstants} for a drivetrain 231 * simulation.</h4> 232 * 233 * <p> 234 * This method processes an array of {@link SwerveModuleConstants} to apply 235 * necessary 236 * adjustments for simulation purposes, ensuring compatibility and avoiding 237 * known bugs. 238 * 239 * @see #regulateModuleConstantForSimulation(SwerveModuleConstants) 240 */ 241 public static SwerveModuleConstants<?, ?, ?>[] regulateModuleConstantsForSimulation( 242 SwerveModuleConstants<?, ?, ?>[] moduleConstants) { 243 for (SwerveModuleConstants<?, ?, ?> moduleConstant : moduleConstants) 244 regulateModuleConstantForSimulation(moduleConstant); 245 return moduleConstants; 246 } 247 248 /** 249 * 250 * 251 * <h4>Regulates the {@link SwerveModuleConstants} for a single module.</h4> 252 * 253 * <p> 254 * This method applies specific adjustments to the {@link SwerveModuleConstants} 255 * for simulation 256 * purposes. These changes have no effect on real robot operations and address 257 * known simulation 258 * bugs: 259 * 260 * <ul> 261 * <li><strong>Inverted Drive Motors:</strong> Prevents drive PID issues caused 262 * by inverted 263 * configurations. 264 * <li><strong>Non-zero CanCoder Offsets:</strong> Fixes potential module state 265 * optimization 266 * issues. 267 * <li><strong>Steer Motor PID:</strong> Adjusts PID values tuned for real 268 * robots to improve 269 * simulation performance. 270 * </ul> 271 * 272 * <h4>Note:This function is skipped when running on a real robot, ensuring no 273 * impact on constants 274 * used on real robot hardware.</h4> 275 */ 276 private static void regulateModuleConstantForSimulation( 277 SwerveModuleConstants<?, ?, ?> moduleConstants) { 278 // Skip regulation if running on a real robot 279 if (RobotBase.isReal()) 280 return; 281 // Apply simulation-specific adjustments to module constants 282 moduleConstants 283 . // Disable encoder offsets 284 withEncoderOffset(0) 285 . // Disable motor inversions for drive and steer motors 286 withDriveMotorInverted(false) 287 .withSteerMotorInverted(false) 288 . // Disable CanCoder inversion 289 withEncoderInverted(false) 290 . // Adjust steer motor PID gains for simulation 291 withSteerMotorGains( 292 new Slot0Configs() 293 .withKP(70) 294 .withKI(0) 295 .withKD(4.5) 296 .withKS(0) 297 .withKV(1.91) 298 .withKA(0) 299 .withStaticFeedforwardSign( 300 StaticFeedforwardSignValue.UseClosedLoopSign)) 301 .withSteerMotorGearRatio(16.0) 302 . // Adjust friction voltages 303 withDriveFrictionVoltage(Volts.of(0.1)) 304 .withSteerFrictionVoltage(Volts.of(0.05)) 305 . // Adjust steer inertia 306 withSteerInertia(KilogramSquareMeters.of(0.05)); 307 } 308}