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

Enable frames to create the same children as bodies.

Browse Source 
PiperOrigin-RevId: 822097998
Change-Id: I8477d8abe2d6f711421126825fd94d39af1a4e07
This commit is contained in:
Alessio Quaglino
2025-10-21 06:56:15 -07:00
committed by Copybara-Service
parent 09682ade4b
commit 056945df20
2 changed files with 40 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/codegen/generate_spec_bindings.py
+26 -1
View File
@@ -52,6 +52,15 @@ SPECS = [
('mjsExclude', 'Spec', False, 'excludes', 'mjOBJ_EXCLUDE'),
('mjsPlugin', 'Spec', False, 'plugins', 'mjOBJ_PLUGIN'),
]
SPECS_ADD = SPECS + [
('mjsBody', 'Frame', True, 'bodies', 'mjOBJ_BODY'),
('mjsSite', 'Frame', True, 'sites', 'mjOBJ_SITE'),
('mjsGeom', 'Frame', True, 'geoms', 'mjOBJ_GEOM'),
('mjsJoint', 'Frame', True, 'joints', 'mjOBJ_JOINT'),
('mjsCamera', 'Frame', True, 'cameras', 'mjOBJ_CAMERA'),
('mjsFrame', 'Frame', True, 'frames', 'mjOBJ_FRAME'),
('mjsLight', 'Frame', True, 'lights', 'mjOBJ_LIGHT'),
]
# pylint: enable=bad-whitespace
@@ -301,7 +310,7 @@ def generate() -> None:
def generate_add() -> None:
"""Generate add constructors with optional keyword arguments."""
for key, parent, default, listname, objtype in SPECS:
for key, parent, default, listname, objtype in SPECS_ADD:
def _field(f: ast_nodes.StructFieldDecl):
if f.type == ast_nodes.PointerType(
@@ -390,6 +399,22 @@ def generate_add() -> None:
raw::MjsDefault* default_, py::kwargs kwargs) -> raw::{titlecase}* {{
auto out = mjs_add{elem}(&self, default_);
"""
elif parent == 'Frame':
if key == 'mjsFrame':
code = f"""
{'mjs' + parent}.def("add_{elemlower}", []({'raw::Mjs' + parent}& self,
raw::MjsFrame* parentframe_, py::kwargs kwargs) -> raw::{titlecase}* {{
raw::MjsBody* body = mjs_getParent(self.element);
auto out = mjs_add{elem}(body, &self);
"""
else:
code = f"""
{'mjs' + parent}.def("add_{elemlower}", []({'raw::Mjs' + parent}& self,
raw::MjsDefault* default_, py::kwargs kwargs) -> raw::{titlecase}* {{
raw::MjsBody* body = mjs_getParent(self.element);
auto out = mjs_add{elem}(body, default_);
mjs_setFrame(out->element, &self);
"""
else:
raise NotImplementedError(f'{parent} parent is not implement.')
python/mujoco/specs_test.py
+14
View File
@@ -318,6 +318,14 @@ class SpecsTest(absltest.TestCase):
np.testing.assert_array_equal(frameb1.pos, frameb0.pos)
np.testing.assert_array_equal(frameb1.quat, frameb0.quat)
# Add frame in frame.
framec0 = frameb0.add_frame(name='framec', pos=[7, 8, 9], quat=[0, 0, 0, 1])
self.assertEqual(framec0.name, 'framec')
self.assertEqual(framec0.parent, frameb0.parent)
self.assertEqual(framec0.frame, frameb0)
np.testing.assert_array_equal(framec0.pos, [7, 8, 9])
np.testing.assert_array_equal(framec0.quat, [0, 0, 0, 1])
# Add joint.
joint = body.add_joint(type=mujoco.mjtJoint.mjJNT_HINGE, axis=[0, 1, 0])
self.assertEqual(joint.type, mujoco.mjtJoint.mjJNT_HINGE)
@@ -340,6 +348,12 @@ class SpecsTest(absltest.TestCase):
light = body.add_light(attenuation=[1, 2, 3])
np.testing.assert_array_equal(light.attenuation, [1, 2, 3])
# Add light in a frame.
light_in_frame = framea0.add_light(cutoff=10)
self.assertEqual(light_in_frame.cutoff, 10)
self.assertEqual(light_in_frame.parent, framea0.parent)
self.assertEqual(light_in_frame.frame, framea0)
# Invalid input for valid keyword argument.
with self.assertRaises(ValueError) as cm:
body.add_geom(pos='pos')