Fix Isaac Lab observation space (#160)

Toni-SM · Jun 22, 2024 · 932a56b · 932a56b
1 parent cf9e6ad
commit 932a56b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/skrl/envs/wrappers/jax/isaaclab_envs.py b/skrl/envs/wrappers/jax/isaaclab_envs.py
@@ -44,6 +44,8 @@ def __init__(self, env: Any) -> None:
         self._reset_once = True
         self._obs_dict = None
 
+        self._observation_space = self._observation_space["policy"]
+
     def step(self, actions: Union[np.ndarray, jax.Array]) -> \
         Tuple[Union[np.ndarray, jax.Array], Union[np.ndarray, jax.Array],
               Union[np.ndarray, jax.Array], Union[np.ndarray, jax.Array], Any]:

diff --git a/skrl/envs/wrappers/torch/isaaclab_envs.py b/skrl/envs/wrappers/torch/isaaclab_envs.py
@@ -17,6 +17,8 @@ def __init__(self, env: Any) -> None:
         self._reset_once = True
         self._obs_dict = None
 
+        self._observation_space = self._observation_space["policy"]
+
     def step(self, actions: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Any]:
         """Perform a step in the environment