Updates mimic to support multi-eef (DexMimicGen) data generation (#287)

This PR updates mimic to support multi-eef (DexMimicgen) data
generation.
It consists of the following major changes:
- Updated mimic code to support environments with multiple end effectors
- Added support for setting subtask constraints based on DexMimicGen
- Updated annotate_demos.py to support annotating subtask term signals
for multiple end effectors
- Updated mimic API target_eef_pose_to_action() to take noise as
dictionary of eef noise values instead of a single value

- New feature (non-breaking change which adds functionality)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- This change requires a documentation update

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] I have made corresponding changes to the documentation
- [ ] 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

scripts/imitation_learning/isaaclab_mimic/generate_dataset.py

@@ -7,7 +7,7 @@
 Main data generation script.
 """
 
-# Launching Isaac Sim Simulator first.
+"""Launch Isaac Sim Simulator first."""
 
 import argparse
 
@@ -45,10 +45,15 @@ simulation_app = app_launcher.app
 
 import asyncio
 import gymnasium as gym
+import inspect
 import numpy as np
 import random
 import torch
 
+import omni
+
+from isaaclab.envs import ManagerBasedRLMimicEnv
+
 import isaaclab_mimic.envs  # noqa: F401
 from isaaclab_mimic.datagen.generation import env_loop, setup_async_generation, setup_env_config
 from isaaclab_mimic.datagen.utils import get_env_name_from_dataset, setup_output_paths
@@ -74,12 +79,22 @@ def main():
     )
 
     # create environment
-    env = gym.make(env_name, cfg=env_cfg)
+    env = gym.make(env_name, cfg=env_cfg).unwrapped
+
+    if not isinstance(env, ManagerBasedRLMimicEnv):
+        raise ValueError("The environment should be derived from ManagerBasedRLMimicEnv")
+
+    # check if the mimic API from this environment contains decprecated signatures
+    if "action_noise_dict" not in inspect.signature(env.target_eef_pose_to_action).parameters:
+        omni.log.warn(
+            f'The "noise" parameter in the "{env_name}" environment\'s mimic API "target_eef_pose_to_action", '
+            "is deprecated. Please update the API to take action_noise_dict instead."
+        )
 
     # 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)
+    random.seed(env.cfg.datagen_config.seed)
+    np.random.seed(env.cfg.datagen_config.seed)
+    torch.manual_seed(env.cfg.datagen_config.seed)
 
     # reset before starting
     env.reset()
@@ -95,7 +110,13 @@ def main():
 
     try:
         asyncio.ensure_future(asyncio.gather(*async_components["tasks"]))
-        env_loop(env, async_components["action_queue"], async_components["info_pool"], async_components["event_loop"])
+        env_loop(
+            env,
+            async_components["reset_queue"],
+            async_components["action_queue"],
+            async_components["info_pool"],
+            async_components["event_loop"],
+        )
     except asyncio.CancelledError:
         print("Tasks were cancelled.")
 

source/isaaclab/docs/CHANGELOG.rst

@@ -118,6 +118,18 @@ Changed
   * ``set_fixed_tendon_limit`` → ``set_fixed_tendon_position_limit``
 
 
+0.34.12 (2025-03-06)
+~~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Updated the mimic API :meth:`target_eef_pose_to_action` in :class:`isaaclab.envs.ManagerBasedRLMimicEnv` to take a dictionary of
+  eef noise values instead of a single noise value.
+* Added support for optional subtask constraints based on DexMimicGen to the mimic configuration class :class:`isaaclab.envs.MimicEnvCfg`.
+* Enabled data compression in HDF5 dataset file handler :class:`isaaclab.utils.datasets.hdf5_dataset_file_handler.HDF5DatasetFileHandler`.
+
+
 0.34.11 (2025-03-04)
 ~~~~~~~~~~~~~~~~~~~~
 

source/isaaclab/isaaclab/envs/manager_based_rl_mimic_env.py

@@ -47,7 +47,11 @@ class ManagerBasedRLMimicEnv(ManagerBasedRLEnv):
         raise NotImplementedError
 
     def target_eef_pose_to_action(
-        self, target_eef_pose_dict: dict, gripper_action_dict: dict, noise: float | None = None, env_id: int = 0
+        self,
+        target_eef_pose_dict: dict,
+        gripper_action_dict: dict,
+        action_noise_dict: dict | None = None,
+        env_id: int = 0,
     ) -> torch.Tensor:
         """
         Takes a target pose and gripper action for the end effector controller and returns an action
@@ -57,7 +61,7 @@ class ManagerBasedRLMimicEnv(ManagerBasedRLEnv):
         Args:
             target_eef_pose_dict: Dictionary of 4x4 target eef pose for each end-effector.
             gripper_action_dict: Dictionary of gripper actions for each end-effector.
-            noise: Noise to add to the action. If None, no noise is added.
+            action_noise_dict: Noise to add to the action. If None, no noise is added.
             env_id: Environment index to compute the action for.
 
         Returns:

source/isaaclab/isaaclab/utils/datasets/hdf5_dataset_file_handler.py

@@ -163,7 +163,7 @@ class HDF5DatasetFileHandler(DatasetFileHandlerBase):
                 for sub_key, sub_value in value.items():
                     create_dataset_helper(key_group, sub_key, sub_value)
             else:
-                group.create_dataset(key, data=value.cpu().numpy())
+                group.create_dataset(key, data=value.cpu().numpy(), compression="gzip")
 
         for key, value in episode.data.items():
             create_dataset_helper(h5_episode_group, key, value)

source/isaaclab_mimic/config/extension.toml

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

source/isaaclab_mimic/docs/CHANGELOG.rst

 Changelog
 ---------
 
-1.0.4 (2025-03-10)
+1.0.5 (2025-03-10)
 ~~~~~~~~~~~~~~~~~~
 
 Changed
@@ -15,6 +15,16 @@ Added
 * Added ``Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0`` environment for blueprint vision stacking.
 
 
+1.0.4 (2025-03-07)
+~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Updated data generator to support environments with multiple end effectors.
+* Updated data generator to support subtask constraints based on DexMimicGen.
+
+
 1.0.3 (2025-03-06)
 ~~~~~~~~~~~~~~~~~~
 

source/isaaclab_mimic/isaaclab_mimic/datagen/datagen_info.py

@@ -6,7 +6,6 @@
 """
 Defines structure of information that is needed from an environment for data generation.
 """
-import torch
 from copy import deepcopy
 
 
@@ -46,40 +45,6 @@ class DatagenInfo:
             gripper_action (torch.Tensor or None): gripper actions of shape [..., D] where D
                 is the dimension of the gripper actuation action for the robot arm
         """
-        # Type checks using assert
-        if eef_pose is not None:
-            assert isinstance(
-                eef_pose, torch.Tensor
-            ), f"Expected 'eef_pose' to be of type torch.Tensor, but got {type(eef_pose)}"
-
-        if object_poses is not None:
-            assert isinstance(
-                object_poses, dict
-            ), f"Expected 'object_poses' to be a dictionary, but got {type(object_poses)}"
-            for k, v in object_poses.items():
-                assert isinstance(
-                    v, torch.Tensor
-                ), f"Expected 'object_poses[{k}]' to be of type torch.Tensor, but got {type(v)}"
-
-        if subtask_term_signals is not None:
-            assert isinstance(
-                subtask_term_signals, dict
-            ), f"Expected 'subtask_term_signals' to be a dictionary, but got {type(subtask_term_signals)}"
-            for k, v in subtask_term_signals.items():
-                assert isinstance(
-                    v, (torch.Tensor, int, float)
-                ), f"Expected 'subtask_term_signals[{k}]' to be of type torch.Tensor, int, or float, but got {type(v)}"
-
-        if target_eef_pose is not None:
-            assert isinstance(
-                target_eef_pose, torch.Tensor
-            ), f"Expected 'target_eef_pose' to be of type torch.Tensor, but got {type(target_eef_pose)}"
-
-        if gripper_action is not None:
-            assert isinstance(
-                gripper_action, torch.Tensor
-            ), f"Expected 'gripper_action' to be of type torch.Tensor, but got {type(gripper_action)}"
-
         self.eef_pose = None
         if eef_pose is not None:
             self.eef_pose = eef_pose

source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py

@@ -8,9 +8,9 @@ import contextlib
 import torch
 from typing import Any
 
-from isaaclab.envs import ManagerBasedEnv
+from isaaclab.envs import ManagerBasedRLMimicEnv
 from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg
-from isaaclab.managers import DatasetExportMode
+from isaaclab.managers import DatasetExportMode, TerminationTermCfg
 
 from isaaclab_mimic.datagen.data_generator import DataGenerator
 from isaaclab_mimic.datagen.datagen_info_pool import DataGenInfoPool
@@ -24,23 +24,32 @@ num_attempts = 0
 
 
 async def run_data_generator(
-    env: ManagerBasedEnv,
+    env: ManagerBasedRLMimicEnv,
     env_id: int,
+    env_reset_queue: asyncio.Queue,
     env_action_queue: asyncio.Queue,
     data_generator: DataGenerator,
-    success_term: Any,
+    success_term: TerminationTermCfg,
     pause_subtask: bool = False,
 ):
-    """Run data generator."""
+    """Run mimic data generation from the given data generator in the specified environment index.
+
+    Args:
+        env: The environment to run the data generator on.
+        env_id: The environment index to run the data generation on.
+        env_reset_queue: The asyncio queue to send environment (for this particular env_id) reset requests to.
+        env_action_queue: The asyncio queue to send actions to for executing actions.
+        data_generator: The data generator instance to use.
+        success_term: The success termination term to use.
+        pause_subtask: Whether to pause the subtask during generation.
+    """
     global num_success, num_failures, num_attempts
     while True:
         results = await data_generator.generate(
             env_id=env_id,
             success_term=success_term,
+            env_reset_queue=env_reset_queue,
             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,
         )
         if bool(results["success"]):
@@ -51,22 +60,40 @@ async def run_data_generator(
 
 
 def env_loop(
-    env: ManagerBasedEnv,
+    env: ManagerBasedRLMimicEnv,
+    env_reset_queue: asyncio.Queue,
     env_action_queue: asyncio.Queue,
     shared_datagen_info_pool: DataGenInfoPool,
     asyncio_event_loop: asyncio.AbstractEventLoop,
-) -> None:
-    """Main loop for the environment."""
+):
+    """Main asyncio loop for the environment.
+
+    Args:
+        env: The environment to run the main step loop on.
+        env_reset_queue: The asyncio queue to handle reset request the environment.
+        env_action_queue: The asyncio queue to handle actions to for executing actions.
+        shared_datagen_info_pool: The shared datagen info pool that stores source demo info.
+        asyncio_event_loop: The main asyncio event loop.
+    """
     global num_success, num_failures, num_attempts
+    env_id_tensor = torch.tensor([0], dtype=torch.int64, device=env.device)
     prev_num_attempts = 0
     # simulate environment -- run everything in inference mode
     with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode():
         while True:
 
-            actions = torch.zeros(env.unwrapped.action_space.shape)
+            # check if any environment needs to be reset while waiting for actions
+            while env_action_queue.qsize() != env.num_envs:
+                asyncio_event_loop.run_until_complete(asyncio.sleep(0))
+                while not env_reset_queue.empty():
+                    env_id_tensor[0] = env_reset_queue.get_nowait()
+                    env.reset(env_ids=env_id_tensor)
+                    env_reset_queue.task_done()
+
+            actions = torch.zeros(env.action_space.shape)
 
             # get actions from all the data generators
-            for i in range(env.unwrapped.num_envs):
+            for i in range(env.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
@@ -75,27 +102,30 @@ def env_loop(
             env.step(actions)
 
             # mark done so the data generators can continue with the step results
-            for i in range(env.unwrapped.num_envs):
+            for i in range(env.num_envs):
                 env_action_queue.task_done()
 
             if prev_num_attempts != num_attempts:
                 prev_num_attempts = num_attempts
+                generated_sucess_rate = 100 * num_success / num_attempts if num_attempts > 0 else 0.0
                 print("")
-                print("*" * 50)
-                print(f"have {num_success} successes out of {num_attempts} trials so far")
-                print(f"have {num_failures} failures out of {num_attempts} trials so far")
-                print("*" * 50)
+                print("*" * 50, "\033[K")
+                print(
+                    f"{num_success}/{num_attempts} ({generated_sucess_rate:.1f}%) successful demos generated by"
+                    " mimic\033[K"
+                )
+                print("*" * 50, "\033[K")
 
                 # termination condition is on enough successes if @guarantee_success or enough attempts otherwise
-                generation_guarantee = env.unwrapped.cfg.datagen_config.generation_guarantee
-                generation_num_trials = env.unwrapped.cfg.datagen_config.generation_num_trials
+                generation_guarantee = env.cfg.datagen_config.generation_guarantee
+                generation_num_trials = env.cfg.datagen_config.generation_num_trials
                 check_val = num_success if generation_guarantee else num_attempts
                 if check_val >= generation_num_trials:
                     print(f"Reached {generation_num_trials} successes/attempts. Exiting.")
                     break
 
             # check that simulation is stopped or not
-            if env.unwrapped.sim.is_stopped():
+            if env.sim.is_stopped():
                 break
 
     env.close()
@@ -175,26 +205,28 @@ def setup_async_generation(
         List of asyncio tasks for data generation
     """
     asyncio_event_loop = asyncio.get_event_loop()
+    env_reset_queue = asyncio.Queue()
     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
-    )
+    shared_datagen_info_pool = DataGenInfoPool(env, env.cfg, env.device, asyncio_lock=shared_datagen_info_pool_lock)
     shared_datagen_info_pool.load_from_dataset_file(input_file)
     print(f"Loaded {shared_datagen_info_pool.num_datagen_infos} to datagen info pool")
 
     # Create and schedule data generator tasks
-    data_generator = DataGenerator(env=env.unwrapped, src_demo_datagen_info_pool=shared_datagen_info_pool)
+    data_generator = DataGenerator(env=env, src_demo_datagen_info_pool=shared_datagen_info_pool)
     data_generator_asyncio_tasks = []
     for i in range(num_envs):
         task = asyncio_event_loop.create_task(
-            run_data_generator(env, i, env_action_queue, data_generator, success_term, pause_subtask=pause_subtask)
+            run_data_generator(
+                env, i, env_reset_queue, env_action_queue, data_generator, success_term, pause_subtask=pause_subtask
+            )
         )
         data_generator_asyncio_tasks.append(task)
 
     return {
         "tasks": data_generator_asyncio_tasks,
         "event_loop": asyncio_event_loop,
+        "reset_queue": env_reset_queue,
         "action_queue": env_action_queue,
         "info_pool": shared_datagen_info_pool,
     }

source/isaaclab_mimic/test/test_generate_dataset.py

@@ -59,10 +59,6 @@ class TestGenerateDataset(unittest.TestCase):
             DATASETS_DOWNLOAD_DIR + "/dataset.hdf5",
             "--output_file",
             DATASETS_DOWNLOAD_DIR + "/annotated_dataset.hdf5",
-            "--signals",
-            "grasp_1",
-            "stack_1",
-            "grasp_2",
             "--auto",
             "--headless",
         ]