chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Change site and frame attach API.

Browse Source 
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
This commit is contained in:
Alessio Quaglino
2025-01-21 07:28:44 -08:00
committed by Copybara-Service
parent bd1bbfe401
commit 4bb6aeb19f
3 changed files with 114 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/python.rst
+8 -7
View File
@@ -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<body-attach>` an :ref:`replicate<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
-------------------
python/mujoco/specs.cc
+63 -43
View File
@@ -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<std::string>& prefix,
std::optional<std::string>& suffix, std::optional<py::object>& site,
std::optional<py::object>& 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<raw::MjsFrame*>();
} catch (const py::cast_error& e) {
frame_ptr =
mjs_findFrame(self.ptr, frame->cast<std::string>().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<raw::MjsSite*>();
} catch (const py::cast_error& e) {
site_ptr = mjs_asSite(mjs_findElement(
self.ptr, mjOBJ_SITE, site->cast<std::string>().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<std::string>& prefix,
std::optional<std::string>& 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<std::string>& prefix,
std::optional<std::string>& 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",
python/mujoco/specs_test.py
+43 -2
View File
@@ -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()