From 41ac783d3af361cffac97123755334633b4e4409 Mon Sep 17 00:00:00 2001
From: Michael Moss <mmoss@google.com>
Date: Thu, 11 Jun 2026 07:09:44 -0700
Subject: [PATCH] Add ImGui context sharing for UI rendering in Python MuJoCo
 Studio

PiperOrigin-RevId: 930507935
Change-Id: I7615931f9a71793b049e8b6b442be6a2d726f0b2
---
 python/mujoco/experimental/studio/native_viewer.cc | 8 +++++++-
 python/mujoco/experimental/studio/native_viewer.py | 5 +++++
 python/mujoco/experimental/studio/ux.cc            | 7 +++++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/python/mujoco/experimental/studio/native_viewer.cc b/python/mujoco/experimental/studio/native_viewer.cc
index 54330338..a93666e3 100644
--- a/python/mujoco/experimental/studio/native_viewer.cc
+++ b/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);
 }
diff --git a/python/mujoco/experimental/studio/native_viewer.py b/python/mujoco/experimental/studio/native_viewer.py
index 2628e16c..ab884015 100644
--- a/python/mujoco/experimental/studio/native_viewer.py
+++ b/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:
diff --git a/python/mujoco/experimental/studio/ux.cc b/python/mujoco/experimental/studio/ux.cc
index f5bb1985..da341b8f 100644
--- a/python/mujoco/experimental/studio/ux.cc
+++ b/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",
       []() {