From a7064e316fafc38f8b68561969d2bb1d6bd598c9 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Thu, 9 Jul 2026 10:52:49 +0200 Subject: [PATCH 1/4] feat: per-step MJCF actuator control + multibody state readback (Python) - GpuMultibodySet::sync_joint_data_from_rapier: refresh every link's joint data (motor targets/gains, limits) from the rapier multibody set in one buffer write, mirroring from_rapier's traversal order. - NexusState::control_multibody_motors: runtime actuation entry point that mutates the rapier joints (e.g. rapier3d-mjcf's apply_controls_multibody) and pushes the refreshed joint data to the GPU, without marking the world dirty (no rebuild). - Viewer::read_multibody_links + links_workspace COPY_SRC: one-readback joint coordinates, link world poses and world-space velocities per env. - Python: NexusState keeps the MjcfRobotHandles from insert_mjcf and exposes actuator_names() / apply_actuator_controls(viewer, ctrl, env) with full MJCF actuator semantics; NexusViewer.read_multibody_links(state, env) returns (coords, positions, quats, linvels, angvels) numpy arrays. Together these make MJCF robots drivable per control step from Python (position-servo PD runs inside the solver) with full state observation -- the two pieces sim-to-sim eval loops need. Co-Authored-By: Claude Fable 5 --- crates/nexus_python3d/src/loaders.rs | 12 ++- crates/nexus_python3d/src/nexus.rs | 76 ++++++++++++++++++- crates/nexus_python3d/src/viewer.rs | 48 ++++++++++++ src/state.rs | 31 ++++++++ .../multibody/multibody_from_rapier.rs | 5 +- src_rbd/dynamics/multibody/multibody_set.rs | 63 +++++++++++++++ src_viewer/viewer.rs | 37 +++++++++ 7 files changed, 265 insertions(+), 7 deletions(-) diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index 5622367..09960f7 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/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>; + pub fn insert_mjcf( state: &mut nexus3d::prelude::NexusState, mut viewer: PyRefMut, scene_path: &std::path::Path, render_colliders: bool, -) -> PyResult { +) -> PyResult<(MjcfSceneInfo, Option)> { use pyo3::exceptions::PyRuntimeError; use rapier3d::parry::bounding_volume::BoundingVolume; // for `Aabb::merge` use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot}; @@ -140,6 +146,7 @@ pub fn insert_mjcf( let mut floor: Option<(glamx::Vec3, glamx::Vec3)> = None; let mut camera: Option<(glamx::Vec3, glamx::Vec3)> = None; + let mut robot_handles: Option = None; match MjcfRobot::from_file(scene_path, options) { Ok((robot, _model)) => { let world = state.rbd_world_mut(0); @@ -224,6 +231,7 @@ pub fn insert_mjcf( let eye = target + glamx::Vec3::new(radius * 2.2, -radius * 2.2, radius * 1.6); camera = Some((eye, target)); } + robot_handles = Some(handles); } Err(e) => { return Err(PyRuntimeError::new_err(format!( @@ -276,5 +284,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)) } diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 0cee52f..b10bb1f 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/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); #[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 { - 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 ``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 { + 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 (`` servos with kp/kv, `` + /// 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, + ctrl: Vec, + 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 ------------------------------------------------------- diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index d5e6983..c19e027 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -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, + env: u32, + ) -> ( + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + ) { + 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) { diff --git a/src/state.rs b/src/state.rs index e4dd821..d191e28 100644 --- a/src/state.rs +++ b/src/state.rs @@ -233,6 +233,37 @@ 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. + pub fn control_multibody_motors( + &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) } diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index cb390b3..ba981b3 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/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 diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 89734fe..a17c691 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -255,6 +255,69 @@ 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 { + &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( diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 770ea0c..78e64db 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -947,6 +947,43 @@ 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`. + pub async fn read_multibody_links( + &mut self, + state: &NexusState, + env: u32, + ) -> Vec { + 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 From c3c78d46e622b35a3dd6c29dc2769626647751ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 16 Jul 2026 14:14:59 +0200 Subject: [PATCH 2/4] chore: cargo fmt --- src_rbd/dynamics/multibody/multibody_set.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index a17c691..6e67a27 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -311,11 +311,7 @@ impl GpuMultibodySet { offset += 1; } } - backend.write_buffer( - self.links_static.buffer_mut(), - 0, - &self.links_static_mirror, - ) + backend.write_buffer(self.links_static.buffer_mut(), 0, &self.links_static_mirror) } /// Upload a new gravity vector. From bff76cc0f1372e4453339bbec181feaf38ebb2d6 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Thu, 16 Jul 2026 17:51:05 +0200 Subject: [PATCH 3/4] fix: gate multibody control/readback on dim3, add missing PyArray2 import The rebase onto upstream main left control_multibody_motors and read_multibody_links ungated while RbdState::multibodies(_mut) is dim3-only, breaking the 2D crates. Also fixes the missing numpy PyArray2 import and an unused_assignments warning in the MJCF loader. Co-Authored-By: Claude Fable 5 --- crates/nexus_python3d/src/loaders.rs | 7 +++---- crates/nexus_python3d/src/viewer.rs | 2 +- src/state.rs | 1 + src_viewer/viewer.rs | 1 + 4 files changed, 6 insertions(+), 5 deletions(-) diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index 09960f7..9c90615 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/crates/nexus_python3d/src/loaders.rs @@ -146,8 +146,7 @@ pub fn insert_mjcf( let mut floor: Option<(glamx::Vec3, glamx::Vec3)> = None; let mut camera: Option<(glamx::Vec3, glamx::Vec3)> = None; - let mut robot_handles: Option = None; - match MjcfRobot::from_file(scene_path, options) { + let robot_handles: Option = 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( @@ -231,7 +230,7 @@ pub fn insert_mjcf( let eye = target + glamx::Vec3::new(radius * 2.2, -radius * 2.2, radius * 1.6); camera = Some((eye, target)); } - robot_handles = Some(handles); + Some(handles) } Err(e) => { return Err(PyRuntimeError::new_err(format!( @@ -239,7 +238,7 @@ pub fn insert_mjcf( scene_path.display() ))); } - } + }; let loaded = camera.is_some(); let v = viewer.rust_mut(); diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index c19e027..25ad62d 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/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::*; diff --git a/src/state.rs b/src/state.rs index d191e28..4b79462 100644 --- a/src/state.rs +++ b/src/state.rs @@ -242,6 +242,7 @@ impl NexusState { /// 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( &mut self, backend: &GpuBackend, diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 78e64db..7d0781d 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -957,6 +957,7 @@ impl NexusViewer { /// 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, From 9f05a312ee653f92d58040eec34fcfb1d957a1df Mon Sep 17 00:00:00 2001 From: haixuantao Date: Tue, 21 Jul 2026 16:40:49 +0200 Subject: [PATCH 4/4] build: patch naga to haixuanTao/naga-fixed (MSL while-loop miscompile) Stock naga 29's MSL writer re-evaluates a loop's break_if condition after the continuing block has advanced the loop phis, dropping the final body iteration on Metal. In the multibody solve kernels the per-lane J.v loops run exactly one iteration per lane, so they executed zero times: zero impulses from the contact and joint/PD sweeps while gravity kept integrating -> robots free-fell through the floor on macOS. Pin the fixed fork (see zealot/docs/metal-contact-bug-proposal.md). Co-Authored-By: Claude Fable 5 --- Cargo.toml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/Cargo.toml b/Cargo.toml index b997730..ac802c5 100644 --- a/Cargo.toml +++ b/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