Fixes device propogation for noise and adds noise tests

source/extensions/omni.isaac.lab/omni/isaac/lab/utils/noise/noise_model.py

@@ -18,7 +18,20 @@
 
 
 def constant_noise(data: torch.Tensor, cfg: noise_cfg.ConstantNoiseCfg) -> torch.Tensor:
-    """Constant noise."""
+    """Applies a constant noise bias to a given data set.
+
+    Args:
+        data: The unmodified data set to apply noise to.
+        cfg: The configuration parameters for constant noise.
+
+    Returns:
+        The data modified by the noise parameters provided.
+    """
+
+    # fix tensor device for bias on first call and update config parameters
+    if isinstance(cfg.bias, torch.Tensor):
+        cfg.bias = cfg.bias.to(device=data.device)
+
     if cfg.operation == "add":
         return data + cfg.bias
     elif cfg.operation == "scale":
@@ -30,7 +43,23 @@ def constant_noise(data: torch.Tensor, cfg: noise_cfg.ConstantNoiseCfg) -> torch
 
 
 def uniform_noise(data: torch.Tensor, cfg: noise_cfg.UniformNoiseCfg) -> torch.Tensor:
-    """Uniform noise."""
+    """Applies a uniform noise to a given data set.
+
+    Args:
+        data: The unmodified data set to apply noise to.
+        cfg: The configuration parameters for uniform noise.
+
+    Returns:
+        The data modified by the noise parameters provided.
+    """
+
+    # fix tensor device for n_max on first call and update config parameters
+    if isinstance(cfg.n_max, torch.Tensor):
+        cfg.n_max = cfg.n_max.to(data.device)
+    # fix tensor device for n_min on first call and update config parameters
+    if isinstance(cfg.n_min, torch.Tensor):
+        cfg.n_min = cfg.n_min.to(data.device)
+
     if cfg.operation == "add":
         return data + torch.rand_like(data) * (cfg.n_max - cfg.n_min) + cfg.n_min
     elif cfg.operation == "scale":
@@ -42,7 +71,23 @@ def uniform_noise(data: torch.Tensor, cfg: noise_cfg.UniformNoiseCfg) -> torch.T
 
 
 def gaussian_noise(data: torch.Tensor, cfg: noise_cfg.GaussianNoiseCfg) -> torch.Tensor:
-    """Gaussian noise."""
+    """Applies a gaussian noise to a given data set.
+
+    Args:
+        data: The unmodified data set to apply noise to.
+        cfg: The configuration parameters for gaussian noise.
+
+    Returns:
+        The data modified by the noise parameters provided.
+    """
+
+    # fix tensor device for mean on first call and update config parameters
+    if isinstance(cfg.mean, torch.Tensor):
+        cfg.mean = cfg.mean.to(data.device)
+    # fix tensor device for std on first call and update config parameters
+    if isinstance(cfg.std, torch.Tensor):
+        cfg.std = cfg.std.to(data.device)
+
     if cfg.operation == "add":
         return data + cfg.mean + cfg.std * torch.randn_like(data)
     elif cfg.operation == "scale":

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

@@ -0,0 +1,118 @@
+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Launch Isaac Sim Simulator first."""
+
+from omni.isaac.lab.app import AppLauncher, run_tests
+
+# launch omniverse app
+app_launcher = AppLauncher(headless=True)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+
+import torch
+import unittest
+
+import omni.isaac.lab.utils.noise as noise
+
+
+class TestNoise(unittest.TestCase):
+    """Test different noise implementations."""
+
+    def test_gaussian_noise(self):
+        """Test guassian_noise function."""
+
+        for device in ["cpu", "cuda"]:
+            for noise_device in ["cpu", "cuda"]:
+                for op in ["add", "scale", "abs"]:
+                    with self.subTest(device=device, noise_device=noise_device, operation=op):
+                        # create random data set
+                        data = torch.rand(10000, 3, device=device)
+                        # define standard deviation and mean
+                        std = torch.tensor([0.1, 0.2, 0.3], device=noise_device)
+                        mean = torch.tensor([0.4, 0.5, 0.6], device=noise_device)
+                        # create noise config
+                        noise_cfg = noise.GaussianNoiseCfg(std=std, mean=mean, operation=op)
+
+                        for i in range(10):
+                            # apply noise
+                            noisy_data = noise_cfg.func(data, cfg=noise_cfg)
+                            # calculate resulting noise compared to original data set
+                            if op == "add":
+                                std_result, mean_result = torch.std_mean(noisy_data - data, dim=0)
+                            elif op == "scale":
+                                std_result, mean_result = torch.std_mean(noisy_data / data, dim=0)
+                            elif op == "abs":
+                                std_result, mean_result = torch.std_mean(noisy_data, dim=0)
+
+                            self.assertTrue(noise_cfg.mean.device, device)
+                            self.assertTrue(noise_cfg.std.device, device)
+                            torch.testing.assert_close(noise_cfg.std, std_result, atol=1e-2, rtol=1e-2)
+                            torch.testing.assert_close(noise_cfg.mean, mean_result, atol=1e-2, rtol=1e-2)
+
+    def test_uniform_noise(self):
+        """Test uniform_noise function."""
+        for device in ["cpu", "cuda"]:
+            for noise_device in ["cpu", "cuda"]:
+                for op in ["add", "scale", "abs"]:
+                    with self.subTest(device=device, noise_device=noise_device, operation=op):
+                        # create random data set
+                        data = torch.rand(10000, 3, device=device)
+                        # define uniform minimum and maximum
+                        n_min = torch.tensor([0.1, 0.2, 0.3], device=noise_device)
+                        n_max = torch.tensor([0.4, 0.5, 0.6], device=noise_device)
+                        # create noise config
+                        noise_cfg = noise.UniformNoiseCfg(n_max=n_max, n_min=n_min, operation=op)
+
+                        for i in range(10):
+                            # apply noise
+                            noisy_data = noise_cfg.func(data, cfg=noise_cfg)
+                            # calculate resulting noise compared to original data set
+                            if op == "add":
+                                min_result, _ = torch.min(noisy_data - data, dim=0)
+                                max_result, _ = torch.max(noisy_data - data, dim=0)
+                            elif op == "scale":
+                                min_result, _ = torch.min(torch.div(noisy_data, data), dim=0)
+                                max_result, _ = torch.max(torch.div(noisy_data, data), dim=0)
+                            elif op == "abs":
+                                min_result, _ = torch.min(noisy_data, dim=0)
+                                max_result, _ = torch.max(noisy_data, dim=0)
+
+                            self.assertTrue(noise_cfg.n_min.device, device)
+                            self.assertTrue(noise_cfg.n_max.device, device)
+                            self.assertTrue(all(torch.le(noise_cfg.n_min, min_result).tolist()))
+                            self.assertTrue(all(torch.ge(noise_cfg.n_max, max_result).tolist()))
+
+    def test_constant_noise(self):
+        """Test constant_noise"""
+        for device in ["cpu", "cuda"]:
+            for noise_device in ["cpu", "cuda"]:
+                for op in ["add", "scale", "abs"]:
+                    with self.subTest(device=device, noise_device=noise_device, operation=op):
+                        # create random data set
+                        data = torch.rand(10000, 3, device=device)
+                        # define a bias
+                        bias = torch.tensor([0.1, 0.2, 0.3], device=noise_device)
+                        # create noise config
+                        noise_cfg = noise.ConstantNoiseCfg(bias=bias, operation=op)
+
+                        for i in range(10):
+                            # apply noise
+                            noisy_data = noise_cfg.func(data, cfg=noise_cfg)
+                            # calculate resulting noise compared to original data set
+                            if op == "add":
+                                bias_result = noisy_data - data
+                            elif op == "scale":
+                                bias_result = noisy_data / data
+                            elif op == "abs":
+                                bias_result = noisy_data
+
+                            self.assertTrue(noise_cfg.bias.device, device)
+                            torch.testing.assert_close(noise_cfg.bias.repeat(data.shape[0], 1), bias_result)
+
+
+if __name__ == "__main__":
+    run_tests()