Skip to content

Separate recorder capture from batched episode storage - #8822

Draft
ooctipus wants to merge 7 commits into
isaac-sim:mainfrom
ooctipus:refactor/recorder-capture-storage
Draft

ooctipus wants to merge 7 commits into
isaac-sim:mainfrom
ooctipus:refactor/recorder-capture-storage

Conversation

@ooctipus

@ooctipus ooctipus commented Oct 26, 2026 •

Copy link
Copy Markdown
Collaborator

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. Removing the old accumulator and unused hooks makes the production Python change 32 lines smaller than the base.

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.423 → 0.0892 ms per batch append.
  • 4,096 worlds: 651× faster, 57.730 → 0.0887 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 unused RecorderTerm.close() hook are removed. RL post-step callbacks capture before reset; observation recording explicitly requests an observation group. Migration guidance is included.

Validation: 16 focused recorder/dataset/replay/lifecycle tests pass after consolidating duplicate coverage; 12 deliberate mutations confirm the remaining tests catch the removed tests' meaningful regressions. Formatting/changelog checks and the warnings-as-errors documentation build passed. Coverage includes pre-reset capture, unchanged policy noise/history/RNG, and the recorder's dependency boundary.

Benchmark methodology and reproduction

Episode-buffer insertion is about 651× faster at 4,096 worlds in a fresh comparison of this PR's exact parent a9fb50e53 and measured head a2ba7a126. These are storage-only timings; the factor does not describe simulation or training throughput.

  • 256 worlds: 2.423 → 0.0892 ms per batch append (27× faster). Trial ranges: 2.272–2.915 ms before; 0.0884–0.0938 ms after.
  • 4,096 worlds: 57.730 → 0.0887 ms per batch append (651× faster). Trial ranges: 55.146–59.798 ms before; 0.0881–0.0900 ms after.

Measured on an NVIDIA GeForce RTX 5090 (driver 580.178.04, PyTorch 2.11.0+cu130, CUDA 13.0) on 2026-10-26. 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 job; other work used the shared host and its other 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.

Earlier graph-enabled integration check: the 4,096-world native G1 recheck used the recorder source (SHA-256 d92bc5ee04e8655fc28ac364813497a67b1641c11863dbf5c292a180971e9c6a). Isolated capture took 0.089 ms, and inserting its ten-field, 3,891,200-byte payload took 0.117 ms. Newton/MuJoCo Warp executed four physics substeps in one CUDA graph launch per control step. That recheck ran on a shared CPU host during a documentation build; these component timings are not an isolated end-to-end throughput claim. The G1 consumer is supplied by #8826, and the original integration evidence remains in #8806.

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 a9fb50e5388724a680bbc5c1f836917271605a99:source/isaaclab/isaaclab/managers/recorder_manager.py > /tmp/recorder_before.py
git show a9fb50e5388724a680bbc5c1f836917271605a99: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; CUDA-graph behavior is covered separately by the G1 integration check above.

"""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

Copy link
Copy Markdown
Collaborator Author

run-ci

@isaaclab-bot 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
"""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.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

too verbose, it's leaking alot of implementation details

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's avoid hasattr

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

recorder_manager.close()


def test_batched_episode_indices_and_ownership(monkeypatch):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Operations.
"""

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we also add mask support?

# 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()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation isaac-mimic Related to Isaac Mimic team

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants