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"
#include "aligator/third-party/polymorphic_cxx14.h"
 
namespace aligator {


template <typename T1, typename T2>
[[nodiscard]] bool assign_no_resize(const std::vector<T1> &lhs,
                                    std::vector<T2> &rhs) {
  static_assert(std::is_base_of_v<Eigen::EigenBase<T1>, T1>,
                "T1 should be an Eigen object!");
  static_assert(std::is_base_of_v<Eigen::EigenBase<T2>, T2>,
                "T2 should be an Eigen object!");
  if (lhs.size() != rhs.size())
    return false;
 
  const auto same_dims = [](auto &x, auto &y) {
    return (x.cols() == y.cols()) && (x.rows() == y.rows());
  };
 
  for (std::size_t i = 0; i < lhs.size(); i++) {
    if (!same_dims(lhs[i], rhs[i]))
      return false;
    rhs[i] = lhs[i];
  }
  return true;
}

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>;
  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;
  }

  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;

  xyz::polymorphic<StageModel>
  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;
};

template <typename Scalar>
void xs_default_init(const TrajOptProblemTpl<Scalar> &problem,
                     std::vector<typename math_types<Scalar>::VectorXs> &xs) {
  const std::size_t nsteps = problem.numSteps();
  xs.resize(nsteps + 1);
  if (problem.initCondIsStateError()) {
    xs[0] = problem.getInitState();
  } else {
    if (problem.stages_.size() > 0) {
      xs[0] = problem.stages_[0]->xspace().neutral();
    } else {
      ALIGATOR_RUNTIME_ERROR(
          "The problem should have either a StateErrorResidual as an initial "
          "condition or at least one stage.");
    }
  }
  for (std::size_t i = 0; i < nsteps; i++) {
    const StageModelTpl<Scalar> &sm = *problem.stages_[i];
    xs[i + 1] = sm.xspace_next().neutral();
  }
}

template <typename Scalar>
void us_default_init(const TrajOptProblemTpl<Scalar> &problem,
                     std::vector<typename math_types<Scalar>::VectorXs> &us) {
  const std::size_t nsteps = problem.numSteps();
  us.resize(nsteps);
  for (std::size_t i = 0; i < nsteps; i++) {
    const StageModelTpl<Scalar> &sm = *problem.stages_[i];
    us[i] = sm.uspace().neutral();
  }
}

template <typename Scalar>
void check_initial_guess_and_assign(
    const TrajOptProblemTpl<Scalar> &problem,
    const typename math_types<Scalar>::VectorOfVectors &xs_in,
    const typename math_types<Scalar>::VectorOfVectors &us_in,
    typename math_types<Scalar>::VectorOfVectors &xs_out,
    typename math_types<Scalar>::VectorOfVectors &us_out) {
  if (xs_in.empty()) {
    problem.initializeSolution(xs_out, us_out);
  } else if (us_in.empty()) {
    us_default_init(problem, us_out);
  } else {
    if (!assign_no_resize(xs_in, xs_out))
      ALIGATOR_RUNTIME_ERROR("warm-start for xs has wrong size!");
    if (!assign_no_resize(us_in, us_out))
      ALIGATOR_RUNTIME_ERROR("warm-start for us has wrong size!");
  }
}
 
namespace internal {
template <typename Scalar>
int problem_last_ndx_helper(const TrajOptProblemTpl<Scalar> &problem) {
  return problem.term_cost_->ndx();
}
} // namespace internal
 
#ifdef ALIGATOR_ENABLE_TEMPLATE_INSTANTIATION
extern template struct TrajOptProblemTpl<context::Scalar>;
#endif
} // namespace aligator