Add ImGui context sharing for UI rendering in Python MuJoCo Studio

PiperOrigin-RevId: 930507935
Change-Id: I7615931f9a71793b049e8b6b442be6a2d726f0b2

python/mujoco/experimental/studio/native_viewer.cc

@@ -22,6 +22,7 @@
 #include <vector>
 
 #include <fstream>
+#include <imgui.h>
 #include <implot.h>
 #include <mujoco/mujoco.h>
 #include <mujoco/experimental/platform/hal/graphics_mode.h>
@@ -180,6 +181,10 @@ class Viewer {
     window_->Present(pixels_);
   }
 
+  intptr_t GetImGuiContext() {
+    return reinterpret_cast<intptr_t>(ImGui::GetCurrentContext());
+  }
+
  private:
   std::unique_ptr<mujoco::platform::Window> window_;
   std::unique_ptr<mujoco::platform::Renderer> renderer_;
@@ -195,7 +200,8 @@ PYBIND11_MODULE(native_viewer_cc, m) {
       .def("Present", &Viewer::Present)
       .def("UploadImage", &Viewer::UploadImage)
       .def("RenderToTexture", &Viewer::RenderToTexture)
-      .def("GetDropFile", &Viewer::GetDropFile);
+      .def("GetDropFile", &Viewer::GetDropFile)
+      .def("GetImGuiContext", &Viewer::GetImGuiContext);
   m.def("IsCrd", &IsCrd);
   m.def("IsCuda", &IsCuda);
 }

python/mujoco/experimental/studio/native_viewer.py

@@ -24,6 +24,7 @@ import mujoco
 from mujoco.experimental.studio import native_viewer_cc as _viewer
 from mujoco.experimental.studio import ux
 from mujoco.experimental.studio import viewer_protocol
+from mujoco.experimental.dear_imgui import dear_imgui as imgui
 
 
 class NativeViewer(viewer_protocol.Viewer):
@@ -75,6 +76,10 @@ class NativeViewer(viewer_protocol.Viewer):
       # Initted to match mujoco/src/engine/engine_vis_init.c
       self.render_flags.flags = [1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1]
 
+    ctx = self._viewer.GetImGuiContext()
+    imgui.SetCurrentContext(ctx)
+    ux.set_imgui_context(ctx)
+
   def _sync_renderer(self, model: mujoco.MjModel) -> None:
     """Re-initializes the renderer if the model object has changed."""
     if id(model) != self._renderer_model_id:

python/mujoco/experimental/studio/ux.cc

@@ -95,6 +95,13 @@ PYBIND11_MODULE(ux, m) {
       },
       py::arg("theme"), "Set up Dear ImGui visual theme.");
 
+  m.def(
+      "set_imgui_context",
+      [](intptr_t ptr) {
+        ImGui::SetCurrentContext(reinterpret_cast<ImGuiContext*>(ptr));
+      },
+      py::arg("ptr"), "Set ImGui context pointer.");
+
   m.def(
       "configure_docking_layout",
       []() {