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

Add basic site support to Mjcf SdfFileFormat plugin.

Browse Source 
We will still need to tag these with a custom API schema to ease round-tripping but that will follow creation of such a schema.

PiperOrigin-RevId: 746421793
Change-Id: I2df2209b4bc90d763718fb4bea8158b8598755c6
This commit is contained in:
Sam Haves
2025-04-11 05:26:59 -07:00
committed by Copybara-Service
parent 440ce5eb7a
commit 446a8f1742
5 changed files with 214 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
+136 -55
View File
@@ -122,6 +122,7 @@ using mujoco::usd::SetAttributeMetadata;
using mujoco::usd::SetLayerMetadata;
using mujoco::usd::SetPrimKind;
using mujoco::usd::SetPrimMetadata;
using mujoco::usd::SetPrimPurpose;
pxr::GfMatrix4d MujocoPosQuatToTransform(double *pos, double *quat) {
pxr::GfQuatd quaternion = pxr::GfQuatd::GetIdentity();
@@ -598,34 +599,57 @@ class ModelWriter {
return subcomponent_path;
}
pxr::SdfPath WriteBoxGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
pxr::SdfPath WriteSiteGeom(const mjsSite *site,
const pxr::SdfPath &body_path) {
auto name =
GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Cube, body_path);
GetAvailablePrimName(*site->name, pxr::UsdGeomTokens->Cube, body_path);
int site_idx = mjs_getId(site->element);
const mjtNum *size = &model_->site_size[site_idx * 3];
pxr::SdfPath site_path;
switch (site->type) {
case mjGEOM_BOX:
site_path = WriteBox(name, size, body_path);
break;
case mjGEOM_SPHERE:
site_path = WriteSphere(name, size, body_path);
break;
case mjGEOM_CAPSULE:
site_path = WriteCapsule(name, size, body_path);
break;
case mjGEOM_CYLINDER:
site_path = WriteCylinder(name, size, body_path);
break;
case mjGEOM_ELLIPSOID:
site_path = WriteEllipsoid(name, size, body_path);
break;
default:
break;
}
return site_path;
}
pxr::SdfPath WriteBox(const pxr::TfToken &name, const mjtNum *size,
const pxr::SdfPath &body_path) {
pxr::SdfPath box_path =
CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Cube);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
// MuJoCo uses half sizes.
pxr::SdfPath size_attr_path =
CreateAttributeSpec(data_, box_path, pxr::UsdGeomTokens->size,
pxr::SdfValueTypeNames->Float);
pxr::GfVec3f scale(static_cast<float>(geom_size[0]),
static_cast<float>(geom_size[1]),
static_cast<float>(geom_size[2]));
pxr::GfVec3f scale(static_cast<float>(size[0]), static_cast<float>(size[1]),
static_cast<float>(size[2]));
SetAttributeDefault(data_, size_attr_path, 2.0);
pxr::SdfPath extent_attr_path =
CreateAttributeSpec(data_, box_path, pxr::UsdGeomTokens->extent,
pxr::SdfValueTypeNames->Float3Array);
SetAttributeDefault(
data_, extent_attr_path,
pxr::VtArray<pxr::GfVec3f>({
pxr::GfVec3f(-geom_size[0], -geom_size[1], -geom_size[2]),
pxr::GfVec3f(geom_size[0], geom_size[1], geom_size[2]),
}));
SetAttributeDefault(data_, extent_attr_path,
pxr::VtArray<pxr::GfVec3f>({
pxr::GfVec3f(-size[0], -size[1], -size[2]),
pxr::GfVec3f(size[0], size[1], size[2]),
}));
WriteScaleXformOp(box_path, scale);
WriteXformOpOrder(box_path,
@@ -633,65 +657,80 @@ class ModelWriter {
return box_path;
}
pxr::SdfPath WriteCapsuleGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Capsule,
body_path);
pxr::SdfPath capsule_path =
CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Capsule);
pxr::SdfPath WriteBoxGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name =
GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Cube, body_path);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
return WriteBox(name, geom_size, body_path);
}
pxr::SdfPath WriteCapsule(const pxr::TfToken name, const mjtNum *size,
const pxr::SdfPath &body_path) {
pxr::SdfPath capsule_path =
CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Capsule);
// MuJoCo uses half sizes.
pxr::SdfPath radius_attr_path =
CreateAttributeSpec(data_, capsule_path, pxr::UsdGeomTokens->radius,
pxr::SdfValueTypeNames->Float);
SetAttributeDefault(data_, radius_attr_path, geom_size[0] * 2);
SetAttributeDefault(data_, radius_attr_path, size[0] * 2);
pxr::SdfPath height_attr_path =
CreateAttributeSpec(data_, capsule_path, pxr::UsdGeomTokens->height,
pxr::SdfValueTypeNames->Float);
SetAttributeDefault(data_, height_attr_path, geom_size[1] * 2);
SetAttributeDefault(data_, height_attr_path, size[1] * 2);
return capsule_path;
}
pxr::SdfPath WriteCapsuleGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Capsule,
body_path);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
return WriteCapsule(name, geom_size, body_path);
}
pxr::SdfPath WriteCylinder(const pxr::TfToken name, const mjtNum *size,
const pxr::SdfPath &body_path) {
pxr::SdfPath cylinder_path =
CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Cylinder);
// MuJoCo uses half sizes.
pxr::SdfPath radius_attr_path =
CreateAttributeSpec(data_, cylinder_path, pxr::UsdGeomTokens->radius,
pxr::SdfValueTypeNames->Float);
SetAttributeDefault(data_, radius_attr_path, size[0] * 2);
pxr::SdfPath height_attr_path =
CreateAttributeSpec(data_, cylinder_path, pxr::UsdGeomTokens->height,
pxr::SdfValueTypeNames->Float);
SetAttributeDefault(data_, height_attr_path, size[1] * 2);
return cylinder_path;
}
pxr::SdfPath WriteCylinderGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Cylinder,
body_path);
pxr::SdfPath cylinder_path =
CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Cylinder);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
// MuJoCo uses half sizes.
pxr::SdfPath radius_attr_path =
CreateAttributeSpec(data_, cylinder_path, pxr::UsdGeomTokens->radius,
pxr::SdfValueTypeNames->Float);
SetAttributeDefault(data_, radius_attr_path, geom_size[0] * 2);
pxr::SdfPath height_attr_path =
CreateAttributeSpec(data_, cylinder_path, pxr::UsdGeomTokens->height,
pxr::SdfValueTypeNames->Float);
SetAttributeDefault(data_, height_attr_path, geom_size[1] * 2);
return cylinder_path;
return WriteCylinder(name, geom_size, body_path);
}
pxr::SdfPath WriteEllipsoidGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Sphere,
body_path);
pxr::SdfPath WriteEllipsoid(const pxr::TfToken name, const mjtNum *size,
const pxr::SdfPath &body_path) {
pxr::SdfPath ellipsoid_path =
CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Sphere);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
pxr::GfVec3f scale = {static_cast<float>(geom_size[0] * 2),
static_cast<float>(geom_size[1] * 2),
static_cast<float>(geom_size[2] * 2)};
pxr::GfVec3f scale = {static_cast<float>(size[0] * 2),
static_cast<float>(size[1] * 2),
static_cast<float>(size[2] * 2)};
// MuJoCo uses half sizes.
pxr::SdfPath radius_attr_path =
@@ -705,28 +744,61 @@ class ModelWriter {
return ellipsoid_path;
}
pxr::SdfPath WriteSphereGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
pxr::SdfPath WriteEllipsoidGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Sphere,
body_path);
pxr::SdfPath sphere_path =
CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Sphere);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
return WriteEllipsoid(name, geom_size, body_path);
}
pxr::SdfPath WriteSphere(const pxr::TfToken name, const mjtNum *size,
const pxr::SdfPath &body_path) {
pxr::SdfPath sphere_path =
CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Sphere);
// MuJoCo uses half sizes.
pxr::SdfPath radius_attr_path =
CreateAttributeSpec(data_, sphere_path, pxr::UsdGeomTokens->radius,
pxr::SdfValueTypeNames->Float);
SetAttributeDefault(data_, radius_attr_path, geom_size[0] * 2);
SetAttributeDefault(data_, radius_attr_path, size[0] * 2);
return sphere_path;
}
pxr::SdfPath WriteSphereGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Sphere,
body_path);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
return WriteSphere(name, geom_size, body_path);
}
void WriteSite(mjsSite *site, const mjsBody *body) {
const int body_id = mjs_getId(body->element);
const auto &body_path = body_paths_[body_id];
auto name =
GetAvailablePrimName(*site->name, pxr::UsdGeomTokens->Xform, body_path);
// Create a geom primitive and set its purpose to guide so it won't be
// rendered.
pxr::SdfPath site_path = WriteSiteGeom(site, body_path);
SetPrimPurpose(data_, site_path, pxr::UsdGeomTokens->guide);
int site_id = mjs_getId(site->element);
auto transform = MujocoPosQuatToTransform(&model_->site_pos[3 * site_id],
&model_->site_quat[4 * site_id]);
WriteTransformXformOp(site_path, transform);
PrependToXformOpOrder(
site_path, pxr::VtArray<pxr::TfToken>{kTokens->xformOpTransform});
}
void WriteGeom(mjsGeom *geom, const mjsBody *body) {
const int body_id = mjs_getId(body->element);
const auto &body_path = body_paths_[body_id];
auto name = GetAvailablePrimName(*geom->name, kTokens->geom, body_path);
pxr::SdfPath geom_path;
int geom_id = mjs_getId(geom->element);
@@ -798,6 +870,14 @@ class ModelWriter {
geom_path, pxr::VtArray<pxr::TfToken>{kTokens->xformOpTransform});
}
void WriteSites(mjsBody *body) {
mjsSite *site = mjs_asSite(mjs_firstChild(body, mjOBJ_SITE, false));
while (site) {
WriteSite(site, body);
site = mjs_asSite(mjs_nextChild(body, site->element, false));
}
}
void WriteGeoms(mjsBody *body) {
mjsGeom *geom = mjs_asGeom(mjs_firstChild(body, mjOBJ_GEOM, false));
while (geom) {
@@ -957,6 +1037,7 @@ class ModelWriter {
if (mjs_getId(body->element) != kWorldIndex) {
WriteBody(body);
}
WriteSites(body);
WriteGeoms(body);
WriteCameras(body);
WriteLights(body);
src/experimental/usd/plugins/mjcf/utils.cc
+9
View File
@@ -26,6 +26,7 @@
#include <pxr/usd/sdf/types.h>
#include <pxr/usd/sdf/valueTypeName.h>
#include <pxr/usd/usd/tokens.h>
#include <pxr/usd/usdGeom/tokens.h>
namespace {
template <typename T>
@@ -195,5 +196,13 @@ void SetPrimKind(pxr::SdfAbstractDataRefPtr& data,
SetPrimMetadata(data, prim_path, pxr::TfToken("kind"), kind);
}
void SetPrimPurpose(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path, pxr::TfToken purpose) {
const pxr::SdfPath attr_path = CreateAttributeSpec(
data, prim_path, pxr::UsdGeomTokens->purpose,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
SetAttributeDefault(data, attr_path, purpose);
}
} // namespace usd
} // namespace mujoco
src/experimental/usd/plugins/mjcf/utils.h
+3
View File
@@ -71,6 +71,9 @@ void ApplyApiSchema(pxr::SdfAbstractDataRefPtr& data,
void SetPrimKind(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path, pxr::TfToken kind);
void SetPrimPurpose(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path, pxr::TfToken purpose);
// Set the value specified by key on any field at field_path.
template <typename T>
void SetField(pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& field_path,
test/experimental/usd/plugins/mjcf/fixture.h
+14
View File
@@ -29,12 +29,26 @@
#define EXPECT_PRIM_VALID(stage, path) \
EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).IsValid());
#define EXPECT_PRIM_IS_A(stage, path, type) \
{ \
EXPECT_TRUE((stage)->GetPrimAtPath(SdfPath(path)).IsA<type>()); \
}
#define EXPECT_PRIM_KIND(stage, path, kind) \
{ \
pxr::TfToken prim_kind; \
pxr::UsdModelAPI::Get(stage, SdfPath(path)).GetKind(&prim_kind); \
EXPECT_EQ(kind, prim_kind); \
}
#define EXPECT_PRIM_PURPOSE(stage, path, purpose) \
{ \
pxr::TfToken prim_purpose; \
pxr::UsdGeomImageable::Get(stage, SdfPath(path)) \
.GetPurposeAttr() \
.Get(&prim_purpose); \
EXPECT_EQ(prim_purpose, purpose); \
}
namespace mujoco {
pxr::SdfLayerRefPtr LoadLayer(const std::string& xml);
test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc
+52
View File
@@ -36,9 +36,14 @@
#include <pxr/usd/usd/modelAPI.h>
#include <pxr/usd/usd/prim.h>
#include <pxr/usd/usd/stage.h>
#include <pxr/usd/usdGeom/capsule.h>
#include <pxr/usd/usdGeom/cube.h>
#include <pxr/usd/usdGeom/cylinder.h>
#include <pxr/usd/usdGeom/mesh.h>
#include <pxr/usd/usdGeom/primvar.h>
#include <pxr/usd/usdGeom/primvarsAPI.h>
#include <pxr/usd/usdGeom/sphere.h>
#include <pxr/usd/usdGeom/tokens.h>
PXR_NAMESPACE_OPEN_SCOPE
// clang-format off
@@ -383,5 +388,52 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestKindAuthoring) {
EXPECT_PRIM_KIND(stage, "/test/root/tet", pxr::KindTokens->subcomponent);
}
static constexpr char kSiteXml[] = R"(
<mujoco model="test">
<worldbody>
<site type="box" name="box_site"/>
<body name="ball">
<site type="sphere" name="sphere_site"/>
<site type="capsule" name="capsule_site"/>
<site type="cylinder" name="cylinder_site"/>
<site type="ellipsoid" name="ellipsoid_site"/>
<geom type="sphere" size="1 1 1"/>
</body>
</worldbody>
</mujoco>
)";
TEST_F(MjcfSdfFileFormatPluginTest, TestSitePrimsAuthored) {
pxr::SdfLayerRefPtr layer = LoadLayer(kSiteXml);
auto stage = pxr::UsdStage::Open(layer);
EXPECT_PRIM_VALID(stage, "/test/box_site");
EXPECT_PRIM_IS_A(stage, "/test/box_site", pxr::UsdGeomCube);
EXPECT_PRIM_VALID(stage, "/test/ball/ball/sphere_site");
EXPECT_PRIM_IS_A(stage, "/test/ball/ball/sphere_site", pxr::UsdGeomSphere);
EXPECT_PRIM_VALID(stage, "/test/ball/ball/capsule_site");
EXPECT_PRIM_IS_A(stage, "/test/ball/ball/capsule_site", pxr::UsdGeomCapsule);
EXPECT_PRIM_VALID(stage, "/test/ball/ball/cylinder_site");
EXPECT_PRIM_IS_A(stage, "/test/ball/ball/cylinder_site",
pxr::UsdGeomCylinder);
EXPECT_PRIM_VALID(stage, "/test/ball/ball/ellipsoid_site");
EXPECT_PRIM_IS_A(stage, "/test/ball/ball/ellipsoid_site", pxr::UsdGeomSphere);
}
TEST_F(MjcfSdfFileFormatPluginTest, TestSitePrimsPurpose) {
pxr::SdfLayerRefPtr layer = LoadLayer(kSiteXml);
auto stage = pxr::UsdStage::Open(layer);
EXPECT_PRIM_PURPOSE(stage, "/test/box_site", pxr::UsdGeomTokens->guide);
EXPECT_PRIM_PURPOSE(stage, "/test/ball/ball/sphere_site",
pxr::UsdGeomTokens->guide);
EXPECT_PRIM_PURPOSE(stage, "/test/ball/ball/capsule_site",
pxr::UsdGeomTokens->guide);
EXPECT_PRIM_PURPOSE(stage, "/test/ball/ball/cylinder_site",
pxr::UsdGeomTokens->guide);
EXPECT_PRIM_PURPOSE(stage, "/test/ball/ball/ellipsoid_site",
pxr::UsdGeomTokens->guide);
}
} // namespace
} // namespace mujoco