Implementation:Google deepmind Mujoco MJX Render Util
| Knowledge Sources | |
|---|---|
| Domains | Physics_Simulation, JAX, Rendering |
| Last Updated | 2026-02-15 04:00 GMT |
Overview
JAX render utilities for unpacking RGB and depth render output from the MuJoCo Warp rendering backend into standard float32 arrays.
Description
This module provides helper functions for converting packed render output from the Warp-based renderer into usable image formats. get_rgb() unpacks uint32 ABGR pixel data into float32 RGB arrays with shape (H, W, 3) and values in [0, 1]. get_depth() extracts and normalizes raw depth data into float32 arrays with shape (H, W), clamped to [0, 1]. Both functions use the RenderContext handle to look up camera-specific pixel address offsets and resolution, then apply JAX operations (dynamic slicing, bitwise extraction) to produce the final images.
Usage
Used after calling the Warp renderer to convert raw packed pixel buffers into standard image formats suitable for visualization, training (e.g., visual RL), or further processing.
Code Reference
Source Location
- Repository: Google_deepmind_Mujoco
- File: mjx/mujoco/mjx/_src/render_util.py
- Lines: 1-97
Key Functions
def get_rgb(rc: Any, rgb_data: jax.Array, cam_id: int) -> jax.Array
def get_depth(rc: Any, depth_data: jax.Array, cam_id: int, depth_scale: float) -> jax.Array
Import
from mujoco.mjx._src.render_util import get_rgb
from mujoco.mjx._src.render_util import get_depth
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| rc | RenderContext | Yes | Warp render context handle with camera resolution and address info |
| rgb_data | jax.Array | Yes | Packed uint32 ABGR pixel data from the renderer (total_pixels,) |
| depth_data | jax.Array | Yes | Raw float32 depth values from the renderer (total_pixels,) |
| cam_id | int | Yes | Camera index to extract from the render buffer |
| depth_scale | float | Yes | Scale factor for normalizing depth values (for get_depth) |
Outputs
| Name | Type | Description |
|---|---|---|
| rgb | jax.Array | Float32 RGB image with shape (H, W, 3), values in [0, 1] |
| depth | jax.Array | Float32 depth image with shape (H, W), clamped to [0, 1] |