warp-lang 1.0.0b2__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/sim/import_mjcf.py

@@ -44,7 +44,7 @@ def parse_mjcf(
     Args:
         mjcf_filename (str): The filename of the MuJoCo file to parse.
         builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
-        xform (wp.transform): The transform to apply to the imported mechanism.
+        xform (:ref:`transform <transform>`): The transform to apply to the imported mechanism.
         density (float): The density of the shapes in kg/m^3 which will be used to calculate the body mass and inertia.
         stiffness (float): The stiffness of the joints.
         damping (float): The damping of the joints.
@@ -150,9 +150,15 @@ def parse_mjcf(
 
     def parse_vec(attrib, key, default):
         if key in attrib:
-            return np.fromstring(attrib[key], sep=" ")
+            out = np.fromstring(attrib[key], sep=" ", dtype=np.float32)
         else:
-            return np.array(default)
+            out = np.array(default, dtype=np.float32)
+
+        length = len(out)
+        if length == 1:
+            return wp.vec(len(default), wp.float32)(out[0], out[0], out[0])
+
+        return wp.vec(length, wp.float32)(out)
 
     def parse_orientation(attrib):
         if "quat" in attrib:
@@ -232,46 +238,51 @@ def parse_mjcf(
         angular_axes = []
         joint_type = None
 
-        joints = body.findall("joint")
-        for i, joint in enumerate(joints):
-            if "joint" in defaults:
-                joint_attrib = merge_attrib(defaults["joint"], joint.attrib)
-            else:
-                joint_attrib = joint.attrib
-
-            # default to hinge if not specified
-            joint_type_str = joint_attrib.get("type", "hinge")
-
-            joint_name.append(joint_attrib["name"])
-            joint_pos.append(parse_vec(joint_attrib, "pos", (0.0, 0.0, 0.0)) * scale)
-            joint_range = parse_vec(joint_attrib, "range", (-3.0, 3.0))
-            joint_armature.append(parse_float(joint_attrib, "armature", armature) * armature_scale)
-
-            if joint_type_str == "free":
-                joint_type = wp.sim.JOINT_FREE
-                break
-            if joint_type_str == "fixed":
-                joint_type = wp.sim.JOINT_FIXED
-                break
-            is_angular = joint_type_str == "hinge"
-            mode = wp.sim.JOINT_MODE_LIMIT
-            if stiffness > 0.0 or "stiffness" in joint_attrib:
-                mode = wp.sim.JOINT_MODE_TARGET_POSITION
-            axis_vec = parse_vec(joint_attrib, "axis", (0.0, 0.0, 0.0))
-            ax = wp.sim.model.JointAxis(
-                axis=axis_vec,
-                limit_lower=(np.deg2rad(joint_range[0]) if is_angular and use_degrees else joint_range[0]),
-                limit_upper=(np.deg2rad(joint_range[1]) if is_angular and use_degrees else joint_range[1]),
-                target_ke=parse_float(joint_attrib, "stiffness", stiffness),
-                target_kd=parse_float(joint_attrib, "damping", damping),
-                limit_ke=limit_ke,
-                limit_kd=limit_kd,
-                mode=mode,
-            )
-            if is_angular:
-                angular_axes.append(ax)
-            else:
-                linear_axes.append(ax)
+        freejoint_tags = body.findall("freejoint")
+        if len(freejoint_tags) > 0:
+            joint_type = wp.sim.JOINT_FREE
+            joint_name.append(freejoint_tags[0].attrib.get("name", f"{body_name}_freejoint"))
+        else:
+            joints = body.findall("joint")
+            for i, joint in enumerate(joints):
+                if "joint" in defaults:
+                    joint_attrib = merge_attrib(defaults["joint"], joint.attrib)
+                else:
+                    joint_attrib = joint.attrib
+
+                # default to hinge if not specified
+                joint_type_str = joint_attrib.get("type", "hinge")
+
+                joint_name.append(joint_attrib["name"])
+                joint_pos.append(parse_vec(joint_attrib, "pos", (0.0, 0.0, 0.0)) * scale)
+                joint_range = parse_vec(joint_attrib, "range", (-3.0, 3.0))
+                joint_armature.append(parse_float(joint_attrib, "armature", armature) * armature_scale)
+
+                if joint_type_str == "free":
+                    joint_type = wp.sim.JOINT_FREE
+                    break
+                if joint_type_str == "fixed":
+                    joint_type = wp.sim.JOINT_FIXED
+                    break
+                is_angular = joint_type_str == "hinge"
+                mode = wp.sim.JOINT_MODE_LIMIT
+                if stiffness > 0.0 or "stiffness" in joint_attrib:
+                    mode = wp.sim.JOINT_MODE_TARGET_POSITION
+                axis_vec = parse_vec(joint_attrib, "axis", (0.0, 0.0, 0.0))
+                ax = wp.sim.model.JointAxis(
+                    axis=axis_vec,
+                    limit_lower=(np.deg2rad(joint_range[0]) if is_angular and use_degrees else joint_range[0]),
+                    limit_upper=(np.deg2rad(joint_range[1]) if is_angular and use_degrees else joint_range[1]),
+                    target_ke=parse_float(joint_attrib, "stiffness", stiffness),
+                    target_kd=parse_float(joint_attrib, "damping", damping),
+                    limit_ke=limit_ke,
+                    limit_kd=limit_kd,
+                    mode=mode,
+                )
+                if is_angular:
+                    angular_axes.append(ax)
+                else:
+                    linear_axes.append(ax)
 
         link = builder.add_body(
             origin=wp.transform(body_pos, body_ori),  # will be evaluated in fk()
@@ -392,19 +403,19 @@ def parse_mjcf(
                 if "fromto" in geom_attrib:
                     geom_fromto = parse_vec(geom_attrib, "fromto", (0.0, 0.0, 0.0, 1.0, 0.0, 0.0))
 
-                    start = geom_fromto[0:3] * scale
-                    end = geom_fromto[3:6] * scale
+                    start = wp.vec3(geom_fromto[0:3]) * scale
+                    end = wp.vec3(geom_fromto[3:6]) * scale
 
                     # compute rotation to align the Warp capsule (along x-axis), with mjcf fromto direction
                     axis = wp.normalize(end - start)
-                    angle = math.acos(np.dot(axis, (0.0, 1.0, 0.0)))
-                    axis = wp.normalize(np.cross(axis, (0.0, 1.0, 0.0)))
+                    angle = math.acos(wp.dot(axis, wp.vec3(0.0, 1.0, 0.0)))
+                    axis = wp.normalize(wp.cross(axis, wp.vec3(0.0, 1.0, 0.0)))
 
                     geom_pos = (start + end) * 0.5
                     geom_rot = wp.quat_from_axis_angle(axis, -angle)
 
                     geom_radius = geom_size[0]
-                    geom_height = np.linalg.norm(end - start) * 0.5
+                    geom_height = wp.length(end - start) * 0.5
                     geom_up_axis = 1
 
                 else:

warp/sim/import_urdf.py

@@ -47,7 +47,7 @@ def parse_urdf(
     Args:
         urdf_filename (str): The filename of the URDF file to parse.
         builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
-        xform (wp.transform): The transform to apply to the root body.
+        xform (:ref:`transform <transform>`): The transform to apply to the root body.
         floating (bool): If True, the root body is a free joint. If False, the root body is connected via a fixed joint to the world, unless a `base_joint` is defined.
         base_joint (Union[str, dict]): The joint by which the root body is connected to the world. This can be either a string defining the joint axes of a D6 joint with comma-separated positional and angular axis names (e.g. "px,py,rz" for a D6 joint with linear axes in x, y and an angular axis in z) or a dict with joint parameters (see :meth:`ModelBuilder.add_joint`).
         density (float): The density of the shapes in kg/m^3 which will be used to calculate the body mass and inertia.
@@ -100,8 +100,8 @@ def parse_urdf(
                 size = [float(x) for x in size.split()]
                 builder.add_shape_box(
                     body=link,
-                    pos=tf.p,
-                    rot=tf.q,
+                    pos=wp.vec3(tf.p),
+                    rot=wp.quat(tf.q),
                     hx=size[0] * 0.5 * scale,
                     hy=size[1] * 0.5 * scale,
                     hz=size[2] * 0.5 * scale,
@@ -117,8 +117,8 @@ def parse_urdf(
             for sphere in geo.findall("sphere"):
                 builder.add_shape_sphere(
                     body=link,
-                    pos=tf.p,
-                    rot=tf.q,
+                    pos=wp.vec3(tf.p),
+                    rot=wp.quat(tf.q),
                     radius=float(sphere.get("radius") or "1") * scale,
                     density=density,
                     ke=shape_ke,
@@ -132,8 +132,8 @@ def parse_urdf(
             for cylinder in geo.findall("cylinder"):
                 builder.add_shape_capsule(
                     body=link,
-                    pos=tf.p,
-                    rot=tf.q,
+                    pos=wp.vec3(tf.p),
+                    rot=wp.quat(tf.q),
                     radius=float(cylinder.get("radius") or "1") * scale,
                     half_height=float(cylinder.get("length") or "1") * 0.5 * scale,
                     density=density,
@@ -180,12 +180,12 @@ def parse_urdf(
                     # multiple meshes are contained in a scene
                     for geom in m.geometry.values():
                         vertices = np.array(geom.vertices, dtype=np.float32) * scaling
-                        faces = np.array(geom.faces, dtype=np.int32)
+                        faces = np.array(geom.faces.flatten(), dtype=np.int32)
                         mesh = Mesh(vertices, faces)
                         builder.add_shape_mesh(
                             body=link,
-                            pos=tf.p,
-                            rot=tf.q,
+                            pos=wp.vec3(tf.p),
+                            rot=wp.quat(tf.q),
                             mesh=mesh,
                             density=density,
                             ke=shape_ke,
@@ -198,7 +198,7 @@ def parse_urdf(
                 else:
                     # a single mesh
                     vertices = np.array(m.vertices, dtype=np.float32) * scaling
-                    faces = np.array(m.faces, dtype=np.int32)
+                    faces = np.array(m.faces.flatten(), dtype=np.int32)
                     mesh = Mesh(vertices, faces)
                     builder.add_shape_mesh(
                         body=link,
@@ -252,22 +252,22 @@ def parse_urdf(
             I_m[2, 0] = I_m[0, 2]
             I_m[2, 1] = I_m[1, 2]
             rot = wp.quat_to_matrix(inertial_frame.q)
-            I_m = rot @ I_m
+            I_m = rot @ wp.mat33(I_m)
             m = float(inertial.find("mass").get("value") or "0")
             builder.body_mass[link] = m
             builder.body_inv_mass[link] = 1.0 / m
             builder.body_com[link] = com
             builder.body_inertia[link] = I_m
-            builder.body_inv_inertia[link] = np.linalg.inv(I_m)
+            builder.body_inv_inertia[link] = wp.inverse(I_m)
         if m == 0.0 and ensure_nonstatic_links:
             # set the mass to something nonzero to ensure the body is dynamic
             m = static_link_mass
             # cube with side length 0.5
-            I_m = np.eye(3) * m / 12.0 * (0.5 * scale) ** 2 * 2.0
+            I_m = wp.mat33(np.eye(3)) * m / 12.0 * (0.5 * scale) ** 2 * 2.0
             builder.body_mass[link] = m
             builder.body_inv_mass[link] = 1.0 / m
             builder.body_inertia[link] = I_m
-            builder.body_inv_inertia[link] = np.linalg.inv(I_m)
+            builder.body_inv_inertia[link] = wp.inverse(I_m)
 
     end_shape_count = len(builder.shape_geo_type)
 

warp/sim/import_usd.py

@@ -10,32 +10,33 @@ import numpy as np
 import re
 
 import warp as wp
-from . import ModelBuilder
 
 
 def parse_usd(
-    filename,
-    builder: ModelBuilder,
+    source,
+    builder,
     default_density=1.0e3,
     only_load_enabled_rigid_bodies=False,
     only_load_enabled_joints=True,
     default_ke=1e5,
     default_kd=250.0,
     default_kf=500.0,
-    default_mu=0.0,
+    default_mu=0.6,
     default_restitution=0.0,
     default_thickness=0.0,
     joint_limit_ke=100.0,
     joint_limit_kd=10.0,
+    invert_rotations=False,
     verbose=False,
     ignore_paths=[],
-    export_usda=False,
 ):
     """
-    Parses a USD file containing UsdPhysics schema definitions for rigid-body articulations and adds the bodies, shapes and joints to the given ModelBuilder.
+    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:
-        filename (str): The file path or URL to the USD file to parse.
+        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.
@@ -48,12 +49,33 @@ def parse_usd(
         default_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 SemiImplicitIntegrator.
         joint_limit_kd (float): The default damping to use for joint limits, only considered by SemiImplicitIntegrator.
+        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.
-        export_usda (bool): If True and the filename is a URL, export the downloaded USD file to a USDA file.
 
     Returns:
-        dict: Dictionary with the following entries: "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 up axis.
+        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.
@@ -61,73 +83,7 @@ def parse_usd(
     try:
         from pxr import Usd, UsdGeom, UsdPhysics
     except ImportError:
-        raise ImportError("Failed to import pxr. Please install USD.")
-
-    if filename.startswith("http://") or filename.startswith("https://"):
-        # download file
-        import requests
-        import os
-        import datetime
-
-        response = requests.get(filename, 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(filename)
-        url_folder = os.path.dirname(filename)
-        base_name = dot.join(base.split(dot)[:-1])
-        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
-        folder_name = os.path.join(".usd_cache", f"{base_name}_{timestamp}")
-        os.makedirs(folder_name, exist_ok=True)
-        target_filename = os.path.join(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(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(folder_name, refdir), exist_ok=True)
-                ref_filename = os.path.join(folder_name, refname)
-                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(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}.")
-
-        filename = target_filename
+        raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).")
 
     def get_attribute(prim, name):
         if "*" in name:
@@ -156,7 +112,10 @@ def parse_usd(
         if not attr or not attr.HasAuthoredValue():
             return default
         val = attr.Get()
-        quat = wp.quat(*val.imaginary, val.real)
+        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
@@ -182,20 +141,31 @@ def parse_usd(
         axis["XYZ".index(s.upper())] = 1.0
         return axis
 
-    stage = Usd.Stage.Open(filename, Usd.Stage.LoadAll)
-    if UsdPhysics.StageHasAuthoredKilogramsPerUnit(stage):
-        mass_unit = UsdPhysics.GetStageKilogramsPerUnit(stage)
-    else:
-        mass_unit = 1.0
-    if UsdGeom.StageHasAuthoredMetersPerUnit(stage):
-        linear_unit = UsdGeom.GetStageMetersPerUnit(stage)
+    if isinstance(source, str):
+        stage = Usd.Stage.Open(source, Usd.Stage.LoadAll)
     else:
-        linear_unit = 1.0
+        stage = source
+
+    mass_unit = 1.0
+    try:
+        if UsdPhysics.StageHasAuthoredKilogramsPerUnit(stage):
+            mass_unit = UsdPhysics.GetStageKilogramsPerUnit(stage)
+    except:
+        pass
+    linear_unit = 1.0
+    try:
+        if UsdGeom.StageHasAuthoredMetersPerUnit(stage):
+            linear_unit = UsdGeom.GetStageMetersPerUnit(stage)
+    except:
+        pass
 
     def parse_xform(prim):
         xform = UsdGeom.Xform(prim)
         mat = np.array(xform.GetLocalTransformation(), dtype=np.float32)
-        rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].flatten()))
+        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():
@@ -276,12 +246,18 @@ def parse_usd(
         else:
             joint_data["linear_axes"].append(axis)
 
-    upaxis = str2axis(UsdGeom.GetStageUpAxis(stage))
+    axis_str = "Y"
+    try:
+        axis_str = UsdGeom.GetStageUpAxis(stage)
+    except:
+        pass
+    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
@@ -376,21 +352,27 @@ def parse_usd(
             materials[path] = material
 
         elif type_name == "PhysicsScene":
-            scene = UsdPhysics.Scene(prim)
-            g_vec = scene.GetGravityDirectionAttr()
-            g_mag = scene.GetGravityMagnitudeAttr()
-            if g_mag.HasAuthoredValue() and np.isfinite(g_mag.Get()):
-                builder.gravity = 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)  # TODO flip sign?
-            else:
-                builder.up_vector = upaxis
+            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:
+                pass
 
     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
@@ -403,13 +385,14 @@ def parse_usd(
                 return
 
         type_name = str(prim.GetTypeName())
-        if (
-            type_name.endswith("Joint")
-            or type_name.endswith("Light")
-            or type_name.endswith("Scene")
-            or type_name.endswith("Material")
-        ):
+        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()
@@ -725,6 +708,7 @@ def parse_usd(
 
             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()
@@ -791,4 +775,92 @@ def parse_usd(
         "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 requests
+    import datetime
+    import os
+
+    try:
+        from pxr import Usd
+    except ImportError:
+        raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).")
+
+    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