From 4e6992ef4b7f17b2e7af0e0b541f3478dcd48281 Mon Sep 17 00:00:00 2001 From: pepijn kooijmans Date: Tue, 21 Jul 2026 16:16:32 +0200 Subject: [PATCH] feat: collision-pair capacity override (port of 81c76bb to main) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_01NTtT75BpS7XH9DoHVbdfwN --- crates/nexus_python3d/src/nexus.rs | 7 +++++++ src/state.rs | 9 +++++++++ 2 files changed, 16 insertions(+) diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 0cee52f..6a56091 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -282,6 +282,13 @@ impl NexusState { self.0.set_rbd_steps_per_frame(steps); } + /// Overrides the per-environment collision-pair capacity (default 4096) + /// used when the GPU state is allocated at `finalize`. Lower it for many + /// small batched envs; pair-keyed buffers scale as capacity × envs. + fn set_rbd_collisions_capacity(&mut self, capacity: u32) { + self.0.set_rbd_collisions_capacity(capacity); + } + fn set_rbd_gravity(&mut self, viewer: PyRef, gravity: Vec3) { self.0 .set_rbd_gravity(viewer.backend(), [gravity.0.x, gravity.0.y, gravity.0.z]); diff --git a/src/state.rs b/src/state.rs index e4dd821..625141e 100644 --- a/src/state.rs +++ b/src/state.rs @@ -160,6 +160,15 @@ impl NexusState { self.rbd_steps_per_frame = steps.max(1); } + /// Overrides the per-environment collision-pair capacity used when the + /// GPU rigid-body state is (re)allocated at `finalize`. The default (4096) + /// is sized for one busy scene, not thousands of small batched envs — + /// pair-keyed workspaces scale as `capacity x num_envs x sizeof(manifold)`, + /// which at 2048 envs binds ~9 GiB unless this is lowered. + pub fn set_rbd_collisions_capacity(&mut self, capacity: u32) { + self.capacities.rbd.collisions_capacity = capacity.max(1); + } + /// Number of rigid-body solver steps per [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call. pub fn rbd_steps_per_frame(&self) -> u32 { self.rbd_steps_per_frame