chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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
This commit is contained in:
Sam Haves
2026-01-15 10:43:29 -08:00
committed by Copybara-Service
parent d86acd7e12
commit c08e1d9e6e
14 changed files with 1004 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plugin/usd_decoder/usd_decoder.cc
+191 -2
View File
@@ -25,6 +25,8 @@
#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/equalityWeldAPI.h>
#include <mujoco/experimental/usd/mjcPhysics/imageableAPI.h>
#include <mujoco/experimental/usd/mjcPhysics/jointAPI.h>
#include <mujoco/experimental/usd/mjcPhysics/keyframe.h>
@@ -39,6 +41,7 @@
#include "material_parsing.h"
#include <pxr/base/gf/declare.h>
#include <pxr/base/gf/matrix4d.h>
#include <pxr/base/gf/matrix4f.h>
#include <pxr/base/gf/rotation.h>
#include <pxr/base/gf/vec3d.h>
#include <pxr/base/tf/token.h>
@@ -1821,16 +1824,197 @@ 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);
pxr::SdfPath body0_path;
if (!targets0.empty()) body0_path = targets0[0];
pxr::SdfPath body1_path;
if (!targets1.empty()) body1_path = targets1[0];
auto stage = prim.GetStage();
auto body0_prim = stage->GetPrimAtPath(body0_path);
auto body1_prim = stage->GetPrimAtPath(body1_path);
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>();
}
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;
}
mjsEquality* eq = mjs_addEquality(spec, nullptr);
eq->type = mjEQ_WELD;
mjs_setName(eq->element, prim.GetPath().GetAsString().c_str());
SetUsdPrimPathUserValue(eq->element, prim.GetPath());
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;
}
// 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)) {
// unable to decompose, emit warning and set scale to identity
mju_warning(
"Unable to decompose matrix for body 1: %s.",
body1_path.GetAsString().c_str());
body1_scale = pxr::GfVec3f(1, 1, 1);
}
}
pxr::GfVec3f localPos1;
joint.GetLocalPos1Attr().Get(&localPos1);
localPos1[0] *= body1_scale[0];
localPos1[1] *= body1_scale[1];
localPos1[2] *= body1_scale[2];
pxr::GfQuatf localRot1;
joint.GetLocalRot1Attr().Get(&localRot1);
pxr::GfVec3f localPos0;
joint.GetLocalPos0Attr().Get(&localPos0);
localPos0[0] *= body0_scale[0];
localPos0[1] *= body0_scale[1];
localPos0[2] *= body0_scale[2];
pxr::GfQuatf localRot0;
joint.GetLocalRot0Attr().Get(&localRot0);
auto relpose_quat = localRot0 * localRot1.GetConjugate();
relpose_quat.Normalize();
auto relpose_pos = localPos0 - relpose_quat.Transform(localPos1);
eq->data[0] = localPos1[0];
eq->data[1] = localPos1[1];
eq->data[2] = localPos1[2];
eq->data[3] = relpose_pos[0];
eq->data[4] = relpose_pos[1];
eq->data[5] = relpose_pos[2];
eq->data[6] = relpose_quat.GetReal();
eq->data[7] = relpose_quat.GetImaginary()[0];
eq->data[8] = relpose_quat.GetImaginary()[1];
eq->data[9] = relpose_quat.GetImaginary()[2];
if (prim.HasAPI<pxr::MjcPhysicsEqualityWeldAPI>()) {
// MjcPhysicsEqualityAPI is always automatically applied
// by MjcPhysicsEqualityWeldAPI.
pxr::MjcPhysicsEqualityAPI equality_api(prim);
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];
}
} else {
mju_warning(
"solref attribute for weld equality %s has incorrect size "
"%zu, expected %d.",
prim.GetPath().GetAsString().c_str(), solref.size(), mjNREF);
}
}
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 {
mju_warning(
"solimp attribute for weld equality %s has incorrect size "
"%zu, expected %d.",
prim.GetPath().GetAsString().c_str(), solimp.size(), mjNIMP);
}
}
pxr::MjcPhysicsEqualityWeldAPI weld_api(prim);
auto torque_scale_attr = weld_api.GetTorqueScaleAttr();
if (torque_scale_attr.HasAuthoredValue()) {
float torque_scale;
torque_scale_attr.Get(&torque_scale);
eq->data[10] = torque_scale;
}
}
}
void ParseUsdPhysicsJoint(mjSpec* spec, const pxr::UsdPrim& prim, mjsBody* body,
pxr::UsdGeomXformCache& xform_cache) {
pxr::UsdPhysicsJoint joint(prim);
// A fixed joint means the bodies are welded.
if (prim.IsA<pxr::UsdPhysicsFixedJoint>()) {
// No joint needed for welded bodies.
return;
}
pxr::UsdPhysicsJoint joint(prim);
mjtJoint type;
if (prim.IsA<pxr::UsdPhysicsRevoluteJoint>()) {
type = mjJNT_HINGE;
@@ -2093,6 +2277,11 @@ void PopulateSpecFromTree(pxr::UsdStageRefPtr stage, mjSpec* spec,
free_joint->type = mjJNT_FREE;
}
for (const auto& constraint_path : current_node->constraints) {
ParseConstraint(spec, stage->GetPrimAtPath(constraint_path),
current_mj_body, caches.xform_cache);
}
pxr::UsdPrim body_prim_for_xform =
current_node->body_path.IsEmpty()
? stage->GetPseudoRoot()