warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/examples/optim/example_trajectory.py

@@ -1,227 +0,0 @@
-# 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

@@ -1,143 +0,0 @@
-# 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()

warp/examples/sim/example_cloth.py

@@ -1,225 +0,0 @@
-# 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 Cloth
-#
-# Shows a simulation of an FEM cloth model colliding against a static
-# rigid body mesh using the wp.sim.ModelBuilder().
-#
-###########################################################################
-
-import math
-import os
-from enum import Enum
-
-import numpy as np
-from pxr import Usd, UsdGeom
-
-import warp as wp
-import warp.examples
-import warp.sim
-import warp.sim.render
-
-
-class IntegratorType(Enum):
-    EULER = "euler"
-    XPBD = "xpbd"
-    VBD = "vbd"
-
-    def __str__(self):
-        return self.value
-
-
-class Example:
-    def __init__(
-        self, stage_path="example_cloth.usd", integrator: IntegratorType = IntegratorType.EULER, height=32, width=64
-    ):
-        self.integrator_type = integrator
-
-        self.sim_height = height
-        self.sim_width = width
-
-        fps = 60
-        self.sim_substeps = 32
-        self.frame_dt = 1.0 / 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,
-            )
-        elif self.integrator_type == IntegratorType.XPBD:
-            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,
-            )
-        else:
-            # VBD
-            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=1e4,
-                tri_ka=1e4,
-                tri_kd=1e-5,
-                edge_ke=100,
-            )
-
-        usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))
-        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/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.VBD:
-            builder.color()
-
-        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()
-        elif self.integrator_type == IntegratorType.XPBD:
-            self.integrator = wp.sim.XPBDIntegrator(iterations=1)
-        else:
-            self.integrator = wp.sim.VBDIntegrator(self.model, iterations=1)
-
-        self.state_0 = self.model.state()
-        self.state_1 = self.model.state()
-
-        if stage_path:
-            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=40.0)
-        else:
-            self.renderer = None
-
-        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):
-        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_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_0)
-            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_cloth.usd",
-        help="Path to the output USD file.",
-    )
-    parser.add_argument("--num_frames", type=int, default=300, help="Total number of frames.")
-    parser.add_argument(
-        "--integrator",
-        help="Type of integrator",
-        type=IntegratorType,
-        choices=list(IntegratorType),
-        default=IntegratorType.EULER,
-    )
-    parser.add_argument("--width", type=int, default=64, help="Cloth resolution in x.")
-    parser.add_argument("--height", type=int, default=32, help="Cloth resolution in y.")
-
-    args = parser.parse_known_args()[0]
-
-    with wp.ScopedDevice(args.device):
-        example = Example(stage_path=args.stage_path, integrator=args.integrator, height=args.height, width=args.width)
-
-        for _i in range(args.num_frames):
-            example.step()
-            example.render()
-
-        frame_times = example.profiler["step"]
-        print(f"\nAverage frame sim time: {sum(frame_times) / len(frame_times):.2f} ms")
-
-        if example.renderer:
-            example.renderer.save()