Fix integration stage for integrators of type ExplicitRungeKuttaSOMixin

[xela-95]

As per title, this PR fixes how the infra-stage integration step is performed on integrators inheriting from ExplicitRungeKuttaSOMixin class

---

📚 Documentation preview 📚: https://jaxsim--192.org.readthedocs.build//192/

[traversaro]

Probably it is worth to highlight what this PR is fixing, as this may have drastically influenced past tests. We can do that by using the RK equations from https://en.wikipedia.org/w/index.php?title=Runge%E2%80%93Kutta_methods&oldid=1227591654#Explicit_Runge%E2%80%93Kutta_methods for reference.

As mentioned in https://en.wikipedia.org/w/index.php?title=Runge%E2%80%93Kutta_methods&oldid=1227591654#Examples, the Forward Euler method can be implemented as a Explicit Runge Kutta method with $A=[0]$, $b=[1]$, $c = [0]$.

By using this equations, the ForwardEuler integration equation boil down to (I use $x$ for state and $\Delta T$ for timestep instead of $y$ and $h$ used by Wikipedia):

$$ x_{t+ \Delta T} = x_{t} + \Delta T \sum_{i=1}^1 b_i k_i = x_{t} + \Delta T k_1 $$

where

$$ k_1 = f(x_t, t) $$

however, this can be written in more general term as:

$$ k_1 = f(x_t + a_{11} \Delta T) $$

It is important to write it in such general form as indeed the $x_t + a_{11} \Delta T = x_t$ equation is computed by the integrate_rk_stage function, that actually is a generic function for all RK stages. Indeed, everything works fine if the regular ExplicitRungeKutta (used by non-SO3 ForwardEuler and RungeKutta*** integrators) class, but if a ***SO3 integrator is used, then a different version of integrate_rk_stage function described in

jaxsim/src/jaxsim/integrators/common.py

Lines 544 to 563 in 7340d43

	def integrate_rk_stage(
	cls, x0: js.ode_data.ODEState, t0: Time, dt: TimeStep, k: js.ode_data.ODEState
	) -> js.ode_data.ODEState:
	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=Quaternion.integration(
	quaternion=W_Q_B_t0,
	dt=dt,
	omega=W_ω_WB_t0,
	omega_in_body_fixed=False,
	),
	)
	)

is used.

Before this PR, that modified SO3 integrate_rk_stage is assuming that the $x$ state can be decomposed as:

$$ x = \begin{bmatrix} o \\ q \\ \vdots \\ v \\ w \\ \vdots \end{bmatrix} $$

where $o \in \mathbb{R}^3$ is the base frame origin (${}^A o_B$), $q in \mathbb{R}^4$ are the coordinates of the quaternion representing the orientation of the base frame orientation w.r.t. to the inertial (quaternion corresponding to ${}^A R_B$), and $v \in \mathbb{R}^3$ and $\omega \in \mathbb{R}^3$ are the the velocity of the base frame in some kind of representation, where the $\omega$ is expressed with the orientation of the inertial frame (i.e. ${}^A \omega_{A,B}$).

So, only for the portion of $x$ corresponding to $q$ (that we can indicate with $x^q$, the equation:

$$ x_t + a_{11} \Delta T $$

was modified to be:

$$ exp^q(\omega \Delta T)*x_t^q $$

where $exp^q(.) : \mathbb{R}^3 \mapsto \mathbb{R}^4$ is the exponential map for quaternions and $*$ is the quaternion multiplication.

Note that this change changes an identity (as $a_{11} is always zero) to use instead $exp^q(\omega \delta T)$ that may be not an identity, as in general $\omega$ is different from zero.

This PR modifies the SO3-specific modification of the integrate_rk_stage to compute instead:

$$ \frac{x_t^q + a_{11} \Delta T}{|x_t^q + a_{11} \Delta T|} $$

that still normalize the quaternions, but remaining an identity.

Note that the problem was not apparent if the system had zero (or small) angular velocity in the base frame, and indeed to properly debug it we set explicitly the initial angular velocity of the system to a really high number (~2000 [rad]/s).

In the long term, we could consider to modify our system to actually normalize the input quaternions to avoid having to manually modifies the internal steps of the integrators, but that is another story and it would not block the improvement proposed in this PR.

fyi @flferretti @DanielePucci @CarlottaSartore

[traversaro]

For reference, this effectively reverts the commit in 1d2049f that apparently was introduced to handle variable step integrator. However, all tests seems to be passing also with this modification.

[S-Dafarra]

By using this equations, the ForwardEuler integration equation boil down to (I use $x$ for state and $\Delta T$ for timestep instead of $y$ and $h$ used by Wikipedia):

$$ x_{t+ \Delta T} = x_{t} + \Delta T \sum_{i=1}^1 b_i k_i = x_{t} + \Delta T k_1 $$

where

$$ k_1 = f(x_t, t) $$

Ok, I started writing this, but apparently in jaxsim there is something different w.r.t. to the formalism used in Wikipedia there is something different, as for some reason also $k_2$ is computed (@diegoferigo @S-Dafarra any idea why or if you have a reference for RK integrators that is more close to jaxsim then wikipedia?):

$$ k_2 = f(x_t + a_{21} k_1 \Delta T, t + c_2 \Delta T) $$

Not sure I understood your comment, but in the case of Forward Euler, there should not be any k2, since there should be only one stage.

[traversaro]

Not sure I understood your comment, but in the case of Forward Euler, there should not be any k2, since there should be only one stage.

That is exactly what is confusing me. While debugging with @xela-95, we found the problem in the call in

jaxsim/src/jaxsim/integrators/common.py

Line 394 in 7340d43

xi = self.integrate_rk_stage(x0, t0, Δt, sum_ak)

, but if RK has only one stage, it is not clear to me why it was called. Perhaps it was called with k_leaf=0 ?

[traversaro]

Yes, I think this was called with k_leaf=0 as in https://github.com/ami-iit/jaxsim/blob/7340d43172d5ec528f3f7547d532ffeb9770355d/src/jaxsim/integrators/common.py#L394C58-L394C64 it is summing a_{11} that is always zero. Great, this simplifies the explanation, I will edit the comment.

[traversaro]

I updated the comment know that @S-Dafarra explained me that $a_{11}$ is always zero, while the Wikipedia article just wrote the equations by always assuming $a_{11} = 0$.

[diegoferigo]

I can chime in with further details since I'm the author of these integrators. Here below the history of the SO(3) integrators:

• 443049f
• ef51c17
• 1d2049f

As @traversaro pointed out, the last commit is the one to blame for the differences seen between the ForwardEuler and ForwardEulerSO3. In fact, when the quaternion is integrated on the manifold also in intermediate RK stages, spurious effect appear due to the fact that the integration rule on SO(3) appear when $a_{i,j}=0$. In fact, on SO(3), the final quaternion $\texttt{Q}(t_f)$ becomes the wrong $\texttt{Q}(t_f) = \exp(\boldsymbol{\omega} t)\texttt{Q}(t_0)$ instead of the no-op $\texttt{Q}(t_f) = \texttt{Q}(t_0)$. This is due to the fact that the integration rule on SO(3) in this case is not the identity.

This workaround was originally introduced to mitigate the effect of losing slowly the norm of the quaternion in RK stages (considering no explicit normalization). This PR basically reverts to this behavior that seems more sound even if less AD-friendly.

Not sure I understood your comment, but in the case of Forward Euler, there should not be any k2, since there should be only one stage.

@S-Dafarra yes you are right but the stage code is generic and it should work for any valid explicit/embedded Butcher tableau. The origin of the problem is the integration scheme of SO(3) performed in intermediate stages that cannot be considered 1:1 with the corresponding FE-like integration (with fractional $\Delta t$ and $\mathbf{k}_i$).

In the long term, we could consider to modify our system to actually normalize the input quaternions to avoid having to manually modifies the internal steps of the integrators, but that is another story and it would not block the improvement proposed in this PR.

The non-SO(3) variants operate on a generic pytree of $\mathbb{R}^n$ arrays. The SO(3) variants have been introduces to specialize them on a pytree that has a quaternion on a leaf, and they are the recommended integrators in JaxSim. Your suggestion, if I understood, already means to use the SO(3) variants.

[diegoferigo]

Before merging I want to write something that might not clear to future readers. For a generic RK scheme described by a valid Butcher tableau (explicit, not necessarily single-output), our variants on $\text{SO}(3)$ perform the following integration:

• Integrate first the ODEState pytree with regular RK stages. Since there is a quaternion, its norm gets normalized to 1 at each stage. This allows to compute the states at intermediate points used by later stages of the active scheme.
• Once we got $\mathbf{x}(t_f)$, we get the $\boldsymbol{\omega}_{W,B}$ resulting by the overall integration (that will be in general a weighted average at different instants and states), and replace the quaternion obtained from the stages (those that are normalized to 1) with a new one computed through a semi-implicit integration on SO(3) (see here). If you unroll that equation with the actual weighted form of $\boldsymbol{\omega}$, it should result to something similar to specialized RK schemes that operate directly on SO(3). Here we cannot use them since our state is $(\mathbf{q}, \boldsymbol{\nu})$, a mixture of SO(3) and $\mathbb{R}^n$.

[traversaro]

In the long term, we could consider to modify our system to actually normalize the input quaternions to avoid having to manually modifies the internal steps of the integrators, but that is another story and it would not block the improvement proposed in this PR.

The non-SO(3) variants operate on a generic pytree of arrays. The SO(3) variants have been introduces to specialize them on a pytree that has a quaternion on a leaf, and they are the recommended integrators in JaxSim. Your suggestion, if I understood, already means to use the SO(3) variants.

Actually what I wanted to suggest is to modify our system_dynamics (sorry, not system, that was a typo) to normalize the quaternion given in input, to avoid the need to manually normalize each quaternion passed to the system_dynamics. This does not solve possible quaternion drift that stop being unit form (for that baumgarte or other similar methods still needs to be used), but it avoids the need (if I get it correctly) to overload integrate_rk_stage (or any similar method for other classes of integrators) to normalize the quaternion produced by the euler integrator in it.

[traversaro]

As @traversaro pointed out, the last commit is the one to blame for the differences seen between the ForwardEuler and ForwardEulerSO3. In fact, when the quaternion is integrated on the manifold also in intermediate RK stages, spurious effect appear due to the fact that the integration rule on SO(3) appear when

Just to clarify: from what I understand, the current computation (before this PR) is wrong also for stages different from 1, as in any case it uses the system angular velocity instead of the angular velocity that would result from doing the computation of $a_{ij}$ and $k_i$, is that correct or I am missing something?

[diegoferigo]

In the long term, we could consider to modify our system to actually normalize the input quaternions to avoid having to manually modifies the internal steps of the integrators, but that is another story and it would not block the improvement proposed in this PR.

The non-SO(3) variants operate on a generic pytree of arrays. The SO(3) variants have been introduces to specialize them on a pytree that has a quaternion on a leaf, and they are the recommended integrators in JaxSim. Your suggestion, if I understood, already means to use the SO(3) variants.

Actually what I wanted to suggest is to modify our system_dynamics (sorry, not system, that was a typo) to normalize the quaternion given in input, to avoid the need to manually normalize each quaternion passed to the system_dynamics. This does not solve possible quaternion drift that stop being unit form (for that baumgarte or other similar methods still needs to be used), but it avoids the need (if I get it correctly) to overload integrate_rk_stage (or any similar method for other classes of integrators) to normalize the quaternion produced by the euler integrator in it.

Few notes. The base quaternion, part of the state vector of the system dynamics, is always normalized when returned to the user (JaxSimModelData.base_orientation). The data has no knowledge of the active integrator, therefore normalizing it is the default since its norm might not be zero in the non-SO(3) integrators are used (i.e. the quaternion is integrated from $\dot{\texttt{Q}}$ and a Baumgarte correction to dynamically adjust its norm).

Regardless to this, there is always the need to perform a kind or projection to unary quaternions in RK stages. It cannot be worked around by simply starting from a unary quaternion. In fact, especially with high-order integrators and sufficiently large $\Delta t$, the quaternion norm can significantly diverge from 1 within the multiple stages that are executed.

I inspected the code used by our RBDAs. I thought that they were normalizing by default the quaternions before using them, but this is actually not the case (they all call Transform.from_quaternion_and_translation or Adjoint.from_quaternion_and_translation). Maybe we can update the default argument to normalize_quaternion=True to be more robust.

As @traversaro pointed out, the last commit is the one to blame for the differences seen between the ForwardEuler and ForwardEulerSO3. In fact, when the quaternion is integrated on the manifold also in intermediate RK stages, spurious effect appear due to the fact that the integration rule on SO(3) appear when

Just to clarify: from what I understand, the current computation (before this PR) is wrong also for stages different from 1, as in any case it uses the system angular velocity instead of the angular velocity that would result from doing the computation of aij and ki, is that correct or I am missing something?

Yes because the derivative of the quaternion was taken from $\mathbf{x}_0$. In order to have a 1:1 alignment with non-the SO(3) variant, we should get $\boldsymbol{\omega}$ from $\mathbf{k}_{i}$, that in case of FE is zero due to $a_{i,j} = 0$. Unfortunately this is not straightforward because $\boldsymbol{\omega}$ is not directly part of $\mathbf{k}_i$ (that is basically a relative of $\dot{\mathbf{x}}$). If we want to do that, we should extract $\boldsymbol{\omega}$ from $\dot{\mathtt{Q}}$ (part of $\dot{\mathbf{x}}$), which is doable. Note that in this way, we might probably keep the integration on the manifold within RK stages. In fact, for FE with $a_{i,j}=0$, the spurious effects would not be injected in this case. Does this make sense to you?

[traversaro]

Yes because the derivative of the quaternion was taken from $\mathbf{x}_0$. In order to have a 1:1 alignment with non-the SO(3) variant, we should get $\boldsymbol{\omega}$ from $\mathbf{k}_{i}$, that in case of FE is zero due to $a_{i,j} = 0$. Unfortunately this is not straightforward because $\boldsymbol{\omega}$ is not directly part of $\mathbf{k}_i$ (that is basically a relative of $\dot{\mathbf{x}}$). If we want to do that, we should extract $\boldsymbol{\omega}$ from $\dot{\mathtt{Q}}$ (part of $\dot{\mathbf{x}}$), which is doable. Note that in this way, we might probably keep the integration on the manifold within RK stages. In fact, for FE with $a_{i,j}=0$, the spurious effects would not be injected in this case. Does this make sense to you?

Yes, that make sense, without this fix you get completely wrong integration steps (for all stages) especially if you a huge velocity in x0 (that is the test case we were using with @xela-95).

Regardless to this, there is always the need to perform a kind or projection to unary quaternions in RK stages. It cannot be worked around by simply starting from a unary quaternion. In fact, especially with high-order integrators and sufficiently large $\Delta t$, the quaternion norm can significantly diverge from 1 within the multiple stages that are executed.

That is clear, but why this is a problem? As long as the quaternion is normalized when it is used (i.e. while computing the system_dynamics), what is the problem of having a non-normalized quaternion being computed by integrate_rk_stage ?

[diegoferigo]

Regardless to this, there is always the need to perform a kind or projection to unary quaternions in RK stages. It cannot be worked around by simply starting from a unary quaternion. In fact, especially with high-order integrators and sufficiently large Δt, the quaternion norm can significantly diverge from 1 within the multiple stages that are executed.

That is clear, but why this is a problem? As long as the quaternion is normalized when it is used (i.e. while computing the system_dynamics), what is the problem of having a non-normalized quaternion being computed by integrate_rk_stage ?

This is not necessarily a problem if there is somewhere a normalization. However, in complex codebases sometimes it can get overlooked and funny things might happen. As in #192 (comment), I forgot that by default quaternions were not normalized in our RBDAs (and this could create problems with with non-SO(3) integrators).

[traversaro]

This is not necessarily a problem if there is somewhere a normalization. However, in complex codebases sometimes it can get overlooked and funny things might happen. As in #192 (comment), I forgot that by default quaternions were not normalized in our RBDAs (and this could create problems with with non-SO(3) integrators).

Just to clarify, I was not suggesting for RBDAs not to normalize quaternions, for clarity I think the normalization should be done only once and in the system dynamics, before the quaternion is passed to RBDA. The idea (that may be wrong) is that for non-experienced users is is much easier to understand what is going on at the system_dynamics level as opposed to some internal function of the integrator (at least for @xela-95 and me that was the case).

Anyhow, this discussion are unrelated and do not block the merge of the PR, @flferretti I guess that if you want to review the PR then we can merge.

[flferretti]

Thanks everybody, LGTM

[DanielePucci]

This was a big and important fix. Well done everyone.