Allow mjx.get_rgb/get_depth to accept batched render output

[tkelestemur]

Summary

• Make mjx.get_rgb() and mjx.get_depth() accept packed buffers with leading batch dimensions, so callers can pass mjx.render() outputs directly for both single-world and batched rendering.
• Preserve any leading dimensions in the unpacked outputs, returning (..., H, W, 3) for RGB and (..., H, W, 1) for depth.
• Update render_util tests to cover direct batched inputs and remove the unnecessary Warp/CUDA dependency from these JAX-side unpacking tests.

Why

mjx.render() always returns packed outputs with a leading world axis, including when nworld == 1, but the public unpacking helpers previously only handled a single packed buffer. This change makes the helper API match the render API and removes the need for downstream callers to manually squeeze or vmap before unpacking.

Testing

• python -m pytest mjx/mujoco/mjx/_src/render_util_test.py -q
• python -m pytest -n auto -v -k 'not IntegrationTest' --pyargs mujoco.mjx
• python -m pytest -v --pyargs mujoco
• ctest -C Release --output-on-failure

[btaba]

can we keep this setup? we have other CI besides github CI, and want to make sure the tests have minimal changes. LGTM otherwise

[tkelestemur]

Restored the existing Warp/CUDA-gated setup and kept the new batched-shape assertions as the only functional test change. Re-ran the MJX suite on this final shape (376 passed, 39 skipped).

[tkelestemur]

Thanks @btaba Good to go!