Use explicit integration of quaternions also in schemes on SO(3)

[diegoferigo]

After the fix performed in #192, I had a deeper look to the $\text{SO}(3)$ integrators and how RBDAs and user code interface with the simulated data. This PR:

• Raises an exception if:
  • Any RBDA receives a quaternion that is not unary.
  • Any $\text{SO}(3)$ integrator receives at $t_0$ a quaternion that is not unary.
• Removes the custom implementation of the RK stage of $\text{SO}(3)$ integrators. This is no longer necessary since the quaternion is always normalized before being passed to RBDAs (and the raise, otherwise).
• Removing these normalizations should be safe also on the user side since they never should never directly read the buffer JaxSimModelData.state.physics_state.base_quaternion (that may be not unary on schemes not on $\text{SO}(3)$ ).

Beyond enhancin the robustness and fixing the explicitness of the integrators, another aim of this PR is to minimize the division by the norm, since it may affect AD performance. It should now be safe enough to remove the normalizations used here and there since all RBDA receive a quaternion from JaxSimModelData.base_orientation(), that already normalizes the quaternion if needed.

cc @traversaro @xela-95 @DanielePucci @flferretti

---

Note that I found this problem by simulating one of the upcoming examples, in particular one simulating the tennis racket theorem (Dzhanibekov effect). In such system, one of the axis should always be constant. With the removed semi-implicit integration of the quaternion on the manifold, I was observing a kind of precession of such axis. This PR fixes the wrong outcome.

Semi-implicit integration on SO(3) (wrong)	Explicit integration on SO(3) (correct)
tennis_racket_no.mp4	tennis_racket_ok.mp4

As you see, before this PR, the motion was getting altered by the integrator. In the wrong version, there is a precession of the rotation axis that disrupt the motion.

---

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

[flferretti]

That's great, thanks Diego! LGTM

[traversaro]

I am not a big fan of raising an error in the RBDA algorithm if the quaternion is not normalized, as it prevents taking a numerical derivative of any RBDA with a step higher then the tolerance of the jax method used to check the unitary norm (i.e. 1e-5, if https://numpy.org/doc/stable/reference/generated/numpy.allclose.html#numpy.allclose and jax are properly aligned). Anyhow, when necessary we could make that check optional.

On the other hand, I am really happy that we remove the custom rk integration step, that make the debug more complex.

to minimize the division by the norm, since it may affect AD performance.

In which way the division by the norm is affecting AD performance?

I have some technical doubt of how quaternions are integrated (see #205 (comment) and https://github.com/ami-iit/jaxsim/pull/205/files#r1675427312) but we can proceed with the merge anyhow and discuss those later.

[diegoferigo]

I am not a big fan of raising an error in the RBDA algorithm if the quaternion is not normalized, as it prevents taking a numerical derivative of any RBDA with a step higher then the tolerance of the jax method used to check the unitary norm

I'm not sure if you followed the commits of this PR, but this is exactly what happened here and that got fixed in f06065e. My take on this is that finite differences should not be used on JAX functions. The only exception is performing tests against AD as in this case, and there is a quick workaround as done in the commit - if needed. Do you think this compromise is good enough? If users have this need, wrapping the function they want to differentiate with FD with a tiny wrapper that normalizes the quaternion should be decent compromise.

In which way the division by the norm is affecting AD performance?

If you take a complex input-output computation involving quaternions, why should we introduce on purpose in the AD chain multiple divisions by the square root? I suspect that this would make the functions returned by jax.jacfwd|jax.jacrev more heavy for no reason. Ideally, we should have only one normalization injected in that chain, does this make sense?

---

Replying here @traversaro's review comments (#205 (comment) and #205 (comment)).

Do you have a note and/or reference explain how this extensions of RK to SO(3) works in math terms? I am afraid it is not obvious.

I may be confused, isn't this replacing all RK machinery for the angular velocity/quaternion with a single step forward euler? What is the gain of this over simply normalizing the quaternion in the dynamics, and use baungmarte to ensure the unit norm of the quaternion is respected (that we already do if I am not wrong).

At a first view, I totally understand the reason that led you write the second comment, and I would totally agree with you if we would integrate the the generalized position $\mathbf{q} \in \mathbb{R}^{7+n}$ semi-implicitly. I'll try to provide some insights here below.

All our ExplicitRungeKutta are... explicit integrators. Even if the dynamics of the system is second order. Given the system state $(\mathbf{q}, \boldsymbol{\nu}) \in \mathbb{R}^{7+n} \times \mathbb{R}^{6+n}$ (let me drop the terrain deformation for simplicity), the regular non- $\text{SO}(3)$ integration scheme is the following:

$$\begin{cases} \boldsymbol{\nu}[k+1] &= \boldsymbol{\nu}[k] + \Delta t \; \dot{\boldsymbol{\nu}}[k] \\\ \mathbf{q}[k+1] &= \mathbf{q}[k] + \Delta t \; \dot{\mathbf{q}}(\mathbf{q}[k], \boldsymbol{\nu}[k]) \end{cases}$$

where $\dot{\mathbf{q}}(\mathbf{q}[k], \boldsymbol{\nu}[k])$ models the extraction of the quaternion derivative ${}^W \dot{\mathtt{Q}}_B$ from the initial state (and this is also where the Baumgarte correction is injected).

An equivalent integrator on $\text{SO}(3)$ would be the following, where I unpacked the terms of the second equation:

$$\begin{cases} \boldsymbol{\nu}[k+1] &= \boldsymbol{\nu}[k] + \Delta t \; \dot{\boldsymbol{\nu}}[k] \\\ {}^W\mathbf{p}_B[k+1] &= {}^W\mathbf{p}_B[k] + \Delta t \; {}^W\dot{\mathbf{p}}_B[k] \\\ {}^W\mathtt{Q}_B[k+1] &= \text{IntegratorSO3}({}^W\mathtt{Q}_B[k], {}^W\boldsymbol{\omega}_B[k]) \\\ \mathbf{s}[k+1] &= \mathbf{s}[k] + \Delta t \; \dot{\mathbf{s}}[k] \\\ \end{cases}$$

As you can notice, the integration on $\text{SO}(3)$, uses the initial angular velocity ${}^W\boldsymbol{\omega}_B[k]$. This velocity is computed in the previous step by integrating the with the active RK scheme the system acceleration.

As we already discussed in #192, we cannot do much on the manifold in the intermediate RK stages. Indeed, the current situation is those stages is not integrate the quaternion not on the manifold (at best, using a Baumgarte regularization) and always normalize the quaternion when needed. Note that this is only necessary in the intermediate RK stages (e.g. when they evaluate again the dynamics). The users, when they use $\text{SO}(3)$ integrators, can expect that the resulting quaternion after the integration is always unary.

Let me know if all of this makes more sense now, and if you notice methodological mistakes.

[traversaro]

I am not a big fan of raising an error in the RBDA algorithm if the quaternion is not normalized, as it prevents taking a numerical derivative of any RBDA with a step higher then the tolerance of the jax method used to check the unitary norm

I'm not sure if you followed the commits of this PR, but this is exactly what happened here and that got fixed in f06065e. My take on this is that we finite differences should not be used on JAX functions. The only exception is performing tests against AD as in this case, and there is a quick workaround as done in the commit - if needed. Do you think this compromise is good enough? If users have this need, wrapping the function they want to differentiate with FD with a tiny wrapper that normalizes the quaternion should be decent compromise.

Good point, I did not realized that we could just normalize the quaternion before calling the RBDA.

In which way the division by the norm is affecting AD performance?

If you take a complex input-output computation involving quaternions, why should we introduce on purpose in the AD chain multiple divisions by the square root? I suspect that this would make the functions returned by jax.jacfwd|jax.jacrev more heavy for no reason. Ideally, we should have only one normalization injected in that chain, does this make sense?

Sure, you want only one normalization, not more, I was just confused by the generate sentence "division by the norm is affecting AD performance" that seemed to refer to division by the norm in general and not by the (repeated) division by the norm.

Let me know if all of this makes more sense now, and if you notice methodological mistakes.

Your example make sense, but I think the confusing part happens when you use an actual RK integrator (not a forward euler like in your example). In that case, somehow only the base state integration continue to use Forward Euler. To keep thing simple, could it make sense to avoid mixing Forward Euler and more complex RK, and just baugarte stabilization and/or quaternion normalization for handling the constraint, instead of manually changing the integration logic? Would we loose some property in that case?

[traversaro]

Again, we can merge without waiting for the outcome of the discussion, as anyhow this improves the current situation.

[diegoferigo]

Your example make sense, but I think the confusing part happens when you use an actual RK integrator (not a forward euler like in your example). In that case, somehow only the base state integration continue to use Forward Euler.

If the integrator in my example would be a higher-order RK scheme, what I described would still be 100% valid. In fact, in that case, the ${}^W\boldsymbol{\omega}_B$ used to integrate the quaternion on the manifold would come from the previous step (and it doesn't matter if it's computed from a higher-order RK scheme). This is correct because this operation doesn't happen on the equation corresponding to the higher order equation of the equivalent ODE.

Seen in this way, I believe that the scheme implemented in this PR is 1:1 compatible with the variants not operating on $\text{SO}(3)$. Am I missing something? We can align f2f if there are other doubts on this.

To keep thing simple, could it make sense to avoid mixing Forward Euler and more complex RK, and just baugarte stabilization and/or quaternion normalization for handling the constraint, instead of manually changing the integration logic? Would we loose some property in that case?

This is already the case of the integration schemes not operating on $\text{SO}(3)$. In these schemes, the quaternion is integrated from ${}^W \dot{\mathtt{Q}}_B$ (possibly, with a Baumgarte correction) and it is normalized before being fed to the RBDAs. So, if users want this option, they can use the already existing integrators not operating on the manifold.

[traversaro]

• This is no longer necessary since the quaternion is always normalized before being passed to RBDAs (and the raise, otherwise).

Just to understand, which part of the code is doing the normalization in the inner RK integration?

[diegoferigo]

• This is no longer necessary since the quaternion is always normalized before being passed to RBDAs (and the raise, otherwise).

Just to understand, which part of the code is doing the normalization in the inner RK integration?

We never access the possibly non-unary quaternion stored into JaxSimModelData.state.physics_model.base_quaternion (and this internal buffer should also not be used by downstream users). In all our computations, we always use the high-level getter that includes the following logic:

jaxsim/src/jaxsim/api/data.py

Lines 372 to 381 in 2056e70

	# Always normalize the quaternion to avoid numerical issues.
	# If the active scheme does not integrate the quaternion on its manifold,
	# we introduce a Baumgarte stabilization to let the quaternion converge to
	# a unit quaternion. In this case, it is not guaranteed that the quaternion
	# stored in the state is a unit quaternion.
	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),
	)

For example, see here below one case with ABA:

jaxsim/src/jaxsim/api/model.py

Lines 674 to 702 in 2056e70

	# Extract the state in inertial-fixed representation.
	with data.switch_velocity_representation(VelRepr.Inertial):
	W_p_B = data.base_position()
	W_v_WB = data.base_velocity()
	W_Q_B = data.base_orientation(dcm=False)
	s = data.joint_positions(model=model, joint_names=joint_names)
	ṡ = data.joint_velocities(model=model, joint_names=joint_names)
	# Extract the inputs in inertial-fixed representation.
	with references.switch_velocity_representation(VelRepr.Inertial):
	W_f_L = references.link_forces(model=model, data=data, link_names=link_names)
	τ = references.joint_force_references(model=model, joint_names=joint_names)
	# ========================
	# Compute forward dynamics
	# ========================
	W_v̇_WB, s̈ = jaxsim.rbda.aba(
	model=model,
	base_position=W_p_B,
	base_quaternion=W_Q_B,
	joint_positions=s,
	base_linear_velocity=W_v_WB[0:3],
	base_angular_velocity=W_v_WB[3:6],
	joint_velocities=ṡ,
	joint_forces=τ,
	link_forces=W_f_L,
	standard_gravity=data.standard_gravity(),
	)

[traversaro]

We can align f2f if there are other doubts on this.

Yes, I think it will be more efficient.

[traversaro]

So, if users want this option, they can use the already existing integrators not operating on the manifold.

Yes, indeed I was proposing to just use that version by default as is is more simple to understand.

[diegoferigo]

So, if users want this option, they can use the already existing integrators not operating on the manifold.

Yes, indeed I was proposing to just use that version by default as is is more simple to understand.

This is a choice of downstream users, in the past I noticed better stability with $\text{SO}(3)$ integrators. But after all these changes, we might need to perform additional tests.

[traversaro]

This is a choice of downstream users, in the past I noticed better stability with $SO(3)$ integrators.

We can discuss in person, but the choice of the default integrator (and/or the one documented in the examples) that most of the users will use is a choice on our side, not a choice of downstream users.

[diegoferigo]

This is a choice of downstream users, in the past I noticed better stability with SO(3) integrators.

We can discuss in person, but the choice of the default integrator (and/or the one documented in the examples) that most of the users will use is a choice on our side, not a choice of downstream users.

There is no such thing as default integrator in JaxSim. The downstream users are responsible for instantiating the integrator they desire to use among those we provide.

About the documentation, instead, we can definitely make recommendations.