From a7f82b2ec621c16951356e2c34fa20a433198ee0 Mon Sep 17 00:00:00 2001
From: Toni-SM <toni.semu@gmail.com>
Date: Sat, 4 Jan 2025 23:25:43 -0500
Subject: [PATCH] Isaac Lab wrapper's multi-agent state retrieval with
 gymnasium 1.0

---
 CHANGELOG.md                              | 1 +
 skrl/envs/wrappers/jax/isaaclab_envs.py   | 5 ++++-
 skrl/envs/wrappers/torch/isaaclab_envs.py | 5 ++++-
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 72b559b7..3e14ded4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Model state dictionary initialization for composite Gymnasium spaces in JAX
 - Add missing `reduction` parameter to Gaussian model instantiator
 - Optax's learning rate schedulers integration in JAX implementation
+- Isaac Lab wrapper's multi-agent state retrieval with gymnasium 1.0
 
 ### Removed
 - Remove OpenAI Gym (`gym`) from dependencies and source code. **skrl** continues to support gym environments,
diff --git a/skrl/envs/wrappers/jax/isaaclab_envs.py b/skrl/envs/wrappers/jax/isaaclab_envs.py
index 39584516..576afa52 100644
--- a/skrl/envs/wrappers/jax/isaaclab_envs.py
+++ b/skrl/envs/wrappers/jax/isaaclab_envs.py
@@ -207,7 +207,10 @@ def state(self) -> Union[np.ndarray, jax.Array, None]:
         :return: State
         :rtype: np.ndarray, jax.Array or None
         """
-        state = self._env.state()
+        try:
+            state = self._env.state()
+        except AttributeError:  # 'OrderEnforcing' object has no attribute 'state'
+            state = self._unwrapped.state()
         if state is not None:
             state = flatten_tensorized_space(tensorize_space(next(iter(self.state_spaces.values())), state))
             return _torch2jax(state, self._jax)
diff --git a/skrl/envs/wrappers/torch/isaaclab_envs.py b/skrl/envs/wrappers/torch/isaaclab_envs.py
index e354c196..a23a1d09 100644
--- a/skrl/envs/wrappers/torch/isaaclab_envs.py
+++ b/skrl/envs/wrappers/torch/isaaclab_envs.py
@@ -148,7 +148,10 @@ def state(self) -> torch.Tensor:
         :return: State
         :rtype: torch.Tensor
         """
-        state = self._env.state()
+        try:
+            state = self._env.state()
+        except AttributeError:  # 'OrderEnforcing' object has no attribute 'state'
+            state = self._unwrapped.state()
         if state is not None:
             return flatten_tensorized_space(tensorize_space(next(iter(self.state_spaces.values())), state))
         return state