Separate recorder capture from batched episode storage by ooctipus · Pull Request #8822 · isaac-sim/IsaacLab

ooctipus · 2026-10-26T09:54:27Z

Separate capture from episode storage in the existing RecorderManager. Terms can stream batches to consumers without retaining episodes; retained data uses preallocated batched tensors with independent world indices.

Capture, retention and export are independent: returning (None, None) skips retention; EXPORT_NONE only disables file export. Terms choose the captured fields, while consumers own subsequent processing.

What makes insertion faster:

Replace each world's Python episode dictionary/lists with one preallocated tensor per field, indexed by time and world.
Use one indexed batch write per field, removing per-world tensor views, list appends and the separate batch clone.
Keep per-world write counters on the GPU; a host capacity bound avoids reading them back on ordinary appends.
Reuse buffers after reset, grow capacity geometrically when needed, and build individual episode snapshots only on read/export.

This removes per-world Python bookkeeping from the frequent append path. Snapshot/export work still occurs when requested.

Storage speed — RTX 5090:

256 worlds: 27× faster, 2.362 → 0.0883 ms per batch append.
4,096 worlds: 697× faster, 61.585 → 0.0884 ms.

These measure steady-state buffer insertion, excluding allocation growth, reset, export and simulation. Sampled stored episodes match byte-for-byte.

Breaking API: get_episode() returns an owned snapshot. Collect through add_to_episodes(); EpisodeData.add() / pre_export() and the RecorderTerm.close() hook are removed. RL post-step callbacks capture before reset; observation recording explicitly requests an observation group. Migration guidance is included.

Validation: 10 recorder/lifecycle tests, CUDA initial-state capture, formatting/changelog checks, and the warnings-as-errors documentation build pass. Coverage includes reordered/empty tuple IDs, preserved slices, snapshot ownership, export modes, and unchanged policy noise/history/RNG. Four mutation checks confirm the selector regression is caught. HDF5/replay checks passed during the rebase; full simulator integration remains unavailable locally without Isaac Sim.

Benchmark methodology and reproduction

Episode-buffer insertion is 697× faster at 4,096 environments in a fresh comparison of base 0ca25c56 and measured head c0931607. These are storage-only timings; the factor does not describe simulation or training throughput.

256 environments: 2.362 → 0.0883 ms per batch append (27× faster). Trial ranges: 2.294–2.871 ms before; 0.0865–0.0899 ms after.
4,096 environments: 61.585 → 0.0884 ms per batch append (697× faster). Trial ranges: 55.874–65.940 ms before; 0.0876–0.1011 ms after.

Measured on an NVIDIA GeForce RTX 5090 (driver 580.178.04, PyTorch 2.11.0+cu130, CUDA 13.0) on 2026-10-29. The benchmark uses identical seven-field payloads: four FP32 tensors with widths 37, 37, 19 and 72, and three Boolean fields. Payload size is 169,728 bytes at 256 worlds and 2,715,648 bytes at 4,096 worlds. Each result is the median of five randomized paired trials, with 20 appends per trial after warmup. Wall time is synchronized before and after each trial and includes Python and GPU work. The tested GPU had no competing compute jobs. The CPU host was shared, including a documentation build; other compute jobs used the second GPU.

The batched buffer starts with capacity 32 and reuses its allocation after reset. Timing covers steady-state insertion only: capture, physics, reset, allocation growth, episode materialization and disk export are outside this measurement. Sampled 20-frame episodes from the first, middle and last world match byte-for-byte at both batch sizes.

Where the gain comes from: the previous implementation clones each field's batch, then creates a tensor view and traverses/appends to a Python episode dictionary/list for every world. The new implementation writes all worlds into a preallocated tensor using per-world indices, once per field. Per-world episode materialization moves to get_episode()/export, and reset reuses storage. Streaming consumers can also skip episode retention, and physical-state recorder terms can avoid an unnecessary observation pass.

Use the normal Isaac Lab GPU environment with this PR's head checked out. Save the Python script below as /tmp/benchmark_recorder.py, then run:

git show 0ca25c56eec6fe7860fdcba44f8848f19693d6f6:source/isaaclab/isaaclab/managers/recorder_manager.py > /tmp/recorder_before.py
git show 0ca25c56eec6fe7860fdcba44f8848f19693d6f6:source/isaaclab/isaaclab/utils/datasets/episode_data.py > /tmp/episode_before.py
uv run python /tmp/benchmark_recorder.py --reference /tmp/recorder_before.py --reference-episode /tmp/episode_before.py --output /tmp/recorder_performance.json --device cuda:0

The script loads the parent recorder and its original EpisodeData alongside the checked-out recorder, uses the same dependencies and payloads for both, synchronizes CUDA, randomizes trial order, checks stored bytes and writes every timing sample. There is no simulation in this storage microbenchmark.

"""Compare PR #8822's exact parent and head on identical episode payloads."""
import argparse
import gc
import hashlib
import importlib.util
import json
import random
import statistics
import sys
import time
from pathlib import Path
from types import SimpleNamespace

import torch
from isaaclab.managers import RecorderTermCfg
from isaaclab.managers import recorder_manager as current

parser = argparse.ArgumentParser()
parser.add_argument('--reference', type=Path, required=True)
parser.add_argument('--reference-episode', type=Path, required=True)
parser.add_argument('--output', type=Path, required=True)
parser.add_argument('--device', default='cuda:0')
args = parser.parse_args()
torch.cuda.set_device(args.device)
spec = importlib.util.spec_from_file_location('isaaclab.managers._recorder_baseline', args.reference)
baseline = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = baseline
spec.loader.exec_module(baseline)
episode_spec = importlib.util.spec_from_file_location('_recorder_baseline_episode', args.reference_episode)
episode_module = importlib.util.module_from_spec(episode_spec)
episode_spec.loader.exec_module(episode_module)
baseline.EpisodeData = episode_module.EpisodeData

def timed(fn, count):
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(count):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - start) * 1000 / count

@torch.no_grad()
def main():
    report = dict(gpu=torch.cuda.get_device_name(), device=args.device, torch=torch.__version__,
                  cuda=torch.version.cuda, reference_sha256=hashlib.sha256(args.reference.read_bytes()).hexdigest(),
                  reference_episode_sha256=hashlib.sha256(args.reference_episode.read_bytes()).hexdigest(),
                  current_sha256=hashlib.sha256(Path(current.__file__).read_bytes()).hexdigest(),
                  current_file=current.__file__, repeats=5, steps=20, cases=[])
    rng = random.Random(42)
    torch.manual_seed(42)
    for worlds in (256, 4096):
        env = SimpleNamespace(num_envs=worlds, device=args.device,
                              sim=SimpleNamespace(is_playing=lambda: True), cfg=SimpleNamespace())
        payload = {name: torch.randn(worlds, width, device=args.device)
                   for name, width in [('state', 37), ('next_state', 37), ('control', 19), ('contacts', 72)]}
        payload.update({name: torch.zeros(worlds, dtype=torch.bool, device=args.device)
                        for name in ('terminated', 'truncated', 'valid')})
        managers = {}
        for name, module in [('baseline', baseline), ('batched', current)]:
            cfg = module.RecorderManagerBaseCfg(dataset_export_mode=module.DatasetExportMode.EXPORT_NONE)
            if name == 'batched':
                cfg.buffer_initial_capacity = 32
            cfg.term = RecorderTermCfg(class_type=module.RecorderTerm)
            manager = managers[name] = module.RecorderManager(cfg, env)
            for _ in range(20):
                manager.add_to_episodes('transition', payload)
            manager.reset()
        values = {name: [] for name in managers}
        orders = []
        for _ in range(5):
            order = list(managers)
            rng.shuffle(order)
            orders.append(order)
            for name in order:
                manager = managers[name]
                manager.reset()
                gc.collect()
                values[name].append(timed(lambda: manager.add_to_episodes('transition', payload), 20))
        for world in (0, worlds // 2, worlds - 1):
            old = managers['baseline'].get_episode(world).data['transition']
            new = managers['batched'].get_episode(world).data['transition']
            assert old.keys() == new.keys()
            for key in old:
                expected = torch.stack(old[key])
                assert expected.shape == new[key].shape and expected.dtype == new[key].dtype
                assert torch.equal(expected.contiguous().view(torch.uint8), new[key].contiguous().view(torch.uint8))
        case = dict(worlds=worlds, fields=len(payload), payload_bytes=sum(t.numel()*t.element_size() for t in payload.values()),
                    append_ms={name: dict(median=statistics.median(v), min=min(v), max=max(v), samples=v)
                               for name, v in values.items()}, order=orders, sampled_episodes_byte_equal=True)
        report['cases'].append(case)
        print(json.dumps(case), flush=True)
        for manager in managers.values():
            manager.close()
    args.output.write_text(json.dumps(report, indent=2) + '\n')

main()

Contribution checklist

Read the contribution guidelines
Ran repository formatting and changelog checks
Updated documentation and migration guidance
Validated changed behavior and applied the test-audit criteria
Added package changelog fragments
Contributor already listed in CONTRIBUTORS.md
No runtime warnings (existing dependency deprecations remain)

ooctipus · 2026-10-26T21:58:24Z

run-ci

StafaH · 2026-10-27T00:52:58Z

-        """Record data at the end of env.step() when all the managers are processed.
+        """Record the completed physical transition.
+
+        In RL environments, this runs after rewards and terminations, before automatic reset.


too verbose, it's leaking alot of implementation details

StafaH · 2026-10-27T00:53:33Z

            The number of successful episodes.
        """
-        if not hasattr(self, "_exported_successful_episode_count"):
+        if not hasattr(self, "_exported_counts"):


let's avoid hasattr

StafaH · 2026-10-27T02:22:53Z

+        if cfg.buffer_initial_capacity < 1:
+            raise ValueError("buffer_initial_capacity must be positive.")
+        self._fields: list[_RecorderField] = []
+        self._env_ids = torch.arange(env.num_envs, device=env.device)


I notice alot of torch, should we migrate to warp?

i'd like to keep this pr small and left that to recorder rewrite

StafaH · 2026-10-27T02:24:53Z

-        if isinstance(env_ids, torch.Tensor):
-            env_ids = env_ids.tolist()
+            env_ids = self._env_ids
+        elif isinstance(env_ids, slice):


In previous PR we had fixed the id vs slice by allowing a union of the types and then letting torch handle it since torch can handle id and slice, so we don't need to convert, can we use the same here?

StafaH · 2026-10-27T02:26:26Z

    recorder_manager.close()
+
+
+def test_batched_episode_indices_and_ownership(monkeypatch):


tests can be cleaned up and minimized, don't think we need these many tests

hujc7 · 2026-10-27T06:21:02Z

    Operations.
    """

    def reset(self, env_ids: Sequence[int] | slice | None = None) -> dict[str, torch.Tensor]:


can we also add mask support?

this recorder will be rewriten at some point, I would just leave this pr speed fix minimum

hujc7 · 2026-10-27T06:38:02Z

-            # update observations for recording if needed
-            self.obs_buf = self.observation_manager.compute()
-            self.recorder_manager.record_post_step()
+        self.recorder_manager.record_post_step()


this might be easy to miss. the assumption for pre and post seems different that pre obs is fresh but post is stale. Might prefer to have the behavior aligned to avoid problem. Or a cfg level switch to opt in/out

same as above, i think recoder manager put a bit too many hooks into managers, a better appoach would be instead of we expose the hook, we just let user compose the step function, and they can add the record in wherever they think its appropriate. this pr will just mostly be speed up

ooctipus · 2026-10-28T06:26:54Z

run-ci

greptile-apps · 2026-10-28T06:31:04Z

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge under its documented breaking API changes.

Summary

Separates recorder capture from retained episode storage.

Recorder fields use reusable batched tensors, with snapshots made only when requested.
RL post-step capture runs before reset without an extra observation pass.
EpisodeData holds tensor episodes instead of collecting each sample.

No actionable issues were established. Review checks were static; no tests or benchmarks were run.

Acknowledged intentional changes: ooctipus explicitly identifies the removed collection methods and mutable episode access as breaking changes and supplies migration guidance. ooctipus also explicitly identifies removal of RecorderTerm.close() as intentional and assigns resource cleanup to applications.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Recorder callback captures batch] --> B{Return field and values?}
    B -->|No| C[Consumer copies published batch]
    B -->|Yes| D[Append to batched field storage]
    D --> E[Reset selected world indices]
    D --> F[get_episode copies owned snapshot]
    D --> G{File export enabled?}
    G -->|Yes| H[Build snapshots and write dataset]
    G -->|No| I[Finalize selected episodes without file writes]

_{Reviews (1) · Last reviewed commit: "Merge upstream dataset regression with c..." · Reviewed by Greptile}

isaaclab-review-bot

Isaac Lab Review Bot

Needs attention: Scope and over-engineering (1), Compatibility and deprecation (3), Style consistency (2), Test quality (1). Other reviewed categories look good.

PR description: The benchmark targets a9fb50e/a2ba7a126, not this PR’s current revisions; rerun the documented workload against 1ab7bc6/6d112c52b and remove stale integration evidence.
Design and architecture: Architecture looks good.
Scope and over-engineering: Minimal implementation: extend RecorderManager, EpisodeData, and both RL step owners; no source modules were added, and unrelated HDF5 rewrites should be reverted; Suggested PR breakdown: 1) batched storage with compatibility bridges and recorder validation; 2) post-step lifecycle migration across both environments with focused observation tests.
API: API changes look good.
Compatibility and deprecation: Breaking changes: get_episode semantics, EpisodeData.add/pre_export removal, and post-step observation timing lack functional deprecation bridges, targeted warnings, removal guidance, and transition coverage.
Implementation: Implementation looks good.
Style consistency: New recorder methods omit required type annotations and Google-style parameter or return sections while replacing established environment terminology with “world”; restore the existing typed documentation pattern.
Test quality: The AST-based storage-boundary test asserts imports, absent attributes, and private loop structure rather than observable behavior; delete it and retain the recorder behavior and performance-boundary tests.

Significant concerns. Posted 7 inline findings.

Automated review; human maintainers own approval decisions.

isaaclab-review-bot · 2026-10-28T07:13:23Z


-        Args:
-            env_id: The environment id.
+    def get_episode(self, env_id: int) -> EpisodeData:


🔴 Critical · Compatibility — get_episode contract changed without bridge

Callers now get a detached snapshot with tensor leaves, so mutations are silently lost and torch.stack(episode.data[key]) fails. Keep the live contract during deprecation and expose snapshots under a new name.

isaaclab-review-bot · 2026-10-28T07:13:23Z

        """Check if the episode data is empty."""
        return not self._data

-    def add(self, key: str, value: torch.Tensor | dict, clone: bool = True):


🔴 Critical · Compatibility — EpisodeData collection API removed without deprecation

add() and pre_export() are public API; removal breaks external episode builders immediately. Keep both delegating to the new storage with a DeprecationWarning and removal guidance for the prescribed window.

isaaclab-review-bot · 2026-10-28T07:13:23Z

-            # update observations for recording if needed
-            self.obs_buf = self.observation_manager.compute()
-            self.recorder_manager.record_post_step()
+        self.recorder_manager.record_post_step()


🟡 Warning · Compatibility — Post-step recorders now see stale observations

Removing the refresh means existing post-step terms reading env.obs_buf silently record the previous step's observations. Gate the new timing behind an opt-in or deprecation cycle covering both env packages.

isaaclab-review-bot · 2026-10-28T07:13:23Z

-        """
-        # Do nothing if no active recorder terms are provided
-        if len(self.active_terms) == 0:
+        """Set selected worlds' success flags from a tensor of shape (selected worlds,) or (selected worlds, 1)."""


🟡 Warning · Style Consistency — Public docstrings drop Args and Returns

set_success_to_episodes and get_episode fold parameters and returns into prose and say "world" where this file and API use environment and env_ids. Restore Google-style sections and existing vocabulary.

isaaclab-review-bot · 2026-10-28T07:13:23Z

    Helper functions.
    """

+    def _reset_buffers(self, env_ids):


🟡 Warning · Style Consistency — New helpers omit type annotations

_reset_buffers, _RecorderField.append and _RecorderField.reset leave env_ids unannotated and omit -> None, unlike every sibling recorder method declaring Sequence[int] | slice | None.

isaaclab-review-bot · 2026-10-28T07:13:23Z

+            assert not bool(file["data/demo_10"].attrs["success"])
+
+
+def test_recorder_storage_boundary():


🟡 Warning · Test Quality — Source-grep test asserts module structure

test_recorder_storage_boundary parses recorder_manager source to assert absent attributes, allowed imports and no loops in private _RecorderField.append. It guards no observable contract and fails on behavior-preserving refactors; delete it.

isaaclab-review-bot · 2026-10-28T07:13:23Z

-
-    # check if the dataset is created
-    assert os.path.exists(os.path.join(temp_dir, dataset_file_name))
+def test_create_dataset_file(tmp_path, monkeypatch):


🟡 Warning · Scope Complexity — Unrelated HDF5 test rewrite

Only create_test_episode needed updating for the removed add()/pre_export(); the tmp_path migration, loop rewrite and merge of test_create_resets_env_args_when_reusing_handler are drive-by churn that enlarge review surface. Revert them.

maxkra15 · 2026-10-29T08:18:31Z

i reproduced a selector regression: reset((0, 1)) and set_success_to_episodes((0, 1), ...) work on the base but now raise IndexError. indexing _success directly with a tuple treats it as multiple dimensions. could we normalize sequence selectors before indexing?

get_episode() also silently stops supporting live edits because it now returns a snapshot. could we add a separate snapshot method and deprecate the old behavior first?

similarly, the removed EpisodeData.add() / pre_export() methods and RecorderTerm.close() hook need a deprecation period. custom close hooks are no longer called.

ooctipus · 2026-10-30T00:02:46Z

i reproduced a selector regression: reset((0, 1)) and set_success_to_episodes((0, 1), ...) work on the base but now raise IndexError. indexing _success directly with a tuple treats it as multiple dimensions. could we normalize sequence selectors before indexing?

get_episode() also silently stops supporting live edits because it now returns a snapshot. could we add a separate snapshot method and deprecate the old behavior first?

similarly, the removed EpisodeData.add() / pre_export() methods and RecorderTerm.close() hook need a deprecation period. custom close hooks are no longer called.

addressed the regression

github-actions Bot added documentation Improvements or additions to documentation isaac-mimic Related to Isaac Mimic team labels Oct 26, 2026

isaaclab-bot Bot added ci:run-docker Trigger the on-demand Docker and GPU CI workflow and removed ci:run-docker Trigger the on-demand Docker and GPU CI workflow labels Oct 26, 2026

StafaH reviewed Oct 27, 2026

View reviewed changes

hujc7 reviewed Oct 27, 2026

View reviewed changes

ooctipus marked this pull request as ready for review October 28, 2026 06:26

ooctipus requested a review from a team October 28, 2026 06:26

ooctipus requested a review from AntoineRichard as a code owner October 28, 2026 06:26

ooctipus requested review from fatimaanes, kellyguo11, pbarejko, peterd-NV and rwiltz as code owners October 28, 2026 06:26

isaaclab-bot Bot added ci:run-docker Trigger the on-demand Docker and GPU CI workflow and removed ci:run-docker Trigger the on-demand Docker and GPU CI workflow labels Oct 28, 2026

isaaclab-review-bot Bot reviewed Oct 28, 2026

View reviewed changes

ooctipus added 5 commits October 28, 2026 22:38

Separate recorder capture from batched episode storage

32a5d44

Document learner-independent recorder integration boundaries

b27227c

Keep recorder documentation and dependency checks generic

09645f3

Make recorder observation capture explicit

4b0d0b1

Remove obsolete recorder accumulation and finalization APIs

81c0eff

ooctipus added 3 commits October 28, 2026 22:38

Consolidate recorder and dataset tests at their owning boundaries

87bc5b2

Simplify recorder documentation and preserve reset selectors

11b8c77

Declare relocated replay fixture as a test-only change

70d4e53

ooctipus force-pushed the refactor/recorder-capture-storage branch from 6d112c5 to 70d4e53 Compare October 29, 2026 05:41

ooctipus requested a review from hougantc-nvda as a code owner October 29, 2026 05:41

Preserve sequence selectors in recorder lifecycle operations

c093160

		recorder_manager.close()


		def test_batched_episode_indices_and_ownership(monkeypatch):

		assert not bool(file["data/demo_10"].attrs["success"])


		def test_recorder_storage_boundary():

Conversation

ooctipus commented Oct 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ooctipus commented Oct 26, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ooctipus commented Oct 28, 2026

Uh oh!

greptile-apps Bot commented Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

isaaclab-review-bot Bot left a comment

Choose a reason for hiding this comment

Isaac Lab Review Bot

Uh oh!

isaaclab-review-bot Bot Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

isaaclab-review-bot Bot Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

isaaclab-review-bot Bot Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

isaaclab-review-bot Bot Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

isaaclab-review-bot Bot Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

isaaclab-review-bot Bot Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

isaaclab-review-bot Bot Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

maxkra15 commented Oct 29, 2026

Uh oh!

ooctipus commented Oct 30, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ooctipus commented Oct 26, 2026 •

edited

Loading

greptile-apps Bot commented Oct 28, 2026 •

edited

Loading