Harvest VBD proxy wrenches from the capped per-body contact list (#4004)

Author: deepspace28

changelog/3795.fixed.md

@@ -0,0 +1 @@
+Cap the `SolverVBD` proxy-body wrench harvest at the per-body soft-contact list capacity, `rigid_body_particle_contact_buffer_size`. The harvest walked the whole soft-contact stream while the destination solve applied only the records the per-body list kept, so a proxy-coupled body carrying more soft contacts than the list holds fed the source solver a reaction it never applied, injecting momentum that grew with the overflow. Both now read the same list. Below the cap nothing changes.

newton/_src/solvers/vbd/solver_vbd.py

@@ -1190,9 +1190,12 @@ def coupling_harvest_proxy_wrenches(
             )
 
         if contacts.soft_contact_max > 0 and self.body_particle_contact_penalty_k.shape[0] >= contacts.soft_contact_max:
+            # Per-body, mirroring accumulate_body_particle_contacts_per_body: the harvest reads the
+            # same truncated adjacency list the destination solve consumed, so the reaction reported
+            # to the source is by construction the one the solve applied.
             wp.launch(
                 _harvest_vbd_body_particle_contact_forces_on_proxy_bodies_kernel,
-                dim=contacts.soft_contact_max,
+                dim=self.model.body_count * _NUM_CONTACT_THREADS_PER_BODY,
                 inputs=[
                     float(dt),
                     body_local_to_proxy_global,
@@ -1216,6 +1219,9 @@ def coupling_harvest_proxy_wrenches(
                     contacts.soft_contact_normal,
                     self.model.shape_margin,
                     self.model.shape_body,
+                    self.body_particle_contact_buffer_pre_alloc,
+                    self.body_particle_contact_counts,
+                    self.body_particle_contact_indices,
                     out_body_f,
                 ],
                 device=self.device,

newton/_src/solvers/vbd/vbd_coupling_kernels.py

@@ -13,7 +13,7 @@
     evaluate_edge_edge_contact_2_vertices,
     evaluate_vertex_triangle_collision_force_hessian_4_vertices,
 )
-from .rigid_vbd_kernels import _eval_body_particle_contact, _eval_soft_ef_contact
+from .rigid_vbd_kernels import _NUM_CONTACT_THREADS_PER_BODY, _eval_body_particle_contact, _eval_soft_ef_contact
 from .tri_mesh_collision import TriMeshCollisionInfo
 
 wp.set_module_options({"enable_backward": False})
@@ -150,58 +150,84 @@ def _harvest_vbd_body_particle_contact_forces_on_proxy_bodies_kernel(
     body_particle_contact_normal: wp.array[wp.vec3],
     shape_margin: wp.array[float],
     shape_body: wp.array[wp.int32],
+    body_particle_contact_buffer_pre_alloc: int,
+    body_particle_contact_counts: wp.array[wp.int32],
+    body_particle_contact_indices: wp.array[wp.int32],
     out_body_f: wp.array[wp.spatial_vector],
 ):
-    contact_idx = wp.tid()
-    if contact_idx >= body_particle_contact_count[0]:
-        return
-
-    shape_idx = body_particle_contact_shape[contact_idx]
-    if shape_idx < 0 or shape_idx >= shape_body.shape[0]:
-        return
-
-    body_idx = shape_body[shape_idx]
-    if body_idx < 0 or body_idx >= body_local_to_proxy_global.shape[0]:
+    """Sum the body-side soft-contact reaction the destination solve applied to each proxy body.
+
+    Walks the per-body adjacency list, and truncates it at the same per-body capacity as
+    ``accumulate_body_particle_contacts_per_body``. A body carrying more soft contacts than the
+    list holds is pushed by only the records the list kept, so harvesting the whole contact
+    stream would hand the source solver a reaction the destination never applied -- momentum the
+    coupled pair would then disagree about, growing with the size of the overflow.
+    """
+    tid = wp.tid()
+    body_idx = tid // _NUM_CONTACT_THREADS_PER_BODY
+    thread_id_within_body = tid % _NUM_CONTACT_THREADS_PER_BODY
+
+    if body_idx >= body_local_to_proxy_global.shape[0]:
         return
 
     proxy_global = body_local_to_proxy_global[body_idx]
     if proxy_global < 0 or proxy_global >= out_body_f.shape[0]:
         return
 
-    corners = soft_contact_indices[contact_idx]
-    if corners[0] < 0 or corners[0] >= particle_q.shape[0]:
+    num_contacts = body_particle_contact_counts[body_idx]
+    if num_contacts > body_particle_contact_buffer_pre_alloc:
+        num_contacts = body_particle_contact_buffer_pre_alloc
+    if num_contacts == 0:
         return
 
-    bary = soft_contact_barycentric[contact_idx]
-
-    force_on_particle, _hess, cp_world = _eval_soft_ef_contact(
-        contact_idx,
-        corners,
-        bary,
-        particle_q,
-        particle_q_prev,
-        particle_radius,
-        body_particle_contact_penalty_k[contact_idx],
-        body_particle_contact_material_kd[contact_idx],
-        body_particle_contact_material_mu[contact_idx],
-        friction_epsilon,
-        shape_body,
-        body_q,
-        body_q_prev,
-        body_qd,
-        body_com,
-        body_particle_contact_shape,
-        body_particle_contact_body_pos,
-        body_particle_contact_body_vel,
-        body_particle_contact_normal,
-        shape_margin,
-        dt,
-    )
-
-    force_on_body = -force_on_particle
+    max_contacts = body_particle_contact_count[0]  # single total soft-contact count
     com_world = wp.transform_point(body_q[body_idx], body_com[body_idx])
-    torque_on_body = wp.cross(cp_world - com_world, force_on_body)
-    wp.atomic_add(out_body_f, proxy_global, wp.spatial_vector(force_on_body, torque_on_body))
+
+    force_acc = wp.vec3(0.0)
+    torque_acc = wp.vec3(0.0)
+
+    i = thread_id_within_body
+    while i < num_contacts:
+        contact_idx = body_particle_contact_indices[body_idx * body_particle_contact_buffer_pre_alloc + i]
+        i += _NUM_CONTACT_THREADS_PER_BODY
+        if contact_idx >= max_contacts:
+            continue
+
+        corners = soft_contact_indices[contact_idx]
+        if corners[0] < 0 or corners[0] >= particle_q.shape[0]:
+            continue
+
+        bary = soft_contact_barycentric[contact_idx]
+
+        force_on_particle, _hess, cp_world = _eval_soft_ef_contact(
+            contact_idx,
+            corners,
+            bary,
+            particle_q,
+            particle_q_prev,
+            particle_radius,
+            body_particle_contact_penalty_k[contact_idx],
+            body_particle_contact_material_kd[contact_idx],
+            body_particle_contact_material_mu[contact_idx],
+            friction_epsilon,
+            shape_body,
+            body_q,
+            body_q_prev,
+            body_qd,
+            body_com,
+            body_particle_contact_shape,
+            body_particle_contact_body_pos,
+            body_particle_contact_body_vel,
+            body_particle_contact_normal,
+            shape_margin,
+            dt,
+        )
+
+        force_on_body = -force_on_particle
+        force_acc += force_on_body
+        torque_acc += wp.cross(cp_world - com_world, force_on_body)
+
+    wp.atomic_add(out_body_f, proxy_global, wp.spatial_vector(force_acc, torque_acc))
 
 
 @wp.func

newton/tests/test_coupled_solver.py

@@ -2341,9 +2341,87 @@ def add_free_body(*, is_kinematic=False):
     np.testing.assert_allclose(state_out.body_qd.numpy()[source_bodies], model_qd[source_bodies], atol=1.0e-5)
 
 
+def _harvest_wrench_on_one_proxy_body(buffer_size, particle_count, depth=0.01, radius=0.05, dt=1.0 / 60.0):
+    """Harvest the proxy wrench for one body pressed by ``particle_count`` identical soft contacts.
+
+    The particles sit at the same depth on a flat, uniform face, so every contact carries the same
+    force and any subset of them of a given size sums to the same linear reaction. The body is
+    heavy enough that a single step leaves that geometry intact.
+    """
+    builder = newton.ModelBuilder(gravity=(0.0, 0.0, 0.0))
+    body = builder.add_body(mass=1.0e6, inertia=wp.mat33(np.eye(3) * 1.0e6))
+    builder.add_shape_box(
+        body=body,
+        hx=5.0,
+        hy=5.0,
+        hz=0.5,
+        xform=wp.transform(wp.vec3(0.0, 0.0, -0.5), wp.quat_identity()),
+    )
+    for i in range(particle_count):
+        builder.add_particle(
+            pos=(2.0 * radius * (i - 0.5 * (particle_count - 1)), 0.0, radius - depth),
+            vel=(0.0, 0.0, 0.0),
+            mass=1.0,
+            radius=radius,
+        )
+    builder.color()
+    model = builder.finalize(device="cpu")
+
+    solver = SolverVBD(
+        model=model,
+        iterations=1,
+        integrate_with_external_rigid_solver=False,
+        rigid_body_particle_contact_buffer_size=buffer_size,
+        rigid_compliant_alm=True,
+    )
+    state_in = model.state()
+    state_out = model.state()
+    collision_pipeline = newton.CollisionPipeline(model)
+    contacts = collision_pipeline.contacts()
+    collision_pipeline.collide(state_in, contacts)
+    solver.step(state_in, state_out, control=None, contacts=contacts, dt=dt)
+
+    out_body_f = wp.zeros(1, dtype=wp.spatial_vector, device=model.device)
+    solver.coupling_harvest_proxy_wrenches(
+        wp.array([0], dtype=int, device=model.device),
+        out_body_f,
+        body_qd_before=state_in.body_qd,
+        state=state_in,
+        state_out=state_out,
+        contacts=contacts,
+        dt=dt,
+    )
+    return {
+        "soft_contact_count": int(contacts.soft_contact_count.numpy()[0]),
+        "listed": int(solver.body_particle_contact_counts.numpy()[0]),
+        "overflow_max": int(solver.body_particle_contact_overflow_max.numpy()[0]),
+        "wrench": out_body_f.numpy()[0],
+    }
+
+
 class TestSolverVBDCouplingHooks(unittest.TestCase):
     """VBD-specific coupling hook behavior."""
 
+    def test_proxy_wrench_harvest_respects_body_particle_contact_buffer_cap(self):
+        particle_count = 12
+        cap = 4
+
+        full = _harvest_wrench_on_one_proxy_body(particle_count, particle_count)
+        self.assertEqual(full["soft_contact_count"], particle_count)
+        self.assertEqual(full["listed"], particle_count)
+        self.assertEqual(full["overflow_max"], 0, "the uncapped reference must not overflow")
+        force_full = full["wrench"][:3]
+        self.assertGreater(abs(force_full[2]), 0.0, "the contacts must actually push the body")
+
+        capped = _harvest_wrench_on_one_proxy_body(cap, particle_count)
+        self.assertEqual(capped["soft_contact_count"], particle_count)
+        self.assertGreater(capped["overflow_max"], cap, "the small buffer must actually overflow")
+
+        # The solve applied only the `cap` records its per-body list kept, so the reaction fed back
+        # to the source is that same fraction of the whole stream -- not all of it.
+        expected = force_full * (cap / particle_count)
+        np.testing.assert_allclose(capped["wrench"][:3], expected, rtol=1.0e-5, atol=1.0e-6)
+
     def test_external_rigid_solver_harvests_particle_soft_contacts(self):
         builder = newton.ModelBuilder(gravity=(0.0, 0.0, 0.0))
         body = builder.add_body(mass=1.0, inertia=wp.mat33(np.eye(3)))

newton/tests/test_solver_vbd_proxy_full_surface.py

@@ -9,6 +9,7 @@
 import warp as wp
 
 import newton
+from newton._src.solvers.vbd.rigid_vbd_kernels import _NUM_CONTACT_THREADS_PER_BODY
 from newton._src.solvers.vbd.vbd_coupling_kernels import (
     _harvest_vbd_body_particle_contact_forces_on_proxy_bodies_kernel,
 )
@@ -59,9 +60,12 @@ def _set(arr, value):
     _set(contacts.soft_contact_normal, list(_NORMAL))
 
     out_body_f = wp.zeros(1, dtype=wp.spatial_vector, device=device)
+    # The harvest walks a per-body adjacency list, so seed body 0 with the one contact.
+    contact_counts = wp.array([1], dtype=wp.int32, device=device)
+    contact_indices = wp.zeros(smax, dtype=wp.int32, device=device)
     wp.launch(
         _harvest_vbd_body_particle_contact_forces_on_proxy_bodies_kernel,
-        dim=smax,
+        dim=_NUM_CONTACT_THREADS_PER_BODY,
         inputs=[
             0.01,  # dt
             wp.array([0], dtype=int, device=device),  # body 0 -> proxy global 0
@@ -85,6 +89,9 @@ def _set(arr, value):
             contacts.soft_contact_normal,
             model.shape_margin,
             model.shape_body,
+            smax,
+            contact_counts,
+            contact_indices,
         ],
         outputs=[out_body_f],
         device=device,