hofee/gen_data_curobo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

release repository

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2023-10-26 04:17:19 -07:00
commit 07e6ccfc91
287 changed files with 70659 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
benchmark/README.md Normal file
View File

@@ -0,0 +1,13 @@
<!--
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.
-->
This folder contains scripts to run the motion planning benchmarks.
Refer to Benchmarks & Profiling instructions in documentation for more information.

727
benchmark/curobo_benchmark.py Normal file
View File

@@ -0,0 +1,727 @@
#
# 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
from copy import deepcopy
from typing import Optional
# Third Party
import numpy as np
import torch
from tqdm import tqdm
# CuRobo
from curobo.geom.sdf.world import CollisionCheckerType, WorldConfig
from curobo.geom.types import Mesh
from curobo.types.base import TensorDeviceType
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_file import (
get_assets_path,
get_robot_configs_path,
get_world_configs_path,
join_path,
load_yaml,
write_yaml,
)
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
# torch.set_num_threads(8)
# torch.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
# torch.backends.cuda.matmul.allow_tf32 = False
# torch.backends.cudnn.allow_tf32 = False
# torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True
np.random.seed(0)
# Standard Library
import argparse
import warnings
from typing import List, Optional
# Third Party
from metrics import CuroboGroupMetrics, CuroboMetrics
from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw
def plot_cost_iteration(cost: torch.Tensor, save_path="cost", title="", log_scale=False):
# Third Party
import matplotlib.pyplot as plt
fig = plt.figure(figsize=(5, 4))
cost = cost.cpu().numpy()
# save to csv:
np.savetxt(save_path + ".csv", cost, delimiter=",")
# if cost.shape[0] > 1:
colormap = plt.cm.winter
plt.gca().set_prop_cycle(plt.cycler("color", colormap(np.linspace(0, 1, cost.shape[0]))))
x = [i for i in range(cost.shape[-1])]
for i in range(cost.shape[0]):
plt.plot(x, cost[i], label="seed_" + str(i))
plt.tight_layout()
# plt.title(title)
plt.xlabel("iteration")
plt.ylabel("cost")
if log_scale:
plt.yscale("log")
plt.grid()
# plt.legend()
plt.tight_layout()
plt.savefig(save_path + ".pdf")
plt.close()
def plot_traj(act_seq: JointState, dt=0.25, title="", save_path="plot.png", sma_filter=False):
# Third Party
import matplotlib.pyplot as plt
fig, ax = plt.subplots(4, 1, figsize=(5, 8), sharex=True)
t_steps = np.linspace(0, act_seq.position.shape[0] * dt, act_seq.position.shape[0])
# compute acceleration from finite difference of velocity:
# act_seq.acceleration = (torch.roll(act_seq.velocity, -1, 0) - act_seq.velocity) / dt
# act_seq.acceleration = ( act_seq.velocity - torch.roll(act_seq.velocity, 1, 0)) / dt
# act_seq.acceleration[0,:] = 0.0
# act_seq.jerk = ( act_seq.acceleration - torch.roll(act_seq.acceleration, 1, 0)) / dt
# act_seq.jerk[0,:] = 0.0
if sma_filter:
kernel = 5
sma = torch.nn.AvgPool1d(kernel_size=kernel, stride=1, padding=2, ceil_mode=False).cuda()
# act_seq.jerk = sma(act_seq.jerk)
# act_seq.acceleration[-1,:] = 0.0
for i in range(act_seq.position.shape[-1]):
ax[0].plot(t_steps, act_seq.position[:, i].cpu(), "-", label=str(i))
# act_seq.velocity[1:-1, i] = sma(act_seq.velocity[:,i].view(1,-1)).squeeze()#@[1:-2]
ax[1].plot(t_steps[: act_seq.velocity.shape[0]], act_seq.velocity[:, i].cpu(), "-")
if sma_filter:
act_seq.acceleration[:, i] = sma(
act_seq.acceleration[:, i].view(1, -1)
).squeeze() # @[1:-2]
ax[2].plot(t_steps[: act_seq.acceleration.shape[0]], act_seq.acceleration[:, i].cpu(), "-")
if sma_filter:
act_seq.jerk[:, i] = sma(act_seq.jerk[:, i].view(1, -1)).squeeze() # @[1:-2]\
ax[3].plot(t_steps[: act_seq.jerk.shape[0]], act_seq.jerk[:, i].cpu(), "-")
ax[0].set_title(title + " dt=" + "{:.3f}".format(dt))
ax[3].set_xlabel("Time(s)")
ax[3].set_ylabel("Jerk rad. s$^{-3}$")
ax[0].set_ylabel("Position rad.")
ax[1].set_ylabel("Velocity rad. s$^{-1}$")
ax[2].set_ylabel("Acceleration rad. s$^{-2}$")
ax[0].grid()
ax[1].grid()
ax[2].grid()
ax[3].grid()
# ax[0].legend(loc="upper right")
ax[0].legend(bbox_to_anchor=(0.5, 1.6), loc="upper center", ncol=4)
plt.tight_layout()
plt.savefig(save_path)
plt.close()
# plt.legend()
def load_curobo(
n_cubes: int,
enable_debug: bool = False,
tsteps: int = 30,
trajopt_seeds: int = 4,
mpinets: bool = False,
graph_mode: bool = True,
mesh_mode: bool = False,
cuda_graph: bool = True,
collision_buffer: float = -0.01,
):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
robot_cfg["kinematics"]["collision_sphere_buffer"] = collision_buffer
robot_cfg["kinematics"]["collision_spheres"] = "spheres/franka_mesh.yml"
robot_cfg["kinematics"]["collision_link_names"].remove("attached_object")
# del robot_cfg["kinematics"]
ik_seeds = 30 # 500
if graph_mode:
trajopt_seeds = 4
if trajopt_seeds >= 14:
ik_seeds = max(100, trajopt_seeds * 4)
if mpinets:
robot_cfg["kinematics"]["lock_joints"] = {
"panda_finger_joint1": 0.025,
"panda_finger_joint2": -0.025,
}
world_cfg = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
).get_obb_world()
interpolation_steps = 2000
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()
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.1
K.position[1, :] += 0.1
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_cfg_instance,
world_cfg,
trajopt_tsteps=tsteps,
collision_checker_type=c_checker,
use_cuda_graph=cuda_graph,
collision_cache=c_cache,
position_threshold=0.005, # 5 mm
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=0.03,
trajopt_dt=0.25,
finetune_dt_scale=1.05, # 1.05,
maximum_trajectory_dt=0.1,
)
mg = MotionGen(motion_gen_config)
mg.warmup(enable_graph=True, warmup_js_trajopt=False)
return mg, robot_cfg
def benchmark_mb(
write_usd=False,
save_log=False,
write_plot=False,
write_benchmark=False,
plot_cost=False,
override_tsteps: Optional[int] = None,
save_kpi=False,
graph_mode=False,
args=None,
):
# load dataset:
interpolation_dt = 0.02
# mpinets_data = True
# if mpinets_data:
file_paths = [motion_benchmaker_raw, mpinets_raw][:]
if args.demo:
file_paths = [demo_raw]
# else:22
# file_paths = [get_mb_dataset_path()][: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
for file_path in file_paths:
all_groups = []
mpinets_data = False
problems = file_path()
if "dresser_task_oriented" in list(problems.keys()):
mpinets_data = True
for key, v in tqdm(problems.items()):
tsteps = og_tsteps
trajopt_seeds = og_trajopt_seeds
if "cage_panda" in key:
trajopt_seeds = 16
# else:
# continue
if "table_under_pick_panda" in key:
tsteps = 44
trajopt_seeds = 28
if "cubby_task_oriented" in key and "merged" not in key:
trajopt_seeds = 16
if "dresser_task_oriented" in key:
trajopt_seeds = 16
scene_problems = problems[key] # [:4] # [:1] # [:20] # [0:10]
n_cubes = check_problems(scene_problems)
# torch.cuda.empty_cache()
mg, robot_cfg = load_curobo(
n_cubes,
enable_debug,
tsteps,
trajopt_seeds,
mpinets_data,
graph_mode,
args.mesh,
not args.disable_cuda_graph,
collision_buffer=args.collision_buffer,
)
m_list = []
i = 0
ik_fail = 0
for problem in tqdm(scene_problems, leave=False):
i += 1
if problem["collision_buffer_ik"] < 0.0:
# print("collision_ik:", problem["collision_buffer_ik"])
continue
# if i != 269: # 226
# continue
plan_config = MotionGenPlanConfig(
max_attempts=100, # 00, # 00, # 100, # 00, # 000,#,00,#00, # 5000,
enable_graph_attempt=3,
enable_finetune_trajopt=True,
partial_ik_opt=False,
enable_graph=graph_mode,
timeout=60,
enable_opt=not graph_mode,
)
# if "table_under_pick_panda" in key:
# plan_config.enable_graph = True
# plan_config.partial_ik_opt = False
q_start = problem["start"]
pose = (
problem["goal_pose"]["position_xyz"] + problem["goal_pose"]["quaternion_wxyz"]
)
problem_name = "d_" + key + "_" + str(i)
# reset planner
mg.reset(reset_seed=False)
if args.mesh:
world = WorldConfig.from_dict(deepcopy(problem["obstacles"])).get_mesh_world()
else:
world = WorldConfig.from_dict(deepcopy(problem["obstacles"])).get_obb_world()
mg.world_coll_checker.clear_cache()
mg.update_world(world)
# continue
# load obstacles
# run planner
start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
goal_pose = Pose.from_list(pose)
result = mg.plan_single(
start_state,
goal_pose,
plan_config,
)
if result.status == "IK Fail":
ik_fail += 1
# rint(plan_config.enable_graph, plan_config.enable_graph_attempt)
problem["solution"] = None
if plan_config.enable_finetune_trajopt:
problem_name = key + "_" + str(i)
else:
problem_name = "noft_" + key + "_" + str(i)
problem_name = "nosw_" + problem_name
if write_usd or save_log:
# CuRobo
from curobo.util.usd_helper import UsdHelper
world.randomize_color(r=[0.5, 0.9], g=[0.2, 0.5], b=[0.0, 0.2])
gripper_mesh = Mesh(
name="target_gripper",
file_path=join_path(
get_assets_path(),
"robot/franka_description/meshes/visual/hand.dae",
),
color=[0.0, 0.8, 0.1, 1.0],
pose=pose,
)
world.add_obstacle(gripper_mesh)
# get costs:
if plot_cost:
dt = 0.5
problem_name = "approx_wolfe_p" + problem_name
if result.optimized_dt is not None:
dt = result.optimized_dt.item()
if "trajopt_result" in result.debug_info:
success = result.success.item()
traj_cost = (
# result.debug_info["trajopt_result"].debug_info["solver"]["cost"][0]
result.debug_info["trajopt_result"].debug_info["solver"]["cost"][-1]
)
# print(traj_cost[0])
traj_cost = torch.cat(traj_cost, dim=-1)
plot_cost_iteration(
traj_cost,
title=problem_name + "_" + str(success) + "_" + str(dt),
save_path=join_path("log/plot/", problem_name + "_cost"),
log_scale=False,
)
if "finetune_trajopt_result" in result.debug_info:
traj_cost = result.debug_info["finetune_trajopt_result"].debug_info[
"solver"
]["cost"][0]
traj_cost = torch.cat(traj_cost, dim=-1)
plot_cost_iteration(
traj_cost,
title=problem_name + "_" + str(success) + "_" + str(dt),
save_path=join_path("log/plot/", problem_name + "_ft_cost"),
log_scale=False,
)
if result.success.item():
# print("GT: ", result.graph_time)
q_traj = result.get_interpolated_plan()
problem["goal_ik"] = q_traj.position.cpu().squeeze().numpy()[-1, :].tolist()
problem["solution"] = {
"position": result.get_interpolated_plan()
.position.cpu()
.squeeze()
.numpy()
.tolist(),
"velocity": result.get_interpolated_plan()
.velocity.cpu()
.squeeze()
.numpy()
.tolist(),
"acceleration": result.get_interpolated_plan()
.acceleration.cpu()
.squeeze()
.numpy()
.tolist(),
"jerk": result.get_interpolated_plan()
.jerk.cpu()
.squeeze()
.numpy()
.tolist(),
"dt": interpolation_dt,
}
# print(problem["solution"]["position"])
# exit()
debug = {
"used_graph": result.used_graph,
"attempts": result.attempts,
"ik_time": result.ik_time,
"graph_time": result.graph_time,
"trajopt_time": result.trajopt_time,
"total_time": result.total_time,
"solve_time": result.solve_time,
"opt_traj": {
"position": result.optimized_plan.position.cpu()
.squeeze()
.numpy()
.tolist(),
"velocity": result.optimized_plan.velocity.cpu()
.squeeze()
.numpy()
.tolist(),
"acceleration": result.optimized_plan.acceleration.cpu()
.squeeze()
.numpy()
.tolist(),
"jerk": result.optimized_plan.jerk.cpu().squeeze().numpy().tolist(),
"dt": result.optimized_dt.item(),
},
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
# print(
# "T: ",
# result.motion_time.item(),
# result.optimized_dt.item(),
# (len(problem["solution"]["position"]) - 1 ) * result.interpolation_dt,
# result.interpolation_dt,
# )
# exit()
reached_pose = mg.compute_kinematics(result.optimized_plan[-1]).ee_pose
rot_error = goal_pose.angular_distance(reached_pose) * 100.0
if args.graph:
solve_time = result.graph_time
else:
solve_time = result.solve_time
current_metrics = CuroboMetrics(
skip=False,
success=True,
time=result.total_time,
collision=False,
joint_limit_violation=False,
self_collision=False,
position_error=result.position_error.item() * 1000.0,
orientation_error=rot_error.item(),
eef_position_path_length=10,
eef_orientation_path_length=10,
attempts=result.attempts,
motion_time=result.motion_time.item(),
solve_time=solve_time,
)
if write_usd:
# CuRobo
q_traj = result.get_interpolated_plan()
UsdHelper.write_trajectory_animation_with_robot_usd(
robot_cfg,
world,
start_state,
q_traj,
dt=result.interpolation_dt,
save_path=join_path("log/usd/", problem_name) + ".usd",
interpolation_steps=1,
write_robot_usd_path="log/usd/assets/",
robot_usd_local_reference="assets/",
base_frame="/world_" + problem_name,
visualize_robot_spheres=True,
)
if write_plot:
problem_name = problem_name
plot_traj(
result.optimized_plan,
result.optimized_dt.item(),
# result.get_interpolated_plan(),
# result.interpolation_dt,
title=problem_name,
save_path=join_path("log/plot/", problem_name + ".pdf"),
)
plot_traj(
# result.optimized_plan,
# result.optimized_dt.item(),
result.get_interpolated_plan(),
result.interpolation_dt,
title=problem_name,
save_path=join_path("log/plot/", problem_name + "_int.pdf"),
)
# exit()
m_list.append(current_metrics)
all_groups.append(current_metrics)
elif result.valid_query:
# print("fail")
# print(result.status)
current_metrics = CuroboMetrics()
debug = {
"used_graph": result.used_graph,
"attempts": result.attempts,
"ik_time": result.ik_time,
"graph_time": result.graph_time,
"trajopt_time": result.trajopt_time,
"total_time": result.total_time,
"solve_time": result.solve_time,
"status": result.status,
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
m_list.append(current_metrics)
all_groups.append(current_metrics)
else:
# print("invalid: " + problem_name)
debug = {
"used_graph": result.used_graph,
"attempts": result.attempts,
"ik_time": result.ik_time,
"graph_time": result.graph_time,
"trajopt_time": result.trajopt_time,
"total_time": result.total_time,
"solve_time": result.solve_time,
"status": result.status,
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
if False:
world.save_world_as_mesh(problem_name + ".obj")
q_traj = start_state # .unsqueeze(0)
# CuRobo
from curobo.util.usd_helper import UsdHelper
UsdHelper.write_trajectory_animation_with_robot_usd(
robot_cfg,
world,
start_state,
q_traj,
dt=result.interpolation_dt,
save_path=join_path("log/usd/", problem_name) + ".usd",
interpolation_steps=1,
write_robot_usd_path="log/usd/assets/",
robot_usd_local_reference="assets/",
base_frame="/world_" + problem_name,
visualize_robot_spheres=True,
)
if save_log: # and not result.success.item():
# print("save log")
UsdHelper.write_motion_gen_log(
result,
robot_cfg,
world,
start_state,
Pose.from_list(pose),
join_path("log/usd/", problem_name) + "_debug",
write_ik=False,
write_trajopt=True,
visualize_robot_spheres=False,
grid_space=2,
)
# exit()
g_m = CuroboGroupMetrics.from_list(m_list)
if not args.kpi:
print(
key,
f"{g_m.success:2.2f}",
# g_m.motion_time,
g_m.time.mean,
# g_m.time.percent_75,
g_m.time.percent_98,
g_m.position_error.percent_98,
# g_m.position_error.median,
g_m.orientation_error.percent_98,
# g_m.orientation_error.median,
) # , g_m.attempts)
print(g_m.attempts)
# print("MT: ", g_m.motion_time)
# $print(ik_fail)
# exit()
# print(g_m.position_error, g_m.orientation_error)
g_m = CuroboGroupMetrics.from_list(all_groups)
if not args.kpi:
print(
"All: ",
f"{g_m.success:2.2f}",
g_m.motion_time.percent_98,
g_m.time.mean,
g_m.time.percent_75,
g_m.position_error.percent_75,
g_m.orientation_error.percent_75,
) # g_m.time, g_m.attempts)
# print("MT: ", g_m.motion_time)
# print(g_m.position_error, g_m.orientation_error)
# exit()
if write_benchmark:
if not mpinets_data:
write_yaml(problems, args.file_name + "_mb_solution.yaml")
else:
write_yaml(problems, args.file_name + "_mpinets_solution.yaml")
all_files += all_groups
g_m = CuroboGroupMetrics.from_list(all_files)
# print(g_m.success, g_m.time, g_m.attempts, g_m.position_error, g_m.orientation_error)
print("######## FULL SET ############")
print("All: ", f"{g_m.success:2.2f}")
print("MT: ", g_m.motion_time)
print("PT:", g_m.time)
print("ST: ", g_m.solve_time)
print("position error (mm): ", g_m.position_error)
print("orientation error(%): ", g_m.orientation_error)
if args.kpi:
kpi_data = {
"Success": g_m.success,
"Planning Time Mean": float(g_m.time.mean),
"Planning Time Std": float(g_m.time.std),
"Planning Time Median": float(g_m.time.median),
"Planning Time 75th": float(g_m.time.percent_75),
"Planning Time 98th": float(g_m.time.percent_98),
}
write_yaml(kpi_data, join_path(args.save_path, args.file_name + ".yml"))
# run on mb dataset:
def check_problems(all_problems):
n_cube = 0
for problem in all_problems:
cache = (
WorldConfig.from_dict(deepcopy(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=".",
help="path to save file",
)
parser.add_argument(
"--file_name",
type=str,
default="mg_curobo_",
help="File name prefix to use to save benchmark results",
)
parser.add_argument(
"--collision_buffer",
type=float,
default=-0.00, # in meters
help="Robot collision buffer",
)
parser.add_argument(
"--graph",
action="store_true",
help="When True, runs only geometric planner",
default=False,
)
parser.add_argument(
"--mesh",
action="store_true",
help="When True, converts obstacles to meshes",
default=False,
)
parser.add_argument(
"--kpi",
action="store_true",
help="When True, saves minimal metrics",
default=False,
)
parser.add_argument(
"--demo",
action="store_true",
help="When True, runs only on small dataaset",
default=False,
)
parser.add_argument(
"--disable_cuda_graph",
action="store_true",
help="When True, disable cuda graph during benchmarking",
default=False,
)
parser.add_argument(
"--write_benchmark",
action="store_true",
help="When True, writes paths to file",
default=False,
)
args = parser.parse_args()
setup_curobo_logger("error")
benchmark_mb(
save_log=False,
write_usd=False,
write_plot=False,
write_benchmark=args.write_benchmark,
plot_cost=False,
save_kpi=args.kpi,
graph_mode=args.graph,
args=args,
)

558
benchmark/curobo_nvblox_benchmark.py Normal file
View File

@@ -0,0 +1,558 @@
#
# 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 matplotlib.pyplot as plt
import numpy as np
import torch
from metrics import CuroboGroupMetrics, CuroboMetrics
from nvblox_torch.datasets.mesh_dataset import MeshDataset
from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw
from tqdm import tqdm
# CuRobo
from curobo.geom.sdf.world import CollisionCheckerType, WorldConfig
from curobo.geom.types import Mesh
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_file import (
get_assets_path,
get_robot_configs_path,
get_world_configs_path,
join_path,
load_yaml,
write_yaml,
)
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 plot_cost_iteration(cost: torch.Tensor, save_path="cost", title="", log_scale=False):
fig = plt.figure(figsize=(5, 4))
cost = cost.cpu().numpy()
# save to csv:
np.savetxt(save_path + ".csv", cost, delimiter=",")
# if cost.shape[0] > 1:
colormap = plt.cm.winter
plt.gca().set_prop_cycle(plt.cycler("color", colormap(np.linspace(0, 1, cost.shape[0]))))
x = [i for i in range(cost.shape[-1])]
for i in range(cost.shape[0]):
plt.plot(x, cost[i], label="seed_" + str(i))
plt.tight_layout()
# plt.title(title)
plt.xlabel("iteration")
plt.ylabel("cost")
if log_scale:
plt.yscale("log")
plt.grid()
# plt.legend()
plt.tight_layout()
plt.savefig(save_path + ".pdf")
plt.close()
def plot_traj(act_seq: JointState, dt=0.25, title="", save_path="plot.png", sma_filter=False):
fig, ax = plt.subplots(4, 1, figsize=(5, 8), sharex=True)
t_steps = np.linspace(0, act_seq.position.shape[0] * dt, act_seq.position.shape[0])
if sma_filter:
kernel = 5
sma = torch.nn.AvgPool1d(kernel_size=kernel, stride=1, padding=2, ceil_mode=False).cuda()
for i in range(act_seq.position.shape[-1]):
ax[0].plot(t_steps, act_seq.position[:, i].cpu(), "-", label=str(i))
ax[1].plot(t_steps[: act_seq.velocity.shape[0]], act_seq.velocity[:, i].cpu(), "-")
if sma_filter:
act_seq.acceleration[:, i] = sma(
act_seq.acceleration[:, i].view(1, -1)
).squeeze() # @[1:-2]
ax[2].plot(t_steps[: act_seq.acceleration.shape[0]], act_seq.acceleration[:, i].cpu(), "-")
if sma_filter:
act_seq.jerk[:, i] = sma(act_seq.jerk[:, i].view(1, -1)).squeeze() # @[1:-2]\
ax[3].plot(t_steps[: act_seq.jerk.shape[0]], act_seq.jerk[:, i].cpu(), "-")
ax[0].set_title(title + " dt=" + "{:.3f}".format(dt))
ax[3].set_xlabel("Time(s)")
ax[3].set_ylabel("Jerk rad. s$^{-3}$")
ax[0].set_ylabel("Position rad.")
ax[1].set_ylabel("Velocity rad. s$^{-1}$")
ax[2].set_ylabel("Acceleration rad. s$^{-2}$")
ax[0].grid()
ax[1].grid()
ax[2].grid()
ax[3].grid()
ax[0].legend(bbox_to_anchor=(0.5, 1.6), loc="upper center", ncol=4)
plt.tight_layout()
plt.savefig(save_path)
plt.close()
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,
):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
robot_cfg["kinematics"]["collision_sphere_buffer"] = -0.015
ik_seeds = 30 # 500
if graph_mode:
trajopt_seeds = 4
if trajopt_seeds >= 14:
ik_seeds = max(100, trajopt_seeds * 4)
if mpinets:
robot_cfg["kinematics"]["lock_joints"] = {
"panda_finger_joint1": 0.025,
"panda_finger_joint2": -0.025,
}
world_cfg = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_nvblox_online.yml"))
)
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.1
K.position[1, :] += 0.1
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_cfg_instance,
world_cfg,
trajopt_tsteps=tsteps,
collision_checker_type=CollisionCheckerType.BLOX,
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=0.025,
state_finite_difference_mode="CENTRAL",
trajopt_dt=0.25,
minimize_jerk=True,
finetune_dt_scale=1.05,
maximum_trajectory_dt=0.1,
)
mg = MotionGen(motion_gen_config)
mg.warmup(enable_graph=True, warmup_js_trajopt=False)
return mg, robot_cfg
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][2:]
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
if args.graph:
og_trajopt_seeds = 4
for file_path in file_paths:
all_groups = []
mpinets_data = False
problems = file_path()
if "dresser_task_oriented" in list(problems.keys()):
mpinets_data = True
mg, robot_cfg = load_curobo(
1,
enable_debug,
og_tsteps,
og_trajopt_seeds,
mpinets_data,
graph_mode,
not args.disable_cuda_graph,
)
for key, v in tqdm(problems.items()):
scene_problems = problems[key][:] # [:1] # [:20] # [0:10]
m_list = []
i = 0
ik_fail = 0
for problem in tqdm(scene_problems, leave=False):
i += 1
plan_config = MotionGenPlanConfig(
max_attempts=10, # 00, # 00, # 100, # 00, # 000,#,00,#00, # 5000,
enable_graph_attempt=3,
enable_finetune_trajopt=True,
partial_ik_opt=False,
enable_graph=graph_mode,
timeout=60,
enable_opt=not graph_mode,
)
q_start = problem["start"]
pose = (
problem["goal_pose"]["position_xyz"] + problem["goal_pose"]["quaternion_wxyz"]
)
problem_name = "d_" + key + "_" + str(i)
# reset planner
mg.reset(reset_seed=False)
world = WorldConfig.from_dict(deepcopy(problem["obstacles"])).get_mesh_world(
merge_meshes=True
)
mesh = world.mesh[0].get_trimesh_mesh()
# world.save_world_as_mesh(problem_name + ".stl")
mg.world_coll_checker.clear_cache()
m_dataset = MeshDataset(
None, n_frames=200, image_size=640, save_data_dir=None, trimesh_mesh=mesh
)
tensor_args = mg.tensor_args
for j in range(len(m_dataset)):
data = m_dataset[j]
cam_obs = CameraObservation(
rgb_image=tensor_args.to_device(data["rgba"]),
depth_image=tensor_args.to_device(data["depth"]),
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=mg.tensor_args.device)),
)
cam_obs = cam_obs
mg.add_camera_frame(cam_obs, "world")
mg.process_camera_frames("world", False)
torch.cuda.synchronize()
mg.world_coll_checker.update_blox_hashes()
torch.cuda.synchronize()
if save_log or write_usd:
# nvblox_obs = mg.world_coll_checker.get_mesh_from_blox_layer("world", mode="nvblox")
# nvblox_obs.save_as_mesh("debug_tsdf.obj")
nvblox_obs = mg.world_coll_checker.get_mesh_from_blox_layer(
"world", mode="voxel"
)
nvblox_obs.color = [0.0, 0.0, 0.8, 0.8]
# nvblox_obs.save_as_mesh("debug_voxel_occ.obj")
# exit()
nvblox_obs.name = "nvblox_world"
world.add_obstacle(nvblox_obs)
# run planner
start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
result = mg.plan_single(
start_state,
Pose.from_list(pose),
plan_config,
)
if result.status == "IK Fail":
ik_fail += 1
problem["solution"] = None
if plan_config.enable_finetune_trajopt:
problem_name = key + "_" + str(i)
else:
problem_name = "noft_" + key + "_" + str(i)
problem_name = "nvblox_" + problem_name
if write_usd or save_log:
# CuRobo
from curobo.util.usd_helper import UsdHelper
world.randomize_color(r=[0.5, 0.9], g=[0.2, 0.5], b=[0.0, 0.2])
if len(world.mesh) > 1:
world.mesh[1].color = [125 / 255, 255 / 255, 70.0 / 255, 1.0]
gripper_mesh = Mesh(
name="target_gripper",
file_path=join_path(
get_assets_path(),
"robot/franka_description/meshes/visual/hand.dae",
),
color=[0.0, 0.8, 0.1, 1.0],
pose=pose,
)
world.add_obstacle(gripper_mesh)
# get costs:
if plot_cost:
dt = 0.5
problem_name = "approx_wolfe_p" + problem_name
if result.optimized_dt is not None:
dt = result.optimized_dt.item()
if "trajopt_result" in result.debug_info:
success = result.success.item()
traj_cost = result.debug_info["trajopt_result"].debug_info["solver"][
"cost"
][-1]
traj_cost = torch.cat(traj_cost, dim=-1)
plot_cost_iteration(
traj_cost,
title=problem_name + "_" + str(success) + "_" + str(dt),
save_path=join_path("log/plot/", problem_name + "_cost"),
log_scale=False,
)
if "finetune_trajopt_result" in result.debug_info:
traj_cost = result.debug_info["finetune_trajopt_result"].debug_info[
"solver"
]["cost"][0]
traj_cost = torch.cat(traj_cost, dim=-1)
plot_cost_iteration(
traj_cost,
title=problem_name + "_" + str(success) + "_" + str(dt),
save_path=join_path("log/plot/", problem_name + "_ft_cost"),
log_scale=False,
)
if result.success.item():
q_traj = result.get_interpolated_plan()
problem["goal_ik"] = q_traj.position.cpu().squeeze().numpy()[-1, :].tolist()
problem["solution"] = {
"position": result.get_interpolated_plan()
.position.cpu()
.squeeze()
.numpy()
.tolist(),
"velocity": result.get_interpolated_plan()
.velocity.cpu()
.squeeze()
.numpy()
.tolist(),
"acceleration": result.get_interpolated_plan()
.acceleration.cpu()
.squeeze()
.numpy()
.tolist(),
"jerk": result.get_interpolated_plan()
.jerk.cpu()
.squeeze()
.numpy()
.tolist(),
"dt": interpolation_dt,
}
debug = {
"used_graph": result.used_graph,
"attempts": result.attempts,
"ik_time": result.ik_time,
"graph_time": result.graph_time,
"trajopt_time": result.trajopt_time,
"total_time": result.total_time,
"solve_time": result.solve_time,
"opt_traj": {
"position": result.optimized_plan.position.cpu()
.squeeze()
.numpy()
.tolist(),
"velocity": result.optimized_plan.velocity.cpu()
.squeeze()
.numpy()
.tolist(),
"acceleration": result.optimized_plan.acceleration.cpu()
.squeeze()
.numpy()
.tolist(),
"jerk": result.optimized_plan.jerk.cpu().squeeze().numpy().tolist(),
"dt": result.optimized_dt.item(),
},
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
current_metrics = CuroboMetrics(
skip=False,
success=True,
time=result.total_time,
collision=False,
joint_limit_violation=False,
self_collision=False,
position_error=result.position_error.item() * 100.0,
orientation_error=result.rotation_error.item() * 100.0,
eef_position_path_length=10,
eef_orientation_path_length=10,
attempts=result.attempts,
motion_time=result.motion_time.item(),
solve_time=result.solve_time,
)
if write_usd:
# CuRobo
q_traj = result.get_interpolated_plan()
UsdHelper.write_trajectory_animation_with_robot_usd(
robot_cfg,
world,
start_state,
q_traj,
dt=result.interpolation_dt,
save_path=join_path("log/usd/", problem_name) + ".usd",
interpolation_steps=1,
write_robot_usd_path="log/usd/assets/",
robot_usd_local_reference="assets/",
base_frame="/world_" + problem_name,
visualize_robot_spheres=True,
)
if write_plot:
problem_name = problem_name
plot_traj(
result.optimized_plan,
result.optimized_dt.item(),
# result.get_interpolated_plan(),
# result.interpolation_dt,
title=problem_name,
save_path=join_path("log/plot/", problem_name + ".pdf"),
)
plot_traj(
# result.optimized_plan,
# result.optimized_dt.item(),
result.get_interpolated_plan(),
result.interpolation_dt,
title=problem_name,
save_path=join_path("log/plot/", problem_name + "_int.pdf"),
)
m_list.append(current_metrics)
all_groups.append(current_metrics)
elif result.valid_query:
current_metrics = CuroboMetrics()
debug = {
"used_graph": result.used_graph,
"attempts": result.attempts,
"ik_time": result.ik_time,
"graph_time": result.graph_time,
"trajopt_time": result.trajopt_time,
"total_time": result.total_time,
"solve_time": result.solve_time,
"status": result.status,
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
m_list.append(current_metrics)
all_groups.append(current_metrics)
else:
print("invalid: " + problem_name)
debug = {
"used_graph": result.used_graph,
"attempts": result.attempts,
"ik_time": result.ik_time,
"graph_time": result.graph_time,
"trajopt_time": result.trajopt_time,
"total_time": result.total_time,
"solve_time": result.solve_time,
"status": result.status,
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
if save_log:
UsdHelper.write_motion_gen_log(
result,
robot_cfg,
world,
start_state,
Pose.from_list(pose),
join_path("log/usd/", problem_name) + "_debug",
write_ik=not result.success.item(),
write_trajopt=True,
visualize_robot_spheres=True,
grid_space=2,
)
g_m = CuroboGroupMetrics.from_list(m_list)
print(
key,
f"{g_m.success:2.2f}",
g_m.time.mean,
g_m.time.percent_98,
g_m.position_error.percent_98,
g_m.orientation_error.percent_98,
)
g_m = CuroboGroupMetrics.from_list(all_groups)
print(
"All: ",
f"{g_m.success:2.2f}",
g_m.motion_time.percent_98,
g_m.time.mean,
g_m.time.percent_75,
g_m.position_error.percent_75,
g_m.orientation_error.percent_75,
)
print(g_m.attempts)
if write_benchmark:
if not mpinets_data:
write_yaml(problems, "mb_curobo_solution_nvblox.yaml")
else:
write_yaml(problems, "mpinets_curobo_solution_nvblox.yaml")
all_files += all_groups
g_m = CuroboGroupMetrics.from_list(all_files)
print("######## FULL SET ############")
print("All: ", f"{g_m.success:2.2f}")
print("MT: ", g_m.motion_time)
print("PT:", g_m.time)
print("ST: ", g_m.solve_time)
print("accuracy: ", g_m.position_error, g_m.orientation_error)
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,
)

193
benchmark/curobo_nvblox_profile.py Normal file
View File

@@ -0,0 +1,193 @@
#
# 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 time
from typing import Any, Dict, List
# Third Party
import numpy as np
import torch
import torch.autograd.profiler as profiler
from robometrics.datasets import demo_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 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_assets_path,
get_robot_configs_path,
get_task_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 nvblox_torch.datasets.mesh_dataset import MeshDataset
# CuRobo
from curobo.types.camera import CameraObservation
def load_curobo(n_cubes: int, enable_log: bool = False):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
robot_cfg["kinematics"]["collision_sphere_buffer"] = -0.015
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_cfg,
"collision_nvblox_online.yml",
trajopt_tsteps=32,
collision_checker_type=CollisionCheckerType.BLOX,
use_cuda_graph=False,
position_threshold=0.005,
rotation_threshold=0.05,
num_ik_seeds=30,
num_trajopt_seeds=12,
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)
# print("warmed up")
# exit()
return mg
def benchmark_mb(write_usd=False, save_log=False):
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, save_log)
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)
world = WorldConfig.from_dict(problem["obstacles"]).get_mesh_world(
merge_meshes=True
)
# clear cache:
mesh = world.mesh[0].get_trimesh_mesh()
mg.clear_world_cache()
obs = []
# get camera_observations:
m_dataset = MeshDataset(
None, n_frames=200, image_size=640, save_data_dir=None, trimesh_mesh=mesh
)
obs = []
tensor_args = mg.tensor_args
for j in range(len(m_dataset)):
with profiler.record_function("nvblox/create_camera_images"):
data = m_dataset[j]
cam_obs = CameraObservation(
rgb_image=tensor_args.to_device(data["rgba"]),
depth_image=tensor_args.to_device(data["depth"]),
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=mg.tensor_args.device)),
)
obs.append(cam_obs)
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
for j in range(len(obs)):
cam_obs = obs[j]
cam_obs.rgb_image = None
with profiler.record_function("nvblox/add_camera_images"):
mg.add_camera_frame(cam_obs, "world")
with profiler.record_function("nvblox/process_camera_images"):
mg.process_camera_frames("world", False)
mg.world_coll_checker.update_blox_hashes()
# run planner
start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
result = mg.plan_single(
start_state,
Pose.from_list(pose),
plan_config,
)
print("Exporting the trace..")
prof.export_chrome_trace("log/trace/trajopt_global_nvblox.json")
print(result.success, result.status)
exit()
def get_metrics_obstacles(obs: Dict[str, List[Any]]):
obs_list = []
if "cylinder" in obs and len(obs["cylinder"].items()) > 0:
for _, vi in enumerate(obs["cylinder"].values()):
obs_list.append(
Cylinder(
np.ravel(vi["pose"][:3]), vi["radius"], vi["height"], np.ravel(vi["pose"][3:])
)
)
if "cuboid" in obs and len(obs["cuboid"].items()) > 0:
for _, vi in enumerate(obs["cuboid"].values()):
obs_list.append(
Cuboid(np.ravel(vi["pose"][:3]), np.ravel(vi["dims"]), np.ravel(vi["pose"][3:]))
)
return obs_list
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__":
setup_curobo_logger("error")
benchmark_mb()

188
benchmark/curobo_profile.py Normal file
View File

@@ -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)

180
benchmark/curobo_python_profile.py Normal file
View File

@@ -0,0 +1,180 @@
#
# 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 cProfile
import time
# Third Party
import torch
from torch.profiler import ProfilerActivity, profile, record_function
# CuRobo
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
from curobo.geom.sdf.world import CollisionCheckerType
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import JointState, RobotConfig
from curobo.util_file import get_robot_configs_path, get_robot_path, join_path, load_yaml
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
def demo_motion_gen():
# Standard Library
st_time = time.time()
tensor_args = TensorDeviceType()
world_file = "collision_table.yml"
robot_file = "franka.yml"
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_file,
world_file,
tensor_args,
trajopt_tsteps=32,
collision_checker_type=CollisionCheckerType.PRIMITIVE,
use_cuda_graph=False,
num_trajopt_seeds=4,
num_graph_seeds=1,
evaluate_interpolated_trajectory=True,
interpolation_dt=0.01,
)
motion_gen = MotionGen(motion_gen_config)
# st_time = time.time()
motion_gen.warmup(batch=50, enable_graph=False, warmup_js_trajopt=False)
print("motion gen time:", time.time() - st_time)
# print(time.time() - st_time)
return
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
retract_cfg = motion_gen.get_retract_config()
print(retract_cfg)
state = motion_gen.rollout_fn.compute_kinematics(
JointState.from_position(retract_cfg.view(1, -1))
)
retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
result = motion_gen.plan(
start_state, retract_pose, enable_graph=True, enable_opt=False, max_attempts=1
)
print(result.optimized_plan.position.shape)
traj = result.get_interpolated_plan() # $.position.view(-1, 7) # optimized plan
print("Trajectory Generated: ", result.success, result.optimized_dt.item())
def demo_basic_robot():
st_time = time.time()
tensor_args = TensorDeviceType()
# load a urdf:
config_file = load_yaml(join_path(get_robot_path(), "franka.yml"))
urdf_file = config_file["robot_cfg"]["kinematics"][
"urdf_path"
] # Send global path starting with "/"
base_link = config_file["robot_cfg"]["kinematics"]["base_link"]
ee_link = config_file["robot_cfg"]["kinematics"]["ee_link"]
robot_cfg = RobotConfig.from_basic(urdf_file, base_link, ee_link, tensor_args)
kin_model = CudaRobotModel(robot_cfg.kinematics)
print("base kin time:", time.time() - st_time)
return
# compute forward kinematics:
# q = torch.rand((10, kin_model.get_dof()), **vars(tensor_args))
# out = kin_model.get_state(q)
# here is the kinematics state:
# print(out)
def demo_full_config_robot(config_file):
st_time = time.time()
tensor_args = TensorDeviceType()
# load a urdf:
robot_cfg = RobotConfig.from_dict(config_file, tensor_args)
# kin_model = CudaRobotModel(robot_cfg.kinematics)
print("full kin time: ", time.time() - st_time)
# compute forward kinematics:
# q = torch.rand((10, kin_model.get_dof()), **vars(tensor_args))
# out = kin_model.get_state(q)
# here is the kinematics state:
# print(out)
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="startup_trace",
help="File name prefix to use to save benchmark results",
)
parser.add_argument(
"--motion_gen",
action="store_true",
help="When True, runs startup for motion generation",
default=False,
)
parser.add_argument(
"--kinematics",
action="store_true",
help="When True, runs startup for kinematics",
default=True,
)
args = parser.parse_args()
# cProfile.run('demo_motion_gen()')
config_file = load_yaml(join_path(get_robot_path(), "franka.yml"))["robot_cfg"]
# Third Party
if args.kinematics:
for _ in range(5):
demo_full_config_robot(config_file)
pr = cProfile.Profile()
pr.enable()
demo_full_config_robot(config_file)
pr.disable()
filename = join_path(args.save_path, args.file_name) + "_kinematics_cprofile.prof"
pr.dump_stats(filename)
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
demo_full_config_robot(config_file)
filename = join_path(args.save_path, args.file_name) + "_kinematics_trace.json"
prof.export_chrome_trace(filename)
if args.motion_gen:
for _ in range(5):
demo_motion_gen()
pr = cProfile.Profile()
pr.enable()
demo_motion_gen()
pr.disable()
filename = join_path(args.save_path, args.file_name) + "_motion_gen_cprofile.prof"
pr.dump_stats(filename)
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
demo_motion_gen()
filename = join_path(args.save_path, args.file_name) + "_motion_gen_trace.json"
prof.export_chrome_trace(filename)

182
benchmark/ik_benchmark.py Normal file
View File

@@ -0,0 +1,182 @@
#
# 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.geom.types import WorldConfig
from curobo.rollout.arm_base import ArmBase, ArmBaseConfig
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import RobotConfig
from curobo.util_file import (
get_motion_gen_robot_list,
get_robot_configs_path,
get_robot_list,
get_task_configs_path,
get_world_configs_path,
join_path,
load_yaml,
write_yaml,
)
from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
def run_full_config_collision_free_ik(
robot_file,
world_file,
batch_size,
use_cuda_graph=False,
collision_free=True,
high_precision=False,
):
tensor_args = TensorDeviceType()
robot_data = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
if not collision_free:
robot_data["kinematics"]["collision_link_names"] = None
robot_data["kinematics"]["lock_joints"] = {}
robot_cfg = RobotConfig.from_dict(robot_data)
world_cfg = WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
position_threshold = 0.005
if high_precision:
position_threshold = 0.001
ik_config = IKSolverConfig.load_from_robot_config(
robot_cfg,
world_cfg,
rotation_threshold=0.05,
position_threshold=position_threshold,
num_seeds=30,
self_collision_check=collision_free,
self_collision_opt=collision_free,
tensor_args=tensor_args,
use_cuda_graph=use_cuda_graph,
high_precision=high_precision,
regularization=False,
# grad_iters=500,
)
ik_solver = IKSolver(ik_config)
for _ in range(3):
q_sample = ik_solver.sample_configs(batch_size)
while q_sample.shape[0] == 0:
q_sample = ik_solver.sample_configs(batch_size)
kin_state = ik_solver.fk(q_sample)
goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
st_time = time.time()
result = ik_solver.solve_batch(goal)
torch.cuda.synchronize()
total_time = (time.time() - st_time) / q_sample.shape[0]
return (
total_time,
100.0 * torch.count_nonzero(result.success).item() / len(q_sample),
np.percentile(result.position_error[result.success].cpu().numpy(), 98).item(),
np.percentile(result.rotation_error[result.success].cpu().numpy(), 98).item(),
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--save_path",
type=str,
default=".",
help="path to save file",
)
parser.add_argument(
"--high_precision",
action="store_true",
help="When True, enables IK for 1 mm precision, when False 5mm precision",
default=False,
)
parser.add_argument(
"--file_name",
type=str,
default="ik",
help="File name prefix to use to save benchmark results",
)
args = parser.parse_args()
b_list = [1, 10, 100, 1000][-1:]
robot_list = get_motion_gen_robot_list()
world_file = "collision_test.yml"
print("running...")
data = {
"robot": [],
"IK time(ms)": [],
"Collision-Free IK time(ms)": [],
"Batch Size": [],
"Success IK": [],
"Success Collision-Free IK": [],
"Position Error(mm)": [],
"Orientation Error": [],
"Position Error Collision-Free IK(mm)": [],
"Orientation Error Collision-Free IK": [],
}
for robot_file in robot_list:
# create a sampler with dof:
for b_size in b_list:
# sample test configs:
dt_cu_ik, succ, p_err, q_err = run_full_config_collision_free_ik(
robot_file,
world_file,
batch_size=b_size,
use_cuda_graph=True,
collision_free=False,
high_precision=args.high_precision,
)
dt_cu_ik_cfree, success, p_err_c, q_err_c = run_full_config_collision_free_ik(
robot_file,
world_file,
batch_size=b_size,
use_cuda_graph=True,
collision_free=True,
)
# print(dt_cu/b_size, dt_cu_cg/b_size)
data["robot"].append(robot_file)
data["IK time(ms)"].append(dt_cu_ik * 1000.0)
data["Batch Size"].append(b_size)
data["Success Collision-Free IK"].append(success)
data["Success IK"].append(succ)
data["Position Error(mm)"].append(p_err * 1000.0)
data["Orientation Error"].append(q_err)
data["Position Error Collision-Free IK(mm)"].append(p_err_c * 1000.0)
data["Orientation Error Collision-Free IK"].append(q_err_c)
data["Collision-Free IK time(ms)"].append(dt_cu_ik_cfree * 1000.0)
write_yaml(data, join_path(args.save_path, args.file_name + ".yml"))
try:
# Third Party
import pandas as pd
df = pd.DataFrame(data)
print(df)
df.to_csv(join_path(args.save_path, args.file_name + ".csv"))
except ImportError:
pass

231
benchmark/kinematics_benchmark.py Normal file
View File

@@ -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

36
benchmark/metrics.py Normal file
View File

@@ -0,0 +1,36 @@
#
# 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
from dataclasses import dataclass
from typing import List
# Third Party
import numpy as np
from robometrics.statistics import Statistic, TrajectoryGroupMetrics, TrajectoryMetrics
@dataclass
class CuroboMetrics(TrajectoryMetrics):
time: float = np.inf
@dataclass
class CuroboGroupMetrics(TrajectoryGroupMetrics):
time: float = np.inf
@classmethod
def from_list(cls, group: List[CuroboMetrics]):
unskipped = [m for m in group if not m.skip]
successes = [m for m in unskipped if m.success]
data = super().from_list(group)
data.time = Statistic.from_list([m.time for m in successes])
return data

612
benchmark/robometrics_benchmark.py Normal file
View File

@@ -0,0 +1,612 @@
#
# 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 copy import deepcopy
from typing import Any, Dict, List, Optional
# Third Party
import numpy as np
import torch
from geometrout.primitive import Cuboid, Cylinder
from geometrout.transform import SE3
from metrics import CuroboGroupMetrics, CuroboMetrics
from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw
from robometrics.evaluator import Evaluator
from robometrics.robot import CollisionSphereConfig, Robot
from tqdm import tqdm
# CuRobo
from curobo.geom.sdf.world import CollisionCheckerType, WorldConfig
from curobo.types.base import TensorDeviceType
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_file import (
get_assets_path,
get_robot_configs_path,
get_task_configs_path,
get_world_configs_path,
join_path,
load_yaml,
write_yaml,
)
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 = True,
):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
robot_cfg["kinematics"]["collision_sphere_buffer"] = 0.0
robot_cfg["kinematics"]["collision_spheres"] = "spheres/franka_mesh.yml"
robot_cfg["kinematics"]["collision_link_names"].remove("attached_object")
ik_seeds = 30 # 500
if graph_mode:
trajopt_seeds = 4
if trajopt_seeds >= 14:
ik_seeds = max(100, trajopt_seeds * 4)
if mpinets:
robot_cfg["kinematics"]["lock_joints"] = {
"panda_finger_joint1": 0.025,
"panda_finger_joint2": -0.025,
}
world_cfg = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
).get_obb_world()
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.1
K.position[1, :] += 0.1
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_cfg_instance,
world_cfg,
trajopt_tsteps=tsteps,
collision_checker_type=CollisionCheckerType.PRIMITIVE,
use_cuda_graph=True,
collision_cache={"obb": n_cubes},
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,
interpolation_steps=interpolation_steps,
collision_activation_distance=0.03,
state_finite_difference_mode="CENTRAL",
trajopt_dt=0.25,
minimize_jerk=True,
finetune_dt_scale=1.05, # 1.05,
maximum_trajectory_dt=0.1,
)
mg = MotionGen(motion_gen_config)
mg.warmup(enable_graph=True, warmup_js_trajopt=False)
robot_cfg["kinematics"]["ee_link"] = "panda_hand"
return mg, robot_cfg
def benchmark_mb(
write_usd=False,
save_log=False,
write_plot=False,
write_benchmark=False,
plot_cost=False,
override_tsteps: Optional[int] = None,
save_kpi=False,
graph_mode=False,
args=None,
):
interpolation_dt = 0.02
enable_debug = False
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
urdf = join_path(get_assets_path(), robot_cfg["kinematics"]["urdf_path"])
spheres = robot_cfg["kinematics"]["collision_spheres"]
if isinstance(spheres, str):
spheres = load_yaml(join_path(get_robot_configs_path(), spheres))["collision_spheres"]
for s in spheres:
for k in spheres[s]:
k["radius"] = max(0.001, k["radius"] - 0.015)
data = {
"collision_spheres": spheres,
"self_collision_ignore": robot_cfg["kinematics"]["self_collision_ignore"],
"self_collision_buffer": robot_cfg["kinematics"]["self_collision_buffer"],
}
config = CollisionSphereConfig.load_from_dictionary(data)
metrics_robot = Robot(urdf, config)
evaluator = Evaluator(metrics_robot)
all_files = []
# modify robot joint limits as some start states in the dataset are at the joint limits:
if True:
for k in evaluator.robot.actuated_joints:
k.limit.lower -= 0.1
k.limit.upper += 0.1
plan_config = MotionGenPlanConfig(
max_attempts=60,
enable_graph_attempt=3,
enable_finetune_trajopt=False,
partial_ik_opt=True,
)
file_paths = [motion_benchmaker_raw, mpinets_raw]
if args.demo:
file_paths = [demo_raw]
# load dataset:
og_tsteps = 32
if override_tsteps is not None:
og_tsteps = override_tsteps
og_trajopt_seeds = 12
for file_path in file_paths:
all_groups = []
mpinets_data = False
problems = file_path()
if "dresser_task_oriented" in list(problems.keys()):
mpinets_data = True
for key, v in tqdm(problems.items()):
tsteps = og_tsteps
trajopt_seeds = og_trajopt_seeds
if "cage_panda" in key:
trajopt_seeds = 16
if "table_under_pick_panda" in key:
tsteps = 44
trajopt_seeds = 28
if "cubby_task_oriented" in key and "merged" not in key:
trajopt_seeds = 16
if "dresser_task_oriented" in key:
trajopt_seeds = 16
scene_problems = problems[key] # [:4] # [:1] # [:20] # [0:10]
n_cubes = check_problems(scene_problems)
mg, robot_cfg = load_curobo(
n_cubes, enable_debug, tsteps, trajopt_seeds, mpinets_data, graph_mode
)
m_list = []
i = 0
ik_fail = 0
for problem in tqdm(scene_problems, leave=False):
i += 1
if problem["collision_buffer_ik"] < 0.0:
continue
plan_config = MotionGenPlanConfig(
max_attempts=100, # 00, # 00, # 100, # 00, # 000,#,00,#00, # 5000,
enable_graph_attempt=3,
enable_finetune_trajopt=True,
partial_ik_opt=False,
enable_graph=graph_mode,
timeout=60,
enable_opt=not graph_mode,
)
q_start = problem["start"]
pose = (
problem["goal_pose"]["position_xyz"] + problem["goal_pose"]["quaternion_wxyz"]
)
problem_name = "d_" + key + "_" + str(i)
# reset planner
mg.reset(reset_seed=False)
world = WorldConfig.from_dict(deepcopy(problem["obstacles"])).get_obb_world()
# world.save_world_as_mesh(problem_name + ".stl")
mg.world_coll_checker.clear_cache()
mg.update_world(world)
# continue
# load obstacles
# run planner
start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
result = mg.plan_single(
start_state,
Pose.from_list(pose),
plan_config,
)
if result.status == "IK Fail":
ik_fail += 1
if result.success.item():
eval_obs = get_metrics_obstacles(problem["obstacles"])
q_traj = result.get_interpolated_plan()
ee_trajectory = mg.compute_kinematics(q_traj)
q_metrics = q_traj.position.cpu().numpy()
ee_m = ee_trajectory.ee_pose
ee_pos = ee_m.position.cpu().numpy()
ee_q = ee_m.quaternion.cpu().numpy()
se3_list = [
SE3(np.ravel(ee_pos[p]), np.ravel(ee_q[p]))
for p in range(ee_m.position.shape[0])
]
# add gripper position:
q_met = np.zeros((q_metrics.shape[0], 8))
q_met[:, :7] = q_metrics
q_met[:, 7] = 0.04
if mpinets_data:
q_met[:, 7] = 0.025
st_time = time.time()
current_metrics = evaluator.evaluate_trajectory(
q_met,
se3_list,
SE3(np.ravel(pose[:3]), np.ravel(pose[3:])),
eval_obs,
result.total_time,
)
# if not current_metrics.success:
# print(current_metrics)
# write_usd = True
# else:
# write_usd = False
# rint(plan_config.enable_graph, plan_config.enable_graph_attempt)
problem["solution"] = None
if plan_config.enable_finetune_trajopt:
problem_name = key + "_" + str(i)
else:
problem_name = "noft_" + key + "_" + str(i)
problem_name = "nosw_" + problem_name
if write_usd or save_log:
# CuRobo
from curobo.util.usd_helper import UsdHelper
world.randomize_color(r=[0.5, 0.9], g=[0.2, 0.5], b=[0.0, 0.2])
gripper_mesh = Mesh(
name="target_gripper",
file_path=join_path(
get_assets_path(),
"robot/franka_description/meshes/visual/hand_ee_link.dae",
),
color=[0.0, 0.8, 0.1, 1.0],
pose=pose,
)
world.add_obstacle(gripper_mesh)
# get costs:
if plot_cost:
dt = 0.5
problem_name = "approx_wolfe_p" + problem_name
if result.optimized_dt is not None:
dt = result.optimized_dt.item()
if "trajopt_result" in result.debug_info:
success = result.success.item()
traj_cost = (
# result.debug_info["trajopt_result"].debug_info["solver"]["cost"][0]
result.debug_info["trajopt_result"].debug_info["solver"]["cost"][-1]
)
# print(traj_cost[0])
traj_cost = torch.cat(traj_cost, dim=-1)
plot_cost_iteration(
traj_cost,
title=problem_name + "_" + str(success) + "_" + str(dt),
save_path=join_path("log/plot/", problem_name + "_cost"),
log_scale=False,
)
if "finetune_trajopt_result" in result.debug_info:
traj_cost = result.debug_info["finetune_trajopt_result"].debug_info[
"solver"
]["cost"][0]
traj_cost = torch.cat(traj_cost, dim=-1)
plot_cost_iteration(
traj_cost,
title=problem_name + "_" + str(success) + "_" + str(dt),
save_path=join_path("log/plot/", problem_name + "_ft_cost"),
log_scale=False,
)
if result.success.item() and current_metrics.success:
q_traj = result.get_interpolated_plan()
problem["goal_ik"] = q_traj.position.cpu().squeeze().numpy()[-1, :].tolist()
problem["solution"] = {
"position": result.get_interpolated_plan()
.position.cpu()
.squeeze()
.numpy()
.tolist(),
"velocity": result.get_interpolated_plan()
.velocity.cpu()
.squeeze()
.numpy()
.tolist(),
"acceleration": result.get_interpolated_plan()
.acceleration.cpu()
.squeeze()
.numpy()
.tolist(),
"jerk": result.get_interpolated_plan()
.jerk.cpu()
.squeeze()
.numpy()
.tolist(),
"dt": interpolation_dt,
}
# print(problem["solution"]["position"])
# exit()
debug = {
"used_graph": result.used_graph,
"attempts": result.attempts,
"ik_time": result.ik_time,
"graph_time": result.graph_time,
"trajopt_time": result.trajopt_time,
"total_time": result.total_time,
"solve_time": result.solve_time,
"opt_traj": {
"position": result.optimized_plan.position.cpu()
.squeeze()
.numpy()
.tolist(),
"velocity": result.optimized_plan.velocity.cpu()
.squeeze()
.numpy()
.tolist(),
"acceleration": result.optimized_plan.acceleration.cpu()
.squeeze()
.numpy()
.tolist(),
"jerk": result.optimized_plan.jerk.cpu().squeeze().numpy().tolist(),
"dt": result.optimized_dt.item(),
},
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
# print(
# "T: ",
# result.motion_time.item(),
# result.optimized_dt.item(),
# (len(problem["solution"]["position"]) - 1 ) * result.interpolation_dt,
# result.interpolation_dt,
# )
# exit()
current_metrics = CuroboMetrics(
skip=False,
success=True,
time=result.total_time,
collision=False,
joint_limit_violation=False,
self_collision=False,
position_error=result.position_error.item() * 100.0,
orientation_error=result.rotation_error.item() * 100.0,
eef_position_path_length=10,
eef_orientation_path_length=10,
attempts=result.attempts,
motion_time=result.motion_time.item(),
solve_time=result.solve_time,
)
if write_usd:
# CuRobo
q_traj = result.get_interpolated_plan()
UsdHelper.write_trajectory_animation_with_robot_usd(
robot_cfg,
world,
start_state,
q_traj,
dt=result.interpolation_dt,
save_path=join_path("log/usd/", problem_name) + ".usd",
interpolation_steps=1,
write_robot_usd_path="log/usd/assets/",
robot_usd_local_reference="assets/",
base_frame="/world_" + problem_name,
visualize_robot_spheres=True,
)
if write_plot:
problem_name = problem_name
plot_traj(
result.optimized_plan,
result.optimized_dt.item(),
# result.get_interpolated_plan(),
# result.interpolation_dt,
title=problem_name,
save_path=join_path("log/plot/", problem_name + ".pdf"),
)
plot_traj(
# result.optimized_plan,
# result.optimized_dt.item(),
result.get_interpolated_plan(),
result.interpolation_dt,
title=problem_name,
save_path=join_path("log/plot/", problem_name + "_int.pdf"),
)
# exit()
m_list.append(current_metrics)
all_groups.append(current_metrics)
elif result.valid_query:
# print("fail")
# print(result.status)
current_metrics = CuroboMetrics()
debug = {
"used_graph": result.used_graph,
"attempts": result.attempts,
"ik_time": result.ik_time,
"graph_time": result.graph_time,
"trajopt_time": result.trajopt_time,
"total_time": result.total_time,
"solve_time": result.solve_time,
"status": result.status,
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
m_list.append(current_metrics)
all_groups.append(current_metrics)
else:
# print("invalid: " + problem_name)
debug = {
"used_graph": result.used_graph,
"attempts": result.attempts,
"ik_time": result.ik_time,
"graph_time": result.graph_time,
"trajopt_time": result.trajopt_time,
"total_time": result.total_time,
"solve_time": result.solve_time,
"status": result.status,
"valid_query": result.valid_query,
}
problem["solution_debug"] = debug
if save_log: # and not result.success.item():
# print("save log")
UsdHelper.write_motion_gen_log(
result,
robot_cfg,
world,
start_state,
Pose.from_list(pose),
join_path("log/usd/", problem_name) + "_debug",
write_ik=False,
write_trajopt=True,
visualize_robot_spheres=False,
grid_space=2,
)
# exit()
g_m = CuroboGroupMetrics.from_list(m_list)
print(
key,
f"{g_m.success:2.2f}",
# g_m.motion_time,
g_m.time.mean,
# g_m.time.percent_75,
g_m.time.percent_98,
g_m.position_error.percent_98,
# g_m.position_error.median,
g_m.orientation_error.percent_98,
# g_m.orientation_error.median,
) # , g_m.attempts)
print(g_m.attempts)
# print("MT: ", g_m.motion_time)
# $print(ik_fail)
# exit()
# print(g_m.position_error, g_m.orientation_error)
g_m = CuroboGroupMetrics.from_list(all_groups)
print(
"All: ",
f"{g_m.success:2.2f}",
g_m.motion_time.percent_98,
g_m.time.mean,
g_m.time.percent_75,
g_m.position_error.percent_75,
g_m.orientation_error.percent_75,
) # g_m.time, g_m.attempts)
# print("MT: ", g_m.motion_time)
# print(g_m.position_error, g_m.orientation_error)
# exit()
if write_benchmark:
if not mpinets_data:
write_yaml(problems, "robometrics_mb_curobo_solution.yaml")
else:
write_yaml(problems, "robometrics_mpinets_curobo_solution.yaml")
all_files += all_groups
g_m = CuroboGroupMetrics.from_list(all_files)
# print(g_m.success, g_m.time, g_m.attempts, g_m.position_error, g_m.orientation_error)
print("######## FULL SET ############")
print("All: ", f"{g_m.success:2.2f}")
print("MT: ", g_m.motion_time)
print("PT:", g_m.time)
print("ST: ", g_m.solve_time)
print("position accuracy: ", g_m.position_error)
print("orientation accuracy: ", g_m.orientation_error)
if args.kpi:
kpi_data = {"Success": g_m.success, "Planning Time": float(g_m.time.mean)}
write_yaml(kpi_data, join_path(args.save_path, args.file_name + ".yml"))
def get_metrics_obstacles(obs: Dict[str, List[Any]]):
obs_list = []
if "cylinder" in obs and len(obs["cylinder"].items()) > 0:
for _, vi in enumerate(obs["cylinder"].values()):
obs_list.append(
Cylinder(
np.ravel(vi["pose"][:3]), vi["radius"], vi["height"], np.ravel(vi["pose"][3:])
)
)
if "cuboid" in obs and len(obs["cuboid"].items()) > 0:
for _, vi in enumerate(obs["cuboid"].values()):
obs_list.append(
Cuboid(np.ravel(vi["pose"][:3]), np.ravel(vi["dims"]), np.ravel(vi["pose"][3:]))
)
return obs_list
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=".",
help="path to save file",
)
parser.add_argument(
"--file_name",
type=str,
default="mg",
help="File name prefix to use to save benchmark results",
)
parser.add_argument(
"--graph",
action="store_true",
help="When True, runs only geometric planner",
default=False,
)
parser.add_argument(
"--kpi",
action="store_true",
help="When True, saves minimal metrics",
default=False,
)
parser.add_argument(
"--demo",
action="store_true",
help="When True, runs only on small dataaset",
default=False,
)
args = parser.parse_args()
setup_curobo_logger("error")
benchmark_mb(args=args, save_kpi=args.kpi)