From e80a62c6d5e9b078bdc8e1ecdb76156c9739e08f Mon Sep 17 00:00:00 2001
From: Sam Haves <shaves@google.com>
Date: Fri, 27 Jun 2025 07:58:27 -0700
Subject: [PATCH] Allow loading USD files into simulate.

PiperOrigin-RevId: 776584803
Change-Id: I66a5277b3752d69ba5d44ed419204951edb18d41
---
 simulate/CMakeLists.txt             | 24 +++++++++++++
 simulate/main.cc                    | 53 +++++++++++++++++++++++++++--
 src/experimental/usd/CMakeLists.txt |  2 ++
 3 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/simulate/CMakeLists.txt b/simulate/CMakeLists.txt
index 85ffaa4e..808de669 100644
--- a/simulate/CMakeLists.txt
+++ b/simulate/CMakeLists.txt
@@ -174,6 +174,14 @@ if(SIMULATE_BUILD_EXECUTABLE)
     lodepng
   )
 
+  if (MUJOCO_WITH_USD)
+    target_link_libraries(
+      simulate
+      mujoco::usd
+      mujoco::usd::mjcf
+    )
+  endif()
+
   target_link_options(simulate PRIVATE ${MUJOCO_SIMULATE_LINK_OPTIONS})
 
   if(APPLE AND MUJOCO_BUILD_MACOS_FRAMEWORKS)
@@ -240,6 +248,22 @@ if(SIMULATE_BUILD_EXECUTABLE)
       MUJOCO_ENABLE_RPATH
     )
 
+    if (MUJOCO_WITH_USD)
+      target_compile_definitions(simulate PUBLIC SIMULATE_WITH_USD)
+
+      # Add support to RPATH for the samples.
+      target_add_rpath(
+        TARGETS
+        simulate
+        INSTALL_DIRECTORY
+        "${CMAKE_INSTALL_PREFIX}/${CMAKE_INSTALL_BINDIR}"
+        LIB_DIRS
+        ${CMAKE_INSTALL_PREFIX}/${CMAKE_INSTALL_LIBDIR}/mujocoUsd ${USD_DIR}/lib
+        DEPENDS
+        MUJOCO_ENABLE_RPATH
+      )
+    endif()
+
     install(
       TARGETS simulate
       EXPORT ${PROJECT_NAME}
diff --git a/simulate/main.cc b/simulate/main.cc
index 1b5ccb65..e3d42a03 100644
--- a/simulate/main.cc
+++ b/simulate/main.cc
@@ -25,6 +25,12 @@
 #include <string>
 #include <thread>
 
+#if defined(SIMULATE_WITH_USD)
+#include <mujoco/experimental/usd/usd.h>
+#include <pxr/pxr.h>
+#include <pxr/usd/usd/common.h>
+#include <pxr/usd/usd/stage.h>
+#endif
 #include <mujoco/mujoco.h>
 #include "glfw_adapter.h"
 #include "simulate.h"
@@ -224,13 +230,49 @@ mjModel* LoadModel(const char* file, mj::Simulate& sim) {
   char loadError[kErrorLength] = "";
   mjModel* mnew = 0;
   auto load_start = mj::Simulate::Clock::now();
-  if (mju::strlen_arr(filename)>4 &&
-      !std::strncmp(filename + mju::strlen_arr(filename) - 4, ".mjb",
-                    mju::sizeof_arr(filename) - mju::strlen_arr(filename)+4)) {
+
+  std::string filename_str(filename);
+  std::string extension;
+  size_t dot_pos = filename_str.rfind('.');
+
+  if (dot_pos != std::string::npos && dot_pos < filename_str.length() - 1) {
+    extension = filename_str.substr(dot_pos);
+  }
+
+  if (extension == ".mjb") {
     mnew = mj_loadModel(filename, nullptr);
     if (!mnew) {
       mju::strcpy_arr(loadError, "could not load binary model");
     }
+#if defined(SIMULATE_WITH_USD)
+  } else if (extension == ".usda" || extension == ".usd" ||
+             extension == ".usdc" || extension == ".usdz" ) {
+    auto stage = pxr::UsdStage::Open(filename);
+    if (!stage) {
+      mju::strcpy_arr(loadError, "could not open USD stage");
+    } else {
+      mjSpec* spec = mj_parseUSDStage(stage);
+      if (!spec) {
+        mju::strcpy_arr(loadError, "could not parse USD stage to mjSpec");
+      } else {
+        mjModel* model = mj_compile(spec, nullptr);
+        if (!model) {
+          mju::strcpy_arr(loadError,
+                          "could not compile USD parsed mjSpec to mjModel:\n");
+          mju::strcat_arr(loadError, mjs_getError(spec));
+        } else {
+          // handle compile warning
+          if (mjs_isWarning(spec)) {
+            mju::strcpy_arr(
+                loadError,
+                "warning while compiling USD parsed mjSpec to mjModel:\n");
+            mju::strcat_arr(loadError, mjs_getError(spec));
+          }
+        }
+        mnew = model;
+      }
+    }
+#endif
   } else {
     mnew = mj_loadXML(filename, nullptr, loadError, kErrorLength);
 
@@ -495,6 +537,11 @@ int main(int argc, char** argv) {
   // scan for libraries in the plugin directory to load additional plugins
   scanPluginLibraries();
 
+#if defined(SIMULATE_WITH_USD)
+  // If USD is used, print the version.
+  std::printf("OpenUSD version v%d.%02d\n", PXR_MINOR_VERSION, PXR_PATCH_VERSION);
+#endif
+
   mjvCamera cam;
   mjv_defaultCamera(&cam);
 
diff --git a/src/experimental/usd/CMakeLists.txt b/src/experimental/usd/CMakeLists.txt
index 4a256f58..96cd9eeb 100644
--- a/src/experimental/usd/CMakeLists.txt
+++ b/src/experimental/usd/CMakeLists.txt
@@ -152,6 +152,7 @@ endif()
 set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
 
 add_library(${MJCF_PLUGIN_TARGET_NAME} SHARED)
+add_library(mujoco::usd::mjcf ALIAS ${MJCF_PLUGIN_TARGET_NAME})
 target_sources(${MJCF_PLUGIN_TARGET_NAME} PRIVATE
   plugins/mjcf/mjcf_file_format.cc
   plugins/mjcf/mjcf_file_format.h
@@ -190,6 +191,7 @@ target_include_directories(${MJC_PHYSICS_PLUGIN_TARGET_NAME} PUBLIC
 ## mj_usd
 
 add_library(${MJ_USD_TARGET_NAME} SHARED)
+add_library(mujoco::usd ALIAS ${MJ_USD_TARGET_NAME})
 target_sources(${MJ_USD_TARGET_NAME} PRIVATE
   writer.cc
   layer_sink.cc