Significantly improved convergence for mesh and cuboid, new ESDF collision.
This commit is contained in:
Balakumar Sundaralingam
305
benchmark/curobo_voxel_profile.py
Normal file
305
benchmark/curobo_voxel_profile.py
Normal file
@@ -0,0 +1,305 @@
|
||||
#
|
||||
# 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
|
||||
from copy import deepcopy
|
||||
from typing import Optional
|
||||
|
||||
# Third Party
|
||||
import numpy as np
|
||||
import torch
|
||||
from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw
|
||||
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 Cuboid
|
||||
from curobo.geom.types import Cuboid as curobo_Cuboid
|
||||
from curobo.geom.types import Mesh, VoxelGrid
|
||||
from curobo.types.base import TensorDeviceType
|
||||
from curobo.types.camera import CameraObservation
|
||||
from curobo.types.math import Pose
|
||||
from curobo.types.robot import RobotConfig
|
||||
from curobo.types.state import JointState
|
||||
from curobo.util.logger import setup_curobo_logger
|
||||
from curobo.util.metrics import CuroboGroupMetrics, CuroboMetrics
|
||||
from curobo.util_file import (
|
||||
get_assets_path,
|
||||
get_robot_configs_path,
|
||||
get_world_configs_path,
|
||||
join_path,
|
||||
load_yaml,
|
||||
write_yaml,
|
||||
)
|
||||
from curobo.wrap.model.robot_world import RobotWorld, RobotWorldConfig
|
||||
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
|
||||
|
||||
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(0)
|
||||
|
||||
|
||||
def load_curobo(
|
||||
n_cubes: int,
|
||||
enable_debug: bool = False,
|
||||
tsteps: int = 30,
|
||||
trajopt_seeds: int = 4,
|
||||
mpinets: bool = False,
|
||||
graph_mode: bool = False,
|
||||
cuda_graph: bool = True,
|
||||
collision_activation_distance: float = 0.025,
|
||||
finetune_dt_scale: float = 1.0,
|
||||
parallel_finetune: bool = True,
|
||||
):
|
||||
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
|
||||
robot_cfg["kinematics"]["collision_sphere_buffer"] = 0.0
|
||||
|
||||
ik_seeds = 30
|
||||
if graph_mode:
|
||||
trajopt_seeds = 4
|
||||
if trajopt_seeds >= 14:
|
||||
ik_seeds = max(100, trajopt_seeds * 2)
|
||||
if mpinets:
|
||||
robot_cfg["kinematics"]["lock_joints"] = {
|
||||
"panda_finger_joint1": 0.025,
|
||||
"panda_finger_joint2": 0.025,
|
||||
}
|
||||
world_cfg = WorldConfig.from_dict(
|
||||
{
|
||||
"voxel": {
|
||||
"base": {
|
||||
"dims": [2.0, 2.0, 3.0],
|
||||
"pose": [0, 0, 0, 1, 0, 0, 0],
|
||||
"voxel_size": 0.01,
|
||||
"feature_dtype": torch.float8_e4m3fn,
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
interpolation_steps = 2000
|
||||
if graph_mode:
|
||||
interpolation_steps = 100
|
||||
robot_cfg_instance = RobotConfig.from_dict(robot_cfg, tensor_args=TensorDeviceType())
|
||||
K = robot_cfg_instance.kinematics.kinematics_config.joint_limits
|
||||
K.position[0, :] -= 0.2
|
||||
K.position[1, :] += 0.2
|
||||
|
||||
motion_gen_config = MotionGenConfig.load_from_robot_config(
|
||||
robot_cfg_instance,
|
||||
world_cfg,
|
||||
trajopt_tsteps=tsteps,
|
||||
collision_checker_type=CollisionCheckerType.VOXEL,
|
||||
use_cuda_graph=cuda_graph,
|
||||
position_threshold=0.005, # 0.5 cm
|
||||
rotation_threshold=0.05,
|
||||
num_ik_seeds=ik_seeds,
|
||||
num_graph_seeds=trajopt_seeds,
|
||||
num_trajopt_seeds=trajopt_seeds,
|
||||
interpolation_dt=0.025,
|
||||
store_ik_debug=enable_debug,
|
||||
store_trajopt_debug=enable_debug,
|
||||
interpolation_steps=interpolation_steps,
|
||||
collision_activation_distance=collision_activation_distance,
|
||||
trajopt_dt=0.25,
|
||||
finetune_dt_scale=finetune_dt_scale,
|
||||
maximum_trajectory_dt=0.16,
|
||||
finetune_trajopt_iters=300,
|
||||
)
|
||||
mg = MotionGen(motion_gen_config)
|
||||
mg.warmup(enable_graph=True, warmup_js_trajopt=False, parallel_finetune=True)
|
||||
# create a ground truth collision checker:
|
||||
config = RobotWorldConfig.load_from_config(
|
||||
robot_cfg_instance,
|
||||
"collision_table.yml",
|
||||
collision_activation_distance=0.0,
|
||||
collision_checker_type=CollisionCheckerType.PRIMITIVE,
|
||||
n_cuboids=50,
|
||||
)
|
||||
robot_world = RobotWorld(config)
|
||||
|
||||
return mg, robot_cfg, robot_world
|
||||
|
||||
|
||||
def benchmark_mb(
|
||||
write_usd=False,
|
||||
save_log=False,
|
||||
write_plot=False,
|
||||
write_benchmark=False,
|
||||
plot_cost=False,
|
||||
override_tsteps: Optional[int] = None,
|
||||
args=None,
|
||||
):
|
||||
# load dataset:
|
||||
graph_mode = args.graph
|
||||
interpolation_dt = 0.02
|
||||
file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][:1]
|
||||
|
||||
enable_debug = save_log or plot_cost
|
||||
all_files = []
|
||||
og_tsteps = 32
|
||||
if override_tsteps is not None:
|
||||
og_tsteps = override_tsteps
|
||||
|
||||
og_trajopt_seeds = 12
|
||||
og_collision_activation_distance = 0.01 # 0.03
|
||||
if args.graph:
|
||||
og_trajopt_seeds = 4
|
||||
for file_path in file_paths:
|
||||
all_groups = []
|
||||
mpinets_data = False
|
||||
problems = file_path()
|
||||
|
||||
for key, v in tqdm(problems.items()):
|
||||
scene_problems = problems[key]
|
||||
m_list = []
|
||||
i = -1
|
||||
ik_fail = 0
|
||||
trajopt_seeds = og_trajopt_seeds
|
||||
tsteps = og_tsteps
|
||||
collision_activation_distance = og_collision_activation_distance
|
||||
finetune_dt_scale = 1.0
|
||||
parallel_finetune = True
|
||||
if "cage_panda" in key:
|
||||
trajopt_seeds = 16
|
||||
finetune_dt_scale = 0.95
|
||||
parallel_finetune = True
|
||||
if "table_under_pick_panda" in key:
|
||||
tsteps = 36
|
||||
trajopt_seeds = 16
|
||||
finetune_dt_scale = 0.95
|
||||
parallel_finetune = True
|
||||
# collision_activation_distance = 0.015
|
||||
|
||||
if "table_pick_panda" in key:
|
||||
collision_activation_distance = 0.005
|
||||
|
||||
if "cubby_task_oriented" in key: # and "merged" not in key:
|
||||
trajopt_seeds = 16
|
||||
finetune_dt_scale = 0.95
|
||||
collision_activation_distance = 0.005
|
||||
parallel_finetune = True
|
||||
if "dresser_task_oriented" in key:
|
||||
trajopt_seeds = 16
|
||||
finetune_dt_scale = 0.95
|
||||
parallel_finetune = True
|
||||
if key in [
|
||||
"tabletop_neutral_start",
|
||||
"merged_cubby_neutral_start",
|
||||
"cubby_neutral_start",
|
||||
"cubby_neutral_goal",
|
||||
"dresser_neutral_start",
|
||||
"tabletop_task_oriented",
|
||||
]:
|
||||
collision_activation_distance = 0.005
|
||||
if "dresser_task_oriented" in list(problems.keys()):
|
||||
mpinets_data = True
|
||||
|
||||
mg, robot_cfg, robot_world = load_curobo(
|
||||
0,
|
||||
enable_debug,
|
||||
tsteps,
|
||||
trajopt_seeds,
|
||||
mpinets_data,
|
||||
graph_mode,
|
||||
not args.disable_cuda_graph,
|
||||
collision_activation_distance=collision_activation_distance,
|
||||
finetune_dt_scale=finetune_dt_scale,
|
||||
parallel_finetune=parallel_finetune,
|
||||
)
|
||||
for problem in tqdm(scene_problems, leave=False):
|
||||
i += 1
|
||||
if problem["collision_buffer_ik"] < 0.0:
|
||||
continue
|
||||
|
||||
plan_config = MotionGenPlanConfig(
|
||||
max_attempts=10,
|
||||
enable_graph_attempt=1,
|
||||
enable_finetune_trajopt=True,
|
||||
partial_ik_opt=False,
|
||||
enable_graph=graph_mode,
|
||||
timeout=60,
|
||||
enable_opt=not graph_mode,
|
||||
parallel_finetune=True,
|
||||
)
|
||||
|
||||
q_start = problem["start"]
|
||||
pose = (
|
||||
problem["goal_pose"]["position_xyz"] + problem["goal_pose"]["quaternion_wxyz"]
|
||||
)
|
||||
|
||||
# reset planner
|
||||
mg.reset(reset_seed=False)
|
||||
world_coll = WorldConfig.from_dict(problem["obstacles"]).get_obb_world()
|
||||
|
||||
robot_world.update_world(world_coll)
|
||||
|
||||
esdf = robot_world.world_model.get_esdf_in_bounding_box(
|
||||
Cuboid(name="base", pose=[0, 0, 0, 1, 0, 0, 0], dims=[2, 2, 3]), voxel_size=0.01
|
||||
)
|
||||
world_voxel_collision = mg.world_coll_checker
|
||||
world_voxel_collision.update_voxel_data(esdf)
|
||||
|
||||
start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
|
||||
for _ in range(2):
|
||||
result = mg.plan_single(
|
||||
start_state,
|
||||
Pose.from_list(pose),
|
||||
plan_config,
|
||||
)
|
||||
print("Profiling...")
|
||||
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
|
||||
torch.cuda.profiler.start()
|
||||
|
||||
result = mg.plan_single(
|
||||
start_state,
|
||||
Pose.from_list(pose),
|
||||
plan_config,
|
||||
)
|
||||
torch.cuda.profiler.stop()
|
||||
|
||||
print("Exporting the trace..")
|
||||
prof.export_chrome_trace("benchmark/log/trace/motion_gen_voxel.json")
|
||||
|
||||
exit()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
setup_curobo_logger("error")
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument(
|
||||
"--graph",
|
||||
action="store_true",
|
||||
help="When True, runs only geometric planner",
|
||||
default=False,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--disable_cuda_graph",
|
||||
action="store_true",
|
||||
help="When True, disable cuda graph during benchmarking",
|
||||
default=False,
|
||||
)
|
||||
args = parser.parse_args()
|
||||
benchmark_mb(
|
||||
save_log=False,
|
||||
write_usd=False,
|
||||
write_plot=False,
|
||||
write_benchmark=False,
|
||||
plot_cost=False,
|
||||
args=args,
|
||||
)
|
||||
Reference in New Issue
Block a user