Python: per-environment MJCF insertion + collision-capacity override

crates/nexus_python3d/src/loaders.rs

@@ -120,6 +120,7 @@ pub fn insert_mjcf(
     mut viewer: PyRefMut<crate::viewer::NexusViewer>,
     scene_path: &std::path::Path,
     render_colliders: bool,
+    env: usize,
 ) -> PyResult<MjcfSceneInfo> {
     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 })
 }

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<NexusViewer>,
         scene_path: std::path::PathBuf,
         render_colliders: bool,
+        env: usize,
     ) -> PyResult<MjcfSceneInfo> {
-        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 -------------------------------------------------------

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) {