From 2e13c6430813a0d267136bf38432ff7260bd2603 Mon Sep 17 00:00:00 2001 From: Sam Haves Date: Thu, 17 Jul 2025 08:18:59 -0700 Subject: [PATCH] Simplify parsing from USD stage to kinematic tree for mjSpec population. Summary of changes: - Move stage parsing from usd_to_mjspec to kinematic_tree - A Node now has lists of paths for prims that we support parsing. - When parsing the stage, we create a list of all the Nodes that we will place in our tree. - Removed the need for expensive SdfPath maps and sorting of bodies in favor of indexed integer arrays. - Simplifies usd_to_mjspec parsing as we can process each node and it's owned prims one at a time. PiperOrigin-RevId: 784187799 Change-Id: I4907cfa1137014af21c45b10fa29ac53a49b6dac --- src/experimental/usd/kinematic_tree.cc | 258 ++++++++++++++++--------- src/experimental/usd/kinematic_tree.h | 21 +- src/experimental/usd/usd_to_mjspec.cc | 208 ++++++-------------- 3 files changed, 244 insertions(+), 243 deletions(-) diff --git a/src/experimental/usd/kinematic_tree.cc b/src/experimental/usd/kinematic_tree.cc index d6314ad6..88273095 100644 --- a/src/experimental/usd/kinematic_tree.cc +++ b/src/experimental/usd/kinematic_tree.cc @@ -14,24 +14,36 @@ #include "experimental/usd/kinematic_tree.h" -#include -#include #include #include -#include #include #include +#include +#include +#include #include #include +#include +#include +#include +#include #include +#include +#include namespace mujoco { namespace usd { -bool GetJointBodies(const pxr::UsdPhysicsJoint& joint, - const pxr::SdfPath& default_prim_path, pxr::SdfPath* from, +bool GetJointBodies(const pxr::UsdPhysicsJoint& joint, pxr::SdfPath* from, pxr::SdfPath* to) { + // Grab the default prim path + pxr::SdfPath default_prim_path; + auto stage = joint.GetPrim().GetStage(); + if (stage->GetDefaultPrim().IsValid()) { + default_prim_path = stage->GetDefaultPrim().GetPath(); + } + pxr::SdfPathVector body1_paths; joint.GetBody1Rel().GetTargets(&body1_paths); if (body1_paths.empty()) { @@ -62,121 +74,187 @@ bool GetJointBodies(const pxr::UsdPhysicsJoint& joint, return true; } -std::unique_ptr BuildKinematicTree( - const std::vector& joints, - const std::vector& all_body_paths, - const pxr::SdfPath& default_prim_path) { - std::map> children_map; - std::map parent_map; - std::map, pxr::SdfPath> - edge_to_joint_map; - std::set all_nodes(all_body_paths.begin(), - all_body_paths.end()); +struct ExtractedPrims { + std::vector> nodes; + std::vector joints; +}; - for (const auto& joint : joints) { - pxr::SdfPath from, to; - if (!GetJointBodies(joint, default_prim_path, &from, &to)) { +ExtractedPrims ExtractPrims(pxr::UsdStageRefPtr stage) { + pxr::UsdPhysicsScene physics_scene; + std::vector> nodes; + std::vector joints; + nodes.push_back(std::make_unique()); + Node* root = nodes.back().get(); + + // ========================================================================= + // PASS 1: Collect Bodies, Joints, and Geoms/Sites/etc. + // ========================================================================= + // A single DFS pass to find all bodies, joints, and determine + // which body owns each geom/site/etc. prim. + std::vector owner_stack; + owner_stack.push_back(root); // Start with the world as owner. + + const auto range = pxr::UsdPrimRange::PreAndPostVisit( + stage->GetPseudoRoot(), pxr::UsdTraverseInstanceProxies()); + + pxr::UsdGeomXformCache xform_cache; + for (auto it = range.begin(); it != range.end(); ++it) { + pxr::UsdPrim prim = *it; + + bool is_body = prim.HasAPI(); + 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; } - 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() ? "" : from.GetString()).c_str(), - to.GetString().c_str(), it->second.GetString().c_str(), - joint.GetPath().GetString().c_str(), it->second.GetString().c_str()); - continue; + pxr::SdfPath prim_path = prim.GetPath(); + Node* current_node = owner_stack.back(); + + if (is_body) { + auto new_node = std::make_unique(); + new_node->body_path = prim_path; + nodes.push_back(std::move(new_node)); + current_node = nodes.back().get(); + } else if (resets) { + current_node = root; // Reset owner to world. } - if (from == to) { - mju_error("Self-loop detected at node %s", to.GetString().c_str()); - return nullptr; + if (is_pushed_to_stack) { + owner_stack.push_back(current_node); } - 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; + + if (prim.IsA() && root->physics_scene.IsEmpty()) { + root->physics_scene = prim_path; } - children_map[from].push_back(to); - parent_map[to] = from; - all_nodes.insert(from); - all_nodes.insert(to); + + if (prim.HasAPI()) { + current_node->colliders.push_back(prim.GetPath()); + } + + if (prim.HasAPI()) { + current_node->sites.push_back(prim.GetPath()); + // Sites should not have children. + it.PruneChildren(); + } + + if (prim.IsA()) { + // We may not know which body this belongs to yet so we'll add it to a + // list and the caller can assign the joints when building the tree. + joints.push_back(pxr::UsdPhysicsJoint(prim)); + // Joints should not have children. + it.PruneChildren(); + } + + if (prim.IsA()) { + root->actuators.push_back(prim_path); + // Joints should not have children. + it.PruneChildren(); + } + + if (prim.IsA()) { + root->keyframes.push_back(prim.GetPath()); + // Keyframes should not have children. + it.PruneChildren(); + } + } + return {.nodes = std::move(nodes), .joints = std::move(joints)}; +} + +std::unique_ptr BuildKinematicTree(const pxr::UsdStageRefPtr stage) { + ExtractedPrims extraction = ExtractPrims(stage); + + std::map body_index; + body_index[pxr::SdfPath()] = 0; + for (int i = 0; i < extraction.nodes.size(); ++i) { + body_index[extraction.nodes[i]->body_path] = i; } - // 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)); - }); + // List of direct children for each body. + std::vector> children( + extraction.nodes.size(), std::vector(extraction.nodes.size())); + // List of joint prim paths associated with a child for each body. + std::vector> parent_joints(extraction.nodes.size()); + for (const pxr::UsdPhysicsJoint& joint : extraction.joints) { + pxr::SdfPath from, to; + if (!GetJointBodies(joint, &from, &to)) { + continue; + } + + int from_idx = body_index[from]; + int to_idx = body_index[to]; + if (from_idx == to_idx) { + mju_error("Cycle detected: self referencing joint at node %s", + to.GetString().c_str()); + return nullptr; + } + + children[from_idx][to_idx] = true; + parent_joints[to_idx].push_back(joint.GetPath()); + // Now that we know all the bodies, we can assign joints to respective + // nodes. + extraction.nodes[to_idx]->joints.push_back(joint.GetPath()); } // The world body is represented by an empty SdfPath. - auto world_root = std::make_unique(); - std::map node_map; - node_map[pxr::SdfPath()] = world_root.get(); + auto world_root = std::move(extraction.nodes[0]); - // 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 q; + std::vector> stack; + stack.emplace_back(0, world_root.get()); - // 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); - } + // A node without any joints from a parent has a free joint. + // Add all free joints as children of the world body. + for (int i = 1; i < extraction.nodes.size(); ++i) { + if (parent_joints[i].empty()) { + children[0][i] = true; } } - while (!q.empty()) { - pxr::SdfPath current_path = q.front(); - q.pop_front(); + std::vector visited(extraction.nodes.size()); + while (!stack.empty()) { + auto [current_body, parent] = stack.back(); + stack.pop_back(); + visited[current_body] = true; - 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(); - 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* current_node = nullptr; + if (current_body > 0) { + parent->children.push_back(std::move(extraction.nodes[current_body])); + current_node = parent->children.back().get(); + } else { + current_node = world_root.get(); } - 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); + // Process children in reverse to maintain DFS. + for (int i = extraction.nodes.size() - 1; i > 0; --i) { + if (!children[current_body][i]) { + continue; } + if (visited[i]) { + mju_error("Cycle detected in the kinematic tree at node %s", + current_node->body_path.GetString().c_str()); + return nullptr; + } + stack.emplace_back( + i, current_body > 0 ? parent->children.back().get() : parent); } } - // 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()); + for (int i = 1; i < visited.size(); ++i) { + if (!visited[i]) { + mju_error("Cycle detected: Node %s is not reachable from the world.", + extraction.nodes[i]->body_path.GetString().c_str()); return nullptr; } } return world_root; } + } // namespace usd } // namespace mujoco diff --git a/src/experimental/usd/kinematic_tree.h b/src/experimental/usd/kinematic_tree.h index fc051ccc..af704ca3 100644 --- a/src/experimental/usd/kinematic_tree.h +++ b/src/experimental/usd/kinematic_tree.h @@ -27,22 +27,27 @@ 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 { +struct Node { pxr::SdfPath body_path; - pxr::SdfPath joint_path; // Joint connecting this node to its parent. - std::vector> children; + pxr::SdfPath physics_scene; + std::vector actuators; + std::vector joints; + std::vector colliders; + std::vector sites; + std::vector keyframes; + std::vector> children; }; -// Builds a single kinematic tree from a list of joints. +// A kinematic edge represents a joint. +using JointVec = std::vector; + +// Builds a single kinematic tree from a list of directed edges. // 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 BuildKinematicTree( - const std::vector& joints, - const std::vector& all_body_paths, - const pxr::SdfPath& default_prim_path); +std::unique_ptr BuildKinematicTree(const pxr::UsdStageRefPtr stage); } // namespace usd } // namespace mujoco diff --git a/src/experimental/usd/usd_to_mjspec.cc b/src/experimental/usd/usd_to_mjspec.cc index 752d8033..d5e6fdb1 100644 --- a/src/experimental/usd/usd_to_mjspec.cc +++ b/src/experimental/usd/usd_to_mjspec.cc @@ -68,7 +68,9 @@ #include #include #include -namespace { + +namespace mujoco { +namespace usd { using pxr::MjcPhysicsTokens; using pxr::TfToken; @@ -1524,19 +1526,13 @@ using BodyPrimMap = std::map>; // Recursively traverses the kinematic tree, creating bodies, joints, and geoms // 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, - UsdCaches& caches, const BodyPrimMap& body_to_prims) { - mjsBody* current_mj_body; + mjsBody* parent_mj_body, const Node* parent_node, + const Node* current_node, + UsdCaches& caches) { + mjsBody* current_mj_body = nullptr; - if (current_node.body_path.IsEmpty()) { - // This is the world root node. - current_mj_body = mjs_findBody(spec, "world"); - } else { - // This is a regular body. - pxr::UsdPrim current_body_prim = - stage->GetPrimAtPath(current_node.body_path); + if (!current_node->body_path.IsEmpty()) { + // This is *not* the world body. pxr::SdfPath parent_body_path = parent_node ? parent_node->body_path : pxr::SdfPath(); pxr::UsdPrim parent_prim_for_xform = @@ -1544,160 +1540,82 @@ void PopulateSpecFromTree(pxr::UsdStageRefPtr stage, mjSpec* spec, : stage->GetPrimAtPath(parent_body_path); current_mj_body = ParseUsdPhysicsRigidbody( - spec, pxr::UsdPhysicsRigidBodyAPI(current_body_prim), + spec, pxr::UsdPhysicsRigidBodyAPI::Get(stage, current_node->body_path), parent_prim_for_xform, parent_mj_body, caches.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, - caches.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; - } + } else { + current_mj_body = mjs_findBody(spec, "world"); } - // Add geoms/sites/etc. belonging to the current body. - auto it_prims = body_to_prims.find(current_node.body_path); - if (it_prims != body_to_prims.end()) { - pxr::UsdPrim body_prim_for_xform = - current_node.body_path.IsEmpty() - ? stage->GetPseudoRoot() - : stage->GetPrimAtPath(current_node.body_path); - for (const auto& gprim_path : it_prims->second) { - pxr::UsdPrim prim = stage->GetPrimAtPath(gprim_path); - if (prim.HasAPI()) { - ParseUsdPhysicsCollider(spec, pxr::UsdPhysicsCollisionAPI(prim), - body_prim_for_xform, current_mj_body, caches); - } - if (prim.HasAPI()) { - ParseMjcPhysicsSite(spec, pxr::MjcPhysicsSiteAPI(prim), - body_prim_for_xform, current_mj_body, caches.xform_cache); - } + if (!current_node->joints.empty()) { + for (const auto& joint_path : current_node->joints) { + ParseUsdPhysicsJoint(spec, stage->GetPrimAtPath(joint_path), + current_mj_body, caches.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; + } + + pxr::UsdPrim body_prim_for_xform = + current_node->body_path.IsEmpty() + ? stage->GetPseudoRoot() + : stage->GetPrimAtPath(current_node->body_path); + + for (const auto& collider_path : current_node->colliders) { + ParseUsdPhysicsCollider( + spec, pxr::UsdPhysicsCollisionAPI(stage->GetPrimAtPath(collider_path)), + body_prim_for_xform, current_mj_body, caches); + } + + for (const auto& site_path : current_node->sites) { + ParseMjcPhysicsSite(spec, + pxr::MjcPhysicsSiteAPI(stage->GetPrimAtPath(site_path)), + body_prim_for_xform, current_mj_body, caches.xform_cache); } // Recurse through children. - for (const auto& child_node : current_node.children) { - PopulateSpecFromTree(stage, spec, current_mj_body, ¤t_node, - *child_node, caches, body_to_prims); + for (const auto& child_node : current_node->children) { + PopulateSpecFromTree(stage, spec, current_mj_body, current_node, + child_node.get(), caches); } } -} // namespace +} // namespace usd +} // namespace mujoco mjSpec* mj_parseUSDStage(const pxr::UsdStageRefPtr stage) { mjSpec* spec = mj_makeSpec(); - std::vector physics_scenes; + std::unique_ptr root = + mujoco::usd::BuildKinematicTree(stage); // Set of caches to use for all queries when parsing. - UsdCaches caches; + mujoco::usd::UsdCaches caches; - // Search for UsdPhysicsScene type prim, use the first one that has - // the MjcPhysicsSceneAPI applied or the first UsdPhysicsScene otherwise. - std::optional physics_scene; - for (auto prim : stage->Traverse()) { - if (prim.IsA()) { - bool has_mjc_physics_api = prim.HasAPI(); - if (!physics_scene.has_value() || has_mjc_physics_api) { - physics_scene = pxr::UsdPhysicsScene(prim); - // If we've found the first scene with MjcPhysicsSceneAPI, we can stop - // searching. - if (has_mjc_physics_api) { - break; - } - } + // First parse the physics scene. + if (!root->physics_scene.IsEmpty()) { + mujoco::usd::ParseUsdPhysicsScene( + spec, pxr::UsdPhysicsScene::Get(stage, root->physics_scene)); + } + + if (!root->keyframes.empty()) { + for (const auto& keyframe : root->keyframes) { + mujoco::usd::ParseMjcPhysicsKeyframe( + spec, pxr::MjcPhysicsKeyframe::Get(stage, keyframe)); } } - if (physics_scene.has_value()) { - ParseUsdPhysicsScene(spec, *physics_scene); - } - - pxr::SdfPath default_prim_path; - if (stage->GetDefaultPrim().IsValid()) { - default_prim_path = stage->GetDefaultPrim().GetPath(); - } - - // Data Structures - std::vector all_joints; - std::vector all_body_paths_vec; - BodyPrimMap body_to_prims; - - // ========================================================================= - // PASS 1: Collect Bodies, Joints, and Geoms/Sites/etc. - // ========================================================================= - // A single DFS pass to find all bodies, joints, and determine - // which body owns each geom/site/etc. prim. - std::vector owner_stack; - owner_stack.push_back(pxr::SdfPath()); // Start with the world as owner. - - const auto range = pxr::UsdPrimRange::PreAndPostVisit( - stage->GetPseudoRoot(), pxr::UsdTraverseInstanceProxies()); - - for (auto it = range.begin(); it != range.end(); ++it) { - pxr::UsdPrim prim = *it; - - bool is_body = prim.HasAPI(); - bool resets = caches.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() || - prim.HasAPI()) { - body_to_prims[prim_owner].push_back(prim_path); - } - - if (prim.IsA()) { - all_joints.push_back(pxr::UsdPhysicsJoint(prim)); - - it.PruneChildren(); - } else if (prim.IsA()) { - ParseMjcPhysicsKeyframe(spec, pxr::MjcPhysicsKeyframe(prim)); - - it.PruneChildren(); - } else if (prim.IsA()) { - ParseMjcPhysicsActuator(spec, pxr::MjcPhysicsActuator(prim)); - - it.PruneChildren(); + if (!root->actuators.empty()) { + for (const auto& actuator : root->actuators) { + mujoco::usd::ParseMjcPhysicsActuator( + spec, pxr::MjcPhysicsActuator::Get(stage, actuator)); } } - // ========================================================================= - // PASS 2: Build the kinematic tree and populate the mjSpec. - // ========================================================================= - std::unique_ptr kinematic_tree = - mujoco::usd::BuildKinematicTree(all_joints, all_body_paths_vec, - default_prim_path); - - if (kinematic_tree) { - PopulateSpecFromTree(stage, spec, /*parent_mj_body=*/nullptr, - /*parent_node=*/nullptr, *kinematic_tree, caches, - body_to_prims); - } + // Then populate the kinematic tree. + pxr::UsdGeomXformCache xform_cache; + PopulateSpecFromTree(stage, spec, /*parent_mj_body=*/nullptr, + /*parent_node=*/nullptr, root.get(), caches); return spec; }