### Launch Madrona Viewer with Franka Emika Panda Example

Source: https://github.com/shacklettbp/madrona_mjx/blob/main/README.md

Generates a Franka Emika Panda robot example and then launches the Madrona viewer with a specific MJCF file. This example showcases loading robotics models and configuring viewer window and batch render dimensions.

```shell
./scripts/create_franka_example

python scripts/viewer.py --mjcf mujoco_menagerie/franka_emika_panda/mjx_single_cube_camera.xml --num-worlds 16 --window-width 2730 --window-height 1536 --batch-render-view-width 64 --batch-render-view-height 64
```

--------------------------------

### Launch Madrona Viewer with Cartpole MJCF

Source: https://github.com/shacklettbp/madrona_mjx/blob/main/README.md

Launches the Madrona viewer script to visualize a cartpole MJCF file. This is a basic example demonstrating how to start the viewer with a specified simulation file.

```shell
python scripts/viewer.py --mjcf data/cartpole.xml
```

--------------------------------

### Enable Madrona Kernel Cache for Faster Initialization

Source: https://github.com/shacklettbp/madrona_mjx/blob/main/README.md

Demonstrates how to use environment variables to enable kernel caching for Madrona. This speeds up initialization by reusing compiled CUDA kernels, significantly reducing load times after the first run. It's recommended to set these variables before running the viewer script.

```shell
MADRONA_MWGPU_KERNEL_CACHE=build/kernel_cache MADRONA_BVH_KERNEL_CACHE=build/bvh_cache python scripts/viewer.py --mjcf mujoco_menagerie/franka_emika_panda/mjx_single_cube_camera.xml
```

--------------------------------

### Initialize and Render Simulation State with MadronaBatchRenderer (Python)

Source: https://github.com/shacklettbp/madrona_mjx/blob/main/docs/WALKTHROUGH.md

This pseudocode demonstrates the typical usage pattern for the MadronaBatchRenderer. It involves initializing the simulator and renderer, stepping through simulation steps, rendering the current state, and retrieving the output tensors. This pattern is designed for users interfacing with Madrona MJX through Python.

```python
# Initialize the physics simulator
my_simulator = Simulator(...)

# Initialize the Madrona renderer
renderer = MadronaBatchRenderer(geometry, num_instances, ...)
renderer.init(my_simulator.initial_state())

for _ in range(num_steps):
    # Step the simulation
    my_simulator.step()

    # Render the output of the current simulation step
    renderer.render(my_simulator.get_state())

    # Retrieve the output tensors
    rendered_output_tensors = renderer.rgbd_tensors()

```

--------------------------------

### Launch Interactive Viewer for Real-time Visualization

Source: https://context7.com/shacklettbp/madrona_mjx/llms.txt

Initializes a GPU state and a BatchRenderer to launch an interactive viewer for real-time visualization of batch-rendered environments. The visualization step function simulates environments, generates random actions, and renders RGB and depth data. Dependencies include madrona_mjx, jax, and MuJoCo. Inputs are simulation data and model. Outputs are visualizer window and rendered frames.

```python
from madrona_mjx.viz import Visualizer, VisualizerGPUState
from madrona_mjx import BatchRenderer
import jax

# Initialize GPU state for viewer window
viz_gpu_state = VisualizerGPUState(
    window_width=2560,
    window_height=1440,
    gpu_id=0
)

# Create renderer with visualizer handles
renderer = BatchRenderer(
    mjx_model, gpu_id=0, num_worlds=16,
    batch_render_view_width=128,
    batch_render_view_height=128,
    viz_gpu_hdls=viz_gpu_state.get_gpu_handles()
)

# Initialize simulation
mjx_data = jax.make_data(mjx_model)
render_token, _, _ = renderer.init(mjx_data, mjx_model)

# Define visualization step function
def vis_step_fn(carry):
    rng, data = carry
    rng, act_rng = jax.random.split(rng)

    # Generate random actions
    actions = jax.random.uniform(act_rng, shape=(16, mjx_model.nu))

    # Step all environments
    @jax.vmap
    def step(d, a):
        d = d.replace(ctrl=a)
        d = mjx.step(mjx_model, d)
        return d

    data = step(data, actions)
    _, rgb, depth = renderer.render(render_token, data)

    return rng, data

# Launch interactive viewer
rng = jax.random.PRNGKey(0)
visualizer = Visualizer(viz_gpu_state, renderer.madrona)
visualizer.loop(renderer.madrona, vis_step_fn, (rng, mjx_data))
```

--------------------------------

### Initialize BatchRenderer with MJX Model Geometry

Source: https://context7.com/shacklettbp/madrona_mjx/llms.txt

Initializes the BatchRenderer with MJX model geometry and configuration. It loads a MuJoCo model from an XML file and prepares it for GPU rendering. The renderer is configured with parameters such as the number of parallel worlds, view dimensions, and enabled geometry groups.

```python
import numpy as np
from mujoco import mjx
from madrona_mjx import BatchRenderer
from madrona_mjx.wrapper import load_model

# Load MuJoCo model from XML file
model = load_model('path/to/scene.xml')
mjx_model = mjx.put_model(model)

# Initialize batch renderer with 1024 parallel worlds
renderer = BatchRenderer(
    m=mjx_model,
    gpu_id=0,
    num_worlds=1024,
    batch_render_view_width=64,
    batch_render_view_height=64,
    enabled_geom_groups=np.array([0, 1, 2]),  # Render geometry groups
    enabled_cameras=None,  # Use all cameras in model
    add_cam_debug_geo=False,
    use_rasterizer=False,  # Use raytracer (recommended)
    viz_gpu_hdls=None
)

# Initialize renderer with first frame
# render_token, init_rgb, init_depth = renderer.init(mjx_data, mjx_model)
# Output shapes: rgb (1024, num_cams, 64, 64, 4), depth (1024, num_cams, 64, 64, 1)
```

--------------------------------

### MadronaWrapper for Brax Environments (Python)

Source: https://context7.com/shacklettbp/madrona_mjx/llms.txt

Wraps Brax environments to enable batched rendering with domain randomization support. This allows for parallel simulation and rendering of multiple environment instances. It takes a base Brax environment, the number of worlds, and a randomization function as input.

```python
from brax.envs import create
from madrona_mjx.wrapper import MadronaWrapper

# Create base Brax environment
env = create('cartpole', backend='mjx')

# Wrap with MadronaWrapper for batched rendering
num_worlds = 512
wrapped_env = MadronaWrapper(
    env=env,
    num_worlds=num_worlds,
    randomization_fn=lambda sys: domain_randomize(sys, num_worlds)
)

# Reset environments
rng = jax.random.PRNGKey(0)
rng_keys = jax.random.split(rng, num_worlds)
state = wrapped_env.reset(rng_keys)

# Step environments
actions = jax.random.uniform(rng, (num_worlds, env.action_size))
state = wrapped_env.step(state, actions)
# State contains observations, rewards, and rendered images

```

--------------------------------

### Enable Kernel Caching for Fast Initialization

Source: https://context7.com/shacklettbp/madrona_mjx/llms.txt

Accelerates subsequent runs by enabling kernel caching, skipping expensive CUDA compilation. This involves setting environment variables for kernel cache directories before running the visualization script. The first run compiles and caches kernels (slow), while subsequent runs reuse cached kernels for fast startup.

```bash
# Set environment variables for kernel caching
export MADRONA_MWGPU_KERNEL_CACHE=./build/kernel_cache
export MADRONA_BVH_KERNEL_CACHE=./build/bvh_cache

# First run compiles and caches kernels (slow)
python scripts/viewer.py --mjcf data/cartpole.xml --num-worlds 1024

# Subsequent runs reuse cached kernels (fast startup)
python scripts/viewer.py --mjcf data/cartpole.xml --num-worlds 1024
```

--------------------------------

### Render MJX Simulation States with BatchRenderer

Source: https://context7.com/shacklettbp/madrona_mjx/llms.txt

Renders RGB and depth images from batched MJX physics states at each simulation step. This involves stepping the physics simulation and then calling the renderer to capture the visual output. JIT compilation is used to optimize the performance of the step-and-render loop.

```python
import jax
import jax.numpy as jp
from mujoco import mjx

# Assuming mjx_model and renderer are already initialized as in the previous example
# mjx_model = ...
# renderer = ...

# Create initial MJX data
mjx_data = mjx.make_data(mjx_model)
mjx_data = mjx.forward(mjx_model, mjx_data)

# Initialize rendering
# render_token, rgb, depth = renderer.init(mjx_data, mjx_model)

# Simulation and rendering loop
def step_and_render(data, action):
    # Step physics simulation
    data = data.replace(ctrl=action)
    data = mjx.step(mjx_model, data)

    # Render current state
    render_token, rgb, depth = renderer.render(render_token, data)
    return data, rgb, depth

# JIT compile for performance
step_and_render_jit = jax.jit(step_and_render)

# Execute rollout
# for _ in range(1000):
#     actions = jp.zeros((mjx_model.nu,))
#     mjx_data, rgb, depth = step_and_render_jit(mjx_data, actions)
    # rgb shape: (num_worlds, num_cams, height, width, 4)
    # depth shape: (num_worlds, num_cams, height, width, 1)
```

--------------------------------

### Taskgraph for GPU Rendering Data Transformation (C++)

Source: https://github.com/shacklettbp/madrona_mjx/blob/main/docs/WALKTHROUGH.md

The code in `sim.cpp` defines the `Taskgraph`, a series of GPU-executed functions. Its primary role is to transform simulator data into a format understood by the Madrona renderer, specifically using components and archetypes represented as contiguous GPU arrays. This is crucial for enabling high-throughput rendering.

```cpp
// In sim.cpp:
// Defines the Taskgraph, a series of GPU functions.
// Transforms simulator data into Madrona renderer's required format (components and archetypes).
// Enables high-throughput rendering by managing contiguous GPU memory buffers.

```

--------------------------------

### Render from Multiple Camera Viewpoints Simultaneously

Source: https://context7.com/shacklettbp/madrona_mjx/llms.txt

Enables rendering from multiple camera viewpoints within each environment simultaneously. This is achieved by specifying an array of camera indices to enable during BatchRenderer initialization. The output RGB and depth arrays will have an additional dimension for the number of cameras, allowing access to individual camera views.

```python
import numpy as np

# Specify which cameras to enable (by index in MuJoCo model)
enabled_cameras = np.array([0, 2, 3])  # Enable cameras 0, 2, and 3

# Initialize renderer with selected cameras
renderer = BatchRenderer(
    mjx_model,
    gpu_id=0,
    num_worlds=256,
    batch_render_view_width=128,
    batch_render_view_height=128,
    enabled_cameras=enabled_cameras
)

# Initialize and render
render_token, rgb, depth = renderer.init(mjx_data, mjx_model)

# Output shapes with multiple cameras:
# rgb: (256, 3, 128, 128, 4)   # 256 worlds, 3 cameras, 128x128 RGBA
# depth: (256, 3, 128, 128, 1) # 256 worlds, 3 cameras, 128x128 depth

# Render next frame
render_token, rgb, depth = renderer.render(render_token, mjx_data)

# Access specific camera views
cam0_rgb = rgb[:, 0]  # First camera across all worlds
cam2_rgb = rgb[:, 2]  # Third camera across all worlds
```

--------------------------------

### MadronaBatchRenderer Class Definition (C++ Bindings for Python)

Source: https://github.com/shacklettbp/madrona_mjx/blob/main/docs/WALKTHROUGH.md

The `MadronaBatchRenderer` class, defined in `bindings.cpp`, acts as a Python interface to Madrona's rendering capabilities. It wraps the C++ `Manager` class and exposes functionalities like initialization (`init`) and rendering (`render`) to Python users. This class handles geometry, instance transforms, camera, and light information.

```cpp
// In bindings.cpp:
// Defines MadronaBatchRenderer class for Python interaction.
// Constructor takes geometry, num_instances, etc.
// Methods: init(...), render(...), rgbd_tensors(), get_instance_transforms().
// Wraps the C++ Manager class.

```

--------------------------------

### Domain Randomization for Visual Properties (Python)

Source: https://context7.com/shacklettbp/madrona_mjx/llms.txt

Randomizes geometry colors, sizes, materials, and lighting across batched simulation worlds using JAX. This function takes a system object and the number of worlds as input, returning the modified system object and in_axes for vmap. It's designed for robust policy training by introducing variability in visual properties.

```python
import jax
import jax.numpy as jp

def domain_randomize(sys, num_worlds):
    """Randomize visual properties across batched worlds."""

    @jax.vmap
    def randomize_per_world(rng):
        rng1, rng2, rng3 = jax.random.split(rng, 3)

        # Randomize geometry colors (material ID -2 uses color override)
        geom_matid = sys.geom_matid.at[:].set(-1)  # Use default materials
        geom_matid = geom_matid.at[0].set(-2)  # Override first geom with color
        new_color = jax.random.uniform(rng1, (1,), minval=0.0, maxval=0.4)
        geom_rgba = sys.geom_rgba.at[0, 0:1].set(new_color)

        # Randomize geometry sizes
        new_size = jax.random.uniform(rng2, (3,), minval=1.0, maxval=5.0)
        geom_size = sys.geom_size.at[0, 0:3].set(new_size)

        # Randomize light position and properties
        light_pos = jax.random.uniform(
            rng3,
            (3,),
            minval=jp.array([-0.5, -0.5, 2.0]),
            maxval=jp.array([0.5, 0.5, 3.0])
        )
        light_cutoff = jax.random.uniform(rng3, (1,), minval=1.0, maxval=1.5)

        return geom_rgba, geom_matid, geom_size, light_pos, light_cutoff

    # Generate random keys for each world
    rng = jax.random.PRNGKey(0)
    rngs = jax.random.split(rng, num_worlds)

    # Apply randomization
    geom_rgba, geom_matid, geom_size, light_pos, light_cutoff = randomize_per_world(rngs)

    # Specify which axes are batched for vmap
    in_axes = jax.tree_util.tree_map(lambda x: None, sys)
    in_axes = in_axes.tree_replace({
        'geom_rgba': 0, 'geom_matid': 0, 'geom_size': 0,
        'light_pos': 0, 'light_cutoff': 0
    })

    # Replace system properties with randomized values
    sys = sys.tree_replace({
        'geom_rgba': geom_rgba, 'geom_matid': geom_matid, 'geom_size': geom_size,
        'light_pos': light_pos, 'light_cutoff': light_cutoff
    })

    return sys, in_axes

```

--------------------------------

### Benchmark High-Throughput Rendering (Python)

Source: https://context7.com/shacklettbp/madrona_mjx/llms.txt

Measures rendering and simulation throughput for performance optimization using MJX and Madrona's BatchRenderer. This function takes an MJX model, number of worlds, steps, and render dimensions as input. It outputs the frames per second (FPS) and the realtime factor.

```python
import jax
import jax.numpy as jp
from mujoco import mjx
from madrona_mjx import BatchRenderer
import time

def benchmark_rendering(mjx_model, num_worlds, num_steps, render_width, render_height):
    """Benchmark MJX simulation with Madrona rendering."""

    # Initialize renderer
    renderer = BatchRenderer(
        mjx_model, gpu_id=0, num_worlds=num_worlds,
        batch_render_view_width=render_width,
        batch_render_view_height=render_height,
        use_rasterizer=False  # Use raytracer
    )

    # Initialize simulation data
    rng = jax.random.PRNGKey(0)
    rngs = jax.random.split(rng, num_worlds)

    @jax.vmap
    def init_world(rng):
        data = mjx.make_data(mjx_model)
        data = data.replace(qpos=0.01 * jax.random.uniform(rng, shape=(mjx_model.nq,)))
        return mjx.forward(mjx_model, data)

    mjx_data = init_world(rngs)
    render_token, _, _ = renderer.init(mjx_data, mjx_model)

    # Define step function
    @jax.jit
    def rollout(data):
        @jax.vmap
        def step(d, _):
            d = mjx.step(mjx_model, d)
            _, rgb, depth = renderer.render(render_token, d)
            return d, None

        d, _ = jax.lax.scan(step, data, None, length=num_steps)
        return d

    # Warmup JIT compilation
    mjx_data = rollout(mjx_data)
    jax.block_until_ready(mjx_data)

    # Measure performance
    start = time.time()
    mjx_data = rollout(mjx_data)
    jax.block_until_ready(mjx_data)
    elapsed = time.time() - start

    total_steps = num_steps * num_worlds
    fps = total_steps / elapsed

    print(f"Total steps: {total_steps}")
    print(f"Time: {elapsed:.2f} s")
    print(f"Steps per second: {fps:.0f}")
    print(f"Realtime factor: {total_steps * mjx_model.opt.timestep / elapsed:.2f}x")

    return fps

# Run benchmark
fps = benchmark_rendering(
    mjx_model=mjx_model,
    num_worlds=1024,
    num_steps=1000,
    render_width=64,
    render_height=64
)

```

--------------------------------

### Manager Class Responsibilities (C++)

Source: https://github.com/shacklettbp/madrona_mjx/blob/main/docs/WALKTHROUGH.md

The `Manager` class in `mgr.cpp` is the central orchestration component in Madrona. It is responsible for initializing Madrona's core systems, processing geometry data, copying simulator state into GPU buffers, invoking the batched renderer, and exposing data to Python via `bindings.cpp`.

```cpp
// In mgr.cpp:
// The Manager class orchestrates Madrona's core components.
// Responsibilities include: initialization, geometry loading/processing,
// state copying, invoking the batch renderer, and data exposure to Python.

```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.