Spaces:
Sleeping
Sleeping
| # 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.sim | |
| import warp.sim.render | |
| wp.init() | |
| class Example: | |
| def __init__(self, stage=None, num_envs=1, enable_rendering=True, print_timers=True): | |
| self.device = wp.get_device() | |
| builder = wp.sim.ModelBuilder() | |
| self.num_envs = num_envs | |
| articulation_builder = wp.sim.ModelBuilder() | |
| wp.sim.parse_urdf( | |
| os.path.join(os.path.dirname(__file__), "assets/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, | |
| shape_ke=1.0e4, | |
| shape_kd=1.0e2, | |
| shape_kf=1.0e2, | |
| shape_mu=1.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] | |
| builder.joint_target[:3] = [0.0, 0.0, 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.enable_rendering = enable_rendering | |
| self.renderer = None | |
| if self.enable_rendering: | |
| 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 | |
| self.graph = None | |
| if self.use_graph: | |
| # create update graph | |
| wp.capture_begin(self.device) | |
| try: | |
| self.update() | |
| finally: | |
| self.graph = wp.capture_end(self.device) | |
| def update(self): | |
| with wp.ScopedTimer("simulate", active=True, print=self.print_timers): | |
| if not self.use_graph or self.graph is None: | |
| for _ in range(self.sim_substeps): | |
| self.state.clear_forces() | |
| self.state = self.integrator.simulate(self.model, self.state, self.state, self.sim_dt) | |
| else: | |
| wp.capture_launch(self.graph) | |
| if not wp.get_device().is_capturing: | |
| self.sim_time += self.frame_dt | |
| def render(self, is_live=False): | |
| if self.enable_rendering: | |
| with wp.ScopedTimer("render", active=True, print=self.print_timers): | |
| time = 0.0 if is_live else self.sim_time | |
| self.renderer.begin_frame(time) | |
| self.renderer.render(self.state) | |
| self.renderer.end_frame() | |
| def run(self): | |
| profiler = {} | |
| with wp.ScopedTimer("simulate", detailed=False, print=False, active=True, dict=profiler): | |
| for _ in range(self.episode_frames): | |
| self.update() | |
| self.render() | |
| wp.synchronize() | |
| if self.enable_rendering: | |
| self.renderer.save() | |
| 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}") | |
| return 1000.0 * float(self.num_envs) / avg_time | |
| if __name__ == "__main__": | |
| profile = False | |
| if profile: | |
| env_count = 2 | |
| env_times = [] | |
| env_size = [] | |
| for i in range(15): | |
| example = Example(num_envs=env_count, enable_rendering=False, print_timers=False) | |
| steps_per_second = example.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 | |
| import matplotlib.pyplot as plt | |
| 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: | |
| stage = os.path.join(os.path.dirname(__file__), "outputs/example_sim_cartpole.usd") | |
| example = Example(stage, num_envs=10) | |
| example.run() | |