From fb47ff01cd2050b2c4b72c2ba7bb77909e824f50 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 13 Feb 2026 11:39:54 -0800 Subject: [PATCH] initial implementation of mjx-warp render. PiperOrigin-RevId: 869822624 Change-Id: Ic20db1dcc2e7287752b273c4ad986066cceccad9 --- mjx/mujoco/mjx/__init__.py | 3 + mjx/mujoco/mjx/_src/bvh.py | 32 +++ mjx/mujoco/mjx/_src/io.py | 43 +++- mjx/mujoco/mjx/_src/render.py | 32 +++ mjx/mujoco/mjx/_src/render_util.py | 73 +++++++ mjx/mujoco/mjx/_src/types.py | 2 +- mjx/mujoco/mjx/warp/bvh.py | 117 +++++++++++ mjx/mujoco/mjx/warp/collision_driver.py | 5 +- mjx/mujoco/mjx/warp/ffi.py | 16 +- mjx/mujoco/mjx/warp/forward.py | 7 +- mjx/mujoco/mjx/warp/render.py | 263 ++++++++++++++++++++++++ mjx/mujoco/mjx/warp/render_test.py | 111 ++++++++++ mjx/mujoco/mjx/warp/smooth.py | 7 +- mjx/mujoco/mjx/warp/testspeed.py | 89 +++++++- mjx/mujoco/mjx/warp/types.py | 50 ++++- mjx/mujoco/mjx/warp/visualize_render.py | 196 ++++++++++++++++++ 16 files changed, 1015 insertions(+), 31 deletions(-) create mode 100644 mjx/mujoco/mjx/_src/bvh.py create mode 100644 mjx/mujoco/mjx/_src/render.py create mode 100644 mjx/mujoco/mjx/_src/render_util.py create mode 100644 mjx/mujoco/mjx/warp/bvh.py create mode 100644 mjx/mujoco/mjx/warp/render.py create mode 100644 mjx/mujoco/mjx/warp/render_test.py create mode 100644 mjx/mujoco/mjx/warp/visualize_render.py diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index d6343c08..35c09673 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -22,6 +22,7 @@ from mujoco.mjx._src.types import Data # pylint:disable=g-importing-member from mujoco.mjx._src.collision_driver import collision +from mujoco.mjx._src.bvh import refit_bvh from mujoco.mjx._src.constraint import make_constraint from mujoco.mjx._src.derivative import deriv_smooth_vel from mujoco.mjx._src.forward import euler @@ -33,6 +34,7 @@ from mujoco.mjx._src.forward import fwd_velocity from mujoco.mjx._src.forward import implicit from mujoco.mjx._src.forward import rungekutta4 from mujoco.mjx._src.inverse import inverse +from mujoco.mjx._src.io import create_render_context from mujoco.mjx._src.io import get_data from mujoco.mjx._src.io import get_data_into from mujoco.mjx._src.io import get_state @@ -43,6 +45,7 @@ from mujoco.mjx._src.io import set_state from mujoco.mjx._src.io import state_size from mujoco.mjx._src.passive import passive from mujoco.mjx._src.ray import ray +from mujoco.mjx._src.render import render from mujoco.mjx._src.sensor import sensor_acc from mujoco.mjx._src.sensor import sensor_pos from mujoco.mjx._src.sensor import sensor_vel diff --git a/mjx/mujoco/mjx/_src/bvh.py b/mjx/mujoco/mjx/_src/bvh.py new file mode 100644 index 00000000..a35f4f6e --- /dev/null +++ b/mjx/mujoco/mjx/_src/bvh.py @@ -0,0 +1,32 @@ +# 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. +# ============================================================================== +"""BVH helpers for MJX.""" + +from typing import Any +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import Impl +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import mujoco.mjx.warp as mjxw + + +def refit_bvh(m: Model, d: Data, ctx: Any): + """Refit the scene BVH for the current pose.""" + if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED: + from mujoco.mjx.warp import bvh as mjxw_bvh # pylint: disable=g-import-not-at-top # pytype: disable=import-error + return mjxw_bvh.refit_bvh(m, d, ctx) + + raise NotImplementedError('refit_bvh only implemented for MuJoCo Warp.') diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 108458e2..8f3ec681 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -17,7 +17,7 @@ import copy import logging import os -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union import warnings import jax @@ -52,7 +52,9 @@ def _is_cuda_gpu_device(device: jax.Device) -> bool: def _check_warp_installed(): if not mjxw.WARP_INSTALLED: - raise RuntimeError('warp-lang is not installed. Cannot use WARP implementation of MJX.') + raise RuntimeError( + 'warp-lang is not installed. Cannot use WARP implementation of MJX.' + ) def _resolve_impl( @@ -121,8 +123,7 @@ def _check_impl_device_compatibility( if impl == types.Impl.WARP: if not _is_cuda_gpu_device(device): raise AssertionError( - 'Warp implementation requires a CUDA GPU device, got ' - f'{device}.' + f'Warp implementation requires a CUDA GPU device, got {device}.' ) _check_warp_installed() @@ -999,8 +1000,7 @@ def make_data( if isinstance(m, types.Model) and m.impl != impl: raise ValueError( - f'Model impl {m.impl} does not match make_data ' - f'implementation {impl}.' + f'Model impl {m.impl} does not match make_data implementation {impl}.' ) if impl == types.Impl.JAX: @@ -1546,8 +1546,7 @@ def _get_data_into( all_fields = types.Data.fields() + types.DataC.fields() else: raise NotImplementedError( - f'get_data_into for implementation "{d.impl}" not implemented' - ' yet.' + f'get_data_into for implementation "{d.impl}" not implemented yet.' ) for field in all_fields: @@ -1699,8 +1698,9 @@ def _get_data_into_cpp( # mjx.Model which we don't have access to in this function. fields_to_check = ['qpos', 'qvel', 'act', 'mocap_pos', 'mocap_quat'] for i in range(batch_size): - d_i: types.Data = jax.tree_util.tree_map( - lambda x, i=i: x[i], d) if batched else d + d_i: types.Data = ( + jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d + ) src_data = mj_data_list[i] needs_syncing = False @@ -1939,3 +1939,26 @@ def set_state( offset += size return d.replace(**updates) + + +def create_render_context( + mjm: mujoco.MjModel, + nworld: int, + **kwargs, +): + """Creates a render context. + + Args: + mjm: the MuJoCo model + nworld: number of worlds to render. We must hardcode the nworld + because Warp creates arrays of size nworld that are not exposed + to JAX. Thus we cannot use JAX transforms like vmap with the + render context. + **kwargs: forwarded to the render context constructor. + + Returns: + 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) diff --git a/mjx/mujoco/mjx/_src/render.py b/mjx/mujoco/mjx/_src/render.py new file mode 100644 index 00000000..5577ffa6 --- /dev/null +++ b/mjx/mujoco/mjx/_src/render.py @@ -0,0 +1,32 @@ +# 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. +# ============================================================================== +"""Render helpers for MJX.""" + +from typing import Any +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import Impl +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import mujoco.mjx.warp as mjxw + + +def render(m: Model, d: Data, ctx: Any) -> Data: + """Render.""" + if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.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.render(m, d, ctx) + + raise NotImplementedError('render only implemented for MuJoCo Warp.') diff --git a/mjx/mujoco/mjx/_src/render_util.py b/mjx/mujoco/mjx/_src/render_util.py new file mode 100644 index 00000000..ac54e13d --- /dev/null +++ b/mjx/mujoco/mjx/_src/render_util.py @@ -0,0 +1,73 @@ +# 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. +# ============================================================================== +"""JAX render utilities for unpacking render output from MuJoCo Warp.""" + +import jax +import jax.numpy as jnp + + +def get_rgb( + 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) + 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] + 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 + rgb = jnp.stack([r, g, b], axis=-1) + nworld = rgb_data.shape[0] + return rgb.reshape(nworld, height, width, 3) + + +def get_depth( + 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) + 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] + nworld = depth_data.shape[0] + depth = jnp.clip(raw / depth_scale, 0.0, 1.0) + return depth.reshape(nworld, height, width) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index d7521a88..fcdc41d6 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -917,7 +917,7 @@ class Model(PyTreeNode): tex_adr: np.ndarray tex_data: np.ndarray mat_rgba: jax.Array - mat_texid: np.ndarray + mat_texid: jax.Array pair_dim: np.ndarray pair_geom1: np.ndarray pair_geom2: np.ndarray diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py new file mode 100644 index 00000000..3c9789e7 --- /dev/null +++ b/mjx/mujoco/mjx/warp/bvh.py @@ -0,0 +1,117 @@ +# 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. +# ============================================================================== +"""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.types import RenderContext +import mujoco.mjx.third_party.mujoco_warp as mjwarp +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} +) + + +@ffi.format_args_for_warp +def _refit_bvh_shim( + # Model + 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 + 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), +): + _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.geom_xmat = geom_xmat + _d.geom_xpos = geom_xpos + _d.flexvert_xpos = flexvert_xpos + _d.nworld = geom_xpos.shape[0] + + render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id] + 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] + + jf = ffi.jax_callable_variadic_tuple( + _refit_bvh_shim, + num_outputs=1, + output_dims={'geom_xpos_out': (nworld, ngeom, 3)}, + vmap_method=None, + ) + out = jf( + 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, + d.geom_xmat, + d.geom_xpos, + d.flexvert_xpos, + ctx.key, + ) + return d.replace(geom_xpos=out[0]) + + +@jax.custom_batching.custom_vmap +@functools.partial(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] diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index f4be0dae..70cb1b47 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -1,4 +1,4 @@ -# Copyright 2025 DeepMind Technologies Limited +# 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. @@ -14,6 +14,7 @@ # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" + import dataclasses import jax from mujoco.mjx._src import types @@ -399,6 +400,8 @@ def _collision_jax_impl(m: types.Model, d: types.Data): @ffi.marshal_jax_warp_callable def collision(m: types.Model, d: types.Data): return _collision_jax_impl(m, d) + + @collision.def_vmap @ffi.marshal_custom_vmap def collision_vmap(unused_axis_size, is_batched, m, d): diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py index fde118d3..18632162 100644 --- a/mjx/mujoco/mjx/warp/ffi.py +++ b/mjx/mujoco/mjx/warp/ffi.py @@ -252,11 +252,11 @@ def _squeeze_dim(leaf_expanded: Any, leaf: Any) -> Any: return leaf_expanded -def marshal_jax_warp_callable(func, raw_output: bool = False): +def marshal_jax_warp_callable(func, skip_output_dim_reshape: bool = False): """Marshal fields into a MuJoCo Warp function.""" @functools.wraps(func) - def wrapper(m, d): + def wrapper(m, d, *extra_args): # Expand dims for Warp implicit vmap before calling into the FFI wrapped # function. m_expanded = jax.tree.map_with_path( @@ -271,9 +271,9 @@ def marshal_jax_warp_callable(func, raw_output: bool = False): ), d, ) - d_expanded_result = func(m_expanded, d_expanded) + d_expanded_result = func(m_expanded, d_expanded, *extra_args) - if raw_output: + if skip_output_dim_reshape: return d_expanded_result d_result = jax.tree.map(_squeeze_dim, d_expanded_result, d) return d_result @@ -360,11 +360,11 @@ def _check_leading_dim( ) -def marshal_custom_vmap(vmap_func, raw_output: bool = False): +def marshal_custom_vmap(vmap_func, skip_output_dim_reshape: bool = False): """Marshal fields for a custom vmap into an MuJoCo Warp function.""" @functools.wraps(vmap_func) - def wrapper(axis_size, is_batched, m, d): + def wrapper(axis_size, is_batched, m, d, *extra_args): # Vmappable data fields may not have been broadcasted if vmap_func is called # within a vmap trace. Since data fields are read/write in warp, we need to # explicitly broadcast them here. @@ -395,9 +395,9 @@ def marshal_custom_vmap(vmap_func, raw_output: bool = False): d_broadcast, ) d_broadcast_flat_result, out_batched = vmap_func( - axis_size, is_batched, m_flat, d_broadcast_flat + axis_size, is_batched, m_flat, d_broadcast_flat, *extra_args ) - if raw_output: + if skip_output_dim_reshape: return d_broadcast_flat_result, out_batched # Explicitly mark MuJoCo Warp data fields as batched after vmapping is done. diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 2604244d..65e69dbc 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -1,4 +1,4 @@ -# Copyright 2025 DeepMind Technologies Limited +# 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. @@ -14,6 +14,7 @@ # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" + import dataclasses import jax from mujoco.mjx._src import types @@ -1809,6 +1810,8 @@ def _forward_jax_impl(m: types.Model, d: types.Data): @ffi.marshal_jax_warp_callable def forward(m: types.Model, d: types.Data): return _forward_jax_impl(m, d) + + @forward.def_vmap @ffi.marshal_custom_vmap def forward_vmap(unused_axis_size, is_batched, m, d): @@ -3604,6 +3607,8 @@ def _step_jax_impl(m: types.Model, d: types.Data): @ffi.marshal_jax_warp_callable def step(m: types.Model, d: types.Data): return _step_jax_impl(m, d) + + @step.def_vmap @ffi.marshal_custom_vmap def step_vmap(unused_axis_size, is_batched, m, d): diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py new file mode 100644 index 00000000..05a2d207 --- /dev/null +++ b/mjx/mujoco/mjx/warp/render.py @@ -0,0 +1,263 @@ +# 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. +# ============================================================================== +"""DO NOT EDIT. This file is auto-generated.""" + +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.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} +) +_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 _render_shim( + # Model + geom_dataid: wp.array(dtype=int), + geom_matid: wp.array2d(dtype=int), + geom_rgba: wp.array2d(dtype=wp.vec4), + geom_size: wp.array2d(dtype=wp.vec3), + geom_type: wp.array(dtype=int), + light_active: wp.array2d(dtype=bool), + light_castshadow: wp.array2d(dtype=bool), + light_type: wp.array2d(dtype=int), + mat_rgba: wp.array2d(dtype=wp.vec4), + mat_texid: wp.array3d(dtype=int), + mat_texrepeat: wp.array2d(dtype=wp.vec2), + mesh_face: wp.array(dtype=wp.vec3i), + mesh_faceadr: wp.array(dtype=int), + ncam: int, + ngeom: int, + nlight: 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), + cam_projection: wp.array(dtype=int), + cam_fovy: wp.array2d(dtype=wp.float32), + cam_sensorsize: wp.array(dtype=wp.vec2), + cam_intrinsic: wp.array2d(dtype=wp.vec4), + # Data + cam_xmat: wp.array2d(dtype=wp.mat33), + cam_xpos: wp.array2d(dtype=wp.vec3), + geom_xmat: wp.array2d(dtype=wp.mat33), + geom_xpos: wp.array2d(dtype=wp.vec3), + light_xdir: wp.array2d(dtype=wp.vec3), + light_xpos: wp.array2d(dtype=wp.vec3), + flexvert_xpos: wp.array2d(dtype=wp.vec3), + # Registry + rc_id: int, + rgb: wp.array3d(dtype=wp.uint32), + depth: wp.array3d(dtype=wp.float32), +): + _m.stat = _s + _m.opt = _o + _d.efc = _e + _d.contact = _c + _m.geom_dataid = geom_dataid + _m.geom_matid = geom_matid + _m.geom_rgba = geom_rgba + _m.geom_size = geom_size + _m.geom_type = geom_type + _m.light_active = light_active + _m.light_castshadow = light_castshadow + _m.light_type = light_type + _m.mat_rgba = mat_rgba + _m.mat_texid = mat_texid + _m.mat_texrepeat = mat_texrepeat + _m.mesh_face = mesh_face + _m.mesh_faceadr = mesh_faceadr + _m.ncam = ncam + _m.ngeom = ngeom + _m.nlight = nlight + _m.cam_projection = cam_projection + _m.cam_fovy = cam_fovy + _m.cam_sensorsize = cam_sensorsize + _m.cam_intrinsic = cam_intrinsic + _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.cam_xmat = cam_xmat + _d.cam_xpos = cam_xpos + _d.geom_xmat = geom_xmat + _d.geom_xpos = geom_xpos + _d.light_xdir = light_xdir + _d.light_xpos = light_xpos + _d.flexvert_xpos = flexvert_xpos + + _d.nworld = cam_xpos.shape[0] + + render_context = _MJX_RENDER_CONTEXT_BUFFERS[rc_id] + 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), + } + + jf = ffi.jax_callable_variadic_tuple( + _render_shim, + num_outputs=2, + output_dims=output_dims, + vmap_method=None, + ) + out = jf( + m.geom_dataid, + m.geom_matid, + m.geom_rgba, + m.geom_size, + m.geom_type, + m.light_active, + m.light_castshadow, + m.light_type, + m.mat_rgba, + m.mat_texid, + m.mat_texrepeat, + m.mesh_face, + m.mesh_faceadr, + m.ncam, + m.ngeom, + m.nlight, + m.nflex, + m.nflexelemdata, + m.nflexvert, + m.flex_dim, + m.flex_elem, + m.flex_elemnum, + m.flex_vertadr, + m.cam_projection, + m.cam_fovy, + m.cam_sensorsize, + m.cam_intrinsic, + d.cam_xmat, + d.cam_xpos, + d.geom_xmat, + d.geom_xpos, + d.light_xdir, + d.light_xpos, + d.flexvert_xpos, + ctx.key, + ) + return out + + +@jax.custom_batching.custom_vmap +@functools.partial(ffi.marshal_jax_warp_callable, skip_output_dim_reshape=True) +def render(m: types.Model, d: types.Data, ctx: RenderContext): + return _render_jax_impl(m, d, ctx) + + +@render.def_vmap +@functools.partial(ffi.marshal_custom_vmap, skip_output_dim_reshape=True) +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) diff --git a/mjx/mujoco/mjx/warp/render_test.py b/mjx/mujoco/mjx/warp/render_test.py new file mode 100644 index 00000000..9e8c53dc --- /dev/null +++ b/mjx/mujoco/mjx/warp/render_test.py @@ -0,0 +1,111 @@ +# 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. +# ============================================================================== +import functools +import os + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import io +from mujoco.mjx._src import forward +import mujoco.mjx.warp as mjxw +from mujoco.mjx.warp import test_util as tu +from mujoco.mjx.warp import warp as wp # pylint: disable=g-importing-member +import numpy as np + +_FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1' + + +class RenderTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + if mjxw.WARP_INSTALLED: + tempdir = '/tmp/wp_kernel_cache_dir_RenderTest' + wp.config.kernel_cache_dir = tempdir + np.random.seed(0) + + @parameterized.product( + xml=( + 'humanoid/humanoid.xml', + ), + batch_size=(1, 16), + ) + def test_render(self, xml: str, batch_size: int): + """Tests MJX render pipeline.""" + if not _FORCE_TEST: + if not mjxw.WARP_INSTALLED: + self.skipTest('Warp not installed.') + if not io.has_cuda_gpu_device(): + self.skipTest('No CUDA GPU device available.') + + m = tu.load_test_file(xml) + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + + mx = mjx.put_model(m, impl='warp') + + worldids = jp.arange(batch_size) + dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids) + + key = jax.random.PRNGKey(0) + keys = jax.random.split(key, batch_size) + qpos0 = jp.array(m.qpos0) + rand_qpos = jax.vmap( + lambda k: qpos0 + jax.random.uniform( + k, (m.nq,), minval=-0.2, maxval=0.05 + ) + )(keys) + dx_batch = jax.vmap( + lambda dx, q: dx.replace(qpos=q) + )(dx_batch, rand_qpos) + + dx_batch = jax.jit( + jax.vmap(forward.forward, in_axes=(None, 0)) + )(mx, dx_batch) + + width, height = 32, 32 + rc = mjx.create_render_context( + mjm=m, + nworld=batch_size, + cam_res=(width, height), + use_textures=True, + use_shadows=True, + render_rgb=True, + render_depth=True, + enabled_geom_groups=[0, 1, 2], + ) + + dx_batch = jax.jit( + jax.vmap(mjx.refit_bvh, in_axes=(None, 0, None)) + )(mx, dx_batch, rc) + out_batch = jax.jit( + jax.vmap(mjx.render, in_axes=(None, 0, None)) + )(mx, dx_batch, rc) + + rgb = np.asarray(out_batch[0]) + depth = np.asarray(out_batch[1]) + + self.assertGreater(np.count_nonzero(rgb), 0) + self.assertGreater(np.count_nonzero(depth), 0) + self.assertNotEqual(np.unique(rgb).shape[0], 1) + self.assertNotEqual(np.unique(depth).shape[0], 1) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index ac5f3470..a3521ba0 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -1,4 +1,4 @@ -# Copyright 2025 DeepMind Technologies Limited +# 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. @@ -14,6 +14,7 @@ # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" + import dataclasses import jax from mujoco.mjx._src import types @@ -277,6 +278,8 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data): @ffi.marshal_jax_warp_callable def kinematics(m: types.Model, d: types.Data): return _kinematics_jax_impl(m, d) + + @kinematics.def_vmap @ffi.marshal_custom_vmap def kinematics_vmap(unused_axis_size, is_batched, m, d): @@ -456,6 +459,8 @@ def _tendon_jax_impl(m: types.Model, d: types.Data): @ffi.marshal_jax_warp_callable def tendon(m: types.Model, d: types.Data): return _tendon_jax_impl(m, d) + + @tendon.def_vmap @ffi.marshal_custom_vmap def tendon_vmap(unused_axis_size, is_batched, m, d): diff --git a/mjx/mujoco/mjx/warp/testspeed.py b/mjx/mujoco/mjx/warp/testspeed.py index 69d42517..5a131e55 100644 --- a/mjx/mujoco/mjx/warp/testspeed.py +++ b/mjx/mujoco/mjx/warp/testspeed.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Run benchmarks.""" +"""Run benchmarks for MJX-Warp.""" import functools import os @@ -23,8 +23,10 @@ from absl import app from absl import flags import jax import jax.numpy as jnp +import jax.tree_util import mujoco from mujoco import mjx +from mujoco.mjx._src 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 @@ -64,6 +66,20 @@ _BENCHMARK = flags.DEFINE_enum( ['jax_warp', 'jax', 'warp'], 'Which benchmark to run.', ) +# Render flags +_RENDER_WIDTH = flags.DEFINE_integer('render_width', 64, 'render width') +_RENDER_HEIGHT = flags.DEFINE_integer('render_height', 64, 'render height') +_RENDER_RGB = flags.DEFINE_boolean('render_rgb', False, 'render RGB image') +_RENDER_DEPTH = flags.DEFINE_boolean( + 'render_depth', False, 'render depth image' +) +_RENDER_USE_TEXTURES = flags.DEFINE_boolean( + 'render_textures', True, 'use textures' +) +_RENDER_USE_SHADOWS = flags.DEFINE_boolean( + 'render_shadows', False, 'use shadows' +) + _COMPILER_OPTIONS = {'xla_gpu_graph_min_graph_size': 1} jax_jit = functools.partial(jax.jit, compiler_options=_COMPILER_OPTIONS) @@ -101,6 +117,7 @@ def benchmark( nstep: int = 1000, nenv: int = 8192, unroll_steps: int = 4, + render: bool = False, ) -> Tuple[float, float, int]: """Benchmark a model.""" @@ -119,13 +136,39 @@ def benchmark( jax.block_until_ready(d) + rc = None + if render: + ncam = max(1, m.ncam) + rc = io.create_render_context( + mjm=m, + nworld=nenv, + cam_res=(_RENDER_WIDTH.value, _RENDER_HEIGHT.value), + use_textures=_RENDER_USE_TEXTURES.value, + use_shadows=_RENDER_USE_SHADOWS.value, + render_rgb=[_RENDER_RGB.value] * ncam, + render_depth=[_RENDER_DEPTH.value] * ncam, + enabled_geom_groups=[0, 1, 2], + ) + @jax_jit def unroll(d): - def fn(d, _): + def fn(carry, _): + d, accum = carry d = d.replace(qpos=d.qpos + 0 * d.qpos) - return step_fn(mx, d), None + d = step_fn(mx, d) - return jax.lax.scan(fn, d, None, length=nstep, unroll=unroll_steps) + if render: + 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 + + return (d, accum), None + + (d_final, accum_final), _ = jax.lax.scan( + fn, (d, jnp.array(0.0)), None, length=nstep, unroll=unroll_steps + ) + return d_final, accum_final jit_time, run_time = _measure(unroll, d) steps = nstep * nenv @@ -147,6 +190,7 @@ def benchmark_raw_warp( nenv: int = 8192, unroll_steps: int = 4, function: str = 'kinematics', + render: bool = False, ): """Benchmarks raw warp.""" del unroll_steps @@ -178,8 +222,27 @@ def benchmark_raw_warp( else: raise NotImplementedError(f'{function} not implemented in speed test.') + ncam = max(1, m.ncam) + rc = None + if render: + mjwarp.forward(mw, dw) + rc = mjwarp.create_render_context( + m, mw, dw, + (_RENDER_WIDTH.value, _RENDER_HEIGHT.value), + [_RENDER_RGB.value] * ncam, + [_RENDER_DEPTH.value] * ncam, + _RENDER_USE_TEXTURES.value, + _RENDER_USE_SHADOWS.value, + ) + + def run_fn(m, d): + fn(m, d) + if render: + mjwarp.refit_bvh(m, d, rc) + mjwarp.render(m, d, rc) + start = time.time() - graph = _compile_fn(fn, mw, dw) + graph = _compile_fn(run_fn, mw, dw) jit_time = time.time() - start start = time.time() @@ -241,17 +304,27 @@ def _main(_: Sequence[str]): print(f' timestep : {m.opt.timestep}') print(f' unroll : {unroll}') print(f' benchmark : {benchmark_type}') - print(f' graph_mode : {_GRAPH_MODE.value}\n') + print(f' graph_mode : {_GRAPH_MODE.value}') + is_render = _RENDER_DEPTH.value or _RENDER_RGB.value + if is_render: + print( + f' render resolution : {_RENDER_WIDTH.value}x{_RENDER_HEIGHT.value}' + ) + print(f' textures : {_RENDER_USE_TEXTURES.value}') + print(f' shadows : {_RENDER_USE_SHADOWS.value}') + print() if benchmark_type == 'jax_warp': - jit_time, run_time, steps = benchmark(m, mw, func_warp, nstep, nenv, unroll) + jit_time, run_time, steps = benchmark( + m, mw, func_warp, nstep, nenv, unroll, is_render + ) 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_ + m, nstep, nenv, unroll, function=function_, render=is_render ) print(' Pure WARP:') else: diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index dac6fa55..d39d4861 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/mjx/mujoco/mjx/warp/types.py @@ -1,4 +1,4 @@ -# Copyright 2025 DeepMind Technologies Limited +# 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. @@ -13,8 +13,10 @@ # limitations under the License. # ============================================================================== """MJX Warp types. + DO NOT EDIT. This file is auto-generated. """ + import dataclasses import typing from typing import Tuple @@ -23,6 +25,7 @@ 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: @@ -32,6 +35,7 @@ else: GraphMode = int PyTreeNode = mjx_dataclasses.PyTreeNode + @dataclasses.dataclass(frozen=True) @tree_util.register_pytree_node_class class TileSet: @@ -64,6 +68,7 @@ class BlockDim: TODO(team): experimental and may be removed """ + actuator_velocity: int cholesky_factorize: int cholesky_factorize_solve: int @@ -91,12 +96,48 @@ class BlockDim: return cls(*children) +@dataclasses.dataclass(frozen=True) +@tree_util.register_pytree_node_class +class RenderContext: + """Render context handle with automatic cleanup. + + The key is static (aux_data) so JAX doesn't trace it, allowing the Warp FFI + to receive a concrete int value. Only the original instance (_owner=True) + cleans up on deletion; JAX copies have _owner=False. + """ + + key: int + _owner: bool = True + + def tree_flatten(self): + return ((), (self.key, False)) + + @classmethod + def tree_unflatten(cls, aux_data, children): + del children + key, owner = aux_data + return cls(key, owner) + + def __del__(self): + if not self._owner: + return + lock = globals().get('_MJX_RENDER_CONTEXT_LOCK') + buffers = globals().get('_MJX_RENDER_CONTEXT_BUFFERS') + if lock is None or buffers is None: + return + with lock: + buffers.pop(self.key, None) + + class StatisticWarp(PyTreeNode): """Derived fields from Statistic.""" + meaninertia: jax.Array + class OptionWarp(PyTreeNode): """Derived fields from Option.""" + broadphase: int broadphase_filter: int ccd_iterations: int @@ -111,8 +152,10 @@ class OptionWarp(PyTreeNode): sdf_initpoints: int sdf_iterations: int + class ModelWarp(PyTreeNode): """Derived fields from Model.""" + M_colind: np.ndarray M_rowadr: np.ndarray M_rownnz: np.ndarray @@ -245,8 +288,10 @@ class ModelWarp(PyTreeNode): wrap_site_adr: np.ndarray wrap_site_pair_adr: np.ndarray + class DataWarp(PyTreeNode): """Derived fields from Data.""" + actuator_moment: jax.Array actuator_velocity: jax.Array cacc: jax.Array @@ -312,6 +357,8 @@ class DataWarp(PyTreeNode): wrap_obj: jax.Array wrap_xpos: jax.Array shape = property(lambda self: self.cacc.shape) + + DATA_NON_VMAP = { 'contact__dim', 'contact__dist', @@ -335,6 +382,7 @@ DATA_NON_VMAP = { 'nworld', } + def _to_elt(cont, _, d, axis): return DataWarp(**{ f.name: ( diff --git a/mjx/mujoco/mjx/warp/visualize_render.py b/mjx/mujoco/mjx/warp/visualize_render.py new file mode 100644 index 00000000..4198c62f --- /dev/null +++ b/mjx/mujoco/mjx/warp/visualize_render.py @@ -0,0 +1,196 @@ +# 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. +# ============================================================================== +"""Visualize batch rendering output from MuJoCo Warp for debugging.""" + +import functools +import os +from typing import Sequence + +from absl import app +from absl import flags +import jax +import jax.numpy as jp +import mediapy as media +import mujoco +from mujoco import mjx +from mujoco.mjx._src import bvh +from mujoco.mjx._src import forward +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 +import numpy as np +import warp as wp + + +_MODELFILE = flags.DEFINE_string( + 'modelfile', + 'humanoid/humanoid.xml', + 'path to model', +) +_NWORLD = flags.DEFINE_integer( + 'nworld', 4, 'number of worlds to render' +) +_WIDTH = flags.DEFINE_integer('width', 512, 'image width') +_HEIGHT = flags.DEFINE_integer('height', 512, 'image height') +_CAMERA_ID = flags.DEFINE_integer( + 'camera_id', 1, 'camera id to visualize' +) +_OUTPUT_DIR = flags.DEFINE_string( + 'output_dir', '/tmp/visualize_render', 'output directory' +) +_USE_TEXTURES = flags.DEFINE_boolean( + 'use_textures', True, 'enable textures' +) +_USE_SHADOWS = flags.DEFINE_boolean( + 'use_shadows', True, 'enable shadows' +) +_WP_KERNEL_CACHE_DIR = flags.DEFINE_string( + 'wp_kernel_cache_dir', + '/tmp/wp_kernel_cache_dir_visualize_render', + 'warp kernel cache directory', +) + +_COMPILER_OPTIONS = {'xla_gpu_graph_min_graph_size': 1} +jax_jit = functools.partial( + jax.jit, compiler_options=_COMPILER_OPTIONS +) + + +def _save_single(rgb, out_path): + """Save first world as a single image.""" + img = np.asarray(rgb[0]) + img_uint8 = (img * 255).astype(np.uint8) + media.write_image(out_path, img_uint8) + print(f' single image: {out_path}') + + +def _save_tiled(rgb, out_path): + """Save all worlds as a tiled grid.""" + nworld, height, width, _ = rgb.shape + cols = int(np.ceil(np.sqrt(nworld))) + rows = int(np.ceil(nworld / cols)) + canvas = np.zeros( + (rows * height, cols * width, 3), dtype=np.uint8 + ) + + for w in range(nworld): + img_uint8 = (np.asarray(rgb[w]) * 255).astype( + np.uint8 + ) + r, c = w // cols, w % cols + y0, y1 = r * height, (r + 1) * height + x0, x1 = c * width, (c + 1) * width + canvas[y0:y1, x0:x1, :] = img_uint8 + + media.write_image(out_path, canvas) + print(f' tiled image: {out_path}') + + +def _main(_: Sequence[str]): + os.environ['MJX_WARP_ENABLED'] = 'true' + + wp.config.kernel_cache_dir = _WP_KERNEL_CACHE_DIR.value + + os.makedirs(_OUTPUT_DIR.value, exist_ok=True) + + try: + m = test_util.load_test_file(_MODELFILE.value) + except Exception: + m = mujoco.MjModel.from_xml_path(_MODELFILE.value) + + print('visualize_render.py:\n') + print(f' modelfile : {_MODELFILE.value}') + print(f' nworld : {_NWORLD.value}') + print(f' resolution : {_WIDTH.value}x{_HEIGHT.value}') + print(f' camera_id : {_CAMERA_ID.value}') + print(f' use_textures: {_USE_TEXTURES.value}') + print(f' use_shadows : {_USE_SHADOWS.value}') + print(f' output_dir : {_OUTPUT_DIR.value}\n') + + mx = mjx.put_model(m, impl='warp') + + worldids = jp.arange(_NWORLD.value) + + @jax.vmap + def init(worldid): + dx = mjx.make_data(m, impl='warp') + rng = jax.random.PRNGKey(worldid) + qpos0 = jp.array(m.qpos0) + qpos = qpos0 + jax.random.uniform( + rng, (m.nq,), minval=-0.2, maxval=0.05 + ) + return dx.replace(qpos=qpos) + + print('initializing data...') + dx_batch = jax_jit(init)(worldids) + + print('running forward...') + dx_batch = jax_jit( + jax.vmap(forward.forward, in_axes=(None, 0)) + )(mx, dx_batch) + + print('creating render context...') + rc = io.create_render_context( + mjm=m, + nworld=_NWORLD.value, + cam_res=(_WIDTH.value, _HEIGHT.value), + use_textures=_USE_TEXTURES.value, + use_shadows=_USE_SHADOWS.value, + render_rgb=True, + render_depth=True, + enabled_geom_groups=[0, 1, 2], + ) + + print('rendering...') + dx_batch = jax_jit( + jax.vmap( + bvh.refit_bvh, in_axes=(None, 0, None) + ) + )(mx, dx_batch, rc) + + out_batch = jax_jit( + jax.vmap( + render.render, in_axes=(None, 0, None) + ) + )(mx, dx_batch, rc) + + rgb_packed = out_batch[0] + print(f' rgb shape: {rgb_packed.shape}\n') + + rgb = render_util.get_rgb( + rgb_packed, _CAMERA_ID.value, _WIDTH.value, _HEIGHT.value + ) + + single_path = os.path.join( + _OUTPUT_DIR.value, f'camera_{_CAMERA_ID.value}.png' + ) + _save_single(rgb, single_path) + + if _NWORLD.value > 1: + tiled_path = os.path.join( + _OUTPUT_DIR.value, f'tiled_{_CAMERA_ID.value}.png' + ) + _save_tiled(rgb, tiled_path) + + print('\ndone.') + + +def main(): + app.run(_main) + + +if __name__ == '__main__': + main()