Intro
Hey MuJoCo team,
Thanks for the great work on MJWarp! I've been messing around with MJX-Warp and it's generally been working well for me. I ran into a tricky tracer leak issue earlier today, and I'm unsure if it's a bug, or if I'm misunderstanding how to design the JAX interface with the Warp data.
Here's a reproducer script that leads to the UnexpectedTracerError in an "autoreset" style interface similar to other RL pipelines. (Looking at BRAX and Playground has been super helpful!) This occurs when jnp.where is called between the reset data and the stepped data, and I'm not quite sure if it's the where or the vmap that is the source of the error.
Disclaimer: Claude helped me track down this error. I still consider myself to be an MJWarp beginner, so bear with me if I'm misunderstanding something fundamental.
My setup
Here's what I'm currently running in my venv:
jax[cuda13]==0.10.2
mujoco == 3.10.0
mujoco-mjx == 3.10.0
warp-lang == 1.13.0
Python: 3.12.3
GPU: RTX 5000 (laptop) with CUDA 13.0
OS: Ubuntu 24.04
What's happening? What did you expect?
[described above]
Steps for reproduction
Here's a reproducer script
import jax
import jax.numpy as jnp
import mujoco
from mujoco import mjx
# Simple setup with a box initialized in the air above a flat plane
XML = """
<mujoco>
<option timestep="0.005"/>
<worldbody>
<geom name="floor" type="plane" size="5 5 0.1"/>
<body name="box" pos="0 0 0.2">
<freejoint/>
<geom name="box" type="box" size="0.1 0.1 0.1"/>
</body>
</worldbody>
</mujoco>
"""
# Flag to apply a somewhat hacky "fix" to this issue
# This was discovered by Claude but it does not really feel like a great solution
# If True, this will always maintain the _impl field of the stepped env rather than the reset env.
# I believe (maybe) this doesn't cause issues if this data is recomputed?
# But also I'm wondering if there is ever a case where the wrong data is applied immediately after a reset
KEEP_STEPPED_IMPL = False
def make_reset(mjx_model, data_template):
def reset(key):
# reset the box height to a somewhat random value
z = 0.2 + 0.1 * jax.random.uniform(key)
qpos = data_template.qpos.at[2].set(z)
d = data_template.replace(qpos=qpos)
return mjx.forward(mjx_model, d)
return reset
def make_autoreset_step(mjx_model, reset_fn):
def step(key, data):
stepped_data = mjx.step(mjx_model, data)
# terminate when the box has fallen below a height (just for example purposes)
done = stepped_data.qpos[2] < 0.15
reset_data = reset_fn(key)
next_data = jax.tree_util.tree_map(
lambda r, s: jnp.where(done, r, s), reset_data, stepped_data
)
if KEEP_STEPPED_IMPL: # Hacky "fix" described above
next_data = next_data.replace(_impl=stepped_data._impl)
return next_data
return step
def main():
NUM_ENVS = 8
NUM_STEPS = 4
mj_model = mujoco.MjModel.from_xml_string(XML)
mjx_model = mjx.put_model(mj_model, impl="warp")
data_template = mjx.make_data(
mj_model, impl="warp", naconmax=16 * NUM_ENVS, njmax=64
)
reset_fn = make_reset(mjx_model, data_template)
step_fn = make_autoreset_step(mjx_model, reset_fn)
keys = jax.random.split(jax.random.PRNGKey(0), NUM_ENVS)
batched_data = jax.vmap(reset_fn)(keys)
def rollout(data, key):
def body(carry, _):
data, key = carry
key, sk = jax.random.split(key)
step_keys = jax.random.split(sk, NUM_ENVS)
data = jax.vmap(step_fn)(step_keys, data)
return (data, key), data.qpos[:, 2]
return jax.lax.scan(body, (data, key), None, length=NUM_STEPS)
# This will fail with an UnexpectedTracerError if KEEP_STEPPED_IMPL = False
out = jax.jit(rollout)(batched_data, jax.random.PRNGKey(1))
if __name__ == "__main__":
main()Minimal model for reproduction
No response
Code required for reproduction
No response
Confirmations
Intro
Hey MuJoCo team,
Thanks for the great work on MJWarp! I've been messing around with MJX-Warp and it's generally been working well for me. I ran into a tricky tracer leak issue earlier today, and I'm unsure if it's a bug, or if I'm misunderstanding how to design the JAX interface with the Warp data.
Here's a reproducer script that leads to the
UnexpectedTracerErrorin an "autoreset" style interface similar to other RL pipelines. (Looking at BRAX and Playground has been super helpful!) This occurs whenjnp.whereis called between the reset data and the stepped data, and I'm not quite sure if it's thewhereor thevmapthat is the source of the error.Disclaimer: Claude helped me track down this error. I still consider myself to be an MJWarp beginner, so bear with me if I'm misunderstanding something fundamental.
My setup
Here's what I'm currently running in my venv:
Python: 3.12.3
GPU: RTX 5000 (laptop) with CUDA 13.0
OS: Ubuntu 24.04
What's happening? What did you expect?
[described above]
Steps for reproduction
Here's a reproducer script
Minimal model for reproduction
No response
Code required for reproduction
No response
Confirmations