From 1ee57a1fcab6279e6bf0ec04925a8645bc367a7b Mon Sep 17 00:00:00 2001
From: Yu Fang <squarefk@gmail.com>
Date: Mon, 23 Sep 2024 12:01:25 -0700
Subject: [PATCH] Fix the initialization of the EGL context in multiprocessing

---
 python/mujoco/egl/__init__.py | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/python/mujoco/egl/__init__.py b/python/mujoco/egl/__init__.py
index 045b2aff..832c1108 100644
--- a/python/mujoco/egl/__init__.py
+++ b/python/mujoco/egl/__init__.py
@@ -62,13 +62,7 @@ def create_initialized_egl_device_display():
   return EGL.EGL_NO_DISPLAY
 
 
-EGL_DISPLAY = create_initialized_egl_device_display()
-if EGL_DISPLAY == EGL.EGL_NO_DISPLAY:
-  raise ImportError(
-      'Cannot initialize a EGL device display. This likely means that your EGL '
-      'driver does not support the PLATFORM_DEVICE extension, which is '
-      'required for creating a headless rendering context.')
-atexit.register(EGL.eglTerminate, EGL_DISPLAY)
+EGL_DISPLAY = None
 
 
 EGL_ATTRIBUTES = (
@@ -95,6 +89,17 @@ class GLContext:
     # ctypes syntax for making an array of length config_size.
     configs = (EGL.EGLConfig * config_size)()
     EGL.eglReleaseThread()
+    global EGL_DISPLAY
+    if EGL_DISPLAY is None:
+      # only initialize for the first time
+      EGL_DISPLAY = create_initialized_egl_device_display()
+      if EGL_DISPLAY == EGL.EGL_NO_DISPLAY:
+        raise ImportError(
+          "Cannot initialize a EGL device display. This likely means that your EGL "
+          "driver does not support the PLATFORM_DEVICE extension, which is "
+          "required for creating a headless rendering context."
+        )
+      atexit.register(EGL.eglTerminate, EGL_DISPLAY)
     EGL.eglChooseConfig(
         EGL_DISPLAY,
         EGL_ATTRIBUTES,
@@ -112,12 +117,14 @@ class GLContext:
       raise RuntimeError('Cannot create an EGL context.')
 
   def make_current(self):
+    global EGL_DISPLAY
     if not EGL.eglMakeCurrent(
         EGL_DISPLAY, EGL.EGL_NO_SURFACE, EGL.EGL_NO_SURFACE, self._context):
       raise RuntimeError('Failed to make the EGL context current.')
 
   def free(self):
     """Frees resources associated with this context."""
+    global EGL_DISPLAY
     if self._context:
       current_context = EGL.eglGetCurrentContext()
       if current_context and self._context.address == current_context.address: