warp-lang 1.0.1__py3-none-manylinux2014_x86_64.whl → 1.1.0__py3-none-manylinux2014_x86_64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of warp-lang might be problematic. Click here for more details.

Files changed (346) hide show
  1. warp/__init__.py +108 -97
  2. warp/__init__.pyi +1 -1
  3. warp/bin/warp-clang.so +0 -0
  4. warp/bin/warp.so +0 -0
  5. warp/build.py +115 -113
  6. warp/build_dll.py +383 -375
  7. warp/builtins.py +3425 -3354
  8. warp/codegen.py +2878 -2792
  9. warp/config.py +40 -36
  10. warp/constants.py +45 -45
  11. warp/context.py +5194 -5102
  12. warp/dlpack.py +442 -442
  13. warp/examples/__init__.py +16 -16
  14. warp/examples/assets/bear.usd +0 -0
  15. warp/examples/assets/bunny.usd +0 -0
  16. warp/examples/assets/cartpole.urdf +110 -110
  17. warp/examples/assets/crazyflie.usd +0 -0
  18. warp/examples/assets/cube.usd +0 -0
  19. warp/examples/assets/nv_ant.xml +92 -92
  20. warp/examples/assets/nv_humanoid.xml +183 -183
  21. warp/examples/assets/quadruped.urdf +267 -267
  22. warp/examples/assets/rocks.nvdb +0 -0
  23. warp/examples/assets/rocks.usd +0 -0
  24. warp/examples/assets/sphere.usd +0 -0
  25. warp/examples/benchmarks/benchmark_api.py +383 -383
  26. warp/examples/benchmarks/benchmark_cloth.py +278 -279
  27. warp/examples/benchmarks/benchmark_cloth_cupy.py +88 -88
  28. warp/examples/benchmarks/benchmark_cloth_jax.py +97 -100
  29. warp/examples/benchmarks/benchmark_cloth_numba.py +146 -142
  30. warp/examples/benchmarks/benchmark_cloth_numpy.py +77 -77
  31. warp/examples/benchmarks/benchmark_cloth_pytorch.py +86 -86
  32. warp/examples/benchmarks/benchmark_cloth_taichi.py +112 -112
  33. warp/examples/benchmarks/benchmark_cloth_warp.py +146 -146
  34. warp/examples/benchmarks/benchmark_launches.py +295 -295
  35. warp/examples/browse.py +29 -28
  36. warp/examples/core/example_dem.py +234 -221
  37. warp/examples/core/example_fluid.py +293 -267
  38. warp/examples/core/example_graph_capture.py +144 -129
  39. warp/examples/core/example_marching_cubes.py +188 -176
  40. warp/examples/core/example_mesh.py +174 -154
  41. warp/examples/core/example_mesh_intersect.py +205 -193
  42. warp/examples/core/example_nvdb.py +176 -169
  43. warp/examples/core/example_raycast.py +105 -89
  44. warp/examples/core/example_raymarch.py +199 -178
  45. warp/examples/core/example_render_opengl.py +185 -141
  46. warp/examples/core/example_sph.py +405 -389
  47. warp/examples/core/example_torch.py +222 -181
  48. warp/examples/core/example_wave.py +263 -249
  49. warp/examples/fem/bsr_utils.py +378 -380
  50. warp/examples/fem/example_apic_fluid.py +407 -391
  51. warp/examples/fem/example_convection_diffusion.py +182 -168
  52. warp/examples/fem/example_convection_diffusion_dg.py +219 -209
  53. warp/examples/fem/example_convection_diffusion_dg0.py +204 -194
  54. warp/examples/fem/example_deformed_geometry.py +177 -159
  55. warp/examples/fem/example_diffusion.py +201 -173
  56. warp/examples/fem/example_diffusion_3d.py +177 -152
  57. warp/examples/fem/example_diffusion_mgpu.py +221 -214
  58. warp/examples/fem/example_mixed_elasticity.py +244 -222
  59. warp/examples/fem/example_navier_stokes.py +259 -243
  60. warp/examples/fem/example_stokes.py +220 -192
  61. warp/examples/fem/example_stokes_transfer.py +265 -249
  62. warp/examples/fem/mesh_utils.py +133 -109
  63. warp/examples/fem/plot_utils.py +292 -287
  64. warp/examples/optim/example_bounce.py +260 -248
  65. warp/examples/optim/example_cloth_throw.py +222 -210
  66. warp/examples/optim/example_diffray.py +566 -535
  67. warp/examples/optim/example_drone.py +864 -835
  68. warp/examples/optim/example_inverse_kinematics.py +176 -169
  69. warp/examples/optim/example_inverse_kinematics_torch.py +185 -170
  70. warp/examples/optim/example_spring_cage.py +239 -234
  71. warp/examples/optim/example_trajectory.py +223 -201
  72. warp/examples/optim/example_walker.py +306 -292
  73. warp/examples/sim/example_cartpole.py +139 -128
  74. warp/examples/sim/example_cloth.py +196 -184
  75. warp/examples/sim/example_granular.py +124 -113
  76. warp/examples/sim/example_granular_collision_sdf.py +197 -185
  77. warp/examples/sim/example_jacobian_ik.py +236 -213
  78. warp/examples/sim/example_particle_chain.py +118 -106
  79. warp/examples/sim/example_quadruped.py +193 -179
  80. warp/examples/sim/example_rigid_chain.py +197 -189
  81. warp/examples/sim/example_rigid_contact.py +189 -176
  82. warp/examples/sim/example_rigid_force.py +127 -126
  83. warp/examples/sim/example_rigid_gyroscopic.py +109 -97
  84. warp/examples/sim/example_rigid_soft_contact.py +134 -124
  85. warp/examples/sim/example_soft_body.py +190 -178
  86. warp/fabric.py +337 -335
  87. warp/fem/__init__.py +60 -27
  88. warp/fem/cache.py +401 -388
  89. warp/fem/dirichlet.py +178 -179
  90. warp/fem/domain.py +262 -263
  91. warp/fem/field/__init__.py +100 -101
  92. warp/fem/field/field.py +148 -149
  93. warp/fem/field/nodal_field.py +298 -299
  94. warp/fem/field/restriction.py +22 -21
  95. warp/fem/field/test.py +180 -181
  96. warp/fem/field/trial.py +183 -183
  97. warp/fem/geometry/__init__.py +15 -19
  98. warp/fem/geometry/closest_point.py +69 -70
  99. warp/fem/geometry/deformed_geometry.py +270 -271
  100. warp/fem/geometry/element.py +744 -744
  101. warp/fem/geometry/geometry.py +184 -186
  102. warp/fem/geometry/grid_2d.py +380 -373
  103. warp/fem/geometry/grid_3d.py +441 -435
  104. warp/fem/geometry/hexmesh.py +953 -953
  105. warp/fem/geometry/partition.py +374 -376
  106. warp/fem/geometry/quadmesh_2d.py +532 -532
  107. warp/fem/geometry/tetmesh.py +840 -840
  108. warp/fem/geometry/trimesh_2d.py +577 -577
  109. warp/fem/integrate.py +1630 -1615
  110. warp/fem/operator.py +190 -191
  111. warp/fem/polynomial.py +214 -213
  112. warp/fem/quadrature/__init__.py +2 -2
  113. warp/fem/quadrature/pic_quadrature.py +243 -245
  114. warp/fem/quadrature/quadrature.py +295 -294
  115. warp/fem/space/__init__.py +294 -292
  116. warp/fem/space/basis_space.py +488 -489
  117. warp/fem/space/collocated_function_space.py +100 -105
  118. warp/fem/space/dof_mapper.py +236 -236
  119. warp/fem/space/function_space.py +148 -145
  120. warp/fem/space/grid_2d_function_space.py +267 -267
  121. warp/fem/space/grid_3d_function_space.py +305 -306
  122. warp/fem/space/hexmesh_function_space.py +350 -352
  123. warp/fem/space/partition.py +350 -350
  124. warp/fem/space/quadmesh_2d_function_space.py +368 -369
  125. warp/fem/space/restriction.py +158 -160
  126. warp/fem/space/shape/__init__.py +13 -15
  127. warp/fem/space/shape/cube_shape_function.py +738 -738
  128. warp/fem/space/shape/shape_function.py +102 -103
  129. warp/fem/space/shape/square_shape_function.py +611 -611
  130. warp/fem/space/shape/tet_shape_function.py +565 -567
  131. warp/fem/space/shape/triangle_shape_function.py +429 -429
  132. warp/fem/space/tetmesh_function_space.py +294 -292
  133. warp/fem/space/topology.py +297 -295
  134. warp/fem/space/trimesh_2d_function_space.py +223 -221
  135. warp/fem/types.py +77 -77
  136. warp/fem/utils.py +495 -495
  137. warp/jax.py +166 -141
  138. warp/jax_experimental.py +341 -339
  139. warp/native/array.h +1072 -1025
  140. warp/native/builtin.h +1560 -1560
  141. warp/native/bvh.cpp +398 -398
  142. warp/native/bvh.cu +525 -525
  143. warp/native/bvh.h +429 -429
  144. warp/native/clang/clang.cpp +495 -464
  145. warp/native/crt.cpp +31 -31
  146. warp/native/crt.h +334 -334
  147. warp/native/cuda_crt.h +1049 -1049
  148. warp/native/cuda_util.cpp +549 -540
  149. warp/native/cuda_util.h +288 -203
  150. warp/native/cutlass_gemm.cpp +34 -34
  151. warp/native/cutlass_gemm.cu +372 -372
  152. warp/native/error.cpp +66 -66
  153. warp/native/error.h +27 -27
  154. warp/native/fabric.h +228 -228
  155. warp/native/hashgrid.cpp +301 -278
  156. warp/native/hashgrid.cu +78 -77
  157. warp/native/hashgrid.h +227 -227
  158. warp/native/initializer_array.h +32 -32
  159. warp/native/intersect.h +1204 -1204
  160. warp/native/intersect_adj.h +365 -365
  161. warp/native/intersect_tri.h +322 -322
  162. warp/native/marching.cpp +2 -2
  163. warp/native/marching.cu +497 -497
  164. warp/native/marching.h +2 -2
  165. warp/native/mat.h +1498 -1498
  166. warp/native/matnn.h +333 -333
  167. warp/native/mesh.cpp +203 -203
  168. warp/native/mesh.cu +293 -293
  169. warp/native/mesh.h +1887 -1887
  170. warp/native/nanovdb/NanoVDB.h +4782 -4782
  171. warp/native/nanovdb/PNanoVDB.h +2553 -2553
  172. warp/native/nanovdb/PNanoVDBWrite.h +294 -294
  173. warp/native/noise.h +850 -850
  174. warp/native/quat.h +1084 -1084
  175. warp/native/rand.h +299 -299
  176. warp/native/range.h +108 -108
  177. warp/native/reduce.cpp +156 -156
  178. warp/native/reduce.cu +348 -348
  179. warp/native/runlength_encode.cpp +61 -61
  180. warp/native/runlength_encode.cu +46 -46
  181. warp/native/scan.cpp +30 -30
  182. warp/native/scan.cu +36 -36
  183. warp/native/scan.h +7 -7
  184. warp/native/solid_angle.h +442 -442
  185. warp/native/sort.cpp +94 -94
  186. warp/native/sort.cu +97 -97
  187. warp/native/sort.h +14 -14
  188. warp/native/sparse.cpp +337 -337
  189. warp/native/sparse.cu +544 -544
  190. warp/native/spatial.h +630 -630
  191. warp/native/svd.h +562 -562
  192. warp/native/temp_buffer.h +30 -30
  193. warp/native/vec.h +1132 -1132
  194. warp/native/volume.cpp +297 -297
  195. warp/native/volume.cu +32 -32
  196. warp/native/volume.h +538 -538
  197. warp/native/volume_builder.cu +425 -425
  198. warp/native/volume_builder.h +19 -19
  199. warp/native/warp.cpp +1057 -1052
  200. warp/native/warp.cu +2943 -2828
  201. warp/native/warp.h +313 -305
  202. warp/optim/__init__.py +9 -9
  203. warp/optim/adam.py +120 -120
  204. warp/optim/linear.py +1104 -939
  205. warp/optim/sgd.py +104 -92
  206. warp/render/__init__.py +10 -10
  207. warp/render/render_opengl.py +3217 -3204
  208. warp/render/render_usd.py +768 -749
  209. warp/render/utils.py +152 -150
  210. warp/sim/__init__.py +52 -59
  211. warp/sim/articulation.py +685 -685
  212. warp/sim/collide.py +1594 -1590
  213. warp/sim/import_mjcf.py +489 -481
  214. warp/sim/import_snu.py +220 -221
  215. warp/sim/import_urdf.py +536 -516
  216. warp/sim/import_usd.py +887 -881
  217. warp/sim/inertia.py +316 -317
  218. warp/sim/integrator.py +234 -233
  219. warp/sim/integrator_euler.py +1956 -1956
  220. warp/sim/integrator_featherstone.py +1910 -1991
  221. warp/sim/integrator_xpbd.py +3294 -3312
  222. warp/sim/model.py +4473 -4314
  223. warp/sim/particles.py +113 -112
  224. warp/sim/render.py +417 -403
  225. warp/sim/utils.py +413 -410
  226. warp/sparse.py +1227 -1227
  227. warp/stubs.py +2109 -2469
  228. warp/tape.py +1162 -225
  229. warp/tests/__init__.py +1 -1
  230. warp/tests/__main__.py +4 -4
  231. warp/tests/assets/torus.usda +105 -105
  232. warp/tests/aux_test_class_kernel.py +26 -26
  233. warp/tests/aux_test_compile_consts_dummy.py +10 -10
  234. warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -21
  235. warp/tests/aux_test_dependent.py +22 -22
  236. warp/tests/aux_test_grad_customs.py +23 -23
  237. warp/tests/aux_test_reference.py +11 -11
  238. warp/tests/aux_test_reference_reference.py +10 -10
  239. warp/tests/aux_test_square.py +17 -17
  240. warp/tests/aux_test_unresolved_func.py +14 -14
  241. warp/tests/aux_test_unresolved_symbol.py +14 -14
  242. warp/tests/disabled_kinematics.py +239 -239
  243. warp/tests/run_coverage_serial.py +31 -31
  244. warp/tests/test_adam.py +157 -157
  245. warp/tests/test_arithmetic.py +1124 -1124
  246. warp/tests/test_array.py +2417 -2326
  247. warp/tests/test_array_reduce.py +150 -150
  248. warp/tests/test_async.py +668 -656
  249. warp/tests/test_atomic.py +141 -141
  250. warp/tests/test_bool.py +204 -149
  251. warp/tests/test_builtins_resolution.py +1292 -1292
  252. warp/tests/test_bvh.py +164 -171
  253. warp/tests/test_closest_point_edge_edge.py +228 -228
  254. warp/tests/test_codegen.py +566 -553
  255. warp/tests/test_compile_consts.py +97 -101
  256. warp/tests/test_conditional.py +246 -246
  257. warp/tests/test_copy.py +232 -215
  258. warp/tests/test_ctypes.py +632 -632
  259. warp/tests/test_dense.py +67 -67
  260. warp/tests/test_devices.py +91 -98
  261. warp/tests/test_dlpack.py +530 -529
  262. warp/tests/test_examples.py +400 -378
  263. warp/tests/test_fabricarray.py +955 -955
  264. warp/tests/test_fast_math.py +62 -54
  265. warp/tests/test_fem.py +1277 -1278
  266. warp/tests/test_fp16.py +130 -130
  267. warp/tests/test_func.py +338 -337
  268. warp/tests/test_generics.py +571 -571
  269. warp/tests/test_grad.py +746 -640
  270. warp/tests/test_grad_customs.py +333 -336
  271. warp/tests/test_hash_grid.py +210 -164
  272. warp/tests/test_import.py +39 -39
  273. warp/tests/test_indexedarray.py +1134 -1134
  274. warp/tests/test_intersect.py +67 -67
  275. warp/tests/test_jax.py +307 -307
  276. warp/tests/test_large.py +167 -164
  277. warp/tests/test_launch.py +354 -354
  278. warp/tests/test_lerp.py +261 -261
  279. warp/tests/test_linear_solvers.py +191 -171
  280. warp/tests/test_lvalue.py +421 -493
  281. warp/tests/test_marching_cubes.py +65 -65
  282. warp/tests/test_mat.py +1801 -1827
  283. warp/tests/test_mat_lite.py +115 -115
  284. warp/tests/test_mat_scalar_ops.py +2907 -2889
  285. warp/tests/test_math.py +126 -193
  286. warp/tests/test_matmul.py +500 -499
  287. warp/tests/test_matmul_lite.py +410 -410
  288. warp/tests/test_mempool.py +188 -190
  289. warp/tests/test_mesh.py +284 -324
  290. warp/tests/test_mesh_query_aabb.py +228 -241
  291. warp/tests/test_mesh_query_point.py +692 -702
  292. warp/tests/test_mesh_query_ray.py +292 -303
  293. warp/tests/test_mlp.py +276 -276
  294. warp/tests/test_model.py +110 -110
  295. warp/tests/test_modules_lite.py +39 -39
  296. warp/tests/test_multigpu.py +163 -163
  297. warp/tests/test_noise.py +248 -248
  298. warp/tests/test_operators.py +250 -250
  299. warp/tests/test_options.py +123 -125
  300. warp/tests/test_peer.py +133 -137
  301. warp/tests/test_pinned.py +78 -78
  302. warp/tests/test_print.py +54 -54
  303. warp/tests/test_quat.py +2086 -2086
  304. warp/tests/test_rand.py +288 -288
  305. warp/tests/test_reload.py +217 -217
  306. warp/tests/test_rounding.py +179 -179
  307. warp/tests/test_runlength_encode.py +190 -190
  308. warp/tests/test_sim_grad.py +243 -0
  309. warp/tests/test_sim_kinematics.py +91 -97
  310. warp/tests/test_smoothstep.py +168 -168
  311. warp/tests/test_snippet.py +305 -266
  312. warp/tests/test_sparse.py +468 -460
  313. warp/tests/test_spatial.py +2148 -2148
  314. warp/tests/test_streams.py +486 -473
  315. warp/tests/test_struct.py +710 -675
  316. warp/tests/test_tape.py +173 -148
  317. warp/tests/test_torch.py +743 -743
  318. warp/tests/test_transient_module.py +87 -87
  319. warp/tests/test_types.py +556 -659
  320. warp/tests/test_utils.py +490 -499
  321. warp/tests/test_vec.py +1264 -1268
  322. warp/tests/test_vec_lite.py +73 -73
  323. warp/tests/test_vec_scalar_ops.py +2099 -2099
  324. warp/tests/test_verify_fp.py +94 -94
  325. warp/tests/test_volume.py +737 -736
  326. warp/tests/test_volume_write.py +255 -265
  327. warp/tests/unittest_serial.py +37 -37
  328. warp/tests/unittest_suites.py +363 -359
  329. warp/tests/unittest_utils.py +603 -578
  330. warp/tests/unused_test_misc.py +71 -71
  331. warp/tests/walkthrough_debug.py +85 -85
  332. warp/thirdparty/appdirs.py +598 -598
  333. warp/thirdparty/dlpack.py +143 -143
  334. warp/thirdparty/unittest_parallel.py +566 -561
  335. warp/torch.py +321 -295
  336. warp/types.py +4504 -4450
  337. warp/utils.py +1008 -821
  338. {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/LICENSE.md +126 -126
  339. {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/METADATA +338 -400
  340. warp_lang-1.1.0.dist-info/RECORD +352 -0
  341. warp/examples/assets/cube.usda +0 -42
  342. warp/examples/assets/sphere.usda +0 -56
  343. warp/examples/assets/torus.usda +0 -105
  344. warp_lang-1.0.1.dist-info/RECORD +0 -352
  345. {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/WHEEL +0 -0
  346. {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/top_level.txt +0 -0
warp/sim/import_usd.py CHANGED
@@ -1,881 +1,887 @@
1
- # Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved.
2
- # NVIDIA CORPORATION and its licensors retain all intellectual property
3
- # and proprietary rights in and to this software, related documentation
4
- # and any modifications thereto. Any use, reproduction, disclosure or
5
- # distribution of this software and related documentation without an express
6
- # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
-
8
-
9
- import numpy as np
10
- import re
11
-
12
- import warp as wp
13
-
14
-
15
- def parse_usd(
16
- source,
17
- builder,
18
- default_density=1.0e3,
19
- only_load_enabled_rigid_bodies=False,
20
- only_load_enabled_joints=True,
21
- contact_ke=1e5, # TODO rename to contact...
22
- contact_kd=250.0,
23
- contact_kf=500.0,
24
- contact_ka=0.0,
25
- contact_mu=0.6,
26
- contact_restitution=0.0,
27
- contact_thickness=0.0,
28
- joint_limit_ke=100.0,
29
- joint_limit_kd=10.0,
30
- armature=0.0,
31
- invert_rotations=False,
32
- verbose=False,
33
- ignore_paths=[],
34
- ):
35
- """
36
- 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.
37
-
38
- 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.
39
-
40
- Args:
41
- source (str | pxr.UsdStage): The file path to the USD file, or an existing USD stage instance.
42
- builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
43
- default_density (float): The default density to use for bodies without a density attribute.
44
- only_load_enabled_rigid_bodies (bool): If True, only rigid bodies which do not have `physics:rigidBodyEnabled` set to False are loaded.
45
- only_load_enabled_joints (bool): If True, only joints which do not have `physics:jointEnabled` set to False are loaded.
46
- contact_ke (float): The default contact stiffness to use, only considered by the Euler integrators.
47
- contact_kd (float): The default contact damping to use, only considered by the Euler integrators.
48
- contact_kf (float): The default friction stiffness to use, only considered by the Euler integrators.
49
- contact_ka (float): The default adhesion distance to use, only considered by the Euler integrators.
50
- contact_mu (float): The default friction coefficient to use if a shape has not friction coefficient defined.
51
- contact_restitution (float): The default coefficient of restitution to use if a shape has not coefficient of restitution defined.
52
- contact_thickness (float): The thickness to add to the shape geometry.
53
- joint_limit_ke (float): The default stiffness to use for joint limits, only considered by the Euler integrators.
54
- joint_limit_kd (float): The default damping to use for joint limits, only considered by the Euler integrators.
55
- armature (float): The armature to use for the bodies.
56
- invert_rotations (bool): If True, inverts any rotations defined in the shape transforms.
57
- verbose (bool): If True, print additional information about the parsed USD file.
58
- ignore_paths (List[str]): A list of regular expressions matching prim paths to ignore.
59
-
60
- Returns:
61
- dict: Dictionary with the following entries:
62
-
63
- .. list-table::
64
- :widths: 25 75
65
-
66
- * - "fps"
67
- - USD stage frames per second
68
- * - "duration"
69
- - Difference between end time code and start time code of the USD stage
70
- * - "up_axis"
71
- - Upper-case string of the stage's up axis ("X", "Y", or "Z")
72
- * - "path_shape_map"
73
- - Mapping from prim path (str) of the UsdGeom to the respective shape index in :class:`ModelBuilder`
74
- * - "path_body_map"
75
- - Mapping from prim path (str) of a rigid body prim (e.g. that implements the PhysicsRigidBodyAPI) to the respective body index in :class:`ModelBuilder`
76
- * - "path_shape_scale"
77
- - Mapping from prim path (str) of the UsdGeom to its respective 3D world scale
78
- * - "mass_unit"
79
- - The stage's Kilograms Per Unit (KGPU) definition (1.0 by default)
80
- * - "linear_unit"
81
- - The stage's Meters Per Unit (MPU) definition (1.0 by default)
82
-
83
-
84
- Note:
85
- This importer is experimental and only supports a subset of the USD Physics schema. Please report any issues you encounter.
86
- """
87
- try:
88
- from pxr import Usd, UsdGeom, UsdPhysics
89
- except ImportError:
90
- raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).")
91
-
92
- def get_attribute(prim, name):
93
- if "*" in name:
94
- regex = name.replace("*", ".*")
95
- for attr in prim.GetAttributes():
96
- if re.match(regex, attr.GetName()):
97
- return attr
98
- else:
99
- return prim.GetAttribute(name)
100
-
101
- def has_attribute(prim, name):
102
- attr = get_attribute(prim, name)
103
- return attr.IsValid() and attr.HasAuthoredValue()
104
-
105
- def parse_float(prim, name, default=None):
106
- attr = get_attribute(prim, name)
107
- if not attr or not attr.HasAuthoredValue():
108
- return default
109
- val = attr.Get()
110
- if np.isfinite(val):
111
- return val
112
- return default
113
-
114
- def parse_quat(prim, name, default=None):
115
- attr = get_attribute(prim, name)
116
- if not attr or not attr.HasAuthoredValue():
117
- return default
118
- val = attr.Get()
119
- if invert_rotations:
120
- quat = wp.quat(*val.imaginary, -val.real)
121
- else:
122
- quat = wp.quat(*val.imaginary, val.real)
123
- l = wp.length(quat)
124
- if np.isfinite(l) and l > 0.0:
125
- return quat
126
- return default
127
-
128
- def parse_vec(prim, name, default=None):
129
- attr = get_attribute(prim, name)
130
- if not attr or not attr.HasAuthoredValue():
131
- return default
132
- val = attr.Get()
133
- if np.isfinite(val).all():
134
- return np.array(val, dtype=np.float32)
135
- return default
136
-
137
- def parse_generic(prim, name, default=None):
138
- attr = get_attribute(prim, name)
139
- if not attr or not attr.HasAuthoredValue():
140
- return default
141
- return attr.Get()
142
-
143
- def str2axis(s: str) -> np.ndarray:
144
- axis = np.zeros(3, dtype=np.float32)
145
- axis["XYZ".index(s.upper())] = 1.0
146
- return axis
147
-
148
- if isinstance(source, str):
149
- stage = Usd.Stage.Open(source, Usd.Stage.LoadAll)
150
- else:
151
- stage = source
152
-
153
- mass_unit = 1.0
154
- try:
155
- if UsdPhysics.StageHasAuthoredKilogramsPerUnit(stage):
156
- mass_unit = UsdPhysics.GetStageKilogramsPerUnit(stage)
157
- except Exception as e:
158
- if verbose:
159
- print(f"Failed to get mass unit: {e}")
160
- linear_unit = 1.0
161
- try:
162
- if UsdGeom.StageHasAuthoredMetersPerUnit(stage):
163
- linear_unit = UsdGeom.GetStageMetersPerUnit(stage)
164
- except Exception as e:
165
- if verbose:
166
- print(f"Failed to get linear unit: {e}")
167
-
168
- def parse_xform(prim):
169
- xform = UsdGeom.Xform(prim)
170
- mat = np.array(xform.GetLocalTransformation(), dtype=np.float32)
171
- if invert_rotations:
172
- rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].T.flatten()))
173
- else:
174
- rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].flatten()))
175
- pos = mat[3, :3] * linear_unit
176
- scale = np.ones(3, dtype=np.float32)
177
- for op in xform.GetOrderedXformOps():
178
- if op.GetOpType() == UsdGeom.XformOp.TypeScale:
179
- scale = np.array(op.Get(), dtype=np.float32)
180
- return wp.transform(pos, rot), scale
181
-
182
- def parse_axis(prim, type, joint_data, is_angular, axis=None):
183
- # parse joint axis data
184
- schemas = prim.GetAppliedSchemas()
185
- schemas_str = "".join(schemas)
186
- if f"DriveAPI:{type}" not in schemas_str and f"PhysicsLimitAPI:{type}" not in schemas_str:
187
- return
188
- drive_type = parse_generic(prim, f"drive:{type}:physics:type", "force")
189
- if drive_type != "force":
190
- print(f"Warning: only force drive type is supported, ignoring drive:{type} for joint {path}")
191
- return
192
- stiffness = parse_float(prim, f"drive:{type}:physics:stiffness", 0.0)
193
- damping = parse_float(prim, f"drive:{type}:physics:damping", 0.0)
194
- low = parse_float(prim, f"limit:{type}:physics:low")
195
- high = parse_float(prim, f"limit:{type}:physics:high")
196
- target_pos = parse_float(prim, f"drive:{type}:physics:targetPosition")
197
- target_vel = parse_float(prim, f"drive:{type}:physics:targetVelocity")
198
- if is_angular:
199
- stiffness *= mass_unit * linear_unit**2
200
- stiffness = np.deg2rad(stiffness)
201
- damping *= mass_unit * linear_unit**2
202
- damping = np.deg2rad(damping)
203
- if target_pos is not None:
204
- target_pos = np.deg2rad(target_pos)
205
- if target_vel is not None:
206
- target_vel = np.deg2rad(target_vel)
207
- if low is None:
208
- low = joint_data["lowerLimit"]
209
- else:
210
- low = np.deg2rad(low)
211
- if high is None:
212
- high = joint_data["upperLimit"]
213
- else:
214
- high = np.deg2rad(high)
215
- else:
216
- stiffness *= mass_unit
217
- damping *= mass_unit
218
- if target_pos is not None:
219
- target_pos *= linear_unit
220
- if target_vel is not None:
221
- target_vel *= linear_unit
222
- if low is None:
223
- low = joint_data["lowerLimit"]
224
- else:
225
- low *= linear_unit
226
- if high is None:
227
- high = joint_data["upperLimit"]
228
- else:
229
- high *= linear_unit
230
-
231
- mode = wp.sim.JOINT_MODE_FORCE
232
- if f"DriveAPI:{type}" in schemas_str:
233
- if target_vel is not None and target_vel != 0.0:
234
- mode = wp.sim.JOINT_MODE_TARGET_VELOCITY
235
- else:
236
- mode = wp.sim.JOINT_MODE_TARGET_POSITION
237
- if low > high:
238
- low = (low + high) / 2
239
- high = low
240
- axis = wp.sim.JointAxis(
241
- axis=(axis or joint_data["axis"]),
242
- limit_lower=low,
243
- limit_upper=high,
244
- action=(target_pos or target_vel or (low + high) / 2),
245
- target_ke=stiffness,
246
- target_kd=damping,
247
- mode=mode,
248
- limit_ke=joint_limit_ke,
249
- limit_kd=joint_limit_kd,
250
- )
251
- if is_angular:
252
- joint_data["angular_axes"].append(axis)
253
- else:
254
- joint_data["linear_axes"].append(axis)
255
-
256
- axis_str = "Y"
257
- try:
258
- axis_str = UsdGeom.GetStageUpAxis(stage)
259
- except Exception as e:
260
- if verbose:
261
- print(f"Failed to parse stage up axis: {e}")
262
- upaxis = str2axis(axis_str)
263
-
264
- shape_types = {"Cube", "Sphere", "Mesh", "Capsule", "Plane", "Cylinder", "Cone"}
265
-
266
- path_body_map = {}
267
- path_shape_map = {}
268
- path_shape_scale = {}
269
- # maps prim path name to its world transform
270
- path_world_poses = {}
271
- # transform from body frame to where the actual joint child frame is
272
- # so that the link's children will use the right parent tf for the joint
273
- prim_joint_xforms = {}
274
- path_collision_filters = set()
275
- no_collision_shapes = set()
276
-
277
- body_density = {} # mapping from body ID to defined density
278
-
279
- # first find all joints and materials
280
- joint_data = {} # mapping from path of child link to joint USD settings
281
- materials = {} # mapping from material path to material USD settings
282
- joint_parents = set() # paths of joint parents
283
- for prim in stage.Traverse():
284
- type_name = str(prim.GetTypeName())
285
- path = str(prim.GetPath())
286
- # if verbose:
287
- # print(path, type_name)
288
- if type_name.endswith("Joint"):
289
- # the type name can sometimes be "DistancePhysicsJoint" or "PhysicsDistanceJoint" ...
290
- type_name = type_name.replace("Physics", "").replace("Joint", "")
291
- child = str(prim.GetRelationship("physics:body1").GetTargets()[0])
292
- pos0 = parse_vec(prim, "physics:localPos0", np.zeros(3, dtype=np.float32)) * linear_unit
293
- pos1 = parse_vec(prim, "physics:localPos1", np.zeros(3, dtype=np.float32)) * linear_unit
294
- rot0 = parse_quat(prim, "physics:localRot0", wp.quat_identity())
295
- rot1 = parse_quat(prim, "physics:localRot1", wp.quat_identity())
296
- joint_data[child] = {
297
- "type": type_name,
298
- "name": str(prim.GetName()),
299
- "parent_tf": wp.transform(pos0, rot0),
300
- "child_tf": wp.transform(pos1, rot1),
301
- "enabled": parse_generic(prim, "physics:jointEnabled", True),
302
- "collisionEnabled": parse_generic(prim, "physics:collisionEnabled", False),
303
- "excludeFromArticulation": parse_generic(prim, "physics:excludeFromArticulation", False),
304
- "axis": str2axis(parse_generic(prim, "physics:axis", "X")),
305
- "breakForce": parse_float(prim, "physics:breakForce", np.inf),
306
- "breakTorque": parse_float(prim, "physics:breakTorque", np.inf),
307
- "linear_axes": [],
308
- "angular_axes": [],
309
- }
310
- if only_load_enabled_joints and not joint_data[child]["enabled"]:
311
- print("Skipping disabled joint", path)
312
- continue
313
- # parse joint limits
314
- lower = parse_float(prim, "physics:lowerLimit", -np.inf)
315
- upper = parse_float(prim, "physics:upperLimit", np.inf)
316
- if type_name == "Distance":
317
- # if distance is negative the joint is not limited
318
- joint_data[child]["lowerLimit"] = parse_float(prim, "physics:minDistance", -1.0) * linear_unit
319
- joint_data[child]["upperLimit"] = parse_float(prim, "physics:maxDistance", -1.0) * linear_unit
320
- elif type_name == "Prismatic":
321
- joint_data[child]["lowerLimit"] = lower * linear_unit
322
- joint_data[child]["upperLimit"] = upper * linear_unit
323
- else:
324
- joint_data[child]["lowerLimit"] = np.deg2rad(lower) if np.isfinite(lower) else lower
325
- joint_data[child]["upperLimit"] = np.deg2rad(upper) if np.isfinite(upper) else upper
326
-
327
- if joint_data[child]["lowerLimit"] > joint_data[child]["upperLimit"]:
328
- joint_data[child]["lowerLimit"] = (
329
- joint_data[child]["lowerLimit"] + joint_data[child]["upperLimit"]
330
- ) / 2
331
- joint_data[child]["upperLimit"] = joint_data[child]["lowerLimit"]
332
- parents = prim.GetRelationship("physics:body0").GetTargets()
333
- if len(parents) > 0:
334
- parent_path = str(parents[0])
335
- joint_data[child]["parent"] = parent_path
336
- joint_parents.add(parent_path)
337
- else:
338
- joint_data[child]["parent"] = None
339
-
340
- # parse joint drive
341
- parse_axis(prim, "angular", joint_data[child], is_angular=True)
342
- parse_axis(prim, "rotX", joint_data[child], is_angular=True, axis=(1.0, 0.0, 0.0))
343
- parse_axis(prim, "rotY", joint_data[child], is_angular=True, axis=(0.0, 1.0, 0.0))
344
- parse_axis(prim, "rotZ", joint_data[child], is_angular=True, axis=(0.0, 0.0, 1.0))
345
- parse_axis(prim, "linear", joint_data[child], is_angular=False)
346
- parse_axis(prim, "transX", joint_data[child], is_angular=False, axis=(1.0, 0.0, 0.0))
347
- parse_axis(prim, "transY", joint_data[child], is_angular=False, axis=(0.0, 1.0, 0.0))
348
- parse_axis(prim, "transZ", joint_data[child], is_angular=False, axis=(0.0, 0.0, 1.0))
349
-
350
- elif type_name == "Material":
351
- material = {}
352
- if has_attribute(prim, "physics:density"):
353
- material["density"] = parse_float(prim, "physics:density") * mass_unit # / (linear_unit**3)
354
- if has_attribute(prim, "physics:restitution"):
355
- material["restitution"] = parse_float(prim, "physics:restitution", contact_restitution)
356
- if has_attribute(prim, "physics:staticFriction"):
357
- material["staticFriction"] = parse_float(prim, "physics:staticFriction", contact_mu)
358
- if has_attribute(prim, "physics:dynamicFriction"):
359
- material["dynamicFriction"] = parse_float(prim, "physics:dynamicFriction", contact_mu)
360
- materials[path] = material
361
-
362
- elif type_name == "PhysicsScene":
363
- try:
364
- scene = UsdPhysics.Scene(prim)
365
- g_vec = scene.GetGravityDirectionAttr()
366
- g_mag = scene.GetGravityMagnitudeAttr()
367
- if g_mag.HasAuthoredValue() and np.isfinite(g_mag.Get()):
368
- builder.gravity = -np.abs(g_mag.Get() * linear_unit)
369
- if g_vec.HasAuthoredValue() and np.linalg.norm(g_vec.Get()) > 0.0:
370
- builder.up_vector = np.array(g_vec.Get(), dtype=np.float32)
371
- if np.any(builder.up_vector < 0.0):
372
- builder.up_vector = -builder.up_vector
373
- else:
374
- builder.up_vector = upaxis
375
- except Exception as e:
376
- if verbose:
377
- print(f"Failed to parse physics scene: {e}")
378
-
379
- def parse_prim(prim, incoming_xform, incoming_scale, parent_body: int = -1):
380
- nonlocal builder
381
- nonlocal joint_data
382
- nonlocal path_body_map
383
- nonlocal path_shape_map
384
- nonlocal path_shape_scale
385
- nonlocal path_world_poses
386
- nonlocal prim_joint_xforms
387
- nonlocal path_collision_filters
388
- nonlocal no_collision_shapes
389
- nonlocal body_density
390
-
391
- path = str(prim.GetPath())
392
- for pattern in ignore_paths:
393
- if re.match(pattern, path):
394
- return
395
-
396
- type_name = str(prim.GetTypeName())
397
- if type_name.endswith("Joint") or type_name.endswith("Light") or type_name.endswith("Material"):
398
- return
399
- if verbose:
400
- print(f"parse_prim {prim.GetPath()} ({type_name})")
401
- if type_name == "PhysicsScene":
402
- # in case the PhysicsScene has bodies as children...
403
- for child in prim.GetChildren():
404
- parse_prim(child, incoming_xform, incoming_scale, parent_body)
405
-
406
- schemas = set(prim.GetAppliedSchemas())
407
- children_refs = prim.GetChildren()
408
-
409
- prim_joint_xforms[path] = wp.transform()
410
-
411
- local_xform, scale = parse_xform(prim)
412
- scale = incoming_scale * scale
413
- xform = wp.mul(incoming_xform, local_xform)
414
- path_world_poses[path] = xform
415
-
416
- geo_tf = local_xform
417
- body_id = parent_body
418
- is_rigid_body = "PhysicsRigidBodyAPI" in schemas and parent_body == -1
419
- create_rigid_body = is_rigid_body or path in joint_parents
420
- if create_rigid_body:
421
- body_id = builder.add_body(
422
- origin=xform,
423
- name=prim.GetName(),
424
- armature=armature,
425
- )
426
- path_body_map[path] = body_id
427
- body_density[body_id] = 0.0
428
-
429
- parent_body = body_id
430
-
431
- geo_tf = wp.transform()
432
-
433
- # set up joints between rigid bodies after the children have been added
434
- if path in joint_data:
435
- joint = joint_data[path]
436
-
437
- joint_params = dict(
438
- child=body_id,
439
- linear_axes=joint["linear_axes"],
440
- angular_axes=joint["angular_axes"],
441
- name=joint["name"],
442
- enabled=joint["enabled"],
443
- parent_xform=joint["parent_tf"],
444
- child_xform=joint["child_tf"],
445
- )
446
-
447
- parent_path = joint["parent"]
448
- if parent_path is None:
449
- joint_params["parent"] = -1
450
- parent_tf = wp.transform()
451
- else:
452
- joint_params["parent"] = path_body_map[parent_path]
453
- parent_tf = path_world_poses[parent_path]
454
-
455
- # the joint to which we are connected will transform this body already
456
- geo_tf = wp.transform()
457
-
458
- if verbose:
459
- print(f"Adding joint {joint['name']} between {joint['parent']} and {path}")
460
- print(" parent_xform", joint["parent_tf"])
461
- print(" child_xform ", joint["child_tf"])
462
- print(" parent_tf ", parent_tf)
463
- print(f" geo_tf at {path} = {geo_tf} (xform was {xform})")
464
-
465
- if joint["type"] == "Revolute":
466
- joint_params["joint_type"] = wp.sim.JOINT_REVOLUTE
467
- if len(joint_params["angular_axes"]) == 0:
468
- joint_params["angular_axes"].append(
469
- wp.sim.JointAxis(
470
- joint["axis"],
471
- limit_lower=joint["lowerLimit"],
472
- limit_upper=joint["upperLimit"],
473
- limit_ke=joint_limit_ke,
474
- limit_kd=joint_limit_kd,
475
- )
476
- )
477
- elif joint["type"] == "Prismatic":
478
- joint_params["joint_type"] = wp.sim.JOINT_PRISMATIC
479
- if len(joint_params["linear_axes"]) == 0:
480
- joint_params["linear_axes"].append(
481
- wp.sim.JointAxis(
482
- joint["axis"],
483
- limit_lower=joint["lowerLimit"],
484
- limit_upper=joint["upperLimit"],
485
- limit_ke=joint_limit_ke,
486
- limit_kd=joint_limit_kd,
487
- )
488
- )
489
- elif joint["type"] == "Spherical":
490
- joint_params["joint_type"] = wp.sim.JOINT_BALL
491
- elif joint["type"] == "Fixed":
492
- joint_params["joint_type"] = wp.sim.JOINT_FIXED
493
- elif joint["type"] == "Distance":
494
- joint_params["joint_type"] = wp.sim.JOINT_DISTANCE
495
- # we have to add a dummy linear X axis to define the joint limits
496
- joint_params["linear_axes"].append(
497
- wp.sim.JointAxis(
498
- (1.0, 0.0, 0.0),
499
- limit_lower=joint["lowerLimit"],
500
- limit_upper=joint["upperLimit"],
501
- limit_ke=joint_limit_ke,
502
- limit_kd=joint_limit_kd,
503
- )
504
- )
505
- elif joint["type"] == "":
506
- joint_params["joint_type"] = wp.sim.JOINT_D6
507
- else:
508
- print(f"Warning: unsupported joint type {joint['type']} for {path}")
509
-
510
- builder.add_joint(**joint_params)
511
-
512
- elif is_rigid_body:
513
- builder.add_joint_free(child=body_id)
514
- # free joint; we set joint_q/qd, not body_q/qd since eval_fk is used after model creation
515
- builder.joint_q[-4:] = xform.q
516
- builder.joint_q[-7:-4] = xform.p
517
- linear_vel = parse_vec(prim, "physics:velocity", np.zeros(3, dtype=np.float32)) * linear_unit
518
- angular_vel = parse_vec(prim, "physics:angularVelocity", np.zeros(3, dtype=np.float32)) * linear_unit
519
- builder.joint_qd[-6:-3] = angular_vel
520
- builder.joint_qd[-3:] = linear_vel
521
-
522
- if verbose:
523
- print(f"added {type_name} body {body_id} ({path}) at {xform}")
524
-
525
- density = None
526
-
527
- material = None
528
- if prim.HasRelationship("material:binding:physics"):
529
- other_paths = prim.GetRelationship("material:binding:physics").GetTargets()
530
- if len(other_paths) > 0:
531
- material = materials[str(other_paths[0])]
532
- if material is not None:
533
- if "density" in material:
534
- density = material["density"]
535
- if has_attribute(prim, "physics:density"):
536
- d = parse_float(prim, "physics:density")
537
- density = d * mass_unit # / (linear_unit**3)
538
-
539
- # assert prim.GetAttribute('orientation').Get() == "rightHanded", "Only right-handed orientations are supported."
540
- enabled = parse_generic(prim, "physics:rigidBodyEnabled", True)
541
- if only_load_enabled_rigid_bodies and not enabled:
542
- if verbose:
543
- print("Skipping disabled rigid body", path)
544
- return
545
- mass = parse_float(prim, "physics:mass")
546
- if is_rigid_body:
547
- if density is None:
548
- density = default_density
549
- body_density[body_id] = density
550
- elif density is None:
551
- if body_id >= 0:
552
- density = body_density[body_id]
553
- else:
554
- density = 0.0
555
-
556
- com = parse_vec(prim, "physics:centerOfMass", np.zeros(3, dtype=np.float32))
557
- i_diag = parse_vec(prim, "physics:diagonalInertia", np.zeros(3, dtype=np.float32))
558
- i_rot = parse_quat(prim, "physics:principalAxes", wp.quat_identity())
559
-
560
- # parse children
561
- if type_name == "Xform":
562
- if prim.IsInstance():
563
- proto = prim.GetPrototype()
564
- for child in proto.GetChildren():
565
- parse_prim(child, xform, scale, parent_body)
566
- else:
567
- for child in children_refs:
568
- parse_prim(child, xform, scale, parent_body)
569
- elif type_name == "Scope":
570
- for child in children_refs:
571
- parse_prim(child, incoming_xform, incoming_scale, parent_body)
572
- elif type_name in shape_types:
573
- # parse shapes
574
- shape_params = dict(
575
- ke=contact_ke,
576
- kd=contact_kd,
577
- kf=contact_kf,
578
- ka=contact_ka,
579
- mu=contact_mu,
580
- restitution=contact_restitution,
581
- )
582
- if material is not None:
583
- if "restitution" in material:
584
- shape_params["restitution"] = material["restitution"]
585
- if "dynamicFriction" in material:
586
- shape_params["mu"] = material["dynamicFriction"]
587
-
588
- if has_attribute(prim, "doubleSided") and not prim.GetAttribute("doubleSided").Get():
589
- print(f"Warning: treating {path} as double-sided because single-sided collisions are not supported.")
590
-
591
- if type_name == "Cube":
592
- size = parse_float(prim, "size", 2.0)
593
- if has_attribute(prim, "extents"):
594
- extents = parse_vec(prim, "extents") * scale
595
- # TODO position geom at extents center?
596
- geo_pos = 0.5 * (extents[0] + extents[1])
597
- extents = extents[1] - extents[0]
598
- else:
599
- extents = scale * size
600
- shape_id = builder.add_shape_box(
601
- body_id,
602
- geo_tf.p,
603
- geo_tf.q,
604
- hx=extents[0] / 2,
605
- hy=extents[1] / 2,
606
- hz=extents[2] / 2,
607
- density=density,
608
- thickness=contact_thickness,
609
- **shape_params,
610
- )
611
- elif type_name == "Sphere":
612
- if not (scale[0] == scale[1] == scale[2]):
613
- print("Warning: Non-uniform scaling of spheres is not supported.")
614
- if has_attribute(prim, "extents"):
615
- extents = parse_vec(prim, "extents") * scale
616
- # TODO position geom at extents center?
617
- geo_pos = 0.5 * (extents[0] + extents[1])
618
- extents = extents[1] - extents[0]
619
- if not (extents[0] == extents[1] == extents[2]):
620
- print("Warning: Non-uniform extents of spheres are not supported.")
621
- radius = extents[0]
622
- else:
623
- radius = parse_float(prim, "radius", 1.0) * scale[0]
624
- shape_id = builder.add_shape_sphere(
625
- body_id, geo_tf.p, geo_tf.q, radius, density=density, **shape_params
626
- )
627
- elif type_name == "Plane":
628
- normal_str = parse_generic(prim, "axis", "Z").upper()
629
- geo_rot = geo_tf.q
630
- if normal_str != "Y":
631
- normal = str2axis(normal_str)
632
- c = np.cross(normal, (0.0, 1.0, 0.0))
633
- angle = np.arcsin(np.linalg.norm(c))
634
- axis = c / np.linalg.norm(c)
635
- geo_rot = wp.mul(geo_rot, wp.quat_from_axis_angle(axis, angle))
636
- width = parse_float(prim, "width", 0.0) * scale[0]
637
- length = parse_float(prim, "length", 0.0) * scale[1]
638
- shape_id = builder.add_shape_plane(
639
- body=body_id,
640
- pos=geo_tf.p,
641
- rot=geo_rot,
642
- width=width,
643
- length=length,
644
- thickness=contact_thickness,
645
- **shape_params,
646
- )
647
- elif type_name == "Capsule":
648
- axis_str = parse_generic(prim, "axis", "Z").upper()
649
- radius = parse_float(prim, "radius", 0.5) * scale[0]
650
- half_height = parse_float(prim, "height", 2.0) / 2 * scale[1]
651
- assert not has_attribute(prim, "extents"), "Capsule extents are not supported."
652
- shape_id = builder.add_shape_capsule(
653
- body_id,
654
- geo_tf.p,
655
- geo_tf.q,
656
- radius,
657
- half_height,
658
- density=density,
659
- up_axis="XYZ".index(axis_str),
660
- **shape_params,
661
- )
662
- elif type_name == "Cylinder":
663
- axis_str = parse_generic(prim, "axis", "Z").upper()
664
- radius = parse_float(prim, "radius", 0.5) * scale[0]
665
- half_height = parse_float(prim, "height", 2.0) / 2 * scale[1]
666
- assert not has_attribute(prim, "extents"), "Cylinder extents are not supported."
667
- shape_id = builder.add_shape_cylinder(
668
- body_id,
669
- geo_tf.p,
670
- geo_tf.q,
671
- radius,
672
- half_height,
673
- density=density,
674
- up_axis="XYZ".index(axis_str),
675
- **shape_params,
676
- )
677
- elif type_name == "Cone":
678
- axis_str = parse_generic(prim, "axis", "Z").upper()
679
- radius = parse_float(prim, "radius", 0.5) * scale[0]
680
- half_height = parse_float(prim, "height", 2.0) / 2 * scale[1]
681
- assert not has_attribute(prim, "extents"), "Cone extents are not supported."
682
- shape_id = builder.add_shape_cone(
683
- body_id,
684
- geo_tf.p,
685
- geo_tf.q,
686
- radius,
687
- half_height,
688
- density=density,
689
- up_axis="XYZ".index(axis_str),
690
- **shape_params,
691
- )
692
- elif type_name == "Mesh":
693
- mesh = UsdGeom.Mesh(prim)
694
- points = np.array(mesh.GetPointsAttr().Get(), dtype=np.float32)
695
- indices = np.array(mesh.GetFaceVertexIndicesAttr().Get(), dtype=np.float32)
696
- counts = mesh.GetFaceVertexCountsAttr().Get()
697
- faces = []
698
- face_id = 0
699
- for count in counts:
700
- if count == 3:
701
- faces.append(indices[face_id : face_id + 3])
702
- elif count == 4:
703
- faces.append(indices[face_id : face_id + 3])
704
- faces.append(indices[[face_id, face_id + 2, face_id + 3]])
705
- else:
706
- # assert False, f"Error while parsing USD mesh {path}: encountered polygon with {count} vertices, but only triangles and quads are supported."
707
- continue
708
- face_id += count
709
- m = wp.sim.Mesh(points, np.array(faces, dtype=np.int32).flatten())
710
- shape_id = builder.add_shape_mesh(
711
- body_id,
712
- geo_tf.p,
713
- geo_tf.q,
714
- scale=scale,
715
- mesh=m,
716
- density=density,
717
- thickness=contact_thickness,
718
- **shape_params,
719
- )
720
- else:
721
- print(f"Warning: Unsupported geometry type {type_name} at {path}.")
722
- return
723
-
724
- path_body_map[path] = body_id
725
- path_shape_map[path] = shape_id
726
- path_shape_scale[path] = scale
727
-
728
- if prim.HasRelationship("physics:filteredPairs"):
729
- other_paths = prim.GetRelationship("physics:filteredPairs").GetTargets()
730
- for other_path in other_paths:
731
- path_collision_filters.add((path, str(other_path)))
732
-
733
- if "PhysicsCollisionAPI" not in schemas or not parse_generic(prim, "physics:collisionEnabled", True):
734
- no_collision_shapes.add(shape_id)
735
-
736
- else:
737
- print(f"Warning: encountered unsupported prim type {type_name}")
738
-
739
- # update mass properties of rigid bodies in cases where properties are defined with higher precedence
740
- if body_id >= 0:
741
- com = parse_vec(prim, "physics:centerOfMass")
742
- if com is not None:
743
- # overwrite COM
744
- builder.body_com[body_id] = com * scale
745
-
746
- if mass is not None and not (is_rigid_body and mass == 0.0):
747
- mass_ratio = mass / builder.body_mass[body_id]
748
- # mass has precedence over density, so we overwrite the mass computed from density
749
- builder.body_mass[body_id] = mass * mass_unit
750
- if mass > 0.0:
751
- builder.body_inv_mass[body_id] = 1.0 / builder.body_mass[body_id]
752
- else:
753
- builder.body_inv_mass[body_id] = 0.0
754
- # update inertia
755
- builder.body_inertia[body_id] *= mass_ratio
756
- if np.array(builder.body_inertia[body_id]).any():
757
- builder.body_inv_inertia[body_id] = wp.inverse(builder.body_inertia[body_id])
758
- else:
759
- builder.body_inv_inertia[body_id] = wp.mat33(*np.zeros((3, 3), dtype=np.float32))
760
-
761
- if np.linalg.norm(i_diag) > 0.0:
762
- rot = np.array(wp.quat_to_matrix(i_rot), dtype=np.float32).reshape(3, 3)
763
- inertia = rot @ np.diag(i_diag) @ rot.T
764
- builder.body_inertia[body_id] = inertia
765
- if inertia.any():
766
- builder.body_inv_inertia[body_id] = wp.inverse(wp.mat33(*inertia))
767
- else:
768
- builder.body_inv_inertia[body_id] = wp.mat33(*np.zeros((3, 3), dtype=np.float32))
769
-
770
- parse_prim(
771
- stage.GetDefaultPrim(), incoming_xform=wp.transform(), incoming_scale=np.ones(3, dtype=np.float32) * linear_unit
772
- )
773
-
774
- shape_count = len(builder.shape_geo_type)
775
-
776
- # apply collision filters now that we have added all shapes
777
- for path1, path2 in path_collision_filters:
778
- shape1 = path_shape_map[path1]
779
- shape2 = path_shape_map[path2]
780
- builder.shape_collision_filter_pairs.add((shape1, shape2))
781
-
782
- # apply collision filters to all shapes that have no collision
783
- for shape_id in no_collision_shapes:
784
- for other_shape_id in range(shape_count):
785
- if other_shape_id != shape_id:
786
- builder.shape_collision_filter_pairs.add((shape_id, other_shape_id))
787
-
788
- # return stage parameters
789
- return {
790
- "fps": stage.GetFramesPerSecond(),
791
- "duration": stage.GetEndTimeCode() - stage.GetStartTimeCode(),
792
- "up_axis": UsdGeom.GetStageUpAxis(stage).upper(),
793
- "path_shape_map": path_shape_map,
794
- "path_body_map": path_body_map,
795
- "path_shape_scale": path_shape_scale,
796
- "mass_unit": mass_unit,
797
- "linear_unit": linear_unit,
798
- }
799
-
800
-
801
- def resolve_usd_from_url(url: str, target_folder_name: str = None, export_usda: bool = False):
802
- """
803
- Downloads a USD file from a URL and resolves all references to other USD files to be downloaded to the given target folder.
804
-
805
- Args:
806
- url (str): URL to the USD file.
807
- target_folder_name (str): Target folder name. If None, a timestamped folder will be created in the current directory.
808
- 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.
809
-
810
- Returns:
811
- str: File path to the downloaded USD file.
812
- """
813
- import requests
814
- import datetime
815
- import os
816
-
817
- try:
818
- from pxr import Usd
819
- except ImportError:
820
- raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).")
821
-
822
- response = requests.get(url, allow_redirects=True)
823
- if response.status_code != 200:
824
- raise RuntimeError(f"Failed to download USD file. Status code: {response.status_code}")
825
- file = response.content
826
- dot = os.path.extsep
827
- base = os.path.basename(url)
828
- url_folder = os.path.dirname(url)
829
- base_name = dot.join(base.split(dot)[:-1])
830
- if target_folder_name is None:
831
- timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
832
- target_folder_name = os.path.join(".usd_cache", f"{base_name}_{timestamp}")
833
- os.makedirs(target_folder_name, exist_ok=True)
834
- target_filename = os.path.join(target_folder_name, base)
835
- with open(target_filename, "wb") as f:
836
- f.write(file)
837
-
838
- stage = Usd.Stage.Open(target_filename, Usd.Stage.LoadNone)
839
- stage_str = stage.GetRootLayer().ExportToString()
840
- print(f"Downloaded USD file to {target_filename}.")
841
- if export_usda:
842
- usda_filename = os.path.join(target_folder_name, base_name + ".usda")
843
- with open(usda_filename, "w") as f:
844
- f.write(stage_str)
845
- print(f"Exported USDA file to {usda_filename}.")
846
-
847
- # parse referenced USD files like `references = @./franka_collisions.usd@`
848
- downloaded = set()
849
- for match in re.finditer(r"references.=.@(.*?)@", stage_str):
850
- refname = match.group(1)
851
- if refname.startswith("./"):
852
- refname = refname[2:]
853
- if refname in downloaded:
854
- continue
855
- try:
856
- response = requests.get(f"{url_folder}/{refname}", allow_redirects=True)
857
- if response.status_code != 200:
858
- print(f"Failed to download reference {refname}. Status code: {response.status_code}")
859
- continue
860
- file = response.content
861
- refdir = os.path.dirname(refname)
862
- if refdir:
863
- os.makedirs(os.path.join(target_folder_name, refdir), exist_ok=True)
864
- ref_filename = os.path.join(target_folder_name, refname)
865
- if not os.path.exists(ref_filename):
866
- with open(ref_filename, "wb") as f:
867
- f.write(file)
868
- downloaded.add(refname)
869
- print(f"Downloaded USD reference {refname} to {ref_filename}.")
870
- if export_usda:
871
- ref_stage = Usd.Stage.Open(ref_filename, Usd.Stage.LoadNone)
872
- ref_stage_str = ref_stage.GetRootLayer().ExportToString()
873
- base = os.path.basename(ref_filename)
874
- base_name = dot.join(base.split(dot)[:-1])
875
- usda_filename = os.path.join(target_folder_name, base_name + ".usda")
876
- with open(usda_filename, "w") as f:
877
- f.write(ref_stage_str)
878
- print(f"Exported USDA file to {usda_filename}.")
879
- except Exception:
880
- print(f"Failed to download {refname}.")
881
- return target_filename
1
+ # Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved.
2
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
3
+ # and proprietary rights in and to this software, related documentation
4
+ # and any modifications thereto. Any use, reproduction, disclosure or
5
+ # distribution of this software and related documentation without an express
6
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
+
8
+
9
+ import re
10
+
11
+ import numpy as np
12
+
13
+ import warp as wp
14
+
15
+
16
+ def parse_usd(
17
+ source,
18
+ builder,
19
+ default_density=1.0e3,
20
+ only_load_enabled_rigid_bodies=False,
21
+ only_load_enabled_joints=True,
22
+ contact_ke=1e5,
23
+ contact_kd=250.0,
24
+ contact_kf=500.0,
25
+ contact_ka=0.0,
26
+ contact_mu=0.6,
27
+ contact_restitution=0.0,
28
+ contact_thickness=0.0,
29
+ joint_limit_ke=100.0,
30
+ joint_limit_kd=10.0,
31
+ armature=0.0,
32
+ invert_rotations=False,
33
+ verbose=False,
34
+ ignore_paths=None,
35
+ ):
36
+ """
37
+ 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.
38
+
39
+ 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.
40
+
41
+ Args:
42
+ source (str | pxr.UsdStage): The file path to the USD file, or an existing USD stage instance.
43
+ builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
44
+ default_density (float): The default density to use for bodies without a density attribute.
45
+ only_load_enabled_rigid_bodies (bool): If True, only rigid bodies which do not have `physics:rigidBodyEnabled` set to False are loaded.
46
+ only_load_enabled_joints (bool): If True, only joints which do not have `physics:jointEnabled` set to False are loaded.
47
+ contact_ke (float): The default contact stiffness to use, only considered by the Euler integrators.
48
+ contact_kd (float): The default contact damping to use, only considered by the Euler integrators.
49
+ contact_kf (float): The default friction stiffness to use, only considered by the Euler integrators.
50
+ contact_ka (float): The default adhesion distance to use, only considered by the Euler integrators.
51
+ contact_mu (float): The default friction coefficient to use if a shape has not friction coefficient defined.
52
+ contact_restitution (float): The default coefficient of restitution to use if a shape has not coefficient of restitution defined.
53
+ contact_thickness (float): The thickness to add to the shape geometry.
54
+ joint_limit_ke (float): The default stiffness to use for joint limits, only considered by the Euler integrators.
55
+ joint_limit_kd (float): The default damping to use for joint limits, only considered by the Euler integrators.
56
+ armature (float): The armature to use for the bodies.
57
+ invert_rotations (bool): If True, inverts any rotations defined in the shape transforms.
58
+ verbose (bool): If True, print additional information about the parsed USD file.
59
+ ignore_paths (List[str]): A list of regular expressions matching prim paths to ignore.
60
+
61
+ Returns:
62
+ dict: Dictionary with the following entries:
63
+
64
+ .. list-table::
65
+ :widths: 25 75
66
+
67
+ * - "fps"
68
+ - USD stage frames per second
69
+ * - "duration"
70
+ - Difference between end time code and start time code of the USD stage
71
+ * - "up_axis"
72
+ - Upper-case string of the stage's up axis ("X", "Y", or "Z")
73
+ * - "path_shape_map"
74
+ - Mapping from prim path (str) of the UsdGeom to the respective shape index in :class:`ModelBuilder`
75
+ * - "path_body_map"
76
+ - Mapping from prim path (str) of a rigid body prim (e.g. that implements the PhysicsRigidBodyAPI) to the respective body index in :class:`ModelBuilder`
77
+ * - "path_shape_scale"
78
+ - Mapping from prim path (str) of the UsdGeom to its respective 3D world scale
79
+ * - "mass_unit"
80
+ - The stage's Kilograms Per Unit (KGPU) definition (1.0 by default)
81
+ * - "linear_unit"
82
+ - The stage's Meters Per Unit (MPU) definition (1.0 by default)
83
+
84
+
85
+ Note:
86
+ This importer is experimental and only supports a subset of the USD Physics schema. Please report any issues you encounter.
87
+ """
88
+ try:
89
+ from pxr import Usd, UsdGeom, UsdPhysics
90
+ except ImportError as e:
91
+ raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).") from e
92
+
93
+ if ignore_paths is None:
94
+ ignore_paths = []
95
+
96
+ def get_attribute(prim, name):
97
+ if "*" in name:
98
+ regex = name.replace("*", ".*")
99
+ for attr in prim.GetAttributes():
100
+ if re.match(regex, attr.GetName()):
101
+ return attr
102
+ else:
103
+ return prim.GetAttribute(name)
104
+
105
+ def has_attribute(prim, name):
106
+ attr = get_attribute(prim, name)
107
+ return attr.IsValid() and attr.HasAuthoredValue()
108
+
109
+ def parse_float(prim, name, default=None):
110
+ attr = get_attribute(prim, name)
111
+ if not attr or not attr.HasAuthoredValue():
112
+ return default
113
+ val = attr.Get()
114
+ if np.isfinite(val):
115
+ return val
116
+ return default
117
+
118
+ def parse_quat(prim, name, default=None):
119
+ attr = get_attribute(prim, name)
120
+ if not attr or not attr.HasAuthoredValue():
121
+ return default
122
+ val = attr.Get()
123
+ if invert_rotations:
124
+ quat = wp.quat(*val.imaginary, -val.real)
125
+ else:
126
+ quat = wp.quat(*val.imaginary, val.real)
127
+ l = wp.length(quat)
128
+ if np.isfinite(l) and l > 0.0:
129
+ return quat
130
+ return default
131
+
132
+ def parse_vec(prim, name, default=None):
133
+ attr = get_attribute(prim, name)
134
+ if not attr or not attr.HasAuthoredValue():
135
+ return default
136
+ val = attr.Get()
137
+ if np.isfinite(val).all():
138
+ return np.array(val, dtype=np.float32)
139
+ return default
140
+
141
+ def parse_generic(prim, name, default=None):
142
+ attr = get_attribute(prim, name)
143
+ if not attr or not attr.HasAuthoredValue():
144
+ return default
145
+ return attr.Get()
146
+
147
+ def str2axis(s: str) -> np.ndarray:
148
+ axis = np.zeros(3, dtype=np.float32)
149
+ axis["XYZ".index(s.upper())] = 1.0
150
+ return axis
151
+
152
+ if isinstance(source, str):
153
+ stage = Usd.Stage.Open(source, Usd.Stage.LoadAll)
154
+ else:
155
+ stage = source
156
+
157
+ mass_unit = 1.0
158
+ try:
159
+ if UsdPhysics.StageHasAuthoredKilogramsPerUnit(stage):
160
+ mass_unit = UsdPhysics.GetStageKilogramsPerUnit(stage)
161
+ except Exception as e:
162
+ if verbose:
163
+ print(f"Failed to get mass unit: {e}")
164
+ linear_unit = 1.0
165
+ try:
166
+ if UsdGeom.StageHasAuthoredMetersPerUnit(stage):
167
+ linear_unit = UsdGeom.GetStageMetersPerUnit(stage)
168
+ except Exception as e:
169
+ if verbose:
170
+ print(f"Failed to get linear unit: {e}")
171
+
172
+ def parse_xform(prim):
173
+ xform = UsdGeom.Xform(prim)
174
+ mat = np.array(xform.GetLocalTransformation(), dtype=np.float32)
175
+ if invert_rotations:
176
+ rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].T.flatten()))
177
+ else:
178
+ rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].flatten()))
179
+ pos = mat[3, :3] * linear_unit
180
+ scale = np.ones(3, dtype=np.float32)
181
+ for op in xform.GetOrderedXformOps():
182
+ if op.GetOpType() == UsdGeom.XformOp.TypeScale:
183
+ scale = np.array(op.Get(), dtype=np.float32)
184
+ return wp.transform(pos, rot), scale
185
+
186
+ def parse_axis(prim, type, joint_data, is_angular, axis=None):
187
+ # parse joint axis data
188
+ schemas = prim.GetAppliedSchemas()
189
+ schemas_str = "".join(schemas)
190
+ if f"DriveAPI:{type}" not in schemas_str and f"PhysicsLimitAPI:{type}" not in schemas_str:
191
+ return
192
+ drive_type = parse_generic(prim, f"drive:{type}:physics:type", "force")
193
+ if drive_type != "force":
194
+ print(f"Warning: only force drive type is supported, ignoring drive:{type} for joint {path}")
195
+ return
196
+ stiffness = parse_float(prim, f"drive:{type}:physics:stiffness", 0.0)
197
+ damping = parse_float(prim, f"drive:{type}:physics:damping", 0.0)
198
+ low = parse_float(prim, f"limit:{type}:physics:low")
199
+ high = parse_float(prim, f"limit:{type}:physics:high")
200
+ target_pos = parse_float(prim, f"drive:{type}:physics:targetPosition")
201
+ target_vel = parse_float(prim, f"drive:{type}:physics:targetVelocity")
202
+ if is_angular:
203
+ stiffness *= mass_unit * linear_unit**2
204
+ stiffness = np.deg2rad(stiffness)
205
+ damping *= mass_unit * linear_unit**2
206
+ damping = np.deg2rad(damping)
207
+ if target_pos is not None:
208
+ target_pos = np.deg2rad(target_pos)
209
+ if target_vel is not None:
210
+ target_vel = np.deg2rad(target_vel)
211
+ if low is None:
212
+ low = joint_data["lowerLimit"]
213
+ else:
214
+ low = np.deg2rad(low)
215
+ if high is None:
216
+ high = joint_data["upperLimit"]
217
+ else:
218
+ high = np.deg2rad(high)
219
+ else:
220
+ stiffness *= mass_unit
221
+ damping *= mass_unit
222
+ if target_pos is not None:
223
+ target_pos *= linear_unit
224
+ if target_vel is not None:
225
+ target_vel *= linear_unit
226
+ if low is None:
227
+ low = joint_data["lowerLimit"]
228
+ else:
229
+ low *= linear_unit
230
+ if high is None:
231
+ high = joint_data["upperLimit"]
232
+ else:
233
+ high *= linear_unit
234
+
235
+ mode = wp.sim.JOINT_MODE_FORCE
236
+ if f"DriveAPI:{type}" in schemas_str:
237
+ if target_vel is not None and target_vel != 0.0:
238
+ mode = wp.sim.JOINT_MODE_TARGET_VELOCITY
239
+ else:
240
+ mode = wp.sim.JOINT_MODE_TARGET_POSITION
241
+ if low > high:
242
+ low = (low + high) / 2
243
+ high = low
244
+ axis = wp.sim.JointAxis(
245
+ axis=(axis or joint_data["axis"]),
246
+ limit_lower=low,
247
+ limit_upper=high,
248
+ action=(target_pos or target_vel or (low + high) / 2),
249
+ target_ke=stiffness,
250
+ target_kd=damping,
251
+ mode=mode,
252
+ limit_ke=joint_limit_ke,
253
+ limit_kd=joint_limit_kd,
254
+ )
255
+ if is_angular:
256
+ joint_data["angular_axes"].append(axis)
257
+ else:
258
+ joint_data["linear_axes"].append(axis)
259
+
260
+ axis_str = "Y"
261
+ try:
262
+ axis_str = UsdGeom.GetStageUpAxis(stage)
263
+ except Exception as e:
264
+ if verbose:
265
+ print(f"Failed to parse stage up axis: {e}")
266
+ upaxis = str2axis(axis_str)
267
+
268
+ shape_types = {"Cube", "Sphere", "Mesh", "Capsule", "Plane", "Cylinder", "Cone"}
269
+
270
+ path_body_map = {}
271
+ path_shape_map = {}
272
+ path_shape_scale = {}
273
+ # maps prim path name to its world transform
274
+ path_world_poses = {}
275
+ # transform from body frame to where the actual joint child frame is
276
+ # so that the link's children will use the right parent tf for the joint
277
+ prim_joint_xforms = {}
278
+ path_collision_filters = set()
279
+ no_collision_shapes = set()
280
+
281
+ body_density = {} # mapping from body ID to defined density
282
+
283
+ # first find all joints and materials
284
+ joint_data = {} # mapping from path of child link to joint USD settings
285
+ materials = {} # mapping from material path to material USD settings
286
+ joint_parents = set() # paths of joint parents
287
+ for prim in stage.Traverse():
288
+ type_name = str(prim.GetTypeName())
289
+ path = str(prim.GetPath())
290
+ # if verbose:
291
+ # print(path, type_name)
292
+ if type_name.endswith("Joint"):
293
+ # the type name can sometimes be "DistancePhysicsJoint" or "PhysicsDistanceJoint" ...
294
+ type_name = type_name.replace("Physics", "").replace("Joint", "")
295
+ child = str(prim.GetRelationship("physics:body1").GetTargets()[0])
296
+ pos0 = parse_vec(prim, "physics:localPos0", np.zeros(3, dtype=np.float32)) * linear_unit
297
+ pos1 = parse_vec(prim, "physics:localPos1", np.zeros(3, dtype=np.float32)) * linear_unit
298
+ rot0 = parse_quat(prim, "physics:localRot0", wp.quat_identity())
299
+ rot1 = parse_quat(prim, "physics:localRot1", wp.quat_identity())
300
+ joint_data[child] = {
301
+ "type": type_name,
302
+ "name": str(prim.GetName()),
303
+ "parent_tf": wp.transform(pos0, rot0),
304
+ "child_tf": wp.transform(pos1, rot1),
305
+ "enabled": parse_generic(prim, "physics:jointEnabled", True),
306
+ "collisionEnabled": parse_generic(prim, "physics:collisionEnabled", False),
307
+ "excludeFromArticulation": parse_generic(prim, "physics:excludeFromArticulation", False),
308
+ "axis": str2axis(parse_generic(prim, "physics:axis", "X")),
309
+ "breakForce": parse_float(prim, "physics:breakForce", np.inf),
310
+ "breakTorque": parse_float(prim, "physics:breakTorque", np.inf),
311
+ "linear_axes": [],
312
+ "angular_axes": [],
313
+ }
314
+ if only_load_enabled_joints and not joint_data[child]["enabled"]:
315
+ print("Skipping disabled joint", path)
316
+ continue
317
+ # parse joint limits
318
+ lower = parse_float(prim, "physics:lowerLimit", -np.inf)
319
+ upper = parse_float(prim, "physics:upperLimit", np.inf)
320
+ if type_name == "Distance":
321
+ # if distance is negative the joint is not limited
322
+ joint_data[child]["lowerLimit"] = parse_float(prim, "physics:minDistance", -1.0) * linear_unit
323
+ joint_data[child]["upperLimit"] = parse_float(prim, "physics:maxDistance", -1.0) * linear_unit
324
+ elif type_name == "Prismatic":
325
+ joint_data[child]["lowerLimit"] = lower * linear_unit
326
+ joint_data[child]["upperLimit"] = upper * linear_unit
327
+ else:
328
+ joint_data[child]["lowerLimit"] = np.deg2rad(lower) if np.isfinite(lower) else lower
329
+ joint_data[child]["upperLimit"] = np.deg2rad(upper) if np.isfinite(upper) else upper
330
+
331
+ if joint_data[child]["lowerLimit"] > joint_data[child]["upperLimit"]:
332
+ joint_data[child]["lowerLimit"] = (
333
+ joint_data[child]["lowerLimit"] + joint_data[child]["upperLimit"]
334
+ ) / 2
335
+ joint_data[child]["upperLimit"] = joint_data[child]["lowerLimit"]
336
+ parents = prim.GetRelationship("physics:body0").GetTargets()
337
+ if len(parents) > 0:
338
+ parent_path = str(parents[0])
339
+ joint_data[child]["parent"] = parent_path
340
+ joint_parents.add(parent_path)
341
+ else:
342
+ joint_data[child]["parent"] = None
343
+
344
+ # parse joint drive
345
+ parse_axis(prim, "angular", joint_data[child], is_angular=True)
346
+ parse_axis(prim, "rotX", joint_data[child], is_angular=True, axis=(1.0, 0.0, 0.0))
347
+ parse_axis(prim, "rotY", joint_data[child], is_angular=True, axis=(0.0, 1.0, 0.0))
348
+ parse_axis(prim, "rotZ", joint_data[child], is_angular=True, axis=(0.0, 0.0, 1.0))
349
+ parse_axis(prim, "linear", joint_data[child], is_angular=False)
350
+ parse_axis(prim, "transX", joint_data[child], is_angular=False, axis=(1.0, 0.0, 0.0))
351
+ parse_axis(prim, "transY", joint_data[child], is_angular=False, axis=(0.0, 1.0, 0.0))
352
+ parse_axis(prim, "transZ", joint_data[child], is_angular=False, axis=(0.0, 0.0, 1.0))
353
+
354
+ elif type_name == "Material":
355
+ material = {}
356
+ if has_attribute(prim, "physics:density"):
357
+ material["density"] = parse_float(prim, "physics:density") * mass_unit # / (linear_unit**3)
358
+ if has_attribute(prim, "physics:restitution"):
359
+ material["restitution"] = parse_float(prim, "physics:restitution", contact_restitution)
360
+ if has_attribute(prim, "physics:staticFriction"):
361
+ material["staticFriction"] = parse_float(prim, "physics:staticFriction", contact_mu)
362
+ if has_attribute(prim, "physics:dynamicFriction"):
363
+ material["dynamicFriction"] = parse_float(prim, "physics:dynamicFriction", contact_mu)
364
+ materials[path] = material
365
+
366
+ elif type_name == "PhysicsScene":
367
+ try:
368
+ scene = UsdPhysics.Scene(prim)
369
+ g_vec = scene.GetGravityDirectionAttr()
370
+ g_mag = scene.GetGravityMagnitudeAttr()
371
+ if g_mag.HasAuthoredValue() and np.isfinite(g_mag.Get()):
372
+ builder.gravity = -np.abs(g_mag.Get() * linear_unit)
373
+ if g_vec.HasAuthoredValue() and np.linalg.norm(g_vec.Get()) > 0.0:
374
+ builder.up_vector = np.array(g_vec.Get(), dtype=np.float32)
375
+ if np.any(builder.up_vector < 0.0):
376
+ builder.up_vector = -builder.up_vector
377
+ else:
378
+ builder.up_vector = upaxis
379
+ except Exception as e:
380
+ if verbose:
381
+ print(f"Failed to parse physics scene: {e}")
382
+
383
+ def parse_prim(prim, incoming_xform, incoming_scale, parent_body: int = -1):
384
+ nonlocal builder
385
+ nonlocal joint_data
386
+ nonlocal path_body_map
387
+ nonlocal path_shape_map
388
+ nonlocal path_shape_scale
389
+ nonlocal path_world_poses
390
+ nonlocal prim_joint_xforms
391
+ nonlocal path_collision_filters
392
+ nonlocal no_collision_shapes
393
+ nonlocal body_density
394
+
395
+ path = str(prim.GetPath())
396
+ for pattern in ignore_paths:
397
+ if re.match(pattern, path):
398
+ return
399
+
400
+ type_name = str(prim.GetTypeName())
401
+ if type_name.endswith("Joint") or type_name.endswith("Light") or type_name.endswith("Material"):
402
+ return
403
+ if verbose:
404
+ print(f"parse_prim {prim.GetPath()} ({type_name})")
405
+ if type_name == "PhysicsScene":
406
+ # in case the PhysicsScene has bodies as children...
407
+ for child in prim.GetChildren():
408
+ parse_prim(child, incoming_xform, incoming_scale, parent_body)
409
+
410
+ schemas = set(prim.GetAppliedSchemas())
411
+ children_refs = prim.GetChildren()
412
+
413
+ prim_joint_xforms[path] = wp.transform()
414
+
415
+ local_xform, scale = parse_xform(prim)
416
+ scale = incoming_scale * scale
417
+ xform = wp.mul(incoming_xform, local_xform)
418
+ path_world_poses[path] = xform
419
+
420
+ geo_tf = local_xform
421
+ body_id = parent_body
422
+ is_rigid_body = "PhysicsRigidBodyAPI" in schemas and parent_body == -1
423
+ create_rigid_body = is_rigid_body or path in joint_parents
424
+ if create_rigid_body:
425
+ body_id = builder.add_body(
426
+ origin=xform,
427
+ name=prim.GetName(),
428
+ armature=armature,
429
+ )
430
+ path_body_map[path] = body_id
431
+ body_density[body_id] = 0.0
432
+
433
+ parent_body = body_id
434
+
435
+ geo_tf = wp.transform()
436
+
437
+ # set up joints between rigid bodies after the children have been added
438
+ if path in joint_data:
439
+ joint = joint_data[path]
440
+
441
+ joint_params = {
442
+ "child": body_id,
443
+ "linear_axes": joint["linear_axes"],
444
+ "angular_axes": joint["angular_axes"],
445
+ "name": joint["name"],
446
+ "enabled": joint["enabled"],
447
+ "parent_xform": joint["parent_tf"],
448
+ "child_xform": joint["child_tf"],
449
+ "armature": armature,
450
+ }
451
+
452
+ parent_path = joint["parent"]
453
+ if parent_path is None:
454
+ joint_params["parent"] = -1
455
+ parent_tf = wp.transform()
456
+ else:
457
+ joint_params["parent"] = path_body_map[parent_path]
458
+ parent_tf = path_world_poses[parent_path]
459
+
460
+ # the joint to which we are connected will transform this body already
461
+ geo_tf = wp.transform()
462
+
463
+ if verbose:
464
+ print(f"Adding joint {joint['name']} between {joint['parent']} and {path}")
465
+ print(" parent_xform", joint["parent_tf"])
466
+ print(" child_xform ", joint["child_tf"])
467
+ print(" parent_tf ", parent_tf)
468
+ print(f" geo_tf at {path} = {geo_tf} (xform was {xform})")
469
+
470
+ if joint["type"] == "Revolute":
471
+ joint_params["joint_type"] = wp.sim.JOINT_REVOLUTE
472
+ if len(joint_params["angular_axes"]) == 0:
473
+ joint_params["angular_axes"].append(
474
+ wp.sim.JointAxis(
475
+ joint["axis"],
476
+ limit_lower=joint["lowerLimit"],
477
+ limit_upper=joint["upperLimit"],
478
+ limit_ke=joint_limit_ke,
479
+ limit_kd=joint_limit_kd,
480
+ )
481
+ )
482
+ elif joint["type"] == "Prismatic":
483
+ joint_params["joint_type"] = wp.sim.JOINT_PRISMATIC
484
+ if len(joint_params["linear_axes"]) == 0:
485
+ joint_params["linear_axes"].append(
486
+ wp.sim.JointAxis(
487
+ joint["axis"],
488
+ limit_lower=joint["lowerLimit"],
489
+ limit_upper=joint["upperLimit"],
490
+ limit_ke=joint_limit_ke,
491
+ limit_kd=joint_limit_kd,
492
+ )
493
+ )
494
+ elif joint["type"] == "Spherical":
495
+ joint_params["joint_type"] = wp.sim.JOINT_BALL
496
+ elif joint["type"] == "Fixed":
497
+ joint_params["joint_type"] = wp.sim.JOINT_FIXED
498
+ elif joint["type"] == "Distance":
499
+ joint_params["joint_type"] = wp.sim.JOINT_DISTANCE
500
+ # we have to add a dummy linear X axis to define the joint limits
501
+ joint_params["linear_axes"].append(
502
+ wp.sim.JointAxis(
503
+ (1.0, 0.0, 0.0),
504
+ limit_lower=joint["lowerLimit"],
505
+ limit_upper=joint["upperLimit"],
506
+ limit_ke=joint_limit_ke,
507
+ limit_kd=joint_limit_kd,
508
+ )
509
+ )
510
+ elif joint["type"] == "":
511
+ joint_params["joint_type"] = wp.sim.JOINT_D6
512
+ else:
513
+ print(f"Warning: unsupported joint type {joint['type']} for {path}")
514
+
515
+ builder.add_joint(**joint_params)
516
+
517
+ elif is_rigid_body:
518
+ builder.add_joint_free(child=body_id)
519
+ # free joint; we set joint_q/qd, not body_q/qd since eval_fk is used after model creation
520
+ builder.joint_q[-4:] = xform.q
521
+ builder.joint_q[-7:-4] = xform.p
522
+ linear_vel = parse_vec(prim, "physics:velocity", np.zeros(3, dtype=np.float32)) * linear_unit
523
+ angular_vel = parse_vec(prim, "physics:angularVelocity", np.zeros(3, dtype=np.float32)) * linear_unit
524
+ builder.joint_qd[-6:-3] = angular_vel
525
+ builder.joint_qd[-3:] = linear_vel
526
+
527
+ if verbose:
528
+ print(f"added {type_name} body {body_id} ({path}) at {xform}")
529
+
530
+ density = None
531
+
532
+ material = None
533
+ if prim.HasRelationship("material:binding:physics"):
534
+ other_paths = prim.GetRelationship("material:binding:physics").GetTargets()
535
+ if len(other_paths) > 0:
536
+ material = materials[str(other_paths[0])]
537
+ if material is not None:
538
+ if "density" in material:
539
+ density = material["density"]
540
+ if has_attribute(prim, "physics:density"):
541
+ d = parse_float(prim, "physics:density")
542
+ density = d * mass_unit # / (linear_unit**3)
543
+
544
+ # assert prim.GetAttribute('orientation').Get() == "rightHanded", "Only right-handed orientations are supported."
545
+ enabled = parse_generic(prim, "physics:rigidBodyEnabled", True)
546
+ if only_load_enabled_rigid_bodies and not enabled:
547
+ if verbose:
548
+ print("Skipping disabled rigid body", path)
549
+ return
550
+ mass = parse_float(prim, "physics:mass")
551
+ if is_rigid_body:
552
+ if density is None:
553
+ density = default_density
554
+ body_density[body_id] = density
555
+ elif density is None:
556
+ if body_id >= 0:
557
+ density = body_density[body_id]
558
+ else:
559
+ density = 0.0
560
+
561
+ com = parse_vec(prim, "physics:centerOfMass", np.zeros(3, dtype=np.float32))
562
+ i_diag = parse_vec(prim, "physics:diagonalInertia", np.zeros(3, dtype=np.float32))
563
+ i_rot = parse_quat(prim, "physics:principalAxes", wp.quat_identity())
564
+
565
+ # parse children
566
+ if type_name == "Xform":
567
+ if prim.IsInstance():
568
+ proto = prim.GetPrototype()
569
+ for child in proto.GetChildren():
570
+ parse_prim(child, xform, scale, parent_body)
571
+ else:
572
+ for child in children_refs:
573
+ parse_prim(child, xform, scale, parent_body)
574
+ elif type_name == "Scope":
575
+ for child in children_refs:
576
+ parse_prim(child, incoming_xform, incoming_scale, parent_body)
577
+ elif type_name in shape_types:
578
+ # parse shapes
579
+ shape_params = {
580
+ "ke": contact_ke,
581
+ "kd": contact_kd,
582
+ "kf": contact_kf,
583
+ "ka": contact_ka,
584
+ "mu": contact_mu,
585
+ "restitution": contact_restitution,
586
+ }
587
+ if material is not None:
588
+ if "restitution" in material:
589
+ shape_params["restitution"] = material["restitution"]
590
+ if "dynamicFriction" in material:
591
+ shape_params["mu"] = material["dynamicFriction"]
592
+
593
+ if has_attribute(prim, "doubleSided") and not prim.GetAttribute("doubleSided").Get():
594
+ print(f"Warning: treating {path} as double-sided because single-sided collisions are not supported.")
595
+
596
+ if type_name == "Cube":
597
+ size = parse_float(prim, "size", 2.0)
598
+ if has_attribute(prim, "extents"):
599
+ extents = parse_vec(prim, "extents") * scale
600
+ # TODO position geom at extents center?
601
+ # geo_pos = 0.5 * (extents[0] + extents[1])
602
+ extents = extents[1] - extents[0]
603
+ else:
604
+ extents = scale * size
605
+ shape_id = builder.add_shape_box(
606
+ body_id,
607
+ geo_tf.p,
608
+ geo_tf.q,
609
+ hx=extents[0] / 2,
610
+ hy=extents[1] / 2,
611
+ hz=extents[2] / 2,
612
+ density=density,
613
+ thickness=contact_thickness,
614
+ **shape_params,
615
+ )
616
+ elif type_name == "Sphere":
617
+ if not (scale[0] == scale[1] == scale[2]):
618
+ print("Warning: Non-uniform scaling of spheres is not supported.")
619
+ if has_attribute(prim, "extents"):
620
+ extents = parse_vec(prim, "extents") * scale
621
+ # TODO position geom at extents center?
622
+ # geo_pos = 0.5 * (extents[0] + extents[1])
623
+ extents = extents[1] - extents[0]
624
+ if not (extents[0] == extents[1] == extents[2]):
625
+ print("Warning: Non-uniform extents of spheres are not supported.")
626
+ radius = extents[0]
627
+ else:
628
+ radius = parse_float(prim, "radius", 1.0) * scale[0]
629
+ shape_id = builder.add_shape_sphere(
630
+ body_id, geo_tf.p, geo_tf.q, radius, density=density, **shape_params
631
+ )
632
+ elif type_name == "Plane":
633
+ normal_str = parse_generic(prim, "axis", "Z").upper()
634
+ geo_rot = geo_tf.q
635
+ if normal_str != "Y":
636
+ normal = str2axis(normal_str)
637
+ c = np.cross(normal, (0.0, 1.0, 0.0))
638
+ angle = np.arcsin(np.linalg.norm(c))
639
+ axis = c / np.linalg.norm(c)
640
+ geo_rot = wp.mul(geo_rot, wp.quat_from_axis_angle(axis, angle))
641
+ width = parse_float(prim, "width", 0.0) * scale[0]
642
+ length = parse_float(prim, "length", 0.0) * scale[1]
643
+ shape_id = builder.add_shape_plane(
644
+ body=body_id,
645
+ pos=geo_tf.p,
646
+ rot=geo_rot,
647
+ width=width,
648
+ length=length,
649
+ thickness=contact_thickness,
650
+ **shape_params,
651
+ )
652
+ elif type_name == "Capsule":
653
+ axis_str = parse_generic(prim, "axis", "Z").upper()
654
+ radius = parse_float(prim, "radius", 0.5) * scale[0]
655
+ half_height = parse_float(prim, "height", 2.0) / 2 * scale[1]
656
+ assert not has_attribute(prim, "extents"), "Capsule extents are not supported."
657
+ shape_id = builder.add_shape_capsule(
658
+ body_id,
659
+ geo_tf.p,
660
+ geo_tf.q,
661
+ radius,
662
+ half_height,
663
+ density=density,
664
+ up_axis="XYZ".index(axis_str),
665
+ **shape_params,
666
+ )
667
+ elif type_name == "Cylinder":
668
+ axis_str = parse_generic(prim, "axis", "Z").upper()
669
+ radius = parse_float(prim, "radius", 0.5) * scale[0]
670
+ half_height = parse_float(prim, "height", 2.0) / 2 * scale[1]
671
+ assert not has_attribute(prim, "extents"), "Cylinder extents are not supported."
672
+ shape_id = builder.add_shape_cylinder(
673
+ body_id,
674
+ geo_tf.p,
675
+ geo_tf.q,
676
+ radius,
677
+ half_height,
678
+ density=density,
679
+ up_axis="XYZ".index(axis_str),
680
+ **shape_params,
681
+ )
682
+ elif type_name == "Cone":
683
+ axis_str = parse_generic(prim, "axis", "Z").upper()
684
+ radius = parse_float(prim, "radius", 0.5) * scale[0]
685
+ half_height = parse_float(prim, "height", 2.0) / 2 * scale[1]
686
+ assert not has_attribute(prim, "extents"), "Cone extents are not supported."
687
+ shape_id = builder.add_shape_cone(
688
+ body_id,
689
+ geo_tf.p,
690
+ geo_tf.q,
691
+ radius,
692
+ half_height,
693
+ density=density,
694
+ up_axis="XYZ".index(axis_str),
695
+ **shape_params,
696
+ )
697
+ elif type_name == "Mesh":
698
+ mesh = UsdGeom.Mesh(prim)
699
+ points = np.array(mesh.GetPointsAttr().Get(), dtype=np.float32)
700
+ indices = np.array(mesh.GetFaceVertexIndicesAttr().Get(), dtype=np.float32)
701
+ counts = mesh.GetFaceVertexCountsAttr().Get()
702
+ faces = []
703
+ face_id = 0
704
+ for count in counts:
705
+ if count == 3:
706
+ faces.append(indices[face_id : face_id + 3])
707
+ elif count == 4:
708
+ faces.append(indices[face_id : face_id + 3])
709
+ faces.append(indices[[face_id, face_id + 2, face_id + 3]])
710
+ else:
711
+ # assert False, f"Error while parsing USD mesh {path}: encountered polygon with {count} vertices, but only triangles and quads are supported."
712
+ continue
713
+ face_id += count
714
+ m = wp.sim.Mesh(points, np.array(faces, dtype=np.int32).flatten())
715
+ shape_id = builder.add_shape_mesh(
716
+ body_id,
717
+ geo_tf.p,
718
+ geo_tf.q,
719
+ scale=scale,
720
+ mesh=m,
721
+ density=density,
722
+ thickness=contact_thickness,
723
+ **shape_params,
724
+ )
725
+ else:
726
+ print(f"Warning: Unsupported geometry type {type_name} at {path}.")
727
+ return
728
+
729
+ path_body_map[path] = body_id
730
+ path_shape_map[path] = shape_id
731
+ path_shape_scale[path] = scale
732
+
733
+ if prim.HasRelationship("physics:filteredPairs"):
734
+ other_paths = prim.GetRelationship("physics:filteredPairs").GetTargets()
735
+ for other_path in other_paths:
736
+ path_collision_filters.add((path, str(other_path)))
737
+
738
+ if "PhysicsCollisionAPI" not in schemas or not parse_generic(prim, "physics:collisionEnabled", True):
739
+ no_collision_shapes.add(shape_id)
740
+
741
+ else:
742
+ print(f"Warning: encountered unsupported prim type {type_name}")
743
+
744
+ # update mass properties of rigid bodies in cases where properties are defined with higher precedence
745
+ if body_id >= 0:
746
+ com = parse_vec(prim, "physics:centerOfMass")
747
+ if com is not None:
748
+ # overwrite COM
749
+ builder.body_com[body_id] = com * scale
750
+
751
+ if mass is not None and not (is_rigid_body and mass == 0.0):
752
+ mass_ratio = mass / builder.body_mass[body_id]
753
+ # mass has precedence over density, so we overwrite the mass computed from density
754
+ builder.body_mass[body_id] = mass * mass_unit
755
+ if mass > 0.0:
756
+ builder.body_inv_mass[body_id] = 1.0 / builder.body_mass[body_id]
757
+ else:
758
+ builder.body_inv_mass[body_id] = 0.0
759
+ # update inertia
760
+ builder.body_inertia[body_id] *= mass_ratio
761
+ if np.array(builder.body_inertia[body_id]).any():
762
+ builder.body_inv_inertia[body_id] = wp.inverse(builder.body_inertia[body_id])
763
+ else:
764
+ builder.body_inv_inertia[body_id] = wp.mat33(*np.zeros((3, 3), dtype=np.float32))
765
+
766
+ if np.linalg.norm(i_diag) > 0.0:
767
+ rot = np.array(wp.quat_to_matrix(i_rot), dtype=np.float32).reshape(3, 3)
768
+ inertia = rot @ np.diag(i_diag) @ rot.T
769
+ builder.body_inertia[body_id] = inertia
770
+ if inertia.any():
771
+ builder.body_inv_inertia[body_id] = wp.inverse(wp.mat33(*inertia))
772
+ else:
773
+ builder.body_inv_inertia[body_id] = wp.mat33(*np.zeros((3, 3), dtype=np.float32))
774
+
775
+ parse_prim(
776
+ stage.GetDefaultPrim(), incoming_xform=wp.transform(), incoming_scale=np.ones(3, dtype=np.float32) * linear_unit
777
+ )
778
+
779
+ shape_count = len(builder.shape_geo_type)
780
+
781
+ # apply collision filters now that we have added all shapes
782
+ for path1, path2 in path_collision_filters:
783
+ shape1 = path_shape_map[path1]
784
+ shape2 = path_shape_map[path2]
785
+ builder.shape_collision_filter_pairs.add((shape1, shape2))
786
+
787
+ # apply collision filters to all shapes that have no collision
788
+ for shape_id in no_collision_shapes:
789
+ for other_shape_id in range(shape_count):
790
+ if other_shape_id != shape_id:
791
+ builder.shape_collision_filter_pairs.add((shape_id, other_shape_id))
792
+
793
+ # return stage parameters
794
+ return {
795
+ "fps": stage.GetFramesPerSecond(),
796
+ "duration": stage.GetEndTimeCode() - stage.GetStartTimeCode(),
797
+ "up_axis": UsdGeom.GetStageUpAxis(stage).upper(),
798
+ "path_shape_map": path_shape_map,
799
+ "path_body_map": path_body_map,
800
+ "path_shape_scale": path_shape_scale,
801
+ "mass_unit": mass_unit,
802
+ "linear_unit": linear_unit,
803
+ }
804
+
805
+
806
+ def resolve_usd_from_url(url: str, target_folder_name: str = None, export_usda: bool = False):
807
+ """
808
+ Downloads a USD file from a URL and resolves all references to other USD files to be downloaded to the given target folder.
809
+
810
+ Args:
811
+ url (str): URL to the USD file.
812
+ target_folder_name (str): Target folder name. If None, a timestamped folder will be created in the current directory.
813
+ 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.
814
+
815
+ Returns:
816
+ str: File path to the downloaded USD file.
817
+ """
818
+ import datetime
819
+ import os
820
+
821
+ import requests
822
+
823
+ try:
824
+ from pxr import Usd
825
+ except ImportError as e:
826
+ raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).") from e
827
+
828
+ response = requests.get(url, allow_redirects=True)
829
+ if response.status_code != 200:
830
+ raise RuntimeError(f"Failed to download USD file. Status code: {response.status_code}")
831
+ file = response.content
832
+ dot = os.path.extsep
833
+ base = os.path.basename(url)
834
+ url_folder = os.path.dirname(url)
835
+ base_name = dot.join(base.split(dot)[:-1])
836
+ if target_folder_name is None:
837
+ timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
838
+ target_folder_name = os.path.join(".usd_cache", f"{base_name}_{timestamp}")
839
+ os.makedirs(target_folder_name, exist_ok=True)
840
+ target_filename = os.path.join(target_folder_name, base)
841
+ with open(target_filename, "wb") as f:
842
+ f.write(file)
843
+
844
+ stage = Usd.Stage.Open(target_filename, Usd.Stage.LoadNone)
845
+ stage_str = stage.GetRootLayer().ExportToString()
846
+ print(f"Downloaded USD file to {target_filename}.")
847
+ if export_usda:
848
+ usda_filename = os.path.join(target_folder_name, base_name + ".usda")
849
+ with open(usda_filename, "w") as f:
850
+ f.write(stage_str)
851
+ print(f"Exported USDA file to {usda_filename}.")
852
+
853
+ # parse referenced USD files like `references = @./franka_collisions.usd@`
854
+ downloaded = set()
855
+ for match in re.finditer(r"references.=.@(.*?)@", stage_str):
856
+ refname = match.group(1)
857
+ if refname.startswith("./"):
858
+ refname = refname[2:]
859
+ if refname in downloaded:
860
+ continue
861
+ try:
862
+ response = requests.get(f"{url_folder}/{refname}", allow_redirects=True)
863
+ if response.status_code != 200:
864
+ print(f"Failed to download reference {refname}. Status code: {response.status_code}")
865
+ continue
866
+ file = response.content
867
+ refdir = os.path.dirname(refname)
868
+ if refdir:
869
+ os.makedirs(os.path.join(target_folder_name, refdir), exist_ok=True)
870
+ ref_filename = os.path.join(target_folder_name, refname)
871
+ if not os.path.exists(ref_filename):
872
+ with open(ref_filename, "wb") as f:
873
+ f.write(file)
874
+ downloaded.add(refname)
875
+ print(f"Downloaded USD reference {refname} to {ref_filename}.")
876
+ if export_usda:
877
+ ref_stage = Usd.Stage.Open(ref_filename, Usd.Stage.LoadNone)
878
+ ref_stage_str = ref_stage.GetRootLayer().ExportToString()
879
+ base = os.path.basename(ref_filename)
880
+ base_name = dot.join(base.split(dot)[:-1])
881
+ usda_filename = os.path.join(target_folder_name, base_name + ".usda")
882
+ with open(usda_filename, "w") as f:
883
+ f.write(ref_stage_str)
884
+ print(f"Exported USDA file to {usda_filename}.")
885
+ except Exception:
886
+ print(f"Failed to download {refname}.")
887
+ return target_filename