Merge remote-tracking branch 'upstream/develop' into esekkin/camera-refactor

Author: nvsekkin

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/docs/CHANGELOG.rst

@@ -34,6 +34,20 @@ Removed
   now managed internally by the renderer backend and are not part of the public API.
 
 
+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

source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py

@@ -78,6 +78,7 @@ def env_loop(
     env_action_queue: asyncio.Queue,
     shared_datagen_info_pool: DataGenInfoPool,
     asyncio_event_loop: asyncio.AbstractEventLoop,
+    data_gen_tasks: asyncio.Future | None = None,
 ):
     """Main asyncio loop for the environment.
 
@@ -87,6 +88,8 @@ def env_loop(
         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.
+        data_gen_tasks: The gathered async data generation future. When provided, the loop
+            will exit early if all tasks finish unexpectedly (e.g. due to an unhandled exception).
     """
     global num_success, num_failures, num_attempts
     env_id_tensor = torch.tensor([0], dtype=torch.int64, device=env.device)
@@ -97,17 +100,20 @@ def env_loop(
             # 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))
+                if data_gen_tasks is not None and data_gen_tasks.done():
+                    exc = data_gen_tasks.exception()
+                    if exc is not None:
+                        raise exc
+                    return
                 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.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())
+            get_tasks = [env_action_queue.get() for _ in range(env.num_envs)]
+            results = asyncio_event_loop.run_until_complete(asyncio.gather(*get_tasks))
+            for env_id, action in results:
                 actions[env_id] = action
 
             # perform action on environment
@@ -152,6 +158,7 @@ def setup_env_config(
     device: str,
     generation_num_trials: int | None = None,
     recorder_cfg: RecorderManagerBaseCfg | None = None,
+    dataset_compression: bool = True,
 ) -> tuple[Any, Any]:
     """Configure the environment for data generation.
 
@@ -162,6 +169,8 @@ def setup_env_config(
         num_envs: Number of environments to run
         device: Device to run on
         generation_num_trials: Optional override for number of trials
+        recorder_cfg: Recorder manager configuration
+        dataset_compression: Whether to enable dataset compression
 
     Returns:
         tuple containing:
@@ -198,6 +207,8 @@ def setup_env_config(
     env_cfg.recorders.dataset_export_dir_path = output_dir
     env_cfg.recorders.dataset_filename = output_file_name
 
+    env_cfg.recorders.dataset_compression = dataset_compression
+
     if env_cfg.datagen_config.generation_keep_failed:
         env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES
     else:
@@ -230,8 +241,7 @@ def setup_async_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, env.cfg, env.device, asyncio_lock=shared_datagen_info_pool_lock)
+    shared_datagen_info_pool = DataGenInfoPool(env, env.cfg, env.device)
     shared_datagen_info_pool.load_from_dataset_file(input_file)
     print(f"Loaded {shared_datagen_info_pool.num_datagen_infos} to datagen info pool")