diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 047002f9..4fcd94b9 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -746,7 +746,7 @@ PYBIND11_MODULE(_specs, m) { }, py::return_value_policy::reference_internal); mjsSite.def( - "attach", + "attach_body", [](raw::MjsSite& self, raw::MjsBody& body, std::optional& prefix, std::optional& suffix) -> raw::MjsBody* { @@ -762,6 +762,28 @@ 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 a7133e77..4090a7cf 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -875,11 +875,26 @@ class SpecsTest(absltest.TestCase): def test_attach_body_to_site(self): child = mujoco.MjSpec() parent = mujoco.MjSpec() - site = parent.worldbody.add_site(pos=[1, 2, 3]) + site = parent.worldbody.add_site(pos=[1, 2, 3], quat=[0, 0, 0, 1]) body = child.worldbody.add_body() - self.assertIsNotNone(site.attach(body, prefix='_')) - model = parent.compile() - np.testing.assert_array_equal(model.body_pos[1], [1, 2, 3]) + + # Attach body to site and compile. + self.assertIsNotNone(site.attach_body(body, prefix='_')) + model1 = parent.compile() + self.assertIsNotNone(model1) + self.assertEqual(model1.nbody, 2) + np.testing.assert_array_equal(model1.body_pos[1], [1, 2, 3]) + np.testing.assert_array_equal(model1.body_quat[1], [0, 0, 0, 1]) + + # Attach entire spec to site and compile again. + self.assertIsNotNone(site.attach(child, prefix='child-')) + model2 = parent.compile() + self.assertIsNotNone(model2) + self.assertEqual(model2.nbody, 3) + np.testing.assert_array_equal(model2.body_pos[1], [1, 2, 3]) + np.testing.assert_array_equal(model2.body_pos[2], [1, 2, 3]) + 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]) def test_body_to_frame(self): spec = mujoco.MjSpec() @@ -890,16 +905,27 @@ class SpecsTest(absltest.TestCase): 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, prefix='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, '') + frame = parent.worldbody.add_frame(pos=[1, 2, 3], quat=[0, 0, 0, 1]) + body = child.worldbody.add_body() + + # Attach body to frame and compile. + self.assertIsNotNone(frame.attach_body(body, prefix='_')) + model1 = parent.compile() + self.assertIsNotNone(model1) + self.assertEqual(model1.nbody, 2) + np.testing.assert_array_equal(model1.body_pos[1], [1, 2, 3]) + np.testing.assert_array_equal(model1.body_quat[1], [0, 0, 0, 1]) + + # Attach entire spec to frame and compile again. + self.assertIsNotNone(frame.attach(child, prefix='child-')) + model2 = parent.compile() + self.assertIsNotNone(model2) + self.assertEqual(model2.nbody, 3) + np.testing.assert_array_equal(model2.body_pos[1], [1, 2, 3]) + np.testing.assert_array_equal(model2.body_pos[2], [1, 2, 3]) + 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]) if __name__ == '__main__':