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/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",
       []() {