TermTaskspacePoseCG.hpp

Go to the documentation of this file.

/**********************************************************************************************************************
This file is part of the Control Toolbox (https://github.com/ethz-adrl/control-toolbox), copyright by ETH Zurich.
Authors:  Michael Neunert, Markus Giftthaler, Markus Stäuble, Diego Pardo, Farbod Farshidian
Licensed under the BSD-2 license (see LICENSE file in main directory)
**********************************************************************************************************************/

#pragma once

#ifdef CPPADCG

#include <ct/optcon/optcon.h>

#include <ct/rbd/robot/Kinematics.h>
#include <ct/rbd/state/RBDState.h>

namespace ct {
namespace rbd {

template <class KINEMATICS, bool FLOATING_BASE, size_t STATE_DIM, size_t CONTROL_DIM>
class TermTaskspacePoseCG : public optcon::TermBase<STATE_DIM, CONTROL_DIM, double, double>
{
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW

    /*
     * The AD parameter vector consists of:
     * - states (STATE_DIM parameters)
     * - controls (CONTROL_DIM parameters)
     * - desired end-effector positions (3 parameters)
     * - desired end-effector orientation (9 parameters)
     * - time (1 parameter)
     */
    static const size_t AD_PARAMETER_DIM = STATE_DIM + CONTROL_DIM + 3 + 9 + 1;

    using BASE = optcon::TermBase<STATE_DIM, CONTROL_DIM, double, double>;

    using DerivativesCppadJIT = core::DerivativesCppadJIT<AD_PARAMETER_DIM, 1>;
    using CGScalar = typename DerivativesCppadJIT::CG_SCALAR;

    using RBDStateTpl = ct::rbd::RBDState<KINEMATICS::NJOINTS, CGScalar>;

    using state_matrix_t = Eigen::Matrix<double, STATE_DIM, STATE_DIM>;
    using control_matrix_t = Eigen::Matrix<double, CONTROL_DIM, CONTROL_DIM>;
    using control_state_matrix_t = Eigen::Matrix<double, CONTROL_DIM, STATE_DIM>;


    TermTaskspacePoseCG() = default;

    TermTaskspacePoseCG(size_t eeInd,
        const Eigen::Matrix3d& Qpos,
        const double& Qrot,
        const core::StateVector<3>& w_pos_des,
        const Eigen::Quaterniond& w_q_des,
        const std::string& name = "TermTaskSpace",
        bool evalControlDerivatives = true)
        : BASE(name), eeInd_(eeInd), Q_pos_(Qpos), Q_rot_(Qrot), evalControlDerivatives_(evalControlDerivatives)
    {
        setReferencePosition(w_pos_des);
        setReferenceOrientation(w_q_des);
        setup();
    }

    TermTaskspacePoseCG(size_t eeInd,
        const Eigen::Matrix3d& Qpos,
        const double& Qrot,
        const ct::rbd::RigidBodyPose& rbdPose,
        const std::string& name = "TermTaskSpace",
        bool evalControlDerivatives = true)
        // delegate constructor
        : TermTaskspacePoseCG(eeInd,
              Qpos,
              Qrot,
              rbdPose.position().toImplementation(),
              rbdPose.getRotationQuaternion().toImplementation(),
              name,
              evalControlDerivatives)
    {
    }

    TermTaskspacePoseCG(size_t eeInd,
        const Eigen::Matrix3d& Qpos,
        const double& Qrot,
        const core::StateVector<3>& w_pos_des,
        const Eigen::Matrix3d& eulerXyz,
        const std::string& name = "TermTaskSpace",
        bool evalControlDerivatives = true)
        // delegate constructor
        : TermTaskspacePoseCG(eeInd,
              Qpos,
              Qrot,
              w_pos_des,
              Eigen::Quaterniond(Eigen::AngleAxisd(eulerXyz(0), Eigen::Vector3d::UnitX()) *
                                 Eigen::AngleAxisd(eulerXyz(1), Eigen::Vector3d::UnitY()) *
                                 Eigen::AngleAxisd(eulerXyz(2), Eigen::Vector3d::UnitZ())),
              name,
              evalControlDerivatives)
    {
    }

    TermTaskspacePoseCG(size_t eeInd,
        const Eigen::Matrix3d& Qpos,
        const double& Qrot,
        const std::string& name = "TermTaskSpace",
        bool evalControlDerivatives = false)
        : BASE(name), eeInd_(eeInd), Q_pos_(Qpos), Q_rot_(Qrot), evalControlDerivatives_(evalControlDerivatives)
    {
        // arbitrary dummy values
        Eigen::Quaterniond w_q_des(0.0, 0.0, 0.0, 1.0);
        setReferenceOrientation(w_q_des);
        setReferencePosition(core::StateVector<3>::Zero());
        setup();
    }

    TermTaskspacePoseCG(const std::string& configFile, const std::string& termName, bool verbose = false)
    {
        loadConfigFile(configFile, termName, verbose);
    }

    TermTaskspacePoseCG(const TermTaskspacePoseCG& arg)
        : BASE(arg),
          eeInd_(arg.eeInd_),
          kinematics_(KINEMATICS()),
          Q_pos_(arg.Q_pos_),
          Q_rot_(arg.Q_rot_),
          costFun_(arg.costFun_),
          adParameterVector_(arg.adParameterVector_),
          evalControlDerivatives_(arg.evalControlDerivatives_)
    {
        derivativesCppadJIT_ = std::shared_ptr<DerivativesCppadJIT>(arg.derivativesCppadJIT_->clone());
    }

    virtual ~TermTaskspacePoseCG() {}
    virtual TermTaskspacePoseCG<KINEMATICS, FLOATING_BASE, STATE_DIM, CONTROL_DIM>* clone() const override
    {
        return new TermTaskspacePoseCG(*this);
    }

    void setup()
    {
        // map term evaluation function to AD function
        costFun_ = [&](const Eigen::Matrix<CGScalar, AD_PARAMETER_DIM, 1>& inputParams) {
            return evalLocal<CGScalar>(inputParams);
        };

        // set up derivatives
        derivativesCppadJIT_ =
            std::shared_ptr<DerivativesCppadJIT>(new DerivativesCppadJIT(costFun_, AD_PARAMETER_DIM, 1));

        ct::core::DerivativesCppadSettings settings;
        settings.createForwardZero_ = true;
        settings.createJacobian_ = true;
        settings.createHessian_ = true;

        std::cout << "compiling JIT for TermTaskspacePoseCG" << std::endl;
        derivativesCppadJIT_->compileJIT(settings, "TermTaskspacePoseCGcosts");
    }


    void setCurrentStateAndControl(const Eigen::Matrix<double, STATE_DIM, 1>& x,
        const Eigen::Matrix<double, CONTROL_DIM, 1>& u,
        const double& t)
    {
        adParameterVector_.template segment<STATE_DIM>(0) = x;
        adParameterVector_.template segment<CONTROL_DIM>(STATE_DIM) = u;
        adParameterVector_(AD_PARAMETER_DIM - 1) = t;
    }


    virtual double evaluate(const Eigen::Matrix<double, STATE_DIM, 1>& x,
        const Eigen::Matrix<double, CONTROL_DIM, 1>& u,
        const double& t) override
    {
        setCurrentStateAndControl(x, u, t);
        return derivativesCppadJIT_->forwardZero(adParameterVector_)(0);
    }


    virtual core::StateVector<STATE_DIM> stateDerivative(const core::StateVector<STATE_DIM>& x,
        const core::ControlVector<CONTROL_DIM>& u,
        const double& t) override
    {
        setCurrentStateAndControl(x, u, t);
        Eigen::Matrix<double, 1, AD_PARAMETER_DIM> jacTot = derivativesCppadJIT_->jacobian(adParameterVector_);
        return jacTot.template leftCols<STATE_DIM>().transpose();
    }

    virtual core::ControlVector<CONTROL_DIM> controlDerivative(const core::StateVector<STATE_DIM>& x,
        const core::ControlVector<CONTROL_DIM>& u,
        const double& t) override
    {
        if (evalControlDerivatives_)
        {
            setCurrentStateAndControl(x, u, t);
            Eigen::Matrix<double, 1, AD_PARAMETER_DIM> jacTot = derivativesCppadJIT_->jacobian(adParameterVector_);
            return jacTot.template block<1, CONTROL_DIM>(0, STATE_DIM).transpose();
        }
        else
            return core::ControlVector<CONTROL_DIM>::Zero();
    }

    virtual state_matrix_t stateSecondDerivative(const core::StateVector<STATE_DIM>& x,
        const core::ControlVector<CONTROL_DIM>& u,
        const double& t) override
    {
        setCurrentStateAndControl(x, u, t);
        Eigen::Matrix<double, 1, 1> w(1.0);
        Eigen::MatrixXd hesTot = derivativesCppadJIT_->hessian(adParameterVector_, w);
        return hesTot.template block<STATE_DIM, STATE_DIM>(0, 0);
    }

    virtual control_matrix_t controlSecondDerivative(const core::StateVector<STATE_DIM>& x,
        const core::ControlVector<CONTROL_DIM>& u,
        const double& t) override
    {
        if (evalControlDerivatives_)
        {
            setCurrentStateAndControl(x, u, t);
            Eigen::Matrix<double, 1, 1> w(1.0);
            Eigen::MatrixXd hesTot = derivativesCppadJIT_->hessian(adParameterVector_, w);
            return hesTot.template block<CONTROL_DIM, CONTROL_DIM>(STATE_DIM, STATE_DIM);
        }
        else
            return control_matrix_t::Zero();
    }

    virtual control_state_matrix_t stateControlDerivative(const core::StateVector<STATE_DIM>& x,
        const core::ControlVector<CONTROL_DIM>& u,
        const double& t) override
    {
        if (evalControlDerivatives_)
        {
            setCurrentStateAndControl(x, u, t);
            Eigen::Matrix<double, 1, 1> w(1.0);
            Eigen::MatrixXd hesTot = derivativesCppadJIT_->hessian(adParameterVector_, w);
            return hesTot.template block<CONTROL_DIM, STATE_DIM>(STATE_DIM, 0);
        }
        else
            return control_state_matrix_t::Zero();
    }


    void loadConfigFile(const std::string& filename, const std::string& termName, bool verbose = false) override
    {
        if (verbose)
            std::cout << "Loading TermTaskspacePoseCG from file " << filename << std::endl;

        ct::optcon::loadScalarCF(filename, "eeId", eeInd_, termName);
        ct::optcon::loadScalarCF(filename, "Q_rot", Q_rot_, termName);

        ct::optcon::loadMatrixCF(filename, "Q_pos", Q_pos_, termName);

        Eigen::Matrix<double, 3, 1> w_p_ref;
        ct::optcon::loadMatrixCF(filename, "x_des", w_p_ref, termName);
        setReferencePosition(w_p_ref);

        ct::optcon::loadScalarCF(filename, "evalControlDerivatives", evalControlDerivatives_, termName);


        // try loading euler angles
        if (verbose)
            std::cout << "trying to load euler angles" << std::endl;
        try
        {
            Eigen::Vector3d eulerXyz;
            ct::optcon::loadMatrixCF(filename, "eulerXyz_des", eulerXyz, termName);
            Eigen::Quaternion<double> quat_des(Eigen::AngleAxisd(eulerXyz(0), Eigen::Vector3d::UnitX()) *
                                               Eigen::AngleAxisd(eulerXyz(1), Eigen::Vector3d::UnitY()) *
                                               Eigen::AngleAxisd(eulerXyz(2), Eigen::Vector3d::UnitZ()));
            setReferenceOrientation(quat_des);
            if (verbose)
                std::cout << "Read desired Euler Angles Xyz as  = \n" << eulerXyz.transpose() << std::endl;
        } catch (const std::exception& e)
        {
            throw std::runtime_error(
                "Failed to load TermTaskspacePoseCG, could not find a desired end effector orientation in file.");
        }

        if (verbose)
        {
            std::cout << "Read eeId as eeId = \n" << eeInd_ << std::endl;
            std::cout << "Read Q_pos as Q_pos = \n" << Q_pos_ << std::endl;
            std::cout << "Read Q_rot as Q_rot = \n" << Q_rot_ << std::endl;
            std::cout << "Read x_des as x_des = \n" << getReferencePosition().transpose() << std::endl;
            std::cout << "Read evalControlDerivatives as evalControlDerivatives = \n"
                      << evalControlDerivatives_ << std::endl;
        }

        // new initialization required
        setup();
    }


    const Eigen::Matrix<double, 3, 1> getReferencePosition() const
    {
        return adParameterVector_.template segment<3>(STATE_DIM + CONTROL_DIM);
    }

    void setReferencePosition(Eigen::Matrix<double, 3, 1> w_p_ref)
    {
        adParameterVector_.template segment<3>(STATE_DIM + CONTROL_DIM) = w_p_ref;
    }

    const Eigen::Quaterniond getReferenceOrientation() const
    {
        return Eigen::Quaterniond(extractReferenceRotationMatrix(adParameterVector_));
    }

    void setReferenceOrientation(const Eigen::Matrix<double, 3, 3>& w_R_ref)
    {
        // transcribe the rotation matrix into the parameter vector
        const Eigen::Matrix<double, 9, 1> matVectorized(
            Eigen::Map<const Eigen::Matrix<double, 9, 1>>(w_R_ref.data(), 9));
        adParameterVector_.template segment<9>(STATE_DIM + CONTROL_DIM + 3) = matVectorized;
    }

    void setReferenceOrientation(const Eigen::Quaterniond& w_q_des)
    {
        setReferenceOrientation(w_q_des.normalized().toRotationMatrix());
    }

    const ct::rbd::RigidBodyPose getReferencePose() const
    {
        return ct::rbd::RigidBodyPose(getReferenceOrientation(), getReferencePosition());
    }


protected:
    virtual RBDStateTpl setStateFromVector(const Eigen::Matrix<CGScalar, STATE_DIM, 1>& x)
    {
        return setStateFromVector_specialized<FLOATING_BASE>(x);
    }

private:
    template <typename SC>
    const Eigen::Matrix<SC, 3, 3> extractReferenceRotationMatrix(
        const Eigen::Matrix<SC, AD_PARAMETER_DIM, 1>& parameterVector) const
    {
        Eigen::Matrix<SC, 9, 1> matVectorized = parameterVector.template segment<9>(STATE_DIM + CONTROL_DIM + 3);
        Eigen::Matrix<SC, 3, 3> w_R_ee(Eigen::Map<Eigen::Matrix<SC, 3, 3>>(matVectorized.data(), 3, 3));
        return w_R_ee;
    }

    template <typename SC>
    const Eigen::Matrix<SC, STATE_DIM, 1> extractStateVector(
        const Eigen::Matrix<SC, AD_PARAMETER_DIM, 1>& parameterVector) const
    {
        return parameterVector.template segment<STATE_DIM>(0);
    }

    template <typename SC>
    const Eigen::Matrix<SC, 3, 1> extractReferencePosition(
        const Eigen::Matrix<SC, AD_PARAMETER_DIM, 1>& parameterVector) const
    {
        return parameterVector.template segment<3>(STATE_DIM + CONTROL_DIM);
    }

    template <bool FB>
    RBDStateTpl setStateFromVector_specialized(const Eigen::Matrix<CGScalar, -1, 1>& x,
        typename std::enable_if<FB, bool>::type = true)
    {
        RBDStateTpl rbdState;
        rbdState.fromStateVectorEulerXyz(x);
        return rbdState;
    }

    template <bool FB>
    RBDStateTpl setStateFromVector_specialized(const Eigen::Matrix<CGScalar, -1, 1>& x,
        typename std::enable_if<!FB, bool>::type = true)
    {
        RBDStateTpl rbdState;
        rbdState.joints().toImplementation() = x.head(STATE_DIM);
        return rbdState;
    }


    template <typename SC>
    Eigen::Matrix<SC, 1, 1> evalLocal(const Eigen::Matrix<SC, AD_PARAMETER_DIM, 1>& adParams)
    {
        // extract state, plus position and rotation references
        Eigen::Matrix<SC, STATE_DIM, 1> x = extractStateVector(adParams);
        Eigen::Matrix<SC, 3, 1> w_p_ref = extractReferencePosition(adParams);
        Eigen::Matrix<SC, 3, 3> w_R_ref = extractReferenceRotationMatrix<SC>(adParams);

        // transform the robot state vector into a CT RBDState
        RBDState<KINEMATICS::NJOINTS, SC> rbdState = setStateFromVector(x);

        // position difference in world frame
        Eigen::Matrix<SC, 3, 1> xCurr =
            kinematics_.getEEPositionInWorld(eeInd_, rbdState.basePose(), rbdState.jointPositions()).toImplementation();
        Eigen::Matrix<SC, 3, 1> xDiff = xCurr - w_p_ref;


        // compute the cost from the position error
        SC pos_cost = (xDiff.transpose() * Q_pos_.template cast<SC>() * xDiff)(0, 0);


        // get current end-effector rotation in world frame
        Eigen::Matrix<SC, 3, 3> ee_rot =
            kinematics_.getEERotInWorld(eeInd_, rbdState.basePose(), rbdState.jointPositions());

        // compute a measure for the difference between current rotation and desired rotation and compute cost based on the orientation error
        // for the intuition behind, consider the following posts:
        // https://math.stackexchange.com/a/87698
        // https://math.stackexchange.com/a/773635
        Eigen::Matrix<SC, 3, 3> ee_R_diff = w_R_ref.transpose() * ee_rot;

        // compute rotation penalty using the Frobenius norm of (R_diff-I)
        SC rot_cost = (SC)Q_rot_ * (ee_R_diff - Eigen::Matrix<SC, 3, 3>::Identity()).squaredNorm();

        Eigen::Matrix<SC, 1, 1> result;
        result(0, 0) = (pos_cost + rot_cost);
        return result;
    }

    size_t eeInd_;

    KINEMATICS kinematics_;

    Eigen::Matrix3d Q_pos_;

    double Q_rot_;

    std::shared_ptr<DerivativesCppadJIT> derivativesCppadJIT_;

    typename DerivativesCppadJIT::FUN_TYPE_CG costFun_;

    Eigen::Matrix<double, AD_PARAMETER_DIM, 1> adParameterVector_;

    bool evalControlDerivatives_;
};


}  // namespace rbd
}  // namespace ct

#endif