jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev2__py3-none-any.whl

jaxsim/api/data.py

@@ -2,28 +2,23 @@ from __future__ import annotations
 
 import dataclasses
 import functools
-from typing import Sequence
+from collections.abc import Sequence
 
 import jax
 import jax.numpy as jnp
+import jax.scipy.spatial.transform
 import jax_dataclasses
-import jaxlie
-import numpy as np
-
-import jaxsim.api
-import jaxsim.physics.algos.aba
-import jaxsim.physics.algos.crba
-import jaxsim.physics.algos.forward_kinematics
-import jaxsim.physics.algos.rnea
-import jaxsim.physics.model.physics_model
-import jaxsim.physics.model.physics_model_state
+
+import jaxsim.api as js
+import jaxsim.math
+import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim.high_level.common import VelRepr
-from jaxsim.physics.algos import soft_contacts
-from jaxsim.simulation.ode_data import ODEState
 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
@@ -34,21 +29,35 @@ except ImportError:
 @jax_dataclasses.pytree_dataclass
 class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     """
-    Class containing the state of a `JaxSimModel` object.
+    Class containing the data of a `JaxSimModel` object.
     """
 
     state: ODEState
 
-    gravity: jtp.Array
+    gravity: jtp.Vector
 
-    soft_contacts_params: soft_contacts.SoftContactsParams = dataclasses.field(
-        repr=False
-    )
-    time_ns: jtp.Int = dataclasses.field(
-        default_factory=lambda: jnp.array(0, dtype=jnp.uint64)
-    )
+    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:
 
-    def valid(self, model: jaxsim.api.model.JaxSimModel | None = None) -> 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.
 
@@ -60,15 +69,16 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         """
 
         valid = True
+        valid = valid and self.standard_gravity() > 0
 
         if model is not None:
-            valid = valid and self.state.valid(physics_model=model.physics_model)
+            valid = valid and self.state.valid(model=model)
 
         return valid
 
     @staticmethod
     def zero(
-        model: jaxsim.api.model.JaxSimModel,
+        model: js.model.JaxSimModel,
         velocity_representation: VelRepr = VelRepr.Inertial,
     ) -> JaxSimModelData:
         """
@@ -88,18 +98,17 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     @staticmethod
     def build(
-        model: jaxsim.api.model.JaxSimModel,
-        base_position: jtp.Vector | None = None,
-        base_quaternion: jtp.Vector | None = None,
-        joint_positions: jtp.Vector | None = None,
-        base_linear_velocity: jtp.Vector | None = None,
-        base_angular_velocity: jtp.Vector | None = None,
-        joint_velocities: jtp.Vector | None = None,
-        gravity: jtp.Vector | None = None,
-        soft_contacts_state: soft_contacts.SoftContactsState | None = None,
-        soft_contacts_params: soft_contacts.SoftContactsParams | None = None,
+        model: js.model.JaxSimModel,
+        base_position: jtp.VectorLike | None = None,
+        base_quaternion: jtp.VectorLike | None = None,
+        joint_positions: jtp.VectorLike | None = None,
+        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,
-        time: jtp.FloatLike | None = None,
+        extended_ode_state: dict[str, jtp.PyTree] | None = None,
     ) -> JaxSimModelData:
         """
         Create a `JaxSimModelData` object with the given state.
@@ -114,97 +123,119 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             base_angular_velocity:
                 The base angular velocity in the selected representation.
             joint_velocities: The joint velocities.
-            gravity: The gravity 3D vector.
-            soft_contacts_state: The state of the soft contacts.
-            soft_contacts_params: The parameters of the soft contacts.
+            standard_gravity: The standard gravity constant.
+            contacts_params: The parameters of the soft contacts.
             velocity_representation: The velocity representation to use.
-            time: The time at which the state is created.
+            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.
 
         Returns:
-            A `JaxSimModelData` object with the given state.
+            A `JaxSimModelData` initialized with the given state.
         """
 
         base_position = jnp.array(
-            base_position if base_position is not None else jnp.zeros(3)
+            base_position if base_position is not None else jnp.zeros(3),
+            dtype=float,
         ).squeeze()
 
         base_quaternion = jnp.array(
-            base_quaternion
-            if base_quaternion is not None
-            else jnp.array([1.0, 0, 0, 0])
+            (
+                base_quaternion
+                if base_quaternion is not None
+                else jnp.array([1.0, 0, 0, 0])
+            ),
+            dtype=float,
         ).squeeze()
 
         base_linear_velocity = jnp.array(
-            base_linear_velocity if base_linear_velocity is not None else jnp.zeros(3)
+            base_linear_velocity if base_linear_velocity is not None else jnp.zeros(3),
+            dtype=float,
         ).squeeze()
 
         base_angular_velocity = jnp.array(
-            base_angular_velocity if base_angular_velocity is not None else jnp.zeros(3)
+            (
+                base_angular_velocity
+                if base_angular_velocity is not None
+                else jnp.zeros(3)
+            ),
+            dtype=float,
         ).squeeze()
 
-        gravity = jnp.array(
-            gravity if gravity is not None else model.physics_model.gravity[0:3]
-        ).squeeze()
+        gravity = jnp.zeros(3).at[2].set(-standard_gravity)
 
         joint_positions = jnp.atleast_1d(
-            joint_positions.squeeze()
-            if joint_positions is not None
-            else jnp.zeros(model.dofs())
+            jnp.array(
+                (
+                    joint_positions
+                    if joint_positions is not None
+                    else jnp.zeros(model.dofs())
+                ),
+                dtype=float,
+            ).squeeze()
         )
 
         joint_velocities = jnp.atleast_1d(
-            joint_velocities.squeeze()
-            if joint_velocities is not None
-            else jnp.zeros(model.dofs())
-        )
-
-        time_ns = (
-            jnp.array(time * 1e9, dtype=jnp.uint64)
-            if time is not None
-            else jnp.array(0, dtype=jnp.uint64)
+            jnp.array(
+                (
+                    joint_velocities
+                    if joint_velocities is not None
+                    else jnp.zeros(model.dofs())
+                ),
+                dtype=float,
+            ).squeeze()
         )
 
-        soft_contacts_params = (
-            soft_contacts_params
-            if soft_contacts_params is not None
-            else jaxsim.api.contact.estimate_good_soft_contacts_parameters(model=model)
+        W_H_B = jaxsim.math.Transform.from_quaternion_and_translation(
+            translation=base_position, quaternion=base_quaternion
         )
 
-        W_H_B = jaxlie.SE3.from_rotation_and_translation(
-            translation=base_position,
-            rotation=jaxlie.SO3.from_quaternion_xyzw(
-                base_quaternion[jnp.array([1, 2, 3, 0])]
-            ),
-        ).as_matrix()
-
         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 {}),
         )
 
-        ode_state = ODEState.build(
-            physics_model=model.physics_model,
-            physics_model_state=jaxsim.physics.model.physics_model_state.PhysicsModelState(
-                base_position=base_position.astype(float),
-                base_quaternion=base_quaternion.astype(float),
-                joint_positions=joint_positions.astype(float),
-                base_linear_velocity=v_WB[0:3].astype(float),
-                base_angular_velocity=v_WB[3:6].astype(float),
-                joint_velocities=joint_velocities.astype(float),
-            ),
-            soft_contacts_state=soft_contacts_state,
-        )
-
-        if not ode_state.valid(physics_model=model.physics_model):
+        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
+                | jaxsim.rbda.contacts.ViscoElasticContacts,
+            ):
+
+                contacts_params = js.contact.estimate_good_contact_parameters(
+                    model=model, standard_gravity=standard_gravity
+                )
+
+            else:
+                contacts_params = model.contact_model._parameters_class()
+
         return JaxSimModelData(
-            time_ns=time_ns,
             state=ode_state,
-            gravity=gravity.astype(float),
-            soft_contacts_params=soft_contacts_params,
+            gravity=gravity,
+            contacts_params=contacts_params,
             velocity_representation=velocity_representation,
         )
 
@@ -212,20 +243,21 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     # Extract quantities
     # ==================
 
-    def time(self) -> jtp.Float:
+    def standard_gravity(self) -> jtp.Float:
         """
-        Get the simulated time.
+        Get the standard gravity constant.
 
         Returns:
-            The simulated time in seconds.
+            The standard gravity constant.
         """
 
-        return self.time_ns.astype(float) / 1e9
+        return -self.gravity[2]
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def joint_positions(
         self,
-        model: jaxsim.api.model.JaxSimModel | None = None,
+        model: js.model.JaxSimModel | None = None,
         joint_names: tuple[str, ...] | None = None,
     ) -> jtp.Vector:
         """
@@ -250,22 +282,30 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         """
 
         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 self.valid(model=model):
+        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_names = joint_names if joint_names is not None else model.joint_names()
+        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[
-            jaxsim.api.joint.names_to_idxs(joint_names=joint_names, model=model)
-        ]
+        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: jaxsim.api.model.JaxSimModel | None = None,
+        model: js.model.JaxSimModel | None = None,
         joint_names: tuple[str, ...] | None = None,
     ) -> jtp.Vector:
         """
@@ -290,18 +330,26 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         """
 
         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 self.valid(model=model):
+        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_names = joint_names if joint_names is not None else model.joint_names()
+        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[
-            jaxsim.api.joint.names_to_idxs(joint_names=joint_names, model=model)
-        ]
+        return self.state.physics_model.joint_velocities[joint_idxs]
 
+    @js.common.named_scope
     @jax.jit
     def base_position(self) -> jtp.Vector:
         """
@@ -313,6 +361,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         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:
         """
@@ -325,29 +374,24 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             The base orientation.
         """
 
+        # Extract the base quaternion.
+        W_Q_B = self.state.physics_model.base_quaternion.squeeze()
+
         # Always normalize the quaternion to avoid numerical issues.
         # If the active scheme does not integrate the quaternion on its manifold,
         # we introduce a Baumgarte stabilization to let the quaternion converge to
         # a unit quaternion. In this case, it is not guaranteed that the quaternion
         # stored in the state is a unit quaternion.
-        base_unit_quaternion = (
-            self.state.physics_model.base_quaternion.squeeze()
-            / jnp.linalg.norm(self.state.physics_model.base_quaternion)
-        )
+        norm = jaxsim.math.safe_norm(W_Q_B)
+        W_Q_B = W_Q_B / (norm + jnp.finfo(float).eps * (norm == 0))
 
-        # Slice to convert quaternion wxyz -> xyzw
-        to_xyzw = np.array([1, 2, 3, 0])
-
-        return (
-            base_unit_quaternion
-            if not dcm
-            else jaxlie.SO3.from_quaternion_xyzw(
-                base_unit_quaternion[to_xyzw]
-            ).as_matrix()
+        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.MatrixJax:
+    def base_transform(self) -> jtp.Matrix:
         """
         Get the base transform.
 
@@ -365,6 +409,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ]
         )
 
+    @js.common.named_scope
     @jax.jit
     def base_velocity(self) -> jtp.Vector:
         """
@@ -394,9 +439,10 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             .astype(float)
         )
 
+    @js.common.named_scope
     @jax.jit
     def generalized_position(self) -> tuple[jtp.Matrix, jtp.Vector]:
-        """
+        r"""
         Get the generalized position
         :math:`\mathbf{q} = ({}^W \mathbf{H}_B, \mathbf{s}) \in \text{SO}(3) \times \mathbb{R}^n`.
 
@@ -406,10 +452,12 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.base_transform(), self.joint_positions()
 
+    @js.common.named_scope
     @jax.jit
     def generalized_velocity(self) -> jtp.Vector:
-        """
-        Get the generalized velocity
+        r"""
+        Get the generalized velocity.
+
         :math:`\boldsymbol{\nu} = (\boldsymbol{v}_{W,B};\, \boldsymbol{\omega}_{W,B};\, \mathbf{s}) \in \mathbb{R}^{6+n}`
 
         Returns:
@@ -426,11 +474,12 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     # Store quantities
     # ================
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def reset_joint_positions(
         self,
         positions: jtp.VectorLike,
-        model: jaxsim.api.model.JaxSimModel | None = None,
+        model: js.model.JaxSimModel | None = None,
         joint_names: tuple[str, ...] | None = None,
     ) -> Self:
         """
@@ -460,23 +509,26 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         if model is None:
             return replace(s=positions)
 
-        if not self.valid(model=model):
+        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_names = joint_names if joint_names is not None else model.joint_names()
+        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[
-                jaxsim.api.joint.names_to_idxs(joint_names=joint_names, model=model)
-            ].set(positions)
+            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: jaxsim.api.model.JaxSimModel | None = None,
+        model: js.model.JaxSimModel | None = None,
         joint_names: tuple[str, ...] | None = None,
     ) -> Self:
         """
@@ -506,18 +558,21 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         if model is None:
             return replace(ṡ=velocities)
 
-        if not self.valid(model=model):
+        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_names = joint_names if joint_names is not None else model.joint_names()
+        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[
-                jaxsim.api.joint.names_to_idxs(joint_names=joint_names, model=model)
-            ].set(velocities)
+            ṡ=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:
         """
@@ -541,6 +596,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ),
         )
 
+    @js.common.named_scope
     @jax.jit
     def reset_base_quaternion(self, base_quaternion: jtp.VectorLike) -> Self:
         """
@@ -553,19 +609,19 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             The updated `JaxSimModelData` object.
         """
 
-        base_quaternion = jnp.array(base_quaternion)
+        W_Q_B = jnp.array(base_quaternion, dtype=float)
+
+        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=jnp.atleast_1d(base_quaternion.squeeze()).astype(
-                        float
-                    )
-                )
+                physics_model=self.state.physics_model.replace(base_quaternion=W_Q_B)
             ),
         )
 
+    @js.common.named_scope
     @jax.jit
     def reset_base_pose(self, base_pose: jtp.MatrixLike) -> Self:
         """
@@ -582,14 +638,13 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         W_p_B = base_pose[0:3, 3]
 
-        to_wxyz = np.array([3, 0, 1, 2])
-        W_R_B: jaxlie.SO3 = jaxlie.SO3.from_matrix(base_pose[0:3, 0:3])  # noqa
-        W_Q_B = W_R_B.as_quaternion_xyzw()[to_wxyz]
+        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
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["velocity_representation"])
     def reset_base_linear_velocity(
         self,
@@ -613,11 +668,15 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.reset_base_velocity(
             base_velocity=jnp.hstack(
-                [linear_velocity.squeeze(), self.base_velocity()[3:6]]
+                [
+                    linear_velocity.squeeze(),
+                    self.base_velocity()[3:6],
+                ]
             ),
             velocity_representation=velocity_representation,
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["velocity_representation"])
     def reset_base_angular_velocity(
         self,
@@ -641,11 +700,15 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.reset_base_velocity(
             base_velocity=jnp.hstack(
-                [self.base_velocity()[0:3], angular_velocity.squeeze()]
+                [
+                    self.base_velocity()[0:3],
+                    angular_velocity.squeeze(),
+                ]
             ),
             velocity_representation=velocity_representation,
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["velocity_representation"])
     def reset_base_velocity(
         self,
@@ -691,8 +754,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         )
 
 
+@functools.partial(jax.jit, static_argnames=["velocity_representation", "base_rpy_seq"])
 def random_model_data(
-    model: jaxsim.api.model.JaxSimModel,
+    model: js.model.JaxSimModel,
     *,
     key: jax.Array | None = None,
     velocity_representation: VelRepr | None = None,
@@ -700,6 +764,18 @@ def random_model_data(
         jtp.FloatLike | Sequence[jtp.FloatLike],
         jtp.FloatLike | Sequence[jtp.FloatLike],
     ] = ((-1, -1, 0.5), 1.0),
+    base_rpy_bounds: tuple[
+        jtp.FloatLike | Sequence[jtp.FloatLike],
+        jtp.FloatLike | Sequence[jtp.FloatLike],
+    ] = (-jnp.pi, jnp.pi),
+    base_rpy_seq: str = "XYZ",
+    joint_pos_bounds: (
+        tuple[
+            jtp.FloatLike | Sequence[jtp.FloatLike],
+            jtp.FloatLike | Sequence[jtp.FloatLike],
+        ]
+        | None
+    ) = None,
     base_vel_lin_bounds: tuple[
         jtp.FloatLike | Sequence[jtp.FloatLike],
         jtp.FloatLike | Sequence[jtp.FloatLike],
@@ -712,6 +788,11 @@ 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.
@@ -721,19 +802,29 @@ def random_model_data(
         key: The random key.
         velocity_representation: The velocity representation to use.
         base_pos_bounds: The bounds for the base position.
+        base_rpy_bounds:
+            The bounds for the euler angles used to build the base orientation.
+        base_rpy_seq:
+            The sequence of axes for rotation (using `Rotation` from scipy).
+        joint_pos_bounds:
+            The bounds for the joint positions (reading the joint limits if None).
         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 = jax.random.split(key, num=6)
+    k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, num=7)
 
     p_min = jnp.array(base_pos_bounds[0], dtype=float)
     p_max = jnp.array(base_pos_bounds[1], dtype=float)
+    rpy_min = jnp.array(base_rpy_bounds[0], dtype=float)
+    rpy_max = jnp.array(base_rpy_bounds[1], dtype=float)
     v_min = jnp.array(base_vel_lin_bounds[0], dtype=float)
     v_max = jnp.array(base_vel_lin_bounds[1], dtype=float)
     ω_min = jnp.array(base_vel_ang_bounds[0], dtype=float)
@@ -749,7 +840,9 @@ def random_model_data(
         ),
     )
 
-    with random_data.mutable_context(mutability=Mutability.MUTABLE):
+    with random_data.mutable_context(
+        mutability=Mutability.MUTABLE, restore_after_exception=False
+    ):
 
         physics_model_state = random_data.state.physics_model
 
@@ -757,24 +850,76 @@ def random_model_data(
             key=k1, shape=(3,), minval=p_min, maxval=p_max
         )
 
-        physics_model_state.base_quaternion = jaxlie.SO3.from_rpy_radians(
-            *jax.random.uniform(key=k2, shape=(3,), minval=0, maxval=2 * jnp.pi)
-        ).as_quaternion_xyzw()[np.array([3, 0, 1, 2])]
-
-        physics_model_state.joint_positions = jaxsim.api.joint.random_joint_positions(
-            model=model, key=k3
+        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()
         )
 
-        physics_model_state.base_linear_velocity = jax.random.uniform(
-            key=k4, shape=(3,), minval=v_min, maxval=v_max
-        )
+        if model.number_of_joints() > 0:
 
-        physics_model_state.base_angular_velocity = jax.random.uniform(
-            key=k5, shape=(3,), minval=ω_min, maxval=ω_max
-        )
+            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
+            )
 
-        physics_model_state.joint_velocities = jax.random.uniform(
-            key=k6, shape=(model.dofs(),), 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],
+                )
+            )
         )
 
+        if contacts_params is None:
+
+            if isinstance(
+                model.contact_model,
+                jaxsim.rbda.contacts.SoftContacts
+                | jaxsim.rbda.contacts.ViscoElasticContacts,
+            ):
+
+                random_data = random_data.replace(
+                    contacts_params=js.contact.estimate_good_contact_parameters(
+                        model=model, standard_gravity=random_data.gravity
+                    ),
+                    validate=False,
+                )
+
+            else:
+                random_data = random_data.replace(
+                    contacts_params=model.contact_model._parameters_class(),
+                    validate=False,
+                )
+
     return random_data