Add support for inertia and mass to Mujoco USD interop.

PiperOrigin-RevId: 763737956
Change-Id: I6cdf267d76536a26aeb6327ce90165842e32bd8a
This commit is contained in:
Sam Haves
2025-05-27 04:57:00 -07:00
committed by Copybara-Service
parent de3dc7c234
commit 441ca838a1
4 changed files with 223 additions and 5 deletions
@@ -26,6 +26,7 @@
#include "mjcf/utils.h"
#include <pxr/base/arch/attributes.h>
#include <pxr/base/gf/matrix4d.h>
#include <pxr/base/gf/quatf.h>
#include <pxr/base/gf/rotation.h>
#include <pxr/base/gf/vec2f.h>
#include <pxr/base/gf/vec3d.h>
@@ -180,7 +181,7 @@ class ModelWriter {
WritePhysicsScene();
// Author mesh scope + mesh prims to be referenced.
WriteMeshes();
WriteMeshes(write_physics);
WriteMaterials();
WriteBodies(write_physics);
}
@@ -298,7 +299,8 @@ class ModelWriter {
SetAttributeDefault(data_, xform_op_order_path, new_order);
}
void WriteMesh(const mjsMesh *mesh, const pxr::SdfPath &parent_path) {
void WriteMesh(const mjsMesh *mesh, const pxr::SdfPath &parent_path,
bool write_physics) {
auto name = GetAvailablePrimName(*mesh->name, pxr::UsdGeomTokens->Mesh,
parent_path);
pxr::SdfPath subcomponent_path =
@@ -308,6 +310,22 @@ class ModelWriter {
pxr::UsdGeomTokens->Mesh);
mesh_paths_[*mesh->name] = subcomponent_path;
if (write_physics) {
ApplyApiSchema(data_, mesh_path, MjcPhysicsTokens->MeshCollisionAPI);
pxr::TfToken inertia = MjcPhysicsTokens->legacy;
if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_EXACT) {
inertia = MjcPhysicsTokens->exact;
} else if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_CONVEX) {
inertia = MjcPhysicsTokens->convex;
} else if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_SHELL) {
inertia = MjcPhysicsTokens->shell;
}
WriteUniformAttribute(mesh_path, pxr::SdfValueTypeNames->Token,
MjcPhysicsTokens->mjcInertia, inertia);
}
// NOTE: The geometry data taken from the spec is the post-compilation
// data after it has been mjCMesh::Compile'd. So don't be surprised if
// things like user defined vertices have moved due to re-centering to
@@ -584,7 +602,7 @@ class ModelWriter {
}
}
void WriteMeshes() {
void WriteMeshes(bool write_physics) {
// Create a scope for the meshes to keep things organized
pxr::SdfPath scope_path =
CreatePrimSpec(data_, body_paths_[kWorldIndex], kTokens->meshScope,
@@ -596,7 +614,7 @@ class ModelWriter {
mjsMesh *mesh = mjs_asMesh(mjs_firstElement(spec_, mjOBJ_MESH));
while (mesh) {
WriteMesh(mesh, scope_path);
WriteMesh(mesh, scope_path, write_physics);
mesh = mjs_asMesh(mjs_nextElement(spec_, mesh->element));
}
}
@@ -1032,12 +1050,43 @@ class ModelWriter {
return;
}
// Apply the PhysicsCollisionAPI schema if we are writing physics and the
// Apply the physics schemas if we are writing physics and the
// geom participates in collisions.
if (write_physics && (model_->geom_contype[geom_id] != 0 ||
model_->geom_conaffinity[geom_id] != 0)) {
ApplyApiSchema(data_, geom_path,
pxr::UsdPhysicsTokens->PhysicsCollisionAPI);
ApplyApiSchema(data_, geom_path, MjcPhysicsTokens->CollisionAPI);
WriteUniformAttribute(
geom_path, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcShellinertia,
geom->typeinertia == mjtGeomInertia::mjINERTIA_SHELL);
if (geom->mass >= mjMINVAL || geom->density >= mjMINVAL) {
ApplyApiSchema(data_, geom_path, pxr::UsdPhysicsTokens->PhysicsMassAPI);
}
if (geom->mass >= mjMINVAL) {
pxr::SdfPath mass_attr = CreateAttributeSpec(
data_, geom_path, pxr::UsdPhysicsTokens->physicsMass,
pxr::SdfValueTypeNames->Float, pxr::SdfVariabilityUniform);
// Make sure to cast to float here since mjtNum might be a double.
SetAttributeDefault(data_, mass_attr, (float)geom->mass);
}
// Even though density is not used for mass computation when mass exists
// we want to retain the information anyways.
if (geom->density >= mjMINVAL) {
pxr::SdfPath density_attr = CreateAttributeSpec(
data_, geom_path, pxr::UsdPhysicsTokens->physicsDensity,
pxr::SdfValueTypeNames->Float, pxr::SdfVariabilityUniform);
// Make sure to cast to float here since mjtNum might be a double.
SetAttributeDefault(data_, density_attr, (float)geom->density);
}
// For meshes, also apply PhysicsMeshCollisionAPI and set the
// approximation attribute.
if (geom->type == mjGEOM_MESH) {
@@ -1238,6 +1287,44 @@ class ModelWriter {
// Apply the PhysicsRigidBodyAPI schema if we are writing physics.
if (write_physics) {
// If the body had a mass specified then it must have either inertia or
// fullinertia specified per inertia element XML documentation.
// Therefore it is sufficient to check if the mass is non-zero to see if
// we should set inertial attributes on the body.
//
// Note that if the user has NOT specified any inertial properties then
// we don't want to pull values from the compiled model since coming back
// into Mujoco would take those values instead of computing them
// automatically from the subtree.
if (body->mass > 0) {
// User might have specified the inertia via fullinertia and the
// compiler has extracted all values properly. So leverage those
// instead of doing the computation ourselves here.
ApplyApiSchema(data_, body_path, pxr::UsdPhysicsTokens->PhysicsMassAPI);
WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Float,
pxr::UsdPhysicsTokens->physicsMass,
(float)model_->body_mass[body_id]);
mjtNum *body_ipos = &model_->body_ipos[body_id * 3];
pxr::GfVec3f inertial_pos(body_ipos[0], body_ipos[1], body_ipos[2]);
WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Point3f,
pxr::UsdPhysicsTokens->physicsCenterOfMass,
inertial_pos);
mjtNum *body_iquat = &model_->body_iquat[body_id * 4];
pxr::GfQuatf inertial_frame(body_iquat[0], body_iquat[1], body_iquat[2],
body_iquat[3]);
WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Quatf,
pxr::UsdPhysicsTokens->physicsPrincipalAxes,
inertial_frame);
mjtNum *inertia = &model_->body_inertia[body_id * 3];
pxr::GfVec3f diag_inertia(inertia[0], inertia[1], inertia[2]);
WriteUniformAttribute(body_path, pxr::SdfValueTypeNames->Float3,
pxr::UsdPhysicsTokens->physicsDiagonalInertia,
diag_inertia);
}
ApplyApiSchema(data_, body_path,
pxr::UsdPhysicsTokens->PhysicsRigidBodyAPI);
@@ -17,6 +17,8 @@
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "src/experimental/usd/mjcPhysics/collisionAPI.h"
#include "src/experimental/usd/mjcPhysics/meshCollisionAPI.h"
#include "src/experimental/usd/mjcPhysics/sceneAPI.h"
#include "src/experimental/usd/mjcPhysics/siteAPI.h"
#include "src/experimental/usd/mjcPhysics/tokens.h"
@@ -50,6 +52,7 @@
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdPhysics/articulationRootAPI.h>
#include <pxr/usd/usdPhysics/collisionAPI.h>
#include <pxr/usd/usdPhysics/massAPI.h>
#include <pxr/usd/usdPhysics/meshCollisionAPI.h>
#include <pxr/usd/usdPhysics/rigidBodyAPI.h>
#include <pxr/usd/usdPhysics/scene.h>
@@ -1106,6 +1109,7 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsColliders) {
EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/ground",
pxr::UsdPhysicsRigidBodyAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/ground", pxr::UsdPhysicsCollisionAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/ground", pxr::MjcPhysicsCollisionAPI);
// body_0/body_0_0 [rigidbody] (Nested body - reparented)
EXPECT_PRIM_VALID(stage, "/test/body_0/body_0_0");
@@ -1113,30 +1117,40 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsColliders) {
pxr::UsdPhysicsRigidBodyAPI);
EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_0/body_0_0",
pxr::UsdPhysicsCollisionAPI);
EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_0/body_0_0",
pxr::MjcPhysicsCollisionAPI);
// body_0/body_0_0/body_0_0_col [collider]
EXPECT_PRIM_VALID(stage, "/test/body_0/body_0_0/body_0_0_col");
EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_0/body_0_0/body_0_0_col",
pxr::UsdPhysicsRigidBodyAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/body_0/body_0_0/body_0_0_col",
pxr::UsdPhysicsCollisionAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/body_0/body_0_0/body_0_0_col",
pxr::MjcPhysicsCollisionAPI);
// body_1 [rigidbody]
EXPECT_PRIM_VALID(stage, "/test/body_1");
EXPECT_PRIM_API_APPLIED(stage, "/test/body_1", pxr::UsdPhysicsRigidBodyAPI);
EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1",
pxr::UsdPhysicsCollisionAPI);
EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1",
pxr::MjcPhysicsCollisionAPI);
// body_1/body_1_col_0 [collider]
EXPECT_PRIM_VALID(stage, "/test/body_1/body_1_col_0");
EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1/body_1_col_0",
pxr::UsdPhysicsRigidBodyAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_0",
pxr::UsdPhysicsCollisionAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_0",
pxr::MjcPhysicsCollisionAPI);
// body_1/body_1_col_1 [collider]
EXPECT_PRIM_VALID(stage, "/test/body_1/body_1_col_1");
EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body_1/body_1_col_1",
pxr::UsdPhysicsRigidBodyAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_1",
pxr::UsdPhysicsCollisionAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/body_1/body_1_col_1",
pxr::MjcPhysicsCollisionAPI);
// body_2 [rigidbody]
EXPECT_PRIM_VALID(stage, "/test/body_2");
@@ -1169,11 +1183,117 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsColliders) {
pxr::UsdPhysicsCollisionAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/body_3/body_3_col/Mesh",
pxr::UsdPhysicsMeshCollisionAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/body_3/body_3_col/Mesh",
pxr::MjcPhysicsMeshCollisionAPI);
ExpectAttributeEqual(stage,
"/test/body_3/body_3_col/Mesh.physics:approximation",
pxr::UsdPhysicsTokens->convexHull);
}
TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsCollisionAPI) {
static constexpr char xml[] = R"(
<mujoco model="test">
<worldbody>
<body name="body">
<geom name="box" type="box" size=".05 .05 .05" mass="0.1" shellinertia="true"/>
</body>
</worldbody>
</mujoco>
)";
auto stage = OpenStageWithPhysics(xml);
ExpectAttributeEqual(stage, "/test/body/box.mjc:shellinertia", true);
}
TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsMeshCollisionAPI) {
static constexpr char xml[] = R"(
<mujoco model="test">
<asset>
<mesh name="tet_legacy" inertia="legacy" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
<mesh name="tet_exact" inertia="exact" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
<mesh name="tet_convex" inertia="convex" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
<mesh name="tet_shell" inertia="shell" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
</asset>
<worldbody>
<body name="body">
<geom name="tet_legacy" type="mesh" mesh="tet_legacy"/>
<geom name="tet_exact" type="mesh" mesh="tet_exact"/>
<geom name="tet_convex" type="mesh" mesh="tet_convex"/>
<geom name="tet_shell" type="mesh" mesh="tet_shell"/>
</body>
</worldbody>
</mujoco>
)";
auto stage = OpenStageWithPhysics(xml);
ExpectAttributeEqual(stage, "/test/body/tet_legacy/Mesh.mjc:inertia",
MjcPhysicsTokens->legacy);
ExpectAttributeEqual(stage, "/test/body/tet_exact/Mesh.mjc:inertia",
MjcPhysicsTokens->exact);
ExpectAttributeEqual(stage, "/test/body/tet_convex/Mesh.mjc:inertia",
MjcPhysicsTokens->convex);
ExpectAttributeEqual(stage, "/test/body/tet_shell/Mesh.mjc:inertia",
MjcPhysicsTokens->shell);
}
TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIApplied) {
static constexpr char xml[] = R"(
<mujoco model="test">
<worldbody>
<body name="body">
<geom name="box" type="box" size=".05 .05 .05" mass="0.1"/>
</body>
</worldbody>
</mujoco>
)";
auto stage = OpenStageWithPhysics(xml);
EXPECT_PRIM_VALID(stage, "/test/body");
EXPECT_PRIM_VALID(stage, "/test/body/box");
EXPECT_PRIM_API_APPLIED(stage, "/test/body/box", pxr::UsdPhysicsMassAPI);
EXPECT_PRIM_API_NOT_APPLIED(stage, "/test/body", pxr::UsdPhysicsMassAPI);
ExpectAttributeEqual(stage, "/test/body/box.physics:mass", 0.1f);
}
TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIAppliedToBody) {
static constexpr char xml[] = R"(
<mujoco model="test">
<worldbody>
<body name="body">
<inertial pos="1 2 3" mass="3" diaginertia="1 1 1"/>
<geom name="box" type="box" size=".05 .05 .05" mass="0.1"/>
</body>
</worldbody>
</mujoco>
)";
auto stage = OpenStageWithPhysics(xml);
EXPECT_PRIM_VALID(stage, "/test/body");
EXPECT_PRIM_VALID(stage, "/test/body/box");
EXPECT_PRIM_API_APPLIED(stage, "/test/body/box", pxr::UsdPhysicsMassAPI);
EXPECT_PRIM_API_APPLIED(stage, "/test/body", pxr::UsdPhysicsMassAPI);
// Make sure that body gets it's inertial elements from the inertial element
// and not from the subtree.
ExpectAttributeEqual(stage, "/test/body.physics:mass", 3.0f);
ExpectAttributeEqual(stage, "/test/body.physics:centerOfMass",
pxr::GfVec3f(1, 2, 3));
}
TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIDensity) {
static constexpr char xml[] = R"(
<mujoco model="test">
<worldbody>
<body name="body">
<geom name="box" type="box" size=".05 .05 .05" density="1234"/>
</body>
</worldbody>
</mujoco>
)";
auto stage = OpenStageWithPhysics(xml);
ExpectAttributeEqual(stage, "/test/body/box.physics:density", 1234.0f);
}
} // namespace
} // namespace usd
} // namespace mujoco
+9
View File
@@ -42,6 +42,15 @@ pxr::SdfLayerRefPtr LoadLayer(
return layer;
}
pxr::UsdStageRefPtr OpenStageWithPhysics(const std::string& xml) {
pxr::SdfFileFormat::FileFormatArguments args;
args["usdMjcfToggleUsdPhysics"] = "true";
pxr::SdfLayerRefPtr layer = LoadLayer(xml, args);
auto stage = pxr::UsdStage::Open(layer);
EXPECT_THAT(stage, testing::NotNull());
return stage;
}
template <>
void ExpectAttributeEqual<pxr::SdfAssetPath>(pxr::UsdStageRefPtr stage,
pxr::SdfPath path,
+2
View File
@@ -73,6 +73,8 @@ pxr::SdfLayerRefPtr LoadLayer(
const std::string& xml,
const pxr::SdfFileFormat::FileFormatArguments& args = {});
pxr::UsdStageRefPtr OpenStageWithPhysics(const std::string& xml);
template <typename T>
void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, pxr::SdfPath path,
const T& value) {