From 06e5af18441d350fc1521e7a24012777192975b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 30 Aug 2026 21:24:27 +0200 Subject: [PATCH 01/25] fix joint control in the mujoco-menagerie demo --- crates/examples3d/rbd_mujoco_menagerie3.rs | 47 +++++----------------- 1 file changed, 10 insertions(+), 37 deletions(-) diff --git a/crates/examples3d/rbd_mujoco_menagerie3.rs b/crates/examples3d/rbd_mujoco_menagerie3.rs index d2159b1f..dc3d553b 100644 --- a/crates/examples3d/rbd_mujoco_menagerie3.rs +++ b/crates/examples3d/rbd_mujoco_menagerie3.rs @@ -2,8 +2,6 @@ use khal::backend::GpuTimestamps; use kiss3d::egui; use nexus_viewer3d::{NexusViewer, RenderMaterial}; use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; -use nexus3d::rbd::dynamics::convert_joint_motor; -use nexus3d::rbd::shaders::dynamics::JointMotor; use rapier3d::prelude::*; use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot, MjcfRobotHandles}; use std::fs; @@ -523,8 +521,14 @@ async fn load_scene( /// Drives the model's actuators toward `ctrl`, scaled by `gain`. /// /// The motor configuration is baked into the GPU state at finalization, so this -/// runs the MJCF actuator model on the CPU-side joints and then pushes each -/// touched motor across. +/// runs the MJCF actuator model on the CPU-side joints and then re-syncs the +/// GPU link records from them. +/// +/// `control_multibody_motors` does that sync itself, in the multibody traversal +/// order the GPU link ids follow. Pushing the motors by hand through +/// `set_motors` would need those same link ids, which are NOT the rapier body +/// indices as soon as the scene holds a body that is not a multibody link +/// (aloha's table, for instance) ahead of the robot. fn apply_controls( state: &mut NexusState, backend: &khal::backend::GpuBackend, @@ -532,45 +536,14 @@ fn apply_controls( ctrl: &[Real], gain: Real, ) { - let mut updates: Vec<(u32, usize, JointMotor)> = Vec::new(); - { - // Untracked: the rapier sets are only the scratch the MJCF actuator - // model writes into. Marking them dirty would rebuild the GPU buffers - // from the authored poses and reset the model every step. - let world = state.rbd_world_mut_untracked(0); + let _ = state.control_multibody_motors(backend, |_, world| { controls.handles.apply_controls_multibody_scaled( &mut world.bodies, &mut world.multibody_joints, ctrl, gain, ); - for ah in &controls.handles.actuators { - let Some(Some(handle)) = ah.joint else { - continue; - }; - let Some((mb, link_id)) = world.multibody_joints.get(handle) else { - continue; - }; - let Some(link) = mb.links().nth(link_id) else { - continue; - }; - // The GPU link id is the body index (see `GpuMultibodySet::set_motor`). - let body_idx = link.rigid_body_handle().into_raw_parts().0; - let axes = link.joint().data.motor_axes.bits(); - for axis in 0..6 { - if axes & (1 << axis) != 0 { - updates.push(( - body_idx, - axis, - convert_joint_motor(link.joint().data.motors[axis]), - )); - } - } - } - } - if let Some(rbd) = state.rbd.as_mut() { - let _ = rbd.multibodies_mut().set_motors(backend, 0, &updates); - } + }); } /// Picks a scene: first runs the cheap DoF pre-check (no mesh I/O); if the model From 81f433a140f7ec52360307955cef28131d2497e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 30 Aug 2026 21:40:24 +0200 Subject: [PATCH 02/25] fix(rbd): make the number of pgs iterations per substep configurable --- crates/examples3d/rbd_mujoco_menagerie3.rs | 27 +++++++++++++++++----- src/state.rs | 12 ++++++++++ src_rbd/pipeline/rbd_state.rs | 20 ++++++++++++++++ src_rbd/pipeline/rbd_state_from_rapier.rs | 1 + src_rbd_shaders/dynamics/sim_params.rs | 7 ++++++ 5 files changed, 61 insertions(+), 6 deletions(-) diff --git a/crates/examples3d/rbd_mujoco_menagerie3.rs b/crates/examples3d/rbd_mujoco_menagerie3.rs index dc3d553b..b03a197f 100644 --- a/crates/examples3d/rbd_mujoco_menagerie3.rs +++ b/crates/examples3d/rbd_mujoco_menagerie3.rs @@ -129,6 +129,9 @@ struct Settings { enable_controls: bool, enable_springs: bool, actuator_strength: f32, + /// Multibody PGS iterations per substep. Servo-driven robots resting on + /// contacts need more than one to stop the motor and contact rows fighting. + pgs_iterations: u32, /// Index into the keyframe picker: 0 is "(none)", `i + 1` is keyframe `i`. keyframe: usize, } @@ -144,6 +147,7 @@ impl Default for Settings { enable_controls: true, enable_springs: true, actuator_strength: 1.0, + pgs_iterations: 4, keyframe: 0, } } @@ -158,8 +162,8 @@ impl Settings { } /// Whether moving from `self` to `next` requires rebuilding the scene. - /// Actuator strength is read live every step, and so is the keyframe while - /// the servos are driving. + /// Actuator strength is read live every step, the PGS iteration count is + /// pushed live, and so is the keyframe while the servos are driving. fn needs_reload(&self, next: &Self) -> bool { self.use_multibody != next.use_multibody || self.render_colliders != next.render_colliders @@ -495,7 +499,9 @@ async fn load_scene( // multibody path instead raises the PGS iterations per substep. Mirrors the // reference example. let mut sim_params = nexus3d::rbd::shaders::dynamics::RbdSimParams::default(); - if !settings.use_multibody { + if settings.use_multibody { + sim_params.num_internal_pgs_iterations = settings.pgs_iterations; + } else { sim_params.dt = 1.0 / 240.0; sim_params.num_solver_iterations = 12; } @@ -504,9 +510,6 @@ async fn load_scene( state.finalize(viewer.backend()).await?; state.set_rbd_gravity(viewer.backend(), [0.0, 0.0, gravity]); if let Some(rbd) = state.rbd.as_mut() { - if settings.use_multibody { - rbd.multibodies_mut().set_num_internal_pgs_iterations(4); - } // MuJoCo-style explicit coriolis: a single plain mass matrix, with // coriolis / gyroscopic forces applied explicitly on the rhs. rbd.set_implicit_coriolis(viewer.backend(), false); @@ -686,6 +689,11 @@ pub async fn run( ui.checkbox(&mut next.disable_collisions, "Disable collisions"); ui.checkbox(&mut next.enable_controls, "Enable joint controls"); ui.checkbox(&mut next.enable_springs, "Enable joint springs"); + ui.add_enabled( + next.use_multibody, + egui::Slider::new(&mut next.pgs_iterations, 1..=16) + .text("PGS iterations / substep"), + ); ui.add( egui::Slider::new(&mut next.actuator_strength, 0.02..=2.0) .text("Actuator strength"), @@ -752,7 +760,14 @@ pub async fn run( let mut reload = false; if let Some(next) = pending_settings.take() { reload = settings.needs_reload(&next); + let pgs_changed = settings.pgs_iterations != next.pgs_iterations; settings = next; + if pgs_changed && !reload { + state.set_rbd_num_internal_pgs_iterations( + viewer.backend(), + settings.pgs_iterations, + ); + } } // A model change always rebuilds, and resets the keyframe to the new // model's default. diff --git a/src/state.rs b/src/state.rs index ca57904d..6f85438b 100644 --- a/src/state.rs +++ b/src/state.rs @@ -344,6 +344,18 @@ impl NexusState { } } + /// Sets how many PGS iterations the multibody solver's biased pass runs per + /// substep. + #[cfg(all(feature = "rbd", feature = "dim3"))] + pub fn set_rbd_num_internal_pgs_iterations(&mut self, backend: &GpuBackend, n: u32) { + for params in &mut self.rbd_sim_params { + params.num_internal_pgs_iterations = n.max(1); + } + if let Some(rbd) = self.rbd.as_mut() { + rbd.set_num_internal_pgs_iterations(backend, n); + } + } + // ── Rigid-body runtime settings ───────────────────────────────────── /// Overrides the per-environment collision-pair capacity used when the GPU diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 05cc8235..9d5f1255 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -417,6 +417,26 @@ impl RbdState { self.sim_params_cpu = params; } + /// Sets how many PGS iterations the multibody solver's biased pass runs per + /// substep, without rebuilding the GPU state. + #[cfg(feature = "dim3")] + pub fn set_num_internal_pgs_iterations(&mut self, backend: &GpuBackend, n: u32) { + let n = n.max(1); + self.multibodies.set_num_internal_pgs_iterations(n); + // Keep the mirror (and the uniform it backs) honest, even though no + // shader reads this field. + let mut params = self.sim_params_cpu; + params.num_internal_pgs_iterations = n; + let _ = backend.write_buffer(self.sim_params.buffer_mut(), 0, &[params]); + self.sim_params_cpu = params; + } + + /// PGS iterations per substep in the multibody solver's biased pass. + #[cfg(feature = "dim3")] + pub fn num_internal_pgs_iterations(&self) -> u32 { + self.multibodies.num_internal_pgs_iterations() + } + /// The gravity uniform shared by every solver kernel. pub fn gravity(&self) -> &Tensor { &self.gravity diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index b5b03916..9a578ecf 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -495,6 +495,7 @@ impl RbdState { // `set_visible_dt` divides by the substep count, so that has to be // in place first or the multibody integrates at the wrong rate. mb.set_num_solver_iterations(num_solver_iterations); + mb.set_num_internal_pgs_iterations(sim_params.num_internal_pgs_iterations); mb.set_visible_dt(backend, multibody_dt); // Soft contact coefficients (rapier TGS-soft) from the substep sim // params, so multibody-vs-floor contacts use the same soft ERP + CFM diff --git a/src_rbd_shaders/dynamics/sim_params.rs b/src_rbd_shaders/dynamics/sim_params.rs index fbf8d662..93c70361 100644 --- a/src_rbd_shaders/dynamics/sim_params.rs +++ b/src_rbd_shaders/dynamics/sim_params.rs @@ -167,6 +167,12 @@ pub struct RbdSimParams { /// `-1.0` merges every manifold of a collider pair regardless of normal: /// cheaper, but one averaged normal then stands in for a ridge or a step. pub contact_merge_cos: f32, + + /// Multibody only: PGS iterations over the joint + contact constraints run + /// per substep, in the biased pass (default: `1`). + /// + /// Host-side only: it is a dispatch count, never read by a shader. + pub num_internal_pgs_iterations: u32, } impl RbdSimParams { @@ -191,6 +197,7 @@ impl RbdSimParams { contact_merge_cos: crate::broad_phase::COS_MERGE_ANGLE, normalized_max_linear_velocity: 400.0, length_unit: 1.0, + num_internal_pgs_iterations: 1, } } } From 849cab1b2976d7e446554247a3624361a1f82a65 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 16:07:00 +0200 Subject: [PATCH 03/25] feat: robot state, sensor cameras and joint control for the Python bindings --- Cargo.toml | 17 +- crates/nexus_python3d/README.md | 38 + crates/nexus_python3d/src/lib.rs | 2 + crates/nexus_python3d/src/loaders.rs | 13 + crates/nexus_python3d/src/math.rs | 23 + crates/nexus_python3d/src/nexus.rs | 750 +++++++++++++++++++- crates/nexus_python3d/src/robot.rs | 291 ++++++++ crates/nexus_python3d/src/viewer.rs | 212 +++++- src/state.rs | 383 ++++++++++ src_rbd/dynamics/multibody/multibody_set.rs | 87 +++ src_rbd/pipeline/rbd_state.rs | 11 + src_viewer/graphics.rs | 139 +++- src_viewer/lib.rs | 6 +- src_viewer/sensors.rs | 228 ++++++ src_viewer/viewer.rs | 249 ++++++- 15 files changed, 2432 insertions(+), 17 deletions(-) create mode 100644 crates/nexus_python3d/src/robot.rs create mode 100644 src_viewer/sensors.rs diff --git a/Cargo.toml b/Cargo.toml index ef0117eb..f6b140f1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -102,11 +102,14 @@ khal-std = { git = "https://github.com/dimforge/khal", branch = "main" } khal-builder = { git = "https://github.com/dimforge/khal", branch = "main" } vortx = { git = "https://github.com/dimforge/vortx", branch = "main" } vortx-shaders = { git = "https://github.com/dimforge/vortx", branch = "main" } -#rapier2d = { path = "../rapier/crates/rapier2d" } -#rapier3d = { path = "../rapier/crates/rapier3d" } -#rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" } -#rapier3d-urdf = { path = "../rapier/crates/rapier3d-urdf" } -#rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" } +# rapier3d-urdf's URDF roll-pitch-yaw fix (fix-urdf-rpy branch), pending a +# rapier release; every rapier crate rides the same source so there is one +# `rapier3d` in the graph. +rapier2d = { path = "../rapier/crates/rapier2d" } +rapier3d = { path = "../rapier/crates/rapier3d" } +rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" } +rapier3d-urdf = { path = "../rapier/crates/rapier3d-urdf" } +rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" } ## Local glam clone with SPIR-V vector-arithmetic intrinsics (Vec3 add/sub/mul/scale). #glam = { path = "../glam-rs" } # 30% faster for loop in P2G @@ -124,7 +127,9 @@ vortx-shaders = { git = "https://github.com/dimforge/vortx", branch = "main" } #vortx-shaders = { git = "https://github.com/dimforge/vortx", branch = "int-reduce" } #glamx = { path = "../glamx" } -#kiss3d = { path = "../kiss3d" } +# kiss3d's shared-manager fix for a second offscreen surface (branch +# fix-shared-window-managers), pending a kiss3d release. +kiss3d = { path = "../kiss3d" } #parry2d = { path = "../parry/crates/parry2d" } #parry3d = { path = "../parry/crates/parry3d" } #rapier2d = { path = "../rapier/crates/rapier2d" } diff --git a/crates/nexus_python3d/README.md b/crates/nexus_python3d/README.md index 74ffc1e9..d30030d6 100644 --- a/crates/nexus_python3d/README.md +++ b/crates/nexus_python3d/README.md @@ -39,6 +39,44 @@ pip install target/wheels/dimforge_nexus3d-*.whl > so a single wheel works across all supported Python versions — you don't need > to match the build interpreter to the run interpreter. +## Robots, state access and sensor cameras + +Beyond the demo-oriented API, the module exposes what a robotics environment +needs to drive Nexus as its physics engine (the AISLE bridge is the first +consumer): + +- `NexusState.load_urdf_robot(env, path, options)` / `load_mjcf_robot(viewer, + env, path)` return a `Robot`: link and joint names in the multibody's own + order, per-DoF limits, and the PD gains (`kp`, `kv`, `max_force`) that + `set_robot_targets` applies as force-based position motors. MJCF position + actuators seed the gains; URDF joints start at zero, and + `set_robot_joint_dynamics` gives them armature and damping before + `finalize`. +- `robot_state(viewer, robot)` reads joint coordinates and link poses and + velocities back from the GPU; `set_robot_qpos` teleports the joints (and + zeroes their velocities) between steps. +- `robot_forward_kinematics` / `robot_inverse_kinematics` run on the CPU + kinematic model (damped least squares, joint limits clamped, a per-axis + world-frame constraint mask). +- `read_body_poses` / `read_body_velocities` / `set_body_pose` / + `set_body_velocity` do the same for free rigid bodies; `set_rbd_timestep` + fixes the step and substep count before `finalize`. +- `NexusViewer.add_sensor_camera(w, h, fov_y_deg)` creates an offscreen + camera; `set_sensor_camera_pose` (OpenGL frame: -Z forward, +Y up) or + `attach_sensor_camera` (follow a body with a mount pose) place it, and + `render_sensor_camera(id, rgb, depth, segmentation)` returns an RGB image, + linear metric depth and per-body segmentation ids. Call + `set_sensor_rendering(True)` before inserting shapes so every body gets its + own render node (the depth and segmentation passes ignore GPU instances), + and `set_body_segmentation_id` to choose the ids. `insert_sensor_shape` + registers a per-body node with an optional UV-mapped texture. +- Several `NexusState`s can share one viewer: `begin_scene()` opens a render + generation for the next scene and `set_active_scene(gen)` switches which + scene's nodes and cameras are live. + +The URDF loader relies on the rapier3d-urdf roll-pitch-yaw fix on the rapier +`fix-urdf-rpy` branch (patched in the workspace `Cargo.toml`). + ## Examples [`examples/`](examples) contains Python ports of the 3D Rust demos in diff --git a/crates/nexus_python3d/src/lib.rs b/crates/nexus_python3d/src/lib.rs index e4fe4623..c6e2591b 100644 --- a/crates/nexus_python3d/src/lib.rs +++ b/crates/nexus_python3d/src/lib.rs @@ -13,6 +13,7 @@ pub mod math; pub mod mpm; pub mod nexus; pub mod rbd; +pub mod robot; pub mod viewer; /// The `nexus3d` Python module. @@ -61,6 +62,7 @@ fn nexus3d(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; // MPM m.add_class::()?; diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index b19602d9..bdc0bbb8 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/crates/nexus_python3d/src/loaders.rs @@ -23,6 +23,9 @@ pub struct UrdfLoaderOptions { pub enable_joint_collisions: bool, pub scale: f32, pub shift: Option, + /// Replace every mesh collider by its convex hull (`False` keeps the + /// triangle mesh). + pub convex_hull: bool, } #[pymethods] @@ -36,7 +39,9 @@ impl UrdfLoaderOptions { enable_joint_collisions=false, scale=1.0, shift=None, + convex_hull=false, ))] + #[allow(clippy::too_many_arguments)] fn new( create_colliders_from_collision_shapes: bool, create_colliders_from_visual_shapes: bool, @@ -45,6 +50,7 @@ impl UrdfLoaderOptions { enable_joint_collisions: bool, scale: f32, shift: Option, + convex_hull: bool, ) -> Self { Self { create_colliders_from_collision_shapes, @@ -54,6 +60,7 @@ impl UrdfLoaderOptions { enable_joint_collisions, scale, shift, + convex_hull, } } } @@ -68,6 +75,12 @@ impl UrdfLoaderOptions { enable_joint_collisions: self.enable_joint_collisions, scale: self.scale, shift: self.shift.map(|p| p.0).unwrap_or(rp::Pose::IDENTITY), + mesh_converter: self + .convex_hull + .then_some(rapier3d::prelude::MeshConverter::ConvexHull), + // Keep every URDF link so `UrdfRobot::links` stays index-aligned + // with `urdf_rs::Robot::links` (link and joint names). + squeeze_empty_fixed_links: false, ..Default::default() } } diff --git a/crates/nexus_python3d/src/math.rs b/crates/nexus_python3d/src/math.rs index 779922b1..35a4c0ab 100644 --- a/crates/nexus_python3d/src/math.rs +++ b/crates/nexus_python3d/src/math.rs @@ -155,6 +155,29 @@ impl Quat { Self(glamx::Quat::from_scaled_axis(axisangle.0)) } + /// A (normalized) quaternion from its `x, y, z, w` components. + #[staticmethod] + fn from_xyzw(x: f32, y: f32, z: f32, w: f32) -> Self { + Self(glamx::Quat::from_xyzw(x, y, z, w).normalize()) + } + + #[getter] + fn x(&self) -> f32 { + self.0.x + } + #[getter] + fn y(&self) -> f32 { + self.0.y + } + #[getter] + fn z(&self) -> f32 { + self.0.z + } + #[getter] + fn w(&self) -> f32 { + self.0.w + } + fn __mul__(&self, rhs: Quat) -> Quat { Quat(self.0 * rhs.0) } diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index caa55932..f5b539ea 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -8,6 +8,7 @@ use crate::rbd::{ Collider, ImpulseJointHandle, JointArg, JointAxis, MultibodyJointHandle, RigidBody, RigidBodyHandle, SharedShape, }; +use crate::robot::{Robot, build_robot, free_axes, joint_axis, pose_from_wxyz, to_wxyz}; use crate::viewer::NexusViewer; use khal::backend::GpuTimestamps as RGpuTimestamps; use nexus3d::mpm::solver::BoundaryCondition as RBoundaryCondition; @@ -15,10 +16,11 @@ use nexus3d::prelude::{ NexusPipeline as RNexusPipeline, NexusPipelineMask, NexusState as RNexusState, RbdCoupling as RRbdCoupling, }; -use numpy::PyArray2; +use numpy::{PyArray1, PyArray2}; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use rapier3d::prelude as rp; +use std::collections::HashMap; /// Maps a GPU backend error to a Python exception. fn gpu_err(e: E) -> PyErr { @@ -466,6 +468,661 @@ impl NexusState { self.0.multibody_links_per_env() } + // --- robots ------------------------------------------------------------- + + /// Loads a URDF robot into environment `env` as one multibody and returns + /// its [`Robot`] (names, DoF layout, limits, render shapes). Register the + /// render shapes with `viewer.insert_visual_shape(env, body, shape, pose)`. + #[pyo3(signature = (env, path, options))] + fn load_urdf_robot( + &mut self, + env: usize, + path: std::path::PathBuf, + options: PyRef, + ) -> PyResult { + use rapier3d_urdf::{UrdfMultibodyOptions, UrdfRobot}; + let opts = options.to_rapier(); + let (robot, urdf) = UrdfRobot::from_file(&path, opts, None).map_err(|e| { + PyRuntimeError::new_err(format!("failed to load URDF {}: {e}", path.display())) + })?; + if env >= self.0.num_environments() { + return Err(PyRuntimeError::new_err(format!( + "environment {env} does not exist" + ))); + } + let world = self.0.rbd_world_mut(env); + let handles = robot.insert_using_multibody_joints( + &mut world.bodies, + &mut world.colliders, + &mut world.multibody_joints, + UrdfMultibodyOptions::DISABLE_SELF_CONTACTS, + ); + let mut body_names = HashMap::new(); + let mut child_joint_names = HashMap::new(); + for (link, urdf_link) in handles.links.iter().zip(&urdf.links) { + body_names.insert(link.body, urdf_link.name.clone()); + } + for (joint, urdf_joint) in handles.joints.iter().zip(&urdf.joints) { + child_joint_names.insert(joint.link2, urdf_joint.name.clone()); + } + let mut render_shapes = Vec::new(); + for link in &handles.links { + for collider in &link.colliders { + let (shape, local_pose) = match &collider.visual { + Some(v) => (v.shape.clone(), v.local_pose), + None => ( + world.colliders[collider.handle].shared_shape().clone(), + rp::Pose::IDENTITY, + ), + }; + render_shapes.push(( + RigidBodyHandle(link.body), + SharedShape(shape), + Pose(local_pose), + )); + } + } + let root = handles + .links + .first() + .map(|l| l.body) + .ok_or_else(|| PyRuntimeError::new_err("URDF has no links"))?; + build_robot( + world, + env, + root, + &body_names, + &child_joint_names, + render_shapes, + ) + } + + /// Loads an MJCF robot into environment `env` as one multibody, registers + /// its visual meshes with `viewer` (environment 0 only) and returns its + /// [`Robot`]. Unlike `insert_mjcf` this adds no floor and moves no camera. + /// Joints driven by MJCF position actuators start with those gains as + /// their PD defaults. + #[pyo3(signature = (viewer, env, path, register_visuals=true))] + fn load_mjcf_robot( + &mut self, + mut viewer: PyRefMut, + env: usize, + path: std::path::PathBuf, + register_visuals: bool, + ) -> PyResult { + use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot}; + let options = MjcfLoaderOptions { + skip_plane_geoms: true, + make_roots_fixed: false, + create_colliders_from_visual_shapes: false, + collider_blueprint: rp::ColliderBuilder::default().density(0.0), + ..Default::default() + }; + let (robot, _model) = MjcfRobot::from_file(&path, options).map_err(|e| { + PyRuntimeError::new_err(format!("failed to load MJCF {}: {e}", path.display())) + })?; + if env >= self.0.num_environments() { + return Err(PyRuntimeError::new_err(format!( + "environment {env} does not exist" + ))); + } + let world = self.0.rbd_world_mut(env); + let handles = robot.clone().insert_using_multibody_joints( + &mut world.bodies, + &mut world.colliders, + &mut world.multibody_joints, + &mut world.impulse_joints, + MjcfMultibodyOptions::DISABLE_SELF_CONTACTS, + ); + // Position servos hold their neutral target from the start. + let ctrl = vec![0.0; handles.actuators.len()]; + handles.apply_controls_multibody(&mut world.bodies, &mut world.multibody_joints, &ctrl); + + let mut body_names = HashMap::new(); + let mut child_joint_names = HashMap::new(); + let mut root = None; + for (i, bh) in handles.bodies.iter().enumerate() { + let Some(bh) = bh else { continue }; + root.get_or_insert(bh.body); + let name = robot.bodies[i] + .name + .clone() + .unwrap_or_else(|| format!("body_{i}")); + body_names.insert(bh.body, name); + } + for (jh, mj) in handles.joints.iter().zip(&robot.joints) { + if let Some(name) = &mj.name { + child_joint_names.insert(jh.link2, name.clone()); + } + } + let root = root.ok_or_else(|| PyRuntimeError::new_err("MJCF has no bodies"))?; + + if register_visuals && env == 0 { + let v = viewer.rust_mut(); + for (i, bh) in handles.bodies.iter().enumerate() { + let Some(bh) = bh else { continue }; + let mjcf_body = &robot.bodies[i]; + if mjcf_body.visual_meshes.is_empty() { + for collider in &bh.colliders { + let c = &world.colliders[collider.handle]; + let local_pose = c + .position_wrt_parent() + .copied() + .unwrap_or(rp::Pose::IDENTITY); + v.insert_visual_shape(0, bh.body, c.shared_shape(), local_pose); + } + continue; + } + for vm in &mjcf_body.visual_meshes { + let textured = vm.texture.is_some(); + let color = vm.rgba.unwrap_or(if textured { + [1.0, 1.0, 1.0, 1.0] + } else { + [0.7, 0.7, 0.75, 1.0] + }); + let material = vm + .material + .as_ref() + .map(|m| nexus_viewer3d::RenderMaterial { + metallic: m.metallic, + roughness: m.roughness, + reflectance: m.reflectance, + emissive: m.emissive, + }); + v.insert_visual_mesh( + 0, + bh.body, + &vm.shape, + vm.local_pose, + color, + vm.uvs.as_deref(), + vm.normals.as_deref(), + vm.texture.as_deref(), + material, + ); + } + } + } + build_robot( + world, + env, + root, + &body_names, + &child_joint_names, + Vec::new(), + ) + } + + /// Sets the per-DoF reflected rotor inertia (`armature`, added to the + /// mass-matrix diagonal) and viscous joint damping of `robot`, one value + /// per robot DoF each (`None` leaves that quantity unchanged). Both are + /// read at the next GPU build, so call before `finalize`. MJCF robots + /// carry theirs from the model; URDF robots start at zero, which leaves a + /// light arm without the stabilizing inertia its servos assume. + #[pyo3(signature = (robot, armature=None, damping=None))] + fn set_robot_joint_dynamics( + &mut self, + robot: PyRef, + armature: Option>, + damping: Option>, + ) -> PyResult<()> { + let robot = robot.clone(); + let n = robot.dof_axes.len(); + for (name, values) in [("armature", &armature), ("damping", &damping)] { + if let Some(v) = values + && v.len() != n + { + return Err(PyRuntimeError::new_err(format!( + "{name} has {} entries for {n} dofs", + v.len() + ))); + } + } + let world = self.0.rbd_world_mut(robot.env); + let rp::PhysicsWorld { + bodies, + multibody_joints, + .. + } = world; + let link_id = *multibody_joints + .rigid_body_link(robot.root) + .ok_or_else(|| PyRuntimeError::new_err("robot is not a multibody"))?; + let mb = multibody_joints + .get_multibody_mut(link_id.multibody) + .ok_or_else(|| PyRuntimeError::new_err("robot multibody not found"))?; + let map = assembly_dof_map(mb, bodies); + if let Some(values) = armature { + let full = to_assembly(&values, &map); + let target = mb.armature_mut(); + for (i, v) in full.iter().enumerate() { + if i < target.len() { + target[i] = *v; + } + } + } + if let Some(values) = damping { + let full = to_assembly(&values, &map); + let target = mb.damping_mut(); + for (i, v) in full.iter().enumerate() { + if i < target.len() { + target[i] = *v; + } + } + } + Ok(()) + } + + /// Current per-DoF `(armature, damping)` of `robot` on the CPU model. + fn robot_joint_dynamics(&self, robot: PyRef) -> PyResult<(Vec, Vec)> { + let world = self.0.rbd_world(robot.env); + let link_id = world + .multibody_joints + .rigid_body_link(robot.root) + .ok_or_else(|| PyRuntimeError::new_err("robot is not a multibody"))?; + let mb = world + .multibody_joints + .get_multibody(link_id.multibody) + .ok_or_else(|| PyRuntimeError::new_err("robot multibody not found"))?; + let map = assembly_dof_map(mb, &world.bodies); + let pick = |v: &[f32]| -> Vec { + map.iter() + .enumerate() + .filter_map(|(i, d)| d.map(|_| v.get(i).copied().unwrap_or(0.0))) + .collect() + }; + Ok(( + pick(mb.armature().as_slice()), + pick(mb.damping().as_slice()), + )) + } + + /// Generalized coordinates of `robot` as the CPU multibody last saw them + /// (the authored pose, or the last `set_robot_qpos`). For the simulated + /// values use `robot_qpos`. + fn robot_cpu_qpos(&self, robot: PyRef) -> Vec { + self.0 + .multibody_joint_positions(robot.env, robot.root) + .unwrap_or_default() + } + + /// Per-link GPU slots of `robot` (indices into `read_multibody_links` rows, + /// before the `env * multibody_links_per_env()` offset), in link order. + fn robot_link_slots(&self, robot: PyRef) -> Vec { + self.link_slots(&robot) + } + + /// Reads `robot`'s simulated state back from the GPU: generalized + /// coordinates `qpos (n_dofs,)`, link positions `(n_links, 3)`, link + /// quaternions `(n_links, 4)` as `(w, x, y, z)`, and link linear and + /// angular velocities `(n_links, 3)`. + #[allow(clippy::type_complexity)] + fn robot_state<'py>( + &self, + py: Python<'py>, + viewer: PyRef, + robot: PyRef, + ) -> PyResult<( + Bound<'py, PyArray1>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + )> { + let links = pollster::block_on(self.0.read_multibody_links(viewer.backend())); + let per_env = self.0.multibody_links_per_env() as usize; + let slots = self.link_slots(&robot); + if slots.len() != robot.link_names.len() { + return Err(PyRuntimeError::new_err( + "robot state unavailable: call finalize() first", + )); + } + let mut qpos = Vec::with_capacity(robot.dof_axes.len()); + let mut pos = Vec::with_capacity(slots.len()); + let mut quat = Vec::with_capacity(slots.len()); + let mut linvel = Vec::with_capacity(slots.len()); + let mut angvel = Vec::with_capacity(slots.len()); + for (link_idx, slot) in slots.iter().enumerate() { + let ws = links + .get(robot.env * per_env + *slot as usize) + .ok_or_else(|| PyRuntimeError::new_err("multibody readback too short"))?; + for (d, axis) in robot.dof_axes.iter().enumerate() { + if robot.dof_links[d] == link_idx { + qpos.push(ws.coords[*axis as usize]); + } + } + let t = ws.local_to_world.translation; + pos.push(vec![t.x, t.y, t.z]); + quat.push(to_wxyz(ws.local_to_world.rotation).to_vec()); + let (l, a) = (ws.rb_vels.linear, ws.rb_vels.angular); + linvel.push(vec![l.x, l.y, l.z]); + angvel.push(vec![a.x, a.y, a.z]); + } + Ok(( + PyArray1::from_vec(py, qpos), + PyArray2::from_vec2(py, &pos).unwrap(), + PyArray2::from_vec2(py, &quat).unwrap(), + PyArray2::from_vec2(py, &linvel).unwrap(), + PyArray2::from_vec2(py, &angvel).unwrap(), + )) + } + + /// Sets `robot`'s generalized coordinates and zeroes its joint velocities + /// on the GPU (between steps, after `finalize`). + fn set_robot_qpos( + &mut self, + viewer: PyRef, + robot: PyRef, + qpos: Vec, + ) -> PyResult<()> { + self.0 + .set_multibody_joint_positions(viewer.backend(), robot.env, robot.root, &qpos) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + /// Drives `robot`'s DoFs toward `targets` with the robot's per-DoF PD + /// gains (`robot.kp`, `robot.kv`, `robot.max_force`; force-based motors). + /// `dofs` selects a subset (default: all, in which case `targets` has one + /// entry per DoF). Motors of unselected DoFs keep their current target. + #[pyo3(signature = (viewer, robot, targets, dofs=None))] + fn set_robot_targets( + &mut self, + viewer: PyRef, + robot: PyRef, + targets: Vec, + dofs: Option>, + ) -> PyResult<()> { + let dofs = dofs.unwrap_or_else(|| (0..robot.dof_axes.len()).collect()); + if dofs.len() != targets.len() { + return Err(PyRuntimeError::new_err(format!( + "{} targets for {} dofs", + targets.len(), + dofs.len() + ))); + } + let robot = robot.clone(); + self.0 + .control_multibody_motors_env(viewer.backend(), robot.env, |world| { + let Some(link_id) = world.multibody_joints.rigid_body_link(robot.root).copied() + else { + return; + }; + let Some(mb) = world.multibody_joints.get_multibody_mut(link_id.multibody) else { + return; + }; + for (d, target) in dofs.iter().zip(&targets) { + let (Some(link_idx), Some(axis)) = + (robot.dof_links.get(*d), robot.dof_axes.get(*d)) + else { + continue; + }; + let Some(link) = mb.link_mut(*link_idx) else { + continue; + }; + let axis = joint_axis(*axis); + link.joint + .data + .set_motor_model(axis, rp::MotorModel::ForceBased); + link.joint + .data + .set_motor_position(axis, *target, robot.kp[*d], robot.kv[*d]); + link.joint + .data + .set_motor_max_force(axis, robot.max_force[*d]); + } + }) + .map_err(gpu_err) + } + + /// Pose of link `link` of `robot` at coordinates `qpos`, from the CPU + /// kinematic model (no GPU access): `(position, (w, x, y, z))`. Leaves the + /// CPU multibody at `qpos`. + fn robot_forward_kinematics( + &mut self, + robot: PyRef, + qpos: Vec, + link: &str, + ) -> PyResult<([f32; 3], [f32; 4])> { + let link_idx = robot.link_index(link)?; + let robot = robot.clone(); + let world = self.0.rbd_world_mut_untracked(robot.env); + let mb = cpu_multibody_at(world, &robot, &qpos)?; + let pose = mb + .link(link_idx) + .map(|l| *l.local_to_world()) + .ok_or_else(|| PyRuntimeError::new_err("link index out of range"))?; + let t = pose.translation; + Ok(([t.x, t.y, t.z], to_wxyz(pose.rotation))) + } + + /// Damped-least-squares inverse kinematics on the CPU model, from + /// `init_qpos`, for link `link` to reach `target_pos` (of `local_point` in + /// the link frame) with orientation `target_quat` (`w, x, y, z`). + /// `constrained_axes` are `[lin_x, lin_y, lin_z, ang_x, ang_y, ang_z]` + /// world-frame error components the solver drives to zero; `dofs` + /// restricts which DoFs may move (default: all). Coordinates are clamped + /// to the joint limits after every iteration. Returns `(qpos, error)` + /// where `error` is `[lin(3), ang(3)]` with unconstrained components zeroed. + #[pyo3(signature = (robot, link, target_pos, target_quat, init_qpos, local_point=None, constrained_axes=None, dofs=None, max_iters=100, damping=0.05, pos_tol=1.0e-4, rot_tol=1.0e-3))] + #[allow(clippy::too_many_arguments)] + fn robot_inverse_kinematics( + &mut self, + robot: PyRef, + link: &str, + target_pos: [f32; 3], + target_quat: [f32; 4], + init_qpos: Vec, + local_point: Option<[f32; 3]>, + constrained_axes: Option<[bool; 6]>, + dofs: Option>, + max_iters: usize, + damping: f32, + pos_tol: f32, + rot_tol: f32, + ) -> PyResult<(Vec, [f32; 6])> { + use rapier3d::dynamics::InverseKinematicsOption; + let link_idx = robot.link_index(link)?; + let robot = robot.clone(); + let constrained = constrained_axes.unwrap_or([true; 6]); + let mut mask = rp::JointAxesMask::empty(); + for (axis, on) in constrained.iter().enumerate() { + if *on { + mask |= rp::JointAxesMask::from_bits_truncate(1 << axis); + } + } + let target = pose_from_wxyz(target_pos, target_quat); + let target = match local_point { + // Aim the link origin so that `local_point` lands on `target_pos`. + Some(p) => rp::Pose::from_parts( + target.translation - target.rotation * glamx::Vec3::from(p), + target.rotation, + ), + None => target, + }; + let movable: Vec = { + let mut m = vec![dofs.is_none(); robot.link_names.len()]; + if let Some(dofs) = &dofs { + for d in dofs { + if let Some(l) = robot.dof_links.get(*d) { + m[*l] = true; + } + } + } + m + }; + let options = InverseKinematicsOption { + damping, + max_iters: 1, + constrained_axes: mask, + epsilon_linear: pos_tol, + epsilon_angular: rot_tol, + }; + + let world = self.0.rbd_world_mut_untracked(robot.env); + let root = robot.root; + cpu_multibody_at(world, &robot, &init_qpos)?; + let rp::PhysicsWorld { + bodies, + multibody_joints, + .. + } = world; + let link_id = *multibody_joints + .rigid_body_link(root) + .ok_or_else(|| PyRuntimeError::new_err("robot is not a multibody"))?; + let mut error = [0.0f32; 6]; + let (map, root_fixed) = { + let mb = multibody_joints + .get_multibody(link_id.multibody) + .unwrap_or_else(|| unreachable!()); + ( + assembly_dof_map(mb, bodies), + RNexusState::multibody_root_is_fixed(bodies, mb), + ) + }; + for _ in 0..max_iters { + let mb = multibody_joints + .get_multibody_mut(link_id.multibody) + .unwrap_or_else(|| unreachable!()); + let mut disp = rapier3d::na::DVector::zeros(mb.ndofs()); + mb.inverse_kinematics( + bodies, + link_idx, + &options, + &target, + |l| { + if l.link_id() == 0 && root_fixed { + return false; + } + movable.get(l.link_id()).copied().unwrap_or(true) + }, + &mut disp, + ); + // Clamp to the joint limits: rapier's solver is unaware of them. + let current = to_assembly(&cpu_qpos(mb, bodies), &map); + let mut disp_vec: Vec = disp.as_slice().to_vec(); + for (a, (q, delta)) in current.iter().zip(disp_vec.iter_mut()).enumerate() { + match map[a] { + Some(d) => { + let clamped = (q + *delta).clamp(robot.dof_lower[d], robot.dof_upper[d]); + *delta = clamped - q; + } + None => *delta = 0.0, + } + } + mb.apply_displacements(&disp_vec); + mb.forward_kinematics(bodies, false); + let pose = *mb + .link(link_idx) + .unwrap_or_else(|| unreachable!()) + .local_to_world(); + let lin = target.translation - pose.translation; + let ang = (target.rotation * pose.rotation.inverse()).to_scaled_axis(); + let raw = [lin.x, lin.y, lin.z, ang.x, ang.y, ang.z]; + for (i, e) in raw.iter().enumerate() { + error[i] = if constrained[i] { *e } else { 0.0 }; + } + let lin_err = (error[0] * error[0] + error[1] * error[1] + error[2] * error[2]).sqrt(); + let ang_err = (error[3] * error[3] + error[4] * error[4] + error[5] * error[5]).sqrt(); + if lin_err <= pos_tol && ang_err <= rot_tol { + break; + } + } + let mb = multibody_joints + .get_multibody(link_id.multibody) + .unwrap_or_else(|| unreachable!()); + Ok((cpu_qpos(mb, bodies), error)) + } + + // --- rigid-body state --------------------------------------------------- + + /// GPU pose slot of `handle` in `env`: the row of `read_body_poses` / + /// `read_body_velocities`. `None` before `finalize`. + fn body_gpu_index(&self, env: usize, handle: RigidBodyHandle) -> Option { + self.0.rigid_body_gpu_index(env, handle.0) + } + + /// Reads every body's world-origin pose from the GPU: positions `(n, 3)` + /// and quaternions `(n, 4)` as `(w, x, y, z)`, rows indexed by + /// `body_gpu_index`. + fn read_body_poses<'py>( + &self, + py: Python<'py>, + viewer: PyRef, + ) -> (Bound<'py, PyArray2>, Bound<'py, PyArray2>) { + let poses = pollster::block_on(self.0.read_rigid_body_poses(viewer.backend())); + let mut pos = Vec::with_capacity(poses.len()); + let mut quat = Vec::with_capacity(poses.len()); + for p in &poses { + let t = p.translation; + pos.push(vec![t.x, t.y, t.z]); + quat.push(to_wxyz(p.rotation).to_vec()); + } + ( + PyArray2::from_vec2(py, &pos).unwrap(), + PyArray2::from_vec2(py, &quat).unwrap(), + ) + } + + /// Reads every body's world-space linear `(n, 3)` and angular `(n, 3)` + /// velocity from the GPU, rows indexed by `body_gpu_index`. + fn read_body_velocities<'py>( + &self, + py: Python<'py>, + viewer: PyRef, + ) -> (Bound<'py, PyArray2>, Bound<'py, PyArray2>) { + let vels = pollster::block_on(self.0.read_rigid_body_velocities(viewer.backend())); + let mut lin = Vec::with_capacity(vels.len()); + let mut ang = Vec::with_capacity(vels.len()); + for v in &vels { + lin.push(vec![v.linear.x, v.linear.y, v.linear.z]); + ang.push(vec![v.angular.x, v.angular.y, v.angular.z]); + } + ( + PyArray2::from_vec2(py, &lin).unwrap(), + PyArray2::from_vec2(py, &ang).unwrap(), + ) + } + + /// Teleports free body `handle` of `env` to `pos` / `quat` (`w, x, y, z`) + /// between steps. Multibody links go through `set_robot_qpos`. + fn set_body_pose( + &mut self, + viewer: PyRef, + env: usize, + handle: RigidBodyHandle, + pos: [f32; 3], + quat: [f32; 4], + ) -> PyResult<()> { + self.0 + .set_rigid_body_pose(viewer.backend(), env, handle.0, pose_from_wxyz(pos, quat)) + .map_err(gpu_err) + } + + /// Sets body `handle`'s world-space linear and angular velocity. + fn set_body_velocity( + &mut self, + viewer: PyRef, + env: usize, + handle: RigidBodyHandle, + linvel: [f32; 3], + angvel: [f32; 3], + ) -> PyResult<()> { + self.0 + .set_rigid_body_velocity( + viewer.backend(), + env, + handle.0, + glamx::Vec3::from(linvel), + glamx::Vec3::from(angvel), + ) + .map_err(gpu_err) + } + + /// Rigid-body timestep of every environment: `dt` seconds per `simulate` + /// step, in `substeps` solver substeps. Call before `finalize`. + fn set_rbd_timestep(&mut self, dt: f32, substeps: u32) { + self.0.set_rbd_timestep(dt, substeps); + } + // --- rbd config ------------------------------------------------------- fn set_rbd_steps_per_frame(&mut self, steps: u32) { @@ -614,3 +1271,94 @@ impl NexusPipeline { pollster::block_on(self.0.simulate(backend, &mut state.0, ts)).map_err(gpu_err) } } + +/// Generalized coordinates of `mb` in assembly order, without a fixed root's +/// locked coordinates (the GPU build's DoF vector, see `Robot`). +fn cpu_qpos(mb: &rapier3d::dynamics::Multibody, bodies: &rp::RigidBodySet) -> Vec { + let root_fixed = RNexusState::multibody_root_is_fixed(bodies, mb); + let mut out = Vec::with_capacity(mb.ndofs()); + for (i, link) in mb.links().enumerate() { + if i == 0 && root_fixed { + continue; + } + let coords = link.joint().coords(); + for axis in free_axes(&link.joint().data) { + out.push(coords[axis as usize]); + } + } + out +} + +/// For each entry of rapier's full assembly displacement vector, the robot DoF +/// it maps to (`None` for a fixed root's locked coordinates). +fn assembly_dof_map( + mb: &rapier3d::dynamics::Multibody, + bodies: &rp::RigidBodySet, +) -> Vec> { + let root_fixed = RNexusState::multibody_root_is_fixed(bodies, mb); + let mut out = Vec::with_capacity(mb.ndofs()); + let mut next = 0usize; + for (i, link) in mb.links().enumerate() { + for _ in 0..link.joint().ndofs() { + if i == 0 && root_fixed { + out.push(None); + } else { + out.push(Some(next)); + next += 1; + } + } + } + out +} + +/// Expands a robot DoF vector into rapier's full assembly vector (zeros for a +/// fixed root's locked coordinates). +fn to_assembly(values: &[f32], map: &[Option]) -> Vec { + map.iter() + .map(|d| d.map(|i| values[i]).unwrap_or(0.0)) + .collect() +} + +/// Moves `robot`'s CPU multibody to `qpos` (forward kinematics included) and +/// returns it. The CPU model is a scratch kinematic model once the GPU owns the +/// simulation, so this never touches the simulated state. +fn cpu_multibody_at<'a>( + world: &'a mut rp::PhysicsWorld, + robot: &Robot, + qpos: &[f32], +) -> PyResult<&'a rapier3d::dynamics::Multibody> { + let rp::PhysicsWorld { + bodies, + multibody_joints, + .. + } = world; + let link_id = *multibody_joints + .rigid_body_link(robot.root) + .ok_or_else(|| PyRuntimeError::new_err("robot is not a multibody"))?; + let mb = multibody_joints + .get_multibody_mut(link_id.multibody) + .ok_or_else(|| PyRuntimeError::new_err("robot multibody not found"))?; + let current = cpu_qpos(mb, bodies); + if current.len() != qpos.len() { + return Err(PyRuntimeError::new_err(format!( + "qpos has {} entries but the robot has {} DoFs", + qpos.len(), + current.len() + ))); + } + let map = assembly_dof_map(mb, bodies); + let delta: Vec = qpos.iter().zip(¤t).map(|(q, c)| q - c).collect(); + mb.apply_displacements(&to_assembly(&delta, &map)); + mb.forward_kinematics(bodies, false); + Ok(&*mb) +} + +impl NexusState { + /// Per-link GPU slots of `robot`, in link order (empty before `finalize`). + fn link_slots(&self, robot: &Robot) -> Vec { + self.0 + .multibody_link_slots(robot.env, robot.root) + .map(|(_, _, links)| links.into_iter().map(|(_, slot)| slot).collect()) + .unwrap_or_default() + } +} diff --git a/crates/nexus_python3d/src/robot.rs b/crates/nexus_python3d/src/robot.rs new file mode 100644 index 00000000..cf661939 --- /dev/null +++ b/crates/nexus_python3d/src/robot.rs @@ -0,0 +1,291 @@ +//! Articulated robots loaded from URDF or MJCF: link and joint metadata in the +//! multibody's own order, CPU-side kinematics (forward and inverse), and the +//! per-DoF PD gains used by `NexusState.set_robot_targets`. + +use crate::math::{Pose, Quat, Vec3}; +use crate::rbd::{RigidBodyHandle, SharedShape}; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use rapier3d::prelude as rp; +use std::collections::HashMap; + +/// A robot inserted as one multibody in one environment. +/// +/// DoF order is the multibody's assembly order: links in multibody order, each +/// link's free linear axes then its free angular axes. Joint `k` is the joint +/// between link `joint_link[k]` and its parent; fixed joints are not listed. +#[pyclass(name = "Robot", from_py_object)] +#[derive(Clone)] +pub struct Robot { + #[pyo3(get)] + pub env: usize, + /// The root link's body, used to look the multibody up. + pub root: rp::RigidBodyHandle, + #[pyo3(get)] + pub link_names: Vec, + pub link_bodies: Vec, + #[pyo3(get)] + pub joint_names: Vec, + /// Index (into `link_names`) of each joint's child link. + #[pyo3(get)] + pub joint_links: Vec, + #[pyo3(get)] + pub joint_dof_offsets: Vec, + #[pyo3(get)] + pub joint_ndofs: Vec, + /// rapier axis index (`0..3` linear, `3..6` angular) of each DoF. + #[pyo3(get)] + pub dof_axes: Vec, + /// Index (into `link_names`) of the link each DoF belongs to. + #[pyo3(get)] + pub dof_links: Vec, + #[pyo3(get)] + pub dof_lower: Vec, + #[pyo3(get)] + pub dof_upper: Vec, + /// Per-DoF PD gains and force limit applied by `set_robot_targets`. + #[pyo3(get, set)] + pub kp: Vec, + #[pyo3(get, set)] + pub kv: Vec, + #[pyo3(get, set)] + pub max_force: Vec, + /// `(body, shape, local_pose)` render shapes (URDF only; MJCF visuals are + /// registered with the viewer by the loader). + #[pyo3(get)] + pub render_shapes: Vec<(RigidBodyHandle, SharedShape, Pose)>, +} + +#[pymethods] +impl Robot { + #[getter] + fn n_dofs(&self) -> usize { + self.dof_axes.len() + } + + #[getter] + fn n_links(&self) -> usize { + self.link_names.len() + } + + #[getter] + fn root_body(&self) -> RigidBodyHandle { + RigidBodyHandle(self.root) + } + + /// Index of the link named `name`. + pub fn link_index(&self, name: &str) -> PyResult { + self.link_names + .iter() + .position(|n| n == name) + .ok_or_else(|| PyValueError::new_err(format!("unknown link {name:?}"))) + } + + /// Index of the (articulated) joint named `name`. + pub fn joint_index(&self, name: &str) -> PyResult { + self.joint_names + .iter() + .position(|n| n == name) + .ok_or_else(|| PyValueError::new_err(format!("unknown joint {name:?}"))) + } + + /// The rigid body of the link named `name`. + fn link_body(&self, name: &str) -> PyResult { + Ok(RigidBodyHandle(self.link_bodies[self.link_index(name)?])) + } + + /// Rigid bodies of every link, in link order. + #[getter] + fn link_body_handles(&self) -> Vec { + self.link_bodies + .iter() + .map(|h| RigidBodyHandle(*h)) + .collect() + } + + /// DoF indices of joint `name`, in DoF order. + fn joint_dofs(&self, name: &str) -> PyResult> { + let j = self.joint_index(name)?; + Ok((self.joint_dof_offsets[j]..self.joint_dof_offsets[j] + self.joint_ndofs[j]).collect()) + } + + /// Sets the PD gains (and optionally the force limit) of the given DoFs. + #[pyo3(signature = (dofs, kp=None, kv=None, max_force=None))] + fn set_pd_gains( + &mut self, + dofs: Vec, + kp: Option>, + kv: Option>, + max_force: Option>, + ) -> PyResult<()> { + for (name, values, target) in [ + ("kp", kp, &mut self.kp), + ("kv", kv, &mut self.kv), + ("max_force", max_force, &mut self.max_force), + ] { + let Some(values) = values else { continue }; + if values.len() != dofs.len() { + return Err(PyValueError::new_err(format!( + "{name} has {} entries for {} dofs", + values.len(), + dofs.len() + ))); + } + for (d, v) in dofs.iter().zip(values) { + let slot = target + .get_mut(*d) + .ok_or_else(|| PyValueError::new_err(format!("dof {d} out of range")))?; + *slot = v; + } + } + Ok(()) + } + + fn __repr__(&self) -> String { + format!( + "Robot(env={}, links={}, dofs={})", + self.env, + self.link_names.len(), + self.dof_axes.len() + ) + } +} + +/// Free axes of a joint in rapier's assembly order (linear, then angular). +pub fn free_axes(joint: &rp::GenericJoint) -> Vec { + let locked = joint.locked_axes.bits(); + (0..6u8).filter(|axis| locked & (1 << axis) == 0).collect() +} + +/// Limits of `axis`, `(-inf, inf)` when the joint has none on that axis. +fn axis_limits(joint: &rp::GenericJoint, axis: u8) -> (f32, f32) { + let mask = rp::JointAxesMask::from_bits_truncate(1 << axis); + if joint.limit_axes.contains(mask) { + let l = &joint.limits[axis as usize]; + (l.min, l.max) + } else { + (f32::NEG_INFINITY, f32::INFINITY) + } +} + +/// Builds a [`Robot`] from the multibody containing `any_link` in `world`, using +/// the loader's names: `body_names` maps each link body to its name, and +/// `child_joint_names` maps each non-root link body to the name of the joint +/// connecting it to its parent. Gains default to the joints' current position +/// motors (zero when there is none). +pub fn build_robot( + world: &rp::PhysicsWorld, + env: usize, + any_link: rp::RigidBodyHandle, + body_names: &HashMap, + child_joint_names: &HashMap, + render_shapes: Vec<(RigidBodyHandle, SharedShape, Pose)>, +) -> PyResult { + let link_id = world + .multibody_joints + .rigid_body_link(any_link) + .ok_or_else(|| PyRuntimeError::new_err("robot root is not a multibody link"))?; + let mb = world + .multibody_joints + .get_multibody(link_id.multibody) + .ok_or_else(|| PyRuntimeError::new_err("robot multibody not found"))?; + let root_is_dynamic = world + .bodies + .get(mb.root().rigid_body_handle()) + .map(|rb| rb.is_dynamic()) + .unwrap_or(false); + + let mut robot = Robot { + env, + root: mb.root().rigid_body_handle(), + link_names: Vec::new(), + link_bodies: Vec::new(), + joint_names: Vec::new(), + joint_links: Vec::new(), + joint_dof_offsets: Vec::new(), + joint_ndofs: Vec::new(), + dof_axes: Vec::new(), + dof_links: Vec::new(), + dof_lower: Vec::new(), + dof_upper: Vec::new(), + kp: Vec::new(), + kv: Vec::new(), + max_force: Vec::new(), + render_shapes, + }; + for (link_idx, link) in mb.links().enumerate() { + let body = link.rigid_body_handle(); + robot.link_names.push( + body_names + .get(&body) + .cloned() + .unwrap_or_else(|| format!("link_{link_idx}")), + ); + robot.link_bodies.push(body); + // A fixed root has no DoFs on the GPU even though rapier models the + // root joint as free; mirror the GPU build. + let axes = if link_idx == 0 && !root_is_dynamic { + Vec::new() + } else { + free_axes(&link.joint().data) + }; + if axes.is_empty() { + continue; + } + robot.joint_names.push( + child_joint_names + .get(&body) + .cloned() + .unwrap_or_else(|| format!("joint_{link_idx}")), + ); + robot.joint_links.push(link_idx); + robot.joint_dof_offsets.push(robot.dof_axes.len()); + robot.joint_ndofs.push(axes.len()); + for axis in axes { + let (lo, hi) = axis_limits(&link.joint().data, axis); + let motor = link.joint().data.motor(joint_axis(axis)).copied(); + robot.dof_axes.push(axis); + robot.dof_links.push(link_idx); + robot.dof_lower.push(lo); + robot.dof_upper.push(hi); + robot.kp.push(motor.map(|m| m.stiffness).unwrap_or(0.0)); + robot.kv.push(motor.map(|m| m.damping).unwrap_or(0.0)); + robot + .max_force + .push(motor.map(|m| m.max_force).unwrap_or(f32::INFINITY)); + } + } + Ok(robot) +} + +/// The single-axis mask of rapier axis index `axis`. +pub fn axis_mask(axis: u8) -> rp::JointAxesMask { + rp::JointAxesMask::from_bits_truncate(1 << axis) +} + +/// The `JointAxis` of rapier axis index `axis`. +pub fn joint_axis(axis: u8) -> rp::JointAxis { + match axis { + 0 => rp::JointAxis::LinX, + 1 => rp::JointAxis::LinY, + 2 => rp::JointAxis::LinZ, + 3 => rp::JointAxis::AngX, + 4 => rp::JointAxis::AngY, + _ => rp::JointAxis::AngZ, + } +} + +/// Converts a `(w, x, y, z)` quaternion. +pub fn quat_wxyz(q: [f32; 4]) -> Quat { + Quat(glamx::Quat::from_xyzw(q[1], q[2], q[3], q[0]).normalize()) +} + +/// `(w, x, y, z)` components of a quaternion. +pub fn to_wxyz(q: glamx::Quat) -> [f32; 4] { + [q.w, q.x, q.y, q.z] +} + +/// A pose from a position and a `(w, x, y, z)` quaternion. +pub fn pose_from_wxyz(pos: [f32; 3], quat: [f32; 4]) -> rp::Pose { + rp::Pose::from_parts(Vec3(glamx::Vec3::from(pos)).0, quat_wxyz(quat).0) +} diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index d5e6983e..d4381107 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -8,9 +8,10 @@ use crate::math::Pose; use crate::math::{Vec3, Vec4}; use crate::nexus::{GpuTimestamps, NexusState}; use crate::rbd::{RigidBodyHandle, SharedShape}; +use crate::robot::{pose_from_wxyz, to_wxyz}; use khal::backend::GpuBackend; -use nexus_viewer3d::NexusViewer as RViewer; -use numpy::{IntoPyArray, PyArray3, PyArrayMethods}; +use nexus_viewer3d::{NexusViewer as RViewer, VisualTexture}; +use numpy::{IntoPyArray, PyArray2, PyArray3, PyArrayMethods}; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; @@ -182,6 +183,213 @@ impl NexusViewer { .insert_visual_shape(env, handle.0, &shape.0, local_pose.0); } + // --- sensor cameras ----------------------------------------------------- + + /// Makes later `insert_shape*` calls register one render node per body + /// instead of GPU instances. Required for bodies that sensor cameras must + /// see in their depth and segmentation passes and for + /// `set_body_segmentation_id`. Call before inserting shapes. + fn set_sensor_rendering(&mut self, enabled: bool) { + self.inner_mut().set_sensor_rendering(enabled); + } + + /// Adds an offscreen sensor camera (`width x height`, vertical field of + /// view `fov_y_deg` in degrees) and returns its id. It starts at the + /// identity pose. + #[pyo3(signature = (width, height, fov_y_deg, znear=0.01, zfar=100.0))] + fn add_sensor_camera( + &mut self, + width: u32, + height: u32, + fov_y_deg: f32, + znear: f32, + zfar: f32, + ) -> usize { + pollster::block_on(self.inner_mut().add_sensor_camera( + width, + height, + fov_y_deg.to_radians(), + znear, + zfar, + )) + } + + /// Sets sensor camera `id`'s pose: position and `(w, x, y, z)` quaternion + /// of the OpenGL camera frame (looks down its -Z axis, +Y up). + fn set_sensor_camera_pose(&mut self, id: usize, pos: [f32; 3], quat: [f32; 4]) { + self.inner_mut() + .set_sensor_camera_pose(id, pose_from_wxyz(pos, quat)); + } + + /// Sensor camera `id`'s pose as `(position, (w, x, y, z))`, OpenGL frame. + fn sensor_camera_pose(&self, id: usize) -> Option<([f32; 3], [f32; 4])> { + let pose = self.inner().sensor_camera(id)?.pose(); + let t = pose.translation; + Some(([t.x, t.y, t.z], to_wxyz(pose.rotation))) + } + + /// Attaches sensor camera `id` to body `handle` of `env` with the mount + /// pose `pos` / `quat` (`w, x, y, z`, body frame to camera frame). The + /// camera follows the body at every `sync`. + fn attach_sensor_camera( + &mut self, + id: usize, + env: u32, + handle: RigidBodyHandle, + pos: [f32; 3], + quat: [f32; 4], + state: PyRef, + ) { + self.inner_mut().attach_sensor_camera( + id, + env, + handle.0, + pose_from_wxyz(pos, quat), + &state.0, + ); + } + + /// Ambient light level of sensor camera `id`'s shaded render. + fn set_sensor_camera_ambient(&mut self, id: usize, ambient: f32) { + if let Some(sensor) = self.inner_mut().sensor_camera_mut(id) { + sensor.set_ambient(ambient); + } + } + + /// Background color of sensor camera `id`'s shaded render. + fn set_sensor_camera_background(&mut self, id: usize, rgba: [f32; 4]) { + if let Some(sensor) = self.inner_mut().sensor_camera_mut(id) { + sensor.set_background_color(rgba); + } + } + + /// Renders sensor camera `id` and returns `(rgb, depth, segmentation)`: + /// `rgb` is `(H, W, 3)` uint8, `depth` `(H, W)` float32 linear metric depth + /// (`0.0` = background), `segmentation` `(H, W)` uint32 per-body ids (`0` = + /// background). Each is `None` unless requested. + #[pyo3(signature = (id, rgb=true, depth=false, segmentation=false))] + #[allow(clippy::type_complexity)] + fn render_sensor_camera<'py>( + &mut self, + py: Python<'py>, + id: usize, + rgb: bool, + depth: bool, + segmentation: bool, + ) -> PyResult<( + Option>>, + Option>>, + Option>>, + )> { + let (w, h) = self + .inner() + .sensor_camera(id) + .map(|s| s.size()) + .ok_or_else(|| PyRuntimeError::new_err(format!("no sensor camera {id}")))?; + let (w, h) = (w as usize, h as usize); + let rgb = if rgb { + let pixels = pollster::block_on(self.inner_mut().render_sensor_rgb(id)) + .ok_or_else(|| PyRuntimeError::new_err("rgb render failed"))?; + Some(Self::to_array(py, w as u32, h as u32, pixels)?) + } else { + None + }; + let depth = if depth { + let values = self + .inner_mut() + .render_sensor_depth(id) + .ok_or_else(|| PyRuntimeError::new_err("depth render failed"))?; + Some( + values + .into_pyarray(py) + .reshape([h, w]) + .map_err(|e| PyRuntimeError::new_err(format!("{e:?}")))?, + ) + } else { + None + }; + let segmentation = if segmentation { + let ids = self + .inner_mut() + .render_sensor_segmentation(id) + .ok_or_else(|| PyRuntimeError::new_err("segmentation render failed"))?; + Some( + ids.into_pyarray(py) + .reshape([h, w]) + .map_err(|e| PyRuntimeError::new_err(format!("{e:?}")))?, + ) + } else { + None + }; + Ok((rgb, depth, segmentation)) + } + + /// Starts a new scene generation: render nodes and sensor cameras created + /// afterwards belong to it and it becomes the active one, and the previous + /// scene's nodes leave the graph for good. Use one generation per + /// `NexusState` sharing this viewer; only the latest can render. + fn begin_scene(&mut self) -> u32 { + self.inner_mut().begin_scene() + } + + /// The active (most recently begun) scene generation. + fn active_scene(&self) -> u32 { + self.inner().active_scene() + } + + /// Tags every render node of body `handle` in `env` with segmentation id + /// `id` (avoid `0`, the background). Returns the number of nodes tagged; + /// `0` means the body has no per-body node (see `set_sensor_rendering`). + fn set_body_segmentation_id(&mut self, env: u32, handle: RigidBodyHandle, id: u32) -> usize { + self.inner_mut().set_body_segmentation_id(env, handle.0, id) + } + + /// Sets the base color (RGBA) of every render node of body `handle` in `env`. + fn set_body_color(&mut self, env: u32, handle: RigidBodyHandle, rgba: [f32; 4]) { + self.inner_mut().set_body_color(env, handle.0, rgba); + } + + /// Ambient light level of the main window's shaded render. + fn set_ambient(&mut self, ambient: f32) { + self.inner_mut().set_ambient(ambient); + } + + /// Registers one render node for body `handle` in `env` drawing `shape` + /// with base color `rgba`, optional per-vertex UVs (trimesh shapes only) + /// and an optional encoded image texture (`texture` bytes, PNG/JPEG, + /// cached under `texture_name`). Unlike `insert_visual_shape`, the node is + /// never instanced, so sensor cameras see it in every pass. + #[pyo3(signature = (env, handle, shape, local_pose, rgba, uvs=None, texture=None, texture_name=None))] + #[allow(clippy::too_many_arguments)] + fn insert_sensor_shape( + &mut self, + env: u32, + handle: RigidBodyHandle, + shape: PyRef, + local_pose: Pose, + rgba: [f32; 4], + uvs: Option>, + texture: Option>, + texture_name: Option, + ) { + let name = texture_name.unwrap_or_else(|| format!("sensor-texture-{env}-{:?}", handle.0)); + let tex = match &texture { + Some(bytes) => VisualTexture::Bytes(bytes, &name), + None => VisualTexture::None, + }; + self.inner_mut().insert_visual_mesh_textured( + env, + handle.0, + &shape.0, + local_pose.0, + rgba, + uvs.as_deref(), + None, + tex, + None, + ); + } + // --- run loop --------------------------------------------------------- /// Renders one frame and processes UI/events. Returns `False` when the diff --git a/src/state.rs b/src/state.rs index 6f85438b..5ad335ce 100644 --- a/src/state.rs +++ b/src/state.rs @@ -119,6 +119,45 @@ pub struct NexusCounts { pub particles: usize, } +/// What [`NexusState::multibody_link_slots`] returns: the multibody's index in +/// its environment, its GPU descriptor, and `(link body, GPU link slot)` per +/// link in link order. +#[cfg(feature = "dim3")] +pub type MultibodySlots = ( + u32, + crate::rbd::shaders::dynamics::MultibodyInfo, + Vec<(RigidBodyHandle, u32)>, +); + +/// Failure of [`NexusState::set_multibody_joint_positions`]. +#[derive(Debug)] +pub enum JointPositionsError { + /// The GPU write failed. + Gpu(GpuBackendError), + /// `qpos` did not have one entry per multibody DoF. + DofMismatch { expected: usize, got: usize }, +} + +impl From for JointPositionsError { + fn from(e: GpuBackendError) -> Self { + JointPositionsError::Gpu(e) + } +} + +impl core::fmt::Display for JointPositionsError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + JointPositionsError::Gpu(e) => write!(f, "{e:?}"), + JointPositionsError::DofMismatch { expected, got } => { + write!( + f, + "qpos has {got} entries but the multibody has {expected} DoFs" + ) + } + } + } +} + /// High-level, GPU-resident state of a multiphysics simulation. /// /// Each sub-state (`rbd`/`mpm`) is lazily allocated the first time content @@ -523,6 +562,35 @@ impl NexusState { Ok(()) } + /// [`Self::control_multibody_motors`] for a single environment: runs `f` on + /// `env`'s rapier world, then pushes that environment's joint data (motor + /// targets and gains, limits) to the GPU. Cheaper than the all-environment + /// variant when only one environment retargets. + #[cfg(feature = "dim3")] + pub fn control_multibody_motors_env( + &mut self, + backend: &GpuBackend, + env: usize, + f: F, + ) -> Result<(), GpuBackendError> + where + F: FnOnce(&mut PhysicsWorld), + { + let Some(world) = self.rbd_envs.get_mut(env) else { + return Ok(()); + }; + 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(()) + } + /// Reads every environment's multibody link workspace back from the GPU in /// one transfer: per link, the generalized joint coordinates, accumulated /// joint rotation, world pose and world-space velocity. @@ -574,6 +642,321 @@ impl NexusState { .unwrap_or(0) } + // --- rigid-body and multibody state access (between steps) ------------- + + /// GPU pose slot of `handle` in environment `env`: the index into + /// [`Self::read_rigid_body_poses`] and [`Self::read_rigid_body_velocities`]. + /// `None` before `finalize` or for a handle this state does not know. + pub fn rigid_body_gpu_index(&self, env: usize, handle: RigidBodyHandle) -> Option { + self.rbd2gpu + .get(env)? + .get(handle.0) + .map(|r| r.gpu_id) + .filter(|id| *id != u32::MAX) + } + + /// Reads every environment's rigid-body world-origin poses back from the + /// GPU, indexed by [`Self::rigid_body_gpu_index`]. Empty before `finalize`. + pub async fn read_rigid_body_poses(&self, backend: &GpuBackend) -> Vec { + let Some(rbd) = self.rbd.as_ref() else { + return Vec::new(); + }; + let mut out = vec![crate::rbd::math::Pose::default(); rbd.body_poses().len() as usize]; + if backend + .slow_read_buffer(rbd.body_poses().buffer(), &mut out) + .await + .is_err() + { + return Vec::new(); + } + out + } + + /// Reads every environment's rigid-body world-space velocities back from + /// the GPU, indexed like [`Self::read_rigid_body_poses`]. + #[cfg(feature = "dim3")] + pub async fn read_rigid_body_velocities( + &self, + backend: &GpuBackend, + ) -> Vec { + let Some(rbd) = self.rbd.as_ref() else { + return Vec::new(); + }; + let mut out = vec![Self::zero_velocity(); rbd.vels().len() as usize]; + if backend + .slow_read_buffer(rbd.vels().buffer(), &mut out) + .await + .is_err() + { + return Vec::new(); + } + out + } + + #[cfg(feature = "dim3")] + fn zero_velocity() -> crate::rbd::shaders::dynamics::Velocity { + crate::rbd::shaders::dynamics::Velocity { + linear: crate::rbd::math::Vector::ZERO, + padding0: 0, + angular: crate::rbd::math::AngVector::ZERO, + padding1: 0, + } + } + + /// Teleports a free rigid body: overwrites its world-origin pose on the GPU. + /// Valid between steps, after `finalize`. A multibody link's pose is + /// re-derived from its joint coordinates every step, so links go through + /// [`Self::set_multibody_joint_positions`] instead. Unknown handles are + /// ignored. + pub fn set_rigid_body_pose( + &mut self, + backend: &GpuBackend, + env: usize, + handle: RigidBodyHandle, + pose: crate::rbd::math::Pose, + ) -> Result<(), GpuBackendError> { + let Some(id) = self.rigid_body_gpu_index(env, handle) else { + return Ok(()); + }; + let Some(rbd) = self.rbd.as_mut() else { + return Ok(()); + }; + backend.write_buffer(rbd.body_poses_mut().buffer_mut(), id as u64, &[pose]) + } + + /// Overwrites a rigid body's world-space linear and angular velocity on the + /// GPU. Same validity rules as [`Self::set_rigid_body_pose`]. + #[cfg(feature = "dim3")] + pub fn set_rigid_body_velocity( + &mut self, + backend: &GpuBackend, + env: usize, + handle: RigidBodyHandle, + linvel: crate::rbd::math::Vector, + angvel: crate::rbd::math::AngVector, + ) -> Result<(), GpuBackendError> { + let Some(id) = self.rigid_body_gpu_index(env, handle) else { + return Ok(()); + }; + let Some(rbd) = self.rbd.as_mut() else { + return Ok(()); + }; + let mut vel = Self::zero_velocity(); + vel.linear = linvel; + vel.angular = angvel; + backend.write_buffer(rbd.vels_mut().buffer_mut(), id as u64, &[vel]) + } + + /// Sets the rigid-body timestep of every environment: `dt` seconds per + /// `simulate` step, split into `substeps` solver substeps. Environments must + /// share one parameter set, so this applies to all of them. Takes effect at + /// the next GPU build, so call it before `finalize`. + pub fn set_rbd_timestep(&mut self, dt: f32, substeps: u32) { + for params in &mut self.rbd_sim_params { + params.dt = dt; + params.num_solver_iterations = substeps.max(1); + } + } + + /// The multibody rooted at (or containing) `body` in environment `env`, as + /// its index in the environment's multibody iteration order, the GPU + /// descriptor of that multibody, and every link's rigid body paired with its + /// per-environment GPU link slot (the stride unit of + /// [`Self::read_multibody_links`]), in link order. `None` when `body` is + /// not part of a multibody or before `finalize`. + #[cfg(feature = "dim3")] + pub fn multibody_link_slots( + &self, + env: usize, + body: RigidBodyHandle, + ) -> Option { + let world = self.rbd_envs.get(env)?; + let link_id = world.multibody_joints.rigid_body_link(body)?; + let target = world.multibody_joints.get_multibody(link_id.multibody)?; + let mb_idx = world + .multibody_joints + .multibodies() + .position(|mb| core::ptr::eq(mb, target))? as u32; + let info = self + .rbd + .as_ref()? + .multibodies() + .multibody_layout(env as u32, mb_idx)?; + let links = target + .links() + .enumerate() + .map(|(i, link)| (link.rigid_body_handle(), info.first_link + i as u32)) + .collect(); + Some((mb_idx, info, links)) + } + + /// Generalized coordinates of the multibody containing `body` in environment + /// `env`, as last written to the CPU multibody (`finalize`'s authored state + /// or the last [`Self::set_multibody_joint_positions`]): assembly order, + /// i.e. links in order, each link's free linear then angular DoFs. The GPU's + /// current coordinates are in [`Self::read_multibody_links`]. + #[cfg(feature = "dim3")] + pub fn multibody_joint_positions(&self, env: usize, body: RigidBodyHandle) -> Option> { + let world = self.rbd_envs.get(env)?; + let link_id = world.multibody_joints.rigid_body_link(body)?; + let mb = world.multibody_joints.get_multibody(link_id.multibody)?; + let root_fixed = Self::multibody_root_is_fixed(&world.bodies, mb); + let mut out = Vec::with_capacity(mb.ndofs()); + for (i, link) in mb.links().enumerate() { + if i == 0 && root_fixed { + continue; + } + Self::push_free_coords(link.joint(), &mut out); + } + Some(out) + } + + /// Whether the multibody's root body is not dynamic. rapier models such a + /// root with a free joint whose six coordinates the GPU build locks, so + /// they are excluded from the joint-position vectors. + #[cfg(feature = "dim3")] + pub fn multibody_root_is_fixed( + bodies: &crate::rapier::dynamics::RigidBodySet, + mb: &crate::rapier::dynamics::Multibody, + ) -> bool { + !bodies + .get(mb.root().rigid_body_handle()) + .map(|rb| rb.is_dynamic()) + .unwrap_or(false) + } + + /// Appends the joint's free coordinates in rapier's assembly order (linear + /// axes first, then angular) to `out`. + #[cfg(feature = "dim3")] + fn push_free_coords(joint: &crate::rapier::dynamics::MultibodyJoint, out: &mut Vec) { + let coords = joint.coords(); + let locked = joint.data.locked_axes.bits(); + for (axis, coord) in coords.iter().enumerate().take(6) { + if locked & (1 << axis) == 0 { + out.push(*coord); + } + } + } + + /// Sets the generalized coordinates of the multibody containing `body` in + /// environment `env` (same order as [`Self::multibody_joint_positions`]) and + /// zeroes its joint velocities, on the CPU multibody and on the GPU: link + /// workspaces, link body poses and velocities, and the generalized + /// velocities. Valid between steps, after `finalize`. Errors when `qpos` + /// does not have one entry per DoF; a `body` outside any multibody is + /// ignored. + #[cfg(feature = "dim3")] + pub fn set_multibody_joint_positions( + &mut self, + backend: &GpuBackend, + env: usize, + body: RigidBodyHandle, + qpos: &[f32], + ) -> Result<(), JointPositionsError> { + let Some((_, info, slots)) = self.multibody_link_slots(env, body) else { + return Ok(()); + }; + let Some(current) = self.multibody_joint_positions(env, body) else { + return Ok(()); + }; + if current.len() != qpos.len() { + return Err(JointPositionsError::DofMismatch { + expected: current.len(), + got: qpos.len(), + }); + } + let world = &mut self.rbd_envs[env]; + let PhysicsWorld { + bodies, + multibody_joints, + .. + } = world; + let link_id = *multibody_joints + .rigid_body_link(body) + .unwrap_or_else(|| unreachable!()); + let mb = multibody_joints + .get_multibody_mut(link_id.multibody) + .unwrap_or_else(|| unreachable!()); + // The displacement spans rapier's full assembly vector; a fixed root's + // (locked) coordinates stay put. + let root_fixed = Self::multibody_root_is_fixed(bodies, mb); + let mut disp = Vec::with_capacity(mb.ndofs()); + let mut next = 0usize; + for (i, link) in mb.links().enumerate() { + let ndofs = link.joint().ndofs(); + if i == 0 && root_fixed { + disp.extend(core::iter::repeat_n(0.0, ndofs)); + continue; + } + for _ in 0..ndofs { + disp.push(qpos[next] - current[next]); + next += 1; + } + } + mb.apply_displacements(&disp); + mb.forward_kinematics(bodies, false); + mb.update_rigid_bodies(bodies, false); + mb.generalized_velocity_mut().fill(0.0); + + let Some(rbd) = self.rbd.as_mut() else { + return Ok(()); + }; + let rbd2gpu = &self.rbd2gpu[env]; + for (k, (link, (handle, slot))) in mb.links().zip(&slots).enumerate() { + let rb = &bodies[*handle]; + let ltw = *link.local_to_world(); + let ltp = *link.local_to_parent(); + let world_com = ltw * rb.mass_properties().local_mprops.local_com; + // The COM shifts mirror the GPU forward-kinematics kernel: parent + // COM to the child anchor, then child anchor to the link COM. + let (shift02, shift23) = match link.parent_id() { + Some(parent_id) if k > 0 => { + let parent = mb.link(parent_id).unwrap_or_else(|| unreachable!()); + let parent_rb = &bodies[parent.rigid_body_handle()]; + let parent_com = *parent.local_to_world() + * parent_rb.mass_properties().local_mprops.local_com; + let anchor = ltw * link.joint().data.local_frame2.translation; + (anchor - parent_com, world_com - anchor) + } + _ => ( + crate::rbd::math::Vector::ZERO, + crate::rbd::math::Vector::ZERO, + ), + }; + rbd.multibodies_mut().set_link_kinematic_state( + backend, + env as u32, + *slot, + link.joint().joint_rot(), + link.joint().coords(), + ltp, + ltw, + shift02, + shift23, + world_com, + )?; + if let Some(id) = rbd2gpu + .get(handle.0) + .map(|r| r.gpu_id) + .filter(|id| *id != u32::MAX) + { + backend.write_buffer(rbd.body_poses_mut().buffer_mut(), id as u64, &[ltw])?; + backend.write_buffer( + rbd.vels_mut().buffer_mut(), + id as u64, + &[Self::zero_velocity()], + )?; + } + } + rbd.multibodies_mut().zero_dof_velocities( + backend, + env as u32, + info.first_dof, + info.ndofs, + )?; + Ok(()) + } + /// Mutable access to environment `env`'s rapier world that does **not** mark /// the rbd state dirty, for use after [`Self::finalize`]. /// diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 6bb9b173..fa6b8da7 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -784,6 +784,93 @@ impl GpuMultibodySet { Ok(()) } + /// Host copy of the descriptor of multibody `mb_idx` in `batch_id` (link + /// and DoF offsets relative to the batch), or `None` when out of range. + pub fn multibody_layout(&self, batch_id: u32, mb_idx: u32) -> Option { + self.info_mirror + .get((batch_id * self.multibodies_per_batch + mb_idx) as usize) + .copied() + } + + /// Overwrites the kinematic part of link slot `k` (per-batch link index) of + /// `batch_id`: joint rotation and generalized coordinates, local-to-parent + /// and local-to-world poses, the COM shifts and the world COM. The link's + /// joint, body and kinematic-acceleration velocities are zeroed; external + /// wrenches are untouched. This is the teleport primitive behind + /// joint-position resets: the next step's forward kinematics re-derives + /// the poses from the coordinates, the poses written here only keep + /// readbacks consistent until then. + #[cfg(feature = "dim3")] + pub fn set_link_kinematic_state( + &mut self, + backend: &GpuBackend, + batch_id: u32, + k: u32, + joint_rot: crate::math::Rotation, + coords: [f32; crate::shaders::dynamics::MAX_JOINT_DOFS], + local_to_parent: Pose, + local_to_world: Pose, + shift02: crate::math::Vector, + shift23: crate::math::Vector, + world_com: crate::math::Vector, + ) -> Result<(), GpuBackendError> { + use crate::shaders::dynamics::{ + Velocity, WS_EXT_FORCE, WS_JOINT_ROT, WS_JOINT_VEL, WS_KIN_ACC, WS_LTP, WS_LTW, + WS_QUADS, WS_RB_VELS, WS_SHIFT02, WS_SHIFT23, WS_WORLD_COM, WsAddr, ws_set_coord, + ws_set_pose, ws_set_rot, ws_set_vec, ws_set_vel, + }; + + // A link's quads are dense in the SoA buffer, so the record is built + // in a scratch one-link, one-batch view and uploaded as a prefix. + let mut quads = vec![glamx::Vec4::ZERO; WS_QUADS as usize]; + let local = WsAddr::new(0, 1, 0); + let zero_vel = Velocity { + linear: crate::math::Vector::ZERO, + padding0: 0, + angular: crate::math::AngVector::ZERO, + padding1: 0, + }; + ws_set_rot(&mut quads, local, 0, WS_JOINT_ROT, joint_rot); + for (i, c) in coords.iter().enumerate() { + ws_set_coord(&mut quads, local, 0, i as u32, *c); + } + ws_set_pose(&mut quads, local, 0, WS_LTP, local_to_parent); + ws_set_pose(&mut quads, local, 0, WS_LTW, local_to_world); + ws_set_vec(&mut quads, local, 0, WS_SHIFT02, shift02); + ws_set_vec(&mut quads, local, 0, WS_SHIFT23, shift23); + ws_set_vel(&mut quads, local, 0, WS_JOINT_VEL, zero_vel); + ws_set_vel(&mut quads, local, 0, WS_RB_VELS, zero_vel); + ws_set_vel(&mut quads, local, 0, WS_KIN_ACC, zero_vel); + ws_set_vec(&mut quads, local, 0, WS_WORLD_COM, world_com); + + let a = WsAddr::new(0, self.num_batches, batch_id); + backend.write_buffer( + self.links_workspace.buffer_mut(), + a.at(k, 0) as u64, + &quads[..WS_EXT_FORCE as usize], + ) + } + + /// Zeroes the generalized velocities of the per-batch DoFs + /// `first_dof..first_dof + ndofs` of `batch_id` (one strided write per DoF). + pub fn zero_dof_velocities( + &mut self, + backend: &GpuBackend, + batch_id: u32, + first_dof: u32, + ndofs: u32, + ) -> Result<(), GpuBackendError> { + let nb = self.num_batches as u64; + for d in first_dof..first_dof + ndofs { + backend.write_buffer( + self.dof_state.buffer_mut(), + d as u64 * nb + batch_id as u64, + &[0.0f32], + )?; + } + Ok(()) + } + /// Per-multibody descriptors (contact counts, dof offsets, ...). pub fn multibody_info(&self) -> &Tensor { &self.multibody_info diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 9d5f1255..82ad80ac 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -371,6 +371,17 @@ impl RbdState { &mut self.body_poses } + /// Per-body world-space velocities, indexed like [`Self::body_poses`]. + pub fn vels(&self) -> &Tensor { + &self.vels + } + + /// Mutable access to the per-body velocities, for teleports and resets. + /// Only valid between steps, like [`Self::body_poses_mut`]. + pub fn vels_mut(&mut self) -> &mut Tensor { + &mut self.vels + } + /// Live collision-pair count (total across all batches) most recently /// harvested by the non-blocking readback in [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers). Lags the GPU by a /// frame or two; `0` until the first readback completes. diff --git a/src_viewer/graphics.rs b/src_viewer/graphics.rs index 4c4d8be4..24df1465 100644 --- a/src_viewer/graphics.rs +++ b/src_viewer/graphics.rs @@ -53,6 +53,25 @@ pub struct RenderMaterial { pub emissive: [f32; 3], } +/// Where a visual mesh's texture comes from. +#[cfg(feature = "dim3")] +#[derive(Copy, Clone, Debug)] +pub enum VisualTexture<'a> { + /// No texture: the base color alone. + None, + /// An image file on disk, cached by path. + File(&'a Path), + /// Encoded image bytes (PNG, JPEG, ...), cached under `name`. + Bytes(&'a [u8], &'a str), +} + +#[cfg(feature = "dim3")] +impl<'a> VisualTexture<'a> { + pub fn is_some(&self) -> bool { + !matches!(self, VisualTexture::None) + } +} + /// A render-only mesh attached to a rigid body, rendered as its own kiss3d node /// (so it can carry a per-mesh texture and PBR material, which the instanced /// path can't). Its pose follows the body's world-origin pose each frame, @@ -68,6 +87,12 @@ pub struct VisualNode { pub local_pose: Pose, /// Cached GPU pose slot (resolved lazily, like the instanced entries). pub pose_index: u32, + /// Scene generation this node belongs to (see `RenderContext::generation`). + /// Nodes of an inactive generation are detached from the scene graph and + /// never re-posed. + pub generation: u32, + /// Whether the node currently hangs off the scene root. + pub attached: bool, /// Single-entry GPU descriptor for the zero-readback path (the node is drawn /// as one compute-written instance). `None` until the first direct sync. desc: Option>, @@ -280,6 +305,21 @@ pub struct RenderContext { /// instanced. Driven by body-origin poses in [`Self::update_visual_nodes`]. #[cfg(feature = "dim3")] pub visual_nodes: Vec, + /// When set, [`Self::insert_shape`] registers one kiss3d node per body + /// (a [`VisualNode`]) instead of sharing an instanced node per shape type. + /// Slower to draw, but the per-object passes (depth, segmentation) and + /// per-body ids only see non-instanced nodes. + #[cfg(feature = "dim3")] + pub prefer_visual_nodes: bool, + /// Scene generation stamped on newly inserted visual nodes. Several + /// scenes can share one viewer over a process's lifetime (a state per + /// scene); bumping the generation and activating one keeps the others' + /// nodes hidden instead of following unrelated bodies. + #[cfg(feature = "dim3")] + pub generation: u32, + /// The generation whose visual nodes are drawn and re-posed. + #[cfg(feature = "dim3")] + pub active_generation: u32, } impl RenderContext { @@ -290,6 +330,12 @@ impl RenderContext { instances: Vec::new(), #[cfg(feature = "dim3")] visual_nodes: Vec::new(), + #[cfg(feature = "dim3")] + prefer_visual_nodes: false, + #[cfg(feature = "dim3")] + generation: 0, + #[cfg(feature = "dim3")] + active_generation: 0, } } @@ -367,6 +413,22 @@ impl RenderContext { _ => Vec3::new(255.0, 127.0, 0.0) * coeff, }; let color = color.unwrap_or(Vec4::new(rgb.x, rgb.y, rgb.z, 1.0)); + #[cfg(feature = "dim3")] + if self.prefer_visual_nodes { + self.insert_visual_mesh( + scene, + env, + handle, + shape, + local_pose, + [color.x, color.y, color.z, color.w], + None, + None, + VisualTexture::None, + None, + ); + return; + } // Translucent instances must go to a dedicated node so they render in the // transparent (OIT) pass rather than the opaque one. let transparent = color.w < 1.0; @@ -787,7 +849,7 @@ impl RenderContext { color: [f32; 4], uvs: Option<&[[f32; 2]]>, normals: Option<&[[f32; 3]]>, - texture: Option<&Path>, + texture: VisualTexture<'_>, material: Option, ) { let Some(mut node) = build_visual_node(scene_3d, shape, uvs, normals, texture.is_some()) @@ -819,10 +881,16 @@ impl RenderContext { if color[3] < 1.0 { node.set_alpha_mode(kiss3d::scene::AlphaMode::Blend); } - if let Some(tex_path) = texture { - // Key the texture cache by the path so the same file uploads once. - let key = tex_path.to_string_lossy(); - node.set_texture_from_file(tex_path, &key); + match texture { + VisualTexture::None => {} + VisualTexture::File(tex_path) => { + // Key the texture cache by the path so the same file uploads once. + let key = tex_path.to_string_lossy(); + node.set_texture_from_file(tex_path, &key); + } + VisualTexture::Bytes(bytes, name) => { + node.set_texture_from_memory(bytes, name); + } } self.visual_nodes.push(VisualNode { @@ -831,12 +899,48 @@ impl RenderContext { handle: handle.0, local_pose, pose_index: u32::MAX, + generation: self.generation, + attached: true, desc: None, count_buf: None, desc_resolved: false, }); } + /// Sets the segmentation id written by the segmentation pass for every + /// visual node of body `handle` in `env`. Returns how many nodes were + /// tagged; instanced shapes cannot carry an id and are not counted. + #[cfg(feature = "dim3")] + pub fn set_body_segmentation_id( + &mut self, + env: u32, + handle: RigidBodyHandle, + id: u32, + ) -> usize { + let mut count = 0; + for visual in &mut self.visual_nodes { + if visual.env == env && visual.handle == handle.0 { + visual + .node + .apply_to_objects_mut_recursive(&mut |o| o.set_segmentation_id(id)); + count += 1; + } + } + count + } + + /// Sets the base color of every visual node of body `handle` in `env`. + #[cfg(feature = "dim3")] + pub fn set_body_color(&mut self, env: u32, handle: RigidBodyHandle, color: [f32; 4]) { + for visual in &mut self.visual_nodes { + if visual.env == env && visual.handle == handle.0 { + visual + .node + .set_color(Color::new(color[0], color[1], color[2], color[3])); + } + } + } + /// Updates each visual node's transform from the body-origin poses (indexed /// by the body's GPU pose slot). Visual-mesh local poses are body-relative, /// so they compose with `body_poses` (not the collider world poses the @@ -844,6 +948,9 @@ impl RenderContext { #[cfg(feature = "dim3")] pub fn update_visual_nodes(&mut self, state: &NexusState, body_poses: &[Pose]) { for visual in &mut self.visual_nodes { + if visual.generation != self.active_generation { + continue; // another scene's node: hidden, not re-posed + } if visual.pose_index == u32::MAX { visual.pose_index = state .rbd2gpu @@ -879,6 +986,9 @@ impl RenderContext { encoder: &mut GpuEncoder, ) -> Result<(), GpuBackendError> { for visual in &mut self.visual_nodes { + if visual.generation != self.active_generation { + continue; // another scene's node: hidden, not re-posed + } if visual.pose_index == u32::MAX { visual.pose_index = state .rbd2gpu @@ -938,6 +1048,25 @@ impl RenderContext { pub fn has_visual_nodes(&self) -> bool { !self.visual_nodes.is_empty() } + + /// Starts a new scene generation: later visual nodes belong to it and it + /// becomes the active one, and every earlier generation's nodes are + /// detached from the scene graph for good (they would otherwise follow + /// unrelated bodies of the new state). Returns the generation id. Nodes + /// are never re-attached: kiss3d draws a node re-added to the graph once + /// per re-attachment, so a superseded scene cannot be rendered again. + #[cfg(feature = "dim3")] + pub fn next_generation(&mut self) -> u32 { + self.generation += 1; + self.active_generation = self.generation; + for visual in &mut self.visual_nodes { + if visual.generation != self.active_generation && visual.attached { + visual.node.detach(); + visual.attached = false; + } + } + self.generation + } } /// Builds a kiss3d node for a visual mesh: a `RenderMesh` carrying authored diff --git a/src_viewer/lib.rs b/src_viewer/lib.rs index b608d285..780f1bc8 100644 --- a/src_viewer/lib.rs +++ b/src_viewer/lib.rs @@ -12,12 +12,16 @@ pub extern crate rapier3d as rapier; mod backend; mod graphics; +#[cfg(feature = "dim3")] +pub mod sensors; mod ui; pub mod viewer; pub use backend::BackendType; #[cfg(feature = "dim3")] -pub use graphics::RenderMaterial; +pub use graphics::{RenderMaterial, VisualTexture}; +#[cfg(feature = "dim3")] +pub use sensors::{SensorAttachment, SensorCamera, SensorCamera3d}; pub use viewer::{NexusViewer, UiState}; #[derive(Copy, Clone, PartialEq, Eq, Debug)] diff --git a/src_viewer/sensors.rs b/src_viewer/sensors.rs new file mode 100644 index 00000000..db642572 --- /dev/null +++ b/src_viewer/sensors.rs @@ -0,0 +1,228 @@ +//! Offscreen sensor cameras. +//! +//! A [`SensorCamera`] renders the viewer's 3D scene on demand into its own +//! [`OffscreenSurface`] (so its resolution is independent of the window), from +//! a [`SensorCamera3d`] whose pose is set explicitly, or follows a rigid body +//! when attached. Besides the shaded RGB image it exposes kiss3d's auxiliary +//! passes: linear metric depth and per-object segmentation ids. +//! +//! Sensors are rendering concerns, so they live in the viewer: the physics +//! state only provides body poses (read back in `NexusViewer::sync`) that the +//! attachment composes with the mount pose. + +use glamx::glam::camera::rh::proj::opengl; +use glamx::{Mat4, Pose3, Vec3}; +use kiss3d::camera::Camera3d; +use kiss3d::event::WindowEvent; +use kiss3d::prelude::Color; +use kiss3d::scene::SceneNode3d; +use kiss3d::window::{Canvas, OffscreenSurface}; +use nexus::rbd::math::Pose; +use rapier::data::Index; + +/// A camera with an explicit pose and fixed intrinsics: a vertical field of +/// view, near/far planes, and the aspect ratio of the surface it renders to. +/// It ignores window events; its owner positions it. +/// +/// The pose uses the OpenGL camera convention (the camera looks down its local +/// -Z axis with +Y up), the frame `Pose3::look_at_rh` inverts. +pub struct SensorCamera3d { + pose: Pose3, + fov_y: f32, + znear: f32, + zfar: f32, + aspect: f32, + view: Pose3, + proj: Mat4, + proj_view: Mat4, + inv_proj_view: Mat4, +} + +impl SensorCamera3d { + /// A camera at the identity pose. `fov_y` is in radians. + pub fn new(fov_y: f32, znear: f32, zfar: f32, aspect: f32) -> Self { + let mut cam = Self { + pose: Pose3::IDENTITY, + fov_y, + znear, + zfar, + aspect, + view: Pose3::IDENTITY, + proj: Mat4::IDENTITY, + proj_view: Mat4::IDENTITY, + inv_proj_view: Mat4::IDENTITY, + }; + cam.refresh(); + cam + } + + /// Sets the camera pose (OpenGL convention, see the type docs). + pub fn set_pose(&mut self, pose: Pose3) { + self.pose = pose; + self.refresh(); + } + + /// The camera pose (OpenGL convention). + pub fn pose(&self) -> Pose3 { + self.pose + } + + /// Vertical field of view, in radians. + pub fn fov_y(&self) -> f32 { + self.fov_y + } + + fn refresh(&mut self) { + self.view = self.pose.inverse(); + self.proj = opengl::perspective(self.fov_y, self.aspect, self.znear, self.zfar); + self.proj_view = self.proj * self.view.to_mat4(); + self.inv_proj_view = self.proj_view.inverse(); + } +} + +impl Camera3d for SensorCamera3d { + fn handle_event(&mut self, _canvas: &Canvas, _event: &WindowEvent) {} + + fn eye(&self) -> Vec3 { + self.pose.translation + } + + fn view_transform(&self) -> Pose3 { + self.view + } + + fn transformation(&self) -> Mat4 { + self.proj_view + } + + fn inverse_transformation(&self) -> Mat4 { + self.inv_proj_view + } + + fn clip_planes(&self) -> (f32, f32) { + (self.znear, self.zfar) + } + + fn update(&mut self, canvas: &Canvas) { + let (w, h) = canvas.size(); + let aspect = w as f32 / h.max(1) as f32; + if aspect != self.aspect { + self.aspect = aspect; + self.refresh(); + } + } + + fn view_transform_pair(&self, _pass: usize) -> (Pose3, Mat4) { + (self.view, self.proj) + } +} + +/// A rigid body a sensor camera follows: `(env, body handle, mount pose)`. The +/// camera pose is `body_pose * mount` after every viewer sync. +#[derive(Copy, Clone, Debug)] +pub struct SensorAttachment { + pub env: u32, + pub handle: Index, + pub local_pose: Pose, +} + +/// An offscreen sensor camera: its own render surface plus a fixed-intrinsics +/// camera, optionally attached to a rigid body. +pub struct SensorCamera { + surface: OffscreenSurface, + camera: SensorCamera3d, + width: u32, + height: u32, + attachment: Option, + /// Scene generation the camera belongs to; only the active generation's + /// attached cameras follow their bodies at sync. + pub generation: u32, +} + +impl SensorCamera { + /// Creates the offscreen surface (sharing the window's GPU context) and a + /// camera with the given vertical field of view (radians) and clip planes. + pub async fn new(width: u32, height: u32, fov_y: f32, znear: f32, zfar: f32) -> Self { + let surface = OffscreenSurface::new(width, height).await; + let camera = SensorCamera3d::new(fov_y, znear, zfar, width as f32 / height.max(1) as f32); + Self { + surface, + camera, + width, + height, + attachment: None, + generation: 0, + } + } + + /// Image size `(width, height)` in pixels. + pub fn size(&self) -> (u32, u32) { + (self.width, self.height) + } + + /// Vertical field of view, in radians. + pub fn fov_y(&self) -> f32 { + self.camera.fov_y() + } + + /// Sets the camera pose (OpenGL convention). Overridden at the next sync + /// while the camera is attached to a body. + pub fn set_pose(&mut self, pose: Pose) { + self.camera.set_pose(pose); + } + + /// The camera pose (OpenGL convention). + pub fn pose(&self) -> Pose { + self.camera.pose() + } + + /// Makes the camera follow a rigid body with a fixed mount pose. + pub fn attach(&mut self, env: u32, handle: Index, local_pose: Pose) { + self.attachment = Some(SensorAttachment { + env, + handle, + local_pose, + }); + } + + /// Stops following a body; the camera keeps its current pose. + pub fn detach(&mut self) { + self.attachment = None; + } + + /// The body this camera follows, if any. + pub fn attachment(&self) -> Option { + self.attachment + } + + /// Ambient light level of the shaded render. + pub fn set_ambient(&mut self, ambient: f32) { + self.surface.set_ambient(ambient); + } + + /// Background color (RGBA) of the shaded render. + pub fn set_background_color(&mut self, rgba: [f32; 4]) { + self.surface + .set_background_color(Color::new(rgba[0], rgba[1], rgba[2], rgba[3])); + } + + /// Renders the shaded scene and returns it as row-major, top-left origin + /// RGB bytes (`width * height * 3`). + pub async fn render_rgb(&mut self, scene: &mut SceneNode3d) -> Vec { + self.surface.render_3d(scene, &mut self.camera).await; + self.surface.snap_image().into_raw() + } + + /// Renders linear eye-space depth in world units, row-major with a top-left + /// origin; background pixels are `0.0`. + pub fn render_depth(&mut self, scene: &mut SceneNode3d) -> Vec { + self.surface.snap_depth_raw(scene, &mut self.camera) + } + + /// Renders the per-pixel segmentation id (`0` for background), row-major + /// with a top-left origin. Ids are the objects' `segmentation_id`s, which + /// the viewer sets per body. + pub fn render_segmentation(&mut self, scene: &mut SceneNode3d) -> Vec { + self.surface.snap_segmentation(scene, &mut self.camera) + } +} diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 32a35ebf..782f917e 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -54,6 +54,10 @@ use rapier::prelude::{RigidBodyHandle, SharedShape}; use crate::backend::BackendType; use crate::graphics::RenderContext; +#[cfg(feature = "dim3")] +use crate::graphics::VisualTexture; +#[cfg(feature = "dim3")] +use crate::sensors::SensorCamera; use crate::{DemoKind, RunState, Transition, UiSections}; /// Per-particle coloring mode for MPM rendering, written into the @@ -235,6 +239,15 @@ pub struct NexusViewer { /// frames so samples keep accumulating while the scene is static. #[cfg(feature = "dim3")] raytracer: Option, + /// Offscreen sensor cameras (see [`crate::sensors`]). While any exists, + /// `sync` takes the readback path so attached cameras can follow their body + /// and the per-object passes see world-posed nodes. + #[cfg(feature = "dim3")] + sensors: Vec, + /// Body-origin poses from the last readback sync, indexed by GPU pose slot. + /// Empty on the zero-readback path. + #[cfg(feature = "dim3")] + body_pose_cache: Vec, pub ui: UiState, } @@ -335,6 +348,10 @@ impl NexusViewer { draw_ui: true, #[cfg(feature = "dim3")] raytracer: None, + #[cfg(feature = "dim3")] + sensors: Vec::new(), + #[cfg(feature = "dim3")] + body_pose_cache: Vec::new(), ui: UiState { run_state: RunState::Paused, run_stats: RunStats::default(), @@ -456,6 +473,7 @@ impl NexusViewer { fn direct_render_path(&self) -> bool { self.webgpu_shared && !self.rt_active() + && !self.sensors_active() && self.ui.backend_type == BackendType::Gpu && matches!(self.webgpu, Some(KhalGpuBackend::WebGpu(_))) } @@ -475,6 +493,20 @@ impl NexusViewer { } } + /// Whether sensor cameras exist. Attached cameras need the CPU body poses + /// and the depth/segmentation passes need world-posed nodes, both of which + /// only the readback path provides. + fn sensors_active(&self) -> bool { + #[cfg(feature = "dim3")] + { + !self.sensors.is_empty() + } + #[cfg(not(feature = "dim3"))] + { + false + } + } + #[cfg(feature = "cuda")] fn init_cuda(&mut self) -> Option { match khal::backend::cuda::Cuda::new(0) { @@ -684,6 +716,30 @@ impl NexusViewer { normals: Option<&[[f32; 3]]>, texture: Option<&std::path::Path>, material: Option, + ) { + let texture = match texture { + Some(path) => VisualTexture::File(path), + None => VisualTexture::None, + }; + self.insert_visual_mesh_textured( + env, handle, shape, local_pose, color, uvs, normals, texture, material, + ); + } + + /// [`Self::insert_visual_mesh`] with any texture source, including encoded + /// image bytes generated at runtime. + #[cfg(feature = "dim3")] + pub fn insert_visual_mesh_textured( + &mut self, + env: u32, + handle: RigidBodyHandle, + shape: &SharedShape, + local_pose: Pose, + color: [f32; 4], + uvs: Option<&[[f32; 2]]>, + normals: Option<&[[f32; 3]]>, + texture: VisualTexture<'_>, + material: Option, ) { self.nexus_render.insert_visual_mesh( &mut self.scene3d, @@ -699,6 +755,192 @@ impl NexusViewer { ); } + // --- sensor cameras ----------------------------------------------------- + + /// Makes every later `insert_shape*` register one node per body instead of + /// GPU instances (see `RenderContext::prefer_visual_nodes`). Required for + /// bodies that sensor cameras must see in their depth and segmentation + /// passes, and for [`Self::set_body_segmentation_id`]. Call before + /// inserting the shapes. + #[cfg(feature = "dim3")] + pub fn set_sensor_rendering(&mut self, enabled: bool) { + self.nexus_render.prefer_visual_nodes = enabled; + } + + /// Adds an offscreen sensor camera of `width x height` pixels with a + /// vertical field of view `fov_y` (radians) and the given clip planes, and + /// returns its index. The camera starts at the identity pose; position it + /// with [`Self::set_sensor_camera_pose`] or [`Self::attach_sensor_camera`]. + #[cfg(feature = "dim3")] + pub async fn add_sensor_camera( + &mut self, + width: u32, + height: u32, + fov_y: f32, + znear: f32, + zfar: f32, + ) -> usize { + let mut sensor = SensorCamera::new(width, height, fov_y, znear, zfar).await; + sensor.generation = self.nexus_render.generation; + self.sensors.push(sensor); + self.sensors.len() - 1 + } + + /// Starts a new scene generation (see `RenderContext::next_generation`): + /// visual nodes and sensor cameras created afterwards belong to it, and it + /// becomes the active one; the previous scene's nodes leave the graph and + /// its attached cameras stop tracking. Call when a new simulation state is + /// built on a viewer that already showed another. + #[cfg(feature = "dim3")] + pub fn begin_scene(&mut self) -> u32 { + self.nexus_render.next_generation() + } + + /// The active (most recently begun) scene generation. + #[cfg(feature = "dim3")] + pub fn active_scene(&self) -> u32 { + self.nexus_render.active_generation + } + + /// Number of sensor cameras. + #[cfg(feature = "dim3")] + pub fn num_sensor_cameras(&self) -> usize { + self.sensors.len() + } + + /// The sensor camera at `id`. + #[cfg(feature = "dim3")] + pub fn sensor_camera(&self, id: usize) -> Option<&SensorCamera> { + self.sensors.get(id) + } + + /// The sensor camera at `id`, mutably (pose, attachment, ambient, ...). + #[cfg(feature = "dim3")] + pub fn sensor_camera_mut(&mut self, id: usize) -> Option<&mut SensorCamera> { + self.sensors.get_mut(id) + } + + /// Sets sensor camera `id`'s pose (OpenGL convention: looking down local + /// -Z, +Y up). Overridden at the next `sync` while the camera is attached. + #[cfg(feature = "dim3")] + pub fn set_sensor_camera_pose(&mut self, id: usize, pose: Pose) { + if let Some(sensor) = self.sensors.get_mut(id) { + sensor.set_pose(pose); + } + } + + /// Attaches sensor camera `id` to body `handle` of environment `env` with + /// the mount pose `local_pose` (body frame to camera frame). The camera pose + /// is refreshed from the readback body poses at every `sync`, and once + /// right away if a readback already happened. + #[cfg(feature = "dim3")] + pub fn attach_sensor_camera( + &mut self, + id: usize, + env: u32, + handle: RigidBodyHandle, + local_pose: Pose, + state: &NexusState, + ) { + if let Some(sensor) = self.sensors.get_mut(id) { + sensor.attach(env, handle.0, local_pose); + } + self.update_sensor_attachments(state); + } + + /// Body-origin pose of `handle` in `env` from the last readback `sync`. + #[cfg(feature = "dim3")] + pub fn cached_body_pose( + &self, + state: &NexusState, + env: u32, + handle: RigidBodyHandle, + ) -> Option { + let slot = state + .rbd2gpu + .get(env as usize)? + .get(handle.0) + .map(|r| r.gpu_id)?; + self.body_pose_cache.get(slot as usize).copied() + } + + /// Re-poses every attached sensor camera from the cached body poses. + #[cfg(feature = "dim3")] + fn update_sensor_attachments(&mut self, state: &NexusState) { + if self.body_pose_cache.is_empty() { + return; + } + let active = self.nexus_render.active_generation; + for sensor in &mut self.sensors { + if sensor.generation != active { + continue; + } + let Some(att) = sensor.attachment() else { + continue; + }; + let slot = state + .rbd2gpu + .get(att.env as usize) + .and_then(|env| env.get(att.handle)) + .map(|r| r.gpu_id); + if let Some(body_pose) = slot.and_then(|s| self.body_pose_cache.get(s as usize)) { + sensor.set_pose(*body_pose * att.local_pose); + } + } + } + + /// Renders sensor camera `id`'s shaded RGB image (row-major, top-left + /// origin, `width * height * 3` bytes). + #[cfg(feature = "dim3")] + pub async fn render_sensor_rgb(&mut self, id: usize) -> Option> { + let sensor = self.sensors.get_mut(id)?; + Some(sensor.render_rgb(&mut self.scene3d).await) + } + + /// Renders sensor camera `id`'s linear metric depth (`0.0` = background). + #[cfg(feature = "dim3")] + pub fn render_sensor_depth(&mut self, id: usize) -> Option> { + let sensor = self.sensors.get_mut(id)?; + Some(sensor.render_depth(&mut self.scene3d)) + } + + /// Renders sensor camera `id`'s per-pixel segmentation ids (`0` = background). + #[cfg(feature = "dim3")] + pub fn render_sensor_segmentation(&mut self, id: usize) -> Option> { + let sensor = self.sensors.get_mut(id)?; + Some(sensor.render_segmentation(&mut self.scene3d)) + } + + /// Tags every visual node of body `handle` in `env` with segmentation id + /// `id` (avoid `0`, the background). Returns the number of nodes tagged. + #[cfg(feature = "dim3")] + pub fn set_body_segmentation_id( + &mut self, + env: u32, + handle: RigidBodyHandle, + id: u32, + ) -> usize { + self.nexus_render.set_body_segmentation_id(env, handle, id) + } + + /// Sets the base color of every visual node of body `handle` in `env`. + #[cfg(feature = "dim3")] + pub fn set_body_color(&mut self, env: u32, handle: RigidBodyHandle, color: [f32; 4]) { + self.nexus_render.set_body_color(env, handle, color) + } + + /// Ambient light level of the main window's shaded render. + pub fn set_ambient(&mut self, ambient: f32) { + self.window.set_ambient(ambient); + } + + /// Adds a directional light to the 3D scene, shared by the window and + /// every sensor camera. + #[cfg(feature = "dim3")] + pub fn add_directional_light(&mut self, direction: glamx::Vec3) { + self.scene3d.add_directional_light(direction); + } + async fn sync_timestamps(&mut self, timestamps: Option<&mut GpuTimestamps>) { if let Some(timestamps) = timestamps && let Some(results) = timestamps.try_take(self.backend()) @@ -735,9 +977,10 @@ impl NexusViewer { self.nexus_render.update_instances_from_poses(state, &cache); // Body-attached visual meshes follow the body-origin poses, since - // their local poses are body-relative. + // their local poses are body-relative. Attached sensor cameras use + // the same poses. #[cfg(feature = "dim3")] - if self.nexus_render.has_visual_nodes() { + if self.nexus_render.has_visual_nodes() || !self.sensors.is_empty() { let body_poses = rbd.body_poses(); let mut body_cache = vec![Pose::default(); body_poses.len() as usize]; let _ = self @@ -745,6 +988,8 @@ impl NexusViewer { .slow_read_buffer(body_poses.buffer(), &mut body_cache) .await; self.nexus_render.update_visual_nodes(state, &body_cache); + self.body_pose_cache = body_cache; + self.update_sensor_attachments(state); } } From 7feee066478a1a7df1ee9781cdab4059e1a71002 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 18:59:54 +0200 Subject: [PATCH 04/25] feat(python): expose the rigid-body contact-solver parameters --- crates/nexus_python3d/src/nexus.rs | 83 ++++++++++++++++++++++++++++++ src/state.rs | 6 +++ 2 files changed, 89 insertions(+) diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index f5b539ea..8814fa27 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -1123,6 +1123,89 @@ impl NexusState { self.0.set_rbd_timestep(dt, substeps); } + /// Contact-solver parameters of every environment, applied at the next GPU + /// build (call before `finalize`). `None` leaves a value unchanged. + /// `contact_natural_frequency` / `contact_damping_ratio` shape the soft + /// contact model (higher frequency = stiffer contacts), the `static_*` + /// pair applies to contacts at rest, `allowed_linear_error` is the + /// tolerated penetration in length units, `max_corrective_velocity` caps + /// penetration recovery, `prediction_distance` is the contact detection + /// margin, and `internal_pgs_iterations` the multibody solver's PGS + /// iterations per substep. + #[pyo3(signature = (contact_natural_frequency=None, contact_damping_ratio=None, static_contact_natural_frequency=None, static_contact_damping_ratio=None, allowed_linear_error=None, max_corrective_velocity=None, prediction_distance=None, internal_pgs_iterations=None))] + #[allow(clippy::too_many_arguments)] + fn set_rbd_solver_params( + &mut self, + contact_natural_frequency: Option, + contact_damping_ratio: Option, + static_contact_natural_frequency: Option, + static_contact_damping_ratio: Option, + allowed_linear_error: Option, + max_corrective_velocity: Option, + prediction_distance: Option, + internal_pgs_iterations: Option, + ) { + for env in 0..self.0.num_environments() { + let Some(mut params) = self.0.rbd_sim_params(env) else { + continue; + }; + if let Some(v) = contact_natural_frequency { + params.contact_natural_frequency = v; + } + if let Some(v) = contact_damping_ratio { + params.contact_damping_ratio = v; + } + if let Some(v) = static_contact_natural_frequency { + params.static_contact_natural_frequency = v; + } + if let Some(v) = static_contact_damping_ratio { + params.static_contact_damping_ratio = v; + } + if let Some(v) = allowed_linear_error { + params.normalized_allowed_linear_error = v; + } + if let Some(v) = max_corrective_velocity { + params.normalized_max_corrective_velocity = v; + } + if let Some(v) = prediction_distance { + params.normalized_prediction_distance = v; + } + if let Some(v) = internal_pgs_iterations { + params.num_internal_pgs_iterations = v.max(1); + } + self.0.set_rbd_sim_params(env, params); + } + } + + /// The contact-solver parameters of environment 0 as a dict (see + /// `set_rbd_solver_params`), for attestation. + fn rbd_solver_params<'py>(&self, py: Python<'py>) -> PyResult> { + use pyo3::types::PyDict; + let dict = PyDict::new(py); + if let Some(p) = self.0.rbd_sim_params(0) { + dict.set_item("dt", p.dt)?; + dict.set_item("substeps", p.num_solver_iterations)?; + dict.set_item("contact_natural_frequency", p.contact_natural_frequency)?; + dict.set_item("contact_damping_ratio", p.contact_damping_ratio)?; + dict.set_item( + "static_contact_natural_frequency", + p.static_contact_natural_frequency, + )?; + dict.set_item( + "static_contact_damping_ratio", + p.static_contact_damping_ratio, + )?; + dict.set_item("allowed_linear_error", p.normalized_allowed_linear_error)?; + dict.set_item( + "max_corrective_velocity", + p.normalized_max_corrective_velocity, + )?; + dict.set_item("prediction_distance", p.normalized_prediction_distance)?; + dict.set_item("internal_pgs_iterations", p.num_internal_pgs_iterations)?; + } + Ok(dict) + } + // --- rbd config ------------------------------------------------------- fn set_rbd_steps_per_frame(&mut self, steps: u32) { diff --git a/src/state.rs b/src/state.rs index 5ad335ce..5c0cb427 100644 --- a/src/state.rs +++ b/src/state.rs @@ -509,6 +509,12 @@ impl NexusState { /// Marks the rbd state dirty so [`Self::finalize`] rebuilds with them. /// Mainly for tests that need to match an external engine's /// `IntegrationParameters` exactly (e.g. `num_solver_iterations = 1`). + /// The rigid-body simulation parameters of environment `env` (applied at + /// the next GPU build). + pub fn rbd_sim_params(&self, env: usize) -> Option { + self.rbd_sim_params.get(env).copied() + } + pub fn set_rbd_sim_params(&mut self, env: usize, params: RbdSimParams) { self.rbd_sim_params[env] = params; self.rbd_dirty = true; From 31b4f2665c71604c6ea63fc749c134a04dbec1de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 27 Aug 2026 20:25:52 +0200 Subject: [PATCH 05/25] feat(python): expose per-step run statistics --- crates/nexus_python3d/src/nexus.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 8814fa27..9d5b9792 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -1117,6 +1117,28 @@ impl NexusState { .map_err(gpu_err) } + /// Timing and solver statistics of the last `simulate` call as a dict: + /// `encoding_time_ms` (CPU command encoding), `gpu_total_time_ms` and + /// `gpu_pass_times` (`{pass label: ms}`) from the GPU timestamp queries + /// (only populated when a `GpuTimestamps` is passed to `simulate` and + /// harvested by `viewer.sync`, which lags a frame or two), plus the + /// constraint-coloring `num_colors` / `coloring_iterations`. + fn run_stats<'py>(&self, py: Python<'py>) -> PyResult> { + use pyo3::types::PyDict; + let stats = &self.0.run_stats; + let dict = PyDict::new(py); + dict.set_item("encoding_time_ms", stats.encoding_time_ms())?; + dict.set_item("gpu_total_time_ms", stats.gpu_total_time_ms)?; + let passes = PyDict::new(py); + for (label, ms) in &stats.gpu_pass_times { + passes.set_item(label, *ms)?; + } + dict.set_item("gpu_pass_times", passes)?; + dict.set_item("num_colors", stats.num_colors)?; + dict.set_item("coloring_iterations", stats.coloring_iterations)?; + Ok(dict) + } + /// Rigid-body timestep of every environment: `dt` seconds per `simulate` /// step, in `substeps` solver substeps. Call before `finalize`. fn set_rbd_timestep(&mut self, dt: f32, substeps: u32) { From 8d0433c4262e92c19b959543e2dac08f333950d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 16:25:29 +0200 Subject: [PATCH 06/25] fix(rbd): solve multibody contacts once and give friction the full PGS budget --- CHANGELOG.md | 27 ++++++ .../dynamics/multibody/multibody_solver.rs | 31 ++++-- src_rbd/dynamics/solver.rs | 96 +++++++++++++------ src_rbd/pipeline/insertion_removal.rs | 4 + src_rbd/pipeline/rbd_state.rs | 16 +++- src_rbd/pipeline/rbd_state_from_rapier.rs | 19 ++++ src_rbd/pipeline/rbd_step.rs | 10 +- .../dynamics/multibody/contact_constraints.rs | 10 +- .../dynamics/multibody/solve_constraints.rs | 37 ++++--- src_rbd_shaders/dynamics/sim_params.rs | 40 +++++++- src_rbd_shaders/dynamics/solver.rs | 85 +++++++++------- src_rbd_shaders/dynamics/solver_utils.rs | 10 +- 12 files changed, 282 insertions(+), 103 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fe769985..4a3706d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,30 @@ +## Unreleased + +### Added + +- `RbdSimParams::friction_in_bias_pass`: also solve the contact friction rows during the biased + PGS pass, instead of only during the once-per-substep stabilization sweep (rapier's + `friction_in_bias_pass`). Off by default; on, friction gets as many iterations as the normal + rows, which keeps pinch grasps and resting stacks from drifting. +- Python: `NexusState.set_rbd_solver_params(friction_in_bias_pass=...)`, + `ColliderBuilder.friction_combine_rule`, and the `debug_contacts` / + `debug_multibody_contact_impulses` GPU readbacks for contact diagnostics. + +### Fixed + +- Contacts between a multibody link and a rigid body were solved twice: by the multibody + contact solver and, again, by the rigid-body solver against a zero-inverse-mass copy of the + link. The second copy saw the link as never moving, so a robot lifting a grasped object had + its friction cancelled by a "ghost" of its own fingers. The rigid-body constraint builder now + leaves multibody-owned manifolds to the multibody solver. + +### Modified + +- `RbdSimParams::num_internal_pgs_iterations` now also drives the rigid-body contact and joint + sweeps of the biased pass, interleaved one iteration at a time with the multibody sweeps (it + used to loop the multibody solver alone, leaving a box pinched by a robot against a board with + eight multibody iterations against one rigid-body iteration per substep). + ## v0.5.0 (16 August 2026) ### Breaking changes diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 90fee1af..c3a44f13 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -4,6 +4,7 @@ use super::multibody_set::*; use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::broad_phase::ContactPlan; +use crate::shaders::dynamics::{BIAS_MODE_BIAS, BIAS_MODE_BIAS_FRICTION}; use crate::shaders::dynamics::{ GpuMbApplyContactRestitution, GpuMbBuildContactDelassus, GpuMbComputeDynamicsPre, GpuMbComputeSolveBounds, GpuMbConsOffsetsScan, GpuMbCountContactConstraints, GpuMbDelayTick, @@ -132,6 +133,9 @@ pub struct MultibodySolverArgs<'a> { /// Per-color-index uniform tensors (`color_uniforms[c]` holds `c`), /// shared with the contact/joint solvers. pub color_uniforms: &'a [Tensor], + /// Solve friction rows during the biased pass too + /// (`RbdSimParams::friction_in_bias_pass`). + pub friction_in_bias_pass: bool, /// GPU-written workgroup grid for the per-multibody contact-constraint /// dispatches: `[multibodies_batch_capacity, num_batches, 1]`. pub mb_sweep_indirect: &'a Tensor<[u32; 3]>, @@ -576,29 +580,36 @@ impl GpuMultibodySolver { } /// P3: one PGS iteration with bias over the joint, contact, and multibody- - /// touching impulse-joint constraints. + /// touching impulse-joint constraints. The rigid-body solver interleaves + /// its own sweep after each call, `num_internal_pgs_iterations` times per + /// substep; `first_iteration` (re)builds the impulse-joint constraints. pub fn substep_solve_with_bias( &self, pass: &mut GpuPass, mb: &mut GpuMultibodySet, args: &mut MultibodySolverArgs<'_>, + first_iteration: bool, ) -> Result<(), GpuBackendError> { if mb.is_empty() { return Ok(()); } // One 64-lane workgroup per multibody with the generalized velocities - // held in workgroup memory (`color_uniforms[1]` holds the constant - // 1 = use_bias). With the Delassus blocks allocated, the contact half - // runs in constraint space instead (joints first, same order). + // held in workgroup memory (`color_uniforms[bias_mode]` holds the + // bias-mode constant, see `decode_bias_mode`). With the Delassus blocks + // allocated, the contact half runs in constraint space instead (joints + // first, same order). + let bias_mode = if args.friction_in_bias_pass { + BIAS_MODE_BIAS_FRICTION as usize + } else { + BIAS_MODE_BIAS as usize + }; let solve_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; - for _ in 0..mb.num_internal_pgs_iterations() { - self.dispatch_solve(pass, mb, args, solve_dispatch, 1)?; - } + self.dispatch_solve(pass, mb, args, solve_dispatch, bias_mode)?; // Multibody-touching impulse joints — generic (rb-mb / mb-mb) - // constraints. - if mb.mb_imp_joints_per_batch > 0 { + // constraints, built on the first iteration of the substep. + if mb.mb_imp_joints_per_batch > 0 && first_iteration { // Flat 1-D sweep over the interleaved joint slots. let imp_dispatch = [mb.mb_imp_joints_per_batch * mb.num_batches, 1, 1]; self.update_impulse_joint_constraints.call( @@ -630,6 +641,8 @@ impl GpuMultibodySolver { &mb.lu_pivots, &mb.links_static, )?; + } + if mb.mb_imp_joints_per_batch > 0 { // Colored PGS iteration WITH bias: one dispatch per color, each // color's joints solved race-free in parallel (graph coloring // done at init in `set_impulse_joints`). diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 06076903..70221443 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -12,6 +12,7 @@ use crate::queries::GpuIndexedContact; use crate::shaders::broad_phase::ContactPlan; #[cfg(feature = "dim3")] use crate::shaders::dynamics::MbContactIndexEntry; +use crate::shaders::dynamics::{BIAS_MODE_BIAS, BIAS_MODE_BIAS_FRICTION}; use crate::shaders::dynamics::{ GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized, GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, GpuSolverInitConstraints, @@ -139,8 +140,14 @@ pub struct SolverArgs<'a> { pub prefix_sum: &'a GpuPrefixSum, /// Number of solver iterations (max across all environments). pub num_solver_iterations: u32, + /// PGS iterations of the biased pass per substep, shared by the + /// rigid-body and multibody sweeps (`RbdSimParams::num_internal_pgs_iterations`). + pub num_internal_pgs_iterations: u32, /// Per-body graph-coloring group id (multibody-aware). pub body_group: &'a Tensor, + /// Per-body flag, 1 for multibody links, whose contacts the rigid-body + /// pipeline leaves to the multibody solver. + pub body_is_multibody: &'a Tensor, /// When `true` (no multibody in the scene), warmstart uses the single /// gather-per-body dispatch instead of one scatter dispatch per color. /// The gather variant looks bodies up by their own id, which is only @@ -154,6 +161,9 @@ pub struct SolverArgs<'a> { pub fused_color_sweeps: bool, /// `true` when every rigid-body contact constraint is provably a no-op. pub rb_contacts_inert: bool, + /// Solve friction rows during the biased pass too + /// (`RbdSimParams::friction_in_bias_pass`). + pub friction_in_bias_pass: bool, /// Shared per-batch indices. pub batch_indices: &'a Tensor, /// The one gravity uniform every rigid-body and multibody kernel reads. @@ -208,6 +218,7 @@ impl GpuSolver { args.constraints, args.constraint_builders, args.contact_plan, + args.body_is_multibody, args.collider_world_poses, args.solver_body_poses, args.vels, @@ -222,7 +233,7 @@ impl GpuSolver { args.contacts_len_indirect, args.contacts, args.body_constraint_counts, - args.body_group, + args.body_is_multibody, args.mprops, args.contact_plan, )?; @@ -245,7 +256,7 @@ impl GpuSolver { args.contacts, args.contact_plan, args.body_constraint_ids, - args.body_group, + args.body_is_multibody, )?; Ok(()) @@ -311,6 +322,7 @@ impl GpuSolver { gravity: args.gravity, color_uniforms: args.color_uniforms, mb_sweep_indirect: args.mb_sweep_indirect, + friction_in_bias_pass: args.friction_in_bias_pass, }; solver.layout_contact_constraints(&mut pass, state, &mut mb_args)?; } @@ -338,12 +350,20 @@ impl GpuSolver { gravity: args.gravity, color_uniforms: args.color_uniforms, mb_sweep_indirect: args.mb_sweep_indirect, + friction_in_bias_pass: args.friction_in_bias_pass, }; solver.$method(&mut pass, state, &mut mb_args $(, $extra)*)?; } }}; } + // Bias-mode uniform of the biased pass (see `decode_bias_mode`). + let bias_mode = if args.friction_in_bias_pass { + BIAS_MODE_BIAS_FRICTION as usize + } else { + BIAS_MODE_BIAS as usize + }; + for substep_id in 0..num_substeps { let is_last_substep = substep_id == num_substeps - 1; // Only consumed by the dim3-only multibody phases. @@ -385,6 +405,7 @@ impl GpuSolver { gravity: args.gravity, color_uniforms: args.color_uniforms, mb_sweep_indirect: args.mb_sweep_indirect, + friction_in_bias_pass: args.friction_in_bias_pass, }; solver.substep_build_constraints( encoder, @@ -456,43 +477,56 @@ impl GpuSolver { } /* - * Solve all joints + contacts with bias. + * Solve all joints + contacts with bias. The multibody and + * rigid-body sweeps interleave, one iteration each, so a body + * squeezed between a multibody link and a rigid body sees both + * sides converge at the same rate. */ - mb_phase!("[RBD] slv/mb-solve-bias", substep_solve_with_bias); - if !skip_rb || !joints_empty { - let mut pass = - encoder.begin_pass("[RBD] slv/rb-solve-bias", timestamps.as_deref_mut()); - let pass = &mut pass; - joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?; - if skip_rb { - // Contact sweeps skipped (inert constraints). - } else if args.fused_color_sweeps { - self.step_gauss_seidel_fused.call( - pass, - [64, args.num_batches, 1], - args.constraints, - args.solver_vels, - args.color_buckets, - args.color_sorted_ids, - &args.color_uniforms[args.num_colors as usize], - args.batch_indices, - // use_bias = 1 (the `color_uniform[1]` contains the value 1) - &args.color_uniforms[1], - )?; - } else { - for c in 1..=args.num_colors { - self.step_gauss_seidel.call( + for iteration in 0..args.num_internal_pgs_iterations.max(1) { + let first_iteration = iteration == 0; + // Only consumed by the dim3-only multibody phase. + #[cfg(not(feature = "dim3"))] + let _ = first_iteration; + mb_phase!( + "[RBD] slv/mb-solve-bias", + substep_solve_with_bias, + first_iteration + ); + if !skip_rb || !joints_empty { + let mut pass = + encoder.begin_pass("[RBD] slv/rb-solve-bias", timestamps.as_deref_mut()); + let pass = &mut pass; + joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?; + if skip_rb { + // Contact sweeps skipped (inert constraints). + } else if args.fused_color_sweeps { + self.step_gauss_seidel_fused.call( pass, - args.contacts_len_indirect, + [64, args.num_batches, 1], args.constraints, args.solver_vels, args.color_buckets, args.color_sorted_ids, - &args.color_uniforms[c as usize], + &args.color_uniforms[args.num_colors as usize], args.batch_indices, - // use_bias = 1 (the `color_uniform[1]` contains the value 1) - &args.color_uniforms[1], + // Biased pass: `color_uniforms[bias_mode]` holds `bias_mode`. + &args.color_uniforms[bias_mode], )?; + } else { + for c in 1..=args.num_colors { + self.step_gauss_seidel.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.solver_vels, + args.color_buckets, + args.color_sorted_ids, + &args.color_uniforms[c as usize], + args.batch_indices, + // Biased pass: `color_uniforms[bias_mode]` holds `bias_mode`. + &args.color_uniforms[bias_mode], + )?; + } } } } diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 80b3d2b9..455ffadf 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -110,6 +110,9 @@ impl RbdState { Tensor::vector(backend, [Point::ZERO.into()], BufferUsages::STORAGE).unwrap(); let index_buffers = Tensor::vector(backend, [0u32, 0, 0], BufferUsages::STORAGE).unwrap(); let body_group = Tensor::vector(backend, &all_body_group, BufferUsages::STORAGE).unwrap(); + // Appended bodies are free bodies: no multibody link flag is ever set. + let body_is_multibody = + Tensor::vector(backend, vec![0u32; num_bodies_total], BufferUsages::STORAGE).unwrap(); // Per-body buffers carry COPY_DST | COPY_SRC so `append_bodies` / // `remove_bodies` can write / relocate slots in place. @@ -256,6 +259,7 @@ impl RbdState { multibodies, gravity: Self::gravity_tensor(backend, [0.0, -9.81, 0.0]), body_group, + body_is_multibody, local_mprops: Tensor::vector(backend, &all_local_mprops, rw).unwrap(), mprops: Tensor::vector(backend, &all_mprops, rw).unwrap(), body_poses: Tensor::vector(backend, &all_poses, rw).unwrap(), diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 82ad80ac..94eedd23 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -263,6 +263,10 @@ pub struct RbdState { /// contacts touching different bodies of the same multibody can never be /// assigned the same color. pub(super) body_group: Tensor, + /// Per-body flag, 1 for the links of a multibody. The rigid-body contact + /// pipeline skips every manifold touching such a body: the multibody + /// solver owns those contacts. + pub(super) body_is_multibody: Tensor, pub(super) prefix_sum_workspace: PrefixSumWorkspace, /// Separate workspace for the color-bucket prefix scan (different length /// than the body-count scan, so sharing one workspace would thrash its @@ -428,8 +432,8 @@ impl RbdState { self.sim_params_cpu = params; } - /// Sets how many PGS iterations the multibody solver's biased pass runs per - /// substep, without rebuilding the GPU state. + /// Sets how many PGS iterations the biased pass runs per substep (rigid-body + /// and multibody sweeps alike), without rebuilding the GPU state. #[cfg(feature = "dim3")] pub fn set_num_internal_pgs_iterations(&mut self, backend: &GpuBackend, n: u32) { let n = n.max(1); @@ -442,7 +446,7 @@ impl RbdState { self.sim_params_cpu = params; } - /// PGS iterations per substep in the multibody solver's biased pass. + /// PGS iterations per substep in the biased pass. #[cfg(feature = "dim3")] pub fn num_internal_pgs_iterations(&self) -> u32 { self.multibodies.num_internal_pgs_iterations() @@ -533,6 +537,12 @@ impl RbdState { &self.contacts } + /// GPU buffer holding the rigid-body contact constraints of the current + /// step (impulses included). For debugging. + pub fn rigid_contact_constraints(&self) -> &Tensor { + &self.new_constraints + } + /// Debug: read back active contacts as `(collider_a, collider_b, body_a, /// body_b, manifold_len)` tuples (only `len > 0` entries). pub fn debug_contact_pairs(&self, backend: &GpuBackend) -> Vec<(u32, u32, u32, u32, u32)> { diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 9a578ecf..42cba532 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -594,6 +594,21 @@ impl RbdState { } } + // Per-body multibody-link flag (batch-major here, interleaved below). + #[cfg_attr(not(feature = "dim3"), allow(unused_mut))] + let mut all_body_is_mb: Vec = vec![0; max_colliders * num_batches as usize]; + #[cfg(feature = "dim3")] + for (batch_idx, (mb_set, body_ids, _)) in multibody_envs.iter().enumerate() { + let base = batch_idx * max_colliders; + for mb in mb_set.multibodies() { + for link in mb.links() { + if let Some(&local) = body_ids.get(&link.rigid_body_handle()) { + all_body_is_mb[base + local as usize] = 1; + } + } + } + } + let num_colliders_per_batch = max_colliders; let num_bodies_total = num_colliders_per_batch * num_batches as usize; @@ -636,6 +651,9 @@ impl RbdState { let all_collider_materials = interleave_batches(&all_collider_materials, nb); let all_body_group = interleave_batches(&all_body_group, nb); let body_group = Tensor::vector(backend, &all_body_group, BufferUsages::STORAGE).unwrap(); + let all_body_is_mb = interleave_batches(&all_body_is_mb, nb); + let body_is_multibody = + Tensor::vector(backend, &all_body_is_mb, BufferUsages::STORAGE).unwrap(); // Initial body velocities were accumulated in body-slot order alongside // `all_poses`; zero-filling here would silently drop each body's initial @@ -800,6 +818,7 @@ impl RbdState { multibodies, gravity: RbdState::gravity_tensor(backend, [0.0, -9.81, 0.0]), body_group, + body_is_multibody, local_mprops: Tensor::vector(backend, &all_local_mprops, storage).unwrap(), mprops: Tensor::vector(backend, &all_mprops, storage).unwrap(), body_poses: Tensor::vector( diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index c15c50d7..aa7b59a3 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -147,7 +147,8 @@ impl RbdPipeline { // Make sure the color index uniforms are up-to-date. // This is the maximum over the colors needed for contacts, joints, and multibodies. { - let mut needed = state.max_colors + 2; + // At least 3: the bias-mode constants 0..=2 (see `decode_bias_mode`). + let mut needed = (state.max_colors + 2).max(3); needed = needed.max(state.joints.num_colors() + 1); #[cfg(feature = "dim3")] { @@ -173,6 +174,7 @@ impl RbdPipeline { color_uniforms: &state.color_uniforms, mb_sweep_indirect: &state.mb_sweep_indirect, gravity: &state.gravity, + friction_in_bias_pass: state.sim_params_cpu.friction_in_bias_pass != 0, }; self.multibody_solver.init_step( &mut *encoder, @@ -369,12 +371,15 @@ impl RbdPipeline { num_batches: state.num_batches, num_colliders: state.num_colliders_per_batch, num_solver_iterations: state.num_solver_iterations, + num_internal_pgs_iterations: state.sim_params_cpu.num_internal_pgs_iterations, body_group: &state.body_group, + body_is_multibody: &state.body_is_multibody, batch_indices: &state.batch_indices, mb_sweep_indirect: &state.mb_sweep_indirect, colorless_warmstart: false, fused_color_sweeps, rb_contacts_inert: state.rb_contacts_inert, + friction_in_bias_pass: state.sim_params_cpu.friction_in_bias_pass != 0, gravity: &state.gravity, }; self.solver.prepare( @@ -520,7 +525,9 @@ impl RbdPipeline { num_batches: state.num_batches, num_colliders: state.num_colliders_per_batch, num_solver_iterations: state.num_solver_iterations, + num_internal_pgs_iterations: state.sim_params_cpu.num_internal_pgs_iterations, body_group: &state.body_group, + body_is_multibody: &state.body_is_multibody, batch_indices: &state.batch_indices, mb_sweep_indirect: &state.mb_sweep_indirect, // The gather warmstart is only valid without multibody grouping; @@ -531,6 +538,7 @@ impl RbdPipeline { colorless_warmstart: true, fused_color_sweeps, rb_contacts_inert: state.rb_contacts_inert, + friction_in_bias_pass: state.sim_params_cpu.friction_in_bias_pass != 0, gravity: &state.gravity, }; diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index fde27b7d..420e0f09 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -726,10 +726,12 @@ pub fn gpu_mb_init_contact_constraints( } }; - // Positional bias along the tangent: pull the two anchors back - // together so friction sticks instead of drifting. No surface - // velocity yet (TODO: conveyor belts), so `rhs_wo_bias` is 0. - let tang_bias = (p1 - p2).dot(mb_tangent) * inv_dt; + // Friction stays velocity-level, as in rapier: no positional + // anchoring term. A rigid tangent anchor over-constrains a + // pinch pressed against a resting surface (fingers, box and + // board all welded together) and locks up the grasp. Surface + // velocity (conveyor belts) would go here; `rhs_wo_bias` is 0. + let tang_bias = 0.0f32; #[cfg(feature = "dim3")] let tang_cons = MultibodyContactConstraint { multibody_id: mb_idx, diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index 7f22a835..d1140b97 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -8,6 +8,7 @@ use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::sync::workgroup_memory_barrier_with_group_sync; use crate::dynamics::body::Velocity; +use crate::dynamics::decode_bias_mode; use crate::gdot; use crate::utils::BatchIndices; use crate::utils::linalg::MAX_MB_DOFS; @@ -117,7 +118,7 @@ pub fn gpu_mb_solve_constraints( if ndofs == 0 { return; } - let use_bias = *use_bias != 0; + let (use_bias, solve_friction) = decode_bias_mode(*use_bias); let v_base = mb.first_dof as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; @@ -232,14 +233,15 @@ pub fn gpu_mb_solve_constraints( }; let cons = contact_constraints.read(cons_idx); let is_tangent = cons.kind == MB_CONTACT_KIND_TANGENT; - // Friction is only solved during the relaxation phase, and in 3D a - // tangent pair is solved by its first row only. + // Friction is solved during the relaxation phase (and during the + // biased pass when `friction_in_bias_pass` is set); in 3D a tangent + // pair is solved by its first row only. #[cfg(feature = "dim3")] let solve = slot_active - && !(use_bias && is_tangent) + && !(is_tangent && !solve_friction) && !(is_tangent && s != cons.normal_constraint_slot + 1); #[cfg(feature = "dim2")] - let solve = slot_active && !(use_bias && is_tangent); + let solve = slot_active && !(is_tangent && !solve_friction); #[cfg(not(feature = "web-compat"))] if !solve { continue; @@ -310,7 +312,12 @@ pub fn gpu_mb_solve_constraints( } } - let cfm_factor = if use_bias { cons.cfm_factor } else { 1.0 }; + // Compliance only softens the normal rows; friction stays rigid. + let cfm_factor = if use_bias && !is_tangent { + cons.cfm_factor + } else { + 1.0 + }; let impulse0 = cons.impulse; let rhs0 = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; let raw0 = cfm_factor * (impulse0 - cons.inv_lhs * (j_dot_v0 + rhs0)); @@ -636,7 +643,7 @@ pub fn gpu_mb_solve_contacts_delassus( return; } let active = in_range && ndofs != 0 && count != 0; - let use_bias = *use_bias != 0; + let (use_bias, solve_friction) = decode_bias_mode(*use_bias); let v_base = mb.first_dof as usize; let cons_base = mb.contact_constraint_start as usize; @@ -667,7 +674,9 @@ pub fn gpu_mb_solve_contacts_delassus( if use_bias { cons.rhs } else { cons.rhs_wo_bias }, ); inv_lhs_shared.write(s as usize, cons.inv_lhs); - cfm_shared.write(s as usize, if use_bias { cons.cfm_factor } else { 1.0 }); + // Compliance only softens the normal rows; friction stays rigid. + let soft = use_bias && cons.kind != MB_CONTACT_KIND_TANGENT; + cfm_shared.write(s as usize, if soft { cons.cfm_factor } else { 1.0 }); friction_shared.write(s as usize, cons.friction_coeff); let is_self = cons.free_body_id == u32::MAX; let free_active = !is_self @@ -723,13 +732,15 @@ pub fn gpu_mb_solve_contacts_delassus( let free_active = (meta >> 24) != 0; let is_tangent = kind == MB_CONTACT_KIND_TANGENT; - // Friction is only solved during the stabilization sweep, and in 3D a - // tangent pair is solved by its first row only. + // Friction is solved during the stabilization sweep (and during the + // biased pass when `friction_in_bias_pass` is set); in 3D a tangent + // pair is solved by its first row only. #[cfg(feature = "dim3")] - let solve = - slot_active && !(use_bias && is_tangent) && !(is_tangent && s != normal_slot + 1); + let solve = slot_active + && !(is_tangent && !solve_friction) + && !(is_tangent && s != normal_slot + 1); #[cfg(feature = "dim2")] - let solve = slot_active && !(use_bias && is_tangent); + let solve = slot_active && !(is_tangent && !solve_friction); #[cfg(not(feature = "web-compat"))] if !solve { continue; diff --git a/src_rbd_shaders/dynamics/sim_params.rs b/src_rbd_shaders/dynamics/sim_params.rs index 93c70361..f0825eed 100644 --- a/src_rbd_shaders/dynamics/sim_params.rs +++ b/src_rbd_shaders/dynamics/sim_params.rs @@ -69,6 +69,21 @@ impl ConstraintSoftness { } } +/// Bias-mode uniform handed to the constraint solve kernels: the unbiased +/// stabilization sweep, the biased pass with friction rows skipped (rapier's +/// default scheduling), or the biased pass with friction rows solved +/// (`RbdSimParams::friction_in_bias_pass`). The values double as indices into +/// the solver's constant uniforms (`color_uniforms[c] == c`). +pub const BIAS_MODE_NONE: u32 = 0; +pub const BIAS_MODE_BIAS: u32 = 1; +pub const BIAS_MODE_BIAS_FRICTION: u32 = 2; + +/// Decodes a bias-mode uniform into `(use_bias, solve_friction)`. +#[inline(always)] +pub fn decode_bias_mode(mode: u32) -> (bool, bool) { + (mode != BIAS_MODE_NONE, mode != BIAS_MODE_BIAS) +} + /// Parameters for a time-step of the physics engine. #[derive(Clone, Copy, PartialEq)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] @@ -168,11 +183,28 @@ pub struct RbdSimParams { /// cheaper, but one averaged normal then stands in for a ridge or a step. pub contact_merge_cos: f32, - /// Multibody only: PGS iterations over the joint + contact constraints run - /// per substep, in the biased pass (default: `1`). + /// PGS iterations over the joint + contact constraints run per substep in + /// the biased pass (default: `1`). The rigid-body and multibody sweeps + /// interleave, one iteration each, so both sides of a rigid-body/multibody + /// contact converge at the same rate. /// /// Host-side only: it is a dispatch count, never read by a shader. pub num_internal_pgs_iterations: u32, + + /// Nonzero: friction rows are also solved during the biased pass instead + /// of only during the unbiased stabilization sweep (default: `0`, matching + /// rapier's `friction_in_bias_pass`). Turning it on gives friction as many + /// PGS iterations as the normal rows, which stiffens grasps and resting + /// contacts at a small cost per iteration. + /// + /// Host-side only: it selects the bias-mode uniform passed to the solve + /// kernels, never read by a shader. + pub friction_in_bias_pass: u32, + // Uniform-layout padding to a 16-byte multiple (scalars: an array member + // here would itself need 16-byte alignment). + pub _padding0: u32, + pub _padding1: u32, + pub _padding2: u32, } impl RbdSimParams { @@ -198,6 +230,10 @@ impl RbdSimParams { normalized_max_linear_velocity: 400.0, length_unit: 1.0, num_internal_pgs_iterations: 1, + friction_in_bias_pass: 0, + _padding0: 0, + _padding1: 0, + _padding2: 0, } } } diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index cebfbd89..01b22c05 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -15,7 +15,7 @@ use khal_std::{ use super::body::{LocalMassProperties, Velocity, WorldMassProperties}; use super::constraint::{TwoBodyConstraint, TwoBodyConstraintBuilder}; -use super::sim_params::RbdSimParams; +use super::sim_params::{RbdSimParams, decode_bias_mode}; use super::solver_utils::warmstart_body; use crate::queries::IndexedManifold; @@ -39,6 +39,7 @@ pub fn gpu_solver_init_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraint_builders: &mut [TwoBodyConstraintBuilder], #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] body_is_multibody: &[u32], #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] collider_world_poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_body_poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 1, binding = 2)] vels: &[Velocity], @@ -52,13 +53,18 @@ pub fn gpu_solver_init_constraints( let solver_body_poses = Slice(solver_body_poses, 0); let vels = Slice(vels, 0); let mprops = Slice(mprops, 0); + let body_is_multibody = Slice(body_is_multibody, 0); for i in StepRng::new(invocation_id.x..total, num_threads) { let i = i as usize; let im = contacts.at(i); - if im.contact.len == 0 { - // Gap or inert slot: clear the (stale) constraint so every flat - // consumer skips it. + // Manifolds touching a multibody link belong to the multibody contact + // solver (`gpu_mb_init_contact_constraints`); building a rigid-body + // constraint for them too would solve the contact twice, the second + // time against a zero-inverse-mass copy of the link that never moves. + if im.contact.len == 0 || touches_multibody(im, &body_is_multibody) { + // Gap, inert or multibody-owned slot: clear the (stale) + // constraint so every flat consumer skips it. constraints.at_mut(i).len = 0; continue; } @@ -83,7 +89,7 @@ pub fn gpu_solver_count_constraints( #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &[IndexedManifold], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_constraint_counts: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_group: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_is_multibody: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mprops: &[WorldMassProperties], #[spirv(uniform, descriptor_set = 0, binding = 4)] contact_plan: &ContactPlan, ) { @@ -93,7 +99,7 @@ pub fn gpu_solver_count_constraints( let total = contact_plan.bound; let contacts = Slice(contacts, 0); let mut body_constraint_counts = SliceMut(body_constraint_counts, 0); - let body_group = Slice(body_group, 0); + let body_is_multibody = Slice(body_is_multibody, 0); let mprops = Slice(mprops, 0); for i in StepRng::new(invocation_id.x..total, num_threads) { @@ -101,27 +107,33 @@ pub fn gpu_solver_count_constraints( if im.contact.len == 0 { continue; } + // Multibody-owned manifolds have no rigid-body constraint (see + // `gpu_solver_init_constraints`). + if touches_multibody(im, &body_is_multibody) { + continue; + } let body1 = im.bodies.x; let body2 = im.bodies.y; - let group1 = body_group[body1 as usize]; - let group2 = body_group[body2 as usize]; - - // Count toward the body's GROUP slot. A body is "active" for the - // graph-coloring graph if it's a free dynamic body (inv_mass != 0) OR - // it's part of a multibody (group != self — the multibody handles its - // own dynamics but its bodies still need correct coloring so contacts - // touching different links of the same multibody never share a color). - let is_mb1 = group1 != body1; - if mprops[body1 as usize].inv_mass != Vector::ZERO || is_mb1 { - atomic_add_u32(&mut body_constraint_counts[group1 as usize], 1); + + // Count toward the body's slot (only free bodies are left here, whose + // graph group is themselves). A body is "active" for the + // graph-coloring graph if it's a free dynamic body (inv_mass != 0). + if mprops[body1 as usize].inv_mass != Vector::ZERO { + atomic_add_u32(&mut body_constraint_counts[body1 as usize], 1); } - let is_mb2 = group2 != body2; - if mprops[body2 as usize].inv_mass != Vector::ZERO || is_mb2 { - atomic_add_u32(&mut body_constraint_counts[group2 as usize], 1); + if mprops[body2 as usize].inv_mass != Vector::ZERO { + atomic_add_u32(&mut body_constraint_counts[body2 as usize], 1); } } } +/// `true` when either body of the manifold is a multibody link: its contacts +/// are solved by the multibody solver, never by the rigid-body one. +#[inline(always)] +fn touches_multibody(im: &IndexedManifold, body_is_multibody: &Slice<'_, u32>) -> bool { + body_is_multibody[im.bodies.x as usize] != 0 || body_is_multibody[im.bodies.y as usize] != 0 +} + /// Updates constraints for a new substep. #[spirv_bindgen] #[spirv(compute(threads(64)))] @@ -201,35 +213,32 @@ pub fn gpu_solver_sort_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contacts: &[IndexedManifold], #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] body_constraint_ids: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] body_group: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] body_is_multibody: &[u32], ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let total = contact_plan.bound; let contacts = Slice(contacts, 0); let mut body_constraint_counts = SliceMut(body_constraint_counts, 0); - let body_group = Slice(body_group, 0); + let body_is_multibody = Slice(body_is_multibody, 0); let mprops = Slice(mprops, 0); let mut body_constraint_ids = SliceMut(body_constraint_ids, 0); for i in StepRng::new(invocation_id.x..total, num_threads) { - if contacts[i as usize].contact.len == 0 { + let im = &contacts[i as usize]; + // Same filter as `gpu_solver_count_constraints`. + if im.contact.len == 0 || touches_multibody(im, &body_is_multibody) { continue; } - let body1 = contacts[i as usize].bodies.x as usize; - let body2 = contacts[i as usize].bodies.y as usize; - let group1 = body_group[body1] as usize; - let group2 = body_group[body2] as usize; - - let is_mb1 = group1 != body1; - if mprops[body1].inv_mass != Vector::ZERO || is_mb1 { - let id1 = atomic_add_u32(&mut body_constraint_counts[group1], 1); + let body1 = im.bodies.x as usize; + let body2 = im.bodies.y as usize; + + if mprops[body1].inv_mass != Vector::ZERO { + let id1 = atomic_add_u32(&mut body_constraint_counts[body1], 1); body_constraint_ids[id1 as usize] = i; } - - let is_mb2 = group2 != body2; - if mprops[body2].inv_mass != Vector::ZERO || is_mb2 { - let id2 = atomic_add_u32(&mut body_constraint_counts[group2], 1); + if mprops[body2].inv_mass != Vector::ZERO { + let id2 = atomic_add_u32(&mut body_constraint_counts[body2], 1); body_constraint_ids[id2 as usize] = i; } } @@ -415,7 +424,7 @@ pub fn gpu_step_gauss_seidel( let color_sorted_ids = Slice(color_sorted_ids, 0); let mut solver_vels = SliceMut(solver_vels, 0); let color = *curr_color; - let use_bias = *use_bias != 0; + let (use_bias, solve_friction) = decode_bias_mode(*use_bias); // Color-major bucket ends; see `gpu_warmstart`. let start = color_starts.read((color * nb - 1) as usize); @@ -433,6 +442,7 @@ pub fn gpu_step_gauss_seidel( &mut solver_vel1, &mut solver_vel2, use_bias, + solve_friction, ); solver_vels[solver_id1] = solver_vel1; @@ -529,7 +539,7 @@ pub fn gpu_step_gauss_seidel_fused( let color_sorted_ids = Slice(color_sorted_ids, 0); let mut solver_vels = SliceMut(solver_vels, 0); let num_colors = *num_colors; - let use_bias = *use_bias != 0; + let (use_bias, solve_friction) = decode_bias_mode(*use_bias); for color in 1..=num_colors { // Empty-color skip: see `gpu_warmstart_fused`. @@ -554,6 +564,7 @@ pub fn gpu_step_gauss_seidel_fused( &mut solver_vel1, &mut solver_vel2, use_bias, + solve_friction, ); solver_vels[solver_id1] = solver_vel1; diff --git a/src_rbd_shaders/dynamics/solver_utils.rs b/src_rbd_shaders/dynamics/solver_utils.rs index 1d199979..ea8f76de 100644 --- a/src_rbd_shaders/dynamics/solver_utils.rs +++ b/src_rbd_shaders/dynamics/solver_utils.rs @@ -536,13 +536,16 @@ impl TwoBodyConstraint { } } - /// Main constraint solver iteration (Projected Gauss-Seidel). + /// Main constraint solver iteration (Projected Gauss-Seidel). `solve_friction` + /// gates the tangent rows: the stabilization sweep always solves them, the + /// biased pass only when `RbdSimParams::friction_in_bias_pass` is set. #[inline(always)] pub fn solve_constraint_gauss_seidel( &mut self, solver_vel1: &mut Velocity, solver_vel2: &mut Velocity, use_bias: bool, + solve_friction: bool, ) { let dir_a = self.dir_a; let friction_coeff = self.limit; @@ -576,8 +579,9 @@ impl TwoBodyConstraint { solver_vel2.angular += ii_torque_dir_b * delta_impulse; } - // Friction is only solved during the stabilization sweep. - if use_bias { + // Friction is solved during the stabilization sweep, and during the + // biased pass only when `friction_in_bias_pass` is set. + if !solve_friction { return; } From 58462872946df4bc4620676f0e219b4af512a792 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 20:38:07 +0200 Subject: [PATCH 07/25] feat(python): contact impulse readbacks, friction combine rule and friction_in_bias_pass --- crates/nexus_python3d/src/nexus.rs | 158 ++++++++++++++++++++++++++++- crates/nexus_python3d/src/rbd.rs | 24 +++++ src/state.rs | 8 +- 3 files changed, 181 insertions(+), 9 deletions(-) diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 9d5b9792..0ebf216c 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -10,7 +10,7 @@ use crate::rbd::{ }; use crate::robot::{Robot, build_robot, free_axes, joint_axis, pose_from_wxyz, to_wxyz}; use crate::viewer::NexusViewer; -use khal::backend::GpuTimestamps as RGpuTimestamps; +use khal::backend::{Backend, GpuTimestamps as RGpuTimestamps}; use nexus3d::mpm::solver::BoundaryCondition as RBoundaryCondition; use nexus3d::prelude::{ NexusPipeline as RNexusPipeline, NexusPipelineMask, NexusState as RNexusState, @@ -1152,9 +1152,13 @@ impl NexusState { /// pair applies to contacts at rest, `allowed_linear_error` is the /// tolerated penetration in length units, `max_corrective_velocity` caps /// penetration recovery, `prediction_distance` is the contact detection - /// margin, and `internal_pgs_iterations` the multibody solver's PGS - /// iterations per substep. - #[pyo3(signature = (contact_natural_frequency=None, contact_damping_ratio=None, static_contact_natural_frequency=None, static_contact_damping_ratio=None, allowed_linear_error=None, max_corrective_velocity=None, prediction_distance=None, internal_pgs_iterations=None))] + /// margin, `internal_pgs_iterations` the biased-pass PGS iterations per + /// substep (rigid-body and multibody sweeps alike), and `friction_in_bias_pass` whether friction + /// rows are solved in every biased PGS iteration instead of only in the + /// per-substep stabilization sweep (rapier's default, `False`); `True` + /// gives friction as many iterations as the normal rows, which holds + /// grasps and resting contacts far more firmly. + #[pyo3(signature = (contact_natural_frequency=None, contact_damping_ratio=None, static_contact_natural_frequency=None, static_contact_damping_ratio=None, allowed_linear_error=None, max_corrective_velocity=None, prediction_distance=None, internal_pgs_iterations=None, friction_in_bias_pass=None))] #[allow(clippy::too_many_arguments)] fn set_rbd_solver_params( &mut self, @@ -1166,6 +1170,7 @@ impl NexusState { max_corrective_velocity: Option, prediction_distance: Option, internal_pgs_iterations: Option, + friction_in_bias_pass: Option, ) { for env in 0..self.0.num_environments() { let Some(mut params) = self.0.rbd_sim_params(env) else { @@ -1195,6 +1200,9 @@ impl NexusState { if let Some(v) = internal_pgs_iterations { params.num_internal_pgs_iterations = v.max(1); } + if let Some(v) = friction_in_bias_pass { + params.friction_in_bias_pass = v as u32; + } self.0.set_rbd_sim_params(env, params); } } @@ -1224,10 +1232,152 @@ impl NexusState { )?; dict.set_item("prediction_distance", p.normalized_prediction_distance)?; dict.set_item("internal_pgs_iterations", p.num_internal_pgs_iterations)?; + dict.set_item("friction_in_bias_pass", p.friction_in_bias_pass != 0)?; } Ok(dict) } + /// Debug readback of the live contact manifolds, one dict per active + /// manifold: `collider_a` / `collider_b` and `body_a` / `body_b` (GPU + /// indices), the combined `friction`, `normal_a` and the `points` as + /// `[x, y, z, dist]` rows, both in collider A's local frame. Blocks on + /// the GPU; for diagnostics, not for control loops. + fn debug_contacts<'py>( + &self, + py: Python<'py>, + viewer: PyRef, + ) -> PyResult>> { + use pyo3::types::PyDict; + let Some(rbd) = self.0.rbd.as_ref() else { + return Ok(Vec::new()); + }; + let contacts: Vec = + pollster::block_on(viewer.backend().slow_read_vec(rbd.contacts().buffer())) + .map_err(gpu_err)?; + let mut out = Vec::new(); + for c in contacts.iter().filter(|c| c.contact.len > 0) { + let dict = PyDict::new(py); + dict.set_item("collider_a", c.colliders.x)?; + dict.set_item("collider_b", c.colliders.y)?; + dict.set_item("body_a", c.bodies.x)?; + dict.set_item("body_b", c.bodies.y)?; + dict.set_item("friction", c.friction)?; + let n = c.contact.normal_a; + dict.set_item("normal_a", [n.x, n.y, n.z])?; + let points: Vec<[f32; 4]> = (0..c.contact.len as usize) + .map(|k| { + let p = c.contact.points_a[k]; + [p.pt.x, p.pt.y, p.pt.z, p.dist] + }) + .collect(); + dict.set_item("points", points)?; + out.push(dict); + } + Ok(out) + } + + /// Debug readback of the multibody contact constraints as left by the + /// last step, one dict per active slot: `multibody` and `link` (batch + /// local), `free_body` (GPU index, `None` for a self-contact), `kind` + /// (`"normal"` or `"tangent"`), `friction`, the accumulated per-substep + /// `impulse` and the free-body jacobian direction `dir`. Blocks on the + /// GPU; for diagnostics only. + fn debug_multibody_contact_impulses<'py>( + &self, + py: Python<'py>, + viewer: PyRef, + ) -> PyResult>> { + use nexus3d::rbd::shaders::dynamics::{ + MB_CONTACT_KIND_NORMAL, MB_CONTACT_KIND_TANGENT, MultibodyContactConstraint, + MultibodyInfo, + }; + use pyo3::types::PyDict; + let Some(rbd) = self.0.rbd.as_ref() else { + return Ok(Vec::new()); + }; + let backend = viewer.backend(); + let mb = rbd.multibodies(); + let cons: Vec = + pollster::block_on(backend.slow_read_vec(mb.contact_constraints().buffer())) + .map_err(gpu_err)?; + let infos: Vec = + pollster::block_on(backend.slow_read_vec(mb.multibody_info().buffer())) + .map_err(gpu_err)?; + let mut out = Vec::new(); + for info in infos.iter().filter(|i| i.ndofs > 0) { + let start = info.contact_constraint_start as usize; + let end = start + info.contact_constraint_count as usize; + for c in cons.iter().take(end.min(cons.len())).skip(start) { + let kind = match c.kind { + MB_CONTACT_KIND_NORMAL => "normal", + MB_CONTACT_KIND_TANGENT => "tangent", + _ => continue, + }; + let dict = PyDict::new(py); + dict.set_item("multibody", c.multibody_id)?; + dict.set_item("link", c.link_id)?; + dict.set_item( + "free_body", + (c.free_body_id != u32::MAX).then_some(c.free_body_id), + )?; + dict.set_item("kind", kind)?; + dict.set_item("friction", c.friction_coeff)?; + dict.set_item("impulse", c.impulse)?; + dict.set_item("dir", [c.lin_jac.x, c.lin_jac.y, c.lin_jac.z])?; + dict.set_item("free_body_inv_mass", c.free_body_im)?; + out.push(dict); + } + } + Ok(out) + } + + /// Debug readback of the rigid-body (two-body) contact constraints as + /// left by the last step, one dict per active manifold: `body_a` / + /// `body_b` (GPU indices), `dir_a` (world normal force direction on + /// body A), the combined `friction`, and per point the accumulated + /// per-substep `normal_impulse` and `tangent_impulse` pair. Blocks on + /// the GPU; for diagnostics only. + fn debug_rigid_contact_impulses<'py>( + &self, + py: Python<'py>, + viewer: PyRef, + ) -> PyResult>> { + use nexus3d::rbd::shaders::dynamics::TwoBodyConstraint; + use pyo3::types::PyDict; + let Some(rbd) = self.0.rbd.as_ref() else { + return Ok(Vec::new()); + }; + let cons: Vec = pollster::block_on( + viewer + .backend() + .slow_read_vec(rbd.rigid_contact_constraints().buffer()), + ) + .map_err(gpu_err)?; + let mut out = Vec::new(); + for c in cons.iter().filter(|c| c.len > 0) { + let dict = PyDict::new(py); + dict.set_item("body_a", c.solver_body_a)?; + dict.set_item("body_b", c.solver_body_b)?; + dict.set_item("dir_a", [c.dir_a.x, c.dir_a.y, c.dir_a.z])?; + dict.set_item("friction", c.limit)?; + dict.set_item("inv_mass_a", c.im_a.x)?; + dict.set_item("inv_mass_b", c.im_b.x)?; + let normal: Vec = (0..c.len as usize) + .map(|k| c.elements[k].normal_part.impulse) + .collect(); + let tangent: Vec<[f32; 2]> = (0..c.len as usize) + .map(|k| { + let t = c.elements[k].tangent_part.impulse; + [t.x, t.y] + }) + .collect(); + dict.set_item("normal_impulse", normal)?; + dict.set_item("tangent_impulse", tangent)?; + out.push(dict); + } + Ok(out) + } + // --- rbd config ------------------------------------------------------- fn set_rbd_steps_per_frame(&mut self, steps: u32) { diff --git a/crates/nexus_python3d/src/rbd.rs b/crates/nexus_python3d/src/rbd.rs index 6ba0430b..5a063709 100644 --- a/crates/nexus_python3d/src/rbd.rs +++ b/crates/nexus_python3d/src/rbd.rs @@ -221,6 +221,14 @@ impl ColliderBuilder { fn friction(&self, friction: f32) -> Self { Self(self.0.clone().friction(friction)) } + /// How this collider's friction merges with the other collider's: + /// `"average"` (default), `"min"`, `"multiply"`, `"max"` or `"sum"` + /// (clamped to `[0, 1]`). The stronger rule of the two colliders wins. + fn friction_combine_rule(&self, rule: &str) -> PyResult { + Ok(Self( + self.0.clone().friction_combine_rule(combine_rule(rule)?), + )) + } fn restitution(&self, restitution: f32) -> Self { Self(self.0.clone().restitution(restitution)) } @@ -460,3 +468,19 @@ impl JointArg { } } } + +/// Parses a `CoefficientCombineRule` name (see `ColliderBuilder.friction_combine_rule`). +fn combine_rule(rule: &str) -> PyResult { + Ok(match rule { + "average" => rp::CoefficientCombineRule::Average, + "min" => rp::CoefficientCombineRule::Min, + "multiply" => rp::CoefficientCombineRule::Multiply, + "max" => rp::CoefficientCombineRule::Max, + "sum" => rp::CoefficientCombineRule::ClampedSum, + other => { + return Err(PyValueError::new_err(format!( + "unknown combine rule {other:?} (expected average, min, multiply, max or sum)" + ))); + } + }) +} diff --git a/src/state.rs b/src/state.rs index 5c0cb427..4d4d9d10 100644 --- a/src/state.rs +++ b/src/state.rs @@ -12,9 +12,7 @@ use crate::rbd::dynamics::{ body::{BodyCoupling, RapierBodyCouplingEntry}, }; use crate::rbd::pipeline::{RbdCapacities, RbdResizePolicy, RbdState, RunStats}; -#[cfg(feature = "dim3")] -use khal::backend::Backend; -use khal::backend::{GpuBackend, GpuBackendError}; +use khal::backend::{Backend, GpuBackend, GpuBackendError}; /// Handle referencing a rigid-body managed by a [`NexusState`]. #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)] @@ -383,8 +381,8 @@ impl NexusState { } } - /// Sets how many PGS iterations the multibody solver's biased pass runs per - /// substep. + /// Sets how many PGS iterations the biased pass runs per substep (rigid-body + /// and multibody sweeps alike). #[cfg(all(feature = "rbd", feature = "dim3"))] pub fn set_rbd_num_internal_pgs_iterations(&mut self, backend: &GpuBackend, n: u32) { for params in &mut self.rbd_sim_params { From febdf8cc4985afd837461988038881ef06a32473 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 29 Aug 2026 14:41:16 +0200 Subject: [PATCH 08/25] feat(viewer): antialiased sensor renders with hard, near-fit shadows --- CHANGELOG.md | 5 ++ crates/nexus_python3d/src/viewer.rs | 45 +++++++++++++++ src_viewer/graphics.rs | 12 ++++ src_viewer/sensors.rs | 25 +++++++- src_viewer/viewer.rs | 89 +++++++++++++++++++++++++++++ 5 files changed, 175 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a3706d8..52cabeb0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ PGS pass, instead of only during the once-per-substep stabilization sweep (rapier's `friction_in_bias_pass`). Off by default; on, friction gets as many iterations as the normal rows, which keeps pinch grasps and resting stacks from drifting. +- Viewer: sensor cameras render with 4x MSAA and hard shadow edges by default + (`set_sensor_antialiasing`, `set_sensor_shadow_softness`); the window's shadows are hard-edged + too (`set_shadow_softness`); the sharpest directional cascade covers the first 3 m instead of + 12 (`set_sensor_shadow_range`, `set_shadow_range`), and `set_body_casts_shadows` excludes a + body from the shadow map. Needs kiss3d's offscreen MSAA and cascade controls. - Python: `NexusState.set_rbd_solver_params(friction_in_bias_pass=...)`, `ColliderBuilder.friction_combine_rule`, and the `debug_contacts` / `debug_multibody_contact_impulses` GPU readbacks for contact diagnostics. diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index d4381107..d7ba812b 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -349,11 +349,56 @@ impl NexusViewer { self.inner_mut().set_body_color(env, handle.0, rgba); } + /// Whether body `handle`'s visual nodes cast shadows (default `True`). A + /// floor slab that does not cast keeps the shadow map fit to the objects + /// above it, so hard shadow edges stay crisp. + fn set_body_casts_shadows(&mut self, env: u32, handle: RigidBodyHandle, casts: bool) { + self.inner_mut() + .set_body_casts_shadows(env, handle.0, casts); + } + /// Ambient light level of the main window's shaded render. fn set_ambient(&mut self, ambient: f32) { self.inner_mut().set_ambient(ambient); } + /// MSAA sample count of the sensor cameras' shaded renders, existing and + /// future ones (`1` disables antialiasing, `4` is the default). + fn set_sensor_antialiasing(&mut self, samples: u32) { + self.inner_mut().set_sensor_antialiasing(samples); + } + + /// Shadow-edge softness of the sensor cameras' shaded renders, existing + /// and future ones (`0.0` hard edges, the default; `1.0` the PCF + /// penumbra). + fn set_sensor_shadow_softness(&mut self, softness: f32) { + self.inner_mut().set_sensor_shadow_softness(softness); + } + + /// Shadow-edge softness of the main window's shaded render (`0.0` hard + /// edges, the default; `1.0` the PCF penumbra). + fn set_shadow_softness(&mut self, softness: f32) { + self.inner_mut().set_shadow_softness(softness); + } + + /// Directional-shadow cascade layout of the sensor cameras' renders, + /// existing and future ones: the sharpest cascade covers the camera's + /// first `first_cascade_far_bound` meters (default 3) and shadows stop at + /// `shadow_distance` meters (default: the camera far plane). + #[pyo3(signature = (first_cascade_far_bound, shadow_distance=f32::INFINITY))] + fn set_sensor_shadow_range(&mut self, first_cascade_far_bound: f32, shadow_distance: f32) { + self.inner_mut() + .set_sensor_shadow_range(first_cascade_far_bound, shadow_distance); + } + + /// Directional-shadow cascade layout of the main window's render (see + /// `set_sensor_shadow_range`). + #[pyo3(signature = (first_cascade_far_bound, shadow_distance=f32::INFINITY))] + fn set_shadow_range(&mut self, first_cascade_far_bound: f32, shadow_distance: f32) { + self.inner_mut() + .set_shadow_range(first_cascade_far_bound, shadow_distance); + } + /// Registers one render node for body `handle` in `env` drawing `shape` /// with base color `rgba`, optional per-vertex UVs (trimesh shapes only) /// and an optional encoded image texture (`texture` bytes, PNG/JPEG, diff --git a/src_viewer/graphics.rs b/src_viewer/graphics.rs index 24df1465..45fbf0ae 100644 --- a/src_viewer/graphics.rs +++ b/src_viewer/graphics.rs @@ -941,6 +941,18 @@ impl RenderContext { } } + /// Whether body `handle`'s visual nodes cast shadows. Turning it off for a + /// large floor keeps the directional shadow frustum fit to the objects + /// that matter, so shadow texels stay small. + #[cfg(feature = "dim3")] + pub fn set_body_casts_shadows(&mut self, env: u32, handle: RigidBodyHandle, casts: bool) { + for visual in &mut self.visual_nodes { + if visual.env == env && visual.handle == handle.0 { + let _ = visual.node.set_casts_shadows(casts); + } + } + } + /// Updates each visual node's transform from the body-origin poses (indexed /// by the body's GPU pose slot). Visual-mesh local poses are body-relative, /// so they compose with `body_poses` (not the collider world poses the diff --git a/src_viewer/sensors.rs b/src_viewer/sensors.rs index db642572..bcaabafe 100644 --- a/src_viewer/sensors.rs +++ b/src_viewer/sensors.rs @@ -16,7 +16,7 @@ use kiss3d::camera::Camera3d; use kiss3d::event::WindowEvent; use kiss3d::prelude::Color; use kiss3d::scene::SceneNode3d; -use kiss3d::window::{Canvas, OffscreenSurface}; +use kiss3d::window::{Canvas, NumSamples, OffscreenSurface}; use nexus::rbd::math::Pose; use rapier::data::Index; @@ -206,6 +206,29 @@ impl SensorCamera { .set_background_color(Color::new(rgba[0], rgba[1], rgba[2], rgba[3])); } + /// MSAA sample count of the shaded render (`1` disables antialiasing; + /// kiss3d supports 1 and 4, other values round down to 1). + pub fn set_samples(&mut self, samples: u32) { + let samples = NumSamples::from_u32(samples).unwrap_or(NumSamples::One); + self.surface.set_samples(samples); + } + + /// Shadow-edge softness of the shaded render: `0.0` hard edges, `1.0` + /// kiss3d's default PCF penumbra. + pub fn set_shadow_softness(&mut self, softness: f32) { + self.surface.set_shadow_softness(softness); + } + + /// Directional-shadow cascade layout of the shaded render: the + /// highest-resolution cascade covers the camera's first + /// `first_cascade_far_bound` meters and shadows stop at `shadow_distance` + /// (`INFINITY` = the camera far plane). + pub fn set_shadow_range(&mut self, first_cascade_far_bound: f32, shadow_distance: f32) { + self.surface + .set_first_cascade_far_bound(first_cascade_far_bound); + self.surface.set_shadow_distance(shadow_distance); + } + /// Renders the shaded scene and returns it as row-major, top-left origin /// RGB bytes (`width * height * 3`). pub async fn render_rgb(&mut self, scene: &mut SceneNode3d) -> Vec { diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 782f917e..1b2cc66a 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -171,6 +171,15 @@ pub type SceneNode = SceneNode3d; /// the freeze; rendering a few real frames first forces the paint. const COMPILE_BANNER_PRESENT_FRAMES: u32 = 10; +/// MSAA sample count of new sensor cameras. +#[cfg(feature = "dim3")] +const DEFAULT_SENSOR_SAMPLES: u32 = 4; +/// Shadow-edge softness of the window and of new sensor cameras (hard edges). +const DEFAULT_SHADOW_SOFTNESS: f32 = 0.0; +/// Far bound (meters) of the highest-resolution directional shadow cascade of +/// the window and of new sensor cameras: desk-sized scenes, not landscapes. +const DEFAULT_SHADOW_FIRST_CASCADE: f32 = 3.0; + pub struct NexusViewer { window: Window, scene2d: SceneNode2d, @@ -244,6 +253,17 @@ pub struct NexusViewer { /// and the per-object passes see world-posed nodes. #[cfg(feature = "dim3")] sensors: Vec, + /// MSAA sample count given to every new sensor camera (default 4). + #[cfg(feature = "dim3")] + sensor_samples: u32, + /// Shadow-edge softness given to every new sensor camera (default 0.0, + /// hard edges; 1.0 is kiss3d's default penumbra). + #[cfg(feature = "dim3")] + sensor_shadow_softness: f32, + /// Directional-shadow cascade layout given to every new sensor camera: + /// `(first cascade far bound, shadow distance)` in meters. + #[cfg(feature = "dim3")] + sensor_shadow_range: (f32, f32), /// Body-origin poses from the last readback sync, indexed by GPU pose slot. /// Empty on the zero-readback path. #[cfg(feature = "dim3")] @@ -292,6 +312,10 @@ impl NexusViewer { // Disable MSAA, this puts extra load on the GPU that ends up // falsifying the gpu physics timestamps. window.set_samples(NumSamples::One); + // Hard shadow edges and a short first cascade: crisper contact shadows + // on the small geometry physics scenes are made of. + window.set_shadow_softness(DEFAULT_SHADOW_SOFTNESS); + window.set_first_cascade_far_bound(DEFAULT_SHADOW_FIRST_CASCADE); #[cfg(feature = "dim2")] let (camera2d, camera3d) = { @@ -351,6 +375,12 @@ impl NexusViewer { #[cfg(feature = "dim3")] sensors: Vec::new(), #[cfg(feature = "dim3")] + sensor_samples: DEFAULT_SENSOR_SAMPLES, + #[cfg(feature = "dim3")] + sensor_shadow_softness: DEFAULT_SHADOW_SOFTNESS, + #[cfg(feature = "dim3")] + sensor_shadow_range: (DEFAULT_SHADOW_FIRST_CASCADE, f32::INFINITY), + #[cfg(feature = "dim3")] body_pose_cache: Vec::new(), ui: UiState { run_state: RunState::Paused, @@ -782,10 +812,61 @@ impl NexusViewer { ) -> usize { let mut sensor = SensorCamera::new(width, height, fov_y, znear, zfar).await; sensor.generation = self.nexus_render.generation; + sensor.set_samples(self.sensor_samples); + sensor.set_shadow_softness(self.sensor_shadow_softness); + sensor.set_shadow_range(self.sensor_shadow_range.0, self.sensor_shadow_range.1); self.sensors.push(sensor); self.sensors.len() - 1 } + /// MSAA sample count of the sensor cameras' shaded renders, existing and + /// future ones (`1` disables antialiasing, `4` is the default). + #[cfg(feature = "dim3")] + pub fn set_sensor_antialiasing(&mut self, samples: u32) { + self.sensor_samples = samples.max(1); + for sensor in &mut self.sensors { + sensor.set_samples(self.sensor_samples); + } + } + + /// Shadow-edge softness of the sensor cameras' shaded renders, existing + /// and future ones (`0.0` hard edges, the default; `1.0` kiss3d's PCF + /// penumbra). + #[cfg(feature = "dim3")] + pub fn set_sensor_shadow_softness(&mut self, softness: f32) { + self.sensor_shadow_softness = softness.max(0.0); + for sensor in &mut self.sensors { + sensor.set_shadow_softness(self.sensor_shadow_softness); + } + } + + /// Shadow-edge softness of the main window's shaded render (`0.0` hard + /// edges, the default; `1.0` kiss3d's PCF penumbra). + pub fn set_shadow_softness(&mut self, softness: f32) { + self.window.set_shadow_softness(softness.max(0.0)); + } + + /// Directional-shadow cascade layout of the sensor cameras' renders, + /// existing and future ones: the highest-resolution cascade covers the + /// camera's first `first_cascade_far_bound` meters (default 3) and shadows + /// stop at `shadow_distance` (default: the camera far plane). A short first + /// cascade keeps hard shadow edges crisp on a desk-sized scene. + #[cfg(feature = "dim3")] + pub fn set_sensor_shadow_range(&mut self, first_cascade_far_bound: f32, shadow_distance: f32) { + self.sensor_shadow_range = (first_cascade_far_bound.max(0.01), shadow_distance.max(0.0)); + for sensor in &mut self.sensors { + sensor.set_shadow_range(self.sensor_shadow_range.0, self.sensor_shadow_range.1); + } + } + + /// Directional-shadow cascade layout of the main window's render (see + /// [`Self::set_sensor_shadow_range`]). + pub fn set_shadow_range(&mut self, first_cascade_far_bound: f32, shadow_distance: f32) { + self.window + .set_first_cascade_far_bound(first_cascade_far_bound.max(0.01)); + self.window.set_shadow_distance(shadow_distance.max(0.0)); + } + /// Starts a new scene generation (see `RenderContext::next_generation`): /// visual nodes and sensor cameras created afterwards belong to it, and it /// becomes the active one; the previous scene's nodes leave the graph and @@ -929,6 +1010,14 @@ impl NexusViewer { self.nexus_render.set_body_color(env, handle, color) } + /// Whether body `handle`'s visual nodes cast shadows (default `true`). + /// Turn it off for a floor slab so the shadow map covers only the objects + /// above it. + #[cfg(feature = "dim3")] + pub fn set_body_casts_shadows(&mut self, env: u32, handle: RigidBodyHandle, casts: bool) { + self.nexus_render.set_body_casts_shadows(env, handle, casts) + } + /// Ambient light level of the main window's shaded render. pub fn set_ambient(&mut self, ambient: f32) { self.window.set_ambient(ambient); From da9c4da8167070f5ff1075b5846c23e50a6292aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 29 Aug 2026 17:14:36 +0200 Subject: [PATCH 09/25] feat(viewer): 4096-texel shadow maps over a four-layer atlas --- CHANGELOG.md | 6 +++-- crates/nexus_python3d/src/viewer.rs | 17 ++++++++++++ src_viewer/sensors.rs | 7 +++++ src_viewer/viewer.rs | 40 +++++++++++++++++++++++++++++ 4 files changed, 68 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 52cabeb0..526536f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,8 +9,10 @@ - Viewer: sensor cameras render with 4x MSAA and hard shadow edges by default (`set_sensor_antialiasing`, `set_sensor_shadow_softness`); the window's shadows are hard-edged too (`set_shadow_softness`); the sharpest directional cascade covers the first 3 m instead of - 12 (`set_sensor_shadow_range`, `set_shadow_range`), and `set_body_casts_shadows` excludes a - body from the shadow map. Needs kiss3d's offscreen MSAA and cascade controls. + 12 (`set_sensor_shadow_range`, `set_shadow_range`), the shadow map is 4096 texels over a + 4-layer atlas instead of 2048 over 16 (`set_sensor_shadow_resolution`, `set_shadow_resolution`), + and `set_body_casts_shadows` excludes a body from the shadow map. Needs kiss3d's offscreen + MSAA, cascade and atlas-layer controls. - Python: `NexusState.set_rbd_solver_params(friction_in_bias_pass=...)`, `ColliderBuilder.friction_combine_rule`, and the `debug_contacts` / `debug_multibody_contact_impulses` GPU readbacks for contact diagnostics. diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index d7ba812b..2336aed1 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -391,6 +391,23 @@ impl NexusViewer { .set_sensor_shadow_range(first_cascade_far_bound, shadow_distance); } + /// Shadow map resolution (texels per atlas layer, default 4096) and atlas + /// layer count (default 4, one directional light's cascades; a point light + /// needs 6) of the sensor cameras' renders, existing and future ones. + /// Memory per camera is `resolution² × layers × 8` bytes. + #[pyo3(signature = (resolution, layers=4))] + fn set_sensor_shadow_resolution(&mut self, resolution: u32, layers: u32) { + self.inner_mut() + .set_sensor_shadow_resolution(resolution, layers); + } + + /// Shadow map resolution and atlas layer count of the main window's render + /// (see `set_sensor_shadow_resolution`). + #[pyo3(signature = (resolution, layers=4))] + fn set_shadow_resolution(&mut self, resolution: u32, layers: u32) { + self.inner_mut().set_shadow_resolution(resolution, layers); + } + /// Directional-shadow cascade layout of the main window's render (see /// `set_sensor_shadow_range`). #[pyo3(signature = (first_cascade_far_bound, shadow_distance=f32::INFINITY))] diff --git a/src_viewer/sensors.rs b/src_viewer/sensors.rs index bcaabafe..f5f83c52 100644 --- a/src_viewer/sensors.rs +++ b/src_viewer/sensors.rs @@ -219,6 +219,13 @@ impl SensorCamera { self.surface.set_shadow_softness(softness); } + /// Shadow map resolution (texels per atlas layer, square) and the number + /// of atlas layers to allocate (one directional light needs four). + pub fn set_shadow_resolution(&mut self, resolution: u32, layers: u32) { + self.surface.set_shadow_atlas_layers(layers); + self.surface.set_shadow_resolution(resolution); + } + /// Directional-shadow cascade layout of the shaded render: the /// highest-resolution cascade covers the camera's first /// `first_cascade_far_bound` meters and shadows stop at `shadow_distance` diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 1b2cc66a..1e70f409 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -179,6 +179,11 @@ const DEFAULT_SHADOW_SOFTNESS: f32 = 0.0; /// Far bound (meters) of the highest-resolution directional shadow cascade of /// the window and of new sensor cameras: desk-sized scenes, not landscapes. const DEFAULT_SHADOW_FIRST_CASCADE: f32 = 3.0; +/// Shadow map texels per atlas layer for the window and new sensor cameras. +const DEFAULT_SHADOW_RESOLUTION: u32 = 4096; +/// Shadow atlas layers for the window and new sensor cameras: one directional +/// light's four cascades. +const DEFAULT_SHADOW_ATLAS_LAYERS: u32 = 4; pub struct NexusViewer { window: Window, @@ -264,6 +269,9 @@ pub struct NexusViewer { /// `(first cascade far bound, shadow distance)` in meters. #[cfg(feature = "dim3")] sensor_shadow_range: (f32, f32), + /// Shadow map `(resolution, atlas layers)` given to every new sensor camera. + #[cfg(feature = "dim3")] + sensor_shadow_resolution: (u32, u32), /// Body-origin poses from the last readback sync, indexed by GPU pose slot. /// Empty on the zero-readback path. #[cfg(feature = "dim3")] @@ -316,6 +324,10 @@ impl NexusViewer { // on the small geometry physics scenes are made of. window.set_shadow_softness(DEFAULT_SHADOW_SOFTNESS); window.set_first_cascade_far_bound(DEFAULT_SHADOW_FIRST_CASCADE); + // Physics scenes light with directional lights only: four cascade + // layers instead of kiss3d's sixteen pay for a finer shadow map. + window.set_shadow_atlas_layers(DEFAULT_SHADOW_ATLAS_LAYERS); + window.set_shadow_resolution(DEFAULT_SHADOW_RESOLUTION); #[cfg(feature = "dim2")] let (camera2d, camera3d) = { @@ -381,6 +393,8 @@ impl NexusViewer { #[cfg(feature = "dim3")] sensor_shadow_range: (DEFAULT_SHADOW_FIRST_CASCADE, f32::INFINITY), #[cfg(feature = "dim3")] + sensor_shadow_resolution: (DEFAULT_SHADOW_RESOLUTION, DEFAULT_SHADOW_ATLAS_LAYERS), + #[cfg(feature = "dim3")] body_pose_cache: Vec::new(), ui: UiState { run_state: RunState::Paused, @@ -815,6 +829,10 @@ impl NexusViewer { sensor.set_samples(self.sensor_samples); sensor.set_shadow_softness(self.sensor_shadow_softness); sensor.set_shadow_range(self.sensor_shadow_range.0, self.sensor_shadow_range.1); + sensor.set_shadow_resolution( + self.sensor_shadow_resolution.0, + self.sensor_shadow_resolution.1, + ); self.sensors.push(sensor); self.sensors.len() - 1 } @@ -859,6 +877,28 @@ impl NexusViewer { } } + /// Shadow map resolution (texels per atlas layer, default 4096) and atlas + /// layer count (default 4: one directional light's cascades; a point light + /// needs 6, a spot light 1) of the sensor cameras' renders, existing and + /// future ones. Memory per camera is `resolution² × layers × 8` bytes. + #[cfg(feature = "dim3")] + pub fn set_sensor_shadow_resolution(&mut self, resolution: u32, layers: u32) { + self.sensor_shadow_resolution = (resolution.max(1), layers.max(1)); + for sensor in &mut self.sensors { + sensor.set_shadow_resolution( + self.sensor_shadow_resolution.0, + self.sensor_shadow_resolution.1, + ); + } + } + + /// Shadow map resolution and atlas layer count of the main window's render + /// (see [`Self::set_sensor_shadow_resolution`]). + pub fn set_shadow_resolution(&mut self, resolution: u32, layers: u32) { + self.window.set_shadow_atlas_layers(layers.max(1)); + self.window.set_shadow_resolution(resolution.max(1)); + } + /// Directional-shadow cascade layout of the main window's render (see /// [`Self::set_sensor_shadow_range`]). pub fn set_shadow_range(&mut self, first_cascade_far_bound: f32, shadow_distance: f32) { From 4e8dad8bcc20df679838f54d844d191420a83521 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 29 Aug 2026 19:18:55 +0200 Subject: [PATCH 10/25] feat(viewer): soft shadow penumbra default, mipmapped anisotropic textures --- CHANGELOG.md | 11 ++++++----- crates/nexus_python3d/src/viewer.rs | 8 ++++---- src_viewer/viewer.rs | 27 +++++++++++++++++---------- 3 files changed, 27 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 526536f0..8a05f32a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,13 +6,14 @@ PGS pass, instead of only during the once-per-substep stabilization sweep (rapier's `friction_in_bias_pass`). Off by default; on, friction gets as many iterations as the normal rows, which keeps pinch grasps and resting stacks from drifting. -- Viewer: sensor cameras render with 4x MSAA and hard shadow edges by default - (`set_sensor_antialiasing`, `set_sensor_shadow_softness`); the window's shadows are hard-edged - too (`set_shadow_softness`); the sharpest directional cascade covers the first 3 m instead of +- Viewer: sensor cameras render with 4x MSAA by default (`set_sensor_antialiasing`) and expose + the shadow-edge softness (`set_sensor_shadow_softness`, `set_shadow_softness` for the window; + kiss3d's penumbra stays the default); the sharpest directional cascade covers the first 3 m instead of 12 (`set_sensor_shadow_range`, `set_shadow_range`), the shadow map is 4096 texels over a 4-layer atlas instead of 2048 over 16 (`set_sensor_shadow_resolution`, `set_shadow_resolution`), - and `set_body_casts_shadows` excludes a body from the shadow map. Needs kiss3d's offscreen - MSAA, cascade and atlas-layer controls. + `set_body_casts_shadows` excludes a body from the shadow map, and textures load with mip + chains and 16x anisotropic filtering. Needs kiss3d's offscreen MSAA, cascade, atlas-layer and + anisotropy controls. - Python: `NexusState.set_rbd_solver_params(friction_in_bias_pass=...)`, `ColliderBuilder.friction_combine_rule`, and the `debug_contacts` / `debug_multibody_contact_impulses` GPU readbacks for contact diagnostics. diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index 2336aed1..d2798fa2 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -369,14 +369,14 @@ impl NexusViewer { } /// Shadow-edge softness of the sensor cameras' shaded renders, existing - /// and future ones (`0.0` hard edges, the default; `1.0` the PCF - /// penumbra). + /// and future ones (`1.0` the PCF penumbra, the default; `0.0` hard + /// edges). fn set_sensor_shadow_softness(&mut self, softness: f32) { self.inner_mut().set_sensor_shadow_softness(softness); } - /// Shadow-edge softness of the main window's shaded render (`0.0` hard - /// edges, the default; `1.0` the PCF penumbra). + /// Shadow-edge softness of the main window's shaded render (`1.0` the + /// PCF penumbra, the default; `0.0` hard edges). fn set_shadow_softness(&mut self, softness: f32) { self.inner_mut().set_shadow_softness(softness); } diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 1e70f409..10376e35 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -174,8 +174,9 @@ const COMPILE_BANNER_PRESENT_FRAMES: u32 = 10; /// MSAA sample count of new sensor cameras. #[cfg(feature = "dim3")] const DEFAULT_SENSOR_SAMPLES: u32 = 4; -/// Shadow-edge softness of the window and of new sensor cameras (hard edges). -const DEFAULT_SHADOW_SOFTNESS: f32 = 0.0; +/// Shadow-edge softness of the window and of new sensor cameras: kiss3d's +/// default PCF penumbra (about five shadow texels wide). +const DEFAULT_SHADOW_SOFTNESS: f32 = 1.0; /// Far bound (meters) of the highest-resolution directional shadow cascade of /// the window and of new sensor cameras: desk-sized scenes, not landscapes. const DEFAULT_SHADOW_FIRST_CASCADE: f32 = 3.0; @@ -261,8 +262,8 @@ pub struct NexusViewer { /// MSAA sample count given to every new sensor camera (default 4). #[cfg(feature = "dim3")] sensor_samples: u32, - /// Shadow-edge softness given to every new sensor camera (default 0.0, - /// hard edges; 1.0 is kiss3d's default penumbra). + /// Shadow-edge softness given to every new sensor camera (default 1.0, + /// kiss3d's PCF penumbra; 0.0 gives hard edges). #[cfg(feature = "dim3")] sensor_shadow_softness: f32, /// Directional-shadow cascade layout given to every new sensor camera: @@ -320,14 +321,20 @@ impl NexusViewer { // Disable MSAA, this puts extra load on the GPU that ends up // falsifying the gpu physics timestamps. window.set_samples(NumSamples::One); - // Hard shadow edges and a short first cascade: crisper contact shadows - // on the small geometry physics scenes are made of. + // A short first cascade: crisper contact shadows on the small geometry + // physics scenes are made of. window.set_shadow_softness(DEFAULT_SHADOW_SOFTNESS); window.set_first_cascade_far_bound(DEFAULT_SHADOW_FIRST_CASCADE); // Physics scenes light with directional lights only: four cascade // layers instead of kiss3d's sixteen pay for a finer shadow map. window.set_shadow_atlas_layers(DEFAULT_SHADOW_ATLAS_LAYERS); window.set_shadow_resolution(DEFAULT_SHADOW_RESOLUTION); + // Textures (visual meshes, labels, floors) get mip chains and + // anisotropic filtering so they stay sharp at grazing angles. + kiss3d::resource::TextureManager::get_global_manager(|tm| { + tm.set_generate_mipmaps(true); + tm.set_anisotropy(16); + }); #[cfg(feature = "dim2")] let (camera2d, camera3d) = { @@ -848,8 +855,8 @@ impl NexusViewer { } /// Shadow-edge softness of the sensor cameras' shaded renders, existing - /// and future ones (`0.0` hard edges, the default; `1.0` kiss3d's PCF - /// penumbra). + /// and future ones (`1.0` kiss3d's PCF penumbra, the default; `0.0` hard + /// edges). #[cfg(feature = "dim3")] pub fn set_sensor_shadow_softness(&mut self, softness: f32) { self.sensor_shadow_softness = softness.max(0.0); @@ -858,8 +865,8 @@ impl NexusViewer { } } - /// Shadow-edge softness of the main window's shaded render (`0.0` hard - /// edges, the default; `1.0` kiss3d's PCF penumbra). + /// Shadow-edge softness of the main window's shaded render (`1.0` + /// kiss3d's PCF penumbra, the default; `0.0` hard edges). pub fn set_shadow_softness(&mut self, softness: f32) { self.window.set_shadow_softness(softness.max(0.0)); } From 09e6e30ef78410a8dccaae75d30b19fa18205470 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 30 Aug 2026 11:14:50 +0200 Subject: [PATCH 11/25] feat(python): expose the implicit Coriolis toggle --- CHANGELOG.md | 2 ++ crates/nexus_python3d/src/nexus.rs | 8 ++++++++ src/state.rs | 30 +++++++++++++++++++++++++++++- 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a05f32a..c40f52f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,8 @@ `set_body_casts_shadows` excludes a body from the shadow map, and textures load with mip chains and 16x anisotropic filtering. Needs kiss3d's offscreen MSAA, cascade, atlas-layer and anisotropy controls. +- `NexusState::set_rbd_implicit_coriolis` (Python: `set_rbd_implicit_coriolis`), also honored + when the state is finalized after the call; `rbd_solver_params` reports it. - Python: `NexusState.set_rbd_solver_params(friction_in_bias_pass=...)`, `ColliderBuilder.friction_combine_rule`, and the `debug_contacts` / `debug_multibody_contact_impulses` GPU readbacks for contact diagnostics. diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 0ebf216c..45da4028 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -1207,6 +1207,13 @@ impl NexusState { } } + /// Implicit (default) or explicit treatment of the robots' Coriolis and + /// gyroscopic terms. Explicit terms refresh the multibody mass matrix once + /// per step instead of every substep, which is cheaper on the GPU. + fn set_rbd_implicit_coriolis(&mut self, viewer: PyRef, enabled: bool) { + self.0.set_rbd_implicit_coriolis(viewer.backend(), enabled); + } + /// The contact-solver parameters of environment 0 as a dict (see /// `set_rbd_solver_params`), for attestation. fn rbd_solver_params<'py>(&self, py: Python<'py>) -> PyResult> { @@ -1233,6 +1240,7 @@ impl NexusState { dict.set_item("prediction_distance", p.normalized_prediction_distance)?; dict.set_item("internal_pgs_iterations", p.num_internal_pgs_iterations)?; dict.set_item("friction_in_bias_pass", p.friction_in_bias_pass != 0)?; + dict.set_item("implicit_coriolis", self.0.rbd_implicit_coriolis())?; } Ok(dict) } diff --git a/src/state.rs b/src/state.rs index 4d4d9d10..2391ae29 100644 --- a/src/state.rs +++ b/src/state.rs @@ -212,6 +212,11 @@ pub struct NexusState { rbd_envs: Vec, /// Per-environment simulation parameters (same length as `rbd_envs`). rbd_sim_params: Vec, + /// Whether the multibody solver folds the Coriolis and gyroscopic terms + /// into the mass matrix (refreshed every substep) or applies them + /// explicitly (mass matrix refreshed once per step). Applied at + /// [`Self::finalize`] and by [`Self::set_rbd_implicit_coriolis`]. + rbd_implicit_coriolis: bool, /// Set whenever the rapier worlds change; consumed by [`Self::finalize`] to /// decide whether the GPU [`RbdState`] needs rebuilding. rbd_dirty: bool, @@ -249,6 +254,7 @@ impl NexusState { run_stats: RunStats::default(), rbd_envs: vec![PhysicsWorld::default()], rbd_sim_params: vec![RbdSimParams::tgs_soft()], + rbd_implicit_coriolis: true, rbd_dirty: false, rbd_steps_per_frame: 1, compute_graphs: false, @@ -393,6 +399,23 @@ impl NexusState { } } + /// Enables or disables the implicit treatment of the multibody Coriolis + /// and gyroscopic terms (default: enabled). Explicit terms skip the + /// per-substep mass-matrix refresh, the main cost of the implicit path. + #[cfg(all(feature = "rbd", feature = "dim3"))] + pub fn set_rbd_implicit_coriolis(&mut self, backend: &GpuBackend, enabled: bool) { + self.rbd_implicit_coriolis = enabled; + if let Some(rbd) = self.rbd.as_mut() { + rbd.set_implicit_coriolis(backend, enabled); + } + } + + /// Whether the multibody Coriolis terms are treated implicitly. + #[cfg(all(feature = "rbd", feature = "dim3"))] + pub fn rbd_implicit_coriolis(&self) -> bool { + self.rbd_implicit_coriolis + } + // ── Rigid-body runtime settings ───────────────────────────────────── /// Overrides the per-environment collision-pair capacity used when the GPU @@ -1364,7 +1387,8 @@ impl NexusState { // reservation (`reserve_rigid_bodies`) the buffers are sized for // spare slots so later `add_rigid_body` calls can append in place; // otherwise the state is sized exactly to the current body count. - let rbd_state = if self.rbd_reserve_per_env > 0 { + #[cfg_attr(not(feature = "dim3"), allow(unused_mut))] + let mut rbd_state = if self.rbd_reserve_per_env > 0 { let num_envs = self.rbd_envs.len() as u32; let max_count = self .rbd_envs @@ -1479,6 +1503,10 @@ impl NexusState { } } } + #[cfg(feature = "dim3")] + if !self.rbd_implicit_coriolis { + rbd_state.set_implicit_coriolis(backend, false); + } self.rbd = Some(rbd_state); self.rbd_dirty = false; } From 2c2933a0207be0e4a1979257c893cfdf67460f4f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 30 Aug 2026 13:36:48 +0200 Subject: [PATCH 12/25] feat(python): expose the multibody substep refresh cadence --- CHANGELOG.md | 5 +++-- crates/nexus_python3d/src/nexus.rs | 14 ++++++++++++++ src/state.rs | 29 +++++++++++++++++++++++++++++ src_rbd/pipeline/rbd_state.rs | 12 ++++++++++++ 4 files changed, 58 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c40f52f9..19e8d471 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,8 +14,9 @@ `set_body_casts_shadows` excludes a body from the shadow map, and textures load with mip chains and 16x anisotropic filtering. Needs kiss3d's offscreen MSAA, cascade, atlas-layer and anisotropy controls. -- `NexusState::set_rbd_implicit_coriolis` (Python: `set_rbd_implicit_coriolis`), also honored - when the state is finalized after the call; `rbd_solver_params` reports it. +- `NexusState::set_rbd_implicit_coriolis` and `set_rbd_substep_refresh` (Python: + `set_rbd_implicit_coriolis`, `set_rbd_substep_refresh`), also honored when the state is + finalized after the call; `rbd_solver_params` reports them. - Python: `NexusState.set_rbd_solver_params(friction_in_bias_pass=...)`, `ColliderBuilder.friction_combine_rule`, and the `debug_contacts` / `debug_multibody_contact_impulses` GPU readbacks for contact diagnostics. diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 45da4028..1b23a457 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -1214,6 +1214,17 @@ impl NexusState { self.0.set_rbd_implicit_coriolis(viewer.backend(), enabled); } + /// Multibody refresh cadence. `refresh` (default `True`) rebuilds the + /// robots' constraints, mass matrix and LU factors every substep; `False` + /// builds them once per step and later substeps only refresh the joint + /// rhs and limit activity, which is cheaper. `light` keeps the constraints + /// per substep but the mass matrix per step (ignored while `refresh`). + #[pyo3(signature = (viewer, refresh, light=false))] + fn set_rbd_substep_refresh(&mut self, viewer: PyRef, refresh: bool, light: bool) { + let _ = viewer; + self.0.set_rbd_substep_refresh(refresh, light); + } + /// The contact-solver parameters of environment 0 as a dict (see /// `set_rbd_solver_params`), for attestation. fn rbd_solver_params<'py>(&self, py: Python<'py>) -> PyResult> { @@ -1241,6 +1252,9 @@ impl NexusState { dict.set_item("internal_pgs_iterations", p.num_internal_pgs_iterations)?; dict.set_item("friction_in_bias_pass", p.friction_in_bias_pass != 0)?; dict.set_item("implicit_coriolis", self.0.rbd_implicit_coriolis())?; + let (refresh, light) = self.0.rbd_substep_refresh(); + dict.set_item("substep_refresh", refresh)?; + dict.set_item("substep_refresh_light", light)?; } Ok(dict) } diff --git a/src/state.rs b/src/state.rs index 2391ae29..974d1a6b 100644 --- a/src/state.rs +++ b/src/state.rs @@ -217,6 +217,9 @@ pub struct NexusState { /// explicitly (mass matrix refreshed once per step). Applied at /// [`Self::finalize`] and by [`Self::set_rbd_implicit_coriolis`]. rbd_implicit_coriolis: bool, + /// Multibody refresh cadence `(every substep, light)`; see + /// [`Self::set_rbd_substep_refresh`]. + rbd_substep_refresh: (bool, bool), /// Set whenever the rapier worlds change; consumed by [`Self::finalize`] to /// decide whether the GPU [`RbdState`] needs rebuilding. rbd_dirty: bool, @@ -255,6 +258,7 @@ impl NexusState { rbd_envs: vec![PhysicsWorld::default()], rbd_sim_params: vec![RbdSimParams::tgs_soft()], rbd_implicit_coriolis: true, + rbd_substep_refresh: (true, false), rbd_dirty: false, rbd_steps_per_frame: 1, compute_graphs: false, @@ -416,6 +420,26 @@ impl NexusState { self.rbd_implicit_coriolis } + /// Multibody refresh cadence. `refresh` (default `true`) rebuilds the + /// constraints, mass matrix and LU factors every substep; off, once per + /// step, with later substeps only refreshing the joint rhs and limit + /// activity (closer to how MuJoCo and Genesis step, and cheaper). `light` + /// keeps the constraints per substep but the mass matrix per step; ignored + /// while `refresh` is on. + #[cfg(all(feature = "rbd", feature = "dim3"))] + pub fn set_rbd_substep_refresh(&mut self, refresh: bool, light: bool) { + self.rbd_substep_refresh = (refresh, light); + if let Some(rbd) = self.rbd.as_mut() { + rbd.set_substep_refresh(refresh, light); + } + } + + /// The multibody refresh cadence as `(every substep, light)`. + #[cfg(all(feature = "rbd", feature = "dim3"))] + pub fn rbd_substep_refresh(&self) -> (bool, bool) { + self.rbd_substep_refresh + } + // ── Rigid-body runtime settings ───────────────────────────────────── /// Overrides the per-environment collision-pair capacity used when the GPU @@ -1507,6 +1531,11 @@ impl NexusState { if !self.rbd_implicit_coriolis { rbd_state.set_implicit_coriolis(backend, false); } + #[cfg(feature = "dim3")] + if self.rbd_substep_refresh != (true, false) { + rbd_state + .set_substep_refresh(self.rbd_substep_refresh.0, self.rbd_substep_refresh.1); + } self.rbd = Some(rbd_state); self.rbd_dirty = false; } diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 94eedd23..40443b28 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -512,6 +512,18 @@ impl RbdState { self.rebuild_batch_indices(backend); } + /// Sets the multibody refresh cadence: `refresh` rebuilds the joint and + /// contact constraints, mass matrix and LU factors every substep (the + /// default); off, they are built once per step and later substeps only + /// refresh the joint rhs and limit activity. `light` (ignored while + /// `refresh` is on) keeps the constraints per substep but the mass matrix + /// per step. Pure dispatch gating: no GPU layout changes. + #[cfg(feature = "dim3")] + pub fn set_substep_refresh(&mut self, refresh: bool, light: bool) { + self.multibodies.set_substep_refresh(refresh); + self.multibodies.set_substep_refresh_light(light); + } + /// Sets the per-DoF dry joint friction (N·m). #[cfg(feature = "dim3")] pub fn set_dof_frictionloss(&mut self, backend: &GpuBackend, values: &[f32]) { From 2e2be809851db29aa1f87b0a7ac4c0c239b50689 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 30 Aug 2026 15:42:27 +0200 Subject: [PATCH 13/25] fix(rbd): warmstart each contact point from its nearest previous point --- CHANGELOG.md | 4 ++++ .../dynamics/multibody/contact_constraints.rs | 21 +++++++++++++----- src_rbd_shaders/dynamics/warmstart.rs | 22 +++++++++++++------ 3 files changed, 34 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 19e8d471..15755bf0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,10 @@ ### Fixed +- Contact warmstarting matched each new contact point to the *first* old point of the same pair + within 10 cm, so every point of a small manifold (a fingertip pad) inherited the first point's + normal and friction impulses and the solver had to redistribute them each step. Both the + rigid-body and multibody transfers now take the nearest old point. - Contacts between a multibody link and a rigid body were solved twice: by the multibody contact solver and, again, by the rigid-body solver against a zero-inverse-mass copy of the link. The second copy saw the link as never moving, so a robot lifting a grasped object had diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 420e0f09..23c2fca0 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -1055,8 +1055,10 @@ pub fn gpu_mb_transfer_contact_warmstart( #[spirv(uniform, descriptor_set = 0, binding = 4)] softness: &ConstraintSoftness, ) { const LANES: u32 = 64; - // Anchors this far apart (in each side's own frame) are taken to be the - // same contact point from one frame to the next. + // Anchors further apart than this (in each side's own frame) are never the + // same contact point from one frame to the next; among the candidates the + // nearest wins, so the points of one small manifold (a fingertip pad) each + // recover their own impulse instead of all copying the first one. const MATCH_DIST: f32 = 1.0e-1; let batch_id = workgroup_id.y; @@ -1095,6 +1097,9 @@ pub fn gpu_mb_transfer_contact_warmstart( continue; } + // Nearest old point of the same link pair within the threshold. + let mut best_j = u32::MAX; + let mut best_sq = sq_threshold; for j in 0..old_count { let old = old_contact_constraints.read(old_base + j as usize); if old.kind != MB_CONTACT_KIND_NORMAL @@ -1106,10 +1111,15 @@ pub fn gpu_mb_transfer_contact_warmstart( } let d1 = old.local_p1 - cons.local_p1; let d2 = old.local_p2 - cons.local_p2; - if d1.dot(d1) >= sq_threshold || d2.dot(d2) >= sq_threshold { - continue; + let sq = d1.dot(d1).max(d2.dot(d2)); + if sq < best_sq { + best_sq = sq; + best_j = j; } - + } + if best_j != u32::MAX { + let j = best_j; + let old = old_contact_constraints.read(old_base + j as usize); cons.impulse = old.impulse * warmstart_coeff; contact_constraints.write(cons_base + s as usize, cons); @@ -1135,7 +1145,6 @@ pub fn gpu_mb_transfer_contact_warmstart( new_t0.impulse = old_t0.impulse * warmstart_coeff; contact_constraints.write(cons_base + (s + 1) as usize, new_t0); } - break; } } } diff --git a/src_rbd_shaders/dynamics/warmstart.rs b/src_rbd_shaders/dynamics/warmstart.rs index d590a026..9925dc7a 100644 --- a/src_rbd_shaders/dynamics/warmstart.rs +++ b/src_rbd_shaders/dynamics/warmstart.rs @@ -202,26 +202,34 @@ pub fn transfer_warmstart_impulses( let dist_threshold = 1.0e-1; // 10cm let sq_threshold = dist_threshold * dist_threshold; - // Try to match each new contact point with old contact points + // Match each new contact point with the NEAREST old contact point + // within the threshold: the points of one small manifold are all + // within 10cm of each other, so taking the first candidate would hand + // every point the same (first) impulse. for k_new in 0..(new_constraints[i].len as usize) { let pt_new_a = new_constraint_builders[i].infos.at(k_new).local_pt_a; let pt_new_b = new_constraint_builders[i].infos.at(k_new).local_pt_b; - // Search through old contact points for a match + let mut best_k_old = usize::MAX; + let mut best_sq = sq_threshold; for k_old in 0..(old_constraints[cid_old].len as usize) { let pt_old_a = old_constraint_builders[cid_old].infos.at(k_old).local_pt_a; let pt_old_b = old_constraint_builders[cid_old].infos.at(k_old).local_pt_b; - - // Compute distance between contact points in local space let dpt_a = pt_old_a - pt_new_a; let dpt_b = pt_old_b - pt_new_b; + let sq = dpt_a.dot(dpt_a).max(dpt_b.dot(dpt_b)); + if sq < best_sq { + best_sq = sq; + best_k_old = k_old; + } + } - // If both points are close enough, consider it a match - if dpt_a.dot(dpt_a) < sq_threshold && dpt_b.dot(dpt_b) < sq_threshold { + { + let k_old = best_k_old; + if k_old != usize::MAX { // Contact point match found! Transfer the accumulated impulse. // The impulse field contains the last substep's impulse, which serves // as the warmstart value for this frame. - // TODO: what if we have multiple matches? (currently uses first match) new_constraints[i] .elements .at_mut(k_new) From bb1182e942f490e3bcbaf36334dabf17d38d6443 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 3 Sep 2026 20:39:10 +0200 Subject: [PATCH 14/25] feat(python): generalized velocity readback and per-channel sensor ambient --- CHANGELOG.md | 6 +++++ Cargo.toml | 4 +-- crates/nexus_python3d/src/nexus.rs | 26 ++++++++++++++++++++ crates/nexus_python3d/src/viewer.rs | 8 ++++++ src/state.rs | 26 ++++++++++++++++++++ src_rbd/dynamics/multibody/multibody_set.rs | 27 ++++++++++++++++++++- src_viewer/sensors.rs | 10 ++++++++ 7 files changed, 104 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 15755bf0..e6837bbc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,12 @@ - Python: `NexusState.set_rbd_solver_params(friction_in_bias_pass=...)`, `ColliderBuilder.friction_combine_rule`, and the `debug_contacts` / `debug_multibody_contact_impulses` GPU readbacks for contact diagnostics. +- `GpuMultibodySet::read_dof_velocities` and `NexusState::multibody_joint_velocities` read the + generalized velocities of a batch (resp. of one multibody) back from the GPU, in the same + assembly order `read_dof_coords` and `multibody_joint_positions` use. Python: `robot_qvel`, + the joint-velocity counterpart of `robot_state`'s `qpos`. +- Viewer: `SensorCamera::set_ambient_color` (Python: `set_sensor_camera_ambient_color`) tints a + sensor camera's ambient fill light, whose brightness `set_sensor_camera_ambient` already set. ### Fixed diff --git a/Cargo.toml b/Cargo.toml index f6b140f1..2da8b83a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,8 +46,8 @@ 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 } +parry2d = { version = "0.31", default-features = false } +parry3d = { version = "0.31", default-features = false } # Viewer / examples deps kiss3d = "0.46.0" diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 1b23a457..1361d8cb 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -806,6 +806,32 @@ impl NexusState { )) } + /// Reads `robot`'s simulated generalized velocities back from the GPU as + /// `qvel (n_dofs,)`, ordered exactly like `robot_state`'s `qpos`. + fn robot_qvel<'py>( + &self, + py: Python<'py>, + viewer: PyRef, + robot: PyRef, + ) -> PyResult>> { + let qvel = pollster::block_on(self.0.multibody_joint_velocities( + viewer.backend(), + robot.env, + robot.root, + )) + .ok_or_else(|| { + PyRuntimeError::new_err("robot velocities unavailable: call finalize() first") + })?; + if qvel.len() != robot.dof_axes.len() { + return Err(PyRuntimeError::new_err(format!( + "{} velocities read back for {} dofs", + qvel.len(), + robot.dof_axes.len() + ))); + } + Ok(PyArray1::from_vec(py, qvel)) + } + /// Sets `robot`'s generalized coordinates and zeroes its joint velocities /// on the GPU (between steps, after `finalize`). fn set_robot_qpos( diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index d2798fa2..6809aa80 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -256,6 +256,14 @@ impl NexusViewer { } } + /// Ambient light color (RGB) of sensor camera `id`'s shaded render, a tint + /// on the fill light `set_sensor_camera_ambient` scales. + fn set_sensor_camera_ambient_color(&mut self, id: usize, rgb: [f32; 3]) { + if let Some(sensor) = self.inner_mut().sensor_camera_mut(id) { + sensor.set_ambient_color(rgb); + } + } + /// Background color of sensor camera `id`'s shaded render. fn set_sensor_camera_background(&mut self, id: usize, rgba: [f32; 4]) { if let Some(sensor) = self.inner_mut().sensor_camera_mut(id) { diff --git a/src/state.rs b/src/state.rs index 974d1a6b..c1ff8068 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1008,6 +1008,32 @@ impl NexusState { Ok(()) } + /// Generalized velocities of the multibody containing `body` in environment + /// `env`, read back from the GPU in the same order as + /// [`Self::multibody_joint_positions`]'s coordinates. `None` before + /// `finalize`, or when `body` is not part of a multibody. + #[cfg(feature = "dim3")] + pub async fn multibody_joint_velocities( + &self, + backend: &GpuBackend, + env: usize, + body: RigidBodyHandle, + ) -> Option> { + let (_, info, _) = self.multibody_link_slots(env, body)?; + // The readback covers the whole batch, so slice out this multibody's + // DoF range. + let vels = self + .rbd + .as_ref()? + .multibodies() + .read_dof_velocities(backend, env as u32) + .await + .ok()?; + let first = info.first_dof as usize; + vels.get(first..first + info.ndofs as usize) + .map(<[f32]>::to_vec) + } + /// Mutable access to environment `env`'s rapier world that does **not** mark /// the rbd state dirty, for use after [`Self::finalize`]. /// diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index fa6b8da7..1c7a2373 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -1271,7 +1271,7 @@ impl GpuMultibodySet { let a = WsAddr::new(0, self.num_batches, batch_id); let mut out = Vec::new(); for k in 0..self.links_per_batch { - let stat = &self.links_static_mirror[(batch_id * self.links_per_batch + k) as usize]; + let stat = &self.links_static_mirror[(k * self.num_batches + batch_id) as usize]; let locked = stat.data.locked_axes; for axis in 0..6u32 { if locked & (1 << axis) == 0 { @@ -1282,6 +1282,31 @@ impl GpuMultibodySet { Ok(out) } + /// Reads back the generalized velocity of every DoF of batch `batch_id`, in + /// assembly order (the same order as [`Self::read_dof_coords`]'s + /// coordinates). Only the velocity section of [`Self::dof_state`] is read, + /// and it is DoF-major, so this de-interleaves the batch for callers. + pub async fn read_dof_velocities( + &self, + backend: &GpuBackend, + batch_id: u32, + ) -> Result, GpuBackendError> { + let state: Vec = backend.slow_read_vec(self.dof_state.buffer()).await?; + let nb = self.num_batches as usize; + let mut out = Vec::new(); + for k in 0..self.links_per_batch { + let stat = &self.links_static_mirror[(k * self.num_batches + batch_id) as usize]; + let locked = stat.data.locked_axes; + for axis in 0..6u32 { + if locked & (1 << axis) == 0 { + let dof = out.len(); + out.push(state[dof * nb + batch_id as usize]); + } + } + } + Ok(out) + } + /// Upload a new integration timestep. pub fn set_dt(&mut self, backend: &GpuBackend, dt: f32) { self.dt = Tensor::scalar( diff --git a/src_viewer/sensors.rs b/src_viewer/sensors.rs index f5f83c52..11af3c32 100644 --- a/src_viewer/sensors.rs +++ b/src_viewer/sensors.rs @@ -200,6 +200,16 @@ impl SensorCamera { self.surface.set_ambient(ambient); } + /// Ambient light color (RGB) of the shaded render. The ambient term is + /// `ambient_color * ambient * albedo * ao`, so this tints the fill light + /// that [`Self::set_ambient`] scales; the surface only forwards the + /// intensity, so the color goes through the underlying window. + pub fn set_ambient_color(&mut self, rgb: [f32; 3]) { + self.surface + .window_mut() + .set_ambient_color(Color::new(rgb[0], rgb[1], rgb[2], 1.0)); + } + /// Background color (RGBA) of the shaded render. pub fn set_background_color(&mut self, rgba: [f32; 4]) { self.surface From f6bb36344814e07059b65df9f017acffc52fceb0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 20 Sep 2026 15:40:00 +0200 Subject: [PATCH 15/25] build: patch rapier and kiss3d from their pushed branches by rev --- Cargo.toml | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2da8b83a..aa5765dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -102,14 +102,15 @@ khal-std = { git = "https://github.com/dimforge/khal", branch = "main" } khal-builder = { git = "https://github.com/dimforge/khal", branch = "main" } vortx = { git = "https://github.com/dimforge/vortx", branch = "main" } vortx-shaders = { git = "https://github.com/dimforge/vortx", branch = "main" } -# rapier3d-urdf's URDF roll-pitch-yaw fix (fix-urdf-rpy branch), pending a -# rapier release; every rapier crate rides the same source so there is one -# `rapier3d` in the graph. -rapier2d = { path = "../rapier/crates/rapier2d" } -rapier3d = { path = "../rapier/crates/rapier3d" } -rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" } -rapier3d-urdf = { path = "../rapier/crates/rapier3d-urdf" } -rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" } +# rapier3d-urdf's URDF roll-pitch-yaw fix and the Python multibody bindings +# (branch python-multibody-control), pending a rapier release; every rapier +# crate rides the same source so there is one `rapier3d` in the graph. Pinned +# by rev, not branch: a moving head would make this build irreproducible. +rapier2d = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } +rapier3d = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } +rapier3d-mjcf = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } +rapier3d-urdf = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } +rapier3d-meshloader = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } ## Local glam clone with SPIR-V vector-arithmetic intrinsics (Vec3 add/sub/mul/scale). #glam = { path = "../glam-rs" } # 30% faster for loop in P2G @@ -127,9 +128,10 @@ rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" } #vortx-shaders = { git = "https://github.com/dimforge/vortx", branch = "int-reduce" } #glamx = { path = "../glamx" } -# kiss3d's shared-manager fix for a second offscreen surface (branch -# fix-shared-window-managers), pending a kiss3d release. -kiss3d = { path = "../kiss3d" } +# kiss3d's shared-manager fix for a second offscreen surface, its offscreen +# MSAA and shadow controls (branch fix-shared-window-managers), pending a +# kiss3d release. Pinned by rev for the same reason as rapier above. +kiss3d = { git = "https://github.com/dimforge/kiss3d", rev = "28cdddde9e920fdb5462e9c24e44b57e55dd7bbe" } #parry2d = { path = "../parry/crates/parry2d" } #parry3d = { path = "../parry/crates/parry3d" } #rapier2d = { path = "../rapier/crates/rapier2d" } From e71563aff0e414811af46309f7afd505846240f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 25 Sep 2026 18:02:00 +0200 Subject: [PATCH 16/25] build: move to the published rapier 0.36.0 --- Cargo.toml | 20 ++++++-------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index aa5765dc..b63db59e 100644 --- a/Cargo.toml +++ b/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.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"] } +rapier2d = { version = "0.36.0", default-features = false } +rapier3d = { version = "0.36.0", default-features = false } +rapier3d-urdf = "0.36.0" +rapier3d-mjcf = { version = "0.36.0", features = ["stl", "wavefront", "msh"] } parry2d = { version = "0.31", default-features = false } parry3d = { version = "0.31", default-features = false } @@ -102,15 +102,6 @@ khal-std = { git = "https://github.com/dimforge/khal", branch = "main" } khal-builder = { git = "https://github.com/dimforge/khal", branch = "main" } vortx = { git = "https://github.com/dimforge/vortx", branch = "main" } vortx-shaders = { git = "https://github.com/dimforge/vortx", branch = "main" } -# rapier3d-urdf's URDF roll-pitch-yaw fix and the Python multibody bindings -# (branch python-multibody-control), pending a rapier release; every rapier -# crate rides the same source so there is one `rapier3d` in the graph. Pinned -# by rev, not branch: a moving head would make this build irreproducible. -rapier2d = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } -rapier3d = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } -rapier3d-mjcf = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } -rapier3d-urdf = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } -rapier3d-meshloader = { git = "https://github.com/dimforge/rapier", rev = "6a047ed3d174e0ddf192ce4a478546a3f5aea0fb" } ## Local glam clone with SPIR-V vector-arithmetic intrinsics (Vec3 add/sub/mul/scale). #glam = { path = "../glam-rs" } # 30% faster for loop in P2G @@ -130,7 +121,8 @@ rapier3d-meshloader = { git = "https://github.com/dimforge/rapier", rev = "6a047 #glamx = { path = "../glamx" } # kiss3d's shared-manager fix for a second offscreen surface, its offscreen # MSAA and shadow controls (branch fix-shared-window-managers), pending a -# kiss3d release. Pinned by rev for the same reason as rapier above. +# kiss3d release. Pinned by rev, not branch: a moving head would make this +# build irreproducible. kiss3d = { git = "https://github.com/dimforge/kiss3d", rev = "28cdddde9e920fdb5462e9c24e44b57e55dd7bbe" } #parry2d = { path = "../parry/crates/parry2d" } #parry3d = { path = "../parry/crates/parry3d" } From fadd1249b770ef3b144cb7af67fe4acd38e7d1b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 25 Sep 2026 18:09:50 +0200 Subject: [PATCH 17/25] feat(viewer): remove sensor cameras and free their GPU resources --- crates/nexus_python3d/src/viewer.rs | 13 ++++++++ src_viewer/viewer.rs | 47 +++++++++++++++++------------ 2 files changed, 41 insertions(+), 19 deletions(-) diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index 6809aa80..8468c6bb 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -214,6 +214,19 @@ impl NexusViewer { )) } + /// Removes sensor camera `id` and frees its GPU resources. The id is not + /// reused; rendering it afterwards raises. Returns whether a camera was + /// removed. Each camera holds its own render targets and shadow atlas, so a + /// process that builds many scenes must release the ones it is done with. + fn remove_sensor_camera(&mut self, id: usize) -> bool { + self.inner_mut().remove_sensor_camera(id) + } + + /// Number of live (not removed) sensor cameras. + fn num_sensor_cameras(&self) -> usize { + self.inner().num_sensor_cameras() + } + /// Sets sensor camera `id`'s pose: position and `(w, x, y, z)` quaternion /// of the OpenGL camera frame (looks down its -Z axis, +Y up). fn set_sensor_camera_pose(&mut self, id: usize, pos: [f32; 3], quat: [f32; 4]) { diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 10376e35..6e3a67f7 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -254,11 +254,12 @@ pub struct NexusViewer { /// frames so samples keep accumulating while the scene is static. #[cfg(feature = "dim3")] raytracer: Option, - /// Offscreen sensor cameras (see [`crate::sensors`]). While any exists, + /// Offscreen sensor cameras (see [`crate::sensors`]), indexed by id; a + /// removed camera leaves `None` so later ids stay stable. While any exists, /// `sync` takes the readback path so attached cameras can follow their body /// and the per-object passes see world-posed nodes. #[cfg(feature = "dim3")] - sensors: Vec, + sensors: Vec>, /// MSAA sample count given to every new sensor camera (default 4). #[cfg(feature = "dim3")] sensor_samples: u32, @@ -550,7 +551,7 @@ impl NexusViewer { fn sensors_active(&self) -> bool { #[cfg(feature = "dim3")] { - !self.sensors.is_empty() + self.sensors.iter().any(Option::is_some) } #[cfg(not(feature = "dim3"))] { @@ -840,16 +841,24 @@ impl NexusViewer { self.sensor_shadow_resolution.0, self.sensor_shadow_resolution.1, ); - self.sensors.push(sensor); + self.sensors.push(Some(sensor)); self.sensors.len() - 1 } + /// Removes sensor camera `id` and frees its GPU resources (render targets, + /// shadow atlas). Its id is not reused, and later calls with it find no + /// camera. Returns whether a camera was removed. + #[cfg(feature = "dim3")] + pub fn remove_sensor_camera(&mut self, id: usize) -> bool { + self.sensors.get_mut(id).and_then(Option::take).is_some() + } + /// MSAA sample count of the sensor cameras' shaded renders, existing and /// future ones (`1` disables antialiasing, `4` is the default). #[cfg(feature = "dim3")] pub fn set_sensor_antialiasing(&mut self, samples: u32) { self.sensor_samples = samples.max(1); - for sensor in &mut self.sensors { + for sensor in self.sensors.iter_mut().flatten() { sensor.set_samples(self.sensor_samples); } } @@ -860,7 +869,7 @@ impl NexusViewer { #[cfg(feature = "dim3")] pub fn set_sensor_shadow_softness(&mut self, softness: f32) { self.sensor_shadow_softness = softness.max(0.0); - for sensor in &mut self.sensors { + for sensor in self.sensors.iter_mut().flatten() { sensor.set_shadow_softness(self.sensor_shadow_softness); } } @@ -879,7 +888,7 @@ impl NexusViewer { #[cfg(feature = "dim3")] pub fn set_sensor_shadow_range(&mut self, first_cascade_far_bound: f32, shadow_distance: f32) { self.sensor_shadow_range = (first_cascade_far_bound.max(0.01), shadow_distance.max(0.0)); - for sensor in &mut self.sensors { + for sensor in self.sensors.iter_mut().flatten() { sensor.set_shadow_range(self.sensor_shadow_range.0, self.sensor_shadow_range.1); } } @@ -891,7 +900,7 @@ impl NexusViewer { #[cfg(feature = "dim3")] pub fn set_sensor_shadow_resolution(&mut self, resolution: u32, layers: u32) { self.sensor_shadow_resolution = (resolution.max(1), layers.max(1)); - for sensor in &mut self.sensors { + for sensor in self.sensors.iter_mut().flatten() { sensor.set_shadow_resolution( self.sensor_shadow_resolution.0, self.sensor_shadow_resolution.1, @@ -930,29 +939,29 @@ impl NexusViewer { self.nexus_render.active_generation } - /// Number of sensor cameras. + /// Number of live (not removed) sensor cameras. #[cfg(feature = "dim3")] pub fn num_sensor_cameras(&self) -> usize { - self.sensors.len() + self.sensors.iter().flatten().count() } /// The sensor camera at `id`. #[cfg(feature = "dim3")] pub fn sensor_camera(&self, id: usize) -> Option<&SensorCamera> { - self.sensors.get(id) + self.sensors.get(id)?.as_ref() } /// The sensor camera at `id`, mutably (pose, attachment, ambient, ...). #[cfg(feature = "dim3")] pub fn sensor_camera_mut(&mut self, id: usize) -> Option<&mut SensorCamera> { - self.sensors.get_mut(id) + self.sensors.get_mut(id)?.as_mut() } /// Sets sensor camera `id`'s pose (OpenGL convention: looking down local /// -Z, +Y up). Overridden at the next `sync` while the camera is attached. #[cfg(feature = "dim3")] pub fn set_sensor_camera_pose(&mut self, id: usize, pose: Pose) { - if let Some(sensor) = self.sensors.get_mut(id) { + if let Some(sensor) = self.sensor_camera_mut(id) { sensor.set_pose(pose); } } @@ -970,7 +979,7 @@ impl NexusViewer { local_pose: Pose, state: &NexusState, ) { - if let Some(sensor) = self.sensors.get_mut(id) { + if let Some(sensor) = self.sensor_camera_mut(id) { sensor.attach(env, handle.0, local_pose); } self.update_sensor_attachments(state); @@ -999,7 +1008,7 @@ impl NexusViewer { return; } let active = self.nexus_render.active_generation; - for sensor in &mut self.sensors { + for sensor in self.sensors.iter_mut().flatten() { if sensor.generation != active { continue; } @@ -1021,21 +1030,21 @@ impl NexusViewer { /// origin, `width * height * 3` bytes). #[cfg(feature = "dim3")] pub async fn render_sensor_rgb(&mut self, id: usize) -> Option> { - let sensor = self.sensors.get_mut(id)?; + let sensor = self.sensors.get_mut(id)?.as_mut()?; Some(sensor.render_rgb(&mut self.scene3d).await) } /// Renders sensor camera `id`'s linear metric depth (`0.0` = background). #[cfg(feature = "dim3")] pub fn render_sensor_depth(&mut self, id: usize) -> Option> { - let sensor = self.sensors.get_mut(id)?; + let sensor = self.sensors.get_mut(id)?.as_mut()?; Some(sensor.render_depth(&mut self.scene3d)) } /// Renders sensor camera `id`'s per-pixel segmentation ids (`0` = background). #[cfg(feature = "dim3")] pub fn render_sensor_segmentation(&mut self, id: usize) -> Option> { - let sensor = self.sensors.get_mut(id)?; + let sensor = self.sensors.get_mut(id)?.as_mut()?; Some(sensor.render_segmentation(&mut self.scene3d)) } @@ -1116,7 +1125,7 @@ impl NexusViewer { // their local poses are body-relative. Attached sensor cameras use // the same poses. #[cfg(feature = "dim3")] - if self.nexus_render.has_visual_nodes() || !self.sensors.is_empty() { + if self.nexus_render.has_visual_nodes() || self.sensors.iter().any(Option::is_some) { let body_poses = rbd.body_poses(); let mut body_cache = vec![Pose::default(); body_poses.len() as usize]; let _ = self From 2e357cb95123fa756a1585e4ba8da5e122ba6c5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 26 Sep 2026 22:32:22 +0200 Subject: [PATCH 18/25] fix(rbd): reset the multibody contact-index cursors after the scatter --- .../dynamics/multibody/contact_constraints.rs | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 23c2fca0..0d41e24d 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -249,8 +249,8 @@ pub fn gpu_mb_count_contact_constraints( /// (`contact_index_start/len`, unclamped: the index buffer is sized like the /// contacts buffer and each contact owns at most one entry). Also publishes /// the total slot demand for the host's auto-resize readback, and re-zeroes -/// `mb_cons_counts` (for the next frame) and `mb_index_counts` (which the -/// scatter pass reuses as its write cursors). Serial in one thread (the +/// `mb_cons_counts` for the next frame (`mb_index_counts` stays: the scatter +/// pass counts it down as its write cursors). Serial in one thread (the /// multibody count per scene is small). #[spirv_bindgen] #[spirv(compute(threads(1)))] @@ -301,17 +301,17 @@ pub fn gpu_mb_cons_offsets_scan( index_acc += mb.contact_index_len; multibody_info.write(i as usize, mb); - // Zeroed for the next frame's count pass / for the scatter cursors. + // Zeroed for the next frame's count pass. `mb_index_counts` is left as + // is: the scatter counts it back down to zero. mb_cons_counts.write(i as usize, 0); - mb_index_counts.write(i as usize, 0); } mb_cons_demand.write(0, demand); } /// Builds the contact→multibody index: one flat sweep over the contacts, /// each contact appending its entry to its owner's segment (laid out by the -/// offsets scan; `mb_index_counts` was re-zeroed there and serves as the -/// per-multibody write cursors). Entry order within a segment follows the +/// offsets scan; `mb_index_counts` still holds the per-multibody counts and +/// serves as write cursors counting down to zero). Entry order within a segment follows the /// atomic race; the emission's warmstart matching is key-based, so the order /// only affects the (already nondeterministic) impulse iteration order. /// Grid: `contacts_indirect`. @@ -339,7 +339,9 @@ pub fn gpu_mb_scatter_contact_index( } let slot = batch_ids.mbi(owner.batch, owner.mb as usize); let mb = multibody_info.read(slot); - let pos = atomic_add_u32(mb_index_counts.at_mut(slot), 1); + // Count the cursor down (`+ u32::MAX` wraps to `- 1`), so the pass + // leaves it at zero for the next frame's count pass. + let pos = atomic_add_u32(mb_index_counts.at_mut(slot), u32::MAX) - 1; mb_contact_index.write( (mb.contact_index_start + pos) as usize, MbContactIndexEntry { From 4b0c2ca20b45b99830844a7d768e214a303a70e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 27 Sep 2026 19:31:38 +0200 Subject: [PATCH 19/25] refactor(rbd): count the contact-index cursors down with atomic_sub_u32 --- .../dynamics/multibody/contact_constraints.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 0d41e24d..faac3aa0 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -18,7 +18,9 @@ use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::iter::StepRng; use khal_std::macros::{spirv, spirv_bindgen}; -use khal_std::sync::{atomic_add_u32, atomic_load_u32, workgroup_memory_barrier_with_group_sync}; +use khal_std::sync::{ + atomic_add_u32, atomic_load_u32, atomic_sub_u32, workgroup_memory_barrier_with_group_sync, +}; use crate::broad_phase::ContactPlan; use crate::dynamics::ConstraintSoftness; @@ -339,9 +341,9 @@ pub fn gpu_mb_scatter_contact_index( } let slot = batch_ids.mbi(owner.batch, owner.mb as usize); let mb = multibody_info.read(slot); - // Count the cursor down (`+ u32::MAX` wraps to `- 1`), so the pass - // leaves it at zero for the next frame's count pass. - let pos = atomic_add_u32(mb_index_counts.at_mut(slot), u32::MAX) - 1; + // Count the cursor down, so the pass leaves it at zero for the next + // frame's count pass. + let pos = atomic_sub_u32(mb_index_counts.at_mut(slot), 1) - 1; mb_contact_index.write( (mb.contact_index_start + pos) as usize, MbContactIndexEntry { From a6a63ace4050e272c1f83626479298ae7a87f25f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 1 Oct 2026 12:50:25 +0200 Subject: [PATCH 20/25] style: rustfmt the mujoco menagerie example --- crates/examples3d/rbd_mujoco_menagerie3.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/crates/examples3d/rbd_mujoco_menagerie3.rs b/crates/examples3d/rbd_mujoco_menagerie3.rs index b03a197f..b001d315 100644 --- a/crates/examples3d/rbd_mujoco_menagerie3.rs +++ b/crates/examples3d/rbd_mujoco_menagerie3.rs @@ -763,10 +763,8 @@ pub async fn run( let pgs_changed = settings.pgs_iterations != next.pgs_iterations; settings = next; if pgs_changed && !reload { - state.set_rbd_num_internal_pgs_iterations( - viewer.backend(), - settings.pgs_iterations, - ); + state + .set_rbd_num_internal_pgs_iterations(viewer.backend(), settings.pgs_iterations); } } // A model change always rebuilds, and resets the keyframe to the new From a1066be822f21745e50307db43b82bcab7a68cd5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 1 Oct 2026 12:54:56 +0200 Subject: [PATCH 21/25] fix: gate 3D-only state and doc links so the 2D crates build without warnings --- src/state.rs | 11 +++++++++-- src_viewer/viewer.rs | 4 ++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/state.rs b/src/state.rs index c1ff8068..6bf568d9 100644 --- a/src/state.rs +++ b/src/state.rs @@ -128,6 +128,7 @@ pub type MultibodySlots = ( ); /// Failure of [`NexusState::set_multibody_joint_positions`]. +#[cfg(feature = "dim3")] #[derive(Debug)] pub enum JointPositionsError { /// The GPU write failed. @@ -136,12 +137,14 @@ pub enum JointPositionsError { DofMismatch { expected: usize, got: usize }, } +#[cfg(feature = "dim3")] impl From for JointPositionsError { fn from(e: GpuBackendError) -> Self { JointPositionsError::Gpu(e) } } +#[cfg(feature = "dim3")] impl core::fmt::Display for JointPositionsError { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { match self { @@ -216,9 +219,11 @@ pub struct NexusState { /// into the mass matrix (refreshed every substep) or applies them /// explicitly (mass matrix refreshed once per step). Applied at /// [`Self::finalize`] and by [`Self::set_rbd_implicit_coriolis`]. + #[cfg(feature = "dim3")] rbd_implicit_coriolis: bool, /// Multibody refresh cadence `(every substep, light)`; see /// [`Self::set_rbd_substep_refresh`]. + #[cfg(feature = "dim3")] rbd_substep_refresh: (bool, bool), /// Set whenever the rapier worlds change; consumed by [`Self::finalize`] to /// decide whether the GPU [`RbdState`] needs rebuilding. @@ -257,7 +262,9 @@ impl NexusState { run_stats: RunStats::default(), rbd_envs: vec![PhysicsWorld::default()], rbd_sim_params: vec![RbdSimParams::tgs_soft()], + #[cfg(feature = "dim3")] rbd_implicit_coriolis: true, + #[cfg(feature = "dim3")] rbd_substep_refresh: (true, false), rbd_dirty: false, rbd_steps_per_frame: 1, @@ -696,7 +703,7 @@ impl NexusState { // --- rigid-body and multibody state access (between steps) ------------- /// GPU pose slot of `handle` in environment `env`: the index into - /// [`Self::read_rigid_body_poses`] and [`Self::read_rigid_body_velocities`]. + /// [`Self::read_rigid_body_poses`] (and, in 3D, `read_rigid_body_velocities`). /// `None` before `finalize` or for a handle this state does not know. pub fn rigid_body_gpu_index(&self, env: usize, handle: RigidBodyHandle) -> Option { self.rbd2gpu @@ -757,7 +764,7 @@ impl NexusState { /// Teleports a free rigid body: overwrites its world-origin pose on the GPU. /// Valid between steps, after `finalize`. A multibody link's pose is /// re-derived from its joint coordinates every step, so links go through - /// [`Self::set_multibody_joint_positions`] instead. Unknown handles are + /// `set_multibody_joint_positions` (3D) instead. Unknown handles are /// ignored. pub fn set_rigid_body_pose( &mut self, diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 6e3a67f7..5e863820 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -909,14 +909,14 @@ impl NexusViewer { } /// Shadow map resolution and atlas layer count of the main window's render - /// (see [`Self::set_sensor_shadow_resolution`]). + /// (see `set_sensor_shadow_resolution`, the 3D sensor-camera counterpart). pub fn set_shadow_resolution(&mut self, resolution: u32, layers: u32) { self.window.set_shadow_atlas_layers(layers.max(1)); self.window.set_shadow_resolution(resolution.max(1)); } /// Directional-shadow cascade layout of the main window's render (see - /// [`Self::set_sensor_shadow_range`]). + /// `set_sensor_shadow_range`, the 3D sensor-camera counterpart). pub fn set_shadow_range(&mut self, first_cascade_far_bound: f32, shadow_distance: f32) { self.window .set_first_cascade_far_bound(first_cascade_far_bound.max(0.01)); From d2d1cda85d470c13746cd680dca541e1d9757d87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 1 Oct 2026 13:14:08 +0200 Subject: [PATCH 22/25] fix(rbd): stop bodies resting on trimeshes from being ejected --- src_rbd/dynamics/coloring.rs | 20 +++-- src_rbd_shaders/broad_phase/narrow_phase.rs | 19 +++-- src_rbd_shaders/dynamics/coloring.rs | 18 ++++ src_rbd_shaders/dynamics/constraint.rs | 11 ++- src_rbd_shaders/dynamics/solver_utils.rs | 3 + src_rbd_shaders/dynamics/warmstart.rs | 25 +++++- src_rbd_shaders/queries/contact.rs | 10 ++- src_rbd_shaders/queries/contact_pfm_pfm.rs | 94 ++++++++++++++++++++- 8 files changed, 177 insertions(+), 23 deletions(-) diff --git a/src_rbd/dynamics/coloring.rs b/src_rbd/dynamics/coloring.rs index fc377717..fa0b2aa4 100644 --- a/src_rbd/dynamics/coloring.rs +++ b/src_rbd/dynamics/coloring.rs @@ -11,9 +11,9 @@ use crate::pipeline::RunStats; use crate::shaders::broad_phase::ContactPlan; use crate::shaders::dynamics::TwoBodyConstraint; use crate::shaders::dynamics::{ - GpuColorBucketsCount, GpuColorBucketsReset, GpuColorBucketsScatter, GpuFixConflictsTopoGc, - GpuResetCompletionFlagTopoGc, GpuResetLuby, GpuResetTopoGc, GpuStepGraphColoringLuby, - GpuStepGraphColoringTopoGc, + GpuClearCompletionFlagTopoGc, GpuColorBucketsCount, GpuColorBucketsReset, + GpuColorBucketsScatter, GpuFixConflictsTopoGc, GpuResetCompletionFlagTopoGc, GpuResetLuby, + GpuResetTopoGc, GpuStepGraphColoringLuby, GpuStepGraphColoringTopoGc, }; use crate::utils::{GpuPrefixSum, PrefixSumWorkspace}; use khal::Shader; @@ -36,6 +36,7 @@ pub struct GpuColoring { /// Detects and fixes conflicts in TOPO-GC coloring. fix_conflicts_topo_gc_kernel: GpuFixConflictsTopoGc, reset_completion_flag_topo_gc: GpuResetCompletionFlagTopoGc, + clear_completion_flag_topo_gc: GpuClearCompletionFlagTopoGc, // Workspace for bucket-sorting constraint ids by color so each color iteration // only touches their own constraint. color_buckets_reset: GpuColorBucketsReset, @@ -314,9 +315,16 @@ impl GpuColoring { mut args: ColoringArgs<'a>, max_colors: u32, ) -> Result<(), GpuBackendError> { - for _ in 0..max_colors { - self.reset_completion_flag_topo_gc - .call(pass, 1u32, args.uncolored)?; + // The first fix-conflicts pass must validate the seeded colors even when the first step + // colors nothing, so it starts from a cleared ("not converged") flag. At least two rounds + // run so the last pass can still record the color count. + self.clear_completion_flag_topo_gc + .call(pass, 1u32, args.uncolored)?; + for i in 0..max_colors.max(2) { + if i > 0 { + self.reset_completion_flag_topo_gc + .call(pass, 1u32, args.uncolored)?; + } self.dispatch_step_topo_gc(pass, &mut args)?; self.dispatch_fix_conflicts_topo_gc(pass, &mut args)?; } diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 2cbb43c9..a64e45a2 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -417,7 +417,8 @@ pub fn gpu_narrow_phase_shape_shape( bodies: UVec2::new(body1, body2), friction: mat1.combined_friction(&mat2), restitution: mat1.combined_restitution(&mat2), - _padding: [0.0; 2], + subshape: 0, + _padding: 0.0, }, ); } else { @@ -525,7 +526,8 @@ pub fn gpu_narrow_phase_shape_shape_deferred( thickness2: sub2.thickness, colliders: pair.colliders, pair_index: t, - _padding: [0; 3], + subshape: 0, + _padding: [0; 2], }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // NOTE: if we exceed capacity, just skip the pair. @@ -675,7 +677,8 @@ fn trimesh_convex( thickness2: sub2.thickness, colliders, pair_index, - _padding: [0; 3], + subshape: idx.shape_index + 1, + _padding: [0; 2], }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // Skip (don’t write) on overflow; the caller resizes and re-runs. @@ -752,7 +755,8 @@ fn polyline_convex( thickness2: sub2.thickness, colliders, pair_index, - _padding: [0; 3], + subshape: idx.shape_index + 1, + _padding: [0; 2], }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // Skip (don’t write) on overflow; the caller resizes and re-runs. @@ -786,7 +790,9 @@ pub struct NarrowPhasePfmPair { /// Index of the originating pair in the flat collision-pair list; the /// per-pair sort key of the contact-reduction path. pair_index: u32, - _padding: [u32; 3], + /// See [`IndexedManifold::subshape`]. + subshape: u32, + _padding: [u32; 2], } /// PFM (GJK/EPA) manifold computation for the deferred work-list entries. @@ -862,7 +868,8 @@ pub fn gpu_narrow_phase_pfm_pfm( bodies: UVec2::new(body1, body2), friction: mat1.combined_friction(&mat2), restitution: mat1.combined_restitution(&mat2), - _padding: [0.0; 2], + subshape: pair.subshape, + _padding: 0.0, }, ); } else { diff --git a/src_rbd_shaders/dynamics/coloring.rs b/src_rbd_shaders/dynamics/coloring.rs index be112356..a8001bf9 100644 --- a/src_rbd_shaders/dynamics/coloring.rs +++ b/src_rbd_shaders/dynamics/coloring.rs @@ -221,6 +221,24 @@ pub fn gpu_reset_completion_flag_topo_gc( } } +/// Clears the convergence flag so the first fix-conflicts pass validates every color. +/// +/// Colors seeded from the previous frame are marked colored, so the first step may color +/// nothing and leave the flag set, which would skip the validation of the seeds. +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_clear_completion_flag_topo_gc( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] num_colors: &mut u32, +) { + if invocation_id.x == 0 { + // NOTE: same trivial-kernel workaround as `gpu_reset_completion_flag_topo_gc`. + for k in 0..1 { + *num_colors = k; + } + } +} + /// Performs one iteration of Topo-GC coloring. /// /// Generates up to 63 colors (color 0 = uncolored). diff --git a/src_rbd_shaders/dynamics/constraint.rs b/src_rbd_shaders/dynamics/constraint.rs index 559dfb80..17097502 100644 --- a/src_rbd_shaders/dynamics/constraint.rs +++ b/src_rbd_shaders/dynamics/constraint.rs @@ -95,21 +95,26 @@ pub struct TwoBodyConstraint { /// Contact normal direction from body A's perspective (points away from A). /// Normal impulses are applied along this direction to prevent penetration. pub dir_a: Vector, // Non-penetration force direction for the first body. + /// Collider A of the source manifold (3D; stored in the `dir_a` padding lane). #[cfg(feature = "dim3")] - pub _padding0: f32, + pub warmstart_collider_a: u32, #[cfg(feature = "dim3")] /// First tangent direction (3D only, orthogonal to normal). /// Used for friction in the contact plane. pub tangent_a: Vector, // One of the friction force directions. + /// Collider B of the source manifold (3D; stored in the `tangent_a` padding lane). #[cfg(feature = "dim3")] - pub _padding1: f32, + pub warmstart_collider_b: u32, /// Inverse mass of body A along each axis. /// Used to compute linear velocity changes from impulses. pub im_a: Vector, + /// [`IndexedManifold::subshape`] of the source manifold (3D; `im_a` padding lane). + /// + /// [`IndexedManifold::subshape`]: crate::queries::IndexedManifold::subshape #[cfg(feature = "dim3")] - pub _padding2: f32, + pub warmstart_subshape: u32, /// Inverse mass of body B along each axis. pub im_b: Vector, diff --git a/src_rbd_shaders/dynamics/solver_utils.rs b/src_rbd_shaders/dynamics/solver_utils.rs index ea8f76de..d56b0d6f 100644 --- a/src_rbd_shaders/dynamics/solver_utils.rs +++ b/src_rbd_shaders/dynamics/solver_utils.rs @@ -148,6 +148,9 @@ impl IndexedManifold { #[cfg(feature = "dim3")] { constraint.tangent_a = tangents1.read(0); + constraint.warmstart_collider_a = self.colliders.x; + constraint.warmstart_collider_b = self.colliders.y; + constraint.warmstart_subshape = self.subshape; } for k in 0..(contact.len as usize) { diff --git a/src_rbd_shaders/dynamics/warmstart.rs b/src_rbd_shaders/dynamics/warmstart.rs index 9925dc7a..cc7b38d3 100644 --- a/src_rbd_shaders/dynamics/warmstart.rs +++ b/src_rbd_shaders/dynamics/warmstart.rs @@ -120,8 +120,20 @@ pub fn gpu_seed_colors_from_warmstart( for j in first_ref..last_ref { let cid_old = old_body_constraint_ids[j] as usize; + // Same manifold identity as `transfer_warmstart_impulses`: the manifolds of one body + // pair had distinct colors, and seeding them all with the first one's color makes + // them conflict. + #[cfg(feature = "dim3")] + let same_manifold = old_constraints[cid_old].warmstart_collider_a + == new_constraints[i].warmstart_collider_a + && old_constraints[cid_old].warmstart_collider_b + == new_constraints[i].warmstart_collider_b + && old_constraints[cid_old].warmstart_subshape == new_constraints[i].warmstart_subshape; + #[cfg(feature = "dim2")] + let same_manifold = true; if old_constraints[cid_old].solver_body_a == body_a && old_constraints[cid_old].solver_body_b == body_b + && same_manifold { let old_color = old_constraints_colors[cid_old]; // Colors 1..64 are the valid topo-gc range; anything else @@ -191,9 +203,20 @@ pub fn transfer_warmstart_impulses( for j in first_constraint_id_ref..last_constraint_id_ref { let cid_old = old_body_constraint_ids[j] as usize; - // Check if this old constraint involves the same body pair + // Check if this old constraint comes from the same manifold: same body pair and, in 3D, + // the same collider pair and sub-shape. A body pair alone is ambiguous (one manifold per + // trimesh triangle, or per collider of a compound body), and matching the first one + // hands the same impulses to every manifold of the pair. + #[cfg(feature = "dim3")] + let same_manifold = old_constraints[cid_old].warmstart_collider_a + == new_constraints[i].warmstart_collider_a + && old_constraints[cid_old].warmstart_collider_b == new_constraints[i].warmstart_collider_b + && old_constraints[cid_old].warmstart_subshape == new_constraints[i].warmstart_subshape; + #[cfg(feature = "dim2")] + let same_manifold = true; if old_constraints[cid_old].solver_body_a == body_a && old_constraints[cid_old].solver_body_b == body_b + && same_manifold { // Body pair match found! Now match individual contact points. // We don't have feature IDs, so matching is done by proximity in local space. diff --git a/src_rbd_shaders/queries/contact.rs b/src_rbd_shaders/queries/contact.rs index 26b8678b..4301a68b 100644 --- a/src_rbd_shaders/queries/contact.rs +++ b/src_rbd_shaders/queries/contact.rs @@ -122,10 +122,12 @@ pub struct IndexedManifold { /// Combined restitution coefficient of the two colliders (see /// [`ColliderMaterial::combined_restitution`]). pub restitution: f32, - /// Padding so the struct size stays a multiple of 16 bytes — std430 storage - /// buffers require the array stride to satisfy the 16-byte alignment of the - /// inner vector members. - pub _padding: [f32; 2], + /// `1 +` the index of the triangle/segment of a trimesh/polyline collider that produced + /// this manifold, or 0 for a whole convex shape. With `colliders`, it identifies the + /// manifold across frames for warmstarting. + pub subshape: u32, + /// Padding so the struct size stays a multiple of 16 bytes (std430 array stride). + pub _padding: f32, } /// Computes the contact between two balls. diff --git a/src_rbd_shaders/queries/contact_pfm_pfm.rs b/src_rbd_shaders/queries/contact_pfm_pfm.rs index f6a405ce..a5a4e22f 100644 --- a/src_rbd_shaders/queries/contact_pfm_pfm.rs +++ b/src_rbd_shaders/queries/contact_pfm_pfm.rs @@ -10,6 +10,8 @@ use crate::queries::gjk::{ }; use crate::queries::polygonal_feature; use crate::shapes::Shape; +#[cfg(feature = "dim3")] +use crate::shapes::{SHAPE_TYPE_TRIANGLE, Triangle}; use crate::{DIM, PaddedVector, Pose, Vector}; use khal_std::index::MaybeIndexUnchecked; @@ -53,6 +55,61 @@ pub fn contact_support_map_support_map( gjk::gjk_result_no_intersection(Vector::X) } +/// Unit normal of `tri`'s plane oriented towards `point`, or zero for degenerate triangles or +/// when `point` (nearly) lies on the plane. Independent of the triangle's winding. +#[cfg(feature = "dim3")] +#[inline] +fn triangle_face_normal_towards(tri: &Triangle, point: Vector) -> Vector { + let n = (tri.b - tri.a).cross(tri.c - tri.a); + let len = n.length(); + if len <= FLT_EPS { + return Vector::ZERO; + } + let n = n / len; + let side = (point - tri.a).dot(n); + if side > FLT_EPS { + n + } else if side < -FLT_EPS { + -n + } else { + Vector::ZERO + } +} + +/// Merges near-coincident points of `manifold` (within a quarter of `prediction`, keeping the +/// deeper one). The clipped features and the appended GJK point often coincide, and duplicated +/// points skew the solver and the per-point warmstart. +#[inline] +fn dedup_manifold_points(manifold: &mut ContactManifold, prediction: f32) { + let eps = (0.25 * prediction).max(1.0e-5); + let eps_sq = eps * eps; + let len = manifold.len; + let mut out = 0u32; + for i in 0..MAX_MANIFOLD_POINTS as u32 { + if i < len { + let p = *manifold.points_a.at(i as usize); + let mut merged = false; + for k in 0..MAX_MANIFOLD_POINTS as u32 { + if k < out && !merged { + let q = *manifold.points_a.at(k as usize); + let d = q.pt - p.pt; + if d.dot(d) < eps_sq { + if p.dist < q.dist { + manifold.points_a.write(k as usize, p); + } + merged = true; + } + } + } + if !merged { + manifold.points_a.write(out as usize, p); + out += 1; + } + } + } + manifold.len = out; +} + /// Computes the contact manifold between two polygonal feature-based shapes. pub fn contact_manifold_pfm_pfm( pose12: Pose, @@ -69,9 +126,36 @@ pub fn contact_manifold_pfm_pfm( match contact.status { CLOSEST_POINTS => { - let p1 = contact.a; - let p2_1 = contact.b; - let local_n1 = contact.dir; + #[allow(unused_mut)] + let mut p1 = contact.a; + #[allow(unused_mut)] + let mut p2_1 = contact.b; + #[allow(unused_mut)] + let mut local_n1 = contact.dir; + #[allow(unused_mut)] + let mut fallback = false; + + // A triangle has no thickness, so EPA can pick its back face when the other shape + // barely touches it. Re-derive the contact along the face normal on the side where + // the other shape sits. + #[cfg(feature = "dim3")] + if pfm1.shape_type() == SHAPE_TYPE_TRIANGLE { + let tri = pfm1.to_triangle(); + let face_n = triangle_face_normal_towards(&tri, pose12.translation); + // `face_n` is zero when undecidable, which fails this test. + if local_n1.dot(face_n) < 0.0 { + let support_dir2 = pose12.inverse_transform_vector(-face_n); + p2_1 = pose12.transform_point(pfm2.local_support_point(support_dir2, vertices)); + let dist = (p2_1 - tri.a).dot(face_n); + if dist > total_prediction { + return ContactManifold::default(); + } + p1 = p2_1 - face_n * dist; + local_n1 = face_n; + fallback = true; + } + } + let local_n2 = pose12.inverse_transform_vector(-local_n1); #[cfg(feature = "dim2")] @@ -93,8 +177,11 @@ pub fn contact_manifold_pfm_pfm( false, ); + // The fallback's projected deepest point may lie outside the triangle, so it is + // only used when feature clipping produced nothing. if manifold.len < MAX_MANIFOLD_POINTS as u32 && (DIM == 3 || (DIM == 2 && manifold.len == 0)) + && !(fallback && manifold.len > 0) { let dist = (p2_1 - p1).dot(local_n1); manifold @@ -111,6 +198,7 @@ pub fn contact_manifold_pfm_pfm( } } + dedup_manifold_points(&mut manifold, prediction); manifold.normal_a = local_n1; manifold } From 96db3d07f1e37f8ffb71c557e94a51282ebc401a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 1 Oct 2026 13:25:07 +0200 Subject: [PATCH 23/25] fix(rbd): enlarge the polyline broad-phase AABB by the segment capsule radius --- src_rbd_shaders/broad_phase/narrow_phase.rs | 2 +- src_rbd_shaders/shapes/mod.rs | 2 +- src_rbd_shaders/shapes/polyline.rs | 6 ++++++ src_rbd_shaders/shapes/shape.rs | 5 ++++- 4 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index a64e45a2..767bdca4 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -720,7 +720,7 @@ fn polyline_convex( } // Get the convex shape's AABB in the polyline's local space, and enlarge with the prediction distance. - let thickness = 0.4; // TODO: make thickness configurable or part of the polyline struct + let thickness = crate::shapes::POLYLINE_THICKNESS; let mut test_aabb = convex.compute_aabb(pose12, vertices); test_aabb.mins -= Vector::splat(prediction + thickness); test_aabb.maxs += Vector::splat(prediction + thickness); diff --git a/src_rbd_shaders/shapes/mod.rs b/src_rbd_shaders/shapes/mod.rs index c54b6f01..c108dd7e 100644 --- a/src_rbd_shaders/shapes/mod.rs +++ b/src_rbd_shaders/shapes/mod.rs @@ -36,7 +36,7 @@ pub use cone::Cone; pub use convex_polyhedron::ConvexPolyhedron; #[cfg(feature = "dim3")] pub use cylinder::Cylinder; -pub use polyline::Polyline; +pub use polyline::{POLYLINE_THICKNESS, Polyline}; pub use segment::Segment; pub use shape::*; #[cfg(feature = "dim3")] diff --git a/src_rbd_shaders/shapes/polyline.rs b/src_rbd_shaders/shapes/polyline.rs index fdcb3950..3d101f6f 100644 --- a/src_rbd_shaders/shapes/polyline.rs +++ b/src_rbd_shaders/shapes/polyline.rs @@ -8,6 +8,12 @@ use crate::bounding_volumes::Aabb; use crate::shapes::segment::Segment; use khal_std::index::MaybeIndexUnchecked; +/// Radius of the capsule each polyline segment is treated as in contact +/// detection. The broad-phase AABB is enlarged by it too, so a pair is found +/// before a body reaches the capsule surface rather than once it is deep inside. +// TODO: make the thickness configurable or part of the polyline struct. +pub const POLYLINE_THICKNESS: f32 = 0.4; + /// A polyline (connected line segments) with BVH acceleration structure. #[derive(Clone, Copy, Default)] #[repr(C)] diff --git a/src_rbd_shaders/shapes/shape.rs b/src_rbd_shaders/shapes/shape.rs index 8fe0a728..78e15927 100644 --- a/src_rbd_shaders/shapes/shape.rs +++ b/src_rbd_shaders/shapes/shape.rs @@ -622,7 +622,10 @@ impl Shape { if ty == SHAPE_TYPE_POLYLINE { let pline = self.to_polyline(); - let local_aabb = pline.aabb(); + let mut local_aabb = pline.aabb(); + // the narrow phase sees each segment as a capsule of this radius + local_aabb.mins -= Vector::splat(crate::shapes::POLYLINE_THICKNESS); + local_aabb.maxs += Vector::splat(crate::shapes::POLYLINE_THICKNESS); return local_aabb.transform_by(pose); } From 2dfab44b3be106e8c2c7fd379cac069b3dd085b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 1 Oct 2026 13:25:07 +0200 Subject: [PATCH 24/25] fix(rbd): match 2D warmstart manifolds by collider pair and segment --- src_rbd/pipeline/mod.rs | 2 + src_rbd/pipeline/test_polyline_warmstart.rs | 174 ++++++++++++++++++++ src_rbd_shaders/dynamics/constraint.rs | 11 +- src_rbd_shaders/dynamics/solver_utils.rs | 6 +- src_rbd_shaders/dynamics/warmstart.rs | 18 +- 5 files changed, 196 insertions(+), 15 deletions(-) create mode 100644 src_rbd/pipeline/test_polyline_warmstart.rs diff --git a/src_rbd/pipeline/mod.rs b/src_rbd/pipeline/mod.rs index ce56e6e9..b62731bd 100644 --- a/src_rbd/pipeline/mod.rs +++ b/src_rbd/pipeline/mod.rs @@ -13,6 +13,8 @@ mod rbd_state_from_rapier; mod rbd_step; #[cfg(all(test, feature = "dim3"))] mod test_batched_stacks; +#[cfg(all(test, feature = "dim2"))] +mod test_polyline_warmstart; #[cfg(feature = "dim3")] pub use rbd_state::RbdSnapshot; diff --git a/src_rbd/pipeline/test_polyline_warmstart.rs b/src_rbd/pipeline/test_polyline_warmstart.rs new file mode 100644 index 00000000..00c7b29b --- /dev/null +++ b/src_rbd/pipeline/test_polyline_warmstart.rs @@ -0,0 +1,174 @@ +//! 2D probe for contact warmstarting on polylines and multi-collider bodies. +//! +//! A polyline ground emits one manifold per segment, and a body with two +//! colliders one per collider, so a body pair alone does not identify a +//! manifold across frames. Bodies resting on a zigzag polyline must stay at +//! rest: no ejection, no fall-through, no resting jitter. Two failures are +//! guarded: a broad phase that ignored the segment capsules found pairs only +//! once bodies were deep inside them and launched every body, and matching +//! manifolds by body pair alone left a resting jitter of a few tenths of a +//! millimeter. Run with +//! `cargo test -p nexus_rbd2d --features metal test_polyline_warmstart -- --nocapture --ignored`. + +use crate::math::Pose; +use crate::pipeline::{RbdCapacities, RbdPipeline, RbdState}; +use crate::rapier::prelude::*; +use crate::shaders::dynamics::RbdSimParams; +use khal::backend::{Backend, GpuBackend}; + +async fn test_backend() -> GpuBackend { + #[cfg(feature = "metal")] + { + GpuBackend::Metal(khal::backend::metal::Metal::new().unwrap()) + } + #[cfg(not(feature = "metal"))] + { + GpuBackend::WebGpu(khal::backend::WebGpu::default().await.unwrap()) + } +} + +/// Number of dynamic bodies per environment. +const NUM_BODIES: usize = 9; + +/// The polyline narrow phase sees each segment as a capsule of this radius, so +/// the ground surface sits this far above the polyline's vertices. +const POLYLINE_THICKNESS: f32 = 0.4; + +/// A zigzag polyline ground (2 cm teeth every 10 cm; with the segment capsules +/// every resting body touches about ten segments, one manifold each) with +/// boxes, balls and two-collider bodies on it. +fn build_env() -> (RigidBodySet, ColliderSet) { + let mut bodies = RigidBodySet::new(); + let mut colliders = ColliderSet::new(); + let ground = bodies.insert(RigidBodyBuilder::fixed()); + let vertices: Vec = (0..=60) + .map(|i| Vector::new(i as f32 * 0.1 - 3.0, 0.02 * (i % 2) as f32)) + .collect(); + let indices: Vec<[u32; 2]> = (0..60).map(|i| [i, i + 1]).collect(); + colliders.insert_with_parent( + ColliderBuilder::polyline(vertices, Some(indices)), + ground, + &mut bodies, + ); + for k in 0..NUM_BODIES { + let x = k as f32 * 0.6 - 2.4; + let y = POLYLINE_THICKNESS + 0.2; + let body = bodies.insert(RigidBodyBuilder::dynamic().translation(Vector::new(x, y))); + match k % 3 { + 0 => { + colliders.insert_with_parent( + ColliderBuilder::cuboid(0.15, 0.08), + body, + &mut bodies, + ); + } + 1 => { + colliders.insert_with_parent(ColliderBuilder::ball(0.1), body, &mut bodies); + } + _ => { + for dx in [-0.09, 0.09] { + colliders.insert_with_parent( + ColliderBuilder::cuboid(0.08, 0.06).translation(Vector::new(dx, 0.0)), + body, + &mut bodies, + ); + } + } + } + } + (bodies, colliders) +} + +/// The worst resting motion over `steps` steps after settling, across every +/// dynamic body of `num_envs` environments. +struct RestReport { + /// Largest height excursion from the settled pose (m). + max_excursion: f32, + /// Largest upward displacement in a single step (m). + max_step_rise: f32, + /// Bodies that left the scene (ejected or fell through). + lost: usize, +} + +async fn rest_report(num_envs: usize, settle: u32, steps: u32) -> RestReport { + let backend = test_backend().await; + let pipeline = RbdPipeline::new(&backend).unwrap(); + let envs: Vec<_> = (0..num_envs).map(|_| build_env()).collect(); + let impulse_joints = ImpulseJointSet::new(); + let multibody_joints = MultibodyJointSet::new(); + let params = RbdSimParams::tgs_soft(); + let refs: Vec<_> = envs + .iter() + .map(|(b, c)| (b, c, &impulse_joints, &multibody_joints, ¶ms)) + .collect(); + let capacities = RbdCapacities { + batches: num_envs as u32, + ..Default::default() + }; + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + + for _ in 0..settle { + pipeline.step(&backend, &mut state, None).unwrap(); + } + let rest: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let mut prev = rest.clone(); + let mut report = RestReport { + max_excursion: 0.0, + max_step_rise: 0.0, + lost: 0, + }; + // GPU pose slot of body `b` in environment `e` is `b * num_envs + e`; body 0 + // is the fixed ground and slots past the bodies are spare capacity + let dynamic = num_envs..(NUM_BODIES + 1) * num_envs; + let mut lost = vec![false; rest.len()]; + for _ in 0..steps { + pipeline.step(&backend, &mut state, None).unwrap(); + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + for k in dynamic.clone() { + let y = poses[k].translation.y; + if !(POLYLINE_THICKNESS - 0.1..=POLYLINE_THICKNESS + 1.0).contains(&y) { + lost[k] = true; + continue; + } + report.max_excursion = report.max_excursion.max((y - rest[k].translation.y).abs()); + report.max_step_rise = report.max_step_rise.max(y - prev[k].translation.y); + } + prev = poses; + } + report.lost = lost.iter().filter(|l| **l).count(); + println!( + "polyline rest, {num_envs} env(s): excursion {:.4} m, step rise {:.4} m, lost {}/{}", + report.max_excursion, + report.max_step_rise, + report.lost, + num_envs * NUM_BODIES + ); + report +} + +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_polyline_warmstart() { + for num_envs in [1, 256] { + let report = rest_report(num_envs, 240, 240).await; + assert_eq!(report.lost, 0, "bodies left a resting polyline scene"); + // body-pair matching measured 0.2 to 0.3 mm here; per-manifold matching 0 + assert!( + report.max_excursion < 1.0e-4, + "resting bodies moved {} m", + report.max_excursion + ); + assert!( + report.max_step_rise < 1.0e-4, + "a resting body popped {} m", + report.max_step_rise + ); + } +} diff --git a/src_rbd_shaders/dynamics/constraint.rs b/src_rbd_shaders/dynamics/constraint.rs index 17097502..5ff6d312 100644 --- a/src_rbd_shaders/dynamics/constraint.rs +++ b/src_rbd_shaders/dynamics/constraint.rs @@ -140,8 +140,17 @@ pub struct TwoBodyConstraint { /// Number of active contact points in this manifold (1-4 in 3D, 1-2 in 2D). pub len: u32, + /// Collider A of the source manifold (2D; appended, no spare padding lane). #[cfg(feature = "dim2")] - pub _padding: u32, + pub warmstart_collider_a: u32, + /// Collider B of the source manifold (2D). + #[cfg(feature = "dim2")] + pub warmstart_collider_b: u32, + /// [`IndexedManifold::subshape`] of the source manifold (2D). + /// + /// [`IndexedManifold::subshape`]: crate::queries::IndexedManifold::subshape + #[cfg(feature = "dim2")] + pub warmstart_subshape: u32, #[cfg(feature = "dim3")] pub _padding: [u32; 3], } diff --git a/src_rbd_shaders/dynamics/solver_utils.rs b/src_rbd_shaders/dynamics/solver_utils.rs index d56b0d6f..eeaff880 100644 --- a/src_rbd_shaders/dynamics/solver_utils.rs +++ b/src_rbd_shaders/dynamics/solver_utils.rs @@ -148,10 +148,10 @@ impl IndexedManifold { #[cfg(feature = "dim3")] { constraint.tangent_a = tangents1.read(0); - constraint.warmstart_collider_a = self.colliders.x; - constraint.warmstart_collider_b = self.colliders.y; - constraint.warmstart_subshape = self.subshape; } + constraint.warmstart_collider_a = self.colliders.x; + constraint.warmstart_collider_b = self.colliders.y; + constraint.warmstart_subshape = self.subshape; for k in 0..(contact.len as usize) { let pt = cpose1 diff --git a/src_rbd_shaders/dynamics/warmstart.rs b/src_rbd_shaders/dynamics/warmstart.rs index cc7b38d3..5141b0ae 100644 --- a/src_rbd_shaders/dynamics/warmstart.rs +++ b/src_rbd_shaders/dynamics/warmstart.rs @@ -123,14 +123,12 @@ pub fn gpu_seed_colors_from_warmstart( // Same manifold identity as `transfer_warmstart_impulses`: the manifolds of one body // pair had distinct colors, and seeding them all with the first one's color makes // them conflict. - #[cfg(feature = "dim3")] let same_manifold = old_constraints[cid_old].warmstart_collider_a == new_constraints[i].warmstart_collider_a && old_constraints[cid_old].warmstart_collider_b == new_constraints[i].warmstart_collider_b - && old_constraints[cid_old].warmstart_subshape == new_constraints[i].warmstart_subshape; - #[cfg(feature = "dim2")] - let same_manifold = true; + && old_constraints[cid_old].warmstart_subshape + == new_constraints[i].warmstart_subshape; if old_constraints[cid_old].solver_body_a == body_a && old_constraints[cid_old].solver_body_b == body_b && same_manifold @@ -203,17 +201,15 @@ pub fn transfer_warmstart_impulses( for j in first_constraint_id_ref..last_constraint_id_ref { let cid_old = old_body_constraint_ids[j] as usize; - // Check if this old constraint comes from the same manifold: same body pair and, in 3D, - // the same collider pair and sub-shape. A body pair alone is ambiguous (one manifold per - // trimesh triangle, or per collider of a compound body), and matching the first one + // Check if this old constraint comes from the same manifold: same body pair, collider pair + // and sub-shape. A body pair alone is ambiguous (one manifold per trimesh triangle or + // polyline segment, or per collider of a compound body), and matching the first one // hands the same impulses to every manifold of the pair. - #[cfg(feature = "dim3")] let same_manifold = old_constraints[cid_old].warmstart_collider_a == new_constraints[i].warmstart_collider_a - && old_constraints[cid_old].warmstart_collider_b == new_constraints[i].warmstart_collider_b + && old_constraints[cid_old].warmstart_collider_b + == new_constraints[i].warmstart_collider_b && old_constraints[cid_old].warmstart_subshape == new_constraints[i].warmstart_subshape; - #[cfg(feature = "dim2")] - let same_manifold = true; if old_constraints[cid_old].solver_body_a == body_a && old_constraints[cid_old].solver_body_b == body_b && same_manifold From d199733bc99f3681bd9e0b6a88b216bc8298abc5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Thu, 1 Oct 2026 20:49:44 +0200 Subject: [PATCH 25/25] ci: install cargo-gpu 0.10.0 to match the rust-gpu 0.10.0 release --- .github/workflows/ci.yaml | 8 ++++---- .github/workflows/python-release.yaml | 2 +- README.md | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 4d01ba72..45392d1d 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -49,7 +49,7 @@ jobs: run: sudo apt-get update; sudo apt-get install --no-install-recommends build-essential curl wget file libssl-dev - name: Install cargo-gpu (rust-gpu shader compiler) - run: cargo install cargo-gpu --version 0.10.0-alpha.1 + run: cargo install cargo-gpu --version 0.10.0 - name: Install the rust-gpu toolchain run: cargo gpu install --auto-install-rust-toolchain @@ -78,7 +78,7 @@ jobs: run: sudo apt-get update; sudo apt-get install --no-install-recommends build-essential curl wget file libssl-dev - name: Install cargo-gpu (rust-gpu shader compiler) - run: cargo install cargo-gpu --version 0.10.0-alpha.1 + run: cargo install cargo-gpu --version 0.10.0 - name: Install the rust-gpu toolchain run: cargo gpu install --auto-install-rust-toolchain @@ -116,7 +116,7 @@ jobs: libegl1-mesa-dev libgl1-mesa-dri libxcb-xfixes0-dev mesa-vulkan-drivers - name: Install cargo-gpu (rust-gpu shader compiler) - run: cargo install cargo-gpu --version 0.10.0-alpha.1 + run: cargo install cargo-gpu --version 0.10.0 - name: Install the rust-gpu toolchain run: cargo gpu install --auto-install-rust-toolchain @@ -159,7 +159,7 @@ jobs: libegl1-mesa-dev libgl1-mesa-dri libxcb-xfixes0-dev mesa-vulkan-drivers - name: Install cargo-gpu (rust-gpu shader compiler) - run: cargo install cargo-gpu --version 0.10.0-alpha.1 + run: cargo install cargo-gpu --version 0.10.0 - name: Install the rust-gpu toolchain run: cargo gpu install --auto-install-rust-toolchain diff --git a/.github/workflows/python-release.yaml b/.github/workflows/python-release.yaml index 62f1c11e..b1b3a475 100644 --- a/.github/workflows/python-release.yaml +++ b/.github/workflows/python-release.yaml @@ -21,7 +21,7 @@ env: # Passed to every `maturin build`. Default features (webgpu + extension-module) # give portable wheels; wgpu picks the native GPU API at runtime. MATURIN_ARGS: --release -m crates/nexus_python3d/Cargo.toml - CARGO_GPU_VERSION: 0.10.0-alpha.1 + CARGO_GPU_VERSION: 0.10.0 jobs: # One abi3 wheel per platform (see the `abi3-py39` pyo3 feature): a single diff --git a/README.md b/README.md index 691a7e8b..e83c6b16 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ Nexus uses [`cargo gpu`](https://github.com/Rust-GPU/cargo-gpu) to compile its R the build. **You must install it before building**, otherwise the shader compilation step will fail: ```sh -cargo install cargo-gpu --version 0.10.0-alpha.1 +cargo install cargo-gpu --version 0.10.0 cargo gpu install # Install the toolchain needed by cargo-gpu ```