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

Make MJCF SdfFileFormat plugin create hierarchy reflecting MJCF instead of sibling rigidbodies.

Browse Source 
PiperOrigin-RevId: 762117447
Change-Id: Id13f7a4f4e383a0c9e670e5272836f2fb1099830
This commit is contained in:
Sam Haves
2025-05-22 14:03:52 -07:00
committed by Copybara-Service
parent 74cc904edc
commit 76ddc30072
2 changed files with 117 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
+25 -54
View File
@@ -154,7 +154,6 @@ class ModelWriter {
ModelWriter(mjSpec *spec, mjModel *model, pxr::SdfAbstractDataRefPtr &data)
: spec_(spec), model_(model), data_(data), class_path_("/Bad_Path") {
body_paths_ = std::vector<pxr::SdfPath>(model->nbody);
body_xforms_ = std::vector<pxr::GfMatrix4d>(model->nbody);
}
~ModelWriter() { mj_deleteModel(model_); }
@@ -165,7 +164,6 @@ class ModelWriter {
// Create the world body.
body_paths_[kWorldIndex] = WriteWorldBody(kWorldIndex);
body_xforms_[kWorldIndex] = pxr::GfMatrix4d().SetIdentity();
SetLayerMetadata(data_, pxr::SdfFieldKeys->Documentation,
"Generated by mujoco model writer.");
@@ -198,8 +196,6 @@ class ModelWriter {
pxr::SdfPath class_path_;
// Mapping from Mujoco body id to SdfPath.
std::vector<pxr::SdfPath> body_paths_;
// Mapping from Mujoco body id to world space transform.
std::vector<pxr::GfMatrix4d> body_xforms_;
// Mapping from mesh names to Mesh prim path.
std::unordered_map<std::string, pxr::SdfPath> mesh_paths_;
@@ -248,44 +244,6 @@ class ModelWriter {
pxr::TfToken body_name;
};
pxr::SdfPath CreateParentIfNotExists(mjsBody *body,
const pxr::SdfPath &world_path,
pxr::SdfAbstractDataRefPtr &data) {
// To allow for easier scene authoring and modification, we want to
// place MJCF bodies belonging to the same kinematic chain under some
// identity parent Xform prim. This allows users to move the entire
// asset.
//
// We cannot simply recreate the MJCF kinematic tree structure
// because in USD it is assumed that children move rigidly with their
// parents. This is not true in MJCF if you have joints. We could perhaps
// use a more complex heuristic where we evaluate a common tree prefix
// in MJCF that is effectively welded together but for now we choose
// simplicity.
// In the trivial case where the parent of body is already the world
// body we want to create a parent xform of the same name.
// So if the MJCF has a child of the world body called "root" we will
// create a parent Xform at /World/root and the actual body will be
// created at /World/root/root.
mjsBody *last_parent = body;
mjsBody *parent = mjs_getParent(body->element);
while (mjs_getId(parent->element) != kWorldIndex) {
last_parent = parent;
parent = mjs_getParent(parent->element);
}
pxr::TfToken last_parent_name = GetValidPrimName(*last_parent->name);
pxr::SdfPath parent_xform_path = world_path.AppendChild(last_parent_name);
if (!data->HasSpec(parent_xform_path)) {
pxr::SdfPath prim_path = CreatePrimSpec(
data, world_path, last_parent_name, pxr::UsdGeomTokens->Xform);
SetPrimKind(data_, prim_path, pxr::KindTokens->component);
}
return parent_xform_path;
}
void WriteScaleXformOp(const pxr::SdfPath &prim_path,
const pxr::GfVec3f &scale) {
pxr::SdfPath scale_attr_path =
@@ -1263,8 +1221,10 @@ class ModelWriter {
void WriteBody(mjsBody *body, bool write_physics) {
int body_id = mjs_getId(body->element);
pxr::SdfPath parent_path =
CreateParentIfNotExists(body, body_paths_[kWorldIndex], data_);
// This should be safe as we process parent bodies before children.
mjsBody *parent = mjs_getParent(body->element);
int parent_id = mjs_getId(parent->element);
pxr::SdfPath parent_path = body_paths_[parent_id];
pxr::TfToken body_name = GetValidPrimName(*body->name);
// Create Xform prim for body.
@@ -1272,12 +1232,25 @@ class ModelWriter {
pxr::UsdGeomTokens->Xform);
// The parent_path will be a component which makes the actual articulated
// bodies subcomponents.
SetPrimKind(data_, body_path, pxr::KindTokens->subcomponent);
auto kind = parent_id == kWorldIndex ? pxr::KindTokens->component
: pxr::KindTokens->subcomponent;
SetPrimKind(data_, body_path, kind);
// Apply the PhysicsRigidBodyAPI schema if we are writing physics.
if (write_physics) {
ApplyApiSchema(data_, body_path,
pxr::UsdPhysicsTokens->PhysicsRigidBodyAPI);
// If the parent is not the world body, but is child of the world body
// then we need to apply the articulation root API.
if (parent_id != kWorldIndex) {
int parent_parent_id =
mjs_getId(mjs_getParent(parent->element)->element);
if (parent_parent_id == kWorldIndex) {
ApplyApiSchema(data_, parent_path,
pxr::UsdPhysicsTokens->PhysicsArticulationRootAPI);
}
}
}
// Create classes if necessary
@@ -1290,17 +1263,15 @@ class ModelWriter {
}
// Create XformOp attribute for body transform.
// Make sure to account for the parent since UsdPhysics doesn't support
// nested bodies!
auto parent_xform = body_xforms_[model_->body_parentid[body_id]];
pxr::SdfPath xform_op_path =
CreateAttributeSpec(data_, body_path, kTokens->xformOpTransform,
pxr::SdfValueTypeNames->Matrix4d);
// mjModel will have all frames already accounted for so no need to worry
// about them here.
body_xforms_[body_id] =
MujocoPosQuatToTransform(&model_->body_pos[body_id * 3],
&model_->body_quat[body_id * 4]) *
parent_xform;
WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Matrix4d,
kTokens->xformOpTransform, body_xforms_[body_id]);
auto body_xform = MujocoPosQuatToTransform(&model_->body_pos[body_id * 3],
&model_->body_quat[body_id * 4]);
SetAttributeDefault(data_, xform_op_path, body_xform);
// Create XformOpOrder attribute for body transform order.
// For us this is simply the transform we authored above.