diff --git a/CHANGELOG.md b/CHANGELOG.md index 7088b4d..01d122f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,10 +10,23 @@ its affiliates is strictly prohibited. --> # Changelog -## Latest Commit +## Version 0.7.1 ### New Features - Add mimic joint parsing and optimization support. Check `ur5e_robotiq_2f_140.yml`. +- Add `finetune_dt_scale` as a parameter to `MotionGenPlanConfig` to dynamically change the +time-optimal scaling on a per problem instance. +- `MotionGen.plan_single()` will now try finetuning in a for-loop, with larger and larger dt +until convergence. This also warm starts from previous failure. +- Add `high_precision` mode to `MotionGenConfig` to support `<1mm` convergence. + +### Changes in default behavior +- collision_sphere_buffer now supports having offset per link. Also, collision_sphere_buffer only +applies to world collision while self_collision_buffer applies for self collision. Previously, +self_collision_buffer was added on top of collision_sphere_buffer. +- `TrajEvaluatorConfig` cannot be initialized without dof as now per-joint jerk and acceleration +limits are used. Use `TrajEvaluatorConfig.from_basic()` to initialize similar to previous behavior. +- `finetune_dt_scale` default value is 0.9 from 0.95. ### BugFixes & Misc. - Fix bug in `WorldVoxelCollision` where `env_query_idx` was being overwritten. @@ -26,6 +39,21 @@ its affiliates is strictly prohibited. - Added flag to sample from ik seeder instead of `rollout_fn` sampler. - Added ik startup profiler to `benchmark/curobo_python_profile.py`. - Reduced branching in Kinematics kernels and added mimic joint computations. +- Add init_cache to WorldVoxelCollision to create cache for Mesh and Cuboid obstacles. +- `TrajEvaluator` now uses per-joint acceleration and jerk limits. +- Fixed regression in `batch_motion_gen_reacher.py` example where robot's position was not being +set correctly. +- Switched from smooth l2 to l2 for BoundCost as that gives better convergence. +- `requires_grad` is explicitly stored in a varaible before `tensor.detach()` in warp kernel calls +as this can get set to False in some instances. +- Fix dt update in `MotionGen.plan_single_js()` where dt was not reset after finetunestep, causing +joint space planner to fail often. +- Improve joint space planner success by changing smooth l2 distance cost to l2 distance. Also, +added fallback to graph planner when linear path is not possible. +- Retuned weigths for IKSolver, now 98th percentile accuracy is 10 micrometers wtih 16 seeds +(vs 24 seeds previously). +- Switch float8 precision check from `const` to macro to avoid compile errors in older nvcc, this +fixes docker build issues for isaac sim 2023.1.0. ## Version 0.7.0 ### Changes in default behavior diff --git a/benchmark/curobo_benchmark.py b/benchmark/curobo_benchmark.py index 548692f..5c582b1 100644 --- a/benchmark/curobo_benchmark.py +++ b/benchmark/curobo_benchmark.py @@ -152,7 +152,6 @@ def load_curobo( robot_cfg["kinematics"]["collision_link_names"].remove("attached_object") robot_cfg["kinematics"]["ee_link"] = "panda_hand" - # del robot_cfg["kinematics"] if ik_seeds is None: ik_seeds = 32 @@ -211,6 +210,7 @@ def load_curobo( trajopt_dt=0.25, finetune_dt_scale=finetune_dt_scale, maximum_trajectory_dt=0.15, + high_precision=args.high_precision, ) mg = MotionGen(motion_gen_config) mg.warmup(enable_graph=True, warmup_js_trajopt=False, parallel_finetune=parallel_finetune) @@ -240,7 +240,7 @@ def benchmark_mb( og_tsteps = 32 if override_tsteps is not None: og_tsteps = override_tsteps - og_finetune_dt_scale = 0.9 + og_finetune_dt_scale = 0.85 og_trajopt_seeds = 4 og_parallel_finetune = True og_collision_activation_distance = 0.01 @@ -252,6 +252,7 @@ def benchmark_mb( if "dresser_task_oriented" in list(problems.keys()): mpinets_data = True for key, v in tqdm(problems.items()): + finetune_dt_scale = og_finetune_dt_scale force_graph = False tsteps = og_tsteps @@ -260,23 +261,12 @@ def benchmark_mb( parallel_finetune = og_parallel_finetune ik_seeds = og_ik_seeds - if "cage_panda" in key: - trajopt_seeds = 8 - - if "table_under_pick_panda" in key: - trajopt_seeds = 8 - finetune_dt_scale = 0.95 - - if key == "cubby_task_oriented": # or key == "merged_cubby_task_oriented": - trajopt_seeds = 16 - finetune_dt_scale = 0.95 - - if "dresser_task_oriented" in key: - trajopt_seeds = 16 - finetune_dt_scale = 0.95 - - scene_problems = problems[key][:] + scene_problems = problems[key] n_cubes = check_problems(scene_problems) + + if "cubby_task_oriented" in key and "merged" not in key: + trajopt_seeds = 8 + mg, robot_cfg = load_curobo( n_cubes, enable_debug, @@ -302,7 +292,7 @@ def benchmark_mb( continue plan_config = MotionGenPlanConfig( - max_attempts=20, + max_attempts=20, # 20, enable_graph_attempt=1, disable_graph_attempt=10, enable_finetune_trajopt=not args.no_finetune, @@ -312,6 +302,7 @@ def benchmark_mb( enable_opt=not graph_mode, need_graph_success=force_graph, parallel_finetune=parallel_finetune, + finetune_dt_scale=finetune_dt_scale, ) q_start = problem["start"] pose = ( @@ -579,16 +570,31 @@ def benchmark_mb( 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, - ) + try: + from tabulate import tabulate + + headers = ["Metric", "Value"] + + table = [ + ["Success %", f"{g_m.success:2.2f}"], + ["Plan Time (s)", g_m.time], + ["Motion Time(s)", g_m.motion_time], + ["Path Length (rad.)", g_m.cspace_path_length], + ["Jerk", g_m.jerk], + ["Position Error (mm)", g_m.position_error], + ] + print(tabulate(table, headers, tablefmt="grid")) + except ImportError: + 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, + ) if write_benchmark: if not mpinets_data: write_yaml(problems, args.file_name + "_mb_solution.yaml") @@ -596,15 +602,32 @@ def benchmark_mb( write_yaml(problems, args.file_name + "_mpinets_solution.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("path-length: ", g_m.cspace_path_length) - 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) - print("jerk: ", g_m.jerk) + + try: + from tabulate import tabulate + + headers = ["Metric", "Value"] + + table = [ + ["Success %", f"{g_m.success:2.2f}"], + ["Plan Time (s)", g_m.time], + ["Motion Time(s)", g_m.motion_time], + ["Path Length (rad.)", g_m.cspace_path_length], + ["Jerk", g_m.jerk], + ["Position Error (mm)", g_m.position_error], + ] + print(tabulate(table, headers, tablefmt="grid")) + except ImportError: + + print("######## FULL SET ############") + print("All: ", f"{g_m.success:2.2f}") + print("MT: ", g_m.motion_time) + print("path-length: ", g_m.cspace_path_length) + 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) + print("jerk: ", g_m.jerk) if args.kpi: kpi_data = { @@ -716,6 +739,12 @@ if __name__ == "__main__": help="When True, runs benchmark with parameters for jetson", default=False, ) + parser.add_argument( + "--high_precision", + action="store_true", + help="When True, runs benchmark with parameters for jetson", + default=False, + ) args = parser.parse_args() diff --git a/benchmark/curobo_nvblox_benchmark.py b/benchmark/curobo_nvblox_benchmark.py index 3e6ffde..6217ef5 100644 --- a/benchmark/curobo_nvblox_benchmark.py +++ b/benchmark/curobo_nvblox_benchmark.py @@ -147,7 +147,7 @@ def load_curobo( "world": { "pose": [0, 0, 0, 1, 0, 0, 0], "integrator_type": "tsdf", - "voxel_size": 0.01, + "voxel_size": 0.02, } } } @@ -177,9 +177,9 @@ def load_curobo( interpolation_steps=interpolation_steps, collision_activation_distance=collision_activation_distance, trajopt_dt=0.25, - finetune_dt_scale=1.0, - maximum_trajectory_dt=0.1, - finetune_trajopt_iters=300, + finetune_dt_scale=0.9, + maximum_trajectory_dt=0.15, + finetune_trajopt_iters=200, ) mg = MotionGen(motion_gen_config) mg.warmup(enable_graph=True, warmup_js_trajopt=False, parallel_finetune=True) @@ -208,7 +208,7 @@ def benchmark_mb( # load dataset: graph_mode = args.graph interpolation_dt = 0.02 - file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][2:] + file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][1:2] enable_debug = save_log or plot_cost all_files = [] @@ -237,8 +237,9 @@ def benchmark_mb( mpinets_data = True if "cage_panda" in key: - trajopt_seeds = 16 - finetune_dt_scale = 0.95 + trajopt_seeds = 8 + else: + continue if "table_under_pick_panda" in key: tsteps = 44 trajopt_seeds = 16 diff --git a/benchmark/curobo_voxel_benchmark.py b/benchmark/curobo_voxel_benchmark.py index bfccc0e..08171c3 100644 --- a/benchmark/curobo_voxel_benchmark.py +++ b/benchmark/curobo_voxel_benchmark.py @@ -133,7 +133,7 @@ def load_curobo( robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"] robot_cfg["kinematics"]["collision_sphere_buffer"] = -0.00 - ik_seeds = 24 + ik_seeds = 32 if graph_mode: trajopt_seeds = 4 if trajopt_seeds >= 14: @@ -253,21 +253,9 @@ def benchmark_mb( collision_activation_distance = og_collision_activation_distance finetune_dt_scale = 0.9 parallel_finetune = True - if "cage_panda" in key: + if "cubby_task_oriented" in key and "merged" not in key: trajopt_seeds = 8 - if "table_under_pick_panda" in key: - trajopt_seeds = 8 - finetune_dt_scale = 0.98 - - if key == "cubby_task_oriented": - trajopt_seeds = 16 - finetune_dt_scale = 0.98 - - if "dresser_task_oriented" in key: - trajopt_seeds = 16 - finetune_dt_scale = 0.98 - mg, robot_cfg, robot_world = load_curobo( 0, enable_debug, @@ -285,16 +273,12 @@ def benchmark_mb( i += 1 if problem["collision_buffer_ik"] < 0.0: continue - plan_config = MotionGenPlanConfig( - max_attempts=10, + max_attempts=20, enable_graph_attempt=1, - enable_finetune_trajopt=True, + disable_graph_attempt=10, partial_ik_opt=False, - enable_graph=graph_mode, timeout=60, - enable_opt=not graph_mode, - parallel_finetune=True, ) q_start = problem["start"] @@ -593,17 +577,31 @@ def benchmark_mb( ) print(g_m.attempts) 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.perception_success, - ) - print(g_m.attempts) + try: + from tabulate import tabulate + + headers = ["Metric", "Value"] + + table = [ + ["Success %", f"{g_m.success:2.2f}"], + ["Plan Time (s)", g_m.time], + ["Motion Time(s)", g_m.motion_time], + ["Path Length (rad.)", g_m.cspace_path_length], + ["Jerk", g_m.jerk], + ["Position Error (mm)", g_m.position_error], + ] + print(tabulate(table, headers, tablefmt="grid")) + except ImportError: + 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, + ) + if write_benchmark: if not mpinets_data: write_yaml(problems, "mb_curobo_solution_voxel.yaml") @@ -612,17 +610,33 @@ def benchmark_mb( all_files += all_groups g_m = CuroboGroupMetrics.from_list(all_files) print("######## FULL SET ############") - print("All: ", f"{g_m.success:2.2f}") - print( - "Perception Success (coarse, interpolated):", - g_m.perception_success, - g_m.perception_interpolated_success, - ) - 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) - print("Jerk: ", g_m.jerk) + try: + from tabulate import tabulate + + headers = ["Metric", "Value"] + + table = [ + ["Success %", f"{g_m.success:2.2f}"], + ["Plan Time (s)", g_m.time], + ["Motion Time(s)", g_m.motion_time], + ["Path Length (rad.)", g_m.cspace_path_length], + ["Jerk", g_m.jerk], + ["Position Error (mm)", g_m.position_error], + ] + print(tabulate(table, headers, tablefmt="grid")) + except ImportError: + + print("All: ", f"{g_m.success:2.2f}") + print( + "Perception Success (coarse, interpolated):", + g_m.perception_success, + g_m.perception_interpolated_success, + ) + 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) + print("Jerk: ", g_m.jerk) if __name__ == "__main__": diff --git a/benchmark/ik_benchmark.py b/benchmark/ik_benchmark.py index b4f127e..83216b8 100644 --- a/benchmark/ik_benchmark.py +++ b/benchmark/ik_benchmark.py @@ -67,7 +67,7 @@ def run_full_config_collision_free_ik( robot_cfg, world_cfg, position_threshold=position_threshold, - num_seeds=30, + num_seeds=16, self_collision_check=collision_free, self_collision_opt=collision_free, tensor_args=tensor_args, @@ -123,7 +123,7 @@ if __name__ == "__main__": args = parser.parse_args() - b_list = [1, 10, 100, 1000][-1:] + b_list = [1, 10, 100, 2000][-1:] robot_list = get_motion_gen_robot_list() + get_multi_arm_robot_list()[:2] world_file = "collision_test.yml" @@ -141,7 +141,7 @@ if __name__ == "__main__": "Position-Error-Collision-Free-IK(mm)": [], "Orientation-Error-Collision-Free-IK": [], } - for robot_file in robot_list[:1]: + for robot_file in robot_list[:-1]: # create a sampler with dof: for b_size in b_list: # sample test configs: @@ -176,13 +176,21 @@ if __name__ == "__main__": 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("Reported errors are 98th percentile") - print(df) df.to_csv(join_path(args.save_path, args.file_name + ".csv")) + try: + from tabulate import tabulate + + print(tabulate(df, headers="keys", tablefmt="grid")) + except ImportError: + print(df) + + pass except ImportError: pass diff --git a/docker/aarch64.dockerfile b/docker/aarch64.dockerfile index 3ea0d5c..152aa19 100644 --- a/docker/aarch64.dockerfile +++ b/docker/aarch64.dockerfile @@ -129,7 +129,7 @@ RUN pip3 install trimesh \ empy # Add cache date to avoid using cached layers older than this -ARG CACHE_DATE=2024-02-20 +ARG CACHE_DATE=2024-04-11 # install warp: # diff --git a/docker/isaac_sim.dockerfile b/docker/isaac_sim.dockerfile index 90c6b84..7ca0e0c 100644 --- a/docker/isaac_sim.dockerfile +++ b/docker/isaac_sim.dockerfile @@ -161,7 +161,7 @@ RUN echo "alias omni_python='/isaac-sim/python.sh'" >> /.bashrc # Add cache date to avoid using cached layers older than this -ARG CACHE_DATE=2023-12-15 +ARG CACHE_DATE=2024-04-11 RUN $omni_python -m pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git" diff --git a/docker/user.dockerfile b/docker/user.dockerfile index b402b83..1ad73d4 100644 --- a/docker/user.dockerfile +++ b/docker/user.dockerfile @@ -15,7 +15,7 @@ FROM curobo_docker:${IMAGE_TAG} # Set variables ARG USERNAME ARG USER_ID -ARG CACHE_DATE=2024-03-18 +ARG CACHE_DATE=2023-04-11 # Set environment variables diff --git a/docker/x86.dockerfile b/docker/x86.dockerfile index 4725bf3..b853afb 100644 --- a/docker/x86.dockerfile +++ b/docker/x86.dockerfile @@ -78,7 +78,7 @@ ENV TORCH_CUDA_ARCH_LIST "7.0+PTX" ENV LD_LIBRARY_PATH="/usr/local/lib:${LD_LIBRARY_PATH}" # Add cache date to avoid using cached layers older than this -ARG CACHE_DATE=2024-02-20 +ARG CACHE_DATE=2024-04-11 RUN pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git" diff --git a/examples/isaac_sim/helper.py b/examples/isaac_sim/helper.py index 10b185a..b6c2f8e 100644 --- a/examples/isaac_sim/helper.py +++ b/examples/isaac_sim/helper.py @@ -23,6 +23,7 @@ from pxr import UsdPhysics # CuRobo from curobo.util.logger import log_warn +from curobo.util.usd_helper import set_prim_transform ISAAC_SIM_23 = False try: @@ -107,8 +108,13 @@ def add_robot_to_scene( robot_p = Robot( prim_path=robot_path + "/" + base_link_name, name=robot_name, - position=position, ) + + robot_prim = robot_p.prim + stage = robot_prim.GetStage() + linkp = stage.GetPrimAtPath(robot_path) + set_prim_transform(linkp, [position[0], position[1], position[2], 1, 0, 0, 0]) + if False and ISAAC_SIM_23: # this doesn't work in isaac sim 2023.1.1 robot_p.set_solver_velocity_iteration_count(0) robot_p.set_solver_position_iteration_count(44) @@ -116,8 +122,6 @@ def add_robot_to_scene( my_world._physics_context.set_solver_type("PGS") if ISAAC_SIM_23: # fix to load robot correctly in isaac sim 2023.1.1 - robot_prim = robot_p.prim - stage = robot_prim.GetStage() linkp = stage.GetPrimAtPath(robot_path + "/" + base_link_name) mass = UsdPhysics.MassAPI(linkp) mass.GetMassAttr().Set(0) diff --git a/examples/isaac_sim/motion_gen_reacher.py b/examples/isaac_sim/motion_gen_reacher.py index 6461fd7..4704a77 100644 --- a/examples/isaac_sim/motion_gen_reacher.py +++ b/examples/isaac_sim/motion_gen_reacher.py @@ -205,7 +205,7 @@ def main(): interpolation_dt = 0.05 if args.reactive: trajopt_tsteps = 40 - trajopt_dt = 0.05 + trajopt_dt = 0.04 optimize_dt = False max_attempts = 1 trim_steps = [1, None] @@ -223,9 +223,6 @@ def main(): trajopt_dt=trajopt_dt, trajopt_tsteps=trajopt_tsteps, trim_steps=trim_steps, - # velocity_scale=0.1, - # self_collision_check=False, - # self_collision_opt=False, ) motion_gen = MotionGen(motion_gen_config) print("warming up...") diff --git a/examples/isaac_sim/multi_arm_reacher.py b/examples/isaac_sim/multi_arm_reacher.py index ee97e3c..4545a89 100644 --- a/examples/isaac_sim/multi_arm_reacher.py +++ b/examples/isaac_sim/multi_arm_reacher.py @@ -136,7 +136,6 @@ def main(): robot_cfg, world_cfg, tensor_args, - trajopt_tsteps=40, collision_checker_type=CollisionCheckerType.MESH, use_cuda_graph=True, num_trajopt_seeds=12, @@ -146,6 +145,7 @@ def main(): collision_activation_distance=0.025, acceleration_scale=1.0, fixed_iters_trajopt=True, + trajopt_tsteps=40, ) motion_gen = MotionGen(motion_gen_config) print("warming up...") @@ -154,7 +154,9 @@ def main(): print("Curobo is Ready") add_extensions(simulation_app, args.headless_mode) plan_config = MotionGenPlanConfig( - enable_graph=False, enable_graph_attempt=4, max_attempts=10, enable_finetune_trajopt=True + enable_graph=False, + enable_graph_attempt=4, + max_attempts=10, ) usd_help.load_stage(my_world.stage) diff --git a/examples/isaac_sim/simple_stacking.py b/examples/isaac_sim/simple_stacking.py index 5d3422f..aaadc47 100644 --- a/examples/isaac_sim/simple_stacking.py +++ b/examples/isaac_sim/simple_stacking.py @@ -162,7 +162,7 @@ class CuroboController(BaseController): pose_metric = None if constrain_grasp_approach: pose_metric = PoseCostMetric.create_grasp_approach_metric( - offset_position=0.1, tstep_fraction=0.6 + offset_position=0.1, tstep_fraction=0.8 ) self.plan_config = MotionGenPlanConfig( diff --git a/examples/motion_gen_example.py b/examples/motion_gen_example.py index 820d8c1..ace6162 100644 --- a/examples/motion_gen_example.py +++ b/examples/motion_gen_example.py @@ -107,7 +107,7 @@ def demo_motion_gen_mesh(): motion_gen = MotionGen(motion_gen_config) 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 = robot_cfg.retract_config + retract_cfg = robot_cfg.cpsace.retract_config state = motion_gen.rollout_fn.compute_kinematics( JointState.from_position(retract_cfg.view(1, -1)) ) diff --git a/examples/motion_gen_profile.py b/examples/motion_gen_profile.py index b64f11d..aeb4145 100644 --- a/examples/motion_gen_profile.py +++ b/examples/motion_gen_profile.py @@ -55,7 +55,7 @@ def demo_motion_gen(): motion_gen = MotionGen(motion_gen_config) 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 = robot_cfg.retract_config + retract_cfg = robot_cfg.cspace.retract_config state = motion_gen.rollout_fn.compute_kinematics( JointState.from_position(retract_cfg.view(1, -1)) ) diff --git a/examples/usd_example.py b/examples/usd_example.py index db4dbf8..d4e1dcc 100644 --- a/examples/usd_example.py +++ b/examples/usd_example.py @@ -51,14 +51,13 @@ def get_trajectory(robot_file="franka.yml", dt=1.0 / 60.0): robot_file, world_file, tensor_args, - trajopt_tsteps=24, + trajopt_tsteps=32, collision_checker_type=CollisionCheckerType.PRIMITIVE, use_cuda_graph=True, num_trajopt_seeds=2, num_graph_seeds=2, evaluate_interpolated_trajectory=True, interpolation_dt=dt, - self_collision_check=True, ) motion_gen = MotionGen(motion_gen_config) motion_gen.warmup() @@ -70,13 +69,15 @@ def get_trajectory(robot_file="franka.yml", dt=1.0 / 60.0): ) retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()) - start_state = JointState.from_position(retract_cfg.view(1, -1).clone() + 0.5) + start_state = JointState.from_position(retract_cfg.view(1, -1).clone()) + start_state.position[..., :-2] += 0.5 result = motion_gen.plan_single(start_state, retract_pose) traj = result.get_interpolated_plan() # optimized plan return traj def save_curobo_robot_world_to_usd(robot_file="franka.yml"): + print(robot_file) tensor_args = TensorDeviceType() world_file = "collision_test.yml" world_model = WorldConfig.from_dict( @@ -87,16 +88,18 @@ def save_curobo_robot_world_to_usd(robot_file="franka.yml"): q_traj = get_trajectory(robot_file, dt) if q_traj is not None: q_start = q_traj[0] - UsdHelper.write_trajectory_animation_with_robot_usd( robot_file, world_model, q_start, q_traj, - save_path="test.usda", + save_path="test.usd", robot_color=[0.5, 0.5, 0.2, 1.0], base_frame="/grid_world_1", + flatten_usd=True, ) + else: + print("failed") def save_curobo_robot_to_usd(robot_file="franka.yml"): @@ -241,5 +244,5 @@ def save_log_motion_gen(robot_file: str = "franka.yml"): if __name__ == "__main__": # save_curobo_world_to_usd() setup_curobo_logger("error") - save_log_motion_gen("franka.yml") - # save_curobo_robot_world_to_usd("franka.yml") + # save_log_motion_gen("franka.yml") + save_curobo_robot_world_to_usd("ur5e_robotiq_2f_140.yml") diff --git a/setup.cfg b/setup.cfg index d5e1758..3b8256c 100644 --- a/setup.cfg +++ b/setup.cfg @@ -80,6 +80,9 @@ ci = sphinx sphinx_rtd_theme graphviz>=0.20.1 + furo + sphinx-copybutton + # this is only available in 3.8+ smooth = @@ -107,6 +110,8 @@ doc = sphinx sphinx_rtd_theme graphviz>=0.20.1 + furo + sphinx-copybutton [options.entry_points] # Add here console scripts like: diff --git a/src/curobo/__init__.py b/src/curobo/__init__.py index 5db6e50..403c5fa 100644 --- a/src/curobo/__init__.py +++ b/src/curobo/__init__.py @@ -10,6 +10,19 @@ # """ +cuRobo provides accelerated modules for robotics which can be used to build high-performance +robotics applications. The library has several modules for numerical optimization, robot kinematics, +geometry processing, collision checking, graph search planning. cuRobo provides high-level APIs for +performing tasks like collision-free inverse kinematics, model predictive control, and motion +planning. + +High-level APIs: + +- Inverse Kinematics: :mod:`curobo.wrap.reacher.ik_solver`. +- Model Predictive Control: :mod:`curobo.wrap.reacher.mpc`. +- Motion Generation / Planning: :mod:`curobo.wrap.reacher.motion_gen`. + + cuRobo package is split into several modules: - :mod:`curobo.opt` contains optimization solvers. @@ -18,7 +31,7 @@ cuRobo package is split into several modules: - :mod:`curobo.geom` contains geometry processing, collision checking and frame transforms. - :mod:`curobo.graph` contains geometric planning with graph search methods. - :mod:`curobo.rollout` contains methods that map actions to costs. This class wraps instances of - :mod:`curobo.cuda_robot_model` and :mod:`geom` to compute costs given trajectory of actions. + :mod:`curobo.cuda_robot_model` and :mod:`curobo.geom` to compute costs given trajectory of actions. - :mod:`curobo.util` contains utility methods. - :mod:`curobo.wrap` adds the user-level api for task programming. Includes implementation of collision-free reacher and batched robot world collision checking. diff --git a/src/curobo/content/configs/robot/template.yml b/src/curobo/content/configs/robot/template.yml index aabdb21..3aee2ed 100644 --- a/src/curobo/content/configs/robot/template.yml +++ b/src/curobo/content/configs/robot/template.yml @@ -29,7 +29,7 @@ robot_cfg: collision_link_names: null # List[str] collision_spheres: null # - collision_sphere_buffer: 0.005 + collision_sphere_buffer: 0.005 # float or Dict[str, float] extra_collision_spheres: {} self_collision_ignore: {} # Dict[str, List[str]] self_collision_buffer: {} # Dict[str, float] diff --git a/src/curobo/content/configs/robot/ur5e.yml b/src/curobo/content/configs/robot/ur5e.yml index 3c17015..15846e6 100644 --- a/src/curobo/content/configs/robot/ur5e.yml +++ b/src/curobo/content/configs/robot/ur5e.yml @@ -11,7 +11,7 @@ robot_cfg: kinematics: - usd_path: "robot/ur_description/ur5e.usd" + usd_path: null usd_robot_root: "/robot" isaac_usd_path: "" usd_flip_joints: {} @@ -95,10 +95,11 @@ robot_cfg: "radius": -0.01 #0.05 - collision_sphere_buffer: 0.005 + collision_sphere_buffer: 0.0 extra_collision_spheres: {} self_collision_ignore: { - "upper_arm_link": ["forearm_link", "shoulder_link"], + "shoulder_link": ["tool0"], + "upper_arm_link": ["forearm_link", "shoulder_link", "tool0"], "forearm_link": ["wrist_1_link", "wrist_2_link", "wrist_3_link"], "wrist_1_link": ["wrist_2_link","wrist_3_link","tool0"], "wrist_2_link": ["wrist_3_link", "tool0"], diff --git a/src/curobo/content/configs/task/base_cfg.yml b/src/curobo/content/configs/task/base_cfg.yml index 4d31c14..bd1eb42 100644 --- a/src/curobo/content/configs/task/base_cfg.yml +++ b/src/curobo/content/configs/task/base_cfg.yml @@ -43,6 +43,7 @@ constraint: classify: True bound_cfg: weight: [5000.0, 5000.0, 5000.0,5000.0] + activation_distance: [0.0,0.0,0.0,0.0] # for position, velocity, acceleration and jerk @@ -54,14 +55,14 @@ convergence: terminal: False link_pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] - weight: [1.0, 1.0] + weight: [0.1, 10.0] vec_convergence: [0.0, 0.0] # orientation, position terminal: False cspace_cfg: weight: 1.0 terminal: True - run_weight: 1.0 + run_weight: 0.0 null_space_cfg: weight: 1.0 terminal: True diff --git a/src/curobo/content/configs/task/finetune_trajopt.yml b/src/curobo/content/configs/task/finetune_trajopt.yml index d692422..6c1c66d 100644 --- a/src/curobo/content/configs/task/finetune_trajopt.yml +++ b/src/curobo/content/configs/task/finetune_trajopt.yml @@ -32,7 +32,7 @@ cost: pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position - weight: [2000,500000.0,30,50] + weight: [2000,500000.0,30,60] vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation terminal: True run_weight: 1.0 @@ -41,10 +41,10 @@ cost: link_pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position - weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40] + weight: [2000,500000.0,30,60] #[150.0, 2000.0, 30, 40] vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation terminal: True - run_weight: 0.001 + run_weight: 1.0 use_metric: True @@ -54,8 +54,8 @@ cost: run_weight: 0.00 #1 bound_cfg: - weight: [5000.0, 50000.0, 50000.0, 50000.0] # needs to be 3 values - smooth_weight: [0.0,5000.0, 50.0, 0.0] #[0.0,5000.0, 50.0, 0.0] # [vel, acc, jerk,] + weight: [10000.0, 500000.0, 500.0, 500.0] + smooth_weight: [0.0,10000.0, 5.0, 0.0] run_weight_velocity: 0.0 run_weight_acceleration: 1.0 run_weight_jerk: 1.0 @@ -63,7 +63,7 @@ cost: null_space_weight: [0.0] primitive_collision_cfg: - weight: 1000000.0 #1000000.0 1000000 + weight: 1000000 use_sweep: True sweep_steps: 4 classify: False @@ -79,11 +79,11 @@ cost: lbfgs: - n_iters: 300 # 400 + n_iters: 300 inner_iters: 25 cold_start_n_iters: null min_iters: 25 - line_search_scale: [0.01, 0.3, 0.7,1.0] # # + line_search_scale: [0.1, 0.3, 0.7,1.0] # # fixed_iters: True cost_convergence: 0.01 cost_delta_threshold: 1.0 diff --git a/src/curobo/content/configs/task/finetune_trajopt_slow.yml b/src/curobo/content/configs/task/finetune_trajopt_slow.yml index 35f1a35..383b903 100644 --- a/src/curobo/content/configs/task/finetune_trajopt_slow.yml +++ b/src/curobo/content/configs/task/finetune_trajopt_slow.yml @@ -54,10 +54,8 @@ cost: run_weight: 0.00 #1 bound_cfg: - weight: [5000.0, 50000.0, 50000.0, 50000.0] # needs to be 3 values - - smooth_weight: [0.0,50000.0, 500.0, 0.0] # [vel, acc, jerk,] - + weight: [10000.0, 100000.0, 500.0, 500.0] + smooth_weight: [0.0,10000.0, 50.0, 0.0] run_weight_velocity: 0.0 run_weight_acceleration: 1.0 run_weight_jerk: 1.0 diff --git a/src/curobo/content/configs/task/gradient_ik.yml b/src/curobo/content/configs/task/gradient_ik.yml index 708fb1e..df3d90d 100644 --- a/src/curobo/content/configs/task/gradient_ik.yml +++ b/src/curobo/content/configs/task/gradient_ik.yml @@ -28,24 +28,24 @@ model: cost: pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] - weight: [200,4000,10,20] - #weight: [5000,500000,5,20] - + weight: [2000,10000,20,40] vec_convergence: [0.0, 0.00] terminal: False use_metric: True run_weight: 1.0 link_pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] - weight: [200,4000,10,20] + weight: [2000,10000,20,40] vec_convergence: [0.00, 0.000] terminal: False use_metric: True + run_weight: 1.0 + cspace_cfg: weight: 0.000 bound_cfg: - weight: 100.0 + weight: 500.0 activation_distance: [0.1] null_space_weight: [1.0] primitive_collision_cfg: @@ -62,7 +62,7 @@ lbfgs: n_iters: 100 #60 inner_iters: 25 cold_start_n_iters: null - min_iters: 25 + min_iters: null line_search_scale: [0.1, 0.3, 0.7, 1.0] fixed_iters: True cost_convergence: 0.001 diff --git a/src/curobo/content/configs/task/gradient_trajopt.yml b/src/curobo/content/configs/task/gradient_trajopt.yml index a40367a..506a25e 100644 --- a/src/curobo/content/configs/task/gradient_trajopt.yml +++ b/src/curobo/content/configs/task/gradient_trajopt.yml @@ -33,7 +33,6 @@ cost: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position weight: [2000,50000.0,30,50] #[150.0, 2000.0, 30, 40] - vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation terminal: True run_weight: 1.0 @@ -45,17 +44,17 @@ cost: weight: [2000,50000.0,30,50] #[150.0, 2000.0, 30, 40] vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation terminal: True - run_weight: 0.001 + run_weight: 1.0 use_metric: True cspace_cfg: - weight: 10000.0 + weight: 1000.0 terminal: True - run_weight: 0.00 #1 + run_weight: 0.0 bound_cfg: - weight: [5000.0, 50000.0, 50000.0,50000.0] # needs to be 3 values - smooth_weight: [0.0,10000.0,10.0, 0.0] # [vel, acc, jerk, alpha_v-not-used] + weight: [50000.0, 50000.0, 50000.0,50000.0] + smooth_weight: [0.0,10000.0,5.0, 0.0] run_weight_velocity: 0.00 run_weight_acceleration: 1.0 run_weight_jerk: 1.0 @@ -84,7 +83,7 @@ lbfgs: inner_iters: 25 cold_start_n_iters: null min_iters: 25 - line_search_scale: [0.01,0.3,0.7,1.0] + line_search_scale: [0.1,0.3,0.7,1.0] fixed_iters: True cost_convergence: 0.01 cost_delta_threshold: 2000.0 diff --git a/src/curobo/content/configs/task/particle_mpc.yml b/src/curobo/content/configs/task/particle_mpc.yml index d10544d..7939b8f 100644 --- a/src/curobo/content/configs/task/particle_mpc.yml +++ b/src/curobo/content/configs/task/particle_mpc.yml @@ -38,7 +38,7 @@ cost: use_metric: True cspace_cfg: - weight: 500.0 + weight: 1000.0 terminal: True run_weight: 1.0 diff --git a/src/curobo/content/configs/task/particle_trajopt.yml b/src/curobo/content/configs/task/particle_trajopt.yml index 23ae39a..268a613 100644 --- a/src/curobo/content/configs/task/particle_trajopt.yml +++ b/src/curobo/content/configs/task/particle_trajopt.yml @@ -34,7 +34,7 @@ cost: weight: [250.0, 5000.0, 20, 20] vec_convergence: [0.0,0.0,1000.0,1000.0] terminal: True - run_weight: 1.0 + run_weight: 0.0 use_metric: True link_pose_cfg: diff --git a/src/curobo/cuda_robot_model/cuda_robot_generator.py b/src/curobo/cuda_robot_model/cuda_robot_generator.py index 4845e82..0fbaf45 100644 --- a/src/curobo/cuda_robot_model/cuda_robot_generator.py +++ b/src/curobo/cuda_robot_model/cuda_robot_generator.py @@ -73,7 +73,7 @@ class CudaRobotGeneratorConfig: collision_spheres: Union[None, str, Dict[str, Any]] = None #: Radius buffer to add to collision spheres as padding. - collision_sphere_buffer: float = 0.0 + collision_sphere_buffer: Union[float, Dict[str, float]] = 0.0 #: Compute jacobian of link poses. Currently not supported. compute_jacobian: bool = False @@ -436,7 +436,6 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): if self._bodies[0].link_name in link_names: store_link_map.append(chain_link_names.index(self._bodies[0].link_name)) ordered_link_names.append(self._bodies[0].link_name) - joint_offset_map.append(self._bodies[0].joint_offset) # get joint types: for i in range(1, len(self._bodies)): body = self._bodies[i] @@ -614,7 +613,6 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): self.joint_names.remove(j) self._n_dofs -= 1 self._active_joints.remove(i) - self._joint_map[self._joint_map < -1] = -1 self._joint_map = self._joint_map.to(device=self.tensor_args.device) self._joint_map_type = self._joint_map_type.to(device=self.tensor_args.device) @@ -628,7 +626,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): self, collision_spheres: Dict, collision_link_names: List[str], - collision_sphere_buffer: float = 0.0, + collision_sphere_buffer: Union[float, Dict[str, float]] = 0.0, ): """ @@ -643,6 +641,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): # we store as [n_link, 7] link_sphere_idx_map = [] cpu_tensor_args = self.tensor_args.cpu() + self_collision_buffer = self.self_collision_buffer.copy() with profiler.record_function("robot_generator/build_collision_spheres"): for j_idx, j in enumerate(collision_link_names): # print(j_idx) @@ -652,11 +651,22 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): ) # find link index in global map: l_idx = self._name_to_idx_map[j] - + offset_radius = 0.0 + if isinstance(collision_sphere_buffer, float): + offset_radius = collision_sphere_buffer + elif j in collision_sphere_buffer: + offset_radius = collision_sphere_buffer[j] + if j in self_collision_buffer: + self_collision_buffer[j] -= offset_radius + else: + self_collision_buffer[j] = -offset_radius for i in range(n_spheres): + padded_radius = collision_spheres[j][i]["radius"] + offset_radius + if padded_radius <= 0.0 and padded_radius > -1.0: + padded_radius = 0.001 link_spheres[i, :] = tensor_sphere( collision_spheres[j][i]["center"], - collision_spheres[j][i]["radius"], + padded_radius, tensor_args=cpu_tensor_args, tensor=link_spheres[i, :], ) @@ -665,10 +675,6 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): self.total_spheres += n_spheres self._link_spheres_tensor = torch.cat(coll_link_spheres, dim=0) - new_radius = self._link_spheres_tensor[..., 3] + collision_sphere_buffer - flag = torch.logical_and(new_radius > -1.0, new_radius <= 0.0) - new_radius[flag] = 0.001 - self._link_spheres_tensor[:, 3] = new_radius self._link_sphere_idx_map = torch.as_tensor( link_sphere_idx_map, dtype=torch.int16, device=cpu_tensor_args.device ) @@ -696,9 +702,9 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): link1_spheres_idx = torch.nonzero(self._link_sphere_idx_map == link1_idx) rad1 = self._link_spheres_tensor[link1_spheres_idx, 3] - if j not in self.self_collision_buffer.keys(): - self.self_collision_buffer[j] = 0.0 - c1 = self.self_collision_buffer[j] + if j not in self_collision_buffer.keys(): + self_collision_buffer[j] = 0.0 + c1 = self_collision_buffer[j] self.self_collision_offset[link1_spheres_idx] = c1 for _, i_name in enumerate(collision_link_names): if i_name == j or i_name in ignore_links: @@ -706,9 +712,9 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): if i_name not in collision_link_names: log_error("Self Collision Link name not found in collision_link_names") # find index of this link name: - if i_name not in self.self_collision_buffer.keys(): - self.self_collision_buffer[i_name] = 0.0 - c2 = self.self_collision_buffer[i_name] + if i_name not in self_collision_buffer.keys(): + self_collision_buffer[i_name] = 0.0 + c2 = self_collision_buffer[i_name] link2_idx = self._name_to_idx_map[i_name] # update collision distance between spheres from these two links: link2_spheres_idx = torch.nonzero(self._link_sphere_idx_map == link2_idx) diff --git a/src/curobo/cuda_robot_model/cuda_robot_model.py b/src/curobo/cuda_robot_model/cuda_robot_model.py index 3fbd31f..3135bb5 100644 --- a/src/curobo/cuda_robot_model/cuda_robot_model.py +++ b/src/curobo/cuda_robot_model/cuda_robot_model.py @@ -143,6 +143,10 @@ class CudaRobotModelConfig: def cspace(self): return self.kinematics_config.cspace + @property + def dof(self) -> int: + return self.kinematics_config.n_dof + class CudaRobotModel(CudaRobotModelConfig): """ @@ -368,6 +372,10 @@ class CudaRobotModel(CudaRobotModelConfig): def get_dof(self) -> int: return self.kinematics_config.n_dof + @property + def dof(self) -> int: + return self.kinematics_config.n_dof + @property def joint_names(self) -> List[str]: return self.kinematics_config.joint_names diff --git a/src/curobo/cuda_robot_model/urdf_kinematics_parser.py b/src/curobo/cuda_robot_model/urdf_kinematics_parser.py index d627c7d..05b067d 100644 --- a/src/curobo/cuda_robot_model/urdf_kinematics_parser.py +++ b/src/curobo/cuda_robot_model/urdf_kinematics_parser.py @@ -68,12 +68,12 @@ class UrdfKinematicsParser(KinematicsParser): "velocity": joint.limit.velocity, } else: - log_warn("Converting continuous joint to revolute with limits[-10,10]") + log_warn("Converting continuous joint to revolute with limits[-6.28,6.28]") joint_type = "revolute" joint_limits = { "effort": joint.limit.effort, - "lower": -10, - "upper": 10, + "lower": -3.14 * 2, + "upper": 3.14 * 2, "velocity": joint.limit.velocity, } return joint_limits, joint_type @@ -181,6 +181,7 @@ class UrdfKinematicsParser(KinematicsParser): log_error("Joint type not supported") if joint_axis[0] == -1 or joint_axis[1] == -1 or joint_axis[2] == -1: joint_offset[0] = -1.0 * joint_offset[0] + joint_axis = [abs(x) for x in joint_axis] body_params["joint_type"] = joint_type body_params["joint_offset"] = joint_offset diff --git a/src/curobo/curobolib/cpp/check_cuda.h b/src/curobo/curobolib/cpp/check_cuda.h index dd48d44..f951605 100644 --- a/src/curobo/curobolib/cpp/check_cuda.h +++ b/src/curobo/curobolib/cpp/check_cuda.h @@ -18,4 +18,12 @@ #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) #define CHECK_FP8 defined(CUDA_VERSION) && CUDA_VERSION >= 11080 && TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 2 -#define CHECK_INPUT_GUARD(x) CHECK_INPUT(x); const at::cuda::OptionalCUDAGuard guard(x.device()) \ No newline at end of file +#define CHECK_INPUT_GUARD(x) CHECK_INPUT(x); const at::cuda::OptionalCUDAGuard guard(x.device()) + +#if CHECK_FP8 + #define FP8_TYPE_MACRO torch::kFloat8_e4m3fn + //constexpr const auto fp8_type = torch::kFloat8_e4m3fn; +#else + #define FP8_TYPE_MACRO torch::kHalf + //const constexpr auto fp8_type = torch::kHalf; +#endif \ No newline at end of file diff --git a/src/curobo/curobolib/cpp/kinematics_fused_kernel.cu b/src/curobo/curobolib/cpp/kinematics_fused_kernel.cu index f64b9c2..8640d62 100644 --- a/src/curobo/curobolib/cpp/kinematics_fused_kernel.cu +++ b/src/curobo/curobolib/cpp/kinematics_fused_kernel.cu @@ -1341,13 +1341,14 @@ std::vectorkin_fused_backward_16t( cudaStream_t stream = at::cuda::getCurrentCUDAStream(); assert(sparsity_opt); + const bool parallel_write = true; if (use_global_cumul) { if (n_joints < 16) { AT_DISPATCH_FLOATING_TYPES( grad_spheres.scalar_type(), "kin_fused_backward_16t", ([&] { - kin_fused_backward_kernel3 + kin_fused_backward_kernel3 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( grad_out.data_ptr(), grad_nlinks_pos.data_ptr(), @@ -1371,7 +1372,7 @@ std::vectorkin_fused_backward_16t( { AT_DISPATCH_FLOATING_TYPES( grad_spheres.scalar_type(), "kin_fused_backward_16t", ([&] { - kin_fused_backward_kernel3 + kin_fused_backward_kernel3 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( grad_out.data_ptr(), grad_nlinks_pos.data_ptr(), @@ -1395,7 +1396,7 @@ std::vectorkin_fused_backward_16t( { AT_DISPATCH_FLOATING_TYPES( grad_spheres.scalar_type(), "kin_fused_backward_16t", ([&] { - kin_fused_backward_kernel3 + kin_fused_backward_kernel3 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( grad_out.data_ptr(), grad_nlinks_pos.data_ptr(), @@ -1423,7 +1424,7 @@ std::vectorkin_fused_backward_16t( // AT_DISPATCH_FLOATING_TYPES( grad_spheres.scalar_type(), "kin_fused_backward_16t", ([&] { - kin_fused_backward_kernel3 + kin_fused_backward_kernel3 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( grad_out.data_ptr(), grad_nlinks_pos.data_ptr(), diff --git a/src/curobo/curobolib/cpp/sphere_obb_kernel.cu b/src/curobo/curobolib/cpp/sphere_obb_kernel.cu index df3da00..da250af 100644 --- a/src/curobo/curobolib/cpp/sphere_obb_kernel.cu +++ b/src/curobo/curobolib/cpp/sphere_obb_kernel.cu @@ -481,7 +481,7 @@ namespace Curobo - + } delta = make_float3(pt.x - sphere.x, pt.y - sphere.y, pt.z - sphere.z); @@ -2746,13 +2746,8 @@ sphere_obb_clpt(const torch::Tensor sphere_position, // batch_size, 3 else { - #if CHECK_FP8 - const auto fp8_type = torch::kFloat8_e4m3fn; - #else - const auto fp8_type = torch::kHalf; - #endif // typename scalar_t, typename dist_scalar_t=float, bool BATCH_ENV_T=true, bool SCALE_METRIC=true, bool SUM_COLLISIONS=true, bool COMPUTE_ESDF=false - AT_DISPATCH_FLOATING_TYPES_AND2(torch::kBFloat16, fp8_type, + AT_DISPATCH_FLOATING_TYPES_AND2(torch::kBFloat16, FP8_TYPE_MACRO, distance.scalar_type(), "SphereObb_clpt", ([&]{ auto distance_kernel = sphere_obb_distance_kernel; if (use_batch_env) @@ -3211,13 +3206,8 @@ sphere_voxel_clpt(const torch::Tensor sphere_position, // batch_size, 3 int blocksPerGrid = (bnh_spheres + threadsPerBlock - 1) / threadsPerBlock; - #if CHECK_FP8 - const auto fp8_type = torch::kFloat8_e4m3fn; - #else - const auto fp8_type = torch::kHalf; - #endif - AT_DISPATCH_FLOATING_TYPES_AND2(torch::kBFloat16, fp8_type, + AT_DISPATCH_FLOATING_TYPES_AND2(torch::kBFloat16, FP8_TYPE_MACRO, grid_features.scalar_type(), "SphereVoxel_clpt", ([&] { @@ -3255,7 +3245,6 @@ sphere_voxel_clpt(const torch::Tensor sphere_position, // batch_size, 3 } } - selected_kernel << < blocksPerGrid, threadsPerBlock, 0, stream >> > ( sphere_position.data_ptr(), @@ -3275,7 +3264,6 @@ sphere_voxel_clpt(const torch::Tensor sphere_position, // batch_size, 3 num_voxels, batch_size, horizon, n_spheres, transform_back); - })); @@ -3346,13 +3334,7 @@ swept_sphere_voxel_clpt(const torch::Tensor sphere_position, // batch_size, 3 int blocksPerGrid = (bnh_spheres + threadsPerBlock - 1) / threadsPerBlock; - #if CHECK_FP8 - const auto fp8_type = torch::kFloat8_e4m3fn; - #else - const auto fp8_type = torch::kHalf; - #endif - - AT_DISPATCH_FLOATING_TYPES_AND2(torch::kBFloat16, fp8_type, + AT_DISPATCH_FLOATING_TYPES_AND2(torch::kBFloat16, FP8_TYPE_MACRO, grid_features.scalar_type(), "SphereVoxel_clpt", ([&] { auto collision_kernel_n = swept_sphere_voxel_distance_jump_kernel; diff --git a/src/curobo/geom/sdf/warp_primitives.py b/src/curobo/geom/sdf/warp_primitives.py index b748122..a5ac772 100644 --- a/src/curobo/geom/sdf/warp_primitives.py +++ b/src/curobo/geom/sdf/warp_primitives.py @@ -525,7 +525,7 @@ class SdfMeshWarpPy(torch.autograd.Function): if env_query_idx is None: use_batch_env = False env_query_idx = n_env_mesh - + requires_grad = query_spheres.requires_grad wp.launch( kernel=get_closest_pt_batch_env, dim=b * h * n, @@ -541,7 +541,7 @@ class SdfMeshWarpPy(torch.autograd.Function): wp.from_torch(mesh_enable.view(-1), dtype=wp.uint8), wp.from_torch(n_env_mesh.view(-1), dtype=wp.int32), wp.from_torch(max_dist, dtype=wp.float32), - query_spheres.requires_grad, + requires_grad, b, h, n, @@ -608,6 +608,7 @@ class SweptSdfMeshWarpPy(torch.autograd.Function): if env_query_idx is None: use_batch_env = False env_query_idx = n_env_mesh + requires_grad = query_spheres.requires_grad wp.launch( kernel=get_swept_closest_pt_batch_env, @@ -625,7 +626,7 @@ class SweptSdfMeshWarpPy(torch.autograd.Function): wp.from_torch(mesh_enable.view(-1), dtype=wp.uint8), wp.from_torch(n_env_mesh.view(-1), dtype=wp.int32), wp.from_torch(max_dist, dtype=wp.float32), - query_spheres.requires_grad, + requires_grad, b, h, n, diff --git a/src/curobo/geom/sdf/world_voxel.py b/src/curobo/geom/sdf/world_voxel.py index b3238d5..9eb9c62 100644 --- a/src/curobo/geom/sdf/world_voxel.py +++ b/src/curobo/geom/sdf/world_voxel.py @@ -42,6 +42,7 @@ class WorldVoxelCollision(WorldMeshCollision): and self.cache["voxel"] not in [None, 0] ): self._create_voxel_cache(self.cache["voxel"]) + return super()._init_cache() def _create_voxel_cache(self, voxel_cache: Dict[str, Any]): n_layers = voxel_cache["layers"] @@ -699,7 +700,6 @@ class WorldVoxelCollision(WorldMeshCollision): - self.max_esdf_distance ).to(dtype=self._voxel_tensor_list[3].dtype) self._env_n_voxels[:] = 0 - print(self._voxel_tensor_list) def get_voxel_grid_shape(self, env_idx: int = 0, obs_idx: int = 0): return self._voxel_tensor_list[3][env_idx, obs_idx].shape diff --git a/src/curobo/rollout/cost/bound_cost.py b/src/curobo/rollout/cost/bound_cost.py index 0c70194..3524cf0 100644 --- a/src/curobo/rollout/cost/bound_cost.py +++ b/src/curobo/rollout/cost/bound_cost.py @@ -355,6 +355,8 @@ class WarpBoundSmoothFunction(torch.autograd.Function): wp_device = wp.device_from_torch(vel.device) # assert smooth_weight.shape[0] == 7 b, h, dof = vel.shape + requires_grad = pos.requires_grad + wp.launch( kernel=forward_bound_smooth_warp, dim=b * h * dof, @@ -383,7 +385,7 @@ class WarpBoundSmoothFunction(torch.autograd.Function): wp.from_torch(out_gv.view(-1), dtype=wp.float32), wp.from_torch(out_ga.view(-1), dtype=wp.float32), wp.from_torch(out_gj.view(-1), dtype=wp.float32), - pos.requires_grad, + requires_grad, b, h, dof, @@ -471,6 +473,7 @@ class WarpBoundFunction(torch.autograd.Function): ): wp_device = wp.device_from_torch(vel.device) b, h, dof = vel.shape + requires_grad = pos.requires_grad wp.launch( kernel=forward_bound_warp, dim=b * h * dof, @@ -494,7 +497,7 @@ class WarpBoundFunction(torch.autograd.Function): wp.from_torch(out_gv.view(-1), dtype=wp.float32), wp.from_torch(out_ga.view(-1), dtype=wp.float32), wp.from_torch(out_gj.view(-1), dtype=wp.float32), - pos.requires_grad, + requires_grad, b, h, dof, @@ -505,6 +508,7 @@ class WarpBoundFunction(torch.autograd.Function): ctx.save_for_backward(out_gp, out_gv, out_ga, out_gj) # out_c = out_cost # out_c = torch.linalg.norm(out_cost, dim=-1) + out_c = torch.sum(out_cost, dim=-1) return out_c @@ -569,11 +573,11 @@ class WarpBoundPosFunction(torch.autograd.Function): ): wp_device = wp.device_from_torch(pos.device) b, h, dof = pos.shape + requires_grad = pos.requires_grad wp.launch( kernel=forward_bound_pos_warp, dim=b * h * dof, inputs=[ - # wp.from_torch(pos.detach().view(-1).contiguous(), dtype=wp.float32), wp.from_torch(pos.detach().view(-1), dtype=wp.float32), wp.from_torch(retract_config.detach().view(-1), dtype=wp.float32), wp.from_torch(retract_idx.detach().view(-1), dtype=wp.int32), @@ -584,7 +588,7 @@ class WarpBoundPosFunction(torch.autograd.Function): wp.from_torch(vec_weight.view(-1), dtype=wp.float32), wp.from_torch(out_cost.view(-1), dtype=wp.float32), wp.from_torch(out_gp.view(-1), dtype=wp.float32), - pos.requires_grad, + requires_grad, b, h, dof, @@ -685,23 +689,33 @@ def forward_bound_pos_warp( c_total = n_w * error * error g_p = 2.0 * n_w * error + # bound cost: + # if c_p < p_l: + # delta = p_l - c_p + # if (delta) > eta_p or eta_p == 0.0: + # c_total += w * (delta - 0.5 * eta_p) + # g_p += -w + # else: + # c_total += w * (0.5 / eta_p) * delta * delta + # g_p += -w * (1.0 / eta_p) * delta + # elif c_p > p_u: + # delta = c_p - p_u + # if (delta) > eta_p or eta_p == 0.0: + # c_total += w * (delta - 0.5 * eta_p) + # g_p += w + # else: + # c_total += w * (0.5 / eta_p) * delta * delta + # g_p += w * (1.0 / eta_p) * delta + # bound cost: if c_p < p_l: - delta = p_l - c_p - if (delta) > eta_p or eta_p == 0.0: - c_total += w * (delta - 0.5 * eta_p) - g_p += -w - else: - c_total += w * (0.5 / eta_p) * delta * delta - g_p += -w * (1.0 / eta_p) * delta + delta = c_p - p_l + c_total += w * delta * delta + g_p += 2.0 * w * delta elif c_p > p_u: delta = c_p - p_u - if (delta) > eta_p or eta_p == 0.0: - c_total += w * (delta - 0.5 * eta_p) - g_p += w - else: - c_total += w * (0.5 / eta_p) * delta * delta - g_p += w * (1.0 / eta_p) * delta + c_total += w * delta * delta + g_p += 2.0 * w * delta out_cost[b_addrs] = c_total @@ -811,73 +825,43 @@ def forward_bound_warp( g_p = 2.0 * n_w * error # bound cost: + delta = 0.0 if c_p < p_l: - delta = p_l - c_p - if (delta) > eta_p or eta_p == 0.0: - c_total += w * (delta - 0.5 * eta_p) - g_p += -w - else: - c_total += w * (0.5 / eta_p) * delta * delta - g_p += -w * (1.0 / eta_p) * delta + delta = c_p - p_l elif c_p > p_u: delta = c_p - p_u - if (delta) > eta_p or eta_p == 0.0: - c_total += w * (delta - 0.5 * eta_p) - g_p += w - else: - c_total += w * (0.5 / eta_p) * delta * delta - g_p += w * (1.0 / eta_p) * delta + c_total += w * delta * delta + g_p += 2.0 * w * delta + + # bound cost: + delta = 0.0 if c_v < v_l: - delta = v_l - c_v - if (delta) > eta_v or eta_v == 0.0: - c_total += b_wv * (delta - 0.5 * eta_v) - g_v = -b_wv - else: - c_total += b_wv * (0.5 / eta_v) * delta * delta - g_v = -b_wv * (1.0 / eta_v) * delta + delta = c_v - v_l elif c_v > v_u: delta = c_v - v_u - if (delta) > eta_v or eta_v == 0.0: - c_total += b_wv * (delta - 0.5 * eta_v) - g_v = b_wv - else: - c_total += b_wv * (0.5 / eta_v) * delta * delta - g_v = b_wv * (1.0 / eta_v) * delta + + c_total += b_wv * delta * delta + g_v = 2.0 * b_wv * delta + + delta = 0.0 if c_a < a_l: - delta = a_l - c_a - if (delta) > eta_a or eta_a == 0.0: - c_total += b_wa * (delta - 0.5 * eta_a) - g_a = -b_wa - else: - c_total += b_wa * (0.5 / eta_a) * delta * delta - g_a = -b_wa * (1.0 / eta_a) * delta + delta = c_a - a_l elif c_a > a_u: delta = c_a - a_u - if (delta) > eta_a or eta_a == 0.0: - c_total += b_wa * (delta - 0.5 * eta_a) - g_a = b_wa - else: - c_total += b_wa * (0.5 / eta_a) * delta * delta - g_a = b_wa * (1.0 / eta_a) * delta + c_total += b_wa * delta * delta + g_a = b_wa * 2.0 * delta + + delta = 0.0 if c_j < j_l: - delta = j_l - c_j - if (delta) > eta_j or eta_j == 0.0: - c_total += b_wj * (delta - 0.5 * eta_j) - g_j = -b_wj - else: - c_total += b_wj * (0.5 / eta_j) * delta * delta - g_j = -b_wj * (1.0 / eta_j) * delta + delta = c_j - j_l elif c_j > j_u: delta = c_j - j_u - if (delta) > eta_j or eta_j == 0.0: - c_total += b_wj * (delta - 0.5 * eta_j) - g_j = b_wj - else: - c_total += b_wj * (0.5 / eta_j) * delta * delta - g_j = b_wj * (1.0 / eta_j) * delta + + c_total += b_wj * delta * delta + g_j = b_wj * delta * 2.0 out_cost[b_addrs] = c_total @@ -1031,75 +1015,45 @@ def forward_bound_smooth_warp( g_p = 2.0 * n_w * error # bound cost: + # bound cost: + delta = 0.0 if c_p < p_l: - delta = p_l - c_p - if (delta) > eta_p or eta_p == 0.0: - c_total += w * (delta - 0.5 * eta_p) - g_p += -w - else: - c_total += w * (0.5 / eta_p) * delta * delta - g_p += -w * (1.0 / eta_p) * delta + delta = c_p - p_l elif c_p > p_u: delta = c_p - p_u - if (delta) > eta_p or eta_p == 0.0: - c_total += w * (delta - 0.5 * eta_p) - g_p += w - else: - c_total += w * (0.5 / eta_p) * delta * delta - g_p += w * (1.0 / eta_p) * delta + c_total += w * delta * delta + g_p += 2.0 * w * delta + + # bound cost: + delta = 0.0 if c_v < v_l: - delta = v_l - c_v - if (delta) > eta_v or eta_v == 0.0: - c_total += b_wv * (delta - 0.5 * eta_v) - g_v = -b_wv - else: - c_total += b_wv * (0.5 / eta_v) * delta * delta - g_v = -b_wv * (1.0 / eta_v) * delta + delta = c_v - v_l elif c_v > v_u: delta = c_v - v_u - if (delta) > eta_v or eta_v == 0.0: - c_total += b_wv * (delta - 0.5 * eta_v) - g_v = b_wv - else: - c_total += b_wv * (0.5 / eta_v) * delta * delta - g_v = b_wv * (1.0 / eta_v) * delta + + c_total += b_wv * delta * delta + g_v = 2.0 * b_wv * delta + + delta = 0.0 + if c_a < a_l: - delta = a_l - c_a - if (delta) > eta_a or eta_a == 0.0: - c_total += b_wa * (delta - 0.5 * eta_a) - g_a = -b_wa - else: - c_total += b_wa * (0.5 / eta_a) * delta * delta - g_a = -b_wa * (1.0 / eta_a) * delta + delta = c_a - a_l elif c_a > a_u: delta = c_a - a_u - if (delta) > eta_a or eta_a == 0.0: - c_total += b_wa * (delta - 0.5 * eta_a) - g_a = b_wa - else: - c_total += b_wa * (0.5 / eta_a) * delta * delta - g_a = b_wa * (1.0 / eta_a) * delta + + c_total += b_wa * delta * delta + g_a = b_wa * 2.0 * delta + + delta = 0.0 if c_j < j_l: - delta = j_l - c_j - if (delta) > eta_j or eta_j == 0.0: - c_total += b_wj * (delta - 0.5 * eta_j) - g_j = -b_wj - else: - c_total += b_wj * (0.5 / eta_j) * delta * delta - g_j = -b_wj * (1.0 / eta_j) * delta + delta = c_j - j_l elif c_j > j_u: delta = c_j - j_u - if (delta) > eta_j or eta_j == 0.0: - c_total += b_wj * (delta - 0.5 * eta_j) - g_j = b_wj - else: - c_total += b_wj * (0.5 / eta_j) * delta * delta - g_j = b_wj * (1.0 / eta_j) * delta - # g_v = -1.0 * g_v - # g_a = -1.0 * g_a - # g_j = - 1.0 + c_total += b_wj * delta * delta + g_j = b_wj * delta * 2.0 + # do l2 regularization for velocity: if r_wv < 1.0: s_v = w_v * r_wv * c_v * c_v diff --git a/src/curobo/rollout/cost/dist_cost.py b/src/curobo/rollout/cost/dist_cost.py index 58d9f4c..543b772 100644 --- a/src/curobo/rollout/cost/dist_cost.py +++ b/src/curobo/rollout/cost/dist_cost.py @@ -128,12 +128,12 @@ def forward_l2_warp( target_p = target[target_id] error = c_p - target_p - if r_w >= 1.0 and w > 100.0: - c_total = w * wp.log2(wp.cosh(50.0 * error)) - g_p = w * 50.0 * wp.sinh(50.0 * error) / (wp.cosh(50.0 * error)) - else: - c_total = w * error * error - g_p = 2.0 * w * error + # if r_w >= 1.0 and w > 100.0: + # c_total = w * wp.log2(wp.cosh(10.0 * error)) + # g_p = w * 10.0 * wp.sinh(10.0 * error) / (wp.cosh(10.0 * error)) + # else: + c_total = w * error * error + g_p = 2.0 * w * error out_cost[b_addrs] = c_total @@ -159,8 +159,7 @@ class L2DistFunction(torch.autograd.Function): ): wp_device = wp.device_from_torch(pos.device) b, h, dof = pos.shape - # print(target) - + requires_grad = pos.requires_grad wp.launch( kernel=forward_l2_warp, dim=b * h * dof, @@ -173,7 +172,7 @@ class L2DistFunction(torch.autograd.Function): wp.from_torch(vec_weight.view(-1), dtype=wp.float32), wp.from_torch(out_cost_v.view(-1), dtype=wp.float32), wp.from_torch(out_gp.view(-1), dtype=wp.float32), - pos.requires_grad, + requires_grad, b, h, dof, @@ -181,11 +180,8 @@ class L2DistFunction(torch.autograd.Function): device=wp_device, stream=wp.stream_from_torch(pos.device), ) - # cost = torch.linalg.norm(out_cost_v, dim=-1) - # if pos.requires_grad: - # out_gp = out_gp * torch.nan_to_num( 1.0/cost.unsqueeze(-1), 0.0) - cost = torch.sum(out_cost_v, dim=-1) + cost = torch.sum(out_cost_v, dim=-1) ctx.save_for_backward(out_gp) return cost @@ -277,7 +273,11 @@ class DistCost(CostBase, DistCostConfig): self._run_weight_vec[:, :-1] *= self.run_weight cost = self._run_weight_vec * dist if RETURN_GOAL_DIST: - return cost, dist / self.weight + dist_scale = torch.nan_to_num( + 1.0 / torch.sqrt((self.weight * self._run_weight_vec)), 0.0 + ) + + return cost, dist * dist_scale return cost def forward_target_idx(self, goal_vec, current_vec, goal_idx, RETURN_GOAL_DIST=False): @@ -292,7 +292,6 @@ class DistCost(CostBase, DistCostConfig): self._run_weight_vec[:, :-1] *= self.run_weight else: raise NotImplementedError("terminal flag needs to be set to true") - if self.dist_type == DistType.L2: # print(goal_idx.shape, goal_vec.shape) cost = L2DistFunction.apply( @@ -306,11 +305,12 @@ class DistCost(CostBase, DistCostConfig): self._out_cv_buffer, self._out_g_buffer, ) - # cost = torch.linalg.norm(cost, dim=-1) else: raise NotImplementedError() - # print(cost.shape, cost[:,-1]) if RETURN_GOAL_DIST: - return cost, (cost / torch.sqrt((self.weight * self._run_weight_vec))) + dist_scale = torch.nan_to_num( + 1.0 / torch.sqrt((self.weight * self._run_weight_vec)), 0.0 + ) + return cost, cost * dist_scale return cost diff --git a/src/curobo/rollout/cost/pose_cost.py b/src/curobo/rollout/cost/pose_cost.py index e232b28..fd42ae3 100644 --- a/src/curobo/rollout/cost/pose_cost.py +++ b/src/curobo/rollout/cost/pose_cost.py @@ -105,9 +105,9 @@ class PoseCostMetric: @classmethod def create_grasp_approach_metric( cls, - offset_position: float = 0.5, + offset_position: float = 0.1, linear_axis: int = 2, - tstep_fraction: float = 0.6, + tstep_fraction: float = 0.8, tensor_args: TensorDeviceType = TensorDeviceType(), ) -> PoseCostMetric: """Enables moving to a pregrasp and then locked orientation movement to final grasp. @@ -203,7 +203,6 @@ class PoseCost(CostBase, PoseCostConfig): self.offset_waypoint[:3].copy_(offset_rotation) self.offset_tstep_fraction[:] = offset_tstep_fraction if self._horizon <= 0: - print(self.weight) log_error( "Updating offset waypoint is only possible after initializing motion gen" + " run motion_gen.warmup() before adding offset_waypoint" diff --git a/src/curobo/util/usd_helper.py b/src/curobo/util/usd_helper.py index 3767351..4d9f8de 100644 --- a/src/curobo/util/usd_helper.py +++ b/src/curobo/util/usd_helper.py @@ -823,7 +823,6 @@ class UsdHelper: config_file["robot_cfg"]["kinematics"], tensor_args=tensor_args ) kin_model = CudaRobotModel(robot_cfg) - m = kin_model.get_robot_link_meshes() offsets = [x.pose for x in m] robot_mesh_model = WorldConfig(mesh=m) diff --git a/src/curobo/wrap/__init__.py b/src/curobo/wrap/__init__.py index a08745f..db172db 100644 --- a/src/curobo/wrap/__init__.py +++ b/src/curobo/wrap/__init__.py @@ -8,3 +8,4 @@ # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. # +""" Module contains high-level APIs for robotics applications. """ diff --git a/src/curobo/wrap/reacher/__init__.py b/src/curobo/wrap/reacher/__init__.py index a08745f..20c97d8 100644 --- a/src/curobo/wrap/reacher/__init__.py +++ b/src/curobo/wrap/reacher/__init__.py @@ -8,3 +8,4 @@ # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. # +""" Module contains commonly used high-level APIs for motion generation. """ diff --git a/src/curobo/wrap/reacher/evaluator.py b/src/curobo/wrap/reacher/evaluator.py index 66926fb..2dfc351 100644 --- a/src/curobo/wrap/reacher/evaluator.py +++ b/src/curobo/wrap/reacher/evaluator.py @@ -8,33 +8,116 @@ # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. # +""" This modules contains heuristics for scoring trajectories. """ + +from __future__ import annotations # Standard Library from dataclasses import dataclass -from typing import Optional +from typing import Optional, Tuple, Union # Third Party import torch import torch.autograd.profiler as profiler # CuRobo +from curobo.types.base import TensorDeviceType from curobo.types.robot import JointState from curobo.types.tensor import T_DOF +from curobo.util.logger import log_error from curobo.util.torch_utils import get_torch_jit_decorator from curobo.util.trajectory import calculate_dt @dataclass class TrajEvaluatorConfig: - max_acc: float = 15.0 - max_jerk: float = 500.0 + """Configurable Parameters for Trajectory Evaluator.""" + + #: Maximum acceleration for each joint. + max_acc: torch.Tensor + + #: Maximum jerk for each joint. + max_jerk: torch.Tensor + + #: Minimum allowed time step for trajectory. Trajectories with time step less than this + #: value will be rejected. + min_dt: torch.Tensor + + #: Maximum allowed time step for trajectory. Trajectories with time step greater than this + #: value will be rejected. + max_dt: torch.Tensor + + #: Weight to scale smoothness cost, total cost = path length + cost_weight * smoothness cost. cost_weight: float = 0.01 - min_dt: float = 0.001 - max_dt: float = 0.1 + + def __post_init__(self): + """Checks if values are of correct type and converts them if possible.""" + if not isinstance(self.max_acc, torch.Tensor): + log_error( + "max_acc should be a torch.Tensor, this was changed recently, use " + + "TrajEvaluatorConfig.from_basic() to create TrajEvaluatorConfig object" + ) + if not isinstance(self.max_jerk, torch.Tensor): + log_error( + "max_jerk should be a torch.Tensor, this was changed recently, use " + + "TrajEvaluatorConfig.from_basic() to create TrajEvaluatorConfig object" + ) + if not isinstance(self.min_dt, torch.Tensor): + self.min_dt = torch.as_tensor( + self.min_dt, device=self.max_acc.device, dtype=self.max_acc.dtype + ) + if not isinstance(self.max_dt, torch.Tensor): + self.max_dt = torch.as_tensor( + self.max_dt, device=self.max_acc.device, dtype=self.max_acc.dtype + ) + + @staticmethod + def from_basic( + dof: int, + max_acc: float = 15.0, + max_jerk: float = 500.0, + cost_weight: float = 0.01, + min_dt: float = 0.001, + max_dt: float = 0.1, + tensor_args: TensorDeviceType = TensorDeviceType(), + ) -> TrajEvaluatorConfig: + """Creates TrajEvaluatorConfig object from basic parameters. + + Args: + dof: number of active joints in the robot. + max_acc: maximum acceleration for all joints. Treats this as same for all joints. + max_jerk: maximum jerk for all joints. Treats this as same for all joints. + cost_weight: weight to scale smoothness cost. + min_dt: minimum allowed time step between waypoints of a trajectory. + max_dt: maximum allowed time step between waypoints of a trajectory. + tensor_args: device and dtype for the tensors. + + Returns: + TrajEvaluatorConfig: Configured Parameters for Trajectory Evaluator. + """ + return TrajEvaluatorConfig( + max_acc=torch.full((dof,), max_acc, device=tensor_args.device, dtype=tensor_args.dtype), + max_jerk=torch.full( + (dof,), max_jerk, device=tensor_args.device, dtype=tensor_args.dtype + ), + cost_weight=cost_weight, + min_dt=torch.as_tensor(min_dt, device=tensor_args.device, dtype=tensor_args.dtype), + max_dt=torch.as_tensor(max_dt, device=tensor_args.device, dtype=tensor_args.dtype), + ) @get_torch_jit_decorator() -def compute_path_length(vel, traj_dt, cspace_distance_weight): +def compute_path_length(vel, traj_dt, cspace_distance_weight) -> torch.Tensor: + """JIT compatible function to compute path length. + + Args: + vel: joint space velocity tensor of shape (batch, horizon, dof). + traj_dt: dt of trajectory tensor of shape (batch, horizon) or (1, horizon) or (1, 1). + cspace_distance_weight: weight tensor of shape (dof). + + Returns: + torch.Tensor: path length tensor of shape (batch). + """ pl = torch.sum( torch.sum(torch.abs(vel) * traj_dt.unsqueeze(-1) * cspace_distance_weight, dim=-1), dim=-1 ) @@ -42,19 +125,35 @@ def compute_path_length(vel, traj_dt, cspace_distance_weight): @get_torch_jit_decorator() -def compute_path_length_cost(vel, cspace_distance_weight): +def compute_path_length_cost(vel, cspace_distance_weight) -> torch.Tensor: + """JIT compatible function to compute path length cost without considering time step. + + Args: + vel: joint space velocity tensor of shape (batch, horizon, dof). + cspace_distance_weight: weight tensor of shape (dof). + + Returns: + torch.Tensor: path length cost tensor of shape (batch). + """ pl = torch.sum(torch.sum(torch.abs(vel) * cspace_distance_weight, dim=-1), dim=-1) return pl @get_torch_jit_decorator() -def smooth_cost(abs_acc, abs_jerk, opt_dt): - # acc = torch.max(torch.max(abs_acc, dim=-1)[0], dim=-1)[0] - # jerk = torch.max(torch.max(abs_jerk, dim=-1)[0], dim=-1)[0] +def smooth_cost(abs_acc, abs_jerk, opt_dt) -> torch.Tensor: + """JIT compatible function to compute smoothness cost. + + Args: + abs_acc: absolute acceleration tensor of shape (batch, horizon, dof). + abs_jerk: absolute jerk tensor of shape (batch, horizon, dof). + opt_dt: optimal time step tensor of shape (batch). + + Returns: + torch.Tensor: smoothness cost tensor of shape (batch). + """ jerk = torch.mean(torch.max(abs_jerk, dim=-1)[0], dim=-1) mean_acc = torch.mean(torch.max(abs_acc, dim=-1)[0], dim=-1) # [0] a = (jerk * 0.001) + 10.0 * opt_dt + (mean_acc * 0.01) - # a = (jerk * 0.001) + 50.0 * opt_dt + (mean_acc * 0.01) return a @@ -65,24 +164,44 @@ def compute_smoothness( acc: torch.Tensor, jerk: torch.Tensor, max_vel: torch.Tensor, - max_acc: float, - max_jerk: float, + max_acc: torch.Tensor, + max_jerk: torch.Tensor, traj_dt: torch.Tensor, min_dt: float, max_dt: float, -): +) -> Tuple[torch.Tensor, torch.Tensor]: + """JIT compatible function to compute smoothness. + + Args: + vel: joint space velocity tensor of shape (batch, horizon, dof). + acc: joint space acceleration tensor of shape (batch, horizon, dof). + jerk: joint space jerk tensor of shape (batch, horizon, dof). + max_vel: maximum velocity limits, used to find scaling factor for dt that pushes at least + one joint to its limit, taking into account acceleration and jerk limits. + max_acc: maximum acceleration limits, used to find scaling factor for dt that pushes at least + one joint to its limit, taking into account velocity and jerk limits. + max_jerk: maximum jerk limits, used to find scaling factor for dt that pushes at least + one joint to its limit, taking into account velocity and acceleration limits. + traj_dt: dt of trajectory tensor of shape (batch, horizon) or (1, horizon) + min_dt: minimum delta time allowed between steps/waypoints in a trajectory. + max_dt: maximum delta time allowed between steps/waypoints in a trajectory. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: success label tensor of shape (batch) and + smoothness cost tensor of shape (batch) + """ # compute scaled dt: - # h = int(vel.shape[-2] / 2) max_v_arr = torch.max(torch.abs(vel), dim=-2)[0] # output is batch, dof scale_dt = (max_v_arr) / (max_vel.view(1, max_v_arr.shape[-1])) # batch,dof max_acc_arr = torch.max(torch.abs(acc), dim=-2)[0] # batch, dof - + max_acc = max_acc.view(1, max_acc_arr.shape[-1]) scale_dt_acc = torch.sqrt(torch.max(max_acc_arr / max_acc, dim=-1)[0]) # batch, 1 max_jerk_arr = torch.max(torch.abs(jerk), dim=-2)[0] # batch, dof + max_jerk = max_jerk.view(1, max_jerk_arr.shape[-1]) scale_dt_jerk = torch.pow(torch.max(max_jerk_arr / max_jerk, dim=-1)[0], 1.0 / 3.0) # batch, 1 dt_score_vel = torch.max(scale_dt, dim=-1)[0] # batch, 1 @@ -102,25 +221,55 @@ def compute_smoothness( # mean_jerk_val = torch.max(torch.mean(abs_jerk, dim=-1), dim=-1)[0] max_jerk_val = torch.max(torch.max(abs_jerk, dim=-1)[0], dim=-1)[0] acc_label = torch.logical_and(max_acc_val <= max_acc, max_jerk_val <= max_jerk) - # print(max_acc_val, max_jerk_val, dt_score) return (acc_label, smooth_cost(abs_acc, abs_jerk, dt_score)) @get_torch_jit_decorator() def compute_smoothness_opt_dt( - vel, acc, jerk, max_vel: torch.Tensor, max_acc: float, max_jerk: float, opt_dt: torch.Tensor -): - abs_acc = torch.abs(acc) - max_acc_val = torch.max(torch.max(abs_acc, dim=-1)[0], dim=-1)[0] - abs_jerk = torch.abs(jerk) - max_jerk_val = torch.max(torch.max(abs_jerk, dim=-1)[0], dim=-1)[0] + vel, + acc, + jerk, + max_vel: torch.Tensor, + max_acc: torch.Tensor, + max_jerk: torch.Tensor, + opt_dt: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """JIT compatible function to compute smoothness with pre-computed optimal time step. - acc_label = torch.logical_and(max_acc_val <= max_acc, max_jerk_val <= max_jerk) + Args: + vel: joint space velocity tensor of shape (batch, horizon, dof), not used in this + implementation. + acc: joint space acceleration tensor of shape (batch, horizon, dof). + jerk: joint space jerk tensor of shape (batch, horizon, dof). + max_vel: maximum velocity limit, not used in this implementation. + max_acc: maximum acceleration limit, used to reject trajectories. + max_jerk: maximum jerk limit, used to reject trajectories. + opt_dt: optimal time step tensor of shape (batch). + + Returns: + Tuple[torch.Tensor, torch.Tensor]: success label tensor of shape (batch) and + smoothness cost tensor of shape (batch). + """ + abs_acc = torch.abs(acc) + abs_jerk = torch.abs(jerk) + + delta_acc = torch.min(torch.min(max_acc - abs_acc, dim=-1)[0], dim=-1)[0] + + max_jerk = max_jerk.view(1, abs_jerk.shape[-1]) + delta_jerk = torch.min(torch.min(max_jerk - abs_jerk, dim=-1)[0], dim=-1)[0] + acc_label = torch.logical_and(delta_acc >= 0.0, delta_jerk >= 0.0) return acc_label, smooth_cost(abs_acc, abs_jerk, opt_dt) class TrajEvaluator(TrajEvaluatorConfig): + """Trajectory Evaluator class that uses heuristics to score trajectories.""" + def __init__(self, config: Optional[TrajEvaluatorConfig] = None): + """Initializes the TrajEvaluator object. + + Args: + config: Configurable parameters for Trajectory Evaluator. + """ if config is None: config = TrajEvaluatorConfig() super().__init__(**vars(config)) @@ -128,17 +277,41 @@ class TrajEvaluator(TrajEvaluatorConfig): def _compute_path_length( self, js: JointState, traj_dt: torch.Tensor, cspace_distance_weight: T_DOF ): + """Compute path length from joint velocities across trajectory and dt between them. + + Args: + js: joint state object with velocity tensor. + traj_dt: time step tensor of shape (batch, horizon) or (1, horizon) or (1, 1). + cspace_distance_weight: weight tensor of shape (dof). + + Returns: + torch.Tensor: path length tensor of shape (batch). + """ path_length = compute_path_length(js.velocity, traj_dt, cspace_distance_weight) return path_length - def _check_smoothness(self, js: JointState, traj_dt: torch.Tensor, max_vel: torch.Tensor): + def _check_smoothness( + self, js: JointState, traj_dt: torch.Tensor, max_vel: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Check smoothness of trajectory. + + Args: + js: joint state object with velocity, acceleration and jerk tensors. + traj_dt: time step tensor of shape (batch, horizon) or (1, horizon) or (1, 1). + max_vel: maximum velocity limits, used to find scaling factor for dt that pushes at + least one joint to its limit, taking into account acceleration and jerk limits. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: success label tensor of shape (batch) and + smoothness cost tensor of shape (batch). + """ if js.jerk is None: js.jerk = ( (torch.roll(js.acceleration, -1, -2) - js.acceleration) * (1 / traj_dt).unsqueeze(-1) )[..., :-1, :] - acc_label, max_acc = compute_smoothness( + smooth_success_label, smooth_cost = compute_smoothness( js.velocity, js.acceleration, js.jerk, @@ -149,7 +322,7 @@ class TrajEvaluator(TrajEvaluatorConfig): self.min_dt, self.max_dt, ) - return acc_label, max_acc + return smooth_success_label, smooth_cost @profiler.record_function("traj_evaluate_smoothness") def evaluate( @@ -158,7 +331,22 @@ class TrajEvaluator(TrajEvaluatorConfig): traj_dt: torch.Tensor, cspace_distance_weight: T_DOF, max_vel: torch.Tensor, - ): + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Evaluate trajectory based on smoothness and path length. + + Args: + js: joint state object with velocity, acceleration and jerk tensors. + traj_dt: time step tensor of shape (batch, horizon) or (1, horizon) or (1, 1) or + (batch, 1). + cspace_distance_weight: weight tensor of shape (dof). + max_vel: maximum velocity limits, used to find scaling factor for dt that pushes at + least one joint to its limit, taking into account acceleration and jerk limits. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: success label tensor of shape (batch) and + total cost tensor of shape (batch) where total cost = (path length + + cost_weight * smoothness cost). + """ label, cost = self._check_smoothness(js, traj_dt, max_vel) pl_cost = self._compute_path_length(js, traj_dt, cspace_distance_weight) return label, pl_cost + self.cost_weight * cost @@ -170,7 +358,21 @@ class TrajEvaluator(TrajEvaluatorConfig): opt_dt: torch.Tensor, cspace_distance_weight: T_DOF, max_vel: torch.Tensor, - ): + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Evaluate trajectory based on smoothness and path length with pre-computed optimal dt. + + Args: + js: joint state object with velocity, acceleration and jerk tensors. + opt_dt: optimal time step tensor of shape (batch). + cspace_distance_weight: weight tensor of shape (dof). + max_vel: maximum velocity limits, used to find scaling factor for dt that pushes at + least one joint to its limit, taking into account acceleration and jerk limits. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: success label tensor of shape (batch) and + total cost tensor of shape (batch) where total cost = (path length + + cost_weight * smoothness cost). + """ label, cost = compute_smoothness_opt_dt( js.velocity, js.acceleration, js.jerk, max_vel, self.max_acc, self.max_jerk, opt_dt ) @@ -185,7 +387,22 @@ class TrajEvaluator(TrajEvaluatorConfig): cspace_distance_weight: T_DOF, max_vel: torch.Tensor, skip_last_tstep: bool = False, - ): + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Evaluate trajectory by first computing velocity, acceleration and jerk from position. + + Args: + js: joint state object with position tensor. + traj_dt: time step tensor of shape (batch, 1) or (1, 1). + cspace_distance_weight: weight tensor of shape (dof). + max_vel: maximum velocity limits, used to find scaling factor for dt that pushes at + least one joint to its limit, taking into account acceleration and jerk limits. + skip_last_tstep: flag to skip last time step in trajectory. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: success label tensor of shape (batch) and + total cost tensor of shape (batch) where total cost = (path length + + cost_weight * smoothness cost). + """ js = js.calculate_fd_from_position(traj_dt) if skip_last_tstep: js.position = js.position[..., :-1, :] @@ -194,6 +411,3 @@ class TrajEvaluator(TrajEvaluatorConfig): js.jerk = js.jerk[..., :-1, :] return self.evaluate(js, traj_dt, cspace_distance_weight, max_vel) - - def calculate_dt(self, js: JointState, max_vel: torch.Tensor, raw_dt: float): - return calculate_dt(js.velocity, max_vel, raw_dt) diff --git a/src/curobo/wrap/reacher/ik_solver.py b/src/curobo/wrap/reacher/ik_solver.py index 5ea033c..0e7120a 100644 --- a/src/curobo/wrap/reacher/ik_solver.py +++ b/src/curobo/wrap/reacher/ik_solver.py @@ -8,6 +8,21 @@ # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. # + +""" +This module contains :meth:`IKSolver` to solve inverse kinematics, along with +:meth:`IKSolverConfig` to configure the solver, and :meth:`IKResult` to store the result. + +.. raw:: html + +

+ +

+This demo is available in :ref:`isaac_ik_reachability`. + +""" +from __future__ import annotations + # Standard Library from dataclasses import dataclass from typing import Any, Dict, List, Optional, Sequence, Union @@ -31,7 +46,7 @@ from curobo.rollout.rollout_base import Goal, RolloutBase, RolloutMetrics from curobo.types.base import TensorDeviceType from curobo.types.math import Pose from curobo.types.robot import JointState, RobotConfig -from curobo.types.tensor import T_BDOF, T_BValue_bool, T_BValue_float +from curobo.types.tensor import T_BDOF, T_DOF, T_BValue_bool, T_BValue_float from curobo.util.logger import log_error, log_warn from curobo.util.sample_lib import HaltonGenerator from curobo.util.torch_utils import get_torch_jit_decorator @@ -48,17 +63,54 @@ from curobo.wrap.wrap_base import WrapBase, WrapConfig, WrapResult @dataclass class IKSolverConfig: + """Configuration for Inverse Kinematics Solver. + + A helper function to load from robot + configuration is provided as :func:`IKSolverConfig.load_from_robot_config`. + + """ + + #: representation of robot, includes kinematic model, link geometry and joint limits. robot_config: RobotConfig + + #: Wrapped solver which has many instances of ArmReacher and two optimization + #: solvers (MPPI, LBFGS) as default. solver: WrapBase + + #: Number of seeds to optimize per IK problem in parallel. num_seeds: int + + #: Position convergence threshold in meters to use to compute success. position_threshold: float + + #: Rotation convergence threshold to use to compute success. Currently this + #: measures the geodesic distance between reached quaternion and target quaternion. + #: A good accuracy threshold is 0.05. Check pose_distance_kernel.cu for the exact + #: implementation. rotation_threshold: float + + #: Reference to an instance of ArmReacher rollout class to use for auxillary functions. rollout_fn: ArmReacher + + #: Undeveloped functionality to use a neural network as seed for IK. ik_nn_seeder: Optional[str] = None + + #: Reference to world collision checker to use for collision avoidance. world_coll_checker: Optional[WorldCollision] = None + + #: Rejection ratio for sampling collision-free configurations. These samples are not + #: used as seeds for solving IK as starting from collision-free seeds did not improve success. sample_rejection_ratio: int = 50 + + #: Device and floating precision to use for IKSolver. tensor_args: TensorDeviceType = TensorDeviceType() + + #: Flag to enable capturing solver iterations as a cuda graph to reduce kernel launch overhead. + #: Setting this to True can give upto 10x speedup while limiting the IKSolver to solving fixed + #: batch size of problems. use_cuda_graph: bool = True + + #: Seed to use in pseudorandom generator used for creating IK seeds. seed: int = 1531 @staticmethod @@ -96,7 +148,71 @@ class IKSolverConfig: high_precision: bool = False, project_pose_to_goal_frame: bool = True, seed: int = 1531, - ): + ) -> IKSolverConfig: + """Helper function to load IKSolver configuration from a robot file and world file. + + Use this function to create an instance of IKSolverConfig and load the config into IKSolver. + + Args: + robot_cfg: representation of robot, includes kinematic model, link geometry, and + joint limits. This can take as input a cuRobo robot configuration yaml file path or + a loaded dictionary of the robot configuration file or an instance of RobotConfig. + Configuration files for some common robots is available at + :ref:`available_robot_list`. For other robots, follow :ref:`tut_robot_configuration` + tutorial to create a configuration file. + world_model: representation of the world for obtaining collision-free IK solutions. + The world can be represented as cuboids, meshes, from depth camera with nvblox, and + as an Euclidean Signed Distance Grid (ESDF). This world model can be loaded from a + dictionary (from disk through yaml) or from :class:`curobo.geom.types.WorldConfig`. + In an application, if collision computations are being used in more than one + instance, it's memory efficient to create one instance of + :class:`curobo.geom.sdf.world.WorldCollision` and share across these class. For + example, if an instance of IKSolver and MotionGen exists in an application, a + :class:`curobo.geom.sdf.world.WorldCollision` object can be created in IKSolver + and then when creating :class:`curobo.wrap.reacher.motion_gen.MotionGenConfig`, this + ``world_coll_checker`` can be passed as reference. :ref:`world_collision` discusses + types of obstacles in more detail. + tensor_args: Device and floating precision to use for IKSolver. + num_seeds: Number of seeds to optimize per IK problem in parallel. + position_threshold: Position convergence threshold in meters to use to compute success. + rotation_threshold: Rotation convergence threshold to use to compute success. See + :meth:`IKSolverConfig.rotation_threshold` for more details. + world_coll_checker: Reference to world collision checker to use for collision avoidance. + base_cfg_file: Base configuration file for IKSolver. This configuration file is used to + measure convergence and collision-free check after optimization is complete. + particle_file: Configuration file for particle optimization (uses MPPI as optimizer). + gradient_file: Configuration file for gradient optimization (uses LBFGS as optimizer). + use_cuda_graph: Flag to enable capturing solver iterations as a cuda graph to reduce + kernel launch overhead. Setting this to True can give upto 10x speedup while + limiting the IKSolver to solving fixed batch size of problems. + self_collision_check: Flag to enable self-collision checking. + self_collision_opt: Flag to enable self-collision cost during optimization. + grad_iters: Number of iterations for gradient optimization. + use_particle_opt: Flag to enable particle optimization. + collision_checker_type: Type of collision checker to use for collision checking. + sync_cuda_time: Flag to enable synchronization of cuda device with host before + measuring compute time. If you set this to False, then measured time will not + include completion of any launched CUDA kernels. + use_gradient_descent: Flag to enable gradient descent optimization instead of LBFGS. + collision_cache: Number of objects to cache for collision checking. + n_collision_envs: Number of collision environments to use for IK. This is useful when + solving IK for multiple robots in different environments in parallel. + ee_link_name: Name of end-effector link to use for IK. + use_es: Flag to enable Evolution Strategies optimization instead of MPPI. + es_learning_rate: Learning rate for Evolution Strategies optimization. + use_fixed_samples: Flag to enable fixed samples for MPPI optimization. + store_debug: Flag to enable storing debug information during optimization. Enabling this + will store solution and cost at every iteration. This will significantly slow down + the optimization as CUDA graph capture is disabled. Only use this for debugging. + regularization: Flag to enable regularization during optimization. + collision_activation_distance: Distance from obstacle at which to activate collision + cost. A good value is 0.01 (1cm). + high_precision: Flag to solve IK at higher pose accuracy. This will increase the number + of LBFGS iterations from 100 to 200. This flag is set to True when + position_threshold is less than or equal to 1mm (0.001). + project_pose_to_goal_frame: Flag to project pose to goal frame when computing distance. + seed: Seed to use in pseudorandom generator used for creating IK seeds. + """ if position_threshold <= 0.001: high_precision = True if high_precision: @@ -106,12 +222,12 @@ class IKSolverConfig: base_config_data = load_yaml(join_path(get_task_configs_path(), base_cfg_file)) config_data = load_yaml(join_path(get_task_configs_path(), particle_file)) grad_config_data = load_yaml(join_path(get_task_configs_path(), gradient_file)) + base_config_data["cost"]["pose_cfg"]["project_distance"] = project_pose_to_goal_frame base_config_data["convergence"]["pose_cfg"]["project_distance"] = project_pose_to_goal_frame config_data["cost"]["pose_cfg"]["project_distance"] = project_pose_to_goal_frame grad_config_data["cost"]["pose_cfg"]["project_distance"] = project_pose_to_goal_frame - if collision_cache is not None: base_config_data["world_collision_checker_cfg"]["cache"] = collision_cache if n_collision_envs is not None: @@ -132,9 +248,7 @@ class IKSolverConfig: robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_cfg))["robot_cfg"] if ee_link_name is not None: if isinstance(robot_cfg, RobotConfig): - raise NotImplementedError("ee link cannot be changed after creating RobotConfig") - - robot_cfg.kinematics.ee_link = ee_link_name + log_error("ee link cannot be changed after creating RobotConfig") else: robot_cfg["kinematics"]["ee_link"] = ee_link_name if isinstance(robot_cfg, dict): @@ -261,21 +375,59 @@ class IKSolverConfig: @dataclass class IKResult(Sequence): + """Solution of Inverse Kinematics problem.""" + + #: Joint state solution for the IK problem. js_solution: JointState + + #: Goal pose used in IK problem. goal_pose: Pose + + #: Joint configuration Solution as tensor solution: T_BDOF + + #: Seed that was selected as the starting point for optimization. This is currently not + #: available in the result and is set to None. seed: T_BDOF + + #: Success tensor for IK problem. If planning for batch, use this to filter js_solution. If + #: planning for single problem, use this to check if the problem was solved successfully. success: T_BValue_bool + + #: Position error between solved pose and goal pose in meters, computed with l-2 norm. position_error: T_BValue_float - #: rotation error is computed as \sqrt(q_des^T * q) + #: Rotation error between solved pose and goal pose, computed with geodesic distance. Roughly + #: \sqrt(q_des^T * q). A good accuracy threshold is 0.05. rotation_error: T_BValue_float + + #: Total error for IK problem. This is the sum of position_error and rotation_error. error: T_BValue_float + + #: Time taken to solve the IK problem in seconds. solve_time: float + + #: Debug information from solver. This can be used to debug solver convergence and tune + #: weights between cost terms. debug_info: Optional[Any] = None + + #: Index of goal in goalset that the IK solver reached. This is useful when solving with + #: :meth:`IKSolver.solve_goalset` (also :meth:`IKSolver.solve_batch_goalset`, + #: :meth:`IKSolver.solve_batch_env_goalset`) where the task is to find an IK solution that + #: reaches 1 pose in set of poses. This index is used to identify which pose was reached by the + #: solution. goalset_index: Optional[torch.Tensor] = None - def __getitem__(self, idx): + def __getitem__(self, idx) -> IKResult: + """Get IKResult for a single problem in batch. + + Args: + idx: Index of the problem in batch. + + Returns: + IKResult for the problem at index idx. + """ + success = self.success[idx] return IKResult( @@ -290,10 +442,24 @@ class IKResult(Sequence): goalset_index=None if self.goalset_index is None else self.goalset_index[idx], ) - def __len__(self): + def __len__(self) -> int: + """Get number of problems in IKResult.""" return self.seed.shape[0] def get_unique_solution(self, roundoff_decimals: int = 2) -> torch.Tensor: + """Get unique solutions from many feasible solutions for the same problem. + + Use this after solving IK with :meth:`IKSolver.solve_single` with return_seeds > 1. This + function will return unique solutions from the set of feasible solutions, filering out + solutions that are within roundoff_decimals of each other. + + Args: + roundoff_decimals: Number of decimal places to round off the solution to measure + uniqueness. + + Returns: + Unique solutions from the set of feasible solutions. + """ in_solution = self.solution[self.success] r_sol = torch.round(in_solution, decimals=roundoff_decimals) @@ -306,6 +472,21 @@ class IKResult(Sequence): return sol def get_batch_unique_solution(self, roundoff_decimals: int = 2) -> List[torch.Tensor]: + """Get unique solutions from many feasible solutions for the same problem in batch. + + Use this after solving IK with :meth:`IKSolver.solve_batch` with return_seeds > 1. This + function will return unique solutions from the set of feasible solutions, filering out + solutions that are within roundoff_decimals of each other. Current implementation is not + efficient as it will run a for loop over the batch as each problem in the batch can have + different number of unique solutions. + + Args: + roundoff_decimals: Number of decimal places to round off the solution to measure + uniqueness. + + Returns: + List of unique solutions from the set of feasible solutions. + """ in_solution = self.solution r_sol = torch.round(in_solution, decimals=roundoff_decimals) if not (len(in_solution.shape) == 3): @@ -319,12 +500,29 @@ class IKResult(Sequence): r_k = r_sol[k][self.success[k]] s, i = torch.unique(r_k, dim=-2, return_inverse=True) sol.append(in_succ[i[: s.shape[0]]]) - # sol.append(s) return sol class IKSolver(IKSolverConfig): + """Inverse Kinematics Solver for reaching Cartesian Pose with end-effector. + + This also supports reaching poses for multiple links of a robot (e.g., a bimanual robot). + This solver creates memory buffers on the GPU, captures CUDAGraphs of the operations during the + very first call to any of the solve functions and reuses them for solving subsequent IK + problems. As such, changing the number of problems, number of seeds, number of seeds or type of + IKProblem to solve will cause existing CUDAGraphs to be invalidated, which currently leads to an + exit with error. Either use multiple instances of IKSolver to solve different types of + IKProblems or disable `use_cuda_graph` to avoid this issue. Disabling `use_cuda_graph` can lead + to a 10x slowdown in solving IK problems. + """ + def __init__(self, config: IKSolverConfig) -> None: + """Initializer for IK Solver. + + Args: + config: Configuration for Inverse Kinematics Solver. + """ + super().__init__(**vars(config)) self.batch_size = -1 self._num_seeds = self.num_seeds @@ -343,8 +541,6 @@ class IKSolver(IKSolverConfig): seed=self.seed, ) - # load ik nn: - # store og outer iters: self.og_newton_iters = self.solver.newton_optimizer.outer_iters self._goal_buffer = Goal() @@ -352,13 +548,26 @@ class IKSolver(IKSolverConfig): self._kin_list = None self._rollout_list = None - def update_goal_buffer( + def _update_goal_buffer( self, solve_state: ReacherSolveState, goal_pose: Pose, retract_config: Optional[T_BDOF] = None, link_poses: Optional[Dict[str, Pose]] = None, ) -> Goal: + """Update goal buffer with new goal pose and retract configuration. + + Args: + solve_state: Current IK problem parameters. + goal_pose: New goal pose to solve for IK. + retract_config: Retract configuration to use for IK. If None, the retract configuration + is set to the retract_config in :meth:`curobo.types.robot.RobotConfig`. + link_poses: New poses for other links to use for IK. This is useful when solving for + bimanual robots or when solving for multiple links of the same robot. + + Returns: + Updated goal buffer. This also updates the internal goal_buffer of IKSolver. + """ self._solve_state, self._goal_buffer, update_reference = solve_state.update_goal_buffer( goal_pose, None, @@ -386,6 +595,554 @@ class IKSolver(IKSolverConfig): return self._goal_buffer + def solve_single( + self, + goal_pose: Pose, + retract_config: Optional[T_BDOF] = None, + seed_config: Optional[T_BDOF] = None, + return_seeds: int = 1, + num_seeds: Optional[int] = None, + use_nn_seed: bool = True, + newton_iters: Optional[int] = None, + link_poses: Optional[Dict[str, Pose]] = None, + ) -> IKResult: + """Solve single IK problem. + + To get the closest IK solution from current joint configuration useful for IK servoing, set + retract_config and seed_config to current joint configuration, also make sure IKSolverConfig + was created with regularization enabled. If the solution is still not sufficiently close to + the current joint configuration, increase the weight for `null_space_cfg` in `convergence` + of `base_cfg.yml` which will select a solution that is closer to the current + joint configuration after optimization. You can also increase the weight of + `null_space_weight` in `bound_cfg` of `gradient_ik.yml` to encourage the optimization to + stay near the current joint configuration during iterations. + + Args: + goal_pose: Pose to reach with end-effector. + retract_config: Retract configuration to use as regularization for IK. For this to work, + :meth:`IKSolverConfig.load_from_robot_config` should have regularization enabled. + This should be of shape (1, dof), where dof is the number of degrees of freedom in + the robot. The order of joints should match :meth:`IKSolver.joint_names`. + seed_config: Initial seed configuration to use for optimization. If None, a random seed + is generated. The n seeds passed should be of shape (n, 1, dof), where dof is + the number of degrees of freedom in the robot. The number of seeds do not have to + match the number of seeds in the IKSolver. The remaining seeds are generated + randomly. The order of joints should match :meth:`IKSolver.joint_names`. + return_seeds: Number of solutions to return for the IK problem. + num_seeds: Number of seeds to optimize per IK problem in parallel. Changing number of + seeds is not allowed when use_cuda_graph is enabled. + use_nn_seed: Flag to use neural network as seed for IK. This is not implemented yet. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number + of iterations is set to the default value in the configuration (`gradient_ik.yml`). + link_poses: Poses for other links to use for IK. This is useful when solving for + bimanual robots or when solving for multiple links of the same robot. The link_poses + should be a dictionary with link name as key and pose as value. + + Returns: + :class:`IKResult` object with solution to the IK problem. Use :meth:`IKResult.success` to check + if the problem was solved successfully. + """ + if num_seeds is None: + num_seeds = self.num_seeds + if return_seeds > num_seeds: + num_seeds = return_seeds + + solve_state = ReacherSolveState( + ReacherSolveType.SINGLE, num_ik_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1 + ) + + return self._solve_from_solve_state( + solve_state, + goal_pose, + num_seeds, + retract_config, + seed_config, + return_seeds, + use_nn_seed, + newton_iters, + link_poses=link_poses, + ) + + def solve_goalset( + self, + goal_pose: Pose, + retract_config: Optional[T_BDOF] = None, + seed_config: Optional[T_BDOF] = None, + return_seeds: int = 1, + num_seeds: Optional[int] = None, + use_nn_seed: bool = True, + newton_iters: Optional[int] = None, + link_poses: Optional[Dict[str, Pose]] = None, + ) -> IKResult: + """Solve IK problem to reach one pose in a set of poses. + + To get the closest IK solution from current joint configuration useful for IK servoing, set + retract_config and seed_config to current joint configuration, also make sure IKSolverConfig + was created with regularization enabled. If the solution is still not sufficiently close to + the current joint configuration, increase the weight for `null_space_cfg` in `convergence` + of `base_cfg.yml` which will select a solution that is closer to the current + joint configuration after optimization. You can also increase the weight of + `null_space_weight` in `bound_cfg` of `gradient_ik.yml` to encourage the optimization to + stay near the current joint configuration during iterations. + + Args: + goal_pose: Pose to reach with end-effector. + retract_config: Retract configuration to use as regularization for IK. For this to work, + :meth:`IKSolverConfig.load_from_robot_config` should have regularization enabled. + This should be of shape (1, dof), where dof is the number of degrees of freedom in + the robot. The order of joints should match :meth:`IKSolver.joint_names`. + seed_config: Initial seed configuration to use for optimization. If None, a random seed + is generated. The n seeds passed should be of shape (n, 1, dof), where dof is + the number of degrees of freedom in the robot. The number of seeds do not have to + match the number of seeds in the IKSolver. The remaining seeds are generated + randomly. The order of joints should match :meth:`IKSolver.joint_names`. + return_seeds: Number of solutions to return for the IK problem. + num_seeds: Number of seeds to optimize per IK problem in parallel. Changing number of + seeds is not allowed when use_cuda_graph is enabled. + use_nn_seed: Flag to use neural network as seed for IK. This is not implemented yet. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number + of iterations is set to the default value in the configuration (`gradient_ik.yml`). + link_poses: Poses for other links to use for IK. This is useful when solving for + bimanual robots or when solving for multiple links of the same robot. The link_poses + should be a dictionary with link name as key and pose as value. + + Returns: + :class:`IKResult` object with solution to the IK problem. Check :meth:`IKResult.goalset_index` to + identify which pose was reached by the solution. Use :meth:`IKResult.success` to check + if the problem was solved successfully. + """ + if num_seeds is None: + num_seeds = self.num_seeds + if return_seeds > num_seeds: + num_seeds = return_seeds + + solve_state = ReacherSolveState( + ReacherSolveType.GOALSET, + num_ik_seeds=num_seeds, + batch_size=1, + n_envs=1, + n_goalset=goal_pose.n_goalset, + ) + return self._solve_from_solve_state( + solve_state, + goal_pose, + num_seeds, + retract_config, + seed_config, + return_seeds, + use_nn_seed, + newton_iters, + link_poses=link_poses, + ) + + def solve_batch( + self, + goal_pose: Pose, + retract_config: Optional[T_BDOF] = None, + seed_config: Optional[T_BDOF] = None, + return_seeds: int = 1, + num_seeds: Optional[int] = None, + use_nn_seed: bool = True, + newton_iters: Optional[int] = None, + link_poses: Optional[Dict[str, Pose]] = None, + ) -> IKResult: + """Solve batch of IK problems. + + Changing number of problems (batch size) is not allowed when use_cuda_graph is enabled. The + number of problems is determined during the first call to this function. + + Args: + goal_pose: Pose to reach with end-effector. + retract_config: Retract configuration to use as regularization for IK. For this to work, + :meth:`IKSolverConfig.load_from_robot_config` should have regularization enabled. + This should be of shape (batch, dof), where dof is the number of degrees of freedom in + the robot. The order of joints should match :meth:`IKSolver.joint_names`. + seed_config: Initial seed configuration to use for optimization. If None, a random seed + is generated. The n seeds passed should be of shape (n, batch, dof), where dof is + the number of degrees of freedom in the robot, and batch is number of problems in + batch. The number of seeds do not have to match the number of seeds in the IKSolver. + The remaining seeds are generated randomly. The order of joints should match + :meth:`IKSolver.joint_names`. + return_seeds: Number of solutions to return per problem in batch. + num_seeds: Number of seeds to optimize per IK problem in parallel. Changing number of + seeds is not allowed when use_cuda_graph is enabled. + use_nn_seed: Flag to use neural network as seed for IK. This is not implemented yet. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number + of iterations is set to the default value in the configuration (`gradient_ik.yml`). + link_poses: Poses for other links to use for IK. This is useful when solving for + bimanual robots or when solving for multiple links of the same robot. The link_poses + should be a dictionary with link name as key and pose as value. + + Returns: + :class:`IKResult` object with solution to the batch of IK problems. Use :meth:`IKResult.success` + to filter successful solutions. + """ + if num_seeds is None: + num_seeds = self.num_seeds + if return_seeds > num_seeds: + num_seeds = return_seeds + + solve_state = ReacherSolveState( + ReacherSolveType.BATCH, + num_ik_seeds=num_seeds, + batch_size=goal_pose.batch, + n_envs=1, + n_goalset=1, + ) + return self._solve_from_solve_state( + solve_state, + goal_pose, + num_seeds, + retract_config, + seed_config, + return_seeds, + use_nn_seed, + newton_iters, + link_poses=link_poses, + ) + + def solve_batch_goalset( + self, + goal_pose: Pose, + retract_config: Optional[T_BDOF] = None, + seed_config: Optional[T_BDOF] = None, + return_seeds: int = 1, + num_seeds: Optional[int] = None, + use_nn_seed: bool = True, + newton_iters: Optional[int] = None, + link_poses: Optional[Dict[str, Pose]] = None, + ) -> IKResult: + """Solve batch of IK problems to reach one pose in a set of poses for a batch of problems. + + Args: + goal_pose: Pose to reach with end-effector. + retract_config: Retract configuration to use as regularization for IK. For this to work, + :meth:`IKSolverConfig.load_from_robot_config` should have regularization enabled. + This should be of shape (batch, dof), where dof is the number of degrees of freedom in + the robot. The order of joints should match :meth:`IKSolver.joint_names`. + seed_config: Initial seed configuration to use for optimization. If None, a random seed + is generated. The n seeds passed should be of shape (n, batch, dof), where dof is + the number of degrees of freedom in the robot, and batch is number of problems in + batch. The number of seeds do not have to match the number of seeds in the IKSolver. + The remaining seeds are generated randomly. The order of joints should match + :meth:`IKSolver.joint_names`. + return_seeds: Number of solutions to return per problem in batch. + num_seeds: Number of seeds to optimize per IK problem in parallel. Changing number of + seeds is not allowed when use_cuda_graph is enabled. + use_nn_seed: Flag to use neural network as seed for IK. This is not implemented yet. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in the configuration (`gradient_ik.yml`). + link_poses: Poses for other links to use for IK. This is useful when solving for + bimanual robots or when solving for multiple links of the same robot. The link_poses + should be a dictionary with link name as key and pose as value. + + Returns: + :class:`IKResult` object with solution to the batch of IK problems. Use :meth:`IKResult.success` + to filter successful solutions. Use :meth:`IKResult.goalset_index` to identify which + pose was reached by each solution. + """ + if num_seeds is None: + num_seeds = self.num_seeds + if return_seeds > num_seeds: + num_seeds = return_seeds + + solve_state = ReacherSolveState( + ReacherSolveType.BATCH_GOALSET, + num_ik_seeds=num_seeds, + batch_size=goal_pose.batch, + n_envs=1, + n_goalset=goal_pose.n_goalset, + ) + return self._solve_from_solve_state( + solve_state, + goal_pose, + num_seeds, + retract_config, + seed_config, + return_seeds, + use_nn_seed, + newton_iters, + link_poses=link_poses, + ) + + def solve_batch_env( + self, + goal_pose: Pose, + retract_config: Optional[T_BDOF] = None, + seed_config: Optional[T_BDOF] = None, + return_seeds: int = 1, + num_seeds: Optional[int] = None, + use_nn_seed: bool = True, + newton_iters: Optional[int] = None, + link_poses: Optional[Dict[str, Pose]] = None, + ) -> IKResult: + """Solve batch of IK problems with each problem in different world environments. + + Args: + goal_pose: Pose to reach with end-effector. + retract_config: Retract configuration to use as regularization for IK. For this to work, + :meth:`IKSolverConfig.load_from_robot_config` should have regularization enabled. + This should be of shape (batch, dof), where dof is the number of degrees of freedom + in the robot. The order of joints should match :meth:`IKSolver.joint_names`. + seed_config: Initial seed configuration to use for optimization. If None, a random seed + is generated. The n seeds passed should be of shape (n, batch, dof), where dof is + the number of degrees of freedom in the robot, and batch is number of problems in + batch. The number of seeds do not have to match the number of seeds in the IKSolver. + The remaining seeds are generated randomly. The order of joints should match + :meth:`IKSolver.joint_names`. + return_seeds: Number of solutions to return per problem in batch. + num_seeds: Number of seeds to optimize per IK problem in parallel. Changing number of + seeds is not allowed when use_cuda_graph is enabled. + use_nn_seed: Flag to use neural network as seed for IK. This is not implemented yet. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in the configuration (`gradient_ik.yml`). + link_poses: Poses for other links to use for IK. This is useful when solving for + bimanual robots or when solving for multiple links of the same robot. The link_poses + should be a dictionary with link name as key and pose as value. + + Returns: + :class:`IKResult` object with solution to the batch of IK problems. Use :meth:`IKResult.success` + to filter successful solutions. + """ + if num_seeds is None: + num_seeds = self.num_seeds + if return_seeds > num_seeds: + num_seeds = return_seeds + + solve_state = ReacherSolveState( + ReacherSolveType.BATCH_ENV, + num_ik_seeds=num_seeds, + batch_size=goal_pose.batch, + n_envs=goal_pose.batch, + n_goalset=1, + ) + return self._solve_from_solve_state( + solve_state, + goal_pose, + num_seeds, + retract_config, + seed_config, + return_seeds, + use_nn_seed, + newton_iters, + link_poses=link_poses, + ) + + def solve_batch_env_goalset( + self, + goal_pose: Pose, + retract_config: Optional[T_BDOF] = None, + seed_config: Optional[T_BDOF] = None, + return_seeds: int = 1, + num_seeds: Optional[int] = None, + use_nn_seed: bool = True, + newton_iters: Optional[int] = None, + link_poses: Optional[Dict[str, Pose]] = None, + ) -> IKResult: + """Solve batch of goalset IK problems with each problem in different world environments. + + Args: + goal_pose: Pose to reach with end-effector. + retract_config: Retract configuration to use as regularization for IK. For this to work, + :meth:`IKSolverConfig.load_from_robot_config` should have regularization enabled. + This should be of shape (batch, dof), where dof is the number of degrees of freedom + in the robot. The order of joints should match :meth:`IKSolver.joint_names`. + seed_config: Initial seed configuration to use for optimization. If None, a random seed + is generated. The n seeds passed should be of shape (n, batch, dof), where dof is + the number of degrees of freedom in the robot, and batch is number of problems in + batch. The number of seeds do not have to match the number of seeds in the IKSolver. + The remaining seeds are generated randomly. The order of joints should match + :meth:`IKSolver.joint_names`. + return_seeds: Number of solutions to return per problem in batch. + num_seeds: Number of seeds to optimize per IK problem in parallel. Changing number of + seeds is not allowed when use_cuda_graph is enabled. + use_nn_seed: Flag to use neural network as seed for IK. This is not implemented yet. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in the configuration (`gradient_ik.yml`). + link_poses: Poses for other links to use for IK. This is useful when solving for + bimanual robots or when solving for multiple links of the same robot. The link_poses + should be a dictionary with link name as key and pose as value. + + Returns: + :class:`IKResult` object with solution to the batch of IK problems. Use :meth:`IKResult.success` + to filter successful solutions. Use :meth:`IKResult.goalset_index` to identify which + pose was reached by each solution. + """ + if num_seeds is None: + num_seeds = self.num_seeds + if return_seeds > num_seeds: + num_seeds = return_seeds + + solve_state = ReacherSolveState( + ReacherSolveType.BATCH_ENV_GOALSET, + num_ik_seeds=num_seeds, + batch_size=goal_pose.batch, + n_envs=goal_pose.batch, + n_goalset=goal_pose.n_goalset, + ) + return self._solve_from_solve_state( + solve_state, + goal_pose, + num_seeds, + retract_config, + seed_config, + return_seeds, + use_nn_seed, + newton_iters, + link_poses=link_poses, + ) + + def _solve_from_solve_state( + self, + solve_state: ReacherSolveState, + goal_pose: Pose, + num_seeds: int, + retract_config: Optional[T_BDOF] = None, + seed_config: Optional[T_BDOF] = None, + return_seeds: int = 1, + use_nn_seed: bool = True, + newton_iters: Optional[int] = None, + link_poses: Optional[Dict[str, Pose]] = None, + ) -> IKResult: + """Solve IK problem from ReacherSolveState. Called by all solve functions. + + Args: + solve_state: ReacherSolveState with information about the type of IK problem to solve. + goal_pose: Pose to reach with end-effector. + num_seeds: Number of seeds to optimize per IK problem in parallel. + retract_config: Retract configuration to use as regularization for IK. For this to work, + :meth:`IKSolverConfig.load_from_robot_config` should have regularization enabled. + Shape of retract_config depends on type of IK problem being solved. + seed_config: Initial seed configuration to use for optimization. If None, a random seed + is generated. The n seeds passed should be of shape (n, batch, dof), where dof is + the number of degrees of freedom in the robot, and batch is number of problems in + batch. The number of seeds do not have to match the number of seeds in the IKSolver. + The remaining seeds are generated randomly. The order of joints should match + :meth:`IKSolver.joint_names`. When solving for single IK problem types, batch==1. + return_seeds: Number of solutions to return per problem in batch. + use_nn_seed: Flag to use neural network as seed for IK. This is not implemented yet. + newton_iters: Number of iterations to run LBFGS optimization. If None, the number of + iterations is set to the default value in the configuration (`gradient_ik.yml`). + link_poses: Poses for other links to use for IK. This is useful when solving for + bimanual robots or when solving for multiple links of the same robot. The link_poses + should be a dictionary with link name as key and pose as value. + + Returns: + :class:`IKResult` object with solution to the IK problem. Use :meth:`IKResult.success` + to check if the problem was solved successfully. + """ + # create goal buffer: + goal_buffer = self._update_goal_buffer(solve_state, goal_pose, retract_config, link_poses) + coord_position_seed = self.get_seed( + num_seeds, goal_buffer.goal_pose, use_nn_seed, seed_config + ) + + if newton_iters is not None: + self.solver.newton_optimizer.outer_iters = newton_iters + self.solver.reset() + result = self.solver.solve(goal_buffer, coord_position_seed) + if newton_iters is not None: + self.solver.newton_optimizer.outer_iters = self.og_newton_iters + ik_result = self._get_result(num_seeds, result, goal_buffer.goal_pose, return_seeds) + if ik_result.goalset_index is not None: + ik_result.goalset_index[ik_result.goalset_index >= goal_pose.n_goalset] = 0 + + return ik_result + + @profiler.record_function("ik/get_result") + def _get_result( + self, num_seeds: int, result: WrapResult, goal_pose: Pose, return_seeds: int + ) -> IKResult: + """Get IKResult from WrapResult after optimization. + + Args: + num_seeds: number of seeds used for optimization. + result: result from optimization. + goal_pose: goal poses used for IK problems. + return_seeds: number of seeds to return per problem. + + Returns: + IKResult object with solutions to the IK problems. + """ + success = self._get_success(result.metrics, num_seeds=num_seeds) + + if result.metrics.null_space_error is not None: + result.metrics.pose_error += result.metrics.null_space_error + if result.metrics.cspace_error is not None: + result.metrics.pose_error += result.metrics.cspace_error + + q_sol, success, position_error, rotation_error, total_error, goalset_index = get_result( + result.metrics.pose_error, + result.metrics.position_error, + result.metrics.rotation_error, + result.metrics.goalset_index, + success, + result.action.position, + self._col, + goal_pose.batch, + return_seeds, + num_seeds, + ) + # check if locked joints exist and create js solution: + + new_js = JointState(q_sol, joint_names=self.rollout_fn.kinematics.joint_names) + sol_js = self.rollout_fn.get_full_dof_from_solution(new_js) + # reindex success to get successful poses? + ik_result = IKResult( + success=success, + goal_pose=goal_pose, + solution=q_sol, + seed=None, + js_solution=sol_js, + # seed=coord_position_seed[idx].view(goal_pose.batch, return_seeds, -1).detach(), + position_error=position_error, + rotation_error=rotation_error, + solve_time=result.solve_time, + error=total_error, + debug_info={"solver": result.debug}, + goalset_index=goalset_index, + ) + return ik_result + + @profiler.record_function("ik/get_seed") + def get_seed( + self, num_seeds: int, goal_pose: Pose, use_nn_seed, seed_config: Optional[T_BDOF] = None + ) -> torch.Tensor: + """Get seed joint configurations for optimization. + + Args: + num_seeds: number of seeds to generate. + goal_pose: goal poses for IK problems. This is used to generate seeds with a + neural network. Not implemented yet. + use_nn_seed: flag to use neural network as seed for IK. This is not implemented yet. + seed_config: seed configuration to use for optimization. If None, random seeds are + generated. seed config should be of shape (batch, n, dof), where n can be lower + or equal to num_seeds. The order of joints should match + :meth:`IKSolver.joint_names`. + + Returns: + seed joint configurations for optimization. + """ + if seed_config is None: + coord_position_seed = self.generate_seed( + num_seeds=num_seeds, + batch=goal_pose.batch, + use_nn_seed=use_nn_seed, + pose=goal_pose, + ) + elif seed_config.shape[1] < num_seeds: + coord_position_seed = self.generate_seed( + num_seeds=num_seeds - seed_config.shape[1], + batch=goal_pose.batch, + use_nn_seed=use_nn_seed, + pose=goal_pose, + ) + coord_position_seed = torch.cat((seed_config, coord_position_seed), dim=1) + else: + coord_position_seed = seed_config + coord_position_seed = coord_position_seed.view(-1, 1, self.dof) + return coord_position_seed + def solve_any( self, solve_type: ReacherSolveType, @@ -398,6 +1155,7 @@ class IKSolver(IKSolverConfig): newton_iters: Optional[int] = None, link_poses: Optional[Dict[str, Pose]] = None, ) -> IKResult: + """Solve IK problem with any solve type.""" if solve_type == ReacherSolveType.SINGLE: return self.solve_single( goal_pose, @@ -461,312 +1219,6 @@ class IKSolver(IKSolverConfig): newton_iters, ) - def solve_single( - self, - goal_pose: Pose, - retract_config: Optional[T_BDOF] = None, - seed_config: Optional[T_BDOF] = None, - return_seeds: int = 1, - num_seeds: Optional[int] = None, - use_nn_seed: bool = True, - newton_iters: Optional[int] = None, - link_poses: Optional[Dict[str, Pose]] = None, - ) -> IKResult: - if num_seeds is None: - num_seeds = self.num_seeds - if return_seeds > num_seeds: - num_seeds = return_seeds - - solve_state = ReacherSolveState( - ReacherSolveType.SINGLE, num_ik_seeds=num_seeds, batch_size=1, n_envs=1, n_goalset=1 - ) - - return self.solve_from_solve_state( - solve_state, - goal_pose, - num_seeds, - retract_config, - seed_config, - return_seeds, - use_nn_seed, - newton_iters, - link_poses=link_poses, - ) - - def solve_goalset( - self, - goal_pose: Pose, - retract_config: Optional[T_BDOF] = None, - seed_config: Optional[T_BDOF] = None, - return_seeds: int = 1, - num_seeds: Optional[int] = None, - use_nn_seed: bool = True, - newton_iters: Optional[int] = None, - link_poses: Optional[Dict[str, Pose]] = None, - ) -> IKResult: - if num_seeds is None: - num_seeds = self.num_seeds - if return_seeds > num_seeds: - num_seeds = return_seeds - - solve_state = ReacherSolveState( - ReacherSolveType.GOALSET, - num_ik_seeds=num_seeds, - batch_size=1, - n_envs=1, - n_goalset=goal_pose.n_goalset, - ) - return self.solve_from_solve_state( - solve_state, - goal_pose, - num_seeds, - retract_config, - seed_config, - return_seeds, - use_nn_seed, - newton_iters, - link_poses=link_poses, - ) - - def solve_batch( - self, - goal_pose: Pose, - retract_config: Optional[T_BDOF] = None, - seed_config: Optional[T_BDOF] = None, - return_seeds: int = 1, - num_seeds: Optional[int] = None, - use_nn_seed: bool = True, - newton_iters: Optional[int] = None, - link_poses: Optional[Dict[str, Pose]] = None, - ) -> IKResult: - if num_seeds is None: - num_seeds = self.num_seeds - if return_seeds > num_seeds: - num_seeds = return_seeds - - solve_state = ReacherSolveState( - ReacherSolveType.BATCH, - num_ik_seeds=num_seeds, - batch_size=goal_pose.batch, - n_envs=1, - n_goalset=1, - ) - return self.solve_from_solve_state( - solve_state, - goal_pose, - num_seeds, - retract_config, - seed_config, - return_seeds, - use_nn_seed, - newton_iters, - link_poses=link_poses, - ) - - def solve_batch_goalset( - self, - goal_pose: Pose, - retract_config: Optional[T_BDOF] = None, - seed_config: Optional[T_BDOF] = None, - return_seeds: int = 1, - num_seeds: Optional[int] = None, - use_nn_seed: bool = True, - newton_iters: Optional[int] = None, - link_poses: Optional[Dict[str, Pose]] = None, - ) -> IKResult: - if num_seeds is None: - num_seeds = self.num_seeds - if return_seeds > num_seeds: - num_seeds = return_seeds - - solve_state = ReacherSolveState( - ReacherSolveType.BATCH_GOALSET, - num_ik_seeds=num_seeds, - batch_size=goal_pose.batch, - n_envs=1, - n_goalset=goal_pose.n_goalset, - ) - return self.solve_from_solve_state( - solve_state, - goal_pose, - num_seeds, - retract_config, - seed_config, - return_seeds, - use_nn_seed, - newton_iters, - link_poses=link_poses, - ) - - def solve_batch_env( - self, - goal_pose: Pose, - retract_config: Optional[T_BDOF] = None, - seed_config: Optional[T_BDOF] = None, - return_seeds: int = 1, - num_seeds: Optional[int] = None, - use_nn_seed: bool = True, - newton_iters: Optional[int] = None, - link_poses: Optional[Dict[str, Pose]] = None, - ) -> IKResult: - if num_seeds is None: - num_seeds = self.num_seeds - if return_seeds > num_seeds: - num_seeds = return_seeds - - solve_state = ReacherSolveState( - ReacherSolveType.BATCH_ENV, - num_ik_seeds=num_seeds, - batch_size=goal_pose.batch, - n_envs=goal_pose.batch, - n_goalset=1, - ) - return self.solve_from_solve_state( - solve_state, - goal_pose, - num_seeds, - retract_config, - seed_config, - return_seeds, - use_nn_seed, - newton_iters, - link_poses=link_poses, - ) - - def solve_batch_env_goalset( - self, - goal_pose: Pose, - retract_config: Optional[T_BDOF] = None, - seed_config: Optional[T_BDOF] = None, - return_seeds: int = 1, - num_seeds: Optional[int] = None, - use_nn_seed: bool = True, - newton_iters: Optional[int] = None, - link_poses: Optional[Dict[str, Pose]] = None, - ) -> IKResult: - if num_seeds is None: - num_seeds = self.num_seeds - if return_seeds > num_seeds: - num_seeds = return_seeds - - solve_state = ReacherSolveState( - ReacherSolveType.BATCH_ENV_GOALSET, - num_ik_seeds=num_seeds, - batch_size=goal_pose.batch, - n_envs=goal_pose.batch, - n_goalset=goal_pose.n_goalset, - ) - return self.solve_from_solve_state( - solve_state, - goal_pose, - num_seeds, - retract_config, - seed_config, - return_seeds, - use_nn_seed, - newton_iters, - link_poses=link_poses, - ) - - def solve_from_solve_state( - self, - solve_state: ReacherSolveState, - goal_pose: Pose, - num_seeds: int, - retract_config: Optional[T_BDOF] = None, - seed_config: Optional[T_BDOF] = None, - return_seeds: int = 1, - use_nn_seed: bool = True, - newton_iters: Optional[int] = None, - link_poses: Optional[Dict[str, Pose]] = None, - ) -> IKResult: - # create goal buffer: - goal_buffer = self.update_goal_buffer(solve_state, goal_pose, retract_config, link_poses) - coord_position_seed = self.get_seed( - num_seeds, goal_buffer.goal_pose, use_nn_seed, seed_config - ) - - if newton_iters is not None: - self.solver.newton_optimizer.outer_iters = newton_iters - self.solver.reset() - result = self.solver.solve(goal_buffer, coord_position_seed) - if newton_iters is not None: - self.solver.newton_optimizer.outer_iters = self.og_newton_iters - ik_result = self.get_result(num_seeds, result, goal_buffer.goal_pose, return_seeds) - if ik_result.goalset_index is not None: - ik_result.goalset_index[ik_result.goalset_index >= goal_pose.n_goalset] = 0 - - return ik_result - - @profiler.record_function("ik/get_result") - def get_result( - self, num_seeds: int, result: WrapResult, goal_pose: Pose, return_seeds: int - ) -> IKResult: - success = self.get_success(result.metrics, num_seeds=num_seeds) - # if result.metrics.cost is not None: - # result.metrics.pose_error += result.metrics.cost * 0.0001 - if result.metrics.null_space_error is not None: - result.metrics.pose_error += result.metrics.null_space_error - if result.metrics.cspace_error is not None: - result.metrics.pose_error += result.metrics.cspace_error - - q_sol, success, position_error, rotation_error, total_error, goalset_index = get_result( - result.metrics.pose_error, - result.metrics.position_error, - result.metrics.rotation_error, - result.metrics.goalset_index, - success, - result.action.position, - self._col, - goal_pose.batch, - return_seeds, - num_seeds, - ) - # check if locked joints exist and create js solution: - - new_js = JointState(q_sol, joint_names=self.rollout_fn.kinematics.joint_names) - sol_js = self.rollout_fn.get_full_dof_from_solution(new_js) - # reindex success to get successful poses? - ik_result = IKResult( - success=success, - goal_pose=goal_pose, - solution=q_sol, - seed=None, - js_solution=sol_js, - # seed=coord_position_seed[idx].view(goal_pose.batch, return_seeds, -1).detach(), - position_error=position_error, - rotation_error=rotation_error, - solve_time=result.solve_time, - error=total_error, - debug_info={"solver": result.debug}, - goalset_index=goalset_index, - ) - return ik_result - - @profiler.record_function("ik/get_seed") - def get_seed( - self, num_seeds: int, goal_pose: Pose, use_nn_seed, seed_config: Optional[T_BDOF] = None - ) -> torch.Tensor: - if seed_config is None: - coord_position_seed = self.generate_seed( - num_seeds=num_seeds, - batch=goal_pose.batch, - use_nn_seed=use_nn_seed, - pose=goal_pose, - ) - elif seed_config.shape[1] < num_seeds: - coord_position_seed = self.generate_seed( - num_seeds=num_seeds - seed_config.shape[1], - batch=goal_pose.batch, - use_nn_seed=use_nn_seed, - pose=goal_pose, - ) - coord_position_seed = torch.cat((seed_config, coord_position_seed), dim=1) - else: - coord_position_seed = seed_config - coord_position_seed = coord_position_seed.view(-1, 1, self.dof) - return coord_position_seed - def solve( self, goal_pose: Pose, @@ -777,20 +1229,7 @@ class IKSolver(IKSolverConfig): use_nn_seed: bool = True, newton_iters: Optional[int] = None, ) -> IKResult: # pragma : no cover - """Ik solver - - Args: - goal_pose (Pose): _description_ - retract_config (Optional[T_BDOF], optional): _description_. Defaults to None. - seed_config (Optional[T_BDOF], optional): _description_. Defaults to None. - return_seeds (int, optional): _description_. Defaults to 1. - num_seeds (Optional[int], optional): _description_. Defaults to None. - use_nn_seed (bool, optional): _description_. Defaults to True. - newton_iters (Optional[int], optional): _description_. Defaults to None. - - Returns: - IKResult: _description_ - """ + """Deprecated API for solving single or batch problems.""" log_warn("IKSolver.solve() is deprecated, use solve_single() or others instead") if goal_pose.batch == 1 and goal_pose.n_goalset == 1: return self.solve_single( @@ -843,20 +1282,7 @@ class IKSolver(IKSolverConfig): use_nn_seed: bool = True, newton_iters: Optional[int] = None, ) -> IKResult: # pragma : no cover - """Ik solver - - Args: - goal_pose (Pose): _description_ - retract_config (Optional[T_BDOF], optional): _description_. Defaults to None. - seed_config (Optional[T_BDOF], optional): _description_. Defaults to None. - return_seeds (int, optional): _description_. Defaults to 1. - num_seeds (Optional[int], optional): _description_. Defaults to None. - use_nn_seed (bool, optional): _description_. Defaults to True. - newton_iters (Optional[int], optional): _description_. Defaults to None. - - Returns: - IKResult: _description_ - """ + """Deprecated API, use solve_batch_env() or solve_batch_env_goalset() instead.""" log_warn( "IKSolver.batch_env_solve() is deprecated, use solve_batch_env() or others instead" ) @@ -883,7 +1309,16 @@ class IKSolver(IKSolverConfig): @torch.no_grad() @profiler.record_function("ik/get_success") - def get_success(self, metrics: RolloutMetrics, num_seeds: int) -> torch.Tensor: + def _get_success(self, metrics: RolloutMetrics, num_seeds: int) -> torch.Tensor: + """Get success of IK optimization. + + Args: + metrics: RolloutMetrics with feasibility, pose error, and other costs. + num_seeds: Number of seeds used for IK optimization. + + Returns: + Success of IK optimization as a boolean tensor of shape (batch, num_seeds). + """ success = get_success( metrics.feasible, metrics.position_error, @@ -904,15 +1339,16 @@ class IKSolver(IKSolverConfig): use_nn_seed: bool = False, pose: Optional[Pose] = None, ) -> torch.Tensor: - """Generate seeds for a batch. Given a pose set, this will create all - the seeds: [batch + batch*random_restarts] + """Generate seeds for IK optimization. Args: - batch (int, optional): [description]. Defaults to 1. - num_seeds (Optional[int], optional): [description]. Defaults to None. + num_seeds: Number of seeds to generate using pseudo-random generator. + batch: Number of problems in batch. + use_nn_seed: Flag to use neural network as seed for IK. This is not implemented yet. + pose: Pose to use for generating seeds. This is not implemented yet. Returns: - [type]: [description] + Seed configurations for IK optimization. """ num_random_seeds = num_seeds seed_list = [] @@ -923,34 +1359,72 @@ class IKSolver(IKSolverConfig): seed_list.append(nn_seed) if num_random_seeds > 0: random_seed = self.q_sample_gen.get_samples( - num_random_seeds * batch, bounded=True + num_random_seeds * batch, + bounded=True, ).view(batch, num_random_seeds, self.dof) - seed_list.append(random_seed) coord_position_seed = torch.cat(seed_list, dim=1) return coord_position_seed - def update_world(self, world: WorldConfig) -> bool: + def update_world(self, world: WorldConfig) -> None: + """Update world in IKSolver. + + If the new world configuration has more obstacles than initial cache, the collision cache + will be recreated, breaking existing cuda graphs. This will lead to an exit with error if + use_cuda_graph is enabled. + + Args: + world: World configuration to update in IKSolver. + """ self.world_coll_checker.load_collision_model(world) - return True def reset_seed(self) -> None: + """Reset seed generator in IKSolver.""" self.q_sample_gen.reset() def check_constraints(self, q: JointState) -> RolloutMetrics: + """Check constraints for joint state. + + Args: + q: Joint state to check constraints for. + + Returns: + RolloutMetrics with feasibility of joint state. + """ metrics = self.rollout_fn.rollout_constraint(q.position.unsqueeze(1)) return metrics def sample_configs( - self, n: int, use_batch_env=False, sample_from_ik_seeder: bool = False + self, + n: int, + use_batch_env=False, + sample_from_ik_seeder: bool = False, + rejection_ratio: Optional[int] = None, ) -> torch.Tensor: + """Sample n feasible joint configurations. Only samples with environment=0. + + Args: + n: Number of joint configurations to sample. + use_batch_env: Flag to sample from batch environments. This is not implemented yet. + sample_from_ik_seeder: Flag to sample from IK seeder. This is not implemented yet. + rejection_ratio: Ratio of samples to generate to get n feasible samples. If None, the + rejection ratio is set to the default value meth:`IKSolver.sample_rejection_ratio`. + + Returns: + Joint configurations sampled from the IK seeder of shape (n, dof). """ - Only works for environment=0 - """ + + if use_batch_env: + log_warn( + "IKSolver.sample_configs() does not work with batch environments," + + " sampling only from env=0" + ) + if rejection_ratio is None: + rejection_ratio = self.sample_rejection_ratio if sample_from_ik_seeder: - samples = self.q_sample_gen.get_samples(n * self.sample_rejection_ratio, bounded=True) + samples = self.q_sample_gen.get_samples(n * rejection_ratio, bounded=True) else: - samples = self.rollout_fn.sample_random_actions(n * self.sample_rejection_ratio) + samples = self.rollout_fn.sample_random_actions(n * rejection_ratio) metrics = self.rollout_fn.rollout_constraint( samples.unsqueeze(1), use_batch_env=use_batch_env ) @@ -958,28 +1432,43 @@ class IKSolver(IKSolverConfig): @property def kinematics(self) -> CudaRobotModel: + """Get kinematics instance in IKSolver.""" return self.rollout_fn.dynamics_model.robot_model - def get_all_rollout_instances(self) -> List[RolloutBase]: + def get_all_rollout_instances(self) -> List[ArmReacher]: + """Get all rollout instances in IKSolver. + + Returns: + List of all rollout instances in IKSolver. + """ if self._rollout_list is None: self._rollout_list = [self.rollout_fn] + self.solver.get_all_rollout_instances() return self._rollout_list def get_all_kinematics_instances(self) -> List[CudaRobotModel]: - if self._kin_list is None: - self._kin_list = [ - i.dynamics_model.robot_model for i in self.get_all_rollout_instances() - ] - return self._kin_list + """Deprecated API, use kinematics instead.""" + log_warn("IKSolver.get_all_kinematics_instances() is deprecated, use kinematics instead") + + return [self.kinematics for _ in range(len(self.get_all_rollout_instances))] def fk(self, q: torch.Tensor) -> CudaRobotModelState: + """Forward kinematics for the robot. + + Args: + q: Joint configuration of the robot, with joint values in order of :meth:`joint_names`. + + Returns: + :class:`CudaRobotModelState` with link poses, and link spheres for the robot. + """ return self.kinematics.get_state(q) def reset_cuda_graph(self) -> None: + """Reset the cuda graph for all rollout instances in IKSolver. Does not work currently.""" self.solver.reset_cuda_graph() self.rollout_fn.reset_cuda_graph() def reset_shape(self): + """Reset the shape of the rollout function and solver to the original shape.""" self.solver.reset_shape() self.rollout_fn.reset_shape() @@ -989,22 +1478,45 @@ class IKSolver(IKSolverConfig): sphere_tensor: Optional[torch.Tensor] = None, link_name: str = "attached_object", ) -> None: - for k in self.get_all_kinematics_instances(): - k.attach_object( - sphere_radius=sphere_radius, sphere_tensor=sphere_tensor, link_name=link_name - ) + """Attach object to robot for collision checking. + + Args: + sphere_radius: Radius of the sphere to attach to robot. + sphere_tensor: Tensor of shape (n, 4) to attach to robot, where n is the number of + spheres for that link. If None, only radius is updated for the existing spheres. + link_name: Name of the link to attach object spheres to. + """ + + # get single kinematics instance: + self.kinematics.kinematics_config.attach_object( + sphere_radius=sphere_radius, sphere_tensor=sphere_tensor, link_name=link_name + ) def detach_object_from_robot(self, link_name: str = "attached_object") -> None: - for k in self.get_all_kinematics_instances(): - k.detach_object(link_name) + """Detach all attached objects from robot. - def get_retract_config(self): + This currently will reset the spheres for link_name to -10, disabling collision checking + for that link. + + Args: + link_name: Name of the link to detach object from. + """ + + self.kinematics.kinematics_config.detach_object(link_name=link_name) + + def get_retract_config(self) -> T_DOF: + """Get retract configuration from robot model.""" return self.rollout_fn.dynamics_model.retract_config def update_pose_cost_metric( self, metric: PoseCostMetric, ): + """Update pose cost metric for all rollout instances in IKSolver. + + Args: + metric: New values for Pose cost metric to update. + """ rollouts = self.get_all_rollout_instances() [ rollout.update_pose_cost_metric(metric) @@ -1012,6 +1524,11 @@ class IKSolver(IKSolverConfig): if isinstance(rollout, ArmReacher) ] + @property + def joint_names(self) -> List[str]: + """Get ordered names of all joints used in optimization with IKSolver.""" + return self.rollout_fn.kinematics.joint_names + @get_torch_jit_decorator() def get_success( @@ -1022,6 +1539,7 @@ def get_success( position_threshold: float, rotation_threshold: float, ): + """JIT compatible function to get the success of IK solutions.""" feasible = feasible.view(-1, num_seeds) converge = torch.logical_and( position_error <= position_threshold, @@ -1044,6 +1562,7 @@ def get_result( return_seeds: int, num_seeds: int, ): + """JIT compatible function to get the best IK solutions.""" error = pose_error.view(-1, num_seeds) error[~success] += 1000.0 _, idx = torch.topk(error, k=return_seeds, largest=False, dim=-1) diff --git a/src/curobo/wrap/reacher/motion_gen.py b/src/curobo/wrap/reacher/motion_gen.py index e3d33ec..1468a5c 100644 --- a/src/curobo/wrap/reacher/motion_gen.py +++ b/src/curobo/wrap/reacher/motion_gen.py @@ -9,6 +9,28 @@ # its affiliates is strictly prohibited. # +""" +This module contains :meth:`MotionGen` class that provides a high-level interface for motion +generation. Motion Generation can take goals either as joint configurations +:meth:`MotionGen.plan_single_js` or as Cartesian poses :meth:`MotionGen.plan_single`. When Cartesian +pose is given as target, inverse kinematics is first done to generate seeds for trajectory +optimization. Motion generation fallback to using a graph planner when linear interpolated +trajectory optimization seeds are not successful. Reaching one Cartesian pose in a goalset is also +supported using :meth:`MotionGen.plan_goalset`. Batched planning in the same world environment ( +:meth:`MotionGen.plan_batch`) and different world environments (:meth:`MotionGen.plan_batch_env`) +is also provided. + +A motion generation request can be configured using :meth:`MotionGenPlanConfig`. The result +of motion generation is returned as a :meth:`MotionGenResult`. + + +.. raw:: html + +

+ +

+ +""" from __future__ import annotations @@ -16,6 +38,7 @@ from __future__ import annotations import math import time from dataclasses import dataclass +from enum import Enum from typing import Any, Dict, List, Optional, Union # Third Party @@ -129,16 +152,23 @@ class MotionGenConfig: #: record compute ops as cuda graphs and replay recorded graphs for upto 10x faster execution. use_cuda_graph: bool = True + #: After 100 iterations of trajectory optimization, a new dt is computed that pushes the + #: velocity, acceleration, or jerk limits to the maximum. This new dt is then used with a + #: reduction :attr:`MotionGenPlanConfig.finetune_dt_scale` to run 200+ iterations of trajectory + #: optimization. If trajectory optimization fails with the new dt, the new dt is increased and + #: tried again until :attr:`MotionGenPlanConfig.finetune_attempts`. + optimize_dt: bool = True + @staticmethod @profiler.record_function("motion_gen_config/load_from_robot_config") def load_from_robot_config( robot_cfg: Union[Union[str, Dict], RobotConfig], world_model: Optional[Union[Union[str, Dict], WorldConfig]] = None, tensor_args: TensorDeviceType = TensorDeviceType(), - num_ik_seeds: int = 30, + num_ik_seeds: int = 32, num_graph_seeds: int = 1, - num_trajopt_seeds: int = 12, - num_batch_ik_seeds: int = 30, + num_trajopt_seeds: int = 4, + num_batch_ik_seeds: int = 32, num_batch_trajopt_seeds: int = 1, num_trajopt_noisy_seeds: int = 1, position_threshold: float = 0.005, @@ -197,7 +227,7 @@ class MotionGenConfig: smooth_weight: List[float] = None, finetune_smooth_weight: Optional[List[float]] = None, state_finite_difference_mode: Optional[str] = None, - finetune_dt_scale: float = 0.95, + finetune_dt_scale: float = 0.9, maximum_trajectory_time: Optional[float] = None, maximum_trajectory_dt: float = 0.1, velocity_scale: Optional[Union[List[float], float]] = None, @@ -207,87 +237,18 @@ class MotionGenConfig: project_pose_to_goal_frame: bool = True, ik_seed: int = 1531, graph_seed: int = 1531, + high_precision: bool = False, ): - """Helper function to create configuration from robot and world configuration. - - Args: - robot_cfg: Robot configuration to use for motion generation. - world_model: World model to use for motion generation. Use None to disable world model. - tensor_args: Tensor device to use for motion generation. Use to change cuda device id. - num_ik_seeds: Number of seeds to use in inverse kinematics (IK) optimization. - num_graph_seeds: Number of graph paths to use as seed in trajectory optimization. - num_trajopt_seeds: Number of seeds to use in trajectory optimization. - num_batch_ik_seeds: Number of seeds to use in batch planning modes for IK. - num_batch_trajopt_seeds: Number of seeds to use in batch planning modes for trajopt. - num_trajopt_noisy_seeds: Number of seeds to use for trajopt. - position_threshold: _description_ - rotation_threshold: _description_ - cspace_threshold: _description_ - world_coll_checker: _description_ - base_cfg_file: _description_ - particle_ik_file: _description_ - gradient_ik_file: _description_ - graph_file: _description_ - particle_trajopt_file: _description_ - gradient_trajopt_file: _description_ - finetune_trajopt_file: _description_ - trajopt_tsteps: _description_ - interpolation_steps: _description_ - interpolation_dt: _description_ - interpolation_type: _description_ - use_cuda_graph: _description_ - self_collision_check: _description_ - self_collision_opt: _description_ - grad_trajopt_iters: _description_ - trajopt_seed_ratio: _description_ - ik_opt_iters: _description_ - ik_particle_opt: _description_ - collision_checker_type: _description_ - sync_cuda_time: _description_ - trajopt_particle_opt: _description_ - traj_evaluator_config: _description_ - traj_evaluator: _description_ - minimize_jerk: _description_ - filter_robot_command: _description_ - use_gradient_descent: _description_ - collision_cache: _description_ - n_collision_envs: _description_ - ee_link_name: _description_ - use_es_ik: _description_ - use_es_trajopt: _description_ - es_ik_learning_rate: _description_ - es_trajopt_learning_rate: _description_ - use_ik_fixed_samples: _description_ - use_trajopt_fixed_samples: _description_ - evaluate_interpolated_trajectory: _description_ - partial_ik_iters: _description_ - fixed_iters_trajopt: _description_ - store_ik_debug: _description_ - store_trajopt_debug: _description_ - graph_trajopt_iters: _description_ - collision_max_outside_distance: _description_ - collision_activation_distance: _description_ - trajopt_dt: _description_ - js_trajopt_dt: _description_ - js_trajopt_tsteps: _description_ - trim_steps: _description_ - store_debug_in_result: _description_ - finetune_trajopt_iters: _description_ - smooth_weight: _description_ - finetune_smooth_weight: _description_ - state_finite_difference_mode: _description_ - finetune_dt_scale: _description_ - maximum_trajectory_time: _description_ - maximum_trajectory_dt: _description_ - velocity_scale: _description_ - acceleration_scale: _description_ - jerk_scale: _description_ - optimize_dt: _description_ - - Returns: - _description_ - """ - + if position_threshold <= 0.001: + high_precision = True + if high_precision: + finetune_trajopt_iters = ( + 300 if finetune_trajopt_iters is None else max(300, finetune_trajopt_iters) + ) + grad_trajopt_iters = 200 if grad_trajopt_iters is None else max(200, grad_trajopt_iters) + position_threshold = min(position_threshold, 0.001) + rotation_threshold = min(rotation_threshold, 0.025) + cspace_threshold = min(cspace_threshold, 0.01) init_warp(tensor_args=tensor_args) if js_trajopt_tsteps is not None: log_warn("js_trajopt_tsteps is deprecated, use trajopt_tsteps instead.") @@ -317,14 +278,6 @@ class MotionGenConfig: ): maximum_trajectory_dt = (1.0 / min(velocity_scale)) * maximum_trajectory_dt - if traj_evaluator_config is None: - if maximum_trajectory_dt is not None: - max_dt = maximum_trajectory_dt - if maximum_trajectory_time is not None: - max_dt = maximum_trajectory_time / trajopt_tsteps - if acceleration_scale is not None: - max_dt = max_dt * (1.0 / np.sqrt(min(acceleration_scale))) - traj_evaluator_config = TrajEvaluatorConfig(min_dt=interpolation_dt, max_dt=max_dt) if maximum_trajectory_dt is not None: if trajopt_dt is None: trajopt_dt = maximum_trajectory_dt @@ -385,10 +338,20 @@ class MotionGenConfig: robot_cfg["kinematics"]["cspace"]["velocity_scale"] = velocity_scale robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args) - traj_evaluator_config.max_acc = ( - robot_cfg.kinematics.get_joint_limits().acceleration[1, 0].item() - ) - traj_evaluator_config.max_jerk = robot_cfg.kinematics.get_joint_limits().jerk[1, 0].item() + + if traj_evaluator_config is None: + if maximum_trajectory_dt is not None: + max_dt = maximum_trajectory_dt + if maximum_trajectory_time is not None: + max_dt = maximum_trajectory_time / trajopt_tsteps + if acceleration_scale is not None: + max_dt = max_dt * (1.0 / np.sqrt(min(acceleration_scale))) + traj_evaluator_config = TrajEvaluatorConfig.from_basic( + min_dt=interpolation_dt, max_dt=max_dt, dof=robot_cfg.kinematics.dof + ) + traj_evaluator_config.max_acc = robot_cfg.kinematics.get_joint_limits().acceleration[1] + + traj_evaluator_config.max_jerk = robot_cfg.kinematics.get_joint_limits().jerk[1] if isinstance(world_model, str): world_model = load_yaml(join_path(get_world_configs_path(), world_model)) @@ -439,6 +402,7 @@ class MotionGenConfig: collision_activation_distance=collision_activation_distance, project_pose_to_goal_frame=project_pose_to_goal_frame, seed=ik_seed, + high_precision=high_precision, ) ik_solver = IKSolver(ik_solver_cfg) @@ -612,15 +576,175 @@ class MotionGenConfig: interpolation_dt=interpolation_dt, finetune_dt_scale=finetune_dt_scale, use_cuda_graph=use_cuda_graph, + optimize_dt=optimize_dt, ) +@dataclass +class MotionGenPlanConfig: + """Configuration for querying motion generation.""" + + #: Use graph planner to generate collision-free seed for trajectory optimization. + enable_graph: bool = False + + #: Enable trajectory optimization. + enable_opt: bool = True + + #: Use neural network IK seed for solving inverse kinematics. Not implemented. + use_nn_ik_seed: bool = False + + #: Run trajectory optimization only if graph planner is successful. + need_graph_success: bool = False + + #: Maximum number of attempts allowed to solve the motion generation problem. + max_attempts: int = 60 + + #: Maximum time in seconds allowed to solve the motion generation problem. + timeout: float = 10.0 + + #: Number of failed attempts at which to fallback to a graph planner for obtaining trajectory + #: seeds. + enable_graph_attempt: Optional[int] = 3 + + #: Number of failed attempts at which to disable graph planner. This has not been found to be + #: useful. + disable_graph_attempt: Optional[int] = None + + #: Number of IK attempts allowed before returning a failure. Set this to a low value (5) to + #: save compute time when an unreachable goal is given. + ik_fail_return: Optional[int] = None + + #: Full IK solving takes 10% of the planning time. Setting this to True will run only 50 + #: iterations of IK instead of 100 and then run trajecrtory optimization without checking if + #: IK is successful. This can reduce the planning time by 5% but generated solutions can + #: have larger motion time and path length. Leave this to False for most cases. + partial_ik_opt: bool = False + + #: Number of seeds to use for solving inverse kinematics. Chanigng this will cause the existing + #: CUDAGraphs to be invalidated. Instead, set num_ik_seeds when creating + #: :meth:`MotionGenConfig`. + num_ik_seeds: Optional[int] = None + + #: Number of seeds to use for graph planner. We found 4 to be a good number for most cases. The + #: default value is set when creating :meth:`MotionGenConfig` so leave this to None. + num_graph_seeds: Optional[int] = None + + #: Number of seeds to use for trajectory optimization. We found 12 to be a good number for most + #: cases. The default value is set when creating :meth:`MotionGenConfig` so leave this to None. + num_trajopt_seeds: Optional[int] = None + + #: Ratio of successful motion generation queries to total queries in batch planning mode. Motion + #: generation is queries upto :attr:`MotionGenPlanConfig.max_attempts` until this ratio is met. + success_ratio: float = 1 + + #: Return a failure if the query is invalid. Set this to False to debug a query that is invalid. + fail_on_invalid_query: bool = True + + #: use start config as regularization for IK instead of + #: :meth:`curobo.types.robot.RobotConfig.kinematics.kinematics_config.retract_config` + use_start_state_as_retract: bool = True + + #: Use a custom pose cost metric for trajectory optimization. This is useful for adding + #: additional constraints to motion generation, such as constraining the end-effector's motion + #: to a plane or a line or hold orientation while moving. This is also useful for only reaching + #: partial pose (e.g., only position). See :meth:`curobo.rollout.cost.pose_cost.PoseCostMetric` + #: for more details. + pose_cost_metric: Optional[PoseCostMetric] = None + + #: Run finetuning trajectory optimization after running 100 iterations of + #: trajectory optimization. This will provide shorter and smoother trajectories. When + # :attr:`MotionGenConfig.optimize_dt` is True, this flag will also scale the trajectory + #: optimization by a new dt. Leave this to True for most cases. If you are not interested in + #: finding time-optimal solutions and only want to use motion generation as a feasibility check, + #: set this to False. Note that when set to False, the resulting trajectory is only guaranteed + #: to be collision-free and within joint limits. When False, it's not guaranteed to be smooth + #: and might not execute on a real robot. + enable_finetune_trajopt: bool = True + + #: Run finetuning trajectory optimization across all trajectory optimization seeds. If this is + #: set to False, then only 1 successful seed per query is selected and finetuned. When False, + #: we have observed that the resulting trajectory is not as optimal as when set to True. + parallel_finetune: bool = True + + #: Scale dt by this value before running finetuning trajectory optimization. This enables + #: trajectory optimization to find shorter paths and also account for smoothness over variable + #: length trajectories. This is only used when :attr:`MotionGenConfig.optimize_dt` is True. + finetune_dt_scale: Optional[float] = 0.85 + + #: Number of attempts to run finetuning trajectory optimization. Every attempt will increase the + #: :attr:`MotionGenPlanConfig.finetune_dt_scale` by + #: :attr:`MotionGenPlanConfig.finetune_dt_decay` as a path couldn't be found with the previous + #: smaller dt. + finetune_attempts: int = 5 + + #: Decay factor used to increase :attr:`MotionGenPlanConfig.finetune_dt_scale` when optimization + #: fails to find a solution. This is only used when :attr:`MotionGenConfig.optimize_dt` is True. + finetune_dt_decay: float = 1.025 + + def __post_init__(self): + if not self.enable_opt and not self.enable_graph: + log_error("Graph search and Optimization are both disabled, enable one") + + def clone(self) -> MotionGenPlanConfig: + return MotionGenPlanConfig( + enable_graph=self.enable_graph, + enable_opt=self.enable_opt, + use_nn_ik_seed=self.use_nn_ik_seed, + need_graph_success=self.need_graph_success, + max_attempts=self.max_attempts, + timeout=self.timeout, + enable_graph_attempt=self.enable_graph_attempt, + disable_graph_attempt=self.disable_graph_attempt, + ik_fail_return=self.ik_fail_return, + partial_ik_opt=self.partial_ik_opt, + num_ik_seeds=self.num_ik_seeds, + num_graph_seeds=self.num_graph_seeds, + num_trajopt_seeds=self.num_trajopt_seeds, + success_ratio=self.success_ratio, + fail_on_invalid_query=self.fail_on_invalid_query, + enable_finetune_trajopt=self.enable_finetune_trajopt, + parallel_finetune=self.parallel_finetune, + use_start_state_as_retract=self.use_start_state_as_retract, + pose_cost_metric=( + None if self.pose_cost_metric is None else self.pose_cost_metric.clone() + ), + finetune_dt_scale=self.finetune_dt_scale, + finetune_attempts=self.finetune_attempts, + ) + + +class MotionGenStatus(Enum): + """Status of motion generation query.""" + + #: Inverse kinematics failed to find a solution. + IK_FAIL = "IK Fail" + + #: Graph planner failed to find a solution. + GRAPH_FAIL = "Graph Fail" + + #: Trajectory optimization failed to find a solution. + TRAJOPT_FAIL = "TrajOpt Fail" + + #: Finetune Trajectory optimization failed to find a solution. + FINETUNE_TRAJOPT_FAIL = "Finetune TrajOpt Fail" + + #: Invalid query was given. The start state is either out of joint limits, in collision with + #: world, or in self-collision. + INVALID_QUERY = "Invalid Query" + + #: Motion generation query was successful. + SUCCESS = "Success" + + #: Motion generation was not attempted. + NOT_ATTEMPTED = "Not Attempted" + + @dataclass class MotionGenResult: """Result obtained from motion generation.""" #: success tensor with index referring to the batch index. - success: Optional[T_BValue_bool] = None + success: Optional[torch.Tensor] = None #: returns true if the start state is not satisfying constraints (e.g., in collision) valid_query: bool = True @@ -672,7 +796,7 @@ class MotionGenResult: debug_info: Any = None #: status of motion generation query. returns [IK Fail, Graph Fail, TrajOpt Fail]. - status: Optional[str] = None + status: Optional[Union[MotionGenStatus, str]] = None #: number of attempts used before returning a solution. attempts: int = 1 @@ -818,63 +942,6 @@ class MotionGenResult: return self.optimized_dt * (self.optimized_plan.position.shape[-2] - 1 - 2 - 1) -@dataclass -class MotionGenPlanConfig: - enable_graph: bool = False - enable_opt: bool = True - use_nn_ik_seed: bool = False - need_graph_success: bool = False - max_attempts: int = 60 - timeout: float = 10.0 - enable_graph_attempt: Optional[int] = 3 - disable_graph_attempt: Optional[int] = None - ik_fail_return: Optional[int] = None - partial_ik_opt: bool = False - num_ik_seeds: Optional[int] = None - num_graph_seeds: Optional[int] = None - num_trajopt_seeds: Optional[int] = None - success_ratio: float = 1 - fail_on_invalid_query: bool = True - - #: enables retiming trajectory optimization, useful for getting low jerk trajectories. - enable_finetune_trajopt: bool = True - parallel_finetune: bool = True - - #: use start config as regularization - use_start_state_as_retract: bool = True - - pose_cost_metric: Optional[PoseCostMetric] = None - - def __post_init__(self): - if not self.enable_opt and not self.enable_graph: - log_error("Graph search and Optimization are both disabled, enable one") - - def clone(self) -> MotionGenPlanConfig: - return MotionGenPlanConfig( - enable_graph=self.enable_graph, - enable_opt=self.enable_opt, - use_nn_ik_seed=self.use_nn_ik_seed, - need_graph_success=self.need_graph_success, - max_attempts=self.max_attempts, - timeout=self.timeout, - enable_graph_attempt=self.enable_graph_attempt, - disable_graph_attempt=self.disable_graph_attempt, - ik_fail_return=self.ik_fail_return, - partial_ik_opt=self.partial_ik_opt, - num_ik_seeds=self.num_ik_seeds, - num_graph_seeds=self.num_graph_seeds, - num_trajopt_seeds=self.num_trajopt_seeds, - success_ratio=self.success_ratio, - fail_on_invalid_query=self.fail_on_invalid_query, - enable_finetune_trajopt=self.enable_finetune_trajopt, - parallel_finetune=self.parallel_finetune, - use_start_state_as_retract=self.use_start_state_as_retract, - pose_cost_metric=( - None if self.pose_cost_metric is None else self.pose_cost_metric.clone() - ), - ) - - class MotionGen(MotionGenConfig): def __init__(self, config: MotionGenConfig): super().__init__(**vars(config)) @@ -1092,7 +1159,6 @@ class MotionGen(MotionGenConfig): link_poses: List[Pose] = None, ): start_time = time.time() - if plan_config.pose_cost_metric is not None: valid_query = self.update_pose_cost_metric( plan_config.pose_cost_metric, start_state, goal_pose @@ -1120,7 +1186,8 @@ class MotionGen(MotionGenConfig): "trajopt_attempts": 0, } best_status = 0 # 0== None, 1==IK Fail, 2== Graph Fail, 3==Opt Fail - + if plan_config.finetune_dt_scale is None: + plan_config.finetune_dt_scale = self.finetune_dt_scale for n in range(plan_config.max_attempts): log_info("MG Iter: " + str(n)) result = self._plan_from_solve_state( @@ -1145,16 +1212,22 @@ class MotionGen(MotionGenConfig): break if result.success[0].item(): break + + if result.status == "Finetune Opt Fail": + plan_config.finetune_dt_scale *= ( + plan_config.finetune_dt_decay**plan_config.finetune_attempts + ) + plan_config.finetune_dt_scale = min(plan_config.finetune_dt_scale, 1.25) if plan_config.enable_graph_attempt is not None and ( n >= plan_config.enable_graph_attempt - 1 - and result.status == "Opt Fail" + and result.status in ["Opt Fail", "Finetune Opt Fail"] and not plan_config.enable_graph ): plan_config.enable_graph = True plan_config.partial_ik_opt = False if plan_config.disable_graph_attempt is not None and ( n >= plan_config.disable_graph_attempt - 1 - and result.status in ["Opt Fail", "Graph Fail"] + and result.status in ["Opt Fail", "Graph Fail", "Finetune Opt Fail"] and not force_graph ): plan_config.enable_graph = False @@ -1192,7 +1265,7 @@ class MotionGen(MotionGenConfig): plan_config: MotionGenPlanConfig = MotionGenPlanConfig(), ): start_time = time.time() - + goal_pose = goal_pose.clone() if plan_config.pose_cost_metric is not None: valid_query = self.update_pose_cost_metric( plan_config.pose_cost_metric, start_state, goal_pose @@ -1646,19 +1719,35 @@ class MotionGen(MotionGenConfig): if plan_config.parallel_finetune: opt_dt = torch.min(opt_dt[traj_result.success]) seed_override = solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds - scaled_dt = torch.clamp( - opt_dt * self.finetune_dt_scale, - self.trajopt_solver.interpolation_dt, - ) - self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item()) - traj_result = self._solve_trajopt_from_solve_state( - goal, - solve_state, - seed_traj, - trajopt_instance=self.finetune_trajopt_solver, - num_seeds_override=seed_override, - ) + finetune_time = 0 + for k in range(plan_config.finetune_attempts): + newton_iters = None + + scaled_dt = torch.clamp( + opt_dt + * plan_config.finetune_dt_scale + * (plan_config.finetune_dt_decay ** (k)), + self.trajopt_solver.interpolation_dt, + ) + if self.optimize_dt: + self.finetune_trajopt_solver.update_solver_dt(scaled_dt.item()) + + traj_result = self._solve_trajopt_from_solve_state( + goal, + solve_state, + seed_traj, + trajopt_instance=self.finetune_trajopt_solver, + num_seeds_override=seed_override, + newton_iters=newton_iters, + ) + finetune_time += traj_result.solve_time + if torch.count_nonzero(traj_result.success) > 0: + break + seed_traj = traj_result.optimized_seeds.detach().clone() + newton_iters = 4 + + traj_result.solve_time = finetune_time result.finetune_time = traj_result.solve_time @@ -1667,13 +1756,15 @@ class MotionGen(MotionGenConfig): result.debug_info["finetune_trajopt_result"] = traj_result elif plan_config.enable_finetune_trajopt: traj_result.success = traj_result.success[0:1] + # if torch.count_nonzero(result.success) == 0: + result.status = "Opt Fail" result.solve_time += traj_result.solve_time + result.finetune_time result.trajopt_time = traj_result.solve_time result.trajopt_attempts = 1 result.success = traj_result.success - if torch.count_nonzero(result.success) == 0: - result.status = "Opt Fail" + if plan_config.enable_finetune_trajopt and torch.count_nonzero(result.success) == 0: + result.status = "Finetune Opt Fail" result.interpolated_plan = traj_result.interpolated_solution.trim_trajectory( 0, traj_result.path_buffer_last_tstep[0] @@ -1696,15 +1787,17 @@ class MotionGen(MotionGenConfig): ) -> MotionGenResult: trajopt_seed_traj = None trajopt_seed_success = None - trajopt_newton_iters = None + trajopt_newton_iters = self.js_trajopt_solver.newton_iters graph_success = 0 if len(start_state.shape) == 1: log_error("Joint state should be not a vector (dof) should be (bxdof)") result = MotionGenResult(cspace_error=torch.zeros((1), device=self.tensor_args.device)) + if self.store_debug_in_result: + result.debug_info = {} # do graph search: - if plan_config.enable_graph: + if plan_config.enable_graph and False: start_config = torch.zeros( (solve_state.num_graph_seeds, self.js_trajopt_solver.dof), device=self.tensor_args.device, @@ -1782,7 +1875,6 @@ class MotionGen(MotionGenConfig): # do trajopt: if plan_config.enable_opt: with profiler.record_function("motion_gen/setup_trajopt_seeds"): - # self._trajopt_goal_config[:, :ik_success] = goal_config goal = Goal( current_state=start_state, @@ -1815,12 +1907,16 @@ class MotionGen(MotionGenConfig): batch_mode=False, seed_success=trajopt_seed_success, ) - trajopt_seed_traj = trajopt_seed_traj.view( - self.js_trajopt_solver.num_seeds * 1, - 1, - self.trajopt_solver.action_horizon, - self._dof, - ).contiguous() + trajopt_seed_traj = ( + trajopt_seed_traj.view( + self.js_trajopt_solver.num_seeds * 1, + 1, + self.trajopt_solver.action_horizon, + self._dof, + ) + .contiguous() + .clone() + ) if plan_config.enable_finetune_trajopt: og_value = self.trajopt_solver.interpolation_type self.js_trajopt_solver.interpolation_type = InterpolateType.LINEAR_CUDA @@ -1830,8 +1926,8 @@ class MotionGen(MotionGenConfig): goal, solve_state, trajopt_seed_traj, - num_seeds_override=solve_state.num_trajopt_seeds * 1, - newton_iters=trajopt_newton_iters, + num_seeds_override=solve_state.num_trajopt_seeds, + newton_iters=trajopt_newton_iters + 2, return_all_solutions=plan_config.enable_finetune_trajopt, trajopt_instance=self.js_trajopt_solver, ) @@ -1853,15 +1949,15 @@ class MotionGen(MotionGenConfig): torch.max(traj_result.optimized_dt[traj_result.success]), self.trajopt_solver.interpolation_dt, ) - og_dt = self.js_trajopt_solver.solver_dt + og_dt = self.js_trajopt_solver.solver_dt.clone() self.js_trajopt_solver.update_solver_dt(scaled_dt.item()) - traj_result = self._solve_trajopt_from_solve_state( goal, solve_state, seed_traj, trajopt_instance=self.js_trajopt_solver, - num_seeds_override=solve_state.num_trajopt_seeds * self.noisy_trajopt_seeds, + num_seeds_override=solve_state.num_trajopt_seeds, + newton_iters=trajopt_newton_iters + 4, ) self.js_trajopt_solver.update_solver_dt(og_dt) @@ -2281,6 +2377,7 @@ class MotionGen(MotionGenConfig): joint_names=self.rollout_fn.joint_names, ).repeat_seeds(batch) state = self.rollout_fn.compute_kinematics(start_state) + link_poses = state.link_pose if n_goalset == -1: retract_pose = Pose(state.ee_pos_seq, quaternion=state.ee_quat_seq) @@ -2355,6 +2452,7 @@ class MotionGen(MotionGenConfig): "graph_time": 0, "trajopt_time": 0, "trajopt_attempts": 0, + "finetune_time": 0, } result = None goal = Goal(goal_state=goal_state, current_state=start_state) @@ -2367,23 +2465,36 @@ class MotionGen(MotionGenConfig): n_envs=1, n_goalset=1, ) + force_graph = plan_config.enable_graph + for n in range(plan_config.max_attempts): - traj_result = self._plan_js_from_solve_state( + result = self._plan_js_from_solve_state( solve_state, start_state, goal_state, plan_config=plan_config ) - time_dict["trajopt_time"] += traj_result.solve_time + time_dict["trajopt_time"] += result.solve_time + time_dict["graph_time"] += result.graph_time + time_dict["finetune_time"] += result.finetune_time time_dict["trajopt_attempts"] = n + if plan_config.enable_graph_attempt is not None and ( + n >= plan_config.enable_graph_attempt - 1 and not plan_config.enable_graph + ): + plan_config.enable_graph = True + if plan_config.disable_graph_attempt is not None and ( + n >= plan_config.disable_graph_attempt - 1 and not force_graph + ): + plan_config.enable_graph = False - if result is None: - result = traj_result - - if traj_result.success.item(): + if result.success.item(): + break + if not result.valid_query: break - result.solve_time = time_dict["trajopt_time"] + result.graph_time = time_dict["graph_time"] + result.finetune_time = time_dict["finetune_time"] + result.trajopt_time = time_dict["trajopt_time"] + result.solve_time = result.trajopt_time + result.graph_time + result.finetune_time result.total_time = result.solve_time - if self.store_debug_in_result: - result.debug_info = {"trajopt_result": traj_result} + result.attempts = n return result diff --git a/src/curobo/wrap/reacher/trajopt.py b/src/curobo/wrap/reacher/trajopt.py index 1ffbf25..8ce1dbf 100644 --- a/src/curobo/wrap/reacher/trajopt.py +++ b/src/curobo/wrap/reacher/trajopt.py @@ -31,17 +31,14 @@ from curobo.opt.particle.parallel_es import ParallelES, ParallelESConfig from curobo.opt.particle.parallel_mppi import ParallelMPPI, ParallelMPPIConfig from curobo.rollout.arm_reacher import ArmReacher, ArmReacherConfig from curobo.rollout.cost.pose_cost import PoseCostMetric -from curobo.rollout.dynamics_model.integration_utils import ( - action_interpolate_kernel, - interpolate_kernel, -) +from curobo.rollout.dynamics_model.integration_utils import interpolate_kernel from curobo.rollout.rollout_base import Goal, RolloutBase, RolloutMetrics from curobo.types.base import TensorDeviceType from curobo.types.robot import JointState, RobotConfig from curobo.types.tensor import T_BDOF, T_DOF, T_BValue_bool, T_BValue_float from curobo.util.helpers import list_idx_if_not_none from curobo.util.logger import log_error, log_info, log_warn -from curobo.util.torch_utils import get_torch_jit_decorator, is_torch_compile_available +from curobo.util.torch_utils import get_torch_jit_decorator from curobo.util.trajectory import ( InterpolateType, calculate_dt_no_clamp, @@ -69,7 +66,7 @@ class TrajOptSolverConfig: num_seeds: int = 1 bias_node: Optional[T_DOF] = None interpolation_dt: float = 0.01 - traj_evaluator_config: TrajEvaluatorConfig = TrajEvaluatorConfig() + traj_evaluator_config: Optional[TrajEvaluatorConfig] = None traj_evaluator: Optional[TrajEvaluator] = None evaluate_interpolated_trajectory: bool = True cspace_threshold: float = 0.1 @@ -109,7 +106,7 @@ class TrajOptSolverConfig: seed_ratio: Dict[str, int] = {"linear": 1.0, "bias": 0.0, "start": 0.0, "end": 0.0}, use_particle_opt: bool = True, collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.MESH, - traj_evaluator_config: TrajEvaluatorConfig = TrajEvaluatorConfig(), + traj_evaluator_config: Optional[TrajEvaluatorConfig] = None, traj_evaluator: Optional[TrajEvaluator] = None, minimize_jerk: bool = True, use_gradient_descent: bool = False, @@ -172,7 +169,6 @@ class TrajOptSolverConfig: config_data["cost"]["pose_cfg"]["project_distance"] = project_pose_to_goal_frame grad_config_data["cost"]["pose_cfg"]["project_distance"] = project_pose_to_goal_frame - config_data["model"]["horizon"] = traj_tsteps grad_config_data["model"]["horizon"] = traj_tsteps if minimize_jerk is not None: @@ -325,6 +321,12 @@ class TrajOptSolverConfig: sync_cuda_time=sync_cuda_time, ) trajopt = WrapBase(cfg) + if traj_evaluator_config is None: + traj_evaluator_config = TrajEvaluatorConfig.from_basic( + min_dt=interpolation_dt, + dof=robot_cfg.kinematics.dof, + tensor_args=tensor_args, + ) trajopt_cfg = TrajOptSolverConfig( robot_config=robot_cfg, rollout_fn=aux_rollout, @@ -375,6 +377,7 @@ class TrajResult(Sequence): raw_solution: Optional[JointState] = None raw_action: Optional[torch.Tensor] = None goalset_index: Optional[torch.Tensor] = None + optimized_seeds: Optional[torch.Tensor] = None def __getitem__(self, idx): # position_error = rotation_error = cspace_error = path_buffer_last_tstep = None @@ -405,6 +408,7 @@ class TrajResult(Sequence): rotation_error=idx_vals[4], cspace_error=idx_vals[5], goalset_index=idx_vals[6], + optimized_seeds=self.optimized_seeds, ) def __len__(self): @@ -918,6 +922,7 @@ class TrajOptSolver(TrajOptSolverConfig): raw_solution=result.action, raw_action=result.raw_action, goalset_index=result.metrics.goalset_index, + optimized_seeds=result.raw_action, ) else: # get path length: @@ -938,79 +943,38 @@ class TrajOptSolver(TrajOptSolverConfig): ) with profiler.record_function("trajopt/best_select"): - if True: # not get_torch_jit_decorator() == torch.jit.script: - # This only works if torch compile is available: - ( - idx, - position_error, - rotation_error, - cspace_error, - goalset_index, - opt_dt, - success, - ) = jit_trajopt_best_select( - success, - smooth_label, - result.metrics.cspace_error, - result.metrics.pose_error, - result.metrics.position_error, - result.metrics.rotation_error, - result.metrics.goalset_index, - result.metrics.cost, - smooth_cost, - batch_mode, - goal.batch, - num_seeds, - self._col, - opt_dt, - ) - if batch_mode: - last_tstep = [last_tstep[i] for i in idx] - else: - last_tstep = [last_tstep[idx.item()]] - best_act_seq = result.action[idx] - best_raw_action = result.raw_action[idx] - interpolated_traj = interpolated_trajs[idx] - + ( + idx, + position_error, + rotation_error, + cspace_error, + goalset_index, + opt_dt, + success, + ) = jit_trajopt_best_select( + success, + smooth_label, + result.metrics.cspace_error, + result.metrics.pose_error, + result.metrics.position_error, + result.metrics.rotation_error, + result.metrics.goalset_index, + result.metrics.cost, + smooth_cost, + batch_mode, + goal.batch, + num_seeds, + self._col, + opt_dt, + ) + if batch_mode: + last_tstep = [last_tstep[i] for i in idx] else: - success[~smooth_label] = False - # get the best solution: - if result.metrics.pose_error is not None: - convergence_error = result.metrics.pose_error[..., -1] - elif result.metrics.cspace_error is not None: - convergence_error = result.metrics.cspace_error[..., -1] - else: - raise ValueError( - "convergence check requires either goal_pose or goal_state" - ) - running_cost = torch.mean(result.metrics.cost, dim=-1) * 0.0001 - error = convergence_error + smooth_cost + running_cost - error[~success] += 10000.0 - if batch_mode: - idx = torch.argmin(error.view(goal.batch, num_seeds), dim=-1) - idx = idx + num_seeds * self._col - last_tstep = [last_tstep[i] for i in idx] - success = success[idx] - else: - idx = torch.argmin(error, dim=0) + last_tstep = [last_tstep[idx.item()]] + best_act_seq = result.action[idx] + best_raw_action = result.raw_action[idx] + interpolated_traj = interpolated_trajs[idx] - last_tstep = [last_tstep[idx.item()]] - success = success[idx : idx + 1] - - best_act_seq = result.action[idx] - best_raw_action = result.raw_action[idx] - interpolated_traj = interpolated_trajs[idx] - goalset_index = position_error = rotation_error = cspace_error = None - if result.metrics.position_error is not None: - position_error = result.metrics.position_error[idx, -1] - if result.metrics.rotation_error is not None: - rotation_error = result.metrics.rotation_error[idx, -1] - if result.metrics.cspace_error is not None: - cspace_error = result.metrics.cspace_error[idx, -1] - if result.metrics.goalset_index is not None: - goalset_index = result.metrics.goalset_index[idx, -1] - - opt_dt = opt_dt[idx] if self.sync_cuda_time: torch.cuda.synchronize() if len(best_act_seq.shape) == 3: @@ -1046,6 +1010,7 @@ class TrajOptSolver(TrajOptSolverConfig): raw_solution=best_act_seq, raw_action=best_raw_action, goalset_index=goalset_index, + optimized_seeds=result.raw_action, ) return traj_result @@ -1307,6 +1272,10 @@ class TrajOptSolver(TrajOptSolverConfig): if isinstance(rollout, ArmReacher) ] + @property + def newton_iters(self): + return self._og_newton_iters + @get_torch_jit_decorator() def jit_feasible_success( diff --git a/src/curobo/wrap/reacher/types.py b/src/curobo/wrap/reacher/types.py index 0f84a7a..913b086 100644 --- a/src/curobo/wrap/reacher/types.py +++ b/src/curobo/wrap/reacher/types.py @@ -8,22 +8,25 @@ # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. # +""" Module contains custom types and dataclasses used across reacher solvers. """ from __future__ import annotations # Standard Library from dataclasses import dataclass from enum import Enum -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, Tuple, Union # CuRobo from curobo.rollout.rollout_base import Goal from curobo.types.base import TensorDeviceType from curobo.types.math import Pose from curobo.types.robot import JointState +from curobo.types.tensor import T_BDOF class ReacherSolveType(Enum): - # TODO: how to differentiate between goal pose and goal config? + """Enum for different types of problems solved with reacher solvers.""" + SINGLE = 0 GOALSET = 1 BATCH = 2 @@ -34,20 +37,46 @@ class ReacherSolveType(Enum): @dataclass class ReacherSolveState: + """Dataclass for storing the current problem type of a reacher solver.""" + + #: Type of problem solved by the reacher solver. solve_type: ReacherSolveType + + #: Number of problems in the batch. batch_size: int + + #: Number of environments in the batch. n_envs: int + + #: Number of goals per problem. Only valid for goalset problems. n_goalset: int = 1 + + #: Flag to indicate if the problems use different world environments in the batch. batch_env: bool = False + + #: Flag to indicate if the problems use different retract configurations in the batch. batch_retract: bool = False + + #: Flag to indicate if there is more than 1 problem to be solved. batch_mode: bool = False + + #: Number of seeds for each problem. num_seeds: Optional[int] = None + + #: Number of seeds for inverse kinematics problems. num_ik_seeds: Optional[int] = None + + #: Number of seeds for graph search problems. num_graph_seeds: Optional[int] = None + + #: Number of seeds for trajectory optimization problems. num_trajopt_seeds: Optional[int] = None + + #: Number of seeds for model predictive control problems. num_mpc_seeds: Optional[int] = None def __post_init__(self): + """Post init method to set default flags based on input values.""" if self.n_envs == 1: self.batch_env = False else: @@ -63,7 +92,8 @@ class ReacherSolveState: if self.num_seeds is None: self.num_seeds = self.num_mpc_seeds - def clone(self): + def clone(self) -> ReacherSolveState: + """Method to create a deep copy of the current reacher solve state.""" return ReacherSolveState( solve_type=self.solve_type, n_envs=self.n_envs, @@ -79,10 +109,12 @@ class ReacherSolveState: num_mpc_seeds=self.num_mpc_seeds, ) - def get_batch_size(self): + def get_batch_size(self) -> int: + """Method to get total number of optimization problems in the batch including seeds.""" return self.num_seeds * self.batch_size - def get_ik_batch_size(self): + def get_ik_batch_size(self) -> int: + """Method to get total number of IK problems in the batch including seeds.""" return self.num_ik_seeds * self.batch_size def create_goal_buffer( @@ -92,8 +124,24 @@ class ReacherSolveState: retract_config: Optional[T_BDOF] = None, link_poses: Optional[Dict[str, Pose]] = None, tensor_args: TensorDeviceType = TensorDeviceType(), - ): - # TODO: Refactor to account for num_ik_seeds or num_trajopt_seeds + ) -> Goal: + """Method to create a goal buffer from goal pose and other problem targets. + + Args: + goal_pose: Pose to reach with the end effector. + goal_state: Joint configuration to reach. If None, the goal is to reach the pose. + retract_config: Joint configuration to use for L2 regularization. If None, + `retract_config` from robot configuration file is used. An alternative value is to + use the start state as the retract configuration. + link_poses: Dictionary of link poses to reach. This is only required for multi-link + pose reaching, where the goal is to reach multiple poses with different links. + tensor_args: Device and floating precision. + + Returns: + Goal buffer with the goal pose, goal state, retract state, and link poses. + """ + + # NOTE: Refactor to account for num_ik_seeds or num_trajopt_seeds batch_retract = True if retract_config is None or retract_config.shape[0] != goal_pose.batch: batch_retract = False @@ -122,7 +170,28 @@ class ReacherSolveState: current_solve_state: Optional[ReacherSolveState] = None, current_goal_buffer: Optional[Goal] = None, tensor_args: TensorDeviceType = TensorDeviceType(), - ): + ) -> Tuple[ReacherSolveState, Goal, bool]: + """Method to update the goal buffer with new goal pose and other problem targets. + + Args: + goal_pose: Pose to reach with the end effector. + goal_state: Joint configuration to reach. If None, the goal is to reach the pose. + retract_config: Joint configuration to use for L2 regularization. If None, + `retract_config` from robot configuration file is used. An alternative value is to + use the start state as the retract configuration. + link_poses: Dictionary of link poses to reach. This is only required for multi-link + pose reaching, where the goal is to reach multiple poses with different links. To + use this, + :var:`curobo.cuda_robot_model.cuda_robot_model.CudaRobotModelConfig.link_names` + should have the link names to reach. + current_solve_state: Current reacher solve state. + current_goal_buffer: Current goal buffer. + tensor_args: Device and floating precision. + + Returns: + Tuple of updated reacher solve state, goal buffer, and a flag indicating if the goal + buffer reference has changed which is useful to break existing CUDA Graphs. + """ solve_state = self # create goal buffer by comparing to existing solve type update_reference = False @@ -167,7 +236,19 @@ class ReacherSolveState: current_solve_state: Optional[ReacherSolveState] = None, current_goal_buffer: Optional[Goal] = None, tensor_args: TensorDeviceType = TensorDeviceType(), - ): + ) -> Tuple[ReacherSolveState, Goal, bool]: + """Method to update the goal buffer with values from new Rollout goal. + + Args: + goal: Rollout goal to update the goal buffer. + current_solve_state: Current reacher solve state. + current_goal_buffer: Current goal buffer. + tensor_args: Device and floating precision. + + Returns: + Tuple of updated reacher solve state, goal buffer, and a flag indicating if the goal + buffer reference has changed which is useful to break existing CUDA Graphs. + """ solve_state = self update_reference = False if ( @@ -204,6 +285,8 @@ class ReacherSolveState: @dataclass class MotionGenSolverState: + """Dataclass for storing the current state of a motion generation solver.""" + solve_type: ReacherSolveType ik_solve_state: ReacherSolveState trajopt_solve_state: ReacherSolveState @@ -215,6 +298,22 @@ def get_padded_goalset( current_goal_buffer: Goal, new_goal_pose: Pose, ) -> Union[Pose, None]: + """Method to pad number of goals in goalset to match the cached goal buffer. + + This allows for creating a goalset problem with large number of goals during the first call, + and subsequent calls can have fewer goals. This function will pad the new goalset with the + first goal to match the cached goal buffer's shape. + + Args: + solve_state: New problem's solve state. + current_solve_state: Current solve state. + current_goal_buffer: Current goal buffer. + new_goal_pose: Padded goal pose to match the cached goal buffer's shape. + + Returns: + Padded goal pose to match the cached goal buffer's shape. If the new goal can't be padded, + returns None. + """ if ( current_solve_state.solve_type == ReacherSolveType.GOALSET and solve_state.solve_type == ReacherSolveType.SINGLE diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..e635c76 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,15 @@ +# +# 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 os + +os.environ["CUROBO_TORCH_COMPILE_DISABLE"] = str(1) diff --git a/tests/motion_gen_api_test.py b/tests/motion_gen_api_test.py index 4c852b3..7c672de 100644 --- a/tests/motion_gen_api_test.py +++ b/tests/motion_gen_api_test.py @@ -80,28 +80,25 @@ def test_motion_gen_lock_js_update(): kin_state = motion_gen_instance.compute_kinematics(start_state) ee_pose = kin_state.ee_pose.clone() - # test motion gen: plan_start = start_state.clone() plan_start.position[..., :-2] += 0.1 - result = motion_gen_instance.plan_single(plan_start, ee_pose) + result = motion_gen_instance.plan_single(plan_start, ee_pose.clone()) assert result.success.item() - lock_js = {"base_x": 1.0, "base_y": 0.0, "base_z": 0.0} + lock_js = {"base_x": 2.0, "base_y": 0.0, "base_z": 0.0} motion_gen_instance.update_locked_joints(lock_js, robot_config) kin_state_new = motion_gen_instance.compute_kinematics(start_state) ee_pose_shift = kin_state_new.ee_pose.clone() - assert torch.norm(ee_pose.position[..., 0] - ee_pose_shift.position[..., 0]).item() == 1.0 + assert 2 - torch.norm(ee_pose.position[..., 0] - ee_pose_shift.position[..., 0]).item() <= 1e-5 assert torch.norm(ee_pose.position[..., 1:] - ee_pose_shift.position[..., 1:]).item() == 0.0 # test motion gen with new lock state: - - result = motion_gen_instance.plan_single(plan_start, ee_pose_shift) + result = motion_gen_instance.plan_single(plan_start, ee_pose_shift.clone()) assert result.success.item() - result = motion_gen_instance.plan_single( - plan_start, ee_pose, MotionGenPlanConfig(max_attempts=3) + plan_start, ee_pose.clone(), MotionGenPlanConfig(max_attempts=3) ) assert result.success.item() == False diff --git a/tests/motion_gen_goalset_test.py b/tests/motion_gen_goalset_test.py index 74f2ef2..a2bddbf 100644 --- a/tests/motion_gen_goalset_test.py +++ b/tests/motion_gen_goalset_test.py @@ -109,32 +109,35 @@ def test_batch_goalset_padded(motion_gen_batch): # run goalset planning motion_gen.reset() - retract_cfg = motion_gen.get_retract_config() + retract_cfg = motion_gen.get_retract_config().clone() state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1))) goal_pose = Pose( state.ee_pos_seq.repeat(3 * 3, 1).view(3, -1, 3).contiguous(), quaternion=state.ee_quat_seq.repeat(3 * 3, 1).view(3, -1, 4).contiguous(), - ) + ).clone() goal_pose.position[0, 1, 1] = 0.2 goal_pose.position[1, 0, 1] = 0.2 goal_pose.position[2, 1, 1] = 0.2 + retract_cfg = motion_gen.get_retract_config().clone() start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.2).repeat_seeds(3) - m_config = MotionGenPlanConfig(False, True, max_attempts=10, enable_graph_attempt=20) - result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config.clone()) + m_config = MotionGenPlanConfig(enable_graph_attempt=100, max_attempts=2) + result = motion_gen.plan_batch_goalset(start_state, goal_pose.clone(), m_config.clone()) # get final solutions: assert torch.count_nonzero(result.success) == result.success.shape[0] - reached_state = motion_gen.compute_kinematics(result.optimized_plan.trim_trajectory(-1)) + reached_state = motion_gen.compute_kinematics( + result.optimized_plan.trim_trajectory(-1).squeeze(1) + ) # goal_position = torch.cat( [ - goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0) + goal_pose.position[x, result.goalset_index[x], :].clone().unsqueeze(0) for x in range(len(result.goalset_index)) ] ) @@ -161,7 +164,7 @@ def test_batch_goalset_padded(motion_gen_batch): start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.2).repeat_seeds(3) - result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config) + result = motion_gen.plan_batch_goalset(start_state, goal_pose.clone(), m_config) # get final solutions: assert torch.count_nonzero(result.success) == result.success.shape[0] @@ -192,7 +195,7 @@ def test_batch_goalset_padded(motion_gen_batch): goal_pose.position[1, 0] -= 0.1 - result = motion_gen.plan_batch(start_state, goal_pose.clone(), m_config) + result = motion_gen.plan_batch(start_state, goal_pose, m_config) assert torch.count_nonzero(result.success) == 3 # get final solutions: diff --git a/tests/motion_gen_js_test.py b/tests/motion_gen_js_test.py new file mode 100644 index 0000000..36a5af8 --- /dev/null +++ b/tests/motion_gen_js_test.py @@ -0,0 +1,44 @@ +# +# 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. +# + + +# Third Party +import torch + +# CuRobo +from curobo.types.robot import JointState +from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig + + +def test_motion_gen_plan_js(): + world_file = "collision_table.yml" + robot_file = "ur5e.yml" + motion_gen_config = MotionGenConfig.load_from_robot_config( + robot_file, + world_file, + trajopt_tsteps=32, + use_cuda_graph=False, + num_trajopt_seeds=4, + fixed_iters_trajopt=True, + evaluate_interpolated_trajectory=True, + ) + motion_gen = MotionGen(motion_gen_config) + motion_gen.warmup(warmup_js_trajopt=True) + + retract_cfg = motion_gen.get_retract_config() + + start_state = JointState.from_position(retract_cfg.view(1, -1).clone()) + goal_state = JointState.from_position(retract_cfg.view(1, -1).clone()) + goal_state.position[:] = torch.as_tensor( + [1.000, -2.2000, 1.9000, -1.3830, -1.5700, 0.0000], device=motion_gen.tensor_args.device + ) + result = motion_gen.plan_single_js(start_state, goal_state, MotionGenPlanConfig(max_attempts=1)) + assert result.success.item() diff --git a/tests/multi_pose_test.py b/tests/multi_pose_test.py index 85168b0..1ca27b5 100644 --- a/tests/multi_pose_test.py +++ b/tests/multi_pose_test.py @@ -42,11 +42,12 @@ def test_multi_pose_franka(b_size: int): world_cfg, rotation_threshold=0.05, position_threshold=0.005, - num_seeds=30, + num_seeds=16, self_collision_check=True, self_collision_opt=True, use_cuda_graph=True, tensor_args=tensor_args, + regularization=False, ) ik_solver = IKSolver(ik_config) @@ -60,3 +61,41 @@ def test_multi_pose_franka(b_size: int): assert ( torch.count_nonzero(success).item() / b_size >= 0.9 ) # we check if atleast 90% are successful + + +@pytest.mark.parametrize( + "b_size", + [1, 10, 100], +) +def test_multi_pose_hand(b_size: int): + tensor_args = TensorDeviceType() + world_file = "collision_table.yml" + + robot_file = "iiwa_allegro.yml" + robot_data = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"] + + robot_cfg = RobotConfig.from_dict(robot_data) + + world_cfg = WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file))) + ik_config = IKSolverConfig.load_from_robot_config( + robot_cfg, + world_cfg, + rotation_threshold=0.05, + position_threshold=0.005, + num_seeds=16, + use_cuda_graph=True, + tensor_args=tensor_args, + regularization=False, + ) + ik_solver = IKSolver(ik_config) + + q_sample = ik_solver.sample_configs(b_size) + kin_state = ik_solver.fk(q_sample) + link_poses = kin_state.link_pose + goal = Pose(kin_state.ee_position, kin_state.ee_quaternion).clone() + result = ik_solver.solve_batch(goal, link_poses=link_poses) + + success = result.success + assert ( + torch.count_nonzero(success).item() / b_size >= 0.9 + ) # we check if atleast 90% are successful