mujofil-warp

A GPU simulation pipeline for vision-based RL: MuJoCo Warp physics + a parallel, high-fidelity rasterization renderer (forked from Google Filament), zero-copy to PyTorch.

mujofil-warp builds an efficient, GPU-parallel rasterization render engine (a fork of Google Filament — PBR materials, image-based lighting, soft shadows, SSAO, reflections) and wires it into a complete simulation pipeline: it plugs MuJoCo Warp's high-throughput GPU physics into that renderer so you get the best of both — MuJoCo Warp's fast, massively parallel dynamics and fast, parallel, photoreal visual frames from the Filament fork — delivered straight to PyTorch as torch.cuda tensors with no CPU round-trip for the pixels.

What that unlocks:

• Drop in any environment. Pull scenes/assets from Sketchfab, Poly Haven and similar sources (glTF / GLB / OBJ / USD) and train your robot's RL policy inside them — photoreal worlds MuJoCo and MuJoCo Warp's built-in raycaster cannot even load.
• Photoreal vision observations (PBR, IBL, reflections) at parallel-batch throughput, so the renderer keeps up with GPU physics instead of bottlenecking it.
• One import. Your code only ever imports mujofil_warp; it drives the MuJoCo Warp physics and the renderer for you.

Positioning, honestly: the physics is MuJoCo Warp (DeepMind + NVIDIA's GPU MuJoCo) — we don't reimplement dynamics. Our work is the parallel rasterization renderer and the zero-copy GPU→PyTorch pipeline that turns those GPU-resident world states into photoreal training observations. It targets the middle of the fidelity/speed spectrum: more realistic than a flat raycaster, far lighter than ray-traced stacks like Omniverse — photoreal-enough RGB that runs on a mid-range GPU.

📖 Full documentation: docs/ — getting started, API guide, feature reference, cookbook & troubleshooting.

🖥️ Running CPU MuJoCo instead? Use the CPU edition, mujofil (photoreal frames as NumPy arrays).

Highlights

• Zero-copy to torch.cuda. Filament renders into GPU memory that CUDA imports directly; observations arrive as torch.cuda tensors with no GPU→CPU→GPU bounce.
• GPU-resident pipeline. MJWarp steps physics on the GPU; only a tiny transform array crosses to the host. Pixels never leave the GPU.
• Photoreal. Full PBR metalness/roughness, IBL, soft shadows, SSAO, MSAA, filmic tone mapping — renders complete GLB environments MJWarp/MuJoCo can't.
• Two backends. An OpenGL single-sync path and a Vulkan shared-device path, selectable at runtime.

Performance (RTX 4060 Laptop, 8 GiB)

All numbers are env-steps/s (= cameras/s), MJWarp GPU physics → torch.cuda.

vs vanilla MuJoCo, same scene, same workload (ours adds PBR + zero-copy):

We beat vanilla MuJoCo by 1.25–2.12× on equal work — the gap widens at higher resolution because zero-copy avoids the CPU readback that scales with pixels.

Full photoreal warehouse (3 GLB meshes + IBL + 16 spotlights + SSAO — geometry vanilla MuJoCo and MJWarp cannot even load): ~3,200 cam/s at 128px, holding flat from N=64 to N=2048.

GL vs Vulkan backend (full warehouse): the GL single-sync path is 1.3× faster and, critically, its sync cost is constant across N (one flushAndWait), where the Vulkan path's grows linearly with batch size.

vs MJWarp's own raycaster: MJWarp scales to ~42,000 cam/s at N=2048 — but that is flat Lambertian on bare objects (no PBR/IBL, no GLB environments). At small N (≤32) mujofil-warp is faster and photoreal; at large N MJWarp wins raw throughput by trading away all visual fidelity. Different categories: MJWarp is a parallel raycaster, this is a photoreal rasterizer.

Quickstart

You only import mujofil_warp. ParallelScene runs the GPU physics (MuJoCo Warp) and renders every world to a zero-copy torch.cuda tensor — no put_model / make_data / host-copy boilerplate:

import mujofil_warp

scene = mujofil_warp.ParallelScene("scene.xml", num_worlds=32,
                                   width=256, height=256, preset="high")

for _ in range(100):
    scene.step()                     # GPU physics (MuJoCo Warp)
    obs = scene.render(camera=0)     # (32, 256, 256, 4) uint8 torch.cuda — zero-copy

Set controls or initial state through scene.data (the MuJoCo Warp Data) and the model through scene.model. See examples/minimal_render.py for a runnable demo.

import mujoco, mujoco_warp as mjw, warp as wp
from mujofil_warp import WarpRenderer

mjm = mujoco.MjModel.from_xml_path("scene.xml")
M = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=32)
host = [mujoco.MjData(mjm) for _ in range(32)]

r = WarpRenderer(width=256, height=256, batch_size=32, preset="high")
r.load_model(mjm)

mjw.step(M, d); wp.synchronize()
gx = d.geom_xpos.numpy(); gm = d.geom_xmat.numpy().reshape(32, mjm.ngeom, 9)
for i, h in enumerate(host):
    h.geom_xpos[:] = gx[i]; h.geom_xmat[:] = gm[i]

obs = r.render_batch(mjm, host, cam_id=0)   # (32, 256, 256, 4) uint8 torch.cuda

Quality toggles

Every fidelity feature is an independent toggle so you can reproduce the throughput/fidelity trade-offs in benchmarks/ on your own hardware:

from mujofil_warp import WarpRenderer, make_config

# keyword toggles
r = WarpRenderer(width=256, batch_size=32, ssao=False, shadows=True, msaa=True)

# or a named preset, optionally overriding individual toggles
r = WarpRenderer(width=256, batch_size=32, preset="fast")          # SSAO off, ~2x
r = WarpRenderer(width=256, batch_size=32, preset="high", bloom=True)

# or an explicit config
cfg = make_config(width=256, height=256, batch_size=32, exposure=1.6)
r = WarpRenderer(config=cfg)

Presets: high (photoreal, default), medium (high-quality SSAO, no cone tracing), fast (SSAO off, ~2×), ultra (8× MSAA + bloom), raw (no AO/shadows/AA, ~3×).

Backends

Select at runtime with MUJOFIL_WARP_BACKEND:

• gl (default) — OpenGL single-sync. Renders N worlds into N imported GL textures bracketed by one flushAndWait, then exports via GL↔CUDA interop. Sync cost is constant in N; fastest in the warehouse. Requires an X display (DISPLAY); when none is available it automatically falls back to Vulkan.
• vulkan — shared Vulkan device + exportable swapchain + CUDA external-memory import. Works fully headless (no X), but the 2-frame in-flight cap makes its sync cost grow with batch size.

# default is gl; force a backend explicitly with the env var:
MUJOFIL_WARP_BACKEND=gl     python examples/minimal_render.py --preset high
MUJOFIL_WARP_BACKEND=vulkan python examples/minimal_render.py --preset high

Installation

pip install mujofil-warp

The wheel is self-contained: Filament and the CUDA runtime are statically baked in, the compiled materials ship inside it, and libc++ is bundled. There is no CUDA toolkit, no Filament, and no mujofil to install — the only hard requirement at runtime is an NVIDIA GPU + driver.

Supported environments

Because the package contains no CUDA device code (only host-side runtime calls), a single wheel is portable across GPUs and driver versions:

Not yet supported: aarch64 (Jetson/Grace), glibc < 2.34 (Ubuntu 20.04 / RHEL 8), non-NVIDIA GPUs. These need a from-source Filament build (planned).

PyTorch (zero-copy target)

torch is an optional dependency (pip install "mujofil-warp[torch]"), and you must install a build that matches your GPU's compute capability — the zero-copy DLPack handoff runs CUDA kernels through your torch, not ours.

• Blackwell (RTX 50-series / sm_120, e.g. 5090): install the CUDA 12.8 torch — pip install torch --index-url https://download.pytorch.org/whl/cu128. A torch+cu124 (or older) build has no sm_120 kernels and fails at runtime with CUDA error: no kernel image is available for execution on the device.
• Ada / Hopper / Ampere (sm_80–sm_90): the default cu124 torch is fine.

warp-lang and mujoco-warp JIT-compile for the local GPU, so they need no such pinning — only torch ships prebuilt device code.

Headless / display

Both backends are fully headless — no X server, no display, nothing extra to install beyond the NVIDIA driver:

• GL (default) uses surfaceless EGL, so it renders headless at full speed on a bare GPU server (cloud, cluster, container). This is the recommended path for vision-RL training.
• Vulkan is also headless (shared device + exportable swapchain).

GL auto-falls back to Vulkan only if the GL module fails to initialize.

Building from source

Most users never need this — pip install mujofil-warp ships prebuilt wheels. Build from source only to hack on the C++ or target an unsupported environment.

Prerequisites (the native modules and Filament are built with Clang + libc++):

Then:

git clone https://github.com/tau-intelligence/mujofil-warp
cd mujofil-warp
CC=clang CXX=clang++ pip install .

How Filament is resolved (the GL backend's headless EGL rendering needs a custom EGL-enabled Filament — Google's prebuilt Linux Filament is GLX-only). CMakeLists.txt tries, in order:

1. FILAMENT_DIR=/path/to/egl-filament if you set it — used as-is (fastest).
2. Download a prebuilt EGL Filament artifact (seconds). The default path.
3. Build from source via packaging/build_filament_egl.sh (~20–30 min) if the download is unavailable — this is the step that needs git/cmake/ninja.

So a plain pip install . is one command; supply FILAMENT_DIR to skip the download/build entirely:

CC=clang CXX=clang++ FILAMENT_DIR=/path/to/egl-filament pip install .

The EGL Filament artifact is reproducible from source:

packaging/build_filament_egl.sh ./_filament_egl   # clone + patch + build

Dev rebuilds (no full reinstall)

For iterating on the C++ without a full pip install, the two helper scripts build the modules in place (point FILAMENT_DIR at the EGL Filament build):

bash native/build_gl.sh   # OpenGL single-sync, headless EGL -> _mujofil_warp_gl
bash native/build.sh      # Vulkan zero-copy                  -> _mujofil_warp

Architecture & porting

mujofil-warp is one core with pluggable rendering backends, so new platforms are added as a backend — not a fork.

mujofil_warp/__init__.py     Python API, presets, backend selection   (shared)
native/render_module.cpp     pybind bindings, batching                (shared)
native/vendor/core/          scene / material / light bridge          (shared)
native/renderer_gl.cpp       Linux: surfaceless EGL  + CUDA interop   (backend)
native/renderer_warp.cpp     Linux: Vulkan device    + CUDA interop   (backend)

Everything platform-specific lives behind the vf_mujoco::Renderer interface (context creation, GPU→tensor interop). Adding macOS or Windows means adding one renderer_*.{cpp,mm} implementing that interface — the scene, material, lighting, Python API, and batching layers are reused unchanged.

• Windows would use a WGL/EGL context + OPAQUE_WIN32 external-memory handles for the CUDA interop.
• macOS is a different target: there is no CUDA on Apple platforms, so a Mac backend would use Filament's Metal backend and export to PyTorch via MPS (MTLBuffer → torch-MPS) rather than torch.cuda.

These are not yet implemented (they need the respective hardware to develop and validate on), but the codebase is structured so they slot in without a fork.

Layout

mujofil_warp/        Python package (WarpRenderer, make_config, presets)
native/              C++ renderer + pybind module + build scripts
  renderer_gl.cpp      OpenGL single-sync zero-copy backend
  renderer_warp.cpp    Vulkan shared-device zero-copy backend
  render_module.cpp    pybind bindings (shared by both backends)
examples/            runnable demos
benchmarks/          the benchmark suite behind the numbers above
spikes/              isolated feasibility proofs (GL↔CUDA, Vulkan↔CUDA, DLPack)
docs/ARCHITECTURE.md design + phased integration plan

Relationship to mujofil

mujofil-warp reuses the CPU-MuJoCo mujofil renderer's scene/material/light source but is a separate build — the published mujofil package is untouched. Use mujofil for high-fidelity CPU-MuJoCo vector-env rendering; use mujofil-warp when you want MJWarp's GPU-resident physics with photoreal, zero-copy observations.

License

Apache-2.0.