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

examples/example_sim_cloth.py

@@ -14,21 +14,20 @@
 ###########################################################################
 
 import argparse
-import os
 import math
+import os
 from enum import Enum
 
 import numpy as np
+from pxr import Usd, UsdGeom
 
 import warp as wp
-
 import warp.sim
 import warp.sim.render
 
-from pxr import Usd, UsdGeom
-
 wp.init()
 
+
 class IntegratorType(Enum):
     EULER = "euler"
     XPBD = "xpbd"
@@ -36,21 +35,12 @@ class IntegratorType(Enum):
     def __str__(self):
         return self.value
 
-class Example:
-    def __init__(self, stage):
-
-        self.parser = argparse.ArgumentParser()
-        self.parser.add_argument(
-            "--integrator",
-            help="Type of integrator",
-            type=IntegratorType,
-            choices=list(IntegratorType),
-            default=IntegratorType.EULER,
-        )
 
-        args = self.parser.parse_args()
+class Example:
+    def __init__(self, stage, integrator=IntegratorType.EULER):
+        self.device = wp.get_device()
 
-        self.integrator_type = args.integrator
+        self.integrator_type = integrator
 
         self.sim_width = 64
         self.sim_height = 32
@@ -59,18 +49,19 @@ class Example:
         self.sim_substeps = 32
         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
+        self.frame_dt = 1.0 / self.sim_fps
+        self.sim_dt = self.frame_dt / self.sim_substeps
         self.sim_time = 0.0
         self.sim_use_graph = wp.get_device().is_cuda
-        self.simulate_timer = wp.ScopedTimer("simulate", dict={})
+        self.profiler = {}
 
         builder = wp.sim.ModelBuilder()
 
         if self.integrator_type == IntegratorType.EULER:
             builder.add_cloth_grid(
-                pos=(0.0, 4.0, 0.0),
-                rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi * 0.5),
-                vel=(0.0, 0.0, 0.0),
+                pos=wp.vec3(0.0, 4.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi * 0.5),
+                vel=wp.vec3(0.0, 0.0, 0.0),
                 dim_x=self.sim_width,
                 dim_y=self.sim_height,
                 cell_x=0.1,
@@ -83,9 +74,9 @@ class Example:
             )
         else:
             builder.add_cloth_grid(
-                pos=(0.0, 4.0, 0.0),
-                rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi * 0.5),
-                vel=(0.0, 0.0, 0.0),
+                pos=wp.vec3(0.0, 4.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi * 0.5),
+                vel=wp.vec3(0.0, 0.0, 0.0),
                 dim_x=self.sim_width,
                 dim_y=self.sim_height,
                 cell_x=0.1,
@@ -109,9 +100,9 @@ class Example:
         builder.add_shape_mesh(
             body=-1,
             mesh=mesh,
-            pos=(1.0, 0.0, 1.0),
-            rot=wp.quat_from_axis_angle((0.0, 1.0, 0.0), math.pi * 0.5),
-            scale=(2.0, 2.0, 2.0),
+            pos=wp.vec3(1.0, 0.0, 1.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), math.pi * 0.5),
+            scale=wp.vec3(2.0, 2.0, 2.0),
             ke=1.0e2,
             kd=1.0e2,
             kf=1.0e1,
@@ -131,35 +122,32 @@ class Example:
         self.state_1 = self.model.state()
 
         self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=40.0)
+        self.graph = None
 
         if self.sim_use_graph:
             # create update graph
-            wp.capture_begin()
-
-            wp.sim.collide(self.model, self.state_0)
-
-            for s in range(self.sim_substeps):
-                self.state_0.clear_forces()
-
-                self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
-
-                # swap states
-                (self.state_0, self.state_1) = (self.state_1, self.state_0)
-
-            self.graph = wp.capture_end()
+            wp.capture_begin(self.device)
+            try:
+                self.update()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
     def update(self):
-        with self.simulate_timer:
-            if self.sim_use_graph:
+        with wp.ScopedTimer("simulate", dict=self.profiler):
+            if self.sim_use_graph and self.graph:
                 wp.capture_launch(self.graph)
+                self.sim_time += self.frame_dt
             else:
                 wp.sim.collide(self.model, self.state_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)
 
+                    if not wp.get_device().is_capturing:
+                        self.sim_time += self.sim_dt
+
                     # swap states
                     (self.state_0, self.state_1) = (self.state_1, self.state_0)
 
@@ -171,19 +159,28 @@ class Example:
             self.renderer.render(self.state_0)
             self.renderer.end_frame()
 
-        self.sim_time += 1.0 / self.sim_fps
-
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--integrator",
+        help="Type of integrator",
+        type=IntegratorType,
+        choices=list(IntegratorType),
+        default=IntegratorType.EULER,
+    )
+
+    args = parser.parse_args()
+
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_cloth.usd")
 
-    example = Example(stage_path)
+    example = Example(stage_path, integrator=args.integrator)
 
     for i in range(example.sim_frames):
         example.update()
         example.render()
 
-    frame_times = example.simulate_timer.dict["simulate"]
+    frame_times = example.profiler["simulate"]
     print("\nAverage frame sim time: {:.2f} ms".format(sum(frame_times) / len(frame_times)))
 
     example.renderer.save()

examples/example_sim_fk_grad.py

@@ -42,8 +42,14 @@ def step_kernel(x: wp.array(dtype=float), grad: wp.array(dtype=float), alpha: fl
     x[tid] = x[tid] - grad[tid] * alpha
 
 
-class Robot:
-    def __init__(self, render=True, num_envs=1, device=None):
+class Example:
+    def __init__(self, stage, device=None, verbose=False):
+        self.verbose = verbose
+
+        self.frame_dt = 1.0 / 60.0
+
+        self.render_time = 0.0
+
         builder = wp.sim.ModelBuilder()
 
         builder.add_articulation()
@@ -78,12 +84,12 @@ class Robot:
 
             if i == chain_length - 1:
                 # create end effector
-                s = builder.add_shape_sphere(pos=(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
+                builder.add_shape_sphere(pos=wp.vec3(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
 
             else:
                 # create shape
-                s = builder.add_shape_box(
-                    pos=(chain_width * 0.5, 0.0, 0.0), hx=chain_width * 0.5, hy=0.1, hz=0.1, density=10.0, body=b
+                builder.add_shape_box(
+                    pos=wp.vec3(chain_width * 0.5, 0.0, 0.0), hx=chain_width * 0.5, hy=0.1, hz=0.1, density=10.0, body=b
                 )
 
         self.device = wp.get_device(device)
@@ -94,16 +100,7 @@ class Robot:
 
         self.state = self.model.state()
 
-        # -----------------------
-        # set up Usd renderer
-        self.renderer = wp.sim.render.SimRenderer(
-            self.model, os.path.join(os.path.dirname(__file__), "outputs/example_sim_fk_grad.usd"), scaling=50.0
-        )
-
-    def run(self, render=True):
-        render_time = 0.0
-        train_iters = 1024
-        train_rate = 0.01
+        self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=50.0)
 
         # optimization variables
         self.loss = wp.zeros(1, dtype=float, device=self.device)
@@ -112,44 +109,54 @@ class Robot:
         self.state.body_q.requires_grad = True
         self.loss.requires_grad = True
 
-        for i in range(train_iters):
-            tape = wp.Tape()
-            with tape:
-                wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state)
-
-                wp.launch(
-                    compute_loss,
-                    dim=1,
-                    inputs=[self.state.body_q, len(self.state.body_q) - 1, self.loss],
-                    device=self.device,
-                )
+        self.train_rate = 0.01
 
-            tape.backward(loss=self.loss)
+    def update(self):
+        tape = wp.Tape()
 
-            print(self.loss)
-            print(tape.gradients[self.model.joint_q])
+        with tape:
+            wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state)
 
-            # gradient descent
             wp.launch(
-                step_kernel,
-                dim=len(self.model.joint_q),
-                inputs=[self.model.joint_q, tape.gradients[self.model.joint_q], train_rate],
+                compute_loss,
+                dim=1,
+                inputs=[self.state.body_q, len(self.state.body_q) - 1, self.loss],
                 device=self.device,
             )
 
-            # zero gradients
-            tape.zero()
+        tape.backward(loss=self.loss)
+
+        if self.verbose:
+            print(self.loss)
+            print(tape.gradients[self.model.joint_q])
+
+        # gradient descent
+        wp.launch(
+            step_kernel,
+            dim=len(self.model.joint_q),
+            inputs=[self.model.joint_q, tape.gradients[self.model.joint_q], self.train_rate],
+            device=self.device,
+        )
+
+        # zero gradients
+        tape.zero()
+
+    def render(self):
+        self.renderer.begin_frame(self.render_time)
+        self.renderer.render(self.state)
+        self.renderer.render_sphere(name="target", pos=TARGET, rot=wp.quat_identity(), radius=0.1)
+        self.renderer.end_frame()
+        self.render_time += self.frame_dt
 
-            # render
-            self.renderer.begin_frame(render_time)
-            self.renderer.render(self.state)
-            self.renderer.render_sphere(name="target", pos=TARGET, rot=wp.quat_identity(), radius=0.1)
-            self.renderer.end_frame()
 
-            render_time += 1.0 / 60.0
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_fk_grad.usd")
+    example = Example(stage_path, device=wp.get_preferred_device(), verbose=True)
 
-        self.renderer.save()
+    train_iters = 512
 
+    for _ in range(train_iters):
+        example.update()
+        example.render()
 
-robot = Robot(render=True, device=wp.get_preferred_device(), num_envs=1)
-robot.run()
+    example.renderer.save()

examples/example_sim_fk_grad_torch.py

@@ -17,14 +17,12 @@
 import os
 
 import numpy as np
-
 import torch
 
 import warp as wp
 import warp.sim
 import warp.sim.render
 
-
 wp.init()
 
 
@@ -68,8 +66,14 @@ class ForwardKinematics(torch.autograd.Function):
         return (wp.to_torch(ctx.tape.gradients[ctx.joint_q]), wp.to_torch(ctx.tape.gradients[ctx.joint_qd]), None)
 
 
-class Robot:
-    def __init__(self, render=True, num_envs=1, device=None):
+class Example:
+    def __init__(self, stage, device=None, verbose=False):
+        self.verbose = verbose
+
+        self.frame_dt = 1.0 / 60.0
+
+        self.render_time = 0.0
+
         builder = wp.sim.ModelBuilder()
 
         builder.add_articulation()
@@ -91,7 +95,7 @@ class Robot:
             builder.add_joint_revolute(
                 parent=parent,
                 child=b,
-                axis=(0.0, 0.0, 1.0),
+                axis=wp.vec3(0.0, 0.0, 1.0),
                 parent_xform=parent_joint_xform,
                 child_xform=wp.transform_identity(),
                 limit_lower=-np.deg2rad(60.0),
@@ -104,12 +108,12 @@ class Robot:
 
             if i == chain_length - 1:
                 # create end effector
-                s = builder.add_shape_sphere(pos=(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
+                builder.add_shape_sphere(pos=wp.vec3(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
 
             else:
                 # create shape
-                s = builder.add_shape_box(
-                    pos=(chain_width * 0.5, 0.0, 0.0), hx=chain_width * 0.5, hy=0.1, hz=0.1, density=10.0, body=b
+                builder.add_shape_box(
+                    pos=wp.vec3(chain_width * 0.5, 0.0, 0.0), hx=chain_width * 0.5, hy=0.1, hz=0.1, density=10.0, body=b
                 )
 
         # finalize model
@@ -118,50 +122,53 @@ class Robot:
 
         self.torch_device = wp.device_to_torch(self.model.device)
 
-        # -----------------------
-        # set up Usd renderer
-        self.renderer = wp.sim.render.SimRenderer(
-            self.model, os.path.join(os.path.dirname(__file__), "outputs/example_sim_fk_grad.usd"), scaling=50.0
-        )
+        self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=50.0)
 
-    def run(self, render=True):
-        render_time = 0.0
-        train_iters = 1024
-        train_rate = 0.01
+        self.target = torch.from_numpy(np.array((2.0, 1.0, 0.0))).to(self.torch_device)
 
-        target = torch.from_numpy(np.array((2.0, 1.0, 0.0))).to(self.torch_device)
+        self.body_q = None
+        self.body_qd = None
 
         # optimization variable
-        joint_q = torch.zeros(len(self.model.joint_q), requires_grad=True, device=self.torch_device)
-        joint_qd = torch.zeros(len(self.model.joint_qd), requires_grad=True, device=self.torch_device)
+        self.joint_q = torch.zeros(len(self.model.joint_q), requires_grad=True, device=self.torch_device)
+        self.joint_qd = torch.zeros(len(self.model.joint_qd), requires_grad=True, device=self.torch_device)
+
+        self.train_rate = 0.01
 
-        for i in range(train_iters):
-            (body_q, body_qd) = ForwardKinematics.apply(joint_q, joint_qd, self.model)
+    def update(self):
+        (self.body_q, self.body_qd) = ForwardKinematics.apply(self.joint_q, self.joint_qd, self.model)
 
-            l = torch.norm(body_q[self.model.body_count - 1][0:3] - target) ** 2.0
-            l.backward()
+        l = torch.norm(self.body_q[self.model.body_count - 1][0:3] - self.target) ** 2.0
+        l.backward()
 
+        if self.verbose:
             print(l)
-            print(joint_q.grad)
+            print(self.joint_q.grad)
+
+        with torch.no_grad():
+            self.joint_q -= self.joint_q.grad * self.train_rate
+            self.joint_q.grad.zero_()
 
-            with torch.no_grad():
-                joint_q -= joint_q.grad * train_rate
-                joint_q.grad.zero_()
+    def render(self):
+        s = self.model.state()
+        s.body_q = wp.from_torch(self.body_q, dtype=wp.transform, requires_grad=False)
+        s.body_qd = wp.from_torch(self.body_qd, dtype=wp.spatial_vector, requires_grad=False)
 
-            # render
-            s = self.model.state()
-            s.body_q = wp.from_torch(body_q, dtype=wp.transform)
-            s.body_qd = wp.from_torch(body_qd, dtype=wp.spatial_vector)
+        self.renderer.begin_frame(self.render_time)
+        self.renderer.render(s)
+        self.renderer.render_sphere(name="target", pos=self.target, rot=wp.quat_identity(), radius=0.1)
+        self.renderer.end_frame()
+        self.render_time += self.frame_dt
 
-            self.renderer.begin_frame(render_time)
-            self.renderer.render(s)
-            self.renderer.render_sphere(name="target", pos=target, rot=wp.quat_identity(), radius=0.1)
-            self.renderer.end_frame()
 
-            render_time += 1.0 / 60.0
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_fk_grad.usd")
+    example = Example(stage_path, device=wp.get_preferred_device(), verbose=True)
 
-        self.renderer.save()
+    train_iters = 512
 
+    for _ in range(train_iters):
+        example.update()
+        example.render()
 
-robot = Robot(render=True, device="cuda", num_envs=1)
-robot.run()
+    example.renderer.save()