Madrona MJX

When using Madrona MJX, adhere to the following technical requirements and constraints:

Data Requirements

• Camera Requirement: An MJCF file must include a camera; otherwise, the application will crash.
• Batching: The MJX Model passed to the Madrona Renderer must be batched such that arrays have the shape [num_worlds, ...]. Non-domain-randomized MJX model attributes must be manually tiled to this shape.
• World Count: The number of worlds must be known and initialized at the start. All worlds are rendered every time; you cannot disable rendering for specific worlds. This implies that training and evaluation must share the same batch size.

Resource and Lifecycle Constraints

• Renderer Singleton: Only one renderer instance can be initialized at a time. You cannot create two separate environments (e.g., one for training and one for evaluation) that each use their own batch renderer instance.
• Memory Management: Use the MADRONA_MWGPU_DEVICE_HEAP_SIZE environment variable to adjust how much memory Madrona pre-allocates on the GPU. If you do not allocate enough memory for your use case, an error will be thrown.

Supported Features

• Rigid body rendering (position/rotation changes).
• Camera position and rotation changes.
• Light position and rotation changes.
• Geometry color, size, and texture domain randomization.
• Light parameter domain randomization.

Unsupported Features

• Deformable bodies.
• Particle systems.
• Muscles, Tendons, and Composites (except for rigid body components).

To integrate a custom simulator, you must modify the geometry loading process in mgr.cpp, specifically the loadRenderObjects function. Madrona's geometry format is hierarchical:

1. Mesh: A collection of arrays (triangles, normals, UVs, etc.) representing a shape.
2. Object: A group of one or more meshes that are rendered together as a single entity.

To implement custom asset loading, your code must produce an array of SourceObject structures. You can achieve this via:

• Using the Madrona-provided AssetImporter class to load files (e.g., .obj, .gltf) from disk, which populates an ImportedAssets structure.
• Manually filling the SourceObject array.
• A hybrid approach using AssetImporter for some assets and manual configuration for others.

Once the SourceObject array is created, pass it to the AssetProcessor's BVH utility functions to generate the final structures required for the renderer's initialization.

Madrona MJX requires building the renderer from source and then installing the Python wrapper.

Prerequisites

1. Install system libraries:

sudo apt install libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev mesa-common-dev 2. Install CUDA 12.5.1 or earlier. 3. Install Jax (Note: versions >= 0.6.0 are incompatible). Use: pip install -U "jax[cuda12]<0.6.0" or pip install "jax[cuda12_local]" to avoid version mismatches. 4. Install MuJoCo and MuJoCo MJX (requires mujoco-mjx >= 3.2.7): pip install -U mujoco mujoco-mjx.

Build Steps

git clone https://github.com/shacklettbp/madrona_mjx.git
cd madrona_mjx
git submodule update --init --recursive
mkdir build
cd build
cmake ..
make -j

cd ..
pip install -e .

Note on Vulkan: If your machine lacks Vulkan and you encounter an assertion failure at madrona::render::vk::Backend::Init::init, use cmake -DLOAD_VULKAN=OFF .. instead of cmake ...

The madrona_mjx Python module acts as a renderer interface that translates MJX simulation state into Madrona's GPU-based ECS state for high-throughput rendering.

To use the renderer, follow this lifecycle:

1. Initialize the physics simulator (e.g., MJX).
2. Initialize the Madrona renderer using MadronaBatchRenderer, providing geometry and instance counts. Pass the simulator's initial state to the renderer's .init() method.
3. Step the simulation in a loop.
4. Render the current step by passing the simulator's current state to renderer.render().
5. Retrieve output using renderer.rgbd_tensors() to get the rendered RGB and depth data.

# 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()