Allow height-field elevation data to be specified in XML.

PiperOrigin-RevId: 602374498
Change-Id: I2d3364d68876a35799b24aca1051e2f5aebdda75

src/user/user_model.cc

@@ -3376,11 +3376,28 @@ bool mjCModel::CopyBack(const mjModel* m) {
   mjCMesh* pm;
   for (int i=0; i<nmesh; i++) {
     pm = meshes[i];
-
     copyvec(pm->GetOffsetPosPtr(), m->mesh_pos+3*i, 3);
     copyvec(pm->GetOffsetQuatPtr(), m->mesh_quat+4*i, 4);
   }
 
+  // heightfield
+  mjCHField* phf;
+  for (int i=0; i<nhfield; i++) {
+    phf = hfields[i];
+    int size = phf->userdata().size();
+    if (size) {
+      int nrow = m->hfield_nrow[i];
+      int ncol = m->hfield_ncol[i];
+      float* userdata = phf->userdata().data();
+      float* modeldata = m->hfield_data + m->hfield_adr[i];
+      // copy back in reverse row order
+      for (int j=0; j<nrow; j++) {
+        int flip = nrow-1-j;
+        copyvec(userdata + flip*ncol, modeldata+j*ncol, ncol);
+      }
+    }
+  }
+
   // sites
   for (int i=0; i<nsite; i++) {
     copyvec(sites[i]->size, m->site_size + 3 * i, 3);

src/user/user_objects.cc

@@ -20,9 +20,11 @@
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <optional>
 #include <sstream>
 #include <string>
 #include <string_view>
+#include <utility>
 #include <vector>
 
 #include "lodepng.h"
@@ -2233,6 +2235,7 @@ mjCHField::mjCHField(mjCModel* _model) {
   nrow = 0;
   ncol = 0;
   data = 0;
+  userdata_.clear();
 }
 
 
@@ -2335,6 +2338,14 @@ void mjCHField::LoadPNG(mjResource* resource) {
 
 
 
+// user data setter
+void mjCHField::set_userdata(std::optional<std::vector<float>>&& userdata) {
+  if (userdata.has_value()) {
+    userdata_ = std::move(userdata.value());
+  }
+}
+
+
 // compiler
 void mjCHField::Compile(const mjVFS* vfs) {
   // check size parameters

src/user/user_objects.h

@@ -843,9 +843,16 @@ class mjCHField : public mjCBase {
   int ncol;                       // number of columns
   float* data;                    // elevation data, row-major format
 
+  // getter for user data
+  std::vector<float>& userdata() { return userdata_; }
+
+  // setter for user data
+  void set_userdata(std::optional<std::vector<float>>&& userdata);
+
  private:
   mjCHField(mjCModel* model);             // constructor
   ~mjCHField();                           // destructor
+  std::vector<float> userdata_;           // user-provided elevation data
   void Compile(const mjVFS* vfs);         // compiler
 
   void LoadCustom(mjResource* resource);  // load from custom format

src/xml/xml_native_reader.cc

@@ -217,7 +217,7 @@ static const char* MJCF[nMJCF][mjXATTRNUM] = {
             "fileright", "fileleft", "fileup", "filedown", "filefront", "fileback",
             "builtin", "rgb1", "rgb2", "mark", "markrgb", "random", "width", "height",
             "hflip", "vflip"},
-        {"hfield", "*", "6", "name", "content_type", "file", "nrow", "ncol", "size"},
+        {"hfield", "*", "7", "name", "content_type", "file", "nrow", "ncol", "size", "elevation"},
         {"mesh", "*", "12", "name", "class", "content_type", "file", "vertex", "normal",
             "texcoord", "face", "refpos", "refquat", "scale", "smoothnormal"},
         {"<"},
@@ -2931,10 +2931,33 @@ void mjXReader::Asset(XMLElement* section) {
       ReadAttrInt(elem, "ncol", &phf->ncol);
       ReadAttr(elem, "size", 4, phf->size, text, true);
 
-      // allocate buffer for dynamic hfield
+      // allocate buffer for dynamic hfield, copy user data if given
       if (phf->file.empty() && phf->nrow>0 && phf->ncol>0) {
-        phf->data = (float*) mju_malloc(phf->nrow*phf->ncol*sizeof(float));
-        memset(phf->data, 0, phf->nrow*phf->ncol*sizeof(float));
+        int nrow = phf->nrow;
+        int ncol = phf->ncol;
+        phf->data = (float*) mju_malloc(nrow*ncol*sizeof(float));
+
+        // read user data
+        phf->set_userdata(ReadAttrVec<float>(elem, "elevation"));
+
+        // user data given, copy into data
+        if (!phf->userdata().empty()) {
+          if (phf->userdata().size() != nrow*ncol) {
+            throw mjXError(elem, "elevation data length must match nrow*ncol");
+          }
+
+          // copy in reverse row order, so XML string is top-to-bottom
+          const float* userdata = phf->userdata().data();
+          for (int i = 0; i < nrow; i++) {
+            int flip = nrow-1-i;
+            memcpy(phf->data + flip*ncol, userdata + i*ncol, ncol*sizeof(float));
+          }
+        }
+
+        // user data not given, set to 0
+        else {
+          memset(phf->data, 0, phf->nrow*phf->ncol*sizeof(float));
+        }
       }
     }
 

src/xml/xml_native_writer.cc

@@ -1449,6 +1449,11 @@ void mjXWriter::Asset(XMLElement* root) {
     } else {
       WriteAttrInt(elem, "nrow", phf->nrow);
       WriteAttrInt(elem, "ncol", phf->ncol);
+      if (!phf->userdata().empty()) {
+        string text;
+        Vector2String(text, phf->userdata(), phf->ncol);
+        WriteAttrTxt(elem, "elevation", text);
+      }
     }
   }
 }

src/xml/xml_util.cc

@@ -875,11 +875,13 @@ void mjXUtil::Vector2String(string& txt, const vector<double>& vec) {
 
 
 // write vector<float> to string
-void mjXUtil::Vector2String(string& txt, const vector<float>& vec) {
+void mjXUtil::Vector2String(string& txt, const vector<float>& vec, int ncol) {
   stringstream strm;
 
   for (size_t i=0; i<vec.size(); i++) {
-    if (i>0) {
+    if (ncol && (i % ncol) == 0) {
+      strm << "\n            ";
+    } else if (i>0) {
       strm << " ";
     }
     strm << vec[i];

src/xml/xml_util.h

@@ -179,7 +179,7 @@ class mjXUtil {
   static void Vector2String(std::string& txt, const std::vector<double>& vec);
 
   // write vector<float> to string
-  static void Vector2String(std::string& txt, const std::vector<float>& vec);
+  static void Vector2String(std::string& txt, const std::vector<float>& vec, int ncol = 0);
 
   // write vector<int> to string
   static void Vector2String(std::string& txt, const std::vector<int>& vec);