warp-lang 1.0.1__py3-none-manylinux2014_x86_64.whl → 1.1.0__py3-none-manylinux2014_x86_64.whl

warp/examples/sim/example_particle_chain.py

@@ -1,106 +1,118 @@
-# 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 Particle Chain
-#
-# Shows how to set up a simple chain of particles connected by springs
-# using wp.sim.ModelBuilder().
-#
-###########################################################################
-
-import math
-import os
-
-import warp as wp
-import warp.sim
-import warp.sim.render
-
-wp.init()
-
-
-class Example:
-    def __init__(self, stage):
-        self.sim_width = 64
-        self.sim_height = 32
-
-        self.sim_fps = 60.0
-        self.frame_dt = 1.0 / self.sim_fps
-        self.sim_substeps = 10
-        self.sim_duration = 5.0
-        self.sim_frames = int(self.sim_duration * self.sim_fps)
-        self.sim_dt = (1.0 / self.sim_fps) / self.sim_substeps
-        self.sim_time = 0.0
-
-        builder = wp.sim.ModelBuilder()
-
-        # anchor
-        builder.add_particle(wp.vec3(0.0, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), 0.0)
-
-        # chain
-        for i in range(1, 10):
-            radius = math.sqrt(i) * 0.2
-            mass = math.pi * radius * radius * radius
-            builder.add_particle(wp.vec3(i, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), mass, radius=radius)
-            builder.add_spring(i - 1, i, 1.0e6, 0.0, 0)
-
-        self.model = builder.finalize()
-        self.model.ground = False
-
-        self.integrator = wp.sim.XPBDIntegrator()
-
-        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=15.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.state_1.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"):
-            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"):
-            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(wp.examples.get_output_directory(), "example_particle_chain.usd")
-
-    example = Example(stage_path)
-
-    for i in range(example.sim_frames):
-        example.step()
-        example.render()
-
-    if example.renderer:
-        example.renderer.save()
+# 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 Particle Chain
+#
+# Shows how to set up a simple chain of particles connected by springs
+# using wp.sim.ModelBuilder().
+#
+###########################################################################
+
+import math
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+wp.init()
+
+
+class Example:
+    def __init__(self, stage_path="example_particle_chain.usd"):
+        self.sim_width = 64
+        self.sim_height = 32
+
+        fps = 60
+        self.frame_dt = 1.0 / fps
+        self.sim_substeps = 10
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_time = 0.0
+
+        builder = wp.sim.ModelBuilder()
+
+        # anchor
+        builder.add_particle(wp.vec3(0.0, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), 0.0)
+
+        # chain
+        for i in range(1, 10):
+            radius = math.sqrt(i) * 0.2
+            mass = math.pi * radius * radius * radius
+            builder.add_particle(wp.vec3(i, 1.0, 0.0), wp.vec3(0.0, 0.0, 0.0), mass, radius=radius)
+            builder.add_spring(i - 1, i, 1.0e6, 0.0, 0)
+
+        self.model = builder.finalize()
+        self.model.ground = False
+
+        self.integrator = wp.sim.XPBDIntegrator()
+
+        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=15.0)
+        else:
+            self.renderer = None
+
+        # 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.state_1.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"):
+            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_particle_chain.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()

warp/examples/sim/example_quadruped.py

@@ -1,179 +1,193 @@
-# 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 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.sim
-import warp.sim.render
-
-wp.init()
-
-
-# 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=None, num_envs=1, print_timers=True):
-        self.num_envs = num_envs
-        articulation_builder = wp.sim.ModelBuilder()
-        wp.sim.parse_urdf(
-            os.path.join(os.path.dirname(__file__), "../assets/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
-        self.frame_dt = 1.0 / 100.0
-
-        episode_duration = 3.0  # seconds
-        self.episode_frames = int(episode_duration / self.frame_dt)
-
-        self.sim_substeps = 5
-        self.sim_dt = self.frame_dt / self.sim_substeps
-
-        offsets = compute_env_offsets(num_envs)
-        for i in range(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.integrator = wp.sim.XPBDIntegrator()
-        # self.integrator = wp.sim.SemiImplicitIntegrator()
-        self.integrator = wp.sim.FeatherstoneIntegrator(self.model)
-
-        self.renderer = None
-        if stage:
-            self.renderer = wp.sim.render.SimRenderer(self.model, stage)
-
-        self.print_timers = print_timers
-
-        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_graph = wp.get_device().is_cuda
-        self.graph = None
-        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()
-            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, 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_0)
-            self.renderer.end_frame()
-
-
-if __name__ == "__main__":
-
-    stage = os.path.join(os.path.dirname(__file__), "example_quadruped.usd")
-
-    example = Example(stage, num_envs=8)
-
-    for _ in range(example.episode_frames):
-        example.step()
-        example.render()
-
-    if example.renderer:
-        example.renderer.save()
+# 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 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
+
+wp.init()
+
+
+# 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 = 5
+        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.integrator = wp.sim.XPBDIntegrator()
+        # self.integrator = wp.sim.SemiImplicitIntegrator()
+        self.integrator = wp.sim.FeatherstoneIntegrator(self.model)
+
+        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()