jaxsim 0.4.3.dev64__py3-none-any.whl → 0.4.3.dev68__py3-none-any.whl

jaxsim/__init__.py

@@ -20,11 +20,6 @@ def _jnp_options() -> None:
         if jnp.empty(0, dtype=float).dtype != jnp.empty(0, dtype=np.float64).dtype:
             logging.warning("Failed to enable 64bit precision in JAX")
 
-    else:
-        logging.warning(
-            "Using 32bit precision in JaxSim is still experimental, please avoid to use variable step integrators."
-        )
-
 
 def _np_options() -> None:
     import numpy as np

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev64'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev64')
+__version__ = version = '0.4.3.dev68'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev68')

jaxsim/api/contact.py

@@ -131,8 +131,7 @@ def collidable_point_dynamics(
     Returns:
         The 6D force applied to each collidable point and additional data based on the contact model configured:
         - Soft: the material deformation rate.
-        - Rigid: no additional data.
-        - QuasiRigid: no additional data.
+        - Rigid: nothing.
 
     Note:
         The material deformation rate is always returned in the mixed frame
@@ -145,10 +144,6 @@ def collidable_point_dynamics(
     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 (
-        RelaxedRigidContacts,
-        RelaxedRigidContactsState,
-    )
     from jaxsim.rbda.contacts.rigid import RigidContacts, RigidContactsState
     from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 
@@ -195,27 +190,6 @@ def collidable_point_dynamics(
 
             aux_data = dict()
 
-        case RelaxedRigidContacts():
-            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,
-                model=model,
-                data=data,
-                link_forces=link_forces,
-            )
-
-            aux_data = dict()
-
         case _:
             raise ValueError(f"Invalid contact model {model.contact_model}")
 

jaxsim/api/data.py

@@ -39,9 +39,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     contacts_params: jaxsim.rbda.ContactsParams = dataclasses.field(repr=False)
 
     time_ns: jtp.Int = dataclasses.field(
-        default_factory=lambda: jnp.array(
-            0, dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32
-        ),
+        default_factory=lambda: jnp.array(0, dtype=jnp.uint64)
     )
 
     def __hash__(self) -> int:
@@ -174,14 +172,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         )
 
         time_ns = (
-            jnp.array(
-                time * 1e9,
-                dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32,
-            )
+            jnp.array(time * 1e9, dtype=jnp.uint64)
             if time is not None
-            else jnp.array(
-                0, dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32
-            )
+            else jnp.array(0, dtype=jnp.uint64)
         )
 
         if isinstance(model.contact_model, SoftContacts):
@@ -593,18 +586,16 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             The updated `JaxSimModelData` object.
         """
 
-        W_Q_B = jnp.array(base_quaternion, dtype=float)
-
-        W_Q_B = jax.lax.select(
-            pred=jnp.allclose(jnp.linalg.norm(W_Q_B), 1.0, atol=1e-6, rtol=0.0),
-            on_true=W_Q_B,
-            on_false=W_Q_B / jnp.linalg.norm(W_Q_B),
-        )
+        base_quaternion = jnp.array(base_quaternion)
 
         return self.replace(
             validate=True,
             state=self.state.replace(
-                physics_model=self.state.physics_model.replace(base_quaternion=W_Q_B)
+                physics_model=self.state.physics_model.replace(
+                    base_quaternion=jnp.atleast_1d(base_quaternion.squeeze()).astype(
+                        float
+                    )
+                )
             ),
         )
 
@@ -746,13 +737,6 @@ def random_model_data(
         jtp.FloatLike | Sequence[jtp.FloatLike],
         jtp.FloatLike | Sequence[jtp.FloatLike],
     ] = ((-1, -1, 0.5), 1.0),
-    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],
@@ -778,8 +762,6 @@ def random_model_data(
         key: The random key.
         velocity_representation: The velocity representation to use.
         base_pos_bounds: The bounds for the base position.
-        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.
@@ -824,19 +806,8 @@ def random_model_data(
         ).wxyz
 
         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_positions = js.joint.random_joint_positions(
+                model=model, key=k3
             )
 
             physics_model_state.joint_velocities = jax.random.uniform(

jaxsim/api/joint.py

@@ -180,77 +180,17 @@ def random_joint_positions(
 
     Args:
         model: The model to consider.
-        joint_names: The names of the considered joints (all if None).
-        key: The random key (initialized from seed 0 if None).
-
-    Note:
-        If the joint range or revolute joints is larger than 2π, their joint positions
-        will be sampled from an interval of size 2π.
+        joint_names: The names of the joints.
+        key: The random key.
 
     Returns:
         The random joint positions.
     """
 
-    # Consider the key corresponding to a zero seed if it was not passed.
     key = key if key is not None else jax.random.PRNGKey(seed=0)
 
-    # Get the joint limits parsed from the model description.
     s_min, s_max = position_limits(model=model, joint_names=joint_names)
 
-    # Get the joint indices.
-    # Note that it will trigger an exception if the given `joint_names` are not valid.
-    joint_names = joint_names if joint_names is not None else model.joint_names()
-    joint_indices = names_to_idxs(model=model, joint_names=joint_names)
-
-    from jaxsim.parsers.descriptions.joint import JointType
-
-    # Filter for revolute joints.
-    is_revolute = jnp.where(
-        jnp.array(model.kin_dyn_parameters.joint_model.joint_types[1:])[joint_indices]
-        == JointType.Revolute,
-        True,
-        False,
-    )
-
-    # Shorthand for π.
-    π = jnp.pi
-
-    # Filter for revolute with full range (or continuous).
-    is_revolute_full_range = jnp.logical_and(is_revolute, s_max - s_min >= 2 * π)
-
-    # Clip the lower limit to -π if the joint range is larger than [-π, π].
-    s_min = jnp.where(
-        jnp.logical_and(
-            is_revolute_full_range, jnp.logical_and(s_min <= -π, s_max >= π)
-        ),
-        -π,
-        s_min,
-    )
-
-    # Clip the upper limit to +π if the joint range is larger than [-π, π].
-    s_max = jnp.where(
-        jnp.logical_and(
-            is_revolute_full_range, jnp.logical_and(s_min <= -π, s_max >= π)
-        ),
-        π,
-        s_max,
-    )
-
-    # Shift the lower limit if the upper limit is smaller than +π.
-    s_min = jnp.where(
-        jnp.logical_and(is_revolute_full_range, s_max < π),
-        s_max - 2 * π,
-        s_min,
-    )
-
-    # Shift the upper limit if the lower limit is larger than -π.
-    s_max = jnp.where(
-        jnp.logical_and(is_revolute_full_range, s_min > -π),
-        s_min + 2 * π,
-        s_max,
-    )
-
-    # Sample the joint positions.
     s_random = jax.random.uniform(
         minval=s_min,
         maxval=s_max,

jaxsim/api/model.py

@@ -1931,22 +1931,11 @@ def step(
         ),
     )
 
-    tf_ns = t0_ns + jnp.array(dt * 1e9, dtype=t0_ns.dtype)
-    tf_ns = jnp.where(tf_ns >= t0_ns, tf_ns, jnp.array(0, dtype=t0_ns.dtype))
-
-    jax.lax.cond(
-        pred=tf_ns < t0_ns,
-        true_fun=lambda: jax.debug.print(
-            "The simulation time overflowed, resetting simulation time to 0."
-        ),
-        false_fun=lambda: None,
-    )
-
     data_tf = (
         # Store the new state of the model and the new time.
         data.replace(
             state=state_tf,
-            time_ns=tf_ns,
+            time_ns=t0_ns + jnp.array(dt * 1e9).astype(jnp.uint64),
         )
     )
 

jaxsim/api/ode.py

@@ -175,15 +175,17 @@ def system_velocity_dynamics(
             forces=W_f_Li_terrain,
             additive=True,
         )
-
-        # Get the link forces in inertial representation
+    # Get the link forces in the data representation
+    with references.switch_velocity_representation(data.velocity_representation):
         f_L_total = references.link_forces(model=model, data=data)
 
-        v̇_WB, s̈ = system_acceleration(
-            model=model, data=data, joint_forces=joint_forces, link_forces=f_L_total
-        )
+    # The following method always returns the inertial-fixed acceleration, and expects
+    # the link_forces expressed in the inertial frame.
+    W_v̇_WB, s̈ = system_acceleration(
+        model=model, data=data, joint_forces=joint_forces, link_forces=f_L_total
+    )
 
-    return v̇_WB, s̈, aux_data
+    return W_v̇_WB, s̈, aux_data
 
 
 def system_acceleration(
@@ -194,7 +196,7 @@ def system_acceleration(
     link_forces: jtp.MatrixLike | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
-    Compute the system acceleration in the active representation.
+    Compute the system acceleration in inertial-fixed representation.
 
     Args:
         model: The model to consider.
@@ -204,7 +206,7 @@ def system_acceleration(
             The 6D forces to apply to the links expressed in the same representation of data.
 
     Returns:
-        A tuple containing the base 6D acceleration in in the active representation
+        A tuple containing the base 6D acceleration in inertial-fixed representation
         and the joint accelerations.
     """
 
@@ -270,15 +272,18 @@ def system_acceleration(
     )
 
     # - Joint accelerations: s̈ ∈ ℝⁿ
-    # - Base acceleration: v̇_WB ∈ ℝ⁶
-    v̇_WB, s̈ = js.model.forward_dynamics_aba(
-        model=model,
-        data=data,
-        joint_forces=references.joint_force_references(model=model),
-        link_forces=references.link_forces(model=model, data=data),
-    )
-
-    return v̇_WB, s̈
+    # - Base inertial-fixed acceleration: W_v̇_WB = (W_p̈_B, W_ω̇_B) ∈ ℝ⁶
+    with (
+        data.switch_velocity_representation(velocity_representation=VelRepr.Inertial),
+        references.switch_velocity_representation(VelRepr.Inertial),
+    ):
+        W_v̇_WB, s̈ = js.model.forward_dynamics_aba(
+            model=model,
+            data=data,
+            joint_forces=references.joint_force_references(),
+            link_forces=references.link_forces(),
+        )
+    return W_v̇_WB, s̈
 
 
 @jax.jit
@@ -348,7 +353,7 @@ def system_dynamics(
         corresponding derivative, and the dictionary of auxiliary data returned
         by the system dynamics evaluation.
     """
-    from jaxsim.rbda.contacts.relaxed_rigid import RelaxedRigidContacts
+
     from jaxsim.rbda.contacts.rigid import RigidContacts
     from jaxsim.rbda.contacts.soft import SoftContacts
 
@@ -366,7 +371,7 @@ def system_dynamics(
         case SoftContacts():
             ode_state_kwargs["tangential_deformation"] = aux_dict["m_dot"]
 
-        case RigidContacts() | RelaxedRigidContacts():
+        case RigidContacts():
             pass
 
         case _:

jaxsim/api/ode_data.py

@@ -6,10 +6,6 @@ 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 (
-    RelaxedRigidContacts,
-    RelaxedRigidContactsState,
-)
 from jaxsim.rbda.contacts.rigid import RigidContacts, RigidContactsState
 from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 from jaxsim.utils import JaxsimDataclass
@@ -177,10 +173,6 @@ class ODEState(JaxsimDataclass):
                 )
             case RigidContacts():
                 contact = RigidContactsState.build()
-
-            case RelaxedRigidContacts():
-                contact = RelaxedRigidContactsState.build()
-
             case _:
                 raise ValueError("Unable to determine contact state class prefix.")
 
@@ -224,9 +216,7 @@ class ODEState(JaxsimDataclass):
 
         # Get the contact model from the `JaxSimModel`.
         match contact:
-            case (
-                SoftContactsState() | RigidContactsState() | RelaxedRigidContactsState()
-            ):
+            case SoftContactsState() | RigidContactsState():
                 pass
             case None:
                 contact = SoftContactsState.zero(model=model)

jaxsim/integrators/common.py

@@ -497,7 +497,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         b: jtp.Matrix,
         c: jtp.Vector,
         index_of_solution: jtp.IntLike = 0,
-    ) -> tuple[bool, int | None]:
+    ) -> [bool, int | None]:
         """
         Check if the Butcher tableau supports the FSAL (first-same-as-last) property.
 

jaxsim/math/inertia.py

@@ -45,7 +45,7 @@ class Inertia:
             M (jtp.Matrix): The 6x6 inertia matrix.
 
         Returns:
-            tuple[jtp.Float, jtp.Vector, jtp.Matrix]: A tuple containing mass, center of mass (3D), and inertia matrix (3x3).
+            Tuple[jtp.Float, jtp.Vector, jtp.Matrix]: A tuple containing mass, center of mass (3D), and inertia matrix (3x3).
 
         Raises:
             ValueError: If the input matrix M has an unexpected shape.

jaxsim/mujoco/loaders.py

@@ -211,7 +211,7 @@ class RodModelToMjcf:
         joints_dict = {j.name: j for j in rod_model.joints()}
 
         # Convert all the joints not considered to fixed joints.
-        for joint_name in {j.name for j in rod_model.joints()} - considered_joints:
+        for joint_name in set(j.name for j in rod_model.joints()) - considered_joints:
             joints_dict[joint_name].type = "fixed"
 
         # Convert the ROD model to URDF.
@@ -289,10 +289,10 @@ class RodModelToMjcf:
         mj_model = mj.MjModel.from_xml_string(xml=urdf_string, assets=assets)
 
         # Get the joint names.
-        mj_joint_names = {
+        mj_joint_names = set(
             mj.mj_id2name(mj_model, mj.mjtObj.mjOBJ_JOINT, idx)
             for idx in range(mj_model.njnt)
-        }
+        )
 
         # Check that the Mujoco model only has the considered joints.
         if mj_joint_names != considered_joints:

jaxsim/parsers/kinematic_graph.py

@@ -394,7 +394,7 @@ class KinematicGraph(Sequence[LinkDescription]):
             return copy.deepcopy(self)
 
         # Check if all considered joints are part of the full kinematic graph
-        if len(set(considered_joints) - {j.name for j in full_graph.joints}) != 0:
+        if len(set(considered_joints) - set(j.name for j in full_graph.joints)) != 0:
             extra_j = set(considered_joints) - {j.name for j in full_graph.joints}
             msg = f"Not all joints to consider are part of the graph ({{{extra_j}}})"
             raise ValueError(msg)
@@ -536,8 +536,8 @@ class KinematicGraph(Sequence[LinkDescription]):
             root_link_name=full_graph.root.name,
         )
 
-        assert {f.name for f in self.frames}.isdisjoint(
-            {f.name for f in unconnected_frames + reduced_frames}
+        assert set(f.name for f in self.frames).isdisjoint(
+            set(f.name for f in unconnected_frames + reduced_frames)
         )
 
         for link in unconnected_links:

jaxsim/parsers/rod/parser.py

@@ -223,7 +223,7 @@ def extract_model_data(
             child=links_dict[j.child],
             jtype=utils.joint_to_joint_type(joint=j),
             axis=(
-                np.array(j.axis.xyz.xyz, dtype=float)
+                np.array(j.axis.xyz.xyz)
                 if j.axis is not None
                 and j.axis.xyz is not None
                 and j.axis.xyz.xyz is not None
@@ -232,43 +232,39 @@ def extract_model_data(
             pose=j.pose.transform() if j.pose is not None else np.eye(4),
             initial_position=0.0,
             position_limit=(
-                float(
-                    j.axis.limit.lower
-                    if j.axis is not None
-                    and j.axis.limit is not None
-                    and j.axis.limit.lower is not None
-                    else jnp.finfo(float).min
+                (
+                    float(j.axis.limit.lower)
+                    if j.axis is not None and j.axis.limit is not None
+                    else np.finfo(float).min
                 ),
-                float(
-                    j.axis.limit.upper
-                    if j.axis is not None
-                    and j.axis.limit is not None
-                    and j.axis.limit.upper is not None
-                    else jnp.finfo(float).max
+                (
+                    float(j.axis.limit.upper)
+                    if j.axis is not None and j.axis.limit is not None
+                    else np.finfo(float).max
                 ),
             ),
-            friction_static=float(
+            friction_static=(
                 j.axis.dynamics.friction
                 if j.axis is not None
                 and j.axis.dynamics is not None
                 and j.axis.dynamics.friction is not None
                 else 0.0
             ),
-            friction_viscous=float(
+            friction_viscous=(
                 j.axis.dynamics.damping
                 if j.axis is not None
                 and j.axis.dynamics is not None
                 and j.axis.dynamics.damping is not None
                 else 0.0
             ),
-            position_limit_damper=float(
+            position_limit_damper=(
                 j.axis.limit.dissipation
                 if j.axis is not None
                 and j.axis.limit is not None
                 and j.axis.limit.dissipation is not None
                 else 0.0
             ),
-            position_limit_spring=float(
+            position_limit_spring=(
                 j.axis.limit.stiffness
                 if j.axis is not None
                 and j.axis.limit is not None
@@ -277,7 +273,7 @@ def extract_model_data(
             ),
         )
         for j in sdf_model.joints()
-        if j.type in {"revolute", "continuous", "prismatic", "fixed"}
+        if j.type in {"revolute", "prismatic", "fixed"}
         and j.parent != "world"
         and j.child in links_dict.keys()
     ]

jaxsim/rbda/contacts/rigid.py

@@ -9,6 +9,7 @@ import jax_dataclasses
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
+from jaxsim import math
 from jaxsim.api.common import ModelDataWithVelocityRepresentation, VelRepr
 from jaxsim.terrain import FlatTerrain, Terrain
 
@@ -271,17 +272,9 @@ class RigidContacts(ContactModel):
             link_forces=link_forces,
         )
 
-        with (
-            references.switch_velocity_representation(VelRepr.Mixed),
-            data.switch_velocity_representation(VelRepr.Mixed),
-        ):
-            BW_ν̇_free = jnp.hstack(
-                js.ode.system_acceleration(
-                    model=model,
-                    data=data,
-                    joint_forces=references.joint_force_references(model=model),
-                    link_forces=references.link_forces(model=model, data=data),
-                )
+        with references.switch_velocity_representation(VelRepr.Mixed):
+            BW_ν̇_free = RigidContacts._compute_mixed_nu_dot_free(
+                model, data, references=references
             )
 
         free_contact_acc = RigidContacts._linear_acceleration_of_collidable_points(
@@ -387,6 +380,43 @@ class RigidContacts(ContactModel):
         n_constraints = 6 * n_collidable_points
         return jnp.zeros(shape=(n_constraints,))
 
+    @staticmethod
+    def _compute_mixed_nu_dot_free(
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        references: js.references.JaxSimModelReferences | None = None,
+    ) -> jtp.Array:
+        references = (
+            references
+            if references is not None
+            else js.references.JaxSimModelReferences.zero(model=model, data=data)
+        )
+
+        with (
+            data.switch_velocity_representation(VelRepr.Mixed),
+            references.switch_velocity_representation(VelRepr.Mixed),
+        ):
+            BW_v_WB = data.base_velocity()
+            W_ṗ_B, W_ω_WB = jnp.split(BW_v_WB, 2)
+            W_v̇_WB, s̈ = js.ode.system_acceleration(
+                model=model,
+                data=data,
+                joint_forces=references.joint_force_references(model=model),
+                link_forces=references.link_forces(model=model, data=data),
+            )
+
+        # Convert the inertial-fixed base acceleration to a mixed base acceleration.
+        W_H_B = data.base_transform()
+        W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
+        BW_X_W = math.Adjoint.from_transform(W_H_BW, inverse=True)
+        term1 = BW_X_W @ W_v̇_WB
+        term2 = jnp.zeros(6).at[0:3].set(jnp.cross(W_ṗ_B, W_ω_WB))
+        BW_v̇_WB = term1 - term2
+
+        BW_ν̇ = jnp.hstack([BW_v̇_WB, s̈])
+
+        return BW_ν̇
+
     @staticmethod
     def _linear_acceleration_of_collidable_points(
         model: js.model.JaxSimModel,

jaxsim/terrain/terrain.py

@@ -46,82 +46,66 @@ class Terrain(abc.ABC):
 @jax_dataclasses.pytree_dataclass
 class FlatTerrain(Terrain):
 
-    _height: float = dataclasses.field(default=0.0, kw_only=True)
+    z: float = dataclasses.field(default=0.0, kw_only=True)
 
     @staticmethod
     def build(height: jtp.FloatLike) -> FlatTerrain:
 
-        return FlatTerrain(_height=float(height))
+        return FlatTerrain(z=float(height))
 
     def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
 
-        return jnp.array(self._height, dtype=float)
-
-    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
-
-        return jnp.array([0.0, 0.0, 1.0], dtype=float)
+        return jnp.array(self.z, dtype=float)
 
     def __hash__(self) -> int:
 
-        return hash(self._height)
+        return hash(self.z)
 
     def __eq__(self, other: FlatTerrain) -> bool:
 
         if not isinstance(other, FlatTerrain):
             return False
 
-        return self._height == other._height
+        return self.z == other.z
 
 
 @jax_dataclasses.pytree_dataclass
 class PlaneTerrain(FlatTerrain):
 
-    _normal: tuple[float, float, float] = jax_dataclasses.field(
+    plane_normal: tuple[float, float, float] = jax_dataclasses.field(
         default=(0.0, 0.0, 1.0), kw_only=True
     )
 
     @staticmethod
-    def build(height: jtp.FloatLike = 0.0, *, normal: jtp.VectorLike) -> PlaneTerrain:
+    def build(
+        plane_normal: jtp.VectorLike, plane_height_over_origin: jtp.FloatLike = 0.0
+    ) -> PlaneTerrain:
         """
         Create a PlaneTerrain instance with a specified plane normal vector.
 
         Args:
-            normal: The normal vector of the terrain plane.
-            height: The height of the plane over the origin.
+            plane_normal: The normal vector of the terrain plane.
+            plane_height_over_origin: The height of the plane over the origin.
 
         Returns:
             PlaneTerrain: A PlaneTerrain instance.
         """
 
-        normal = jnp.array(normal, dtype=float)
-        height = jnp.array(height, dtype=float)
+        plane_normal = jnp.array(plane_normal, dtype=float)
+        plane_height_over_origin = jnp.array(plane_height_over_origin, dtype=float)
 
-        if normal.shape != (3,):
+        if plane_normal.shape != (3,):
             msg = "Expected a 3D vector for the plane normal, got '{}'."
-            raise ValueError(msg.format(normal.shape))
+            raise ValueError(msg.format(plane_normal.shape))
 
         # Make sure that the plane normal is a unit vector.
-        normal = normal / jnp.linalg.norm(normal)
+        plane_normal = plane_normal / jnp.linalg.norm(plane_normal)
 
         return PlaneTerrain(
-            _height=height.item(),
-            _normal=tuple(normal.tolist()),
+            z=float(plane_height_over_origin),
+            plane_normal=tuple(plane_normal.tolist()),
         )
 
-    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
-        """
-        Compute the normal vector of the terrain at a specific (x, y) location.
-
-        Args:
-            x: The x-coordinate of the location.
-            y: The y-coordinate of the location.
-
-        Returns:
-            The normal vector of the terrain surface at the specified location.
-        """
-
-        return jnp.array(self._normal, dtype=float)
-
     def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
         """
         Compute the height of the terrain at a specific (x, y) location on a plane.
@@ -139,10 +123,10 @@ class PlaneTerrain(FlatTerrain):
         # The height over the origin: -D/C
 
         # Get the plane equation coefficients from the terrain normal.
-        A, B, C = self._normal
+        A, B, C = self.plane_normal
 
         # Compute the final coefficient D considering the terrain height.
-        D = -C * self._height
+        D = -C * self.z
 
         # Invert the plane equation to get the height at the given (x, y) coordinates.
         return jnp.array(-(A * x + B * y + D) / C).astype(float)
@@ -153,9 +137,9 @@ class PlaneTerrain(FlatTerrain):
 
         return hash(
             (
-                hash(self._height),
+                hash(self.z),
                 HashedNumpyArray.hash_of_array(
-                    array=jnp.array(self._normal, dtype=float)
+                    array=jnp.array(self.plane_normal, dtype=float)
                 ),
             )
         )
@@ -166,10 +150,10 @@ class PlaneTerrain(FlatTerrain):
             return False
 
         if not (
-            np.allclose(self._height, other._height)
+            np.allclose(self.z, other.z)
             and np.allclose(
-                np.array(self._normal, dtype=float),
-                np.array(other._normal, dtype=float),
+                np.array(self.plane_normal, dtype=float),
+                np.array(other.plane_normal, dtype=float),
             )
         ):
             return False

jaxsim/typing.py

@@ -16,7 +16,7 @@ Int = Scalar
 Bool = Scalar
 Float = Scalar
 
-PyTree: object = (
+PyTree = (
     dict[Hashable, TypeVar("PyTree")]
     | list[TypeVar("PyTree")]
     | tuple[TypeVar("PyTree")]

jaxsim/utils/jaxsim_dataclass.py

@@ -135,10 +135,9 @@ class JaxsimDataclass(abc.ABC):
         """
 
         return tuple(
-            map(
-                lambda leaf: getattr(leaf, "shape", None),
-                jax.tree_util.tree_leaves(tree),
-            )
+            leaf.shape if hasattr(leaf, "shape") else None
+            for leaf in jax.tree_util.tree_leaves(tree)
+            if hasattr(leaf, "shape")
         )
 
     @staticmethod
@@ -155,10 +154,9 @@ class JaxsimDataclass(abc.ABC):
         """
 
         return tuple(
-            map(
-                lambda leaf: getattr(leaf, "dtype", None),
-                jax.tree_util.tree_leaves(tree),
-            )
+            leaf.dtype if hasattr(leaf, "dtype") else None
+            for leaf in jax.tree_util.tree_leaves(tree)
+            if hasattr(leaf, "dtype")
         )
 
     @staticmethod
@@ -174,10 +172,9 @@ class JaxsimDataclass(abc.ABC):
         """
 
         return tuple(
-            map(
-                lambda leaf: getattr(leaf, "weak_type", None),
-                jax.tree_util.tree_leaves(tree),
-            )
+            leaf.weak_type if hasattr(leaf, "weak_type") else False
+            for leaf in jax.tree_util.tree_leaves(tree)
+            if hasattr(leaf, "weak_type")
         )
 
     @staticmethod

jaxsim/utils/wrappers.py

@@ -110,7 +110,7 @@ class HashedNumpyArray:
             return np.allclose(
                 self.array,
                 other.array,
-                **(dict(atol=self.precision) if self.precision is not None else {}),
+                **({dict(atol=self.precision)} if self.precision is not None else {}),
             )
 
         return hash(self) == hash(other)

{jaxsim-0.4.3.dev64.dist-info → jaxsim-0.4.3.dev68.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev64
+Version: 0.4.3.dev68
 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>
@@ -61,7 +61,6 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: coloredlogs
 Requires-Dist: jax>=0.4.13
-Requires-Dist: jaxopt>=0.8.0
 Requires-Dist: jaxlib>=0.4.13
 Requires-Dist: jaxlie>=1.3.0
 Requires-Dist: jax-dataclasses>=1.4.0

{jaxsim-0.4.3.dev64.dist-info → jaxsim-0.4.3.dev68.dist-info}/RECORD

@@ -1,29 +1,29 @@
-jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=lLNskxtfHW1HqvnLRuhux3LlK89fMiZFUWknSYopw7k,426
+jaxsim/__init__.py,sha256=ixsS4dYMPex2wOUUp_rkPnwrPhYzkRh1xO_YuMj3Cr4,2626
+jaxsim/_version.py,sha256=XDf5LPSlhAhH48AO29kysLP_4FTR5VWOpS0LrK5RSfo,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
-jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
+jaxsim/typing.py,sha256=IbFx3UkEXi-cm7UBqMPi58rJAFV_HbZ9E_K4JwfNvVM,753
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=m-p3EJDhpnMTlXKplfbZE_aH9NqX_VyLlAE3vUhc6l4,13642
 jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
-jaxsim/api/contact.py,sha256=C_PgMjWYYiqpA7Oz3IxHeFgrp855-xG6AQr6Ze98CtI,21863
-jaxsim/api/data.py,sha256=mFUw2mj8AIXduW6HnkGN7eooZHfJhwnWbtYZfLF6gk4,28206
+jaxsim/api/contact.py,sha256=HyEAjF7BySDDOlRahN0l7V15IPB0HPXuoM0twamuEW0,20913
+jaxsim/api/data.py,sha256=CUh9lvhVk3_clNQ26BUBGpjvFSsK_PrVWVMEWpMdHRM,27206
 jaxsim/api/frame.py,sha256=KS8A5wRfjxhe9NgcVo2QA516iP5zky7UVnWxG7nTa7c,12911
-jaxsim/api/joint.py,sha256=lksT1Doxz2jknHyhb4ls20z6f6dofpZSzBJtVacZXAE,7129
+jaxsim/api/joint.py,sha256=L81bQe-noPT6_54KOSF7KBjRmEPAS433ULn2EcXI8vI,5115
 jaxsim/api/kin_dyn_parameters.py,sha256=CcfSg5Mc8qb1mZeMQ4AK_ffZIsK5yOl7tu397pFhcDA,29369
 jaxsim/api/link.py,sha256=qPRtc8qqMRjZxUCZYXJMygbB6huDXBfIT1b1b8Durkw,18631
-jaxsim/api/model.py,sha256=K0q8-j-04f6B3MEXsctDGtWiuWlN3HbDrsS7zoPYStk,65871
-jaxsim/api/ode.py,sha256=VuOLvCFoyGLmhNf2vFP5BI9BAPz78V_RW5tJ4hrizsw,13041
-jaxsim/api/ode_data.py,sha256=7RSoBhfCJdP6P9InQbDwdBVpClPMMuetewI-6AWm-_0,20276
+jaxsim/api/model.py,sha256=HXoqCtQ3KStGoxhgvFm8P_Sc-lbEM4l5No2MoHzNlOk,65558
+jaxsim/api/ode.py,sha256=Vb2sN4zwpXnaJDD9-ziz2qvfmfa4jvIQ0fONbBIRGmU,13368
+jaxsim/api/ode_data.py,sha256=U7F6TL6bENAxpQQl4PupPoDG7d7VfTTFqDAs3xwu6Hs,20003
 jaxsim/api/references.py,sha256=XOVKuQXRmjPoP-T5JWGSbqIGX5DzOkeGafqRpj0ZQEM,20771
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
-jaxsim/integrators/common.py,sha256=XIrJVJDO0ldaZ93WgoGNlFoRvazsRJTpO3DrK9kIXqM,20437
+jaxsim/integrators/common.py,sha256=ntjflaV3qWaFH_E65pAGZ6QipdnFsgQDasKtIKpxTe4,20432
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
 jaxsim/integrators/variable_step.py,sha256=5StkFh9oQba34zlkIoXG2fUN78gbxkHePWbrpQ-QZOI,21274
 jaxsim/math/__init__.py,sha256=8oPITEoGwgRcOeG8KxtqxPQ8b5uku1HNRMokpCoi9Tc,352
 jaxsim/math/adjoint.py,sha256=o1FCipkGwPtMbN2gFNIyUV8ADF3TX5fxElpTEXK0bIs,4377
 jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
-jaxsim/math/inertia.py,sha256=01hz6wMFreN2jBA0rVoBS1YMVh77KvwuzXSOpI3pxNk,1614
+jaxsim/math/inertia.py,sha256=_hNpoeyEpAGr9ExDQJjckbjhk39luJFF-jv0SKqefnQ,1614
 jaxsim/math/joint_model.py,sha256=EzAveaG5B6ZnCFNUzN30KEQUVesd83lfWXJarYR-kUw,9989
 jaxsim/math/quaternion.py,sha256=_WA7W3iv7px83sWO1V1n0-J78hqAlO4SL1-jofE-UZ4,4754
 jaxsim/math/rotation.py,sha256=k-nwT79zmWrys3NNAB-lGWxat7Kqm_6JnFRoimJ8rBg,2156
@@ -31,18 +31,18 @@ jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
 jaxsim/math/transform.py,sha256=KXzQgOnCfAtbXCwxhplpJ3F0JT3oEyeLVby1_uRAryQ,2892
 jaxsim/mujoco/__init__.py,sha256=Zo5GAlN1DYKvX8s1hu1j6HntKIbBMLB9Puv9ouaNAZ8,158
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
-jaxsim/mujoco/loaders.py,sha256=_8Af_5Yo0-lWHE-46BBMcrqSJnDNxr3peyc519DExtA,25322
+jaxsim/mujoco/loaders.py,sha256=XB-fgXuWMTFiaand5MZlLFQ5__Sh8MK5CJsxIU34MBk,25328
 jaxsim/mujoco/model.py,sha256=AQksXemXWACJ3yvefV2G5HLwwBU9ISoJrOD1wlxdY5w,16386
 jaxsim/mujoco/visualizer.py,sha256=T1vU-w4NKSmgEkZ0FqVcGmIvYrYO0len2UBSsU4MOZ0,6978
 jaxsim/parsers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxsim/parsers/kinematic_graph.py,sha256=wT2bgaCS8VQJTHy2H9sENkVPDOiMkRikxEF1t_WaahQ,34748
+jaxsim/parsers/kinematic_graph.py,sha256=KijMWKyhTLKSNUmOOk4sYQMgPh_OkA_brncL7gBRHaY,34757
 jaxsim/parsers/descriptions/__init__.py,sha256=PbIlunVfb59pB5jSX97YVpMAANRZPRkJ0X-hS14rzv4,221
 jaxsim/parsers/descriptions/collision.py,sha256=BQeIG-TKi4SVny23w6riDrQ5itC6VRwEMBX6HgAXHxA,3973
 jaxsim/parsers/descriptions/joint.py,sha256=VSb6C0FBBKMqwrHBKfc-Bbn4rl_J0RzUxMQlhIEvOPM,5185
 jaxsim/parsers/descriptions/link.py,sha256=Eh0W5qL7_Uw0GV-BkNKXhm9Q2dRTfIWCX5D-87zQkxA,3711
 jaxsim/parsers/descriptions/model.py,sha256=I2Vsbv8Josl4Le7b5rIvhqA2k9Bbv5JxMqwytayxds0,9833
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
-jaxsim/parsers/rod/parser.py,sha256=9EigYv2oGn4bfIY1q0Cd_55yVKfN2rXP_MuZSZqGxYM,13681
+jaxsim/parsers/rod/parser.py,sha256=HskeCqDsbtwH2BDk3vfxvx391wUTVGLaUXNvBrdNo-4,13486
 jaxsim/parsers/rod/utils.py,sha256=5DsF3OeePZGidOJ5GiFSZx-51uIdnFvMW9EK6SgOW6Q,5698
 jaxsim/rbda/__init__.py,sha256=H7DhXpxkPOi9lpUvg31IMHFfRafke1UoJLc5GQIdyhA,387
 jaxsim/rbda/aba.py,sha256=w7ciyxB0IsmueatT0C7PcBQEl9dyiH9oqJgIi3xeTUE,8983
@@ -54,17 +54,16 @@ jaxsim/rbda/rnea.py,sha256=CLfqs9XFVaD-hvkLABshDAfdw5bm_AMV3UVAQ_IvURQ,7542
 jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
 jaxsim/rbda/contacts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxsim/rbda/contacts/common.py,sha256=VwAs742futAmLnDgbaOuLzNDBFiKDfYItdEZ4UcFgzE,2467
-jaxsim/rbda/contacts/relaxed_rigid.py,sha256=9YkPLbK6Kk0wPkuj47r7NBqY2tARyJsiCbrvDlOWHSI,12700
-jaxsim/rbda/contacts/rigid.py,sha256=fbZk7sC6YOnTs_tzQRfsyBpHyT22XF-wB-EvOSZmhos,14746
+jaxsim/rbda/contacts/rigid.py,sha256=8Vbnxng-ERZ5ka_eZGIBuhBDr2PNjc7m-Or255AfEw4,15862
 jaxsim/rbda/contacts/soft.py,sha256=_wvb5iZDjGcVg6rNQelN4LZN7qSC2NIp0HdKvZmlGfk,15647
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
-jaxsim/terrain/terrain.py,sha256=xUQg47yGxIOcTkLPbnO3sruEGBhoCd16j1evTGlmNjI,5010
+jaxsim/terrain/terrain.py,sha256=ctyNANIFSM3tZmamprjaEDcWgUSP0oNJbmT1zw9RjPs,4565
 jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
-jaxsim/utils/jaxsim_dataclass.py,sha256=FSiUvdnq4Y1T9Jaa_mw4ZBQJe8H7deLr3Kupxtlh4iI,11322
+jaxsim/utils/jaxsim_dataclass.py,sha256=5xJbY0G8d7C0OTNIW9T4vQxiDak6TGZT9gpNOvRykFI,11373
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
-jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev64.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev64.dist-info/METADATA,sha256=0-JS1eJjFMSaMzwqbCSpWYU2GcrZkxT1LBDo7lhWICo,17276
-jaxsim-0.4.3.dev64.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-jaxsim-0.4.3.dev64.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev64.dist-info/RECORD,,
+jaxsim/utils/wrappers.py,sha256=JhLUh1g8iU-lhjbuZRfkscPZhYlLCOorVM2Xl3ulRBI,4054
+jaxsim-0.4.3.dev68.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev68.dist-info/METADATA,sha256=IrZMXHUptvvLA5YgloveNIge4OdEBjT-DxhdHBrn_WM,17247
+jaxsim-0.4.3.dev68.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+jaxsim-0.4.3.dev68.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev68.dist-info/RECORD,,

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -1,384 +0,0 @@
-from __future__ import annotations
-
-import dataclasses
-from typing import Any
-
-import jax
-import jax.numpy as jnp
-import jax_dataclasses
-import jaxopt
-
-import jaxsim.api as js
-import jaxsim.typing as jtp
-from jaxsim.api.common import VelRepr
-from jaxsim.math import Adjoint
-from jaxsim.terrain.terrain import FlatTerrain, Terrain
-
-from .common import ContactModel, ContactsParams, ContactsState
-
-
-@jax_dataclasses.pytree_dataclass
-class RelaxedRigidContactsParams(ContactsParams):
-    """Parameters of the relaxed rigid contacts model."""
-
-    # Time constant
-    time_constant: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(0.01, dtype=float)
-    )
-
-    # Adimensional damping coefficient
-    damping_coefficient: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(1.0, dtype=float)
-    )
-
-    # Minimum impedance
-    d_min: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(0.9, dtype=float)
-    )
-
-    # Maximum impedance
-    d_max: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(0.95, dtype=float)
-    )
-
-    # Width
-    width: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(0.0001, dtype=float)
-    )
-
-    # Midpoint
-    midpoint: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(0.1, dtype=float)
-    )
-
-    # Power exponent
-    power: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(1.0, dtype=float)
-    )
-
-    # Stiffness
-    stiffness: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(0.0, dtype=float)
-    )
-
-    # Damping
-    damping: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(0.0, dtype=float)
-    )
-
-    # Friction coefficient
-    mu: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(0.5, dtype=float)
-    )
-
-    # Maximum number of iterations
-    max_iterations: jtp.Int = dataclasses.field(
-        default_factory=lambda: jnp.array(50, dtype=int)
-    )
-
-    # Solver tolerance
-    tolerance: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(1e-6, dtype=float)
-    )
-
-    def __hash__(self) -> int:
-        from jaxsim.utils.wrappers import HashedNumpyArray
-
-        return hash(
-            (
-                HashedNumpyArray(self.time_constant),
-                HashedNumpyArray(self.damping_coefficient),
-                HashedNumpyArray(self.d_min),
-                HashedNumpyArray(self.d_max),
-                HashedNumpyArray(self.width),
-                HashedNumpyArray(self.midpoint),
-                HashedNumpyArray(self.power),
-                HashedNumpyArray(self.stiffness),
-                HashedNumpyArray(self.damping),
-                HashedNumpyArray(self.mu),
-                HashedNumpyArray(self.max_iterations),
-                HashedNumpyArray(self.tolerance),
-            )
-        )
-
-    def __eq__(self, other: RelaxedRigidContactsParams) -> bool:
-        return hash(self) == hash(other)
-
-    @classmethod
-    def build(
-        cls,
-        time_constant: jtp.FloatLike | None = None,
-        damping_coefficient: jtp.FloatLike | None = None,
-        d_min: jtp.FloatLike | None = None,
-        d_max: jtp.FloatLike | None = None,
-        width: jtp.FloatLike | None = None,
-        midpoint: jtp.FloatLike | None = None,
-        power: jtp.FloatLike | None = None,
-        stiffness: jtp.FloatLike | None = None,
-        damping: jtp.FloatLike | None = None,
-        mu: jtp.FloatLike | None = None,
-        max_iterations: jtp.IntLike | None = None,
-        tolerance: jtp.FloatLike | None = None,
-    ) -> RelaxedRigidContactsParams:
-        """Create a `RelaxedRigidContactsParams` instance"""
-
-        return cls(
-            **{
-                field: jnp.array(locals().get(field, default), dtype=default.dtype)
-                for field, default in map(
-                    lambda f: (f, cls.__dataclass_fields__[f].default),
-                    filter(lambda f: f != "__mutability__", cls.__dataclass_fields__),
-                )
-            }
-        )
-
-    def valid(self) -> bool:
-        return bool(
-            jnp.all(self.time_constant >= 0.0)
-            and jnp.all(self.damping_coefficient > 0.0)
-            and jnp.all(self.d_min >= 0.0)
-            and jnp.all(self.d_max <= 1.0)
-            and jnp.all(self.d_min <= self.d_max)
-            and jnp.all(self.width >= 0.0)
-            and jnp.all(self.midpoint >= 0.0)
-            and jnp.all(self.power >= 0.0)
-            and jnp.all(self.mu >= 0.0)
-            and jnp.all(self.max_iterations > 0)
-            and jnp.all(self.tolerance > 0.0)
-        )
-
-
-@jax_dataclasses.pytree_dataclass
-class RelaxedRigidContactsState(ContactsState):
-    """Class storing the state of the relaxed rigid contacts model."""
-
-    def __eq__(self, other: RelaxedRigidContactsState) -> bool:
-        return hash(self) == hash(other)
-
-    @staticmethod
-    def build() -> RelaxedRigidContactsState:
-        """Create a `RelaxedRigidContactsState` instance"""
-
-        return RelaxedRigidContactsState()
-
-    @staticmethod
-    def zero(model: js.model.JaxSimModel) -> RelaxedRigidContactsState:
-        """Build a zero `RelaxedRigidContactsState` instance from a `JaxSimModel`."""
-        return RelaxedRigidContactsState.build()
-
-    def valid(self, model: js.model.JaxSimModel) -> bool:
-        return True
-
-
-@jax_dataclasses.pytree_dataclass
-class RelaxedRigidContacts(ContactModel):
-    """Relaxed rigid contacts model."""
-
-    parameters: RelaxedRigidContactsParams = dataclasses.field(
-        default_factory=RelaxedRigidContactsParams
-    )
-
-    terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
-        default_factory=FlatTerrain
-    )
-
-    def compute_contact_forces(
-        self,
-        position: jtp.Vector,
-        velocity: jtp.Vector,
-        model: js.model.JaxSimModel,
-        data: js.data.JaxSimModelData,
-        link_forces: jtp.MatrixLike | None = None,
-    ) -> tuple[jtp.Vector, tuple[Any, ...]]:
-
-        link_forces = (
-            link_forces
-            if link_forces is not None
-            else jnp.zeros((model.number_of_links(), 6))
-        )
-
-        references = js.references.JaxSimModelReferences.build(
-            model=model,
-            data=data,
-            velocity_representation=data.velocity_representation,
-            link_forces=link_forces,
-        )
-
-        def _detect_contact(x: jtp.Array, y: jtp.Array, z: jtp.Array) -> jtp.Array:
-            x, y, z = jax.tree_map(jnp.squeeze, (x, y, z))
-
-            n̂ = self.terrain.normal(x=x, y=y).squeeze()
-            h = jnp.array([0, 0, z - model.terrain.height(x=x, y=y)])
-
-            return jnp.dot(h, n̂)
-
-        # Compute the activation state of the collidable points
-        δ = jax.vmap(_detect_contact)(*position.T)
-
-        with (
-            references.switch_velocity_representation(VelRepr.Mixed),
-            data.switch_velocity_representation(VelRepr.Mixed),
-        ):
-            M = js.model.free_floating_mass_matrix(model=model, data=data)
-            Jl_WC = jnp.vstack(
-                jax.vmap(lambda J, height: J * (height < 0))(
-                    js.contact.jacobian(model=model, data=data)[:, :3, :], δ
-                )
-            )
-            W_H_C = js.contact.transforms(model=model, data=data)
-            BW_ν̇_free = jnp.hstack(
-                js.ode.system_acceleration(
-                    model=model,
-                    data=data,
-                    link_forces=references.link_forces(model=model, data=data),
-                )
-            )
-            BW_ν = data.generalized_velocity()
-            J̇_WC = jnp.vstack(
-                jax.vmap(lambda J̇, height: J̇ * (height < 0))(
-                    js.contact.jacobian_derivative(model=model, data=data)[:, :3], δ
-                ),
-            )
-
-            a_ref, R, K, D = self._regularizers(
-                model=model,
-                penetration=δ,
-                velocity=velocity,
-                parameters=self.parameters,
-            )
-
-        G = Jl_WC @ jnp.linalg.lstsq(M, Jl_WC.T)[0]
-        CW_al_free_WC = Jl_WC @ BW_ν̇_free + J̇_WC @ BW_ν
-
-        # Calculate quantities for the linear optimization problem.
-        A = G + R
-        b = CW_al_free_WC - a_ref
-
-        objective = lambda x: jnp.sum(jnp.square(A @ x + b))
-
-        # Compute the 3D linear force in C[W] frame
-        opt = jaxopt.LBFGS(
-            fun=objective,
-            maxiter=self.parameters.max_iterations,
-            tol=self.parameters.tolerance,
-            maxls=30,
-            history_size=10,
-            max_stepsize=100.0,
-        )
-
-        init_params = (
-            K[:, jnp.newaxis] * jnp.zeros_like(position).at[:, 2].set(δ)
-            + D[:, jnp.newaxis] * velocity
-        ).flatten()
-
-        CW_f_Ci = opt.run(init_params=init_params).params.reshape(-1, 3)
-
-        def mixed_to_inertial(W_H_C: jax.Array, CW_fl: jax.Array) -> jax.Array:
-            W_Xf_CW = Adjoint.from_transform(
-                W_H_C.at[0:3, 0:3].set(jnp.eye(3)),
-                inverse=True,
-            ).T
-            return W_Xf_CW @ jnp.hstack([CW_fl, jnp.zeros(3)])
-
-        W_f_C = jax.vmap(mixed_to_inertial)(W_H_C, CW_f_Ci)
-
-        return W_f_C, (None,)
-
-    @staticmethod
-    def _regularizers(
-        model: js.model.JaxSimModel,
-        penetration: jtp.Array,
-        velocity: jtp.Array,
-        parameters: RelaxedRigidContactsParams,
-    ) -> tuple:
-        """
-        Compute the contact jacobian and the reference acceleration.
-
-        Args:
-            model: The jaxsim model.
-            penetration: The penetration of the collidable points.
-            velocity: The velocity of the collidable points.
-            parameters: The parameters of the relaxed rigid contacts model.
-
-        Returns:
-            A tuple containing the reference acceleration, the regularization matrix, the stiffness, and the damping.
-        """
-
-        Ω, ζ, ξ_min, ξ_max, width, mid, p, K, D, μ, *_ = jax_dataclasses.astuple(
-            parameters
-        )
-
-        def _imp_aref(
-            penetration: jtp.Array,
-            velocity: jtp.Array,
-        ) -> tuple[jtp.Array, jtp.Array]:
-            """
-            Calculates impedance and offset acceleration in constraint frame.
-
-            Args:
-                penetration: penetration in constraint frame
-                velocity: velocity in constraint frame
-
-            Returns:
-                a_ref: offset acceleration in constraint frame
-                R: regularization matrix
-                K: computed stiffness
-                D: computed damping
-            """
-            position = jnp.zeros(shape=(3,)).at[2].set(penetration)
-
-            imp_x = jnp.abs(position) / width
-            imp_a = (1.0 / jnp.power(mid, p - 1)) * jnp.power(imp_x, p)
-
-            imp_b = 1 - (1.0 / jnp.power(1 - mid, p - 1)) * jnp.power(1 - imp_x, p)
-
-            imp_y = jnp.where(imp_x < mid, imp_a, imp_b)
-
-            imp = jnp.clip(ξ_min + imp_y * (ξ_max - ξ_min), ξ_min, ξ_max)
-            imp = jnp.atleast_1d(jnp.where(imp_x > 1.0, ξ_max, imp))
-
-            # When passing negative values, K and D represent a spring and damper, respectively.
-            K_f = jnp.where(K < 0, -K / ξ_max**2, 1 / (ξ_max * Ω * ζ) ** 2)
-            D_f = jnp.where(D < 0, -D / ξ_max, 2 / (ξ_max * Ω))
-
-            a_ref = -jnp.atleast_1d(D_f * velocity + K_f * imp * position)
-
-            return imp, a_ref, jnp.atleast_1d(K_f), jnp.atleast_1d(D_f)
-
-        def _compute_row(
-            *,
-            link_idx: jtp.Float,
-            penetration: jtp.Array,
-            velocity: jtp.Array,
-        ) -> tuple[jtp.Array, jtp.Array]:
-
-            # Compute the reference acceleration.
-            ξ, a_ref, K, D = _imp_aref(
-                penetration=penetration,
-                velocity=velocity,
-            )
-
-            # Compute the regularization terms.
-            R = (
-                (2 * μ**2 * (1 - ξ) / (ξ + 1e-12))
-                * (1 + μ**2)
-                @ jnp.linalg.inv(M_L[link_idx, :3, :3])
-            )
-
-            return jax.tree.map(lambda x: x * (penetration < 0), (a_ref, R, K, D))
-
-        M_L = js.model.link_spatial_inertia_matrices(model=model)
-
-        a_ref, R, K, D = jax.tree.map(
-            jnp.concatenate,
-            (
-                *jax.vmap(_compute_row)(
-                    link_idx=jnp.array(
-                        model.kin_dyn_parameters.contact_parameters.body
-                    ),
-                    penetration=penetration,
-                    velocity=velocity,
-                ),
-            ),
-        )
-        return a_ref, jnp.diag(R), K, D