[Sprint] Simplify contact model usage in Jaxsim

environment.yml

@@ -27,6 +27,7 @@ dependencies:
   # [testing]
   - idyntree >= 12.2.1
   - pytest
+  - pytest-benchmark
   - pytest-icdiff
   - robot_descriptions
   - icub-models

examples/jaxsim_as_physics_engine.ipynb

@@ -175,7 +175,7 @@
     "\n",
     "# Simulate\n",
     "for _t in T:\n",
-    "    data, _ = js.model.step(\n",
+    "    data = js.model.step(\n",
     "        model=model,\n",
     "        data=data,\n",
     "        link_forces=None,\n",
@@ -239,7 +239,7 @@
     "\n",
     "data = data_batch_t0\n",
     "for _t in T:\n",
-    "    data, _ = step_parallel(model, data)"
+    "    data = step_parallel(model, data)"
    ]
   }
  ],
@@ -266,7 +266,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.0"
+   "version": "3.13.1"
   }
  },
  "nbformat": 4,

examples/jaxsim_as_physics_engine_advanced.ipynb

@@ -63,7 +63,6 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import jaxsim.api as js\n",
-    "import jaxsim\n",
     "import rod\n",
     "from jaxsim import logging\n",
     "from rod.builder.primitives import SphereBuilder\n",
@@ -163,7 +162,6 @@
     "model = js.model.JaxSimModel.build_from_model_description(\n",
     "    model_description=model_sdf_string,\n",
     "    time_step=0.001,\n",
-    "    integrator=jaxsim.integrators.fixed_step.Heun2,\n",
     ")\n",
     "\n",
     "# Create the data of a single model.\n",

src/jaxsim/__init__.py

@@ -114,5 +114,5 @@ def _get_default_logging_level(env_var: str) -> logging.LoggingLevel:
 del _is_editable
 
 from . import terrain  # isort:skip
-from . import api, integrators, logging, math, rbda
+from . import api, logging, math, rbda
 from .api.common import VelRepr

src/jaxsim/api/__init__.py

@@ -3,7 +3,9 @@
 from . import (
     com,
     contact,
+    contact_model,
     frame,
+    integrators,
     joint,
     kin_dyn_parameters,
     link,

src/jaxsim/api/contact.py

@@ -95,143 +95,6 @@ def collidable_point_velocities(
     return W_ṗ_Ci
 
 
-@jax.jit
-@js.common.named_scope
-def collidable_point_forces(
-    model: js.model.JaxSimModel,
-    data: js.data.JaxSimModelData,
-    link_forces: jtp.MatrixLike | None = None,
-    joint_force_references: jtp.VectorLike | None = None,
-    **kwargs,
-) -> jtp.Matrix:
-    """
-    Compute the 6D forces 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.
-        joint_force_references:
-            The joint force references to apply to the joints.
-        kwargs: Additional keyword arguments to pass to the active contact model.
-
-    Returns:
-        The 6D forces applied to each collidable point expressed in the frame
-        corresponding to the active representation.
-    """
-
-    f_Ci, _ = collidable_point_dynamics(
-        model=model,
-        data=data,
-        link_forces=link_forces,
-        joint_force_references=joint_force_references,
-        **kwargs,
-    )
-
-    return f_Ci
-
-
-@jax.jit
-@js.common.named_scope
-def collidable_point_dynamics(
-    model: js.model.JaxSimModel,
-    data: js.data.JaxSimModelData,
-    link_forces: jtp.MatrixLike | None = None,
-    joint_force_references: jtp.VectorLike | None = None,
-    **kwargs,
-) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
-    r"""
-    Compute the 6D force applied to each enabled 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.
-        joint_force_references:
-            The joint force references to apply to the joints.
-        kwargs: Additional keyword arguments to pass to the active contact model.
-
-    Returns:
-        The 6D force applied to each enabled collidable point and additional data based
-        on the contact model configured:
-        - Soft: the material deformation rate.
-        - Rigid: no additional data.
-        - QuasiRigid: no additional data.
-
-    Note:
-        The material deformation rate is always returned in the mixed frame
-        `C[W] = ({}^W \mathbf{p}_C, [W])`. This is convenient for integration purpose.
-        Instead, the 6D forces are returned in the active representation.
-    """
-
-    # Build the common kw arguments to pass to the computation of the contact forces.
-    common_kwargs = dict(
-        link_forces=link_forces,
-        joint_force_references=joint_force_references,
-    )
-
-    # Build the additional kwargs to pass to the computation of the contact forces.
-    match model.contact_model:
-
-        case contacts.SoftContacts():
-
-            kwargs_contact_model = {}
-
-        case contacts.RigidContacts():
-
-            kwargs_contact_model = common_kwargs | kwargs
-
-        case contacts.RelaxedRigidContacts():
-
-            kwargs_contact_model = common_kwargs | kwargs
-
-        case contacts.ViscoElasticContacts():
-
-            kwargs_contact_model = common_kwargs | dict(dt=model.time_step) | kwargs
-
-        case _:
-            raise ValueError(f"Invalid contact model: {model.contact_model}")
-
-    # Compute the contact forces with the active contact model.
-    W_f_C, aux_data = model.contact_model.compute_contact_forces(
-        model=model,
-        data=data,
-        **kwargs_contact_model,
-    )
-
-    # Compute the transforms of the implicit frames `C[L] = (W_p_C, [L])`
-    # associated to the enabled collidable point.
-    # In inertial-fixed representation, the computation of these transforms
-    # is not necessary and the conversion below becomes a no-op.
-
-    # Get the indices of the enabled collidable points.
-    indices_of_enabled_collidable_points = (
-        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
-    )
-
-    W_H_C = (
-        js.contact.transforms(model=model, data=data)
-        if data.velocity_representation is not VelRepr.Inertial
-        else jnp.stack([jnp.eye(4)] * len(indices_of_enabled_collidable_points))
-    )
-
-    # Convert the 6D forces to the active representation.
-    f_Ci = jax.vmap(
-        lambda W_f_C, W_H_C: data.inertial_to_other_representation(
-            array=W_f_C,
-            other_representation=data.velocity_representation,
-            transform=W_H_C,
-            is_force=True,
-        )
-    )(W_f_C, W_H_C)
-
-    return f_Ci, aux_data
-
-
 @functools.partial(jax.jit, static_argnames=["link_names"])
 @js.common.named_scope
 def in_contact(
@@ -351,7 +214,7 @@ def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
 
         zero_data = js.data.JaxSimModelData.build(
             model=model,
-            contacts_params=jaxsim.rbda.contacts.SoftContactsParams(),
+            contacts_params=jaxsim.rbda.contacts.RelaxedRigidContactsParams(),
         )
 
         W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
@@ -362,63 +225,17 @@ def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
 
         return 2 * W_pz_CoM
 
-    max_δ = (
+    max_δ = (  # noqa: F841
         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)
     )
 
-    nc = number_of_active_collidable_points_steady_state
+    nc = number_of_active_collidable_points_steady_state  # noqa: F841
 
     match model.contact_model:
 
-        case contacts.SoftContacts():
-            assert isinstance(model.contact_model, contacts.SoftContacts)
-
-            parameters = contacts.SoftContactsParams.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,
-                **kwargs,
-            )
-
-        case contacts.ViscoElasticContacts():
-            assert isinstance(model.contact_model, contacts.ViscoElasticContacts)
-
-            parameters = (
-                contacts.ViscoElasticContactsParams.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,
-                    **kwargs,
-                )
-            )
-
-        case contacts.RigidContacts():
-            assert isinstance(model.contact_model, contacts.RigidContacts)
-
-            # Disable Baumgarte stabilization by default since it does not play
-            # well with the forward Euler integrator.
-            K = kwargs.get("K", 0.0)
-
-            parameters = contacts.RigidContactsParams.build(
-                mu=static_friction_coefficient,
-                **(
-                    dict(
-                        K=K,
-                        D=2 * jnp.sqrt(K),
-                    )
-                    | kwargs
-                ),
-            )
-
         case contacts.RelaxedRigidContacts():
             assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
 

src/jaxsim/api/contact_model.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+
+
+@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_force_references: jtp.VectorLike | None = None,
+    **kwargs,
+) -> 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_force_references:
+            The joint force references to apply to the joints.
+        kwargs: Additional keyword arguments to pass to the active contact model..
+
+    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, _ = model.contact_model.compute_contact_forces(
+        model=model,
+        data=data,
+        link_forces=link_forces,
+        joint_force_references=joint_force_references,
+    )
+
+    # 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, data=data, contact_forces=W_f_C
+    )
+
+    return W_f_L
+
+
+@staticmethod
+def link_forces_from_contact_forces(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    contact_forces: jtp.MatrixLike,
+) -> jtp.Matrix:
+    """
+    Compute the link forces from the contact forces.
+
+    Args:
+        model: The robot model considered by the contact model.
+        data: The data of the considered 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