Fix cartpole-camera frame-stacking performance regression#5849
Fix cartpole-camera frame-stacking performance regression#5849jmart-nv wants to merge 11 commits into
Conversation
…isaac-sim#5574. - circular_buffer.py: replace torch.roll with O(1) write-index append; add CPU mirror of max_len and CPU bool gate to skip GPU sync on the steady-state path; rotate on read via the `buffer` property. - observations.py (stacked_image): drop trailing `.clone()` from the permute().reshape() chain (reshape on non-contig already allocates). - cartpole_camera_presets_env.py: same `.clone()` drop in the direct env. - delay_buffer.py: same `.clone()` drop on the DelayBuffer.compute return.
Skip per-frame /255 + mean-subtract on the stacking path; apply once on the stacked output. Ring buffer now holds native uint8 (4x cheaper per-step copies). Math is identical because K frames live in disjoint channel slices. Affects stacked_image MDP term and the cartpole DirectRLEnv parent (new `normalize: bool = True` kwarg) + subclass.
Replace the three-pass PyTorch normalize on the camera-observation hot path with a single fused Warp kernel + small mean reduction. Consolidate the per-data-type dispatch into a shared helper used by image(), stacked_image, CartpoleCameraEnv, and CartpoleCameraPresetsEnv. Frame-stacking subclasses pre-allocate the float32 output once and reuse it across steps, eliminating the per-step transient that the previous deferred-normalize path incurred.
The Warp kernel produces int32 partial sums along H; PyTorch collapses the partials and divides once. Channels-innermost launch shape gives coalesced reads on src's contiguous trailing dim. Partials scratch is cached per (shape, device).
- stacked_image / CartpoleCameraPresetsEnv: don't reuse a pre-allocated normalize-output buffer across env.step; the previous-iteration obs (held by the RL trainer) was overwritten before record_transition could read it, causing PPO to record degenerate transitions - normalize_image_uint8: docstring warning on the out= aliasing hazard - CircularBuffer.reset: setitem (= 0.0) instead of advanced-getitem + .zero_() so partial-batch resets actually zero the storage - CircularBuffer.stacked: drop unreachable empty-buffer guard - _uint8_sum_partials_cache: store the tensor directly, not a tuple - Tests: regressions for both bugs, image() clone kwarg, and DelayBuffer.compute() no-aliasing contract
![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 addresses a significant 42% throughput regression introduced by frame-stacking in #5574 for the cartpole-camera task. The fix implements several well-designed optimizations: (1) storing uint8 frames in the ring buffer instead of float32 (4× smaller), (2) deferring normalization past the frame-stack buffer, (3) replacing torch.roll with memmove-style shifting, (4) adding a fused Warp kernel for RGB normalize, and (5) providing a free contiguous stacked view via the new stack_dim parameter. The result is a 24% improvement over the pre-framestack baseline (67.87 FPS vs 54.62 FPS) with identical PPO convergence.
Incremental Update (a86c6a7)
The author addressed review feedback in commit a86c6a7:
✅ Docstring improvements — stack_dim parameter now includes a concrete example (stack_dim=-1 on (B, H, W, C) → (B, H, W, K*C)) making the API much more intuitive.
✅ In-place subtract in normalize fallback — images.py now uses images -= torch.mean(...) instead of allocating a new tensor, addressing the memory optimization suggestion.
✅ Cache growth documented — The _uint8_sum_partials_cache comment now states it "Typically holds one entry per training run (camera resolution and device are fixed)," making the bounded growth clear.
✅ Tile size rationale — _UINT8_SUM_TILE_HW comment now documents it was "Tuned on L40; robust across modern NVIDIA arches (Ampere/Ada/Hopper) at R256," addressing the hardcoded value concern.
✅ Shift loop documented — Added "Cheap at the typical frame-stack K=2-4" comment clarifying expected usage.
✅ Minor fix — Corrected TiledCameraCfg → CameraCfg in is_rgb_like docstring.
Test Coverage
✅ Excellent coverage (unchanged from previous push). The PR includes comprehensive tests for the circular buffer stack_dim mode, normalize dispatch paths, Warp kernel correctness, and critically a consecutive-call aliasing regression test.
CI Status
⏳ In progress — 9 checks passed, 6 running, 2 skipped on latest commit. No failures.
Verdict
Approve — All review suggestions have been addressed cleanly in the follow-up commit. The changes are documentation and minor optimization improvements that don't affect correctness. Ready to merge once CI completes.
…y-up) - images.py: TiledCameraCfg -> CameraCfg in is_rgb_like docstring (deprecated class) - images.py: PyTorch normalize fallback uses in-place subtract (one fewer alloc on the cold path) - circular_buffer.py: clarified stack_dim docstring with a (B,H,W,C) example; noted the typical K=2-4 range on the append shift loop - ops.py: noted the realistic 1-entry bound on _uint8_sum_partials_cache; noted L40 tuning provenance of _UINT8_SUM_TILE_HW=32
jmart-nv
requested review from
jtigue-bdai,
kellyguo11,
ooctipus and
pascal-roth
as code owners
Greptile SummaryThis PR fixes a 42% throughput regression and 2.7× VRAM inflation from frame-stacking (#5574), delivering 24% faster throughput than the pre-framestack baseline. The rewrite defers RGB normalization past the ring buffer (uint8 storage), replaces
Confidence Score: 3/5Safe at the stated R256 target; _uint8_spatial_mean produces silently wrong normalization for images wider than ~2902 pixels per side, which would corrupt training observations at larger camera resolutions without raising any error. The _uint8_spatial_mean function accumulates per-channel pixel sums in int32 and collapses them with partials.sum(dim=1).float(). When H × W × 255 exceeds INT32_MAX the total wraps around silently — PPO receives incorrect zero-mean observations and trains on corrupted data with no exception. The function is now callable from the public normalize_camera_image API. The fix is a single dtype= argument. All other logic — ring-buffer memmove, warmup-fill across partial resets, aliasing prevention via fresh allocation, and the clone=False path — checks out correctly. source/isaaclab/isaaclab/utils/warp/ops.py — the _uint8_spatial_mean return line needs dtype=torch.int64 promotion before this function is used at resolutions above R256. Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if self._stack_dim_internal is None: | ||
| raise RuntimeError("stacked is only available when CircularBuffer was created with stack_dim set.") |
There was a problem hiding this comment.
RuntimeError when stack_dim set but buffer uninitialized
When constructed with stack_dim set but append() never called, _stack_dim_internal is still None. Accessing .stacked raises RuntimeError("stacked is only available when CircularBuffer was created with stack_dim set.") — the opposite of the truth. Differentiating the two None cases avoids the misleading diagnostic.
| if self._stack_dim_internal is None: | |
| raise RuntimeError("stacked is only available when CircularBuffer was created with stack_dim set.") | |
| if self._stack_dim_internal is None: | |
| if self._stack_dim_arg is None: | |
| raise RuntimeError("stacked is only available when CircularBuffer was created with stack_dim set.") | |
| raise RuntimeError("stacked is not yet available: call append() at least once to initialize the buffer.") |
| self._cartpole.set_joint_effort_target_index(target=self.actions, joint_ids=self._cart_dof_idx) | ||
|
|
||
| def _get_observations(self) -> dict: | ||
| def _get_observations(self, clone: bool = True, normalize: bool = True) -> dict: |
There was a problem hiding this comment.
DirectRLEnv._get_observations() is called without arguments by the framework; adding clone and normalize with defaults preserves compatibility, but these parameters exist only to serve CartpoleCameraPresetsEnv's internal frame-stack optimization. Future subclasses and introspective tooling will see them as public API. Factoring the camera-data extraction into a private helper (_get_camera_data) and keeping _get_observations() zero-argument would be cleaner.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
|
Moving this back to draft while I rebase onto the cartpole refactor |
2 participants
Description
Frame-stacking was added to cartpole-camera by #5574 to make the task solvable on the Newton+Warp backend. Single-frame RGB observations don't carry pole velocity, so a single-image policy can only recover ~100/300 mean episode reward. On PhysX+RTX, implicit damping in the dynamics plus temporal antialiasing in the renderer leak enough adjacent-frame correlation that a single-frame policy still hits ~296. Newton+Warp has neither, so explicit frame stacking is required for the task to be solvable on this backend.
The implementation regressed throughput by 42% (54.6 → 31.7 FPS) and inflated peak VRAM 2.7× (3.13 → 8.37 GB) at R256, 1024 envs. PR #5574 added a
CircularBufferring + astacked_imageMDP term that:x/255 - mean(x)) before the buffer, forcing the buffer to be float32.torch.rollon every append → ring-sized temporary per env.step.permute + reshape + clone→ fresh(B, H, W, K*C)float32 tensor per env.step.This PR rewrites the frame-stacking hot path; the result runs 24% faster than the pre-framestack baseline with identical PPO convergence.
x/255 − meannormalize past the frame-stack buffer for RGB-like data types.CircularBuffer.stack_dimarg +.stackedproperty arrange storage so the channel-stacked output is a contiguous view of internal storage.(B, H, W, K*C)-sized allocation and a permute kernel per env.step.torch.rollon append. Replace with a front-to-back memmove.normalize_image_uint8). One kernel:(uint8 / 255) − per-(batch, channel) mean → float32.spatial_sum_uint8_tiled). Warp kernel writes(B, NUM_TILES, C)int32 partials along H; PyTorch collapses the partial dim.isaaclab.utils.images.normalize_camera_imagedispatch. One helper used by bothimage()(DirectRLEnv) andstacked_image.__call__(ManagerBasedEnv); routes uint8 contiguous inputs through the Warp fast path and falls back to PyTorch otherwise.image(clone=False)opt-out. Callers that immediately copy into their own storage skip the defensive clone._get_observationsallocates the normalize output fresh each call (no persistentout=buffer).record_transitionreturns; reusing one buffer across env.step boundaries aliases the trainer's prior obs. Fresh allocation lets PyTorch's caching allocator hand out a different block while the prior is still referenced.Type of change
Results
R256 (256×256), 1024 envs, PPO via skrl 2.1.0, 32k training steps, seed 42, perf measured over n=3 warm trials:
Screenshots
Environment
NVIDIA L40 (48 GB), driver 570.158.01. Newton+Warp backend (
presets=newton_mjwarp,newton_renderer). skrl 2.1.0 trainer.Checklist
pre-commitchecks with./isaaclab.sh --formatconfig/extension.tomlfileCONTRIBUTORS.mdor my name already exists there