Refactor model loading APIs.

There are three ways to load a model: from a file, from a buffer, or an
"empty" model. Instead of overloading a single function to do all three,
provide three separate functions.

Also disable model loading functions that will not work for emscripten.

PiperOrigin-RevId: 859051341
Change-Id: I71557d31d0e7c9c6882895c36116d300f19e255d

src/experimental/studio/app.cc

@@ -63,7 +63,11 @@ static void ToggleWindow(bool& window) {
   ImGui::GetIO().WantSaveIniSettings = true;
 }
 
-static void ShowPopup(bool& popup) { popup = true; }
+static void ShowFileDialog(bool& dialog) {
+  #ifndef __EMSCRIPTEN__
+    dialog = true;
+  #endif
+}
 
 static void SelectParentPerturb(const mjModel* model, mjvPerturb& perturb) {
   if (perturb.select > 0) {
@@ -141,71 +145,114 @@ void App::ClearModel() {
     }
   }
 
+  window_->SetTitle("MuJoCo Studio");
+
   step_control_.SetSpeed(100.f);
   profiler_.Clear();
   tmp_ = UiTempState();
-  error_ = "";
+  load_error_ = "";
+  step_error_ = "";
 }
 
 void App::RequestModelLoad(std::string model_file) {
-  if (model_file.starts_with('[') || model_file.ends_with(']')) {
+  pending_load_ = std::move(model_file);
-    pending_load_ = "";
+}
-  } else {
+
-    pending_load_ = std::move(model_file);
+void App::RequestModelReload() {
+  if (model_kind_ == kModelFromFile) {
+    pending_load_ = model_path_;
   }
 }
 
-void App::LoadModel(std::string data, ContentType type) {
+void App::InitEmptyModel() {
-  // Delete the existing mjModel and mjData.
+  mjSpec* spec = mj_makeSpec();
+  mjModel* model = mj_compile(spec, nullptr);
+  InitModel(model, spec, nullptr, "", kEmptyModel);
+}
+
+void App::LoadModelFromFile(const std::string& filepath) {
+  mjModel* model = nullptr;
+  mjSpec* spec = nullptr;
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+
+  const std::string resolved_file =
+      platform::ResolveFile(filepath, search_paths_);
+
+  char err[1000] = "";
+  if (resolved_file.ends_with(".mjb")) {
+    model = mj_loadModel(resolved_file.c_str(), &vfs);
+  } else {
+    spec = mj_parse(resolved_file.c_str(), nullptr, &vfs, err, sizeof(err));
+  }
+  if (err[0]) {
+    SetLoadError(err);
+    return;
+  }
+
+  InitModel(model, spec, &vfs, filepath, kModelFromFile);
+  UpdateFilePaths(resolved_file);
+  window_->SetTitle("MuJoCo Studio : " + filepath);
+
+  mj_deleteVFS(&vfs);
+}
+
+void App::LoadModelFromBuffer(std::span<const std::byte> buffer,
+                              std::string_view content_type,
+                              std::string_view filename) {
+  mjModel* model = nullptr;
+  mjSpec* spec = nullptr;
+
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+
+  char err[1000] = "";
+  if (content_type == "text/xml") {
+    const char* ptr = reinterpret_cast<const char*>(buffer.data());
+    spec = mj_parseXMLString(ptr, nullptr, err, sizeof(err));
+  } else if (content_type == "application/mjb") {
+    model = mj_loadModelBuffer(buffer.data(), buffer.size());
+  } else {
+    SetLoadError(
+        "Unknown content type; expected text/xml or application/mjb");
+    return;
+  }
+  if (err[0]) {
+    SetLoadError(err);
+    return;
+  }
+
+  InitModel(model, spec, &vfs, "", kModelFromBuffer);
+
+  mj_deleteVFS(&vfs);
+}
+
+void App::InitModel(mjModel* model, mjSpec* spec, mjVFS* vfs,
+                    std::string filename, ModelKind model_kind) {
   ClearModel();
 
-  if (!data.empty()) {
+  model_path_ = std::move(filename);
-    char err[1000] = "";
+  model_kind_ = model_kind;
-    if (type == ContentType::kFilepath) {
-      // Store the file path as the model name. Note that we use this model name
-      // to perform reload operations.
-      model_name_ = std::move(data);
-      const std::string resolved_file =
-        platform::ResolveFile(model_name_, search_paths_);
-      if (resolved_file.ends_with(".mjb")) {
-        model_ = mj_loadModel(resolved_file.c_str(), 0);
-      } else {
-        spec_ =
-            mj_parse(resolved_file.c_str(), nullptr, nullptr, err, sizeof(err));
-        if (spec_ && err[0] == 0) {
-          model_ = mj_compile(spec_, nullptr);
-        }
-      }
-    } else if (type == ContentType::kModelXml) {
-      model_name_ = "[xml]";
-      spec_ = mj_parseXMLString(data.c_str(), nullptr, err, sizeof(err));
-      if (spec_ && err[0] == 0) {
-        model_ = mj_compile(spec_, nullptr);
-      }
-    } else if (type == ContentType::kModelMjb) {
-      model_name_ = "[mjb]";
-      model_ = mj_loadModelBuffer(data.data(), data.size());
-    }
 
-    if (err[0]) {
+  spec_ = spec;
-      error_ = err;
+  model_ = model;
+
+  // If we have a spec but not a model, we need to compile the model from the spec.
+  if (spec_ && !model_) {
+    model_ = mj_compile(spec_, vfs);
+    if (!model_) {
+      SetLoadError("Error compiling model from spec.");
+      return;
     }
   }
 
-  // If no mjModel was loaded, load an empty mjModel.
-  if (model_name_.empty() || model_ == nullptr) {
-    spec_ = mj_makeSpec();
-    model_ = mj_compile(spec_, 0);
-    model_name_ = "";
-  }
   if (!model_) {
-    mju_error("Error loading model: %s", error_.c_str());
+    mju_error("Error making data for model: %s", model_path_.c_str());
   }
 
-  // Create the mjData for the mjModel.
   data_ = mj_makeData(model_);
   if (!data_) {
-    mju_error("Error making data for model: %s", error_.c_str());
+    mju_error("Error making data for model: %s", model_path_.c_str());
   }
 
   // Reset/reinitialize everything that depends on the new mjModel.
@@ -224,45 +271,35 @@ void App::LoadModel(std::string data, ContentType type) {
       SetSpeedIndex(i);
     }
   }
-
-  // Update the window title and update the file paths for saving files related
-  // to the loaded model.
-  std::string base_path = "/";
-  std::string model_name = "model";
-  if (!model_name_.empty() &&
-      (model_name_.ends_with(".xml") || model_name_.ends_with(".mjb"))) {
-    window_->SetTitle("MuJoCo Studio : " + model_name_);
-    tmp_.last_load_file = std::string(model_name_);
-    std::filesystem::path path(model_name_);
-    base_path = path.parent_path().string() + "/";
-    model_name = path.stem().string();
-  } else {
-    window_->SetTitle("MuJoCo Studio");
-    tmp_.last_load_file = base_path;
-  }
-
-  tmp_.last_save_mjb_file = base_path + model_name + "_saved.mjb";
-  tmp_.last_save_xml_file = base_path + model_name + "_saved.xml";
-  tmp_.last_print_model_file = base_path + model_name + "_MJMODEL.TXT";
-  tmp_.last_print_data_file = base_path + model_name + "_MJDATA.TXT";
-  tmp_.last_save_screenshot_file = base_path + "screenshot.webp";
 }
 
-bool App::IsModelLoaded() const { return !model_name_.empty(); }
+void App::UpdateFilePaths(const std::string& resolved_path) {
+  if (!resolved_path.empty()) {
+    const std::filesystem::path path(resolved_path);
+    const std::string base_path = path.parent_path().string() + "/";
+    const std::string model_name = path.stem().string();
+    tmp_.last_load_file = base_path + model_name + ".xml";
+    tmp_.last_save_mjb_file = base_path + model_name + "_saved.mjb";
+    tmp_.last_save_xml_file = base_path + model_name + "_saved.xml";
+    tmp_.last_print_model_file = base_path + model_name + "_MJMODEL.TXT";
+    tmp_.last_print_data_file = base_path + model_name + "_MJDATA.TXT";
+    tmp_.last_save_screenshot_file = base_path + "screenshot.webp";
+  }
+}
+
+void App::SetLoadError(std::string error) {
+  InitEmptyModel();
+  load_error_ = std::move(error);
+}
 
 void App::ResetPhysics() {
   mj_resetData(model_, data_);
   mj_forward(model_, data_);
-  error_ = "";
+  step_error_ = "";
 }
 
 void App::UpdatePhysics() {
-  if (pending_load_.has_value()) {
+  if (model_ == nullptr) {
-    std::string model_file = std::move(pending_load_.value());
-    pending_load_.reset();
-    LoadModel(model_file, ContentType::kFilepath);
-  }
-  if (!IsModelLoaded()) {
     return;
   }
 
@@ -289,14 +326,14 @@ void App::UpdatePhysics() {
   } else if (status == platform::StepControl::Status::kOk) {
     stepped = true;
     // If we are adding to the history we didn't have a divergence error
-    error_ = "";
+    step_error_ = "";
   } else if (status == platform::StepControl::Status::kAutoReset) {
     ResetPhysics();
   } else if (status == platform::StepControl::Status::kDiverged) {
     stepped = true;
     for (mjtWarning w : platform::StepControl::kDivergedWarnings) {
       if (data_->warning[w].number > 0) {
-        error_ = mju_warningText(w, data_->warning[w].lastinfo);
+        step_error_ = mju_warningText(w, data_->warning[w].lastinfo);
       }
     }
   }
@@ -334,6 +371,18 @@ bool App::Update() {
     RequestModelLoad(drop_file);
   }
 
+  // Check to see if we need to load a new model.
+  if (pending_load_.has_value()) {
+    std::string model_file = std::move(pending_load_.value());
+    pending_load_.reset();
+
+    if (model_file.empty()) {
+      InitEmptyModel();
+    } else {
+      LoadModelFromFile(model_file);
+    }
+  }
+
   // Only update the simulation if a popup window is not open. Note that the
   // simulation itself will only update if it is not paused.
   if (!tmp_.modal_open) {
@@ -487,22 +536,22 @@ void App::HandleKeyboardEvents() {
 
   // Menu shortcuts.
   if (ImGui_IsChordJustPressed(ImGuiKey_O | ImGuiMod_Ctrl)) {
-    ShowPopup(tmp_.load_popup);
+    ShowFileDialog(tmp_.load_popup);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_S | ImGuiMode_CtrlShift)) {
-    ShowPopup(tmp_.save_mjb_popup);
+    ShowFileDialog(tmp_.save_mjb_popup);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_S | ImGuiMod_Ctrl)) {
-    ShowPopup(tmp_.save_xml_popup);
+    ShowFileDialog(tmp_.save_xml_popup);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_M | ImGuiMod_Ctrl)) {
-    ShowPopup(tmp_.print_model_popup);
+    ShowFileDialog(tmp_.print_model_popup);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_D | ImGuiMod_Ctrl)) {
-    ShowPopup(tmp_.print_data_popup);
+    ShowFileDialog(tmp_.print_data_popup);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_P | ImGuiMod_Ctrl)) {
-    ShowPopup(tmp_.save_screenshot_popup);
+    ShowFileDialog(tmp_.save_screenshot_popup);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_C | ImGuiMod_Ctrl)) {
     std::string keyframe = platform::KeyframeToString(model_, data_, false);
     platform::MaybeSaveToClipboard(keyframe);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_L | ImGuiMod_Ctrl)) {
-    RequestModelLoad(model_name_);
+    RequestModelReload();
   } else if (ImGui_IsChordJustPressed(ImGuiKey_Q | ImGuiMod_Ctrl)) {
     tmp_.should_exit = true;
   } else if (ImGui_IsChordJustPressed(ImGuiKey_A | ImGuiMod_Ctrl)) {
@@ -820,11 +869,12 @@ void App::BuildGui() {
   }
 
   // Display a drag-and-drop message if no model is loaded.
-  if (!IsModelLoaded()) {
+  if (model_kind_ == kEmptyModel) {
+    #ifndef __EMSCRIPTEN__
     const char* text = "Load model file or drag-and-drop model file here.";
 
-    const float width = window_->GetWidth();
+    const float width = window_->GetWidth() * ImGui::GetWindowDpiScale();
-    const float height = window_->GetHeight();
+    const float height = window_->GetHeight() * ImGui::GetWindowDpiScale();
     const ImVec2 text_size = ImGui::CalcTextSize(text);
 
     ImGui::SetNextWindowPos(ImVec2((width - text_size.x) / 2, height / 2),
@@ -835,9 +885,12 @@ void App::BuildGui() {
         ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoBackground |
         ImGuiWindowFlags_NoInputs | ImGuiWindowFlags_NoDocking;
     if (ImGui::Begin("##Overlay", 0, kOverlayFlags)) {
+      platform::ScopedStyle style;
+      style.Color(ImGuiCol_Text, ImVec4(1, 1, 1, 1));
       ImGui::Text("%s", text);
     }
     ImGui::End();
+    #endif  // !__EMSCRIPTEN__
   }
 
   FileDialogGui();
@@ -1231,7 +1284,7 @@ void App::ToolBarGui() {
     ImGui::SameLine();
     style.Color(ImGuiCol_ButtonHovered, red);
     if (ImGui::Button(ICON_UNLOAD_MODEL, button_size)) {
-      RequestModelLoad("");
+      InitEmptyModel();
     }
     ImGui::SetItemTooltip("%s", "Unload");
     style.Reset();
@@ -1239,7 +1292,7 @@ void App::ToolBarGui() {
     // Reload button.
     ImGui::SameLine();
     if (ImGui::Button(ICON_RELOAD_MODEL, button_size)) {
-      RequestModelLoad(model_name_);
+      RequestModelReload();
     }
     ImGui::SetItemTooltip("%s", "Reload");
 
@@ -1385,10 +1438,8 @@ void App::StatusBarGui() {
 
     ImGui::TableNextColumn();
 
-    if (!IsModelLoaded()) {
+    if (model_ == nullptr) {
-      ImGui::Text("Not loaded");
+      ImGui::Text("No model loaded");
-    } else if (model_ == nullptr) {
-      ImGui::Text("Not loaded");
     } else if (step_control_.IsPaused()) {
       ImGui::Text("Paused");
     } else {
@@ -1405,9 +1456,12 @@ void App::StatusBarGui() {
       }
     }
 
-    if (!error_.empty()) {
+    if (!step_error_.empty()) {
       ImGui::SameLine();
-      ImGui::Text(" | Error: %s", error_.c_str());
+      ImGui::Text(" | Step Error: %s", step_error_.c_str());
+    } else if (!load_error_.empty()) {
+      ImGui::SameLine();
+      ImGui::Text(" | Load Error: %s", load_error_.c_str());
     }
 
     ImGui::TableNextColumn();
@@ -1457,34 +1511,38 @@ void App::StatusBarGui() {
 void App::MainMenuGui() {
   if (ImGui::BeginMainMenuBar()) {
     if (ImGui::BeginMenu("File")) {
+      #ifndef __EMSCRIPTEN__
       if (ImGui::MenuItem("Open Model File", "Ctrl+O")) {
-        ShowPopup(tmp_.load_popup);
+        ShowFileDialog(tmp_.load_popup);
       }
       ImGui::Separator();
       if (ImGui::MenuItem("Save XML", "Ctrl+S")) {
-        ShowPopup(tmp_.save_xml_popup);
+        ShowFileDialog(tmp_.save_xml_popup);
       }
       if (ImGui::MenuItem("Save MJB", "Ctrl+Shift+S")) {
-        ShowPopup(tmp_.save_mjb_popup);
+        ShowFileDialog(tmp_.save_mjb_popup);
       }
       if (ImGui::MenuItem("Save Screenshot", "Ctrl+P")) {
-        ShowPopup(tmp_.save_screenshot_popup);
+        ShowFileDialog(tmp_.save_screenshot_popup);
       }
       ImGui::Separator();
       if (ImGui::MenuItem("Print Model", "Ctrl+M")) {
-        ShowPopup(tmp_.print_model_popup);
+        ShowFileDialog(tmp_.print_model_popup);
       }
       if (ImGui::MenuItem("Print Data", "Ctrl+D")) {
-        ShowPopup(tmp_.print_data_popup);
+        ShowFileDialog(tmp_.print_data_popup);
       }
       ImGui::Separator();
+      #endif  // !__EMSCRIPTEN__
       if (ImGui::MenuItem("Unload", "Ctrl+U")) {
-        RequestModelLoad("");
+        InitEmptyModel();
       }
+      #ifndef __EMSCRIPTEN__
       ImGui::Separator();
       if (ImGui::MenuItem("Quit", "Ctrl+Q")) {
         tmp_.should_exit = true;
       }
+      #endif  // !__EMSCRIPTEN__
       ImGui::EndMenu();
     }
     if (ImGui::BeginMenu("Simulation")) {
@@ -1495,7 +1553,7 @@ void App::MainMenuGui() {
         ResetPhysics();
       }
       if (ImGui::MenuItem("Reload", "Ctrl+L")) {
-        RequestModelLoad(model_name_);
+        RequestModelReload();
       }
       ImGui::Separator();
       if (ImGui::BeginMenu("Keyframes")) {

src/experimental/studio/app.h

@@ -15,15 +15,17 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_STUDIO_APP_H_
 #define MUJOCO_SRC_EXPERIMENTAL_STUDIO_APP_H_
 
+#include <cstddef>
 #include <memory>
 #include <optional>
+#include <span>
 #include <string>
+#include <string_view>
 #include <unordered_map>
 #include <vector>
 
 #include <mujoco/mujoco.h>
 #include "experimental/platform/gui.h"
-#include "experimental/platform/helpers.h"
 #include "experimental/platform/interaction.h"
 #include "experimental/platform/picture_gui.h"
 #include "experimental/platform/renderer.h"
@@ -39,17 +41,18 @@ class App {
  public:
   App(int width, int height, std::string ini_path);
 
-  enum ContentType {
+  // Loads an empty mjModel.
-    kFilepath,  // Path to a model file.
+  void InitEmptyModel();
-    kModelXml,  // XML model string.
-    kModelMjb,  // Binary model payload.
-  };
 
-  // Loads a model into the simulation.
+  // Loads an mjModel from the given file. This extension should be one of:
-  //
+  // .xml or .mjb.
-  // Note: Do not call this function from within Update() (i.e. while drawing
+  void LoadModelFromFile(const std::string& filepath);
-  // the UX). Call RequestModelLoad() instead.
+
-  void LoadModel(std::string data, ContentType type);
+  // Loads an mjModel from the given memory buffer. The content_type should be
+  // one of: "text/xml" or "application/mjb".
+  void LoadModelFromBuffer(std::span<const std::byte> buffer,
+                           std::string_view content_type,
+                           std::string_view name);
 
   // Processes window events and advances the state of the simulation.
   bool Update();
@@ -62,6 +65,13 @@ class App {
   void Render();
 
  private:
+  // The kind of model that is currently loaded.
+  enum ModelKind {
+    kEmptyModel,
+    kModelFromFile,
+    kModelFromBuffer,
+  };
+
   // UI state that is persisted across application runs
   struct UiState {
     char watch_field[1000] = "qpos";
@@ -131,11 +141,23 @@ class App {
     std::string last_print_data_file;
   };
 
-  void ClearModel();
+  // Requests that the model be loaded from the given file at the next update.
-  void ProcessPendingLoad();
-  bool IsModelLoaded() const;
   void RequestModelLoad(std::string model_file);
 
+  // Requests that the currently loaded model be reloaded at the next update.
+  void RequestModelReload();
+
+  // Clears the currently loaded model and all associated state.
+  void ClearModel();
+
+  // Updates the currently loaded model to the given model. If model is null,
+  // then compile the spec to a model.
+  void InitModel(mjModel* model, mjSpec* spec, mjVFS* vfs, std::string filename,
+                 ModelKind model_kind);
+
+  void SetLoadError(std::string error);
+  void UpdateFilePaths(const std::string& resolved_path);
+
   void ResetPhysics();
   void UpdatePhysics();
 
@@ -165,10 +187,12 @@ class App {
   float GetExpectedLabelWidth();
   std::vector<const char*> GetCameraNames();
 
-  std::string error_;
   std::string ini_path_;
-  std::string model_name_;
+  std::string model_path_;
+  std::string load_error_;
+  std::string step_error_;
   std::optional<std::string> pending_load_;
+  ModelKind model_kind_ = kEmptyModel;
 
   std::unique_ptr<platform::Window> window_;
   std::unique_ptr<platform::Renderer> renderer_;

src/experimental/studio/index.html

@@ -130,14 +130,7 @@
           reader.onload = (e) => {
             const buffer = e.target.result;
             try {
-              if (file.name.endsWith('.mjb')) {
+              Module.loadFile(file.name, buffer);
-                Module.loadMjb(buffer);
-              } else if (file.name.endsWith('.xml')) {
-                Module.loadXml(buffer);
-              } else {
-                console.error('Unsupported file type:', file.name);
-              }
-              console.log('Model loaded from file:', file.name);
             } catch (error) {
               console.error('Failed to load model from file:', error);
             }
@@ -145,7 +138,7 @@
           reader.onerror = (e) => {
             console.error('Error reading file:', e);
           };
-          reader.readAsArrayBuffer(file); // Reads as text, suitable for XML.
+          reader.readAsArrayBuffer(file);
         });
       });
     </script>

src/experimental/studio/main.cc

@@ -105,7 +105,13 @@ int main(int argc, char** argv, char** envp) {
   // If the model file is not specified, try to load it from the first argument
   std::string model_file = absl::GetFlag(FLAGS_model_file);
   if (model_file.empty() && argc > 1 && argv[1][0] != '-') model_file = argv[1];
-  app.LoadModel(model_file, mujoco::studio::App::ContentType::kFilepath);
+
+  if (model_file.empty()) {
+    app.InitEmptyModel();
+  } else {
+    app.LoadModelFromFile(model_file);
+  }
+
   while (app.Update()) {
     app.BuildGui();
     app.Render();

src/experimental/studio/wasm.cc

@@ -92,21 +92,26 @@ void Init() {
   mjp_registerResourceProvider(&resource_provider);
 
   g_app = new mujoco::studio::App(width, height, ini_path);
-  g_app->LoadModel("", mujoco::studio::App::ContentType::kModelXml);
+  g_app->InitEmptyModel();
 }
 
-// Javascript-facing function to load a model from a MJB file.
+// Javascript-facing function to load a model from an uploaded file.
-void LoadMjb(const std::string& src) {
+void LoadFile(const std::string& filename, const std::string& data) {
-  if (g_app) {
+  if (!g_app) {
-    g_app->LoadModel(src, mujoco::studio::App::ContentType::kModelMjb);
+    return;
   }
-}
 
-// Javascript-facing function to load a model from an XML file.
+  std::string content_type;
-void LoadXml(const std::string& src) {
+  if (filename.ends_with(".mjb")) {
-  if (g_app) {
+    content_type = "application/mjb";
-    g_app->LoadModel(src, mujoco::studio::App::ContentType::kModelXml);
+  } else if (filename.ends_with(".xml")) {
+    content_type = "text/xml";
+  } else {
+    return;
   }
+
+  const auto ptr = reinterpret_cast<const std::byte*>(data.data());
+  g_app->LoadModelFromBuffer({ptr, ptr + data.size()}, content_type, filename);
 }
 
 // Javascript-facing function to render a single frame.
@@ -128,8 +133,7 @@ void Deinit() {
 EMSCRIPTEN_BINDINGS(studio_bindings) {
   emscripten::function("registerAsset", &RegisterAsset);
   emscripten::function("init", &Init);
-  emscripten::function("loadMjb", &LoadMjb);
+  emscripten::function("loadFile", &LoadFile);
-  emscripten::function("loadXml", &LoadXml);
   emscripten::function("renderFrame", &RenderFrame);
   emscripten::function("deinit", &Deinit);
 }