Skip to content

Instantly share code, notes, and snippets.

@edwhu
Last active March 7, 2024 12:55
Show Gist options
  • Save edwhu/8e9de0c8e880dad560fc728a3975ade8 to your computer and use it in GitHub Desktop.
Save edwhu/8e9de0c8e880dad560fc728a3975ade8 to your computer and use it in GitHub Desktop.
Grokking Mujoco GPU rendering
import mujoco_py
from mujoco_py import MjSim, load_model_from_xml, MjRenderContext
from time import time
import argparse
BASIC_MODEL_XML = """
<mujoco>
<asset>
<texture name="t1" width="33" height="36" type="2d" builtin="flat" />
<texture name="t2" width="34" height="39" type="2d" builtin="flat" />
<texture name="t3" width="31" height="37" type="2d" builtin="flat" />
<texture name="t4" width="38" height="34" type="2d" builtin="flat" />
<material name="m1" texture="t1" />
<material name="m2" texture="t2" />
<material name="m3" texture="t3" />
<material name="m4" texture="t4" />
</asset>
<worldbody>
<light diffuse=".5 .5 .5" pos="0 0 5" dir="0 0 -1" />
<camera name="topcam" pos="0 0 2.5" zaxis="0 0 1" />
<geom name="g1" pos="-0.5 0.5 0" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" material="m1" />
<geom name="g2" pos="0.5 0.5 0" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" material="m2" />
<geom name="g3" pos="0.5 -0.5 0" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" material="m3" />
<geom name="g4" pos="-0.5 -0.5 0" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" material="m4" />
<!-- The following boxes will show up in reflections -->
<geom pos="0.5 0.5 4" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" />
<geom pos="0.5 -0.5 4" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" />
<geom pos="-0.5 -0.5 4" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" />
<geom pos="-0.5 0.5 4" type="box" size="0.4 0.4 0.1" rgba="1 1 1 1" />
</worldbody>
</mujoco>
"""
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
parser = argparse.ArgumentParser()
parser.add_argument('device_id', type=int, help='device id for rendering')
parser.add_argument('--backend', type=str, default=None, choices=['glfw', None])
args = parser.parse_args()
render_context = MjRenderContext(sim, True, args.device_id, args.backend)
print(mujoco_py.cymj)
start = time()
for i in range(10000):
img = sim.render(201, 205, camera_name="topcam", device_id=args.device_id)
print("took", time() - start)

Grokking Mujoco GPU Rendering

Mujoco has GLFW, EGL rendering backends for GPU. They are around 10-100x faster than rendering on the CPU, so it is important to enable GPU rendering if you render a lot.

Mujoco GPU rendering is really annoying to setup, so I thought to consolidate my knowledge here after many hours of getting it to work on a Ubuntu 20.04 cluster.

Install the pip package gpustat to monitor gpu usage. In a tmux pane, run gpustat -i to get a live dashboard of the GPU memory usage.

Enable Mujoco GPU rendering

Next, ensure that mujoco_py is built for GPU. When you import mujoco_py, print mujoco_py.cymj and ensure that it says libgpu instead of libcpu. If it is libcpu, you must force mujoco_py to build for GPU. There are two ways to force mujoco to build GPU.

One is to modify the source code of mujoco_py directly, and then compile it from source. This is the cleanest way to do it. openai/mujoco-py#512

The easiest way though, is to make an empty folder /usr/lib/nvidia-000. Mujoco_py looks for a folder in /usr/lib/ with the pattern nvidia-XXX, and if so, it will build for GPU.

Testing Mujoco EGL Rendering

First, make sure your system has the libnvidia GL, encode, decode packages installed. Otherwise, hardware rendering will not work. See https://imgur.com/a/LZvhkhS for what we have installed.

For our server, it is libnvidia-gl-460-server.

Make sure the env var LD_PRELOAD is unset. Then run python test_mj.py 0 and check your gpustat. You should see some activity pop up on the 0th gpu.

You can see how fast it is in comparison to the CPU rendering. To enable cpu rendering, MUJOCO_PY_FORCE_CPU=1 python test_mj.py 0. Make sure to remove the MUJOCO_PY_FORCE_CPU env var when you are done, since mujoco_py only checks the existence, not value, of MUJOCO_PY_FORCE_CPU to enable CPU rendering.

EGL rendering is nice on Mujoco-py since you don't actually need this line: render_context = MjRenderContext(sim, True, args.device_id, args.backend). If you look at the sim.render code, it lazily initializes MjRenderContext with backend=None on the first call. https://github.com/openai/mujoco-py/blob/4830435a169c1f3e3b5f9b58a7c3d9c39bdf4acb/mujoco_py/mjsim.pyx#L131

Testing Mujoco GLFW Rendering

For the GLFW rendering, Mujoco renders onto a virtual screen. So make sure you have X running, and set the DISPLAY env var to the virtual screen.

Next, your LD_PRELOAD must contain a symlink to libGLEW. For the server, it was export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libGLEW.so. But you must check where your libGLEW is. If the file is not there, install libglew through apt.

Once the two preconditions of LD_PRELOAD and virtual display are satisfied, we can do GLFW rendering. Run DISPLAY=:0 python test_mj.py 0 --backend glfw. If you read the mujoco_py code, you must specify glfw to the MjRenderContext code for it to work. Otherwise, it will default to EGL rendering.

GLFW rendering has the additional limitation that it will only render to the 0th GPU. You cannot select which GPU to do the rendering. See https://github.com/openai/mujoco-py/blob/4830435a169c1f3e3b5f9b58a7c3d9c39bdf4acb/mujoco_py/mjrendercontext.pyx#L96 for more of the logic.

Also, you must explicitly declare the MjRenderContext with the backend argument set to glfw beforehand instead of relying on sim.render to lazily construct it for you. This is because the lazy constructor dsets backend=None, which goes to EGL / CPU rendering.

Benchmarks

This is benchmarked on a server with Nvidia 2080ti, driver 460.

  • CPU time: this script will take forever. Ctrl C.
  • EGL: Around 10-15 seconds.
  • GLFW: Around 10-15 seconds.

Additional Resources

I found reading the mujoco_py code to be helpful for understanding what's happening. https://github.com/openai/mujoco-py/blob/4830435a169c1f3e3b5f9b58a7c3d9c39bdf4acb/mujoco_py/mjrendercontext.pyx

Finally, I found deepmind_control easier to work with. You can specify MUJOCO_GL to egl, glfw, osmesa to choose between the rendering methods, and LD_PRELOAD can be anything. https://github.com/deepmind/dm_control#rendering

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment