Fixes `wrap_to_pi` function in math utilities (#771)

# Description

The previous implementation of `wrap_to_pi` was rather incorrect since
it would map -PI to PI but -3PI to -PI. Following the general convention
of this function (based on MATLAB), all odd positive multiples of PI
should map to PI, and all negative multiples should map to -PI.

This MR fixes the function and also adds a unit test for it.

Fixes #770

## Type of change

- Bug fix (non-breaking change which fixes an issue)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.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 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.lab/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.20.0"
+version = "0.20.1"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

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

 Changelog
 ---------
 
+0.20.1 (2024-07-26)
+~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Fixed the :meth:`omni.isaac.lab.utils.math.wrap_to_pi` method to handle the wrapping of angles correctly.
+  Earlier, the method was not wrapping the angles to the range [-pi, pi] correctly when the angles were outside
+  the range [-2*pi, 2*pi].
+
+
 0.20.0 (2024-07-26)
 ~~~~~~~~~~~~~~~~~~~
 

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

@@ -93,18 +93,27 @@ def normalize(x: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
 
 @torch.jit.script
 def wrap_to_pi(angles: torch.Tensor) -> torch.Tensor:
-    """Wraps input angles (in radians) to the range [-pi, pi].
+    r"""Wraps input angles (in radians) to the range :math:`[-\pi, \pi]`.
+
+    This function wraps angles in radians to the range :math:`[-\pi, \pi]`, such that
+    :math:`\pi` maps to :math:`\pi`, and :math:`-\pi` maps to :math:`-\pi`. In general,
+    odd positive multiples of :math:`\pi` are mapped to :math:`\pi`, and odd negative
+    multiples of :math:`\pi` are mapped to :math:`-\pi`.
+
+    The function behaves similar to MATLAB's `wrapToPi <https://www.mathworks.com/help/map/ref/wraptopi.html>`_
+    function.
 
     Args:
         angles: Input angles of any shape.
 
     Returns:
-        Angles in the range [-pi, pi].
+        Angles in the range :math:`[-\pi, \pi]`.
     """
-    angles = angles.clone()
-    angles %= 2 * torch.pi
-    angles -= 2 * torch.pi * (angles > torch.pi)
-    return angles
+    # wrap to [0, 2*pi)
+    wrapped_angle = (angles + torch.pi) % (2 * torch.pi)
+    # map to [-pi, pi]
+    # we check for zero in wrapped angle to make it go to pi when input angle is odd multiple of pi
+    return torch.where((wrapped_angle == 0) & (angles > 0), torch.pi, wrapped_angle - torch.pi)
 
 
 @torch.jit.script
@@ -121,8 +130,8 @@ def copysign(mag: float, other: torch.Tensor) -> torch.Tensor:
     Returns:
         The output tensor.
     """
-    mag = torch.tensor(mag, device=other.device, dtype=torch.float).repeat(other.shape[0])
-    return torch.abs(mag) * torch.sign(other)
+    mag_torch = torch.tensor(mag, device=other.device, dtype=torch.float).repeat(other.shape[0])
+    return torch.abs(mag_torch) * torch.sign(other)
 
 
 """

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

@@ -5,7 +5,6 @@
 
 import torch
 import unittest
-from math import pi as PI
 
 """Launch Isaac Sim Simulator first.
 
@@ -18,6 +17,10 @@ from omni.isaac.lab.app import AppLauncher, run_tests
 simulation_app = AppLauncher(headless=True).app
 
 
+"""Rest everything follows."""
+
+from math import pi as PI
+
 import omni.isaac.lab.utils.math as math_utils
 
 
@@ -190,6 +193,40 @@ class TestMathUtilities(unittest.TestCase):
         torch.testing.assert_close(error_3, error_4)
         torch.testing.assert_close(error_4, error_1)
 
+    def test_wrap_to_pi(self):
+        """Test wrap_to_pi method."""
+        # Define test cases
+        # Each tuple contains: angle, expected wrapped angle
+        test_cases = [
+            # No wrapping needed
+            (torch.Tensor([0.0]), torch.Tensor([0.0])),
+            # Positive angle
+            (torch.Tensor([PI]), torch.Tensor([PI])),
+            # Negative angle
+            (torch.Tensor([-PI]), torch.Tensor([-PI])),
+            # Multiple angles
+            (torch.Tensor([3 * PI, -3 * PI, 4 * PI, -4 * PI]), torch.Tensor([PI, -PI, 0.0, 0.0])),
+            # Multiple angles from MATLAB docs
+            # fmt: off
+            (
+                torch.Tensor([-2 * PI, - PI - 0.1, -PI, -2.8, 3.1, PI, PI + 0.001, PI + 1, 2 * PI, 2 * PI + 0.1]),
+                torch.Tensor([0.0, PI - 0.1, -PI, -2.8, 3.1 , PI, -PI + 0.001, -PI + 1 , 0.0, 0.1])
+            ),
+            # fmt: on
+        ]
+
+        # Iterate over test cases
+        for device in ["cpu", "cuda:0"]:
+            for angle, expected_angle in test_cases:
+                with self.subTest(angle=angle, device=device):
+                    # move to the device
+                    angle = angle.to(device)
+                    expected_angle = expected_angle.to(device)
+                    # Compute the wrapped angle
+                    wrapped_angle = math_utils.wrap_to_pi(angle)
+                    # Check that the wrapped angle is close to the expected value
+                    torch.testing.assert_close(wrapped_angle, expected_angle)
+
 
 if __name__ == "__main__":
     run_tests()