Add special case support for MJB and TXT files to mj_encode.

PiperOrigin-RevId: 939875974
Change-Id: I57fc753871b348a7f6342aca0a892e74acb74173

doc/APIreference/functions.rst

@@ -66,11 +66,20 @@ Parse spec from a file.
 
 .. mujoco-include:: mj_encode
 
-Encode spec/model to a file using a registered encoder.
+Encode :ref:`mjSpec` or :ref:`mjModel` to a file. The output format is determined by the file extension (case insensitive) or
+``content_type``. Returns the number of bytes written on success, -1 on failure.
 
-Returns the number of bytes written on success, -1 on failure.
+The following formats are supported natively, without a registered encoder:
 
-*Nullable:* ``m``, ``vfs``, ``error``
+- **MJCF XML** — extension: ``.xml``, content_type: ``text/xml``. If an :ref:`mjSpec` is provided, saves via
+  :ref:`mj_saveXML`. Otherwise falls back to :ref:`mj_saveLastXML`, which requires a compiled :ref:`mjModel`.
+- **MJB** — extension: ``.mjb``. MuJoCo binary format. Requires a compiled :ref:`mjModel`.
+- **TXT** — extension: ``.txt``, content_type: ``text/plain``. Human-readable text dump via :ref:`mj_printModel`.
+  Requires a compiled :ref:`mjModel`.
+
+For all other formats, a registered encoder is looked up via :ref:`mjp_findEncoder`.
+
+*Nullable:* ``s``, ``m``, ``vfs``, ``error``
 
 .. _mj_compile:
 

doc/APIreference/functions_override.rst

@@ -72,6 +72,23 @@ the required size. XML saving automatically compiles the spec before saving.
 
 Save spec to XML file, return 0 on success, -1 otherwise. XML saving requires that the spec first be compiled.
 
+.. _mj_encode:
+
+Encode :ref:`mjSpec` or :ref:`mjModel` to a file. The output format is determined by the file extension (case insensitive) or
+``content_type``. Returns the number of bytes written on success, -1 on failure.
+
+The following formats are supported natively, without a registered encoder:
+
+- **MJCF XML** — extension: ``.xml``, content_type: ``text/xml``. If an :ref:`mjSpec` is provided, saves via
+  :ref:`mj_saveXML`. Otherwise falls back to :ref:`mj_saveLastXML`, which requires a compiled :ref:`mjModel`.
+- **MJB** — extension: ``.mjb``. MuJoCo binary format. Requires a compiled :ref:`mjModel`.
+- **TXT** — extension: ``.txt``, content_type: ``text/plain``. Human-readable text dump via :ref:`mj_printModel`.
+  Requires a compiled :ref:`mjModel`.
+
+For all other formats, a registered encoder is looked up via :ref:`mjp_findEncoder`.
+
+*Nullable:* ``s``, ``m``, ``vfs``, ``error``
+
 .. _Mainsimulation:
 
 These are the main entry points to the simulator. Most users will only need to call :ref:`mj_step`, which computes

doc/changelog.rst

@@ -10,6 +10,8 @@ General
 - Added Nesterov momentum extrapolation with adaptive gradient restart (O'Donoghue-Candès) to the PGS solver,
   significantly improving convergence. Overall PGS now requires ~2x fewer iterations.
 
+* :ref:`mj_encode` now supports encoding of MJB and TXT files.
+
 .. admonition:: Breaking API changes
    :class: attention
 

src/user/user_api.cc

@@ -15,6 +15,7 @@
 #include "user/user_api.h"
 
 #include <algorithm>
+#include <cctype>
 #include <cmath>
 #include <cstddef>
 #include <cstdio>
@@ -145,20 +146,97 @@ mjSpec* mj_parse(const char* filename, const char* content_type,
   return spec;
 }
 
+// Encode spec or model as MJCF XML.
+//
+// If a spec is provided, it is saved directly to XML. This preserves the original
+// structure and any user modifications in the spec.
+// If spec is null, the model must be provided, and it is saved using mj_saveLastXML.
+// mj_saveLastXML reconstructs the XML from the compiled model state, which may
+// differ from the original XML (e.g., losing comments, reordering elements) but
+// reflects the actual compiled model.
+//
+// Returns file size in bytes on success, -1 on failure.
+static mjtSize encode_xml(const mjSpec* s, const mjModel* m, const char* filename,
+                          char* error, int error_sz) {
+  if (s) {
+    // Save directly from the spec
+    if (mj_saveXML(s, filename, error, error_sz) < 0) {
+      return -1;
+    }
+  } else {
+    if (!m) {
+      if (error) {
+        strncpy(error, "model is required for XML encoding when spec is null",
+                error_sz);
+        error[error_sz - 1] = '\0';
+      }
+      return -1;
+    }
+    // Reconstruct XML from the compiled model, this will copy values back
+    // from the mjModel into the last compiled spec and write that out
+    // to disk. If there was no last compiled spec, such as when loading from
+    // MJB, this will return fail and we return -1.
+    if (!mj_saveLastXML(filename, m, error, error_sz)) {
+      return -1;
+    }
+  }
+  return static_cast<mjtSize>(std::filesystem::file_size(filename));
+}
+
+// Encode model as MJB (MuJoCo binary format).
+// Requires a compiled model; spec-only encoding is not supported.
+// Returns file size in bytes on success, -1 on failure.
+static mjtSize encode_mjb(const mjModel* m, const char* filename, char* error, int error_sz) {
+  if (!m) {
+    if (error) {
+      strncpy(error, "model is required for MJB encoding", error_sz);
+      error[error_sz - 1] = '\0';
+    }
+    return -1;
+  }
+  mj_saveModel(m, filename, nullptr, 0);
+  return static_cast<mjtSize>(std::filesystem::file_size(filename));
+}
+
+// Encode model as human-readable TXT (via mj_printModel).
+// Requires a compiled model; spec-only encoding is not supported.
+// Returns file size in bytes on success, -1 on failure.
+static mjtSize encode_txt(const mjModel* m, const char* filename, char* error,
+                          int error_sz) {
+  if (!m) {
+    if (error) {
+      strncpy(error, "model is required for TXT encoding", error_sz);
+      error[error_sz - 1] = '\0';
+    }
+    return -1;
+  }
+  mj_printModel(m, filename);
+  return static_cast<mjtSize>(std::filesystem::file_size(filename));
+}
+
 // encode spec/model to file
 mjtSize mj_encode(const mjSpec* s, const mjModel* m, const char* filename,
                   const char* content_type, const mjVFS* vfs, char* error,
                   int error_sz) {
-  // TODO(shaves) Move MJCF and URDF to encoders/decoders.
+  // special case handling
+  // TODO(shaves) write encoder/decoder paths for MJCF, TXT, MJB
   auto filepath = mujoco::user::FilePath(filename);
-  if (filepath.Ext() == ".xml" ||
-      (content_type && std::strcmp(content_type, "text/xml") == 0)) {
-    int result = mj_saveXML(s, filename, error, error_sz);
-    if (result < 0) {
-      return -1;
-    }
+  std::string ext = filepath.Ext();
+  std::transform(ext.begin(), ext.end(), ext.begin(),
+                 [](unsigned char c) { return std::tolower(c); });
 
-    return static_cast<mjtSize>(std::filesystem::file_size(filename));
+  if (ext == ".xml" ||
+      (content_type && std::strcmp(content_type, "text/xml") == 0)) {
+    return encode_xml(s, m, filename, error, error_sz);
+  }
+
+  if (ext == ".mjb") {
+    return encode_mjb(m, filename, error, error_sz);
+  }
+
+  if (ext == ".txt" ||
+      (content_type && std::strcmp(content_type, "text/plain") == 0)) {
+    return encode_txt(m, filename, error, error_sz);
   }
 
   const mjpEncoder* encoder = mjp_findEncoder(filename, content_type);

test/user/user_api_test.cc

@@ -18,6 +18,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <filesystem>
 #include <functional>
 #include <map>
 #include <memory>
@@ -3537,5 +3538,69 @@ TEST_F(MujocoTest, CompilationWarningsClearedOnRecompile) {
   mj_deleteSpec(spec);
 }
 
+TEST_F(MujocoTest, MjEncodeNativeFormats) {
+  // simple test spec
+  mjSpec* spec = mj_makeSpec();
+  mjsBody* world = mjs_findBody(spec, "world");
+  mjsBody* body = mjs_addBody(world, nullptr);
+  mjsGeom* geom = mjs_addGeom(body, nullptr);
+  geom->size[0] = 1.0;
+  geom->size[1] = 1.0;
+  geom->size[2] = 1.0;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, NotNull());
+
+  std::filesystem::path tmp_dir = std::filesystem::temp_directory_path();
+  std::string xml_path = (tmp_dir / "test_encode.xml").string();
+  std::string mjb_path = (tmp_dir / "test_encode.mjb").string();
+  std::string txt_path = (tmp_dir / "test_encode.txt").string();
+
+  char error[1000];
+
+  // XML Encoding with spec
+  EXPECT_GT(mj_encode(spec, model, xml_path.c_str(), nullptr, nullptr, error,
+                      sizeof(error)),
+            0);
+  EXPECT_TRUE(std::filesystem::exists(xml_path));
+  EXPECT_GT(std::filesystem::file_size(xml_path), 0);
+
+  // XML Encoding without spec (should fail because no XML loaded in global
+  // spec)
+  std::filesystem::remove(xml_path);
+  // free global spec
+  mj_freeLastXML();
+  EXPECT_EQ(mj_encode(nullptr, model, xml_path.c_str(), nullptr, nullptr, error,
+                      sizeof(error)),
+            -1);
+  EXPECT_THAT(error, HasSubstr("No XML model loaded"));
+  EXPECT_TRUE(std::filesystem::exists(xml_path));
+  EXPECT_EQ(std::filesystem::file_size(xml_path), 0);
+  std::filesystem::remove(xml_path);
+
+  // MJB Encoding (spec can be null)
+  EXPECT_GT(mj_encode(nullptr, model, mjb_path.c_str(), nullptr, nullptr, error,
+                      sizeof(error)),
+            0);
+  EXPECT_TRUE(std::filesystem::exists(mjb_path));
+  EXPECT_GT(std::filesystem::file_size(mjb_path), 0);
+
+  mjModel* model2 = mj_loadModel(mjb_path.c_str(), nullptr);
+  EXPECT_THAT(model2, NotNull());
+  mj_deleteModel(model2);
+
+  // TXT Encoding (spec can be null)
+  EXPECT_GT(mj_encode(nullptr, model, txt_path.c_str(), nullptr, nullptr, error,
+                      sizeof(error)),
+            0);
+  EXPECT_TRUE(std::filesystem::exists(txt_path));
+  EXPECT_GT(std::filesystem::file_size(txt_path), 0);
+
+  // Clean up
+  std::filesystem::remove(mjb_path);
+  std::filesystem::remove(txt_path);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+}
 }  // namespace
 }  // namespace mujoco