warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.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 (170) hide show
  1. warp/__init__.py +8 -0
  2. warp/bin/warp-clang.so +0 -0
  3. warp/bin/warp.so +0 -0
  4. warp/build.py +7 -6
  5. warp/build_dll.py +70 -79
  6. warp/builtins.py +10 -6
  7. warp/codegen.py +51 -19
  8. warp/config.py +7 -8
  9. warp/constants.py +3 -0
  10. warp/context.py +948 -245
  11. warp/dlpack.py +198 -113
  12. warp/examples/assets/bunny.usd +0 -0
  13. warp/examples/assets/cartpole.urdf +110 -0
  14. warp/examples/assets/crazyflie.usd +0 -0
  15. warp/examples/assets/cube.usda +42 -0
  16. warp/examples/assets/nv_ant.xml +92 -0
  17. warp/examples/assets/nv_humanoid.xml +183 -0
  18. warp/examples/assets/quadruped.urdf +268 -0
  19. warp/examples/assets/rocks.nvdb +0 -0
  20. warp/examples/assets/rocks.usd +0 -0
  21. warp/examples/assets/sphere.usda +56 -0
  22. warp/examples/assets/torus.usda +105 -0
  23. warp/examples/benchmarks/benchmark_api.py +383 -0
  24. warp/examples/benchmarks/benchmark_cloth.py +279 -0
  25. warp/examples/benchmarks/benchmark_cloth_cupy.py +88 -0
  26. warp/examples/benchmarks/benchmark_cloth_jax.py +100 -0
  27. warp/examples/benchmarks/benchmark_cloth_numba.py +142 -0
  28. warp/examples/benchmarks/benchmark_cloth_numpy.py +77 -0
  29. warp/examples/benchmarks/benchmark_cloth_pytorch.py +86 -0
  30. warp/examples/benchmarks/benchmark_cloth_taichi.py +112 -0
  31. warp/examples/benchmarks/benchmark_cloth_warp.py +146 -0
  32. warp/examples/benchmarks/benchmark_launches.py +295 -0
  33. warp/examples/core/example_dem.py +221 -0
  34. warp/examples/core/example_fluid.py +267 -0
  35. warp/examples/core/example_graph_capture.py +129 -0
  36. warp/examples/core/example_marching_cubes.py +177 -0
  37. warp/examples/core/example_mesh.py +154 -0
  38. warp/examples/core/example_mesh_intersect.py +193 -0
  39. warp/examples/core/example_nvdb.py +169 -0
  40. warp/examples/core/example_raycast.py +89 -0
  41. warp/examples/core/example_raymarch.py +178 -0
  42. warp/examples/core/example_render_opengl.py +141 -0
  43. warp/examples/core/example_sph.py +389 -0
  44. warp/examples/core/example_torch.py +181 -0
  45. warp/examples/core/example_wave.py +249 -0
  46. warp/examples/fem/bsr_utils.py +380 -0
  47. warp/examples/fem/example_apic_fluid.py +391 -0
  48. warp/examples/fem/example_convection_diffusion.py +168 -0
  49. warp/examples/fem/example_convection_diffusion_dg.py +209 -0
  50. warp/examples/fem/example_convection_diffusion_dg0.py +194 -0
  51. warp/examples/fem/example_deformed_geometry.py +159 -0
  52. warp/examples/fem/example_diffusion.py +173 -0
  53. warp/examples/fem/example_diffusion_3d.py +152 -0
  54. warp/examples/fem/example_diffusion_mgpu.py +214 -0
  55. warp/examples/fem/example_mixed_elasticity.py +222 -0
  56. warp/examples/fem/example_navier_stokes.py +243 -0
  57. warp/examples/fem/example_stokes.py +192 -0
  58. warp/examples/fem/example_stokes_transfer.py +249 -0
  59. warp/examples/fem/mesh_utils.py +109 -0
  60. warp/examples/fem/plot_utils.py +287 -0
  61. warp/examples/optim/example_bounce.py +248 -0
  62. warp/examples/optim/example_cloth_throw.py +210 -0
  63. warp/examples/optim/example_diffray.py +535 -0
  64. warp/examples/optim/example_drone.py +850 -0
  65. warp/examples/optim/example_inverse_kinematics.py +169 -0
  66. warp/examples/optim/example_inverse_kinematics_torch.py +170 -0
  67. warp/examples/optim/example_spring_cage.py +234 -0
  68. warp/examples/optim/example_trajectory.py +201 -0
  69. warp/examples/sim/example_cartpole.py +128 -0
  70. warp/examples/sim/example_cloth.py +184 -0
  71. warp/examples/sim/example_granular.py +113 -0
  72. warp/examples/sim/example_granular_collision_sdf.py +185 -0
  73. warp/examples/sim/example_jacobian_ik.py +213 -0
  74. warp/examples/sim/example_particle_chain.py +106 -0
  75. warp/examples/sim/example_quadruped.py +179 -0
  76. warp/examples/sim/example_rigid_chain.py +191 -0
  77. warp/examples/sim/example_rigid_contact.py +176 -0
  78. warp/examples/sim/example_rigid_force.py +126 -0
  79. warp/examples/sim/example_rigid_gyroscopic.py +97 -0
  80. warp/examples/sim/example_rigid_soft_contact.py +124 -0
  81. warp/examples/sim/example_soft_body.py +178 -0
  82. warp/fabric.py +29 -20
  83. warp/fem/cache.py +0 -1
  84. warp/fem/dirichlet.py +0 -2
  85. warp/fem/integrate.py +0 -1
  86. warp/jax.py +45 -0
  87. warp/jax_experimental.py +339 -0
  88. warp/native/builtin.h +12 -0
  89. warp/native/bvh.cu +18 -18
  90. warp/native/clang/clang.cpp +8 -3
  91. warp/native/cuda_util.cpp +94 -5
  92. warp/native/cuda_util.h +35 -6
  93. warp/native/cutlass_gemm.cpp +1 -1
  94. warp/native/cutlass_gemm.cu +4 -1
  95. warp/native/error.cpp +66 -0
  96. warp/native/error.h +27 -0
  97. warp/native/mesh.cu +2 -2
  98. warp/native/reduce.cu +4 -4
  99. warp/native/runlength_encode.cu +2 -2
  100. warp/native/scan.cu +2 -2
  101. warp/native/sparse.cu +0 -1
  102. warp/native/temp_buffer.h +2 -2
  103. warp/native/warp.cpp +95 -60
  104. warp/native/warp.cu +1053 -218
  105. warp/native/warp.h +49 -32
  106. warp/optim/linear.py +33 -16
  107. warp/render/render_opengl.py +202 -101
  108. warp/render/render_usd.py +82 -40
  109. warp/sim/__init__.py +13 -4
  110. warp/sim/articulation.py +4 -5
  111. warp/sim/collide.py +320 -175
  112. warp/sim/import_mjcf.py +25 -30
  113. warp/sim/import_urdf.py +94 -63
  114. warp/sim/import_usd.py +51 -36
  115. warp/sim/inertia.py +3 -2
  116. warp/sim/integrator.py +233 -0
  117. warp/sim/integrator_euler.py +447 -469
  118. warp/sim/integrator_featherstone.py +1991 -0
  119. warp/sim/integrator_xpbd.py +1420 -640
  120. warp/sim/model.py +765 -487
  121. warp/sim/particles.py +2 -1
  122. warp/sim/render.py +35 -13
  123. warp/sim/utils.py +222 -11
  124. warp/stubs.py +8 -0
  125. warp/tape.py +16 -1
  126. warp/tests/aux_test_grad_customs.py +23 -0
  127. warp/tests/test_array.py +190 -1
  128. warp/tests/test_async.py +656 -0
  129. warp/tests/test_bool.py +50 -0
  130. warp/tests/test_dlpack.py +164 -11
  131. warp/tests/test_examples.py +166 -74
  132. warp/tests/test_fem.py +8 -1
  133. warp/tests/test_generics.py +15 -5
  134. warp/tests/test_grad.py +1 -1
  135. warp/tests/test_grad_customs.py +172 -12
  136. warp/tests/test_jax.py +254 -0
  137. warp/tests/test_large.py +29 -6
  138. warp/tests/test_launch.py +25 -0
  139. warp/tests/test_linear_solvers.py +20 -3
  140. warp/tests/test_matmul.py +61 -16
  141. warp/tests/test_matmul_lite.py +13 -13
  142. warp/tests/test_mempool.py +186 -0
  143. warp/tests/test_multigpu.py +3 -0
  144. warp/tests/test_options.py +16 -2
  145. warp/tests/test_peer.py +137 -0
  146. warp/tests/test_print.py +3 -1
  147. warp/tests/test_quat.py +23 -0
  148. warp/tests/test_sim_kinematics.py +97 -0
  149. warp/tests/test_snippet.py +126 -3
  150. warp/tests/test_streams.py +108 -79
  151. warp/tests/test_torch.py +16 -8
  152. warp/tests/test_utils.py +32 -27
  153. warp/tests/test_verify_fp.py +65 -0
  154. warp/tests/test_volume.py +1 -1
  155. warp/tests/unittest_serial.py +2 -0
  156. warp/tests/unittest_suites.py +12 -0
  157. warp/tests/unittest_utils.py +14 -7
  158. warp/thirdparty/unittest_parallel.py +15 -3
  159. warp/torch.py +10 -8
  160. warp/types.py +363 -246
  161. warp/utils.py +143 -19
  162. warp_lang-1.0.0.dist-info/LICENSE.md +126 -0
  163. warp_lang-1.0.0.dist-info/METADATA +394 -0
  164. {warp_lang-0.11.0.dist-info → warp_lang-1.0.0.dist-info}/RECORD +167 -86
  165. warp/sim/optimizer.py +0 -138
  166. warp_lang-0.11.0.dist-info/LICENSE.md +0 -36
  167. warp_lang-0.11.0.dist-info/METADATA +0 -238
  168. /warp/tests/{walkthough_debug.py → walkthrough_debug.py} +0 -0
  169. {warp_lang-0.11.0.dist-info → warp_lang-1.0.0.dist-info}/WHEEL +0 -0
  170. {warp_lang-0.11.0.dist-info → warp_lang-1.0.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,185 @@
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
+ # Example Sim Granular Collision SDF
10
+ #
11
+ # Shows how to set up a particle-based granular material model using the
12
+ # wp.sim.ModelBuilder(). This version shows how to create collision geometry
13
+ # objects from SDFs.
14
+ #
15
+ # Note: requires a CUDA-capable device
16
+ ###########################################################################
17
+
18
+ import math
19
+ import os
20
+
21
+ import numpy as np
22
+
23
+ import warp as wp
24
+ import warp.sim
25
+ import warp.sim.render
26
+
27
+ wp.init()
28
+
29
+
30
+ class Example:
31
+ def __init__(self, stage):
32
+ self.frame_dt = 1.0 / 60
33
+ self.frame_count = 400
34
+
35
+ self.sim_substeps = 64
36
+ self.sim_dt = self.frame_dt / self.sim_substeps
37
+ self.sim_steps = self.frame_count * self.sim_substeps
38
+ self.sim_time = 0.0
39
+
40
+ self.radius = 0.1
41
+
42
+ builder = wp.sim.ModelBuilder()
43
+ builder.default_particle_radius = self.radius
44
+
45
+ builder.add_particle_grid(
46
+ dim_x=16,
47
+ dim_y=32,
48
+ dim_z=16,
49
+ cell_x=self.radius * 2.0,
50
+ cell_y=self.radius * 2.0,
51
+ cell_z=self.radius * 2.0,
52
+ pos=wp.vec3(0.0, 20.0, 0.0),
53
+ rot=wp.quat_identity(),
54
+ vel=wp.vec3(2.0, 0.0, 0.0),
55
+ mass=0.1,
56
+ jitter=self.radius * 0.1,
57
+ )
58
+ rock_file = open(os.path.join(os.path.dirname(__file__), "../assets/rocks.nvdb"), "rb")
59
+ rock_vdb = wp.Volume.load_from_nvdb(rock_file.read())
60
+ rock_file.close()
61
+
62
+ rock_sdf = wp.sim.SDF(rock_vdb)
63
+
64
+ builder.add_shape_sdf(
65
+ ke=1.0e4,
66
+ kd=1000.0,
67
+ kf=1000.0,
68
+ mu=0.5,
69
+ sdf=rock_sdf,
70
+ body=-1,
71
+ pos=wp.vec3(0.0, 0.0, 0.0),
72
+ rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -0.5 * math.pi),
73
+ scale=wp.vec3(0.01, 0.01, 0.01),
74
+ )
75
+
76
+ mins = np.array([-3.0, -3.0, -3.0])
77
+ voxel_size = 0.2
78
+ maxs = np.array([3.0, 3.0, 3.0])
79
+ nums = np.ceil((maxs - mins) / (voxel_size)).astype(dtype=int)
80
+ center = np.array([0.0, 0.0, 0.0])
81
+ rad = 2.5
82
+ sphere_sdf_np = np.zeros(tuple(nums))
83
+ for x in range(nums[0]):
84
+ for y in range(nums[1]):
85
+ for z in range(nums[2]):
86
+ pos = mins + voxel_size * np.array([x, y, z])
87
+ dis = np.linalg.norm(pos - center)
88
+ sphere_sdf_np[x, y, z] = dis - rad
89
+
90
+ sphere_vdb = wp.Volume.load_from_numpy(sphere_sdf_np, mins, voxel_size, rad + 3.0 * voxel_size)
91
+ sphere_sdf = wp.sim.SDF(sphere_vdb)
92
+
93
+ self.sphere_pos = wp.vec3(3.0, 15.0, 0.0)
94
+ self.sphere_scale = 1.0
95
+ self.sphere_radius = rad
96
+ builder.add_shape_sdf(
97
+ ke=1.0e4,
98
+ kd=1000.0,
99
+ kf=1000.0,
100
+ mu=0.5,
101
+ sdf=sphere_sdf,
102
+ body=-1,
103
+ pos=self.sphere_pos,
104
+ scale=wp.vec3(self.sphere_scale, self.sphere_scale, self.sphere_scale),
105
+ )
106
+
107
+ self.model = builder.finalize()
108
+ self.model.particle_kf = 25.0
109
+
110
+ self.model.soft_contact_kd = 100.0
111
+ self.model.soft_contact_kf *= 2.0
112
+
113
+ self.state_0 = self.model.state()
114
+ self.state_1 = self.model.state()
115
+
116
+ self.integrator = wp.sim.SemiImplicitIntegrator()
117
+
118
+ self.renderer = None
119
+ if stage:
120
+ self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=20.0)
121
+
122
+ self.use_graph = wp.get_device().is_cuda
123
+ if self.use_graph:
124
+ with wp.ScopedCapture() as capture:
125
+ self.simulate()
126
+ self.graph = capture.graph
127
+
128
+ def simulate(self):
129
+ for _ in range(self.sim_substeps):
130
+ self.state_0.clear_forces()
131
+ wp.sim.collide(self.model, self.state_0)
132
+ self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
133
+
134
+ # swap states
135
+ (self.state_0, self.state_1) = (self.state_1, self.state_0)
136
+
137
+ def step(self):
138
+ with wp.ScopedTimer("step", active=True):
139
+ self.model.particle_grid.build(self.state_0.particle_q, self.radius * 2.0)
140
+ if self.use_graph:
141
+ wp.capture_launch(self.graph)
142
+ else:
143
+ self.simulate()
144
+
145
+ self.sim_time += self.frame_dt
146
+
147
+ def render(self):
148
+ if self.renderer is None:
149
+ return
150
+
151
+ with wp.ScopedTimer("render", active=True):
152
+ self.renderer.begin_frame(self.sim_time)
153
+
154
+ # Note the extra wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi) is because .usd is oriented differently from .nvdb
155
+ self.renderer.render_ref(
156
+ name="collision",
157
+ path=os.path.join(os.path.dirname(__file__), "../assets/rocks.usd"),
158
+ pos=wp.vec3(0.0, 0.0, 0.0),
159
+ rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -0.5 * math.pi)
160
+ * wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi),
161
+ scale=wp.vec3(0.01, 0.01, 0.01),
162
+ )
163
+
164
+ self.renderer.render_sphere(
165
+ name="sphere",
166
+ pos=self.sphere_pos,
167
+ radius=self.sphere_scale * self.sphere_radius,
168
+ rot=wp.quat(0.0, 0.0, 0.0, 1.0),
169
+ )
170
+
171
+ self.renderer.render(self.state_0)
172
+ self.renderer.end_frame()
173
+
174
+
175
+ if __name__ == "__main__":
176
+ stage_path = os.path.join(os.path.dirname(__file__), "example_granular_collision_sdf.usd")
177
+
178
+ example = Example(stage_path)
179
+
180
+ for _ in range(example.frame_count):
181
+ example.step()
182
+ example.render()
183
+
184
+ if example.renderer:
185
+ example.renderer.save()
@@ -0,0 +1,213 @@
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
+ # Example Jacobian
10
+ #
11
+ # Demonstrates how to compute the Jacobian of a multi-valued function.
12
+ # Here, we use the simulation of a cartpole to differentiate
13
+ # through the kinematics function. We instantiate multiple copies of the
14
+ # cartpole and compute the Jacobian of the state of each cartpole in parallel
15
+ # in order to perform inverse kinematics via Jacobian transpose.
16
+ #
17
+ ###########################################################################
18
+
19
+ import math
20
+ import os
21
+
22
+ import numpy as np
23
+
24
+ import warp as wp
25
+ import warp.sim
26
+ import warp.sim.render
27
+
28
+ wp.init()
29
+
30
+
31
+ @wp.kernel
32
+ def compute_endeffector_position(
33
+ body_q: wp.array(dtype=wp.transform),
34
+ num_links: int,
35
+ ee_link_index: int,
36
+ ee_link_offset: wp.vec3,
37
+ ee_pos: wp.array(dtype=wp.vec3),
38
+ ):
39
+ tid = wp.tid()
40
+ ee_pos[tid] = wp.transform_point(body_q[tid * num_links + ee_link_index], ee_link_offset)
41
+
42
+
43
+ class Example:
44
+ def __init__(self, stage, num_envs=1):
45
+ builder = wp.sim.ModelBuilder()
46
+
47
+ self.num_envs = num_envs
48
+
49
+ self.frame_dt = 1.0 / 60.0
50
+
51
+ self.render_time = 0.0
52
+
53
+ # step size to use for the IK updates
54
+ self.step_size = 0.1
55
+
56
+ articulation_builder = wp.sim.ModelBuilder()
57
+
58
+ wp.sim.parse_urdf(
59
+ os.path.join(os.path.dirname(__file__), "../assets/cartpole.urdf"),
60
+ articulation_builder,
61
+ xform=wp.transform_identity(),
62
+ floating=False,
63
+ )
64
+
65
+ builder = wp.sim.ModelBuilder()
66
+
67
+ self.num_links = len(articulation_builder.joint_type)
68
+ # use the last link as the end-effector
69
+ self.ee_link_index = self.num_links - 1
70
+ self.ee_link_offset = wp.vec3(0.0, 0.0, 1.0)
71
+
72
+ self.dof = len(articulation_builder.joint_q)
73
+
74
+ self.target_origin = []
75
+ for i in range(num_envs):
76
+ builder.add_builder(
77
+ articulation_builder,
78
+ xform=wp.transform(
79
+ wp.vec3(i * 2.0, 4.0, 0.0), wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)
80
+ ),
81
+ )
82
+ self.target_origin.append((i * 2.0, 4.0, 0.0))
83
+ # joint initial positions
84
+ builder.joint_q[-3:] = np.random.uniform(-0.5, 0.5, size=3)
85
+ self.target_origin = np.array(self.target_origin)
86
+
87
+ # finalize model
88
+ self.model = builder.finalize()
89
+ self.model.ground = True
90
+
91
+ self.model.joint_q.requires_grad = True
92
+ self.model.body_q.requires_grad = True
93
+
94
+ self.model.joint_attach_ke = 1600.0
95
+ self.model.joint_attach_kd = 20.0
96
+
97
+ self.integrator = wp.sim.SemiImplicitIntegrator()
98
+
99
+ self.renderer = None
100
+ if stage:
101
+ self.renderer = wp.sim.render.SimRenderer(self.model, stage)
102
+
103
+ self.ee_pos = wp.zeros(self.num_envs, dtype=wp.vec3, requires_grad=True)
104
+
105
+ self.state = self.model.state(requires_grad=True)
106
+
107
+ self.targets = self.target_origin.copy()
108
+
109
+ self.profiler = {}
110
+
111
+ def compute_ee_position(self):
112
+ # computes the end-effector position from the current joint angles
113
+ wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state)
114
+ wp.launch(
115
+ compute_endeffector_position,
116
+ dim=self.num_envs,
117
+ inputs=[self.state.body_q, self.num_links, self.ee_link_index, self.ee_link_offset],
118
+ outputs=[self.ee_pos]
119
+ )
120
+ return self.ee_pos
121
+
122
+ def compute_jacobian(self):
123
+ # our function has 3 outputs (EE position), so we need a 3xN jacobian per environment
124
+ jacobians = np.empty((self.num_envs, 3, self.dof), dtype=np.float32)
125
+ tape = wp.Tape()
126
+ with tape:
127
+ self.compute_ee_position()
128
+ for output_index in range(3):
129
+ # select which row of the Jacobian we want to compute
130
+ select_index = np.zeros(3)
131
+ select_index[output_index] = 1.0
132
+ e = wp.array(np.tile(select_index, self.num_envs), dtype=wp.vec3)
133
+ tape.backward(grads={self.ee_pos: e})
134
+ q_grad_i = tape.gradients[self.model.joint_q]
135
+ jacobians[:, output_index, :] = q_grad_i.numpy().reshape(self.num_envs, self.dof)
136
+ tape.zero()
137
+ return jacobians
138
+
139
+ def compute_fd_jacobian(self, eps=1e-4):
140
+ jacobians = np.zeros((self.num_envs, 3, self.dof), dtype=np.float32)
141
+ q0 = self.model.joint_q.numpy().copy()
142
+ for e in range(self.num_envs):
143
+ for i in range(self.dof):
144
+ q = q0.copy()
145
+ q[e * self.dof + i] += eps
146
+ self.model.joint_q.assign(q)
147
+ self.compute_ee_position()
148
+ f_plus = self.ee_pos.numpy()[e].copy()
149
+ q[e * self.dof + i] -= 2 * eps
150
+ self.model.joint_q.assign(q)
151
+ self.compute_ee_position()
152
+ f_minus = self.ee_pos.numpy()[e].copy()
153
+ jacobians[e, :, i] = (f_plus - f_minus) / (2 * eps)
154
+ return jacobians
155
+
156
+ def step(self):
157
+ with wp.ScopedTimer("jacobian", print=False, active=True, dict=self.profiler):
158
+ # compute jacobian
159
+ jacobians = self.compute_jacobian()
160
+
161
+ # compute error
162
+ self.ee_pos_np = self.compute_ee_position().numpy()
163
+ error = self.targets - self.ee_pos_np
164
+ self.error = error.reshape(self.num_envs, 3, 1)
165
+
166
+ # compute Jacobian transpose update
167
+ delta_q = np.matmul(jacobians.transpose(0, 2, 1), self.error)
168
+
169
+ self.model.joint_q = wp.array(
170
+ self.model.joint_q.numpy() + self.step_size * delta_q.flatten(),
171
+ dtype=wp.float32,
172
+ requires_grad=True,
173
+ )
174
+
175
+ def render(self):
176
+ if self.renderer is None:
177
+ return
178
+
179
+ self.renderer.begin_frame(self.render_time)
180
+ self.renderer.render(self.state)
181
+ self.renderer.render_points("targets", self.targets, radius=0.05)
182
+ self.renderer.render_points("ee_pos", self.ee_pos_np, radius=0.05)
183
+ self.renderer.end_frame()
184
+ self.render_time += self.frame_dt
185
+
186
+
187
+ if __name__ == "__main__":
188
+ stage_path = os.path.join(os.path.dirname(__file__), "example_jacobian_ik.usd")
189
+
190
+ example = Example(stage_path, num_envs=10)
191
+
192
+ print("autodiff:")
193
+ print(example.compute_jacobian())
194
+ print("finite diff:")
195
+ print(example.compute_fd_jacobian())
196
+
197
+ for _ in range(5):
198
+ # select new random target points
199
+ example.targets = example.target_origin.copy()
200
+ example.targets[:, 1:] += np.random.uniform(-0.5, 0.5, size=(example.num_envs, 2))
201
+
202
+ for iter in range(50):
203
+ example.step()
204
+ example.render()
205
+ print("iter:", iter, "error:", example.error.mean())
206
+
207
+ if example.renderer:
208
+ example.renderer.save()
209
+
210
+ avg_time = np.array(example.profiler["jacobian"]).mean()
211
+ avg_steps_second = 1000.0 * float(example.num_envs) / avg_time
212
+
213
+ print(f"envs: {example.num_envs} steps/second {avg_steps_second} avg_time {avg_time}")
@@ -0,0 +1,106 @@
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
+ # Example Sim Particle Chain
10
+ #
11
+ # Shows how to set up a simple chain of particles connected by springs
12
+ # using wp.sim.ModelBuilder().
13
+ #
14
+ ###########################################################################
15
+
16
+ import math
17
+ import os
18
+
19
+ import warp as wp
20
+ import warp.sim
21
+ import warp.sim.render
22
+
23
+ wp.init()
24
+
25
+
26
+ class Example:
27
+ def __init__(self, stage):
28
+ self.sim_width = 64
29
+ self.sim_height = 32
30
+
31
+ self.sim_fps = 60.0
32
+ self.frame_dt = 1.0 / self.sim_fps
33
+ self.sim_substeps = 10
34
+ self.sim_duration = 5.0
35
+ self.sim_frames = int(self.sim_duration * self.sim_fps)
36
+ self.sim_dt = (1.0 / self.sim_fps) / self.sim_substeps
37
+ self.sim_time = 0.0
38
+
39
+ builder = wp.sim.ModelBuilder()
40
+
41
+ # anchor
42
+ builder.add_particle(wp.vec3(0.0, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), 0.0)
43
+
44
+ # chain
45
+ for i in range(1, 10):
46
+ radius = math.sqrt(i) * 0.2
47
+ mass = math.pi * radius * radius * radius
48
+ builder.add_particle(wp.vec3(i, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), mass, radius=radius)
49
+ builder.add_spring(i - 1, i, 1.0e6, 0.0, 0)
50
+
51
+ self.model = builder.finalize()
52
+ self.model.ground = False
53
+
54
+ self.integrator = wp.sim.XPBDIntegrator()
55
+
56
+ self.state_0 = self.model.state()
57
+ self.state_1 = self.model.state()
58
+
59
+ self.renderer = None
60
+ if stage:
61
+ self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=15.0)
62
+
63
+ self.use_graph = wp.get_device().is_cuda
64
+ if self.use_graph:
65
+ with wp.ScopedCapture() as capture:
66
+ self.simulate()
67
+ self.graph = capture.graph
68
+
69
+ def simulate(self):
70
+ for _ in range(self.sim_substeps):
71
+ self.state_0.clear_forces()
72
+ self.state_1.clear_forces()
73
+ self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
74
+
75
+ # swap states
76
+ (self.state_0, self.state_1) = (self.state_1, self.state_0)
77
+
78
+ def step(self):
79
+ with wp.ScopedTimer("step"):
80
+ if self.use_graph:
81
+ wp.capture_launch(self.graph)
82
+ else:
83
+ self.simulate()
84
+ self.sim_time += self.frame_dt
85
+
86
+ def render(self):
87
+ if self.renderer is None:
88
+ return
89
+
90
+ with wp.ScopedTimer("render"):
91
+ self.renderer.begin_frame(self.sim_time)
92
+ self.renderer.render(self.state_0)
93
+ self.renderer.end_frame()
94
+
95
+
96
+ if __name__ == "__main__":
97
+ stage_path = os.path.join(os.path.dirname(__file__), "example_particle_chain.usd")
98
+
99
+ example = Example(stage_path)
100
+
101
+ for i in range(example.sim_frames):
102
+ example.step()
103
+ example.render()
104
+
105
+ if example.renderer:
106
+ example.renderer.save()
@@ -0,0 +1,179 @@
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
+ # Example Sim Quadruped
10
+ #
11
+ # Shows how to set up a simulation of a rigid-body quadruped articulation
12
+ # from a URDF using the wp.sim.ModelBuilder().
13
+ # Note this example does not include a trained policy.
14
+ #
15
+ ###########################################################################
16
+
17
+ import math
18
+ import os
19
+
20
+ import numpy as np
21
+
22
+ import warp as wp
23
+ import warp.sim
24
+ import warp.sim.render
25
+
26
+ wp.init()
27
+
28
+
29
+ # Taken from env/environment.py
30
+ def compute_env_offsets(num_envs, env_offset=(5.0, 0.0, 5.0), up_axis="Y"):
31
+ # compute positional offsets per environment
32
+ env_offset = np.array(env_offset)
33
+ nonzeros = np.nonzero(env_offset)[0]
34
+ num_dim = nonzeros.shape[0]
35
+ if num_dim > 0:
36
+ side_length = int(np.ceil(num_envs ** (1.0 / num_dim)))
37
+ env_offsets = []
38
+ else:
39
+ env_offsets = np.zeros((num_envs, 3))
40
+ if num_dim == 1:
41
+ for i in range(num_envs):
42
+ env_offsets.append(i * env_offset)
43
+ elif num_dim == 2:
44
+ for i in range(num_envs):
45
+ d0 = i // side_length
46
+ d1 = i % side_length
47
+ offset = np.zeros(3)
48
+ offset[nonzeros[0]] = d0 * env_offset[nonzeros[0]]
49
+ offset[nonzeros[1]] = d1 * env_offset[nonzeros[1]]
50
+ env_offsets.append(offset)
51
+ elif num_dim == 3:
52
+ for i in range(num_envs):
53
+ d0 = i // (side_length * side_length)
54
+ d1 = (i // side_length) % side_length
55
+ d2 = i % side_length
56
+ offset = np.zeros(3)
57
+ offset[0] = d0 * env_offset[0]
58
+ offset[1] = d1 * env_offset[1]
59
+ offset[2] = d2 * env_offset[2]
60
+ env_offsets.append(offset)
61
+ env_offsets = np.array(env_offsets)
62
+ min_offsets = np.min(env_offsets, axis=0)
63
+ correction = min_offsets + (np.max(env_offsets, axis=0) - min_offsets) / 2.0
64
+ if isinstance(up_axis, str):
65
+ up_axis = "XYZ".index(up_axis.upper())
66
+ correction[up_axis] = 0.0 # ensure the envs are not shifted below the ground plane
67
+ env_offsets -= correction
68
+ return env_offsets
69
+
70
+
71
+ class Example:
72
+ def __init__(self, stage=None, num_envs=1, print_timers=True):
73
+ self.num_envs = num_envs
74
+ articulation_builder = wp.sim.ModelBuilder()
75
+ wp.sim.parse_urdf(
76
+ os.path.join(os.path.dirname(__file__), "../assets/quadruped.urdf"),
77
+ articulation_builder,
78
+ xform=wp.transform([0.0, 0.7, 0.0], wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)),
79
+ floating=True,
80
+ density=1000,
81
+ armature=0.01,
82
+ stiffness=200,
83
+ damping=1,
84
+ contact_ke=1.0e4,
85
+ contact_kd=1.0e2,
86
+ contact_kf=1.0e2,
87
+ contact_mu=1.0,
88
+ limit_ke=1.0e4,
89
+ limit_kd=1.0e1,
90
+ )
91
+
92
+ builder = wp.sim.ModelBuilder()
93
+
94
+ self.sim_time = 0.0
95
+ self.frame_dt = 1.0 / 100.0
96
+
97
+ episode_duration = 5.0 # seconds
98
+ self.episode_frames = int(episode_duration / self.frame_dt)
99
+
100
+ self.sim_substeps = 5
101
+ self.sim_dt = self.frame_dt / self.sim_substeps
102
+
103
+ offsets = compute_env_offsets(num_envs)
104
+ for i in range(num_envs):
105
+ builder.add_builder(articulation_builder, xform=wp.transform(offsets[i], wp.quat_identity()))
106
+
107
+ builder.joint_q[-12:] = [0.2, 0.4, -0.6, -0.2, -0.4, 0.6, -0.2, 0.4, -0.6, 0.2, -0.4, 0.6]
108
+
109
+ builder.joint_axis_mode = [wp.sim.JOINT_MODE_TARGET_POSITION] * len(builder.joint_axis_mode)
110
+ builder.joint_act[-12:] = [0.2, 0.4, -0.6, -0.2, -0.4, 0.6, -0.2, 0.4, -0.6, 0.2, -0.4, 0.6]
111
+
112
+ np.set_printoptions(suppress=True)
113
+ # finalize model
114
+ self.model = builder.finalize()
115
+ self.model.ground = True
116
+ # self.model.gravity = 0.0
117
+
118
+ self.model.joint_attach_ke = 16000.0
119
+ self.model.joint_attach_kd = 200.0
120
+
121
+ # self.integrator = wp.sim.XPBDIntegrator()
122
+ # self.integrator = wp.sim.SemiImplicitIntegrator()
123
+ self.integrator = wp.sim.FeatherstoneIntegrator(self.model)
124
+
125
+ self.renderer = None
126
+ if stage:
127
+ self.renderer = wp.sim.render.SimRenderer(self.model, stage)
128
+
129
+ self.print_timers = print_timers
130
+
131
+ self.state_0 = self.model.state()
132
+ self.state_1 = self.model.state()
133
+
134
+ wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state_0)
135
+
136
+ self.use_graph = wp.get_device().is_cuda
137
+ self.graph = None
138
+ if self.use_graph:
139
+ with wp.ScopedCapture() as capture:
140
+ self.simulate()
141
+ self.graph = capture.graph
142
+
143
+ def simulate(self):
144
+ for _ in range(self.sim_substeps):
145
+ self.state_0.clear_forces()
146
+ wp.sim.collide(self.model, self.state_0)
147
+ self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
148
+ self.state_0, self.state_1 = self.state_1, self.state_0
149
+
150
+ def step(self):
151
+ with wp.ScopedTimer("step", active=True, print=self.print_timers):
152
+ if self.use_graph:
153
+ wp.capture_launch(self.graph)
154
+ else:
155
+ self.simulate()
156
+ self.sim_time += self.frame_dt
157
+
158
+ def render(self):
159
+ if self.renderer is None:
160
+ return
161
+
162
+ with wp.ScopedTimer("render", active=True, print=self.print_timers):
163
+ self.renderer.begin_frame(self.sim_time)
164
+ self.renderer.render(self.state_0)
165
+ self.renderer.end_frame()
166
+
167
+
168
+ if __name__ == "__main__":
169
+
170
+ stage = os.path.join(os.path.dirname(__file__), "example_quadruped.usd")
171
+
172
+ example = Example(stage, num_envs=25)
173
+
174
+ for _ in range(example.episode_frames):
175
+ example.step()
176
+ example.render()
177
+
178
+ if example.renderer:
179
+ example.renderer.save()