warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.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 (271) hide show
  1. docs/conf.py +17 -5
  2. examples/env/env_ant.py +1 -1
  3. examples/env/env_cartpole.py +1 -1
  4. examples/env/env_humanoid.py +1 -1
  5. examples/env/env_usd.py +4 -1
  6. examples/env/environment.py +8 -9
  7. examples/example_dem.py +34 -33
  8. examples/example_diffray.py +364 -337
  9. examples/example_fluid.py +32 -23
  10. examples/example_jacobian_ik.py +97 -93
  11. examples/example_marching_cubes.py +6 -16
  12. examples/example_mesh.py +6 -16
  13. examples/example_mesh_intersect.py +16 -14
  14. examples/example_nvdb.py +14 -16
  15. examples/example_raycast.py +14 -13
  16. examples/example_raymarch.py +16 -23
  17. examples/example_render_opengl.py +19 -10
  18. examples/example_sim_cartpole.py +82 -78
  19. examples/example_sim_cloth.py +45 -48
  20. examples/example_sim_fk_grad.py +51 -44
  21. examples/example_sim_fk_grad_torch.py +47 -40
  22. examples/example_sim_grad_bounce.py +108 -133
  23. examples/example_sim_grad_cloth.py +99 -113
  24. examples/example_sim_granular.py +5 -6
  25. examples/{example_sim_sdf_shape.py → example_sim_granular_collision_sdf.py} +37 -26
  26. examples/example_sim_neo_hookean.py +51 -55
  27. examples/example_sim_particle_chain.py +4 -4
  28. examples/example_sim_quadruped.py +126 -81
  29. examples/example_sim_rigid_chain.py +54 -61
  30. examples/example_sim_rigid_contact.py +66 -70
  31. examples/example_sim_rigid_fem.py +3 -3
  32. examples/example_sim_rigid_force.py +1 -1
  33. examples/example_sim_rigid_gyroscopic.py +3 -4
  34. examples/example_sim_rigid_kinematics.py +28 -39
  35. examples/example_sim_trajopt.py +112 -110
  36. examples/example_sph.py +9 -8
  37. examples/example_wave.py +7 -7
  38. examples/fem/bsr_utils.py +30 -17
  39. examples/fem/example_apic_fluid.py +85 -69
  40. examples/fem/example_convection_diffusion.py +97 -93
  41. examples/fem/example_convection_diffusion_dg.py +142 -149
  42. examples/fem/example_convection_diffusion_dg0.py +141 -136
  43. examples/fem/example_deformed_geometry.py +146 -0
  44. examples/fem/example_diffusion.py +115 -84
  45. examples/fem/example_diffusion_3d.py +116 -86
  46. examples/fem/example_diffusion_mgpu.py +102 -79
  47. examples/fem/example_mixed_elasticity.py +139 -100
  48. examples/fem/example_navier_stokes.py +175 -162
  49. examples/fem/example_stokes.py +143 -111
  50. examples/fem/example_stokes_transfer.py +186 -157
  51. examples/fem/mesh_utils.py +59 -97
  52. examples/fem/plot_utils.py +138 -17
  53. tools/ci/publishing/build_nodes_info.py +54 -0
  54. warp/__init__.py +4 -3
  55. warp/__init__.pyi +1 -0
  56. warp/bin/warp-clang.dll +0 -0
  57. warp/bin/warp.dll +0 -0
  58. warp/build.py +5 -3
  59. warp/build_dll.py +29 -9
  60. warp/builtins.py +836 -492
  61. warp/codegen.py +864 -553
  62. warp/config.py +3 -1
  63. warp/context.py +389 -172
  64. warp/fem/__init__.py +24 -6
  65. warp/fem/cache.py +318 -25
  66. warp/fem/dirichlet.py +7 -3
  67. warp/fem/domain.py +14 -0
  68. warp/fem/field/__init__.py +30 -38
  69. warp/fem/field/field.py +149 -0
  70. warp/fem/field/nodal_field.py +244 -138
  71. warp/fem/field/restriction.py +8 -6
  72. warp/fem/field/test.py +127 -59
  73. warp/fem/field/trial.py +117 -60
  74. warp/fem/geometry/__init__.py +5 -1
  75. warp/fem/geometry/deformed_geometry.py +271 -0
  76. warp/fem/geometry/element.py +24 -1
  77. warp/fem/geometry/geometry.py +86 -14
  78. warp/fem/geometry/grid_2d.py +112 -54
  79. warp/fem/geometry/grid_3d.py +134 -65
  80. warp/fem/geometry/hexmesh.py +953 -0
  81. warp/fem/geometry/partition.py +85 -33
  82. warp/fem/geometry/quadmesh_2d.py +532 -0
  83. warp/fem/geometry/tetmesh.py +451 -115
  84. warp/fem/geometry/trimesh_2d.py +197 -92
  85. warp/fem/integrate.py +534 -268
  86. warp/fem/operator.py +58 -31
  87. warp/fem/polynomial.py +11 -0
  88. warp/fem/quadrature/__init__.py +1 -1
  89. warp/fem/quadrature/pic_quadrature.py +150 -58
  90. warp/fem/quadrature/quadrature.py +209 -57
  91. warp/fem/space/__init__.py +230 -53
  92. warp/fem/space/basis_space.py +489 -0
  93. warp/fem/space/collocated_function_space.py +105 -0
  94. warp/fem/space/dof_mapper.py +49 -2
  95. warp/fem/space/function_space.py +90 -39
  96. warp/fem/space/grid_2d_function_space.py +149 -496
  97. warp/fem/space/grid_3d_function_space.py +173 -538
  98. warp/fem/space/hexmesh_function_space.py +352 -0
  99. warp/fem/space/partition.py +129 -76
  100. warp/fem/space/quadmesh_2d_function_space.py +369 -0
  101. warp/fem/space/restriction.py +46 -34
  102. warp/fem/space/shape/__init__.py +15 -0
  103. warp/fem/space/shape/cube_shape_function.py +738 -0
  104. warp/fem/space/shape/shape_function.py +103 -0
  105. warp/fem/space/shape/square_shape_function.py +611 -0
  106. warp/fem/space/shape/tet_shape_function.py +567 -0
  107. warp/fem/space/shape/triangle_shape_function.py +429 -0
  108. warp/fem/space/tetmesh_function_space.py +132 -1039
  109. warp/fem/space/topology.py +295 -0
  110. warp/fem/space/trimesh_2d_function_space.py +104 -742
  111. warp/fem/types.py +13 -11
  112. warp/fem/utils.py +335 -60
  113. warp/native/array.h +120 -34
  114. warp/native/builtin.h +101 -72
  115. warp/native/bvh.cpp +73 -325
  116. warp/native/bvh.cu +406 -23
  117. warp/native/bvh.h +22 -40
  118. warp/native/clang/clang.cpp +1 -0
  119. warp/native/crt.h +2 -0
  120. warp/native/cuda_util.cpp +8 -3
  121. warp/native/cuda_util.h +1 -0
  122. warp/native/exports.h +1522 -1243
  123. warp/native/intersect.h +19 -4
  124. warp/native/intersect_adj.h +8 -8
  125. warp/native/mat.h +76 -17
  126. warp/native/mesh.cpp +33 -108
  127. warp/native/mesh.cu +114 -18
  128. warp/native/mesh.h +395 -40
  129. warp/native/noise.h +272 -329
  130. warp/native/quat.h +51 -8
  131. warp/native/rand.h +44 -34
  132. warp/native/reduce.cpp +1 -1
  133. warp/native/sparse.cpp +4 -4
  134. warp/native/sparse.cu +163 -155
  135. warp/native/spatial.h +2 -2
  136. warp/native/temp_buffer.h +18 -14
  137. warp/native/vec.h +103 -21
  138. warp/native/warp.cpp +2 -1
  139. warp/native/warp.cu +28 -3
  140. warp/native/warp.h +4 -3
  141. warp/render/render_opengl.py +261 -109
  142. warp/sim/__init__.py +1 -2
  143. warp/sim/articulation.py +385 -185
  144. warp/sim/import_mjcf.py +59 -48
  145. warp/sim/import_urdf.py +15 -15
  146. warp/sim/import_usd.py +174 -102
  147. warp/sim/inertia.py +17 -18
  148. warp/sim/integrator_xpbd.py +4 -3
  149. warp/sim/model.py +330 -250
  150. warp/sim/render.py +1 -1
  151. warp/sparse.py +625 -152
  152. warp/stubs.py +341 -309
  153. warp/tape.py +9 -6
  154. warp/tests/__main__.py +3 -6
  155. warp/tests/assets/curlnoise_golden.npy +0 -0
  156. warp/tests/assets/pnoise_golden.npy +0 -0
  157. warp/tests/{test_class_kernel.py → aux_test_class_kernel.py} +9 -1
  158. warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -0
  159. warp/tests/{test_dependent.py → aux_test_dependent.py} +2 -2
  160. warp/tests/{test_reference.py → aux_test_reference.py} +1 -1
  161. warp/tests/aux_test_unresolved_func.py +14 -0
  162. warp/tests/aux_test_unresolved_symbol.py +14 -0
  163. warp/tests/disabled_kinematics.py +239 -0
  164. warp/tests/run_coverage_serial.py +31 -0
  165. warp/tests/test_adam.py +103 -106
  166. warp/tests/test_arithmetic.py +94 -74
  167. warp/tests/test_array.py +82 -101
  168. warp/tests/test_array_reduce.py +57 -23
  169. warp/tests/test_atomic.py +64 -28
  170. warp/tests/test_bool.py +22 -12
  171. warp/tests/test_builtins_resolution.py +1292 -0
  172. warp/tests/test_bvh.py +18 -18
  173. warp/tests/test_closest_point_edge_edge.py +54 -57
  174. warp/tests/test_codegen.py +165 -134
  175. warp/tests/test_compile_consts.py +28 -20
  176. warp/tests/test_conditional.py +108 -24
  177. warp/tests/test_copy.py +10 -12
  178. warp/tests/test_ctypes.py +112 -88
  179. warp/tests/test_dense.py +21 -14
  180. warp/tests/test_devices.py +98 -0
  181. warp/tests/test_dlpack.py +75 -75
  182. warp/tests/test_examples.py +237 -0
  183. warp/tests/test_fabricarray.py +22 -24
  184. warp/tests/test_fast_math.py +15 -11
  185. warp/tests/test_fem.py +1034 -124
  186. warp/tests/test_fp16.py +23 -16
  187. warp/tests/test_func.py +187 -86
  188. warp/tests/test_generics.py +194 -49
  189. warp/tests/test_grad.py +123 -181
  190. warp/tests/test_grad_customs.py +176 -0
  191. warp/tests/test_hash_grid.py +35 -34
  192. warp/tests/test_import.py +10 -23
  193. warp/tests/test_indexedarray.py +24 -25
  194. warp/tests/test_intersect.py +18 -9
  195. warp/tests/test_large.py +141 -0
  196. warp/tests/test_launch.py +14 -41
  197. warp/tests/test_lerp.py +64 -65
  198. warp/tests/test_lvalue.py +493 -0
  199. warp/tests/test_marching_cubes.py +12 -13
  200. warp/tests/test_mat.py +517 -2898
  201. warp/tests/test_mat_lite.py +115 -0
  202. warp/tests/test_mat_scalar_ops.py +2889 -0
  203. warp/tests/test_math.py +103 -9
  204. warp/tests/test_matmul.py +304 -69
  205. warp/tests/test_matmul_lite.py +410 -0
  206. warp/tests/test_mesh.py +60 -22
  207. warp/tests/test_mesh_query_aabb.py +21 -25
  208. warp/tests/test_mesh_query_point.py +111 -22
  209. warp/tests/test_mesh_query_ray.py +12 -24
  210. warp/tests/test_mlp.py +30 -22
  211. warp/tests/test_model.py +92 -89
  212. warp/tests/test_modules_lite.py +39 -0
  213. warp/tests/test_multigpu.py +88 -114
  214. warp/tests/test_noise.py +12 -11
  215. warp/tests/test_operators.py +16 -20
  216. warp/tests/test_options.py +11 -11
  217. warp/tests/test_pinned.py +17 -18
  218. warp/tests/test_print.py +32 -11
  219. warp/tests/test_quat.py +275 -129
  220. warp/tests/test_rand.py +18 -16
  221. warp/tests/test_reload.py +38 -34
  222. warp/tests/test_rounding.py +50 -43
  223. warp/tests/test_runlength_encode.py +168 -20
  224. warp/tests/test_smoothstep.py +9 -11
  225. warp/tests/test_snippet.py +143 -0
  226. warp/tests/test_sparse.py +261 -63
  227. warp/tests/test_spatial.py +276 -243
  228. warp/tests/test_streams.py +110 -85
  229. warp/tests/test_struct.py +268 -63
  230. warp/tests/test_tape.py +39 -21
  231. warp/tests/test_torch.py +90 -86
  232. warp/tests/test_transient_module.py +10 -12
  233. warp/tests/test_types.py +363 -0
  234. warp/tests/test_utils.py +451 -0
  235. warp/tests/test_vec.py +354 -2050
  236. warp/tests/test_vec_lite.py +73 -0
  237. warp/tests/test_vec_scalar_ops.py +2099 -0
  238. warp/tests/test_volume.py +418 -376
  239. warp/tests/test_volume_write.py +124 -134
  240. warp/tests/unittest_serial.py +35 -0
  241. warp/tests/unittest_suites.py +291 -0
  242. warp/tests/unittest_utils.py +342 -0
  243. warp/tests/{test_misc.py → unused_test_misc.py} +13 -5
  244. warp/tests/{test_debug.py → walkthough_debug.py} +3 -17
  245. warp/thirdparty/appdirs.py +36 -45
  246. warp/thirdparty/unittest_parallel.py +589 -0
  247. warp/types.py +622 -211
  248. warp/utils.py +54 -393
  249. warp_lang-1.0.0b6.dist-info/METADATA +238 -0
  250. warp_lang-1.0.0b6.dist-info/RECORD +409 -0
  251. {warp_lang-1.0.0b2.dist-info → warp_lang-1.0.0b6.dist-info}/WHEEL +1 -1
  252. examples/example_cache_management.py +0 -40
  253. examples/example_multigpu.py +0 -54
  254. examples/example_struct.py +0 -65
  255. examples/fem/example_stokes_transfer_3d.py +0 -210
  256. warp/bin/warp-clang.so +0 -0
  257. warp/bin/warp.so +0 -0
  258. warp/fem/field/discrete_field.py +0 -80
  259. warp/fem/space/nodal_function_space.py +0 -233
  260. warp/tests/test_all.py +0 -223
  261. warp/tests/test_array_scan.py +0 -60
  262. warp/tests/test_base.py +0 -208
  263. warp/tests/test_unresolved_func.py +0 -7
  264. warp/tests/test_unresolved_symbol.py +0 -7
  265. warp_lang-1.0.0b2.dist-info/METADATA +0 -26
  266. warp_lang-1.0.0b2.dist-info/RECORD +0 -380
  267. /warp/tests/{test_compile_consts_dummy.py → aux_test_compile_consts_dummy.py} +0 -0
  268. /warp/tests/{test_reference_reference.py → aux_test_reference_reference.py} +0 -0
  269. /warp/tests/{test_square.py → aux_test_square.py} +0 -0
  270. {warp_lang-1.0.0b2.dist-info → warp_lang-1.0.0b6.dist-info}/LICENSE.md +0 -0
  271. {warp_lang-1.0.0b2.dist-info → warp_lang-1.0.0b6.dist-info}/top_level.txt +0 -0
warp/sim/import_mjcf.py CHANGED
@@ -44,7 +44,7 @@ def parse_mjcf(
44
44
  Args:
45
45
  mjcf_filename (str): The filename of the MuJoCo file to parse.
46
46
  builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
47
- xform (wp.transform): The transform to apply to the imported mechanism.
47
+ xform (:ref:`transform <transform>`): The transform to apply to the imported mechanism.
48
48
  density (float): The density of the shapes in kg/m^3 which will be used to calculate the body mass and inertia.
49
49
  stiffness (float): The stiffness of the joints.
50
50
  damping (float): The damping of the joints.
@@ -150,9 +150,15 @@ def parse_mjcf(
150
150
 
151
151
  def parse_vec(attrib, key, default):
152
152
  if key in attrib:
153
- return np.fromstring(attrib[key], sep=" ")
153
+ out = np.fromstring(attrib[key], sep=" ", dtype=np.float32)
154
154
  else:
155
- return np.array(default)
155
+ out = np.array(default, dtype=np.float32)
156
+
157
+ length = len(out)
158
+ if length == 1:
159
+ return wp.vec(len(default), wp.float32)(out[0], out[0], out[0])
160
+
161
+ return wp.vec(length, wp.float32)(out)
156
162
 
157
163
  def parse_orientation(attrib):
158
164
  if "quat" in attrib:
@@ -232,46 +238,51 @@ def parse_mjcf(
232
238
  angular_axes = []
233
239
  joint_type = None
234
240
 
235
- joints = body.findall("joint")
236
- for i, joint in enumerate(joints):
237
- if "joint" in defaults:
238
- joint_attrib = merge_attrib(defaults["joint"], joint.attrib)
239
- else:
240
- joint_attrib = joint.attrib
241
-
242
- # default to hinge if not specified
243
- joint_type_str = joint_attrib.get("type", "hinge")
244
-
245
- joint_name.append(joint_attrib["name"])
246
- joint_pos.append(parse_vec(joint_attrib, "pos", (0.0, 0.0, 0.0)) * scale)
247
- joint_range = parse_vec(joint_attrib, "range", (-3.0, 3.0))
248
- joint_armature.append(parse_float(joint_attrib, "armature", armature) * armature_scale)
249
-
250
- if joint_type_str == "free":
251
- joint_type = wp.sim.JOINT_FREE
252
- break
253
- if joint_type_str == "fixed":
254
- joint_type = wp.sim.JOINT_FIXED
255
- break
256
- is_angular = joint_type_str == "hinge"
257
- mode = wp.sim.JOINT_MODE_LIMIT
258
- if stiffness > 0.0 or "stiffness" in joint_attrib:
259
- mode = wp.sim.JOINT_MODE_TARGET_POSITION
260
- axis_vec = parse_vec(joint_attrib, "axis", (0.0, 0.0, 0.0))
261
- ax = wp.sim.model.JointAxis(
262
- axis=axis_vec,
263
- limit_lower=(np.deg2rad(joint_range[0]) if is_angular and use_degrees else joint_range[0]),
264
- limit_upper=(np.deg2rad(joint_range[1]) if is_angular and use_degrees else joint_range[1]),
265
- target_ke=parse_float(joint_attrib, "stiffness", stiffness),
266
- target_kd=parse_float(joint_attrib, "damping", damping),
267
- limit_ke=limit_ke,
268
- limit_kd=limit_kd,
269
- mode=mode,
270
- )
271
- if is_angular:
272
- angular_axes.append(ax)
273
- else:
274
- linear_axes.append(ax)
241
+ freejoint_tags = body.findall("freejoint")
242
+ if len(freejoint_tags) > 0:
243
+ joint_type = wp.sim.JOINT_FREE
244
+ joint_name.append(freejoint_tags[0].attrib.get("name", f"{body_name}_freejoint"))
245
+ else:
246
+ joints = body.findall("joint")
247
+ for i, joint in enumerate(joints):
248
+ if "joint" in defaults:
249
+ joint_attrib = merge_attrib(defaults["joint"], joint.attrib)
250
+ else:
251
+ joint_attrib = joint.attrib
252
+
253
+ # default to hinge if not specified
254
+ joint_type_str = joint_attrib.get("type", "hinge")
255
+
256
+ joint_name.append(joint_attrib["name"])
257
+ joint_pos.append(parse_vec(joint_attrib, "pos", (0.0, 0.0, 0.0)) * scale)
258
+ joint_range = parse_vec(joint_attrib, "range", (-3.0, 3.0))
259
+ joint_armature.append(parse_float(joint_attrib, "armature", armature) * armature_scale)
260
+
261
+ if joint_type_str == "free":
262
+ joint_type = wp.sim.JOINT_FREE
263
+ break
264
+ if joint_type_str == "fixed":
265
+ joint_type = wp.sim.JOINT_FIXED
266
+ break
267
+ is_angular = joint_type_str == "hinge"
268
+ mode = wp.sim.JOINT_MODE_LIMIT
269
+ if stiffness > 0.0 or "stiffness" in joint_attrib:
270
+ mode = wp.sim.JOINT_MODE_TARGET_POSITION
271
+ axis_vec = parse_vec(joint_attrib, "axis", (0.0, 0.0, 0.0))
272
+ ax = wp.sim.model.JointAxis(
273
+ axis=axis_vec,
274
+ limit_lower=(np.deg2rad(joint_range[0]) if is_angular and use_degrees else joint_range[0]),
275
+ limit_upper=(np.deg2rad(joint_range[1]) if is_angular and use_degrees else joint_range[1]),
276
+ target_ke=parse_float(joint_attrib, "stiffness", stiffness),
277
+ target_kd=parse_float(joint_attrib, "damping", damping),
278
+ limit_ke=limit_ke,
279
+ limit_kd=limit_kd,
280
+ mode=mode,
281
+ )
282
+ if is_angular:
283
+ angular_axes.append(ax)
284
+ else:
285
+ linear_axes.append(ax)
275
286
 
276
287
  link = builder.add_body(
277
288
  origin=wp.transform(body_pos, body_ori), # will be evaluated in fk()
@@ -392,19 +403,19 @@ def parse_mjcf(
392
403
  if "fromto" in geom_attrib:
393
404
  geom_fromto = parse_vec(geom_attrib, "fromto", (0.0, 0.0, 0.0, 1.0, 0.0, 0.0))
394
405
 
395
- start = geom_fromto[0:3] * scale
396
- end = geom_fromto[3:6] * scale
406
+ start = wp.vec3(geom_fromto[0:3]) * scale
407
+ end = wp.vec3(geom_fromto[3:6]) * scale
397
408
 
398
409
  # compute rotation to align the Warp capsule (along x-axis), with mjcf fromto direction
399
410
  axis = wp.normalize(end - start)
400
- angle = math.acos(np.dot(axis, (0.0, 1.0, 0.0)))
401
- axis = wp.normalize(np.cross(axis, (0.0, 1.0, 0.0)))
411
+ angle = math.acos(wp.dot(axis, wp.vec3(0.0, 1.0, 0.0)))
412
+ axis = wp.normalize(wp.cross(axis, wp.vec3(0.0, 1.0, 0.0)))
402
413
 
403
414
  geom_pos = (start + end) * 0.5
404
415
  geom_rot = wp.quat_from_axis_angle(axis, -angle)
405
416
 
406
417
  geom_radius = geom_size[0]
407
- geom_height = np.linalg.norm(end - start) * 0.5
418
+ geom_height = wp.length(end - start) * 0.5
408
419
  geom_up_axis = 1
409
420
 
410
421
  else:
warp/sim/import_urdf.py CHANGED
@@ -47,7 +47,7 @@ def parse_urdf(
47
47
  Args:
48
48
  urdf_filename (str): The filename of the URDF file to parse.
49
49
  builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
50
- xform (wp.transform): The transform to apply to the root body.
50
+ xform (:ref:`transform <transform>`): The transform to apply to the root body.
51
51
  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.
52
52
  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`).
53
53
  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(
100
100
  size = [float(x) for x in size.split()]
101
101
  builder.add_shape_box(
102
102
  body=link,
103
- pos=tf.p,
104
- rot=tf.q,
103
+ pos=wp.vec3(tf.p),
104
+ rot=wp.quat(tf.q),
105
105
  hx=size[0] * 0.5 * scale,
106
106
  hy=size[1] * 0.5 * scale,
107
107
  hz=size[2] * 0.5 * scale,
@@ -117,8 +117,8 @@ def parse_urdf(
117
117
  for sphere in geo.findall("sphere"):
118
118
  builder.add_shape_sphere(
119
119
  body=link,
120
- pos=tf.p,
121
- rot=tf.q,
120
+ pos=wp.vec3(tf.p),
121
+ rot=wp.quat(tf.q),
122
122
  radius=float(sphere.get("radius") or "1") * scale,
123
123
  density=density,
124
124
  ke=shape_ke,
@@ -132,8 +132,8 @@ def parse_urdf(
132
132
  for cylinder in geo.findall("cylinder"):
133
133
  builder.add_shape_capsule(
134
134
  body=link,
135
- pos=tf.p,
136
- rot=tf.q,
135
+ pos=wp.vec3(tf.p),
136
+ rot=wp.quat(tf.q),
137
137
  radius=float(cylinder.get("radius") or "1") * scale,
138
138
  half_height=float(cylinder.get("length") or "1") * 0.5 * scale,
139
139
  density=density,
@@ -180,12 +180,12 @@ def parse_urdf(
180
180
  # multiple meshes are contained in a scene
181
181
  for geom in m.geometry.values():
182
182
  vertices = np.array(geom.vertices, dtype=np.float32) * scaling
183
- faces = np.array(geom.faces, dtype=np.int32)
183
+ faces = np.array(geom.faces.flatten(), dtype=np.int32)
184
184
  mesh = Mesh(vertices, faces)
185
185
  builder.add_shape_mesh(
186
186
  body=link,
187
- pos=tf.p,
188
- rot=tf.q,
187
+ pos=wp.vec3(tf.p),
188
+ rot=wp.quat(tf.q),
189
189
  mesh=mesh,
190
190
  density=density,
191
191
  ke=shape_ke,
@@ -198,7 +198,7 @@ def parse_urdf(
198
198
  else:
199
199
  # a single mesh
200
200
  vertices = np.array(m.vertices, dtype=np.float32) * scaling
201
- faces = np.array(m.faces, dtype=np.int32)
201
+ faces = np.array(m.faces.flatten(), dtype=np.int32)
202
202
  mesh = Mesh(vertices, faces)
203
203
  builder.add_shape_mesh(
204
204
  body=link,
@@ -252,22 +252,22 @@ def parse_urdf(
252
252
  I_m[2, 0] = I_m[0, 2]
253
253
  I_m[2, 1] = I_m[1, 2]
254
254
  rot = wp.quat_to_matrix(inertial_frame.q)
255
- I_m = rot @ I_m
255
+ I_m = rot @ wp.mat33(I_m)
256
256
  m = float(inertial.find("mass").get("value") or "0")
257
257
  builder.body_mass[link] = m
258
258
  builder.body_inv_mass[link] = 1.0 / m
259
259
  builder.body_com[link] = com
260
260
  builder.body_inertia[link] = I_m
261
- builder.body_inv_inertia[link] = np.linalg.inv(I_m)
261
+ builder.body_inv_inertia[link] = wp.inverse(I_m)
262
262
  if m == 0.0 and ensure_nonstatic_links:
263
263
  # set the mass to something nonzero to ensure the body is dynamic
264
264
  m = static_link_mass
265
265
  # cube with side length 0.5
266
- I_m = np.eye(3) * m / 12.0 * (0.5 * scale) ** 2 * 2.0
266
+ I_m = wp.mat33(np.eye(3)) * m / 12.0 * (0.5 * scale) ** 2 * 2.0
267
267
  builder.body_mass[link] = m
268
268
  builder.body_inv_mass[link] = 1.0 / m
269
269
  builder.body_inertia[link] = I_m
270
- builder.body_inv_inertia[link] = np.linalg.inv(I_m)
270
+ builder.body_inv_inertia[link] = wp.inverse(I_m)
271
271
 
272
272
  end_shape_count = len(builder.shape_geo_type)
273
273
 
warp/sim/import_usd.py CHANGED
@@ -10,32 +10,33 @@ import numpy as np
10
10
  import re
11
11
 
12
12
  import warp as wp
13
- from . import ModelBuilder
14
13
 
15
14
 
16
15
  def parse_usd(
17
- filename,
18
- builder: ModelBuilder,
16
+ source,
17
+ builder,
19
18
  default_density=1.0e3,
20
19
  only_load_enabled_rigid_bodies=False,
21
20
  only_load_enabled_joints=True,
22
21
  default_ke=1e5,
23
22
  default_kd=250.0,
24
23
  default_kf=500.0,
25
- default_mu=0.0,
24
+ default_mu=0.6,
26
25
  default_restitution=0.0,
27
26
  default_thickness=0.0,
28
27
  joint_limit_ke=100.0,
29
28
  joint_limit_kd=10.0,
29
+ invert_rotations=False,
30
30
  verbose=False,
31
31
  ignore_paths=[],
32
- export_usda=False,
33
32
  ):
34
33
  """
35
- Parses a USD file containing UsdPhysics schema definitions for rigid-body articulations and adds the bodies, shapes and joints to the given ModelBuilder.
34
+ 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.
35
+
36
+ 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.
36
37
 
37
38
  Args:
38
- filename (str): The file path or URL to the USD file to parse.
39
+ source (str | pxr.UsdStage): The file path to the USD file, or an existing USD stage instance.
39
40
  builder (ModelBuilder): The :class:`ModelBuilder` to add the bodies and joints to.
40
41
  default_density (float): The default density to use for bodies without a density attribute.
41
42
  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(
48
49
  default_thickness (float): The thickness to add to the shape geometry.
49
50
  joint_limit_ke (float): The default stiffness to use for joint limits, only considered by SemiImplicitIntegrator.
50
51
  joint_limit_kd (float): The default damping to use for joint limits, only considered by SemiImplicitIntegrator.
52
+ invert_rotations (bool): If True, inverts any rotations defined in the shape transforms.
51
53
  verbose (bool): If True, print additional information about the parsed USD file.
52
54
  ignore_paths (List[str]): A list of regular expressions matching prim paths to ignore.
53
- export_usda (bool): If True and the filename is a URL, export the downloaded USD file to a USDA file.
54
55
 
55
56
  Returns:
56
- 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.
57
+ dict: Dictionary with the following entries:
58
+
59
+ .. list-table::
60
+ :widths: 25 75
61
+
62
+ * - "fps"
63
+ - USD stage frames per second
64
+ * - "duration"
65
+ - Difference between end time code and start time code of the USD stage
66
+ * - "up_axis"
67
+ - Upper-case string of the stage's up axis ("X", "Y", or "Z")
68
+ * - "path_shape_map"
69
+ - Mapping from prim path (str) of the UsdGeom to the respective shape index in :class:`ModelBuilder`
70
+ * - "path_body_map"
71
+ - Mapping from prim path (str) of a rigid body prim (e.g. that implements the PhysicsRigidBodyAPI) to the respective body index in :class:`ModelBuilder`
72
+ * - "path_shape_scale"
73
+ - Mapping from prim path (str) of the UsdGeom to its respective 3D world scale
74
+ * - "mass_unit"
75
+ - The stage's Kilograms Per Unit (KGPU) definition (1.0 by default)
76
+ * - "linear_unit"
77
+ - The stage's Meters Per Unit (MPU) definition (1.0 by default)
78
+
57
79
 
58
80
  Note:
59
81
  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(
61
83
  try:
62
84
  from pxr import Usd, UsdGeom, UsdPhysics
63
85
  except ImportError:
64
- raise ImportError("Failed to import pxr. Please install USD.")
65
-
66
- if filename.startswith("http://") or filename.startswith("https://"):
67
- # download file
68
- import requests
69
- import os
70
- import datetime
71
-
72
- response = requests.get(filename, allow_redirects=True)
73
- if response.status_code != 200:
74
- raise RuntimeError(f"Failed to download USD file. Status code: {response.status_code}")
75
- file = response.content
76
- dot = os.path.extsep
77
- base = os.path.basename(filename)
78
- url_folder = os.path.dirname(filename)
79
- base_name = dot.join(base.split(dot)[:-1])
80
- timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
81
- folder_name = os.path.join(".usd_cache", f"{base_name}_{timestamp}")
82
- os.makedirs(folder_name, exist_ok=True)
83
- target_filename = os.path.join(folder_name, base)
84
- with open(target_filename, "wb") as f:
85
- f.write(file)
86
-
87
- stage = Usd.Stage.Open(target_filename, Usd.Stage.LoadNone)
88
- stage_str = stage.GetRootLayer().ExportToString()
89
- print(f"Downloaded USD file to {target_filename}.")
90
- if export_usda:
91
- usda_filename = os.path.join(folder_name, base_name + ".usda")
92
- with open(usda_filename, "w") as f:
93
- f.write(stage_str)
94
- print(f"Exported USDA file to {usda_filename}.")
95
-
96
- # parse referenced USD files like `references = @./franka_collisions.usd@`
97
- downloaded = set()
98
- for match in re.finditer(r"references.=.@(.*?)@", stage_str):
99
- refname = match.group(1)
100
- if refname.startswith("./"):
101
- refname = refname[2:]
102
- if refname in downloaded:
103
- continue
104
- try:
105
- response = requests.get(f"{url_folder}/{refname}", allow_redirects=True)
106
- if response.status_code != 200:
107
- print(f"Failed to download reference {refname}. Status code: {response.status_code}")
108
- continue
109
- file = response.content
110
- refdir = os.path.dirname(refname)
111
- if refdir:
112
- os.makedirs(os.path.join(folder_name, refdir), exist_ok=True)
113
- ref_filename = os.path.join(folder_name, refname)
114
- with open(ref_filename, "wb") as f:
115
- f.write(file)
116
- downloaded.add(refname)
117
- print(f"Downloaded USD reference {refname} to {ref_filename}.")
118
- if export_usda:
119
- ref_stage = Usd.Stage.Open(ref_filename, Usd.Stage.LoadNone)
120
- ref_stage_str = ref_stage.GetRootLayer().ExportToString()
121
- base = os.path.basename(ref_filename)
122
- base_name = dot.join(base.split(dot)[:-1])
123
- usda_filename = os.path.join(folder_name, base_name + ".usda")
124
- with open(usda_filename, "w") as f:
125
- f.write(ref_stage_str)
126
- print(f"Exported USDA file to {usda_filename}.")
127
- except Exception:
128
- print(f"Failed to download {refname}.")
129
-
130
- filename = target_filename
86
+ raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).")
131
87
 
132
88
  def get_attribute(prim, name):
133
89
  if "*" in name:
@@ -156,7 +112,10 @@ def parse_usd(
156
112
  if not attr or not attr.HasAuthoredValue():
157
113
  return default
158
114
  val = attr.Get()
159
- quat = wp.quat(*val.imaginary, val.real)
115
+ if invert_rotations:
116
+ quat = wp.quat(*val.imaginary, -val.real)
117
+ else:
118
+ quat = wp.quat(*val.imaginary, val.real)
160
119
  l = wp.length(quat)
161
120
  if np.isfinite(l) and l > 0.0:
162
121
  return quat
@@ -182,20 +141,31 @@ def parse_usd(
182
141
  axis["XYZ".index(s.upper())] = 1.0
183
142
  return axis
184
143
 
185
- stage = Usd.Stage.Open(filename, Usd.Stage.LoadAll)
186
- if UsdPhysics.StageHasAuthoredKilogramsPerUnit(stage):
187
- mass_unit = UsdPhysics.GetStageKilogramsPerUnit(stage)
188
- else:
189
- mass_unit = 1.0
190
- if UsdGeom.StageHasAuthoredMetersPerUnit(stage):
191
- linear_unit = UsdGeom.GetStageMetersPerUnit(stage)
144
+ if isinstance(source, str):
145
+ stage = Usd.Stage.Open(source, Usd.Stage.LoadAll)
192
146
  else:
193
- linear_unit = 1.0
147
+ stage = source
148
+
149
+ mass_unit = 1.0
150
+ try:
151
+ if UsdPhysics.StageHasAuthoredKilogramsPerUnit(stage):
152
+ mass_unit = UsdPhysics.GetStageKilogramsPerUnit(stage)
153
+ except:
154
+ pass
155
+ linear_unit = 1.0
156
+ try:
157
+ if UsdGeom.StageHasAuthoredMetersPerUnit(stage):
158
+ linear_unit = UsdGeom.GetStageMetersPerUnit(stage)
159
+ except:
160
+ pass
194
161
 
195
162
  def parse_xform(prim):
196
163
  xform = UsdGeom.Xform(prim)
197
164
  mat = np.array(xform.GetLocalTransformation(), dtype=np.float32)
198
- rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].flatten()))
165
+ if invert_rotations:
166
+ rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].T.flatten()))
167
+ else:
168
+ rot = wp.quat_from_matrix(wp.mat33(mat[:3, :3].flatten()))
199
169
  pos = mat[3, :3] * linear_unit
200
170
  scale = np.ones(3, dtype=np.float32)
201
171
  for op in xform.GetOrderedXformOps():
@@ -276,12 +246,18 @@ def parse_usd(
276
246
  else:
277
247
  joint_data["linear_axes"].append(axis)
278
248
 
279
- upaxis = str2axis(UsdGeom.GetStageUpAxis(stage))
249
+ axis_str = "Y"
250
+ try:
251
+ axis_str = UsdGeom.GetStageUpAxis(stage)
252
+ except:
253
+ pass
254
+ upaxis = str2axis(axis_str)
280
255
 
281
256
  shape_types = {"Cube", "Sphere", "Mesh", "Capsule", "Plane", "Cylinder", "Cone"}
282
257
 
283
258
  path_body_map = {}
284
259
  path_shape_map = {}
260
+ path_shape_scale = {}
285
261
  # maps prim path name to its world transform
286
262
  path_world_poses = {}
287
263
  # transform from body frame to where the actual joint child frame is
@@ -376,21 +352,27 @@ def parse_usd(
376
352
  materials[path] = material
377
353
 
378
354
  elif type_name == "PhysicsScene":
379
- scene = UsdPhysics.Scene(prim)
380
- g_vec = scene.GetGravityDirectionAttr()
381
- g_mag = scene.GetGravityMagnitudeAttr()
382
- if g_mag.HasAuthoredValue() and np.isfinite(g_mag.Get()):
383
- builder.gravity = g_mag.Get() * linear_unit
384
- if g_vec.HasAuthoredValue() and np.linalg.norm(g_vec.Get()) > 0.0:
385
- builder.up_vector = np.array(g_vec.Get(), dtype=np.float32) # TODO flip sign?
386
- else:
387
- builder.up_vector = upaxis
355
+ try:
356
+ scene = UsdPhysics.Scene(prim)
357
+ g_vec = scene.GetGravityDirectionAttr()
358
+ g_mag = scene.GetGravityMagnitudeAttr()
359
+ if g_mag.HasAuthoredValue() and np.isfinite(g_mag.Get()):
360
+ builder.gravity = -np.abs(g_mag.Get() * linear_unit)
361
+ if g_vec.HasAuthoredValue() and np.linalg.norm(g_vec.Get()) > 0.0:
362
+ builder.up_vector = np.array(g_vec.Get(), dtype=np.float32)
363
+ if np.any(builder.up_vector < 0.0):
364
+ builder.up_vector = -builder.up_vector
365
+ else:
366
+ builder.up_vector = upaxis
367
+ except:
368
+ pass
388
369
 
389
370
  def parse_prim(prim, incoming_xform, incoming_scale, parent_body: int = -1):
390
371
  nonlocal builder
391
372
  nonlocal joint_data
392
373
  nonlocal path_body_map
393
374
  nonlocal path_shape_map
375
+ nonlocal path_shape_scale
394
376
  nonlocal path_world_poses
395
377
  nonlocal prim_joint_xforms
396
378
  nonlocal path_collision_filters
@@ -403,13 +385,14 @@ def parse_usd(
403
385
  return
404
386
 
405
387
  type_name = str(prim.GetTypeName())
406
- if (
407
- type_name.endswith("Joint")
408
- or type_name.endswith("Light")
409
- or type_name.endswith("Scene")
410
- or type_name.endswith("Material")
411
- ):
388
+ if type_name.endswith("Joint") or type_name.endswith("Light") or type_name.endswith("Material"):
412
389
  return
390
+ if verbose:
391
+ print(f"parse_prim {prim.GetPath()} ({type_name})")
392
+ if type_name == "PhysicsScene":
393
+ # in case the PhysicsScene has bodies as children...
394
+ for child in prim.GetChildren():
395
+ parse_prim(child, incoming_xform, incoming_scale, parent_body)
413
396
 
414
397
  schemas = set(prim.GetAppliedSchemas())
415
398
  children_refs = prim.GetChildren()
@@ -725,6 +708,7 @@ def parse_usd(
725
708
 
726
709
  path_body_map[path] = body_id
727
710
  path_shape_map[path] = shape_id
711
+ path_shape_scale[path] = scale
728
712
 
729
713
  if prim.HasRelationship("physics:filteredPairs"):
730
714
  other_paths = prim.GetRelationship("physics:filteredPairs").GetTargets()
@@ -791,4 +775,92 @@ def parse_usd(
791
775
  "fps": stage.GetFramesPerSecond(),
792
776
  "duration": stage.GetEndTimeCode() - stage.GetStartTimeCode(),
793
777
  "up_axis": UsdGeom.GetStageUpAxis(stage).upper(),
778
+ "path_shape_map": path_shape_map,
779
+ "path_body_map": path_body_map,
780
+ "path_shape_scale": path_shape_scale,
781
+ "mass_unit": mass_unit,
782
+ "linear_unit": linear_unit,
794
783
  }
784
+
785
+
786
+ def resolve_usd_from_url(url: str, target_folder_name: str = None, export_usda: bool = False):
787
+ """
788
+ Downloads a USD file from a URL and resolves all references to other USD files to be downloaded to the given target folder.
789
+
790
+ Args:
791
+ url (str): URL to the USD file.
792
+ target_folder_name (str): Target folder name. If None, a timestamped folder will be created in the current directory.
793
+ 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.
794
+
795
+ Returns:
796
+ str: File path to the downloaded USD file.
797
+ """
798
+ import requests
799
+ import datetime
800
+ import os
801
+
802
+ try:
803
+ from pxr import Usd
804
+ except ImportError:
805
+ raise ImportError("Failed to import pxr. Please install USD (e.g. via `pip install usd-core`).")
806
+
807
+ response = requests.get(url, allow_redirects=True)
808
+ if response.status_code != 200:
809
+ raise RuntimeError(f"Failed to download USD file. Status code: {response.status_code}")
810
+ file = response.content
811
+ dot = os.path.extsep
812
+ base = os.path.basename(url)
813
+ url_folder = os.path.dirname(url)
814
+ base_name = dot.join(base.split(dot)[:-1])
815
+ if target_folder_name is None:
816
+ timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
817
+ target_folder_name = os.path.join(".usd_cache", f"{base_name}_{timestamp}")
818
+ os.makedirs(target_folder_name, exist_ok=True)
819
+ target_filename = os.path.join(target_folder_name, base)
820
+ with open(target_filename, "wb") as f:
821
+ f.write(file)
822
+
823
+ stage = Usd.Stage.Open(target_filename, Usd.Stage.LoadNone)
824
+ stage_str = stage.GetRootLayer().ExportToString()
825
+ print(f"Downloaded USD file to {target_filename}.")
826
+ if export_usda:
827
+ usda_filename = os.path.join(target_folder_name, base_name + ".usda")
828
+ with open(usda_filename, "w") as f:
829
+ f.write(stage_str)
830
+ print(f"Exported USDA file to {usda_filename}.")
831
+
832
+ # parse referenced USD files like `references = @./franka_collisions.usd@`
833
+ downloaded = set()
834
+ for match in re.finditer(r"references.=.@(.*?)@", stage_str):
835
+ refname = match.group(1)
836
+ if refname.startswith("./"):
837
+ refname = refname[2:]
838
+ if refname in downloaded:
839
+ continue
840
+ try:
841
+ response = requests.get(f"{url_folder}/{refname}", allow_redirects=True)
842
+ if response.status_code != 200:
843
+ print(f"Failed to download reference {refname}. Status code: {response.status_code}")
844
+ continue
845
+ file = response.content
846
+ refdir = os.path.dirname(refname)
847
+ if refdir:
848
+ os.makedirs(os.path.join(target_folder_name, refdir), exist_ok=True)
849
+ ref_filename = os.path.join(target_folder_name, refname)
850
+ if not os.path.exists(ref_filename):
851
+ with open(ref_filename, "wb") as f:
852
+ f.write(file)
853
+ downloaded.add(refname)
854
+ print(f"Downloaded USD reference {refname} to {ref_filename}.")
855
+ if export_usda:
856
+ ref_stage = Usd.Stage.Open(ref_filename, Usd.Stage.LoadNone)
857
+ ref_stage_str = ref_stage.GetRootLayer().ExportToString()
858
+ base = os.path.basename(ref_filename)
859
+ base_name = dot.join(base.split(dot)[:-1])
860
+ usda_filename = os.path.join(target_folder_name, base_name + ".usda")
861
+ with open(usda_filename, "w") as f:
862
+ f.write(ref_stage_str)
863
+ print(f"Exported USDA file to {usda_filename}.")
864
+ except Exception:
865
+ print(f"Failed to download {refname}.")
866
+ return target_filename