# Note - needs to run on a GPU?
import equinox as eqx
import jax.numpy as jnp
import mujoco
from mujoco import mjx
import jax
import mujoco.mjx.warp.types as _wt
#minimal sphere model
XML = r"""
<mujoco>
<worldbody>
<body>
<freejoint/>
<geom size=".15" mass="1" type="sphere"/>
</body>
</worldbody>
</mujoco>
"""
model = mujoco.MjModel.from_xml_string(XML) #make minimal model
mjx_model = mjx.put_model(model, impl="warp") #establish warp backend
N = 4 #number of environments to map over
def make_state(_):
d = mjx.make_data(model, impl="warp", naconmax=100, njmax=100) #make minimal data with warp backend
return mjx.forward(mjx_model, d)
def step(d):
return mjx.step(mjx_model, d) #simple step function
def custom_filtering_function(x):
"""Build a ``eqx.filter_vmap`` ``in_axes`` prefix for a (batched) warp state.
Args:
x: a (batched) state pytree
Returns:
A pytree mirroring ``x`` whose leaves are ``0`` for batched arrays (and
for the ``DataWarp``) and ``None`` for non-array leaves
and non-vmap data. Pass it as the ``in_axes`` for the state argument of ``eqx.filter_vmap``.
"""
def _is_warp_leaf(v):
return _wt is not None and isinstance(v, _wt.DataWarp)
return jax.tree.map(
lambda leaf: 0 if eqx.is_array(leaf) or _is_warp_leaf(leaf) else None,
x,
is_leaf=_is_warp_leaf,
)
if __name__ == "__main__":
batched = jax.vmap(make_state)(jnp.arange(N))
print("built batched Data OK:", batched.qpos.shape) # (N, 7)
# ---------------------------------------------------------------------------
# BUG: eqx.filter_vmap cannot map over the Warp Data.
# ValueError: vmap in_axes specification must be a tree prefix ...
# ---------------------------------------------------------------------------
print("eqx.filter_vmap(step)(batched) ...")
out = eqx.filter_vmap(step, in_axes=(custom_filtering_function(batched),), out_axes=(custom_filtering_function(batched)))(batched) # <-- raises ValueError
print(" unexpectedly OK:", out.qpos.shape)
# NOTE: plain jax.vmap respects the vmappable and works.
out = jax.vmap(step, in_axes=(custom_filtering_function(batched),))(batched)
print("jax.vmap(step)(batched) OK:", out.qpos.shape)
MJX-warp 3.9.
.
.
Intro
We use an equinox-based RL stack with Mujoco and have recently been testing switching to MJX-warp from MJX-Jax, but ran into a compatibility issue, mainly stemming from MJX-warp's use of a non-public custom vmap API. I was wondering why you're using that one instead of the publicly documented one? If there's a good reason for it I'll see if I can update
eqx.filter_vmap, otherwise I can submit a PR here to switch to the public API?Analog issue in equinox here.
Minimally reproducable script:
My setup
MJX-warp 3.9.
What's happening? What did you expect?
.
Steps for reproduction
.
Minimal model for reproduction
No response
Code required for reproduction
No response
Confirmations