From bbfe7fe0e22863d1dcb19bc74b76e020d924c2c8 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 13 Feb 2026 13:48:57 -0800 Subject: [PATCH] Update mjx-warp render with partial codegen. PiperOrigin-RevId: 869874195 Change-Id: I71a16a6c192873786a09211dc4d326c3a4115ad9 --- mjx/mujoco/mjx/_src/io.py | 4 +- mjx/mujoco/mjx/_src/render_util.py | 42 +++++++---- mjx/mujoco/mjx/warp/bvh.py | 94 +++++++++++++++---------- mjx/mujoco/mjx/warp/collision_driver.py | 11 +-- mjx/mujoco/mjx/warp/forward.py | 26 +++---- mjx/mujoco/mjx/warp/io.py | 60 ++++++++++++++++ mjx/mujoco/mjx/warp/render.py | 75 +++----------------- mjx/mujoco/mjx/warp/smooth.py | 24 ++++--- mjx/mujoco/mjx/warp/testspeed.py | 4 +- mjx/mujoco/mjx/warp/visualize_render.py | 3 +- 10 files changed, 195 insertions(+), 148 deletions(-) create mode 100644 mjx/mujoco/mjx/warp/io.py diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 8f3ec681..1e1101a5 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -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) diff --git a/mjx/mujoco/mjx/_src/render_util.py b/mjx/mujoco/mjx/_src/render_util.py index ac54e13d..f4a05d9b 100644 --- a/mjx/mujoco/mjx/_src/render_util.py +++ b/mjx/mujoco/mjx/_src/render_util.py @@ -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) + diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py index 3c9789e7..b6b627cd 100644 --- a/mjx/mujoco/mjx/warp/bvh.py +++ b/mjx/mujoco/mjx/warp/bvh.py @@ -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] diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index 70cb1b47..a1934fb1 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -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( diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 65e69dbc..9b813d23 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -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( diff --git a/mjx/mujoco/mjx/warp/io.py b/mjx/mujoco/mjx/warp/io.py new file mode 100644 index 00000000..b4f6c517 --- /dev/null +++ b/mjx/mujoco/mjx/warp/io.py @@ -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) diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index 05a2d207..cde132ec 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/mjx/mujoco/mjx/warp/render.py @@ -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) diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py index a3521ba0..f645f999 100644 --- a/mjx/mujoco/mjx/warp/smooth.py +++ b/mjx/mujoco/mjx/warp/smooth.py @@ -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( diff --git a/mjx/mujoco/mjx/warp/testspeed.py b/mjx/mujoco/mjx/warp/testspeed.py index 5a131e55..2c3a7792 100644 --- a/mjx/mujoco/mjx/warp/testspeed.py +++ b/mjx/mujoco/mjx/warp/testspeed.py @@ -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 diff --git a/mjx/mujoco/mjx/warp/visualize_render.py b/mjx/mujoco/mjx/warp/visualize_render.py index 4198c62f..f036f5af 100644 --- a/mjx/mujoco/mjx/warp/visualize_render.py +++ b/mjx/mujoco/mjx/warp/visualize_render.py @@ -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(