Optimizes `yaw_quat` implementation (#2247)

# Description

Added optimizations to the yaw_quat function and added a function to
test_math to make sure it works as intended.

**BENCHMARKS:**

Device: `cpu`

```
Time for yaw_quat: <torch.utils.benchmark.utils.common.Measurement object at 0x7a342d1f8640>
math_utils.yaw_quat(q_rand)  `511.79 us`
1 measurement, 1000 runs , 1 thread

Time for iter_yaw_quat: <torch.utils.benchmark.utils.common.Measurement object at 0x7a342d1f9720>
iter_yaw_quat(q_rand)  `2.18 ms`
1 measurement, 1000 runs , 1 thread

Time for iter_old_yaw_quat: <torch.utils.benchmark.utils.common.Measurement object at 0x7a342d1f8640>
iter_old_yaw_quat(q_rand)  `2.36 ms`
1 measurement, 1000 runs , 1 thread
```

---

Device: `cuda:0`

```
Time for yaw_quat: <torch.utils.benchmark.utils.common.Measurement object at 0x7a342d1f8730>
math_utils.yaw_quat(q_rand)  `246.16 us`
1 measurement, 1000 runs , 1 thread

Time for iter_yaw_quat: <torch.utils.benchmark.utils.common.Measurement object at 0x7a342d1f8100>
iter_yaw_quat(q_rand)  `2.97 ms`
1 measurement, 1000 runs , 1 thread

Time for iter_old_yaw_quat: <torch.utils.benchmark.utils.common.Measurement object at 0x7a342d1f8730>
iter_old_yaw_quat(q_rand)  `3.46 ms`
1 measurement, 1000 runs , 1 thread
```

## Type of change

- New feature (non-breaking change which adds functionality)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] 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
- [ ] 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/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 @@ class TestMathUtilities(unittest.TestCase):
                     # 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.