Unverified Commit 6f59b88f authored by rebeccazhang0707's avatar rebeccazhang0707 Committed by GitHub

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: 's avatarooctipus <zhengyuz@nvidia.com>
parent 6acdd82a
[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"
......
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)
~~~~~~~~~~~~~~~~~~~
......
......@@ -442,12 +442,16 @@ class RecorderManager(ManagerBase):
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 @@ class RecorderManager(ManagerBase):
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 @@ class RecorderManager(ManagerBase):
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 @@ class RecorderManager(ManagerBase):
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:
......
......@@ -136,18 +136,27 @@ class HDF5DatasetFileHandler(DatasetFileHandlerBase):
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,6 +185,8 @@ class HDF5DatasetFileHandler(DatasetFileHandlerBase):
# increment total step counts
self._hdf5_data_group.attrs["total"] += h5_episode_group.attrs["num_samples"]
# Only increment demo count if using default indexing
if demo_id is None:
# increment total demo counts
self._demo_count += 1
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment