Add write_physics_ as a state variable on ModelWriter.

This way we don't have to worry about passing the write_physics param everywhere.

PiperOrigin-RevId: 764394280
Change-Id: I2c879ae57bdc30f50189960864097c932f967330
This commit is contained in:
Robin Alazard
2025-05-28 13:31:59 -07:00
committed by Copybara-Service
parent fc69ef1084
commit 6b919162d1
@@ -166,6 +166,9 @@ class ModelWriter {
// Create the world body.
body_paths_[kWorldIndex] = WriteWorldBody(kWorldIndex);
// Set working parameters.
write_physics_ = write_physics;
SetLayerMetadata(data_, pxr::SdfFieldKeys->Documentation,
"Generated by mujoco model writer.");
// Mujoco is Z up by default.
@@ -181,9 +184,9 @@ class ModelWriter {
WritePhysicsScene();
// Author mesh scope + mesh prims to be referenced.
WriteMeshes(write_physics);
WriteMeshes();
WriteMaterials();
WriteBodies(write_physics);
WriteBodies();
}
private:
@@ -199,6 +202,8 @@ class ModelWriter {
std::vector<pxr::SdfPath> body_paths_;
// Mapping from mesh names to Mesh prim path.
std::unordered_map<std::string, pxr::SdfPath> mesh_paths_;
// Whether to write physics data.
bool write_physics_ = false;
// Given a name index and a parent prim path this returns a
// token such that appending it to the parent prim path does not
@@ -299,8 +304,7 @@ class ModelWriter {
SetAttributeDefault(data_, xform_op_order_path, new_order);
}
void WriteMesh(const mjsMesh *mesh, const pxr::SdfPath &parent_path,
bool write_physics) {
void WriteMesh(const mjsMesh *mesh, const pxr::SdfPath &parent_path) {
auto name = GetAvailablePrimName(*mesh->name, pxr::UsdGeomTokens->Mesh,
parent_path);
pxr::SdfPath subcomponent_path =
@@ -310,7 +314,7 @@ class ModelWriter {
pxr::UsdGeomTokens->Mesh);
mesh_paths_[*mesh->name] = subcomponent_path;
if (write_physics) {
if (write_physics_) {
ApplyApiSchema(data_, mesh_path, MjcPhysicsTokens->MeshCollisionAPI);
pxr::TfToken inertia = MjcPhysicsTokens->legacy;
@@ -602,7 +606,7 @@ class ModelWriter {
}
}
void WriteMeshes(bool write_physics) {
void WriteMeshes() {
// Create a scope for the meshes to keep things organized
pxr::SdfPath scope_path =
CreatePrimSpec(data_, body_paths_[kWorldIndex], kTokens->meshScope,
@@ -614,7 +618,7 @@ class ModelWriter {
mjsMesh *mesh = mjs_asMesh(mjs_firstElement(spec_, mjOBJ_MESH));
while (mesh) {
WriteMesh(mesh, scope_path, write_physics);
WriteMesh(mesh, scope_path);
mesh = mjs_asMesh(mjs_nextElement(spec_, mesh->element));
}
}
@@ -1016,7 +1020,7 @@ class ModelWriter {
site_path, pxr::VtArray<pxr::TfToken>{kTokens->xformOpTransform});
}
void WriteGeom(mjsGeom *geom, const mjsBody *body, bool write_physics) {
void WriteGeom(mjsGeom *geom, const mjsBody *body) {
const int body_id = mjs_getId(body->element);
const auto &body_path = body_paths_[body_id];
@@ -1052,8 +1056,8 @@ class ModelWriter {
// Apply the physics schemas if we are writing physics and the
// geom participates in collisions.
if (write_physics && (model_->geom_contype[geom_id] != 0 ||
model_->geom_conaffinity[geom_id] != 0)) {
if (write_physics_ && (model_->geom_contype[geom_id] != 0 ||
model_->geom_conaffinity[geom_id] != 0)) {
ApplyApiSchema(data_, geom_path,
pxr::UsdPhysicsTokens->PhysicsCollisionAPI);
ApplyApiSchema(data_, geom_path, MjcPhysicsTokens->CollisionAPI);
@@ -1178,10 +1182,10 @@ class ModelWriter {
}
}
void WriteGeoms(mjsBody *body, bool write_physics) {
void WriteGeoms(mjsBody *body) {
mjsGeom *geom = mjs_asGeom(mjs_firstChild(body, mjOBJ_GEOM, false));
while (geom) {
WriteGeom(geom, body, write_physics);
WriteGeom(geom, body);
geom = mjs_asGeom(mjs_nextChild(body, geom->element, false));
}
}
@@ -1268,7 +1272,7 @@ class ModelWriter {
}
}
void WriteBody(mjsBody *body, bool write_physics) {
void WriteBody(mjsBody *body) {
int body_id = mjs_getId(body->element);
// This should be safe as we process parent bodies before children.
mjsBody *parent = mjs_getParent(body->element);
@@ -1286,7 +1290,7 @@ class ModelWriter {
SetPrimKind(data_, body_path, kind);
// Apply the PhysicsRigidBodyAPI schema if we are writing physics.
if (write_physics) {
if (write_physics_) {
// If the body had a mass specified then it must have either inertia or
// fullinertia specified per inertia element XML documentation.
// Therefore it is sufficient to check if the mass is non-zero to see if
@@ -1373,17 +1377,17 @@ class ModelWriter {
body_paths_[body_id] = body_path;
}
void WriteBodies(bool write_physics) {
void WriteBodies() {
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, write_physics);
WriteBody(body);
}
WriteSites(body);
WriteGeoms(body, write_physics);
WriteGeoms(body);
WriteCameras(body);
WriteLights(body);
body = mjs_asBody(mjs_nextElement(spec_, body->element));