chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 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
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()