parallel-solver.hpp

Go to the documentation of this file.

#pragma once
 
#include "aligator/gar/riccati-base.hpp"
#include "aligator/gar/riccati-kernel.hpp"
#include "aligator/gar/lqr-problem.hpp"
#include "aligator/tracy.hpp"
 
namespace aligator {
namespace gar {
 
#ifdef ALIGATOR_MULTITHREADING
template <typename _Scalar>
class ParallelRiccatiSolver : public RiccatiSolverBase<_Scalar> {
public:
  using Scalar = _Scalar;
  ALIGATOR_DYNAMIC_TYPEDEFS_WITH_ROW_TYPES(Scalar);
  using Base = RiccatiSolverBase<Scalar>;
  using StageFactorVec = std::vector<StageFactor<Scalar>>;
  StageFactorVec datas;
 
  using Kernel = ProximalRiccatiKernel<Scalar>;
  using KnotType = LqrKnotTpl<Scalar>;
 
  using BlkMat = BlkMatrix<MatrixXs, -1, -1>;
  using BlkVec = BlkMatrix<VectorXs, -1, 1>;
 
  explicit ParallelRiccatiSolver(LqrProblemTpl<Scalar> &problem,
                                 const uint num_threads);
 
  void allocateLeg(uint start, uint end, bool last_leg);
 
  static void setupKnot(KnotType &knot, const Scalar mudyn) {
    ALIGATOR_TRACY_ZONE_SCOPED;
    ALIGATOR_NOMALLOC_SCOPED;
    knot.Gx = knot.A.to_const_map().transpose();
    knot.Gu = knot.B.to_const_map().transpose();
    knot.Gth.setZero();
    knot.Gth.diagonal().setConstant(-mudyn);
    knot.gamma = knot.f;
  }
 
  bool backward(const Scalar mudyn, const Scalar mueq);
 
  inline void collapseFeedback() {
    using RowMatrix = Eigen::Matrix<Scalar, -1, -1, Eigen::RowMajor>;
    StageFactor<Scalar> &d = datas[0];
    Eigen::Ref<RowMatrix> K = d.fb.blockRow(0);
    Eigen::Ref<RowMatrix> Kth = d.fth.blockRow(0);
 
    // condensedSystem.subdiagonal contains the 'U' factors in the
    // block-tridiag UDUt decomposition
    // and ∂Xi+1 = -Ui+1.t ∂Xi
    auto &Up1t = condensedKktSystem.subdiagonal[1];
    K.noalias() -= Kth * Up1t;
  }
 
  struct condensed_system_t {
    std::vector<MatrixXs> subdiagonal;
    std::vector<MatrixXs> diagonal;
    std::vector<MatrixXs> superdiagonal;
  };
 
  struct condensed_system_factor {
    std::vector<MatrixXs> diagonalFacs; //< diagonal factors
    std::vector<MatrixXs> upFacs;       //< transposed U factors
    std::vector<Eigen::BunchKaufman<MatrixXs>> ldlt;
  };

  void assembleCondensedSystem(const Scalar mudyn);
 
  bool forward(VectorOfVectors &xs, VectorOfVectors &us, VectorOfVectors &vs,
               VectorOfVectors &lbdas,
               const std::optional<ConstVectorRef> & = std::nullopt) const;
 
  void cycleAppend(const KnotType &knot);
  VectorRef getFeedforward(size_t i) { return datas[i].ff.matrix(); }
  RowMatrixRef getFeedback(size_t i) { return datas[i].fb.matrix(); }

  uint numThreads;

  condensed_system_t condensedKktSystem;
  condensed_system_factor condensedFacs;
  BlkVec condensedKktRhs, condensedKktSolution, condensedErr;
 
  Scalar condensedThreshold{1e-11};

  void initializeTridiagSystem(const std::vector<long> &dims);
 
protected:
  LqrProblemTpl<Scalar> *problem_;
};
#endif
 
} // namespace gar
} // namespace aligator
 
#ifdef ALIGATOR_ENABLE_TEMPLATE_INSTANTIATION
#include "aligator/gar/parallel-solver.txx"
#endif