From d7e872b6f9f812664489002946588557311596af Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Sat, 30 May 2026 03:40:59 -0700
Subject: [PATCH] Add UI control for MuJoCo threadpool size.

PiperOrigin-RevId: 923860258
Change-Id: I11dc5d7904f1c3657d7e1dd2b0e3c7864c0c7ead
---
 src/experimental/studio/app.cc | 18 ++++++++++++++++++
 src/experimental/studio/app.h  |  2 ++
 2 files changed, 20 insertions(+)

diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc
index d24df9a8..3c6a64f2 100644
--- a/src/experimental/studio/app.cc
+++ b/src/experimental/studio/app.cc
@@ -28,6 +28,7 @@
 #include <span>
 #include <string>
 #include <string_view>
+#include <thread>
 #include <utility>
 #include <vector>
 
@@ -227,6 +228,7 @@ void App::OnModelLoaded(std::string filename, ModelKind model_kind) {
       plugin->post_model_loaded(plugin, model_path_.c_str());
     }
   });
+  tmp_.update_threadpool = true;
 }
 
 void App::UpdateFilePaths(const std::string& resolved_path) {
@@ -268,6 +270,11 @@ void App::UpdatePhysics() {
     return;
   }
 
+  if (tmp_.update_threadpool) {
+    mju_threadpool(data(), ui_.nthread);
+    tmp_.update_threadpool = false;
+  }
+
   bool stepped = false;
   platform::ForEachPlugin<platform::ModelPlugin>([&](auto* plugin) {
     if (plugin->do_update) {
@@ -1457,6 +1464,15 @@ void App::ToolBarGui() {
     ImGui::SameLine(0, separator_width);
     platform::StepControlGui(model(), &step_control_, tmp_.speed_index);
 
+    ImGui::SameLine(0, separator_width);
+    ImGui::SetNextItemWidth(120);
+    ImGui::BeginDisabled(std::thread::hardware_concurrency() <= 1);
+    if (ImGui::SliderInt("##NumThreads", &ui_.nthread, 0, 8, "%d threads")) {
+      tmp_.update_threadpool = true;
+    }
+    ImGui::EndDisabled();
+    ImGui::SetItemTooltip("%s", "Number of threads in threadpool");
+
     ImGui::TableNextColumn();
 
     platform::CameraSelectionGui(model(), data(), camera_, ui_.camera_idx);
@@ -1823,6 +1839,7 @@ App::UiState::Dict App::UiState::ToDict() const {
       {"font_scale", std::to_string(font_scale)},
       {"window_width", std::to_string(window_width)},
       {"window_height", std::to_string(window_height)},
+      {"nthread", std::to_string(nthread)},
   };
 }
 
@@ -1834,5 +1851,6 @@ void App::UiState::FromDict(const Dict& dict) {
   window_width = ReadIniValue(dict, "window_width", window_width);
   window_height = ReadIniValue(dict, "window_height", window_height);
   font_scale = ReadIniValue(dict, "font_scale", font_scale);
+  nthread = ReadIniValue(dict, "nthread", nthread);
 }
 }  // namespace mujoco::studio
diff --git a/src/experimental/studio/app.h b/src/experimental/studio/app.h
index 9819c059..fa678ad2 100644
--- a/src/experimental/studio/app.h
+++ b/src/experimental/studio/app.h
@@ -114,6 +114,7 @@ class App {
     float font_scale = 1.0f;
     int window_width = 0;
     int window_height = 0;
+    int nthread = 0;
 
     using Dict = std::unordered_map<std::string, std::string>;
     Dict ToDict() const;
@@ -124,6 +125,7 @@ class App {
   struct UiTempState {
     bool should_exit = false;
     bool first_frame = true;
+    bool update_threadpool = false;
 
     // Windows.
     bool help = false;