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

Change site.attach to accept a spec instead of a body.

Browse Source 
Add site.attach_body for attaching a body to a site.

PiperOrigin-RevId: 703047770
Change-Id: Ibc1f916f8741872135220461cda7cebc4b418913
This commit is contained in:
Alessio Quaglino
2024-12-05 03:09:17 -08:00
committed by Copybara-Service
parent 10239a673b
commit 503e5e1815
2 changed files with 62 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/specs.cc
+23 -1
View File
@@ -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<std::string>& prefix,
std::optional<std::string>& 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<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
+39 -13
View File
@@ -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__':