Update mjx-warp render with partial codegen.

PiperOrigin-RevId: 869874195
Change-Id: I71a16a6c192873786a09211dc4d326c3a4115ad9

Author: btaba

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)

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)
+

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]

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(

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(