Add headless EGL/OSMesa OpenGL UI context rendering for web viewer

PiperOrigin-RevId: 957060514
Change-Id: I8a0df22255c5a938d18e172c638a8c83dfe64b1b

python/mujoco/experimental/studio/web/headless_ui.cc

@@ -0,0 +1,302 @@
+// Copyright 2026 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// The web viewer's headless Studio UI.
+//
+// This pybind11 module provides a HeadlessUi class that:
+// 1. Creates a headless ImGui context (no window, no renderer).
+// 2. Connects as a netimgui client, streaming ImGui draw data to a remote
+//    viewer (the browser's web_client, bridged through web_server.py).
+// 3. Receives input events from the remote viewer and injects them into
+//    the ImGui context.
+
+#include <chrono>
+#include <cstdint>
+#include <string>
+#include <string_view>
+#include <thread>
+#include <vector>
+
+#include <imgui.h>
+#include <implot.h>
+#include "experimental/platform/ux/fonts.h"
+#include <NetImgui_Api.h>
+#include "google/logging.h"
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+// Headless ImGui + netimgui viewer for the MuJoCo web viewer.
+//
+// This class manages a headless ImGui context that streams its draw data
+// via the netimgui protocol to a remote viewer. Unlike the NativeViewer's
+// Viewer class (native_viewer.cc), this does NOT create a window, initialize
+// a renderer, or handle mouse/keyboard input directly. All rendering and
+// input handling happens on the remote client (web_client.cc in the browser).
+//
+// The lifecycle follows the SampleNoBackend pattern:
+//   Client_Startup()   — create context, load fonts, init NetImgui
+//   Client_Connect()   — manage connection state (called each frame)
+//   Client_Shutdown()  — release resources
+
+// Initialize the Dear ImGui Context and the NetImgui library.
+// Based on SampleNoBackend::Client_Startup() from
+// netimgui/Code/Sample/SampleNoBackend/SampleNoBackend.cpp
+static bool Client_Startup(ImGuiContext*& context,
+                           const std::string& assets_dir) {
+  IMGUI_CHECKVERSION();
+  context = ImGui::CreateContext();
+  ImGui::SetCurrentContext(context);
+  ImPlot::CreateContext();
+
+  ImGuiIO& io = ImGui::GetIO();
+  io.ConfigFlags |= ImGuiConfigFlags_DockingEnable;
+  io.BackendFlags |= ImGuiBackendFlags_HasGamepad;
+  io.IniFilename = nullptr;
+  io.ConfigDpiScaleFonts = true;
+  io.ConfigDpiScaleViewports = true;
+  io.DisplaySize = ImVec2(1400, 720);
+
+  // Disable ImGui auto-repeat in headless mode since the browser handles key
+  // repeat at the OS level, avoiding the frame amplification issue where many
+  // frames with identical key-down states trigger unwanted repeats.
+  io.KeyRepeatDelay = 9999.0f;
+
+  // Initialize the main viewport's DPI scale. Without this, the headless
+  // context leaves DpiScale at 0.0f (no platform backend sets it), which
+  // triggers an assertion in ImGui::SetCurrentViewport() when
+  // BeginMainMenuBar() is called.
+  ImGuiViewport* main_vp = ImGui::GetMainViewport();
+  if (main_vp) {
+    main_vp->DpiScale = 1.0f;
+  }
+
+  ImGui::StyleColorsLight();
+
+  // The Studio font set shared with the native viewer and the browser client.
+  mujoco::platform::AddStudioFonts([&assets_dir](std::string_view filename) {
+    return mujoco::platform::LoadFontAsset(assets_dir, filename);
+  });
+
+  if (!NetImgui::Startup()) {
+    LOG(Error, "NetImgui::Startup() failed");
+    return false;
+  }
+
+  return true;
+}
+
+// Release resources.
+// Based on SampleNoBackend::Client_Shutdown() from
+// netimgui/Code/Sample/SampleNoBackend/SampleNoBackend.cpp
+static void Client_Shutdown(ImGuiContext*& context) {
+  NetImgui::Shutdown();
+  ImPlot::DestroyContext();
+  if (context) {
+    ImGui::DestroyContext(context);
+    context = nullptr;
+  }
+}
+
+// Manage connection to the netimgui proxy.
+// Based on SampleNoBackend::Client_Connect() from
+// netimgui/Code/Sample/SampleNoBackend/SampleNoBackend.cpp
+static void Client_Connect(const char* title, int port) {
+  bool connected = NetImgui::IsConnected();
+  bool pending = NetImgui::IsConnectionPending();
+
+  if (!connected && !pending) {
+    static std::chrono::steady_clock::time_point last_reconnect_time =
+        std::chrono::steady_clock::now();
+    const std::chrono::steady_clock::time_point now =
+        std::chrono::steady_clock::now();
+    if (now - last_reconnect_time > std::chrono::seconds(1)) {
+      last_reconnect_time = now;
+      VLOG(1, "Retrying ConnectToApp...");
+      NetImgui::ConnectToApp(title, "127.0.0.1", port);
+    }
+  }
+
+  static bool last_connected = false;
+  if (connected != last_connected) {
+    VLOG(1, "Status change: Connected=%s", connected ? "true" : "false");
+    last_connected = connected;
+  }
+}
+
+class HeadlessUi {
+ public:
+  HeadlessUi(const std::string& title, int port, const std::string& assets_dir)
+      : title_(title), port_(port) {
+    if (!Client_Startup(context_, assets_dir)) {
+      return;
+    }
+
+    VLOG(1, "Calling ConnectToApp('%s', '127.0.0.1', %d)", title_.c_str(),
+         port_);
+    bool connect_result =
+        NetImgui::ConnectToApp(title_.c_str(), "127.0.0.1", port_);
+    VLOG(1, "ConnectToApp returned: %s", connect_result ? "true" : "false");
+    VLOG(1, "IsConnected: %s, IsConnectionPending: %s",
+         NetImgui::IsConnected() ? "true" : "false",
+         NetImgui::IsConnectionPending() ? "true" : "false");
+  }
+
+  ~HeadlessUi() { Client_Shutdown(context_); }
+
+  bool NewFrame() {
+    py::gil_scoped_release no_gil;
+    ImGui::SetCurrentContext(context_);
+    static int frame_count = 0;
+    frame_count++;
+
+    // Returns true with an active ImGui frame once a browser is connected and
+    // ready. Returns false (no frame) while no browser is connected so the
+    // caller's loop keeps running and can drain messages (e.g., to respond to
+    // an ExitEvent for shutdown).
+    std::chrono::steady_clock::time_point last_signal_check =
+        std::chrono::steady_clock::now();
+    while (true) {
+      // While no browser is connected, periodically check for
+      // Python signals so Ctrl+C interrupts the wait instead of hanging.
+      const std::chrono::steady_clock::time_point now =
+          std::chrono::steady_clock::now();
+      if (now - last_signal_check > std::chrono::milliseconds(200)) {
+        last_signal_check = now;
+        py::gil_scoped_acquire acquire;
+        if (PyErr_CheckSignals() != 0) {
+          throw py::error_already_set();
+        }
+      }
+
+      Client_Connect(title_.c_str(), port_);
+
+      if (!NetImgui::IsConnected()) {
+        // Not connected; return false so the caller's loop is not blocked.
+        is_drawing_remote_ = false;
+        std::this_thread::sleep_for(std::chrono::milliseconds(10));
+        return false;
+      }
+
+      bool new_frame_result = NetImgui::NewFrame(false);
+      is_drawing_remote_ = new_frame_result;
+      if (!new_frame_result) {
+        // Connected but NetImgui not ready for a draw — wait and retry.
+        std::this_thread::sleep_for(std::chrono::milliseconds(1));
+        continue;
+      }
+
+      break;
+    }
+
+    VLOG(1, "Frame %d: DrawingRemote=%s", frame_count,
+         is_drawing_remote_ ? "Y" : "N");
+    return true;
+  }
+
+  void EndFrame() {
+    py::gil_scoped_release no_gil;
+    ImGui::SetCurrentContext(context_);
+    static int end_frame_count = 0;
+    end_frame_count++;
+    if (!is_drawing_remote_) {
+      // No frame was started — nothing to end.
+      return;
+    }
+    VLOG(1, "EndFrame %d: sending remote draw data", end_frame_count);
+    NetImgui::EndFrame();
+  }
+
+  // Uploads an RGB/RGBA image to the browser over the NetImgui texture
+  // channel so handlers can display it with imgui.Image(). Returns the
+  // texture id to use (allocates one when tex_id == 0). This is the UI-link
+  // counterpart of the native renderer's UploadImage.
+  uintptr_t UploadImage(uintptr_t tex_id, const py::bytes& pixels, int width,
+                        int height, int bpp) {
+    if (tex_id == 0) {
+      tex_id = next_tex_id_++;
+    }
+    std::string data = pixels;
+    if (width <= 0 || height <= 0 ||
+        data.size() < static_cast<size_t>(width) * height * bpp) {
+      LOG(Error, "UploadImage: bad dimensions %dx%dx%d for %zu bytes", width,
+          height, bpp, data.size());
+      return tex_id;
+    }
+
+    // NetImgui transfers RGBA8; expand RGB if needed.
+    std::vector<uint8_t> rgba;
+    const void* upload_data = data.data();
+    if (bpp == 3) {
+      rgba.resize(static_cast<size_t>(width) * height * 4);
+      for (size_t p = 0; p < static_cast<size_t>(width) * height; ++p) {
+        rgba[p * 4 + 0] = data[p * 3 + 0];
+        rgba[p * 4 + 1] = data[p * 3 + 1];
+        rgba[p * 4 + 2] = data[p * 3 + 2];
+        rgba[p * 4 + 3] = 255;
+      }
+      upload_data = rgba.data();
+    } else if (bpp != 4) {
+      LOG(Error, "UploadImage: unsupported bpp %d (expected 3 or 4)", bpp);
+      return tex_id;
+    }
+
+    py::gil_scoped_release no_gil;
+    ImGui::SetCurrentContext(context_);
+    NetImgui::SendDataTexture(
+        static_cast<ImTextureID>(tex_id), const_cast<void*>(upload_data),
+        static_cast<uint16_t>(width), static_cast<uint16_t>(height),
+        NetImgui::eTexFormat::kTexFmtRGBA8);
+    return tex_id;
+  }
+
+  intptr_t GetContext() const { return reinterpret_cast<intptr_t>(context_); }
+
+  // The ImPlot context created alongside the ImGui context. Python must pass
+  // this to ux.set_implot_context: extension modules each hold their own copy
+  // of the ImPlot globals, so the context pointer has to be shared explicitly
+  // (same pattern as get_context/set_imgui_context).
+  intptr_t GetImPlotContext() const {
+    return reinterpret_cast<intptr_t>(ImPlot::GetCurrentContext());
+  }
+
+ private:
+  std::string title_;
+  int port_;
+  ImGuiContext* context_ = nullptr;
+  bool is_drawing_remote_ = false;
+  // User texture ids start well above the ids ImGui's managed texture system
+  // (font atlas) hands out, so the two can never collide in the browser's
+  // texture map.
+  uintptr_t next_tex_id_ = 0x10000;
+};
+
+PYBIND11_MODULE(headless_ui, m, pybind11::mod_gil_not_used()) {
+  m.doc() = "MuJoCo web viewer headless Studio UI, streamed via NetImgui";
+
+  py::class_<HeadlessUi>(m, "HeadlessUi")
+      .def(py::init<const std::string&, int, const std::string&>(),
+           py::arg("title"), py::arg("port") = 8888, py::arg("assets_dir") = "")
+      .def("new_frame", &HeadlessUi::NewFrame,
+           "Starts a headless ImGui frame, returning True once the frame is "
+           "active. When a browser is viewing the page (via the URL printed "
+           "at startup), this returns at the browser's requested frame rate. "
+           "When no browser is viewing, it returns False (no frame) after a "
+           "short wait, so the caller's loop keeps running and can shut down.")
+      .def("end_frame", &HeadlessUi::EndFrame)
+      .def("get_context", &HeadlessUi::GetContext)
+      .def("get_implot_context", &HeadlessUi::GetImPlotContext)
+      .def("upload_image", &HeadlessUi::UploadImage);
+}

src/experimental/platform/CMakeLists.txt

@@ -44,6 +44,8 @@ target_sources(${MUJOCO_PLATFORM_TARGET_NAME}
         sim/step_control.h
         ux/enum_utils.h
         ux/file_dialog.h
+        ux/fonts.cc
+        ux/fonts.h
         ux/gui.cc
         ux/gui.h
         ux/gui_spec.cc

src/experimental/platform/hal/window.cc

@@ -32,6 +32,7 @@
 #include <imgui.h>
 #include <mujoco/mujoco.h>
 #include "experimental/platform/hal/graphics_mode.h"
+#include "experimental/platform/ux/fonts.h"
 #include "user/user_resource.h"
 
 // Because X11/Xlib.h defines Status.
@@ -63,36 +64,7 @@ static void InitImGui(SDL_Window* window, float content_scale,
   style.FontScaleDpi = content_scale;
 
   if (load_fonts) {
-    mjResource* font = nullptr;
-    int size = 0;
-    void* data = nullptr;
-
-    ImFontConfig main_cfg;
-    main_cfg.FontDataOwnedByAtlas = false;
-    font =
-        mju_openResource("", "font:AtkinsonHyperlegibleNext[wght].ttf",
-                         nullptr, nullptr, 0);
-    size = mju_readResource(font, const_cast<const void**>(&data));
-    io.Fonts->AddFontFromMemoryTTF(data, size, 16.f, &main_cfg);
-
-    ImFontConfig icon_cfg;
-    icon_cfg.FontDataOwnedByAtlas = false;
-    icon_cfg.MergeMode = true;
-    font = mju_openResource("", "font:fontawesome-webfont.ttf", nullptr,
-                            nullptr, 0);
-    size = mju_readResource(font, const_cast<const void**>(&data));
-    constexpr ImWchar icon_ranges[] = {0xf000, 0xf3ff, 0x000};
-    io.Fonts->AddFontFromMemoryTTF(data, size, 13.f, &icon_cfg, icon_ranges);
-
-    ImFontConfig mono_cfg;
-    mono_cfg.FontDataOwnedByAtlas = false;
-    font = mju_openResource("", "font:AtkinsonHyperlegibleMono-Regular.ttf", nullptr,
-                            nullptr, 0);
-    size = mju_readResource(font, const_cast<const void**>(&data));
-    io.Fonts->AddFontFromMemoryTTF(data, size, 14.f, &mono_cfg);
-
-    // Note: we purposefully do not "close" the font resources as ImGui may
-    // need them again to resize fonts.
+    AddStudioFonts(LoadFontResource);
   }
 }
 

src/experimental/platform/ux/fonts.cc

@@ -0,0 +1,106 @@
+// Copyright 2026 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "experimental/platform/ux/fonts.h"
+
+#include <cstddef>
+#include <cstring>
+#include <filesystem>  // NOLINT(build/c++17)
+#include <fstream>
+#include <ios>
+#include <string>
+#include <string_view>
+#include <vector>
+
+#include <imgui.h>
+#include <mujoco/mujoco.h>
+
+namespace mujoco::platform {
+
+namespace {
+
+// Copies `data` into ImGui-owned memory and adds it to the current context's
+// atlas; the atlas frees the copy with the context.
+ImFont* AddFontCopy(const std::vector<std::byte>& data, float size_pixels,
+                    const ImFontConfig* config = nullptr,
+                    const ImWchar* glyph_ranges = nullptr) {
+  void* copy = ImGui::MemAlloc(data.size());
+  memcpy(copy, data.data(), data.size());
+  return ImGui::GetIO().Fonts->AddFontFromMemoryTTF(
+      copy, static_cast<int>(data.size()), size_pixels, config, glyph_ranges);
+}
+
+}  // namespace
+
+std::vector<std::byte> LoadFontAsset(const std::string& assets_dir,
+                                     std::string_view filename) {
+  if (assets_dir.empty()) return {};
+  const std::string file_path =
+      (std::filesystem::path(assets_dir) / filename).string();
+
+  std::ifstream file(file_path, std::ios::binary | std::ios::ate);
+  if (!file.is_open()) {
+    return {};
+  }
+  const std::streamsize file_size = file.tellg();
+  file.seekg(0, std::ios::beg);
+  std::vector<std::byte> buffer(file_size);
+  if (!file.read(reinterpret_cast<char*>(buffer.data()), file_size)) {
+    return {};
+  }
+  return buffer;
+}
+
+std::vector<std::byte> LoadFontResource(std::string_view filename) {
+  const std::string name = std::string("font:") + std::string(filename);
+  mjResource* resource =
+      mju_openResource("", name.c_str(), nullptr, nullptr, 0);
+  if (resource == nullptr) {
+    return {};
+  }
+  const void* data = nullptr;
+  const int size = mju_readResource(resource, &data);
+  std::vector<std::byte> buffer;
+  if (size > 0) {
+    const std::byte* bytes = static_cast<const std::byte*>(data);
+    buffer.assign(bytes, bytes + size);
+  }
+  // Resource can be closed because the font atlas owns a copy of the data.
+  mju_closeResource(resource);
+  return buffer;
+}
+
+void AddStudioFonts(const FontLoaderFn& load) {
+  const std::vector<std::byte> main_data = load(kMainFontFile);
+  if (!main_data.empty()) {
+    AddFontCopy(main_data, 16.0f);
+  }
+
+  // The icon font merges into the font added above; without a base font the
+  // merge has nothing to attach to, so skip it.
+  const std::vector<std::byte> icon_data = load(kIconFontFile);
+  if (!icon_data.empty() && !main_data.empty()) {
+    ImFontConfig icon_config;
+    icon_config.MergeMode = true;
+    static constexpr ImWchar kIconRanges[] = {0xf000, 0xf3ff, 0};
+    AddFontCopy(icon_data, 13.0f, &icon_config, kIconRanges);
+  }
+
+  const std::vector<std::byte> mono_data = load(kMonoFontFile);
+  if (!mono_data.empty()) {
+    AddFontCopy(mono_data, 14.0f);
+  }
+}
+
+}  // namespace mujoco::platform

src/experimental/platform/ux/fonts.h

@@ -0,0 +1,48 @@
+// Copyright 2026 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_EXPERIMENTAL_PLATFORM_UX_FONTS_H_
+#define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_UX_FONTS_H_
+
+#include <cstddef>
+#include <functional>
+#include <string>
+#include <string_view>
+#include <vector>
+
+namespace mujoco::platform {
+
+// The Studio UI font files (every viewer loads this same files).
+inline constexpr char kMainFontFile[] = "AtkinsonHyperlegibleNext[wght].ttf";
+inline constexpr char kIconFontFile[] = "fontawesome-webfont.ttf";
+inline constexpr char kMonoFontFile[] = "AtkinsonHyperlegibleMono-Regular.ttf";
+
+// Maps a Studio font filename to its TTF bytes; empty means unavailable.
+using FontLoaderFn = std::function<std::vector<std::byte>(std::string_view)>;
+
+// A FontLoaderFn that reads font `filename` from an `assets_dir` folder.
+std::vector<std::byte> LoadFontAsset(const std::string& assets_dir,
+                                     std::string_view filename);
+
+// A FontLoaderFn that reads font `font:<filename>` using a resource provider.
+std::vector<std::byte> LoadFontResource(std::string_view filename);
+
+// Adds the Studio fonts to the current ImGui context's atlas, fetching font
+// data via `load`. The ImGui font atlas will own a copy of the data so it can
+// rebuild at new DPI scales without buffers used by `load` staying alive.
+void AddStudioFonts(const FontLoaderFn& load);
+
+}  // namespace mujoco::platform
+
+#endif  // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_UX_FONTS_H_