Allow MJB and XML content to be loaded directly.

PiperOrigin-RevId: 843180720
Change-Id: Id4f2bf81507b971dfbbf250d775d7fdfc76517df
This commit is contained in:
Haroon Qureshi
2025-12-11 05:40:36 -08:00
committed by Copybara-Service
parent 1ff74ba810
commit 87d333c64a
3 changed files with 62 additions and 40 deletions
+49 -37
View File
@@ -199,45 +199,53 @@ void App::ClearModel() {
error_ = "";
}
void App::LoadModel(std::string model_file) {
pending_load_ = std::move(model_file);
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);
}
}
void App::ProcessPendingLoad() {
if (!pending_load_.has_value()) {
return;
}
// Note that a non-empty model_file_ implies that a model was successfully
// loaded.
model_file_ = std::move(pending_load_.value());
pending_load_.reset();
void App::LoadModel(std::string data, ContentType type) {
// Delete the existing mjModel and mjData.
ClearModel();
// Try to load the requested mjModel.
char err[1000] = "";
if (model_file_.ends_with(".mjb")) {
model_ = mj_loadModel(model_file_.c_str(), 0);
} else if (model_file_.ends_with(".xml")) {
spec_ = mj_parseXML(model_file_.c_str(), nullptr, err, sizeof(err));
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);
if (model_name_.ends_with(".mjb")) {
model_ = mj_loadModel(model_name_.c_str(), 0);
} else if (model_name_.ends_with(".xml")) {
spec_ = mj_parseXML(model_name_.c_str(), nullptr, err, sizeof(err));
if (spec_ && err[0] == 0) {
model_ = mj_compile(spec_, nullptr);
}
} else {
error_ = "Unknown model file type; expected .mjb or .xml.";
}
} 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 {
error_ = "Unknown model file type; expected .mjb or .xml.";
} else if (type == ContentType::kModelMjb) {
model_name_ = "[mjb]";
model_ = mj_loadModelBuffer(data.data(), data.size());
}
if (err[0]) {
error_ = err;
fprintf(stderr, "Error loading model: %s\n", error_.c_str());
}
// If no mjModel was loaded, load an empty mjModel.
if (model_file_.empty() || model_ == nullptr) {
if (model_name_.empty() || model_ == nullptr) {
spec_ = mj_makeSpec();
model_ = mj_compile(spec_, 0);
model_file_ = "";
model_name_ = "";
}
if (!model_) {
mju_error("Error loading model: %s", error_.c_str());
@@ -270,11 +278,11 @@ void App::ProcessPendingLoad() {
// to the loaded model.
std::string base_path = "/";
std::string model_name = "model";
if (!model_file_.empty() &&
(model_file_.ends_with(".xml") || model_file_.ends_with(".mjb"))) {
window_->SetTitle("MuJoCo Studio : " + model_file_);
tmp_.last_load_file = std::string(model_file_);
std::filesystem::path path(model_file_);
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 {
@@ -289,7 +297,7 @@ void App::ProcessPendingLoad() {
tmp_.last_save_screenshot_file = base_path + "screenshot.webp";
}
bool App::IsModelLoaded() const { return !model_file_.empty(); }
bool App::IsModelLoaded() const { return !model_name_.empty(); }
void App::ResetPhysics() {
mj_resetData(model_, data_);
@@ -298,7 +306,11 @@ void App::ResetPhysics() {
}
void App::UpdatePhysics() {
ProcessPendingLoad();
if (pending_load_.has_value()) {
std::string model_file = std::move(pending_load_.value());
pending_load_.reset();
LoadModel(model_file, ContentType::kFilepath);
}
if (!IsModelLoaded()) {
return;
}
@@ -372,7 +384,7 @@ bool App::Update() {
// Check to see if a model was dropped on the window.
const std::string drop_file = window_->GetDropFile();
if (!drop_file.empty()) {
LoadModel(drop_file);
RequestModelLoad(drop_file);
}
// Only update the simulation if a popup window is not open. Note that the
@@ -564,7 +576,7 @@ void App::HandleKeyboardEvents() {
std::string keyframe = platform::KeyframeToString(model_, data_, false);
platform::MaybeSaveToClipboard(keyframe);
} else if (ImGui_IsChordJustPressed(ImGuiKey_L | ImGuiMod_Ctrl)) {
LoadModel(model_file_);
RequestModelLoad(model_name_);
} else if (ImGui_IsChordJustPressed(ImGuiKey_Q | ImGuiMod_Ctrl)) {
tmp_.should_exit = true;
} else if (ImGui_IsChordJustPressed(ImGuiKey_A | ImGuiMod_Ctrl)) {
@@ -1208,14 +1220,14 @@ void App::ToolBarGui() {
// Reset/Reload/Unload.
style.Color(ImGuiCol_ButtonHovered, ImColor(220, 40, 40, 255));
if (ImGui::Button(ICON_UNLOAD_MODEL, ImVec2(48, 32))) {
LoadModel("");
RequestModelLoad("");
}
ImGui::SetItemTooltip("%s", "Unload");
style.Reset();
ImGui::SameLine();
if (ImGui::Button(ICON_RELOAD_MODEL, ImVec2(48, 32))) {
LoadModel(model_file_);
RequestModelLoad(model_name_);
}
ImGui::SetItemTooltip("%s", "Reload");
@@ -1324,7 +1336,7 @@ void App::StatusBarGui() {
ImGui::TableNextColumn();
if (model_file_.empty()) {
if (!IsModelLoaded()) {
ImGui::Text("Not loaded");
} else if (model_ == nullptr) {
ImGui::Text("Not loaded");
@@ -1416,7 +1428,7 @@ void App::MainMenuGui() {
}
ImGui::Separator();
if (ImGui::MenuItem("Unload", "Ctrl+U")) {
LoadModel("");
RequestModelLoad("");
}
ImGui::Separator();
if (ImGui::MenuItem("Quit", "Ctrl+Q")) {
@@ -1432,7 +1444,7 @@ void App::MainMenuGui() {
ResetPhysics();
}
if (ImGui::MenuItem("Reload", "Ctrl+L")) {
LoadModel(model_file_);
RequestModelLoad(model_name_);
}
ImGui::Separator();
if (ImGui::BeginMenu("Keyframes")) {
@@ -1561,7 +1573,7 @@ void App::FileDialogGui() {
if (ImGui::BeginPopupModal("LoadModel", NULL,
ImGuiWindowFlags_AlwaysAutoResize)) {
if (platform::ImGui_FileDialog(tmp_.filename, sizeof(tmp_.filename))) {
LoadModel(tmp_.filename);
RequestModelLoad(tmp_.filename);
tmp_.last_load_file = tmp_.filename;
}
ImGui::EndPopup();
+12 -2
View File
@@ -48,8 +48,17 @@ class App {
App(int width, int height, std::string ini_path,
const platform::LoadAssetFn& load_asset_fn);
enum ContentType {
kFilepath, // Path to a model file.
kModelXml, // XML model string.
kModelMjb, // Binary model payload.
};
// Loads a model into the simulation.
void LoadModel(std::string model_file);
//
// 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);
// Processes window events and advances the state of the simulation.
bool Update();
@@ -132,6 +141,7 @@ class App {
void ClearModel();
void ProcessPendingLoad();
bool IsModelLoaded() const;
void RequestModelLoad(std::string model_file);
void ResetPhysics();
void UpdatePhysics();
@@ -164,7 +174,7 @@ class App {
std::string error_;
std::string ini_path_;
std::string model_file_;
std::string model_name_;
std::optional<std::string> pending_load_;
std::unique_ptr<platform::Window> window_;
+1 -1
View File
@@ -56,7 +56,7 @@ 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);
app.LoadModel(model_file, mujoco::studio::App::ContentType::kFilepath);
while (app.Update()) {
app.BuildGui();
app.Render();