-
Notifications
You must be signed in to change notification settings - Fork 9
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Use explicit integration of quaternions also in schemes on SO(3) #205
Use explicit integration of quaternions also in schemes on SO(3) #205
Conversation
RBDAs raise an exception if they receive a non-unary quaternion. When checking gradients with finite differences, the quaternion norm is not enforced to be 1.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That's great, thanks Diego! LGTM
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.
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. |
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.
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 Replying here @traversaro's review comments (#205 (comment) and #205 (comment)).
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 All our where An equivalent integrator on As you can notice, the integration on 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 Let me know if all of this makes more sense now, and if you notice methodological mistakes. |
Good point, I did not realized that we could just normalize the quaternion before calling the RBDA.
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.
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? |
Again, we can merge without waiting for the outcome of the discussion, as anyhow this improves the current situation. |
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 Seen in this way, I believe that the scheme implemented in this PR is 1:1 compatible with the variants not operating on
This is already the case of the integration schemes not operating on |
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 Lines 372 to 381 in 2056e70
For example, see here below one case with ABA: jaxsim/src/jaxsim/api/model.py Lines 674 to 702 in 2056e70
|
Yes, I think it will be more efficient. |
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 |
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. |
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:
JaxSimModelData.state.physics_state.base_quaternion
(that may be not unary on schemes not onBeyond 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.
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/