warp-lang 1.0.2__py3-none-win_amd64.whl → 1.2.0__py3-none-win_amd64.whl

warp/examples/sim/example_cartpole.py

@@ -1,129 +1,137 @@
-# 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.examples
-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(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
-        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_path = "example_cartpole.usd"
-
-    example = Example(stage_path, num_envs=10)
-
-    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 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):
+        builder = wp.sim.ModelBuilder()
+
+        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,186 +1,194 @@
-# 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 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
-
-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(warp.examples.get_asset_directory(), "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__":
-    import argparse
-
-    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 = "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()
+# 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 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"
+
+    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,
+            )
+        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(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.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()
+
+        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("\nAverage frame sim time: {:.2f} ms".format(sum(frame_times) / len(frame_times)))
+
+        if example.renderer:
+            example.renderer.save()