Allow to attach an entire mjSpec in mjs_attach.

PiperOrigin-RevId: 742737713
Change-Id: Ie172568b8c35232e1f13d75aebee20346e76e95b

Author: quagla

doc/XMLreference.rst

@@ -3716,8 +3716,9 @@ all attachments will appear in the saved XML file.
 
 .. _body-attach-body:
 
-:at:`body`: :at-val:`string, required`
-   Name of the body in the sub-model to attach here. The body and its subtree will be attached.
+:at:`body`: :at-val:`string, optional`
+   Name of the body in the sub-model to attach here. The body and its subtree will be attached. If this attribute is not
+   specified, the contents of the world body will be attached in a new :ref:`frame<body-frame>`.
 
 .. _body-attach-prefix:
 

python/mujoco/specs.cc

@@ -269,16 +269,6 @@ py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype, bool recursive) {
   return list;  // list of pointers, so they can be copied
 }
 
-void SetFrame(raw::MjsBody* body, mjtObj objtype, raw::MjsFrame* frame) {
-  mjsElement* el = mjs_firstChild(body, objtype, 0);
-  while (el) {
-    if (frame->element != el && mjs_getFrame(el) == nullptr) {
-      mjs_setFrame(el, frame);
-    }
-    el = mjs_nextChild(body, el, 0);
-  }
-}
-
 PYBIND11_MODULE(_specs, m) {
   auto structs_m = py::module::import("mujoco._structs");
   py::function mjmodel_from_raw_ptr =
@@ -526,18 +516,6 @@ PYBIND11_MODULE(_specs, m) {
           throw pybind11::value_error(
               "Only one of frame or site can be specified.");
         }
-        auto worldbody = mjs_findBody(child.ptr, "world");
-        if (!worldbody) {
-          throw pybind11::value_error("Child does not have a world body.");
-        }
-        auto worldframe = mjs_addFrame(worldbody, nullptr);
-        SetFrame(worldbody, mjOBJ_BODY, worldframe);
-        SetFrame(worldbody, mjOBJ_SITE, worldframe);
-        SetFrame(worldbody, mjOBJ_FRAME, worldframe);
-        SetFrame(worldbody, mjOBJ_JOINT, worldframe);
-        SetFrame(worldbody, mjOBJ_GEOM, worldframe);
-        SetFrame(worldbody, mjOBJ_LIGHT, worldframe);
-        SetFrame(worldbody, mjOBJ_CAMERA, worldframe);
         const char* p = prefix.has_value() ? prefix.value().c_str() : "";
         const char* s = suffix.has_value() ? suffix.value().c_str() : "";
         raw::MjsElement* attached_frame = nullptr;
@@ -556,18 +534,11 @@ PYBIND11_MODULE(_specs, m) {
             throw pybind11::value_error(
                 "Frame spec does not match parent spec.");
           }
-          raw::MjsBody* parent_body = mjs_getParent(frame_ptr->element);
-          if (!parent_body) {
-            throw pybind11::value_error("Frame does not have a parent body.");
-          }
           attached_frame =
-              mjs_attach(parent_body->element, worldframe->element, p, s);
+              mjs_attach(frame_ptr->element, child.ptr->element, p, s);
           if (!attached_frame) {
             throw pybind11::value_error(mjs_getError(self.ptr));
           }
-          if (mjs_setFrame(attached_frame, frame_ptr) != 0) {
-            throw pybind11::value_error(mjs_getError(self.ptr));
-          }
         }
         if (site.has_value()) {
           raw::MjsSite* site_ptr = nullptr;
@@ -585,7 +556,7 @@ PYBIND11_MODULE(_specs, m) {
                 "Site spec does not match parent spec.");
           }
           attached_frame =
-              mjs_attach(site_ptr->element, worldframe->element, p, s);
+              mjs_attach(site_ptr->element, child.ptr->element, p, s);
           if (!attached_frame) {
             throw pybind11::value_error(mjs_getError(self.ptr));
           }

src/user/user_api.cc

@@ -119,6 +119,18 @@ mjModel* mj_compile(mjSpec* s, const mjVFS* vfs) {
 }
 
 
+// set frame for all elements of a body
+static void SetFrame(mjsBody* body, mjtObj objtype, mjsFrame* frame) {
+  mjsElement* el = mjs_firstChild(body, objtype, 0);
+  while (el) {
+    if (frame->element != el && mjs_getFrame(el) == nullptr) {
+      mjs_setFrame(el, frame);
+    }
+    el = mjs_nextChild(body, el, 0);
+  }
+}
+
+
 
 // attach body to a frame of the parent
 static mjsElement* attachBody(mjCFrame* parent, const mjCBody* child,
@@ -207,13 +219,44 @@ mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child,
     return nullptr;
   }
   mjCModel* model = static_cast<mjCModel*>(mjs_getSpec(parent)->element);
+  if (child->elemtype == mjOBJ_MODEL) {
+    mjCModel* child_model = static_cast<mjCModel*>((mjsElement*)child);
+    mjsBody* worldbody = mjs_findBody(&child_model->spec, "world");
+    if (!worldbody) {
+      model->SetError(mjCError(0, "Child does not have a world body."));
+      return nullptr;
+    }
+    mjsFrame* worldframe = mjs_addFrame(worldbody, nullptr);
+    SetFrame(worldbody, mjOBJ_BODY, worldframe);
+    SetFrame(worldbody, mjOBJ_SITE, worldframe);
+    SetFrame(worldbody, mjOBJ_FRAME, worldframe);
+    SetFrame(worldbody, mjOBJ_JOINT, worldframe);
+    SetFrame(worldbody, mjOBJ_GEOM, worldframe);
+    SetFrame(worldbody, mjOBJ_LIGHT, worldframe);
+    SetFrame(worldbody, mjOBJ_CAMERA, worldframe);
+    child = worldframe->element;
+  }
   switch (parent->elemtype) {
     case mjOBJ_FRAME:
       if (child->elemtype == mjOBJ_BODY) {
         return attachBody(static_cast<mjCFrame*>(parent),
                           static_cast<const mjCBody*>(child), prefix, suffix);
+      } else if (child->elemtype == mjOBJ_FRAME) {
+        mjsBody* parent_body = mjs_getParent(parent);
+        if (!parent_body) {
+          model->SetError(mjCError(0, "Frame does not have a parent body."));
+          return nullptr;
+        }
+        mjCFrame* frame = static_cast<mjCFrame*>(parent);
+        mjsElement* attached_frame =
+            attachFrame(static_cast<mjCBody*>(parent_body->element),
+                        static_cast<const mjCFrame*>(child), prefix, suffix);
+        if (mjs_setFrame(attached_frame, &frame->spec)) {
+          return nullptr;
+        }
+        return attached_frame;
       } else {
-        model->SetError(mjCError(0, "child element is not a body"));
+        model->SetError(mjCError(0, "child element is not a body or frame"));
         return nullptr;
       }
     case mjOBJ_BODY:

src/xml/xml_native_reader.cc

@@ -3627,27 +3627,33 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame,
     else if (name == "attach") {
       string model_name, body_name, prefix;
       ReadAttrTxt(elem, "model", model_name, /*required=*/true);
-      ReadAttrTxt(elem, "body", body_name, /*required=*/true);
+      ReadAttrTxt(elem, "body", body_name, /*required=*/false);
       ReadAttrTxt(elem, "prefix", prefix, /*required=*/true);
 
-      mjsBody* child = mjs_findBody(spec, (prefix+body_name).c_str());
+      mjsBody* child_body = mjs_findBody(spec, (prefix+body_name).c_str());
       mjsFrame* pframe = frame ? frame : mjs_addFrame(body, nullptr);
 
-      if (!child) {
+      if (!child_body) {
         mjSpec* asset = mjs_findSpec(spec, model_name.c_str());
         if (!asset) {
           throw mjXError(elem, "could not find model '%s'", model_name.c_str());
         }
-        child = mjs_findBody(asset, body_name.c_str());
-        if (!child) {
-          throw mjXError(elem, "could not find body '%s''%s'", body_name.c_str());
+        mjsElement* child;
+        if (body_name.empty()) {
+          child = asset->element;
+        } else {
+          child_body = mjs_findBody(asset, body_name.c_str());
+          if (!child_body) {
+            throw mjXError(elem, "could not find body '%s''%s'", body_name.c_str());
+          }
+          child = child_body->element;
         }
-        if (!mjs_attach(pframe->element, child->element, prefix.c_str(), "")) {
+        if (!mjs_attach(pframe->element, child, prefix.c_str(), "")) {
           throw mjXError(elem, mjs_getError(spec));
         }
       } else {
         // only set frame to existing body
-        if (mjs_setFrame(child->element, pframe)) {
+        if (mjs_setFrame(child_body->element, pframe)) {
           throw mjXError(elem, mjs_getError(spec));
         }
       }