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

examples/example_sim_trajopt.py

@@ -23,15 +23,15 @@ import warp.sim
 import warp.sim.render
 from warp.optim import Adam
 
-import matplotlib.pyplot as plt
-
 wp.init()
 
 
 @wp.kernel
-def loss_l2(states: wp.array(dtype=wp.float32), targets: wp.array(dtype=wp.float32), loss: wp.array(dtype=wp.float32)):
-    i = wp.tid()
-    diff = states[i] - targets[i]
+def loss_l2(
+    states: wp.array2d(dtype=wp.float32), targets: wp.array2d(dtype=wp.float32), loss: wp.array(dtype=wp.float32)
+):
+    i, j = wp.tid()
+    diff = states[i, j] - targets[i, j]
     l = diff * diff
     wp.atomic_add(loss, 0, l)
 
@@ -44,23 +44,25 @@ def apply_torque(torques: wp.array(dtype=wp.float32), start_index: int, body_f:
 
 
 @wp.kernel
-def save_state(body_q: wp.array(dtype=wp.transform), start_index: int, states: wp.array(dtype=wp.float32)):
+def save_state(body_q: wp.array(dtype=wp.transform), write_index: int, states: wp.array2d(dtype=wp.float32)):
     pos = wp.transform_get_translation(body_q[0])
-    states[start_index + 0] = pos[0]
-    states[start_index + 1] = pos[2]
+    states[write_index, 0] = pos[0]
+    states[write_index, 1] = pos[2]
+
 
+class Example:
+    def __init__(self, stage, device=None, verbose=False):
+        self.verbose = verbose
+        self.frame_dt = 1.0 / 60.0
+        self.episode_frames = 100
 
-class Environment:
-    frame_dt = 1.0 / 60.0
-    episode_frames = 100
+        self.sim_substeps = 1
+        self.sim_dt = self.frame_dt / self.sim_substeps
 
-    sim_substeps = 1
-    sim_dt = frame_dt / sim_substeps
+        self.render_time = 0.0
 
-    sim_time = 0.0
-    render_time = 0.0
+        self.iter = 0
 
-    def __init__(self, device="cpu"):
         builder = wp.sim.ModelBuilder()
 
         self.device = device
@@ -69,7 +71,7 @@ class Environment:
         builder = wp.sim.ModelBuilder(gravity=0.0)
         builder.add_articulation()
         b = builder.add_body(origin=wp.transform())
-        s = builder.add_shape_box(pos=(0.0, 0.0, 0.0), hx=0.5, hy=0.5, hz=0.5, density=100.0, body=b)
+        builder.add_shape_box(pos=wp.vec3(0.0, 0.0, 0.0), hx=0.5, hy=0.5, hz=0.5, density=100.0, body=b)
 
         # compute reference trajectory
         rad = np.linspace(0.0, np.pi * 2, self.episode_frames)
@@ -78,7 +80,8 @@ class Environment:
         # set initial joint configuration to first reference state
         builder.body_q[0] = wp.transform(p=[self.ref_traj[0][0], 0.0, self.ref_traj[0][1]])
 
-        self.ref_traj = wp.array(self.ref_traj.flatten(), dtype=wp.float32, device=self.device, requires_grad=True)
+        self.ref_traj = wp.array(self.ref_traj, dtype=wp.float32, device=self.device, requires_grad=True)
+        self.last_traj = wp.empty_like(self.ref_traj)
 
         # finalize model
         self.model = builder.finalize(device, requires_grad=True)
@@ -93,127 +96,126 @@ class Environment:
         self.action_dim = 2
         self.state_dim = 2
 
-        assert len(self.ref_traj) == self.episode_frames * self.state_dim
+        assert self.ref_traj.shape == (self.episode_frames, self.state_dim)
 
         self.integrator = wp.sim.SemiImplicitIntegrator()
 
-    def simulate(self, state: wp.sim.State, action: wp.array, action_index: int, requires_grad=False) -> wp.sim.State:
+        # initial guess
+        self.actions = wp.array(
+            np.zeros(self.episode_frames * self.action_dim) * 100.0,
+            dtype=wp.float32,
+            device=self.device,
+            requires_grad=True,
+        )
+
+        self.optimizer = Adam([self.actions], lr=1e2)
+        self.loss = wp.zeros(1, dtype=wp.float32, device=self.device, requires_grad=True)
+
+        self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=100.0)
+
+        # allocate sim states for trajectory
+        self.states = []
+        for _ in range(self.episode_frames + 1):
+            self.states.append(self.model.state(requires_grad=True))
+
+    def compute_loss(self):
         """
-        Simulate the system for the given states.
+        Advances the system dynamics given the rigid-body state in maximal coordinates and generalized joint torques
+        [body_q, body_qd, tau].
         """
 
-        for _ in range(self.sim_substeps):
-            if requires_grad:
+        self.last_traj.zero_()
+
+        for i in range(self.episode_frames):
+            state = self.states[i]
+
+            for _ in range(self.sim_substeps):
                 next_state = self.model.state(requires_grad=True)
-            else:
-                next_state = state
-                next_state.clear_forces()
-
-            wp.sim.collide(self.model, state)
-            # apply generalized torques to rigid body here, instead of planar joints
-            wp.launch(apply_torque, 1, inputs=[action, action_index], outputs=[state.body_f], device=action.device)
-            state = self.integrator.simulate(self.model, state, next_state, self.sim_dt, requires_grad=requires_grad)
-        return state
-
-    def _render(self, state: wp.sim.State):
-        self.renderer.begin_frame(self.render_time)
-        self.renderer.render(state)
-        self.renderer.end_frame()
-        self.render_time += self.frame_dt
-        self.renderer.save()
-
-    def forward(self, actions: wp.array, requires_grad=False, loss=None, render=False):
-        """
-        Advances the system dynamics given the rigid-body state in maximal coordinates and generalized joint torques [body_q, body_qd, tau].
-        Simulates for the set number of substeps and returns the next state in maximal and (optional) generalized coordinates [body_q_next, body_qd_next, joint_q_next, joint_qd_next].
-        """
 
-        actions.requires_grad = requires_grad
-        state = self.model.state(requires_grad=requires_grad)
-        if render:
-            # set up Usd renderer
-            self.renderer = wp.sim.render.SimRenderer(
-                self.model, os.path.join(os.path.dirname(__file__), "outputs/example_sim_trajopt.usd"), scaling=100.0
-            )
-            self.render_time = 0.0
-            self._render(state)
+                wp.sim.collide(self.model, state)
 
-        states = wp.zeros(
-            self.episode_frames * self.state_dim, dtype=wp.float32, device=self.device, requires_grad=requires_grad
-        )
+                # apply generalized torques to rigid body here, instead of planar joints
+                wp.launch(
+                    apply_torque,
+                    1,
+                    inputs=[self.actions, i * self.action_dim],
+                    outputs=[state.body_f],
+                    device=self.device,
+                )
 
-        for i in range(self.episode_frames):
-            # simulate
-            next_state = self.simulate(state, actions, action_index=i * self.action_dim, requires_grad=requires_grad)
+                state = self.integrator.simulate(self.model, state, next_state, self.sim_dt, requires_grad=True)
+
+            self.states[i + 1] = state
 
             # save state
             wp.launch(
-                save_state, dim=1, inputs=[next_state.body_q, i * self.state_dim], outputs=[states], device=self.device
+                save_state,
+                dim=1,
+                inputs=[self.states[i + 1].body_q, i],
+                outputs=[self.last_traj],
+                device=self.device,
             )
 
-            # update state
-            state = next_state
-
-            if render:
-                self._render(state)
-
         # compute loss
-        if loss is None:
-            loss = wp.zeros(1, dtype=wp.float32, device=self.device, requires_grad=requires_grad)
         wp.launch(
             loss_l2,
-            dim=self.state_dim * self.episode_frames,
-            inputs=[states, self.ref_traj],
-            outputs=[loss],
+            dim=self.last_traj.shape,
+            inputs=[self.last_traj, self.ref_traj],
+            outputs=[self.loss],
             device=self.device,
         )
 
-        return states
+    def update(self):
+        """Runs a single optimizer iteration"""
+        self.loss.zero_()
+        tape = wp.Tape()
+        with tape:
+            self.compute_loss()
 
-    def optimize(self, num_iter=100, lr=0.01):
-        # initial guess
-        actions = wp.array(
-            np.zeros(self.episode_frames * self.action_dim) * 100.0,
-            dtype=wp.float32,
-            device=self.device,
-            requires_grad=True,
-        )
+        if self.verbose and (self.iter + 1) % 10 == 0:
+            print(f"Iter {self.iter+1} Loss: {self.loss.numpy()[0]:.3f}")
 
-        optimizer = Adam([actions], lr=lr)
-        loss = wp.zeros(1, dtype=wp.float32, device=self.device, requires_grad=True)
+        tape.backward(loss=self.loss)
 
-        # optimize
-        for i in range(1, num_iter + 1):
-            loss.zero_()
-            tape = wp.Tape()
-            with tape:
-                self.forward(actions, requires_grad=True, loss=loss)
+        # print("action grad", self.actions.grad.numpy())
+        assert not np.isnan(self.actions.grad.numpy()).any(), "NaN in gradient"
 
-            if i % 10 == 0:
-                print(f"iter {i}/{num_iter} loss: {loss.numpy()[0]:.3f}")
+        self.optimizer.step([self.actions.grad])
+        tape.zero()
+        self.iter = self.iter + 1
+
+    def render(self):
+        for i in range(self.episode_frames):
+            self.renderer.begin_frame(self.render_time)
+            self.renderer.render(self.states[i + 1])
+            self.renderer.end_frame()
+            self.render_time += self.frame_dt
 
-            tape.backward(loss=loss)
-            # print("action grad", actions.grad.numpy())
-            assert not np.isnan(actions.grad.numpy()).any(), "NaN in gradient"
-            optimizer.step([actions.grad])
-            tape.zero()
 
-        return actions
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
 
+    stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_trajopt.usd")
+    example = Example(stage_path, device=wp.get_preferred_device(), verbose=True)
 
-np.set_printoptions(precision=4, linewidth=2000, suppress=True)
+    # Optimize
+    num_iter = 250
 
-sim = Environment(device=wp.get_preferred_device())
+    for i in range(num_iter):
+        example.update()
 
-best_actions = sim.optimize(num_iter=250, lr=1e2)
-# render
-opt_traj = sim.forward(best_actions, render=True)
+        # Render every 25 iters
+        if i % 25 == 0:
+            example.render()
 
-np_states = opt_traj.numpy().reshape((-1, 2))
-np_ref = sim.ref_traj.numpy().reshape((-1, 2))
-plt.plot(np_ref[:, 0], np_ref[:, 1], label="reference")
-plt.plot(np_states[:, 0], np_states[:, 1], label="optimized")
-plt.grid()
-plt.legend()
-plt.axis("equal")
-plt.show()
+    example.renderer.save()
+
+    np_states = example.last_traj.numpy()
+    np_ref = example.ref_traj.numpy()
+    plt.plot(np_ref[:, 0], np_ref[:, 1], label="Reference Trajectory")
+    plt.plot(np_states[:, 0], np_states[:, 1], label="Optimized Trajectory")
+    plt.grid()
+    plt.legend()
+    plt.axis("equal")
+    plt.show()
+#

examples/example_sph.py

@@ -6,26 +6,27 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 ###########################################################################
-# Example Smooth Particle Hydrodynamics
+# Example Smoothed Particle Hydrodynamics
+#
+# Shows how to implement a SPH fluid simulation.
 #
-# Shows how to implement a SPH
 # Neighbors are found using the wp.HashGrid class, and
 # wp.hash_grid_query(), wp.hash_grid_query_next() kernel methods.
 #
 # Reference Publication
-# Müller, Matthias, David Charypar, and Markus H. Gross.
+# Matthias Müller, David Charypar, and Markus H. Gross.
 # "Particle-based fluid simulation for interactive applications."
 # Symposium on Computer animation. Vol. 2. 2003.
 #
 ###########################################################################
 
+import os
+
 import numpy as np
 
 import warp as wp
 import warp.render
 
-import os
-
 wp.init()
 
 
@@ -314,7 +315,7 @@ class Example:
 
     def update(self):
         with wp.ScopedTimer("simulate", active=True):
-            for s in range(self.sim_step_to_frame_ratio):
+            for _ in range(self.sim_step_to_frame_ratio):
                 with wp.ScopedTimer("grid build", active=False):
                     # build grid
                     self.grid.build(self.x, self.smoothing_length)
@@ -359,6 +360,8 @@ class Example:
                     # drift
                     wp.launch(kernel=drift, dim=self.n, inputs=[self.x, self.v, self.dt])
 
+            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
@@ -367,8 +370,6 @@ class Example:
             self.renderer.render_points(points=self.x.numpy(), radius=self.smoothing_length, name="points")
             self.renderer.end_frame()
 
-        self.sim_time += self.frame_dt
-
 
 if __name__ == "__main__":
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sph.usd")

examples/example_wave.py

@@ -13,8 +13,8 @@
 #
 ###########################################################################
 
-import os
 import math
+import os
 
 import warp as wp
 import warp.render
@@ -111,9 +111,9 @@ def grid_update(heights: wp.array(dtype=float), vertices: wp.array(dtype=wp.vec3
 
 
 class Example:
-    def __init__(self, stage):
-        self.sim_width = 128
-        self.sim_height = 128
+    def __init__(self, stage, resx, resy):
+        self.sim_width = resx
+        self.sim_height = resy
 
         self.sim_fps = 60.0
         self.sim_substeps = 16
@@ -163,14 +163,14 @@ class Example:
         self.sim_grid1 = wp.zeros(self.sim_width * self.sim_height, dtype=float)
         self.sim_verts = wp.array(vertices, dtype=wp.vec3)
 
-        # create surface displacment around a point
+        # create surface displacement around a point
         self.cx = self.sim_width / 2 + math.sin(self.sim_time) * self.sim_width / 3
         self.cy = self.sim_height / 2 + math.cos(self.sim_time) * self.sim_height / 3
 
     def update(self):
         with wp.ScopedTimer("simulate", active=True):
             for s in range(self.sim_substeps):
-                # create surface displacment around a point
+                # create surface displacement around a point
                 self.cx = self.sim_width / 2 + math.sin(self.sim_time) * self.sim_width / 3
                 self.cy = self.sim_height / 2 + math.cos(self.sim_time) * self.sim_height / 3
 
@@ -235,7 +235,7 @@ class Example:
 if __name__ == "__main__":
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_wave.usd")
 
-    example = Example(stage_path)
+    example = Example(stage_path, resx=128, resy=128)
 
     for i in range(example.sim_frames):
         example.update()

examples/fem/bsr_utils.py

@@ -6,17 +6,14 @@ import warp.types
 from warp.sparse import BsrMatrix, bsr_zeros, bsr_get_diag, bsr_mv
 from warp.utils import array_inner
 
-try:
-    from scipy.sparse import csr_array, bsr_array
-except ImportError:
-    # WAR for older scipy
+
+def bsr_to_scipy(matrix: BsrMatrix) -> "scipy.sparse.bsr_array":
     try:
-        from scipy.sparse import csr_matrix as csr_array, bsr_matrix as bsr_array
+        from scipy.sparse import csr_array, bsr_array
     except ImportError:
-        pass
-
+        # WAR for older scipy
+        from scipy.sparse import csr_matrix as csr_array, bsr_matrix as bsr_array
 
-def bsr_to_scipy(matrix: BsrMatrix) -> "bsr_array":
     if matrix.block_shape == (1, 1):
         return csr_array(
             (
@@ -37,7 +34,13 @@ def bsr_to_scipy(matrix: BsrMatrix) -> "bsr_array":
     )
 
 
-def scipy_to_bsr(sp: Union["bsr_array", "csr_array"], device=None, dtype=None) -> BsrMatrix:
+def scipy_to_bsr(sp: Union["scipy.sparse.bsr_array", "scipy.sparse.csr_array"], device=None, dtype=None) -> BsrMatrix:
+    try:
+        from scipy.sparse import csr_array
+    except ImportError:
+        # WAR for older scipy
+        from scipy.sparse import csr_matrix as csr_array
+
     if dtype is None:
         dtype = warp.types.np_dtype_to_warp_type[sp.dtype]
 
@@ -68,9 +71,11 @@ def _bsr_cg_kernel_1(
     Ap: wp.array(dtype=Any),
 ):
     i = wp.tid()
-    alpha = rs_old[0] / p_Ap[0]
-    x[i] = x[i] + alpha * p[i]
-    r[i] = r[i] - alpha * Ap[i]
+
+    if p_Ap[0] != 0.0:
+        alpha = rs_old[0] / p_Ap[0]
+        x[i] = x[i] + alpha * p[i]
+        r[i] = r[i] - alpha * Ap[i]
 
 
 @wp.kernel
@@ -133,6 +138,7 @@ def bsr_cg(
     use_diag_precond=True,
     mv_routine=bsr_mv,
     device=None,
+    quiet=False,
 ) -> Tuple[float, int]:
     """Solves the linear system A x = b using the Conjugate Gradient method, optionally with diagonal preconditioning
 
@@ -217,7 +223,8 @@ def bsr_cg(
 
             if ((i + 1) % check_every) == 0:
                 err = rz_new.numpy()[0]
-                print(f"At iteration {i} error = \t {err}  \t tol: {tol_sq}")
+                if not quiet:
+                    print(f"At iteration {i} error = \t {err}  \t tol: {tol_sq}")
                 if err <= tol_sq:
                     end_iter = i
                     break
@@ -231,17 +238,23 @@ def bsr_cg(
 
         err = rz_old.numpy()[0]
 
-    print(f"Terminated after {end_iter} iterations with error = \t {err}")
+    if not quiet:
+        print(f"Terminated after {end_iter} iterations with error = \t {err}")
     return err, end_iter
 
 
 def invert_diagonal_bsr_mass_matrix(A: BsrMatrix):
     """Inverts each block of a block-diagonal mass matrix"""
 
-    wp.launch(kernel=_block_diagonal_mass_invert, dim=A.nrow, inputs=[A.values], device=A.values.device)
+    scale = A.scalar_type(A.block_shape[0])
+    values = A.values
+    if not wp.types.type_is_matrix(values.dtype):
+        values = values.view(dtype=wp.mat(shape=(1, 1), dtype=A.scalar_type))
+
+    wp.launch(kernel=_block_diagonal_mass_invert, dim=A.nrow, inputs=[values, scale], device=values.device)
 
 
 @wp.kernel
-def _block_diagonal_mass_invert(values: wp.array(dtype=Any)):
+def _block_diagonal_mass_invert(values: wp.array(dtype=Any), scale: Any):
     i = wp.tid()
-    values[i] = values[i] / wp.ddot(values[i], values[i])
+    values[i] = scale * values[i] / wp.ddot(values[i], values[i])