jaxsim 0.6.2.dev194__py3-none-any.whl → 0.6.2.dev228__py3-none-any.whl

jaxsim/rbda/contacts/__init__.py

@@ -1,5 +1,9 @@
-from . import relaxed_rigid
+from . import relaxed_rigid, rigid, soft
 from .common import ContactModel, ContactsParams
 from .relaxed_rigid import RelaxedRigidContacts, RelaxedRigidContactsParams
+from .rigid import RigidContacts, RigidContactsParams
+from .soft import SoftContacts, SoftContactsParams
 
-ContactParamsTypes = RelaxedRigidContactsParams
+ContactParamsTypes = (
+    SoftContactsParams | RigidContactsParams | RelaxedRigidContactsParams
+)

jaxsim/rbda/contacts/common.py

@@ -9,6 +9,7 @@ import jax.numpy as jnp
 import jaxsim.api as js
 import jaxsim.terrain
 import jaxsim.typing as jtp
+from jaxsim.math import STANDARD_GRAVITY
 from jaxsim.utils import JaxsimDataclass
 
 try:
@@ -80,6 +81,86 @@ class ContactsParams(JaxsimDataclass):
         """
         pass
 
+    def build_default_from_jaxsim_model(
+        self: type[Self],
+        model: js.model.JaxSimModel,
+        *,
+        stiffness: jtp.FloatLike | None = None,
+        damping: jtp.FloatLike | None = None,
+        standard_gravity: jtp.FloatLike = STANDARD_GRAVITY,
+        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,
+        **kwargs,
+    ) -> Self:
+        """
+        Create a `ContactsParams` instance with default parameters.
+
+        Args:
+            model: The robot model considered by the contact model.
+            stiffness: The stiffness of the contact model.
+            damping: The damping of the contact model.
+            standard_gravity: The standard gravity acceleration.
+            static_friction_coefficient: The static friction coefficient.
+            max_penetration: The maximum penetration depth.
+            number_of_active_collidable_points_steady_state:
+                The number of active collidable points in steady state.
+            damping_ratio: The damping ratio.
+            p: The first parameter of the contact model.
+            q: The second parameter of the contact model.
+            **kwargs: Optional additional arguments.
+
+        Returns:
+            The `ContactsParams` instance.
+
+        Note:
+            The `stiffness` is intended as the terrain stiffness in the Soft Contacts model,
+            while it is the Baumgarte stabilization stiffness in the Rigid Contacts model.
+
+            The `damping` is intended as the terrain damping in the Soft Contacts model,
+            while it is the Baumgarte stabilization damping in the Rigid Contacts model.
+
+            The `damping_ratio` parameter allows to operate on the following conditions:
+            - ξ > 1.0: over-damped
+            - ξ = 1.0: critically damped
+            - ξ < 1.0: under-damped
+        """
+
+        # Use symbols for input parameters.
+        ξ = damping_ratio
+        δ_max = max_penetration
+        μc = static_friction_coefficient
+
+        # Compute the total mass of the model.
+        m = jnp.array(model.kin_dyn_parameters.link_parameters.mass).sum()
+
+        # Rename the standard gravity.
+        g = standard_gravity
+
+        # Compute the average support force on each collidable point.
+        f_average = m * g / number_of_active_collidable_points_steady_state
+
+        # Compute the stiffness to get the desired steady-state penetration.
+        # Note that this is dependent on the non-linear exponent used in
+        # the damping term of the Hunt/Crossley model.
+        K = f_average / jnp.power(δ_max, 1 + p) if stiffness is None else stiffness
+
+        # Compute the damping using the damping ratio.
+        critical_damping = 2 * jnp.sqrt(K * m)
+        D = ξ * critical_damping if damping is None else damping
+
+        return self.build(
+            K=K,
+            D=D,
+            mu=μc,
+            p=p,
+            q=q,
+            **kwargs,
+        )
+
     @abc.abstractmethod
     def valid(self, **kwargs) -> jtp.BoolLike:
         """
@@ -156,7 +237,7 @@ class ContactModel(JaxsimDataclass):
         return {}
 
     @property
-    def _parameters_class(cls) -> type[ContactsParams]:
+    def _parameters_class(self) -> type[ContactsParams]:
         """
         Return the class of the contact parameters.
 
@@ -168,8 +249,37 @@ class ContactModel(JaxsimDataclass):
         return getattr(
             importlib.import_module("jaxsim.rbda.contacts"),
             (
-                cls.__name__ + "Params"
-                if isinstance(cls, type)
-                else cls.__class__.__name__ + "Params"
+                self.__name__ + "Params"
+                if isinstance(self, type)
+                else self.__class__.__name__ + "Params"
             ),
         )
+
+    @abc.abstractmethod
+    def update_contact_state(
+        self: type[Self], old_contact_state: dict[str, jtp.Array]
+    ) -> dict[str, jtp.Array]:
+        """
+        Update the contact state.
+
+        Args:
+            old_contact_state: The old contact state.
+
+        Returns:
+            The updated contact state.
+        """
+
+    @abc.abstractmethod
+    def update_velocity_after_impact(
+        self: type[Self], model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+    ) -> js.data.JaxSimModelData:
+        """
+        Update the velocity after an impact.
+
+        Args:
+            model: The robot model considered by the contact model.
+            data: The data of the considered model.
+
+        Returns:
+            The updated data of the considered model.
+        """

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import dataclasses
-import functools
 from collections.abc import Callable
 from typing import Any
 
@@ -11,12 +10,10 @@ import jax_dataclasses
 import optax
 
 import jaxsim.api as js
-import jaxsim.rbda.contacts
 import jaxsim.typing as jtp
 from jaxsim.api.common import ModelDataWithVelocityRepresentation, VelRepr
-from jaxsim.terrain.terrain import Terrain
 
-from . import common
+from . import common, soft
 
 try:
     from typing import Self
@@ -64,12 +61,12 @@ class RelaxedRigidContactsParams(common.ContactsParams):
     )
 
     # Stiffness
-    stiffness: jtp.Float = dataclasses.field(
+    K: jtp.Float = dataclasses.field(
         default_factory=lambda: jnp.array(0.0, dtype=float)
     )
 
     # Damping
-    damping: jtp.Float = dataclasses.field(
+    D: jtp.Float = dataclasses.field(
         default_factory=lambda: jnp.array(0.0, dtype=float)
     )
 
@@ -90,13 +87,16 @@ class RelaxedRigidContactsParams(common.ContactsParams):
                 HashedNumpyArray(self.width),
                 HashedNumpyArray(self.midpoint),
                 HashedNumpyArray(self.power),
-                HashedNumpyArray(self.stiffness),
-                HashedNumpyArray(self.damping),
+                HashedNumpyArray(self.K),
+                HashedNumpyArray(self.D),
                 HashedNumpyArray(self.mu),
             )
         )
 
     def __eq__(self, other: RelaxedRigidContactsParams) -> bool:
+        if not isinstance(other, RelaxedRigidContactsParams):
+            return False
+
         return hash(self) == hash(other)
 
     @classmethod
@@ -110,9 +110,10 @@ class RelaxedRigidContactsParams(common.ContactsParams):
         width: jtp.FloatLike | None = None,
         midpoint: jtp.FloatLike | None = None,
         power: jtp.FloatLike | None = None,
-        stiffness: jtp.FloatLike | None = None,
-        damping: jtp.FloatLike | None = None,
+        K: jtp.FloatLike | None = None,
+        D: jtp.FloatLike | None = None,
         mu: jtp.FloatLike | None = None,
+        **kwargs,
     ) -> Self:
         """Create a `RelaxedRigidContactsParams` instance."""
 
@@ -151,13 +152,11 @@ class RelaxedRigidContactsParams(common.ContactsParams):
             power=jnp.array(
                 power if power is not None else default("power"), dtype=float
             ),
-            stiffness=jnp.array(
-                stiffness if stiffness is not None else default("stiffness"),
+            K=jnp.array(
+                K if K is not None else default("K"),
                 dtype=float,
             ),
-            damping=jnp.array(
-                damping if damping is not None else default("damping"), dtype=float
-            ),
+            D=jnp.array(D if D is not None else default("D"), dtype=float),
             mu=jnp.array(mu if mu is not None else default("mu"), dtype=float),
         )
 
@@ -243,6 +242,37 @@ class RelaxedRigidContacts(common.ContactModel):
             **kwargs,
         )
 
+    def update_contact_state(
+        self: type[Self], old_contact_state: dict[str, jtp.Array]
+    ) -> dict[str, jtp.Array]:
+        """
+        Update the contact state.
+
+        Args:
+            old_contact_state: The old contact state.
+
+        Returns:
+            The updated contact state.
+        """
+
+        return {}
+
+    def update_velocity_after_impact(
+        self: type[Self], model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+    ) -> js.data.JaxSimModelData:
+        """
+        Update the velocity after an impact.
+
+        Args:
+            model: The robot model considered by the contact model.
+            data: The data of the considered model.
+
+        Returns:
+            The updated data of the considered model.
+        """
+
+        return data
+
     @jax.jit
     def compute_contact_forces(
         self,
@@ -314,11 +344,11 @@ class RelaxedRigidContacts(common.ContactModel):
             BW_ν = data.generalized_velocity
 
             BW_ν̇_free = jnp.hstack(
-                js.ode.system_acceleration(
+                js.model.forward_dynamics_aba(
                     model=model,
                     data=data,
                     link_forces=references.link_forces(model=model, data=data),
-                    joint_torques=references.joint_force_references(model=model),
+                    joint_forces=references.joint_force_references(model=model),
                 )
             )
 
@@ -341,7 +371,7 @@ class RelaxedRigidContacts(common.ContactModel):
             model=model,
             position_constraint=position_constraint,
             velocity_constraint=velocity,
-            parameters=model.contacts_params,
+            parameters=model.contact_params,
         )
 
         # Compute the Delassus matrix and the free mixed linear acceleration of
@@ -425,7 +455,7 @@ class RelaxedRigidContacts(common.ContactModel):
 
         # Initialize the optimized forces with a linear Hunt/Crossley model.
         init_params = jax.vmap(
-            lambda p, v: self._hunt_crossley_contact_model(
+            lambda p, v: soft.SoftContacts.hunt_crossley_contact_model(
                 position=p,
                 velocity=v,
                 terrain=model.terrain,
@@ -501,8 +531,8 @@ class RelaxedRigidContacts(common.ContactModel):
                 "width",
                 "midpoint",
                 "power",
-                "stiffness",
-                "damping",
+                "K",
+                "D",
                 "mu",
             )
         )
@@ -598,149 +628,3 @@ class RelaxedRigidContacts(common.ContactModel):
         )
 
         return a_ref, jnp.diag(R), K, D
-
-    @staticmethod
-    @functools.partial(jax.jit, static_argnames=("terrain",))
-    def _hunt_crossley_contact_model(
-        position: jtp.VectorLike,
-        velocity: jtp.VectorLike,
-        tangential_deformation: jtp.VectorLike,
-        terrain: Terrain,
-        K: jtp.FloatLike,
-        D: jtp.FloatLike,
-        mu: jtp.FloatLike,
-        p: jtp.FloatLike = 0.5,
-        q: jtp.FloatLike = 0.5,
-    ) -> tuple[jtp.Vector, jtp.Vector]:
-        """
-        Compute the contact force using the Hunt/Crossley model.
-
-        Args:
-            position: The position of the collidable point.
-            velocity: The velocity of the collidable point.
-            tangential_deformation: The material deformation of the collidable point.
-            terrain: The terrain model.
-            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 tuple containing the computed contact force and the derivative of the
-            material deformation.
-        """
-
-        # 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()
-
-        # Use symbol for the static friction.
-        μ = mu
-
-        # Compute the penetration depth, its rate, and the considered terrain normal.
-        δ, δ̇, n̂ = common.compute_penetration_data(p=W_p_C, v=W_ṗ_C, terrain=terrain)
-
-        # 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)
-
-        # ========================
-        # Compute the normal force
-        # ========================
-
-        # 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) * δ̇
-
-        # Depending on the magnitude of δ̇, the normal force could be negative.
-        force_normal_mag = jnp.maximum(0.0, force_normal_mag)
-
-        # 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 normal and tangential components of the material deformation.
-        m_normal = jnp.dot(m, n̂) * n̂
-        m_tangential = m - jnp.dot(m, n̂) * n̂
-
-        # Compute the tangential force in the sticking case.
-        # Using the tangential component of the material deformation 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, especially for highly uneven terrains.
-        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 direction of the tangential force.
-        # To prevent dividing by zero, we use a switch statement.
-        norm = jaxsim.math.safe_norm(f_tangential, axis=-1)
-        f_tangential_direction = f_tangential / (
-            norm + jnp.finfo(float).eps * (norm == 0)
-        )
-
-        # Project the tangential force to the friction cone if slipping.
-        f_tangential = jnp.where(
-            sticking,
-            f_tangential,
-            jnp.minimum(μ * force_normal_mag, norm) * f_tangential_direction,
-        )
-
-        # Set the tangential force to zero if there is no contact.
-        f_tangential = jnp.where(δ <= 0, jnp.zeros(3), f_tangential)
-
-        # =====================================
-        # Compute the material deformation rate
-        # =====================================
-
-        # Compute the derivative of the material deformation.
-        # Note that we included an additional relaxation of `m_normal` in the
-        # sticking case, so that the normal deformation that could have accumulated
-        # from a previous slipping phase can relax to zero.
-        ṁ_no_contact = -(K / D) * m
-        ṁ_sticking = v_tangential - (K / D) * m_normal
-        ṁ_slipping = -(f_tangential + (K * δp) * m_tangential) / (D * δq)
-
-        # Compute the contact status:
-        # 0: slipping
-        # 1: sticking
-        # 2: no contact
-        contact_status = sticking.astype(int)
-        contact_status += (δ <= 0).astype(int)
-
-        # Select the right material deformation rate depending on the contact status.
-        ṁ = jax.lax.select_n(contact_status, ṁ_slipping, ṁ_sticking, ṁ_no_contact)
-
-        # ==========================================
-        # Compute and return the final contact force
-        # ==========================================
-
-        # Sum the normal and tangential forces.
-        CW_fl = f_normal + f_tangential
-
-        return CW_fl, ṁ