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

jaxsim/api/integrators.py

@@ -22,7 +22,7 @@ def semi_implicit_euler_integration(
     with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
 
         # Compute the system acceleration
-        W_v̇_WB, s̈ = js.ode.system_acceleration(
+        W_v̇_WB, s̈, contact_state_derivative = js.ode.system_acceleration(
             model=model,
             data=data,
             link_forces=link_forces,
@@ -58,12 +58,18 @@ def semi_implicit_euler_integration(
         W_p_B = data.base_position + dt * W_ṗ_B
         W_Q_B = data.base_orientation + dt * W_Q̇_B
 
-        base_quaternion_norm = jaxsim.math.safe_norm(W_Q_B)
+        base_quaternion_norm = jaxsim.math.safe_norm(W_Q_B, axis=-1)
 
         W_Q_B = W_Q_B / jnp.where(base_quaternion_norm == 0, 1.0, base_quaternion_norm)
 
         s = data.joint_positions + dt * ṡ
 
+        integrated_contact_state = jax.tree.map(
+            lambda x, x_dot: x + dt * x_dot,
+            data.contact_state,
+            contact_state_derivative,
+        )
+
     # TODO: Avoid double replace, e.g. by computing cached value here
     data = dataclasses.replace(
         data,
@@ -73,6 +79,7 @@ def semi_implicit_euler_integration(
         _joint_velocities=ṡ,
         _base_linear_velocity=W_v_B[0:3],
         _base_angular_velocity=W_ω_WB,
+        contact_state=integrated_contact_state,
     )
 
     # Update the cached computations.
@@ -104,7 +111,7 @@ def rk4_integration(
                 joint_torques=joint_torques,
             )
 
-    base_quaternion_norm = jaxsim.math.safe_norm(data._base_quaternion)
+    base_quaternion_norm = jaxsim.math.safe_norm(data._base_quaternion, axis=-1)
     base_quaternion = data._base_quaternion / jnp.where(
         base_quaternion_norm == 0, 1.0, base_quaternion_norm
     )
@@ -116,6 +123,7 @@ def rk4_integration(
         base_linear_velocity=data._base_linear_velocity,
         base_angular_velocity=data._base_angular_velocity,
         joint_velocities=data._joint_velocities,
+        contact_state=data.contact_state,
     )
 
     euler_mid = lambda x, dxdt: x + (0.5 * dt) * dxdt
@@ -136,14 +144,13 @@ def rk4_integration(
 
     data_tf = dataclasses.replace(
         data,
-        **{
-            "_base_position": x_tf["base_position"],
-            "_base_quaternion": x_tf["base_quaternion"],
-            "_joint_positions": x_tf["joint_positions"],
-            "_base_linear_velocity": x_tf["base_linear_velocity"],
-            "_base_angular_velocity": x_tf["base_angular_velocity"],
-            "_joint_velocities": x_tf["joint_velocities"],
-        },
+        _base_position=x_tf["base_position"],
+        _base_quaternion=x_tf["base_quaternion"],
+        _joint_positions=x_tf["joint_positions"],
+        _base_linear_velocity=x_tf["base_linear_velocity"],
+        _base_angular_velocity=x_tf["base_angular_velocity"],
+        _joint_velocities=x_tf["joint_velocities"],
+        contact_state=x_tf["contact_state"],
     )
 
     return data_tf.replace(model=model)

jaxsim/api/model.py

@@ -47,13 +47,13 @@ class JaxSimModel(JaxsimDataclass):
         default_factory=jaxsim.terrain.FlatTerrain.build, repr=False
     )
 
-    gravity: Static[float] = jaxsim.math.STANDARD_GRAVITY
+    gravity: Static[float] = -jaxsim.math.STANDARD_GRAVITY
 
     contact_model: Static[jaxsim.rbda.contacts.ContactModel | None] = dataclasses.field(
         default=None, repr=False
     )
 
-    contacts_params: Static[jaxsim.rbda.contacts.ContactsParams] = dataclasses.field(
+    contact_params: Static[jaxsim.rbda.contacts.ContactsParams] = dataclasses.field(
         default=None, repr=False
     )
 
@@ -177,9 +177,9 @@ class JaxSimModel(JaxsimDataclass):
             time_step=time_step,
             terrain=terrain,
             contact_model=contact_model,
-            contacts_params=contact_params,
+            contact_params=contact_params,
             integrator=integrator,
-            gravity=gravity,
+            gravity=-gravity,
         )
 
         # Store the origin of the model, in case downstream logic needs it.
@@ -197,7 +197,7 @@ class JaxSimModel(JaxsimDataclass):
         time_step: jtp.FloatLike | None = None,
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
-        contacts_params: jaxsim.rbda.contacts.ContactsParams | None = None,
+        contact_params: jaxsim.rbda.contacts.ContactsParams | None = None,
         integrator: IntegratorType | None = None,
         gravity: jtp.FloatLike = jaxsim.math.STANDARD_GRAVITY,
     ) -> JaxSimModel:
@@ -217,8 +217,8 @@ class JaxSimModel(JaxsimDataclass):
                 The optional name of the model overriding the physics model name.
             contact_model:
                 The contact model to consider.
-                If not specified, a soft contacts model is used.
-            contacts_params: The parameters of the soft contacts.
+                If not specified, a relaxed-constraints rigid contacts model is used.
+            contact_params: The parameters of the contact model.
             integrator: The integrator to use for the simulation.
             gravity: The gravity constant.
 
@@ -252,8 +252,8 @@ class JaxSimModel(JaxsimDataclass):
             else jaxsim.rbda.contacts.RelaxedRigidContacts.build()
         )
 
-        if contacts_params is None:
-            contacts_params = contact_model._parameters_class()
+        if contact_params is None:
+            contact_params = contact_model._parameters_class()
 
         # Consider the default integrator if not specified.
         integrator = (
@@ -271,7 +271,7 @@ class JaxSimModel(JaxsimDataclass):
             time_step=time_step,
             terrain=terrain,
             contact_model=contact_model,
-            contacts_params=contacts_params,
+            contact_params=contact_params,
             integrator=integrator,
             gravity=gravity,
             # The following is wrapped as hashless since it's a static argument, and we
@@ -473,7 +473,7 @@ def reduce(
         time_step=model.time_step,
         terrain=model.terrain,
         contact_model=model.contact_model,
-        contacts_params=model.contacts_params,
+        contact_params=model.contact_params,
         gravity=model.gravity,
         integrator=model.integrator,
     )
@@ -2069,14 +2069,12 @@ def step(
     )
 
     # Get the external forces in inertial-fixed representation.
-    W_f_L_external = jax.vmap(
-        lambda f_L, W_H_L: js.data.JaxSimModelData.other_representation_to_inertial(
-            f_L,
-            other_representation=data.velocity_representation,
-            transform=W_H_L,
-            is_force=True,
-        )
-    )(O_f_L_external, data._link_transforms)
+    W_f_L_external = js.data.JaxSimModelData.other_representation_to_inertial(
+        O_f_L_external,
+        other_representation=data.velocity_representation,
+        transform=data._link_transforms,
+        is_force=True,
+    )
 
     τ_references = jnp.atleast_1d(
         jnp.array(joint_force_references, dtype=float).squeeze()
@@ -2092,29 +2090,6 @@ def step(
         model, data, joint_force_references=τ_references
     )
 
-    # ======================
-    # Compute contact forces
-    # ======================
-
-    W_f_L_terrain = jnp.zeros_like(W_f_L_external)
-
-    if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
-
-        # Compute the 6D forces W_f ∈ ℝ^{n_L × 6} applied to links due to contact
-        # with the terrain.
-        W_f_L_terrain = js.contact_model.link_contact_forces(
-            model=model,
-            data=data,
-            link_forces=W_f_L_external,
-            joint_torques=τ_total,
-        )
-
-    # ==============================
-    # Compute the total link forces
-    # ==============================
-
-    W_f_L_total = W_f_L_external + W_f_L_terrain
-
     # =============================
     # Advance the simulation state
     # =============================
@@ -2124,7 +2099,14 @@ def step(
     integrator_fn = _INTEGRATORS_MAP[model.integrator]
 
     data_tf = integrator_fn(
-        model=model, data=data, link_forces=W_f_L_total, joint_torques=τ_total
+        model=model,
+        data=data,
+        link_forces=W_f_L_external,
+        joint_torques=τ_total,
+    )
+
+    data_tf = model.contact_model.update_velocity_after_impact(
+        model=model, data=data_tf
     )
 
     return data_tf

jaxsim/api/ode.py

@@ -46,12 +46,36 @@ def system_acceleration(
         else jnp.zeros((model.number_of_links(), 6))
     ).astype(float)
 
+    # ======================
+    # Compute contact forces
+    # ======================
+
+    W_f_L_terrain = jnp.zeros_like(f_L)
+    contact_state_derivative = {}
+
+    if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
+
+        # Compute the 6D forces W_f ∈ ℝ^{n_L × 6} applied to links due to contact
+        # with the terrain.
+        W_f_L_terrain, contact_state_derivative = js.contact.link_contact_forces(
+            model=model,
+            data=data,
+            link_forces=f_L,
+            joint_torques=joint_torques,
+        )
+
+    W_f_L_total = f_L + W_f_L_terrain
+
+    # Update the contact state data. This is necessary only for the contact models
+    # that require propagation and integration of contact state.
+    contact_state = model.contact_model.update_contact_state(contact_state_derivative)
+
     # Store the link forces in a references object.
     references = js.references.JaxSimModelReferences.build(
         model=model,
         data=data,
         velocity_representation=data.velocity_representation,
-        link_forces=f_L,
+        link_forces=W_f_L_total,
     )
 
     # Compute forward dynamics.
@@ -68,13 +92,12 @@ def system_acceleration(
         link_forces=references.link_forces(model=model, data=data),
     )
 
-    return v̇_WB, s̈
+    return v̇_WB, s̈, contact_state
 
 
 @jax.jit
 @js.common.named_scope
 def system_position_dynamics(
-    model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[jtp.Vector, jtp.Vector, jtp.Vector]:
@@ -82,7 +105,6 @@ def system_position_dynamics(
     Compute the dynamics of the system position.
 
     Args:
-        model: The model to consider.
         data: The data of the considered model.
         baumgarte_quaternion_regularization:
             The Baumgarte regularization coefficient for adjusting the quaternion norm.
@@ -144,7 +166,7 @@ def system_dynamics(
     """
 
     with data.switch_velocity_representation(velocity_representation=VelRepr.Inertial):
-        W_v̇_WB, s̈ = system_acceleration(
+        W_v̇_WB, s̈, contact_state_derivative = system_acceleration(
             model=model,
             data=data,
             joint_torques=joint_torques,
@@ -152,7 +174,6 @@ def system_dynamics(
         )
 
         W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
-            model=model,
             data=data,
             baumgarte_quaternion_regularization=baumgarte_quaternion_regularization,
         )
@@ -164,4 +185,5 @@ def system_dynamics(
         base_linear_velocity=W_v̇_WB[0:3],
         base_angular_velocity=W_v̇_WB[3:6],
         joint_velocities=s̈,
+        contact_state=contact_state_derivative,
     )

jaxsim/api/references.py

@@ -434,24 +434,17 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         if not_tracing(forces) and not data.valid(model=model):
             raise ValueError("The provided data is not valid for the model")
 
-        # Helper function to convert a single 6D force to the inertial representation
-        # considering as body the link (i.e. L_f_L and LW_f_L).
-        def convert_using_link_frame(
-            f_L: jtp.MatrixLike, W_H_L: jtp.ArrayLike
-        ) -> jtp.Matrix:
-
-            return jax.vmap(
-                lambda f_L, W_H_L: JaxSimModelReferences.other_representation_to_inertial(
-                    array=f_L,
-                    other_representation=self.velocity_representation,
-                    transform=W_H_L,
-                    is_force=True,
-                )
-            )(f_L, W_H_L)
+        W_H_L = data._link_transforms
 
+        # Convert a single 6D force to the inertial representation
+        # considering as body the link (i.e. L_f_L and LW_f_L).
         # The f_L input is either L_f_L or LW_f_L, depending on the representation.
-        W_H_L = data._link_transforms
-        W_f_L = convert_using_link_frame(f_L=f_L, W_H_L=W_H_L[link_idxs, :, :])
+        W_f_L = JaxSimModelReferences.other_representation_to_inertial(
+            array=f_L,
+            other_representation=self.velocity_representation,
+            transform=W_H_L[link_idxs] if model.number_of_links() > 1 else W_H_L,
+            is_force=True,
+        )
 
         return replace(forces=self._link_forces.at[link_idxs, :].set(W_f0_L + W_f_L))
 

jaxsim/math/__init__.py

@@ -11,4 +11,4 @@ from .joint_model import JointModel, supported_joint_motion  # isort:skip
 
 
 # Define the default standard gravity constant.
-STANDARD_GRAVITY = -9.81
+STANDARD_GRAVITY = 9.81

jaxsim/math/adjoint.py

@@ -55,13 +55,13 @@ class Adjoint:
             The 6x6 adjoint matrix.
         """
 
-        A_H_B = jnp.reshape(transform, (-1, 4, 4))
+        A_H_B = transform
 
         return (
             jaxlie.SE3.from_matrix(matrix=A_H_B).adjoint()
             if not inverse
             else jaxlie.SE3.from_matrix(matrix=A_H_B).inverse().adjoint()
-        ).reshape(transform.shape[:-2] + (6, 6))
+        )
 
     @staticmethod
     def from_rotation_and_translation(

jaxsim/math/transform.py

@@ -36,8 +36,8 @@ class Transform:
         W_Q_B = jnp.array(quaternion).astype(float)
         W_p_B = jnp.array(translation).astype(float)
 
-        assert W_p_B.size == 3
-        assert W_Q_B.size == 4
+        assert W_p_B.shape[-1] == 3
+        assert W_Q_B.shape[-1] == 4
 
         A_R_B = jaxlie.SO3(wxyz=W_Q_B)
         A_R_B = A_R_B if not normalize_quaternion else A_R_B.normalize()

jaxsim/math/utils.py

@@ -3,7 +3,7 @@ import jax.numpy as jnp
 import jaxsim.typing as jtp
 
 
-def safe_norm(array: jtp.ArrayLike, axis=None) -> jtp.Array:
+def safe_norm(array: jtp.ArrayLike, *, axis=None, keepdims: bool = False) -> jtp.Array:
     """
     Compute an array norm handling NaNs and making sure that
     it is safe to get the gradient.
@@ -11,6 +11,7 @@ def safe_norm(array: jtp.ArrayLike, axis=None) -> jtp.Array:
     Args:
         array: The array for which to compute the norm.
         axis: The axis for which to compute the norm.
+        keepdims: Whether to keep the dimensions of the input
 
     Returns:
         The norm of the array with handling for zero arrays to avoid NaNs.
@@ -24,7 +25,7 @@ def safe_norm(array: jtp.ArrayLike, axis=None) -> jtp.Array:
     array = jnp.where(is_zero, jnp.ones_like(array), array)
 
     # Compute the norm of the array along the specified axis.
-    norm = jnp.linalg.norm(array, axis=axis)
+    norm = jnp.linalg.norm(array, axis=axis, keepdims=keepdims)
 
     # Use `jnp.where` to set the norm to 0.0 where the input array was all zeros.
     # This usage supports potential batch processing for future scalability.

jaxsim/mujoco/visualizer.py

@@ -64,7 +64,7 @@ class MujocoVideoRecorder:
 
         self.frames = []
 
-        self.data = [data] if data is not None else self.data
+        self.data = data if data is not None else self.data
         self.data = self.data if isinstance(self.data, list) else [self.data]
 
         self.model = model if model is not None else self.model

jaxsim/parsers/kinematic_graph.py

@@ -973,7 +973,7 @@ class KinematicGraphTransforms:
 
         if frame.parent_name in self.graph.links_dict:
             return frame.parent_name
-        elif frame.parent_name in self.graph.frames_dict:
+        if frame.parent_name in self.graph.frames_dict:
             return self.find_parent_link_of_frame(name=frame.parent_name)
 
         msg = f"Failed to find parent element of frame '{name}' with name '{frame.parent_name}'"

jaxsim/rbda/__init__.py

@@ -2,7 +2,7 @@ from . import contacts
 from .aba import aba
 from .collidable_points import collidable_points_pos_vel
 from .crba import crba
-from .forward_kinematics import forward_kinematics, forward_kinematics_model
+from .forward_kinematics import forward_kinematics_model
 from .jacobian import (
     jacobian,
     jacobian_derivative_full_doubly_left,

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.
+        """