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

Modify mujoco_to_usd to operate on SdfLayer intead of SdfAbstractData.

Browse Source 
Since we don't do any on demand data loading for USD and always convert everything it should be safe to use the SdfLayer API instead of the SdfAbstractData API. This enables other codepaths such as the USD writer class to reuse the same logic to generate USD specification from an mjSpec.

PiperOrigin-RevId: 827431722
Change-Id: I60a581e0c9577164a4a66229ea9dcb8912afcf74
This commit is contained in:
Sam Haves
2025-11-03 04:15:31 -08:00
committed by Copybara-Service
parent 3a7aa84e53
commit 9319956a16
5 changed files with 621 additions and 705 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/usd/plugins/mjcf/mjcf_file_format.cc
+5 -2
View File
@@ -36,6 +36,7 @@
#include <pxr/usd/ar/asset.h>
#include <pxr/usd/ar/resolvedPath.h>
#include <pxr/usd/ar/resolver.h>
#include <pxr/usd/sdf/changeBlock.h>
#include <pxr/usd/sdf/declareHandles.h>
#include <pxr/usd/sdf/fileFormat.h>
#include <pxr/usd/sdf/layer.h>
@@ -236,13 +237,15 @@ bool UsdMjcfFileFormat::ReadImpl(pxr::SdfLayer *layer, mjSpec *spec) const {
auto args = layer->GetFileFormatArguments();
auto data = InitData(args);
auto success = mujoco::usd::WriteSpecToData(spec, data);
pxr::SdfChangeBlock block;
pxr::SdfLayerRefPtr spec_layer = pxr::SdfLayer::CreateAnonymous();
auto success = mujoco::usd::WriteSpecToData(spec, spec_layer);
mj_deleteSpec(spec);
if (!success) {
return false;
}
_SetLayerData(layer, data);
layer->TransferContent(spec_layer);
return true;
}
src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
+514 -516
View File
File diff suppressed because it is too large Load Diff
src/experimental/usd/plugins/mjcf/mujoco_to_usd.h
+4 -4
View File
@@ -16,16 +16,16 @@
#define MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_MUJOCO_TO_USD_H_
#include <mujoco/mujoco.h>
#include <pxr/usd/sdf/abstractData.h>
#include <pxr/usd/sdf/declareHandles.h>
namespace mujoco {
namespace usd {
// Given an mjSpec, write it to a SdfAbstractData.
// Given an mjSpec, writes a USD representation of the data to the given layer.
//
// Args:
// spec: mjSpec built programmatically or via parsed XML.
// data: SdfAbstractDataRefPtr that will be written to.
bool WriteSpecToData(mjSpec* spec, pxr::SdfAbstractDataRefPtr& data);
// layer: SdfLayerRefPtr that will be written to.
bool WriteSpecToData(mjSpec* spec, pxr::SdfLayerRefPtr layer);
} // namespace usd
} // namespace mujoco
src/experimental/usd/plugins/mjcf/utils.cc
+60 -119
View File
@@ -15,13 +15,16 @@
#include "mjcf/utils.h"
#include <string>
#include <vector>
#include <pxr/base/tf/token.h>
#include <pxr/usd/sdf/abstractData.h>
#include <pxr/usd/sdf/attributeSpec.h>
#include <pxr/usd/sdf/declareHandles.h>
#include <pxr/usd/sdf/layer.h>
#include <pxr/usd/sdf/listOp.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/usd/sdf/primSpec.h>
#include <pxr/usd/sdf/reference.h>
#include <pxr/usd/sdf/relationshipSpec.h>
#include <pxr/usd/sdf/schema.h>
#include <pxr/usd/sdf/types.h>
#include <pxr/usd/sdf/valueTypeName.h>
@@ -30,184 +33,122 @@
namespace {
template <typename T>
void AppendChild(pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& specPath,
const pxr::TfToken& childKey, const T& child) {
// Get existing children.
std::vector<T> children;
pxr::SdfAbstractDataTypedValue<std::vector<T>> getter(&children);
data->Has(specPath, childKey, &getter);
children.push_back(child);
data->Set(specPath, childKey,
pxr::SdfAbstractDataConstTypedValue<std::vector<T>>(&children));
}
template <typename T>
void AppendListOp(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& spec_path, const pxr::TfToken& field,
void AppendListOp(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec, const pxr::TfToken& field,
const T& item) {
// Get existing list op.
pxr::SdfListOp<T> list_op;
pxr::SdfAbstractDataTypedValue<pxr::SdfListOp<T>> getter(&list_op);
data->Has(spec_path, field, &getter);
auto items = list_op.GetExplicitItems();
pxr::SdfListOp<T> listOp =
prim_spec->GetInfo(field).Get<pxr::SdfListOp<T>>();
auto items = listOp.GetExplicitItems();
items.push_back(item);
list_op.SetExplicitItems(items);
data->Set(spec_path, field,
pxr::SdfAbstractDataConstTypedValue<pxr::SdfListOp<T>>(&list_op));
listOp.SetExplicitItems(items);
prim_spec->SetInfo(field, pxr::VtValue::Take(listOp));
}
template <typename T>
void PrependListOp(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& spec_path, const pxr::TfToken& field,
void PrependListOp(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec, const pxr::TfToken& field,
const T& item) {
// Get existing list op.
pxr::SdfListOp<T> listOp;
pxr::SdfAbstractDataTypedValue<pxr::SdfListOp<T>> getter(&listOp);
data->Has(spec_path, field, &getter);
pxr::SdfListOp<T> listOp =
prim_spec->GetInfo(field).Get<pxr::SdfListOp<T>>();
auto prependedItems = listOp.GetPrependedItems();
prependedItems.insert(prependedItems.begin(), item);
listOp.SetPrependedItems(prependedItems);
data->Set(spec_path, field,
pxr::SdfAbstractDataConstTypedValue<pxr::SdfListOp<T>>(&listOp));
prim_spec->SetInfo(field, pxr::VtValue::Take(listOp));
}
} // namespace
namespace mujoco {
namespace usd {
pxr::SdfPath CreatePrimSpec(pxr::SdfAbstractDataRefPtr& data,
pxr::SdfPrimSpecHandle CreatePrimSpec(pxr::SdfLayerRefPtr layer,
const pxr::SdfPath& parent_path,
const pxr::TfToken& name, const pxr::TfToken& type,
pxr::SdfSpecifier specifier) {
const pxr::SdfPath prim_path = parent_path.AppendChild(name);
data->CreateSpec(prim_path, pxr::SdfSpecTypePrim);
data->Set(prim_path, pxr::SdfFieldKeys->Specifier,
pxr::SdfAbstractDataConstTypedValue<pxr::SdfSpecifier>(&specifier));
pxr::SdfPrimSpecHandle prim_spec = pxr::SdfCreatePrimInLayer(layer, prim_path);
layer->SetField(prim_path, pxr::SdfFieldKeys->Specifier, specifier);
if (!type.IsEmpty()) {
data->Set(prim_path, pxr::SdfFieldKeys->TypeName,
pxr::SdfAbstractDataConstTypedValue<const pxr::TfToken>(&type));
layer->SetField(prim_path, pxr::SdfFieldKeys->TypeName, type);
}
AppendChild(data, parent_path, pxr::SdfChildrenKeys->PrimChildren, name);
return prim_path;
return prim_spec;
}
pxr::SdfPath CreateAttributeSpec(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path,
pxr::SdfAttributeSpecHandle CreateAttributeSpec(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::TfToken& name,
const pxr::SdfValueTypeName& type_name,
pxr::SdfVariability variability) {
const pxr::SdfPath propertyPath = prim_path.AppendProperty(name);
const pxr::SdfPath propertyPath = prim_spec->GetPath().AppendProperty(name);
// Early exit if the attribute spec already exists.
if (data->HasSpec(propertyPath)) {
return propertyPath;
if (layer->HasSpec(propertyPath)) {
return layer->GetAttributeAtPath(propertyPath);
}
data->CreateSpec(propertyPath, pxr::SdfSpecTypeAttribute);
auto spec = pxr::SdfAttributeSpec::New(prim_spec,
name, type_name, variability);
pxr::TfToken typeNameToken = type_name.GetAsToken();
data->Set(propertyPath, pxr::SdfFieldKeys->TypeName,
pxr::SdfAbstractDataConstTypedValue<pxr::TfToken>(&typeNameToken));
if (variability != pxr::SdfVariabilityVarying) {
data->Set(
propertyPath, pxr::SdfFieldKeys->Variability,
pxr::SdfAbstractDataConstTypedValue<pxr::SdfVariability>(&variability));
}
AppendChild(data, prim_path, pxr::SdfChildrenKeys->PropertyChildren, name);
return propertyPath;
return spec;
}
pxr::SdfPath CreateRelationshipSpec(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path,
pxr::SdfRelationshipSpecHandle CreateRelationshipSpec(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::TfToken& relationship_name,
const pxr::SdfPath& relationship_path,
pxr::SdfVariability variability) {
pxr::SdfPath prop_path = prim_path.AppendProperty(relationship_name);
data->CreateSpec(prop_path, pxr::SdfSpecTypeRelationship);
if (variability != pxr::SdfVariabilityVarying) {
data->Set(
prop_path, pxr::SdfFieldKeys->Variability,
pxr::SdfAbstractDataConstTypedValue<pxr::SdfVariability>(&variability));
}
auto spec = pxr::SdfRelationshipSpec::New(prim_spec, relationship_name);
spec->GetTargetPathList().Append(relationship_path);
AppendChild(data, prim_path, pxr::SdfChildrenKeys->PropertyChildren,
relationship_name);
AppendChild(data, prop_path, pxr::SdfChildrenKeys->RelationshipTargetChildren,
relationship_path);
AppendListOp(data, prop_path, pxr::SdfFieldKeys->TargetPaths,
relationship_path);
pxr::SdfPath target_path = prop_path.AppendTarget(relationship_path);
data->CreateSpec(target_path, pxr::SdfSpecTypeRelationshipTarget);
return prop_path;
return spec;
}
pxr::SdfPath CreateClassSpec(pxr::SdfAbstractDataRefPtr& data,
pxr::SdfPrimSpecHandle CreateClassSpec(pxr::SdfLayerRefPtr layer,
const pxr::SdfPath& prim_path,
const pxr::TfToken& class_name) {
pxr::SdfPath class_path = prim_path.AppendChild(class_name);
pxr::SdfSpecifier class_specifier = pxr::SdfSpecifier::SdfSpecifierClass;
data->CreateSpec(class_path, pxr::SdfSpecTypePrim);
data->Set(
class_path, pxr::SdfFieldKeys->Specifier,
pxr::SdfAbstractDataConstTypedValue<pxr::SdfSpecifier>(&class_specifier));
pxr::SdfPrimSpecHandle prim_spec = pxr::SdfCreatePrimInLayer(layer, class_path);
layer->SetField(class_path, pxr::SdfFieldKeys->Specifier, pxr::SdfSpecifier::SdfSpecifierClass);
AppendChild(data, prim_path, pxr::SdfChildrenKeys->PrimChildren, class_name);
return class_path;
return prim_spec;
}
void AddAttributeConnection(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& attribute_path,
const pxr::SdfPath& target_attribute_path) {
AppendChild(data, attribute_path, pxr::SdfChildrenKeys->ConnectionChildren,
target_attribute_path);
AppendListOp(data, attribute_path, pxr::SdfFieldKeys->ConnectionPaths,
target_attribute_path);
data->CreateSpec(attribute_path.AppendTarget(target_attribute_path),
pxr::SdfSpecTypeConnection);
void AddAttributeConnection(pxr::SdfLayerRefPtr layer,
const pxr::SdfAttributeSpecHandle& attribute_spec,
const pxr::SdfAttributeSpecHandle& target_attribute_spec) {
attribute_spec->GetConnectionPathList().GetExplicitItems() = {target_attribute_spec->GetPath()};
}
void AddPrimReference(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path,
void AddPrimReference(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::SdfPath& referenced_prim_path) {
PrependListOp(data, prim_path, pxr::SdfFieldKeys->References,
PrependListOp(layer, prim_spec, pxr::SdfFieldKeys->References,
pxr::SdfReference("", referenced_prim_path));
}
void AddPrimInherit(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path,
void AddPrimInherit(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::SdfPath& class_path) {
PrependListOp(data, prim_path, pxr::SdfFieldKeys->InheritPaths, class_path);
PrependListOp(layer, prim_spec, pxr::SdfFieldKeys->InheritPaths, class_path);
}
void ApplyApiSchema(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path,
void ApplyApiSchema(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::TfToken& schema_name) {
PrependListOp(data, prim_path, pxr::UsdTokens->apiSchemas, schema_name);
PrependListOp(layer, prim_spec, pxr::UsdTokens->apiSchemas, schema_name);
}
void SetPrimKind(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path, pxr::TfToken kind) {
SetPrimMetadata(data, prim_path, pxr::TfToken("kind"), kind);
void SetPrimKind(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec, pxr::TfToken kind) {
prim_spec->SetField(pxr::SdfFieldKeys->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,
void SetPrimPurpose(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec, pxr::TfToken purpose) {
const pxr::SdfAttributeSpecHandle& attr = CreateAttributeSpec(
layer, prim_spec, pxr::UsdGeomTokens->purpose,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
SetAttributeDefault(data, attr_path, purpose);
SetAttributeDefault(layer, attr, purpose);
}
} // namespace usd
src/experimental/usd/plugins/mjcf/utils.h
+38 -64
View File
@@ -15,11 +15,11 @@
#ifndef MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_UTILS_H_
#define MUJOCO_SRC_EXPERIMENTAL_USD_PLUGINS_MJCF_UTILS_H_
#include <type_traits>
#include <pxr/base/tf/token.h>
#include <pxr/imaging/hd/primTypeIndex.h>
#include <pxr/usd/sdf/abstractData.h>
#include <pxr/usd/sdf/attributeSpec.h>
#include <pxr/usd/sdf/declareHandles.h>
#include <pxr/usd/sdf/layer.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/usd/sdf/schema.h>
#include <pxr/usd/sdf/types.h>
@@ -29,123 +29,97 @@ namespace mujoco {
namespace usd {
// Create a prim spec and append it as a child of parent_path.
pxr::SdfPath CreatePrimSpec(
pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& parent_path,
pxr::SdfPrimSpecHandle CreatePrimSpec(
pxr::SdfLayerRefPtr layer, const pxr::SdfPath& parent_path,
const pxr::TfToken& name, const pxr::TfToken& type = pxr::TfToken(),
pxr::SdfSpecifier specifier = pxr::SdfSpecifier::SdfSpecifierDef);
// Create an attribute spec and append it as a child of parent_path.
// By default the attribute will be varying.
pxr::SdfPath CreateAttributeSpec(
pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& prim_path,
pxr::SdfAttributeSpecHandle CreateAttributeSpec(
pxr::SdfLayerRefPtr layer, const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::TfToken& name, const pxr::SdfValueTypeName& type_name,
pxr::SdfVariability variability = pxr::SdfVariabilityVarying);
// Create a relationship spec and append it as a child of prim_path.
pxr::SdfPath CreateRelationshipSpec(
pxr::SdfAbstractDataRefPtr& data, const pxr::SdfPath& prim_path,
pxr::SdfRelationshipSpecHandle CreateRelationshipSpec(
pxr::SdfLayerRefPtr layer, const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::TfToken& relationship_name,
const pxr::SdfPath& relationship_path,
pxr::SdfVariability variability = pxr::SdfVariabilityVarying);
pxr::SdfPath CreateClassSpec(pxr::SdfAbstractDataRefPtr& data,
pxr::SdfPrimSpecHandle CreateClassSpec(pxr::SdfLayerRefPtr layer,
const pxr::SdfPath& prim_path,
const pxr::TfToken& class_name);
void AddAttributeConnection(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& attribute_path,
const pxr::SdfPath& target_attribute_path);
void AddAttributeConnection(pxr::SdfLayerRefPtr layer,
const pxr::SdfAttributeSpecHandle& attribute_spec,
const pxr::SdfAttributeSpecHandle& target_attribute_spec);
void AddPrimReference(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path,
void AddPrimReference(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::SdfPath& referenced_prim_path);
void AddPrimInherit(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path,
void AddPrimInherit(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::SdfPath& class_path);
void ApplyApiSchema(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path,
void ApplyApiSchema(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec,
const pxr::TfToken& schema_name);
void SetPrimKind(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path, pxr::TfToken kind);
void SetPrimKind(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec, pxr::TfToken kind);
void SetPrimPurpose(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path, pxr::TfToken purpose);
void SetPrimPurpose(pxr::SdfLayerRefPtr layer,
const pxr::SdfPrimSpecHandle& prim_spec, 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,
void SetField(pxr::SdfLayerRefPtr layer, const pxr::SdfPath& field_path,
const pxr::TfToken key, T&& value) {
using Deduced = typename std::remove_reference_t<T>;
const auto typed_val = pxr::SdfAbstractDataConstTypedValue<Deduced>(&value);
const pxr::SdfAbstractDataConstValue& untyped_val = typed_val;
data->Set(field_path, key, untyped_val);
layer->SetField(field_path, key, value);
}
// Set the value specified by key on any field at field_path.
template <typename T>
void SetFieldTimeSample(pxr::SdfAbstractDataRefPtr& data,
void SetFieldTimeSample(pxr::SdfLayerRefPtr layer,
const pxr::SdfPath& field_path, double time,
T&& value) {
using Deduced = typename std::remove_reference_t<T>;
const auto typed_val = pxr::SdfAbstractDataConstTypedValue<Deduced>(&value);
const pxr::SdfAbstractDataConstValue& untyped_val = typed_val;
pxr::VtValue vt_value;
untyped_val.GetValue(&vt_value);
// NOTE: SetTimeSample doesn't accept an SdfAbstractDataConstValue yet.
data->SetTimeSample(field_path, time, vt_value);
layer->SetTimeSample(field_path, time, value);
}
// Set the value specified by key on an attribute spec at attribute_path.
template <typename T>
void SetAttribute(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& attribute_path, const pxr::TfToken key,
T&& value) {
SetField(data, attribute_path, key, value);
}
// Set the value specified by key on a prim spec at prim_path.
template <typename T>
void SetPrimMetadata(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& prim_path, const pxr::TfToken key,
T&& value) {
SetAttribute(data, prim_path, key, value);
}
// Set the value specified by key on an attribute spec at attribute_path.
template <typename T>
void SetAttributeMetadata(pxr::SdfAbstractDataRefPtr& data,
void SetAttributeMetadata(pxr::SdfLayerRefPtr layer,
const pxr::SdfPath& attribute_path,
const pxr::TfToken key, T&& value) {
SetAttribute(data, attribute_path, key, value);
SetAttribute(layer, attribute_path, key, value);
}
// Set the default value on an attribute spec at attribute_path.
template <typename T>
void SetAttributeDefault(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& attribute_path,
void SetAttributeDefault(pxr::SdfLayerRefPtr layer,
const pxr::SdfAttributeSpecHandle& attribute,
T&& default_value) {
SetAttribute(data, attribute_path, pxr::SdfFieldKeys->Default, default_value);
attribute->SetField(pxr::SdfFieldKeys->Default, default_value);
}
// Set the default value on an attribute spec at attribute_path.
template <typename T>
void SetAttributeTimeSample(pxr::SdfAbstractDataRefPtr& data,
const pxr::SdfPath& attribute_path,
void SetAttributeTimeSample(pxr::SdfLayerRefPtr layer,
const pxr::SdfAttributeSpecHandle& attr_spec,
double time,
T&& default_value) {
SetFieldTimeSample(data, attribute_path, time, default_value);
layer->SetTimeSample(attr_spec->GetPath(), time, default_value);
}
// Set the value specified by key on the root layer.
template <typename T>
void SetLayerMetadata(pxr::SdfAbstractDataRefPtr& data, const pxr::TfToken& key,
void SetLayerMetadata(pxr::SdfLayerRefPtr layer, const pxr::TfToken& key,
T&& value) {
SetAttribute(data, pxr::SdfPath::AbsoluteRootPath(), key, value);
layer->SetField(pxr::SdfPath::AbsoluteRootPath(), key, value);
}
} // namespace usd