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

@@ -38,6 +38,7 @@ class Writer {
   void AddSink(PoseSinkFn sink_fn);
 
  private:
+  void WriteSpecsWithoutUSDOriginToLayer();
   void BuildMjUsdMapping();
 
   pxr::UsdStageRefPtr stage_;

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

@@ -23,6 +23,7 @@
 #include <vector>
 
 #include <mujoco/experimental/usd/mjcPhysics/tokens.h>
+#include <mujoco/experimental/usd/utils.h>
 #include <mujoco/mjspec.h>
 #include <mujoco/mujoco.h>
 #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<pxr::SdfPath>(model->nbody);
     site_paths_ = std::vector<pxr::SdfPath>(model->nsite);
     joint_paths_ = std::vector<pxr::SdfPath>(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<pxr::VtArray<pxr::TfToken>>();
         prim_spec->GetField(pxr::UsdGeomTokens->xformOpOrder).Get<pxr::VtArray<pxr::TfToken>>();
 
     pxr::VtArray<pxr::TfToken> 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<std::string, std::vector<mjsKey *>> 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;
 }

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
 

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<float> 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<pxr::MjcPhysicsCollisionAPI>()) {
     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>()) {
     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());

src/experimental/usd/writer.cc

@@ -18,9 +18,11 @@
 #include <vector>
 
 #include <mujoco/experimental/usd/utils.h>
+#include "experimental/usd/plugins/mjcf/mujoco_to_usd.h"
 #include <mujoco/mujoco.h>
 #include <pxr/usd/sdf/path.h>
 #include <pxr/usd/usd/common.h>
+#include <pxr/usd/usd/stage.h>
 
 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());