Expose a handle for the Python viewer.

This change also requires user scripts to explicitly synchronize changes to physics state to the viewer. The Simulate class was reconfigured so that certain UI events are handled during this sync operation, outside of the render loop on the main thread. These correspond to operations that require access to the full mjModel/mjData.

To support other, more interactive operations (e.g. camera movements), a new mjvSceneState struct is introduced which captures only the portion of the physics state required for scene re-rendering. The mjvSceneState is updated from mjModel/mjData during the viewer sync operation, and is significantly cheaper than a full mj_copyModel and mj_copyData.

Fixes https://github.com/deepmind/mujoco/issues/796

PiperOrigin-RevId: 525723636
Change-Id: Id08d0210a2c067d5afe85e2bf104f276aeddd75e

python/mujoco/mjpython/mjpython.mm

@@ -52,24 +52,28 @@ struct {
 #define CPYTHON_FN(fname) decltype(&::fname) fname
 
 #if PY_MINOR_VERSION >= 8
-  CPYTHON_FN(Py_InitializeFromConfig);
-  CPYTHON_FN(Py_RunMain);
+  // go/keep-sorted start
   CPYTHON_FN(PyConfig_Clear);
   CPYTHON_FN(PyConfig_InitPythonConfig);
   CPYTHON_FN(PyConfig_SetBytesArgv);
+  CPYTHON_FN(Py_InitializeFromConfig);
+  CPYTHON_FN(Py_RunMain);
+  // go/keep-sorted end
 #else
+  // go/keep-sorted start
+  CPYTHON_FN(PyMem_RawFree);
   CPYTHON_FN(Py_DecodeLocale);
   CPYTHON_FN(Py_Initialize);
   CPYTHON_FN(Py_Main);
-  CPYTHON_FN(PyMem_RawFree);
   CPYTHON_FN(Py_SetProgramName);
+  // go/keep-sorted end
 #endif
 
   // go/keep-sorted start
-  CPYTHON_FN(Py_FinalizeEx);
   CPYTHON_FN(PyGILState_Ensure);
   CPYTHON_FN(PyGILState_Release);
   CPYTHON_FN(PyRun_SimpleStringFlags);
+  CPYTHON_FN(Py_FinalizeEx);
   // go/keep-sorted end
 
 #undef CPYTHON_FN
@@ -131,16 +135,16 @@ class _MjPythonImpl(mujoco.viewer._MjPythonBase):
 
   def __init__(self):
     self._cond = threading.Condition()
-    self._model_data = None
+    self._task = None
     self._termination = self.__class__.NOT_TERMINATED
     self._busy = False
 
-  def launch_on_ui_thread(self, model, data):
+  def launch_on_ui_thread(self, model, data, handle_return):
     with self._cond:
-      if self._busy or self._model_data is not None:
+      if self._busy or self._task is not None:
         raise RuntimeError('another MuJoCo viewer is already open')
       else:
-        self._model_data = (model, data)
+        self._task = (model, data, handle_return)
         self._cond.notify()
 
   def terminate(self):
@@ -153,17 +157,17 @@ class _MjPythonImpl(mujoco.viewer._MjPythonBase):
   def get(self):
     with self._cond:
       self._cond.wait_for(
-          lambda: self._model_data is not None or self._termination)
+          lambda: self._task 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
+      task = self._task
       self._busy = True
-      self._model_data = None
-      return model_data
+      self._task = None
+      return task
 
   def done(self):
     with self._cond:
@@ -257,24 +261,28 @@ int main(int argc, char** argv) {
   }
 
 #if PY_MINOR_VERSION >= 8
-  CPYTHON_INITFN(Py_InitializeFromConfig);
-  CPYTHON_INITFN(Py_RunMain);
+  // go/keep-sorted start
   CPYTHON_INITFN(PyConfig_Clear);
   CPYTHON_INITFN(PyConfig_InitPythonConfig);
   CPYTHON_INITFN(PyConfig_SetBytesArgv);
+  CPYTHON_INITFN(Py_InitializeFromConfig);
+  CPYTHON_INITFN(Py_RunMain);
+  // go/keep-sorted end
 #else
+  // go/keep-sorted start
+  CPYTHON_INITFN(PyMem_RawFree);
   CPYTHON_INITFN(Py_DecodeLocale);
   CPYTHON_INITFN(Py_Initialize);
   CPYTHON_INITFN(Py_Main);
-  CPYTHON_INITFN(PyMem_RawFree);
   CPYTHON_INITFN(Py_SetProgramName);
+  // go/keep-sorted end
 #endif
 
   // go/keep-sorted start
-  CPYTHON_INITFN(Py_FinalizeEx);
   CPYTHON_INITFN(PyGILState_Ensure);
   CPYTHON_INITFN(PyGILState_Release);
   CPYTHON_INITFN(PyRun_SimpleStringFlags);
+  CPYTHON_INITFN(Py_FinalizeEx);
   // go/keep-sorted end
 
 #undef CPYTHON_INITFN
@@ -327,17 +335,18 @@ with cond:
 while True:
   try:
     # Wait for an incoming payload.
-    payload = mujoco.viewer._MJPYTHON.get()
+    task = mujoco.viewer._MJPYTHON.get()
 
     # None means that we are exiting.
-    if payload is None:
+    if task is None:
       glfw.terminate()
       break
 
     # Otherwise, launch the viewer.
-    model, data = payload
+    model, data, handle_return = task
     ctypes.CDLL(None).mjpython_show_dock_icon()
-    mujoco.viewer._launch_internal(model, data, run_physics_thread=False)
+    mujoco.viewer._launch_internal(
+        model, data, run_physics_thread=False, handle_return=handle_return)
     ctypes.CDLL(None).mjpython_hide_dock_icon()
 
   finally:

python/mujoco/render.cc

@@ -250,6 +250,7 @@ PYBIND11_MODULE(_render, pymodule) {
   Def<traits::mjr_changeFont>(pymodule);
   Def<traits::mjr_addAux>(pymodule);
   // Skipped: mjr_freeContext (have MjrContext.__del__)
+  Def<traits::mjr_resizeOffscreen>(pymodule);
   Def<traits::mjr_uploadTexture>(pymodule);
   Def<traits::mjr_uploadMesh>(pymodule);
   Def<traits::mjr_uploadHField>(pymodule);

python/mujoco/simulate.cc

@@ -12,29 +12,76 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <cstdint>
+#include <atomic>
 #include <cstring>
 #include <memory>
 #include <string>
+#include <utility>
 
 #include <glfw_adapter.h>
 #include <glfw_dispatch.h>
 #include <simulate.h>
 #include "structs.h"
+#include <pybind11/gil.h>
 #include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
 
 namespace mujoco::python {
 namespace {
+namespace py = ::pybind11;
+
 template <typename T, int N>
 constexpr inline std::size_t sizeof_arr(const T(&arr)[N]) {
   return sizeof(arr);
 }
 
-PYBIND11_MODULE(_simulate, pymodule) {
-  namespace py = ::pybind11;
-  using SimulateMutex = decltype(mujoco::Simulate::mtx);
+class SimulateWrapper : public mujoco::Simulate {
+ public:
+  SimulateWrapper(std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
+                  py::object scn, py::object cam,
+                  py::object opt, py::object pert,  bool fully_managed)
+      : Simulate(std::move(platform_ui_adapter),
+                 scn.cast<MjvSceneWrapper&>().get(),
+                 cam.cast<MjvCameraWrapper&>().get(),
+                 opt.cast<MjvOptionWrapper&>().get(),
+                 pert.cast<MjvPerturbWrapper&>().get(),
+                 fully_managed),
+        m_(py::none()),
+        d_(py::none()),
+        scn_(scn),
+        cam_(cam),
+        opt_(opt),
+        pert_(pert) {}
 
-  py::class_<SimulateMutex>(pymodule, "SimulateMutex")
+  void Load(py::object m, py::object d, const std::string& path) {
+    mjModel* m_raw = m.cast<MjModelWrapper&>().get();
+    mjData* d_raw = d.cast<MjDataWrapper&>().get();
+    {
+      py::gil_scoped_release no_gil;
+      Simulate::Load(m_raw, d_raw, path.c_str());
+    }
+    m_ = m;
+    d_ = d;
+    m_raw_ = m_raw;
+    d_raw_ = d_raw;
+  }
+
+ private:
+  // Hold references to keep these Python objects alive for as long as the
+  // simulate object.
+  py::object m_;
+  py::object d_;
+  py::object scn_;
+  py::object cam_;
+  py::object opt_;
+  py::object pert_;
+
+  mjModel* m_raw_ = nullptr;
+  mjData* d_raw_ = nullptr;
+};
+
+PYBIND11_MODULE(_simulate, pymodule) {
+  py::class_<SimulateMutex>(pymodule, "Mutex")
       .def(
           "__enter__", [](SimulateMutex& mtx) { mtx.lock(); },
           py::call_guard<py::gil_scoped_release>())
@@ -45,36 +92,29 @@ PYBIND11_MODULE(_simulate, pymodule) {
           },
           py::call_guard<py::gil_scoped_release>());
 
-  py::class_<mujoco::Simulate>(pymodule, "Simulate")
-      .def(py::init([]() {
-        return std::make_unique<mujoco::Simulate>(
-            std::make_unique<mujoco::GlfwAdapter>());
+  py::class_<SimulateWrapper>(pymodule, "Simulate")
+      .def_readonly_static("MAX_GEOM", &mujoco::Simulate::kMaxGeom)
+      .def(py::init([](py::object scn, py::object cam, py::object opt,
+                       py::object pert, bool fully_managed) {
+        return std::make_unique<SimulateWrapper>(
+            std::make_unique<mujoco::GlfwAdapter>(), scn, cam, opt, pert,
+            fully_managed);
       }))
-      .def(
-          "render_loop",
-          [](mujoco::Simulate& simulate) { simulate.RenderLoop(); },
-          py::call_guard<py::gil_scoped_release>())
-      .def(
-          "load",
-          [](mujoco::Simulate& simulate, MjModelWrapper& m, MjDataWrapper& d,
-             const std::string& path) {
-            simulate.Load(m.get(), d.get(), path.c_str());
-          },
-          py::call_guard<py::gil_scoped_release>())
-      .def("apply_pose_perturbations",
-           &mujoco::Simulate::ApplyPosePerturbations,
-           py::call_guard<py::gil_scoped_release>())
-      .def("apply_force_perturbations",
-           &mujoco::Simulate::ApplyForcePerturbations,
+      .def("load", &SimulateWrapper::Load)
+      .def("sync", &mujoco::Simulate::Sync,
            py::call_guard<py::gil_scoped_release>())
 
+      .def(
+          "render_loop",
+          [](SimulateWrapper& simulate) { simulate.RenderLoop(); },
+          py::call_guard<py::gil_scoped_release>())
       .def(
           "lock",
-          [](mujoco::Simulate& simulate) -> SimulateMutex& {
+          [](SimulateWrapper& simulate) -> SimulateMutex& {
             return simulate.mtx;
           },
           py::call_guard<py::gil_scoped_release>(),
-          py::return_value_policy::reference)
+          py::return_value_policy::reference_internal)
       .def_readonly("ctrl_noise_std", &mujoco::Simulate::ctrl_noise_std,
                     py::call_guard<py::gil_scoped_release>())
       .def_readonly("ctrl_noise_rate", &mujoco::Simulate::ctrl_noise_rate,
@@ -96,54 +136,52 @@ PYBIND11_MODULE(_simulate, pymodule) {
 
       .def_property(
           "exitrequest",
-          [](mujoco::Simulate& simulate) {
-            return simulate.exitrequest.load();
-          },
-          [](mujoco::Simulate& simulate, bool exitrequest) {
+          [](SimulateWrapper& simulate) { return simulate.exitrequest.load(); },
+          [](SimulateWrapper& simulate, int exitrequest) {
             simulate.exitrequest.store(exitrequest);
           },
           py::call_guard<py::gil_scoped_release>())
 
       .def_property_readonly(
           "uiloadrequest",
-          [](mujoco::Simulate& simulate) {
+          [](SimulateWrapper& simulate) {
             return simulate.uiloadrequest.load();
           },
           py::call_guard<py::gil_scoped_release>())
       .def(
           "uiloadrequest_decrement",
-          [](mujoco::Simulate& simulate) {
+          [](SimulateWrapper& simulate) {
             simulate.uiloadrequest.fetch_sub(1);
           },
           py::call_guard<py::gil_scoped_release>())
 
       .def_property(
           "droploadrequest",
-          [](mujoco::Simulate& simulate) {
+          [](SimulateWrapper& simulate) {
             return simulate.droploadrequest.load();
           },
-          [](mujoco::Simulate& simulate, bool droploadrequest) {
+          [](SimulateWrapper& simulate, bool droploadrequest) {
             simulate.droploadrequest.store(droploadrequest);
           },
           py::call_guard<py::gil_scoped_release>())
       .def_property_readonly(
           "dropfilename",
-          [](mujoco::Simulate& simulate) -> std::string {
+          [](SimulateWrapper& simulate) -> std::string {
             return simulate.dropfilename;
           },
           py::call_guard<py::gil_scoped_release>())
       .def_property_readonly(
           "filename",
-          [](mujoco::Simulate& simulate) -> std::string {
+          [](SimulateWrapper& simulate) -> std::string {
             return simulate.filename;
           },
           py::call_guard<py::gil_scoped_release>())
       .def_property(
           "load_error",
-          [](mujoco::Simulate& simulate) -> std::string {
+          [](SimulateWrapper& simulate) -> std::string {
             return simulate.load_error;
           },
-          [](mujoco::Simulate& simulate, const std::string& error) {
+          [](SimulateWrapper& simulate, const std::string& error) {
             const auto max_length = sizeof_arr(simulate.load_error);
             std::strncpy(simulate.load_error, error.c_str(), max_length - 1);
             simulate.load_error[max_length - 1] = '\0';

python/mujoco/structs.cc

@@ -1305,7 +1305,6 @@ PYBIND11_MODULE(_structs, m) {
   X(realtime);
   X(offwidth);
   X(offheight);
-  X(treedepth);
   X(ellipsoidinertia);
 #undef X
 
@@ -2134,6 +2133,7 @@ This is useful for example when the MJB is not available as a file on disk.)"));
       })
   X(label);
   X(frame);
+  X(bvh_depth);
 #undef X
 
 #define X(var) DefinePyArray(mjvOption, #var, &MjvOptionWrapper::var)

python/mujoco/viewer.py

@@ -16,14 +16,15 @@
 
 import abc
 import atexit
-import code
-import inspect
+import contextlib
 import math
 import os
+import queue
 import sys
 import threading
 import time
 from typing import Callable, Optional, Tuple, Union
+import weakref
 
 import glfw
 import mujoco
@@ -56,7 +57,70 @@ LoaderType = Callable[[], Tuple[mujoco.MjModel, mujoco.MjData]]
 _LoaderWithPathType = Callable[[], Tuple[mujoco.MjModel, mujoco.MjData, str]]
 _InternalLoaderType = Union[LoaderType, _LoaderWithPathType]
 
-Simulate = _simulate.Simulate
+_Simulate = _simulate.Simulate
+
+
+class Handle:
+  """A handle for interacting with a MuJoCo viewer."""
+
+  def __init__(
+      self,
+      sim: _Simulate,
+      scn: mujoco.MjvScene,
+      cam: mujoco.MjvCamera,
+      opt: mujoco.MjvOption,
+      pert: mujoco.MjvPerturb,
+  ):
+    self._sim = weakref.ref(sim)
+    self._scn = scn
+    self._cam = cam
+    self._opt = opt
+    self._pert = pert
+
+  @property
+  def scn(self):
+    return self._scn
+
+  @property
+  def cam(self):
+    return self._cam
+
+  @property
+  def opt(self):
+    return self._opt
+
+  @property
+  def perturb(self):
+    return self._pert
+
+  def close(self):
+    sim = self._sim()
+    if sim is not None:
+      sim.exitrequest = 1
+
+  def is_running(self) -> bool:
+    sim = self._sim()
+    if sim is not None:
+      return sim.exitrequest < 2
+    return False
+
+  def lock(self):
+    sim = self._sim()
+    if sim is not None:
+      return sim.lock()
+    return contextlib.nullcontext()
+
+  def sync(self):
+    sim = self._sim()
+    if sim is not None:
+      with sim.lock():
+        sim.sync()
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    self.close()
 
 
 # Abstract base dispatcher class for systems that require UI calls to be made
@@ -83,7 +147,8 @@ def _file_loader(path: str) -> _LoaderWithPathType:
 
 
 def _reload(
-    simulate: Simulate, loader: _InternalLoaderType
+    simulate: _Simulate, loader: _InternalLoaderType,
+    notify_loaded: Optional[Callable[[], None]] = None
 ) -> Optional[Tuple[mujoco.MjModel, mujoco.MjData]]:
   """Internal function for reloading a model in the viewer."""
   try:
@@ -102,10 +167,13 @@ def _reload(
     path = load_tuple[2] if len(load_tuple) == 3 else ''
     simulate.load(m, d, path)
 
+    if notify_loaded:
+      notify_loaded()
+
     return m, d
 
 
-def _physics_loop(simulate: Simulate, loader: Optional[_InternalLoaderType]):
+def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
   """Physics loop for the GUI, to be run in a separate thread."""
   m: mujoco.MjModel = None
   d: mujoco.MjData = None
@@ -181,11 +249,6 @@ def _physics_loop(simulate: Simulate, loader: Optional[_InternalLoaderType]):
             syncsim = d.time
             simulate.speed_changed = False
 
-            # Clear old perturbations, apply new.
-            d.xfrc_applied[:, :] = 0
-            simulate.apply_pose_perturbations(0)  # Move mocap bodies only.
-            simulate.apply_force_perturbations()
-
             # Run single step, let next iteration deal with timing.
             mujoco.mj_step(m, d)
 
@@ -203,11 +266,6 @@ def _physics_loop(simulate: Simulate, loader: Optional[_InternalLoaderType]):
                 simulate.measured_slowdown = elapsedcpu / elapsedsim
                 measured = True
 
-              # Clear old perturbations, apply new.
-              d.xfrc_applied[:, :] = 0
-              simulate.applyposepertubations(0)  # Move mocap bodies only.
-              simulate.applyforceperturbations()
-
               # Call mj_step.
               mujoco.mj_step(m, d)
 
@@ -215,19 +273,19 @@ def _physics_loop(simulate: Simulate, loader: Optional[_InternalLoaderType]):
               if d.time < prevsim:
                 break
         else:  # simulate.run is False: GUI is paused.
-          # Apply pose perturbation.
-          simulate.applyposepertubations(1)  # Move mocap and dynamic bodies.
 
           # Run mj_forward, to update rendering and joint sliders.
           mujoco.mj_forward(m, d)
 
 
-def _launch_internal(model: Optional[mujoco.MjModel] = None,
-                     data: Optional[mujoco.MjData] = None,
-                     *,
-                     run_physics_thread: bool = True,
-                     loader: Optional[_InternalLoaderType] = None,
-                     simulate: Optional[Simulate] = None) -> None:
+def _launch_internal(
+    model: Optional[mujoco.MjModel] = None,
+    data: Optional[mujoco.MjData] = None,
+    *,
+    run_physics_thread: bool,
+    loader: Optional[_InternalLoaderType] = None,
+    handle_return: Optional['queue.Queue[Handle]'] = None,
+) -> None:
   """Internal API, so that the public API has more readable type annotations."""
   if model is None and data is not None:
     raise ValueError('mjData is specified but mjModel is not')
@@ -236,6 +294,8 @@ def _launch_internal(model: Optional[mujoco.MjModel] = None,
         'mjData should not be specified when an mjModel loader is used')
   elif loader is not None and model is not None:
     raise ValueError('model and loader are both specified')
+  elif run_physics_thread and handle_return is not None:
+    raise ValueError('run_physics_thread and handle_return are both specified')
 
   if loader is None and model is not None:
 
@@ -246,9 +306,14 @@ def _launch_internal(model: Optional[mujoco.MjModel] = None,
 
     loader = _loader
 
-  # The simulate object encapsulates the UI.
-  if simulate is None:
-    simulate = Simulate()
+  if model and not run_physics_thread:
+    scn = mujoco.MjvScene(model, _Simulate.MAX_GEOM)
+  else:
+    scn = mujoco.MjvScene()
+  cam = mujoco.MjvCamera()
+  opt = mujoco.MjvOption()
+  pert = mujoco.MjvPerturb()
+  simulate = _Simulate(scn, cam, opt, pert, run_physics_thread)
 
   # Initialize GLFW if not using mjpython.
   if _MJPYTHON is None:
@@ -256,13 +321,18 @@ def _launch_internal(model: Optional[mujoco.MjModel] = None,
       raise mujoco.FatalError('could not initialize GLFW')
     atexit.register(glfw.terminate)
 
+  notify_loaded = None
+  if handle_return:
+    notify_loaded = (
+        lambda: handle_return.put_nowait(Handle(simulate, scn, cam, opt, pert)))
+
   side_thread = None
   if run_physics_thread:
     side_thread = threading.Thread(
         target=_physics_loop, args=(simulate, loader))
   else:
     side_thread = threading.Thread(
-        target=_reload, args=(simulate, loader))
+        target=_reload, args=(simulate, loader, notify_loaded))
 
   def make_exit_requester(simulate):
     def exit_requester():
@@ -281,18 +351,15 @@ def _launch_internal(model: Optional[mujoco.MjModel] = None,
 def launch(model: Optional[mujoco.MjModel] = None,
            data: Optional[mujoco.MjData] = None,
            *,
-           run_physics_thread: bool = True,
            loader: Optional[LoaderType] = None) -> None:
   """Launches the Simulate GUI."""
-  if not run_physics_thread:
-    mujoco.mj_forward(model, data)
   _launch_internal(
-      model, data, run_physics_thread=run_physics_thread, loader=loader)
+      model, data, run_physics_thread=True, loader=loader)
 
 
 def launch_from_path(path: str) -> None:
   """Launches the Simulate GUI from file path."""
-  _launch_internal(loader=_file_loader(path))
+  _launch_internal(run_physics_thread=True, loader=_file_loader(path))
 
 
 def launch_passive(model: mujoco.MjModel, data: mujoco.MjData) -> None:
@@ -303,12 +370,13 @@ def launch_passive(model: mujoco.MjModel, data: mujoco.MjData) -> None:
     raise ValueError(f'`data` is not a mujoco.MjData: got {data!r}')
 
   mujoco.mj_forward(model, data)
+  handle_return = queue.Queue(1)
 
   if sys.platform != 'darwin':
     thread = threading.Thread(
         target=_launch_internal,
         args=(model, data),
-        kwargs=dict(run_physics_thread=False),
+        kwargs=dict(run_physics_thread=False, handle_return=handle_return),
     )
     thread.daemon = True
     thread.start()
@@ -316,85 +384,10 @@ def launch_passive(model: mujoco.MjModel, data: mujoco.MjData) -> None:
     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)
+          '`mjpython` on macOS')
+    _MJPYTHON.launch_on_ui_thread(model, data, handle_return)
 
-
-def launch_repl(model: mujoco.MjModel, data: mujoco.MjData) -> None:
-  """Launches the Simulate GUI in REPL mode."""
-  ipython_shell = None
-  try:
-    import IPython  # pylint: disable=g-import-not-at-top
-    ipython_shell = IPython.get_ipython()
-    ipython_is_terminal_interactive_shell = isinstance(
-        ipython_shell,
-        IPython.terminal.interactiveshell.TerminalInteractiveShell)
-  except ImportError:
-    ipython_is_terminal_interactive_shell = False
-
-  simulate = Simulate()
-  viewer_is_running = True
-
-  def start_shell(global_variables):
-    if ipython_is_terminal_interactive_shell:
-      ipython_shell.execution_count += 1
-
-      # A SQLite connection can only be used on the same thread that opened it.
-      # We cache the existing connection and reopen on the current thread.
-      old_db = ipython_shell.history_manager.db
-      ipython_shell.history_manager.init_db()
-      ipython_shell.history_manager.new_session()
-
-      try:
-        # Replicate IPython main loop without exiting on keyboard interrupt,
-        # unless the viewer window has already been closed.
-        # (https://github.com/ipython/ipython/blob/8.9.0/IPython/terminal/interactiveshell.py#L701)
-        while viewer_is_running and ipython_shell.keep_running:
-          print(ipython_shell.separate_in, end='')
-          try:
-            c = ipython_shell.prompt_for_code()
-          except EOFError:
-            if not ipython_shell.confirm_exit or ipython_shell.ask_yes_no(
-                'Do you really want to exit ([y]/n)?', 'y', 'n'):
-              ipython_shell.ask_exit()
-            if not ipython_shell.keep_running and simulate is not None:
-              simulate.exitrequest = True
-          else:
-            if c:
-              ipython_shell.run_cell(c, store_history=True)
-      finally:
-        # Close the temporary history DB connection and restore the old one.
-        ipython_shell.history_manager.end_session()
-        ipython_shell.history_manager.db.close()
-        ipython_shell.history_manager.db = old_db
-        ipython_shell.execution_count -= 1
-    else:
-      code.InteractiveConsole(locals=global_variables).interact()
-
-  # End IPython history session on the main thread. We will need to open
-  # a new session in the REPL thread.
-  if ipython_is_terminal_interactive_shell:
-    ipython_shell.history_manager.end_session()
-
-  try:
-    # Continue the IPython REPL session in a separate thread.
-    repl_thread = threading.Thread(
-        target=start_shell, args=(inspect.stack()[1][0].f_globals,))
-    repl_thread.start()
-
-    # Launch the viewer on the main thread.
-    mujoco.mj_forward(model, data)
-    _launch_internal(
-        model, data, run_physics_thread=False, simulate=simulate)
-    simulate = None
-
-    # Wait until the REPL thread quits, then restore the IPython history
-    # DB session on the main thread.
-    viewer_is_running = False
-    repl_thread.join()
-  finally:
-    if ipython_is_terminal_interactive_shell:
-      ipython_shell.history_manager.new_session()
+  return handle_return.get()
 
 
 if __name__ == '__main__':