Allow self-attach in MJCF

PiperOrigin-RevId: 941676717
Change-Id: I9ff8fc89716bc14b438a41e2b4075bbe5fdb83cf

Author: yuvaltassa

doc/XMLreference.rst

@@ -3996,16 +3996,18 @@ Associate this body with an :ref:`engine plugin<exPlugin>`. Either :at:`plugin`
 :el-prefix:`body/` |-| **attach** |*|
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-The :el:`attach` element is used to insert elements from another (child) model into this (parent) model's kinematic tree.
-Unlike :ref:`include<include>`, which is implemented in the parser and is equivalent to copying and pasting XML from one
-file into another, :el:`attach` is implemented in the model compiler. In order to use this element, the sub-model must
-first be defined as an :ref:`asset<asset-model>`. When creating an attachment, a frame, body or the entire child model in the
-child model is specified, and all referencing elements outside the kinematic tree (e.g., sensors and actuators), are also copied into
-the parent model. Additionally, any elements referenced from within the attached subtree (e.g. defaults and assets)
-will be copied in to the parent model. :el:`attach` is a :ref:`meta-element`, so upon saving all attachments will
-appear in the saved XML file. Note that this element is a subset of the functionality of the procedural
-:ref:`attachment<meAttachment>` functionality. As such, it shares the same limitations as described there. See example `here
-<https://github.com/google-deepmind/mujoco/blob/main/test/xml/testdata/parent.xml>`__.
+The :el:`attach` element is used to insert elements from another (child) model, or from the current model itself
+(self-attachment), into this (parent) model's kinematic tree. Unlike :ref:`include<include>`, which is implemented in
+the parser and is equivalent to copying and pasting XML from one file into another, :el:`attach` is implemented in the
+model compiler. In order to use this element to import from another model, the sub-model must first be defined as an
+:ref:`asset<asset-model>`. When creating an attachment, a frame, body or the entire child model in the child model is
+specified, and all referencing elements outside the kinematic tree (e.g., sensors and actuators), are also copied into
+the parent model. Additionally, any elements referenced from within the attached subtree (e.g. defaults and assets) will
+be copied in to the parent model. For self-attaching within the same model, the :at:`model` attribute is omitted, and a
+body or frame must be specified. :el:`attach` is a :ref:`meta-element`, so upon saving all attachments will appear in
+the saved XML file. Note that this element is a subset of the functionality of the procedural
+:ref:`attachment<meAttachment>` functionality. As such, it shares the same limitations as described there. See example
+`here <https://github.com/google-deepmind/mujoco/blob/main/test/xml/testdata/parent.xml>`__.
 
 .. admonition:: Known issues
    :class: note
@@ -4020,8 +4022,9 @@ appear in the saved XML file. Note that this element is a subset of the function
 
 .. _body-attach-model:
 
-:at:`model`: :at-val:`string, required`
+:at:`model`: :at-val:`string, optional`
    The child model from which to attach a subtree or a frame.
+   If omitted, the attachment is performed within the current model (self-attachment).
 
 .. _body-attach-body:
 

doc/changelog.rst

@@ -9,8 +9,10 @@ General
 ^^^^^^^
 - Added Nesterov momentum extrapolation with adaptive gradient restart (O'Donoghue-Candès) to the PGS solver,
   significantly improving convergence. Overall PGS now requires ~2x fewer iterations.
-
-* :ref:`mj_encode` now supports encoding of MJB and TXT files.
+- :ref:`mj_encode` now supports encoding of MJB and TXT files.
+- The :el:`attach` element now supports self-attachment (attaching elements of the current model to itself) by omitting
+  the :at:`model` attribute. It also supports attaching a frame via the new :at:`frame` attribute, which is mutually
+  exclusive with :at:`body`.
 
 .. admonition:: Breaking API changes
    :class: attention

src/xml/xml_native_reader.cc

@@ -3942,39 +3942,69 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame,
     // attachment
     else if (name == "attach") {
       string model_name, child_name, prefix;
-      ReadAttrTxt(elem, "model", model_name, /*required=*/true);
+      bool has_model = ReadAttrTxt(elem, "model", model_name, /*required=*/false);
       bool has_body = ReadAttrTxt(elem, "body", child_name, /*required=*/false);
       bool has_frame = ReadAttrTxt(elem, "frame", child_name, /*required=*/false);
       ReadAttrTxt(elem, "prefix", prefix, /*required=*/true);
 
       if (has_body && has_frame) {
         throw mjXError(elem, "only one of body or frame can be specified in attach");
       }
-      mjtObj type = has_body ? mjOBJ_BODY : mjOBJ_FRAME;
+      mjtObj type = mjOBJ_UNKNOWN;
+      if (has_body) type = mjOBJ_BODY;
+      else if (has_frame) type = mjOBJ_FRAME;
+
+      mjsElement* source_elem = nullptr;
+      if (!has_model) { // Self-attach
+        if (type == mjOBJ_UNKNOWN) {
+          throw mjXError(elem, "either 'body' or 'frame' attribute must be specified for self-attach");
+        }
 
-      string full_name = prefix+child_name;
-      if (mjs_findElement(spec, type, full_name.c_str())) {
-        throw mjXError(elem, "cannot attach: element %s already exists", full_name.c_str());
-      }
+        // check for name collision in the current spec
+        string full_name = prefix + child_name;
+        if (mjs_findElement(spec, type, full_name.c_str())) {
+          throw mjXError(elem, "cannot self-attach: element %s already exists", full_name.c_str());
+        }
+        source_elem = mjs_findElement(spec, type, child_name.c_str());
+        if (!source_elem) {
+          throw mjXError(elem, "%s",
+                         (string("could not find ") + mju_type2Str(type) + " '" + child_name +
+                          "' in the current model for self-attachment").c_str());
+        }
+      } else { // Attach from external model asset
+        // Check for name collision in the current spec
+        if (!child_name.empty()) {
+            string full_name = prefix + child_name;
+            if (mjs_findElement(spec, type, full_name.c_str())) {
+                throw mjXError(elem, "%s",
+                               (string("cannot attach: element ") + child_name +
+                                " already exists with prefix " + prefix).c_str());
+            }
+        }
 
-      mjSpec* asset = mjs_findSpec(spec, model_name.c_str());
-      if (!asset) {
-        throw mjXError(elem, "could not find model '%s'", model_name.c_str());
-      }
+        mjSpec* asset = mjs_findSpec(spec, model_name.c_str());
+        if (!asset) {
+          throw mjXError(elem, "could not find model '%s'", model_name.c_str());
+        }
 
-      mjsElement* child;
-      if (child_name.empty()) {
-        child = asset->element;
-      } else {
-        child = mjs_findElement(asset, type, child_name.c_str());
-        if (!child) {
-          throw mjXError(elem, "could not find %s",
-                          (string(mju_type2Str(type)) + " '" + child_name + "'").c_str());
+        if (type == mjOBJ_UNKNOWN) { // Attach world body contents
+          source_elem = asset->element;
+        } else { // Attach specific body or frame
+          source_elem = mjs_findElement(asset, type, child_name.c_str());
+          if (!source_elem) {
+            throw mjXError(elem, "%s",
+                           (string("could not find ") + mju_type2Str(type) + " '" + child_name +
+                            "' in model asset '" + model_name + "'").c_str());
+          }
         }
       }
 
       mjsFrame* pframe = frame ? frame : mjs_addFrame(body, nullptr);
-      if (!mjs_attach(pframe->element, child, prefix.c_str(), "")) {
+      // Set default for the new frame from the current context
+      mjs_setDefault(pframe->element, mjs_getDefault(frame ? frame->element : body->element));
+      mjs_setString(pframe->info, ("line = " + std::to_string(elem->GetLineNum())).c_str());
+
+      if (!mjs_attach(pframe->element, source_elem, prefix.c_str(), "")) {
         throw mjXError(elem, "%s", stripError(mjs_getError(spec)));
       }
     }

test/xml/xml_native_reader_test.cc

@@ -521,7 +521,8 @@ TEST_F(XMLReaderTest, InvalidDoubleOrientation) {
         if (orient1 == orient2) continue;
         std::string xml = prefix + field + orient1 + orient2 + suffix;
         std::array<char, 1024> error;
-        MjModelPtr model = LoadModelFromString(xml.c_str(), error.data(), error.size());
+        MjModelPtr model =
+            LoadModelFromString(xml.c_str(), error.data(), error.size());
         ASSERT_THAT(model.get(), IsNull());
         EXPECT_THAT(
             error.data(),
@@ -1836,7 +1837,8 @@ TEST_F(XMLReaderTest, AttachSpecAssets) {
   mj_addBufferVFS(vfs.get(), "xml_child.xml", xml_child, sizeof(xml_child));
 
   std::array<char, 1024> er;
-  MjModelPtr model = LoadModelFromString(xml_parent, er.data(), er.size(), vfs.get());
+  MjModelPtr model =
+      LoadModelFromString(xml_parent, er.data(), er.size(), vfs.get());
   EXPECT_THAT(model.get(), NotNull()) << er.data();
 
   MjModelPtr expected = LoadModelFromString(xml_expected, er.data(), er.size());
@@ -1890,7 +1892,8 @@ TEST_F(XMLReaderTest, InvalidAttach) {
   mj_addBufferVFS(vfs.get(), "child.xml", xml_child, sizeof(xml_child));
 
   std::array<char, 1024> er;
-  MjModelPtr model = LoadModelFromString(xml_parent, er.data(), er.size(), vfs.get());
+  MjModelPtr model =
+      LoadModelFromString(xml_parent, er.data(), er.size(), vfs.get());
 
   EXPECT_THAT(model.get(), IsNull()) << er.data();
   EXPECT_THAT(er.data(), HasSubstr("repeated name '_actuator' in actuator"));
@@ -1999,7 +2002,8 @@ TEST_F(XMLReaderTest, ResizeKeyframeAfterParsing) {
   mj_addBufferVFS(vfs.get(), "child.xml", child_xml, sizeof(child_xml));
 
   std::array<char, 1024> error;
-  MjModelPtr m = LoadModelFromString(parent_xml, error.data(), error.size(), vfs.get());
+  MjModelPtr m =
+      LoadModelFromString(parent_xml, error.data(), error.size(), vfs.get());
   EXPECT_THAT(m.get(), NotNull()) << error.data();
   mj_deleteVFS(vfs.get());
 }
@@ -3725,7 +3729,7 @@ TEST_F(ActuatorParseTest, ActuatorDelayParsed) {
   ASSERT_EQ(model->nu, 3);
   // actuator_history[2*i] = nsample, actuator_history[2*i+1] = interp
   EXPECT_EQ(model->actuator_history[0], 0);   // jnt1 nsample
-  EXPECT_EQ(model->actuator_history[1], 0);   // jnt1 interp (no buffer, default 0)
+  EXPECT_EQ(model->actuator_history[1], 0);   // jnt1 interp (no buffer, def 0)
   EXPECT_EQ(model->actuator_history[2], 3);   // jnt2 nsample
   EXPECT_EQ(model->actuator_history[3], 0);   // jnt2 interp (ZOH)
   EXPECT_EQ(model->actuator_history[4], 10);  // jnt3 nsample
@@ -3785,6 +3789,7 @@ TEST_F(ActuatorParseTest, ActuatorDelayRequiresHistory) {
   EXPECT_THAT(error.data(),
               HasSubstr("setting delay > 0 without a history buffer"));
 }
+
 TEST_F(ActuatorParseTest, DampingArmatureDefaultsPropagate) {
   static constexpr char xml[] = R"(
   <mujoco>
@@ -3884,5 +3889,106 @@ TEST_F(XMLReaderTest, AttachConflictXMLMergeUnmergableError) {
               HasSubstr("gravity: parent has 0 0 -10, child has 0 0 0"));
 }
 
+TEST_F(XMLReaderTest, SelfAttach) {
+  static constexpr char xml[] = R"(
+    <mujoco model="self-attach-test">
+      <worldbody>
+        <body name="body1">
+          <geom name="geom1" size="1"/>
+        </body>
+        <body name="body2">
+          <attach body="body1" prefix="attached_"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+  std::array<char, 1024> error;
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error.data(), error.size());
+  ASSERT_THAT(spec, NotNull()) << error.data();
+
+  mjModel* m = mj_compile(spec, nullptr);
+  ASSERT_THAT(m, NotNull());
+
+  int body2_id = mj_name2id(m, mjOBJ_BODY, "body2");
+  int attached_body1_id = mj_name2id(m, mjOBJ_BODY, "attached_body1");
+
+  EXPECT_GE(body2_id, 0);
+  EXPECT_GE(attached_body1_id, 0);
+  EXPECT_EQ(m->body_parentid[attached_body1_id], body2_id);
+
+  mj_deleteModel(m);
+  mj_deleteSpec(spec);
+}
+
+TEST_F(XMLReaderTest, SelfAttachCollisionError) {
+  static constexpr char xml[] = R"(
+    <mujoco model="self-attach-collision">
+      <worldbody>
+        <body name="body1"/>
+        <body name="attached_body1"/>
+        <body name="body2">
+          <attach body="body1" prefix="attached_"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+  std::array<char, 1024> error;
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error.data(), error.size());
+  EXPECT_THAT(spec, IsNull());
+  EXPECT_THAT(
+      error.data(),
+      HasSubstr("cannot self-attach: element attached_body1 already exists"));
+}
+
+TEST_F(XMLReaderTest, SelfAttachMissingError) {
+  static constexpr char xml[] = R"(
+    <mujoco model="self-attach-missing">
+      <worldbody>
+        <body name="body1">
+          <attach body="nonexistent" prefix="attached_"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+  std::array<char, 1024> error;
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error.data(), error.size());
+  EXPECT_THAT(spec, IsNull());
+  EXPECT_THAT(error.data(), HasSubstr("could not find body 'nonexistent' in "
+                                      "the current model for self-attachment"));
+}
+
+TEST_F(XMLReaderTest, SelfAttachFrame) {
+  static constexpr char xml[] = R"(
+    <mujoco model="self-attach-frame-test">
+      <worldbody>
+        <frame name="frame1">
+          <body name="body1" pos="-1 0 0">
+            <geom name="geom1" type="box" size="1 1 1"/>
+          </body>
+        </frame>
+        <body name="body2">
+          <attach frame="frame1" prefix="attached_"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+  std::array<char, 1024> error;
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error.data(), error.size());
+  ASSERT_THAT(spec, NotNull()) << error.data();
+
+  mjModel* m = mj_compile(spec, nullptr);
+  ASSERT_THAT(m, NotNull()) << mjs_getError(spec);
+
+  int body2_id = mj_name2id(m, mjOBJ_BODY, "body2");
+  int attached_body1_id = mj_name2id(m, mjOBJ_BODY, "attached_body1");
+
+  EXPECT_GE(body2_id, 0);
+  EXPECT_GE(attached_body1_id, 0);
+  EXPECT_EQ(m->body_parentid[attached_body1_id], body2_id);
+
+  mj_deleteModel(m);
+  mj_deleteSpec(spec);
+}
+
 }  // namespace
 }  // namespace mujoco