Add MjcEqualityAPI as a base class for MuJoCo constraints and MjcEqualityWeldAPI for representing weld constraints.

This change introduces two new USD API schema: * MjcEqualityAPI - base API for representing MuJoCo equality constraints in USD, this will typically be applied to joints such as UsdPhysicsFixedJoint. * MjcEqualityWeldAPI - extends MjcEqualityAPI for MuJoCo's weld equality constraints with an attribute for torqueScale. The usd_decoder is updated to parse UsdPhysicsFixedJoints and convert the new schema attributes into the appropriate mjsEquality data fields. PiperOrigin-RevId: 856723634 Change-Id: I0982fd3168b643f8e91d87bb5ae5933e526b63a1
2026-01-15 10:43:29 -08:00
parent d86acd7e12
commit c08e1d9e6e
14 changed files with 1004 additions and 12 deletions
@@ -127,6 +127,8 @@ add_library(${MJC_PHYSICS_PLUGIN_TARGET_NAME} SHARED)
 target_sources(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PRIVATE
  mjcPhysics/actuator.cpp
  mjcPhysics/collisionAPI.cpp
+  mjcPhysics/equalityAPI.cpp
+  mjcPhysics/equalityWeldAPI.cpp
  mjcPhysics/imageableAPI.cpp
  mjcPhysics/jointAPI.cpp
  mjcPhysics/keyframe.cpp
@@ -162,19 +164,19 @@ target_compile_options(

 # Configure RPATH for macOS and UNIX.
 if(APPLE)
-  set_target_properties($(MJCF_PLUGIN_TARGET_NAME) PROPERTIES
+  set_target_properties(${MJCF_PLUGIN_TARGET_NAME} PROPERTIES
    BUILD_RPATH "${USD_INSTALL_ROOT}"
    INSTALL_RPATH "@loader_path")

-  set_target_properties($(MJC_PHYSICS_PLUGIN_TARGET_NAME) PROPERTIES
+  set_target_properties(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PROPERTIES
    BUILD_RPATH "${USD_INSTALL_ROOT}"
    INSTALL_RPATH "@loader_path")
 elseif(UNIX)
-  set_target_properties($(MJCF_PLUGIN_TARGET_NAME) PROPERTIES
+  set_target_properties(${MJCF_PLUGIN_TARGET_NAME} PROPERTIES
    BUILD_RPATH "${USD_INSTALL_ROOT}"
    INSTALL_RPATH "$ORIGIN")

-  set_target_properties($(MJC_PHYSICS_PLUGIN_TARGET_NAME) PROPERTIES
+  set_target_properties(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PROPERTIES
    BUILD_RPATH "${USD_INSTALL_ROOT}"
    INSTALL_RPATH "$ORIGIN")
 endif()
@@ -0,0 +1,144 @@
+// 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/equalityAPI.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<MjcPhysicsEqualityAPI, TfType::Bases<UsdAPISchemaBase> >();
+}
+
+/* virtual */
+MjcPhysicsEqualityAPI::~MjcPhysicsEqualityAPI() {}
+
+/* static */
+MjcPhysicsEqualityAPI MjcPhysicsEqualityAPI::Get(const UsdStagePtr& stage,
+                                                 const SdfPath& path) {
+  if (!stage) {
+    TF_CODING_ERROR("Invalid stage");
+    return MjcPhysicsEqualityAPI();
+  }
+  return MjcPhysicsEqualityAPI(stage->GetPrimAtPath(path));
+}
+
+/* virtual */
+UsdSchemaKind MjcPhysicsEqualityAPI::_GetSchemaKind() const {
+  return MjcPhysicsEqualityAPI::schemaKind;
+}
+
+/* static */
+bool MjcPhysicsEqualityAPI::CanApply(const UsdPrim& prim, std::string* whyNot) {
+  return prim.CanApplyAPI<MjcPhysicsEqualityAPI>(whyNot);
+}
+
+/* static */
+MjcPhysicsEqualityAPI MjcPhysicsEqualityAPI::Apply(const UsdPrim& prim) {
+  if (prim.ApplyAPI<MjcPhysicsEqualityAPI>()) {
+    return MjcPhysicsEqualityAPI(prim);
+  }
+  return MjcPhysicsEqualityAPI();
+}
+
+/* static */
+const TfType& MjcPhysicsEqualityAPI::_GetStaticTfType() {
+  static TfType tfType = TfType::Find<MjcPhysicsEqualityAPI>();
+  return tfType;
+}
+
+/* static */
+bool MjcPhysicsEqualityAPI::_IsTypedSchema() {
+  static bool isTyped = _GetStaticTfType().IsA<UsdTyped>();
+  return isTyped;
+}
+
+/* virtual */
+const TfType& MjcPhysicsEqualityAPI::_GetTfType() const {
+  return _GetStaticTfType();
+}
+
+UsdAttribute MjcPhysicsEqualityAPI::GetSolRefAttr() const {
+  return GetPrim().GetAttribute(MjcPhysicsTokens->mjcSolref);
+}
+
+UsdAttribute MjcPhysicsEqualityAPI::CreateSolRefAttr(
+    VtValue const& defaultValue, bool writeSparsely) const {
+  return UsdSchemaBase::_CreateAttr(
+      MjcPhysicsTokens->mjcSolref, SdfValueTypeNames->DoubleArray,
+      /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely);
+}
+
+UsdAttribute MjcPhysicsEqualityAPI::GetSolImpAttr() const {
+  return GetPrim().GetAttribute(MjcPhysicsTokens->mjcSolimp);
+}
+
+UsdAttribute MjcPhysicsEqualityAPI::CreateSolImpAttr(
+    VtValue const& defaultValue, bool writeSparsely) const {
+  return UsdSchemaBase::_CreateAttr(
+      MjcPhysicsTokens->mjcSolimp, SdfValueTypeNames->DoubleArray,
+      /* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely);
+}
+
+UsdRelationship MjcPhysicsEqualityAPI::GetMjcTargetRel() const {
+  return GetPrim().GetRelationship(MjcPhysicsTokens->mjcTarget);
+}
+
+UsdRelationship MjcPhysicsEqualityAPI::CreateMjcTargetRel() const {
+  return GetPrim().CreateRelationship(MjcPhysicsTokens->mjcTarget,
+                                      /* custom = */ false);
+}
+
+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& MjcPhysicsEqualityAPI::GetSchemaAttributeNames(
+    bool includeInherited) {
+  static TfTokenVector localNames = {
+      MjcPhysicsTokens->mjcSolref,
+      MjcPhysicsTokens->mjcSolimp,
+  };
+  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)--
@@ -0,0 +1,125 @@
+// 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/equalityWeldAPI.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<MjcPhysicsEqualityWeldAPI, TfType::Bases<UsdAPISchemaBase> >();
+}
+
+/* virtual */
+MjcPhysicsEqualityWeldAPI::~MjcPhysicsEqualityWeldAPI() {}
+
+/* static */
+MjcPhysicsEqualityWeldAPI MjcPhysicsEqualityWeldAPI::Get(
+    const UsdStagePtr& stage, const SdfPath& path) {
+  if (!stage) {
+    TF_CODING_ERROR("Invalid stage");
+    return MjcPhysicsEqualityWeldAPI();
+  }
+  return MjcPhysicsEqualityWeldAPI(stage->GetPrimAtPath(path));
+}
+
+/* virtual */
+UsdSchemaKind MjcPhysicsEqualityWeldAPI::_GetSchemaKind() const {
+  return MjcPhysicsEqualityWeldAPI::schemaKind;
+}
+
+/* static */
+bool MjcPhysicsEqualityWeldAPI::CanApply(const UsdPrim& prim,
+                                         std::string* whyNot) {
+  return prim.CanApplyAPI<MjcPhysicsEqualityWeldAPI>(whyNot);
+}
+
+/* static */
+MjcPhysicsEqualityWeldAPI MjcPhysicsEqualityWeldAPI::Apply(
+    const UsdPrim& prim) {
+  if (prim.ApplyAPI<MjcPhysicsEqualityWeldAPI>()) {
+    return MjcPhysicsEqualityWeldAPI(prim);
+  }
+  return MjcPhysicsEqualityWeldAPI();
+}
+
+/* static */
+const TfType& MjcPhysicsEqualityWeldAPI::_GetStaticTfType() {
+  static TfType tfType = TfType::Find<MjcPhysicsEqualityWeldAPI>();
+  return tfType;
+}
+
+/* static */
+bool MjcPhysicsEqualityWeldAPI::_IsTypedSchema() {
+  static bool isTyped = _GetStaticTfType().IsA<UsdTyped>();
+  return isTyped;
+}
+
+/* virtual */
+const TfType& MjcPhysicsEqualityWeldAPI::_GetTfType() const {
+  return _GetStaticTfType();
+}
+
+UsdAttribute MjcPhysicsEqualityWeldAPI::GetTorqueScaleAttr() const {
+  return GetPrim().GetAttribute(MjcPhysicsTokens->mjcTorqueScale);
+}
+
+UsdAttribute MjcPhysicsEqualityWeldAPI::CreateTorqueScaleAttr(
+    VtValue const& defaultValue, bool writeSparsely) const {
+  return UsdSchemaBase::_CreateAttr(
+      MjcPhysicsTokens->mjcTorqueScale, SdfValueTypeNames->Float,
+      /* 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& MjcPhysicsEqualityWeldAPI::GetSchemaAttributeNames(
+    bool includeInherited) {
+  static TfTokenVector localNames = {
+      MjcPhysicsTokens->mjcTorqueScale,
+  };
+  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)--
@@ -550,6 +550,34 @@ class "MjcMaterialAPI" (
    )
 }

+class "MjcEqualityAPI" (
+    doc = "Base API for equality constraints."
+)
+{
+    uniform double[] mjc:solimp = [0.9, 0.95, 0.001, 0.5, 2] (
+        displayName = "SolImp"
+        doc = "Constraint solver parameter for equality constraint simulation."
+    )
+    uniform double[] mjc:solref = [0.02, 1] (
+        displayName = "SolRef"
+        doc = "Constraint solver parameter for equality constraint simulation."
+    )
+    rel mjc:target (
+        doc = "Secondary target of the equality constraint."
+    )
+}
+
+class "MjcEqualityWeldAPI" (
+    apiSchemas = ["MjcEqualityAPI"]
+    doc = "API providing extension attributes to represent equality/weld constraints."
+)
+{
+    uniform float mjc:torqueScale = 1 (
+        displayName = "Torque Scale"
+        doc = "A constant that scales the angular residual (angular constraint violation). Notionally in units of torque/force = length. Intuitively this coefficient defines how much the weld “cares” about rotational displacements vs. translational displacements."
+    )
+}
+
 class MjcTendon "MjcTendon" (
    doc = "Type describing fixed and spatial tendons."
 )
@@ -23,6 +23,26 @@
            ],
            "schemaKind": "singleApplyAPI"
          },
+          "MjcPhysicsEqualityAPI": {
+            "alias": {
+              "UsdSchemaBase": "MjcEqualityAPI"
+            },
+            "autoGenerated": true,
+            "bases": [
+              "UsdAPISchemaBase"
+            ],
+            "schemaKind": "singleApplyAPI"
+          },
+          "MjcPhysicsEqualityWeldAPI": {
+            "alias": {
+              "UsdSchemaBase": "MjcEqualityWeldAPI"
+            },
+            "autoGenerated": true,
+            "bases": [
+              "UsdAPISchemaBase"
+            ],
+            "schemaKind": "singleApplyAPI"
+          },
          "MjcPhysicsImageableAPI": {
            "alias": {
              "UsdSchemaBase": "MjcImageableAPI"
@@ -1057,6 +1057,55 @@ class "MjcMaterialAPI"
    )
 }

+class "MjcEqualityAPI" (
+    customData = {
+        string className = "EqualityAPI"
+    }
+    doc = """Base API for equality constraints."""
+
+    inherits = </APISchemaBase>
+)
+{
+    rel mjc:target (
+      doc = "Secondary target of the equality constraint."
+    )
+
+    uniform double[] mjc:solref = [0.02, 1.0] (
+        customData = {
+            string apiName = "SolRef"
+        }
+        displayName = "SolRef"
+        doc = """Constraint solver parameter for equality constraint simulation."""
+    )
+
+    uniform double[] mjc:solimp = [0.9, 0.95, 0.001, 0.5, 2.0] (
+        customData = {
+            string apiName = "SolImp"
+        }
+        displayName = "SolImp"
+        doc = """Constraint solver parameter for equality constraint simulation."""
+    )
+}
+
+class "MjcEqualityWeldAPI" (
+    customData = {
+        string className = "EqualityWeldAPI"
+    }
+    doc = """API providing extension attributes to represent equality/weld constraints."""
+
+    prepend apiSchemas = ["MjcEqualityAPI"]
+    inherits = </APISchemaBase>
+)
+{
+    uniform float mjc:torqueScale = 1.0 (
+        customData = {
+            string apiName = "TorqueScale"
+        }
+        displayName = "Torque Scale"
+        doc = """A constant that scales the angular residual (angular constraint violation). Notionally in units of torque/force = length. Intuitively this coefficient defines how much the weld “cares” about rotational displacements vs. translational displacements."""
+    )
+}
+
 class MjcTendon "MjcTendon"
 (
    customData = {
@@ -175,6 +175,7 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
      mjcSpringref("mjc:springref", TfToken::Immortal),
      mjcStiffness("mjc:stiffness", TfToken::Immortal),
      mjcTarget("mjc:target", TfToken::Immortal),
+      mjcTorqueScale("mjc:torqueScale", TfToken::Immortal),
      mjcTorsionalfriction("mjc:torsionalfriction", TfToken::Immortal),
      mjcType("mjc:type", TfToken::Immortal),
      mjcWidth("mjc:width", TfToken::Immortal),
@@ -192,6 +193,8 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
      user("user", TfToken::Immortal),
      MjcActuator("MjcActuator", TfToken::Immortal),
      MjcCollisionAPI("MjcCollisionAPI", TfToken::Immortal),
+      MjcEqualityAPI("MjcEqualityAPI", TfToken::Immortal),
+      MjcEqualityWeldAPI("MjcEqualityWeldAPI", TfToken::Immortal),
      MjcImageableAPI("MjcImageableAPI", TfToken::Immortal),
      MjcJointAPI("MjcJointAPI", TfToken::Immortal),
      MjcKeyframe("MjcKeyframe", TfToken::Immortal),
@@ -351,6 +354,7 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
                 mjcSpringref,
                 mjcStiffness,
                 mjcTarget,
+                 mjcTorqueScale,
                 mjcTorsionalfriction,
                 mjcType,
                 mjcWidth,
@@ -368,6 +372,8 @@ MjcPhysicsTokensType::MjcPhysicsTokensType()
                 user,
                 MjcActuator,
                 MjcCollisionAPI,
+                 MjcEqualityAPI,
+                 MjcEqualityWeldAPI,
                 MjcImageableAPI,
                 MjcJointAPI,
                 MjcKeyframe,