diff --git a/CHANGELOG.md b/CHANGELOG.md index 99d8626..7088b4d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,23 @@ its affiliates is strictly prohibited. --> # Changelog +## Latest Commit + +### New Features +- Add mimic joint parsing and optimization support. Check `ur5e_robotiq_2f_140.yml`. + +### BugFixes & Misc. +- Fix bug in `WorldVoxelCollision` where `env_query_idx` was being overwritten. +- Fix bug in `WorldVoxelCollision` where parent collision types were not getting called in some cases. +- Change voxelization dimensions to include 1 extra voxel per dim. +- Added `seed` parameter to `IKSolverConfig`. +- Added `sampler_seed` parameter `RolloutConfig`. +- Fixed bug in `links_goal_pose` where tensor could be non contiguous. +- Improved `ik_solver` success by removing gaussian projection of seed samples. +- 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. + ## Version 0.7.0 ### Changes in default behavior - Increased default collision cache to 50 in RobotWorld. diff --git a/benchmark/curobo_benchmark.py b/benchmark/curobo_benchmark.py index 74cd07f..548692f 100644 --- a/benchmark/curobo_benchmark.py +++ b/benchmark/curobo_benchmark.py @@ -154,7 +154,7 @@ def load_curobo( # del robot_cfg["kinematics"] if ik_seeds is None: - ik_seeds = 24 + ik_seeds = 32 if graph_mode: trajopt_seeds = 4 @@ -231,7 +231,7 @@ def benchmark_mb( # load dataset: force_graph = False - file_paths = [motion_benchmaker_raw, mpinets_raw] + file_paths = [motion_benchmaker_raw, mpinets_raw][:] if args.demo: file_paths = [demo_raw] @@ -500,8 +500,8 @@ def benchmark_mb( write_robot_usd_path="benchmark/log/usd/assets/", robot_usd_local_reference="assets/", base_frame="/world_" + problem_name, - visualize_robot_spheres=False, - flatten_usd=False, + visualize_robot_spheres=True, + flatten_usd=True, ) if write_plot: # and result.optimized_dt.item() > 0.06: @@ -554,12 +554,12 @@ def benchmark_mb( start_state, Pose.from_list(pose), join_path("benchmark/log/usd/", problem_name) + "_debug", - write_ik=False, + write_ik=True, write_trajopt=True, - visualize_robot_spheres=False, + visualize_robot_spheres=True, grid_space=2, write_robot_usd_path="benchmark/log/usd/assets/", - # flatten_usd=True, + flatten_usd=True, ) print(result.status) diff --git a/benchmark/curobo_python_profile.py b/benchmark/curobo_python_profile.py index 797bac3..4a2457e 100644 --- a/benchmark/curobo_python_profile.py +++ b/benchmark/curobo_python_profile.py @@ -24,39 +24,52 @@ from curobo.geom.sdf.world import CollisionCheckerType from curobo.types.base import TensorDeviceType from curobo.types.math import Pose from curobo.types.robot import JointState, RobotConfig +from curobo.util.logger import setup_curobo_logger from curobo.util_file import get_robot_configs_path, get_robot_path, join_path, load_yaml +from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig -def demo_motion_gen(): - # Standard Library - +def demo_motion_gen(robot_file, motion_gen=None): st_time = time.time() - tensor_args = TensorDeviceType() - world_file = "collision_table.yml" - robot_file = "franka.yml" - motion_gen_config = MotionGenConfig.load_from_robot_config( - robot_file, - world_file, - tensor_args, - trajopt_tsteps=44, - collision_checker_type=CollisionCheckerType.PRIMITIVE, - use_cuda_graph=False, - num_trajopt_seeds=24, - num_graph_seeds=24, - evaluate_interpolated_trajectory=True, - interpolation_dt=0.02, - ) - motion_gen = MotionGen(motion_gen_config) - # st_time = time.time() - motion_gen.warmup(enable_graph=True, warmup_js_trajopt=False) - print("motion gen time:", time.time() - st_time) + if motion_gen is None: + setup_curobo_logger("warn") + # Standard Library + + tensor_args = TensorDeviceType() + world_file = "collision_table.yml" + robot_file = "ur5e.yml" + motion_gen_config = MotionGenConfig.load_from_robot_config( + robot_file, + world_file, + tensor_args, + trajopt_tsteps=32, + collision_checker_type=CollisionCheckerType.PRIMITIVE, + use_cuda_graph=True, + num_trajopt_seeds=4, + num_graph_seeds=4, + evaluate_interpolated_trajectory=True, + interpolation_dt=0.02, + ) + motion_gen = MotionGen(motion_gen_config) + + # st_time = time.time() + torch.cuda.synchronize() + print("LOAD TIME: ", time.time() - st_time) + st_time = time.time() + + motion_gen.warmup(enable_graph=True, warmup_js_trajopt=False) + + torch.cuda.synchronize() + + print("warmup TIME: ", time.time() - st_time) + + return motion_gen # print(time.time() - st_time) # return retract_cfg = motion_gen.get_retract_config() - print(retract_cfg) state = motion_gen.rollout_fn.compute_kinematics( JointState.from_position(retract_cfg.view(1, -1)) ) @@ -66,15 +79,58 @@ def demo_motion_gen(): result = motion_gen.plan( start_state, retract_pose, - enable_graph=True, + enable_graph=False, enable_opt=True, max_attempts=1, - need_graph_success=True, + # need_graph_success=True, ) - print(result.status) - print(result.optimized_plan.position.shape) traj = result.get_interpolated_plan() # $.position.view(-1, 7) # optimized plan - print("Trajectory Generated: ", result.success, result.optimized_dt.item()) + print("Trajectory Generated: ", result.success, time.time() - st_time) + return motion_gen + + +def demo_basic_ik(config_file="ur10e.yml"): + st_time = time.time() + + tensor_args = TensorDeviceType() + + config_file = load_yaml(join_path(get_robot_configs_path(), config_file)) + urdf_file = config_file["kinematics"]["urdf_path"] # Send global path starting with "/" + base_link = config_file["kinematics"]["base_link"] + ee_link = config_file["kinematics"]["ee_link"] + robot_cfg = RobotConfig.from_basic(urdf_file, base_link, ee_link, tensor_args) + + ik_config = IKSolverConfig.load_from_robot_config( + robot_cfg, + None, + rotation_threshold=0.05, + position_threshold=0.005, + num_seeds=20, + self_collision_check=False, + self_collision_opt=False, + tensor_args=tensor_args, + use_cuda_graph=False, + ) + ik_solver = IKSolver(ik_config) + torch.cuda.synchronize() + print("IK load time:", time.time() - st_time) + st_time = time.time() + # print(kin_state) + + q_sample = ik_solver.sample_configs(100) + kin_state = ik_solver.fk(q_sample) + goal = Pose(kin_state.ee_position, kin_state.ee_quaternion) + + torch.cuda.synchronize() + print("FK time:", time.time() - st_time) + + st_time = time.time() + result = ik_solver.solve_batch(goal) + torch.cuda.synchronize() + print( + "Cold Start Solve Time(s) ", + result.solve_time, + ) def demo_basic_robot(): @@ -122,7 +178,7 @@ if __name__ == "__main__": parser.add_argument( "--save_path", type=str, - default="log/trace", + default="benchmark/log/trace", help="path to save file", ) parser.add_argument( @@ -137,12 +193,24 @@ if __name__ == "__main__": help="When True, runs startup for motion generation", default=False, ) + parser.add_argument( + "--motion_gen_plan", + action="store_true", + help="When True, runs startup for motion generation", + default=False, + ) parser.add_argument( "--kinematics", action="store_true", help="When True, runs startup for kinematics", default=False, ) + parser.add_argument( + "--inverse_kinematics", + action="store_true", + help="When True, runs startup for kinematics", + default=False, + ) parser.add_argument( "--motion_gen_once", action="store_true", @@ -172,21 +240,59 @@ if __name__ == "__main__": filename = join_path(args.save_path, args.file_name) + "_kinematics_trace.json" prof.export_chrome_trace(filename) - if args.motion_gen_once: - demo_motion_gen() - - if args.motion_gen: - for _ in range(5): - demo_motion_gen() + if args.inverse_kinematics: + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + demo_basic_ik(config_file) + filename = join_path(args.save_path, args.file_name) + "_inverse_kinematics_trace.json" + prof.export_chrome_trace(filename) pr = cProfile.Profile() pr.enable() - demo_motion_gen() + demo_basic_ik(config_file) + pr.disable() + filename = join_path(args.save_path, args.file_name) + "_inverse_kinematics_cprofile.prof" + pr.dump_stats(filename) + + if args.motion_gen_once: + pr = cProfile.Profile() + pr.enable() + demo_motion_gen(config_file) pr.disable() filename = join_path(args.save_path, args.file_name) + "_motion_gen_cprofile.prof" pr.dump_stats(filename) with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: - demo_motion_gen() + demo_motion_gen(config_file) + filename = join_path(args.save_path, args.file_name) + "_motion_gen_trace.json" + prof.export_chrome_trace(filename) + + if args.motion_gen_plan: + motion_gen = demo_motion_gen(config_file) + + pr = cProfile.Profile() + pr.enable() + demo_motion_gen(config_file, motion_gen) + pr.disable() + filename = join_path(args.save_path, args.file_name) + "_motion_gen_plan_cprofile.prof" + pr.dump_stats(filename) + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + demo_motion_gen(config_file, motion_gen) + filename = join_path(args.save_path, args.file_name) + "_motion_gen_plan_trace.json" + prof.export_chrome_trace(filename) + + if args.motion_gen: + for _ in range(5): + demo_motion_gen(config_file) + + pr = cProfile.Profile() + pr.enable() + demo_motion_gen(config_file) + pr.disable() + filename = join_path(args.save_path, args.file_name) + "_motion_gen_cprofile.prof" + pr.dump_stats(filename) + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + demo_motion_gen(config_file) filename = join_path(args.save_path, args.file_name) + "_motion_gen_trace.json" prof.export_chrome_trace(filename) diff --git a/benchmark/ik_benchmark.py b/benchmark/ik_benchmark.py index 2eb75ba..b4f127e 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=20, + num_seeds=30, self_collision_check=collision_free, self_collision_opt=collision_free, tensor_args=tensor_args, diff --git a/src/curobo/content/assets/robot/franka_description/franka_panda_mobile_xy_tz.urdf b/src/curobo/content/assets/robot/franka_description/franka_panda_mobile_xy_tz.urdf index a1ee12e..416da76 100644 --- a/src/curobo/content/assets/robot/franka_description/franka_panda_mobile_xy_tz.urdf +++ b/src/curobo/content/assets/robot/franka_description/franka_panda_mobile_xy_tz.urdf @@ -324,7 +324,6 @@ - diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/LICENSE b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/LICENSE new file mode 100644 index 0000000..316f610 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/LICENSE @@ -0,0 +1,23 @@ +Copyright (c) 2013, ROS-Industrial +All rights reserved. + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + + Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + Redistributions in binary form must reproduce the above copyright notice, this + list of conditions and the following disclaimer in the documentation and/or + other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/README.md b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/README.md new file mode 100644 index 0000000..0db6472 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/README.md @@ -0,0 +1,12 @@ +This is the README file from Robotiq's package for the 2F-140 gripper. The package's license file can also be found in this folder. +We used commit `4e3196c1c6a83f54acfa3dd8cf02b575ebba3e53` from [Robotiq's ROS Industrial Github repository](https://github.com/ros-industrial/robotiq) + +# Robotiq 140mm 2-Finger-Adaptive-Gripper + +This package contains the URDF files describing the 140mm stroke gripper from robotiq, also known as series **C3**. + +## Robot Visual +![140](https://user-images.githubusercontent.com/8356912/49428409-463f8580-f7a6-11e8-8278-5246acdc5c14.png) + +## Robot Collision +![1402](https://user-images.githubusercontent.com/8356912/49428407-463f8580-f7a6-11e8-9c4e-df69e478f107.png) diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_inner_finger.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_inner_finger.stl new file mode 100644 index 0000000..c15965c --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_inner_finger.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:043901d9c22b38d09b30d11326108ab3ef1445b4bd655ef313f94199f25f57ed +size 7284 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_inner_knuckle.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_inner_knuckle.stl new file mode 100644 index 0000000..1375e4e --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_inner_knuckle.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d341d986aa8aca7262565c039c04fb6b4c0f2f5c93443519eb7ca9bfc67ba17c +size 5484 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_outer_finger.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_outer_finger.stl new file mode 100644 index 0000000..4f19c84 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_outer_finger.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bd7da1b31b73f1aa1d61b26a582cea10f16f17396d54b8937890fa51547d26b0 +size 11684 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_outer_knuckle.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_outer_knuckle.stl new file mode 100644 index 0000000..c2818a6 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_140_outer_knuckle.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:37c1779f71fb5504a5898048a36f9f4bfce0f5e7039ff1dce53808b95c229777 +size 9784 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_base_link.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_base_link.stl new file mode 100644 index 0000000..3ef56dc --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_base_link.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:111e37f13a664989dd54226f80f521b32ea0b71c975282a16696b14be7cc9249 +size 86384 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_coupling.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_coupling.stl new file mode 100644 index 0000000..8958e11 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/collision/robotiq_arg2f_coupling.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ca9ffc28ed04193854b005358599dd9c3dc6fa92c8403e661fda94732d9ac25 +size 21184 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_inner_finger.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_inner_finger.stl new file mode 100644 index 0000000..dc22f55 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_inner_finger.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:be21c5e18c901d4747cbc8fa1a2af15dad7173ff701482a8517e33278432bb21 +size 33984 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_inner_knuckle.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_inner_knuckle.stl new file mode 100644 index 0000000..33365cf --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_inner_knuckle.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d61e9c5304d8015333856e4a26d99e32b695b6ada993253986b1afe8e396ab11 +size 43484 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_outer_finger.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_outer_finger.stl new file mode 100644 index 0000000..ff31816 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_outer_finger.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:666a92ee075f6f320ffb13b39995a5b374657cee90ded4c52c23ede20a812f34 +size 76084 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_outer_knuckle.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_outer_knuckle.stl new file mode 100644 index 0000000..d511e73 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_140_outer_knuckle.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f042edc44ede9773e7cad00f7d7354d0b7c1c6a6353fe897b2ca2e6ac71107fc +size 78384 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_base_link.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_base_link.stl new file mode 100644 index 0000000..6c39b86 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_base_link.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:74a62de75ae10cf77c60f2c49749b5d11f4c265f8624bbe7697a941fa86f6b3b +size 1054984 diff --git a/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_coupling.stl b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_coupling.stl new file mode 100644 index 0000000..e2289a1 --- /dev/null +++ b/src/curobo/content/assets/robot/kinova/kortex_description/grippers/robotiq_2f_140/meshes/visual/robotiq_arg2f_coupling.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4281e83002a25c20dc07c68b8d77da30a13e9a8401f157f6848ed8287d7cce44 +size 160684 diff --git a/src/curobo/content/assets/robot/simple/simple_mimic_robot.urdf b/src/curobo/content/assets/robot/simple/simple_mimic_robot.urdf new file mode 100644 index 0000000..825e1b2 --- /dev/null +++ b/src/curobo/content/assets/robot/simple/simple_mimic_robot.urdf @@ -0,0 +1,108 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/src/curobo/content/assets/robot/ur_description/ur5e_robotiq_2f_140.urdf b/src/curobo/content/assets/robot/ur_description/ur5e_robotiq_2f_140.urdf new file mode 100644 index 0000000..03f0ee0 --- /dev/null +++ b/src/curobo/content/assets/robot/ur_description/ur5e_robotiq_2f_140.urdf @@ -0,0 +1,691 @@ + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + 1 + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + 1 + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + 1 + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + 1 + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + 1 + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + 1 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/curobo/content/configs/robot/simple_mimic_robot.yml b/src/curobo/content/configs/robot/simple_mimic_robot.yml new file mode 100644 index 0000000..307ea40 --- /dev/null +++ b/src/curobo/content/configs/robot/simple_mimic_robot.yml @@ -0,0 +1,47 @@ +## +## 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. +## + +robot_cfg: + kinematics: + use_usd_kinematics: False + isaac_usd_path: null + usd_path: "robot/simple/three_link_mimic.usda" + usd_robot_root: "/base_link" + urdf_path: "robot/simple/simple_mimic_robot.urdf" + asset_root_path: "robot/simple" + base_link: "base_link" + ee_link: "ee_link" + collision_link_names: null + collision_spheres: null + collision_sphere_buffer: 0.004 # 0.0025 + extra_collision_spheres: null + use_global_cumul: True + self_collision_ignore: null + + + self_collision_buffer: null + mesh_link_names: + [ + "base_link", + "chain_1_link_1", + "chain_1_link_2", + "ee_link", + + ] + lock_joints: {"chain_1_active_joint_1": 0.2} + extra_links: null + cspace: + joint_names: ["chain_1_active_joint_1", "active_joint_2"] + retract_config: [0.3, 0.0] + null_space_weight: [1,1] + cspace_distance_weight: [1,1] + max_acceleration: 15.0 + max_jerk: 500.0 \ No newline at end of file diff --git a/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml b/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml new file mode 100644 index 0000000..e146757 --- /dev/null +++ b/src/curobo/content/configs/robot/ur5e_robotiq_2f_140.yml @@ -0,0 +1,265 @@ +## +## 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. +## + +robot_cfg: + kinematics: + usd_path: "robot/ur_description/ur5e_robotiq_2f_140.usd" + usd_robot_root: "/robot" + isaac_usd_path: "" + usd_flip_joints: {} + usd_flip_joint_limits: [] + + urdf_path: "robot/ur_description/ur5e_robotiq_2f_140.urdf" + asset_root_path: "robot/ur_description" + + base_link: "base_link" + ee_link: "grasp_frame" + link_names: null + lock_joints: {'finger_joint': 0.7} + + extra_links: null #{"attached_object":{"parent_link_name": "grasp_frame" , + #"link_name": "attached_object", "fixed_transform": [0,0,0,1,0,0,0], "joint_type":"FIXED", + #"joint_name": "attach_joint" }} + extra_collision_spheres: null #{"attached_object": 4} + + collision_link_names: ['shoulder_link','upper_arm_link', 'forearm_link', 'wrist_1_link', + 'wrist_2_link' ,'wrist_3_link', 'tool0', 'robotiq_arg2f_base_link', + 'left_outer_knuckle', + 'left_inner_knuckle', + 'left_outer_finger', + 'left_inner_finger', + 'left_inner_finger_pad', + 'right_outer_knuckle', + 'right_inner_knuckle', + 'right_outer_finger' , + 'right_inner_finger', + 'right_inner_finger_pad', + ] # List[str] + collision_spheres: + shoulder_link: + - "center": [0.0, 0.0, 0.0] + "radius": 0.1 + upper_arm_link: + - "center": [-0.416, -0.0, 0.143] + "radius": 0.078 + - "center": [-0.015, 0.0, 0.134] + "radius": 0.077 + - "center": [-0.14, 0.0, 0.138] + "radius": 0.062 + - "center": [-0.285, -0.001, 0.139] + "radius": 0.061 + - "center": [-0.376, 0.001, 0.138] + "radius": 0.077 + - "center": [-0.222, 0.001, 0.139] + "radius": 0.061 + - "center": [-0.055, 0.008, 0.14] + "radius": 0.07 + - "center": [-0.001, -0.002, 0.143] + "radius": 0.08 + forearm_link: + - "center": [-0.01, 0.002, 0.031] + "radius": 0.072 + - "center": [-0.387, 0.0, 0.014] + "radius": 0.057 + - "center": [-0.121, -0.0, 0.006] + "radius": 0.057 + - "center": [-0.206, 0.001, 0.007] + "radius": 0.057 + - "center": [-0.312, -0.001, 0.006] + "radius": 0.056 + - "center": [-0.057, 0.003, 0.008] + "radius": 0.065 + - "center": [-0.266, 0.0, 0.006] + "radius": 0.057 + - "center": [-0.397, -0.001, -0.018] + "radius": 0.052 + - "center": [-0.164, -0.0, 0.007] + "radius": 0.057 + wrist_1_link: + - "center": [-0.0, 0.0, -0.009] + "radius": 0.047 + - "center": [-0.0, 0.0, -0.052] + "radius": 0.047 + - "center": [-0.002, 0.027, -0.001] + "radius": 0.045 + - "center": [0.001, -0.01, 0.0] + "radius": 0.046 + wrist_2_link: + - "center": [0.0, -0.01, -0.001] + "radius": 0.047 + - "center": [0.0, 0.008, -0.001] + "radius": 0.047 + - "center": [0.001, -0.001, -0.036] + "radius": 0.047 + - "center": [0.001, -0.03, -0.0] + "radius": 0.047 + wrist_3_link: + - "center": [0.001, 0.001, -0.029] + "radius": 0.043 + tool0: + - "center": [0.001, 0.001, 0.05] + "radius": -0.01 #0.05 + robotiq_arg2f_base_link: + - "center": [0.0, -0.0, 0.04] + "radius": 0.035 + - "center": [0.0, -0.0, 0.02] + "radius": 0.035 + left_outer_finger: + - "center": [0.0, -0.01, 0.0] + "radius": 0.01 + - "center": [0.0, 0.015, -0.01] + "radius": 0.01 + + left_inner_finger_pad: + - "center": [0.0, -0.0, 0.005] + "radius": 0.01 + - "center": [0.0, 0.02, 0.005] + "radius": 0.01 + - "center": [0.0, -0.02, 0.005] + "radius": 0.01 + left_inner_knuckle: + - "center": [0.0, 0.04, -0.0] + "radius": 0.01 + - "center": [0.0, 0.06, -0.0] + "radius": 0.01 + - "center": [0.0, 0.08, -0.0] + "radius": 0.01 + left_inner_finger: + - "center": [0.0, -0.0, -0.0] + "radius": 0.01 + + left_outer_knuckle: + - "center": [0.0, 0.055, 0.01] + "radius": 0.01 + - "center": [0.0, 0.08, 0.005] + "radius": 0.01 + + + right_outer_finger: + - "center": [0.0, -0.01, 0.0] + "radius": 0.01 + - "center": [0.0, 0.015, -0.01] + "radius": 0.01 + + right_inner_finger_pad: + - "center": [0.0, -0.0, 0.005] + "radius": 0.01 + - "center": [0.0, 0.02, 0.005] + "radius": 0.01 + - "center": [0.0, -0.02, 0.005] + "radius": 0.01 + right_inner_knuckle: + - "center": [0.0, 0.04, -0.0] + "radius": 0.01 + - "center": [0.0, 0.06, -0.0] + "radius": 0.01 + - "center": [0.0, 0.08, -0.0] + "radius": 0.01 + right_inner_finger: + - "center": [0.0, -0.0, -0.0] + "radius": 0.01 + + right_outer_knuckle: + - "center": [0.0, 0.055, 0.01] + "radius": 0.01 + - "center": [0.0, 0.08, 0.005] + "radius": 0.01 + + + collision_sphere_buffer: 0.005 + + self_collision_ignore: { + "upper_arm_link": ["forearm_link", "shoulder_link"], + "forearm_link": ["wrist_1_link", "wrist_2_link", "wrist_3_link"], + "wrist_1_link": ["wrist_2_link","wrist_3_link","tool0", "robotiq_arg2f_base_link"], + "wrist_2_link": ["wrist_3_link", "tool0", "robotiq_arg2f_base_link"], + "wrist_3_link": ["tool0", "robotiq_arg2f_base_link"], + + "tool0": ['robotiq_arg2f_base_link', 'left_outer_finger', 'left_inner_finger_pad', + 'left_inner_knuckle', 'left_inner_finger', 'left_outer_knuckle', + 'right_outer_finger', 'right_inner_finger_pad', + 'right_inner_knuckle', 'right_inner_finger', 'right_outer_knuckle'], + + "robotiq_arg2f_base_link": ['left_outer_finger', 'left_inner_finger_pad', + 'left_inner_knuckle', 'left_inner_finger', 'left_outer_knuckle', + 'right_outer_finger', 'right_inner_finger_pad', + 'right_inner_knuckle', 'right_inner_finger', 'right_outer_knuckle'], + + "left_outer_finger": ['left_inner_finger_pad', + 'left_inner_knuckle', 'left_inner_finger', 'left_outer_knuckle', + 'right_outer_finger', 'right_inner_finger_pad', + 'right_inner_knuckle', 'right_inner_finger', 'right_outer_knuckle'], + + "left_inner_finger_pad": [ + 'left_inner_knuckle', 'left_inner_finger', 'left_outer_knuckle', + 'right_outer_finger', 'right_inner_finger_pad', + 'right_inner_knuckle', 'right_inner_finger', 'right_outer_knuckle'], + + "left_inner_knuckle": ['left_inner_finger', 'left_outer_knuckle', + 'right_outer_finger', 'right_inner_finger_pad', + 'right_inner_knuckle', 'right_inner_finger', 'right_outer_knuckle'], + + "left_inner_finger": ['left_outer_knuckle', + 'right_outer_finger', 'right_inner_finger_pad', + 'right_inner_knuckle', 'right_inner_finger', 'right_outer_knuckle'], + + "left_outer_knuckle": [ + 'right_outer_finger', 'right_inner_finger_pad', + 'right_inner_knuckle', 'right_inner_finger', 'right_outer_knuckle'], + + + + "right_outer_finger": ['right_inner_finger_pad', + 'right_inner_knuckle', 'right_inner_finger', 'right_outer_knuckle'], + + "right_inner_finger_pad": [ + 'right_inner_knuckle', 'right_inner_finger', 'right_outer_knuckle'], + + "right_inner_knuckle": [ 'right_inner_finger', 'right_outer_knuckle'], + + "right_inner_finger": [ 'right_outer_knuckle'], + + + + } + self_collision_buffer: { + 'shoulder_link': 0.01, + 'upper_arm_link': 0, + 'forearm_link': 0, + 'wrist_1_link': 0, + 'wrist_2_link': 0, + 'wrist_3_link' : 0, + 'tool0': 0.02, + } + + use_global_cumul: True + mesh_link_names: ['shoulder_link','upper_arm_link', 'forearm_link', 'wrist_1_link', + 'wrist_2_link' ,'wrist_3_link', 'robotiq_arg2f_base_link', + 'left_outer_knuckle', + 'left_inner_knuckle', + 'left_outer_finger', + 'left_inner_finger', + 'right_outer_knuckle', + 'right_inner_knuckle', + 'right_outer_finger' , + 'right_inner_finger', + ] # List[str] + + cspace: + joint_names: ['shoulder_pan_joint', 'shoulder_lift_joint', 'elbow_joint', 'wrist_1_joint', + 'wrist_2_joint', 'wrist_3_joint', + 'finger_joint'] + retract_config: [0.0, -2.2, 1.0, -1.383, -1.57, 0.00, 0.6] + null_space_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0,1.0] + cspace_distance_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + max_jerk: 500.0 + max_acceleration: 12.0 + position_limit_clip: 0.0 diff --git a/src/curobo/content/configs/task/gradient_ik.yml b/src/curobo/content/configs/task/gradient_ik.yml index c81517b..708fb1e 100644 --- a/src/curobo/content/configs/task/gradient_ik.yml +++ b/src/curobo/content/configs/task/gradient_ik.yml @@ -28,14 +28,16 @@ model: cost: pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] - weight: [1000, 20000, 30, 50] + weight: [200,4000,10,20] + #weight: [5000,500000,5,20] + 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: [1000, 20000, 30, 50] + weight: [200,4000,10,20] vec_convergence: [0.00, 0.000] terminal: False use_metric: True @@ -57,15 +59,15 @@ cost: lbfgs: - n_iters: 80 #60 - inner_iters: 20 + n_iters: 100 #60 + inner_iters: 25 cold_start_n_iters: null - min_iters: 20 + min_iters: 25 line_search_scale: [0.1, 0.3, 0.7, 1.0] fixed_iters: True - cost_convergence: 1e-7 - cost_delta_threshold: 1e-6 #0.0001 - cost_relative_threshold: 1.0 + cost_convergence: 0.001 + cost_delta_threshold: 1.0 #0.0001 + cost_relative_threshold: 0.999 epsilon: 0.01 #0.01 # used only in stable_mode history: 6 horizon: 1 @@ -73,6 +75,7 @@ lbfgs: n_problems: 1 store_debug: False use_cuda_kernel: True + use_shared_buffers_kernel: True stable_mode: True line_search_type: "approx_wolfe" #"wolfe" use_cuda_line_search_kernel: True diff --git a/src/curobo/content/configs/task/gradient_trajopt.yml b/src/curobo/content/configs/task/gradient_trajopt.yml index 02d3966..a40367a 100644 --- a/src/curobo/content/configs/task/gradient_trajopt.yml +++ b/src/curobo/content/configs/task/gradient_trajopt.yml @@ -32,7 +32,8 @@ 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,50000.0,30,60] #[150.0, 2000.0, 30, 40] + 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 @@ -62,7 +63,7 @@ cost: null_space_weight: [0.0] primitive_collision_cfg: - weight: 100000.0 + weight: 100000 use_sweep: True sweep_steps: 4 classify: False diff --git a/src/curobo/content/configs/task/particle_ik.yml b/src/curobo/content/configs/task/particle_ik.yml index 414f873..9a8f11c 100644 --- a/src/curobo/content/configs/task/particle_ik.yml +++ b/src/curobo/content/configs/task/particle_ik.yml @@ -34,7 +34,7 @@ cost: run_weight: 1.0 link_pose_cfg: vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] - weight: [0, 50, 10, 10] #[20.0, 100.0] + weight: [30, 50, 10, 10] #[20.0, 100.0] vec_convergence: [0.00, 0.000] # orientation, position terminal: False use_metric: True diff --git a/src/curobo/cuda_robot_model/cuda_robot_generator.py b/src/curobo/cuda_robot_model/cuda_robot_generator.py index 60dbea8..4845e82 100644 --- a/src/curobo/cuda_robot_model/cuda_robot_generator.py +++ b/src/curobo/cuda_robot_model/cuda_robot_generator.py @@ -314,6 +314,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): link_map=self._link_map, joint_map=self._joint_map, joint_map_type=self._joint_map_type, + joint_offset_map=self._joint_offset_map, store_link_map=self._store_link_map, link_chain_map=self._link_chain_map, link_names=self.link_names, @@ -332,6 +333,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): base_link=self.base_link, ee_link=self.ee_link, lock_jointstate=self.lock_jointstate, + mimic_joints=self._mimic_joint_data, ) if self.asset_root_path != "": self._kinematics_parser.add_absolute_path_to_link_meshes(self.asset_root_path) @@ -368,9 +370,8 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): @profiler.record_function("robot_generator/build_chain") def _build_chain(self, base_link, ee_link, other_links, link_names): self._n_dofs = 0 - self._controlled_joints = [] + self._controlled_links = [] self._bodies = [] - self._name_to_idx_map = dict() self.base_link = base_link self.ee_link = ee_link @@ -403,37 +404,61 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): self.non_fixed_joint_names = self.joint_names.copy() return chain_link_names + def _get_mimic_joint_data(self): + # get joint types: + mimic_joint_data = {} + for i in range(1, len(self._bodies)): + body = self._bodies[i] + if i in self._controlled_links: + if body.mimic_joint_name is not None: + if body.joint_name not in mimic_joint_data: + mimic_joint_data[body.joint_name] = [] + mimic_joint_data[body.joint_name].append(i) + return mimic_joint_data + @profiler.record_function("robot_generator/build_kinematics_tensors") def _build_kinematics_tensors(self, base_link, ee_link, link_names, chain_link_names): + self._active_joints = [] + self._mimic_joint_data = {} link_map = [0 for i in range(len(self._bodies))] store_link_map = [] # [-1 for i in range(len(self._bodies))] joint_map = [ - -1 if i not in self._controlled_joints else i for i in range(len(self._bodies)) + -1 if i not in self._controlled_links else i for i in range(len(self._bodies)) ] # joint_map_type = [ - -1 if i not in self._controlled_joints else i for i in range(len(self._bodies)) + -1 if i not in self._controlled_links else i for i in range(len(self._bodies)) ] all_joint_names = [] - j_count = 0 ordered_link_names = [] + joint_offset_map = [[1.0, 0.0]] # add body 0 details: 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] parent_name = body.parent_link_name link_map[i] = self._name_to_idx_map[parent_name] - all_joint_names.append(body.joint_name) + joint_offset_map.append(body.joint_offset) + joint_map_type[i] = body.joint_type.value if body.link_name in link_names: store_link_map.append(chain_link_names.index(body.link_name)) ordered_link_names.append(body.link_name) - if i in self._controlled_joints: - joint_map[i] = j_count - joint_map_type[i] = body.joint_type.value - j_count += 1 + if body.joint_name not in all_joint_names: + all_joint_names.append(body.joint_name) + if i in self._controlled_links: + joint_map[i] = self.joint_names.index(body.joint_name) + if body.mimic_joint_name is not None: + if body.joint_name not in self._mimic_joint_data: + self._mimic_joint_data[body.joint_name] = [] + self._mimic_joint_data[body.joint_name].append( + {"joint_offset": body.joint_offset, "joint_name": body.mimic_joint_name} + ) + else: + self._active_joints.append(i) self.link_names = ordered_link_names # do a for loop to get link matrix: link_chain_map = torch.eye( @@ -458,6 +483,10 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): self._store_link_map = torch.as_tensor( store_link_map, device=self.tensor_args.device, dtype=torch.int16 ) + self._joint_offset_map = torch.as_tensor( + joint_offset_map, device=self.tensor_args.device, dtype=torch.float32 + ) + self._joint_offset_map = self._joint_offset_map.view(-1).contiguous() self._link_chain_map = link_chain_map.to(device=self.tensor_args.device) self._fixed_transform = torch.cat((self._fixed_transform), dim=0).to( device=self.tensor_args.device @@ -489,12 +518,18 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): joint_data = self._get_joint_links(lock_joint_names) lock_links = list( - set( - [joint_data[j]["parent"] for j in joint_data.keys()] - + [joint_data[j]["child"] for j in joint_data.keys()] - ) + [joint_data[j]["parent"] for j in joint_data.keys()] + + [joint_data[j]["child"] for j in joint_data.keys()] ) - new_link_names = link_names + lock_links + + for k in lock_joint_names: + if "mimic" in joint_data[k]: + mimic_link_names = [[x["parent"], x["child"]] for x in joint_data[k]["mimic"]] + mimic_link_names = [x for xs in mimic_link_names for x in xs] + lock_links += mimic_link_names + lock_links = list(set(lock_links)) + + new_link_names = list(set(link_names + lock_links)) # rebuild kinematic tree with link names added to link pose computation: self._build_kinematics_tensors(base_link, ee_link, new_link_names, chain_link_names) @@ -520,6 +555,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): link_map=self._link_map, joint_map=self._joint_map, joint_map_type=self._joint_map_type, + joint_offset_map=self._joint_offset_map, store_link_map=self._store_link_map, link_chain_map=self._link_chain_map, link_names=self.link_names, @@ -531,14 +567,12 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): total_spheres=self.total_spheres, ) link_poses = self._get_link_poses(q, lock_links, kinematics_config) - # remove lock links from store map: store_link_map = [chain_link_names.index(l) for l in link_names] self._store_link_map = torch.as_tensor( store_link_map, device=self.tensor_args.device, dtype=torch.int16 ) self.link_names = link_names - # compute a fixed transform for fixing joints: with profiler.record_function("cuda_robot_generator/fix_locked_joints"): # convert tensors to cpu: @@ -554,14 +588,33 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): .multiply(link_poses.get_index(0, w_child)) ) # Make this joint as fixed - i = self._all_joint_names.index(j) + 1 + i = joint_data[j]["link_index"] + self._fixed_transform[i] = parent_t_child.get_matrix() + + if "mimic" in joint_data[j]: + for mimic_joint in joint_data[j]["mimic"]: + w_parent = lock_links.index(mimic_joint["parent"]) + w_child = lock_links.index(mimic_joint["child"]) + parent_t_child = ( + link_poses.get_index(0, w_parent) + .inverse() + .multiply(link_poses.get_index(0, w_child)) + ) + i_q = mimic_joint["link_index"] + self._fixed_transform[i_q] = parent_t_child.get_matrix() + self._controlled_links.remove(i_q) + self._joint_map_type[i_q] = -1 + self._joint_map[i_q] = -1 + + i = joint_data[j]["link_index"] self._joint_map_type[i] = -1 self._joint_map[i:] -= 1 self._joint_map[i] = -1 - self._n_dofs -= 1 - self._fixed_transform[i] = parent_t_child.get_matrix() - self._controlled_joints.remove(i) + self._controlled_links.remove(i) 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) @@ -715,7 +768,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): if not valid_data: break if torch.max(coll_cpu[i]) == -torch.inf: - print("skip", i) + log_info("skip" + str(i)) for j in range(i + 1, n_spheres): if sl_idx > thread_loc.shape[0] - 1: valid_data = False @@ -776,9 +829,10 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): rigid_body_params = self._kinematics_parser.get_link_parameters(link_name, base=base) self._bodies.append(rigid_body_params) if rigid_body_params.joint_type != JointType.FIXED: - self._controlled_joints.append(body_idx) - self.joint_names.append(rigid_body_params.joint_name) - self._n_dofs = self._n_dofs + 1 + self._controlled_links.append(body_idx) + if rigid_body_params.joint_name not in self.joint_names: + self.joint_names.append(rigid_body_params.joint_name) + self._n_dofs = self._n_dofs + 1 self._fixed_transform.append( torch.as_tensor( rigid_body_params.fixed_transform, @@ -790,21 +844,34 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): def _get_joint_links(self, joint_names: List[str]): j_data = {} + for j in joint_names: - for b in self._bodies: + for bi, b in enumerate(self._bodies): if b.joint_name == j: - j_data[j] = {"parent": b.parent_link_name, "child": b.link_name} + if j not in j_data: + j_data[j] = {} + if b.mimic_joint_name is None: + j_data[j]["parent"] = b.parent_link_name + j_data[j]["child"] = b.link_name + j_data[j]["link_index"] = bi + else: + if "mimic" not in j_data[j]: + j_data[j]["mimic"] = [] + j_data[j]["mimic"].append( + { + "parent": b.parent_link_name, + "child": b.link_name, + "link_index": bi, + "joint_offset": b.joint_offset, + } + ) + return j_data @profiler.record_function("robot_generator/get_link_poses") def _get_link_poses( self, q: torch.Tensor, link_names: List[str], kinematics_config: KinematicsTensorConfig ) -> Pose: - if q.is_contiguous(): - q_in = q.view(-1) - else: - q_in = q.reshape(-1) # .reshape(-1) - # q_in = q.reshape(-1) link_pos_seq = torch.zeros( (1, len(self.link_names), 3), device=self.tensor_args.device, @@ -831,12 +898,11 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): dtype=self.tensor_args.dtype, ) link_pos_seq, link_quat_seq, _ = get_cuda_kinematics( - # self._link_mat_seq, # data will be stored here link_pos_seq, link_quat_seq, batch_robot_spheres, global_cumul_mat, - q_in, + q, kinematics_config.fixed_transforms, kinematics_config.link_spheres, kinematics_config.link_map, # tells which link is attached to which link i @@ -845,6 +911,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): kinematics_config.store_link_map, kinematics_config.link_sphere_idx_map, # sphere idx map kinematics_config.link_chain_map, + kinematics_config.joint_offset_map, grad_out_q, self.use_global_cumul, ) @@ -873,7 +940,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): def _get_joint_position_velocity_limits(self): joint_limits = {"position": [[], []], "velocity": [[], []]} - for idx in self._controlled_joints: + for idx in self._active_joints: joint_limits["position"][0].append(self._bodies[idx].joint_limits[0]) joint_limits["position"][1].append(self._bodies[idx].joint_limits[1]) joint_limits["velocity"][0].append(self._bodies[idx].joint_velocity_limits[0]) @@ -897,7 +964,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig): ] ) # clip joint position: - # TODO: change this to be per joint + # NOTE: change this to be per joint joint_limits["position"][0] += self.cspace.position_limit_clip joint_limits["position"][1] -= self.cspace.position_limit_clip joint_limits["velocity"][0] *= self.cspace.velocity_scale diff --git a/src/curobo/cuda_robot_model/cuda_robot_model.py b/src/curobo/cuda_robot_model/cuda_robot_model.py index 016816b..3fbd31f 100644 --- a/src/curobo/cuda_robot_model/cuda_robot_model.py +++ b/src/curobo/cuda_robot_model/cuda_robot_model.py @@ -33,6 +33,7 @@ from curobo.curobolib.kinematics import get_cuda_kinematics from curobo.geom.types import Mesh, Sphere from curobo.types.base import TensorDeviceType from curobo.types.math import Pose +from curobo.types.state import JointState from curobo.util.logger import log_error from curobo.util_file import get_robot_path, join_path, load_yaml @@ -322,7 +323,6 @@ class CudaRobotModel(CudaRobotModelConfig): def _cuda_forward(self, q): link_pos, link_quat, robot_spheres = get_cuda_kinematics( - # self._link_mat_seq, # data will be stored here self._link_pos_seq, self._link_quat_seq, self._batch_robot_spheres, @@ -336,11 +336,10 @@ class CudaRobotModel(CudaRobotModelConfig): self.kinematics_config.store_link_map, self.kinematics_config.link_sphere_idx_map, # sphere idx map self.kinematics_config.link_chain_map, + self.kinematics_config.joint_offset_map, self._grad_out_q, self.use_global_cumul, ) - # if(robot_spheres.shape[0]<10): - # print(robot_spheres) return link_pos, link_quat, robot_spheres @property @@ -381,6 +380,31 @@ class CudaRobotModel(CudaRobotModelConfig): def lock_jointstate(self): return self.kinematics_config.lock_jointstate + def get_full_js(self, js: JointState): + all_joint_names = self.all_articulated_joint_names + lock_joint_state = self.lock_jointstate + + new_js = js.get_augmented_joint_state(all_joint_names, lock_joint_state) + return new_js + + def get_mimic_js(self, js: JointState): + if self.kinematics_config.mimic_joints is None: + return None + extra_joints = {"position": [], "joint_names": []} + # for every joint in mimic_joints, get active joint name + for j in self.kinematics_config.mimic_joints: + active_q = js.position[..., js.joint_names.index(j)] + for k in self.kinematics_config.mimic_joints[j]: + extra_joints["joint_names"].append(k["joint_name"]) + extra_joints["position"].append( + k["joint_offset"][0] * active_q + k["joint_offset"][1] + ) + extra_js = JointState.from_position( + position=torch.stack(extra_joints["position"]), joint_names=extra_joints["joint_names"] + ) + new_js = js.get_augmented_joint_state(js.joint_names + extra_js.joint_names, extra_js) + return new_js + @property def ee_link(self): return self.kinematics_config.ee_link diff --git a/src/curobo/cuda_robot_model/kinematics_parser.py b/src/curobo/cuda_robot_model/kinematics_parser.py index a247498..f1d5bbe 100644 --- a/src/curobo/cuda_robot_model/kinematics_parser.py +++ b/src/curobo/cuda_robot_model/kinematics_parser.py @@ -11,7 +11,7 @@ # Standard Library from dataclasses import dataclass, field -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional # Third Party import numpy as np @@ -21,7 +21,6 @@ from curobo.cuda_robot_model.types import JointType from curobo.geom.types import Mesh from curobo.types.base import TensorDeviceType from curobo.types.math import Pose -from curobo.util.logger import log_error @dataclass @@ -35,6 +34,8 @@ class LinkParams: joint_axis: Optional[np.ndarray] = None joint_id: Optional[int] = None joint_velocity_limits: List[float] = field(default_factory=lambda: [-2.0, 2.0]) + joint_offset: List[float] = field(default_factory=lambda: [1.0, 0.0]) + mimic_joint_name: Optional[str] = None @staticmethod def from_dict(dict_data): @@ -56,7 +57,7 @@ class KinematicsParser: # add extra links to parent: if self.extra_links is not None and len(list(self.extra_links.keys())) > 0: for i in self.extra_links: - self._parent_map[i] = self.extra_links[i].parent_link_name + self._parent_map[i] = {"parent": self.extra_links[i].parent_link_name} def build_link_parent(self): """Build a map of parent links to each link in the kinematic tree. @@ -78,7 +79,7 @@ class KinematicsParser: chain_links = [ee_link] link = ee_link while link != base_link: - link = self._parent_map[link] + link = self._parent_map[link]["parent"] # add link to chain: chain_links.append(link) chain_links.reverse() diff --git a/src/curobo/cuda_robot_model/types.py b/src/curobo/cuda_robot_model/types.py index f9ba859..8c4c20e 100644 --- a/src/curobo/cuda_robot_model/types.py +++ b/src/curobo/cuda_robot_model/types.py @@ -228,28 +228,87 @@ class CSpaceConfig: @dataclass class KinematicsTensorConfig: + """Stores robot's kinematics parameters as Tensors to use in Kinematics computations. + + Use :meth:`curobo.cuda_robot_model.cuda_robot_generator.CudaRobotGenerator` to generate this + configuration from a urdf or usd. + + """ + + #: Static Homogenous Transform from parent link to child link for all links [n_links,4,4]. fixed_transforms: torch.Tensor + + #: index of fixed_transform given link index [n_links]. link_map: torch.Tensor + + #: joint index given link index [n_links]. joint_map: torch.Tensor + + #: type of joint given link index [n_links]. joint_map_type: torch.Tensor + + joint_offset_map: torch.Tensor + + #: index of link to write out pose [n_store_links]. store_link_map: torch.Tensor + + #: Mapping between each link to every other link, this is used to check + #: if a link is part of a serial chain formed by another link [n_links, n_links]. link_chain_map: torch.Tensor + + #: Name of links whose pose will be stored [n_store_links]. link_names: List[str] + + #: Joint limits joint_limits: JointLimits + + #: Name of joints which are not fixed. non_fixed_joint_names: List[str] + + #: Number of joints that are active. Each joint is only actuated along 1 dimension. n_dof: int + + #: Name of links which have a mesh. Currently only used for debugging and rendering. mesh_link_names: Optional[List[str]] = None + + #: Name of all actuated joints. joint_names: Optional[List[str]] = None + + #: lock_jointstate: Optional[JointState] = None + + #: + mimic_joints: Optional[dict] = None + + #: link_spheres: Optional[torch.Tensor] = None + + #: link_sphere_idx_map: Optional[torch.Tensor] = None + + #: link_name_to_idx_map: Optional[Dict[str, int]] = None + + #: total number of spheres that represent the robot's geometry. total_spheres: int = 0 + + #: Additional debug parameters. debug: Optional[Any] = None + + #: index of end-effector in stored link poses. ee_idx: int = 0 + + #: Cspace configuration cspace: Optional[CSpaceConfig] = None + + #: Name of base link. This is the root link from which all kinematic parameters were computed. base_link: str = "base_link" + + #: Name of end-effector link for which the Cartesian pose will be computed. ee_link: str = "ee_link" + + #: A copy of link spheres that is used as reference, in case the link_spheres get modified at + #: runtime. reference_link_spheres: Optional[torch.Tensor] = None def __post_init__(self): @@ -260,7 +319,7 @@ class KinematicsTensorConfig: if self.link_spheres is not None and self.reference_link_spheres is None: self.reference_link_spheres = self.link_spheres.clone() - def copy_(self, new_config: KinematicsTensorConfig): + def copy_(self, new_config: KinematicsTensorConfig) -> KinematicsTensorConfig: self.fixed_transforms.copy_(new_config.fixed_transforms) self.link_map.copy_(new_config.link_map) self.joint_map.copy_(new_config.joint_map) @@ -268,6 +327,7 @@ class KinematicsTensorConfig: self.store_link_map.copy_(new_config.store_link_map) self.link_chain_map.copy_(new_config.link_chain_map) self.joint_limits.copy_(new_config.joint_limits) + self.joint_offset_map.copy_(new_config.joint_offset_map) if new_config.link_spheres is not None and self.link_spheres is not None: self.link_spheres.copy_(new_config.link_spheres) if new_config.link_sphere_idx_map is not None and self.link_sphere_idx_map is not None: @@ -321,7 +381,8 @@ class KinematicsTensorConfig: ): """Update sphere parameters - #NOTE: This currently does not update self collision distances. + NOTE: This currently does not update self collision distances. + Args: link_name: _description_ sphere_position_radius: _description_ diff --git a/src/curobo/cuda_robot_model/urdf_kinematics_parser.py b/src/curobo/cuda_robot_model/urdf_kinematics_parser.py index f733798..d627c7d 100644 --- a/src/curobo/cuda_robot_model/urdf_kinematics_parser.py +++ b/src/curobo/cuda_robot_model/urdf_kinematics_parser.py @@ -9,7 +9,7 @@ # its affiliates is strictly prohibited. # # Standard Library -from typing import Dict, Optional +from typing import Dict, List, Optional # Third Party import numpy as np @@ -33,12 +33,13 @@ class UrdfKinematicsParser(KinematicsParser): load_meshes: bool = False, mesh_root: str = "", extra_links: Optional[Dict[str, LinkParams]] = None, + build_scene_graph: bool = False, ) -> None: # load robot from urdf: self._robot = yourdfpy.URDF.load( urdf_path, load_meshes=load_meshes, - build_scene_graph=False, + build_scene_graph=build_scene_graph, mesh_dir=mesh_root, filename_handler=yourdfpy.filename_handler_null, ) @@ -46,63 +47,114 @@ class UrdfKinematicsParser(KinematicsParser): def build_link_parent(self): self._parent_map = {} - for j in self._robot.joint_map: - self._parent_map[self._robot.joint_map[j].child] = self._robot.joint_map[j].parent - - def _find_parent_joint_of_link(self, link_name): - for j_idx, j in enumerate(self._robot.joint_map): - if self._robot.joint_map[j].child == link_name: - return j_idx, j - log_error("Link is not attached to any joint") + for jid, j in enumerate(self._robot.joint_map): + self._parent_map[self._robot.joint_map[j].child] = { + "parent": self._robot.joint_map[j].parent, + "jid": jid, + "joint_name": j, + } def _get_joint_name(self, idx): joint = self._robot.joint_names[idx] return joint + def _get_joint_limits(self, joint): + joint_type = joint.type + if joint_type != "continuous": + joint_limits = { + "effort": joint.limit.effort, + "lower": joint.limit.lower, + "upper": joint.limit.upper, + "velocity": joint.limit.velocity, + } + else: + log_warn("Converting continuous joint to revolute with limits[-10,10]") + joint_type = "revolute" + joint_limits = { + "effort": joint.limit.effort, + "lower": -10, + "upper": 10, + "velocity": joint.limit.velocity, + } + return joint_limits, joint_type + def get_link_parameters(self, link_name: str, base=False) -> LinkParams: link_params = self._get_from_extra_links(link_name) if link_params is not None: return link_params body_params = {} body_params["link_name"] = link_name - + mimic_joint_name = None if base: body_params["parent_link_name"] = None joint_transform = np.eye(4) joint_name = "base_joint" + active_joint_name = joint_name joint_type = "fixed" joint_limits = None joint_axis = None body_params["joint_id"] = 0 - else: - body_params["parent_link_name"] = self._parent_map[link_name] + body_params["joint_type"] = JointType.FIXED - jid, joint_name = self._find_parent_joint_of_link(link_name) + else: + parent_data = self._parent_map[link_name] + body_params["parent_link_name"] = parent_data["parent"] + + jid, joint_name = parent_data["jid"], parent_data["joint_name"] body_params["joint_id"] = jid joint = self._robot.joint_map[joint_name] + active_joint_name = joint_name joint_transform = joint.origin if joint_transform is None: joint_transform = np.eye(4) joint_type = joint.type joint_limits = None joint_axis = None + body_params["joint_type"] = JointType.FIXED + if joint_type != "fixed": - if joint_type != "continuous": - joint_limits = { - "effort": joint.limit.effort, - "lower": joint.limit.lower, - "upper": joint.limit.upper, - "velocity": joint.limit.velocity, - } - else: - log_warn("Converting continuous joint to revolute") - joint_type = "revolute" - joint_limits = { - "effort": joint.limit.effort, - "lower": -3.14 * 2, - "upper": 3.14 * 2, - "velocity": joint.limit.velocity, - } + joint_offset = [1.0, 0.0] + joint_limits, joint_type = self._get_joint_limits(joint) + + if joint.mimic is not None: + joint_offset = [joint.mimic.multiplier, joint.mimic.offset] + # read joint limits of active joint: + mimic_joint_name = joint_name + active_joint_name = joint.mimic.joint + active_joint = self._robot.joint_map[active_joint_name] + active_joint_limits, active_joint_type = self._get_joint_limits(active_joint) + # check to make sure mimic joint limits are not larger than active joint: + if ( + active_joint_limits["lower"] * joint_offset[0] + joint_offset[1] + < joint_limits["lower"] + ): + log_error( + "mimic joint can go out of it's lower limit as active joint has larger range " + + "FIX: make mimic joint's lower limit even lower " + + active_joint_name + + " " + + mimic_joint_name + ) + if ( + active_joint_limits["upper"] * joint_offset[0] + joint_offset[1] + > joint_limits["upper"] + ): + log_error( + "mimic joint can go out of it's upper limit as active joint has larger range " + + "FIX: make mimic joint's upper limit higher" + + active_joint_name + + " " + + mimic_joint_name + ) + if active_joint_limits["velocity"] * joint_offset[0] > joint_limits["velocity"]: + log_error( + "mimic joint can move at higher velocity than active joint," + + "increase velocity limit for mimic joint" + + active_joint_name + + " " + + mimic_joint_name + ) + joint_limits = active_joint_limits joint_axis = joint.axis @@ -111,45 +163,33 @@ class UrdfKinematicsParser(KinematicsParser): -1.0 * joint_limits["velocity"], joint_limits["velocity"], ] + if joint_type == "prismatic": + if abs(joint_axis[0]) == 1: + joint_type = JointType.X_PRISM + if abs(joint_axis[1]) == 1: + joint_type = JointType.Y_PRISM + if abs(joint_axis[2]) == 1: + joint_type = JointType.Z_PRISM + elif joint_type == "revolute": + if abs(joint_axis[0]) == 1: + joint_type = JointType.X_ROT + if abs(joint_axis[1]) == 1: + joint_type = JointType.Y_ROT + if abs(joint_axis[2]) == 1: + joint_type = JointType.Z_ROT + else: + 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] + body_params["joint_type"] = joint_type + body_params["joint_offset"] = joint_offset body_params["fixed_transform"] = joint_transform - body_params["joint_name"] = joint_name + body_params["joint_name"] = active_joint_name body_params["joint_axis"] = joint_axis + body_params["mimic_joint_name"] = mimic_joint_name - if joint_type == "fixed": - joint_type = JointType.FIXED - elif joint_type == "prismatic": - if joint_axis[0] == 1: - joint_type = JointType.X_PRISM - if joint_axis[1] == 1: - joint_type = JointType.Y_PRISM - if joint_axis[2] == 1: - joint_type = JointType.Z_PRISM - if joint_axis[0] == -1: - joint_type = JointType.X_PRISM_NEG - if joint_axis[1] == -1: - joint_type = JointType.Y_PRISM_NEG - if joint_axis[2] == -1: - joint_type = JointType.Z_PRISM_NEG - - elif joint_type == "revolute": - if joint_axis[0] == 1: - joint_type = JointType.X_ROT - if joint_axis[1] == 1: - joint_type = JointType.Y_ROT - if joint_axis[2] == 1: - joint_type = JointType.Z_ROT - if joint_axis[0] == -1: - joint_type = JointType.X_ROT_NEG - if joint_axis[1] == -1: - joint_type = JointType.Y_ROT_NEG - if joint_axis[2] == -1: - joint_type = JointType.Z_ROT_NEG - else: - log_error("Joint type not supported") - - body_params["joint_type"] = joint_type link_params = LinkParams(**body_params) return link_params @@ -194,3 +234,16 @@ class UrdfKinematicsParser(KinematicsParser): scale=m.scale, file_path=m.filename, ) + + @property + def root_link(self): + return self._robot.base_link + + def get_controlled_joint_names(self) -> List[str]: + j_list = [] + for k in self._parent_map.keys(): + joint_name = self._parent_map[k]["joint_name"] + joint = self._robot.joint_map[joint_name] + if joint.type != "fixed" and joint.mimic is None: + j_list.append(joint_name) + return j_list diff --git a/src/curobo/cuda_robot_model/usd_kinematics_parser.py b/src/curobo/cuda_robot_model/usd_kinematics_parser.py index 2715b0c..806324f 100644 --- a/src/curobo/cuda_robot_model/usd_kinematics_parser.py +++ b/src/curobo/cuda_robot_model/usd_kinematics_parser.py @@ -71,7 +71,7 @@ class UsdKinematicsParser(KinematicsParser): for l in all_joints: parent, child = get_links_for_joint(l) if child is not None and parent is not None: - self._parent_map[child.GetName()] = parent.GetName() + self._parent_map[child.GetName()] = {"parent": parent.GetName()} self._parent_joint_map[child.GetName()] = l # store joint prim def get_link_parameters(self, link_name: str, base: bool = False) -> LinkParams: @@ -100,7 +100,7 @@ class UsdKinematicsParser(KinematicsParser): joint_type = JointType.FIXED else: - parent_link_name = self._parent_map[link_name] + parent_link_name = self._parent_map[link_name]["parent"] joint_prim = self._parent_joint_map[link_name] # joint prim connects link joint_transform = self._get_joint_transform(joint_prim) joint_axis = None diff --git a/src/curobo/curobolib/cpp/check_cuda.h b/src/curobo/curobolib/cpp/check_cuda.h new file mode 100644 index 0000000..dd48d44 --- /dev/null +++ b/src/curobo/curobolib/cpp/check_cuda.h @@ -0,0 +1,21 @@ +/* + * 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. + */ +#include +#include + + +// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. +#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), # x " must be contiguous") +#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 diff --git a/src/curobo/curobolib/cpp/cuda_precisions.h b/src/curobo/curobolib/cpp/cuda_precisions.h new file mode 100644 index 0000000..5ff8057 --- /dev/null +++ b/src/curobo/curobolib/cpp/cuda_precisions.h @@ -0,0 +1,17 @@ +/* + * 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. + */ +#include "check_cuda.h" + +#include +#include +#if CHECK_FP8 +#include +#endif diff --git a/src/curobo/curobolib/cpp/geom_cuda.cpp b/src/curobo/curobolib/cpp/geom_cuda.cpp index 6ea44ac..2ff5222 100644 --- a/src/curobo/curobolib/cpp/geom_cuda.cpp +++ b/src/curobo/curobolib/cpp/geom_cuda.cpp @@ -13,7 +13,7 @@ #include #include - +#include "check_cuda.h" // CUDA forward declarations std::vectorself_collision_distance( @@ -168,13 +168,6 @@ backward_pose_distance(torch::Tensor out_grad_p, // C++ interface -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) \ - AT_ASSERTM(x.is_contiguous(), # x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) std::vectorself_collision_distance_wrapper( torch::Tensor out_distance, torch::Tensor out_vec, diff --git a/src/curobo/curobolib/cpp/kinematics_fused_cuda.cpp b/src/curobo/curobolib/cpp/kinematics_fused_cuda.cpp index b7dda16..23453ea 100644 --- a/src/curobo/curobolib/cpp/kinematics_fused_cuda.cpp +++ b/src/curobo/curobolib/cpp/kinematics_fused_cuda.cpp @@ -13,6 +13,8 @@ #include #include +#include "check_cuda.h" + // CUDA forward declarations std::vector @@ -33,7 +35,9 @@ std::vectorkin_fused_forward( const torch::Tensor joint_map_type, const torch::Tensor store_link_map, const torch::Tensor link_sphere_map, + const torch::Tensor joint_offset_map, const int batch_size, + const int n_joints, const int n_spheres, const bool use_global_cumul = false); @@ -52,76 +56,20 @@ std::vectorkin_fused_backward_16t( const torch::Tensor store_link_map, const torch::Tensor link_sphere_map, const torch::Tensor link_chain_map, + const torch::Tensor joint_offset_map, const int batch_size, + const int n_joints, const int n_spheres, const bool sparsity_opt = true, const bool use_global_cumul = false); // C++ interface -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) \ - AT_ASSERTM(x.is_contiguous(), # x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -std::vectorkin_forward_wrapper( - torch::Tensor link_pos, torch::Tensor link_quat, - torch::Tensor batch_robot_spheres, torch::Tensor global_cumul_mat, - const torch::Tensor joint_vec, const torch::Tensor fixed_transform, - const torch::Tensor robot_spheres, const torch::Tensor link_map, - const torch::Tensor joint_map, const torch::Tensor joint_map_type, - const torch::Tensor store_link_map, const torch::Tensor link_sphere_map, - const int batch_size, const int n_spheres, - const bool use_global_cumul = false) -{ - const at::cuda::OptionalCUDAGuard guard(joint_vec.device()); - - // TODO: add check input - return kin_fused_forward( - link_pos, link_quat, batch_robot_spheres, global_cumul_mat, joint_vec, - fixed_transform, robot_spheres, link_map, joint_map, joint_map_type, - store_link_map, link_sphere_map, batch_size, n_spheres, use_global_cumul); -} - -std::vectorkin_backward_wrapper( - torch::Tensor grad_out, const torch::Tensor grad_nlinks_pos, - const torch::Tensor grad_nlinks_quat, const torch::Tensor grad_spheres, - const torch::Tensor global_cumul_mat, const torch::Tensor joint_vec, - const torch::Tensor fixed_transform, const torch::Tensor robot_spheres, - const torch::Tensor link_map, const torch::Tensor joint_map, - const torch::Tensor joint_map_type, const torch::Tensor store_link_map, - const torch::Tensor link_sphere_map, const torch::Tensor link_chain_map, - const int batch_size, const int n_spheres, const bool sparsity_opt = true, - const bool use_global_cumul = false) -{ - const at::cuda::OptionalCUDAGuard guard(joint_vec.device()); - - return kin_fused_backward_16t( - grad_out, grad_nlinks_pos, grad_nlinks_quat, grad_spheres, - global_cumul_mat, joint_vec, fixed_transform, robot_spheres, link_map, - joint_map, joint_map_type, store_link_map, link_sphere_map, - link_chain_map, batch_size, n_spheres, sparsity_opt, use_global_cumul); -} - -std::vector -matrix_to_quaternion_wrapper(torch::Tensor out_quat, - const torch::Tensor in_rot // batch_size, 3 - ) -{ - const at::cuda::OptionalCUDAGuard guard(in_rot.device()); - - CHECK_INPUT(in_rot); - CHECK_INPUT(out_quat); - return matrix_to_quaternion(out_quat, in_rot); -} PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("forward", &kin_forward_wrapper, "Kinematics fused forward (CUDA)"); - m.def("backward", &kin_backward_wrapper, "Kinematics fused backward (CUDA)"); - m.def("matrix_to_quaternion", &matrix_to_quaternion_wrapper, + m.def("forward", &kin_fused_forward, "Kinematics fused forward (CUDA)"); + m.def("backward", &kin_fused_backward_16t, "Kinematics fused backward (CUDA)"); + m.def("matrix_to_quaternion", &matrix_to_quaternion, "Rotation Matrix to Quaternion (CUDA)"); } diff --git a/src/curobo/curobolib/cpp/kinematics_fused_kernel.cu b/src/curobo/curobolib/cpp/kinematics_fused_kernel.cu index 4b71f57..f64b9c2 100644 --- a/src/curobo/curobolib/cpp/kinematics_fused_kernel.cu +++ b/src/curobo/curobolib/cpp/kinematics_fused_kernel.cu @@ -10,12 +10,13 @@ */ #include + #include #include #include #include "helper_math.h" -#include +#include "check_cuda.h" #include #define M 4 @@ -28,13 +29,6 @@ #define Y_ROT 4 #define Z_ROT 5 -#define X_PRISM_NEG 6 -#define Y_PRISM_NEG 7 -#define Z_PRISM_NEG 8 -#define X_ROT_NEG 9 -#define Y_ROT_NEG 10 -#define Z_ROT_NEG 11 - #define MAX_BATCH_PER_BLOCK 32 // tunable parameter for improving occupancy #define MAX_BW_BATCH_PER_BLOCK 16 // tunable parameter for improving occupancy @@ -46,6 +40,8 @@ namespace Curobo { namespace Kinematics { + + template __device__ __forceinline__ void scale_cross_sum(float3 a, float3 b, float3 scale, psum_t& sum_out) @@ -293,36 +289,11 @@ namespace Curobo } } - __device__ __forceinline__ void update_joint_type_direction(int& j_type, int8_t& axis_sign) - { - // Assume that input j_type >= 0 . Check fixed joint outside of this function. - - // Don't do anything if j_type < 6. j_type range is [0, 11] - // divergence here. - axis_sign = 1; - - if (j_type >= 6) - { - j_type -= 6; - axis_sign = -1; - } - } - - __device__ __forceinline__ void update_joint_type_direction(int& j_type) - { - // Assume that input j_type >= 0 . Check fixed joint outside of this function. - - // Don't do anything if j_type < 6. j_type range is [0, 11] - // divergence here. - if (j_type >= 6) - { - j_type -= 6; - } - } __device__ __forceinline__ void update_axis_direction( float& angle, - int & j_type) + int & j_type, + const float2 &j_offset) { // Assume that input j_type >= 0 . Check fixed joint outside of this function. // sign should be +ve <= 5 and -ve >5 @@ -330,9 +301,7 @@ namespace Curobo // cuda code treats -1.0 * 0.0 as negative. Hence we subtract 6. If in future, -1.0 * 0.0 = // +ve, // then this code should be j_type - 5. - angle = -1 * copysignf(1.0, j_type - 6) * angle; - - update_joint_type_direction(j_type); + angle = j_offset.x * angle + j_offset.y; } // In the following versions of rot_fn, some non-nan values may become nan as we @@ -462,7 +431,7 @@ namespace Curobo template __device__ __forceinline__ void rot_backward_translation(const float3& vec, float *cumul_mat, float *l_pos, - const float3& loc_grad, psum_t& grad_q, const int8_t axis_sign = 1) + const float3& loc_grad, psum_t& grad_q, const float axis_sign = 1) { float3 e_pos, j_pos; @@ -481,7 +450,7 @@ namespace Curobo rot_backward_rotation(const float3 vec, const float3 grad_vec, psum_t & grad_q, - const int8_t axis_sign = 1) + const float axis_sign = 1) { grad_q += axis_sign * dot(vec, grad_vec); } @@ -489,7 +458,7 @@ namespace Curobo template __device__ __forceinline__ void prism_backward_translation(const float3 vec, const float3 grad_vec, - psum_t& grad_q, const int8_t axis_sign = 1) + psum_t& grad_q, const float axis_sign = 1) { grad_q += axis_sign * dot(vec, grad_vec); } @@ -497,7 +466,7 @@ namespace Curobo template __device__ __forceinline__ void z_rot_backward(float *link_cumul_mat, float *l_pos, float3& loc_grad_position, - float3& loc_grad_orientation, psum_t& grad_q, const int8_t axis_sign = 1) + float3& loc_grad_orientation, psum_t& grad_q, const float axis_sign = 1) { float3 vec = make_float3(link_cumul_mat[2], link_cumul_mat[6], link_cumul_mat[10]); @@ -512,7 +481,7 @@ namespace Curobo template __device__ __forceinline__ void x_rot_backward(float *link_cumul_mat, float *l_pos, float3& loc_grad_position, - float3& loc_grad_orientation, psum_t& grad_q, const int8_t axis_sign = 1) + float3& loc_grad_orientation, psum_t& grad_q, const float axis_sign = 1) { float3 vec = make_float3(link_cumul_mat[0], link_cumul_mat[4], link_cumul_mat[8]); @@ -527,7 +496,7 @@ namespace Curobo template __device__ __forceinline__ void y_rot_backward(float *link_cumul_mat, float *l_pos, float3& loc_grad_position, - float3& loc_grad_orientation, psum_t& grad_q, const int8_t axis_sign = 1) + float3& loc_grad_orientation, psum_t& grad_q, const float axis_sign = 1) { float3 vec = make_float3(link_cumul_mat[1], link_cumul_mat[5], link_cumul_mat[9]); @@ -542,7 +511,7 @@ namespace Curobo template __device__ __forceinline__ void xyz_prism_backward_translation(float *cumul_mat, float3& loc_grad, - psum_t& grad_q, int xyz, const int8_t axis_sign = 1) + psum_t& grad_q, int xyz, const float axis_sign = 1) { prism_backward_translation( make_float3(cumul_mat[0 + xyz], cumul_mat[4 + xyz], cumul_mat[8 + xyz]), @@ -553,7 +522,7 @@ namespace Curobo __device__ __forceinline__ void x_prism_backward_translation(float *cumul_mat, float3 & loc_grad, psum_t & grad_q, - const int8_t axis_sign = 1) + const float axis_sign = 1) { // get rotation vector: prism_backward_translation( @@ -564,7 +533,7 @@ namespace Curobo __device__ __forceinline__ void y_prism_backward_translation(float *cumul_mat, float3 & loc_grad, psum_t & grad_q, - const int8_t axis_sign = 1) + const float axis_sign = 1) { // get rotation vector: prism_backward_translation( @@ -575,16 +544,15 @@ namespace Curobo __device__ __forceinline__ void z_prism_backward_translation(float *cumul_mat, float3 & loc_grad, psum_t & grad_q, - const int8_t axis_sign = 1) + const float axis_sign = 1) { // get rotation vector: prism_backward_translation( make_float3(cumul_mat[2], cumul_mat[6], cumul_mat[10]), loc_grad, grad_q, axis_sign); } - __device__ __forceinline__ void xyz_rot_backward_translation(float *cumul_mat, float *l_pos, float3& loc_grad, - float& grad_q, int xyz, const int8_t axis_sign = 1) + float& grad_q, int xyz, const float axis_sign = 1) { // get rotation vector: rot_backward_translation( @@ -592,10 +560,11 @@ namespace Curobo &cumul_mat[0], &l_pos[0], loc_grad, grad_q, axis_sign); } + template __device__ __forceinline__ void x_rot_backward_translation(float *cumul_mat, float *l_pos, float3& loc_grad, - psum_t& grad_q, const int8_t axis_sign = 1) + psum_t& grad_q, const float axis_sign = 1) { // get rotation vector: rot_backward_translation( @@ -606,7 +575,7 @@ namespace Curobo template __device__ __forceinline__ void y_rot_backward_translation(float *cumul_mat, float *l_pos, float3& loc_grad, - psum_t& grad_q, const int8_t axis_sign = 1) + psum_t& grad_q, const float axis_sign = 1) { // get rotation vector: rot_backward_translation( @@ -617,7 +586,7 @@ namespace Curobo template __device__ __forceinline__ void z_rot_backward_translation(float *cumul_mat, float *l_pos, float3& loc_grad, - psum_t& grad_q, const int8_t axis_sign = 1) + psum_t& grad_q, const float axis_sign = 1) { // get rotation vector: rot_backward_translation( @@ -629,18 +598,19 @@ namespace Curobo // This one should be about 10% faster. template __global__ void - kin_fused_warp_kernel2(scalar_t *link_pos, // batchSize xz store_n_links x M x M - scalar_t *link_quat, // batchSize x store_n_links x M x M + kin_fused_warp_kernel2(float *link_pos, // batchSize xz store_n_links x M x M + float *link_quat, // batchSize x store_n_links x M x M scalar_t *b_robot_spheres, // batchSize x nspheres x M - scalar_t *global_cumul_mat, - const scalar_t *q, // batchSize x njoints - const scalar_t *fixedTransform, // nlinks x M x M - const scalar_t *robot_spheres, // nspheres x M + float *global_cumul_mat, // batchSize x nlinks x M x M + const float *q, // batchSize x njoints + const float *fixedTransform, // nlinks x M x M + const float *robot_spheres, // nspheres x M const int8_t *jointMapType, // nlinks const int16_t *jointMap, // nlinks const int16_t *linkMap, // nlinks const int16_t *storeLinkMap, // store_n_links const int16_t *linkSphereMap, // nspheres + const float *jointOffset, // nlinks const int batchSize, const int nspheres, const int nlinks, const int njoints, const int store_n_links) @@ -688,7 +658,8 @@ namespace Curobo else { float angle = q[batch * njoints + jointMap[l]]; - update_axis_direction(angle, j_type); + float2 angle_offset = *(float2 *)&jointOffset[l*2]; + update_axis_direction(angle, j_type, angle_offset); if (j_type <= Z_PRISM) { @@ -794,20 +765,21 @@ namespace Curobo template __global__ void kin_fused_backward_kernel3( - scalar_t *grad_out_link_q, // batchSize * njoints - const scalar_t *grad_nlinks_pos, // batchSize * store_n_links * 16 - const scalar_t *grad_nlinks_quat, + float *grad_out_link_q, // batchSize * njoints + const float *grad_nlinks_pos, // batchSize * store_n_links * 16 + const float *grad_nlinks_quat, const scalar_t *grad_spheres, // batchSize * nspheres * 4 - const scalar_t *global_cumul_mat, - const scalar_t *q, // batchSize * njoints - const scalar_t *fixedTransform, // nlinks * 16 - const scalar_t *robotSpheres, // batchSize * nspheres * 4 + const float *global_cumul_mat, + const float *q, // batchSize * njoints + const float *fixedTransform, // nlinks * 16 + const float *robotSpheres, // batchSize * nspheres * 4 const int8_t *jointMapType, // nlinks const int16_t *jointMap, // nlinks const int16_t *linkMap, // nlinks const int16_t *storeLinkMap, // store_n_links const int16_t *linkSphereMap, // nspheres const int16_t *linkChainMap, // nlinks*nlinks + const float *jointOffset, // nlinks*2 const int batchSize, const int nspheres, const int nlinks, const int njoints, const int store_n_links) { @@ -819,8 +791,7 @@ namespace Curobo if (batch >= batchSize) return; - - // Each thread computes one element of the cumul_mat. + // Each thread computes one element of the cumul_mat. // first 4 threads compute a row of the output; const int elem_idx = threadIdx.x % 16; const int col_idx = elem_idx % 4; @@ -832,6 +803,7 @@ namespace Curobo for (int l = 0; l < nlinks; l++) { int outAddrStart = matAddrBase + l * M * M; // + (t % M) * M; + cumul_mat[outAddrStart + elem_idx] = global_cumul_mat[batch * nlinks * M * M + l * M * M + elem_idx]; } @@ -857,8 +829,8 @@ namespace Curobo else { float angle = q[batch * njoints + jointMap[l]]; - - update_axis_direction(angle, j_type); + float2 angle_offset = *(float2 *)&jointOffset[l*2]; + update_axis_direction(angle, j_type, angle_offset); if (j_type <= Z_PRISM) { @@ -949,15 +921,11 @@ namespace Curobo { continue; } - int8_t axis_sign = 1; + float axis_sign = jointOffset[j*2]; int j_type = jointMapType[j]; - if (j_type != FIXED) - { - update_joint_type_direction(j_type, axis_sign); - } - + if (j_type == Z_ROT) { float result = 0.0; @@ -1046,12 +1014,13 @@ namespace Curobo // for (int16_t k = joints_per_thread; k >= 0; k--) for (int16_t k = 0; k < joints_per_thread; k++) { - // int16_t j = elem_idx * joints_per_thread + k; - int16_t j = elem_idx + k * 16; + int16_t j = elem_idx * joints_per_thread + k; + //int16_t j = elem_idx + k * 16; + // int16_t j = elem_idx + k * 16; // (threadidx.x % 16) + k * 16 (0 to 16) // int16_t j = k * M + elem_idx; - if ((j > max_lmap) || (j < 0)) - continue; + if ((j > max_lmap)) + break; // This can be spread across threads as they are not sequential? if (linkChainMap[l_map * nlinks + j] == 0.0) @@ -1061,12 +1030,8 @@ namespace Curobo int16_t j_idx = jointMap[j]; int j_type = jointMapType[j]; - int8_t axis_sign = 1; + float axis_sign = jointOffset[j*2]; - if (j_type != FIXED) - { - update_joint_type_direction(j_type, axis_sign); - } // get rotation vector: if (j_type == Z_ROT) @@ -1090,8 +1055,9 @@ namespace Curobo g_position, g_orientation, psum_grad[j_idx], axis_sign); } } - } + } + __syncthreads(); if (PARALLEL_WRITE) { // accumulate the partial sums across the 16 threads @@ -1114,8 +1080,8 @@ namespace Curobo for (int16_t j = 0; j < joints_per_thread; j++) { - // const int16_t j_idx = elem_idx * joints_per_thread + j; - const int16_t j_idx = elem_idx + j * 16; + const int16_t j_idx = elem_idx * joints_per_thread + j; + //const int16_t j_idx = elem_idx + j * 16; if (j_idx >= njoints) { @@ -1151,7 +1117,7 @@ namespace Curobo { { grad_out_link_q[batch * njoints + j] = - psum_grad[j]; // write the sum to memory + (float) psum_grad[j]; // write the sum to memory } } } @@ -1195,11 +1161,27 @@ std::vectorkin_fused_forward( const torch::Tensor robot_spheres, const torch::Tensor link_map, const torch::Tensor joint_map, const torch::Tensor joint_map_type, const torch::Tensor store_link_map, const torch::Tensor link_sphere_map, - const int batch_size, const int n_spheres, + const torch::Tensor joint_offset_map, + const int batch_size, + const int n_joints,const int n_spheres, const bool use_global_cumul = false) { using namespace Curobo::Kinematics; - const int n_joints = joint_vec.size(0) / batch_size; + CHECK_INPUT_GUARD(joint_vec); + CHECK_INPUT(link_pos); + CHECK_INPUT(link_quat); + CHECK_INPUT(global_cumul_mat); + CHECK_INPUT(batch_robot_spheres); + CHECK_INPUT(fixed_transform); + CHECK_INPUT(robot_spheres); + CHECK_INPUT(link_map); + CHECK_INPUT(joint_map); + CHECK_INPUT(joint_map_type); + CHECK_INPUT(store_link_map); + CHECK_INPUT(link_sphere_map); + //CHECK_INPUT(link_chain_map); + CHECK_INPUT(joint_offset_map); + const int n_links = link_map.size(0); const int store_n_links = link_pos.size(1); assert(joint_map.dtype() == torch::kInt16); @@ -1238,38 +1220,42 @@ std::vectorkin_fused_forward( if (use_global_cumul) { AT_DISPATCH_FLOATING_TYPES( - link_pos.scalar_type(), "kin_fused_forward", ([&] { + batch_robot_spheres.scalar_type(), "kin_fused_forward", ([&] { kin_fused_warp_kernel2 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( - link_pos.data_ptr(), link_quat.data_ptr(), + link_pos.data_ptr(), link_quat.data_ptr(), batch_robot_spheres.data_ptr(), - global_cumul_mat.data_ptr(), - joint_vec.data_ptr(), - fixed_transform.data_ptr(), - robot_spheres.data_ptr(), + global_cumul_mat.data_ptr(), + joint_vec.data_ptr(), + fixed_transform.data_ptr(), + robot_spheres.data_ptr(), joint_map_type.data_ptr(), joint_map.data_ptr(), link_map.data_ptr(), store_link_map.data_ptr(), - link_sphere_map.data_ptr(), batch_size, n_spheres, + link_sphere_map.data_ptr(), + joint_offset_map.data_ptr(), + batch_size, n_spheres, n_links, n_joints, store_n_links); })); } else { AT_DISPATCH_FLOATING_TYPES( - link_pos.scalar_type(), "kin_fused_forward", ([&] { + batch_robot_spheres.scalar_type(), "kin_fused_forward", ([&] { kin_fused_warp_kernel2 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( - link_pos.data_ptr(), link_quat.data_ptr(), + link_pos.data_ptr(), link_quat.data_ptr(), batch_robot_spheres.data_ptr(), - global_cumul_mat.data_ptr(), - joint_vec.data_ptr(), - fixed_transform.data_ptr(), - robot_spheres.data_ptr(), + global_cumul_mat.data_ptr(), + joint_vec.data_ptr(), + fixed_transform.data_ptr(), + robot_spheres.data_ptr(), joint_map_type.data_ptr(), joint_map.data_ptr(), link_map.data_ptr(), store_link_map.data_ptr(), - link_sphere_map.data_ptr(), batch_size, n_spheres, + link_sphere_map.data_ptr(), + joint_offset_map.data_ptr(), + batch_size, n_spheres, n_links, n_joints, store_n_links); })); } @@ -1292,15 +1278,29 @@ std::vectorkin_fused_backward_16t( const torch::Tensor link_map, const torch::Tensor joint_map, const torch::Tensor joint_map_type, const torch::Tensor store_link_map, const torch::Tensor link_sphere_map, const torch::Tensor link_chain_map, - const int batch_size, const int n_spheres, const bool sparsity_opt = true, + const torch::Tensor joint_offset_map, + const int batch_size, const int n_joints, const int n_spheres, const bool sparsity_opt = true, const bool use_global_cumul = false) { using namespace Curobo::Kinematics; - - const int n_joints = joint_vec.size(0) / batch_size; + CHECK_INPUT_GUARD(joint_vec); + CHECK_INPUT(grad_out); + CHECK_INPUT(grad_nlinks_pos); + CHECK_INPUT(grad_nlinks_quat); + CHECK_INPUT(global_cumul_mat); + CHECK_INPUT(fixed_transform); + CHECK_INPUT(robot_spheres); + CHECK_INPUT(link_map); + CHECK_INPUT(joint_map); + CHECK_INPUT(joint_map_type); + CHECK_INPUT(store_link_map); + CHECK_INPUT(link_sphere_map); + CHECK_INPUT(link_chain_map); + CHECK_INPUT(joint_offset_map); + const int n_links = link_map.size(0); const int store_n_links = store_link_map.size(0); - + // assert(n_links < 128); assert(n_joints < 128); // for larger num. of joints, change kernel3's // MAX_JOINTS template value. @@ -1341,72 +1341,77 @@ std::vectorkin_fused_backward_16t( cudaStream_t stream = at::cuda::getCurrentCUDAStream(); assert(sparsity_opt); - if (use_global_cumul) { if (n_joints < 16) { AT_DISPATCH_FLOATING_TYPES( - grad_nlinks_pos.scalar_type(), "kin_fused_backward_16t", ([&] { - kin_fused_backward_kernel3 + grad_spheres.scalar_type(), "kin_fused_backward_16t", ([&] { + kin_fused_backward_kernel3 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( - grad_out.data_ptr(), - grad_nlinks_pos.data_ptr(), - grad_nlinks_quat.data_ptr(), + grad_out.data_ptr(), + grad_nlinks_pos.data_ptr(), + grad_nlinks_quat.data_ptr(), grad_spheres.data_ptr(), - global_cumul_mat.data_ptr(), - joint_vec.data_ptr(), - fixed_transform.data_ptr(), - robot_spheres.data_ptr(), + global_cumul_mat.data_ptr(), + joint_vec.data_ptr(), + fixed_transform.data_ptr(), + robot_spheres.data_ptr(), joint_map_type.data_ptr(), joint_map.data_ptr(), link_map.data_ptr(), store_link_map.data_ptr(), link_sphere_map.data_ptr(), - link_chain_map.data_ptr(), batch_size, n_spheres, + link_chain_map.data_ptr(), + joint_offset_map.data_ptr(), + batch_size, n_spheres, n_links, n_joints, store_n_links); })); } else if (n_joints < 64) { AT_DISPATCH_FLOATING_TYPES( - grad_nlinks_pos.scalar_type(), "kin_fused_backward_16t", ([&] { - kin_fused_backward_kernel3 + grad_spheres.scalar_type(), "kin_fused_backward_16t", ([&] { + kin_fused_backward_kernel3 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( - grad_out.data_ptr(), - grad_nlinks_pos.data_ptr(), - grad_nlinks_quat.data_ptr(), + grad_out.data_ptr(), + grad_nlinks_pos.data_ptr(), + grad_nlinks_quat.data_ptr(), grad_spheres.data_ptr(), - global_cumul_mat.data_ptr(), - joint_vec.data_ptr(), - fixed_transform.data_ptr(), - robot_spheres.data_ptr(), + global_cumul_mat.data_ptr(), + joint_vec.data_ptr(), + fixed_transform.data_ptr(), + robot_spheres.data_ptr(), joint_map_type.data_ptr(), joint_map.data_ptr(), link_map.data_ptr(), store_link_map.data_ptr(), link_sphere_map.data_ptr(), - link_chain_map.data_ptr(), batch_size, n_spheres, + link_chain_map.data_ptr(), + joint_offset_map.data_ptr(), + batch_size, n_spheres, n_links, n_joints, store_n_links); })); } else { AT_DISPATCH_FLOATING_TYPES( - grad_nlinks_pos.scalar_type(), "kin_fused_backward_16t", ([&] { - kin_fused_backward_kernel3 + grad_spheres.scalar_type(), "kin_fused_backward_16t", ([&] { + kin_fused_backward_kernel3 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( - grad_out.data_ptr(), - grad_nlinks_pos.data_ptr(), - grad_nlinks_quat.data_ptr(), + grad_out.data_ptr(), + grad_nlinks_pos.data_ptr(), + grad_nlinks_quat.data_ptr(), grad_spheres.data_ptr(), - global_cumul_mat.data_ptr(), - joint_vec.data_ptr(), - fixed_transform.data_ptr(), - robot_spheres.data_ptr(), + global_cumul_mat.data_ptr(), + joint_vec.data_ptr(), + fixed_transform.data_ptr(), + robot_spheres.data_ptr(), joint_map_type.data_ptr(), joint_map.data_ptr(), link_map.data_ptr(), store_link_map.data_ptr(), link_sphere_map.data_ptr(), - link_chain_map.data_ptr(), batch_size, n_spheres, + link_chain_map.data_ptr(), + joint_offset_map.data_ptr(), + batch_size, n_spheres, n_links, n_joints, store_n_links); })); } @@ -1417,22 +1422,25 @@ std::vectorkin_fused_backward_16t( { // AT_DISPATCH_FLOATING_TYPES( - grad_nlinks_pos.scalar_type(), "kin_fused_backward_16t", ([&] { - kin_fused_backward_kernel3 + grad_spheres.scalar_type(), "kin_fused_backward_16t", ([&] { + kin_fused_backward_kernel3 << < blocksPerGrid, threadsPerBlock, sharedMemSize, stream >> > ( - grad_out.data_ptr(), - grad_nlinks_pos.data_ptr(), - grad_nlinks_quat.data_ptr(), + grad_out.data_ptr(), + grad_nlinks_pos.data_ptr(), + grad_nlinks_quat.data_ptr(), grad_spheres.data_ptr(), - global_cumul_mat.data_ptr(), - joint_vec.data_ptr(), - fixed_transform.data_ptr(), - robot_spheres.data_ptr(), + global_cumul_mat.data_ptr(), + joint_vec.data_ptr(), + fixed_transform.data_ptr(), + robot_spheres.data_ptr(), joint_map_type.data_ptr(), - joint_map.data_ptr(), link_map.data_ptr(), + joint_map.data_ptr(), + link_map.data_ptr(), store_link_map.data_ptr(), link_sphere_map.data_ptr(), - link_chain_map.data_ptr(), batch_size, n_spheres, + link_chain_map.data_ptr(), + joint_offset_map.data_ptr(), + batch_size, n_spheres, n_links, n_joints, store_n_links); })); } @@ -1447,7 +1455,8 @@ matrix_to_quaternion(torch::Tensor out_quat, ) { using namespace Curobo::Kinematics; - + CHECK_INPUT(out_quat); + CHECK_INPUT_GUARD(in_rot); // we compute the warp threads based on number of boxes: // TODO: verify this math diff --git a/src/curobo/curobolib/cpp/lbfgs_step_kernel.cu b/src/curobo/curobolib/cpp/lbfgs_step_kernel.cu index 503aaba..be14450 100644 --- a/src/curobo/curobolib/cpp/lbfgs_step_kernel.cu +++ b/src/curobo/curobolib/cpp/lbfgs_step_kernel.cu @@ -807,6 +807,7 @@ lbfgs_cuda_fuse(torch::Tensor step_vec, torch::Tensor rho_buffer, { using namespace Curobo::Optimization; + // call first kernel: //const bool use_experimental = true; const bool use_fixed_m = true; diff --git a/src/curobo/curobolib/cpp/sphere_obb_kernel.cu b/src/curobo/curobolib/cpp/sphere_obb_kernel.cu index 698aa0b..df3da00 100644 --- a/src/curobo/curobolib/cpp/sphere_obb_kernel.cu +++ b/src/curobo/curobolib/cpp/sphere_obb_kernel.cu @@ -15,23 +15,13 @@ #include #include "helper_math.h" -#include -#include + #include #include #include +#include "check_cuda.h" +#include "cuda_precisions.h" -// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. -#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) \ - AT_ASSERTM(x.is_contiguous(), # x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -#if defined(CUDA_VERSION) && CUDA_VERSION >= 11080 && TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 2 -#include -#endif #define M 4 #define VOXEL_DEBUG true #define VOXEL_UNOBSERVED_DISTANCE -1000.0 @@ -63,8 +53,7 @@ namespace Curobo } - #if defined(CUDA_VERSION) && CUDA_VERSION >= 11080 && TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 2 - + #if CHECK_FP8 __device__ __forceinline__ float get_array_value(const at::Float8_e4m3fn *grid_features, const int voxel_idx) { @@ -704,16 +693,21 @@ float4 &sum_pt) int3 &xyz_grid, float &interpolated_distance, int &voxel_idx) - { + { // convert location to index: can use floor to cast to int. // to account for negative values, add 0.5 * bounds. const float3 loc_grid = make_float3(loc_grid_params.x, loc_grid_params.y, loc_grid_params.z); const float3 sphere = make_float3(loc_sphere.x, loc_sphere.y, loc_sphere.z); - // xyz_loc = make_int3(floorf((sphere + 0.5 * loc_grid) / loc_grid_params.w)); - xyz_loc = make_int3(floorf((sphere) / loc_grid_params.w) + (0.5 * loc_grid/loc_grid_params.w)); - xyz_grid = make_int3(floorf(loc_grid / loc_grid_params.w)); + //xyz_loc = make_int3(floorf((sphere + 0.5 * loc_grid) / loc_grid_params.w)); + const float inv_voxel_size = 1.0/loc_grid_params.w; + //xyz_loc = make_int3(sphere * inv_voxel_size) + make_int3(0.5 * loc_grid * inv_voxel_size); + + xyz_loc = make_int3((sphere + 0.5 * loc_grid) * inv_voxel_size); + + //xyz_loc = make_int3(sphere / loc_grid_params.w) + make_int3(floorf(0.5 * loc_grid/loc_grid_params.w)); + xyz_grid = make_int3((loc_grid * inv_voxel_size)) + 1; // find next nearest voxel to current point and then do weighted interpolation: voxel_idx = xyz_loc.x * xyz_grid.y * xyz_grid.z + xyz_loc.y * xyz_grid.z + xyz_loc.z; @@ -732,8 +726,8 @@ float4 &sum_pt) float3 delta = sphere - voxel_origin; - int3 next_loc = make_int3(floorf((make_float3(xyz_loc) + normalize(delta)))); - float ratio = length(delta)/loc_grid_params.w; + int3 next_loc = make_int3(((make_float3(xyz_loc) + normalize(delta)))); + float ratio = length(delta) * inv_voxel_size; int next_voxel_idx = next_loc.x * xyz_grid.y * xyz_grid.z + next_loc.y * xyz_grid.z + next_loc.z; @@ -2752,7 +2746,7 @@ sphere_obb_clpt(const torch::Tensor sphere_position, // batch_size, 3 else { - #if defined(CUDA_VERSION) && CUDA_VERSION >= 11080 && TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 2 + #if CHECK_FP8 const auto fp8_type = torch::kFloat8_e4m3fn; #else const auto fp8_type = torch::kHalf; @@ -3217,7 +3211,7 @@ sphere_voxel_clpt(const torch::Tensor sphere_position, // batch_size, 3 int blocksPerGrid = (bnh_spheres + threadsPerBlock - 1) / threadsPerBlock; - #if defined(CUDA_VERSION) && CUDA_VERSION >= 11080 && TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 2 + #if CHECK_FP8 const auto fp8_type = torch::kFloat8_e4m3fn; #else const auto fp8_type = torch::kHalf; @@ -3352,7 +3346,7 @@ swept_sphere_voxel_clpt(const torch::Tensor sphere_position, // batch_size, 3 int blocksPerGrid = (bnh_spheres + threadsPerBlock - 1) / threadsPerBlock; - #if defined(CUDA_VERSION) && CUDA_VERSION >= 11080 && TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 2 + #if CHECK_FP8 const auto fp8_type = torch::kFloat8_e4m3fn; #else const auto fp8_type = torch::kHalf; diff --git a/src/curobo/curobolib/kinematics.py b/src/curobo/curobolib/kinematics.py index 1fcb696..009afd5 100644 --- a/src/curobo/curobolib/kinematics.py +++ b/src/curobo/curobolib/kinematics.py @@ -43,100 +43,10 @@ def rotation_matrix_to_quaternion(in_mat, out_quat): class KinematicsFusedFunction(Function): - @staticmethod - def _call_cuda( - link_pos_out, - link_quat_out, - robot_sphere_out, - global_cumul_mat, - angle, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - ): - b_shape = link_pos_out.shape - b_size = b_shape[0] - n_spheres = robot_sphere_out.shape[1] - r = kinematics_fused_cu.forward( - link_pos_out, - link_quat_out, - robot_sphere_out, - global_cumul_mat, - angle, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - b_size, - n_spheres, - False, - ) - out_link_pos = r[0] - out_link_quat = r[1] - out_spheres = r[2] - global_cumul_mat = r[3] - return out_link_pos, out_link_quat, out_spheres, global_cumul_mat - - @staticmethod - def _call_backward_cuda( - grad_out, - link_pos_out, - link_quat_out, - robot_sphere_out, - global_cumul_mat, - angle, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - link_chain_map, - sparsity_opt=True, - ): - b_shape = link_pos_out.shape - b_size = b_shape[0] - n_spheres = robot_sphere_out.shape[1] - if grad_out.is_contiguous(): - grad_out = grad_out.view(-1) - else: - grad_out = grad_out.reshape(-1) - r = kinematics_fused_cu.backward( - grad_out, - link_pos_out, - link_quat_out, - robot_sphere_out, - global_cumul_mat, - angle, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - link_chain_map, - b_size, - n_spheres, - sparsity_opt, - False, - ) - out_q = r[0].view(b_size, -1) - - return out_q @staticmethod def forward( ctx, - # link_mats: torch.Tensor, link_pos: torch.Tensor, link_quat: torch.Tensor, b_robot_spheres: torch.tensor, @@ -150,9 +60,17 @@ class KinematicsFusedFunction(Function): store_link_map: torch.Tensor, link_sphere_map: torch.Tensor, link_chain_map: torch.Tensor, + joint_offset_map: torch.Tensor, grad_out: torch.Tensor, + use_global_cumul: bool = True, ): - link_pos, link_quat, b_robot_spheres, global_cumul_mat = KinematicsFusedFunction._call_cuda( + ctx.use_global_cumul = use_global_cumul + b_shape = link_pos.shape + b_size = b_shape[0] + n_spheres = b_robot_spheres.shape[1] + n_joints = joint_seq.shape[-1] + + r = kinematics_fused_cu.forward( link_pos, link_quat, b_robot_spheres, @@ -165,7 +83,17 @@ class KinematicsFusedFunction(Function): joint_map_type, store_link_map, link_sphere_map, + joint_offset_map, # offset_joint_map + b_size, + n_joints, + n_spheres, + True, ) + out_link_pos = r[0] + out_link_quat = r[1] + out_spheres = r[2] + global_cumul_mat = r[3] + ctx.save_for_backward( joint_seq, fixed_transform, @@ -176,11 +104,11 @@ class KinematicsFusedFunction(Function): store_link_map, link_sphere_map, link_chain_map, + joint_offset_map, grad_out, global_cumul_mat, ) - - return link_pos, link_quat, b_robot_spheres + return out_link_pos, out_link_quat, out_spheres @staticmethod def backward(ctx, grad_out_link_pos, grad_out_link_quat, grad_out_spheres): @@ -197,9 +125,11 @@ class KinematicsFusedFunction(Function): store_link_map, link_sphere_map, link_chain_map, + joint_offset_map, grad_out, global_cumul_mat, ) = ctx.saved_tensors + grad_joint = KinematicsFusedFunction._call_backward_cuda( grad_out, grad_out_link_pos, @@ -215,7 +145,9 @@ class KinematicsFusedFunction(Function): store_link_map, link_sphere_map, link_chain_map, + joint_offset_map, True, + use_global_cumul=ctx.use_global_cumul, ) return ( @@ -233,53 +165,10 @@ class KinematicsFusedFunction(Function): None, None, None, + None, + None, ) - -class KinematicsFusedGlobalCumulFunction(Function): - @staticmethod - def _call_cuda( - link_pos_out, - link_quat_out, - robot_sphere_out, - global_cumul_mat, - angle, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - ): - b_shape = link_pos_out.shape - b_size = b_shape[0] - n_spheres = robot_sphere_out.shape[1] - # print(n_spheres) - # print(angle) - r = kinematics_fused_cu.forward( - link_pos_out, - link_quat_out, - robot_sphere_out, - global_cumul_mat, - angle, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - b_size, - n_spheres, - True, - ) - out_link_pos = r[0] - out_link_quat = r[1] - out_spheres = r[2] - global_cumul_mat = r[3] - return out_link_pos, out_link_quat, out_spheres, global_cumul_mat - @staticmethod def _call_backward_cuda( grad_out, @@ -296,18 +185,18 @@ class KinematicsFusedGlobalCumulFunction(Function): store_link_map, link_sphere_map, link_chain_map, + joint_offset_map, sparsity_opt=True, + use_global_cumul=False, ): - b_shape = link_pos_out.shape + b_shape = grad_out.shape b_size = b_shape[0] - # b_size = link_mat_out.shape[0] n_spheres = robot_sphere_out.shape[1] + n_joints = angle.shape[-1] if grad_out.is_contiguous(): grad_out = grad_out.view(-1) else: grad_out = grad_out.reshape(-1) - # create backward tensors? - r = kinematics_fused_cu.backward( grad_out, link_pos_out, @@ -323,126 +212,17 @@ class KinematicsFusedGlobalCumulFunction(Function): store_link_map, link_sphere_map, link_chain_map, + joint_offset_map, # offset_joint_map b_size, + n_joints, n_spheres, sparsity_opt, - True, + use_global_cumul, ) - out_q = r[0] # .view(angle.shape) - out_q = out_q.view(b_size, -1) + out_q = r[0].view(b_size, -1) + return out_q - @staticmethod - def forward( - ctx, - # link_mats: torch.Tensor, - link_pos: torch.Tensor, - link_quat: torch.Tensor, - b_robot_spheres: torch.tensor, - global_cumul_mat: torch.Tensor, - joint_seq: torch.Tensor, - fixed_transform: torch.tensor, - robot_spheres: torch.tensor, - link_map: torch.tensor, - joint_map: torch.Tensor, - joint_map_type: torch.Tensor, - store_link_map: torch.Tensor, - link_sphere_map: torch.Tensor, - link_chain_map: torch.Tensor, - grad_out: torch.Tensor, - ): - if joint_seq.is_contiguous(): - joint_seq = joint_seq.view(-1) - else: - joint_seq = joint_seq.reshape(-1) - ( - link_pos, - link_quat, - b_robot_spheres, - global_cumul_mat, - ) = KinematicsFusedGlobalCumulFunction._call_cuda( - link_pos, - link_quat, - b_robot_spheres, - global_cumul_mat, - joint_seq, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - ) - ctx.save_for_backward( - joint_seq, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - link_chain_map, - grad_out, - global_cumul_mat, - ) - - return link_pos, link_quat, b_robot_spheres - - @staticmethod - def backward(ctx, grad_out_link_pos, grad_out_link_quat, grad_out_spheres): - grad_joint = None - - if ctx.needs_input_grad[4]: - ( - joint_seq, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - link_chain_map, - grad_out, - global_cumul_mat, - ) = ctx.saved_tensors - grad_joint = KinematicsFusedGlobalCumulFunction._call_backward_cuda( - grad_out, - grad_out_link_pos, - grad_out_link_quat, - grad_out_spheres, - global_cumul_mat, - joint_seq, - fixed_transform, - robot_spheres, - link_map, - joint_map, - joint_map_type, - store_link_map, - link_sphere_map, - link_chain_map, - True, - ) - - return ( - None, - None, - None, - None, - grad_joint, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - def get_cuda_kinematics( link_pos_seq, @@ -458,41 +238,28 @@ def get_cuda_kinematics( store_link_map, link_sphere_idx_map, # sphere idx map link_chain_map, + joint_offset_map, grad_out_q, use_global_cumul: bool = True, ): - if use_global_cumul: - link_pos, link_quat, robot_spheres = KinematicsFusedGlobalCumulFunction.apply( - link_pos_seq, - link_quat_seq, - batch_robot_spheres, - global_cumul_mat, - q_in, - fixed_transform, - link_spheres_tensor, - link_map, # tells which link is attached to which link i - joint_map, # tells which joint is attached to a link i - joint_map_type, # joint type - store_link_map, - link_sphere_idx_map, # sphere idx map - link_chain_map, - grad_out_q, - ) - else: - link_pos, link_quat, robot_spheres = KinematicsFusedFunction.apply( - link_pos_seq, - link_quat_seq, - batch_robot_spheres, - global_cumul_mat, - q_in, - fixed_transform, - link_spheres_tensor, - link_map, # tells which link is attached to which link i - joint_map, # tells which joint is attached to a link i - joint_map_type, # joint type - store_link_map, - link_sphere_idx_map, # sphere idx map - link_chain_map, - grad_out_q, - ) + # if not q_in.is_contiguous(): + # q_in = q_in.contiguous() + link_pos, link_quat, robot_spheres = KinematicsFusedFunction.apply( + link_pos_seq, + link_quat_seq, + batch_robot_spheres, + global_cumul_mat, + q_in, + fixed_transform, + link_spheres_tensor, + link_map, # tells which link is attached to which link i + joint_map, # tells which joint is attached to a link i + joint_map_type, # joint type + store_link_map, + link_sphere_idx_map, # sphere idx map + link_chain_map, + joint_offset_map, + grad_out_q, + use_global_cumul, + ) return link_pos, link_quat, robot_spheres diff --git a/src/curobo/geom/sdf/world.py b/src/curobo/geom/sdf/world.py index ae1ec18..03d6550 100644 --- a/src/curobo/geom/sdf/world.py +++ b/src/curobo/geom/sdf/world.py @@ -263,13 +263,16 @@ class WorldCollisionConfig: cache: Optional[Dict[Obstacle, int]] = None n_envs: int = 1 checker_type: CollisionCheckerType = CollisionCheckerType.PRIMITIVE - max_distance: Union[torch.Tensor, float] = 0.01 + max_distance: Union[torch.Tensor, float] = 0.1 + max_esdf_distance: Union[torch.Tensor, float] = 1000.0 def __post_init__(self): if self.world_model is not None and isinstance(self.world_model, list): self.n_envs = len(self.world_model) if isinstance(self.max_distance, float): self.max_distance = self.tensor_args.to_device([self.max_distance]) + if isinstance(self.max_esdf_distance, float): + self.max_esdf_distance = self.tensor_args.to_device([self.max_esdf_distance]) @staticmethod def load_from_dict( diff --git a/src/curobo/geom/sdf/world_blox.py b/src/curobo/geom/sdf/world_blox.py index 4eb6576..3340406 100644 --- a/src/curobo/geom/sdf/world_blox.py +++ b/src/curobo/geom/sdf/world_blox.py @@ -18,7 +18,7 @@ import torch.autograd.profiler as profiler # CuRobo from curobo.geom.sdf.world import CollisionQueryBuffer, WorldCollisionConfig -from curobo.geom.sdf.world_mesh import WorldMeshCollision +from curobo.geom.sdf.world_voxel import WorldVoxelCollision from curobo.geom.types import Cuboid, Mesh, Sphere, SphereFitType, WorldConfig from curobo.types.camera import CameraObservation from curobo.types.math import Pose @@ -33,7 +33,7 @@ except ImportError: from abc import ABC as Mapper -class WorldBloxCollision(WorldMeshCollision): +class WorldBloxCollision(WorldVoxelCollision): """World Collision Representaiton using Nvidia's nvblox library. This class depends on pytorch wrapper for nvblox. @@ -127,6 +127,7 @@ class WorldBloxCollision(WorldMeshCollision): collision_query_buffer: CollisionQueryBuffer, weight, activation_distance, + compute_esdf: bool = False, ): d = self._blox_mapper.query_sphere_sdf_cost( query_spheres, @@ -196,6 +197,7 @@ class WorldBloxCollision(WorldMeshCollision): collision_query_buffer, weight=weight, activation_distance=activation_distance, + compute_esdf=compute_esdf, ) if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and ( @@ -227,6 +229,7 @@ class WorldBloxCollision(WorldMeshCollision): activation_distance: torch.Tensor, env_query_idx=None, return_loss: bool = False, + **kwargs, ): if "blox" not in self.collision_types or not self.collision_types["blox"]: return super().get_sphere_collision( diff --git a/src/curobo/geom/sdf/world_voxel.py b/src/curobo/geom/sdf/world_voxel.py index 4fb6be4..b3238d5 100644 --- a/src/curobo/geom/sdf/world_voxel.py +++ b/src/curobo/geom/sdf/world_voxel.py @@ -48,11 +48,10 @@ class WorldVoxelCollision(WorldMeshCollision): dims = voxel_cache["dims"] voxel_size = voxel_cache["voxel_size"] feature_dtype = voxel_cache["feature_dtype"] - n_voxels = int( - math.floor(dims[0] / voxel_size) - * math.floor(dims[1] / voxel_size) - * math.floor(dims[2] / voxel_size) - ) + grid_shape = VoxelGrid( + "test", pose=[0, 0, 0, 1, 0, 0, 0], dims=dims, voxel_size=voxel_size + ).get_grid_shape()[0] + n_voxels = grid_shape[0] * grid_shape[1] * grid_shape[2] voxel_params = torch.zeros( (self.n_envs, n_layers, 4), @@ -77,6 +76,12 @@ class WorldVoxelCollision(WorldMeshCollision): dtype=feature_dtype, ) + if feature_dtype in [torch.float32, torch.float16, torch.bfloat16]: + voxel_features[:] = -1.0 * self.max_esdf_distance + else: + voxel_features = (voxel_features.to(dtype=torch.float16) - self.max_esdf_distance).to( + dtype=feature_dtype + ) self._voxel_tensor_list = [voxel_params, voxel_pose, voxel_enable, voxel_features] self.collision_types["voxel"] = True self._env_voxel_names = [[None for _ in range(n_layers)] for _ in range(self.n_envs)] @@ -84,14 +89,14 @@ class WorldVoxelCollision(WorldMeshCollision): def load_collision_model( self, world_model: WorldConfig, env_idx=0, fix_cache_reference: bool = False ): - self._load_collision_model_in_cache( + self._load_voxel_collision_model_in_cache( world_model, env_idx, fix_cache_reference=fix_cache_reference ) return super().load_collision_model( world_model, env_idx=env_idx, fix_cache_reference=fix_cache_reference ) - def _load_collision_model_in_cache( + def _load_voxel_collision_model_in_cache( self, world_config: WorldConfig, env_idx: int = 0, fix_cache_reference: bool = False ): """TODO: @@ -396,9 +401,10 @@ class WorldVoxelCollision(WorldMeshCollision): b, h, n, _ = query_sphere.shape # This can be read from collision query buffer use_batch_env = True + env_query_idx_voxel = env_query_idx if env_query_idx is None: use_batch_env = False - env_query_idx = self._env_n_voxels + env_query_idx_voxel = self._env_n_voxels dist = SdfSphereVoxel.apply( query_sphere, collision_query_buffer.voxel_collision_buffer.distance_buffer, @@ -406,13 +412,13 @@ class WorldVoxelCollision(WorldMeshCollision): collision_query_buffer.voxel_collision_buffer.sparsity_index_buffer, weight, activation_distance, - self.max_distance, + self.max_esdf_distance, self._voxel_tensor_list[3], self._voxel_tensor_list[0], self._voxel_tensor_list[1], self._voxel_tensor_list[2], self._env_n_voxels, - env_query_idx, + env_query_idx_voxel, self._voxel_tensor_list[0].shape[1], b, h, @@ -424,12 +430,8 @@ class WorldVoxelCollision(WorldMeshCollision): sum_collisions, compute_esdf, ) - - if ( - "primitive" not in self.collision_types - or not self.collision_types["primitive"] - or "mesh" not in self.collision_types - or not self.collision_types["mesh"] + if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and ( + "mesh" not in self.collision_types or not self.collision_types["mesh"] ): return dist d_prim = super().get_sphere_distance( @@ -443,9 +445,10 @@ class WorldVoxelCollision(WorldMeshCollision): compute_esdf=compute_esdf, ) if compute_esdf: + d_val = torch.maximum(dist.view(d_prim.shape), d_prim) else: - d_val = d_val.view(d_prim.shape) + d_prim + d_val = dist.view(d_prim.shape) + d_prim return d_val @@ -473,9 +476,10 @@ class WorldVoxelCollision(WorldMeshCollision): raise ValueError("cannot return loss for classification, use get_sphere_distance") b, h, n, _ = query_sphere.shape use_batch_env = True + env_query_idx_voxel = env_query_idx if env_query_idx is None: use_batch_env = False - env_query_idx = self._env_n_voxels + env_query_idx_voxel = self._env_n_voxels dist = SdfSphereVoxel.apply( query_sphere, collision_query_buffer.voxel_collision_buffer.distance_buffer, @@ -483,13 +487,13 @@ class WorldVoxelCollision(WorldMeshCollision): collision_query_buffer.voxel_collision_buffer.sparsity_index_buffer, weight, activation_distance, - self.max_distance, + self.max_esdf_distance, self._voxel_tensor_list[3], self._voxel_tensor_list[0], self._voxel_tensor_list[1], self._voxel_tensor_list[2], self._env_n_voxels, - env_query_idx, + env_query_idx_voxel, self._voxel_tensor_list[0].shape[1], b, h, @@ -501,11 +505,8 @@ class WorldVoxelCollision(WorldMeshCollision): True, False, ) - if ( - "primitive" not in self.collision_types - or not self.collision_types["primitive"] - or "mesh" not in self.collision_types - or not self.collision_types["mesh"] + if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and ( + "mesh" not in self.collision_types or not self.collision_types["mesh"] ): return dist d_prim = super().get_sphere_collision( @@ -552,9 +553,10 @@ class WorldVoxelCollision(WorldMeshCollision): ) b, h, n, _ = query_sphere.shape use_batch_env = True + env_query_idx_voxel = env_query_idx if env_query_idx is None: use_batch_env = False - env_query_idx = self._env_n_voxels + env_query_idx_voxel = self._env_n_voxels dist = SdfSweptSphereVoxel.apply( query_sphere, @@ -563,14 +565,14 @@ class WorldVoxelCollision(WorldMeshCollision): collision_query_buffer.voxel_collision_buffer.sparsity_index_buffer, weight, activation_distance, - self.max_distance, + self.max_esdf_distance, speed_dt, self._voxel_tensor_list[3], self._voxel_tensor_list[0], self._voxel_tensor_list[1], self._voxel_tensor_list[2], self._env_n_voxels, - env_query_idx, + env_query_idx_voxel, self._voxel_tensor_list[0].shape[1], b, h, @@ -583,12 +585,8 @@ class WorldVoxelCollision(WorldMeshCollision): return_loss, sum_collisions, ) - - if ( - "primitive" not in self.collision_types - or not self.collision_types["primitive"] - or "mesh" not in self.collision_types - or not self.collision_types["mesh"] + if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and ( + "mesh" not in self.collision_types or not self.collision_types["mesh"] ): return dist d_prim = super().get_swept_sphere_distance( @@ -641,9 +639,10 @@ class WorldVoxelCollision(WorldMeshCollision): b, h, n, _ = query_sphere.shape use_batch_env = True + env_query_idx_voxel = env_query_idx if env_query_idx is None: use_batch_env = False - env_query_idx = self._env_n_voxels + env_query_idx_voxel = self._env_n_voxels dist = SdfSweptSphereVoxel.apply( query_sphere, collision_query_buffer.voxel_collision_buffer.distance_buffer, @@ -651,14 +650,14 @@ class WorldVoxelCollision(WorldMeshCollision): collision_query_buffer.voxel_collision_buffer.sparsity_index_buffer, weight, activation_distance, - self.max_distance, + self.max_esdf_distance, speed_dt, self._voxel_tensor_list[3], self._voxel_tensor_list[0], self._voxel_tensor_list[1], self._voxel_tensor_list[2], self._env_n_voxels, - env_query_idx, + env_query_idx_voxel, self._voxel_tensor_list[0].shape[1], b, h, @@ -671,11 +670,8 @@ class WorldVoxelCollision(WorldMeshCollision): return_loss, True, ) - if ( - "primitive" not in self.collision_types - or not self.collision_types["primitive"] - or "mesh" not in self.collision_types - or not self.collision_types["mesh"] + if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and ( + "mesh" not in self.collision_types or not self.collision_types["mesh"] ): return dist d_prim = super().get_swept_sphere_collision( @@ -695,5 +691,15 @@ class WorldVoxelCollision(WorldMeshCollision): def clear_cache(self): if self._voxel_tensor_list is not None: self._voxel_tensor_list[2][:] = 0 - self._voxel_tensor_list[-1][:] = -1.0 * self.max_distance + if self._voxel_tensor_list[3].dtype in [torch.float32, torch.float16, torch.bfloat16]: + self._voxel_tensor_list[3][:] = -1.0 * self.max_esdf_distance + else: + self._voxel_tensor_list[3][:] = ( + self._voxel_tensor_list[3].to(dtype=torch.float16) * 0.0 + - 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/geom/types.py b/src/curobo/geom/types.py index 3d01099..8a15488 100644 --- a/src/curobo/geom/types.py +++ b/src/curobo/geom/types.py @@ -577,22 +577,24 @@ class VoxelGrid(Obstacle): if self.feature_tensor is not None: self.feature_dtype = self.feature_tensor.dtype - def create_xyzr_tensor( - self, transform_to_origin: bool = False, tensor_args: TensorDeviceType = TensorDeviceType() - ): + def get_grid_shape(self): bounds = self.dims low = [-bounds[0] / 2, -bounds[1] / 2, -bounds[2] / 2] high = [bounds[0] / 2, bounds[1] / 2, bounds[2] / 2] - trange = [h - l for l, h in zip(low, high)] - x = torch.linspace( - low[0], high[0], int(math.floor(trange[0] / self.voxel_size)), device=tensor_args.device - ) - y = torch.linspace( - low[1], high[1], int(math.floor(trange[1] / self.voxel_size)), device=tensor_args.device - ) - z = torch.linspace( - low[2], high[2], int(math.floor(trange[2] / self.voxel_size)), device=tensor_args.device - ) + grid_shape = [ + 1 + int(high[i] / self.voxel_size) - (int(low[i] / self.voxel_size)) + for i in range(len(low)) + ] + return grid_shape, low, high + + def create_xyzr_tensor( + self, transform_to_origin: bool = False, tensor_args: TensorDeviceType = TensorDeviceType() + ): + trange, low, high = self.get_grid_shape() + + x = torch.linspace(low[0], high[0], trange[0], device=tensor_args.device) + y = torch.linspace(low[1], high[1], trange[1], device=tensor_args.device) + z = torch.linspace(low[2], high[2], trange[2], device=tensor_args.device) w, l, h = x.shape[0], y.shape[0], z.shape[0] xyz = ( torch.stack(torch.meshgrid(x, y, z, indexing="ij")).permute((1, 2, 3, 0)).reshape(-1, 3) @@ -603,11 +605,12 @@ class VoxelGrid(Obstacle): xyz = pose.transform_points(xyz.contiguous()) r = torch.zeros_like(xyz[:, 0:1]) + (self.voxel_size * 0.5) xyzr = torch.cat([xyz, r], dim=1) + return xyzr def get_occupied_voxels(self, feature_threshold: Optional[float] = None): if feature_threshold is None: - feature_threshold = -1.0 * self.voxel_size + feature_threshold = -0.5 * self.voxel_size if self.xyzr_tensor is None or self.feature_tensor is None: log_error("Feature tensor or xyzr tensor is empty") xyzr = self.xyzr_tensor.clone() diff --git a/src/curobo/graph/graph_base.py b/src/curobo/graph/graph_base.py index c58acae..72c243e 100644 --- a/src/curobo/graph/graph_base.py +++ b/src/curobo/graph/graph_base.py @@ -144,6 +144,7 @@ class GraphConfig: graph_file: str = "graph.yml", self_collision_check: bool = True, use_cuda_graph: bool = True, + seed: Optional[int] = None, ): graph_data = load_yaml(join_path(get_task_configs_path(), graph_file)) base_config_data = load_yaml(join_path(get_task_configs_path(), base_cfg_file)) @@ -182,6 +183,8 @@ class GraphConfig: arm_base_cg_rollout = ArmBase(cfg_cg) else: arm_base_cg_rollout = arm_base + if seed is not None: + graph_data["graph"]["seed"] = seed graph_cfg = GraphConfig.from_dict( graph_data["graph"], tensor_args, @@ -911,7 +914,7 @@ class GraphPlanBase(GraphConfig): i = self.i if x_set is not None: if x_set.shape[0] == 0: - log_warn("no valid configuration found") + log_info("no valid configuration found") return if connect_mode == "radius": diff --git a/src/curobo/rollout/arm_reacher.py b/src/curobo/rollout/arm_reacher.py index 095e89c..77b970a 100644 --- a/src/curobo/rollout/arm_reacher.py +++ b/src/curobo/rollout/arm_reacher.py @@ -274,7 +274,7 @@ class ArmReacher(ArmBase, ArmReacherConfig): current_fn = self._link_pose_costs[k] if current_fn.enabled: # get link pose - current_pose = link_poses[k] + current_pose = link_poses[k].contiguous() current_pos = current_pose.position current_quat = current_pose.quaternion diff --git a/src/curobo/rollout/rollout_base.py b/src/curobo/rollout/rollout_base.py index 2b40cc1..1a811a6 100644 --- a/src/curobo/rollout/rollout_base.py +++ b/src/curobo/rollout/rollout_base.py @@ -417,6 +417,7 @@ class Goal(Sequence): class RolloutConfig: tensor_args: TensorDeviceType sum_horizon: bool = False + sampler_seed: int = 1312 class RolloutBase: @@ -436,7 +437,7 @@ class RolloutBase: self.tensor_args, up_bounds=self.action_bound_highs, low_bounds=self.action_bound_lows, - seed=1312, + seed=self.sampler_seed, ) @abstractmethod diff --git a/src/curobo/types/math.py b/src/curobo/types/math.py index cd624c9..98908e6 100644 --- a/src/curobo/types/math.py +++ b/src/curobo/types/math.py @@ -261,7 +261,6 @@ class Pose(Sequence): position=clone_if_not_none(self.position), quaternion=clone_if_not_none(self.quaternion), normalize_rotation=False, - # rotation=clone_if_not_none(self.rotation), ) def to( @@ -452,6 +451,13 @@ class Pose(Sequence): def compute_local_pose(self, world_pose: Pose) -> Pose: return self.inverse().multiply(world_pose) + def contiguous(self) -> Pose: + return Pose( + position=self.position.contiguous() if self.position is not None else None, + quaternion=self.quaternion.contiguous() if self.quaternion is not None else None, + normalize_rotation=False, + ) + def quat_multiply(q1, q2, q_res): a_w = q1[..., 0] diff --git a/src/curobo/types/state.py b/src/curobo/types/state.py index 6faf75f..2655e46 100644 --- a/src/curobo/types/state.py +++ b/src/curobo/types/state.py @@ -512,6 +512,26 @@ class JointState(State): def shape(self): return self.position.shape + def index_dof(self, idx: int): + # get index of joint names: + velocity = acceleration = jerk = None + new_index = idx + position = torch.index_select(self.position, -1, new_index) + if self.velocity is not None: + velocity = torch.index_select(self.velocity, -1, new_index) + if self.acceleration is not None: + acceleration = torch.index_select(self.acceleration, -1, new_index) + if self.jerk is not None: + jerk = torch.index_select(self.jerk, -1, new_index) + joint_names = [self.joint_names[x] for x in new_index] + return JointState( + position=position, + joint_names=joint_names, + velocity=velocity, + acceleration=acceleration, + jerk=jerk, + ) + @get_torch_jit_decorator() def jit_js_scale( diff --git a/src/curobo/util/sample_lib.py b/src/curobo/util/sample_lib.py index 44a0ea4..9c2140f 100644 --- a/src/curobo/util/sample_lib.py +++ b/src/curobo/util/sample_lib.py @@ -499,7 +499,7 @@ class HaltonGenerator: def get_samples(self, num_samples, bounded=False): samples = self._get_samples(num_samples) if bounded: - samples = samples * self.range_b + self.low_bounds + samples = bound_samples(samples, self.range_b, self.low_bounds) return samples @profiler.record_function("halton_generator/gaussian_samples") @@ -512,6 +512,12 @@ class HaltonGenerator: return gaussian_halton_samples +@get_torch_jit_decorator() +def bound_samples(samples: torch.Tensor, range_b: torch.Tensor, low_bounds: torch.Tensor): + samples = samples * range_b + low_bounds + return samples + + @get_torch_jit_decorator() def gaussian_transform( uniform_samples: torch.Tensor, proj_mat: torch.Tensor, i_mat: torch.Tensor, std_dev: float diff --git a/src/curobo/util/usd_helper.py b/src/curobo/util/usd_helper.py index d53d396..3767351 100644 --- a/src/curobo/util/usd_helper.py +++ b/src/curobo/util/usd_helper.py @@ -842,7 +842,7 @@ class UsdHelper: usd_helper.add_world_to_stage(world_model, base_frame=base_frame) animation_links = kin_model.kinematics_config.mesh_link_names - animation_poses = kin_model.get_link_poses(q_traj.position, animation_links) + animation_poses = kin_model.get_link_poses(q_traj.position.contiguous(), animation_links) # add offsets for visual mesh: for i, ival in enumerate(offsets): offset_pose = Pose.from_list(ival) diff --git a/src/curobo/wrap/reacher/ik_solver.py b/src/curobo/wrap/reacher/ik_solver.py index 0ad5635..5ea033c 100644 --- a/src/curobo/wrap/reacher/ik_solver.py +++ b/src/curobo/wrap/reacher/ik_solver.py @@ -59,6 +59,7 @@ class IKSolverConfig: sample_rejection_ratio: int = 50 tensor_args: TensorDeviceType = TensorDeviceType() use_cuda_graph: bool = True + seed: int = 1531 @staticmethod @profiler.record_function("ik_solver/load_from_robot_config") @@ -94,6 +95,7 @@ class IKSolverConfig: collision_activation_distance: Optional[float] = None, high_precision: bool = False, project_pose_to_goal_frame: bool = True, + seed: int = 1531, ): if position_threshold <= 0.001: high_precision = True @@ -104,6 +106,11 @@ 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 @@ -247,6 +254,7 @@ class IKSolverConfig: rollout_fn=aux_rollout, tensor_args=tensor_args, use_cuda_graph=use_cuda_graph, + seed=seed, ) return ik_cfg @@ -318,24 +326,21 @@ class IKResult(Sequence): class IKSolver(IKSolverConfig): def __init__(self, config: IKSolverConfig) -> None: super().__init__(**vars(config)) - # self._solve_ self.batch_size = -1 self._num_seeds = self.num_seeds self.init_state = JointState.from_position( self.solver.rollout_fn.retract_state.unsqueeze(0) ) self.dof = self.solver.safety_rollout.d_action - self._col = None # torch.arange(0, 1, device=self.tensor_args.device, dtype=torch.long) + self._col = None - # self.fixed_seeds = self.solver.safety_rollout.sample_random_actions(100 * 200) # create random seeder: self.q_sample_gen = HaltonGenerator( self.dof, self.tensor_args, up_bounds=self.solver.safety_rollout.action_bound_highs, low_bounds=self.solver.safety_rollout.action_bound_lows, - seed=1531, - # store_buffer=1000, + seed=self.seed, ) # load ik nn: @@ -913,21 +918,13 @@ class IKSolver(IKSolverConfig): seed_list = [] if use_nn_seed and self.ik_nn_seeder is not None: num_random_seeds = num_seeds - 1 - in_data = torch.cat( - (pose.position, pose.quaternion), dim=-1 - ) # .to(dtype=torch.float32) - nn_seed = self.ik_nn_seeder(in_data).view( - batch, 1, self.dof - ) # .to(dtype=self.tensor_args.dtype) + in_data = torch.cat((pose.position, pose.quaternion), dim=-1) + nn_seed = self.ik_nn_seeder(in_data).view(batch, 1, self.dof) seed_list.append(nn_seed) - # print("using nn seed") if num_random_seeds > 0: - random_seed = self.q_sample_gen.get_gaussian_samples(num_random_seeds * batch).view( - batch, num_random_seeds, self.dof - ) - - # random_seed = self.fixed_seeds[:num_random_seeds*batch].view(batch, num_random_seeds, - # self.solver.safety_rollout.d_action) + random_seed = self.q_sample_gen.get_samples( + 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) @@ -944,11 +941,16 @@ class IKSolver(IKSolverConfig): metrics = self.rollout_fn.rollout_constraint(q.position.unsqueeze(1)) return metrics - def sample_configs(self, n: int, use_batch_env=False) -> torch.Tensor: + def sample_configs( + self, n: int, use_batch_env=False, sample_from_ik_seeder: bool = False + ) -> torch.Tensor: """ Only works for environment=0 """ - samples = self.rollout_fn.sample_random_actions(n * self.sample_rejection_ratio) + if sample_from_ik_seeder: + samples = self.q_sample_gen.get_samples(n * self.sample_rejection_ratio, bounded=True) + else: + samples = self.rollout_fn.sample_random_actions(n * self.sample_rejection_ratio) metrics = self.rollout_fn.rollout_constraint( samples.unsqueeze(1), use_batch_env=use_batch_env ) diff --git a/src/curobo/wrap/reacher/motion_gen.py b/src/curobo/wrap/reacher/motion_gen.py index 35c6906..e3d33ec 100644 --- a/src/curobo/wrap/reacher/motion_gen.py +++ b/src/curobo/wrap/reacher/motion_gen.py @@ -205,6 +205,8 @@ class MotionGenConfig: jerk_scale: Optional[Union[List[float], float]] = None, optimize_dt: bool = True, project_pose_to_goal_frame: bool = True, + ik_seed: int = 1531, + graph_seed: int = 1531, ): """Helper function to create configuration from robot and world configuration. @@ -436,6 +438,7 @@ class MotionGenConfig: store_debug=store_ik_debug, collision_activation_distance=collision_activation_distance, project_pose_to_goal_frame=project_pose_to_goal_frame, + seed=ik_seed, ) ik_solver = IKSolver(ik_solver_cfg) @@ -448,6 +451,7 @@ class MotionGenConfig: base_config_data, graph_file, use_cuda_graph=use_cuda_graph, + seed=graph_seed, ) graph_cfg.interpolation_dt = interpolation_dt graph_cfg.interpolation_steps = interpolation_steps diff --git a/src/curobo/wrap/reacher/trajopt.py b/src/curobo/wrap/reacher/trajopt.py index 67f1019..1ffbf25 100644 --- a/src/curobo/wrap/reacher/trajopt.py +++ b/src/curobo/wrap/reacher/trajopt.py @@ -172,10 +172,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 - 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 config_data["model"]["horizon"] = traj_tsteps grad_config_data["model"]["horizon"] = traj_tsteps diff --git a/src/curobo/wrap/reacher/types.py b/src/curobo/wrap/reacher/types.py index 7e202a9..0f84a7a 100644 --- a/src/curobo/wrap/reacher/types.py +++ b/src/curobo/wrap/reacher/types.py @@ -108,6 +108,9 @@ class ReacherSolveState: goal_buffer.retract_state = retract_config goal_buffer.goal_state = goal_state goal_buffer.links_goal_pose = link_poses + if goal_buffer.links_goal_pose is not None: + for k in goal_buffer.links_goal_pose.keys(): + goal_buffer.links_goal_pose[k] = goal_buffer.links_goal_pose[k].contiguous() return goal_buffer def update_goal_buffer( @@ -154,7 +157,7 @@ class ReacherSolveState: current_goal_buffer.goal_state.copy_(goal_state) if link_poses is not None: for k in link_poses.keys(): - current_goal_buffer.links_goal_pose[k].copy_(link_poses[k]) + current_goal_buffer.links_goal_pose[k].copy_(link_poses[k].contiguous()) return current_solve_state, current_goal_buffer, update_reference diff --git a/tests/kinematics_parser_test.py b/tests/kinematics_parser_test.py index 78bdb11..005f655 100644 --- a/tests/kinematics_parser_test.py +++ b/tests/kinematics_parser_test.py @@ -73,7 +73,7 @@ def usd_parser(robot_params_all): def urdf_parser(robot_params_all): robot_params = robot_params_all["kinematics"] robot_urdf = join_path(get_assets_path(), robot_params["urdf_path"]) - kinematics_parser = UrdfKinematicsParser(robot_urdf) + kinematics_parser = UrdfKinematicsParser(robot_urdf, build_scene_graph=True) return kinematics_parser @@ -140,3 +140,8 @@ def test_basic_ee_pose(usd_cuda_robot, urdf_cuda_robot, retract_state): p_d, q_d = usd_pose.distance(urdf_pose) assert p_d < 1e-5 assert q_d < 1e-5 + + +def test_urdf_parser(urdf_parser): + assert urdf_parser.root_link == "base_link" + assert len(urdf_parser.get_controlled_joint_names()) == 9 diff --git a/tests/mimic_joint_test.py b/tests/mimic_joint_test.py new file mode 100644 index 0000000..4e3af85 --- /dev/null +++ b/tests/mimic_joint_test.py @@ -0,0 +1,38 @@ +# +# 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. +# + +# CuRobo +from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig +from curobo.types.state import JointState + + +def test_mimic_config(): + cfg = CudaRobotModelConfig.from_robot_yaml_file("simple_mimic_robot.yml", "ee_link") + + # print(cfg.kinematics_config.fixed_transforms) + robot_model = CudaRobotModel(cfg) + q = JointState.from_position(robot_model.retract_config, joint_names=robot_model.joint_names) + q = robot_model.get_full_js(q) + + # print(q) + q_mimic = robot_model.get_mimic_js(q) + assert len(q_mimic) == 3 + + +def test_robotiq_mimic_config(): + cfg = CudaRobotModelConfig.from_robot_yaml_file("ur5e_robotiq_2f_140.yml", "grasp_frame") + + robot_model = CudaRobotModel(cfg) + q = JointState.from_position(robot_model.retract_config, joint_names=robot_model.joint_names) + q = robot_model.get_full_js(q) + + q_mimic = robot_model.get_mimic_js(q) + assert len(q_mimic.joint_names) == 12 diff --git a/tests/multi_pose_test.py b/tests/multi_pose_test.py index 1757031..85168b0 100644 --- a/tests/multi_pose_test.py +++ b/tests/multi_pose_test.py @@ -10,6 +10,7 @@ # # Third Party +import pytest import torch # CuRobo @@ -21,7 +22,11 @@ from curobo.util_file import get_robot_configs_path, get_world_configs_path, joi from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig -def test_multi_pose_franka(): +@pytest.mark.parametrize( + "b_size", + [1, 10, 100], +) +def test_multi_pose_franka(b_size: int): tensor_args = TensorDeviceType() world_file = "collision_cubby.yml" @@ -44,13 +49,14 @@ def test_multi_pose_franka(): tensor_args=tensor_args, ) ik_solver = IKSolver(ik_config) - b_size = 1 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) - result = ik_solver.solve_single(goal, link_poses=link_poses) + result = ik_solver.solve_batch(goal, link_poses=link_poses) success = result.success - assert torch.count_nonzero(success).item() >= 1.0 # we check if atleast 90% are successful + assert ( + torch.count_nonzero(success).item() / b_size >= 0.9 + ) # we check if atleast 90% are successful diff --git a/tests/voxel_collision_test.py b/tests/voxel_collision_test.py index ddf3a3d..e63623e 100644 --- a/tests/voxel_collision_test.py +++ b/tests/voxel_collision_test.py @@ -184,7 +184,7 @@ def test_primitive_voxel_sphere_distance(world_collision): error = torch.abs(d_cuboid.view(-1) - d_voxel.view(-1)) - assert torch.max(error) < voxel_grid.voxel_size + assert torch.max(error) - voxel_grid.voxel_size < 1e-3 @pytest.mark.parametrize( @@ -243,12 +243,11 @@ def test_primitive_voxel_sphere_gradient(world_collision): error = torch.abs(d_cuboid.view(-1) - d_voxel.view(-1)) - assert torch.max(error) < voxel_grid.voxel_size + assert torch.max(error) - voxel_grid.voxel_size < 1e-3 cuboid_gradient = cuboid_collision_buffer.get_gradient_buffer() voxel_gradient = voxel_collision_buffer.get_gradient_buffer() error = torch.linalg.norm(cuboid_gradient - voxel_gradient, dim=-1) - print(cuboid_gradient) - print(voxel_gradient) - assert torch.max(error) < voxel_grid.voxel_size + + assert torch.max(error) - voxel_grid.voxel_size < 1e-3 diff --git a/tests/voxelization_test.py b/tests/voxelization_test.py index 336c94a..c79de05 100644 --- a/tests/voxelization_test.py +++ b/tests/voxelization_test.py @@ -107,12 +107,12 @@ def test_voxels_from_world(world_collision): assert voxels.shape[0] > 4 -@pytest.mark.skip(reason="Not ready yet.") +# @pytest.mark.skip(reason="Not ready yet.") @pytest.mark.parametrize( "world_collision", [ (True, True), - (False, True), + # (False, True), ], indirect=True, ) @@ -156,9 +156,8 @@ def test_esdf_prim_mesh(world_collision, world_collision_primitive): voxel_size = 0.1 esdf = world_collision.get_esdf_in_bounding_box(voxel_size=voxel_size).clone() esdf_prim = world_collision_primitive.get_esdf_in_bounding_box(voxel_size=voxel_size).clone() - print(esdf.voxel_size, esdf_prim.voxel_size) - voxels = esdf.get_occupied_voxels() - voxels_prim = esdf_prim.get_occupied_voxels() + voxels = esdf.get_occupied_voxels(voxel_size) + voxels_prim = esdf_prim.get_occupied_voxels(voxel_size) assert voxels.shape == voxels_prim.shape