chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Simplify parsing from USD stage to kinematic tree for mjSpec population.

Browse Source 
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
This commit is contained in:
Sam Haves
2025-07-17 08:18:59 -07:00
committed by Copybara-Service
parent e23226509c
commit 2e13c64308
3 changed files with 244 additions and 243 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/usd/kinematic_tree.cc
+168 -90
View File
@@ -14,24 +14,36 @@
#include "experimental/usd/kinematic_tree.h" #include "experimental/usd/kinematic_tree.h"
#include <algorithm>
#include <deque>
#include <map> #include <map>
#include <memory> #include <memory>
#include <set>
#include <utility> #include <utility>
#include <vector> #include <vector>
#include <mujoco/experimental/usd/mjcPhysics/actuator.h>
#include <mujoco/experimental/usd/mjcPhysics/keyframe.h>
#include <mujoco/experimental/usd/mjcPhysics/siteAPI.h>
#include <mujoco/mujoco.h> #include <mujoco/mujoco.h>
#include <pxr/usd/sdf/path.h> #include <pxr/usd/sdf/path.h>
#include <pxr/usd/usd/common.h>
#include <pxr/usd/usd/primRange.h>
#include <pxr/usd/usdGeom/xformCache.h>
#include <pxr/usd/usdPhysics/collisionAPI.h>
#include <pxr/usd/usdPhysics/joint.h> #include <pxr/usd/usdPhysics/joint.h>
#include <pxr/usd/usdPhysics/rigidBodyAPI.h>
#include <pxr/usd/usdPhysics/scene.h>
namespace mujoco { namespace mujoco {
namespace usd { namespace usd {
bool GetJointBodies(const pxr::UsdPhysicsJoint& joint, bool GetJointBodies(const pxr::UsdPhysicsJoint& joint, pxr::SdfPath* from,
const pxr::SdfPath& default_prim_path, pxr::SdfPath* from,
pxr::SdfPath* to) { 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; pxr::SdfPathVector body1_paths;
joint.GetBody1Rel().GetTargets(&body1_paths); joint.GetBody1Rel().GetTargets(&body1_paths);
if (body1_paths.empty()) { if (body1_paths.empty()) {
@@ -62,121 +74,187 @@ bool GetJointBodies(const pxr::UsdPhysicsJoint& joint,
return true; return true;
} }
std::unique_ptr<KinematicNode> BuildKinematicTree( struct ExtractedPrims {
const std::vector<pxr::UsdPhysicsJoint>& joints, std::vector<std::unique_ptr<Node>> nodes;
const std::vector<pxr::SdfPath>& all_body_paths, std::vector<pxr::UsdPhysicsJoint> joints;
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) { ExtractedPrims ExtractPrims(pxr::UsdStageRefPtr stage) {
pxr::SdfPath from, to; pxr::UsdPhysicsScene physics_scene;
if (!GetJointBodies(joint, default_prim_path, &from, &to)) { std::vector<std::unique_ptr<Node>> nodes;
std::vector<pxr::UsdPhysicsJoint> joints;
nodes.push_back(std::make_unique<Node>());
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<Node*> 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<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; continue;
} }
auto edge_key = std::make_pair(from, to); pxr::SdfPath prim_path = prim.GetPath();
auto it = edge_to_joint_map.find(edge_key); Node* current_node = owner_stack.back();
if (it == edge_to_joint_map.end()) {
edge_to_joint_map[edge_key] = joint.GetPath(); if (is_body) {
} else { auto new_node = std::make_unique<Node>();
mju_warning( new_node->body_path = prim_path;
"Multiple explicit joints defined between body %s and body %s. " nodes.push_back(std::move(new_node));
"Joint1: %s, Joint2: %s. Keeping the first one found: %s", current_node = nodes.back().get();
(from.IsEmpty() ? "<worldbody>" : from.GetString()).c_str(), } else if (resets) {
to.GetString().c_str(), it->second.GetString().c_str(), current_node = root; // Reset owner to world.
joint.GetPath().GetString().c_str(), it->second.GetString().c_str());
continue;
} }
if (from == to) { if (is_pushed_to_stack) {
mju_error("Self-loop detected at node %s", to.GetString().c_str()); owner_stack.push_back(current_node);
return nullptr;
} }
if (parent_map.count(to)) {
mju_error("Node %s has multiple parents ('%s' and '%s').", if (prim.IsA<pxr::UsdPhysicsScene>() && root->physics_scene.IsEmpty()) {
to.GetString().c_str(), parent_map.at(to).GetString().c_str(), root->physics_scene = prim_path;
from.GetString().c_str());
return nullptr;
} }
children_map[from].push_back(to);
parent_map[to] = from; if (prim.HasAPI<pxr::UsdPhysicsCollisionAPI>()) {
all_nodes.insert(from); current_node->colliders.push_back(prim.GetPath());
all_nodes.insert(to); }
if (prim.HasAPI<pxr::MjcPhysicsSiteAPI>()) {
current_node->sites.push_back(prim.GetPath());
// Sites should not have children.
it.PruneChildren();
}
if (prim.IsA<pxr::UsdPhysicsJoint>()) {
// 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<pxr::MjcPhysicsActuator>()) {
root->actuators.push_back(prim_path);
// Joints should not have children.
it.PruneChildren();
}
if (prim.IsA<pxr::MjcPhysicsKeyframe>()) {
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<Node> BuildKinematicTree(const pxr::UsdStageRefPtr stage) {
ExtractedPrims extraction = ExtractPrims(stage);
std::map<pxr::SdfPath, int> 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. // List of direct children for each body.
for (auto& [_, children] : children_map) { std::vector<std::vector<bool>> children(
std::sort( extraction.nodes.size(), std::vector<bool>(extraction.nodes.size()));
children.begin(), children.end(), // List of joint prim paths associated with a child for each body.
[&v = all_body_paths](const auto& a, const auto& b) { std::vector<std::vector<pxr::SdfPath>> parent_joints(extraction.nodes.size());
return std::distance(v.begin(), std::find(v.begin(), v.end(), a)) < for (const pxr::UsdPhysicsJoint& joint : extraction.joints) {
std::distance(v.begin(), std::find(v.begin(), v.end(), b)); 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. // The world body is represented by an empty SdfPath.
auto world_root = std::make_unique<KinematicNode>(); auto world_root = std::move(extraction.nodes[0]);
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 std::vector<std::pair<int, Node*>> stack;
// to the front to perform a DFS on each root's tree. stack.emplace_back(0, world_root.get());
std::deque<pxr::SdfPath> q;
// Add roots (floating-base bodies and children of the world) to the queue, // A node without any joints from a parent has a free joint.
// preserving the DFS order from the USD stage. // Add all free joints as children of the world body.
for (const auto& body_path : all_body_paths) { for (int i = 1; i < extraction.nodes.size(); ++i) {
if (!body_path.IsEmpty()) { if (parent_joints[i].empty()) {
const auto it = parent_map.find(body_path); children[0][i] = true;
// 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()) { std::vector<bool> visited(extraction.nodes.size());
pxr::SdfPath current_path = q.front(); while (!stack.empty()) {
q.pop_front(); auto [current_body, parent] = stack.back();
stack.pop_back();
visited[current_body] = true;
pxr::SdfPath parent_path = parent_map.count(current_path) Node* current_node = nullptr;
? parent_map.at(current_path) if (current_body > 0) {
: pxr::SdfPath(); parent->children.push_back(std::move(extraction.nodes[current_body]));
KinematicNode* parent_node = node_map.at(parent_path); current_node = parent->children.back().get();
} else {
auto new_node = std::make_unique<KinematicNode>(); current_node = world_root.get();
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)) { // Process children in reverse to maintain DFS.
const auto& children = children_map.at(current_path); for (int i = extraction.nodes.size() - 1; i > 0; --i) {
// Add children to the front of the queue in reverse order to ensure if (!children[current_body][i]) {
// they are processed in the correct order by the DFS. continue;
for (auto it = children.rbegin(); it != children.rend(); ++it) {
q.push_front(*it);
} }
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. for (int i = 1; i < visited.size(); ++i) {
// Unvisited nodes at this point imply a cycle. if (!visited[i]) {
for (const auto& node : all_nodes) { mju_error("Cycle detected: Node %s is not reachable from the world.",
if (!node.IsEmpty() && !node_map.count(node)) { extraction.nodes[i]->body_path.GetString().c_str());
mju_error("Cycle detected involving node %s.", node.GetString().c_str());
return nullptr; return nullptr;
} }
} }
return world_root; return world_root;
} }
} // namespace usd } // namespace usd
} // namespace mujoco } // namespace mujoco
src/experimental/usd/kinematic_tree.h
+13 -8
View File
@@ -27,22 +27,27 @@ namespace usd {
// A struct to represent a node in the kinematic tree. // A struct to represent a node in the kinematic tree.
// Using a struct with a vector of children preserves the order of bodies, // Using a struct with a vector of children preserves the order of bodies,
// which is important for things like keyframes and policy compatibility. // which is important for things like keyframes and policy compatibility.
struct KinematicNode { struct Node {
pxr::SdfPath body_path; pxr::SdfPath body_path;
pxr::SdfPath joint_path; // Joint connecting this node to its parent. pxr::SdfPath physics_scene;
std::vector<std::unique_ptr<KinematicNode>> children; std::vector<pxr::SdfPath> actuators;
std::vector<pxr::SdfPath> joints;
std::vector<pxr::SdfPath> colliders;
std::vector<pxr::SdfPath> sites;
std::vector<pxr::SdfPath> keyframes;
std::vector<std::unique_ptr<Node>> children;
}; };
// Builds a single kinematic tree from a list of joints. // A kinematic edge represents a joint.
using JointVec = std::vector<pxr::UsdPhysicsJoint>;
// 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 // The DFS order of bodies in the tree is determined by the order of bodies in
// `all_body_paths`. // `all_body_paths`.
// All bodies, including static and floating-base bodies, are organized under a // 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. // 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. // Returns the root of the kinematic tree, or `nullptr` for invalid structures.
std::unique_ptr<KinematicNode> BuildKinematicTree( std::unique_ptr<Node> BuildKinematicTree(const pxr::UsdStageRefPtr stage);
const std::vector<pxr::UsdPhysicsJoint>& joints,
const std::vector<pxr::SdfPath>& all_body_paths,
const pxr::SdfPath& default_prim_path);
} // namespace usd } // namespace usd
} // namespace mujoco } // namespace mujoco
src/experimental/usd/usd_to_mjspec.cc
+63 -145
View File
@@ -68,7 +68,9 @@
#include <pxr/usd/usdPhysics/scene.h> #include <pxr/usd/usdPhysics/scene.h>
#include <pxr/usd/usdShade/material.h> #include <pxr/usd/usdShade/material.h>
#include <pxr/usd/usdShade/materialBindingAPI.h> #include <pxr/usd/usdShade/materialBindingAPI.h>
namespace {
namespace mujoco {
namespace usd {
using pxr::MjcPhysicsTokens; using pxr::MjcPhysicsTokens;
using pxr::TfToken; using pxr::TfToken;
@@ -1524,19 +1526,13 @@ using BodyPrimMap = std::map<pxr::SdfPath, std::vector<pxr::SdfPath>>;
// Recursively traverses the kinematic tree, creating bodies, joints, and geoms // Recursively traverses the kinematic tree, creating bodies, joints, and geoms
// in the mjSpec. // in the mjSpec.
void PopulateSpecFromTree(pxr::UsdStageRefPtr stage, mjSpec* spec, void PopulateSpecFromTree(pxr::UsdStageRefPtr stage, mjSpec* spec,
mjsBody* parent_mj_body, mjsBody* parent_mj_body, const Node* parent_node,
const mujoco::usd::KinematicNode* parent_node, const Node* current_node,
const mujoco::usd::KinematicNode& current_node, UsdCaches& caches) {
UsdCaches& caches, const BodyPrimMap& body_to_prims) { mjsBody* current_mj_body = nullptr;
mjsBody* current_mj_body;
if (current_node.body_path.IsEmpty()) { if (!current_node->body_path.IsEmpty()) {
// This is the world root node. // This is *not* the world body.
current_mj_body = mjs_findBody(spec, "world");
} else {
// This is a regular body.
pxr::UsdPrim current_body_prim =
stage->GetPrimAtPath(current_node.body_path);
pxr::SdfPath parent_body_path = pxr::SdfPath parent_body_path =
parent_node ? parent_node->body_path : pxr::SdfPath(); parent_node ? parent_node->body_path : pxr::SdfPath();
pxr::UsdPrim parent_prim_for_xform = pxr::UsdPrim parent_prim_for_xform =
@@ -1544,160 +1540,82 @@ void PopulateSpecFromTree(pxr::UsdStageRefPtr stage, mjSpec* spec,
: stage->GetPrimAtPath(parent_body_path); : stage->GetPrimAtPath(parent_body_path);
current_mj_body = ParseUsdPhysicsRigidbody( 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); parent_prim_for_xform, parent_mj_body, caches.xform_cache);
} else {
if (!current_node.joint_path.IsEmpty()) { current_mj_body = mjs_findBody(spec, "world");
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;
}
} }
// Add geoms/sites/etc. belonging to the current body. if (!current_node->joints.empty()) {
auto it_prims = body_to_prims.find(current_node.body_path); for (const auto& joint_path : current_node->joints) {
if (it_prims != body_to_prims.end()) { ParseUsdPhysicsJoint(spec, stage->GetPrimAtPath(joint_path),
pxr::UsdPrim body_prim_for_xform = current_mj_body, caches.xform_cache);
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<pxr::UsdPhysicsCollisionAPI>()) {
ParseUsdPhysicsCollider(spec, pxr::UsdPhysicsCollisionAPI(prim),
body_prim_for_xform, current_mj_body, caches);
}
if (prim.HasAPI<pxr::MjcPhysicsSiteAPI>()) {
ParseMjcPhysicsSite(spec, pxr::MjcPhysicsSiteAPI(prim),
body_prim_for_xform, 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. // Recurse through children.
for (const auto& child_node : current_node.children) { for (const auto& child_node : current_node->children) {
PopulateSpecFromTree(stage, spec, current_mj_body, &current_node, PopulateSpecFromTree(stage, spec, current_mj_body, current_node,
*child_node, caches, body_to_prims); child_node.get(), caches);
} }
} }
} // namespace } // namespace usd
} // namespace mujoco
mjSpec* mj_parseUSDStage(const pxr::UsdStageRefPtr stage) { mjSpec* mj_parseUSDStage(const pxr::UsdStageRefPtr stage) {
mjSpec* spec = mj_makeSpec(); mjSpec* spec = mj_makeSpec();
std::vector<pxr::UsdPhysicsScene> physics_scenes; std::unique_ptr<mujoco::usd::Node> root =
mujoco::usd::BuildKinematicTree(stage);
// Set of caches to use for all queries when parsing. // 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 // First parse the physics scene.
// the MjcPhysicsSceneAPI applied or the first UsdPhysicsScene otherwise. if (!root->physics_scene.IsEmpty()) {
std::optional<pxr::UsdPhysicsScene> physics_scene; mujoco::usd::ParseUsdPhysicsScene(
for (auto prim : stage->Traverse()) { spec, pxr::UsdPhysicsScene::Get(stage, root->physics_scene));
if (prim.IsA<pxr::UsdPhysicsScene>()) { }
bool has_mjc_physics_api = prim.HasAPI<pxr::MjcPhysicsSceneAPI>();
if (!physics_scene.has_value() || has_mjc_physics_api) { if (!root->keyframes.empty()) {
physics_scene = pxr::UsdPhysicsScene(prim); for (const auto& keyframe : root->keyframes) {
// If we've found the first scene with MjcPhysicsSceneAPI, we can stop mujoco::usd::ParseMjcPhysicsKeyframe(
// searching. spec, pxr::MjcPhysicsKeyframe::Get(stage, keyframe));
if (has_mjc_physics_api) {
break;
}
}
} }
} }
if (physics_scene.has_value()) { if (!root->actuators.empty()) {
ParseUsdPhysicsScene(spec, *physics_scene); for (const auto& actuator : root->actuators) {
} mujoco::usd::ParseMjcPhysicsActuator(
spec, pxr::MjcPhysicsActuator::Get(stage, actuator));
pxr::SdfPath default_prim_path;
if (stage->GetDefaultPrim().IsValid()) {
default_prim_path = stage->GetDefaultPrim().GetPath();
}
// Data Structures
std::vector<pxr::UsdPhysicsJoint> all_joints;
std::vector<pxr::SdfPath> 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<pxr::SdfPath> 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<pxr::UsdPhysicsRigidBodyAPI>();
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<pxr::UsdPhysicsCollisionAPI>() ||
prim.HasAPI<pxr::MjcPhysicsSiteAPI>()) {
body_to_prims[prim_owner].push_back(prim_path);
}
if (prim.IsA<pxr::UsdPhysicsJoint>()) {
all_joints.push_back(pxr::UsdPhysicsJoint(prim));
it.PruneChildren();
} else if (prim.IsA<pxr::MjcPhysicsKeyframe>()) {
ParseMjcPhysicsKeyframe(spec, pxr::MjcPhysicsKeyframe(prim));
it.PruneChildren();
} else if (prim.IsA<pxr::MjcPhysicsActuator>()) {
ParseMjcPhysicsActuator(spec, pxr::MjcPhysicsActuator(prim));
it.PruneChildren();
} }
} }
// ========================================================================= // Then populate the kinematic tree.
// PASS 2: Build the kinematic tree and populate the mjSpec. pxr::UsdGeomXformCache xform_cache;
// ========================================================================= PopulateSpecFromTree(stage, spec, /*parent_mj_body=*/nullptr,
std::unique_ptr<mujoco::usd::KinematicNode> kinematic_tree = /*parent_node=*/nullptr, root.get(), caches);
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);
}
return spec; return spec;
} }