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

add find_geom property to mjSpec

Browse Source 
PiperOrigin-RevId: 716225290
Change-Id: Ie6e229b0e802e7df06148f994d598334592e8ad1
This commit is contained in:
Silvia Cruciani
2025-01-16 07:19:03 -08:00
committed by Copybara-Service
parent f899e717d4
commit 4903d5321c
2 changed files with 24 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/specs.cc
+6
View File
@@ -402,6 +402,12 @@ PYBIND11_MODULE(_specs, m) {
return mjs_findDefault(self.ptr, classname.c_str());
},
py::return_value_policy::reference_internal);
mjSpec.def(
"find_geom",
[](MjSpec& self, std::string& name) -> raw::MjsGeom* {
return mjs_asGeom(mjs_findElement(self.ptr, mjOBJ_GEOM, name.c_str()));
},
py::return_value_policy::reference_internal);
mjSpec.def("compile", [mjmodel_from_spec_ptr](MjSpec& self) -> py::object {
if (self.assets.empty()) {
return mjmodel_from_spec_ptr(reinterpret_cast<uintptr_t>(self.ptr));
python/mujoco/specs_test.py
+18
View File
@@ -691,6 +691,24 @@ class SpecsTest(absltest.TestCase):
body4.name = 'body4_new'
self.assertEqual(spec.bodies[4].name, 'body4_new')
def test_geom_list(self):
main_xml = """
<mujoco>
<worldbody>
<body name="body1"/>
</worldbody>
</mujoco>
"""
spec = mujoco.MjSpec.from_string(main_xml)
geom1 = spec.worldbody.add_geom(name='geom1')
geom2 = spec.worldbody.add_geom(name='geom2')
geom3 = spec.find_body('body1').add_geom(name='geom3')
self.assertEqual(spec.geoms, [geom1, geom2, geom3])
self.assertEqual(spec.find_geom('geom1'), geom1)
self.assertEqual(spec.find_geom('geom2'), geom2)
self.assertEqual(spec.find_geom('geom3'), geom3)
def test_iterators(self):
spec = mujoco.MjSpec()
geom1 = spec.worldbody.add_geom()