Initial support for writing UsdPhysics data starting with UsdPhysicsRigidbodyAPI support.

The SDF_FORMAT_ARGS `usdMjcfToggleUsdPhysics` is introduced to toggle writing UsdPhysics or not when converting. Off by default

PiperOrigin-RevId: 746467942
Change-Id: I26cb9a4c4ae2eef96163dab967a74f14940a2a6d

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

@@ -26,6 +26,7 @@
 #include <pxr/base/tf/pathUtils.h>
 #include <pxr/base/tf/registryManager.h>
 #include <pxr/base/tf/staticTokens.h>
+#include <pxr/base/tf/stringUtils.h>
 #include <pxr/base/tf/type.h>
 #include <pxr/base/work/loops.h>
 #include <pxr/pxr.h>
@@ -145,9 +146,18 @@ bool UsdMjcfFileFormat::CanRead(const std::string &filePath) const {
 }
 
 bool UsdMjcfFileFormat::ReadImpl(pxr::SdfLayer *layer, mjSpec *spec) const {
-  auto data = InitData(layer->GetFileFormatArguments());
+  auto args = layer->GetFileFormatArguments();
 
-  auto success = mujoco::usd::WriteSpecToData(spec, data);
+  bool toggleUsdPhysics = false;
+  const auto it =
+      args.find(UsdMjcfFileFormatTokens->ToggleUsdPhysicsArg.GetString());
+  if (it != args.end()) {
+    toggleUsdPhysics = pxr::TfUnstringify<bool>(it->second);
+  }
+
+  auto data = InitData(args);
+
+  auto success = mujoco::usd::WriteSpecToData(spec, data, toggleUsdPhysics);
   mj_deleteSpec(spec);
   if (!success) {
     return false;

src/experimental/usd/plugins/mjcf/mjcf_file_format.h

@@ -27,10 +27,15 @@
 
 PXR_NAMESPACE_OPEN_SCOPE
 
+// clang-format off
 // The Id should realistically be mjcf, but the id and extension need to match.
 // So near term it just assumes the only .xml file we would import is MJCF.
-#define USD_MJCF_FILE_FORMAT_TOKENS \
+#define USD_MJCF_FILE_FORMAT_TOKENS  \
-  ((Id, "xml"))((Version, "1.0"))((Target, "usd"))
+    ((Id, "xml"))                    \
+    ((Version, "1.0"))               \
+    ((Target, "usd"))                \
+    ((ToggleUsdPhysicsArg, "usdMjcfToggleUsdPhysics"))
+// clang-format on
 
 TF_DECLARE_PUBLIC_TOKENS(UsdMjcfFileFormatTokens, USD_MJCF_FILE_FORMAT_TOKENS);
 

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

@@ -49,6 +49,7 @@
 #include <pxr/usd/usdGeom/metrics.h>
 #include <pxr/usd/usdGeom/tokens.h>
 #include <pxr/usd/usdLux/tokens.h>
+#include <pxr/usd/usdPhysics/tokens.h>
 #include <pxr/usd/usdShade/tokens.h>
 #include <pxr/usdImaging/usdImaging/tokens.h>
 
@@ -149,7 +150,7 @@ class ModelWriter {
   }
   ~ModelWriter() { mj_deleteModel(model_); }
 
-  void Write() {
+  void Write(bool write_physics) {
     // Create top level class holder.
     class_path_ = CreateClassSpec(data_, pxr::SdfPath::AbsoluteRootPath(),
                                   pxr::TfToken("__class__"));
@@ -173,7 +174,7 @@ class ModelWriter {
     // Author mesh scope + mesh prims to be referenced.
     WriteMeshes();
     WriteMaterials();
-    WriteBodies();
+    WriteBodies(write_physics);
   }
 
  private:
@@ -977,7 +978,7 @@ class ModelWriter {
     }
   }
 
-  void WriteBody(mjsBody *body) {
+  void WriteBody(mjsBody *body, bool write_physics) {
     int body_id = mjs_getId(body->element);
     pxr::SdfPath parent_path =
         CreateParentIfNotExists(body, body_paths_[kWorldIndex], data_);
@@ -990,6 +991,12 @@ class ModelWriter {
     // bodies subcomponents.
     SetPrimKind(data_, body_path, pxr::KindTokens->subcomponent);
 
+    // Apply the PhysicsRigidBodyAPI schema if we are writing physics.
+    if (write_physics) {
+      ApplyApiSchema(data_, body_path,
+                     pxr::UsdPhysicsTokens->PhysicsRigidBodyAPI);
+    }
+
     // Create classes if necessary
     mjsDefault *spec_default = mjs_getDefault(body->element);
 
@@ -1028,14 +1035,14 @@ class ModelWriter {
     body_paths_[body_id] = body_path;
   }
 
-  void WriteBodies() {
+  void WriteBodies(bool write_physics) {
     mjsBody *body = mjs_asBody(mjs_firstElement(spec_, mjOBJ_BODY));
     while (body) {
       // Only write a rigidbody if we are not the world body.
       // We fall through since the world body might have static
       // geom children.
       if (mjs_getId(body->element) != kWorldIndex) {
-        WriteBody(body);
+        WriteBody(body, write_physics);
       }
       WriteSites(body);
       WriteGeoms(body);
@@ -1061,7 +1068,8 @@ class ModelWriter {
 namespace mujoco {
 namespace usd {
 
-bool WriteSpecToData(mjSpec *spec, pxr::SdfAbstractDataRefPtr &data) {
+bool WriteSpecToData(mjSpec *spec, pxr::SdfAbstractDataRefPtr &data,
+                     bool write_physics) {
   // Create pseudo root first.
   data->CreateSpec(pxr::SdfPath::AbsoluteRootPath(),
                    pxr::SdfSpecTypePseudoRoot);
@@ -1072,7 +1080,7 @@ bool WriteSpecToData(mjSpec *spec, pxr::SdfAbstractDataRefPtr &data) {
     return false;
   }
 
-  ModelWriter(spec, model, data).Write();
+  ModelWriter(spec, model, data).Write(write_physics);
 
   return true;
 }

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

@@ -25,7 +25,9 @@ namespace usd {
 // Args:
 //   spec: mjSpec built programmatically or via parsed XML.
 //   data: SdfAbstractDataRefPtr that will be written to.
-bool WriteSpecToData(mjSpec* spec, pxr::SdfAbstractDataRefPtr& data);
+//   write_physics: Whether to write physics data.
+bool WriteSpecToData(mjSpec* spec, pxr::SdfAbstractDataRefPtr& data,
+                     bool write_physics);
 }  // namespace usd
 }  // namespace mujoco
 

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

@@ -32,9 +32,11 @@ namespace mujoco {
 
 using pxr::SdfPath;
 
-pxr::SdfLayerRefPtr LoadLayer(const std::string& xml) {
+pxr::SdfLayerRefPtr LoadLayer(
+    const std::string& xml,
+    const pxr::SdfFileFormat::FileFormatArguments& args) {
   auto layer = pxr::SdfLayer::CreateAnonymous(
-      "test_layer", pxr::SdfFileFormat::FindByExtension("xml"));
+      "test_layer", pxr::SdfFileFormat::FindByExtension("xml"), args);
   layer->ImportFromString(xml);
   EXPECT_THAT(layer, testing::NotNull());
   return layer;

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

@@ -21,6 +21,7 @@
 #include "test/fixture.h"
 #include <pxr/usd/sdf/assetPath.h>
 #include <pxr/usd/sdf/declareHandles.h>
+#include <pxr/usd/sdf/fileFormat.h>
 #include <pxr/usd/sdf/path.h>
 #include <pxr/usd/usd/common.h>
 #include <pxr/usd/usd/modelAPI.h>
@@ -34,6 +35,16 @@
     EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).IsA<type>()); \
   }
 
+#define EXPECT_PRIM_API_APPLIED(stage, path, api)                     \
+  {                                                                   \
+    EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).HasAPI<api>()); \
+  }
+
+#define EXPECT_PRIM_API_NOT_APPLIED(stage, path, api)                  \
+  {                                                                    \
+    EXPECT_FALSE((stage)->GetPrimAtPath(SdfPath(path)).HasAPI<api>()); \
+  }
+
 #define EXPECT_PRIM_KIND(stage, path, kind)                          \
   {                                                                  \
     pxr::TfToken prim_kind;                                          \
@@ -51,7 +62,9 @@
   }
 namespace mujoco {
 
-pxr::SdfLayerRefPtr LoadLayer(const std::string& xml);
+pxr::SdfLayerRefPtr LoadLayer(
+    const std::string& xml,
+    const pxr::SdfFileFormat::FileFormatArguments& args = {});
 
 template <typename T>
 void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, const char* path,

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

@@ -44,6 +44,7 @@
 #include <pxr/usd/usdGeom/primvarsAPI.h>
 #include <pxr/usd/usdGeom/sphere.h>
 #include <pxr/usd/usdGeom/tokens.h>
+#include <pxr/usd/usdPhysics/rigidBodyAPI.h>
 
 PXR_NAMESPACE_OPEN_SCOPE
 // clang-format off
@@ -434,6 +435,72 @@ 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);
+
+  // Test that the default is no physics.
+  auto stage_no_physics = pxr::UsdStage::Open(xml_path);
+  EXPECT_THAT(stage_no_physics, testing::NotNull());
+  EXPECT_PRIM_VALID(stage_no_physics, "/mesh_test/test_body/test_body");
+  EXPECT_PRIM_API_NOT_APPLIED(stage_no_physics,
+                              "/mesh_test/test_body/test_body",
+                              pxr::UsdPhysicsRigidBodyAPI);
+
+  // Then test that the physics flag enables physics.
+  std::string xml_path_physics_flag =
+      xml_path + ":SDF_FORMAT_ARGS:usdMjcfToggleUsdPhysics=true";
+  auto stage_with_physics = pxr::UsdStage::Open(xml_path_physics_flag);
+  EXPECT_THAT(stage_with_physics, testing::NotNull());
+
+  EXPECT_PRIM_VALID(stage_with_physics, "/mesh_test/test_body/test_body");
+  EXPECT_PRIM_API_APPLIED(stage_with_physics, "/mesh_test/test_body/test_body",
+                          pxr::UsdPhysicsRigidBodyAPI);
+}
+
+TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsRigidBody) {
+  static constexpr char kXml[] = R"(
+    <mujoco model="physics_test">
+      <worldbody>
+        <body name="test_body" pos="0 0 0">
+          <geom name="test_geom" type="sphere" size="1"/>
+          <body name="test_body_2" pos="2 0 0">
+            <geom name="test_geom_2" type="sphere" size="1"/>
+          </body>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+
+  pxr::SdfFileFormat::FileFormatArguments args;
+  args["usdMjcfToggleUsdPhysics"] = "true";
+  pxr::SdfLayerRefPtr layer = LoadLayer(kXml, args);
+  auto stage = pxr::UsdStage::Open(layer);
+
+  EXPECT_THAT(stage, testing::NotNull());
+  EXPECT_PRIM_VALID(stage, "/physics_test");
+  EXPECT_PRIM_VALID(stage, "/physics_test/test_body");
+  EXPECT_PRIM_VALID(stage, "/physics_test/test_body/test_body");
+  // USD does not allow nested rigidbodies so we put them as siblings to the
+  // first body in the hierarchy.
+  EXPECT_PRIM_VALID(stage, "/physics_test/test_body/test_body_2");
+
+  // The parent containing the body should not have the RigidBodyAPI applied.
+  EXPECT_PRIM_API_NOT_APPLIED(stage, "/physics_test/test_body",
+                              pxr::UsdPhysicsRigidBodyAPI);
+
+  EXPECT_PRIM_API_APPLIED(stage, "/physics_test/test_body/test_body",
+                          pxr::UsdPhysicsRigidBodyAPI);
+  EXPECT_PRIM_API_APPLIED(stage, "/physics_test/test_body/test_body_2",
+                          pxr::UsdPhysicsRigidBodyAPI);
+
+  // Geoms should not have RigidBodyAPI applied either.
+  EXPECT_PRIM_API_NOT_APPLIED(stage,
+                              "/physics_test/test_body/test_body/test_geom",
+                              pxr::UsdPhysicsRigidBodyAPI);
+  EXPECT_PRIM_API_NOT_APPLIED(stage,
+                              "/physics_test/test_body/test_body_2/test_geom_2",
+                              pxr::UsdPhysicsRigidBodyAPI);
+}
 
 }  // namespace
 }  // namespace mujoco