Puts all tasks unit tests into subTest context (#500)

# Description

This MR moves all environment tasks names into `with.subTest()` context
manager to report per sub test as well.

## Type of change

- Bug fix (non-breaking change which fixes an issue)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] 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
- [ ] 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

---------
Co-authored-by: David Hoeller <dhoeller@ethz.ch>

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

@@ -75,7 +75,7 @@ def convert_to_torch(
         tensor = array
     elif isinstance(array, np.ndarray):
         if array.dtype == np.uint32:
-            array = array.astype(np.int64)
+            array = array.astype(np.int32)
         # need to deal with object arrays (np.void) separately
         tensor = torch.from_numpy(array)
     elif isinstance(array, wp.array):

source/extensions/omni.isaac.orbit/test/sensors/test_camera.py

@@ -309,11 +309,9 @@ class TestCamera(unittest.TestCase):
         self.assertEqual(output["rgb"].shape, hw_3c_shape)
         self.assertEqual(output["distance_to_image_plane"].shape, hw_1c_shape)
         self.assertEqual(output["normals"].shape, hw_3c_shape)
-        # FIXME: No idea why it does not work here. The raw buffers are of type int64 than int32 -> need to investigate
-        #   It works fine when run_usd_camera.py tutorial is run.
-        # self.assertEqual(output["semantic_segmentation"].shape, hw_3c_shape)
-        # self.assertEqual(output["instance_segmentation_fast"].shape, hw_3c_shape)
-        # self.assertEqual(output["instance_id_segmentation_fast"].shape, hw_3c_shape)
+        self.assertEqual(output["semantic_segmentation"].shape, hw_3c_shape)
+        self.assertEqual(output["instance_segmentation_fast"].shape, hw_3c_shape)
+        self.assertEqual(output["instance_id_segmentation_fast"].shape, hw_3c_shape)
 
         # access image data and compare dtype
         output = camera.data.output
@@ -368,11 +366,9 @@ class TestCamera(unittest.TestCase):
         self.assertEqual(output["rgb"].dtype, torch.uint8)
         self.assertEqual(output["distance_to_image_plane"].dtype, torch.float)
         self.assertEqual(output["normals"].dtype, torch.float)
-        # FIXME: No idea why it does not work here. The raw buffers are of type int64 than int32 -> need to investigate
-        #   It works fine when run_usd_camera.py tutorial is run.
-        # self.assertEqual(output["semantic_segmentation"].dtype, torch.int32)
-        # self.assertEqual(output["instance_segmentation_fast"].dtype, torch.int32)
-        # self.assertEqual(output["instance_id_segmentation_fast"].dtype, torch.int32)
+        self.assertEqual(output["semantic_segmentation"].dtype, torch.int32)
+        self.assertEqual(output["instance_segmentation_fast"].dtype, torch.int32)
+        self.assertEqual(output["instance_id_segmentation_fast"].dtype, torch.int32)
 
     def test_throughput(self):
         """Checks that the single camera gets created properly with a rig."""

source/extensions/omni.isaac.orbit_tasks/test/test_environments.py

@@ -54,12 +54,13 @@ class TestEnvironments(unittest.TestCase):
         use_gpu = True
         # iterate over all registered environments
         for task_name in self.registered_tasks:
-            print(f">>> Running test for environment: {task_name}")
-            # check environment
-            self._check_random_actions(task_name, use_gpu, num_envs, num_steps=100)
-            # close the environment
-            print(f">>> Closing environment: {task_name}")
-            print("-" * 80)
+            with self.subTest(task_name=task_name):
+                print(f">>> Running test for environment: {task_name}")
+                # check environment
+                self._check_random_actions(task_name, use_gpu, num_envs, num_steps=100)
+                # close the environment
+                print(f">>> Closing environment: {task_name}")
+                print("-" * 80)
 
     def test_single_instance_gpu(self):
         """Run all environments with single instance and check environments return valid signals."""
@@ -68,12 +69,13 @@ class TestEnvironments(unittest.TestCase):
         use_gpu = True
         # iterate over all registered environments
         for task_name in self.registered_tasks:
-            print(f">>> Running test for environment: {task_name}")
-            # check environment
-            self._check_random_actions(task_name, use_gpu, num_envs, num_steps=100)
-            # close the environment
-            print(f">>> Closing environment: {task_name}")
-            print("-" * 80)
+            with self.subTest(task_name=task_name):
+                print(f">>> Running test for environment: {task_name}")
+                # check environment
+                self._check_random_actions(task_name, use_gpu, num_envs, num_steps=100)
+                # close the environment
+                print(f">>> Closing environment: {task_name}")
+                print("-" * 80)
 
     """
     Helper functions.

source/extensions/omni.isaac.orbit_tasks/test/test_record_video.py

@@ -22,7 +22,7 @@ import unittest
 
 import omni.usd
 
-from omni.isaac.orbit.envs import RLTaskEnv, RLTaskEnvCfg
+from omni.isaac.orbit.envs import RLTaskEnvCfg
 
 import omni.isaac.orbit_tasks  # noqa: F401
 from omni.isaac.orbit_tasks.utils import parse_env_cfg
@@ -56,35 +56,40 @@ class TestRecordVideoWrapper(unittest.TestCase):
     def test_record_video(self):
         """Run random actions agent with recording of videos."""
         for task_name in self.registered_tasks:
-            print(f">>> Running test for environment: {task_name}")
-            # create a new stage
-            omni.usd.get_context().new_stage()
-
-            # parse configuration
-            env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
-
-            # create environment
-            env: RLTaskEnv = gym.make(task_name, cfg=env_cfg, render_mode="rgb_array")
-
-            # directory to save videos
-            videos_dir = os.path.join(self.videos_dir, task_name)
-            # wrap environment to record videos
-            env = gym.wrappers.RecordVideo(
-                env, videos_dir, step_trigger=self.step_trigger, video_length=self.video_length, disable_logger=True
-            )
-
-            # reset environment
-            env.reset()
-            # simulate environment
-            with torch.inference_mode():
-                for _ in range(500):
-                    # compute zero actions
-                    actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1
-                    # apply actions
-                    _ = env.step(actions)
-
-            # close the simulator
-            env.close()
+            with self.subTest(task_name=task_name):
+                print(f">>> Running test for environment: {task_name}")
+                # create a new stage
+                omni.usd.get_context().new_stage()
+
+                # parse configuration
+                env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
+
+                # create environment
+                env = gym.make(task_name, cfg=env_cfg, render_mode="rgb_array")
+
+                # directory to save videos
+                videos_dir = os.path.join(self.videos_dir, task_name)
+                # wrap environment to record videos
+                env = gym.wrappers.RecordVideo(
+                    env,
+                    videos_dir,
+                    step_trigger=self.step_trigger,
+                    video_length=self.video_length,
+                    disable_logger=True,
+                )
+
+                # reset environment
+                env.reset()
+                # simulate environment
+                with torch.inference_mode():
+                    for _ in range(500):
+                        # compute zero actions
+                        actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1
+                        # apply actions
+                        _ = env.step(actions)
+
+                # close the simulator
+                env.close()
 
 
 if __name__ == "__main__":

source/extensions/omni.isaac.orbit_tasks/test/wrappers/test_rl_games_wrapper.py

@@ -41,8 +41,8 @@ class TestRlGamesVecEnvWrapper(unittest.TestCase):
                 cls.registered_tasks.append(task_spec.id)
         # sort environments by name
         cls.registered_tasks.sort()
-        # only pick the first three environments to test
-        cls.registered_tasks = cls.registered_tasks[:3]
+        # only pick the first four environments to test
+        cls.registered_tasks = cls.registered_tasks[:4]
         # print all existing task names
         print(">>> All registered environments:", cls.registered_tasks)
 
@@ -54,36 +54,37 @@ class TestRlGamesVecEnvWrapper(unittest.TestCase):
     def test_random_actions(self):
         """Run random actions and check environments return valid signals."""
         for task_name in self.registered_tasks:
-            print(f">>> Running test for environment: {task_name}")
-            # create a new stage
-            omni.usd.get_context().new_stage()
-            # parse configuration
-            env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
-
-            # create environment
-            env = gym.make(task_name, cfg=env_cfg)
-            # wrap environment
-            env = RlGamesVecEnvWrapper(env, "cuda:0", 100, 100)
-
-            # reset environment
-            obs = env.reset()
-            # check signal
-            self.assertTrue(self._check_valid_tensor(obs))
-
-            # simulate environment for 100 steps
-            with torch.inference_mode():
-                for _ in range(100):
-                    # sample actions from -1 to 1
-                    actions = 2 * torch.rand(env.num_envs, *env.action_space.shape, device=env.device) - 1
-                    # apply actions
-                    transition = env.step(actions)
-                    # check signals
-                    for data in transition:
-                        self.assertTrue(self._check_valid_tensor(data), msg=f"Invalid data: {data}")
-
-            # close the environment
-            print(f">>> Closing environment: {task_name}")
-            env.close()
+            with self.subTest(task_name=task_name):
+                print(f">>> Running test for environment: {task_name}")
+                # create a new stage
+                omni.usd.get_context().new_stage()
+                # parse configuration
+                env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
+
+                # create environment
+                env = gym.make(task_name, cfg=env_cfg)
+                # wrap environment
+                env = RlGamesVecEnvWrapper(env, "cuda:0", 100, 100)
+
+                # reset environment
+                obs = env.reset()
+                # check signal
+                self.assertTrue(self._check_valid_tensor(obs))
+
+                # simulate environment for 100 steps
+                with torch.inference_mode():
+                    for _ in range(100):
+                        # sample actions from -1 to 1
+                        actions = 2 * torch.rand(env.num_envs, *env.action_space.shape, device=env.device) - 1
+                        # apply actions
+                        transition = env.step(actions)
+                        # check signals
+                        for data in transition:
+                            self.assertTrue(self._check_valid_tensor(data), msg=f"Invalid data: {data}")
+
+                # close the environment
+                print(f">>> Closing environment: {task_name}")
+                env.close()
 
     """
     Helper functions.

source/extensions/omni.isaac.orbit_tasks/test/wrappers/test_rsl_rl_wrapper.py

@@ -41,8 +41,8 @@ class TestRslRlVecEnvWrapper(unittest.TestCase):
                 cls.registered_tasks.append(task_spec.id)
         # sort environments by name
         cls.registered_tasks.sort()
-        # only pick the first three environments to test
-        cls.registered_tasks = cls.registered_tasks[:3]
+        # only pick the first four environments to test
+        cls.registered_tasks = cls.registered_tasks[:4]
         # print all existing task names
         print(">>> All registered environments:", cls.registered_tasks)
 
@@ -54,72 +54,76 @@ class TestRslRlVecEnvWrapper(unittest.TestCase):
     def test_random_actions(self):
         """Run random actions and check environments return valid signals."""
         for task_name in self.registered_tasks:
-            print(f">>> Running test for environment: {task_name}")
-            # create a new stage
-            omni.usd.get_context().new_stage()
-            # parse configuration
-            env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
-
-            # create environment
-            env = gym.make(task_name, cfg=env_cfg)
-            # wrap environment
-            env = RslRlVecEnvWrapper(env)
-
-            # reset environment
-            obs, extras = env.reset()
-            # check signal
-            self.assertTrue(self._check_valid_tensor(obs))
-            self.assertTrue(self._check_valid_tensor(extras))
-
-            # simulate environment for 1000 steps
-            with torch.inference_mode():
-                for _ in range(1000):
-                    # sample actions from -1 to 1
-                    actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1
-                    # apply actions
-                    transition = env.step(actions)
-                    # check signals
-                    for data in transition:
-                        self.assertTrue(self._check_valid_tensor(data), msg=f"Invalid data: {data}")
-
-            # close the environment
-            print(f">>> Closing environment: {task_name}")
-            env.close()
+            with self.subTest(task_name=task_name):
+                print(f">>> Running test for environment: {task_name}")
+                # create a new stage
+                omni.usd.get_context().new_stage()
+                # parse configuration
+                env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
+
+                # create environment
+                env = gym.make(task_name, cfg=env_cfg)
+                # wrap environment
+                env = RslRlVecEnvWrapper(env)
+
+                # reset environment
+                obs, extras = env.reset()
+                # check signal
+                self.assertTrue(self._check_valid_tensor(obs))
+                self.assertTrue(self._check_valid_tensor(extras))
+
+                # simulate environment for 1000 steps
+                with torch.inference_mode():
+                    for _ in range(1000):
+                        # sample actions from -1 to 1
+                        actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1
+                        # apply actions
+                        transition = env.step(actions)
+                        # check signals
+                        for data in transition:
+                            self.assertTrue(self._check_valid_tensor(data), msg=f"Invalid data: {data}")
+
+                # close the environment
+                print(f">>> Closing environment: {task_name}")
+                env.close()
 
     def test_no_time_outs(self):
         """Check that environments with finite horizon do not send time-out signals."""
-        for task_name in self.registered_tasks[0:5]:
-            print(f">>> Running test for environment: {task_name}")
-            # create a new stage
-            omni.usd.get_context().new_stage()
-            # parse configuration
-            env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
-            # change to finite horizon
-            env_cfg.is_finite_horizon = True
-
-            # create environment
-            env = gym.make(task_name, cfg=env_cfg)
-            # wrap environment
-            env = RslRlVecEnvWrapper(env)
-
-            # reset environment
-            _, extras = env.reset()
-            # check signal
-            self.assertNotIn("time_outs", extras, msg="Time-out signal found in finite horizon environment.")
-
-            # simulate environment for 10 steps
-            with torch.inference_mode():
-                for _ in range(10):
-                    # sample actions from -1 to 1
-                    actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1
-                    # apply actions
-                    extras = env.step(actions)[-1]
-                    # check signals
-                    self.assertNotIn("time_outs", extras, msg="Time-out signal found in finite horizon environment.")
-
-            # close the environment
-            print(f">>> Closing environment: {task_name}")
-            env.close()
+        for task_name in self.registered_tasks:
+            with self.subTest(task_name=task_name):
+                print(f">>> Running test for environment: {task_name}")
+                # create a new stage
+                omni.usd.get_context().new_stage()
+                # parse configuration
+                env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
+                # change to finite horizon
+                env_cfg.is_finite_horizon = True
+
+                # create environment
+                env = gym.make(task_name, cfg=env_cfg)
+                # wrap environment
+                env = RslRlVecEnvWrapper(env)
+
+                # reset environment
+                _, extras = env.reset()
+                # check signal
+                self.assertNotIn("time_outs", extras, msg="Time-out signal found in finite horizon environment.")
+
+                # simulate environment for 10 steps
+                with torch.inference_mode():
+                    for _ in range(10):
+                        # sample actions from -1 to 1
+                        actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1
+                        # apply actions
+                        extras = env.step(actions)[-1]
+                        # check signals
+                        self.assertNotIn(
+                            "time_outs", extras, msg="Time-out signal found in finite horizon environment."
+                        )
+
+                # close the environment
+                print(f">>> Closing environment: {task_name}")
+                env.close()
 
     """
     Helper functions.

source/extensions/omni.isaac.orbit_tasks/test/wrappers/test_sb3_wrapper.py

@@ -43,8 +43,8 @@ class TestStableBaselines3VecEnvWrapper(unittest.TestCase):
                 cls.registered_tasks.append(task_spec.id)
         # sort environments by name
         cls.registered_tasks.sort()
-        # only pick the first three environments to test
-        cls.registered_tasks = cls.registered_tasks[:3]
+        # only pick the first four environments to test
+        cls.registered_tasks = cls.registered_tasks[:4]
         # print all existing task names
         print(">>> All registered environments:", cls.registered_tasks)
 
@@ -56,36 +56,37 @@ class TestStableBaselines3VecEnvWrapper(unittest.TestCase):
     def test_random_actions(self):
         """Run random actions and check environments return valid signals."""
         for task_name in self.registered_tasks:
-            print(f">>> Running test for environment: {task_name}")
-            # create a new stage
-            omni.usd.get_context().new_stage()
-            # parse configuration
-            env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
-
-            # create environment
-            env = gym.make(task_name, cfg=env_cfg)
-            # wrap environment
-            env = Sb3VecEnvWrapper(env)
-
-            # reset environment
-            obs = env.reset()
-            # check signal
-            self.assertTrue(self._check_valid_array(obs))
-
-            # simulate environment for 1000 steps
-            with torch.inference_mode():
-                for _ in range(1000):
-                    # sample actions from -1 to 1
-                    actions = 2 * np.random.rand(env.num_envs, *env.action_space.shape) - 1
-                    # apply actions
-                    transition = env.step(actions)
-                    # check signals
-                    for data in transition:
-                        self.assertTrue(self._check_valid_array(data), msg=f"Invalid data: {data}")
-
-            # close the environment
-            print(f">>> Closing environment: {task_name}")
-            env.close()
+            with self.subTest(task_name=task_name):
+                print(f">>> Running test for environment: {task_name}")
+                # create a new stage
+                omni.usd.get_context().new_stage()
+                # parse configuration
+                env_cfg: RLTaskEnvCfg = parse_env_cfg(task_name, use_gpu=self.use_gpu, num_envs=self.num_envs)
+
+                # create environment
+                env = gym.make(task_name, cfg=env_cfg)
+                # wrap environment
+                env = Sb3VecEnvWrapper(env)
+
+                # reset environment
+                obs = env.reset()
+                # check signal
+                self.assertTrue(self._check_valid_array(obs))
+
+                # simulate environment for 1000 steps
+                with torch.inference_mode():
+                    for _ in range(1000):
+                        # sample actions from -1 to 1
+                        actions = 2 * np.random.rand(env.num_envs, *env.action_space.shape) - 1
+                        # apply actions
+                        transition = env.step(actions)
+                        # check signals
+                        for data in transition:
+                            self.assertTrue(self._check_valid_array(data), msg=f"Invalid data: {data}")
+
+                # close the environment
+                print(f">>> Closing environment: {task_name}")
+                env.close()
 
     """
     Helper functions.

tools/tests_to_skip.py

@@ -13,6 +13,4 @@ TESTS_TO_SKIP = [
     # orbit_tasks
     "test_data_collector.py",  # Failing
     "test_record_video.py",  # Failing
-    "test_rsl_rl_wrapper.py",  # Timing out (10 minutes)
-    "test_sb3_wrapper.py",  # Timing out (10 minutes)
 ]