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

Initial support for writing UsdPhysics data starting with UsdPhysicsRigidbodyAPI support.

Browse Source 
The SDF_FORMAT_ARGS `usdMjcfToggleUsdPhysics` is introduced to toggle writing UsdPhysics or not when converting. Off by default

PiperOrigin-RevId: 746467942
Change-Id: I26cb9a4c4ae2eef96163dab967a74f14940a2a6d
This commit is contained in:
Google DeepMind
2025-04-11 08:18:11 -07:00
committed by Copybara-Service
parent 006c18b1d8
commit 8e7457d7c6
7 changed files with 122 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/usd/plugins/mjcf/mjcf_file_format.cc
+12 -2
View File
@@ -26,6 +26,7 @@
#include <pxr/base/tf/pathUtils.h>
#include <pxr/base/tf/registryManager.h>
#include <pxr/base/tf/staticTokens.h>
#include <pxr/base/tf/stringUtils.h>
#include <pxr/base/tf/type.h>
#include <pxr/base/work/loops.h>
#include <pxr/pxr.h>
@@ -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<bool>(it->second);
}
auto data = InitData(args);
auto success = mujoco::usd::WriteSpecToData(spec, data, toggleUsdPhysics);
mj_deleteSpec(spec);
if (!success) {
return false;
src/experimental/usd/plugins/mjcf/mjcf_file_format.h
+7 -2
View File
@@ -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);
src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
+15 -7
View File
@@ -49,6 +49,7 @@
#include <pxr/usd/usdGeom/metrics.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdLux/tokens.h>
#include <pxr/usd/usdPhysics/tokens.h>
#include <pxr/usd/usdShade/tokens.h>
#include <pxr/usdImaging/usdImaging/tokens.h>
@@ -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;
}
src/experimental/usd/plugins/mjcf/mujoco_to_usd.h
+3 -1
View File
@@ -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
test/experimental/usd/plugins/mjcf/fixture.cc
+4 -2
View File
@@ -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;
test/experimental/usd/plugins/mjcf/fixture.h
+14 -1
View File
@@ -21,6 +21,7 @@
#include "test/fixture.h"
#include <pxr/usd/sdf/assetPath.h>
#include <pxr/usd/sdf/declareHandles.h>
#include <pxr/usd/sdf/fileFormat.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/usd/usd/common.h>
#include <pxr/usd/usd/modelAPI.h>
@@ -34,6 +35,16 @@
EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).IsA<type>()); \
}
#define EXPECT_PRIM_API_APPLIED(stage, path, api) \
{ \
EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).HasAPI<api>()); \
}
#define EXPECT_PRIM_API_NOT_APPLIED(stage, path, api) \
{ \
EXPECT_FALSE((stage)->GetPrimAtPath(SdfPath(path)).HasAPI<api>()); \
}
#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 <typename T>
void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, const char* path,
test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc
+67
View File
@@ -44,6 +44,7 @@
#include <pxr/usd/usdGeom/primvarsAPI.h>
#include <pxr/usd/usdGeom/sphere.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdPhysics/rigidBodyAPI.h>
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"(
<mujoco model="physics_test">
<worldbody>
<body name="test_body" pos="0 0 0">
<geom name="test_geom" type="sphere" size="1"/>
<body name="test_body_2" pos="2 0 0">
<geom name="test_geom_2" type="sphere" size="1"/>
</body>
</body>
</worldbody>
</mujoco>
)";
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