yaw_quat function optimization and test

source/isaaclab/isaaclab/utils/math.py

@@ -529,13 +529,13 @@ def yaw_quat(quat: torch.Tensor) -> torch.Tensor:
         A quaternion with only yaw component.
     """
     shape = quat.shape
-    quat_yaw = quat.clone().view(-1, 4)
+    quat_yaw = quat.view(-1, 4)
     qw = quat_yaw[:, 0]
     qx = quat_yaw[:, 1]
     qy = quat_yaw[:, 2]
     qz = quat_yaw[:, 3]
     yaw = torch.atan2(2 * (qw * qz + qx * qy), 1 - 2 * (qy * qy + qz * qz))
-    quat_yaw[:] = 0.0
+    quat_yaw = torch.zeros_like(quat_yaw)
     quat_yaw[:, 3] = torch.sin(yaw / 2)
     quat_yaw[:, 0] = torch.cos(yaw / 2)
     quat_yaw = normalize(quat_yaw)

source/isaaclab/test/utils/test_math.py

@@ -275,6 +275,26 @@ def test_wrap_to_pi(self):
                     # Check that the wrapped angle is close to the expected value
                     torch.testing.assert_close(wrapped_angle, expected_angle)
 
+    def test_yaw_quat(self):
+        """
+        Test for yaw_quat methods.
+        """
+        # 90-degree (n/2 radians) rotations about the Y-axis
+        quat_input = torch.Tensor([0.7071, 0, 0.7071, 0])
+        cloned_quat_input = quat_input.clone()
+
+        # Calculated output that the function should return
+        expected_output = torch.Tensor([1, 0, 0, 0])
+
+        # Compute the result using the existing implementation
+        result = math_utils.yaw_quat(quat_input)
+
+        # Verify original quat is not being modified
+        torch.testing.assert_close(quat_input, cloned_quat_input)
+
+        # check that the output is equivalent to the expected output
+        torch.testing.assert_close(result, expected_output)
+
     def test_quat_rotate_and_quat_rotate_inverse(self):
         """Test for quat_rotate and quat_rotate_inverse methods.