LQRController.hpp

Go to the documentation of this file.

// Copyright (c) 2026 Yet Another Software Suite
// SPDX-License-Identifier: LGPL-3.0-or-later
 
#pragma once
 
#include <memory>
#include <optional>
#include <utility>
#include <variant>
#include <wpi/math/controller/LinearQuadraticRegulator.hpp>
#include <wpi/math/estimator/KalmanFilter.hpp>
#include <wpi/math/system/LinearSystem.hpp>
#include <wpi/units/angle.hpp>
#include <wpi/units/angular_velocity.hpp>
#include <wpi/units/length.hpp>
#include <wpi/units/time.hpp>
#include <wpi/units/velocity.hpp>
#include <wpi/units/voltage.hpp>
 
#include "LQRConfig.hpp"
 
namespace yams::math {
 
class LQRController {
 public:
  explicit LQRController(LQRConfig config);
 
  void UpdateConfig(LQRConfig config);
 
  void Reset(wpi::units::radian_t angle, wpi::units::radians_per_second_t velocity);
 
  void Reset(wpi::units::meter_t distance, wpi::units::meters_per_second_t velocity);
 
  wpi::units::volt_t Calculate(wpi::units::radian_t measured, wpi::units::radian_t position,
                               wpi::units::radians_per_second_t velocity);
 
  wpi::units::volt_t Calculate(wpi::units::meter_t measured, wpi::units::meter_t position,
                               wpi::units::meters_per_second_t velocity);
 
  wpi::units::volt_t Calculate(wpi::units::radians_per_second_t measured,
                               wpi::units::radians_per_second_t velocity);
 
  wpi::units::volt_t Calculate(wpi::units::meters_per_second_t measured,
                               wpi::units::meters_per_second_t velocity);
 
  LQRConfig::LQRType GetType() const;
 
  const std::optional<LQRConfig>& GetConfig() const;
 
 private:
  // wpi::math::LinearSystemLoop stores raw pointers to the LQR and Kalman filter passed to its
  // constructor. These bundles heap-pin all three together so the pointers are never dangled.
  struct FlywheelBundle {
    wpi::math::LinearQuadraticRegulator<1, 1> controller;
    wpi::math::KalmanFilter<1, 1, 1> observer;
    LQRConfig::Loop1 loop;
 
    FlywheelBundle(wpi::math::LinearSystem<1, 1, 1> plant,
                   wpi::math::LinearQuadraticRegulator<1, 1> ctrl,
                   wpi::math::KalmanFilter<1, 1, 1> obs, wpi::units::volt_t maxV,
                   wpi::units::second_t dt)
        : controller{std::move(ctrl)},
          observer{std::move(obs)},
          loop{plant, controller, observer, maxV, dt} {}
 
    FlywheelBundle(const FlywheelBundle&) = delete;
    FlywheelBundle& operator=(const FlywheelBundle&) = delete;
    FlywheelBundle(FlywheelBundle&&) = delete;
    FlywheelBundle& operator=(FlywheelBundle&&) = delete;
  };
 
  struct ArmElevatorBundle {
    wpi::math::LinearQuadraticRegulator<2, 1> controller;
    wpi::math::KalmanFilter<2, 1, 1> observer;
    LQRConfig::Loop2 loop;
 
    ArmElevatorBundle(wpi::math::LinearSystem<2, 1, 1> plant,
                      wpi::math::LinearQuadraticRegulator<2, 1> ctrl,
                      wpi::math::KalmanFilter<2, 1, 1> obs, wpi::units::volt_t maxV,
                      wpi::units::second_t dt)
        : controller{std::move(ctrl)},
          observer{std::move(obs)},
          loop{plant, controller, observer, maxV, dt} {}
 
    ArmElevatorBundle(const ArmElevatorBundle&) = delete;
    ArmElevatorBundle& operator=(const ArmElevatorBundle&) = delete;
    ArmElevatorBundle(ArmElevatorBundle&&) = delete;
    ArmElevatorBundle& operator=(ArmElevatorBundle&&) = delete;
  };
 
  static std::unique_ptr<FlywheelBundle> BuildFlywheel(const LQRConfig& cfg);
  static std::unique_ptr<ArmElevatorBundle> BuildArmElevator(const LQRConfig& cfg);
 
  std::optional<LQRConfig> m_config;
  LQRConfig::LQRType m_type;
  std::variant<std::unique_ptr<FlywheelBundle>, std::unique_ptr<ArmElevatorBundle>> m_bundle;
  wpi::units::second_t m_period;
};
 
}  // namespace yams::math