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 diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index 5622367..9c90615 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,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 = 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,6 +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)); } + Some(handles) } Err(e) => { return Err(PyRuntimeError::new_err(format!( @@ -231,7 +238,7 @@ pub fn insert_mjcf( 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)) } 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..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::*; @@ -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..4b79462 100644 --- a/src/state.rs +++ b/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( + &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..6e67a27 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/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 { + &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..7d0781d 100644 --- a/src_viewer/viewer.rs +++ b/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 { + 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