jaxsim 0.6.2.dev182__py3-none-any.whl → 0.6.2.dev225__py3-none-any.whl

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,
@@ -460,16 +490,12 @@ class RelaxedRigidContacts(common.ContactModel):
         CW_fl_C = solution.reshape(-1, 3)
 
         # Convert the contact forces from mixed to inertial-fixed representation.
-        W_f_C = jax.vmap(
-            lambda CW_fl_C, W_H_C: (
-                ModelDataWithVelocityRepresentation.other_representation_to_inertial(
-                    array=jnp.zeros(6).at[0:3].set(CW_fl_C),
-                    transform=W_H_C,
-                    other_representation=VelRepr.Mixed,
-                    is_force=True,
-                )
-            ),
-        )(CW_fl_C, W_H_C)
+        W_f_C = ModelDataWithVelocityRepresentation.other_representation_to_inertial(
+            array=jnp.zeros((W_H_C.shape[0], 6)).at[:, :3].set(CW_fl_C),
+            transform=W_H_C,
+            other_representation=VelRepr.Mixed,
+            is_force=True,
+        )
 
         return W_f_C, {}
 
@@ -505,8 +531,8 @@ class RelaxedRigidContacts(common.ContactModel):
                 "width",
                 "midpoint",
                 "power",
-                "stiffness",
-                "damping",
+                "K",
+                "D",
                 "mu",
             )
         )
@@ -602,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)
-        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, ṁ