jaxsim 0.6.2.dev2__py3-none-any.whl → 0.6.2.dev105__py3-none-any.whl

jaxsim/api/data.py

@@ -4,6 +4,11 @@ import dataclasses
 import functools
 from collections.abc import Sequence
 
+try:
+    from typing import override
+except ImportError:
+    from typing_extensions import override
+
 import jax
 import jax.numpy as jnp
 import jax.scipy.spatial.transform
@@ -13,12 +18,9 @@ import jaxsim.api as js
 import jaxsim.math
 import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim.utils import Mutability
-from jaxsim.utils.tracing import not_tracing
 
 from . import common
 from .common import VelRepr
-from .ode_data import ODEState
 
 try:
     from typing import Self
@@ -29,72 +31,38 @@ except ImportError:
 @jax_dataclasses.pytree_dataclass
 class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     """
-    Class containing the data of a `JaxSimModel` object.
+    Class storing the state of the physics model dynamics.
+
+    Attributes:
+        joint_positions: The vector of joint positions.
+        joint_velocities: The vector of joint velocities.
+        base_position: The 3D position of the base link.
+        base_quaternion: The quaternion defining the orientation of the base link.
+        base_linear_velocity:
+            The linear velocity of the base link in inertial-fixed representation.
+        base_angular_velocity:
+            The angular velocity of the base link in inertial-fixed representation.
+        base_transform: The base transform.
+        joint_transforms: The joint transforms.
+        link_transforms: The link transforms.
+        link_velocities: The link velocities in inertial-fixed representation.
     """
 
-    state: ODEState
-
-    gravity: jtp.Vector
-
-    contacts_params: jaxsim.rbda.contacts.ContactsParams = dataclasses.field(repr=False)
-
-    def __hash__(self) -> int:
-
-        from jaxsim.utils.wrappers import HashedNumpyArray
-
-        return hash(
-            (
-                hash(self.state),
-                HashedNumpyArray.hash_of_array(self.gravity),
-                hash(self.contacts_params),
-            )
-        )
-
-    def __eq__(self, other: JaxSimModelData) -> bool:
-
-        if not isinstance(other, JaxSimModelData):
-            return False
-
-        return hash(self) == hash(other)
-
-    def valid(self, model: js.model.JaxSimModel | None = None) -> bool:
-        """
-        Check if the current state is valid for the given model.
-
-        Args:
-            model: The model to check against.
-
-        Returns:
-            `True` if the current state is valid for the given model, `False` otherwise.
-        """
+    # Joint state
+    _joint_positions: jtp.Vector
+    _joint_velocities: jtp.Vector
 
-        valid = True
-        valid = valid and self.standard_gravity() > 0
+    # Base state
+    _base_quaternion: jtp.Vector
+    _base_linear_velocity: jtp.Vector
+    _base_angular_velocity: jtp.Vector
+    _base_position: jtp.Vector
 
-        if model is not None:
-            valid = valid and self.state.valid(model=model)
-
-        return valid
-
-    @staticmethod
-    def zero(
-        model: js.model.JaxSimModel,
-        velocity_representation: VelRepr = VelRepr.Inertial,
-    ) -> JaxSimModelData:
-        """
-        Create a `JaxSimModelData` object with zero state.
-
-        Args:
-            model: The model for which to create the zero state.
-            velocity_representation: The velocity representation to use.
-
-        Returns:
-            A `JaxSimModelData` object with zero state.
-        """
-
-        return JaxSimModelData.build(
-            model=model, velocity_representation=velocity_representation
-        )
+    # Cached computations.
+    _base_transform: jtp.Matrix = dataclasses.field(repr=False, default=None)
+    _joint_transforms: jtp.Matrix = dataclasses.field(repr=False, default=None)
+    _link_transforms: jtp.Matrix = dataclasses.field(repr=False, default=None)
+    _link_velocities: jtp.Matrix = dataclasses.field(repr=False, default=None)
 
     @staticmethod
     def build(
@@ -105,10 +73,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,
-        standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
-        contacts_params: jaxsim.rbda.contacts.ContactsParams | None = None,
-        velocity_representation: VelRepr = VelRepr.Inertial,
-        extended_ode_state: dict[str, jtp.PyTree] | None = None,
+        velocity_representation: VelRepr = VelRepr.Mixed,
     ) -> JaxSimModelData:
         """
         Create a `JaxSimModelData` object with the given state.
@@ -123,13 +88,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             base_angular_velocity:
                 The base angular velocity in the selected representation.
             joint_velocities: The joint velocities.
-            standard_gravity: The standard gravity constant.
-            contacts_params: The parameters of the soft contacts.
-            velocity_representation: The velocity representation to use.
-            extended_ode_state:
-                Additional user-defined state variables that are not part of the
-                standard `ODEState` object. Useful to extend the system dynamics
-                considered by default in JaxSim.
+            velocity_representation: The velocity representation to use. It defaults to mixed if not provided.
 
         Returns:
             A `JaxSimModelData` initialized with the given state.
@@ -163,8 +122,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             dtype=float,
         ).squeeze()
 
-        gravity = jnp.zeros(3).at[2].set(-standard_gravity)
-
         joint_positions = jnp.atleast_1d(
             jnp.array(
                 (
@@ -191,166 +148,104 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             translation=base_position, quaternion=base_quaternion
         )
 
-        v_WB = JaxSimModelData.other_representation_to_inertial(
+        W_v_WB = JaxSimModelData.other_representation_to_inertial(
             array=jnp.hstack([base_linear_velocity, base_angular_velocity]),
             other_representation=velocity_representation,
             transform=W_H_B,
             is_force=False,
         ).astype(float)
 
-        ode_state = ODEState.build_from_jaxsim_model(
-            model=model,
-            base_position=base_position,
-            base_quaternion=base_quaternion,
-            joint_positions=joint_positions,
-            base_linear_velocity=v_WB[0:3],
-            base_angular_velocity=v_WB[3:6],
-            joint_velocities=joint_velocities,
-            # Unpack all the additional ODE states. If the contact model requires an
-            # additional state that is not explicitly passed to this builder, ODEState
-            # automatically populates that state with zeroed variables.
-            # This is not true for any other custom state that the user might want to
-            # pass to the integrator.
-            **(extended_ode_state if extended_ode_state else {}),
+        joint_transforms = model.kin_dyn_parameters.joint_transforms(
+            joint_positions=joint_positions, base_transform=W_H_B
         )
 
-        if not ode_state.valid(model=model):
-            raise ValueError(ode_state)
+        link_transforms, link_velocities_inertial = (
+            jaxsim.rbda.forward_kinematics_model(
+                model=model,
+                base_position=base_position,
+                base_quaternion=base_quaternion,
+                joint_positions=joint_positions,
+                base_linear_velocity_inertial=W_v_WB[0:3],
+                base_angular_velocity_inertial=W_v_WB[3:6],
+                joint_velocities=joint_velocities,
+            )
+        )
 
-        if contacts_params is None:
+        model_data = JaxSimModelData(
+            velocity_representation=velocity_representation,
+            _base_quaternion=base_quaternion,
+            _base_position=base_position,
+            _joint_positions=joint_positions,
+            _base_linear_velocity=W_v_WB[0:3],
+            _base_angular_velocity=W_v_WB[3:6],
+            _joint_velocities=joint_velocities,
+            _base_transform=W_H_B,
+            _joint_transforms=joint_transforms,
+            _link_transforms=link_transforms,
+            _link_velocities=link_velocities_inertial,
+        )
 
-            if isinstance(
-                model.contact_model,
-                jaxsim.rbda.contacts.SoftContacts
-                | jaxsim.rbda.contacts.ViscoElasticContacts,
-            ):
+        if not model_data.valid(model=model):
+            raise ValueError(
+                "The built state is not compatible with the model.", model_data
+            )
 
-                contacts_params = js.contact.estimate_good_contact_parameters(
-                    model=model, standard_gravity=standard_gravity
-                )
+        return model_data
 
-            else:
-                contacts_params = model.contact_model._parameters_class()
+    @staticmethod
+    def zero(
+        model: js.model.JaxSimModel,
+        velocity_representation: VelRepr = VelRepr.Mixed,
+    ) -> JaxSimModelData:
+        """
+        Create a `JaxSimModelData` object with zero state.
 
-        return JaxSimModelData(
-            state=ode_state,
-            gravity=gravity,
-            contacts_params=contacts_params,
-            velocity_representation=velocity_representation,
+        Args:
+            model: The model for which to create the state.
+            velocity_representation: The velocity representation to use. It defaults to mixed if not provided.
+
+        Returns:
+            A `JaxSimModelData` initialized with zero state.
+        """
+        return JaxSimModelData.build(
+            model=model, velocity_representation=velocity_representation
         )
 
     # ==================
     # Extract quantities
     # ==================
 
-    def standard_gravity(self) -> jtp.Float:
+    @property
+    def joint_positions(self) -> jtp.Vector:
         """
-        Get the standard gravity constant.
+        Get the joint positions.
 
         Returns:
-            The standard gravity constant.
+            The joint positions.
         """
+        return self._joint_positions
 
-        return -self.gravity[2]
-
-    @js.common.named_scope
-    @functools.partial(jax.jit, static_argnames=["joint_names"])
-    def joint_positions(
-        self,
-        model: js.model.JaxSimModel | None = None,
-        joint_names: tuple[str, ...] | None = None,
-    ) -> jtp.Vector:
+    @property
+    def joint_velocities(self) -> jtp.Vector:
         """
-        Get the joint positions.
-
-        Args:
-            model: The model to consider.
-            joint_names:
-                The names of the joints for which to get the positions. If `None`,
-                the positions of all joints are returned.
+        Get the joint velocities.
 
         Returns:
-            If no model and no joint names are provided, the joint positions as a
-            `(DoFs,)` vector corresponding to the serialization of the original
-            model used to build the data object.
-            If a model is provided and no joint names are provided, the joint positions
-            as a `(DoFs,)` vector corresponding to the serialization of the
-            provided model.
-            If a model and joint names are provided, the joint positions as a
-            `(len(joint_names),)` vector corresponding to the serialization of
-            the passed joint names vector.
+            The joint velocities.
         """
+        return self._joint_velocities
 
-        if model is None:
-            if joint_names is not None:
-                raise ValueError("Joint names cannot be provided without a model")
-
-            return self.state.physics_model.joint_positions
-
-        if not_tracing(self.state.physics_model.joint_positions) and not self.valid(
-            model=model
-        ):
-            msg = "The data object is not compatible with the provided model"
-            raise ValueError(msg)
-
-        joint_idxs = (
-            js.joint.names_to_idxs(joint_names=joint_names, model=model)
-            if joint_names is not None
-            else jnp.arange(model.number_of_joints())
-        )
-
-        return self.state.physics_model.joint_positions[joint_idxs]
-
-    @js.common.named_scope
-    @functools.partial(jax.jit, static_argnames=["joint_names"])
-    def joint_velocities(
-        self,
-        model: js.model.JaxSimModel | None = None,
-        joint_names: tuple[str, ...] | None = None,
-    ) -> jtp.Vector:
+    @property
+    def base_quaternion(self) -> jtp.Vector:
         """
-        Get the joint velocities.
-
-        Args:
-            model: The model to consider.
-            joint_names:
-                The names of the joints for which to get the velocities. If `None`,
-                the velocities of all joints are returned.
+        Get the base quaternion.
 
         Returns:
-            If no model and no joint names are provided, the joint velocities as a
-            `(DoFs,)` vector corresponding to the serialization of the original
-            model used to build the data object.
-            If a model is provided and no joint names are provided, the joint velocities
-            as a `(DoFs,)` vector corresponding to the serialization of the
-            provided model.
-            If a model and joint names are provided, the joint velocities as a
-            `(len(joint_names),)` vector corresponding to the serialization of
-            the passed joint names vector.
+            The base quaternion.
         """
+        return self._base_quaternion
 
-        if model is None:
-            if joint_names is not None:
-                raise ValueError("Joint names cannot be provided without a model")
-
-            return self.state.physics_model.joint_velocities
-
-        if not_tracing(self.state.physics_model.joint_velocities) and not self.valid(
-            model=model
-        ):
-            msg = "The data object is not compatible with the provided model"
-            raise ValueError(msg)
-
-        joint_idxs = (
-            js.joint.names_to_idxs(joint_names=joint_names, model=model)
-            if joint_names is not None
-            else jnp.arange(model.number_of_joints())
-        )
-
-        return self.state.physics_model.joint_velocities[joint_idxs]
-
-    @js.common.named_scope
-    @jax.jit
+    @property
     def base_position(self) -> jtp.Vector:
         """
         Get the base position.
@@ -358,24 +253,19 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         Returns:
             The base position.
         """
+        return self._base_position
 
-        return self.state.physics_model.base_position.squeeze()
-
-    @js.common.named_scope
-    @functools.partial(jax.jit, static_argnames=["dcm"])
-    def base_orientation(self, dcm: jtp.BoolLike = False) -> jtp.Vector | jtp.Matrix:
+    @property
+    def base_orientation(self) -> jtp.Matrix:
         """
         Get the base orientation.
 
-        Args:
-            dcm: Whether to return the orientation as a SO(3) matrix or quaternion.
-
         Returns:
             The base orientation.
         """
 
         # Extract the base quaternion.
-        W_Q_B = self.state.physics_model.base_quaternion.squeeze()
+        W_Q_B = self.base_quaternion.squeeze()
 
         # Always normalize the quaternion to avoid numerical issues.
         # If the active scheme does not integrate the quaternion on its manifold,
@@ -384,33 +274,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         # stored in the state is a unit quaternion.
         norm = jaxsim.math.safe_norm(W_Q_B)
         W_Q_B = W_Q_B / (norm + jnp.finfo(float).eps * (norm == 0))
+        return W_Q_B
 
-        return (W_Q_B if not dcm else jaxsim.math.Quaternion.to_dcm(W_Q_B)).astype(
-            float
-        )
-
-    @js.common.named_scope
-    @jax.jit
-    def base_transform(self) -> jtp.Matrix:
-        """
-        Get the base transform.
-
-        Returns:
-            The base transform as an SE(3) matrix.
-        """
-
-        W_R_B = self.base_orientation(dcm=True)
-        W_p_B = jnp.vstack(self.base_position())
-
-        return jnp.vstack(
-            [
-                jnp.block([W_R_B, W_p_B]),
-                jnp.array([0, 0, 0, 1]),
-            ]
-        )
-
-    @js.common.named_scope
-    @jax.jit
+    @property
     def base_velocity(self) -> jtp.Vector:
         """
         Get the base 6D velocity.
@@ -421,12 +287,12 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         W_v_WB = jnp.hstack(
             [
-                self.state.physics_model.base_linear_velocity,
-                self.state.physics_model.base_angular_velocity,
+                self._base_linear_velocity,
+                self._base_angular_velocity,
             ]
         )
 
-        W_H_B = self.base_transform()
+        W_H_B = self._base_transform
 
         return (
             JaxSimModelData.inertial_to_other_representation(
@@ -439,8 +305,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             .astype(float)
         )
 
-    @js.common.named_scope
-    @jax.jit
+    @property
     def generalized_position(self) -> tuple[jtp.Matrix, jtp.Vector]:
         r"""
         Get the generalized position
@@ -450,10 +315,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             A tuple containing the base transform and the joint positions.
         """
 
-        return self.base_transform(), self.joint_positions()
+        return self._base_transform, self.joint_positions
 
-    @js.common.named_scope
-    @jax.jit
+    @property
     def generalized_velocity(self) -> jtp.Vector:
         r"""
         Get the generalized velocity.
@@ -465,136 +329,24 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         """
 
         return (
-            jnp.hstack([self.base_velocity(), self.joint_velocities()])
+            jnp.hstack([self.base_velocity, self.joint_velocities])
             .squeeze()
             .astype(float)
         )
 
-    # ================
-    # Store quantities
-    # ================
-
-    @js.common.named_scope
-    @functools.partial(jax.jit, static_argnames=["joint_names"])
-    def reset_joint_positions(
-        self,
-        positions: jtp.VectorLike,
-        model: js.model.JaxSimModel | None = None,
-        joint_names: tuple[str, ...] | None = None,
-    ) -> Self:
-        """
-        Reset the joint positions.
-
-        Args:
-            positions: The joint positions.
-            model: The model to consider.
-            joint_names: The names of the joints for which to set the positions.
-
-        Returns:
-            The updated `JaxSimModelData` object.
-        """
-
-        positions = jnp.array(positions)
-
-        def replace(s: jtp.VectorLike) -> JaxSimModelData:
-            return self.replace(
-                validate=True,
-                state=self.state.replace(
-                    physics_model=self.state.physics_model.replace(
-                        joint_positions=jnp.atleast_1d(s.squeeze()).astype(float)
-                    )
-                ),
-            )
-
-        if model is None:
-            return replace(s=positions)
-
-        if not_tracing(positions) and not self.valid(model=model):
-            msg = "The data object is not compatible with the provided model"
-            raise ValueError(msg)
-
-        joint_idxs = (
-            js.joint.names_to_idxs(joint_names=joint_names, model=model)
-            if joint_names is not None
-            else jnp.arange(model.number_of_joints())
-        )
-
-        return replace(
-            s=self.state.physics_model.joint_positions.at[joint_idxs].set(positions)
-        )
-
-    @js.common.named_scope
-    @functools.partial(jax.jit, static_argnames=["joint_names"])
-    def reset_joint_velocities(
-        self,
-        velocities: jtp.VectorLike,
-        model: js.model.JaxSimModel | None = None,
-        joint_names: tuple[str, ...] | None = None,
-    ) -> Self:
-        """
-        Reset the joint velocities.
-
-        Args:
-            velocities: The joint velocities.
-            model: The model to consider.
-            joint_names: The names of the joints for which to set the velocities.
-
-        Returns:
-            The updated `JaxSimModelData` object.
-        """
-
-        velocities = jnp.array(velocities)
-
-        def replace(ṡ: jtp.VectorLike) -> JaxSimModelData:
-            return self.replace(
-                validate=True,
-                state=self.state.replace(
-                    physics_model=self.state.physics_model.replace(
-                        joint_velocities=jnp.atleast_1d(ṡ.squeeze()).astype(float)
-                    )
-                ),
-            )
-
-        if model is None:
-            return replace(ṡ=velocities)
-
-        if not_tracing(velocities) and not self.valid(model=model):
-            msg = "The data object is not compatible with the provided model"
-            raise ValueError(msg)
-
-        joint_idxs = (
-            js.joint.names_to_idxs(joint_names=joint_names, model=model)
-            if joint_names is not None
-            else jnp.arange(model.number_of_joints())
-        )
-
-        return replace(
-            ṡ=self.state.physics_model.joint_velocities.at[joint_idxs].set(velocities)
-        )
-
-    @js.common.named_scope
-    @jax.jit
-    def reset_base_position(self, base_position: jtp.VectorLike) -> Self:
+    @property
+    def base_transform(self) -> jtp.Matrix:
         """
-        Reset the base position.
-
-        Args:
-            base_position: The base position.
+        Get the base transform.
 
         Returns:
-            The updated `JaxSimModelData` object.
+            The base transform.
         """
+        return self._base_transform
 
-        base_position = jnp.array(base_position)
-
-        return self.replace(
-            validate=True,
-            state=self.state.replace(
-                physics_model=self.state.physics_model.replace(
-                    base_position=jnp.atleast_1d(base_position.squeeze()).astype(float)
-                )
-            ),
-        )
+    # ================
+    # Store quantities
+    # ================
 
     @js.common.named_scope
     @jax.jit
@@ -614,12 +366,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         norm = jaxsim.math.safe_norm(W_Q_B)
         W_Q_B = W_Q_B / (norm + jnp.finfo(float).eps * (norm == 0))
 
-        return self.replace(
-            validate=True,
-            state=self.state.replace(
-                physics_model=self.state.physics_model.replace(base_quaternion=W_Q_B)
-            ),
-        )
+        return self.replace(validate=True, base_quaternion=W_Q_B)
 
     @js.common.named_scope
     @jax.jit
@@ -635,123 +382,116 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         """
 
         base_pose = jnp.array(base_pose)
-
         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.reset_base_position(base_position=W_p_B).reset_base_quaternion(
-            base_quaternion=W_Q_B
+        return self.replace(
+            base_position=W_p_B,
+            base_quaternion=W_Q_B,
         )
 
-    @js.common.named_scope
-    @functools.partial(jax.jit, static_argnames=["velocity_representation"])
-    def reset_base_linear_velocity(
+    @override
+    def replace(
         self,
-        linear_velocity: jtp.VectorLike,
-        velocity_representation: VelRepr | None = None,
+        model: js.model.JaxSimModel,
+        joint_positions: jtp.Vector | None = None,
+        joint_velocities: jtp.Vector | None = None,
+        base_quaternion: jtp.Vector | None = None,
+        base_linear_velocity: jtp.Vector | None = None,
+        base_angular_velocity: jtp.Vector | None = None,
+        base_position: jtp.Vector | None = None,
+        validate: bool = False,
     ) -> Self:
         """
-        Reset the base linear velocity.
-
-        Args:
-            linear_velocity: The base linear velocity as a 3D array.
-            velocity_representation:
-                The velocity representation in which the base velocity is expressed.
-                If `None`, the active representation is considered.
-
-        Returns:
-            The updated `JaxSimModelData` object.
+        Replace the attributes of the `JaxSimModelData` object.
         """
+        if joint_positions is None:
+            joint_positions = self.joint_positions
+        if joint_velocities is None:
+            joint_velocities = self.joint_velocities
+        if base_quaternion is None:
+            base_quaternion = self.base_quaternion
+        if base_position is None:
+            base_position = self.base_position
 
-        linear_velocity = jnp.array(linear_velocity)
+        joint_positions = jnp.atleast_1d(joint_positions.squeeze()).astype(float)
+        joint_velocities = jnp.atleast_1d(joint_velocities.squeeze()).astype(float)
+        base_quaternion = jnp.atleast_1d(base_quaternion.squeeze()).astype(float)
+        base_position = jnp.atleast_1d(base_position.squeeze()).astype(float)
 
-        return self.reset_base_velocity(
-            base_velocity=jnp.hstack(
-                [
-                    linear_velocity.squeeze(),
-                    self.base_velocity()[3:6],
-                ]
-            ),
-            velocity_representation=velocity_representation,
+        base_transform = jaxsim.math.Transform.from_quaternion_and_translation(
+            translation=base_position, quaternion=base_quaternion
+        )
+        joint_transforms = model.kin_dyn_parameters.joint_transforms(
+            joint_positions=joint_positions, base_transform=base_transform
         )
 
-    @js.common.named_scope
-    @functools.partial(jax.jit, static_argnames=["velocity_representation"])
-    def reset_base_angular_velocity(
-        self,
-        angular_velocity: jtp.VectorLike,
-        velocity_representation: VelRepr | None = None,
-    ) -> Self:
-        """
-        Reset the base angular velocity.
-
-        Args:
-            angular_velocity: The base angular velocity as a 3D array.
-            velocity_representation:
-                The velocity representation in which the base velocity is expressed.
-                If `None`, the active representation is considered.
-
-        Returns:
-            The updated `JaxSimModelData` object.
-        """
+        if base_linear_velocity is None and base_angular_velocity is None:
+            base_linear_velocity = self._base_linear_velocity
+            base_angular_velocity = self._base_angular_velocity
+        else:
+            if base_linear_velocity is None:
+                base_linear_velocity = self.base_velocity[:3]
+            if base_angular_velocity is None:
+                base_angular_velocity = self.base_velocity[3:]
+            base_linear_velocity = jnp.atleast_1d(base_linear_velocity.squeeze())
+            base_angular_velocity = jnp.atleast_1d(base_angular_velocity.squeeze())
+            W_v_WB = JaxSimModelData.other_representation_to_inertial(
+                array=jnp.hstack([base_linear_velocity, base_angular_velocity]),
+                other_representation=self.velocity_representation,
+                transform=base_transform,
+                is_force=False,
+            ).astype(float)
+            base_linear_velocity, base_angular_velocity = W_v_WB[:3], W_v_WB[3:]
 
-        angular_velocity = jnp.array(angular_velocity)
+        link_transforms, link_velocities = jaxsim.rbda.forward_kinematics_model(
+            model=model,
+            base_position=base_position,
+            base_quaternion=base_quaternion,
+            joint_positions=joint_positions,
+            joint_velocities=joint_velocities,
+            base_linear_velocity_inertial=base_linear_velocity,
+            base_angular_velocity_inertial=base_angular_velocity,
+        )
 
-        return self.reset_base_velocity(
-            base_velocity=jnp.hstack(
-                [
-                    self.base_velocity()[0:3],
-                    angular_velocity.squeeze(),
-                ]
-            ),
-            velocity_representation=velocity_representation,
+        return super().replace(
+            _joint_positions=joint_positions,
+            _joint_velocities=joint_velocities,
+            _base_quaternion=base_quaternion,
+            _base_linear_velocity=base_linear_velocity,
+            _base_angular_velocity=base_angular_velocity,
+            _base_position=base_position,
+            _base_transform=base_transform,
+            _joint_transforms=joint_transforms,
+            _link_transforms=link_transforms,
+            _link_velocities=link_velocities,
+            validate=validate,
         )
 
-    @js.common.named_scope
-    @functools.partial(jax.jit, static_argnames=["velocity_representation"])
-    def reset_base_velocity(
-        self,
-        base_velocity: jtp.VectorLike,
-        velocity_representation: VelRepr | None = None,
-    ) -> Self:
+    def valid(self, model: js.model.JaxSimModel) -> bool:
         """
-        Reset the base 6D velocity.
+        Check if the `JaxSimModelData` is valid for a given `JaxSimModel`.
 
         Args:
-            base_velocity: The base 6D velocity in the active representation.
-            velocity_representation:
-                The velocity representation in which the base velocity is expressed.
-                If `None`, the active representation is considered.
+            model: The `JaxSimModel` to validate the `JaxSimModelData` against.
 
         Returns:
-            The updated `JaxSimModelData` object.
+            `True` if the `JaxSimModelData` is valid for the given model,
+            `False` otherwise.
         """
+        if self._joint_positions.shape != (model.dofs(),):
+            return False
+        if self._joint_velocities.shape != (model.dofs(),):
+            return False
+        if self._base_position.shape != (3,):
+            return False
+        if self._base_quaternion.shape != (4,):
+            return False
+        if self._base_linear_velocity.shape != (3,):
+            return False
+        if self._base_angular_velocity.shape != (3,):
+            return False
 
-        base_velocity = jnp.array(base_velocity)
-
-        velocity_representation = (
-            velocity_representation
-            if velocity_representation is not None
-            else self.velocity_representation
-        )
-
-        W_v_WB = self.other_representation_to_inertial(
-            array=jnp.atleast_1d(base_velocity.squeeze()).astype(float),
-            other_representation=velocity_representation,
-            transform=self.base_transform(),
-            is_force=False,
-        )
-
-        return self.replace(
-            validate=True,
-            state=self.state.replace(
-                physics_model=self.state.physics_model.replace(
-                    base_linear_velocity=W_v_WB[0:3].squeeze().astype(float),
-                    base_angular_velocity=W_v_WB[3:6].squeeze().astype(float),
-                )
-            ),
-        )
+        return True
 
 
 @functools.partial(jax.jit, static_argnames=["velocity_representation", "base_rpy_seq"])
@@ -788,11 +528,6 @@ def random_model_data(
         jtp.FloatLike | Sequence[jtp.FloatLike],
         jtp.FloatLike | Sequence[jtp.FloatLike],
     ] = (-1.0, 1.0),
-    contacts_params: jaxsim.rbda.contacts.ContactsParams | None = None,
-    standard_gravity_bounds: tuple[jtp.FloatLike, jtp.FloatLike] = (
-        jaxsim.math.StandardGravity,
-        jaxsim.math.StandardGravity,
-    ),
 ) -> JaxSimModelData:
     """
     Randomly generate a `JaxSimModelData` object.
@@ -811,15 +546,13 @@ def random_model_data(
         base_vel_lin_bounds: The bounds for the base linear velocity.
         base_vel_ang_bounds: The bounds for the base angular velocity.
         joint_vel_bounds: The bounds for the joint velocities.
-        contacts_params: The parameters of the contact model.
-        standard_gravity_bounds: The bounds for the standard gravity.
 
     Returns:
         A `JaxSimModelData` object with random data.
     """
 
     key = key if key is not None else jax.random.PRNGKey(seed=0)
-    k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, num=7)
+    k1, k2, k3, k4, k5, k6 = jax.random.split(key, num=6)
 
     p_min = jnp.array(base_pos_bounds[0], dtype=float)
     p_max = jnp.array(base_pos_bounds[1], dtype=float)
@@ -831,95 +564,64 @@ def random_model_data(
     ω_max = jnp.array(base_vel_ang_bounds[1], dtype=float)
     ṡ_min, ṡ_max = joint_vel_bounds
 
-    random_data = JaxSimModelData.zero(
-        model=model,
-        **(
-            dict(velocity_representation=velocity_representation)
-            if velocity_representation is not None
-            else {}
-        ),
-    )
+    base_position = jax.random.uniform(key=k1, shape=(3,), minval=p_min, maxval=p_max)
 
-    with random_data.mutable_context(
-        mutability=Mutability.MUTABLE, restore_after_exception=False
-    ):
+    base_quaternion = jaxsim.math.Quaternion.to_wxyz(
+        xyzw=jax.scipy.spatial.transform.Rotation.from_euler(
+            seq=base_rpy_seq,
+            angles=jax.random.uniform(
+                key=k2, shape=(3,), minval=rpy_min, maxval=rpy_max
+            ),
+        ).as_quat()
+    )
 
-        physics_model_state = random_data.state.physics_model
+    (
+        joint_positions,
+        joint_velocities,
+        base_linear_velocity,
+        base_angular_velocity,
+    ) = (None,) * 4
 
-        physics_model_state.base_position = jax.random.uniform(
-            key=k1, shape=(3,), minval=p_min, maxval=p_max
-        )
+    if model.number_of_joints() > 0:
 
-        physics_model_state.base_quaternion = jaxsim.math.Quaternion.to_wxyz(
-            xyzw=jax.scipy.spatial.transform.Rotation.from_euler(
-                seq=base_rpy_seq,
-                angles=jax.random.uniform(
-                    key=k2, shape=(3,), minval=rpy_min, maxval=rpy_max
-                ),
-            ).as_quat()
+        s_min, s_max = (
+            jnp.array(joint_pos_bounds, dtype=float)
+            if joint_pos_bounds is not None
+            else (None, None)
         )
 
-        if model.number_of_joints() > 0:
-
-            s_min, s_max = (
-                jnp.array(joint_pos_bounds, dtype=float)
-                if joint_pos_bounds is not None
-                else (None, None)
-            )
-
-            physics_model_state.joint_positions = (
-                js.joint.random_joint_positions(model=model, key=k3)
-                if (s_min is None or s_max is None)
-                else jax.random.uniform(
-                    key=k3, shape=(model.dofs(),), minval=s_min, maxval=s_max
-                )
-            )
-
-            physics_model_state.joint_velocities = jax.random.uniform(
-                key=k4, shape=(model.dofs(),), minval=ṡ_min, maxval=ṡ_max
-            )
-
-        if model.floating_base():
-            physics_model_state.base_linear_velocity = jax.random.uniform(
-                key=k5, shape=(3,), minval=v_min, maxval=v_max
-            )
-
-            physics_model_state.base_angular_velocity = jax.random.uniform(
-                key=k6, shape=(3,), minval=ω_min, maxval=ω_max
-            )
-
-        random_data.gravity = (
-            jnp.zeros(3, dtype=random_data.gravity.dtype)
-            .at[2]
-            .set(
-                -jax.random.uniform(
-                    key=k7,
-                    shape=(),
-                    minval=standard_gravity_bounds[0],
-                    maxval=standard_gravity_bounds[1],
-                )
+        joint_positions = (
+            js.joint.random_joint_positions(model=model, key=k3)
+            if (s_min is None or s_max is None)
+            else jax.random.uniform(
+                key=k3, shape=(model.dofs(),), minval=s_min, maxval=s_max
             )
         )
 
-        if contacts_params is None:
-
-            if isinstance(
-                model.contact_model,
-                jaxsim.rbda.contacts.SoftContacts
-                | jaxsim.rbda.contacts.ViscoElasticContacts,
-            ):
+        joint_velocities = jax.random.uniform(
+            key=k4, shape=(model.dofs(),), minval=ṡ_min, maxval=ṡ_max
+        )
 
-                random_data = random_data.replace(
-                    contacts_params=js.contact.estimate_good_contact_parameters(
-                        model=model, standard_gravity=random_data.gravity
-                    ),
-                    validate=False,
-                )
+    if model.floating_base():
+        base_linear_velocity = jax.random.uniform(
+            key=k5, shape=(3,), minval=v_min, maxval=v_max
+        )
 
-            else:
-                random_data = random_data.replace(
-                    contacts_params=model.contact_model._parameters_class(),
-                    validate=False,
-                )
+        base_angular_velocity = jax.random.uniform(
+            key=k6, shape=(3,), minval=ω_min, maxval=ω_max
+        )
 
-    return random_data
+    return JaxSimModelData.build(
+        model=model,
+        base_position=base_position,
+        base_quaternion=base_quaternion,
+        joint_positions=joint_positions,
+        joint_velocities=joint_velocities,
+        base_linear_velocity=base_linear_velocity,
+        base_angular_velocity=base_angular_velocity,
+        **(
+            {"velocity_representation": velocity_representation}
+            if velocity_representation is not None
+            else {}
+        ),
+    )