chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 4 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 <algorithm>
#include <deque>
#include <map>
#include <memory>
#include <set>
#include <utility>
#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 <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/rigidBodyAPI.h>
#include <pxr/usd/usdPhysics/scene.h>
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<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());
struct ExtractedPrims {
std::vector<std::unique_ptr<Node>> nodes;
std::vector<pxr::UsdPhysicsJoint> 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<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;
}
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;
pxr::SdfPath prim_path = prim.GetPath();
Node* current_node = owner_stack.back();
if (is_body) {
auto new_node = std::make_unique<Node>();
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<pxr::UsdPhysicsScene>() && 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<pxr::UsdPhysicsCollisionAPI>()) {
current_node->colliders.push_back(prim.GetPath());
}
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.
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<std::vector<bool>> children(
extraction.nodes.size(), std::vector<bool>(extraction.nodes.size()));
// List of joint prim paths associated with a child for each body.
std::vector<std::vector<pxr::SdfPath>> 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<KinematicNode>();
std::map<pxr::SdfPath, KinematicNode*> 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<pxr::SdfPath> q;
std::vector<std::pair<int, Node*>> 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<bool> 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<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* 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
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.
// 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<std::unique_ptr<KinematicNode>> children;
pxr::SdfPath physics_scene;
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
// `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);
std::unique_ptr<Node> BuildKinematicTree(const pxr::UsdStageRefPtr stage);
} // namespace usd
} // 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/usdShade/material.h>
#include <pxr/usd/usdShade/materialBindingAPI.h>
namespace {
namespace mujoco {
namespace usd {
using pxr::MjcPhysicsTokens;
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
// 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<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);
}
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, &current_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<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.
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<pxr::UsdPhysicsScene> physics_scene;
for (auto prim : stage->Traverse()) {
if (prim.IsA<pxr::UsdPhysicsScene>()) {
bool has_mjc_physics_api = prim.HasAPI<pxr::MjcPhysicsSceneAPI>();
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<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();
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<mujoco::usd::KinematicNode> 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;
}