Add viewer.launch_passive and a mjpython launcher for macOS.

The `launch_passive` function launches the GUI viewer in a non-blocking manner, allowing the Python script or REPL to continue execution. The viewer is automatically kept up to date with any subsequent modifications to mjModel and mjData.

Note that when run inside a REPL (including IPython), `launch_passive` is functionally identical to `launch_repl`.

On Linux and Windows, this is achieved by spawning a new thread and launching the GUI there.

On macOS, this is not possible as all Cocoa API calls must be made on the "macOS main thread", which is always the first thread launched in a process and carries the `com.apple.main-thread` dispatch queue. We also cannot simply trampoline from a Python script on the main thread into the user's script on a side thread because CPython's signal handler can only be installed on the "Python main thread". Putting the user's script in a side thread means that it cannot e.g. gracefully handle SIGINT by catching a KeyboardInterrupt exception. To work around this, we ship a custom Python launcher on macOS called `mjpython`. This launcher is a native binary that spawns a pthread and initialize the Python interpreter on that thread, thus allowing "Python main thread" and "macOS main thread" to represent two distinct threads. From Python's point of view, the "macOS main thread" is a secondary thread that runs a loop that continuously empties a Queue of (mjModel, mjData) and launches a viewer.

PiperOrigin-RevId: 517167868
Change-Id: Icac9d2126bbb4760d47e0b9300e0a979cffa4338

python/MANIFEST.in

@@ -1,3 +1,4 @@
 include LICENSE *.md
 recursive-include mujoco *.h *.cc *.mm CMakeLists.txt Simulate*.cmake
 recursive-include cmake *.cmake
+recursive-include mujoco/mjpython mjpython.* Info.plist

python/mujoco/CMakeLists.txt

@@ -22,6 +22,9 @@ set(CMAKE_POLICY_DEFAULT_CMP0069 NEW)
 
 enable_language(C)
 enable_language(CXX)
+if(APPLE)
+  enable_language(OBJCXX)
+endif()
 
 if(MSVC AND MSVC_VERSION GREATER_EQUAL 1927)
   set(CMAKE_CXX_STANDARD 20) # For forceinline lambdas.
@@ -407,6 +410,13 @@ set(LIBRARIES_FOR_WHEEL
     "$<TARGET_FILE:mujoco>"
 )
 
+if(APPLE)
+  add_executable(mjpython mjpython/mjpython.mm)
+  target_include_directories(mjpython PRIVATE "${Python3_INCLUDE_DIRS}")
+  target_link_libraries(mjpython PRIVATE "-framework Cocoa")
+  set(LIBRARIES_FOR_WHEEL "${LIBRARIES_FOR_WHEEL}" "$<TARGET_FILE:mjpython>")
+endif()
+
 if(MUJOCO_PYTHON_MAKE_WHEEL)
   add_custom_target(
     wheel ALL
@@ -434,4 +444,7 @@ if(MUJOCO_PYTHON_MAKE_WHEEL)
     _structs
     mujoco
   )
+  if(APPLE)
+    add_dependencies(wheel mjpython)
+  endif()
 endif()

python/mujoco/mjpython/Info.plist

@@ -0,0 +1,34 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+    <key>CFBundleName</key>
+    <string>mjpython</string>
+    <key>CFBundleIdentifier</key>
+    <string>org.mujoco.mjpython</string>
+    <key>CFBundleVersion</key>
+    <string>2.3.2</string>
+    <key>CFBundleGetInfoString</key>
+    <string>2.3.2</string>
+    <key>CFBundleLongVersionString</key>
+    <string>2.3.2</string>
+    <key>CFBundleShortVersionString</key>
+    <string>2.3.2</string>
+    <key>CFBundleExecutable</key>
+    <string>mjpython</string>
+    <key>CFBundleIconFile</key>
+    <string>mjpython.icns</string>
+    <key>CFBundlePackageType</key>
+    <string>APPL</string>
+    <key>NSHumanReadableCopyright</key>
+    <string>Copyright 2023 DeepMind Technologies Limited.</string>
+    <key>CFBundleInfoDictionaryVersion</key>
+    <string>6.0</string>
+    <key>CFBundleDevelopmentRegion</key>
+    <string>en</string>
+    <key>CFBundleSupportedPlatforms</key>
+    <array>
+        <string>MacOSX</string>
+    </array>
+</dict>
+</plist>

python/mujoco/mjpython/mjpython.mm

@@ -0,0 +1,298 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// 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 <atomic>
+#import <cstdlib>
+#import <iostream>
+
+#import <dlfcn.h>
+#import <pthread.h>
+#import <sys/resource.h>
+
+#import <Cocoa/Cocoa.h>
+#import <Python.h>
+
+// Wrap Objective-C Cocoa calls into C-style functions with default visibility,
+// so that we can dlsym and call them from Python via ctypes.
+extern "C" {
+__attribute__((used)) void mjpython_hide_dock_icon() {
+  [NSApp setActivationPolicy:NSApplicationActivationPolicyAccessory];
+}
+__attribute__((used)) void mjpython_show_dock_icon() {
+  [NSApp setActivationPolicy:NSApplicationActivationPolicyRegular];
+}
+}
+
+// TODO(b/273744079): Remove Python 3.7 code after end-of-life (27 Jun 2023).
+namespace {
+struct {
+#define CPYTHON_FN(fname) decltype(&::fname) fname
+
+#if PY_MINOR_VERSION >= 8
+  CPYTHON_FN(Py_InitializeFromConfig);
+  CPYTHON_FN(Py_RunMain);
+  CPYTHON_FN(PyConfig_Clear);
+  CPYTHON_FN(PyConfig_InitPythonConfig);
+  CPYTHON_FN(PyConfig_SetBytesArgv);
+#else
+  CPYTHON_FN(Py_DecodeLocale);
+  CPYTHON_FN(Py_Initialize);
+  CPYTHON_FN(Py_Main);
+  CPYTHON_FN(PyMem_RawFree);
+  CPYTHON_FN(Py_SetProgramName);
+#endif
+
+  // go/keep-sorted start
+  CPYTHON_FN(Py_FinalizeEx);
+  CPYTHON_FN(PyGILState_Ensure);
+  CPYTHON_FN(PyGILState_Release);
+  CPYTHON_FN(PyRun_SimpleStringFlags);
+  // go/keep-sorted end
+
+#undef CPYTHON_FN
+} cpython;
+
+std::atomic_bool py_initialized = false;
+
+struct Args {
+  int argc;
+  char** argv;
+};
+
+// The Python main thread (distinct from the macOS main thread, which executes the main function).
+void* mjpython_pymain(void* vargs) {
+  Args* args = static_cast<Args*>(vargs);
+  PyGILState_STATE gil;
+
+  // Initialize the Python interpreter.
+#if PY_MINOR_VERSION >= 8
+  PyConfig config;
+  cpython.PyConfig_InitPythonConfig(&config);
+  cpython.PyConfig_SetBytesArgv(&config, args->argc, args->argv);
+  cpython.Py_InitializeFromConfig(&config);
+  cpython.PyConfig_Clear(&config);
+#else
+  // Convert each argv to wchar_t* (needed for Py_Main).
+  wchar_t** wargv = static_cast<wchar_t**>(std::calloc(args->argc, sizeof(wchar_t*)));
+  for (int i = 0; i < args->argc; ++i) {
+    wargv[i] = cpython.Py_DecodeLocale(args->argv[i], nullptr);
+  }
+  cpython.Py_SetProgramName(wargv[0]);
+  cpython.Py_Initialize();
+#endif
+
+  // Set up the condition variable to pass control back to the macOS main thread.
+  gil = cpython.PyGILState_Ensure();
+  cpython.PyRun_SimpleStringFlags("import threading; cond = threading.Condition()", nullptr);
+  py_initialized.store(true);
+
+  // Wait until GLFW is initialized on macOS main thread, set up the queue and an atexit hook
+  // to enqueue a termination flag upon exit.
+  cpython.PyRun_SimpleStringFlags(R"(
+import atexit
+
+# The mujoco.viewer module should only be imported here after glfw.init() in the macOS main thread.
+with cond:
+  cond.wait()
+import mujoco.viewer
+
+# Similar to a queue.Queue(maxsize=1), but where only one active task is allowed at a time.
+# With queue.Queue(1), another item is allowed to be enqueued before task_done is called.
+class _MjPythonImpl(mujoco.viewer._MjPythonBase):
+
+  # Termination statuses
+  NOT_TERMINATED = 0
+  TERMINATION_REQUESTED = 1
+  TERMINATION_ACCEPTED = 2
+  TERMINATED = 3
+
+  def __init__(self):
+    self._cond = threading.Condition()
+    self._model_data = None
+    self._termination = self.__class__.NOT_TERMINATED
+    self._busy = False
+
+  def launch_on_ui_thread(self, model, data):
+    with self._cond:
+      if self._busy or self._model_data is not None:
+        raise RuntimeError('another MuJoCo viewer is already open')
+      else:
+        self._model_data = (model, data)
+        self._cond.notify()
+
+  def terminate(self):
+    with self._cond:
+      self._termination = self.__class__.TERMINATION_REQUESTED
+      self._cond.notify()
+      self._cond.wait_for(
+          lambda: self._termination == self.__class__.TERMINATED)
+
+  def get(self):
+    with self._cond:
+      self._cond.wait_for(
+          lambda: self._model_data is not None or self._termination)
+
+      if self._termination:
+        if self._termination == self.__class__.TERMINATION_REQUESTED:
+          self._termination = self.__class__.TERMINATION_ACCEPTED
+        return None
+
+      model_data = self._model_data
+      self._busy = True
+      self._model_data = None
+      return model_data
+
+  def done(self):
+    with self._cond:
+      self._busy = False
+      if self._termination == self.__class__.TERMINATION_ACCEPTED:
+        self._termination = self.__class__.TERMINATED
+      self._cond.notify()
+
+
+mujoco.viewer._MJPYTHON = _MjPythonImpl()
+atexit.register(mujoco.viewer._MJPYTHON.terminate)
+del _MjPythonImpl  # Don't pollute globals for user script.
+
+with cond:
+  cond.notify()
+del cond  # Don't pollute globals for user script.
+)", nullptr);
+
+  // Run the Python interpreter main loop.
+#if PY_MINOR_VERSION >= 8
+  cpython.Py_RunMain();
+#else
+  cpython.Py_Main(args->argc, wargv);
+#endif
+
+  // Tear down the interpreter.
+  cpython.Py_FinalizeEx();
+#if PY_MINOR_VERSION < 8
+  for (int i = 0; i < args->argc; ++i) {
+    cpython.PyMem_RawFree(wargv[i]);
+    wargv[i] = nullptr;
+  }
+  std::free(wargv);
+  wargv = nullptr;
+#endif
+  return nullptr;
+}
+}  // namespace
+
+int main(int argc, char** argv) {
+  const char* libpython_path = getenv("MJPYTHON_LIBPYTHON");
+  if (!libpython_path || !libpython_path[0]) {
+    std::cerr << "This binary must be launched via the mjpython.py script.\n";
+    return 1;
+  }
+
+  // Resolve libpython at runtime to prevent linking against the wrong dylib. The correct libpython
+  // path is passed from a Python trampoline script, which ran inside the desired interpreter and
+  // exec'd this binary.
+  void* libpython = dlopen(libpython_path, RTLD_NOW | RTLD_GLOBAL);
+
+  // Look up required CPython API functions from table of symbols already loaded into the process.
+#define CPYTHON_INITFN(fname) \
+    cpython.fname = reinterpret_cast<decltype(cpython.fname)>(dlsym(libpython, #fname))
+
+#if PY_MINOR_VERSION >= 8
+  CPYTHON_INITFN(Py_InitializeFromConfig);
+  CPYTHON_INITFN(Py_RunMain);
+  CPYTHON_INITFN(PyConfig_Clear);
+  CPYTHON_INITFN(PyConfig_InitPythonConfig);
+  CPYTHON_INITFN(PyConfig_SetBytesArgv);
+#else
+  CPYTHON_INITFN(Py_DecodeLocale);
+  CPYTHON_INITFN(Py_Initialize);
+  CPYTHON_INITFN(Py_Main);
+  CPYTHON_INITFN(PyMem_RawFree);
+  CPYTHON_INITFN(Py_SetProgramName);
+#endif
+
+  // go/keep-sorted start
+  CPYTHON_INITFN(Py_FinalizeEx);
+  CPYTHON_INITFN(PyGILState_Ensure);
+  CPYTHON_INITFN(PyGILState_Release);
+  CPYTHON_INITFN(PyRun_SimpleStringFlags);
+  // go/keep-sorted end
+
+#undef CPYTHON_INITFN
+
+  // Package up argc and argv together to pass to pthread_create.
+  Args args{argc, argv};
+
+  // Create a thread to be used as the "Python main thread".
+  pthread_t pymain_thread = [&args]() {
+    // Set the stack size of the Python main thread to be the same as the OS main thread.
+    // (e.g. the default pthread stack size is too small to import NumPy)
+    rlimit limit;
+    getrlimit(RLIMIT_STACK, &limit);
+
+    pthread_attr_t attr;
+    pthread_attr_init(&attr);
+    pthread_attr_setstacksize(&attr, limit.rlim_cur);
+
+    pthread_t thread;
+    pthread_create(&thread, &attr, &mjpython_pymain, &args);
+    return thread;
+  }();
+
+  // Busy-wait until Python interpreter is initialized.
+  while (!py_initialized.load()) {}
+
+  // Initialize GLFW on the macOS main thread, yield control to Python main thread and wait for it
+  // to finish setting up _MJPYTHON, then serve incoming viewer launch requests.
+  PyGILState_STATE gil = cpython.PyGILState_Ensure();
+  cpython.PyRun_SimpleStringFlags(R"(
+import ctypes
+
+# GLFW must be initialized on the OS main thread (i.e. here).
+import glfw
+import mujoco.viewer
+
+glfw.init()
+glfw.poll_events()
+ctypes.CDLL(None).mjpython_hide_dock_icon()
+
+# Wait for Python main thread to finish setting up _MJPYTHON
+with cond:
+  cond.notify()
+  cond.wait()
+
+while True:
+  try:
+    # Wait for an incoming payload.
+    payload = mujoco.viewer._MJPYTHON.get()
+
+    # None means that we are exiting.
+    if payload is None:
+      glfw.terminate()
+      break
+
+    # Otherwise, launch the viewer.
+    model, data = payload
+    ctypes.CDLL(None).mjpython_show_dock_icon()
+    mujoco.viewer._launch_internal(model, data, run_physics_thread=False)
+    ctypes.CDLL(None).mjpython_hide_dock_icon()
+
+  finally:
+    mujoco.viewer._MJPYTHON.done()
+)", nullptr);
+  cpython.PyGILState_Release(gil);
+
+  // Tear everything down.
+  pthread_join(pymain_thread, nullptr);
+  dlclose(libpython);
+}