When parsing USD, infer gravity from stage metadata if UsdPhysicsScene missing or undefined.

PiperOrigin-RevId: 813705210
Change-Id: I13741af6dc87ecd9efb83219e978f5f06c72d021

Author: havess

src/experimental/usd/usd_to_mjspec.cc

@@ -16,6 +16,7 @@
 #include <cstddef>
 #include <map>
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -51,6 +52,7 @@
 #include <pxr/usd/usdGeom/cylinder.h>
 #include <pxr/usd/usdGeom/gprim.h>
 #include <pxr/usd/usdGeom/mesh.h>
+#include <pxr/usd/usdGeom/metrics.h>
 #include <pxr/usd/usdGeom/plane.h>
 #include <pxr/usd/usdGeom/primvar.h>
 #include <pxr/usd/usdGeom/primvarsAPI.h>
@@ -383,20 +385,55 @@ mjsMesh* ParseUsdMesh(mjSpec* spec, const pxr::UsdPrim& prim, mjsGeom* geom,
   return mesh;
 }
 
+void SetGravityAttributes(
+    mjSpec* spec, const pxr::UsdStageRefPtr stage,
+    std::optional<pxr::GfVec3f> gravity_direction = std::nullopt,
+    std::optional<float> gravity_magnitude = std::nullopt) {
+  // Parse gravity and gravity direction.
+  if (!gravity_direction.has_value()) {
+    TfToken up_axis = pxr::UsdGeomGetStageUpAxis(stage);
+    if (up_axis == pxr::UsdGeomTokens->y) {
+      gravity_direction = pxr::GfVec3f(0, -1, 0);
+    } else if (up_axis == pxr::UsdGeomTokens->z) {
+      gravity_direction = pxr::GfVec3f(0, 0, -1);
+    } else {
+      mju_error("Invalid stage up axis token %s", up_axis.GetString().c_str());
+    }
+  }
+
+  if (!gravity_magnitude.has_value()) {
+    gravity_magnitude = 9.81f / pxr::UsdGeomGetStageMetersPerUnit(stage);
+  }
+
+  pxr::GfVec3f gravity = gravity_direction.value() * gravity_magnitude.value();
+
+  spec->option.gravity[0] = gravity[0];
+  spec->option.gravity[1] = gravity[1];
+  spec->option.gravity[2] = gravity[2];
+}
+
 void ParseUsdPhysicsScene(mjSpec* spec,
                           const pxr::UsdPhysicsScene& physics_scene) {
-  // Parse gravity and gravity direction.
-  pxr::GfVec3f gravity_direction;
-  physics_scene.GetGravityDirectionAttr().Get(&gravity_direction);
+  std::optional<pxr::GfVec3f> gravity_direction = std::nullopt;
+  std::optional<float> gravity_magnitude = std::nullopt;
+  auto stage = physics_scene.GetPrim().GetStage();
 
-  float gravity_magnitude;
-  physics_scene.GetGravityMagnitudeAttr().Get(&gravity_magnitude);
+  // Parse gravity and gravity direction.
+  auto gravity_direction_attr = physics_scene.GetGravityDirectionAttr();
+  if (gravity_direction_attr.HasAuthoredValue()) {
+    pxr::GfVec3f authored_gravity_dir;
+    gravity_direction_attr.Get(&authored_gravity_dir);
+    gravity_direction = authored_gravity_dir;
+  }
 
-  gravity_direction *= gravity_magnitude;
+  auto gravity_magnitude_attr = physics_scene.GetGravityMagnitudeAttr();
+  if (gravity_magnitude_attr.HasAuthoredValue()) {
+    float authored_gravity_magnitude;
+    gravity_magnitude_attr.Get(&authored_gravity_magnitude);
+    gravity_magnitude = authored_gravity_magnitude;
+  }
 
-  spec->option.gravity[0] = gravity_direction[0];
-  spec->option.gravity[1] = gravity_direction[1];
-  spec->option.gravity[2] = gravity_direction[2];
+  SetGravityAttributes(spec, stage, gravity_direction, gravity_magnitude);
 
   // Early exit if theres no MjcPhysicsSceneAPI applied.
   if (!physics_scene.GetPrim().HasAPI<pxr::MjcPhysicsSceneAPI>()) {
@@ -1724,6 +1761,10 @@ mjSpec* mj_parseUSDStage(const pxr::UsdStageRefPtr stage) {
   if (!root->physics_scene.IsEmpty()) {
     mujoco::usd::ParseUsdPhysicsScene(
         spec, pxr::UsdPhysicsScene::Get(stage, root->physics_scene));
+  } else {
+    // If there is no physics scene we still need to infer the gravity vector
+    // from the stage up axis and units per meter metadata.
+    mujoco::usd::SetGravityAttributes(spec, stage);
   }
 
   if (!root->keyframes.empty()) {
@@ -1754,4 +1795,3 @@ MJAPI mjSpec* mj_parseUSD(const char* identifier, const mjVFS* vfs, char* error,
   auto stage = pxr::UsdStage::Open(identifier);
   return mj_parseUSDStage(stage);
 }
-