From e0224c6440cc4ab41af0656d90a44c7f7e986c34 Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Tue, 21 Jul 2026 23:40:32 -0700 Subject: [PATCH] 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 | 14 ++++++++ .../experimental/studio/native_viewer.cc | 34 +++++++++++-------- .../experimental/studio/native_viewer.py | 4 +++ python/mujoco/experimental/studio/renderer.cc | 9 ++--- python/mujoco/experimental/studio/ux.cc | 12 +++++++ 5 files changed, 54 insertions(+), 19 deletions(-) diff --git a/python/mujoco/experimental/implot/implot.cc b/python/mujoco/experimental/implot/implot.cc index 7f287f32..ed7eca3c 100644 --- a/python/mujoco/experimental/implot/implot.cc +++ b/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(ptr)); + }, + py::arg("ptr"), "Set ImGui context pointer."); + m.def( + "set_implot_context", + [](intptr_t ptr) { + ImPlot::SetCurrentContext(reinterpret_cast(ptr)); + }, + py::arg("ptr"), "Set ImPlot context pointer."); + // Types. py::class_(m, "Point") .def(py::init<>()) diff --git a/python/mujoco/experimental/studio/native_viewer.cc b/python/mujoco/experimental/studio/native_viewer.cc index 1436e767..77ef189b 100644 --- a/python/mujoco/experimental/studio/native_viewer.cc +++ b/python/mujoco/experimental/studio/native_viewer.cc @@ -22,7 +22,6 @@ #include #include - #include #include #include @@ -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& render_flags, - const std::vector& 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& render_flags, + const std::vector& extra_geoms = {}) { std::vector 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(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(ImPlot::GetCurrentContext()); + } + private: std::unique_ptr window_; std::unique_ptr 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); } diff --git a/python/mujoco/experimental/studio/native_viewer.py b/python/mujoco/experimental/studio/native_viewer.py index eefa3b43..dceaba8f 100644 --- a/python/mujoco/experimental/studio/native_viewer.py +++ b/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)) diff --git a/python/mujoco/experimental/studio/renderer.cc b/python/mujoco/experimental/studio/renderer.cc index 05b7c211..ba3c98a6 100644 --- a/python/mujoco/experimental/studio/renderer.cc +++ b/python/mujoco/experimental/studio/renderer.cc @@ -55,10 +55,11 @@ class Renderer { std::vector 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()); } diff --git a/python/mujoco/experimental/studio/ux.cc b/python/mujoco/experimental/studio/ux.cc index 08d8750d..c6e30516 100644 --- a/python/mujoco/experimental/studio/ux.cc +++ b/python/mujoco/experimental/studio/ux.cc @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -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(ptr)); + }, + py::arg("ptr"), "Set ImPlot context pointer."); + m.def( "configure_docking_layout", []() {