From 4bb6aeb19f65920791e4907cdb6a6f29bcfa44ec Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Tue, 21 Jan 2025 07:28:44 -0800 Subject: [PATCH] Change site and frame attach API. Old API: ``` site.attach(child_spec) frame.attach(child_spec) ``` New API: ``` parent_spec.attach(child_spec, frame=frame_name_or_object) parent_spec.attach(site_spec, site=site_name_or_object) ``` This enables accessing the parent object during attach, which will be used in a follow-up CL to automatically append the child assets to the parent. PiperOrigin-RevId: 717908589 Change-Id: I0fb27e99694c954cb8ba8a8c98484ffd53a4d6db --- doc/python.rst | 15 ++--- python/mujoco/specs.cc | 106 +++++++++++++++++++++--------------- python/mujoco/specs_test.py | 45 ++++++++++++++- 3 files changed, 114 insertions(+), 52 deletions(-) diff --git a/doc/python.rst b/doc/python.rst index 95bd409a..5a479d98 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -534,11 +534,12 @@ It is possible to combine multiple specs by using attachments. The following opt the reference to the attached body, which should be identical to the body used as input. - Attach a frame from the child spec to a body in the parent spec: ``body.attach_frame(frame, prefix, suffix)``, returns the reference to the attached frame, which should be identical to the frame used as input. -- Attach a body from the child spec to a site in the parent spec: ``site.attach(body, prefix, suffix)``, returns the - reference to the attached body, which should be identical to the body used as input. -- 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. +- Attach a child spec to a site in the parent spec: ``spec.attach(child_spec, site=site_name_or_obj)``, returns the + reference to a frame, which is the attached worldbody transformed into a frame. The site must belong to the child + spec. Prefix and suffix can also be specified as keyword arguments. +- Attach a child spec to a frame in the parent spec: ``parent_spec.attach(child_spec, frame=frame_name_or_obj)``, + returns the reference to a frame, which is the attached worldbody transformed into a frame. The frame must belong to + the child spec. Prefix and suffix can also be specified as keyword arguments. Attaching does not copy, so all the child reference are still valid in the parent and therefore modifying the child will modify the parent. This is not true for the attach :ref:`attach` an :ref:`replicate` @@ -562,8 +563,8 @@ meta-elements in MJCF, which create deep copies while attaching. # 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-', '') + worldframe_in_site = parent.attach(child, site=site, prefix='child-') + worldframe_in_frame = parent.attach(child, frame=frame, prefix='child-') Convenience methods ------------------- diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 560d5ba6..08716d1b 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -495,6 +495,69 @@ PYBIND11_MODULE(_specs, m) { mjSpec.def("detach_body", [](MjSpec& self, raw::MjsBody& body) { mjs_detachBody(self.ptr, &body); }); + mjSpec.def( + "attach", + [](MjSpec& self, MjSpec& child, std::optional& prefix, + std::optional& suffix, std::optional& site, + std::optional& frame) -> raw::MjsFrame* { + if (!frame.has_value() && !site.has_value()) { + throw pybind11::value_error( + "One of frame or site must be specified."); + } + if (frame.has_value() && site.has_value()) { + throw pybind11::value_error( + "Only one of frame or site can be specified."); + } + auto world = mjs_findBody(child.ptr, "world"); + if (!world) { + throw pybind11::value_error("Child does not have a world body."); + } + const char* p = prefix.has_value() ? prefix.value().c_str() : ""; + const char* s = suffix.has_value() ? suffix.value().c_str() : ""; + raw::MjsBody* attached_world = nullptr; + if (frame.has_value()) { + raw::MjsFrame* frame_ptr = nullptr; + try { + frame_ptr = frame->cast(); + } catch (const py::cast_error& e) { + frame_ptr = + mjs_findFrame(self.ptr, frame->cast().c_str()); + } + if (!frame_ptr) { + throw pybind11::value_error("Frame not found."); + } + if (mjs_getSpec(frame_ptr->element) != self.ptr) { + throw pybind11::value_error( + "Frame spec does not match parent spec."); + } + attached_world = mjs_attachBody(frame_ptr, world, p, s); + } + if (site.has_value()) { + raw::MjsSite* site_ptr = nullptr; + try { + site_ptr = site->cast(); + } catch (const py::cast_error& e) { + site_ptr = mjs_asSite(mjs_findElement( + self.ptr, mjOBJ_SITE, site->cast().c_str())); + } + if (!site_ptr) { + throw pybind11::value_error("Site not found."); + } + if (mjs_getSpec(site_ptr->element) != self.ptr) { + throw pybind11::value_error( + "Site spec does not match parent spec."); + } + attached_world = mjs_attachToSite(site_ptr, world, p, s); + } + if (!attached_world) { + throw pybind11::value_error(mjs_getError(self.ptr)); + } + return mjs_bodyToFrame(&attached_world); + }, + py::arg("child"), py::arg("prefix") = py::none(), + py::arg("suffix") = py::none(), py::arg("site") = py::none(), + py::arg("frame") = py::none(), + py::return_value_policy::reference_internal); // ============================= MJSBODY ===================================== mjsBody.def( @@ -777,27 +840,6 @@ PYBIND11_MODULE(_specs, m) { py::arg("body"), py::arg("prefix") = py::none(), py::arg("suffix") = py::none(), py::return_value_policy::reference_internal); - mjsFrame.def( - "attach", - [](raw::MjsFrame& self, MjSpec& spec, std::optional& prefix, - std::optional& suffix) -> raw::MjsFrame* { - auto world = mjs_findBody(spec.ptr, "world"); - if (!world) { - throw pybind11::value_error( - mjs_getError(mjs_getSpec(self.element))); - } - const char* p = prefix.has_value() ? prefix.value().c_str() : ""; - const char* s = suffix.has_value() ? suffix.value().c_str() : ""; - auto attached_world = mjs_attachBody(&self, world, p, s); - if (!attached_world) { - throw pybind11::value_error( - mjs_getError(mjs_getSpec(self.element))); - } - return mjs_bodyToFrame(&attached_world); - }, - py::arg("spec"), py::arg("prefix") = py::none(), - py::arg("suffix") = py::none(), - py::return_value_policy::reference_internal); // ============================= MJSGEOM ===================================== mjsGeom.def("delete", [](raw::MjsGeom& self) { mjs_delete(self.element); }); @@ -878,28 +920,6 @@ PYBIND11_MODULE(_specs, m) { py::arg("body"), py::arg("prefix") = py::none(), py::arg("suffix") = py::none(), py::return_value_policy::reference_internal); - mjsSite.def( - "attach", - [](raw::MjsSite& self, MjSpec& spec, - std::optional& prefix, - std::optional& suffix) -> raw::MjsFrame* { - auto world = mjs_findBody(spec.ptr, "world"); - if (!world) { - throw pybind11::value_error( - mjs_getError(mjs_getSpec(self.element))); - } - const char* p = prefix.has_value() ? prefix.value().c_str() : ""; - const char* s = suffix.has_value() ? suffix.value().c_str() : ""; - auto attached_world = mjs_attachToSite(&self, world, p, s); - if (!attached_world) { - throw pybind11::value_error( - mjs_getError(mjs_getSpec(self.element))); - } - return mjs_bodyToFrame(&attached_world); - }, - py::arg("body"), py::arg("prefix") = py::none(), - py::arg("suffix") = py::none(), - py::return_value_policy::reference_internal); // ============================= MJSCAMERA =================================== mjsCamera.def("delete", diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 993fc2d4..6465619b 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -939,6 +939,7 @@ class SpecsTest(absltest.TestCase): def test_attach_to_site(self): parent = mujoco.MjSpec() site = parent.worldbody.add_site(pos=[1, 2, 3], quat=[0, 0, 0, 1]) + site.name = 'site' # Attach body to site and compile. child1 = mujoco.MjSpec() @@ -954,7 +955,7 @@ class SpecsTest(absltest.TestCase): # Attach entire spec to site and compile again. child2 = mujoco.MjSpec() body2 = child2.worldbody.add_body(name='body') - self.assertIsNotNone(site.attach(child2, prefix='child-')) + self.assertIsNotNone(parent.attach(child2, site=site, prefix='child2-')) body2.pos = [-1, -1, -1] model2 = parent.compile() self.assertIsNotNone(model2) @@ -964,6 +965,26 @@ class SpecsTest(absltest.TestCase): np.testing.assert_array_equal(model2.body_quat[1], [0, 0, 0, 1]) np.testing.assert_array_equal(model2.body_quat[2], [0, 0, 0, 1]) + # Attach another spec to site (referenced by name) and compile again. + child3 = mujoco.MjSpec() + body3 = child3.worldbody.add_body(name='body') + self.assertIsNotNone(parent.attach(child3, site='site', prefix='child3-')) + body3.pos = [-2, -2, -2] + model3 = parent.compile() + self.assertIsNotNone(model3) + self.assertEqual(model3.nbody, 4) + np.testing.assert_array_equal(model3.body_pos[1], [0, 1, 4]) + np.testing.assert_array_equal(model3.body_pos[2], [2, 3, 2]) + np.testing.assert_array_equal(model3.body_pos[3], [3, 4, 1]) + np.testing.assert_array_equal(model3.body_quat[1], [0, 0, 0, 1]) + np.testing.assert_array_equal(model3.body_quat[2], [0, 0, 0, 1]) + np.testing.assert_array_equal(model3.body_quat[3], [0, 0, 0, 1]) + + # Fail to attach to a site that does not exist. + child4 = mujoco.MjSpec() + with self.assertRaisesRegex(ValueError, 'Site not found.'): + parent.attach(child4, site='invalid_site', prefix='child3-') + def test_body_to_frame(self): spec = mujoco.MjSpec() body = spec.worldbody.add_body(pos=[1, 2, 3]) @@ -974,6 +995,7 @@ class SpecsTest(absltest.TestCase): def test_attach_to_frame(self): parent = mujoco.MjSpec() frame = parent.worldbody.add_frame(pos=[1, 2, 3], quat=[0, 0, 0, 1]) + frame.name = 'frame' # Attach body to frame and compile. child1 = mujoco.MjSpec() @@ -989,7 +1011,7 @@ class SpecsTest(absltest.TestCase): # Attach entire spec to frame and compile again. child2 = mujoco.MjSpec() body2 = child2.worldbody.add_body(name='body') - self.assertIsNotNone(frame.attach(child2, prefix='child-')) + self.assertIsNotNone(parent.attach(child2, frame=frame, prefix='child-')) body2.pos = [-1, -1, -1] model2 = parent.compile() self.assertIsNotNone(model2) @@ -999,6 +1021,25 @@ class SpecsTest(absltest.TestCase): np.testing.assert_array_equal(model2.body_quat[1], [0, 0, 0, 1]) np.testing.assert_array_equal(model2.body_quat[2], [0, 0, 0, 1]) + # Attach another spec to frame (referenced by name) and compile again. + child3 = mujoco.MjSpec() + body3 = child3.worldbody.add_body(name='body') + self.assertIsNotNone(parent.attach(child3, frame='frame', prefix='child3-')) + body3.pos = [-2, -2, -2] + model3 = parent.compile() + self.assertIsNotNone(model3) + self.assertEqual(model3.nbody, 4) + np.testing.assert_array_equal(model3.body_pos[1], [0, 1, 4]) + np.testing.assert_array_equal(model3.body_pos[2], [2, 3, 2]) + np.testing.assert_array_equal(model3.body_pos[3], [3, 4, 1]) + np.testing.assert_array_equal(model3.body_quat[1], [0, 0, 0, 1]) + np.testing.assert_array_equal(model3.body_quat[2], [0, 0, 0, 1]) + np.testing.assert_array_equal(model3.body_quat[3], [0, 0, 0, 1]) + + # Fail to attach to a frame that does not exist. + child4 = mujoco.MjSpec() + with self.assertRaisesRegex(ValueError, 'Frame not found.'): + parent.attach(child4, frame='invalid_frame', prefix='child3-') if __name__ == '__main__': absltest.main()