robot_model.hpp

Go to the documentation of this file.

#pragma once

#include "hebi.h"

#include <memory>
#include <vector>

#include "Eigen/Eigen"
#include "util.hpp"

using namespace Eigen;

namespace hebi {
namespace robot_model {

using Matrix4dVector = std::vector<Matrix4d, Eigen::aligned_allocator<Eigen::Matrix4d>>;
using MatrixXdVector = std::vector<MatrixXd, Eigen::aligned_allocator<Eigen::MatrixXd>>;
// The result of an IK operation.  More fields will be added to this structure
// in future API releases.
struct IKResult {
  HebiStatusCode result; // Success or failure
};

class RobotModel;

class Objective {
  friend RobotModel;

public:
  virtual ~Objective() {}

protected:
  virtual HebiStatusCode addObjective(HebiIKPtr ik) const = 0;
};

class EndEffectorPositionObjective final : public Objective {
public:
  EndEffectorPositionObjective(const Eigen::Vector3d&);
  EndEffectorPositionObjective(double weight, const Eigen::Vector3d&);

private:
  HebiStatusCode addObjective(HebiIKPtr ik) const override;
  double _weight, _x, _y, _z;
};

class EndEffectorSO3Objective final : public Objective {
public:
  EndEffectorSO3Objective(const Eigen::Matrix3d&);
  EndEffectorSO3Objective(double weight, const Eigen::Matrix3d&);

private:
  HebiStatusCode addObjective(HebiIKPtr ik) const override;
  double _weight;
  const double _matrix[9];
};

class EndEffectorTipAxisObjective final : public Objective {
public:
  EndEffectorTipAxisObjective(const Eigen::Vector3d&);
  EndEffectorTipAxisObjective(double weight, const Eigen::Vector3d&);

private:
  HebiStatusCode addObjective(HebiIKPtr ik) const override;
  double _weight, _x, _y, _z;
};

class JointLimitConstraint final : public Objective {
public:
  JointLimitConstraint(const Eigen::VectorXd& min_positions, const Eigen::VectorXd& max_positions);
  JointLimitConstraint(double weight, const Eigen::VectorXd& min_positions, const Eigen::VectorXd& max_positions);

private:
  HebiStatusCode addObjective(HebiIKPtr ik) const override;
  double _weight;
  Eigen::VectorXd _min_positions;
  Eigen::VectorXd _max_positions;

public:
  // Allow Eigen member variables:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
};

template<size_t T>
inline void custom_objective_callback_wrapper(void* user_data, size_t num_positions, const double* positions,
                                              double* errors);

template<size_t N>
class CustomObjective final : public Objective {
public:
  // This function is called with the following parameters:
  // const std::vector<double>& positions
  // std::array<double, N>& errors (NOTE: FILL THIS IN VIA THE CALLBACK FUNCTION)
  using ObjectiveCallback = std::function<void(const std::vector<double>&, std::array<double, N>&)>;

  CustomObjective(ObjectiveCallback error_function) : _weight(1.0f), _callback(error_function) {}
  CustomObjective(double weight, ObjectiveCallback error_function) : _weight(weight), _callback(error_function) {}

  // Note -- internal function to be called only from
  // custom_objective_callback_wrapper.
  void callCallback(void*, size_t num_positions, const double* positions, double* errors) const {
    // Note -- the data here is copied to/from the C-style arrays.  This isn't
    // optimally efficient, but allows us to use type-checked C++ vectors and
    // arrays.  For performance critical applications, this could be modified
    // to support user callbacks using the raw C-style arrays directly.

    // Process data into C++ structures.
    std::vector<double> positions_array(num_positions);
    for (size_t i = 0; i < num_positions; ++i)
      positions_array[i] = positions[i];
    // Note -- std::array is not guaranteed to be layout-compatible with a
    // C-style array, even for POD, so we must copy here it we want the type
    // safety of std::array.
    std::array<double, N> errors_array;
    for (size_t i = 0; i < N; ++i)
      errors_array[i] = errors[i];

    _callback(positions_array, errors_array);

    for (size_t i = 0; i < N; ++i)
      errors[i] = errors_array[i];
  }

private:
  HebiStatusCode addObjective(HebiIKPtr ik) const override {
    return hebiIKAddObjectiveCustom(ik, _weight, N, &custom_objective_callback_wrapper<N>,
                                    const_cast<CustomObjective*>(this));
  }
  double _weight;
  ObjectiveCallback _callback;
};

template<size_t T>
inline void custom_objective_callback_wrapper(void* user_data, size_t num_positions, const double* positions,
                                              double* errors) {
  reinterpret_cast<CustomObjective<T>*>(user_data)->callCallback(user_data, num_positions, positions, errors);
}

class ActuatorMetadata;
class BracketMetadata;
class JointMetadata;
class LinkMetadata;
class RigidBodyMetadata;

enum class ElementType {
  Other = HebiRobotModelElementTypeOther,
  Actuator = HebiRobotModelElementTypeActuator,
  Bracket = HebiRobotModelElementTypeBracket,
  Joint = HebiRobotModelElementTypeJoint,
  Link = HebiRobotModelElementTypeLink,
  RigidBody = HebiRobotModelElementTypeRigidBody,
  EndEffector = HebiRobotModelElementTypeEndEffector
};

enum class FrameType {
  CenterOfMass = HebiFrameTypeCenterOfMass,
  Output = HebiFrameTypeOutput,
  EndEffector = HebiFrameTypeEndEffector,
  Input = HebiFrameTypeInput,
  Mesh = HebiFrameTypeMesh
};

enum class JointType {
  RotationX = HebiJointTypeRotationX,
  RotationY = HebiJointTypeRotationY,
  RotationZ = HebiJointTypeRotationZ,
  TranslationX = HebiJointTypeTranslationX,
  TranslationY = HebiJointTypeTranslationY,
  TranslationZ = HebiJointTypeTranslationZ
};

enum class ActuatorType {
  X5_1 = HebiActuatorTypeX5_1,
  X5_4 = HebiActuatorTypeX5_4,
  X5_9 = HebiActuatorTypeX5_9,
  X8_3 = HebiActuatorTypeX8_3,
  X8_9 = HebiActuatorTypeX8_9,
  X8_16 = HebiActuatorTypeX8_16,
  R8_3 = HebiActuatorTypeR8_3,
  R8_9 = HebiActuatorTypeR8_9,
  R8_16 = HebiActuatorTypeR8_16
};

enum class LinkType {
  X5 = HebiLinkTypeX5,
  X8 = HebiLinkTypeR8
};

enum class LinkInputType {
  RightAngle = HebiLinkInputTypeRightAngle,
  Inline = HebiLinkInputTypeInline
};

enum class LinkOutputType {
  RightAngle = HebiLinkOutputTypeRightAngle,
  Inline = HebiLinkOutputTypeInline
};

enum class BracketType {
  X5LightLeft = HebiBracketTypeX5LightLeft,
  X5LightRight = HebiBracketTypeX5LightRight,
  X5HeavyLeftInside = HebiBracketTypeX5HeavyLeftInside,
  X5HeavyLeftOutside = HebiBracketTypeX5HeavyLeftOutside,
  X5HeavyRightInside = HebiBracketTypeX5HeavyRightInside,
  X5HeavyRightOutside = HebiBracketTypeX5HeavyRightOutside,
  R8LightLeft = HebiBracketTypeR8LightLeft,
  R8LightRight = HebiBracketTypeR8LightRight,
  R8HeavyLeftInside = HebiBracketTypeR8HeavyLeftInside,
  R8HeavyLeftOutside = HebiBracketTypeR8HeavyLeftOutside,
  R8HeavyRightInside = HebiBracketTypeR8HeavyRightInside,
  R8HeavyRightOutside = HebiBracketTypeR8HeavyRightOutside
};

enum class EndEffectorType {
  Custom = HebiEndEffectorTypeCustom,
  X5Parallel = HebiEndEffectorTypeX5Parallel,
  R8Parallel = HebiEndEffectorTypeR8Parallel
};


class MetadataBase {
  friend class RobotModel;
public:
  MetadataBase() = default;
  MetadataBase(const MetadataBase&) = delete;
  MetadataBase& operator=(const MetadataBase&) = delete;
  MetadataBase(MetadataBase&&) = default;
  MetadataBase& operator=(MetadataBase&&) = default;

  ElementType elementType() const {
    return static_cast<ElementType>(metadata_.element_type_);
  }

  const ActuatorMetadata* asActuator() const {
    if (elementType() == ElementType::Actuator) {
      return reinterpret_cast<const ActuatorMetadata*>(this);
    }
    return nullptr;
  }

  const BracketMetadata* asBracket() const {
    if (elementType() == ElementType::Bracket) {
      return reinterpret_cast<const BracketMetadata*>(this);
    }
    return nullptr;
  }

  const JointMetadata* asJoint() const {
    if (elementType() == ElementType::Joint) {
      return reinterpret_cast<const JointMetadata*>(this);
    }
    return nullptr;
  }

  const LinkMetadata* asLink() const {
    if (elementType() == ElementType::Link) {
      return reinterpret_cast<const LinkMetadata*>(this);
    }
    return nullptr;
  }

  const RigidBodyMetadata* asRigidBody() const {
    if (elementType() == ElementType::RigidBody) {
      return reinterpret_cast<const RigidBodyMetadata*>(this);
    }
    return nullptr;
  }

protected:
  // Raw data can only be accessed by the sub classes
  const HebiRobotModelElementMetadata& metadata() const { return metadata_; }
private:
  HebiRobotModelElementMetadata metadata_;
};

class ActuatorMetadata : public MetadataBase {
public:
  ActuatorType actuatorType() const { return static_cast<ActuatorType>(metadata().actuator_type_); }
};

class BracketMetadata : public MetadataBase {
public:
  BracketType bracketType() const { return static_cast<BracketType>(metadata().bracket_type_); }
};

class JointMetadata : public MetadataBase {
public:
  JointType jointType() const { return static_cast<JointType>(metadata().joint_type_); }
};

class LinkMetadata : public MetadataBase {
public:
  LinkType linkType() const { return static_cast<LinkType>(metadata().link_type_); }

  float extension() const { return metadata().extension_; }

  float twist() const { return metadata().twist_; }
};

class RigidBodyMetadata : public MetadataBase {};

class RobotModel final {
  friend Objective;

private:
  HebiRobotModelPtr const internal_;

  template<typename T>
  HebiStatusCode addObjectives(HebiIKPtr ik, const T& objective) const {
    static_assert(std::is_convertible<T*, Objective*>::value,
                  "Must pass arguments of base type hebi::robot_model::Objective to the variable args of solveIK!");
    return static_cast<const Objective*>(&objective)->addObjective(ik);
  }

  template<typename T, typename... Args>
  HebiStatusCode addObjectives(HebiIKPtr ik, const T& objective, Args... args) const {
    static_assert(std::is_convertible<T*, Objective*>::value,
                  "Must pass arguments of base type hebi::robot_model::Objective to the variable args of solveIK!");
    auto res = static_cast<const Objective*>(&objective)->addObjective(ik);
    if (res != HebiStatusSuccess)
      return res;
    return addObjectives(ik, args...);
  }

  bool tryAdd(HebiRobotModelElementPtr element);

  RobotModel(HebiRobotModelPtr);

public:

  using ActuatorType [[deprecated ("Replace with hebi::robot_model::ActuatorType")]] = robot_model::ActuatorType;
  using BracketType [[deprecated ("Replace with hebi::robot_model::BracketType")]] = robot_model::BracketType;
  using LinkType [[deprecated ("Replace with hebi::robot_model::LinkType")]] = robot_model::LinkType;

  RobotModel();

  static std::unique_ptr<RobotModel> loadHRDF(const std::string& file);

  static std::unique_ptr<RobotModel> loadHRDFString(const std::string& string);

  std::unique_ptr<RobotModel> createSubtree(size_t element_index);

  ~RobotModel() noexcept;

  void setBaseFrame(const Eigen::Matrix4d& base_frame);

  Eigen::Matrix4d getBaseFrame() const;

  size_t getFrameCount(FrameType frame_type) const;

  [[deprecated ("Replaced by hebi::robot_model::RobotModel::getFrameCount(hebi::robot_model::FrameType)")]]
  size_t getFrameCount(HebiFrameType frame_type) const { return getFrameCount(static_cast<FrameType>(frame_type)); }

  size_t getDoFCount() const;

  std::string getMeshPath(size_t mesh_frame_index) const;

  bool addRigidBody(const Eigen::Matrix4d& com, const Eigen::VectorXd& inertia, double mass,
                    const Eigen::Matrix4d& output);

  bool addJoint(JointType joint_type);

  [[deprecated ("Replaced by hebi::robot_model::RobotModel::addJoint(hebi::robot_model::JointType)")]]
  bool addJoint(HebiJointType joint_type) { return addJoint(static_cast<JointType>(joint_type)); }

  bool addActuator(robot_model::ActuatorType actuator_type);

  bool addLink(robot_model::LinkType link_type, double extension, double twist,
               LinkInputType input_type = LinkInputType::RightAngle,
               LinkOutputType output_type = LinkOutputType::RightAngle);

  bool addBracket(robot_model::BracketType bracket_type);

  bool addEndEffector(EndEffectorType end_effector_type);
  bool addEndEffector(double x, double y, double z);
  bool addEndEffector(const Eigen::Matrix4d& com, const Eigen::VectorXd& inertia, double mass,
                      const Eigen::Matrix4d& output);
  void getForwardKinematics(FrameType, const Eigen::VectorXd& positions, Matrix4dVector& frames) const;

  [[deprecated ("Change first argument form HebiFrameType to hebi::robot_model::FrameType")]]
  void getForwardKinematics(HebiFrameType frame_type, const Eigen::VectorXd& positions, Matrix4dVector& frames) const {
    getForwardKinematics(static_cast<FrameType>(frame_type), positions, frames);
  }

  void getFK(FrameType, const Eigen::VectorXd& positions, Matrix4dVector& frames) const;

  [[deprecated ("Change first argument form HebiFrameType to hebi::robot_model::FrameType")]]
  void getFK(HebiFrameType frame_type, const Eigen::VectorXd& positions, Matrix4dVector& frames) const {
    getFK(static_cast<FrameType>(frame_type), positions, frames);
  }

  void getEndEffector(const Eigen::VectorXd& positions, Eigen::Matrix4d& transform) const;

  template<typename... Args>
  IKResult solveInverseKinematics(const Eigen::VectorXd& initial_positions, Eigen::VectorXd& result,
                                  Args... args) const {
    return solveIK(initial_positions, result, args...);
  }

  template<typename... Args>
  IKResult solveIK(const Eigen::VectorXd& initial_positions, Eigen::VectorXd& result, Args... objectives) const {
    // Create a HEBI C library IK object
    auto ik = hebiIKCreate();

    // (Try) to add objectives to the IK object
    IKResult res;
    res.result = addObjectives(ik, objectives...);
    if (res.result != HebiStatusSuccess) {
      hebiIKRelease(ik);
      return res;
    }

    // Transfer/initialize from Eigen to C arrays
    auto positions_array = new double[initial_positions.size()];
    {
      Map<Eigen::VectorXd> tmp(positions_array, initial_positions.size());
      tmp = initial_positions;
    }
    auto result_array = new double[initial_positions.size()];

    // Call into C library to solve
    res.result = hebiIKSolve(ik, internal_, positions_array, result_array, nullptr);

    // Transfer/cleanup from C arrays to Eigen
    delete[] positions_array;
    {
      Map<Eigen::VectorXd> tmp(result_array, initial_positions.size());
      result = tmp;
    }
    delete[] result_array;

    hebiIKRelease(ik);

    return res;
  }

  void getJacobians(FrameType, const Eigen::VectorXd& positions, MatrixXdVector& jacobians) const;
  
  [[deprecated]]
  void getJacobians(HebiFrameType frame_type, const Eigen::VectorXd& positions, MatrixXdVector& jacobians) const {
    getJacobians(static_cast<FrameType>(frame_type), positions, jacobians);
  }

  void getJ(FrameType, const Eigen::VectorXd& positions, MatrixXdVector& jacobians) const;

  [[deprecated ("Change first argument form HebiFrameType to hebi::robot_model::FrameType")]]
  void getJ(HebiFrameType frame_type, const Eigen::VectorXd& positions, MatrixXdVector& jacobians) const {
    getJ(static_cast<FrameType>(frame_type), positions, jacobians);
  }

  void getJacobianEndEffector(const Eigen::VectorXd& positions, Eigen::MatrixXd& jacobian) const;
  void getJEndEffector(const Eigen::VectorXd& positions, Eigen::MatrixXd& jacobian) const;

  void getMasses(Eigen::VectorXd& masses) const;

  void getMetadata(std::vector<MetadataBase>& metadata) const;
private:
  HEBI_DISABLE_COPY_MOVE(RobotModel)
};

} // namespace robot_model
} // namespace hebi