Merge pull request #29 from ARISE-Initiative/gpu_id_specification

GPU id specification

LICENSE

@@ -19,3 +19,10 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
+
+This software includes the partial implementation of Deepmind Mujoco https://github.com/deepmind/mujoco. 
+Deepmind Mujoco is licensed under the Apache License, Version 2.0 (the "License");
+you may not use the files except in compliance with the License. 
+
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0

robosuite/macros.py

@@ -32,7 +32,6 @@
 # so we skip concatenating all images together by default, unless this flag is set to True
 CONCATENATE_IMAGES = False
 
-
 try:
     from robosuite.macros_private import *
 except ImportError:
@@ -43,3 +42,5 @@
         "\nIt is recommended to use a private macro file"
         "\nTo setup, run: python robosuite/scripts/setup_macros.py"
     )
+
+MUJOCO_GPU_RENDERING = True

robosuite/renderers/context/egl_context.py

@@ -0,0 +1,155 @@
+# Modifications Copyright 2022 The robosuite Authors
+# Original Copyright 2018 The dm_control Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import atexit
+import ctypes
+import os
+
+PYOPENGL_PLATFORM = os.environ.get("PYOPENGL_PLATFORM")
+
+if not PYOPENGL_PLATFORM:
+    os.environ["PYOPENGL_PLATFORM"] = "egl"
+elif PYOPENGL_PLATFORM.lower() != "egl":
+    raise ImportError(
+        "Cannot use EGL rendering platform. "
+        "The PYOPENGL_PLATFORM environment variable is set to {!r} "
+        "(should be either unset or 'egl')."
+    )
+
+from mujoco.egl import egl_ext as EGL
+from OpenGL import error
+
+
+def create_initialized_egl_device_display(device_id=0):
+    """Creates an initialized EGL display directly on a device."""
+    all_devices = EGL.eglQueryDevicesEXT()
+    selected_device = (
+        os.environ.get("CUDA_VISIBLE_DEVICES", None)
+        if os.environ.get("MUJOCO_EGL_DEVICE_ID", None) is None
+        else os.environ.get("MUJOCO_EGL_DEVICE_ID", None)
+    )
+    if selected_device is None:
+        candidates = all_devices
+        if device_id == -1:
+            device_idx = 0
+        else:
+            device_idx = device_id
+    else:
+        if not selected_device.isdigit():
+            device_inds = [int(x) for x in selected_device.split(",")]
+            if device_id == -1:
+                device_idx = device_inds[0]
+            else:
+                assert(device_id in device_inds), "specified device id is not made visible in environment variables."
+                device_idx = device_id
+        else:
+            device_idx = int(selected_device)
+        if not 0 <= device_idx < len(all_devices):
+            raise RuntimeError(
+                f"The MUJOCO_EGL_DEVICE_ID environment variable must be an integer "
+                f"between 0 and {len(all_devices)-1} (inclusive), got {device_idx}."
+            )
+    candidates = all_devices[device_idx : device_idx + 1]
+    for device in candidates:
+        display = EGL.eglGetPlatformDisplayEXT(EGL.EGL_PLATFORM_DEVICE_EXT, device, None)
+        if display != EGL.EGL_NO_DISPLAY and EGL.eglGetError() == EGL.EGL_SUCCESS:
+            # `eglInitialize` may or may not raise an exception on failure depending
+            # on how PyOpenGL is configured. We therefore catch a `GLError` and also
+            # manually check the output of `eglGetError()` here.
+            try:
+                initialized = EGL.eglInitialize(display, None, None)
+            except error.GLError:
+                pass
+            else:
+                if initialized == EGL.EGL_TRUE and EGL.eglGetError() == EGL.EGL_SUCCESS:
+                    return display
+    return EGL.EGL_NO_DISPLAY
+
+
+global EGL_DISPLAY
+EGL_DISPLAY = None
+
+EGL_ATTRIBUTES = (
+    EGL.EGL_RED_SIZE,
+    8,
+    EGL.EGL_GREEN_SIZE,
+    8,
+    EGL.EGL_BLUE_SIZE,
+    8,
+    EGL.EGL_ALPHA_SIZE,
+    8,
+    EGL.EGL_DEPTH_SIZE,
+    24,
+    EGL.EGL_STENCIL_SIZE,
+    8,
+    EGL.EGL_COLOR_BUFFER_TYPE,
+    EGL.EGL_RGB_BUFFER,
+    EGL.EGL_SURFACE_TYPE,
+    EGL.EGL_PBUFFER_BIT,
+    EGL.EGL_RENDERABLE_TYPE,
+    EGL.EGL_OPENGL_BIT,
+    EGL.EGL_NONE,
+)
+
+
+class EGLGLContext:
+    """An EGL context for headless accelerated OpenGL rendering on GPU devices."""
+
+    def __init__(self, max_width, max_height, device_id=0):
+
+        del max_width, max_height  # unused
+        num_configs = ctypes.c_long()
+        config_size = 1
+        config = EGL.EGLConfig()
+        EGL.eglReleaseThread()
+        global EGL_DISPLAY
+        if EGL_DISPLAY is None:
+            # only initialize for the first time
+            EGL_DISPLAY = create_initialized_egl_device_display(device_id=device_id)
+            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, ctypes.byref(config), config_size, num_configs)
+        if num_configs.value < 1:
+            raise RuntimeError(
+                "EGL failed to find a framebuffer configuration that matches the "
+                "desired attributes: {}".format(EGL_ATTRIBUTES)
+            )
+        EGL.eglBindAPI(EGL.EGL_OPENGL_API)
+        self._context = EGL.eglCreateContext(EGL_DISPLAY, config, EGL.EGL_NO_CONTEXT, None)
+        if not self._context:
+            raise RuntimeError("Cannot create an EGL context.")
+
+    def make_current(self):
+        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."""
+        if self._context:
+            current_context = EGL.eglGetCurrentContext()
+            if current_context and self._context.address == current_context.address:
+                EGL.eglMakeCurrent(EGL_DISPLAY, EGL.EGL_NO_SURFACE, EGL.EGL_NO_SURFACE, EGL.EGL_NO_CONTEXT)
+            EGL.eglDestroyContext(EGL_DISPLAY, self._context)
+            EGL.eglReleaseThread()
+        self._context = None
+
+    def __del__(self):
+        self.free()

robosuite/renderers/context/glfw_context.py

@@ -0,0 +1,24 @@
+# Copyright 2017 The dm_control Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""An OpenGL context created via GLFW."""
+
+from mujoco.glfw import GLContext
+
+
+class GLFWGLContext(GLContext):
+    """An OpenGL context created via GLFW."""
+
+    def __init__(self, max_width, max_height, device_id=0):
+        super().__init__(max_width, max_height)

robosuite/renderers/context/osmesa_context.py

@@ -0,0 +1,26 @@
+# Copyright 2018 The dm_control Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""An OSMesa context for software-based OpenGL rendering."""
+
+import os
+
+from mujoco.osmesa import GLContext
+
+
+class OSMesaGLContext(GLContext):
+    """An OSMesa context for software-based OpenGL rendering."""
+
+    def __init__(self, max_width, max_height, device_id=-1):
+        super().__init__(max_width, max_height)

robosuite/utils/binding_utils.py

@@ -14,6 +14,48 @@
 
 _MjSim_render_lock = Lock()
 
+import ctypes
+import ctypes.util
+import os
+import platform
+import subprocess
+
+import robosuite.macros as macros
+
+_SYSTEM = platform.system()
+if _SYSTEM == "Windows":
+    ctypes.WinDLL(os.path.join(os.path.dirname(__file__), "mujoco.dll"))
+
+CUDA_VISIBLE_DEVICES = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+if CUDA_VISIBLE_DEVICES != "":
+    MUJOCO_EGL_DEVICE_ID = os.environ.get("MUJOCO_EGL_DEVICE_ID", None)
+    if MUJOCO_EGL_DEVICE_ID is not None:
+        assert MUJOCO_EGL_DEVICE_ID.isdigit() and (
+            MUJOCO_EGL_DEVICE_ID in CUDA_VISIBLE_DEVICES
+        ), "MUJOCO_EGL_DEVICE_ID needs to be set to one of the device id specified in CUDA_VISIBLE_DEVICES"
+
+if macros.MUJOCO_GPU_RENDERING and os.environ.get("MUJOCO_GL", None) not in ["osmesa", "glx"]:
+    # If gpu rendering is specified in macros, then we enforce gpu
+    # option for rendering
+    os.environ["MUJOCO_GL"] = "egl"
+_MUJOCO_GL = os.environ.get("MUJOCO_GL", "").lower().strip()
+if _MUJOCO_GL not in ("disable", "disabled", "off", "false", "0"):
+    _VALID_MUJOCO_GL = ("enable", "enabled", "on", "true", "1", "glfw", "")
+    if _SYSTEM == "Linux":
+        _VALID_MUJOCO_GL += ("glx", "egl", "osmesa")
+    elif _SYSTEM == "Windows":
+        _VALID_MUJOCO_GL += ("wgl",)
+    elif _SYSTEM == "Darwin":
+        _VALID_MUJOCO_GL += ("cgl",)
+    if _MUJOCO_GL not in _VALID_MUJOCO_GL:
+        raise RuntimeError(f"invalid value for environment variable MUJOCO_GL: {_MUJOCO_GL}")
+    if _SYSTEM == "Linux" and _MUJOCO_GL == "osmesa":
+        from robosuite.renderers.context.osmesa_context import OSMesaGLContext as GLContext
+    elif _SYSTEM == "Linux" and _MUJOCO_GL == "egl":
+        from robosuite.renderers.context.egl_context import EGLGLContext as GLContext
+    else:
+        from robosuite.renderers.context.glfw_context import GLFWGLContext as GLContext
+
 
 class MjRenderContext:
     """
@@ -23,21 +65,14 @@ class MjRenderContext:
     See https://github.com/openai/mujoco-py/blob/4830435a169c1f3e3b5f9b58a7c3d9c39bdf4acb/mujoco_py/mjrendercontext.pyx
     """
 
-    def __init__(self, sim, offscreen=True, device_id=-1):
+    def __init__(self, sim, offscreen=True, device_id=-1, max_width=640, max_height=480):
         assert offscreen, "only offscreen supported for now"
         self.sim = sim
         self.offscreen = offscreen
         self.device_id = device_id
 
-        if offscreen:
-            if self.device_id is not None and self.device_id >= 0:
-                os.environ["MUJOCO_GL"] = "egl"
-                os.environ["MUJOCO_EGL_DEVICE_ID"] = str(self.device_id)
-            else:
-                os.environ["MUJOCO_GL"] = "osmesa"
-
         # setup GL context with defaults for now
-        self.gl_ctx = mujoco.GLContext(max_width=640, max_height=480)
+        self.gl_ctx = GLContext(max_width=max_width, max_height=max_height, device_id=self.device_id)
         self.gl_ctx.make_current()
 
         # Ensure the model data has been updated so that there
@@ -163,8 +198,8 @@ def __del__(self):
 
 
 class MjRenderContextOffscreen(MjRenderContext):
-    def __init__(self, sim, device_id):
-        super().__init__(sim, offscreen=True, device_id=device_id)
+    def __init__(self, sim, device_id, max_width=640, max_height=480):
+        super().__init__(sim, offscreen=True, device_id=device_id, max_width=max_width, max_height=max_height)
 
 
 class MjSimState: