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

PiperOrigin-RevId: 676412967
Change-Id: I8cc1bfcc443327fde33246981400b375f8d4ff3d

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
+      if f.type == ast_nodes.PointerType(
-      unsupported = (
+          inner_type=ast_nodes.ValueType(name='mjsElement')
-          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.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

@@ -497,8 +497,8 @@ PYBIND11_MODULE(_specs, m) {
             }
           } else {
             throw pybind11::type_error(
-                "Invalid '" + key +
+                "Invalid " + key +
-                "' keyword argument. Valid options are: align, group, name.");
+                " keyword argument. Valid options are: align, group, name.");
           }
         }
         return out;

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):
+    with self.assertRaises(TypeError) as cm:
-      body.add_freejoint(axis=[1, 2, 3])  # invalid keyword argument
+      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):
+    with self.assertRaises(ValueError) as cm:
-      body.add_geom(pos='pos')  # wrong type for array
+      body.add_geom(pos='pos')
+    self.assertEqual(
+        str(cm.exception),
+        'pos should be a list/array.',
+    )
 
-    with self.assertRaises(ValueError):
+    with self.assertRaises(ValueError) as cm:
-      body.add_geom(pos=[0, 1])  # wrong size
+      body.add_geom(pos=[0, 1])
+    self.assertEqual(
+        str(cm.exception),
+        'pos should be a list/array of size 3.',
+    )
 
-    with self.assertRaises(ValueError):
+    with self.assertRaises(ValueError) as cm:
-      body.add_geom(type='type')  # wrong type for value
+      body.add_geom(type='type')
+    self.assertEqual(
+        str(cm.exception),
+        'type is the wrong type.',
+    )
 
-    with self.assertRaises(ValueError):
+    with self.assertRaises(ValueError) as cm:
-      body.add_geom(userdata='')  # wrong type of vector
+      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'