Add mjs_bodyToFrame() to convert a body to a frame.

Also add test for attaching the world body.

PiperOrigin-RevId: 685603933
Change-Id: I3cfe7d226d6f032b5a587e44d0f4f18207977188

doc/APIreference/functions.rst

@@ -4400,6 +4400,15 @@ mjs_resolveOrientation
 
 Resolve alternative orientations to quat, return error if any.
 
+.. _mjs_bodyToFrame:
+
+mjs_bodyToFrame
+~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_bodyToFrame
+
+Transform body into a frame.
+
 .. _ElementInitialization:
 
 Element initialization

doc/XMLreference.rst

@@ -3848,16 +3848,16 @@ defaults and assets) will be copied in to the top-level model. :el:`attach` is a
 all attachments will appear in the saved XML file.
 
 .. admonition:: Known issues
-   :class: attention
+   :class: note
 
-   The :el:`attach` meta-element is new and not well tested. Please report any issues you encounter to the development
+   The following known limitations exist, to be addressed in a future release:
-   team. Additionally, the following known limitations exist, to be addressed in a future release:
 
-   - The world body cannot be attached.
    - An entire model cannot be attached (i.e. including all elements, referenced or not).
    - All assets from the child model will be copied in, whether they are referenced or not.
-   - Self-attach or circular references are not checked for and will lead to infinite loops.
+   - Circular references are not checked for and will lead to infinite loops.
-   - :ref:`Keyframes<keyframe>` are attached once, so they are not replicated in nested attachments.
+   - When attaching a model with :ref:`keyframes<keyframe>`, model compilation is required for the re-indexing to be
+     finalized. If a second attachment is performed without compilation, the keyframes from the first attachment will be
+     lost.
 
 .. _body-attach-model:
 

doc/includes/references.h

@@ -3628,6 +3628,7 @@ void mjs_setDefault(mjsElement* element, mjsDefault* def);
 void mjs_setFrame(mjsElement* dest, mjsFrame* frame);
 const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence,
                                    const mjsOrientation* orientation);
+mjsFrame* mjs_bodyToFrame(mjsBody** body);
 void mjs_defaultSpec(mjSpec* spec);
 void mjs_defaultOrientation(mjsOrientation* orient);
 void mjs_defaultBody(mjsBody* body);

include/mujoco/mujoco.h

@@ -1625,6 +1625,8 @@ MJAPI void mjs_setFrame(mjsElement* dest, mjsFrame* frame);
 MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence,
                                          const mjsOrientation* orientation);
 
+// Transform body into a frame.
+MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body);
 
 //---------------------------------- Element initialization  ---------------------------------------
 

introspect/functions.py

@@ -10327,6 +10327,24 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Resolve alternative orientations to quat, return error if any.',
      )),
+    ('mjs_bodyToFrame',
+     FunctionDecl(
+         name='mjs_bodyToFrame',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsFrame'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='body',
+                 type=PointerType(
+                     inner_type=PointerType(
+                         inner_type=ValueType(name='mjsBody'),
+                     ),
+                 ),
+             ),
+         ),
+         doc='Transform body into a frame.',
+     )),
     ('mjs_defaultSpec',
      FunctionDecl(
          name='mjs_defaultSpec',

python/mujoco/specs.cc

@@ -573,6 +573,17 @@ PYBIND11_MODULE(_specs, m) {
         return new_frame;
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def(
+    "to_frame",
+      [](raw::MjsBody* self) -> raw::MjsFrame* {
+        raw::MjsFrame* frame = mjs_bodyToFrame(&self);
+        if (!frame) {
+          throw pybind11::value_error(mjs_getError(mjs_getSpec(self->element)));
+        }
+        // TODO: set the body to py::none
+        return frame;
+      },
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSFRAME ====================================
   mjsFrame.def_property_readonly(

python/mujoco/specs_test.py

@@ -866,6 +866,13 @@ class SpecsTest(absltest.TestCase):
     model = parent.compile()
     np.testing.assert_array_equal(model.body_pos[1], [1, 2, 3])
 
+  def test_body_to_frame(self):
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body(pos=[1, 2, 3])
+    spec.compile()
+    frame = body.to_frame()
+    np.testing.assert_array_equal(frame.pos, [1, 2, 3])
+
 
 if __name__ == '__main__':
   absltest.main()

src/user/user_api.cc

@@ -666,6 +666,17 @@ const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* s
 
 
 
+// Transform body into a frame.
+mjsFrame* mjs_bodyToFrame(mjsBody** body) {
+  mjCBody* bodyC = static_cast<mjCBody*>((*body)->element);
+  mjCFrame* frameC = bodyC->ToFrame();
+  delete bodyC;
+  *body = nullptr;
+  return &frameC->spec;
+}
+
+
+
 // get id
 int mjs_getId(mjsElement* element) {
   if (!element) {

src/user/user_api.h

@@ -365,6 +365,9 @@ MJAPI void mjs_setFrame(mjsElement* dest, mjsFrame* frame);
 MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence,
                                          const mjsOrientation* orientation);
 
+// Transform body into a frame.
+MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body);
+
 
 //---------------------------------- Initialization  -----------------------------------------------
 

src/user/user_objects.cc

@@ -1165,6 +1165,39 @@ mjCLight* mjCBody::AddLight(mjCDef* _def) {
 
 
 
+// create a frame in the parent body and move all contents of this body into it
+mjCFrame* mjCBody::ToFrame() {
+  if (parentid < 0) {
+    // TODO: store the parent pointer instead of using the id
+    throw mjCError(this, "parent body is not defined, please compile the model first");
+  }
+  mjCBody* parent = model->Bodies()[parentid];
+  mjCFrame* newframe = parent->AddFrame(frame);
+  mjuu_copyvec(newframe->spec.pos, spec.pos, 3);
+  mjuu_copyvec(newframe->spec.quat, spec.quat, 4);
+  parent->bodies.insert(parent->bodies.end(), bodies.begin(), bodies.end());
+  parent->geoms.insert(parent->geoms.end(), geoms.begin(), geoms.end());
+  parent->joints.insert(parent->joints.end(), joints.begin(), joints.end());
+  parent->sites.insert(parent->sites.end(), sites.begin(), sites.end());
+  parent->cameras.insert(parent->cameras.end(), cameras.begin(), cameras.end());
+  parent->lights.insert(parent->lights.end(), lights.begin(), lights.end());
+  parent->frames.insert(parent->frames.end(), frames.begin(), frames.end());
+  bodies.clear();
+  geoms.clear();
+  joints.clear();
+  sites.clear();
+  cameras.clear();
+  lights.clear();
+  frames.clear();
+  parent->bodies.erase(
+      std::remove_if(parent->bodies.begin(), parent->bodies.end(),
+                     [this](mjCBody* body) { return body == this; }),
+      parent->bodies.end());
+  return newframe;
+}
+
+
+
 // get number of objects of specified type
 int mjCBody::NumObjects(mjtObj type) {
   switch (type) {

src/user/user_objects.h

@@ -331,6 +331,9 @@ class mjCBody : public mjCBody_, private mjsBody {
   // reset keyframe references for allowing self-attach
   void ForgetKeyframes() const;
 
+  // create a frame and move all contents of this body into it
+  mjCFrame* ToFrame();
+
   // get mocap position and quaternion
   mjtNum* mpos(const std::string& state_name);
   mjtNum* mquat(const std::string& state_name);

test/user/user_api_test.cc

@@ -1152,6 +1152,71 @@ TEST_F(MujocoTest, AttachToSite) {
   mj_deleteModel(expected);
 }
 
+TEST_F(MujocoTest, AttachWorld) {
+  std::array<char, 1000> er;
+  mjtNum tol = 0;
+  std::string field = "";
+
+  static constexpr char xml_parent[] = R"(
+  <mujoco>
+    <worldbody>
+      <frame name="frame"/>
+    </worldbody>
+  </mujoco>)";
+
+  static constexpr char xml_child[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="sphere">
+        <joint type="slide"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+  </mujoco>)";
+
+  static constexpr char xml_result[] = R"(
+  <mujoco>
+    <worldbody>
+      <frame name="frame">
+        <frame name="attached-world-1">
+          <body name="attached-sphere-1">
+            <joint type="slide"/>
+            <geom size=".1"/>
+          </body>
+        </frame>
+      </frame>
+    </worldbody>
+  </mujoco>)";
+
+  mjSpec* parent = mj_parseXMLString(xml_parent, 0, er.data(), er.size());
+  EXPECT_THAT(parent, NotNull()) << er.data();
+  mjSpec* child = mj_parseXMLString(xml_child, 0, er.data(), er.size());
+  EXPECT_THAT(child, NotNull()) << er.data();
+
+  mjsFrame* frame = mjs_findFrame(parent, "frame");
+  EXPECT_THAT(frame, NotNull());
+  mjsBody* world = mjs_findBody(child, "world");
+  EXPECT_THAT(world, NotNull());
+  mjsBody* child_world = mjs_attachBody(frame, world, "attached-", "-1");
+  EXPECT_THAT(child_world, NotNull());
+  mjsFrame* frame_world = mjs_bodyToFrame(&child_world);
+  EXPECT_THAT(frame_world, NotNull());
+  EXPECT_THAT(child_world, IsNull());
+
+  mjModel* model = mj_compile(parent, 0);
+  EXPECT_THAT(model, NotNull());
+  mjModel* expected = LoadModelFromString(xml_result, er.data(), er.size());
+  EXPECT_THAT(expected, NotNull()) << er.data();
+  EXPECT_LE(CompareModel(model, expected, field), tol)
+            << "Expected and attached models are different!\n"
+            << "Different field: " << field << '\n';
+
+  mj_deleteSpec(parent);
+  mj_deleteSpec(child);
+  mj_deleteModel(model);
+  mj_deleteModel(expected);
+}
+
 TEST_F(MujocoTest, PreserveState) {
   std::array<char, 1000> er;
   std::string field = "";