Add support for user-defined figures in the MuJoCo viewer.

PiperOrigin-RevId: 726983516
Change-Id: Ifb821f3cafb163d4f3077d76b0668037b4bc5a9c

python/mujoco/simulate.cc

@@ -28,6 +28,7 @@
 #include <pybind11/gil.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/pytypes.h>
+#include <pybind11/stl.h>
 
 namespace mujoco::python {
 namespace {
@@ -69,11 +70,10 @@ class UIAdapterWithPyCallback : public Adapter {
 class SimulateWrapper {
  public:
   SimulateWrapper(std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
-                  py::object cam, py::object opt,
-                  py::object pert, py::object user_scn, bool is_passive)
+                  py::object cam, py::object opt, py::object pert,
+                  py::object user_scn, bool is_passive)
       : simulate_(new mujoco::Simulate(
-            std::move(platform_ui_adapter),
-            cam.cast<MjvCameraWrapper&>().get(),
+            std::move(platform_ui_adapter), cam.cast<MjvCameraWrapper&>().get(),
             opt.cast<MjvOptionWrapper&>().get(),
             pert.cast<MjvPerturbWrapper&>().get(), is_passive)),
         m_(py::none()),
@@ -126,6 +126,28 @@ class SimulateWrapper {
   py::object GetModel() const { return m_; }
   py::object GetData() const { return d_; }
 
+  mjrRect GetViewport() const {
+    // Return the viewport corresponding to the 3D view, i.e. the viewer window
+    // without the UI elements.
+    return simulate_->uistate.rect[3];
+  }
+
+  void SetFigures(
+      const std::vector<std::pair<mjrRect, py::object>>& viewports_figures) {
+    // Pairs of [viewport, figure], where viewport corresponds to the location
+    // of the figure on the viewer window.
+    std::vector<std::pair<mjrRect, mjvFigure>> user_figures;
+    for (const auto& [viewport, figure] : viewports_figures) {
+      mjvFigure casted_figure = *figure.cast<MjvFigureWrapper&>().get();
+      user_figures.push_back(std::make_pair(viewport, casted_figure));
+    }
+
+    // Set them all at once to prevent figure flickering.
+    simulate_->user_figures_ = user_figures;
+  }
+
+  void ClearFigures() { simulate_->user_figures_.clear(); }
+
  private:
   mujoco::Simulate* simulate_;
   std::atomic_int destroyed_ = 0;
@@ -173,14 +195,14 @@ inline auto CallIfNotNull(void (mujoco::Simulate::*func)(Args...)) {
 }
 
 template <typename T>
-inline auto GetIfNotNull(T mujoco::Simulate::*member) {
+inline auto GetIfNotNull(T mujoco::Simulate::* member) {
   return [member](SimulateWrapper& wrapper) -> T& {
     return SimulateRefOrThrow(wrapper).*member;
   };
 }
 
 template <typename T, typename... Args>
-inline auto SetIfNotNull(T mujoco::Simulate::*member) {
+inline auto SetIfNotNull(T mujoco::Simulate::* member) {
   return [member](SimulateWrapper& wrapper, const T& value) -> void {
     SimulateRefOrThrow(wrapper).*member = value;
   };
@@ -205,8 +227,7 @@ PYBIND11_MODULE(_simulate, pymodule) {
                        py::object key_callback) {
         bool is_passive = !run_physics_thread;
         return std::make_unique<SimulateWrapper>(
-            std::make_unique<UIAdapterWithPyCallback<UIAdapter>>(
-                key_callback),
+            std::make_unique<UIAdapterWithPyCallback<UIAdapter>>(key_callback),
             scn, cam, opt, pert, is_passive);
       }))
       .def("destroy", &SimulateWrapper::Destroy)
@@ -225,8 +246,12 @@ PYBIND11_MODULE(_simulate, pymodule) {
       .def("lock", GetIfNotNull(&mujoco::Simulate::mtx),
            py::call_guard<py::gil_scoped_release>(),
            py::return_value_policy::reference_internal)
+      .def("set_figures", &SimulateWrapper::SetFigures,
+           py::arg("viewports_figures"))
+      .def("clear_figures", &SimulateWrapper::ClearFigures)
       .def_property_readonly("m", &SimulateWrapper::GetModel)
       .def_property_readonly("d", &SimulateWrapper::GetData)
+      .def_property_readonly("viewport", &SimulateWrapper::GetViewport)
       .def_property_readonly("ctrl_noise_std",
                              GetIfNotNull(&mujoco::Simulate::ctrl_noise_std),
                              py::call_guard<py::gil_scoped_release>())
@@ -260,18 +285,17 @@ PYBIND11_MODULE(_simulate, pymodule) {
                                return sim.exitrequest.load();
                              }),
                              py::call_guard<py::gil_scoped_release>())
-      .def(
-          "exit",
-          [](SimulateWrapper& wrapper) {
-            mujoco::Simulate* sim = wrapper.simulate();
-            if (!sim) {
-              return;
-            }
+      .def("exit",
+           [](SimulateWrapper& wrapper) {
+             mujoco::Simulate* sim = wrapper.simulate();
+             if (!sim) {
+               return;
+             }
 
-            int value = 0;
-            sim->exitrequest.compare_exchange_strong(value, 1);
-            wrapper.WaitUntilExit();
-          })
+             int value = 0;
+             sim->exitrequest.compare_exchange_strong(value, 1);
+             wrapper.WaitUntilExit();
+           })
 
       .def_property_readonly("uiloadrequest",
                              CallIfNotNull(+[](mujoco::Simulate& sim) {

python/mujoco/viewer.py

@@ -108,6 +108,23 @@ class Handle:
       return sim.d
     return None
 
+  @property
+  def viewport(self):
+    sim = self._sim()
+    if sim is not None:
+      return sim.viewport
+    return None
+
+  def set_figures(self, viewports_figures):
+    sim = self._sim()
+    if sim is not None:
+      sim.set_figures(viewports_figures)
+
+  def clear_figures(self):
+    sim = self._sim()
+    if sim is not None:
+      sim.clear_figures()
+
   def close(self):
     sim = self._sim()
     if sim is not None: