kinematics refactor with mimic joint support

This commit is contained in:
Balakumar Sundaralingam
2024-04-03 18:42:01 -07:00
parent b481ee201a
commit 774dcfd609
60 changed files with 2177 additions and 810 deletions

View File

@@ -154,7 +154,7 @@ def load_curobo(
# del robot_cfg["kinematics"]
if ik_seeds is None:
ik_seeds = 24
ik_seeds = 32
if graph_mode:
trajopt_seeds = 4
@@ -231,7 +231,7 @@ def benchmark_mb(
# load dataset:
force_graph = False
file_paths = [motion_benchmaker_raw, mpinets_raw]
file_paths = [motion_benchmaker_raw, mpinets_raw][:]
if args.demo:
file_paths = [demo_raw]
@@ -500,8 +500,8 @@ def benchmark_mb(
write_robot_usd_path="benchmark/log/usd/assets/",
robot_usd_local_reference="assets/",
base_frame="/world_" + problem_name,
visualize_robot_spheres=False,
flatten_usd=False,
visualize_robot_spheres=True,
flatten_usd=True,
)
if write_plot: # and result.optimized_dt.item() > 0.06:
@@ -554,12 +554,12 @@ def benchmark_mb(
start_state,
Pose.from_list(pose),
join_path("benchmark/log/usd/", problem_name) + "_debug",
write_ik=False,
write_ik=True,
write_trajopt=True,
visualize_robot_spheres=False,
visualize_robot_spheres=True,
grid_space=2,
write_robot_usd_path="benchmark/log/usd/assets/",
# flatten_usd=True,
flatten_usd=True,
)
print(result.status)

View File

@@ -24,39 +24,52 @@ 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.logger import setup_curobo_logger
from curobo.util_file import get_robot_configs_path, get_robot_path, join_path, load_yaml
from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
def demo_motion_gen():
# Standard Library
def demo_motion_gen(robot_file, motion_gen=None):
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=44,
collision_checker_type=CollisionCheckerType.PRIMITIVE,
use_cuda_graph=False,
num_trajopt_seeds=24,
num_graph_seeds=24,
evaluate_interpolated_trajectory=True,
interpolation_dt=0.02,
)
motion_gen = MotionGen(motion_gen_config)
# st_time = time.time()
motion_gen.warmup(enable_graph=True, warmup_js_trajopt=False)
print("motion gen time:", time.time() - st_time)
if motion_gen is None:
setup_curobo_logger("warn")
# Standard Library
tensor_args = TensorDeviceType()
world_file = "collision_table.yml"
robot_file = "ur5e.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=True,
num_trajopt_seeds=4,
num_graph_seeds=4,
evaluate_interpolated_trajectory=True,
interpolation_dt=0.02,
)
motion_gen = MotionGen(motion_gen_config)
# st_time = time.time()
torch.cuda.synchronize()
print("LOAD TIME: ", time.time() - st_time)
st_time = time.time()
motion_gen.warmup(enable_graph=True, warmup_js_trajopt=False)
torch.cuda.synchronize()
print("warmup TIME: ", time.time() - st_time)
return motion_gen
# print(time.time() - st_time)
# return
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))
)
@@ -66,15 +79,58 @@ def demo_motion_gen():
result = motion_gen.plan(
start_state,
retract_pose,
enable_graph=True,
enable_graph=False,
enable_opt=True,
max_attempts=1,
need_graph_success=True,
# need_graph_success=True,
)
print(result.status)
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())
print("Trajectory Generated: ", result.success, time.time() - st_time)
return motion_gen
def demo_basic_ik(config_file="ur10e.yml"):
st_time = time.time()
tensor_args = TensorDeviceType()
config_file = load_yaml(join_path(get_robot_configs_path(), config_file))
urdf_file = config_file["kinematics"]["urdf_path"] # Send global path starting with "/"
base_link = config_file["kinematics"]["base_link"]
ee_link = config_file["kinematics"]["ee_link"]
robot_cfg = RobotConfig.from_basic(urdf_file, base_link, ee_link, tensor_args)
ik_config = IKSolverConfig.load_from_robot_config(
robot_cfg,
None,
rotation_threshold=0.05,
position_threshold=0.005,
num_seeds=20,
self_collision_check=False,
self_collision_opt=False,
tensor_args=tensor_args,
use_cuda_graph=False,
)
ik_solver = IKSolver(ik_config)
torch.cuda.synchronize()
print("IK load time:", time.time() - st_time)
st_time = time.time()
# print(kin_state)
q_sample = ik_solver.sample_configs(100)
kin_state = ik_solver.fk(q_sample)
goal = Pose(kin_state.ee_position, kin_state.ee_quaternion)
torch.cuda.synchronize()
print("FK time:", time.time() - st_time)
st_time = time.time()
result = ik_solver.solve_batch(goal)
torch.cuda.synchronize()
print(
"Cold Start Solve Time(s) ",
result.solve_time,
)
def demo_basic_robot():
@@ -122,7 +178,7 @@ if __name__ == "__main__":
parser.add_argument(
"--save_path",
type=str,
default="log/trace",
default="benchmark/log/trace",
help="path to save file",
)
parser.add_argument(
@@ -137,12 +193,24 @@ if __name__ == "__main__":
help="When True, runs startup for motion generation",
default=False,
)
parser.add_argument(
"--motion_gen_plan",
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=False,
)
parser.add_argument(
"--inverse_kinematics",
action="store_true",
help="When True, runs startup for kinematics",
default=False,
)
parser.add_argument(
"--motion_gen_once",
action="store_true",
@@ -172,21 +240,59 @@ if __name__ == "__main__":
filename = join_path(args.save_path, args.file_name) + "_kinematics_trace.json"
prof.export_chrome_trace(filename)
if args.motion_gen_once:
demo_motion_gen()
if args.motion_gen:
for _ in range(5):
demo_motion_gen()
if args.inverse_kinematics:
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
demo_basic_ik(config_file)
filename = join_path(args.save_path, args.file_name) + "_inverse_kinematics_trace.json"
prof.export_chrome_trace(filename)
pr = cProfile.Profile()
pr.enable()
demo_motion_gen()
demo_basic_ik(config_file)
pr.disable()
filename = join_path(args.save_path, args.file_name) + "_inverse_kinematics_cprofile.prof"
pr.dump_stats(filename)
if args.motion_gen_once:
pr = cProfile.Profile()
pr.enable()
demo_motion_gen(config_file)
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()
demo_motion_gen(config_file)
filename = join_path(args.save_path, args.file_name) + "_motion_gen_trace.json"
prof.export_chrome_trace(filename)
if args.motion_gen_plan:
motion_gen = demo_motion_gen(config_file)
pr = cProfile.Profile()
pr.enable()
demo_motion_gen(config_file, motion_gen)
pr.disable()
filename = join_path(args.save_path, args.file_name) + "_motion_gen_plan_cprofile.prof"
pr.dump_stats(filename)
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
demo_motion_gen(config_file, motion_gen)
filename = join_path(args.save_path, args.file_name) + "_motion_gen_plan_trace.json"
prof.export_chrome_trace(filename)
if args.motion_gen:
for _ in range(5):
demo_motion_gen(config_file)
pr = cProfile.Profile()
pr.enable()
demo_motion_gen(config_file)
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(config_file)
filename = join_path(args.save_path, args.file_name) + "_motion_gen_trace.json"
prof.export_chrome_trace(filename)

View File

@@ -67,7 +67,7 @@ def run_full_config_collision_free_ik(
robot_cfg,
world_cfg,
position_threshold=position_threshold,
num_seeds=20,
num_seeds=30,
self_collision_check=collision_free,
self_collision_opt=collision_free,
tensor_args=tensor_args,