From 407be2fc0bce6b4caa5c5a890a9423d6a40fa0d6 Mon Sep 17 00:00:00 2001 From: Martin Schuck Date: Mon, 29 Jun 2026 10:44:54 -0700 Subject: [PATCH] Copybara import of the project: -- 83a17d2844770fc2bbff37eda73b82df56076414 by Martin Schuck : Fix overflow cast -- d12211c6665e1d791f77338daf8c62fa2374e3c8 by Martin Schuck : 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 | 4 ++-- mjx/mujoco/mjx/_src/collision_driver_test.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py index 0d60c4c9..d4662a09 100644 --- a/mjx/mujoco/mjx/_src/collision_convex.py +++ b/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 diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index ea05e11c..c392a5ed 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/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 = """