jaxsim 0.4.3.dev159__py3-none-any.whl → 0.4.3.dev169__py3-none-any.whl

jaxsim/rbda/contacts/visco_elastic.py

@@ -0,0 +1,1053 @@
+from __future__ import annotations
+
+import dataclasses
+import functools
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import jax_dataclasses
+
+import jaxsim
+import jaxsim.api as js
+import jaxsim.exceptions
+import jaxsim.typing as jtp
+from jaxsim import logging
+from jaxsim.math import StandardGravity
+from jaxsim.terrain import FlatTerrain, Terrain
+
+from . import common
+from .soft import SoftContacts, SoftContactsParams
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+
+@jax_dataclasses.pytree_dataclass
+class ViscoElasticContactsParams(common.ContactsParams):
+    """Parameters of the visco-elastic contacts model."""
+
+    K: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(1e6, dtype=float)
+    )
+
+    D: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(2000, dtype=float)
+    )
+
+    static_friction: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.5, dtype=float)
+    )
+
+    p: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.5, dtype=float)
+    )
+
+    q: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.5, dtype=float)
+    )
+
+    @classmethod
+    def build(
+        cls: type[Self],
+        K: jtp.FloatLike = 1e6,
+        D: jtp.FloatLike = 2_000,
+        static_friction: jtp.FloatLike = 0.5,
+        p: jtp.FloatLike = 0.5,
+        q: jtp.FloatLike = 0.5,
+    ) -> Self:
+        """
+        Create a SoftContactsParams instance with specified parameters.
+
+        Args:
+            K: The stiffness parameter.
+            D: The damping parameter of the soft contacts model.
+            static_friction: The static friction coefficient.
+            p:
+                The exponent p corresponding to the damping-related non-linearity
+                of the Hunt/Crossley model.
+            q:
+                The exponent q corresponding to the spring-related non-linearity
+                of the Hunt/Crossley model.
+
+        Returns:
+            A ViscoElasticParams instance with the specified parameters.
+        """
+
+        return ViscoElasticContactsParams(
+            K=jnp.array(K, dtype=float),
+            D=jnp.array(D, dtype=float),
+            static_friction=jnp.array(static_friction, dtype=float),
+            p=jnp.array(p, dtype=float),
+            q=jnp.array(q, dtype=float),
+        )
+
+    @classmethod
+    def build_default_from_jaxsim_model(
+        cls: type[Self],
+        model: js.model.JaxSimModel,
+        *,
+        standard_gravity: jtp.FloatLike = StandardGravity,
+        static_friction_coefficient: jtp.FloatLike = 0.5,
+        max_penetration: jtp.FloatLike = 0.001,
+        number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
+        damping_ratio: jtp.FloatLike = 1.0,
+        p: jtp.FloatLike = 0.5,
+        q: jtp.FloatLike = 0.5,
+    ) -> Self:
+        """
+        Create a ViscoElasticContactsParams instance with good default parameters.
+
+        Args:
+            model: The target model.
+            standard_gravity: The standard gravity constant.
+            static_friction_coefficient:
+                The static friction coefficient between the model and the terrain.
+            max_penetration: The maximum penetration depth.
+            number_of_active_collidable_points_steady_state:
+                The number of contacts supporting the weight of the model
+                in steady state.
+            damping_ratio: The ratio controlling the damping behavior.
+            p:
+                The exponent p corresponding to the damping-related non-linearity
+                of the Hunt/Crossley model.
+            q:
+                The exponent q corresponding to the spring-related non-linearity
+                of the Hunt/Crossley model.
+
+        Returns:
+            A `ViscoElasticContactsParams` instance with the specified parameters.
+
+        Note:
+            The `damping_ratio` parameter allows to operate on the following conditions:
+            - ξ > 1.0: over-damped
+            - ξ = 1.0: critically damped
+            - ξ < 1.0: under-damped
+        """
+
+        # Call the SoftContact builder instead of duplicating the logic.
+        soft_contacts_params = SoftContactsParams.build_default_from_jaxsim_model(
+            model=model,
+            standard_gravity=standard_gravity,
+            static_friction_coefficient=static_friction_coefficient,
+            max_penetration=max_penetration,
+            number_of_active_collidable_points_steady_state=number_of_active_collidable_points_steady_state,
+            damping_ratio=damping_ratio,
+        )
+
+        return ViscoElasticContactsParams.build(
+            K=soft_contacts_params.K,
+            D=soft_contacts_params.D,
+            static_friction=soft_contacts_params.mu,
+            p=p,
+            q=q,
+        )
+
+    def valid(self) -> jtp.BoolLike:
+        """
+        Check if the parameters are valid.
+
+        Returns:
+            `True` if the parameters are valid, `False` otherwise.
+        """
+
+        return (
+            jnp.all(self.K >= 0.0)
+            and jnp.all(self.D >= 0.0)
+            and jnp.all(self.static_friction >= 0.0)
+            and jnp.all(self.p >= 0.0)
+            and jnp.all(self.q >= 0.0)
+        )
+
+    def __hash__(self) -> int:
+
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
+        return hash(
+            (
+                HashedNumpyArray.hash_of_array(self.K),
+                HashedNumpyArray.hash_of_array(self.D),
+                HashedNumpyArray.hash_of_array(self.static_friction),
+                HashedNumpyArray.hash_of_array(self.p),
+                HashedNumpyArray.hash_of_array(self.q),
+            )
+        )
+
+    def __eq__(self, other: ViscoElasticContactsParams) -> bool:
+
+        if not isinstance(other, ViscoElasticContactsParams):
+            return False
+
+        return hash(self) == hash(other)
+
+
+@jax_dataclasses.pytree_dataclass
+class ViscoElasticContacts(common.ContactModel):
+    """Visco-elastic contacts model."""
+
+    parameters: ViscoElasticContactsParams = dataclasses.field(
+        default_factory=ViscoElasticContactsParams
+    )
+
+    terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
+        default_factory=FlatTerrain
+    )
+
+    max_squarings: jax_dataclasses.Static[int] = 25
+
+    @classmethod
+    def build(
+        cls: type[Self],
+        parameters: SoftContactsParams | None = None,
+        terrain: Terrain | None = None,
+        model: js.model.JaxSimModel | None = None,
+        max_squarings: jtp.IntLike | None = None,
+        **kwargs,
+    ) -> Self:
+        """
+        Create a `ViscoElasticContacts` instance with specified parameters.
+
+        Args:
+            parameters: The parameters of the soft contacts model.
+            terrain: The considered terrain.
+            model:
+                The robot model considered by the contact model.
+                If passed, it is used to estimate good default parameters.
+            max_squarings:
+                The maximum number of squarings performed in the matrix exponential.
+
+        Returns:
+            The `ViscoElasticContacts` instance.
+        """
+
+        if len(kwargs) != 0:
+            logging.debug(msg=f"Ignoring extra arguments: {kwargs}")
+
+        # Build the contact parameters if not provided. Use the model to estimate
+        # good default parameters, if passed. Users can later override these default
+        # parameters with their own values -- possibly tuned better.
+        if parameters is None:
+            parameters = (
+                ViscoElasticContactsParams.build_default_from_jaxsim_model(model=model)
+                if model is not None
+                else cls.__dataclass_fields__["parameters"].default_factory()
+            )
+
+        return ViscoElasticContacts(
+            parameters=parameters,
+            terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
+            max_squarings=int(
+                max_squarings or cls.__dataclass_fields__["max_squarings"].default()
+            ),
+        )
+
+    @classmethod
+    def zero_state_variables(cls, model: js.model.JaxSimModel) -> dict[str, jtp.Array]:
+        """
+        Build zero state variables of the contact model.
+        """
+
+        # Initialize the material deformation to zero.
+        tangential_deformation = jnp.zeros(
+            shape=(len(model.kin_dyn_parameters.contact_parameters.body), 3),
+            dtype=float,
+        )
+
+        return {"tangential_deformation": tangential_deformation}
+
+    @jax.jit
+    def compute_contact_forces(
+        self,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        *,
+        dt: jtp.FloatLike,
+        link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
+    ) -> tuple[jtp.Vector, tuple[Any, ...]]:
+        """
+        Compute the contact forces.
+
+        Args:
+            model: The robot model considered by the contact model.
+            data: The data of the considered model.
+            dt: The integration time step.
+            link_forces:
+                The 6D forces to apply to the links expressed in the frame corresponding
+                to the velocity representation of `data`.
+            joint_force_references: The joint force references to apply.
+
+        Note:
+            This contact model, contrarily to most other contact models, requires the
+            knowledge of the integration step. It is not straightforward to assess how
+            this contact model behaves when used with high-order Runge-Kutta schemes.
+            For the time being, it is recommended to use a simple forward Euler scheme.
+            The main benefit of this model is that the stiff contact dynamics is computed
+            separately from the rest of the system dynamics, which allows to use simple
+            integration schemes without altering significantly the simulation stability.
+
+        Returns:
+            A tuple containing as first element the computed 6D contact force applied to
+            the contact point and expressed in the world frame, and as second element
+            a tuple of optional additional information.
+        """
+
+        # Initialize the model and data this contact model is operating on.
+        # This will raise an exception if either the contact model or the
+        # contact parameters are not compatible.
+        model, data = self.initialize_model_and_data(model=model, data=data)
+        assert isinstance(data.contacts_params, ViscoElasticContactsParams)
+
+        # Extract the indices corresponding to the enabled collidable points.
+        indices_of_enabled_collidable_points = (
+            model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+        )
+
+        # Compute the average contact linear forces in mixed representation by
+        # integrating the contact dynamics in the continuous time domain.
+        CW_f̅l, CW_fl̿, m_tf = (
+            ViscoElasticContacts._compute_contact_forces_with_exponential_integration(
+                model=model,
+                data=data,
+                dt=dt,
+                joint_force_references=joint_force_references,
+                link_forces=link_forces,
+                indices_of_enabled_collidable_points=indices_of_enabled_collidable_points,
+                max_squarings=self.max_squarings,
+            )
+        )
+
+        # ============================================
+        # Compute the inertial-fixed 6D contact forces
+        # ============================================
+
+        # Compute the transforms of the mixed frames `C[W] = (W_p_C, [W])`
+        # associated to each collidable point.
+        W_H_C = js.contact.transforms(model=model, data=data)[
+            indices_of_enabled_collidable_points, :, :
+        ]
+
+        # Vmapped transformation from mixed to inertial-fixed representation.
+        compute_forces_inertial_fixed_vmap = jax.vmap(
+            lambda CW_fl_C, W_H_C: data.other_representation_to_inertial(
+                array=jnp.zeros(6).at[0:3].set(CW_fl_C),
+                other_representation=jaxsim.VelRepr.Mixed,
+                transform=W_H_C,
+                is_force=True,
+            )
+        )
+
+        # Express the linear contact forces in the inertial-fixed frame.
+        W_f̅_C, W_f̿_C = jax.vmap(
+            lambda CW_fl: compute_forces_inertial_fixed_vmap(CW_fl, W_H_C)
+        )(jnp.stack([CW_f̅l, CW_fl̿]))
+
+        return W_f̅_C, (W_f̿_C, m_tf)
+
+    @staticmethod
+    @functools.partial(jax.jit, static_argnames=("max_squarings",))
+    def _compute_contact_forces_with_exponential_integration(
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        *,
+        dt: jtp.FloatLike,
+        link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
+        indices_of_enabled_collidable_points: jtp.VectorLike | None = None,
+        max_squarings: int = 25,
+    ) -> tuple[jtp.Matrix, jtp.Matrix, jtp.Matrix]:
+        """
+        Compute the average contact forces by integrating the contact dynamics.
+
+        Args:
+            model: The robot model considered by the contact model.
+            data: The data of the considered model.
+            dt: The integration time step.
+            link_forces: The 6D forces to apply to the links.
+            joint_force_references: The joint force references to apply.
+            indices_of_enabled_collidable_points:
+                The indices of the enabled collidable points.
+            max_squarings:
+                The maximum number of squarings performed in the matrix exponential.
+
+        Returns:
+            A tuple containing:
+            - The average contact forces.
+            - The average of the average contact forces.
+            - The tangential deformation at the final state.
+        """
+
+        # ==========================
+        # Populate missing arguments
+        # ==========================
+
+        indices = (
+            indices_of_enabled_collidable_points
+            if indices_of_enabled_collidable_points is not None
+            else jnp.arange(
+                len(model.kin_dyn_parameters.contact_parameters.body)
+            ).astype(int)
+        )
+
+        # ==================================
+        # Compute the contact point dynamics
+        # ==================================
+
+        p_t0 = js.contact.collidable_point_positions(model, data)[indices, :]
+        v_t0 = js.contact.collidable_point_velocities(model, data)[indices, :]
+        m_t0 = data.state.extended["tangential_deformation"][indices, :]
+
+        # Compute the linearized contact dynamics.
+        # Note that it linearizes the (non-linear) contact model at (p, v, m)[t0].
+        A, b, A_sc, b_sc = ViscoElasticContacts._contact_points_dynamics(
+            model=model,
+            data=data,
+            joint_force_references=joint_force_references,
+            link_forces=link_forces,
+            indices_of_enabled_collidable_points=indices,
+            p_t0=p_t0,
+            v_t0=v_t0,
+            m_t0=m_t0,
+        )
+
+        # =============================================
+        # Compute the integrals of the contact dynamics
+        # =============================================
+
+        # Pack the initial state of the contact points.
+        x_t0 = jnp.hstack([p_t0.flatten(), v_t0.flatten(), m_t0.flatten()])
+
+        # Pack the augmented matrix used to compute the single and double integral
+        # of the exponential integration.
+        A̅ = jnp.vstack(
+            [
+                jnp.hstack(
+                    [
+                        A,
+                        jnp.vstack(b),
+                        jnp.vstack(x_t0),
+                        jnp.vstack(jnp.zeros_like(x_t0)),
+                    ]
+                ),
+                jnp.hstack([jnp.zeros(A.shape[1]), 0, 1, 0]),
+                jnp.hstack([jnp.zeros(A.shape[1]), 0, 0, 1]),
+                jnp.hstack([jnp.zeros(A.shape[1]), 0, 0, 0]),
+            ]
+        )
+
+        # Compute the matrix exponential.
+        exp_tA = jax.scipy.linalg.expm(
+            (dt * A̅).astype(float), max_squarings=max_squarings
+        )
+
+        # Integrate the contact dynamics in the continuous time domain.
+        x_int, x_int2 = (
+            jnp.hstack([jnp.eye(A.shape[0]), jnp.zeros(shape=(A.shape[0], 3))])
+            @ exp_tA
+            @ jnp.vstack([jnp.zeros(shape=(A.shape[0] + 1, 2)), jnp.eye(2)])
+        ).T
+
+        jaxsim.exceptions.raise_runtime_error_if(
+            condition=jnp.isnan(x_int).any(),
+            msg="NaN integration, try to increase `max_squarings` or decreasing `dt`",
+        )
+
+        # ==========================
+        # Compute the contact forces
+        # ==========================
+
+        # Compute the average contact forces.
+        CW_f̅, _ = jnp.split(
+            (A_sc @ x_int / dt + b_sc).reshape(-1, 3),
+            indices_or_sections=2,
+        )
+
+        # Compute the average of the average contact forces.
+        CW_f̿, _ = jnp.split(
+            (A_sc @ x_int2 * 2 / (dt**2) + b_sc).reshape(-1, 3),
+            indices_or_sections=2,
+        )
+
+        # Extract the tangential deformation at the final state.
+        x_tf = x_int / dt
+        m_tf = jnp.split(x_tf, 3)[2].reshape(-1, 3)
+
+        return CW_f̅, CW_f̿, m_tf
+
+    @staticmethod
+    @jax.jit
+    def _contact_points_dynamics(
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        *,
+        link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
+        indices_of_enabled_collidable_points: jtp.VectorLike | None = None,
+        p_t0: jtp.MatrixLike | None = None,
+        v_t0: jtp.MatrixLike | None = None,
+        m_t0: jtp.MatrixLike | None = None,
+    ) -> tuple[jtp.Matrix, jtp.Vector, jtp.Matrix, jtp.Vector]:
+        """
+        Compute the dynamics of the contact points.
+
+        Note:
+            This function projects the system dynamics to the contact space and
+            returns the matrices of a linear system to simulate its evolution.
+            Since the active contact model can be non-linear, this function also
+            linearizes the contact model at the initial state.
+
+        Args:
+            model: The robot model considered by the contact model.
+            data: The data of the considered model.
+            link_forces: The 6D forces to apply to the links.
+            joint_force_references: The joint force references to apply.
+            indices_of_enabled_collidable_points:
+                The indices of the enabled collidable points.
+            p_t0: The initial position of the collidable points.
+            v_t0: The initial velocity of the collidable points.
+            m_t0: The initial tangential deformation of the collidable points.
+
+        Returns:
+            A tuple containing:
+            - The `A` matrix of the linear system that models the contact dynamics.
+            - The `b` vector of the linear system that models the contact dynamics.
+            - The `A_sc` matrix of the linear system that approximates the contact model.
+            - The `b_sc` vector of the linear system that approximates the contact model.
+        """
+
+        indices_of_enabled_collidable_points = (
+            indices_of_enabled_collidable_points
+            if indices_of_enabled_collidable_points is not None
+            else jnp.arange(
+                len(model.kin_dyn_parameters.contact_parameters.body)
+            ).astype(int)
+        )
+
+        p_t0 = jnp.atleast_2d(
+            p_t0
+            if p_t0 is not None
+            else js.contact.collidable_point_positions(model=model, data=data)[
+                indices_of_enabled_collidable_points, :
+            ]
+        )
+
+        v_t0 = jnp.atleast_2d(
+            v_t0
+            if v_t0 is not None
+            else js.contact.collidable_point_velocities(model=model, data=data)[
+                indices_of_enabled_collidable_points, :
+            ]
+        )
+
+        m_t0 = jnp.atleast_2d(
+            m_t0
+            if m_t0 is not None
+            else data.state.extended["tangential_deformation"][
+                indices_of_enabled_collidable_points, :
+            ]
+        )
+
+        # We expect that the 6D forces of the `link_forces` argument are expressed
+        # in the frame corresponding to the velocity representation of `data`.
+        references = js.references.JaxSimModelReferences.build(
+            model=model,
+            link_forces=link_forces,
+            joint_force_references=joint_force_references,
+            data=data,
+            velocity_representation=data.velocity_representation,
+        )
+
+        # ===========================
+        # Linearize the contact model
+        # ===========================
+
+        # Linearize the contact model at the initial state of all considered
+        # contact points.
+        A_sc_points, b_sc_points = jax.vmap(
+            lambda p, v, m: ViscoElasticContacts._linearize_contact_model(
+                position=p,
+                velocity=v,
+                tangential_deformation=m,
+                parameters=data.contacts_params,
+                terrain=model.terrain,
+            )
+        )(p_t0, v_t0, m_t0)
+
+        # Since x = [p1, p2, ..., v1, v2, ..., m1, m2, ...], we need to split the A_sc of
+        # individual points since otherwise we'd get x = [ p1, v1, m1, p2, v2, m2, ...].
+        A_sc_p, A_sc_v, A_sc_m = jnp.split(A_sc_points, indices_or_sections=3, axis=-1)
+
+        # We want to have in output first the forces and then the material deformation rates.
+        # Therefore, we need to extract the components is A_sc_* separately.
+        A_sc = jnp.vstack(
+            [
+                jnp.hstack(
+                    [
+                        jax.scipy.linalg.block_diag(*A_sc_p[:, 0:3, :]),
+                        jax.scipy.linalg.block_diag(*A_sc_v[:, 0:3, :]),
+                        jax.scipy.linalg.block_diag(*A_sc_m[:, 0:3, :]),
+                    ],
+                ),
+                jnp.hstack(
+                    [
+                        jax.scipy.linalg.block_diag(*A_sc_p[:, 3:6, :]),
+                        jax.scipy.linalg.block_diag(*A_sc_v[:, 3:6, :]),
+                        jax.scipy.linalg.block_diag(*A_sc_m[:, 3:6, :]),
+                    ]
+                ),
+            ]
+        )
+
+        # We need to do the same for the b_sc.
+        b_sc = jnp.hstack(
+            [b_sc_points[:, 0:3].flatten(), b_sc_points[:, 3:6].flatten()]
+        )
+
+        # ===========================================================
+        # Compute the A and b matrices of the contact points dynamics
+        # ===========================================================
+
+        with data.switch_velocity_representation(jaxsim.VelRepr.Mixed):
+
+            BW_ν = data.generalized_velocity()
+
+            M = js.model.free_floating_mass_matrix(model=model, data=data)
+
+            CW_Jl_WC = js.contact.jacobian(
+                model=model,
+                data=data,
+                output_vel_repr=jaxsim.VelRepr.Mixed,
+            )[indices_of_enabled_collidable_points, 0:3, :]
+
+            CW_J̇l_WC = js.contact.jacobian_derivative(
+                model=model, data=data, output_vel_repr=jaxsim.VelRepr.Mixed
+            )[indices_of_enabled_collidable_points, 0:3, :]
+
+        # Compute the Delassus matrix.
+        ψ = jnp.vstack(CW_Jl_WC) @ jnp.linalg.lstsq(M, jnp.vstack(CW_Jl_WC).T)[0]
+
+        I_nc = jnp.eye(v_t0.flatten().size)
+        O_nc = jnp.zeros(shape=(p_t0.flatten().size, p_t0.flatten().size))
+
+        # Pack the A matrix.
+        A = jnp.vstack(
+            [
+                jnp.hstack([O_nc, I_nc, O_nc]),
+                ψ @ jnp.split(A_sc, 2, axis=0)[0],
+                jnp.split(A_sc, 2, axis=0)[1],
+            ]
+        )
+
+        # Short names for few variables.
+        ν = BW_ν
+        J = jnp.vstack(CW_Jl_WC)
+        J̇ = jnp.vstack(CW_J̇l_WC)
+
+        # Compute the free system acceleration components.
+        with (
+            data.switch_velocity_representation(jaxsim.VelRepr.Mixed),
+            references.switch_velocity_representation(jaxsim.VelRepr.Mixed),
+        ):
+
+            BW_v̇_free_WB, s̈_free = js.ode.system_acceleration(
+                model=model,
+                data=data,
+                joint_force_references=references.joint_force_references(model=model),
+                link_forces=references.link_forces(model=model, data=data),
+            )
+
+        # Pack the free system acceleration in mixed representation.
+        ν̇_free = jnp.hstack([BW_v̇_free_WB, s̈_free])
+
+        # Compute the acceleration of collidable points.
+        # This is the true derivative of ṗ only in mixed representation.
+        p̈ = J @ ν̇_free + J̇ @ ν
+
+        # Pack the b array.
+        b = jnp.hstack(
+            [
+                jnp.zeros_like(p_t0.flatten()),
+                p̈ + ψ @ jnp.split(b_sc, indices_or_sections=2)[0],
+                jnp.split(b_sc, indices_or_sections=2)[1],
+            ]
+        )
+
+        return A, b, A_sc, b_sc
+
+    @staticmethod
+    @functools.partial(jax.jit, static_argnames=("terrain",))
+    def _linearize_contact_model(
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
+        tangential_deformation: jtp.VectorLike,
+        parameters: ViscoElasticContactsParams,
+        terrain: Terrain,
+    ) -> tuple[jtp.Matrix, jtp.Vector]:
+        """"""
+
+        # Initialize the state at which the model is linearized.
+        p0 = jnp.array(position, dtype=float).squeeze()
+        v0 = jnp.array(velocity, dtype=float).squeeze()
+        m0 = jnp.array(tangential_deformation, dtype=float).squeeze()
+
+        # ============
+        # Compute A_sc
+        # ============
+
+        compute_contact_force_non_linear_model = functools.partial(
+            ViscoElasticContacts._compute_contact_force_non_linear_model,
+            parameters=parameters,
+            terrain=terrain,
+        )
+
+        # Compute with AD the functions to get the Jacobians of CW_fl.
+        df_dp_fun, df_dv_fun, df_dm_fun = (
+            jax.jacrev(
+                lambda p0, v0, m0: compute_contact_force_non_linear_model(
+                    position=p0, velocity=v0, tangential_deformation=m0
+                )[0],
+                argnums=num,
+            )
+            for num in (0, 1, 2)
+        )
+
+        # Compute with AD the functions to get the Jacobians of ṁ.
+        dṁ_dp_fun, dṁ_dv_fun, dṁ_dm_fun = (
+            jax.jacrev(
+                lambda p0, v0, m0: compute_contact_force_non_linear_model(
+                    position=p0, velocity=v0, tangential_deformation=m0
+                )[1],
+                argnums=num,
+            )
+            for num in (0, 1, 2)
+        )
+
+        # Compute the Jacobians of the contact forces w.r.t. the state.
+        df_dp = jnp.vstack(df_dp_fun(p0, v0, m0))
+        df_dv = jnp.vstack(df_dv_fun(p0, v0, m0))
+        df_dm = jnp.vstack(df_dm_fun(p0, v0, m0))
+
+        # Compute the Jacobians of the material deformation rate w.r.t. the state.
+        dṁ_dp = jnp.vstack(dṁ_dp_fun(p0, v0, m0))
+        dṁ_dv = jnp.vstack(dṁ_dv_fun(p0, v0, m0))
+        dṁ_dm = jnp.vstack(dṁ_dm_fun(p0, v0, m0))
+
+        # Pack the A matrix.
+        A_sc = jnp.vstack(
+            [
+                jnp.hstack([df_dp, df_dv, df_dm]),
+                jnp.hstack([dṁ_dp, dṁ_dv, dṁ_dm]),
+            ]
+        )
+
+        # ============
+        # Compute b_sc
+        # ============
+
+        # Compute the output of the non-linear model at the initial state.
+        x0 = jnp.hstack([p0, v0, m0])
+        f0, ṁ0 = compute_contact_force_non_linear_model(
+            position=p0, velocity=v0, tangential_deformation=m0
+        )
+
+        # Pack the b vector.
+        b_sc = jnp.hstack([f0, ṁ0]) - A_sc @ x0
+
+        return A_sc, b_sc
+
+    @staticmethod
+    @functools.partial(jax.jit, static_argnames=("terrain",))
+    def _compute_contact_force_non_linear_model(
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
+        tangential_deformation: jtp.VectorLike,
+        parameters: ViscoElasticContactsParams,
+        terrain: Terrain,
+    ) -> tuple[jtp.Vector, jtp.Vector]:
+        """
+        Compute the contact forces using the non-linear Hunt/Crossley model.
+
+        Args:
+            position: The position of the contact point.
+            velocity: The velocity of the contact point.
+            tangential_deformation: The tangential deformation of the contact point.
+            parameters: The parameters of the contact model.
+            terrain: The considered terrain.
+
+        Returns:
+            A tuple containing:
+            - The linear contact force in the mixed contact frame.
+            - The rate of material deformation.
+        """
+
+        # Compute the linear contact force in mixed representation using
+        # the non-linear Hunt/Crossley model.
+        # The following function also returns the rate of material deformation.
+        CW_fl, ṁ = SoftContacts.hunt_crossley_contact_model(
+            position=position,
+            velocity=velocity,
+            tangential_deformation=tangential_deformation,
+            terrain=terrain,
+            K=parameters.K,
+            D=parameters.D,
+            mu=parameters.static_friction,
+            p=parameters.p,
+            q=parameters.q,
+        )
+
+        return CW_fl, ṁ
+
+    @staticmethod
+    @jax.jit
+    def integrate_data_with_average_contact_forces(
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        *,
+        dt: jtp.FloatLike,
+        link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
+        average_link_contact_forces_inertial: jtp.MatrixLike | None = None,
+        average_of_average_link_contact_forces_mixed: jtp.MatrixLike | None = None,
+    ) -> js.data.JaxSimModelData:
+        """
+        Advance the system state by integrating the dynamics.
+
+        Args:
+            model: The model to consider.
+            data: The data of the considered model.
+            dt: The integration time step.
+            link_forces:
+                The 6D forces to apply to the links expressed in the frame corresponding
+                to the velocity representation of `data`.
+            joint_force_references: The joint force references to apply.
+            average_link_contact_forces_inertial:
+                The average contact forces computed with the exponential integrator and
+                expressed in the inertial-fixed frame.
+            average_of_average_link_contact_forces_mixed:
+                The average of the average contact forces computed with the exponential
+                integrator and expressed in the mixed frame.
+
+        Returns:
+            The data object storing the system state at the final time.
+        """
+
+        s_t0 = data.joint_positions()
+        W_p_B_t0 = data.base_position()
+        W_Q_B_t0 = data.base_orientation(dcm=False)
+
+        ṡ_t0 = data.joint_velocities()
+        with data.switch_velocity_representation(jaxsim.VelRepr.Mixed):
+            W_ṗ_B_t0 = data.base_velocity()[0:3]
+            W_ω_WB_t0 = data.base_velocity()[3:6]
+
+        with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
+            W_ν_t0 = data.generalized_velocity()
+
+        # We expect that the 6D forces of the `link_forces` argument are expressed
+        # in the frame corresponding to the velocity representation of `data`.
+        references = js.references.JaxSimModelReferences.build(
+            model=model,
+            link_forces=link_forces,
+            joint_force_references=joint_force_references,
+            data=data,
+            velocity_representation=data.velocity_representation,
+        )
+
+        W_f̅_L = (
+            jnp.array(average_link_contact_forces_inertial)
+            if average_link_contact_forces_inertial is not None
+            else jnp.zeros_like(references.input.physics_model.f_ext)
+        ).astype(float)
+
+        LW_f̿_L = (
+            jnp.array(average_of_average_link_contact_forces_mixed)
+            if average_of_average_link_contact_forces_mixed is not None
+            else W_f̅_L
+        ).astype(float)
+
+        # Compute the system inertial acceleration, used to integrate the system velocity.
+        # It considers the average contact forces computed with the exponential integrator.
+        with (
+            data.switch_velocity_representation(jaxsim.VelRepr.Inertial),
+            references.switch_velocity_representation(jaxsim.VelRepr.Inertial),
+        ):
+
+            W_ν̇_pr = jnp.hstack(
+                js.ode.system_acceleration(
+                    model=model,
+                    data=data,
+                    joint_force_references=references.joint_force_references(
+                        model=model
+                    ),
+                    link_forces=W_f̅_L + references.link_forces(model=model, data=data),
+                )
+            )
+
+        # Compute the system mixed acceleration, used to integrate the system position.
+        # It considers the average of the average contact forces computed with the
+        # exponential integrator.
+        with (
+            data.switch_velocity_representation(jaxsim.VelRepr.Mixed),
+            references.switch_velocity_representation(jaxsim.VelRepr.Mixed),
+        ):
+
+            BW_ν̇_pr2 = jnp.hstack(
+                js.ode.system_acceleration(
+                    model=model,
+                    data=data,
+                    joint_force_references=references.joint_force_references(
+                        model=model
+                    ),
+                    link_forces=LW_f̿_L + references.link_forces(model=model, data=data),
+                )
+            )
+
+        # Integrate the system velocity using the inertial-fixed acceleration.
+        W_ν_plus = W_ν_t0 + dt * W_ν̇_pr
+
+        # Integrate the system position using the mixed velocity.
+        q_plus = jnp.hstack(
+            [
+                # Note: here both ṗ and p̈ -> need mixed representation.
+                W_p_B_t0 + dt * W_ṗ_B_t0 + 0.5 * dt**2 * BW_ν̇_pr2[0:3],
+                jaxsim.math.Quaternion.integration(
+                    dt=dt,
+                    quaternion=W_Q_B_t0,
+                    omega=(W_ω_WB_t0 + 0.5 * dt * BW_ν̇_pr2[3:6]),
+                    omega_in_body_fixed=False,
+                ).squeeze(),
+                s_t0 + dt * ṡ_t0 + 0.5 * dt**2 * BW_ν̇_pr2[6:],
+            ]
+        )
+
+        # Create the data at the final time.
+        with data.editable(validate=True) as data_tf:
+            data_tf: js.data.JaxSimModelData
+            data_tf.time_ns = data.time_ns + (dt * 1e9).astype(data.time_ns.dtype)
+
+        data_tf = data_tf.reset_joint_positions(q_plus[7:])
+        data_tf = data_tf.reset_base_position(q_plus[0:3])
+        data_tf = data_tf.reset_base_quaternion(q_plus[3:7])
+
+        data_tf = data_tf.reset_joint_velocities(W_ν_plus[6:])
+        data_tf = data_tf.reset_base_velocity(
+            W_ν_plus[0:6], velocity_representation=jaxsim.VelRepr.Inertial
+        )
+
+        return data_tf.replace(
+            velocity_representation=data.velocity_representation, validate=False
+        )
+
+
+@jax.jit
+def step(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    dt: jtp.FloatLike,
+    link_forces: jtp.MatrixLike | None = None,
+    joint_force_references: jtp.VectorLike | None = None,
+) -> tuple[js.data.JaxSimModelData, dict[str, Any]]:
+    """
+    Step the system dynamics with the visco-elastic contact model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        dt: The time step to consider.
+        link_forces:
+            The 6D forces to apply to the links expressed in the frame corresponding to
+            the velocity representation of `data`.
+        joint_force_references: The joint force references to consider.
+
+    Returns:
+        A tuple containing the new data of the model
+        and an empty dictionary of auxiliary data.
+    """
+
+    assert isinstance(model.contact_model, ViscoElasticContacts)
+    assert isinstance(data.contacts_params, ViscoElasticContactsParams)
+
+    # Compute the contact forces with the exponential integrator.
+    W_f̅_C, (W_f̿_C, m_tf) = model.contact_model.compute_contact_forces(
+        model=model,
+        data=data,
+        dt=dt,
+        link_forces=link_forces,
+        joint_force_references=joint_force_references,
+    )
+
+    # ===============================
+    # Compute the link contact forces
+    # ===============================
+
+    # Extract the indices corresponding to the enabled collidable points.
+    # The visco-elastic contact model computed only their contact forces.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    # Compute the link transforms.
+    W_H_L = js.model.forward_kinematics(model=model, data=data)
+
+    # 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(
+        model.kin_dyn_parameters.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
+    W_f̿_L = mask.T @ W_f̿_C
+
+    # For integration purpose, we need these average of averages expressed in
+    # mixed representation.
+    LW_f̿_L = jax.vmap(
+        lambda W_f_L, W_H_L: data.inertial_to_other_representation(
+            array=W_f_L,
+            other_representation=jaxsim.VelRepr.Mixed,
+            transform=W_H_L,
+            is_force=True,
+        )
+    )(W_f̿_L, W_H_L)
+
+    # ==========================
+    # Integrate the system state
+    # ==========================
+
+    # Integrate the system dynamics using the average contact forces.
+    data_tf: js.data.JaxSimModelData = (
+        model.contact_model.integrate_data_with_average_contact_forces(
+            model=model,
+            data=data,
+            dt=dt,
+            link_forces=link_forces,
+            joint_force_references=joint_force_references,
+            average_link_contact_forces_inertial=W_f̅_L,
+            average_of_average_link_contact_forces_mixed=LW_f̿_L,
+        )
+    )
+
+    # Store the tangential deformation at the final state.
+    # Note that this was integrated in the continuous time domain, therefore it should
+    # be much more accurate than the one computed with the discrete soft contacts.
+    with data_tf.mutable_context():
+
+        data_tf.state.extended |= {
+            "tangential_deformation": data_tf.state.extended["tangential_deformation"]
+            .at[indices_of_enabled_collidable_points]
+            .set(m_tf)
+        }
+
+    return data_tf, {}