Fix overflow cast in MJX collisions for box-box and convex-convex

[amacati]

Fixes the overflow cast in MJX collisions encountered in #3368 and possibly #2226.

To repeat the central part of #3368: The current implementation of MJX throws a warning for box-box and convex-convex collisions. The exact warning is lax_numpy.py:2799: RuntimeWarning: overflow encountered in cast.

The reason is that the collision checking functions for box-box and convex-convex collisions use a jp.where selection to set infinite values to the maximum possible float value instead. This is done via jp.where(jp.isinf(dist), jp.finfo(float).max, dist). jax uses float32 by default, and so dist is a float32 array. jp.finfo(float) uses float64 however, because jax treats Python's float built-in as double precision. The dtype mismatch then results in the overflow error, because we are setting elements of a float32 array to float64 (or at least attempt it). The fix is to replace the jp.finfo with jp.finfo(dist.dtype). This PR adds a minimal test and fixes the error in the two affected functions.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[amacati]

@btaba I spotted a caveat in the tests: When running the full suite, jax caches the traced collision function and does not recompile. Later, we then skip the warning in the new tests even if the cast is still overflowing, because the code for the warning does not get executed. Clearing the cache and forcing a recompile seemed the best solution to me, so I added that in the latest push.