Copybara import of the project:

--
83a17d2844770fc2bbff37eda73b82df56076414 by Martin Schuck <martin.schuck@tum.de>:

Fix overflow cast

--
d12211c6665e1d791f77338daf8c62fa2374e3c8 by Martin Schuck <martin.schuck@tum.de>:

Prevent skipping warnings from cached jax functions by clearning the cache before invokation

COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/3369 from amacati:fix.overflow_cast d12211c6665e1d791f77338daf8c62fa2374e3c8
PiperOrigin-RevId: 939914544
Change-Id: I0ab6e9ad1e7c45f352d2f49642049ad930955f0b

mjx/mujoco/mjx/_src/collision_convex.py

@@ -714,7 +714,7 @@ def _box_box(b1: ConvexInfo, b2: ConvexInfo) -> Collision:
   # Go back to world frame.
   pos = b2.pos + pos @ b2.mat.T
   n = normal @ b2.mat.T
-  dist = jp.where(jp.isinf(dist), jp.finfo(float).max, dist)
+  dist = jp.where(jp.isinf(dist), jp.finfo(dist.dtype).max, dist)
 
   return dist, pos, n
 
@@ -927,7 +927,7 @@ def _convex_convex(c1: ConvexInfo, c2: ConvexInfo) -> Collision:
   pos = c2.pos + pos @ c2.mat.T
   n = normal @ c2.mat.T
   n = -n if swapped else n
-  dist = jp.where(jp.isinf(dist), jp.finfo(float).max, dist)
+  dist = jp.where(jp.isinf(dist), jp.finfo(dist.dtype).max, dist)
 
   return dist, pos, n
 

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -16,6 +16,7 @@
 
 import dataclasses
 from typing import Dict, Optional, Tuple
+import warnings
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -797,6 +798,20 @@ class ConvexTest(absltest.TestCase):
     c = dx._impl.contact
     self.assertTrue((c.dist > 0).all())
 
+  def test_no_overflow_warning_f32(self):
+    """Tests box-box and convex-convex don't overflow finfo.max in float32."""
+    directory = epath.resource_path('mujoco.mjx')
+    assets = {
+        'meshes/dodecahedron.stl': (
+            directory / 'test_data' / 'meshes/dodecahedron.stl'
+        ).read_bytes(),
+    }
+    with warnings.catch_warnings():  # Regression test for #3368
+      warnings.simplefilter('error', RuntimeWarning)
+      jax.clear_caches()  # force a re-trace so the cast (and warning) re-runs
+      _collide(self._BOX_BOX, keyframe=0)
+      _collide(self._CONVEX_CONVEX, assets=assets)
+
 
 class HFieldTest(absltest.TestCase):
   _HFIELD = """