From a6011c655060f5068878dc35d5cf93087d2b6f0f Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 4 Jul 2024 17:13:45 +0200 Subject: [PATCH] Use `Quaternion.integrate` for SO3 mixin --- src/jaxsim/integrators/common.py | 25 +++++++------------------ 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/src/jaxsim/integrators/common.py b/src/jaxsim/integrators/common.py index 71900957..b578120a 100644 --- a/src/jaxsim/integrators/common.py +++ b/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, )