Use spaces utils to process actions in Isaac Gym preview wrapper

Toni-SM · Oct 2, 2024 · b5e27b2 · b5e27b2
1 parent 7e41f75
commit b5e27b2
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/skrl/envs/wrappers/torch/gymnasium_envs.py b/skrl/envs/wrappers/torch/gymnasium_envs.py
@@ -2,7 +2,6 @@
 
 import gymnasium
 
-import numpy as np
 import torch
 
 from skrl import logger

diff --git a/skrl/envs/wrappers/torch/isaacgym_envs.py b/skrl/envs/wrappers/torch/isaacgym_envs.py
@@ -5,7 +5,12 @@
 import torch
 
 from skrl.envs.wrappers.torch.base import Wrapper
-from skrl.utils.spaces.torch import convert_gym_space, flatten_tensorized_space, tensorize_space
+from skrl.utils.spaces.torch import (
+    convert_gym_space,
+    flatten_tensorized_space,
+    tensorize_space,
+    unflatten_tensorized_space
+)
 
 
 class IsaacGymPreview2Wrapper(Wrapper):
@@ -42,7 +47,7 @@ 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
         """
-        observations, reward, terminated, self._info = self._env.step(actions)
+        observations, reward, terminated, self._info = self._env.step(unflatten_tensorized_space(self.action_space, actions))
         self._observations = flatten_tensorized_space(tensorize_space(self.observation_space, observations))
         truncated = self._info["time_outs"] if "time_outs" in self._info else torch.zeros_like(terminated)
         return self._observations, reward.view(-1, 1), terminated.view(-1, 1), truncated.view(-1, 1), self._info
@@ -115,7 +120,7 @@ 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
         """
-        observations, reward, terminated, self._info = self._env.step(actions)
+        observations, reward, terminated, self._info = self._env.step(unflatten_tensorized_space(self.action_space, actions))
         self._observations = flatten_tensorized_space(tensorize_space(self.observation_space, observations["obs"]))
         truncated = self._info["time_outs"] if "time_outs" in self._info else torch.zeros_like(terminated)
         return self._observations, reward.view(-1, 1), terminated.view(-1, 1), truncated.view(-1, 1), self._info