diff --git a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc index f1c8e5ee..44c8481f 100644 --- a/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc +++ b/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc @@ -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(geom_size[0]), - static_cast(geom_size[1]), - static_cast(geom_size[2])); + pxr::GfVec3f scale(static_cast(size[0]), static_cast(size[1]), + static_cast(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(-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(-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(geom_size[0] * 2), - static_cast(geom_size[1] * 2), - static_cast(geom_size[2] * 2)}; + pxr::GfVec3f scale = {static_cast(size[0] * 2), + static_cast(size[1] * 2), + static_cast(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{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{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); diff --git a/src/experimental/usd/plugins/mjcf/utils.cc b/src/experimental/usd/plugins/mjcf/utils.cc index 880cce55..e919a672 100644 --- a/src/experimental/usd/plugins/mjcf/utils.cc +++ b/src/experimental/usd/plugins/mjcf/utils.cc @@ -26,6 +26,7 @@ #include #include #include +#include namespace { template @@ -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 diff --git a/src/experimental/usd/plugins/mjcf/utils.h b/src/experimental/usd/plugins/mjcf/utils.h index 4a2643ec..9f41a697 100644 --- a/src/experimental/usd/plugins/mjcf/utils.h +++ b/src/experimental/usd/plugins/mjcf/utils.h @@ -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 void SetField(pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& field_path, diff --git a/test/experimental/usd/plugins/mjcf/fixture.h b/test/experimental/usd/plugins/mjcf/fixture.h index 64985fae..adfddbc4 100644 --- a/test/experimental/usd/plugins/mjcf/fixture.h +++ b/test/experimental/usd/plugins/mjcf/fixture.h @@ -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()); \ + } + #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); 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 8ca9bdba..df04db21 100644 --- a/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc +++ b/test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc @@ -36,9 +36,14 @@ #include #include #include +#include +#include +#include #include #include #include +#include +#include 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"( + + + + + + + + + + + + + )"; + +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