frc971/control_loops/state_feedback_loop.h

#ifndef FRC971_CONTROL_LOOPS_STATEFEEDBACKLOOP_H_
#define FRC971_CONTROL_LOOPS_STATEFEEDBACKLOOP_H_

#include <assert.h>

#include <vector>
#include <iostream>

#include "Eigen/Dense"

#include "aos/common/logging/logging.h"
#include "aos/common/macros.h"

// TODO(brians): There are a lot of public member variables in here that should
// be made private and have (const) accessors instead (and have _s at the end of
// their names).

template <int number_of_states, int number_of_inputs, int number_of_outputs>
class StateFeedbackPlantCoefficients {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

  const Eigen::Matrix<double, number_of_states, number_of_states> A;
  const Eigen::Matrix<double, number_of_states, number_of_inputs> B;
  const Eigen::Matrix<double, number_of_outputs, number_of_states> C;
  const Eigen::Matrix<double, number_of_outputs, number_of_inputs> D;
  const Eigen::Matrix<double, number_of_inputs, 1> U_min;
  const Eigen::Matrix<double, number_of_inputs, 1> U_max;

  StateFeedbackPlantCoefficients(const StateFeedbackPlantCoefficients &other)
      : A(other.A),
        B(other.B),
        C(other.C),
        D(other.D),
        U_min(other.U_min),
        U_max(other.U_max) {
  }

  StateFeedbackPlantCoefficients(
      const Eigen::Matrix<double, number_of_states, number_of_states> &A,
      const Eigen::Matrix<double, number_of_states, number_of_inputs> &B,
      const Eigen::Matrix<double, number_of_outputs, number_of_states> &C,
      const Eigen::Matrix<double, number_of_outputs, number_of_inputs> &D,
      const Eigen::Matrix<double, number_of_outputs, 1> &U_max,
      const Eigen::Matrix<double, number_of_outputs, 1> &U_min)
      : A(A),
        B(B),
        C(C),
        D(D),
        U_min(U_min),
        U_max(U_max) {
  }

 protected:
  // these are accessible from non-templated subclasses
  static const int kNumStates = number_of_states;
  static const int kNumOutputs = number_of_outputs;
  static const int kNumInputs = number_of_inputs;
};

template <int number_of_states, int number_of_inputs, int number_of_outputs>
class StateFeedbackPlant {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

  StateFeedbackPlant(
      const ::std::vector<StateFeedbackPlantCoefficients<
          number_of_states, number_of_inputs,
          number_of_outputs> *> &coefficients)
      : coefficients_(coefficients),
        plant_index_(0) {
    Reset();
  }

  StateFeedbackPlant(StateFeedbackPlant &&other)
      : plant_index_(other.plant_index_) {
    ::std::swap(coefficients_, other.coefficients_);
    X.swap(other.X);
    Y.swap(other.Y);
    U.swap(other.U);
  }

  virtual ~StateFeedbackPlant() {
    for (auto c : coefficients_) {
      delete c;
    }
  }

  ::std::vector<StateFeedbackPlantCoefficients<
      number_of_states, number_of_inputs, number_of_outputs> *> coefficients_;

  const Eigen::Matrix<double, number_of_states, number_of_states> &A() const {
    return coefficients().A;
  }
  double A(int i, int j) const { return A()(i, j); }
  const Eigen::Matrix<double, number_of_states, number_of_inputs> &B() const {
    return coefficients().B;
  }
  double B(int i, int j) const { return B()(i, j); }
  const Eigen::Matrix<double, number_of_outputs, number_of_states> &C() const {
    return coefficients().C;
  }
  double C(int i, int j) const { return C()(i, j); }
  const Eigen::Matrix<double, number_of_outputs, number_of_inputs> &D() const {
    return coefficients().D;
  }
  double D(int i, int j) const { return D()(i, j); }
  const Eigen::Matrix<double, number_of_inputs, 1> &U_min() const {
    return coefficients().U_min;
  }
  double U_min(int i, int j) const { return U_min()(i, j); }
  const Eigen::Matrix<double, number_of_inputs, 1> &U_max() const {
    return coefficients().U_max;
  }
  double U_max(int i, int j) const { return U_max()(i, j); }

  const StateFeedbackPlantCoefficients<
      number_of_states, number_of_inputs, number_of_outputs>
          &coefficients() const {
    return *coefficients_[plant_index_];
  }

  int plant_index() const { return plant_index_; }
  void set_plant_index(int plant_index) {
    if (plant_index < 0) {
      plant_index_ = 0;
    } else if (plant_index >= static_cast<int>(coefficients_.size())) {
      plant_index_ = static_cast<int>(coefficients_.size()) - 1;
    } else {
      plant_index_ = plant_index;
    }
  }

  void Reset() {
    X.setZero();
    Y.setZero();
    U.setZero();
  }

  Eigen::Matrix<double, number_of_states, 1> X;
  Eigen::Matrix<double, number_of_outputs, 1> Y;
  Eigen::Matrix<double, number_of_inputs, 1> U;

  // Assert that U is within the hardware range.
  virtual void CheckU() {
    for (int i = 0; i < kNumOutputs; ++i) {
      assert(U(i, 0) <= U_max(i, 0) + 0.00001);
      assert(U(i, 0) >= U_min(i, 0) - 0.00001);
    }
  }

  // Computes the new X and Y given the control input.
  void Update() {
    // Powers outside of the range are more likely controller bugs than things
    // that the plant should deal with.
    CheckU();
    X = A() * X + B() * U;
    Y = C() * X + D() * U;
  }

 protected:
  // these are accessible from non-templated subclasses
  static const int kNumStates = number_of_states;
  static const int kNumOutputs = number_of_outputs;
  static const int kNumInputs = number_of_inputs;

 private:
  int plant_index_;

  DISALLOW_COPY_AND_ASSIGN(StateFeedbackPlant);
};

// A Controller is a structure which holds a plant and the K and L matrices.
// This is designed such that multiple controllers can share one set of state to
// support gain scheduling easily.
template <int number_of_states, int number_of_inputs, int number_of_outputs>
struct StateFeedbackController {
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
  const Eigen::Matrix<double, number_of_states, number_of_outputs> L;
  const Eigen::Matrix<double, number_of_outputs, number_of_states> K;
  StateFeedbackPlantCoefficients<number_of_states, number_of_inputs,
                                 number_of_outputs> plant;

  StateFeedbackController(
      const Eigen::Matrix<double, number_of_states, number_of_outputs> &L,
      const Eigen::Matrix<double, number_of_outputs, number_of_states> &K,
      const StateFeedbackPlantCoefficients<number_of_states, number_of_inputs,
                                           number_of_outputs> &plant)
      : L(L),
        K(K),
        plant(plant) {
  }
};

template <int number_of_states, int number_of_inputs, int number_of_outputs>
class StateFeedbackLoop {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

  StateFeedbackLoop(const StateFeedbackController<
      number_of_states, number_of_inputs, number_of_outputs> &controller)
      : controller_index_(0) {
    controllers_.push_back(new StateFeedbackController<
        number_of_states, number_of_inputs, number_of_outputs>(controller));
    Reset();
  }

  StateFeedbackLoop(const ::std::vector<StateFeedbackController<
      number_of_states, number_of_inputs, number_of_outputs> *> &controllers)
      : controllers_(controllers), controller_index_(0) {
    Reset();
  }

  StateFeedbackLoop(StateFeedbackLoop &&other) {
    X_hat.swap(other.X_hat);
    R.swap(other.R);
    U.swap(other.U);
    U_uncapped.swap(other.U_uncapped);
    U_ff.swap(other.U_ff);
    ::std::swap(controllers_, other.controllers_);
    Y_.swap(other.Y_);
    new_y_ = other.new_y_;
    controller_index_ = other.controller_index_;
  }

  StateFeedbackLoop(
      const Eigen::Matrix<double, number_of_states, number_of_outputs> &L,
      const Eigen::Matrix<double, number_of_outputs, number_of_states> &K,
      const StateFeedbackPlantCoefficients<number_of_states, number_of_inputs,
                               number_of_outputs> &plant)
      : controller_index_(0) {
    controllers_.push_back(
        new StateFeedbackController<number_of_states, number_of_inputs,
                                    number_of_outputs>(L, K, plant));

    Reset();
  }

  virtual ~StateFeedbackLoop() {
    for (auto c : controllers_) {
      delete c;
    }
  }

  const Eigen::Matrix<double, number_of_states, number_of_states> &A() const {
    return controller().plant.A;
  }
  double A(int i, int j) const { return A()(i, j); }
  const Eigen::Matrix<double, number_of_states, number_of_inputs> &B() const {
    return controller().plant.B;
  }
  double B(int i, int j) const { return B()(i, j); }
  const Eigen::Matrix<double, number_of_outputs, number_of_states> &C() const {
    return controller().plant.C;
  }
  double C(int i, int j) const { return C()(i, j); }
  const Eigen::Matrix<double, number_of_outputs, number_of_inputs> &D() const {
    return controller().plant.D;
  }
  double D(int i, int j) const { return D()(i, j); }
  const Eigen::Matrix<double, number_of_outputs, number_of_states> &K() const {
    return controller().K;
  }
  double K(int i, int j) const { return K()(i, j); }
  const Eigen::Matrix<double, number_of_states, number_of_outputs> &L() const {
    return controller().L;
  }
  double L(int i, int j) const { return L()(i, j); }
  const Eigen::Matrix<double, number_of_inputs, 1> &U_min() const {
    return controller().plant.U_min;
  }
  double U_min(int i, int j) const { return U_min()(i, j); }
  const Eigen::Matrix<double, number_of_inputs, 1> &U_max() const {
    return controller().plant.U_max;
  }
  double U_max(int i, int j) const { return U_max()(i, j); }

  const StateFeedbackController<number_of_states, number_of_inputs,
                                number_of_outputs> &controller() const {
    return *controllers_[controller_index_];
  }

  const StateFeedbackController<number_of_states, number_of_inputs,
                                number_of_outputs> &controller(
                                    int index) const {
    return *controllers_[index];
  }

  void Reset() {
    X_hat.setZero();
    R.setZero();
    U.setZero();
    U_uncapped.setZero();
    U_ff.setZero();
  }

  virtual void FeedForward() {
    for (int i = 0; i < number_of_outputs; ++i) {
      U_ff[i] = 0.0;
    }
  }

  // If U is outside the hardware range, limit it before the plant tries to use
  // it.
  virtual void CapU() {
    for (int i = 0; i < kNumOutputs; ++i) {
      if (U(i, 0) > U_max(i, 0)) {
        U(i, 0) = U_max(i, 0);
      } else if (U(i, 0) < U_min(i, 0)) {
        U(i, 0) = U_min(i, 0);
      }
    }
  }

  // Corrects X_hat given the observation in Y.
  void Correct(const Eigen::Matrix<double, number_of_outputs, 1> &Y) {
  /*
    auto eye =
        Eigen::Matrix<double, number_of_states, number_of_states>::Identity();
    //LOG(DEBUG, "X_hat(2, 0) = %f\n", X_hat(2, 0));
    ::std::cout << "Identity " << eye << ::std::endl;
    ::std::cout << "X_hat " << X_hat << ::std::endl;
    ::std::cout << "LC " << L() * C() << ::std::endl;
    ::std::cout << "L " << L() << ::std::endl;
    ::std::cout << "C " << C() << ::std::endl;
    ::std::cout << "y " << Y << ::std::endl;
    ::std::cout << "z " << (Y - C() * X_hat) << ::std::endl;
    ::std::cout << "correction " << L() * (Y - C() * X_hat) << ::std::endl;
    X_hat = (eye - L() * C()) * X_hat + L() * Y;
    ::std::cout << "X_hat after " << X_hat << ::std::endl;
    ::std::cout << ::std::endl;
    */
    //LOG(DEBUG, "X_hat(2, 0) = %f\n", X_hat(2, 0));
    Y_ = Y;
    new_y_ = true;
  }

  // stop_motors is whether or not to output all 0s.
  void Update(bool stop_motors) {
    if (stop_motors) {
      U.setZero();
      U_uncapped.setZero();
    } else {
      U = U_uncapped = K() * (R - X_hat);
      CapU();
    }

    //::std::cout << "Predict xhat before " << X_hat;
    //::std::cout << "Measurement error is " << Y_ - C() * X_hat;
    //X_hat = A() * X_hat + B() * U;
    //::std::cout << "Predict xhat after " << X_hat;
    UpdateObserver();
  }

  void UpdateObserver() {
    if (new_y_) {
      X_hat = (A() - L() * C()) * X_hat + L() * Y_ + B() * U;
      new_y_ = false;
    } else {
      X_hat = A() * X_hat + B() * U;
    }
  }

  // Sets the current controller to be index and verifies that it isn't out of
  // range.
  void set_controller_index(int index) {
    if (index < 0) {
      controller_index_ = 0;
    } else if (index >= static_cast<int>(controllers_.size())) {
      controller_index_ = static_cast<int>(controllers_.size()) - 1;
    } else {
      controller_index_ = index;
    }
  }

  int controller_index() const { return controller_index_; }

  Eigen::Matrix<double, number_of_states, 1> X_hat;
  Eigen::Matrix<double, number_of_states, 1> R;
  Eigen::Matrix<double, number_of_inputs, 1> U;
  Eigen::Matrix<double, number_of_inputs, 1> U_uncapped;
  Eigen::Matrix<double, number_of_outputs, 1> U_ff;

 protected:
  ::std::vector<StateFeedbackController<number_of_states, number_of_inputs,
                                        number_of_outputs> *> controllers_;

  // these are accessible from non-templated subclasses
  static const int kNumStates = number_of_states;
  static const int kNumOutputs = number_of_outputs;
  static const int kNumInputs = number_of_inputs;

  // Temporary storage for a measurement until I can figure out why I can't
  // correct when the measurement is made.
  Eigen::Matrix<double, number_of_outputs, 1> Y_;
  bool new_y_ = false;

  int controller_index_;

 private:
  DISALLOW_COPY_AND_ASSIGN(StateFeedbackLoop);
};

#endif  // FRC971_CONTROL_LOOPS_STATEFEEDBACKLOOP_H_