Write new Mujoco-specific JointAPI data to USD.

PiperOrigin-RevId: 778035661
Change-Id: I945aa38fc58973023d78ab0ab5fad7783d23e1ef

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

@@ -557,13 +557,13 @@ class ModelWriter {
                           MjcPhysicsTokens->mjcOptionMagnetic, magnetic);
 
     pxr::VtArray<double> o_solref(spec_->option.o_solref,
-                                  spec_->option.o_solref + 2);
+                                  spec_->option.o_solref + mjNREF);
     WriteUniformAttribute(physics_scene_path,
                           pxr::SdfValueTypeNames->DoubleArray,
                           MjcPhysicsTokens->mjcOptionO_solref, o_solref);
 
     pxr::VtArray<double> o_solimp(spec_->option.o_solimp,
-                                  spec_->option.o_solimp + 5);
+                                  spec_->option.o_solimp + mjNIMP);
     WriteUniformAttribute(physics_scene_path,
                           pxr::SdfValueTypeNames->DoubleArray,
                           MjcPhysicsTokens->mjcOptionO_solimp, o_solimp);
@@ -1716,6 +1716,80 @@ class ModelWriter {
               upper_limit);
         }
       }
+
+      // Finally write the mjcPhysicsJointAPI attributes.
+      ApplyApiSchema(data_, joint_path, MjcPhysicsTokens->PhysicsJointsAPI);
+
+      WriteUniformAttribute(
+          joint_path, pxr::SdfValueTypeNames->DoubleArray,
+          MjcPhysicsTokens->mjcSpringdamper,
+          pxr::VtArray<double>(joint->springdamper, joint->springdamper + 2));
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->DoubleArray,
+                            MjcPhysicsTokens->mjcSolreflimit,
+                            pxr::VtArray<double>(joint->solref_limit,
+                                                 joint->solref_limit + mjNREF));
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->DoubleArray,
+                            MjcPhysicsTokens->mjcSolimplimit,
+                            pxr::VtArray<double>(joint->solimp_limit,
+                                                 joint->solimp_limit + mjNIMP));
+
+      WriteUniformAttribute(
+          joint_path, pxr::SdfValueTypeNames->DoubleArray,
+          MjcPhysicsTokens->mjcSolreffriction,
+          pxr::VtArray<double>(joint->solref_friction,
+                               joint->solref_friction + mjNREF));
+
+      WriteUniformAttribute(
+          joint_path, pxr::SdfValueTypeNames->DoubleArray,
+          MjcPhysicsTokens->mjcSolimpfriction,
+          pxr::VtArray<double>(joint->solimp_friction,
+                               joint->solimp_friction + mjNIMP));
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Double,
+                            MjcPhysicsTokens->mjcStiffness, joint->stiffness);
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Double,
+                            MjcPhysicsTokens->mjcActuatorfrcrangeMin,
+                            joint->actfrcrange[0]);
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Double,
+                            MjcPhysicsTokens->mjcActuatorfrcrangeMax,
+                            joint->actfrcrange[1]);
+
+      pxr::TfToken actuatorfrclimited_token = MjcPhysicsTokens->auto_;
+      if (joint->actfrclimited == mjLIMITED_TRUE) {
+        actuatorfrclimited_token = MjcPhysicsTokens->true_;
+      } else if (joint->actfrclimited == mjLIMITED_FALSE) {
+        actuatorfrclimited_token = MjcPhysicsTokens->false_;
+      }
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Token,
+                            MjcPhysicsTokens->mjcActuatorfrclimited,
+                            actuatorfrclimited_token);
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Bool,
+                            MjcPhysicsTokens->mjcActuatorgravcomp,
+                            static_cast<bool>(joint->actgravcomp));
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Double,
+                            MjcPhysicsTokens->mjcMargin, joint->margin);
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Double,
+                            MjcPhysicsTokens->mjcRef, joint->ref);
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Double,
+                            MjcPhysicsTokens->mjcSpringref, joint->springref);
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Double,
+                            MjcPhysicsTokens->mjcArmature, joint->armature);
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Double,
+                            MjcPhysicsTokens->mjcDamping, joint->damping);
+
+      WriteUniformAttribute(joint_path, pxr::SdfValueTypeNames->Double,
+                            MjcPhysicsTokens->mjcFrictionloss,
+                            joint->frictionloss);
     }
     if (joint_id >= 0) {
       joint_paths_[joint_id] = joint_path;
@@ -1824,8 +1898,7 @@ class ModelWriter {
     // 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 =
+      int parent_parent_id = mjs_getId(mjs_getParent(parent->element)->element);
-          mjs_getId(mjs_getParent(parent->element)->element);
       if (parent_parent_id == kWorldIndex) {
         ApplyApiSchema(data_, parent_path,
                        pxr::UsdPhysicsTokens->PhysicsArticulationRootAPI);

test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc

@@ -19,6 +19,7 @@
 #include <gtest/gtest.h>
 #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/meshCollisionAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/sceneAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/siteAPI.h>
@@ -1169,13 +1170,13 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsRigidBody) {
   // test_body_3 is a child of the world but has no children so should not be
   // an articulation root.
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/physics_test/test_body_3",
-                          pxr::UsdPhysicsArticulationRootAPI);
+                              pxr::UsdPhysicsArticulationRootAPI);
 
   // Articulation root is not applied to other bodies or world body.
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/physics_test",
-                          pxr::UsdPhysicsArticulationRootAPI);
+                              pxr::UsdPhysicsArticulationRootAPI);
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/physics_test/test_body/test_body_2",
-                          pxr::UsdPhysicsArticulationRootAPI);
+                              pxr::UsdPhysicsArticulationRootAPI);
 
   // Geoms should not have RigidBodyAPI applied either.
   EXPECT_PRIM_API_NOT_APPLIED(stage, "/physics_test/test_body/test_geom",
@@ -1589,6 +1590,69 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsSliderCrankActuator) {
   ExpectAttributeEqual(stage, "/test/body/crank.mjc:crankLength", 1.23);
 }
 
+TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsJointAPI) {
+  static constexpr char xml[] = R"(
+  <mujoco model="test">
+    <worldbody>
+      <body name="parent">
+        <body name="child">
+          <joint name="my_joint" type="hinge"
+            springdamper="1 2"
+            solreflimit="0.1 0.2"
+            solimplimit="0.3 0.4 0.5 0.6 0.7"
+            solreffriction="0.8 0.9"
+            solimpfriction="1.0 1.1 1.2 1.3 1.4"
+            stiffness="1.5"
+            actuatorfrcrange="-1.6 1.7"
+            actuatorfrclimited="true"
+            actuatorgravcomp="true"
+            margin="1.8"
+            ref="1.9"
+            springref="2.0"
+            armature="2.1"
+            damping="2.2"
+            frictionloss="2.3"
+          />
+          <geom type="sphere" size="1"/>
+        </body>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  auto stage = OpenStageWithPhysics(xml);
+
+  const SdfPath joint_path("/test/parent/child/my_joint");
+  EXPECT_PRIM_API_APPLIED(stage, joint_path, pxr::MjcPhysicsJointAPI);
+
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:springdamper",
+                       pxr::VtArray<double>({1, 2}));
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:solreflimit",
+                       pxr::VtArray<double>({0.1, 0.2}));
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:solimplimit",
+                       pxr::VtArray<double>({0.3, 0.4, 0.5, 0.6, 0.7}));
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:solreffriction",
+                       pxr::VtArray<double>({0.8, 0.9}));
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:solimpfriction",
+                       pxr::VtArray<double>({1.0, 1.1, 1.2, 1.3, 1.4}));
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:stiffness", 1.5);
+  ExpectAttributeEqual(
+      stage, "/test/parent/child/my_joint.mjc:actuatorfrcrange:min", -1.6);
+  ExpectAttributeEqual(
+      stage, "/test/parent/child/my_joint.mjc:actuatorfrcrange:max", 1.7);
+  ExpectAttributeEqual(stage,
+                       "/test/parent/child/my_joint.mjc:actuatorfrclimited",
+                       MjcPhysicsTokens->true_);
+  ExpectAttributeEqual(
+      stage, "/test/parent/child/my_joint.mjc:actuatorgravcomp", true);
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:margin", 1.8);
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:ref", 1.9);
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:springref", 2.0);
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:armature", 2.1);
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:damping", 2.2);
+  ExpectAttributeEqual(stage, "/test/parent/child/my_joint.mjc:frictionloss",
+                       2.3);
+}
+
 TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsFloatingAndFixedBaseBody) {
   static constexpr char kXml[] = R"(
     <mujoco model="test">