From 67b6ab0589263839f1486101190127592553f048 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 26 Feb 2026 17:25:48 -0800 Subject: [PATCH] Update mjx-warp docs. PiperOrigin-RevId: 875966048 Change-Id: I333e8562f3db333f6c64af67cee1f735cb21232d --- doc/mjx.rst | 8 ++++++++ mjx/mujoco/mjx/_src/render_util.py | 4 ++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/doc/mjx.rst b/doc/mjx.rst index b6a79671..48247b93 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -262,6 +262,14 @@ volume hierarchy (BVH) and executing the raycaster: rgb, d = render_fn(mx, d, rc.pytree()) +.. WARNING:: + The batch dimension ``nworld`` is fixed when the render context is created via + :func:`~mujoco.mjx.create_render_context` since the underlying Warp render context allocates + buffers for ``nworld`` environments that are not visible to JAX. :func:`~mujoco.mjx.render` + will always return outputs with a leading batch dimension of size ``nworld``. Because of this, + there is a known issue where :func:`~mujoco.mjx.render` does not play nice with a + ``jax.vmap(jax.lax.scan)``. + Multi-GPU rendering with ``pmap`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/render_util.py b/mjx/mujoco/mjx/_src/render_util.py index 54ab6af7..0ef94281 100644 --- a/mjx/mujoco/mjx/_src/render_util.py +++ b/mjx/mujoco/mjx/_src/render_util.py @@ -32,7 +32,7 @@ def get_rgb( """Unpack uint32 ABGR pixel data into float32 RGB. Args: - rc: RenderContext or RenderContextPytree (both have .key). + rc: RenderContextPytree. cam_id: Camera index to extract. rgb_data: Packed render output, shape (total_pixels,) as uint32. @@ -78,7 +78,7 @@ def get_depth( """Extract and normalize depth data for a camera. Args: - rc: RenderContext or RenderContextPytree (both have .key). + rc: RenderContextPytree. cam_id: Camera index to extract. depth_data: Raw depth output, shape (total_pixels,) as float32. depth_scale: Scale factor for normalizing depth values.