studio: share ImGui and ImPlot contexts across python child extension modules.

Each pybind extension module holds its own copy of the ImGui/ImPlot globals (extension modules are loaded RTLD_LOCAL), so the viewer that owns the contexts must share the pointers explicitly. Without this, client plotting GUIs and user plotting scripts (e.g. the implot sample) crash on null context pointers.

PiperOrigin-RevId: 951917287
Change-Id: I639a173c2ae9f1af4feaf57aff7b5c81854b7393

python/mujoco/experimental/implot/implot.cc

@@ -33,6 +33,20 @@ PYBIND11_MODULE(implot, m) {
   // Import dear_imgui to make types like ImVec2 available.
   py::module_::import("mujoco.experimental.dear_imgui.dear_imgui");
 
+  // Each pybind module holds its own ImGui/ImPlot context globals
+  m.def(
+      "set_imgui_context",
+      [](intptr_t ptr) {
+        ImGui::SetCurrentContext(reinterpret_cast<ImGuiContext*>(ptr));
+      },
+      py::arg("ptr"), "Set ImGui context pointer.");
+  m.def(
+      "set_implot_context",
+      [](intptr_t ptr) {
+        ImPlot::SetCurrentContext(reinterpret_cast<ImPlotContext*>(ptr));
+      },
+      py::arg("ptr"), "Set ImPlot context pointer.");
+
   // Types.
   py::class_<ImPlotPoint>(m, "Point")
       .def(py::init<>())

python/mujoco/experimental/studio/native_viewer.cc

@@ -22,7 +22,6 @@
 #include <string_view>
 #include <vector>
 
-
 #include <fstream>
 #include <imgui.h>
 #include <implot.h>
@@ -155,18 +154,16 @@ class Viewer {
                                   bytes_per_pixel);
   }
 
-  std::string GetDropFile() {
-    return window_->GetDropFile();
-  }
+  std::string GetDropFile() { return window_->GetDropFile(); }
 
-  void Present(const mujoco::python::MjModelWrapper& model,
-               mujoco::python::MjDataWrapper& data,
-               mujoco::python::MjvPerturbWrapper& perturb,
-               mujoco::python::MjvCameraWrapper& camera,
-               mujoco::python::MjvOptionWrapper& vis_options,
-               const std::vector<uint8_t>& render_flags,
-               const std::vector<mujoco::python::MjvGeomWrapper>& extra_geoms =
-                   {}) {
+  void Present(
+      const mujoco::python::MjModelWrapper& model,
+      mujoco::python::MjDataWrapper& data,
+      mujoco::python::MjvPerturbWrapper& perturb,
+      mujoco::python::MjvCameraWrapper& camera,
+      mujoco::python::MjvOptionWrapper& vis_options,
+      const std::vector<uint8_t>& render_flags,
+      const std::vector<mujoco::python::MjvGeomWrapper>& extra_geoms = {}) {
     std::vector<mjvGeom> geoms;
     geoms.reserve(extra_geoms.size());
     for (const auto& geom_wrapper : extra_geoms) {
@@ -194,8 +191,8 @@ class Viewer {
     }
 
     renderer_->Render(model.get(), data.get(), perturb.get(), camera.get(),
-                      vis_options.get(), width * scale, height * scale,
-                      pixels_, geoms);
+                      vis_options.get(), width * scale, height * scale, pixels_,
+                      geoms);
 
     window_->EndFrame();
     window_->Present(pixels_);
@@ -205,6 +202,12 @@ class Viewer {
     return reinterpret_cast<intptr_t>(ImGui::GetCurrentContext());
   }
 
+  // See ux.set_implot_context: extension modules each hold their own copy of
+  // the ImPlot globals, so the context pointer must be shared explicitly.
+  intptr_t GetImPlotContext() {
+    return reinterpret_cast<intptr_t>(ImPlot::GetCurrentContext());
+  }
+
  private:
   std::unique_ptr<mujoco::platform::Window> window_;
   std::unique_ptr<mujoco::platform::Renderer> renderer_;
@@ -226,7 +229,8 @@ PYBIND11_MODULE(native_viewer_cc, m, pybind11::mod_gil_not_used()) {
       .def("UploadImage", &Viewer::UploadImage)
       .def("RenderToTexture", &Viewer::RenderToTexture)
       .def("GetDropFile", &Viewer::GetDropFile)
-      .def("GetImGuiContext", &Viewer::GetImGuiContext);
+      .def("GetImGuiContext", &Viewer::GetImGuiContext)
+      .def("GetImPlotContext", &Viewer::GetImPlotContext);
   m.def("IsCrd", &IsCrd);
   m.def("IsCuda", &IsCuda);
 }

python/mujoco/experimental/studio/native_viewer.py

@@ -27,6 +27,7 @@ from mujoco.experimental.studio import ux
 from mujoco.experimental.studio import viewer_protocol
 
 from mujoco.experimental.dear_imgui import dear_imgui as imgui
+from mujoco.experimental.implot import implot
 
 
 class NativeViewer(viewer_protocol.Viewer):
@@ -85,6 +86,9 @@ class NativeViewer(viewer_protocol.Viewer):
     ctx = self._viewer.GetImGuiContext()
     imgui.SetCurrentContext(ctx)
     ux.set_imgui_context(ctx)
+    ux.set_implot_context(self._viewer.GetImPlotContext())
+    implot.set_imgui_context(ctx)
+    implot.set_implot_context(self._viewer.GetImPlotContext())
 
     # Dispatch lifecycle event so handlers can cache the viewer reference.
     self.dispatch(viewer_protocol.ViewerInitEvent(viewer=self))

python/mujoco/experimental/studio/renderer.cc

@@ -55,10 +55,11 @@ class Renderer {
     std::vector<std::byte> pixels(width * height * 3);
     {
       py::gil_scoped_release no_gil;
-      impl_->Render(
-          model.get(), data.get(), perturb ? perturb.value().get() : nullptr,
-          camera ? camera.value().get() : nullptr,
-          vis_option ? vis_option.value().get() : nullptr, width, height, pixels);
+      impl_->Render(model.get(), data.get(),
+                    perturb ? perturb.value().get() : nullptr,
+                    camera ? camera.value().get() : nullptr,
+                    vis_option ? vis_option.value().get() : nullptr, width,
+                    height, pixels);
     }
     return pybind11::bytes((const char*)pixels.data(), pixels.size());
   }

python/mujoco/experimental/studio/ux.cc

@@ -20,6 +20,7 @@
 #include <vector>
 
 #include <imgui.h>
+#include <implot.h>
 #include <mujoco/mujoco.h>
 #include <mujoco/experimental/platform/helpers.h>
 #include <mujoco/experimental/platform/sim/step_control.h>
@@ -96,6 +97,17 @@ PYBIND11_MODULE(ux, m, pybind11::mod_gil_not_used()) {
       },
       py::arg("ptr"), "Set ImGui context pointer.");
 
+  // Each pybind extension module holds its own copy of the ImGui/ImPlot
+  // globals (extension modules are loaded RTLD_LOCAL), so the viewer that
+  // owns the contexts must share the pointers with this module. Without
+  // this, the plotting GUIs would dereference a null ImPlot context.
+  m.def(
+      "set_implot_context",
+      [](intptr_t ptr) {
+        ImPlot::SetCurrentContext(reinterpret_cast<ImPlotContext*>(ptr));
+      },
+      py::arg("ptr"), "Set ImPlot context pointer.");
+
   m.def(
       "configure_docking_layout",
       []() {