jaxsim 0.4.3.dev105__py3-none-any.whl → 0.4.3.dev118__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev105'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev105')
+__version__ = version = '0.4.3.dev118'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev118')

jaxsim/api/contact.py

@@ -186,7 +186,9 @@ def collidable_point_dynamics(
             # Note that the material deformation rate is always returned in the mixed frame
             # C[W] = (W_p_C, [W]). This is convenient for integration purpose.
             W_f_Ci, (CW_ṁ,) = jax.vmap(soft_contacts.compute_contact_forces)(
-                W_p_Ci, W_ṗ_Ci, data.state.contact.tangential_deformation
+                position=W_p_Ci,
+                velocity=W_ṗ_Ci,
+                tangential_deformation=data.state.contact.tangential_deformation,
             )
             aux_data = dict(m_dot=CW_ṁ)
 

jaxsim/rbda/contacts/__init__.py

@@ -0,0 +1,8 @@
+from .common import ContactModel, ContactsParams, ContactsState
+from .relaxed_rigid import (
+    RelaxedRigidContacts,
+    RelaxedRigidContactsParams,
+    RelaxedRigidContactsState,
+)
+from .rigid import RigidContacts, RigidContactsParams, RigidContactsState
+from .soft import SoftContacts, SoftContactsParams, SoftContactsState

jaxsim/rbda/contacts/common.py

@@ -7,15 +7,20 @@ import jaxsim.terrain
 import jaxsim.typing as jtp
 from jaxsim.utils import JaxsimDataclass
 
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
 
-class ContactsState(abc.ABC):
+
+class ContactsState(JaxsimDataclass):
     """
     Abstract class storing the state of the contacts model.
     """
 
     @classmethod
     @abc.abstractmethod
-    def build(cls, **kwargs) -> ContactsState:
+    def build(cls: type[Self], **kwargs) -> Self:
         """
         Build the contact state object.
 
@@ -26,7 +31,7 @@ class ContactsState(abc.ABC):
 
     @classmethod
     @abc.abstractmethod
-    def zero(cls, **kwargs) -> ContactsState:
+    def zero(cls: type[Self], **kwargs) -> Self:
         """
         Build a zero contact state.
 
@@ -36,7 +41,7 @@ class ContactsState(abc.ABC):
         pass
 
     @abc.abstractmethod
-    def valid(self, **kwargs) -> bool:
+    def valid(self, **kwargs) -> jtp.BoolLike:
         """
         Check if the contacts state is valid.
         """
@@ -50,18 +55,20 @@ class ContactsParams(JaxsimDataclass):
 
     @classmethod
     @abc.abstractmethod
-    def build(cls) -> ContactsParams:
+    def build(cls: type[Self], **kwargs) -> Self:
         """
         Create a `ContactsParams` instance with specified parameters.
+
         Returns:
             The `ContactsParams` instance.
         """
         pass
 
     @abc.abstractmethod
-    def valid(self, *args, **kwargs) -> bool:
+    def valid(self, **kwargs) -> jtp.BoolLike:
         """
         Check if the parameters are valid.
+
         Returns:
             True if the parameters are valid, False otherwise.
         """
@@ -83,8 +90,8 @@ class ContactModel(JaxsimDataclass):
     @abc.abstractmethod
     def compute_contact_forces(
         self,
-        position: jtp.Vector,
-        velocity: jtp.Vector,
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
         **kwargs,
     ) -> tuple[jtp.Vector, tuple[Any, ...]]:
         """

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -16,6 +16,11 @@ from jaxsim.terrain.terrain import FlatTerrain, Terrain
 
 from .common import ContactModel, ContactsParams, ContactsState
 
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
 
 @jax_dataclasses.pytree_dataclass
 class RelaxedRigidContactsParams(ContactsParams):
@@ -106,7 +111,8 @@ class RelaxedRigidContactsParams(ContactsParams):
 
     @classmethod
     def build(
-        cls,
+        cls: type[Self],
+        *,
         time_constant: jtp.FloatLike | None = None,
         damping_coefficient: jtp.FloatLike | None = None,
         d_min: jtp.FloatLike | None = None,
@@ -119,7 +125,7 @@ class RelaxedRigidContactsParams(ContactsParams):
         mu: jtp.FloatLike | None = None,
         max_iterations: jtp.IntLike | None = None,
         tolerance: jtp.FloatLike | None = None,
-    ) -> RelaxedRigidContactsParams:
+    ) -> Self:
         """Create a `RelaxedRigidContactsParams` instance"""
 
         return cls(
@@ -132,7 +138,8 @@ class RelaxedRigidContactsParams(ContactsParams):
             }
         )
 
-    def valid(self) -> bool:
+    def valid(self) -> jtp.BoolLike:
+
         return bool(
             jnp.all(self.time_constant >= 0.0)
             and jnp.all(self.damping_coefficient > 0.0)
@@ -155,18 +162,19 @@ class RelaxedRigidContactsState(ContactsState):
     def __eq__(self, other: RelaxedRigidContactsState) -> bool:
         return hash(self) == hash(other)
 
-    @staticmethod
-    def build() -> RelaxedRigidContactsState:
+    @classmethod
+    def build(cls: type[Self]) -> Self:
         """Create a `RelaxedRigidContactsState` instance"""
 
-        return RelaxedRigidContactsState()
+        return cls()
 
-    @staticmethod
-    def zero(model: js.model.JaxSimModel) -> RelaxedRigidContactsState:
+    @classmethod
+    def zero(cls: type[Self], **kwargs) -> Self:
         """Build a zero `RelaxedRigidContactsState` instance from a `JaxSimModel`."""
-        return RelaxedRigidContactsState.build()
 
-    def valid(self, model: js.model.JaxSimModel) -> bool:
+        return cls.build()
+
+    def valid(self, **kwargs) -> jtp.BoolLike:
         return True
 
 
@@ -182,10 +190,12 @@ class RelaxedRigidContacts(ContactModel):
         default_factory=FlatTerrain
     )
 
+    @jax.jit
     def compute_contact_forces(
         self,
-        position: jtp.Vector,
-        velocity: jtp.Vector,
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
+        *,
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
         link_forces: jtp.MatrixLike | None = None,

jaxsim/rbda/contacts/rigid.py

@@ -14,6 +14,11 @@ from jaxsim.terrain import FlatTerrain, Terrain
 
 from .common import ContactModel, ContactsParams, ContactsState
 
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
 
 @jax_dataclasses.pytree_dataclass
 class RigidContactsParams(ContactsParams):
@@ -50,19 +55,22 @@ class RigidContactsParams(ContactsParams):
 
     @classmethod
     def build(
-        cls,
+        cls: type[Self],
+        *,
         mu: jtp.FloatLike | None = None,
         K: jtp.FloatLike | None = None,
         D: jtp.FloatLike | None = None,
-    ) -> RigidContactsParams:
+    ) -> Self:
         """Create a `RigidContactParams` instance"""
-        return RigidContactsParams(
+
+        return cls(
             mu=mu or cls.__dataclass_fields__["mu"].default,
             K=K or cls.__dataclass_fields__["K"].default,
             D=D or cls.__dataclass_fields__["D"].default,
         )
 
-    def valid(self) -> bool:
+    def valid(self) -> jtp.BoolLike:
+
         return bool(
             jnp.all(self.mu >= 0.0)
             and jnp.all(self.K >= 0.0)
@@ -77,18 +85,19 @@ class RigidContactsState(ContactsState):
     def __eq__(self, other: RigidContactsState) -> bool:
         return hash(self) == hash(other)
 
-    @staticmethod
-    def build(**kwargs) -> RigidContactsState:
+    @classmethod
+    def build(cls: type[Self]) -> Self:
         """Create a `RigidContactsState` instance"""
 
-        return RigidContactsState()
+        return cls()
 
-    @staticmethod
-    def zero(**kwargs) -> RigidContactsState:
+    @classmethod
+    def zero(cls: type[Self], **kwargs) -> Self:
         """Build a zero `RigidContactsState` instance from a `JaxSimModel`."""
-        return RigidContactsState.build()
 
-    def valid(self, **kwargs) -> bool:
+        return cls.build()
+
+    def valid(self, **kwargs) -> jtp.BoolLike:
         return True
 
 
@@ -117,7 +126,8 @@ class RigidContacts(ContactModel):
             terrain_height: The height of the terrain at the collidable point position.
 
         Returns:
-            A tuple containing the activation state of the collidable points and the contact penetration depth h.
+            A tuple containing the activation state of the collidable points
+            and the contact penetration depth h.
         """
 
         # TODO: reduce code duplication with js.contact.in_contact
@@ -154,8 +164,8 @@ class RigidContacts(ContactModel):
 
         Args:
             inactive_collidable_points: The activation state of the collidable points.
-            M: The mass matrix of the system.
-            J_WC: The Jacobian matrix of the collidable points.
+            M: The mass matrix of the system (in mixed representation).
+            J_WC: The Jacobian matrix of the collidable points (in mixed representation).
             data: The `JaxSimModelData` instance.
         """
 
@@ -206,10 +216,12 @@ class RigidContacts(ContactModel):
 
         return BW_ν_post_impact
 
+    @jax.jit
     def compute_contact_forces(
         self,
-        position: jtp.Vector,
-        velocity: jtp.Vector,
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
+        *,
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
         link_forces: jtp.MatrixLike | None = None,

jaxsim/rbda/contacts/soft.py

@@ -7,12 +7,18 @@ import jax.numpy as jnp
 import jax_dataclasses
 
 import jaxsim.api as js
+import jaxsim.math
 import jaxsim.typing as jtp
-from jaxsim.math import Skew, StandardGravity
+from jaxsim.math import StandardGravity
 from jaxsim.terrain import FlatTerrain, Terrain
 
 from .common import ContactModel, ContactsParams, ContactsState
 
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
 
 @jax_dataclasses.pytree_dataclass
 class SoftContactsParams(ContactsParams):
@@ -30,6 +36,14 @@ class SoftContactsParams(ContactsParams):
         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)
+    )
+
     def __hash__(self) -> int:
 
         from jaxsim.utils.wrappers import HashedNumpyArray
@@ -39,6 +53,8 @@ class SoftContactsParams(ContactsParams):
                 HashedNumpyArray.hash_of_array(self.K),
                 HashedNumpyArray.hash_of_array(self.D),
                 HashedNumpyArray.hash_of_array(self.mu),
+                HashedNumpyArray.hash_of_array(self.p),
+                HashedNumpyArray.hash_of_array(self.q),
             )
         )
 
@@ -49,10 +65,16 @@ class SoftContactsParams(ContactsParams):
 
         return hash(self) == hash(other)
 
-    @staticmethod
+    @classmethod
     def build(
-        K: jtp.FloatLike = 1e6, D: jtp.FloatLike = 2_000, mu: jtp.FloatLike = 0.5
-    ) -> SoftContactsParams:
+        cls: type[Self],
+        *,
+        K: jtp.FloatLike = 1e6,
+        D: jtp.FloatLike = 2_000,
+        mu: jtp.FloatLike = 0.5,
+        p: jtp.FloatLike = 0.5,
+        q: jtp.FloatLike = 0.5,
+    ) -> Self:
         """
         Create a SoftContactsParams instance with specified parameters.
 
@@ -60,6 +82,12 @@ class SoftContactsParams(ContactsParams):
             K: The stiffness parameter.
             D: The damping parameter of the soft contacts model.
             mu: 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 SoftContactsParams instance with the specified parameters.
@@ -69,10 +97,13 @@ class SoftContactsParams(ContactsParams):
             K=jnp.array(K, dtype=float),
             D=jnp.array(D, dtype=float),
             mu=jnp.array(mu, dtype=float),
+            p=jnp.array(p, dtype=float),
+            q=jnp.array(q, dtype=float),
         )
 
-    @staticmethod
+    @classmethod
     def build_default_from_jaxsim_model(
+        cls: type[Self],
         model: js.model.JaxSimModel,
         *,
         standard_gravity: jtp.FloatLike = StandardGravity,
@@ -80,6 +111,8 @@ class SoftContactsParams(ContactsParams):
         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,
     ) -> SoftContactsParams:
         """
         Create a SoftContactsParams instance with good default parameters.
@@ -94,6 +127,12 @@ class SoftContactsParams(ContactsParams):
                 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 `SoftContactsParams` instance with the specified parameters.
@@ -126,9 +165,9 @@ class SoftContactsParams(ContactsParams):
         critical_damping = 2 * jnp.sqrt(K * m)
         D = ξ * critical_damping
 
-        return SoftContactsParams.build(K=K, D=D, mu=μc)
+        return SoftContactsParams.build(K=K, D=D, mu=μc, p=p, q=q)
 
-    def valid(self) -> bool:
+    def valid(self) -> jtp.BoolLike:
         """
         Check if the parameters are valid.
 
@@ -136,11 +175,15 @@ class SoftContactsParams(ContactsParams):
             `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.mu >= 0.0)
-        )
+        return jnp.hstack(
+            [
+                self.K >= 0.0,
+                self.D >= 0.0,
+                self.mu >= 0.0,
+                self.p >= 0.0,
+                self.q >= 0.0,
+            ]
+        ).all()
 
 
 @jax_dataclasses.pytree_dataclass
@@ -157,179 +200,158 @@ class SoftContacts(ContactModel):
 
     def compute_contact_forces(
         self,
-        position: jtp.Vector,
-        velocity: jtp.Vector,
-        tangential_deformation: jtp.Vector,
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
+        *,
+        tangential_deformation: jtp.VectorLike,
     ) -> tuple[jtp.Vector, tuple[jtp.Vector]]:
-        """
-        Compute the contact forces and material deformation rate.
 
-        Args:
-            position: The position of the collidable point.
-            velocity: The linear velocity of the collidable point.
-            tangential_deformation: The tangential deformation.
-
-        Returns:
-            A tuple containing the contact force and material deformation rate.
-        """
+        # Convert the input vectors to arrays.
+        W_p_C = jnp.array(position, dtype=float).squeeze()
+        W_ṗ_C = jnp.array(velocity, dtype=float).squeeze()
+        m = jnp.array(tangential_deformation, dtype=float).squeeze()
 
-        # Short name of parameters
+        # Short name of parameters.
         K = self.parameters.K
         D = self.parameters.D
         μ = self.parameters.mu
 
-        # Material 3D tangential deformation and its derivative
-        m = tangential_deformation.squeeze()
-        ṁ = jnp.zeros_like(m)
+        # Compute the penetration depth, its rate, and the considered terrain normal.
+        δ, δ̇, n̂ = self.compute_penetration_data(p=W_p_C, v=W_ṗ_C, terrain=self.terrain)
 
-        # Note: all the small hardcoded tolerances in this method have been introduced
-        # to allow jax differentiating through this algorithm. They should not affect
-        # the accuracy of the simulation, although they might make it less readable.
+        # Get the exponents of the Hunt/Crossley model non-linear terms.
+        p = self.parameters.p
+        q = self.parameters.q
+
+        # There are few operations like computing the norm of a vector with zero length
+        # or computing the square root of zero that are problematic in an AD context.
+        # To avoid these issues, we introduce a small tolerance ε to their arguments
+        # and make sure that we do not check them against zero directly.
+        ε = jnp.finfo(float).eps
+
+        # Compute the powers of the penetration depth.
+        # Inject ε to address AD issues in differentiating the square root when
+        #  p and q are fractional.
+        δp = jnp.power(δ + ε, p)
+        δq = jnp.power(δ + ε, q)
 
         # ========================
-        # Normal force computation
+        # Compute the normal force
         # ========================
 
-        # Unpack the position of the collidable point.
-        px, py, pz = W_p_C = position.squeeze()
-        W_ṗ_C = velocity.squeeze()
+        # Non-linear spring-damper model (Hunt/Crossley model).
+        # This is the force magnitude along the direction normal to the terrain.
+        force_normal_mag = (K * δp) * δ + (D * δq) * δ̇
 
-        # Compute the terrain normal and the contact depth.
-        n̂ = self.terrain.normal(x=px, y=py).squeeze()
-        h = jnp.array([0, 0, self.terrain.height(x=px, y=py) - pz])
+        # Depending on the magnitude of δ̇, the normal force could be negative.
+        force_normal_mag = jnp.maximum(0.0, force_normal_mag)
 
-        # Compute the penetration depth normal to the terrain.
-        δ = jnp.maximum(0.0, jnp.dot(h, n̂))
+        # Compute the 3D linear force in C[W] frame.
+        f_normal = force_normal_mag * n̂
+
+        # ============================
+        # Compute the tangential force
+        # ============================
+
+        # Extract the tangential component of the velocity.
+        v_tangential = W_ṗ_C - jnp.dot(W_ṗ_C, n̂) * n̂
+
+        # Extract the tangential component of the material deformation.
+        # This should not be necessary if the sticking-slipping transition occurs
+        # in a terrain area with a locally constant normal. However, this assumption
+        # is not true in general for highly uneven terrains.
+        m_normal = jnp.dot(m, n̂) * n̂
+        m_tangential = m - jnp.dot(m, n̂) * n̂
+
+        # Compute the tangential force in the sticking case.
+        f_tangential = -((K * δp) * m_tangential + (D * δq) * v_tangential)
+
+        # Detect the contact type (sticking or slipping).
+        # Note that if there is no contact, sticking is set to True, and this detail
+        # is exploited in the computation of the `contact_status` variable.
+        sticking = jnp.logical_or(
+            δ <= 0, f_tangential.dot(f_tangential) <= (μ * force_normal_mag) ** 2
+        )
 
-        # Compute the penetration normal velocity.
-        δ̇ = -jnp.dot(W_ṗ_C, n̂)
+        # Compute the direction of the tangential force.
+        # To prevent dividing by zero, we use a switch statement.
+        # The ε, instead, is needed to make AD happy.
+        f_tangential_direction = jnp.where(
+            f_tangential.dot(f_tangential) != 0,
+            f_tangential / jnp.linalg.norm(f_tangential + ε),
+            jnp.zeros(3),
+        )
 
-        # Non-linear spring-damper model.
-        # This is the force magnitude along the direction normal to the terrain.
-        force_normal_mag = jax.lax.select(
-            pred=δ >= 1e-9,
-            on_true=jnp.sqrt(δ + 1e-12) * (K * δ + D * δ̇),
-            on_false=jnp.array(0.0),
+        # Project the tangential force to the friction cone if slipping.
+        f_tangential = jnp.where(
+            sticking,
+            f_tangential,
+            jnp.minimum(μ * force_normal_mag, jnp.linalg.norm(f_tangential + ε))
+            * f_tangential_direction,
         )
 
-        # Prevent negative normal forces that might occur when δ̇ is largely negative.
-        force_normal_mag = jnp.maximum(0.0, force_normal_mag)
+        # Set the tangential force to zero if there is no contact.
+        f_tangential = jnp.where(δ <= 0, jnp.zeros(3), f_tangential)
 
-        # Compute the 3D linear force in C[W] frame.
-        force_normal = force_normal_mag * n̂
+        # =====================================
+        # Compute the material deformation rate
+        # =====================================
 
-        # ====================================
-        # No friction and no tangential forces
-        # ====================================
+        # Compute the derivative of the material deformation.
+        ṁ_no_contact = -(K / D) * m
+        ṁ_sticking = v_tangential - (K / D) * m_normal
+        ṁ_slipping = -(f_tangential + (K * δp) * m_tangential) / (D * δq)
 
-        # Compute the adjoint C[W]->W for transforming 6D forces from mixed to inertial.
-        # Note: this is equal to the 6D velocities transform: CW_X_W.transpose().
-        W_Xf_CW = jnp.vstack(
-            [
-                jnp.block([jnp.eye(3), jnp.zeros(shape=(3, 3))]),
-                jnp.block([Skew.wedge(W_p_C), jnp.eye(3)]),
-            ]
-        )
+        # Compute the contact status:
+        # 0: slipping
+        # 1: sticking
+        # 2: no contact
+        contact_status = sticking.astype(int)
+        contact_status += (δ <= 0).astype(int)
 
-        def with_no_friction():
-            # Compute 6D mixed force in C[W].
-            CW_f_lin = force_normal
-            CW_f = jnp.hstack([force_normal, jnp.zeros_like(CW_f_lin)])
-
-            # Compute lin-ang 6D forces (inertial representation).
-            W_f = W_Xf_CW @ CW_f
-
-            return W_f, (ṁ,)
-
-        # =========================
-        # Compute tangential forces
-        # =========================
-
-        def with_friction():
-            # Initialize the tangential deformation rate ṁ.
-            # For inactive contacts with m≠0, this is the dynamics of the material
-            # relaxation converging exponentially to steady state.
-            ṁ = (-K / D) * m
-
-            # Check if the collidable point is below ground.
-            # Note: when δ=0, we consider the point still not it contact such that
-            #       we prevent divisions by 0 in the computations below.
-            active_contact = pz < self.terrain.height(x=px, y=py)
-
-            def above_terrain():
-                return jnp.zeros(6), (ṁ,)
-
-            def below_terrain():
-                # Decompose the velocity in normal and tangential components.
-                v_normal = jnp.dot(W_ṗ_C, n̂) * n̂
-                v_tangential = W_ṗ_C - v_normal
-
-                # Compute the tangential force. If inside the friction cone, the contact.
-                f_tangential = -jnp.sqrt(δ + 1e-12) * (K * m + D * v_tangential)
-
-                def sticking_contact():
-                    # Sum the normal and tangential forces, and create the 6D force.
-                    CW_f_stick = force_normal + f_tangential
-                    CW_f = jnp.hstack([CW_f_stick, jnp.zeros(3)])
-
-                    # In this case the 3D material deformation is the tangential velocity.
-                    ṁ = v_tangential
-
-                    # Return the 6D force in the contact frame and
-                    # the deformation derivative.
-                    return CW_f, ṁ
-
-                def slipping_contact():
-                    # Project the force to the friction cone boundary.
-                    f_tangential_projected = (μ * force_normal_mag) * (
-                        f_tangential / jnp.maximum(jnp.linalg.norm(f_tangential), 1e-9)
-                    )
-
-                    # Sum the normal and tangential forces, and create the 6D force.
-                    CW_f_slip = force_normal + f_tangential_projected
-                    CW_f = jnp.hstack([CW_f_slip, jnp.zeros(3)])
-
-                    # Correct the material deformation derivative for slipping contacts.
-                    # Basically we compute ṁ such that we get `f_tangential` on the cone
-                    # given the current (m, δ).
-                    ε = 1e-9
-                    δε = jnp.maximum(δ, ε)
-                    α = -K * jnp.sqrt(δε)
-                    β = -D * jnp.sqrt(δε)
-                    ṁ = (f_tangential_projected - α * m) / β
-
-                    # Return the 6D force in the contact frame and
-                    # the deformation derivative.
-                    return CW_f, ṁ
-
-                CW_f, ṁ = jax.lax.cond(
-                    pred=f_tangential.dot(f_tangential) > (μ * force_normal_mag) ** 2,
-                    true_fun=lambda _: slipping_contact(),
-                    false_fun=lambda _: sticking_contact(),
-                    operand=None,
-                )
-
-                # Express the 6D force in the world frame.
-                W_f = W_Xf_CW @ CW_f
-
-                # Return the 6D force in the world frame and the deformation derivative.
-                return W_f, (ṁ,)
-
-            # (W_f, (ṁ,))
-            return jax.lax.cond(
-                pred=active_contact,
-                true_fun=lambda _: below_terrain(),
-                false_fun=lambda _: above_terrain(),
-                operand=None,
-            )
+        # Select the right material deformation rate depending on the contact status.
+        ṁ = jax.lax.select_n(contact_status, ṁ_slipping, ṁ_sticking, ṁ_no_contact)
 
-        # (W_f, (ṁ,))
-        return jax.lax.cond(
-            pred=(μ == 0.0),
-            true_fun=lambda _: with_no_friction(),
-            false_fun=lambda _: with_friction(),
-            operand=None,
-        )
+        # ==========================================
+        # Compute and return the final contact force
+        # ==========================================
+
+        # Sum the normal and tangential forces and create a mixed 6D force.
+        CW_f = jnp.hstack([f_normal + f_tangential, jnp.zeros(3)])
+
+        # Compute the 6D force transform from the mixed to the inertial-fixed frame.
+        W_Xf_CW = jaxsim.math.Adjoint.from_quaternion_and_translation(
+            translation=W_p_C, inverse=True
+        ).T
+
+        return W_Xf_CW @ CW_f, (ṁ,)
+
+    @staticmethod
+    @jax.jit
+    def compute_penetration_data(
+        p: jtp.VectorLike,
+        v: jtp.VectorLike,
+        terrain: jaxsim.terrain.Terrain,
+    ) -> tuple[jtp.Float, jtp.Float, jtp.Vector]:
+
+        # Pre-process the position and the linear velocity of the collidable point.
+        W_ṗ_C = jnp.array(v).squeeze()
+        px, py, pz = jnp.array(p).squeeze()
+
+        # Compute the terrain normal and the contact depth.
+        n̂ = terrain.normal(x=px, y=py).squeeze()
+        h = jnp.array([0, 0, terrain.height(x=px, y=py) - pz])
+
+        # Compute the penetration depth normal to the terrain.
+        δ = jnp.maximum(0.0, jnp.dot(h, n̂))
+
+        # Compute the penetration normal velocity.
+        δ̇ = -jnp.dot(W_ṗ_C, n̂)
+
+        # Enforce the penetration rate to be zero when the penetration depth is zero.
+        δ̇ = jnp.where(δ > 0, δ̇, 0.0)
+
+        return δ, δ̇, n̂
 
 
 @jax_dataclasses.pytree_dataclass
@@ -346,21 +368,24 @@ class SoftContactsState(ContactsState):
     tangential_deformation: jtp.Matrix
 
     def __hash__(self) -> int:
+
         return hash(
             tuple(jnp.atleast_1d(self.tangential_deformation.flatten()).tolist())
         )
 
-    def __eq__(self, other: SoftContactsState) -> bool:
-        if not isinstance(other, SoftContactsState):
+    def __eq__(self: Self, other: Self) -> bool:
+
+        if not isinstance(other, type(self)):
             return False
 
         return hash(self) == hash(other)
 
-    @staticmethod
+    @classmethod
     def build_from_jaxsim_model(
+        cls: type[Self],
         model: js.model.JaxSimModel | None = None,
-        tangential_deformation: jtp.Matrix | None = None,
-    ) -> SoftContactsState:
+        tangential_deformation: jtp.MatrixLike | None = None,
+    ) -> Self:
         """
         Build a `SoftContactsState` from a `JaxSimModel`.
 
@@ -376,18 +401,20 @@ class SoftContactsState(ContactsState):
             `JaxSimModel` and initialized to zero.
         """
 
-        return SoftContactsState.build(
+        return cls.build(
             tangential_deformation=tangential_deformation,
             number_of_collidable_points=len(
                 model.kin_dyn_parameters.contact_parameters.body
             ),
         )
 
-    @staticmethod
+    @classmethod
     def build(
-        tangential_deformation: jtp.Matrix | None = None,
+        cls: type[Self],
+        *,
+        tangential_deformation: jtp.MatrixLike | None = None,
         number_of_collidable_points: int | None = None,
-    ) -> SoftContactsState:
+    ) -> Self:
         """
         Create a `SoftContactsState`.
 
@@ -402,10 +429,10 @@ class SoftContactsState(ContactsState):
         """
 
         tangential_deformation = (
-            tangential_deformation
+            jnp.atleast_2d(tangential_deformation)
             if tangential_deformation is not None
             else jnp.zeros(shape=(number_of_collidable_points, 3))
-        )
+        ).astype(float)
 
         if tangential_deformation.shape[1] != 3:
             raise RuntimeError("The tangential deformation matrix must have 3 columns.")
@@ -418,12 +445,10 @@ class SoftContactsState(ContactsState):
             msg += "in the tangential deformation matrix."
             raise RuntimeError(msg)
 
-        return SoftContactsState(
-            tangential_deformation=jnp.array(tangential_deformation).astype(float)
-        )
+        return cls(tangential_deformation=tangential_deformation)
 
-    @staticmethod
-    def zero(model: js.model.JaxSimModel) -> SoftContactsState:
+    @classmethod
+    def zero(cls: type[Self], *, model: js.model.JaxSimModel) -> Self:
         """
         Build a zero `SoftContactsState` from a `JaxSimModel`.
 
@@ -434,9 +459,9 @@ class SoftContactsState(ContactsState):
             A zero `SoftContactsState` instance.
         """
 
-        return SoftContactsState.build_from_jaxsim_model(model=model)
+        return cls.build_from_jaxsim_model(model=model)
 
-    def valid(self, model: js.model.JaxSimModel) -> bool:
+    def valid(self, *, model: js.model.JaxSimModel) -> jtp.BoolLike:
         """
         Check if the `SoftContactsState` is valid for a given `JaxSimModel`.
 

{jaxsim-0.4.3.dev105.dist-info → jaxsim-0.4.3.dev118.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev105
+Version: 0.4.3.dev118
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.3.dev105.dist-info → jaxsim-0.4.3.dev118.dist-info}/RECORD

@@ -1,12 +1,12 @@
 jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=1Pw2WIpSfIcMXUmFaqKP5hkuhIK6fZ2PSjrJqsDco98,428
+jaxsim/_version.py,sha256=SpzKi91xA9hgMqjpOts0HnzKduXJ_CeCiMYlA0iOEgY,428
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=m-p3EJDhpnMTlXKplfbZE_aH9NqX_VyLlAE3vUhc6l4,13642
 jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
-jaxsim/api/contact.py,sha256=BQMIMHBFYiHe_LVx_bwxKCpy20uiy0V-NljHfYXWhI0,23013
+jaxsim/api/contact.py,sha256=ocwsVS1jaBfrd81990hcgfS0-2xD8VVzDq7gdPguAUg,23087
 jaxsim/api/data.py,sha256=QldUHniJqKrdNtAcXuRaS9UyeslJ0Rjvb17UA0Ca5Tw,29008
 jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
 jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
@@ -52,19 +52,19 @@ jaxsim/rbda/forward_kinematics.py,sha256=2GmEoWsrioVl_SAbKRKfhOLz57pY4aR81PKRdul
 jaxsim/rbda/jacobian.py,sha256=p0EV_8cLzLVV-93VKznT7VPuRj8W7h7rQWkPlWJXfCA,11023
 jaxsim/rbda/rnea.py,sha256=CLfqs9XFVaD-hvkLABshDAfdw5bm_AMV3UVAQ_IvURQ,7542
 jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
-jaxsim/rbda/contacts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxsim/rbda/contacts/common.py,sha256=VwAs742futAmLnDgbaOuLzNDBFiKDfYItdEZ4UcFgzE,2467
-jaxsim/rbda/contacts/relaxed_rigid.py,sha256=deTC0M2a_RER7iwVpxLCfuSlgBLqkTmHQdOJ4169IR4,13646
-jaxsim/rbda/contacts/rigid.py,sha256=BYvQgoaG5yKo2N1SLlXgjP6cb1OrJ1BawGXkJf0Hhi0,15060
-jaxsim/rbda/contacts/soft.py,sha256=_wvb5iZDjGcVg6rNQelN4LZN7qSC2NIp0HdKvZmlGfk,15647
+jaxsim/rbda/contacts/__init__.py,sha256=Y1yT2zdgFa0zviZseI09wNaMcydH8TeoaWr6ehqzwdc,328
+jaxsim/rbda/contacts/common.py,sha256=CEmLS_PT44AOWKJ0bWrJJBqm2Q9v9LiqvL0rht63-ic,2605
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=VCFU2WG7MoDA6eo72VOLzNhXd1ZEy78F-52JVNASAcU,13696
+jaxsim/rbda/contacts/rigid.py,sha256=QBjgBIJR6jz3w_dd-ZEex6K1c9_Cwwl5xrZALHzz5Zo,15224
+jaxsim/rbda/contacts/soft.py,sha256=-d7zbMdKNq0aRT2zRXIu_Dbh8BL4VUnMDprz4Ddfwj0,16276
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=xUQg47yGxIOcTkLPbnO3sruEGBhoCd16j1evTGlmNjI,5010
 jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev105.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev105.dist-info/METADATA,sha256=jYNJIkcAgei9qg7OLD7-E0HHYZRb3h4qFdmaVnyHXGo,17277
-jaxsim-0.4.3.dev105.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev105.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev105.dist-info/RECORD,,
+jaxsim-0.4.3.dev118.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev118.dist-info/METADATA,sha256=SH1vKEjwiZxKVzUXu2BqZTNkEjIgXylaoaNuYwqiQQY,17277
+jaxsim-0.4.3.dev118.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev118.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev118.dist-info/RECORD,,