chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support for mjcPhysics keyframes to USD interop.

Browse Source 
PiperOrigin-RevId: 778552679
Change-Id: I04026599f7a2836e3fe1cf0d553bcc86a5795f98
This commit is contained in:
Sam Haves
2025-07-02 10:16:42 -07:00
committed by Copybara-Service
parent 072b327430
commit 1d613b9b9a
7 changed files with 236 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
+94
View File
@@ -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
+6
View File
@@ -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
+24
View File
@@ -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
+64
View File
@@ -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);