From 89aea5e21794a0eb615785df66ee769f2dd6fd63 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino
Date: Tue, 1 Oct 2024 08:56:03 -0700
Subject: [PATCH] Add `body.find_all()` to Python bindings.
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
---
doc/APIreference/functions.rst | 3 +-
doc/includes/references.h | 4 +-
include/mujoco/mujoco.h | 7 ++--
introspect/functions.py | 12 +++++-
python/mujoco/specs.cc | 72 +++++++++++++++++++++++++++-------
python/mujoco/specs_test.py | 36 ++++++++++++++---
src/user/user_api.cc | 18 +++++++--
src/user/user_api.h | 7 ++--
src/user/user_model.cc | 7 ++++
src/user/user_objects.cc | 3 ++
test/user/user_api_test.cc | 27 +++++++------
11 files changed, 150 insertions(+), 46 deletions(-)
diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst
index ed317aea..fbec781a 100644
--- a/doc/APIreference/functions.rst
+++ b/doc/APIreference/functions.rst
@@ -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:
diff --git a/doc/includes/references.h b/doc/includes/references.h
index 5f5ecafd..19ea8a36 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -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);
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index e9a41e0d..0df9b7e9 100644
--- a/include/mujoco/mujoco.h
+++ b/include/mujoco/mujoco.h
@@ -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);
diff --git a/introspect/functions.py b/introspect/functions.py
index 08f3d44c..25f21010 100644
--- a/introspect/functions.py
+++ b/introspect/functions.py
@@ -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(
diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc
index 01e9aaf8..8a1e7ecc 100644
--- a/python/mujoco/specs.cc
+++ b/python/mujoco/specs.cc
@@ -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(
diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
index 5c9c160e..68e0f0cb 100644
--- a/python/mujoco/specs_test.py
+++ b/python/mujoco/specs_test.py
@@ -580,7 +580,9 @@ class SpecsTest(absltest.TestCase):
-
+
+
+
@@ -588,11 +590,33 @@ class SpecsTest(absltest.TestCase):
"""
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()
diff --git a/src/user/user_api.cc b/src/user/user_api.cc
index a401f131..55472226 100644
--- a/src/user/user_api.cc
+++ b/src/user/user_api.cc
@@ -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(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(body->element);
- return bodyC->NextChild(child);
+ try {
+ return bodyC->NextChild(child, child->elemtype, recurse);
+ } catch(mjCError& e) {
+ bodyC->model->SetError(e);
+ return nullptr;
+ }
}
diff --git a/src/user/user_api.h b/src/user/user_api.h
index 43cff0c5..8c55c0ff 100644
--- a/src/user/user_api.h
+++ b/src/user/user_api.h
@@ -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);
diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index a69dfbc4..1a6bfbd7 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -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:
diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc
index 35bd3645..833a8106 100644
--- a/src/user/user_objects.cc
+++ b/src/user/user_objects.cc
@@ -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;
}
diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc
index 6e7fd45f..a4a7d425 100644
--- a/test/user/user_api_test.cc
+++ b/test/user/user_api_test.cc
@@ -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);