Add support for plane geometry.

PiperOrigin-RevId: 750994019
Change-Id: Ia3b085bed5b8e2fcdb00ffc69dd5551a7fa3c30d
This commit is contained in:
Google DeepMind
2025-04-24 07:53:58 -07:00
committed by Copybara-Service
parent 9e72874801
commit 8df3614781
3 changed files with 101 additions and 2 deletions
@@ -777,6 +777,43 @@ class ModelWriter {
return WriteSphere(name, geom_size, body_path);
}
pxr::SdfPath WritePlane(const pxr::TfToken &name, const mjtNum *size,
const pxr::SdfPath &body_path) {
pxr::SdfPath plane_path =
CreatePrimSpec(data_, body_path, name, pxr::UsdGeomTokens->Plane);
// MuJoCo uses half sizes.
// Note that UsdGeomPlane is infinite for simulation purposes but can have
// width/length for visualization, same as MuJoCo.
double width = size[0] * 2.0;
double length = size[1] * 2.0;
pxr::SdfPath width_attr_path =
CreateAttributeSpec(data_, plane_path, pxr::UsdGeomTokens->width,
pxr::SdfValueTypeNames->Double);
SetAttributeDefault(data_, width_attr_path, width);
pxr::SdfPath length_attr_path =
CreateAttributeSpec(data_, plane_path, pxr::UsdGeomTokens->length,
pxr::SdfValueTypeNames->Double);
SetAttributeDefault(data_, length_attr_path, length);
// MuJoCo plane is always a XY plane with +Z up.
// UsdGeomPlane is also a XY plane if axis is 'Z', which is default.
// So no need to set axis attribute explicitly.
return plane_path;
}
pxr::SdfPath WritePlaneGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name =
GetAvailablePrimName(*geom->name, pxr::UsdGeomTokens->Plane, body_path);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
return WritePlane(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];
@@ -804,6 +841,9 @@ class ModelWriter {
pxr::SdfPath geom_path;
int geom_id = mjs_getId(geom->element);
switch (geom->type) {
case mjGEOM_PLANE:
geom_path = WritePlaneGeom(geom, body_path);
break;
case mjGEOM_MESH:
geom_path = WriteMeshGeom(geom, body_path);
break;
+3 -2
View File
@@ -70,10 +70,11 @@ template <typename T>
void ExpectAttributeEqual(pxr::UsdStageRefPtr stage, const char* path,
const T& value) {
auto attr = stage->GetAttributeAtPath(pxr::SdfPath(path));
EXPECT_TRUE(attr.IsValid());
EXPECT_TRUE(attr.IsValid()) << "Attribute " << path << " is not valid";
T attr_value;
attr.Get(&attr_value);
EXPECT_EQ(attr_value, value);
EXPECT_EQ(attr_value, value) << "Attribute " << path << " has value "
<< attr_value << ". Expected: " << value;
}
// Specialization for SdfAssetPath, so that we can compare only the asset path
@@ -40,6 +40,7 @@
#include <pxr/usd/usdGeom/cube.h>
#include <pxr/usd/usdGeom/cylinder.h>
#include <pxr/usd/usdGeom/mesh.h>
#include <pxr/usd/usdGeom/plane.h>
#include <pxr/usd/usdGeom/primvar.h>
#include <pxr/usd/usdGeom/primvarsAPI.h>
#include <pxr/usd/usdGeom/sphere.h>
@@ -389,6 +390,62 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestKindAuthoring) {
EXPECT_PRIM_KIND(stage, "/test/root/tet", pxr::KindTokens->subcomponent);
}
TEST_F(MjcfSdfFileFormatPluginTest, TestGeomsPrims) {
static constexpr char kXml[] = R"(
<mujoco model="test">
<worldbody>
<geom type="plane" name="plane_geom" size="10 20 0.1"/>
<geom type="box" name="box_geom" size="10 20 30"/>
<geom type="sphere" name="sphere_geom" size="10 20 30"/>
<geom type="capsule" name="capsule_geom" size="10 20 30"/>
<geom type="cylinder" name="cylinder_geom" size="10 20 30"/>
<geom type="ellipsoid" name="ellipsoid_geom" size="10 20 30"/>
</worldbody>
</mujoco>
)";
pxr::SdfLayerRefPtr layer = LoadLayer(kXml);
auto stage = pxr::UsdStage::Open(layer);
// Note that all sizes are multiplied by 2 because Mujoco uses half sizes.
// Plane
EXPECT_PRIM_VALID(stage, "/test/plane_geom");
EXPECT_PRIM_IS_A(stage, "/test/plane_geom", pxr::UsdGeomPlane);
ExpectAttributeEqual(stage, "/test/plane_geom.width", 2 * 10.0);
ExpectAttributeEqual(stage, "/test/plane_geom.length", 2 * 20.0);
// Box
EXPECT_PRIM_VALID(stage, "/test/box_geom");
EXPECT_PRIM_IS_A(stage, "/test/box_geom", pxr::UsdGeomCube);
// Box is a special case, it uses a UsdGeomCube and scales it with
// xformOp:scale. The radius is always set to 2.
ExpectAttributeEqual(stage, "/test/box_geom.size", 2.0);
ExpectAttributeEqual(stage, "/test/box_geom.xformOp:scale",
pxr::GfVec3f(10.0, 20.0, 30.0));
// Sphere
EXPECT_PRIM_VALID(stage, "/test/sphere_geom");
EXPECT_PRIM_IS_A(stage, "/test/sphere_geom", pxr::UsdGeomSphere);
ExpectAttributeEqual(stage, "/test/sphere_geom.radius", 2 * 10.0);
// Capsule
EXPECT_PRIM_VALID(stage, "/test/capsule_geom");
EXPECT_PRIM_IS_A(stage, "/test/capsule_geom", pxr::UsdGeomCapsule);
ExpectAttributeEqual(stage, "/test/capsule_geom.radius", 2 * 10.0);
ExpectAttributeEqual(stage, "/test/capsule_geom.height", 2 * 20.0);
// Cylinder
EXPECT_PRIM_VALID(stage, "/test/cylinder_geom");
EXPECT_PRIM_IS_A(stage, "/test/cylinder_geom", pxr::UsdGeomCylinder);
ExpectAttributeEqual(stage, "/test/cylinder_geom.radius", 2 * 10.0);
ExpectAttributeEqual(stage, "/test/cylinder_geom.height", 2 * 20.0);
// Ellipsoid
EXPECT_PRIM_VALID(stage, "/test/ellipsoid_geom");
// Ellipsoid is a special case, it uses a UsdGeomSphere and scales it with
// xformOp:scale. The radius is always set to 1.
EXPECT_PRIM_IS_A(stage, "/test/ellipsoid_geom", pxr::UsdGeomSphere);
ExpectAttributeEqual(stage, "/test/ellipsoid_geom.radius", 1.0);
ExpectAttributeEqual(stage, "/test/ellipsoid_geom.xformOp:scale",
pxr::GfVec3f(2.0 * 10.0, 2.0 * 20.0, 2.0 * 30.0));
}
static constexpr char kSiteXml[] = R"(
<mujoco model="test">
<worldbody>
@@ -435,6 +492,7 @@ 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);