Add support for mjcPhysics keyframes to USD interop.

PiperOrigin-RevId: 778552679
Change-Id: I04026599f7a2836e3fe1cf0d553bcc86a5795f98

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

@@ -84,6 +84,9 @@ TF_DEFINE_PRIVATE_TOKENS(kTokens,
                          ((meshScope, "MeshSources"))
                          ((materialsScope, "Materials"))
                          ((previewSurface, "PreviewSurface"))
+                         ((keyframesScope, "Keyframes"))
+                         ((keyframe, "Keyframe"))
+                         ((surface, "PreviewSurface"))
                          ((world, "World"))
                          ((xformOpTransform, "xformOp:transform"))
                          ((xformOpScale, "xformOp:scale"))
@@ -144,6 +147,7 @@ using mujoco::usd::CreatePrimSpec;
 using mujoco::usd::CreateRelationshipSpec;
 using mujoco::usd::SetAttributeDefault;
 using mujoco::usd::SetAttributeMetadata;
+using mujoco::usd::SetAttributeTimeSample;
 using mujoco::usd::SetLayerMetadata;
 using mujoco::usd::SetPrimKind;
 using mujoco::usd::SetPrimMetadata;
@@ -207,6 +211,7 @@ class ModelWriter {
     if (write_physics_) {
       WriteActuators();
     }
+    WriteKeyframes();
   }
 
  private:
@@ -929,6 +934,95 @@ class ModelWriter {
     }
   }
 
+  void WriteKeyframesWithName(const std::string &name,
+                              const std::vector<mjsKey *> &keyframes,
+                              const pxr::SdfPath &parent_path) {
+    if (keyframes.empty()) {
+      return;
+    }
+    const auto keyframe_name = pxr::TfToken(pxr::TfMakeValidIdentifier(
+        name.empty() ? MjcPhysicsTokens->Keyframe : name));
+    pxr::SdfPath keyframe_path = parent_path.AppendChild(keyframe_name);
+    if (!data_->HasSpec(keyframe_path)) {
+      CreatePrimSpec(data_, parent_path, keyframe_name,
+                     pxr::MjcPhysicsTokens->Keyframe);
+    }
+    auto set_attribute_data = [&](const pxr::SdfPath &attr_path,
+                                  const pxr::VtDoubleArray &value,
+                                  mjsKey *keyframe) {
+      // If the keyframe time is the default, and there are no other keyframes
+      // set the attribute at the default time code.
+      if (keyframe->time == 0 && keyframes.size() == 1) {
+        SetAttributeDefault(data_, attr_path, value);
+      } else {
+        SetAttributeTimeSample(data_, attr_path, keyframe->time, value);
+      }
+    };
+
+    for (auto *keyframe : keyframes) {
+      pxr::SdfPath qpos_attr_path =
+          CreateAttributeSpec(data_, keyframe_path, MjcPhysicsTokens->mjcQpos,
+                              pxr::SdfValueTypeNames->DoubleArray);
+      set_attribute_data(qpos_attr_path,
+                         pxr::VtDoubleArray(keyframe->qpos->begin(),
+                                            keyframe->qpos->end()), keyframe);
+
+      pxr::SdfPath qvel_attr_path =
+          CreateAttributeSpec(data_, keyframe_path, MjcPhysicsTokens->mjcQvel,
+                              pxr::SdfValueTypeNames->DoubleArray);
+      set_attribute_data(qvel_attr_path,
+                         pxr::VtDoubleArray(keyframe->qvel->begin(),
+                                            keyframe->qvel->end()), keyframe);
+
+      pxr::SdfPath act_attr_path =
+          CreateAttributeSpec(data_, keyframe_path, MjcPhysicsTokens->mjcAct,
+                              pxr::SdfValueTypeNames->DoubleArray);
+      set_attribute_data(act_attr_path,
+                         pxr::VtDoubleArray(keyframe->act->begin(),
+                                            keyframe->act->end()), keyframe);
+
+      pxr::SdfPath ctrl_attr_path =
+          CreateAttributeSpec(data_, keyframe_path, MjcPhysicsTokens->mjcCtrl,
+                              pxr::SdfValueTypeNames->DoubleArray);
+      set_attribute_data(ctrl_attr_path,
+                         pxr::VtDoubleArray(keyframe->ctrl->begin(),
+                                            keyframe->ctrl->end()), keyframe);
+
+      pxr::SdfPath mpos_attr_path =
+          CreateAttributeSpec(data_, keyframe_path, MjcPhysicsTokens->mjcMpos,
+                              pxr::SdfValueTypeNames->DoubleArray);
+      set_attribute_data(mpos_attr_path,
+                         pxr::VtDoubleArray(keyframe->mpos->begin(),
+                                            keyframe->mpos->end()), keyframe);
+
+      pxr::SdfPath mquat_attr_path =
+          CreateAttributeSpec(data_, keyframe_path, MjcPhysicsTokens->mjcMquat,
+                              pxr::SdfValueTypeNames->DoubleArray);
+      set_attribute_data(mquat_attr_path,
+                         pxr::VtDoubleArray(keyframe->mquat->begin(),
+                                            keyframe->mquat->end()), keyframe);
+    }
+  }
+
+  void WriteKeyframes() {
+    std::unordered_map<std::string, std::vector<mjsKey *>> keyframes_map;
+    mjsKey *keyframe = mjs_asKey(mjs_firstElement(spec_, mjOBJ_KEY));
+    while (keyframe) {
+      std::string keyframe_name = keyframe->name->empty()
+                                      ? MjcPhysicsTokens->Keyframe
+                                      : *keyframe->name;
+      keyframes_map[keyframe_name].push_back(keyframe);
+      keyframe = mjs_asKey(mjs_nextElement(spec_, keyframe->element));
+    }
+
+    pxr::SdfPath scope_path =
+        CreatePrimSpec(data_, body_paths_[kWorldIndex], kTokens->keyframesScope,
+                       pxr::UsdGeomTokens->Scope);
+    for (const auto &[keyframe_name, keyframes] : keyframes_map) {
+      WriteKeyframesWithName(keyframe_name, keyframes, scope_path);
+    }
+  }
+
   void WriteActuator(mjsActuator *actuator) {
     pxr::SdfPath transmission_path;
     if (actuator->trntype == mjtTrn::mjTRN_BODY) {

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

@@ -102,6 +102,12 @@ pxr::SdfPath CreateAttributeSpec(pxr::SdfAbstractDataRefPtr& data,
                                  const pxr::SdfValueTypeName& type_name,
                                  pxr::SdfVariability variability) {
   const pxr::SdfPath propertyPath = prim_path.AppendProperty(name);
+
+  // Early exit if the attribute spec already exists.
+  if (data->HasSpec(propertyPath)) {
+    return propertyPath;
+  }
+
   data->CreateSpec(propertyPath, pxr::SdfSpecTypeAttribute);
 
   pxr::TfToken typeNameToken = type_name.GetAsToken();

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

@@ -85,6 +85,21 @@ void SetField(pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& field_path,
   data->Set(field_path, key, untyped_val);
 }
 
+// Set the value specified by key on any field at field_path.
+template <typename T>
+void SetFieldTimeSample(pxr::SdfAbstractDataRefPtr& data,
+                        const pxr::SdfPath& field_path, double time,
+                        T&& value) {
+  using Deduced = typename std::remove_reference_t<T>;
+  const auto typed_val = pxr::SdfAbstractDataConstTypedValue<Deduced>(&value);
+  const pxr::SdfAbstractDataConstValue& untyped_val = typed_val;
+
+  pxr::VtValue vt_value;
+  untyped_val.GetValue(&vt_value);
+  // NOTE: SetTimeSample doesn't accept an SdfAbstractDataConstValue yet.
+  data->SetTimeSample(field_path, time, vt_value);
+}
+
 // Set the value specified by key on an attribute spec at attribute_path.
 template <typename T>
 void SetAttribute(pxr::SdfAbstractDataRefPtr& data,
@@ -117,6 +132,15 @@ void SetAttributeDefault(pxr::SdfAbstractDataRefPtr& data,
   SetAttribute(data, attribute_path, pxr::SdfFieldKeys->Default, default_value);
 }
 
+// Set the default value on an attribute spec at attribute_path.
+template <typename T>
+void SetAttributeTimeSample(pxr::SdfAbstractDataRefPtr& data,
+                         const pxr::SdfPath& attribute_path,
+                         double time,
+                         T&& default_value) {
+  SetFieldTimeSample(data, attribute_path, time, default_value);
+}
+
 // Set the value specified by key on the root layer.
 template <typename T>
 void SetLayerMetadata(pxr::SdfAbstractDataRefPtr& data, const pxr::TfToken& key,

src/experimental/usd/usd_to_mjspec.cc

@@ -14,6 +14,7 @@
 
 #include <algorithm>
 #include <cmath>
+#include <cstddef>
 #include <deque>
 #include <iterator>
 #include <map>
@@ -26,6 +27,7 @@
 #include <mujoco/experimental/usd/mjcPhysics/actuatorAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/collisionAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/jointAPI.h>
+#include <mujoco/experimental/usd/mjcPhysics/keyframe.h>
 #include <mujoco/experimental/usd/mjcPhysics/meshCollisionAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/sceneAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/siteAPI.h>
@@ -1154,6 +1156,65 @@ void ParseMjcPhysicsSite(mjSpec* spec, const pxr::MjcPhysicsSiteAPI& site_api,
   }
 }
 
+void ParseMjcPhysicsKeyframe(mjSpec* spec,
+                             const pxr::MjcPhysicsKeyframe& keyframe) {
+  auto prim = keyframe.GetPrim();
+  auto qpos_attr = keyframe.GetMjcQposAttr();
+  auto qvel_attr = keyframe.GetMjcQvelAttr();
+  auto act_attr = keyframe.GetMjcActAttr();
+  auto ctrl_attr = keyframe.GetMjcCtrlAttr();
+  auto mpos_attr = keyframe.GetMjcMposAttr();
+  auto mquat_attr = keyframe.GetMjcMquatAttr();
+
+  auto setKeyframeData = [](mjsKey* key, const pxr::UsdAttribute& attr,
+                            std::vector<double>** key_data, double* time = nullptr) {
+    if (attr.HasAuthoredValue()) {
+      pxr::VtDoubleArray data;
+      if (time == nullptr) {
+        attr.Get(&data);
+      } else {
+        attr.Get(&data, *time);
+      }
+      *key_data = new std::vector<double>(data.begin(), data.end());
+    }
+  };
+
+  size_t n_time_samples = 0;
+  if (qpos_attr.HasAuthoredValue()) {
+    n_time_samples = qpos_attr.GetNumTimeSamples();
+  }
+
+  if (n_time_samples == 0) {
+    // If no time samples, we create a single keyframe.
+    mjsKey* key = mjs_addKey(spec);
+    mjs_setString(key->name, prim.GetName().GetString().c_str());
+    setKeyframeData(key, qpos_attr, &key->qpos);
+    setKeyframeData(key, qvel_attr, &key->qvel);
+    setKeyframeData(key, act_attr, &key->act);
+    setKeyframeData(key, ctrl_attr, &key->ctrl);
+    setKeyframeData(key, mpos_attr, &key->mpos);
+    setKeyframeData(key, mquat_attr, &key->mquat);
+  } else {
+    // If time samples, we create a keyframe for each time sample.
+    std::vector<double> times;
+    qpos_attr.GetTimeSamples(&times);
+    int keyframe_id = 0;
+    for (double time : times) {
+      mjsKey* key = mjs_addKey(spec);
+      std::string key_name =
+          prim.GetName().GetString() + "_" + std::to_string(keyframe_id++);
+      mjs_setString(key->name, key_name.c_str());
+      key->time = time;
+      setKeyframeData(key, qpos_attr, &key->qpos, &time);
+      setKeyframeData(key, qvel_attr, &key->qvel, &time);
+      setKeyframeData(key, act_attr, &key->act, &time);
+      setKeyframeData(key, ctrl_attr, &key->ctrl, &time);
+      setKeyframeData(key, mpos_attr, &key->mpos, &time);
+      setKeyframeData(key, mquat_attr, &key->mquat, &time);
+    }
+  }
+}
+
 mjsBody* ParseUsdPhysicsRigidbody(
     mjSpec* spec, const pxr::UsdPhysicsRigidBodyAPI& rigidbody_api,
     const pxr::UsdPrim& parent_prim, mjsBody* parent,
@@ -1563,6 +1624,9 @@ mjSpec* mj_parseUSDStage(const pxr::UsdStageRefPtr stage) {
       }
 
       AddEdge(edges, body0_path, body1_path, joint_path);
+    } else if (prim.IsA<pxr::MjcPhysicsKeyframe>()) {
+      ParseMjcPhysicsKeyframe(spec, pxr::MjcPhysicsKeyframe(prim));
+      continue;
     } else if (prim.HasAPI<pxr::UsdPhysicsRigidBodyAPI>()) {
       pxr::SdfPath body_path = prim.GetPath();
       body_paths.push_back(body_path);

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

@@ -48,6 +48,7 @@
 #include <pxr/usd/usd/prim.h>
 #include <pxr/usd/usd/primRange.h>  // IWYU pragma: keep, used for TraverseAll
 #include <pxr/usd/usd/stage.h>
+#include <pxr/usd/usd/timeCode.h>
 #include <pxr/usd/usdGeom/capsule.h>
 #include <pxr/usd/usdGeom/cube.h>
 #include <pxr/usd/usdGeom/cylinder.h>
@@ -1993,6 +1994,38 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsUnsupportedJoint) {
 
   EXPECT_PRIM_INVALID(stage, "/test/parent/ball_joint");
 }
+
+TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsKeyframe) {
+  static constexpr char xml[] = R"(
+  <mujoco model="test">
+      <worldbody>
+        <frame name="frame"/>
+        <body name="body">
+          <joint/>
+          <geom size="0.1"/>
+        </body>
+      </worldbody>
+      <keyframe>
+        <key name="home" qpos="1"/>
+        <key time="1" qpos="2"/>
+        <key time="2" qpos="3"/>
+      </keyframe>
+    </mujoco>)";
+  auto stage = OpenStageWithPhysics(xml);
+
+  EXPECT_PRIM_VALID(stage, "/test/Keyframes/home");
+  EXPECT_PRIM_VALID(stage, "/test/Keyframes/Keyframe");
+  ExpectAttributeEqual(stage, "/test/Keyframes/home.mjc:qpos",
+                       pxr::VtDoubleArray({1}));
+
+  // Check time samples are correctly authored.
+  ExpectAttributeEqual(stage, "/test/Keyframes/Keyframe.mjc:qpos",
+                       pxr::VtDoubleArray({2}), pxr::UsdTimeCode(1.0));
+
+  ExpectAttributeEqual(stage, "/test/Keyframes/Keyframe.mjc:qpos",
+                       pxr::VtDoubleArray({3}), pxr::UsdTimeCode(2.0));
+}
+
 }  // namespace
 }  // namespace usd
 }  // namespace mujoco

test/experimental/usd/test_utils.cc

@@ -30,6 +30,7 @@
 #include <pxr/usd/usd/common.h>
 #include <pxr/usd/usd/modelAPI.h>
 #include <pxr/usd/usd/stage.h>
+#include <pxr/usd/usd/timeCode.h>
 #include <pxr/usd/usdGeom/mesh.h>
 #include <pxr/usd/usdGeom/primvarsAPI.h>
 namespace mujoco {
@@ -59,7 +60,8 @@ pxr::UsdStageRefPtr OpenStageWithPhysics(const std::string& xml) {
 template <>
 void ExpectAttributeEqual<pxr::SdfAssetPath>(pxr::UsdStageRefPtr stage,
                                              pxr::SdfPath path,
-                                             const pxr::SdfAssetPath& value) {
+                                             const pxr::SdfAssetPath& value,
+                                             const pxr::UsdTimeCode time) {
   auto attr = stage->GetAttributeAtPath(path);
   EXPECT_TRUE(attr.IsValid());
   pxr::SdfAssetPath attr_value;

test/experimental/usd/test_utils.h

@@ -26,6 +26,7 @@
 #include <pxr/usd/usd/common.h>
 #include <pxr/usd/usd/modelAPI.h>
 #include <pxr/usd/usd/stage.h>
+#include <pxr/usd/usd/timeCode.h>
 
 #define EXPECT_PRIM_VALID(stage, path) \
   EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).IsValid());
@@ -101,20 +102,22 @@ pxr::SdfLayerRefPtr LoadLayer(
 pxr::UsdStageRefPtr OpenStageWithPhysics(const std::string& xml);
 
 template <typename T>
-void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, pxr::SdfPath path,
-                          const T& value) {
+void ExpectAttributeEqual(
+    pxr::UsdStageRefPtr stage, pxr::SdfPath path, const T& value,
+    const pxr::UsdTimeCode time = pxr::UsdTimeCode::Default()) {
   auto attr = stage->GetAttributeAtPath(pxr::SdfPath(path));
   EXPECT_TRUE(attr.IsValid()) << "Attribute " << path << " is not valid";
   T attr_value;
-  attr.Get(&attr_value);
+  attr.Get(&attr_value, time);
   EXPECT_EQ(attr_value, value) << "Attribute " << path << " has value "
                                << attr_value << ". Expected: " << value;
 }
 
 template <typename T>
-void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, const char* path,
-                          const T& value) {
-  ExpectAttributeEqual(stage, pxr::SdfPath(path), value);
+void ExpectAttributeEqual(
+    pxr::UsdStageRefPtr stage, const char* path, const T& value,
+    const pxr::UsdTimeCode time = pxr::UsdTimeCode::Default()) {
+  ExpectAttributeEqual(stage, pxr::SdfPath(path), value, time);
 }
 
 // Specialization for SdfAssetPath, so that we can compare only the asset path
@@ -122,9 +125,9 @@ void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, const char* path,
 // Otherwise the default operator== would fail because it tests for equality of
 // the asset path AND the resolved path.
 template <>
-void ExpectAttributeEqual<pxr::SdfAssetPath>(pxr::UsdStageRefPtr stage,
-                                             pxr::SdfPath,
-                                             const pxr::SdfAssetPath& value);
+void ExpectAttributeEqual<pxr::SdfAssetPath>(
+    pxr::UsdStageRefPtr stage, pxr::SdfPath, const pxr::SdfAssetPath& value,
+    pxr::UsdTimeCode time);
 
 void ExpectAttributeHasConnection(pxr::UsdStageRefPtr stage, const char* path,
                                   const char* connection_path);