change how we test for quaternion equality

Author: Peter

src/pytorch_kinematics/transforms/math.py

@@ -10,6 +10,16 @@
 import torch
 
 
+def quaternion_close(q1: torch.Tensor, q2: torch.Tensor, eps: float = 1e-4):
+    """
+    Returns true if two quaternions are close to each other. Assumes the quaternions are normalized.
+    Based on: https://math.stackexchange.com/a/90098/516340
+
+    """
+    dist = 1 - torch.square(torch.sum(q1*q2, dim=-1))
+    return torch.all(dist < eps)
+
+
 def acos_linear_extrapolation(
         x: torch.Tensor,
         bound: Union[float, Tuple[float, float]] = 1.0 - 1e-4,

src/pytorch_kinematics/transforms/rotation_conversions.py

@@ -1,11 +1,9 @@
 # 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
 
@@ -715,25 +713,22 @@ def pos_rot_to_matrix(pos, rot):
 _TUPLE2AXES = {v: k for k, v in _AXES2TUPLE.items()}
 
 
-def quaternion_from_euler(ai, aj, ak, axes='sxyz'):
+def quaternion_from_euler(rpy, 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
 
+    ai, aj, ak = torch.unbind(rpy, -1)
     i = firstaxis + 1
     j = _NEXT_AXIS[i + parity - 1] + 1
     k = _NEXT_AXIS[i - parity] + 1
@@ -746,29 +741,29 @@ def quaternion_from_euler(ai, aj, ak, axes='sxyz'):
     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)
+    ci = torch.cos(ai)
+    si = torch.sin(ai)
+    cj = torch.cos(aj)
+    sj = torch.sin(aj)
+    ck = torch.cos(ak)
+    sk = torch.sin(ak)
     cc = ci * ck
     cs = ci * sk
     sc = si * ck
     ss = si * sk
 
-    q = numpy.empty((4,))
+    q = torch.zeros([*rpy.shape[:-1], 4]).to(rpy)
     if repetition:
-        q[0] = cj * (cc - ss)
-        q[i] = cj * (cs + sc)
-        q[j] = sj * (cc + ss)
-        q[k] = sj * (cs - sc)
+        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
+        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
+        q[..., j] *= -1.0
 
     return q

src/pytorch_kinematics/urdf.py

@@ -14,8 +14,8 @@ def _convert_transform(origin):
     if origin is None:
         return tf.Transform3d()
     else:
-        return tf.Transform3d(rot=torch.tensor(tf.quaternion_from_euler(*origin.rpy, "sxyz"), dtype=torch.float32),
-                              pos=origin.xyz)
+        rpy = torch.tensor(origin.rpy, dtype=torch.float32)
+        return tf.Transform3d(rot=tf.quaternion_from_euler(rpy, "sxyz"), pos=origin.xyz)
 
 
 def _convert_visual(visual):

tests/test_kinematics.py

@@ -5,6 +5,7 @@
 import torch
 
 import pytorch_kinematics as pk
+from pytorch_kinematics.transforms.math import quaternion_close
 
 TEST_DIR = os.path.dirname(__file__)
 
@@ -16,11 +17,6 @@ def quat_pos_from_transform3d(tg):
     return pos, rot
 
 
-def quaternion_equality(a, b, rtol=1e-5):
-    # negative of a quaternion is the same rotation
-    return torch.allclose(a, b, rtol=rtol) or torch.allclose(a, -b, rtol=rtol)
-
-
 # test more complex robot and the MJCF parser
 def test_fk_mjcf():
     chain = pk.build_chain_from_mjcf(open(os.path.join(TEST_DIR, "ant.xml")).read())
@@ -33,11 +29,11 @@ def test_fk_mjcf():
     ret = chain.forward_kinematics(th)
     tg = ret['aux_1']
     pos, rot = quat_pos_from_transform3d(tg)
-    assert quaternion_equality(rot, torch.tensor([0.87758256, 0., 0., 0.47942554], dtype=torch.float64))
+    assert quaternion_close(rot, torch.tensor([0.87758256, 0., 0., 0.47942554], dtype=torch.float64))
     assert torch.allclose(pos, torch.tensor([0.2, 0.2, 0.75], dtype=torch.float64))
     tg = ret['front_left_foot']
     pos, rot = quat_pos_from_transform3d(tg)
-    assert quaternion_equality(rot, torch.tensor([0.77015115, -0.4600326, 0.13497724, 0.42073549], dtype=torch.float64))
+    assert quaternion_close(rot, torch.tensor([0.77015115, -0.4600326, 0.13497724, 0.42073549], dtype=torch.float64))
     assert torch.allclose(pos, torch.tensor([0.13976626, 0.47635466, 0.75], dtype=torch.float64))
     print(ret)
 
@@ -47,7 +43,7 @@ def test_fk_serial_mjcf():
     chain = chain.to(dtype=torch.float64)
     tg = chain.forward_kinematics([1.0, 1.0])
     pos, rot = quat_pos_from_transform3d(tg)
-    assert quaternion_equality(rot, torch.tensor([0.77015115, -0.4600326, 0.13497724, 0.42073549], dtype=torch.float64))
+    assert quaternion_close(rot, torch.tensor([0.77015115, -0.4600326, 0.13497724, 0.42073549], dtype=torch.float64))
     assert torch.allclose(pos, torch.tensor([0.13976626, 0.47635466, 0.75], dtype=torch.float64))
 
 
@@ -72,7 +68,7 @@ def test_fkik():
     tg = chain.forward_kinematics(th1)
     pos, rot = quat_pos_from_transform3d(tg)
     assert torch.allclose(pos, torch.tensor([[1.91081784, 0.41280851, 0.0000]]))
-    assert quaternion_equality(rot, torch.tensor([[0.95521418, 0.0000, 0.0000, 0.2959153]]))
+    assert quaternion_close(rot, torch.tensor([[0.95521418, 0.0000, 0.0000, 0.2959153]]))
     N = 20
     th_batch = torch.rand(N, 2)
     tg_batch = chain.forward_kinematics(th_batch)
@@ -98,22 +94,20 @@ def test_urdf():
     ret = chain.forward_kinematics(th)
     tg = ret['lbr_iiwa_link_7']
     pos, rot = quat_pos_from_transform3d(tg)
-    assert quaternion_equality(rot, torch.tensor([7.07106781e-01, 0, -7.07106781e-01, 0], dtype=torch.float64))
-    assert torch.allclose(pos, torch.tensor([-6.60827561e-01, 0, 3.74142136e-01], dtype=torch.float64))
+    assert quaternion_close(rot, torch.tensor([7.07106781e-01, 0, -7.07106781e-01, 0], dtype=torch.float64))
+    assert torch.allclose(pos, torch.tensor([-6.60827561e-01, 0, 3.74142136e-01], dtype=torch.float64), atol=1e-6)
 
 
 def test_urdf_serial():
     chain = pk.build_serial_chain_from_urdf(open(os.path.join(TEST_DIR, "kuka_iiwa.urdf")).read(), "lbr_iiwa_link_7")
     chain.to(dtype=torch.float64)
-    print(chain)
-    print(chain.get_joint_parameter_names())
     th = [0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0]
 
     ret = chain.forward_kinematics(th, end_only=False)
     tg = ret['lbr_iiwa_link_7']
     pos, rot = quat_pos_from_transform3d(tg)
-    assert quaternion_equality(rot, torch.tensor([7.07106781e-01, 0, -7.07106781e-01, 0], dtype=torch.float64))
-    assert torch.allclose(pos, torch.tensor([-6.60827561e-01, 0, 3.74142136e-01], dtype=torch.float64))
+    assert quaternion_close(rot, torch.tensor([7.07106781e-01, 0, -7.07106781e-01, 0], dtype=torch.float64))
+    assert torch.allclose(pos, torch.tensor([-6.60827561e-01, 0, 3.74142136e-01], dtype=torch.float64), atol=1e-6)
 
     N = 1000
     d = "cuda" if torch.cuda.is_available() else "cpu"
@@ -162,7 +156,7 @@ def test_fk_simple_arm():
     })
     tg = ret['arm_wrist_roll']
     pos, rot = quat_pos_from_transform3d(tg)
-    assert quaternion_equality(rot, torch.tensor([0.70710678, 0., 0., 0.70710678], dtype=torch.float64))
+    assert quaternion_close(rot, torch.tensor([0.70710678, 0., 0., 0.70710678], dtype=torch.float64))
     assert torch.allclose(pos, torch.tensor([1.05, 0.55, 0.5], dtype=torch.float64))
 
     N = 100
@@ -176,7 +170,7 @@ def test_sdf_serial_chain():
     chain = chain.to(dtype=torch.float64)
     tg = chain.forward_kinematics([0., math.pi / 2.0, -0.5, 0.])
     pos, rot = quat_pos_from_transform3d(tg)
-    assert quaternion_equality(rot, torch.tensor([0.70710678, 0., 0., 0.70710678], dtype=torch.float64))
+    assert quaternion_close(rot, torch.tensor([0.70710678, 0., 0., 0.70710678], dtype=torch.float64))
     assert torch.allclose(pos, torch.tensor([1.05, 0.55, 0.5], dtype=torch.float64))
 
 
@@ -201,7 +195,7 @@ def test_cuda():
         })
         tg = ret['arm_wrist_roll']
         pos, rot = quat_pos_from_transform3d(tg)
-        assert quaternion_equality(rot, torch.tensor([0.70710678, 0., 0., 0.70710678], dtype=dtype, device=d))
+        assert quaternion_close(rot, torch.tensor([0.70710678, 0., 0., 0.70710678], dtype=dtype, device=d))
         assert torch.allclose(pos, torch.tensor([1.05, 0.55, 0.5], dtype=dtype, device=d))
 
         data = '<robot name="test_robot">' \
@@ -256,7 +250,7 @@ def test_fk_val():
     tg = ret['drive45']
     pos, rot = quat_pos_from_transform3d(tg)
     torch.set_printoptions(precision=6, sci_mode=False)
-    assert quaternion_equality(rot, torch.tensor([0.5, 0.5, -0.5, 0.5], dtype=torch.float64), rtol=1e-4)
+    assert quaternion_close(rot, torch.tensor([0.5, 0.5, -0.5, 0.5], dtype=torch.float64))
     assert torch.allclose(pos, torch.tensor([-0.225692, 0.259045, 0.262139], dtype=torch.float64))
 
 

tests/test_rotation_conversions.py

@@ -3,6 +3,7 @@
 import numpy as np
 import torch
 
+from pytorch_kinematics.transforms.math import quaternion_close
 from pytorch_kinematics.transforms.rotation_conversions import axis_and_angle_to_matrix_33, axis_angle_to_matrix, \
     pos_rot_to_matrix, matrix_to_pos_rot, random_rotations, quaternion_from_euler
 
@@ -24,18 +25,27 @@ def test_axis_angle_to_matrix_perf():
 
 
 def test_quaternion_from_euler():
-    q = quaternion_from_euler(0, 0, 0)
-    np.testing.assert_allclose(q, np.array([1., 0, 0, 0]))
+    q = quaternion_from_euler(torch.tensor([0., 0, 0]))
+    assert quaternion_close(q, torch.tensor([1., 0, 0, 0]))
     root2_over_2 = np.sqrt(2) / 2
 
-    q = quaternion_from_euler(0, 0, np.pi / 2)
-    np.testing.assert_allclose(q, np.array([root2_over_2, 0, 0, root2_over_2]))
+    q = quaternion_from_euler(torch.tensor([0, 0, np.pi / 2]))
+    assert quaternion_close(q, torch.tensor([root2_over_2, 0, 0, root2_over_2], dtype=q.dtype))
 
-    q = quaternion_from_euler(-np.pi / 2, 0, 0)
-    np.testing.assert_allclose(q, np.array([root2_over_2, -root2_over_2, 0, 0]))
+    q = quaternion_from_euler(torch.tensor([-np.pi / 2, 0, 0]))
+    assert quaternion_close(q, torch.tensor([root2_over_2, -root2_over_2, 0, 0], dtype=q.dtype))
 
-    q = quaternion_from_euler(0, np.pi / 2, 0)
-    np.testing.assert_allclose(q, np.array([root2_over_2, 0, root2_over_2, 0]))
+    q = quaternion_from_euler(torch.tensor([0, np.pi / 2, 0]))
+    assert quaternion_close(q, torch.tensor([root2_over_2, 0, root2_over_2, 0], dtype=q.dtype))
+
+    # Test batched
+    b = 32
+    rpy = torch.tensor([0, np.pi / 2, 0])
+    rpy_batch = torch.tile(rpy[None], (b, 1))
+    q_batch = quaternion_from_euler(rpy_batch)
+    q_expected = torch.tensor([root2_over_2, 0, root2_over_2, 0], dtype=q.dtype)
+    q_expected_batch = torch.tile(q_expected[None], (b, 1))
+    assert quaternion_close(q_batch, q_expected_batch)
 
 
 def test_pos_rot_conversion():