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

examples/example_nvdb.py

@@ -15,10 +15,11 @@
 #
 ###########################################################################
 
-import os
 import math
+import os
 
 import numpy as np
+
 import warp as wp
 import warp.render
 
@@ -49,7 +50,6 @@ def simulate(
     positions: wp.array(dtype=wp.vec3),
     velocities: wp.array(dtype=wp.vec3),
     volume: wp.uint64,
-    restitution: float,
     margin: float,
     dt: float,
 ):
@@ -62,12 +62,12 @@ def simulate(
     xpred = x + v * dt
     xpred_local = wp.volume_world_to_index(volume, xpred)
 
-    #d = wp.volume_sample_f(volume, xpred_local, wp.Volume.LINEAR)
+    # d = wp.volume_sample_f(volume, xpred_local, wp.Volume.LINEAR)
     n = wp.vec3()
     d = wp.volume_sample_grad_f(volume, xpred_local, wp.Volume.LINEAR, n)
 
     if d < margin:
-        #n = volume_grad(volume, xpred)
+        # n = volume_grad(volume, xpred)
         n = wp.normalize(n)
         err = d - margin
 
@@ -91,14 +91,14 @@ class Example:
         self.num_particles = 10000
 
         self.sim_steps = 1000
-        self.sim_dt = 1.0 / 60.0
+        frame_dt = 1.0 / 60.0
         self.sim_substeps = 3
+        self.sim_dt = frame_dt / self.sim_substeps
 
         self.sim_time = 0.0
         self.sim_timers = {}
         self.sim_render = True
 
-        self.sim_restitution = 0.0
         self.sim_margin = 15.0
 
         self.renderer = wp.render.UsdRenderer(stage, up_axis="z")
@@ -116,9 +116,11 @@ class Example:
         # create Volume object
         self.volume = wp.Volume.load_from_nvdb(file)
 
+        file.close()
+
     def update(self):
         with wp.ScopedTimer("simulate", detailed=False, dict=self.sim_timers):
-            for s in range(self.sim_substeps):
+            for _ in range(self.sim_substeps):
                 wp.launch(
                     kernel=simulate,
                     dim=self.num_particles,
@@ -126,13 +128,11 @@ class Example:
                         self.positions,
                         self.velocities,
                         self.volume.id,
-                        self.sim_restitution,
                         self.sim_margin,
-                        self.sim_dt / float(self.sim_substeps),
+                        self.sim_dt,
                     ],
                 )
-
-            wp.synchronize_device()
+                self.sim_time += self.sim_dt
 
     def render(self, is_live=False):
         with wp.ScopedTimer("render", detailed=False):
@@ -143,16 +143,14 @@ class Example:
             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), math.pi),
-                scale=(1.0, 1.0, 1.0),
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi),
+                scale=wp.vec3(1.0, 1.0, 1.0),
             )
             self.renderer.render_points(name="points", points=self.positions.numpy(), radius=self.sim_margin)
 
             self.renderer.end_frame()
 
-        self.sim_time += self.sim_dt
-
 
 if __name__ == "__main__":
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_nvdb.usd")

examples/example_raycast.py

@@ -13,13 +13,12 @@
 #
 ##############################################################################
 
-import matplotlib.pyplot as plt
-from pxr import Usd, UsdGeom
+import os
 
-import warp as wp
 import numpy as np
+from pxr import Usd, UsdGeom
 
-import os
+import warp as wp
 
 wp.init()
 
@@ -54,7 +53,7 @@ def draw(mesh: wp.uint64, cam_pos: wp.vec3, width: int, height: int, pixels: wp.
 
 
 class Example:
-    def __init__(self):
+    def __init__(self, **kwargs):
         self.width = 1024
         self.height = 1024
         self.cam_pos = (0.0, 1.0, 2.0)
@@ -75,7 +74,7 @@ class Example:
     def update(self):
         pass
 
-    def render(self, is_live=False):
+    def render(self):
         with wp.ScopedTimer("render"):
             wp.launch(
                 kernel=draw,
@@ -83,14 +82,16 @@ class Example:
                 inputs=[self.mesh.id, self.cam_pos, self.width, self.height, self.pixels],
             )
 
-            wp.synchronize_device()
-
-        plt.imshow(
-            self.pixels.numpy().reshape((self.height, self.width, 3)), origin="lower", interpolation="antialiased"
-        )
-        plt.show()
-
 
 if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+
     example = Example()
     example.render()
+
+    wp.synchronize_device()
+
+    plt.imshow(
+        example.pixels.numpy().reshape((example.height, example.width, 3)), origin="lower", interpolation="antialiased"
+    )
+    plt.show()

examples/example_raymarch.py

@@ -15,20 +15,16 @@
 ###########################################################################
 
 
-import matplotlib.pyplot as plt
-
 import warp as wp
 
 wp.init()
 
 
-# signed sphere
 @wp.func
 def sdf_sphere(p: wp.vec3, r: float):
     return wp.length(p) - r
 
 
-# signed box
 @wp.func
 def sdf_box(upper: wp.vec3, p: wp.vec3):
     qx = wp.abs(p[0]) - upper[0]
@@ -45,19 +41,16 @@ def sdf_plane(p: wp.vec3, plane: wp.vec4):
     return plane[0] * p[0] + plane[1] * p[1] + plane[2] * p[2] + plane[3]
 
 
-# union
 @wp.func
 def op_union(d1: float, d2: float):
     return wp.min(d1, d2)
 
 
-# subtraction
 @wp.func
 def op_subtract(d1: float, d2: float):
     return wp.max(-d1, d2)
 
 
-# intersection
 @wp.func
 def op_intersect(d1: float, d2: float):
     return wp.max(d1, d2)
@@ -66,12 +59,9 @@ def op_intersect(d1: float, d2: float):
 # simple scene
 @wp.func
 def sdf(p: wp.vec3):
-    # intersection of two spheres
-    sphere_1 = wp.vec3(0.0, 0.0, 0.0)
-    sphere_2 = wp.vec3(0.0, 0.75, 0.0)
+    sphere_1 = wp.vec3(0.0, 0.75, 0.0)
 
-    d = op_subtract(sdf_sphere(p - sphere_2, 0.75), sdf_box(wp.vec3(1.0, 0.5, 0.5), p))
-    # sdf_sphere(p + sphere_1, 1.0))
+    d = op_subtract(sdf_sphere(p - sphere_1, 0.75), sdf_box(wp.vec3(1.0, 0.5, 0.5), p))
 
     # ground plane
     d = op_union(d, sdf_plane(p, wp.vec4(0.0, 1.0, 0.0, 1.0)))
@@ -96,7 +86,7 @@ def shadow(ro: wp.vec3, rd: wp.vec3):
     t = float(0.0)
     s = float(1.0)
 
-    for i in range(64):
+    for _ in range(64):
         d = sdf(ro + t * rd)
         t = t + wp.clamp(d, 0.0001, 2.0)
 
@@ -127,7 +117,7 @@ def draw(cam_pos: wp.vec3, cam_rot: wp.quat, width: int, height: int, pixels: wp
     t = float(0.0)
 
     # ray march
-    for i in range(128):
+    for _ in range(128):
         d = sdf(ro + rd * t)
         t = t + d
 
@@ -156,7 +146,7 @@ def draw(cam_pos: wp.vec3, cam_rot: wp.quat, width: int, height: int, pixels: wp
 
 
 class Example:
-    def __init__(self):
+    def __init__(self, **kwargs):
         self.width = 2048
         self.height = 1024
         self.cam_pos = (-1.25, 1.0, 2.0)
@@ -164,7 +154,10 @@ class Example:
 
         self.pixels = wp.zeros(self.width * self.height, dtype=wp.vec3)
 
-    def render(self, is_live=False):
+    def update(self):
+        pass
+
+    def render(self):
         with wp.ScopedTimer("render"):
             wp.launch(
                 kernel=draw,
@@ -172,14 +165,14 @@ class Example:
                 inputs=[self.cam_pos, self.cam_rot, self.width, self.height, self.pixels],
             )
 
-            wp.synchronize_device()
-
-        plt.imshow(
-            self.pixels.numpy().reshape((self.height, self.width, 3)), origin="lower", interpolation="antialiased"
-        )
-        plt.show()
-
 
 if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+
     example = Example()
     example.render()
+
+    plt.imshow(
+        example.pixels.numpy().reshape((example.height, example.width, 3)), origin="lower", interpolation="antialiased"
+    )
+    plt.show()

examples/example_render_opengl.py

@@ -27,6 +27,8 @@ split_up_tiles = True
 custom_tile_arrangement = False
 # whether to display the pixels in a matplotlib figure
 show_plot = True
+# whether to render depth image to a Warp array
+render_mode = "depth"
 
 renderer = wp.render.OpenGLRenderer(vsync=False)
 instance_ids = []
@@ -51,11 +53,12 @@ if num_tiles > 1:
 
 renderer.render_ground()
 
+channels = 1 if render_mode == "depth" else 3
 if show_plot:
     import matplotlib.pyplot as plt
 
     if split_up_tiles:
-        pixels = wp.zeros((num_tiles, renderer.tile_height, renderer.tile_width, 3), dtype=wp.float32)
+        pixels = wp.zeros((num_tiles, renderer.tile_height, renderer.tile_width, channels), dtype=wp.float32)
         ncols = int(np.ceil(np.sqrt(num_tiles)))
         nrows = int(np.ceil(num_tiles / float(ncols)))
         img_plots = []
@@ -70,17 +73,23 @@ if show_plot:
             sharey=True,
             num=1,
         )
-        tile_temp = np.zeros((renderer.tile_height, renderer.tile_width, 3), dtype=np.float32)
+        tile_temp = np.zeros((renderer.tile_height, renderer.tile_width, channels), dtype=np.float32)
         for dim in range(ncols * nrows):
             ax = axes[dim // ncols, dim % ncols]
             if dim >= num_tiles:
                 ax.axis("off")
                 continue
-            img_plots.append(ax.imshow(tile_temp))
+            if render_mode == "depth":
+                img_plots.append(ax.imshow(tile_temp, vmin=renderer.camera_near_plane, vmax=renderer.camera_far_plane))
+            else:
+                img_plots.append(ax.imshow(tile_temp))
     else:
         fig = plt.figure(1)
-        pixels = wp.zeros((renderer.screen_height, renderer.screen_width, 3), dtype=wp.float32)
-        img_plot = plt.imshow(pixels.numpy())
+        pixels = wp.zeros((renderer.screen_height, renderer.screen_width, channels), dtype=wp.float32)
+        if render_mode == "depth":
+            img_plot = plt.imshow(pixels.numpy(), vmin=renderer.camera_near_plane, vmax=renderer.camera_far_plane)
+        else:
+            img_plot = plt.imshow(pixels.numpy())
 
     plt.ion()
     plt.show()
@@ -95,14 +104,14 @@ while renderer.is_running():
     renderer.render_cylinder(
         "cylinder",
         [3.2, 1.0, np.sin(time + 0.5)],
-        np.array(wp.quat_from_axis_angle((1.0, 0.0, 0.0), np.sin(time + 0.5))),
+        np.array(wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), wp.sin(time + 0.5))),
         radius=0.5,
         half_height=0.8,
     )
     renderer.render_cone(
         "cone",
         [-1.2, 1.0, 0.0],
-        np.array(wp.quat_from_axis_angle((0.707, 0.707, 0.0), time)),
+        np.array(wp.quat_from_axis_angle(wp.vec3(0.707, 0.707, 0.0), time)),
         radius=0.5,
         half_height=0.8,
     )
@@ -110,15 +119,15 @@ while renderer.is_running():
 
     if show_plot and plt.fignum_exists(1):
         if split_up_tiles:
-            pixel_shape = (num_tiles, renderer.tile_height, renderer.tile_width, 3)
+            pixel_shape = (num_tiles, renderer.tile_height, renderer.tile_width, channels)
         else:
-            pixel_shape = (renderer.screen_height, renderer.screen_width, 3)
+            pixel_shape = (renderer.screen_height, renderer.screen_width, channels)
 
         if pixel_shape != pixels.shape:
             # make sure we resize the pixels array to the right dimensions if the user resizes the window
             pixels = wp.zeros(pixel_shape, dtype=wp.float32)
 
-        renderer.get_pixels(pixels, split_up_tiles=split_up_tiles)
+        renderer.get_pixels(pixels, split_up_tiles=split_up_tiles, mode=render_mode)
 
         if split_up_tiles:
             pixels_np = pixels.numpy()

examples/example_sim_cartpole.py

@@ -14,8 +14,8 @@
 #
 ###########################################################################
 
-import os
 import math
+import os
 
 import numpy as np
 
@@ -27,22 +27,11 @@ wp.init()
 
 
 class Example:
-    frame_dt = 1.0 / 60.0
-
-    episode_duration = 20.0  # seconds
-    episode_frames = int(episode_duration / frame_dt)
-
-    sim_substeps = 10
-    sim_dt = frame_dt / sim_substeps
-    sim_steps = int(episode_duration / sim_dt)
+    def __init__(self, stage=None, num_envs=1, enable_rendering=True, print_timers=True):
+        self.device = wp.get_device()
 
-    sim_time = 0.0
-
-    def __init__(self, stage=None, render=True, num_envs=1):
         builder = wp.sim.ModelBuilder()
 
-        self.enable_rendering = render
-
         self.num_envs = num_envs
 
         articulation_builder = wp.sim.ModelBuilder()
@@ -50,7 +39,7 @@ class Example:
         wp.sim.parse_urdf(
             os.path.join(os.path.dirname(__file__), "assets/cartpole.urdf"),
             articulation_builder,
-            xform=wp.transform(np.zeros(3), wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)),
+            xform=wp.transform(wp.vec3(), wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)),
             floating=False,
             density=100,
             armature=0.1,
@@ -67,6 +56,15 @@ class Example:
 
         builder = wp.sim.ModelBuilder()
 
+        self.sim_time = 0.0
+        self.frame_dt = 1.0 / 60.0
+
+        episode_duration = 20.0  # seconds
+        self.episode_frames = int(episode_duration / self.frame_dt)
+
+        self.sim_substeps = 10
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
         for i in range(num_envs):
             builder.add_builder(
                 articulation_builder, xform=wp.transform(np.array((i * 2.0, 4.0, 0.0)), wp.quat_identity())
@@ -86,57 +84,62 @@ class Example:
 
         self.integrator = wp.sim.SemiImplicitIntegrator()
 
-        # -----------------------
-        # set up Usd renderer
+        self.enable_rendering = enable_rendering
         self.renderer = None
-        if render:
-            self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=15.0)
-
-    def update(self):
-        for _ in range(self.sim_substeps):
-            self.state.clear_forces()
-            self.state = self.integrator.simulate(self.model, self.state, self.state, self.sim_dt)
-
-    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)
-        self.renderer.end_frame()
+        if self.enable_rendering:
+            self.renderer = wp.sim.render.SimRenderer(path=stage, model=self.model, scaling=15.0)
 
-    def run(self, render=True):
-        # ---------------
-        # run simulation
+        self.print_timers = print_timers
 
-        self.sim_time = 0.0
         self.state = self.model.state()
 
         wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state)
 
-        profiler = {}
+        self.use_graph = wp.get_device().is_cuda
+        self.graph = None
 
-        # create update graph
-        wp.capture_begin()
+        if self.use_graph:
+            # create update graph
+            wp.capture_begin(self.device)
+            try:
+                self.update()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
-        # simulate
-        self.update()
+    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.clear_forces()
+                    self.state = self.integrator.simulate(self.model, self.state, self.state, self.sim_dt)
+            else:
+                wp.capture_launch(self.graph)
+
+            if not wp.get_device().is_capturing:
+                self.sim_time += self.frame_dt
 
-        graph = wp.capture_end()
+    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
 
-        # 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
+                self.renderer.begin_frame(time)
+                self.renderer.render(self.state)
+                self.renderer.end_frame()
+
+    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
 
@@ -145,38 +148,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_cartpole.usd")
-    robot = Example(stage, render=True, num_envs=10)
-    robot.run()
+    else:
+        stage = os.path.join(os.path.dirname(__file__), "outputs/example_sim_cartpole.usd")
+        example = Example(stage, num_envs=10)
+        example.run()