001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.math.controller;
006
007import edu.wpi.first.math.geometry.Pose2d;
008import edu.wpi.first.math.geometry.Rotation2d;
009import edu.wpi.first.math.kinematics.ChassisSpeeds;
010import edu.wpi.first.math.trajectory.Trajectory;
011
012/**
013 * This holonomic drive controller can be used to follow trajectories using a holonomic drivetrain
014 * (i.e. swerve or mecanum). Holonomic trajectory following is a much simpler problem to solve
015 * compared to skid-steer style drivetrains because it is possible to individually control forward,
016 * sideways, and angular velocity.
017 *
018 * <p>The holonomic drive controller takes in one PID controller for each direction, forward and
019 * sideways, and one profiled PID controller for the angular direction. Because the heading dynamics
020 * are decoupled from translations, users can specify a custom heading that the drivetrain should
021 * point toward. This heading reference is profiled for smoothness.
022 */
023@SuppressWarnings("MemberName")
024public class HolonomicDriveController {
025  private Pose2d m_poseError = new Pose2d();
026  private Rotation2d m_rotationError = new Rotation2d();
027  private Pose2d m_poseTolerance = new Pose2d();
028  private boolean m_enabled = true;
029
030  private final PIDController m_xController;
031  private final PIDController m_yController;
032  private final ProfiledPIDController m_thetaController;
033
034  private boolean m_firstRun = true;
035
036  /**
037   * Constructs a holonomic drive controller.
038   *
039   * @param xController A PID Controller to respond to error in the field-relative x direction.
040   * @param yController A PID Controller to respond to error in the field-relative y direction.
041   * @param thetaController A profiled PID controller to respond to error in angle.
042   */
043  @SuppressWarnings("ParameterName")
044  public HolonomicDriveController(
045      PIDController xController, PIDController yController, ProfiledPIDController thetaController) {
046    m_xController = xController;
047    m_yController = yController;
048    m_thetaController = thetaController;
049  }
050
051  /**
052   * Returns true if the pose error is within tolerance of the reference.
053   *
054   * @return True if the pose error is within tolerance of the reference.
055   */
056  public boolean atReference() {
057    final var eTranslate = m_poseError.getTranslation();
058    final var eRotate = m_rotationError;
059    final var tolTranslate = m_poseTolerance.getTranslation();
060    final var tolRotate = m_poseTolerance.getRotation();
061    return Math.abs(eTranslate.getX()) < tolTranslate.getX()
062        && Math.abs(eTranslate.getY()) < tolTranslate.getY()
063        && Math.abs(eRotate.getRadians()) < tolRotate.getRadians();
064  }
065
066  /**
067   * Sets the pose error which is considered tolerance for use with atReference().
068   *
069   * @param tolerance The pose error which is tolerable.
070   */
071  public void setTolerance(Pose2d tolerance) {
072    m_poseTolerance = tolerance;
073  }
074
075  /**
076   * Returns the next output of the holonomic drive controller.
077   *
078   * @param currentPose The current pose.
079   * @param poseRef The desired pose.
080   * @param linearVelocityRefMeters The linear velocity reference.
081   * @param angleRef The angular reference.
082   * @return The next output of the holonomic drive controller.
083   */
084  @SuppressWarnings("LocalVariableName")
085  public ChassisSpeeds calculate(
086      Pose2d currentPose, Pose2d poseRef, double linearVelocityRefMeters, Rotation2d angleRef) {
087    // If this is the first run, then we need to reset the theta controller to the current pose's
088    // heading.
089    if (m_firstRun) {
090      m_thetaController.reset(currentPose.getRotation().getRadians());
091      m_firstRun = false;
092    }
093
094    // Calculate feedforward velocities (field-relative).
095    double xFF = linearVelocityRefMeters * poseRef.getRotation().getCos();
096    double yFF = linearVelocityRefMeters * poseRef.getRotation().getSin();
097    double thetaFF =
098        m_thetaController.calculate(currentPose.getRotation().getRadians(), angleRef.getRadians());
099
100    m_poseError = poseRef.relativeTo(currentPose);
101    m_rotationError = angleRef.minus(currentPose.getRotation());
102
103    if (!m_enabled) {
104      return ChassisSpeeds.fromFieldRelativeSpeeds(xFF, yFF, thetaFF, currentPose.getRotation());
105    }
106
107    // Calculate feedback velocities (based on position error).
108    double xFeedback = m_xController.calculate(currentPose.getX(), poseRef.getX());
109    double yFeedback = m_yController.calculate(currentPose.getY(), poseRef.getY());
110
111    // Return next output.
112    return ChassisSpeeds.fromFieldRelativeSpeeds(
113        xFF + xFeedback, yFF + yFeedback, thetaFF, currentPose.getRotation());
114  }
115
116  /**
117   * Returns the next output of the holonomic drive controller.
118   *
119   * @param currentPose The current pose.
120   * @param desiredState The desired trajectory state.
121   * @param angleRef The desired end-angle.
122   * @return The next output of the holonomic drive controller.
123   */
124  public ChassisSpeeds calculate(
125      Pose2d currentPose, Trajectory.State desiredState, Rotation2d angleRef) {
126    return calculate(
127        currentPose, desiredState.poseMeters, desiredState.velocityMetersPerSecond, angleRef);
128  }
129
130  /**
131   * Enables and disables the controller for troubleshooting problems. When calculate() is called on
132   * a disabled controller, only feedforward values are returned.
133   *
134   * @param enabled If the controller is enabled or not.
135   */
136  public void setEnabled(boolean enabled) {
137    m_enabled = enabled;
138  }
139}