Adapts certain utils.math functions to process tensors of higher dimensions (#441)

# Description

As reported in [this
issue](https://github.com/NVIDIA-Omniverse/orbit/issues/252), we have a
mismatch between the behavior described in the documentation and the
actual behavior of `orbit.utils.math.axis_angle_from_quat`. We currently
only accept tensors of the form (N,4), but this PR allows us to accept
(...,4) as described. A corresponding change has been made to
`orbit.utils.math.quat_error_magnitude` as it is a reasonable extension
requested in the original issue.

I have also added tests for these functions.

## Type of change

- Bug fix (non-breaking change which fixes an issue)
- New feature (non-breaking change which adds functionality)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] I have run all the tests with `./orbit.sh --test` and they pass
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

source/extensions/omni.isaac.orbit/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.11.2"
+version = "0.11.3"
 
 # Description
 title = "ORBIT framework for Robot Learning"

source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst

 Changelog
 ---------
 
+0.11.3 (2024-03-04)
+~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Corrects the functions :func:`omni.isaac.orbit.utils.math.axis_angle_from_quat` and :func:`omni.isaac.orbit.utils.math.quat_error_magnitude`
+  to accept tensors of the form (..., 4) instead of (N, 4). This brings us in line with our documentation and also upgrades one of our functions
+  to handle higher dimensions.
+
+
 0.11.2 (2024-03-04)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/math.py

@@ -634,7 +634,7 @@ def axis_angle_from_quat(quat: torch.Tensor, eps: float = 1.0e-6) -> torch.Tenso
     # When theta = 0, (sin(theta/2) / theta) is undefined
     # However, as theta --> 0, we can use the Taylor approximation 1/2 - theta^2 / 48
     quat = quat * (1.0 - 2.0 * (quat[..., 0:1] < 0.0))
-    mag = torch.linalg.norm(quat[..., 1:], dim=1)
+    mag = torch.linalg.norm(quat[..., 1:], dim=-1)
     half_angle = torch.atan2(mag, quat[..., 0])
     angle = 2.0 * half_angle
     # check whether to apply Taylor approximation
@@ -649,14 +649,14 @@ def quat_error_magnitude(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
     """Computes the rotation difference between two quaternions.
 
     Args:
-        q1: The first quaternion in (w, x, y, z). Shape is (N, 4).
-        q2: The second quaternion in (w, x, y, z). Shape is (N, 4).
+        q1: The first quaternion in (w, x, y, z). Shape is (..., 4).
+        q2: The second quaternion in (w, x, y, z). Shape is (..., 4).
 
     Returns:
         Angular error between input quaternions in radians.
     """
     quat_diff = quat_mul(q1, quat_conjugate(q2))
-    return torch.norm(axis_angle_from_quat(quat_diff), dim=1)
+    return torch.norm(axis_angle_from_quat(quat_diff), dim=-1)
 
 
 @torch.jit.script

source/extensions/omni.isaac.orbit/test/utils/test_math.py

@@ -8,6 +8,7 @@ from __future__ import annotations
 import torch
 import traceback
 import unittest
+from math import pi as PI
 
 """Launch Isaac Sim Simulator first.
 
@@ -50,6 +51,67 @@ class TestMathUtilities(unittest.TestCase):
 
         self.assertFalse(math_utils.is_identity_pose(identity_pos_multi_row, identity_rot_multi_row))
 
+    def test_axis_angle_from_quat(self):
+        """Test axis_angle_from_quat method."""
+        # Quaternions of the form (2,4) and (2,2,4)
+        quats = [
+            torch.Tensor([[1.0, 0.0, 0.0, 0.0], [0.8418536, 0.142006, 0.0, 0.5206887]]),
+            torch.Tensor([
+                [[1.0, 0.0, 0.0, 0.0], [0.8418536, 0.142006, 0.0, 0.5206887]],
+                [[1.0, 0.0, 0.0, 0.0], [0.9850375, 0.0995007, 0.0995007, 0.0995007]],
+            ]),
+        ]
+
+        # Angles of the form (2,3) and (2,2,3)
+        angles = [
+            torch.Tensor([[0.0, 0.0, 0.0], [0.3, 0.0, 1.1]]),
+            torch.Tensor([[[0.0, 0.0, 0.0], [0.3, 0.0, 1.1]], [[0.0, 0.0, 0.0], [0.2, 0.2, 0.2]]]),
+        ]
+
+        for quat, angle in zip(quats, angles):
+            with self.subTest(quat=quat, angle=angle):
+                self.assertTrue(torch.allclose(math_utils.axis_angle_from_quat(quat), angle, atol=1e-7))
+
+    def test_axis_angle_from_quat_approximation(self):
+        """Test Taylor approximation from axis_angle_from_quat method
+        for unstable conversions where theta is very small."""
+        # Generate a small rotation quaternion
+        # Small angle
+        theta = torch.Tensor([0.0000001])
+        # Arbitrary normalized axis of rotation in rads, (x,y,z)
+        axis = [-0.302286, 0.205494, -0.930803]
+        # Generate quaternion
+        qw = torch.cos(theta / 2)
+        quat_vect = [qw] + [d * torch.sin(theta / 2) for d in axis]
+        quaternion = torch.tensor(quat_vect, dtype=torch.float32)
+
+        # Convert quaternion to axis-angle
+        axis_angle_computed = math_utils.axis_angle_from_quat(quaternion)
+
+        # Expected axis-angle representation
+        axis_angle_expected = torch.tensor([theta * d for d in axis], dtype=torch.float32)
+
+        # Assert that the computed values are close to the expected values
+        self.assertTrue(torch.allclose(axis_angle_computed, axis_angle_expected, atol=1e-7))
+
+    def test_quat_error_magnitude(self):
+        # Define test cases
+        test_cases = [
+            # q1, q2, expected error
+            # No rotation
+            (torch.Tensor([1, 0, 0, 0]), torch.Tensor([1, 0, 0, 0]), torch.Tensor([0.0])),
+            # PI/2 rotation
+            (torch.Tensor([1.0, 0, 0.0, 0]), torch.Tensor([0.7071068, 0.7071068, 0, 0]), torch.Tensor([PI / 2])),
+            # PI rotation
+            (torch.Tensor([1.0, 0, 0.0, 0]), torch.Tensor([0.0, 0.0, 1.0, 0]), torch.Tensor([PI])),
+        ]
+        # Test higher dimension
+        test_cases += tuple([(torch.stack(tensors) for tensors in zip(*test_cases))])
+        for q1, q2, expected_diff in test_cases:
+            with self.subTest(q1=q1, q2=q2):
+                q12_diff = math_utils.quat_error_magnitude(q1, q2)
+                self.assertTrue(torch.allclose(q12_diff, torch.flatten(expected_diff), atol=1e-7))
+
 
 if __name__ == "__main__":
     try: