From 1d2049f17637a20ab53ac194e07aef86768cbe0e Mon Sep 17 00:00:00 2001 From: diegoferigo Date: Thu, 21 Mar 2024 09:42:59 +0100 Subject: [PATCH] Integrate quaternions in SO(3) also in RK stages --- src/jaxsim/integrators/fixed_step.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/jaxsim/integrators/fixed_step.py b/src/jaxsim/integrators/fixed_step.py index eb7426bfa..90fd90759 100644 --- a/src/jaxsim/integrators/fixed_step.py +++ b/src/jaxsim/integrators/fixed_step.py @@ -6,6 +6,7 @@ import jaxlie import jaxsim.api as js +from jaxsim.math import Quaternion from .common import ExplicitRungeKutta, PyTreeType, Time, TimeStep @@ -103,11 +104,18 @@ def integrate_rk_stage( op = lambda x0_leaf, k_leaf: x0_leaf + dt * k_leaf xf: js.ode_data.ODEState = jax.tree_util.tree_map(op, x0, k) + W_Q_B_t0 = x0.physics_model.base_quaternion + W_ω_WB_t0 = x0.physics_model.base_angular_velocity + return xf.replace( physics_model=xf.physics_model.replace( - base_quaternion=xf.physics_model.base_quaternion - / jnp.linalg.norm(xf.physics_model.base_quaternion) - ), + base_quaternion=Quaternion.integration( + quaternion=W_Q_B_t0, + dt=dt, + omega=W_ω_WB_t0, + omega_in_body_fixed=False, + ), + ) ) @classmethod