Implement RBDA with jax.lax.scan

[diegoferigo]

All our Rigid Body Dynamics Algorithms have been implemented either using plain for loops (that are unrolled during JIT compilation incurring in long build times) or jax.experimental.loops (that provided a nice syntactic sugar over the low-level jax.lax.{scan|fori_loop|while_loop}) that unfortunately have been removed in jax-ml/jax#11607 and no longer part of JAX starting from v0.3.16.

This issue tracks the activity of updating all the usage of the removed jax.experimental.loops. At this point, it seems wise to exploit jax.lax.scan throughout the algorithms. Readability will be definitely affected, but at least in this way we can also ensure that the code is forward and backward differentiable (#4).

• ABA (ABA and RNEA with jax.lax.scan, Python 3.11, fixed bugs in other representations #38)
• CRBA (Implement CRBA with jax.lax.scan #14)
• RNEA (ABA and RNEA with jax.lax.scan, Python 3.11, fixed bugs in other representations #38)
• Jacobian (Implement Jacobian with jax.lax.scan #16)
• Forward Kinematics (Implement Forward Kinematics with jax.lax.scan #15)
• Contact points kinematics (Rewrite soft-contacts algorithm #34)

[diegoferigo]

Reminder: as soon as we purge all jax.experimental.loops invocations, we can remove the following pinning:

jaxsim/setup.cfg

Line 59 in 4ad8f12

jax < 0.3.16

[traversaro]

cc @fl-ferr