Adds optimizations for Isaac Lab Mimic data generation performance#5141
Adds optimizations for Isaac Lab Mimic data generation performance#5141peterd-NV merged 33 commits intoisaac-sim:developfrom
Conversation
![isaaclab-review-bot[bot]](https://avatars.githubusercontent.com/u/157846462?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 Isaac Lab Review Bot
Summary
This PR optimizes Mimic data generation by: (1) removing redundant per-step accumulation of states/observations/actions in memory, (2) reducing tensor cloning overhead in recorder_manager/episode_data, (3) switching to batched asyncio.gather for action collection, (4) removing the shared asyncio lock for the standard generation flow, (5) switching visuomotor cameras from Camera to TiledCamera, (6) adding optional HDF5 dataset compression control, and (7) lowering gzip compression level. The core memory optimization (eliminating generated_states/obs/actions lists) is sound since this data is already captured by the recorder manager.
Design Assessment
Design is sound. The core insight is correct: the generate() method was accumulating states, observations, and actions in lists that were never consumed (only success was used from the results dict). These were redundant with what the recorder manager already stores. Removing them eliminates the largest memory sink during long generation runs.
The lock removal is justified — in single-threaded asyncio with no await points inside the critical section, the lock was a no-op from a correctness standpoint. The _optional_lock pattern correctly supports both the lock-free generate_dataset.py path and the lock-requiring consolidated_demo.py path.
The Camera → TiledCamera switch is a standard optimization for multi-env scenarios in Isaac Lab.
One concern: The disable_dataset_compression parameter is threaded through 3 layers (CLI → config → recorder manager → file handler) as a runtime flag rather than being part of the file handler configuration. A cleaner design would put compression settings on the file handler or its config, but the current approach is acceptable short-term given the small scope.
Architecture Impact
DatasetFileHandlerBase.write_episodebase class signature is now further out of sync withHDF5DatasetFileHandler(pre-existing issue, worsened).EpisodeData.addhas a newcloneparameter — backwards compatible (defaultTrue).MultiWaypoint.executereturn type changed fromdict→bool— the only caller (data_generator.py) is updated accordingly.DataGenerator.generatereturn dict dropsstates,observations,actionskeys — the only consumer (generation.py:run_data_generator) only usesresults["success"].- gzip compression level changed from default (4) to 2 for ALL datasets, not just mimic — this is a behavioral change visible to any user of
HDF5DatasetFileHandler.
Implementation Verdict
Minor fixes needed — Correct approach, 3 issues to address.
Test Coverage
No new tests for the disable_dataset_compression feature. The existing test_hdf5_dataset_file_handler.py tests call write_episode(test_episode) without the new parameter, so they only cover the default (compression enabled) path. A simple test that writes with disable_dataset_compression=True and verifies the dataset has no compression would be valuable. That said, the existing isaaclab_mimic CI check passes, which covers the end-to-end generation flow.
CI Status
Most checks still in progress (just pushed ~5 min ago). Passed so far: pre-commit ✅, isaaclab_mimic ✅, isaaclab_assets ✅, isaaclab_ov ✅, isaaclab_teleop ✅, labeler ✅, Build Base Docker Image ✅. Remaining checks pending.
Findings
🟡 Warning: hdf5_dataset_file_handler.py:257 — Compression level silently reduced for all users
The compression_opts=2 change reduces gzip compression from the h5py default (level 4) to level 2 for ALL HDF5 dataset writes, not just when mimic is running. This is an undocumented behavioral change that produces larger files. If the intent is to speed up mimic generation, this should either be (a) tied to the disable_dataset_compression flag as a "fast compression" mode, or (b) documented as an intentional project-wide change. As-is, users who rely on the default compression level will get larger files without knowing why.
🟡 Warning: dataset_file_handler_base.py — Base class write_episode signature diverges further
DatasetFileHandlerBase.write_episode(self, episode) doesn't match HDF5DatasetFileHandler.write_episode(self, episode, demo_id, disable_dataset_compression). The demo_id mismatch was pre-existing, but adding disable_dataset_compression widens the gap. Any future DatasetFileHandler subclass won't know about these parameters.
🔵 Improvement: recorder_manager.py:58 — Use positive naming for the compression config flag
disable_dataset_compression leads to confusing double-negation patterns (if not disable_dataset_compression). A positive name like dataset_compression: bool = True reads more naturally: if self.cfg.dataset_compression.
🔵 Improvement: Missing CHANGELOG entries
This PR modifies public APIs in both isaaclab (new clone param on EpisodeData.add, new disable_dataset_compression on RecorderManagerBaseCfg, changed write_episode signature) and isaaclab_mimic (changed return type of MultiWaypoint.execute, changed generate() return dict). Per project conventions, both source/isaaclab/docs/CHANGELOG.rst and source/isaaclab_mimic/docs/CHANGELOG.rst should be updated.
🔵 Improvement: PR description is empty
The PR body uses the default template without filling in the description, motivation, or type of change. For a performance optimization PR, including before/after measurements (memory usage, throughput, wall-clock time) would help reviewers evaluate the tradeoff.
| export_in_close: bool = False | ||
| """Whether to export episodes in the close call.""" | ||
|
|
||
| disable_dataset_compression: bool = False |
There was a problem hiding this comment.
Negative-sense booleans create confusing double-negation patterns downstream (if not disable_dataset_compression). The other boolean fields in this class use positive naming (export_in_record_pre_reset, export_in_close). This should follow the same convention.
| disable_dataset_compression: bool = False | |
| dataset_compression: bool = True | |
| """Whether to compress data when writing to the dataset file.""" |
This requires updating the references in recorder_manager.py:export_episodes, generation.py:setup_env_config, generate_dataset.py, and hdf5_dataset_file_handler.py accordingly (passing not self.cfg.dataset_compression or inverting the logic).
peterd-NV
changed the title
There was a problem hiding this comment.
🤖 Isaac Lab Review Bot — Follow-up
New commits (a225857→2f829a28): Added 3 test files for dataset generation (nutpour GR1T2, pickplace GR1T2, visuomotor).
1 new issue found in the added tests.
There was a problem hiding this comment.
🤖 Isaac Lab Review Bot — Follow-up
New commits (3ae515f→a18c6ff4): Fixed TODO_FILL_IN_ENV_NAME placeholders in test files with actual Nucleus paths.
1 new issue found in the updated paths.
peterd-NV
changed the title
peterd-NV
requested review from
Mayankm96,
hhansen-bdai,
jtigue-bdai,
ooctipus and
xyao-nv
as code owners
Greptile SummaryThis PR adds performance optimizations (batch-level tensor cloning, optional dataset compression, improved async execution) and additional test coverage (Franka visuomotor, GR1T2 pick-place, GR1T2 nut-pour) for the Mimic data generation pipeline. One P1 bug was found: line 98 of Confidence Score: 4/5Safe to merge after fixing the context manager composition bug on line 98 of generation.py One P1 logic bug found: source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py line 98 Important Files Changed
Sequence DiagramsequenceDiagram
participant M as main()
participant S as setup_env_config()
participant A as setup_async_generation()
participant L as env_loop()
participant DG as run_data_generator()
participant E as env
M->>S: env_name, compression, num_trials
S-->>M: env_cfg, success_term
M->>E: gym.make(env_name, cfg)
E-->>M: env instance
M->>A: env, input_file, success_term
A-->>M: tasks, event_loop, queues
M->>L: env, queues, data_gen_tasks
loop Until num_trials reached
DG->>A: put action in env_action_queue
L->>E: env.step(actions)
L->>DG: task_done (unblock join)
DG->>DG: check success_term
end
L-->>M: return
M->>M: cancel tasks
M->>E: env.close()
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: peterd-NV <peterd@nvidia.com>
kellyguo11
left a comment
kellyguo11
left a comment
There was a problem hiding this comment.
🤖 Isaac Lab Review Bot — Follow-up
New commits (ebf5a4a→36c7013c): Added h5py==3.15.1 pin to isaaclab_mimic and isaaclab_rl, new CHANGELOG entries.
2 new issues found in the added changes.
Co-authored-by: Kelly Guo <kellyg@nvidia.com> Signed-off-by: peterd-NV <peterd@nvidia.com>
3 participants
Description
This PR adds performance improvements and additional test case coverage for Isaac Lab Mimic data generation.
Performance improvements:
executeand return task success onlyenv_loopaddmethod of episode_data.py allowing for RecorderManager to clone contiguous tensors before callingaddAdditional test coverage:
Other changes:
env_loopto prevent event loop from spinning indefinitely when exception occurs in data generation.Type of change
Checklist
pre-commitchecks with./isaaclab.sh --formatconfig/extension.tomlfileCONTRIBUTORS.mdor my name already exists there