add 6D Actuator, IK and FK Solvers (specify rotation and orientation of end-effector)

tinyik/core.py

@@ -5,7 +5,7 @@
 import autograd.numpy as np
 
 from .component import Link, Joint
-from .solver import FKSolver, IKSolver
+from .solver import FKSolver, IKSolver, FKSolver6D, IKSolver6D
 from .optimizer import ScipyOptimizer
 
 
@@ -17,24 +17,23 @@ def __init__(self, tokens, optimizer=None):
         components = []
         for t in tokens:
             if isinstance(t, Number):
-                components.append(Link([t, 0., 0.]))
+                components.append(Link([t, 0.0, 0.0]))
             elif isinstance(t, list) or isinstance(t, np.ndarray):
                 components.append(Link(t))
-            elif isinstance(t, str) and t in {'x', 'y', 'z'}:
+            elif isinstance(t, str) and t in {"x", "y", "z"}:
                 components.append(Joint(t))
             else:
                 raise ValueError(
-                    'the arguments need to be '
-                    'link length or joint axis: {}'.format(t)
+                    "the arguments need to be "
+                    "link length or joint axis: {}".format(t)
                 )
 
         self.fk = FKSolver(components)
         self.ik = IKSolver(
-            self.fk, ScipyOptimizer() if optimizer is None else optimizer)
-
-        self.angles = [0.] * len(
-            [c for c in components if isinstance(c, Joint)]
+            self.fk, ScipyOptimizer() if optimizer is None else optimizer
         )
+
+        self.angles = [0.0] * len([c for c in components if isinstance(c, Joint)])
         self.components = components
 
     @property
@@ -54,3 +53,35 @@ def ee(self):
     @ee.setter
     def ee(self, position):
         self.angles = self.ik.solve(self.angles, position)
+
+
+class Actuator6D(Actuator):
+    """
+    Drop-in replacement for tinyik.Actuator that supports 6D targets.
+    """
+
+    def __init__(self, tokens, optimizer=None):
+        super().__init__(tokens, optimizer)
+
+        self.fk = FKSolver6D(self.components)
+        self.ik = IKSolver6D(
+            self.fk, ScipyOptimizer() if optimizer is None else optimizer
+        )
+
+    @property
+    def ee(self):
+        """Returns current [x, y, z] position (for compatibility)."""
+        return self.fk.solve(self.angles)
+
+    @ee.setter
+    def ee(self, values):
+        """
+        Sets target.
+        If values has 3 elements -> Solves Position Only (legacy behavior)
+        If values has 6 elements -> Solves Position + Orientation [x, y, z, rx, ry, rz]
+        """
+        if len(values) == 3:
+            # Default to 0 rotation if only 3 are given.
+            values = list(values) + [0, 0, 0]
+
+        self.angles = self.ik.solve(self.angles, values)

tinyik/solver.py

@@ -13,9 +13,7 @@ class FKSolver(object):
 
     def __init__(self, components):
         """Generate a FK solver from link and joint instances."""
-        joint_indexes = [
-            i for i, c in enumerate(components) if isinstance(c, Joint)
-        ]
+        joint_indexes = [i for i, c in enumerate(components) if isinstance(c, Joint)]
 
         def matrices(angles):
             joints = dict(zip(joint_indexes, angles))
@@ -29,7 +27,7 @@ def solve(self, angles):
         return reduce(
             lambda a, m: np.dot(m, a),
             reversed(self._matrices(angles)),
-            np.array([0., 0., 0., 1.])
+            np.array([0.0, 0.0, 0.0, 1.0]),
         )[:3]
 
 
@@ -38,6 +36,7 @@ class IKSolver(object):
 
     def __init__(self, fk_solver, optimizer):
         """Generate an IK solver from a FK solver instance."""
+
         def distance_squared(angles, target):
             x = target - fk_solver.solve(angles)
             return np.sum(np.power(x, 2))
@@ -53,9 +52,7 @@ def solve(self, angles0, target):
 class CCDFKSolver(object):
 
     def __init__(self, components):
-        joint_indexes = [
-            i for i, c in enumerate(components) if isinstance(c, Joint)
-        ]
+        joint_indexes = [i for i, c in enumerate(components) if isinstance(c, Joint)]
 
         def matrices(angles):
             joints = dict(zip(joint_indexes, angles))
@@ -68,13 +65,13 @@ def matrices(angles):
 
     def solve(self, angles, p=None, index=None):
         if p is None:
-            p = [0., 0., 0., 1.]
+            p = [0.0, 0.0, 0.0, 1.0]
         if index is None:
             index = len(self.components) - 1
         return reduce(
             lambda a, m: np.dot(m, a),
-            reversed(self._matrices(angles)[:index + 1]),
-            np.array(p)
+            reversed(self._matrices(angles)[: index + 1]),
+            np.array(p),
         )[:3]
 
 
@@ -89,16 +86,16 @@ def solve(self, angles0, target):
         joint_indexes = list(reversed(self._fk_solver.joint_indexes))
         angles = angles0[:]
         pmap = {
-            'x': [1., 0., 0., 0.],
-            'y': [0., 1., 0., 0.],
-            'z': [0., 0., 1., 0.]}
+            "x": [1.0, 0.0, 0.0, 0.0],
+            "y": [0.0, 1.0, 0.0, 0.0],
+            "z": [0.0, 0.0, 1.0, 0.0],
+        }
         prev_dist = sys.float_info.max
         for _ in range(self.maxiter):
             for i, idx in enumerate(joint_indexes):
                 axis = self._fk_solver.solve(
-                    angles,
-                    p=pmap[self._fk_solver.components[idx].axis],
-                    index=idx)
+                    angles, p=pmap[self._fk_solver.components[idx].axis], index=idx
+                )
                 pj = self._fk_solver.solve(angles, index=idx)
                 ee = self._fk_solver.solve(angles)
                 eorp = self.p_on_rot_plane(ee, pj, axis)
@@ -107,7 +104,7 @@ def solve(self, angles0, target):
                 vt = (torp - pj) / np.linalg.norm(torp - pj)
                 a = np.arccos(np.dot(vt, ve))
                 sign = 1 if np.dot(axis, np.cross(ve, vt)) > 0 else -1
-                angles[-(i + 1)] += (a * sign)
+                angles[-(i + 1)] += a * sign
 
             ee = self._fk_solver.solve(angles)
             dist = np.linalg.norm(target - ee)
@@ -121,3 +118,65 @@ def solve(self, angles0, target):
     def p_on_rot_plane(self, p, joint_pos, joint_axis):
         ua = joint_axis / np.linalg.norm(joint_axis)
         return p - (np.dot(p - joint_pos, ua) * ua)
+
+
+class FKSolver6D(FKSolver):
+    """
+    Extends FKSolver to return the full 4x4 transformation matrix
+    instead of just the position vector.
+    """
+
+    def solve_matrix(self, angles):
+        # Multiply all link/joint matrices together to get the full end-effector pose
+        return reduce(
+            lambda a, m: np.dot(m, a), reversed(self._matrices(angles)), np.eye(4)
+        )
+
+
+class IKSolver6D(IKSolver):
+    """
+    Extends IKSolver to optimize for both Position and Orientation (6D).
+    Target format: [x, y, z, rx, ry, rz] (Euler ZYX convention in radians)
+    """
+
+    def __init__(self, fk_solver, optimizer):
+        # We assume the user works in mm, so we weight rotation heavily
+        # to ensure it isn't ignored by the optimizer.
+        # 1 radian error ~= 100mm error
+        ROTATION_WEIGHT = 100.0
+
+        def objective(angles, target):
+            # 1. Unpack Target (6D vector)
+            target_pos = target[:3]
+            rx, ry, rz = target[3:]
+
+            # 2. Construct Target Rotation Matrix (Euler ZYX)
+            cx, sx = np.cos(rx), np.sin(rx)
+            cy, sy = np.cos(ry), np.sin(ry)
+            cz, sz = np.cos(rz), np.sin(rz)
+
+            # Standard Rotation Matrices
+            Rx = np.array([[1.0, 0.0, 0.0], [0.0, cx, -sx], [0.0, sx, cx]])
+            Ry = np.array([[cy, 0.0, sy], [0.0, 1.0, 0.0], [-sy, 0.0, cy]])
+            Rz = np.array([[cz, -sz, 0.0], [sz, cz, 0.0], [0.0, 0.0, 1.0]])
+
+            # Combined Target Rotation (Rz * Ry * Rx)
+            R_target = np.dot(Rz, np.dot(Ry, Rx))
+
+            # 3. Compute Current Pose from Angles
+            M_current = fk_solver.solve_matrix(angles)
+            P_current = M_current[:3, 3]
+            R_current = M_current[:3, :3]
+
+            # 4. Calculate Errors
+            # Position Error (Euclidean Distance Squared)
+            pos_err = np.sum((P_current - target_pos) ** 2)
+
+            # Orientation Error (Frobenius Norm Squared of matrix difference)
+            # This is a robust way to measure "how different" two rotation matrices are
+            rot_err = np.sum((R_current - R_target) ** 2)
+
+            return pos_err + (rot_err * ROTATION_WEIGHT)
+
+        optimizer.prepare(objective)
+        self.optimizer = optimizer