From 8e7457d7c6c1cbda42ec39d2297d4f90b3396040 Mon Sep 17 00:00:00 2001 From: Google DeepMind Date: Fri, 11 Apr 2025 08:18:11 -0700 Subject: [PATCH] Initial support for writing UsdPhysics data starting with UsdPhysicsRigidbodyAPI support. The SDF_FORMAT_ARGS `usdMjcfToggleUsdPhysics` is introduced to toggle writing UsdPhysics or not when converting. Off by default PiperOrigin-RevId: 746467942 Change-Id: I26cb9a4c4ae2eef96163dab967a74f14940a2a6d --- .../usd/plugins/mjcf/mjcf_file_format.cc | 14 +++- .../usd/plugins/mjcf/mjcf_file_format.h | 9 ++- .../usd/plugins/mjcf/mujoco_to_usd.cc | 22 ++++-- .../usd/plugins/mjcf/mujoco_to_usd.h | 4 +- test/experimental/usd/plugins/mjcf/fixture.cc | 6 +- test/experimental/usd/plugins/mjcf/fixture.h | 15 ++++- .../usd/plugins/mjcf/mjcf_file_format_test.cc | 67 +++++++++++++++++++ 7 files changed, 122 insertions(+), 15 deletions(-) diff --git a/src/experimental/usd/plugins/mjcf/mjcf_file_format.cc b/src/experimental/usd/plugins/mjcf/mjcf_file_format.cc index 8c5858c9..904c84fd 100644 --- a/src/experimental/usd/plugins/mjcf/mjcf_file_format.cc +++ b/src/experimental/usd/plugins/mjcf/mjcf_file_format.cc @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -145,9 +146,18 @@ bool UsdMjcfFileFormat::CanRead(const std::string &filePath) const { } bool UsdMjcfFileFormat::ReadImpl(pxr::SdfLayer *layer, mjSpec *spec) const { - auto data = InitData(layer->GetFileFormatArguments()); + auto args = layer->GetFileFormatArguments(); - auto success = mujoco::usd::WriteSpecToData(spec, data); + bool toggleUsdPhysics = false; + const auto it = + args.find(UsdMjcfFileFormatTokens->ToggleUsdPhysicsArg.GetString()); + if (it != args.end()) { + toggleUsdPhysics = pxr::TfUnstringify(it->second); + } + + auto data = InitData(args); + + auto success = mujoco::usd::WriteSpecToData(spec, data, toggleUsdPhysics); mj_deleteSpec(spec); if (!success) { return false; diff --git a/src/experimental/usd/plugins/mjcf/mjcf_file_format.h b/src/experimental/usd/plugins/mjcf/mjcf_file_format.h index c5776ec4..5bc72337 100644 --- a/src/experimental/usd/plugins/mjcf/mjcf_file_format.h +++ b/src/experimental/usd/plugins/mjcf/mjcf_file_format.h @@ -27,10 +27,15 @@ PXR_NAMESPACE_OPEN_SCOPE +// clang-format off // The Id should realistically be mjcf, but the id and extension need to match. // So near term it just assumes the only .xml file we would import is MJCF. -#define USD_MJCF_FILE_FORMAT_TOKENS \ - ((Id, "xml"))((Version, "1.0"))((Target, "usd")) +#define USD_MJCF_FILE_FORMAT_TOKENS \ + ((Id, "xml")) \ + ((Version, "1.0")) \ + ((Target, "usd")) \ + ((ToggleUsdPhysicsArg, "usdMjcfToggleUsdPhysics")) +// clang-format on TF_DECLARE_PUBLIC_TOKENS(UsdMjcfFileFormatTokens, USD_MJCF_FILE_FORMAT_TOKENS); diff --git a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc index 44c8481f..effc26a3 100644 --- a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc +++ b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc @@ -49,6 +49,7 @@ #include #include #include +#include #include #include @@ -149,7 +150,7 @@ class ModelWriter { } ~ModelWriter() { mj_deleteModel(model_); } - void Write() { + void Write(bool write_physics) { // Create top level class holder. class_path_ = CreateClassSpec(data_, pxr::SdfPath::AbsoluteRootPath(), pxr::TfToken("__class__")); @@ -173,7 +174,7 @@ class ModelWriter { // Author mesh scope + mesh prims to be referenced. WriteMeshes(); WriteMaterials(); - WriteBodies(); + WriteBodies(write_physics); } private: @@ -977,7 +978,7 @@ class ModelWriter { } } - void WriteBody(mjsBody *body) { + void WriteBody(mjsBody *body, bool write_physics) { int body_id = mjs_getId(body->element); pxr::SdfPath parent_path = CreateParentIfNotExists(body, body_paths_[kWorldIndex], data_); @@ -990,6 +991,12 @@ class ModelWriter { // bodies subcomponents. SetPrimKind(data_, body_path, pxr::KindTokens->subcomponent); + // Apply the PhysicsRigidBodyAPI schema if we are writing physics. + if (write_physics) { + ApplyApiSchema(data_, body_path, + pxr::UsdPhysicsTokens->PhysicsRigidBodyAPI); + } + // Create classes if necessary mjsDefault *spec_default = mjs_getDefault(body->element); @@ -1028,14 +1035,14 @@ class ModelWriter { body_paths_[body_id] = body_path; } - void WriteBodies() { + void WriteBodies(bool write_physics) { mjsBody *body = mjs_asBody(mjs_firstElement(spec_, mjOBJ_BODY)); while (body) { // Only write a rigidbody if we are not the world body. // We fall through since the world body might have static // geom children. if (mjs_getId(body->element) != kWorldIndex) { - WriteBody(body); + WriteBody(body, write_physics); } WriteSites(body); WriteGeoms(body); @@ -1061,7 +1068,8 @@ class ModelWriter { namespace mujoco { namespace usd { -bool WriteSpecToData(mjSpec *spec, pxr::SdfAbstractDataRefPtr &data) { +bool WriteSpecToData(mjSpec *spec, pxr::SdfAbstractDataRefPtr &data, + bool write_physics) { // Create pseudo root first. data->CreateSpec(pxr::SdfPath::AbsoluteRootPath(), pxr::SdfSpecTypePseudoRoot); @@ -1072,7 +1080,7 @@ bool WriteSpecToData(mjSpec *spec, pxr::SdfAbstractDataRefPtr &data) { return false; } - ModelWriter(spec, model, data).Write(); + ModelWriter(spec, model, data).Write(write_physics); return true; } diff --git a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h index 4c5934c6..9c0db6cf 100644 --- a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h +++ b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.h @@ -25,7 +25,9 @@ namespace usd { // Args: // spec: mjSpec built programmatically or via parsed XML. // data: SdfAbstractDataRefPtr that will be written to. -bool WriteSpecToData(mjSpec* spec, pxr::SdfAbstractDataRefPtr& data); +// write_physics: Whether to write physics data. +bool WriteSpecToData(mjSpec* spec, pxr::SdfAbstractDataRefPtr& data, + bool write_physics); } // namespace usd } // namespace mujoco diff --git a/test/experimental/usd/plugins/mjcf/fixture.cc b/test/experimental/usd/plugins/mjcf/fixture.cc index 19cd8e6e..74063da2 100644 --- a/test/experimental/usd/plugins/mjcf/fixture.cc +++ b/test/experimental/usd/plugins/mjcf/fixture.cc @@ -32,9 +32,11 @@ namespace mujoco { using pxr::SdfPath; -pxr::SdfLayerRefPtr LoadLayer(const std::string& xml) { +pxr::SdfLayerRefPtr LoadLayer( + const std::string& xml, + const pxr::SdfFileFormat::FileFormatArguments& args) { auto layer = pxr::SdfLayer::CreateAnonymous( - "test_layer", pxr::SdfFileFormat::FindByExtension("xml")); + "test_layer", pxr::SdfFileFormat::FindByExtension("xml"), args); layer->ImportFromString(xml); EXPECT_THAT(layer, testing::NotNull()); return layer; diff --git a/test/experimental/usd/plugins/mjcf/fixture.h b/test/experimental/usd/plugins/mjcf/fixture.h index adfddbc4..5c125930 100644 --- a/test/experimental/usd/plugins/mjcf/fixture.h +++ b/test/experimental/usd/plugins/mjcf/fixture.h @@ -21,6 +21,7 @@ #include "test/fixture.h" #include #include +#include #include #include #include @@ -34,6 +35,16 @@ EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).IsA()); \ } +#define EXPECT_PRIM_API_APPLIED(stage, path, api) \ + { \ + EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).HasAPI()); \ + } + +#define EXPECT_PRIM_API_NOT_APPLIED(stage, path, api) \ + { \ + EXPECT_FALSE((stage)->GetPrimAtPath(SdfPath(path)).HasAPI()); \ + } + #define EXPECT_PRIM_KIND(stage, path, kind) \ { \ pxr::TfToken prim_kind; \ @@ -51,7 +62,9 @@ } namespace mujoco { -pxr::SdfLayerRefPtr LoadLayer(const std::string& xml); +pxr::SdfLayerRefPtr LoadLayer( + const std::string& xml, + const pxr::SdfFileFormat::FileFormatArguments& args = {}); template void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, const char* path, diff --git a/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc b/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc index df04db21..0b3dca5d 100644 --- a/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc +++ b/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc @@ -44,6 +44,7 @@ #include #include #include +#include PXR_NAMESPACE_OPEN_SCOPE // clang-format off @@ -434,6 +435,72 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestSitePrimsPurpose) { EXPECT_PRIM_PURPOSE(stage, "/test/ball/ball/ellipsoid_site", pxr::UsdGeomTokens->guide); } +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsToggleSdfFormatArg) { + std::string xml_path = GetTestDataFilePath(kMeshObjPath); + + // Test that the default is no physics. + auto stage_no_physics = pxr::UsdStage::Open(xml_path); + EXPECT_THAT(stage_no_physics, testing::NotNull()); + EXPECT_PRIM_VALID(stage_no_physics, "/mesh_test/test_body/test_body"); + EXPECT_PRIM_API_NOT_APPLIED(stage_no_physics, + "/mesh_test/test_body/test_body", + pxr::UsdPhysicsRigidBodyAPI); + + // Then test that the physics flag enables physics. + std::string xml_path_physics_flag = + xml_path + ":SDF_FORMAT_ARGS:usdMjcfToggleUsdPhysics=true"; + auto stage_with_physics = pxr::UsdStage::Open(xml_path_physics_flag); + EXPECT_THAT(stage_with_physics, testing::NotNull()); + + EXPECT_PRIM_VALID(stage_with_physics, "/mesh_test/test_body/test_body"); + EXPECT_PRIM_API_APPLIED(stage_with_physics, "/mesh_test/test_body/test_body", + pxr::UsdPhysicsRigidBodyAPI); +} + +TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsRigidBody) { + static constexpr char kXml[] = R"( + + + + + + + + + + + )"; + + pxr::SdfFileFormat::FileFormatArguments args; + args["usdMjcfToggleUsdPhysics"] = "true"; + pxr::SdfLayerRefPtr layer = LoadLayer(kXml, args); + auto stage = pxr::UsdStage::Open(layer); + + EXPECT_THAT(stage, testing::NotNull()); + EXPECT_PRIM_VALID(stage, "/physics_test"); + EXPECT_PRIM_VALID(stage, "/physics_test/test_body"); + EXPECT_PRIM_VALID(stage, "/physics_test/test_body/test_body"); + // USD does not allow nested rigidbodies so we put them as siblings to the + // first body in the hierarchy. + EXPECT_PRIM_VALID(stage, "/physics_test/test_body/test_body_2"); + + // The parent containing the body should not have the RigidBodyAPI applied. + EXPECT_PRIM_API_NOT_APPLIED(stage, "/physics_test/test_body", + pxr::UsdPhysicsRigidBodyAPI); + + EXPECT_PRIM_API_APPLIED(stage, "/physics_test/test_body/test_body", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_APPLIED(stage, "/physics_test/test_body/test_body_2", + pxr::UsdPhysicsRigidBodyAPI); + + // Geoms should not have RigidBodyAPI applied either. + EXPECT_PRIM_API_NOT_APPLIED(stage, + "/physics_test/test_body/test_body/test_geom", + pxr::UsdPhysicsRigidBodyAPI); + EXPECT_PRIM_API_NOT_APPLIED(stage, + "/physics_test/test_body/test_body_2/test_geom_2", + pxr::UsdPhysicsRigidBodyAPI); +} } // namespace } // namespace mujoco