Add equalityJointAPI to mjcPhysics.

PiperOrigin-RevId: 860081702
Change-Id: Ifd961dcb2a8e0f5d885e9e9dadc0153b03cb5f4d

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,