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

Add mjsPlugin and mjsOrientation related keyword argument support for mjSpec Python bindings.

Browse Source 
PiperOrigin-RevId: 676412967
Change-Id: I8cc1bfcc443327fde33246981400b375f8d4ff3d
This commit is contained in:
Taylor Howell
2024-09-19 07:54:24 -07:00
committed by Copybara-Service
parent f4381e12a2
commit 6c7e1095c5
3 changed files with 265 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/codegen/generate_spec_bindings.py
+129 -13
View File
@@ -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<double> array = kwargs[str].cast<std::vector<double>>();
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<raw::MjsPlugin> input = kwargs["plugin"].cast<raw::MjsPlugin>();
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<char>* des) {
if (kwargs.contains(str)) {
@@ -393,7 +506,10 @@ def generate_add() -> None:
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) + ".");
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) {
python/mujoco/specs.cc
+2 -2
View File
@@ -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;
python/mujoco/specs_test.py
+134 -13
View File
@@ -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'