pgs-solver.hpp

//
// Copyright (c) 2022-2024 INRIA
//

#pragma once

// IWYU pragma: begin_keep
#include <cassert>
#include <cstddef>
#include <optional>
#include <limits>
#include <vector>

#include <Eigen/Core>

#include <boost/fusion/container/vector/vector.hpp>

#include "pinocchio/macros.hpp"
#include "pinocchio/eigen-common.hpp"

#include "pinocchio/utils/check.hpp"
#include "pinocchio/utils/reference.hpp"

#include "pinocchio/container/eigen-storage.hpp"

#include "pinocchio/math.hpp"

#include "pinocchio/constraints.hpp"

#include "pinocchio/algorithm/solvers/fwd.hpp"
#include "pinocchio/algorithm/solvers/constraint-solver-base.hpp"
#include "pinocchio/algorithm/solvers/constraint-solver-utils.hpp"
// IWYU pragma: end_keep

namespace pinocchio
{
 // fwd declarations for PGS-internal structs
 // user-api structs are fwd delclared in solvers/fwd.hpp.
 // see below for definitions
 namespace internal
 {
 template<typename Scalar, int Options>
 struct PGSSolverWorkspaceTpl;
 }

 template<typename _Scalar, int _Options>
 struct traits<PGSConstraintSolverTpl<_Scalar, _Options>>
 {
 typedef _Scalar Scalar;
 static constexpr int Options = _Options;

 typedef PGSSolverSettingsTpl<Scalar> SolverSettings;
 typedef PGSSolverResultTpl<Scalar, Options> SolverResult;
 };

 template<typename _Scalar, int _Options>
 struct PGSConstraintSolverTpl : ConstraintSolverBase<PGSConstraintSolverTpl<_Scalar, _Options>>
 {
 typedef _Scalar Scalar;
 static constexpr int Options = _Options;
 typedef PGSConstraintSolverTpl Self;
 typedef ConstraintSolverBase<Self> Base;
 typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1, Options> VectorXs;
 typedef Eigen::Ref<const VectorXs> RefConstVectorXs;

 typedef internal::PGSSolverWorkspaceTpl<Scalar, Options> PGSSolverWorkspace;
 typedef PGSSolverSettingsTpl<Scalar> PGSSolverSettings;
 typedef PGSSolverResultTpl<Scalar, Options> PGSSolverResult;
 typedef PGSSolverStatsTpl<Scalar> PGSSolverStats;

 using Base::reset;
 using Base::timerStart;
 using Base::timerStop;

 explicit PGSConstraintSolverTpl(std::size_t max_problem_size = 0)
 : Base()
 , stats()
 , m_workspace(max_problem_size)
 , m_is_valid(false)
 {
 // we need to call reset - the solver needs to look as if it never ran
 reset();
 }

 bool isValid() const
 {
 return m_is_valid;
 }

 template<
 typename DelassusDerived,
 typename VectorLike,
 typename ConstraintModel,
 typename ConstraintModelAllocator,
 typename ConstraintData,
 typename ConstraintDataAllocator>
 bool solveImpl(
 DelassusOperatorBase<DelassusDerived> & delassus,
 const Eigen::MatrixBase<VectorLike> & g,
 const std::vector<ConstraintModel, ConstraintModelAllocator> & constraint_models,
 const std::vector<ConstraintData, ConstraintDataAllocator> & constraint_datas,
 const PGSSolverSettings & settings,
 PGSSolverResult & result);

 void resetImpl()
 {
 stats.reset();
 m_workspace.reset();
 m_is_valid = false;
 }

 PGSSolverStats stats;

 protected:
 PGSSolverWorkspace m_workspace;

 bool m_is_valid;

 }; // struct PGSConstraintSolverTpl

 template<typename _Scalar>
 struct traits<PGSSolverSettingsTpl<_Scalar>>
 {
 typedef _Scalar Scalar;
 };

 template<typename _Scalar>
 struct PGSSolverSettingsTpl : ConstraintSolverSettingsBase<PGSSolverSettingsTpl<_Scalar>>
 {
 typedef _Scalar Scalar;
 typedef PGSSolverSettingsTpl Self;
 typedef ConstraintSolverSettingsBase<Self> Base;
 typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1> VectorXs;

 PGSSolverSettingsTpl(
 std::size_t max_iterations = 1000,
 Scalar absolute_feasibility_tol = Scalar(1e-6),
 Scalar relative_feasibility_tol = Scalar(1e-6),
 Scalar absolute_complementarity_tol = Scalar(1e-6),
 Scalar relative_complementarity_tol = Scalar(1e-6),
 bool solve_ncp = true,
 bool measure_timings = false,
 bool stat_record = false,
 Scalar over_relaxation = Scalar(1))
 : Base(
 max_iterations,
 absolute_feasibility_tol,
 relative_feasibility_tol,
 absolute_complementarity_tol,
 relative_complementarity_tol,
 solve_ncp,
 measure_timings,
 stat_record)
 , over_relaxation(over_relaxation)
 {
 }

 void checkValidityImpl() const
 {
 PINOCCHIO_CHECK_INPUT_ARGUMENT(
 over_relaxation < Scalar(2) && over_relaxation > Scalar(0),
 "over_relaxation should lie in ]0,2[.");
 }

 // ----------------------
 // General settings

 using Base::max_iterations;

 using Base::absolute_feasibility_tol;

 using Base::relative_feasibility_tol;

 using Base::absolute_complementarity_tol;

 using Base::relative_complementarity_tol;

 using Base::solve_ncp;

 using Base::measure_timings;

 using Base::stat_record;

 // ----------------------
 // PGS specific settings

 Scalar over_relaxation;
 }; // struct PGSSolverSettingsTpl

 template<typename _Scalar, int _Options>
 struct traits<PGSSolverResultTpl<_Scalar, _Options>>
 {
 typedef _Scalar Scalar;
 static constexpr int Options = _Options;
 };

 template<typename _Scalar, int _Options>
 struct PGSSolverResultTpl : ConstraintSolverResultBase<PGSSolverResultTpl<_Scalar, _Options>>
 {
 typedef _Scalar Scalar;
 static constexpr int Options = _Options;
 typedef PGSSolverResultTpl Self;
 typedef ConstraintSolverResultBase<Self> Base;

 typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1, Options> VectorXs;
 typedef Eigen::Ref<const VectorXs> RefConstVectorXs;
 typedef internal::EigenStorageTpl<VectorXs> VectorXsStorage;

 using Base::constraintSize;
 using Base::setConstraintImpulseGuess;
 using Base::setConstraintVelocityGuess;

 PGSSolverResultTpl()
 : Base()
 , problem_size(0)
 , impulse_guess(std::nullopt)
 {
 }

 PGSSolverResultTpl(const PGSSolverResultTpl & other)
 {
 *this = other;
 }

 PGSSolverResultTpl & operator=(const PGSSolverResultTpl & other)
 {
 if (this != &other)
 {
 Base::operator=(other);
 problem_size = other.problem_size;
 // Since some members are maps reference on EigenStorage, we cannot simply copy them.
 // Thus we need to explicitly say we copy the storage, and the maps will automatically point
 // to the new storage.
 x_storage = other.x_storage;
 y_storage = other.y_storage;

 // special care must be taken for the optional guesses
 if (other.impulse_guess)
 {
 setConstraintImpulseGuess(other.impulse_guess.value());
 }
 else
 {
 clearConstraintImpulseGuessImpl();
 }
 }
 return *this;
 }

 int constraintSizeImpl() const
 {
 return static_cast<int>(problem_size);
 }

 void resetImpl()
 {
 impulse_guess.reset();

 // set solution to nan - solver has not run
 x.setConstant(std::numeric_limits<Scalar>::quiet_NaN());
 y.setConstant(std::numeric_limits<Scalar>::quiet_NaN());
 }

 void resize(std::size_t problem_size_)
 {
 problem_size = problem_size_;

 Eigen::Index np = static_cast<Eigen::Index>(problem_size);
 x_storage.resize(np);
 y_storage.resize(np);
 }

 template<typename VectorLike>
 void retrievePrimalSolution(const Eigen::MatrixBase<VectorLike> & primal_solution_) const
 {
 auto & primal_solution = primal_solution_.const_cast_derived();
 primal_solution = x;
 }

 template<typename VectorLike>
 void retrieveDualSolution(const Eigen::MatrixBase<VectorLike> & dual_solution_) const
 {
 auto & dual_solution = dual_solution_.const_cast_derived();
 dual_solution = y;
 }

 template<typename VectorLike>
 void setConstraintImpulseGuessImpl(const Eigen::MatrixBase<VectorLike> & impulse_guess_in)
 {
 m_impulse_guess_storage.resize(impulse_guess_in.size());
 m_impulse_guess = impulse_guess_in;
 impulse_guess.emplace(m_impulse_guess);
 }

 void clearConstraintImpulseGuessImpl()
 {
 m_impulse_guess_storage.resize(0);
 impulse_guess.reset();
 }

 template<typename VectorLike>
 void setConstraintVelocityGuessImpl(const Eigen::MatrixBase<VectorLike> & /*velocity_guess_*/)
 {
 // do nothing, no velocity guess for PGS
 }

 void clearConstraintVelocityGuessImpl()
 {
 // do nothing, no velocity guess for PGS
 }

 template<typename VectorLike>
 void
 retrieveConstraintImpulsesImpl(const Eigen::MatrixBase<VectorLike> & constraint_impulses_) const
 {
 auto & constraint_impulses = constraint_impulses_.const_cast_derived();
 constraint_impulses = x;
 }

 template<typename VectorLike>
 void retrieveConstraintVelocitiesImpl(
 const Eigen::MatrixBase<VectorLike> & constraint_velocities_) const
 {
 auto & constraint_velocities = constraint_velocities_.const_cast_derived();
 constraint_velocities = y;
 }

 using Base::iterations;

 using Base::converged;

 using Base::primal_feasibility;

 using Base::dual_feasibility;

 using Base::complementarity;

 std::size_t problem_size;

 // ----------------------
 // Solution warmstart

 std::optional<RefConstVectorXs> impulse_guess;

 // ----------------------
 // Solution - output of the solver

 VectorXsStorage x_storage;
 typename VectorXsStorage::RefMapType x = x_storage.map();

 VectorXsStorage y_storage;
 typename VectorXsStorage::RefMapType y = y_storage.map();

 VectorXsStorage m_impulse_guess_storage;
 typename VectorXsStorage::RefMapType m_impulse_guess = m_impulse_guess_storage.map();
 }; // struct PGSSolverResultTpl

 template<typename _Scalar>
 struct traits<PGSSolverStatsTpl<_Scalar>>
 {
 typedef _Scalar Scalar;
 };

 template<typename _Scalar>
 struct PGSSolverStatsTpl : ConstraintSolverStatsBase<PGSSolverStatsTpl<_Scalar>>
 {
 typedef _Scalar Scalar;
 typedef PGSSolverStatsTpl Self;
 typedef ConstraintSolverStatsBase<Self> Base;

 using Base::reserve;

 PGSSolverStatsTpl()
 : Base()
 {
 }

 explicit PGSSolverStatsTpl(std::size_t max_iterations)
 : Base(max_iterations)
 {
 reserve(max_iterations);
 }

 void reserveImpl(std::size_t /*max_iterations*/)
 {
 // No extra fields to reserve for PGS stats.
 }

 void resetImpl()
 {
 // No extra fields to reset for PGS stats.
 }

 using Base::iterations;

 using Base::primal_feasibility;

 using Base::dual_feasibility;

 using Base::dual_feasibility_ncp;

 using Base::complementarity;
 }; // struct PGSSolverStatsTpl

 namespace internal
 {
 template<typename _Scalar, int _Options>
 struct PGSSolverWorkspaceTpl
 {
 typedef _Scalar Scalar;
 static constexpr int Options = _Options;
 typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1, Options> VectorXs;
 typedef internal::EigenStorageTpl<VectorXs> VectorXsStorage;
 typedef Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Options> MatrixXs;
 typedef internal::EigenStorageTpl<MatrixXs> MatrixXsStorage;

 PGSSolverWorkspaceTpl(std::size_t problem_size = 0)
 : problem_size(problem_size)
 {
 resize(problem_size);
 reset();
 }

 PGSSolverWorkspaceTpl(const PGSSolverWorkspaceTpl & other)
 : PGSSolverWorkspaceTpl(0)
 {
 *this = other;
 }

 PGSSolverWorkspaceTpl & operator=(const PGSSolverWorkspaceTpl & other)
 {
 if (this != &other)
 {
 problem_size = other.problem_size;

 // Since some members are maps reference on EigenStorage, we cannot simply copy them.
 // Thus we need to explicitly say we copy the storage, and the maps will automatically
 // point to the new storage.
 delassus_matrix_storage = other.delassus_matrix_storage;
 x_storage = other.x_storage;
 x_previous_storage = other.x_previous_storage;
 y_storage = other.y_storage;
 rhs_storage = other.rhs_storage;
 tmp_storage = other.tmp_storage;
 }
 return *this;
 }

 void reset()
 {
 resize(problem_size);

#ifndef NDEBUG
 delassus_matrix.setConstant(std::numeric_limits<Scalar>::quiet_NaN());
 x.setConstant(std::numeric_limits<Scalar>::quiet_NaN());
 x_previous.setConstant(std::numeric_limits<Scalar>::quiet_NaN());
 y.setConstant(std::numeric_limits<Scalar>::quiet_NaN());
 tmp.setConstant(std::numeric_limits<Scalar>::quiet_NaN());
 rhs.setConstant(std::numeric_limits<Scalar>::quiet_NaN());
#endif
 }

 void resize(std::size_t problem_size_)
 {
 problem_size = problem_size_;

 Eigen::Index np = static_cast<Eigen::Index>(problem_size);
 delassus_matrix_storage.resize(np, np);
 x_storage.resize(np);
 x_previous_storage.resize(np);
 y_storage.resize(np);
 tmp_storage.resize(np);
 rhs_storage.resize(np);
 }

 std::size_t problem_size;

 MatrixXsStorage delassus_matrix_storage;
 typename MatrixXsStorage::RefMapType delassus_matrix = delassus_matrix_storage.map();

 VectorXsStorage x_storage;
 typename VectorXsStorage::RefMapType x = x_storage.map();

 VectorXsStorage x_previous_storage;
 typename VectorXsStorage::RefMapType x_previous = x_previous_storage.map();

 VectorXsStorage y_storage;
 typename VectorXsStorage::RefMapType y = y_storage.map();

 VectorXsStorage rhs_storage;
 typename VectorXsStorage::RefMapType rhs = rhs_storage.map();

 VectorXsStorage tmp_storage;
 typename VectorXsStorage::RefMapType tmp = tmp_storage.map();
 }; // struct PGSSolverWorkspaceTpl
 } // namespace internal

} // namespace pinocchio

// IWYU pragma: begin_exports
#include "pinocchio/src/algorithm/solvers/pgs-solver.hxx"
// IWYU pragma: end_exports