Rework USD scene traversal, ensure DFS order of bodies is respected, remove implicit welding of non-jointed nested bodies

PiperOrigin-RevId: 779121640
Change-Id: Ic3281ac10ce51975cd681e9342b2fffd5f24376f

src/experimental/usd/kinematic_tree.cc

@@ -0,0 +1,182 @@
+// 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/kinematic_tree.h"
+
+#include <algorithm>
+#include <deque>
+#include <map>
+#include <memory>
+#include <set>
+#include <utility>
+#include <vector>
+
+#include <mujoco/mujoco.h>
+#include <pxr/usd/sdf/path.h>
+#include <pxr/usd/usdPhysics/joint.h>
+
+namespace mujoco {
+namespace usd {
+
+bool GetJointBodies(const pxr::UsdPhysicsJoint& joint,
+                    const pxr::SdfPath& default_prim_path, pxr::SdfPath* from,
+                    pxr::SdfPath* to) {
+  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;
+}
+
+std::unique_ptr<KinematicNode> BuildKinematicTree(
+    const std::vector<pxr::UsdPhysicsJoint>& joints,
+    const std::vector<pxr::SdfPath>& all_body_paths,
+    const pxr::SdfPath& default_prim_path) {
+  std::map<pxr::SdfPath, std::vector<pxr::SdfPath>> children_map;
+  std::map<pxr::SdfPath, pxr::SdfPath> parent_map;
+  std::map<std::pair<pxr::SdfPath, pxr::SdfPath>, pxr::SdfPath>
+      edge_to_joint_map;
+  std::set<pxr::SdfPath> all_nodes(all_body_paths.begin(),
+                                   all_body_paths.end());
+
+  for (const auto& joint : joints) {
+    pxr::SdfPath from, to;
+    if (!GetJointBodies(joint, default_prim_path, &from, &to)) {
+      continue;
+    }
+
+    auto edge_key = std::make_pair(from, to);
+    auto it = edge_to_joint_map.find(edge_key);
+    if (it == edge_to_joint_map.end()) {
+      edge_to_joint_map[edge_key] = joint.GetPath();
+    } else {
+      mju_warning(
+          "Multiple explicit joints defined between body %s and body %s. "
+          "Joint1: %s, Joint2: %s. Keeping the first one found: %s",
+          (from.IsEmpty() ? "<worldbody>" : from.GetString()).c_str(),
+          to.GetString().c_str(), it->second.GetString().c_str(),
+          joint.GetPath().GetString().c_str(), it->second.GetString().c_str());
+      continue;
+    }
+
+    if (from == to) {
+      mju_error("Self-loop detected at node %s", to.GetString().c_str());
+      return nullptr;
+    }
+    if (parent_map.count(to)) {
+      mju_error("Node %s has multiple parents ('%s' and '%s').",
+                to.GetString().c_str(), parent_map.at(to).GetString().c_str(),
+                from.GetString().c_str());
+      return nullptr;
+    }
+    children_map[from].push_back(to);
+    parent_map[to] = from;
+    all_nodes.insert(from);
+    all_nodes.insert(to);
+  }
+
+  // Sort children in children_map to respect the DFS order from the stage.
+  for (auto& [_, children] : children_map) {
+    std::sort(
+        children.begin(), children.end(),
+        [&v = all_body_paths](const auto& a, const auto& b) {
+          return std::distance(v.begin(), std::find(v.begin(), v.end(), a)) <
+                 std::distance(v.begin(), std::find(v.begin(), v.end(), b));
+        });
+  }
+
+  // The world body is represented by an empty SdfPath.
+  auto world_root = std::make_unique<KinematicNode>();
+  std::map<pxr::SdfPath, KinematicNode*> node_map;
+  node_map[pxr::SdfPath()] = world_root.get();
+
+  // Use a deque for traversal. We will add roots to the back and children
+  // to the front to perform a DFS on each root's tree.
+  std::deque<pxr::SdfPath> q;
+
+  // Add roots (floating-base bodies and children of the world) to the queue,
+  // preserving the DFS order from the USD stage.
+  for (const auto& body_path : all_body_paths) {
+    if (!body_path.IsEmpty()) {
+      const auto it = parent_map.find(body_path);
+      // A root is a body that has no parent, or its parent is the world.
+      if (it == parent_map.end() || it->second.IsEmpty()) {
+        q.push_back(body_path);
+      }
+    }
+  }
+
+  while (!q.empty()) {
+    pxr::SdfPath current_path = q.front();
+    q.pop_front();
+
+    pxr::SdfPath parent_path = parent_map.count(current_path)
+                                   ? parent_map.at(current_path)
+                                   : pxr::SdfPath();
+    KinematicNode* parent_node = node_map.at(parent_path);
+
+    auto new_node = std::make_unique<KinematicNode>();
+    new_node->body_path = current_path;
+    if (edge_to_joint_map.count({parent_path, current_path})) {
+      new_node->joint_path = edge_to_joint_map.at({parent_path, current_path});
+    }
+    node_map[current_path] = new_node.get();
+    parent_node->children.push_back(std::move(new_node));
+
+    if (children_map.count(current_path)) {
+      const auto& children = children_map.at(current_path);
+      // Add children to the front of the queue in reverse order to ensure
+      // they are processed in the correct order by the DFS.
+      for (auto it = children.rbegin(); it != children.rend(); ++it) {
+        q.push_front(*it);
+      }
+    }
+  }
+
+  // After traversal, check for unvisited nodes.
+  // Unvisited nodes at this point imply a cycle.
+  for (const auto& node : all_nodes) {
+    if (!node.IsEmpty() && !node_map.count(node)) {
+      mju_error("Cycle detected involving node %s.", node.GetString().c_str());
+      return nullptr;
+    }
+  }
+
+  return world_root;
+}
+}  // namespace usd
+}  // namespace mujoco

src/experimental/usd/kinematic_tree.h

@@ -0,0 +1,50 @@
+// 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_KINEMATIC_TREE_H_
+#define MUJOCO_SRC_EXPERIMENTAL_USD_KINEMATIC_TREE_H_
+
+#include <memory>
+#include <vector>
+
+#include <pxr/usd/sdf/path.h>
+#include <pxr/usd/usdPhysics/joint.h>
+
+namespace mujoco {
+namespace usd {
+
+// 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 KinematicNode {
+  pxr::SdfPath body_path;
+  pxr::SdfPath joint_path;  // Joint connecting this node to its parent.
+  std::vector<std::unique_ptr<KinematicNode>> children;
+};
+
+// Builds a single kinematic tree from a list of joints.
+// 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<KinematicNode> BuildKinematicTree(
+    const std::vector<pxr::UsdPhysicsJoint>& joints,
+    const std::vector<pxr::SdfPath>& all_body_paths,
+    const pxr::SdfPath& default_prim_path);
+
+}  // namespace usd
+}  // namespace mujoco
+
+#endif  // MUJOCO_SRC_EXPERIMENTAL_USD_KINEMATIC_TREE_H_

src/experimental/usd/usd_to_mjspec.cc

@@ -12,14 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <algorithm>
 #include <cmath>
 #include <cstddef>
-#include <deque>
-#include <iterator>
 #include <map>
+#include <memory>
 #include <optional>
-#include <set>
 #include <string>
 #include <utility>
 #include <vector>
@@ -35,6 +32,7 @@
 #include <mujoco/experimental/usd/usd.h>
 #include <mujoco/experimental/usd/utils.h>
 #include <mujoco/mujoco.h>
+#include "experimental/usd/kinematic_tree.h"
 #include <pxr/base/gf/declare.h>
 #include <pxr/base/gf/matrix4d.h>
 #include <pxr/base/gf/vec3d.h>
@@ -887,7 +885,7 @@ void ParseMjcPhysicsJointAPI(mjsJoint* mj_joint,
 
 void ParseUsdPhysicsCollider(mjSpec* spec,
                              const pxr::UsdPhysicsCollisionAPI& collision_api,
-                             const pxr::UsdPrim& parent_prim, mjsBody* parent,
+                             const pxr::UsdPrim& body_prim, mjsBody* parent,
                              pxr::UsdGeomXformCache& xform_cache) {
   pxr::UsdPrim prim = collision_api.GetPrim();
   // UsdPhysicsCollisionAPI can only be applied to gprim primitives.
@@ -938,7 +936,7 @@ void ParseUsdPhysicsCollider(mjSpec* spec,
     }
   }
 
-  SetLocalPoseFromPrim(prim, parent_prim, geom, xform_cache);
+  SetLocalPoseFromPrim(prim, body_prim, geom, xform_cache);
 
   if (!MaybeParseGeomPrimitive(prim, geom, xform_cache)) {
     if (prim.IsA<pxr::UsdGeomMesh>()) {
@@ -1139,7 +1137,8 @@ void ParseMjcPhysicsSite(mjSpec* spec, const pxr::MjcPhysicsSiteAPI& site_api,
                          pxr::UsdGeomXformCache& xform_cache) {
   auto prim = site_api.GetPrim();
   mjsSite* site = mjs_addSite(parent, 0);
-  mjs_setName(site->element, site_api.GetPrim().GetPath().GetAsString().c_str());
+  mjs_setName(site->element,
+              site_api.GetPrim().GetPath().GetAsString().c_str());
   SetLocalPoseFromPrim(site_api.GetPrim(), parent_prim, site, xform_cache);
 
   // Convert USD type to MuJoCo type.
@@ -1170,7 +1169,8 @@ void ParseMjcPhysicsKeyframe(mjSpec* spec,
   auto mquat_attr = keyframe.GetMjcMquatAttr();
 
   auto setKeyframeData = [](mjsKey* key, const pxr::UsdAttribute& attr,
-                            std::vector<double>** key_data, double* time = nullptr) {
+                            std::vector<double>** key_data,
+                            double* time = nullptr) {
     if (attr.HasAuthoredValue()) {
       pxr::VtDoubleArray data;
       if (time == nullptr) {
@@ -1257,291 +1257,78 @@ bool IsObjectInPhysicsScene(const T& object,
   return false;
 };
 
-// A struct to hold the forest representation.
+// Helper type to store all the prims that belong to a body.
-// The forest is a map from a root path to its tree.
+using BodyPrimMap = std::map<pxr::SdfPath, std::vector<pxr::SdfPath>>;
-// A tree is an adjacency list, mapping a parent path to its children paths.
-using MjUsdForest =
-    std::map<pxr::SdfPath, std::map<pxr::SdfPath, std::vector<pxr::SdfPath>>>;
 
-// A directed edge.
+// Recursively traverses the kinematic tree, creating bodies, joints, and geoms
-using Edge = std::pair<pxr::SdfPath, pxr::SdfPath>;
+// in the mjSpec.
+void PopulateSpecFromTree(pxr::UsdStageRefPtr stage, mjSpec* spec,
+                          mjsBody* parent_mj_body,
+                          const mujoco::usd::KinematicNode* parent_node,
+                          const mujoco::usd::KinematicNode& current_node,
+                          pxr::UsdGeomXformCache& xform_cache,
+                          const BodyPrimMap& body_to_prims) {
+  mjsBody* current_mj_body;
 
-// A map from a directed edge to the path of the joint representing that edge.
+  if (current_node.body_path.IsEmpty()) {
-using EdgesMap = std::map<Edge, pxr::SdfPath>;
+    // This is the world root node.
-
+    current_mj_body = mjs_findBody(spec, "world");
-// Constructs and validates a forest (a collection of disjoint trees)
-// from a list of directed edges. It also considers all rigid bodies in the
-// scene, treating those not involved in any edge as isolated, free-floating
-// bodies.
-//
-// An empty 'from' path represents the world body. Returns `std::nullopt` for
-// invalid forest structures (e.g. cycles, multiple parents).
-std::optional<MjUsdForest> BuildForestFromEdges(
-    const EdgesMap& edges, const std::vector<pxr::SdfPath>& all_body_paths) {
-  if (edges.empty() && all_body_paths.empty()) {
-    return MjUsdForest{};
-  }
-
-  std::map<pxr::SdfPath, std::vector<pxr::SdfPath>> children_map;
-  std::map<pxr::SdfPath, pxr::SdfPath> parent_map;
-  std::set<pxr::SdfPath> all_nodes(all_body_paths.begin(),
-                                   all_body_paths.end());
-
-  // 1. Initial Pass: Build maps and perform local validation
-  for (const auto& edge_pair : edges) {
-    const auto& [from, to] = edge_pair.first;
-    if (from == to) {
-      mju_error("Self-loop detected at node %s", to.GetString().c_str());
-      return std::nullopt;
-    }
-    if (parent_map.count(to)) {
-      mju_error("Node %s has multiple parents ('%s' and '%s').",
-                to.GetString().c_str(), parent_map.at(to).GetString().c_str(),
-                from.GetString().c_str());
-      return std::nullopt;
-    }
-
-    children_map[from].push_back(to);
-    parent_map[to] = from;
-    all_nodes.insert(from);
-    all_nodes.insert(to);
-  }
-
-  // 2. Find all root nodes
-  // A root is a node that is not a child of any other node. This includes
-  // roots of kinematic trees and isolated rigid bodies.
-  std::set<pxr::SdfPath> roots = all_nodes;
-  for (const auto& pair : parent_map) {
-    roots.erase(pair.first);  // `pair.first` is a child node.
-  }
-
-  if (roots.empty() && !all_nodes.empty()) {
-    mju_error("No root nodes found, but edges exist. A cycle is present.");
-    return std::nullopt;
-  }
-
-  // 3. Build the forest, claiming nodes for each tree
-  MjUsdForest forest;
-  std::set<pxr::SdfPath> claimed_nodes;
-
-  for (const auto& root : roots) {
-    // Perform a traversal (BFS) to find all nodes in this tree
-    std::deque<pxr::SdfPath> q;
-    q.push_back(root);
-    std::set<pxr::SdfPath> nodes_in_this_tree;
-
-    while (!q.empty()) {
-      pxr::SdfPath current_node = q.front();
-      q.pop_front();
-
-      if (claimed_nodes.count(current_node)) {
-        mju_error("Node %s is shared between multiple trees.",
-                  current_node.GetString().c_str());
-        return std::nullopt;
-      }
-
-      nodes_in_this_tree.insert(current_node);
-      claimed_nodes.insert(current_node);
-
-      if (children_map.count(current_node)) {
-        for (const auto& child : children_map.at(current_node)) {
-          q.push_back(child);
-        }
-      }
-    }
-
-    // Construct the adjacency list for this specific tree
-    std::map<pxr::SdfPath, std::vector<pxr::SdfPath>> tree_adj_list;
-    for (const auto& node : nodes_in_this_tree) {
-      if (children_map.count(node)) {
-        tree_adj_list[node] = children_map.at(node);
-      }
-    }
-    forest[root] = tree_adj_list;
-  }
-
-  // 4. Final check for cycles (unclaimed nodes)
-  std::set<pxr::SdfPath> unclaimed_nodes;  // all_nodes - claimed_nodes
-  std::set_difference(all_nodes.begin(), all_nodes.end(), claimed_nodes.begin(),
-                      claimed_nodes.end(),
-                      std::inserter(unclaimed_nodes, unclaimed_nodes.begin()));
-
-  if (!unclaimed_nodes.empty()) {
-    std::string unclaimed_str;
-    for (const auto& node : unclaimed_nodes) {
-      unclaimed_str += "'" + node.GetString() + "' ";
-    }
-    mju_error(
-        "Cycle detected. The following nodes are part of a cycle "
-        "and not reachable from any root: %s",
-        unclaimed_str.c_str());
-    return std::nullopt;
-  }
-
-  return forest;
-}
-
-void TraverseAndBuildTree(
-    pxr::UsdStageRefPtr stage, mjSpec* spec, mjsBody* parent_mj_body,
-    const pxr::SdfPath& parent_body_path, const pxr::SdfPath& current_body_path,
-    const std::map<pxr::SdfPath, std::vector<pxr::SdfPath>>& tree,
-    const EdgesMap& edges, pxr::UsdGeomXformCache& xform_cache);
-
-// Traverses the prim and all its descendants in the USD hierarchy and parses
-// supported entities like colliders and sites, attaching them to the given
-// mjBody. The traversal for a given branch stops when:
-// - a descendant with a RigidBodyAPI is found, as that will be handled by
-// TraverseAndBuildTree.
-// - a descendant with a ResetXformStack is found, as that will be handled by
-// the top level traversal of independent prims.
-void ParseCurrentAndDescendants(mjSpec* spec, const pxr::UsdPrim& prim,
-                                const pxr::UsdPrim& parent_prim, mjsBody* body,
-                                pxr::UsdGeomXformCache& xform_cache) {
-  if (prim.HasAPI<pxr::UsdPhysicsCollisionAPI>()) {
-    ParseUsdPhysicsCollider(spec, pxr::UsdPhysicsCollisionAPI(prim),
-                            parent_prim, body, xform_cache);
-  }
-  if (prim.HasAPI<pxr::MjcPhysicsSiteAPI>()) {
-    ParseMjcPhysicsSite(spec, pxr::MjcPhysicsSiteAPI(prim), parent_prim, body,
-                        xform_cache);
-  }
-
-  // Make sure we traverse into instance proxies to ensure we support
-  // instanceable references.
-  // See https://openusd.org/dev/api/_usd__page__scenegraph_instancing.html
-  for (const auto& child :
-       prim.GetFilteredChildren(pxr::UsdTraverseInstanceProxies())) {
-    if (child.HasAPI<pxr::UsdPhysicsRigidBodyAPI>()) {
-      continue;
-    }
-    if (xform_cache.GetResetXformStack(child)) {
-      continue;
-    }
-    ParseCurrentAndDescendants(spec, child, prim, body, xform_cache);
-  }
-}
-
-// Recursively traverses a kinematic tree, creating bodies and joints in the
-// mjSpec.
-void TraverseAndBuildTree(
-    pxr::UsdStageRefPtr stage, mjSpec* spec, mjsBody* parent_mj_body,
-    const pxr::SdfPath& parent_body_path, const pxr::SdfPath& current_body_path,
-    const std::map<pxr::SdfPath, std::vector<pxr::SdfPath>>& tree,
-    const EdgesMap& edges, pxr::UsdGeomXformCache& xform_cache) {
-  pxr::UsdPrim current_body_prim = stage->GetPrimAtPath(current_body_path);
-  pxr::UsdPrim parent_prim_for_xform =
-      parent_body_path.IsEmpty() ? stage->GetPseudoRoot()
-                                 : stage->GetPrimAtPath(parent_body_path);
-
-  mjsBody* current_mj_body = ParseUsdPhysicsRigidbody(
-      spec, pxr::UsdPhysicsRigidBodyAPI(current_body_prim),
-      parent_prim_for_xform, parent_mj_body, xform_cache);
-
-  auto edge_key = std::make_pair(parent_body_path, current_body_path);
-  auto it_edge = edges.find(edge_key);
-
-  if (it_edge != edges.end()) {
-    // An edge exists, indicating a connection to the parent. This body is
-    // either world-attached or part of a larger articulation.
-    const pxr::SdfPath& joint_path = it_edge->second;
-    if (!joint_path.IsEmpty()) {
-      // An explicit joint prim exists, so we parse it.
-      pxr::UsdPrim joint_prim = stage->GetPrimAtPath(joint_path);
-      ParseUsdPhysicsJoint(spec, joint_prim, current_mj_body, xform_cache);
-    }
-    // If joint_path is empty, it's an implicit fixed joint. No joint is created
-    // in the mjSpec, effectively welding the body to its parent.
   } else {
-    // No edge found. This condition is met for the root of a floating-base
+    // This is a regular body.
-    // tree, which has no defined joint connecting it to the world.
+    pxr::UsdPrim current_body_prim =
-    if (parent_mj_body == mjs_findBody(spec, "world")) {
+        stage->GetPrimAtPath(current_node.body_path);
-      // We explicitly create a free joint to make it a floating-base body.
+    pxr::SdfPath parent_body_path =
+        parent_node ? parent_node->body_path : pxr::SdfPath();
+    pxr::UsdPrim parent_prim_for_xform =
+        parent_body_path.IsEmpty() ? stage->GetPseudoRoot()
+                                   : stage->GetPrimAtPath(parent_body_path);
+
+    current_mj_body = ParseUsdPhysicsRigidbody(
+        spec, pxr::UsdPhysicsRigidBodyAPI(current_body_prim),
+        parent_prim_for_xform, parent_mj_body, xform_cache);
+
+    if (!current_node.joint_path.IsEmpty()) {
+      pxr::UsdPrim joint_prim = stage->GetPrimAtPath(current_node.joint_path);
+      ParseUsdPhysicsJoint(spec, joint_prim, current_mj_body, xform_cache);
+    } else if (parent_mj_body == mjs_findBody(spec, "world")) {
+      // No joint to parent, and parent is world: this is a floating body.
       mjsJoint* free_joint = mjs_addJoint(current_mj_body, nullptr);
       free_joint->type = mjJNT_FREE;
     }
   }
 
-  // Parse all geoms/sites that are found on this body in the USD hierarchy.
+  // Add geoms/sites/etc. belonging to the current body.
-  ParseCurrentAndDescendants(spec, current_body_prim, parent_prim_for_xform,
+  auto it_prims = body_to_prims.find(current_node.body_path);
-                             current_mj_body, xform_cache);
+  if (it_prims != body_to_prims.end()) {
-  // Recurse through the kinematic tree.
+    pxr::UsdPrim body_prim_for_xform =
-  auto it_tree = tree.find(current_body_path);
+        current_node.body_path.IsEmpty()
-  if (it_tree != tree.end()) {
+            ? stage->GetPseudoRoot()
-    const auto& children_paths = it_tree->second;
+            : stage->GetPrimAtPath(current_node.body_path);
-    for (const auto& child_path : children_paths) {
+    for (const auto& gprim_path : it_prims->second) {
-      TraverseAndBuildTree(stage, spec, current_mj_body, current_body_path,
+      pxr::UsdPrim prim = stage->GetPrimAtPath(gprim_path);
-                           child_path, tree, edges, xform_cache);
+      if (prim.HasAPI<pxr::UsdPhysicsCollisionAPI>()) {
-    }
+        ParseUsdPhysicsCollider(spec, pxr::UsdPhysicsCollisionAPI(prim),
-  }
+                                body_prim_for_xform, current_mj_body,
-}
+                                xform_cache);
-
+      }
-// Adds a new edge to the edges map, dealing with duplicates:
+      if (prim.HasAPI<pxr::MjcPhysicsSiteAPI>()) {
-// - explicit joints always replace implicit joints
+        ParseMjcPhysicsSite(spec, pxr::MjcPhysicsSiteAPI(prim),
-// - more than one explicit joint is unsupported so we print a warning and
+                            body_prim_for_xform, current_mj_body, xform_cache);
-// keep the first one found
-void AddEdge(EdgesMap& edges, const pxr::SdfPath& from, const pxr::SdfPath& to,
-             const pxr::SdfPath& joint) {
-  auto edge_key = std::make_pair(from, to);
-  auto it = edges.find(edge_key);
-
-  if (it == edges.end()) {
-    // No existing edge, add the new one.
-    edges[edge_key] = joint;
-  } else {
-    // Edge already exists.
-    pxr::SdfPath& existing_joint = it->second;
-    bool new_is_explicit = !joint.IsEmpty();
-    bool existing_is_explicit = !existing_joint.IsEmpty();
-
-    if (new_is_explicit) {
-      if (existing_is_explicit) {
-        // Both are explicit: this is an error condition.
-        mju_warning(
-            "Multiple explicit joints defined between body %s and body %s. "
-            "Joint1: %s, Joint2: %s. Keeping the first one found: %s",
-            (from.IsEmpty() ? "<worldbody>" : from.GetString()).c_str(),
-            to.GetString().c_str(), existing_joint.GetString().c_str(),
-            joint.GetString().c_str(), existing_joint.GetString().c_str());
-
-      } else {
-        // New is explicit, existing is implicit: replace.
-        existing_joint = joint;
       }
     }
-    // If new is implicit, and an edge already exists (either explicit or
-    // implicit), we keep the existing one. No action needed.
   }
-}
 
-// Returns the nesting body prim (or an invalid prim if there's
+  // Recurse through children.
-// any resets_xform_stack, or we've reached the end).
+  for (const auto& child_node : current_node.children) {
-pxr::UsdPrim GetNestingBodyPrim(const pxr::UsdPrim& prim,
+    PopulateSpecFromTree(stage, spec, current_mj_body, &current_node,
-                                pxr::UsdGeomXformCache& xform_cache) {
+                         *child_node, xform_cache, body_to_prims);
-  if (xform_cache.GetResetXformStack(prim)) {
-    return pxr::UsdPrim();
   }
-  pxr::UsdPrim previous_prim = prim.GetParent();
-  while (previous_prim.IsValid()) {
-    // 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();
 }
 }  // namespace
 
 mjSpec* mj_parseUSDStage(const pxr::UsdStageRefPtr stage) {
   mjSpec* spec = mj_makeSpec();
 
-  mjsBody* world = mjs_findBody(spec, "world");
-
   std::vector<pxr::UsdPhysicsScene> physics_scenes;
 
   // Xform cache to use for all queries when parsing.
@@ -1573,135 +1360,79 @@ mjSpec* mj_parseUSDStage(const pxr::UsdStageRefPtr stage) {
     default_prim_path = stage->GetDefaultPrim().GetPath();
   }
 
-  std::vector<pxr::SdfPath> body_paths;
+  // Data Structures
-  EdgesMap edges;
+  std::vector<pxr::UsdPhysicsJoint> all_joints;
-  // TODO(robinalazard): Re-introduce properly adding objects to their
+  std::vector<pxr::SdfPath> all_body_paths_vec;
-  // respective physics scene, or default.
+  BodyPrimMap body_to_prims;
 
-  // Traverse all prims under the pseudo-root.
+  // =========================================================================
-  // We ensure to traverse into instance proxies to ensure we support
+  // PASS 1: Collect Bodies, Joints, and Geoms/Sites/etc.
-  // instanceable references.
+  // =========================================================================
-  // See https://openusd.org/dev/api/_usd__page__scenegraph_instancing.html
+  // A single DFS pass to find all bodies, joints, and determine
-  for (auto prim : stage->Traverse(pxr::UsdTraverseInstanceProxies())) {
+  // which body owns each geom/site/etc. prim.
-    // When traversing the whole scene, if we encounter a rigidbody or a joint,
+  std::vector<pxr::SdfPath> owner_stack;
-    // then we populate the edges map and the list of body paths, which will be
+  owner_stack.push_back(pxr::SdfPath());  // Start with the world as owner.
-    // processed later to build the articulation trees.
+
-    //
+  const auto range = pxr::UsdPrimRange::PreAndPostVisit(
-    // If we encounter _anything else_:
+      stage->GetPseudoRoot(), pxr::UsdTraverseInstanceProxies());
-    // - if we find that it's an independent prim (e.g. a static collider or
+
-    // site) not belonging to any rigidbody, then we add it directly to the
+  for (auto it = range.begin(); it != range.end(); ++it) {
-    // world.
+    pxr::UsdPrim prim = *it;
-    // - otherwise, they will be handled when building the articulation trees.
+
+    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();
+    pxr::SdfPath prim_owner = owner_stack.back();
+
+    if (is_body) {
+      all_body_paths_vec.push_back(prim_path);
+      prim_owner = prim_path;
+    } else if (resets) {
+      prim_owner = pxr::SdfPath();  // Reset owner to world.
+    }
+
+    if (is_pushed_to_stack) {
+      owner_stack.push_back(prim_owner);
+    }
+
+    if (prim.HasAPI<pxr::UsdPhysicsCollisionAPI>() ||
+        prim.HasAPI<pxr::MjcPhysicsSiteAPI>()) {
+      body_to_prims[prim_owner].push_back(prim_path);
+    }
 
     if (prim.IsA<pxr::UsdPhysicsJoint>()) {
-      pxr::SdfPath joint_path = prim.GetPath();
+      all_joints.push_back(pxr::UsdPhysicsJoint(prim));
-      pxr::UsdPhysicsJoint joint(prim);
 
-      pxr::SdfPath body1_path;
+      it.PruneChildren();
-      pxr::SdfPathVector body1_paths;
-      joint.GetBody1Rel().GetTargets(&body1_paths);
-      if (body1_paths.empty()) {
-        mju_warning("Joint %s does not have body1 rel. Skipping.",
-                    prim.GetPath().GetAsString().c_str());
-        continue;
-      } else if (body1_paths.size() > 1) {
-        mju_warning("Joint %s has multiple body1 rels. Skipping.",
-                    prim.GetPath().GetAsString().c_str());
-        continue;
-      }
-      body1_path = body1_paths[0];
-
-      pxr::SdfPath body0_path;
-      pxr::SdfPathVector body0_paths;
-      joint.GetBody0Rel().GetTargets(&body0_paths);
-      if (body0_paths.size() > 1) {
-        mju_warning("Joint %s has multiple body0 rels. Skipping.",
-                    prim.GetPath().GetAsString().c_str());
-        continue;
-      }
-      // 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) {
-        body0_path = pxr::SdfPath();
-      } else {
-        body0_path = body0_paths[0];
-      }
-
-      AddEdge(edges, body0_path, body1_path, joint_path);
     } else if (prim.IsA<pxr::MjcPhysicsKeyframe>()) {
       ParseMjcPhysicsKeyframe(spec, pxr::MjcPhysicsKeyframe(prim));
-      continue;
-    } else if (prim.HasAPI<pxr::UsdPhysicsRigidBodyAPI>()) {
-      pxr::SdfPath body_path = prim.GetPath();
-      body_paths.push_back(body_path);
 
-      // Find whether we are nested under a parent body.
+      it.PruneChildren();
-      // Note: if any xform in between (including the current prim) resets the
-      // xform stack then we are not nested.
-      pxr::UsdPrim nesting_body_prim = GetNestingBodyPrim(prim, xform_cache);
-      if (nesting_body_prim.IsValid()) {
-        AddEdge(edges, nesting_body_prim.GetPath(), body_path, pxr::SdfPath());
-      }
-    } else {
-      // TODO(robinalazard): the way we handle independent prims right now means
-      // their relative transforms will be ignored be always pass
-      // stage->GetPseudoRoot() as the parent. It works for most scene
-      // realistically. But we should fix it.
-
-      if (prim.HasAPI<pxr::UsdPhysicsCollisionAPI>()) {
-        // Find whether the collider belongs to a body. If yes, it will be we
-        // handled later when building the articulation trees. Otherwise, it's a
-        // static collider and we add it directly to the world.
-        pxr::UsdPrim nesting_body_prim = GetNestingBodyPrim(prim, xform_cache);
-        if (!nesting_body_prim.IsValid()) {
-          ParseUsdPhysicsCollider(spec, pxr::UsdPhysicsCollisionAPI(prim),
-                                  stage->GetPseudoRoot(), world, xform_cache);
-        }
-      }
-      if (prim.HasAPI<pxr::MjcPhysicsSiteAPI>()) {
-        // Find whether the site belongs to a body. If yes, it will be we
-        // handled later when building the articulation trees. Otherwise, it's a
-        // static site and we add it directly to the world.
-        pxr::UsdPrim nesting_body_prim = GetNestingBodyPrim(prim, xform_cache);
-        if (!nesting_body_prim.IsValid()) {
-          ParseMjcPhysicsSite(spec, pxr::MjcPhysicsSiteAPI(prim),
-                              stage->GetPseudoRoot(), world, xform_cache);
-        }
-      }
     }
   }
 
-  std::optional<MjUsdForest> forest = BuildForestFromEdges(edges, body_paths);
+  // =========================================================================
-  if (forest.has_value()) {
+  // PASS 2: Build the kinematic tree and populate the mjSpec.
-    // Now that we have the forest, we can walk through it and add the bodies
+  // =========================================================================
-    // and joints to the spec. From each body we can also visit and parse all
+  std::unique_ptr<mujoco::usd::KinematicNode> kinematic_tree =
-    // the children to also add the corresponding colliders to the body in the
+      mujoco::usd::BuildKinematicTree(all_joints, all_body_paths_vec,
-    // spec.
+                                      default_prim_path);
-    for (const auto& [root_path, tree] : *forest) {
+
-      if (root_path.IsEmpty()) {
+  if (kinematic_tree) {
-        // This root is the world. This case handles all kinematic trees that
+    PopulateSpecFromTree(stage, spec, /*parent_mj_body=*/nullptr,
-        // are attached to the world (fixed-base articulations).
+                         /*parent_node=*/nullptr, *kinematic_tree, xform_cache,
-        // We iterate through its direct children which are the root bodies of
+                         body_to_prims);
-        // each world-attached tree. We iterate through these children and begin
-        // the recursive build from there.
-        const auto& children_of_world = tree.at(pxr::SdfPath());
-        for (const auto& child_path : children_of_world) {
-          TraverseAndBuildTree(stage, spec, world, root_path, child_path, tree,
-                               edges, xform_cache);
-        }
-      } else {
-        // Conversely, this case handles all the remaining top-level root bodies
-        // which are their own roots and are not attached to the world
-        // (floating-base articulations). This includes isolated bodies.
-        // We directly begin the recursive build from the toplevel root.
-        //
-        // Note: the absence of an edge in the `edges` map connecting the world
-        // to this root is what signals to `TraverseAndBuildTree` that this is a
-        // floating base, prompting the creation of a free joint.
-        TraverseAndBuildTree(stage, spec, world, pxr::SdfPath(), root_path,
-                             tree, edges, xform_cache);
-      }
-    }
   }
 
   return spec;