y2022/control_loops/superstructure/catapult/catapult.cc

#include "y2022/control_loops/superstructure/catapult/catapult.h"

#include "Eigen/Dense"
#include "Eigen/Sparse"
#include "glog/logging.h"

#include "aos/realtime.h"
#include "aos/time/time.h"
#include "osqp++.h"
#include "osqp.h"
#include "y2022/control_loops/superstructure/catapult/catapult_plant.h"

namespace y2022::control_loops::superstructure::catapult {
namespace chrono = std::chrono;

namespace {
osqp::OsqpInstance MakeInstance(
    size_t horizon, Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> P) {
  osqp::OsqpInstance instance;
  instance.objective_matrix = P.sparseView();

  instance.constraint_matrix =
      Eigen::SparseMatrix<double, Eigen::ColMajor, osqp::c_int>(horizon,
                                                                horizon);
  instance.constraint_matrix.setIdentity();

  instance.lower_bounds =
      Eigen::Matrix<double, Eigen::Dynamic, 1>::Zero(horizon, 1);
  instance.upper_bounds =
      Eigen::Matrix<double, Eigen::Dynamic, 1>::Ones(horizon, 1) * 12.0;
  return instance;
}
}  // namespace

MPCProblem::MPCProblem(size_t horizon,
                       Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> P,
                       Eigen::Matrix<double, Eigen::Dynamic, 1> accel_q,
                       Eigen::Matrix<double, 2, 2> Af,
                       Eigen::Matrix<double, Eigen::Dynamic, 2> final_q)
    : horizon_(horizon),
      accel_q_(std::move(accel_q)),
      Af_(std::move(Af)),
      final_q_(std::move(final_q)),
      instance_(MakeInstance(horizon, std::move(P))) {
  // Start with a representative problem.
  Eigen::Matrix<double, 2, 1> X_initial(0.0, 0.0);
  Eigen::Matrix<double, 2, 1> X_final(2.0, 25.0);

  objective_vector_ =
      X_initial(1, 0) * accel_q_ + final_q_ * (Af_ * X_initial - X_final);
  instance_.objective_vector = objective_vector_;
  settings_.max_iter = 25;
  settings_.check_termination = 5;
  settings_.warm_start = 1;
  // TODO(austin): Do we need this scaling thing?  It makes it not solve
  // sometimes... I'm pretty certain by giving it a decently formed problem to
  // initialize with, it will not try doing crazy things with the scaling
  // internally.
  settings_.scaling = 0;
  auto status = solver_.Init(instance_, settings_);
  CHECK(status.ok()) << status;
}

void MPCProblem::SetState(Eigen::Matrix<double, 2, 1> X_initial,
                          Eigen::Matrix<double, 2, 1> X_final) {
  X_initial_ = X_initial;
  X_final_ = X_final;
  // If we mark this noalias(), it won't re-allocate the vector each time.
  objective_vector_.noalias() =
      X_initial(1, 0) * accel_q_ + final_q_ * (Af_ * X_initial - X_final);

  auto status = solver_.SetObjectiveVector(objective_vector_);
  CHECK(status.ok()) << status;
}

bool MPCProblem::Solve() {
  const aos::monotonic_clock::time_point start_time =
      aos::monotonic_clock::now();
  osqp::OsqpExitCode exit_code = solver_.Solve();
  const aos::monotonic_clock::time_point end_time = aos::monotonic_clock::now();
  VLOG(1) << "OSQP solved in "
          << std::chrono::duration<double>(end_time - start_time).count();
  solve_time_ = std::chrono::duration<double>(end_time - start_time).count();
  // TODO(austin): Dump the exit codes out as an enum for logging.
  //
  // TODO(austin): The dual problem doesn't appear to be converging on all
  // problems.  Are we phrasing something wrong?

  // TODO(austin): Set a time limit so we can't run forever, and signal back
  // when we hit our limit.
  return exit_code == osqp::OsqpExitCode::kOptimal;
}

void MPCProblem::WarmStart(const MPCProblem &p) {
  CHECK_GE(p.horizon(), horizon())
      << ": Can only copy a bigger problem's solution into a smaller problem.";
  auto status = solver_.SetPrimalWarmStart(p.solver_.primal_solution().block(
      p.horizon() - horizon(), 0, horizon(), 1));
  CHECK(status.ok()) << status;
  status = solver_.SetDualWarmStart(p.solver_.dual_solution().block(
      p.horizon() - horizon(), 0, horizon(), 1));
  CHECK(status.ok()) << status;
}

CatapultProblemGenerator::CatapultProblemGenerator(size_t horizon)
    : plant_(MakeCatapultPlant()),
      horizon_(horizon),
      Q_final_(
          (Eigen::DiagonalMatrix<double, 2>().diagonal() << 10000.0, 10000.0)
              .finished()),
      As_(MakeAs()),
      Bs_(MakeBs()),
      m_(Makem()),
      M_(MakeM()),
      W_(MakeW()),
      w_(Makew()),
      Pi_(MakePi()),
      WM_(W_ * M_),
      Wmpw_(W_ * m_ + w_) {}

std::unique_ptr<MPCProblem> CatapultProblemGenerator::MakeProblem(
    size_t horizon) {
  return std::make_unique<MPCProblem>(
      horizon, P(horizon), accel_q(horizon), Af(horizon),
      (2.0 * Q_final_ * Bf(horizon)).transpose());
}

const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
CatapultProblemGenerator::P(size_t horizon) {
  CHECK_GT(horizon, 0u);
  CHECK_LE(horizon, horizon_);
  return 2.0 * (WM_.block(0, 0, horizon, horizon).transpose() * Pi(horizon) *
                    WM_.block(0, 0, horizon, horizon) +
                Bf(horizon).transpose() * Q_final_ * Bf(horizon));
}

const Eigen::Matrix<double, Eigen::Dynamic, 1> CatapultProblemGenerator::q(
    size_t horizon, Eigen::Matrix<double, 2, 1> X_initial,
    Eigen::Matrix<double, 2, 1> X_final) {
  CHECK_GT(horizon, 0u);
  CHECK_LE(horizon, horizon_);
  return 2.0 * X_initial(1, 0) * accel_q(horizon) +
         2.0 * ((Af(horizon) * X_initial - X_final).transpose() * Q_final_ *
                Bf(horizon))
                   .transpose();
}

const Eigen::Matrix<double, Eigen::Dynamic, 1>
CatapultProblemGenerator::accel_q(size_t horizon) {
  return 2.0 * ((Wmpw_.block(0, 0, horizon, 1)).transpose() * Pi(horizon) *
                WM_.block(0, 0, horizon, horizon))
                   .transpose();
}

const Eigen::Matrix<double, 2, 2> CatapultProblemGenerator::Af(size_t horizon) {
  CHECK_GT(horizon, 0u);
  CHECK_LE(horizon, horizon_);
  return As_.block<2, 2>(2 * (horizon - 1), 0);
}

const Eigen::Matrix<double, 2, Eigen::Dynamic> CatapultProblemGenerator::Bf(
    size_t horizon) {
  CHECK_GT(horizon, 0u);
  CHECK_LE(horizon, horizon_);
  return Bs_.block(2 * (horizon - 1), 0, 2, horizon);
}

const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
CatapultProblemGenerator::Pi(size_t horizon) {
  CHECK_GT(horizon, 0u);
  CHECK_LE(horizon, horizon_);
  return Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>(Pi_).block(
      horizon_ - horizon, horizon_ - horizon, horizon, horizon);
}

Eigen::Matrix<double, Eigen::Dynamic, 2> CatapultProblemGenerator::MakeAs() {
  Eigen::Matrix<double, Eigen::Dynamic, 2> As =
      Eigen::Matrix<double, Eigen::Dynamic, 2>::Zero(horizon_ * 2, 2);
  for (size_t i = 0; i < horizon_; ++i) {
    if (i == 0) {
      As.block<2, 2>(0, 0) = plant_.A();
    } else {
      As.block<2, 2>(i * 2, 0) = plant_.A() * As.block<2, 2>((i - 1) * 2, 0);
    }
  }
  return As;
}

Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
CatapultProblemGenerator::MakeBs() {
  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Bs =
      Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>::Zero(horizon_ * 2,
                                                                  horizon_);
  for (size_t i = 0; i < horizon_; ++i) {
    for (size_t j = 0; j < i + 1; ++j) {
      if (i == j) {
        Bs.block<2, 1>(i * 2, j) = plant_.B();
      } else {
        Bs.block<2, 1>(i * 2, j) =
            As_.block<2, 2>((i - j - 1) * 2, 0) * plant_.B();
      }
    }
  }
  return Bs;
}

Eigen::Matrix<double, Eigen::Dynamic, 1> CatapultProblemGenerator::Makem() {
  Eigen::Matrix<double, Eigen::Dynamic, 1> m =
      Eigen::Matrix<double, Eigen::Dynamic, 1>::Zero(horizon_, 1);
  for (size_t i = 0; i < horizon_; ++i) {
    m(i, 0) = As_(1 + 2 * i, 1);
  }
  return m;
}

Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
CatapultProblemGenerator::MakeM() {
  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> M =
      Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>::Zero(horizon_,
                                                                  horizon_);
  for (size_t i = 0; i < horizon_; ++i) {
    for (size_t j = 0; j < horizon_; ++j) {
      M(i, j) = Bs_(2 * i + 1, j);
    }
  }
  return M;
}

Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
CatapultProblemGenerator::MakeW() {
  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> W =
      Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>::Identity(horizon_,
                                                                      horizon_);
  for (size_t i = 0; i < horizon_ - 1; ++i) {
    W(i + 1, i) = -1.0;
  }
  W /= std::chrono::duration<double>(plant_.dt()).count();
  return W;
}

Eigen::Matrix<double, Eigen::Dynamic, 1> CatapultProblemGenerator::Makew() {
  Eigen::Matrix<double, Eigen::Dynamic, 1> w =
      Eigen::Matrix<double, Eigen::Dynamic, 1>::Zero(horizon_, 1);
  w(0, 0) = -1.0 / std::chrono::duration<double>(plant_.dt()).count();
  return w;
}

Eigen::DiagonalMatrix<double, Eigen::Dynamic>
CatapultProblemGenerator::MakePi() {
  Eigen::DiagonalMatrix<double, Eigen::Dynamic> Pi(horizon_);
  for (size_t i = 0; i < horizon_; ++i) {
    Pi.diagonal()(i) =
        std::pow(0.01, 2.0) +
        std::pow(0.02 * std::max(0.0, (20 - ((int)horizon_ - (int)i)) / 20.),
                 2.0);
  }
  return Pi;
}

Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>
CatapultProblemGenerator::MakeP() {
  return 2.0 * (M_.transpose() * W_.transpose() * Pi_ * W_ * M_ +
                Bf(horizon_).transpose() * Q_final_ * Bf(horizon_));
}

CatapultController::CatapultController(size_t horizon) : generator_(horizon) {
  problems_.reserve(generator_.horizon());
  for (size_t i = generator_.horizon(); i > 0; --i) {
    problems_.emplace_back(generator_.MakeProblem(i));
  }

  Reset();
}

void CatapultController::Reset() {
  current_controller_ = 0;
  solve_time_ = 0.0;
}

void CatapultController::SetState(Eigen::Matrix<double, 2, 1> X_initial,
                                  Eigen::Matrix<double, 2, 1> X_final) {
  if (current_controller_ >= problems_.size()) {
    return;
  }
  problems_[current_controller_]->SetState(X_initial, X_final);
}

bool CatapultController::Solve() {
  if (current_controller_ >= problems_.size()) {
    return true;
  }
  const bool result = problems_[current_controller_]->Solve();
  solve_time_ = problems_[current_controller_]->solve_time();
  return result;
}

std::optional<double> CatapultController::Next() {
  if (current_controller_ >= problems_.size()) {
    return std::nullopt;
  }

  double u;
  size_t solution_number = 0;
  if (current_controller_ == 0u) {
    while (solution_number < problems_[current_controller_]->horizon() &&
           problems_[current_controller_]->U(solution_number) < 0.01) {
      u = problems_[current_controller_]->U(solution_number);
      ++solution_number;
    }
  }
  u = problems_[current_controller_]->U(solution_number);

  if (current_controller_ + 1u + solution_number < problems_.size()) {
    problems_[current_controller_ + solution_number + 1]->WarmStart(
        *problems_[current_controller_]);
  }
  current_controller_ += 1u + solution_number;
  return u;
}

const flatbuffers::Offset<
    frc971::control_loops::PotAndAbsoluteEncoderProfiledJointStatus>
Catapult::Iterate(const CatapultGoal *catapult_goal, const Position *position,
                  double battery_voltage, double *catapult_voltage, bool fire,
                  flatbuffers::FlatBufferBuilder *fbb) {
  const frc971::control_loops::StaticZeroingSingleDOFProfiledSubsystemGoal
      *return_goal =
          catapult_goal != nullptr && catapult_goal->has_return_position()
              ? catapult_goal->return_position()
              : nullptr;

  const bool catapult_disabled = catapult_.Correct(
      return_goal, position->catapult(), catapult_voltage == nullptr);

  if (catapult_disabled) {
    catapult_state_ = CatapultState::PROFILE;
  } else if (catapult_.running() && catapult_goal != nullptr && fire &&
             !last_firing_) {
    catapult_state_ = CatapultState::FIRING;
    latched_shot_position = catapult_goal->shot_position();
    latched_shot_velocity = catapult_goal->shot_velocity();
  }

  // Don't update last_firing_ if the catapult is disabled, so that we actually
  // end up firing once it's enabled
  if (catapult_.running() && !catapult_disabled) {
    last_firing_ = fire;
  }

  use_profile_ = true;

  switch (catapult_state_) {
    case CatapultState::FIRING: {
      // Select the ball controller.  We should only be firing if we have a
      // ball, or at least should only care about the shot accuracy.
      catapult_.set_controller_index(0);
      // Ok, so we've now corrected.  Next step is to run the MPC.
      //
      // Since there is a unit delay between when we ask for a U and the
      // hardware applies it, we need to run the optimizer for the position at
      // the *next* control loop cycle.

      Eigen::Vector3d next_X = catapult_.estimated_state();
      for (int i = catapult_.controller().plant().coefficients().delayed_u;
           i > 1; --i) {
        next_X = catapult_.controller().plant().A() * next_X +
                 catapult_.controller().plant().B() *
                     catapult_.controller().observer().last_U(i - 1);
      }

      catapult_mpc_.SetState(
          next_X.block<2, 1>(0, 0),
          Eigen::Vector2d(latched_shot_position, latched_shot_velocity));

      const bool solved = catapult_mpc_.Solve();
      current_horizon_ = catapult_mpc_.current_horizon();
      const bool started = catapult_mpc_.started();
      if (solved || started) {
        std::optional<double> solution = catapult_mpc_.Next();

        if (!solution.has_value()) {
          CHECK_NOTNULL(catapult_voltage);
          *catapult_voltage = 0.0;
          if (catapult_mpc_.started()) {
            ++shot_count_;
            // Finished the catapult, time to fire.
            catapult_state_ = CatapultState::RESETTING;
          }
        } else {
          // TODO(austin): Voltage error?
          CHECK_NOTNULL(catapult_voltage);
          if (current_horizon_ == 1) {
            battery_voltage = 12.0;
          }
          *catapult_voltage = std::max(
              0.0, std::min(12.0, (*solution - 0.0 * next_X(2, 0)) * 12.0 /
                                      std::max(battery_voltage, 8.0)));
          use_profile_ = false;
        }
      } else {
        if (!fire) {
          // Eh, didn't manage to solve before it was time to fire.  Give up.
          catapult_state_ = CatapultState::PROFILE;
        }
      }

      if (!use_profile_) {
        catapult_.ForceGoal(catapult_.estimated_position(),
                            catapult_.estimated_velocity());
      }
    }
      if (catapult_state_ != CatapultState::RESETTING) {
        break;
      } else {
        [[fallthrough]];
      }

    case CatapultState::RESETTING:
      if (catapult_.controller().R(1, 0) > 7.0) {
        catapult_.AdjustProfile(7.0, 2000.0);
      } else if (catapult_.controller().R(1, 0) > 0.0) {
        catapult_.AdjustProfile(7.0, 1000.0);
      } else {
        catapult_state_ = CatapultState::PROFILE;
      }
      [[fallthrough]];

    case CatapultState::PROFILE:
      break;
  }

  if (use_profile_) {
    if (catapult_state_ != CatapultState::FIRING) {
      catapult_mpc_.Reset();
    }
    // Select the controller designed for when we have no ball.
    catapult_.set_controller_index(1);

    current_horizon_ = 0u;
    const double output_voltage = catapult_.UpdateController(catapult_disabled);
    if (catapult_voltage != nullptr) {
      *catapult_voltage = output_voltage;
    }
  }

  catapult_.UpdateObserver(catapult_voltage != nullptr
                               ? (*catapult_voltage * battery_voltage / 12.0)
                               : 0.0);

  return catapult_.MakeStatus(fbb);
}

}  // namespace y2022::control_loops::superstructure::catapult