Add support for ORM maps in mujoco file format plugin

The ORM texture is a packed-channel map, with ambient occlusion in Red, roughness in Green, and metallic in Blue channel.

PiperOrigin-RevId: 772367369
Change-Id: I88a4dee2477daa1e3c0b00197bfbe2b276ba9937

src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc

@@ -92,7 +92,12 @@ TF_DEFINE_PRIVATE_TOKENS(kTokens,
                          ((inputsWrapT, "inputs:wrapT"))
                          ((inputsDiffuseColor, "inputs:diffuseColor"))
                          ((outputsRgb, "outputs:rgb"))
+                         ((outputsR, "outputs:r"))
+                         ((outputsG, "outputs:g"))
+                         ((outputsB, "outputs:b"))
                          ((inputsMetallic, "inputs:metallic"))
+                         ((inputsOcclusion, "inputs:occlusion"))
+                         ((inputsRoughness, "inputs:roughness"))
                          (repeat)
                          ((sourceMesh, pxr::UsdGeomTokens->Mesh))
                          ((inputsNormal, "inputs:normal"))
@@ -655,10 +660,10 @@ class ModelWriter {
     return uvmap_st_output_attr;
   }
 
-  pxr::SdfPath AddTextureShader(const pxr::SdfPath &material_path,
+  std::vector<pxr::SdfPath> AddTextureShader(
-                                const char *texture_file,
+      const pxr::SdfPath &material_path, const char *texture_file,
-                                const pxr::TfToken &name,
+      const pxr::TfToken &name, const pxr::SdfPath &uvmap_st_output_attr,
-                                const pxr::SdfPath &uvmap_st_output_attr) {
+      const std::vector<pxr::TfToken> &output_channels) {
     pxr::SdfPath texture_shader_path =
         CreatePrimSpec(data_, material_path, name, pxr::UsdShadeTokens->Shader);
     pxr::SdfPath texture_info_id_attr = CreateAttributeSpec(
@@ -688,11 +693,19 @@ class ModelWriter {
                             pxr::SdfValueTypeNames->Token);
     SetAttributeDefault(data_, texture_wrap_t_attr, kTokens->repeat);
 
-    pxr::SdfPath texture_rgb_output_attr =
+    std::vector<pxr::SdfPath> texture_output_attrs;
-        CreateAttributeSpec(data_, texture_shader_path, kTokens->outputsRgb,
+    for (const auto &output_channel : output_channels) {
-                            pxr::SdfValueTypeNames->Float3);
+      pxr::SdfValueTypeName value_type;
-
+      if (output_channel == kTokens->outputsRgb) {
-    return texture_rgb_output_attr;
+        value_type = pxr::SdfValueTypeNames->Float3;
+      } else {
+        // Assume the other specified channels are outputR, outputG, outputB.
+        value_type = pxr::SdfValueTypeNames->Float;
+      }
+      texture_output_attrs.push_back(CreateAttributeSpec(
+          data_, texture_shader_path, output_channel, value_type));
+    }
+    return texture_output_attrs;
   }
 
   void WriteMaterial(mjsMaterial *material, const pxr::SdfPath &parent_path) {
@@ -722,9 +735,44 @@ class ModelWriter {
 
     const pxr::SdfPath &uvmap_st_output_attr =
         AddUVTextureShader(material_path, pxr::TfToken("uvmap"));
-
-    // Find the normal texture if specified.
     const mjStringVec &textures = *(material->textures);
+
+    // Set the values of metallic and roughness. These can come from an ORM
+    // texture or as a value defined in mjsMaterial_. Occlusion is only present
+    // in the ORM texture.
+    pxr::SdfPath metallic_attr = CreateAttributeSpec(
+        data_, preview_surface_shader_path, kTokens->inputsMetallic,
+        pxr::SdfValueTypeNames->Float);
+    pxr::SdfPath roughness_attr = CreateAttributeSpec(
+        data_, preview_surface_shader_path, kTokens->inputsRoughness,
+        pxr::SdfValueTypeNames->Float);
+    // Find the ORM (occlusion, roughness, metallic) packed-channel texture if
+    // specified.
+    if (mjTEXROLE_ORM < textures.size()) {
+      std::string orm_texture_name = textures[mjTEXROLE_ORM];
+      mjsTexture *orm_texture = mjs_asTexture(
+          mjs_findElement(spec_, mjOBJ_TEXTURE, orm_texture_name.c_str()));
+      if (orm_texture) {
+        // Create the ORM shader and connect its output to the preview
+        // surface ORM attrs.
+        const std::vector<pxr::SdfPath> orm_output_attrs = AddTextureShader(
+            material_path, orm_texture->file->c_str(),
+            pxr::TfToken("orm_packed"), uvmap_st_output_attr,
+            {kTokens->outputsR, kTokens->outputsG, kTokens->outputsB});
+        if (orm_output_attrs.size() == 3) {
+          pxr::SdfPath occlusion_attr = CreateAttributeSpec(
+              data_, preview_surface_shader_path, kTokens->inputsOcclusion,
+              pxr::SdfValueTypeNames->Float);
+          AddAttributeConnection(data_, occlusion_attr, orm_output_attrs[0]);
+          AddAttributeConnection(data_, roughness_attr, orm_output_attrs[1]);
+          AddAttributeConnection(data_, metallic_attr, orm_output_attrs[2]);
+        }
+      } else {
+        SetAttributeDefault(data_, metallic_attr, material->metallic);
+        SetAttributeDefault(data_, roughness_attr, material->roughness);
+      }
+    }
+    // Find the normal texture if specified.
     if (mjTEXROLE_NORMAL < textures.size()) {
       std::string normal_texture_name = textures[mjTEXROLE_NORMAL];
       mjsTexture *normal_texture = mjs_asTexture(
@@ -735,10 +783,14 @@ class ModelWriter {
             pxr::SdfValueTypeNames->Normal3f);
         // Create the normal map shader and connect its output to the preview
         // surface normal attr.
-        pxr::SdfPath normal_map_output_attr =
+        const std::vector<pxr::SdfPath> normal_map_output_attrs =
             AddTextureShader(material_path, normal_texture->file->c_str(),
-                             pxr::TfToken("normal"), uvmap_st_output_attr);
+                             pxr::TfToken("normal"), uvmap_st_output_attr,
-        AddAttributeConnection(data_, normal_attr, normal_map_output_attr);
+                             {kTokens->outputsRgb});
+        if (normal_map_output_attrs.size() == 1) {
+          AddAttributeConnection(data_, normal_attr,
+                                 normal_map_output_attrs[0]);
+        }
       }
     }
 
@@ -753,11 +805,14 @@ class ModelWriter {
     if (main_texture) {
       // Create the texture shader and connect it to the diffuse color
       // attribute.
-      pxr::SdfPath texture_diffuse_output_attr =
+      const std::vector<pxr::SdfPath> texture_diffuse_output_attrs =
           AddTextureShader(material_path, main_texture->file->c_str(),
-                           pxr::TfToken("diffuse"), uvmap_st_output_attr);
+                           pxr::TfToken("diffuse"), uvmap_st_output_attr,
-      AddAttributeConnection(data_, diffuse_color_attr,
+                           {kTokens->outputsRgb});
-                             texture_diffuse_output_attr);
+      if (texture_diffuse_output_attrs.size() == 1) {
+        AddAttributeConnection(data_, diffuse_color_attr,
+                               texture_diffuse_output_attrs[0]);
+      }
     } else {
       // If no texture is specified, use the rgba diffuse color.
       SetAttributeDefault(data_, diffuse_color_attr,
@@ -765,11 +820,6 @@ class ModelWriter {
                                        material->rgba[2]));
     }
 
-    pxr::SdfPath metallic_attr = CreateAttributeSpec(
-        data_, preview_surface_shader_path, kTokens->inputsMetallic,
-        pxr::SdfValueTypeNames->Float);
-    SetAttributeDefault(data_, metallic_attr, material->metallic);
-
     pxr::SdfPath material_surface_output_attr = CreateAttributeSpec(
         data_, material_path, pxr::UsdShadeTokens->outputsSurface,
         pxr::SdfValueTypeNames->Token);

test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc

@@ -161,6 +161,7 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestMaterials) {
   EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered");
   EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered/uvmap");
   EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered/diffuse");
+
   EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered/normal");
   ExpectAttributeHasConnection(
     stage,
@@ -170,6 +171,33 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestMaterials) {
   ExpectAttributeEqual(
     stage, "/mesh_test/Materials/material_layered/normal.inputs:file",
     pxr::SdfAssetPath("textures/normal.png"));
+
+  EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_layered/orm_packed");
+  ExpectAttributeHasConnection(
+    stage,
+    "/mesh_test/Materials/material_layered/"
+    "PreviewSurface.inputs:occlusion",
+    "/mesh_test/Materials/material_layered/orm_packed.outputs:r");
+  ExpectAttributeHasConnection(
+    stage,
+    "/mesh_test/Materials/material_layered/"
+    "PreviewSurface.inputs:roughness",
+    "/mesh_test/Materials/material_layered/orm_packed.outputs:g");
+  ExpectAttributeHasConnection(
+    stage,
+    "/mesh_test/Materials/material_layered/"
+    "PreviewSurface.inputs:metallic",
+    "/mesh_test/Materials/material_layered/orm_packed.outputs:b");
+  ExpectAttributeEqual(
+    stage, "/mesh_test/Materials/material_layered/orm_packed.inputs:file",
+    pxr::SdfAssetPath("textures/orm.png"));
+
+  EXPECT_PRIM_VALID(stage, "/mesh_test/Materials/material_metallic");
+  EXPECT_PRIM_VALID(stage,
+                    "/mesh_test/Materials/material_metallic/PreviewSurface");
+  ExpectAttributeEqual(stage,
+    "/mesh_test/Materials/material_metallic/PreviewSurface.inputs:metallic",
+    0.6f);
 }
 
 TEST_F(MjcfSdfFileFormatPluginTest, TestGeomRgba) {

test/experimental/usd/plugins/mjcf/testdata/materials.xml

@@ -2,11 +2,14 @@
   <asset>
     <texture name="diffuse_tex" file="textures/cube.png" type="2d" />
     <texture name="normal_tex" file="textures/normal.png" type="2d" />
+    <texture name="orm_tex" file="textures/orm.png" type="2d" />
     <material name="material_red" rgba="0.8 0 0 1" />
     <material name="material_texture" texture="diffuse_tex" />
+    <material name="material_metallic" rgba="0.8 0 0 1" metallic="0.6"/>
     <material name="material_layered">
       <layer texture="diffuse_tex" role="rgb"/>
       <layer texture="normal_tex" role="normal"/>
+      <layer texture="orm_tex" role="orm"/>
     </material>
     2d texture="diffuse" normal="normal" />
     <mesh name="tetrahedron" vertex="0 0 0  1 0 0  0 1 0  0 0 1"/>