Add alternative rigid contact model

environment.yml

@@ -12,6 +12,7 @@ dependencies:
   - jaxlie >= 1.3.0
   - jax-dataclasses >= 1.4.0
   - pptree
+  - qpax
   - rod >= 0.3.0
   - typing_extensions # python<3.12
   # ====================================

pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
     "jaxlie >= 1.3.0",
     "jax_dataclasses >= 1.4.0",
     "pptree",
+    "qpax",
     "rod >= 0.3.0",
     "typing_extensions ; python_version < '3.12'",
 ]
@@ -187,6 +188,7 @@ mediapy = "*"
 mujoco = "*"
 notebook = "*"
 pptree = "*"
+qpax = "*"
 rod = "*"
 sdformat14 = "*"
 typing_extensions = "*"

src/jaxsim/api/contact.py

@@ -114,19 +114,24 @@ def collidable_point_forces(
 
 @jax.jit
 def collidable_point_dynamics(
-    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
-) -> tuple[jtp.Matrix, jtp.Matrix]:
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    link_forces: jtp.MatrixLike | None = None,
+) -> tuple[jtp.Matrix, dict[str, jtp.Array]]:
     r"""
-    Compute the 6D force applied to each collidable point and the corresponding
-    material deformation rate.
+    Compute the 6D force applied to each collidable point.
 
     Args:
         model: The model to consider.
         data: The data of the considered model.
+        link_forces:
+            The 6D external forces to apply to the links expressed in the same
+            representation of data.
 
     Returns:
-        The 6D force applied to each collidable point and the corresponding
-        material deformation rate.
+        The 6D force applied to each collidable point and additional data based on the contact model configured:
+        - Soft: the material deformation rate.
+        - Rigid: nothing.
 
     Note:
         The material deformation rate is always returned in the mixed frame
@@ -138,7 +143,8 @@ def collidable_point_dynamics(
     # all collidable points belonging to the robot.
     W_p_Ci, W_ṗ_Ci = js.contact.collidable_point_kinematics(model=model, data=data)
 
-    # Import privately the soft contacts classes.
+    # Import privately the contacts classes.
+    from jaxsim.rbda.contacts.rigid import RigidContacts, RigidContactsState
     from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 
     # Build the soft contact model.
@@ -161,6 +167,28 @@ def collidable_point_dynamics(
             W_f_Ci, (CW_ṁ,) = jax.vmap(soft_contacts.compute_contact_forces)(
                 W_p_Ci, W_ṗ_Ci, data.state.contact.tangential_deformation
             )
+            aux_data = dict(m_dot=CW_ṁ)
+
+        case RigidContacts():
+            assert isinstance(model.contact_model, RigidContacts)
+            assert isinstance(data.state.contact, RigidContactsState)
+
+            # Build the contact model.
+            rigid_contacts = RigidContacts(
+                parameters=data.contacts_params, terrain=model.terrain
+            )
+
+            # Compute the 6D force expressed in the inertial frame and applied to each
+            # collidable point.
+            W_f_Ci, _ = rigid_contacts.compute_contact_forces(
+                position=W_p_Ci,
+                velocity=W_ṗ_Ci,
+                model=model,
+                data=data,
+                link_forces=link_forces,
+            )
+
+            aux_data = dict()
 
         case _:
             raise ValueError(f"Invalid contact model {model.contact_model}")
@@ -175,7 +203,7 @@ def collidable_point_dynamics(
         )
     )(W_f_Ci)
 
-    return f_Ci, CW_ṁ
+    return f_Ci, aux_data
 
 
 @functools.partial(jax.jit, static_argnames=["link_names"])

src/jaxsim/api/model.py

@@ -14,6 +14,7 @@
 from jax_dataclasses import Static
 
 import jaxsim.api as js
+import jaxsim.exceptions
 import jaxsim.terrain
 import jaxsim.typing as jtp
 from jaxsim.math import Adjoint, Cross
@@ -1890,6 +1891,8 @@ def step(
         and the new state of the integrator.
     """
 
+    from jaxsim.rbda.contacts.rigid import RigidContacts
+
     # Extract the integrator kwargs.
     # The following logic allows using integrators having kwargs colliding with the
     # kwargs of this step function.
@@ -1901,12 +1904,12 @@ def step(
 
     # Extract the initial resources.
     t0_ns = data.time_ns
-    state_x0 = data.state
+    state_t0 = data.state
     integrator_state_x0 = integrator_state
 
     # Step the dynamics forward.
-    state_xf, integrator_state_xf = integrator.step(
-        x0=state_x0,
+    state_tf, integrator_state_tf = integrator.step(
+        x0=state_t0,
         t0=jnp.array(t0_ns / 1e9).astype(float),
         dt=dt,
         params=integrator_state_x0,
@@ -1928,11 +1931,61 @@ def step(
         ),
     )
 
-    return (
+    data_tf = (
         # Store the new state of the model and the new time.
         data.replace(
-            state=state_xf,
+            state=state_tf,
             time_ns=t0_ns + jnp.array(dt * 1e9).astype(jnp.uint64),
-        ),
-        integrator_state_xf,
+        )
+    )
+
+    # Post process the simulation state, if needed.
+    match model.contact_model:
+
+        # Rigid contact models use an impact model that produces a discontinuous model velocity.
+        # Hence here we need to reset the velocity after each impact to guarantee that
+        # the linear velocity of the active collidable points is zero.
+        case RigidContacts():
+            # Raise runtime error for not supported case in which Rigid contacts and Baumgarte stabilization
+            # enabled are used with ForwardEuler integrator.
+            jaxsim.exceptions.raise_runtime_error_if(
+                condition=jnp.logical_and(
+                    isinstance(
+                        integrator,
+                        jaxsim.integrators.fixed_step.ForwardEuler
+                        | jaxsim.integrators.fixed_step.ForwardEulerSO3,
+                    ),
+                    jnp.array(
+                        [data_tf.contacts_params.K, data_tf.contacts_params.D]
+                    ).any(),
+                ),
+                msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
+            )
+
+            with data_tf.switch_velocity_representation(VelRepr.Mixed):
+                W_p_C = js.contact.collidable_point_positions(model, data_tf)
+                M = js.model.free_floating_mass_matrix(model, data_tf)
+                J_WC = js.contact.jacobian(model, data_tf)
+                px, py, _ = W_p_C.T
+                terrain_height = jax.vmap(model.terrain.height)(px, py)
+                inactive_collidable_points, _ = RigidContacts.detect_contacts(
+                    W_p_C=W_p_C,
+                    terrain_height=terrain_height,
+                )
+                BW_nu_post_impact = RigidContacts.compute_impact_velocity(
+                    data=data_tf,
+                    inactive_collidable_points=inactive_collidable_points,
+                    M=M,
+                    J_WC=J_WC,
+                )
+                data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
+                data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
+            # Restore the input velocity representation.
+            data_tf = data_tf.replace(
+                velocity_representation=data.velocity_representation, validate=False
+            )
+
+    return (
+        data_tf,
+        integrator_state_tf,
     )