Add keyboard event callback for Python passive viewer.

Fixes #766
Fixes #846

PiperOrigin-RevId: 549449372
Change-Id: I37d17f1162c66d0e8482d402aae22d9a16b59deb

python/mujoco/mjpython/mjpython.mm

@@ -115,12 +115,12 @@ class _MjPythonImpl(mujoco.viewer._MjPythonBase):
     self._termination = self.__class__.NOT_TERMINATED
     self._busy = False
 
-  def launch_on_ui_thread(self, model, data, handle_return):
+  def launch_on_ui_thread(self, model, data, handle_return, key_callback):
     with self._cond:
       if self._busy or self._task is not None:
         raise RuntimeError('another MuJoCo viewer is already open')
       else:
-        self._task = (model, data, handle_return)
+        self._task = (model, data, handle_return, key_callback)
         self._cond.notify()
 
   def terminate(self):
@@ -294,10 +294,11 @@ while True:
       break
 
     # Otherwise, launch the viewer.
-    model, data, handle_return = task
+    model, data, handle_return, key_callback = task
     ctypes.CDLL(None).mjpython_show_dock_icon()
     mujoco.viewer._launch_internal(
-        model, data, run_physics_thread=False, handle_return=handle_return)
+        model, data, run_physics_thread=False, handle_return=handle_return,
+        key_callback=key_callback)
     ctypes.CDLL(None).mjpython_hide_dock_icon()
 
   finally:

python/mujoco/simulate.cc

@@ -38,6 +38,33 @@ constexpr inline std::size_t sizeof_arr(const T (&arr)[N]) {
   return sizeof(arr);
 }
 
+template <typename Adapter>
+class UIAdapterWithPyCallback : public Adapter {
+ public:
+  template <typename... Args>
+  UIAdapterWithPyCallback(py::handle key_callback, Args&&... args)
+      : Adapter(std::forward<Args>(args)...) {
+    if (!key_callback.is_none()) {
+      Py_XINCREF(key_callback.ptr());
+      key_callback_ = key_callback.ptr();
+    }
+  }
+
+  ~UIAdapterWithPyCallback() override { Py_XDECREF(key_callback_); }
+
+ protected:
+  void OnKey(int key, int scancode, int act) override {
+    Adapter::OnKey(key, scancode, act);
+    if (this->IsKeyDownEvent(act) && key_callback_) {
+      py::gil_scoped_acquire gil;
+      (py::handle(key_callback_))(this->last_key_);
+    }
+  }
+
+ private:
+  PyObject* key_callback_ = nullptr;
+};
+
 class SimulateWrapper {
  public:
   SimulateWrapper(std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
@@ -166,10 +193,12 @@ PYBIND11_MODULE(_simulate, pymodule) {
   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) {
+                       py::object pert, bool fully_managed,
+                       py::object key_callback) {
         return std::make_unique<SimulateWrapper>(
-            std::make_unique<mujoco::GlfwAdapter>(), scn, cam, opt, pert,
-            fully_managed);
+            std::make_unique<UIAdapterWithPyCallback<mujoco::GlfwAdapter>>(
+                key_callback),
+            scn, cam, opt, pert, fully_managed);
       }))
       .def("destroy", &SimulateWrapper::Destroy,
            py::call_guard<py::gil_scoped_release>())

python/mujoco/viewer.py

@@ -52,6 +52,7 @@ SIM_REFRESH_FRACTION = 0.7
 
 CallbackType = Callable[[mujoco.MjModel, mujoco.MjData], None]
 LoaderType = Callable[[], Tuple[mujoco.MjModel, mujoco.MjData]]
+KeyCallbackType = Callable[[int], None]
 
 # Loader function that also returns a file path for the GUI to display.
 _LoaderWithPathType = Callable[[], Tuple[mujoco.MjModel, mujoco.MjData, str]]
@@ -142,9 +143,16 @@ class Handle:
 # 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):
+  def launch_on_ui_thread(
+      self,
+      model: mujoco.MjModel,
+      data: mujoco.MjData,
+      handle_return: Optional['queue.Queue[Handle]'],
+      key_callback: Optional[KeyCallbackType],
+  ):
     pass
 
+
 # When running under mjpython, the launcher initializes this object.
 _MJPYTHON: Optional[_MjPythonBase] = None
 
@@ -299,6 +307,7 @@ def _launch_internal(
     run_physics_thread: bool,
     loader: Optional[_InternalLoaderType] = None,
     handle_return: Optional['queue.Queue[Handle]'] = None,
+    key_callback: Optional[KeyCallbackType] = None,
 ) -> None:
   """Internal API, so that the public API has more readable type annotations."""
   if model is None and data is not None:
@@ -327,7 +336,7 @@ def _launch_internal(
   cam = mujoco.MjvCamera()
   opt = mujoco.MjvOption()
   pert = mujoco.MjvPerturb()
-  simulate = _Simulate(scn, cam, opt, pert, run_physics_thread)
+  simulate = _Simulate(scn, cam, opt, pert, run_physics_thread, key_callback)
 
   # Initialize GLFW if not using mjpython.
   if _MJPYTHON is None:
@@ -377,12 +386,20 @@ def launch_from_path(path: str) -> None:
   _launch_internal(run_physics_thread=True, loader=_file_loader(path))
 
 
-def launch_passive(model: mujoco.MjModel, data: mujoco.MjData) -> Handle:
+def launch_passive(
+    model: mujoco.MjModel,
+    data: mujoco.MjData,
+    *,
+    key_callback: Optional[KeyCallbackType] = None,
+) -> Handle:
   """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 key_callback is not None and not callable(key_callback):
+    raise ValueError(
+        f'`key_callback` is not callable: got {key_callback!r}')
 
   mujoco.mj_forward(model, data)
   handle_return = queue.Queue(1)
@@ -391,7 +408,11 @@ def launch_passive(model: mujoco.MjModel, data: mujoco.MjData) -> Handle:
     thread = threading.Thread(
         target=_launch_internal,
         args=(model, data),
-        kwargs=dict(run_physics_thread=False, handle_return=handle_return),
+        kwargs=dict(
+            run_physics_thread=False,
+            handle_return=handle_return,
+            key_callback=key_callback,
+        ),
     )
     thread.daemon = True
     thread.start()
@@ -400,7 +421,7 @@ def launch_passive(model: mujoco.MjModel, data: mujoco.MjData) -> Handle:
       raise RuntimeError(
           '`launch_passive` requires that the Python script be run under '
           '`mjpython` on macOS')
-    _MJPYTHON.launch_on_ui_thread(model, data, handle_return)
+    _MJPYTHON.launch_on_ui_thread(model, data, handle_return, key_callback)
 
   return handle_return.get()