Release GIL in MuJoCo Studio Python bindings.

This change adds py::gil_scoped_release to various functions in the parser, renderer, sim and ux modules of the MuJoCo Studio Python bindings. This allows other Python threads to execute while these C++ functions, which can sometimes be time-consuming, are running. The two heaviest operations Present() (rendering) and Advance() (physics) now both release the GIL in C++.

PiperOrigin-RevId: 930645361
Change-Id: Ib239fb54edc7ac0d8b1d94bec8c57a447d5f4d27

python/mujoco/experimental/studio/parser.cc

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <memory>
 #include <string_view>
 
 #include <mujoco/mujoco.h>
@@ -24,7 +25,11 @@ namespace mujoco::python {
 // Loads, parses, and compiles a MuJoCo model from the given file. Returns the
 // python mjData object (which also contains the compiled mjModel).
 py::object Parse(std::string_view filepath) {
-  auto holder = platform::ModelHolder::FromFile(filepath);
+  std::unique_ptr<platform::ModelHolder> holder;
+  {
+    py::gil_scoped_release no_gil;
+    holder = platform::ModelHolder::FromFile(filepath);
+  }
   if (!holder->ok()) {
     throw py::value_error(
         std::string("Failed to load model from '") +

python/mujoco/experimental/studio/renderer.cc

@@ -36,12 +36,16 @@ class Renderer {
   using GraphicsMode = mujoco::platform::GraphicsMode;
 
   Renderer(const std::string& graphics_mode_str) {
+    py::gil_scoped_release no_gil;
     const GraphicsMode mode = mujoco::platform::GraphicsModeFromString(
         graphics_mode_str, GraphicsMode::FilamentOpenGl);
     impl_ = std::make_unique<RendererImpl>(nullptr, mode);
   }
 
-  void Init(const MjModelWrapper& model) { impl_->Init(model.get()); }
+  void Init(const MjModelWrapper& model) {
+    py::gil_scoped_release no_gil;
+    impl_->Init(model.get());
+  }
 
   pybind11::bytes Render(const MjModelWrapper& model, MjDataWrapper& data,
                          std::optional<MjvPerturbWrapper>& perturb,
@@ -49,10 +53,13 @@ class Renderer {
                          std::optional<MjvOptionWrapper>& vis_option, int width,
                          int height) {
     std::vector<std::byte> pixels(width * height * 3);
-    impl_->Render(
+    {
-        model.get(), data.get(), perturb ? perturb.value().get() : nullptr,
+      py::gil_scoped_release no_gil;
-        camera ? camera.value().get() : nullptr,
+      impl_->Render(
-        vis_option ? vis_option.value().get() : nullptr, width, height, pixels);
+          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/sim.cc

@@ -48,6 +48,7 @@ PYBIND11_MODULE(sim, m) {
             auto& model = py::cast<mujoco::python::MjModelWrapper&>(model_obj);
             auto& data = py::cast<mujoco::python::MjDataWrapper&>(data_obj);
             if (step_fn.is_none()) {
+              py::gil_scoped_release no_gil;
               return self.Advance(model.get(), data.get());
             } else {
               return self.Advance(

python/mujoco/experimental/studio/ux.cc

@@ -91,6 +91,7 @@ PYBIND11_MODULE(ux, m) {
   m.def(
       "setup_theme",
       [](mujoco::platform::GuiTheme theme) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::SetupTheme(theme);
       },
       py::arg("theme"), "Set up Dear ImGui visual theme.");
@@ -105,6 +106,7 @@ PYBIND11_MODULE(ux, m) {
   m.def(
       "configure_docking_layout",
       []() {
+        py::gil_scoped_release no_gil;
         ImVec4 r = mujoco::platform::ConfigureDockingLayout();
         return std::make_tuple(r.x, r.y, r.z, r.w);
       },
@@ -115,6 +117,7 @@ PYBIND11_MODULE(ux, m) {
       "step_control_gui",
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::platform::StepControl* step_control, UxState& ux_state) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::StepControlGui(model.get(), step_control,
                                          ux_state.speed_index);
       },
@@ -125,6 +128,7 @@ PYBIND11_MODULE(ux, m) {
   m.def(
       "theme_select_gui",
       [](mujoco::platform::GuiTheme theme) {
+        py::gil_scoped_release no_gil;
         bool changed = mujoco::platform::ThemeSelectGui(&theme);
         return std::make_tuple(changed, theme);
       },
@@ -134,6 +138,7 @@ PYBIND11_MODULE(ux, m) {
   m.def(
       "label_selection_gui",
       [](mujoco::python::MjvOptionWrapper& vis_options) {
+        py::gil_scoped_release no_gil;
         return mujoco::platform::LabelSelectionGui(vis_options.get());
       },
       py::arg("vis_options"), "Render the visualization label selection GUI.");
@@ -141,6 +146,7 @@ PYBIND11_MODULE(ux, m) {
   m.def(
       "frame_selection_gui",
       [](mujoco::python::MjvOptionWrapper& vis_options) {
+        py::gil_scoped_release no_gil;
         return mujoco::platform::FrameSelectionGui(vis_options.get());
       },
       py::arg("vis_options"), "Render the visualization frame selection GUI.");
@@ -150,6 +156,7 @@ PYBIND11_MODULE(ux, m) {
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data,
          mujoco::python::MjvCameraWrapper& camera, UxState& ux_state) {
+        py::gil_scoped_release no_gil;
         bool changed = mujoco::platform::CameraSelectionGui(
             model.get(), data.get(), *camera.get(), ux_state.camera_index);
         return changed;
@@ -162,6 +169,7 @@ PYBIND11_MODULE(ux, m) {
       "set_camera",
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjvCameraWrapper& camera, int request_idx) {
+        py::gil_scoped_release no_gil;
         return mujoco::platform::SetCamera(model.get(), camera.get(), request_idx);
       },
       py::arg("model"), py::arg("camera"), py::arg("request_idx"),
@@ -170,6 +178,7 @@ PYBIND11_MODULE(ux, m) {
   m.def(
       "set_speed_index",
       [](mujoco::platform::StepControl* step_control, UxState& ux_state, int idx) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::SetSpeedIndex(step_control, ux_state.speed_index, idx);
       },
       py::arg("step_control"), py::arg("ux_state"), py::arg("idx"),
@@ -178,6 +187,7 @@ PYBIND11_MODULE(ux, m) {
   m.def(
       "physics_gui",
       [](mujoco::python::MjModelWrapper& model, float min_width) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::PhysicsGui(model.get(), min_width);
       },
       py::arg("model"), py::arg("min_width") = 150.0f,
@@ -188,6 +198,7 @@ PYBIND11_MODULE(ux, m) {
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjvOptionWrapper& vis_options,
          RenderFlags& render_flags) {
+        py::gil_scoped_release no_gil;
         mjtByte flags[mjNRNDFLAG] = {0};
         for (int i = 0; i < mjNRNDFLAG; ++i) {
           flags[i] = render_flags.flags[i];
@@ -205,6 +216,7 @@ PYBIND11_MODULE(ux, m) {
       "groups_gui",
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjvOptionWrapper& vis_options, float min_width) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::GroupsGui(model.get(), vis_options.get(), min_width);
       },
       py::arg("model"), py::arg("vis_options"), py::arg("min_width") = 150.0f,
@@ -215,6 +227,7 @@ PYBIND11_MODULE(ux, m) {
       [](mujoco::python::MjModelWrapper& model,
          mujoco::python::MjvOptionWrapper& vis_options,
          mujoco::python::MjvCameraWrapper& camera, float min_width) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::VisualizationGui(model.get(), vis_options.get(),
                                            camera.get(), min_width);
       },
@@ -226,6 +239,7 @@ PYBIND11_MODULE(ux, m) {
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data,
          mujoco::python::MjvOptionWrapper& vis_options) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::ControlsGui(model.get(), data.get(),
                                       vis_options.get());
       },
@@ -237,6 +251,7 @@ PYBIND11_MODULE(ux, m) {
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data,
          mujoco::python::MjvOptionWrapper& vis_options) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::JointsGui(model.get(), data.get(), vis_options.get());
       },
       py::arg("model"), py::arg("data"), py::arg("vis_options"),
@@ -246,6 +261,7 @@ PYBIND11_MODULE(ux, m) {
       "sensor_gui",
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::SensorGui(model.get(), data.get());
       },
       py::arg("model"), py::arg("data"), "Render the sensor data plot.");
@@ -255,6 +271,7 @@ PYBIND11_MODULE(ux, m) {
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data, UxState& ux_state,
          float min_width) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::StateGui(model.get(), data.get(), ux_state.state,
                                    ux_state.state_sig, min_width);
       },
@@ -266,6 +283,7 @@ PYBIND11_MODULE(ux, m) {
       "watch_gui",
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data, UxState& ux_state) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::WatchGui(
             model.get(), data.get(), ux_state.watch_field_name,
             sizeof(ux_state.watch_field_name), ux_state.watch_field_index);
@@ -278,6 +296,7 @@ PYBIND11_MODULE(ux, m) {
       "noise_gui",
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data, UxState& ux_state) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::NoiseGui(model.get(), data.get(),
                                    ux_state.noise_scale, ux_state.noise_rate);
       },
@@ -289,6 +308,7 @@ PYBIND11_MODULE(ux, m) {
       "convergence_gui",
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::ConvergenceGui(model.get(), data.get());
       },
       py::arg("model"), py::arg("data"),
@@ -298,6 +318,7 @@ PYBIND11_MODULE(ux, m) {
       "counts_gui",
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::CountsGui(model.get(), data.get());
       },
       py::arg("model"), py::arg("data"), "Render the solver counts chart.");
@@ -306,6 +327,7 @@ PYBIND11_MODULE(ux, m) {
       "stats_gui",
       [](const mujoco::python::MjModelWrapper& model,
          mujoco::python::MjDataWrapper& data, bool paused, float fps) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::StatsGui(model.get(), data.get(), paused, fps);
       },
       py::arg("model"), py::arg("data"), py::arg("paused"), py::arg("fps"),
@@ -331,6 +353,7 @@ PYBIND11_MODULE(ux, m) {
          const mujoco::python::MjDataWrapper& data,
          mujoco::python::MjvCameraWrapper& cam,
          mujoco::platform::CameraMotion motion, mjtNum dx, mjtNum dy) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::MoveCamera(model.get(), data.get(), cam.get(), motion,
                                      dx, dy);
       },
@@ -343,6 +366,7 @@ PYBIND11_MODULE(ux, m) {
          const mujoco::python::MjDataWrapper& data,
          const mujoco::python::MjvCameraWrapper& cam,
          mujoco::python::MjvPerturbWrapper& pert, int active) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::InitPerturb(model.get(), data.get(), cam.get(),
                                       pert.get(),
                                       static_cast<mjtPertBit>(active));
@@ -357,6 +381,7 @@ PYBIND11_MODULE(ux, m) {
          const mujoco::python::MjvCameraWrapper& cam,
          mujoco::python::MjvPerturbWrapper& pert, int action, mjtNum reldx,
          mjtNum reldy) {
+        py::gil_scoped_release no_gil;
         mujoco::platform::MovePerturb(model.get(), data.get(), cam.get(),
                                       pert.get(), static_cast<mjtMouse>(action),
                                       reldx, reldy);
@@ -388,6 +413,7 @@ PYBIND11_MODULE(ux, m) {
          const mujoco::python::MjDataWrapper& data,
          const mujoco::python::MjvCameraWrapper& cam, float x, float y,
          float aspect_ratio, const mujoco::python::MjvOptionWrapper& opt) {
+        py::gil_scoped_release no_gil;
         return mujoco::platform::Pick(model.get(), data.get(), cam.get(), x, y,
                                       aspect_ratio, opt.get());
       },
@@ -399,6 +425,7 @@ PYBIND11_MODULE(ux, m) {
       "camera_to_string",
       [](const mujoco::python::MjDataWrapper& data,
          const mujoco::python::MjvCameraWrapper& camera) {
+        py::gil_scoped_release no_gil;
         return mujoco::platform::CameraToString(data.get(), camera.get());
       },
       py::arg("data"), py::arg("camera"),