Refactor contact models to use shape transforms

Author: flferretti

src/jaxsim/api/contact.py

@@ -36,18 +36,28 @@ def contact_point_kinematics(
         the linear component of the mixed 6D frame velocity.
     """
 
-    _, _, _, W_p_Ci, W_ṗ_Ci = jax.vmap(
-        jaxsim.rbda.contacts.common.compute_penetration_data, in_axes=(None,)
+    _, _, _, W_p_Ci, W_ṗ_Ci = jax.vmap(
+        lambda shape_transform, shape_type, shape_size, link_transform, link_velocity: jaxsim.rbda.contacts.common.compute_penetration_data(
+            model,
+            shape_transform=shape_transform,
+            shape_type=shape_type,
+            shape_size=shape_size,
+            link_transforms=link_transform,
+            link_velocities=link_velocity,
+        )
     )(
-        model,
-        shape_offset=model.kin_dyn_parameters.contact_parameters.center,
-        shape_type=model.kin_dyn_parameters.contact_parameters.shape_type,
-        shape_size=model.kin_dyn_parameters.contact_parameters.shape_size,
-        link_transforms=data._link_transforms,
-        link_velocities=data._link_velocities,
+        model.kin_dyn_parameters.contact_parameters.transform,
+        model.kin_dyn_parameters.contact_parameters.shape_type,
+        model.kin_dyn_parameters.contact_parameters.shape_size,
+        data._link_transforms[
+            jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+        ],
+        data._link_velocities[
+            jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+        ],
     )
 
-    return W_p_Ci, W_ṗ_Ci
+    return W_p_Ci, W_ṗ_Ci
 
 
 @jax.jit
@@ -241,21 +251,29 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
     # Get the transforms of the parent link of all collidable points.
     W_H_L = data._link_transforms
 
-    def _process_single_shape(shape_type, shape_size, W_H_Li):
+    # Index transforms by the body (parent link) of each collision shape
+    body_indices = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+    W_H_L_indexed = W_H_L[body_indices]
+
+    def _process_single_shape(shape_type, shape_size, shape_transform, W_H_Li):
+        # Apply the collision shape transform to get W_H_S
+        W_H_S = W_H_Li @ shape_transform
+
         _, W_H_C = jax.lax.switch(
             shape_type,
             (detection.box_plane, detection.cylinder_plane, detection.sphere_plane),
             model.terrain,
             shape_size,
-            W_H_Li,
+            W_H_S,
         )
 
         return W_H_C
 
     return jax.vmap(_process_single_shape)(
         model.kin_dyn_parameters.contact_parameters.shape_type,
         model.kin_dyn_parameters.contact_parameters.shape_size,
-        W_H_L,
+        model.kin_dyn_parameters.contact_parameters.transform,
+        W_H_L_indexed,
     )
 
 
@@ -294,13 +312,17 @@ def jacobian(
         model=model, data=data, output_vel_repr=VelRepr.Inertial
     )
 
-    # Compute contact transforms (n_links, n_contacts, 4, 4)
+    # Compute contact transforms (n_shapes, n_contacts_per_shape, 4, 4)
     W_H_C = transforms(model=model, data=data)
 
-    # Flatten link × contact axes for single-batch processing (n_links*n_contacts, 6, 6+n)
-    W_J_WC_flat = jnp.repeat(W_J_WL, 3, axis=0)
+    # Index Jacobians by the body (parent link) of each collision shape
+    body_indices = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+    W_J_WL_indexed = W_J_WL[body_indices]  # (n_shapes, 6, 6+n)
+
+    # Repeat for each contact point per shape: (n_shapes*n_contacts_per_shape, 6, 6+n)
+    W_J_WC_flat = jnp.repeat(W_J_WL_indexed, 3, axis=0)
 
-    # Flatten contact transforms (n_links*n_contacts, 4, 4)
+    # Flatten contact transforms (n_shapes*n_contacts_per_shape, 4, 4)
     W_H_C_flat = W_H_C.reshape(-1, 4, 4)
 
     # Transform Jacobian based on velocity representation
@@ -357,7 +379,11 @@ def jacobian_derivative(
     # Get the link velocities.
     W_v_WL = data._link_velocities
 
-    # Compute the contact transforms (n_links, n_contacts, 4, 4)
+    # Index link velocities by body (parent link) of each collision shape
+    body_indices = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+    W_v_WL_indexed = W_v_WL[body_indices]  # (n_shapes, 6)
+
+    # Compute the contact transforms (n_shapes, n_contacts, 4, 4)
     W_H_C = transforms(model=model, data=data)
 
     # =====================================================
@@ -408,6 +434,10 @@ def compute_Ṫ(Ẋ: jtp.Matrix) -> jtp.Matrix:
             model=model, data=data
         )
 
+    # Index Jacobians by body (parent link) of each collision shape
+    W_J_WL_W_indexed = W_J_WL_W[body_indices]  # (n_shapes, 6, 6+n)
+    W_J̇_WL_W_indexed = W_J̇_WL_W[body_indices]  # (n_shapes, 6, 6+n)
+
     def compute_O_J̇_WC_I(W_H_C, W_v_WL, W_J_WL_W, W_J̇_WL_W) -> jtp.Matrix:
         match output_vel_repr:
             case VelRepr.Inertial:
@@ -430,15 +460,15 @@ def compute_O_J̇_WC_I(W_H_C, W_v_WL, W_J_WL_W, W_J̇_WL_W) -> jtp.Matrix:
 
         return O_J̇_WC_I
 
-    O_J̇_per_link = jax.vmap(
-        lambda H_C_link, v_WL_link, J_WL_link, J̇_WL_link: jax.vmap(
+    O_J̇_per_shape = jax.vmap(
+        lambda H_C_shape, v_WL_shape, J_WL_shape, J̇_WL_shape: jax.vmap(
             compute_O_J̇_WC_I,
             in_axes=(0, None, None, None),  # Map over contacts for W_H_C only
-        )(H_C_link, v_WL_link, J_WL_link, J̇_WL_link),
-        in_axes=(0, 0, 0, 0),  # Map over links
-    )(W_H_C, W_v_WL, W_J_WL_W, W_J̇_WL_W)
+        )(H_C_shape, v_WL_shape, J_WL_shape, J̇_WL_shape),
+        in_axes=(0, 0, 0, 0),  # Map over shapes
+    )(W_H_C, W_v_WL_indexed, W_J_WL_W_indexed, W_J̇_WL_W_indexed)
 
-    O_J̇_WC = O_J̇_per_link.reshape(-1, 6, 6 + model.dofs())
+    O_J̇_WC = O_J̇_per_shape.reshape(-1, 6, 6 + model.dofs())
 
     return O_J̇_WC
 

src/jaxsim/rbda/contacts/relaxed_rigid.py

@@ -328,15 +328,25 @@ def compute_contact_forces(
 
         # Compute the penetration depth and velocity of the collidable points.
         # Note that this function considers the penetration in the normal direction.
-        δ, δ̇, n̂, W_p_C, CW_ṗ_C = jax.vmap(
-            common.compute_penetration_data, in_axes=(None,)
+        δ, δ̇, n̂, W_p_C, CW_ṗ_C = jax.vmap(
+            lambda shape_transform, shape_type, shape_size, link_transform, link_velocity: common.compute_penetration_data(
+                model,
+                shape_transform=shape_transform,
+                shape_type=shape_type,
+                shape_size=shape_size,
+                link_transforms=link_transform,
+                link_velocities=link_velocity,
+            )
         )(
-            model,
-            shape_offset=model.kin_dyn_parameters.contact_parameters.center,
-            shape_type=model.kin_dyn_parameters.contact_parameters.shape_type,
-            shape_size=model.kin_dyn_parameters.contact_parameters.shape_size,
-            link_transforms=data._link_transforms,
-            link_velocities=data._link_velocities,
+            model.kin_dyn_parameters.contact_parameters.transform,
+            model.kin_dyn_parameters.contact_parameters.shape_type,
+            model.kin_dyn_parameters.contact_parameters.shape_size,
+            data._link_transforms[
+                jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+            ],
+            data._link_velocities[
+                jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+            ],
         )
 
         # Compute the position in the constraint frame.
@@ -346,7 +356,7 @@ def compute_contact_forces(
         a_ref, r, *_ = self._regularizers(
             model=model,
             position_constraint=position_constraint,
-            velocity_constraint=CW_ṗ_C,
+            velocity_constraint=CW_ṗ_C,
             parameters=model.contact_params,
         )
 
@@ -529,13 +539,21 @@ def to_inertial(force, H_C):
 
         # Compute the contact forces in inertial representation for
         # each link and contact point.
-        # Nested vmap: inner over contacts, outer over links
-        W_f_C = jax.vmap(lambda f_link, H_link: jax.vmap(to_inertial)(f_link, H_link))(
-            CW_fl_per_link, W_H_C
+        # Nested vmap: inner over contacts, outer over shapes
+        W_f_C = jax.vmap(
+            lambda f_shape, H_shape: jax.vmap(to_inertial)(f_shape, H_shape)
+        )(CW_fl_per_link, W_H_C)
+
+        # Sum over contacts for each shape: (n_shapes, 6)
+        W_f_per_shape = W_f_C.sum(axis=1)
+
+        # Accumulate forces by parent link using segment_sum
+        body_indices = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+        W_f_per_link = jax.ops.segment_sum(
+            W_f_per_shape, body_indices, num_segments=model.number_of_links()
         )
 
-        # Sum over contacts for each link
-        return W_f_C.sum(axis=1), {}
+        return W_f_per_link, {}
 
     @staticmethod
     def _regularizers(
@@ -576,7 +594,11 @@ def _regularizers(
         )
 
         # Compute the 6D inertia matrices of all links.
-        M_L = js.model.link_spatial_inertia_matrices(model=model)[:, :3, :3]
+        M_L_all = js.model.link_spatial_inertia_matrices(model=model)[:, :3, :3]
+
+        # Index M_L by the body (parent link) of each collision shape
+        body_indices = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+        M_L = M_L_all[body_indices]
 
         def imp_aref(
             pos: jtp.Vector,

src/jaxsim/rbda/contacts/soft.py

@@ -419,15 +419,25 @@ def compute_contact_forces(
         # Compute the position and linear velocities (mixed representation) of
         # all the collidable shapes belonging to the robot and extract the ones
         # for the enabled collidable shapes.
-        δ, δ̇, n̂, W_p_C, CW_ṗ_C = jax.vmap(
-            common.compute_penetration_data, in_axes=(None,)
+        δ, δ̇, n̂, W_p_C, CW_ṗ_C = jax.vmap(
+            lambda shape_transform, shape_type, shape_size, link_transform, link_velocity: common.compute_penetration_data(
+                model,
+                shape_transform=shape_transform,
+                shape_type=shape_type,
+                shape_size=shape_size,
+                link_transforms=link_transform,
+                link_velocities=link_velocity,
+            )
         )(
-            model,
-            shape_offset=model.kin_dyn_parameters.contact_parameters.center,
-            shape_type=model.kin_dyn_parameters.contact_parameters.shape_type,
-            shape_size=model.kin_dyn_parameters.contact_parameters.shape_size,
-            link_transforms=data._link_transforms,
-            link_velocities=data._link_velocities,
+            model.kin_dyn_parameters.contact_parameters.transform,
+            model.kin_dyn_parameters.contact_parameters.shape_type,
+            model.kin_dyn_parameters.contact_parameters.shape_size,
+            data._link_transforms[
+                jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+            ],
+            data._link_velocities[
+                jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+            ],
         )
 
         # Extract the material deformation corresponding to the collidable shapes.
@@ -441,9 +451,15 @@ def compute_contact_forces(
         # We exploit two levels of vmap to vectorize over both the shapes and the points.
         # The outer vmap vectorizes over the shapes, while the inner vmap vectorizes
         # over the maximum points (3) belonging to each shape.
-        W_f, ṁ = jax.vmap(
+        W_f_per_shape, ṁ = jax.vmap(
             SoftContacts.compute_contact_force,
             in_axes=(0, 0, 0, 0, 0, 0, None),  # vectorize over shapes
-        )(δ, δ̇, W_p_C, CW_ṗ_C, n̂, m, model.contact_params)
+        )(δ, δ̇, W_p_C, CW_ṗ_C, n̂, m, model.contact_params)
+
+        # Accumulate forces by parent link using segment_sum
+        body_indices = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
+        W_f = jax.ops.segment_sum(
+            W_f_per_shape, body_indices, num_segments=model.number_of_links()
+        )
 
-        return W_f, {"m_dot": ṁ}
+        return W_f, {"m_dot": ṁ}