lqr-problem.hpp

Go to the documentation of this file.

#pragma once
 
#include "aligator/context.hpp"
#include "aligator/math.hpp"
#include "aligator/gar/fwd.hpp"
#include "aligator/memory/arena-matrix.hpp"
#include <fmt/format.h>
 
#include <optional>
 
namespace aligator {
namespace gar {

template <typename Scalar> struct LqrKnotTpl {
  ALIGATOR_DYNAMIC_TYPEDEFS(Scalar);
  static constexpr int Alignment = Eigen::AlignedMax;
  using allocator_type = polymorphic_allocator;
 
  uint nx;
  uint nu;
  uint nc;
  uint nx2;
  uint nth;
 
  ArenaMatrix<MatrixXs> Q, S, R;
  ArenaMatrix<VectorXs> q, r;
  ArenaMatrix<MatrixXs> A, B;
  ArenaMatrix<VectorXs> f;
  ArenaMatrix<MatrixXs> C, D;
  ArenaMatrix<VectorXs> d;
 
  ArenaMatrix<MatrixXs> Gth;   //< \f$\theta^\top G_\theta\theta\f$ term
  ArenaMatrix<MatrixXs> Gx;    //< \f$x^\top G_x \theta\f$ term in Lagrangian
  ArenaMatrix<MatrixXs> Gu;    //< \f$u^\top G_x \theta\f$ term in Lagrangian
  ArenaMatrix<MatrixXs> Gv;    //< \f$\nu^\top G_x \theta\f$ term in Lagrangian
  ArenaMatrix<VectorXs> gamma; //< \f$\gamma^\top \theta\f$ term in Lagrangian
 
  LqrKnotTpl() = default;
  explicit LqrKnotTpl(const allocator_type &alloc);
 
  LqrKnotTpl(uint nx, uint nu, uint nc, uint nx2, uint nth,
             allocator_type alloc = {});

  LqrKnotTpl(uint nx, uint nu, uint nc, uint nx2, allocator_type alloc = {})
      : LqrKnotTpl(nx, nu, nc, nx2, 0, alloc) {}

  LqrKnotTpl(uint nx, uint nu, uint nc, allocator_type alloc = {})
      : LqrKnotTpl(nx, nu, nc, nx, 0, alloc) {}

  LqrKnotTpl(const LqrKnotTpl &other, allocator_type alloc = {});
  LqrKnotTpl(LqrKnotTpl &&other) noexcept;
  LqrKnotTpl(LqrKnotTpl &&other, const allocator_type &alloc);
  LqrKnotTpl &operator=(const LqrKnotTpl &other);
  LqrKnotTpl &operator=(LqrKnotTpl &&);
 
  ~LqrKnotTpl() = default;

  void assign(const LqrKnotTpl<Scalar> &other);
 
  // reallocates entire buffer for contigousness
  LqrKnotTpl &addParameterization(uint nth);
 
  bool isApprox(const LqrKnotTpl &other,
                Scalar prec = std::numeric_limits<Scalar>::epsilon()) const;
 
  friend bool operator==(const LqrKnotTpl &lhs, const LqrKnotTpl &rhs) {
    return lhs.isApprox(rhs);
  }
 
  allocator_type get_allocator() const { return m_allocator; }
 
private:
  allocator_type m_allocator;
};
 
template <typename Scalar> struct LqrProblemTpl {
  ALIGATOR_DYNAMIC_TYPEDEFS(Scalar);
  static constexpr int Alignment = Eigen::AlignedMax;
  using KnotType = LqrKnotTpl<Scalar>;
  using KnotVector = std::pmr::vector<KnotType>;
  using allocator_type = polymorphic_allocator;
  ArenaMatrix<MatrixXs> G0;
  ArenaMatrix<VectorXs> g0;
  KnotVector stages;
 
  inline int horizon() const noexcept { return (int)stages.size() - 1; }
  inline uint nc0() const noexcept { return (uint)g0.rows(); }
 
  explicit LqrProblemTpl(allocator_type alloc = {})
      : G0(alloc)
      , g0(alloc)
      , stages(alloc) {
    assert(check_allocators());
  }

  LqrProblemTpl(const KnotVector &knots, long nc0, allocator_type alloc = {});
  LqrProblemTpl(KnotVector &&knots, long nc0);

  LqrProblemTpl(const LqrProblemTpl &other, allocator_type alloc = {})
      : LqrProblemTpl(other.stages, other.nc0(), alloc) {
    this->G0 = other.G0;
    this->g0 = other.g0;
  }

  LqrProblemTpl(LqrProblemTpl &&other)
      : LqrProblemTpl(std::move(other.stages), other.nc0()) {
    this->G0 = other.G0;
    this->g0 = other.g0;
  }
 
  LqrProblemTpl &operator=(LqrProblemTpl &&other) {
    this->G0 = std::move(other.G0);
    this->g0 = std::move(other.g0);
    this->stages = std::move(other.stages);
    return *this;
  }
 
  ~LqrProblemTpl() = default;
 
  void addParameterization(uint nth) {
    if (stages.empty())
      return;
    for (uint i = 0; i <= (uint)horizon(); i++) {
      stages[i].addParameterization(nth);
    }
  }
 
  inline bool isParameterized() const {
    return !stages.empty() && (stages[0].nth > 0);
  }
 
  inline bool isInitialized() const { return !stages.empty(); }
 
  inline uint ntheta() const { return stages[0].nth; }
 
  [[nodiscard]] bool isApprox(const LqrProblemTpl &other) {
    if (horizon() != other.horizon() || !G0.isApprox(other.G0) ||
        !g0.isApprox(other.g0))
      return false;
    for (uint i = 0; i < uint(horizon()); i++) {
      if (!stages[i].isApprox(other.stages[i]))
        return false;
    }
    return true;
  }

  [[nodiscard]] Scalar
  evaluate(const VectorOfVectors &xs, const VectorOfVectors &us,
           const std::optional<ConstVectorRef> &theta) const;
 
  allocator_type get_allocator() const { return G0.get_allocator(); }
 
private:
  [[nodiscard]] bool check_allocators() const {
    return get_allocator() == g0.get_allocator() &&
           get_allocator() == stages.get_allocator();
  }
};
 
template <typename Scalar>
[[nodiscard]] bool lqrKnotsSameDim(const LqrKnotTpl<Scalar> &lhs,
                                   const LqrKnotTpl<Scalar> &rhs) {
  return (lhs.nx == rhs.nx) && (lhs.nu == rhs.nu) && (lhs.nc == rhs.nc) &&
         (lhs.nx2 == rhs.nx2) && (lhs.nth == rhs.nth);
}
 
template <typename Scalar>
std::ostream &operator<<(std::ostream &oss, const LqrKnotTpl<Scalar> &self) {
  oss << "LqrKnot {";
  oss << fmt::format("\n  nx:  {:d}", self.nx) //
      << fmt::format("\n  nu:  {:d}", self.nu) //
      << fmt::format("\n  nc:  {:d}", self.nc);
  if (self.nth > 0) {
    oss << fmt::format("\n  nth: {:d}", self.nth);
  }
#ifndef NDEBUG
  oss << eigenPrintWithPreamble(self.Q, "\n  Q: ") //
      << eigenPrintWithPreamble(self.S, "\n  S: ") //
      << eigenPrintWithPreamble(self.R, "\n  R: ") //
      << eigenPrintWithPreamble(self.q, "\n  q: ") //
      << eigenPrintWithPreamble(self.r, "\n  r: ");
 
  oss << eigenPrintWithPreamble(self.A, "\n  A: ") //
      << eigenPrintWithPreamble(self.B, "\n  B: ") //
      << eigenPrintWithPreamble(self.f, "\n  f: ");
 
  oss << eigenPrintWithPreamble(self.C, "\n  C: ") //
      << eigenPrintWithPreamble(self.D, "\n  D: ") //
      << eigenPrintWithPreamble(self.d, "\n  d: ");
  if (self.nth > 0) {
    oss << eigenPrintWithPreamble(self.Gth, "\n  Gth: ") //
        << eigenPrintWithPreamble(self.Gx, "\n  Gx: ")   //
        << eigenPrintWithPreamble(self.Gu, "\n  Gu: ")   //
        << eigenPrintWithPreamble(self.gamma, "\n  gamma: ");
  }
#endif
  oss << "\n}";
  return oss;
}
 
#ifdef ALIGATOR_ENABLE_TEMPLATE_INSTANTIATION
extern template struct LqrKnotTpl<context::Scalar>;
extern template struct LqrProblemTpl<context::Scalar>;
#endif
 
} // namespace gar
} // namespace aligator