release repository

benchmark/curobo_profile.py

@@ -0,0 +1,188 @@
+#
+# 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
+from typing import Any, Dict, List
+
+# Third Party
+import numpy as np
+import torch
+
+# from geometrout.primitive import Cuboid, Cylinder
+# from geometrout.transform import SE3
+# from robometrics.robot import CollisionSphereConfig, Robot
+from torch.profiler import ProfilerActivity, profile, record_function
+from tqdm import tqdm
+
+# CuRobo
+from curobo.geom.sdf.world import CollisionCheckerType, WorldConfig
+from curobo.geom.types import Mesh
+from curobo.types.math import Pose
+from curobo.types.state import JointState
+from curobo.util.logger import setup_curobo_logger
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
+
+# torch.set_num_threads(8)
+# ttorch.use_deterministic_algorithms(True)
+torch.manual_seed(0)
+
+torch.backends.cudnn.benchmark = True
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+np.random.seed(10)
+# Third Party
+from robometrics.datasets import demo_raw
+
+
+def load_curobo(
+    n_cubes: int, enable_log: bool = False, mesh_mode: bool = False, cuda_graph: bool = False
+):
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
+    robot_cfg["kinematics"]["collision_sphere_buffer"] = -0.015
+    world_cfg = WorldConfig.from_dict(
+        load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
+    ).get_obb_world()
+
+    c_checker = CollisionCheckerType.PRIMITIVE
+    c_cache = {"obb": n_cubes}
+    if mesh_mode:
+        c_checker = CollisionCheckerType.MESH
+        c_cache = {"mesh": n_cubes}
+        world_cfg = world_cfg.get_mesh_world()
+
+    motion_gen_config = MotionGenConfig.load_from_robot_config(
+        robot_cfg,
+        world_cfg,
+        trajopt_tsteps=32,
+        collision_checker_type=c_checker,
+        use_cuda_graph=cuda_graph,
+        collision_cache=c_cache,
+        ee_link_name="panda_hand",
+        position_threshold=0.005,
+        rotation_threshold=0.05,
+        num_ik_seeds=30,
+        num_trajopt_seeds=10,
+        interpolation_dt=0.02,
+        # grad_trajopt_iters=200,
+        store_ik_debug=enable_log,
+        store_trajopt_debug=enable_log,
+    )
+    mg = MotionGen(motion_gen_config)
+    mg.warmup(enable_graph=False)
+    return mg
+
+
+def benchmark_mb(args):
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
+    spheres = robot_cfg["kinematics"]["collision_spheres"]
+    if isinstance(spheres, str):
+        spheres = load_yaml(join_path(get_robot_configs_path(), spheres))["collision_spheres"]
+
+    plan_config = MotionGenPlanConfig(
+        max_attempts=2,
+        enable_graph_attempt=3,
+        enable_finetune_trajopt=True,
+        partial_ik_opt=False,
+        enable_graph=False,
+    )
+    # load dataset:
+
+    file_paths = [demo_raw]
+    all_files = []
+    for file_path in file_paths:
+        all_groups = []
+
+        problems = file_path()
+
+        for key, v in tqdm(problems.items()):
+            # if key not in ["table_under_pick_panda"]:
+            #    continue
+            scene_problems = problems[key]  # [:2]
+            n_cubes = check_problems(scene_problems)
+            mg = load_curobo(n_cubes, False, args.mesh, args.cuda_graph)
+            m_list = []
+            i = 0
+            for problem in tqdm(scene_problems, leave=False):
+                q_start = problem["start"]
+
+                pose = (
+                    problem["goal_pose"]["position_xyz"] + problem["goal_pose"]["quaternion_wxyz"]
+                )
+
+                # reset planner
+                mg.reset(reset_seed=False)
+                if args.mesh:
+                    world = WorldConfig.from_dict(problem["obstacles"]).get_mesh_world()
+
+                else:
+                    world = WorldConfig.from_dict(problem["obstacles"]).get_obb_world()
+
+                mg.update_world(world)
+                # continue
+                # load obstacles
+
+                # run planner
+                start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
+                with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+                    result = mg.plan_single(
+                        start_state,
+                        Pose.from_list(pose),
+                        plan_config,
+                    )
+                print("Exporting the trace..")
+                prof.export_chrome_trace(join_path(args.save_path, args.file_name) + ".json")
+                print(result.success, result.status)
+                exit()
+
+
+def check_problems(all_problems):
+    n_cube = 0
+    for problem in all_problems:
+        cache = WorldConfig.from_dict(problem["obstacles"]).get_obb_world().get_cache_dict()
+        n_cube = max(n_cube, cache["obb"])
+    return n_cube
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--save_path",
+        type=str,
+        default="log/trace",
+        help="path to save file",
+    )
+    parser.add_argument(
+        "--file_name",
+        type=str,
+        default="motion_gen_trace",
+        help="File name prefix to use to save benchmark results",
+    )
+    parser.add_argument(
+        "--mesh",
+        action="store_true",
+        help="When True, converts obstacles to meshes",
+        default=False,
+    )
+    parser.add_argument(
+        "--cuda_graph",
+        action="store_true",
+        help="When True, uses cuda graph during profiing",
+        default=False,
+    )
+
+    args = parser.parse_args()
+
+    setup_curobo_logger("error")
+    benchmark_mb(args)