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

Add body.find_all() to Python bindings.

Browse Source 
This enables to generate a list of all body children of a given type.

Also, add recursive option to mjs_firstChild and mjs_nextChild. Fixes #2112.

PiperOrigin-RevId: 681039348
Change-Id: Iebbb26fd5ec314150206edfdddd88633baa0bec6
This commit is contained in:
Alessio Quaglino
2024-10-01 08:56:03 -07:00
committed by Copybara-Service
parent 8a825e6050
commit 89aea5e217
11 changed files with 150 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+2 -1
View File
@@ -4205,7 +4205,7 @@ mjs_firstChild
.. mujoco-include:: mjs_firstChild
Return body's first child of given type.
Return body's first child of given type. If recurse is nonzero, also search the body's subtree.
.. _mjs_nextChild:
@@ -4215,6 +4215,7 @@ mjs_nextChild
.. mujoco-include:: mjs_nextChild
Return body's next child of the same type; return NULL if child is last.
If recurse is nonzero, also search the body's subtree.
.. _mjs_firstElement:
doc/includes/references.h
+2 -2
View File
@@ -3598,8 +3598,8 @@ mjsDefault* mjs_getDefault(mjsElement* element);
mjsDefault* mjs_findDefault(mjSpec* s, const char* classname);
mjsDefault* mjs_getSpecDefault(mjSpec* s);
int mjs_getId(mjsElement* element);
mjsElement* mjs_firstChild(mjsBody* body, mjtObj type);
mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child);
mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse);
mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child, int recurse);
mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);
mjsElement* mjs_nextElement(mjSpec* s, mjsElement* element);
void mjs_setBuffer(mjByteVec* dest, const void* array, int size);
include/mujoco/mujoco.h
+4 -3
View File
@@ -1550,11 +1550,12 @@ MJAPI mjsDefault* mjs_getSpecDefault(mjSpec* s);
// Get element id.
MJAPI int mjs_getId(mjsElement* element);
// Return body's first child of given type.
MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type);
// Return body's first child of given type. If recurse is nonzero, also search the body's subtree.
MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse);
// Return body's next child of the same type; return NULL if child is last.
MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child);
// If recurse is nonzero, also search the body's subtree.
MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child, int recurse);
// Return spec's first element of selected type.
MJAPI mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);
introspect/functions.py
+10 -2
View File
@@ -9859,8 +9859,12 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
name='type',
type=ValueType(name='mjtObj'),
),
FunctionParameterDecl(
name='recurse',
type=ValueType(name='int'),
),
),
doc="Return body's first child of given type.",
doc="Return body's first child of given type. If recurse is nonzero, also search the body's subtree.", # pylint: disable=line-too-long
)),
('mjs_nextChild',
FunctionDecl(
@@ -9881,8 +9885,12 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
inner_type=ValueType(name='mjsElement'),
),
),
FunctionParameterDecl(
name='recurse',
type=ValueType(name='int'),
),
),
doc="Return body's next child of the same type; return NULL if child is last.", # pylint: disable=line-too-long
doc="Return body's next child of the same type; return NULL if child is last. If recurse is nonzero, also search the body's subtree.", # pylint: disable=line-too-long
)),
('mjs_firstElement',
FunctionDecl(
python/mujoco/specs.cc
+58 -14
View File
@@ -385,6 +385,50 @@ PYBIND11_MODULE(_specs, m) {
return mjs_getDefault(self.element);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"find_all",
[](raw::MjsBody& self, mjtObj objtype) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstChild(&self, objtype, true);
std::string error = mjs_getError(mjs_getSpec(&self));
if (!el && !error.empty()) {
throw pybind11::value_error(error);
}
while (el) {
switch (objtype) {
case mjOBJ_BODY:
list.append(mjs_asBody(el));
break;
case mjOBJ_CAMERA:
list.append(mjs_asCamera(el));
break;
case mjOBJ_FRAME:
list.append(mjs_asFrame(el));
break;
case mjOBJ_GEOM:
list.append(mjs_asGeom(el));
break;
case mjOBJ_JOINT:
list.append(mjs_asJoint(el));
break;
case mjOBJ_LIGHT:
list.append(mjs_asLight(el));
break;
case mjOBJ_SITE:
list.append(mjs_asSite(el));
break;
default:
// this should never happen
throw pybind11::value_error(
"body.find_all supports the types: body, frame, geom, site, "
"light, camera.");
break;
}
el = mjs_nextChild(&self, el, true);
}
return list;
},
py::return_value_policy::reference_internal);
mjsBody.def(
"find_child",
[](raw::MjsBody& self, std::string& name) -> raw::MjsBody* {
@@ -394,85 +438,85 @@ PYBIND11_MODULE(_specs, m) {
mjsBody.def(
"first_body",
[](raw::MjsBody& self) -> raw::MjsBody* {
return mjs_asBody(mjs_firstChild(&self, mjOBJ_BODY));
return mjs_asBody(mjs_firstChild(&self, mjOBJ_BODY, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"next_body",
[](raw::MjsBody& self, raw::MjsBody& child) -> raw::MjsBody* {
return mjs_asBody(mjs_nextChild(&self, child.element));
return mjs_asBody(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_camera",
[](raw::MjsBody& self) -> raw::MjsCamera* {
return mjs_asCamera(mjs_firstChild(&self, mjOBJ_CAMERA));
return mjs_asCamera(mjs_firstChild(&self, mjOBJ_CAMERA, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"next_camera",
[](raw::MjsBody& self, raw::MjsCamera& child) -> raw::MjsCamera* {
return mjs_asCamera(mjs_nextChild(&self, child.element));
return mjs_asCamera(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_light",
[](raw::MjsBody& self) -> raw::MjsLight* {
return mjs_asLight(mjs_firstChild(&self, mjOBJ_LIGHT));
return mjs_asLight(mjs_firstChild(&self, mjOBJ_LIGHT, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"next_light",
[](raw::MjsBody& self, raw::MjsLight& child) -> raw::MjsLight* {
return mjs_asLight(mjs_nextChild(&self, child.element));
return mjs_asLight(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_joint",
[](raw::MjsBody& self) -> raw::MjsJoint* {
return mjs_asJoint(mjs_firstChild(&self, mjOBJ_JOINT));
return mjs_asJoint(mjs_firstChild(&self, mjOBJ_JOINT, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"next_joint",
[](raw::MjsBody& self, raw::MjsJoint& child) -> raw::MjsJoint* {
return mjs_asJoint(mjs_nextChild(&self, child.element));
return mjs_asJoint(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_geom",
[](raw::MjsBody& self) -> raw::MjsGeom* {
return mjs_asGeom(mjs_firstChild(&self, mjOBJ_GEOM));
return mjs_asGeom(mjs_firstChild(&self, mjOBJ_GEOM, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"next_geom",
[](raw::MjsBody& self, raw::MjsGeom& child) -> raw::MjsGeom* {
return mjs_asGeom(mjs_nextChild(&self, child.element));
return mjs_asGeom(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_site",
[](raw::MjsBody& self) -> raw::MjsSite* {
return mjs_asSite(mjs_firstChild(&self, mjOBJ_SITE));
return mjs_asSite(mjs_firstChild(&self, mjOBJ_SITE, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"next_site",
[](raw::MjsBody& self, raw::MjsSite& child) -> raw::MjsSite* {
return mjs_asSite(mjs_nextChild(&self, child.element));
return mjs_asSite(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"first_frame",
[](raw::MjsBody& self) -> raw::MjsFrame* {
return mjs_asFrame(mjs_firstChild(&self, mjOBJ_FRAME));
return mjs_asFrame(mjs_firstChild(&self, mjOBJ_FRAME, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
"next_frame",
[](raw::MjsBody& self, raw::MjsFrame& child) -> raw::MjsFrame* {
return mjs_asFrame(mjs_nextChild(&self, child.element));
return mjs_asFrame(mjs_nextChild(&self, child.element, false));
},
py::return_value_policy::reference_internal);
mjsBody.def(
python/mujoco/specs_test.py
+30 -6
View File
@@ -580,7 +580,9 @@ class SpecsTest(absltest.TestCase):
<worldbody>
<body name="body1">
<body name="body3">
<body name="body4"/>
<body name="body4">
<site name="site"/>
</body>
</body>
</body>
<body name="body2"/>
@@ -588,11 +590,33 @@ class SpecsTest(absltest.TestCase):
</mujoco>
"""
spec = mujoco.MjSpec.from_string(main_xml)
self.assertLen(spec.bodies, 4)
self.assertEqual(spec.bodies[0].name, 'body1')
self.assertEqual(spec.bodies[1].name, 'body2')
self.assertEqual(spec.bodies[2].name, 'body3')
self.assertEqual(spec.bodies[3].name, 'body4')
bodytype = mujoco.mjtObj.mjOBJ_BODY
sitetype = mujoco.mjtObj.mjOBJ_SITE
self.assertLen(spec.bodies, 5)
self.assertEqual(spec.bodies[1].name, 'body1')
self.assertEqual(spec.bodies[2].name, 'body2')
self.assertEqual(spec.bodies[3].name, 'body3')
self.assertEqual(spec.bodies[4].name, 'body4')
self.assertLen(spec.worldbody.find_all(bodytype), 4)
self.assertLen(spec.bodies[1].find_all(bodytype), 2)
self.assertLen(spec.bodies[3].find_all(bodytype), 1)
self.assertEqual(spec.worldbody.find_all(bodytype)[0].name, 'body1')
self.assertEqual(spec.worldbody.find_all(bodytype)[1].name, 'body2')
self.assertEqual(spec.worldbody.find_all(bodytype)[2].name, 'body3')
self.assertEqual(spec.worldbody.find_all(bodytype)[3].name, 'body4')
self.assertEqual(spec.bodies[1].find_all(bodytype)[0].name, 'body3')
self.assertEqual(spec.bodies[1].find_all(bodytype)[1].name, 'body4')
self.assertEqual(spec.bodies[3].find_all(bodytype)[0].name, 'body4')
self.assertEmpty(spec.bodies[2].find_all(bodytype))
self.assertEmpty(spec.bodies[4].find_all(bodytype))
self.assertEqual(spec.worldbody.find_all(sitetype)[0].name, 'site')
with self.assertRaises(ValueError) as cm:
spec.worldbody.find_all(mujoco.mjtObj.mjOBJ_ACTUATOR)
self.assertEqual(
str(cm.exception),
'Error: Body.NextChild supports the types: body, frame, geom, site,'
' light, camera\nElement name \'world\', id 0',
)
def test_iterators(self):
spec = mujoco.MjSpec()
src/user/user_api.cc
+14 -4
View File
@@ -655,17 +655,27 @@ void mjs_setDefault(mjsElement* element, mjsDefault* defspec) {
// return first child of selected type
mjsElement* mjs_firstChild(mjsBody* body, mjtObj type) {
mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse) {
mjCBody* bodyC = static_cast<mjCBody*>(body->element);
return bodyC->NextChild(NULL, type);
try {
return bodyC->NextChild(NULL, type, recurse);
} catch (mjCError& e) {
bodyC->model->SetError(e);
return nullptr;
}
}
// return body's next child; return NULL if child is last
mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child) {
mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child, int recurse) {
mjCBody* bodyC = static_cast<mjCBody*>(body->element);
return bodyC->NextChild(child);
try {
return bodyC->NextChild(child, child->elemtype, recurse);
} catch(mjCError& e) {
bodyC->model->SetError(e);
return nullptr;
}
}
src/user/user_api.h
+4 -3
View File
@@ -218,11 +218,12 @@ MJAPI int mjs_getId(mjsElement* element);
//---------------------------------- Tree traversal ------------------------------------------------
// Return body's first child of given type.
MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type);
// Return body's first child of given type. If recurse is nonzero, also search the body's subtree.
MJAPI mjsElement* mjs_firstChild(mjsBody* body, mjtObj type, int recurse);
// Return body's next child of the same type; return NULL if child is last.
MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child);
// If recurse is nonzero, also search the body's subtree.
MJAPI mjsElement* mjs_nextChild(mjsBody* body, mjsElement* child, int recurse);
// Return spec's first element of selected type.
MJAPI mjsElement* mjs_firstElement(mjSpec* s, mjtObj type);
src/user/user_model.cc
+7
View File
@@ -921,6 +921,13 @@ mjsElement* mjCModel::NextObject(mjsElement* object, mjtObj type) {
switch (type) {
case mjOBJ_BODY:
if (!object) {
return bodies_[0]->spec.element;
} else if (object == bodies_[0]->spec.element) {
return bodies_[0]->NextChild(NULL, type, /*recursive=*/true);
} else {
return bodies_[0]->NextChild(object, type, /*recursive=*/true);
}
case mjOBJ_SITE:
case mjOBJ_GEOM:
case mjOBJ_JOINT:
src/user/user_objects.cc
+3
View File
@@ -1338,6 +1338,9 @@ mjsElement* mjCBody::NextChild(mjsElement* child, mjtObj type, bool recursive) {
candidate = GetNext(frames, child, recursive);
break;
default:
throw mjCError(this,
"Body.NextChild supports the types: body, frame, geom, "
"site, light, camera");
break;
}
test/user/user_api_test.cc
+16 -11
View File
@@ -88,20 +88,25 @@ TEST_F(MujocoTest, TreeTraversal) {
mjs_setString(geom3->name, "geom3");
mjs_setString(site4->name, "site4");
bool recursive = false;
mjsElement* b_el0 = mjs_firstElement(spec, mjOBJ_BODY);
mjsElement* b_el1 = mjs_nextElement(spec, b_el0);
mjsElement* a_el0 = mjs_firstElement(spec, mjOBJ_ACTUATOR);
mjsElement* l_el0 = mjs_firstElement(spec, mjOBJ_LIGHT);
mjsElement* c_el0 = mjs_firstChild(body, mjOBJ_CAMERA);
mjsElement* t_el0 = mjs_firstChild(body, mjOBJ_TENDON);
mjsElement* s_el1 = mjs_firstChild(body, mjOBJ_SITE);
mjsElement* s_el2 = mjs_nextChild(body, s_el1);
mjsElement* s_el3 = mjs_nextChild(body, s_el2);
mjsElement* s_el0 = mjs_nextChild(body, s_el3);
mjsElement* s_el4 = mjs_firstChild(body1, mjOBJ_SITE);
mjsElement* g_el1 = mjs_firstChild(body, mjOBJ_GEOM);
mjsElement* g_el2 = mjs_nextChild(body, g_el1);
mjsElement* g_el3 = mjs_nextChild(body, g_el2);
mjsElement* g_el0 = mjs_nextChild(body, g_el3);
mjsElement* c_el0 = mjs_firstChild(body, mjOBJ_CAMERA, recursive);
mjsElement* t_el0 = mjs_firstChild(body, mjOBJ_TENDON, recursive);
mjsElement* s_el1 = mjs_firstChild(body, mjOBJ_SITE, recursive);
mjsElement* s_el2 = mjs_nextChild(body, s_el1, recursive);
mjsElement* s_el3 = mjs_nextChild(body, s_el2, recursive);
mjsElement* s_el0 = mjs_nextChild(body, s_el3, recursive);
mjsElement* s_el4 = mjs_firstChild(body1, mjOBJ_SITE, recursive);
mjsElement* g_el1 = mjs_firstChild(body, mjOBJ_GEOM, recursive);
mjsElement* g_el2 = mjs_nextChild(body, g_el1, recursive);
mjsElement* g_el3 = mjs_nextChild(body, g_el2, recursive);
mjsElement* g_el0 = mjs_nextChild(body, g_el3, recursive);
EXPECT_EQ(b_el0, world->element);
EXPECT_EQ(b_el1, body->element);
EXPECT_EQ(a_el0, nullptr);
EXPECT_EQ(l_el0, nullptr);
EXPECT_EQ(c_el0, nullptr);