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/consolidated_demo.py

+# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Script to record teleoperated demos and run mimic dataset generation in real-time.
+"""
+
+# Launching Isaac Sim Simulator first.
+
+import argparse
+
+from omni.isaac.lab.app import AppLauncher
+
+# add argparse arguments
+parser = argparse.ArgumentParser(
+    description="Record demonstrations and run mimic dataset generation for Isaac Lab environments."
+)
+parser.add_argument("--task", type=str, default=None, help="Name of the task.")
+parser.add_argument(
+    "--num_demos", type=int, default=0, help="Number of demonstrations to record. Set to 0 for infinite."
+)
+parser.add_argument(
+    "--num_envs",
+    type=int,
+    default=5,
+    help=(
+        "Number of environments to instantiate to test recording and generating datasets. The environment specified by"
+        " `teleop_env_index` will be used for teleoperation and recording while the remaining environments will be used"
+        " for real-time data generation. Default is 5."
+    ),
+)
+parser.add_argument(
+    "--teleop_env_index",
+    type=int,
+    default=0,
+    help="Index of the environment to be used for teleoperation. Set -1 for disabling the teleop robot. Default is 0.",
+)
+parser.add_argument("--teleop_device", type=str, default="keyboard", help="Device for interacting with environment.")
+parser.add_argument(
+    "--step_hz", type=int, default=0, help="Environment stepping rate in Hz. Set to 0 for maximum speed."
+)
+parser.add_argument("--input_file", type=str, default=None, help="File path to the source demo dataset file.")
+parser.add_argument(
+    "--output_file",
+    type=str,
+    default="./datasets/output_dataset.hdf5",
+    help="File path to export recorded episodes.",
+)
+parser.add_argument(
+    "--generated_output_file",
+    type=str,
+    default=None,
+    help="File path to export generated episodes by mimic.",
+)
+# append AppLauncher cli args
+AppLauncher.add_app_launcher_args(parser)
+# parse the arguments
+args_cli = parser.parse_args()
+
+# launch the simulator
+app_launcher = AppLauncher(args_cli)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+
+import asyncio
+import contextlib
+import gymnasium as gym
+import numpy as np
+import os
+import random
+import time
+import torch
+
+import omni.isaac.lab_mimic.envs
+from omni.isaac.lab_mimic.datagen.data_generator import DataGenerator
+from omni.isaac.lab_mimic.datagen.datagen_info_pool import DataGenInfoPool
+
+from omni.isaac.lab.devices import Se3Keyboard, Se3SpaceMouse
+from omni.isaac.lab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg
+from omni.isaac.lab.managers import DatasetExportMode, RecorderTerm, RecorderTermCfg
+from omni.isaac.lab.utils import configclass
+from omni.isaac.lab.utils.datasets import HDF5DatasetFileHandler
+
+import omni.isaac.lab_tasks  # noqa: F401
+from omni.isaac.lab_tasks.utils.parse_cfg import parse_env_cfg
+
+# global variable to keep track of the data generation statistics
+num_recorded = 0
+num_success = 0
+num_failures = 0
+num_attempts = 0
+
+
+class PreStepDatagenInfoRecorder(RecorderTerm):
+    """Recorder term that records the datagen info data in each step."""
+
+    def record_pre_step(self):
+        datagen_info = {
+            "object_pose": self._env.scene.get_state(is_relative=True)["rigid_object"],
+            "target_eef_pose": self._env.action_to_target_eef_pos(self._env.action_manager.action),
+        }
+        return "obs", datagen_info
+
+
+@configclass
+class PreStepDatagenInfoRecorderCfg(RecorderTermCfg):
+    """Configuration for the datagen info recorder term."""
+
+    class_type: type[RecorderTerm] = PreStepDatagenInfoRecorder
+
+
+class PreStepSubtaskTermsObservationsRecorder(RecorderTerm):
+    """Recorder term that records the subtask completion observations in each step."""
+
+    def record_pre_step(self):
+        return "obs/subtask_term_signals", self._env.obs_buf["subtask_terms"]
+
+
+@configclass
+class PreStepSubtaskTermsObservationsRecorderCfg(RecorderTermCfg):
+    """Configuration for the step subtask terms observation recorder term."""
+
+    class_type: type[RecorderTerm] = PreStepSubtaskTermsObservationsRecorder
+
+
+@configclass
+class MimicRecorderManagerCfg(ActionStateRecorderManagerCfg):
+    """Mimic specific recorder terms."""
+
+    record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg()
+    record_pre_step_subtask_term_signals = PreStepSubtaskTermsObservationsRecorderCfg()
+
+
+class RateLimiter:
+    """Convenience class for enforcing rates in loops."""
+
+    def __init__(self, hz):
+        """
+        Args:
+            hz (int): frequency to enforce
+        """
+        self.hz = hz
+        self.last_time = time.time()
+        self.sleep_duration = 1.0 / hz
+        self.render_period = min(0.033, self.sleep_duration)
+
+    def sleep(self, env):
+        """Attempt to sleep at the specified rate in hz."""
+        next_wakeup_time = self.last_time + self.sleep_duration
+        while time.time() < next_wakeup_time:
+            time.sleep(self.render_period)
+            env.unwrapped.sim.render()
+
+        self.last_time = self.last_time + self.sleep_duration
+
+        # detect time jumping forwards (e.g. loop is too slow)
+        if self.last_time < time.time():
+            while self.last_time < time.time():
+                self.last_time += self.sleep_duration
+
+
+def pre_process_actions(delta_pose: torch.Tensor, gripper_command: bool) -> torch.Tensor:
+    """Pre-process actions for the environment."""
+    # compute actions based on environment
+    if "Reach" in args_cli.task:
+        # note: reach is the only one that uses a different action space
+        # compute actions
+        return delta_pose
+    else:
+        # resolve gripper command
+        gripper_vel = torch.zeros((delta_pose.shape[0], 1), dtype=torch.float, device=delta_pose.device)
+        gripper_vel[:] = -1 if gripper_command else 1
+        # compute actions
+        return torch.concat([delta_pose, gripper_vel], dim=1)
+
+
+async def run_teleop_robot(
+    env, env_id, env_action_queue, shared_datagen_info_pool, exported_dataset_path, teleop_interface=None
+):
+    """Run teleop robot."""
+    global num_recorded
+    should_reset_teleop_instance = False
+    # create controller if needed
+    if teleop_interface is None:
+        if args_cli.teleop_device.lower() == "keyboard":
+            teleop_interface = Se3Keyboard(pos_sensitivity=0.2, rot_sensitivity=0.5)
+        elif args_cli.teleop_device.lower() == "spacemouse":
+            teleop_interface = Se3SpaceMouse(pos_sensitivity=0.2, rot_sensitivity=0.5)
+        else:
+            raise ValueError(
+                f"Invalid device interface '{args_cli.teleop_device}'. Supported: 'keyboard', 'spacemouse'."
+            )
+
+    # add teleoperation key for reset current recording instance
+    def reset_teleop_instance():
+        nonlocal should_reset_teleop_instance
+        should_reset_teleop_instance = True
+
+    teleop_interface.add_callback("R", reset_teleop_instance)
+
+    teleop_interface.reset()
+    print(teleop_interface)
+
+    recorded_episode_dataset_file_handler = HDF5DatasetFileHandler()
+    recorded_episode_dataset_file_handler.create(exported_dataset_path, env_name=env.unwrapped.cfg.env_name)
+
+    env_id_tensor = torch.tensor([env_id], dtype=torch.int64, device=env.device)
+    success_delay_step_count = 10
+    has_succeeded = False
+    num_recorded = 0
+    while True:
+        if should_reset_teleop_instance:
+            env.unwrapped.recorder_manager.reset(env_id_tensor)
+            env.unwrapped.reset(env_ids=env_id_tensor)
+            should_reset_teleop_instance = False
+
+        # get keyboard command
+        delta_pose, gripper_command = teleop_interface.advance()
+        # convert to torch
+        delta_pose = torch.tensor(delta_pose, dtype=torch.float, device=env.device).repeat(1, 1)
+        # compute actions based on environment
+        teleop_action = pre_process_actions(delta_pose, gripper_command)
+
+        await env_action_queue.put((env_id, teleop_action))
+        await env_action_queue.join()
+
+        if has_succeeded:
+            if success_delay_step_count > 0:
+                success_delay_step_count -= 1
+                continue
+
+            if bool(env.is_success()[env_id]):
+                env.recorder_manager.set_success_to_episodes(
+                    env_id_tensor, torch.tensor([[True]], dtype=torch.bool, device=env.device)
+                )
+                teleop_episode = env.unwrapped.recorder_manager.get_episode(env_id)
+                await shared_datagen_info_pool.add_episode(teleop_episode)
+
+                recorded_episode_dataset_file_handler.write_episode(teleop_episode)
+                recorded_episode_dataset_file_handler.flush()
+                env.recorder_manager.reset(env_id_tensor)
+                num_recorded += 1
+                should_reset_teleop_instance = True
+            else:
+                has_succeeded = False
+        else:
+            if bool(env.is_success()[env_id]):
+                has_succeeded = True
+                success_delay_step_count = 10
+
+
+async def run_data_generator(
+    env, env_id, env_action_queue, shared_datagen_info_pool, pause_subtask=False, export_demo=True
+):
+    """Run data generator."""
+    global num_success, num_failures, num_attempts
+    data_generator = DataGenerator(env=env.unwrapped, src_demo_datagen_info_pool=shared_datagen_info_pool)
+    idle_action = torch.zeros(env.unwrapped.action_space.shape)[0]
+    while True:
+        while data_generator.src_demo_datagen_info_pool.num_datagen_infos < 1:
+            await env_action_queue.put((env_id, idle_action))
+            await env_action_queue.join()
+
+        results = await data_generator.generate(
+            env_id=env_id,
+            env_action_queue=env_action_queue,
+            select_src_per_subtask=env.unwrapped.cfg.datagen_config.generation_select_src_per_subtask,
+            transform_first_robot_pose=env.unwrapped.cfg.datagen_config.generation_transform_first_robot_pose,
+            interpolate_from_last_target_pose=env.unwrapped.cfg.datagen_config.generation_interpolate_from_last_target_pose,
+            pause_subtask=pause_subtask,
+            export_demo=export_demo,
+        )
+        if bool(results["success"]):
+            num_success += 1
+        else:
+            num_failures += 1
+        num_attempts += 1
+
+
+def env_loop(env, env_action_queue, shared_datagen_info_pool, asyncio_event_loop):
+    """Main loop for the environment."""
+    global num_recorded, num_success, num_failures, num_attempts
+    prev_num_attempts = 0
+    prev_num_recorded = 0
+
+    rate_limiter = None
+    if args_cli.step_hz > 0:
+        rate_limiter = RateLimiter(args_cli.step_hz)
+
+    # simulate environment -- run everything in inference mode
+    is_first_print = True
+    with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode():
+        while True:
+
+            actions = torch.zeros(env.unwrapped.action_space.shape)
+
+            # get actions from all the data generators
+            for i in range(env.unwrapped.num_envs):
+                # an async-blocking call to get an action from a data generator
+                env_id, action = asyncio_event_loop.run_until_complete(env_action_queue.get())
+                actions[env_id] = action
+
+            # perform action on environment
+            env.step(actions)
+
+            # mark done so the data generators can continue with the step results
+            for i in range(env.unwrapped.num_envs):
+                env_action_queue.task_done()
+
+            if prev_num_attempts != num_attempts or prev_num_recorded != num_recorded:
+                prev_num_attempts = num_attempts
+                prev_num_recorded = num_recorded
+                generated_sucess_rate = 100 * num_success / num_attempts if num_attempts > 0 else 0.0
+                if is_first_print:
+                    is_first_print = False
+                else:
+                    print("\r", "\033[F" * 5, end="")
+                print("")
+                print("*" * 50, "\033[K")
+                print(f"{num_recorded} teleoperated demos recorded\033[K")
+                print(
+                    f"{num_success}/{num_attempts} ({generated_sucess_rate:.1f}%) successful demos generated by"
+                    " mimic\033[K"
+                )
+                print("*" * 50, "\033[K")
+
+                if args_cli.num_demos > 0 and num_recorded >= args_cli.num_demos:
+                    print(f"All {args_cli.num_demos} demonstrations recorded. Exiting the app.")
+                    break
+
+            # check that simulation is stopped or not
+            if env.unwrapped.sim.is_stopped():
+                break
+
+            if rate_limiter:
+                rate_limiter.sleep(env.unwrapped)
+    env.close()
+
+
+def main():
+    num_envs = args_cli.num_envs
+
+    # create output directory for recorded episodes if it does not exist
+    recorded_output_dir = os.path.dirname(args_cli.output_file)
+    if not os.path.exists(recorded_output_dir):
+        os.makedirs(recorded_output_dir)
+
+    # check if the given input dataset file exists
+    if args_cli.input_file and not os.path.exists(args_cli.input_file):
+        raise FileNotFoundError(f"The dataset file {args_cli.input_file} does not exist.")
+
+    # get the environment name
+    if args_cli.task is not None:
+        env_name = args_cli.task
+    elif args_cli.input_file:
+        # if the environment name is not specified, try to get it from the dataset file
+        dataset_file_handler = HDF5DatasetFileHandler()
+        dataset_file_handler.open(args_cli.input_file)
+        env_name = dataset_file_handler.get_env_name()
+    else:
+        raise ValueError("Task/env name was not specified nor found in the dataset.")
+
+    # parse configuration
+    env_cfg = parse_env_cfg(env_name, device=args_cli.device, num_envs=num_envs)
+    env_cfg.env_name = env_name
+
+    # modify configuration such that the environment runs indefinitely
+    env_cfg.terminations.time_out = None
+    # data generator is in charge of resetting the environment
+    env_cfg.terminations.success = None
+
+    env_cfg.observations.policy.concatenate_terms = False
+
+    env_cfg.recorders = MimicRecorderManagerCfg()
+
+    env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_NONE
+    if args_cli.generated_output_file:
+        # create output directory for generated episodes if it does not exist
+        generated_output_dir = os.path.dirname(args_cli.generated_output_file)
+        if not os.path.exists(generated_output_dir):
+            os.makedirs(generated_output_dir)
+        generated_output_file_name = os.path.splitext(os.path.basename(args_cli.generated_output_file))[0]
+        env_cfg.recorders.dataset_export_dir_path = generated_output_dir
+        env_cfg.recorders.dataset_filename = generated_output_file_name
+        env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY
+
+    # create environment
+    env = gym.make(env_name, cfg=env_cfg)
+
+    # set seed for generation
+    random.seed(env.unwrapped.cfg.datagen_config.seed)
+    np.random.seed(env.unwrapped.cfg.datagen_config.seed)
+    torch.manual_seed(env.unwrapped.cfg.datagen_config.seed)
+
+    # reset before starting
+    env.reset()
+
+    # Set up asyncio stuff
+    asyncio_event_loop = asyncio.get_event_loop()
+    env_action_queue = asyncio.Queue()
+
+    shared_datagen_info_pool_lock = asyncio.Lock()
+    shared_datagen_info_pool = DataGenInfoPool(
+        env.unwrapped, env.unwrapped.cfg, env.unwrapped.device, asyncio_lock=shared_datagen_info_pool_lock
+    )
+    if args_cli.input_file:
+        shared_datagen_info_pool.load_from_dataset_file(args_cli.input_file)
+        print(f"Loaded {shared_datagen_info_pool.num_datagen_infos} to datagen info pool")
+
+    # make data generator object
+    data_generator_asyncio_tasks = []
+    for i in range(num_envs):
+        if args_cli.teleop_env_index is not None and i == args_cli.teleop_env_index:
+            data_generator_asyncio_tasks.append(
+                asyncio_event_loop.create_task(
+                    run_teleop_robot(env, i, env_action_queue, shared_datagen_info_pool, args_cli.output_file)
+                )
+            )
+            continue
+        data_generator_asyncio_tasks.append(
+            asyncio_event_loop.create_task(
+                run_data_generator(
+                    env, i, env_action_queue, shared_datagen_info_pool, export_demo=bool(args_cli.generated_output_file)
+                )
+            )
+        )
+
+    try:
+        asyncio.ensure_future(asyncio.gather(*data_generator_asyncio_tasks))
+    except asyncio.CancelledError:
+        print("Tasks were cancelled.")
+
+    env_loop(env, env_action_queue, shared_datagen_info_pool, asyncio_event_loop)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProgram interrupted by user. Exiting...")
+    # close sim app
+    simulation_app.close()

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(