Runtime Warning for MJX

[lucattyy]

Intro

Hi!

I am a student and I am using MuJoCo for my research on Robotics.

My setup

MJX 3.2.4
MUJOCO 3.2.4
Ubuntu 22.04

My question

When I defined the step function in the RL environment, a warning appeared during JIT compilation and execution, so I used jax.debug.print in the code to find the source of the problem. I found that the problem occurred in the step function of mjx.The code and the warning are attached below.

Here is my code.

def step():
    ....
    jax.config.update("jax_disable_jit", False)
    pipeline_state = state.pipeline_state
    pipeline_state = pipeline_state.replace(xfrc_applied=xfrc_applied)
    state = state.replace(pipeline_state = pipeline_state)
    jax.debug.print('flag1')
    jax.config.update("jax_disable_jit", True)
    jax.config.update("jax_disable_jit", False)

    target_torque = jp.clip(target_torque, -1000, 1000)
    pipeline_state = self.pipeline_step(state.pipeline_state, target_torque)
    jax.debug.print('flag2')
    jax.config.update("jax_disable_jit", False)
    x, xd = pipeline_state.x, pipeline_state.xd

    done |= positive_transpos(x.pos[0])
    jax.debug.print('flag3 done')
    jax.config.update("jax_disable_jit", True)
    .....

Here is the warning.

flag1
/home/.conda/envs/jax/lib/python3.9/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
/home/.conda/envs/jax/lib/python3.9/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
/home/.conda/envs/jax/lib/python3.9/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
/home/.conda/envs/jax/lib/python3.9/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
/home/.conda/envs/jax/lib/python3.9/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
/home/.conda/envs/jax/lib/python3.9/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
/home/.conda/envs/jax/lib/python3.9/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
/home/.conda/envs/jax/lib/python3.9/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast
  return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))
flag2
flag3 done

This warning will only appear in the first two steps and will not appear again afterwards. Considering that the mjx step has no connection with the code, does this mean that the problem comes from the xml? How can I determine the source of the problem later?

Minimal model and/or code that explain my question

No response

Confirmations

• I searched the latest documentation thoroughly before posting.
• I searched previous Issues and Discussions, I am certain this has not been raised before.