Improves recorder manager to support customized demo indices. (#3552)

# 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.
-->

Current manager-based workflow leveraged recorder_manager to store demo
obs and actions to hdf5 file.
- But it only outputs successful demos and accumulates the demo index
- Not supporting customized demo index (the use case is like: replay an
existing hdf5, and store successful demos while keeping its original
demo index)

The modification in this PR:
- keeps original function, but extends to store demo with specific demo
index

<!-- 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. -->

## Type of change

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

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


## Screenshots

| Before | After |
| ------ | ----- |
|<img width="100" height="150" alt="Screenshot from 2025-09-25 10-20-31"
src="https://github.com/user-attachments/assets/b4af24df-2781-4ba2-8693-fd246875012b"
/>|<img width="100" height="150" alt="Screenshot from 2025-09-25
10-20-47"
src="https://github.com/user-attachments/assets/e86d3210-e205-4d6b-b83e-cf69a585743b"
/> |


## 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
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] 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: ooctipus <[email protected]>

Author: rebeccazhang0707

source/isaaclab/config/extension.toml

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

source/isaaclab/docs/CHANGELOG.rst

@@ -1,6 +1,23 @@
 Changelog
 ---------
 
+0.47.4 (2025-10-30)
+~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Enhanced :meth:`~isaaclab.managers.RecorderManager.export_episodes` method to support customizable sequence of demo IDs:
+
+  - Added argument ``demo_ids`` to :meth:`~isaaclab.managers.RecorderManager.export_episodes` to accept a sequence of integers
+    for custom episode identifiers.
+
+* Enhanced :meth:`~isaaclab.utils.datasets.HDF5DatasetFileHandler.write_episode` method to support customizable episode identifiers:
+
+  - Added argument ``demo_id`` to :meth:`~isaaclab.utils.datasets.HDF5DatasetFileHandler.write_episode` to accept a custom integer
+    for episode identifier.
+
+
 0.47.3 (2025-10-22)
 ~~~~~~~~~~~~~~~~~~~
 

source/isaaclab/isaaclab/managers/recorder_manager.py

@@ -442,12 +442,16 @@ def get_ep_meta(self) -> dict:
         ep_meta = self._env.cfg.get_ep_meta()
         return ep_meta
 
-    def export_episodes(self, env_ids: Sequence[int] | None = None) -> None:
+    def export_episodes(self, env_ids: Sequence[int] | None = None, demo_ids: Sequence[int] | None = None) -> None:
         """Concludes and exports the episodes for the given environment ids.
 
         Args:
             env_ids: The environment ids. Defaults to None, in which case
                 all environments are considered.
+            demo_ids: Custom identifiers for the exported episodes.
+                If provided, episodes will be named "demo_{demo_id}" in the dataset.
+                Should have the same length as env_ids if both are provided.
+                If None, uses the default sequential naming scheme. Defaults to None.
         """
         # Do nothing if no active recorder terms are provided
         if len(self.active_terms) == 0:
@@ -458,6 +462,17 @@ def export_episodes(self, env_ids: Sequence[int] | None = None) -> None:
         if isinstance(env_ids, torch.Tensor):
             env_ids = env_ids.tolist()
 
+        # Handle demo_ids processing
+        if demo_ids is not None:
+            if isinstance(demo_ids, torch.Tensor):
+                demo_ids = demo_ids.tolist()
+            if len(demo_ids) != len(env_ids):
+                raise ValueError(f"Length of demo_ids ({len(demo_ids)}) must match length of env_ids ({len(env_ids)})")
+            # Check for duplicate demo_ids
+            if len(set(demo_ids)) != len(demo_ids):
+                duplicates = [x for i, x in enumerate(demo_ids) if demo_ids.index(x) != i]
+                raise ValueError(f"demo_ids must be unique. Found duplicates: {list(set(duplicates))}")
+
         # Export episode data through dataset exporter
         need_to_flush = False
 
@@ -468,7 +483,7 @@ def export_episodes(self, env_ids: Sequence[int] | None = None) -> None:
             if self._failed_episode_dataset_file_handler is not None:
                 self._failed_episode_dataset_file_handler.add_env_args(ep_meta)
 
-        for env_id in env_ids:
+        for i, env_id in enumerate(env_ids):
             if env_id in self._episodes and not self._episodes[env_id].is_empty():
                 self._episodes[env_id].pre_export()
 
@@ -484,7 +499,9 @@ def export_episodes(self, env_ids: Sequence[int] | None = None) -> None:
                     else:
                         target_dataset_file_handler = self._failed_episode_dataset_file_handler
                 if target_dataset_file_handler is not None:
-                    target_dataset_file_handler.write_episode(self._episodes[env_id])
+                    # 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)
                     need_to_flush = True
                 # Update episode count
                 if episode_succeeded:

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

@@ -136,18 +136,27 @@ def load_dataset_helper(group):
 
         return episode
 
-    def write_episode(self, episode: EpisodeData):
+    def write_episode(self, episode: EpisodeData, demo_id: int | None = None):
         """Add an episode to the dataset.
 
         Args:
             episode: The episode data to add.
+            demo_id: Custom index for the episode. If None, uses default index.
         """
         self._raise_if_not_initialized()
         if episode.is_empty():
             return
 
-        # create episode group based on demo count
-        h5_episode_group = self._hdf5_data_group.create_group(f"demo_{self._demo_count}")
+        # Use custom demo id if provided, otherwise use default naming
+        if demo_id is not None:
+            episode_group_name = f"demo_{demo_id}"
+        else:
+            episode_group_name = f"demo_{self._demo_count}"
+
+        # create episode group with the specified name
+        if episode_group_name in self._hdf5_data_group:
+            raise ValueError(f"Episode group '{episode_group_name}' already exists in the dataset")
+        h5_episode_group = self._hdf5_data_group.create_group(episode_group_name)
 
         # store number of steps taken
         if "actions" in episode.data:
@@ -176,8 +185,10 @@ def create_dataset_helper(group, key, value):
         # increment total step counts
         self._hdf5_data_group.attrs["total"] += h5_episode_group.attrs["num_samples"]
 
-        # increment total demo counts
-        self._demo_count += 1
+        # Only increment demo count if using default indexing
+        if demo_id is None:
+            # increment total demo counts
+            self._demo_count += 1
 
     def flush(self):
         """Flush the episode data to disk."""