diff --git a/src/experimental/platform/gui_spec.cc b/src/experimental/platform/gui_spec.cc index 5784b777..27dbbba3 100644 --- a/src/experimental/platform/gui_spec.cc +++ b/src/experimental/platform/gui_spec.cc @@ -21,6 +21,7 @@ #include #include #include "experimental/platform/imgui_widgets.h" +#include "experimental/platform/spec_editor.h" // Define the mujoco X macros to add fields to the ImGui_DataTable. // We limit the fields to the ones with a matching element by comparing the @@ -63,7 +64,7 @@ static int GetElementIndexInSpec(mjsElement* element) { // Returns a name for the element; either the element has a name, or we // construct a unique name from the element's id (using mjs_getId) or index // (using GetElementIndexInSpec). -std::string ElementName(mjsElement* element) { +static std::string GetElementName(mjsElement* element) { const mjString* name = mjs_getName(element); std::string label = *name; if (label.empty()) { @@ -77,118 +78,105 @@ std::string ElementName(mjsElement* element) { return label; } -static bool AddDeleteButton(mjsElement* element) { - // Right-align the delete button. - const float button_width = ImGui::CalcTextSize(ICON_FA_TRASH_CAN).x + - ImGui::GetStyle().FramePadding.x * 2.0f; - ImGui::SameLine(ImGui::GetWindowContentRegionMax().x - button_width); +static void AddEditorButtons(mjsElement* element, mjsElement** selected_element, + SpecEditor& editor) { + // Right-align the buttons. + float x = ImGui::GetWindowContentRegionMax().x; + x -= ImGui::GetStyle().FramePadding.x * 2; + x -= ImGui::CalcTextSize(ICON_FA_TRASH_CAN).x; + if (element->elemtype == mjOBJ_BODY) { + x -= ImGui::GetStyle().FramePadding.x * 2; + x -= ImGui::CalcTextSize(ICON_FA_PLUS).x; + } + ImGui::SameLine(x); + if (element->elemtype == mjOBJ_BODY) { + if (ImGui::SmallButton(ICON_FA_PLUS)) { + ImGui::OpenPopupOnItemClick("BodyAddChild", 0); + } + ImGui::SameLine(); + } if (ImGui::SmallButton(ICON_FA_TRASH_CAN)) { - mjs_delete(mjs_getSpec(element), element); - return true; + editor.DeleteActiveElement(); + *selected_element = nullptr; } - return false; -} -static bool AddBodyAddChildButton(mjsElement* element, mjsElement** selected_element) { - // Right-align the add button. - const float button_width = ImGui::CalcTextSize(ICON_FA_TRASH_CAN).x + - ImGui::CalcTextSize(ICON_FA_PLUS).x + - ImGui::GetStyle().FramePadding.x * 4.0f; - ImGui::SameLine(ImGui::GetWindowContentRegionMax().x - button_width); - if (ImGui::SmallButton(ICON_FA_PLUS)) { - ImGui::OpenPopupOnItemClick("BodyAddChild", 0); - } - bool modified = false; if (ImGui::BeginPopupContextItem("BodyAddChild")) { mjsBody* body = mjs_asBody(element); - auto option = [&](const char* label, auto fn) { + auto option = [&](const char* label, mjtObj type) { if (ImGui::Selectable(label)) { - *selected_element = fn()->element; - mjs_setName(*selected_element, ElementName(*selected_element).c_str()); - modified = true; + *selected_element = editor.AddBodyElement(body, type); } }; - option("Camera", [&]() { return mjs_addCamera(body, nullptr); }); - option("Frame", [&]() { return mjs_addFrame(body, nullptr); }); - option("Geom", [&]() { return mjs_addGeom(body, nullptr); }); - option("Joint", [&]() { return mjs_addJoint(body, nullptr); }); - option("Light", [&]() { return mjs_addLight(body, nullptr); }); - option("Site", [&]() { return mjs_addSite(body, nullptr); }); + option("Camera", mjOBJ_CAMERA); + option("Frame", mjOBJ_FRAME); + option("Geom", mjOBJ_GEOM); + option("Joint", mjOBJ_JOINT); + option("Light", mjOBJ_LIGHT); + option("Site", mjOBJ_SITE); ImGui::EndPopup(); } - return modified; } -static bool SelectableElement(mjsElement* element, +static void SelectableElement(mjsElement* element, mjsElement** selected_element, - SpecEditMode mode) { + SpecEditor* editor) { constexpr ImGuiSelectableFlags flags = ImGuiSelectableFlags_AllowOverlap; - const std::string name = ElementName(element); + const std::string name = GetElementName(element); const bool selected = (element == *selected_element); if (ImGui::Selectable(name.c_str(), selected, flags)) { *selected_element = element; } - - bool modified = false; - if (selected && mode == SpecEditMode::kEdit) { - if (AddDeleteButton(element)) { - *selected_element = nullptr; - modified = true; - } + if (selected && editor) { + AddEditorButtons(element, selected_element, *editor); } - return modified; } -static bool BodyChildrenGui(const char* heading, mjtObj type, +static void BodyChildrenGui(const char* heading, mjtObj type, mjsElement** element, mjsBody* body, - SpecEditMode mode) { + SpecEditor* editor) { mjsElement* iter = mjs_firstChild(body, type, 0); if (!iter) { - return false; + return; } - bool modified = false; constexpr ImGuiTreeNodeFlags tree_flags = ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_DrawLinesFull; if (ImGui::TreeNodeEx(heading, tree_flags)) { while (iter) { mjsElement* next = mjs_nextChild(body, iter, 0); - modified |= SelectableElement(iter, element, mode); + SelectableElement(iter, element, editor); iter = next; } ImGui::TreePop(); } - return modified; } -static bool ElementListGui(const char* heading, mjtObj type, +static void ElementListGui(const char* heading, mjtObj type, mjsElement** element, mjSpec* spec, - SpecEditMode mode) { + SpecEditor* editor) { mjsElement* iter = mjs_firstElement(spec, type); if (!iter) { - return false; + return; } - bool modified = false; constexpr ImGuiTreeNodeFlags tree_flags = ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed; if (ImGui::TreeNodeEx(heading, tree_flags)) { while (iter) { mjsElement* next = mjs_nextElement(spec, iter); - modified |= SelectableElement(iter, element, mode); + SelectableElement(iter, element, editor); iter = next; } ImGui::TreePop(); } - return modified; } -static bool BodyTreeGuiRecursive(mjsElement** element, mjsBody* body, - SpecEditMode mode) { - const std::string label = ElementName(body->element); +static void BodyTreeGuiRecursive(mjsElement** element, mjsBody* body, + SpecEditor* editor) { + const std::string label = GetElementName(body->element); - ImGui::PushID(body); + ImGui::PushID(label.c_str()); ImGuiTreeNodeFlags flags = ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed | @@ -198,40 +186,34 @@ static bool BodyTreeGuiRecursive(mjsElement** element, mjsBody* body, flags |= ImGuiTreeNodeFlags_Selected; } - bool modified = false; const bool tree_open = ImGui::TreeNodeEx(label.c_str(), flags); if (ImGui::IsItemClicked()) { *element = body->element; } - if (*element == body->element && mode == SpecEditMode::kEdit) { - modified |= AddBodyAddChildButton(body->element, element); - modified |= AddDeleteButton(body->element); + if (editor && *element == body->element) { + AddEditorButtons(body->element, element, *editor); } if (tree_open) { mjsElement* iter = mjs_firstChild(body, mjOBJ_BODY, 0); while (iter) { mjsElement* next = mjs_nextChild(body, iter, 0); - modified |= BodyTreeGuiRecursive(element, mjs_asBody(iter), mode); + BodyTreeGuiRecursive(element, mjs_asBody(iter), editor); iter = next; } - modified |= BodyChildrenGui("Frames", mjOBJ_FRAME, element, body, mode); - modified |= BodyChildrenGui("Sites", mjOBJ_SITE, element, body, mode); - modified |= BodyChildrenGui("Joints", mjOBJ_JOINT, element, body, mode); - modified |= BodyChildrenGui("Geoms", mjOBJ_GEOM, element, body, mode); - modified |= BodyChildrenGui("Lights", mjOBJ_LIGHT, element, body, mode); - modified |= BodyChildrenGui("Cameras", mjOBJ_CAMERA, element, body, mode); - + BodyChildrenGui("Frames", mjOBJ_FRAME, element, body, editor); + BodyChildrenGui("Sites", mjOBJ_SITE, element, body, editor); + BodyChildrenGui("Joints", mjOBJ_JOINT, element, body, editor); + BodyChildrenGui("Geoms", mjOBJ_GEOM, element, body, editor); + BodyChildrenGui("Lights", mjOBJ_LIGHT, element, body, editor); + BodyChildrenGui("Cameras", mjOBJ_CAMERA, element, body, editor); ImGui::TreePop(); } - ImGui::PopID(); - return modified; } -bool SpecTreeGui(mjsElement** element, mjSpec* spec, SpecEditMode mode) { - bool modified = false; +void SpecTreeGui(mjsElement** element, mjSpec* spec, SpecEditor* editor) { const ImGuiTreeNodeFlags flags = ImGuiTreeNodeFlags_SpanAvailWidth | ImGuiTreeNodeFlags_Framed; @@ -245,17 +227,17 @@ bool SpecTreeGui(mjsElement** element, mjSpec* spec, SpecEditMode mode) { if (root) { mjsBody* body = mjs_asBody(root); if (body) { - modified |= BodyTreeGuiRecursive(element, body, mode); + BodyTreeGuiRecursive(element, body, editor); } } ImGui::TreePop(); } auto list = [&](const char* heading, mjtObj type) { - modified |= ElementListGui(heading, type, element, spec, mode); + ElementListGui(heading, type, element, spec, editor); }; - ImGui::PushID(spec); + ImGui::PushID("$spec$"); // Non-tree elements. if (ImGui::TreeNodeEx("Elements", flags)) { @@ -285,22 +267,20 @@ bool SpecTreeGui(mjsElement** element, mjSpec* spec, SpecEditMode mode) { } ImGui::PopID(); - return modified; } -bool ElementSpecGui(mjsElement* element, mjsElement* ref_element, - SpecEditMode mode) { +void ElementSpecGui(mjsElement* element, SpecEditor* editor) { if (element == nullptr) { - return false; - } - if (ref_element == nullptr) { - ref_element = element; + return; } + // if editor, assert that element == editor->GetActiveElement(); + mjsElement* ref_element = editor ? editor->GetRefElement() : element; + #define FIELD(NAME, TIP) table(#NAME, elem->NAME, ref->NAME, TIP); #define QFIELD(NAME, ALT, TIP) table(#NAME, #ALT, elem->NAME, ref->NAME, elem->ALT, ref->ALT, TIP); - ImGui_SpecElementTable table(mode == SpecEditMode::kPlay); + ImGui_SpecElementTable table(editor == nullptr); switch (element->elemtype) { case mjOBJ_BODY: { mjsBody* elem = mjs_asBody(element); @@ -725,7 +705,10 @@ bool ElementSpecGui(mjsElement* element, mjsElement* ref_element, } #undef FIELD - return table.WasModified(); + + if (editor && table.WasModified()) { + editor->CommitChanges(element); + } } void ElementModelGui(const mjModel* model, mjsElement* element) { diff --git a/src/experimental/platform/gui_spec.h b/src/experimental/platform/gui_spec.h index 44cf22b6..6ff5b0cc 100644 --- a/src/experimental/platform/gui_spec.h +++ b/src/experimental/platform/gui_spec.h @@ -15,31 +15,24 @@ #ifndef MUJOCO_SRC_EXPERIMENTAL_PLATFORM_GUI_SPEC_H_ #define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_GUI_SPEC_H_ -#include - #include +#include "experimental/platform/spec_editor.h" namespace mujoco::platform { -// The mode of the spec visualizer. -enum class SpecEditMode { - kPlay, - kEdit, -}; - -// Returns the name of the given element. -std::string ElementName(mjsElement* element); - // Displaying the mjSpec as a tree. `element` is the currently selected -// element and will be updated if a new element is selected. The function -// returns true if the spec was modified in any way. -bool SpecTreeGui(mjsElement** element, mjSpec* spec, SpecEditMode mode); +// element and will be updated if a new element is selected. The `editor` is +// optional; if provided, it will allow the user to modify the spec (e.g. +// add/delete elements) on the assumption that `element` and `spec` were +// obtained from the editor. +void SpecTreeGui(mjsElement** element, mjSpec* spec, + SpecEditor* editor = nullptr); -// Displays the properties of the given element in the table. Returns true if -// any value in the element was changed. The `ref_element` is used to highlight -// when a value differs from a reference element. -bool ElementSpecGui(mjsElement* element, mjsElement* ref_element, - SpecEditMode mode); +// Displays the properties of the given element in the table. If `editor` is +// provided, the user will be able to modify the properties of the element, +// updating the editor of any changes on the assumption that `element` was +// obtained from the editor. +void ElementSpecGui(mjsElement* element, SpecEditor* editor = nullptr); // Displays a (read-only) data table of the mjData values that correspond to the // given element. diff --git a/src/experimental/platform/spec_editor.cc b/src/experimental/platform/spec_editor.cc new file mode 100644 index 00000000..dab71b49 --- /dev/null +++ b/src/experimental/platform/spec_editor.cc @@ -0,0 +1,377 @@ +// 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 +// +// http://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/spec_editor.h" + +#include +#include +#include +#include "experimental/platform/model_holder.h" + +namespace mujoco::platform { + +SpecEditor::SpecEditor(int history_size) : capacity_(history_size) {} + +void SpecEditor::Reset(const mjSpec& spec) { + active_element_key_ = kInvalidElementKey; + active_element_ = nullptr; + ref_element_ = nullptr; + + ref_spec_ = Copy(&spec); + active_spec_ = Copy(&spec); + + // Build the initial element maps. + for (int i = mjOBJ_UNKNOWN + 1; i < mjNOBJECT; ++i) { + const mjtObj type = static_cast(i); + mjsElement* element = mjs_firstElement(ref_spec_.get(), type); + while (element) { + ref_map_.Append(type, next_element_key_++); + element = mjs_nextElement(ref_spec_.get(), element); + } + } + active_map_ = ref_map_; + + // Seed the history buffer with the initial spec. + history_.clear(); + history_.push_back(HistoryEntry{ + .spec = Copy(ref_spec_.get()), + .op = kInitialize, + .key = kInvalidElementKey, + .type_index = kInvalidTypeIndex, + }); + cursor_ = 1; +} + +std::unique_ptr SpecEditor::Compile() { + auto holder = ModelHolder::FromSpec(mj_copySpec(active_spec_.get())); + if (holder->ok()) { + ref_spec_ = Copy(active_spec_.get()); + ref_map_ = active_map_; + } + return holder; +} + +mjSpec* SpecEditor::GetActiveSpec() const { + return active_spec_.get(); +} + +mjsElement* SpecEditor::AddElement(mjtObj type) { + // TODO: check that type is only for spec elements. + mjsElement* element = AddElementToSpec(active_spec_.get(), type); + if (element) { + const ElementKey key = next_element_key_++; + const int index = active_map_.Append(element->elemtype, key); + AppendHistory(HistoryEntry{ + .spec = Copy(active_spec_.get()), + .op = kAdd, + .key = key, + .type_index = {element->elemtype, index}, + }); + } + return element; +} + +mjsElement* SpecEditor::AddBodyElement(mjsBody* body, mjtObj type) { + // TODO: check that type is only for body elements. + mjsElement* element = AddElementToSpec(active_spec_.get(), type, body); + if (element) { + const ElementKey key = next_element_key_++; + const int index = active_map_.Append(element->elemtype, key); + AppendHistory(HistoryEntry{ + .spec = Copy(active_spec_.get()), + .op = kAdd, + .key = key, + .type_index = {element->elemtype, index}, + }); + } + return element; +} + +void SpecEditor::DeleteActiveElement() { + if (active_element_) { + const mjtObj type = active_element_->elemtype; + const int index = active_map_.Remove(active_element_key_); + mjs_delete(active_spec_.get(), active_element_); + AppendHistory(HistoryEntry{ + .spec = Copy(active_spec_.get()), + .op = kDelete, + .key = active_element_key_, + .type_index = {type, index}, + }); + + ref_element_ = nullptr; + active_element_ = nullptr; + active_element_key_ = kInvalidElementKey; + } +} + +void SpecEditor::SetActiveElement(mjsElement* element) { + if (element == nullptr) { + ref_element_ = nullptr; + active_element_ = nullptr; + active_element_key_ = kInvalidElementKey; + return; + } + + mjSpec* spec = mjs_getSpec(element); + if (spec != active_spec_.get()) { + mju_warning("Element is not owned by the active spec."); + return; + } + + const mjtObj type = element->elemtype; + active_element_ = element; + + int index = 0; + for (mjsElement* iter = mjs_firstElement(spec, type); iter != nullptr; + iter = mjs_nextElement(spec, iter), ++index) { + if (iter == element) { + active_element_key_ = active_map_.LookupElementKey({type, index}); + if (active_element_key_ == kInvalidElementKey) { + mju_warning("Element not found in active spec."); + } + UpdateReferenceElement(); + return; + } + } + + mju_warning("Element not found in active spec."); + ref_element_ = nullptr; + active_element_ = nullptr; + active_element_key_ = kInvalidElementKey; +} + +void SpecEditor::UpdateReferenceElement() { + if (active_element_ == nullptr) { + ref_element_ = nullptr; + } else { + const mjtObj type = active_element_->elemtype; + ref_element_ = ref_map_.Resolve(ref_spec_.get(), active_element_key_); + if (ref_element_ == nullptr) { + ref_element_ = mjs_firstElement(dummy_spec_.get(), type); + } + if (ref_element_ == nullptr) { + ref_element_ = AddElementToSpec(dummy_spec_.get(), type); + } + if (ref_element_ == nullptr) { + mju_warning("Failed to create reference element."); + ref_element_ = active_element_; + } + } +} + +mjsElement* SpecEditor::GetActiveElement() const { + return active_element_; +} + +mjsElement* SpecEditor::GetRefElement() const { + return ref_element_; +} + +void SpecEditor::CommitChanges(mjsElement* element) { + if (element == nullptr) { + mju_warning("Element is null."); + return; + } + if (mjs_getSpec(element) != active_spec_.get()) { + mju_warning("Element is not owned by the active spec."); + return; + } + + AppendHistory({ + .spec = Copy(active_spec_.get()), + .op = kModify, + .key = active_element_key_, + .type_index = active_map_.LookupTypeIndex(active_element_key_), + }); +} + +void SpecEditor::Undo() { + if (CanUndo()) { + --cursor_; + + auto& entry = history_[cursor_]; + active_spec_ = Copy(entry.spec.get()); + if (entry.op == kAdd) { + active_map_.Remove(entry.key); + } else if (entry.op == kDelete) { + active_map_.Insert(entry.key, entry.type_index); + } + + active_element_ = active_map_.Resolve(active_spec_.get(), active_element_key_); + UpdateReferenceElement(); + } +} + +bool SpecEditor::CanUndo() const { + return cursor_ > 0; +} + +void SpecEditor::Redo() { + if (CanRedo()) { + ++cursor_; + + auto& entry = history_[cursor_]; + active_spec_ = Copy(entry.spec.get()); + if (entry.op == kAdd) { + active_map_.Insert(entry.key, entry.type_index); + } else if (entry.op == kDelete) { + active_map_.Remove(entry.key); + } + + active_element_ = active_map_.Resolve(active_spec_.get(), active_element_key_); + UpdateReferenceElement(); + } +} + +bool SpecEditor::CanRedo() const { + return cursor_ < history_.size() - 1; +} + +void SpecEditor::AppendHistory(HistoryEntry entry) { + ++cursor_; + while (history_.size() > cursor_) { + history_.pop_back(); + } + history_.push_back(std::move(entry)); + if (cursor_ > capacity_) { + history_.pop_front(); + --cursor_; + } +} + +int SpecEditor::ElementKeyMap::Append(mjtObj type, ElementKey key) { + const int index = keys_[type].size(); + keys_[type].push_back(key); + return index; +} + +int SpecEditor::ElementKeyMap::Remove(ElementKey key) { + const TypeIndex type_index = LookupTypeIndex(key); + auto& list = keys_[type_index.type]; + list.erase(list.begin() + type_index.index); + return type_index.index; +} + +void SpecEditor::ElementKeyMap::Insert(ElementKey key, TypeIndex type_index) { + auto& list = keys_[type_index.type]; + list.insert(list.begin() + type_index.index, key); +} + +SpecEditor::TypeIndex SpecEditor::ElementKeyMap::LookupTypeIndex( + ElementKey key) const { + for (int type = mjOBJ_UNKNOWN + 1; type < mjNOBJECT; ++type) { + const auto& list = keys_[type]; + for (int index = 0; index < list.size(); ++index) { + if (list[index] == key) { + return {static_cast(type), index}; + } + } + } + return kInvalidTypeIndex; +} + +SpecEditor::ElementKey SpecEditor::ElementKeyMap::LookupElementKey( + TypeIndex type_index) const { + if (type_index.type == mjOBJ_UNKNOWN) { + return kInvalidElementKey; + } + auto& list = keys_[type_index.type]; + if (type_index.index >= list.size()) { + return kInvalidElementKey; + } + return list[type_index.index]; +} + +mjsElement* SpecEditor::ElementKeyMap::Resolve(mjSpec* spec, + ElementKey key) const { + const TypeIndex type_index = LookupTypeIndex(key); + if (type_index.type == mjOBJ_UNKNOWN) { + return nullptr; + } + + mjsElement* element = mjs_firstElement(spec, type_index.type); + for (int i = 0; i < type_index.index && element; ++i) { + element = mjs_nextElement(spec, element); + } + return element; +} + +SpecEditor::SpecPtr SpecEditor::Copy(const mjSpec* spec) { + return SpecPtr(mj_copySpec(spec), mj_deleteSpec); +} + +mjsElement* SpecEditor::AddElementToSpec(mjSpec* spec, mjtObj type, mjsBody* body) { + if (spec == nullptr || type == mjOBJ_UNKNOWN) { + return nullptr; + } + if (body == nullptr) { + body = mjs_asBody(mjs_firstElement(spec, mjOBJ_BODY)); + } + + const mjsDefault* def = nullptr; + switch (type) { + case mjOBJ_ACTUATOR: + return mjs_addActuator(spec, def)->element; + case mjOBJ_EQUALITY: + return mjs_addEquality(spec, def)->element; + case mjOBJ_EXCLUDE: + return mjs_addExclude(spec)->element; + case mjOBJ_FLEX: + return mjs_addFlex(spec)->element; + case mjOBJ_HFIELD: + return mjs_addHField(spec)->element; + case mjOBJ_KEY: + return mjs_addKey(spec)->element; + case mjOBJ_MATERIAL: + return mjs_addMaterial(spec, def)->element; + case mjOBJ_MESH: + return mjs_addMesh(spec, def)->element; + case mjOBJ_NUMERIC: + return mjs_addNumeric(spec)->element; + case mjOBJ_PAIR: + return mjs_addPair(spec, def)->element; + case mjOBJ_SENSOR: + return mjs_addSensor(spec)->element; + case mjOBJ_SKIN: + return mjs_addSkin(spec)->element; + case mjOBJ_TENDON: + return mjs_addTendon(spec, def)->element; + case mjOBJ_TEXT: + return mjs_addText(spec)->element; + case mjOBJ_TEXTURE: + return mjs_addTexture(spec)->element; + case mjOBJ_TUPLE: + return mjs_addTuple(spec)->element; + case mjOBJ_PLUGIN: + return mjs_addPlugin(spec)->element; + case mjOBJ_BODY: + return mjs_addBody(body, def)->element; + case mjOBJ_SITE: + return mjs_addSite(body, def)->element; + case mjOBJ_JOINT: + return mjs_addJoint(body, def)->element; + case mjOBJ_GEOM: + return mjs_addGeom(body, def)->element; + case mjOBJ_CAMERA: + return mjs_addCamera(body, def)->element; + case mjOBJ_LIGHT: + return mjs_addLight(body, def)->element; + default: + mju_warning("Unsupported element type: %d", type); + return nullptr; + } +} + +} // namespace mujoco::platform diff --git a/src/experimental/platform/spec_editor.h b/src/experimental/platform/spec_editor.h new file mode 100644 index 00000000..b9338c94 --- /dev/null +++ b/src/experimental/platform/spec_editor.h @@ -0,0 +1,177 @@ +// 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 +// +// http://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_SPEC_EDITOR_H_ +#define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_SPEC_EDITOR_H_ + +#include +#include +#include +#include +#include + +#include +#include "experimental/platform/model_holder.h" + +namespace mujoco::platform { + +// Manages everything related to mjSpec editing including compiling, undo/redo, +// etc. +// +// The primary responsibility of the SpecEditor is to manage an "active" spec +// and element. This is the spec/element that is currently being edited. It +// also keeps track of the most recently compiled spec which is used as a +// reference spec for editing. Finally, it tracks a history of specs in +// order to support undo/redo operations. +class SpecEditor { + public: + explicit SpecEditor(int history_size = 256); + + // Resets the editor to manage the given spec, clearing all cached data. + // Assumes the given spec is correctly formed (i.e. compilable). + void Reset(const mjSpec& spec); + + // Attempts to compile the active spec, returning a ModelHolder. If the + // compilation was successfully, it will also update the reference spec to be + // the source spec. + std::unique_ptr Compile(); + + // Returns the active spec being edited. That said, users should not directly + // modify this spec. Instead, they should use the add/delete operations below. + mjSpec* GetActiveSpec() const; + + // Adds an element to the active spec. + mjsElement* AddElement(mjtObj type); + + // Adds an element to the given body of the active spec. + mjsElement* AddBodyElement(mjsBody* body, mjtObj type); + + // Deletes the currently active element from the active spec. + void DeleteActiveElement(); + + // Marks the given element as the active element being edited. This element + // must be owned by the active spec (i.e. the spec at the current point in + // the history buffer). Note that some operations (e.g. DeleteElement) will + // set the active element to nullptr. + void SetActiveElement(mjsElement* element); + + // Returns the element that is currently being edited. If this element is + // modified, then the caller should call CommitChanges() below (otherwise + // the change will not be saved in the history buffer). + mjsElement* GetActiveElement() const; + + // Commits any changes made to the given element. This effectively updates + // the history buffer to support undo/redo operations. + void CommitChanges(mjsElement* element); + + // Undoes the last change, making a new spec the active spec. + void Undo(); + bool CanUndo() const; + + // Redoes the last change, making a new spec the active spec. + void Redo(); + bool CanRedo() const; + + // Returns the reference element against which the active element is being + // edited. This is either an element from the "reference" spec that matches + // the active element, or, in the case of a newly added element, a default + // element of the same type. + mjsElement* GetRefElement() const; + + private: + // A unique pointer to an mjSpec that will be deleted by mj_deleteSpec. + using SpecPtr = std::unique_ptr; + + // A simple tuple of mjtObj type and int index. + struct TypeIndex { + mjtObj type; + int index; + }; + static constexpr TypeIndex kInvalidTypeIndex = {mjOBJ_UNKNOWN, 0}; + + // Every element encountered by this class is assigned a unique ElementKey. + // This allows us to track elements across copies of the spec. + using ElementKey = std::size_t; + static constexpr ElementKey kInvalidElementKey = 0; + + // A bidirectional mapping of ElementKey to TypeIndex. This allows us to + // efficiently track elements across copies of a spec. + // + // The map should be initialized from a base spec, assigning a unique key + // to every element. When elements are added or removed (usually from copies + // of the spec), this mapping can be updated to reflect the changes. We can + // then "resolve" an element from a spec based on its key. + class ElementKeyMap { + public: + int Append(mjtObj type, ElementKey key); + + int Remove(ElementKey key); + void Insert(ElementKey key, TypeIndex type_index); + + TypeIndex LookupTypeIndex(ElementKey key) const; + ElementKey LookupElementKey(TypeIndex type_index) const; + + // Returns the element in the spec that corresponds to the given key. + mjsElement* Resolve(mjSpec* spec, ElementKey key) const; + + private: + std::array, mjNOBJECT> keys_; + }; + + // The operation that was performed on the spec, used for undo/redo. + enum Operation { + kInitialize, + kModify, + kAdd, + kDelete, + }; + + // A single entry in the history buffer. + struct HistoryEntry { + SpecPtr spec; // A full copy of a spec. + Operation op; // The operation performed on the spec. + ElementKey key; // The key of the element being modified. + TypeIndex type_index; // The type/index of the element after the change. + }; + + // Appends a new HistoryEntry to the history buffer. + void AppendHistory(HistoryEntry entry); + + // Updates the reference element to be the element in the reference spec that + // corresponds the active element. + void UpdateReferenceElement(); + + // Adds an element to the spec with the given type and (optional) body. + static mjsElement* AddElementToSpec(mjSpec* spec, mjtObj type, + mjsBody* body = nullptr); + + // Creates a copy of the given spec. + static SpecPtr Copy(const mjSpec* spec); + + std::deque history_; + SpecPtr ref_spec_ = SpecPtr(nullptr, mj_deleteSpec); + SpecPtr active_spec_ = SpecPtr(nullptr, mj_deleteSpec); + SpecPtr dummy_spec_ = SpecPtr(mj_makeSpec(), mj_deleteSpec); + ElementKeyMap active_map_; + ElementKeyMap ref_map_; + mjsElement* active_element_ = nullptr; + mjsElement* ref_element_ = nullptr; + ElementKey active_element_key_ = kInvalidElementKey; + ElementKey next_element_key_ = 1; + int capacity_ = 0; + int cursor_ = 0; +}; +} // namespace mujoco::platform + +#endif // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_SPEC_EDITOR_H_ diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index 8bfcd879..137797f0 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -171,6 +171,7 @@ void App::RequestModelReload() { void App::InitEmptyModel() { model_holder_ = platform::ModelHolder::FromSpec(mj_makeSpec()); OnModelLoaded("", kEmptyModel); + spec_editor_.Reset(*spec()); } void App::LoadModelFromFile(const std::string& filepath) { @@ -179,6 +180,7 @@ void App::LoadModelFromFile(const std::string& filepath) { model_holder_ = platform::ModelHolder::FromFile(resolved_file); if (model_holder_->ok()) { OnModelLoaded(filepath, kModelFromFile); + spec_editor_.Reset(*spec()); UpdateFilePaths(resolved_file); window_->SetTitle("MuJoCo Studio : " + filepath); } else { @@ -196,6 +198,7 @@ void App::LoadModelFromBuffer(std::span buffer, } else { SetLoadError(std::string(model_holder_->error())); } + spec_editor_.Reset(*spec()); } void App::OnModelLoaded(std::string filename, ModelKind model_kind) { @@ -215,9 +218,6 @@ void App::OnModelLoaded(std::string filename, ModelKind model_kind) { const int state_size = mj_stateSize(model, mjSTATE_INTEGRATION); history_.Init(state_size); - // Create a copy of the spec for editing. - CopyLoadedSpecForEditing(); - if (!preserve_camera_on_load_) { const int model_cam = model->vis.global.cameraid; if (model_cam >= 0 && model_cam < model->ncam) { @@ -404,9 +404,9 @@ void App::ProcessPendingLoads() { } } - if (spec_op_) { - spec_op_(); - spec_op_ = nullptr; + if (pending_op_) { + pending_op_(); + pending_op_ = nullptr; } // Allow plugins to edit the spec as well. @@ -438,56 +438,6 @@ void App::ProcessPendingLoads() { }); } -void App::CopyLoadedSpecForEditing() { - if (scratch_spec_ != nullptr) { - mj_deleteSpec(scratch_spec_); - scratch_spec_ = nullptr; - } - - scratch_spec_modified_ = false; - tmp_.curr_edit_element = nullptr; - - if (has_spec()) { - scratch_spec_ = mj_copySpec(spec()); - - auto add_ref_elements = [&](mjtObj type) { - mjsElement* elem = mjs_firstElement(spec(), type); - mjsElement* scratch = mjs_firstElement(scratch_spec_, type); - while (elem != nullptr && scratch != nullptr) { - scratch_to_spec_[scratch] = elem; - spec_to_scratch_[elem] = scratch; - elem = mjs_nextElement(spec(), elem); - scratch = mjs_nextElement(scratch_spec_, scratch); - } - }; - add_ref_elements(mjOBJ_BODY); - add_ref_elements(mjOBJ_XBODY); - add_ref_elements(mjOBJ_JOINT); - add_ref_elements(mjOBJ_DOF); - add_ref_elements(mjOBJ_GEOM); - add_ref_elements(mjOBJ_SITE); - add_ref_elements(mjOBJ_CAMERA); - add_ref_elements(mjOBJ_LIGHT); - add_ref_elements(mjOBJ_FLEX); - add_ref_elements(mjOBJ_MESH); - add_ref_elements(mjOBJ_SKIN); - add_ref_elements(mjOBJ_HFIELD); - add_ref_elements(mjOBJ_TEXTURE); - add_ref_elements(mjOBJ_MATERIAL); - add_ref_elements(mjOBJ_PAIR); - add_ref_elements(mjOBJ_EXCLUDE); - add_ref_elements(mjOBJ_EQUALITY); - add_ref_elements(mjOBJ_TENDON); - add_ref_elements(mjOBJ_ACTUATOR); - add_ref_elements(mjOBJ_SENSOR); - add_ref_elements(mjOBJ_NUMERIC); - add_ref_elements(mjOBJ_TEXT); - add_ref_elements(mjOBJ_TUPLE); - add_ref_elements(mjOBJ_KEY); - add_ref_elements(mjOBJ_PLUGIN); - } -} - void App::HandleWindowEvents() { const std::string drop_file = window_->GetDropFile(); if (!drop_file.empty()) { @@ -1196,7 +1146,7 @@ void App::SpecExplorerGui() { bool open = element != nullptr; if (platform::ImGui_BeginHSplit("SpecExplorerTree", &tmp_.explorer_split, &open)) { - platform::SpecTreeGui(&element, spec(), SpecEditMode::kPlay); + platform::SpecTreeGui(&element, spec()); if (element != tmp_.curr_element) { tmp_.curr_element = element; @@ -1214,50 +1164,34 @@ void App::SpecExplorerGui() { if (platform::ImGui_HSplit("SpecExplorerProperties", &tmp_.explorer_split, &open)) { - platform::ScopedStyle style; - ImGui::Text("%s", mju_type2Str(tmp_.curr_element->elemtype)); ImGui::SameLine(); ImGui::Text("(%d)", mjs_getId(tmp_.curr_element)); ImGui::SameLine(ImGui::GetContentRegionAvail().x - 100.0f); - if (tmp_.spec_prop_mode == SpecPropertiesMode::kSpec) { - style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); - } - style.Var(ImGuiStyleVar_FramePadding, ImVec2(0, 0)); - if (ImGui::Button("S", ImVec2(24.0f, 20.0f))) { - tmp_.spec_prop_mode = SpecPropertiesMode::kSpec; - } - ImGui::SetItemTooltip("Spec"); - style.Reset(); + auto mode_button = [&](const char* label, const char* tooltip, + SpecPropertiesMode mode) { + platform::ScopedStyle style; + style.Var(ImGuiStyleVar_FramePadding, ImVec2(0, 0)); + if (tmp_.spec_prop_mode == mode) { + style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); + } + if (ImGui::Button(label, ImVec2(24.0f, 20.0f))) { + tmp_.spec_prop_mode = mode; + } + ImGui::SetItemTooltip("%s", tooltip); + }; + mode_button("S", "Spec", SpecPropertiesMode::kSpec); ImGui::SameLine(); - - if (tmp_.spec_prop_mode == SpecPropertiesMode::kModel) { - style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); - } - style.Var(ImGuiStyleVar_FramePadding, ImVec2(0, 0)); - if (ImGui::Button("M", ImVec2(24.0f, 20.0f))) { - tmp_.spec_prop_mode = SpecPropertiesMode::kModel; - } - ImGui::SetItemTooltip("Model"); - style.Reset(); - + mode_button("M", "Model", SpecPropertiesMode::kModel); ImGui::SameLine(); + mode_button("D", "Data", SpecPropertiesMode::kData); - if (tmp_.spec_prop_mode == SpecPropertiesMode::kData) { - style.Color(ImGuiCol_Button, ImGuiCol_ButtonActive); - } - style.Var(ImGuiStyleVar_FramePadding, ImVec2(0, 0)); - if (ImGui::Button("D", ImVec2(24.0f, 20.0f))) { - tmp_.spec_prop_mode = SpecPropertiesMode::kData; - } - ImGui::SetItemTooltip("Data"); - style.Reset(); ImGui::Separator(); if (tmp_.spec_prop_mode == SpecPropertiesMode::kSpec) { - platform::ElementSpecGui(element, element, SpecEditMode::kPlay); + platform::ElementSpecGui(element); } else if (tmp_.spec_prop_mode == SpecPropertiesMode::kModel) { platform::ElementModelGui(model(), tmp_.curr_element); } else { @@ -1274,36 +1208,32 @@ void App::SpecExplorerGui() { void App::SpecEditorGui() { if (ImGui::BeginChild("SpecEditor", ImVec2(-1, 36))) { if (ImGui::BeginTable("##SpecEditorHeader", 3)) { - ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, 70.0f); + ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, 50.0f); ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthStretch); ImGui::TableSetupColumn("", ImGuiTableColumnFlags_WidthFixed, 30.0f); ImGui::TableNextColumn(); - if (ImGui::Button(ICON_RELOAD_SPEC)) { - CopyLoadedSpecForEditing(); - } - ImGui::BeginDisabled(true); - ImGui::SameLine(); + ImGui::BeginDisabled(!spec_editor_.CanUndo()); if (ImGui::Button(ICON_UNDO_SPEC)) { - // TODO... + spec_editor_.Undo(); } + ImGui::EndDisabled(); ImGui::SameLine(); + ImGui::BeginDisabled(!spec_editor_.CanRedo()); if (ImGui::Button(ICON_REDO_SPEC)) { - // TODO... + spec_editor_.Redo(); } ImGui::EndDisabled(); ImGui::TableNextColumn(); ImGui::PushStyleColor(ImGuiCol_Button, ImColor(40, 180, 40, 255).Value); if (ImGui::Button("Compile and Reload", ImVec2(-1, 0))) { - spec_op_ = [this]() { - auto tmp_holder = platform::ModelHolder::FromSpec(scratch_spec_); + pending_op_ = [this]() { + auto tmp_holder = spec_editor_.Compile(); if (tmp_holder->ok()) { model_holder_ = std::move(tmp_holder); - scratch_spec_ = nullptr; OnModelLoaded(model_name_, model_kind_); } else { - scratch_spec_ = tmp_holder->ReleaseSpec(); load_error_ = std::move(tmp_holder->error()); } }; @@ -1315,30 +1245,29 @@ void App::SpecEditorGui() { ImGui::OpenPopupOnItemClick("SpecAddElement", 0); } if (ImGui::BeginPopupContextItem("SpecAddElement")) { - auto option = [&](const char* label, auto fn) { + auto option = [&](const char* label, mjtObj type) { if (ImGui::Selectable(label)) { - tmp_.curr_edit_element = fn()->element; - mjs_setName(tmp_.curr_edit_element, - platform::ElementName(tmp_.curr_edit_element).c_str()); - scratch_spec_modified_ = true; + mjsElement* element = spec_editor_.AddElement(type); + spec_editor_.SetActiveElement(element); } }; - option("Actuator", [&]() { return mjs_addActuator(scratch_spec_, 0); }); - option("Equality", [&]() { return mjs_addEquality(scratch_spec_, 0); }); - option("Exclude", [&]() { return mjs_addExclude(scratch_spec_); }); - option("Flex", [&]() { return mjs_addFlex(scratch_spec_); }); - option("Height Field", [&]() { return mjs_addHField(scratch_spec_); }); - option("Key", [&]() { return mjs_addKey(scratch_spec_); }); - option("Material", [&]() { return mjs_addMaterial(scratch_spec_, 0); }); - option("Mesh", [&]() { return mjs_addMesh(scratch_spec_, 0); }); - option("Numeric", [&]() { return mjs_addNumeric(scratch_spec_); }); - option("Pair", [&]() { return mjs_addPair(scratch_spec_, 0); }); - option("Sensor", [&]() { return mjs_addSensor(scratch_spec_); }); - option("Skin", [&]() { return mjs_addSkin(scratch_spec_); }); - option("Tendon", [&]() { return mjs_addTendon(scratch_spec_, 0); }); - option("Text", [&]() { return mjs_addText(scratch_spec_); }); - option("Texture", [&]() { return mjs_addTexture(scratch_spec_); }); - option("Tuple", [&]() { return mjs_addTuple(scratch_spec_); }); + + option("Actuator", mjOBJ_ACTUATOR); + option("Equality", mjOBJ_EQUALITY); + option("Exclude", mjOBJ_EXCLUDE); + option("Flex", mjOBJ_FLEX); + option("Height Field", mjOBJ_HFIELD); + option("Key", mjOBJ_KEY); + option("Material", mjOBJ_MATERIAL); + option("Mesh", mjOBJ_MESH); + option("Numeric", mjOBJ_NUMERIC); + option("Pair", mjOBJ_PAIR); + option("Sensor", mjOBJ_SENSOR); + option("Skin", mjOBJ_SKIN); + option("Tendon", mjOBJ_TENDON); + option("Text", mjOBJ_TEXT); + option("Texture", mjOBJ_TEXTURE); + option("Tuple", mjOBJ_TUPLE); ImGui::EndPopup(); } @@ -1354,32 +1283,28 @@ void App::SpecEditorGui() { } tmp_.editor_split = std::clamp(tmp_.editor_split, 20.0f, region.y - 40.0f); - bool open = tmp_.curr_edit_element != nullptr; + mjsElement* element = spec_editor_.GetActiveElement(); + + bool open = element != nullptr; if (platform::ImGui_BeginHSplit("SpecEditorTree", &tmp_.editor_split, &open)) { - if (platform::SpecTreeGui(&tmp_.curr_edit_element, scratch_spec_, - SpecEditMode::kEdit)) { - scratch_spec_modified_ = true; - } - open = tmp_.curr_edit_element != nullptr; + platform::SpecTreeGui(&element, spec_editor_.GetActiveSpec(), + &spec_editor_); + open = element != nullptr; + spec_editor_.SetActiveElement(element); if (platform::ImGui_HSplit("SpecEditorProperties", &tmp_.editor_split, &open)) { - ImGui::Text("%s", mju_type2Str(tmp_.curr_edit_element->elemtype)); + ImGui::Text("%s", mju_type2Str(element->elemtype)); ImGui::SameLine(); - ImGui::Text("(%d)", mjs_getId(tmp_.curr_edit_element)); + ImGui::Text("(%d)", mjs_getId(element)); ImGui::Separator(); - mjsElement* ref = scratch_to_spec_[tmp_.curr_edit_element]; - if (platform::ElementSpecGui(tmp_.curr_edit_element, ref, - SpecEditMode::kEdit)) { - scratch_spec_modified_ = true; - } + platform::ElementSpecGui(element, &spec_editor_); } platform::ImGui_EndHSplit(open); - if (!open) { - tmp_.curr_edit_element = nullptr; + spec_editor_.SetActiveElement(nullptr); } } } @@ -1939,7 +1864,7 @@ void App::MainMenuGui() { mode = platform::GraphicsMode::FilamentVulkan; } if (mode.has_value()) { - spec_op_ = [=, this]() { + pending_op_ = [=, this]() { const int width = window_->GetWidth(); const int height = window_->GetHeight(); SwitchGraphicsMode(width, height, *mode); diff --git a/src/experimental/studio/app.h b/src/experimental/studio/app.h index c11080fb..8edf1825 100644 --- a/src/experimental/studio/app.h +++ b/src/experimental/studio/app.h @@ -35,6 +35,7 @@ #include "experimental/platform/renderer.h" #include "experimental/platform/sim_history.h" #include "experimental/platform/sim_profiler.h" +#include "experimental/platform/spec_editor.h" #include "experimental/platform/step_control.h" #include "experimental/platform/window.h" @@ -84,8 +85,6 @@ class App { void Render(); private: - using SpecEditMode = platform::SpecEditMode; - // The kind of model that is currently loaded. enum ModelKind { kEmptyModel, @@ -148,7 +147,6 @@ class App { // Spec editing. SpecPropertiesMode spec_prop_mode = SpecPropertiesMode::kSpec; mjsElement* curr_element = nullptr; - mjsElement* curr_edit_element = nullptr; // State. int state_sig = 0; @@ -226,7 +224,6 @@ class App { float GetExpectedLabelWidth(); std::vector GetCameraNames(); - void CopyLoadedSpecForEditing(); mjSpec* spec() { return model_holder_->spec(); } mjModel* model() { return model_holder_->model(); } @@ -241,7 +238,9 @@ class App { std::string load_error_; std::string step_error_; std::string edit_error_; + std::optional pending_load_; + std::function pending_op_; bool preserve_camera_on_load_ = false; ModelKind model_kind_ = kEmptyModel; platform::GraphicsMode gfx_mode_ = platform::GraphicsMode::FilamentVulkan; @@ -250,21 +249,10 @@ class App { std::unique_ptr renderer_; std::unique_ptr model_holder_; - // Spec editing. We keep a separate copy of the loaded spec that we can edit. - // Once we're done editing, we will (re)compile the spec and update the - // active model and data. - mjSpec* scratch_spec_ = nullptr; - // Whether or not the scratch spec differs from the loaded spec. - bool scratch_spec_modified_ = false; - // We keep a mapping of the elements between the loaded spec and the scratch - // spec. - std::unordered_map spec_to_scratch_; - std::unordered_map scratch_to_spec_; - std::function spec_op_; - platform::StepControl step_control_; platform::SimProfiler profiler_; platform::SimHistory history_; + platform::SpecEditor spec_editor_; std::vector search_paths_; std::vector pixels_;