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

warp/examples/env/env_usd.py

@@ -1,97 +0,0 @@
-# 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.
-
-###########################################################################
-# USD Environment
-#
-# Shows how to load a USD file containing USD Physics schema definitions.
-#
-###########################################################################
-
-import warp as wp
-import warp.sim
-
-from .environment import Environment, run_env
-
-
-class UsdEnvironment(Environment):
-    sim_name = "env_usd"
-    opengl_render_settings = dict(scaling=10.0, draw_grid=True)
-    usd_render_settings = dict(scaling=100.0)
-
-    episode_duration = 2.0
-
-    sim_substeps_euler = 64
-    sim_substeps_xpbd = 8
-
-    xpbd_settings = dict(
-        iterations=10,
-        enable_restitution=True,
-        joint_linear_relaxation=0.8,
-        joint_angular_relaxation=0.45,
-        rigid_contact_relaxation=1.0,
-        rigid_contact_con_weighting=True,
-    )
-
-    # USD files define their own ground plane if necessary
-    activate_ground_plane = False
-    num_envs = 1
-
-    plot_body_coords = False
-
-    def create_articulation(self, builder):
-        usd_filename = wp.sim.resolve_usd_from_url(
-            "http://omniverse-content-staging.s3-us-west-2.amazonaws.com/Assets/Isaac/2022.2.1/Isaac/Robots/Franka/franka_instanceable.usd",
-            target_folder_name=".panda_usd_files")
-        settings = wp.sim.parse_usd(
-            usd_filename,
-            builder,
-            default_thickness=0.01,
-            # ignore collision meshes from Franka robot
-            ignore_paths=[".*collisions.*"],
-            default_ke=1e6,
-        )
-
-        self.frame_dt = 1.0 / settings["fps"]
-        if settings["duration"] > 0.0:
-            self.episode_duration = settings["duration"]
-        self.sim_substeps = 10
-        self.sim_dt = self.frame_dt / self.sim_substeps
-        self.episode_frames = int(self.episode_duration / self.frame_dt)
-        self.sim_steps = int(self.episode_duration / self.sim_dt)
-        self.sim_time = 0.0
-        self.render_time = 0.0
-        self.up_axis = settings["up_axis"]
-
-    def before_simulate(self):
-        # print some information about the loaded model
-        if self.model.shape_count > 0:
-            print("shape_transform", self.model.shape_transform.numpy())
-            print("geo_scale", self.model.shape_geo.scale.numpy())
-        if self.model.joint_count > 0:
-            print("joint parent", self.model.joint_parent.numpy())
-            print("joint child", self.model.joint_child.numpy())
-            if len(self.model.joint_q) > 0:
-                print("joint q", self.model.joint_q.numpy())
-            if len(self.model.joint_axis) > 0:
-                print("joint axis", self.model.joint_axis.numpy())
-                print("joint target", self.model.joint_target.numpy())
-                print("joint target ke", self.model.joint_target_ke.numpy())
-                print("joint target kd", self.model.joint_target_kd.numpy())
-                print("joint limit lower", self.model.joint_limit_lower.numpy())
-                print("joint limit upper", self.model.joint_limit_upper.numpy())
-            print("joint_X_p", self.model.joint_X_p.numpy())
-            print("joint_X_c", self.model.joint_X_c.numpy())
-        if self.model.body_count > 0:
-            print("COM", self.model.body_com.numpy())
-            print("Mass", self.model.body_mass.numpy())
-            print("Inertia", self.model.body_inertia.numpy())
-            print("body_q", self.state.body_q.numpy())
-
-
-if __name__ == "__main__":
-    run_env(UsdEnvironment)

warp/examples/env/environment.py

@@ -1,526 +0,0 @@
-# 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.
-
-import argparse
-import os
-from enum import Enum
-from typing import Tuple
-
-import numpy as np
-
-import warp as wp
-import warp.sim
-import warp.sim.render
-
-wp.init()
-
-
-class RenderMode(Enum):
-    NONE = "none"
-    OPENGL = "opengl"
-    USD = "usd"
-
-    def __str__(self):
-        return self.value
-
-
-class IntegratorType(Enum):
-    EULER = "euler"
-    XPBD = "xpbd"
-
-    def __str__(self):
-        return self.value
-
-
-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 Environment:
-    sim_name: str = "Environment"
-
-    frame_dt = 1.0 / 60.0
-
-    episode_duration = 5.0  # seconds
-
-    # whether to play the simulation indefinitely when using the OpenGL renderer
-    continuous_opengl_render: bool = True
-
-    sim_substeps_euler: int = 16
-    sim_substeps_xpbd: int = 5
-
-    euler_settings = dict()
-    xpbd_settings = dict()
-
-    render_mode: RenderMode = RenderMode.OPENGL
-    opengl_render_settings = dict()
-    usd_render_settings = dict(scaling=10.0)
-    show_rigid_contact_points = False
-    contact_points_radius = 1e-3
-    show_joints = False
-    # whether OpenGLRenderer should render each environment in a separate tile
-    use_tiled_rendering = False
-
-    # whether to apply model.joint_q, joint_qd to bodies before simulating
-    eval_fk: bool = True
-
-    profile: bool = False
-
-    use_graph_capture: bool = wp.get_preferred_device().is_cuda
-
-    num_envs: int = 100
-
-    activate_ground_plane: bool = True
-
-    integrator_type: IntegratorType = IntegratorType.XPBD
-
-    up_axis: str = "Y"
-    gravity: float = -9.81
-    env_offset: Tuple[float, float, float] = (1.0, 0.0, 1.0)
-
-    # stiffness and damping for joint attachment dynamics used by Euler
-    joint_attach_ke: float = 32000.0
-    joint_attach_kd: float = 50.0
-
-    # distance threshold at which contacts are generated
-    rigid_contact_margin: float = 0.05
-
-    # whether each environment should have its own collision group
-    # to avoid collisions between environments
-    separate_collision_group_per_env: bool = True
-
-    plot_body_coords: bool = False
-    plot_joint_coords: bool = False
-
-    requires_grad: bool = False
-
-    # control-related definitions, to be updated by derived classes
-    control_dim: int = 0
-
-    def __init__(self):
-        self.parser = argparse.ArgumentParser()
-        self.parser.add_argument(
-            "--integrator",
-            help="Type of integrator",
-            type=IntegratorType,
-            choices=list(IntegratorType),
-            default=self.integrator_type.value,
-        )
-        self.parser.add_argument(
-            "--visualizer",
-            help="Type of renderer",
-            type=RenderMode,
-            choices=list(RenderMode),
-            default=self.render_mode.value,
-        )
-        self.parser.add_argument(
-            "--num_envs", help="Number of environments to simulate", type=int, default=self.num_envs
-        )
-        self.parser.add_argument("--profile", help="Enable profiling", type=bool, default=self.profile)
-
-    def parse_args(self):
-        args = self.parser.parse_args()
-        self.integrator_type = args.integrator
-        self.render_mode = args.visualizer
-        self.num_envs = args.num_envs
-        self.profile = args.profile
-
-    def init(self):
-        if self.integrator_type == IntegratorType.EULER:
-            self.sim_substeps = self.sim_substeps_euler
-        elif self.integrator_type == IntegratorType.XPBD:
-            self.sim_substeps = self.sim_substeps_xpbd
-
-        self.episode_frames = int(self.episode_duration / self.frame_dt)
-        self.sim_dt = self.frame_dt / self.sim_substeps
-        self.sim_steps = int(self.episode_duration / self.sim_dt)
-
-        if self.use_tiled_rendering and self.render_mode == RenderMode.OPENGL:
-            # no environment offset when using tiled rendering
-            self.env_offset = (0.0, 0.0, 0.0)
-
-        builder = wp.sim.ModelBuilder()
-        builder.rigid_contact_margin = self.rigid_contact_margin
-        try:
-            articulation_builder = wp.sim.ModelBuilder()
-            self.create_articulation(articulation_builder)
-            env_offsets = compute_env_offsets(self.num_envs, self.env_offset, self.up_axis)
-            for i in range(self.num_envs):
-                xform = wp.transform(env_offsets[i], wp.quat_identity())
-                builder.add_builder(
-                    articulation_builder, xform, separate_collision_group=self.separate_collision_group_per_env
-                )
-            self.bodies_per_env = len(articulation_builder.body_q)
-        except NotImplementedError:
-            # custom simulation setup where something other than an articulation is used
-            self.setup(builder)
-            self.bodies_per_env = len(builder.body_q)
-
-        self.model = builder.finalize()
-        self.device = self.model.device
-        if not self.device.is_cuda:
-            self.use_graph_capture = False
-        self.model.ground = self.activate_ground_plane
-
-        self.model.joint_attach_ke = self.joint_attach_ke
-        self.model.joint_attach_kd = self.joint_attach_kd
-
-        # set up current and next state to be used by the integrator
-        self.state_0 = None
-        self.state_1 = None
-
-        if self.integrator_type == IntegratorType.EULER:
-            self.integrator = wp.sim.SemiImplicitIntegrator(**self.euler_settings)
-        elif self.integrator_type == IntegratorType.XPBD:
-            self.integrator = wp.sim.XPBDIntegrator(**self.xpbd_settings)
-
-        self.renderer = None
-        if self.profile:
-            self.render_mode = RenderMode.NONE
-        if self.render_mode == RenderMode.OPENGL:
-            self.renderer = wp.sim.render.SimRendererOpenGL(
-                self.model,
-                self.sim_name,
-                up_axis=self.up_axis,
-                show_rigid_contact_points=self.show_rigid_contact_points,
-                contact_points_radius=self.contact_points_radius,
-                show_joints=self.show_joints,
-                **self.opengl_render_settings,
-            )
-            if self.use_tiled_rendering and self.num_envs > 1:
-                floor_id = self.model.shape_count - 1
-                # all shapes except the floor
-                instance_ids = np.arange(floor_id, dtype=np.int32).tolist()
-                shapes_per_env = floor_id // self.num_envs
-                additional_instances = []
-                if self.activate_ground_plane:
-                    additional_instances.append(floor_id)
-                self.renderer.setup_tiled_rendering(
-                    instances=[
-                        instance_ids[i * shapes_per_env : (i + 1) * shapes_per_env] + additional_instances
-                        for i in range(self.num_envs)
-                    ]
-                )
-        elif self.render_mode == RenderMode.USD:
-            filename = os.path.join(os.path.dirname(__file__), "..", "outputs", self.sim_name + ".usd")
-            self.renderer = wp.sim.render.SimRendererUsd(
-                self.model,
-                filename,
-                up_axis=self.up_axis,
-                show_rigid_contact_points=self.show_rigid_contact_points,
-                **self.usd_render_settings,
-            )
-
-    def create_articulation(self, builder):
-        raise NotImplementedError
-
-    def setup(self, builder):
-        pass
-
-    def customize_model(self, model):
-        pass
-
-    def before_simulate(self):
-        pass
-
-    def after_simulate(self):
-        pass
-
-    def custom_update(self):
-        pass
-
-    @property
-    def state(self):
-        # shortcut to current state
-        return self.state_0
-
-    def update(self):
-        for i in range(self.sim_substeps):
-            self.state_0.clear_forces()
-            self.custom_update()
-            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 render(self, state=None):
-        if self.renderer is not None:
-            with wp.ScopedTimer("render", False):
-                self.render_time += self.frame_dt
-                self.renderer.begin_frame(self.render_time)
-                # render state 1 (swapped with state 0 just before)
-                self.renderer.render(state or self.state_1)
-                self.renderer.end_frame()
-
-    def run(self):
-        # ---------------
-        # run simulation
-
-        self.sim_time = 0.0
-        self.render_time = 0.0
-        self.state_0 = self.model.state()
-        self.state_1 = self.model.state()
-
-        if self.eval_fk:
-            wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.state_0)
-
-        self.before_simulate()
-
-        if self.renderer is not None:
-            self.render(self.state_0)
-
-            if self.render_mode == RenderMode.OPENGL:
-                self.renderer.paused = True
-
-        profiler = {}
-
-        if self.use_graph_capture:
-            # create update graph
-            wp.capture_begin()
-            try:
-                self.update()
-            finally:
-                graph = wp.capture_end()
-
-        if self.plot_body_coords:
-            q_history = []
-            q_history.append(self.state_0.body_q.numpy().copy())
-            qd_history = []
-            qd_history.append(self.state_0.body_qd.numpy().copy())
-            delta_history = []
-            delta_history.append(self.state_0.body_deltas.numpy().copy())
-            num_con_history = []
-            num_con_history.append(self.model.rigid_contact_inv_weight.numpy().copy())
-        if self.plot_joint_coords:
-            joint_q_history = []
-            joint_q = wp.zeros_like(self.model.joint_q)
-            joint_qd = wp.zeros_like(self.model.joint_qd)
-
-        # simulate
-        with wp.ScopedTimer("simulate", detailed=False, print=False, active=True, dict=profiler):
-            running = True
-            while running:
-                for f in range(self.episode_frames):
-                    if self.use_graph_capture:
-                        wp.capture_launch(graph)
-                        self.sim_time += self.frame_dt
-                    else:
-                        self.update()
-                        self.sim_time += self.frame_dt
-
-                        if not self.profile:
-                            if self.plot_body_coords:
-                                q_history.append(self.state_0.body_q.numpy().copy())
-                                qd_history.append(self.state_0.body_qd.numpy().copy())
-                                delta_history.append(self.state_0.body_deltas.numpy().copy())
-                                num_con_history.append(self.model.rigid_contact_inv_weight.numpy().copy())
-
-                            if self.plot_joint_coords:
-                                wp.sim.eval_ik(self.model, self.state_0, joint_q, joint_qd)
-                                joint_q_history.append(joint_q.numpy().copy())
-
-                    self.render()
-                    if self.render_mode == RenderMode.OPENGL and self.renderer.has_exit:
-                        running = False
-                        break
-
-                if not self.continuous_opengl_render or self.render_mode != RenderMode.OPENGL:
-                    break
-
-            wp.synchronize()
-
-        self.after_simulate()
-
-        avg_time = np.array(profiler["simulate"]).mean() / self.episode_frames
-        avg_steps_second = 1000.0 * float(self.num_envs) / avg_time
-
-        print(f"envs: {self.num_envs} steps/second {avg_steps_second} avg_time {avg_time}")
-
-        if self.renderer is not None:
-            self.renderer.save()
-
-        if self.plot_body_coords:
-            import matplotlib.pyplot as plt
-
-            q_history = np.array(q_history)
-            qd_history = np.array(qd_history)
-            delta_history = np.array(delta_history)
-            num_con_history = np.array(num_con_history)
-
-            # find bodies with non-zero mass
-            body_indices = np.where(self.model.body_mass.numpy() > 0)[0]
-            body_indices = body_indices[:5]  # limit number of bodies to plot
-
-            fig, ax = plt.subplots(len(body_indices), 7, figsize=(10, 10), squeeze=False)
-            fig.subplots_adjust(hspace=0.2, wspace=0.2)
-            for i, j in enumerate(body_indices):
-                ax[i, 0].set_title(f"Body {j} Position")
-                ax[i, 0].grid()
-                ax[i, 1].set_title(f"Body {j} Orientation")
-                ax[i, 1].grid()
-                ax[i, 2].set_title(f"Body {j} Linear Velocity")
-                ax[i, 2].grid()
-                ax[i, 3].set_title(f"Body {j} Angular Velocity")
-                ax[i, 3].grid()
-                ax[i, 4].set_title(f"Body {j} Linear Delta")
-                ax[i, 4].grid()
-                ax[i, 5].set_title(f"Body {j} Angular Delta")
-                ax[i, 5].grid()
-                ax[i, 6].set_title(f"Body {j} Num Contacts")
-                ax[i, 6].grid()
-                ax[i, 0].plot(q_history[:, j, :3])
-                ax[i, 1].plot(q_history[:, j, 3:])
-                ax[i, 2].plot(qd_history[:, j, 3:])
-                ax[i, 3].plot(qd_history[:, j, :3])
-                ax[i, 4].plot(delta_history[:, j, 3:])
-                ax[i, 5].plot(delta_history[:, j, :3])
-                ax[i, 6].plot(num_con_history[:, j])
-                ax[i, 0].set_xlim(0, self.sim_steps)
-                ax[i, 1].set_xlim(0, self.sim_steps)
-                ax[i, 2].set_xlim(0, self.sim_steps)
-                ax[i, 3].set_xlim(0, self.sim_steps)
-                ax[i, 4].set_xlim(0, self.sim_steps)
-                ax[i, 5].set_xlim(0, self.sim_steps)
-                ax[i, 6].set_xlim(0, self.sim_steps)
-                ax[i, 6].yaxis.get_major_locator().set_params(integer=True)
-            plt.show()
-
-        if self.plot_joint_coords:
-            import matplotlib.pyplot as plt
-
-            joint_q_history = np.array(joint_q_history)
-            dof_q = joint_q_history.shape[1]
-            ncols = int(np.ceil(np.sqrt(dof_q)))
-            nrows = int(np.ceil(dof_q / float(ncols)))
-            fig, axes = plt.subplots(
-                ncols=ncols,
-                nrows=nrows,
-                constrained_layout=True,
-                figsize=(ncols * 3.5, nrows * 3.5),
-                squeeze=False,
-                sharex=True,
-            )
-
-            joint_id = 0
-            joint_type_names = {
-                wp.sim.JOINT_BALL: "ball",
-                wp.sim.JOINT_REVOLUTE: "hinge",
-                wp.sim.JOINT_PRISMATIC: "slide",
-                wp.sim.JOINT_UNIVERSAL: "universal",
-                wp.sim.JOINT_COMPOUND: "compound",
-                wp.sim.JOINT_FREE: "free",
-                wp.sim.JOINT_FIXED: "fixed",
-                wp.sim.JOINT_DISTANCE: "distance",
-                wp.sim.JOINT_D6: "D6",
-            }
-            joint_lower = self.model.joint_limit_lower.numpy()
-            joint_upper = self.model.joint_limit_upper.numpy()
-            joint_type = self.model.joint_type.numpy()
-            while joint_id < len(joint_type) - 1 and joint_type[joint_id] == wp.sim.JOINT_FIXED:
-                # skip fixed joints
-                joint_id += 1
-            q_start = self.model.joint_q_start.numpy()
-            qd_start = self.model.joint_qd_start.numpy()
-            qd_i = qd_start[joint_id]
-            for dim in range(ncols * nrows):
-                ax = axes[dim // ncols, dim % ncols]
-                if dim >= dof_q:
-                    ax.axis("off")
-                    continue
-                ax.grid()
-                ax.plot(joint_q_history[:, dim])
-                if joint_type[joint_id] != wp.sim.JOINT_FREE:
-                    lower = joint_lower[qd_i]
-                    if abs(lower) < 2 * np.pi:
-                        ax.axhline(lower, color="red")
-                    upper = joint_upper[qd_i]
-                    if abs(upper) < 2 * np.pi:
-                        ax.axhline(upper, color="red")
-                joint_name = joint_type_names[joint_type[joint_id]]
-                ax.set_title(f"$\\mathbf{{q_{{{dim}}}}}$ ({self.model.joint_name[joint_id]} / {joint_name} {joint_id})")
-                if joint_id < self.model.joint_count - 1 and q_start[joint_id + 1] == dim + 1:
-                    joint_id += 1
-                    qd_i = qd_start[joint_id]
-                else:
-                    qd_i += 1
-            plt.tight_layout()
-            plt.show()
-
-        return 1000.0 * float(self.num_envs) / avg_time
-
-
-def run_env(Demo):
-    demo = Demo()
-    demo.parse_args()
-    if demo.profile:
-        import matplotlib.pyplot as plt
-
-        env_count = 2
-        env_times = []
-        env_size = []
-
-        for i in range(15):
-            demo.num_envs = env_count
-            demo.init()
-            steps_per_second = demo.run()
-
-            env_size.append(env_count)
-            env_times.append(steps_per_second)
-
-            env_count *= 2
-
-        # dump times
-        for i in range(len(env_times)):
-            print(f"envs: {env_size[i]} steps/second: {env_times[i]}")
-
-        # plot
-        plt.figure(1)
-        plt.plot(env_size, env_times)
-        plt.xscale("log")
-        plt.xlabel("Number of Envs")
-        plt.yscale("log")
-        plt.ylabel("Steps/Second")
-        plt.show()
-    else:
-        demo.init()
-        return demo.run()