release repository

benchmark/kinematics_benchmark.py

@@ -0,0 +1,231 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+# Standard Library
+import argparse
+import time
+
+# Third Party
+import numpy as np
+import torch
+
+# CuRobo
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
+from curobo.rollout.arm_base import ArmBase, ArmBaseConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.robot import RobotConfig
+from curobo.util_file import (
+    get_robot_configs_path,
+    get_robot_list,
+    get_task_configs_path,
+    get_world_configs_path,
+    join_path,
+    load_yaml,
+    write_yaml,
+)
+
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+
+def load_curobo(robot_file, world_file):
+    # load curobo arm base?
+
+    world_cfg = load_yaml(join_path(get_world_configs_path(), world_file))
+
+    base_config_data = load_yaml(join_path(get_task_configs_path(), "base_cfg.yml"))
+    graph_config_data = load_yaml(join_path(get_task_configs_path(), "graph.yml"))
+    # base_config_data["constraint"]["self_collision_cfg"]["weight"] = 0.0
+    # if not compute_distance:
+    #    base_config_data["constraint"]["primitive_collision_cfg"]["classify"] = False
+    robot_config_data = load_yaml(join_path(get_robot_configs_path(), robot_file))
+
+    arm_base = ArmBaseConfig.from_dict(
+        robot_config_data["robot_cfg"],
+        graph_config_data["model"],
+        base_config_data["cost"],
+        base_config_data["constraint"],
+        base_config_data["convergence"],
+        base_config_data["world_collision_checker_cfg"],
+        world_cfg,
+    )
+    arm_base = ArmBase(arm_base)
+    return arm_base
+
+
+def bench_collision_curobo(robot_file, world_file, q_test, use_cuda_graph=True):
+    arm_base = load_curobo(robot_file, world_file)
+    # load graph module:
+    tensor_args = TensorDeviceType()
+    q_test = tensor_args.to_device(q_test).unsqueeze(1)
+
+    tensor_args = TensorDeviceType()
+    q_warm = q_test + 0.5
+
+    if not use_cuda_graph:
+        for _ in range(10):
+            out = arm_base.rollout_constraint(q_warm)
+            torch.cuda.synchronize()
+
+        torch.cuda.synchronize()
+
+        st_time = time.time()
+        out = arm_base.rollout_constraint(q_test)
+
+        torch.cuda.synchronize()
+        dt = time.time() - st_time
+    else:
+        q = q_warm.clone()
+
+        g = torch.cuda.CUDAGraph()
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s):
+            for i in range(3):
+                out = arm_base.rollout_constraint(q_warm)
+        torch.cuda.current_stream().wait_stream(s)
+        with torch.cuda.graph(g):
+            out = arm_base.rollout_constraint(q_warm)
+
+        for _ in range(10):
+            q.copy_(q_warm.detach())
+            g.replay()
+            a = out.feasible
+            # print(a)
+            # a = ee_mat.clone()
+        # q_new = torch.rand((b_size, robot_model.get_dof()), **vars(tensor_args))
+
+        torch.cuda.synchronize()
+        st_time = time.time()
+
+        q.copy_(q_test.detach().requires_grad_(False))
+        g.replay()
+        a = out.feasible
+        # print(a)
+        # a = ee_mat.clone()
+        torch.cuda.synchronize()
+        dt = time.time() - st_time
+    return dt
+
+
+def bench_kin_curobo(robot_file, q_test, use_cuda_graph=True, use_coll_spheres=True):
+    robot_data = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    if not use_coll_spheres:
+        robot_data["kinematics"]["collision_spheres"] = None
+        robot_data["kinematics"]["collision_link_names"] = None
+        robot_data["kinematics"]["use_global_cumul"] = False
+        robot_data["kinematics"]["lock_joints"] = {}
+    tensor_args = TensorDeviceType()
+    robot_data["kinematics"]["use_global_cumul"] = False
+
+    robot_cfg = RobotConfig.from_dict(robot_data)
+    robot_model = CudaRobotModel(robot_cfg.kinematics)
+
+    if not use_cuda_graph:
+        for _ in range(10):
+            q = torch.rand((b_size, robot_model.get_dof()), **vars(tensor_args))
+            out = robot_model.forward(q)
+            torch.cuda.synchronize()
+
+        q = torch.rand((b_size, robot_model.get_dof()), **vars(tensor_args))
+        torch.cuda.synchronize()
+
+        st_time = time.time()
+        out = robot_model.forward(q)
+        torch.cuda.synchronize()
+        dt = time.time() - st_time
+    else:
+        q = torch.rand((b_size, robot_model.get_dof()), **vars(tensor_args))
+
+        g = torch.cuda.CUDAGraph()
+        q = q.detach()
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s):
+            for i in range(3):
+                ee_mat, _, _, _, _, _, _ = robot_model.forward(q=q)
+        torch.cuda.current_stream().wait_stream(s)
+        with torch.cuda.graph(g):
+            ee_mat, _, _, _, _, _, _ = robot_model.forward(q=q)
+        q_new = torch.rand((b_size, robot_model.get_dof()), **vars(tensor_args))
+
+        for _ in range(10):
+            q.copy_(q_new.detach().requires_grad_(False))
+            g.replay()
+            # a = ee_mat.clone()
+        q_new = torch.rand((b_size, robot_model.get_dof()), **vars(tensor_args))
+
+        torch.cuda.synchronize()
+        st_time = time.time()
+
+        q.copy_(q_new.detach().requires_grad_(False))
+        g.replay()
+        # a = ee_mat.clone()
+        torch.cuda.synchronize()
+        dt = time.time() - st_time
+    return dt
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--save_path",
+        type=str,
+        default=".",
+        help="path to save file",
+    )
+    parser.add_argument(
+        "--file_name",
+        type=str,
+        default="kinematics",
+        help="File name prefix to use to save benchmark results",
+    )
+
+    args = parser.parse_args()
+    b_list = [1, 10, 100, 1000, 10000]
+
+    robot_list = get_robot_list()
+
+    world_file = "collision_test.yml"
+
+    print("running...")
+    data = {"robot": [], "Kinematics": [], "Collision Checking": [], "Batch Size": []}
+    for robot_file in robot_list:
+        arm_sampler = load_curobo(robot_file, world_file)
+
+        # create a sampler with dof:
+        for b_size in b_list:
+            # sample test configs:
+            q_test = arm_sampler.sample_random_actions(b_size).cpu().numpy()
+
+            dt_cu_cg = bench_collision_curobo(
+                robot_file,
+                world_file,
+                q_test,
+                use_cuda_graph=True,
+            )
+            dt_kin_cg = bench_kin_curobo(
+                robot_file, q_test, use_cuda_graph=True, use_coll_spheres=True
+            )
+
+            data["robot"].append(robot_file)
+            data["Collision Checking"].append(dt_cu_cg)
+            data["Kinematics"].append(dt_kin_cg)
+            data["Batch Size"].append(b_size)
+    write_yaml(data, join_path(args.save_path, args.file_name + ".yml"))
+    try:
+        # Third Party
+        import pandas as pd
+
+        df = pd.DataFrame(data)
+        df.to_csv(join_path(args.save_path, args.file_name + ".csv"))
+    except ImportError:
+        pass