chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
8eca9926dc16200dbcae50587f11f235f7d10d94
Mujoco_WASM/src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
T
Matija Kecman 8eca9926dc Fix various MuJoCo Studio build issues on Windows, MacOS and Linux 
PiperOrigin-RevId: 833313973
Change-Id: I7ec181374e0bb49c32420d025044c31f6fc3c657
2025-11-17 06:35:43 -08:00

2374 lines
101 KiB
C++
Raw Blame History

// Copyright 2025 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "mjcf/mujoco_to_usd.h"
#include <algorithm>
#include <cstddef>
#include <numbers>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include <mujoco/experimental/usd/mjcPhysics/tokens.h>
#include <mujoco/experimental/usd/utils.h>
#include <mujoco/mjspec.h>
#include <mujoco/mujoco.h>
#include "mjcf/utils.h"
#include <pxr/base/arch/attributes.h>
#include <pxr/base/gf/matrix4d.h>
#include <pxr/base/gf/quatf.h>
#include <pxr/base/gf/rotation.h>
#include <pxr/base/gf/vec2f.h>
#include <pxr/base/gf/vec3d.h>
#include <pxr/base/gf/vec3f.h>
#include <pxr/base/tf/callContext.h>
#include <pxr/base/tf/diagnostic.h>
#include <pxr/base/tf/diagnosticHelper.h>
#include <pxr/base/tf/enum.h>
#include <pxr/base/tf/registryManager.h>
#include <pxr/base/tf/staticData.h>
#include <pxr/base/tf/staticTokens.h>
#include <pxr/base/tf/stringUtils.h>
#include <pxr/base/tf/token.h>
#include <pxr/base/vt/array.h>
#include <pxr/base/vt/dictionary.h>
#include <pxr/base/vt/types.h>
#include <pxr/base/vt/value.h>
#include <pxr/usd/kind/registry.h>
#include <pxr/usd/sdf/abstractData.h>
#include <pxr/usd/sdf/assetPath.h>
#include <pxr/usd/sdf/declareHandles.h>
#include <pxr/usd/sdf/path.h>
#include <pxr/usd/sdf/schema.h>
#include <pxr/usd/sdf/types.h>
#include <pxr/usd/sdf/valueTypeName.h>
#include <pxr/usd/usdGeom/metrics.h>
#include <pxr/usd/usdGeom/tokens.h>
#include <pxr/usd/usdLux/tokens.h>
#include <pxr/usd/usdPhysics/fixedJoint.h>
#include <pxr/usd/usdPhysics/joint.h>
#include <pxr/usd/usdPhysics/prismaticJoint.h>
#include <pxr/usd/usdPhysics/revoluteJoint.h>
#include <pxr/usd/usdPhysics/sphericalJoint.h>
#include <pxr/usd/usdPhysics/tokens.h>
#include <pxr/usd/usdShade/tokens.h>
#include <pxr/usdImaging/usdImaging/tokens.h>
namespace {
// The ID of the World in mjModel and mjData.
static constexpr int kWorldIndex = 0;
// Using to satisfy TF_DEFINE_PRIVATE_TOKENS macro below and avoid operating in
// PXR_NS.
using pxr::TfToken;
template <typename T>
using TfStaticData = pxr::TfStaticData<T>;
// clang-format off
TF_DEFINE_PRIVATE_TOKENS(kTokens,
((body, "Body"))
((body_name, "mujoco:body_name"))
((geom, "Geom"))
((light, "Light"))
((meshScope, "MeshSources"))
((materialsScope, "Materials"))
((physicsMaterialsScope, "PhysicsMaterials"))
((previewSurface, "PreviewSurface"))
((keyframesScope, "Keyframes"))
((actuatorsScope, "Actuators"))
((keyframe, "Keyframe"))
((surface, "PreviewSurface"))
((world, "World"))
((xformOpTransform, "xformOp:transform"))
((xformOpScale, "xformOp:scale"))
(st)
((primvarsSt, "primvars:st"))
((outputsSt, "outputs:st"))
((inputsSt, "inputs:st"))
((inputsVarname, "inputs:varname"))
((inputsFile, "inputs:file"))
((inputsWrapS, "inputs:wrapS"))
((inputsWrapT, "inputs:wrapT"))
((inputsDiffuseColor, "inputs:diffuseColor"))
((inputsEmissiveColor, "inputs:emissiveColor"))
((outputsRgb, "outputs:rgb"))
((outputsR, "outputs:r"))
((outputsG, "outputs:g"))
((outputsB, "outputs:b"))
((inputsMetallic, "inputs:metallic"))
((inputsOcclusion, "inputs:occlusion"))
((inputsRoughness, "inputs:roughness"))
(repeat)
((sourceMesh, pxr::UsdGeomTokens->Mesh))
((inputsNormal, "inputs:normal"))
((joint, "Joint"))
);
// Using to satisfy TF_REGISTRY_FUNCTION macro below and avoid operating in PXR_NS.
using pxr::TfEnum;
using pxr::Tf_RegistryStaticInit;
using pxr::Tf_RegistryInit;
using pxr::TfEnum;
template <typename T>
using Arch_PerLibInit = pxr::Arch_PerLibInit<T>;
#if defined(ARCH_OS_DARWIN)
using Arch_ConstructorEntry = pxr::Arch_ConstructorEntry;
#endif
enum ErrorCodes { UnsupportedActuatorTypeError, UnsupportedGeomTypeError, MujocoCompilationError };
TF_REGISTRY_FUNCTION(pxr::TfEnum) {
TF_ADD_ENUM_NAME(UnsupportedGeomTypeError, "UsdGeom type is unsupported.")
TF_ADD_ENUM_NAME(MujocoCompilationError, "Mujoco spec failed to compile.")
}
// Usings to satisfy TF_ERROR macro.
using pxr::TfCallContext;
using pxr::Tf_PostErrorHelper;
// clang-format on
using pxr::MjcPhysicsTokens;
using mujoco::usd::AddAttributeConnection;
using mujoco::usd::AddPrimInherit;
using mujoco::usd::AddPrimReference;
using mujoco::usd::ApplyApiSchema;
using mujoco::usd::CreateAttributeSpec;
using mujoco::usd::CreateClassSpec;
using mujoco::usd::CreatePrimSpec;
using mujoco::usd::CreateRelationshipSpec;
using mujoco::usd::SetAttributeDefault;
using mujoco::usd::SetAttributeTimeSample;
using mujoco::usd::SetLayerMetadata;
using mujoco::usd::SetPrimKind;
using mujoco::usd::SetPrimPurpose;
pxr::GfMatrix4d MujocoPosQuatToTransform(mjtNum *pos, mjtNum *quat) {
pxr::GfQuatd quaternion = pxr::GfQuatd::GetIdentity();
quaternion.SetReal(quat[0]);
quaternion.SetImaginary(quat[1], quat[2], quat[3]);
pxr::GfRotation rotation(quaternion);
pxr::GfVec3d translation(0.0, 0.0, 0.0);
translation.Set(pos[0], pos[1], pos[2]);
pxr::GfMatrix4d transform;
transform.SetTransform(rotation, translation);
return transform;
}
} // namespace
class ModelWriter {
public:
ModelWriter(mjSpec* spec, mjModel* model, pxr::SdfLayerRefPtr layer,
bool skip_elems_from_usd)
: spec_(spec),
model_(model),
layer_(layer),
skip_elems_from_usd_(skip_elems_from_usd),
class_path_("/Bad_Path") {
body_paths_ = std::vector<pxr::SdfPath>(model->nbody);
site_paths_ = std::vector<pxr::SdfPath>(model->nsite);
joint_paths_ = std::vector<pxr::SdfPath>(model->njnt);
}
~ModelWriter() { mj_deleteModel(model_); }
bool ShouldWrite(mjsElement* element) {
// If we've been told to skip over spec elements that originated from USD,
// check the user value and if it has a primpath then it came from decoding
// a USD stage into an MjSpec.
return !skip_elems_from_usd_ || mujoco::usd::GetUsdPrimPathUserValue(element).IsEmpty();
}
void Write() {
// Create top level class holder.
class_path_ = CreateClassSpec(layer_, pxr::SdfPath::AbsoluteRootPath(),
pxr::TfToken("__class__"))->GetPath();
// Create the world body.
body_paths_[kWorldIndex] = WriteWorldBody(kWorldIndex)->GetPath();
layer_->SetDocumentation("Generated by mujoco model writer.");
// Set the world body to be the default prim for referencing/payloads.
layer_->SetDefaultPrim(body_paths_[kWorldIndex].GetNameToken());
// Mujoco is Z up by default.
SetLayerMetadata(layer_, pxr::UsdGeomTokens->upAxis, pxr::UsdGeomTokens->z);
// Mujoco is authored in meters by default.
SetLayerMetadata(layer_, pxr::UsdGeomTokens->metersPerUnit,
pxr::UsdGeomLinearUnits::meters);
WritePhysicsScene();
// Author mesh scope + mesh prims to be referenced.
WriteMeshes();
WriteMaterials();
WriteBodies();
WriteActuators();
WriteKeyframes();
}
private:
mjSpec *spec_;
mjModel *model_;
pxr::SdfLayerRefPtr layer_;
bool skip_elems_from_usd_ = false;
// Path to top level class spec that all classes should be children of.
pxr::SdfPath class_path_;
// Mapping from Mujoco body id to SdfPath.
std::vector<pxr::SdfPath> body_paths_;
// Mapping from Mujoco site id to SdfPath.
std::vector<pxr::SdfPath> site_paths_;
// Mapping from Mujoco joint id to SdfPath.
std::vector<pxr::SdfPath> joint_paths_;
// Mapping from mesh names to Mesh prim path.
std::unordered_map<std::string, pxr::SdfPath> mesh_paths_;
// Set of body ids that have had the articulation root API applied.
std::unordered_set<int> articulation_roots_;
// Given a name index and a parent prim path this returns a
// token such that appending it to the parent prim path does not
// identify an existing prim spec.
//
// This is necessary since mujoco does not require names for elements
// so we must differentiate between elements of the same type.
//
// For example:
// <mujoco model="model with two planes">
// <worldbody>
// <geom type="plane"/>
// <geom type="plane"/>
// </worldbody>
// </mujoco>
//
// We expect that the occurrence of this happens little enough that linear
// searching is plenty efficient.
pxr::TfToken GetAvailablePrimName(const std::string base_name,
const pxr::TfToken fallback_name,
const pxr::SdfPath &parent_path) {
const auto valid_base_name = pxr::TfMakeValidIdentifier(
base_name.empty() ? fallback_name : base_name);
std::string name = valid_base_name;
pxr::SdfPath test_path = parent_path.AppendChild(pxr::TfToken(name));
int count = 1;
while (layer_->HasSpec(test_path) &&
layer_->GetSpecType(test_path) == pxr::SdfSpecType::SdfSpecTypePrim) {
name = pxr::TfStringPrintf("%s_%d", valid_base_name.c_str(), count++);
test_path = parent_path.AppendChild(pxr::TfToken(name));
}
return pxr::TfToken(name);
}
// This function, conversely to GetAvailablePrimName will not handle
// collisions. This is useful when looking up a prim path that might exist or
// a path that you know must be unique.
pxr::TfToken GetValidPrimName(const std::string name) {
return pxr::TfToken(pxr::TfMakeValidIdentifier(name));
}
struct BodyPathComponents {
pxr::SdfPath parent_path;
pxr::TfToken body_name;
};
void WriteScaleXformOp(const pxr::SdfPrimSpecHandle prim_spec,
const pxr::GfVec3f &scale) {
auto scale_attr_spec =
CreateAttributeSpec(layer_, prim_spec, kTokens->xformOpScale,
pxr::SdfValueTypeNames->Float3);
SetAttributeDefault(layer_, scale_attr_spec, scale);
}
void WriteTransformXformOp(const pxr::SdfPrimSpecHandle prim_spec,
const pxr::GfMatrix4d &transform) {
auto transform_op_spec =
CreateAttributeSpec(layer_, prim_spec, kTokens->xformOpTransform,
pxr::SdfValueTypeNames->Matrix4d);
SetAttributeDefault(layer_, transform_op_spec, transform);
}
void WriteXformOpOrder(const pxr::SdfPrimSpecHandle prim_spec,
const pxr::VtArray<pxr::TfToken> &order) {
auto xform_op_order_spec =
CreateAttributeSpec(layer_, prim_spec, pxr::UsdGeomTokens->xformOpOrder,
pxr::SdfValueTypeNames->TokenArray);
SetAttributeDefault(layer_, xform_op_order_spec, order);
}
template <typename T>
void WriteUniformAttribute(const pxr::SdfPrimSpecHandle prim_spec,
const pxr::SdfValueTypeName &value_type_name,
const pxr::TfToken &token, const T &value) {
auto attr_spec = CreateAttributeSpec(
layer_, prim_spec, token, value_type_name, pxr::SdfVariabilityUniform);
SetAttributeDefault(layer_, attr_spec, value);
}
template <typename T>
void WriteColorAndOpacityAttributes(const pxr::SdfPrimSpecHandle prim_spec,
const T &element) {
// If rgba is not the default (0.5, 0.5, 0.5, 1), then set the
// displayColor attribute.
// No effort is made to properly handle the interaction between rgba
// and the material if both are specified.
if (element->rgba[0] != 0.5f || element->rgba[1] != 0.5f ||
element->rgba[2] != 0.5f || element->rgba[3] != 1.0f) {
// Set the displayColor attribute.
auto display_color_attr = CreateAttributeSpec(
layer_, prim_spec, pxr::UsdGeomTokens->primvarsDisplayColor,
pxr::SdfValueTypeNames->Color3fArray);
SetAttributeDefault(
layer_, display_color_attr,
pxr::VtArray<pxr::GfVec3f>{
{element->rgba[0], element->rgba[1], element->rgba[2]}});
// Set the displayOpacity attribute, only if the opacity is not 1.
if (element->rgba[3] != 1.0f) {
auto display_opacity_attr = CreateAttributeSpec(
layer_, prim_spec, pxr::UsdGeomTokens->primvarsDisplayOpacity,
pxr::SdfValueTypeNames->FloatArray);
SetAttributeDefault(layer_, display_opacity_attr,
pxr::VtArray<float>{element->rgba[3]});
}
}
}
void PrependToXformOpOrder(const pxr::SdfPrimSpecHandle prim_spec,
const pxr::VtArray<pxr::TfToken> &order) {
auto xform_op_order_path =
prim_spec->GetPath().AppendProperty(pxr::UsdGeomTokens->xformOpOrder);
if (!prim_spec->HasField(pxr::UsdGeomTokens->xformOpOrder)) {
WriteXformOpOrder(prim_spec, order);
return;
}
auto existing_order =
prim_spec->GetField(pxr::UsdGeomTokens->xformOpOrder).Get<pxr::VtArray<pxr::TfToken>>();
pxr::VtArray<pxr::TfToken> new_order(order.size() + existing_order.size());
std::copy(order.begin(), order.end(), new_order.begin());
std::copy(existing_order.begin(), existing_order.end(),
new_order.begin() + order.size());
SetAttributeDefault(layer_, layer_->GetAttributeAtPath(xform_op_order_path), new_order);
}
void WriteMesh(const mjsMesh *mesh, const pxr::SdfPath &parent_path) {
auto name = GetAvailablePrimName(*mjs_getName(mesh->element),
pxr::UsdGeomTokens->Mesh, parent_path);
auto subcomponent_spec =
CreatePrimSpec(layer_, parent_path, name, pxr::UsdGeomTokens->Xform);
auto mesh_spec =
CreatePrimSpec(layer_, subcomponent_spec->GetPath(), kTokens->sourceMesh,
pxr::UsdGeomTokens->Mesh);
mesh_paths_[*mjs_getName(mesh->element)] = subcomponent_spec->GetPath();
ApplyApiSchema(layer_, mesh_spec, MjcPhysicsTokens->MjcMeshCollisionAPI);
pxr::TfToken inertia = MjcPhysicsTokens->legacy;
if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_EXACT) {
inertia = MjcPhysicsTokens->exact;
} else if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_CONVEX) {
inertia = MjcPhysicsTokens->convex;
} else if (mesh->inertia == mjtMeshInertia::mjMESH_INERTIA_SHELL) {
inertia = MjcPhysicsTokens->shell;
}
WriteUniformAttribute(mesh_spec, pxr::SdfValueTypeNames->Token,
MjcPhysicsTokens->mjcInertia, inertia);
WriteUniformAttribute(mesh_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcMaxhullvert, mesh->maxhullvert);
// NOTE: The geometry data taken from the spec is the post-compilation
// data after it has been mjCMesh::Compile'd. So don't be surprised if
// things like user defined vertices have moved due to re-centering to
// CoM and other modifications (see mjCMesh::Process for other xforms).
int mesh_id = mjs_getId(mesh->element);
int vert_start_offset = model_->mesh_vertadr[mesh_id] * 3;
int nvert = model_->mesh_vertnum[mesh_id];
pxr::VtArray<pxr::GfVec3f> points;
points.reserve(nvert);
for (int i = vert_start_offset; i < vert_start_offset + nvert * 3; i += 3) {
points.emplace_back(&model_->mesh_vert[i]);
}
auto points_attr =
CreateAttributeSpec(layer_, mesh_spec, pxr::UsdGeomTokens->points,
pxr::SdfValueTypeNames->Vector3fArray);
SetAttributeDefault(layer_, points_attr, points);
// NOTE: nface is never 0.
int nface = model_->mesh_facenum[mesh_id];
pxr::VtArray<int> faces;
faces.reserve(nface * 3);
int face_start_offset = model_->mesh_faceadr[mesh_id] * 3;
for (int i = face_start_offset; i < face_start_offset + nface * 3; i += 3) {
faces.push_back(model_->mesh_face[i]);
faces.push_back(model_->mesh_face[i + 1]);
faces.push_back(model_->mesh_face[i + 2]);
}
auto face_vertex_idx_attr= CreateAttributeSpec(
layer_, mesh_spec, pxr::UsdGeomTokens->faceVertexIndices,
pxr::SdfValueTypeNames->IntArray);
SetAttributeDefault(layer_, face_vertex_idx_attr, faces);
pxr::VtArray<int> vertex_counts;
for (int i = 0; i < nface; ++i) {
// Mujoco is always triangles.
vertex_counts.push_back(3);
}
auto face_vertex_counts_attr= CreateAttributeSpec(
layer_, mesh_spec, pxr::UsdGeomTokens->faceVertexCounts,
pxr::SdfValueTypeNames->IntArray);
SetAttributeDefault(layer_, face_vertex_counts_attr, vertex_counts);
if (model_->mesh_normalnum[mesh_id]) {
// We have to convert from Mujoco's indexed normals to USD's faceVarying
// normals.
pxr::VtArray<pxr::GfVec3f> normals;
normals.reserve(nface * 3);
int normal_start_adr = model_->mesh_normaladr[mesh_id];
int face_start_offset = model_->mesh_faceadr[mesh_id] * 3;
for (int i = face_start_offset; i < face_start_offset + nface * 3; ++i) {
int normal_adr = normal_start_adr + model_->mesh_facenormal[i];
normals.emplace_back(&model_->mesh_normal[normal_adr * 3]);
}
auto normals_attr=
CreateAttributeSpec(layer_, mesh_spec, pxr::UsdGeomTokens->normals,
pxr::SdfValueTypeNames->Vector3fArray);
SetAttributeDefault(layer_, normals_attr, normals);
normals_attr->SetField(pxr::UsdGeomTokens->interpolation,
pxr::UsdGeomTokens->faceVarying);
}
if (model_->mesh_texcoordnum[mesh_id]) {
// We have to convert from Mujoco's indexed texcoords to USD's faceVarying
// texcoords.
pxr::VtArray<pxr::GfVec2f> texcoords;
texcoords.reserve(nface * 3);
int texcoord_start_adr = model_->mesh_texcoordadr[mesh_id];
int face_start_offset = model_->mesh_faceadr[mesh_id] * 3;
for (int i = face_start_offset; i < face_start_offset + nface * 3; ++i) {
int texcoord_adr = texcoord_start_adr + model_->mesh_facetexcoord[i];
// Invert the V coordinate, Mujoco assumes OpenGL 0,0 is top left.
// But USD UVs use image bottom left 0,0 convention.
pxr::GfVec2f uv(&model_->mesh_texcoord[texcoord_adr * 2]);
uv[1] = 1.0f - uv[1];
texcoords.push_back(uv);
}
auto texcoords_attr=
CreateAttributeSpec(layer_, mesh_spec, kTokens->primvarsSt,
pxr::SdfValueTypeNames->TexCoord2fArray);
SetAttributeDefault(layer_, texcoords_attr, texcoords);
texcoords_attr->SetField(pxr::UsdGeomTokens->interpolation,
pxr::UsdGeomTokens->faceVarying);
}
// Default subdivision scheme is catmull clark so explicitly set it
// to none here.
auto subdivision_scheme= CreateAttributeSpec(
layer_, mesh_spec, pxr::UsdGeomTokens->subdivisionScheme,
pxr::SdfValueTypeNames->Token);
SetAttributeDefault(layer_, subdivision_scheme,
pxr::UsdGeomTokens->none);
}
void WritePhysicsScene() {
auto physics_scene_spec = CreatePrimSpec(
layer_, body_paths_[kWorldIndex], pxr::UsdPhysicsTokens->PhysicsScene,
pxr::UsdPhysicsTokens->PhysicsScene);
ApplyApiSchema(layer_, physics_scene_spec, MjcPhysicsTokens->MjcSceneAPI);
const std::vector<std::pair<pxr::TfToken, double>>
option_double_attributes = {
{MjcPhysicsTokens->mjcOptionTimestep, spec_->option.timestep},
{MjcPhysicsTokens->mjcOptionTolerance, spec_->option.tolerance},
{MjcPhysicsTokens->mjcOptionLs_tolerance,
spec_->option.ls_tolerance},
{MjcPhysicsTokens->mjcOptionNoslip_tolerance,
spec_->option.noslip_tolerance},
{MjcPhysicsTokens->mjcOptionCcd_tolerance,
spec_->option.ccd_tolerance},
{MjcPhysicsTokens->mjcOptionImpratio, spec_->option.impratio},
{MjcPhysicsTokens->mjcOptionDensity, spec_->option.density},
{MjcPhysicsTokens->mjcOptionViscosity, spec_->option.viscosity},
{MjcPhysicsTokens->mjcOptionO_margin, spec_->option.o_margin},
};
for (const auto &[token, value] : option_double_attributes) {
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Double,
token, value);
}
const std::vector<std::pair<pxr::TfToken, int>> option_int_attributes = {
{MjcPhysicsTokens->mjcOptionIterations, spec_->option.iterations},
{MjcPhysicsTokens->mjcOptionLs_iterations, spec_->option.ls_iterations},
{MjcPhysicsTokens->mjcOptionNoslip_iterations,
spec_->option.noslip_iterations},
{MjcPhysicsTokens->mjcOptionCcd_iterations,
spec_->option.ccd_iterations},
{MjcPhysicsTokens->mjcOptionSdf_iterations,
spec_->option.sdf_iterations},
{MjcPhysicsTokens->mjcOptionSdf_initpoints,
spec_->option.sdf_initpoints},
};
for (const auto &[token, value] : option_int_attributes) {
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Int,
token, value);
}
auto cone_attr = CreateAttributeSpec(
layer_, physics_scene_spec, MjcPhysicsTokens->mjcOptionCone,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
switch (spec_->option.cone) {
case mjCONE_PYRAMIDAL:
SetAttributeDefault(layer_, cone_attr, MjcPhysicsTokens->pyramidal);
break;
case mjCONE_ELLIPTIC:
SetAttributeDefault(layer_, cone_attr, MjcPhysicsTokens->elliptic);
break;
default:
break;
}
auto jacobian_attr = CreateAttributeSpec(
layer_, physics_scene_spec, MjcPhysicsTokens->mjcOptionJacobian,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
switch (spec_->option.jacobian) {
case mjJAC_AUTO:
SetAttributeDefault(layer_, jacobian_attr, MjcPhysicsTokens->auto_);
break;
case mjJAC_DENSE:
SetAttributeDefault(layer_, jacobian_attr, MjcPhysicsTokens->dense);
break;
case mjJAC_SPARSE:
SetAttributeDefault(layer_, jacobian_attr, MjcPhysicsTokens->sparse);
break;
default:
break;
}
auto solver_attr = CreateAttributeSpec(
layer_, physics_scene_spec, MjcPhysicsTokens->mjcOptionSolver,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
switch (spec_->option.solver) {
case mjSOL_NEWTON:
SetAttributeDefault(layer_, solver_attr, MjcPhysicsTokens->newton);
break;
case mjSOL_PGS:
SetAttributeDefault(layer_, solver_attr, MjcPhysicsTokens->pgs);
break;
case mjSOL_CG:
SetAttributeDefault(layer_, solver_attr, MjcPhysicsTokens->cg);
break;
default:
break;
}
pxr::GfVec3f gravity(spec_->option.gravity[0], spec_->option.gravity[1],
spec_->option.gravity[2]);
// Normalize will normalize gravity in place and return the magnitude before
// normalization.
float gravity_magnitude = gravity.Normalize();
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Float,
pxr::UsdPhysicsTokens->physicsGravityMagnitude,
gravity_magnitude);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Vector3f,
pxr::UsdPhysicsTokens->physicsGravityDirection,
gravity);
pxr::GfVec3d wind(spec_->option.wind[0], spec_->option.wind[1],
spec_->option.wind[2]);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Double3,
MjcPhysicsTokens->mjcOptionWind, wind);
pxr::GfVec3d magnetic(spec_->option.magnetic[0], spec_->option.magnetic[1],
spec_->option.magnetic[2]);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Double3,
MjcPhysicsTokens->mjcOptionMagnetic, magnetic);
pxr::VtArray<double> o_solref(spec_->option.o_solref,
spec_->option.o_solref + mjNREF);
WriteUniformAttribute(physics_scene_spec,
pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcOptionO_solref, o_solref);
pxr::VtArray<double> o_solimp(spec_->option.o_solimp,
spec_->option.o_solimp + mjNIMP);
WriteUniformAttribute(physics_scene_spec,
pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcOptionO_solimp, o_solimp);
pxr::VtArray<double> o_friction(spec_->option.o_friction,
spec_->option.o_friction + 5);
WriteUniformAttribute(physics_scene_spec,
pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcOptionO_friction, o_friction);
auto integrator_attr = CreateAttributeSpec(
layer_, physics_scene_spec, MjcPhysicsTokens->mjcOptionIntegrator,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
switch (spec_->option.integrator) {
case mjINT_EULER:
SetAttributeDefault(layer_, integrator_attr, MjcPhysicsTokens->euler);
break;
case mjINT_RK4:
SetAttributeDefault(layer_, integrator_attr, MjcPhysicsTokens->rk4);
break;
default:
break;
}
auto create_flag_attr = [&](pxr::TfToken token, int flag, bool enable) {
int flags =
enable ? spec_->option.enableflags : spec_->option.disableflags;
bool value = enable ? (flags & flag) : !(flags & flag);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
token, value);
};
const std::vector<std::pair<pxr::TfToken, int>> enable_flags = {
{MjcPhysicsTokens->mjcFlagMulticcd, mjENBL_MULTICCD},
{MjcPhysicsTokens->mjcFlagFwdinv, mjENBL_FWDINV},
{MjcPhysicsTokens->mjcFlagEnergy, mjENBL_ENERGY},
{MjcPhysicsTokens->mjcFlagOverride, mjENBL_OVERRIDE},
{MjcPhysicsTokens->mjcFlagInvdiscrete, mjENBL_INVDISCRETE}};
for (const auto &[token, flag] : enable_flags) {
create_flag_attr(token, flag, true);
}
const std::vector<std::pair<pxr::TfToken, int>> disable_flags = {
{MjcPhysicsTokens->mjcFlagConstraint, mjDSBL_CONSTRAINT},
{MjcPhysicsTokens->mjcFlagEquality, mjDSBL_EQUALITY},
{MjcPhysicsTokens->mjcFlagFrictionloss, mjDSBL_FRICTIONLOSS},
{MjcPhysicsTokens->mjcFlagLimit, mjDSBL_LIMIT},
{MjcPhysicsTokens->mjcFlagContact, mjDSBL_CONTACT},
{MjcPhysicsTokens->mjcFlagSpring, mjDSBL_SPRING},
{MjcPhysicsTokens->mjcFlagDamper, mjDSBL_DAMPER},
{MjcPhysicsTokens->mjcFlagGravity, mjDSBL_GRAVITY},
{MjcPhysicsTokens->mjcFlagClampctrl, mjDSBL_CLAMPCTRL},
{MjcPhysicsTokens->mjcFlagWarmstart, mjDSBL_WARMSTART},
{MjcPhysicsTokens->mjcFlagFilterparent, mjDSBL_FILTERPARENT},
{MjcPhysicsTokens->mjcFlagActuation, mjDSBL_ACTUATION},
{MjcPhysicsTokens->mjcFlagRefsafe, mjDSBL_REFSAFE},
{MjcPhysicsTokens->mjcFlagSensor, mjDSBL_SENSOR},
{MjcPhysicsTokens->mjcFlagMidphase, mjDSBL_MIDPHASE},
{MjcPhysicsTokens->mjcFlagEulerdamp, mjDSBL_EULERDAMP},
{MjcPhysicsTokens->mjcFlagAutoreset, mjDSBL_AUTORESET},
{MjcPhysicsTokens->mjcFlagNativeccd, mjDSBL_NATIVECCD},
{MjcPhysicsTokens->mjcFlagIsland, mjDSBL_ISLAND}};
for (const auto &[token, flag] : disable_flags) {
create_flag_attr(token, flag, false);
}
// Compiler attributes
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcCompilerAutoLimits,
(bool)spec_->compiler.autolimits);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcCompilerBoundMass,
spec_->compiler.boundmass);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcCompilerBoundInertia,
spec_->compiler.boundinertia);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcCompilerSetTotalMass,
spec_->compiler.settotalmass);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcCompilerUseThread,
(bool)spec_->compiler.usethread);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcCompilerBalanceInertia,
(bool)spec_->compiler.balanceinertia);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Token,
MjcPhysicsTokens->mjcCompilerAngle,
spec_->compiler.degree ? MjcPhysicsTokens->degree
: MjcPhysicsTokens->radian);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcCompilerFitAABB,
(bool)spec_->compiler.fitaabb);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcCompilerFuseStatic,
(bool)spec_->compiler.fusestatic);
pxr::TfToken inertiafromgeom_token = MjcPhysicsTokens->auto_;
if (spec_->compiler.inertiafromgeom ==
mjINERTIAFROMGEOM_TRUE) { // mjINERTIA_TRUE
inertiafromgeom_token = MjcPhysicsTokens->true_;
} else if (spec_->compiler.inertiafromgeom ==
mjINERTIAFROMGEOM_FALSE) { // mjINERTIA_FALSE
inertiafromgeom_token = MjcPhysicsTokens->false_;
}
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Token,
MjcPhysicsTokens->mjcCompilerInertiaFromGeom,
inertiafromgeom_token);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcCompilerAlignFree,
(bool)spec_->compiler.alignfree);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcCompilerInertiaGroupRangeMin,
spec_->compiler.inertiagrouprange[0]);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcCompilerInertiaGroupRangeMax,
spec_->compiler.inertiagrouprange[1]);
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcCompilerSaveInertial,
(bool)spec_->compiler.saveinertial);
}
void WriteMeshes() {
// Create a scope for the meshes to keep things organized
auto scope =
CreatePrimSpec(layer_, body_paths_[kWorldIndex], kTokens->meshScope,
pxr::UsdGeomTokens->Scope);
// Make the mesh scope invisible since they will be referenced by the bits
// that should be visible.
scope->SetField(pxr::SdfFieldKeys->Active, false);
mjsMesh *mesh = mjs_asMesh(mjs_firstElement(spec_, mjOBJ_MESH));
while (mesh) {
if (ShouldWrite(mesh->element)) {
WriteMesh(mesh, scope->GetPath());
}
mesh = mjs_asMesh(mjs_nextElement(spec_, mesh->element));
}
}
pxr::SdfAttributeSpecHandle AddUVTextureShader(
const pxr::SdfPrimSpecHandle material_spec, const pxr::TfToken &name) {
auto uvmap_shader_spec =
CreatePrimSpec(layer_, material_spec->GetPath(), name, pxr::UsdShadeTokens->Shader);
auto uvmap_info_id_attr = CreateAttributeSpec(
layer_, uvmap_shader_spec, pxr::UsdShadeTokens->infoId,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
SetAttributeDefault(layer_, uvmap_info_id_attr,
pxr::UsdImagingTokens->UsdPrimvarReader_float2);
auto uvmap_varname_attr =
CreateAttributeSpec(layer_, uvmap_shader_spec, kTokens->inputsVarname,
pxr::SdfValueTypeNames->Token);
SetAttributeDefault(layer_, uvmap_varname_attr, kTokens->st);
auto uvmap_st_output_attr =
CreateAttributeSpec(layer_, uvmap_shader_spec, kTokens->outputsSt,
pxr::SdfValueTypeNames->Float2);
return uvmap_st_output_attr;
}
std::vector<pxr::SdfAttributeSpecHandle> AddTextureShader(
const pxr::SdfPrimSpecHandle material_spec, const char *texture_file,
const pxr::TfToken &name,
const pxr::SdfAttributeSpecHandle uvmap_st_output_attr,
const std::vector<pxr::TfToken> &output_channels) {
auto texture_shader_spec =
CreatePrimSpec(layer_, material_spec->GetPath(), name, pxr::UsdShadeTokens->Shader);
auto texture_info_id_attr = CreateAttributeSpec(
layer_, texture_shader_spec, pxr::UsdShadeTokens->infoId,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
SetAttributeDefault(layer_, texture_info_id_attr,
pxr::UsdImagingTokens->UsdUVTexture);
auto texture_file_attr =
CreateAttributeSpec(layer_, texture_shader_spec, kTokens->inputsFile,
pxr::SdfValueTypeNames->Asset);
SetAttributeDefault(layer_, texture_file_attr,
pxr::SdfAssetPath(texture_file));
auto texture_st_input_attr =
CreateAttributeSpec(layer_, texture_shader_spec, kTokens->inputsSt,
pxr::SdfValueTypeNames->Float2);
AddAttributeConnection(layer_, texture_st_input_attr, uvmap_st_output_attr);
auto texture_wrap_s_attr =
CreateAttributeSpec(layer_, texture_shader_spec, kTokens->inputsWrapS,
pxr::SdfValueTypeNames->Token);
SetAttributeDefault(layer_, texture_wrap_s_attr, kTokens->repeat);
auto texture_wrap_t_attr =
CreateAttributeSpec(layer_, texture_shader_spec, kTokens->inputsWrapT,
pxr::SdfValueTypeNames->Token);
SetAttributeDefault(layer_, texture_wrap_t_attr, kTokens->repeat);
std::vector<pxr::SdfAttributeSpecHandle> texture_output_attrs;
for (const auto &output_channel : output_channels) {
pxr::SdfValueTypeName value_type;
if (output_channel == kTokens->outputsRgb) {
value_type = pxr::SdfValueTypeNames->Float3;
} else {
// Assume the other specified channels are outputR, outputG, outputB.
value_type = pxr::SdfValueTypeNames->Float;
}
texture_output_attrs.push_back(CreateAttributeSpec(
layer_, texture_shader_spec, output_channel, value_type));
}
return texture_output_attrs;
}
pxr::SdfPrimSpecHandle WritePhysicsMaterial(mjsGeom *geom) {
pxr::SdfPath scope_path =
body_paths_[kWorldIndex].AppendChild(kTokens->physicsMaterialsScope);
pxr::SdfPrimSpecHandle scope_spec = layer_->GetPrimAtPath(scope_path);
if (!scope_spec) {
scope_spec = CreatePrimSpec(layer_, body_paths_[kWorldIndex],
kTokens->physicsMaterialsScope,
pxr::UsdGeomTokens->Scope);
}
auto name = GetAvailablePrimName(*mjs_getName(geom->element),
pxr::UsdShadeTokens->Material, scope_path);
auto material_spec =
CreatePrimSpec(layer_, scope_spec->GetPath(), name, pxr::UsdShadeTokens->Material);
ApplyApiSchema(layer_, material_spec,
pxr::UsdPhysicsTokens->PhysicsMaterialAPI);
ApplyApiSchema(layer_, material_spec, MjcPhysicsTokens->MjcMaterialAPI);
mjsGeom *geom_default = mjs_getDefault(geom->element)->geom;
if (geom->friction[0] != geom_default->friction[0]) {
// Since MuJoCo has no concept of static friction, only write dynamic
// friction to remain truthful to how MuJoCo perceives the data.
WriteUniformAttribute(material_spec, pxr::SdfValueTypeNames->Float,
pxr::UsdPhysicsTokens->physicsDynamicFriction,
(float)geom->friction[0]);
}
if (geom->friction[1] != geom_default->friction[1]) {
WriteUniformAttribute(material_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcTorsionalfriction,
geom->friction[1]);
}
if (geom->friction[2] != geom_default->friction[2]) {
WriteUniformAttribute(material_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcRollingfriction,
geom->friction[2]);
}
return material_spec;
}
void WriteMaterial(mjsMaterial *material, const pxr::SdfPath &parent_path) {
// Create a Material prim.
auto name =
GetAvailablePrimName(*mjs_getName(material->element),
pxr::UsdShadeTokens->Material, parent_path);
auto material_spec =
CreatePrimSpec(layer_, parent_path, name, pxr::UsdShadeTokens->Material);
// Create a Shader prim "PreviewSurface" under the Material prim.
auto preview_surface_shader_spec =
CreatePrimSpec(layer_, material_spec->GetPath(), kTokens->previewSurface,
pxr::UsdShadeTokens->Shader);
// Set the Shader'sinfoId attribute to UsdPreviewSurface, a standard surface
// shader.
auto info_id_attr = CreateAttributeSpec(
layer_, preview_surface_shader_spec, pxr::UsdShadeTokens->infoId,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
SetAttributeDefault(layer_, info_id_attr,
pxr::UsdImagingTokens->UsdPreviewSurface);
// Connect material's surface output to the preview surface's surface
// output.
auto surface_output_attr = CreateAttributeSpec(
layer_, preview_surface_shader_spec, pxr::UsdShadeTokens->outputsSurface,
pxr::SdfValueTypeNames->Token);
auto material_surface_output_attr = CreateAttributeSpec(
layer_, material_spec, pxr::UsdShadeTokens->outputsSurface,
pxr::SdfValueTypeNames->Token);
AddAttributeConnection(layer_, material_surface_output_attr,
surface_output_attr);
// Connect material's displacement output to the preview surface's
// displacement output.
auto displacement_output_attr =
CreateAttributeSpec(layer_, preview_surface_shader_spec,
pxr::UsdShadeTokens->outputsDisplacement,
pxr::SdfValueTypeNames->Token);
auto material_displacement_output_attr = CreateAttributeSpec(
layer_, material_spec, pxr::UsdShadeTokens->outputsDisplacement,
pxr::SdfValueTypeNames->Token);
AddAttributeConnection(layer_, material_displacement_output_attr,
displacement_output_attr);
// Add an st (uv) Shader, a prim var reader for the UV coordinates.
const pxr::SdfAttributeSpecHandle &uvmap_st_output_attr =
AddUVTextureShader(material_spec, pxr::TfToken("uvmap"));
const mjStringVec &textures = *(material->textures);
// Set the values of metallic, roughness and occlusion. These can come from
// an ORM packed texture, as individual textures, or as a values defined in
// mjsMaterial_ (with the exception of occlusion).
auto metallic_attr = CreateAttributeSpec(
layer_, preview_surface_shader_spec, kTokens->inputsMetallic,
pxr::SdfValueTypeNames->Float);
auto roughness_attr = CreateAttributeSpec(
layer_, preview_surface_shader_spec, kTokens->inputsRoughness,
pxr::SdfValueTypeNames->Float);
// Find the occlusion, roughness, and metallic textures.
if (mjTEXROLE_ORM < textures.size()) {
std::string orm_texture_name = textures[mjTEXROLE_ORM];
mjsTexture *orm_texture = mjs_asTexture(
mjs_findElement(spec_, mjOBJ_TEXTURE, orm_texture_name.c_str()));
std::string occlusion_texture_name = textures[mjTEXROLE_OCCLUSION];
mjsTexture *occlusion_texture = mjs_asTexture(mjs_findElement(
spec_, mjOBJ_TEXTURE, occlusion_texture_name.c_str()));
std::string roughness_texture_name = textures[mjTEXROLE_ROUGHNESS];
mjsTexture *roughness_texture = mjs_asTexture(mjs_findElement(
spec_, mjOBJ_TEXTURE, roughness_texture_name.c_str()));
std::string metallic_texture_name = textures[mjTEXROLE_METALLIC];
mjsTexture *metallic_texture = mjs_asTexture(
mjs_findElement(spec_, mjOBJ_TEXTURE, metallic_texture_name.c_str()));
if (orm_texture) {
// Create the ORM shader and connect its output to the preview
// surface ORM attrs.
const std::vector<pxr::SdfAttributeSpecHandle> orm_output_attrs =
AddTextureShader(
material_spec, orm_texture->file->c_str(),
pxr::TfToken("orm_packed"), uvmap_st_output_attr,
{kTokens->outputsR, kTokens->outputsG, kTokens->outputsB});
if (orm_output_attrs.size() == 3) {
auto occlusion_attr = CreateAttributeSpec(
layer_, preview_surface_shader_spec, kTokens->inputsOcclusion,
pxr::SdfValueTypeNames->Float);
AddAttributeConnection(layer_, occlusion_attr, orm_output_attrs[0]);
AddAttributeConnection(layer_, roughness_attr, orm_output_attrs[1]);
AddAttributeConnection(layer_, metallic_attr, orm_output_attrs[2]);
}
} else {
if (metallic_texture) {
const std::vector<pxr::SdfAttributeSpecHandle> metallic_output_attrs =
AddTextureShader(material_spec, metallic_texture->file->c_str(),
pxr::TfToken("metallic"), uvmap_st_output_attr,
{kTokens->outputsRgb});
if (metallic_output_attrs.size() == 1) {
AddAttributeConnection(layer_, metallic_attr,
metallic_output_attrs[0]);
}
} else {
SetAttributeDefault(layer_, metallic_attr, material->metallic);
}
if (roughness_texture) {
const std::vector<pxr::SdfAttributeSpecHandle> roughness_output_attrs =
AddTextureShader(material_spec, roughness_texture->file->c_str(),
pxr::TfToken("roughness"), uvmap_st_output_attr,
{kTokens->outputsRgb});
if (roughness_output_attrs.size() == 1) {
AddAttributeConnection(layer_, roughness_attr,
roughness_output_attrs[0]);
}
} else {
SetAttributeDefault(layer_, roughness_attr, material->roughness);
}
if (occlusion_texture) {
auto occlusion_attr = CreateAttributeSpec(
layer_, preview_surface_shader_spec, kTokens->inputsOcclusion,
pxr::SdfValueTypeNames->Float);
const std::vector<pxr::SdfAttributeSpecHandle> occlusion_output_attrs =
AddTextureShader(material_spec, occlusion_texture->file->c_str(),
pxr::TfToken("occlusion"), uvmap_st_output_attr,
{kTokens->outputsRgb});
if (occlusion_output_attrs.size() == 1) {
AddAttributeConnection(layer_, occlusion_attr,
occlusion_output_attrs[0]);
}
}
}
}
// Find the normal texture if specified.
if (mjTEXROLE_NORMAL < textures.size()) {
std::string normal_texture_name = textures[mjTEXROLE_NORMAL];
mjsTexture *normal_texture = mjs_asTexture(
mjs_findElement(spec_, mjOBJ_TEXTURE, normal_texture_name.c_str()));
if (normal_texture) {
auto normal_attr = CreateAttributeSpec(
layer_, preview_surface_shader_spec, kTokens->inputsNormal,
pxr::SdfValueTypeNames->Normal3f);
// Create the normal map shader and connect its output to the preview
// surface normal attr.
const std::vector<pxr::SdfAttributeSpecHandle> normal_map_output_attrs =
AddTextureShader(material_spec, normal_texture->file->c_str(),
pxr::TfToken("normal"), uvmap_st_output_attr,
{kTokens->outputsRgb});
if (normal_map_output_attrs.size() == 1) {
AddAttributeConnection(layer_, normal_attr,
normal_map_output_attrs[0]);
}
}
}
// Connect an emissive texture if specified.
if (mjTEXROLE_EMISSIVE < textures.size()) {
std::string emissive_texture_name = textures[mjTEXROLE_EMISSIVE];
mjsTexture *emissive_texture = mjs_asTexture(
mjs_findElement(spec_, mjOBJ_TEXTURE, emissive_texture_name.c_str()));
if (emissive_texture) {
auto emissive_attr = CreateAttributeSpec(
layer_, preview_surface_shader_spec, kTokens->inputsEmissiveColor,
pxr::SdfValueTypeNames->Color3f);
const std::vector<pxr::SdfAttributeSpecHandle> emissive_map_output_attrs =
AddTextureShader(material_spec, emissive_texture->file->c_str(),
pxr::TfToken("emissive"), uvmap_st_output_attr,
{kTokens->outputsRgb});
if (emissive_map_output_attrs.size() == 1) {
AddAttributeConnection(layer_, emissive_attr,
emissive_map_output_attrs[0]);
}
}
}
// Set the value of diffuse color. This can come from a diffuse texture
// or as a value defined in mjsMaterial_.
auto diffuse_color_attr = CreateAttributeSpec(
layer_, preview_surface_shader_spec, kTokens->inputsDiffuseColor,
pxr::SdfValueTypeNames->Color3f);
// Find the main texture if specified.
std::string main_texture_name = textures[mjTEXROLE_RGB];
mjsTexture *main_texture = mjs_asTexture(
mjs_findElement(spec_, mjOBJ_TEXTURE, main_texture_name.c_str()));
if (main_texture) {
// Create the texture shader and connect it to the diffuse color
// attribute.
const std::vector<pxr::SdfAttributeSpecHandle>
texture_diffuse_output_attrs =
AddTextureShader(material_spec, main_texture->file->c_str(),
pxr::TfToken("diffuse"), uvmap_st_output_attr,
{kTokens->outputsRgb});
if (texture_diffuse_output_attrs.size() == 1) {
AddAttributeConnection(layer_, diffuse_color_attr,
texture_diffuse_output_attrs[0]);
}
} else {
// If no texture is specified, use the rgba diffuse color.
SetAttributeDefault(layer_, diffuse_color_attr,
pxr::GfVec3f(material->rgba[0], material->rgba[1],
material->rgba[2]));
}
}
void WriteMaterials() {
// Create a scope for the meshes to keep things organized
auto scope =
CreatePrimSpec(layer_, body_paths_[kWorldIndex], kTokens->materialsScope,
pxr::UsdGeomTokens->Scope);
mjsMaterial *material =
mjs_asMaterial(mjs_firstElement(spec_, mjOBJ_MATERIAL));
while (material) {
WriteMaterial(material, scope->GetPath());
material = mjs_asMaterial(mjs_nextElement(spec_, material->element));
}
}
void WriteKeyframesWithName(const std::string &name,
const std::vector<mjsKey *> &keyframes,
const pxr::SdfPath &parent_path) {
if (keyframes.empty()) {
return;
}
const auto keyframe_name = pxr::TfToken(pxr::TfMakeValidIdentifier(
name.empty() ? MjcPhysicsTokens->MjcKeyframe : name));
pxr::SdfPath keyframe_path = parent_path.AppendChild(keyframe_name);
auto keyframe_spec =
layer_->HasSpec(keyframe_path)
? layer_->GetPrimAtPath(keyframe_path)
: CreatePrimSpec(layer_, parent_path, keyframe_name,
pxr::MjcPhysicsTokens->MjcKeyframe);
auto set_attribute_data = [&](const pxr::SdfAttributeSpecHandle &attr_spec,
const pxr::VtDoubleArray &value,
mjsKey *keyframe) {
// If the keyframe time is the default, and there are no other keyframes
// set the attribute at the default time code.
if (keyframe->time == 0 && keyframes.size() == 1) {
SetAttributeDefault(layer_, attr_spec, value);
} else {
SetAttributeTimeSample(layer_, attr_spec, keyframe->time, value);
}
};
for (auto *keyframe : keyframes) {
auto qpos_attr =
CreateAttributeSpec(layer_, keyframe_spec, MjcPhysicsTokens->mjcQpos,
pxr::SdfValueTypeNames->DoubleArray);
set_attribute_data(
qpos_attr,
pxr::VtDoubleArray(keyframe->qpos->begin(), keyframe->qpos->end()),
keyframe);
auto qvel_attr =
CreateAttributeSpec(layer_, keyframe_spec, MjcPhysicsTokens->mjcQvel,
pxr::SdfValueTypeNames->DoubleArray);
set_attribute_data(
qvel_attr,
pxr::VtDoubleArray(keyframe->qvel->begin(), keyframe->qvel->end()),
keyframe);
auto act_attr =
CreateAttributeSpec(layer_, keyframe_spec, MjcPhysicsTokens->mjcAct,
pxr::SdfValueTypeNames->DoubleArray);
set_attribute_data(
act_attr,
pxr::VtDoubleArray(keyframe->act->begin(), keyframe->act->end()),
keyframe);
auto ctrl_attr =
CreateAttributeSpec(layer_, keyframe_spec, MjcPhysicsTokens->mjcCtrl,
pxr::SdfValueTypeNames->DoubleArray);
set_attribute_data(
ctrl_attr,
pxr::VtDoubleArray(keyframe->ctrl->begin(), keyframe->ctrl->end()),
keyframe);
auto mpos_attr =
CreateAttributeSpec(layer_, keyframe_spec, MjcPhysicsTokens->mjcMpos,
pxr::SdfValueTypeNames->DoubleArray);
set_attribute_data(
mpos_attr,
pxr::VtDoubleArray(keyframe->mpos->begin(), keyframe->mpos->end()),
keyframe);
auto mquat_attr_path =
CreateAttributeSpec(layer_, keyframe_spec, MjcPhysicsTokens->mjcMquat,
pxr::SdfValueTypeNames->DoubleArray);
set_attribute_data(
mquat_attr_path,
pxr::VtDoubleArray(keyframe->mquat->begin(), keyframe->mquat->end()),
keyframe);
}
}
void WriteKeyframes() {
std::unordered_map<std::string, std::vector<mjsKey *>> keyframes_map;
mjsKey *keyframe = mjs_asKey(mjs_firstElement(spec_, mjOBJ_KEY));
while (keyframe) {
if (ShouldWrite(keyframe->element)) {
std::string keyframe_name = mjs_getName(keyframe->element)->empty()
? kTokens->keyframe
: *mjs_getName(keyframe->element);
keyframes_map[keyframe_name].push_back(keyframe);
}
keyframe = mjs_asKey(mjs_nextElement(spec_, keyframe->element));
}
auto scope =
CreatePrimSpec(layer_, body_paths_[kWorldIndex], kTokens->keyframesScope,
pxr::UsdGeomTokens->Scope);
for (const auto &[keyframe_name, keyframes] : keyframes_map) {
WriteKeyframesWithName(keyframe_name, keyframes, scope->GetPath());
}
}
void WriteActuator(mjsActuator *actuator, const pxr::SdfPath &parent_path) {
const pxr::TfToken valid_name = GetValidPrimName(*mjs_getName(actuator->element));
const pxr::SdfPath actuator_path = parent_path.AppendChild(valid_name);
pxr::SdfPrimSpecHandle actuator_spec;
if (layer_->HasSpec(actuator_path)) {
actuator_spec = layer_->GetPrimAtPath(actuator_path);
} else {
actuator_spec = CreatePrimSpec(layer_, parent_path, valid_name,
pxr::MjcPhysicsTokens->MjcActuator);
}
pxr::SdfPath target_path;
if (actuator->trntype == mjtTrn::mjTRN_BODY) {
int body_id = mj_name2id(model_, mjOBJ_BODY, actuator->target->c_str());
target_path = body_paths_[body_id];
} else if (actuator->trntype == mjtTrn::mjTRN_SITE ||
actuator->trntype == mjtTrn::mjTRN_SLIDERCRANK) {
int site_id = mj_name2id(model_, mjOBJ_SITE, actuator->target->c_str());
target_path = site_paths_[site_id];
} else if (actuator->trntype == mjtTrn::mjTRN_JOINT) {
int joint_id = mj_name2id(model_, mjOBJ_JOINT, actuator->target->c_str());
target_path = joint_paths_[joint_id];
} else {
TF_WARN(UnsupportedActuatorTypeError,
"Unsupported transmission type for actuator %d",
mjs_getId(actuator->element));
return;
}
CreateRelationshipSpec(layer_, actuator_spec, MjcPhysicsTokens->mjcTarget,
target_path, pxr::SdfVariabilityUniform);
WriteUniformAttribute(actuator_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcGroup, actuator->group);
if (!actuator->refsite->empty()) {
int refsite_id =
mj_name2id(model_, mjOBJ_SITE, actuator->refsite->c_str());
pxr::SdfPath refsite_path = site_paths_[refsite_id];
CreateRelationshipSpec(layer_, actuator_spec, MjcPhysicsTokens->mjcRefSite,
refsite_path, pxr::SdfVariabilityUniform);
}
if (!actuator->slidersite->empty()) {
int slidersite_id =
mj_name2id(model_, mjOBJ_SITE, actuator->slidersite->c_str());
pxr::SdfPath slidersite_path = site_paths_[slidersite_id];
CreateRelationshipSpec(layer_, actuator_spec,
MjcPhysicsTokens->mjcSliderSite, slidersite_path,
pxr::SdfVariabilityUniform);
}
const std::vector<std::pair<pxr::TfToken, int>> limited_attributes = {
{MjcPhysicsTokens->mjcCtrlLimited, actuator->ctrllimited},
{MjcPhysicsTokens->mjcForceLimited, actuator->forcelimited},
{MjcPhysicsTokens->mjcActLimited, actuator->actlimited},
};
for (const auto &[token, value] : limited_attributes) {
pxr::TfToken limited_token = pxr::MjcPhysicsTokens->auto_;
if (value == mjLIMITED_TRUE) {
limited_token = pxr::MjcPhysicsTokens->true_;
} else if (value == mjLIMITED_FALSE) {
limited_token = pxr::MjcPhysicsTokens->false_;
}
WriteUniformAttribute(actuator_spec, pxr::SdfValueTypeNames->Token, token,
limited_token);
}
const std::vector<std::pair<pxr::TfToken, double>>
actuator_double_attributes = {
{MjcPhysicsTokens->mjcCtrlRangeMin, actuator->ctrlrange[0]},
{MjcPhysicsTokens->mjcCtrlRangeMax, actuator->ctrlrange[1]},
{MjcPhysicsTokens->mjcForceRangeMin, actuator->forcerange[0]},
{MjcPhysicsTokens->mjcForceRangeMax, actuator->forcerange[1]},
{MjcPhysicsTokens->mjcActRangeMin, actuator->actrange[0]},
{MjcPhysicsTokens->mjcActRangeMax, actuator->actrange[1]},
{MjcPhysicsTokens->mjcLengthRangeMin, actuator->lengthrange[0]},
{MjcPhysicsTokens->mjcLengthRangeMax, actuator->lengthrange[1]},
{MjcPhysicsTokens->mjcCrankLength, actuator->cranklength},
};
for (const auto &[token, value] : actuator_double_attributes) {
WriteUniformAttribute(actuator_spec, pxr::SdfValueTypeNames->Double,
token, value);
}
WriteUniformAttribute(actuator_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcActDim, actuator->actdim);
WriteUniformAttribute(actuator_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcActEarly,
(bool)actuator->actearly);
WriteUniformAttribute(actuator_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcInheritRange,
actuator->inheritrange);
WriteUniformAttribute(
actuator_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcGear,
pxr::VtDoubleArray(actuator->gear, actuator->gear + 6));
pxr::TfToken dyn_type;
if (actuator->dyntype == mjtDyn::mjDYN_NONE) {
dyn_type = MjcPhysicsTokens->none;
} else if (actuator->dyntype == mjtDyn::mjDYN_INTEGRATOR) {
dyn_type = MjcPhysicsTokens->integrator;
} else if (actuator->dyntype == mjtDyn::mjDYN_FILTER) {
dyn_type = MjcPhysicsTokens->filter;
} else if (actuator->dyntype == mjtDyn::mjDYN_FILTEREXACT) {
dyn_type = MjcPhysicsTokens->filterexact;
} else if (actuator->dyntype == mjtDyn::mjDYN_MUSCLE) {
dyn_type = MjcPhysicsTokens->muscle;
} else if (actuator->dyntype == mjtDyn::mjDYN_USER) {
dyn_type = MjcPhysicsTokens->user;
}
WriteUniformAttribute(actuator_spec, pxr::SdfValueTypeNames->Token,
MjcPhysicsTokens->mjcDynType, dyn_type);
WriteUniformAttribute(
actuator_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcDynPrm,
pxr::VtDoubleArray(actuator->dynprm, actuator->dynprm + 10));
pxr::TfToken gain_type;
if (actuator->gaintype == mjtGain::mjGAIN_FIXED) {
gain_type = MjcPhysicsTokens->fixed;
} else if (actuator->gaintype == mjtGain::mjGAIN_AFFINE) {
gain_type = MjcPhysicsTokens->affine;
} else if (actuator->gaintype == mjtGain::mjGAIN_MUSCLE) {
gain_type = MjcPhysicsTokens->muscle;
} else if (actuator->gaintype == mjtGain::mjGAIN_USER) {
gain_type = MjcPhysicsTokens->user;
}
WriteUniformAttribute(actuator_spec, pxr::SdfValueTypeNames->Token,
MjcPhysicsTokens->mjcGainType, gain_type);
WriteUniformAttribute(
actuator_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcGainPrm,
pxr::VtDoubleArray(actuator->gainprm, actuator->gainprm + 10));
pxr::TfToken bias_type;
if (actuator->biastype == mjtBias::mjBIAS_NONE) {
bias_type = MjcPhysicsTokens->fixed;
} else if (actuator->biastype == mjtBias::mjBIAS_AFFINE) {
bias_type = MjcPhysicsTokens->affine;
} else if (actuator->biastype == mjtBias::mjBIAS_MUSCLE) {
bias_type = MjcPhysicsTokens->muscle;
} else if (actuator->biastype == mjtBias::mjBIAS_USER) {
bias_type = MjcPhysicsTokens->user;
}
WriteUniformAttribute(actuator_spec, pxr::SdfValueTypeNames->Token,
MjcPhysicsTokens->mjcBiasType, bias_type);
WriteUniformAttribute(
actuator_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcBiasPrm,
pxr::VtDoubleArray(actuator->biasprm, actuator->biasprm + 10));
}
void WriteActuators() {
auto scope =
CreatePrimSpec(layer_, body_paths_[kWorldIndex], kTokens->actuatorsScope,
pxr::UsdGeomTokens->Scope);
mjsActuator *actuator =
mjs_asActuator(mjs_firstElement(spec_, mjOBJ_ACTUATOR));
while (actuator) {
if (ShouldWrite(actuator->element)) {
WriteActuator(actuator, scope->GetPath());
}
actuator = mjs_asActuator(mjs_nextElement(spec_, actuator->element));
}
}
pxr::SdfPrimSpecHandle WriteMeshGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
std::string mj_name = mjs_getName(geom->element)->empty()
? *geom->meshname
: *mjs_getName(geom->element);
auto name =
GetAvailablePrimName(mj_name, pxr::UsdGeomTokens->Mesh, body_path);
auto subcomponent_spec =
CreatePrimSpec(layer_, body_path, name, pxr::UsdGeomTokens->Xform);
// Reference the mesh asset written in WriteMeshes.
AddPrimReference(layer_, subcomponent_spec, mesh_paths_[*geom->meshname]);
// We want to use instancing with meshes, and it requires creating a parent
// scope to be referenced, with the Mesh prim as a child.
// To be able to actually manipulate the Mesh prim, we need to create and
// return the corresponding `over` prim as a child of the referencing prim.
auto over_mesh_spec =
CreatePrimSpec(layer_, subcomponent_spec->GetPath(), kTokens->sourceMesh,
pxr::UsdGeomTokens->Mesh, pxr::SdfSpecifierOver);
return over_mesh_spec;
}
pxr::SdfPrimSpecHandle WriteSiteGeom(const mjsSite *site,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*mjs_getName(site->element),
pxr::UsdGeomTokens->Cube, body_path);
int site_idx = mjs_getId(site->element);
const mjtNum *size = &model_->site_size[site_idx * 3];
pxr::SdfPrimSpecHandle site_spec;
switch (site->type) {
case mjGEOM_BOX:
site_spec = WriteBox(name, size, body_path);
break;
case mjGEOM_SPHERE:
site_spec = WriteSphere(name, size, body_path);
break;
case mjGEOM_CAPSULE:
site_spec = WriteCapsule(name, size, body_path);
break;
case mjGEOM_CYLINDER:
site_spec = WriteCylinder(name, size, body_path);
break;
case mjGEOM_ELLIPSOID:
site_spec = WriteEllipsoid(name, size, body_path);
break;
default:
break;
}
return site_spec;
}
pxr::SdfPrimSpecHandle WriteBox(const pxr::TfToken &name, const mjtNum *size,
const pxr::SdfPath &body_path) {
pxr::SdfPrimSpecHandle box_spec =
CreatePrimSpec(layer_, body_path, name, pxr::UsdGeomTokens->Cube);
// MuJoCo uses half sizes. Always set size to 2 (and correspondingly extent
// from -1 to 1), and let scale determine the actual size.
auto size_attr_spec =
CreateAttributeSpec(layer_, box_spec, pxr::UsdGeomTokens->size,
pxr::SdfValueTypeNames->Double);
SetAttributeDefault(layer_, size_attr_spec, 2.0);
auto extent_attr_spec =
CreateAttributeSpec(layer_, box_spec, pxr::UsdGeomTokens->extent,
pxr::SdfValueTypeNames->Float3Array);
SetAttributeDefault(layer_, extent_attr_spec,
pxr::VtArray<pxr::GfVec3f>(
{pxr::GfVec3f(-1, -1, -1), pxr::GfVec3f(1, 1, 1)}));
pxr::GfVec3f scale(static_cast<float>(size[0]), static_cast<float>(size[1]),
static_cast<float>(size[2]));
WriteScaleXformOp(box_spec, scale);
WriteXformOpOrder(box_spec,
pxr::VtArray<pxr::TfToken>{kTokens->xformOpScale});
return box_spec;
}
pxr::SdfPrimSpecHandle WriteBoxGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*mjs_getName(geom->element),
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::SdfPrimSpecHandle WriteCapsule(const pxr::TfToken name, const mjtNum *size,
const pxr::SdfPath &body_path) {
auto capsule_spec =
CreatePrimSpec(layer_, body_path, name, pxr::UsdGeomTokens->Capsule);
auto radius_attr_spec =
CreateAttributeSpec(layer_, capsule_spec, pxr::UsdGeomTokens->radius,
pxr::SdfValueTypeNames->Double);
SetAttributeDefault(layer_, radius_attr_spec, (double)size[0]);
auto height_attr_spec =
CreateAttributeSpec(layer_, capsule_spec, pxr::UsdGeomTokens->height,
pxr::SdfValueTypeNames->Double);
// MuJoCo uses half sizes.
SetAttributeDefault(layer_, height_attr_spec, (double)(size[1] * 2));
return capsule_spec;
}
pxr::SdfPrimSpecHandle WriteCapsuleGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*mjs_getName(geom->element),
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::SdfPrimSpecHandle WriteCylinder(const pxr::TfToken name, const mjtNum *size,
const pxr::SdfPath &body_path) {
auto cylinder_spec =
CreatePrimSpec(layer_, body_path, name, pxr::UsdGeomTokens->Cylinder);
auto radius_attr_spec =
CreateAttributeSpec(layer_, cylinder_spec, pxr::UsdGeomTokens->radius,
pxr::SdfValueTypeNames->Double);
SetAttributeDefault(layer_, radius_attr_spec, (double)size[0]);
auto height_attr_spec =
CreateAttributeSpec(layer_, cylinder_spec, pxr::UsdGeomTokens->height,
pxr::SdfValueTypeNames->Double);
// MuJoCo uses half sizes.
SetAttributeDefault(layer_, height_attr_spec, (double)(size[1] * 2));
return cylinder_spec;
}
pxr::SdfPrimSpecHandle WriteCylinderGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*mjs_getName(geom->element),
pxr::UsdGeomTokens->Cylinder, body_path);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
return WriteCylinder(name, geom_size, body_path);
}
pxr::SdfPrimSpecHandle WriteEllipsoid(const pxr::TfToken name, const mjtNum *size,
const pxr::SdfPath &body_path) {
auto ellipsoid_spec =
CreatePrimSpec(layer_, body_path, name, pxr::UsdGeomTokens->Sphere);
pxr::GfVec3f scale = {static_cast<float>(size[0]),
static_cast<float>(size[1]),
static_cast<float>(size[2])};
auto radius_attr_spec =
CreateAttributeSpec(layer_, ellipsoid_spec, pxr::UsdGeomTokens->radius,
pxr::SdfValueTypeNames->Double);
SetAttributeDefault(layer_, radius_attr_spec, 1.0);
WriteScaleXformOp(ellipsoid_spec, scale);
WriteXformOpOrder(ellipsoid_spec,
pxr::VtArray<pxr::TfToken>{kTokens->xformOpScale});
return ellipsoid_spec;
}
pxr::SdfPrimSpecHandle WriteEllipsoidGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*mjs_getName(geom->element),
pxr::UsdGeomTokens->Sphere, body_path);
int geom_idx = mjs_getId(geom->element);
mjtNum *geom_size = &model_->geom_size[geom_idx * 3];
return WriteEllipsoid(name, geom_size, body_path);
}
pxr::SdfPrimSpecHandle WriteSphere(const pxr::TfToken name, const mjtNum *size,
const pxr::SdfPath &body_path) {
auto sphere_spec =
CreatePrimSpec(layer_, body_path, name, pxr::UsdGeomTokens->Sphere);
auto radius_attr_spec =
CreateAttributeSpec(layer_, sphere_spec, pxr::UsdGeomTokens->radius,
pxr::SdfValueTypeNames->Double);
SetAttributeDefault(layer_, radius_attr_spec, (double)size[0]);
return sphere_spec;
}
pxr::SdfPrimSpecHandle WriteSphereGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*mjs_getName(geom->element),
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);
}
pxr::SdfPrimSpecHandle WritePlane(const pxr::TfToken &name, const mjtNum *size,
const pxr::SdfPath &body_path) {
auto plane_spec =
CreatePrimSpec(layer_, 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;
auto width_attr_spec =
CreateAttributeSpec(layer_, plane_spec, pxr::UsdGeomTokens->width,
pxr::SdfValueTypeNames->Double);
SetAttributeDefault(layer_, width_attr_spec, width);
auto length_attr_spec =
CreateAttributeSpec(layer_, plane_spec, pxr::UsdGeomTokens->length,
pxr::SdfValueTypeNames->Double);
SetAttributeDefault(layer_, length_attr_spec, 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_spec;
}
pxr::SdfPrimSpecHandle WritePlaneGeom(const mjsGeom *geom,
const pxr::SdfPath &body_path) {
auto name = GetAvailablePrimName(*mjs_getName(geom->element),
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];
auto name = GetAvailablePrimName(*mjs_getName(site->element),
pxr::UsdGeomTokens->Xform, body_path);
// Create a geom primitive and set its purpose to guide so it won't be
// rendered.
auto site_spec = WriteSiteGeom(site, body_path);
SetPrimPurpose(layer_, site_spec, pxr::UsdGeomTokens->guide);
ApplyApiSchema(layer_, site_spec, MjcPhysicsTokens->MjcSiteAPI);
WriteUniformAttribute(site_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcGroup, site->group);
WriteColorAndOpacityAttributes(site_spec, site);
int site_id = mjs_getId(site->element);
auto transform = MujocoPosQuatToTransform(&model_->site_pos[3 * site_id],
&model_->site_quat[4 * site_id]);
WriteTransformXformOp(site_spec, transform);
PrependToXformOpOrder(
site_spec, pxr::VtArray<pxr::TfToken>{kTokens->xformOpTransform});
site_paths_[site_id] = site_spec->GetPath();
}
void WriteGeom(mjsGeom *geom, const mjsBody *body) {
const int body_id = mjs_getId(body->element);
const auto &body_path = body_paths_[body_id];
pxr::SdfPrimSpecHandle geom_spec;
int geom_id = mjs_getId(geom->element);
switch (geom->type) {
case mjGEOM_PLANE:
geom_spec = WritePlaneGeom(geom, body_path);
break;
case mjGEOM_MESH:
geom_spec = WriteMeshGeom(geom, body_path);
break;
case mjGEOM_BOX:
geom_spec = WriteBoxGeom(geom, body_path);
break;
case mjGEOM_CAPSULE:
geom_spec = WriteCapsuleGeom(geom, body_path);
break;
case mjGEOM_CYLINDER:
geom_spec = WriteCylinderGeom(geom, body_path);
break;
case mjGEOM_ELLIPSOID:
geom_spec = WriteEllipsoidGeom(geom, body_path);
break;
case mjGEOM_SPHERE:
geom_spec = WriteSphereGeom(geom, body_path);
break;
default:
TF_WARN(UnsupportedGeomTypeError, "Unsupported geom type for geom %d",
geom_id);
return;
}
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcGroup, geom->group);
if (model_->geom_contype[geom_id] == 0 &&
model_->geom_conaffinity[geom_id] == 0) {
// If the geom is purely visual, apply the imageable API.
ApplyApiSchema(layer_, geom_spec, MjcPhysicsTokens->MjcImageableAPI);
}
// Apply the physics schemas if we are writing physics and the
// geom participates in collisions.
if (model_->geom_contype[geom_id] != 0 ||
model_->geom_conaffinity[geom_id] != 0) {
ApplyApiSchema(layer_, geom_spec,
pxr::UsdPhysicsTokens->PhysicsCollisionAPI);
ApplyApiSchema(layer_, geom_spec, MjcPhysicsTokens->MjcCollisionAPI);
WriteUniformAttribute(
geom_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcShellinertia,
geom->typeinertia == mjtGeomInertia::mjINERTIA_SHELL);
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcPriority, geom->priority);
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcCondim, geom->condim);
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcSolmix, geom->solmix);
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcSolmix, geom->solmix);
WriteUniformAttribute(
geom_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcSolref,
pxr::VtArray<double>(geom->solref, geom->solref + mjNREF));
WriteUniformAttribute(
geom_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcSolimp,
pxr::VtArray<double>(geom->solimp, geom->solimp + mjNIMP));
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcMargin, geom->margin);
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcGap, geom->gap);
if (geom->mass >= mjMINVAL || geom->density >= mjMINVAL) {
ApplyApiSchema(layer_, geom_spec, pxr::UsdPhysicsTokens->PhysicsMassAPI);
}
if (geom->mass >= mjMINVAL) {
auto mass_attr_spec = CreateAttributeSpec(
layer_, geom_spec, pxr::UsdPhysicsTokens->physicsMass,
pxr::SdfValueTypeNames->Float, pxr::SdfVariabilityUniform);
// Make sure to cast to float here since mjtNum might be a double.
SetAttributeDefault(layer_, mass_attr_spec, (float)geom->mass);
}
// Even though density is not used for mass computation when mass exists
// we want to retain the information anyways.
if (geom->density >= mjMINVAL) {
auto density_attr_spec = CreateAttributeSpec(
layer_, geom_spec, pxr::UsdPhysicsTokens->physicsDensity,
pxr::SdfValueTypeNames->Float, pxr::SdfVariabilityUniform);
// Make sure to cast to float here since mjtNum might be a double.
SetAttributeDefault(layer_, density_attr_spec, (float)geom->density);
}
mjsDefault *geom_default = mjs_getDefault(geom->element);
if (geom->friction[0] != geom_default->geom->friction[0] ||
geom->friction[1] != geom_default->geom->friction[1] ||
geom->friction[2] != geom_default->geom->friction[2]) {
auto physics_material_spec = WritePhysicsMaterial(geom);
ApplyApiSchema(layer_, geom_spec, pxr::UsdShadeTokens->MaterialBindingAPI);
// Bind the material to this geom.
CreateRelationshipSpec(
layer_, geom_spec,
pxr::UsdShadeTokens->materialBinding, physics_material_spec->GetPath(),
pxr::SdfVariabilityUniform);
}
// For meshes, also apply PhysicsMeshCollisionAPI and set the
// approximation attribute.
if (geom->type == mjGEOM_MESH) {
ApplyApiSchema(layer_, geom_spec,
pxr::UsdPhysicsTokens->PhysicsMeshCollisionAPI);
// Note: MuJoCo documentation states that for collision purposes, meshes
// are always replaced with their convex hulls. Therefore, we set the
// approximation attribute to convexHull explicitly.
auto approximation_attr_spec = CreateAttributeSpec(
layer_, geom_spec, pxr::UsdPhysicsTokens->physicsApproximation,
pxr::SdfValueTypeNames->Token, pxr::SdfVariabilityUniform);
SetAttributeDefault(layer_, approximation_attr_spec,
pxr::UsdPhysicsTokens->convexHull);
}
} else {
// Currently imageable only has a group API. But since it's the same
// naming in MjcCollisionsAPI we've already set it earlier in this
// function.
ApplyApiSchema(layer_, geom_spec, MjcPhysicsTokens->MjcImageableAPI);
}
mjsDefault *spec_default = mjs_getDefault(geom->element);
pxr::TfToken valid_class_name =
GetValidPrimName(*mjs_getName(spec_default->element));
pxr::SdfPath geom_class_path = class_path_.AppendChild(valid_class_name);
if (!layer_->HasSpec(geom_class_path)) {
auto class_spec =
CreateClassSpec(layer_, class_path_, valid_class_name);
auto visibility_attr_spec =
CreateAttributeSpec(layer_, class_spec, pxr::UsdGeomTokens->visibility,
pxr::SdfValueTypeNames->Token);
SetAttributeDefault(layer_, visibility_attr_spec,
pxr::UsdGeomTokens->inherited);
}
// Bind material if it exists.
if (!geom->material->empty()) {
pxr::SdfPath material_path =
body_paths_[kWorldIndex]
.AppendChild(kTokens->materialsScope)
.AppendChild(GetValidPrimName(*geom->material));
if (layer_->HasSpec(material_path)) {
ApplyApiSchema(layer_, geom_spec,
pxr::UsdShadeTokens->MaterialBindingAPI);
// Bind the material to this geom.
CreateRelationshipSpec(layer_, geom_spec,
pxr::UsdShadeTokens->materialBinding,
material_path, pxr::SdfVariabilityUniform);
}
}
WriteColorAndOpacityAttributes(geom_spec, geom);
if (body_id == kWorldIndex) {
SetPrimKind(layer_, geom_spec, pxr::KindTokens->component);
}
// Inherit from class.
AddPrimInherit(layer_, geom_spec, geom_class_path);
auto transform = MujocoPosQuatToTransform(&model_->geom_pos[3 * geom_id],
&model_->geom_quat[4 * geom_id]);
WriteTransformXformOp(geom_spec, transform);
PrependToXformOpOrder(
geom_spec, pxr::VtArray<pxr::TfToken>{kTokens->xformOpTransform});
}
void WriteSites(mjsBody *body) {
mjsSite *site = mjs_asSite(mjs_firstChild(body, mjOBJ_SITE, false));
while (site) {
if (ShouldWrite(site->element)) {
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) {
if (ShouldWrite(geom->element)) {
WriteGeom(geom, body);
}
geom = mjs_asGeom(mjs_nextChild(body, geom->element, false));
}
}
void WriteJoints(mjsBody *body) {
int body_id = mjs_getId(body->element);
if (body_id == kWorldIndex) return;
mjsJoint *joint = mjs_asJoint(mjs_firstChild(body, mjOBJ_JOINT, false));
if (!joint) {
// If no joint is found, then we pass nullptr to create a FixedJoint.
// WriteJoint properly handles the case where the parent is the worldbody.
WriteJoint(nullptr, body);
} else {
if (ShouldWrite(joint->element)) {
WriteJoint(joint, body);
}
if (mjs_asJoint(mjs_nextChild(body, joint->element, false))) {
TF_WARN(
"Multiple joints found for body %d. Only writing the first one.",
body_id);
}
}
}
// Write the joint. If null, then a FixedJoint is created.
void WriteJoint(mjsJoint *joint, const mjsBody *parent_mj_body) {
// Default to fixed joint if joint is null.
pxr::TfToken joint_prim_type = pxr::UsdPhysicsTokens->PhysicsFixedJoint;
int joint_id = -1;
if (joint) {
joint_id = mjs_getId(joint->element);
mjtJoint type = (mjtJoint)model_->jnt_type[joint_id];
switch (type) {
case mjJNT_FREE:
// Free joints are guaranteed to only ever be on the top-level body so
// we just write no joint. As a top-level body with no joint it will
// be considered a floating-base body.
return;
case mjJNT_HINGE:
joint_prim_type = pxr::UsdPhysicsTokens->PhysicsRevoluteJoint;
break;
case mjJNT_SLIDE:
joint_prim_type = pxr::UsdPhysicsTokens->PhysicsPrismaticJoint;
break;
default:
TF_WARN("Unsupported joint type '%d' for joint '%s'. Skipping.",
(int)type, mjs_getName(joint->element)->c_str());
return;
}
}
int body_id = mjs_getId(parent_mj_body->element);
// the joint connects the current body as body1, to its parent body as
// body0.
int body1_id_usd = body_id;
int body0_id_usd = model_->body_parentid[body_id];
const pxr::SdfPath &body1_path_usd = body_paths_[body1_id_usd];
auto joint_name = joint ? *mjs_getName(joint->element) : "FixedJoint";
pxr::TfToken joint_name_token =
GetAvailablePrimName(joint_name, kTokens->joint, body1_path_usd);
auto joint_spec = CreatePrimSpec(layer_, body1_path_usd,
joint_name_token, joint_prim_type);
// Set body0 and body1 relationships
// For the initial joints that connect to the world, we signal this by
// keeping the body0 relationship empty.
if (body0_id_usd != kWorldIndex) {
const pxr::SdfPath &body0_path_usd = body_paths_[body0_id_usd];
CreateRelationshipSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsBody0,
body0_path_usd, pxr::SdfVariabilityUniform);
}
CreateRelationshipSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsBody1, body1_path_usd,
pxr::SdfVariabilityUniform);
// Joint frame in MuJoCo is defined by jnt_pos and jnt_axis in body1's frame
// For FixedJoint, these are both unity.
pxr::GfVec3d mj_jnt_pos = pxr::GfVec3d(0.0);
pxr::GfVec3d mj_jnt_axis = pxr::GfVec3d(0.0, 0.0, 1.0);
if (joint) {
mj_jnt_pos = pxr::GfVec3d(&model_->jnt_pos[joint_id * 3]);
mj_jnt_axis = pxr::GfVec3d(&model_->jnt_axis[joint_id * 3]);
}
// Local joint frame for body1
pxr::GfVec3f local_pos1(mj_jnt_pos);
pxr::GfQuatf local_rot1(
pxr::GfRotation(pxr::GfVec3f::ZAxis(), mj_jnt_axis).GetQuat());
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsLocalPos1,
pxr::SdfValueTypeNames->Float3),
local_pos1);
if (joint_prim_type == pxr::UsdPhysicsTokens->PhysicsRevoluteJoint ||
joint_prim_type == pxr::UsdPhysicsTokens->PhysicsPrismaticJoint) {
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsLocalRot1,
pxr::SdfValueTypeNames->Quatf),
local_rot1);
} else {
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsLocalRot1,
pxr::SdfValueTypeNames->Quatf),
pxr::GfQuatf::GetIdentity());
}
// Calculate local joint frame for body0
pxr::GfMatrix4d body1_transform_local =
MujocoPosQuatToTransform(&model_->body_pos[body1_id_usd * 3],
&model_->body_quat[body1_id_usd * 4]);
pxr::GfVec3d jnt_pos_parent_local =
body1_transform_local.Transform(mj_jnt_pos);
pxr::GfVec3d jnt_axis_parent_local =
body1_transform_local.TransformDir(mj_jnt_axis);
pxr::GfVec3f local_pos0(jnt_pos_parent_local);
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsLocalPos0,
pxr::SdfValueTypeNames->Float3),
local_pos0);
if (joint_prim_type == pxr::UsdPhysicsTokens->PhysicsRevoluteJoint ||
joint_prim_type == pxr::UsdPhysicsTokens->PhysicsPrismaticJoint) {
pxr::GfQuatf other_rot0(
pxr::GfRotation(pxr::GfVec3f::ZAxis(), jnt_axis_parent_local)
.GetQuat());
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsLocalRot0,
pxr::SdfValueTypeNames->Quatf),
other_rot0);
} else {
// Fixed joints have no frame and no axis per se. We simply need the
// rotation quaternion of the body its on.
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsLocalRot0,
pxr::SdfValueTypeNames->Quatf),
body1_transform_local.ExtractRotationQuat());
}
if (joint) {
mjtJoint type = (mjtJoint)model_->jnt_type[joint_id];
// Joint-specific attributes
if (type == mjJNT_HINGE || type == mjJNT_SLIDE) {
// The joint motion occurs around/along the Z-axis of the joint frame
// established by localRot0/1.
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsAxis,
pxr::SdfValueTypeNames->Token),
pxr::UsdPhysicsTokens->z); // "Z" axis
}
if (model_->jnt_limited[joint_id]) {
float lower_limit = model_->jnt_range[joint_id * 2];
float upper_limit = model_->jnt_range[joint_id * 2 + 1];
if (type == mjJNT_HINGE) {
// Convert radians to degrees for USD
// As per the XML Reference, "mjModel always uses radians"
lower_limit *= (180.0 / mjPI);
upper_limit *= (180.0 / mjPI);
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsLowerLimit,
pxr::SdfValueTypeNames->Float),
lower_limit);
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsUpperLimit,
pxr::SdfValueTypeNames->Float),
upper_limit);
} else if (type == mjJNT_SLIDE) {
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsLowerLimit,
pxr::SdfValueTypeNames->Float),
lower_limit);
SetAttributeDefault(
layer_,
CreateAttributeSpec(layer_, joint_spec,
pxr::UsdPhysicsTokens->physicsUpperLimit,
pxr::SdfValueTypeNames->Float),
upper_limit);
}
}
// Finally write the mjcPhysicsJointAPI attributes.
ApplyApiSchema(layer_, joint_spec, MjcPhysicsTokens->MjcJointAPI);
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcGroup, joint->group);
WriteUniformAttribute(
joint_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcSpringdamper,
pxr::VtArray<double>(joint->springdamper, joint->springdamper + 2));
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcSolreflimit,
pxr::VtArray<double>(joint->solref_limit,
joint->solref_limit + mjNREF));
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcSolimplimit,
pxr::VtArray<double>(joint->solimp_limit,
joint->solimp_limit + mjNIMP));
WriteUniformAttribute(
joint_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcSolreffriction,
pxr::VtArray<double>(joint->solref_friction,
joint->solref_friction + mjNREF));
WriteUniformAttribute(
joint_spec, pxr::SdfValueTypeNames->DoubleArray,
MjcPhysicsTokens->mjcSolimpfriction,
pxr::VtArray<double>(joint->solimp_friction,
joint->solimp_friction + mjNIMP));
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcStiffness, joint->stiffness);
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcActuatorfrcrangeMin,
joint->actfrcrange[0]);
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcActuatorfrcrangeMax,
joint->actfrcrange[1]);
pxr::TfToken actuatorfrclimited_token = MjcPhysicsTokens->auto_;
if (joint->actfrclimited == mjLIMITED_TRUE) {
actuatorfrclimited_token = MjcPhysicsTokens->true_;
} else if (joint->actfrclimited == mjLIMITED_FALSE) {
actuatorfrclimited_token = MjcPhysicsTokens->false_;
}
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Token,
MjcPhysicsTokens->mjcActuatorfrclimited,
actuatorfrclimited_token);
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcActuatorgravcomp,
static_cast<bool>(joint->actgravcomp));
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcMargin, joint->margin);
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcRef, joint->ref);
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcSpringref, joint->springref);
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcArmature, joint->armature);
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcDamping, joint->damping);
WriteUniformAttribute(joint_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcFrictionloss,
joint->frictionloss);
}
if (joint_id >= 0) {
joint_paths_[joint_id] = joint_spec->GetPath();
}
}
void WriteCamera(mjsCamera *spec_cam, const mjsBody *body) {
const auto &body_path = body_paths_[mjs_getId(body->element)];
auto name = GetAvailablePrimName(*mjs_getName(spec_cam->element),
pxr::UsdGeomTokens->Camera, body_path);
// Create a root Xform for the world body with the model name if it exists
// otherwise called 'World'.
auto camera_spec =
CreatePrimSpec(layer_, body_path, name, pxr::UsdGeomTokens->Camera);
int cam_id = mjs_getId(spec_cam->element);
auto transform = MujocoPosQuatToTransform(&model_->cam_pos[3 * cam_id],
&model_->cam_quat[4 * cam_id]);
WriteTransformXformOp(camera_spec, transform);
WriteXformOpOrder(camera_spec,
pxr::VtArray<pxr::TfToken>{kTokens->xformOpTransform});
// If the camera intrinsics are specified, then it is important that we
// reproduce the code in mujoco/src/engine/engine_vis_visualize.c
const float *cam_sensorsize = &model_->cam_sensorsize[cam_id * 2];
bool use_intrinsic = cam_sensorsize[1] > 0.0f;
float znear = spec_->visual.map.znear * model_->stat.extent * 100;
float zfar = spec_->visual.map.zfar * model_->stat.extent * 100;
mjtNum fovy = model_->cam_fovy[cam_id];
const float *cam_intrinsic = &model_->cam_intrinsic[cam_id * 4];
const float aspect_ratio =
use_intrinsic ? cam_sensorsize[0] / cam_sensorsize[1] : 4.0f / 3;
float vertical_apperture =
2 * znear *
(use_intrinsic ? 1.0f / cam_intrinsic[1] *
(cam_sensorsize[1] / 2.f - cam_intrinsic[3])
: mju_tan((fovy / 2) * (std::numbers::pi / 180.0)));
float horizontal_aperture =
use_intrinsic ? 2 * znear / cam_intrinsic[0] *
(cam_sensorsize[0] / 2.f - cam_intrinsic[2])
: vertical_apperture * aspect_ratio;
WriteUniformAttribute(camera_spec, pxr::SdfValueTypeNames->Float2,
pxr::UsdGeomTokens->clippingRange,
pxr::GfVec2f(znear, zfar));
WriteUniformAttribute(camera_spec, pxr::SdfValueTypeNames->Float,
pxr::UsdGeomTokens->focalLength, znear);
WriteUniformAttribute(camera_spec, pxr::SdfValueTypeNames->Float,
pxr::UsdGeomTokens->verticalAperture,
vertical_apperture);
WriteUniformAttribute(camera_spec, pxr::SdfValueTypeNames->Float,
pxr::UsdGeomTokens->horizontalAperture,
horizontal_aperture);
}
void WriteCameras(mjsBody *body) {
mjsCamera *cam = mjs_asCamera(mjs_firstChild(body, mjOBJ_CAMERA, false));
while (cam) {
if (ShouldWrite(cam->element)) {
WriteCamera(cam, body);
}
cam = mjs_asCamera(mjs_nextChild(body, cam->element, false));
}
}
void WriteLight(mjsLight *light, const mjsBody *body) {
const auto &body_path = body_paths_[mjs_getId(body->element)];
auto name = GetAvailablePrimName(*mjs_getName(light->element),
kTokens->light, body_path);
// Create a root Xform for the world body with the model name if it exists
// otherwise called 'World'.
auto light_spec =
CreatePrimSpec(layer_, body_path, name, pxr::UsdLuxTokens->SphereLight);
int light_id = mjs_getId(light->element);
auto transform = MujocoPosQuatToTransform(&model_->light_pos[3 * light_id],
&model_->light_dir[4 * light_id]);
WriteTransformXformOp(light_spec, transform);
WriteXformOpOrder(light_spec,
pxr::VtArray<pxr::TfToken>{kTokens->xformOpTransform});
}
void WriteLights(mjsBody *body) {
mjsLight *light = mjs_asLight(mjs_firstChild(body, mjOBJ_LIGHT, false));
while (light) {
if (ShouldWrite(light->element)) {
WriteLight(light, body);
}
light = mjs_asLight(mjs_nextChild(body, light->element, false));
}
}
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);
int parent_id = mjs_getId(parent->element);
pxr::SdfPath parent_path = body_paths_[parent_id];
pxr::TfToken body_name = GetValidPrimName(*mjs_getName(body->element));
// Create Xform prim for body.
auto body_spec = CreatePrimSpec(layer_, parent_path, body_name,
pxr::UsdGeomTokens->Xform);
// The parent_path will be a component which makes the actual articulated
// bodies subcomponents.
auto kind = parent_id == kWorldIndex ? pxr::KindTokens->component
: pxr::KindTokens->subcomponent;
SetPrimKind(layer_, body_spec, kind);
// If the parent is not the world body, but is child of the world body
// then we need to apply the articulation root API.
if (parent_id != kWorldIndex) {
int parent_parent_id = mjs_getId(mjs_getParent(parent->element)->element);
// We guard against applying the API more than once, which can happen when
// there are multiple children.
if (parent_parent_id == kWorldIndex &&
articulation_roots_.find(parent_id) == articulation_roots_.end()) {
ApplyApiSchema(layer_, layer_->GetPrimAtPath(parent_path),
pxr::UsdPhysicsTokens->PhysicsArticulationRootAPI);
articulation_roots_.insert(parent_id);
}
}
// 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
// we should set inertial attributes on the body.
//
// Note that if the user has NOT specified any inertial properties then
// we don't want to pull values from the compiled model since coming back
// into Mujoco would take those values instead of computing them
// automatically from the subtree.
if (body->mass > 0) {
// User might have specified the inertia via fullinertia and the
// compiler has extracted all values properly. So leverage those
// instead of doing the computation ourselves here.
ApplyApiSchema(layer_, body_spec, pxr::UsdPhysicsTokens->PhysicsMassAPI);
WriteUniformAttribute(body_spec, pxr::SdfValueTypeNames->Float,
pxr::UsdPhysicsTokens->physicsMass,
(float)model_->body_mass[body_id]);
mjtNum *body_ipos = &model_->body_ipos[body_id * 3];
pxr::GfVec3f inertial_pos(body_ipos[0], body_ipos[1], body_ipos[2]);
WriteUniformAttribute(body_spec, pxr::SdfValueTypeNames->Point3f,
pxr::UsdPhysicsTokens->physicsCenterOfMass,
inertial_pos);
mjtNum *body_iquat = &model_->body_iquat[body_id * 4];
pxr::GfQuatf inertial_frame(body_iquat[0], body_iquat[1], body_iquat[2],
body_iquat[3]);
WriteUniformAttribute(body_spec, pxr::SdfValueTypeNames->Quatf,
pxr::UsdPhysicsTokens->physicsPrincipalAxes,
inertial_frame);
mjtNum *inertia = &model_->body_inertia[body_id * 3];
pxr::GfVec3f diag_inertia(inertia[0], inertia[1], inertia[2]);
WriteUniformAttribute(body_spec, pxr::SdfValueTypeNames->Float3,
pxr::UsdPhysicsTokens->physicsDiagonalInertia,
diag_inertia);
}
ApplyApiSchema(layer_, body_spec,
pxr::UsdPhysicsTokens->PhysicsRigidBodyAPI);
// Create classes if necessary
mjsDefault *spec_default = mjs_getDefault(body->element);
pxr::TfToken body_class_name =
GetValidPrimName(*mjs_getName(spec_default->element));
pxr::SdfPath body_class_path = class_path_.AppendChild(body_class_name);
if (!layer_->HasSpec(body_class_path)) {
CreateClassSpec(layer_, class_path_, body_class_name);
}
// Create XformOp attribute for body transform.
auto xform_op_spec =
CreateAttributeSpec(layer_, body_spec, kTokens->xformOpTransform,
pxr::SdfValueTypeNames->Matrix4d);
// mjModel will have all frames already accounted for so no need to worry
// about them here.
auto body_xform = MujocoPosQuatToTransform(&model_->body_pos[body_id * 3],
&model_->body_quat[body_id * 4]);
SetAttributeDefault(layer_, xform_op_spec, body_xform);
// Create XformOpOrder attribute for body transform order.
// For us this is simply the transform we authored above.
WriteXformOpOrder(body_spec,
pxr::VtArray<pxr::TfToken>{kTokens->xformOpTransform});
pxr::VtDictionary customData;
customData[kTokens->body_name] = *mjs_getName(body->element);
body_spec->SetField(pxr::SdfFieldKeys->CustomData, customData);
body_paths_[body_id] = body_spec->GetPath();
mujoco::usd::SetUsdPrimPathUserValue(body->element, body_spec->GetPath());
}
void WriteBodies() {
mjsBody *body = mjs_asBody(mjs_firstElement(spec_, mjOBJ_BODY));
while (body) {
// If this body originally came from USD, don't try and write it
// again.
if (ShouldWrite(body->element)) {
// 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);
}
WriteSites(body);
WriteGeoms(body);
WriteJoints(body);
WriteCameras(body);
WriteLights(body);
}
body = mjs_asBody(mjs_nextElement(spec_, body->element));
}
}
pxr::SdfPrimSpecHandle WriteWorldBody(const size_t body_index) {
// Create a root Xform for the world body with the model name if it exists
// otherwise called 'World'.
auto name = GetAvailablePrimName(*spec_->modelname, kTokens->world,
pxr::SdfPath::AbsoluteRootPath());
auto world_group_spec =
CreatePrimSpec(layer_, pxr::SdfPath::AbsoluteRootPath(), name,
pxr::UsdGeomTokens->Xform);
SetPrimKind(layer_, world_group_spec, pxr::KindTokens->group);
return world_group_spec;
}
};
namespace mujoco {
namespace usd {
bool WriteSpecToData(mjSpec *spec, pxr::SdfLayerRefPtr layer, bool skip_elems_from_usd) {
mjModel *model = mj_compile(spec, nullptr);
if (model == nullptr) {
TF_ERROR(MujocoCompilationError, "%s", mjs_getError(spec));
return false;
}
ModelWriter(spec, model, layer, skip_elems_from_usd).Write();
return true;
}
} // namespace usd
} // namespace mujoco
Reference in New Issue View Git Blame Copy Permalink