Expose ability to go from USD stage to MjSpec.

PiperOrigin-RevId: 776184388 Change-Id: I0325b347c9896183d148a196897e1e5d0c4bbde2
2025-06-26 10:08:45 -07:00
parent 672746f509
commit afc35451a0
9 changed files with 1995 additions and 163 deletions
@@ -42,14 +42,26 @@ option(MUJOCO_BUILD_EXAMPLES "Build samples for MuJoCo" ON)
 option(MUJOCO_BUILD_SIMULATE "Build simulate library for MuJoCo" ON)
 option(MUJOCO_BUILD_TESTS "Build tests for MuJoCo" ON)
 option(MUJOCO_TEST_PYTHON_UTIL "Build and test utility libraries for Python bindings" ON)
-option(MUJOCO_BUILD_USD_PLUGINS "Build OpenUSD plugins" OFF)
+option(MUJOCO_WITH_USD "Build with OpenUSD" OFF)

-# USD libs to compile against.
-set(MUJOCO_USD_ALLOWED_TARGET_VALUES "USD" "Houdini")
-set(MUJOCO_USD_TARGET "USD" CACHE STRING "Select the USD target for the project.")
-set_property(CACHE MUJOCO_USD_TARGET
-  PROPERTY STRINGS ${MUJOCO_USD_ALLOWED_TARGET_VALUES}
-)
+# Option to provide a path to an existing USD build directory or to Houdini HFS directory.
+set(USD_DIR "" CACHE PATH "Path to an existing USD build directory.")
+set(HOUDINI_HFS_DIR "" CACHE PATH "Path to Houdini HFS directory to build USD plugins against.")
+if(USD_DIR)
+  if(EXISTS "${USD_DIR}")
+    # If the path is provided, set MUJOCO_WITH_USD to ON and add the path to CMAKE_PREFIX_PATH.
+    set(MUJOCO_WITH_USD ON CACHE BOOL "Build with OpenUSD" FORCE)
+    list(PREPEND CMAKE_PREFIX_PATH "${USD_DIR}")
+    message(STATUS "Using custom USD build directory: ${USD_DIR}")
+  else()
+    message(WARNING "Invalid path provided for USD_DIR: ${USD_DIR}")
+  endif()
+elseif(HOUDINI_HFS_DIR)
+  if(EXISTS "${HOUDINI_HFS_DIR}")
+    set(MUJOCO_WITH_USD ON CACHE BOOL "Build with OpenUSD" FORCE)
+    message(STATUS "Building Mujoco USD against Houdini HFS: ${HOUDINI_HFS_DIR}")
+  endif()
+endif()

 if(APPLE AND (MUJOCO_BUILD_EXAMPLES OR MUJOCO_BUILD_SIMULATE))
  enable_language(OBJC)
@@ -189,8 +201,8 @@ if(MUJOCO_BUILD_EXAMPLES)
  add_subdirectory(sample)
 endif()

-if(MUJOCO_BUILD_USD_PLUGINS)
-  add_subdirectory(src/experimental/usd/plugins)
+if(MUJOCO_WITH_USD)
+  add_subdirectory(src/experimental/usd)
 endif()

 if(BUILD_TESTING AND MUJOCO_BUILD_TESTS)
@@ -0,0 +1,28 @@
+// 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_EXPERIMENTAL_SRC_USD_USD_TO_MJSPEC_H_
+#define MUJOCO_EXPERIMENTAL_SRC_USD_USD_TO_MJSPEC_H_
+
+#include <mujoco/mujoco.h>
+#include <pxr/usd/usd/common.h>
+
+// Given a USD stage, this function will do a best effort conversion to
+// mjSpec.
+//
+// Particular care is taken for physics data to be lossless but visual
+// data such as materials may be lossy.
+MJAPI mjSpec* mj_parseUSDStage(pxr::UsdStageRefPtr stage);
+
+#endif  // MUJOCO_EXPERIMENTAL_SRC_USD_USD_TO_MJSPEC_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
+#
+#   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.
+
+# --- Global Configuration ---
+# Plugin target names (used for library and plugInfo.json)
+set(MJCF_PLUGIN_TARGET_NAME usdMjcf)
+set(MJC_PHYSICS_PLUGIN_TARGET_NAME mjcPhysics)
+set(MJ_USD_TARGET_NAME mj_usd)
+
+# Installation directory for USD plugins
+set(MJ_USD_INSTALL_DIR_LIB ${CMAKE_INSTALL_LIBDIR}/mujocoUsd)
+
+# Common OpenUSD libraries
+set(OPENUSD_CORE_LIBS
+  usd ar kind tf gf vt usdShade usdLux usdGeom usdImaging usdPhysics
+)
+
+# We need to set the visibility to default until core type symbol visibility
+# is resolved in OpenUSD https://github.com/PixarAnimationStudios/OpenUSD/issues/1475
+set(DEFAULT_CXX_VISIBILITY_PROPS
+  CXX_VISIBILITY_PRESET default
+)
+
+# --- Helper Functions ---
+function(configure_and_install_usd_plugin_info plugin_name source_subpath install_base_dir)
+  # Determine shared library prefix
+  if(CMAKE_SHARED_LIBRARY_PREFIX)
+    set(LIB_PREFIX ${CMAKE_SHARED_LIBRARY_PREFIX})
+  else()
+    set(LIB_PREFIX "")
+  endif()
+
+  set(PLUG_INFO_LIBRARY_PATH "../../${LIB_PREFIX}${plugin_name}${CMAKE_SHARED_LIBRARY_SUFFIX}")
+  set(OUTPUT_DIR "${CMAKE_BINARY_DIR}/${install_base_dir}/resources/${plugin_name}")
+  set(OUTPUT_FILE "${OUTPUT_DIR}/plugInfo.json")
+
+  # Create the output directory if it doesn't exist
+  file(MAKE_DIRECTORY "${OUTPUT_DIR}")
+
+  # Configure the plugInfo.json file to replace @PLUG_INFO_LIBRARY_PATH@ instances.
+  configure_file(
+    "${CMAKE_CURRENT_SOURCE_DIR}/${source_subpath}/plugInfo.json"
+    "${OUTPUT_FILE}"
+    @ONLY
+  )
+
+  install(FILES "${OUTPUT_FILE}"
+    DESTINATION "${install_base_dir}/resources/${plugin_name}"
+  )
+endfunction()
+
+# Function to apply Houdini-specific link directories and include paths
+function(configure_houdini_target target_name link_scope include_scope)
+  set(HOUDINI_LIBS_DIR "${HOUDINI_HFS_DIR}/../Libraries")
+  get_filename_component(HOUDINI_LIBS_DIR "${HOUDINI_LIBS_DIR}" ABSOLUTE) # Normalize the path
+
+  # Assume everyone using Houdini on 3.11 for now.
+  set(USD_MJCF_PYTHON_LIB python3.11)
+  set(USD_MJCF_PYTHON_LIB_NUMBER python311) # For boost library names
+  set(PYTHON_LIB "${HOUDINI_HFS_DIR}/Frameworks/Python.framework/Versions/3.11/Python")
+  set(PXR_LIB_PREFIX "pxr_")
+
+  # Houdini-specific PXR and other libraries
+  set(HOUDINI_SPECIFIC_LIBS
+    ${PXR_LIB_PREFIX}usd
+    ${PXR_LIB_PREFIX}ar
+    ${PXR_LIB_PREFIX}kind
+    ${PXR_LIB_PREFIX}tf
+    ${PXR_LIB_PREFIX}gf
+    ${PXR_LIB_PREFIX}vt
+    ${PXR_LIB_PREFIX}sdf
+    ${PXR_LIB_PREFIX}usdShade
+    ${PXR_LIB_PREFIX}usdLux
+    ${PXR_LIB_PREFIX}usdGeom
+    ${PXR_LIB_PREFIX}usdImaging
+    ${PXR_LIB_PREFIX}usdPhysics
+    tbb
+    hboost_${USD_MJCF_PYTHON_LIB_NUMBER}
+    ${PYTHON_LIB}
+  )
+
+  target_link_directories(${target_name} ${link_scope} ${HOUDINI_LIBS_DIR})
+  target_include_directories(${target_name} ${include_scope}
+    "${HOUDINI_HFS_DIR}/toolkit/include"
+    "${HOUDINI_HFS_DIR}/toolkit/include/python3.11"
+  )
+  target_link_libraries(${target_name} ${link_scope}
+    ${HOUDINI_SPECIFIC_LIBS}
+    mujoco
+    tinyxml2
+  )
+endfunction()
+
+## usdMjcf
+set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE)
+
+# The list of RPATHs:
+# 1. $ORIGIN: For libmjcPhysics.so (since it's in the same install dir)
+# 2. Path to OpenUSD libraries: Typically ${USD_DIR}/lib
+# 3. Path to MuJoCo library: Assuming it's in ${MUJOCO_INSTALL_DIR}/lib
+if(APPLE)
+    set(RPATH_LIST "@loader_path")
+    # Add OpenUSD lib to RPATH for macOS
+    if(DEFINED USD_DIR AND EXISTS "${USD_DIR}/lib")
+        list(APPEND RPATH_LIST "${USD_DIR}/lib")
+    endif()
+    # Add MuJoCo lib to RPATH for macOS
+    if(DEFINED MUJOCO_INSTALL_DIR AND EXISTS "${MUJOCO_INSTALL_DIR}/lib")
+        list(APPEND RPATH_LIST "${MUJOCO_INSTALL_DIR}/lib")
+    endif()
+
+    string(JOIN ":" CMAKE_INSTALL_RPATH "${RPATH_LIST}") # macOS uses colon for RPATH separator
+    message(STATUS "Final CMAKE_INSTALL_RPATH: ${CMAKE_INSTALL_RPATH}")
+elseif(UNIX)
+
+    # For adjacent libs.
+    set(RPATH_LIST "\$ORIGIN")
+
+    # Add OpenUSD library directory
+    if(DEFINED USD_DIR AND EXISTS "${USD_DIR}/lib")
+        list(APPEND RPATH_LIST "${USD_DIR}/lib")
+        message(STATUS "Adding OpenUSD lib to RPATH: ${USD_DIR}/lib")
+    else()
+        message(WARNING "USD_DIR is not defined or ${USD_DIR}/lib does not exist. OpenUSD libraries might not be found via RPATH.")
+    endif()
+
+    # Add MuJoCo library directory
+    if(DEFINED MUJOCO_INSTALL_DIR AND EXISTS "${MUJOCO_INSTALL_DIR}/lib")
+        list(APPEND RPATH_LIST "${MUJOCO_INSTALL_DIR}/lib")
+        message(STATUS "Adding MuJoCo lib to RPATH: ${MUJOCO_INSTALL_DIR}/lib")
+    elseif(DEFINED MUJOCO_INSTALL_DIR AND EXISTS "${MUJOCO_INSTALL_DIR}/bin")
+        list(APPEND RPATH_LIST "${MUJOCO_INSTALL_DIR}/bin")
+        message(STATUS "Adding MuJoCo bin to RPATH: ${MUJOCO_INSTALL_DIR}/bin")
+    endif()
+
+    # Set the final RPATH for installation
+    string(JOIN ";" CMAKE_INSTALL_RPATH "${RPATH_LIST}")
+    message(STATUS "Final CMAKE_INSTALL_RPATH: ${CMAKE_INSTALL_RPATH}")
+endif()
+
+set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
+
+add_library(${MJCF_PLUGIN_TARGET_NAME} SHARED)
+target_sources(${MJCF_PLUGIN_TARGET_NAME} PRIVATE
+  plugins/mjcf/mjcf_file_format.cc
+  plugins/mjcf/mjcf_file_format.h
+  plugins/mjcf/mujoco_to_usd.cc
+  plugins/mjcf/mujoco_to_usd.h
+  plugins/mjcf/utils.cc
+  plugins/mjcf/utils.h
+)
+set_target_properties(${MJCF_PLUGIN_TARGET_NAME} PROPERTIES
+  OUTPUT_NAME ${MJCF_PLUGIN_TARGET_NAME}
+  ${DEFAULT_CXX_VISIBILITY_PROPS}
+)
+target_include_directories(${MJCF_PLUGIN_TARGET_NAME} PRIVATE
+  "${CMAKE_CURRENT_SOURCE_DIR}/plugins"
+)
+
+## mjcPhysics
+
+add_library(${MJC_PHYSICS_PLUGIN_TARGET_NAME} SHARED)
+target_sources(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PRIVATE
+  mjcPhysics/actuatorAPI.cpp
+  mjcPhysics/collisionAPI.cpp
+  mjcPhysics/meshCollisionAPI.cpp
+  mjcPhysics/sceneAPI.cpp
+  mjcPhysics/siteAPI.cpp
+  mjcPhysics/tokens.cpp
+)
+set_target_properties(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PROPERTIES
+  OUTPUT_NAME ${MJC_PHYSICS_PLUGIN_TARGET_NAME}
+  ${DEFAULT_CXX_VISIBILITY_PROPS}
+)
+target_include_directories(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PUBLIC
+  "${CMAKE_CURRENT_SOURCE_DIR}/mjcPhysics"
+)
+
+## mj_usd
+
+add_library(${MJ_USD_TARGET_NAME} SHARED)
+target_sources(${MJ_USD_TARGET_NAME} PRIVATE
+  writer.cc
+  layer_sink.cc
+  usd_to_mjspec.cc
+  utils.cc
+)
+
+
+if (USD_DIR)
+  find_package(pxr REQUIRED)
+
+  target_link_libraries(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PRIVATE
+    ${OPENUSD_CORE_LIBS}
+    mujoco
+    tinyxml2
+  )
+
+  target_link_libraries(${MJCF_PLUGIN_TARGET_NAME} PRIVATE
+    ${OPENUSD_CORE_LIBS}
+    mujoco
+    tinyxml2
+    ${MJC_PHYSICS_PLUGIN_TARGET_NAME}
+  )
+
+  target_link_libraries(${MJ_USD_TARGET_NAME} PUBLIC
+    ${OPENUSD_CORE_LIBS}
+    mujoco
+    tinyxml2
+    ${MJC_PHYSICS_PLUGIN_TARGET_NAME}
+  )
+elseif (HOUDINI_HFS_DIR)
+  message(STATUS "Configuring for Houdini target...")
+
+  # In Houdini, the Houdini package we would typically use via find_package
+  # does not have all the USD dependencies that we need (namely UsdPhysics)
+  # so we need to manually link all the required libraries.
+  configure_houdini_target(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PRIVATE PRIVATE)
+  configure_houdini_target(${MJCF_PLUGIN_TARGET_NAME} PRIVATE PRIVATE)
+  configure_houdini_target(${MJ_USD_TARGET_NAME} PUBLIC PUBLIC)
+
+  # Re-add inter-plugin dependencies after Houdini-specific linking
+  target_link_libraries(${MJCF_PLUGIN_TARGET_NAME} PRIVATE ${MJC_PHYSICS_PLUGIN_TARGET_NAME})
+  target_link_libraries(${MJ_USD_TARGET_NAME} PUBLIC ${MJC_PHYSICS_PLUGIN_TARGET_NAME})
+
+endif()
+
+## Installation
+
+message(STATUS "Mujoco USD plugins will be installed to: ${CMAKE_INSTALL_PREFIX}/${MJ_USD_INSTALL_DIR_LIB}")
+message(STATUS "Make sure PXR_PLUGINPATH_NAME includes: ${CMAKE_INSTALL_PREFIX}/${MJ_USD_INSTALL_DIR_LIB}/*/resources")
+
+# Generate and install plugInfo.json for each plugin
+configure_and_install_usd_plugin_info(
+  ${MJCF_PLUGIN_TARGET_NAME}
+  plugins/mjcf
+  ${MJ_USD_INSTALL_DIR_LIB}
+)
+configure_and_install_usd_plugin_info(
+  ${MJC_PHYSICS_PLUGIN_TARGET_NAME}
+  mjcPhysics
+  ${MJ_USD_INSTALL_DIR_LIB}
+)
+
+# Install specific resources that are not plugInfo.json (like generatedSchema.usda)
+install(FILES
+  mjcPhysics/generatedSchema.usda
+  DESTINATION ${MJ_USD_INSTALL_DIR_LIB}/resources/mjcPhysics
+)
+
+# Install shared libraries
+install(TARGETS
+  ${MJCF_PLUGIN_TARGET_NAME}
+  ${MJC_PHYSICS_PLUGIN_TARGET_NAME}
+  ${MJ_USD_TARGET_NAME}
+  LIBRARY DESTINATION ${MJ_USD_INSTALL_DIR_LIB}
+)
@@ -0,0 +1,132 @@
+// 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 <mujoco/experimental/usd/layer_sink.h>
+
+#include <vector>
+
+#include <mujoco/mujoco.h>
+#include <pxr/base/gf/declare.h>
+#include <pxr/base/gf/matrix4d.h>
+#include <pxr/base/gf/quaternion.h>
+#include <pxr/base/gf/rotation.h>
+#include <pxr/base/gf/transform.h>
+#include <pxr/base/tf/staticTokens.h>
+#include <pxr/base/tf/token.h>
+#include <pxr/base/vt/types.h>
+#include <pxr/base/vt/value.h>
+#include <pxr/pxr.h>
+#include <pxr/usd/sdf/attributeSpec.h>
+#include <pxr/usd/sdf/changeBlock.h>
+#include <pxr/usd/sdf/childrenPolicies.h>
+#include <pxr/usd/sdf/primSpec.h>
+#include <pxr/usd/sdf/types.h>
+#include <pxr/usd/sdf/valueTypeName.h>
+#include <pxr/usd/usd/common.h>
+#include <pxr/usd/usd/stage.h>
+#include <pxr/usd/usd/timeCode.h>
+#include <pxr/usd/usdGeom/tokens.h>
+#include <pxr/usd/usdGeom/xformCache.h>
+
+PXR_NAMESPACE_OPEN_SCOPE
+TF_DEFINE_PRIVATE_TOKENS(_layer_sink_tokens,
+                         ((xformOpTransform,
+                           "xformOp:transform:mujoco"))(Xform));
+PXR_NAMESPACE_CLOSE_SCOPE
+
+namespace {
+pxr::SdfAttributeSpecHandle GetOrCreateAttribute(
+    pxr::SdfLayerHandle layer, pxr::SdfPrimSpecHandle prim_spec,
+    const pxr::SdfPath& path, const pxr::TfToken& token,
+    const pxr::SdfValueTypeName& type, pxr::SdfVariability variability) {
+  // Find or create transform attribute.
+  pxr::SdfAttributeSpecHandle attr =
+      layer->GetAttributeAtPath(path.AppendProperty(token));
+
+  if (!attr) {
+    attr = pxr::SdfAttributeSpec::New(prim_spec, token, type, variability);
+  }
+  return attr;
+}
+}  // namespace
+
+namespace mujoco {
+namespace usd {
+LayerSink::LayerSink(pxr::UsdStageRefPtr stage) : stage_(stage) {}
+
+void LayerSink::Update(const mjData* const data,
+                       std::vector<pxr::SdfPath> body_paths) {
+  pxr::SdfLayerHandle layer = stage_->GetEditTarget().GetLayer();
+
+  pxr::UsdTimeCode timecode = data->time * layer->GetTimeCodesPerSecond();
+  pxr::UsdGeomXformCache xform_cache(timecode);
+
+  pxr::SdfChangeBlock change_block;
+  // TODO(shaves) this can be parallelised.
+  for (int i = 0; i < body_paths.size(); i++) {
+    const pxr::SdfPath& path = body_paths[i];
+
+    // Author transform timesample
+    // Convert pose data.
+    mjtNum* xpos = &data->xpos[i * 3];
+    mjtNum* xquat = &data->xquat[i * 4];
+    pxr::GfVec3d translation(xpos[0], xpos[1], xpos[2]);
+    pxr::GfVec3d imag_part(xquat[1], xquat[2], xquat[3]);
+    // Some bodies such as the world body don't have mapped USD prims to
+    // transform.
+    if (path.IsEmpty()) {
+      continue;
+    }
+
+    pxr::SdfPrimSpecHandle prim_spec = pxr::SdfCreatePrimInLayer(layer, path);
+
+    // Find or create transform attribute.
+    pxr::SdfAttributeSpecHandle xform_attr = GetOrCreateAttribute(
+        layer, prim_spec, path, pxr::_layer_sink_tokens->xformOpTransform,
+        pxr::SdfValueTypeNames->Matrix4d, pxr::SdfVariabilityVarying);
+
+    pxr::GfQuaternion quat(xquat[0], imag_part);
+    pxr::GfRotation rotation(quat);
+
+    // Identities
+    pxr::GfVec3d scale(1.0);
+    pxr::GfVec3d pivot_position(0);
+    pxr::GfRotation pivot_rotation(pxr::GfQuaternion::GetIdentity());
+
+    // Set time sample
+    pxr::GfTransform xform(translation, rotation, scale, pivot_position,
+                           pivot_rotation);
+    pxr::GfMatrix4d parent_xform =
+        xform_cache.GetParentToWorldTransform(stage_->GetPrimAtPath(path));
+    pxr::GfMatrix4d relative_xform =
+        xform.GetMatrix() * parent_xform.GetInverse();
+    layer->SetTimeSample(xform_attr->GetPath(), timecode.GetValue(),
+                         relative_xform);
+
+    // Create xformop_order attribute if missing.
+    pxr::SdfAttributeSpecHandle op_order_attr = GetOrCreateAttribute(
+        layer, prim_spec, path, pxr::UsdGeomTokens->xformOpOrder,
+        pxr::SdfValueTypeNames->TokenArray, pxr::SdfVariabilityUniform);
+
+    // Override all local transforms so that this is the absolute position.
+    // We don't reset the xform stack here, so parent transforms will still
+    // affect the world pose.
+    const pxr::VtTokenArray xformop_order(
+        {pxr::_layer_sink_tokens->xformOpTransform});
+
+    op_order_attr->SetDefaultValue(pxr::VtValue(xformop_order));
+  }
+}
+}  // namespace usd
+}  // namespace mujoco
@@ -1,137 +0,0 @@
-# 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
-#
-#     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.
-
-
-# Plugin target name (used for library and plugInfo.json)
-set(MJCF_PLUGIN_TARGET_NAME usdMjcf)
-
-add_library(${MJCF_PLUGIN_TARGET_NAME} SHARED
-  mjcf/mjcf_file_format.cc
-  mjcf/mjcf_file_format.h
-  mjcf/mujoco_to_usd.cc
-  mjcf/mujoco_to_usd.h
-  mjcf/utils.cc
-  mjcf/utils.h
-)
-
-# We need to set the visibility to default until core type symbol visibility
-# is resolved in OpenUSD https://github.com/PixarAnimationStudios/OpenUSD/issues/1475
-# Otherwise we will run into issues during composition on MacOS due to std::type_info
-# comparisons failing for pxr::TfTokenVector and the like that we place in SdfAbstractData.
-set_target_properties(${MJCF_PLUGIN_TARGET_NAME} PROPERTIES
-    OUTPUT_NAME ${MJCF_PLUGIN_TARGET_NAME}
-    CXX_VISIBILITY_PRESET default
-)
-
-target_include_directories(${MJCF_PLUGIN_TARGET_NAME} PRIVATE
-  "${CMAKE_CURRENT_SOURCE_DIR}"
-)
-
-if (MUJOCO_USD_TARGET STREQUAL "USD")
-  find_package(pxr REQUIRED)
-
-  # --- Link Dependencies ---
-  # Link against the necessary OpenUSD components
-  target_link_libraries(${MJCF_PLUGIN_TARGET_NAME} PRIVATE
-    usd
-    ar
-    kind
-    tf
-    gf
-    vt
-    usdShade
-    usdLux
-    usdGeom
-    usdImaging
-    usdPhysics
-    mujoco
-    tinyxml2
-  )
-elseif (MUJOCO_USD_TARGET STREQUAL "Houdini")
-
-  if (NOT DEFINED ENV{HFS})
-    message(FATAL_ERROR "Environment variable 'HFS' is not defined: $ENV{HFS}. Please run houdini_setup.")
-  endif()
-
-  # In Houdini, the Houdini package we would typically use via find_package
-  # does not have all the USD dependencies that we need (namely UsdPhysics)
-  # so we need to manually link all the required libraries.
-
-  set(HFS_ENV "$ENV{HFS}")
-  set(HOUDINI_LIBS "${HFS_ENV}/../Libraries")
-  get_filename_component(HOUDINI_LIBS "${HOUDINI_LIBS}" ABSOLUTE) # Normalize the path
-  target_link_directories(${MJCF_PLUGIN_TARGET_NAME} PRIVATE ${HOUDINI_LIBS})
-
-  target_include_directories(${MJCF_PLUGIN_TARGET_NAME} PRIVATE
-    "${HFS_ENV}/toolkit/include"
-    "${HFS_ENV}/toolkit/include/python3.11"
-  )
-
-  # Assume everyone using Houdini on 3.11 for now.
-  set(USD_MJCF_PYTHON_LIB python3.11)
-  set(USD_MJCF_PYTHON_LIB_NUMBER python311)
-  set(PYTHON_LIB "${HFS_ENV}/Frameworks/Python.framework/Versions/3.11/Python")
-  set(PXR_LIB_PREFIX "pxr_")
-
-  # --- Link Dependencies ---
-  # Link against the necessary OpenUSD components
-  target_link_libraries(${MJCF_PLUGIN_TARGET_NAME} PRIVATE
-    ${PXR_LIB_PREFIX}usd
-    ${PXR_LIB_PREFIX}ar
-    ${PXR_LIB_PREFIX}kind
-    ${PXR_LIB_PREFIX}tf
-    ${PXR_LIB_PREFIX}gf
-    ${PXR_LIB_PREFIX}vt
-    ${PXR_LIB_PREFIX}sdf
-    ${PXR_LIB_PREFIX}usdShade
-    ${PXR_LIB_PREFIX}usdLux
-    ${PXR_LIB_PREFIX}usdGeom
-    ${PXR_LIB_PREFIX}usdImaging
-    ${PXR_LIB_PREFIX}usdPhysics
-    tbb
-    hboost_${USD_MJCF_PYTHON_LIB_NUMBER}
-    ${PYTHON_LIB}
-    mujoco
-    tinyxml2
-  )
-endif()
-
-
-# --- Generate plugInfo.json ---
-if(CMAKE_SHARED_LIBRARY_PREFIX)
-  set(LIB_PREFIX ${CMAKE_SHARED_LIBRARY_PREFIX}) # Usually "lib" on Unix
-else()
-  set(LIB_PREFIX "")
-endif()
-set(PLUG_INFO_LIBRARY_PATH "${LIB_PREFIX}${MJCF_PLUGIN_TARGET_NAME}${CMAKE_SHARED_LIBRARY_SUFFIX}")
-
-# --- Installation ---
-
-set(USD_PLUGIN_INSTALL_DIR_LIB ${CMAKE_INSTALL_LIBDIR}/usdMjcf)
-
-message(STATUS "Copying plugInfo.json to ${CMAKE_BINARY_DIR}/${USD_PLUGIN_INSTALL_DIR_LIB}/plugInfo.json")
-configure_file(
-  mjcf/plugInfo.json
-  ${CMAKE_BINARY_DIR}/${USD_PLUGIN_INSTALL_DIR_LIB}/plugInfo.json
-)
-
-install(FILES ${CMAKE_BINARY_DIR}/${USD_PLUGIN_INSTALL_DIR_LIB}/plugInfo.json DESTINATION ${USD_PLUGIN_INSTALL_DIR_LIB})
-
-# Install shared lib and plugInfo to same location for simplicity.
-install(TARGETS ${MJCF_PLUGIN_TARGET_NAME}
-    LIBRARY DESTINATION ${USD_PLUGIN_INSTALL_DIR_LIB}
-)
-
-message(STATUS "USD MJCF Plugin will be installed to: ${CMAKE_INSTALL_PREFIX}/${USD_PLUGIN_INSTALL_DIR_LIB}")
-message(STATUS "Make sure PXR_PLUGINPATH_NAME includes: ${CMAKE_INSTALL_PREFIX}/${USD_PLUGIN_INSTALL_DIR_LIB}")
@@ -0,0 +1,68 @@
+// 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 <mujoco/experimental/usd/writer.h>
+
+#include <utility>
+#include <vector>
+
+#include <mujoco/experimental/usd/utils.h>
+#include <mujoco/mujoco.h>
+#include <pxr/usd/sdf/path.h>
+#include <pxr/usd/usd/common.h>
+
+namespace mujoco {
+namespace usd {
+
+// Writer is the main entry point for mj_usd. It is responsible
+// for converting the USD scene into a Mujoco model and exposing stepping
+// API.
+Writer::Writer(pxr::UsdStageRefPtr stage, mjSpec* spec, mjModel_* model)
+    : stage_(stage), spec_(spec), model_(model) {
+  BuildMjUsdMapping();
+}
+
+Writer::~Writer() = default;
+
+void Writer::BuildMjUsdMapping() {
+  body_id_to_path_.assign(model_->nbody, pxr::SdfPath());
+
+  // World body (ID 0) maps to an empty SdfPath.
+  if (model_->nbody > 0) {
+    body_id_to_path_[0] = pxr::SdfPath();
+  }
+
+  // Iterate over MuJoCo body IDs from the mjModel.
+  for (int body_idx = 1; body_idx < model_->nbody; ++body_idx) {
+    const char* body_name = mj_id2name(model_, mjOBJ_BODY, body_idx);
+    if (body_name) {
+      mjsBody* spec_body = mjs_findBody(spec_, body_name);
+      body_id_to_path_[body_idx] =
+          mujoco::usd::GetUsdPrimPathUserValue(spec_body->element);
+    }
+  }
+}
+
+void Writer::Update(const mjData* const data) {
+  for (const auto& sink_fn : pose_sinks_) {
+    sink_fn(data, body_id_to_path_);
+  }
+}
+
+void Writer::AddSink(PoseSinkFn sink_fn) {
+  pose_sinks_.push_back(std::move(sink_fn));
+}
+
+}  // namespace usd
+}  // namespace mujoco
@@ -12,7 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-if(MUJOCO_BUILD_USD_PLUGINS AND MUJOCO_USD_TARGET STREQUAL "USD")
+# Only support testing when building against a standalone version of USD.
+if(MUJOCO_WITH_USD AND USD_DIR)
+  find_package(pxr REQUIRED)
+
+  add_library(usd_test_utils STATIC usd/test_utils.h usd/test_utils.cc)
+  target_include_directories(usd_test_utils PUBLIC ${MUJOCO_TEST_INCLUDE})
+  target_compile_definitions(usd_test_utils PUBLIC MJSTATIC)
+  target_link_libraries(
+    usd_test_utils
+    PUBLIC usd
+           tf
+           gtest
+           gmock
+           mujoco
+  )
+
  add_subdirectory(usd/plugins/mjcf)
 endif()

@@ -12,19 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-find_package(pxr REQUIRED)
-
-add_library(usd_fixture STATIC fixture.h fixture.cc)
-target_include_directories(usd_fixture PUBLIC ${MUJOCO_TEST_INCLUDE})
-target_compile_definitions(usd_fixture PUBLIC MJSTATIC)
-
-target_link_libraries(
-  usd_fixture
-  PUBLIC usd
-         tf
-         gtest
-         gmock
-         mujoco
-)
-
-mujoco_test(mjcf_file_format_test ADDITIONAL_LINK_LIBRARIES usd tf usdGeom usdImaging usdPhysics usdShade usd_fixture)
+mujoco_test(mjcf_file_format_test ADDITIONAL_LINK_LIBRARIES
+  usd
+  tf
+  usdGeom
+  usdImaging
+  usdPhysics
+  usdShade
+  usd_test_utils
+  mjcPhysics
+  )