cleaner benchmarking code

haosulab · Jan 23, 2024 · 7b27eb6 · 7b27eb6
1 parent 18b176c
commit 7b27eb6
Show file tree

Hide file tree

Showing 7 changed files with 158 additions and 63 deletions.
diff --git a/examples/benchmarking/README.md b/examples/benchmarking/README.md
@@ -1,3 +1,11 @@
 # Benchmarking
 
-Code here is used to benchmark difference frameworks/simulators.
+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
+```
diff --git 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()
+    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
@@ -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
@@ -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,35 +53,49 @@ 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
         video_interval = lambda step: step % cfg.recording_interval == 0
         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
@@ -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,26 +78,34 @@ 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)
 
         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()
+        simulation_app.close()