chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor model loading APIs.

Browse Source 
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
This commit is contained in:
Haroon Qureshi
2026-01-21 05:24:18 -08:00
committed by Copybara-Service
parent e2180ff744
commit 2cc2e579f1
5 changed files with 227 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/studio/app.cc
+162 -104
View File
@@ -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_ = "";
} else {
pending_load_ = std::move(model_file);
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) {
// Delete the existing mjModel and mjData.
void App::InitEmptyModel() {
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()) {
char err[1000] = "";
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());
}
model_path_ = std::move(filename);
model_kind_ = model_kind;
if (err[0]) {
error_ = err;
spec_ = spec;
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()) {
std::string model_file = std::move(pending_load_.value());
pending_load_.reset();
LoadModel(model_file, ContentType::kFilepath);
}
if (!IsModelLoaded()) {
if (model_ == nullptr) {
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 height = window_->GetHeight();
const float width = window_->GetWidth() * ImGui::GetWindowDpiScale();
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()) {
ImGui::Text("Not loaded");
} else if (model_ == nullptr) {
ImGui::Text("Not loaded");
if (model_ == nullptr) {
ImGui::Text("No model 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
+40 -16
View File
@@ -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 {
kFilepath, // Path to a model file.
kModelXml, // XML model string.
kModelMjb, // Binary model payload.
};
// Loads an empty mjModel.
void InitEmptyModel();
// Loads a model into the simulation.
//
// Note: Do not call this function from within Update() (i.e. while drawing
// the UX). Call RequestModelLoad() instead.
void LoadModel(std::string data, ContentType type);
// Loads an mjModel from the given file. This extension should be one of:
// .xml or .mjb.
void LoadModelFromFile(const std::string& filepath);
// 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();
void ProcessPendingLoad();
bool IsModelLoaded() const;
// Requests that the model be loaded from the given file at the next update.
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
+2 -9
View File
@@ -130,14 +130,7 @@
reader.onload = (e) => {
const buffer = e.target.result;
try {
if (file.name.endsWith('.mjb')) {
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);
Module.loadFile(file.name, buffer);
} 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
+7 -1
View File
@@ -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
+16 -12
View File
@@ -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.
void LoadMjb(const std::string& src) {
if (g_app) {
g_app->LoadModel(src, mujoco::studio::App::ContentType::kModelMjb);
// Javascript-facing function to load a model from an uploaded file.
void LoadFile(const std::string& filename, const std::string& data) {
if (!g_app) {
return;
}
}
// Javascript-facing function to load a model from an XML file.
void LoadXml(const std::string& src) {
if (g_app) {
g_app->LoadModel(src, mujoco::studio::App::ContentType::kModelXml);
std::string content_type;
if (filename.ends_with(".mjb")) {
content_type = "application/mjb";
} 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("loadXml", &LoadXml);
emscripten::function("loadFile", &LoadFile);
emscripten::function("renderFrame", &RenderFrame);
emscripten::function("deinit", &Deinit);
}