warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/sim/import_usd.py

@@ -0,0 +1,894 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import re
+
+import numpy as np
+
+import warp as wp
+
+
+def parse_usd(
+    source,
+    builder,
+    default_density=1.0e3,
+    only_load_enabled_rigid_bodies=False,
+    only_load_enabled_joints=True,
+    contact_ke=1e5,
+    contact_kd=250.0,
+    contact_kf=500.0,
+    contact_ka=0.0,
+    contact_mu=0.6,
+    contact_restitution=0.0,
+    contact_thickness=0.0,
+    joint_limit_ke=100.0,
+    joint_limit_kd=10.0,
+    armature=0.0,
+    invert_rotations=False,
+    verbose=False,
+    ignore_paths=None,
+):
+    """
+    Parses a Universal Scene Description (USD) stage containing UsdPhysics schema definitions for rigid-body articulations and adds the bodies, shapes and joints to the given ModelBuilder.
+
+    The USD description has to be either a path (file name or URL), or an existing USD stage instance that implements the `UsdStage <https://openusd.org/dev/api/class_usd_stage.html>`_ interface.
+
+    Args:
+        source (str | pxr.UsdStage): The file path to the USD file, or an existing USD stage instance.
+        builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
+        default_density (float): The default density to use for bodies without a density attribute.
+        only_load_enabled_rigid_bodies (bool): If True, only rigid bodies which do not have `physics:rigidBodyEnabled` set to False are loaded.
+        only_load_enabled_joints (bool): If True, only joints which do not have `physics:jointEnabled` set to False are loaded.
+        contact_ke (float): The default contact stiffness to use, only considered by the Euler integrators.
+        contact_kd (float): The default contact damping to use, only considered by the Euler integrators.
+        contact_kf (float): The default friction stiffness to use, only considered by the Euler integrators.
+        contact_ka (float): The default adhesion distance to use, only considered by the Euler integrators.
+        contact_mu (float): The default friction coefficient to use if a shape has not friction coefficient defined.
+        contact_restitution (float): The default coefficient of restitution to use if a shape has not coefficient of restitution defined.
+        contact_thickness (float): The thickness to add to the shape geometry.
+        joint_limit_ke (float): The default stiffness to use for joint limits, only considered by the Euler integrators.
+        joint_limit_kd (float): The default damping to use for joint limits, only considered by the Euler integrators.
+        armature (float): The armature to use for the bodies.
+        invert_rotations (bool): If True, inverts any rotations defined in the shape transforms.
+        verbose (bool): If True, print additional information about the parsed USD file.
+        ignore_paths (List[str]): A list of regular expressions matching prim paths to ignore.
+
+    Returns:
+        dict: Dictionary with the following entries:
+
+        .. list-table::
+            :widths: 25 75
+
+            * - "fps"
+              - USD stage frames per second
+            * - "duration"
+              - Difference between end time code and start time code of the USD stage
+            * - "up_axis"
+              - Upper-case string of the stage's up axis ("X", "Y", or "Z")
+            * - "path_shape_map"
+              - Mapping from prim path (str) of the UsdGeom to the respective shape index in :class:`ModelBuilder`
+            * - "path_body_map"
+              - Mapping from prim path (str) of a rigid body prim (e.g. that implements the PhysicsRigidBodyAPI) to the respective body index in :class:`ModelBuilder`
+            * - "path_shape_scale"
+              - Mapping from prim path (str) of the UsdGeom to its respective 3D world scale
+            * - "mass_unit"
+              - The stage's Kilograms Per Unit (KGPU) definition (1.0 by default)
+            * - "linear_unit"
+              - The stage's Meters Per Unit (MPU) definition (1.0 by default)
+
+
+    Note:
+        This importer is experimental and only supports a subset of the USD Physics schema. Please report any issues you encounter.
+    """
+    try:
+        from pxr import Usd, UsdGeom, UsdPhysics
+    except ImportError as e:
+        raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).") from e
+
+    if ignore_paths is None:
+        ignore_paths = []
+
+    def get_attribute(prim, name):
+        if "*" in name:
+            regex = name.replace("*", ".*")
+            for attr in prim.GetAttributes():
+                if re.match(regex, attr.GetName()):
+                    return attr
+        else:
+            return prim.GetAttribute(name)
+
+    def has_attribute(prim, name):
+        attr = get_attribute(prim, name)
+        return attr.IsValid() and attr.HasAuthoredValue()
+
+    def parse_float(prim, name, default=None):
+        attr = get_attribute(prim, name)
+        if not attr or not attr.HasAuthoredValue():
+            return default
+        val = attr.Get()
+        if np.isfinite(val):
+            return val
+        return default
+
+    def parse_quat(prim, name, default=None):
+        attr = get_attribute(prim, name)
+        if not attr or not attr.HasAuthoredValue():
+            return default
+        val = attr.Get()
+        if invert_rotations:
+            quat = wp.quat(*val.imaginary, -val.real)
+        else:
+            quat = wp.quat(*val.imaginary, val.real)
+        l = wp.length(quat)
+        if np.isfinite(l) and l > 0.0:
+            return quat
+        return default
+
+    def parse_vec(prim, name, default=None):
+        attr = get_attribute(prim, name)
+        if not attr or not attr.HasAuthoredValue():
+            return default
+        val = attr.Get()
+        if np.isfinite(val).all():
+            return np.array(val, dtype=np.float32)
+        return default
+
+    def parse_generic(prim, name, default=None):
+        attr = get_attribute(prim, name)
+        if not attr or not attr.HasAuthoredValue():
+            return default
+        return attr.Get()
+
+    def str2axis(s: str) -> np.ndarray:
+        axis = np.zeros(3, dtype=np.float32)
+        axis["XYZ".index(s.upper())] = 1.0
+        return axis
+
+    if isinstance(source, str):
+        stage = Usd.Stage.Open(source, Usd.Stage.LoadAll)
+    else:
+        stage = source
+
+    mass_unit = 1.0
+    try:
+        if UsdPhysics.StageHasAuthoredKilogramsPerUnit(stage):
+            mass_unit = UsdPhysics.GetStageKilogramsPerUnit(stage)
+    except Exception as e:
+        if verbose:
+            print(f"Failed to get mass unit: {e}")
+    linear_unit = 1.0
+    try:
+        if UsdGeom.StageHasAuthoredMetersPerUnit(stage):
+            linear_unit = UsdGeom.GetStageMetersPerUnit(stage)
+    except Exception as e:
+        if verbose:
+            print(f"Failed to get linear unit: {e}")
+
+    def parse_xform(prim):
+        xform = UsdGeom.Xform(prim)
+        mat = np.array(xform.GetLocalTransformation(), dtype=np.float32)
+        if invert_rotations:
+            rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].T.flatten()))
+        else:
+            rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].flatten()))
+        pos = mat[3, :3] * linear_unit
+        scale = np.ones(3, dtype=np.float32)
+        for op in xform.GetOrderedXformOps():
+            if op.GetOpType() == UsdGeom.XformOp.TypeScale:
+                scale = np.array(op.Get(), dtype=np.float32)
+        return wp.transform(pos, rot), scale
+
+    def parse_axis(prim, type, joint_data, is_angular, axis=None):
+        # parse joint axis data
+        schemas = prim.GetAppliedSchemas()
+        schemas_str = "".join(schemas)
+        if f"DriveAPI:{type}" not in schemas_str and f"PhysicsLimitAPI:{type}" not in schemas_str:
+            return
+        drive_type = parse_generic(prim, f"drive:{type}:physics:type", "force")
+        if drive_type != "force":
+            print(f"Warning: only force drive type is supported, ignoring drive:{type} for joint {path}")
+            return
+        stiffness = parse_float(prim, f"drive:{type}:physics:stiffness", 0.0)
+        damping = parse_float(prim, f"drive:{type}:physics:damping", 0.0)
+        low = parse_float(prim, f"limit:{type}:physics:low")
+        high = parse_float(prim, f"limit:{type}:physics:high")
+        target_pos = parse_float(prim, f"drive:{type}:physics:targetPosition")
+        target_vel = parse_float(prim, f"drive:{type}:physics:targetVelocity")
+        if is_angular:
+            stiffness *= mass_unit * linear_unit**2
+            stiffness = np.deg2rad(stiffness)
+            damping *= mass_unit * linear_unit**2
+            damping = np.deg2rad(damping)
+            if target_pos is not None:
+                target_pos = np.deg2rad(target_pos)
+            if target_vel is not None:
+                target_vel = np.deg2rad(target_vel)
+            if low is None:
+                low = joint_data["lowerLimit"]
+            else:
+                low = np.deg2rad(low)
+            if high is None:
+                high = joint_data["upperLimit"]
+            else:
+                high = np.deg2rad(high)
+        else:
+            stiffness *= mass_unit
+            damping *= mass_unit
+            if target_pos is not None:
+                target_pos *= linear_unit
+            if target_vel is not None:
+                target_vel *= linear_unit
+            if low is None:
+                low = joint_data["lowerLimit"]
+            else:
+                low *= linear_unit
+            if high is None:
+                high = joint_data["upperLimit"]
+            else:
+                high *= linear_unit
+
+        mode = wp.sim.JOINT_MODE_FORCE
+        if f"DriveAPI:{type}" in schemas_str:
+            if target_vel is not None and target_vel != 0.0:
+                mode = wp.sim.JOINT_MODE_TARGET_VELOCITY
+            else:
+                mode = wp.sim.JOINT_MODE_TARGET_POSITION
+        if low > high:
+            low = (low + high) / 2
+            high = low
+        axis = wp.sim.JointAxis(
+            axis=(axis or joint_data["axis"]),
+            limit_lower=low,
+            limit_upper=high,
+            action=(target_pos or target_vel or (low + high) / 2),
+            target_ke=stiffness,
+            target_kd=damping,
+            mode=mode,
+            limit_ke=joint_limit_ke,
+            limit_kd=joint_limit_kd,
+        )
+        if is_angular:
+            joint_data["angular_axes"].append(axis)
+        else:
+            joint_data["linear_axes"].append(axis)
+
+    axis_str = "Y"
+    try:
+        axis_str = UsdGeom.GetStageUpAxis(stage)
+    except Exception as e:
+        if verbose:
+            print(f"Failed to parse stage up axis: {e}")
+    upaxis = str2axis(axis_str)
+
+    shape_types = {"Cube", "Sphere", "Mesh", "Capsule", "Plane", "Cylinder", "Cone"}
+
+    path_body_map = {}
+    path_shape_map = {}
+    path_shape_scale = {}
+    # maps prim path name to its world transform
+    path_world_poses = {}
+    # transform from body frame to where the actual joint child frame is
+    # so that the link's children will use the right parent tf for the joint
+    prim_joint_xforms = {}
+    path_collision_filters = set()
+    no_collision_shapes = set()
+
+    body_density = {}  # mapping from body ID to defined density
+
+    # first find all joints and materials
+    joint_data = {}  # mapping from path of child link to joint USD settings
+    materials = {}  # mapping from material path to material USD settings
+    joint_parents = set()  # paths of joint parents
+    for prim in stage.Traverse():
+        type_name = str(prim.GetTypeName())
+        path = str(prim.GetPath())
+        # if verbose:
+        #     print(path, type_name)
+        if type_name.endswith("Joint"):
+            # the type name can sometimes be "DistancePhysicsJoint" or "PhysicsDistanceJoint" ...
+            type_name = type_name.replace("Physics", "").replace("Joint", "")
+            child = str(prim.GetRelationship("physics:body1").GetTargets()[0])
+            pos0 = parse_vec(prim, "physics:localPos0", np.zeros(3, dtype=np.float32)) * linear_unit
+            pos1 = parse_vec(prim, "physics:localPos1", np.zeros(3, dtype=np.float32)) * linear_unit
+            rot0 = parse_quat(prim, "physics:localRot0", wp.quat_identity())
+            rot1 = parse_quat(prim, "physics:localRot1", wp.quat_identity())
+            joint_data[child] = {
+                "type": type_name,
+                "name": str(prim.GetName()),
+                "parent_tf": wp.transform(pos0, rot0),
+                "child_tf": wp.transform(pos1, rot1),
+                "enabled": parse_generic(prim, "physics:jointEnabled", True),
+                "collisionEnabled": parse_generic(prim, "physics:collisionEnabled", False),
+                "excludeFromArticulation": parse_generic(prim, "physics:excludeFromArticulation", False),
+                "axis": str2axis(parse_generic(prim, "physics:axis", "X")),
+                "breakForce": parse_float(prim, "physics:breakForce", np.inf),
+                "breakTorque": parse_float(prim, "physics:breakTorque", np.inf),
+                "linear_axes": [],
+                "angular_axes": [],
+            }
+            if only_load_enabled_joints and not joint_data[child]["enabled"]:
+                print("Skipping disabled joint", path)
+                continue
+            # parse joint limits
+            lower = parse_float(prim, "physics:lowerLimit", -np.inf)
+            upper = parse_float(prim, "physics:upperLimit", np.inf)
+            if type_name == "Distance":
+                # if distance is negative the joint is not limited
+                joint_data[child]["lowerLimit"] = parse_float(prim, "physics:minDistance", -1.0) * linear_unit
+                joint_data[child]["upperLimit"] = parse_float(prim, "physics:maxDistance", -1.0) * linear_unit
+            elif type_name == "Prismatic":
+                joint_data[child]["lowerLimit"] = lower * linear_unit
+                joint_data[child]["upperLimit"] = upper * linear_unit
+            else:
+                joint_data[child]["lowerLimit"] = np.deg2rad(lower) if np.isfinite(lower) else lower
+                joint_data[child]["upperLimit"] = np.deg2rad(upper) if np.isfinite(upper) else upper
+
+            if joint_data[child]["lowerLimit"] > joint_data[child]["upperLimit"]:
+                joint_data[child]["lowerLimit"] = (
+                    joint_data[child]["lowerLimit"] + joint_data[child]["upperLimit"]
+                ) / 2
+                joint_data[child]["upperLimit"] = joint_data[child]["lowerLimit"]
+            parents = prim.GetRelationship("physics:body0").GetTargets()
+            if len(parents) > 0:
+                parent_path = str(parents[0])
+                joint_data[child]["parent"] = parent_path
+                joint_parents.add(parent_path)
+            else:
+                joint_data[child]["parent"] = None
+
+            # parse joint drive
+            parse_axis(prim, "angular", joint_data[child], is_angular=True)
+            parse_axis(prim, "rotX", joint_data[child], is_angular=True, axis=(1.0, 0.0, 0.0))
+            parse_axis(prim, "rotY", joint_data[child], is_angular=True, axis=(0.0, 1.0, 0.0))
+            parse_axis(prim, "rotZ", joint_data[child], is_angular=True, axis=(0.0, 0.0, 1.0))
+            parse_axis(prim, "linear", joint_data[child], is_angular=False)
+            parse_axis(prim, "transX", joint_data[child], is_angular=False, axis=(1.0, 0.0, 0.0))
+            parse_axis(prim, "transY", joint_data[child], is_angular=False, axis=(0.0, 1.0, 0.0))
+            parse_axis(prim, "transZ", joint_data[child], is_angular=False, axis=(0.0, 0.0, 1.0))
+
+        elif type_name == "Material":
+            material = {}
+            if has_attribute(prim, "physics:density"):
+                material["density"] = parse_float(prim, "physics:density") * mass_unit  # / (linear_unit**3)
+            if has_attribute(prim, "physics:restitution"):
+                material["restitution"] = parse_float(prim, "physics:restitution", contact_restitution)
+            if has_attribute(prim, "physics:staticFriction"):
+                material["staticFriction"] = parse_float(prim, "physics:staticFriction", contact_mu)
+            if has_attribute(prim, "physics:dynamicFriction"):
+                material["dynamicFriction"] = parse_float(prim, "physics:dynamicFriction", contact_mu)
+            materials[path] = material
+
+        elif type_name == "PhysicsScene":
+            try:
+                scene = UsdPhysics.Scene(prim)
+                g_vec = scene.GetGravityDirectionAttr()
+                g_mag = scene.GetGravityMagnitudeAttr()
+                if g_mag.HasAuthoredValue() and np.isfinite(g_mag.Get()):
+                    builder.gravity = -np.abs(g_mag.Get() * linear_unit)
+                if g_vec.HasAuthoredValue() and np.linalg.norm(g_vec.Get()) > 0.0:
+                    builder.up_vector = np.array(g_vec.Get(), dtype=np.float32)
+                    if np.any(builder.up_vector < 0.0):
+                        builder.up_vector = -builder.up_vector
+                else:
+                    builder.up_vector = upaxis
+            except Exception as e:
+                if verbose:
+                    print(f"Failed to parse physics scene: {e}")
+
+    def parse_prim(prim, incoming_xform, incoming_scale, parent_body: int = -1):
+        nonlocal builder
+        nonlocal joint_data
+        nonlocal path_body_map
+        nonlocal path_shape_map
+        nonlocal path_shape_scale
+        nonlocal path_world_poses
+        nonlocal prim_joint_xforms
+        nonlocal path_collision_filters
+        nonlocal no_collision_shapes
+        nonlocal body_density
+
+        path = str(prim.GetPath())
+        for pattern in ignore_paths:
+            if re.match(pattern, path):
+                return
+
+        type_name = str(prim.GetTypeName())
+        if type_name.endswith("Joint") or type_name.endswith("Light") or type_name.endswith("Material"):
+            return
+        if verbose:
+            print(f"parse_prim {prim.GetPath()} ({type_name})")
+        if type_name == "PhysicsScene":
+            # in case the PhysicsScene has bodies as children...
+            for child in prim.GetChildren():
+                parse_prim(child, incoming_xform, incoming_scale, parent_body)
+
+        schemas = set(prim.GetAppliedSchemas())
+        children_refs = prim.GetChildren()
+
+        prim_joint_xforms[path] = wp.transform()
+
+        local_xform, scale = parse_xform(prim)
+        scale = incoming_scale * scale
+        xform = wp.mul(incoming_xform, local_xform)
+        path_world_poses[path] = xform
+
+        geo_tf = local_xform
+        body_id = parent_body
+        is_rigid_body = "PhysicsRigidBodyAPI" in schemas and parent_body == -1
+        create_rigid_body = is_rigid_body or path in joint_parents
+        if create_rigid_body:
+            body_id = builder.add_body(
+                origin=xform,
+                name=prim.GetName(),
+                armature=armature,
+            )
+            path_body_map[path] = body_id
+            body_density[body_id] = 0.0
+
+            parent_body = body_id
+
+            geo_tf = wp.transform()
+
+            # set up joints between rigid bodies after the children have been added
+            if path in joint_data:
+                joint = joint_data[path]
+
+                joint_params = {
+                    "child": body_id,
+                    "linear_axes": joint["linear_axes"],
+                    "angular_axes": joint["angular_axes"],
+                    "name": joint["name"],
+                    "enabled": joint["enabled"],
+                    "parent_xform": joint["parent_tf"],
+                    "child_xform": joint["child_tf"],
+                    "armature": armature,
+                }
+
+                parent_path = joint["parent"]
+                if parent_path is None:
+                    joint_params["parent"] = -1
+                    parent_tf = wp.transform()
+                else:
+                    joint_params["parent"] = path_body_map[parent_path]
+                    parent_tf = path_world_poses[parent_path]
+
+                # the joint to which we are connected will transform this body already
+                geo_tf = wp.transform()
+
+                if verbose:
+                    print(f"Adding joint {joint['name']} between {joint['parent']} and {path}")
+                    print("  parent_xform", joint["parent_tf"])
+                    print("  child_xform ", joint["child_tf"])
+                    print("  parent_tf   ", parent_tf)
+                    print(f"  geo_tf at {path} = {geo_tf}  (xform was {xform})")
+
+                if joint["type"] == "Revolute":
+                    joint_params["joint_type"] = wp.sim.JOINT_REVOLUTE
+                    if len(joint_params["angular_axes"]) == 0:
+                        joint_params["angular_axes"].append(
+                            wp.sim.JointAxis(
+                                joint["axis"],
+                                limit_lower=joint["lowerLimit"],
+                                limit_upper=joint["upperLimit"],
+                                limit_ke=joint_limit_ke,
+                                limit_kd=joint_limit_kd,
+                            )
+                        )
+                elif joint["type"] == "Prismatic":
+                    joint_params["joint_type"] = wp.sim.JOINT_PRISMATIC
+                    if len(joint_params["linear_axes"]) == 0:
+                        joint_params["linear_axes"].append(
+                            wp.sim.JointAxis(
+                                joint["axis"],
+                                limit_lower=joint["lowerLimit"],
+                                limit_upper=joint["upperLimit"],
+                                limit_ke=joint_limit_ke,
+                                limit_kd=joint_limit_kd,
+                            )
+                        )
+                elif joint["type"] == "Spherical":
+                    joint_params["joint_type"] = wp.sim.JOINT_BALL
+                elif joint["type"] == "Fixed":
+                    joint_params["joint_type"] = wp.sim.JOINT_FIXED
+                elif joint["type"] == "Distance":
+                    joint_params["joint_type"] = wp.sim.JOINT_DISTANCE
+                    # we have to add a dummy linear X axis to define the joint limits
+                    joint_params["linear_axes"].append(
+                        wp.sim.JointAxis(
+                            (1.0, 0.0, 0.0),
+                            limit_lower=joint["lowerLimit"],
+                            limit_upper=joint["upperLimit"],
+                            limit_ke=joint_limit_ke,
+                            limit_kd=joint_limit_kd,
+                        )
+                    )
+                elif joint["type"] == "":
+                    joint_params["joint_type"] = wp.sim.JOINT_D6
+                else:
+                    print(f"Warning: unsupported joint type {joint['type']} for {path}")
+
+                builder.add_joint(**joint_params)
+
+            elif is_rigid_body:
+                builder.add_joint_free(child=body_id)
+                # free joint; we set joint_q/qd, not body_q/qd since eval_fk is used after model creation
+                builder.joint_q[-4:] = xform.q
+                builder.joint_q[-7:-4] = xform.p
+                linear_vel = parse_vec(prim, "physics:velocity", np.zeros(3, dtype=np.float32)) * linear_unit
+                angular_vel = parse_vec(prim, "physics:angularVelocity", np.zeros(3, dtype=np.float32)) * linear_unit
+                builder.joint_qd[-6:-3] = angular_vel
+                builder.joint_qd[-3:] = linear_vel
+
+        if verbose:
+            print(f"added {type_name} body {body_id} ({path}) at {xform}")
+
+        density = None
+
+        material = None
+        if prim.HasRelationship("material:binding:physics"):
+            other_paths = prim.GetRelationship("material:binding:physics").GetTargets()
+            if len(other_paths) > 0:
+                material = materials[str(other_paths[0])]
+        if material is not None:
+            if "density" in material:
+                density = material["density"]
+        if has_attribute(prim, "physics:density"):
+            d = parse_float(prim, "physics:density")
+            density = d * mass_unit  # / (linear_unit**3)
+
+        # assert prim.GetAttribute('orientation').Get() == "rightHanded", "Only right-handed orientations are supported."
+        enabled = parse_generic(prim, "physics:rigidBodyEnabled", True)
+        if only_load_enabled_rigid_bodies and not enabled:
+            if verbose:
+                print("Skipping disabled rigid body", path)
+            return
+        mass = parse_float(prim, "physics:mass")
+        if is_rigid_body:
+            if density is None:
+                density = default_density
+            body_density[body_id] = density
+        elif density is None:
+            if body_id >= 0:
+                density = body_density[body_id]
+            else:
+                density = 0.0
+
+        com = parse_vec(prim, "physics:centerOfMass", np.zeros(3, dtype=np.float32))
+        i_diag = parse_vec(prim, "physics:diagonalInertia", np.zeros(3, dtype=np.float32))
+        i_rot = parse_quat(prim, "physics:principalAxes", wp.quat_identity())
+
+        # parse children
+        if type_name == "Xform":
+            if prim.IsInstance():
+                proto = prim.GetPrototype()
+                for child in proto.GetChildren():
+                    parse_prim(child, xform, scale, parent_body)
+            else:
+                for child in children_refs:
+                    parse_prim(child, xform, scale, parent_body)
+        elif type_name == "Scope":
+            for child in children_refs:
+                parse_prim(child, incoming_xform, incoming_scale, parent_body)
+        elif type_name in shape_types:
+            # parse shapes
+            shape_params = {
+                "ke": contact_ke,
+                "kd": contact_kd,
+                "kf": contact_kf,
+                "ka": contact_ka,
+                "mu": contact_mu,
+                "restitution": contact_restitution,
+            }
+            if material is not None:
+                if "restitution" in material:
+                    shape_params["restitution"] = material["restitution"]
+                if "dynamicFriction" in material:
+                    shape_params["mu"] = material["dynamicFriction"]
+
+            if has_attribute(prim, "doubleSided") and not prim.GetAttribute("doubleSided").Get():
+                print(f"Warning: treating {path} as double-sided because single-sided collisions are not supported.")
+
+            if type_name == "Cube":
+                size = parse_float(prim, "size", 2.0)
+                if has_attribute(prim, "extents"):
+                    extents = parse_vec(prim, "extents") * scale
+                    # TODO position geom at extents center?
+                    # geo_pos = 0.5 * (extents[0] + extents[1])
+                    extents = extents[1] - extents[0]
+                else:
+                    extents = scale * size
+                shape_id = builder.add_shape_box(
+                    body_id,
+                    geo_tf.p,
+                    geo_tf.q,
+                    hx=extents[0] / 2,
+                    hy=extents[1] / 2,
+                    hz=extents[2] / 2,
+                    density=density,
+                    thickness=contact_thickness,
+                    **shape_params,
+                )
+            elif type_name == "Sphere":
+                if not (scale[0] == scale[1] == scale[2]):
+                    print("Warning: Non-uniform scaling of spheres is not supported.")
+                if has_attribute(prim, "extents"):
+                    extents = parse_vec(prim, "extents") * scale
+                    # TODO position geom at extents center?
+                    # geo_pos = 0.5 * (extents[0] + extents[1])
+                    extents = extents[1] - extents[0]
+                    if not (extents[0] == extents[1] == extents[2]):
+                        print("Warning: Non-uniform extents of spheres are not supported.")
+                    radius = extents[0]
+                else:
+                    radius = parse_float(prim, "radius", 1.0) * scale[0]
+                shape_id = builder.add_shape_sphere(
+                    body_id, geo_tf.p, geo_tf.q, radius, density=density, **shape_params
+                )
+            elif type_name == "Plane":
+                normal_str = parse_generic(prim, "axis", "Z").upper()
+                geo_rot = geo_tf.q
+                if normal_str != "Y":
+                    normal = str2axis(normal_str)
+                    c = np.cross(normal, (0.0, 1.0, 0.0))
+                    angle = np.arcsin(np.linalg.norm(c))
+                    axis = c / np.linalg.norm(c)
+                    geo_rot = wp.mul(geo_rot, wp.quat_from_axis_angle(axis, angle))
+                width = parse_float(prim, "width", 0.0) * scale[0]
+                length = parse_float(prim, "length", 0.0) * scale[1]
+                shape_id = builder.add_shape_plane(
+                    body=body_id,
+                    pos=geo_tf.p,
+                    rot=geo_rot,
+                    width=width,
+                    length=length,
+                    thickness=contact_thickness,
+                    **shape_params,
+                )
+            elif type_name == "Capsule":
+                axis_str = parse_generic(prim, "axis", "Z").upper()
+                radius = parse_float(prim, "radius", 0.5) * scale[0]
+                half_height = parse_float(prim, "height", 2.0) / 2 * scale[1]
+                assert not has_attribute(prim, "extents"), "Capsule extents are not supported."
+                shape_id = builder.add_shape_capsule(
+                    body_id,
+                    geo_tf.p,
+                    geo_tf.q,
+                    radius,
+                    half_height,
+                    density=density,
+                    up_axis="XYZ".index(axis_str),
+                    **shape_params,
+                )
+            elif type_name == "Cylinder":
+                axis_str = parse_generic(prim, "axis", "Z").upper()
+                radius = parse_float(prim, "radius", 0.5) * scale[0]
+                half_height = parse_float(prim, "height", 2.0) / 2 * scale[1]
+                assert not has_attribute(prim, "extents"), "Cylinder extents are not supported."
+                shape_id = builder.add_shape_cylinder(
+                    body_id,
+                    geo_tf.p,
+                    geo_tf.q,
+                    radius,
+                    half_height,
+                    density=density,
+                    up_axis="XYZ".index(axis_str),
+                    **shape_params,
+                )
+            elif type_name == "Cone":
+                axis_str = parse_generic(prim, "axis", "Z").upper()
+                radius = parse_float(prim, "radius", 0.5) * scale[0]
+                half_height = parse_float(prim, "height", 2.0) / 2 * scale[1]
+                assert not has_attribute(prim, "extents"), "Cone extents are not supported."
+                shape_id = builder.add_shape_cone(
+                    body_id,
+                    geo_tf.p,
+                    geo_tf.q,
+                    radius,
+                    half_height,
+                    density=density,
+                    up_axis="XYZ".index(axis_str),
+                    **shape_params,
+                )
+            elif type_name == "Mesh":
+                mesh = UsdGeom.Mesh(prim)
+                points = np.array(mesh.GetPointsAttr().Get(), dtype=np.float32)
+                indices = np.array(mesh.GetFaceVertexIndicesAttr().Get(), dtype=np.float32)
+                counts = mesh.GetFaceVertexCountsAttr().Get()
+                faces = []
+                face_id = 0
+                for count in counts:
+                    if count == 3:
+                        faces.append(indices[face_id : face_id + 3])
+                    elif count == 4:
+                        faces.append(indices[face_id : face_id + 3])
+                        faces.append(indices[[face_id, face_id + 2, face_id + 3]])
+                    else:
+                        # assert False, f"Error while parsing USD mesh {path}: encountered polygon with {count} vertices, but only triangles and quads are supported."
+                        continue
+                    face_id += count
+                m = wp.sim.Mesh(points, np.array(faces, dtype=np.int32).flatten())
+                shape_id = builder.add_shape_mesh(
+                    body_id,
+                    geo_tf.p,
+                    geo_tf.q,
+                    scale=scale,
+                    mesh=m,
+                    density=density,
+                    thickness=contact_thickness,
+                    **shape_params,
+                )
+            else:
+                print(f"Warning: Unsupported geometry type {type_name} at {path}.")
+                return
+
+            path_body_map[path] = body_id
+            path_shape_map[path] = shape_id
+            path_shape_scale[path] = scale
+
+            if prim.HasRelationship("physics:filteredPairs"):
+                other_paths = prim.GetRelationship("physics:filteredPairs").GetTargets()
+                for other_path in other_paths:
+                    path_collision_filters.add((path, str(other_path)))
+
+            if "PhysicsCollisionAPI" not in schemas or not parse_generic(prim, "physics:collisionEnabled", True):
+                no_collision_shapes.add(shape_id)
+
+        else:
+            print(f"Warning: encountered unsupported prim type {type_name}")
+
+        # update mass properties of rigid bodies in cases where properties are defined with higher precedence
+        if body_id >= 0:
+            com = parse_vec(prim, "physics:centerOfMass")
+            if com is not None:
+                # overwrite COM
+                builder.body_com[body_id] = com * scale
+
+            if mass is not None and not (is_rigid_body and mass == 0.0):
+                mass_ratio = mass / builder.body_mass[body_id]
+                # mass has precedence over density, so we overwrite the mass computed from density
+                builder.body_mass[body_id] = mass * mass_unit
+                if mass > 0.0:
+                    builder.body_inv_mass[body_id] = 1.0 / builder.body_mass[body_id]
+                else:
+                    builder.body_inv_mass[body_id] = 0.0
+                # update inertia
+                builder.body_inertia[body_id] *= mass_ratio
+                if np.array(builder.body_inertia[body_id]).any():
+                    builder.body_inv_inertia[body_id] = wp.inverse(builder.body_inertia[body_id])
+                else:
+                    builder.body_inv_inertia[body_id] = wp.mat33(*np.zeros((3, 3), dtype=np.float32))
+
+            if np.linalg.norm(i_diag) > 0.0:
+                rot = np.array(wp.quat_to_matrix(i_rot), dtype=np.float32).reshape(3, 3)
+                inertia = rot @ np.diag(i_diag) @ rot.T
+                builder.body_inertia[body_id] = inertia
+                if inertia.any():
+                    builder.body_inv_inertia[body_id] = wp.inverse(wp.mat33(*inertia))
+                else:
+                    builder.body_inv_inertia[body_id] = wp.mat33(*np.zeros((3, 3), dtype=np.float32))
+
+    parse_prim(
+        stage.GetDefaultPrim(), incoming_xform=wp.transform(), incoming_scale=np.ones(3, dtype=np.float32) * linear_unit
+    )
+
+    shape_count = len(builder.shape_geo_type)
+
+    # apply collision filters now that we have added all shapes
+    for path1, path2 in path_collision_filters:
+        shape1 = path_shape_map[path1]
+        shape2 = path_shape_map[path2]
+        builder.shape_collision_filter_pairs.add((shape1, shape2))
+
+    # apply collision filters to all shapes that have no collision
+    for shape_id in no_collision_shapes:
+        for other_shape_id in range(shape_count):
+            if other_shape_id != shape_id:
+                builder.shape_collision_filter_pairs.add((shape_id, other_shape_id))
+
+    # return stage parameters
+    return {
+        "fps": stage.GetFramesPerSecond(),
+        "duration": stage.GetEndTimeCode() - stage.GetStartTimeCode(),
+        "up_axis": UsdGeom.GetStageUpAxis(stage).upper(),
+        "path_shape_map": path_shape_map,
+        "path_body_map": path_body_map,
+        "path_shape_scale": path_shape_scale,
+        "mass_unit": mass_unit,
+        "linear_unit": linear_unit,
+    }
+
+
+def resolve_usd_from_url(url: str, target_folder_name: str = None, export_usda: bool = False):
+    """
+    Downloads a USD file from a URL and resolves all references to other USD files to be downloaded to the given target folder.
+
+    Args:
+        url (str): URL to the USD file.
+        target_folder_name (str): Target folder name. If None, a timestamped folder will be created in the current directory.
+        export_usda (bool): If True, converts each downloaded USD file to USDA and saves the additional USDA file in the target folder with the same base name as the original USD file.
+
+    Returns:
+        str: File path to the downloaded USD file.
+    """
+    import datetime
+    import os
+
+    import requests
+
+    try:
+        from pxr import Usd
+    except ImportError as e:
+        raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).") from e
+
+    response = requests.get(url, allow_redirects=True)
+    if response.status_code != 200:
+        raise RuntimeError(f"Failed to download USD file. Status code: {response.status_code}")
+    file = response.content
+    dot = os.path.extsep
+    base = os.path.basename(url)
+    url_folder = os.path.dirname(url)
+    base_name = dot.join(base.split(dot)[:-1])
+    if target_folder_name is None:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+        target_folder_name = os.path.join(".usd_cache", f"{base_name}_{timestamp}")
+    os.makedirs(target_folder_name, exist_ok=True)
+    target_filename = os.path.join(target_folder_name, base)
+    with open(target_filename, "wb") as f:
+        f.write(file)
+
+    stage = Usd.Stage.Open(target_filename, Usd.Stage.LoadNone)
+    stage_str = stage.GetRootLayer().ExportToString()
+    print(f"Downloaded USD file to {target_filename}.")
+    if export_usda:
+        usda_filename = os.path.join(target_folder_name, base_name + ".usda")
+        with open(usda_filename, "w") as f:
+            f.write(stage_str)
+            print(f"Exported USDA file to {usda_filename}.")
+
+    # parse referenced USD files like `references = @./franka_collisions.usd@`
+    downloaded = set()
+    for match in re.finditer(r"references.=.@(.*?)@", stage_str):
+        refname = match.group(1)
+        if refname.startswith("./"):
+            refname = refname[2:]
+        if refname in downloaded:
+            continue
+        try:
+            response = requests.get(f"{url_folder}/{refname}", allow_redirects=True)
+            if response.status_code != 200:
+                print(f"Failed to download reference {refname}. Status code: {response.status_code}")
+                continue
+            file = response.content
+            refdir = os.path.dirname(refname)
+            if refdir:
+                os.makedirs(os.path.join(target_folder_name, refdir), exist_ok=True)
+            ref_filename = os.path.join(target_folder_name, refname)
+            if not os.path.exists(ref_filename):
+                with open(ref_filename, "wb") as f:
+                    f.write(file)
+            downloaded.add(refname)
+            print(f"Downloaded USD reference {refname} to {ref_filename}.")
+            if export_usda:
+                ref_stage = Usd.Stage.Open(ref_filename, Usd.Stage.LoadNone)
+                ref_stage_str = ref_stage.GetRootLayer().ExportToString()
+                base = os.path.basename(ref_filename)
+                base_name = dot.join(base.split(dot)[:-1])
+                usda_filename = os.path.join(target_folder_name, base_name + ".usda")
+                with open(usda_filename, "w") as f:
+                    f.write(ref_stage_str)
+                    print(f"Exported USDA file to {usda_filename}.")
+        except Exception:
+            print(f"Failed to download {refname}.")
+    return target_filename