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

Add stage_in and stage_out for MJX-Warp.

Browse Source 
PiperOrigin-RevId: 868180029
Change-Id: I0d1b17b6023d32498766b5d1d9e5624bd31cda3d
This commit is contained in:
Baruch Tabanpour
2026-02-10 09:17:58 -08:00
committed by Copybara-Service
parent 7ef3ae6ad1
commit 3c107b3047
7 changed files with 730 additions and 583 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/warp/collision_driver.py
+37 -22
View File
@@ -233,8 +233,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'collision_pair': d._impl.collision_pair.shape,
'collision_pairid': d._impl.collision_pairid.shape,
'collision_worldid': d._impl.collision_worldid.shape,
'geom_xmat': d.geom_xmat.shape,
'geom_xpos': d.geom_xpos.shape,
'nacon': d._impl.nacon.shape,
'ncollision': d._impl.ncollision.shape,
'contact__dim': d._impl.contact__dim.shape,
@@ -253,15 +251,13 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_collision_shim,
num_outputs=20,
num_outputs=18,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
'collision_pair',
'collision_pairid',
'collision_worldid',
'geom_xmat',
'geom_xpos',
'nacon',
'ncollision',
'contact__dim',
@@ -278,6 +274,27 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'contact__type',
'contact__worldid',
},
stage_in_argnames={
'geom_aabb',
'geom_friction',
'geom_gap',
'geom_margin',
'geom_rbound',
'geom_size',
'geom_solimp',
'geom_solmix',
'geom_solref',
'geom_xmat',
'geom_xpos',
'hfield_data',
'pair_friction',
'pair_gap',
'pair_margin',
'pair_solimp',
'pair_solref',
'pair_solreffriction',
},
stage_out_argnames={},
graph_mode=m.opt._impl.graph_mode,
)
out = jf(
@@ -373,23 +390,21 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'_impl.collision_pair': out[0],
'_impl.collision_pairid': out[1],
'_impl.collision_worldid': out[2],
'geom_xmat': out[3],
'geom_xpos': out[4],
'_impl.nacon': out[5],
'_impl.ncollision': out[6],
'_impl.contact__dim': out[7],
'_impl.contact__dist': out[8],
'_impl.contact__frame': out[9],
'_impl.contact__friction': out[10],
'_impl.contact__geom': out[11],
'_impl.contact__geomcollisionid': out[12],
'_impl.contact__includemargin': out[13],
'_impl.contact__pos': out[14],
'_impl.contact__solimp': out[15],
'_impl.contact__solref': out[16],
'_impl.contact__solreffriction': out[17],
'_impl.contact__type': out[18],
'_impl.contact__worldid': out[19],
'_impl.nacon': out[3],
'_impl.ncollision': out[4],
'_impl.contact__dim': out[5],
'_impl.contact__dist': out[6],
'_impl.contact__frame': out[7],
'_impl.contact__friction': out[8],
'_impl.contact__geom': out[9],
'_impl.contact__geomcollisionid': out[10],
'_impl.contact__includemargin': out[11],
'_impl.contact__pos': out[12],
'_impl.contact__solimp': out[13],
'_impl.contact__solref': out[14],
'_impl.contact__solreffriction': out[15],
'_impl.contact__type': out[16],
'_impl.contact__worldid': out[17],
})
return d
mjx/mujoco/mjx/warp/ffi.py
+4
View File
@@ -101,6 +101,8 @@ def jax_callable_variadic_tuple(
vmap_method: Optional[str] = None,
output_dims: Optional[dict[str, tuple[int, ...]]] = None,
in_out_argnames: Optional[Sequence[str]] = None,
stage_in_argnames: Optional[Sequence[str]] = None,
stage_out_argnames: Optional[Sequence[str]] = None,
):
"""Wraps a JAX callable to support variadic tuples and dataclasses."""
@@ -130,6 +132,8 @@ def jax_callable_variadic_tuple(
vmap_method=vmap_method,
output_dims=output_dims,
in_out_argnames=in_out_argnames,
stage_in_argnames=stage_in_argnames,
stage_out_argnames=stage_out_argnames,
)
flat_args, in_tree = jax.tree.flatten(args)
mjx/mujoco/mjx/warp/forward.py
+570 -295
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/warp/forward_test.py
+1 -1
View File
@@ -218,7 +218,7 @@ class StepTest(parameterized.TestCase):
batch_size=(1, 7),
# NOTE: GraphMode.JAX is incompatible with MuJoCo Warp at the moment,
# even when setting graph_conditional=False.
graph_mode=('WARP',),
graph_mode=('WARP', 'WARP_STAGED'),
)
def test_step(self, xml: str, batch_size: int, graph_mode: str):
if not _FORCE_TEST:
mjx/mujoco/mjx/warp/smooth.py
+65 -60
View File
@@ -136,9 +136,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
output_dims = {
'geom_xmat': d.geom_xmat.shape,
'geom_xpos': d.geom_xpos.shape,
'mocap_pos': d.mocap_pos.shape,
'mocap_quat': d.mocap_quat.shape,
'qpos': d.qpos.shape,
'site_xmat': d.site_xmat.shape,
'site_xpos': d.site_xpos.shape,
'xanchor': d.xanchor.shape,
@@ -151,15 +148,52 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_kinematics_shim,
num_outputs=14,
num_outputs=11,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
'geom_xmat',
'geom_xpos',
'site_xmat',
'site_xpos',
'xanchor',
'xaxis',
'ximat',
'xipos',
'xmat',
'xpos',
'xquat',
},
stage_in_argnames={
'body_ipos',
'body_iquat',
'body_pos',
'body_quat',
'geom_pos',
'geom_quat',
'geom_xmat',
'geom_xpos',
'jnt_axis',
'jnt_pos',
'mocap_pos',
'mocap_quat',
'qpos',
'qpos0',
'site_pos',
'site_quat',
'site_xmat',
'site_xpos',
'xanchor',
'xaxis',
'ximat',
'xipos',
'xmat',
'xpos',
'xquat',
},
stage_out_argnames={
'geom_xmat',
'geom_xpos',
'site_xmat',
'site_xpos',
'xanchor',
@@ -216,18 +250,15 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
d = d.tree_replace({
'geom_xmat': out[0],
'geom_xpos': out[1],
'mocap_pos': out[2],
'mocap_quat': out[3],
'qpos': out[4],
'site_xmat': out[5],
'site_xpos': out[6],
'xanchor': out[7],
'xaxis': out[8],
'ximat': out[9],
'xipos': out[10],
'xmat': out[11],
'xpos': out[12],
'xquat': out[13],
'site_xmat': out[2],
'site_xpos': out[3],
'xanchor': out[4],
'xaxis': out[5],
'ximat': out[6],
'xipos': out[7],
'xmat': out[8],
'xpos': out[9],
'xquat': out[10],
})
return d
@@ -243,25 +274,6 @@ def kinematics_vmap(unused_axis_size, is_batched, m, d):
return d, is_batched[1]
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
)
_d = mjwarp.Data(
**{f.name: None for f in dataclasses.fields(mjwarp.Data) if f.init}
)
_o = mjwarp.Option(
**{f.name: None for f in dataclasses.fields(mjwarp.Option) if f.init}
)
_s = mjwarp.Statistic(
**{f.name: None for f in dataclasses.fields(mjwarp.Statistic) if f.init}
)
_c = mjwarp.Contact(
**{f.name: None for f in dataclasses.fields(mjwarp.Contact) if f.init}
)
_e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
)
@ffi.format_args_for_warp
def _tendon_shim(
# Model
@@ -348,12 +360,6 @@ def _tendon_shim(
def _tendon_jax_impl(m: types.Model, d: types.Data):
output_dims = {
'cdof': d.cdof.shape,
'geom_xmat': d.geom_xmat.shape,
'geom_xpos': d.geom_xpos.shape,
'qpos': d.qpos.shape,
'site_xpos': d.site_xpos.shape,
'subtree_com': d.subtree_com.shape,
'ten_J': d._impl.ten_J.shape,
'ten_length': d.ten_length.shape,
'ten_wrapadr': d._impl.ten_wrapadr.shape,
@@ -363,16 +369,10 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
}
jf = ffi.jax_callable_variadic_tuple(
_tendon_shim,
num_outputs=12,
num_outputs=6,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
'cdof',
'geom_xmat',
'geom_xpos',
'qpos',
'site_xpos',
'subtree_com',
'ten_J',
'ten_length',
'ten_wrapadr',
@@ -380,6 +380,17 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
'wrap_obj',
'wrap_xpos',
},
stage_in_argnames={
'cdof',
'geom_size',
'geom_xmat',
'geom_xpos',
'qpos',
'site_xpos',
'subtree_com',
'ten_length',
},
stage_out_argnames={'ten_length'},
graph_mode=m.opt._impl.graph_mode,
)
out = jf(
@@ -421,18 +432,12 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
d._impl.wrap_xpos,
)
d = d.tree_replace({
'cdof': out[0],
'geom_xmat': out[1],
'geom_xpos': out[2],
'qpos': out[3],
'site_xpos': out[4],
'subtree_com': out[5],
'_impl.ten_J': out[6],
'ten_length': out[7],
'_impl.ten_wrapadr': out[8],
'_impl.ten_wrapnum': out[9],
'_impl.wrap_obj': out[10],
'_impl.wrap_xpos': out[11],
'_impl.ten_J': out[0],
'ten_length': out[1],
'_impl.ten_wrapadr': out[2],
'_impl.ten_wrapnum': out[3],
'_impl.wrap_obj': out[4],
'_impl.wrap_xpos': out[5],
})
return d
mjx/mujoco/mjx/warp/testspeed.py
+52 -204
View File
@@ -22,6 +22,7 @@ from typing import Any, Callable, Sequence, Tuple
from absl import app
from absl import flags
import jax
import jax.numpy as jnp
import mujoco
from mujoco import mjx
from mujoco.mjx._src import test_util
@@ -29,6 +30,7 @@ from mujoco.mjx.warp import collision_driver as wp_collision
from mujoco.mjx.warp import forward as wp_forward
from mujoco.mjx.warp import smooth as wp_smooth
import mujoco.mjx.third_party.mujoco_warp as mjwarp
import numpy as np
import warp as wp
from mujoco.mjx.third_party.warp._src.jax_experimental import ffi as warp_ffi
@@ -50,6 +52,18 @@ _WP_KERNEL_CACHE_DIR = flags.DEFINE_string(
None,
'Path to the Warp kernel cache directory.',
)
_GRAPH_MODE = flags.DEFINE_enum(
'graph_mode',
'WARP',
['NONE', 'WARP', 'WARP_STAGED', 'WARP_STAGED_EX'],
'Graph capture mode for JAX WARP FFI benchmark.',
)
_BENCHMARK = flags.DEFINE_enum(
'benchmark',
'jax_warp',
['jax_warp', 'jax', 'warp'],
'Which benchmark to run.',
)
_COMPILER_OPTIONS = {'xla_gpu_graph_min_graph_size': 1}
jax_jit = functools.partial(jax.jit, compiler_options=_COMPILER_OPTIONS)
@@ -95,10 +109,14 @@ def benchmark(
d = mjx.make_data(
m, impl=mx.impl, naconmax=_NACONMAX.value, njmax=_NJMAX.value
)
# Initialize from first keyframe if available
if m.nkey > 0:
d = d.replace(qpos=m.key_qpos[0], ctrl=m.key_ctrl[0])
return d
key = jax.random.split(jax.random.key(0), nenv)
d = jax_jit(init)(key)
jax.block_until_ready(d)
@jax_jit
@@ -115,177 +133,6 @@ def benchmark(
return jit_time, run_time, steps
def benchmark_raw_jax_warp(
m: mujoco.MjModel,
nstep: int = 1000,
nenv: int = 8192,
unroll_steps: int = 4,
function: str = 'kinematics',
):
if function not in ('kinematics', 'forward', 'step', 'collision'):
raise NotImplementedError(
f'{function} is not implemented for raw warp speed test.'
)
def warp_fn(
qpos_in: wp.array2d(dtype=wp.float32),
qpos_out: wp.array2d(dtype=wp.float32),
time_: wp.array(dtype=wp.float32),
xpos: wp.array2d(dtype=wp.vec3),
xquat: wp.array2d(dtype=wp.quat),
xmat: wp.array2d(dtype=wp.mat33),
xipos: wp.array2d(dtype=wp.vec3),
ximat: wp.array2d(dtype=wp.mat33),
xanchor: wp.array2d(dtype=wp.vec3),
xaxis: wp.array2d(dtype=wp.vec3),
geom_xpos: wp.array2d(dtype=wp.vec3),
geom_xmat: wp.array2d(dtype=wp.mat33),
):
wp.copy(d.qpos, qpos_in)
if function == 'kinematics':
mjwarp.kinematics(mw, d)
elif function == 'forward':
mjwarp.forward(mw, d)
elif function == 'step':
mjwarp.step(mw, d)
elif function == 'collision':
mjwarp.collision(mw, d)
else:
raise NotImplementedError(f'{function} not implemented in speed test.')
wp.copy(qpos_out, d.qpos)
wp.copy(time_, d.time)
wp.copy(xpos, d.xpos)
wp.copy(xquat, d.xquat)
wp.copy(xmat, d.xmat)
wp.copy(xipos, d.xipos)
wp.copy(ximat, d.ximat)
wp.copy(xanchor, d.xanchor)
wp.copy(xaxis, d.xaxis)
wp.copy(geom_xpos, d.geom_xpos)
wp.copy(geom_xmat, d.geom_xmat)
def unroll(
qpos,
qpos_out,
time_,
xpos,
xquat,
xmat,
xipos,
ximat,
xanchor,
xaxis,
geom_xpos,
geom_xmat,
):
def step(carry, _):
qpos_, *_ = carry
out = warp_fn_jax(qpos_ + 0.0 * qpos_)
out = tuple(out)
return (out[0],) + out, None
(
qpos,
qpos_out,
time_,
xpos,
xquat,
xmat,
xipos,
ximat,
xanchor,
xaxis,
geom_xpos,
geom_xmat,
), _ = jax.lax.scan(
step,
(
qpos,
qpos_out,
time_,
xpos,
xquat,
xmat,
xipos,
ximat,
xanchor,
xaxis,
geom_xpos,
geom_xmat,
),
length=nstep,
unroll=unroll_steps,
)
return (
qpos,
qpos_out,
time_,
xpos,
xquat,
xmat,
xipos,
ximat,
xanchor,
xaxis,
geom_xpos,
geom_xmat,
)
output_dims = {
'time_': (nenv,),
'qpos_out': (nenv, m.nq),
'xpos': (nenv, m.nbody, 3),
'xquat': (nenv, m.nbody, 4),
'xmat': (nenv, m.nbody, 3, 3),
'xipos': (nenv, m.nbody, 3),
'ximat': (nenv, m.nbody, 3, 3),
'xanchor': (nenv, m.njnt, 3),
'xaxis': (nenv, m.njnt, 3),
'geom_xpos': (nenv, m.ngeom, 3),
'geom_xmat': (nenv, m.ngeom, 3, 3),
}
warp_fn_jax = warp_ffi.jax_callable(
warp_fn,
num_outputs=11,
output_dims=output_dims,
graph_mode=warp_ffi.GraphMode.JAX,
)
@jax.vmap
def init(key):
d = mjx.make_data(m, impl='jax')
return d
key = jax.random.split(jax.random.key(0), nenv)
dx = jax_jit(init)(key)
d_ = mujoco.MjData(m)
mw = mjwarp.put_model(m)
d = mjwarp.put_data(
m, d_, nworld=nenv, naconmax=_NACONMAX.value, njmax=_NJMAX.value
)
jax_unroll_fn = jax_jit(unroll)
jit_time, run_time = _measure(
jax_unroll_fn,
dx.qpos,
dx.qpos,
dx.time,
dx.xpos,
dx.xquat,
dx.xmat,
dx.xipos,
dx.ximat,
dx.xanchor,
dx.xaxis,
dx.geom_xpos,
dx.geom_xmat,
)
steps = nstep * nenv
return jit_time, run_time, steps
def _compile_fn(fn, m, d):
fn(m, d)
fn(m, d)
@@ -310,9 +157,16 @@ def benchmark_raw_warp(
mw = mjwarp.put_model(m)
dw = mjwarp.make_data(
m, nworld=nenv, nconmax=_NACONMAX.value, njmax=_NJMAX.value
m, nworld=nenv, naconmax=_NACONMAX.value, njmax=_NJMAX.value
)
# Initialize from first keyframe if available
if m.nkey > 0:
qpos_init = np.tile(m.key_qpos[0], (nenv, 1)).astype(np.float32)
ctrl_init = np.tile(m.key_ctrl[0], (nenv, 1)).astype(np.float32)
wp.copy(dw.qpos, wp.from_numpy(qpos_init))
wp.copy(dw.ctrl, wp.from_numpy(ctrl_init))
if function == 'kinematics':
fn = mjwarp.kinematics
elif function == 'forward':
@@ -353,8 +207,16 @@ def _main(_: Sequence[str]):
except Exception as _:
m = mujoco.MjModel.from_xml_path(modelfile)
mx = mjx.put_model(m, impl='jax')
mw = mjx.put_model(m, impl='warp')
benchmark_type = _BENCHMARK.value
graph_mode = getattr(warp_ffi.GraphMode, _GRAPH_MODE.value)
# Only allocate the model needed for the specific benchmark
mx = None
mw = None
if benchmark_type == 'jax_warp':
mw = mjx.put_model(m, impl='warp', graph_mode=graph_mode)
elif benchmark_type == 'jax':
mx = mjx.put_model(m, impl='jax')
if function_ == 'kinematics':
func_warp = jax.vmap(wp_smooth.kinematics, in_axes=(None, 0))
@@ -377,38 +239,24 @@ def _main(_: Sequence[str]):
print(f' nenv : {nenv}')
print(f' nstep : {nstep}')
print(f' timestep : {m.opt.timestep}')
print(f' unroll : {unroll}\n')
print(f' unroll : {unroll}')
print(f' benchmark : {benchmark_type}')
print(f' graph_mode : {_GRAPH_MODE.value}\n')
for name, mx_, op in (
('JAX WARP FFI', mw, func_warp),
('Pure JAX', mx, func_jax),
):
if op is not None:
jit_time, run_time, steps = benchmark(m, mx_, op, nstep, nenv, unroll)
if benchmark_type == 'jax_warp':
jit_time, run_time, steps = benchmark(m, mw, func_warp, nstep, nenv, unroll)
print(f' JAX WARP FFI (GraphMode: {_GRAPH_MODE.value}):')
elif benchmark_type == 'jax':
jit_time, run_time, steps = benchmark(m, mx, func_jax, nstep, nenv, unroll)
print(' Pure JAX:')
elif benchmark_type == 'warp':
jit_time, run_time, steps = benchmark_raw_warp(
m, nstep, nenv, unroll, function=function_
)
print(' Pure WARP:')
else:
raise ValueError(f'Unknown benchmark type: {benchmark_type}')
print(f' {name}:')
print(f' JIT time : {jit_time:.2f} s')
print(f' simulation time : {run_time:.2f} s')
print(f' steps per second : {steps / run_time:,.0f}')
print(
f' realtime factor : {steps * m.opt.timestep / run_time:.2f} x'
)
print(f' time per step : {1e6 * run_time / steps:.2f} µs\n')
jit_time, run_time, steps = benchmark_raw_jax_warp(
m, nstep, nenv, unroll, function=function_
)
print(' Pure JAX-WARP:')
print(f' JIT time : {jit_time:.2f} s')
print(f' simulation time : {run_time:.2f} s')
print(f' steps per second : {steps / run_time:,.0f}')
print(f' realtime factor : {steps * m.opt.timestep / run_time:.2f} x')
print(f' time per step : {1e6 * run_time / steps:.2f} µs\n')
jit_time, run_time, steps = benchmark_raw_warp(
m, nstep, nenv, unroll, function=function_
)
print(' Pure WARP:')
print(f' JIT time : {jit_time:.2f} s')
print(f' simulation time : {run_time:.2f} s')
print(f' steps per second : {steps / run_time:,.0f}')
mjx/mujoco/mjx/warp/types.py
+1 -1
View File
@@ -23,7 +23,6 @@ from jax import tree_util
from jax.interpreters import batching
from mujoco.mjx._src import dataclasses as mjx_dataclasses
import numpy as np
if typing.TYPE_CHECKING:
GraphMode = int
else:
@@ -44,6 +43,7 @@ class TileSet:
adr: address of each tile in the set
size: size of all the tiles in this set
"""
adr: np.ndarray
size: int