Reclaim deleted mjSpec elements automatically via reference counting

Fixes #2882.

Previously, `mjCModel::operator-=(mjsElement*)` pushed deleted elements onto `detached_` without calling `Release()`, deferring reclamation until `~mjCModel()` (`mj_deleteSpec`) so Python handles to deleted elements would not become dangling pointers when passed to `model.bind` / `data.bind`. In workflows that repeatedly added and deleted elements or subtrees on a long-lived `mjSpec`, memory grew unbounded.

- Remove `mjCModel::detached_` and call `Release()` immediately when removing elements in `mjCModel::operator-=(mjsElement*)`.
- Set `id = -1` before `Release()` in `RemoveFromList`, `RemovePlugins`, and `operator-=(const mjCDef&)` so implicitly removed elements also report invalid IDs.
- Add intrusive reference counting (`AddRef` / `Release`) to `mjCDef` and use `Release()` on `mjCWrap` in `~mjCTendon()`.
- Add `MjsElementRef<T>` custom pybind11 holder for `Mjs*` wrapper types in `specs.cc` using internal `_mjPRIVATE_addRefElement` / `_mjPRIVATE_releaseElement` entry points so C++ elements stay alive while referenced by Python variables and are freed automatically upon Python GC.

PiperOrigin-RevId: 985105555
Change-Id: I8ffc0064ff56920c96135feb73e75cf6073d12e6

Author: yuvaltassa

doc/changelog.rst

@@ -57,6 +57,8 @@ Bug fixes
   register the forces that tendons apply to bodies (:issue:`832`).
 - :ref:`mjd_transitionFD` and :ref:`mjd_inverseFD` now raise an error when :ref:`sleeping<Sleeping>` is enabled.
   Previously, their repeated evaluations changed the sleep state, leading to internal errors or wrong derivatives.
+- Fixed a memory leak in :ref:`mjs_delete` where deleted elements and subtrees were retained in ``mjSpec`` until
+  :ref:`mj_deleteSpec` (:issue:`2882`).
 - Fixed an out-of-bounds read when parsing the header of a :ref:`GMSH file<gmsh-file-docs>` loaded by
   :ref:`flexcomp<body-flexcomp>`. Truncated headers are now reported as an error.
 

python/mujoco/private.h

@@ -17,6 +17,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
+#include <mujoco/mjspec.h>
 #include <mujoco/mjtype.h>
 
 // DO NOT USE THESE FUNCTIONS ELSEWHERE.
@@ -25,6 +26,8 @@ extern "C" {
 MJAPI mjfLogHandler _mjPRIVATE_setTlsLogHandler(mjfLogHandler handler);
 MJAPI mjfLogHandler _mjPRIVATE_getGlobalLogHandler(void);
 MJAPI void* mj_arenaAllocByte(mjData* d, int bytes, int alignment);
+MJAPI void _mjPRIVATE_addRefElement(mjsElement* element);
+MJAPI void _mjPRIVATE_releaseElement(mjsElement* element);
 }
 
 #endif  // MUJOCO_PYTHON_PRIVATE_H_

python/mujoco/specs.cc

@@ -2380,5 +2380,47 @@ def test_bind_rejects_stale_model_after_first_equality(self):
     model = spec.compile()
     self.assertEqual(model.bind(equality).id, 0)
 
+  def test_delete_reclaims_memory_and_invalidates_bind(self):
+    spec = mujoco.MjSpec()
+    for i in range(100):
+      body = spec.worldbody.add_body(name=f'body_{i}')
+      joint = body.add_joint(
+          name=f'joint_{i}', type=mujoco.mjtJoint.mjJNT_HINGE
+      )
+      body.add_geom(size=[0.1, 0.1, 0.1])
+      act = spec.add_actuator(
+          name=f'act_{i}',
+          target=f'joint_{i}',
+          trntype=mujoco.mjtTrn.mjTRN_JOINT,
+      )
+      spec.delete(body)
+
+    model = spec.compile()
+    data = mujoco.MjData(model)
+    self.assertEqual(model.nbody, 1)
+    self.assertEqual(model.nu, 0)
+
+    # A Python reference to a deleted body or implicitly deleted actuator
+    # remains memory-safe for attribute access and raises IndexError on bind.
+    self.assertEqual(body.name, 'body_99')
+    self.assertIsNone(body.parent)
+    self.assertEqual(joint.name, 'joint_99')
+    self.assertEqual(act.name, 'act_99')
+    with self.assertRaisesRegex(IndexError, 'Invalid index -1'):
+      data.bind(body)
+    with self.assertRaisesRegex(IndexError, 'Invalid index -1'):
+      data.bind(joint)
+    with self.assertRaisesRegex(IndexError, 'Invalid index -1'):
+      data.bind(act)
+
+    # A child obtained via spec.joints (whose reference_internal parent is spec,
+    # not the body) stays alive with parent cleared to None once its body is
+    # deleted and freed.
+    spec.worldbody.add_body(name='temp_body').add_joint(name='orphan_joint')
+    orphan = spec.joints[-1]
+    spec.delete(spec.bodies[-1])
+    self.assertEqual(orphan.name, 'orphan_joint')
+    self.assertIsNone(orphan.parent)
+
 if __name__ == '__main__':
   absltest.main()

python/mujoco/specs_test.py

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "private.h"
 #include "raw.h"
 #include <pybind11/cast.h>
 #include <pybind11/pybind11.h>
@@ -22,6 +23,69 @@ namespace py = ::pybind11;
 
 namespace mujoco::python {
 
+template <typename T>
+class MjsElementRef {
+ public:
+  MjsElementRef() : ptr_(nullptr), element_(nullptr) {}
+  explicit MjsElementRef(T* ptr) : ptr_(ptr), element_(GetElement(ptr)) {
+    if (element_) {
+      _mjPRIVATE_addRefElement(element_);
+    }
+  }
+  MjsElementRef(const MjsElementRef& other)
+      : ptr_(other.ptr_), element_(other.element_) {
+    if (element_) {
+      _mjPRIVATE_addRefElement(element_);
+    }
+  }
+  MjsElementRef(MjsElementRef&& other) noexcept
+      : ptr_(other.ptr_), element_(other.element_) {
+    other.ptr_ = nullptr;
+    other.element_ = nullptr;
+  }
+  MjsElementRef& operator=(const MjsElementRef& other) {
+    if (this != &other) {
+      if (other.element_) {
+        _mjPRIVATE_addRefElement(other.element_);
+      }
+      if (element_) {
+        _mjPRIVATE_releaseElement(element_);
+      }
+      ptr_ = other.ptr_;
+      element_ = other.element_;
+    }
+    return *this;
+  }
+  MjsElementRef& operator=(MjsElementRef&& other) noexcept {
+    if (this != &other) {
+      if (element_) {
+        _mjPRIVATE_releaseElement(element_);
+      }
+      ptr_ = other.ptr_;
+      element_ = other.element_;
+      other.ptr_ = nullptr;
+      other.element_ = nullptr;
+    }
+    return *this;
+  }
+  ~MjsElementRef() {
+    if (element_) {
+      _mjPRIVATE_releaseElement(element_);
+    }
+  }
+  T* get() const { return ptr_; }
+
+ private:
+  static raw::MjsElement* GetElement(raw::MjsElement* ptr) { return ptr; }
+  template <typename U>
+  static raw::MjsElement* GetElement(U* ptr) {
+    return ptr ? ptr->element : nullptr;
+  }
+
+  T* ptr_;
+  raw::MjsElement* element_;
+};
+
 struct MjSpec {
   MjSpec();
   MjSpec(raw::MjSpec* ptr, const py::dict& assets_ = {});
@@ -43,3 +107,5 @@ struct MjSpec {
   MjSpec* parent = nullptr;
 };
 }  // namespace mujoco::python
+
+PYBIND11_DECLARE_HOLDER_TYPE(T, mujoco::python::MjsElementRef<T>, true);

python/mujoco/specs_wrapper.h

@@ -1738,19 +1738,15 @@ mjsBody* mjs_findChild(const mjsBody* bodyspec, const char* name) {
 mjsBody* mjs_getParent(const mjsElement* element) {
   switch (element->elemtype) {
     case mjOBJ_BODY:
-      return &(static_cast<const mjCBody*>(element)->GetParent()->spec);
     case mjOBJ_FRAME:
-      return &(static_cast<const mjCFrame*>(element)->GetParent()->spec);
     case mjOBJ_JOINT:
-      return &(static_cast<const mjCJoint*>(element)->GetParent()->spec);
     case mjOBJ_GEOM:
-      return &(static_cast<const mjCGeom*>(element)->GetParent()->spec);
     case mjOBJ_SITE:
-      return &(static_cast<const mjCSite*>(element)->GetParent()->spec);
     case mjOBJ_CAMERA:
-      return &(static_cast<const mjCCamera*>(element)->GetParent()->spec);
-    case mjOBJ_LIGHT:
-      return &(static_cast<const mjCLight*>(element)->GetParent()->spec);
+    case mjOBJ_LIGHT: {
+      mjCBody* parent = static_cast<mjCBody*>(static_cast<const mjCBase*>(element)->GetParent());
+      return parent ? &parent->spec : nullptr;
+    }
     default:
       return nullptr;
   }
@@ -2621,3 +2617,29 @@ void mjs_setAuthored(const void* elem_ptr, const void* field_ptr, int authored)
 #undef SET_FIELD
 #undef SET_FIELD_VEC
 }
+
+
+// increment reference count of element (for internal bindings use only)
+void _mjPRIVATE_addRefElement(mjsElement* element) {
+  if (!element) { return; }
+  if (element->elemtype == mjOBJ_DEFAULT) {
+    static_cast<mjCDef*>(element)->AddRef();
+    return;
+  }
+  // embedded sub-elements inside mjCDef are by-value members with model == nullptr
+  mjCBase* base = static_cast<mjCBase*>(element);
+  if (base->model) { base->AddRef(); }
+}
+
+
+// decrement reference count of element (for internal bindings use only)
+void _mjPRIVATE_releaseElement(mjsElement* element) {
+  if (!element) { return; }
+  if (element->elemtype == mjOBJ_DEFAULT) {
+    static_cast<mjCDef*>(element)->Release();
+    return;
+  }
+  // embedded sub-elements inside mjCDef are by-value members with model == nullptr
+  mjCBase* base = static_cast<mjCBase*>(element);
+  if (base->model) { base->Release(); }
+}

src/user/user_api.cc

@@ -620,6 +620,10 @@ MJAPI void mj_clearCache(mjCache* cache);
 // Get the internal asset cache used by the compiler.
 MJAPI mjCache* mj_getCache(void);
 
+// Internal element reference counting for language bindings.
+MJAPI void _mjPRIVATE_addRefElement(mjsElement* element);
+MJAPI void _mjPRIVATE_releaseElement(mjsElement* element);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif

src/user/user_api.h

@@ -508,6 +508,7 @@ void mjCModel::RemoveFromList(std::vector<T*>& list, const mjCModel& other) {
       element->ResolveReferences(this);
     } catch (mjCError err) {
       ids[element->elemtype].erase(element->name);
+      element->id = -1;
       element->Release();
       list.erase(list.begin() + i);
       nlist--;
@@ -554,6 +555,7 @@ void mjCModel::RemovePlugins() {
     if (plugins_[i]->name.empty()) { continue; }
     if (instances.find(plugins_[i]->name) == instances.end()) {
       ids[plugins_[i]->elemtype].erase(plugins_[i]->name);
+      plugins_[i]->id = -1;
       plugins_[i]->Release();
       plugins_.erase(plugins_.begin() + i);
       nlist--;
@@ -659,7 +661,8 @@ mjCModel& mjCModel::operator-=(const mjCDef& subtree) {
   std::sort(default_ids_to_remove.begin(), default_ids_to_remove.end(), std::greater<int>());
 
   for (int id : default_ids_to_remove) {
-    delete defaults_[id];
+    defaults_[id]->id = -1;
+    defaults_[id]->Release();
     defaults_.erase(defaults_.begin() + id);
   }
 
@@ -693,14 +696,6 @@ void mjCModel::DeleteSubtreePlugin(mjCBody* subtree) {
 
 // remove the element from the model
 void mjCModel::operator-=(mjsElement* el) {
-  if (el->elemtype == mjOBJ_BODY) {
-    mjCBody* body  = static_cast<mjCBody*>(el);
-    *this         -= *body;
-  }
-
-  detached_.push_back(static_cast<mjCBase*>(el));
-  ResetTreeLists();
-
   if (el->elemtype != mjOBJ_DEFAULT) {
     if (static_cast<mjCBase*>(el)->model != this) {
       throw mjCError(nullptr, "element is not in this model");
@@ -711,6 +706,13 @@ void mjCModel::operator-=(mjsElement* el) {
     }
   }
 
+  if (el->elemtype == mjOBJ_BODY) {
+    mjCBody* body  = static_cast<mjCBody*>(el);
+    *this         -= *body;
+  }
+
+  ResetTreeLists();
+
   switch (el->elemtype) {
     case mjOBJ_BODY: {
       MakeTreeLists();  // rebuild lists that were reset at the beginning of the function
@@ -783,6 +785,8 @@ void mjCModel::operator-=(mjsElement* el) {
 
   // update signature after we updated everything
   spec.element->signature = Signature();
+
+  static_cast<mjCBase*>(el)->Release();
 }
 
 
@@ -1071,10 +1075,9 @@ mjCModel::~mjCModel() {
   for (int i = 0; i < texts_.size(); i++) texts_[i]->Release();
   for (int i = 0; i < tuples_.size(); i++) tuples_[i]->Release();
   for (int i = 0; i < keys_.size(); i++) keys_[i]->Release();
-  for (int i = 0; i < defaults_.size(); i++) delete defaults_[i];
+  for (int i = 0; i < defaults_.size(); i++) defaults_[i]->Release();
   for (int i = 0; i < specs_.size(); i++) mj_deleteSpec(specs_[i]);
   for (int i = 0; i < plugins_.size(); i++) plugins_[i]->Release();
-  for (int i = 0; i < detached_.size(); i++) detached_[i]->Release();
 
   // clear sizes and pointer lists created in Compile
   Clear();

src/user/user_model.cc

@@ -535,6 +535,5 @@ class mjCModel : public mjCModel_, private mjSpec {
   bool                   deepcopy_;          // copy objects when attaching
   bool                   attached_ = false;  // true if model is attached to a parent model
   std::unordered_map<const mjsCompiler*, mjSpec*> compiler2spec_;  // map from compiler to spec
-  std::vector<mjCBase*>                           detached_;       // list of detached objects
 };
 #endif  // MUJOCO_SRC_USER_USER_MODEL_H_

src/user/user_model.h

@@ -1793,6 +1793,8 @@ void mjCBody::CopyList(std::vector<T*>&          dst,
 mjCBody& mjCBody::operator-=(const mjCBody& subtree) {
   for (int i = 0; i < bodies.size(); i++) {
     if (bodies[i] == &subtree) {
+      bodies[i]->SetParent(nullptr);
+      bodies[i]->frame = nullptr;
       bodies.erase(bodies.begin() + i);
       break;
     }
@@ -1889,13 +1891,20 @@ void mjCBody::CopyPlugin() {
 
 // destructor
 mjCBody::~mjCBody() {
-  for (int i = 0; i < bodies.size(); i++) bodies[i]->Release();
-  for (int i = 0; i < geoms.size(); i++) geoms[i]->Release();
-  for (int i = 0; i < frames.size(); i++) frames[i]->Release();
-  for (int i = 0; i < joints.size(); i++) joints[i]->Release();
-  for (int i = 0; i < sites.size(); i++) sites[i]->Release();
-  for (int i = 0; i < cameras.size(); i++) cameras[i]->Release();
-  for (int i = 0; i < lights.size(); i++) lights[i]->Release();
+  auto release_children = [](auto& list) {
+    for (auto* child : list) {
+      child->SetParent(nullptr);
+      child->frame = nullptr;
+      child->Release();
+    }
+  };
+  release_children(bodies);
+  release_children(geoms);
+  release_children(frames);
+  release_children(joints);
+  release_children(sites);
+  release_children(cameras);
+  release_children(lights);
 }
 
 
@@ -6175,7 +6184,7 @@ void mjCTendon::CopyFromSpec() {
 // desctructor
 mjCTendon::~mjCTendon() {
   // delete objects allocated here
-  for (unsigned int i = 0; i < path.size(); i++) { delete path[i]; }
+  for (unsigned int i = 0; i < path.size(); i++) { path[i]->Release(); }
 
   path.clear();
 }