diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index 2ac03612..97854082 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -132,7 +132,7 @@ def _ptr_binding_code( vartype == 'mjDoubleVec' or vartype == 'mjFloatVec' or vartype == 'mjIntVec' - ): + ): vartype = vartype.replace('mj', '').replace('Vec', '').lower() return f"""\ {classname}.def_property( @@ -232,10 +232,171 @@ def generate() -> None: print(code) +def generate_body_add() -> None: + """Generate add functions for bodies.""" + for key in [ + 'mjsSite', + 'mjsGeom', + 'mjsJoint', + 'mjsLight', + 'mjsCamera', + 'mjsBody', + ]: + + 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: + return '', '', '' + elif f.type == ast_nodes.PointerType( + inner_type=ast_nodes.ValueType(name='mjString') + ): + return f'set_string("{f.name}", out->{f.name});', 'string', f.name + elif isinstance(f.type, ast_nodes.PointerType): + return f'set_vec("{f.name}", out->{f.name});', 'vec', f.name + elif isinstance(f.type, ast_nodes.ArrayType): + return ( + ( + f'set_array("{f.name}", out->{f.name},' + f' {f.type.extents[0]});' + ), + 'array', f.name + ) + elif isinstance(f.type, ast_nodes.ValueType): + return f'set_value("{f.name}", out->{f.name});', 'value', f.name + else: + return '', '', '' + + code_field = '' + set_types = [] + names = [] + for field in structs.STRUCTS[key].fields: + line, set_type, name = _field(field) + if line: + code_field = code_field + '\n ' + line + set_types.append(set_type) + names.append(name) + + # assemble + elem = key.removeprefix('mjs') + elemlower = elem.lower() + titlecase = 'Mjs' + elem + + # function definition and call to mjs_add_ + code = f""" + mjsBody.def("add_{elemlower}", [](raw::MjsBody& self, raw::MjsDefault* default_, py::kwargs kwargs) -> raw::{titlecase}* {{ + auto out = mjs_add{elem}(&self, default_); + """ + + # check for valid kwargs + code += '\n std::set valid_kwargs = {' + valid_kwargs = '' + for i, name in enumerate(names): + valid_kwargs += f'"{name}"' + if i != len(names) - 1: + valid_kwargs += ', ' + code += valid_kwargs + '};' + + code += f"""\n + py::dict kwarg_dict = kwargs; + 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)}."); + }} + }} + """ + + # include helper functions + if set_types: + for t in set(set_types): + if t == 'string': + code += """\n + auto set_string = [&kwargs](const char* str, std::basic_string* des) { + if (kwargs.contains(str)) { + try { + *des = kwargs[str].cast(); + } catch (const py::cast_error &e) { + throw pybind11::value_error(std::string(str) + " should be a string."); + } + } + }; + """ + elif t == 'vec': + code += """\n + auto set_vec = [&kwargs](const char* str, auto&& des) { + if (kwargs.contains(str)) { + try { + using T = typename std::decay_t::value_type; + std::vector vec = kwargs[str].cast>(); + des->clear(); + des->reserve(vec.size()); + for (auto val : vec) { + des->push_back(val); + } + } catch (const py::cast_error &e) { + throw pybind11::value_error(std::string(str) + " has the wrong type."); + } + } + }; + """ + elif t == 'array': + code += """\n + auto set_array = [&kwargs](const char* str, auto&& des, int size) { + if (kwargs.contains(str)) { + try { + 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) + "."); + } + 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."); + } + } + }; + """ + elif t == 'value': + code += """\n + auto set_value = [&kwargs](const char* str, auto&& des) { + if (kwargs.contains(str)) { + try { + using T = std::decay_t; + des = kwargs[str].cast(); + } catch (const py::cast_error &e) { + throw pybind11::value_error(std::string(str) + " is the wrong type."); + } + } + }; + """ + + code += code_field + code += """\n + return out; + }, + py::arg_v("default", nullptr), + py::return_value_policy::reference_internal); + """ + + print(code) + + def main(argv: Sequence[str]) -> None: if len(argv) > 1: raise app.UsageError('Too many command-line arguments.') generate() + generate_body_add() + if __name__ == '__main__': app.run(main) diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 27e7d819..4fd9e97f 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -556,13 +556,6 @@ PYBIND11_MODULE(_specs, m) { // ============================= MJSBODY ===================================== mjsBody.def_property_readonly( "id", [](raw::MjsBody& self) -> int { return mjs_getId(self.element); }); - mjsBody.def( - "add_body", - [](raw::MjsBody& self, raw::MjsDefault* default_) -> raw::MjsBody* { - return mjs_addBody(&self, default_); - }, - py::arg_v("default", nullptr), - py::return_value_policy::reference_internal); mjsBody.def( "add_frame", [](raw::MjsBody& self, raw::MjsFrame* parentframe_) -> raw::MjsFrame* { @@ -570,47 +563,12 @@ PYBIND11_MODULE(_specs, m) { }, py::arg_v("default", nullptr), py::return_value_policy::reference_internal); - mjsBody.def( - "add_geom", - [](raw::MjsBody& self, raw::MjsDefault* default_) -> raw::MjsGeom* { - return mjs_addGeom(&self, default_); - }, - py::arg_v("default", nullptr), - py::return_value_policy::reference_internal); - mjsBody.def( - "add_joint", - [](raw::MjsBody& self, raw::MjsDefault* default_) -> raw::MjsJoint* { - return mjs_addJoint(&self, default_); - }, - py::arg_v("default", nullptr), - py::return_value_policy::reference_internal); mjsBody.def( "add_freejoint", [](raw::MjsBody& self) -> raw::MjsJoint* { return mjs_addFreeJoint(&self); }, py::return_value_policy::reference_internal); - mjsBody.def( - "add_light", - [](raw::MjsBody& self, raw::MjsDefault* default_) -> raw::MjsLight* { - return mjs_addLight(&self, default_); - }, - py::arg_v("default", nullptr), - py::return_value_policy::reference_internal); - mjsBody.def( - "add_site", - [](raw::MjsBody& self, raw::MjsDefault* default_) -> raw::MjsSite* { - return mjs_addSite(&self, default_); - }, - py::arg_v("default", nullptr), - py::return_value_policy::reference_internal); - mjsBody.def( - "add_camera", - [](raw::MjsBody& self, raw::MjsDefault* default_) -> raw::MjsCamera* { - return mjs_addCamera(&self, default_); - }, - py::arg_v("default", nullptr), - py::return_value_policy::reference_internal); mjsBody.def("set_frame", [](raw::MjsBody& self, raw::MjsFrame& frame) -> void { mjs_setFrame(self.element, &frame); diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index c440d91d..15434c00 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -98,6 +98,88 @@ class SpecsTest(absltest.TestCase): """),) + def test_kwarg(self): + # Create a spec. + spec = mujoco.MjSpec() + + # Add a body. + body = spec.worldbody.add_body( + name='body', pos=[1, 2, 3], quat=[0, 0, 0, 1] + ) + self.assertEqual(body.name, 'body') + np.testing.assert_array_equal(body.pos, [1, 2, 3]) + np.testing.assert_array_equal(body.quat, [0, 0, 0, 1]) + + # Add a geom. + geom = body.add_geom( + name='geom', + pos=[3, 2, 1], + fromto=[1, 2, 3, 4, 5, 6], + contype=3, + ) + + self.assertEqual(geom.name, 'geom') + np.testing.assert_array_equal(geom.pos, [3, 2, 1]) + np.testing.assert_array_equal(geom.fromto, [1, 2, 3, 4, 5, 6]) + self.assertEqual(geom.contype, 3) + + # Add a site to the body with user data and read it back. + site = body.add_site( + name='sitename', + pos=[0, 1, 2], + quat=[1, 0, 0, 0], + fromto=[0, 1, 2, 3, 4, 5], + size=[3, 2, 1], + type=mujoco.mjtGeom.mjGEOM_BOX, + material='material', + group=7, + rgba=[1, 1, 1, 0.5], + userdata=[1, 2, 3, 4, 5, 6], + info='info', + ) + self.assertEqual(site.name, 'sitename') + np.testing.assert_array_equal(site.pos, [0, 1, 2]) + np.testing.assert_array_equal(site.quat, [1, 0, 0, 0]) + np.testing.assert_array_equal(site.fromto, [0, 1, 2, 3, 4, 5]) + np.testing.assert_array_equal(site.size, [3, 2, 1]) + self.assertEqual(site.type, mujoco.mjtGeom.mjGEOM_BOX) + self.assertEqual(site.material, 'material') + self.assertEqual(site.group, 7) + np.testing.assert_array_equal(site.rgba, [1, 1, 1, 0.5]) + np.testing.assert_array_equal(site.userdata, [1, 2, 3, 4, 5, 6]) + self.assertEqual(site.info, 'info') + + # Add camera. + cam = body.add_camera(orthographic=1, resolution=[10, 20]) + self.assertEqual(cam.orthographic, 1) + np.testing.assert_array_equal(cam.resolution, [10, 20]) + + # Add joint. + jnt = body.add_joint(type=mujoco.mjtJoint.mjJNT_HINGE, axis=[0, 1, 0]) + self.assertEqual(jnt.type, mujoco.mjtJoint.mjJNT_HINGE) + np.testing.assert_array_equal(jnt.axis, [0, 1, 0]) + + # 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): + body.add_geom(pos=[0, 1]) # wrong size + + with self.assertRaises(ValueError): + body.add_geom(type='type') # wrong type for value + + with self.assertRaises(ValueError): + body.add_geom(userdata='') # wrong type of vector + + # Invalid keyword argument. + with self.assertRaises(TypeError): + body.add_geom(vel='vel') + def test_load_xml(self): filename = '../../test/testdata/model.xml' state_type = mujoco.mjtState.mjSTATE_INTEGRATION