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

examples/example_sim_grad_bounce.py

@@ -20,45 +20,60 @@
 import os
 
 import numpy as np
-import warp as wp
 
+import warp as wp
 import warp.sim
 import warp.sim.render
 
 wp.init()
 
 
-class Bounce:
-    # seconds
-    sim_duration = 0.6
+@wp.kernel
+def loss_kernel(pos: wp.array(dtype=wp.vec3), target: wp.vec3, loss: wp.array(dtype=float)):
+    # distance to target
+    delta = pos[0] - target
+    loss[0] = wp.dot(delta, delta)
 
-    # control frequency
-    frame_dt = 1.0 / 60.0
-    frame_steps = int(sim_duration / frame_dt)
 
-    # sim frequency
-    sim_substeps = 8
-    sim_steps = frame_steps * sim_substeps
-    sim_dt = frame_dt / sim_substeps
-    sim_time = 0.0
+@wp.kernel
+def step_kernel(x: wp.array(dtype=wp.vec3), grad: wp.array(dtype=wp.vec3), alpha: float):
+    tid = wp.tid()
 
-    render_time = 0.0
+    # gradient descent step
+    x[tid] = x[tid] - grad[tid] * alpha
 
-    train_iters = 250
-    train_rate = 0.02
 
-    ke = 1.0e4
-    kf = 0.0
-    kd = 1.0e1
-    mu = 0.25
+class Example:
+    def __init__(self, stage=None, enable_rendering=True, profile=False, adapter=None, verbose=False):
+        self.device = wp.get_device()
+        self.verbose = verbose
 
-    def __init__(self, render=True, profile=False, adapter=None):
-        builder = wp.sim.ModelBuilder()
+        # seconds
+        sim_duration = 0.6
+
+        # control frequency
+        self.frame_dt = 1.0 / 60.0
+        frame_steps = int(sim_duration / self.frame_dt)
+
+        # sim frequency
+        self.sim_substeps = 8
+        self.sim_steps = frame_steps * self.sim_substeps
+        self.sim_dt = self.frame_dt / self.sim_substeps
 
-        builder.add_particle(pos=(-0.5, 1.0, 0.0), vel=(5.0, -5.0, 0.0), mass=1.0)
-        builder.add_shape_box(
-            body=-1, pos=(2.0, 1.0, 0.0), hx=0.25, hy=1.0, hz=1.0, ke=self.ke, kf=self.kf, kd=self.kd, mu=self.mu
-        )
+        self.iter = 0
+        self.render_time = 0.0
+
+        self.train_iters = 250
+        self.train_rate = 0.02
+
+        ke = 1.0e4
+        kf = 0.0
+        kd = 1.0e1
+        mu = 0.2
+
+        builder = wp.sim.ModelBuilder()
+        builder.add_particle(pos=wp.vec3(-0.5, 1.0, 0.0), vel=wp.vec3(5.0, -5.0, 0.0), mass=1.0)
+        builder.add_shape_box(body=-1, pos=wp.vec3(2.0, 1.0, 0.0), hx=0.25, hy=1.0, hz=1.0, ke=ke, kf=kf, kd=kd, mu=mu)
 
         self.device = wp.get_device(adapter)
         self.profile = profile
@@ -66,10 +81,10 @@ class Bounce:
         self.model = builder.finalize(self.device)
         self.model.ground = True
 
-        self.model.soft_contact_ke = self.ke
-        self.model.soft_contact_kf = self.kf
-        self.model.soft_contact_kd = self.kd
-        self.model.soft_contact_mu = self.mu
+        self.model.soft_contact_ke = ke
+        self.model.soft_contact_kf = kf
+        self.model.soft_contact_kd = kd
+        self.model.soft_contact_mu = mu
         self.model.soft_contact_margin = 10.0
         self.model.soft_contact_restitution = 1.0
 
@@ -86,24 +101,20 @@ class Bounce:
         # one-shot contact creation (valid if we're doing simple collision against a constant normal plane)
         wp.sim.collide(self.model, self.states[0])
 
-        self.stage = None
-        if render:
-            self.stage = wp.sim.render.SimRendererOpenGL(
-                self.model, os.path.join(os.path.dirname(__file__), "outputs/example_sim_grad_bounce.usd"), scaling=1.0
-            )
-
-    @wp.kernel
-    def loss_kernel(pos: wp.array(dtype=wp.vec3), target: wp.vec3, loss: wp.array(dtype=float)):
-        # distance to target
-        delta = pos[0] - target
-        loss[0] = wp.dot(delta, delta)
+        self.enable_rendering = enable_rendering
+        self.renderer = None
+        if self.enable_rendering:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=1.0)
 
-    @wp.kernel
-    def step_kernel(x: wp.array(dtype=wp.vec3), grad: wp.array(dtype=wp.vec3), alpha: float):
-        tid = wp.tid()
-
-        # gradient descent step
-        x[tid] = x[tid] - grad[tid] * alpha
+        # capture forward/backward passes
+        wp.capture_begin(self.device)
+        try:
+            self.tape = wp.Tape()
+            with self.tape:
+                self.compute_loss()
+            self.tape.backward(self.loss)
+        finally:
+            self.graph = wp.capture_end(self.device)
 
     def compute_loss(self):
         # run control loop
@@ -113,38 +124,53 @@ class Bounce:
             self.integrator.simulate(self.model, self.states[i], self.states[i + 1], self.sim_dt)
 
         # compute loss on final state
-        wp.launch(
-            self.loss_kernel, dim=1, inputs=[self.states[-1].particle_q, self.target, self.loss], device=self.device
-        )
+        wp.launch(loss_kernel, dim=1, inputs=[self.states[-1].particle_q, self.target, self.loss], device=self.device)
 
         return self.loss
 
-    def render(self, iter):
-        if self.stage is None:
-            return
+    def update(self):
+        with wp.ScopedTimer("Step", active=self.profile):
+            # forward + backward
+            wp.capture_launch(self.graph)
 
-        # render every 16 iters
-        if iter % 16 > 0:
-            return
+            # gradient descent step
+            x = self.states[0].particle_qd
+            wp.launch(step_kernel, dim=len(x), inputs=[x, x.grad, self.train_rate], device=self.device)
 
-        # draw trajectory
-        traj_verts = [self.states[0].particle_q.numpy()[0].tolist()]
+            x_grad = self.tape.gradients[self.states[0].particle_qd]
 
-        for i in range(0, self.sim_steps, self.sim_substeps):
-            traj_verts.append(self.states[i].particle_q.numpy()[0].tolist())
+            if self.verbose:
+                print(f"Iter: {self.iter} Loss: {self.loss}")
+                print(f"    x: {x} g: {x_grad}")
 
-            self.stage.begin_frame(self.render_time)
-            self.stage.render(self.states[i])
-            self.stage.render_box(pos=self.target, rot=wp.quat_identity(), extents=(0.1, 0.1, 0.1), name="target")
-            self.stage.render_line_strip(
-                vertices=traj_verts,
-                color=wp.render.bourke_color_map(0.0, 7.0, self.loss.numpy()[0]),
-                radius=0.02,
-                name=f"traj_{iter}",
-            )
-            self.stage.end_frame()
+            # clear grads for next iteration
+            self.tape.zero()
 
-            self.render_time += self.frame_dt
+            self.iter = self.iter + 1
+
+    def render(self):
+        if self.enable_rendering:
+            with wp.ScopedTimer("Render", active=self.profile):
+                # draw trajectory
+                traj_verts = [self.states[0].particle_q.numpy()[0].tolist()]
+
+                for i in range(0, self.sim_steps, self.sim_substeps):
+                    traj_verts.append(self.states[i].particle_q.numpy()[0].tolist())
+
+                    self.renderer.begin_frame(self.render_time)
+                    self.renderer.render(self.states[i])
+                    self.renderer.render_box(
+                        pos=self.target, rot=wp.quat_identity(), extents=(0.1, 0.1, 0.1), name="target"
+                    )
+                    self.renderer.render_line_strip(
+                        vertices=traj_verts,
+                        color=wp.render.bourke_color_map(0.0, 7.0, self.loss.numpy()[0]),
+                        radius=0.02,
+                        name=f"traj_{self.iter-1}",
+                    )
+                    self.renderer.end_frame()
+
+                    self.render_time += self.frame_dt
 
     def check_grad(self):
         param = self.states[0].particle_qd
@@ -191,71 +217,20 @@ class Bounce:
 
         tape.zero()
 
-    def train(self):
-        for i in range(self.train_iters):
-            tape = wp.Tape()
-
-            with wp.ScopedTimer("Forward", active=self.profile):
-                with tape:
-                    self.compute_loss()
-
-            with wp.ScopedTimer("Backward", active=self.profile):
-                tape.backward(self.loss)
-
-            with wp.ScopedTimer("Render", active=self.profile):
-                self.render(i)
-
-            with wp.ScopedTimer("Step", active=self.profile):
-                x = self.states[0].particle_qd
-                x_grad = tape.gradients[self.states[0].particle_qd]
-
-                print(f"Iter: {i} Loss: {self.loss}")
-                print(f"   x: {x} g: {x_grad}")
-
-                wp.launch(self.step_kernel, dim=len(x), inputs=[x, x_grad, self.train_rate], device=self.device)
-
-            tape.zero()
-
-        if self.stage is not None:
-            self.stage.save()
-
-    def train_graph(self):
-        # capture forward/backward passes
-        wp.capture_begin()
-
-        tape = wp.Tape()
-        with tape:
-            self.compute_loss()
-
-        tape.backward(self.loss)
-
-        self.graph = wp.capture_end()
-
+    def run(self):
         # replay and optimize
         for i in range(self.train_iters):
-            with wp.ScopedTimer("Step", active=self.profile):
-                # forward + backward
-                wp.capture_launch(self.graph)
-
-                # gradient descent step
-                x = self.states[0].particle_qd
-                wp.launch(self.step_kernel, dim=len(x), inputs=[x, x.grad, self.train_rate], device=self.device)
-
-                x_grad = tape.gradients[self.states[0].particle_qd]
-
-                print(f"Iter: {i} Loss: {self.loss}")
-                print(f"   x: {x} g: {x_grad}")
-
-                # clear grads for next iteration
-                tape.zero()
-
-            with wp.ScopedTimer("Render", active=self.profile):
-                self.render(i)
+            self.update()
+            # render every 16 iters
+            if i % 16 == 0:
+                self.render()
 
-        if self.stage is not None:
-            self.stage.save()
+        if self.enable_rendering:
+            self.renderer.save()
 
 
-bounce = Bounce(profile=False, render=True)
-bounce.check_grad()
-bounce.train_graph()
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_grad_bounce.usd")
+    example = Example(stage_path, profile=False, enable_rendering=True, verbose=True)
+    example.check_grad()
+    example.run()

examples/example_sim_grad_cloth.py

@@ -9,7 +9,7 @@
 # Example Sim Grad Cloth
 #
 # Shows how to use Warp to optimize the initial velocities of a piece of
-# cloth such that it's center of mass hits a target after a specified time.
+# cloth such that its center of mass hits a target after a specified time.
 #
 # This example uses the built-in wp.Tape() object to compute gradients of
 # the distance to target (loss) w.r.t the initial velocity, followed by
@@ -17,8 +17,8 @@
 #
 ###########################################################################
 
-import os
 import math
+import os
 
 import warp as wp
 import warp.sim
@@ -27,26 +27,53 @@ import warp.sim.render
 wp.init()
 
 
-class Cloth:
-    # seconds
-    sim_duration = 2.0
+@wp.kernel
+def com_kernel(positions: wp.array(dtype=wp.vec3), n: int, com: wp.array(dtype=wp.vec3)):
+    tid = wp.tid()
+
+    # compute center of mass
+    wp.atomic_add(com, 0, positions[tid] / float(n))
+
+
+@wp.kernel
+def loss_kernel(com: wp.array(dtype=wp.vec3), target: wp.vec3, loss: wp.array(dtype=float)):
+    # sq. distance to target
+    delta = com[0] - target
+
+    loss[0] = wp.dot(delta, delta)
+
+
+@wp.kernel
+def step_kernel(x: wp.array(dtype=wp.vec3), grad: wp.array(dtype=wp.vec3), alpha: float):
+    tid = wp.tid()
 
-    # control frequency
-    frame_dt = 1.0 / 60.0
-    frame_steps = int(sim_duration / frame_dt)
+    # gradient descent step
+    x[tid] = x[tid] - grad[tid] * alpha
 
-    # sim frequency
-    sim_substeps = 16
-    sim_steps = frame_steps * sim_substeps
-    sim_dt = frame_dt / sim_substeps
-    sim_time = 0.0
 
-    render_time = 0.0
+class Example:
+    def __init__(self, stage, profile=False, adapter=None, verbose=False):
+        self.device = wp.get_device()
+        self.verbose = verbose
 
-    train_iters = 64
-    train_rate = 5.0
+        # seconds
+        sim_duration = 2.0
+
+        # control frequency
+        self.frame_dt = 1.0 / 60.0
+        frame_steps = int(sim_duration / self.frame_dt)
+
+        # sim frequency
+        self.sim_substeps = 16
+        self.sim_steps = frame_steps * self.sim_substeps
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
+        self.iter = 0
+        self.render_time = 0.0
+
+        self.train_iters = 64
+        self.train_rate = 5.0
 
-    def __init__(self, render=True, profile=False, adapter=None):
         builder = wp.sim.ModelBuilder()
         builder.default_particle_radius = 0.01
 
@@ -54,9 +81,9 @@ class Cloth:
         dim_y = 16
 
         builder.add_cloth_grid(
-            pos=(0.0, 0.0, 0.0),
-            vel=(0.1, 0.1, 0.0),
-            rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.25),
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            vel=wp.vec3(0.1, 0.1, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.25),
             dim_x=dim_x,
             dim_y=dim_y,
             cell_x=1.0 / dim_x,
@@ -86,31 +113,17 @@ class Cloth:
         for i in range(self.sim_steps + 1):
             self.states.append(self.model.state(requires_grad=True))
 
-        if self.render:
-            self.stage = wp.sim.render.SimRendererOpenGL(
-                self.model, os.path.join(os.path.dirname(__file__), "outputs/example_sim_grad_cloth.usd"), scaling=4.0
-            )
-
-    @wp.kernel
-    def com_kernel(positions: wp.array(dtype=wp.vec3), n: int, com: wp.array(dtype=wp.vec3)):
-        tid = wp.tid()
-
-        # compute center of mass
-        wp.atomic_add(com, 0, positions[tid] / float(n))
-
-    @wp.kernel
-    def loss_kernel(com: wp.array(dtype=wp.vec3), target: wp.vec3, loss: wp.array(dtype=float)):
-        # sq. distance to target
-        delta = com[0] - target
-
-        loss[0] = wp.dot(delta, delta)
+        self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=4.0)
 
-    @wp.kernel
-    def step_kernel(x: wp.array(dtype=wp.vec3), grad: wp.array(dtype=wp.vec3), alpha: float):
-        tid = wp.tid()
-
-        # gradient descent step
-        x[tid] = x[tid] - grad[tid] * alpha
+        # capture forward/backward passes
+        wp.capture_begin(self.device)
+        try:
+            self.tape = wp.Tape()
+            with self.tape:
+                self.compute_loss()
+            self.tape.backward(self.loss)
+        finally:
+            self.graph = wp.capture_end(self.device)
 
     def compute_loss(self):
         # run control loop
@@ -122,92 +135,65 @@ class Cloth:
         # compute loss on final state
         self.com.zero_()
         wp.launch(
-            self.com_kernel,
+            com_kernel,
             dim=self.model.particle_count,
             inputs=[self.states[-1].particle_q, self.model.particle_count, self.com],
             device=self.device,
         )
-        wp.launch(self.loss_kernel, dim=1, inputs=[self.com, self.target, self.loss], device=self.device)
+        wp.launch(loss_kernel, dim=1, inputs=[self.com, self.target, self.loss], device=self.device)
 
         return self.loss
 
-    def render(self, iter):
-        # render every 4 iters
-        if iter % 4 > 0:
-            return
-
-        # draw trajectory
-        traj_verts = [self.states[0].particle_q.numpy().mean(axis=0)]
-
-        for i in range(0, self.sim_steps, self.sim_substeps):
-            traj_verts.append(self.states[i].particle_q.numpy().mean(axis=0))
-
-            self.stage.begin_frame(self.render_time)
-            self.stage.render(self.states[i])
-            self.stage.render_box(pos=self.target, rot=wp.quat_identity(), extents=(0.1, 0.1, 0.1), name="target")
-            self.stage.render_line_strip(
-                vertices=traj_verts,
-                color=wp.render.bourke_color_map(0.0, 269.0, self.loss.numpy()[0]),
-                radius=0.02,
-                name=f"traj_{iter}",
-            )
-            self.stage.end_frame()
-
-            self.render_time += self.frame_dt
-
-    def train(self, mode="gd"):
-        tape = wp.Tape()
-
-        for i in range(self.train_iters):
-            with wp.ScopedTimer("Forward", active=self.profile):
-                with tape:
-                    self.compute_loss()
+    def update(self):
+        with wp.ScopedTimer("Step", active=self.profile):
+            wp.capture_launch(self.graph)
 
-            with wp.ScopedTimer("Backward", active=self.profile):
-                tape.backward(self.loss)
+            # gradient descent step
+            x = self.states[0].particle_qd
 
-            with wp.ScopedTimer("Render", active=self.profile):
-                self.render(i)
+            if self.verbose:
+                print(f"Iter: {self.iter} Loss: {self.loss}")
 
-            with wp.ScopedTimer("Step", active=self.profile):
-                x = self.states[0].particle_qd
-                x_grad = tape.gradients[self.states[0].particle_qd]
+            wp.launch(step_kernel, dim=len(x), inputs=[x, x.grad, self.train_rate], device=self.device)
 
-                print(f"Iter: {i} Loss: {self.loss}")
+            # clear grads for next iteration
+            self.tape.zero()
 
-                wp.launch(self.step_kernel, dim=len(x), inputs=[x, x_grad, self.train_rate], device=self.device)
+            self.iter = self.iter + 1
 
-            tape.reset()
+    def render(self):
+        with wp.ScopedTimer("Render", active=self.profile):
+            # draw trajectory
+            traj_verts = [self.states[0].particle_q.numpy().mean(axis=0)]
 
-    def train_graph(self, mode="gd"):
-        wp.capture_begin()
+            for i in range(0, self.sim_steps, self.sim_substeps):
+                traj_verts.append(self.states[i].particle_q.numpy().mean(axis=0))
 
-        tape = wp.Tape()
-        with tape:
-            self.compute_loss()
+                self.renderer.begin_frame(self.render_time)
+                self.renderer.render(self.states[i])
+                self.renderer.render_box(
+                    pos=self.target, rot=wp.quat_identity(), extents=(0.1, 0.1, 0.1), name="target"
+                )
+                self.renderer.render_line_strip(
+                    vertices=traj_verts,
+                    color=wp.render.bourke_color_map(0.0, 269.0, self.loss.numpy()[0]),
+                    radius=0.02,
+                    name=f"traj_{self.iter-1}",
+                )
+                self.renderer.end_frame()
 
-        tape.backward(self.loss)
+                self.render_time += self.frame_dt
 
-        self.graph = wp.capture_end()
 
-        for i in range(self.train_iters):
-            with wp.ScopedTimer("Replay", active=self.profile):
-                wp.capture_launch(self.graph)
-
-            with wp.ScopedTimer("Render", active=self.profile):
-                self.render(i)
-
-            with wp.ScopedTimer("Step", active=self.profile):
-                x = self.states[0].particle_qd
-
-                print(f"Iter: {i} Loss: {self.loss}")
-
-                wp.launch(self.step_kernel, dim=len(x), inputs=[x, x.grad, self.train_rate], device=self.device)
-
-            tape.zero()
-
-        self.stage.save()
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_grad_cloth.usd")
+    example = Example(stage_path, profile=False, verbose=True)
 
+    # replay and optimize
+    for i in range(example.train_iters):
+        example.update()
+        # render every 4 iters
+        if i % 4 == 0:
+            example.render()
 
-bounce = Cloth(profile=False, render=True)
-bounce.train_graph("gd")
+    example.renderer.save()

examples/example_sim_granular.py

@@ -44,9 +44,9 @@ class Example:
             cell_x=self.radius * 2.0,
             cell_y=self.radius * 2.0,
             cell_z=self.radius * 2.0,
-            pos=(0.0, 1.0, 0.0),
+            pos=wp.vec3(0.0, 1.0, 0.0),
             rot=wp.quat_identity(),
-            vel=(5.0, 0.0, 0.0),
+            vel=wp.vec3(5.0, 0.0, 0.0),
             mass=0.1,
             jitter=self.radius * 0.1,
         )
@@ -68,10 +68,11 @@ 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()
 
                 self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+                self.sim_time += self.frame_dt
 
                 # swap states
                 (self.state_0, self.state_1) = (self.state_1, self.state_0)
@@ -84,15 +85,13 @@ class Example:
             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_granular.usd")
 
     example = Example(stage_path)
 
-    for i in range(example.frame_count):
+    for _ in range(example.frame_count):
         example.update()
         example.render()