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.
- warp/__init__.py +8 -0
- warp/bin/warp-clang.so +0 -0
- warp/bin/warp.so +0 -0
- warp/build.py +7 -6
- warp/build_dll.py +70 -79
- warp/builtins.py +10 -6
- warp/codegen.py +51 -19
- warp/config.py +7 -8
- warp/constants.py +3 -0
- warp/context.py +948 -245
- warp/dlpack.py +198 -113
- warp/examples/assets/bunny.usd +0 -0
- warp/examples/assets/cartpole.urdf +110 -0
- warp/examples/assets/crazyflie.usd +0 -0
- warp/examples/assets/cube.usda +42 -0
- warp/examples/assets/nv_ant.xml +92 -0
- warp/examples/assets/nv_humanoid.xml +183 -0
- warp/examples/assets/quadruped.urdf +268 -0
- warp/examples/assets/rocks.nvdb +0 -0
- warp/examples/assets/rocks.usd +0 -0
- warp/examples/assets/sphere.usda +56 -0
- warp/examples/assets/torus.usda +105 -0
- warp/examples/benchmarks/benchmark_api.py +383 -0
- warp/examples/benchmarks/benchmark_cloth.py +279 -0
- warp/examples/benchmarks/benchmark_cloth_cupy.py +88 -0
- warp/examples/benchmarks/benchmark_cloth_jax.py +100 -0
- warp/examples/benchmarks/benchmark_cloth_numba.py +142 -0
- warp/examples/benchmarks/benchmark_cloth_numpy.py +77 -0
- warp/examples/benchmarks/benchmark_cloth_pytorch.py +86 -0
- warp/examples/benchmarks/benchmark_cloth_taichi.py +112 -0
- warp/examples/benchmarks/benchmark_cloth_warp.py +146 -0
- warp/examples/benchmarks/benchmark_launches.py +295 -0
- warp/examples/core/example_dem.py +221 -0
- warp/examples/core/example_fluid.py +267 -0
- warp/examples/core/example_graph_capture.py +129 -0
- warp/examples/core/example_marching_cubes.py +177 -0
- warp/examples/core/example_mesh.py +154 -0
- warp/examples/core/example_mesh_intersect.py +193 -0
- warp/examples/core/example_nvdb.py +169 -0
- warp/examples/core/example_raycast.py +89 -0
- warp/examples/core/example_raymarch.py +178 -0
- warp/examples/core/example_render_opengl.py +141 -0
- warp/examples/core/example_sph.py +389 -0
- warp/examples/core/example_torch.py +181 -0
- warp/examples/core/example_wave.py +249 -0
- warp/examples/fem/bsr_utils.py +380 -0
- warp/examples/fem/example_apic_fluid.py +391 -0
- warp/examples/fem/example_convection_diffusion.py +168 -0
- warp/examples/fem/example_convection_diffusion_dg.py +209 -0
- warp/examples/fem/example_convection_diffusion_dg0.py +194 -0
- warp/examples/fem/example_deformed_geometry.py +159 -0
- warp/examples/fem/example_diffusion.py +173 -0
- warp/examples/fem/example_diffusion_3d.py +152 -0
- warp/examples/fem/example_diffusion_mgpu.py +214 -0
- warp/examples/fem/example_mixed_elasticity.py +222 -0
- warp/examples/fem/example_navier_stokes.py +243 -0
- warp/examples/fem/example_stokes.py +192 -0
- warp/examples/fem/example_stokes_transfer.py +249 -0
- warp/examples/fem/mesh_utils.py +109 -0
- warp/examples/fem/plot_utils.py +287 -0
- warp/examples/optim/example_bounce.py +248 -0
- warp/examples/optim/example_cloth_throw.py +210 -0
- warp/examples/optim/example_diffray.py +535 -0
- warp/examples/optim/example_drone.py +850 -0
- warp/examples/optim/example_inverse_kinematics.py +169 -0
- warp/examples/optim/example_inverse_kinematics_torch.py +170 -0
- warp/examples/optim/example_spring_cage.py +234 -0
- warp/examples/optim/example_trajectory.py +201 -0
- warp/examples/sim/example_cartpole.py +128 -0
- warp/examples/sim/example_cloth.py +184 -0
- warp/examples/sim/example_granular.py +113 -0
- warp/examples/sim/example_granular_collision_sdf.py +185 -0
- warp/examples/sim/example_jacobian_ik.py +213 -0
- warp/examples/sim/example_particle_chain.py +106 -0
- warp/examples/sim/example_quadruped.py +179 -0
- warp/examples/sim/example_rigid_chain.py +191 -0
- warp/examples/sim/example_rigid_contact.py +176 -0
- warp/examples/sim/example_rigid_force.py +126 -0
- warp/examples/sim/example_rigid_gyroscopic.py +97 -0
- warp/examples/sim/example_rigid_soft_contact.py +124 -0
- warp/examples/sim/example_soft_body.py +178 -0
- warp/fabric.py +29 -20
- warp/fem/cache.py +0 -1
- warp/fem/dirichlet.py +0 -2
- warp/fem/integrate.py +0 -1
- warp/jax.py +45 -0
- warp/jax_experimental.py +339 -0
- warp/native/builtin.h +12 -0
- warp/native/bvh.cu +18 -18
- warp/native/clang/clang.cpp +8 -3
- warp/native/cuda_util.cpp +94 -5
- warp/native/cuda_util.h +35 -6
- warp/native/cutlass_gemm.cpp +1 -1
- warp/native/cutlass_gemm.cu +4 -1
- warp/native/error.cpp +66 -0
- warp/native/error.h +27 -0
- warp/native/mesh.cu +2 -2
- warp/native/reduce.cu +4 -4
- warp/native/runlength_encode.cu +2 -2
- warp/native/scan.cu +2 -2
- warp/native/sparse.cu +0 -1
- warp/native/temp_buffer.h +2 -2
- warp/native/warp.cpp +95 -60
- warp/native/warp.cu +1053 -218
- warp/native/warp.h +49 -32
- warp/optim/linear.py +33 -16
- warp/render/render_opengl.py +202 -101
- warp/render/render_usd.py +82 -40
- warp/sim/__init__.py +13 -4
- warp/sim/articulation.py +4 -5
- warp/sim/collide.py +320 -175
- warp/sim/import_mjcf.py +25 -30
- warp/sim/import_urdf.py +94 -63
- warp/sim/import_usd.py +51 -36
- warp/sim/inertia.py +3 -2
- warp/sim/integrator.py +233 -0
- warp/sim/integrator_euler.py +447 -469
- warp/sim/integrator_featherstone.py +1991 -0
- warp/sim/integrator_xpbd.py +1420 -640
- warp/sim/model.py +765 -487
- warp/sim/particles.py +2 -1
- warp/sim/render.py +35 -13
- warp/sim/utils.py +222 -11
- warp/stubs.py +8 -0
- warp/tape.py +16 -1
- warp/tests/aux_test_grad_customs.py +23 -0
- warp/tests/test_array.py +190 -1
- warp/tests/test_async.py +656 -0
- warp/tests/test_bool.py +50 -0
- warp/tests/test_dlpack.py +164 -11
- warp/tests/test_examples.py +166 -74
- warp/tests/test_fem.py +8 -1
- warp/tests/test_generics.py +15 -5
- warp/tests/test_grad.py +1 -1
- warp/tests/test_grad_customs.py +172 -12
- warp/tests/test_jax.py +254 -0
- warp/tests/test_large.py +29 -6
- warp/tests/test_launch.py +25 -0
- warp/tests/test_linear_solvers.py +20 -3
- warp/tests/test_matmul.py +61 -16
- warp/tests/test_matmul_lite.py +13 -13
- warp/tests/test_mempool.py +186 -0
- warp/tests/test_multigpu.py +3 -0
- warp/tests/test_options.py +16 -2
- warp/tests/test_peer.py +137 -0
- warp/tests/test_print.py +3 -1
- warp/tests/test_quat.py +23 -0
- warp/tests/test_sim_kinematics.py +97 -0
- warp/tests/test_snippet.py +126 -3
- warp/tests/test_streams.py +108 -79
- warp/tests/test_torch.py +16 -8
- warp/tests/test_utils.py +32 -27
- warp/tests/test_verify_fp.py +65 -0
- warp/tests/test_volume.py +1 -1
- warp/tests/unittest_serial.py +2 -0
- warp/tests/unittest_suites.py +12 -0
- warp/tests/unittest_utils.py +14 -7
- warp/thirdparty/unittest_parallel.py +15 -3
- warp/torch.py +10 -8
- warp/types.py +363 -246
- warp/utils.py +143 -19
- warp_lang-1.0.0.dist-info/LICENSE.md +126 -0
- warp_lang-1.0.0.dist-info/METADATA +394 -0
- {warp_lang-0.11.0.dist-info → warp_lang-1.0.0.dist-info}/RECORD +167 -86
- warp/sim/optimizer.py +0 -138
- warp_lang-0.11.0.dist-info/LICENSE.md +0 -36
- warp_lang-0.11.0.dist-info/METADATA +0 -238
- /warp/tests/{walkthough_debug.py → walkthrough_debug.py} +0 -0
- {warp_lang-0.11.0.dist-info → warp_lang-1.0.0.dist-info}/WHEEL +0 -0
- {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()
|