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

mj_usd plugin: add joints export.

Browse Source 
Note that ball joint is missing, will come later.

PiperOrigin-RevId: 776057523
Change-Id: I8d13c3ea9e7f3d894af73a4640eb93202f8a2a03
This commit is contained in:
Robin Alazard
2025-06-26 03:01:31 -07:00
committed by Copybara-Service
parent 48205fb570
commit 3441e0aa03
4 changed files with 626 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
+232 -15
View File
@@ -21,9 +21,9 @@
#include <utility>
#include <vector>
#include <mujoco/experimental/usd/mjcPhysics/tokens.h>
#include <mujoco/mjspec.h>
#include <mujoco/mujoco.h>
#include <mujoco/experimental/usd/mjcPhysics/tokens.h>
#include "mjcf/utils.h"
#include <pxr/base/arch/attributes.h>
#include <pxr/base/gf/matrix4d.h>
@@ -55,6 +55,11 @@
#include <pxr/usd/usdGeom/metrics.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdLux/tokens.h>
#include <pxr/usd/usdPhysics/fixedJoint.h>
#include <pxr/usd/usdPhysics/joint.h>
#include <pxr/usd/usdPhysics/prismaticJoint.h>
#include <pxr/usd/usdPhysics/revoluteJoint.h>
#include <pxr/usd/usdPhysics/sphericalJoint.h>
#include <pxr/usd/usdPhysics/tokens.h>
#include <pxr/usd/usdShade/tokens.h>
#include <pxr/usdImaging/usdImaging/tokens.h>
@@ -102,6 +107,7 @@ TF_DEFINE_PRIVATE_TOKENS(kTokens,
(repeat)
((sourceMesh, pxr::UsdGeomTokens->Mesh))
((inputsNormal, "inputs:normal"))
((joint, "Joint"))
);
// Using to satisfy TF_REGISTRY_FUNCTION macro below and avoid operating in PXR_NS.
@@ -169,6 +175,9 @@ class ModelWriter {
~ModelWriter() { mj_deleteModel(model_); }
void Write(bool write_physics) {
// Set working parameters.
write_physics_ = write_physics;
// Create top level class holder.
class_path_ = CreateClassSpec(data_, pxr::SdfPath::AbsoluteRootPath(),
pxr::TfToken("__class__"));
@@ -176,9 +185,6 @@ class ModelWriter {
// Create the world body.
body_paths_[kWorldIndex] = WriteWorldBody(kWorldIndex);
// Set working parameters.
write_physics_ = write_physics;
SetLayerMetadata(data_, pxr::SdfFieldKeys->Documentation,
"Generated by mujoco model writer.");
// Mujoco is Z up by default.
@@ -1494,6 +1500,219 @@ class ModelWriter {
}
}
void WriteJoints(mjsBody *body) {
if (!write_physics_) return;
int body_id = mjs_getId(body->element);
if (body_id == kWorldIndex) return;
mjsJoint *joint = mjs_asJoint(mjs_firstChild(body, mjOBJ_JOINT, false));
if (!joint) {
// If no joint is found, then we pass nullptr to create a FixedJoint.
// WriteJoint properly handles the case where the parent is the worldbody.
WriteJoint(nullptr, body);
} else {
WriteJoint(joint, body);
if (mjs_asJoint(mjs_nextChild(body, joint->element, false))) {
TF_WARN(
"Multiple joints found for body %d. Only writing the first one.",
body_id);
}
}
}
// Write the joint. If null, then a FixedJoint is created.
void WriteJoint(mjsJoint *joint, const mjsBody *parent_mj_body) {
// Default to fixed joint if joint is null.
pxr::TfToken joint_prim_type = pxr::UsdPhysicsTokens->PhysicsFixedJoint;
int joint_id = -1;
if (joint) {
joint_id = mjs_getId(joint->element);
mjtJoint type = (mjtJoint)model_->jnt_type[joint_id];
switch (type) {
case mjJNT_FREE:
// Free joints are guaranteed to only ever be on the top-level body so
// we just write no joint. As a top-level body with no joint it will
// be considered a floating-base body.
return;
case mjJNT_HINGE:
joint_prim_type = pxr::UsdPhysicsTokens->PhysicsRevoluteJoint;
break;
case mjJNT_SLIDE:
joint_prim_type = pxr::UsdPhysicsTokens->PhysicsPrismaticJoint;
break;
default:
TF_WARN("Unsupported joint type '%d' for joint '%s'. Skipping.",
(int)type, joint->name->c_str());
return;
}
}
int body_id = mjs_getId(parent_mj_body->element);
// the joint connects the current body as body1, to its parent body as
// body0.
int body1_id_usd = body_id;
int body0_id_usd = model_->body_parentid[body_id];
const pxr::SdfPath &body1_path_usd = body_paths_[body1_id_usd];
auto joint_name = joint ? *joint->name : "FixedJoint";
pxr::TfToken joint_name_token =
GetAvailablePrimName(joint_name, kTokens->joint, body1_path_usd);
pxr::SdfPath joint_path = CreatePrimSpec(data_, body1_path_usd,
joint_name_token, joint_prim_type);
// Set body0 and body1 relationships
// For the initial joints that connect to the world, we signal this by
// keeping the body0 relationship empty.
if (body0_id_usd != kWorldIndex) {
const pxr::SdfPath &body0_path_usd = body_paths_[body0_id_usd];
CreateRelationshipSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsBody0,
body0_path_usd, pxr::SdfVariabilityUniform);
}
CreateRelationshipSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsBody1, body1_path_usd,
pxr::SdfVariabilityUniform);
// Joint frame in MuJoCo is defined by jnt_pos and jnt_axis in body1's frame
// For FixedJoint, these are both unity.
pxr::GfVec3d mj_jnt_pos = pxr::GfVec3d(0.0);
pxr::GfVec3d mj_jnt_axis = pxr::GfVec3d(0.0, 0.0, 1.0);
if (joint) {
mj_jnt_pos = pxr::GfVec3d(&model_->jnt_pos[joint_id * 3]);
mj_jnt_axis = pxr::GfVec3d(&model_->jnt_axis[joint_id * 3]);
}
// Local joint frame for body1
pxr::GfVec3f local_pos1(mj_jnt_pos);
pxr::GfRotation().SetRotateInto(pxr::GfVec3f::ZAxis(), mj_jnt_axis);
pxr::GfQuatf local_rot1(
pxr::GfRotation()
.SetRotateInto(pxr::GfVec3f::ZAxis(), mj_jnt_axis)
.GetQuat());
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsLocalPos1,
pxr::SdfValueTypeNames->Float3),
local_pos1);
if (joint_prim_type == pxr::UsdPhysicsTokens->PhysicsRevoluteJoint ||
joint_prim_type == pxr::UsdPhysicsTokens->PhysicsPrismaticJoint) {
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsLocalRot1,
pxr::SdfValueTypeNames->Quatf),
local_rot1);
} else {
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsLocalRot1,
pxr::SdfValueTypeNames->Quatf),
pxr::GfQuatf::GetIdentity());
}
// Calculate local joint frame for body0
pxr::GfMatrix4d body1_transform_local =
MujocoPosQuatToTransform(&model_->body_pos[body1_id_usd * 3],
&model_->body_quat[body1_id_usd * 4]);
pxr::GfVec3d jnt_pos_parent_local =
body1_transform_local.Transform(mj_jnt_pos);
pxr::GfVec3d jnt_axis_parent_local =
body1_transform_local.TransformDir(mj_jnt_axis);
pxr::GfVec3f local_pos0(jnt_pos_parent_local);
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsLocalPos0,
pxr::SdfValueTypeNames->Float3),
local_pos0);
if (joint_prim_type == pxr::UsdPhysicsTokens->PhysicsRevoluteJoint ||
joint_prim_type == pxr::UsdPhysicsTokens->PhysicsPrismaticJoint) {
pxr::GfQuatf other_rot0(
pxr::GfRotation()
.SetRotateInto(pxr::GfVec3f::ZAxis(), jnt_axis_parent_local)
.GetQuat());
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsLocalRot0,
pxr::SdfValueTypeNames->Quatf),
other_rot0);
} else {
// Fixed joints have no frame and no axis per se. We simply need the
// rotation quaternion of the body its on.
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsLocalRot0,
pxr::SdfValueTypeNames->Quatf),
body1_transform_local.ExtractRotationQuat());
}
if (joint) {
mjtJoint type = (mjtJoint)model_->jnt_type[joint_id];
// Joint-specific attributes
if (type == mjJNT_HINGE || type == mjJNT_SLIDE) {
// The joint motion occurs around/along the Z-axis of the joint frame
// established by localRot0/1.
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsAxis,
pxr::SdfValueTypeNames->Token),
pxr::UsdPhysicsTokens->z); // "Z" axis
}
if (model_->jnt_limited[joint_id]) {
float lower_limit = model_->jnt_range[joint_id * 2];
float upper_limit = model_->jnt_range[joint_id * 2 + 1];
if (type == mjJNT_HINGE) {
// Convert radians to degrees for USD
// As per the XML Reference, "mjModel always uses radians"
lower_limit *= (180.0 / mjPI);
upper_limit *= (180.0 / mjPI);
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsLowerLimit,
pxr::SdfValueTypeNames->Float),
lower_limit);
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsUpperLimit,
pxr::SdfValueTypeNames->Float),
upper_limit);
} else if (type == mjJNT_SLIDE) {
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsLowerLimit,
pxr::SdfValueTypeNames->Float),
lower_limit);
SetAttributeDefault(
data_,
CreateAttributeSpec(data_, joint_path,
pxr::UsdPhysicsTokens->physicsUpperLimit,
pxr::SdfValueTypeNames->Float),
upper_limit);
}
}
}
}
void WriteCamera(mjsCamera *spec_cam, const mjsBody *body) {
const auto &body_path = body_paths_[mjs_getId(body->element)];
auto name = GetAvailablePrimName(*spec_cam->name,
@@ -1635,17 +1854,6 @@ class ModelWriter {
ApplyApiSchema(data_, body_path,
pxr::UsdPhysicsTokens->PhysicsRigidBodyAPI);
// 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 =
mjs_getId(mjs_getParent(parent->element)->element);
if (parent_parent_id == kWorldIndex) {
ApplyApiSchema(data_, parent_path,
pxr::UsdPhysicsTokens->PhysicsArticulationRootAPI);
}
}
}
// Create classes if necessary
@@ -1692,6 +1900,7 @@ class ModelWriter {
}
WriteSites(body);
WriteGeoms(body);
WriteJoints(body);
WriteCameras(body);
WriteLights(body);
body = mjs_asBody(mjs_nextElement(spec_, body->element));
@@ -1707,6 +1916,14 @@ class ModelWriter {
CreatePrimSpec(data_, pxr::SdfPath::AbsoluteRootPath(), name,
pxr::UsdGeomTokens->Xform);
SetPrimKind(data_, world_group_path, pxr::KindTokens->group);
if (write_physics_) {
// Apply the PhysicsArticulationRootAPI to the world body so that
// everything under it is automatically considered an articulation.
ApplyApiSchema(data_, world_group_path,
pxr::UsdPhysicsTokens->PhysicsArticulationRootAPI);
}
return world_group_path;
}
};