Added: geodesic IK

Needs further testing

heimdall/geometry/distances.py

@@ -3,7 +3,9 @@
 from heimdall.datatypes.pose import SO3, Pose
 
 
-def rotational_geodesic_distance(source: SO3, target: SO3) -> torch.Tensor:
+def rotational_geodesic_distance(
+    source: SO3, target: SO3, eps: float = 1e-6
+) -> torch.Tensor:
     if source.matrix.ndim != target.matrix.ndim:
         raise ValueError(
             "Source and target SO3 matrices must have the same number of dimensions."
@@ -12,8 +14,7 @@ def rotational_geodesic_distance(source: SO3, target: SO3) -> torch.Tensor:
     relative_rotation = source.matrix @ target.matrix.transpose(-2, -1)
 
     trace = torch.sum(torch.diagonal(relative_rotation, dim1=-2, dim2=-1), dim=-1)
-    cos_theta = (trace - 1.0) * 0.5
-    cos_theta = torch.clip(cos_theta, -1.0, 1.0)
+    cos_theta = torch.clip((trace - 1.0) * 0.5, -1.0 + eps, 1.0 - eps)
     return torch.acos(cos_theta)
 
 

heimdall/geometry/transformations.py

@@ -8,17 +8,18 @@ def relative_pose(source_pose: Pose, target_pose: Pose) -> Pose:
 
 
 def kabsch_transform(source: torch.Tensor, target: torch.Tensor) -> Pose:
-    if source.ndim != target.ndim:
-        raise ValueError(
-            "Source and target must have same number of dimensions. Preferably N x 3 or B x N x 3"
-        )
+    if source.shape[-1] != 3 or target.shape[-1] != 3 or source.ndim != target.ndim:
+        raise ValueError("Source and target must have shape (N, 3) or (B, N, 3)")
 
     if source.ndim == 2 and target.ndim == 2:
-        source = source[None, ...]
-        target = target[None, ...]
+        batched_source = source[None, ...]
+        batched_target = target[None, ...]
+    else:
+        batched_source = source
+        batched_target = target
 
-    centered_source = source - source.mean(dim=-2)
-    centered_target = target - target.mean(dim=-2)
+    centered_source = batched_source - batched_source.mean(dim=-2, keepdim=True)
+    centered_target = batched_target - batched_target.mean(dim=-2, keepdim=True)
 
     covariance_matrix = centered_source.transpose(-2, -1) @ centered_target
     U, _, VT = torch.linalg.svd(covariance_matrix)
@@ -29,6 +30,9 @@ def kabsch_transform(source: torch.Tensor, target: torch.Tensor) -> Pose:
     V[:, -1] *= sign[..., None]
 
     rotation_matrix = V @ UT
-    translation = target.mean(dim=-2) - source.mean(dim=-2)
+    translation = batched_target.mean(dim=-2) - batched_source.mean(dim=-2)
 
-    return Pose(translation=translation.squeeze(-1), SO3=SO3(rotation_matrix))
+    if source.ndim == 2 and target.ndim == 2:
+        return Pose(translation=translation[0], SO3=SO3(rotation_matrix[0]))
+
+    return Pose(translation=translation, SO3=SO3(rotation_matrix))

heimdall/robotics/inverse_kinematics.py

@@ -3,6 +3,9 @@
 import numpy as np
 import torch
 
+from heimdall.datatypes.pose import Pose
+from heimdall.geometry.distances import pose_geodesic_distance
+from heimdall.geometry.transformations import relative_pose
 from heimdall.logger import OverWatch
 from heimdall.utils import convert_numpy_to_tensor
 
@@ -54,7 +57,7 @@ def compute_stable_inverse_jacobian(
 @convert_numpy_to_tensor
 def compute_planar_trajectory(
     intial_joint_position: np.ndarray | torch.Tensor,
-    target_tcp_position: np.ndarray | torch.Tensor,
+    target_tcp_position: Pose,
     jacobian_function: callable,
     forward_kinematics_function: callable,
     interpolation: str = "linear",
@@ -108,3 +111,44 @@ def compute_planar_trajectory(
             return [joint.cpu().numpy() for joint in joint_trajectory]
 
     return joint_trajectory
+
+
+@convert_numpy_to_tensor
+def compute_geodesic_trajectory(
+    intial_joint_position: np.ndarray | torch.Tensor,
+    target_tcp_pose: Pose,
+    jacobian_function: callable,
+    forward_kinematics_function: callable,
+    max_iterations: int = 100,
+    logging: bool = False,
+    **kwargs,
+) -> List[np.ndarray] | List[torch.Tensor]:
+    q_init = intial_joint_position
+    p_target = target_tcp_pose
+
+    joint_trajectory: List[np.ndarray] | List[torch.Tensor] = [q_init]
+
+    for _ in range(1, max_iterations + 1):
+        q_current = joint_trajectory[-1]
+        jacobian = jacobian_function(q_current)
+        p_current = forward_kinematics_function(q_current)
+
+        jacobian_inverse = compute_stable_inverse_jacobian(jacobian)
+        q_next = q_current + jacobian_inverse @ relative_pose(p_current, p_target)
+
+        joint_trajectory.append(q_next)
+
+        pose_distance = pose_geodesic_distance(p_target, p_current)
+        if torch.allclose(pose_distance, torch.zeros_like(pose_distance), atol=1e-4):
+            break
+
+    if logging:
+        p_current = forward_kinematics_function(joint_trajectory[-1])
+        error = pose_geodesic_distance(p_target, p_current)
+        logger.info(f"Geodesic trajectory generation completed with error: {error:.4f}")
+
+    if kwargs.get("numpy_found"):
+        with torch.no_grad():
+            return [joint.cpu().numpy() for joint in joint_trajectory]
+
+    return joint_trajectory