Update mjx-warp docs.

PiperOrigin-RevId: 875966048
Change-Id: I333e8562f3db333f6c64af67cee1f735cb21232d

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``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

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.