Add support for inertia and mass to Mujoco USD interop.

PiperOrigin-RevId: 763737956
Change-Id: I6cdf267d76536a26aeb6327ce90165842e32bd8a

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

@@ -26,6 +26,7 @@
 #include "mjcf/utils.h"
 #include <pxr/base/arch/attributes.h>
 #include <pxr/base/gf/matrix4d.h>
+#include <pxr/base/gf/quatf.h>
 #include <pxr/base/gf/rotation.h>
 #include <pxr/base/gf/vec2f.h>
 #include <pxr/base/gf/vec3d.h>
@@ -180,7 +181,7 @@ class ModelWriter {
     WritePhysicsScene();
 
     // Author mesh scope + mesh prims to be referenced.
-    WriteMeshes();
+    WriteMeshes(write_physics);
     WriteMaterials();
     WriteBodies(write_physics);
   }
@@ -298,7 +299,8 @@ class ModelWriter {
     SetAttributeDefault(data_, xform_op_order_path, new_order);
   }
 
-  void WriteMesh(const mjsMesh *mesh, const pxr::SdfPath &parent_path) {
+  void WriteMesh(const mjsMesh *mesh, const pxr::SdfPath &parent_path,
+                 bool write_physics) {
     auto name = GetAvailablePrimName(*mesh->name, pxr::UsdGeomTokens->Mesh,
                                      parent_path);
     pxr::SdfPath subcomponent_path =
@@ -308,6 +310,22 @@ class ModelWriter {
                        pxr::UsdGeomTokens->Mesh);
     mesh_paths_[*mesh->name] = subcomponent_path;
 
+    if (write_physics) {
+      ApplyApiSchema(data_, mesh_path, MjcPhysicsTokens->MeshCollisionAPI);
+
+      pxr::TfToken inertia = MjcPhysicsTokens->legacy;
+      if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_EXACT) {
+        inertia = MjcPhysicsTokens->exact;
+      } else if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_CONVEX) {
+        inertia = MjcPhysicsTokens->convex;
+      } else if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_SHELL) {
+        inertia = MjcPhysicsTokens->shell;
+      }
+
+      WriteUniformAttribute(mesh_path, pxr::SdfValueTypeNames->Token,
+                            MjcPhysicsTokens->mjcInertia, inertia);
+    }
+
     // NOTE: The geometry data taken from the spec is the post-compilation
     // data after it has been mjCMesh::Compile'd. So don't be surprised if
     // things like user defined vertices have moved due to re-centering to
@@ -584,7 +602,7 @@ class ModelWriter {
     }
   }
 
-  void WriteMeshes() {
+  void WriteMeshes(bool write_physics) {
     // Create a scope for the meshes to keep things organized
     pxr::SdfPath scope_path =
         CreatePrimSpec(data_, body_paths_[kWorldIndex], kTokens->meshScope,
@@ -596,7 +614,7 @@ class ModelWriter {
 
     mjsMesh *mesh = mjs_asMesh(mjs_firstElement(spec_, mjOBJ_MESH));
     while (mesh) {
-      WriteMesh(mesh, scope_path);
+      WriteMesh(mesh, scope_path, write_physics);
       mesh = mjs_asMesh(mjs_nextElement(spec_, mesh->element));
     }
   }
@@ -1032,12 +1050,43 @@ class ModelWriter {
         return;
     }
 
-    // Apply the PhysicsCollisionAPI schema if we are writing physics and the
+    // Apply the physics schemas if we are writing physics and the
     // geom participates in collisions.
     if (write_physics && (model_->geom_contype[geom_id] != 0 ||
                           model_->geom_conaffinity[geom_id] != 0)) {
       ApplyApiSchema(data_, geom_path,
                      pxr::UsdPhysicsTokens->PhysicsCollisionAPI);
+      ApplyApiSchema(data_, geom_path, MjcPhysicsTokens->CollisionAPI);
+
+      WriteUniformAttribute(
+          geom_path, pxr::SdfValueTypeNames->Bool,
+          MjcPhysicsTokens->mjcShellinertia,
+          geom->typeinertia == mjtGeomInertia::mjINERTIA_SHELL);
+
+      if (geom->mass >= mjMINVAL || geom->density >= mjMINVAL) {
+        ApplyApiSchema(data_, geom_path, pxr::UsdPhysicsTokens->PhysicsMassAPI);
+      }
+
+      if (geom->mass >= mjMINVAL) {
+        pxr::SdfPath mass_attr = CreateAttributeSpec(
+            data_, geom_path, pxr::UsdPhysicsTokens->physicsMass,
+            pxr::SdfValueTypeNames->Float, pxr::SdfVariabilityUniform);
+
+        // Make sure to cast to float here since mjtNum might be a double.
+        SetAttributeDefault(data_, mass_attr, (float)geom->mass);
+      }
+
+      // Even though density is not used for mass computation when mass exists
+      // we want to retain the information anyways.
+      if (geom->density >= mjMINVAL) {
+        pxr::SdfPath density_attr = CreateAttributeSpec(
+            data_, geom_path, pxr::UsdPhysicsTokens->physicsDensity,
+            pxr::SdfValueTypeNames->Float, pxr::SdfVariabilityUniform);
+
+        // Make sure to cast to float here since mjtNum might be a double.
+        SetAttributeDefault(data_, density_attr, (float)geom->density);
+      }
+
       // For meshes, also apply PhysicsMeshCollisionAPI and set the
       // approximation attribute.
       if (geom->type == mjGEOM_MESH) {
@@ -1238,6 +1287,44 @@ class ModelWriter {
 
     // Apply the PhysicsRigidBodyAPI schema if we are writing physics.
     if (write_physics) {
+      // If the body had a mass specified then it must have either inertia or
+      // fullinertia specified per inertia element XML documentation.
+      // Therefore it is sufficient to check if the mass is non-zero to see if
+      // we should set inertial attributes on the body.
+      //
+      // Note that if the user has NOT specified any inertial properties then
+      // we don't want to pull values from the compiled model since coming back
+      // into Mujoco would take those values instead of computing them
+      // automatically from the subtree.
+      if (body->mass > 0) {
+        // User might have specified the inertia via fullinertia and the
+        // compiler has extracted all values properly. So leverage those
+        // instead of doing the computation ourselves here.
+        ApplyApiSchema(data_, body_path, pxr::UsdPhysicsTokens->PhysicsMassAPI);
+        WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Float,
+                              pxr::UsdPhysicsTokens->physicsMass,
+                              (float)model_->body_mass[body_id]);
+
+        mjtNum *body_ipos = &model_->body_ipos[body_id * 3];
+        pxr::GfVec3f inertial_pos(body_ipos[0], body_ipos[1], body_ipos[2]);
+        WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Point3f,
+                              pxr::UsdPhysicsTokens->physicsCenterOfMass,
+                              inertial_pos);
+
+        mjtNum *body_iquat = &model_->body_iquat[body_id * 4];
+        pxr::GfQuatf inertial_frame(body_iquat[0], body_iquat[1], body_iquat[2],
+                                    body_iquat[3]);
+        WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Quatf,
+                              pxr::UsdPhysicsTokens->physicsPrincipalAxes,
+                              inertial_frame);
+
+        mjtNum *inertia = &model_->body_inertia[body_id * 3];
+        pxr::GfVec3f diag_inertia(inertia[0], inertia[1], inertia[2]);
+        WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Float3,
+                              pxr::UsdPhysicsTokens->physicsDiagonalInertia,
+                              diag_inertia);
+      }
+
       ApplyApiSchema(data_, body_path,
                      pxr::UsdPhysicsTokens->PhysicsRigidBodyAPI);
 

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

@@ -17,6 +17,8 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "src/experimental/usd/mjcPhysics/collisionAPI.h"
+#include "src/experimental/usd/mjcPhysics/meshCollisionAPI.h"
 #include "src/experimental/usd/mjcPhysics/sceneAPI.h"
 #include "src/experimental/usd/mjcPhysics/siteAPI.h"
 #include "src/experimental/usd/mjcPhysics/tokens.h"
@@ -50,6 +52,7 @@
 #include <pxr/usd/usdGeom/tokens.h>
 #include <pxr/usd/usdPhysics/articulationRootAPI.h>
 #include <pxr/usd/usdPhysics/collisionAPI.h>
+#include <pxr/usd/usdPhysics/massAPI.h>
 #include <pxr/usd/usdPhysics/meshCollisionAPI.h>
 #include <pxr/usd/usdPhysics/rigidBodyAPI.h>
 #include <pxr/usd/usdPhysics/scene.h>
@@ -1106,6 +1109,7 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsColliders) {
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/ground",
                               pxr::UsdPhysicsRigidBodyAPI);
   EXPECT_PRIM_API_APPLIED(stage, "/test/ground", pxr::UsdPhysicsCollisionAPI);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/ground", pxr::MjcPhysicsCollisionAPI);
 
   // body_0/body_0_0 [rigidbody] (Nested body - reparented)
   EXPECT_PRIM_VALID(stage, "/test/body_0/body_0_0");
@@ -1113,30 +1117,40 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsColliders) {
                           pxr::UsdPhysicsRigidBodyAPI);
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_0/body_0_0",
                               pxr::UsdPhysicsCollisionAPI);
+  EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_0/body_0_0",
+                              pxr::MjcPhysicsCollisionAPI);
   //   body_0/body_0_0/body_0_0_col [collider]
   EXPECT_PRIM_VALID(stage, "/test/body_0/body_0_0/body_0_0_col");
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_0/body_0_0/body_0_0_col",
                               pxr::UsdPhysicsRigidBodyAPI);
   EXPECT_PRIM_API_APPLIED(stage, "/test/body_0/body_0_0/body_0_0_col",
                           pxr::UsdPhysicsCollisionAPI);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/body_0/body_0_0/body_0_0_col",
+                          pxr::MjcPhysicsCollisionAPI);
 
   // body_1 [rigidbody]
   EXPECT_PRIM_VALID(stage, "/test/body_1");
   EXPECT_PRIM_API_APPLIED(stage, "/test/body_1", pxr::UsdPhysicsRigidBodyAPI);
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1",
                               pxr::UsdPhysicsCollisionAPI);
+  EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1",
+                              pxr::MjcPhysicsCollisionAPI);
   //   body_1/body_1_col_0 [collider]
   EXPECT_PRIM_VALID(stage, "/test/body_1/body_1_col_0");
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1/body_1_col_0",
                               pxr::UsdPhysicsRigidBodyAPI);
   EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_0",
                           pxr::UsdPhysicsCollisionAPI);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_0",
+                          pxr::MjcPhysicsCollisionAPI);
   //   body_1/body_1_col_1 [collider]
   EXPECT_PRIM_VALID(stage, "/test/body_1/body_1_col_1");
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1/body_1_col_1",
                               pxr::UsdPhysicsRigidBodyAPI);
   EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_1",
                           pxr::UsdPhysicsCollisionAPI);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_1",
+                          pxr::MjcPhysicsCollisionAPI);
 
   // body_2 [rigidbody]
   EXPECT_PRIM_VALID(stage, "/test/body_2");
@@ -1169,11 +1183,117 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsColliders) {
                           pxr::UsdPhysicsCollisionAPI);
   EXPECT_PRIM_API_APPLIED(stage, "/test/body_3/body_3_col/Mesh",
                           pxr::UsdPhysicsMeshCollisionAPI);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/body_3/body_3_col/Mesh",
+                          pxr::MjcPhysicsMeshCollisionAPI);
   ExpectAttributeEqual(stage,
                        "/test/body_3/body_3_col/Mesh.physics:approximation",
                        pxr::UsdPhysicsTokens->convexHull);
 }
 
+TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsCollisionAPI) {
+  static constexpr char xml[] = R"(
+  <mujoco model="test">
+    <worldbody>
+      <body name="body">
+        <geom name="box" type="box" size=".05 .05 .05" mass="0.1" shellinertia="true"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  auto stage = OpenStageWithPhysics(xml);
+
+  ExpectAttributeEqual(stage, "/test/body/box.mjc:shellinertia", true);
+}
+
+TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsMeshCollisionAPI) {
+  static constexpr char xml[] = R"(
+  <mujoco model="test">
+    <asset>
+      <mesh name="tet_legacy" inertia="legacy" vertex="0 0 0  1 0 0  0 1 0  0 0 1"/>
+      <mesh name="tet_exact" inertia="exact" vertex="0 0 0  1 0 0  0 1 0  0 0 1"/>
+      <mesh name="tet_convex" inertia="convex" vertex="0 0 0  1 0 0  0 1 0  0 0 1"/>
+      <mesh name="tet_shell" inertia="shell" vertex="0 0 0  1 0 0  0 1 0  0 0 1"/>
+    </asset>
+    <worldbody>
+      <body name="body">
+        <geom name="tet_legacy" type="mesh" mesh="tet_legacy"/>
+        <geom name="tet_exact" type="mesh" mesh="tet_exact"/>
+        <geom name="tet_convex" type="mesh" mesh="tet_convex"/>
+        <geom name="tet_shell" type="mesh" mesh="tet_shell"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  auto stage = OpenStageWithPhysics(xml);
+
+  ExpectAttributeEqual(stage, "/test/body/tet_legacy/Mesh.mjc:inertia",
+                       MjcPhysicsTokens->legacy);
+  ExpectAttributeEqual(stage, "/test/body/tet_exact/Mesh.mjc:inertia",
+                       MjcPhysicsTokens->exact);
+  ExpectAttributeEqual(stage, "/test/body/tet_convex/Mesh.mjc:inertia",
+                       MjcPhysicsTokens->convex);
+  ExpectAttributeEqual(stage, "/test/body/tet_shell/Mesh.mjc:inertia",
+                       MjcPhysicsTokens->shell);
+}
+
+TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIApplied) {
+  static constexpr char xml[] = R"(
+  <mujoco model="test">
+    <worldbody>
+      <body name="body">
+        <geom name="box" type="box" size=".05 .05 .05" mass="0.1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  auto stage = OpenStageWithPhysics(xml);
+
+  EXPECT_PRIM_VALID(stage, "/test/body");
+  EXPECT_PRIM_VALID(stage, "/test/body/box");
+  EXPECT_PRIM_API_APPLIED(stage, "/test/body/box", pxr::UsdPhysicsMassAPI);
+  EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body", pxr::UsdPhysicsMassAPI);
+  ExpectAttributeEqual(stage, "/test/body/box.physics:mass", 0.1f);
+}
+
+TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIAppliedToBody) {
+  static constexpr char xml[] = R"(
+  <mujoco model="test">
+    <worldbody>
+      <body name="body">
+        <inertial pos="1 2 3" mass="3" diaginertia="1 1 1"/>
+        <geom name="box" type="box" size=".05 .05 .05" mass="0.1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  auto stage = OpenStageWithPhysics(xml);
+
+  EXPECT_PRIM_VALID(stage, "/test/body");
+  EXPECT_PRIM_VALID(stage, "/test/body/box");
+  EXPECT_PRIM_API_APPLIED(stage, "/test/body/box", pxr::UsdPhysicsMassAPI);
+  EXPECT_PRIM_API_APPLIED(stage, "/test/body", pxr::UsdPhysicsMassAPI);
+  // Make sure that body gets it's inertial elements from the inertial element
+  // and not from the subtree.
+  ExpectAttributeEqual(stage, "/test/body.physics:mass", 3.0f);
+  ExpectAttributeEqual(stage, "/test/body.physics:centerOfMass",
+                       pxr::GfVec3f(1, 2, 3));
+}
+
+TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIDensity) {
+  static constexpr char xml[] = R"(
+  <mujoco model="test">
+    <worldbody>
+      <body name="body">
+        <geom name="box" type="box" size=".05 .05 .05" density="1234"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  auto stage = OpenStageWithPhysics(xml);
+
+  ExpectAttributeEqual(stage, "/test/body/box.physics:density", 1234.0f);
+}
+
 }  // namespace
 }  // namespace usd
 }  // namespace mujoco

test/experimental/usd/test_utils.cc

@@ -42,6 +42,15 @@ pxr::SdfLayerRefPtr LoadLayer(
   return layer;
 }
 
+pxr::UsdStageRefPtr OpenStageWithPhysics(const std::string& xml) {
+  pxr::SdfFileFormat::FileFormatArguments args;
+  args["usdMjcfToggleUsdPhysics"] = "true";
+  pxr::SdfLayerRefPtr layer = LoadLayer(xml, args);
+  auto stage = pxr::UsdStage::Open(layer);
+  EXPECT_THAT(stage, testing::NotNull());
+  return stage;
+}
+
 template <>
 void ExpectAttributeEqual<pxr::SdfAssetPath>(pxr::UsdStageRefPtr stage,
                                              pxr::SdfPath path,

test/experimental/usd/test_utils.h

@@ -73,6 +73,8 @@ pxr::SdfLayerRefPtr LoadLayer(
     const std::string& xml,
     const pxr::SdfFileFormat::FileFormatArguments& args = {});
 
+pxr::UsdStageRefPtr OpenStageWithPhysics(const std::string& xml);
+
 template <typename T>
 void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, pxr::SdfPath path,
                           const T& value) {