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..37eceaa 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -263,17 +263,26 @@ 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); + } + + /// 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, 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..587667b 100644 --- a/src/state.rs +++ b/src/state.rs @@ -154,6 +154,15 @@ impl NexusState { // ── Rigid-body runtime settings ───────────────────────────────────── + /// 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); + } + /// 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) {