Add relaxed rigid contacts model

environment.yml

@@ -8,6 +8,7 @@ dependencies:
   - python >= 3.12.0
   - coloredlogs
   - jax >= 0.4.13
+  - jaxopt >= 0.8.0
   - jaxlib >= 0.4.13
   - jaxlie >= 1.3.0
   - jax-dataclasses >= 1.4.0

pyproject.toml

@@ -45,6 +45,7 @@ classifiers = [
 dependencies = [
     "coloredlogs",
     "jax >= 0.4.13",
+    "jaxopt >= 0.8.0",
     "jaxlib >= 0.4.13",
     "jaxlie >= 1.3.0",
     "jax_dataclasses >= 1.4.0",
@@ -181,6 +182,7 @@ platforms = ["linux-64", "osx-arm64", "osx-64"]
 coloredlogs = "*"
 jax = "*"
 jax-dataclasses = "*"
+jaxopt = "*"
 jaxlib = "*"
 jaxlie = "*"
 lxml = "*"

src/jaxsim/api/contact.py

@@ -131,7 +131,8 @@ 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: nothing.
+        - Rigid: no additional data.
+        - QuasiRigid: no additional data.
 
     Note:
         The material deformation rate is always returned in the mixed frame
@@ -144,6 +145,10 @@ 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
 
@@ -190,6 +195,27 @@ 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}")
 

src/jaxsim/api/data.py

@@ -593,16 +593,18 @@ def reset_base_quaternion(self, base_quaternion: jtp.VectorLike) -> Self:
             The updated `JaxSimModelData` object.
         """
 
-        base_quaternion = jnp.array(base_quaternion)
+        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),
+        )
 
         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)
             ),
         )
 

src/jaxsim/api/model.py

@@ -1935,7 +1935,7 @@ def step(
     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,
+        pred=tf_ns < t0_ns,
         true_fun=lambda: jax.debug.print(
             "The simulation time overflowed, resetting simulation time to 0."
         ),

src/jaxsim/api/ode.py

@@ -175,17 +175,15 @@ def system_velocity_dynamics(
             forces=W_f_Li_terrain,
             additive=True,
         )
-    # Get the link forces in the data representation
-    with references.switch_velocity_representation(data.velocity_representation):
+
+        # Get the link forces in inertial representation
         f_L_total = references.link_forces(model=model, data=data)
 
-    # 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
-    )
+        v̇_WB, s̈ = system_acceleration(
+            model=model, data=data, joint_forces=joint_forces, link_forces=f_L_total
+        )
 
-    return W_v̇_WB, s̈, aux_data
+    return v̇_WB, s̈, aux_data
 
 
 def system_acceleration(
@@ -196,7 +194,7 @@ def system_acceleration(
     link_forces: jtp.MatrixLike | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
-    Compute the system acceleration in inertial-fixed representation.
+    Compute the system acceleration in the active representation.
 
     Args:
         model: The model to consider.
@@ -206,7 +204,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 inertial-fixed representation
+        A tuple containing the base 6D acceleration in in the active representation
         and the joint accelerations.
     """
 
@@ -272,18 +270,15 @@ def system_acceleration(
     )
 
     # - Joint accelerations: 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̈
+    # - 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̈
 
 
 @jax.jit
@@ -353,7 +348,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
 
@@ -371,7 +366,7 @@ def system_dynamics(
         case SoftContacts():
             ode_state_kwargs["tangential_deformation"] = aux_dict["m_dot"]
 
-        case RigidContacts():
+        case RigidContacts() | RelaxedRigidContacts():
             pass
 
         case _:

src/jaxsim/api/ode_data.py

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

src/jaxsim/integrators/common.py

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

src/jaxsim/math/inertia.py

@@ -45,7 +45,7 @@ def to_params(M: jtp.Matrix) -> tuple[jtp.Float, jtp.Vector, jtp.Matrix]:
             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.

src/jaxsim/mujoco/loaders.py

@@ -211,7 +211,7 @@ def convert(
         joints_dict = {j.name: j for j in rod_model.joints()}
 
         # Convert all the joints not considered to fixed joints.
-        for joint_name in set(j.name for j in rod_model.joints()) - considered_joints:
+        for joint_name in {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 @@ def convert(
         mj_model = mj.MjModel.from_xml_string(xml=urdf_string, assets=assets)
 
         # Get the joint names.
-        mj_joint_names = set(
+        mj_joint_names = {
             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:

src/jaxsim/parsers/kinematic_graph.py

@@ -394,7 +394,7 @@ def reduce(self, considered_joints: Sequence[str]) -> KinematicGraph:
             return copy.deepcopy(self)
 
         # Check if all considered joints are part of the full kinematic graph
-        if len(set(considered_joints) - set(j.name for j in full_graph.joints)) != 0:
+        if len(set(considered_joints) - {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 @@ def reduce(self, considered_joints: Sequence[str]) -> KinematicGraph:
             root_link_name=full_graph.root.name,
         )
 
-        assert set(f.name for f in self.frames).isdisjoint(
-            set(f.name for f in unconnected_frames + reduced_frames)
+        assert {f.name for f in self.frames}.isdisjoint(
+            {f.name for f in unconnected_frames + reduced_frames}
         )
 
         for link in unconnected_links: