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

Allow keyword arguments in attach.

Browse Source 
PiperOrigin-RevId: 702718597
Change-Id: I1f1a5e9c88c8e102d5cbe5984511e7923ec7be3e
This commit is contained in:
Alessio Quaglino
2024-12-04 07:30:33 -08:00
committed by Copybara-Service
parent afc86ac1b3
commit d1b02556ef
2 changed files with 34 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/specs.cc
+31 -16
View File
@@ -625,15 +625,19 @@ PYBIND11_MODULE(_specs, m) {
py::return_value_policy::reference_internal);
mjsBody.def(
"attach_frame",
[](raw::MjsBody& self, raw::MjsFrame& frame, std::string& prefix,
std::string& suffix) -> raw::MjsFrame* {
auto new_frame =
mjs_attachFrame(&self, &frame, prefix.c_str(), suffix.c_str());
[](raw::MjsBody& self, raw::MjsFrame& frame,
std::optional<std::string>& prefix,
std::optional<std::string>& suffix) -> raw::MjsFrame* {
const char* p = prefix.has_value() ? prefix.value().c_str() : "";
const char* s = suffix.has_value() ? suffix.value().c_str() : "";
auto new_frame = mjs_attachFrame(&self, &frame, p, s);
if (!new_frame) {
throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element)));
}
return new_frame;
},
py::arg("frame"), py::arg("prefix") = py::none(),
py::arg("suffix") = py::none(),
py::return_value_policy::reference_internal);
mjsBody.def(
"to_frame",
@@ -654,34 +658,41 @@ PYBIND11_MODULE(_specs, m) {
});
mjsFrame.def(
"attach_body",
[](raw::MjsFrame& self, raw::MjsBody& body, std::string& prefix,
std::string& suffix) -> raw::MjsBody* {
auto new_body =
mjs_attachBody(&self, &body, prefix.c_str(), suffix.c_str());
[](raw::MjsFrame& self, raw::MjsBody& body,
std::optional<std::string>& prefix,
std::optional<std::string>& suffix) -> raw::MjsBody* {
const char* p = prefix.has_value() ? prefix.value().c_str() : "";
const char* s = suffix.has_value() ? suffix.value().c_str() : "";
auto new_body = mjs_attachBody(&self, &body, p, s);
if (!new_body) {
throw pybind11::value_error(
mjs_getError(mjs_getSpec(self.element)));
}
return new_body;
},
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::string& prefix,
std::string& suffix) -> raw::MjsFrame* {
[](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)));
}
auto attached_world =
mjs_attachBody(&self, world, prefix.c_str(), suffix.c_str());
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 =====================================
@@ -730,16 +741,20 @@ PYBIND11_MODULE(_specs, m) {
py::return_value_policy::reference_internal);
mjsSite.def(
"attach",
[](raw::MjsSite& self, raw::MjsBody& body, std::string& prefix,
std::string& suffix) -> raw::MjsBody* {
auto new_body =
mjs_attachToSite(&self, &body, prefix.c_str(), suffix.c_str());
[](raw::MjsSite& self, raw::MjsBody& body,
std::optional<std::string>& prefix,
std::optional<std::string>& suffix) -> raw::MjsBody* {
const char* p = prefix.has_value() ? prefix.value().c_str() : "";
const char* s = suffix.has_value() ? suffix.value().c_str() : "";
auto new_body = mjs_attachToSite(&self, &body, p, s);
if (!new_body) {
throw pybind11::value_error(
mjs_getError(mjs_getSpec(self.element)));
}
return new_body;
},
py::arg("body"), py::arg("prefix") = py::none(),
py::arg("suffix") = py::none(),
py::return_value_policy::reference_internal);
// ============================= MJSCAMERA ===================================
python/mujoco/specs_test.py
+3 -3
View File
@@ -867,7 +867,7 @@ class SpecsTest(absltest.TestCase):
parent.compiler.degree = not child.compiler.degree
body = child.worldbody.add_body(euler=[90, 0, 0])
frame = parent.worldbody.add_frame(euler=[-mujoco.mjPI / 2, 0, 0])
frame.attach_body(body, 'child-', '')
frame.attach_body(body, prefix='child-')
model = parent.compile()
np.testing.assert_almost_equal(model.body_quat[1], [1, 0, 0, 0])
@@ -876,7 +876,7 @@ class SpecsTest(absltest.TestCase):
parent = mujoco.MjSpec()
site = parent.worldbody.add_site(pos=[1, 2, 3])
body = child.worldbody.add_body()
self.assertIsNotNone(site.attach(body, '_', ''))
self.assertIsNotNone(site.attach(body, prefix='_'))
model = parent.compile()
np.testing.assert_array_equal(model.body_pos[1], [1, 2, 3])
@@ -892,7 +892,7 @@ class SpecsTest(absltest.TestCase):
child.worldbody.add_camera(name='camera')
parent = mujoco.MjSpec()
frame = parent.worldbody.add_frame(name='frame')
frame.attach(child, 'child-', '')
frame.attach(child, prefix='child-')
self.assertLen(child.cameras, 1)
self.assertLen(parent.bodies, 1)
self.assertLen(parent.frames, 2)