warp-lang 0.15.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/examples/core/example_dem.py

@@ -148,12 +148,12 @@ class Example:
         if stage is not None:
             self.renderer = wp.render.UsdRenderer(stage)
             self.renderer.render_ground()
-        
+
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
-            wp.capture_begin()
-            self.simulate()
-            self.graph = wp.capture_end()
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
 
     def simulate(self):
         for _ in range(self.sim_substeps):
@@ -182,12 +182,12 @@ class Example:
         with wp.ScopedTimer("step", active=True):
             with wp.ScopedTimer("grid build", active=False):
                 self.grid.build(self.x, self.grid_cell_size)
-            
+
             if self.use_graph:
                 wp.capture_launch(self.graph)
             else:
                 self.simulate()
-            
+
             self.sim_time += self.frame_dt
 
     def render(self):

warp/examples/core/example_fluid.py

@@ -187,9 +187,9 @@ class Example:
         # capture pressure solve as a CUDA graph
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
-            wp.capture_begin()
-            self.pressure_iterations()
-            self.graph = wp.capture_end()
+            with wp.ScopedCapture() as capture:
+                self.pressure_iterations()
+            self.graph = capture.graph
 
     def step(self):
         with wp.ScopedTimer("step"):

warp/examples/core/example_graph_capture.py

@@ -10,7 +10,7 @@
 #
 # Shows how to implement CUDA graph capture using Python's try-finally
 # pattern. The finally block ensures wp.capture_end() gets called, even
-# if an exception occurs during capture, which would otherwise 
+# if an exception occurs during capture, which would otherwise
 # trap the stream in a capturing state.
 #
 ###########################################################################
@@ -72,12 +72,9 @@ class Example:
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
             # record launches
-            wp.capture_begin()
-            try:
+            with wp.ScopedCapture() as capture:
                 self.fbm()
-            finally:
-                # end capture and return a graph object
-                self.graph = wp.capture_end()
+            self.graph = capture.graph
 
     def fbm(self):
         for _ in range(16):

warp/examples/optim/example_bounce.py

@@ -76,7 +76,8 @@ class Example:
 
         self.profile = profile
 
-        self.model = builder.finalize()
+        # use `requires_grad=True` to create a model for differentiable simulation
+        self.model = builder.finalize(requires_grad=True)
         self.model.ground = True
 
         self.model.soft_contact_ke = ke
@@ -94,7 +95,7 @@ class Example:
         # allocate sim states for trajectory
         self.states = []
         for i in range(self.sim_steps + 1):
-            self.states.append(self.model.state(requires_grad=True))
+            self.states.append(self.model.state())
 
         # one-shot contact creation (valid if we're doing simple collision against a constant normal plane)
         wp.sim.collide(self.model, self.states[0])
@@ -106,12 +107,12 @@ class Example:
         # capture forward/backward passes
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
-            wp.capture_begin()
-            self.tape = wp.Tape()
-            with self.tape:
-                self.forward()
-            self.tape.backward(self.loss)
-            self.graph = wp.capture_end()
+            with wp.ScopedCapture() as capture:
+                self.tape = wp.Tape()
+                with self.tape:
+                    self.forward()
+                self.tape.backward(self.loss)
+            self.graph = capture.graph
 
     def forward(self):
         # run control loop

warp/examples/optim/example_cloth_throw.py

@@ -118,12 +118,12 @@ class Example:
         # capture forward/backward passes
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
-            wp.capture_begin()
-            self.tape = wp.Tape()
-            with self.tape:
-                self.forward()
-            self.tape.backward(self.loss)
-            self.graph = wp.capture_end()
+            with wp.ScopedCapture() as capture:
+                self.tape = wp.Tape()
+                with self.tape:
+                    self.forward()
+                self.tape.backward(self.loss)
+            self.graph = capture.graph
 
     def forward(self):
         # run control loop
@@ -137,7 +137,7 @@ class Example:
         wp.launch(
             com_kernel,
             dim=self.model.particle_count,
-            inputs=[self.states[-1].particle_q, self.model.particle_count, self.com]
+            inputs=[self.states[-1].particle_q, self.model.particle_count, self.com],
         )
         wp.launch(loss_kernel, dim=1, inputs=[self.com, self.target, self.loss])
 
@@ -178,7 +178,11 @@ class Example:
                 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", color=(1.0, 0.0, 0.0),
+                    pos=self.target,
+                    rot=wp.quat_identity(),
+                    extents=(0.1, 0.1, 0.1),
+                    name="target",
+                    color=(1.0, 0.0, 0.0),
                 )
                 self.renderer.render_line_strip(
                     vertices=traj_verts,

warp/examples/optim/example_diffray.py

@@ -411,12 +411,12 @@ class Example:
         # capture graph
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
-            wp.capture_begin()
-            self.tape = wp.Tape()
-            with self.tape:
-                self.forward()
-            self.tape.backward(self.loss)
-            self.graph = wp.capture_end()
+            with wp.ScopedCapture() as capture:
+                self.tape = wp.Tape()
+                with self.tape:
+                    self.forward()
+                self.tape.backward(self.loss)
+            self.graph = capture.graph
 
         self.optimizer = SGD(
             [self.render_mesh.rot],
@@ -440,23 +440,21 @@ class Example:
                 self.rays,
                 self.lights,
                 self.render_mode,
-            ]
+            ],
         )
 
         # downsample
         wp.launch(
             kernel=downsample_kernel,
             dim=self.num_pixels,
-            inputs=[self.rays, self.pixels, self.rays_width, pow(2, self.num_samples)]
+            inputs=[self.rays, self.pixels, self.rays_width, pow(2, self.num_samples)],
         )
 
     def forward(self):
         self.ray_cast()
 
         # compute pixel loss
-        wp.launch(
-            loss_kernel, dim=self.num_pixels, inputs=[self.pixels, self.target_pixels, self.loss]
-        )
+        wp.launch(loss_kernel, dim=self.num_pixels, inputs=[self.pixels, self.target_pixels, self.loss])
 
     def step(self):
         if self.use_graph:
@@ -529,7 +527,7 @@ if __name__ == "__main__":
 
         if i % example.period == 0:
             example.render()
-    
+
     final_image = example.get_image()
     img.imsave(output_dir + "/example_diffray_final_image.png", final_image)
 

warp/examples/optim/example_drone.py

@@ -9,7 +9,8 @@
 # Example Drone
 #
 # A drone and its 4 propellers is simulated with the goal of reaching
-# different targets by finding the best trajectory at each point in time.
+# different targets via model-predictive control (MPC) that continuously
+# optimizes the control trajectory.
 #
 ###########################################################################
 
@@ -415,15 +416,18 @@ class Drone:
 
         # Initialize the required simulation states.
         if requires_grad:
-            self.states = tuple(self.model.state(requires_grad=True) for _ in range(state_count + 1))
+            self.states = tuple(self.model.state() for _ in range(state_count + 1))
+            self.controls = tuple(self.model.control() for _ in range(state_count))
         else:
             # When only running a forward simulation, we don't need to store
             # the history of the states at each step, instead we use double
             # buffering to represent the previous and next states.
             self.states = [self.model.state(), self.model.state()]
+            self.controls = (self.model.control(),)
 
-        for state in self.states:
-            state.prop_controls = wp.zeros(len(self.props), dtype=float, requires_grad=requires_grad)
+        # create array for the propeller controls
+        for control in self.controls:
+            control.prop_controls = wp.zeros(len(self.props), dtype=float, requires_grad=requires_grad)
 
         # Define the trajectories as arrays of control points.
         # The point data has an additional item to support linear interpolation.
@@ -446,6 +450,10 @@ class Drone:
     def next_state(self) -> wp.sim.State:
         return self.states[self.sim_tick + 1 if self.requires_grad else 1]
 
+    @property
+    def control(self) -> wp.sim.Control:
+        return self.controls[min(len(self.controls) - 1, self.sim_tick) if self.requires_grad else 0]
+
 
 class Example:
     def __init__(
@@ -602,7 +610,7 @@ class Example:
                 drone.sim_tick / (self.sim_substep_count * self.control_point_step),
                 self.control_dim,
             ),
-            outputs=(drone.state.prop_controls,),
+            outputs=(drone.control.prop_controls,),
         )
 
         wp.sim.collide(drone.model, drone.state)
@@ -612,7 +620,7 @@ class Example:
             dim=len(drone.props),
             inputs=(
                 drone.props,
-                drone.state.prop_controls,
+                drone.control.prop_controls,
                 drone.state.body_q,
                 drone.model.body_com,
             ),
@@ -624,6 +632,7 @@ class Example:
             drone.state,
             drone.next_state,
             self.sim_dt,
+            drone.control,
         )
 
         drone.sim_tick += 1
@@ -657,7 +666,7 @@ class Example:
                     self.rollouts.state.body_q,
                     self.rollouts.state.body_qd,
                     self.targets[self.target_idx],
-                    self.rollouts.state.prop_controls,
+                    self.rollouts.control.prop_controls,
                     i,
                     self.rollout_step_count,
                     1e3,
@@ -711,15 +720,13 @@ class Example:
             self.optim_graph = None
 
         if self.use_cuda_graph and self.optim_graph is None:
-            wp.capture_begin()
-            try:
+            with wp.ScopedCapture() as capture:
                 self.tape = wp.Tape()
                 with self.tape:
                     self.forward()
                 self.rollout_costs.grad.fill_(1.0)
                 self.tape.backward()
-            finally:
-                self.optim_graph = wp.capture_end()
+            self.optim_graph = capture.graph
 
         # Sample control waypoints around the nominal trajectory.
         self.seed.zero_()
@@ -817,17 +824,27 @@ class Example:
 
 
 if __name__ == "__main__":
-    stage_path = os.path.join(os.path.dirname(__file__), "example_drone.usd")
-    drone_path = os.path.join(os.path.dirname(__file__), "..", "assets", "crazyflie.usd")
+    this_dir = os.path.realpath(os.path.dirname(__file__))
+    stage_path = os.path.join(this_dir, "example_drone.usd")
+    drone_path = os.path.join(this_dir, "..", "assets", "crazyflie.usd")
 
     example = Example(stage_path, drone_path, verbose=True)
     for i in range(example.frame_count):
+        if i > 0 and i % int((example.frame_count / len(example.targets))) == 0:
+            example.target_idx += 1
+            if example.verbose:
+                print(f"Choosing new flight target: {example.target_idx + 1}")
+
+            # Force recapturing the CUDA graph for the optimisation pass
+            # by invalidating it.
+            example.optim_graph = None
+
         example.step()
         example.render()
 
         if example.verbose:
             loss = np.min(example.rollout_costs.numpy())
-            print(f"[{example.frame:3d}] loss={loss:.8f}")
+            print(f"[{i+1:3d}/{example.frame_count}] loss={loss:.8f}")
 
     if example.renderer is not None:
         example.renderer.save()

warp/examples/optim/example_spring_cage.py

@@ -113,9 +113,7 @@ class Example:
         self.integrator = warp.sim.SemiImplicitIntegrator()
 
         # Initialize a state for each simulation step.
-        self.states = tuple(
-            self.model.state(requires_grad=True) for _ in range(self.frame_count * self.sim_substep_count + 1)
-        )
+        self.states = tuple(self.model.state() for _ in range(self.frame_count * self.sim_substep_count + 1))
 
         # Initialize a loss value that will represent the distance of the main
         # particle to the target position. It needs to be defined as an array
@@ -126,7 +124,7 @@ class Example:
             import warp.sim.render
 
             # Helper to render the physics scene as a USD file.
-            self.renderer = warp.sim.render.SimRenderer(self.model, stage, fps=self.fps)
+            self.renderer = warp.sim.render.SimRenderer(self.model, stage, fps=self.fps, scaling=10.0)
 
             # Allows rendering one simulation to USD every N training iterations.
             self.render_iteration_steps = 2
@@ -140,12 +138,12 @@ class Example:
         if self.use_graph:
             # Capture all the kernel launches into a CUDA graph so that they can
             # all be run in a single graph launch, which helps with performance.
-            wp.capture_begin()
-            self.tape = wp.Tape()
-            with self.tape:
-                self.forward()
-            self.tape.backward(loss=self.loss)
-            self.graph = wp.capture_end()
+            with wp.ScopedCapture() as capture:
+                self.tape = wp.Tape()
+                with self.tape:
+                    self.forward()
+                self.tape.backward(loss=self.loss)
+            self.graph = capture.graph
 
         self.verbose = verbose
 
@@ -220,7 +218,7 @@ if __name__ == "__main__":
     stage_path = os.path.join(os.path.dirname(__file__), "example_spring_cage.usd")
 
     example = Example(stage_path, verbose=True)
-    
+
     for iteration in range(example.train_iters):
         example.step()
 
@@ -229,10 +227,7 @@ if __name__ == "__main__":
         if example.verbose:
             print(f"[{iteration:3d}] loss={loss:.8f}")
 
-        if (
-            iteration == example.train_iters - 1
-            or iteration % example.render_iteration_steps == 0
-        ):
+        if iteration == example.train_iters - 1 or iteration % example.render_iteration_steps == 0:
             example.render()
 
     if example.renderer:

warp/examples/optim/example_trajectory.py

@@ -100,9 +100,7 @@ class Example:
 
         # initial guess
         self.actions = wp.array(
-            np.zeros(self.episode_frames * self.action_dim) * 100.0,
-            dtype=wp.float32,
-            requires_grad=True
+            np.zeros(self.episode_frames * self.action_dim) * 100.0, dtype=wp.float32, requires_grad=True
         )
 
         self.optimizer = Adam([self.actions], lr=1e2)
@@ -115,7 +113,7 @@ class Example:
         # allocate sim states for trajectory
         self.states = []
         for _ in range(self.episode_frames + 1):
-            self.states.append(self.model.state(requires_grad=True))
+            self.states.append(self.model.state())
 
     def forward(self):
         """
@@ -134,32 +132,17 @@ class Example:
                 wp.sim.collide(self.model, state)
 
                 # 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]
-                )
+                wp.launch(apply_torque, 1, inputs=[self.actions, i * self.action_dim], outputs=[state.body_f])
 
-                state = self.integrator.simulate(self.model, state, next_state, self.sim_dt, requires_grad=True)
+                state = self.integrator.simulate(self.model, state, next_state, self.sim_dt)
 
             self.states[i + 1] = state
 
             # save state
-            wp.launch(
-                save_state,
-                dim=1,
-                inputs=[self.states[i + 1].body_q, i],
-                outputs=[self.last_traj]
-            )
+            wp.launch(save_state, dim=1, inputs=[self.states[i + 1].body_q, i], outputs=[self.last_traj])
 
         # compute loss
-        wp.launch(
-            loss_l2,
-            dim=self.last_traj.shape,
-            inputs=[self.last_traj, self.ref_traj],
-            outputs=[self.loss]
-        )
+        wp.launch(loss_l2, dim=self.last_traj.shape, inputs=[self.last_traj, self.ref_traj], outputs=[self.loss])
 
     def step(self):
         """Runs a single optimizer iteration"""
@@ -168,7 +151,7 @@ class Example:
         with tape:
             self.forward()
         tape.backward(loss=self.loss)
-        
+
         if self.verbose and (self.iter + 1) % 10 == 0:
             print(f"Iter {self.iter+1} Loss: {self.loss.numpy()[0]:.3f}")
 

warp/examples/sim/example_cartpole.py

@@ -43,10 +43,6 @@ class Example:
             armature=0.1,
             stiffness=0.0,
             damping=0.0,
-            shape_ke=1.0e4,
-            shape_kd=1.0e2,
-            shape_kf=1.0e2,
-            shape_mu=1.0,
             limit_ke=1.0e4,
             limit_kd=1.0e1,
             enable_self_collisions=False,
@@ -71,8 +67,6 @@ class Example:
             # joint initial positions
             builder.joint_q[-3:] = [0.0, 0.3, 0.0]
 
-            builder.joint_target[:3] = [0.0, 0.0, 0.0]
-
         # finalize model
         self.model = builder.finalize()
         self.model.ground = False
@@ -94,9 +88,9 @@ class Example:
 
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
-            wp.capture_begin()
-            self.simulate()
-            self.graph = wp.capture_end()
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
 
     def simulate(self):
         for _ in range(self.sim_substeps):

warp/examples/sim/example_cloth.py

@@ -105,28 +105,28 @@ class Example:
             kf=1.0e1,
         )
 
-        self.model = builder.finalize()
-        self.model.ground = True
-        self.model.soft_contact_ke = 1.0e4
-        self.model.soft_contact_kd = 1.0e2
-
         if self.integrator_type == IntegratorType.EULER:
             self.integrator = wp.sim.SemiImplicitIntegrator()
         else:
             self.integrator = wp.sim.XPBDIntegrator(iterations=1)
 
+        self.model = builder.finalize()
+        self.model.ground = True
+        self.model.soft_contact_ke = 1.0e4
+        self.model.soft_contact_kd = 1.0e2
+
         self.state_0 = self.model.state()
         self.state_1 = self.model.state()
 
         self.renderer = None
         if stage:
             self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=40.0)
-        
+
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
-            wp.capture_begin()
-            self.simulate()
-            self.graph = wp.capture_end()
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
 
     def simulate(self):
         wp.sim.collide(self.model, self.state_0)
@@ -150,7 +150,7 @@ class Example:
     def render(self):
         if self.renderer is None:
             return
-        
+
         with wp.ScopedTimer("render", active=True):
             self.renderer.begin_frame(self.sim_time)
             self.renderer.render(self.state_0)

warp/examples/sim/example_granular.py

@@ -68,9 +68,9 @@ class Example:
 
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
-            wp.capture_begin()
-            self.simulate()
-            self.graph = wp.capture_end()
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
 
     def simulate(self):
         for _ in range(self.sim_substeps):

warp/examples/sim/example_granular_collision_sdf.py

@@ -119,9 +119,11 @@ class Example:
         if stage:
             self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=20.0)
 
-        wp.capture_begin()
-        self.simulate()
-        self.graph = wp.capture_end()            
+        self.use_graph = wp.get_device().is_cuda
+        if self.use_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
 
     def simulate(self):
         for _ in range(self.sim_substeps):
@@ -135,7 +137,10 @@ class Example:
     def step(self):
         with wp.ScopedTimer("step", active=True):
             self.model.particle_grid.build(self.state_0.particle_q, self.radius * 2.0)
-            wp.capture_launch(self.graph)
+            if self.use_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.simulate()
 
         self.sim_time += self.frame_dt
 

warp/examples/sim/example_jacobian_ik.py

@@ -60,16 +60,6 @@ class Example:
             articulation_builder,
             xform=wp.transform_identity(),
             floating=False,
-            density=0,
-            armature=0.1,
-            stiffness=0.0,
-            damping=0.0,
-            shape_ke=1.0e4,
-            shape_kd=1.0e2,
-            shape_kf=1.0e2,
-            shape_mu=1.0,
-            limit_ke=1.0e4,
-            limit_kd=1.0e1,
         )
 
         builder = wp.sim.ModelBuilder()

warp/examples/sim/example_particle_chain.py

@@ -62,9 +62,9 @@ class Example:
 
         self.use_graph = wp.get_device().is_cuda
         if self.use_graph:
-            wp.capture_begin()
-            self.simulate()
-            self.graph = wp.capture_end()     
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
 
     def simulate(self):
         for _ in range(self.sim_substeps):
@@ -80,7 +80,7 @@ class Example:
             if self.use_graph:
                 wp.capture_launch(self.graph)
             else:
-                self.simulate()                    
+                self.simulate()
         self.sim_time += self.frame_dt
 
     def render(self):