Merge pull request #9 from FRC3636/drivetrain

src/main/java/com/frcteam3636/frc2024/RobotContainer.kt

@@ -1,6 +1,8 @@
 package com.frcteam3636.frc2024
 
+import com.frcteam3636.frc2024.subsystems.drivetrain.Drivetrain
 import com.frcteam3636.frc2024.subsystems.shooter.Shooter
+import edu.wpi.first.wpilibj.Joystick
 import edu.wpi.first.wpilibj2.command.Command
 import edu.wpi.first.wpilibj2.command.button.CommandXboxController
 
@@ -19,16 +21,19 @@ import edu.wpi.first.wpilibj2.command.button.CommandXboxController
  * directly reference the (single instance of the) object.
  */
 object RobotContainer  {
+    private val controller = CommandXboxController(0)
+    private val joystickLeft = Joystick(0)
+    private val joystickRight = Joystick(1)
+
     init {
+        Drivetrain
         configureBindings()
-        Shooter
     }
 
-    private val controller = CommandXboxController(0)
-
     /** Use this method to define your `trigger->command` mappings. */
     private fun configureBindings() {
-        controller.b().whileTrue(Shooter.shootCommand())
+//        controller.b().whileTrue(Shooter.shootCommand())
+        Drivetrain.defaultCommand = Drivetrain.driveWithJoysticks(translationJoystick = joystickLeft, rotationJoystick = joystickRight)
     }
 
     fun getAutonomousCommand(): Command? {

src/main/java/com/frcteam3636/frc2024/can/CAN.kt

@@ -0,0 +1,32 @@
+package com.frcteam3636.frc2024.can
+
+import com.ctre.phoenix6.hardware.TalonFX
+import com.revrobotics.CANSparkFlex
+import com.revrobotics.CANSparkLowLevel
+import com.revrobotics.CANSparkMax
+
+// This module contains one enum for each (device type, manufacturer) pair we use.
+
+enum class REVMotorControllerId(val num: Int) {
+    FrontLeftTurningMotor(1),
+    BackLeftTurningMotor(2),
+    BackRightTurningMotor(3),
+    FrontRightTurningMotor(4),
+    LeftShooterFlywheel(0),
+    RightShooterFlywheel(0),
+    OverTheBumperIntakeArm(0),
+    OverTheBumperIntakeFeed(0),
+    UnderTheBumperIntakeRoller(0),
+}
+
+fun CANSparkMax(id: REVMotorControllerId, type: CANSparkLowLevel.MotorType) = CANSparkMax(id.num, type)
+fun CANSparkFlex(id: REVMotorControllerId, type: CANSparkLowLevel.MotorType) = CANSparkFlex(id.num, type)
+
+enum class CTREMotorControllerId(val num: Int, val bus: String) {
+    FrontLeftDrivingMotor(1, "*"),
+    BackLeftDrivingMotor(2, "*"),
+    BackRightDrivingMotor(3, "*"),
+    FrontRightDrivingMotor(4, "*"),
+}
+
+fun TalonFX(id: CTREMotorControllerId) = TalonFX(id.num, id.bus)

src/main/java/com/frcteam3636/frc2024/subsystems/drivetrain/Drivetrain.kt

@@ -0,0 +1,186 @@
+package com.frcteam3636.frc2024.subsystems.drivetrain
+
+import edu.wpi.first.math.estimator.SwerveDrivePoseEstimator
+import edu.wpi.first.math.geometry.Pose2d
+import edu.wpi.first.math.kinematics.ChassisSpeeds
+import edu.wpi.first.math.kinematics.SwerveDriveKinematics
+import edu.wpi.first.math.kinematics.SwerveModuleState
+import edu.wpi.first.wpilibj.RobotBase
+import edu.wpi.first.wpilibj.smartdashboard.SmartDashboard
+import edu.wpi.first.wpilibj2.command.CommandScheduler
+import edu.wpi.first.wpilibj2.command.Subsystem
+import com.frcteam3636.frc2024.utils.swerve.PerCorner
+import org.littletonrobotics.junction.Logger
+import edu.wpi.first.math.util.Units
+import com.frcteam3636.frc2024.utils.swerve.*
+import com.frcteam3636.frc2024.can.*
+import edu.wpi.first.math.geometry.Rotation2d
+import edu.wpi.first.math.geometry.Rotation3d
+import edu.wpi.first.math.geometry.Translation2d
+import edu.wpi.first.math.kinematics.SwerveModulePosition
+import edu.wpi.first.wpilibj.Joystick
+import edu.wpi.first.wpilibj.smartdashboard.Field2d
+import edu.wpi.first.wpilibj2.command.Command
+import org.littletonrobotics.junction.AutoLogOutput
+import org.littletonrobotics.junction.LogTable
+import org.littletonrobotics.junction.inputs.LoggableInputs
+
+// A singleton object representing the drivetrain.
+object Drivetrain : Subsystem {
+    private val io = if (RobotBase.isReal()) {
+        DrivetrainIOReal()
+    } else {
+        DrivetrainIOSim()
+    }
+    private val inputs = DrivetrainIO.Inputs()
+
+    // Create swerve drivetrain kinematics using the translation parts of the module positions.
+    private val kinematics =
+        SwerveDriveKinematics(*MODULE_POSITIONS.map(Pose2d::getTranslation).toList().toTypedArray())
+
+    private val poseEstimator = SwerveDrivePoseEstimator(
+        kinematics, // swerve drive kinematics
+        inputs.gyroRotation.toRotation2d(),// initial gyro rotation
+        inputs.measuredPositions.toTypedArray(), // initial module positions
+        Pose2d() // initial pose
+        // TODO: add odometry standard deviation
+    )
+
+    init {
+        CommandScheduler.getInstance().registerSubsystem(this)
+    }
+
+    override fun periodic() {
+        io.updateInputs(inputs)
+        Logger.processInputs("Drivetrain", inputs)
+
+        poseEstimator.update(
+            inputs.gyroRotation.toRotation2d(),
+            inputs.measuredPositions.toTypedArray()
+        )
+
+        Logger.recordOutput("Drivetrain/EstimatedPose", estimatedPose)
+    }
+
+    // The rotation of the robot as measured by the gyro.
+    var gyroRotation
+        get() = inputs.gyroRotation
+        set(value) {
+            io.resetGyro(value)
+        }
+
+    private var moduleStates: PerCorner<SwerveModuleState>
+        // Get the measured module states from the inputs.
+        get() = inputs.measuredStates
+        // Set the desired module states.
+        set(value) {
+            io.setDesiredStates(value)
+            Logger.recordOutput("Drivetrain/DesiredStates", *value.toTypedArray())
+        }
+
+    // The current speed of chassis relative to the ground, assuming that the wheels have perfect traction with the ground.
+    var chassisSpeeds: ChassisSpeeds
+        // Get the chassis speeds from the measured module states.
+        get() = kinematics.cornerStatesToChassisSpeeds(inputs.measuredStates)
+        // Set the drivetrain to move with the given chassis speeds.
+        //
+        // Note that the speeds are relative to the chassis, not the field.
+        set(value) {
+            val states = kinematics.toCornerSwerveModuleStates(value)
+
+            // TODO: desaturate states
+
+            moduleStates = states
+        }
+
+    // Set the drivetrain to an X-formation to passively prevent movement in any direction.
+    fun brake() {
+        // set the modules to radiate outwards from the chassis origin
+        moduleStates = MODULE_POSITIONS.map { position -> SwerveModuleState(0.0, position.rotation) }
+    }
+
+    // Get the estimated pose of the drivetrain using the pose estimator.
+    val estimatedPose: Pose2d
+        get() = poseEstimator.estimatedPosition
+
+    fun driveWithJoysticks(translationJoystick: Joystick, rotationJoystick: Joystick): Command = run {
+        chassisSpeeds = ChassisSpeeds.fromFieldRelativeSpeeds(
+            translationJoystick.y, translationJoystick.x, rotationJoystick.x, gyroRotation.toRotation2d()
+        )
+    }
+}
+
+abstract class DrivetrainIO {
+    abstract val gyro: Gyro
+    abstract val modules: PerCorner<out SwerveModule>
+
+    class Inputs : LoggableInputs {
+        var gyroRotation: Rotation3d = Rotation3d()
+        var measuredStates: PerCorner<SwerveModuleState> = PerCorner.generate { SwerveModuleState() }
+        var measuredPositions: PerCorner<SwerveModulePosition> = PerCorner.generate { SwerveModulePosition() }
+
+        override fun fromLog(table: LogTable?) {
+            gyroRotation = table?.get("GyroRotation", gyroRotation)!![0]
+            measuredStates =
+                PerCorner.fromConventionalArray(table.get("MeasuredStates", *measuredStates.toTypedArray()))
+            measuredPositions =
+                PerCorner.fromConventionalArray(table.get("MeasuredPositions", *measuredPositions.toTypedArray()))
+        }
+
+        override fun toLog(table: LogTable?) {
+            table?.put("GyroRotation", gyroRotation)
+            table?.put("MeasuredStates", *measuredStates.toTypedArray())
+            table?.put("MeasuredPositions", *measuredPositions.toTypedArray())
+        }
+    }
+
+    fun updateInputs(inputs: Inputs) {
+        gyro.periodic()
+        modules.forEach(SwerveModule::periodic)
+
+        inputs.gyroRotation = gyro.rotation
+        inputs.measuredStates = modules.map { it.state }
+        inputs.measuredPositions = modules.map { it.position }
+    }
+
+    fun resetGyro(rotation: Rotation3d) {
+        gyro.rotation = rotation
+    }
+
+    fun setDesiredStates(states: PerCorner<SwerveModuleState>) {
+        modules.zip(states).forEach { (module, state) ->
+            module.desiredState = state
+        }
+    }
+}
+
+class DrivetrainIOReal : DrivetrainIO() {
+    override val gyro = GyroNavX()
+    override val modules = MODULE_CAN_IDS.zip(MODULE_POSITIONS).map { (can, pose) ->
+        val (driving, turning) = can
+        MAXSwerveModule(driving, turning, pose.rotation)
+    }
+}
+
+class DrivetrainIOSim : DrivetrainIO() {
+    override val modules = PerCorner.generate { SimSwerveModule() }
+    override val gyro = GyroSim(modules.map { it })
+}
+
+// Constants
+internal val WHEEL_BASE: Double = Units.inchesToMeters(13.0)
+internal val TRACK_WIDTH: Double = Units.inchesToMeters(14.0)
+
+internal val MODULE_POSITIONS = PerCorner(
+    frontLeft = Pose2d(Translation2d(WHEEL_BASE, TRACK_WIDTH) / 2.0, Rotation2d.fromDegrees(-90.0)),
+    frontRight = Pose2d(Translation2d(WHEEL_BASE, -TRACK_WIDTH) / 2.0, Rotation2d.fromDegrees(180.0)),
+    backRight = Pose2d(Translation2d(-WHEEL_BASE, TRACK_WIDTH) / 2.0, Rotation2d.fromDegrees(0.0)),
+    backLeft = Pose2d(Translation2d(-WHEEL_BASE, -TRACK_WIDTH) / 2.0, Rotation2d.fromDegrees(90.0))
+)
+
+internal val MODULE_CAN_IDS = PerCorner(
+    frontLeft = Pair(CTREMotorControllerId.FrontLeftDrivingMotor, REVMotorControllerId.FrontLeftTurningMotor),
+    frontRight = Pair(CTREMotorControllerId.FrontRightDrivingMotor, REVMotorControllerId.FrontRightTurningMotor),
+    backRight = Pair(CTREMotorControllerId.BackRightDrivingMotor, REVMotorControllerId.BackRightTurningMotor),
+    backLeft = Pair(CTREMotorControllerId.BackLeftDrivingMotor, REVMotorControllerId.BackLeftTurningMotor),
+)

src/main/java/com/frcteam3636/frc2024/subsystems/drivetrain/Gyro.kt

@@ -0,0 +1,48 @@
+package com.frcteam3636.frc2024.subsystems.drivetrain
+
+import com.frcteam3636.frc2024.Robot
+import com.frcteam3636.frc2024.utils.swerve.PerCorner
+import com.kauailabs.navx.frc.AHRS
+import edu.wpi.first.math.geometry.Quaternion
+import edu.wpi.first.math.geometry.Rotation3d
+import edu.wpi.first.math.geometry.Translation2d
+import kotlin.math.sign
+
+
+interface Gyro {
+    var rotation: Rotation3d
+
+    fun periodic() {}
+}
+
+
+class GyroNavX(private var offset: Rotation3d = Rotation3d()) : Gyro {
+    private val ahrs = AHRS()
+
+    override var rotation: Rotation3d
+        get() = offset + Rotation3d(
+            Quaternion(
+                ahrs.quaternionW.toDouble(),
+                ahrs.quaternionX.toDouble(),
+                ahrs.quaternionY.toDouble(),
+                ahrs.quaternionZ.toDouble()
+            )
+        )
+        set(value) {
+            offset = value - rotation
+        }
+}
+
+class GyroSim(private val modules: PerCorner<SwerveModule>) : Gyro {
+    override var rotation = Rotation3d()
+
+    override fun periodic() {
+        val moduleVelocities = modules.map { Translation2d(it.state.speedMetersPerSecond, it.state.angle) }
+        val translationVelocity = moduleVelocities.reduce(Translation2d::plus) / 4.0;
+        val rotationalVelocities = moduleVelocities.map { it - translationVelocity }
+        val yawVelocity =
+            sign(rotationalVelocities.frontLeft.y) * rotationalVelocities.frontLeft.norm / MODULE_POSITIONS.frontLeft.translation.norm
+
+        rotation += Rotation3d(0.0, 0.0, yawVelocity) * Robot.period
+    }
+}