Fixes the initial state recorder term for multi-env use cases

isaac-sim · Dec 5, 2024 · 72b4a33 · 72b4a33
1 parent b0312db
commit 72b4a33
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 2 deletions.
diff --git a/source/extensions/omni.isaac.lab/config/extension.toml b/source/extensions/omni.isaac.lab/config/extension.toml
@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.27.18"
+version = "0.27.19"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

diff --git a/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst b/source/extensions/omni.isaac.lab/docs/CHANGELOG.rst
@@ -1,6 +1,15 @@
 Changelog
 ---------
 
+0.27.19 (2024-12-05)
+~~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Fixed the initial state recorder term in :class:`omni.isaac.lab.envs.mdp.recorders.InitialStateRecorder` to return only the states of the specified environment IDs.
+
+
 0.27.18 (2024-12-04)
 ~~~~~~~~~~~~~~~~~~~~
 

diff --git a/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/recorders.py b/source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/recorders/recorders.py
@@ -14,7 +14,13 @@ class InitialStateRecorder(RecorderTerm):
     """Recorder term that records the initial state of the environment after reset."""
 
     def record_post_reset(self, env_ids: Sequence[int] | None):
-        return "initial_state", self._env.scene.get_state(is_relative=True)
+        def extract_env_ids_values(value):
+            nonlocal env_ids
+            if isinstance(value, dict):
+                return {k: extract_env_ids_values(v) for k, v in value.items()}
+            return value[env_ids]
+
+        return "initial_state", extract_env_ids_values(self._env.scene.get_state(is_relative=True))
 
 
 class PostStepStatesRecorder(RecorderTerm):