warp-lang 1.5.0__py3-none-manylinux2014_aarch64.whl → 1.6.0__py3-none-manylinux2014_aarch64.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.
- warp/__init__.py +5 -0
- warp/autograd.py +414 -191
- warp/bin/warp-clang.so +0 -0
- warp/bin/warp.so +0 -0
- warp/build.py +40 -12
- warp/build_dll.py +13 -6
- warp/builtins.py +1124 -497
- warp/codegen.py +261 -136
- warp/config.py +1 -1
- warp/context.py +357 -119
- warp/examples/assets/square_cloth.usd +0 -0
- warp/examples/benchmarks/benchmark_gemm.py +27 -18
- warp/examples/benchmarks/benchmark_interop_paddle.py +3 -3
- warp/examples/benchmarks/benchmark_interop_torch.py +3 -3
- warp/examples/core/example_torch.py +18 -34
- warp/examples/fem/example_apic_fluid.py +1 -0
- warp/examples/fem/example_mixed_elasticity.py +1 -1
- warp/examples/optim/example_bounce.py +1 -1
- warp/examples/optim/example_cloth_throw.py +1 -1
- warp/examples/optim/example_diffray.py +4 -15
- warp/examples/optim/example_drone.py +1 -1
- warp/examples/optim/example_softbody_properties.py +392 -0
- warp/examples/optim/example_trajectory.py +1 -3
- warp/examples/optim/example_walker.py +5 -0
- warp/examples/sim/example_cartpole.py +0 -2
- warp/examples/sim/example_cloth.py +3 -1
- warp/examples/sim/example_cloth_self_contact.py +260 -0
- warp/examples/sim/example_granular_collision_sdf.py +4 -5
- warp/examples/sim/example_jacobian_ik.py +0 -2
- warp/examples/sim/example_quadruped.py +5 -2
- warp/examples/tile/example_tile_cholesky.py +79 -0
- warp/examples/tile/example_tile_convolution.py +2 -2
- warp/examples/tile/example_tile_fft.py +2 -2
- warp/examples/tile/example_tile_filtering.py +3 -3
- warp/examples/tile/example_tile_matmul.py +4 -4
- warp/examples/tile/example_tile_mlp.py +12 -12
- warp/examples/tile/example_tile_nbody.py +180 -0
- warp/examples/tile/example_tile_walker.py +319 -0
- warp/fem/geometry/geometry.py +0 -2
- warp/math.py +147 -0
- warp/native/array.h +12 -0
- warp/native/builtin.h +0 -1
- warp/native/bvh.cpp +149 -70
- warp/native/bvh.cu +287 -68
- warp/native/bvh.h +195 -85
- warp/native/clang/clang.cpp +5 -1
- warp/native/coloring.cpp +5 -1
- warp/native/cuda_util.cpp +91 -53
- warp/native/cuda_util.h +5 -0
- warp/native/exports.h +40 -40
- warp/native/intersect.h +17 -0
- warp/native/mat.h +41 -0
- warp/native/mathdx.cpp +19 -0
- warp/native/mesh.cpp +25 -8
- warp/native/mesh.cu +153 -101
- warp/native/mesh.h +482 -403
- warp/native/quat.h +40 -0
- warp/native/solid_angle.h +7 -0
- warp/native/sort.cpp +85 -0
- warp/native/sort.cu +34 -0
- warp/native/sort.h +3 -1
- warp/native/spatial.h +11 -0
- warp/native/tile.h +1187 -669
- warp/native/tile_reduce.h +8 -6
- warp/native/vec.h +41 -0
- warp/native/warp.cpp +8 -1
- warp/native/warp.cu +263 -40
- warp/native/warp.h +19 -5
- warp/optim/linear.py +22 -4
- warp/render/render_opengl.py +130 -64
- warp/sim/__init__.py +6 -1
- warp/sim/collide.py +270 -26
- warp/sim/import_urdf.py +8 -8
- warp/sim/integrator_euler.py +25 -7
- warp/sim/integrator_featherstone.py +154 -35
- warp/sim/integrator_vbd.py +842 -40
- warp/sim/model.py +134 -72
- warp/sparse.py +1 -1
- warp/stubs.py +265 -132
- warp/tape.py +28 -30
- warp/tests/aux_test_module_unload.py +15 -0
- warp/tests/{test_sim_grad.py → flaky_test_sim_grad.py} +104 -63
- warp/tests/test_array.py +74 -0
- warp/tests/test_assert.py +242 -0
- warp/tests/test_codegen.py +14 -61
- warp/tests/test_collision.py +2 -2
- warp/tests/test_coloring.py +12 -2
- warp/tests/test_examples.py +12 -1
- warp/tests/test_func.py +21 -4
- warp/tests/test_grad_debug.py +87 -2
- warp/tests/test_hash_grid.py +1 -1
- warp/tests/test_ipc.py +116 -0
- warp/tests/test_lerp.py +13 -87
- warp/tests/test_mat.py +138 -167
- warp/tests/test_math.py +47 -1
- warp/tests/test_matmul.py +17 -16
- warp/tests/test_matmul_lite.py +10 -15
- warp/tests/test_mesh.py +84 -60
- warp/tests/test_mesh_query_aabb.py +165 -0
- warp/tests/test_mesh_query_point.py +328 -286
- warp/tests/test_mesh_query_ray.py +134 -121
- warp/tests/test_mlp.py +2 -2
- warp/tests/test_operators.py +43 -0
- warp/tests/test_overwrite.py +47 -2
- warp/tests/test_quat.py +77 -0
- warp/tests/test_reload.py +29 -0
- warp/tests/test_sim_grad_bounce_linear.py +204 -0
- warp/tests/test_smoothstep.py +17 -83
- warp/tests/test_static.py +19 -3
- warp/tests/test_tape.py +25 -0
- warp/tests/test_tile.py +178 -191
- warp/tests/test_tile_load.py +356 -0
- warp/tests/test_tile_mathdx.py +61 -8
- warp/tests/test_tile_mlp.py +17 -17
- warp/tests/test_tile_reduce.py +24 -18
- warp/tests/test_tile_shared_memory.py +66 -17
- warp/tests/test_tile_view.py +165 -0
- warp/tests/test_torch.py +35 -0
- warp/tests/test_utils.py +36 -24
- warp/tests/test_vec.py +110 -0
- warp/tests/unittest_suites.py +29 -4
- warp/tests/unittest_utils.py +30 -13
- warp/thirdparty/unittest_parallel.py +2 -2
- warp/types.py +411 -101
- warp/utils.py +10 -7
- {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/METADATA +92 -69
- {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/RECORD +130 -119
- {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/WHEEL +1 -1
- warp/examples/benchmarks/benchmark_tile.py +0 -179
- warp/native/tile_gemm.h +0 -341
- {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/LICENSE.md +0 -0
- {warp_lang-1.5.0.dist-info → warp_lang-1.6.0.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,392 @@
|
|
|
1
|
+
# Copyright (c) 2025 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 FEM Bounce
|
|
10
|
+
#
|
|
11
|
+
# Shows how to use Warp to optimize for the material parameters of a soft body,
|
|
12
|
+
# such that it bounces off the wall and floor in order to hit a target.
|
|
13
|
+
#
|
|
14
|
+
# This example uses the built-in wp.Tape() object to compute gradients of
|
|
15
|
+
# the distance to target (loss) w.r.t the material parameters, followed by
|
|
16
|
+
# a simple gradient-descent optimization step.
|
|
17
|
+
#
|
|
18
|
+
###########################################################################
|
|
19
|
+
|
|
20
|
+
import numpy as np
|
|
21
|
+
|
|
22
|
+
import warp as wp
|
|
23
|
+
import warp.optim
|
|
24
|
+
import warp.sim
|
|
25
|
+
import warp.sim.render
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
@wp.kernel
|
|
29
|
+
def assign_param(params: wp.array(dtype=wp.float32), tet_materials: wp.array2d(dtype=wp.float32)):
|
|
30
|
+
tid = wp.tid()
|
|
31
|
+
params_idx = 2 * wp.tid() % params.shape[0]
|
|
32
|
+
tet_materials[tid, 0] = params[params_idx]
|
|
33
|
+
tet_materials[tid, 1] = params[params_idx + 1]
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
@wp.kernel
|
|
37
|
+
def com_kernel(particle_q: wp.array(dtype=wp.vec3), com: wp.array(dtype=wp.vec3)):
|
|
38
|
+
tid = wp.tid()
|
|
39
|
+
point = particle_q[tid]
|
|
40
|
+
a = point / wp.float32(particle_q.shape[0])
|
|
41
|
+
|
|
42
|
+
# Atomically add the point coordinates to the accumulator
|
|
43
|
+
wp.atomic_add(com, 0, a)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
@wp.kernel
|
|
47
|
+
def loss_kernel(
|
|
48
|
+
target: wp.vec3,
|
|
49
|
+
com: wp.array(dtype=wp.vec3),
|
|
50
|
+
pos_error: wp.array(dtype=float),
|
|
51
|
+
loss: wp.array(dtype=float),
|
|
52
|
+
):
|
|
53
|
+
diff = com[0] - target
|
|
54
|
+
pos_error[0] = wp.dot(diff, diff)
|
|
55
|
+
norm = pos_error[0]
|
|
56
|
+
loss[0] = norm
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
@wp.kernel
|
|
60
|
+
def enforce_constraint_kernel(lower_bound: wp.float32, upper_bound: wp.float32, x: wp.array(dtype=wp.float32)):
|
|
61
|
+
tid = wp.tid()
|
|
62
|
+
if x[tid] < lower_bound:
|
|
63
|
+
x[tid] = lower_bound
|
|
64
|
+
elif x[tid] > upper_bound:
|
|
65
|
+
x[tid] = upper_bound
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
class Example:
|
|
69
|
+
def __init__(
|
|
70
|
+
self,
|
|
71
|
+
stage_path="example_softbody_properties.usd",
|
|
72
|
+
material_behavior="anisotropic",
|
|
73
|
+
verbose=False,
|
|
74
|
+
):
|
|
75
|
+
self.verbose = verbose
|
|
76
|
+
self.material_behavior = material_behavior
|
|
77
|
+
|
|
78
|
+
# seconds
|
|
79
|
+
sim_duration = 1.0
|
|
80
|
+
|
|
81
|
+
# control frequency
|
|
82
|
+
fps = 60
|
|
83
|
+
self.frame_dt = 1.0 / fps
|
|
84
|
+
frame_steps = int(sim_duration / self.frame_dt)
|
|
85
|
+
|
|
86
|
+
# sim frequency
|
|
87
|
+
self.sim_substeps = 16
|
|
88
|
+
self.sim_steps = frame_steps * self.sim_substeps
|
|
89
|
+
self.sim_dt = self.frame_dt / self.sim_substeps
|
|
90
|
+
|
|
91
|
+
self.iter = 0
|
|
92
|
+
self.render_time = 0.0
|
|
93
|
+
|
|
94
|
+
self.train_rate = 1e7
|
|
95
|
+
|
|
96
|
+
self.losses = []
|
|
97
|
+
|
|
98
|
+
self.hard_lower_bound = wp.float32(500.0)
|
|
99
|
+
self.hard_upper_bound = wp.float32(4e6)
|
|
100
|
+
|
|
101
|
+
# Create FEM model.
|
|
102
|
+
self.cell_dim = 2
|
|
103
|
+
self.cell_size = 0.1
|
|
104
|
+
center = self.cell_size * self.cell_dim * 0.5
|
|
105
|
+
self.grid_origin = wp.vec3(-0.5, 1.0, -center)
|
|
106
|
+
self.create_model()
|
|
107
|
+
|
|
108
|
+
self.integrator = wp.sim.SemiImplicitIntegrator()
|
|
109
|
+
|
|
110
|
+
self.target = wp.vec3(-1.0, 1.5, 0.0)
|
|
111
|
+
# Initialize material parameters
|
|
112
|
+
if self.material_behavior == "anisotropic":
|
|
113
|
+
# Different Lame parameters for each tet
|
|
114
|
+
self.material_params = wp.array(
|
|
115
|
+
self.model.tet_materials.numpy()[:, :2].flatten(),
|
|
116
|
+
dtype=wp.float32,
|
|
117
|
+
requires_grad=True,
|
|
118
|
+
)
|
|
119
|
+
else:
|
|
120
|
+
# Same Lame parameters for all tets
|
|
121
|
+
self.material_params = wp.array(
|
|
122
|
+
self.model.tet_materials.numpy()[0, :2].flatten(),
|
|
123
|
+
dtype=wp.float32,
|
|
124
|
+
requires_grad=True,
|
|
125
|
+
)
|
|
126
|
+
|
|
127
|
+
self.optimizer = wp.optim.SGD(
|
|
128
|
+
[self.material_params],
|
|
129
|
+
lr=self.train_rate,
|
|
130
|
+
nesterov=False,
|
|
131
|
+
)
|
|
132
|
+
|
|
133
|
+
self.com = wp.array([wp.vec3(0.0, 0.0, 0.0)], dtype=wp.vec3, requires_grad=True)
|
|
134
|
+
self.pos_error = wp.zeros(1, dtype=wp.float32, requires_grad=True)
|
|
135
|
+
self.loss = wp.zeros(1, dtype=wp.float32, requires_grad=True)
|
|
136
|
+
|
|
137
|
+
# allocate sim states for trajectory
|
|
138
|
+
self.states = []
|
|
139
|
+
for _i in range(self.sim_steps + 1):
|
|
140
|
+
self.states.append(self.model.state())
|
|
141
|
+
|
|
142
|
+
if stage_path:
|
|
143
|
+
self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=1.0)
|
|
144
|
+
else:
|
|
145
|
+
self.renderer = None
|
|
146
|
+
|
|
147
|
+
# capture forward/backward passes
|
|
148
|
+
self.use_cuda_graph = wp.get_device().is_cuda
|
|
149
|
+
if self.use_cuda_graph:
|
|
150
|
+
with wp.ScopedCapture() as capture:
|
|
151
|
+
self.tape = wp.Tape()
|
|
152
|
+
with self.tape:
|
|
153
|
+
self.forward()
|
|
154
|
+
self.tape.backward(self.loss)
|
|
155
|
+
self.graph = capture.graph
|
|
156
|
+
|
|
157
|
+
def create_model(self):
|
|
158
|
+
builder = wp.sim.ModelBuilder()
|
|
159
|
+
builder.default_particle_radius = 0.0005
|
|
160
|
+
|
|
161
|
+
total_mass = 0.2
|
|
162
|
+
num_particles = (self.cell_dim + 1) ** 3
|
|
163
|
+
particle_mass = total_mass / num_particles
|
|
164
|
+
particle_density = particle_mass / (self.cell_size**3)
|
|
165
|
+
if self.verbose:
|
|
166
|
+
print(f"Particle density: {particle_density}")
|
|
167
|
+
|
|
168
|
+
young_mod = 1.5 * 1e4
|
|
169
|
+
poisson_ratio = 0.3
|
|
170
|
+
k_mu = 0.5 * young_mod / (1.0 + poisson_ratio)
|
|
171
|
+
k_lambda = young_mod * poisson_ratio / ((1 + poisson_ratio) * (1 - 2 * poisson_ratio))
|
|
172
|
+
|
|
173
|
+
builder.add_soft_grid(
|
|
174
|
+
pos=self.grid_origin,
|
|
175
|
+
rot=wp.quat_identity(),
|
|
176
|
+
vel=wp.vec3(5.0, -5.0, 0.0),
|
|
177
|
+
dim_x=self.cell_dim,
|
|
178
|
+
dim_y=self.cell_dim,
|
|
179
|
+
dim_z=self.cell_dim,
|
|
180
|
+
cell_x=self.cell_size,
|
|
181
|
+
cell_y=self.cell_size,
|
|
182
|
+
cell_z=self.cell_size,
|
|
183
|
+
density=particle_density,
|
|
184
|
+
k_mu=k_mu,
|
|
185
|
+
k_lambda=k_lambda,
|
|
186
|
+
k_damp=0.0,
|
|
187
|
+
tri_ke=1e-4,
|
|
188
|
+
tri_ka=1e-4,
|
|
189
|
+
tri_kd=1e-4,
|
|
190
|
+
tri_drag=0.0,
|
|
191
|
+
tri_lift=0.0,
|
|
192
|
+
fix_bottom=False,
|
|
193
|
+
)
|
|
194
|
+
|
|
195
|
+
ke = 1.0e3
|
|
196
|
+
kf = 0.0
|
|
197
|
+
kd = 1.0e0
|
|
198
|
+
mu = 0.2
|
|
199
|
+
builder.add_shape_box(
|
|
200
|
+
body=-1,
|
|
201
|
+
pos=wp.vec3(2.0, 1.0, 0.0),
|
|
202
|
+
hx=0.25,
|
|
203
|
+
hy=1.0,
|
|
204
|
+
hz=1.0,
|
|
205
|
+
ke=ke,
|
|
206
|
+
kf=kf,
|
|
207
|
+
kd=kd,
|
|
208
|
+
mu=mu,
|
|
209
|
+
)
|
|
210
|
+
|
|
211
|
+
# use `requires_grad=True` to create a model for differentiable simulation
|
|
212
|
+
self.model = builder.finalize(requires_grad=True)
|
|
213
|
+
self.model.ground = True
|
|
214
|
+
|
|
215
|
+
self.model.soft_contact_ke = ke
|
|
216
|
+
self.model.soft_contact_kf = kf
|
|
217
|
+
self.model.soft_contact_kd = kd
|
|
218
|
+
self.model.soft_contact_mu = mu
|
|
219
|
+
self.model.soft_contact_margin = 0.001
|
|
220
|
+
self.model.soft_contact_restitution = 1.0
|
|
221
|
+
|
|
222
|
+
def forward(self):
|
|
223
|
+
wp.launch(
|
|
224
|
+
kernel=assign_param,
|
|
225
|
+
dim=self.model.tet_count,
|
|
226
|
+
inputs=(self.material_params,),
|
|
227
|
+
outputs=(self.model.tet_materials,),
|
|
228
|
+
)
|
|
229
|
+
# run control loop
|
|
230
|
+
for i in range(self.sim_steps):
|
|
231
|
+
wp.sim.collide(self.model, self.states[i])
|
|
232
|
+
self.states[i].clear_forces()
|
|
233
|
+
|
|
234
|
+
self.integrator.simulate(self.model, self.states[i], self.states[i + 1], self.sim_dt)
|
|
235
|
+
|
|
236
|
+
# Update loss
|
|
237
|
+
# Compute the center of mass for the last time step.
|
|
238
|
+
wp.launch(
|
|
239
|
+
kernel=com_kernel,
|
|
240
|
+
dim=self.model.particle_count,
|
|
241
|
+
inputs=(self.states[-1].particle_q,),
|
|
242
|
+
outputs=(self.com,),
|
|
243
|
+
)
|
|
244
|
+
|
|
245
|
+
# calculate loss
|
|
246
|
+
wp.launch(
|
|
247
|
+
kernel=loss_kernel,
|
|
248
|
+
dim=1,
|
|
249
|
+
inputs=(
|
|
250
|
+
self.target,
|
|
251
|
+
self.com,
|
|
252
|
+
),
|
|
253
|
+
outputs=(self.pos_error, self.loss),
|
|
254
|
+
)
|
|
255
|
+
|
|
256
|
+
return self.loss
|
|
257
|
+
|
|
258
|
+
def step(self):
|
|
259
|
+
with wp.ScopedTimer("step"):
|
|
260
|
+
if self.use_cuda_graph:
|
|
261
|
+
wp.capture_launch(self.graph)
|
|
262
|
+
else:
|
|
263
|
+
self.tape = wp.Tape()
|
|
264
|
+
with self.tape:
|
|
265
|
+
self.forward()
|
|
266
|
+
self.tape.backward(loss=self.loss)
|
|
267
|
+
|
|
268
|
+
if self.verbose:
|
|
269
|
+
self.log_step()
|
|
270
|
+
|
|
271
|
+
self.optimizer.step([self.material_params.grad])
|
|
272
|
+
|
|
273
|
+
wp.launch(
|
|
274
|
+
kernel=enforce_constraint_kernel,
|
|
275
|
+
dim=self.material_params.shape[0],
|
|
276
|
+
inputs=(
|
|
277
|
+
self.hard_lower_bound,
|
|
278
|
+
self.hard_upper_bound,
|
|
279
|
+
),
|
|
280
|
+
outputs=(self.material_params,),
|
|
281
|
+
)
|
|
282
|
+
|
|
283
|
+
self.losses.append(self.loss.numpy()[0])
|
|
284
|
+
|
|
285
|
+
# clear grads for next iteration
|
|
286
|
+
self.tape.zero()
|
|
287
|
+
self.loss.zero_()
|
|
288
|
+
self.com.zero_()
|
|
289
|
+
self.pos_error.zero_()
|
|
290
|
+
|
|
291
|
+
self.iter = self.iter + 1
|
|
292
|
+
|
|
293
|
+
def log_step(self):
|
|
294
|
+
x = self.material_params.numpy().reshape(-1, 2)
|
|
295
|
+
x_grad = self.material_params.grad.numpy().reshape(-1, 2)
|
|
296
|
+
|
|
297
|
+
print(f"Iter: {self.iter} Loss: {self.loss.numpy()[0]}")
|
|
298
|
+
|
|
299
|
+
print(f"Pos error: {np.sqrt(self.pos_error.numpy()[0])}")
|
|
300
|
+
|
|
301
|
+
print(
|
|
302
|
+
f"Max Mu: {np.max(x[:, 0])}, Min Mu: {np.min(x[:, 0])}, "
|
|
303
|
+
f"Max Lambda: {np.max(x[:, 1])}, Min Lambda: {np.min(x[:, 1])}"
|
|
304
|
+
)
|
|
305
|
+
|
|
306
|
+
print(
|
|
307
|
+
f"Max Mu Grad: {np.max(x_grad[:, 0])}, Min Mu Grad: {np.min(x_grad[:, 0])}, "
|
|
308
|
+
f"Max Lambda Grad: {np.max(x_grad[:, 1])}, Min Lambda Grad: {np.min(x_grad[:, 1])}"
|
|
309
|
+
)
|
|
310
|
+
|
|
311
|
+
def render(self):
|
|
312
|
+
if self.renderer is None:
|
|
313
|
+
return
|
|
314
|
+
|
|
315
|
+
with wp.ScopedTimer("render"):
|
|
316
|
+
# draw trajectory
|
|
317
|
+
traj_verts = [np.mean(self.states[0].particle_q.numpy(), axis=0).tolist()]
|
|
318
|
+
for i in range(0, self.sim_steps, self.sim_substeps):
|
|
319
|
+
traj_verts.append(np.mean(self.states[i].particle_q.numpy(), axis=0).tolist())
|
|
320
|
+
|
|
321
|
+
self.renderer.begin_frame(self.render_time)
|
|
322
|
+
self.renderer.render(self.states[i])
|
|
323
|
+
self.renderer.render_box(
|
|
324
|
+
pos=self.target,
|
|
325
|
+
rot=wp.quat_identity(),
|
|
326
|
+
extents=(0.1, 0.1, 0.1),
|
|
327
|
+
name="target",
|
|
328
|
+
color=(0.0, 0.0, 0.0),
|
|
329
|
+
)
|
|
330
|
+
self.renderer.render_line_strip(
|
|
331
|
+
vertices=traj_verts,
|
|
332
|
+
color=wp.render.bourke_color_map(0.0, self.losses[0], self.losses[-1]),
|
|
333
|
+
radius=0.02,
|
|
334
|
+
name=f"traj_{self.iter - 1}",
|
|
335
|
+
)
|
|
336
|
+
self.renderer.end_frame()
|
|
337
|
+
|
|
338
|
+
from pxr import Gf, UsdGeom
|
|
339
|
+
|
|
340
|
+
particles_prim = self.renderer.stage.GetPrimAtPath("/root/particles")
|
|
341
|
+
particles = UsdGeom.Points.Get(self.renderer.stage, particles_prim.GetPath())
|
|
342
|
+
particles.CreateDisplayColorAttr().Set([Gf.Vec3f(1.0, 1.0, 1.0)], time=self.renderer.time)
|
|
343
|
+
|
|
344
|
+
self.render_time += self.frame_dt
|
|
345
|
+
|
|
346
|
+
|
|
347
|
+
if __name__ == "__main__":
|
|
348
|
+
import argparse
|
|
349
|
+
|
|
350
|
+
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
351
|
+
parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
|
|
352
|
+
parser.add_argument(
|
|
353
|
+
"--stage_path",
|
|
354
|
+
type=lambda x: None if x == "None" else str(x),
|
|
355
|
+
default="example_softbody_properties.usd",
|
|
356
|
+
help="Path to the output USD file.",
|
|
357
|
+
)
|
|
358
|
+
parser.add_argument(
|
|
359
|
+
"--train_iters",
|
|
360
|
+
type=int,
|
|
361
|
+
default=300,
|
|
362
|
+
help="Total number of training iterations.",
|
|
363
|
+
)
|
|
364
|
+
parser.add_argument(
|
|
365
|
+
"--material_behavior",
|
|
366
|
+
default="anisotropic",
|
|
367
|
+
choices=["anisotropic", "isotropic"],
|
|
368
|
+
help="Set material behavior to be Anisotropic or Isotropic.",
|
|
369
|
+
)
|
|
370
|
+
parser.add_argument(
|
|
371
|
+
"--verbose",
|
|
372
|
+
action="store_true",
|
|
373
|
+
help="Print out additional status messages during execution.",
|
|
374
|
+
)
|
|
375
|
+
|
|
376
|
+
args = parser.parse_known_args()[0]
|
|
377
|
+
|
|
378
|
+
with wp.ScopedDevice(args.device):
|
|
379
|
+
example = Example(
|
|
380
|
+
stage_path=args.stage_path,
|
|
381
|
+
material_behavior=args.material_behavior,
|
|
382
|
+
verbose=args.verbose,
|
|
383
|
+
)
|
|
384
|
+
|
|
385
|
+
# replay and optimize
|
|
386
|
+
for i in range(args.train_iters):
|
|
387
|
+
example.step()
|
|
388
|
+
if i == 0 or i % 50 == 0 or i == args.train_iters - 1:
|
|
389
|
+
example.render()
|
|
390
|
+
|
|
391
|
+
if example.renderer:
|
|
392
|
+
example.renderer.save()
|
|
@@ -60,8 +60,6 @@ class Example:
|
|
|
60
60
|
|
|
61
61
|
self.iter = 0
|
|
62
62
|
|
|
63
|
-
builder = wp.sim.ModelBuilder()
|
|
64
|
-
|
|
65
63
|
# add planar joints
|
|
66
64
|
builder = wp.sim.ModelBuilder(gravity=0.0)
|
|
67
65
|
builder.add_articulation()
|
|
@@ -153,7 +151,7 @@ class Example:
|
|
|
153
151
|
tape.backward(loss=self.loss)
|
|
154
152
|
|
|
155
153
|
if self.verbose and (self.iter + 1) % 10 == 0:
|
|
156
|
-
print(f"Iter {self.iter+1} Loss: {self.loss.numpy()[0]:.3f}")
|
|
154
|
+
print(f"Iter {self.iter + 1} Loss: {self.loss.numpy()[0]:.3f}")
|
|
157
155
|
|
|
158
156
|
assert not np.isnan(self.actions.grad.numpy()).any(), "NaN in gradient"
|
|
159
157
|
|
|
@@ -15,6 +15,11 @@
|
|
|
15
15
|
# simulated forward in time and then evaluated based on the center of mass
|
|
16
16
|
# momentum of the mesh.
|
|
17
17
|
#
|
|
18
|
+
# This example uses the deprecated wp.matmul() for matrix multiplication,
|
|
19
|
+
# which will be removed in a future version. See the updated version of
|
|
20
|
+
# this example, example_tile_walker.py, in examples/tile for the new
|
|
21
|
+
# approach to GEMMs using Warp's tile API.
|
|
22
|
+
#
|
|
18
23
|
###########################################################################
|
|
19
24
|
|
|
20
25
|
import math
|
|
@@ -100,7 +100,6 @@ class Example:
|
|
|
100
100
|
tri_ka=1e4,
|
|
101
101
|
tri_kd=1e-5,
|
|
102
102
|
edge_ke=100,
|
|
103
|
-
color_particles=True,
|
|
104
103
|
)
|
|
105
104
|
|
|
106
105
|
usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))
|
|
@@ -122,6 +121,9 @@ class Example:
|
|
|
122
121
|
kf=1.0e1,
|
|
123
122
|
)
|
|
124
123
|
|
|
124
|
+
if self.integrator_type == IntegratorType.VBD:
|
|
125
|
+
builder.color()
|
|
126
|
+
|
|
125
127
|
self.model = builder.finalize()
|
|
126
128
|
self.model.ground = True
|
|
127
129
|
self.model.soft_contact_ke = 1.0e4
|
|
@@ -0,0 +1,260 @@
|
|
|
1
|
+
# Copyright (c) 2024 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 Cloth Self Contact
|
|
10
|
+
#
|
|
11
|
+
# This simulation demonstrates twisting an FEM cloth model using the VBD
|
|
12
|
+
# integrator, showcasing its ability to handle complex self-contacts while
|
|
13
|
+
# ensuring it remains intersection-free.
|
|
14
|
+
#
|
|
15
|
+
###########################################################################
|
|
16
|
+
|
|
17
|
+
import math
|
|
18
|
+
import os
|
|
19
|
+
|
|
20
|
+
import numpy as np
|
|
21
|
+
from pxr import Usd, UsdGeom
|
|
22
|
+
|
|
23
|
+
import warp as wp
|
|
24
|
+
import warp.examples
|
|
25
|
+
import warp.sim
|
|
26
|
+
import warp.sim.render
|
|
27
|
+
from warp.sim.model import PARTICLE_FLAG_ACTIVE
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
@wp.kernel
|
|
31
|
+
def initialize_rotation(
|
|
32
|
+
particle_indices_to_rot: wp.array(dtype=wp.int32),
|
|
33
|
+
pos: wp.array(dtype=wp.vec3),
|
|
34
|
+
rot_centers: wp.array(dtype=wp.vec3),
|
|
35
|
+
rot_axes: wp.array(dtype=wp.vec3),
|
|
36
|
+
roots: wp.array(dtype=wp.vec3),
|
|
37
|
+
roots_to_ps: wp.array(dtype=wp.vec3),
|
|
38
|
+
):
|
|
39
|
+
tid = wp.tid()
|
|
40
|
+
particle_index = particle_indices_to_rot[wp.tid()]
|
|
41
|
+
|
|
42
|
+
p = pos[particle_index]
|
|
43
|
+
rot_center = rot_centers[tid]
|
|
44
|
+
rot_axis = rot_axes[tid]
|
|
45
|
+
op = p - rot_center
|
|
46
|
+
|
|
47
|
+
root = wp.dot(op, rot_axis) * rot_axis
|
|
48
|
+
|
|
49
|
+
root_to_p = p - root
|
|
50
|
+
|
|
51
|
+
roots[tid] = root
|
|
52
|
+
roots_to_ps[tid] = root_to_p
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
@wp.kernel
|
|
56
|
+
def apply_rotation(
|
|
57
|
+
time: float,
|
|
58
|
+
angular_velocity: float,
|
|
59
|
+
particle_indices_to_rot: wp.array(dtype=wp.int32),
|
|
60
|
+
rot_centers: wp.array(dtype=wp.vec3),
|
|
61
|
+
rot_axes: wp.array(dtype=wp.vec3),
|
|
62
|
+
roots: wp.array(dtype=wp.vec3),
|
|
63
|
+
roots_to_ps: wp.array(dtype=wp.vec3),
|
|
64
|
+
pos_0: wp.array(dtype=wp.vec3),
|
|
65
|
+
pos_1: wp.array(dtype=wp.vec3),
|
|
66
|
+
):
|
|
67
|
+
tid = wp.tid()
|
|
68
|
+
particle_index = particle_indices_to_rot[wp.tid()]
|
|
69
|
+
|
|
70
|
+
rot_axis = rot_axes[tid]
|
|
71
|
+
|
|
72
|
+
ux = rot_axis[0]
|
|
73
|
+
uy = rot_axis[1]
|
|
74
|
+
uz = rot_axis[2]
|
|
75
|
+
|
|
76
|
+
theta = time * angular_velocity
|
|
77
|
+
|
|
78
|
+
R = wp.mat33(
|
|
79
|
+
wp.cos(theta) + ux * ux * (1.0 - wp.cos(theta)),
|
|
80
|
+
ux * uy * (1.0 - wp.cos(theta)) - uz * wp.sin(theta),
|
|
81
|
+
ux * uz * (1.0 - wp.cos(theta)) + uy * wp.sin(theta),
|
|
82
|
+
uy * ux * (1.0 - wp.cos(theta)) + uz * wp.sin(theta),
|
|
83
|
+
wp.cos(theta) + uy * uy * (1.0 - wp.cos(theta)),
|
|
84
|
+
uy * uz * (1.0 - wp.cos(theta)) - ux * wp.sin(theta),
|
|
85
|
+
uz * ux * (1.0 - wp.cos(theta)) - uy * wp.sin(theta),
|
|
86
|
+
uz * uy * (1.0 - wp.cos(theta)) + ux * wp.sin(theta),
|
|
87
|
+
wp.cos(theta) + uz * uz * (1.0 - wp.cos(theta)),
|
|
88
|
+
)
|
|
89
|
+
|
|
90
|
+
root = roots[tid]
|
|
91
|
+
root_to_p = roots_to_ps[tid]
|
|
92
|
+
root_to_p_rot = R * root_to_p
|
|
93
|
+
p_rot = root + root_to_p_rot
|
|
94
|
+
|
|
95
|
+
pos_0[particle_index] = p_rot
|
|
96
|
+
pos_1[particle_index] = p_rot
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
class Example:
|
|
100
|
+
def __init__(self, stage_path="example_cloth_self_contact.usd", num_frames=1500):
|
|
101
|
+
fps = 60
|
|
102
|
+
self.frame_dt = 1.0 / fps
|
|
103
|
+
self.num_substeps = 10
|
|
104
|
+
self.iterations = 10
|
|
105
|
+
self.dt = self.frame_dt / self.num_substeps
|
|
106
|
+
|
|
107
|
+
self.num_frames = num_frames
|
|
108
|
+
self.sim_time = 0.0
|
|
109
|
+
self.profiler = {}
|
|
110
|
+
|
|
111
|
+
self.rot_angular_velocity = math.pi / 6
|
|
112
|
+
self.rot_end_time = 21
|
|
113
|
+
|
|
114
|
+
usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "square_cloth.usd"))
|
|
115
|
+
usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/cloth/cloth"))
|
|
116
|
+
|
|
117
|
+
mesh_points = np.array(usd_geom.GetPointsAttr().Get())
|
|
118
|
+
mesh_indices = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
|
|
119
|
+
|
|
120
|
+
self.input_scale_factor = 1.0
|
|
121
|
+
self.renderer_scale_factor = 0.01
|
|
122
|
+
|
|
123
|
+
vertices = [wp.vec3(v) * self.input_scale_factor for v in mesh_points]
|
|
124
|
+
self.faces = mesh_indices.reshape(-1, 3)
|
|
125
|
+
|
|
126
|
+
builder = wp.sim.ModelBuilder()
|
|
127
|
+
builder.add_cloth_mesh(
|
|
128
|
+
pos=wp.vec3(0.0, 0.0, 0.0),
|
|
129
|
+
rot=wp.quat_identity(),
|
|
130
|
+
scale=1.0,
|
|
131
|
+
vertices=vertices,
|
|
132
|
+
indices=mesh_indices,
|
|
133
|
+
vel=wp.vec3(0.0, 0.0, 0.0),
|
|
134
|
+
density=0.02,
|
|
135
|
+
tri_ke=1.0e5,
|
|
136
|
+
tri_ka=1.0e5,
|
|
137
|
+
tri_kd=3.0e-5,
|
|
138
|
+
)
|
|
139
|
+
builder.color()
|
|
140
|
+
self.model = builder.finalize()
|
|
141
|
+
self.model.ground = False
|
|
142
|
+
self.model.soft_contact_ke = 1.0e5
|
|
143
|
+
self.model.soft_contact_kd = 1.0e-6
|
|
144
|
+
self.model.soft_contact_mu = 0.2
|
|
145
|
+
|
|
146
|
+
# set up contact query and contact detection distances
|
|
147
|
+
self.model.soft_contact_radius = 0.2
|
|
148
|
+
self.model.soft_contact_margin = 0.35
|
|
149
|
+
|
|
150
|
+
cloth_size = 50
|
|
151
|
+
left_side = [cloth_size - 1 + i * cloth_size for i in range(cloth_size)]
|
|
152
|
+
right_side = [i * cloth_size for i in range(cloth_size)]
|
|
153
|
+
rot_point_indices = left_side + right_side
|
|
154
|
+
|
|
155
|
+
if len(rot_point_indices):
|
|
156
|
+
flags = self.model.particle_flags.numpy()
|
|
157
|
+
for fixed_vertex_id in rot_point_indices:
|
|
158
|
+
flags[fixed_vertex_id] = wp.uint32(int(flags[fixed_vertex_id]) & ~int(PARTICLE_FLAG_ACTIVE))
|
|
159
|
+
|
|
160
|
+
self.model.particle_flags = wp.array(flags)
|
|
161
|
+
|
|
162
|
+
self.integrator = wp.sim.VBDIntegrator(
|
|
163
|
+
self.model,
|
|
164
|
+
self.iterations,
|
|
165
|
+
handle_self_contact=True,
|
|
166
|
+
)
|
|
167
|
+
self.state0 = self.model.state()
|
|
168
|
+
self.state1 = self.model.state()
|
|
169
|
+
|
|
170
|
+
rot_axes = [[1, 0, 0]] * len(right_side) + [[-1, 0, 0]] * len(left_side)
|
|
171
|
+
|
|
172
|
+
self.rot_point_indices = wp.array(rot_point_indices, dtype=int)
|
|
173
|
+
self.rot_centers = wp.zeros(len(rot_point_indices), dtype=wp.vec3)
|
|
174
|
+
self.rot_axes = wp.array(rot_axes, dtype=wp.vec3)
|
|
175
|
+
|
|
176
|
+
self.roots = wp.zeros_like(self.rot_centers)
|
|
177
|
+
self.roots_to_ps = wp.zeros_like(self.rot_centers)
|
|
178
|
+
|
|
179
|
+
wp.launch(
|
|
180
|
+
kernel=initialize_rotation,
|
|
181
|
+
dim=self.rot_point_indices.shape[0],
|
|
182
|
+
inputs=[
|
|
183
|
+
self.rot_point_indices,
|
|
184
|
+
self.state0.particle_q,
|
|
185
|
+
self.rot_centers,
|
|
186
|
+
self.rot_axes,
|
|
187
|
+
self.roots,
|
|
188
|
+
self.roots_to_ps,
|
|
189
|
+
],
|
|
190
|
+
)
|
|
191
|
+
|
|
192
|
+
if stage_path:
|
|
193
|
+
self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=40.0)
|
|
194
|
+
else:
|
|
195
|
+
self.renderer = None
|
|
196
|
+
|
|
197
|
+
def step(self):
|
|
198
|
+
with wp.ScopedTimer("step", print=False, dict=self.profiler):
|
|
199
|
+
for _ in range(self.num_substeps):
|
|
200
|
+
if self.sim_time < self.rot_end_time:
|
|
201
|
+
wp.launch(
|
|
202
|
+
kernel=apply_rotation,
|
|
203
|
+
dim=self.rot_point_indices.shape[0],
|
|
204
|
+
inputs=[
|
|
205
|
+
self.sim_time,
|
|
206
|
+
self.rot_angular_velocity,
|
|
207
|
+
self.rot_point_indices,
|
|
208
|
+
self.rot_centers,
|
|
209
|
+
self.rot_axes,
|
|
210
|
+
self.roots,
|
|
211
|
+
self.roots_to_ps,
|
|
212
|
+
self.state0.particle_q,
|
|
213
|
+
self.state1.particle_q,
|
|
214
|
+
],
|
|
215
|
+
)
|
|
216
|
+
|
|
217
|
+
self.integrator.simulate(self.model, self.state0, self.state1, self.dt)
|
|
218
|
+
|
|
219
|
+
(self.state0, self.state1) = (self.state1, self.state0)
|
|
220
|
+
|
|
221
|
+
self.sim_time += self.dt
|
|
222
|
+
|
|
223
|
+
def render(self):
|
|
224
|
+
if self.renderer is None:
|
|
225
|
+
return
|
|
226
|
+
|
|
227
|
+
with wp.ScopedTimer("render", print=False):
|
|
228
|
+
self.renderer.begin_frame(self.sim_time)
|
|
229
|
+
self.renderer.render(self.state0)
|
|
230
|
+
self.renderer.end_frame()
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
if __name__ == "__main__":
|
|
234
|
+
import argparse
|
|
235
|
+
|
|
236
|
+
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
237
|
+
parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
|
|
238
|
+
parser.add_argument(
|
|
239
|
+
"--stage_path",
|
|
240
|
+
type=lambda x: None if x == "None" else str(x),
|
|
241
|
+
default="example_cloth_self_contact.usd",
|
|
242
|
+
help="Path to the output USD file.",
|
|
243
|
+
)
|
|
244
|
+
parser.add_argument("--num_frames", type=int, default=1500, help="Total number of frames.")
|
|
245
|
+
|
|
246
|
+
args = parser.parse_known_args()[0]
|
|
247
|
+
|
|
248
|
+
with wp.ScopedDevice(args.device):
|
|
249
|
+
example = Example(stage_path=args.stage_path, num_frames=args.num_frames)
|
|
250
|
+
|
|
251
|
+
for i in range(example.num_frames):
|
|
252
|
+
example.step()
|
|
253
|
+
example.render()
|
|
254
|
+
print(f"[{i:4d}/{example.num_frames}]")
|
|
255
|
+
|
|
256
|
+
frame_times = example.profiler["step"]
|
|
257
|
+
print("\nAverage frame sim time: {:.2f} ms".format(sum(frame_times) / len(frame_times)))
|
|
258
|
+
|
|
259
|
+
if example.renderer:
|
|
260
|
+
example.renderer.save()
|