From 7b27eb6e0248e31eb9afdd23456c7e729f0a0368 Mon Sep 17 00:00:00 2001 From: StoneT2000 Date: Mon, 22 Jan 2024 16:35:57 -0800 Subject: [PATCH] cleaner benchmarking code --- examples/benchmarking/README.md | 10 ++- examples/benchmarking/benchmark_cpu_sim.py | 16 ++-- examples/benchmarking/benchmark_gpu_sim.py | 89 ++++++++++++++----- .../benchmark_omniisaacgymenvs_sim.py | 45 +++++++--- examples/benchmarking/benchmark_orbit_sim.py | 34 ++++--- mani_skill2/envs/pick_and_place/pick_cube.py | 4 +- mani_skill2/envs/sapien_env.py | 23 +++-- 7 files changed, 158 insertions(+), 63 deletions(-) diff --git a/examples/benchmarking/README.md b/examples/benchmarking/README.md index a0f2ca7e3..4628b19fa 100644 --- a/examples/benchmarking/README.md +++ b/examples/benchmarking/README.md @@ -1,3 +1,11 @@ # Benchmarking -Code here is used to benchmark difference frameworks/simulators. \ No newline at end of file +Code here is used to benchmark the performance of various simulators/benchmarks + +To benchmark ManiSkill + SAPIEN, run + +``` +python benchmark_gpu_sim.py --num-envs=1024 --obs-mode=state # test just state simulation +python benchmark_gpu_sim.py --num-envs=128 --obs-mode=rgbd # test state sim + parallel rendering +python benchmark_gpu_sim.py --num-envs=128 --obs-mode=state --save-video # save a video showing all 128 visual observations +``` \ No newline at end of file diff --git a/examples/benchmarking/benchmark_cpu_sim.py b/examples/benchmarking/benchmark_cpu_sim.py index 46bf24336..b87b0f05f 100644 --- a/examples/benchmarking/benchmark_cpu_sim.py +++ b/examples/benchmarking/benchmark_cpu_sim.py @@ -1,12 +1,14 @@ import time + import gymnasium as gym +import numpy as np import sapien -import mani_skill2.envs import sapien.physx -import tqdm -import numpy as np import sapien.render -import sapien.physx +import tqdm + +import mani_skill2.envs + if __name__ == "__main__": num_envs = 12 env_id = "PickCube-v0" @@ -36,5 +38,7 @@ print("RESET") dtime = time.time() - stime FPS = num_envs * N / dtime - print(f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs with step+reset") - env.close() \ No newline at end of file + print( + f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs with step+reset" + ) + env.close() diff --git a/examples/benchmarking/benchmark_gpu_sim.py b/examples/benchmarking/benchmark_gpu_sim.py index eb5f61f77..5d3c0e9a9 100644 --- a/examples/benchmarking/benchmark_gpu_sim.py +++ b/examples/benchmarking/benchmark_gpu_sim.py @@ -2,73 +2,118 @@ # python manualtest/benchmark_orbit_sim.py --task "Isaac-Lift-Cube-Franka-v0" --num_envs 512 --headless import argparse import time + import gymnasium as gym +import numpy as np import sapien -import torch -import mani_skill2.envs import sapien.physx -import tqdm -import numpy as np import sapien.render -import sapien.physx +import torch +import tqdm + +import mani_skill2.envs from mani_skill2.utils.visualization.misc import images_to_video, tile_images + def main(args): num_envs = args.num_envs - sapien.physx.set_gpu_memory_config(found_lost_pairs_capacity=2**26, max_rigid_patch_count=120000) - env = gym.make(args.env_id, num_envs=num_envs, obs_mode=args.obs_mode, enable_shadow=True, render_mode=args.render_mode, control_mode="pd_joint_delta_pos", sim_freq=100, control_freq=50) - print(f"[INFO]: Gym observation space: {env.observation_space}") - print(f"[INFO]: Gym action space: {env.action_space}") - + # TODO (stao): we need to auto set this gpu memory config somehow + sapien.physx.set_gpu_memory_config( + found_lost_pairs_capacity=2**26, max_rigid_patch_count=120000 + ) + env = gym.make( + args.env_id, + num_envs=num_envs, + obs_mode=args.obs_mode, + enable_shadow=True, + render_mode=args.render_mode, + control_mode=args.control_mode, + ) + print( + "# -------------------------------------------------------------------------- #" + ) + print( + f"Benchmarking ManiSkill GPU Simulation with {num_envs} parallel environments" + ) + print( + f"env_id={args.env_id}, obs_mode={args.obs_mode}, control_mode={args.control_mode}" + ) + print(f"render_mode={args.render_mode}, save_video={args.save_video}") + print(f"sim_freq={env.unwrapped.sim_freq}, control_freq={env.unwrapped.control_freq}") + print(f"observation space: {env.observation_space}") + print(f"action space: {env.action_space}") + print( + "# -------------------------------------------------------------------------- #" + ) images = [] - video_nrows=int(np.sqrt(num_envs)) + video_nrows = int(np.sqrt(num_envs)) with torch.inference_mode(): env.reset(seed=2022) - env.step(env.action_space.sample()) # warmup? it seems first step here is slow for some reason + env.step(env.action_space.sample()) # warmup step env.reset(seed=2022) if args.save_video: images.append(env.render()) N = 100 stime = time.time() for i in tqdm.tqdm(range(N)): - actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + actions = ( + 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + ) obs, rew, terminated, truncated, info = env.step(actions) if args.save_video: images.append(env.render()) dtime = time.time() - stime FPS = num_envs * N / dtime - print(f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs") + print(f"{FPS=:0.3f}. {N=} steps in {dtime:0.3f}s with {num_envs} parallel envs") if args.save_video: - images = [tile_images(rgbs, nrows=video_nrows).cpu().numpy() for rgbs in images] - images_to_video(images, output_dir="./videos/benchmark", video_name=f"mani_skill_gpu_sim-num_envs={num_envs}-obs_mode={args.obs_mode}-render_mode={args.render_mode}", fps=30) + images = [ + tile_images(rgbs, nrows=video_nrows).cpu().numpy() for rgbs in images + ] + images_to_video( + images, + output_dir="./videos/benchmark", + video_name=f"mani_skill_gpu_sim-num_envs={num_envs}-obs_mode={args.obs_mode}-render_mode={args.render_mode}", + fps=30, + ) del images env.reset(seed=2022) N = 1000 stime = time.time() for i in tqdm.tqdm(range(N)): - actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + actions = ( + 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + ) obs, rew, terminated, truncated, info = env.step(actions) if i % 200 == 0 and i != 0: env.reset() - print("RESET") dtime = time.time() - stime FPS = num_envs * N / dtime - print(f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs with step+reset") + print( + f"{FPS=:0.3f}. {N=} steps in {dtime:0.3f}s with {num_envs} parallel envs with step+reset" + ) env.close() + + def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("-e", "--env-id", type=str, default="PickCube-v0") - parser.add_argument("-o", "--obs-mode", type=str, default="none") - parser.add_argument("-n", "--num-envs", type=int, default=256) + parser.add_argument("-o", "--obs-mode", type=str, default="state") + parser.add_argument("-c", "--control-mode", type=str, default="pd_joint_delta_pos") + parser.add_argument("-n", "--num-envs", type=int, default=1024) parser.add_argument( - "--render-mode", type=str, default="rgb_array" + "--render-mode", + type=str, + default="cameras", + help="which set of cameras/sensors to render for video saving. 'cameras' value will save a video showing all sensor/camera data in the observation, e.g. rgb and depth. 'rgb_array' value will show a higher quality render of the environment running.", ), parser.add_argument( "--save-video", action="store_true", help="whether to save videos" ) args = parser.parse_args() return args + + if __name__ == "__main__": main(parse_args()) diff --git a/examples/benchmarking/benchmark_omniisaacgymenvs_sim.py b/examples/benchmarking/benchmark_omniisaacgymenvs_sim.py index 657f570bd..6ae4b34d4 100644 --- a/examples/benchmarking/benchmark_omniisaacgymenvs_sim.py +++ b/examples/benchmarking/benchmark_omniisaacgymenvs_sim.py @@ -27,16 +27,15 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -import gym -import hydra -from omegaconf import DictConfig import os import time +import gym +import hydra import numpy as np -import torch - import omniisaacgymenvs +import torch +from omegaconf import DictConfig from omniisaacgymenvs.envs.vec_env_rlgames import VecEnvRLGames from omniisaacgymenvs.utils.config_utils.path_utils import get_experience from omniisaacgymenvs.utils.hydra_cfg.hydra_utils import * @@ -46,7 +45,6 @@ @hydra.main(version_base=None, config_name="config", config_path="../cfg") def parse_hydra_configs(cfg: DictConfig): - cfg_dict = omegaconf_to_dict(cfg) print_dict(cfg_dict) @@ -55,22 +53,30 @@ def parse_hydra_configs(cfg: DictConfig): enable_viewport = "enable_cameras" in cfg.task.sim and cfg.task.sim.enable_cameras # select kit app file - experience = get_experience(headless, cfg.enable_livestream, enable_viewport, cfg.enable_recording, cfg.kit_app) + experience = get_experience( + headless, + cfg.enable_livestream, + enable_viewport, + cfg.enable_recording, + cfg.kit_app, + ) env = VecEnvRLGames( headless=headless, sim_device=cfg.device_id, enable_livestream=cfg.enable_livestream, enable_viewport=enable_viewport or cfg.enable_recording, - experience=experience + experience=experience, ) # parse experiment directory - module_path = os.path.abspath(os.path.join(os.path.dirname(omniisaacgymenvs.__file__))) + module_path = os.path.abspath( + os.path.join(os.path.dirname(omniisaacgymenvs.__file__)) + ) experiment_dir = os.path.join(module_path, "runs", cfg.train.params.config.name) # use gym RecordVideo wrapper for viewport recording if cfg.enable_recording: - if cfg.recording_dir == '': + if cfg.recording_dir == "": videos_dir = os.path.join(experiment_dir, "videos") else: videos_dir = cfg.recording_dir @@ -78,12 +84,18 @@ def parse_hydra_configs(cfg: DictConfig): video_length = cfg.recording_length env.is_vector_env = True if env.metadata is None: - env.metadata = {"render_modes": ["rgb_array"], "render_fps": cfg.recording_fps} + env.metadata = { + "render_modes": ["rgb_array"], + "render_fps": cfg.recording_fps, + } else: env.metadata["render_modes"] = ["rgb_array"] env.metadata["render_fps"] = cfg.recording_fps env = gym.wrappers.RecordVideo( - env, video_folder=videos_dir, step_trigger=video_interval, video_length=video_length + env, + video_folder=videos_dir, + step_trigger=video_interval, + video_length=video_length, ) # sets seed. if seed is -1 will pick a random one @@ -99,6 +111,7 @@ def parse_hydra_configs(cfg: DictConfig): num_envs = env.num_envs import tqdm + with torch.inference_mode(): # reset environment env.reset(seed=2022) @@ -111,7 +124,9 @@ def parse_hydra_configs(cfg: DictConfig): obs, rew, done, info = env.step(actions) dtime = time.time() - stime FPS = num_envs * N / dtime - print(f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs") + print( + f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs" + ) env.reset(seed=2022) torch.manual_seed(0) @@ -126,7 +141,9 @@ def parse_hydra_configs(cfg: DictConfig): print("RESET") dtime = time.time() - stime FPS = num_envs * N / dtime - print(f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs with step+reset") + print( + f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs with step+reset" + ) env.simulation_app.close() diff --git a/examples/benchmarking/benchmark_orbit_sim.py b/examples/benchmarking/benchmark_orbit_sim.py index b7476bf45..0668dc646 100644 --- a/examples/benchmarking/benchmark_orbit_sim.py +++ b/examples/benchmarking/benchmark_orbit_sim.py @@ -6,6 +6,7 @@ """Script to run an environment with zero action agent.""" from __future__ import annotations + import time import tqdm @@ -15,6 +16,7 @@ import argparse import logging + import carb logging.getLogger("omni.hydra").setLevel(logging.ERROR) @@ -25,12 +27,17 @@ # add argparse arguments parser = argparse.ArgumentParser(description="Zero agent for Orbit environments.") -parser.add_argument("--cpu", action="store_true", default=False, help="Use CPU pipeline.") -parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument( + "--cpu", action="store_true", default=False, help="Use CPU pipeline." +) +parser.add_argument( + "--num_envs", type=int, default=None, help="Number of environments to simulate." +) parser.add_argument("--task", type=str, default=None, help="Name of the task.") # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) import os + app_experience = f"{os.environ['EXP_PATH']}/omni.isaac.sim.python.gym.headless.kit" # parse the arguments args_cli = parser.parse_args() @@ -41,14 +48,13 @@ """Rest everything follows.""" -import gymnasium as gym -import torch import traceback import carb - +import gymnasium as gym import omni.isaac.contrib_tasks # noqa: F401 import omni.isaac.orbit_tasks # noqa: F401 +import torch from omni.isaac.orbit_tasks.utils import parse_env_cfg @@ -72,11 +78,15 @@ def main(): N = 100 stime = time.time() for i in tqdm.tqdm(range(N)): - actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + actions = ( + 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + ) obs, rew, terminated, truncated, info = env.step(actions) dtime = time.time() - stime FPS = num_envs * N / dtime - print(f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs") + print( + f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs" + ) env.reset(seed=2022) torch.manual_seed(0) @@ -84,14 +94,18 @@ def main(): N = 1000 stime = time.time() for i in tqdm.tqdm(range(N)): - actions = 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + actions = ( + 2 * torch.rand(env.action_space.shape, device=env.unwrapped.device) - 1 + ) obs, rew, terminated, truncated, info = env.step(actions) if i % 200 == 0 and i != 0: env.reset() print("RESET") dtime = time.time() - stime FPS = num_envs * N / dtime - print(f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs with step+reset") + print( + f"{FPS=:0.3f}. {N=} frames in {dtime:0.3f}s with {num_envs} parallel envs with step+reset" + ) # close the simulator env.close() @@ -107,4 +121,4 @@ def main(): raise finally: # close sim app - simulation_app.close() \ No newline at end of file + simulation_app.close() diff --git a/mani_skill2/envs/pick_and_place/pick_cube.py b/mani_skill2/envs/pick_and_place/pick_cube.py index 77bed1c70..8accba4db 100644 --- a/mani_skill2/envs/pick_and_place/pick_cube.py +++ b/mani_skill2/envs/pick_and_place/pick_cube.py @@ -38,11 +38,11 @@ def _initialize_actors(self): xyz[..., 2] = self.cube_half_size[2] qs = [1, 0, 0, 0] if self.obj_init_rot_z: - qs = [] + qs = np.zeros((self.num_envs, 4)) for i in range(self.num_envs): ori = self._episode_rng.uniform(0, 2 * np.pi) q = euler2quat(0, 0, ori) - qs.append(q) + qs[i] = q qs = to_tensor(qs) # to set a batch of poses, use the Pose object or provide a raw tensor obj_pose = Pose.create_from_pq(p=xyz, q=qs) diff --git a/mani_skill2/envs/sapien_env.py b/mani_skill2/envs/sapien_env.py index a1a7f6e57..5559e9836 100644 --- a/mani_skill2/envs/sapien_env.py +++ b/mani_skill2/envs/sapien_env.py @@ -51,8 +51,8 @@ class BaseEnv(gym.Env): control_mode: control mode of the agent. "*" represents all registered controllers, and the action space will be a dict. render_mode: render mode registered in @SUPPORTED_RENDER_MODES. - sim_freq (int): simulation frequency (Hz) - control_freq (int): control frequency (Hz) + sim_freq (int): simulation frequency (Hz). Default is 500 for CPU simulation, 100 for GPU simulation + control_freq (int): control frequency (Hz). Default is 20 for CPU simulation, 20 for GPU simulation renderer (str): type of renderer. "sapien" or "client". renderer_kwargs (dict): kwargs to initialize the renderer. Example kwargs for `SapienRenderer` (renderer_type=='sapien'): @@ -103,8 +103,8 @@ def __init__( reward_mode: str = None, control_mode: str = None, render_mode: str = None, - sim_freq: int = 500, - control_freq: int = 20, + sim_freq: int = None, + control_freq: int = None, renderer: str = "sapien", renderer_kwargs: dict = None, shader_dir: str = "default", @@ -125,9 +125,16 @@ def __init__( sapien.set_cuda_tensor_backend("torch") self.gpu_sim_cfgs = gpu_sim_cfgs - # TODO(jigu): Change to `warning` after lighting in VecEnv is fixed. - # TODO Ms2 set log level. What to do now? - # self._engine.set_log_level(os.getenv("MS2_SIM_LOG_LEVEL", "error")) + if sim_freq is None: + if physx.is_gpu_enabled(): + sim_freq = 100 + else: + sim_freq = 500 + if control_freq is None: + if physx.is_gpu_enabled(): + control_freq = 20 + else: + control_freq = 20 # Create SAPIEN renderer self._renderer_type = renderer @@ -629,7 +636,7 @@ def step(self, action: Union[None, np.ndarray, Dict]): reward = reward[0] if physx.is_gpu_enabled(): - return obs, reward, terminated, torch.Tensor(False), info + return obs, reward, terminated, False, info else: return unbatch(obs, reward, terminated.item(), False, to_numpy(info))