Use Quaternion.integrate for SO3 mixin

src/jaxsim/integrators/common.py

@@ -5,11 +5,11 @@
 import jax
 import jax.numpy as jnp
 import jax_dataclasses
-import jaxlie
 from jax_dataclasses import Static
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
+from jaxsim.math import Quaternion
 from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass, Mutability
 
 try:
@@ -560,27 +560,16 @@ def post_process_state(
         cls, x0: js.ode_data.ODEState, t0: Time, xf: js.ode_data.ODEState, dt: TimeStep
     ) -> js.ode_data.ODEState:
 
-        # Indices to convert quaternions between serializations.
-        to_xyzw = jnp.array([1, 2, 3, 0])
-
         # Get the initial rotation.
-        W_R_B_t0 = jaxlie.SO3.from_quaternion_xyzw(
-            xyzw=x0.physics_model.base_quaternion[to_xyzw]
+        W_Q_B_tf = Quaternion.integration(
+            quaternion=x0.physics_model.base_quaternion,
+            dt=dt,
+            omega=x0.physics_model.base_angular_velocity,
+            omega_in_body_fixed=True,
         )
 
-        # Get the final angular velocity.
-        # This is already computed by averaging the kᵢ in RK-based schemes.
-        # Therefore, by using the ω at tf, we obtain a RK scheme operating
-        # on the SO(3) manifold.
-        W_ω_WB_tf = xf.physics_model.base_angular_velocity
-
-        # Integrate the orientation on SO(3).
-        # Note that we left-multiply with the exponential map since the angular
-        # velocity is expressed in the inertial frame.
-        W_R_B_tf = jaxlie.SO3.exp(tangent=dt * W_ω_WB_tf) @ W_R_B_t0
-
         # Replace the quaternion in the final state.
         return xf.replace(
-            physics_model=xf.physics_model.replace(base_quaternion=W_R_B_tf.wxyz),
+            physics_model=xf.physics_model.replace(base_quaternion=W_Q_B_tf),
             validate=True,
         )