Reintroduce alternative contact models and streamline changes to the new API

.github/workflows/gpu_benchmark.yml

@@ -0,0 +1,67 @@
+name: GPU Benchmarks
+
+on:
+  pull_request:
+    types: [opened, reopened, synchronize]
+  workflow_dispatch:
+  schedule:
+    - cron: "0 0  * * 1" # Run At 00:00 on Monday
+
+permissions:
+  pull-requests: write
+  deployments: write
+  contents: write
+
+jobs:
+  benchmark:
+    runs-on: self-hosted
+    container:
+      image: ghcr.io/nvidia/jax:jax
+      options: --rm --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up environment
+        run: |
+          pip install .[all]
+
+      - name: Install Gazebo SDF
+        run: |
+            apt update && apt install -y lsb-release wget
+            sh -c 'echo "deb http://packages.osrfoundation.org/gazebo/ubuntu-stable `lsb_release -cs` main" > /etc/apt/sources.list.d/gazebo-stable.list'
+            wget http://packages.osrfoundation.org/gazebo.key -O - | apt-key add -
+            apt-get update
+            apt install -y libsdformat14-dev libsdformat14
+
+      - name: Performance regression check
+        uses: actions/checkout@v4
+
+      - name: Run benchmark and store result
+        run: |
+            pytest tests/test_benchmark.py -k 'not test_rigid_contact_model and not test_soft_contact_model' --gpu-only --batch-size 128 --benchmark-only --benchmark-json output.json
+
+      - name: Download previous benchmark data
+        uses: actions/cache@v4
+        with:
+          path: ./cache
+          key: ${{ runner.os }}-benchmark
+
+      - name: Store benchmark result
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          tool: 'pytest'
+          output-file-path: output.json
+          external-data-json-path: ./cache/benchmark-data.json
+          fail-on-alert: true
+          summary-always: true
+          comment-always: true
+          alert-threshold: 150%
+          auto-push: false
+          gh-pages-branch: gh-pages
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Push benchmark result
+        if: github.event_name != 'pull_request'
+        run: git push 'https://ami-iit:${{ secrets.GITHUB_TOKEN }}@github.com/ami-iit/jaxsim.git' gh-pages:gh-pages

src/jaxsim/api/__init__.py

@@ -4,7 +4,6 @@
     actuation_model,
     com,
     contact,
-    contact_model,
     frame,
     integrators,
     joint,

src/jaxsim/api/com.py

@@ -301,9 +301,7 @@ def other_representation_to_body(
             C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL  # noqa: F841
             C_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
 
-            L_H_C = L_H_W = jax.vmap(  # noqa: F841
-                lambda W_H_L: jaxsim.math.Transform.inverse(W_H_L)
-            )(W_H_L)
+            L_H_C = L_H_W = jax.vmap(jaxsim.math.Transform.inverse)(W_H_L)  # noqa: F841
 
             L_v_LC = L_v_LW = jax.vmap(  # noqa: F841
                 lambda i: -js.link.velocity(

src/jaxsim/api/contact.py

@@ -11,7 +11,7 @@
 import jaxsim.typing as jtp
 from jaxsim import logging
 from jaxsim.math import Adjoint, Cross, Transform
-from jaxsim.rbda import contacts
+from jaxsim.rbda.contacts import SoftContacts
 
 from .common import VelRepr
 
@@ -37,14 +37,11 @@ def collidable_point_kinematics(
         the linear component of the mixed 6D frame velocity.
     """
 
-    # Switch to inertial-fixed since the RBDAs expect velocities in this representation.
-    with data.switch_velocity_representation(VelRepr.Inertial):
-
-        W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
-            model=model,
-            link_transforms=data._link_transforms,
-            link_velocities=data._link_velocities,
-        )
+    W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
+        model=model,
+        link_transforms=data._link_transforms,
+        link_velocities=data._link_velocities,
+    )
 
     return W_p_Ci, W_ṗ_Ci
 
@@ -164,18 +161,23 @@ def estimate_good_soft_contacts_parameters(
 def estimate_good_contact_parameters(
     model: js.model.JaxSimModel,
     *,
+    standard_gravity: jtp.FloatLike = jaxsim.math.STANDARD_GRAVITY,
     static_friction_coefficient: jtp.FloatLike = 0.5,
-    **kwargs,
+    number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
+    damping_ratio: jtp.FloatLike = 1.0,
+    max_penetration: jtp.FloatLike | None = None,
 ) -> jaxsim.rbda.contacts.ContactParamsTypes:
     """
     Estimate good contact parameters.
 
     Args:
         model: The model to consider.
+        standard_gravity: The standard gravity acceleration.
         static_friction_coefficient: The static friction coefficient.
-        kwargs:
-            Additional model-specific parameters passed to the builder method of
-            the parameters class.
+        number_of_active_collidable_points_steady_state:
+            The number of active collidable points in steady state.
+        damping_ratio: The damping ratio.
+        max_penetration: The maximum penetration allowed.
 
     Returns:
         The estimated good contacts parameters.
@@ -190,20 +192,41 @@ def estimate_good_contact_parameters(
         specific application.
     """
 
-    match model.contact_model:
+    def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
+        """
+        Displacement between the CoM and the lowest collidable point using zero
+        joint positions.
+        """
+
+        zero_data = js.data.JaxSimModelData.build(
+            model=model,
+        )
+
+        W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
 
-        case contacts.RelaxedRigidContacts():
-            assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
+        if model.floating_base():
+            W_pz_C = collidable_point_positions(model=model, data=zero_data)[:, -1]
+            return 2 * (W_pz_CoM - W_pz_C.min())
 
-            parameters = contacts.RelaxedRigidContactsParams.build(
-                mu=static_friction_coefficient,
-                **kwargs,
-            )
+        return 2 * W_pz_CoM
 
-        case _:
-            raise ValueError(f"Invalid contact model: {model.contact_model}")
+    max_δ = (
+        max_penetration
+        if max_penetration is not None
+        # Consider as default a 0.5% of the model height.
+        else 0.005 * estimate_model_height(model=model)
+    )
 
-    return parameters
+    nc = number_of_active_collidable_points_steady_state
+
+    return model.contact_model._parameters_class().build_default_from_jaxsim_model(
+        model=model,
+        standard_gravity=standard_gravity,
+        static_friction_coefficient=static_friction_coefficient,
+        max_penetration=max_δ,
+        number_of_active_collidable_points_steady_state=nc,
+        damping_ratio=damping_ratio,
+    )
 
 
 @jax.jit
@@ -244,7 +267,7 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
 
     # Build the link-to-point transform from the displacement between the link frame L
     # and the implicit contact frame C.
-    L_H_C = jax.vmap(lambda L_p_C: jnp.eye(4).at[0:3, 3].set(L_p_C))(L_p_Ci)
+    L_H_C = jax.vmap(jnp.eye(4).at[0:3, 3].set)(L_p_Ci)
 
     # Compose the work-to-link and link-to-point transforms.
     return jax.vmap(lambda W_H_Li, L_H_Ci: W_H_Li @ L_H_Ci)(W_H_L, L_H_C)
@@ -504,3 +527,96 @@ def compute_O_J̇_WC_I(
     )
 
     return O_J̇_WC
+
+
+@jax.jit
+@js.common.named_scope
+def link_contact_forces(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_forces: jtp.MatrixLike | None = None,
+    joint_torques: jtp.VectorLike | None = None,
+) -> tuple[jtp.Matrix, dict[str, jtp.Matrix]]:
+    """
+    Compute the 6D contact forces of all links of the model in inertial representation.
+
+    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 inertial representation
+        joint_torques:
+            The joint torques acting on the joints.
+
+    Returns:
+        A `(nL, 6)` array containing the stacked 6D contact forces of the links,
+        expressed in inertial representation.
+    """
+
+    # Compute the contact forces for each collidable point with the active contact model.
+    W_f_C, aux_dict = model.contact_model.compute_contact_forces(
+        model=model,
+        data=data,
+        **(
+            dict(link_forces=link_forces, joint_force_references=joint_torques)
+            if not isinstance(model.contact_model, SoftContacts)
+            else {}
+        ),
+    )
+
+    # Compute the 6D forces applied to the links equivalent to the forces applied
+    # to the frames associated to the collidable points.
+    W_f_L = link_forces_from_contact_forces(model=model, contact_forces=W_f_C)
+
+    return W_f_L, aux_dict
+
+
+@staticmethod
+def link_forces_from_contact_forces(
+    model: js.model.JaxSimModel,
+    *,
+    contact_forces: jtp.MatrixLike,
+) -> jtp.Matrix:
+    """
+    Compute the link forces from the contact forces.
+
+    Args:
+        model: The robot model considered by the contact model.
+        contact_forces: The contact forces computed by the contact model.
+
+    Returns:
+        The 6D contact forces applied to the links and expressed in the frame of
+        the velocity representation of data.
+    """
+
+    # Get the object storing the contact parameters of the model.
+    contact_parameters = model.kin_dyn_parameters.contact_parameters
+
+    # Extract the indices corresponding to the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    # Convert the contact forces to a JAX array.
+    W_f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
+
+    # Construct the vector defining the parent link index of each collidable point.
+    # We use this vector to sum the 6D forces of all collidable points rigidly
+    # attached to the same link.
+    parent_link_index_of_collidable_points = jnp.array(
+        contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
+    # Create the mask that associate each collidable point to their parent link.
+    # We use this mask to sum the collidable points to the right link.
+    mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
+        model.number_of_links()
+    )
+
+    # Sum the forces of all collidable points rigidly attached to a body.
+    # Since the contact forces W_f_C are expressed in the world frame,
+    # we don't need any coordinate transformation.
+    W_f_L = mask.T @ W_f_C
+
+    return W_f_L