mpc.hpp

// Copyright (C) 2022 Petter Nilsson. MIT License.
 
#pragma once
 
#include <memory>
 
#include <smooth/concepts/lie_group.hpp>
#include <smooth/lie_sparse.hpp>
 
#include "ocp_to_qp.hpp"
#include "qp_solver.hpp"
#include "time.hpp"
#include "utils/sparse.hpp"
 
namespace smooth::feedback {
 
namespace detail {
 
template<Time T, LieGroup X>
struct XDes
{
  T t0;
  std::function<X(T)> xdes           = [](T) -> X { return Default<X>(); };
  std::function<Tangent<X>(T)> dxdes = [](T) -> Tangent<X> { return Tangent<X>::Zero(); };
 
  X operator()(const double t_rel) const
  {
    const T t_abs = time_trait<T>::plus(t0, t_rel);
    return xdes(t_abs);
  };
 
  Tangent<X> jacobian(const double t_rel) const
  {
    const T t_abs = time_trait<T>::plus(t0, t_rel);
    return dxdes(t_abs);
  }
};
 
template<Time T, Manifold U>
struct UDes
{
  T t0;
  std::function<U(T)> udes = [](T) -> U { return Default<U>(); };
 
  U operator()(const double t_rel) const
  {
    const T t_abs = time_trait<T>::plus(t0, t_rel);
    return udes(t_abs);
  };
};
 
template<LieGroup X>
struct MPCObj
{
  static constexpr auto Nx = Dof<X>;
 
  // public members
  X xf_des                          = Default<X>();
  Eigen::Matrix<double, Nx, Nx> Qtf = Eigen::Matrix<double, Nx, Nx>::Identity();
 
  // private members
  Eigen::SparseMatrix<double> hess{1 + 2 * Nx + 1, 1 + 2 * Nx + 1};
 
  // functor members
 
  // function f(t, x0, xf, q) = (1/2) xf' Q xf + q_0
  double operator()(
    const double, const X &, [[maybe_unused]] const X & xf, [[maybe_unused]] const Eigen::Vector<double, 1> & q) const
  {
    assert(q(0) == 1.);
    assert(xf_des.isApprox(xf, 1e-4));
    if constexpr (false) {
      const auto e = rminus(xf, xf_des);
      return 0.5 * e.dot(Qtf * e) + q(0);
    } else {
      return 1.;
    }
  }
 
  Eigen::RowVector<double, 1 + 2 * Nx + 1>
  jacobian(const double, const X &, [[maybe_unused]] const X & xf, const Eigen::Vector<double, 1> &)
  {
    assert(xf_des.isApprox(xf, 1e-4));
    return Eigen::RowVector<double, 1 + 2 * Nx + 1>::Unit(1 + 2 * Nx);
  }
 
  std::reference_wrapper<const Eigen::SparseMatrix<double>>
  hessian(const double, const X &, [[maybe_unused]] const X & xf, const Eigen::Vector<double, 1> &)
  {
    assert(xf_des.isApprox(xf, 1e-4));
 
    set_zero(hess);
 
    for (auto i = 0u; i < Nx; ++i) {
      for (auto j = 0u; j < Nx; ++j) {
        if (Qtf(i, j) != 0) { hess.coeffRef(1 + i, 1 + j) = Qtf(i, j); }
      }
    }
 
    hess.makeCompressed();
    return hess;
  }
};
 
template<Time T, LieGroup X, Manifold U, typename F, diff::Type DT>
struct MPCDyn
{
  static constexpr auto Nx = Dof<X>;
  static constexpr auto Nu = Dof<U>;
 
  // public members
  F f;
  T t0{};
 
  // functor members
  Tangent<X> operator()(const double t, const X & x, const U & u)
  {
    if constexpr (requires(F & fvar, T tvar) { fvar.set_time(tvar); }) { f.set_time(time_trait<T>::plus(t0, t)); }
 
    return f(x, u);
  }
 
  Eigen::Matrix<double, Nx, 1 + Nx + Nu> jacobian(const double t, const X & x, const U & u)
  {
    if constexpr (requires(F & fvar, T tvar) { fvar.set_time(tvar); }) { f.set_time(time_trait<T>::plus(t0, t)); }
 
    const auto & [fval, df] = diff::dr<1, DT>(f, smooth::wrt(x, u));
 
    Eigen::Matrix<double, Nx, 1 + Nx + Nu> ret;
    ret << Eigen::Vector<double, Nx>::Zero(), df;
    return ret;
  }
};
 
template<Time T, LieGroup X, Manifold U>
struct MPCIntegrand
{
  static constexpr auto Nx = Dof<X>;
  static constexpr auto Nu = Dof<U>;
 
  // public members
  std::shared_ptr<XDes<T, X>> xdes;
  std::shared_ptr<UDes<T, U>> udes;
 
  Eigen::Matrix<double, Nx, Nx> Q = Eigen::Matrix<double, Nx, Nx>::Identity();
  Eigen::Matrix<double, Nu, Nu> R = Eigen::Matrix<double, Nu, Nu>::Identity();
 
  // private members
  Eigen::SparseMatrix<double> hess{1 + Nx + Nu, 1 + Nx + Nu};
 
  // functor members
 
  // function f(x, t, u) = (1/2) * (x' Q x + u' R u)
  Eigen::Vector<double, 1>
  operator()([[maybe_unused]] const double t_rel, [[maybe_unused]] const X & x, [[maybe_unused]] const U & u) const
  {
    assert((*xdes)(t_rel).isApprox(x, 1e-4));
    assert((*udes)(t_rel).isApprox(u, 1e-4));
    if constexpr (false) {
      const auto ex = rminus(x, (*xdes)(t_rel));
      const auto eu = rminus(u, (*udes)(t_rel));
      return Eigen::Vector<double, 1>{0.5 * ex.dot(Q * ex) + 0.5 * eu.dot(R * eu)};
    } else {
      return Eigen::Vector<double, 1>{0.};
    }
  }
 
  Eigen::RowVector<double, 1 + Nx + Nu>
  jacobian([[maybe_unused]] const double t_rel, [[maybe_unused]] const X & x, [[maybe_unused]] const U & u)
  {
    assert((*xdes)(t_rel).isApprox(x, 1e-4));
    assert((*udes)(t_rel).isApprox(u, 1e-4));
    return Eigen::RowVector<double, 1 + Nx + Nu>::Zero();
  }
 
  std::reference_wrapper<const Eigen::SparseMatrix<double>>
  hessian([[maybe_unused]] const double t_rel, [[maybe_unused]] const X & x, [[maybe_unused]] const U & u)
  {
    assert((*xdes)(t_rel).isApprox(x, 1e-4));
    assert((*udes)(t_rel).isApprox(u, 1e-4));
 
    set_zero(hess);
 
    for (auto i = 0u; i < Nx; ++i) {
      for (auto j = 0u; j < Nx; ++j) {
        if (Q(i, j) != 0) { hess.coeffRef(1 + i, 1 + j) = Q(i, j); }
      }
    }
    for (auto i = 0u; i < Nu; ++i) {
      for (auto j = 0u; j < Nu; ++j) {
        if (Q(i, j) != 0) { hess.coeffRef(1 + Nx + i, 1 + Nx + j) = R(i, j); }
      }
    }
 
    hess.makeCompressed();
    return hess;
  }
};
 
template<Time T, LieGroup X, Manifold U, typename F, diff::Type DT>
struct MPCCR
{
  static constexpr auto Nx  = Dof<X>;
  static constexpr auto Nu  = Dof<U>;
  static constexpr auto Ncr = std::invoke_result_t<F, X, U>::SizeAtCompileTime;
 
  // public members
  F f;
  T t0{};
 
  // private members
  Eigen::SparseMatrix<double> jac{Ncr, 1 + Nx + Nu};
 
  // functor members
  Eigen::Vector<double, Ncr> operator()(const double t, const X & x, const U & u)
  {
    if constexpr (requires(F & fvar, T tvar) { fvar.set_time(tvar); }) { f.set_time(time_trait<T>::plus(t0, t)); }
 
    return f(x, u);
  }
 
  std::reference_wrapper<const Eigen::SparseMatrix<double>> jacobian(const double t, const X & x, const U & u)
  {
    if constexpr (requires(F & fvar, T tvar) { fvar.set_time(tvar); }) { f.set_time(time_trait<T>::plus(t0, t)); }
 
    const auto & [fval, df] = diff::dr<1, DT>(f, smooth::wrt(x, u));
    block_write(jac, 0, 1, df);
 
    jac.makeCompressed();
    return jac;
  }
};
 
template<LieGroup X>
struct MPCCE
{
  static constexpr auto Nx = Dof<X>;
 
  // public members
  X x0_fix = Default<X>();
 
  // private members
  Eigen::SparseMatrix<double> dexpinv_tmp = d_exp_sparse_pattern<X>;
  Eigen::SparseMatrix<double> jac{Nx, 1 + 2 * Nx + 1};
 
  // functor members
  Tangent<X> operator()(const double, const X & x0, const X &, const Eigen::Vector<double, 1> &) const
  {
    return rminus(x0, x0_fix);
  }
 
  std::reference_wrapper<const Eigen::SparseMatrix<double>>
  jacobian(const double, const X & x0, const X &, const Eigen::Vector<double, 1> &)
  {
    dr_expinv_sparse<X>(dexpinv_tmp, rminus(x0, x0_fix));
 
    jac.middleCols(1, Nx) = dexpinv_tmp;
    jac.makeCompressed();
    return jac;
  }
};
 
}  // namespace detail
 
struct MPCParams
{
  std::size_t K{10};
 
  double tf{1};
 
  bool warmstart{true};
 
  QPSolverParams qp{};
};
 
template<LieGroup X, Manifold U>
struct MPCWeights
{
  static constexpr auto Nx = Dof<X>;
  static constexpr auto Nu = Dof<U>;
 
  Eigen::Matrix<double, Nx, Nx> Q = Eigen::Matrix<double, Nx, Nx>::Identity();
  Eigen::Matrix<double, Nx, Nx> Qtf = Eigen::Matrix<double, Nx, Nx>::Identity();
  Eigen::Matrix<double, Nu, Nu> R = Eigen::Matrix<double, Nu, Nu>::Identity();
};
 
template<
  Time T,
  LieGroup X,
  Manifold U,
  typename F,
  typename CR,
  std::size_t Kmesh = 4,
  diff::Type DT     = diff::Type::Default>
class MPC
{
  static constexpr auto Ncr = std::invoke_result_t<CR, X, U>::SizeAtCompileTime;
 
public:
  inline MPC(
    F && f, CR && cr, Eigen::Vector<double, Ncr> && crl, Eigen::Vector<double, Ncr> && cru, MPCParams && prm = {})
      : xdes_{std::make_shared<detail::XDes<T, X>>()}, udes_{std::make_shared<detail::UDes<T, U>>()},
        mesh_{(prm.K + Kmesh - 1) / Kmesh},
        ocp_{
          .theta = {},
          .f     = {.f = std::forward<F>(f)},
          .g     = {.xdes = xdes_, .udes = udes_},
          .cr    = {.f = std::forward<CR>(cr)},
          .crl   = std::move(crl),
          .cru   = std::move(cru),
          .ce    = {},
          .cel   = Eigen::Vector<double, Dof<X>>::Zero(),
          .ceu   = Eigen::Vector<double, Dof<X>>::Zero(),
        },
        prm_{std::move(prm)}, qp_solver_{prm_.qp}
  {
    detail::ocp_to_qp_allocate<DT>(qp_, work_, ocp_, mesh_);
    ocp_to_qp_update<diff::Type::Analytic>(qp_, work_, ocp_, mesh_, prm_.tf, *xdes_, *udes_);
    qp_solver_.analyze(qp_);
  }
  inline MPC(
    const F & f,
    const CR & cr,
    const Eigen::Vector<double, Ncr> & crl,
    const Eigen::Vector<double, Ncr> & cru,
    const MPCParams & prm = {})
      : MPC(F(f), CR(cr), Eigen::Vector<double, Ncr>(crl), Eigen::Vector<double, Ncr>(cru), MPCParams(prm))
  {}
  inline MPC() = default;
  inline MPC(const MPC &) = default;
  inline MPC(MPC &&) = default;
  inline MPC & operator=(const MPC &) = default;
  inline MPC & operator=(MPC &&) = default;
  inline ~MPC() = default;
 
  inline std::pair<U, QPSolutionStatus> operator()(
    const T & t,
    const X & x,
    std::optional<std::reference_wrapper<std::vector<U>>> u_traj = std::nullopt,
    std::optional<std::reference_wrapper<std::vector<X>>> x_traj = std::nullopt)
  {
    static constexpr auto Nx = Dof<X>;
    static constexpr auto Nu = Dof<U>;
 
    const auto N      = mesh_.N_colloc();
    const auto xvar_L = Nx * (N + 1);
    const auto xvar_B = 0u;
    const auto uvar_B = xvar_L;
 
    // update problem
    xdes_->t0         = t;
    udes_->t0         = t;
    ocp_.theta.xf_des = (*xdes_)(prm_.tf);
    ocp_.f.t0         = t;
    ocp_.cr.t0        = t;
    ocp_.ce.x0_fix    = x;
 
    // transcribe to QP
    ocp_to_qp_update_dyn<diff::Type::Analytic>(qp_, work_, ocp_, mesh_, prm_.tf, *xdes_, *udes_);
    if constexpr (requires(CR & crvar, T tvar) { crvar.set_time(tvar); }) {
      // only update if running constraints are time-dependent
      ocp_to_qp_update_cr<diff::Type::Analytic>(qp_, work_, ocp_, mesh_, prm_.tf, *xdes_, *udes_);
    }
    ocp_to_qp_update_ce<diff::Type::Analytic>(qp_, work_, ocp_, mesh_, prm_.tf, *xdes_, *udes_);
    qp_.A.makeCompressed();
    qp_.P.makeCompressed();
 
    // solve QP
    const auto & sol = qp_solver_.solve(qp_, warmstart_);
 
    // output solution trajectories
    if (u_traj.has_value()) {
      u_traj.value().get().resize(N);
      for (const auto & [i, tau] : zip(std::views::iota(0u, N), mesh_.all_nodes())) {
        const double t_rel      = prm_.tf * tau;
        u_traj.value().get()[i] = (*udes_)(t_rel) + sol.primal.template segment<Nu>(uvar_B + i * Nu);
      }
    }
    if (x_traj.has_value()) {
      x_traj.value().get().resize(N + 1);
      for (const auto & [i, tau] : zip(std::views::iota(0u, N + 1), mesh_.all_nodes())) {
        const double t_rel      = prm_.tf * tau;
        x_traj.value().get()[i] = (*xdes_)(t_rel) + sol.primal.template segment<Nx>(xvar_B + i * Nx);
      }
    }
 
    // save solution to warmstart next iteration
    if (prm_.warmstart) {
      // clang-format off
      if (sol.code == QPSolutionStatus::Optimal || sol.code == QPSolutionStatus::MaxTime || sol.code == QPSolutionStatus::MaxIterations) {
        warmstart_ = sol;
      }
      // clang-format on
    }
 
    return {rplus((*udes_)(0), sol.primal.template segment<Nu>(uvar_B)), sol.code};
  }
 
  inline void set_udes(std::function<U(T)> && u_des) { udes_->udes = std::move(u_des); }
 
  inline void set_udes(const std::function<U(T)> & u_des) { set_udes(std::function<U(T)>(u_des)); }
 
  template<typename Fun>
    requires(std::is_same_v<std::invoke_result_t<Fun, Scalar<U>>, U>)
  inline void set_udes_rel(Fun && f, T t0 = T(0))
  {
    set_udes([t0 = t0, f = std::forward<Fun>(f)](T t_abs) -> U {
      const double t_rel = time_trait<T>::minus(t_abs, t0);
      return std::invoke(f, t_rel);
    });
  }
 
  inline void set_xdes(std::function<X(T)> && x_des, std::function<Tangent<X>(T)> && dx_des)
  {
    xdes_->xdes  = std::move(x_des);
    xdes_->dxdes = std::move(dx_des);
  }
 
  inline void set_xdes(const std::function<X(T)> & x_des, const std::function<Tangent<X>(T)> & dx_des)
  {
    set_xdes(std::function<X(T)>(x_des), std::function<Tangent<X>(T)>(dx_des));
  }
 
  template<typename Fun>
    requires(std::is_same_v<std::invoke_result_t<Fun, Scalar<X>>, X>)
  inline void set_xdes_rel(Fun && f, T t0 = T(0))
  {
    std::function<X(T)> x_des = [t0 = t0, f = f](T t_abs) -> X {
      const double t_rel = time_trait<T>::minus(t_abs, t0);
      return std::invoke(f, t_rel);
    };
 
    std::function<Tangent<X>(T)> dx_des = [t0 = t0, f = f](T t_abs) -> Tangent<X> {
      const double t_rel = time_trait<T>::minus(t_abs, t0);
      return std::get<1>(diff::dr<1, DT>(f, wrt(t_rel)));
    };
 
    set_xdes(std::move(x_des), std::move(dx_des));
  }
 
  inline void set_weights(const MPCWeights<X, U> & weights)
  {
    ocp_.g.R       = weights.R;
    ocp_.g.Q       = weights.Q;
    ocp_.theta.Qtf = weights.Qtf;
  }
 
  inline void reset_warmstart() { warmstart_ = {}; }
 
private:
  // linearization
  std::shared_ptr<detail::XDes<T, X>> xdes_;
  std::shared_ptr<detail::UDes<T, U>> udes_;
 
  // collocation mesh
  Mesh<Kmesh, Kmesh> mesh_{};
 
  // internal optimal control problem
  OCP<
    X,
    U,
    detail::MPCObj<X>,
    detail::MPCDyn<T, X, U, F, DT>,
    detail::MPCIntegrand<T, X, U>,
    detail::MPCCR<T, X, U, CR, DT>,
    detail::MPCCE<X>>
    ocp_;
 
  // parameters
  MPCParams prm_{};
 
  // internal allocation
  detail::OcpToQpWorkmemory work_;
  QuadraticProgramSparse<double> qp_;
 
  // internal QP solver
  QPSolver<QuadraticProgramSparse<double>> qp_solver_;
 
  // last solution stored for warmstarting
  std::optional<QPSolution<-1, -1, double>> warmstart_{};
};
 
}  // namespace smooth::feedback