From f909d63cdf9033d7ce13744ab631c099aa19b753 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Sun, 24 Nov 2024 07:34:16 -0800
Subject: [PATCH] Enable to attach a spec to a frame.

PiperOrigin-RevId: 699703988
Change-Id: Idbddcbce1e68286248458eb0bd24d5080451a58a
---
 doc/modeling.rst            |  2 +-
 doc/python.rst              | 36 ++++++++++++++++++++++++++++++++++++
 python/mujoco/specs.cc      | 18 ++++++++++++++++++
 python/mujoco/specs_test.py | 13 +++++++++++++
 4 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/doc/modeling.rst b/doc/modeling.rst
index b68d5e99..fdf6017f 100644
--- a/doc/modeling.rst
+++ b/doc/modeling.rst
@@ -1284,7 +1284,7 @@ elastic structures.
 
 
 The box type, as well as the cylinder and ellipsoid types, are now deprecated in favor of 3D flex :ref:`deformable
-objects <CDeformable>``. element.
+objects <CDeformable>`. element.
 
 .. _CDeformable:
 
diff --git a/doc/python.rst b/doc/python.rst
index 9c824b86..d5e4a165 100644
--- a/doc/python.rst
+++ b/doc/python.rst
@@ -503,6 +503,42 @@ The ``MjSpec`` object wraps the :ref:`mjSpec` struct and can be constructed in t
 
 Note the ``from_string()`` and ``from_file()`` methods can only be called at construction time.
 
+Attachments
+-----------
+
+It is possible to combine multiple specs by using attachments. The following options are possible:
+
+-   Attach a body from the child spec to a frame in the parent spec: ``body.attach_body(body, prefix, suffix)``, returns
+    the newly createdbody in the parent spec.
+-   Attach a frame from the child spec to a body in the parent spec: ``body.attach_frame(frame, prefix, suffix)``,
+    returns the newly created frame in the parent spec.
+-   Attach a body from the child spec to a site in the parent spec: ``site.attach(body, prefix, suffix)``, returns the
+    newly created body in the parent spec.
+-   Attach the worldbody from the child spec to a frame in the parent spec and transform it to a frame:
+    ``body.attach(spec, prefix, suffix)``, returns the newly created frame that the child worldbody was transformed
+    into.
+
+.. code-block:: python
+
+   import mujoco
+
+   # Create the parent spec.
+   parent = mujoco.MjSpec()
+   body = parent.worldbody.add_body()
+   frame = parent.worldbody.add_frame()
+   site = parent.worldbody.add_site()
+
+   # Create the child spec.
+   child = mujoco.MjSpec()
+   child_body = child.worldbody.add_body()
+   child_frame = child.worldbody.add_frame()
+
+   # Attach the child to the parent in different ways.
+   body_in_frame = frame.attach_body(child_body, 'child-', '')
+   frame_in_body = body.attach_frame(child_frame, 'child-', '')
+   body_in_site = site.attach(child_body, 'child-', '')
+   worldframe_in_frame = frame.attach(child, 'child-', '')
+
 Convenience methods
 -------------------
 
diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc
index 725894c6..f97ed6fd 100644
--- a/python/mujoco/specs.cc
+++ b/python/mujoco/specs.cc
@@ -665,6 +665,24 @@ PYBIND11_MODULE(_specs, m) {
         return new_body;
       },
       py::return_value_policy::reference_internal);
+  mjsFrame.def(
+      "attach",
+      [](raw::MjsFrame& self, MjSpec& spec, std::string& prefix,
+         std::string& suffix) -> raw::MjsFrame* {
+        auto world = mjs_findBody(spec.ptr, "world");
+        if (!world) {
+          throw pybind11::value_error(
+              mjs_getError(mjs_getSpec(self.element)));
+        }
+        auto attached_world =
+            mjs_attachBody(&self, world, prefix.c_str(), suffix.c_str());
+        if (!attached_world) {
+          throw pybind11::value_error(
+              mjs_getError(mjs_getSpec(self.element)));
+        }
+        return mjs_bodyToFrame(&attached_world);
+      },
+      py::return_value_policy::reference_internal);
 
   // ============================= MJSGEOM =====================================
   mjsGeom.def("delete", [](raw::MjsGeom& self) { mjs_delete(self.element); });
diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
index 74a34648..42686f1b 100644
--- a/python/mujoco/specs_test.py
+++ b/python/mujoco/specs_test.py
@@ -884,6 +884,19 @@ class SpecsTest(absltest.TestCase):
     frame = body.to_frame()
     np.testing.assert_array_equal(frame.pos, [1, 2, 3])
 
+  def test_attach_spec_to_frame(self):
+    child = mujoco.MjSpec()
+    child.worldbody.add_camera(name='camera')
+    parent = mujoco.MjSpec()
+    frame = parent.worldbody.add_frame(name='frame')
+    frame.attach(child, 'child-', '')
+    self.assertLen(child.cameras, 1)
+    self.assertLen(parent.bodies, 1)
+    self.assertLen(parent.frames, 2)
+    self.assertEqual(parent.cameras[0].name, 'child-camera')
+    self.assertEqual(parent.frames[0].name, 'frame')
+    self.assertEqual(parent.frames[1].name, '')
+
 
 if __name__ == '__main__':
   absltest.main()