Add support for plane geometry.

PiperOrigin-RevId: 750994019
Change-Id: Ia3b085bed5b8e2fcdb00ffc69dd5551a7fa3c30d

Author: Google DeepMind

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

@@ -777,6 +777,43 @@ class ModelWriter {
     return WriteSphere(name, geom_size, body_path);
   }
 
+  pxr::SdfPath WritePlane(const pxr::TfToken &name, const mjtNum *size,
+                          const pxr::SdfPath &body_path) {
+    pxr::SdfPath plane_path =
+        CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Plane);
+
+    // MuJoCo uses half sizes.
+    // Note that UsdGeomPlane is infinite for simulation purposes but can have
+    // width/length for visualization, same as MuJoCo.
+    double width = size[0] * 2.0;
+    double length = size[1] * 2.0;
+
+    pxr::SdfPath width_attr_path =
+        CreateAttributeSpec(data_, plane_path, pxr::UsdGeomTokens->width,
+                            pxr::SdfValueTypeNames->Double);
+    SetAttributeDefault(data_, width_attr_path, width);
+
+    pxr::SdfPath length_attr_path =
+        CreateAttributeSpec(data_, plane_path, pxr::UsdGeomTokens->length,
+                            pxr::SdfValueTypeNames->Double);
+    SetAttributeDefault(data_, length_attr_path, length);
+
+    // MuJoCo plane is always a XY plane with +Z up.
+    // UsdGeomPlane is also a XY plane if axis is 'Z', which is default.
+    // So no need to set axis attribute explicitly.
+
+    return plane_path;
+  }
+
+  pxr::SdfPath WritePlaneGeom(const mjsGeom *geom,
+                              const pxr::SdfPath &body_path) {
+    auto name =
+        GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Plane, body_path);
+    int geom_idx = mjs_getId(geom->element);
+    mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
+    return WritePlane(name, geom_size, body_path);
+  }
+
   void WriteSite(mjsSite *site, const mjsBody *body) {
     const int body_id = mjs_getId(body->element);
     const auto &body_path = body_paths_[body_id];
@@ -804,6 +841,9 @@ class ModelWriter {
     pxr::SdfPath geom_path;
     int geom_id = mjs_getId(geom->element);
     switch (geom->type) {
+      case mjGEOM_PLANE:
+        geom_path = WritePlaneGeom(geom, body_path);
+        break;
       case mjGEOM_MESH:
         geom_path = WriteMeshGeom(geom, body_path);
         break;

test/experimental/usd/plugins/mjcf/fixture.h

@@ -70,10 +70,11 @@ template <typename T>
 void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, const char* path,
                           const T& value) {
   auto attr = stage->GetAttributeAtPath(pxr::SdfPath(path));
-  EXPECT_TRUE(attr.IsValid());
+  EXPECT_TRUE(attr.IsValid()) << "Attribute " << path << " is not valid";
   T attr_value;
   attr.Get(&attr_value);
-  EXPECT_EQ(attr_value, value);
+  EXPECT_EQ(attr_value, value) << "Attribute " << path << " has value "
+                               << attr_value << ". Expected: " << value;
 }
 
 // Specialization for SdfAssetPath, so that we can compare only the asset path

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

@@ -40,6 +40,7 @@
 #include <pxr/usd/usdGeom/cube.h>
 #include <pxr/usd/usdGeom/cylinder.h>
 #include <pxr/usd/usdGeom/mesh.h>
+#include <pxr/usd/usdGeom/plane.h>
 #include <pxr/usd/usdGeom/primvar.h>
 #include <pxr/usd/usdGeom/primvarsAPI.h>
 #include <pxr/usd/usdGeom/sphere.h>
@@ -389,6 +390,62 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestKindAuthoring) {
   EXPECT_PRIM_KIND(stage, "/test/root/tet", pxr::KindTokens->subcomponent);
 }
 
+TEST_F(MjcfSdfFileFormatPluginTest, TestGeomsPrims) {
+  static constexpr char kXml[] = R"(
+    <mujoco model="test">
+      <worldbody>
+        <geom type="plane" name="plane_geom" size="10 20 0.1"/>
+        <geom type="box" name="box_geom" size="10 20 30"/>
+        <geom type="sphere" name="sphere_geom" size="10 20 30"/>
+        <geom type="capsule" name="capsule_geom" size="10 20 30"/>
+        <geom type="cylinder" name="cylinder_geom" size="10 20 30"/>
+        <geom type="ellipsoid" name="ellipsoid_geom" size="10 20 30"/>
+      </worldbody>
+    </mujoco>
+  )";
+
+  pxr::SdfLayerRefPtr layer = LoadLayer(kXml);
+  auto stage = pxr::UsdStage::Open(layer);
+
+  // Note that all sizes are multiplied by 2 because Mujoco uses half sizes.
+
+  // Plane
+  EXPECT_PRIM_VALID(stage, "/test/plane_geom");
+  EXPECT_PRIM_IS_A(stage, "/test/plane_geom", pxr::UsdGeomPlane);
+  ExpectAttributeEqual(stage, "/test/plane_geom.width", 2 * 10.0);
+  ExpectAttributeEqual(stage, "/test/plane_geom.length", 2 * 20.0);
+  // Box
+  EXPECT_PRIM_VALID(stage, "/test/box_geom");
+  EXPECT_PRIM_IS_A(stage, "/test/box_geom", pxr::UsdGeomCube);
+  // Box is a special case, it uses a UsdGeomCube and scales it with
+  // xformOp:scale. The radius is always set to 2.
+  ExpectAttributeEqual(stage, "/test/box_geom.size", 2.0);
+  ExpectAttributeEqual(stage, "/test/box_geom.xformOp:scale",
+                       pxr::GfVec3f(10.0, 20.0, 30.0));
+  // Sphere
+  EXPECT_PRIM_VALID(stage, "/test/sphere_geom");
+  EXPECT_PRIM_IS_A(stage, "/test/sphere_geom", pxr::UsdGeomSphere);
+  ExpectAttributeEqual(stage, "/test/sphere_geom.radius", 2 * 10.0);
+  // Capsule
+  EXPECT_PRIM_VALID(stage, "/test/capsule_geom");
+  EXPECT_PRIM_IS_A(stage, "/test/capsule_geom", pxr::UsdGeomCapsule);
+  ExpectAttributeEqual(stage, "/test/capsule_geom.radius", 2 * 10.0);
+  ExpectAttributeEqual(stage, "/test/capsule_geom.height", 2 * 20.0);
+  // Cylinder
+  EXPECT_PRIM_VALID(stage, "/test/cylinder_geom");
+  EXPECT_PRIM_IS_A(stage, "/test/cylinder_geom", pxr::UsdGeomCylinder);
+  ExpectAttributeEqual(stage, "/test/cylinder_geom.radius", 2 * 10.0);
+  ExpectAttributeEqual(stage, "/test/cylinder_geom.height", 2 * 20.0);
+  // Ellipsoid
+  EXPECT_PRIM_VALID(stage, "/test/ellipsoid_geom");
+  // Ellipsoid is a special case, it uses a UsdGeomSphere and scales it with
+  // xformOp:scale. The radius is always set to 1.
+  EXPECT_PRIM_IS_A(stage, "/test/ellipsoid_geom", pxr::UsdGeomSphere);
+  ExpectAttributeEqual(stage, "/test/ellipsoid_geom.radius", 1.0);
+  ExpectAttributeEqual(stage, "/test/ellipsoid_geom.xformOp:scale",
+                       pxr::GfVec3f(2.0 * 10.0, 2.0 * 20.0, 2.0 * 30.0));
+}
+
 static constexpr char kSiteXml[] = R"(
     <mujoco model="test">
       <worldbody>
@@ -435,6 +492,7 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestSitePrimsPurpose) {
   EXPECT_PRIM_PURPOSE(stage, "/test/ball/ball/ellipsoid_site",
                       pxr::UsdGeomTokens->guide);
 }
+
 TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsToggleSdfFormatArg) {
   std::string xml_path = GetTestDataFilePath(kMeshObjPath);