warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/optim/example_spring_cage.py

@@ -0,0 +1,245 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# 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 numpy as np
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+
+@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_path="example_spring_cage.usd", num_frames=30, train_iters=25):
+        # Number of frames per second.
+        self.fps = 30
+
+        # Duration of a single simulation iteration in number of frames.
+        self.num_frames = num_frames
+
+        # 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 = train_iters
+
+        # 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 = wp.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 = wp.sim.SemiImplicitIntegrator()
+
+        # Initialize a state for each simulation step.
+        self.states = tuple(self.model.state() for _ in range(self.num_frames * 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_path:
+            # Helper to render the physics scene as a USD file.
+            self.renderer = warp.sim.render.SimRenderer(self.model, stage_path, 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_cuda_graph = wp.get_device().is_cuda
+        if self.use_cuda_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
+
+    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):
+        with wp.ScopedTimer("step"):
+            if self.use_cuda_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
+
+        with wp.ScopedTimer("render"):
+            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.num_frames):
+                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__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument(
+        "--stage_path",
+        type=lambda x: None if x == "None" else str(x),
+        default="example_spring_cage.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=30, help="Total number of frames per training iteration.")
+    parser.add_argument("--train_iters", type=int, default=25, help="Total number of training iterations.")
+    parser.add_argument("--verbose", action="store_true", help="Print out additional status messages during execution.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path, num_frames=args.num_frames, train_iters=args.train_iters)
+
+        for iteration in range(args.train_iters):
+            example.step()
+
+            loss = example.loss.numpy()[0]
+
+            if args.verbose:
+                print(f"[{iteration:3d}] loss={loss:.8f}")
+
+            if example.renderer and (
+                iteration == example.train_iters - 1 or iteration % example.render_iteration_steps == 0
+            ):
+                example.render()
+
+        if example.renderer:
+            example.renderer.save()

warp/examples/optim/example_trajectory.py

@@ -0,0 +1,227 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Trajectory Optimization
+#
+# Shows how to optimize torque trajectories for a simple planar environment
+# using Warp's provided Adam optimizer.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+from warp.optim import Adam
+
+
+@wp.kernel
+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)
+
+
+@wp.kernel
+def apply_torque(torques: wp.array(dtype=wp.float32), start_index: int, body_f: wp.array(dtype=wp.spatial_vector)):
+    fx = torques[start_index + 0]
+    fz = torques[start_index + 1]
+    body_f[0] = wp.spatial_vector(0.0, 0.0, 0.0, fx, 0.0, fz)
+
+
+@wp.kernel
+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[write_index, 0] = pos[0]
+    states[write_index, 1] = pos[2]
+
+
+class Example:
+    def __init__(self, stage_path="example_trajectory.usd", verbose=False, num_frames=100):
+        self.verbose = verbose
+
+        fps = 60
+        self.frame_dt = 1.0 / fps
+        self.num_frames = num_frames
+
+        self.sim_substeps = 1
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
+        self.render_time = 0.0
+
+        self.iter = 0
+
+        # add planar joints
+        builder = wp.sim.ModelBuilder(gravity=0.0)
+        builder.add_articulation()
+        b = builder.add_body(origin=wp.transform())
+        builder.add_shape_sphere(pos=wp.vec3(0.0, 0.0, 0.0), radius=0.1, density=100.0, body=b)
+
+        # compute reference trajectory
+        rad = np.linspace(0.0, np.pi * 2, self.num_frames)
+        self.ref_traj = np.stack([np.cos(rad), np.sin(rad)], axis=1)
+
+        # 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, dtype=wp.float32, requires_grad=True)
+        self.last_traj = wp.empty_like(self.ref_traj)
+
+        # finalize model
+        self.model = builder.finalize(requires_grad=True)
+
+        self.builder = builder
+        self.model.ground = False
+
+        self.dof_q = self.model.joint_coord_count
+        self.dof_qd = self.model.joint_dof_count
+        self.num_bodies = self.model.body_count
+
+        self.action_dim = 2
+        self.state_dim = 2
+
+        assert self.ref_traj.shape == (self.num_frames, self.state_dim)
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        # initial guess
+        self.actions = wp.array(
+            np.zeros(self.num_frames * self.action_dim) * 100.0, dtype=wp.float32, requires_grad=True
+        )
+
+        self.optimizer = Adam([self.actions], lr=1e2)
+        self.loss = wp.zeros(1, dtype=wp.float32, requires_grad=True)
+
+        if stage_path:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=100.0)
+        else:
+            self.renderer = None
+
+        # allocate sim states for trajectory
+        self.states = []
+        for _ in range(self.num_frames + 1):
+            self.states.append(self.model.state())
+
+    def forward(self):
+        """
+        Advances the system dynamics given the rigid-body state in maximal coordinates and generalized joint torques
+        [body_q, body_qd, tau].
+        """
+
+        self.last_traj.zero_()
+
+        for i in range(self.num_frames):
+            state = self.states[i]
+
+            for _ in range(self.sim_substeps):
+                next_state = self.model.state(requires_grad=True)
+
+                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])
+
+                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])
+
+        # compute 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"""
+
+        with wp.ScopedTimer("step"):
+            self.loss.zero_()
+            tape = wp.Tape()
+            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}")
+
+            assert not np.isnan(self.actions.grad.numpy()).any(), "NaN in gradient"
+
+            self.optimizer.step([self.actions.grad])
+            tape.zero()
+            self.iter = self.iter + 1
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            for i in range(self.num_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
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument(
+        "--stage_path",
+        type=lambda x: None if x == "None" else str(x),
+        default="example_trajectory.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=100, help="Total number of frames per training iteration.")
+    parser.add_argument("--train_iters", type=int, default=250, help="Total number of training iterations.")
+    parser.add_argument("--verbose", action="store_true", help="Print out additional status messages during execution.")
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path, verbose=args.verbose, num_frames=args.num_frames)
+
+        for i in range(args.train_iters):
+            example.step()
+
+            if i % 25 == 0:
+                example.render()
+
+        if example.renderer:
+            example.renderer.save()
+
+        np_states = example.last_traj.numpy()
+        np_ref = example.ref_traj.numpy()
+
+        if not args.headless:
+            import matplotlib.pyplot as plt
+
+            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()

warp/examples/sim/example_cartpole.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Sim Cartpole
+#
+# Shows how to set up a simulation of a rigid-body cartpole articulation
+# from a URDF using the wp.sim.ModelBuilder().
+# Note this example does not include a trained policy.
+#
+###########################################################################
+
+import math
+import os
+
+import numpy as np
+
+import warp as wp
+import warp.examples
+import warp.sim
+import warp.sim.render
+
+
+class Example:
+    def __init__(self, stage_path="example_cartpole.usd", num_envs=8):
+        self.num_envs = num_envs
+
+        articulation_builder = wp.sim.ModelBuilder()
+
+        wp.sim.parse_urdf(
+            os.path.join(warp.examples.get_asset_directory(), "cartpole.urdf"),
+            articulation_builder,
+            xform=wp.transform(wp.vec3(), wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)),
+            floating=False,
+            density=100,
+            armature=0.1,
+            stiffness=0.0,
+            damping=0.0,
+            limit_ke=1.0e4,
+            limit_kd=1.0e1,
+            enable_self_collisions=False,
+        )
+
+        builder = wp.sim.ModelBuilder()
+
+        self.sim_time = 0.0
+        fps = 60
+        self.frame_dt = 1.0 / fps
+
+        self.sim_substeps = 10
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
+        for i in range(self.num_envs):
+            builder.add_builder(
+                articulation_builder, xform=wp.transform(np.array((i * 2.0, 4.0, 0.0)), wp.quat_identity())
+            )
+
+            # joint initial positions
+            builder.joint_q[-3:] = [0.0, 0.3, 0.0]
+
+        # finalize model
+        self.model = builder.finalize()
+        self.model.ground = False
+
+        self.model.joint_attach_ke = 1600.0
+        self.model.joint_attach_kd = 20.0
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        self.renderer = None
+        if stage_path:
+            self.renderer = wp.sim.render.SimRenderer(path=stage_path, model=self.model, scaling=15.0)
+
+        self.state = self.model.state()
+
+        wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state)
+
+        self.use_cuda_graph = wp.get_device().is_cuda
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+
+    def simulate(self):
+        for _ in range(self.sim_substeps):
+            self.state.clear_forces()
+            self.state = self.integrator.simulate(self.model, self.state, self.state, self.sim_dt)
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            if self.use_cuda_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.simulate()
+        self.sim_time += self.frame_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state)
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument(
+        "--stage_path",
+        type=lambda x: None if x == "None" else str(x),
+        default="example_cartpole.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=1200, help="Total number of frames.")
+    parser.add_argument("--num_envs", type=int, default=8, help="Total number of simulated environments.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path, num_envs=args.num_envs)
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()
+
+        if example.renderer:
+            example.renderer.save()