diff --git a/python/mujoco/testdata/checkerboard.png b/python/mujoco/testdata/checkerboard.png new file mode 100644 index 00000000..87fb3579 Binary files /dev/null and b/python/mujoco/testdata/checkerboard.png differ diff --git a/python/mujoco/testdata/simple_usd_preview_surface.usda b/python/mujoco/testdata/simple_usd_preview_surface.usda new file mode 100644 index 00000000..ede01fec --- /dev/null +++ b/python/mujoco/testdata/simple_usd_preview_surface.usda @@ -0,0 +1,81 @@ +#usda 1.0 +( + endTimeCode = 1 + framesPerSecond = 24 + metersPerUnit = 1 + startTimeCode = 1 + timeCodesPerSecond = 24 + upAxis = "Y" +) + +def Xform "World" +{ + def Cube "Cube" ( + prepend apiSchemas = ["MaterialBindingAPI"] + ) + { + float3[] extent = [(-1, -1, -1), (1, 1, 1)] + rel material:binding = + double size = 2 + matrix4d xformOp:transform = ( (1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1) ) + uniform token[] xformOpOrder = ["xformOp:transform"] + } + + def Cube "Cube2" ( + prepend apiSchemas = ["MaterialBindingAPI"] + ) + { + float3[] extent = [(-1, -1, -1), (1, 1, 1)] + rel material:binding = + double size = 2 + matrix4d xformOp:transform = ( (1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1) ) + matrix4d xformOp:transform:transform1 = ( (1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, -2.9009591341018677, 0, 1) ) + uniform token[] xformOpOrder = ["xformOp:transform", "xformOp:transform:transform1"] + } +} + +def Scope "materials" +{ + def Material "simple_usd_preview_surface" ( + prepend inherits = + ) + { + token outputs:surface.connect = + + def Shader "usdpreviewsurface" + { + uniform token info:id = "UsdPreviewSurface" + color3f inputs:diffuseColor = (0.1, 0.2, 0.3) + color3f inputs:emissiveColor = (0.4, 0.5, 0.6) + float inputs:metallic = 0.123 + float inputs:roughness = 0.456 + token outputs:surface + } + } + + def Material "simple_usd_texture" ( + prepend inherits = + ) + { + token outputs:surface.connect = + + def Shader "usdpreviewsurface" + { + uniform token info:id = "UsdPreviewSurface" + color3f inputs:diffuseColor = (0.1, 0.2, 0.3) + color3f inputs:diffuseColor.connect = + color3f inputs:emissiveColor = (0.4, 0.5, 0.6) + float inputs:metallic = 0.123 + float inputs:roughness = 0.456 + token outputs:surface + } + + def Shader "usduvtexture1" + { + uniform token info:id = "UsdUVTexture" + asset inputs:file = @./checkerboard.png@ + vector3f outputs:rgb + } + } +} + diff --git a/src/experimental/usd/CMakeLists.txt b/src/experimental/usd/CMakeLists.txt index 536aa12a..59f8e845 100644 --- a/src/experimental/usd/CMakeLists.txt +++ b/src/experimental/usd/CMakeLists.txt @@ -196,11 +196,20 @@ target_sources(mujoco PRIVATE kinematic_tree.cc kinematic_tree.h layer_sink.cc + material_parsing.cc + material_parsing.h usd_to_mjspec.cc utils.cc writer.cc ) +# pxr/usd/tf/hashset.h might include deprecated headers so disable this check +# when compiling with USD enabled. +target_compile_options( + mujoco + PRIVATE -Wno-deprecated +) + if (USD_DIR) find_package(pxr REQUIRED) diff --git a/src/experimental/usd/material_parsing.cc b/src/experimental/usd/material_parsing.cc new file mode 100644 index 00000000..9b26dcdd --- /dev/null +++ b/src/experimental/usd/material_parsing.cc @@ -0,0 +1,315 @@ +// 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 "experimental/usd/material_parsing.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "lodepng.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mujoco { +namespace usd { + +// Using to satisfy TF_DEFINE_PRIVATE_TOKENS macro below and avoid operating in +// PXR_NS. +using pxr::TfToken; +template +using TfStaticData = pxr::TfStaticData; + +// 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 +struct ResolvedShaderInput { + // The value of the input if it is a constant. + std::optional value; + + // The value of the input if it is a texture. + std::optional 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 +std::optional 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 +ResolvedShaderInput ReadUsdUVTexture(mjSpec* spec, + const pxr::UsdShadeShader shader, + unsigned nchannels) { + ResolvedShaderInput 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); + if (extension != "png") { + mju_error("MuJoCo USD Parsing only supports PNG textures: %s", + resolved_texture_path.c_str()); + return out; + } + + FILE* fp = fopen(resolved_texture_path.c_str(), "r"); + if (fp) { + mjs_setString(texture->content_type, "image/png"); + mjs_setString(texture->file, resolved_texture_path.c_str()); + out.sampler = texture; + return out; + } + + std::shared_ptr texture_asset = + resolver.OpenAsset(pxr::ArResolvedPath(resolved_texture_path)); + + size_t texture_size = texture_asset->GetSize(); + + uint32_t width, height; + std::vector image; + lodepng::State state; + if (nchannels == 3) { + state.info_raw.colortype = LCT_RGB; + } else if (nchannels == 1) { + state.info_raw.colortype = LCT_GREY; + } else { + mju_error("MuJoCo USD Parsing only supports 1 or 3 channel textures."); + return out; + } + + unsigned error = lodepng::decode( + image, width, height, state, + reinterpret_cast(texture_asset->GetBuffer().get()), + texture_size); + + // check for errors + if (error) { + mju_error("LodePNG error %u: %s", error, lodepng_error_text(error)); + return out; + } + + texture->width = width; + texture->height = height; + texture->nchannel = nchannels; + mjs_setBuffer(texture->data, image.data(), image.size()); + + 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 +ResolvedShaderInput ReadShaderInput( + mjSpec* spec, const pxr::UsdShadeConnectableAPI source, + const pxr::TfToken source_name, + const pxr::UsdShadeAttributeType source_type) { + ResolvedShaderInput 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(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 +ResolvedShaderInput ReadShaderInput(mjSpec* spec, pxr::UsdShadeInput input) { + ResolvedShaderInput 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(spec, source, source_name, source_type); + } else { + out.value = ReadInput(input); + } + return out; +} + +template +void AssignShaderInput(mjSpec* spec, mjsMaterial* material, T* val, + mjtTextureRole texrole, + ResolvedShaderInput resolved_input) { + if (resolved_input.value.has_value()) { + if constexpr (std::is_same_v) { + *val = resolved_input.value.value(); + } else if constexpr (std::is_same_v) { + 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) { + 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( + spec, shader.GetInput(kTokens->diffuseColor)); + AssignShaderInput(spec, material, &material->rgba, mjTEXROLE_RGB, diffuse); + + auto metallic = + ReadShaderInput(spec, shader.GetInput(kTokens->metallic)); + AssignShaderInput(spec, material, &material->metallic, mjTEXROLE_METALLIC, + metallic); + + auto roughness = + ReadShaderInput(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; +} +} // namespace usd +} // namespace mujoco diff --git a/src/experimental/usd/material_parsing.h b/src/experimental/usd/material_parsing.h new file mode 100644 index 00000000..48103676 --- /dev/null +++ b/src/experimental/usd/material_parsing.h @@ -0,0 +1,27 @@ +// 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_SRC_EXPERIMENTAL_USD_MATERIAL_PARSING_H_ +#define MUJOCO_SRC_EXPERIMENTAL_USD_MATERIAL_PARSING_H_ + +#include +#include + +namespace mujoco { +namespace usd { +mjsMaterial* ParseMaterial(mjSpec* spec, const pxr::UsdShadeMaterial &material); +} // namespace usd +} // namespace mujoco + +#endif // MUJOCO_SRC_EXPERIMENTAL_USD_MATERIAL_PARSING_H_ diff --git a/src/experimental/usd/usd_to_mjspec.cc b/src/experimental/usd/usd_to_mjspec.cc index e6b56834..c2e667ea 100644 --- a/src/experimental/usd/usd_to_mjspec.cc +++ b/src/experimental/usd/usd_to_mjspec.cc @@ -32,6 +32,7 @@ #include #include #include "experimental/usd/kinematic_tree.h" +#include "experimental/usd/material_parsing.h" #include #include #include @@ -77,6 +78,7 @@ struct UsdCaches { pxr::UsdGeomXformCache xform_cache; pxr::UsdShadeMaterialBindingAPI::BindingsCache bindings_cache; pxr::UsdShadeMaterialBindingAPI::CollectionQueryCache collection_query_cache; + std::map parsed_materials; }; void SetDoubleArrFromGfVec3d(double* to, const pxr::GfVec3d& from) { @@ -1333,6 +1335,21 @@ void ParseUsdGeomGprim(mjSpec* spec, const pxr::UsdPrim& gprim, if (!MaybeParseGeomPrimitive(gprim, geom, caches.xform_cache)) { ParseUsdMesh(spec, gprim, geom, caches.xform_cache); } + pxr::UsdShadeMaterial bound_material = + pxr::UsdShadeMaterialBindingAPI(gprim).ComputeBoundMaterial( + &caches.bindings_cache, &caches.collection_query_cache); + if (bound_material) { + pxr::SdfPath material_path = bound_material.GetPrim().GetPath(); + mjsMaterial* material = nullptr; + if (caches.parsed_materials.find(material_path) != + caches.parsed_materials.end()) { + material = caches.parsed_materials[material_path]; + } else { + material = ParseMaterial(spec, bound_material); + caches.parsed_materials[material_path] = material; + } + mjs_setString(geom->material, mjs_getName(material->element)->c_str()); + } } void ParseUsdPhysicsCollider(mjSpec* spec,