add support for non-rgb textures in xml_native_writer

PiperOrigin-RevId: 736221430
Change-Id: I83c19e06f0f7be054a88c3e894807e7b90ee012a

src/xml/xml_base.h

@@ -60,6 +60,7 @@ extern const mjMap jac_map[];
 extern const mjMap solver_map[];
 extern const mjMap equality_map[];
 extern const mjMap texture_map[];
+extern const mjMap texrole_map[];
 extern const mjMap builtin_map[];
 extern const mjMap mark_map[];
 extern const mjMap dyn_map[];

src/xml/xml_native_writer.cc

@@ -323,6 +323,7 @@ void mjXWriter::OneMaterial(XMLElement* elem, const mjCMaterial* material, mjCDe
   }
 
   // defaults and regular
+  // check if we have non-rgb textures
   bool has_non_rgb = false;
   for (int i=1; i<mjNTEXROLE; i++) {
     if (!material->textures_[i].empty()) {
@@ -330,14 +331,23 @@ void mjXWriter::OneMaterial(XMLElement* elem, const mjCMaterial* material, mjCDe
         has_non_rgb = true;
       }
     }
-    if (material->textures_[i] != def->Material().textures_[i]) {
-      WriteAttrTxt(elem, "texture", material->get_texture(i));
+  }
+
+  // if we have non-rgb textures, write them as layers
+  if (has_non_rgb) {
+    for (int i=1; i<mjNTEXROLE; i++) {
+      if (!material->textures_[i].empty()) {
+        XMLElement * child_elem = InsertEnd(elem, "layer");
+        WriteAttrTxt(child_elem, "texture", material->textures_[i]);
+        WriteAttrTxt(child_elem, "role", FindValue(texrole_map, 9, i));
+      }
+    }
+  } else {
+    if (material->textures_[mjTEXROLE_RGB] != def->Material().textures_[mjTEXROLE_RGB]) {
+      WriteAttrTxt(elem, "texture", material->get_texture(mjTEXROLE_RGB));
     }
   }
-  if (has_non_rgb) {
-  //   // TODO elem = InsertEnd(section, "role");
-    throw mjXError(0, "no support for non-RGB textures.");
-  }
+
   WriteAttrKey(elem, "texuniform", bool_map, 2, material->texuniform, def->Material().texuniform);
   WriteAttr(elem, "texrepeat", 2, material->texrepeat, def->Material().texrepeat);
   WriteAttr(elem, "emission", 1, &material->emission, &def->Material().emission);

test/xml/testdata/hfield_png_nonrgb.xml

@@ -0,0 +1,17 @@
+<mujoco>
+  <asset>
+    <hfield name="hfield" file="hfield.png" size="0.5 0.5 0.05 0.1"/>
+
+    <texture name="hfield" type="2d" file="hfield.png"/>
+    <material name="hfield">
+      <layer texture="hfield" role="rgb"/>
+      <layer texture="hfield" role="orm"/>
+      <layer texture="hfield" role="normal"/>
+    </material>
+  </asset>
+
+  <worldbody>
+    <light pos="-.4 .6 1"/>
+    <geom type="hfield" hfield="hfield" pos="-.4 .6 .05" material="hfield"/>
+  </worldbody>
+</mujoco>

test/xml/xml_native_writer_test.cc

@@ -46,6 +46,9 @@ using ::testing::FloatEq;
 
 using XMLWriterTest = PluginTest;
 
+static const char* const kNonRgbTextureXMLPath =
+    "xml/testdata/hfield_png_nonrgb.xml";
+
 TEST_F(XMLWriterTest, EmptyModel) {
   static constexpr char xml[] = "<mujoco/>";
   mjModel* model = LoadModelFromString(xml);
@@ -1326,6 +1329,24 @@ TEST_F(XMLWriterLocaleTest, IgnoresLocale) {
   EXPECT_EQ(std::string(formatted), "3,9375");
 }
 
+TEST_F(XMLWriterTest, NonRGBTextures) {
+  const std::string xml_path = GetTestDataFilePath(kNonRgbTextureXMLPath);
+  std::array<char, 1024> error;
+  mjModel* model = mj_loadXML(xml_path.c_str(), 0, error.data(), error.size());
+  ASSERT_THAT(model, NotNull()) << error.data();
+
+  std::string saved_xml = SaveAndReadXml(model);
+  EXPECT_FALSE(saved_xml.empty());
+
+  // check that layers are written correctly
+  EXPECT_THAT(saved_xml, HasSubstr("<material name=\"hfield\">"));
+  EXPECT_THAT(saved_xml, HasSubstr("<layer texture=\"hfield\" role=\"rgb\"/>"));
+  EXPECT_THAT(saved_xml, HasSubstr("<layer texture=\"hfield\" role=\"orm\"/>"));
+  EXPECT_THAT(saved_xml,
+              HasSubstr("<layer texture=\"hfield\" role=\"normal\"/>"));
+
+  mj_deleteModel(model);
+}
 
 // ------------------- test loading and saving multiple files ------------------
 TEST_F(XMLWriterTest, WriteReadCompare) {