swerve_drive_subsystem.h

Go to the documentation of this file.

 
 
#pragma once
 
#include <choreo/trajectory/SwerveSample.h>
#include <frc/Timer.h>
#include <frc/controller/HolonomicDriveController.h>
#include <frc/controller/PIDController.h>
#include <frc/controller/ProfiledPIDController.h>
#include <frc/geometry/Pose2d.h>
#include <frc/geometry/Rotation2d.h>
#include <frc/kinematics/ChassisSpeeds.h>
#include <frc/kinematics/SwerveDriveKinematics.h>
#include <frc/kinematics/SwerveDriveOdometry.h>
#include <frc/kinematics/SwerveModulePosition.h>
#include <frc/kinematics/SwerveModuleState.h>
#include <frc/smartdashboard/Field2d.h>
#include <frc2/command/SubsystemBase.h>
#include <wpi/DataLog.h>
 
#include <array>
#include <memory>
#include <mutex>
#include <thread>
 
#include <ctre/phoenix6/CANcoder.hpp>
#include <ctre/phoenix6/Pigeon2.hpp>
#include <ctre/phoenix6/TalonFX.hpp>
 
#include "argos_lib/config/config_types.h"
#include "argos_lib/general/nt_motor_pid_tuner.h"
#include "argos_lib/general/nt_subscriber.h"
#include "argos_lib/homing/fs_homing.h"
#include "constants/feature_flags.h"
#include "ctre/phoenix6/sim/CANcoderSimState.hpp"
#include "ctre/phoenix6/sim/TalonFXSimState.hpp"
#include "frc/StateSpaceUtil.h"
#include "frc/estimator/SwerveDrivePoseEstimator.h"
#include "frc/geometry/Transform3d.h"
 
class SwerveModule {
 public:
  // MOTORS
  ctre::phoenix6::hardware::TalonFX m_drive;
  ctre::phoenix6::hardware::TalonFX m_turn;
  // ENCODER
  ctre::phoenix6::hardware::CANcoder m_encoder;
 
  SwerveModule(const argos_lib::CANAddress& driveAddr,
               const argos_lib::CANAddress& turnAddr,
               const argos_lib::CANAddress& encoderAddr);
 
  frc::SwerveModuleState GetState();
  frc::SwerveModulePosition GetPosition();
  void SimulationPeriodic(const frc::SwerveModuleState& desiredState, units::second_t dt);
 
 private:
  // Internal state to track simulated drive position (in sensor units).
  units::turn_t m_simDrivePos = 0.0_tr;
};
 
class SwerveDriveSubsystem : public frc2::SubsystemBase {
 public:
  enum DriveControlMode { fieldCentricControl, robotCentricControl };
 
  explicit SwerveDriveSubsystem(const argos_lib::RobotInstance instance);
 
  virtual ~SwerveDriveSubsystem();
 
  void SimulationPeriodic() override;
 
  void Disable();
 
  void SwerveDrive(const double fwVelocity, const double sideVelocity, const double rotVelocity);
 
  void SwerveDrive(const units::degree_t& velAngle, const double& velocity, const double& rotVelocity);
 
  void SwerveDrive(const units::degree_t& velAngle, const double& velocity);
 
  void SwerveDrive(frc::ChassisSpeeds desiredChassisSpeed);
 
  void SwerveDrive(const choreo::SwerveSample& sample);
 
  void StopDrive();
 
  // choreolib::ChoreoControllerFunction GetChoreoControllerFunction();
 
  void Home(const units::degree_t& angle);
 
  void FieldHome(units::degree_t homeAngle = 0_deg, bool updateOdometry = true);
 
  void FlipFieldHome();
 
  void InitializeOdometry(const frc::Pose2d& currentPose);
 
  frc::Rotation2d GetContinuousOdometryAngle();
 
  frc::Rotation2d GetRawOdometryAngle();
 
  frc::Rotation2d GetNearestSquareAngle();
 
  frc::Pose2d GetContinuousOdometry();
 
  frc::Pose2d GetRawOdometry();
 
  void UpdateEstimatedPose();
 
  units::degree_t GetFieldCentricAngle();
 
  frc::Pose2d GetPoseEstimate(const frc::Pose2d& robotPose, const units::millisecond_t& latency);
 
  void UpdateVisionMeasurement(const frc::Pose2d& poseEstimate,
                               units::second_t timestamp,
                               const wpi::array<double, 3>& visionMeasurementStdDevs);
 
  void SetControlMode(SwerveDriveSubsystem::DriveControlMode controlMode);
 
  void InitializeMotors();
 
  void UpdateFollowerLinearPIDParams(double kP, double kI, double kD);
 
  void UpdateFollowerRotationalPIDParams(double kP, double kI, double kD);
 
  void CancelDrivingProfile();
 
  bool ProfileIsComplete() const;
 
  bool IsFollowingProfile() const;
 
  frc::ChassisSpeeds GetChassisVelocity();
 
  void LockWheels();
 
  bool GetManualOverride();
 
  units::degree_t GetRobotPitch() { return m_pigeonIMU.GetRoll().GetValue(); }
 
  units::degrees_per_second_t GetRobotPitchRate();
 
  units::degree_t GetIMUYaw();
  units::degrees_per_second_t GetIMUYawRate();
 
  void SimDrive();
 
  void SetTrajectoryDisplay(const std::vector<frc::Pose2d>& trajectory);
 
 private:
  argos_lib::RobotInstance m_instance;
 
  DriveControlMode m_controlMode;  
 
  SwerveModule m_frontLeft;   
  SwerveModule m_frontRight;  
  SwerveModule m_backRight;   
  SwerveModule m_backLeft;    
 
  // GYROSCOPIC SENSORS
  ctre::phoenix6::hardware::Pigeon2 m_pigeonIMU;
 
  units::degree_t m_fieldHomeOffset;  
 
  struct Velocities {
    const double fwVelocity;
    const double sideVelocity;
    const double rotVelocity;
  };
 
  frc::SwerveDriveKinematics<4> m_swerveDriveKinematics;  
 
  frc::SwerveDriveOdometry<4> m_odometry;      
  units::degree_t m_prevOdometryAngle;         
  units::degree_t m_continuousOdometryOffset;  
 
  frc::SwerveDrivePoseEstimator<4> m_poseEstimator;  
  std::mutex m_poseEstimatorLock;
  std::thread m_odometryThread;  
  std::chrono::time_point<std::chrono::steady_clock>
      m_odometryResetTime;  
  bool m_stillRunning;      
 
  // std::FILE SYSTEM HOMING STORAGE
  argos_lib::SwerveFSHomingStorage m_fsStorage;  
 
  bool m_followingProfile;  
  bool m_profileComplete;   
  bool m_manualOverride;
  std::chrono::time_point<std::chrono::steady_clock> m_swerveProfileStartTime;  
  frc::PIDController m_xPID;      
  frc::PIDController m_yPID;      
  frc::PIDController m_thetaPID;  
 
  std::unique_ptr<argos_lib::NTMotorPIDTuner> m_driveMotorPIDTuner;  
  std::unique_ptr<argos_lib::NTSubscriber> m_linearFollowerTuner_P;
  std::unique_ptr<argos_lib::NTSubscriber> m_linearFollowerTuner_I;
  std::unique_ptr<argos_lib::NTSubscriber> m_linearFollowerTuner_D;
  std::unique_ptr<argos_lib::NTSubscriber> m_rotationalFollowerTuner_P;
  std::unique_ptr<argos_lib::NTSubscriber> m_rotationalFollowerTuner_I;
  std::unique_ptr<argos_lib::NTSubscriber> m_rotationalFollowerTuner_D;
 
  wpi::array<frc::SwerveModuleState, 4> GetRawModuleStates(frc::ChassisSpeeds velocities);
  wpi::array<frc::SwerveModuleState, 4> GetRawModuleStates(SwerveDriveSubsystem::Velocities velocities);
 
  wpi::array<frc::SwerveModuleState, 4> GetCurrentModuleStates();
 
  wpi::array<frc::SwerveModulePosition, 4> GetCurrentModulePositions();
 
  void HomeToFS(const units::degree_t& angle);
 
  void InitializeMotorsFromFS();
 
  void ResetIMUYaw();
 
  wpi::array<frc::SwerveModuleState, 4> OptimizeAllModules(wpi::array<frc::SwerveModuleState, 4> rawStates);
 
  void ClosedLoopDrive(wpi::array<frc::SwerveModuleState, 4> moduleStates);
 
  wpi::log::StructLogEntry<frc::Pose2d> m_poseEstimateLogger;
  wpi::log::StructArrayLogEntry<frc::SwerveModuleState> m_setpointLogger;
  wpi::log::StructArrayLogEntry<frc::SwerveModuleState> m_stateLogger;
 
  struct SimVelocities {
    double fwVelocity;
    double sideVelocity;
    double rotVelocity;
  };
  SimVelocities m_simVelocities;
 
  double m_simulatedHeading;
 
  // Field2D for visualization
  frc::Field2d m_field;
};