Add mjs_attachToSite for attaching a body to a site.

PiperOrigin-RevId: 684432183
Change-Id: I7100bebfb6521f9dbfbb07f31490b2b9298bf18c

doc/APIreference/functions.rst

@@ -3813,6 +3813,15 @@ mjs_attachFrame
 
 Attach child frame to a parent body, return the attached frame if success or NULL otherwise.
 
+.. _mjs_attachToSite:
+
+mjs_attachToSite
+~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_attachToSite
+
+Attach child body to a parent site, return the attached body if success or NULL otherwise.
+
 .. _mjs_detachBody:
 
 mjs_detachBody
@@ -4126,15 +4135,6 @@ mjs_getSpec
 
 Get spec from body.
 
-.. _mjs_getSpecFromFrame:
-
-mjs_getSpecFromFrame
-~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjs_getSpecFromFrame
-
-Get spec from frame.
-
 .. _mjs_findBody:
 
 mjs_findBody

doc/includes/references.h

@@ -3564,6 +3564,8 @@ mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
                         const char* prefix, const char* suffix);
 mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
                           const char* prefix, const char* suffix);
+mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
+                          const char* prefix, const char* suffix);
 int mjs_detachBody(mjSpec* s, mjsBody* b);
 mjsBody* mjs_addBody(mjsBody* body, mjsDefault* def);
 mjsSite* mjs_addSite(mjsBody* body, mjsDefault* def);
@@ -3596,8 +3598,7 @@ mjsHField* mjs_addHField(mjSpec* s);
 mjsSkin* mjs_addSkin(mjSpec* s);
 mjsTexture* mjs_addTexture(mjSpec* s);
 mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def);
-mjSpec* mjs_getSpec(mjsBody* body);
+mjSpec* mjs_getSpec(mjsElement* element);
-mjSpec* mjs_getSpecFromFrame(mjsFrame* frame);
 mjsBody* mjs_findBody(mjSpec* s, const char* name);
 mjsElement* mjs_findElement(mjSpec* s, mjtObj type, const char* name);
 mjsBody* mjs_findChild(mjsBody* body, const char* name);

include/mujoco/mujoco.h

@@ -1415,6 +1415,10 @@ MJAPI mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
 MJAPI mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
                                 const char* prefix, const char* suffix);
 
+// Attach child body to a parent site, return the attached body if success or NULL otherwise.
+MJAPI mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
+                                const char* prefix, const char* suffix);
+
 // Detach body from mjSpec, remove all references and delete the body, return 0 on success.
 MJAPI int mjs_detachBody(mjSpec* s, mjsBody* b);
 
@@ -1524,10 +1528,7 @@ MJAPI mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def);
 //---------------------------------- Find and get utilities ----------------------------------------
 
 // Get spec from body.
-MJAPI mjSpec* mjs_getSpec(mjsBody* body);
+MJAPI mjSpec* mjs_getSpec(mjsElement* element);
-
-// Get spec from frame.
-MJAPI mjSpec* mjs_getSpecFromFrame(mjsFrame* frame);
 
 // Find body in spec by name.
 MJAPI mjsBody* mjs_findBody(mjSpec* s, const char* name);

introspect/functions.py

@@ -9030,6 +9030,40 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Attach child frame to a parent body, return the attached frame if success or NULL otherwise.',  # pylint: disable=line-too-long
      )),
+    ('mjs_attachToSite',
+     FunctionDecl(
+         name='mjs_attachToSite',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsBody'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='parent',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsSite'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='child',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsBody', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='prefix',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='suffix',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+             ),
+         ),
+         doc='Attach child body to a parent site, return the attached body if success or NULL otherwise.',  # pylint: disable=line-too-long
+     )),
     ('mjs_detachBody',
      FunctionDecl(
          name='mjs_detachBody',
@@ -9674,30 +9708,14 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          parameters=(
              FunctionParameterDecl(
-                 name='body',
+                 name='element',
                  type=PointerType(
-                     inner_type=ValueType(name='mjsBody'),
+                     inner_type=ValueType(name='mjsElement'),
                  ),
              ),
          ),
          doc='Get spec from body.',
      )),
-    ('mjs_getSpecFromFrame',
-     FunctionDecl(
-         name='mjs_getSpecFromFrame',
-         return_type=PointerType(
-             inner_type=ValueType(name='mjSpec'),
-         ),
-         parameters=(
-             FunctionParameterDecl(
-                 name='frame',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjsFrame'),
-                 ),
-             ),
-         ),
-         doc='Get spec from frame.',
-     )),
     ('mjs_findBody',
      FunctionDecl(
          name='mjs_findBody',

python/mujoco/specs.cc

@@ -426,7 +426,7 @@ PYBIND11_MODULE(_specs, m) {
       [](raw::MjsBody& self, mjtObj objtype) -> py::list {
         py::list list;
         raw::MjsElement* el = mjs_firstChild(&self, objtype, true);
-        std::string error = mjs_getError(mjs_getSpec(&self));
+        std::string error = mjs_getError(mjs_getSpec(self.element));
         if (!el && !error.empty()) {
           throw pybind11::value_error(error);
         }
@@ -557,7 +557,9 @@ PYBIND11_MODULE(_specs, m) {
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "spec",
-      [](raw::MjsBody& self) -> raw::MjSpec* { return mjs_getSpec(&self); },
+      [](raw::MjsBody& self) -> raw::MjSpec* {
+        return mjs_getSpec(self.element);
+      },
       py::return_value_policy::reference_internal);
   mjsBody.def(
       "attach_frame",
@@ -566,7 +568,7 @@ PYBIND11_MODULE(_specs, m) {
         auto new_frame =
             mjs_attachFrame(&self, &frame, prefix.c_str(), suffix.c_str());
         if (!new_frame) {
-          throw pybind11::value_error(mjs_getError(mjs_getSpec(&self)));
+          throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element)));
         }
         return new_frame;
       },
@@ -587,7 +589,7 @@ PYBIND11_MODULE(_specs, m) {
             mjs_attachBody(&self, &body, prefix.c_str(), suffix.c_str());
         if (!new_body) {
           throw pybind11::value_error(
-              mjs_getError(mjs_getSpecFromFrame(&self)));
+              mjs_getError(mjs_getSpec(self.element)));
         }
         return new_body;
       },
@@ -643,6 +645,19 @@ PYBIND11_MODULE(_specs, m) {
         return mjs_getDefault(self.element);
       },
       py::return_value_policy::reference_internal);
+  mjsSite.def(
+      "attach",
+      [](raw::MjsSite& self, raw::MjsBody& body, std::string& prefix,
+         std::string& suffix) -> raw::MjsBody* {
+        auto new_body =
+            mjs_attachToSite(&self, &body, prefix.c_str(), suffix.c_str());
+        if (!new_body) {
+          throw pybind11::value_error(
+              mjs_getError(mjs_getSpec(self.element)));
+        }
+        return new_body;
+      },
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSCAMERA ===================================
   mjsCamera.def_property_readonly("id", [](raw::MjsCamera& self) -> int {

python/mujoco/specs_test.py

@@ -857,6 +857,15 @@ class SpecsTest(absltest.TestCase):
         ' attribute',
     )
 
+  def test_attach_body_to_site(self):
+    child = mujoco.MjSpec()
+    parent = mujoco.MjSpec()
+    site = parent.worldbody.add_site(pos=[1, 2, 3])
+    body = child.worldbody.add_body()
+    self.assertIsNotNone(site.attach(body, '', ''))
+    model = parent.compile()
+    np.testing.assert_array_equal(model.body_pos[1], [1, 2, 3])
+
 
 if __name__ == '__main__':
   absltest.main()

src/user/user_api.cc

@@ -29,6 +29,7 @@
 #include "user/user_model.h"
 #include "user/user_objects.h"
 #include "user/user_cache.h"
+#include "user/user_util.h"
 
 namespace {
 
@@ -160,6 +161,29 @@ mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
 
 
 
+// attach child body to a parent site
+mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
+                           const char* prefix, const char* suffix) {
+  if (!parent) {
+    mju_error("parent site is null");
+    return nullptr;
+  }
+  mjCSite* site = static_cast<mjCSite*>(parent->element);
+  mjCBody* body = site->Body();
+  mjCFrame* frame = body->AddFrame(site->frame);
+  frame->spec.pos[0] = site->spec.pos[0];
+  frame->spec.pos[1] = site->spec.pos[1];
+  frame->spec.pos[2] = site->spec.pos[2];
+  frame->spec.quat[0] = site->spec.quat[0];
+  frame->spec.quat[1] = site->spec.quat[1];
+  frame->spec.quat[2] = site->spec.quat[2];
+  frame->spec.quat[3] = site->spec.quat[3];
+  frame->SetParent(body);
+  return mjs_attachBody(&frame->spec, child, prefix, suffix);
+}
+
+
+
 // get error message from model
 const char* mjs_getError(mjSpec* s) {
   mjCModel* modelC = static_cast<mjCModel*>(s->element);
@@ -521,15 +545,8 @@ mjsDefault* mjs_addDefault(mjSpec* s, const char* classname, const mjsDefault* p
 
 
 // get spec from body
-mjSpec* mjs_getSpec(mjsBody* body) {
+mjSpec* mjs_getSpec(mjsElement* element) {
-  return &(static_cast<mjCBody*>(body->element)->model->spec);
+  return &(static_cast<mjCBase*>(element)->model->spec);
-}
-
-
-
-// get spec from frame
-mjSpec* mjs_getSpecFromFrame(mjsFrame* frame) {
-  return &(static_cast<mjCFrame*>(frame->element)->model->spec);
 }
 
 

src/user/user_api.h

@@ -71,11 +71,15 @@ MJAPI int mjs_activatePlugin(mjSpec* s, const char* name);
 
 // Attach child body to a parent frame, return the attached body if success or NULL otherwise.
 MJAPI mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
-                         const char* prefix, const char* suffix);
+                              const char* prefix, const char* suffix);
 
 // Attach child frame to a parent body, return the attached frame if success or NULL otherwise.
 MJAPI mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
-                          const char* prefix, const char* suffix);
+                                const char* prefix, const char* suffix);
+
+// Attach child body to a parent site, return the attached body if success or NULL otherwise.
+MJAPI mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
+                                const char* prefix, const char* suffix);
 
 // Detach body from mjSpec, remove all references and delete the body, return 0 on success.
 MJAPI int mjs_detachBody(mjSpec* s, mjsBody* b);
@@ -186,10 +190,7 @@ MJAPI mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def);
 //---------------------------------- Find/get utilities --------------------------------------------
 
 // Get spec from body.
-MJAPI mjSpec* mjs_getSpec(mjsBody* body);
+MJAPI mjSpec* mjs_getSpec(mjsElement* element);
-
-// Get spec from frame.
-MJAPI mjSpec* mjs_getSpecFromFrame(mjsFrame* frame);
 
 // Find spec (model asset) by name.
 MJAPI mjSpec* mjs_findSpec(mjSpec* spec, const char* name);

src/user/user_objects.h

@@ -587,6 +587,9 @@ class mjCSite : public mjCSite_, private mjsSite {
 
   mjsSite spec;                   // variables set by user
 
+  // site's body
+  mjCBody* Body() const { return body; }
+
   // use strings from mjCBase rather than mjStrings from mjsSite
   using mjCBase::name;
   using mjCBase::info;

test/user/user_api_test.cc

@@ -1079,6 +1079,69 @@ TEST_F(MujocoTest, DetachBody) {
   TestDetachBody(/*compile=*/true);
 }
 
+TEST_F(MujocoTest, AttachToSite) {
+  std::array<char, 1000> er;
+  mjtNum tol = 0;
+  std::string field = "";
+
+  static constexpr char xml_parent[] = R"(
+  <mujoco>
+    <worldbody>
+      <site name="site" pos="1 0 0" quat="0 1 0 0"/>
+    </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>
+      <site name="site" pos="1 0 0" quat="0 1 0 0"/>
+      <frame pos="1 0 0" quat="0 1 0 0">
+        <body name="attached-sphere-1">
+          <joint type="slide"/>
+          <geom size=".1"/>
+        </body>
+      </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();
+
+  mjsBody* world = mjs_findBody(parent, "world");
+  EXPECT_THAT(world, NotNull());
+  mjsSite* site = mjs_asSite(mjs_firstChild(world, mjOBJ_SITE, 0));
+  EXPECT_THAT(site, NotNull());
+  mjsBody* body = mjs_findBody(child, "sphere");
+  EXPECT_THAT(body, NotNull());
+  mjsBody* attached = mjs_attachToSite(site, body, "attached-", "-1");
+  EXPECT_THAT(attached, NotNull());
+
+  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 = "";