chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update MJX-W testspeed.

Browse Source 
PiperOrigin-RevId: 794549657
Change-Id: I0ea12274da77704be30301d33228066bb8952fef
This commit is contained in:
Baruch Tabanpour
2025-08-13 06:23:21 -07:00
committed by Copybara-Service
parent da9037e2b0
commit 6a9e68a542
1 changed files with 1 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/warp/testspeed.py
+1 -4
View File
@@ -249,6 +249,7 @@ def benchmark_raw_jax_warp(
warp_fn, warp_fn,
num_outputs=11, num_outputs=11,
output_dims=output_dims, output_dims=output_dims,
graph_mode=warp_ffi.GraphMode.JAX,
) )
@jax.vmap @jax.vmap
@@ -260,7 +261,6 @@ def benchmark_raw_jax_warp(
dx = jax_jit(init)(key) dx = jax_jit(init)(key)
d_ = mujoco.MjData(m) d_ = mujoco.MjData(m)
mw = mjwarp.put_model(m) mw = mjwarp.put_model(m)
mw.opt.graph_conditional = False
d = mjwarp.put_data( d = mjwarp.put_data(
m, d_, nworld=nenv, nconmax=_NCONMAX.value, njmax=_NJMAX.value m, d_, nworld=nenv, nconmax=_NCONMAX.value, njmax=_NJMAX.value
) )
@@ -309,9 +309,6 @@ def benchmark_raw_warp(
) )
mw = mjwarp.put_model(m) mw = mjwarp.put_model(m)
# TODO(btaba): re-enable graph conditional once JAX supports it, for fair
# comparison.
mw.opt.graph_conditional = False
dw = mjwarp.make_data( dw = mjwarp.make_data(
m, nworld=nenv, nconmax=_NCONMAX.value, njmax=_NJMAX.value m, nworld=nenv, nconmax=_NCONMAX.value, njmax=_NJMAX.value
) )