Refactor spec editing logic into its own class.

The SpecEditor improves spec editing by keeping
track of what's being edited and how it changes
over time.

PiperOrigin-RevId: 878410011
Change-Id: I7c323eeb30178dad31946095a727a21b3955f19a

src/experimental/platform/gui_spec.cc

@@ -21,6 +21,7 @@
 #include <mujoco/mjxmacro.h>
 #include <mujoco/mujoco.h>
 #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) {

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 <string>
-
 #include <mujoco/mujoco.h>
+#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.

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 <memory>
+#include <utility>
+#include <mujoco/mujoco.h>
+#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<mjtObj>(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<ModelHolder> 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<mjtObj>(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

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 <array>
+#include <cstddef>
+#include <deque>
+#include <memory>
+#include <vector>
+
+#include <mujoco/mujoco.h>
+#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<ModelHolder> 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<mjSpec, decltype(&mj_deleteSpec)>;
+
+  // 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<std::vector<ElementKey>, 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<HistoryEntry> 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_