From 7d7648f7ed4b10a7a11e30f39b14d55cb0c47631 Mon Sep 17 00:00:00 2001 From: Sam Haves Date: Mon, 3 Nov 2025 13:28:36 -0800 Subject: [PATCH] Update UsdWriter to write any parts of MjSpec that are not from USD to the EditTarget. This enables two workflows: 1) Users that parsed some other format into MjSpec, or built it programmatically can export animated USD easily with as minimal as possible interaction with USD. 2) Creating a spec from USD, then modifying the spec, should not cause a duplication of data when writing to USD. PiperOrigin-RevId: 827627520 Change-Id: I02c8c14a18569547710e54f3704d52ac0bf70ce6 --- include/mujoco/experimental/usd/writer.h | 1 + .../usd/plugins/mjcf/mujoco_to_usd.cc | 89 +++++++++++++------ .../usd/plugins/mjcf/mujoco_to_usd.h | 4 +- src/experimental/usd/usd_to_mjspec.cc | 20 +++++ src/experimental/usd/writer.cc | 11 +++ 5 files changed, 96 insertions(+), 29 deletions(-) diff --git a/include/mujoco/experimental/usd/writer.h b/include/mujoco/experimental/usd/writer.h index f3d191d1..b866deaf 100644 --- a/include/mujoco/experimental/usd/writer.h +++ b/include/mujoco/experimental/usd/writer.h @@ -38,6 +38,7 @@ class Writer { void AddSink(PoseSinkFn sink_fn); private: + void WriteSpecsWithoutUSDOriginToLayer(); void BuildMjUsdMapping(); pxr::UsdStageRefPtr stage_; diff --git a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc index 8c268618..456394a7 100644 --- a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc +++ b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc @@ -23,6 +23,7 @@ #include #include +#include #include #include #include "mjcf/utils.h" @@ -150,13 +151,12 @@ using mujoco::usd::CreateClassSpec; 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::SetPrimPurpose; -pxr::GfMatrix4d MujocoPosQuatToTransform(double *pos, double *quat) { +pxr::GfMatrix4d MujocoPosQuatToTransform(mjtNum *pos, mjtNum *quat) { pxr::GfQuatd quaternion = pxr::GfQuatd::GetIdentity(); quaternion.SetReal(quat[0]); quaternion.SetImaginary(quat[1], quat[2], quat[3]); @@ -174,14 +174,26 @@ pxr::GfMatrix4d MujocoPosQuatToTransform(double *pos, double *quat) { class ModelWriter { public: - ModelWriter(mjSpec *spec, mjModel *model, pxr::SdfLayerRefPtr layer) - : spec_(spec), model_(model), layer_(layer), class_path_("/Bad_Path") { + ModelWriter(mjSpec* spec, mjModel* model, pxr::SdfLayerRefPtr layer, + bool skip_elems_from_usd) + : spec_(spec), + model_(model), + layer_(layer), + skip_elems_from_usd_(skip_elems_from_usd), + class_path_("/Bad_Path") { body_paths_ = std::vector(model->nbody); site_paths_ = std::vector(model->nsite); joint_paths_ = std::vector(model->njnt); } ~ModelWriter() { mj_deleteModel(model_); } + bool ShouldWrite(mjsElement* element) { + // If we've been told to skip over spec elements that originated from USD, + // check the user value and if it has a primpath then it came from decoding + // a USD stage into an MjSpec. + return !skip_elems_from_usd_ || mujoco::usd::GetUsdPrimPathUserValue(element).IsEmpty(); + } + void Write() { // Create top level class holder. class_path_ = CreateClassSpec(layer_, pxr::SdfPath::AbsoluteRootPath(), @@ -214,6 +226,7 @@ class ModelWriter { mjModel *model_; pxr::SdfLayerRefPtr layer_; + bool skip_elems_from_usd_ = false; // Path to top level class spec that all classes should be children of. pxr::SdfPath class_path_; // Mapping from Mujoco body id to SdfPath. @@ -343,7 +356,6 @@ class ModelWriter { } auto existing_order = - prim_spec->GetField(pxr::UsdGeomTokens->xformOpOrder).UncheckedGet>(); prim_spec->GetField(pxr::UsdGeomTokens->xformOpOrder).Get>(); pxr::VtArray new_order(order.size() + existing_order.size()); @@ -746,7 +758,9 @@ class ModelWriter { mjsMesh *mesh = mjs_asMesh(mjs_firstElement(spec_, mjOBJ_MESH)); while (mesh) { - WriteMesh(mesh, scope->GetPath()); + if (ShouldWrite(mesh->element)) { + WriteMesh(mesh, scope->GetPath()); + } mesh = mjs_asMesh(mjs_nextElement(spec_, mesh->element)); } } @@ -1161,10 +1175,12 @@ class ModelWriter { std::unordered_map> keyframes_map; mjsKey *keyframe = mjs_asKey(mjs_firstElement(spec_, mjOBJ_KEY)); while (keyframe) { - std::string keyframe_name = mjs_getName(keyframe->element)->empty() - ? kTokens->keyframe - : *mjs_getName(keyframe->element); - keyframes_map[keyframe_name].push_back(keyframe); + if (ShouldWrite(keyframe->element)) { + std::string keyframe_name = mjs_getName(keyframe->element)->empty() + ? kTokens->keyframe + : *mjs_getName(keyframe->element); + keyframes_map[keyframe_name].push_back(keyframe); + } keyframe = mjs_asKey(mjs_nextElement(spec_, keyframe->element)); } @@ -1338,7 +1354,9 @@ class ModelWriter { mjsActuator *actuator = mjs_asActuator(mjs_firstElement(spec_, mjOBJ_ACTUATOR)); while (actuator) { - WriteActuator(actuator, scope->GetPath()); + if (ShouldWrite(actuator->element)) { + WriteActuator(actuator, scope->GetPath()); + } actuator = mjs_asActuator(mjs_nextElement(spec_, actuator->element)); } } @@ -1800,7 +1818,9 @@ class ModelWriter { void WriteSites(mjsBody *body) { mjsSite *site = mjs_asSite(mjs_firstChild(body, mjOBJ_SITE, false)); while (site) { - WriteSite(site, body); + if (ShouldWrite(site->element)) { + WriteSite(site, body); + } site = mjs_asSite(mjs_nextChild(body, site->element, false)); } } @@ -1808,7 +1828,9 @@ class ModelWriter { void WriteGeoms(mjsBody *body) { mjsGeom *geom = mjs_asGeom(mjs_firstChild(body, mjOBJ_GEOM, false)); while (geom) { - WriteGeom(geom, body); + if (ShouldWrite(geom->element)) { + WriteGeom(geom, body); + } geom = mjs_asGeom(mjs_nextChild(body, geom->element, false)); } } @@ -1824,7 +1846,9 @@ class ModelWriter { // WriteJoint properly handles the case where the parent is the worldbody. WriteJoint(nullptr, body); } else { - WriteJoint(joint, body); + if (ShouldWrite(joint->element)) { + WriteJoint(joint, body); + } if (mjs_asJoint(mjs_nextChild(body, joint->element, false))) { TF_WARN( "Multiple joints found for body %d. Only writing the first one.", @@ -2153,7 +2177,9 @@ class ModelWriter { void WriteCameras(mjsBody *body) { mjsCamera *cam = mjs_asCamera(mjs_firstChild(body, mjOBJ_CAMERA, false)); while (cam) { - WriteCamera(cam, body); + if (ShouldWrite(cam->element)) { + WriteCamera(cam, body); + } cam = mjs_asCamera(mjs_nextChild(body, cam->element, false)); } } @@ -2178,7 +2204,9 @@ class ModelWriter { void WriteLights(mjsBody *body) { mjsLight *light = mjs_asLight(mjs_firstChild(body, mjOBJ_LIGHT, false)); while (light) { - WriteLight(light, body); + if (ShouldWrite(light->element)) { + WriteLight(light, body); + } light = mjs_asLight(mjs_nextChild(body, light->element, false)); } } @@ -2286,22 +2314,27 @@ class ModelWriter { body_spec->SetField(pxr::SdfFieldKeys->CustomData, customData); body_paths_[body_id] = body_spec->GetPath(); + mujoco::usd::SetUsdPrimPathUserValue(body->element, body_spec->GetPath()); } void WriteBodies() { mjsBody *body = mjs_asBody(mjs_firstElement(spec_, mjOBJ_BODY)); while (body) { - // Only write a rigidbody if we are not the world body. - // We fall through since the world body might have static - // geom children. - if (mjs_getId(body->element) != kWorldIndex) { - WriteBody(body); + // If this body originally came from USD, don't try and write it + // again. + if (ShouldWrite(body->element)) { + // Only write a rigidbody if we are not the world body. + // We fall through since the world body might have static + // geom children. + if (mjs_getId(body->element) != kWorldIndex) { + WriteBody(body); + } + WriteSites(body); + WriteGeoms(body); + WriteJoints(body); + WriteCameras(body); + WriteLights(body); } - WriteSites(body); - WriteGeoms(body); - WriteJoints(body); - WriteCameras(body); - WriteLights(body); body = mjs_asBody(mjs_nextElement(spec_, body->element)); } } @@ -2323,14 +2356,14 @@ class ModelWriter { namespace mujoco { namespace usd { -bool WriteSpecToData(mjSpec *spec, pxr::SdfLayerRefPtr layer) { +bool WriteSpecToData(mjSpec *spec, pxr::SdfLayerRefPtr layer, bool skip_elems_from_usd) { mjModel *model = mj_compile(spec, nullptr); if (model == nullptr) { TF_ERROR(MujocoCompilationError, "%s", mjs_getError(spec)); return false; } - ModelWriter(spec, model, layer).Write(); + ModelWriter(spec, model, layer, skip_elems_from_usd).Write(); return true; } diff --git a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h index 6c4205be..9256bba7 100644 --- a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h +++ b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h @@ -25,7 +25,9 @@ namespace usd { // Args: // spec: mjSpec built programmatically or via parsed XML. // layer: SdfLayerRefPtr that will be written to. -bool WriteSpecToData(mjSpec* spec, pxr::SdfLayerRefPtr layer); +// skip_elems_from_usd: If true, skips over elements in the spec that have a usd_prim_path +// custom attribute set. +bool WriteSpecToData(mjSpec* spec, pxr::SdfLayerRefPtr layer, bool skip_elems_from_usd = false); } // namespace usd } // namespace mujoco diff --git a/src/experimental/usd/usd_to_mjspec.cc b/src/experimental/usd/usd_to_mjspec.cc index 4cbc1823..8ae885c5 100644 --- a/src/experimental/usd/usd_to_mjspec.cc +++ b/src/experimental/usd/usd_to_mjspec.cc @@ -278,6 +278,8 @@ mjsMesh* ParseUsdMesh(mjSpec* spec, const pxr::UsdPrim& prim, mjsGeom* geom, } mjsMesh* mesh = mjs_addMesh(spec, nullptr); + mujoco::usd::SetUsdPrimPathUserValue(mesh->element, prim.GetPath()); + geom->type = mjGEOM_MESH; pxr::UsdGeomMesh usd_mesh(prim); std::vector uservert; @@ -928,6 +930,8 @@ void ParseMjcPhysicsTendon(mjSpec* spec, mjsTendon* mj_tendon = mjs_addTendon(spec, nullptr); mjs_setName(mj_tendon->element, prim.GetPath().GetAsString().c_str()); + mujoco::usd::SetUsdPrimPathUserValue(mj_tendon->element, prim.GetPath()); + pxr::TfToken type; tendon.GetTypeAttr().Get(&type); @@ -1221,6 +1225,8 @@ void ParseMjcPhysicsActuator(mjSpec* spec, mjsActuator* mj_act = mjs_addActuator(spec, nullptr); mjs_setName(mj_act->element, prim.GetPath().GetAsString().c_str()); + mujoco::usd::SetUsdPrimPathUserValue(mj_act->element, prim.GetPath()); + auto group_attr = tran.GetGroupAttr(); if (group_attr.HasAuthoredValue()) { group_attr.Get(&mj_act->group); @@ -1664,6 +1670,8 @@ void ParseUsdGeomGprim(mjSpec* spec, const pxr::UsdPrim& gprim, geom->contype = 0; geom->conaffinity = 0; + mujoco::usd::SetUsdPrimPathUserValue(geom->element, gprim.GetPath()); + ParseDisplayColorAndOpacity(gprim, geom); SetLocalPoseFromPrim(gprim, body_prim, geom, caches.xform_cache); if (!MaybeParseGeomPrimitive(gprim, geom, caches.xform_cache)) { @@ -1719,6 +1727,8 @@ void ParseUsdPhysicsCollider(mjSpec* spec, geom->contype = 1; geom->conaffinity = 1; + mujoco::usd::SetUsdPrimPathUserValue(geom->element, prim.GetPath()); + if (prim.HasAPI()) { ParseMjcPhysicsCollisionAPI(geom, pxr::MjcPhysicsCollisionAPI(prim)); } @@ -1799,6 +1809,8 @@ void ParseUsdPhysicsJoint(mjSpec* spec, const pxr::UsdPrim& prim, mjsBody* body, mj_joint->type = type; mjs_setName(mj_joint->element, prim.GetPath().GetAsString().c_str()); + mujoco::usd::SetUsdPrimPathUserValue(mj_joint->element, prim.GetPath()); + if (prim.IsA()) { pxr::UsdPhysicsRevoluteJoint revolute(prim); TfToken axis; @@ -1888,6 +1900,8 @@ void ParseMjcPhysicsSite(mjSpec* spec, const pxr::MjcPhysicsSiteAPI& site_api, site_api.GetPrim().GetPath().GetAsString().c_str()); SetLocalPoseFromPrim(site_api.GetPrim(), parent_prim, site, xform_cache); + mujoco::usd::SetUsdPrimPathUserValue(site->element, prim.GetPath()); + auto group_attr = site_api.GetGroupAttr(); if (group_attr.HasAuthoredValue()) { group_attr.Get(&site->group); @@ -1933,6 +1947,9 @@ void ParseMjcPhysicsKeyframe(mjSpec* spec, if (n_time_samples == 0) { // If no time samples, we create a single keyframe. mjsKey* key = mjs_addKey(spec); + + mujoco::usd::SetUsdPrimPathUserValue(key->element, prim.GetPath()); + mjs_setName(key->element, prim.GetName().GetString().c_str()); setKeyframeData(key, qpos_attr, &key->qpos); setKeyframeData(key, qvel_attr, &key->qvel); @@ -1947,6 +1964,9 @@ void ParseMjcPhysicsKeyframe(mjSpec* spec, int keyframe_id = 0; for (double time : times) { mjsKey* key = mjs_addKey(spec); + + mujoco::usd::SetUsdPrimPathUserValue(key->element, prim.GetPath()); + std::string key_name = prim.GetName().GetString() + "_" + std::to_string(keyframe_id++); mjs_setName(key->element, key_name.c_str()); diff --git a/src/experimental/usd/writer.cc b/src/experimental/usd/writer.cc index 55f196fe..52a9122c 100644 --- a/src/experimental/usd/writer.cc +++ b/src/experimental/usd/writer.cc @@ -18,9 +18,11 @@ #include #include +#include "experimental/usd/plugins/mjcf/mujoco_to_usd.h" #include #include #include +#include namespace mujoco { namespace usd { @@ -30,11 +32,20 @@ namespace usd { // API. Writer::Writer(pxr::UsdStageRefPtr stage, mjSpec* spec, mjModel_* model) : stage_(stage), spec_(spec), model_(model) { + WriteSpecsWithoutUSDOriginToLayer(); BuildMjUsdMapping(); } Writer::~Writer() = default; +void Writer::WriteSpecsWithoutUSDOriginToLayer() { + // The user may have constructed the spec from sources other than USD. + // In this case, we will write write all those spec elements to a new layer in + // the stage. This allows us to keep the same writing API for a simulation regardless + // of how a spec was produced. + mujoco::usd::WriteSpecToData(spec_, stage_->GetEditTarget().GetLayer(), true); +} + void Writer::BuildMjUsdMapping() { body_id_to_path_.assign(model_->nbody, pxr::SdfPath());