Update Omniverse Isaac Gym wrapper to use space utils

skrl/envs/wrappers/torch/omniverse_isaacgym_envs.py

@@ -3,6 +3,7 @@
 import torch
 
 from skrl.envs.wrappers.torch.base import Wrapper
+from skrl.utils.spaces.torch import flatten_tensorized_space, tensorize_space
 
 
 class OmniverseIsaacGymWrapper(Wrapper):
@@ -36,9 +37,10 @@ def step(self, actions: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch
         :return: Observation, reward, terminated, truncated, info
         :rtype: tuple of torch.Tensor and any other info
         """
-        self._observations, reward, terminated, info = self._env.step(actions)
+        observations, reward, terminated, info = self._env.step(actions)
+        self._observations = flatten_tensorized_space(tensorize_space(self.observation_space, observations["obs"]))
         truncated = info["time_outs"] if "time_outs" in info else torch.zeros_like(terminated)
-        return self._observations["obs"], reward.view(-1, 1), terminated.view(-1, 1), truncated.view(-1, 1), info
+        return self._observations, reward.view(-1, 1), terminated.view(-1, 1), truncated.view(-1, 1), info
 
     def reset(self) -> Tuple[torch.Tensor, Any]:
         """Reset the environment
@@ -47,9 +49,10 @@ def reset(self) -> Tuple[torch.Tensor, Any]:
         :rtype: torch.Tensor and any other info
         """
         if self._reset_once:
-            self._observations = self._env.reset()
+            observations = self._env.reset()
+            self._observations = flatten_tensorized_space(tensorize_space(self.observation_space, observations["obs"]))
             self._reset_once = False
-        return self._observations["obs"], {}
+        return self._observations, {}
 
     def render(self, *args, **kwargs) -> None:
         """Render the environment