From 072b32743034ad6890e5079a64bae42fae624037 Mon Sep 17 00:00:00 2001
From: Robin Alazard <robinalazard@google.com>
Date: Tue, 1 Jul 2025 11:06:13 -0700
Subject: [PATCH] Fix the parsing of USD scenes that use ResetXformStack

Also add a new test for it

PiperOrigin-RevId: 778114814
Change-Id: Ib019768ac45ecf7ec6d437199a3d1427557b2f06
---
 src/experimental/usd/usd_to_mjspec.cc | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/experimental/usd/usd_to_mjspec.cc b/src/experimental/usd/usd_to_mjspec.cc
index da4e4f2a..7a25e156 100644
--- a/src/experimental/usd/usd_to_mjspec.cc
+++ b/src/experimental/usd/usd_to_mjspec.cc
@@ -63,7 +63,6 @@
 #include <pxr/usd/usdPhysics/revoluteJoint.h>
 #include <pxr/usd/usdPhysics/rigidBodyAPI.h>
 #include <pxr/usd/usdPhysics/scene.h>
-#include <pxr/usd/usdPhysics/sphericalJoint.h>
 namespace {
 
 using pxr::MjcPhysicsTokens;
@@ -1350,7 +1349,7 @@ void ParseCurrentAndDescendants(mjSpec* spec, const pxr::UsdPrim& prim,
     if (child.HasAPI<pxr::UsdPhysicsRigidBodyAPI>()) {
       continue;
     }
-    if (xform_cache.GetResetXformStack(prim)) {
+    if (xform_cache.GetResetXformStack(child)) {
       continue;
     }
     ParseCurrentAndDescendants(spec, child, prim, body, xform_cache);
@@ -1458,12 +1457,15 @@ pxr::UsdPrim GetNestingBodyPrim(const pxr::UsdPrim& prim,
   }
   pxr::UsdPrim previous_prim = prim.GetParent();
   while (previous_prim.IsValid()) {
-    if (xform_cache.GetResetXformStack(previous_prim)) {
-      return pxr::UsdPrim();
-    }
+    // If we find a rigid body, this is our answer. The prim is nested.
     if (previous_prim.HasAPI<pxr::UsdPhysicsRigidBodyAPI>()) {
       return previous_prim;
     }
+    // If we encounter a prim that resets the transform stack *before* finding
+    // a rigid body, the chain is broken. The prim is not nested.
+    if (xform_cache.GetResetXformStack(previous_prim)) {
+      return pxr::UsdPrim();
+    }
     previous_prim = previous_prim.GetParent();
   }
   return pxr::UsdPrim();