feat: global ik with q7 sampling (#3)

- added global ik with q7 sampling
- moved ik sampling solution selection logic to python
- added joint weights

Author: juelg

src/frankik/__init__.py

@@ -1,5 +1,3 @@
-import typing
-
 import numpy as np
 
 from frankik._core import (
@@ -48,9 +46,7 @@ def __init__(self, robot_type: str = RobotType.FR3):
         self.q_home = kQDefault
 
     @staticmethod
-    def pose_inverse(
-        T: np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]]
-    ) -> np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]]:
+    def pose_inverse(T: np.ndarray) -> np.ndarray:
         """Compute the inverse of a homogeneous transformation matrix.
         Args:
             T (np.ndarray): A 4x4 homogeneous transformation matrix.
@@ -66,9 +62,9 @@ def pose_inverse(
 
     def forward(
         self,
-        q0: np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]],
-        tcp_offset: np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]] | None = None,
-    ) -> np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]]:
+        q0: np.ndarray,
+        tcp_offset: np.ndarray | None = None,
+    ) -> np.ndarray:
         """Compute the forward kinematics for the given joint configuration.
         Args:
             q0 (np.ndarray): A 7-element array representing joint angles.
@@ -83,11 +79,15 @@ def forward(
 
     def inverse(
         self,
-        pose: np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]],
-        q0: np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]] | None = None,
-        tcp_offset: np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]] | None = None,
+        pose: np.ndarray,
+        q0: np.ndarray | None = None,
+        tcp_offset: np.ndarray | None = None,
         q7: float | None = None,
-    ) -> np.ndarray[tuple[typing.Literal[7]], np.dtype[np.float64]] | None:
+        global_solution: bool = False,
+        joint_weight: np.ndarray | None = None,
+        q7_sample_interval=40,
+        q7_sample_size=60,
+    ) -> np.ndarray | None:
         """Compute the inverse kinematics for the given end-effector pose.
 
         Args:
@@ -96,15 +96,33 @@ def inverse(
             tcp_offset (np.ndarray, optional): A 4x4 homogeneous transformation matrix representing
                 the tool center point offset. Defaults to None.
             q7 (float, optional): The angle of the seventh joint, used for FR3 robot IK. If None then it will be sampled. Defaults to None.
+            global_solution (bool, optional): Whether to consider global ik solutions. Defaults to False.
+            joint_weight (np.ndarray, optional): Weights for calculating the distance between the solution and q0. Defaults to None.
+            q7_sample_interval (int, optional): The interval for sampling q7. Defaults to 40.
+            q7_sample_size (int, optional): The number of samples for q7. Defaults to 60.
 
         Returns:
             np.ndarray | None: A 7-element array representing the joint angles if a solution is found; otherwise, None.
         """
+        if joint_weight is None:
+            joint_weight = np.ones(7)
+        if q0 is None:
+            q0 = self.q_home
+
         new_pose = pose @ self.pose_inverse(tcp_offset) if tcp_offset is not None else pose
         if q7 is not None:
             q = ik(new_pose, q0, q7, is_fr3=(self.robot_type == RobotType.FR3))  # type: ignore
         else:
-            q = ik_sample_q7(new_pose, q0, is_fr3=(self.robot_type == RobotType.FR3))  # type: ignore
+            qs = ik_sample_q7(
+                new_pose,
+                q0,
+                is_fr3=(self.robot_type == RobotType.FR3),
+                sample_size=q7_sample_size,
+                sample_interval=q7_sample_interval,
+                full_ik=global_solution,
+            )  # type: ignore
+            q_diffs = np.sum((np.array(qs) - q0) * joint_weight, axis=1) ** 2
+            q = qs[np.argmin(q_diffs)]
 
         if np.isnan(q).any():
             return None

src/frankik/_core.pyi

@@ -73,23 +73,27 @@ def ik_full(
 
 def ik_sample_q7(
     O_T_EE: numpy.ndarray[tuple[typing.Literal[4], typing.Literal[4]], numpy.dtype[numpy.float64]],
-    q_actual: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]],
+    q_actual: numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]] | None = None,
     is_fr3: bool = False,
-    sample_size: int = 30,
-    sample_interval: float = 10,
-) -> numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]]:
+    sample_size: int = 60,
+    sample_interval: float = 40,
+    full_ik: bool = False,
+) -> list[numpy.ndarray[tuple[typing.Literal[7]], numpy.dtype[numpy.float64]]]:
     """
     Compute one inverse kinematics solution for Franka with sampling of joint q7.
 
     Args:
-        O_T_EE (Eigen::Matrix<double, 4, 4>): Desired end-effector pose.
-        q_actual (Vector7d): Current joint angles.
-        is_fr3 (bool, optional): Whether to use FR3 joint limits. Defaults to false.
+        O_T_EE (np.array): Desired end-effector pose. Shape (4, 4).
+        q_actual (np.array, optional): Current joint angles. Shape (7,).
+
+        is_fr3 (bool, optional): Whether to use FR3 joint limits. Defaults to False.
         sample_size (int, optional): How many sample to try for q7. Defaults to 20.
         sample_interval (int, optional): Sample interval for q7 in degree. Defaults to 90.
+        full_ik (bool, optional): Whether to use full IK. Defaults to False.
+    degree. Defaults to False.
 
     Returns:
-        Vector7d: One IK solution. NaN if no solution.
+        list[np.array]: One IK solution. Empty if no solution was found. Array shape (7,).
     """
 
 __version__: str = "0.1"

src/pybind/bindings.cpp

@@ -58,21 +58,29 @@ PYBIND11_MODULE(_core, m) {
         "Returns:\n"
         "    Vector7d: One IK solution. NaN if no solution.");
   m.def("ik_sample_q7", &frankik::ik_sample_q7, py::arg("O_T_EE"),
-        py::arg("q_actual"), py::arg("is_fr3") = false,
-        py::arg("sample_size") = 30, py::arg("sample_interval") = 10,
+        py::arg("q_actual") = std::nullopt, py::arg("is_fr3") = false,
+        py::arg("sample_size") = 60, py::arg("sample_interval") = 40,
+        py::arg("full_ik") = false,
         "Compute one inverse kinematics solution for Franka with sampling of "
         "joint q7.\n\n"
         "Args:\n"
-        "    O_T_EE (Eigen::Matrix<double, 4, 4>): Desired end-effector pose.\n"
-        "    q_actual (Vector7d): Current joint angles.\n"
+        "    O_T_EE (np.array): Desired end-effector pose. Shape "
+        "(4, 4).\n"
+        "    q_actual (np.array, optional): Current joint angles. Shape "
+        "(7,).\n"
+        "\n"
         "    is_fr3 (bool, optional): Whether to use FR3 joint limits. "
-        "Defaults to false.\n"
+        "Defaults to False.\n"
         "    sample_size (int, optional): How many sample to try for q7. "
         "Defaults to 20.\n"
         "    sample_interval (int, optional): Sample interval for q7 in "
-        "degree. Defaults to 90.\n\n"
+        "degree. Defaults to 90.\n"
+        "    full_ik (bool, optional): Whether to use full IK. Defaults to "
+        "False.\n"
+        "degree. Defaults to False.\n\n"
         "Returns:\n"
-        "    Vector7d: One IK solution. NaN if no solution.");
+        "    list[np.array]: One IK solution. Empty if no solution was found. "
+        "Array shape (7,).");
   m.def("fk", &frankik::fk, py::arg("q"),
         "Compute forward kinematics for Franka robot.\n\n"
         "Args:\n"

src/pybind/frankik.h

@@ -396,23 +396,23 @@ Vector7d ik(Eigen::Matrix<double, 4, 4> O_T_EE,
   return q;
 }
 
-Vector7d ik_sample_q7(Eigen::Matrix<double, 4, 4> O_T_EE, Vector7d q_actual,
-                      bool is_fr3 = false, const size_t sample_size = 30,
-                      const double sample_interval = 10) {
+std::vector<Vector7d> ik_sample_q7(
+    const Eigen::Matrix<double, 4, 4> O_T_EE,
+    const std::optional<Vector7d> q_actual = std::nullopt, bool is_fr3 = false,
+    size_t sample_size = 30, double sample_interval = 10,
+    bool full_ik = false) {
+  Vector7d q_actual_array = q_actual.value_or(kQDefault);
   const std::array<double, 7> q_min = is_fr3 ? q_min_fr3 : q_min_panda;
   const std::array<double, 7> q_max = is_fr3 ? q_max_fr3 : q_max_panda;
 
-  auto q7 = q_actual[6];
+  auto q7 = q_actual_array[6];
 
   // sample interval in degree
   auto low7 = q7 - (sample_interval / 2.0) / 180.0 * M_PI;
   auto high7 = q7 + (sample_interval / 2.0) / 180.0 * M_PI;
 
   auto sample_step_size = (high7 - low7) / sample_size;
-
-  Vector7d ik_solution =
-      Vector7d::Constant(std::numeric_limits<double>::quiet_NaN());
-  double loss = std::numeric_limits<double>::max();
+  std::vector<Vector7d> ik_solutions;
 
   // shift if interval goes over the edges
   if (low7 < q_min[6]) {
@@ -426,18 +426,24 @@ Vector7d ik_sample_q7(Eigen::Matrix<double, 4, 4> O_T_EE, Vector7d q_actual,
   }
   for (size_t i = 0; i <= sample_size; ++i) {
     // last step is to try the current q7 position
-    q7 = i < sample_size ? low7 + sample_step_size * i : q_actual[6];
-    auto tmp_ik_sol = ik(O_T_EE, q_actual, q7, is_fr3);
-    if (std::isnan(tmp_ik_sol[0])) {
-      continue;
-    }
-    auto tmp_loss = (tmp_ik_sol - q_actual).norm();
-    if (tmp_loss < loss) {
-      loss = tmp_loss;
-      ik_solution = tmp_ik_sol;
+    q7 = i < sample_size ? low7 + sample_step_size * i : q_actual_array[6];
+    if (!full_ik) {
+      auto tmp_ik_sol = ik(O_T_EE, q_actual_array, q7, is_fr3);
+      if (std::isnan(tmp_ik_sol[0])) {
+        continue;
+      }
+      ik_solutions.push_back(tmp_ik_sol);
+    } else {
+      auto tmp_ik_sols = ik_full(O_T_EE, q_actual_array, q7, is_fr3);
+      for (size_t j = 0; j <= 4; ++j) {
+        if (std::isnan(tmp_ik_sols.row(j)[0])) {
+          continue;
+        }
+        ik_solutions.push_back(tmp_ik_sols.row(j));
+      }
     }
   }
-  return ik_solution;
+  return ik_solutions;
 }
 
 Eigen::Matrix<double, 4, 4> fk(const Eigen::Matrix<double, 7, 1>& q) {

src/tests/test_example.py

@@ -8,6 +8,6 @@
 def test_example(robot_type):
     kinematics = FrankaKinematics(robot_type=robot_type)
     q_home = np.array([0.0, -np.pi / 4, 0.0, -3 * np.pi / 4, 0.0, np.pi / 2, np.pi / 4])
-    pose_home = kinematics.forward(q_home, tcp_offset=kinematics.FrankaHandTCPOffset)  # type: ignore
-    q = kinematics.inverse(pose_home, tcp_offset=kinematics.FrankaHandTCPOffset, q0=q_home)  # type: ignore
+    pose_home = kinematics.forward(q_home, tcp_offset=kinematics.FrankaHandTCPOffset)
+    q = kinematics.inverse(pose_home, tcp_offset=kinematics.FrankaHandTCPOffset, q0=q_home)
     assert np.allclose(q, q_home)  # type: ignore