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;

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

  HebiRobotModelElementType elementType() const {
    return metadata_.element_type_;
  }

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

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

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

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

  const RigidBodyMetadata* asRigidBody() const {
    if (elementType() == HebiRobotModelElementTypeRigidBody) {
      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:
  HebiActuatorType actuatorType() const { return metadata().actuator_type_; }
};

class BracketMetadata : public MetadataBase {
public:
  HebiBracketType bracketType() const { return metadata().bracket_type_; }
};

class JointMetadata : public MetadataBase {
public:
  HebiJointType jointType() const { return metadata().joint_type_; }
};

class LinkMetadata : public MetadataBase {
public:
  HebiLinkType linkType() const { return 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, bool combine);

  RobotModel(HebiRobotModelPtr);

public:
  enum class ActuatorType { X5_1, X5_4, X5_9, X8_3, X8_9, X8_16 };

  enum class LinkType { X5 };

  enum class BracketType {
    X5LightLeft,
    X5LightRight,
    X5HeavyLeftInside,
    X5HeavyLeftOutside,
    X5HeavyRightInside,
    X5HeavyRightOutside
  };

  RobotModel();

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

  ~RobotModel() noexcept;

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

  Eigen::Matrix4d getBaseFrame() const;

  size_t getFrameCount(HebiFrameType frame_type) const;

  size_t getDoFCount() const;

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

  bool addJoint(HebiJointType joint_type, bool combine);

  bool addActuator(ActuatorType actuator_type);

  bool addLink(LinkType link_type, double extension, double twist);

  bool addBracket(BracketType bracket_type);

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

  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(HebiFrameType, const Eigen::VectorXd& positions, MatrixXdVector& jacobians) const;
  void getJ(HebiFrameType, const Eigen::VectorXd& positions, MatrixXdVector& jacobians) const;

  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