Intro
Hi!
I am a graduate student at TU Munich, where I use MJX for RL training.
My setup
Latest MJX from the main branch, Ubuntu24 on x86
What's happening? What did you expect?
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). Since dist is float32, but jp.finfo(float) uses float64, this results in the overflow error. The fix is to replace the jp.finfo with jp.finfo(dist.dtype). I have taken the liberty to open a PR that adds a minimal test and fixes the error in the two affected functions.
I know that mjx is getting slowly deprecated in favor of warp, but given the minuscule size of this fix, I hope you do consider taking it in.
This is likely related to discussion #2226 and fixes it.
Steps for reproduction
- Run the code below.
- Code errors because of the warning.
Minimal model for reproduction
No response
Code required for reproduction
import warnings
import jax
import mujoco
from mujoco import mjx
# Two penetrating boxes -> exercises the box_box collider.
XML = """
<mujoco>
<worldbody>
<geom type="box" size="0.025 0.025 0.025" pos="0 0 0.025"/>
<body pos="0 0 0.04">
<freejoint/>
<geom size="0.048 0.01 0.01" type="box"/>
</body>
</worldbody>
</mujoco>
"""
m = mujoco.MjModel.from_xml_string(XML)
mx = mjx.put_model(m)
dx = mjx.put_data(m, mujoco.MjData(m))
with warnings.catch_warnings():
warnings.simplefilter("error", RuntimeWarning) # turn the warning into a hard error
dx = jax.jit(mjx.kinematics)(mx, dx)
dx = jax.jit(mjx.collision)(mx, dx) # box_box -> finfo(float64).max cast into f32
print("done, no warning") # Not reached, code errors out because of the warningConfirmations
Intro
Hi!
I am a graduate student at TU Munich, where I use MJX for RL training.
My setup
Latest MJX from the main branch, Ubuntu24 on x86
What's happening? What did you expect?
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.whereselection to set infinite values to the maximum possible float value instead. This is done viajp.where(jp.isinf(dist), jp.finfo(float).max, dist). Since dist is float32, butjp.finfo(float)uses float64, this results in the overflow error. The fix is to replace thejp.finfowithjp.finfo(dist.dtype). I have taken the liberty to open a PR that adds a minimal test and fixes the error in the two affected functions.I know that
mjxis getting slowly deprecated in favor ofwarp, but given the minuscule size of this fix, I hope you do consider taking it in.This is likely related to discussion #2226 and fixes it.
Steps for reproduction
Minimal model for reproduction
No response
Code required for reproduction
Confirmations