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

warp/examples/optim/example_trajectory.py

@@ -0,0 +1,201 @@
+# 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 Trajectory Optimization
+#
+# Shows how to optimize torque trajectories for a simple planar environment
+# using Warp's provided Adam optimizer.
+#
+###########################################################################
+
+
+import os
+
+import numpy as np
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+from warp.optim import Adam
+
+wp.init()
+
+
+@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, verbose=False):
+        self.verbose = verbose
+        self.frame_dt = 1.0 / 60.0
+        self.episode_frames = 100
+
+        self.sim_substeps = 1
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
+        self.render_time = 0.0
+
+        self.iter = 0
+
+        builder = wp.sim.ModelBuilder()
+
+        # 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.episode_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.episode_frames, self.state_dim)
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        # initial guess
+        self.actions = wp.array(
+            np.zeros(self.episode_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)
+
+        self.renderer = None
+        if stage:
+            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())
+
+    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.episode_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"""
+        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
+
+        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
+
+
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+
+    stage_path = os.path.join(os.path.dirname(__file__), "example_trajectory.usd")
+
+    example = Example(stage_path, verbose=True)
+
+    # Optimize
+    num_iter = 250
+
+    for i in range(num_iter):
+        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()
+    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,128 @@
+# 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 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.sim
+import warp.sim.render
+
+wp.init()
+
+
+class Example:
+    def __init__(self, stage=None, num_envs=1, print_timers=True):
+        builder = wp.sim.ModelBuilder()
+
+        self.num_envs = num_envs
+
+        articulation_builder = wp.sim.ModelBuilder()
+
+        wp.sim.parse_urdf(
+            os.path.join(os.path.dirname(__file__), "../assets/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
+        self.frame_dt = 1.0 / 60.0
+
+        episode_duration = 20.0  # seconds
+        self.episode_frames = int(episode_duration / self.frame_dt)
+
+        self.sim_substeps = 10
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
+        for i in range(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:
+            self.renderer = wp.sim.render.SimRenderer(path=stage, model=self.model, scaling=15.0)
+
+        self.print_timers = print_timers
+
+        self.state = self.model.state()
+
+        wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state)
+
+        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):
+            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", active=True, print=self.print_timers):
+            if self.use_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", active=True, print=self.print_timers):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state)
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    stage = os.path.join(os.path.dirname(__file__), "example_cartpole.usd")
+
+    example = Example(stage, num_envs=10)
+
+    for _ in range(example.episode_frames):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()

warp/examples/sim/example_cloth.py

@@ -0,0 +1,184 @@
+# 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 Cloth
+#
+# Shows a simulation of an FEM cloth model colliding against a static
+# rigid body mesh using the wp.sim.ModelBuilder().
+#
+###########################################################################
+
+import argparse
+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
+
+wp.init()
+
+
+class IntegratorType(Enum):
+    EULER = "euler"
+    XPBD = "xpbd"
+
+    def __str__(self):
+        return self.value
+
+
+class Example:
+    def __init__(self, stage, integrator=IntegratorType.EULER):
+        self.integrator_type = integrator
+
+        self.sim_width = 64
+        self.sim_height = 32
+
+        self.sim_fps = 60.0
+        self.sim_substeps = 32
+        self.sim_duration = 5.0
+        self.sim_frames = int(self.sim_duration * self.sim_fps)
+        self.frame_dt = 1.0 / self.sim_fps
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_time = 0.0
+        self.profiler = {}
+
+        builder = wp.sim.ModelBuilder()
+
+        if self.integrator_type == IntegratorType.EULER:
+            builder.add_cloth_grid(
+                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,
+                cell_y=0.1,
+                mass=0.1,
+                fix_left=True,
+                tri_ke=1.0e3,
+                tri_ka=1.0e3,
+                tri_kd=1.0e1,
+            )
+        else:
+            builder.add_cloth_grid(
+                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,
+                cell_y=0.1,
+                mass=0.1,
+                fix_left=True,
+                edge_ke=1.0e2,
+                add_springs=True,
+                spring_ke=1.0e3,
+                spring_kd=0.0,
+            )
+
+        usd_stage = Usd.Stage.Open(os.path.join(os.path.dirname(__file__), "../assets/bunny.usd"))
+        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/bunny/bunny"))
+
+        mesh_points = np.array(usd_geom.GetPointsAttr().Get())
+        mesh_indices = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
+
+        mesh = wp.sim.Mesh(mesh_points, mesh_indices)
+
+        builder.add_shape_mesh(
+            body=-1,
+            mesh=mesh,
+            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,
+        )
+
+        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:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+
+    def simulate(self):
+        wp.sim.collide(self.model, self.state_0)
+
+        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)
+
+            # swap states
+            (self.state_0, self.state_1) = (self.state_1, self.state_0)
+
+    def step(self):
+        with wp.ScopedTimer("step", dict=self.profiler):
+            if self.use_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", active=True):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state_0)
+            self.renderer.end_frame()
+
+
+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__), "example_cloth.usd")
+
+    example = Example(stage_path, integrator=args.integrator)
+
+    for i in range(example.sim_frames):
+        example.step()
+        example.render()
+
+    frame_times = example.profiler["step"]
+    print("\nAverage frame sim time: {:.2f} ms".format(sum(frame_times) / len(frame_times)))
+
+    if example.renderer:
+        example.renderer.save()

warp/examples/sim/example_granular.py

@@ -0,0 +1,113 @@
+# 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 Granular
+#
+# Shows how to set up a particle-based granular material model using the
+# wp.sim.ModelBuilder().
+#
+###########################################################################
+
+import os
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+wp.init()
+
+
+class Example:
+    def __init__(self, stage):
+        self.frame_dt = 1.0 / 60
+        self.frame_count = 400
+
+        self.sim_substeps = 64
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_steps = self.frame_count * self.sim_substeps
+        self.sim_time = 0.0
+
+        self.radius = 0.1
+
+        builder = wp.sim.ModelBuilder()
+        builder.default_particle_radius = self.radius
+
+        builder.add_particle_grid(
+            dim_x=16,
+            dim_y=32,
+            dim_z=16,
+            cell_x=self.radius * 2.0,
+            cell_y=self.radius * 2.0,
+            cell_z=self.radius * 2.0,
+            pos=wp.vec3(0.0, 1.0, 0.0),
+            rot=wp.quat_identity(),
+            vel=wp.vec3(5.0, 0.0, 0.0),
+            mass=0.1,
+            jitter=self.radius * 0.1,
+        )
+
+        self.model = builder.finalize()
+        self.model.particle_kf = 25.0
+
+        self.model.soft_contact_kd = 100.0
+        self.model.soft_contact_kf *= 2.0
+
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        self.renderer = None
+        if stage:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=20.0)
+
+        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):
+            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)
+
+    def step(self):
+        with wp.ScopedTimer("step", active=True):
+            self.model.particle_grid.build(self.state_0.particle_q, self.radius * 2.0)
+            if self.use_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", active=True):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render(self.state_0)
+            self.renderer.end_frame()
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_granular.usd")
+
+    example = Example(stage_path)
+
+    for _ in range(example.frame_count):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()