Add optional kwarg inputs to mjSpec Python bindings for mjsBody add constructors.
PiperOrigin-RevId: 675141078 Change-Id: I8d735ae7d0a018ce7cd1aca56c899d1455586015
This commit is contained in:
Taylor Howell
committed by
Copybara-Service
Copybara-Service
parent
2ffa0fe942
commit
b0c7101f1c
@@ -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<std::string> 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<char>* des) {
|
||||
if (kwargs.contains(str)) {
|
||||
try {
|
||||
*des = kwargs[str].cast<std::string>();
|
||||
} 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<decltype(*des)>::value_type;
|
||||
std::vector<T> vec = kwargs[str].cast<std::vector<T>>();
|
||||
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::decay_t<decltype(des)>>;
|
||||
std::vector<T> array = kwargs[str].cast<std::vector<T>>();
|
||||
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<decltype(des)>;
|
||||
des = kwargs[str].cast<T>();
|
||||
} 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)
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -98,6 +98,88 @@ class SpecsTest(absltest.TestCase):
|
||||
</mujoco>
|
||||
"""),)
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user