traj-opt-problem.hpp

Go to the documentation of this file.

#pragma once
 
#include "aligator/core/stage-model.hpp"
#include "aligator/modelling/state-error.hpp"
 
namespace aligator {

template <typename _Scalar> struct TrajOptProblemTpl {
  using Self = TrajOptProblemTpl;
  using Scalar = _Scalar;
 
  ALIGATOR_DYNAMIC_TYPEDEFS(Scalar);
 
  using StageModel = StageModelTpl<Scalar>;
  using StageFunction = StageFunctionTpl<Scalar>;
  using UnaryFunction = UnaryFunctionTpl<Scalar>;
  using Data = TrajOptDataTpl<Scalar>;
  using Manifold = ManifoldAbstractTpl<Scalar>;
  using CostAbstract = CostAbstractTpl<Scalar>;
  using ConstraintSet = ConstraintSetTpl<Scalar>;
  using StateErrorResidual = StateErrorResidualTpl<Scalar>;
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
  using StageConstraint = StageConstraintTpl<Scalar>;
#pragma GCC diagnostic pop
  using InitializationStrategy =
      std::function<void(const Self &, std::vector<VectorXs> &)>;


  xyz::polymorphic<UnaryFunction> init_constraint_;
  std::vector<xyz::polymorphic<StageModel>> stages_;
  xyz::polymorphic<CostAbstract> term_cost_;
  ConstraintStackTpl<Scalar> term_cstrs_;
  VectorXs unone_;


  TrajOptProblemTpl(xyz::polymorphic<UnaryFunction> init_constraint,
                    const std::vector<xyz::polymorphic<StageModel>> &stages,
                    xyz::polymorphic<CostAbstract> term_cost);

  TrajOptProblemTpl(const ConstVectorRef &x0,
                    const std::vector<xyz::polymorphic<StageModel>> &stages,
                    xyz::polymorphic<CostAbstract> term_cost);


  TrajOptProblemTpl(xyz::polymorphic<UnaryFunction> init_constraint,
                    xyz::polymorphic<CostAbstract> term_cost);

  TrajOptProblemTpl(const ConstVectorRef &x0, const int nu,
                    xyz::polymorphic<Manifold> space,
                    xyz::polymorphic<CostAbstract> term_cost);
 
  bool initCondIsStateError() const {
    assert(init_cond_is_state_error_ == checkInitCondIsStateError());
    return init_cond_is_state_error_;
  }

  void addStage(const xyz::polymorphic<StageModel> &stage);

  ConstVectorRef getInitState() const {
    if (!initCondIsStateError()) {
      ALIGATOR_RUNTIME_ERROR(
          "Initial condition is not a StateErrorResidual.\n");
    }
    return static_cast<StateErrorResidual const *>(&*init_constraint_)->target_;
  }

  void setInitState(const ConstVectorRef &x0) {
    if (!initCondIsStateError()) {
      ALIGATOR_RUNTIME_ERROR(
          "Initial condition is not a StateErrorResidual.\n");
    }
    static_cast<StateErrorResidual *>(&*init_constraint_)->target_ = x0;
  }

  ALIGATOR_DEPRECATED void addTerminalConstraint(const StageConstraint &cstr);
  void addTerminalConstraint(const xyz::polymorphic<StageFunction> &func,
                             const xyz::polymorphic<ConstraintSet> &set) {
    this->term_cstrs_.pushBack(func, set);
  }

  void removeTerminalConstraints() { term_cstrs_.clear(); }
 
  [[nodiscard]] std::size_t numSteps() const;

  Scalar evaluate(const std::vector<VectorXs> &xs,
                  const std::vector<VectorXs> &us, Data &prob_data,
                  std::size_t num_threads = 1) const;

  void computeDerivatives(const std::vector<VectorXs> &xs,
                          const std::vector<VectorXs> &us, Data &prob_data,
                          std::size_t num_threads = 1,
                          bool compute_second_order = true) const;

  void replaceStageCircular(const xyz::polymorphic<StageModel> &model);
 
  bool checkIntegrity() const;

  //
  template <typename Callable> void setInitializationStrategy(Callable &&func) {
    static_assert(std::is_convertible_v<Callable, InitializationStrategy>);
    this->xs_init_strategy_ = std::forward<Callable>(func);
  }

  void initializeSolution(std::vector<VectorXs> &xs,
                          std::vector<VectorXs> &us) const {
    xs_init_strategy_(*this, xs);
    us_default_init(*this, us);
  }

  void initializeSolution(std::vector<VectorXs> &xs, std::vector<VectorXs> &us,
                          std::vector<VectorXs> &vs,
                          std::vector<VectorXs> &lbdas) const {
    const size_t nsteps = numSteps();
    initializeSolution(xs, us);
    // initialize multipliers...
    vs.resize(nsteps + 1);
    lbdas.resize(nsteps + 1);
    lbdas[0].setZero(init_constraint_->nr);
    for (size_t i = 0; i < nsteps; i++) {
      const StageModelTpl<Scalar> &sm = *stages_[i];
      lbdas[i + 1].setZero(sm.ndx2());
      vs[i].setZero(sm.nc());
    }
 
    if (!term_cstrs_.empty()) {
      vs[nsteps].setZero(term_cstrs_.totalDim());
    }
  }

  [[nodiscard]] auto initializeSolution() const {
    std::vector<VectorXs> xs, us, vs, lbdas;
    initializeSolution(xs, us, vs, lbdas);
    return std::make_tuple(std::move(xs), std::move(us), std::move(vs),
                           std::move(lbdas));
  }
 
private:
  InitializationStrategy xs_init_strategy_;
 
  // Check if the initial state is a StateErrorResidual.
  // Since this is a costly operation (dynamic_cast), we cache the result.
  bool checkInitCondIsStateError() const;
  bool init_cond_is_state_error_ = false;
};
 
namespace internal {
template <typename Scalar>
auto problem_last_state_space_helper(const TrajOptProblemTpl<Scalar> &problem) {
  return problem.term_cost_->space;
}

template <typename Scalar>
int problem_last_ndx_helper(const TrajOptProblemTpl<Scalar> &problem) {
  return problem_last_state_space_helper(problem)->ndx();
}
} // namespace internal
 
} // namespace aligator
 
#ifdef ALIGATOR_ENABLE_TEMPLATE_INSTANTIATION
#include "aligator/core/traj-opt-problem.txx"
#endif