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

Add sibling traversal for tree elements.

Browse Source 
PiperOrigin-RevId: 691057957
Change-Id: I3f30cf175b5ad0842f3fbbf7fdb42ba7a2212682
This commit is contained in:
Alessio Quaglino
2024-10-29 09:54:07 -07:00
committed by Copybara-Service
parent b941e994d9
commit e835f139c8
2 changed files with 68 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/specs.cc
+47 -5
View File
@@ -155,9 +155,9 @@ void DefineArray(py::module& m, const std::string& typestr) {
}, py::keep_alive<0, 1>(), py::return_value_policy::reference_internal);
};
py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype) {
py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype, bool recursive) {
py::list list;
raw::MjsElement* el = mjs_firstChild(&body, objtype, true);
raw::MjsElement* el = mjs_firstChild(&body, objtype, recursive);
std::string error = mjs_getError(mjs_getSpec(body.element));
if (!el && !error.empty()) {
throw pybind11::value_error(error);
@@ -192,7 +192,7 @@ py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype) {
"light, camera.");
break;
}
el = mjs_nextChild(&body, el, true);
el = mjs_nextChild(&body, el, recursive);
}
return list; // list of pointers, so they can be copied
}
@@ -460,7 +460,7 @@ PYBIND11_MODULE(_specs, m) {
mjsBody.def(
"find_all",
[](raw::MjsBody& self, mjtObj objtype) -> py::list {
return FindAllImpl(self, objtype);
return FindAllImpl(self, objtype, true);
},
py::return_value_policy::reference_internal);
mjsBody.def(
@@ -484,7 +484,7 @@ PYBIND11_MODULE(_specs, m) {
"body.find_all supports the types: body, frame, geom, site, "
"light, camera.");
}
return FindAllImpl(self, objtype);
return FindAllImpl(self, objtype, true);
},
py::return_value_policy::reference_internal);
mjsBody.def(
@@ -505,6 +505,12 @@ PYBIND11_MODULE(_specs, m) {
return mjs_asBody(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def_property_readonly(
"bodies",
[](raw::MjsBody& self) -> py::list {
return FindAllImpl(self, mjOBJ_BODY, false);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_camera",
[](raw::MjsBody& self) -> raw::MjsCamera* {
@@ -517,6 +523,12 @@ PYBIND11_MODULE(_specs, m) {
return mjs_asCamera(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def_property_readonly(
"cameras",
[](raw::MjsBody& self) -> py::list {
return FindAllImpl(self, mjOBJ_CAMERA, false);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_light",
[](raw::MjsBody& self) -> raw::MjsLight* {
@@ -529,6 +541,12 @@ PYBIND11_MODULE(_specs, m) {
return mjs_asLight(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def_property_readonly(
"lights",
[](raw::MjsBody& self) -> py::list {
return FindAllImpl(self, mjOBJ_LIGHT, false);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_joint",
[](raw::MjsBody& self) -> raw::MjsJoint* {
@@ -541,6 +559,12 @@ PYBIND11_MODULE(_specs, m) {
return mjs_asJoint(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def_property_readonly(
"joints",
[](raw::MjsBody& self) -> py::list {
return FindAllImpl(self, mjOBJ_JOINT, false);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_geom",
[](raw::MjsBody& self) -> raw::MjsGeom* {
@@ -553,6 +577,12 @@ PYBIND11_MODULE(_specs, m) {
return mjs_asGeom(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def_property_readonly(
"geoms",
[](raw::MjsBody& self) -> py::list {
return FindAllImpl(self, mjOBJ_GEOM, false);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_site",
[](raw::MjsBody& self) -> raw::MjsSite* {
@@ -565,6 +595,12 @@ PYBIND11_MODULE(_specs, m) {
return mjs_asSite(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def_property_readonly(
"sites",
[](raw::MjsBody& self) -> py::list {
return FindAllImpl(self, mjOBJ_SITE, false);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_frame",
[](raw::MjsBody& self) -> raw::MjsFrame* {
@@ -577,6 +613,12 @@ PYBIND11_MODULE(_specs, m) {
return mjs_asFrame(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def_property_readonly(
"frames",
[](raw::MjsBody& self) -> py::list {
return FindAllImpl(self, mjOBJ_FRAME, false);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"spec",
[](raw::MjsBody& self) -> raw::MjSpec* {
python/mujoco/specs_test.py
+21 -2
View File
@@ -602,8 +602,12 @@ class SpecsTest(absltest.TestCase):
<worldbody>
<body name="body1">
<body name="body3">
<site name="site1"/>
<site name="site2"/>
<site name="site3"/>
<site name="site4"/>
<body name="body4">
<site name="site"/>
<site name="site5"/>
</body>
</body>
</body>
@@ -614,6 +618,9 @@ class SpecsTest(absltest.TestCase):
spec = mujoco.MjSpec.from_string(main_xml)
bodytype = mujoco.mjtObj.mjOBJ_BODY
self.assertLen(spec.bodies, 5)
self.assertLen(spec.sites, 5)
self.assertLen(spec.worldbody.find_all('body'), 4)
self.assertLen(spec.worldbody.find_all('site'), 5)
self.assertEqual(spec.bodies[1].name, 'body1')
self.assertEqual(spec.bodies[2].name, 'body2')
self.assertEqual(spec.bodies[3].name, 'body3')
@@ -630,7 +637,19 @@ class SpecsTest(absltest.TestCase):
self.assertEqual(spec.bodies[3].find_all('body')[0].name, 'body4')
self.assertEmpty(spec.bodies[2].find_all('body'))
self.assertEmpty(spec.bodies[4].find_all('body'))
self.assertEqual(spec.worldbody.find_all('site')[0].name, 'site')
self.assertEqual(spec.worldbody.find_all('site')[0].name, 'site1')
self.assertEqual(spec.worldbody.find_all('site')[1].name, 'site2')
self.assertEqual(spec.worldbody.find_all('site')[2].name, 'site3')
self.assertEqual(spec.worldbody.find_all('site')[3].name, 'site4')
self.assertEqual(spec.worldbody.find_all('site')[4].name, 'site5')
self.assertEmpty(spec.bodies[2].sites)
self.assertLen(spec.bodies[3].sites, 4)
self.assertLen(spec.bodies[4].sites, 1)
self.assertEqual(spec.bodies[3].sites[0].name, 'site1')
self.assertEqual(spec.bodies[3].sites[1].name, 'site2')
self.assertEqual(spec.bodies[3].sites[2].name, 'site3')
self.assertEqual(spec.bodies[3].sites[3].name, 'site4')
self.assertEqual(spec.bodies[4].sites[0].name, 'site5')
with self.assertRaises(ValueError) as cm:
spec.worldbody.find_all('actuator')
self.assertEqual(