diff --git a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc index 0a5ee0e2..8a92fa6d 100644 --- a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc +++ b/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 &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> 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) { diff --git a/src/experimental/usd/plugins/mjcf/utils.cc b/src/experimental/usd/plugins/mjcf/utils.cc index e919a672..db2eae85 100644 --- a/src/experimental/usd/plugins/mjcf/utils.cc +++ b/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(); diff --git a/src/experimental/usd/plugins/mjcf/utils.h b/src/experimental/usd/plugins/mjcf/utils.h index 9f41a697..5b7ad71c 100644 --- a/src/experimental/usd/plugins/mjcf/utils.h +++ b/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 +void SetFieldTimeSample(pxr::SdfAbstractDataRefPtr& data, + const pxr::SdfPath& field_path, double time, + T&& value) { + using Deduced = typename std::remove_reference_t; + const auto typed_val = pxr::SdfAbstractDataConstTypedValue(&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 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 +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 void SetLayerMetadata(pxr::SdfAbstractDataRefPtr& data, const pxr::TfToken& key, diff --git a/src/experimental/usd/usd_to_mjspec.cc b/src/experimental/usd/usd_to_mjspec.cc index 7a25e156..b9f6b741 100644 --- a/src/experimental/usd/usd_to_mjspec.cc +++ b/src/experimental/usd/usd_to_mjspec.cc @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -26,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -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** 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(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 times; + qpos_attr.GetTimeSamples(×); + 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()) { + ParseMjcPhysicsKeyframe(spec, pxr::MjcPhysicsKeyframe(prim)); + continue; } else if (prim.HasAPI()) { pxr::SdfPath body_path = prim.GetPath(); body_paths.push_back(body_path); diff --git a/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc b/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc index 98fafdf9..2fc83460 100644 --- a/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc +++ b/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc @@ -48,6 +48,7 @@ #include #include // IWYU pragma: keep, used for TraverseAll #include +#include #include #include #include @@ -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"( + + + + + + + + + + + + + + )"; + 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 diff --git a/test/experimental/usd/test_utils.cc b/test/experimental/usd/test_utils.cc index eb008a68..e051e447 100644 --- a/test/experimental/usd/test_utils.cc +++ b/test/experimental/usd/test_utils.cc @@ -30,6 +30,7 @@ #include #include #include +#include #include #include namespace mujoco { @@ -59,7 +60,8 @@ pxr::UsdStageRefPtr OpenStageWithPhysics(const std::string& xml) { template <> void ExpectAttributeEqual(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; diff --git a/test/experimental/usd/test_utils.h b/test/experimental/usd/test_utils.h index 46abaec2..db2e6a75 100644 --- a/test/experimental/usd/test_utils.h +++ b/test/experimental/usd/test_utils.h @@ -26,6 +26,7 @@ #include #include #include +#include #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 -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 -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::UsdStageRefPtr stage, - pxr::SdfPath, - const pxr::SdfAssetPath& value); +void ExpectAttributeEqual( + pxr::UsdStageRefPtr stage, pxr::SdfPath, const pxr::SdfAssetPath& value, + pxr::UsdTimeCode time); void ExpectAttributeHasConnection(pxr::UsdStageRefPtr stage, const char* path, const char* connection_path);