Add equalityJointAPI to mjcPhysics.

PiperOrigin-RevId: 860081702
Change-Id: Ifd961dcb2a8e0f5d885e9e9dadc0153b03cb5f4d

include/mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h

@@ -0,0 +1,198 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MJCPHYSICS_GENERATED_EQUALITYJOINTAPI_H
+#define MJCPHYSICS_GENERATED_EQUALITYJOINTAPI_H
+
+/// \file mjcPhysics/equalityJointAPI.h
+
+#include <mujoco/experimental/usd/mjcPhysics/api.h>
+#include <mujoco/experimental/usd/mjcPhysics/tokens.h>
+#include <pxr/base/gf/matrix4d.h>
+#include <pxr/base/gf/vec3d.h>
+#include <pxr/base/gf/vec3f.h>
+#include <pxr/base/tf/token.h>
+#include <pxr/base/tf/type.h>
+#include <pxr/base/vt/value.h>
+#include <pxr/pxr.h>
+#include <pxr/usd/usd/apiSchemaBase.h>
+#include <pxr/usd/usd/prim.h>
+#include <pxr/usd/usd/stage.h>
+
+PXR_NAMESPACE_OPEN_SCOPE
+
+class SdfAssetPath;
+
+// -------------------------------------------------------------------------- //
+// MJCEQUALITYJOINTAPI                                                        //
+// -------------------------------------------------------------------------- //
+
+/// \class MjcPhysicsEqualityJointAPI
+///
+/// API providing extension attributes to represent equality/joint constraints.
+/// This API is applied to a joint prim which acts as the constrained joint
+/// (joint1 in MuJoCo terminology). The target relationship points to another
+/// joint prim which is the reference joint (joint2 in MuJoCo terminology). The
+/// constrained joint's position or angle is constrained to be a quartic
+/// polynomial of the reference joint's position or angle. Only scalar joint
+/// types (slide and hinge) can be used.
+///
+class MjcPhysicsEqualityJointAPI : public UsdAPISchemaBase {
+ public:
+  /// Compile time constant representing what kind of schema this class is.
+  ///
+  /// \sa UsdSchemaKind
+  static const UsdSchemaKind schemaKind = UsdSchemaKind::SingleApplyAPI;
+
+  /// Construct a MjcPhysicsEqualityJointAPI on UsdPrim \p prim .
+  /// Equivalent to MjcPhysicsEqualityJointAPI::Get(prim.GetStage(),
+  /// prim.GetPath()) for a \em valid \p prim, but will not immediately throw an
+  /// error for an invalid \p prim
+  explicit MjcPhysicsEqualityJointAPI(const UsdPrim& prim = UsdPrim())
+      : UsdAPISchemaBase(prim) {}
+
+  /// Construct a MjcPhysicsEqualityJointAPI on the prim held by \p schemaObj .
+  /// Should be preferred over MjcPhysicsEqualityJointAPI(schemaObj.GetPrim()),
+  /// as it preserves SchemaBase state.
+  explicit MjcPhysicsEqualityJointAPI(const UsdSchemaBase& schemaObj)
+      : UsdAPISchemaBase(schemaObj) {}
+
+  /// Destructor.
+  MJCPHYSICS_API
+  virtual ~MjcPhysicsEqualityJointAPI();
+
+  /// Return a vector of names of all pre-declared attributes for this schema
+  /// class and all its ancestor classes.  Does not include attributes that
+  /// may be authored by custom/extended methods of the schemas involved.
+  MJCPHYSICS_API
+  static const TfTokenVector& GetSchemaAttributeNames(
+      bool includeInherited = true);
+
+  /// Return a MjcPhysicsEqualityJointAPI holding the prim adhering to this
+  /// schema at \p path on \p stage.  If no prim exists at \p path on
+  /// \p stage, or if the prim at that path does not adhere to this schema,
+  /// return an invalid schema object.  This is shorthand for the following:
+  ///
+  /// \code
+  /// MjcPhysicsEqualityJointAPI(stage->GetPrimAtPath(path));
+  /// \endcode
+  ///
+  MJCPHYSICS_API
+  static MjcPhysicsEqualityJointAPI Get(const UsdStagePtr& stage,
+                                        const SdfPath& path);
+
+  /// Returns true if this <b>single-apply</b> API schema can be applied to
+  /// the given \p prim. If this schema can not be a applied to the prim,
+  /// this returns false and, if provided, populates \p whyNot with the
+  /// reason it can not be applied.
+  ///
+  /// Note that if CanApply returns false, that does not necessarily imply
+  /// that calling Apply will fail. Callers are expected to call CanApply
+  /// before calling Apply if they want to ensure that it is valid to
+  /// apply a schema.
+  ///
+  /// \sa UsdPrim::GetAppliedSchemas()
+  /// \sa UsdPrim::HasAPI()
+  /// \sa UsdPrim::CanApplyAPI()
+  /// \sa UsdPrim::ApplyAPI()
+  /// \sa UsdPrim::RemoveAPI()
+  ///
+  MJCPHYSICS_API
+  static bool CanApply(const UsdPrim& prim, std::string* whyNot = nullptr);
+
+  /// Applies this <b>single-apply</b> API schema to the given \p prim.
+  /// This information is stored by adding "MjcEqualityJointAPI" to the
+  /// token-valued, listOp metadata \em apiSchemas on the prim.
+  ///
+  /// \return A valid MjcPhysicsEqualityJointAPI object is returned upon
+  /// success. An invalid (or empty) MjcPhysicsEqualityJointAPI object is
+  /// returned upon failure. See \ref UsdPrim::ApplyAPI() for conditions
+  /// resulting in failure.
+  ///
+  /// \sa UsdPrim::GetAppliedSchemas()
+  /// \sa UsdPrim::HasAPI()
+  /// \sa UsdPrim::CanApplyAPI()
+  /// \sa UsdPrim::ApplyAPI()
+  /// \sa UsdPrim::RemoveAPI()
+  ///
+  MJCPHYSICS_API
+  static MjcPhysicsEqualityJointAPI Apply(const UsdPrim& prim);
+
+ protected:
+  /// Returns the kind of schema this class belongs to.
+  ///
+  /// \sa UsdSchemaKind
+  MJCPHYSICS_API
+  UsdSchemaKind _GetSchemaKind() const override;
+
+ private:
+  // needs to invoke _GetStaticTfType.
+  friend class UsdSchemaRegistry;
+  MJCPHYSICS_API
+  static const TfType& _GetStaticTfType();
+
+  static bool _IsTypedSchema();
+
+  // override SchemaBase virtuals.
+  MJCPHYSICS_API
+  const TfType& _GetTfType() const override;
+
+ public:
+  // --------------------------------------------------------------------- //
+  // POLYCOEF
+  // --------------------------------------------------------------------- //
+  /// Coefficients a0 through a4 of the quartic polynomial. If the joint values
+  /// of the constrained joint and reference joint are respectively y and x, and
+  /// their reference positions (corresponding to the joint values in the
+  /// initial model configuration) are y0 and x0, the constraint is: y = y0 + a0
+  /// + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4. Omitting the target
+  /// joint (joint2) is equivalent to setting x = x0, in which case the
+  /// constraint is y = y0 + a0. The default [0, 1, 0, 0, 0] creates a simple
+  /// 1:1 mimic constraint where y tracks x with the same offset from their
+  /// references.
+  ///
+  /// | ||
+  /// | -- | -- |
+  /// | Declaration | `uniform double[] mjc:polycoef = [0, 1, 0, 0, 0]` |
+  /// | C++ Type | VtArray<double> |
+  /// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->DoubleArray |
+  /// | \ref SdfVariability "Variability" | SdfVariabilityUniform |
+  MJCPHYSICS_API
+  UsdAttribute GetPolycoefAttr() const;
+
+  /// See GetPolycoefAttr(), and also
+  /// \ref Usd_Create_Or_Get_Property for when to use Get vs Create.
+  /// If specified, author \p defaultValue as the attribute's default,
+  /// sparsely (when it makes sense to do so) if \p writeSparsely is \c true -
+  /// the default for \p writeSparsely is \c false.
+  MJCPHYSICS_API
+  UsdAttribute CreatePolycoefAttr(VtValue const& defaultValue = VtValue(),
+                                  bool writeSparsely = false) const;
+
+ public:
+  // ===================================================================== //
+  // Feel free to add custom code below this line, it will be preserved by
+  // the code generator.
+  //
+  // Just remember to:
+  //  - Close the class declaration with };
+  //  - Close the namespace with PXR_NAMESPACE_CLOSE_SCOPE
+  //  - Close the include guard with #endif
+  // ===================================================================== //
+  // --(BEGIN CUSTOM CODE)--
+};
+
+PXR_NAMESPACE_CLOSE_SCOPE
+
+#endif

include/mujoco/experimental/usd/mjcPhysics/tokens.h

@@ -580,6 +580,10 @@ struct MjcPhysicsTokensType {
   ///
   /// MjcPhysicsTendon
   const TfToken mjcPathSegments;
+  /// \brief "mjc:polycoef"
+  ///
+  /// MjcPhysicsEqualityJointAPI
+  const TfToken mjcPolycoef;
   /// \brief "mjc:priority"
   ///
   /// MjcPhysicsCollisionAPI
@@ -772,6 +776,10 @@ struct MjcPhysicsTokensType {
   ///
   /// Schema identifier and family for MjcPhysicsEqualityAPI
   const TfToken MjcEqualityAPI;
+  /// \brief "MjcEqualityJointAPI"
+  ///
+  /// Schema identifier and family for MjcPhysicsEqualityJointAPI
+  const TfToken MjcEqualityJointAPI;
   /// \brief "MjcEqualityWeldAPI"
   ///
   /// Schema identifier and family for MjcPhysicsEqualityWeldAPI

plugin/usd_decoder/kinematic_tree.cc

@@ -20,6 +20,7 @@
 #include <vector>
 
 #include <mujoco/experimental/usd/mjcPhysics/actuator.h>
+#include <mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/keyframe.h>
 #include <mujoco/experimental/usd/mjcPhysics/siteAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/tendon.h>
@@ -227,6 +228,12 @@ std::unique_ptr<Node> BuildKinematicTree(const pxr::UsdStageRefPtr stage) {
     // Now that we know all the bodies, we can assign joints to respective
     // nodes.
     extraction.nodes[to_idx]->joints.push_back(joint.GetPath());
+
+    // If the joint has MjcPhysicsEqualityJointAPI, also add it to constraints
+    // so that ParseConstraint is called to create the equality constraint.
+    if (joint.GetPrim().HasAPI<pxr::MjcPhysicsEqualityJointAPI>()) {
+      extraction.nodes[to_idx]->constraints.push_back(joint.GetPath());
+    }
   }
 
   // The world body is represented by an empty SdfPath.

plugin/usd_decoder/usd_decoder.cc

@@ -26,6 +26,7 @@
 #include <mujoco/experimental/usd/mjcPhysics/actuator.h>
 #include <mujoco/experimental/usd/mjcPhysics/collisionAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/equalityAPI.h>
+#include <mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/equalityWeldAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/imageableAPI.h>
 #include <mujoco/experimental/usd/mjcPhysics/jointAPI.h>
@@ -1826,97 +1827,157 @@ void ParseUsdPhysicsCollider(mjSpec* spec,
 
 void ParseConstraint(mjSpec* spec, const pxr::UsdPrim& prim, mjsBody* body,
                           pxr::UsdGeomXformCache& xform_cache) {
-  if (!prim.IsA<pxr::UsdPhysicsFixedJoint>()) {
-    mju_warning("Constraint %s is not a fixed joint, skipping.",
-                prim.GetPath().GetAsString().c_str());
-    return;
-  }
-  pxr::UsdPhysicsJoint joint(prim);
-  // A fixed joint means the bodies are welded.
-  pxr::UsdRelationship body0_rel = joint.GetBody0Rel();
-  pxr::UsdRelationship body1_rel = joint.GetBody1Rel();
-  pxr::SdfPathVector targets0, targets1;
-  body0_rel.GetTargets(&targets0);
-  body1_rel.GetTargets(&targets1);
+  if (prim.HasAPI<pxr::MjcPhysicsEqualityJointAPI>()) {
+    // Handle MjcPhysicsEqualityJointAPI on revolute/prismatic joints.
+    pxr::MjcPhysicsEqualityJointAPI eq_joint_api(prim);
+    mjsEquality* eq = mjs_addEquality(spec, nullptr);
+    eq->type = mjEQ_JOINT;
+    mjs_setName(eq->element, prim.GetPath().GetAsString().c_str());
+    SetUsdPrimPathUserValue(eq->element, prim.GetPath());
 
-  pxr::SdfPath body0_path;
-  if (!targets0.empty()) body0_path = targets0[0];
-  pxr::SdfPath body1_path;
-  if (!targets1.empty()) body1_path = targets1[0];
+    // The prim this API is applied to is the constrained joint (joint1).
+    eq->objtype = mjOBJ_JOINT;
+    mjs_setString(eq->name1, prim.GetPath().GetAsString().c_str());
 
-  auto stage = prim.GetStage();
+    // Get the target joint (joint2) from the MjcEqualityAPI target
+    // relationship.
+    pxr::MjcPhysicsEqualityAPI equality_api(prim);
+    pxr::UsdRelationship target_rel = equality_api.GetMjcTargetRel();
+    pxr::SdfPathVector targets;
+    target_rel.GetTargets(&targets);
+    if (!targets.empty()) {
+      mjs_setString(eq->name2, targets[0].GetAsString().c_str());
+    }
+    // If no target, name2 remains empty, meaning joint1 is fixed to a constant.
 
-  auto body0_prim = stage->GetPrimAtPath(body0_path);
-  auto body1_prim = stage->GetPrimAtPath(body1_path);
+    // Parse polycoef attribute for the quartic polynomial coefficients.
+    auto polycoef_attr = eq_joint_api.GetPolycoefAttr();
+    if (polycoef_attr.HasAuthoredValue()) {
+      pxr::VtDoubleArray polycoef;
+      polycoef_attr.Get(&polycoef);
+      size_t num_coefs = std::min(polycoef.size(), static_cast<size_t>(5));
+      for (size_t i = 0; i < num_coefs; ++i) {
+        eq->data[i] = polycoef[i];
+      }
+    } else {
+      // Default polycoef [0, 1, 0, 0, 0] for 1:1 mimic.
+      eq->data[0] = 0;
+      eq->data[1] = 1;
+      eq->data[2] = 0;
+      eq->data[3] = 0;
+      eq->data[4] = 0;
+    }
 
-  bool body0_is_site = false;
-  if (!body0_path.IsEmpty()) {
-    body0_is_site = body0_prim.HasAPI<pxr::MjcPhysicsSiteAPI>();
-  }
-  bool body1_is_site = false;
-  if (!body1_path.IsEmpty()) {
-    body1_is_site = body1_prim.HasAPI<pxr::MjcPhysicsSiteAPI>();
-  }
+    // Parse solver parameters from MjcEqualityAPI.
+    auto solref_attr = equality_api.GetSolRefAttr();
+    if (solref_attr.HasAuthoredValue()) {
+      pxr::VtDoubleArray solref;
+      solref_attr.Get(&solref);
+      if (solref.size() == mjNREF) {
+        for (int i = 0; i < mjNREF; ++i) {
+          eq->solref[i] = solref[i];
+        }
+      }
+    }
 
-  if (body0_is_site != body1_is_site) {
-    mju_warning(
-        "Weld constraint %s has mismatch between site and body targets, "
-        "skipping",
-        prim.GetPath().GetAsString().c_str());
-    return;
-  }
+    auto solimp_attr = equality_api.GetSolImpAttr();
+    if (solimp_attr.HasAuthoredValue()) {
+      pxr::VtDoubleArray solimp;
+      solimp_attr.Get(&solimp);
+      if (solimp.size() == mjNIMP) {
+        for (int i = 0; i < mjNIMP; ++i) {
+          eq->solimp[i] = solimp[i];
+        }
+      }
+    }
+  } else if (prim.IsA<pxr::UsdPhysicsFixedJoint>()) {
+    // Handle fixed joints as weld constraints.
+    pxr::UsdPhysicsJoint joint(prim);
+    // A fixed joint means the bodies are welded.
+    pxr::UsdRelationship body0_rel = joint.GetBody0Rel();
+    pxr::UsdRelationship body1_rel = joint.GetBody1Rel();
+    pxr::SdfPathVector targets0, targets1;
+    body0_rel.GetTargets(&targets0);
+    body1_rel.GetTargets(&targets1);
 
-  mjsEquality* eq = mjs_addEquality(spec, nullptr);
-  eq->type = mjEQ_WELD;
-  mjs_setName(eq->element, prim.GetPath().GetAsString().c_str());
-  SetUsdPrimPathUserValue(eq->element, prim.GetPath());
+    pxr::SdfPath body0_path;
+    if (!targets0.empty()) body0_path = targets0[0];
+    pxr::SdfPath body1_path;
+    if (!targets1.empty()) body1_path = targets1[0];
 
-  if (body0_is_site) {
-    mjs_setString(eq->name1, body0_path.GetAsString().c_str());
-    mjs_setString(eq->name2, body1_path.GetAsString().c_str());
-    eq->objtype = mjOBJ_SITE;
-  } else {
-    mjs_setString(eq->name1, body0_path.GetAsString().c_str());
-    mjs_setString(eq->name2, body1_path.GetAsString().c_str());
-    eq->objtype = mjOBJ_BODY;
-  }
+    auto stage = prim.GetStage();
 
-  // In USD, joints have a reference frame that is shared between the two
-  // connecting bodies. This reference frame is defined relative to both
-  // bodies, in localPos/Rot 0 and 1. A fixed joint removes all degrees of
-  // freedom for the joint, ensuring the reference frame is fixed in place
-  // This means the bodies should also be fixed, but relative to the joint
-  // depending on their respective localPos/Rot.
-  // In MuJoCo fixed joint frames are not explicitly defined relative to their
-  // connecting bodies. Instead, we define a weld constraint providing the
-  // weld point (anchor) relative to body 2 and then we specify the position
-  // of body 2 relative to body 1.
+    auto body0_prim = stage->GetPrimAtPath(body0_path);
+    auto body1_prim = stage->GetPrimAtPath(body1_path);
 
-  // Here is the mapping of terms concretely:
-  // T(bodyX) = transform of bodyX relative to joint frame (localPos/Rot).
-  // anchor = localPos1 (position of the weld point relative to mjc body 2)
-  // relpose = T(body0)*T(body1)^-1
+    bool body0_is_site = false;
+    if (!body0_path.IsEmpty()) {
+      body0_is_site = body0_prim.HasAPI<pxr::MjcPhysicsSiteAPI>();
+    }
+    bool body1_is_site = false;
+    if (!body1_path.IsEmpty()) {
+      body1_is_site = body1_prim.HasAPI<pxr::MjcPhysicsSiteAPI>();
+    }
 
-  // relpose is float[7], pos(3) + quat(4)
-  // anchor is float[3], pos(3)
-  // in mjsEquality data: anchor 0-2, relpose 3-9.
+    if (body0_is_site != body1_is_site) {
+      mju_warning(
+          "Weld constraint %s has mismatch between site and body targets, "
+          "skipping",
+          prim.GetPath().GetAsString().c_str());
+      return;
+    }
 
-  auto body0_xform = xform_cache.GetLocalToWorldTransform(body0_prim);
-  auto body1_xform = xform_cache.GetLocalToWorldTransform(body1_prim);
+    mjsEquality* eq = mjs_addEquality(spec, nullptr);
+    eq->type = mjEQ_WELD;
+    mjs_setName(eq->element, prim.GetPath().GetAsString().c_str());
+    SetUsdPrimPathUserValue(eq->element, prim.GetPath());
 
-  pxr::GfVec3d body0_scale, body1_scale;
+    if (body0_is_site) {
+      mjs_setString(eq->name1, body0_path.GetAsString().c_str());
+      mjs_setString(eq->name2, body1_path.GetAsString().c_str());
+      eq->objtype = mjOBJ_SITE;
+    } else {
+      mjs_setString(eq->name1, body0_path.GetAsString().c_str());
+      mjs_setString(eq->name2, body1_path.GetAsString().c_str());
+      eq->objtype = mjOBJ_BODY;
+    }
 
-  {
-    pxr::GfMatrix4d scale_orient, rot, persp;
-    pxr::GfVec3d  translation;
-    if (!body0_xform.Factor(&scale_orient, &body0_scale, &rot,
-                  &translation, &persp)) {
+    // In USD, joints have a reference frame that is shared between the two
+    // connecting bodies. This reference frame is defined relative to both
+    // bodies, in localPos/Rot 0 and 1. A fixed joint removes all degrees of
+    // freedom for the joint, ensuring the reference frame is fixed in place
+    // This means the bodies should also be fixed, but relative to the joint
+    // depending on their respective localPos/Rot.
+    // In MuJoCo fixed joint frames are not explicitly defined relative to their
+    // connecting bodies. Instead, we define a weld constraint providing the
+    // weld point (anchor) relative to body 2 and then we specify the position
+    // of body 2 relative to body 1.
+
+    // Here is the mapping of terms concretely:
+    // T(bodyX) = transform of bodyX relative to joint frame (localPos/Rot).
+    // anchor = localPos1 (position of the weld point relative to mjc body 2)
+    // relpose = T(body0)*T(body1)^-1
+
+    // relpose is float[7], pos(3) + quat(4)
+    // anchor is float[3], pos(3)
+    // in mjsEquality data: anchor 0-2, relpose 3-9.
+
+    auto body0_xform = xform_cache.GetLocalToWorldTransform(body0_prim);
+    auto body1_xform = xform_cache.GetLocalToWorldTransform(body1_prim);
+
+    pxr::GfVec3d body0_scale, body1_scale;
+
+    {
+      pxr::GfMatrix4d scale_orient, rot, persp;
+      pxr::GfVec3d translation;
+      if (!body0_xform.Factor(&scale_orient, &body0_scale, &rot, &translation,
+                              &persp)) {
         // unable to decompose, emit warning and set scale to identity
         mju_warning(
               "Unable to decompose matrix for body 0: %s.",
               body0_path.GetAsString().c_str());
         body0_scale = pxr::GfVec3f(1, 1, 1);
-    }
+      }
 
     if (!body1_xform.Factor(&scale_orient, &body1_scale, &rot,
                   &translation, &persp)) {
@@ -1926,7 +1987,7 @@ void ParseConstraint(mjSpec* spec, const pxr::UsdPrim& prim, mjsBody* body,
               body1_path.GetAsString().c_str());
         body1_scale = pxr::GfVec3f(1, 1, 1);
     }
-  }
+    }
 
   pxr::GfVec3f localPos1;
   joint.GetLocalPos1Attr().Get(&localPos1);
@@ -2003,8 +2064,11 @@ void ParseConstraint(mjSpec* spec, const pxr::UsdPrim& prim, mjsBody* body,
       eq->data[10] = torque_scale;
     }
   }
+  } else {
+    mju_warning("Constraint %s is not a supported constraint type, skipping.",
+                prim.GetPath().GetAsString().c_str());
+  }
 }
-
 void ParseUsdPhysicsJoint(mjSpec* spec, const pxr::UsdPrim& prim, mjsBody* body,
                           pxr::UsdGeomXformCache& xform_cache) {
   // A fixed joint means the bodies are welded.

src/experimental/usd/CMakeLists.txt

@@ -128,6 +128,7 @@ target_sources(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PRIVATE
   mjcPhysics/actuator.cpp
   mjcPhysics/collisionAPI.cpp
   mjcPhysics/equalityAPI.cpp
+  mjcPhysics/equalityJointAPI.cpp
   mjcPhysics/equalityWeldAPI.cpp
   mjcPhysics/imageableAPI.cpp
   mjcPhysics/jointAPI.cpp

src/experimental/usd/mjcPhysics/equalityJointAPI.cpp

@@ -0,0 +1,126 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h>
+
+#include <pxr/usd/sdf/assetPath.h>
+#include <pxr/usd/sdf/types.h>
+#include <pxr/usd/usd/schemaRegistry.h>
+#include <pxr/usd/usd/typed.h>
+
+PXR_NAMESPACE_OPEN_SCOPE
+
+// Register the schema with the TfType system.
+TF_REGISTRY_FUNCTION(TfType) {
+  TfType::Define<MjcPhysicsEqualityJointAPI,
+                 TfType::Bases<UsdAPISchemaBase> >();
+}
+
+/* virtual */
+MjcPhysicsEqualityJointAPI::~MjcPhysicsEqualityJointAPI() {}
+
+/* static */
+MjcPhysicsEqualityJointAPI MjcPhysicsEqualityJointAPI::Get(
+    const UsdStagePtr& stage, const SdfPath& path) {
+  if (!stage) {
+    TF_CODING_ERROR("Invalid stage");
+    return MjcPhysicsEqualityJointAPI();
+  }
+  return MjcPhysicsEqualityJointAPI(stage->GetPrimAtPath(path));
+}
+
+/* virtual */
+UsdSchemaKind MjcPhysicsEqualityJointAPI::_GetSchemaKind() const {
+  return MjcPhysicsEqualityJointAPI::schemaKind;
+}
+
+/* static */
+bool MjcPhysicsEqualityJointAPI::CanApply(const UsdPrim& prim,
+                                          std::string* whyNot) {
+  return prim.CanApplyAPI<MjcPhysicsEqualityJointAPI>(whyNot);
+}
+
+/* static */
+MjcPhysicsEqualityJointAPI MjcPhysicsEqualityJointAPI::Apply(
+    const UsdPrim& prim) {
+  if (prim.ApplyAPI<MjcPhysicsEqualityJointAPI>()) {
+    return MjcPhysicsEqualityJointAPI(prim);
+  }
+  return MjcPhysicsEqualityJointAPI();
+}
+
+/* static */
+const TfType& MjcPhysicsEqualityJointAPI::_GetStaticTfType() {
+  static TfType tfType = TfType::Find<MjcPhysicsEqualityJointAPI>();
+  return tfType;
+}
+
+/* static */
+bool MjcPhysicsEqualityJointAPI::_IsTypedSchema() {
+  static bool isTyped = _GetStaticTfType().IsA<UsdTyped>();
+  return isTyped;
+}
+
+/* virtual */
+const TfType& MjcPhysicsEqualityJointAPI::_GetTfType() const {
+  return _GetStaticTfType();
+}
+
+UsdAttribute MjcPhysicsEqualityJointAPI::GetPolycoefAttr() const {
+  return GetPrim().GetAttribute(MjcPhysicsTokens->mjcPolycoef);
+}
+
+UsdAttribute MjcPhysicsEqualityJointAPI::CreatePolycoefAttr(
+    VtValue const& defaultValue, bool writeSparsely) const {
+  return UsdSchemaBase::_CreateAttr(
+      MjcPhysicsTokens->mjcPolycoef, SdfValueTypeNames->DoubleArray,
+      /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely);
+}
+
+namespace {
+static inline TfTokenVector _ConcatenateAttributeNames(
+    const TfTokenVector& left, const TfTokenVector& right) {
+  TfTokenVector result;
+  result.reserve(left.size() + right.size());
+  result.insert(result.end(), left.begin(), left.end());
+  result.insert(result.end(), right.begin(), right.end());
+  return result;
+}
+}  // namespace
+
+/*static*/
+const TfTokenVector& MjcPhysicsEqualityJointAPI::GetSchemaAttributeNames(
+    bool includeInherited) {
+  static TfTokenVector localNames = {
+      MjcPhysicsTokens->mjcPolycoef,
+  };
+  static TfTokenVector allNames = _ConcatenateAttributeNames(
+      UsdAPISchemaBase::GetSchemaAttributeNames(true), localNames);
+
+  if (includeInherited)
+    return allNames;
+  else
+    return localNames;
+}
+
+PXR_NAMESPACE_CLOSE_SCOPE
+
+// ===================================================================== //
+// Feel free to add custom code below this line. It will be preserved by
+// the code generator.
+//
+// Just remember to wrap code in the appropriate delimiters:
+// 'PXR_NAMESPACE_OPEN_SCOPE', 'PXR_NAMESPACE_CLOSE_SCOPE'.
+// ===================================================================== //
+// --(BEGIN CUSTOM CODE)--

src/experimental/usd/mjcPhysics/generatedSchema.usda

@@ -578,6 +578,29 @@ class "MjcEqualityWeldAPI" (
     )
 }
 
+class "MjcEqualityJointAPI" (
+    apiSchemas = ["MjcEqualityAPI"]
+    doc = """API providing extension attributes to represent equality/joint constraints.
+    This API is applied to a joint prim which acts as the constrained joint (joint1 in
+    MuJoCo terminology). The target relationship points to another joint prim which is
+    the reference joint (joint2 in MuJoCo terminology). The constrained joint's position
+    or angle is constrained to be a quartic polynomial of the reference joint's position
+    or angle. Only scalar joint types (slide and hinge) can be used."""
+)
+{
+    uniform double[] mjc:polycoef = [0, 1, 0, 0, 0] (
+        displayName = "Polynomial Coefficients"
+        doc = """Coefficients a0 through a4 of the quartic polynomial. If the joint values
+        of the constrained joint and reference joint are respectively y and x, and their
+        reference positions (corresponding to the joint values in the initial model
+        configuration) are y0 and x0, the constraint is:
+        y = y0 + a0 + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4.
+        Omitting the target joint (joint2) is equivalent to setting x = x0, in which case
+        the constraint is y = y0 + a0. The default [0, 1, 0, 0, 0] creates a simple 1:1
+        mimic constraint where y tracks x with the same offset from their references."""
+    )
+}
+
 class MjcTendon "MjcTendon" (
     doc = "Type describing fixed and spatial tendons."
 )

src/experimental/usd/mjcPhysics/plugInfo.json

@@ -33,6 +33,16 @@
             ],
             "schemaKind": "singleApplyAPI"
           },
+          "MjcPhysicsEqualityJointAPI": {
+            "alias": {
+              "UsdSchemaBase": "MjcEqualityJointAPI"
+            },
+            "autoGenerated": true,
+            "bases": [
+              "UsdAPISchemaBase"
+            ],
+            "schemaKind": "singleApplyAPI"
+          },
           "MjcPhysicsEqualityWeldAPI": {
             "alias": {
               "UsdSchemaBase": "MjcEqualityWeldAPI"

src/experimental/usd/mjcPhysics/schema.usda

@@ -1106,6 +1106,37 @@ class "MjcEqualityWeldAPI" (
     )
 }
 
+class "MjcEqualityJointAPI" (
+    customData = {
+        string className = "EqualityJointAPI"
+    }
+    doc = """API providing extension attributes to represent equality/joint constraints.
+    This API is applied to a joint prim which acts as the constrained joint (joint1 in
+    MuJoCo terminology). The target relationship points to another joint prim which is
+    the reference joint (joint2 in MuJoCo terminology). The constrained joint's position
+    or angle is constrained to be a quartic polynomial of the reference joint's position
+    or angle. Only scalar joint types (slide and hinge) can be used."""
+
+    prepend apiSchemas = ["MjcEqualityAPI"]
+    inherits = </APISchemaBase>
+)
+{
+    uniform double[] mjc:polycoef = [0, 1, 0, 0, 0] (
+        customData = {
+            string apiName = "Polycoef"
+        }
+        displayName = "Polynomial Coefficients"
+        doc = """Coefficients a0 through a4 of the quartic polynomial. If the joint values
+        of the constrained joint and reference joint are respectively y and x, and their
+        reference positions (corresponding to the joint values in the initial model
+        configuration) are y0 and x0, the constraint is:
+        y = y0 + a0 + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4.
+        Omitting the target joint (joint2) is equivalent to setting x = x0, in which case
+        the constraint is y = y0 + a0. The default [0, 1, 0, 0, 0] creates a simple 1:1
+        mimic constraint where y tracks x with the same offset from their references."""
+    )
+}
+
 class MjcTendon "MjcTendon"
 (
     customData = {

src/experimental/usd/mjcPhysics/tokens.cpp

@@ -150,6 +150,7 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
       mjcPathDivisors("mjc:path:divisors", TfToken::Immortal),
       mjcPathIndices("mjc:path:indices", TfToken::Immortal),
       mjcPathSegments("mjc:path:segments", TfToken::Immortal),
+      mjcPolycoef("mjc:polycoef", TfToken::Immortal),
       mjcPriority("mjc:priority", TfToken::Immortal),
       mjcQpos("mjc:qpos", TfToken::Immortal),
       mjcQvel("mjc:qvel", TfToken::Immortal),
@@ -194,6 +195,7 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
       MjcActuator("MjcActuator", TfToken::Immortal),
       MjcCollisionAPI("MjcCollisionAPI", TfToken::Immortal),
       MjcEqualityAPI("MjcEqualityAPI", TfToken::Immortal),
+      MjcEqualityJointAPI("MjcEqualityJointAPI", TfToken::Immortal),
       MjcEqualityWeldAPI("MjcEqualityWeldAPI", TfToken::Immortal),
       MjcImageableAPI("MjcImageableAPI", TfToken::Immortal),
       MjcJointAPI("MjcJointAPI", TfToken::Immortal),
@@ -329,6 +331,7 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
                  mjcPathDivisors,
                  mjcPathIndices,
                  mjcPathSegments,
+                 mjcPolycoef,
                  mjcPriority,
                  mjcQpos,
                  mjcQvel,
@@ -373,6 +376,7 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
                  MjcActuator,
                  MjcCollisionAPI,
                  MjcEqualityAPI,
+                 MjcEqualityJointAPI,
                  MjcEqualityWeldAPI,
                  MjcImageableAPI,
                  MjcJointAPI,