Add displayColor authoring to site geometry when converting MJCF to USD.

PiperOrigin-RevId: 789311878
Change-Id: I70027cd5d8d02f4ee26043f3b9fb03311e38984b

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

@@ -317,6 +317,34 @@ class ModelWriter {
     SetAttributeDefault(data_, attr_path, value);
   }
 
+  template <typename T>
+  void WriteColorAndOpacityAttributes(const pxr::SdfPath &prim_path,
+                                      const T &element) {
+    // If rgba is not the default (0.5, 0.5, 0.5, 1), then set the
+    // displayColor attribute.
+    // No effort is made to properly handle the interaction between rgba
+    // and the material if both are specified.
+    if (element->rgba[0] != 0.5f || element->rgba[1] != 0.5f ||
+        element->rgba[2] != 0.5f || element->rgba[3] != 1.0f) {
+      // Set the displayColor attribute.
+      pxr::SdfPath display_color_attr = CreateAttributeSpec(
+          data_, prim_path, pxr::UsdGeomTokens->primvarsDisplayColor,
+          pxr::SdfValueTypeNames->Color3fArray);
+      SetAttributeDefault(
+          data_, display_color_attr,
+          pxr::VtArray<pxr::GfVec3f>{
+              {element->rgba[0], element->rgba[1], element->rgba[2]}});
+      // Set the displayOpacity attribute, only if the opacity is not 1.
+      if (element->rgba[3] != 1.0f) {
+        pxr::SdfPath display_opacity_attr = CreateAttributeSpec(
+            data_, prim_path, pxr::UsdGeomTokens->primvarsDisplayOpacity,
+            pxr::SdfValueTypeNames->FloatArray);
+        SetAttributeDefault(data_, display_opacity_attr,
+                            pxr::VtArray<float>{element->rgba[3]});
+      }
+    }
+  }
+
   void PrependToXformOpOrder(const pxr::SdfPath &prim_path,
                              const pxr::VtArray<pxr::TfToken> &order) {
     auto xform_op_order_path =
@@ -1573,6 +1601,8 @@ class ModelWriter {
     WriteUniformAttribute(site_path, pxr::SdfValueTypeNames->Int,
                           MjcPhysicsTokens->mjcGroup, site->group);
 
+    WriteColorAndOpacityAttributes(site_path, site);
+
     int site_id = mjs_getId(site->element);
     auto transform = MujocoPosQuatToTransform(&model_->site_pos[3 * site_id],
                                               &model_->site_quat[4 * site_id]);
@@ -1759,28 +1789,7 @@ class ModelWriter {
       }
     }
 
-    // If geom rgba is not the default (0.5, 0.5, 0.5, 1), then set the
+    WriteColorAndOpacityAttributes(geom_path, geom);
-    // displayColor attribute.
-    // No effort is made to properly handle the interaction between geom rgba
-    // and the material if both are specified.
-    if (geom->rgba[0] != 0.5f || geom->rgba[1] != 0.5f ||
-        geom->rgba[2] != 0.5f || geom->rgba[3] != 1.0f) {
-      // Set the displayColor attribute.
-      pxr::SdfPath display_color_attr = CreateAttributeSpec(
-          data_, geom_path, pxr::UsdGeomTokens->primvarsDisplayColor,
-          pxr::SdfValueTypeNames->Color3fArray);
-      SetAttributeDefault(data_, display_color_attr,
-                          pxr::VtArray<pxr::GfVec3f>{
-                              {geom->rgba[0], geom->rgba[1], geom->rgba[2]}});
-      // Set the displayOpacity attribute, only if the opacity is not 1.
-      if (geom->rgba[3] != 1.0f) {
-        pxr::SdfPath display_opacity_attr = CreateAttributeSpec(
-            data_, geom_path, pxr::UsdGeomTokens->primvarsDisplayOpacity,
-            pxr::SdfValueTypeNames->FloatArray);
-        SetAttributeDefault(data_, display_opacity_attr,
-                            pxr::VtArray<float>{geom->rgba[3]});
-      }
-    }
 
     if (body_id == kWorldIndex) {
       SetPrimKind(data_, geom_path, pxr::KindTokens->component);

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

@@ -350,6 +350,25 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestGeomRgba) {
                        pxr::VtArray<float>{0.9});
 }
 
+TEST_F(MjcfSdfFileFormatPluginTest, TestSiteRgba) {
+  static constexpr char kXml[] = R"(
+    <mujoco model="test">
+      <worldbody>
+        <site type="sphere" name="sphere_red" size="1" rgba="1 0 0 1"/>
+      </worldbody>
+    </mujoco>
+  )";
+
+  pxr::SdfLayerRefPtr layer = LoadLayer(kXml);
+  auto stage = pxr::UsdStage::Open(layer);
+
+  EXPECT_PRIM_VALID(stage, "/test/sphere_red");
+  ExpectAttributeEqual(stage, "/test/sphere_red.primvars:displayColor",
+                       pxr::VtArray<pxr::GfVec3f>{{1, 0, 0}});
+  EXPECT_ATTRIBUTE_HAS_NO_VALUE(stage,
+                                "/test/sphere_red.primvars:displayOpacity");
+}
+
 TEST_F(MjcfSdfFileFormatPluginTest, TestFaceVaryingMeshSourcesSimpleMjcfMesh) {
   static constexpr char kXml[] = R"(
     <mujoco model="mesh test">