warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

examples/{example_sim_sdf_shape.py → example_sim_granular_collision_sdf.py}

@@ -6,20 +6,22 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 ###########################################################################
-# Example Sim Granular
+# Example Sim Granular Collision SDF
 #
 # Shows how to set up a particle-based granular material model using the
-# wp.sim.ModelBuilder().
+# wp.sim.ModelBuilder(). This version shows how to create collision geometry
+# objects from SDFs.
 #
 ###########################################################################
 
+import math
 import os
 
+import numpy as np
+
 import warp as wp
 import warp.sim
 import warp.sim.render
-import math
-import numpy as np
 
 wp.init()
 
@@ -46,25 +48,29 @@ class Example:
             cell_x=self.radius * 2.0,
             cell_y=self.radius * 2.0,
             cell_z=self.radius * 2.0,
-            pos=(0.0, 20.0, 0.0),
+            pos=wp.vec3(0.0, 20.0, 0.0),
             rot=wp.quat_identity(),
-            vel=(2.0, 0.0, 0.0),
+            vel=wp.vec3(2.0, 0.0, 0.0),
             mass=0.1,
             jitter=self.radius * 0.1,
         )
-        rock_vdb = wp.Volume.load_from_nvdb(
-            open(os.path.join(os.path.dirname(__file__), "assets/rocks.nvdb"), "rb").read())
+        rock_file = open(os.path.join(os.path.dirname(__file__), "assets/rocks.nvdb"), "rb")
+        rock_vdb = wp.Volume.load_from_nvdb(rock_file.read())
+        rock_file.close()
+
         rock_sdf = wp.sim.SDF(rock_vdb)
-        s = builder.add_shape_sdf(
+
+        builder.add_shape_sdf(
             ke=1.0e4,
             kd=1000.0,
             kf=1000.0,
             mu=0.5,
             sdf=rock_sdf,
             body=-1,
-            pos=(0.0, 0.0, 0.0),
-            rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), -0.5 * math.pi),
-            scale=(0.01, 0.01, 0.01))
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -0.5 * math.pi),
+            scale=wp.vec3(0.01, 0.01, 0.01),
+        )
 
         mins = np.array([-3.0, -3.0, -3.0])
         voxel_size = 0.2
@@ -83,10 +89,10 @@ class Example:
         sphere_vdb = wp.Volume.load_from_numpy(sphere_sdf_np, mins, voxel_size, rad + 3.0 * voxel_size)
         sphere_sdf = wp.sim.SDF(sphere_vdb)
 
-        self.sphere_pos = (3.0, 15.0, 0.0)
+        self.sphere_pos = wp.vec3(3.0, 15.0, 0.0)
         self.sphere_scale = 1.0
         self.sphere_radius = rad
-        s = builder.add_shape_sdf(
+        builder.add_shape_sdf(
             ke=1.0e4,
             kd=1000.0,
             kf=1000.0,
@@ -94,7 +100,8 @@ class Example:
             sdf=sphere_sdf,
             body=-1,
             pos=self.sphere_pos,
-            scale=(self.sphere_scale, self.sphere_scale, self.sphere_scale))
+            scale=wp.vec3(self.sphere_scale, self.sphere_scale, self.sphere_scale),
+        )
 
         self.model = builder.finalize()
         self.model.particle_kf = 25.0
@@ -113,7 +120,7 @@ class Example:
         with wp.ScopedTimer("simulate", active=True):
             self.model.particle_grid.build(self.state_0.particle_q, self.radius * 2.0)
 
-            for s in range(self.sim_substeps):
+            for _ in range(self.sim_substeps):
                 self.state_0.clear_forces()
                 wp.sim.collide(self.model, self.state_0)
                 self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
@@ -121,37 +128,41 @@ class Example:
                 # swap states
                 (self.state_0, self.state_1) = (self.state_1, self.state_0)
 
+            self.sim_time += self.frame_dt
+
     def render(self, is_live=False):
         with wp.ScopedTimer("render", active=True):
             time = 0.0 if is_live else self.sim_time
 
             self.renderer.begin_frame(time)
 
-            # Note the extra wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi) is because .usd is oriented differently from .nvdb
+            # 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
             self.renderer.render_ref(
                 name="collision",
                 path=os.path.join(os.path.dirname(__file__), "assets/rocks.usd"),
-                pos=(0.0, 0.0, 0.0),
-                rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), -0.5 * math.pi) *
-                wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi),
-                scale=(0.01, 0.01, 0.01),
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -0.5 * math.pi)
+                * wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi),
+                scale=wp.vec3(0.01, 0.01, 0.01),
             )
 
-            self.renderer.render_sphere(name="sphere", pos=self.sphere_pos,
-                                        radius=self.sphere_scale * self.sphere_radius, rot=(0.0, 0.0, 0.0, 1.0))
+            self.renderer.render_sphere(
+                name="sphere",
+                pos=self.sphere_pos,
+                radius=self.sphere_scale * self.sphere_radius,
+                rot=wp.quat(0.0, 0.0, 0.0, 1.0),
+            )
 
             self.renderer.render(self.state_0)
             self.renderer.end_frame()
 
-        self.sim_time += self.frame_dt
-
 
 if __name__ == "__main__":
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_sdf_shape.usd")
 
     example = Example(stage_path)
 
-    for i in range(example.frame_count):
+    for _ in range(example.frame_count):
         example.update()
         example.render()
 

examples/example_sim_neo_hookean.py

@@ -13,8 +13,8 @@
 #
 ###########################################################################
 
-import os
 import math
+import os
 
 import warp as wp
 import warp.sim
@@ -23,22 +23,54 @@ import warp.sim.render
 wp.init()
 
 
+@wp.kernel
+def twist_points(
+    rest: wp.array(dtype=wp.vec3), points: wp.array(dtype=wp.vec3), mass: wp.array(dtype=float), xform: wp.transform
+):
+    tid = wp.tid()
+
+    r = rest[tid]
+    p = points[tid]
+    m = mass[tid]
+
+    # twist the top layer of particles in the beam
+    if m == 0 and p[1] != 0.0:
+        points[tid] = wp.transform_point(xform, r)
+
+
+@wp.kernel
+def compute_volume(points: wp.array(dtype=wp.vec3), indices: wp.array2d(dtype=int), volume: wp.array(dtype=float)):
+    tid = wp.tid()
+
+    i = indices[tid, 0]
+    j = indices[tid, 1]
+    k = indices[tid, 2]
+    l = indices[tid, 3]
+
+    x0 = points[i]
+    x1 = points[j]
+    x2 = points[k]
+    x3 = points[l]
+
+    x10 = x1 - x0
+    x20 = x2 - x0
+    x30 = x3 - x0
+
+    v = wp.dot(x10, wp.cross(x20, x30)) / 6.0
+
+    wp.atomic_add(volume, 0, v)
+
+
 class Example:
     def __init__(self, stage):
-        self.sim_width = 8
-        self.sim_height = 8
-
-        self.sim_fps = 60.0
+        sim_fps = 60.0
         self.sim_substeps = 64
-        self.sim_duration = 5.0
-        self.sim_frames = int(self.sim_duration * self.sim_fps)
-        self.sim_dt = (1.0 / self.sim_fps) / self.sim_substeps
+        sim_duration = 5.0
+        self.sim_frames = int(sim_duration * sim_fps)
+        self.sim_dt = (1.0 / sim_fps) / self.sim_substeps
         self.sim_time = 0.0
-        self.sim_render = True
-        self.sim_iterations = 1
-        self.sim_relaxation = 1.0
-        self.lift_speed = 2.5 / self.sim_duration * 2.0  # from Smith et al.
-        self.rot_speed = math.pi / self.sim_duration
+        self.lift_speed = 2.5 / sim_duration * 2.0  # from Smith et al.
+        self.rot_speed = math.pi / sim_duration
 
         builder = wp.sim.ModelBuilder()
 
@@ -48,9 +80,9 @@ class Example:
         center = cell_size * cell_dim * 0.5
 
         builder.add_soft_grid(
-            pos=(-center, 0.0, -center),
+            pos=wp.vec3(-center, 0.0, -center),
             rot=wp.quat_identity(),
-            vel=(0.0, 0.0, 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
             dim_x=cell_dim,
             dim_y=cell_dim,
             dim_z=cell_dim,
@@ -85,15 +117,15 @@ class Example:
         with wp.ScopedTimer("simulate"):
             xform = wp.transform(
                 (0.0, self.lift_speed * self.sim_time, 0.0),
-                wp.quat_from_axis_angle((0.0, 1.0, 0.0), self.rot_speed * self.sim_time),
+                wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), self.rot_speed * self.sim_time),
             )
             wp.launch(
-                kernel=self.twist_points,
+                kernel=twist_points,
                 dim=len(self.state_0.particle_q),
                 inputs=[self.rest.particle_q, self.state_0.particle_q, self.model.particle_mass, xform],
             )
 
-            for s in range(self.sim_substeps):
+            for _ in range(self.sim_substeps):
                 self.state_0.clear_forces()
                 self.state_1.clear_forces()
 
@@ -105,7 +137,7 @@ class Example:
 
             self.volume.zero_()
             wp.launch(
-                kernel=self.compute_volume,
+                kernel=compute_volume,
                 dim=self.model.tet_count,
                 inputs=[self.state_0.particle_q, self.model.tet_indices, self.volume],
             )
@@ -118,42 +150,6 @@ class Example:
             self.renderer.render(self.state_0)
             self.renderer.end_frame()
 
-    @wp.kernel
-    def twist_points(
-        rest: wp.array(dtype=wp.vec3), points: wp.array(dtype=wp.vec3), mass: wp.array(dtype=float), xform: wp.transform
-    ):
-        tid = wp.tid()
-
-        r = rest[tid]
-        p = points[tid]
-        m = mass[tid]
-
-        # twist the top layer of particles in the beam
-        if m == 0 and p[1] != 0.0:
-            points[tid] = wp.transform_point(xform, r)
-
-    @wp.kernel
-    def compute_volume(points: wp.array(dtype=wp.vec3), indices: wp.array2d(dtype=int), volume: wp.array(dtype=float)):
-        tid = wp.tid()
-
-        i = indices[tid, 0]
-        j = indices[tid, 1]
-        k = indices[tid, 2]
-        l = indices[tid, 3]
-
-        x0 = points[i]
-        x1 = points[j]
-        x2 = points[k]
-        x3 = points[l]
-
-        x10 = x1 - x0
-        x20 = x2 - x0
-        x30 = x3 - x0
-
-        v = wp.dot(x10, wp.cross(x20, x30)) / 6.0
-
-        wp.atomic_add(volume, 0, v)
-
 
 if __name__ == "__main__":
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_neo_hookean.usd")

examples/example_sim_particle_chain.py

@@ -13,8 +13,8 @@
 #
 ###########################################################################
 
-import os
 import math
+import os
 
 import warp as wp
 import warp.sim
@@ -38,13 +38,13 @@ class Example:
         builder = wp.sim.ModelBuilder()
 
         # anchor
-        builder.add_particle((0.0, 1.0, 0.0), (0.0, 0.0, 0.0), 0.0)
+        builder.add_particle(wp.vec3(0.0, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), 0.0)
 
         # chain
         for i in range(1, 10):
             radius = math.sqrt(i) * 0.2
             mass = math.pi * radius * radius * radius
-            builder.add_particle((i, 1.0, 0.0), (0.0, 0.0, 0.0), mass, radius=radius)
+            builder.add_particle(wp.vec3(i, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), mass, radius=radius)
             builder.add_spring(i - 1, i, 1.0e6, 0.0, 0)
 
         self.model = builder.finalize()
@@ -59,7 +59,7 @@ class Example:
 
     def update(self):
         with wp.ScopedTimer("simulate"):
-            for s in range(self.sim_substeps):
+            for _ in range(self.sim_substeps):
                 self.state_0.clear_forces()
                 self.state_1.clear_forces()
 

examples/example_sim_quadruped.py

@@ -15,39 +15,68 @@
 ###########################################################################
 
 import math
-import numpy as np
 import os
 
+import numpy as np
+
 import warp as wp
 import warp.sim
 import warp.sim.render
-from env.environment import compute_env_offsets
 
 wp.init()
 
 
-class Example:
-    frame_dt = 1.0 / 100.0
-
-    episode_duration = 5.0  # seconds
-    episode_frames = int(episode_duration / frame_dt)
-
-    sim_substeps = 5
-    sim_dt = frame_dt / sim_substeps
-    sim_steps = int(episode_duration / sim_dt)
-
-    sim_time = 0.0
-    render_time = 0.0
+# Taken from env/environment.py
+def compute_env_offsets(num_envs, env_offset=(5.0, 0.0, 5.0), up_axis="Y"):
+    # compute positional offsets per environment
+    env_offset = np.array(env_offset)
+    nonzeros = np.nonzero(env_offset)[0]
+    num_dim = nonzeros.shape[0]
+    if num_dim > 0:
+        side_length = int(np.ceil(num_envs ** (1.0 / num_dim)))
+        env_offsets = []
+    else:
+        env_offsets = np.zeros((num_envs, 3))
+    if num_dim == 1:
+        for i in range(num_envs):
+            env_offsets.append(i * env_offset)
+    elif num_dim == 2:
+        for i in range(num_envs):
+            d0 = i // side_length
+            d1 = i % side_length
+            offset = np.zeros(3)
+            offset[nonzeros[0]] = d0 * env_offset[nonzeros[0]]
+            offset[nonzeros[1]] = d1 * env_offset[nonzeros[1]]
+            env_offsets.append(offset)
+    elif num_dim == 3:
+        for i in range(num_envs):
+            d0 = i // (side_length * side_length)
+            d1 = (i // side_length) % side_length
+            d2 = i % side_length
+            offset = np.zeros(3)
+            offset[0] = d0 * env_offset[0]
+            offset[1] = d1 * env_offset[1]
+            offset[2] = d2 * env_offset[2]
+            env_offsets.append(offset)
+    env_offsets = np.array(env_offsets)
+    min_offsets = np.min(env_offsets, axis=0)
+    correction = min_offsets + (np.max(env_offsets, axis=0) - min_offsets) / 2.0
+    if isinstance(up_axis, str):
+        up_axis = "XYZ".index(up_axis.upper())
+    correction[up_axis] = 0.0  # ensure the envs are not shifted below the ground plane
+    env_offsets -= correction
+    return env_offsets
 
-    def __init__(self, stage=None, render=True, num_envs=1):
-        self.enable_rendering = render
 
+class Example:
+    def __init__(self, stage=None, num_envs=1, enable_rendering=True, print_timers=True):
+        self.device = wp.get_device()
         self.num_envs = num_envs
         articulation_builder = wp.sim.ModelBuilder()
         wp.sim.parse_urdf(
             os.path.join(os.path.dirname(__file__), "assets/quadruped.urdf"),
             articulation_builder,
-            xform=wp.transform([0.0, 0.7, 0.0], wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)),
+            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)),
             floating=True,
             density=1000,
             armature=0.01,
@@ -62,6 +91,16 @@ class Example:
         )
 
         builder = wp.sim.ModelBuilder()
+
+        self.sim_time = 0.0
+        self.frame_dt = 1.0 / 100.0
+
+        episode_duration = 5.0  # seconds
+        self.episode_frames = int(episode_duration / self.frame_dt)
+
+        self.sim_substeps = 5
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
         offsets = compute_env_offsets(num_envs)
         for i in range(num_envs):
             builder.add_builder(articulation_builder, xform=wp.transform(offsets[i], wp.quat_identity()))
@@ -80,60 +119,65 @@ class Example:
 
         self.integrator = wp.sim.XPBDIntegrator()
 
-        # -----------------------
-        # set up Usd renderer
+        self.enable_rendering = enable_rendering
         self.renderer = None
-        if render:
+        if self.enable_rendering:
             self.renderer = wp.sim.render.SimRenderer(self.model, stage)
 
-    def update(self):
-        for _ in range(self.sim_substeps):
-            self.state_0.clear_forces()
-            wp.sim.collide(self.model, self.state_0)
-            self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
-            self.state_0, self.state_1 = self.state_1, self.state_0
-
-    def render(self, is_live=False):
-        time = 0.0 if is_live else self.sim_time
-
-        self.renderer.begin_frame(time)
-        self.renderer.render(self.state_0)
-        self.renderer.end_frame()
+        self.print_timers = print_timers
 
-    def run(self, render=True):
-        # ---------------
-        # run simulation
-
-        self.sim_time = 0.0
         self.state_0 = self.model.state()
         self.state_1 = self.model.state()
 
         wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state_0)
 
-        profiler = {}
+        self.use_graph = wp.get_device().is_cuda
+        self.graph = None
+
+        if self.use_graph:
+            # create update graph
+            wp.capture_begin(self.device)
+            try:
+                self.update()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
-        # create update graph
-        wp.capture_begin()
+    def update(self):
+        with wp.ScopedTimer("simulate", active=True, print=self.print_timers):
+            if not self.use_graph or self.graph is None:
+                for _ in range(self.sim_substeps):
+                    self.state_0.clear_forces()
+                    wp.sim.collide(self.model, self.state_0)
+                    self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+                    self.state_0, self.state_1 = self.state_1, self.state_0
+            else:
+                wp.capture_launch(self.graph)
+
+            if not wp.get_device().is_capturing:
+                self.sim_time += self.frame_dt
 
-        # simulate
-        self.update()
+    def render(self, is_live=False):
+        if self.enable_rendering:
+            with wp.ScopedTimer("render", active=True, print=self.print_timers):
+                time = 0.0 if is_live else self.sim_time
 
-        graph = wp.capture_end()
+                self.renderer.begin_frame(time)
+                self.renderer.render(self.state_0)
+                self.renderer.end_frame()
 
-        # simulate
-        with wp.ScopedTimer("simulate", detailed=False, print=False, active=True, dict=profiler):
-            for f in range(0, self.episode_frames):
-                with wp.ScopedTimer("simulate", active=True):
-                    wp.capture_launch(graph)
-                self.sim_time += self.frame_dt
+    def run(self):
+        profiler = {}
 
-                if self.enable_rendering:
-                    with wp.ScopedTimer("render", active=True):
-                        self.render()
-                    self.renderer.save()
+        with wp.ScopedTimer("simulate", detailed=False, print=False, active=True, dict=profiler):
+            for _ in range(self.episode_frames):
+                self.update()
+                self.render()
 
             wp.synchronize()
 
+        if self.enable_rendering:
+            self.renderer.save()
+
         avg_time = np.array(profiler["simulate"]).mean() / self.episode_frames
         avg_steps_second = 1000.0 * float(self.num_envs) / avg_time
 
@@ -142,38 +186,39 @@ class Example:
         return 1000.0 * float(self.num_envs) / avg_time
 
 
-profile = False
+if __name__ == "__main__":
+    profile = False
 
-if profile:
-    env_count = 2
-    env_times = []
-    env_size = []
+    if profile:
+        env_count = 2
+        env_times = []
+        env_size = []
 
-    for i in range(15):
-        robot = Example(render=False, num_envs=env_count)
-        steps_per_second = robot.run()
+        for i in range(15):
+            example = Example(num_envs=env_count, enable_rendering=False, print_timers=False)
+            steps_per_second = example.run()
 
-        env_size.append(env_count)
-        env_times.append(steps_per_second)
+            env_size.append(env_count)
+            env_times.append(steps_per_second)
 
-        env_count *= 2
+            env_count *= 2
 
-    # dump times
-    for i in range(len(env_times)):
-        print(f"envs: {env_size[i]} steps/second: {env_times[i]}")
+        # dump times
+        for i in range(len(env_times)):
+            print(f"envs: {env_size[i]} steps/second: {env_times[i]}")
 
-    # plot
-    import matplotlib.pyplot as plt
+        # plot
+        import matplotlib.pyplot as plt
 
-    plt.figure(1)
-    plt.plot(env_size, env_times)
-    plt.xscale("log")
-    plt.xlabel("Number of Envs")
-    plt.yscale("log")
-    plt.ylabel("Steps/Second")
-    plt.show()
+        plt.figure(1)
+        plt.plot(env_size, env_times)
+        plt.xscale("log")
+        plt.xlabel("Number of Envs")
+        plt.yscale("log")
+        plt.ylabel("Steps/Second")
+        plt.show()
 
-else:
-    stage = os.path.join(os.path.dirname(__file__), "outputs/example_sim_quadruped.usd")
-    robot = Example(stage, render=True, num_envs=25)
-    robot.run()
+    else:
+        stage = os.path.join(os.path.dirname(__file__), "outputs/example_sim_quadruped.usd")
+        example = Example(stage, num_envs=25, enable_rendering=True)
+        example.run()