solver-proxddp.hpp

Go to the documentation of this file.

#pragma once
 
#include "aligator/overloads.hpp"
#include "aligator/core/linesearch-armijo.hpp"
#include "aligator/core/linesearch-nonmonotone.hpp"
#include "aligator/core/filter.hpp"
#include "aligator/core/callback-base.hpp"
#include "aligator/core/enums.hpp"
#include "aligator/utils/logger.hpp"
#include "aligator/gar/riccati-base.hpp"
 
#include "workspace.hpp"
#include "results.hpp"
 
#include <boost/unordered_map.hpp>
#include <variant>
 
namespace aligator {
 
enum class LQSolverChoice { SERIAL, PARALLEL, STAGEDENSE };

template <typename _Scalar> struct SolverProxDDPTpl {
public:
  using Scalar = _Scalar;
  ALIGATOR_DYNAMIC_TYPEDEFS(Scalar);
  using Problem = TrajOptProblemTpl<Scalar>;
  using Workspace = WorkspaceTpl<Scalar>;
  using Results = ResultsTpl<Scalar>;
  using StageFunctionData = StageFunctionDataTpl<Scalar>;
  using CostData = CostDataAbstractTpl<Scalar>;
  using StageModel = StageModelTpl<Scalar>;
  using StageData = StageDataTpl<Scalar>;
  using CallbackPtr = shared_ptr<CallbackBaseTpl<Scalar>>;
  using CallbackMap = boost::unordered_map<std::string, CallbackPtr>;
  using ConstraintStack = ConstraintStackTpl<Scalar>;
  using CstrSet = ConstraintSetTpl<Scalar>;
  using TrajOptData = TrajOptDataTpl<Scalar>;
  using LinearSolverPtr = std::unique_ptr<gar::RiccatiSolverBase<Scalar>>;
 
  struct LinesearchVariant {
    using VariantType = std::variant<std::monostate, ArmijoLinesearch<Scalar>,
                                     NonmonotoneLinesearch<Scalar>>;
 
    Scalar run(const std::function<Scalar(Scalar)> &fun, const Scalar phi0,
               const Scalar dphi0, Scalar &alpha_try) {
      return std::visit(
          overloads{[](std::monostate &) {
                      return std::numeric_limits<Scalar>::quiet_NaN();
                    },
                    [&](auto &method) {
                      return method.run(fun, phi0, dphi0, alpha_try);
                    }},
          impl_);
    }
 
    void reset() {
      std::visit(overloads{[](std::monostate &) {},
                           [&](auto &method) { method.reset(); }},
                 impl_);
    }
 
    bool isValid() const { return impl_.index() > 0ul; }
 
    operator const VariantType &() const { return impl_; }
 
  private:
    explicit LinesearchVariant() {}
    void init(StepAcceptanceStrategy strat,
              const LinesearchOptions<Scalar> &options) {
      switch (strat) {
      case StepAcceptanceStrategy::LINESEARCH_ARMIJO:
        impl_ = ArmijoLinesearch(options);
        break;
      case StepAcceptanceStrategy::LINESEARCH_NONMONOTONE:
        impl_ = NonmonotoneLinesearch(options);
        break;
      default:
        ALIGATOR_RUNTIME_ERROR(
            "Provided StepAcceptanceStrategy value is invalid.");
        break;
      }
    }
    friend SolverProxDDPTpl;
    VariantType impl_;
  };
 
  struct AlmParams {
    Scalar prim_alpha = 0.1;
    Scalar prim_beta = 0.9;
    Scalar dual_alpha = 1.;
    Scalar dual_beta = 1.;
    Scalar mu_update_factor = 0.01;
    Scalar mu_lower_bound = 1e-8;
  };

  Scalar inner_tol_;
  Scalar prim_tol_;
  Scalar target_tol_;
  Scalar mu_init_; //< Initial AL parameter

 
  Scalar reg_min = 1e-10;         //< Minimal nonzero regularization
  Scalar reg_max = 1e9;           //< Maximum regularization value
  Scalar reg_init = 1e-9;         //< Initial regularization value (can be zero)
  Scalar reg_inc_k_ = 10.;        //< Regularization increase factor
  Scalar reg_inc_first_k_ = 100.; //< Regularization increase (critical)
  Scalar reg_dec_k_ = 1. / 3.;    //< Regularization decrease factor
  Scalar preg_ = reg_init;        //< Primal regularization value
  Scalar preg_last_ = 0.;         //< Last "good" regularization value

 
  Scalar inner_tol0 = 1.; //< Initial BCL inner subproblem tolerance
  Scalar prim_tol0 = 1.;  //< Initial BCL constraint infeasibility tolerance

  Logger logger{};

  VerboseLevel verbose_;
  LQSolverChoice linear_solver_choice = LQSolverChoice::SERIAL;
  HessianApprox hess_approx_;
  LinesearchOptions<Scalar> ls_params;
  MultiplierUpdateMode multiplier_update_mode = MultiplierUpdateMode::NEWTON;
  LinesearchMode ls_mode = LinesearchMode::PRIMAL;
  Scalar dual_weight = 1.0;
  RolloutType rollout_type_ = RolloutType::LINEAR;
  AlmParams bcl_params;
  StepAcceptanceStrategy sa_strategy_;
 
  bool force_initial_condition_ = true; //< Force initial condition @f$ x_0 @f$.
  size_t max_refinement_steps_ = 0;     //< Max KKT system refinement iters.
  Scalar refinement_threshold_ = 1e-13; //< Target tol. for the KKT system.
  size_t max_iters;                     //< Max number of Newton iterations.
  size_t max_al_iters = 100;            //< Maximum number of ALM iterations.
 
  polymorphic_allocator allocator_; //< Main allocator
  Workspace workspace_;             //< Solver workspace
  Results results_;                 //< Solver results
  LinearSolverPtr linear_solver_;   //< LQR subproblem solver
  FilterTpl<Scalar> filter_;        //< Filter linesearch
  LinesearchVariant linesearch_;    //< Linesearch routine
 
protected:
  Scalar target_dual_tol_;
  bool sync_dual_tol_;
  CallbackMap callbacks_;
  size_t num_threads_ = 1;
  Scalar mu_penal_ = mu_init_;
 
public:
  SolverProxDDPTpl(const Scalar tol = 1e-6, const Scalar mu_init = 0.01,
                   const size_t max_iters = 1000,
                   VerboseLevel verbose = VerboseLevel::QUIET,
                   StepAcceptanceStrategy sa_strategy =
                       StepAcceptanceStrategy::LINESEARCH_NONMONOTONE,
                   HessianApprox hess_approx = HessianApprox::GAUSS_NEWTON);
 
  inline size_t getNumThreads() const { return num_threads_; }
  void setNumThreads(const size_t num_threads);
 
  Scalar getDualTolerance() const { return target_dual_tol_; }
  void setDualTolerance(const Scalar tol) {
    target_dual_tol_ = tol;
    sync_dual_tol_ = false;
  }

  Scalar tryLinearStep(const Problem &problem, const Scalar alpha);

  Scalar tryNonlinearRollout(const Problem &problem, const Scalar alpha);
 
  Scalar forwardPass(const Problem &problem, const Scalar alpha);
 
  void updateLQSubproblem();

  void setup(const Problem &problem);
  void cycleProblem(const Problem &problem, const shared_ptr<StageData> &data);

  bool run(const Problem &problem, const std::vector<VectorXs> &xs_init = {},
           const std::vector<VectorXs> &us_init = {},
           const std::vector<VectorXs> &vs_init = {},
           const std::vector<VectorXs> &lams_init = {});

  bool innerLoop(const Problem &problem);

  void computeCriterion();


  void registerCallback(const std::string &name, CallbackPtr cb) {
    callbacks_.insert_or_assign(name, cb);
  }

  void clearCallbacks() noexcept { callbacks_.clear(); }
 
  const CallbackMap &getCallbacks() const { return callbacks_; }
 
  [[nodiscard]] bool removeCallback(const std::string &name) {
    return callbacks_.erase(name);
  }
 
  [[nodiscard]] auto getCallbackNames() const {
    std::vector<std::string> keys;
    for (const auto &item : callbacks_) {
      keys.push_back(item.first);
    }
    return keys;
  }
 
  [[nodiscard]] CallbackPtr getCallback(const std::string &name) const {
    auto cb = callbacks_.find(name);
    if (cb != end(callbacks_)) {
      return cb->second;
    }
    return nullptr;
  }

  void invokeCallbacks() {
    for (const auto &cb : callbacks_) {
      cb.second->call(workspace_, results_);
    }
  }


  bool computeMultipliers(const Problem &problem,
                          const std::vector<VectorXs> &xs,
                          const std::vector<VectorXs> &lams,
                          const std::vector<VectorXs> &vs);
 
  ALIGATOR_INLINE Scalar mu() const { return mu_penal_; }
  ALIGATOR_INLINE Scalar mu_inv() const { return 1. / mu_penal_; }
  ALIGATOR_INLINE Scalar mu_dyn() const { return 0.1 * mu_penal_; }
 
protected:
  void updateTolsOnFailure() noexcept {
    const Scalar arg = std::min(mu_penal_, 0.99);
    prim_tol_ = prim_tol0 * std::pow(arg, bcl_params.prim_alpha);
    inner_tol_ = inner_tol0 * std::pow(arg, bcl_params.dual_alpha);
  }
 
  void updateTolsOnSuccess() noexcept {
    const Scalar arg = std::min(mu_penal_, 0.99);
    prim_tol_ = prim_tol_ * std::pow(arg, bcl_params.prim_beta);
    inner_tol_ = inner_tol_ * std::pow(arg, bcl_params.dual_beta);
  }

  ALIGATOR_INLINE void setAlmPenalty(Scalar new_mu) noexcept {
    mu_penal_ = std::max(new_mu, bcl_params.mu_lower_bound);
  }
 
  // See sec. 3.1 of the IPOPT paper [Wächter, Biegler 2006]
  // called before first bwd pass attempt
  inline void initializeRegularization() noexcept {
    if (preg_last_ == 0.) {
      // this is the 1st iteration
      preg_ = std::max(reg_init, reg_min);
    } else {
      // attempt decrease from last "good" value
      preg_ = std::max(reg_min, preg_last_ * reg_dec_k_);
    }
  }
 
  inline void increaseRegularization() noexcept {
    if (preg_last_ == 0.)
      preg_ *= reg_inc_first_k_;
    else
      preg_ *= reg_inc_k_;
  }
};
 
#ifdef ALIGATOR_ENABLE_TEMPLATE_INSTANTIATION
extern template struct SolverProxDDPTpl<context::Scalar>;
#endif
} // namespace aligator