kkt.cpp

// BSD 3-Clause License
//
// Copyright (C) 2019-2020, LAAS-CNRS, New York University, Max Planck
// Gesellschaft,
// University of Edinburgh
// Copyright note valid unless otherwise stated in individual files.
// All rights reserved.

#include "crocoddyl/core/solvers/kkt.hpp"

namespace crocoddyl {

SolverKKT::SolverKKT(std::shared_ptr<ShootingProblem> problem)
 : SolverAbstract(problem),
 reg_incfactor_(10.),
 reg_decfactor_(10.),
 reg_min_(1e-9),
 reg_max_(1e9),
 cost_try_(0.),
 th_grad_(1e-12),
 was_feasible_(false) {
 allocateData();
 const std::size_t n_alphas = 10;
 preg_ = 0.;
 dreg_ = 0.;
 alphas_.resize(n_alphas);
 for (std::size_t n = 0; n < n_alphas; ++n) {
 alphas_[n] = 1. / pow(2., (double)n);
 }
}

SolverKKT::~SolverKKT() {}

bool SolverKKT::solve(const std::vector<Eigen::VectorXd>& init_xs,
 const std::vector<Eigen::VectorXd>& init_us,
 const std::size_t maxiter, const bool is_feasible,
 const double) {
 setCandidate(init_xs, init_us, is_feasible);
 bool recalc = true;
 for (iter_ = 0; iter_ < maxiter; ++iter_) {
 while (true) {
 try {
 computeDirection(recalc);
 } catch (std::exception& e) {
 recalc = false;
 if (preg_ == reg_max_) {
 return false;
 } else {
 continue;
 }
 }
 break;
 }

 expectedImprovement();
 for (std::vector<double>::const_iterator it = alphas_.begin();
 it != alphas_.end(); ++it) {
 steplength_ = *it;
 try {
 dV_ = tryStep(steplength_);
 } catch (std::exception& e) {
 continue;
 }
 dVexp_ = steplength_ * d_[0] + 0.5 * steplength_ * steplength_ * d_[1];
 if (d_[0] < th_grad_ || !is_feasible_ || dV_ > th_acceptstep_ * dVexp_) {
 was_feasible_ = is_feasible_;
 setCandidate(xs_try_, us_try_, true);
 cost_ = cost_try_;
 break;
 }
 }
 stoppingCriteria();
 const std::size_t n_callbacks = callbacks_.size();
 if (n_callbacks != 0) {
 for (std::size_t c = 0; c < n_callbacks; ++c) {
 CallbackAbstract& callback = *callbacks_[c];
 callback(*this);
 }
 }
 if (was_feasible_ && stop_ < th_stop_) {
 return true;
 }
 }
 return false;
}

void SolverKKT::computeDirection(const bool recalc) {
 const std::size_t T = problem_->get_T();
 if (recalc) {
 calcDiff();
 }
 computePrimalDual();
 const Eigen::VectorBlock<Eigen::VectorXd, Eigen::Dynamic> p_x =
 primal_.segment(0, ndx_);
 const Eigen::VectorBlock<Eigen::VectorXd, Eigen::Dynamic> p_u =
 primal_.segment(ndx_, nu_);

 std::size_t ix = 0;
 std::size_t iu = 0;
 const std::vector<std::shared_ptr<ActionModelAbstract> >& models =
 problem_->get_runningModels();
 for (std::size_t t = 0; t < T; ++t) {
 const std::size_t ndxi = models[t]->get_state()->get_ndx();
 const std::size_t nui = models[t]->get_nu();
 dxs_[t] = p_x.segment(ix, ndxi);
 dus_[t] = p_u.segment(iu, nui);
 lambdas_[t] = dual_.segment(ix, ndxi);
 ix += ndxi;
 iu += nui;
 }
 const std::size_t ndxi =
 problem_->get_terminalModel()->get_state()->get_ndx();
 dxs_.back() = p_x.segment(ix, ndxi);
 lambdas_.back() = dual_.segment(ix, ndxi);
}

double SolverKKT::tryStep(const double steplength) {
 const std::size_t T = problem_->get_T();
 const std::vector<std::shared_ptr<ActionModelAbstract> >& models =
 problem_->get_runningModels();
 for (std::size_t t = 0; t < T; ++t) {
 const std::shared_ptr<ActionModelAbstract>& m = models[t];

 m->get_state()->integrate(xs_[t], steplength * dxs_[t], xs_try_[t]);
 if (m->get_nu() != 0) {
 us_try_[t] = us_[t];
 us_try_[t] += steplength * dus_[t];
 }
 }
 const std::shared_ptr<ActionModelAbstract> m = problem_->get_terminalModel();
 m->get_state()->integrate(xs_[T], steplength * dxs_[T], xs_try_[T]);
 cost_try_ = problem_->calc(xs_try_, us_try_);
 return cost_ - cost_try_;
}

double SolverKKT::stoppingCriteria() {
 const std::size_t T = problem_->get_T();
 std::size_t ix = 0;
 std::size_t iu = 0;
 const std::vector<std::shared_ptr<ActionModelAbstract> >& models =
 problem_->get_runningModels();
 const std::vector<std::shared_ptr<ActionDataAbstract> >& datas =
 problem_->get_runningDatas();
 for (std::size_t t = 0; t < T; ++t) {
 const std::shared_ptr<ActionDataAbstract>& d = datas[t];
 const std::size_t ndxi = models[t]->get_state()->get_ndx();
 const std::size_t nui = models[t]->get_nu();

 dF.segment(ix, ndxi) = lambdas_[t];
 dF.segment(ix, ndxi).noalias() -= d->Fx.transpose() * lambdas_[t + 1];
 dF.segment(ndx_ + iu, nui).noalias() = -lambdas_[t + 1].transpose() * d->Fu;
 ix += ndxi;
 iu += nui;
 }
 const std::size_t ndxi =
 problem_->get_terminalModel()->get_state()->get_ndx();
 dF.segment(ix, ndxi) = lambdas_.back();
 stop_ = (kktref_.segment(0, ndx_ + nu_) + dF).squaredNorm() +
 kktref_.segment(ndx_ + nu_, ndx_).squaredNorm();
 return stop_;
}

const Eigen::Vector2d& SolverKKT::expectedImprovement() {
 d_ = Eigen::Vector2d::Zero();
 // -grad^T.primal
 d_(0) = -kktref_.segment(0, ndx_ + nu_).dot(primal_);
 // -(hessian.primal)^T.primal
 kkt_primal_.noalias() = kkt_.block(0, 0, ndx_ + nu_, ndx_ + nu_) * primal_;
 d_(1) = -kkt_primal_.dot(primal_);
 return d_;
}

const Eigen::MatrixXd& SolverKKT::get_kkt() const { return kkt_; }

const Eigen::VectorXd& SolverKKT::get_kktref() const { return kktref_; }

const Eigen::VectorXd& SolverKKT::get_primaldual() const { return primaldual_; }

const std::vector<Eigen::VectorXd>& SolverKKT::get_dxs() const { return dxs_; }

const std::vector<Eigen::VectorXd>& SolverKKT::get_dus() const { return dus_; }

const std::vector<Eigen::VectorXd>& SolverKKT::get_lambdas() const {
 return lambdas_;
}

std::size_t SolverKKT::get_nx() const { return nx_; }

std::size_t SolverKKT::get_ndx() const { return ndx_; }

std::size_t SolverKKT::get_nu() const { return nu_; }

double SolverKKT::calcDiff() {
 cost_ = problem_->calc(xs_, us_);
 cost_ = problem_->calcDiff(xs_, us_);

 // offset on constraint xnext = f(x,u) due to x0 = ref.
 const std::size_t cx0 =
 problem_->get_runningModels()[0]->get_state()->get_ndx();

 std::size_t ix = 0;
 std::size_t iu = 0;
 const std::size_t T = problem_->get_T();
 kkt_.block(ndx_ + nu_, 0, ndx_, ndx_).setIdentity();
 for (std::size_t t = 0; t < T; ++t) {
 const std::shared_ptr<ActionModelAbstract>& m =
 problem_->get_runningModels()[t];
 const std::shared_ptr<ActionDataAbstract>& d =
 problem_->get_runningDatas()[t];
 const std::size_t ndxi = m->get_state()->get_ndx();
 const std::size_t nui = m->get_nu();

 // Computing the gap at the initial state
 if (t == 0) {
 m->get_state()->diff(problem_->get_x0(), xs_[0],
 kktref_.segment(ndx_ + nu_, ndxi));
 }

 // Filling KKT matrix
 kkt_.block(ix, ix, ndxi, ndxi) = d->Lxx;
 kkt_.block(ix, ndx_ + iu, ndxi, nui) = d->Lxu;
 kkt_.block(ndx_ + iu, ix, nui, ndxi) = d->Lxu.transpose();
 kkt_.block(ndx_ + iu, ndx_ + iu, nui, nui) = d->Luu;
 kkt_.block(ndx_ + nu_ + cx0 + ix, ix, ndxi, ndxi) = -d->Fx;
 kkt_.block(ndx_ + nu_ + cx0 + ix, ndx_ + iu, ndxi, nui) = -d->Fu;

 // Filling KKT vector
 kktref_.segment(ix, ndxi) = d->Lx;
 kktref_.segment(ndx_ + iu, nui) = d->Lu;
 m->get_state()->diff(d->xnext, xs_[t + 1],
 kktref_.segment(ndx_ + nu_ + cx0 + ix, ndxi));

 ix += ndxi;
 iu += nui;
 }
 const std::shared_ptr<ActionDataAbstract>& df = problem_->get_terminalData();
 const std::size_t ndxf =
 problem_->get_terminalModel()->get_state()->get_ndx();
 kkt_.block(ix, ix, ndxf, ndxf) = df->Lxx;
 kktref_.segment(ix, ndxf) = df->Lx;
 kkt_.block(0, ndx_ + nu_, ndx_ + nu_, ndx_) =
 kkt_.block(ndx_ + nu_, 0, ndx_, ndx_ + nu_).transpose();
 return cost_;
}

void SolverKKT::computePrimalDual() {
 primaldual_ = kkt_.lu().solve(-kktref_);
 primal_ = primaldual_.segment(0, ndx_ + nu_);
 dual_ = primaldual_.segment(ndx_ + nu_, ndx_);
}

void SolverKKT::increaseRegularization() {
 preg_ *= reg_incfactor_;
 if (preg_ > reg_max_) {
 preg_ = reg_max_;
 }
 dreg_ = preg_;
}

void SolverKKT::decreaseRegularization() {
 preg_ /= reg_decfactor_;
 if (preg_ < reg_min_) {
 preg_ = reg_min_;
 }
 dreg_ = preg_;
}

void SolverKKT::allocateData() {
 const std::size_t T = problem_->get_T();
 dxs_.resize(T + 1);
 dus_.resize(T);
 lambdas_.resize(T + 1);
 xs_try_.resize(T + 1);
 us_try_.resize(T);

 nx_ = 0;
 ndx_ = 0;
 nu_ = 0;
 const std::size_t nx = problem_->get_nx();
 const std::size_t ndx = problem_->get_ndx();
 const std::vector<std::shared_ptr<ActionModelAbstract> >& models =
 problem_->get_runningModels();
 for (std::size_t t = 0; t < T; ++t) {
 const std::shared_ptr<ActionModelAbstract>& model = models[t];
 if (t == 0) {
 xs_try_[t] = problem_->get_x0();
 } else {
 xs_try_[t] = Eigen::VectorXd::Constant(nx, NAN);
 }
 const std::size_t nu = model->get_nu();
 us_try_[t] = Eigen::VectorXd::Constant(nu, NAN);
 dxs_[t] = Eigen::VectorXd::Zero(ndx);
 dus_[t] = Eigen::VectorXd::Zero(nu);
 lambdas_[t] = Eigen::VectorXd::Zero(ndx);
 nx_ += nx;
 ndx_ += ndx;
 nu_ += nu;
 }
 nx_ += nx;
 ndx_ += ndx;
 xs_try_.back() = problem_->get_terminalModel()->get_state()->zero();
 dxs_.back() = Eigen::VectorXd::Zero(ndx);
 lambdas_.back() = Eigen::VectorXd::Zero(ndx);

 // Set dimensions for kkt matrix and kkt_ref vector
 kkt_.resize(2 * ndx_ + nu_, 2 * ndx_ + nu_);
 kkt_.setZero();
 kktref_.resize(2 * ndx_ + nu_);
 kktref_.setZero();
 primaldual_.resize(2 * ndx_ + nu_);
 primaldual_.setZero();
 primal_.resize(ndx_ + nu_);
 primal_.setZero();
 kkt_primal_.resize(ndx_ + nu_);
 kkt_primal_.setZero();
 dual_.resize(ndx_);
 dual_.setZero();
 dF.resize(ndx_ + nu_);
 dF.setZero();
}

} // namespace crocoddyl