jaxsim 0.4.3.dev115__py3-none-any.whl → 0.4.3.dev129__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.dev115'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev115')
+__version__ = version = '0.4.3.dev129'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev129')

jaxsim/api/contact.py

@@ -9,7 +9,6 @@ import jaxsim.api as js
 import jaxsim.terrain
 import jaxsim.typing as jtp
 from jaxsim.math import Adjoint, Cross, Transform
-from jaxsim.rbda.contacts.soft import SoftContactsParams
 
 from .common import VelRepr
 
@@ -156,56 +155,43 @@ def collidable_point_dynamics(
         Instead, the 6D forces are returned in the active representation.
     """
 
-    # Compute the position and linear velocities (mixed representation) of
-    # all collidable points belonging to the robot.
-    W_p_Ci, W_ṗ_Ci = js.contact.collidable_point_kinematics(model=model, data=data)
-
     # Import privately the contacts classes.
-    from jaxsim.rbda.contacts.relaxed_rigid import (
+    from jaxsim.rbda.contacts import (
         RelaxedRigidContacts,
         RelaxedRigidContactsState,
+        RigidContacts,
+        RigidContactsState,
+        SoftContacts,
+        SoftContactsState,
     )
-    from jaxsim.rbda.contacts.rigid import RigidContacts, RigidContactsState
-    from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 
     # Build the soft contact model.
     match model.contact_model:
 
         case SoftContacts():
-
             assert isinstance(model.contact_model, SoftContacts)
             assert isinstance(data.state.contact, SoftContactsState)
 
-            # Build the contact model.
-            soft_contacts = SoftContacts(
-                parameters=data.contacts_params, terrain=model.terrain
-            )
-
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point, and the corresponding material deformation rate.
             # 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)(
-                position=W_p_Ci,
-                velocity=W_ṗ_Ci,
-                tangential_deformation=data.state.contact.tangential_deformation,
+            W_f_Ci, (CW_ṁ,) = model.contact_model.compute_contact_forces(
+                model=model, data=data
             )
+
+            # Create the dictionary of auxiliary data.
+            # This contact model considers the material deformation as additional state
+            # of the ODE system. We need to pass its dynamics to the integrator.
             aux_data = dict(m_dot=CW_ṁ)
 
         case RigidContacts():
             assert isinstance(model.contact_model, RigidContacts)
             assert isinstance(data.state.contact, RigidContactsState)
 
-            # Build the contact model.
-            rigid_contacts = RigidContacts(
-                parameters=data.contacts_params, terrain=model.terrain
-            )
-
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point.
-            W_f_Ci, _ = rigid_contacts.compute_contact_forces(
-                position=W_p_Ci,
-                velocity=W_ṗ_Ci,
+            W_f_Ci, _ = model.contact_model.compute_contact_forces(
                 model=model,
                 data=data,
                 link_forces=link_forces,
@@ -219,16 +205,9 @@ def collidable_point_dynamics(
             assert isinstance(model.contact_model, RelaxedRigidContacts)
             assert isinstance(data.state.contact, RelaxedRigidContactsState)
 
-            # Build the contact model.
-            relaxed_rigid_contacts = RelaxedRigidContacts(
-                parameters=data.contacts_params, terrain=model.terrain
-            )
-
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point.
-            W_f_Ci, _ = relaxed_rigid_contacts.compute_contact_forces(
-                position=W_p_Ci,
-                velocity=W_ṗ_Ci,
+            W_f_Ci, _ = model.contact_model.compute_contact_forces(
                 model=model,
                 data=data,
                 link_forces=link_forces,
@@ -318,7 +297,7 @@ def estimate_good_soft_contacts_parameters(
     number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
     damping_ratio: jtp.FloatLike = 1.0,
     max_penetration: jtp.FloatLike | None = None,
-) -> SoftContactsParams:
+) -> jaxsim.rbda.contacts.SoftContactsParams:
     """
     Estimate good soft contacts parameters for the given model.
 
@@ -342,14 +321,13 @@ def estimate_good_soft_contacts_parameters(
         The user is encouraged to fine-tune the parameters based on the
         specific application.
     """
-    from jaxsim.rbda.contacts.soft import SoftContactsParams
 
     def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
         """"""
 
         zero_data = js.data.JaxSimModelData.build(
             model=model,
-            contacts_params=SoftContactsParams(),
+            contacts_params=jaxsim.rbda.contacts.SoftContactsParams(),
         )
 
         W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
@@ -368,16 +346,26 @@ def estimate_good_soft_contacts_parameters(
 
     nc = number_of_active_collidable_points_steady_state
 
-    sc_parameters = SoftContactsParams.build_default_from_jaxsim_model(
-        model=model,
-        standard_gravity=standard_gravity,
-        static_friction_coefficient=static_friction_coefficient,
-        max_penetration=max_δ,
-        number_of_active_collidable_points_steady_state=nc,
-        damping_ratio=damping_ratio,
-    )
+    match model.contact_model:
+
+        case jaxsim.rbda.contacts.SoftContacts():
+            assert isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts)
+
+            parameters = (
+                jaxsim.rbda.contacts.SoftContactsParams.build_default_from_jaxsim_model(
+                    model=model,
+                    standard_gravity=standard_gravity,
+                    static_friction_coefficient=static_friction_coefficient,
+                    max_penetration=max_δ,
+                    number_of_active_collidable_points_steady_state=nc,
+                    damping_ratio=damping_ratio,
+                )
+            )
+
+        case _:
+            parameters = model.contact_model.parameters
 
-    return sc_parameters
+    return parameters
 
 
 @jax.jit

jaxsim/api/data.py

@@ -13,7 +13,7 @@ import jaxsim.api as js
 import jaxsim.math
 import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim.rbda.contacts.soft import SoftContacts
+from jaxsim.rbda.contacts import SoftContacts
 from jaxsim.utils import Mutability
 from jaxsim.utils.tracing import not_tracing
 
@@ -37,7 +37,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     gravity: jtp.Array
 
-    contacts_params: jaxsim.rbda.ContactsParams = dataclasses.field(repr=False)
+    contacts_params: jaxsim.rbda.contacts.ContactsParams = dataclasses.field(repr=False)
 
     time_ns: jtp.Int = dataclasses.field(
         default_factory=lambda: jnp.array(
@@ -114,8 +114,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         base_angular_velocity: jtp.Vector | None = None,
         joint_velocities: jtp.Vector | None = None,
         standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
-        contact: jaxsim.rbda.ContactsState | None = None,
-        contacts_params: jaxsim.rbda.ContactsParams | None = None,
+        contact: jaxsim.rbda.contacts.ContactsState | None = None,
+        contacts_params: jaxsim.rbda.contacts.ContactsParams | None = None,
         velocity_representation: VelRepr = VelRepr.Inertial,
         time: jtp.FloatLike | None = None,
     ) -> JaxSimModelData:
@@ -185,17 +185,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             )
         )
 
-        if isinstance(model.contact_model, SoftContacts):
-            contacts_params = (
-                contacts_params
-                if contacts_params is not None
-                else js.contact.estimate_good_soft_contacts_parameters(
-                    model=model, standard_gravity=standard_gravity
-                )
-            )
-        else:
-            contacts_params = model.contact_model.parameters
-
         W_H_B = jaxsim.math.Transform.from_quaternion_and_translation(
             translation=base_position, quaternion=base_quaternion
         )
@@ -225,6 +214,15 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         if not ode_state.valid(model=model):
             raise ValueError(ode_state)
 
+        if contacts_params is None:
+
+            if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
+                contacts_params = js.contact.estimate_good_soft_contacts_parameters(
+                    model=model, standard_gravity=standard_gravity
+                )
+            else:
+                contacts_params = model.contact_model.parameters
+
         return JaxSimModelData(
             time_ns=time_ns,
             state=ode_state,

jaxsim/api/model.py

@@ -36,7 +36,7 @@ class JaxSimModel(JaxsimDataclass):
         default=jaxsim.terrain.FlatTerrain(), repr=False
     )
 
-    contact_model: jaxsim.rbda.ContactModel | None = dataclasses.field(
+    contact_model: jaxsim.rbda.contacts.ContactModel | None = dataclasses.field(
         default=None, repr=False
     )
 
@@ -89,7 +89,7 @@ class JaxSimModel(JaxsimDataclass):
         model_name: str | None = None,
         *,
         terrain: jaxsim.terrain.Terrain | None = None,
-        contact_model: jaxsim.rbda.ContactModel | None = None,
+        contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
         is_urdf: bool | None = None,
         considered_joints: Sequence[str] | None = None,
     ) -> JaxSimModel:
@@ -150,7 +150,7 @@ class JaxSimModel(JaxsimDataclass):
         model_name: str | None = None,
         *,
         terrain: jaxsim.terrain.Terrain | None = None,
-        contact_model: jaxsim.rbda.ContactModel | None = None,
+        contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
     ) -> JaxSimModel:
         """
         Build a Model object from an intermediate model description.
@@ -169,14 +169,15 @@ class JaxSimModel(JaxsimDataclass):
         Returns:
             The built Model object.
         """
-        from jaxsim.rbda.contacts.soft import SoftContacts
 
         # Set the model name (if not provided, use the one from the model description).
         model_name = model_name if model_name is not None else model_description.name
 
         # Set the terrain (if not provided, use the default flat terrain).
         terrain = terrain or JaxSimModel.__dataclass_fields__["terrain"].default
-        contact_model = contact_model or SoftContacts(terrain=terrain)
+        contact_model = contact_model or jaxsim.rbda.contacts.SoftContacts(
+            terrain=terrain
+        )
 
         # Build the model.
         model = JaxSimModel(
@@ -1930,8 +1931,6 @@ def step(
         and the new state of the integrator.
     """
 
-    from jaxsim.rbda.contacts.rigid import RigidContacts
-
     # Extract the integrator kwargs.
     # The following logic allows using integrators having kwargs colliding with the
     # kwargs of this step function.
@@ -1992,12 +1991,16 @@ def step(
     # Post process the simulation state, if needed.
     match model.contact_model:
 
-        # Rigid contact models use an impact model that produces a discontinuous model velocity.
-        # Hence here we need to reset the velocity after each impact to guarantee that
+        # Rigid contact models use an impact model that produces discontinuous model velocities.
+        # Hence, here we need to reset the velocity after each impact to guarantee that
         # the linear velocity of the active collidable points is zero.
-        case RigidContacts():
-            # Raise runtime error for not supported case in which Rigid contacts and Baumgarte stabilization
-            # enabled are used with ForwardEuler integrator.
+        case jaxsim.rbda.contacts.RigidContacts():
+            assert isinstance(
+                data_tf.contacts_params, jaxsim.rbda.contacts.RigidContactsParams
+            )
+
+            # Raise runtime error for not supported case in which Rigid contacts and
+            # Baumgarte stabilization are enabled and used with ForwardEuler integrator.
             jaxsim.exceptions.raise_runtime_error_if(
                 condition=jnp.logical_and(
                     isinstance(
@@ -2013,23 +2016,38 @@ def step(
             )
 
             with data_tf.switch_velocity_representation(VelRepr.Mixed):
-                W_p_C = js.contact.collidable_point_positions(model, data_tf)
-                M = js.model.free_floating_mass_matrix(model, data_tf)
+
                 J_WC = js.contact.jacobian(model, data_tf)
+                M = js.model.free_floating_mass_matrix(model, data_tf)
+                W_p_C = js.contact.collidable_point_positions(model, data_tf)
+
+                # Compute the height of the terrain below each collidable point.
                 px, py, _ = W_p_C.T
                 terrain_height = jax.vmap(model.terrain.height)(px, py)
-                inactive_collidable_points, _ = RigidContacts.detect_contacts(
-                    W_p_C=W_p_C,
-                    terrain_height=terrain_height,
+
+                # Compute the contact state.
+                inactive_collidable_points, _ = (
+                    jaxsim.rbda.contacts.RigidContacts.detect_contacts(
+                        W_p_C=W_p_C,
+                        terrain_height=terrain_height,
+                    )
                 )
-                BW_nu_post_impact = RigidContacts.compute_impact_velocity(
-                    data=data_tf,
-                    inactive_collidable_points=inactive_collidable_points,
-                    M=M,
-                    J_WC=J_WC,
+
+                # Compute the impact velocity.
+                # It may be discontinuous in case new contacts are made.
+                BW_nu_post_impact = (
+                    jaxsim.rbda.contacts.RigidContacts.compute_impact_velocity(
+                        data=data_tf,
+                        inactive_collidable_points=inactive_collidable_points,
+                        M=M,
+                        J_WC=J_WC,
+                    )
                 )
+
+                # Reset the generalized velocity.
                 data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
                 data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
+
             # Restore the input velocity representation.
             data_tf = data_tf.replace(
                 velocity_representation=data.velocity_representation, validate=False

jaxsim/api/ode_data.py

@@ -5,13 +5,15 @@ import jax_dataclasses
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.rbda import ContactsState
-from jaxsim.rbda.contacts.relaxed_rigid import (
+from jaxsim.rbda.contacts import (
+    ContactsState,
     RelaxedRigidContacts,
     RelaxedRigidContactsState,
+    RigidContacts,
+    RigidContactsState,
+    SoftContacts,
+    SoftContactsState,
 )
-from jaxsim.rbda.contacts.rigid import RigidContacts, RigidContactsState
-from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 from jaxsim.utils import JaxsimDataclass
 
 # =============================================================================
@@ -165,8 +167,11 @@ class ODEState(JaxsimDataclass):
 
         # Get the contact model from the `JaxSimModel`.
         match model.contact_model:
+
             case SoftContacts():
+
                 tangential_deformation = kwargs.get("tangential_deformation", None)
+
                 contact = SoftContactsState.build_from_jaxsim_model(
                     model=model,
                     **(
@@ -182,7 +187,7 @@ class ODEState(JaxsimDataclass):
                 contact = RelaxedRigidContactsState.build()
 
             case _:
-                raise ValueError("Unable to determine contact state class prefix.")
+                raise ValueError("Unsupported contact model.")
 
         return ODEState.build(
             model=model,

jaxsim/rbda/__init__.py

@@ -1,6 +1,6 @@
+from . import contacts
 from .aba import aba
 from .collidable_points import collidable_points_pos_vel
-from .contacts.common import ContactModel, ContactsParams, ContactsState
 from .crba import crba
 from .forward_kinematics import forward_kinematics, forward_kinematics_model
 from .jacobian import (

jaxsim/rbda/contacts/__init__.py

@@ -1,3 +1,4 @@
+from . import relaxed_rigid, rigid, soft
 from .common import ContactModel, ContactsParams, ContactsState
 from .relaxed_rigid import (
     RelaxedRigidContacts,

jaxsim/rbda/contacts/common.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import abc
 from typing import Any
 
+import jaxsim.api as js
 import jaxsim.terrain
 import jaxsim.typing as jtp
 from jaxsim.utils import JaxsimDataclass
@@ -90,20 +91,48 @@ class ContactModel(JaxsimDataclass):
     @abc.abstractmethod
     def compute_contact_forces(
         self,
-        position: jtp.VectorLike,
-        velocity: jtp.VectorLike,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
         **kwargs,
     ) -> tuple[jtp.Vector, tuple[Any, ...]]:
         """
         Compute the contact forces.
 
         Args:
-            position: The position of the collidable point w.r.t. the world frame.
-            velocity:
-                The linear velocity of the collidable point (linear component of the mixed 6D velocity).
+            model: The model to consider.
+            data: The data of the considered model.
 
         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.
         """
+
         pass
+
+    def initialize_model_and_data(
+        self,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        validate: bool = True,
+    ) -> tuple[js.model.JaxSimModel, js.data.JaxSimModelData]:
+        """
+        Helper function to initialize the active model and data objects.
+
+        Args:
+            model: The robot model considered by the contact model.
+            data: The data of the considered robot model.
+            validate:
+                Whether to validate if the model and data objects have been
+                initialized with the current contact model.
+
+        Returns:
+            The initialized model and data objects.
+        """
+
+        with model.editable(validate=validate) as model_out:
+            model_out.contact_model = self
+
+        with data.editable(validate=validate) as data_out:
+            data_out.contacts_params = data.contacts_params
+
+        return model_out, data_out

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -169,12 +169,12 @@ class RelaxedRigidContactsState(ContactsState):
         return cls()
 
     @classmethod
-    def zero(cls: type[Self]) -> Self:
+    def zero(cls: type[Self], **kwargs) -> Self:
         """Build a zero `RelaxedRigidContactsState` instance from a `JaxSimModel`."""
 
         return cls.build()
 
-    def valid(self, *, model: js.model.JaxSimModel) -> jtp.BoolLike:
+    def valid(self, **kwargs) -> jtp.BoolLike:
         return True
 
 
@@ -193,11 +193,9 @@ class RelaxedRigidContacts(ContactModel):
     @jax.jit
     def compute_contact_forces(
         self,
-        position: jtp.VectorLike,
-        velocity: jtp.VectorLike,
-        *,
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
+        *,
         link_forces: jtp.MatrixLike | None = None,
         joint_force_references: jtp.VectorLike | None = None,
     ) -> tuple[jtp.Vector, tuple[Any, ...]]:
@@ -205,10 +203,8 @@ class RelaxedRigidContacts(ContactModel):
         Compute the contact forces.
 
         Args:
-            position: The position of the collidable point.
-            velocity: The linear velocity of the collidable point.
-            model: The `JaxSimModel` instance.
-            data: The `JaxSimModelData` instance.
+            model: The model to consider.
+            data: The data of the considered model.
             link_forces:
                 Optional `(n_links, 6)` matrix of external forces acting on the links,
                 expressed in the same representation of data.
@@ -219,6 +215,11 @@ class RelaxedRigidContacts(ContactModel):
             A tuple containing the contact forces.
         """
 
+        # 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)
+
         link_forces = (
             link_forces
             if link_forces is not None
@@ -247,6 +248,12 @@ class RelaxedRigidContacts(ContactModel):
 
             return jnp.dot(h, n̂)
 
+        # Compute the position and linear velocities (mixed representation) of
+        # all collidable points belonging to the robot.
+        position, velocity = js.contact.collidable_point_kinematics(
+            model=model, data=data
+        )
+
         # Compute the activation state of the collidable points
         δ = jax.vmap(_detect_contact)(*position.T)
 

jaxsim/rbda/contacts/rigid.py

@@ -92,12 +92,12 @@ class RigidContactsState(ContactsState):
         return cls()
 
     @classmethod
-    def zero(cls: type[Self]) -> Self:
+    def zero(cls: type[Self], **kwargs) -> Self:
         """Build a zero `RigidContactsState` instance from a `JaxSimModel`."""
 
         return cls.build()
 
-    def valid(self) -> jtp.BoolLike:
+    def valid(self, **kwargs) -> jtp.BoolLike:
         return True
 
 
@@ -219,11 +219,9 @@ class RigidContacts(ContactModel):
     @jax.jit
     def compute_contact_forces(
         self,
-        position: jtp.VectorLike,
-        velocity: jtp.VectorLike,
-        *,
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
+        *,
         link_forces: jtp.MatrixLike | None = None,
         joint_force_references: jtp.VectorLike | None = None,
         regularization_term: jtp.FloatLike = 1e-6,
@@ -233,10 +231,8 @@ class RigidContacts(ContactModel):
         Compute the contact forces.
 
         Args:
-            position: The position of the collidable point.
-            velocity: The linear velocity of the collidable point.
-            model: The `JaxSimModel` instance.
-            data: The `JaxSimModelData` instance.
+            model: The model to consider.
+            data: The data of the considered model.
             link_forces:
                 Optional `(n_links, 6)` matrix of external forces acting on the links,
                 expressed in the same representation of data.
@@ -245,11 +241,17 @@ class RigidContacts(ContactModel):
             regularization_term:
                 The regularization term to add to the diagonal of the Delassus
                 matrix for better numerical conditioning.
+            solver_tol: The convergence tolerance to consider in the QP solver.
 
         Returns:
             A tuple containing the contact forces.
         """
 
+        # 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)
+
         # Import qpax just in this method
         import qpax
 
@@ -273,6 +275,12 @@ class RigidContacts(ContactModel):
             J̇_WC_BW = js.contact.jacobian_derivative(model=model, data=data)
             BW_ν = data.generalized_velocity()
 
+        # Compute the position and linear velocities (mixed representation) of
+        # all collidable points belonging to the robot.
+        position, velocity = js.contact.collidable_point_kinematics(
+            model=model, data=data
+        )
+
         terrain_height = jax.vmap(self.terrain.height)(position[:, 0], position[:, 1])
         n_collidable_points = model.kin_dyn_parameters.contact_parameters.point.shape[0]
 

jaxsim/rbda/contacts/soft.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
+import functools
 
 import jax
 import jax.numpy as jnp
@@ -198,30 +199,32 @@ class SoftContacts(ContactModel):
         default_factory=FlatTerrain
     )
 
-    def compute_contact_forces(
-        self,
+    @staticmethod
+    @functools.partial(jax.jit, static_argnames=("terrain",))
+    def hunt_crossley_contact_model(
         position: jtp.VectorLike,
         velocity: jtp.VectorLike,
-        *,
         tangential_deformation: jtp.VectorLike,
-    ) -> tuple[jtp.Vector, tuple[jtp.Vector]]:
+        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]:
 
         # 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.
-        K = self.parameters.K
-        D = self.parameters.D
-        μ = self.parameters.mu
+        # Use symbol for the static friction.
+        μ = mu
 
         # 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)
-
-        # Get the exponents of the Hunt/Crossley model non-linear terms.
-        p = self.parameters.p
-        q = self.parameters.q
+        δ, δ̇, n̂ = SoftContacts.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.
@@ -256,14 +259,15 @@ class SoftContacts(ContactModel):
         # 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.
+        # 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).
@@ -298,6 +302,9 @@ class SoftContacts(ContactModel):
         # =====================================
 
         # 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)
@@ -316,15 +323,79 @@ class SoftContacts(ContactModel):
         # 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)])
+        # Sum the normal and tangential forces.
+        CW_fl = f_normal + f_tangential
+
+        return CW_fl, ṁ
+
+    @staticmethod
+    @functools.partial(jax.jit, static_argnames=("terrain",))
+    def compute_contact_force(
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
+        tangential_deformation: jtp.VectorLike,
+        parameters: SoftContactsParams,
+        terrain: Terrain,
+    ) -> tuple[jtp.Vector, jtp.Vector]:
+
+        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.mu,
+            p=parameters.p,
+            q=parameters.q,
+        )
+
+        # Pack a mixed 6D force.
+        CW_f = jnp.hstack([CW_fl, 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
+            translation=jnp.array(position), inverse=True
         ).T
 
-        return W_Xf_CW @ CW_f, (ṁ,)
+        # Compute the 6D force in the inertial-fixed frame.
+        W_f = W_Xf_CW @ CW_f
+
+        return W_f, ṁ
+
+    @jax.jit
+    def compute_contact_forces(
+        self,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+    ) -> tuple[jtp.Vector, tuple[jtp.Vector]]:
+
+        # 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)
+
+        # Compute the position and linear velocities (mixed representation) of
+        # all collidable points belonging to the robot.
+        W_p_C, W_ṗ_C = js.contact.collidable_point_kinematics(model=model, data=data)
+
+        # Extract the material deformation corresponding to the collidable points.
+        assert isinstance(data.state.contact, SoftContactsState)
+        m = data.state.contact.tangential_deformation
+
+        # Compute the contact forces for all collidable points.
+        # Since we treat them as independent, we can vmap the computation.
+        W_f, ṁ = jax.vmap(
+            lambda p, v, m: SoftContacts.compute_contact_force(
+                position=p,
+                velocity=v,
+                tangential_deformation=m,
+                parameters=self.parameters,
+                terrain=self.terrain,
+            )
+        )(W_p_C, W_ṗ_C, m)
+
+        return W_f, (ṁ,)
 
     @staticmethod
     @jax.jit

jaxsim/terrain/terrain.py

@@ -155,7 +155,7 @@ class PlaneTerrain(FlatTerrain):
             (
                 hash(self._height),
                 HashedNumpyArray.hash_of_array(
-                    array=jnp.array(self._normal, dtype=float)
+                    array=np.array(self._normal, dtype=float)
                 ),
             )
         )

{jaxsim-0.4.3.dev115.dist-info → jaxsim-0.4.3.dev129.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev115
+Version: 0.4.3.dev129
 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.dev115.dist-info → jaxsim-0.4.3.dev129.dist-info}/RECORD

@@ -1,20 +1,20 @@
 jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=5euYiJAuz4T6MQwktcrK-oclprll5Dr6ScaKSKchsa8,428
+jaxsim/_version.py,sha256=TZYv9WjsK9pTtCnxiHUO_BF5LO5oh7sVdFFBwKnF98k,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=ocwsVS1jaBfrd81990hcgfS0-2xD8VVzDq7gdPguAUg,23087
-jaxsim/api/data.py,sha256=QldUHniJqKrdNtAcXuRaS9UyeslJ0Rjvb17UA0Ca5Tw,29008
+jaxsim/api/contact.py,sha256=ktol-c2UHDqt3Hd6NpMhK86NlyadIyB5eY7chA5GhkY,22568
+jaxsim/api/data.py,sha256=bSuBVKKksYtANLnxCkMQ3u2JGbn5mud7g8G5aLZawls,28988
 jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
 jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
 jaxsim/api/kin_dyn_parameters.py,sha256=ElahFk_RCcLvjTidH2qDOsY-m1gN1hXitCv4SvfgGYY,29260
 jaxsim/api/link.py,sha256=LAA6ZMQXkWomXeptURBtc7z3_xDZ2BBnBMhVrohh0bE,18621
-jaxsim/api/model.py,sha256=FLCk3fKUfLp3eCPrFEgPljTDp3zgIJnVW5Gwx0fWAog,67637
+jaxsim/api/model.py,sha256=hDb8lVoRMhAWOMn1Q-MFHAa4wJY431-_nuLTKY0E3-Q,68200
 jaxsim/api/ode.py,sha256=gYSbtHWGCDP-IkUzQlH3t0fBKnK8qmxwhIvsbLG9lwU,13616
-jaxsim/api/ode_data.py,sha256=7RSoBhfCJdP6P9InQbDwdBVpClPMMuetewI-6AWm-_0,20276
+jaxsim/api/ode_data.py,sha256=k1hVU1x8vuTVYdkf0cLhZ-oqeGJocXN2lLey-pS1_vo,20166
 jaxsim/api/references.py,sha256=XOVKuQXRmjPoP-T5JWGSbqIGX5DzOkeGafqRpj0ZQEM,20771
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
 jaxsim/integrators/common.py,sha256=78MBs89GxsL0wU2yAexjvBZt3HEtfZoGVIN9f0a8yTc,20305
@@ -44,7 +44,7 @@ jaxsim/parsers/descriptions/model.py,sha256=I2Vsbv8Josl4Le7b5rIvhqA2k9Bbv5JxMqwy
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
 jaxsim/parsers/rod/parser.py,sha256=9zBYTQF2vC4NO6HEZBvV8VTaGuZSdbd3v66BAcMpuVI,13923
 jaxsim/parsers/rod/utils.py,sha256=5DsF3OeePZGidOJ5GiFSZx-51uIdnFvMW9EK6SgOW6Q,5698
-jaxsim/rbda/__init__.py,sha256=H7DhXpxkPOi9lpUvg31IMHFfRafke1UoJLc5GQIdyhA,387
+jaxsim/rbda/__init__.py,sha256=kmy4G9aMkrqPNGdLSaSV3k15dpF52vBEUQXDFDuKIxU,337
 jaxsim/rbda/aba.py,sha256=w7ciyxB0IsmueatT0C7PcBQEl9dyiH9oqJgIi3xeTUE,8983
 jaxsim/rbda/collidable_points.py,sha256=Rmf1DhflhOTYh9mDalv0agS0CGSbmfoOybwP2KzKuJ0,4883
 jaxsim/rbda/crba.py,sha256=bXkXESnVbv-lxhU1Y_i0rViEcQA4z2t2_jHwdVj5CBo,5049
@@ -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=Y1yT2zdgFa0zviZseI09wNaMcydH8TeoaWr6ehqzwdc,328
-jaxsim/rbda/contacts/common.py,sha256=CEmLS_PT44AOWKJ0bWrJJBqm2Q9v9LiqvL0rht63-ic,2605
-jaxsim/rbda/contacts/relaxed_rigid.py,sha256=f4g_qjoiml1gU9mqy6gEb8ljI3OSr4aoD2Li0J8dqQs,13708
-jaxsim/rbda/contacts/rigid.py,sha256=SP9qwkf35ybK0ZMiNhfIrP_tiiFN6jiG0ewBSBtR2Mc,15204
-jaxsim/rbda/contacts/soft.py,sha256=-d7zbMdKNq0aRT2zRXIu_Dbh8BL4VUnMDprz4Ddfwj0,16276
+jaxsim/rbda/contacts/__init__.py,sha256=0UnO9ZR3BwdjQa276jOFbPi90pporr32LSc0qa9UUm4,369
+jaxsim/rbda/contacts/common.py,sha256=-eM8d1kvJ2E_2_kAgZJk4s3x8vDZHNSyOAinwPmRmEk,3469
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=VBD1FPjwpuKTfQ1bFKDlD_4xlEox4QhiatsjcrjDMNk,14025
+jaxsim/rbda/contacts/rigid.py,sha256=6cU8kM8LMjEFbt8dtSg5nnz_uh4aD50sKw_svCzYUms,15633
+jaxsim/rbda/contacts/soft.py,sha256=NzzCYw5rvK8Fx_qH3fiMzPgey-KoxmRe9xkF3fluidE,18866
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
-jaxsim/terrain/terrain.py,sha256=xUQg47yGxIOcTkLPbnO3sruEGBhoCd16j1evTGlmNjI,5010
+jaxsim/terrain/terrain.py,sha256=Y0TGnUAGPuaeeSN8vbaSFjMXP5GYy3nxMCETjpUIMSA,5009
 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.dev115.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev115.dist-info/METADATA,sha256=LG0pWBHUEVPV3MGVXaApN8SdiEUrmKTgJuD7PQmmyhI,17277
-jaxsim-0.4.3.dev115.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev115.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev115.dist-info/RECORD,,
+jaxsim-0.4.3.dev129.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev129.dist-info/METADATA,sha256=oWWfWqv_MyqMsPtAZmGGEFc8qm_cuzzxTRGa7J9Vdhg,17277
+jaxsim-0.4.3.dev129.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev129.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev129.dist-info/RECORD,,