From d3bc0544d39c0933067e493f4c1c7ebfd15ae9f3 Mon Sep 17 00:00:00 2001
From: Google DeepMind <noreply@deepmind.com>
Date: Thu, 24 Apr 2025 08:58:05 -0700
Subject: [PATCH] Add basic support for geom rgba.

Adds primvar:displayColor/primvar:displayOpacity to the geom prim if its rgba attribute differs from the default (0.5, 0.5, 0.5, 1.0).  This sets the color and opacity directly on the prim with no need for an external material, making it consistent with how rgba is used in mujoco.

PiperOrigin-RevId: 751013128
Change-Id: I059943ec5b266e62c3696d9a0ac4f2033db9e07f
---
 .../usd/plugins/mjcf/mujoco_to_usd.cc         | 23 +++++++++
 test/experimental/usd/plugins/mjcf/fixture.h  |  7 +++
 .../usd/plugins/mjcf/mjcf_file_format_test.cc | 47 +++++++++++++++++++
 3 files changed, 77 insertions(+)

diff --git a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
index 25757654..e443149c 100644
--- a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
+++ b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
@@ -897,6 +897,29 @@ class ModelWriter {
       }
     }
 
+    // If geom 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 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);
     }
diff --git a/test/experimental/usd/plugins/mjcf/fixture.h b/test/experimental/usd/plugins/mjcf/fixture.h
index 9ab5ffe1..d7b8cda1 100644
--- a/test/experimental/usd/plugins/mjcf/fixture.h
+++ b/test/experimental/usd/plugins/mjcf/fixture.h
@@ -60,6 +60,13 @@
         .Get(&prim_purpose);                         \
     EXPECT_EQ(prim_purpose, purpose);                \
   }
+
+#define EXPECT_ATTRIBUTE_HAS_VALUE(stage, path) \
+  EXPECT_TRUE((stage)->GetAttributeAtPath(SdfPath(path)).HasValue());
+
+#define EXPECT_ATTRIBUTE_HAS_NO_VALUE(stage, path) \
+  EXPECT_FALSE((stage)->GetAttributeAtPath(SdfPath(path)).HasValue());
+
 namespace mujoco {
 
 pxr::SdfLayerRefPtr LoadLayer(
diff --git a/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc b/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc
index f5caeb19..f144da6e 100644
--- a/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc
+++ b/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc
@@ -140,6 +140,53 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestMaterials) {
       pxr::SdfAssetPath("textures/cube.png"));
 }
 
+TEST_F(MjcfSdfFileFormatPluginTest, TestGeomRgba) {
+  static constexpr char kXml[] = R"(
+    <mujoco model="test">
+      <worldbody>
+        <geom type="sphere" name="sphere_red" size="1" rgba="1 0 0 1"/>
+        <geom type="sphere" name="sphere_default" size="1"/>
+        <geom type="sphere" name="sphere_also_default" size="1" rgba="0.5 0.5 0.5 1"/>
+        <geom type="sphere" name="sphere_almost_default" size="1" rgba="0.5 0.5 0.5 0.9"/>
+      </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");
+
+  // There's no mechanism in Mujoco to specify whether an attribute was set
+  // explicitly or not. We do the same as Mujoco does, which is to compare with
+  // the default value.
+  // Which explains why not setting rgba is the same as setting it to the
+  // default value of (0.5, 0.5, 0.5, 1).
+  EXPECT_PRIM_VALID(stage, "/test/sphere_default");
+  EXPECT_ATTRIBUTE_HAS_NO_VALUE(stage,
+                                "/test/sphere_default.primvars:displayColor");
+  EXPECT_ATTRIBUTE_HAS_NO_VALUE(stage,
+                                "/test/sphere_default.primvars:displayOpacity");
+
+  EXPECT_PRIM_VALID(stage, "/test/sphere_also_default");
+  EXPECT_ATTRIBUTE_HAS_NO_VALUE(
+      stage, "/test/sphere_also_default.primvars:displayColor");
+  EXPECT_ATTRIBUTE_HAS_NO_VALUE(
+      stage, "/test/sphere_also_default.primvars:displayOpacity");
+
+  EXPECT_PRIM_VALID(stage, "/test/sphere_almost_default");
+  ExpectAttributeEqual(stage,
+                       "/test/sphere_almost_default.primvars:displayColor",
+                       pxr::VtArray<pxr::GfVec3f>{{0.5, 0.5, 0.5}});
+  ExpectAttributeEqual(stage,
+                       "/test/sphere_almost_default.primvars:displayOpacity",
+                       pxr::VtArray<float>{0.9});
+}
+
 TEST_F(MjcfSdfFileFormatPluginTest, TestFaceVaryingMeshSourcesSimpleMjcfMesh) {
   static constexpr char kXml[] = R"(
     <mujoco model="mesh test">