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

warp/examples/optim/example_inverse_kinematics.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# Example Sim Rigid Kinematics
+#
+# Tests rigid body forward and backwards kinematics through the
+# wp.sim.eval_ik() and wp.sim.eval_fk() methods.
+#
+###########################################################################
+
+import os
+
+import numpy as np
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+wp.init()
+
+TARGET = wp.constant(wp.vec3(2.0, 1.0, 0.0))
+
+
+@wp.kernel
+def compute_loss(body_q: wp.array(dtype=wp.transform), body_index: int, loss: wp.array(dtype=float)):
+    x = wp.transform_get_translation(body_q[body_index])
+
+    delta = x - TARGET
+    loss[0] = wp.dot(delta, delta)
+
+
+@wp.kernel
+def step_kernel(x: wp.array(dtype=float), grad: wp.array(dtype=float), alpha: float):
+    tid = wp.tid()
+
+    # gradient descent step
+    x[tid] = x[tid] - grad[tid] * alpha
+
+
+class Example:
+    def __init__(self, stage, device=None, verbose=False):
+        self.verbose = verbose
+        if device is None:
+            self.device = wp.get_device()
+        else:
+            self.device = device
+
+        self.frame_dt = 1.0 / 60.0
+        self.render_time = 0.0
+
+        builder = wp.sim.ModelBuilder()
+        builder.add_articulation()
+
+        chain_length = 4
+        chain_width = 1.0
+
+        for i in range(chain_length):
+            if i == 0:
+                parent = -1
+                parent_joint_xform = wp.transform([0.0, 0.0, 0.0], wp.quat_identity())
+            else:
+                parent = builder.joint_count - 1
+                parent_joint_xform = wp.transform([chain_width, 0.0, 0.0], wp.quat_identity())
+
+            # create body
+            b = builder.add_body(origin=wp.transform([i, 0.0, 0.0], wp.quat_identity()), armature=0.1)
+
+            builder.add_joint_revolute(
+                parent=parent,
+                child=b,
+                axis=(0.0, 0.0, 1.0),
+                parent_xform=parent_joint_xform,
+                child_xform=wp.transform_identity(),
+                limit_lower=-np.deg2rad(60.0),
+                limit_upper=np.deg2rad(60.0),
+                target_ke=0.0,
+                target_kd=0.0,
+                limit_ke=30.0,
+                limit_kd=30.0,
+            )
+
+            if i == chain_length - 1:
+                # create end effector
+                builder.add_shape_sphere(pos=wp.vec3(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
+
+            else:
+                # create shape
+                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
+        self.model = builder.finalize(self.device)
+        self.model.ground = False
+
+        self.state = self.model.state()
+
+        self.renderer = None
+        if stage:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=50.0)
+
+        # optimization variables
+        self.loss = wp.zeros(1, dtype=float, device=self.device)
+
+        self.model.joint_q.requires_grad = True
+        self.state.body_q.requires_grad = True
+        self.loss.requires_grad = True
+
+        self.train_rate = 0.01
+
+    def forward(self):
+        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,
+        )
+
+    def step(self):
+        tape = wp.Tape()
+        with tape:
+            self.forward()
+        tape.backward(loss=self.loss)
+
+        if self.verbose:
+            print(f"loss: {self.loss}")
+            print(f"joint_grad: {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):
+        if self.renderer is None:
+            return
+
+        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, color=(1.0, 0.0, 0.0))
+        self.renderer.end_frame()
+        self.render_time += self.frame_dt
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_inverse_kinematics.usd")
+    example = Example(stage_path, device=wp.get_preferred_device(), verbose=True)
+
+    train_iters = 512
+
+    for _ in range(train_iters):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()

warp/examples/optim/example_inverse_kinematics_torch.py

@@ -0,0 +1,170 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# Example Sim Rigid Kinematics
+#
+# Tests rigid body forward and backwards kinematics through the
+# wp.sim.eval_ik() and wp.sim.eval_fk() methods. Shows how to connect
+# gradients from Warp to PyTorch, through custom autograd nodes.
+#
+###########################################################################
+
+import os
+
+import numpy as np
+import torch
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+wp.init()
+
+
+class ForwardKinematics(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, joint_q, joint_qd, model):
+
+        ctx.tape = wp.Tape()
+        ctx.model = model
+        ctx.joint_q = wp.from_torch(joint_q)
+        ctx.joint_qd = wp.from_torch(joint_qd)
+
+        # allocate output
+        ctx.state = model.state()
+
+        with ctx.tape:
+            wp.sim.eval_fk(model, ctx.joint_q, ctx.joint_qd, None, ctx.state)
+
+        return (wp.to_torch(ctx.state.body_q), wp.to_torch(ctx.state.body_qd))
+
+    @staticmethod
+    def backward(ctx, adj_body_q, adj_body_qd):
+
+        # map incoming Torch grads to our output variables
+        ctx.state.body_q.grad = wp.from_torch(adj_body_q, dtype=wp.transform)
+        ctx.state.body_qd.grad = wp.from_torch(adj_body_qd, dtype=wp.spatial_vector)
+
+        ctx.tape.backward()
+
+        # return adjoint w.r.t. inputs
+        return (wp.to_torch(ctx.tape.gradients[ctx.joint_q]), wp.to_torch(ctx.tape.gradients[ctx.joint_qd]), 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()
+
+        chain_length = 4
+        chain_width = 1.0
+
+        for i in range(chain_length):
+            if i == 0:
+                parent = -1
+                parent_joint_xform = wp.transform([0.0, 0.0, 0.0], wp.quat_identity())
+            else:
+                parent = builder.joint_count - 1
+                parent_joint_xform = wp.transform([chain_width, 0.0, 0.0], wp.quat_identity())
+
+            # create body
+            b = builder.add_body(origin=wp.transform([i, 0.0, 0.0], wp.quat_identity()), armature=0.1)
+
+            builder.add_joint_revolute(
+                parent=parent,
+                child=b,
+                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),
+                limit_upper=np.deg2rad(60.0),
+                target_ke=0.0,
+                target_kd=0.0,
+                limit_ke=30.0,
+                limit_kd=30.0,
+            )
+
+            if i == chain_length - 1:
+                # create end effector
+                builder.add_shape_sphere(pos=wp.vec3(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
+
+            else:
+                # create shape
+                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
+        self.model = builder.finalize(device)
+        self.model.ground = False
+
+        self.torch_device = wp.device_to_torch(self.model.device)
+
+        self.renderer = None
+        if stage:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=50.0)
+
+        self.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
+        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
+
+    def forward(self):
+        (self.body_q, self.body_qd) = ForwardKinematics.apply(self.joint_q, self.joint_qd, self.model)
+        self.loss = torch.norm(self.body_q[self.model.body_count - 1][0:3] - self.target) ** 2.0
+
+    def step(self):
+        self.forward()
+        self.loss.backward()
+
+        if self.verbose:
+            print(f"loss: {self.loss}")
+            print(f"loss: {self.joint_q.grad}")
+
+        with torch.no_grad():
+            self.joint_q -= self.joint_q.grad * self.train_rate
+            self.joint_q.grad.zero_()
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        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)
+
+        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, color=(1.0, 0.0, 0.0))
+        self.renderer.end_frame()
+        self.render_time += self.frame_dt
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_inverse_kinematics_torch.usd")
+    example = Example(stage_path, device=wp.get_preferred_device(), verbose=True)
+
+    train_iters = 512
+
+    for _ in range(train_iters):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()

warp/examples/optim/example_spring_cage.py

@@ -0,0 +1,234 @@
+# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# Example Diff Sim Spring Cage
+#
+# A single particle is attached with springs to each point of a cage.
+# The objective is to optimize the rest length of the springs in order
+# for the particle to be pulled towards a target position.
+#
+###########################################################################
+
+import os
+
+import numpy as np
+import warp as wp
+
+
+wp.init()
+
+
+@wp.kernel
+def compute_loss_kernel(
+    pos: wp.array(dtype=wp.vec3),
+    target_pos: wp.vec3,
+    loss: wp.array(dtype=float),
+):
+    loss[0] = wp.length_sq(pos[0] - target_pos)
+
+
+@wp.kernel(enable_backward=False)
+def apply_gradient_kernel(
+    spring_rest_lengths_grad: wp.array(dtype=float),
+    train_rate: float,
+    spring_rest_lengths: wp.array(dtype=float),
+):
+    tid = wp.tid()
+
+    spring_rest_lengths[tid] -= spring_rest_lengths_grad[tid] * train_rate
+
+
+class Example:
+    def __init__(
+        self,
+        stage=None,
+        verbose=False,
+    ):
+        import warp.sim
+
+        # Duration of the simulation, in seconds.
+        duration = 1.0
+
+        # Number of frames per second.
+        self.fps = 30.0
+
+        # Duration of a single simulation iteration in number of frames.
+        self.frame_count = int(duration * self.fps)
+
+        # Number of simulation steps to take per frame.
+        self.sim_substep_count = 1
+
+        # Delta time between each simulation substep.
+        self.sim_dt = 1.0 / (self.fps * self.sim_substep_count)
+
+        # Target position that we want the main particle to reach by optimising
+        # the rest lengths of the springs.
+        self.target_pos = (0.125, 0.25, 0.375)
+
+        # Number of training iterations.
+        self.train_iters = 100
+
+        # Factor by which the rest lengths of the springs are adjusted after each
+        # iteration, relatively to the corresponding gradients. Lower values
+        # converge more slowly but have less chances to miss the local minimum.
+        self.train_rate = 0.5
+
+        # Initialize the helper to build a physics scene.
+        builder = warp.sim.ModelBuilder()
+
+        # Define the main particle at the origin.
+        particle_mass = 1.0
+        builder.add_particle((0.0, 0.0, 0.0), (0.0, 0.0, 0.0), particle_mass)
+
+        # Define the cage made of points that will be pulling our main particle
+        # using springs.
+        # fmt: off
+        builder.add_particle((-0.7,  0.8,  0.2), (0.0, 0.0, 0.0), 0.0)
+        builder.add_particle(( 0.0,  0.2,  1.1), (0.0, 0.0, 0.0), 0.0)
+        builder.add_particle(( 0.1,  0.1, -1.2), (0.0, 0.0, 0.0), 0.0)
+        builder.add_particle(( 0.6,  0.4,  0.4), (0.0, 0.0, 0.0), 0.0)
+        builder.add_particle(( 0.7, -0.9, -0.2), (0.0, 0.0, 0.0), 0.0)
+        builder.add_particle((-0.8, -0.8,  0.1), (0.0, 0.0, 0.0), 0.0)
+        builder.add_particle((-0.9,  0.2, -0.8), (0.0, 0.0, 0.0), 0.0)
+        builder.add_particle(( 1.0,  0.4, -0.1), (0.0, 0.0, 0.0), 0.0)
+        # fmt: on
+
+        # Define the spring constraints between the main particle and the cage points.
+        spring_elastic_stiffness = 100.0
+        spring_elastic_damping = 10.0
+        for i in range(1, builder.particle_count):
+            builder.add_spring(0, i, spring_elastic_stiffness, spring_elastic_damping, 0)
+
+        # Build the model and set-up its properties.
+        self.model = builder.finalize(requires_grad=True)
+        self.model.gravity = np.array((0.0, 0.0, 0.0))
+        self.model.ground = False
+
+        # Use the Euler integrator for stepping through the simulation.
+        self.integrator = warp.sim.SemiImplicitIntegrator()
+
+        # Initialize a state for each simulation step.
+        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
+        # so that it can be written out by a kernel.
+        self.loss = wp.zeros(1, dtype=float, requires_grad=True)
+
+        if stage:
+            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, scaling=10.0)
+
+            # Allows rendering one simulation to USD every N training iterations.
+            self.render_iteration_steps = 2
+
+            # Frame number used to render the simulation iterations onto the USD file.
+            self.render_frame = 0
+        else:
+            self.renderer = None
+
+        self.use_graph = wp.get_device().is_cuda
+        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.
+            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
+
+    def forward(self):
+        for i in range(1, len(self.states)):
+            prev = self.states[i - 1]
+            curr = self.states[i]
+            prev.clear_forces()
+            self.integrator.simulate(
+                self.model,
+                prev,
+                curr,
+                self.sim_dt,
+            )
+
+        last_state = self.states[-1]
+        wp.launch(
+            compute_loss_kernel,
+            dim=1,
+            inputs=(
+                last_state.particle_q,
+                self.target_pos,
+            ),
+            outputs=(self.loss,),
+        )
+
+    def step(self):
+        if self.use_graph:
+            wp.capture_launch(self.graph)
+        else:
+            self.tape = wp.Tape()
+            with self.tape:
+                self.forward()
+            self.tape.backward(loss=self.loss)
+
+        wp.launch(
+            apply_gradient_kernel,
+            dim=self.model.spring_count,
+            inputs=(
+                self.model.spring_rest_length.grad,
+                self.train_rate,
+            ),
+            outputs=(self.model.spring_rest_length,),
+        )
+
+        self.tape.zero()
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        self.renderer.begin_frame(0.0)
+        self.renderer.render_box(
+            name="target",
+            pos=self.target_pos,
+            rot=wp.quat_identity(),
+            extents=(0.1, 0.1, 0.1),
+            color=(1.0, 0.0, 0.0),
+        )
+        self.renderer.end_frame()
+
+        for frame in range(self.frame_count):
+            self.renderer.begin_frame(self.render_frame / self.fps)
+            self.renderer.render(self.states[frame * self.sim_substep_count])
+            self.renderer.end_frame()
+
+            self.render_frame += 1
+
+
+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()
+
+        loss = example.loss.numpy()[0]
+
+        if example.verbose:
+            print(f"[{iteration:3d}] loss={loss:.8f}")
+
+        if iteration == example.train_iters - 1 or iteration % example.render_iteration_steps == 0:
+            example.render()
+
+    if example.renderer:
+        example.renderer.save()