Add consolidated demo script for showcasing recording and mimic dataset...

Add consolidated demo script for showcasing recording and mimic dataset generation in real-time in one simulation script (#189)

This PR adds `consolidated_demo.py` script that runs teleop-recording
and real-time mimic dataset generation in the same simulation with
multi-env to showcase the mimic workflow in one script.

It includes changes and fixes needed to enable recording and running
mimic at the same time under multi-env setting.

<!-- As a practice, it is recommended to open an issue to have
discussions on the proposed pull request.
This makes it easier for the community to keep track of what is being
developed or added, and if a given feature
is demanded by more than one party. -->

<!-- As you go through the list, delete the ones that are not
applicable. -->

- Bug fix (non-breaking change which fixes an issue)
- New feature (non-breaking change which adds functionality)

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Co-authored-by: Kelly Guo <kellyg@nvidia.com>

source/extensions/omni.isaac.lab/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.32.7"
+version = "0.32.8"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

source/extensions/omni.isaac.lab/docs/CHANGELOG.rst

 Changelog
 ---------
 
-0.32.7 (2025-01-30)
+0.32.8 (2025-01-30)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -12,7 +12,7 @@ Fixed
   to the event being triggered at the wrong time after the reset.
 
 
-0.32.6 (2025-01-17)
+0.32.7 (2025-01-17)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -42,7 +42,7 @@ Fixed
   the :class:`omni.isaac.lab.assets.RigidObjectCollection` class.
 
 
-0.32.5 (2025-01-14)
+0.32.6 (2025-01-14)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -51,7 +51,7 @@ Fixed
 * Fixed the respawn of only wrong object samples in :func:`repeated_objects_terrain` of :mod:`omni.isaac.lab.terrains.trimesh` module. Previously, the function was respawning all objects in the scene instead of only the wrong object samples, which in worst case could lead to infinite respawn loop.
 
 
-0.32.4 (2025-01-08)
+0.32.5 (2025-01-08)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -61,7 +61,7 @@ Fixed
   In body properties sections, the second dimension should be num_bodies but was documented as 1.
 
 
-0.32.3 (2025-01-02)
+0.32.4 (2025-01-02)
 ~~~~~~~~~~~~~~~~~~~
 
 Added
@@ -70,7 +70,7 @@ Added
 * Added body tracking as an origin type to :class:`omni.isaac.lab.envs.ViewerCfg` and :class:`omni.isaac.lab.envs.ui.ViewportCameraController`.
 
 
-0.32.2 (2024-12-22)
+0.32.3 (2024-12-22)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -79,7 +79,7 @@ Fixed
 * Fixed populating default_joint_stiffness and default_joint_damping values for ImplicitActuator instances in :class:`omni.isaac.lab.assets.Articulation`
 
 
-0.32.1 (2024-12-17)
+0.32.2 (2024-12-17)
 ~~~~~~~~~~~~~~~~~~~
 
 Added
@@ -93,6 +93,19 @@ Added
   :class:`omni.isaac.lab.envs.mdp.actions.OperationalSpaceControllerAction` class.
 
 
+0.32.1 (2024-12-17)
+~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Added a default and generic implementation of the :meth:`get_object_poses` function
+  in the :class:`ManagerBasedRLMimicEnv` class.
+* Added a ``EXPORT_NONE`` mode in the :class:`DatasetExportMode` class and updated
+  :class:`~omni.isaac.lab.managers.RecorderManager` to enable recording without exporting
+  the data to a file.
+
+
 0.32.0 (2024-12-16)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_mimic_env.py

@@ -3,6 +3,7 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+import omni.isaac.lab.utils.math as PoseUtils
 from omni.isaac.lab.envs import ManagerBasedRLEnv
 
 
@@ -86,7 +87,13 @@ class ManagerBasedRLMimicEnv(ManagerBasedRLEnv):
         Returns:
             object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 torch.Tensor)
         """
-        raise NotImplementedError
+        rigid_object_states = self.scene.get_state(is_relative=True)["rigid_object"]
+        object_pose_matrix = dict()
+        for obj_name, obj_state in rigid_object_states.items():
+            object_pose_matrix[obj_name] = PoseUtils.make_pose(
+                obj_state["root_pose"][env_ind, :3], PoseUtils.matrix_from_quat(obj_state["root_pose"][env_ind, 3:7])
+            )
+        return object_pose_matrix
 
     def get_subtask_term_signals(self, env_ind=0):
         """

source/extensions/omni.isaac.lab/omni/isaac/lab/managers/recorder_manager.py

@@ -27,9 +27,10 @@ if TYPE_CHECKING:
 class DatasetExportMode(enum.IntEnum):
     """The mode to handle episode exports."""
 
-    EXPORT_ALL = 0  # Export all episodes to a single dataset file
-    EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES = 1  # Export succeeded and failed episodes in separate files
-    EXPORT_SUCCEEDED_ONLY = 2  # Export only succeeded episodes to a single dataset file
+    EXPORT_NONE = 0  # Export none of the episodes
+    EXPORT_ALL = 1  # Export all episodes to a single dataset file
+    EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES = 2  # Export succeeded and failed episodes in separate files
+    EXPORT_SUCCEEDED_ONLY = 3  # Export only succeeded episodes to a single dataset file
 
 
 @configclass
@@ -157,6 +158,8 @@ class RecorderManager(ManagerBase):
 
         env_name = getattr(env.cfg, "env_name", None)
 
+        self._dataset_file_handler = None
+        if cfg.dataset_export_mode != DatasetExportMode.EXPORT_NONE:
             self._dataset_file_handler = cfg.dataset_file_handler_class_type()
             self._dataset_file_handler.create(
                 os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename), env_name=env_name
@@ -198,6 +201,9 @@ class RecorderManager(ManagerBase):
         if self._dataset_file_handler is not None:
             self._dataset_file_handler.close()
 
+        if self._failed_episode_dataset_file_handler is not None:
+            self._failed_episode_dataset_file_handler.close()
+
     """
     Properties.
     """
@@ -441,6 +447,7 @@ class RecorderManager(ManagerBase):
             self._episodes[env_id] = EpisodeData()
 
         if need_to_flush:
+            if self._dataset_file_handler is not None:
                 self._dataset_file_handler.flush()
             if self._failed_episode_dataset_file_handler is not None:
                 self._failed_episode_dataset_file_handler.flush()

source/extensions/omni.isaac.lab_mimic/config/extension.toml

 [package]
 
 # Semantic Versioning is used: https://semver.org/
-version = "1.0.0"
+version = "1.0.1"
 
 # Description
 category = "isaaclab"

source/extensions/omni.isaac.lab_mimic/docs/CHANGELOG.rst

 Changelog
 ---------
 
+1.0.1 (2024-12-16)
+~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Removed the custom :meth:`get_object_poses` function in the:class:`FrankaCubeStackIKRelMimicEnv`
+  class to use the default implementation from the :class:`ManagerBasedRLMimicEnv` class.
+
+
 1.0.0 (2024-12-06)
 ~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab_mimic/omni/isaac/lab_mimic/datagen/data_generator.py

@@ -194,6 +194,7 @@ class DataGenerator:
         transform_first_robot_pose=False,
         interpolate_from_last_target_pose=True,
         pause_subtask=False,
+        export_demo=True,
     ):
         """
         Attempt to generate a new demonstration.
@@ -415,6 +416,7 @@ class DataGenerator:
         self.env.recorder_manager.set_success_to_episodes(
             env_id_tensor, torch.tensor([[generated_success]], dtype=torch.bool, device=self.env.device)
         )
+        if export_demo:
             self.env.recorder_manager.export_episodes(env_id_tensor)
 
         results = dict(

source/extensions/omni.isaac.lab_mimic/omni/isaac/lab_mimic/envs/franka_stack_ik_rel_mimic_env.py

@@ -124,31 +124,6 @@ class FrankaCubeStackIKRelMimicEnv(ManagerBasedRLMimicEnv):
         # last dimension is gripper action
         return action[:, -1:]
 
-    def get_object_poses(self, env_ind=0):
-        """
-        Gets the pose of each object relevant to Isaac Lab Mimic data generation in the current scene.
-
-        Returns:
-            object_poses (dict): dictionary that maps object name (str) to object pose matrix (4x4 torch.Tensor)
-        """
-        # three relevant objects - three cubes
-        # Retrieve end effector pose from the observation buffer
-        cube_positions = self.obs_buf["policy"]["cube_positions"][env_ind]
-        cube_orientations = self.obs_buf["policy"]["cube_orientations"][env_ind]
-        cube_1_pos = cube_positions[:3]
-        cube_1_rot = cube_orientations[:4]
-        cube_2_pos = cube_positions[3:6]
-        cube_2_rot = cube_orientations[4:8]
-        cube_3_pos = cube_positions[6:]
-        cube_3_rot = cube_orientations[8:]
-
-        # Quaternion format is w,x,y,z
-        return dict(
-            cube_1=PoseUtils.make_pose(cube_1_pos, PoseUtils.matrix_from_quat(cube_1_rot)),
-            cube_2=PoseUtils.make_pose(cube_2_pos, PoseUtils.matrix_from_quat(cube_2_rot)),
-            cube_3=PoseUtils.make_pose(cube_3_pos, PoseUtils.matrix_from_quat(cube_3_rot)),
-        )
-
     def get_subtask_term_signals(self, env_ind=0):
         """
         Gets a dictionary of binary flags for each subtask in a task. The flag is 1

source/extensions/omni.isaac.lab_tasks/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.10.20"
+version = "0.10.21"
 
 # Description
 title = "Isaac Lab Environments"

source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst

 Changelog
 ---------
 
-0.10.20 (2025-01-03)
+0.10.21 (2025-01-03)
 ~~~~~~~~~~~~~~~~~~~
 
 Fixed
@@ -10,7 +10,7 @@ Fixed
 * Fixed the reset of the actions in the function overriding of the low level observations of :class:`omni.isaac.lab_tasks.manager_based.navigation.mdp.PreTrainedPolicyAction`.
 
 
-0.10.19 (2024-12-17)
+0.10.20 (2024-12-17)
 ~~~~~~~~~~~~~~~~~~~~
 
 Changed
@@ -21,6 +21,16 @@ Changed
   inside the ``Isaac-Reach-Franka-OSC-v0`` environment to enable nullspace control.
 
 
+0.10.19 (2024-12-17)
+~~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Fixed :meth:`omni.isaac.lab_tasks.manager_based.manipulation.stack.mdp.ee_frame_pos` to output
+  ``ee_frame_pos`` with respect to the environment's origin.
+
+
 0.10.18 (2024-12-16)
 ~~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/stack/mdp/observations.py

@@ -244,7 +244,7 @@ def instance_randomize_object_obs(
 
 def ee_frame_pos(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor:
     ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name]
-    ee_frame_pos = ee_frame.data.target_pos_w[:, 0, :]
+    ee_frame_pos = ee_frame.data.target_pos_w[:, 0, :] - env.scene.env_origins[:, 0:3]
 
     return ee_frame_pos
 

source/standalone/workflows/isaac_lab_mimic/annotate_demos.py

@@ -3,16 +3,18 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Script to add mimic annotations to demos to be used as source demos for mimic dataset generation."""
+"""
+Script to add mimic annotations to demos to be used as source demos for mimic dataset generation.
+"""
 
-"""Launch Isaac Sim Simulator first."""
+# Launching Isaac Sim Simulator first.
 
 import argparse
 
 from omni.isaac.lab.app import AppLauncher
 
 # add argparse arguments
-parser = argparse.ArgumentParser(description="Collect demonstrations for Isaac Lab environments.")
+parser = argparse.ArgumentParser(description="Annotate demonstrations for Isaac Lab environments.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
 parser.add_argument(
     "--input_file", type=str, default="./datasets/dataset.hdf5", help="File name of the dataset to be annotated."

source/standalone/workflows/isaac_lab_mimic/generate_dataset.py

@@ -14,7 +14,7 @@ import argparse
 from omni.isaac.lab.app import AppLauncher
 
 # add argparse arguments
-parser = argparse.ArgumentParser(description="Record demonstrations for Isaac Lab environments.")
+parser = argparse.ArgumentParser(description="Generate demonstrations for Isaac Lab environments.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
 parser.add_argument("--generation_num_trials", type=int, help="Number of demos to be generated.", default=None)
 parser.add_argument(