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
This commit is contained in:
Saran Tunyasuvunakool
2023-03-16 10:52:22 -07:00
committed by Copybara-Service
parent 9a97674e1e
commit 230e2780de
8 changed files with 509 additions and 8 deletions
+1
View File
@@ -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
+13
View File
@@ -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()
+34
View File
@@ -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>
Binary file not shown.
+298
View File
@@ -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);
}
+44
View File
@@ -0,0 +1,44 @@
#!/usr/bin/env python
# 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.
# ==============================================================================
"""Python interpreter trampoline for macOS to support non-block Cocoa GUI.
This script executes a native binary that runs the CPython interpreter entry
point in a separate thread, thus leaving the macOS main thread free for Cocoa
GUI calls without blocking the user's Python script. In other words, Python's
idea of the "main thread" is different from the thread that holds the
com.apple.main-thread DispatchQueue.
"""
import importlib.util
import os
import sys
import sysconfig
def main(argv):
os.environ['MJPYTHON_LIBPYTHON'] = os.path.join(
sysconfig.get_config_var('PYTHONFRAMEWORKPREFIX'),
sysconfig.get_config_var('INSTSONAME'),
)
argv[0] = sys.executable
mujoco_dir = os.path.dirname(importlib.util.find_spec('mujoco').origin)
os.execve(
os.path.join(mujoco_dir, 'MuJoCo (mjpython).app/Contents/MacOS/mjpython'),
argv, os.environ)
if __name__ == '__main__':
main(sys.argv)
+50 -5
View File
@@ -14,6 +14,7 @@
# ==============================================================================
"""Interactive GUI viewer for MuJoCo."""
import abc
import atexit
import code
import inspect
@@ -58,6 +59,18 @@ _InternalLoaderType = Union[LoaderType, _LoaderWithPathType]
Simulate = _simulate.Simulate
# Abstract base dispatcher class for systems that require UI calls to be made
# on a specific thread (e.g. macOS). This is subclassed by system-specific
# Python launcher (mjpython) to implement the required dispatching mechanism.
class _MjPythonBase(metaclass=abc.ABCMeta):
def launch_on_ui_thread(self, model: mujoco.MjModel, data: mujoco.MjData):
pass
# When running under mjpython, the launcher initializes this object.
_MJPYTHON: Optional[_MjPythonBase] = None
def _file_loader(path: str) -> _LoaderWithPathType:
"""Loads an MJCF model from file path."""
@@ -238,11 +251,11 @@ def _launch_internal(model: Optional[mujoco.MjModel] = None,
if simulate is None:
simulate = Simulate()
# Initialize GLFW.
if not glfw.init():
raise mujoco.FatalError('could not initialize GLFW')
atexit.register(glfw.terminate)
# Initialize GLFW if not using mjpython.
if _MJPYTHON is None:
if not glfw.init():
raise mujoco.FatalError('could not initialize GLFW')
atexit.register(glfw.terminate)
side_thread = None
if run_physics_thread:
@@ -251,8 +264,18 @@ def _launch_internal(model: Optional[mujoco.MjModel] = None,
else:
side_thread = threading.Thread(
target=_reload, args=(simulate, loader))
def make_exit_requester(simulate):
def exit_requester():
simulate.exitrequest = True
return exit_requester
exit_requester = make_exit_requester(simulate)
atexit.register(exit_requester)
side_thread.start()
simulate.renderloop()
atexit.unregister(exit_requester)
side_thread.join()
@@ -271,6 +294,28 @@ def launch_from_path(path: str) -> None:
_launch_internal(loader=_file_loader(path))
def launch_passive(model: mujoco.MjModel, data: mujoco.MjData) -> None:
"""Launches a passive Simulate GUI without blocking the running thread."""
if not isinstance(model, mujoco.MjModel):
raise ValueError(f'`model` is not a mujoco.MjModel: got {model!r}')
if not isinstance(data, mujoco.MjData):
raise ValueError(f'`data` is not a mujoco.MjData: got {data!r}')
if sys.platform != 'darwin':
thread = threading.Thread(
target=_launch_internal,
args=(model, data),
kwargs=dict(run_physics_thread=False),
)
thread.daemon = True
thread.start()
else:
if not isinstance(_MJPYTHON, _MjPythonBase):
raise RuntimeError(
'`launch_passive` requires that the Python script be run under '
'`mjpython`')
_MJPYTHON.launch_on_ui_thread(model, data)
def launch_repl(model: mujoco.MjModel, data: mujoco.MjData) -> None:
"""Launches the Simulate GUI in REPL mode."""
ipython_shell = None
+69 -3
View File
@@ -15,6 +15,7 @@
"""Install script for MuJoCo."""
import fnmatch
import logging
import os
import platform
import random
@@ -29,6 +30,7 @@ import setuptools
from setuptools import find_packages
from setuptools import setup
from setuptools.command import build_ext
from setuptools.command import install_scripts
__version__ = '2.3.2'
@@ -70,6 +72,7 @@ def get_external_lib_patterns():
else:
return ['libmujoco.so.*']
def get_plugin_lib_patterns():
if platform.system() == 'Windows':
return ['*.dll']
@@ -159,6 +162,8 @@ class BuildCMakeExtension(build_ext.build_ext):
self._copy_external_libraries()
self._copy_mujoco_headers()
self._copy_plugin_libraries()
if self._is_apple:
self._copy_mjpython()
def _find_mujoco(self):
if MUJOCO_PATH not in os.environ:
@@ -213,6 +218,26 @@ class BuildCMakeExtension(build_ext.build_ext):
shutil.copyfile(os.path.join(directory, filename),
os.path.join(dst, filename))
def _copy_mjpython(self):
src_dir = os.path.join(os.path.dirname(__file__), 'mujoco/mjpython')
dst_contents_dir = os.path.join(
os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),
'MuJoCo (mjpython).app/Contents')
os.makedirs(dst_contents_dir)
shutil.copyfile(os.path.join(src_dir, 'Info.plist'),
os.path.join(dst_contents_dir, 'Info.plist'))
dst_bin_dir = os.path.join(dst_contents_dir, 'MacOS')
os.makedirs(dst_bin_dir)
shutil.copyfile(os.path.join(self.build_temp, 'mjpython'),
os.path.join(dst_bin_dir, 'mjpython'))
os.chmod(os.path.join(dst_bin_dir, 'mjpython'), 0o755)
dst_resources_dir = os.path.join(dst_contents_dir, 'Resources')
os.makedirs(dst_resources_dir)
shutil.copyfile(os.path.join(src_dir, 'mjpython.icns'),
os.path.join(dst_resources_dir, 'mjpython.icns'))
def _configure_cmake(self):
"""Check for CMake."""
cmake = os.environ.get(MUJOCO_CMAKE, 'cmake')
@@ -276,6 +301,40 @@ class BuildCMakeExtension(build_ext.build_ext):
build_path = os.path.join(self.build_temp, os.path.basename(dest_path))
shutil.copyfile(build_path, dest_path)
class InstallScripts(install_scripts.install_scripts):
"""Strips file extension from executable scripts whose names end in `.py`."""
def run(self):
super().run()
oldfiles = self.outfiles
files = set(oldfiles)
self.outfiles = []
for oldfile in oldfiles:
if oldfile.endswith('.py'):
newfile = oldfile[:-3]
else:
newfile = oldfile
renamed = False
if newfile not in files and not os.path.exists(newfile):
if not self.dry_run:
os.rename(oldfile, newfile)
renamed = True
if renamed:
logging.info(
'Renaming %s script to %s',
os.path.basename(oldfile),
os.path.basename(newfile),
)
self.outfiles.append(newfile)
files.remove(oldfile)
files.add(newfile)
else:
self.outfiles.append(oldfile)
def find_data_files(package_dir, patterns):
"""Recursively finds files whose names match the given shell patterns."""
paths = set()
@@ -287,8 +346,7 @@ def find_data_files(package_dir, patterns):
paths.add(os.path.join(relative_dirpath, filename))
return list(paths)
setup(
SETUP_KWARGS = dict(
name='mujoco',
version=__version__,
author='DeepMind',
@@ -312,7 +370,10 @@ setup(
'Programming Language :: Python :: 3.11',
'Topic :: Scientific/Engineering',
],
cmdclass=dict(build_ext=BuildCMakeExtension),
cmdclass=dict(
build_ext=BuildCMakeExtension,
install_scripts=InstallScripts,
),
ext_modules=[
CMakeExtension('mujoco._callbacks'),
CMakeExtension('mujoco._constants'),
@@ -351,3 +412,8 @@ setup(
]),
},
)
if platform.system() == 'Darwin':
SETUP_KWARGS['scripts'] = ['mujoco/mjpython/mjpython.py']
setup(**SETUP_KWARGS)