Headless support

CMakeLists.txt

@@ -44,7 +44,6 @@ set(BUILD_DOCUMENTATION OFF)
 include(FetchContent)
 
 find_package(Eigen3 REQUIRED)
-find_package(glfw3 REQUIRED)
 find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
 find_package(MuJoCo REQUIRED)
 
@@ -67,29 +66,6 @@ FetchContent_Declare(pybind11
     GIT_PROGRESS TRUE
     EXCLUDE_FROM_ALL
 )
-set(MuJoCo_UI_URL https://raw.githubusercontent.com/google-deepmind/mujoco/${MuJoCo_VERSION}/simulate/)
-set(MuJoCo_UI_FILES platform_ui_adapter.cc platform_ui_adapter.h glfw_adapter.cc glfw_adapter.h glfw_adapter.cc glfw_corevideo.h glfw_corevideo.mm glfw_dispatch.cc glfw_dispatch.h)
-# Set the destination directory for the downloaded files
-set(MuJoCo_UI_DIR "${CMAKE_BINARY_DIR}/mujoco_ui")
-
-# Create the directory if it doesn't exist
-file(MAKE_DIRECTORY ${MuJoCo_UI_DIR})
-
-# Loop through each file and download it only if it doesn't already exist
-foreach(file ${MuJoCo_UI_FILES})
-  set(output_file "${MuJoCo_UI_DIR}/${file}")
-  if (NOT EXISTS ${output_file})
-    message(STATUS "Downloading ${file}...")
-    file(DOWNLOAD ${MuJoCo_UI_URL}${file} ${output_file})
-  else()
-    message(STATUS "File ${file} already exists, skipping download.")
-  endif()
-endforeach()
-
-# Add the include directories or link the files as needed
-add_library(mujoco_ui_lib ${MuJoCo_UI_DIR}/glfw_adapter.cc ${MuJoCo_UI_DIR}/platform_ui_adapter.cc ${MuJoCo_UI_DIR}/glfw_dispatch.cc)
-target_link_libraries(mujoco_ui_lib MuJoCo::MuJoCo glfw)
-target_include_directories(mujoco_ui_lib INTERFACE ${MuJoCo_UI_DIR})
 
 FetchContent_MakeAvailable(libfranka rl pybind11)
 include(compile_scenes)

python/rcs/_core/common.pyi

@@ -232,6 +232,7 @@ def FrankaHandTCPOffset() -> numpy.ndarray[tuple[typing.Literal[4], typing.Liter
 def IdentityRotMatrix() -> numpy.ndarray[tuple[typing.Literal[3], typing.Literal[3]], numpy.dtype[numpy.float64]]: ...
 def IdentityRotQuatVec() -> numpy.ndarray[typing.Literal[4], numpy.dtype[numpy.float64]]: ...
 def IdentityTranslation() -> numpy.ndarray[typing.Literal[3], numpy.dtype[numpy.float64]]: ...
+def _bootstrap_egl(fn_addr: int, display: int, context: int) -> None: ...
 def robots_meta_config(robot_type: RobotType) -> RobotMetaConfig: ...
 
 FR3: RobotType  # value = <RobotType.FR3: 0>

python/rcs/_core/sim.pyi

@@ -12,6 +12,7 @@ import rcs._core.common
 __all__ = [
     "CameraType",
     "FrameSet",
+    "GuiClient",
     "Sim",
     "SimCameraConfig",
     "SimCameraSet",
@@ -24,7 +25,6 @@ __all__ = [
     "default_free",
     "fixed",
     "free",
-    "open_gui_window",
     "tracking",
 ]
 M = typing.TypeVar("M", bound=int)
@@ -73,6 +73,12 @@ class FrameSet:
     @property
     def timestamp(self) -> float: ...
 
+class GuiClient:
+    def __init__(self, id: str) -> None: ...
+    def get_model_bytes(self) -> bytes: ...
+    def set_model_and_data(self, arg0: int, arg1: int) -> None: ...
+    def sync(self) -> None: ...
+
 class Sim:
     def __init__(self, mjmdl: int, mjdata: int) -> None: ...
     def _start_gui_server(self, id: str) -> None: ...
@@ -170,8 +176,6 @@ class SimRobotState(rcs._core.common.RobotState):
     @property
     def target_angles(self) -> numpy.ndarray[M, numpy.dtype[numpy.float64]]: ...
 
-def open_gui_window(uuid: str) -> None: ...
-
 default_free: CameraType  # value = <CameraType.default_free: 3>
 fixed: CameraType  # value = <CameraType.fixed: 2>
 free: CameraType  # value = <CameraType.free: 0>

python/rcs/camera/hw.py

@@ -9,7 +9,8 @@
 import cv2
 import numpy as np
 from rcs._core.common import BaseCameraConfig
-from rcs.camera.interface import BaseCameraSet, Frame, FrameSet, SimpleFrameRate
+from rcs.camera.interface import BaseCameraSet, Frame, FrameSet
+from rcs.utils import SimpleFrameRate
 
 
 class HardwareCamera(typing.Protocol):

python/rcs/camera/interface.py

@@ -1,7 +1,6 @@
 import logging
 from dataclasses import dataclass
 from datetime import datetime
-from time import sleep, time
 from typing import Any, Protocol
 
 from rcs._core.common import BaseCameraConfig
@@ -10,35 +9,6 @@
 logger.setLevel(logging.INFO)
 
 
-class SimpleFrameRate:
-    def __init__(self, frame_rate: int | float):
-        self.t: float | None = None
-        self._last_print: float | None = None
-        self.frame_rate = frame_rate
-
-    def reset(self):
-        self.t = None
-
-    def __call__(self):
-        if self.t is None:
-            self.t = time()
-            self._last_print = self.t
-            sleep(1 / self.frame_rate if isinstance(self.frame_rate, int) else self.frame_rate)
-            return
-        sleep_time = (
-            1 / self.frame_rate - (time() - self.t)
-            if isinstance(self.frame_rate, int)
-            else self.frame_rate - (time() - self.t)
-        )
-        if sleep_time > 0:
-            sleep(sleep_time)
-        if self._last_print is None or time() - self._last_print > 10:
-            self._last_print = time()
-            logger.info(f"FPS: {1 / (time() - self.t)}")
-
-        self.t = time()
-
-
 @dataclass(kw_only=True)
 class DataFrame:
     data: Any

python/rcs/sim/__init__.py

@@ -0,0 +1,20 @@
+from rcs._core.sim import (
+    SimGripper,
+    SimGripperConfig,
+    SimGripperState,
+    SimRobot,
+    SimRobotConfig,
+    SimRobotState,
+)
+from rcs.sim.sim import Sim, gui_loop
+
+__all__ = [
+    "Sim",
+    "SimRobot",
+    "SimRobotConfig",
+    "SimRobotState",
+    "SimGripper",
+    "SimGripperConfig",
+    "SimGripperState",
+    "gui_loop",
+]

python/rcs/sim/egl_bootstrap.py

@@ -0,0 +1,33 @@
+"""
+Load the EGL library, create a persistent GLContext, and register it with the C++ backend.
+
+Globals prevent the library and context from being garbage-collected.  
+Call `bootstrap()` to complete initialization.
+"""
+
+import ctypes
+import ctypes.util
+import os
+
+import mujoco.egl
+import rcs._core as _cxx
+from mujoco.egl import GLContext
+
+name = ctypes.util.find_library("EGL")
+if name is None:
+    msg = "libEGL not found"
+    raise OSError(msg)
+
+_egl = ctypes.CDLL(name, mode=os.RTLD_LOCAL | os.RTLD_NOW)
+
+addr_make_current = ctypes.cast(_egl.eglMakeCurrent, ctypes.c_void_p).value
+
+ctx = GLContext(max_width=3840, max_height=2160)
+
+egl_display = mujoco.egl.EGL_DISPLAY.address
+egl_context = ctx._context.address
+
+
+def bootstrap():
+    assert addr_make_current is not None
+    _cxx.common._bootstrap_egl(addr_make_current, egl_display, egl_context)

python/rcs/sim/sim.py

@@ -0,0 +1,81 @@
+import atexit
+import multiprocessing as mp
+import uuid
+from logging import getLogger
+from multiprocessing.synchronize import Event as EventClass
+from os import PathLike
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Optional
+
+import mujoco as mj
+import mujoco.viewer
+from rcs._core.sim import GuiClient as _GuiClient
+from rcs._core.sim import Sim as _Sim
+from rcs.sim import egl_bootstrap
+from rcs.utils import SimpleFrameRate
+
+egl_bootstrap.bootstrap()
+logger = getLogger(__name__)
+
+
+# Target frames per second
+FPS = 60
+
+
+def gui_loop(gui_uuid: str, close_event):
+    frame_rate = SimpleFrameRate(FPS, "gui_loop")
+    gui_client = _GuiClient(gui_uuid)
+    model_bytes = gui_client.get_model_bytes()
+    with NamedTemporaryFile(mode="wb") as f:
+        f.write(model_bytes)
+        model = mujoco.MjModel.from_binary_path(f.name)
+    data = mujoco.MjData(model)
+    gui_client.set_model_and_data(model._address, data._address)
+    mujoco.mj_step(model, data)
+    with mujoco.viewer.launch_passive(model, data) as viewer:
+        while not close_event.is_set():
+            mujoco.mj_step(model, data)
+            viewer.sync()
+            gui_client.sync()
+            frame_rate()
+
+
+class Sim(_Sim):
+    def __init__(self, mjmdl: str | PathLike):
+        mjmdl = Path(mjmdl)
+        if mjmdl.suffix == ".xml":
+            self.model = mj.MjModel.from_xml_path(str(mjmdl))
+        elif mjmdl.suffix == ".mjb":
+            self.model = mj.MjModel.from_binary_path(str(mjmdl))
+        else:
+            msg = f"Filetype {mjmdl.suffix} is unknown"
+            logger.error(msg)
+        self.data = mj.MjData(self.model)
+        super().__init__(self.model._address, self.data._address)
+        self._mp_context = mp.get_context("spawn")
+        self._gui_uuid: Optional[str] = None
+        self._gui_client: Optional[_GuiClient] = None
+        self._gui_process: Optional[mp.context.SpawnProcess] = None
+        self._stop_event: Optional[EventClass] = None
+
+    def close_gui(self):
+        if self._stop_event is not None:
+            self._stop_event.set()
+        if self._gui_process is not None:
+            self._gui_process.join()
+        self._stop_gui_server()
+
+    def open_gui(self):
+        if self._gui_uuid is None:
+            self._gui_uuid = "rcs_" + str(uuid.uuid4())
+            self._start_gui_server(self._gui_uuid)
+        if self._gui_client is None:
+            ctx = mp.get_context("spawn")
+            self._stop_event = ctx.Event()
+            self._gui_process = ctx.Process(
+                target=gui_loop,
+                args=(self._gui_uuid, self._stop_event),
+            )
+            self._gui_process.start()
+        atexit.register(self.close_gui)

python/rcs/utils.py

@@ -0,0 +1,41 @@
+import logging
+from time import perf_counter, sleep
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+class SimpleFrameRate:
+    def __init__(self, frame_rate: int | float, loop_name: str = "SimpleFrameRate"):
+        """SimpleFrameRate is a utility class to manage frame rates in a simple way.
+        It allows you to call it in a loop, and it will sleep the necessary time to maintain the desired frame rate.
+
+        Args:
+            frame_rate (int | float): The desired frame rate in frames per second (int) or as a time interval in seconds (float).
+        """
+        self.t: float | None = None
+        self._last_print: float | None = None
+        self.frame_rate = frame_rate
+        self.loop_name = loop_name
+
+    def reset(self):
+        self.t = None
+
+    def __call__(self):
+        if self.t is None:
+            self.t = perf_counter()
+            self._last_print = self.t
+            sleep(1 / self.frame_rate if isinstance(self.frame_rate, int) else self.frame_rate)
+            return
+        sleep_time = (
+            1 / self.frame_rate - (perf_counter() - self.t)
+            if isinstance(self.frame_rate, int)
+            else self.frame_rate - (perf_counter() - self.t)
+        )
+        if sleep_time > 0:
+            sleep(sleep_time)
+        if self._last_print is None or perf_counter() - self._last_print > 10:
+            self._last_print = perf_counter()
+            logger.info(f"FPS {self.loop_name}: {1 / (perf_counter() - self.t)}")
+
+        self.t = perf_counter()

src/common/CMakeLists.txt

@@ -1,3 +1,3 @@
 add_library(common)
-target_sources(common PRIVATE Pose.cpp Robot.cpp IK.cpp)
+target_sources(common PRIVATE Pose.cpp Robot.cpp IK.cpp utils.cpp)
 target_link_libraries(common PUBLIC Eigen3::Eigen mdl)

src/common/utils.cpp

@@ -0,0 +1,25 @@
+#include "utils.h"
+
+#include <EGL/egl.h>
+
+#include <iostream>
+
+namespace rcs {
+namespace common {
+static PFNEGLMAKECURRENTPROC g_makeCurrent = nullptr;
+static EGLDisplay g_display = EGL_NO_DISPLAY;
+static EGLSurface g_surface = EGL_NO_SURFACE;
+static EGLContext g_context = EGL_NO_CONTEXT;
+
+void bootstrap_egl(uintptr_t fn_addr, uintptr_t dpy, uintptr_t ctx) {
+  g_makeCurrent = reinterpret_cast<PFNEGLMAKECURRENTPROC>(fn_addr);
+  g_display = reinterpret_cast<EGLDisplay>(dpy);
+  g_context = reinterpret_cast<EGLContext>(ctx);
+}
+
+void ensure_current() {
+  if (!g_makeCurrent(g_display, g_surface, g_surface, g_context))
+    throw std::runtime_error("eglMakeCurrent failed");
+}
+}  // namespace common
+}  // namespace rcs