jaxsim 0.6.2.dev194__tar.gz → 0.6.2.dev225__tar.gz

{jaxsim-0.6.2.dev194 → jaxsim-0.6.2.dev225}/.github/workflows/gpu_benchmark.yml

@@ -51,7 +51,7 @@ jobs:
 
       - name: Run benchmark and store result
         run: |
-            pytest tests/test_benchmark.py -k 'not test_rigid_contact_model and not test_soft_contact_model' --gpu-only --batch-size 128 --benchmark-only --benchmark-json output.json
+            pytest tests/test_benchmark.py --gpu-only --batch-size 128 --benchmark-only --benchmark-json output.json
 
       - name: Compare benchmark results with main branch
         uses: benchmark-action/github-action-benchmark@v1.20.4

{jaxsim-0.6.2.dev194 → jaxsim-0.6.2.dev225}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.6.2.dev194
+Version: 0.6.2.dev225
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>
@@ -89,6 +89,7 @@ Requires-Dist: mujoco>=3.0.0; extra == "viz"
 Requires-Dist: scipy>=1.14.0; extra == "viz"
 Provides-Extra: all
 Requires-Dist: jaxsim[style,testing,viz]; extra == "all"
+Dynamic: license-file
 
 # JaxSim
 

{jaxsim-0.6.2.dev194 → jaxsim-0.6.2.dev225}/src/jaxsim/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.6.2.dev194'
-__version_tuple__ = version_tuple = (0, 6, 2, 'dev194')
+__version__ = version = '0.6.2.dev225'
+__version_tuple__ = version_tuple = (0, 6, 2, 'dev225')

{jaxsim-0.6.2.dev194 → jaxsim-0.6.2.dev225}/src/jaxsim/api/__init__.py

@@ -4,7 +4,6 @@ from . import (
     actuation_model,
     com,
     contact,
-    contact_model,
     frame,
     integrators,
     joint,

{jaxsim-0.6.2.dev194 → jaxsim-0.6.2.dev225}/src/jaxsim/api/com.py

@@ -301,9 +301,7 @@ def bias_acceleration(
             C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL  # noqa: F841
             C_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
 
-            L_H_C = L_H_W = jax.vmap(  # noqa: F841
-                lambda W_H_L: jaxsim.math.Transform.inverse(W_H_L)
-            )(W_H_L)
+            L_H_C = L_H_W = jax.vmap(jaxsim.math.Transform.inverse)(W_H_L)  # noqa: F841
 
             L_v_LC = L_v_LW = jax.vmap(  # noqa: F841
                 lambda i: -js.link.velocity(

{jaxsim-0.6.2.dev194 → jaxsim-0.6.2.dev225}/src/jaxsim/api/contact.py

@@ -11,7 +11,7 @@ import jaxsim.terrain
 import jaxsim.typing as jtp
 from jaxsim import logging
 from jaxsim.math import Adjoint, Cross, Transform
-from jaxsim.rbda import contacts
+from jaxsim.rbda.contacts import SoftContacts
 
 from .common import VelRepr
 
@@ -37,14 +37,11 @@ def collidable_point_kinematics(
         the linear component of the mixed 6D frame velocity.
     """
 
-    # Switch to inertial-fixed since the RBDAs expect velocities in this representation.
-    with data.switch_velocity_representation(VelRepr.Inertial):
-
-        W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
-            model=model,
-            link_transforms=data._link_transforms,
-            link_velocities=data._link_velocities,
-        )
+    W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
+        model=model,
+        link_transforms=data._link_transforms,
+        link_velocities=data._link_velocities,
+    )
 
     return W_p_Ci, W_ṗ_Ci
 
@@ -164,18 +161,23 @@ def estimate_good_soft_contacts_parameters(
 def estimate_good_contact_parameters(
     model: js.model.JaxSimModel,
     *,
+    standard_gravity: jtp.FloatLike = jaxsim.math.STANDARD_GRAVITY,
     static_friction_coefficient: jtp.FloatLike = 0.5,
-    **kwargs,
+    number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
+    damping_ratio: jtp.FloatLike = 1.0,
+    max_penetration: jtp.FloatLike | None = None,
 ) -> jaxsim.rbda.contacts.ContactParamsTypes:
     """
     Estimate good contact parameters.
 
     Args:
         model: The model to consider.
+        standard_gravity: The standard gravity acceleration.
         static_friction_coefficient: The static friction coefficient.
-        kwargs:
-            Additional model-specific parameters passed to the builder method of
-            the parameters class.
+        number_of_active_collidable_points_steady_state:
+            The number of active collidable points in steady state.
+        damping_ratio: The damping ratio.
+        max_penetration: The maximum penetration allowed.
 
     Returns:
         The estimated good contacts parameters.
@@ -190,20 +192,41 @@ def estimate_good_contact_parameters(
         specific application.
     """
 
-    match model.contact_model:
+    def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
+        """
+        Displacement between the CoM and the lowest collidable point using zero
+        joint positions.
+        """
+
+        zero_data = js.data.JaxSimModelData.build(
+            model=model,
+        )
+
+        W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
 
-        case contacts.RelaxedRigidContacts():
-            assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
+        if model.floating_base():
+            W_pz_C = collidable_point_positions(model=model, data=zero_data)[:, -1]
+            return 2 * (W_pz_CoM - W_pz_C.min())
 
-            parameters = contacts.RelaxedRigidContactsParams.build(
-                mu=static_friction_coefficient,
-                **kwargs,
-            )
+        return 2 * W_pz_CoM
 
-        case _:
-            raise ValueError(f"Invalid contact model: {model.contact_model}")
+    max_δ = (
+        max_penetration
+        if max_penetration is not None
+        # Consider as default a 0.5% of the model height.
+        else 0.005 * estimate_model_height(model=model)
+    )
 
-    return parameters
+    nc = number_of_active_collidable_points_steady_state
+
+    return model.contact_model._parameters_class().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,
+    )
 
 
 @jax.jit
@@ -244,7 +267,7 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
 
     # Build the link-to-point transform from the displacement between the link frame L
     # and the implicit contact frame C.
-    L_H_C = jax.vmap(lambda L_p_C: jnp.eye(4).at[0:3, 3].set(L_p_C))(L_p_Ci)
+    L_H_C = jax.vmap(jnp.eye(4).at[0:3, 3].set)(L_p_Ci)
 
     # Compose the work-to-link and link-to-point transforms.
     return jax.vmap(lambda W_H_Li, L_H_Ci: W_H_Li @ L_H_Ci)(W_H_L, L_H_C)
@@ -504,3 +527,96 @@ def jacobian_derivative(
     )
 
     return O_J̇_WC
+
+
+@jax.jit
+@js.common.named_scope
+def link_contact_forces(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_forces: jtp.MatrixLike | None = None,
+    joint_torques: jtp.VectorLike | None = None,
+) -> tuple[jtp.Matrix, dict[str, jtp.Matrix]]:
+    """
+    Compute the 6D contact forces of all links of the model in inertial representation.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        link_forces:
+            The 6D external forces to apply to the links expressed in inertial representation
+        joint_torques:
+            The joint torques acting on the joints.
+
+    Returns:
+        A `(nL, 6)` array containing the stacked 6D contact forces of the links,
+        expressed in inertial representation.
+    """
+
+    # Compute the contact forces for each collidable point with the active contact model.
+    W_f_C, aux_dict = model.contact_model.compute_contact_forces(
+        model=model,
+        data=data,
+        **(
+            dict(link_forces=link_forces, joint_force_references=joint_torques)
+            if not isinstance(model.contact_model, SoftContacts)
+            else {}
+        ),
+    )
+
+    # Compute the 6D forces applied to the links equivalent to the forces applied
+    # to the frames associated to the collidable points.
+    W_f_L = link_forces_from_contact_forces(model=model, contact_forces=W_f_C)
+
+    return W_f_L, aux_dict
+
+
+@staticmethod
+def link_forces_from_contact_forces(
+    model: js.model.JaxSimModel,
+    *,
+    contact_forces: jtp.MatrixLike,
+) -> jtp.Matrix:
+    """
+    Compute the link forces from the contact forces.
+
+    Args:
+        model: The robot model considered by the contact model.
+        contact_forces: The contact forces computed by the contact model.
+
+    Returns:
+        The 6D contact forces applied to the links and expressed in the frame of
+        the velocity representation of data.
+    """
+
+    # Get the object storing the contact parameters of the model.
+    contact_parameters = model.kin_dyn_parameters.contact_parameters
+
+    # Extract the indices corresponding to the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    # Convert the contact forces to a JAX array.
+    W_f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
+
+    # Construct the vector defining the parent link index of each collidable point.
+    # We use this vector to sum the 6D forces of all collidable points rigidly
+    # attached to the same link.
+    parent_link_index_of_collidable_points = jnp.array(
+        contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
+    # Create the mask that associate each collidable point to their parent link.
+    # We use this mask to sum the collidable points to the right link.
+    mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
+        model.number_of_links()
+    )
+
+    # Sum the forces of all collidable points rigidly attached to a body.
+    # Since the contact forces W_f_C are expressed in the world frame,
+    # we don't need any coordinate transformation.
+    W_f_L = mask.T @ W_f_C
+
+    return W_f_L

{jaxsim-0.6.2.dev194 → jaxsim-0.6.2.dev225}/src/jaxsim/api/data.py

@@ -5,9 +5,9 @@ import functools
 from collections.abc import Sequence
 
 try:
-    from typing import override
+    from typing import Self, override
 except ImportError:
-    from typing_extensions import override
+    from typing_extensions import override, Self
 
 import jax
 import jax.numpy as jnp
@@ -22,11 +22,6 @@ import jaxsim.typing as jtp
 from . import common
 from .common import VelRepr
 
-try:
-    from typing import Self
-except ImportError:
-    from typing_extensions import Self
-
 
 @jax_dataclasses.pytree_dataclass
 class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
@@ -64,6 +59,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     _link_transforms: jtp.Matrix = dataclasses.field(repr=False, default=None)
     _link_velocities: jtp.Matrix = dataclasses.field(repr=False, default=None)
 
+    # Extended state for soft and rigid contact models.
+    contact_state: dict[str, jtp.Array] = dataclasses.field(default=None)
+
     @staticmethod
     def build(
         model: js.model.JaxSimModel,
@@ -73,6 +71,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         base_linear_velocity: jtp.VectorLike | None = None,
         base_angular_velocity: jtp.VectorLike | None = None,
         joint_velocities: jtp.VectorLike | None = None,
+        contact_state: dict[str, jtp.Array] | None = None,
         velocity_representation: VelRepr = VelRepr.Mixed,
     ) -> JaxSimModelData:
         """
@@ -89,6 +88,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 The base angular velocity in the selected representation.
             joint_velocities: The joint velocities.
             velocity_representation: The velocity representation to use. It defaults to mixed if not provided.
+            contact_state: The optional contact state.
 
         Returns:
             A `JaxSimModelData` initialized with the given state.
@@ -171,6 +171,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             )
         )
 
+        contact_state = contact_state or {}
+
+        if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
+            contact_state.setdefault(
+                "tangential_deformation",
+                jnp.zeros_like(model.kin_dyn_parameters.contact_parameters.point),
+            )
+
         model_data = JaxSimModelData(
             velocity_representation=velocity_representation,
             _base_quaternion=base_quaternion,
@@ -183,6 +191,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             _joint_transforms=joint_transforms,
             _link_transforms=link_transforms,
             _link_velocities=link_velocities_inertial,
+            contact_state=contact_state,
         )
 
         if not model_data.valid(model=model):
@@ -347,11 +356,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     @js.common.named_scope
     @jax.jit
-    def reset_base_quaternion(self, base_quaternion: jtp.VectorLike) -> Self:
+    def reset_base_quaternion(
+        self, model: js.model.JaxSimModel, base_quaternion: jtp.VectorLike
+    ) -> Self:
         """
         Reset the base quaternion.
 
         Args:
+            model: The JaxSim model to use.
             base_quaternion: The base orientation as a quaternion.
 
         Returns:
@@ -363,15 +375,18 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         norm = jaxsim.math.safe_norm(W_Q_B, axis=-1)
         W_Q_B = W_Q_B / (norm + jnp.finfo(float).eps * (norm == 0))
 
-        return self.replace(validate=True, base_quaternion=W_Q_B)
+        return self.replace(model=model, base_quaternion=W_Q_B)
 
     @js.common.named_scope
     @jax.jit
-    def reset_base_pose(self, base_pose: jtp.MatrixLike) -> Self:
+    def reset_base_pose(
+        self, model: js.model.JaxSimModel, base_pose: jtp.MatrixLike
+    ) -> Self:
         """
         Reset the base pose.
 
         Args:
+            model: The JaxSim model to use.
             base_pose: The base pose as an SE(3) matrix.
 
         Returns:
@@ -382,6 +397,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         W_p_B = base_pose[0:3, 3]
         W_Q_B = jaxsim.math.Quaternion.from_dcm(dcm=base_pose[0:3, 0:3])
         return self.replace(
+            model=model,
             base_position=W_p_B,
             base_quaternion=W_Q_B,
         )
@@ -396,6 +412,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         base_linear_velocity: jtp.Vector | None = None,
         base_angular_velocity: jtp.Vector | None = None,
         base_position: jtp.Vector | None = None,
+        *,
+        contact_state: dict[str, jtp.Array] | None = None,
         validate: bool = False,
     ) -> Self:
         """
@@ -415,6 +433,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             base_quaternion = self.base_quaternion
         if base_position is None:
             base_position = self.base_position
+        if contact_state is None:
+            contact_state = self.contact_state
+
+        if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
+            contact_state.setdefault(
+                "tangential_deformation",
+                jnp.zeros_like(model.kin_dyn_parameters.contact_parameters.point),
+            )
 
         # Normalize the quaternion to avoid numerical issues.
         base_quaternion_norm = jaxsim.math.safe_norm(
@@ -486,8 +512,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         # Adjust the output shapes.
         if batch_size == 1:
-            link_transforms = link_transforms.reshape(link_transforms.shape[1:])
-            link_velocities = link_velocities.reshape(link_velocities.shape[1:])
+            link_transforms = link_transforms.reshape(self._link_transforms.shape)
+            link_velocities = link_velocities.reshape(self._link_velocities.shape)
+            joint_transforms = joint_transforms.reshape(self._joint_transforms.shape)
 
         return super().replace(
             _joint_positions=joint_positions,

{jaxsim-0.6.2.dev194 → jaxsim-0.6.2.dev225}/src/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,
@@ -64,6 +64,12 @@ def semi_implicit_euler_integration(
 
         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.
@@ -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-0.6.2.dev194 → jaxsim-0.6.2.dev225}/src/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,
     )
@@ -2090,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
     # =============================
@@ -2122,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-0.6.2.dev194 → jaxsim-0.6.2.dev225}/src/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,
     )