feat: add features required by RL

Cargo.toml

@@ -42,10 +42,10 @@ bvh = "0.7"
 
 # Physics / geometry. default-features off for the no_std shader crates; host
 # crates opt back in with `features = ["default"]`.
-rapier2d = { version = "0.35.2", default-features = false }
-rapier3d = { version = "0.35.2", default-features = false }
-rapier3d-urdf = "0.35"
-rapier3d-mjcf = { version = "0.35", features = ["stl", "wavefront", "msh"] }
+rapier2d = { version = "0.35.3", default-features = false }
+rapier3d = { version = "0.35.3", default-features = false }
+rapier3d-urdf = "0.35.3"
+rapier3d-mjcf = { version = "0.35.3", features = ["stl", "wavefront", "msh"] }
 parry2d = { version = "0.30", default-features = false }
 parry3d = { version = "0.30", default-features = false }
 
@@ -93,8 +93,6 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = [
 ] }
 
 [patch.crates-io]
-# Compare against the rapier checkout the reference example runs, not the
-# crates.io release (their solver defaults differ).
 #rapier2d = { path = "../rapier/crates/rapier2d" }
 #rapier3d = { path = "../rapier/crates/rapier3d" }
 #rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" }

crates/examples3d/rbd_joint_ball3.rs

@@ -42,7 +42,7 @@ pub async fn run(
             let collider = if status == RigidBodyType::Fixed {
                 ColliderBuilder::cuboid(rad, rad, rad).build()
             } else {
-                ColliderBuilder::ball(rad).density(10.0).build()
+                ColliderBuilder::ball(rad).density(40.0).build()
             };
             let shape = collider.shared_shape().clone();
             let child_handle = state.insert_rigid_body(rigid_body, collider, no_coupling);

crates/nexus_python3d/src/loaders.rs

@@ -115,12 +115,19 @@ struct VisualMeshReg {
 /// floor, camera and light with the viewer. Mirrors the Rust `mujoco_menagerie3`
 /// example's `load_scene` (minus the runtime model picker). Gravity is left to
 /// the caller (set after `finalize`).
+/// Handles of the robot loaded by [`insert_mjcf`], kept by the Python
+/// `NexusState` so per-step actuator control (`apply_actuator_controls`) can
+/// reuse `rapier3d-mjcf`'s MJCF actuator semantics.
+pub type MjcfHandles =
+    rapier3d_mjcf::MjcfRobotHandles<Option<rapier3d::prelude::MultibodyJointHandle>>;
+
 pub fn insert_mjcf(
     state: &mut nexus3d::prelude::NexusState,
     mut viewer: PyRefMut<crate::viewer::NexusViewer>,
     scene_path: &std::path::Path,
     render_colliders: bool,
-) -> PyResult<MjcfSceneInfo> {
+    env: usize,
+) -> PyResult<(MjcfSceneInfo, Option<MjcfHandles>)> {
     use nexus3d::prelude::RbdCoupling;
     use pyo3::exceptions::PyRuntimeError;
     use rapier3d::parry::bounding_volume::BoundingVolume; // for `Aabb::merge`
@@ -141,9 +148,9 @@ pub fn insert_mjcf(
     let mut floor: Option<(glamx::Vec3, glamx::Vec3)> = None;
     let mut camera: Option<(glamx::Vec3, glamx::Vec3)> = None;
 
-    match MjcfRobot::from_file(scene_path, options) {
+    let robot_handles: Option<MjcfHandles> = 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,
@@ -225,14 +232,15 @@ pub fn insert_mjcf(
                 let eye = target + glamx::Vec3::new(radius * 2.2, -radius * 2.2, radius * 1.6);
                 camera = Some((eye, target));
             }
+            Some(handles)
         }
         Err(e) => {
             return Err(PyRuntimeError::new_err(format!(
                 "failed to load MJCF {}: {e}",
                 scene_path.display()
             )));
         }
-    }
+    };
 
     let loaded = camera.is_some();
     let v = viewer.rust_mut();
@@ -242,11 +250,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, RbdCoupling::None);
-        v.insert_shape(handle, &shape, rp::Pose::IDENTITY);
+        let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None);
+        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);
         }
@@ -271,11 +283,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 })
+    Ok((MjcfSceneInfo { z_up: true, loaded }, robot_handles))
 }

crates/nexus_python3d/src/nexus.rs

@@ -15,6 +15,7 @@ use nexus3d::prelude::{
     NexusPipeline as RNexusPipeline, NexusPipelineMask, NexusState as RNexusState,
     RbdCoupling as RRbdCoupling,
 };
+use numpy::PyArray2;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use rapier3d::prelude as rp;
@@ -95,15 +96,17 @@ impl GpuTimestamps {
 }
 
 /// The GPU-resident state of a multiphysics simulation
-/// (`nexus3d::prelude::NexusState`).
+/// (`nexus3d::prelude::NexusState`). The second field keeps the
+/// `rapier3d-mjcf` robot handles of the last `insert_mjcf`, so
+/// `apply_actuator_controls` can drive the robot's actuators per step.
 #[pyclass(name = "NexusState", unsendable)]
-pub struct NexusState(pub RNexusState);
+pub struct NexusState(pub RNexusState, pub Option<crate::loaders::MjcfHandles>);
 
 #[pymethods]
 impl NexusState {
     #[new]
     fn new() -> Self {
-        NexusState(RNexusState::default())
+        NexusState(RNexusState::default(), None)
     }
 
     // --- rigid bodies -----------------------------------------------------
@@ -322,17 +325,145 @@ 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 it 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)
+        let (info, handles) =
+            crate::loaders::insert_mjcf(&mut self.0, viewer, &scene_path, render_colliders, env)?;
+        self.1 = handles;
+        Ok(info)
+    }
+
+    // --- MJCF actuation -----------------------------------------------------
+
+    /// Names of the MJCF `<actuator>`s of the robot loaded by `insert_mjcf`, in
+    /// actuator (control-vector) order. Unnamed actuators fall back to the name
+    /// of the joint they drive. Empty before `insert_mjcf`.
+    fn actuator_names(&self) -> Vec<String> {
+        self.1
+            .as_ref()
+            .map(|h| {
+                h.actuators
+                    .iter()
+                    .map(|a| {
+                        a.actuator
+                            .name
+                            .clone()
+                            .or_else(|| a.actuator.joint.clone())
+                            .unwrap_or_default()
+                    })
+                    .collect()
+            })
+            .unwrap_or_default()
+    }
+
+    /// Applies one MJCF control vector (one entry per actuator, in
+    /// `actuator_names` order) to every environment's copy of the robot loaded
+    /// by `insert_mjcf`, with full MJCF actuator semantics (`<position>` servos
+    /// with kp/kv, `<motor>` force/gear, force limits), and pushes the resulting
+    /// joint-motor state to the GPU.
+    ///
+    /// Call once per control step, after `finalize`; the next
+    /// `NexusPipeline.simulate` steps the solver against the new targets.
+    #[pyo3(signature = (viewer, ctrl))]
+    fn apply_actuator_controls(
+        &mut self,
+        viewer: PyRef<NexusViewer>,
+        ctrl: Vec<f32>,
+    ) -> PyResult<()> {
+        let Some(handles) = self.1.as_ref() else {
+            return Err(PyRuntimeError::new_err(
+                "no MJCF robot loaded (call insert_mjcf first)",
+            ));
+        };
+        if ctrl.len() != handles.actuators.len() {
+            return Err(PyRuntimeError::new_err(format!(
+                "ctrl has {} entries but the robot has {} actuators",
+                ctrl.len(),
+                handles.actuators.len()
+            )));
+        }
+        let handles = handles.clone();
+        self.0
+            .control_multibody_motors(viewer.backend(), |_, world| {
+                handles.apply_controls_multibody(
+                    &mut world.bodies,
+                    &mut world.multibody_joints,
+                    &ctrl,
+                );
+            })
+            .map_err(gpu_err)
+    }
+
+    /// Reads every environment's multibody link states back from the GPU in one
+    /// transfer. Returns five float32 numpy arrays with
+    /// `num_environments * multibody_links_per_env` rows, environment-major;
+    /// links follow the GPU build's traversal order (multibodies, then links,
+    /// parent before child), the same order `apply_actuator_controls` drives:
+    ///
+    /// - `coords (n, 6)`: generalized joint coordinates (only the joint's DOF
+    ///   count is meaningful; a revolute joint's angle is `coords[5]`),
+    /// - `positions (n, 3)` / `quats (n, 4)`: link world pose (`w, x, y, z`),
+    /// - `linvels (n, 3)` / `angvels (n, 3)`: world-space velocities, valid
+    ///   after the first simulated step.
+    ///
+    /// Use `multibody_links_per_env()` to slice a single environment out.
+    #[allow(clippy::type_complexity)]
+    fn read_multibody_links<'py>(
+        &self,
+        py: Python<'py>,
+        viewer: PyRef<NexusViewer>,
+    ) -> (
+        Bound<'py, PyArray2<f32>>,
+        Bound<'py, PyArray2<f32>>,
+        Bound<'py, PyArray2<f32>>,
+        Bound<'py, PyArray2<f32>>,
+        Bound<'py, PyArray2<f32>>,
+    ) {
+        let links = pollster::block_on(self.0.read_multibody_links(viewer.backend()));
+        let mut coords = Vec::with_capacity(links.len());
+        let mut positions = Vec::with_capacity(links.len());
+        let mut quats = Vec::with_capacity(links.len());
+        let mut linvels = Vec::with_capacity(links.len());
+        let mut angvels = Vec::with_capacity(links.len());
+        for ws in &links {
+            coords.push(ws.coords.to_vec());
+            let (t, q) = (ws.local_to_world.translation, ws.local_to_world.rotation);
+            positions.push(vec![t.x, t.y, t.z]);
+            quats.push(vec![q.w, q.x, q.y, q.z]);
+            let (l, a) = (ws.rb_vels.linear, ws.rb_vels.angular);
+            linvels.push(vec![l.x, l.y, l.z]);
+            angvels.push(vec![a.x, a.y, a.z]);
+        }
+        (
+            PyArray2::from_vec2(py, &coords).unwrap(),
+            PyArray2::from_vec2(py, &positions).unwrap(),
+            PyArray2::from_vec2(py, &quats).unwrap(),
+            PyArray2::from_vec2(py, &linvels).unwrap(),
+            PyArray2::from_vec2(py, &angvels).unwrap(),
+        )
+    }
+
+    /// Number of link slots per environment, the stride of
+    /// `read_multibody_links`.
+    fn multibody_links_per_env(&self) -> u32 {
+        self.0.multibody_links_per_env()
     }
 
     // --- rbd config -------------------------------------------------------

src/pipeline.rs

@@ -1,3 +1,4 @@
+#[cfg(feature = "mpm")]
 use crate::mpm::pipeline::MpmPipeline;
 use crate::rbd::pipeline::RbdPipeline;
 use crate::state::NexusState;
@@ -15,6 +16,7 @@ bitflags::bitflags! {
 #[derive(Default)]
 pub struct NexusPipeline {
     pub rbd_pipeline: Option<RbdPipeline>,
+    #[cfg(feature = "mpm")]
     pub mpm_pipeline: Option<MpmPipeline>,
 }
 
@@ -27,6 +29,7 @@ impl NexusPipeline {
         if pipelines.contains(NexusPipelineMask::RBD) && self.rbd_pipeline.is_none() {
             self.rbd_pipeline = Some(RbdPipeline::new(backend)?);
         }
+        #[cfg(feature = "mpm")]
         if pipelines.contains(NexusPipelineMask::MPM) && self.mpm_pipeline.is_none() {
             self.mpm_pipeline = Some(MpmPipeline::new(backend)?);
         }
@@ -73,6 +76,7 @@ impl NexusPipeline {
         }
 
         // MPM pipeline
+        #[cfg(feature = "mpm")]
         if let Some(mpm) = state.mpm.as_mut() {
             self.preload_pipelines(backend, NexusPipelineMask::MPM)?;
             let pipeline = self.mpm_pipeline.as_mut().unwrap_or_else(|| unreachable!());
@@ -91,6 +95,7 @@ impl NexusPipeline {
         // bodies as static. Push that copy back so rendering and the next
         // step's broad phase see a boundary that actually moved.
         // FIXME: the RBD pipeline should remain in charge of moving the bodies.
+        #[cfg(feature = "mpm")]
         if let (Some(rbd), Some(mpm)) = (state.rbd.as_mut(), state.mpm.as_ref()) {
             let pipeline = self.mpm_pipeline.as_ref().unwrap_or_else(|| unreachable!());
             pipeline.writeback_body_poses(backend, mpm, rbd.body_poses_mut())?;