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

Update mjx-warp render with partial codegen.

Browse Source 
PiperOrigin-RevId: 869874195
Change-Id: I71a16a6c192873786a09211dc4d326c3a4115ad9
This commit is contained in:
Baruch Tabanpour
2026-02-13 13:48:57 -08:00
committed by Copybara-Service
parent faf55f7020
commit bbfe7fe0e2
10 changed files with 195 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+2 -2
View File
@@ -1960,5 +1960,5 @@ def create_render_context(
Render context object that is JAX compatible.
"""
_check_warp_installed()
from mujoco.mjx.warp import render as mjxw_render # pylint: disable=g-import-not-at-top # pytype: disable=import-error
return mjxw_render.create_render_context(mjm, nworld=nworld, **kwargs)
from mujoco.mjx.warp import io as mjxw_io # pylint: disable=g-import-not-at-top # pytype: disable=import-error
return mjxw_io.create_render_context(mjm, nworld=nworld, **kwargs)
mjx/mujoco/mjx/_src/render_util.py
+30 -12
View File
@@ -14,30 +14,40 @@
# ==============================================================================
"""JAX render utilities for unpacking render output from MuJoCo Warp."""
import typing
from typing import Any
import jax
import jax.numpy as jnp
def get_rgb(
rc: Any,
rgb_data: jax.Array,
cam_id: int,
width: int,
height: int,
) -> jax.Array:
"""Unpack uint32 ABGR pixel data into float32 RGB.
Args:
rgb_data: Packed render output, shape (nworld, ncam, H*W)
rc: The RenderContext handle.
rgb_data: Packed render output, shape (nworld, total_pixels)
as uint32.
cam_id: Camera index to extract.
width: Image width.
height: Image height.
Returns:
Float32 RGB array with shape (nworld, H, W, 3), values
in [0, 1].
"""
packed = rgb_data[:, cam_id]
import mujoco.mjx.warp.render as mjxw_render # pylint: disable=g-import-not-at-top
warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[rc.key]
rgb_adr = int(warp_rc.rgb_adr.numpy()[cam_id])
width = int(warp_rc.cam_res.numpy()[cam_id][0])
height = int(warp_rc.cam_res.numpy()[cam_id][1])
packed = jax.lax.dynamic_slice_in_dim(
rgb_data, rgb_adr, width * height, axis=1
)
r = (packed & 0xFF).astype(jnp.float32) / 255.0
g = ((packed >> 8) & 0xFF).astype(jnp.float32) / 255.0
b = ((packed >> 16) & 0xFF).astype(jnp.float32) / 255.0
@@ -47,27 +57,35 @@ def get_rgb(
def get_depth(
rc: Any,
depth_data: jax.Array,
cam_id: int,
width: int,
height: int,
depth_scale: float,
) -> jax.Array:
"""Extract and normalize depth data for a camera.
Args:
depth_data: Raw depth output, shape (nworld, ncam, H*W)
rc: The RenderContext handle.
depth_data: Raw depth output, shape (nworld, total_pixels)
as float32.
cam_id: Camera index to extract.
width: Image width.
height: Image height.
depth_scale: Scale factor for normalizing depth values.
Returns:
Float32 depth array with shape (nworld, H, W), clamped
to [0, 1].
"""
raw = depth_data[:, cam_id]
import mujoco.mjx.warp.render as mjxw_render # pylint: disable=g-import-not-at-top
warp_rc = mjxw_render._MJX_RENDER_CONTEXT_BUFFERS[rc.key]
depth_adr = int(warp_rc.depth_adr.numpy()[cam_id])
width = int(warp_rc.cam_res.numpy()[cam_id][0])
height = int(warp_rc.cam_res.numpy()[cam_id][1])
raw = jax.lax.dynamic_slice_in_dim(
depth_data, depth_adr, width * height, axis=1
)
nworld = depth_data.shape[0]
depth = jnp.clip(raw / depth_scale, 0.0, 1.0)
return depth.reshape(nworld, height, width)
mjx/mujoco/mjx/warp/bvh.py
+58 -36
View File
@@ -12,106 +12,128 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
import mujoco
from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
# Re-use the render registry
from mujoco.mjx.warp.render import _MJX_RENDER_CONTEXT_BUFFERS
from mujoco.mjx.warp.io import _MJX_RENDER_CONTEXT_BUFFERS
from mujoco.mjx.warp.types import RenderContext
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
import warp as wp
_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 _refit_bvh_shim(
# Model
nworld: int,
flex_dim: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
geom_dataid: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
geom_type: wp.array(dtype=int),
nflex: int,
nflexelemdata: int,
nflexvert: int,
flex_dim: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
# Data
flexvert_xpos: wp.array2d(dtype=wp.vec3),
geom_xmat: wp.array2d(dtype=wp.mat33),
geom_xpos: wp.array2d(dtype=wp.vec3),
flexvert_xpos: wp.array2d(dtype=wp.vec3),
# Registry
rc_id: int,
geom_xpos_out: wp.array2d(dtype=wp.vec3),
# Dummy output
dummy: wp.array(dtype=int),
):
_m.stat = _s
_m.opt = _o
_d.efc = _e
_d.contact = _c
_m.flex_dim = flex_dim
_m.flex_elem = flex_elem
_m.flex_elemnum = flex_elemnum
_m.flex_vertadr = flex_vertadr
_m.geom_dataid = geom_dataid
_m.geom_size = geom_size
_m.geom_type = geom_type
_m.nflex = nflex
_m.nflexelemdata = nflexelemdata
_m.nflexvert = nflexvert
_m.flex_dim = flex_dim
_m.flex_elem = flex_elem
_m.flex_elemnum = flex_elemnum
_m.flex_vertadr = flex_vertadr
_d.flexvert_xpos = flexvert_xpos
_d.geom_xmat = geom_xmat
_d.geom_xpos = geom_xpos
_d.flexvert_xpos = flexvert_xpos
_d.nworld = geom_xpos.shape[0]
_d.nworld = nworld
render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id]
dummy.zero_()
mjwarp.refit_bvh(_m, _d, render_context)
wp.copy(geom_xpos_out, geom_xpos)
def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
nworld = d.qpos.shape[0]
ngeom = d.geom_xpos.shape[1]
output_dims = {'dummy': (d.qpos.shape[0],)}
jf = ffi.jax_callable_variadic_tuple(
_refit_bvh_shim,
num_outputs=1,
output_dims={'geom_xpos_out': (nworld, ngeom, 3)},
output_dims=output_dims,
vmap_method=None,
in_out_argnames=set([]),
stage_in_argnames=set(['geom_size', 'geom_xmat', 'geom_xpos']),
stage_out_argnames=set([]),
graph_mode=m.opt._impl.graph_mode,
)
out = jf(
d.qpos.shape[0],
m._impl.flex_dim,
m._impl.flex_elem,
m._impl.flex_elemnum,
m._impl.flex_vertadr,
m.geom_dataid,
m.geom_size,
m.geom_type,
m.nflex,
m.nflexelemdata,
m.nflexvert,
m.flex_dim,
m.flex_elem,
m.flex_elemnum,
m.flex_vertadr,
m._impl.nflex,
m._impl.nflexelemdata,
m._impl.nflexvert,
d._impl.flexvert_xpos,
d.geom_xmat,
d.geom_xpos,
d.flexvert_xpos,
ctx.key,
)
return d.replace(geom_xpos=out[0])
d = d.tree_replace({'time': d.time + out[0]})
return d
@jax.custom_batching.custom_vmap
@functools.partial(ffi.marshal_jax_warp_callable)
@ffi.marshal_jax_warp_callable
def refit_bvh(m: types.Model, d: types.Data, ctx: RenderContext):
return _refit_bvh_jax_impl(m, d, ctx)
@refit_bvh.def_vmap
@functools.partial(ffi.marshal_custom_vmap)
def refit_bvh_vmap(unused_axis_size, is_batched, m, d, ctx):
out = refit_bvh(m, d, ctx)
return out, is_batched[1]
@ffi.marshal_custom_vmap
def refit_bvh_vmap(unused_axis_size, is_batched, m, d, ctx: RenderContext):
d = refit_bvh(m, d, ctx)
return d, is_batched[1]
mjx/mujoco/mjx/warp/collision_driver.py
+6 -5
View File
@@ -223,6 +223,7 @@ def _collision_shim(
_d.naconmax = naconmax
_d.ncollision = ncollision
_d.nworld = nworld
mjwarp.collision(_m, _d)
@@ -249,7 +250,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
num_outputs=15,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
in_out_argnames=set([
'nacon',
'ncollision',
'contact__dim',
@@ -265,8 +266,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'contact__solreffriction',
'contact__type',
'contact__worldid',
},
stage_in_argnames={
]),
stage_in_argnames=set([
'geom_aabb',
'geom_friction',
'geom_gap',
@@ -285,8 +286,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
'pair_solimp',
'pair_solref',
'pair_solreffriction',
},
stage_out_argnames={},
]),
stage_out_argnames=set([]),
graph_mode=m.opt._impl.graph_mode,
)
out = jf(
mjx/mujoco/mjx/warp/forward.py
+14 -12
View File
@@ -909,6 +909,7 @@ def _forward_shim(
_d.xpos = xpos
_d.xquat = xquat
_d.nworld = nworld
mjwarp.forward(_m, _d)
@@ -1012,7 +1013,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
num_outputs=92,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
in_out_argnames=set([
'act_dot',
'actuator_force',
'actuator_length',
@@ -1105,8 +1106,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'efc__state',
'efc__type',
'efc__vel',
},
stage_in_argnames={
]),
stage_in_argnames=set([
'act',
'act_dot',
'actuator_acc0',
@@ -1242,8 +1243,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'xmat',
'xpos',
'xquat',
},
stage_out_argnames={
]),
stage_out_argnames=set([
'act_dot',
'actuator_force',
'actuator_length',
@@ -1276,7 +1277,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'xmat',
'xpos',
'xquat',
},
]),
graph_mode=m.opt._impl.graph_mode,
)
out = jf(
@@ -2688,6 +2689,7 @@ def _step_shim(
_d.xpos = xpos
_d.xquat = xquat
_d.nworld = nworld
mjwarp.step(_m, _d)
@@ -2795,7 +2797,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
num_outputs=96,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
in_out_argnames=set([
'act',
'act_dot',
'actuator_force',
@@ -2892,8 +2894,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'efc__state',
'efc__type',
'efc__vel',
},
stage_in_argnames={
]),
stage_in_argnames=set([
'act',
'act_dot',
'actuator_acc0',
@@ -3029,8 +3031,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'xmat',
'xpos',
'xquat',
},
stage_out_argnames={
]),
stage_out_argnames=set([
'act',
'act_dot',
'actuator_force',
@@ -3067,7 +3069,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'xmat',
'xpos',
'xquat',
},
]),
graph_mode=m.opt._impl.graph_mode,
)
out = jf(
mjx/mujoco/mjx/warp/io.py
+60
View File
@@ -0,0 +1,60 @@
# Copyright 2026 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""I/O functions for MJX Warp."""
import threading
import mujoco
from mujoco.mjx.warp.types import RenderContext
import mujoco.mjx.third_party.mujoco_warp as mjw
_MJX_RENDER_CONTEXT_COUNTER = 0
_MJX_RENDER_CONTEXT_LOCK = threading.Lock()
_MJX_RENDER_CONTEXT_BUFFERS = {}
def create_render_context(
mjm: mujoco.MjModel,
nworld: int,
**kwargs,
):
# NOTE: MuJoCo Warp render context expects a Warp Model and Data.
# We create them here but throw them away right after. Preferably,
# the render context should only rely on mujoco.MjModel so we
# do not have to pay the cost of creating dummy Warp Model and Data.
# Some assumptions may be violated if the downstream render context
# builder holds onto the memory of m and d. The API on the MuJoCo
# Warp side needs to be cleaned up.
m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=nworld)
mjw.forward(m, d)
rc = mjw.create_render_context(
mjm=mjm,
m=m,
d=d,
**kwargs,
)
rc.rgb_data_shape = rc.rgb_data.shape
rc.depth_data_shape = rc.depth_data.shape
rc.rgb_data = None
rc.depth_data = None
global _MJX_RENDER_CONTEXT_COUNTER
with _MJX_RENDER_CONTEXT_LOCK:
_MJX_RENDER_CONTEXT_COUNTER += 1
key = _MJX_RENDER_CONTEXT_COUNTER
_MJX_RENDER_CONTEXT_BUFFERS[key] = rc
return RenderContext(key, _owner=True)
mjx/mujoco/mjx/warp/render.py
+8 -67
View File
@@ -16,19 +16,15 @@
import dataclasses
import functools
import threading
import jax
import mujoco
from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
from mujoco.mjx.warp import mujoco_warp as mjwarp
from mujoco.mjx.warp.io import _MJX_RENDER_CONTEXT_BUFFERS
from mujoco.mjx.warp.types import RenderContext
import warp as wp
_MJX_RENDER_CONTEXT_BUFFERS = {}
_MJX_RENDER_CONTEXT_LOCK = threading.Lock()
_MJX_RENDER_CONTEXT_COUNTER = 0
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -90,8 +86,8 @@ def _render_shim(
flexvert_xpos: wp.array2d(dtype=wp.vec3),
# Registry
rc_id: int,
rgb: wp.array3d(dtype=wp.uint32),
depth: wp.array3d(dtype=wp.float32),
rgb: wp.array2d(dtype=wp.uint32),
depth: wp.array2d(dtype=wp.float32),
):
_m.stat = _s
_m.opt = _o
@@ -136,25 +132,17 @@ def _render_shim(
_d.nworld = cam_xpos.shape[0]
render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id]
render_context.rgb_data = rgb
render_context.depth_data = depth
mjwarp.render(_m, _d, render_context)
wp.copy(rgb, render_context.rgb_data)
wp.copy(depth, render_context.depth_data)
def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[ctx.key]
nrender = render_ctx.nrender
nworld = d.qpos.shape[0]
# get width, height from cam_res
width = int(render_ctx.cam_res.numpy()[0][0])
height = int(render_ctx.cam_res.numpy()[0][1])
output_dims = {
'rgb': (nworld, nrender, height * width),
'depth': (nworld, nrender, height * width),
'rgb': render_ctx.rgb_data_shape,
'depth': render_ctx.depth_data_shape,
}
jf = ffi.jax_callable_variadic_tuple(
@@ -214,50 +202,3 @@ def render(m: types.Model, d: types.Data, ctx: RenderContext):
def render_vmap(unused_axis_size, is_batched, m, d, ctx):
out = render(m, d, ctx)
return out, [True, True]
def create_render_context(
mjm: mujoco.MjModel,
nworld: int,
cam_res: list[tuple[int, int]] | tuple[int, int] | None = None,
render_rgb: list[bool] | bool | None = None,
render_depth: list[bool] | bool | None = None,
use_textures: bool = True,
use_shadows: bool = False,
enabled_geom_groups: list[int] = [0, 1, 2],
cam_active: list[bool] | None = None,
flex_render_smooth: bool = True,
):
from mujoco.mjx.warp import mujoco_warp as mjw
# NOTE: MuJoCo Warp render context expects a Warp Model and Data.
# We create them here but throw them away right after. Preferably,
# the render context should only rely on mujoco.MjModel so we
# do not have to pay the cost of creating dummy Warp Model and Data.
# Some assumptions may be violated if the downstream render context
# builder holds onto the memory of m and d. The API on the MuJoCo
# Warp side needs to be cleaned up.
m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=nworld)
mjw.forward(m, d)
rc = mjw.create_render_context(
mjm=mjm,
m=m,
d=d,
cam_res=cam_res,
use_textures=use_textures,
use_shadows=use_shadows,
render_rgb=render_rgb,
render_depth=render_depth,
enabled_geom_groups=enabled_geom_groups,
cam_active=cam_active,
flex_render_smooth=flex_render_smooth,
)
global _MJX_RENDER_CONTEXT_COUNTER
with _MJX_RENDER_CONTEXT_LOCK:
_MJX_RENDER_CONTEXT_COUNTER += 1
key = _MJX_RENDER_CONTEXT_COUNTER
_MJX_RENDER_CONTEXT_BUFFERS[key] = rc
return RenderContext(key, _owner=True)
mjx/mujoco/mjx/warp/smooth.py
+13 -11
View File
@@ -137,6 +137,7 @@ def _kinematics_shim(
_d.xpos = xpos
_d.xquat = xquat
_d.nworld = nworld
mjwarp.kinematics(_m, _d)
@@ -159,7 +160,7 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
num_outputs=11,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
in_out_argnames=set([
'geom_xmat',
'geom_xpos',
'site_xmat',
@@ -171,8 +172,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
'xmat',
'xpos',
'xquat',
},
stage_in_argnames={
]),
stage_in_argnames=set([
'body_ipos',
'body_iquat',
'body_pos',
@@ -198,8 +199,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
'xmat',
'xpos',
'xquat',
},
stage_out_argnames={
]),
stage_out_argnames=set([
'geom_xmat',
'geom_xpos',
'site_xmat',
@@ -211,7 +212,7 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
'xmat',
'xpos',
'xquat',
},
]),
graph_mode=m.opt._impl.graph_mode,
)
out = jf(
@@ -368,6 +369,7 @@ def _tendon_shim(
_d.wrap_obj = wrap_obj
_d.wrap_xpos = wrap_xpos
_d.nworld = nworld
mjwarp.tendon(_m, _d)
@@ -385,15 +387,15 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
num_outputs=6,
output_dims=output_dims,
vmap_method=None,
in_out_argnames={
in_out_argnames=set([
'ten_J',
'ten_length',
'ten_wrapadr',
'ten_wrapnum',
'wrap_obj',
'wrap_xpos',
},
stage_in_argnames={
]),
stage_in_argnames=set([
'cdof',
'geom_size',
'geom_xmat',
@@ -402,8 +404,8 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
'site_xpos',
'subtree_com',
'ten_length',
},
stage_out_argnames={'ten_length'},
]),
stage_out_argnames=set(['ten_length']),
graph_mode=m.opt._impl.graph_mode,
)
out = jf(
mjx/mujoco/mjx/warp/testspeed.py
+2 -2
View File
@@ -26,7 +26,7 @@ import jax.numpy as jnp
import jax.tree_util
import mujoco
from mujoco import mjx
from mujoco.mjx._src import io
from mujoco.mjx.warp import io
from mujoco.mjx._src import test_util
from mujoco.mjx.warp import collision_driver as wp_collision
from mujoco.mjx.warp import forward as wp_forward
@@ -161,7 +161,7 @@ def benchmark(
d = mjx.refit_bvh(mx, d, rc)
pixels = mjx.render(mx, d, rc)
leaves = jax.tree_util.tree_leaves(pixels)
accum += sum(x[0, 0, 0] for x in leaves) if leaves else 0.0
accum += sum(x[0, 0] for x in leaves if x.size > 0) if leaves else 0.0
return (d, accum), None
mjx/mujoco/mjx/warp/visualize_render.py
+2 -1
View File
@@ -31,6 +31,7 @@ from mujoco.mjx._src import io
from mujoco.mjx._src import render
from mujoco.mjx._src import render_util
from mujoco.mjx._src import test_util
from mujoco.mjx.warp import io
import numpy as np
import warp as wp
@@ -171,7 +172,7 @@ def _main(_: Sequence[str]):
print(f' rgb shape: {rgb_packed.shape}\n')
rgb = render_util.get_rgb(
rgb_packed, _CAMERA_ID.value, _WIDTH.value, _HEIGHT.value
rc, rgb_packed, _CAMERA_ID.value
)
single_path = os.path.join(