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.