feat: per-environment collision-pair capacity override#24
Open
haixuanTao wants to merge 1 commit into
Open
Conversation
Ports set_rbd_collisions_capacity() from perf/flat-narrow-phase (81c76bb) to the unified line: NexusState grows a setter that overrides the per-environment collision-pair floor used when the GPU state is allocated at finalize, plus the Python binding. Why: the 4096-pairs/env default is sized for one busy scene, not thousands of small batched envs — pair-keyed buffers scale as capacity x envs, so batching hits wgpu's 4 GiB max-buffer limit near 512 envs (the failing allocation is 4.7 GiB at 1024) and OOMs a 32 GiB card at 4096 envs on CUDA. The capacity is a floor, not a cap: the default RbdResizePolicy::Grow resizes off measured pair counts, so contacts are never lost — the LeRobot-legs biped MJCF produces bit-identical 400-step end frames at capacity 4, 16, and 4096. Measured on the RTX 5090 (capacity 16, 800-step windows, env-steps/s): cube 16,384 envs WebGPU 5.67M | CUDA+graph 5.28M (was: OOM) biped-pendulum 16,384 envs WebGPU 3.14M | CUDA+graph 2.93M (was: OOM) previous ceiling on main: ~170k env-steps/s at <=2048 envs Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NTtT75BpS7XH9DoHVbdfwN
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
NexusState::set_rbd_collisions_capacity()(+ Python binding): overrides the per-environment collision-pair floor used when the GPU rigid-body state is allocated atfinalize. 16 lines, no behavior change unless called.Why. The 4096-pairs/env default sizes every pair-keyed buffer as
capacity × envs. That's right for one busy scene, but batched RL-style workloads (one small robot per env) hit wgpu's 4 GiB max-buffer limit near 512 envs — the failing allocation is a single 4.7 GiB buffer at 1,024 envs for a 2-body scene — and OOM a 32 GiB card at 4,096 envs on the CUDA backend, while actual usage is ~7 pairs/env.Safety. The value is a floor, not a cap: the default
RbdResizePolicy::Growkeeps resizing off measured pair counts, so contacts are never dropped. Verified with a 12-DOF biped MJCF (real ground contacts throughout): bit-identical 400-step end frames at capacity 4, 16, and the default 4096.Measured (RTX 5090, capacity 16, 800-step sustained windows): the cube-per-env scene goes from a hard failure above ~512 envs to 5.67M env-steps/s at 16,384 envs (WebGPU) / 5.28M (CUDA+graph); a 10-link multibody pendulum per env reaches 3.14M / 2.93M. Numbers were taken on haixuanTao/nexus main (includes recent mb perf work), but the memory scaling this PR fixes is independent of those changes. Merged on the fork as haixuanTao#1; original implementation of the idea in the fork's
perf/flat-narrow-phaseline (81c76bb).🤖 Generated with Claude Code
https://claude.ai/code/session_01NTtT75BpS7XH9DoHVbdfwN