feat: bring multibodies on par with Rapier

Cargo.toml

@@ -38,12 +38,12 @@ 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.34", default-features = false }
-rapier3d = { version = "0.34", default-features = false }
-rapier3d-urdf = "0.34"
-rapier3d-mjcf = { version = "0.34", features = ["stl", "wavefront", "msh"] }
-parry2d = { version = "0.29", default-features = false }
-parry3d = { version = "0.29", default-features = false }
+rapier2d = { version = "0.35", default-features = false }
+rapier3d = { version = "0.35", default-features = false }
+rapier3d-urdf = "0.35"
+rapier3d-mjcf = { version = "0.35", features = ["stl", "wavefront", "msh"] }
+parry2d = { version = "0.30", default-features = false }
+parry3d = { version = "0.30", default-features = false }
 
 # Viewer / examples deps
 kiss3d = "0.45.1"
@@ -85,6 +85,13 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = [
 ] }
 
 [patch.crates-io]
+# Compare against the rapier checkout the reference example runs, not the
+# crates.io release (their solver defaults differ).
+#rapier2d = { path = "../rapier/crates/rapier2d" }
+#rapier3d = { path = "../rapier/crates/rapier3d" }
+#rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" }
+#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

src/state.rs

@@ -141,11 +141,12 @@ impl NexusState {
         // TODO: resize the GPU buffers too.
     }
 
-    /// Sets the rigid-body multibody gravity vector, e.g. `[0.0, 0.0, -9.81]`
-    /// for a Z-up scene. No-op until the rigid-body state is built, so call it
-    /// after [`Self::finalize`]. Free (non-multibody) bodies keep the solver's
-    /// fixed gravity.
-    #[cfg(all(feature = "dim3", feature = "rbd"))]
+    /// Sets the rigid-body gravity vector, e.g. `[0.0, 0.0, -9.81]` for a Z-up
+    /// scene. Every solver path reads the same uniform, so this applies to free
+    /// rigid-bodies and multibody links alike (in 2D the third component is
+    /// ignored). No-op until the rigid-body state is built, so call it after
+    /// [`Self::finalize`].
+    #[cfg(feature = "rbd")]
     pub fn set_rbd_gravity(&mut self, backend: &GpuBackend, gravity: [f32; 3]) {
         if let Some(rbd) = self.rbd.as_mut() {
             rbd.set_gravity(backend, gravity);
@@ -233,6 +234,20 @@ impl NexusState {
         &mut self.rbd_envs[env]
     }
 
+    /// Mutable access to environment `env`'s rapier world that does **not** mark
+    /// the rbd state dirty, for use after [`Self::finalize`].
+    ///
+    /// Nothing written here reaches the GPU on its own: the rapier sets are the
+    /// build-time source the GPU buffers were baked from, and marking them dirty
+    /// would rebuild those buffers and snap the simulation back to the authored
+    /// state. Use this to run rapier-side helpers whose output you then push
+    /// through a runtime setter — e.g. driving an MJCF actuator model and
+    /// forwarding the resulting motors with
+    /// `GpuMultibodySet::set_motors` (3D only, hence no intra-doc link here).
+    pub fn rbd_world_mut_untracked(&mut self, env: usize) -> &mut PhysicsWorld {
+        &mut self.rbd_envs[env]
+    }
+
     pub fn insert_rigid_body(&mut self, body: RigidBody, collider: Collider) -> RigidBodyHandle {
         self.insert_rigid_body_in(0, body, collider)
     }

src_rbd/broad_phase/lbvh.rs

@@ -292,11 +292,9 @@ impl Lbvh {
         // One thread per live collider (leaf); padding slots aren't in the tree.
         let colliders_per_batch = active_per_batch;
 
-        self.shaders.reset_collision_pairs.call(
-            pass,
-            [num_batches, 1, 1],
-            collision_pairs_len,
-        )?;
+        self.shaders
+            .reset_collision_pairs
+            .call(pass, [num_batches, 1, 1], collision_pairs_len)?;
         self.shaders.find_collision_pairs.call(
             pass,
             [colliders_per_batch, num_batches, 1],
@@ -318,7 +316,7 @@ impl Lbvh {
 
     /// Brute-force O(n²) replacement for [`Self::update_tree`] +
     /// [`Self::find_pairs`], used when each batch holds at most
-    /// [`Self::BRUTE_FORCE_MAX_COLLIDERS`] colliders. One AABB pass and one
+    /// [`BRUTE_FORCE_MAX_COLLIDERS`] colliders. One AABB pass and one
     /// all-pairs pass emit the same pair set as the whole tree pipeline.
     #[allow(clippy::too_many_arguments)]
     pub fn brute_force_pairs(
@@ -350,11 +348,9 @@ impl Lbvh {
             batch_indices,
             vertex_buffers,
         )?;
-        self.shaders.reset_collision_pairs.call(
-            pass,
-            [num_batches, 1, 1],
-            collision_pairs_len,
-        )?;
+        self.shaders
+            .reset_collision_pairs
+            .call(pass, [num_batches, 1, 1], collision_pairs_len)?;
         self.shaders.bf_find_pairs.call(
             pass,
             [active_per_batch * active_per_batch * num_batches, 1, 1],

src_rbd/dynamics/joint.rs

@@ -5,9 +5,9 @@
 
 use crate::math::Pose;
 use crate::shaders::dynamics::{
-    GpuInitJointConstraints,
-    GpuSolveJointConstraints, GpuUpdateJointConstraints, ImpulseJoint, JointConstraint,
-    JointConstraintBuilder, LocalMassProperties, RbdSimParams, Velocity, WorldMassProperties,
+    GpuInitJointConstraints, GpuSolveJointConstraints, GpuUpdateJointConstraints, ImpulseJoint,
+    JointConstraint, JointConstraintBuilder, LocalMassProperties, RbdSimParams, Velocity,
+    WorldMassProperties,
 };
 use bytemuck::Zeroable;
 use khal::Shader;
@@ -33,7 +33,10 @@ fn convert_joint_limits(limits: RapierJointLimits<f32>) -> JointLimits {
     }
 }
 
-fn convert_joint_motor(motor: RapierJointMotor) -> JointMotor {
+/// Converts a rapier joint motor into the GPU representation. Public so callers
+/// driving actuators at runtime can push a freshly configured motor through
+/// `GpuMultibodySet::set_motor`.
+pub fn convert_joint_motor(motor: RapierJointMotor) -> JointMotor {
     JointMotor {
         target_vel: motor.target_vel,
         target_pos: motor.target_pos,

src_rbd/dynamics/mod.rs

@@ -2,7 +2,7 @@
 
 pub use crate::shaders::dynamics::RbdSimParams;
 pub use coloring::{ColorBucketsArgs, ColoringArgs, GpuColoring};
-pub use joint::{GpuImpulseJointSet, GpuJointSolver, JointSolverArgs};
+pub use joint::{GpuImpulseJointSet, GpuJointSolver, JointSolverArgs, convert_joint_motor};
 pub use mprops_update::{GpuMpropsUpdate, GpuSyncColliderPosesShader};
 #[cfg(feature = "dim3")]
 pub use multibody::{GpuMultibodySet, GpuMultibodySolver, MultibodySolverArgs};

src_rbd/dynamics/multibody/loop_closing_joints.rs

@@ -113,31 +113,31 @@ impl GpuMultibodySet {
                     continue; // Both sides static — no constraint to solve.
                 }
 
-                // Mirror rapier `GenericJoint::transform_to_solver_body_space`:
-                // shift the anchor frame's translation into COM space — but ONLY
-                // for FREE-BODY sides, whose solver pose IS the center of mass.
-                // A multibody-link side is positioned by its `local_to_world`,
-                // which is the link ORIGIN frame (not the COM), so the shift must
-                // NOT be applied there — the anchor stays origin-relative and the
-                // lever arm is taken against the COM separately (see
-                // `world_com` in `update_one_joint`). Applying the shift to MB
-                // links offsets the anchor by `local_com` (≈0.25 m for Cassie's
-                // rods), producing a huge spurious loop-closure violation. This
-                // matches rapier's `generic_joint_constraint_builder` (the shift
-                // is applied to `LinkOrBody::Body` sides only). Fixed-side fold
-                // is still a TODO mirroring rapier's `is_fixed` branch.
+                // Move each anchor frame into the space the solver resolves it
+                // against. A free-body side is positioned by its COM-centered
+                // solver pose, so its anchor shifts by `-local_com`. A fixed
+                // side has no solver pose at all (the kernel resolves it against
+                // the identity), so its world transform is folded in here. A
+                // multibody-link side is positioned by its `local_to_world`, an
+                // origin frame, so the anchor stays as-is — shifting it there
+                // would offset it by `local_com` and fabricate a large loop
+                // violation.
                 let mut joint_data = convert_generic_joint(joint.data);
-                if side_a_kind == SIDE_KIND_BODY
-                    && let Some(rb) = rb1
-                {
-                    let com = rb.mass_properties().local_mprops.local_com;
-                    joint_data.local_frame_a.translation -= com;
+                if let Some(rb) = rb1 {
+                    if side_a_kind == SIDE_KIND_FIXED {
+                        joint_data.local_frame_a = *rb.position() * joint_data.local_frame_a;
+                    } else if side_a_kind == SIDE_KIND_BODY {
+                        joint_data.local_frame_a.translation -=
+                            rb.mass_properties().local_mprops.local_com;
+                    }
                 }
-                if side_b_kind == SIDE_KIND_BODY
-                    && let Some(rb) = rb2
-                {
-                    let com = rb.mass_properties().local_mprops.local_com;
-                    joint_data.local_frame_b.translation -= com;
+                if let Some(rb) = rb2 {
+                    if side_b_kind == SIDE_KIND_FIXED {
+                        joint_data.local_frame_b = *rb.position() * joint_data.local_frame_b;
+                    } else if side_b_kind == SIDE_KIND_BODY {
+                        joint_data.local_frame_b.translation -=
+                            rb.mass_properties().local_mprops.local_com;
+                    }
                 }
 
                 // Per-axis stride = 2 * (ndofs_a + ndofs_b); reserve

src_rbd/dynamics/multibody/multibody_from_rapier.rs

@@ -2,13 +2,11 @@
 
 use super::multibody_set::*;
 use crate::shaders::dynamics::{
-    ConstraintSoftness, MAX_AXIS_CONSTRAINTS, MAX_MB_CONTACT_CONSTRAINTS_PER_MB,
-    MbDofCoupling, MbImpulseJointBuilder, MbImpulseJointConstraint, MultibodyContactConstraint,
-    MultibodyInfo, MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace,
-    RbdSimParams,
+    ConstraintSoftness, MAX_AXIS_CONSTRAINTS, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MbDofCoupling,
+    MbImpulseJointBuilder, MbImpulseJointConstraint, MultibodyContactConstraint, MultibodyInfo,
+    MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams,
 };
 use crate::shaders::utils::linalg::MAX_MB_DOFS;
-use glamx::Vec4;
 use khal::BufferUsages;
 use khal::backend::GpuBackend;
 use vortx::tensor::Tensor;
@@ -31,7 +29,6 @@ impl GpuMultibodySet {
             &HashMap<RigidBodyHandle, u32>,
             &RigidBodySet,
         )],
-        gravity: [f32; 3],
         colliders_per_batch: u32,
     ) -> Self {
         let num_batches = environments.len() as u32;
@@ -245,6 +242,11 @@ impl GpuMultibodySet {
                     let mut ws = make_workspace_init();
                     ws.coords = link.joint.coords();
                     ws.joint_rot = link.joint.joint_rot();
+                    if let Some(rb) = bodies.get(link.rigid_body_handle()) {
+                        ws.gravity_scale = rb.gravity_scale();
+                        ws.external_force = rb.user_force();
+                        ws.external_torque = rb.user_torque();
+                    }
 
                     // For free joints at the root, copy the rigid-body pose directly.
                     if link.joint.data.locked_axes.is_empty()
@@ -467,6 +469,7 @@ impl GpuMultibodySet {
             out
         }
         let nb = num_batches as usize;
+        let info_mirror = all_infos.clone();
         let all_infos = interleave(&all_infos, mb_cap, nb);
         let all_statics = interleave(&all_statics, links_cap, nb);
         let all_dof_vals = interleave(&all_dof_vals, dofs_cap, nb);
@@ -487,21 +490,22 @@ impl GpuMultibodySet {
             mass_matrix_entries_per_batch: mm_cap,
             coriolis_entries_per_batch: cor_cap,
             i_coriolis_dt_entries_per_batch: icdt_cap,
-            // Rapier's scheme (the default): the coriolis-melded matrix only
-            // drives the free-acceleration solve, constraints see the plain
-            // matrix, and everything is built once per step. The implicit
-            // legacy mode (single melded matrix, per-substep rebuilds) stays
-            // available via `set_implicit_coriolis(true)`.
-            implicit_coriolis: false,
+            // Default: implicit coriolis. The acceleration solve uses a mass
+            // matrix augmented with the coriolis/gyroscopic derivatives, while
+            // constraints keep the plain one. `set_implicit_coriolis(false)`
+            // falls back to a single plain matrix with explicit-only
+            // coriolis/gyroscopic forces (cheaper, but less stable).
+            implicit_coriolis: true,
             has_joint_constraints: all_infos.iter().any(|info| info.max_constraints > 0),
 
             multibody_info: Tensor::vector(backend, &all_infos, storage).unwrap(),
             links_static: Tensor::vector(backend, &all_statics, storage | BufferUsages::COPY_DST)
                 .unwrap(),
             links_static_mirror: all_statics.clone(),
+            info_mirror,
             links_workspace: Tensor::vector(
                 backend,
-                &crate::shaders::dynamics::ws_soa_from_structs(&all_ws, links_cap, num_batches),
+                crate::shaders::dynamics::ws_soa_from_structs(&all_ws, links_cap, num_batches),
                 storage,
             )
             .unwrap(),
@@ -581,6 +585,15 @@ impl GpuMultibodySet {
                 storage,
             )
             .unwrap(),
+            old_contact_constraints: Tensor::vector(
+                backend,
+                vec![
+                    MultibodyContactConstraint::default();
+                    (contact_cons_cap * num_batches) as usize
+                ],
+                storage,
+            )
+            .unwrap(),
             contact_constraint_jacs: Tensor::vector(
                 backend,
                 vec![0.0f32; (contact_cons_col_cap * num_batches) as usize],
@@ -601,7 +614,11 @@ impl GpuMultibodySet {
                 // Sized by the capacity stride (the kernels index blocks by
                 // `batch · multibodies_batch_capacity + mb_idx`).
                 let total_mbs = mb_cap * num_batches;
-                if global_max_mb > 0 && total_mbs <= MAX_DELASSUS_MULTIBODIES {
+                // `MAX_DELASSUS_MULTIBODIES` is currently 0, which disables the
+                // path; the bound is kept so raising the constant re-enables it.
+                #[allow(clippy::absurd_extreme_comparisons)]
+                let use_delassus = global_max_mb > 0 && total_mbs <= MAX_DELASSUS_MULTIBODIES;
+                if use_delassus {
                     let block = (MAX_MB_CONTACT_CONSTRAINTS_PER_MB
                         * MAX_MB_CONTACT_CONSTRAINTS_PER_MB)
                         as usize;
@@ -667,14 +684,8 @@ impl GpuMultibodySet {
             contact_constraint_columns_per_batch: contact_cons_col_cap,
 
             num_solver_iterations: 4,
+            num_internal_pgs_iterations: 1,
 
-            // FIXME: should be read from the simulation settings.
-            gravity: Tensor::scalar(
-                backend,
-                Vec4::new(gravity[0], gravity[1], gravity[2], 0.0),
-                BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST,
-            )
-            .unwrap(),
             dt: Tensor::scalar(
                 backend,
                 1.0f32 / 60.0,
@@ -689,6 +700,7 @@ impl GpuMultibodySet {
                 BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST,
             )
             .unwrap(),
+            warmstart_coefficient: RbdSimParams::default().warmstart_coefficient,
         }
     }
 }