warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/sim/example_quadruped.py

@@ -0,0 +1,203 @@
+# 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 Quadruped
+#
+# Shows how to set up a simulation of a rigid-body quadruped 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
+
+
+# Taken from env/environment.py
+def compute_env_offsets(num_envs, env_offset=(5.0, 0.0, 5.0), up_axis="Y"):
+    # compute positional offsets per environment
+    env_offset = np.array(env_offset)
+    nonzeros = np.nonzero(env_offset)[0]
+    num_dim = nonzeros.shape[0]
+    if num_dim > 0:
+        side_length = int(np.ceil(num_envs ** (1.0 / num_dim)))
+        env_offsets = []
+    else:
+        env_offsets = np.zeros((num_envs, 3))
+    if num_dim == 1:
+        for i in range(num_envs):
+            env_offsets.append(i * env_offset)
+    elif num_dim == 2:
+        for i in range(num_envs):
+            d0 = i // side_length
+            d1 = i % side_length
+            offset = np.zeros(3)
+            offset[nonzeros[0]] = d0 * env_offset[nonzeros[0]]
+            offset[nonzeros[1]] = d1 * env_offset[nonzeros[1]]
+            env_offsets.append(offset)
+    elif num_dim == 3:
+        for i in range(num_envs):
+            d0 = i // (side_length * side_length)
+            d1 = (i // side_length) % side_length
+            d2 = i % side_length
+            offset = np.zeros(3)
+            offset[0] = d0 * env_offset[0]
+            offset[1] = d1 * env_offset[1]
+            offset[2] = d2 * env_offset[2]
+            env_offsets.append(offset)
+    env_offsets = np.array(env_offsets)
+    min_offsets = np.min(env_offsets, axis=0)
+    correction = min_offsets + (np.max(env_offsets, axis=0) - min_offsets) / 2.0
+    if isinstance(up_axis, str):
+        up_axis = "XYZ".index(up_axis.upper())
+    correction[up_axis] = 0.0  # ensure the envs are not shifted below the ground plane
+    env_offsets -= correction
+    return env_offsets
+
+
+class Example:
+    def __init__(self, stage_path="example_quadruped.usd", num_envs=8):
+        articulation_builder = wp.sim.ModelBuilder()
+        wp.sim.parse_urdf(
+            os.path.join(warp.examples.get_asset_directory(), "quadruped.urdf"),
+            articulation_builder,
+            xform=wp.transform([0.0, 0.7, 0.0], wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)),
+            floating=True,
+            density=1000,
+            armature=0.01,
+            stiffness=200,
+            damping=1,
+            contact_ke=1.0e4,
+            contact_kd=1.0e2,
+            contact_kf=1.0e2,
+            contact_mu=1.0,
+            limit_ke=1.0e4,
+            limit_kd=1.0e1,
+        )
+
+        builder = wp.sim.ModelBuilder()
+
+        self.sim_time = 0.0
+        fps = 100
+        self.frame_dt = 1.0 / fps
+
+        self.sim_substeps = 10
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
+        self.num_envs = num_envs
+
+        offsets = compute_env_offsets(self.num_envs)
+        for i in range(self.num_envs):
+            builder.add_builder(articulation_builder, xform=wp.transform(offsets[i], wp.quat_identity()))
+
+            builder.joint_q[-12:] = [0.2, 0.4, -0.6, -0.2, -0.4, 0.6, -0.2, 0.4, -0.6, 0.2, -0.4, 0.6]
+
+            builder.joint_axis_mode = [wp.sim.JOINT_MODE_TARGET_POSITION] * len(builder.joint_axis_mode)
+            builder.joint_act[-12:] = [0.2, 0.4, -0.6, -0.2, -0.4, 0.6, -0.2, 0.4, -0.6, 0.2, -0.4, 0.6]
+
+        np.set_printoptions(suppress=True)
+        # finalize model
+        self.model = builder.finalize()
+        self.model.ground = True
+        # self.model.gravity = 0.0
+
+        self.model.joint_attach_ke = 16000.0
+        self.model.joint_attach_kd = 200.0
+        self.use_tile_gemm = False
+        self.fuse_cholesky = False
+
+        # self.integrator = wp.sim.XPBDIntegrator()
+        # self.integrator = wp.sim.SemiImplicitIntegrator()
+        self.integrator = wp.sim.FeatherstoneIntegrator(
+            self.model, use_tile_gemm=self.use_tile_gemm, fuse_cholesky=self.fuse_cholesky
+        )
+
+        if stage_path:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path)
+        else:
+            self.renderer = None
+
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+
+        wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state_0)
+
+        # simulate() allocates memory via a clone, so we can't use graph capture if the device does not support mempools
+        self.use_cuda_graph = wp.get_device().is_cuda and wp.is_mempool_enabled(wp.get_device())
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+        else:
+            self.graph = None
+
+    def simulate(self):
+        for _ in range(self.sim_substeps):
+            self.state_0.clear_forces()
+            wp.sim.collide(self.model, self.state_0)
+            self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+            self.state_0, self.state_1 = self.state_1, self.state_0
+
+    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_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_quadruped.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("--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_rigid_chain.py

@@ -0,0 +1,203 @@
+# 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 Rigid Chain
+#
+# Shows how to set up a chain of rigid bodies connected by different joint
+# types using wp.sim.ModelBuilder(). There is one chain for each joint
+# type, including fixed joints which act as a flexible beam.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+
+class Example:
+    def __init__(self, stage_path="example_rigid_chain.usd"):
+        self.chain_length = 8
+        self.chain_width = 1.0
+        self.chain_types = [
+            wp.sim.JOINT_REVOLUTE,
+            wp.sim.JOINT_FIXED,
+            wp.sim.JOINT_BALL,
+            wp.sim.JOINT_UNIVERSAL,
+            wp.sim.JOINT_COMPOUND,
+        ]
+
+        builder = wp.sim.ModelBuilder()
+
+        self.sim_time = 0.0
+        fps = 100
+        self.frame_dt = 1.0 / fps
+
+        self.sim_substeps = 10
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
+        for c, t in enumerate(self.chain_types):
+            # start a new articulation
+            builder.add_articulation()
+
+            for i in range(self.chain_length):
+                if i == 0:
+                    parent = -1
+                    parent_joint_xform = wp.transform([0.0, 0.0, c * 1.0], wp.quat_identity())
+                else:
+                    parent = builder.joint_count - 1
+                    parent_joint_xform = wp.transform([self.chain_width, 0.0, 0.0], wp.quat_identity())
+
+                # create body
+                b = builder.add_body(origin=wp.transform([i, 0.0, c * 1.0], wp.quat_identity()), armature=0.1)
+
+                # create shape
+                builder.add_shape_box(
+                    pos=wp.vec3(self.chain_width * 0.5, 0.0, 0.0),
+                    hx=self.chain_width * 0.5,
+                    hy=0.1,
+                    hz=0.1,
+                    density=10.0,
+                    body=b,
+                )
+
+                joint_type = t
+
+                if joint_type == wp.sim.JOINT_REVOLUTE:
+                    joint_limit_lower = -np.deg2rad(60.0)
+                    joint_limit_upper = np.deg2rad(60.0)
+                    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=joint_limit_lower,
+                        limit_upper=joint_limit_upper,
+                        target_ke=0.0,
+                        target_kd=0.0,
+                        limit_ke=1e5,
+                        limit_kd=1.0,
+                    )
+
+                elif joint_type == wp.sim.JOINT_UNIVERSAL:
+                    builder.add_joint_universal(
+                        parent=parent,
+                        child=b,
+                        axis_0=wp.sim.JointAxis((1.0, 0.0, 0.0), -np.deg2rad(60.0), np.deg2rad(60.0)),
+                        axis_1=wp.sim.JointAxis((0.0, 0.0, 1.0), -np.deg2rad(60.0), np.deg2rad(60.0)),
+                        parent_xform=parent_joint_xform,
+                        child_xform=wp.transform_identity(),
+                    )
+
+                elif joint_type == wp.sim.JOINT_BALL:
+                    builder.add_joint_ball(
+                        parent=parent,
+                        child=b,
+                        parent_xform=parent_joint_xform,
+                        child_xform=wp.transform_identity(),
+                    )
+
+                elif joint_type == wp.sim.JOINT_FIXED:
+                    builder.add_joint_fixed(
+                        parent=parent,
+                        child=b,
+                        parent_xform=parent_joint_xform,
+                        child_xform=wp.transform_identity(),
+                    )
+
+                elif joint_type == wp.sim.JOINT_COMPOUND:
+                    builder.add_joint_compound(
+                        parent=parent,
+                        child=b,
+                        axis_0=wp.sim.JointAxis((1.0, 0.0, 0.0), -np.deg2rad(60.0), np.deg2rad(60.0)),
+                        axis_1=wp.sim.JointAxis((0.0, 1.0, 0.0), -np.deg2rad(60.0), np.deg2rad(60.0)),
+                        axis_2=wp.sim.JointAxis((0.0, 0.0, 1.0), -np.deg2rad(60.0), np.deg2rad(60.0)),
+                        parent_xform=parent_joint_xform,
+                        child_xform=wp.transform_identity(),
+                    )
+
+        self.model = builder.finalize()
+        self.model.ground = False
+
+        self.integrator = wp.sim.FeatherstoneIntegrator(self.model)
+
+        if stage_path:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=20.0)
+        else:
+            self.renderer = None
+
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+
+        wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state_0)
+
+        # simulate() allocates memory via a clone, so we can't use graph capture if the device does not support mempools
+        self.use_cuda_graph = wp.get_device().is_cuda and wp.is_mempool_enabled(wp.get_device())
+        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_0.clear_forces()
+            self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+            self.state_0, self.state_1 = self.state_1, self.state_0
+
+    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_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_rigid_chain.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=500, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path)
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()
+
+        if example.renderer:
+            example.renderer.save()

warp/examples/sim/example_rigid_contact.py

@@ -0,0 +1,195 @@
+# 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 Rigid Contact
+#
+# Shows how to set up free rigid bodies with different shape types falling
+# and colliding against each other and the ground using wp.sim.ModelBuilder().
+#
+###########################################################################
+
+import math
+import os
+
+import numpy as np
+from pxr import Usd, UsdGeom
+
+import warp as wp
+import warp.examples
+import warp.sim
+import warp.sim.render
+
+
+class Example:
+    def __init__(self, stage_path="example_rigid_contact.usd"):
+        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
+
+        self.num_bodies = 8
+        self.scale = 0.8
+        self.ke = 1.0e5
+        self.kd = 250.0
+        self.kf = 500.0
+
+        # boxes
+        for i in range(self.num_bodies):
+            b = builder.add_body(origin=wp.transform((i, 1.0, 0.0), wp.quat_identity()))
+
+            builder.add_shape_box(
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                hx=0.5 * self.scale,
+                hy=0.2 * self.scale,
+                hz=0.2 * self.scale,
+                body=i,
+                ke=self.ke,
+                kd=self.kd,
+                kf=self.kf,
+            )
+
+        # spheres
+        for i in range(self.num_bodies):
+            b = builder.add_body(origin=wp.transform((i, 1.0, 2.0), wp.quat_identity()))
+
+            builder.add_shape_sphere(
+                pos=wp.vec3(0.0, 0.0, 0.0), radius=0.25 * self.scale, body=b, ke=self.ke, kd=self.kd, kf=self.kf
+            )
+
+        # capsules
+        for i in range(self.num_bodies):
+            b = builder.add_body(origin=wp.transform((i, 1.0, 6.0), wp.quat_identity()))
+
+            builder.add_shape_capsule(
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                radius=0.25 * self.scale,
+                half_height=self.scale * 0.5,
+                up_axis=0,
+                body=b,
+                ke=self.ke,
+                kd=self.kd,
+                kf=self.kf,
+            )
+
+        # initial spin
+        for i in range(len(builder.body_qd)):
+            builder.body_qd[i] = (0.0, 2.0, 10.0, 0.0, 0.0, 0.0)
+
+        # meshes
+        bunny = self.load_mesh(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"), "/root/bunny")
+        for i in range(self.num_bodies):
+            b = builder.add_body(
+                origin=wp.transform(
+                    (i * 0.5 * self.scale, 1.0 + i * 1.7 * self.scale, 4.0 + i * 0.5 * self.scale),
+                    wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), math.pi * 0.1 * i),
+                )
+            )
+
+            builder.add_shape_mesh(
+                body=b,
+                mesh=bunny,
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                scale=wp.vec3(self.scale, self.scale, self.scale),
+                ke=self.ke,
+                kd=self.kd,
+                kf=self.kf,
+                density=1e3,
+            )
+
+        # finalize model
+        self.model = builder.finalize()
+        self.model.ground = True
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        if stage_path:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=0.5)
+        else:
+            self.renderer = None
+
+        self.state_0 = self.model.state()
+        self.state_1 = self.model.state()
+
+        wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state_0)
+
+        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 load_mesh(self, filename, path):
+        asset_stage = Usd.Stage.Open(filename)
+        mesh_geom = UsdGeom.Mesh(asset_stage.GetPrimAtPath(path))
+
+        points = np.array(mesh_geom.GetPointsAttr().Get())
+        indices = np.array(mesh_geom.GetFaceVertexIndicesAttr().Get()).flatten()
+
+        return wp.sim.Mesh(points, indices)
+
+    def simulate(self):
+        for _ in range(self.sim_substeps):
+            self.state_0.clear_forces()
+            wp.sim.collide(self.model, self.state_0)
+            self.integrator.simulate(self.model, self.state_0, self.state_1, self.sim_dt)
+            self.state_0, self.state_1 = self.state_1, self.state_0
+
+    def step(self):
+        with wp.ScopedTimer("step", active=True):
+            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", active=True):
+            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_rigid_contact.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=300, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path)
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()
+
+        if example.renderer:
+            example.renderer.save()