warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__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.
- warp/__init__.py +301 -287
- warp/__init__.pyi +2220 -313
- warp/_src/__init__.py +14 -0
- warp/_src/autograd.py +1075 -0
- warp/_src/build.py +618 -0
- warp/_src/build_dll.py +640 -0
- warp/{builtins.py → _src/builtins.py} +1497 -226
- warp/_src/codegen.py +4359 -0
- warp/{config.py → _src/config.py} +178 -169
- warp/_src/constants.py +57 -0
- warp/_src/context.py +8294 -0
- warp/_src/dlpack.py +462 -0
- warp/_src/fabric.py +355 -0
- warp/_src/fem/__init__.py +14 -0
- warp/_src/fem/adaptivity.py +508 -0
- warp/_src/fem/cache.py +687 -0
- warp/_src/fem/dirichlet.py +188 -0
- warp/{fem → _src/fem}/domain.py +40 -30
- warp/_src/fem/field/__init__.py +131 -0
- warp/_src/fem/field/field.py +701 -0
- warp/{fem → _src/fem}/field/nodal_field.py +30 -15
- warp/{fem → _src/fem}/field/restriction.py +1 -1
- warp/{fem → _src/fem}/field/virtual.py +53 -27
- warp/_src/fem/geometry/__init__.py +32 -0
- warp/{fem → _src/fem}/geometry/adaptive_nanogrid.py +77 -163
- warp/_src/fem/geometry/closest_point.py +97 -0
- warp/{fem → _src/fem}/geometry/deformed_geometry.py +14 -22
- warp/{fem → _src/fem}/geometry/element.py +32 -10
- warp/{fem → _src/fem}/geometry/geometry.py +48 -20
- warp/{fem → _src/fem}/geometry/grid_2d.py +12 -23
- warp/{fem → _src/fem}/geometry/grid_3d.py +12 -23
- warp/{fem → _src/fem}/geometry/hexmesh.py +40 -63
- warp/{fem → _src/fem}/geometry/nanogrid.py +255 -248
- warp/{fem → _src/fem}/geometry/partition.py +121 -63
- warp/{fem → _src/fem}/geometry/quadmesh.py +26 -45
- warp/{fem → _src/fem}/geometry/tetmesh.py +40 -63
- warp/{fem → _src/fem}/geometry/trimesh.py +26 -45
- warp/{fem → _src/fem}/integrate.py +164 -158
- warp/_src/fem/linalg.py +383 -0
- warp/_src/fem/operator.py +396 -0
- warp/_src/fem/polynomial.py +229 -0
- warp/{fem → _src/fem}/quadrature/pic_quadrature.py +15 -20
- warp/{fem → _src/fem}/quadrature/quadrature.py +95 -47
- warp/_src/fem/space/__init__.py +248 -0
- warp/{fem → _src/fem}/space/basis_function_space.py +20 -11
- warp/_src/fem/space/basis_space.py +679 -0
- warp/{fem → _src/fem}/space/dof_mapper.py +3 -3
- warp/{fem → _src/fem}/space/function_space.py +14 -13
- warp/{fem → _src/fem}/space/grid_2d_function_space.py +4 -7
- warp/{fem → _src/fem}/space/grid_3d_function_space.py +4 -4
- warp/{fem → _src/fem}/space/hexmesh_function_space.py +4 -10
- warp/{fem → _src/fem}/space/nanogrid_function_space.py +3 -9
- warp/{fem → _src/fem}/space/partition.py +117 -60
- warp/{fem → _src/fem}/space/quadmesh_function_space.py +4 -10
- warp/{fem → _src/fem}/space/restriction.py +66 -33
- warp/_src/fem/space/shape/__init__.py +152 -0
- warp/{fem → _src/fem}/space/shape/cube_shape_function.py +9 -9
- warp/{fem → _src/fem}/space/shape/shape_function.py +8 -9
- warp/{fem → _src/fem}/space/shape/square_shape_function.py +6 -6
- warp/{fem → _src/fem}/space/shape/tet_shape_function.py +3 -3
- warp/{fem → _src/fem}/space/shape/triangle_shape_function.py +3 -3
- warp/{fem → _src/fem}/space/tetmesh_function_space.py +3 -9
- warp/_src/fem/space/topology.py +459 -0
- warp/{fem → _src/fem}/space/trimesh_function_space.py +3 -9
- warp/_src/fem/types.py +112 -0
- warp/_src/fem/utils.py +486 -0
- warp/_src/jax.py +186 -0
- warp/_src/jax_experimental/__init__.py +14 -0
- warp/_src/jax_experimental/custom_call.py +387 -0
- warp/_src/jax_experimental/ffi.py +1284 -0
- warp/_src/jax_experimental/xla_ffi.py +656 -0
- warp/_src/marching_cubes.py +708 -0
- warp/_src/math.py +414 -0
- warp/_src/optim/__init__.py +14 -0
- warp/_src/optim/adam.py +163 -0
- warp/_src/optim/linear.py +1606 -0
- warp/_src/optim/sgd.py +112 -0
- warp/_src/paddle.py +406 -0
- warp/_src/render/__init__.py +14 -0
- warp/_src/render/imgui_manager.py +289 -0
- warp/_src/render/render_opengl.py +3636 -0
- warp/_src/render/render_usd.py +937 -0
- warp/_src/render/utils.py +160 -0
- warp/_src/sparse.py +2716 -0
- warp/_src/tape.py +1206 -0
- warp/{thirdparty → _src/thirdparty}/unittest_parallel.py +9 -2
- warp/_src/torch.py +391 -0
- warp/_src/types.py +5870 -0
- warp/_src/utils.py +1693 -0
- warp/autograd.py +12 -1054
- warp/bin/warp-clang.dll +0 -0
- warp/bin/warp.dll +0 -0
- warp/build.py +8 -588
- warp/build_dll.py +6 -471
- warp/codegen.py +6 -4246
- warp/constants.py +6 -39
- warp/context.py +12 -7851
- warp/dlpack.py +6 -444
- warp/examples/distributed/example_jacobi_mpi.py +4 -5
- warp/examples/fem/example_adaptive_grid.py +1 -1
- warp/examples/fem/example_apic_fluid.py +1 -1
- warp/examples/fem/example_burgers.py +8 -8
- warp/examples/fem/example_diffusion.py +1 -1
- warp/examples/fem/example_distortion_energy.py +1 -1
- warp/examples/fem/example_mixed_elasticity.py +2 -2
- warp/examples/fem/example_navier_stokes.py +1 -1
- warp/examples/fem/example_nonconforming_contact.py +7 -7
- warp/examples/fem/example_stokes.py +1 -1
- warp/examples/fem/example_stokes_transfer.py +1 -1
- warp/examples/fem/utils.py +2 -2
- warp/examples/interop/example_jax_callable.py +1 -1
- warp/examples/interop/example_jax_ffi_callback.py +1 -1
- warp/examples/interop/example_jax_kernel.py +3 -2
- warp/examples/tile/example_tile_mcgp.py +191 -0
- warp/fabric.py +6 -337
- warp/fem/__init__.py +159 -97
- warp/fem/adaptivity.py +7 -489
- warp/fem/cache.py +9 -648
- warp/fem/dirichlet.py +6 -184
- warp/fem/field/__init__.py +8 -109
- warp/fem/field/field.py +7 -652
- warp/fem/geometry/__init__.py +7 -18
- warp/fem/geometry/closest_point.py +11 -77
- warp/fem/linalg.py +18 -366
- warp/fem/operator.py +11 -369
- warp/fem/polynomial.py +9 -209
- warp/fem/space/__init__.py +5 -211
- warp/fem/space/basis_space.py +6 -662
- warp/fem/space/shape/__init__.py +41 -118
- warp/fem/space/topology.py +6 -437
- warp/fem/types.py +6 -81
- warp/fem/utils.py +11 -444
- warp/jax.py +8 -165
- warp/jax_experimental/__init__.py +14 -1
- warp/jax_experimental/custom_call.py +8 -342
- warp/jax_experimental/ffi.py +17 -853
- warp/jax_experimental/xla_ffi.py +5 -596
- warp/marching_cubes.py +5 -689
- warp/math.py +16 -393
- warp/native/array.h +385 -37
- warp/native/builtin.h +316 -39
- warp/native/bvh.cpp +43 -9
- warp/native/bvh.cu +62 -27
- warp/native/bvh.h +310 -309
- warp/native/clang/clang.cpp +102 -97
- warp/native/coloring.cpp +0 -1
- warp/native/crt.h +208 -0
- warp/native/exports.h +156 -0
- warp/native/hashgrid.cu +2 -0
- warp/native/intersect.h +24 -1
- warp/native/intersect_tri.h +44 -35
- warp/native/mat.h +1456 -276
- warp/native/mesh.cpp +4 -4
- warp/native/mesh.cu +4 -2
- warp/native/mesh.h +176 -61
- warp/native/quat.h +0 -52
- warp/native/scan.cu +2 -0
- warp/native/sort.cu +22 -13
- warp/native/sort.h +2 -0
- warp/native/sparse.cu +7 -3
- warp/native/spatial.h +12 -0
- warp/native/tile.h +837 -70
- warp/native/tile_radix_sort.h +1 -1
- warp/native/tile_reduce.h +394 -46
- warp/native/tile_scan.h +4 -4
- warp/native/vec.h +469 -53
- warp/native/version.h +23 -0
- warp/native/volume.cpp +1 -1
- warp/native/volume.cu +1 -0
- warp/native/volume.h +1 -1
- warp/native/volume_builder.cu +2 -0
- warp/native/warp.cpp +60 -32
- warp/native/warp.cu +313 -201
- warp/native/warp.h +14 -11
- warp/optim/__init__.py +6 -3
- warp/optim/adam.py +6 -145
- warp/optim/linear.py +14 -1585
- warp/optim/sgd.py +6 -94
- warp/paddle.py +6 -388
- warp/render/__init__.py +8 -4
- warp/render/imgui_manager.py +7 -267
- warp/render/render_opengl.py +6 -3616
- warp/render/render_usd.py +6 -918
- warp/render/utils.py +6 -142
- warp/sparse.py +37 -2563
- warp/tape.py +6 -1188
- warp/tests/__main__.py +1 -1
- warp/tests/cuda/test_async.py +4 -4
- warp/tests/cuda/test_conditional_captures.py +1 -1
- warp/tests/cuda/test_multigpu.py +1 -1
- warp/tests/cuda/test_streams.py +58 -1
- warp/tests/geometry/test_bvh.py +157 -22
- warp/tests/geometry/test_hash_grid.py +38 -0
- warp/tests/geometry/test_marching_cubes.py +0 -1
- warp/tests/geometry/test_mesh.py +5 -3
- warp/tests/geometry/test_mesh_query_aabb.py +5 -12
- warp/tests/geometry/test_mesh_query_point.py +5 -2
- warp/tests/geometry/test_mesh_query_ray.py +15 -3
- warp/tests/geometry/test_volume_write.py +5 -5
- warp/tests/interop/test_dlpack.py +14 -14
- warp/tests/interop/test_jax.py +1382 -79
- warp/tests/interop/test_paddle.py +1 -1
- warp/tests/test_adam.py +0 -1
- warp/tests/test_arithmetic.py +9 -9
- warp/tests/test_array.py +529 -100
- warp/tests/test_array_reduce.py +3 -3
- warp/tests/test_atomic.py +12 -8
- warp/tests/test_atomic_bitwise.py +209 -0
- warp/tests/test_atomic_cas.py +4 -4
- warp/tests/test_bool.py +2 -2
- warp/tests/test_builtins_resolution.py +5 -571
- warp/tests/test_codegen.py +34 -15
- warp/tests/test_conditional.py +1 -1
- warp/tests/test_context.py +6 -6
- warp/tests/test_copy.py +242 -161
- warp/tests/test_ctypes.py +3 -3
- warp/tests/test_devices.py +24 -2
- warp/tests/test_examples.py +16 -84
- warp/tests/test_fabricarray.py +35 -35
- warp/tests/test_fast_math.py +0 -2
- warp/tests/test_fem.py +60 -14
- warp/tests/test_fixedarray.py +3 -3
- warp/tests/test_func.py +8 -5
- warp/tests/test_generics.py +1 -1
- warp/tests/test_indexedarray.py +24 -24
- warp/tests/test_intersect.py +39 -9
- warp/tests/test_large.py +1 -1
- warp/tests/test_lerp.py +3 -1
- warp/tests/test_linear_solvers.py +1 -1
- warp/tests/test_map.py +49 -4
- warp/tests/test_mat.py +52 -62
- warp/tests/test_mat_constructors.py +4 -5
- warp/tests/test_mat_lite.py +1 -1
- warp/tests/test_mat_scalar_ops.py +121 -121
- warp/tests/test_math.py +34 -0
- warp/tests/test_module_aot.py +4 -4
- warp/tests/test_modules_lite.py +28 -2
- warp/tests/test_print.py +11 -11
- warp/tests/test_quat.py +93 -58
- warp/tests/test_runlength_encode.py +1 -1
- warp/tests/test_scalar_ops.py +38 -10
- warp/tests/test_smoothstep.py +1 -1
- warp/tests/test_sparse.py +126 -15
- warp/tests/test_spatial.py +105 -87
- warp/tests/test_special_values.py +6 -6
- warp/tests/test_static.py +7 -7
- warp/tests/test_struct.py +13 -2
- warp/tests/test_triangle_closest_point.py +48 -1
- warp/tests/test_tuple.py +96 -0
- warp/tests/test_types.py +82 -9
- warp/tests/test_utils.py +52 -52
- warp/tests/test_vec.py +29 -29
- warp/tests/test_vec_constructors.py +5 -5
- warp/tests/test_vec_scalar_ops.py +97 -97
- warp/tests/test_version.py +75 -0
- warp/tests/tile/test_tile.py +239 -0
- warp/tests/tile/test_tile_atomic_bitwise.py +403 -0
- warp/tests/tile/test_tile_cholesky.py +7 -4
- warp/tests/tile/test_tile_load.py +26 -2
- warp/tests/tile/test_tile_mathdx.py +3 -3
- warp/tests/tile/test_tile_matmul.py +1 -1
- warp/tests/tile/test_tile_mlp.py +2 -4
- warp/tests/tile/test_tile_reduce.py +214 -13
- warp/tests/unittest_suites.py +6 -14
- warp/tests/unittest_utils.py +10 -9
- warp/tests/walkthrough_debug.py +3 -1
- warp/torch.py +6 -373
- warp/types.py +29 -5750
- warp/utils.py +10 -1659
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/METADATA +47 -103
- warp_lang-1.10.0rc2.dist-info/RECORD +468 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/Gaia-LICENSE.txt +6 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/appdirs-LICENSE.txt +22 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/asset_pixel_jpg-LICENSE.txt +3 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/cuda-LICENSE.txt +1582 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/dlpack-LICENSE.txt +201 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/fp16-LICENSE.txt +28 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/libmathdx-LICENSE.txt +220 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/llvm-LICENSE.txt +279 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/moller-LICENSE.txt +16 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/nanovdb-LICENSE.txt +2 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/nvrtc-LICENSE.txt +1592 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/svd-LICENSE.txt +23 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/unittest_parallel-LICENSE.txt +21 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/usd-LICENSE.txt +213 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/windingnumber-LICENSE.txt +21 -0
- warp/examples/assets/cartpole.urdf +0 -110
- warp/examples/assets/crazyflie.usd +0 -0
- warp/examples/assets/nv_ant.xml +0 -92
- warp/examples/assets/nv_humanoid.xml +0 -183
- warp/examples/assets/quadruped.urdf +0 -268
- warp/examples/optim/example_bounce.py +0 -266
- warp/examples/optim/example_cloth_throw.py +0 -228
- warp/examples/optim/example_drone.py +0 -870
- warp/examples/optim/example_inverse_kinematics.py +0 -182
- warp/examples/optim/example_inverse_kinematics_torch.py +0 -191
- warp/examples/optim/example_softbody_properties.py +0 -400
- warp/examples/optim/example_spring_cage.py +0 -245
- warp/examples/optim/example_trajectory.py +0 -227
- warp/examples/sim/example_cartpole.py +0 -143
- warp/examples/sim/example_cloth.py +0 -225
- warp/examples/sim/example_cloth_self_contact.py +0 -316
- warp/examples/sim/example_granular.py +0 -130
- warp/examples/sim/example_granular_collision_sdf.py +0 -202
- warp/examples/sim/example_jacobian_ik.py +0 -244
- warp/examples/sim/example_particle_chain.py +0 -124
- warp/examples/sim/example_quadruped.py +0 -203
- warp/examples/sim/example_rigid_chain.py +0 -203
- warp/examples/sim/example_rigid_contact.py +0 -195
- warp/examples/sim/example_rigid_force.py +0 -133
- warp/examples/sim/example_rigid_gyroscopic.py +0 -115
- warp/examples/sim/example_rigid_soft_contact.py +0 -140
- warp/examples/sim/example_soft_body.py +0 -196
- warp/examples/tile/example_tile_walker.py +0 -327
- warp/sim/__init__.py +0 -74
- warp/sim/articulation.py +0 -793
- warp/sim/collide.py +0 -2570
- warp/sim/graph_coloring.py +0 -307
- warp/sim/import_mjcf.py +0 -791
- warp/sim/import_snu.py +0 -227
- warp/sim/import_urdf.py +0 -579
- warp/sim/import_usd.py +0 -898
- warp/sim/inertia.py +0 -357
- warp/sim/integrator.py +0 -245
- warp/sim/integrator_euler.py +0 -2000
- warp/sim/integrator_featherstone.py +0 -2101
- warp/sim/integrator_vbd.py +0 -2487
- warp/sim/integrator_xpbd.py +0 -3295
- warp/sim/model.py +0 -4821
- warp/sim/particles.py +0 -121
- warp/sim/render.py +0 -431
- warp/sim/utils.py +0 -431
- warp/tests/sim/disabled_kinematics.py +0 -244
- warp/tests/sim/test_cloth.py +0 -863
- warp/tests/sim/test_collision.py +0 -743
- warp/tests/sim/test_coloring.py +0 -347
- warp/tests/sim/test_inertia.py +0 -161
- warp/tests/sim/test_model.py +0 -226
- warp/tests/sim/test_sim_grad.py +0 -287
- warp/tests/sim/test_sim_grad_bounce_linear.py +0 -212
- warp/tests/sim/test_sim_kinematics.py +0 -98
- warp/thirdparty/__init__.py +0 -0
- warp_lang-1.9.0.dist-info/RECORD +0 -456
- /warp/{fem → _src/fem}/quadrature/__init__.py +0 -0
- /warp/{tests/sim → _src/thirdparty}/__init__.py +0 -0
- /warp/{thirdparty → _src/thirdparty}/appdirs.py +0 -0
- /warp/{thirdparty → _src/thirdparty}/dlpack.py +0 -0
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/WHEEL +0 -0
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/licenses/LICENSE.md +0 -0
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/top_level.txt +0 -0
|
@@ -1,140 +0,0 @@
|
|
|
1
|
-
# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
-
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
-
#
|
|
4
|
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
-
# you may not use this file except in compliance with the License.
|
|
6
|
-
# You may obtain a copy of the License at
|
|
7
|
-
#
|
|
8
|
-
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
-
#
|
|
10
|
-
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
-
# See the License for the specific language governing permissions and
|
|
14
|
-
# limitations under the License.
|
|
15
|
-
|
|
16
|
-
###########################################################################
|
|
17
|
-
# Example Sim Rigid FEM
|
|
18
|
-
#
|
|
19
|
-
# Shows how to set up a rigid sphere colliding with an FEM beam
|
|
20
|
-
# using wp.sim.ModelBuilder().
|
|
21
|
-
#
|
|
22
|
-
###########################################################################
|
|
23
|
-
|
|
24
|
-
import warp as wp
|
|
25
|
-
import warp.sim
|
|
26
|
-
import warp.sim.render
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
class Example:
|
|
30
|
-
def __init__(self, stage_path="example_rigid_soft_contact.usd"):
|
|
31
|
-
self.sim_width = 8
|
|
32
|
-
self.sim_height = 8
|
|
33
|
-
|
|
34
|
-
fps = 60
|
|
35
|
-
self.frame_dt = 1.0 / fps
|
|
36
|
-
self.sim_substeps = 32
|
|
37
|
-
self.sim_dt = self.frame_dt / self.sim_substeps
|
|
38
|
-
self.sim_time = 0.0
|
|
39
|
-
self.sim_iterations = 1
|
|
40
|
-
self.sim_relaxation = 1.0
|
|
41
|
-
self.profiler = {}
|
|
42
|
-
|
|
43
|
-
builder = wp.sim.ModelBuilder()
|
|
44
|
-
builder.default_particle_radius = 0.01
|
|
45
|
-
|
|
46
|
-
builder.add_soft_grid(
|
|
47
|
-
pos=wp.vec3(0.0, 0.0, 0.0),
|
|
48
|
-
rot=wp.quat_identity(),
|
|
49
|
-
vel=wp.vec3(0.0, 0.0, 0.0),
|
|
50
|
-
dim_x=20,
|
|
51
|
-
dim_y=10,
|
|
52
|
-
dim_z=10,
|
|
53
|
-
cell_x=0.1,
|
|
54
|
-
cell_y=0.1,
|
|
55
|
-
cell_z=0.1,
|
|
56
|
-
density=100.0,
|
|
57
|
-
k_mu=50000.0,
|
|
58
|
-
k_lambda=20000.0,
|
|
59
|
-
k_damp=0.0,
|
|
60
|
-
)
|
|
61
|
-
|
|
62
|
-
b = builder.add_body(origin=wp.transform((0.5, 2.5, 0.5), wp.quat_identity()))
|
|
63
|
-
builder.add_shape_sphere(body=b, radius=0.75, density=100.0)
|
|
64
|
-
|
|
65
|
-
self.model = builder.finalize()
|
|
66
|
-
self.model.ground = True
|
|
67
|
-
self.model.soft_contact_ke = 1.0e3
|
|
68
|
-
self.model.soft_contact_kd = 0.0
|
|
69
|
-
self.model.soft_contact_kf = 1.0e3
|
|
70
|
-
|
|
71
|
-
self.integrator = wp.sim.SemiImplicitIntegrator()
|
|
72
|
-
|
|
73
|
-
self.state_0 = self.model.state()
|
|
74
|
-
self.state_1 = self.model.state()
|
|
75
|
-
|
|
76
|
-
if stage_path:
|
|
77
|
-
self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=1.0)
|
|
78
|
-
else:
|
|
79
|
-
self.renderer = None
|
|
80
|
-
|
|
81
|
-
self.use_cuda_graph = wp.get_device().is_cuda
|
|
82
|
-
if self.use_cuda_graph:
|
|
83
|
-
with wp.ScopedCapture() as capture:
|
|
84
|
-
self.simulate()
|
|
85
|
-
self.graph = capture.graph
|
|
86
|
-
|
|
87
|
-
def simulate(self):
|
|
88
|
-
for _s in range(self.sim_substeps):
|
|
89
|
-
wp.sim.collide(self.model, self.state_0)
|
|
90
|
-
|
|
91
|
-
self.state_0.clear_forces()
|
|
92
|
-
self.state_1.clear_forces()
|
|
93
|
-
|
|
94
|
-
self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
|
|
95
|
-
|
|
96
|
-
# swap states
|
|
97
|
-
(self.state_0, self.state_1) = (self.state_1, self.state_0)
|
|
98
|
-
|
|
99
|
-
def step(self):
|
|
100
|
-
with wp.ScopedTimer("step", dict=self.profiler):
|
|
101
|
-
if self.use_cuda_graph:
|
|
102
|
-
wp.capture_launch(self.graph)
|
|
103
|
-
else:
|
|
104
|
-
self.simulate()
|
|
105
|
-
self.sim_time += self.frame_dt
|
|
106
|
-
|
|
107
|
-
def render(self):
|
|
108
|
-
if self.renderer is None:
|
|
109
|
-
return
|
|
110
|
-
|
|
111
|
-
with wp.ScopedTimer("render"):
|
|
112
|
-
self.renderer.begin_frame(self.sim_time)
|
|
113
|
-
self.renderer.render(self.state_0)
|
|
114
|
-
self.renderer.end_frame()
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
if __name__ == "__main__":
|
|
118
|
-
import argparse
|
|
119
|
-
|
|
120
|
-
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
121
|
-
parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
|
|
122
|
-
parser.add_argument(
|
|
123
|
-
"--stage_path",
|
|
124
|
-
type=lambda x: None if x == "None" else str(x),
|
|
125
|
-
default="example_rigid_soft_contact.usd",
|
|
126
|
-
help="Path to the output USD file.",
|
|
127
|
-
)
|
|
128
|
-
parser.add_argument("--num_frames", type=int, default=300, help="Total number of frames.")
|
|
129
|
-
|
|
130
|
-
args = parser.parse_known_args()[0]
|
|
131
|
-
|
|
132
|
-
with wp.ScopedDevice(args.device):
|
|
133
|
-
example = Example(stage_path=args.stage_path)
|
|
134
|
-
|
|
135
|
-
for _ in range(args.num_frames):
|
|
136
|
-
example.step()
|
|
137
|
-
example.render()
|
|
138
|
-
|
|
139
|
-
if example.renderer:
|
|
140
|
-
example.renderer.save()
|
|
@@ -1,196 +0,0 @@
|
|
|
1
|
-
# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
-
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
-
#
|
|
4
|
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
-
# you may not use this file except in compliance with the License.
|
|
6
|
-
# You may obtain a copy of the License at
|
|
7
|
-
#
|
|
8
|
-
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
-
#
|
|
10
|
-
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
-
# See the License for the specific language governing permissions and
|
|
14
|
-
# limitations under the License.
|
|
15
|
-
|
|
16
|
-
###########################################################################
|
|
17
|
-
# Example Sim Neo-Hookean
|
|
18
|
-
#
|
|
19
|
-
# Shows a simulation of an Neo-Hookean FEM beam being twisted through a
|
|
20
|
-
# 180 degree rotation.
|
|
21
|
-
#
|
|
22
|
-
###########################################################################
|
|
23
|
-
import math
|
|
24
|
-
|
|
25
|
-
import warp as wp
|
|
26
|
-
import warp.sim
|
|
27
|
-
import warp.sim.render
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
@wp.kernel
|
|
31
|
-
def twist_points(
|
|
32
|
-
rest: wp.array(dtype=wp.vec3), points: wp.array(dtype=wp.vec3), mass: wp.array(dtype=float), xform: wp.transform
|
|
33
|
-
):
|
|
34
|
-
tid = wp.tid()
|
|
35
|
-
|
|
36
|
-
r = rest[tid]
|
|
37
|
-
p = points[tid]
|
|
38
|
-
m = mass[tid]
|
|
39
|
-
|
|
40
|
-
# twist the top layer of particles in the beam
|
|
41
|
-
if m == 0 and p[1] != 0.0:
|
|
42
|
-
points[tid] = wp.transform_point(xform, r)
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
@wp.kernel
|
|
46
|
-
def compute_volume(points: wp.array(dtype=wp.vec3), indices: wp.array2d(dtype=int), volume: wp.array(dtype=float)):
|
|
47
|
-
tid = wp.tid()
|
|
48
|
-
|
|
49
|
-
i = indices[tid, 0]
|
|
50
|
-
j = indices[tid, 1]
|
|
51
|
-
k = indices[tid, 2]
|
|
52
|
-
l = indices[tid, 3]
|
|
53
|
-
|
|
54
|
-
x0 = points[i]
|
|
55
|
-
x1 = points[j]
|
|
56
|
-
x2 = points[k]
|
|
57
|
-
x3 = points[l]
|
|
58
|
-
|
|
59
|
-
x10 = x1 - x0
|
|
60
|
-
x20 = x2 - x0
|
|
61
|
-
x30 = x3 - x0
|
|
62
|
-
|
|
63
|
-
v = wp.dot(x10, wp.cross(x20, x30)) / 6.0
|
|
64
|
-
|
|
65
|
-
wp.atomic_add(volume, 0, v)
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
class Example:
|
|
69
|
-
def __init__(self, stage_path="example_soft_body.usd", num_frames=300):
|
|
70
|
-
self.sim_substeps = 64
|
|
71
|
-
self.num_frames = num_frames
|
|
72
|
-
fps = 60
|
|
73
|
-
sim_duration = self.num_frames / fps
|
|
74
|
-
self.frame_dt = 1.0 / fps
|
|
75
|
-
self.sim_dt = self.frame_dt / self.sim_substeps
|
|
76
|
-
self.sim_time = 0.0
|
|
77
|
-
self.lift_speed = 2.5 / sim_duration * 2.0 # from Smith et al.
|
|
78
|
-
self.rot_speed = math.pi / sim_duration
|
|
79
|
-
|
|
80
|
-
builder = wp.sim.ModelBuilder()
|
|
81
|
-
|
|
82
|
-
cell_dim = 15
|
|
83
|
-
cell_size = 2.0 / cell_dim
|
|
84
|
-
|
|
85
|
-
center = cell_size * cell_dim * 0.5
|
|
86
|
-
|
|
87
|
-
builder.add_soft_grid(
|
|
88
|
-
pos=wp.vec3(-center, 0.0, -center),
|
|
89
|
-
rot=wp.quat_identity(),
|
|
90
|
-
vel=wp.vec3(0.0, 0.0, 0.0),
|
|
91
|
-
dim_x=cell_dim,
|
|
92
|
-
dim_y=cell_dim,
|
|
93
|
-
dim_z=cell_dim,
|
|
94
|
-
cell_x=cell_size,
|
|
95
|
-
cell_y=cell_size,
|
|
96
|
-
cell_z=cell_size,
|
|
97
|
-
density=100.0,
|
|
98
|
-
fix_bottom=True,
|
|
99
|
-
fix_top=True,
|
|
100
|
-
k_mu=1000.0,
|
|
101
|
-
k_lambda=5000.0,
|
|
102
|
-
k_damp=0.0,
|
|
103
|
-
)
|
|
104
|
-
|
|
105
|
-
self.model = builder.finalize()
|
|
106
|
-
self.model.ground = False
|
|
107
|
-
self.model.gravity[1] = 0.0
|
|
108
|
-
|
|
109
|
-
self.integrator = wp.sim.SemiImplicitIntegrator()
|
|
110
|
-
|
|
111
|
-
self.rest = self.model.state()
|
|
112
|
-
self.rest_vol = (cell_size * cell_dim) ** 3
|
|
113
|
-
|
|
114
|
-
self.state_0 = self.model.state()
|
|
115
|
-
self.state_1 = self.model.state()
|
|
116
|
-
|
|
117
|
-
self.volume = wp.zeros(1, dtype=wp.float32)
|
|
118
|
-
|
|
119
|
-
if stage_path:
|
|
120
|
-
self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=20.0)
|
|
121
|
-
else:
|
|
122
|
-
self.renderer = None
|
|
123
|
-
|
|
124
|
-
self.use_cuda_graph = wp.get_device().is_cuda
|
|
125
|
-
if self.use_cuda_graph:
|
|
126
|
-
with wp.ScopedCapture() as capture:
|
|
127
|
-
self.simulate()
|
|
128
|
-
self.graph = capture.graph
|
|
129
|
-
|
|
130
|
-
def simulate(self):
|
|
131
|
-
for _ in range(self.sim_substeps):
|
|
132
|
-
self.state_0.clear_forces()
|
|
133
|
-
self.state_1.clear_forces()
|
|
134
|
-
|
|
135
|
-
self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
|
|
136
|
-
|
|
137
|
-
# swap states
|
|
138
|
-
(self.state_0, self.state_1) = (self.state_1, self.state_0)
|
|
139
|
-
|
|
140
|
-
def step(self):
|
|
141
|
-
with wp.ScopedTimer("step"):
|
|
142
|
-
xform = wp.transform(
|
|
143
|
-
(0.0, self.lift_speed * self.sim_time, 0.0),
|
|
144
|
-
wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), self.rot_speed * self.sim_time),
|
|
145
|
-
)
|
|
146
|
-
wp.launch(
|
|
147
|
-
kernel=twist_points,
|
|
148
|
-
dim=len(self.state_0.particle_q),
|
|
149
|
-
inputs=[self.rest.particle_q, self.state_0.particle_q, self.model.particle_mass, xform],
|
|
150
|
-
)
|
|
151
|
-
if self.use_cuda_graph:
|
|
152
|
-
wp.capture_launch(self.graph)
|
|
153
|
-
else:
|
|
154
|
-
self.simulate()
|
|
155
|
-
self.volume.zero_()
|
|
156
|
-
wp.launch(
|
|
157
|
-
kernel=compute_volume,
|
|
158
|
-
dim=self.model.tet_count,
|
|
159
|
-
inputs=[self.state_0.particle_q, self.model.tet_indices, self.volume],
|
|
160
|
-
)
|
|
161
|
-
self.sim_time += self.frame_dt
|
|
162
|
-
|
|
163
|
-
def render(self):
|
|
164
|
-
if self.renderer is None:
|
|
165
|
-
return
|
|
166
|
-
|
|
167
|
-
with wp.ScopedTimer("render"):
|
|
168
|
-
self.renderer.begin_frame(self.sim_time)
|
|
169
|
-
self.renderer.render(self.state_0)
|
|
170
|
-
self.renderer.end_frame()
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
if __name__ == "__main__":
|
|
174
|
-
import argparse
|
|
175
|
-
|
|
176
|
-
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
177
|
-
parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
|
|
178
|
-
parser.add_argument(
|
|
179
|
-
"--stage_path",
|
|
180
|
-
type=lambda x: None if x == "None" else str(x),
|
|
181
|
-
default="example_soft_body.usd",
|
|
182
|
-
help="Path to the output USD file.",
|
|
183
|
-
)
|
|
184
|
-
parser.add_argument("--num_frames", type=int, default=300, help="Total number of frames.")
|
|
185
|
-
|
|
186
|
-
args = parser.parse_known_args()[0]
|
|
187
|
-
|
|
188
|
-
with wp.ScopedDevice(args.device):
|
|
189
|
-
example = Example(stage_path=args.stage_path, num_frames=args.num_frames)
|
|
190
|
-
|
|
191
|
-
for _ in range(args.num_frames):
|
|
192
|
-
example.step()
|
|
193
|
-
example.render()
|
|
194
|
-
|
|
195
|
-
if example.renderer:
|
|
196
|
-
example.renderer.save()
|
|
@@ -1,327 +0,0 @@
|
|
|
1
|
-
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
-
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
-
#
|
|
4
|
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
-
# you may not use this file except in compliance with the License.
|
|
6
|
-
# You may obtain a copy of the License at
|
|
7
|
-
#
|
|
8
|
-
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
-
#
|
|
10
|
-
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
-
# See the License for the specific language governing permissions and
|
|
14
|
-
# limitations under the License.
|
|
15
|
-
|
|
16
|
-
###########################################################################
|
|
17
|
-
# Example Tile Walker
|
|
18
|
-
#
|
|
19
|
-
# Trains a tetrahedral mesh quadruped to run. Feeds 8 time-varying input
|
|
20
|
-
# phases as inputs into a single layer fully connected network with a tanh
|
|
21
|
-
# activation function. Interprets the output of the network as tet
|
|
22
|
-
# activations, which are fed into the wp.sim soft mesh model. This is
|
|
23
|
-
# simulated forward in time and then evaluated based on the center of mass
|
|
24
|
-
# momentum of the mesh.
|
|
25
|
-
#
|
|
26
|
-
# This example uses the Warp tile API, which as of Warp 1.6 is the
|
|
27
|
-
# recommended way to handle matrix multiplication. example_walker.py in
|
|
28
|
-
# examples/optim demonstrates the old way of doing matrix multiplication,
|
|
29
|
-
# wp.matmul(), which will be deprecated in a future version.
|
|
30
|
-
#
|
|
31
|
-
###########################################################################
|
|
32
|
-
|
|
33
|
-
import math
|
|
34
|
-
import os
|
|
35
|
-
|
|
36
|
-
import numpy as np
|
|
37
|
-
from pxr import Gf, Usd, UsdGeom
|
|
38
|
-
|
|
39
|
-
import warp as wp
|
|
40
|
-
import warp.examples
|
|
41
|
-
import warp.optim
|
|
42
|
-
import warp.sim
|
|
43
|
-
import warp.sim.render
|
|
44
|
-
|
|
45
|
-
PHASE_COUNT = 8
|
|
46
|
-
PHASE_STEP = wp.constant((2.0 * math.pi) / PHASE_COUNT)
|
|
47
|
-
PHASE_FREQ = wp.constant(5.0)
|
|
48
|
-
ACTIVATION_STRENGTH = wp.constant(0.3)
|
|
49
|
-
|
|
50
|
-
TILE_TETS = wp.constant(8)
|
|
51
|
-
TILE_THREADS = 64
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
@wp.kernel
|
|
55
|
-
def loss_kernel(com: wp.array(dtype=wp.vec3), loss: wp.array(dtype=float)):
|
|
56
|
-
tid = wp.tid()
|
|
57
|
-
vx = com[tid][0]
|
|
58
|
-
vy = com[tid][1]
|
|
59
|
-
vz = com[tid][2]
|
|
60
|
-
delta = wp.sqrt(vx * vx) + wp.sqrt(vy * vy) - vz
|
|
61
|
-
|
|
62
|
-
wp.atomic_add(loss, 0, delta)
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
@wp.kernel
|
|
66
|
-
def com_kernel(velocities: wp.array(dtype=wp.vec3), n: int, com: wp.array(dtype=wp.vec3)):
|
|
67
|
-
tid = wp.tid()
|
|
68
|
-
v = velocities[tid]
|
|
69
|
-
a = v / wp.float32(n)
|
|
70
|
-
wp.atomic_add(com, 0, a)
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
@wp.kernel
|
|
74
|
-
def compute_phases(phases: wp.array(dtype=float), sim_time: float):
|
|
75
|
-
tid = wp.tid()
|
|
76
|
-
phases[tid] = wp.sin(PHASE_FREQ * sim_time + wp.float32(tid) * PHASE_STEP)
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
@wp.func
|
|
80
|
-
def tanh(x: float):
|
|
81
|
-
return wp.tanh(x) * ACTIVATION_STRENGTH
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
@wp.kernel
|
|
85
|
-
def network(
|
|
86
|
-
phases: wp.array2d(dtype=float), weights: wp.array2d(dtype=float), tet_activations: wp.array2d(dtype=float)
|
|
87
|
-
):
|
|
88
|
-
# output tile index
|
|
89
|
-
i = wp.tid()
|
|
90
|
-
|
|
91
|
-
# GEMM
|
|
92
|
-
p = wp.tile_load(phases, shape=(PHASE_COUNT, 1))
|
|
93
|
-
w = wp.tile_load(weights, shape=(TILE_TETS, PHASE_COUNT), offset=(i * TILE_TETS, 0))
|
|
94
|
-
out = wp.tile_matmul(w, p)
|
|
95
|
-
|
|
96
|
-
# activation
|
|
97
|
-
activations = wp.tile_map(tanh, out)
|
|
98
|
-
wp.tile_store(tet_activations, activations, offset=(i * TILE_TETS, 0))
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
class Example:
|
|
102
|
-
def __init__(self, stage_path="example_tile_walker.usd", verbose=False, num_frames=300):
|
|
103
|
-
self.verbose = verbose
|
|
104
|
-
|
|
105
|
-
fps = 60
|
|
106
|
-
self.frame_dt = 1.0 / fps
|
|
107
|
-
self.num_frames = num_frames
|
|
108
|
-
|
|
109
|
-
self.sim_substeps = 80
|
|
110
|
-
self.sim_dt = self.frame_dt / self.sim_substeps
|
|
111
|
-
self.sim_time = 0.0
|
|
112
|
-
|
|
113
|
-
self.iter = 0
|
|
114
|
-
self.train_rate = 0.025
|
|
115
|
-
|
|
116
|
-
self.phase_count = PHASE_COUNT
|
|
117
|
-
|
|
118
|
-
self.render_time = 0.0
|
|
119
|
-
|
|
120
|
-
# bear
|
|
121
|
-
asset_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bear.usd"))
|
|
122
|
-
|
|
123
|
-
geom = UsdGeom.Mesh(asset_stage.GetPrimAtPath("/root/bear"))
|
|
124
|
-
points = geom.GetPointsAttr().Get()
|
|
125
|
-
|
|
126
|
-
xform = Gf.Matrix4f(geom.ComputeLocalToWorldTransform(0.0))
|
|
127
|
-
for i in range(len(points)):
|
|
128
|
-
points[i] = xform.Transform(points[i])
|
|
129
|
-
|
|
130
|
-
self.points = [wp.vec3(point) for point in points]
|
|
131
|
-
self.tet_indices = geom.GetPrim().GetAttribute("tetraIndices").Get()
|
|
132
|
-
|
|
133
|
-
# sim model
|
|
134
|
-
builder = wp.sim.ModelBuilder()
|
|
135
|
-
builder.add_soft_mesh(
|
|
136
|
-
pos=wp.vec3(0.0, 0.5, 0.0),
|
|
137
|
-
rot=wp.quat_identity(),
|
|
138
|
-
scale=1.0,
|
|
139
|
-
vel=wp.vec3(0.0, 0.0, 0.0),
|
|
140
|
-
vertices=self.points,
|
|
141
|
-
indices=self.tet_indices,
|
|
142
|
-
density=1.0,
|
|
143
|
-
k_mu=2000.0,
|
|
144
|
-
k_lambda=2000.0,
|
|
145
|
-
k_damp=2.0,
|
|
146
|
-
tri_ke=0.0,
|
|
147
|
-
tri_ka=1e-8,
|
|
148
|
-
tri_kd=0.0,
|
|
149
|
-
tri_drag=0.0,
|
|
150
|
-
tri_lift=0.0,
|
|
151
|
-
)
|
|
152
|
-
|
|
153
|
-
# finalize model
|
|
154
|
-
self.model = builder.finalize(requires_grad=True)
|
|
155
|
-
self.control = self.model.control()
|
|
156
|
-
|
|
157
|
-
self.model.soft_contact_ke = 2.0e3
|
|
158
|
-
self.model.soft_contact_kd = 0.1
|
|
159
|
-
self.model.soft_contact_kf = 10.0
|
|
160
|
-
self.model.soft_contact_mu = 0.7
|
|
161
|
-
|
|
162
|
-
radii = wp.zeros(self.model.particle_count, dtype=float)
|
|
163
|
-
radii.fill_(0.05)
|
|
164
|
-
self.model.particle_radius = radii
|
|
165
|
-
self.model.ground = True
|
|
166
|
-
|
|
167
|
-
# allocate sim states
|
|
168
|
-
self.states = []
|
|
169
|
-
for _i in range(self.num_frames * self.sim_substeps + 1):
|
|
170
|
-
self.states.append(self.model.state(requires_grad=True))
|
|
171
|
-
|
|
172
|
-
# initialize the integrator.
|
|
173
|
-
self.integrator = wp.sim.SemiImplicitIntegrator()
|
|
174
|
-
|
|
175
|
-
# model input
|
|
176
|
-
self.phases = []
|
|
177
|
-
for _i in range(self.num_frames):
|
|
178
|
-
self.phases.append(wp.zeros(self.phase_count, dtype=float, requires_grad=True))
|
|
179
|
-
|
|
180
|
-
# weights matrix for linear network
|
|
181
|
-
rng = np.random.default_rng(42)
|
|
182
|
-
k = 1.0 / self.phase_count
|
|
183
|
-
weights = rng.uniform(-np.sqrt(k), np.sqrt(k), (self.model.tet_count, self.phase_count))
|
|
184
|
-
self.weights = wp.array(weights, dtype=float, requires_grad=True)
|
|
185
|
-
|
|
186
|
-
# tanh activation layer array
|
|
187
|
-
self.tet_activations = []
|
|
188
|
-
for _i in range(self.num_frames):
|
|
189
|
-
self.tet_activations.append(wp.zeros(self.model.tet_count, dtype=float, requires_grad=True))
|
|
190
|
-
|
|
191
|
-
# optimization
|
|
192
|
-
self.loss = wp.zeros(1, dtype=float, requires_grad=True)
|
|
193
|
-
self.coms = []
|
|
194
|
-
for _i in range(self.num_frames):
|
|
195
|
-
self.coms.append(wp.zeros(1, dtype=wp.vec3, requires_grad=True))
|
|
196
|
-
self.optimizer = warp.optim.Adam([self.weights.flatten()], lr=self.train_rate)
|
|
197
|
-
|
|
198
|
-
# rendering
|
|
199
|
-
if stage_path:
|
|
200
|
-
self.renderer = wp.sim.render.SimRenderer(self.model, stage_path)
|
|
201
|
-
else:
|
|
202
|
-
self.renderer = None
|
|
203
|
-
|
|
204
|
-
# capture forward/backward passes
|
|
205
|
-
self.use_cuda_graph = wp.get_device().is_cuda
|
|
206
|
-
if self.use_cuda_graph:
|
|
207
|
-
with wp.ScopedCapture() as capture:
|
|
208
|
-
self.tape = wp.Tape()
|
|
209
|
-
with self.tape:
|
|
210
|
-
for i in range(self.num_frames):
|
|
211
|
-
self.forward(i)
|
|
212
|
-
self.tape.backward(self.loss)
|
|
213
|
-
self.graph = capture.graph
|
|
214
|
-
|
|
215
|
-
def forward(self, frame):
|
|
216
|
-
with wp.ScopedTimer("network", active=self.verbose):
|
|
217
|
-
# build sinusoidal input phases
|
|
218
|
-
wp.launch(kernel=compute_phases, dim=self.phase_count, inputs=[self.phases[frame], self.sim_time])
|
|
219
|
-
|
|
220
|
-
# apply linear network with tanh activation
|
|
221
|
-
wp.launch_tiled(
|
|
222
|
-
kernel=network,
|
|
223
|
-
dim=math.ceil(self.model.tet_count / TILE_TETS),
|
|
224
|
-
inputs=[self.phases[frame].reshape((self.phase_count, 1)), self.weights],
|
|
225
|
-
outputs=[self.tet_activations[frame].reshape((self.model.tet_count, 1))],
|
|
226
|
-
block_dim=TILE_THREADS,
|
|
227
|
-
)
|
|
228
|
-
self.control.tet_activations = self.tet_activations[frame]
|
|
229
|
-
|
|
230
|
-
with wp.ScopedTimer("simulate", active=self.verbose):
|
|
231
|
-
# run simulation loop
|
|
232
|
-
for i in range(self.sim_substeps):
|
|
233
|
-
self.states[frame * self.sim_substeps + i].clear_forces()
|
|
234
|
-
self.integrator.simulate(
|
|
235
|
-
self.model,
|
|
236
|
-
self.states[frame * self.sim_substeps + i],
|
|
237
|
-
self.states[frame * self.sim_substeps + i + 1],
|
|
238
|
-
self.sim_dt,
|
|
239
|
-
self.control,
|
|
240
|
-
)
|
|
241
|
-
self.sim_time += self.sim_dt
|
|
242
|
-
|
|
243
|
-
with wp.ScopedTimer("loss", active=self.verbose):
|
|
244
|
-
# compute center of mass velocity
|
|
245
|
-
wp.launch(
|
|
246
|
-
com_kernel,
|
|
247
|
-
dim=self.model.particle_count,
|
|
248
|
-
inputs=[
|
|
249
|
-
self.states[(frame + 1) * self.sim_substeps].particle_qd,
|
|
250
|
-
self.model.particle_count,
|
|
251
|
-
self.coms[frame],
|
|
252
|
-
],
|
|
253
|
-
outputs=[],
|
|
254
|
-
)
|
|
255
|
-
# compute loss
|
|
256
|
-
wp.launch(loss_kernel, dim=1, inputs=[self.coms[frame], self.loss], outputs=[])
|
|
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
|
-
for i in range(self.num_frames):
|
|
266
|
-
self.forward(i)
|
|
267
|
-
self.tape.backward(self.loss)
|
|
268
|
-
|
|
269
|
-
# optimization
|
|
270
|
-
x = self.weights.grad.flatten()
|
|
271
|
-
self.optimizer.step([x])
|
|
272
|
-
|
|
273
|
-
loss = self.loss.numpy()
|
|
274
|
-
if self.verbose:
|
|
275
|
-
print(f"Iteration {self.iter}: {loss}")
|
|
276
|
-
|
|
277
|
-
# reset sim
|
|
278
|
-
self.sim_time = 0.0
|
|
279
|
-
self.states[0] = self.model.state(requires_grad=True)
|
|
280
|
-
|
|
281
|
-
# clear grads and zero arrays for next iteration
|
|
282
|
-
self.tape.zero()
|
|
283
|
-
self.loss.zero_()
|
|
284
|
-
for i in range(self.num_frames):
|
|
285
|
-
self.coms[i].zero_()
|
|
286
|
-
|
|
287
|
-
self.iter += 1
|
|
288
|
-
|
|
289
|
-
def render(self):
|
|
290
|
-
if self.renderer is None:
|
|
291
|
-
return
|
|
292
|
-
|
|
293
|
-
with wp.ScopedTimer("render"):
|
|
294
|
-
for i in range(self.num_frames + 1):
|
|
295
|
-
self.renderer.begin_frame(self.render_time)
|
|
296
|
-
self.renderer.render(self.states[i * self.sim_substeps])
|
|
297
|
-
self.renderer.end_frame()
|
|
298
|
-
|
|
299
|
-
self.render_time += self.frame_dt
|
|
300
|
-
|
|
301
|
-
|
|
302
|
-
if __name__ == "__main__":
|
|
303
|
-
import argparse
|
|
304
|
-
|
|
305
|
-
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
306
|
-
parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
|
|
307
|
-
parser.add_argument(
|
|
308
|
-
"--stage_path",
|
|
309
|
-
type=lambda x: None if x == "None" else str(x),
|
|
310
|
-
default="example_tile_walker.usd",
|
|
311
|
-
help="Path to the output USD file.",
|
|
312
|
-
)
|
|
313
|
-
parser.add_argument("--num_frames", type=int, default=300, help="Total number of frames per training iteration.")
|
|
314
|
-
parser.add_argument("--train_iters", type=int, default=30, help="Total number of training iterations.")
|
|
315
|
-
parser.add_argument("--verbose", action="store_true", help="Print out additional status messages during execution.")
|
|
316
|
-
|
|
317
|
-
args = parser.parse_known_args()[0]
|
|
318
|
-
|
|
319
|
-
with wp.ScopedDevice(args.device):
|
|
320
|
-
example = Example(stage_path=args.stage_path, verbose=args.verbose, num_frames=args.num_frames)
|
|
321
|
-
|
|
322
|
-
for _ in range(args.train_iters):
|
|
323
|
-
example.step()
|
|
324
|
-
example.render()
|
|
325
|
-
|
|
326
|
-
if example.renderer:
|
|
327
|
-
example.renderer.save()
|