diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index 06d217a5..e6983263 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -262,16 +262,26 @@ def generate_add() -> None: ]: def _field(f: ast_nodes.StructFieldDecl): - # TODO(taylorhowell): add support for mjsOrientation and mjsPlugin - unsupported = ( - ast_nodes.PointerType( - inner_type=ast_nodes.ValueType(name='mjsElement') - ), - ast_nodes.ValueType(name='mjsOrientation'), - ast_nodes.ValueType(name='mjsPlugin'), - ) - if f.type in unsupported: + if f.type == ast_nodes.PointerType( + inner_type=ast_nodes.ValueType(name='mjsElement') + ): return '', '', '' + elif f.type == ast_nodes.ValueType(name='mjsPlugin'): + return f'set_plugin(out->{f.name});', 'plugin', f.name + elif f.type == ast_nodes.ValueType(name='mjsOrientation'): + return ( + ( + f'set_orientation(out->{f.name},' + f' "{"iaxisangle" if f.name == "ialt" else "axisangle"}",' + f' "{"ixyaxes" if f.name == "ialt" else "xyaxes"}",' + f' "{"izaxis" if f.name == "ialt" else "zaxis"}",' + f' "{"ieuler" if f.name == "ialt" else "euler"}");' + ), + 'orientation', + ['iaxisangle', 'ixyaxes', 'izaxis', 'ieuler'] + if f.name == 'ialt' + else ['axisangle', 'xyaxes', 'zaxis', 'euler'], + ) elif f.type == ast_nodes.PointerType( inner_type=ast_nodes.ValueType(name='mjString') ): @@ -297,7 +307,10 @@ def generate_add() -> None: if line: code_field = code_field + '\n ' + line set_types.append(set_type) - names.append(name) + if set_type == 'orientation': + names.extend(name) + else: + names.append(name) # assemble elem = key.removeprefix('mjs') @@ -347,7 +360,9 @@ def generate_add() -> None: for (auto item: kwarg_dict) {{ std::string key = py::str(item.first); if (valid_kwargs.count(key) == 0) {{ - throw pybind11::type_error("Invalid '" + key + "' keyword argument. Valid options are: {", ".join(names)}."); + throw pybind11::type_error("Invalid " + + key + + " keyword argument. Valid options are: {", ".join(names)}."); }} }} """ @@ -355,7 +370,105 @@ def generate_add() -> None: # include helper functions if set_types: for t in set(set_types): - if t == 'string': + if t == 'orientation': + code += """\n + auto set_orientation = [&kwargs](raw::MjsOrientation& orientation, + const char* axisangle, + const char* xyaxes, + const char* zaxis, + const char* euler) { + int nrepresentation = 0; + bool has_axisangle = kwargs.contains(axisangle); + nrepresentation += has_axisangle; + bool has_xyaxes = kwargs.contains(xyaxes); + nrepresentation += has_xyaxes; + bool has_zaxis = kwargs.contains(zaxis); + nrepresentation += has_zaxis; + bool has_euler = kwargs.contains(euler); + nrepresentation += has_euler; + + if (nrepresentation == 0) { + return; + } else if (nrepresentation > 1) { + throw pybind11::value_error("Only one of: " + + std::string(axisangle) + ", " + + std::string(xyaxes) + ", " + + std::string(zaxis) + + ", or" + + std::string(euler) + + " can be set."); + } + + auto set_array = [&kwargs](const char* str, double* des, int size) { + try { + std::vector array = kwargs[str].cast>(); + if (array.size() != size) { + throw pybind11::value_error(std::string(str) + + " should be a list/array of size " + + std::to_string(size) + + "."); + } + int idx = 0; + for (auto val : array) { + des[idx++] = val; + } + } catch (const py::cast_error &e) { + throw pybind11::value_error(std::string(str) + + " should be a list/array."); + } + }; + + if (has_axisangle) { + set_array(axisangle, orientation.axisangle, 4); + orientation.type = mjORIENTATION_AXISANGLE; + } else if (has_xyaxes) { + set_array(xyaxes, orientation.xyaxes, 6); + orientation.type = mjORIENTATION_XYAXES; + } else if (has_zaxis) { + set_array(zaxis, orientation.zaxis, 3); + orientation.type = mjORIENTATION_ZAXIS; + } else if (has_euler) { + set_array(euler, orientation.euler, 3); + orientation.type = mjORIENTATION_EULER; + } + }; + """ + elif t == 'plugin': + code += """\n + auto set_plugin = [&kwargs](raw::MjsPlugin& plugin) { + if (kwargs.contains("plugin")) { + std::optional input = kwargs["plugin"].cast(); + if (input.has_value()) { + try { + plugin.name = input->name; + } catch (const py::cast_error &e) { + throw pybind11::value_error("plugin.name should be a string."); + } + try { + plugin.instance_name = input->instance_name; + } catch (const py::cast_error &e) { + throw pybind11::value_error("plugin.instance_name should be a string."); + } + try { + plugin.plugin_slot = input->plugin_slot; + } catch (const py::cast_error &e) { + throw pybind11::value_error("plugin.plugin_slot should be an int."); + } + try { + plugin.active = input->active; + } catch (const py::cast_error &e) { + throw pybind11::value_error("plugin.active should be an mjtByte."); + } + try { + plugin.info = input->info; + } catch (const py::cast_error &e) { + throw pybind11::value_error("plugin.info should be a string."); + } + } + } + }; + """ + elif t == 'string': code += """\n auto set_string = [&kwargs](const char* str, std::basic_string* des) { if (kwargs.contains(str)) { @@ -393,7 +506,10 @@ def generate_add() -> None: using T = std::remove_pointer_t>; std::vector array = kwargs[str].cast>(); if (array.size() != size) { - throw pybind11::value_error(std::string(str) + " should be a list/array of size " + std::to_string(size) + "."); + throw pybind11::value_error(std::string(str) + + " should be a list/array of size " + + std::to_string(size) + + "."); } int idx = 0; for (auto val : array) { diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index db35d826..5dd4621e 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -497,8 +497,8 @@ PYBIND11_MODULE(_specs, m) { } } else { throw pybind11::type_error( - "Invalid '" + key + - "' keyword argument. Valid options are: align, group, name."); + "Invalid " + key + + " keyword argument. Valid options are: align, group, name."); } } return out; diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 33363b99..fb3a1321 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -181,9 +181,18 @@ class SpecsTest(absltest.TestCase): self.assertEqual(sensor.objtype, mujoco.mjtObj.mjOBJ_SITE) # Add plugin. - plugin = spec.add_plugin(plugin_slot=7, instance_name='plugin') + plugin = spec.add_plugin( + name='name', + instance_name='instance', + plugin_slot=7, + active=True, + info='info', + ) + self.assertEqual(plugin.name, 'name') + self.assertEqual(plugin.instance_name, 'instance') self.assertEqual(plugin.plugin_slot, 7) - self.assertEqual(plugin.instance_name, 'plugin') + self.assertEqual(plugin.active, True) + self.assertEqual(plugin.info, 'info') # Add a body. body = spec.worldbody.add_body( @@ -193,6 +202,14 @@ class SpecsTest(absltest.TestCase): np.testing.assert_array_equal(body.pos, [1, 2, 3]) np.testing.assert_array_equal(body.quat, [0, 0, 0, 1]) + # Add a body with a plugin. + body_with_plugin = spec.worldbody.add_body(plugin=plugin) + self.assertEqual(body_with_plugin.plugin.name, 'name') + self.assertEqual(body_with_plugin.plugin.instance_name, 'instance') + self.assertEqual(body_with_plugin.plugin.plugin_slot, 7) + self.assertEqual(body_with_plugin.plugin.active, True) + self.assertEqual(body_with_plugin.plugin.info, 'info') + # Add a geom. geom = body.add_geom( name='geom', @@ -267,29 +284,133 @@ class SpecsTest(absltest.TestCase): freejoint_align = body.add_freejoint(align=True) self.assertEqual(freejoint_align.align, True) - with self.assertRaises(TypeError): - body.add_freejoint(axis=[1, 2, 3]) # invalid keyword argument + with self.assertRaises(TypeError) as cm: + body.add_freejoint(axis=[1, 2, 3]) + self.assertEqual( + str(cm.exception), + 'Invalid axis keyword argument. Valid options are: align, group, name.', + ) # Add light. light = body.add_light(attenuation=[1, 2, 3]) np.testing.assert_array_equal(light.attenuation, [1, 2, 3]) # Invalid input for valid keyword argument. - with self.assertRaises(ValueError): - body.add_geom(pos='pos') # wrong type for array + with self.assertRaises(ValueError) as cm: + body.add_geom(pos='pos') + self.assertEqual( + str(cm.exception), + 'pos should be a list/array.', + ) - with self.assertRaises(ValueError): - body.add_geom(pos=[0, 1]) # wrong size + with self.assertRaises(ValueError) as cm: + body.add_geom(pos=[0, 1]) + self.assertEqual( + str(cm.exception), + 'pos should be a list/array of size 3.', + ) - with self.assertRaises(ValueError): - body.add_geom(type='type') # wrong type for value + with self.assertRaises(ValueError) as cm: + body.add_geom(type='type') + self.assertEqual( + str(cm.exception), + 'type is the wrong type.', + ) - with self.assertRaises(ValueError): - body.add_geom(userdata='') # wrong type of vector + with self.assertRaises(ValueError) as cm: + body.add_geom(userdata='') + self.assertEqual( + str(cm.exception), + 'userdata has the wrong type.', + ) # Invalid keyword argument. - with self.assertRaises(TypeError): + with self.assertRaises(TypeError) as cm: body.add_geom(vel='vel') + self.assertEqual( + str(cm.exception), + 'Invalid vel keyword argument. Valid options are: name, type, pos,' + ' quat, axisangle, xyaxes, zaxis, euler, fromto, size, contype,' + ' conaffinity, condim, priority, friction, solmix, solref, solimp,' + ' margin, gap, mass, density, typeinertia, fluid_ellipsoid,' + ' fluid_coefs, material, rgba, group, hfieldname, meshname, fitscale,' + ' userdata, plugin, info.', + ) + + # Orientation keyword arguments. + geom_axisangle = body.add_geom(axisangle=[1, 2, 3, 4]) + geom_xyaxes = body.add_geom(xyaxes=[1, 2, 3, 4, 5, 6]) + geom_zaxis = body.add_geom(zaxis=[1, 2, 3]) + geom_euler = body.add_geom(euler=[1, 2, 3]) + + self.assertEqual( + geom_axisangle.alt.type, mujoco.mjtOrientation.mjORIENTATION_AXISANGLE + ) + self.assertEqual( + geom_xyaxes.alt.type, mujoco.mjtOrientation.mjORIENTATION_XYAXES + ) + self.assertEqual( + geom_zaxis.alt.type, mujoco.mjtOrientation.mjORIENTATION_ZAXIS + ) + self.assertEqual( + geom_euler.alt.type, mujoco.mjtOrientation.mjORIENTATION_EULER + ) + np.testing.assert_array_equal(geom_axisangle.alt.axisangle, [1, 2, 3, 4]) + np.testing.assert_array_equal(geom_xyaxes.alt.xyaxes, [1, 2, 3, 4, 5, 6]) + np.testing.assert_array_equal(geom_zaxis.alt.zaxis, [1, 2, 3]) + np.testing.assert_array_equal(geom_euler.alt.euler, [1, 2, 3]) + + body_iaxisangle = spec.worldbody.add_body(iaxisangle=[1, 2, 3, 4]) + body_ixyaxes = spec.worldbody.add_body(ixyaxes=[1, 2, 3, 4, 5, 6]) + body_izaxis = spec.worldbody.add_body(izaxis=[1, 2, 3]) + body_ieuler = spec.worldbody.add_body(ieuler=[1, 2, 3]) + body_euler_ieuler = spec.worldbody.add_body( + euler=[1, 2, 3], ieuler=[4, 5, 6] + ) + np.testing.assert_array_equal(body_iaxisangle.ialt.axisangle, [1, 2, 3, 4]) + np.testing.assert_array_equal(body_ixyaxes.ialt.xyaxes, [1, 2, 3, 4, 5, 6]) + np.testing.assert_array_equal(body_izaxis.ialt.zaxis, [1, 2, 3]) + np.testing.assert_array_equal(body_ieuler.ialt.euler, [1, 2, 3]) + np.testing.assert_array_equal(body_euler_ieuler.alt.euler, [1, 2, 3]) + np.testing.assert_array_equal(body_euler_ieuler.ialt.euler, [4, 5, 6]) + + # Test invalid orientation keyword arguments. + with self.assertRaises(ValueError) as cm: + body.add_geom(axisangle=[1, 2, 3]) + self.assertEqual( + str(cm.exception), + 'axisangle should be a list/array of size 4.', + ) + with self.assertRaises(ValueError) as cm: + body.add_geom(xyaxes=[1, 2, 3, 4, 5]) + self.assertEqual( + str(cm.exception), + 'xyaxes should be a list/array of size 6.', + ) + with self.assertRaises(ValueError) as cm: + body.add_geom(zaxis=[1, 2, 3, 4]) + self.assertEqual( + str(cm.exception), + 'zaxis should be a list/array of size 3.', + ) + with self.assertRaises(ValueError) as cm: + body.add_geom(euler=[1]) + self.assertEqual( + str(cm.exception), + 'euler should be a list/array of size 3.', + ) + with self.assertRaises(ValueError) as cm: + body.add_geom(axisangle=[1, 2, 3, 4], euler=[1, 2, 3]) + self.assertEqual( + str(cm.exception), + 'Only one of: axisangle, xyaxes, zaxis, oreuler can be set.', + ) + with self.assertRaises(ValueError) as cm: + spec.worldbody.add_body(iaxisangle=[1, 2, 3, 4], ieuler=[1, 2, 3]) + self.assertEqual( + str(cm.exception), + 'Only one of: iaxisangle, ixyaxes, izaxis, orieuler can be set.', + ) def test_load_xml(self): filename = '../../test/testdata/model.xml'