We're (me and @saloneeverma) internally testing using the warp-based backend for MuJoCo using the mjx frontend, but we've had some issues getting it to work with equinox.filter_vmap. I'm describing the problems here so you can decide if it's worth fixing upstream, cause I think it's fairly niche (mixing vmap and non-vmap Arrays in a single object). We've attached a small reproduceable script at the end, but I don't think it's necessary to understand the issue.
MJX produces a PyTreeNode datastructure called Data with inside an object called DataWarp - which is the core of the issue. DataWarp has three types of attributes:
- ints - set to static by the PyTreeNode
- Arrays which are vmapped over
- Arrays which are NOT vmapped over
The division of whats vmappable and not is decided by MJX and registered through the (non-public?) interpreters.batching API.
This works well - jax.vmap works without any complicated in-axes magic. eqx.filter_vmap however does not know how to handle this. Because the ints are set to static by the PyTreeNode, doing a tree.map to set their axis to None doesnt work (plus, this is done by the interpreter, I guess?). Hence, the solution is to treat DataWarp as a leaf, but that fails too, because filter_vmap always traverses the entire tree (example). It seems like this issue would be solved by allowing an is_leaf input to filter_vmap, but I'm not sure how useful that is for everyone else?
TL;DR: Equinox doesn't seem to be able to handle PytreeNodes defined by other libraries, but we can't traverse them either as some attributes are set to static.
# 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)
We're (me and @saloneeverma) internally testing using the warp-based backend for MuJoCo using the mjx frontend, but we've had some issues getting it to work with
equinox.filter_vmap. I'm describing the problems here so you can decide if it's worth fixing upstream, cause I think it's fairly niche (mixing vmap and non-vmap Arrays in a single object). We've attached a small reproduceable script at the end, but I don't think it's necessary to understand the issue.MJX produces a PyTreeNode datastructure called
Datawith inside an object calledDataWarp- which is the core of the issue.DataWarphas three types of attributes:The division of whats vmappable and not is decided by MJX and registered through the (non-public?)
interpreters.batchingAPI.This works well -
jax.vmapworks without any complicated in-axes magic.eqx.filter_vmaphowever does not know how to handle this. Because the ints are set to static by the PyTreeNode, doing atree.mapto set their axis toNonedoesnt work (plus, this is done by the interpreter, I guess?). Hence, the solution is to treatDataWarpas a leaf, but that fails too, becausefilter_vmapalways traverses the entire tree (example). It seems like this issue would be solved by allowing anis_leafinput tofilter_vmap, but I'm not sure how useful that is for everyone else?TL;DR: Equinox doesn't seem to be able to handle PytreeNodes defined by other libraries, but we can't traverse them either as some attributes are set to static.