From aaf24e25524051cbc935bb4bd35c06eb014a6f28 Mon Sep 17 00:00:00 2001
From: Prateek Machiraju <prateek.machiraju@gmail.com>
Date: Mon, 19 Apr 2021 00:00:11 -0400
Subject: [PATCH] [wpilib] Fix initial heading behavior in
 HolonomicDriveController (#3290)

---
 .../cpp/controller/HolonomicDriveController.cpp   |  7 +++++++
 .../frc/controller/HolonomicDriveController.h     |  2 ++
 .../controller/HolonomicDriveControllerTest.cpp   | 14 ++++++++++++++
 .../controller/HolonomicDriveController.java      |  9 +++++++++
 .../controller/HolonomicDriveControllerTest.java  | 15 +++++++++++++++
 5 files changed, 47 insertions(+)

diff --git a/wpilibc/src/main/native/cpp/controller/HolonomicDriveController.cpp b/wpilibc/src/main/native/cpp/controller/HolonomicDriveController.cpp
index d5b3833fe5..a4b11f247e 100644
--- a/wpilibc/src/main/native/cpp/controller/HolonomicDriveController.cpp
+++ b/wpilibc/src/main/native/cpp/controller/HolonomicDriveController.cpp
@@ -34,6 +34,13 @@ void HolonomicDriveController::SetTolerance(const Pose2d& tolerance) {
 ChassisSpeeds HolonomicDriveController::Calculate(
     const Pose2d& currentPose, const Pose2d& poseRef,
     units::meters_per_second_t linearVelocityRef, const Rotation2d& angleRef) {
+  // If this is the first run, then we need to reset the theta controller to the
+  // current pose's heading.
+  if (m_firstRun) {
+    m_thetaController.Reset(currentPose.Rotation().Radians());
+    m_firstRun = false;
+  }
+
   // Calculate feedforward velocities (field-relative)
   auto xFF = linearVelocityRef * poseRef.Rotation().Cos();
   auto yFF = linearVelocityRef * poseRef.Rotation().Sin();
diff --git a/wpilibc/src/main/native/include/frc/controller/HolonomicDriveController.h b/wpilibc/src/main/native/include/frc/controller/HolonomicDriveController.h
index a02fd4e445..b85eac51ed 100644
--- a/wpilibc/src/main/native/include/frc/controller/HolonomicDriveController.h
+++ b/wpilibc/src/main/native/include/frc/controller/HolonomicDriveController.h
@@ -105,5 +105,7 @@ class HolonomicDriveController {
   frc2::PIDController m_xController;
   frc2::PIDController m_yController;
   ProfiledPIDController<units::radian> m_thetaController;
+
+  bool m_firstRun = true;
 };
 }  // namespace frc
diff --git a/wpilibc/src/test/native/cpp/controller/HolonomicDriveControllerTest.cpp b/wpilibc/src/test/native/cpp/controller/HolonomicDriveControllerTest.cpp
index d12eaf76cb..388c346c31 100644
--- a/wpilibc/src/test/native/cpp/controller/HolonomicDriveControllerTest.cpp
+++ b/wpilibc/src/test/native/cpp/controller/HolonomicDriveControllerTest.cpp
@@ -47,3 +47,17 @@ TEST(HolonomicDriveControllerTest, ReachesReference) {
   EXPECT_NEAR_UNITS(frc::AngleModulus(robotPose.Rotation().Radians()), 0_rad,
                     kAngularTolerance);
 }
+
+TEST(HolonomicDriveControllerTest, DoesNotRotateUnnecessarily) {
+  frc::HolonomicDriveController controller{
+      frc2::PIDController{1, 0, 0}, frc2::PIDController{1, 0, 0},
+      frc::ProfiledPIDController<units::radian>{
+          1, 0, 0,
+          frc::TrapezoidProfile<units::radian>::Constraints{
+              4_rad_per_s, 2_rad_per_s / 1_s}}};
+
+  frc::ChassisSpeeds speeds = controller.Calculate(
+      frc::Pose2d(0_m, 0_m, 1.57_rad), frc::Pose2d(), 0_mps, 1.57_rad);
+
+  EXPECT_EQ(0, speeds.omega.to<double>());
+}
diff --git a/wpilibj/src/main/java/edu/wpi/first/wpilibj/controller/HolonomicDriveController.java b/wpilibj/src/main/java/edu/wpi/first/wpilibj/controller/HolonomicDriveController.java
index 0072d36157..f1d7f81da7 100644
--- a/wpilibj/src/main/java/edu/wpi/first/wpilibj/controller/HolonomicDriveController.java
+++ b/wpilibj/src/main/java/edu/wpi/first/wpilibj/controller/HolonomicDriveController.java
@@ -31,6 +31,8 @@ public class HolonomicDriveController {
   private final PIDController m_yController;
   private final ProfiledPIDController m_thetaController;
 
+  private boolean m_firstRun = true;
+
   /**
    * Constructs a holonomic drive controller.
    *
@@ -82,6 +84,13 @@ public class HolonomicDriveController {
   @SuppressWarnings("LocalVariableName")
   public ChassisSpeeds calculate(
       Pose2d currentPose, Pose2d poseRef, double linearVelocityRefMeters, Rotation2d angleRef) {
+    // If this is the first run, then we need to reset the theta controller to the current pose's
+    // heading.
+    if (m_firstRun) {
+      m_thetaController.reset(currentPose.getRotation().getRadians());
+      m_firstRun = false;
+    }
+
     // Calculate feedforward velocities (field-relative).
     double xFF = linearVelocityRefMeters * poseRef.getRotation().getCos();
     double yFF = linearVelocityRefMeters * poseRef.getRotation().getSin();
diff --git a/wpilibj/src/test/java/edu/wpi/first/wpilibj/controller/HolonomicDriveControllerTest.java b/wpilibj/src/test/java/edu/wpi/first/wpilibj/controller/HolonomicDriveControllerTest.java
index e63b6b1c25..c7610b8329 100644
--- a/wpilibj/src/test/java/edu/wpi/first/wpilibj/controller/HolonomicDriveControllerTest.java
+++ b/wpilibj/src/test/java/edu/wpi/first/wpilibj/controller/HolonomicDriveControllerTest.java
@@ -72,4 +72,19 @@ class HolonomicDriveControllerTest {
                 MathUtil.angleModulus(finalRobotPose.getRotation().getRadians()),
                 kAngularTolerance));
   }
+
+  @Test
+  void testDoesNotRotateUnnecessarily() {
+    var controller =
+        new HolonomicDriveController(
+            new PIDController(1, 0, 0),
+            new PIDController(1, 0, 0),
+            new ProfiledPIDController(1, 0, 0, new TrapezoidProfile.Constraints(4, 2)));
+
+    ChassisSpeeds speeds =
+        controller.calculate(
+            new Pose2d(0, 0, new Rotation2d(1.57)), new Pose2d(), 0, new Rotation2d(1.57));
+
+    assertEquals(0.0, speeds.omegaRadiansPerSecond);
+  }
 }