Add an "offscreen" mode that allows users to render studio to a texture

then uses SDL to blit the texture into the window.

PiperOrigin-RevId: 861697047
Change-Id: Ib949b4ade0c79d867ba74ce013bc2b6fa2571f56

src/experimental/platform/window.cc

@@ -14,13 +14,22 @@
 
 #include "experimental/platform/window.h"
 
+#include <algorithm>
+#include <cstddef>
+#include <span>
 #include <string>
 #include <string_view>
 
+#ifdef MUJOCO_STUDIO_EGL_SUPPORTED
+#include "third_party/GL/gl/include/EGL/egl.h"
+#include "third_party/GL/gl/include/EGL/eglext.h"
+#endif
 #include <SDL.h>
 #include <SDL_error.h>
 #include <SDL_events.h>
 #include <SDL_hints.h>
+#include <SDL_render.h>
+#include <SDL_surface.h>
 #include <SDL_syswm.h>
 #include <SDL_version.h>
 #include <SDL_video.h>
@@ -89,22 +98,26 @@ static void InitImGui(SDL_Window* window, float content_scale, bool load_fonts,
 
 Window::Window(std::string_view title, int width, int height, Config config)
     : width_(width), height_(height), config_(config) {
-  SDL_SetHint(SDL_HINT_FRAMEBUFFER_ACCELERATION, "1");
-  SDL_GL_SetAttribute(SDL_GL_MULTISAMPLEBUFFERS, 1);
-  SDL_GL_SetAttribute(SDL_GL_MULTISAMPLESAMPLES, 16);
+  const RenderConfig render_config = config_.render_config;
+  if (config_.offscreen_mode) {
+    SDL_SetHint(SDL_HINT_RENDER_DRIVER, "software");
+    SDL_SetHint(SDL_HINT_FRAMEBUFFER_ACCELERATION, "0");
+  } else {
+    SDL_SetHint(SDL_HINT_FRAMEBUFFER_ACCELERATION, "1");
+    SDL_GL_SetAttribute(SDL_GL_MULTISAMPLEBUFFERS, 1);
+    SDL_GL_SetAttribute(SDL_GL_MULTISAMPLESAMPLES, 16);
+  }
+
   if (SDL_Init(SDL_INIT_VIDEO | SDL_INIT_EVENTS) != 0) {
     mju_error("Error initializing SDL: %s", SDL_GetError());
   }
-
-  int window_flags = SDL_WINDOW_RESIZABLE | SDL_WINDOW_ALLOW_HIGHDPI;
-
   if (!config_.enable_keyboard) {
     SDL_EventState(SDL_TEXTINPUT, SDL_DISABLE);
     SDL_EventState(SDL_KEYDOWN, SDL_DISABLE);
     SDL_EventState(SDL_KEYUP, SDL_DISABLE);
   }
 
-  RenderConfig render_config = config_.render_config;
+  int window_flags = SDL_WINDOW_RESIZABLE | SDL_WINDOW_ALLOW_HIGHDPI;
   if (render_config == kFilamentVulkan) {
     window_flags |= SDL_WINDOW_VULKAN;
   } else if (render_config == kFilamentWebGL) {
@@ -112,7 +125,8 @@ Window::Window(std::string_view title, int width, int height, Config config)
     SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 3);
     SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 0);
     SDL_GL_SetAttribute(SDL_GL_CONTEXT_PROFILE_MASK, SDL_GL_CONTEXT_PROFILE_ES);
-  } else if (render_config == kClassicOpenGL || render_config == kFilamentOpenGL) {
+  } else if (render_config == kClassicOpenGL ||
+             render_config == kFilamentOpenGL) {
     window_flags |= SDL_WINDOW_OPENGL;
     SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 3);
     SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 0);
@@ -122,8 +136,8 @@ Window::Window(std::string_view title, int width, int height, Config config)
 
   const float content_scale = ImGui_ImplSDL2_GetContentScaleForDisplay(0);
   sdl_window_ = SDL_CreateWindow(title.data(), SDL_WINDOWPOS_UNDEFINED,
-                                 SDL_WINDOWPOS_UNDEFINED, width * content_scale,
-                                 height * content_scale, window_flags);
+                                 SDL_WINDOWPOS_UNDEFINED, width,
+                                 height, window_flags);
   if (!sdl_window_) {
     mju_error("Error creating window: %s", SDL_GetError());
   }
@@ -131,23 +145,33 @@ Window::Window(std::string_view title, int width, int height, Config config)
   InitImGui(sdl_window_, content_scale, config.load_fonts,
             (render_config != kClassicOpenGL));
 
-  if (render_config == kFilamentWebGL || render_config == kClassicOpenGL) {
+  if (render_config == kFilamentWebGL ||
+      (render_config == kClassicOpenGL && !config.offscreen_mode)) {
     SDL_GLContext gl_context = SDL_GL_CreateContext(sdl_window_);
     SDL_GL_MakeCurrent(sdl_window_, gl_context);
   }
 
+  if (config.offscreen_mode) {
+    sdl_renderer_ = SDL_CreateRenderer(sdl_window_, -1, SDL_RENDERER_SOFTWARE);
+    if (config.render_config == kClassicOpenGL) {
+      InitOffscreenEglContext();
+    }
+  }
+
   SDL_SysWMinfo wmi;
   SDL_VERSION(&wmi.version);
   SDL_GetWindowWMInfo(sdl_window_, &wmi);
 
-  #if defined(__linux__)
-    native_window_ = reinterpret_cast<void*>(wmi.info.x11.window);
-  #elif defined(__WIN32__)
-    native_window_ = reinterpret_cast<void*>(wmi.info.win.window);
-  #elif defined(__APPLE__)
-    native_window_ =
-        GetNativeWindowOsx(reinterpret_cast<void*>(wmi.info.cocoa.window));
-  #endif
+  if (!config.offscreen_mode) {
+    #if defined(__linux__)
+      native_window_ = reinterpret_cast<void*>(wmi.info.x11.window);
+    #elif defined(__WIN32__)
+      native_window_ = reinterpret_cast<void*>(wmi.info.win.window);
+    #elif defined(__APPLE__)
+      native_window_ =
+          GetNativeWindowOsx(reinterpret_cast<void*>(wmi.info.cocoa.window));
+    #endif
+  }
 
   int drawable_width = width;
   int drawable_height = height;
@@ -156,6 +180,13 @@ Window::Window(std::string_view title, int width, int height, Config config)
 }
 
 Window::~Window() {
+#ifdef MUJOCO_STUDIO_EGL_SUPPORTED
+  if (egl_display_ && egl_context_) {
+    eglMakeCurrent(egl_display_, nullptr, nullptr, EGL_NO_CONTEXT);
+    eglDestroyContext(egl_display_, egl_context_);
+    eglTerminate(egl_display_);
+  }
+#endif
   SDL_DestroyWindow(sdl_window_);
   SDL_Quit();
 }
@@ -214,12 +245,133 @@ void Window::EndFrame() {
   ImGui::EndFrame();
 }
 
-void Window::Present() {
-  // Filament (with the exception of WebGL) handles the swapchain internally.
-  if (config_.render_config != kFilamentVulkan
+void Window::Present(std::span<const std::byte> pixels) {
+  if (config_.offscreen_mode) {
+    if (pixels.size() != width_ * height_ * 3) {
+      mju_error("Offscreen mode expects RGB buffer of size %d",
+                width_ * height_ * 3);
+    }
+
+    SDL_Surface* surface = SDL_GetWindowSurface(sdl_window_);
+    const unsigned char* src =
+        reinterpret_cast<const unsigned char*>(pixels.data());
+    unsigned char* dst =
+        static_cast<unsigned char*>(surface->pixels);
+
+    for (int i = 0; i < height_; ++i) {
+      for (int j = 0; j < width_; ++j) {
+        // Swap the RGB channels and add an alpha value of 1.0
+        dst[0] = src[2];
+        dst[1] = src[1];
+        dst[2] = src[0];
+        dst[3] = 0xff;
+        dst += 4;
+        src += 3;
+      }
+    }
+
+    if (config_.render_config == kClassicOpenGL) {
+      dst = static_cast<unsigned char*>(surface->pixels);
+      for (int r = 0; r < height_ / 2; ++r) {
+        unsigned char* top_row = &dst[4 * width_ * r];
+        unsigned char* bottom_row = &dst[4 * width_ * (height_ - 1 - r)];
+        std::swap_ranges(top_row, top_row + 4 * width_, bottom_row);
+      }
+    }
+
+    SDL_RenderPresent(sdl_renderer_);
+  } else if (config_.render_config != kFilamentVulkan
      && config_.render_config != kFilamentOpenGL) {
     SDL_GL_SwapWindow(sdl_window_);
   }
 }
 
+bool Window::IsOffscreenMode() const {
+  return config_.offscreen_mode;
+}
+
+void Window::InitOffscreenEglContext() {
+#ifdef MUJOCO_STUDIO_EGL_SUPPORTED
+  auto eglQueryDevicesEXT =
+      (PFNEGLQUERYDEVICESEXTPROC)eglGetProcAddress("eglQueryDevicesEXT");
+  auto eglGetPlatformDisplayEXT =
+      (PFNEGLGETPLATFORMDISPLAYEXTPROC)eglGetProcAddress(
+          "eglGetPlatformDisplayEXT");
+
+  if (!eglQueryDevicesEXT || !eglGetPlatformDisplayEXT) {
+    mju_error("Failed to load EGL functions");
+  }
+
+  // Query all available EGL devices.
+  constexpr int kMaxDevices = 32;
+  EGLDeviceEXT egl_devices[kMaxDevices];
+  int ndevice = 0;
+  if (!eglQueryDevicesEXT(kMaxDevices, egl_devices, &ndevice)) {
+    mju_error("eglQueryDevices error: 0x%x", eglGetError());
+  }
+
+  // Initialize the first valid EGL display.
+  for (int i = 0; i < ndevice; ++i) {
+    egl_display_ = eglGetPlatformDisplayEXT(EGL_PLATFORM_DEVICE_EXT,
+                                            egl_devices[i], nullptr);
+    if (egl_display_ != EGL_NO_DISPLAY) {
+      int major, minor;
+      EGLBoolean initialized = eglInitialize(egl_display_, &major, &minor);
+      if (!initialized) {
+        egl_display_ = EGL_NO_DISPLAY;
+      } else {
+        break;
+      }
+    }
+  }
+  if (egl_display_ == EGL_NO_DISPLAY) {
+    mju_error("Failed to create and initialize a valid EGL display!");
+  }
+
+  // Choose an EGL config.
+  constexpr EGLint config_attribs[] = {
+      EGL_RED_SIZE, 8,
+      EGL_GREEN_SIZE, 8,
+      EGL_BLUE_SIZE, 8,
+      EGL_ALPHA_SIZE, 8,
+      EGL_DEPTH_SIZE, 24,
+      EGL_STENCIL_SIZE, 8,
+      EGL_COLOR_BUFFER_TYPE, EGL_RGB_BUFFER,
+      EGL_SURFACE_TYPE, EGL_PBUFFER_BIT,
+      EGL_RENDERABLE_TYPE, EGL_OPENGL_BIT,
+      EGL_NONE
+  };
+  EGLint nconfig;
+  EGLConfig config;
+  if (!eglChooseConfig(egl_display_, config_attribs, &config, 1, &nconfig)) {
+    mju_error("eglChooseConfig error: 0x%x", eglGetError());
+  }
+
+  // Bind the OpenGL API to the EGL.
+  if (!eglBindAPI(EGL_OPENGL_API)) {
+    mju_error("eglBindAPI error: 0x%x", eglGetError());
+  }
+
+  // Create an EGL context.
+  constexpr EGLint context_attribs[] = {
+      EGL_CONTEXT_MAJOR_VERSION, 1,
+      EGL_CONTEXT_MINOR_VERSION, 5,
+      EGL_CONTEXT_OPENGL_PROFILE_MASK,
+      EGL_CONTEXT_OPENGL_COMPATIBILITY_PROFILE_BIT,
+      EGL_NONE};
+  egl_context_ =
+      eglCreateContext(egl_display_, config, EGL_NO_CONTEXT, context_attribs);
+  if (egl_context_ == EGL_NO_CONTEXT) {
+    mju_error("eglCreateContext error: 0x%x", eglGetError());
+  }
+
+  // Make the EGL context current.
+  if (!eglMakeCurrent(egl_display_, nullptr, nullptr, egl_context_)) {
+    mju_error("eglMakeCurrent error: 0x%x", eglGetError());
+  }
+#else
+  mju_error("EGL not supported on this platform");
+#endif
+}
+
 }  // namespace mujoco::platform

src/experimental/platform/window.h

@@ -15,11 +15,17 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_PLATFORM_WINDOW_H_
 #define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_WINDOW_H_
 
+#include <cstddef>
+#include <span>
 #include <string>
 #include <string_view>
 
-#include "experimental/platform/helpers.h"
+#if __has_include("third_party/GL/gl/include/EGL/egl.h")
+#define MUJOCO_STUDIO_EGL_SUPPORTED 1
+#include "third_party/GL/gl/include/EGL/egl.h"
+#endif
 #include <SDL_video.h>
+#include <SDL_render.h>
 
 namespace mujoco::platform {
 
@@ -41,6 +47,7 @@ class Window {
     RenderConfig render_config = kClassicOpenGL;
     bool enable_keyboard = true;
     bool load_fonts = true;
+    bool offscreen_mode = false;
   };
 
   Window(std::string_view title, int width, int height, Config config);
@@ -63,7 +70,7 @@ class Window {
   void EndFrame();
 
   // Swaps and presents the window buffer.
-  void Present();
+  void Present(std::span<const std::byte> pixels = {});
 
   // Sets the title of the window.
   void SetTitle(std::string_view title);
@@ -82,6 +89,9 @@ class Window {
   // the value will be cleared until the next time a file is dropped.
   std::string GetDropFile();
 
+  // Returns true if the window is in offscreen mode.
+  bool IsOffscreenMode() const;
+
   // Enables window resizing.
   void EnableWindowResizing();
 
@@ -92,14 +102,21 @@ class Window {
   void* GetNativeWindowHandle() { return native_window_; }
 
  private:
+  void InitOffscreenEglContext();
+
   int width_ = 0;
   int height_ = 0;
   float scale_ = 1.0f;
   Config config_;
   void* native_window_ = nullptr;
   SDL_Window* sdl_window_ = nullptr;
+  SDL_Renderer* sdl_renderer_ = nullptr;
   bool should_exit_ = false;
   std::string drop_file_;
+  #ifdef MUJOCO_STUDIO_EGL_SUPPORTED
+  EGLDisplay egl_display_ = EGL_NO_DISPLAY;
+  EGLContext egl_context_ = EGL_NO_CONTEXT;
+  #endif
 };
 
 }  // namespace mujoco::platform

src/experimental/studio/app.cc

@@ -46,17 +46,16 @@
 
 namespace mujoco::studio {
 
-static constexpr platform::Window::Config kWindowConfig = {
+static constexpr platform::Window::RenderConfig kRenderConfig =
 #if defined(__EMSCRIPTEN__)
-    .render_config = platform::Window::RenderConfig::kFilamentWebGL,
+  platform::Window::RenderConfig::kFilamentWebGL;
 #elif defined(USE_FILAMENT_VULKAN)
-    .render_config = platform::Window::RenderConfig::kFilamentVulkan,
+  platform::Window::RenderConfig::kFilamentVulkan;
 #elif defined(USE_FILAMENT_OPENGL)
-    .render_config = platform::Window::RenderConfig::kFilamentOpenGL,
+  platform::Window::RenderConfig::kFilamentOpenGL;
 #elif defined(USE_CLASSIC_OPENGL)
-    .render_config = platform::Window::RenderConfig::kClassicOpenGL,
+  platform::Window::RenderConfig::kClassicOpenGL;
 #endif
-};
 
 static void ToggleFlag(mjtByte& flag) { flag = flag ? 0 : 1; }
 
@@ -113,10 +112,13 @@ static constexpr std::array<const char*, 31> kPercentRealTime = {
 };
 // clang-format on
 
-App::App(int width, int height, std::string ini_path)
-    : ini_path_(std::move(ini_path)) {
-  window_ = std::make_unique<platform::Window>("MuJoCo Studio", width, height,
-                                               kWindowConfig);
+App::App(Config config)
+    : ini_path_(std::move(config.ini_path)) {
+  platform::Window::Config window_config;
+  window_config.render_config = kRenderConfig;
+  window_config.offscreen_mode = config.offscreen_mode;
+  window_ = std::make_unique<platform::Window>("MuJoCo Studio", config.width,
+                                               config.height, window_config);
   renderer_ =
       std::make_unique<platform::Renderer>(window_->GetNativeWindowHandle());
 
@@ -428,11 +430,16 @@ void App::Render() {
   const float height = window_->GetHeight();
   const float scale = window_->GetScale();
 
+  if (window_->IsOffscreenMode()) {
+    pixels_.resize(width * height * 3);
+  } else {
+    pixels_.clear();
+  }
   renderer_->Render(model_, data_, &perturb_, &camera_, &vis_options_,
-                    width * scale, height * scale);
+                    width * scale, height * scale, pixels_);
 
   window_->EndFrame();
-  window_->Present();
+  window_->Present(pixels_);
 
   if (data_) {
     for (int i = 0; i < mjNTIMER; i++) {

src/experimental/studio/app.h

@@ -39,7 +39,22 @@ namespace mujoco::studio {
 // Owns, updates, and renders a MuJoCo simulation.
 class App {
  public:
-  App(int width, int height, std::string ini_path);
+  // Configuration/initialization options for the application.
+  struct Config {
+    // The original width and height of the window.
+    int width = 0;
+    int height = 0;
+
+    // The path to the ini file containing the user settings.
+    std::string ini_path;
+
+    // By default, we render directly to the window surface. However, in some
+    // cases, we may want to render to an (offscreen) texture and blit the
+    // texture to the window surface.
+    bool offscreen_mode = false;
+  };
+
+  explicit App(Config config);
 
   // Loads an empty mjModel.
   void InitEmptyModel();
@@ -211,6 +226,7 @@ class App {
   mjSpec* spec_ = nullptr;
   mjModel* model_ = nullptr;
   mjData* data_ = nullptr;
+  std::vector<std::byte> pixels_;
 
   mjvCamera camera_;
   mjvPerturb perturb_;

src/experimental/studio/main.cc

@@ -31,6 +31,7 @@
 ABSL_FLAG(int, window_width, 1400, "Window width");
 ABSL_FLAG(int, window_height, 720, "Window height");
 ABSL_FLAG(std::string, model_file, "", "MuJoCo model file.");
+ABSL_FLAG(bool, offscreen_mode, false, "Offscreen mode");
 
 std::string Resolve(std::string_view path) {
   std::string_view subpath = path.substr(path.find(':') + 1);
@@ -100,7 +101,13 @@ int main(int argc, char** argv, char** envp) {
 
   const int width = absl::GetFlag(FLAGS_window_width);
   const int height = absl::GetFlag(FLAGS_window_height);
-  mujoco::studio::App app(width, height, ini_path);
+  const bool offscreen_mode = absl::GetFlag(FLAGS_offscreen_mode);
+  mujoco::studio::App app({
+    .width = width,
+    .height = height,
+    .ini_path = ini_path,
+    .offscreen_mode = offscreen_mode,
+  });
 
   // If the model file is not specified, try to load it from the first argument
   std::string model_file = absl::GetFlag(FLAGS_model_file);

src/experimental/studio/wasm.cc

@@ -91,7 +91,11 @@ void Init() {
   resource_provider.prefix = "filament";
   mjp_registerResourceProvider(&resource_provider);
 
-  g_app = new mujoco::studio::App(width, height, ini_path);
+  g_app = new mujoco::studio::App({
+    .width = width,
+    .height = height,
+    .ini_path = ini_path,
+  });
   g_app->InitEmptyModel();
 }