From 0447bf07a2af0bc10eb043e5ecc8bcd908c55ac3 Mon Sep 17 00:00:00 2001 From: Haixuan Xavier Tao Date: Wed, 15 Jul 2026 15:19:43 +0200 Subject: [PATCH 1/2] feat(python): per-environment MJCF insertion + collision-capacity override MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit insert_mjcf gains an env= kwarg (default 0, fully backward compatible): the robot, its auto-floor and the camera-fit AABB all target rbd_world_mut(env), and viewer registration (visual meshes, camera, light) only runs for env 0 since the viewer draws environment 0 — batch environments are physics-only. add_environment() was already exposed, so one robot per env is now reachable from Python. Why: loading N robots into env 0 is quadratic — every contact-constraint slot stores a dense M^-1 column sized to the WHOLE env's DOF count (contact_cons_col_cap = num_mbs * 192 * dofs_cap), which at 2048 robots requests ~58 GiB twice and caps single-env packing at 96 robots (121 MiB, just under wgpu's 128 MiB binding default). One robot per env keeps dofs_cap at the robot's own DOFs and scales linearly with batches. Also adds NexusState::set_rbd_collisions_capacity(): the 4096-pairs/env default is sized for one busy scene, not thousands of small envs — pair-keyed workspaces bind capacity x envs x sizeof(manifold) ~ 9.4 GiB at 2048 envs for a scene that generates ~7 pairs per env. Lowering it to 64 also halved wall-clock at 256 envs (17.6k -> 36.1k env-steps/s). Measured (RTX 5090, headless WebGPU, 12-DOF LeRobot biped, passive, 200-step window): 2,048 envs = 138,640 env-steps/s; linear from 256 (36k). Note batched finalize requires identical collider counts per env. Co-Authored-By: Claude Fable 5 --- crates/nexus_python3d/src/loaders.rs | 23 +++++++++++++++-------- crates/nexus_python3d/src/nexus.rs | 12 ++++++++++-- src/state.rs | 9 +++++++++ 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index 5622367..ec7ad22 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/crates/nexus_python3d/src/loaders.rs @@ -120,6 +120,7 @@ pub fn insert_mjcf( mut viewer: PyRefMut, scene_path: &std::path::Path, render_colliders: bool, + env: usize, ) -> PyResult { use pyo3::exceptions::PyRuntimeError; use rapier3d::parry::bounding_volume::BoundingVolume; // for `Aabb::merge` @@ -142,7 +143,7 @@ pub fn insert_mjcf( match MjcfRobot::from_file(scene_path, options) { Ok((robot, _model)) => { - let world = state.rbd_world_mut(0); + let world = state.rbd_world_mut(env); let handles = robot.clone().insert_using_multibody_joints( &mut world.bodies, &mut world.colliders, @@ -241,11 +242,15 @@ pub fn insert_mjcf( let body = rp::RigidBodyBuilder::fixed().translation(center).build(); let collider = rp::ColliderBuilder::cuboid(he.x, he.y, he.z).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); - v.insert_shape(handle, &shape, rp::Pose::IDENTITY); + let handle = state.insert_rigid_body_in(env, body, collider); + if env == 0 { + v.insert_shape(handle, &shape, rp::Pose::IDENTITY); + } } - if render_colliders { + if env != 0 { + // Batch environments are physics-only: the viewer draws environment 0. + } else if render_colliders { for (body, shape, local_pose, _) in &collider_shapes { v.insert_visual_shape(0, *body, shape, *local_pose); } @@ -270,11 +275,13 @@ pub fn insert_mjcf( } } - if let Some((eye, target)) = camera { - v.set_camera(eye, target); + if env == 0 { + if let Some((eye, target)) = camera { + v.set_camera(eye, target); + } + v.scene3d_mut() + .add_directional_light(glamx::Vec3::new(-1.0, 1.0, -1.0)); } - v.scene3d_mut() - .add_directional_light(glamx::Vec3::new(-1.0, 1.0, -1.0)); Ok(MjcfSceneInfo { z_up: true, loaded }) } diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 0cee52f..9c295df 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -266,14 +266,22 @@ impl NexusState { /// Loads a MuJoCo MJCF scene into environment 0 as multibodies, registering /// its render shapes (and a sized floor) with `viewer`. Returns scene info /// (suggested camera + whether the scene is Z-up). Call `finalize` after. - #[pyo3(signature = (viewer, scene_path, render_colliders=false))] + /// Per-environment collision-pair capacity (default 4096). Lower this + /// before `finalize` when batching many small environments — pair-keyed + /// GPU workspaces scale with `capacity x num_envs`. + fn set_rbd_collisions_capacity(&mut self, capacity: u32) { + self.0.set_rbd_collisions_capacity(capacity); + } + + #[pyo3(signature = (viewer, scene_path, render_colliders=false, env=0))] fn insert_mjcf( &mut self, viewer: PyRefMut, scene_path: std::path::PathBuf, render_colliders: bool, + env: usize, ) -> PyResult { - crate::loaders::insert_mjcf(&mut self.0, viewer, &scene_path, render_colliders) + crate::loaders::insert_mjcf(&mut self.0, viewer, &scene_path, render_colliders, env) } // --- rbd config ------------------------------------------------------- diff --git a/src/state.rs b/src/state.rs index e4dd821..0478e47 100644 --- a/src/state.rs +++ b/src/state.rs @@ -156,6 +156,15 @@ impl NexusState { /// Sets the number of rigid-body solver steps advanced per /// [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call (default 1). Acts as a simulation-speed control. + /// 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); + } + pub fn set_rbd_steps_per_frame(&mut self, steps: u32) { self.rbd_steps_per_frame = steps.max(1); } From 1aa4958f283ac8e0cb3e23b291f392493d5f16a1 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Thu, 16 Jul 2026 15:36:35 +0200 Subject: [PATCH 2/2] docs: reattach doc comments to the right functions Review feedback: the set_rbd_steps_per_frame doc was merged onto set_rbd_collisions_capacity, and insert_mjcf's doc onto the Python set_rbd_collisions_capacity binding. Co-Authored-By: Claude Fable 5 --- crates/nexus_python3d/src/nexus.rs | 7 ++++--- src/state.rs | 4 ++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 9c295df..37eceaa 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -263,9 +263,6 @@ impl NexusState { }) } - /// Loads a MuJoCo MJCF scene into environment 0 as multibodies, registering - /// its render shapes (and a sized floor) with `viewer`. Returns scene info - /// (suggested camera + whether the scene is Z-up). Call `finalize` after. /// Per-environment collision-pair capacity (default 4096). Lower this /// before `finalize` when batching many small environments — pair-keyed /// GPU workspaces scale with `capacity x num_envs`. @@ -273,6 +270,10 @@ impl NexusState { self.0.set_rbd_collisions_capacity(capacity); } + /// Loads a MuJoCo MJCF scene into environment `env` as multibodies, + /// registering its render shapes (and a sized floor) with `viewer`. Returns + /// scene info (suggested camera + whether the scene is Z-up). Call + /// `finalize` after. #[pyo3(signature = (viewer, scene_path, render_colliders=false, env=0))] fn insert_mjcf( &mut self, diff --git a/src/state.rs b/src/state.rs index 0478e47..587667b 100644 --- a/src/state.rs +++ b/src/state.rs @@ -154,8 +154,6 @@ impl NexusState { // ── Rigid-body runtime settings ───────────────────────────────────── - /// Sets the number of rigid-body solver steps advanced per - /// [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call (default 1). Acts as a simulation-speed control. /// 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 — @@ -165,6 +163,8 @@ impl NexusState { self.capacities.rbd.collisions_capacity = capacity.max(1); } + /// Sets the number of rigid-body solver steps advanced per + /// [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call (default 1). Acts as a simulation-speed control. pub fn set_rbd_steps_per_frame(&mut self, steps: u32) { self.rbd_steps_per_frame = steps.max(1); }