Add MjSpec.encode method.

PiperOrigin-RevId: 907177387
Change-Id: I65800298798037c4a0c0cd79ec86706be67db51e

python/mujoco/specs.cc

@@ -431,6 +431,51 @@ PYBIND11_MODULE(_specs, m) {
       throw FatalError(std::string(err.data()));
     }
   });
+  mjSpec.def(
+      "encode",
+      [](MjSpec& self, std::string filename,
+         std::optional<py::object> model,
+         std::optional<std::string> content_type) -> int {
+        raw::MjModel* m = nullptr;
+        if (model.has_value() && !model->is_none()) {
+          auto& wrapper =
+              py::cast<_impl::MjModelWrapper&>(*model);
+          m = wrapper.get();
+        }
+
+        mjVFS vfs;
+        mjVFS* vfs_ptr = nullptr;
+        if (!self.assets.empty()) {
+          mj_defaultVFS(&vfs);
+          vfs_ptr = &vfs;
+          for (const auto& asset : self.assets) {
+            std::string buffer_name =
+                py::cast<std::string>(asset.first);
+            std::string buffer =
+                py::cast<std::string>(asset.second);
+            mj_addBufferVFS(vfs_ptr, buffer_name.c_str(),
+                            buffer.c_str(), buffer.size());
+          }
+        }
+
+        std::array<char, 1024> err;
+        err[0] = '\0';
+        const char* ct =
+            content_type.has_value() ? content_type->c_str() : nullptr;
+        int nbytes = mj_encode(self.ptr, m, filename.c_str(), ct,
+                               vfs_ptr, err.data(), err.size());
+
+        if (vfs_ptr) {
+          mj_deleteVFS(vfs_ptr);
+        }
+
+        if (nbytes < 0) {
+          throw FatalError(std::string(err.data()));
+        }
+        return nbytes;
+      },
+      py::arg("filename"), py::arg("model") = py::none(),
+      py::arg("content_type") = py::none());
   mjSpec.def(
       "add_default",
       [](MjSpec* spec, std::string& classname,

python/mujoco/specs_test.py

@@ -22,6 +22,7 @@ import textwrap
 import typing
 import zipfile  # pylint: disable=unused-import
 
+from absl import flags
 from absl.testing import absltest
 from etils import epath
 import mujoco
@@ -34,6 +35,11 @@ def get_linenumber():
 
 
 class SpecsTest(absltest.TestCase):
+  def setUp(self):
+    super().setUp()
+    # Mark flags as parsed to avoid pytest errors about unparsed flags.
+    # This is needed for `create_tempdir()` calls below.
+    flags.FLAGS.mark_as_parsed()
 
   def test_typing(self):
     spec = mujoco.MjSpec()
@@ -1964,5 +1970,49 @@ class SpecsTest(absltest.TestCase):
     self.assertGreater(cam_sd[0], cam_sd[1])  # dist > depth
     self.assertAlmostEqual(cam_sd[1], 2.0, places=6)  # depth is still 2.0
 
+  def test_encode_xml(self):
+    # Create a simple spec and compile.
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body()
+    geom = body.add_geom()
+    geom.size[0] = 1
+    model = spec.compile()
+
+    # Encode to XML.
+    filename = os.path.join(self.create_tempdir().full_path, 'output.xml')
+    nbytes = spec.encode(filename, model)
+    self.assertGreater(nbytes, 0)
+
+    # Verify the output is valid XML that can be loaded.
+    reloaded = mujoco.MjSpec.from_file(filename)
+    reloaded_model = reloaded.compile()
+    self.assertEqual(reloaded_model.ngeom, model.ngeom)
+
+  def test_encode_xml_without_model(self):
+    # Create a simple spec and compile so XML can be written.
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body()
+    geom = body.add_geom()
+    geom.size[0] = 1
+    spec.compile()
+
+    # Encode to XML without passing a model explicitly.
+    filename = os.path.join(self.create_tempdir().full_path, 'output.xml')
+    nbytes = spec.encode(filename)
+    self.assertGreater(nbytes, 0)
+
+  def test_encode_no_encoder_raises(self):
+    # Create a simple spec and compile.
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body()
+    geom = body.add_geom()
+    geom.size[0] = 1
+    model = spec.compile()
+
+    # Encode with an unknown extension should fail.
+    filename = os.path.join(self.create_tempdir().full_path, 'output.unknown')
+    with self.assertRaises(mujoco.FatalError):
+      spec.encode(filename, model)
+
 if __name__ == '__main__':
   absltest.main()