Adds optimizations for Isaac Lab Mimic data generation performance (#5141)

# Description

<!--
Thank you for your interest in sending a pull request. Please make sure
to check the contribution guidelines.

Link:
https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html

💡 Please try to keep PRs small and focused. Large PRs are harder to
review and merge.
-->

This PR adds performance improvements and additional test case coverage
for Isaac Lab Mimic data generation.

**Performance improvements:**

- Removal of usage of large state tensors in waypoint.py `execute` and
return task success only
- Change asyncio_lock usage to be optional for data generation flows
which does not require it
- Improve asynchronous execution in generation.py `env_loop`
- Add toggle to clone data in `add` method of episode_data.py allowing
for RecorderManager to clone contiguous tensors before calling `add`
- Add option to toggle dataset compression (reduces env step time for
state-only data generation)
- Use tilted cameras for visuomotor envs (Franka & GR1 nut pouring)

**Additional test coverage:**

- Data generation test for Franka visuomotor env
- Data generation test for GR1T2 pick place env
- Data generation test for GR1T2 nut pouring env

**Other changes:**

- Increase data generation test time outs to prevent test failure due to
insufficient time
- Check for and return exception in generation.py `env_loop` to prevent
event loop from spinning indefinitely when exception occurs in data
generation.
- Locks h5py dep version to 3.15.1 (latest version 3.16.0 has import
errors on Windows)

## Type of change

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

- New feature (non-breaking change which adds functionality)

## Checklist

- [x] I have read and understood the [contribution
guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html)
- [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
- [x] 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
-->

---------

Signed-off-by: peterd-NV <peterd@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Kelly Guo <kellyg@nvidia.com>

Author: 3 people

scripts/imitation_learning/isaaclab_mimic/generate_dataset.py

@@ -30,14 +30,21 @@
 parser.add_argument(
     "--pause_subtask",
     action="store_true",
-    help="pause after every subtask during generation for debugging - only useful with render flag",
+    help="Pause after every subtask during generation for debugging - only useful with render flag",
 )
 parser.add_argument(
     "--use_skillgen",
     action="store_true",
     default=False,
-    help="use skillgen to generate motion trajectories",
+    help="Use skillgen to generate motion trajectories",
 )
+parser.add_argument(
+    "--disable_dataset_compression",
+    action="store_true",
+    default=False,
+    help="Disables dataset compression",
+)
+
 # append AppLauncher cli args
 AppLauncher.add_app_launcher_args(parser)
 # parse the arguments
@@ -88,6 +95,7 @@ def main():
         num_envs=num_envs,
         device=args_cli.device,
         generation_num_trials=args_cli.generation_num_trials,
+        dataset_compression=not args_cli.disable_dataset_compression,
     )
 
     # Create environment
@@ -158,6 +166,7 @@ def main():
                 async_components["action_queue"],
                 async_components["info_pool"],
                 async_components["event_loop"],
+                data_gen_tasks=data_gen_tasks,
             )
         except asyncio.CancelledError:
             print("Tasks were cancelled.")

source/isaaclab/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "4.5.28"
+version = "4.5.29"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

source/isaaclab/docs/CHANGELOG.rst

@@ -1,6 +1,20 @@
 Changelog
 ---------
 
+4.5.29 (2026-04-10)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added flag to toggle dataset compression in RecorderManager and dataset file handler.
+
+Changed
+^^^^^^^
+
+* Changed RecorderManager to clone value tensors before adding to episode data, removing multiple clones in ``episodes.add()`` and replacing with a single clone.
+
+
 4.5.28 (2026-04-10)
 ~~~~~~~~~~~~~~~~~~~
 

source/isaaclab/isaaclab/managers/recorder_manager.py

@@ -55,6 +55,9 @@ class RecorderManagerBaseCfg:
     export_in_close: bool = False
     """Whether to export episodes in the close call."""
 
+    dataset_compression: bool = True
+    """Enable dataset compression."""
+
 
 class RecorderTerm(ManagerTermBase):
     """Base class for recorder terms.
@@ -335,13 +338,14 @@ def add_to_episodes(self, key: str, value: torch.Tensor | dict, env_ids: Sequenc
                 self.add_to_episodes(f"{key}/{sub_key}", sub_value, env_ids)
             return
 
+        if isinstance(value, wp.array):
+            value = wp.to_torch(value)
+        value = value.clone()  # Clone once for all envs
         for value_index, env_id in enumerate(env_ids):
             if env_id not in self._episodes:
                 self._episodes[env_id] = EpisodeData()
                 self._episodes[env_id].env_id = env_id
-            if isinstance(value, wp.array):
-                value = wp.to_torch(value)
-            self._episodes[env_id].add(key, value[value_index])
+            self._episodes[env_id].add(key, value[value_index], clone=False)
 
     def set_success_to_episodes(self, env_ids: Sequence[int] | None, success_values: torch.Tensor):
         """Sets the task success values to the episodes for the given environment ids.
@@ -513,7 +517,9 @@ def export_episodes(self, env_ids: Sequence[int] | None = None, demo_ids: Sequen
                 if target_dataset_file_handler is not None:
                     # Use corresponding demo_id if provided, otherwise None
                     current_demo_id = demo_ids[i] if demo_ids is not None else None
-                    target_dataset_file_handler.write_episode(self._episodes[env_id], current_demo_id)
+                    target_dataset_file_handler.write_episode(
+                        self._episodes[env_id], current_demo_id, self.cfg.dataset_compression
+                    )
                     need_to_flush = True
                 # Update episode count
                 if episode_succeeded:
@@ -567,6 +573,7 @@ def _prepare_terms(self):
                 "dataset_export_mode",
                 "export_in_record_pre_reset",
                 "export_in_close",
+                "dataset_compression",
             ]:
                 continue
             # check if term config is None

source/isaaclab/isaaclab/utils/datasets/episode_data.py

@@ -3,11 +3,6 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
-# All rights reserved.
-#
-# SPDX-License-Identifier: BSD-3-Clause
-
 from __future__ import annotations
 
 import torch
@@ -90,7 +85,7 @@ def is_empty(self):
         """Check if the episode data is empty."""
         return not bool(self._data)
 
-    def add(self, key: str, value: torch.Tensor | dict):
+    def add(self, key: str, value: torch.Tensor | dict, clone: bool = True):
         """Add a key-value pair to the dataset.
 
         The key can be nested by using the "/" character.
@@ -99,23 +94,25 @@ def add(self, key: str, value: torch.Tensor | dict):
         Args:
             key: The key name.
             value: The corresponding value of tensor type or of dict type.
+            clone: Whether to clone the tensor value before storing it in the episode data.
         """
         # check datatype
         if isinstance(value, dict):
             for sub_key, sub_value in value.items():
-                self.add(f"{key}/{sub_key}", sub_value)
+                self.add(f"{key}/{sub_key}", sub_value, clone=clone)
             return
 
+        stored = value.clone() if (clone and isinstance(value, torch.Tensor)) else value
         sub_keys = key.split("/")
         current_dataset_pointer = self._data
         for sub_key_index in range(len(sub_keys)):
             if sub_key_index == len(sub_keys) - 1:
                 # Add value to the final dict layer
                 # Use lists to prevent slow tensor copy during concatenation
                 if sub_keys[sub_key_index] not in current_dataset_pointer:
-                    current_dataset_pointer[sub_keys[sub_key_index]] = [value.clone()]
+                    current_dataset_pointer[sub_keys[sub_key_index]] = [stored]
                 else:
-                    current_dataset_pointer[sub_keys[sub_key_index]].append(value.clone())
+                    current_dataset_pointer[sub_keys[sub_key_index]].append(stored)
                 break
             # key index
             if sub_keys[sub_key_index] not in current_dataset_pointer:

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

@@ -3,11 +3,6 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
-# All rights reserved.
-#
-# SPDX-License-Identifier: BSD-3-Clause
-
 from __future__ import annotations
 
 import json
@@ -210,7 +205,7 @@ def load_dataset_helper(group, path=""):
 
         return episode
 
-    def write_episode(self, episode: EpisodeData, demo_id: int | None = None):
+    def write_episode(self, episode: EpisodeData, demo_id: int | None = None, dataset_compression: bool = True):
         """Add an episode to the dataset.
 
         Args:
@@ -251,7 +246,10 @@ def create_dataset_helper(group, key, value):
                 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(), compression="gzip")
+                if dataset_compression:
+                    group.create_dataset(key, data=value.cpu().numpy(), compression="gzip", compression_opts=2)
+                else:
+                    group.create_dataset(key, data=value.cpu().numpy())
 
         for key, value in episode.data.items():
             create_dataset_helper(h5_episode_group, key, value)

source/isaaclab_mimic/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Semantic Versioning is used: https://semver.org/
-version = "1.2.3"
+version = "1.2.4"
 
 # Description
 category = "isaaclab"

source/isaaclab_mimic/docs/CHANGELOG.rst

@@ -1,6 +1,21 @@
 Changelog
 ---------
 
+1.2.4 (2026-04-06)
+~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Made performance enhancing changes to data generation pipeline (elimate large tensor usage, reduce asyncio overhead and blocking)
+* Locked h5py dependency to last stable version 3.15.1 to prevent package import errors on Windows with version 3.16.0.
+
+Added
+^^^^^
+
+* Added data generation test cases for all tasks (single and multi environment).
+
+
 1.2.3 (2026-03-12)
 ~~~~~~~~~~~~~~~~~~~
 

source/isaaclab_mimic/isaaclab_mimic/datagen/data_generator.py

@@ -6,6 +6,7 @@
 """Base class for data generator."""
 
 import asyncio
+import contextlib
 import copy
 import logging
 from typing import Any
@@ -33,6 +34,16 @@
 from .datagen_info_pool import DataGenInfoPool
 
 
+@contextlib.asynccontextmanager
+async def _optional_lock(lock):
+    """Async context manager that acquires the lock only if it is not None."""
+    if lock is not None:
+        async with lock:
+            yield
+    else:
+        yield
+
+
 def transform_source_data_segment_using_delta_object_pose(
     src_eef_poses: torch.Tensor,
     delta_obj_pose: torch.Tensor,
@@ -664,10 +675,7 @@ async def generate(  # noqa: C901
         for subtask_constraint in self.env_cfg.task_constraint_configs:
             runtime_subtask_constraints_dict.update(subtask_constraint.generate_runtime_subtask_constraints())
 
-        # save generated data in these variables
-        generated_states = []
-        generated_obs = []
-        generated_actions = []
+        # Track if the generated trajectory was successful
         generated_success = False
 
         # some eef-specific state variables used during generation
@@ -694,7 +702,8 @@ async def generate(  # noqa: C901
 
         # While loop that runs per time step
         while True:
-            async with self.src_demo_datagen_info_pool.asyncio_lock:
+            await asyncio.sleep(0)
+            async with _optional_lock(self.src_demo_datagen_info_pool.asyncio_lock):
                 if len(self.src_demo_datagen_info_pool.datagen_infos) > prev_src_demo_datagen_info_pool_size:
                     # src_demo_datagen_info_pool at this point may be updated with new demos,
                     # So we need to update subtask boundaries again
@@ -871,20 +880,18 @@ async def generate(  # noqa: C901
                 eef_waypoint_dict[eef_name] = waypoint
             multi_waypoint = MultiWaypoint(eef_waypoint_dict)
 
+            await asyncio.sleep(0)
+
             # Execute the next waypoints for all eefs
-            exec_results = await multi_waypoint.execute(
+            exec_success = await multi_waypoint.execute(
                 env=self.env,
                 success_term=success_term,
                 env_id=env_id,
                 env_action_queue=env_action_queue,
             )
 
-            # Update execution state buffers
-            if len(exec_results["states"]) > 0:
-                generated_states.extend(exec_results["states"])
-                generated_obs.extend(exec_results["observations"])
-                generated_actions.extend(exec_results["actions"])
-                generated_success = generated_success or exec_results["success"]
+            # Update success state
+            generated_success = generated_success or exec_success
 
             # Get the navigation state
             if self.env_cfg.datagen_config.use_navigation_controller:
@@ -983,10 +990,6 @@ async def generate(  # noqa: C901
             if all(eef_subtasks_done.values()):
                 break
 
-        # Merge numpy arrays
-        if len(generated_actions) > 0:
-            generated_actions = torch.cat(generated_actions, dim=0)
-
         # Set success to the recorded episode data and export to file
         self.env.recorder_manager.set_success_to_episodes(
             env_id_tensor, torch.tensor([[generated_success]], dtype=torch.bool, device=self.env.device)
@@ -996,9 +999,6 @@ async def generate(  # noqa: C901
 
         results = dict(
             initial_state=new_initial_state,
-            states=generated_states,
-            observations=generated_obs,
-            actions=generated_actions,
             success=generated_success,
         )
         return results