de-duplicate axis_angle conversion functions

Also remove dependency on transformations, since we only use one
function from it, and it had compatibility issues. I couldn't get it to
install on my laptop.

Author: Peter

pyproject.toml

@@ -45,7 +45,6 @@ dependencies = [
     'numpy',
     'pyyaml',
     'torch',
-    'transformations',
 ]
 
 [project.optional-dependencies]

src/pytorch_kinematics/chain.py

@@ -1,6 +1,6 @@
 from functools import lru_cache
-from pytorch_kinematics.transforms.rotation_conversions import tensor_axis_and_angle_to_matrix
-from pytorch_kinematics.transforms.rotation_conversions import tensor_axis_and_d_to_pris_matrix
+from pytorch_kinematics.transforms.rotation_conversions import axis_and_angle_to_matrix
+from pytorch_kinematics.transforms.rotation_conversions import axis_and_d_to_pris_matrix
 from typing import Optional, Sequence
 
 import numpy as np
@@ -294,8 +294,8 @@ def forward_kinematics(self, th, frame_indices: Optional = None):
         # compute all joint transforms at once first
         # in order to handle multiple joint types without branching, we create all possible transforms
         # for all joint types and then select the appropriate one for each joint.
-        rev_jnt_transform = tensor_axis_and_angle_to_matrix(axes_expanded, th)
-        pris_jnt_transform = tensor_axis_and_d_to_pris_matrix(axes_expanded, th)
+        rev_jnt_transform = axis_and_angle_to_matrix(axes_expanded, th)
+        pris_jnt_transform = axis_and_d_to_pris_matrix(axes_expanded, th)
 
         frame_transforms = {}
         b = th.shape[0]
@@ -456,7 +456,7 @@ def jacobian(self, th, locations=None):
 
     def forward_kinematics(self, th, end_only: bool = True):
         """ Like the base class, except `th` only needs to contain the joints in the SerialChain, not all joints. """
-        frame_indices, th = self.convert_serial_inputs_to_chain_inputs(end_only, th)
+        frame_indices, th = self.convert_serial_inputs_to_chain_inputs(th, end_only)
 
         mat = super().forward_kinematics(th, frame_indices)
 
@@ -465,7 +465,7 @@ def forward_kinematics(self, th, end_only: bool = True):
         else:
             return mat
 
-    def convert_serial_inputs_to_chain_inputs(self, end_only, th):
+    def convert_serial_inputs_to_chain_inputs(self, th, end_only: bool):
         if end_only:
             frame_indices = self.get_frame_indices(self._serial_frames[-1].name)
         else:

src/pytorch_kinematics/transforms/__init__.py

@@ -1,7 +1,7 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
+from pytorch_kinematics.transforms.perturbation import sample_perturbations
 from .rotation_conversions import (
-    axis_and_angle_to_matrix,
     axis_angle_to_quaternion,
     euler_angles_to_matrix,
     matrix_to_axis_angle,
@@ -13,13 +13,14 @@
     quaternion_multiply,
     quaternion_raw_multiply,
     quaternion_to_matrix,
+    quaternion_from_euler,
     random_quaternions,
     random_rotation,
     random_rotations,
     rotation_6d_to_matrix,
     standardize_quaternion,
-    tensor_axis_and_angle_to_matrix,
-    tensor_axis_and_d_to_pris_matrix,
+    axis_and_angle_to_matrix,
+    axis_and_d_to_pris_matrix,
     wxyz_to_xyzw,
     xyzw_to_wxyz,
 )
@@ -30,6 +31,5 @@
     so3_rotation_angle,
 )
 from .transform3d import Rotate, RotateAxisAngle, Scale, Transform3d, Translate
-from pytorch_kinematics.transforms.perturbation import sample_perturbations
 
 __all__ = [k for k in globals().keys() if not k.startswith("_")]

src/pytorch_kinematics/transforms/rotation_conversions.py

@@ -1,9 +1,11 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
 import functools
+import math
 from typing import Optional
 from warnings import warn
 
+import numpy
 import torch
 import torch.nn.functional as F
 
@@ -450,7 +452,7 @@ def quaternion_apply(quaternion, point):
     return out[..., 1:]
 
 
-def tensor_axis_and_d_to_pris_matrix(axis, d):
+def axis_and_d_to_pris_matrix(axis, d):
     """
     Creates a 4x4 matrix that represents a translation along an axis of a distance d
     Works with any number of batch dimensions.
@@ -470,7 +472,7 @@ def tensor_axis_and_d_to_pris_matrix(axis, d):
     return mat44
 
 
-def tensor_axis_and_angle_to_matrix(axis, theta):
+def axis_and_angle_to_matrix(axis, theta):
     """
     Creates a 4x4 matrix that represents a rotation around an axis by an angle theta.
     Works with any number of batch dimensions.
@@ -505,27 +507,6 @@ def tensor_axis_and_angle_to_matrix(axis, theta):
     return mat44
 
 
-def axis_and_angle_to_matrix(axis, theta):
-    # based on https://ai.stackexchange.com/questions/14041/, and checked against wikipedia
-    c = torch.cos(theta)  # NOTE: cos is not that precise for float32, you may want to use float64
-    one_minus_c = 1 - c
-    s = torch.sin(theta)
-    kx, ky, kz = torch.unbind(axis, -1)
-    r00 = c + kx * kx * one_minus_c
-    r01 = kx * ky * one_minus_c - kz * s
-    r02 = kx * kz * one_minus_c + ky * s
-    r10 = ky * kx * one_minus_c + kz * s
-    r11 = c + ky * ky * one_minus_c
-    r12 = ky * kz * one_minus_c - kx * s
-    r20 = kz * kx * one_minus_c - ky * s
-    r21 = kz * ky * one_minus_c + kx * s
-    r22 = c + kz * kz * one_minus_c
-    rot = torch.stack([torch.cat([r00, r01, r02], -1),
-                       torch.cat([r10, r11, r12], -1),
-                       torch.cat([r20, r21, r22], -1)], -2)
-    return rot
-
-
 def axis_angle_to_matrix(axis_angle):
     """
     Convert rotations given as axis/angle to rotation matrices.
@@ -541,7 +522,7 @@ def axis_angle_to_matrix(axis_angle):
     Returns:
         Rotation matrices as tensor of shape (..., 3, 3).
     """
-    warn('This is deprecated because it is slow. Use axis_and_angle_to_matrix or zpk_cpp.axis_and_angle_to_matrix',
+    warn('This is deprecated because it is slow. Use axis_and_angle_to_matrix',
          DeprecationWarning, stacklevel=2)
     return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
 
@@ -682,3 +663,96 @@ def pos_rot_to_matrix(pos, rot):
     m[..., :3, 3] = pos
     m[..., :3, :3] = rot
     return m
+
+
+# axis sequences for Euler angles
+_NEXT_AXIS = [1, 2, 0, 1]
+
+# map axes strings to/from tuples of inner axis, parity, repetition, frame
+_AXES2TUPLE = {
+    'sxyz': (0, 0, 0, 0),
+    'sxyx': (0, 0, 1, 0),
+    'sxzy': (0, 1, 0, 0),
+    'sxzx': (0, 1, 1, 0),
+    'syzx': (1, 0, 0, 0),
+    'syzy': (1, 0, 1, 0),
+    'syxz': (1, 1, 0, 0),
+    'syxy': (1, 1, 1, 0),
+    'szxy': (2, 0, 0, 0),
+    'szxz': (2, 0, 1, 0),
+    'szyx': (2, 1, 0, 0),
+    'szyz': (2, 1, 1, 0),
+    'rzyx': (0, 0, 0, 1),
+    'rxyx': (0, 0, 1, 1),
+    'ryzx': (0, 1, 0, 1),
+    'rxzx': (0, 1, 1, 1),
+    'rxzy': (1, 0, 0, 1),
+    'ryzy': (1, 0, 1, 1),
+    'rzxy': (1, 1, 0, 1),
+    'ryxy': (1, 1, 1, 1),
+    'ryxz': (2, 0, 0, 1),
+    'rzxz': (2, 0, 1, 1),
+    'rxyz': (2, 1, 0, 1),
+    'rzyz': (2, 1, 1, 1),
+}
+
+_TUPLE2AXES = {v: k for k, v in _AXES2TUPLE.items()}
+
+
+def quaternion_from_euler(ai, aj, ak, axes='sxyz'):
+    """
+    Return quaternion from Euler angles and axis sequence.
+    Taken from https://github.com/cgohlke/transformations/blob/master/transformations/transformations.py#L1238
+
+    ai, aj, ak : Euler's roll, pitch and yaw angles
+    axes : One of 24 axis sequences as string or encoded tuple
+
+    >>> q = quaternion_from_euler(1, 2, 3, 'ryxz')
+    >>> numpy.allclose(q, [0.435953, 0.310622, -0.718287, 0.444435])
+    True
+
+    """
+    try:
+        firstaxis, parity, repetition, frame = _AXES2TUPLE[axes.lower()]
+    except (AttributeError, KeyError):
+        _TUPLE2AXES[axes]  # noqa: validation
+        firstaxis, parity, repetition, frame = axes
+
+    i = firstaxis + 1
+    j = _NEXT_AXIS[i + parity - 1] + 1
+    k = _NEXT_AXIS[i - parity] + 1
+
+    if frame:
+        ai, ak = ak, ai
+    if parity:
+        aj = -aj
+
+    ai /= 2.0
+    aj /= 2.0
+    ak /= 2.0
+    ci = math.cos(ai)
+    si = math.sin(ai)
+    cj = math.cos(aj)
+    sj = math.sin(aj)
+    ck = math.cos(ak)
+    sk = math.sin(ak)
+    cc = ci * ck
+    cs = ci * sk
+    sc = si * ck
+    ss = si * sk
+
+    q = numpy.empty((4,))
+    if repetition:
+        q[0] = cj * (cc - ss)
+        q[i] = cj * (cs + sc)
+        q[j] = sj * (cc + ss)
+        q[k] = sj * (cs - sc)
+    else:
+        q[0] = cj * cc + sj * ss
+        q[i] = cj * sc - sj * cs
+        q[j] = cj * ss + sj * cc
+        q[k] = cj * cs - sj * sc
+    if parity:
+        q[j] *= -1.0
+
+    return q

src/pytorch_kinematics/urdf.py

@@ -3,8 +3,6 @@
 from . import chain
 import torch
 import pytorch_kinematics.transforms as tf
-# has better RPY to quaternion transformation
-import transformations as tf2
 
 JOINT_TYPE_MAP = {'revolute':   'revolute',
                   'continuous': 'revolute',
@@ -16,7 +14,7 @@ def _convert_transform(origin):
     if origin is None:
         return tf.Transform3d()
     else:
-        return tf.Transform3d(rot=torch.tensor(tf2.quaternion_from_euler(*origin.rpy, "sxyz"), dtype=torch.float32),
+        return tf.Transform3d(rot=torch.tensor(tf.quaternion_from_euler(*origin.rpy, "sxyz"), dtype=torch.float32),
                               pos=origin.xyz)