Fix running AD through jaxsim.math.Rotation.from_axis_angle

[diegoferigo]

In these days, I tried to solve an unconstrained non-linear optimization problem using the JaxSim functions to construct the cost function. I used the L-BFGS solver from optax, and it works flawlessly 🎉

What I'm trying to solve is quite usual in AMI. I want to compute the initial configuration of a humanoid robot such that the knees and the elbows are slightly bent, the soles of the feet are in contact, and the feet are slightly more distant w.r.t. the configuration with zero joint positions. The optimized variables are the joint positions and the base height. The corresponding optimization problem computes FK in the cost function (exploiting AD support through the JaxSim function).

I had to fix the gradient propagation as in this PR to prevent having problems either with zero division in the jitted function or the same problem in the gradient computed with AD. The trick I used it the double-where workaround, as described in:

• https://jax.readthedocs.io/en/latest/faq.html#gradients-contain-nan-where-using-where
• Best practice to handle division by zero in auto-differentiation jax-ml/jax#5039 (comment)
• gradients through np.where when one of branches is nan, #347 except with arctan2 jax-ml/jax#1052 (comment)

---

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

[diegoferigo]

Note: probably a similar trick can be used to remove the small tolerances that we always had in SoftContacts:

jaxsim/src/jaxsim/rbda/contacts/soft.py

Lines 288 to 292 in 35cf336

	# There are few operations like computing the norm of a vector with zero length
	# or computing the square root of zero that are problematic in an AD context.
	# To avoid these issues, we introduce a small tolerance ε to their arguments
	# and make sure that we do not check them against zero directly.
	ε = jnp.finfo(float).eps

[flferretti]

Thanks Diego, LGTM!