jaxsim 0.4.3.dev115__tar.gz → 0.4.3.dev129__tar.gz

{jaxsim-0.4.3.dev115 → jaxsim-0.4.3.dev129}/PKG-INFO

@@ -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 → jaxsim-0.4.3.dev129}/src/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-0.4.3.dev115 → jaxsim-0.4.3.dev129}/src/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-0.4.3.dev115 → jaxsim-0.4.3.dev129}/src/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-0.4.3.dev115 → jaxsim-0.4.3.dev129}/src/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-0.4.3.dev115 → jaxsim-0.4.3.dev129}/src/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-0.4.3.dev115 → jaxsim-0.4.3.dev129}/src/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-0.4.3.dev115 → jaxsim-0.4.3.dev129}/src/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-0.4.3.dev115 → jaxsim-0.4.3.dev129}/src/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-0.4.3.dev115 → jaxsim-0.4.3.dev129}/src/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-0.4.3.dev115 → jaxsim-0.4.3.dev129}/src/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]