Python: per-step MJCF actuator control + multibody state readback

Cargo.toml

@@ -85,6 +85,12 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = [
 ] }
 
 [patch.crates-io]
+# naga 29 MSL backend miscompiles rust-gpu `while` loops on Metal: the hoisted
+# continuing block re-evaluates the `break_if` condition after the loop phis
+# were already advanced, dropping the final body iteration (multibody solve
+# produced zero impulses -> free fall on macOS). Fixed fork; see
+# zealot/docs/metal-contact-bug-proposal.md.
+naga = { git = "https://github.com/haixuanTao/naga-fixed", rev = "7c56094" }
 ## Local glam clone with SPIR-V vector-arithmetic intrinsics (Vec3 add/sub/mul/scale).
 #glam = { path = "../glam-rs" }
 # 30% faster for loop in P2G

crates/nexus_python3d/src/loaders.rs

@@ -115,12 +115,18 @@ 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> {
+) -> PyResult<(MjcfSceneInfo, Option<MjcfHandles>)> {
     use pyo3::exceptions::PyRuntimeError;
     use rapier3d::parry::bounding_volume::BoundingVolume; // for `Aabb::merge`
     use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot};
@@ -140,7 +146,7 @@ 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 handles = robot.clone().insert_using_multibody_joints(
@@ -224,14 +230,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();
@@ -276,5 +283,5 @@ pub fn insert_mjcf(
     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

@@ -52,15 +52,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 -----------------------------------------------------
@@ -273,7 +275,73 @@ impl NexusState {
         scene_path: std::path::PathBuf,
         render_colliders: bool,
     ) -> 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)?;
+        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 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 in one buffer write.
+    ///
+    /// Call once per control step, after `finalize`; the next
+    /// `NexusPipeline.simulate` steps the solver against the new targets. This
+    /// is the GPU counterpart of stepping rapier natively with actuators.
+    #[pyo3(signature = (viewer, ctrl, env=0))]
+    fn apply_actuator_controls(
+        &mut self,
+        viewer: PyRef<NexusViewer>,
+        ctrl: Vec<f32>,
+        env: usize,
+    ) -> 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(), env, |world| {
+                handles.apply_controls_multibody(
+                    &mut world.bodies,
+                    &mut world.multibody_joints,
+                    &ctrl,
+                );
+            })
+            .map_err(gpu_err)
     }
 
     // --- rbd config -------------------------------------------------------

crates/nexus_python3d/src/viewer.rs

@@ -10,7 +10,7 @@ use crate::nexus::{GpuTimestamps, NexusState};
 use crate::rbd::{RigidBodyHandle, SharedShape};
 use khal::backend::GpuBackend;
 use nexus_viewer3d::NexusViewer as RViewer;
-use numpy::{IntoPyArray, PyArray3, PyArrayMethods};
+use numpy::{IntoPyArray, PyArray2, PyArray3, PyArrayMethods};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 
@@ -282,6 +282,54 @@ impl NexusViewer {
             .map_err(|e| PyRuntimeError::new_err(format!("{e:?}")))
     }
 
+    /// Reads back environment `env`'s multibody link states from the GPU in one
+    /// readback. Returns five float32 numpy arrays, one row per link (in the
+    /// GPU build's traversal order — multibodies, then links, parent before
+    /// child; the same order `NexusState.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).
+    #[pyo3(signature = (state, env=0))]
+    #[allow(clippy::type_complexity)]
+    fn read_multibody_links<'py>(
+        &mut self,
+        py: Python<'py>,
+        state: PyRef<NexusState>,
+        env: u32,
+    ) -> (
+        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.inner_mut().read_multibody_links(&state.0, env));
+        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(),
+        )
+    }
+
     // --- misc -------------------------------------------------------------
 
     fn clear_scene(&mut self) {

src/state.rs

@@ -233,6 +233,38 @@ impl NexusState {
         &mut self.rbd_envs[env]
     }
 
+    /// Runtime actuation entry point: mutates environment `env`'s rapier
+    /// multibody joints through `f` (e.g. `rapier3d-mjcf`'s
+    /// `apply_controls_multibody`, which implements MJCF actuator semantics),
+    /// then pushes the refreshed joint data — motor targets/gains, limits — to
+    /// the GPU multibody links in one buffer write.
+    ///
+    /// Unlike [`Self::rbd_world_mut`] this does NOT mark the world dirty: motor
+    /// updates are per-step control, not a topology change, so no GPU rebuild
+    /// is triggered. Call after [`Self::finalize`]; a no-op before it.
+    #[cfg(all(feature = "dim3", feature = "rbd"))]
+    pub fn control_multibody_motors<F>(
+        &mut self,
+        backend: &GpuBackend,
+        env: usize,
+        f: F,
+    ) -> Result<(), GpuBackendError>
+    where
+        F: FnOnce(&mut PhysicsWorld),
+    {
+        let world = &mut self.rbd_envs[env];
+        f(world);
+        if let Some(rbd) = self.rbd.as_mut() {
+            rbd.multibodies_mut().sync_joint_data_from_rapier(
+                backend,
+                env as u32,
+                &world.multibody_joints,
+                &world.bodies,
+            )?;
+        }
+        Ok(())
+    }
+
     pub fn insert_rigid_body(&mut self, body: RigidBody, collider: Collider) -> RigidBodyHandle {
         self.insert_rigid_body_in(0, body, collider)
     }

src_rbd/dynamics/multibody/multibody_from_rapier.rs

@@ -385,7 +385,10 @@ impl GpuMultibodySet {
             links_static: Tensor::vector(backend, &all_statics, storage | BufferUsages::COPY_DST)
                 .unwrap(),
             links_static_mirror: all_statics.clone(),
-            links_workspace: Tensor::vector(backend, &all_ws, storage).unwrap(),
+            // COPY_SRC so hosts can read joint/link state back (observation
+            // pipelines); see `GpuMultibodySet::links_workspace`.
+            links_workspace: Tensor::vector(backend, &all_ws, storage | BufferUsages::COPY_SRC)
+                .unwrap(),
             dof_values: Tensor::vector(backend, &all_dof_vals, storage).unwrap(),
             dof_state: {
                 // Pack [velocities (N), damping (N), armature (N)] back-to-back

src_rbd/dynamics/multibody/multibody_set.rs

@@ -255,6 +255,65 @@ impl GpuMultibodySet {
         )
     }
 
+    /// Per-batch per-step link workspace (generalized coordinates, joint
+    /// rotations, world-space link velocities). Read it back with
+    /// `slow_read_buffer` for joint/base state observation; entries are laid out
+    /// `env * links_per_batch + link`, in [`from_rapier`](Self::from_rapier)'s
+    /// link traversal order.
+    pub fn links_workspace(&self) -> &Tensor<MultibodyLinkWorkspace> {
+        &self.links_workspace
+    }
+
+    /// Number of link slots per environment (the stride of
+    /// [`Self::links_workspace`] and `links_static`).
+    pub fn links_per_batch(&self) -> u32 {
+        self.links_per_batch
+    }
+
+    /// Refreshes every link's joint parameters (motor targets/gains, limits) of
+    /// environment `env` from a rapier multibody set laid out identically to the
+    /// one this GPU set was built from (same multibody/link traversal order as
+    /// [`from_rapier`](Self::from_rapier)), then uploads the `links_static`
+    /// buffer in one write.
+    ///
+    /// This is the per-step control path for actuated robots: mutate the motors
+    /// on the CPU rapier joints (e.g. via `rapier3d-mjcf`'s
+    /// `apply_controls_multibody`, which implements the MJCF actuator
+    /// semantics), then call this to push the new motor state to the GPU. Only
+    /// joint data is refreshed — coordinates, velocities and mass properties are
+    /// untouched, so this cannot be used to teleport links.
+    pub fn sync_joint_data_from_rapier(
+        &mut self,
+        backend: &GpuBackend,
+        env: u32,
+        set: &crate::rapier::dynamics::MultibodyJointSet,
+        bodies: &crate::rapier::dynamics::RigidBodySet,
+    ) -> Result<(), GpuBackendError> {
+        let base = (env * self.links_per_batch) as usize;
+        let mut offset = 0usize;
+        for mb in set.multibodies() {
+            // Mirror `from_rapier`'s fixed-root handling: a non-dynamic root has
+            // all 6 DOFs locked on the GPU even though rapier models it as free.
+            let root_is_dynamic = mb
+                .link(0)
+                .and_then(|r| bodies.get(r.rigid_body_handle()))
+                .map(|rb| rb.is_dynamic())
+                .unwrap_or(false);
+            for (link_idx, link) in mb.links().enumerate() {
+                let Some(entry) = self.links_static_mirror.get_mut(base + offset) else {
+                    return Ok(());
+                };
+                let mut data = convert_generic_joint(link.joint().data);
+                if link_idx == 0 && !root_is_dynamic {
+                    data.locked_axes = 0x3f;
+                }
+                entry.data = data;
+                offset += 1;
+            }
+        }
+        backend.write_buffer(self.links_static.buffer_mut(), 0, &self.links_static_mirror)
+    }
+
     /// Upload a new gravity vector.
     pub fn set_gravity(&mut self, backend: &GpuBackend, g: [f32; 3]) {
         self.gravity = Tensor::scalar(

src_viewer/viewer.rs

@@ -947,6 +947,44 @@ impl NexusViewer {
             .set_denoise(enabled);
     }
 
+    /// Reads back environment `env`'s multibody link workspaces from the GPU in
+    /// one readback: per link, the generalized joint coordinates, accumulated
+    /// joint rotation, world pose, and world-space velocity. Links are in the
+    /// GPU build's traversal order (multibodies, then links, parent before
+    /// child) — the same order `NexusState::control_multibody_motors` targets.
+    /// Empty when no multibody state exists.
+    ///
+    /// Velocities are only meaningful after the first simulated step (the
+    /// forward-kinematics pass fills them); coordinates and poses are valid
+    /// from `finalize`.
+    #[cfg(feature = "dim3")]
+    pub async fn read_multibody_links(
+        &mut self,
+        state: &NexusState,
+        env: u32,
+    ) -> Vec<nexus::rbd::shaders::dynamics::MultibodyLinkWorkspace> {
+        let Some(rbd) = state.rbd.as_ref() else {
+            return Vec::new();
+        };
+        let mbs = rbd.multibodies();
+        let stride = mbs.links_per_batch() as usize;
+        if stride == 0 {
+            return Vec::new();
+        }
+        let mut all = bytemuck::zeroed_vec(mbs.links_workspace().len() as usize);
+        if self
+            .backend()
+            .slow_read_buffer(mbs.links_workspace().buffer(), &mut all)
+            .await
+            .is_err()
+        {
+            return Vec::new();
+        }
+        let start = (env as usize * stride).min(all.len());
+        let end = (start + stride).min(all.len());
+        all[start..end].to_vec()
+    }
+
     /// Draws example-specific egui widgets into the current frame's UI pass.
     ///
     /// Call this once per frame, after [`Self::render_frame`], to overlay a