Add mj_copyBack for copying real-valued arrays from mjModel back to mjSpec.

Also add `SaveAndReadXML` function to test fixture using `mjSpec` as input.

PiperOrigin-RevId: 765393821
Change-Id: Ic257addd2fc89678fc11b226c168c077626cc51e

doc/APIreference/functions.rst

@@ -51,6 +51,15 @@ Compile :ref:`mjSpec` to :ref:`mjModel`. A spec can be edited and compiled multi
 :ref:`mjModel` instance that takes the edits into account.
 If compilation fails, :ref:`mj_compile` returns ``NULL``; the error can be read with :ref:`mjs_getError`.
 
+.. _mj_copyBack:
+
+`mj_copyBack <#mj_copyBack>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mj_copyBack
+
+Copy real-valued arrays from model to spec, returns 1 on success.
+
 .. _mj_recompile:
 
 `mj_recompile <#mj_recompile>`__

doc/changelog.rst

@@ -17,6 +17,7 @@ General
 - Added new sub-component :ref:`mj_makeM` which combines the :ref:`mj_crb` call with additional logic to support the
   introduction in 3.3.1 of :ref:`tendon armature<tendon-spatial-armature>`. In addition to the traditional
   ``mjData.qM``, :ref:`mj_makeM` also computes ``mjData.M``, a CSR representation of the same matrix.
+- Added a new function :ref:`mj_copyBack` to copy real-valued arrays in an mjModel to a compatible mjSpec.
 
 Simulate
 ^^^^^^^^

doc/includes/references.h

@@ -2992,6 +2992,7 @@ mjModel* mj_loadXML(const char* filename, const mjVFS* vfs, char* error, int err
 mjSpec* mj_parseXML(const char* filename, const mjVFS* vfs, char* error, int error_sz);
 mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, int error_sz);
 mjModel* mj_compile(mjSpec* s, const mjVFS* vfs);
+int mj_copyBack(mjSpec* s, const mjModel* m);
 int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d);
 int mj_saveLastXML(const char* filename, const mjModel* m, char* error, int error_sz);
 void mj_freeLastXML(void);

include/mujoco/mujoco.h

@@ -107,6 +107,9 @@ MJAPI mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error,
 // Compile spec to model.
 MJAPI mjModel* mj_compile(mjSpec* s, const mjVFS* vfs);
 
+// Copy real-valued arrays from model to spec, returns 1 on success.
+MJAPI int mj_copyBack(mjSpec* s, const mjModel* m);
+
 // Recompile spec to model, preserving the state, return 0 on success.
 MJAPI int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d);
 

python/mujoco/introspect/functions.py

@@ -248,6 +248,26 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Compile spec to model.',
      )),
+    ('mj_copyBack',
+     FunctionDecl(
+         name='mj_copyBack',
+         return_type=ValueType(name='int'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='s',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjSpec'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='m',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjModel', is_const=True),
+                 ),
+             ),
+         ),
+         doc='Copy real-valued arrays from model to spec, returns 1 on success.',  # pylint: disable=line-too-long
+     )),
     ('mj_recompile',
      FunctionDecl(
          name='mj_recompile',

src/user/user_api.cc

@@ -375,6 +375,14 @@ int mjs_setDeepCopy(mjSpec* s, int deepcopy) {
 
 
 
+// copy real-valued arrays from model to spec, returns 1 on success
+int mj_copyBack(mjSpec* s, const mjModel* m) {
+  mjCModel* model = static_cast<mjCModel*>(s->element);
+  return model->CopyBack(m);
+}
+
+
+
 // delete object, return 0 on success
 int mjs_delete(mjsElement* element) {
   mjCModel* model;

src/user/user_api.h

@@ -66,6 +66,9 @@ MJAPI int mjs_activatePlugin(mjSpec* s, const char* name);
 // Turn deep copy on or off attach. Returns 0 on success.
 MJAPI int mjs_setDeepCopy(mjSpec* s, int deepcopy);
 
+// Copy real-valued arrays from model to spec, returns 1 on success.
+MJAPI int mj_copyBack(mjSpec* s, const mjModel* m);
+
 
 //---------------------------------- Attachment ----------------------------------------------------
 

src/user/user_model.cc

@@ -4759,6 +4759,11 @@ bool mjCModel::CopyBack(const mjModel* m) {
     return false;
   }
 
+  if (spec.element->signature != m->signature) {
+    errInfo = mjCError(0, "incompatible signatures in CopyBack");
+    return false;
+  }
+
   // option and visual
   option = m->opt;
   visual = m->vis;

src/xml/xml.cc

@@ -395,7 +395,7 @@ mjSpec* ParseSpecFromString(std::string_view xml, const mjVFS* vfs, char* error,
 }
 
 // Main writer function - calls mjXWrite
-std::string WriteXML(const mjModel* m, const mjSpec* spec, char* error, int nerror) {
+std::string WriteXML(const mjModel* m, mjSpec* spec, char* error, int nerror) {
   LocaleOverride locale_override;
 
   // check for empty model

src/xml/xml.h

@@ -30,6 +30,6 @@ mjSpec* ParseSpecFromString(std::string_view xml, const mjVFS* vfs = nullptr,
                             char* error = nullptr, int nerror = 0);
 
 // Main writer function
-std::string WriteXML(const mjModel* m, const mjSpec* spec, char* error, int nerror);
+std::string WriteXML(const mjModel* m, mjSpec* spec, char* error, int nerror);
 
 #endif  // MUJOCO_SRC_XML_XML_H_

src/xml/xml_api.cc

@@ -230,7 +230,8 @@ mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, int er
 
 // save spec to XML file, return 0 on success, -1 otherwise
 int mj_saveXML(const mjSpec* s, const char* filename, char* error, int error_sz) {
-  std::string result = WriteXML(NULL, s, error, error_sz);
+  // cast to mjSpec since WriteXML can in principle perform mj_copyBack (not here)
+  std::string result = WriteXML(NULL, (mjSpec*)s, error, error_sz);
   if (result.empty()) {
     return -1;
   }
@@ -247,7 +248,7 @@ int mj_saveXML(const mjSpec* s, const char* filename, char* error, int error_sz)
 // save spec to XML string, return 0 on success, -1 on failure
 // if length of the output buffer is too small, returns the required size
 int mj_saveXMLString(const mjSpec* s, char* xml, int xml_sz, char* error, int error_sz) {
-  std::string result = WriteXML(NULL, s, error, error_sz);
+  std::string result = WriteXML(NULL, (mjSpec*)s, error, error_sz);
   if (result.empty()) {
     return -1;
   } else if (result.size() >= xml_sz) {

src/xml/xml_base.cc

@@ -44,8 +44,8 @@ mjXBase::mjXBase() {
 
 
 // set model field
-void mjXBase::SetModel(const mjSpec* _model, const mjModel* m) {
+void mjXBase::SetModel(mjSpec* _model, const mjModel* m) {
-  spec = (mjSpec*)_model;
+  spec = _model;
 }
 
 

src/xml/xml_base.h

@@ -92,7 +92,7 @@ class mjXBase : public mjXUtil {
   };
 
   // set the model allocated externally
-  virtual void SetModel(const mjSpec*, const mjModel* = nullptr);
+  virtual void SetModel(mjSpec*, const mjModel* = nullptr);
 
   // read alternative orientation specification
   static int ReadAlternative(tinyxml2::XMLElement* elem, mjsOrientation& alt);

src/xml/xml_native_writer.cc

@@ -885,12 +885,12 @@ mjXWriter::mjXWriter(void) {
 
 
 // cast model
-void mjXWriter::SetModel(const mjSpec* _spec, const mjModel* m) {
+void mjXWriter::SetModel(mjSpec* _spec, const mjModel* m) {
   if (_spec) {
     model = static_cast<mjCModel*>(_spec->element);
   }
   if (m) {
-    model->CopyBack(m);
+    mj_copyBack(&model->spec, m);
   }
 }
 

src/xml/xml_native_writer.h

@@ -29,7 +29,7 @@ class mjXWriter : public mjXBase {
  public:
   mjXWriter();                                        // constructor
   virtual ~mjXWriter() = default;                     // destructor
-  void SetModel(const mjSpec* _spec, const mjModel* m = nullptr);
+  void SetModel(mjSpec* _spec, const mjModel* m = nullptr);
 
   // write XML document to string
   std::string Write(char *error, std::size_t error_sz);

test/fixture.cc

@@ -130,9 +130,7 @@ std::string GetFileContents(const char* path) {
   return sstream.str();
 }
 
-std::string SaveAndReadXml(const mjModel* model) {
+std::string SaveAndReadXmlImpl(const mjModel* model, const mjSpec* spec) {
-  EXPECT_THAT(model, testing::NotNull());
-
   constexpr int kMaxPathLen = 1024;
   std::string path_template =
       std::filesystem::temp_directory_path().append("tmp.XXXXXX").string();
@@ -148,7 +146,11 @@ std::string SaveAndReadXml(const mjModel* model) {
   EXPECT_NE(_mktemp_s(filepath), EINVAL);
 #endif
 
-  mj_saveLastXML(filepath, model, nullptr, 0);
+  if (spec) {
+    mj_saveXML(spec, filepath, nullptr, 0);
+  } else if (model) {
+    mj_saveLastXML(filepath, model, nullptr, 0);
+  }
   std::string contents = GetFileContents(filepath);
 
 #if defined(_POSIX_VERSION) && _POSIX_VERSION >= 200112L
@@ -159,6 +161,16 @@ std::string SaveAndReadXml(const mjModel* model) {
   return contents;
 }
 
+std::string SaveAndReadXml(const mjModel* model) {
+  EXPECT_THAT(model, testing::NotNull());
+  return SaveAndReadXmlImpl(model, nullptr);
+}
+
+std::string SaveAndReadXml(const mjSpec* spec) {
+  EXPECT_THAT(spec, testing::NotNull());
+  return SaveAndReadXmlImpl(nullptr, spec);
+}
+
 std::vector<mjtNum> GetCtrlNoise(const mjModel* m, int nsteps,
                                  mjtNum ctrlnoise) {
   std::vector<mjtNum> ctrl;

test/fixture.h

@@ -102,6 +102,9 @@ mjModel* LoadModelFromPath(const char* model_path);
 // Returns a string loaded from first saving the model given an input.
 std::string SaveAndReadXml(const mjModel* model);
 
+// Returns a string loaded from first saving the spec given an input.
+std::string SaveAndReadXml(const mjSpec* spec);
+
 // Adds control noise.
 std::vector<mjtNum> GetCtrlNoise(const mjModel* m, int nsteps,
                                  mjtNum ctrlnoise = 0.01);

test/xml/xml_native_writer_test.cc

@@ -1498,6 +1498,65 @@ TEST_F(DecompilerTest, SavesStatistics) {
   mj_deleteModel(model);
 }
 
+TEST_F(DecompilerTest, SaveAndReadXml) {
+  static constexpr char xml1[] = R"(
+  <mujoco>
+    <worldbody>
+      <geom size="1"/>
+      <geom size="2"/>
+    </worldbody>
+  </mujoco>
+  )";
+  static constexpr char xml2[] = R"(
+  <mujoco>
+    <worldbody>
+      <geom size="1"/>
+      <geom size="2"/>
+      <geom size="3"/>
+    </worldbody>
+  </mujoco>
+  )";
+  std::array<char, 1024> error;
+  mjModel* m1 = LoadModelFromString(xml1, error.data(), error.size());
+  ASSERT_THAT(m1, NotNull()) << error.data();
+  m1->geom_size[0] = 10;
+  m1->geom_size[3] = 20;
+  std::string saved_xml = SaveAndReadXml(m1);
+  EXPECT_THAT(saved_xml, HasSubstr("geom size=\"10\""));
+  EXPECT_THAT(saved_xml, HasSubstr("geom size=\"20\""));
+
+  // parse the mjSpec, save it and read it back
+  mjSpec* spec = mj_parseXMLString(xml2, nullptr, error.data(), error.size());
+  EXPECT_THAT(spec, NotNull()) << error.data();
+  mjModel* m2 = mj_compile(spec, nullptr);
+  std::string saved_xml1 = SaveAndReadXml(spec);
+  EXPECT_THAT(saved_xml1, HasSubstr("geom size=\"1\""));
+  EXPECT_THAT(saved_xml1, HasSubstr("geom size=\"2\""));
+  EXPECT_THAT(saved_xml1, HasSubstr("geom size=\"3\""));
+
+  // modify the mjModel, save it and read it back
+  m2->geom_size[0] = .1;
+  m2->geom_size[3] = .2;
+  m2->geom_size[6] = .3;
+  EXPECT_EQ(mj_copyBack(spec, m1), 0);
+  EXPECT_THAT(mjs_getError(spec), HasSubstr("CopyBack"));
+  EXPECT_EQ(mj_copyBack(spec, m2), 1);
+  std::string saved_xml2 = SaveAndReadXml(spec);
+  EXPECT_THAT(saved_xml2, HasSubstr("geom size=\"0.1\""));
+  EXPECT_THAT(saved_xml2, HasSubstr("geom size=\"0.2\""));
+  EXPECT_THAT(saved_xml2, HasSubstr("geom size=\"0.3\""));
+
+  // check that using mjModel as argument writes in the wrong mjSpec
+  std::string saved_xml3 = SaveAndReadXml(m2);
+  EXPECT_THAT(saved_xml3, Not(HasSubstr("geom size=\"0.1\"")));
+  EXPECT_THAT(saved_xml3, Not(HasSubstr("geom size=\"0.2\"")));
+  EXPECT_THAT(saved_xml3, Not(HasSubstr("geom size=\"0.3\"")));
+
+  mj_deleteSpec(spec);
+  mj_deleteModel(m1);
+  mj_deleteModel(m2);
+}
+
 TEST_F(DecompilerTest, DoesntSaveInferredStatistics) {
   static constexpr char xml[] = R"(
   <mujoco>