Modify Python simulate so that the Reload button works.

Also allow a `Callable[[], [MjModel, MjData]]` to be provided as an argument to `launch` instead of a fixed MjModel/MjData instance. This allows users to customise the behaviour of the Reload button.

Context: https://github.com/deepmind/mujoco/discussions/576#discussioncomment-4161429
PiperOrigin-RevId: 492296268
Change-Id: Ic4b797d2ea8a6967f7e3e7316f2184ea94a86d8b

python/mujoco/simulate.cc

@@ -17,9 +17,7 @@
 #include <string>
 
 #include <simulate.h>
-#include "raw.h"
 #include "structs.h"
-#include <pybind11/detail/common.h>
 #include <pybind11/pybind11.h>
 
 namespace mujoco::python {
@@ -47,8 +45,9 @@ PYBIND11_MODULE(_simulate, pymodule) {
           py::call_guard<py::gil_scoped_release>())
       .def(
           "load",
-          [](mujoco::Simulate& simulate, MjModelWrapper& m, MjDataWrapper& d) {
+          [](mujoco::Simulate& simulate, const std::string& path,
-            simulate.load("", m.get(), d.get());
+             MjModelWrapper& m, MjDataWrapper& d) {
+            simulate.load(path.c_str(), m.get(), d.get());
           },
           py::call_guard<py::gil_scoped_release>())
       .def("applyposepertubations", &mujoco::Simulate::applyposepertubations,

python/mujoco/simulate.py

@@ -21,7 +21,7 @@ import math
 import os
 import threading
 import time
-from typing import Callable, Optional
+from typing import Callable, Optional, Union
 
 import glfw
 import mujoco
@@ -48,35 +48,32 @@ MAX_SYNC_MISALIGN = 0.1
 SIM_REFRESH_FRACTION = 0.7
 
 CallbackType = Callable[[mujoco.MjModel, mujoco.MjData], None]
+LoaderType = Callable[[], tuple[mujoco.MjModel, mujoco.MjData]]
+
+# Loader function that also returns a file path for the GUI to display.
+_LoaderWithPathType = Callable[[], tuple[mujoco.MjModel, mujoco.MjData, str]]
+_InternalLoaderType = Union[LoaderType, _LoaderWithPathType]
 
 Simulate = _simulate.Simulate
 
 
-def _reload_from_file(simulate: Simulate, filename: str):
+def _file_loader(path: str) -> _LoaderWithPathType:
-  """Loads an MJCF model into the Simulate GUI."""
+  """Loads an MJCF model from file path."""
-  try:
-    m = mujoco.MjModel.from_xml_path(filename)
-  except mujoco.FatalError as e:
-    m = None
-    simulate.load_error = str(e)
 
-  if m is not None:
+  def load(path=path) -> tuple[mujoco.MjModel, mujoco.MjData, str]:
+    m = mujoco.MjModel.from_xml_path(path)
     d = mujoco.MjData(m)
-    simulate.load(m, d)
+    return m, d, path
-    mujoco.mj_forward(m, d)
-  else:
-    d = None
 
-  return m, d
+  return load
 
 
-def _physics_loop(simulate: Simulate,
+def _physics_loop(simulate: Simulate, loader: Optional[_InternalLoaderType]):
-                  m: Optional[mujoco.MjModel],
-                  d: Optional[mujoco.MjData]):
   """Physics loop for the Simulate GUI, to be run in a separate thread."""
-  ctrlnoise = None
+  m: mujoco.MjModel = None
-  if m is not None:
+  d: mujoco.MjData = None
-    ctrlnoise = np.zeros((m.nu,))
+  ctrlnoise = np.array([])
+  reload = True
 
   # CPU-sim synchronization point.
   synccpu = 0.0
@@ -86,20 +83,40 @@ def _physics_loop(simulate: Simulate,
   while not simulate.exitrequest:
     if simulate.droploadrequest:
       simulate.droploadrequest = 0
-      new_m, new_d = _reload_from_file(simulate, simulate.dropfilename)
+      loader = _file_loader(simulate.dropfilename)
-      if new_m is not None:
+      reload = True
-        m = new_m
-        d = new_d
-        ctrlnoise = np.zeros((m.nu,))
 
     if simulate.uiloadrequest:
       simulate.uiloadrequest_decrement()
-      new_m, new_d = _reload_from_file(simulate, simulate.dropfilename)
+      reload = True
-      if new_m is not None:
+
-        m = new_m
+    if reload and loader is not None:
-        d = new_d
+      try:
+        load_tuple = loader()
+      except Exception as e:  # pylint: disable=broad-except
+        simulate.load_error = str(e)
+      else:
+        # Do not assign to m and d until simulate.load is done!
+        # This is because simulate.load needs to clean up mjvScene and
+        # mjrContext. This cleanup logic requires access to the old m and d.
+        new_m, new_d = load_tuple[:2]
+
+        # If the loader does not raise an exception then we assume that it
+        # successfully created mjModel and mjData. This is specified in the type
+        # annotation, but we perform a runtime assertion here as well to prevent
+        # possible segmentation faults.
+        assert new_m is not None and new_d is not None
+
+        path = load_tuple[2] if len(load_tuple) == 3 else ''
+        simulate.load(path, new_m, new_d)
+
+        # We can now allow the old m and d to be deleted.
+        m, d = new_m, new_d
+        mujoco.mj_forward(m, d)
         ctrlnoise = np.zeros((m.nu,))
 
+    reload = False
+
     # Sleep for 1 ms or yield, to let main thread run.
     if simulate.run != 0 and simulate.busywait != 0:
       time.sleep(0)
@@ -186,15 +203,28 @@ def _physics_loop(simulate: Simulate,
           mujoco.mj_forward(m, d)
 
 
-def launch(model: Optional[mujoco.MjModel] = None,
+def _launch_internal(model: Optional[mujoco.MjModel] = None,
-           data: Optional[mujoco.MjData] = None,
+                     data: Optional[mujoco.MjData] = None,
-           *,
+                     *,
-           run_physics_thread: bool = True) -> None:
+                     run_physics_thread: bool = True,
-  """Launches the Simulate GUI."""
+                     loader: Optional[_InternalLoaderType] = 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')
-  elif model is not None and data is None:
+  elif callable(model) and data is not None:
-    data = mujoco.MjData(model)
+    raise ValueError(
+        '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')
+
+  if loader is None and model is not None:
+
+    def _loader(m=model, d=data) -> tuple[mujoco.MjModel, mujoco.MjData]:
+      if d is None:
+        d = mujoco.MjData(m)
+      return m, d
+
+    loader = _loader
 
   # The simulate object encapsulates the UI.
   simulate = Simulate()
@@ -207,21 +237,30 @@ def launch(model: Optional[mujoco.MjModel] = None,
 
   if run_physics_thread:
     physics_thread = threading.Thread(
-        target=_physics_loop, args=(simulate, model, data))
+        target=_physics_loop, args=(simulate, loader))
     physics_thread.start()
 
-  # Load the initial model, if one is given.
-  if model is not None:
-    t = threading.Thread(target=simulate.load, args=(model, data))
-    t.start()
-    del t
-
   simulate.renderloop()
 
   if run_physics_thread:
     physics_thread.join()
 
 
+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."""
+  _launch_internal(
+      model, data, run_physics_thread=run_physics_thread, loader=loader)
+
+
+def launch_from_path(path: str) -> None:
+  """Launches the Simulate GUI from file path."""
+  _launch_internal(loader=_file_loader(path))
+
+
 def launch_repl(model: mujoco.MjModel, data: mujoco.MjData) -> None:
   """EXPERIMENTAL FEATURE: Launches the Simulate GUI in REPL mode."""
   try:
@@ -253,8 +292,7 @@ if __name__ == '__main__':
   def main(argv) -> None:
     del argv
     if _MJCF_PATH.value is not None:
-      model = mujoco.MjModel.from_xml_path(os.path.expanduser(_MJCF_PATH.value))
+      launch_from_path(os.path.expanduser(_MJCF_PATH.value))
-      launch(model)
     else:
       launch()