Repository navigation
Conversation
|
run-ci |
| """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. |
There was a problem hiding this comment.
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"): |
| 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) |
There was a problem hiding this comment.
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): |
There was a problem hiding this comment.
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): |
There was a problem hiding this comment.
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]: |
There was a problem hiding this comment.
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() |
There was a problem hiding this comment.
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
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_NONEonly 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:
This removes per-world Python bookkeeping from the frequent append path. Snapshot/export work still occurs when requested.
Storage speed — RTX 5090:
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 throughadd_to_episodes();EpisodeData.add()/pre_export()and the unusedRecorderTerm.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
a9fb50e53and measured heada2ba7a126. These are storage-only timings; the factor does not describe simulation or training throughput.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:The script loads the parent recorder and its original
EpisodeDataalongside 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.Contribution checklist