Move OpenUSD parsing to usd_decoder
Add usd_decoder to mujoco/plugin. - Previously was src/experimental/usd/usd_to_mjspec.cc, now that same spec is returned by the plugin. PiperOrigin-RevId: 854350962 Change-Id: Ia980190a557c50ce8197980922a2594b613859b4
This commit is contained in:
Sam Haves
committed by
Copybara-Service
Copybara-Service
parent
64a2345c07
commit
a5dc57c0c3
@@ -0,0 +1,53 @@
|
||||
# Copyright 2026 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
|
||||
#
|
||||
# https://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.
|
||||
find_package(pxr REQUIRED)
|
||||
|
||||
set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE)
|
||||
set(USD_DECODER_SRCS
|
||||
usd_decoder.cc
|
||||
kinematic_tree.cc
|
||||
kinematic_tree.h
|
||||
material_parsing.cc
|
||||
material_parsing.h
|
||||
utils.h
|
||||
)
|
||||
|
||||
add_library(usd_decoder_plugin SHARED ${USD_DECODER_SRCS})
|
||||
target_include_directories(usd_decoder_plugin PRIVATE
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/../..
|
||||
)
|
||||
target_link_libraries(usd_decoder_plugin PRIVATE
|
||||
mujoco
|
||||
mjcPhysics
|
||||
)
|
||||
target_compile_options(usd_decoder_plugin PRIVATE
|
||||
${AVX_COMPILE_OPTIONS}
|
||||
${MUJOCO_MACOS_COMPILE_OPTIONS}
|
||||
${EXTRA_COMPILE_OPTIONS}
|
||||
${MUJOCO_CXX_FLAGS}
|
||||
-Wno-deprecated # pxr Tf lib uses deprecated header
|
||||
)
|
||||
target_link_options(usd_decoder_plugin PRIVATE
|
||||
${MUJOCO_MACOS_LINK_OPTIONS}
|
||||
${EXTRA_LINK_OPTIONS}
|
||||
)
|
||||
|
||||
# Install to mujoco_plugin directory in bin location so that it is picked up by simulate
|
||||
# on startup.
|
||||
install(
|
||||
TARGETS usd_decoder_plugin
|
||||
LIBRARY DESTINATION "${CMAKE_INSTALL_BINDIR}/mujoco_plugin"
|
||||
)
|
||||
|
||||
target_link_libraries(usd_decoder_plugin PRIVATE usd usdGeom usdPhysics usdShade gf tf ar vt kind)
|
||||
@@ -0,0 +1,273 @@
|
||||
// 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 "kinematic_tree.h"
|
||||
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include <mujoco/experimental/usd/mjcPhysics/actuator.h>
|
||||
#include <mujoco/experimental/usd/mjcPhysics/keyframe.h>
|
||||
#include <mujoco/experimental/usd/mjcPhysics/siteAPI.h>
|
||||
#include <mujoco/experimental/usd/mjcPhysics/tendon.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#include <pxr/usd/sdf/path.h>
|
||||
#include <pxr/usd/usd/common.h>
|
||||
#include <pxr/usd/usd/primRange.h>
|
||||
#include <pxr/usd/usdGeom/gprim.h>
|
||||
#include <pxr/usd/usdGeom/xformCache.h>
|
||||
#include <pxr/usd/usdPhysics/collisionAPI.h>
|
||||
#include <pxr/usd/usdPhysics/joint.h>
|
||||
#include <pxr/usd/usdPhysics/rigidBodyAPI.h>
|
||||
#include <pxr/usd/usdPhysics/scene.h>
|
||||
|
||||
bool GetJointBodies(const pxr::UsdPhysicsJoint& joint, pxr::SdfPath* from,
|
||||
pxr::SdfPath* to) {
|
||||
// Grab the default prim path
|
||||
pxr::SdfPath default_prim_path;
|
||||
auto stage = joint.GetPrim().GetStage();
|
||||
if (stage->GetDefaultPrim().IsValid()) {
|
||||
default_prim_path = stage->GetDefaultPrim().GetPath();
|
||||
}
|
||||
|
||||
pxr::SdfPathVector body1_paths;
|
||||
joint.GetBody1Rel().GetTargets(&body1_paths);
|
||||
if (body1_paths.empty()) {
|
||||
mju_warning("Joint %s does not have body1 rel. Skipping.",
|
||||
joint.GetPath().GetAsString().c_str());
|
||||
return false;
|
||||
} else if (body1_paths.size() > 1) {
|
||||
mju_warning("Joint %s has multiple body1 rels. Skipping.",
|
||||
joint.GetPath().GetAsString().c_str());
|
||||
return false;
|
||||
}
|
||||
*to = body1_paths[0];
|
||||
|
||||
pxr::SdfPathVector body0_paths;
|
||||
joint.GetBody0Rel().GetTargets(&body0_paths);
|
||||
if (body0_paths.size() > 1) {
|
||||
mju_warning("Joint %s has multiple body0 rels. Skipping.",
|
||||
joint.GetPath().GetAsString().c_str());
|
||||
return false;
|
||||
}
|
||||
// Empty body0, or body0 pointing to the default prim means we'll attach
|
||||
// to the worldbody.
|
||||
if (body0_paths.empty() || body0_paths[0] == default_prim_path) {
|
||||
*from = pxr::SdfPath();
|
||||
} else {
|
||||
*from = body0_paths[0];
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
struct ExtractedPrims {
|
||||
std::vector<std::unique_ptr<Node>> nodes;
|
||||
std::vector<pxr::UsdPhysicsJoint> joints;
|
||||
};
|
||||
|
||||
ExtractedPrims ExtractPrims(pxr::UsdStageRefPtr stage) {
|
||||
pxr::UsdPhysicsScene physics_scene;
|
||||
std::vector<std::unique_ptr<Node>> nodes;
|
||||
std::vector<pxr::UsdPhysicsJoint> joints;
|
||||
nodes.push_back(std::make_unique<Node>());
|
||||
Node* root = nodes.back().get();
|
||||
|
||||
// =========================================================================
|
||||
// PASS 1: Collect Bodies, Joints, and Geoms/Sites/etc.
|
||||
// =========================================================================
|
||||
// A single DFS pass to find all bodies, joints, and determine
|
||||
// which body owns each geom/site/etc. prim.
|
||||
std::vector<Node*> owner_stack;
|
||||
owner_stack.push_back(root); // Start with the world as owner.
|
||||
|
||||
const auto range = pxr::UsdPrimRange::PreAndPostVisit(
|
||||
stage->GetPseudoRoot(), pxr::UsdTraverseInstanceProxies());
|
||||
|
||||
pxr::UsdGeomXformCache xform_cache;
|
||||
for (auto it = range.begin(); it != range.end(); ++it) {
|
||||
pxr::UsdPrim prim = *it;
|
||||
|
||||
bool is_body = prim.HasAPI<pxr::UsdPhysicsRigidBodyAPI>();
|
||||
bool resets = xform_cache.GetResetXformStack(prim);
|
||||
// Only update (push/pop) the owner stack for bodies (becomes new owner) and
|
||||
// resetXformStack (reset owner to world).
|
||||
bool is_pushed_to_stack = is_body || resets;
|
||||
|
||||
if (it.IsPostVisit()) {
|
||||
if (is_pushed_to_stack) {
|
||||
owner_stack.pop_back();
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
pxr::SdfPath prim_path = prim.GetPath();
|
||||
Node* current_node = owner_stack.back();
|
||||
|
||||
if (is_body) {
|
||||
auto new_node = std::make_unique<Node>();
|
||||
new_node->body_path = prim_path;
|
||||
nodes.push_back(std::move(new_node));
|
||||
current_node = nodes.back().get();
|
||||
} else if (resets) {
|
||||
current_node = root; // Reset owner to world.
|
||||
}
|
||||
|
||||
if (is_pushed_to_stack) {
|
||||
owner_stack.push_back(current_node);
|
||||
}
|
||||
|
||||
if (prim.IsA<pxr::UsdPhysicsScene>() && root->physics_scene.IsEmpty()) {
|
||||
root->physics_scene = prim_path;
|
||||
}
|
||||
|
||||
if (prim.IsA<pxr::UsdGeomGprim>()) {
|
||||
bool has_collision_api = prim.HasAPI<pxr::UsdPhysicsCollisionAPI>();
|
||||
if (has_collision_api) {
|
||||
bool collision_enabled = false;
|
||||
pxr::UsdPhysicsCollisionAPI(prim).GetCollisionEnabledAttr().Get(
|
||||
&collision_enabled);
|
||||
if (collision_enabled) {
|
||||
current_node->colliders.push_back(prim_path);
|
||||
} else {
|
||||
current_node->visual_gprims.push_back(prim_path);
|
||||
}
|
||||
} else {
|
||||
current_node->visual_gprims.push_back(prim_path);
|
||||
}
|
||||
}
|
||||
|
||||
if (prim.HasAPI<pxr::MjcPhysicsSiteAPI>()) {
|
||||
current_node->sites.push_back(prim_path);
|
||||
// Sites should not have children.
|
||||
it.PruneChildren();
|
||||
}
|
||||
|
||||
if (prim.IsA<pxr::UsdPhysicsJoint>()) {
|
||||
// We may not know which body this belongs to yet so we'll add it to a
|
||||
// list and the caller can assign the joints when building the tree.
|
||||
joints.push_back(pxr::UsdPhysicsJoint(prim));
|
||||
// Joints should not have children.
|
||||
it.PruneChildren();
|
||||
}
|
||||
|
||||
if (prim.IsA<pxr::MjcPhysicsActuator>()) {
|
||||
root->actuators.push_back(prim_path);
|
||||
// Joints should not have children.
|
||||
it.PruneChildren();
|
||||
}
|
||||
|
||||
if (prim.IsA<pxr::MjcPhysicsTendon>()) {
|
||||
root->tendons.push_back(prim_path);
|
||||
// Tendons should not have children of interest.
|
||||
it.PruneChildren();
|
||||
}
|
||||
|
||||
if (prim.IsA<pxr::MjcPhysicsKeyframe>()) {
|
||||
root->keyframes.push_back(prim_path);
|
||||
// Keyframes should not have children.
|
||||
it.PruneChildren();
|
||||
}
|
||||
}
|
||||
return {std::move(nodes), std::move(joints)};
|
||||
}
|
||||
|
||||
std::unique_ptr<Node> BuildKinematicTree(const pxr::UsdStageRefPtr stage) {
|
||||
ExtractedPrims extraction = ExtractPrims(stage);
|
||||
|
||||
std::map<pxr::SdfPath, int> body_index;
|
||||
body_index[pxr::SdfPath()] = 0;
|
||||
for (int i = 0; i < extraction.nodes.size(); ++i) {
|
||||
body_index[extraction.nodes[i]->body_path] = i;
|
||||
}
|
||||
|
||||
// List of direct children for each body.
|
||||
std::vector<std::vector<bool>> children(
|
||||
extraction.nodes.size(), std::vector<bool>(extraction.nodes.size()));
|
||||
// List of joint prim paths associated with a child for each body.
|
||||
std::vector<std::vector<pxr::SdfPath>> parent_joints(extraction.nodes.size());
|
||||
for (const pxr::UsdPhysicsJoint& joint : extraction.joints) {
|
||||
pxr::SdfPath from, to;
|
||||
if (!GetJointBodies(joint, &from, &to)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
int from_idx = body_index[from];
|
||||
int to_idx = body_index[to];
|
||||
if (from_idx == to_idx) {
|
||||
mju_error("Cycle detected: self referencing joint at node %s",
|
||||
to.GetString().c_str());
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
children[from_idx][to_idx] = true;
|
||||
parent_joints[to_idx].push_back(joint.GetPath());
|
||||
// Now that we know all the bodies, we can assign joints to respective
|
||||
// nodes.
|
||||
extraction.nodes[to_idx]->joints.push_back(joint.GetPath());
|
||||
}
|
||||
|
||||
// The world body is represented by an empty SdfPath.
|
||||
auto world_root = std::move(extraction.nodes[0]);
|
||||
|
||||
std::vector<std::pair<int, Node*>> stack;
|
||||
stack.emplace_back(0, world_root.get());
|
||||
|
||||
// A node without any joints from a parent has a free joint.
|
||||
// Add all free joints as children of the world body.
|
||||
for (int i = 1; i < extraction.nodes.size(); ++i) {
|
||||
if (parent_joints[i].empty()) {
|
||||
children[0][i] = true;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<bool> visited(extraction.nodes.size());
|
||||
while (!stack.empty()) {
|
||||
auto [current_body, parent] = stack.back();
|
||||
stack.pop_back();
|
||||
visited[current_body] = true;
|
||||
|
||||
Node* current_node = nullptr;
|
||||
if (current_body > 0) {
|
||||
parent->children.push_back(std::move(extraction.nodes[current_body]));
|
||||
current_node = parent->children.back().get();
|
||||
} else {
|
||||
current_node = world_root.get();
|
||||
}
|
||||
|
||||
// Process children in reverse to maintain DFS.
|
||||
for (int i = extraction.nodes.size() - 1; i > 0; --i) {
|
||||
if (!children[current_body][i]) {
|
||||
continue;
|
||||
}
|
||||
if (visited[i]) {
|
||||
mju_error("Cycle detected in the kinematic tree at node %s",
|
||||
current_node->body_path.GetString().c_str());
|
||||
return nullptr;
|
||||
}
|
||||
stack.emplace_back(
|
||||
i, current_body > 0 ? parent->children.back().get() : parent);
|
||||
}
|
||||
}
|
||||
|
||||
for (int i = 1; i < visited.size(); ++i) {
|
||||
if (!visited[i]) {
|
||||
mju_error("Cycle detected: Node %s is not reachable from the world.",
|
||||
extraction.nodes[i]->body_path.GetString().c_str());
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
return world_root;
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
// 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.
|
||||
|
||||
#ifndef MUJOCO_PLUGIN_USD_DECODER_KINEMATIC_TREE_H_
|
||||
#define MUJOCO_PLUGIN_USD_DECODER_KINEMATIC_TREE_H_
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include <pxr/usd/sdf/path.h>
|
||||
#include <pxr/usd/usdPhysics/joint.h>
|
||||
|
||||
// A struct to represent a node in the kinematic tree.
|
||||
// Using a struct with a vector of children preserves the order of bodies,
|
||||
// which is important for things like keyframes and policy compatibility.
|
||||
struct Node {
|
||||
pxr::SdfPath body_path;
|
||||
pxr::SdfPath physics_scene;
|
||||
std::vector<pxr::SdfPath> actuators;
|
||||
std::vector<pxr::SdfPath> joints;
|
||||
std::vector<pxr::SdfPath> visual_gprims;
|
||||
std::vector<pxr::SdfPath> colliders;
|
||||
std::vector<pxr::SdfPath> sites;
|
||||
std::vector<pxr::SdfPath> tendons;
|
||||
std::vector<pxr::SdfPath> keyframes;
|
||||
std::vector<std::unique_ptr<Node>> children;
|
||||
};
|
||||
|
||||
// A kinematic edge represents a joint.
|
||||
using JointVec = std::vector<pxr::UsdPhysicsJoint>;
|
||||
|
||||
// Builds a single kinematic tree from a list of directed edges.
|
||||
// The DFS order of bodies in the tree is determined by the order of bodies in
|
||||
// `all_body_paths`.
|
||||
// All bodies, including static and floating-base bodies, are organized under a
|
||||
// single world root. An empty 'from' path in an edge represents the world body.
|
||||
// Returns the root of the kinematic tree, or `nullptr` for invalid structures.
|
||||
std::unique_ptr<Node> BuildKinematicTree(const pxr::UsdStageRefPtr stage);
|
||||
|
||||
#endif // MUJOCO_PLUGIN_USD_DECODER_KINEMATIC_TREE_H_
|
||||
@@ -0,0 +1,268 @@
|
||||
// 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 "material_parsing.h"
|
||||
|
||||
#include <cstdio>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
|
||||
#include <mujoco/mujoco.h>
|
||||
#include <pxr/base/gf/vec3f.h>
|
||||
#include <pxr/base/gf/vec4f.h>
|
||||
#include <pxr/base/tf/staticData.h>
|
||||
#include <pxr/base/tf/staticTokens.h>
|
||||
#include <pxr/base/tf/token.h>
|
||||
#include <pxr/usd/ar/asset.h>
|
||||
#include <pxr/usd/ar/resolvedPath.h>
|
||||
#include <pxr/usd/ar/resolver.h>
|
||||
#include <pxr/usd/usd/prim.h>
|
||||
#include <pxr/usd/usdShade/input.h>
|
||||
#include <pxr/usd/usdShade/material.h>
|
||||
#include <pxr/usd/usdShade/shader.h>
|
||||
#include <pxr/usd/usdShade/types.h>
|
||||
#include <pxr/usd/usdShade/udimUtils.h>
|
||||
|
||||
// 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,
|
||||
((auto_, "auto"))
|
||||
((diffuseColor, "diffuseColor"))
|
||||
((file, "file"))
|
||||
((metallic, "metallic"))
|
||||
((r, "r"))
|
||||
((raw, "raw"))
|
||||
((rgb, "rgb"))
|
||||
((roughness, "roughness"))
|
||||
((sourceColorSpace, "sourceColorSpace"))
|
||||
((srgb, "sRGB"))
|
||||
((UsdPreviewSurface, "UsdPreviewSurface"))
|
||||
((UsdUVTexture, "UsdUVTexture"))
|
||||
);
|
||||
// clang-format on
|
||||
|
||||
pxr::GfVec3f AsVec3f(pxr::GfVec3f val) { return val; }
|
||||
|
||||
pxr::GfVec3f AsVec3f(float val) { return pxr::GfVec3f(val, val, val); }
|
||||
|
||||
pxr::GfVec4f AsVec4f(pxr::GfVec4f val) { return val; }
|
||||
|
||||
pxr::GfVec4f AsVec4f(float val) { return pxr::GfVec4f(val, val, val, 1.0f); }
|
||||
|
||||
pxr::GfVec4f AsVec4f(pxr::GfVec3f vec3) {
|
||||
return pxr::GfVec4f(vec3[0], vec3[1], vec3[2], 1.0f);
|
||||
}
|
||||
|
||||
pxr::TfToken GetShaderId(const pxr::UsdShadeShader& shader) {
|
||||
pxr::TfToken shader_id;
|
||||
shader.GetShaderId(&shader_id);
|
||||
return shader_id;
|
||||
}
|
||||
|
||||
// Data read from a shader input.
|
||||
template <typename T>
|
||||
struct ResolvedShaderInput {
|
||||
// The value of the input if it is a constant.
|
||||
std::optional<T> value;
|
||||
|
||||
// The value of the input if it is a texture.
|
||||
std::optional<mjsTexture*> sampler;
|
||||
};
|
||||
|
||||
// Attempts to read an input from a USD ShadeInput as a specific type. Returns
|
||||
// nullopt if the input is not found or is not of the given type.
|
||||
template <typename T>
|
||||
std::optional<T> ReadInput(pxr::UsdShadeInput usd_input) {
|
||||
T val;
|
||||
if (!usd_input.Get(&val)) {
|
||||
return std::nullopt;
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
// Attempts to read a UsdUVTexture shader and return a sampler.
|
||||
template <typename T>
|
||||
ResolvedShaderInput<T> ReadUsdUVTexture(mjSpec* spec,
|
||||
const pxr::UsdShadeShader shader,
|
||||
unsigned nchannels) {
|
||||
ResolvedShaderInput<T> out;
|
||||
mjsTexture* texture = mjs_addTexture(spec);
|
||||
texture->type = mjtTexture::mjTEXTURE_2D;
|
||||
mjs_setName(texture->element, shader.GetPath().GetAsString().c_str());
|
||||
|
||||
pxr::TfToken source_color_space = kTokens->auto_;
|
||||
if (auto color_space_input = shader.GetInput(kTokens->sourceColorSpace)) {
|
||||
color_space_input.Get(&source_color_space);
|
||||
}
|
||||
|
||||
if (source_color_space == kTokens->raw) {
|
||||
texture->colorspace = mjtColorSpace::mjCOLORSPACE_LINEAR;
|
||||
} else if (source_color_space == kTokens->srgb) {
|
||||
texture->colorspace = mjtColorSpace::mjCOLORSPACE_SRGB;
|
||||
} else if (source_color_space == kTokens->auto_) {
|
||||
texture->colorspace = mjtColorSpace::mjCOLORSPACE_AUTO;
|
||||
}
|
||||
|
||||
pxr::ArResolver& resolver = pxr::ArGetResolver();
|
||||
pxr::SdfAssetPath resolved_texture_asset_path;
|
||||
if (auto file_input = shader.GetInput(kTokens->file)) {
|
||||
file_input.Get(&resolved_texture_asset_path);
|
||||
} else {
|
||||
mju_error("UsdUVTexture missing inputs:file.");
|
||||
return out;
|
||||
}
|
||||
|
||||
std::string resolved_texture_path =
|
||||
resolved_texture_asset_path.GetResolvedPath();
|
||||
|
||||
if (pxr::UsdShadeUdimUtils::IsUdimIdentifier(resolved_texture_path)) {
|
||||
mju_error("MuJoCo does not support UDIM textures: %s",
|
||||
resolved_texture_path.c_str());
|
||||
return out;
|
||||
}
|
||||
|
||||
auto extension = resolver.GetExtension(resolved_texture_path);
|
||||
|
||||
FILE* fp = fopen(resolved_texture_path.c_str(), "r");
|
||||
if (fp == nullptr) {
|
||||
mju_error(
|
||||
"USD decoder only supports assets that are available on the file "
|
||||
"system");
|
||||
return out;
|
||||
}
|
||||
texture->nchannel = nchannels;
|
||||
mjs_setString(texture->file, resolved_texture_path.c_str());
|
||||
out.sampler = texture;
|
||||
return out;
|
||||
}
|
||||
|
||||
// Given an input to a shader, attempts to bake the shader and it's inputs
|
||||
// into a singular value or texture sampler (ResolvedShaderInput).
|
||||
template <typename T>
|
||||
ResolvedShaderInput<T> ReadShaderInput(
|
||||
mjSpec* spec, const pxr::UsdShadeConnectableAPI source,
|
||||
const pxr::TfToken source_name,
|
||||
const pxr::UsdShadeAttributeType source_type) {
|
||||
ResolvedShaderInput<T> out;
|
||||
pxr::UsdShadeShader shader(source.GetPrim());
|
||||
pxr::TfToken shader_id = GetShaderId(shader);
|
||||
if (shader_id == kTokens->UsdUVTexture) {
|
||||
unsigned nchannels = -1;
|
||||
if (source_name == kTokens->rgb) {
|
||||
nchannels = 3;
|
||||
} else if (source_name == kTokens->r) {
|
||||
nchannels = 1;
|
||||
} else {
|
||||
mju_error("Unsupported texture channel: %s", source_name.GetText());
|
||||
return out;
|
||||
}
|
||||
|
||||
return ReadUsdUVTexture<T>(spec, shader, nchannels);
|
||||
} else {
|
||||
mju_warning("Unsupported shader type: %s", shader_id.GetText());
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
// Reads UsdShadeInput as a value of type T or a texture sampler.
|
||||
template <typename T>
|
||||
ResolvedShaderInput<T> ReadShaderInput(mjSpec* spec, pxr::UsdShadeInput input) {
|
||||
ResolvedShaderInput<T> out;
|
||||
if (!input.GetPrim().IsValid()) {
|
||||
return out;
|
||||
}
|
||||
|
||||
pxr::UsdShadeConnectableAPI source;
|
||||
pxr::TfToken source_name;
|
||||
pxr::UsdShadeAttributeType source_type;
|
||||
if (input.GetConnectedSource(&source, &source_name, &source_type)) {
|
||||
out = ReadShaderInput<T>(spec, source, source_name, source_type);
|
||||
} else {
|
||||
out.value = ReadInput<T>(input);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
template <typename T, typename V>
|
||||
void AssignShaderInput(mjSpec* spec, mjsMaterial* material, T* val,
|
||||
mjtTextureRole texrole,
|
||||
ResolvedShaderInput<V> resolved_input) {
|
||||
if (resolved_input.value.has_value()) {
|
||||
if constexpr (std::is_same_v<T, float>) {
|
||||
*val = resolved_input.value.value();
|
||||
} else if constexpr (std::is_same_v<T, float[3]>) {
|
||||
const pxr::GfVec3f resolved_value = AsVec3f(resolved_input.value.value());
|
||||
(*val)[0] = resolved_value[0];
|
||||
(*val)[1] = resolved_value[1];
|
||||
(*val)[2] = resolved_value[2];
|
||||
} else if constexpr (std::is_same_v<T, float[4]>) {
|
||||
const pxr::GfVec4f resolved_value = AsVec4f(resolved_input.value.value());
|
||||
(*val)[0] = resolved_value[0];
|
||||
(*val)[1] = resolved_value[1];
|
||||
(*val)[2] = resolved_value[2];
|
||||
(*val)[3] = resolved_value[3];
|
||||
}
|
||||
} else if (resolved_input.sampler.has_value()) {
|
||||
mjsTexture* texture = resolved_input.sampler.value();
|
||||
auto name = mjs_getName(texture->element);
|
||||
mjs_setInStringVec(material->textures, texrole, name->c_str());
|
||||
} else {
|
||||
mju_warning("No value or texture for shader input.");
|
||||
}
|
||||
}
|
||||
|
||||
void ParsePreviewSurface(mjSpec* spec, mjsMaterial* material,
|
||||
const pxr::UsdShadeShader& shader) {
|
||||
auto diffuse = ReadShaderInput<pxr::GfVec3f>(
|
||||
spec, shader.GetInput(kTokens->diffuseColor));
|
||||
AssignShaderInput(spec, material, &material->rgba, mjTEXROLE_RGB, diffuse);
|
||||
|
||||
auto metallic =
|
||||
ReadShaderInput<float>(spec, shader.GetInput(kTokens->metallic));
|
||||
AssignShaderInput(spec, material, &material->metallic, mjTEXROLE_METALLIC,
|
||||
metallic);
|
||||
|
||||
auto roughness =
|
||||
ReadShaderInput<float>(spec, shader.GetInput(kTokens->roughness));
|
||||
AssignShaderInput(spec, material, &material->roughness, mjTEXROLE_ROUGHNESS,
|
||||
roughness);
|
||||
}
|
||||
|
||||
mjsMaterial* ParseMaterial(mjSpec* spec,
|
||||
const pxr::UsdShadeMaterial& material) {
|
||||
pxr::UsdShadeShader surface_shader = material.ComputeSurfaceSource();
|
||||
if (!surface_shader.GetPrim().IsValid()) {
|
||||
mju_warning("Material %s has no surface output.",
|
||||
material.GetPath().GetAsString().c_str());
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
pxr::TfToken surface_shader_id = GetShaderId(surface_shader);
|
||||
if (surface_shader_id != kTokens->UsdPreviewSurface) {
|
||||
mju_warning("Mujoco only supports UsdPreviewSurface as surface output.");
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
mjsMaterial* mj_mat = mjs_addMaterial(spec, nullptr);
|
||||
mjs_setName(mj_mat->element, material.GetPath().GetAsString().c_str());
|
||||
ParsePreviewSurface(spec, mj_mat, surface_shader);
|
||||
|
||||
return mj_mat;
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
// 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.
|
||||
|
||||
#ifndef MUJOCO_PLUGIN_USD_DECODER_MATERIAL_PARSING_H_
|
||||
#define MUJOCO_PLUGIN_USD_DECODER_MATERIAL_PARSING_H_
|
||||
|
||||
#include <pxr/usd/usdShade/material.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
|
||||
mjsMaterial* ParseMaterial(mjSpec* spec, const pxr::UsdShadeMaterial &material);
|
||||
|
||||
#endif // MUJOCO_PLUGIN_USD_DECODER_MATERIAL_PARSING_H_
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user