Merge pull request #161 from frc937/heading-snapping

Heading snapping
frc937 · Mar 21, 2024 · e14c944 · e14c944
2 parents 834f888 + 2a13bd7
commit e14c944
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/src/main/java/frc/robot/commands/drive/DriveFieldOrientedHeadingSnapping.java b/src/main/java/frc/robot/commands/drive/DriveFieldOrientedHeadingSnapping.java
@@ -11,6 +11,7 @@
 
 package frc.robot.commands.drive;
 
+import edu.wpi.first.math.geometry.Rotation2d;
 import edu.wpi.first.math.geometry.Translation2d;
 import edu.wpi.first.math.kinematics.ChassisSpeeds;
 import edu.wpi.first.wpilibj2.command.Command;
@@ -28,6 +29,8 @@ public class DriveFieldOrientedHeadingSnapping extends Command {
   private final Supplier<Double> xSupplier, ySupplier, zSupplier;
   private final Supplier<Boolean> upSupplier, downSupplier, leftSupplier, rightSupplier;
 
+  private boolean resetHeading = false;
+
   /**
    * Drives the robot field oriented with heading snapping. All values are suppliers to make the
    * command more versatile.
@@ -65,6 +68,8 @@ public DriveFieldOrientedHeadingSnapping(
   public void initialize() {
     drive.setDriveMode("Field Oriented (Heading Snapping)");
     drive.setHeadingCorrection(true);
+
+    resetHeading = true;
   }
 
   // Called every time the scheduler runs while the command is scheduled.
@@ -85,11 +90,24 @@ public void execute() {
       headingX = -1;
     }
 
+    /* Allows us to adjust the heading using the joystick without it immediately snapping back */
+    if (resetHeading) {
+      if (headingX == 0 && headingY == 0 && Math.abs(zSupplier.get()) > 0) {
+        Rotation2d currentHeading = drive.getHeading();
+
+        headingX = currentHeading.getSin();
+        headingY = currentHeading.getCos();
+      }
+
+      resetHeading = false;
+    }
+
     ChassisSpeeds desiredSpeeds =
         drive.getTargetSpeeds(xSupplier.get(), ySupplier.get(), headingX, headingY);
 
     Translation2d translation = SwerveController.getTranslation2d(desiredSpeeds);
     if (headingX == 0 && headingY == 0 && Math.abs(zSupplier.get()) > 0) {
+      resetHeading = true;
       drive.driveRobot(translation, zSupplier.get() * drive.getMaximumAngularSpeed(), true);
     } else {
       drive.driveRobot(translation, desiredSpeeds.omegaRadiansPerSecond, true);

diff --git a/src/main/java/frc/robot/subsystems/Drive.java b/src/main/java/frc/robot/subsystems/Drive.java
@@ -14,6 +14,7 @@
 import com.pathplanner.lib.auto.AutoBuilder;
 import com.pathplanner.lib.util.HolonomicPathFollowerConfig;
 import com.pathplanner.lib.util.ReplanningConfig;
+import edu.wpi.first.math.geometry.Rotation2d;
 import edu.wpi.first.math.geometry.Translation2d;
 import edu.wpi.first.math.kinematics.ChassisSpeeds;
 import edu.wpi.first.math.util.Units;
@@ -181,6 +182,10 @@ public void setDriveMode(String driveMode) {
     driveLabelEntry.setString(driveMode);
   }
 
+  public Rotation2d getHeading() {
+    return drive.getPose().getRotation();
+  }
+
   /** Runs every scheduler run. */
   @Override
   public void periodic() {