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

constrained planning, robot segmentation

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2024-02-22 21:45:47 -08:00
parent 88eac64edc
commit bafdf80c05
102 changed files with 12440 additions and 8112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -201,3 +201,4 @@ nvidia_curobo*/
*tfevents* *tfevents*
*.csv *.csv
docs/*

60
CHANGELOG.md
View File

@@ -9,12 +9,62 @@ without an express license agreement from NVIDIA CORPORATION or
its affiliates is strictly prohibited. its affiliates is strictly prohibited.
--> -->
# Changelog # Changelog
## Latest Commit
### BugFixes ## Version 0.6.3
- refactored wp.index() instances to `[]` to avoid errors in using with warp-lang>=0.11. ### Changes in default behavior
- Fixed bug in gaussian transformation to ensure values are not -1 or +1. - Increased default collision cache to 50 in RobotWorld.
- Changed `CSpaceConfig.position_limit_clip` default to 0 as previous default of 0.01 can make
default start state in examples be out of bounds.
- MotionGen uses parallel_finetune by default. To get previous motion gen behavior, pass
`warmup(parallel_finetune=False)` and `MotionGenPlanConfig(parallel_finetune=False)`.
- MotionGen loads Mesh Collision checker instead of Primitive by default.
### Breaking Changes
- Renamed `copy_if_not_none` to `clone_if_not_none` to be more descriptive. Now `copy_if_not_none`
will try to copy data into reference.
- Renamed `n_envs` in curobo.opt module to avoid confusion between parallel environments and
parallel problems in optimization.
- Added more inputs to pose distance kernels. Check `curobolib/geom.py`.
- Pose cost `run_vec_weight` should now be `[0,0,0,0,0,0]` instead of `[1,1,1,1,1,1]`
### New Features
- Add function to disable and enable collision for specific links in KinematicsTensorConfig.
- Add goal index to reacher results to return index of goal reached when goalset planning.
- Add locked joint state update api in MotionGen class.
- Add goalset warmup padding to handle varied number of goals during goalset planning and also when
calling plan_single after warmup of goalset.
- Add new trajopt config to allow for smooth solutions at slow speeds (`velocity_scale<=0.25`). Also
add error when `velocity_scale<0.1`.
- Add experimental robot image segmentation module to enable robot removal in depth images.
- Add constrained planning mode to motion_gen.
### BugFixes & Misc.
- refactored wp.index() instances to `[]` to avoid errors in future releases of warp.
- Fix bug in gaussian transformation to ensure values are not -1 or +1.
- Fix bug in ik_solver loading ee_link_name from argument.
- Fix bug in batch_goalset planning, where pose cost was selected as GOALSET instead of
BATCH_GOALSET.
- Added package data to also export `.so` files.
- Fixed bug in transforming link visual mesh offset when reading from urdf. - Fixed bug in transforming link visual mesh offset when reading from urdf.
- Fixed bug in MotionGenPlanConfig.clone() that didn't clone the state of parallel_finetune. - Fixed bug in MotionGenPlanConfig.clone() that didn't clone the state of parallel_finetune.
- Increased weighting from 1.0 to 5.0 for optimized_dt in TrajEvaluator to select shorter
trajectories.
- Improved determinism by setting global seed for random in `graph_nx.py`.
- Added option to clear obstacles in WorldPrimitiveCollision.
- Raise error when reference of tensors change in MotionGen, IKSolver, and TrajOpt when cuda graph
is enabled.
- plan_single will get converted to plan_goalset when a plan_goalset was used to initialize cuda
graph.
- plan_goalset will pad for extra goals when called with less number of goal than initial creation.
- Improved API documentation for Optimizer class.
- Improved benchmark timings, now within 15ms of results reported in technical report. Added
numbers to benchmark [webpage](https://curobo.org/source/getting_started/4_benchmarks.html) for
easy reference.
- Set `use_cuda_graph` to `True` as default from `None` in `MotionGenConfig.load_from_robot_config`
### Known Bugs (WIP)
- Examples don't run in Isaac Sim 2023.1.1 due to behavior change in urdf importer.
## Version 0.6.2 ## Version 0.6.2
### New Features ### New Features
- Added support for actuated axis to be negative (i.e., urdf joints with `<axis xyz="0 -1 0"/>` are - Added support for actuated axis to be negative (i.e., urdf joints with `<axis xyz="0 -1 0"/>` are
@@ -35,6 +85,7 @@ Run `benchmark/ik_benchmark.py` to get the latest results.
- Added `external_asset_path` to robot configuration to help in loading urdf and meshes from an - Added `external_asset_path` to robot configuration to help in loading urdf and meshes from an
external directory. external directory.
### BugFixes & Misc. ### BugFixes & Misc.
- Update nvblox wrappers to work with v0.0.5 without segfaults. Significantly improves stability. - Update nvblox wrappers to work with v0.0.5 without segfaults. Significantly improves stability.
- Remove mimic joints in franka panda to maintain compatibility with Isaac Sim 2023.1.0 and 2022.2.1 - Remove mimic joints in franka panda to maintain compatibility with Isaac Sim 2023.1.0 and 2022.2.1
@@ -60,6 +111,7 @@ planning time, 20ms slower on 4090 and 40ms on ORIN MAXN. We will address this s
release. One way to avoid this regression is to set `finetune_dt_scale=1.05` in release. One way to avoid this regression is to set `finetune_dt_scale=1.05` in
`MotionGenConfig.load_from_robot_config()`. `MotionGenConfig.load_from_robot_config()`.
## Version 0.6.1 ## Version 0.6.1
- Added changes to `examples/isaac_sim` to support Isaac Sim 2023.1.0 - Added changes to `examples/isaac_sim` to support Isaac Sim 2023.1.0

1
MANIFEST.in
View File

@@ -16,3 +16,4 @@
# graft <path to files or directories to add to the project> # graft <path to files or directories to add to the project>
graft src/curobo/content graft src/curobo/content
global-include py.typed global-include py.typed
graft src/curobo/curobolib/*.so

43
benchmark/README.md
View File

@@ -12,10 +12,47 @@ This folder contains scripts to run the motion planning benchmarks.
Refer to Benchmarks & Profiling instructions: https://curobo.org/source/getting_started/4_benchmarks.html. Refer to Benchmarks & Profiling instructions: https://curobo.org/source/getting_started/4_benchmarks.html.
### Note
Results in the arxiv paper were obtained from v0.6.0. Results in the arxiv paper were obtained from v0.6.0.
v0.6.2+ has significant changes to improve motion quality with lower motion time, lower path length, higher pose accuracy (<1mm). v0.6.2+ sacrifices 30ms (RTX 4090) of compute time to achieve signficantly better solutions. The new results are yet to be added to the technical report. For now, to get the latest benchmark results follow instructions here: https://curobo.org/source/getting_started/4_benchmarks.html. v0.6.2+ has significant changes to improve motion quality with lower motion time, lower path length, higher pose accuracy (<1mm). v0.6.2+ sacrifices 15ms (RTX 4090) of compute time to achieve signficantly better solutions. The new results are yet to be added to the technical report. For now, to get the latest benchmark results follow instructions here: https://curobo.org/source/getting_started/4_benchmarks.html.
To get results similar to in the technical report, pass `--report_edition` to `curobo_benchmark.py`. To get results similar to in the technical report, pass `--report_edition` to `curobo_benchmark.py`.
# Latest Results (Feb 2024)
Motion Generation on 2600 problems from motion benchmaker and motion policy networks, gives the
following values on a RTX 4090:
| Metric | Value |
|-------------------|--------------------------------------------------------------|
|Success % | 99.84 |
|Plan Time (s) | mean: 0.068 ± 0.158 median:0.042 75%: 0.055 98%: 0.246 |
|Motion Time (s) | mean: 1.169 ± 0.360 median:1.140 75%: 1.381 98%: 2.163 |
|Path Length (rad.) | mean: 3.177 ± 1.072 median:3.261 75%: 3.804 98%: 5.376 |
|Jerk | mean: 97.700 ± 48.630 median:88.993 75%: 126.092 98%: 199.540|
|Position Error (mm)| mean: 0.119 ± 0.341 median:0.027 75%: 0.091 98%: 1.039 |
## Motion Benchmaker (800 problems):
| Metric | Value |
|-------------------|--------------------------------------------------------------|
|Success % | 100 |
|Plan Time (s) | mean: 0.063 ± 0.137 median:0.042 75%: 0.044 98%: 0.206 |
|Motion Time (s) | mean: 1.429 ± 0.330 median:1.392 75%: 1.501 98%: 2.473 |
|Path Length (rad.) | mean: 3.956 ± 0.783 median:3.755 75%: 4.352 98%: 6.041 |
|Jerk | mean: 67.284 ± 27.788 median:61.853 75%: 83.337 98%: 143.118 |
|Position Error (mm)| mean: 0.079 ± 0.139 median:0.032 75%: 0.077 98%: 0.472 |
## Motion Policy Networks (1800 problems):
| Metric | Value |
|-------------------|-----------------------------------------------------------------|
|Success % | 99.77 |
|Plan Time (s) | mean: 0.068 ± 0.117 median:0.042 75%: 0.059 98%: 0.243 |
|Motion Time (s) | mean: 1.051 ± 0.306 median:1.016 75%: 1.226 98%: 1.760 |
|Path Length (rad.) | mean: 2.829 ± 1.000 median:2.837 75%: 3.482 98%: 4.905 |
|Jerk | mean: 110.610 ± 47.586 median:105.271 75%: 141.517 98%: 217.158 |
|Position Error (mm)| mean: 0.122 ± 0.343 median:0.024 75%: 0.095 98%: 1.114 |

134
benchmark/curobo_benchmark.py
View File

@@ -10,12 +10,15 @@
# #
# Standard Library # Standard Library
import argparse
import random
from copy import deepcopy from copy import deepcopy
from typing import Optional from typing import Optional
# Third Party # Third Party
import numpy as np import numpy as np
import torch import torch
from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw
from tqdm import tqdm from tqdm import tqdm
# CuRobo # CuRobo
@@ -26,6 +29,7 @@ from curobo.types.math import Pose
from curobo.types.robot import RobotConfig from curobo.types.robot import RobotConfig
from curobo.types.state import JointState from curobo.types.state import JointState
from curobo.util.logger import setup_curobo_logger from curobo.util.logger import setup_curobo_logger
from curobo.util.metrics import CuroboGroupMetrics, CuroboMetrics
from curobo.util_file import ( from curobo.util_file import (
get_assets_path, get_assets_path,
get_robot_configs_path, get_robot_configs_path,
@@ -36,28 +40,16 @@ from curobo.util_file import (
) )
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
# torch.set_num_threads(8) # set seeds
# torch.use_deterministic_algorithms(True) torch.manual_seed(2)
torch.manual_seed(0)
torch.backends.cudnn.benchmark = True torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = True torch.backends.cudnn.allow_tf32 = False
# torch.backends.cuda.matmul.allow_tf32 = False np.random.seed(2)
# torch.backends.cudnn.allow_tf32 = False random.seed(2)
# torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True
np.random.seed(0)
# Standard Library
import argparse
import warnings
from typing import List, Optional
# Third Party
from metrics import CuroboGroupMetrics, CuroboMetrics
from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw
def plot_cost_iteration(cost: torch.Tensor, save_path="cost", title="", log_scale=False): def plot_cost_iteration(cost: torch.Tensor, save_path="cost", title="", log_scale=False):
@@ -174,7 +166,7 @@ def load_curobo(
world_cfg = WorldConfig.from_dict( world_cfg = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml")) load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
).get_obb_world() ).get_obb_world()
interpolation_steps = 2000 interpolation_steps = 1000
c_checker = CollisionCheckerType.PRIMITIVE c_checker = CollisionCheckerType.PRIMITIVE
c_cache = {"obb": n_cubes} c_cache = {"obb": n_cubes}
if mesh_mode: if mesh_mode:
@@ -189,7 +181,7 @@ def load_curobo(
K = robot_cfg_instance.kinematics.kinematics_config.joint_limits K = robot_cfg_instance.kinematics.kinematics_config.joint_limits
K.position[0, :] -= 0.2 K.position[0, :] -= 0.2
K.position[1, :] += 0.2 K.position[1, :] += 0.2
finetune_iters = None finetune_iters = 300
grad_iters = None grad_iters = None
if args.report_edition: if args.report_edition:
finetune_iters = 200 finetune_iters = 200
@@ -215,7 +207,7 @@ def load_curobo(
collision_activation_distance=collision_activation_distance, collision_activation_distance=collision_activation_distance,
trajopt_dt=0.25, trajopt_dt=0.25,
finetune_dt_scale=finetune_dt_scale, finetune_dt_scale=finetune_dt_scale,
maximum_trajectory_dt=0.1, maximum_trajectory_dt=0.16,
) )
mg = MotionGen(motion_gen_config) mg = MotionGen(motion_gen_config)
mg.warmup(enable_graph=True, warmup_js_trajopt=False, parallel_finetune=parallel_finetune) mg.warmup(enable_graph=True, warmup_js_trajopt=False, parallel_finetune=parallel_finetune)
@@ -238,7 +230,7 @@ def benchmark_mb(
interpolation_dt = 0.02 interpolation_dt = 0.02
# mpinets_data = True # mpinets_data = True
# if mpinets_data: # if mpinets_data:
file_paths = [motion_benchmaker_raw, mpinets_raw][:] # [1:] file_paths = [motion_benchmaker_raw, mpinets_raw][:]
if args.demo: if args.demo:
file_paths = [demo_raw] file_paths = [demo_raw]
@@ -269,13 +261,13 @@ def benchmark_mb(
if "cage_panda" in key: if "cage_panda" in key:
trajopt_seeds = 16 trajopt_seeds = 16
finetune_dt_scale = 0.95 finetune_dt_scale = 0.95
collision_activation_distance = 0.01
parallel_finetune = True parallel_finetune = True
if "table_under_pick_panda" in key: if "table_under_pick_panda" in key:
tsteps = 44 tsteps = 40
trajopt_seeds = 24 trajopt_seeds = 24
finetune_dt_scale = 0.98 finetune_dt_scale = 0.95
parallel_finetune = True parallel_finetune = True
if "table_pick_panda" in key: if "table_pick_panda" in key:
collision_activation_distance = 0.005 collision_activation_distance = 0.005
@@ -297,15 +289,19 @@ def benchmark_mb(
"cubby_neutral_start", "cubby_neutral_start",
"cubby_neutral_goal", "cubby_neutral_goal",
"dresser_neutral_start", "dresser_neutral_start",
"dresser_neutral_goal",
"tabletop_task_oriented", "tabletop_task_oriented",
]: ]:
collision_activation_distance = 0.005 collision_activation_distance = 0.005
if key in ["dresser_neutral_goal"]: if key in ["dresser_neutral_goal"]:
trajopt_seeds = 24 # 24 trajopt_seeds = 24 # 24
tsteps = 36 # tsteps = 36
collision_activation_distance = 0.005 collision_activation_distance = 0.01
scene_problems = problems[key] # trajopt_seeds = 12
scene_problems = problems[key] # [:10]
n_cubes = check_problems(scene_problems) n_cubes = check_problems(scene_problems)
# print(n_cubes)
# continue
mg, robot_cfg = load_curobo( mg, robot_cfg = load_curobo(
n_cubes, n_cubes,
enable_debug, enable_debug,
@@ -327,11 +323,10 @@ def benchmark_mb(
for problem in tqdm(scene_problems, leave=False): for problem in tqdm(scene_problems, leave=False):
i += 1 i += 1
if problem["collision_buffer_ik"] < 0.0: if problem["collision_buffer_ik"] < 0.0:
# print("collision_ik:", problem["collision_buffer_ik"])
continue continue
plan_config = MotionGenPlanConfig( plan_config = MotionGenPlanConfig(
max_attempts=100, # 100, # 00, # 00, # 100, # 00, # 000,#,00,#00, # 5000, max_attempts=100,
enable_graph_attempt=1, enable_graph_attempt=1,
disable_graph_attempt=20, disable_graph_attempt=20,
enable_finetune_trajopt=True, enable_finetune_trajopt=True,
@@ -349,24 +344,13 @@ def benchmark_mb(
problem_name = "d_" + key + "_" + str(i) problem_name = "d_" + key + "_" + str(i)
# reset planner # reset planner
mg.reset(reset_seed=False) mg.reset(reset_seed=True)
if args.mesh: if args.mesh:
world = WorldConfig.from_dict(deepcopy(problem["obstacles"])).get_mesh_world() world = WorldConfig.from_dict(deepcopy(problem["obstacles"])).get_mesh_world()
else: else:
world = WorldConfig.from_dict(deepcopy(problem["obstacles"])).get_obb_world() world = WorldConfig.from_dict(deepcopy(problem["obstacles"])).get_obb_world()
mg.world_coll_checker.clear_cache() mg.world_coll_checker.clear_cache()
mg.update_world(world) mg.update_world(world)
# from curobo.geom.types import Cuboid as curobo_Cuboid
# coll_mesh = mg.world_coll_checker.get_mesh_in_bounding_box(
# curobo_Cuboid(name="test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1, 1, 1]),
# voxel_size=0.01,
# )
#
# coll_mesh.save_as_mesh(problem_name + "debug_curobo.obj")
# exit()
# continue
# load obstacles
# run planner # run planner
start_state = JointState.from_position(mg.tensor_args.to_device([q_start])) start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
@@ -379,7 +363,6 @@ def benchmark_mb(
) )
if result.status == "IK Fail": if result.status == "IK Fail":
ik_fail += 1 ik_fail += 1
# rint(plan_config.enable_graph, plan_config.enable_graph_attempt)
problem["solution"] = None problem["solution"] = None
problem_name = key + "_" + str(i) problem_name = key + "_" + str(i)
@@ -432,7 +415,6 @@ def benchmark_mb(
log_scale=False, log_scale=False,
) )
if result.success.item(): if result.success.item():
# print("GT: ", result.graph_time)
q_traj = result.get_interpolated_plan() q_traj = result.get_interpolated_plan()
problem["goal_ik"] = q_traj.position.cpu().squeeze().numpy()[-1, :].tolist() problem["goal_ik"] = q_traj.position.cpu().squeeze().numpy()[-1, :].tolist()
problem["solution"] = { problem["solution"] = {
@@ -458,8 +440,7 @@ def benchmark_mb(
.tolist(), .tolist(),
"dt": interpolation_dt, "dt": interpolation_dt,
} }
# print(problem["solution"]["position"])
# exit()
debug = { debug = {
"used_graph": result.used_graph, "used_graph": result.used_graph,
"attempts": result.attempts, "attempts": result.attempts,
@@ -487,14 +468,7 @@ def benchmark_mb(
"valid_query": result.valid_query, "valid_query": result.valid_query,
} }
problem["solution_debug"] = debug problem["solution_debug"] = debug
# print(
# "T: ",
# result.motion_time.item(),
# result.optimized_dt.item(),
# (len(problem["solution"]["position"]) - 1 ) * result.interpolation_dt,
# result.interpolation_dt,
# )
# exit()
reached_pose = mg.compute_kinematics(result.optimized_plan[-1]).ee_pose reached_pose = mg.compute_kinematics(result.optimized_plan[-1]).ee_pose
rot_error = goal_pose.angular_distance(reached_pose) * 100.0 rot_error = goal_pose.angular_distance(reached_pose) * 100.0
if args.graph: if args.graph:
@@ -526,6 +500,7 @@ def benchmark_mb(
motion_time=result.motion_time.item(), motion_time=result.motion_time.item(),
solve_time=solve_time, solve_time=solve_time,
cspace_path_length=path_length, cspace_path_length=path_length,
jerk=torch.max(torch.abs(result.optimized_plan.jerk)).item(),
) )
if write_usd: if write_usd:
@@ -552,20 +527,9 @@ def benchmark_mb(
plot_traj( plot_traj(
result.optimized_plan, result.optimized_plan,
result.optimized_dt.item(), result.optimized_dt.item(),
# result.get_interpolated_plan(),
# result.interpolation_dt,
title=problem_name, title=problem_name,
save_path=join_path("benchmark/log/plot/", problem_name + ".png"), save_path=join_path("benchmark/log/plot/", problem_name + ".png"),
) )
# plot_traj(
# # result.optimized_plan,
# # result.optimized_dt.item(),
# result.get_interpolated_plan(),
# result.interpolation_dt,
# title=problem_name,
# save_path=join_path("benchmark/log/plot/", problem_name + "_int.pdf"),
# )
# exit()
m_list.append(current_metrics) m_list.append(current_metrics)
all_groups.append(current_metrics) all_groups.append(current_metrics)
@@ -587,7 +551,6 @@ def benchmark_mb(
m_list.append(current_metrics) m_list.append(current_metrics)
all_groups.append(current_metrics) all_groups.append(current_metrics)
else: else:
# print("invalid: " + problem_name)
debug = { debug = {
"used_graph": result.used_graph, "used_graph": result.used_graph,
"attempts": result.attempts, "attempts": result.attempts,
@@ -601,28 +564,7 @@ def benchmark_mb(
} }
problem["solution_debug"] = debug problem["solution_debug"] = debug
if False:
world.save_world_as_mesh(problem_name + ".obj")
q_traj = start_state # .unsqueeze(0)
# CuRobo
from curobo.util.usd_helper import UsdHelper
UsdHelper.write_trajectory_animation_with_robot_usd(
robot_cfg,
world,
start_state,
q_traj,
dt=result.interpolation_dt,
save_path=join_path("benchmark/log/usd/", problem_name) + ".usd",
interpolation_steps=1,
write_robot_usd_path="benchmark/log/usd/assets/",
robot_usd_local_reference="assets/",
base_frame="/world_" + problem_name,
visualize_robot_spheres=True,
)
if save_log and not result.success.item(): if save_log and not result.success.item():
# print("save log")
UsdHelper.write_motion_gen_log( UsdHelper.write_motion_gen_log(
result, result,
robot_cfg, robot_cfg,
@@ -637,8 +579,7 @@ def benchmark_mb(
write_robot_usd_path="benchmark/log/usd/assets/", write_robot_usd_path="benchmark/log/usd/assets/",
# flatten_usd=True, # flatten_usd=True,
) )
# print(result.status) print(result.status)
# exit()
g_m = CuroboGroupMetrics.from_list(m_list) g_m = CuroboGroupMetrics.from_list(m_list)
if not args.kpi: if not args.kpi:
@@ -653,10 +594,6 @@ def benchmark_mb(
g_m.motion_time.percent_98, g_m.motion_time.percent_98,
) )
print(g_m.attempts) print(g_m.attempts)
# print("MT: ", g_m.motion_time)
# $print(ik_fail)
# exit()
# print(g_m.position_error, g_m.orientation_error)
g_m = CuroboGroupMetrics.from_list(all_groups) g_m = CuroboGroupMetrics.from_list(all_groups)
if not args.kpi: if not args.kpi:
@@ -668,12 +605,8 @@ def benchmark_mb(
g_m.time.percent_75, g_m.time.percent_75,
g_m.position_error.percent_75, g_m.position_error.percent_75,
g_m.orientation_error.percent_75, g_m.orientation_error.percent_75,
) # g_m.time, g_m.attempts) )
# print("MT: ", g_m.motion_time)
# print(g_m.position_error, g_m.orientation_error)
# exit()
if write_benchmark: if write_benchmark:
if not mpinets_data: if not mpinets_data:
write_yaml(problems, args.file_name + "_mb_solution.yaml") write_yaml(problems, args.file_name + "_mb_solution.yaml")
@@ -681,7 +614,6 @@ def benchmark_mb(
write_yaml(problems, args.file_name + "_mpinets_solution.yaml") write_yaml(problems, args.file_name + "_mpinets_solution.yaml")
all_files += all_groups all_files += all_groups
g_m = CuroboGroupMetrics.from_list(all_files) g_m = CuroboGroupMetrics.from_list(all_files)
# print(g_m.success, g_m.time, g_m.attempts, g_m.position_error, g_m.orientation_error)
print("######## FULL SET ############") print("######## FULL SET ############")
print("All: ", f"{g_m.success:2.2f}") print("All: ", f"{g_m.success:2.2f}")
print("MT: ", g_m.motion_time) print("MT: ", g_m.motion_time)
@@ -690,6 +622,7 @@ def benchmark_mb(
print("ST: ", g_m.solve_time) print("ST: ", g_m.solve_time)
print("position error (mm): ", g_m.position_error) print("position error (mm): ", g_m.position_error)
print("orientation error(%): ", g_m.orientation_error) print("orientation error(%): ", g_m.orientation_error)
print("jerk: ", g_m.jerk)
if args.kpi: if args.kpi:
kpi_data = { kpi_data = {
@@ -702,8 +635,6 @@ def benchmark_mb(
} }
write_yaml(kpi_data, join_path(args.save_path, args.file_name + ".yml")) write_yaml(kpi_data, join_path(args.save_path, args.file_name + ".yml"))
# run on mb dataset:
def check_problems(all_problems): def check_problems(all_problems):
n_cube = 0 n_cube = 0
@@ -801,6 +732,7 @@ if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
setup_curobo_logger("error") setup_curobo_logger("error")
for _ in range(1):
benchmark_mb( benchmark_mb(
save_log=False, save_log=False,
write_usd=args.save_usd, write_usd=args.save_usd,

198
benchmark/curobo_nvblox_benchmark.py
View File

@@ -18,7 +18,6 @@ from typing import Optional
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import numpy as np import numpy as np
import torch import torch
from metrics import CuroboGroupMetrics, CuroboMetrics
from nvblox_torch.datasets.mesh_dataset import MeshDataset from nvblox_torch.datasets.mesh_dataset import MeshDataset
from nvblox_torch.datasets.sun3d_dataset import Sun3dDataset from nvblox_torch.datasets.sun3d_dataset import Sun3dDataset
from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw from robometrics.datasets import demo_raw, motion_benchmaker_raw, mpinets_raw
@@ -34,6 +33,7 @@ from curobo.types.math import Pose
from curobo.types.robot import RobotConfig from curobo.types.robot import RobotConfig
from curobo.types.state import JointState from curobo.types.state import JointState
from curobo.util.logger import setup_curobo_logger from curobo.util.logger import setup_curobo_logger
from curobo.util.metrics import CuroboGroupMetrics, CuroboMetrics
from curobo.util_file import ( from curobo.util_file import (
get_assets_path, get_assets_path,
get_robot_configs_path, get_robot_configs_path,
@@ -126,11 +126,12 @@ def load_curobo(
mpinets: bool = False, mpinets: bool = False,
graph_mode: bool = False, graph_mode: bool = False,
cuda_graph: bool = True, cuda_graph: bool = True,
collision_activation_distance: float = 0.025,
): ):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"] robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
robot_cfg["kinematics"]["collision_sphere_buffer"] = -0.0 robot_cfg["kinematics"]["collision_sphere_buffer"] = 0.0
ik_seeds = 30 # 500 ik_seeds = 30
if graph_mode: if graph_mode:
trajopt_seeds = 4 trajopt_seeds = 4
if trajopt_seeds >= 14: if trajopt_seeds >= 14:
@@ -146,7 +147,7 @@ def load_curobo(
"world": { "world": {
"pose": [0, 0, 0, 1, 0, 0, 0], "pose": [0, 0, 0, 1, 0, 0, 0],
"integrator_type": "tsdf", "integrator_type": "tsdf",
"voxel_size": 0.014, "voxel_size": 0.01,
} }
} }
} }
@@ -174,16 +175,17 @@ def load_curobo(
store_ik_debug=enable_debug, store_ik_debug=enable_debug,
store_trajopt_debug=enable_debug, store_trajopt_debug=enable_debug,
interpolation_steps=interpolation_steps, interpolation_steps=interpolation_steps,
collision_activation_distance=0.01, collision_activation_distance=collision_activation_distance,
trajopt_dt=0.25, trajopt_dt=0.25,
finetune_dt_scale=1.0, finetune_dt_scale=1.0,
maximum_trajectory_dt=0.1, maximum_trajectory_dt=0.1,
finetune_trajopt_iters=300,
) )
mg = MotionGen(motion_gen_config) mg = MotionGen(motion_gen_config)
mg.warmup(enable_graph=True, warmup_js_trajopt=False) mg.warmup(enable_graph=True, warmup_js_trajopt=False, parallel_finetune=True)
# create a ground truth collision checker: # create a ground truth collision checker:
config = RobotWorldConfig.load_from_config( config = RobotWorldConfig.load_from_config(
robot_cfg, robot_cfg_instance,
"collision_table.yml", "collision_table.yml",
collision_activation_distance=0.0, collision_activation_distance=0.0,
collision_checker_type=CollisionCheckerType.PRIMITIVE, collision_checker_type=CollisionCheckerType.PRIMITIVE,
@@ -206,7 +208,7 @@ def benchmark_mb(
# load dataset: # load dataset:
graph_mode = args.graph graph_mode = args.graph
interpolation_dt = 0.02 interpolation_dt = 0.02
file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][1:] file_paths = [demo_raw, motion_benchmaker_raw, mpinets_raw][2:]
enable_debug = save_log or plot_cost enable_debug = save_log or plot_cost
all_files = [] all_files = []
@@ -215,43 +217,73 @@ def benchmark_mb(
og_tsteps = override_tsteps og_tsteps = override_tsteps
og_trajopt_seeds = 12 og_trajopt_seeds = 12
og_collision_activation_distance = 0.03 # 0.03
if args.graph: if args.graph:
og_trajopt_seeds = 4 og_trajopt_seeds = 4
for file_path in file_paths: for file_path in file_paths:
all_groups = [] all_groups = []
mpinets_data = False mpinets_data = False
problems = file_path() problems = file_path()
if "dresser_task_oriented" in list(problems.keys()):
mpinets_data = True
mg, robot_cfg, robot_world = load_curobo(
1,
enable_debug,
og_tsteps,
og_trajopt_seeds,
mpinets_data,
graph_mode,
not args.disable_cuda_graph,
)
for key, v in tqdm(problems.items()): for key, v in tqdm(problems.items()):
scene_problems = problems[key] scene_problems = problems[key]
m_list = [] m_list = []
i = 0 i = -1
ik_fail = 0 ik_fail = 0
trajopt_seeds = og_trajopt_seeds
tsteps = og_tsteps
collision_activation_distance = og_collision_activation_distance
if "dresser_task_oriented" in list(problems.keys()):
mpinets_data = True
if "cage_panda" in key:
trajopt_seeds = 16
finetune_dt_scale = 0.95
if "table_under_pick_panda" in key:
tsteps = 44
trajopt_seeds = 16
finetune_dt_scale = 0.98
if "cubby_task_oriented" in key and "merged" not in key:
trajopt_seeds = 24
finetune_dt_scale = 0.95
collision_activation_distance = 0.035
if "dresser_task_oriented" in key:
trajopt_seeds = 24
finetune_dt_scale = 0.95
collision_activation_distance = 0.035 # 0.035
if "merged_cubby_task_oriented" in key:
collision_activation_distance = 0.025 # 0.035
if "tabletop_task_oriented" in key:
collision_activation_distance = 0.025 # 0.035
if key in ["dresser_neutral_goal"]:
trajopt_seeds = 24
collision_activation_distance = og_collision_activation_distance
mg, robot_cfg, robot_world = load_curobo(
1,
enable_debug,
tsteps,
trajopt_seeds,
mpinets_data,
graph_mode,
not args.disable_cuda_graph,
collision_activation_distance=collision_activation_distance,
)
for problem in tqdm(scene_problems, leave=False): for problem in tqdm(scene_problems, leave=False):
i += 1 i += 1
if problem["collision_buffer_ik"] < 0.0:
continue
plan_config = MotionGenPlanConfig( plan_config = MotionGenPlanConfig(
max_attempts=10, # 00, # 00, # 100, # 00, # 000,#,00,#00, # 5000, max_attempts=10,
enable_graph_attempt=3, enable_graph_attempt=1,
disable_graph_attempt=20,
enable_finetune_trajopt=True, enable_finetune_trajopt=True,
partial_ik_opt=False, partial_ik_opt=False,
enable_graph=graph_mode, enable_graph=graph_mode,
timeout=60, timeout=60,
enable_opt=not graph_mode, enable_opt=not graph_mode,
parallel_finetune=True,
) )
q_start = problem["start"] q_start = problem["start"]
@@ -265,23 +297,15 @@ def benchmark_mb(
mg.reset(reset_seed=False) mg.reset(reset_seed=False)
world = WorldConfig.from_dict(problem["obstacles"]) world = WorldConfig.from_dict(problem["obstacles"])
# .get_mesh_world( mg.world_coll_checker.clear_cache()
# # merge_meshes=True
# )
# mesh = world.mesh[0].get_trimesh_mesh()
# world.save_world_as_mesh(problem_name + ".stl")
mg.world_coll_checker.update_blox_hashes() mg.world_coll_checker.update_blox_hashes()
mg.world_coll_checker.clear_cache()
save_path = "benchmark/log/nvblox/" + key + "_" + str(i) save_path = "benchmark/log/nvblox/" + key + "_" + str(i)
m_dataset = Sun3dDataset(save_path) m_dataset = Sun3dDataset(save_path)
# m_dataset = MeshDataset(
# None, n_frames=100, image_size=640, save_data_dir=None, trimesh_mesh=mesh
# )
tensor_args = mg.tensor_args tensor_args = mg.tensor_args
for j in tqdm(range(len(m_dataset)), leave=False): for j in tqdm(range(min(1, len(m_dataset))), leave=False):
data = m_dataset[j] data = m_dataset[j]
cam_obs = CameraObservation( cam_obs = CameraObservation(
rgb_image=tensor_args.to_device(data["rgba"]) rgb_image=tensor_args.to_device(data["rgba"])
@@ -300,35 +324,12 @@ def benchmark_mb(
torch.cuda.synchronize() torch.cuda.synchronize()
mg.world_coll_checker.update_blox_hashes() mg.world_coll_checker.update_blox_hashes()
torch.cuda.synchronize() torch.cuda.synchronize()
# if i > 2:
# mg.world_coll_checker.clear_cache()
# mg.world_coll_checker.update_blox_hashes()
# mg.world_coll_checker.save_layer("world", "test.nvblx") # mg.world_coll_checker.save_layer("world", "test.nvblx")
if save_log or write_usd:
world.randomize_color(r=[0.5, 0.9], g=[0.2, 0.5], b=[0.0, 0.2])
nvblox_obs = mg.world_coll_checker.get_mesh_from_blox_layer(
"world",
)
# nvblox_obs.vertex_colors = None
if nvblox_obs.vertex_colors is not None:
nvblox_obs.vertex_colors = nvblox_obs.vertex_colors.cpu().numpy()
else:
nvblox_obs.color = [0.0, 0.0, 0.8, 0.8]
nvblox_obs.name = "nvblox_mesh_world"
world.add_obstacle(nvblox_obs)
coll_mesh = mg.world_coll_checker.get_mesh_in_bounding_box(
curobo_Cuboid(name="test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1.5, 1.5, 1]),
voxel_size=0.005,
)
coll_mesh.color = [0.0, 0.8, 0.8, 0.8]
coll_mesh.name = "nvblox_voxel_world"
world.add_obstacle(coll_mesh)
# exit()
# run planner
start_state = JointState.from_position(mg.tensor_args.to_device([q_start])) start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
result = mg.plan_single( result = mg.plan_single(
start_state, start_state,
@@ -338,20 +339,6 @@ def benchmark_mb(
if result.status == "IK Fail": if result.status == "IK Fail":
ik_fail += 1 ik_fail += 1
problem["solution"] = None problem["solution"] = None
if write_usd or save_log:
# CuRobo
from curobo.util.usd_helper import UsdHelper
gripper_mesh = Mesh(
name="target_gripper_1",
file_path=join_path(
get_assets_path(),
"robot/franka_description/meshes/visual/hand.dae",
),
color=[0.0, 0.8, 0.1, 1.0],
pose=pose,
)
world.add_obstacle(gripper_mesh)
# get costs: # get costs:
if plot_cost: if plot_cost:
@@ -437,9 +424,12 @@ def benchmark_mb(
"valid_query": result.valid_query, "valid_query": result.valid_query,
} }
problem["solution_debug"] = debug problem["solution_debug"] = debug
world_coll = WorldConfig.from_dict(problem["obstacles"]).get_obb_world()
# check if path is collision free w.r.t. ground truth mesh: # check if path is collision free w.r.t. ground truth mesh:
robot_world.world_model.clear_cache() # robot_world.world_model.clear_cache()
robot_world.update_world(world) robot_world.update_world(world_coll)
q_int_traj = result.get_interpolated_plan().position.unsqueeze(0) q_int_traj = result.get_interpolated_plan().position.unsqueeze(0)
d_int_mask = ( d_int_mask = (
torch.count_nonzero(~robot_world.validate_trajectory(q_int_traj)) == 0 torch.count_nonzero(~robot_world.validate_trajectory(q_int_traj)) == 0
@@ -449,7 +439,14 @@ def benchmark_mb(
d_mask = ( d_mask = (
torch.count_nonzero(~robot_world.validate_trajectory(q_traj)) == 0 torch.count_nonzero(~robot_world.validate_trajectory(q_traj)) == 0
).item() ).item()
# d_world, _ = robot_world.get_world_self_collision_distance_from_joint_trajectory(
# q_traj)
# thres_dist = robot_world.contact_distance
# in_collision = d_world.squeeze(0) > thres_dist
# d_mask = not torch.any(in_collision, dim=-1).item()
# if not d_mask:
# write_usd = True
# #print(torch.max(d_world).item(), problem_name)
current_metrics = CuroboMetrics( current_metrics = CuroboMetrics(
skip=False, skip=False,
success=True, success=True,
@@ -466,10 +463,40 @@ def benchmark_mb(
attempts=result.attempts, attempts=result.attempts,
motion_time=result.motion_time.item(), motion_time=result.motion_time.item(),
solve_time=result.solve_time, solve_time=result.solve_time,
jerk=torch.max(torch.abs(result.optimized_plan.jerk)).item(),
) )
if save_log or write_usd:
world.randomize_color(r=[0.5, 0.9], g=[0.2, 0.5], b=[0.0, 0.2])
nvblox_obs = mg.world_coll_checker.get_mesh_from_blox_layer(
"world",
)
# nvblox_obs.vertex_colors = None
if nvblox_obs.vertex_colors is not None:
nvblox_obs.vertex_colors = nvblox_obs.vertex_colors.cpu().numpy()
else:
nvblox_obs.color = [0.0, 0.0, 0.8, 0.8]
nvblox_obs.name = "nvblox_mesh_world"
world.add_obstacle(nvblox_obs)
coll_mesh = mg.world_coll_checker.get_mesh_in_bounding_box(
curobo_Cuboid(
name="test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1.5, 1.5, 1]
),
voxel_size=0.005,
)
coll_mesh.color = [0.0, 0.8, 0.8, 0.8]
coll_mesh.name = "nvblox_voxel_world"
world.add_obstacle(coll_mesh)
# run planner
if write_usd: if write_usd:
# CuRobo # CuRobo
from curobo.util.usd_helper import UsdHelper
q_traj = result.get_interpolated_plan() q_traj = result.get_interpolated_plan()
UsdHelper.write_trajectory_animation_with_robot_usd( UsdHelper.write_trajectory_animation_with_robot_usd(
@@ -486,7 +513,8 @@ def benchmark_mb(
visualize_robot_spheres=True, visualize_robot_spheres=True,
# flatten_usd=True, # flatten_usd=True,
) )
exit() # write_usd = False
# exit()
if write_plot: if write_plot:
problem_name = problem_name problem_name = problem_name
plot_traj( plot_traj(
@@ -540,6 +568,9 @@ def benchmark_mb(
} }
problem["solution_debug"] = debug problem["solution_debug"] = debug
if save_log and not result.success.item(): if save_log and not result.success.item():
# CuRobo
from curobo.util.usd_helper import UsdHelper
UsdHelper.write_motion_gen_log( UsdHelper.write_motion_gen_log(
result, result,
robot_cfg, robot_cfg,
@@ -553,7 +584,7 @@ def benchmark_mb(
grid_space=2, grid_space=2,
# flatten_usd=True, # flatten_usd=True,
) )
# exit() exit()
g_m = CuroboGroupMetrics.from_list(m_list) g_m = CuroboGroupMetrics.from_list(m_list)
print( print(
key, key,
@@ -566,7 +597,7 @@ def benchmark_mb(
g_m.orientation_error.percent_98, g_m.orientation_error.percent_98,
g_m.cspace_path_length.percent_98, g_m.cspace_path_length.percent_98,
g_m.motion_time.percent_98, g_m.motion_time.percent_98,
g_m.perception_success, g_m.perception_interpolated_success,
# g_m.orientation_error.median, # g_m.orientation_error.median,
) )
print(g_m.attempts) print(g_m.attempts)
@@ -600,6 +631,7 @@ def benchmark_mb(
print("PT:", g_m.time) print("PT:", g_m.time)
print("ST: ", g_m.solve_time) print("ST: ", g_m.solve_time)
print("accuracy: ", g_m.position_error, g_m.orientation_error) print("accuracy: ", g_m.position_error, g_m.orientation_error)
print("Jerk: ", g_m.jerk)
if __name__ == "__main__": if __name__ == "__main__":

15
benchmark/curobo_profile.py
View File

@@ -73,8 +73,8 @@ def load_curobo(
position_threshold=0.005, position_threshold=0.005,
rotation_threshold=0.05, rotation_threshold=0.05,
num_ik_seeds=30, num_ik_seeds=30,
num_trajopt_seeds=10, num_trajopt_seeds=12,
interpolation_dt=0.02, interpolation_dt=0.025,
# grad_trajopt_iters=200, # grad_trajopt_iters=200,
store_ik_debug=enable_log, store_ik_debug=enable_log,
store_trajopt_debug=enable_log, store_trajopt_debug=enable_log,
@@ -91,7 +91,7 @@ def benchmark_mb(args):
spheres = load_yaml(join_path(get_robot_configs_path(), spheres))["collision_spheres"] spheres = load_yaml(join_path(get_robot_configs_path(), spheres))["collision_spheres"]
plan_config = MotionGenPlanConfig( plan_config = MotionGenPlanConfig(
max_attempts=2, max_attempts=1,
enable_graph_attempt=3, enable_graph_attempt=3,
enable_finetune_trajopt=True, enable_finetune_trajopt=True,
partial_ik_opt=False, partial_ik_opt=False,
@@ -130,11 +130,18 @@ def benchmark_mb(args):
world = WorldConfig.from_dict(problem["obstacles"]).get_obb_world() world = WorldConfig.from_dict(problem["obstacles"]).get_obb_world()
mg.update_world(world) mg.update_world(world)
start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
result = mg.plan_single(
start_state,
Pose.from_list(pose),
plan_config,
)
print(result.total_time, result.solve_time)
# continue # continue
# load obstacles # load obstacles
# run planner # run planner
start_state = JointState.from_position(mg.tensor_args.to_device([q_start]))
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
result = mg.plan_single( result = mg.plan_single(
start_state, start_state,

2
benchmark/generate_nvblox_images.py
View File

@@ -43,7 +43,7 @@ def generate_images():
for key, v in tqdm(problems.items()): for key, v in tqdm(problems.items()):
scene_problems = problems[key] scene_problems = problems[key]
i = 0 i = -1
for problem in tqdm(scene_problems, leave=False): for problem in tqdm(scene_problems, leave=False):
i += 1 i += 1

3
benchmark/metrics.py
View File

@@ -29,6 +29,7 @@ class CuroboMetrics(TrajectoryMetrics):
cspace_path_length: float = 0.0 cspace_path_length: float = 0.0
perception_success: bool = False perception_success: bool = False
perception_interpolated_success: bool = False perception_interpolated_success: bool = False
jerk: float = np.inf
@dataclass @dataclass
@@ -37,6 +38,7 @@ class CuroboGroupMetrics(TrajectoryGroupMetrics):
cspace_path_length: Optional[Statistic] = None cspace_path_length: Optional[Statistic] = None
perception_success: float = 0.0 perception_success: float = 0.0
perception_interpolated_success: float = 0.0 perception_interpolated_success: float = 0.0
jerk: float = np.inf
@classmethod @classmethod
def from_list(cls, group: List[CuroboMetrics]): def from_list(cls, group: List[CuroboMetrics]):
@@ -49,5 +51,6 @@ class CuroboGroupMetrics(TrajectoryGroupMetrics):
data.perception_interpolated_success = percent_true( data.perception_interpolated_success = percent_true(
[m.perception_interpolated_success for m in group] [m.perception_interpolated_success for m in group]
) )
data.jerk = Statistic.from_list([m.jerk for m in successes])
return data return data

17
docker/aarch64.dockerfile
View File

@@ -129,24 +129,15 @@ RUN pip3 install trimesh \
empy empy
# Add cache date to avoid using cached layers older than this # Add cache date to avoid using cached layers older than this
ARG CACHE_DATE=2023-12-15 ARG CACHE_DATE=2024-02-20
# warp from https://github.com/NVIDIA/warp needs to be compiled locally and then
# placed in curobo/docker/pkgs.
# Run the following from your terminal:
# cd curobo/docker && mkdir pkgs && cd pkgs && git clone https://github.com/NVIDIA/warp.git
# cd warp && python build_libs.py
#
# copy pkgs directory:
COPY pkgs /pkgs
# install warp: # install warp:
# #
RUN cd /pkgs/warp && pip3 install . RUN pip3 install warp-lang
# install curobo: # install curobo:
RUN cd /pkgs && git clone https://github.com/NVlabs/curobo.git RUN mkdir /pkgs && git clone https://github.com/NVlabs/curobo.git
ENV TORCH_CUDA_ARCH_LIST "7.0+PTX" ENV TORCH_CUDA_ARCH_LIST "7.0+PTX"
@@ -157,6 +148,7 @@ WORKDIR /pkgs/curobo
# Optionally install nvblox: # Optionally install nvblox:
ENV PYOPENGL_PLATFORM=egl ENV PYOPENGL_PLATFORM=egl
RUN apt-get update && \ RUN apt-get update && \
apt-get install -y libgoogle-glog-dev libgtest-dev curl libsqlite3-dev libbenchmark-dev && \ apt-get install -y libgoogle-glog-dev libgtest-dev curl libsqlite3-dev libbenchmark-dev && \
cd /usr/src/googletest && cmake . && cmake --build . --target install && \ cd /usr/src/googletest && cmake . && cmake --build . --target install && \
@@ -173,6 +165,7 @@ RUN cd /pkgs && git clone https://github.com/nvlabs/nvblox_torch.git && \
sh install.sh $(python3 -c 'import torch.utils; print(torch.utils.cmake_prefix_path)') && \ sh install.sh $(python3 -c 'import torch.utils; print(torch.utils.cmake_prefix_path)') && \
python3 -m pip install -e . python3 -m pip install -e .
RUN python -m pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git" RUN python -m pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git"
# upgrade typing extensions: # upgrade typing extensions:

1
docker/user.dockerfile
View File

@@ -15,6 +15,7 @@ FROM curobo_docker:${IMAGE_TAG}
# Set variables # Set variables
ARG USERNAME ARG USERNAME
ARG USER_ID ARG USER_ID
ARG CACHE_DATE=2024-02-20
# Set environment variables # Set environment variables

2
docker/x86.dockerfile
View File

@@ -78,7 +78,7 @@ ENV TORCH_CUDA_ARCH_LIST "7.0+PTX"
ENV LD_LIBRARY_PATH="/usr/local/lib:${LD_LIBRARY_PATH}" ENV LD_LIBRARY_PATH="/usr/local/lib:${LD_LIBRARY_PATH}"
# Add cache date to avoid using cached layers older than this # Add cache date to avoid using cached layers older than this
ARG CACHE_DATE=2023-12-15 ARG CACHE_DATE=2024-02-20
RUN pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git" RUN pip install "robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git"

355
examples/isaac_sim/constrained_reacher.py Normal file
View File

@@ -0,0 +1,355 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Third Party
import torch
a = torch.zeros(4, device="cuda:0")
# Third Party
from omni.isaac.kit import SimulationApp
simulation_app = SimulationApp(
{
"headless": False,
"width": "1920",
"height": "1080",
}
)
# Third Party
import numpy as np
import torch
from helper import add_extensions, add_robot_to_scene
# CuRobo
from curobo.geom.sdf.world import CollisionCheckerType
from curobo.geom.types import Cuboid, WorldConfig
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import JointState, RobotConfig
from curobo.types.state import JointState
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
from curobo.wrap.model.robot_world import RobotWorld, RobotWorldConfig
from curobo.wrap.reacher.motion_gen import (
MotionGen,
MotionGenConfig,
MotionGenPlanConfig,
PoseCostMetric,
)
simulation_app.update()
# Standard Library
import argparse
# Third Party
import carb
from omni.isaac.core import World
from omni.isaac.core.materials import OmniGlass, OmniPBR
from omni.isaac.core.objects import cuboid, sphere
from omni.isaac.core.utils.types import ArticulationAction
# CuRobo
from curobo.util.usd_helper import UsdHelper
parser = argparse.ArgumentParser()
parser.add_argument("--robot", type=str, default="franka.yml", help="robot configuration to load")
args = parser.parse_args()
if __name__ == "__main__":
my_world = World(stage_units_in_meters=1.0)
stage = my_world.stage
n_obstacle_cuboids = 10
n_obstacle_mesh = 10
stage = my_world.stage
my_world.scene.add_default_ground_plane()
xform = stage.DefinePrim("/World", "Xform")
stage.SetDefaultPrim(xform)
target_material = OmniPBR("/World/looks/t", color=np.array([0, 1, 0]))
target_material_2 = OmniPBR("/World/looks/t2", color=np.array([0, 1, 0]))
target_material_plane = OmniGlass(
"/World/looks/t3", color=np.array([0, 1, 0]), ior=1.25, depth=0.001, thin_walled=False
)
target_material_line = OmniGlass(
"/World/looks/t4", color=np.array([0, 1, 0]), ior=1.25, depth=0.001, thin_walled=True
)
# target_orient = [0,0,0.707,0.707]
target_orient = [0.5, -0.5, 0.5, 0.5]
target = cuboid.VisualCuboid(
"/World/target_1",
position=np.array([0.55, -0.3, 0.5]),
orientation=np.array(target_orient),
size=0.04,
visual_material=target_material,
)
# Make a target to follow
target_2 = cuboid.VisualCuboid(
"/World/target_2",
position=np.array([0.55, 0.4, 0.5]),
orientation=np.array(target_orient),
size=0.04,
visual_material=target_material_2,
)
x_plane = cuboid.VisualCuboid(
"/World/constraint_plane",
position=np.array([0.55, 0.05, 0.5]),
orientation=np.array(target_orient),
scale=[1.1, 0.001, 1.0],
visual_material=target_material_plane,
)
xz_line = cuboid.VisualCuboid(
"/World/constraint_line",
position=np.array([0.55, 0.05, 0.5]),
orientation=np.array(target_orient),
scale=[0.04, 0.04, 0.65],
visual_material=target_material_line,
)
collision_checker_type = CollisionCheckerType.BLOX
tensor_args = TensorDeviceType()
robot_cfg = load_yaml(join_path(get_robot_configs_path(), args.robot))["robot_cfg"]
j_names = robot_cfg["kinematics"]["cspace"]["joint_names"]
default_config = robot_cfg["kinematics"]["cspace"]["retract_config"]
robot, robot_prim_path = add_robot_to_scene(robot_cfg, my_world, "/World/world_robot/")
world_cfg_table = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
)
world_cfg_table.cuboid[0].pose[2] -= 0.01
world_cfg1 = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
).get_mesh_world()
world_cfg1.mesh[0].name += "_mesh"
world_cfg1.mesh[0].pose[2] = -10.5
world_cfg = WorldConfig(cuboid=world_cfg_table.cuboid, mesh=world_cfg1.mesh)
usd_help = UsdHelper()
usd_help.load_stage(my_world.stage)
usd_help.add_world_to_stage(world_cfg_table.get_mesh_world(), base_frame="/World")
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_cfg,
world_cfg,
tensor_args,
collision_checker_type=CollisionCheckerType.MESH,
collision_cache={"obb": n_obstacle_cuboids, "mesh": n_obstacle_mesh},
velocity_scale=0.75,
interpolation_dt=0.02,
ee_link_name="right_gripper",
)
motion_gen = MotionGen(motion_gen_config)
print("warming up..")
motion_gen.warmup(warmup_js_trajopt=False)
world_model = motion_gen.world_collision
i = 0
tensor_args = TensorDeviceType()
target_list = [target, target_2]
target_material_list = [target_material, target_material_2]
for material in target_material_list:
material.set_color(np.array([0.1, 0.1, 0.1]))
target_material_plane.set_color(np.array([1, 1, 1]))
target_material_line.set_color(np.array([1, 1, 1]))
target_idx = 0
cmd_idx = 0
cmd_plan = None
articulation_controller = robot.get_articulation_controller()
plan_config = MotionGenPlanConfig(
enable_graph=False, enable_graph_attempt=4, max_attempts=2, enable_finetune_trajopt=True
)
plan_idx = 0
cmd_step_idx = 0
pose_cost_metric = None
while simulation_app.is_running():
my_world.step(render=True)
if not my_world.is_playing():
if i % 100 == 0:
print("**** Click Play to start simulation *****")
i += 1
# if step_index == 0:
# my_world.play()
continue
step_index = my_world.current_time_step_index
if step_index <= 2:
my_world.reset()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
robot._articulation_view.set_max_efforts(
values=np.array([5000 for i in range(len(idx_list))]), joint_indices=idx_list
)
if False and step_index % 50 == 0.0: # No obstacle update
obstacles = usd_help.get_obstacles_from_stage(
# only_paths=[obstacles_path],
reference_prim_path=robot_prim_path,
ignore_substring=[
robot_prim_path,
"/World/target",
"/World/defaultGroundPlane",
"/curobo",
],
).get_collision_check_world()
# print(len(obstacles.objects))
motion_gen.update_world(obstacles)
# print("Updated World")
carb.log_info("Synced CuRobo world from stage.")
linear_color = np.ravel([249, 87, 56]) / 255.0
orient_color = np.ravel([103, 148, 54]) / 255.0
disable_color = np.ravel([255, 255, 255]) / 255.0
if cmd_plan is None and step_index % 10 == 0:
if plan_idx == 4:
print("Constrained: Holding tool linear-y")
pose_cost_metric = PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([0, 0, 0, 0, 1, 0]),
)
target_material_plane.set_color(linear_color)
if plan_idx == 8:
print("Constrained: Holding tool Orientation and linear-y")
pose_cost_metric = PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([1, 1, 1, 0, 1, 0]),
)
target_material_plane.set_color(orient_color)
if plan_idx == 12:
print("Constrained: Holding tool linear-y, linear-x")
pose_cost_metric = PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([0, 0, 0, 1, 1, 0]),
)
target_material_line.set_color(linear_color)
target_material_plane.set_color(disable_color)
if plan_idx == 16:
print("Constrained: Holding tool Orientation and linear-y, linear-x")
pose_cost_metric = PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([1, 1, 1, 1, 1, 0]),
)
target_material_line.set_color(orient_color)
target_material_plane.set_color(disable_color)
if plan_idx > 20:
plan_idx = 0
if plan_idx == 0:
print("Constrained: Reset")
target_material_line.set_color(disable_color)
target_material_plane.set_color(disable_color)
pose_cost_metric = None
plan_config.pose_cost_metric = pose_cost_metric
# motion generation:
for ks in range(len(target_material_list)):
if ks == target_idx:
target_material_list[ks].set_color(np.ravel([0, 1.0, 0]))
else:
target_material_list[ks].set_color(np.ravel([0.1, 0.1, 0.1]))
sim_js = robot.get_joints_state()
sim_js_names = robot.dof_names
cu_js = JointState(
position=tensor_args.to_device(sim_js.positions),
velocity=tensor_args.to_device(sim_js.velocities) * 0.0,
acceleration=tensor_args.to_device(sim_js.velocities) * 0.0,
jerk=tensor_args.to_device(sim_js.velocities) * 0.0,
joint_names=sim_js_names,
)
cu_js = cu_js.get_ordered_joint_state(motion_gen.kinematics.joint_names)
cube_position, cube_orientation = target_list[target_idx].get_world_pose()
# Set EE teleop goals, use cube for simple non-vr init:
ee_translation_goal = cube_position
ee_orientation_teleop_goal = cube_orientation
# compute curobo solution:
ik_goal = Pose(
position=tensor_args.to_device(ee_translation_goal),
quaternion=tensor_args.to_device(ee_orientation_teleop_goal),
)
result = motion_gen.plan_single(cu_js.unsqueeze(0), ik_goal, plan_config)
# ik_result = ik_solver.solve_single(ik_goal, cu_js.position.view(1,-1), cu_js.position.view(1,1,-1))
succ = result.success.item() # ik_result.success.item()
plan_idx += 1
if succ:
cmd_plan = result.get_interpolated_plan()
cmd_plan = motion_gen.get_full_js(cmd_plan)
# get only joint names that are in both:
idx_list = []
common_js_names = []
for x in sim_js_names:
if x in cmd_plan.joint_names:
idx_list.append(robot.get_dof_index(x))
common_js_names.append(x)
# idx_list = [robot.get_dof_index(x) for x in sim_js_names]
cmd_plan = cmd_plan.get_ordered_joint_state(common_js_names)
cmd_idx = 0
target_idx += 1
if target_idx >= len(target_list):
target_idx = 0
else:
carb.log_warn("Plan did not converge to a solution. No action is being taken.")
if cmd_plan is not None:
cmd_state = cmd_plan[cmd_idx]
# get full dof state
art_action = ArticulationAction(
cmd_state.position.cpu().numpy(),
cmd_state.velocity.cpu().numpy(),
joint_indices=idx_list,
)
# set desired joint angles obtained from IK:
articulation_controller.apply_action(art_action)
cmd_step_idx += 1
if cmd_step_idx == 2:
cmd_idx += 1
cmd_step_idx = 0
# for _ in range(2):
# my_world.step(render=False)
if cmd_idx >= len(cmd_plan.position):
cmd_idx = 0
cmd_plan = None
print("finished program")
simulation_app.close()

80
examples/isaac_sim/motion_gen_reacher.py
View File

@@ -51,6 +51,32 @@ parser.add_argument(
help="When True, runs in reactive mode", help="When True, runs in reactive mode",
default=False, default=False,
) )
parser.add_argument(
"--constrain_grasp_approach",
action="store_true",
help="When True, approaches grasp with fixed orientation and motion only along z axis.",
default=False,
)
parser.add_argument(
"--reach_partial_pose",
nargs=6,
metavar=("qx", "qy", "qz", "x", "y", "z"),
help="Reach partial pose",
type=float,
default=None,
)
parser.add_argument(
"--hold_partial_pose",
nargs=6,
metavar=("qx", "qy", "qz", "x", "y", "z"),
help="Hold partial pose while moving to goal",
type=float,
default=None,
)
args = parser.parse_args() args = parser.parse_args()
############################################################ ############################################################
@@ -97,7 +123,12 @@ from curobo.util_file import (
join_path, join_path,
load_yaml, load_yaml,
) )
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig from curobo.wrap.reacher.motion_gen import (
MotionGen,
MotionGenConfig,
MotionGenPlanConfig,
PoseCostMetric,
)
############################################################ ############################################################
@@ -107,7 +138,7 @@ from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGen
def main(): def main():
# create a curobo motion gen instance: # create a curobo motion gen instance:
num_targets = 0
# assuming obstacles are in objects_path: # assuming obstacles are in objects_path:
my_world = World(stage_units_in_meters=1.0) my_world = World(stage_units_in_meters=1.0)
stage = my_world.stage stage = my_world.stage
@@ -152,12 +183,12 @@ def main():
robot, robot_prim_path = add_robot_to_scene(robot_cfg, my_world) robot, robot_prim_path = add_robot_to_scene(robot_cfg, my_world)
articulation_controller = robot.get_articulation_controller() articulation_controller = None
world_cfg_table = WorldConfig.from_dict( world_cfg_table = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml")) load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
) )
world_cfg_table.cuboid[0].pose[2] -= 0.04 world_cfg_table.cuboid[0].pose[2] -= 0.02
world_cfg1 = WorldConfig.from_dict( world_cfg1 = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml")) load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
).get_mesh_world() ).get_mesh_world()
@@ -171,12 +202,14 @@ def main():
trajopt_tsteps = 32 trajopt_tsteps = 32
trim_steps = None trim_steps = None
max_attempts = 4 max_attempts = 4
interpolation_dt = 0.05
if args.reactive: if args.reactive:
trajopt_tsteps = 36 trajopt_tsteps = 40
trajopt_dt = 0.05 trajopt_dt = 0.05
optimize_dt = False optimize_dt = False
max_attemtps = 1 max_attempts = 1
trim_steps = [1, None] trim_steps = [1, None]
interpolation_dt = trajopt_dt
motion_gen_config = MotionGenConfig.load_from_robot_config( motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_cfg, robot_cfg,
world_cfg, world_cfg,
@@ -184,12 +217,15 @@ def main():
collision_checker_type=CollisionCheckerType.MESH, collision_checker_type=CollisionCheckerType.MESH,
num_trajopt_seeds=12, num_trajopt_seeds=12,
num_graph_seeds=12, num_graph_seeds=12,
interpolation_dt=0.05, interpolation_dt=interpolation_dt,
collision_cache={"obb": n_obstacle_cuboids, "mesh": n_obstacle_mesh}, collision_cache={"obb": n_obstacle_cuboids, "mesh": n_obstacle_mesh},
optimize_dt=optimize_dt, optimize_dt=optimize_dt,
trajopt_dt=trajopt_dt, trajopt_dt=trajopt_dt,
trajopt_tsteps=trajopt_tsteps, trajopt_tsteps=trajopt_tsteps,
trim_steps=trim_steps, trim_steps=trim_steps,
# velocity_scale=0.1,
# self_collision_check=False,
# self_collision_opt=False,
) )
motion_gen = MotionGen(motion_gen_config) motion_gen = MotionGen(motion_gen_config)
print("warming up...") print("warming up...")
@@ -216,6 +252,9 @@ def main():
i = 0 i = 0
spheres = None spheres = None
past_cmd = None past_cmd = None
target_orientation = None
past_orientation = None
pose_metric = None
while simulation_app.is_running(): while simulation_app.is_running():
my_world.step(render=True) my_world.step(render=True)
if not my_world.is_playing(): if not my_world.is_playing():
@@ -228,6 +267,9 @@ def main():
step_index = my_world.current_time_step_index step_index = my_world.current_time_step_index
# print(step_index) # print(step_index)
if articulation_controller is None:
# robot.initialize()
articulation_controller = robot.get_articulation_controller()
if step_index < 2: if step_index < 2:
my_world.reset() my_world.reset()
robot._articulation_view.initialize() robot._articulation_view.initialize()
@@ -265,6 +307,10 @@ def main():
past_pose = cube_position past_pose = cube_position
if target_pose is None: if target_pose is None:
target_pose = cube_position target_pose = cube_position
if target_orientation is None:
target_orientation = cube_orientation
if past_orientation is None:
past_orientation = cube_orientation
sim_js = robot.get_joints_state() sim_js = robot.get_joints_state()
sim_js_names = robot.dof_names sim_js_names = robot.dof_names
@@ -313,8 +359,12 @@ def main():
if (np.max(np.abs(sim_js.velocities)) < 0.2) or args.reactive: if (np.max(np.abs(sim_js.velocities)) < 0.2) or args.reactive:
robot_static = True robot_static = True
if ( if (
(
np.linalg.norm(cube_position - target_pose) > 1e-3 np.linalg.norm(cube_position - target_pose) > 1e-3
or np.linalg.norm(cube_orientation - target_orientation) > 1e-3
)
and np.linalg.norm(past_pose - cube_position) == 0.0 and np.linalg.norm(past_pose - cube_position) == 0.0
and np.linalg.norm(past_orientation - cube_orientation) == 0.0
and robot_static and robot_static
): ):
# Set EE teleop goals, use cube for simple non-vr init: # Set EE teleop goals, use cube for simple non-vr init:
@@ -326,12 +376,24 @@ def main():
position=tensor_args.to_device(ee_translation_goal), position=tensor_args.to_device(ee_translation_goal),
quaternion=tensor_args.to_device(ee_orientation_teleop_goal), quaternion=tensor_args.to_device(ee_orientation_teleop_goal),
) )
plan_config.pose_cost_metric = pose_metric
result = motion_gen.plan_single(cu_js.unsqueeze(0), ik_goal, plan_config) result = motion_gen.plan_single(cu_js.unsqueeze(0), ik_goal, plan_config)
# ik_result = ik_solver.solve_single(ik_goal, cu_js.position.view(1,-1), cu_js.position.view(1,1,-1)) # ik_result = ik_solver.solve_single(ik_goal, cu_js.position.view(1,-1), cu_js.position.view(1,1,-1))
succ = result.success.item() # ik_result.success.item() succ = result.success.item() # ik_result.success.item()
if num_targets == 1:
if args.constrain_grasp_approach:
pose_metric = PoseCostMetric.create_grasp_approach_metric()
if args.reach_partial_pose is not None:
reach_vec = motion_gen.tensor_args.to_device(args.reach_partial_pose)
pose_metric = PoseCostMetric(
reach_partial_pose=True, reach_vec_weight=reach_vec
)
if args.hold_partial_pose is not None:
hold_vec = motion_gen.tensor_args.to_device(args.hold_partial_pose)
pose_metric = PoseCostMetric(hold_partial_pose=True, hold_vec_weight=hold_vec)
if succ: if succ:
num_targets += 1
cmd_plan = result.get_interpolated_plan() cmd_plan = result.get_interpolated_plan()
cmd_plan = motion_gen.get_full_js(cmd_plan) cmd_plan = motion_gen.get_full_js(cmd_plan)
# get only joint names that are in both: # get only joint names that are in both:
@@ -350,7 +412,9 @@ def main():
else: else:
carb.log_warn("Plan did not converge to a solution. No action is being taken.") carb.log_warn("Plan did not converge to a solution. No action is being taken.")
target_pose = cube_position target_pose = cube_position
target_orientation = cube_orientation
past_pose = cube_position past_pose = cube_position
past_orientation = cube_orientation
if cmd_plan is not None: if cmd_plan is not None:
cmd_state = cmd_plan[cmd_idx] cmd_state = cmd_plan[cmd_idx]
past_cmd = cmd_state.clone() past_cmd = cmd_state.clone()

42
examples/isaac_sim/simple_stacking.py
View File

@@ -34,6 +34,13 @@ parser.add_argument(
default=None, default=None,
help="To run headless, use one of [native, websocket], webrtc might not work.", help="To run headless, use one of [native, websocket], webrtc might not work.",
) )
parser.add_argument(
"--constrain_grasp_approach",
action="store_true",
help="When True, approaches grasp with fixed orientation and motion only along z axis.",
default=False,
)
args = parser.parse_args() args = parser.parse_args()
# Third Party # Third Party
@@ -78,6 +85,7 @@ from curobo.wrap.reacher.motion_gen import (
MotionGenConfig, MotionGenConfig,
MotionGenPlanConfig, MotionGenPlanConfig,
MotionGenResult, MotionGenResult,
PoseCostMetric,
) )
@@ -87,6 +95,7 @@ class CuroboController(BaseController):
my_world: World, my_world: World,
my_task: BaseStacking, my_task: BaseStacking,
name: str = "curobo_controller", name: str = "curobo_controller",
constrain_grasp_approach: bool = False,
) -> None: ) -> None:
BaseController.__init__(self, name=name) BaseController.__init__(self, name=name)
self._save_log = False self._save_log = False
@@ -145,13 +154,16 @@ class CuroboController(BaseController):
collision_cache={"obb": n_obstacle_cuboids, "mesh": n_obstacle_mesh}, collision_cache={"obb": n_obstacle_cuboids, "mesh": n_obstacle_mesh},
store_ik_debug=self._save_log, store_ik_debug=self._save_log,
store_trajopt_debug=self._save_log, store_trajopt_debug=self._save_log,
state_finite_difference_mode="CENTRAL", velocity_scale=0.75,
acceleration_scale=0.5,
fixed_iters_trajopt=True,
) )
self.motion_gen = MotionGen(motion_gen_config) self.motion_gen = MotionGen(motion_gen_config)
print("warming up...") print("warming up...")
self.motion_gen.warmup() self.motion_gen.warmup(parallel_finetune=True)
pose_metric = None
if constrain_grasp_approach:
pose_metric = PoseCostMetric.create_grasp_approach_metric(
offset_position=0.1, tstep_fraction=0.6
)
self.plan_config = MotionGenPlanConfig( self.plan_config = MotionGenPlanConfig(
enable_graph=False, enable_graph=False,
@@ -159,6 +171,8 @@ class CuroboController(BaseController):
enable_graph_attempt=None, enable_graph_attempt=None,
enable_finetune_trajopt=True, enable_finetune_trajopt=True,
partial_ik_opt=False, partial_ik_opt=False,
parallel_finetune=True,
pose_cost_metric=pose_metric,
) )
self.usd_help.load_stage(self.my_world.stage) self.usd_help.load_stage(self.my_world.stage)
self.cmd_plan = None self.cmd_plan = None
@@ -185,7 +199,7 @@ class CuroboController(BaseController):
cu_js, cu_js,
[cube_name], [cube_name],
sphere_fit_type=SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE, sphere_fit_type=SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE,
world_objects_pose_offset=Pose.from_list([0, 0, 0.005, 1, 0, 0, 0], self.tensor_args), world_objects_pose_offset=Pose.from_list([0, 0, 0.01, 1, 0, 0, 0], self.tensor_args),
) )
def detach_obj(self) -> None: def detach_obj(self) -> None:
@@ -259,7 +273,7 @@ class CuroboController(BaseController):
# get full dof state # get full dof state
art_action = ArticulationAction( art_action = ArticulationAction(
cmd_state.position.cpu().numpy(), cmd_state.position.cpu().numpy(),
# cmd_state.velocity.cpu().numpy(), cmd_state.velocity.cpu().numpy() * 0.0,
joint_indices=self.idx_list, joint_indices=self.idx_list,
) )
if self.cmd_idx >= len(self.cmd_plan.position): if self.cmd_idx >= len(self.cmd_plan.position):
@@ -417,7 +431,9 @@ my_world.add_task(my_task)
my_world.reset() my_world.reset()
robot_name = my_task.get_params()["robot_name"]["value"] robot_name = my_task.get_params()["robot_name"]["value"]
my_franka = my_world.scene.get_object(robot_name) my_franka = my_world.scene.get_object(robot_name)
my_controller = CuroboController(my_world=my_world, my_task=my_task) my_controller = CuroboController(
my_world=my_world, my_task=my_task, constrain_grasp_approach=args.constrain_grasp_approach
)
articulation_controller = my_franka.get_articulation_controller() articulation_controller = my_franka.get_articulation_controller()
set_camera_view(eye=[2, 0, 1], target=[0.00, 0.00, 0.00], camera_prim_path="/OmniverseKit_Persp") set_camera_view(eye=[2, 0, 1], target=[0.00, 0.00, 0.00], camera_prim_path="/OmniverseKit_Persp")
wait_steps = 8 wait_steps = 8
@@ -439,17 +455,17 @@ print(articulation_controller.get_gains())
print(articulation_controller.get_max_efforts()) print(articulation_controller.get_max_efforts())
robot = my_franka robot = my_franka
print("**********************") print("**********************")
robot.enable_gravity() if True:
articulation_controller.set_gains( robot.enable_gravity()
articulation_controller.set_gains(
kps=np.array( kps=np.array(
[100000000, 6000000.0, 10000000, 600000.0, 25000.0, 15000.0, 50000.0, 6000.0, 6000.0] [100000000, 6000000.0, 10000000, 600000.0, 25000.0, 15000.0, 50000.0, 6000.0, 6000.0]
) )
) )
articulation_controller.set_max_efforts(
articulation_controller.set_max_efforts(
values=np.array([100000, 52.199997, 100000, 52.199997, 7.2, 7.2, 7.2, 50.0, 50]) values=np.array([100000, 52.199997, 100000, 52.199997, 7.2, 7.2, 7.2, 50.0, 50])
) )
print("Updated gains:") print("Updated gains:")
print(articulation_controller.get_gains()) print(articulation_controller.get_gains())

2
examples/kinematics_example.py
View File

@@ -15,7 +15,7 @@
import torch import torch
# CuRobo # CuRobo
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
from curobo.types.base import TensorDeviceType from curobo.types.base import TensorDeviceType
from curobo.types.robot import RobotConfig from curobo.types.robot import RobotConfig
from curobo.util.logger import setup_curobo_logger from curobo.util.logger import setup_curobo_logger

88
examples/motion_gen_example.py
View File

@@ -42,8 +42,9 @@ def plot_traj(trajectory, dt):
axs[3].plot(timesteps, qddd[:, i], label=str(i)) axs[3].plot(timesteps, qddd[:, i], label=str(i))
plt.legend() plt.legend()
# plt.savefig("test.png") plt.savefig("test.png")
plt.show() plt.close()
# plt.show()
def plot_iters_traj(trajectory, d_id=1, dof=7, seed=0): def plot_iters_traj(trajectory, d_id=1, dof=7, seed=0):
@@ -140,37 +141,54 @@ def demo_motion_gen(js=False):
world_file, world_file,
tensor_args, tensor_args,
interpolation_dt=0.01, interpolation_dt=0.01,
# trajopt_dt=0.15,
# velocity_scale=0.1,
use_cuda_graph=True,
# finetune_dt_scale=2.5,
interpolation_steps=10000,
) )
motion_gen = MotionGen(motion_gen_config) motion_gen = MotionGen(motion_gen_config)
motion_gen.warmup(parallel_finetune=True)
motion_gen.warmup(enable_graph=True, warmup_js_trajopt=js) # motion_gen.warmup(enable_graph=True, warmup_js_trajopt=js, parallel_finetune=True)
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"] # robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args) # robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
retract_cfg = motion_gen.get_retract_config() retract_cfg = motion_gen.get_retract_config()
state = motion_gen.rollout_fn.compute_kinematics( state = motion_gen.rollout_fn.compute_kinematics(
JointState.from_position(retract_cfg.view(1, -1)) JointState.from_position(retract_cfg.view(1, -1))
) )
retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()) retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.1) start_state = JointState.from_position(retract_cfg.view(1, -1))
goal_state = start_state.clone() goal_state = start_state.clone()
goal_state.position[..., 3] -= 0.1
start_state.position[0, 0] += 0.25
# goal_state.position[0,0] += 0.5
if js: if js:
result = motion_gen.plan_single_js( result = motion_gen.plan_single_js(
start_state, start_state,
goal_state, goal_state,
MotionGenPlanConfig( MotionGenPlanConfig(max_attempts=1, parallel_finetune=True),
max_attempts=1, enable_graph=False, enable_opt=True, enable_finetune_trajopt=True
),
) )
else: else:
result = motion_gen.plan_single( result = motion_gen.plan_single(
start_state, retract_pose, MotionGenPlanConfig(max_attempts=1) start_state,
retract_pose,
MotionGenPlanConfig(
max_attempts=1, parallel_finetune=True, enable_finetune_trajopt=True
),
)
print(
"Trajectory Generated: ",
result.success,
result.solve_time,
result.status,
result.optimized_dt,
) )
traj = result.get_interpolated_plan()
print("Trajectory Generated: ", result.success, result.solve_time, result.status)
if PLOT and result.success.item(): if PLOT and result.success.item():
traj = result.get_interpolated_plan()
plot_traj(traj, result.interpolation_dt) plot_traj(traj, result.interpolation_dt)
# plt.save("test.png") # plt.save("test.png")
# plt.close() # plt.close()
@@ -261,6 +279,45 @@ def demo_motion_gen_batch():
plot_traj(traj[1, : result.path_buffer_last_tstep[1], :].cpu().numpy()) plot_traj(traj[1, : result.path_buffer_last_tstep[1], :].cpu().numpy())
def demo_motion_gen_goalset():
tensor_args = TensorDeviceType()
world_file = "collision_cubby.yml"
robot_file = "franka.yml"
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_file,
world_file,
tensor_args,
collision_checker_type=CollisionCheckerType.PRIMITIVE,
use_cuda_graph=True,
num_trajopt_seeds=12,
num_graph_seeds=1,
num_ik_seeds=30,
)
motion_gen = MotionGen(motion_gen_config)
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.rollout_fn.compute_kinematics(
JointState.from_position(retract_cfg.view(1, -1))
)
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.6)
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.repeat(2, 1).view(1, -1, 3),
quaternion=state.ee_quat_seq.repeat(2, 1).view(1, -1, 4),
)
goal_pose.position[0, 0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
m_config = MotionGenPlanConfig(False, True, num_trajopt_seeds=10)
result = motion_gen.plan_goalset(start_state, goal_pose, m_config)
def demo_motion_gen_api(): def demo_motion_gen_api():
tensor_args = TensorDeviceType(device=torch.device("cuda:0")) tensor_args = TensorDeviceType(device=torch.device("cuda:0"))
interpolation_dt = 0.02 interpolation_dt = 0.02
@@ -373,8 +430,9 @@ def demo_motion_gen_batch_env(n_envs: int = 10):
if __name__ == "__main__": if __name__ == "__main__":
setup_curobo_logger("error") setup_curobo_logger("error")
# demo_motion_gen(js=True) demo_motion_gen(js=False)
demo_motion_gen_batch() # demo_motion_gen_batch()
# demo_motion_gen_goalset()
# n = [2, 10] # n = [2, 10]
# for n_envs in n: # for n_envs in n:
# demo_motion_gen_batch_env(n_envs=n_envs) # demo_motion_gen_batch_env(n_envs=n_envs)

238
examples/robot_image_segmentation_example.py Normal file
View File

@@ -0,0 +1,238 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
""" This example shows how to use cuRobo's kinematics to generate a mask. """
# Standard Library
import time
# Third Party
import imageio
import numpy as np
import torch
import torch.autograd.profiler as profiler
from nvblox_torch.datasets.mesh_dataset import MeshDataset
from torch.profiler import ProfilerActivity, profile, record_function
# CuRobo
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
from curobo.geom.types import PointCloud, WorldConfig
from curobo.types.base import TensorDeviceType
from curobo.types.camera import CameraObservation
from curobo.types.math import Pose
from curobo.types.robot import RobotConfig
from curobo.types.state import JointState
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
from curobo.wrap.model.robot_segmenter import RobotSegmenter
from curobo.wrap.model.robot_world import RobotWorld, RobotWorldConfig
torch.manual_seed(30)
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
def create_render_dataset(robot_file, save_debug_data: bool = False):
# load robot:
robot_dict = load_yaml(join_path(get_robot_configs_path(), robot_file))
robot_dict["robot_cfg"]["kinematics"]["load_link_names_with_mesh"] = True
robot_dict["robot_cfg"]["kinematics"]["load_meshes"] = True
robot_cfg = RobotConfig.from_dict(robot_dict["robot_cfg"])
kin_model = CudaRobotModel(robot_cfg.kinematics)
q = kin_model.retract_config
meshes = kin_model.get_robot_as_mesh(q)
world = WorldConfig(mesh=meshes[:])
world_table = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_test.yml"))
)
world_table.cuboid[0].dims = [0.5, 0.5, 0.1]
world.add_obstacle(world_table.objects[0])
world.add_obstacle(world_table.objects[1])
if save_debug_data:
world.save_world_as_mesh("scene.stl", process_color=False)
robot_mesh = (
WorldConfig.create_merged_mesh_world(world, process_color=False).mesh[0].get_trimesh_mesh()
)
mesh_dataset = MeshDataset(
None,
n_frames=20,
image_size=480,
save_data_dir=None,
trimesh_mesh=robot_mesh,
)
q_js = JointState(position=q, joint_names=kin_model.joint_names)
return mesh_dataset, q_js
def mask_image(robot_file="ur5e.yml"):
save_debug_data = False
# create robot segmenter:
tensor_args = TensorDeviceType()
curobo_segmenter = RobotSegmenter.from_robot_file(
robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
)
mesh_dataset, q_js = create_render_dataset(robot_file, save_debug_data)
if save_debug_data:
visualize_scale = 10.0
data = mesh_dataset[0]
cam_obs = CameraObservation(
depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
)
# save depth image
imageio.imwrite(
"camera_depth.png",
(cam_obs.depth_image * visualize_scale)
.squeeze()
.detach()
.cpu()
.numpy()
.astype(np.uint16),
)
# save robot spheres in current joint configuration
robot_kinematics = curobo_segmenter._robot_world.kinematics
sph = robot_kinematics.get_robot_as_spheres(q_js.position)
WorldConfig(sphere=sph[0]).save_world_as_mesh("robot_spheres.stl")
# save world pointcloud in robot origin
pc = cam_obs.get_pointcloud()
pc_obs = PointCloud("world", pose=cam_obs.pose.to_list(), points=pc)
pc_obs.save_as_mesh("camera_pointcloud.stl", transform_with_pose=True)
# run segmentation:
depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
cam_obs,
q_js,
)
# save robot points as mesh
robot_mask = cam_obs.clone()
robot_mask.depth_image[~depth_mask] = 0.0
robot_mesh = PointCloud(
"world", pose=robot_mask.pose.to_list(), points=robot_mask.get_pointcloud()
)
robot_mesh.save_as_mesh("robot_segmented.stl", transform_with_pose=True)
# save depth image
imageio.imwrite(
"robot_depth.png",
(robot_mask.depth_image * visualize_scale)
.detach()
.squeeze()
.cpu()
.numpy()
.astype(np.uint16),
)
# save world points as mesh
world_mask = cam_obs.clone()
world_mask.depth_image[depth_mask] = 0.0
world_mesh = PointCloud(
"world", pose=world_mask.pose.to_list(), points=world_mask.get_pointcloud()
)
world_mesh.save_as_mesh("world_segmented.stl", transform_with_pose=True)
imageio.imwrite(
"world_depth.png",
(world_mask.depth_image * visualize_scale)
.detach()
.squeeze()
.cpu()
.numpy()
.astype(np.uint16),
)
dt_list = []
for i in range(len(mesh_dataset)):
data = mesh_dataset[i]
cam_obs = CameraObservation(
depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
)
if not curobo_segmenter.ready:
curobo_segmenter.update_camera_projection(cam_obs)
st_time = time.time()
depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
cam_obs,
q_js,
)
torch.cuda.synchronize()
dt_list.append(time.time() - st_time)
print("Segmentation Time (ms), (hz)", np.mean(dt_list[5:]) * 1000.0, 1.0 / np.mean(dt_list[5:]))
def profile_mask_image(robot_file="ur5e.yml"):
# create robot segmenter:
tensor_args = TensorDeviceType()
curobo_segmenter = RobotSegmenter.from_robot_file(
robot_file, collision_sphere_buffer=0.0, distance_threshold=0.05, use_cuda_graph=False
)
mesh_dataset, q_js = create_render_dataset(robot_file)
dt_list = []
data = mesh_dataset[0]
cam_obs = CameraObservation(
depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
)
if not curobo_segmenter.ready:
curobo_segmenter.update_camera_projection(cam_obs)
depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(cam_obs, q_js)
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
for i in range(len(mesh_dataset)):
with profiler.record_function("get_data"):
data = mesh_dataset[i]
cam_obs = CameraObservation(
depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
)
st_time = time.time()
with profiler.record_function("segmentation"):
depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
cam_obs, q_js
)
print("Exporting the trace..")
prof.export_chrome_trace("segmentation.json")
if __name__ == "__main__":
mask_image()
# profile_mask_image()

2
setup.py
View File

@@ -87,4 +87,6 @@ ext_modules = [
setuptools.setup( setuptools.setup(
ext_modules=ext_modules, ext_modules=ext_modules,
cmdclass={"build_ext": BuildExtension}, cmdclass={"build_ext": BuildExtension},
package_data={"": ["*.so"]},
include_package_data=True,
) )

17
src/curobo/__init__.py
View File

@@ -9,7 +9,22 @@
# its affiliates is strictly prohibited. # its affiliates is strictly prohibited.
# #
"""CuRobo package.""" """
cuRobo package is split into several modules:
- :mod:`curobo.opt` contains optimization solvers.
- :mod:`curobo.cuda_robot_model` contains robot kinematics.
- :mod:`curobo.curobolib` contains the cuda kernels and python bindings for them.
- :mod:`curobo.geom` contains geometry processing, collision checking and frame transforms.
- :mod:`curobo.graph` contains geometric planning with graph search methods.
- :mod:`curobo.rollout` contains methods that map actions to costs. This class wraps instances of
:mod:`curobo.cuda_robot_model` and :mod:`geom` to compute costs given trajectory of actions.
- :mod:`curobo.util` contains utility methods.
- :mod:`curobo.wrap` adds the user-level api for task programming. Includes implementation of
collision-free reacher and batched robot world collision checking.
- :mod:`curobo.types` contains custom dataclasses for common data types in robotics, including
:meth:`types.state.JointState`, :meth:`types.camera.CameraObservation`, :meth:`types.math.Pose`.
"""
# NOTE (roflaherty): This is inspired by how matplotlib does creates its version value. # NOTE (roflaherty): This is inspired by how matplotlib does creates its version value.

11
src/curobo/content/assets/robot/franka_description/franka_panda.urdf
View File

@@ -172,7 +172,7 @@
<dynamics damping="10.0"/> <dynamics damping="10.0"/>
<limit effort="12" lower="-2.8973" upper="2.8973" velocity="2.6100"/> <limit effort="12" lower="-2.8973" upper="2.8973" velocity="2.6100"/>
</joint> </joint>
<!--
<link name="panda_link8"/> <link name="panda_link8"/>
<joint name="panda_joint8" type="fixed"> <joint name="panda_joint8" type="fixed">
<origin rpy="0 0 0" xyz="0 0 0.107"/> <origin rpy="0 0 0" xyz="0 0 0.107"/>
@@ -180,18 +180,11 @@
<child link="panda_link8"/> <child link="panda_link8"/>
<axis xyz="0 0 0"/> <axis xyz="0 0 0"/>
</joint> </joint>
Removing this joint seems to help with some stability things
-->
<joint name="panda_hand_joint" type="fixed"> <joint name="panda_hand_joint" type="fixed">
<!--
<parent link="panda_link8"/> <parent link="panda_link8"/>
-->
<parent link="panda_link7"/>
<child link="panda_hand"/> <child link="panda_hand"/>
<origin rpy="0 0 -0.785398163397" xyz="0 0 0.107"/>
<!--
<origin rpy="0 0 -0.785398163397" xyz="0 0 0"/> <origin rpy="0 0 -0.785398163397" xyz="0 0 0"/>
-->
</joint> </joint>
<link name="panda_hand"> <link name="panda_hand">
<visual> <visual>

6
src/curobo/content/configs/robot/franka.yml
View File

@@ -13,7 +13,7 @@ robot_cfg:
kinematics: kinematics:
use_usd_kinematics: False use_usd_kinematics: False
isaac_usd_path: "/Isaac/Robots/Franka/franka.usd" isaac_usd_path: "/Isaac/Robots/Franka/franka.usd"
usd_path: "robot/non_shipping/franka/franka_panda_meters.usda" usd_path: "robot/franka/franka_panda.usda"
usd_robot_root: "/panda" usd_robot_root: "/panda"
usd_flip_joints: ["panda_joint1","panda_joint2","panda_joint3","panda_joint4", "panda_joint5", usd_flip_joints: ["panda_joint1","panda_joint2","panda_joint3","panda_joint4", "panda_joint5",
"panda_joint6","panda_joint7","panda_finger_joint1", "panda_finger_joint2"] "panda_joint6","panda_joint7","panda_finger_joint1", "panda_finger_joint2"]
@@ -32,7 +32,7 @@ robot_cfg:
usd_flip_joint_limits: ["panda_finger_joint2"] usd_flip_joint_limits: ["panda_finger_joint2"]
urdf_path: "robot/franka_description/franka_panda.urdf" urdf_path: "robot/franka_description/franka_panda.urdf"
asset_root_path: "robot/franka_description" asset_root_path: "robot/franka_description"
base_link: "base_link" base_link: "panda_link0"
ee_link: "panda_hand" ee_link: "panda_hand"
collision_link_names: collision_link_names:
[ [
@@ -50,7 +50,7 @@ robot_cfg:
"attached_object", "attached_object",
] ]
collision_spheres: "spheres/franka_mesh.yml" collision_spheres: "spheres/franka_mesh.yml"
collision_sphere_buffer: 0.0025 collision_sphere_buffer: 0.004 # 0.0025
extra_collision_spheres: {"attached_object": 4} extra_collision_spheres: {"attached_object": 4}
use_global_cumul: True use_global_cumul: True
self_collision_ignore: self_collision_ignore:

2
src/curobo/content/configs/robot/franka_mobile.yml
View File

@@ -109,7 +109,7 @@ robot_cfg:
"base_x", "base_y", "base_z", "base_x", "base_y", "base_z",
"panda_joint1","panda_joint2","panda_joint3","panda_joint4", "panda_joint5", "panda_joint1","panda_joint2","panda_joint3","panda_joint4", "panda_joint5",
"panda_joint6","panda_joint7","panda_finger_joint1", "panda_finger_joint2"] "panda_joint6","panda_joint7","panda_finger_joint1", "panda_finger_joint2"]
retract_config: [0,0,0,0.0, -1.3, 0.0, -2.5, 0.0, 1.0, 0., 0.04, 0.04] retract_config: [0,0,0,0, -1.3, 0.0, -2.5, 0.0, 1.0, 0., 0.04, 0.04]
null_space_weight: [1,1,1,1,1,1,1,1,1,1,1,1] null_space_weight: [1,1,1,1,1,1,1,1,1,1,1,1]
cspace_distance_weight: [1,1,1,1,1,1,1,1,1,1,1,1] cspace_distance_weight: [1,1,1,1,1,1,1,1,1,1,1,1]
max_acceleration: 15.0 #15.0 max_acceleration: 15.0 #15.0

2
src/curobo/content/configs/robot/tm12.yml
View File

@@ -12,7 +12,7 @@
robot_cfg: robot_cfg:
kinematics: kinematics:
use_usd_kinematics: False use_usd_kinematics: False
usd_path: "robot/non_shipping/techman/tm12.usd" usd_path: "robot/techman/tm12.usd"
usd_robot_root: "/tm12" usd_robot_root: "/tm12"
usd_flip_joints: usd_flip_joints:
{ {

2
src/curobo/content/configs/robot/ur5e.yml
View File

@@ -49,7 +49,7 @@ robot_cfg:
- "center": [-0.055, 0.008, 0.14] - "center": [-0.055, 0.008, 0.14]
"radius": 0.07 "radius": 0.07
- "center": [-0.001, -0.002, 0.143] - "center": [-0.001, -0.002, 0.143]
"radius": 0.076 "radius": 0.08
forearm_link: forearm_link:
- "center": [-0.01, 0.002, 0.031] - "center": [-0.01, 0.002, 0.031]
"radius": 0.072 "radius": 0.072

21
src/curobo/content/configs/task/finetune_trajopt.yml
View File

@@ -31,22 +31,23 @@ model:
cost: cost:
pose_cfg: pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps
run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position
weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40] weight: [2000,500000.0,30,50]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True terminal: True
run_weight: 0.0 #0.05 run_weight: 1.0
use_metric: True use_metric: True
link_pose_cfg: link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps
run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position
weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40] weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True terminal: True
run_weight: 0.0 run_weight: 0.001
use_metric: True use_metric: True
cspace_cfg: cspace_cfg:
weight: 10000.0 weight: 10000.0
terminal: True terminal: True
@@ -54,7 +55,7 @@ cost:
bound_cfg: bound_cfg:
weight: [5000.0, 50000.0, 50000.0, 50000.0] # needs to be 3 values weight: [5000.0, 50000.0, 50000.0, 50000.0] # needs to be 3 values
smooth_weight: [0.0,2000.0, 50.0, 0.0] # [vel, acc, jerk,] smooth_weight: [0.0,5000.0, 50.0, 0.0] # [vel, acc, jerk,]
run_weight_velocity: 0.0 run_weight_velocity: 0.0
run_weight_acceleration: 1.0 run_weight_acceleration: 1.0
run_weight_jerk: 1.0 run_weight_jerk: 1.0
@@ -78,11 +79,11 @@ cost:
lbfgs: lbfgs:
n_iters: 400 n_iters: 400 # 400
inner_iters: 25 inner_iters: 25
cold_start_n_iters: 200 cold_start_n_iters: null
min_iters: 25 min_iters: 25
line_search_scale: [0.01, 0.3,0.7,1.0] # # line_search_scale: [0.01, 0.3, 0.7,1.0] # #
fixed_iters: True fixed_iters: True
cost_convergence: 0.01 cost_convergence: 0.01
cost_delta_threshold: 1.0 cost_delta_threshold: 1.0
@@ -90,7 +91,7 @@ lbfgs:
epsilon: 0.01 epsilon: 0.01
history: 15 #15 history: 15 #15
use_cuda_graph: True use_cuda_graph: True
n_envs: 1 n_problems: 1
store_debug: False store_debug: False
use_cuda_kernel: True use_cuda_kernel: True
stable_mode: True stable_mode: True

109
src/curobo/content/configs/task/finetune_trajopt_slow.yml Normal file
View File

@@ -0,0 +1,109 @@
##
## 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.
##
model:
horizon: 32
state_filter_cfg:
filter_coeff:
position: 1.0
velocity: 1.0
acceleration: 1.0
enable: False
dt_traj_params:
base_dt: 0.2
base_ratio: 1.0
max_dt: 0.2
vel_scale: 1.0
control_space: 'POSITION'
state_finite_difference_mode: "CENTRAL"
teleport_mode: False
return_full_act_buffer: True
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: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position
weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
run_weight: 0.0 #0.05
use_metric: True
link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps
run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position
weight: [2000,500000.0,30,50] #[150.0, 2000.0, 30, 40]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True
run_weight: 0.001
use_metric: True
cspace_cfg:
weight: 10000.0
terminal: True
run_weight: 0.00 #1
bound_cfg:
weight: [5000.0, 50000.0, 50000.0, 50000.0] # needs to be 3 values
smooth_weight: [0.0,50000.0, 500.0, 0.0] # [vel, acc, jerk,]
run_weight_velocity: 0.0
run_weight_acceleration: 1.0
run_weight_jerk: 1.0
activation_distance: [0.1,0.1,0.1,10.0] # for position, velocity, acceleration and jerk
null_space_weight: [0.0]
primitive_collision_cfg:
weight: 1000000.0
use_sweep: True
sweep_steps: 6
classify: False
use_sweep_kernel: True
use_speed_metric: True
speed_dt: 0.01 # used only for speed metric
activation_distance: 0.025
self_collision_cfg:
weight: 5000.0
classify: False
lbfgs:
n_iters: 400 # 400
inner_iters: 25
cold_start_n_iters: 200
min_iters: 25
line_search_scale: [0.01, 0.3, 0.7,1.0] # #
fixed_iters: True
cost_convergence: 0.01
cost_delta_threshold: 1.0
cost_relative_threshold: 0.999 #0.999
epsilon: 0.01
history: 15 #15
use_cuda_graph: True
n_problems: 1
store_debug: False
use_cuda_kernel: True
stable_mode: True
line_search_type: "approx_wolfe"
use_cuda_line_search_kernel: True
use_cuda_update_best_kernel: True
use_temporal_smooth: False
sync_cuda_time: True
step_scale: 1.0
last_best: 10
use_coo_sparse: True
debug_info:
visual_traj : null #'ee_pos_seq'

5
src/curobo/content/configs/task/gradient_ik.yml
View File

@@ -32,6 +32,7 @@ cost:
vec_convergence: [0.0, 0.00] vec_convergence: [0.0, 0.00]
terminal: False terminal: False
use_metric: True use_metric: True
run_weight: 1.0
link_pose_cfg: link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [1000, 20000, 30, 50] weight: [1000, 20000, 30, 50]
@@ -58,7 +59,7 @@ cost:
lbfgs: lbfgs:
n_iters: 100 #60 n_iters: 100 #60
inner_iters: 25 inner_iters: 25
cold_start_n_iters: 100 cold_start_n_iters: null
min_iters: 20 min_iters: 20
line_search_scale: [0.01, 0.3, 0.7, 1.0] line_search_scale: [0.01, 0.3, 0.7, 1.0]
fixed_iters: True fixed_iters: True
@@ -69,7 +70,7 @@ lbfgs:
history: 6 history: 6
horizon: 1 horizon: 1
use_cuda_graph: True use_cuda_graph: True
n_envs: 1 n_problems: 1
store_debug: False store_debug: False
use_cuda_kernel: True use_cuda_kernel: True
stable_mode: True stable_mode: True

2
src/curobo/content/configs/task/gradient_mpc.yml
View File

@@ -83,7 +83,7 @@ lbfgs:
epsilon: 0.01 epsilon: 0.01
history: 15 #15 history: 15 #15
use_cuda_graph: True use_cuda_graph: True
n_envs: 1 n_problems: 1
store_debug: False store_debug: False
use_cuda_kernel: True use_cuda_kernel: True
stable_mode: True stable_mode: True

14
src/curobo/content/configs/task/gradient_trajopt.yml
View File

@@ -31,11 +31,11 @@ model:
cost: cost:
pose_cfg: pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] # orientation, position for all timesteps
run_vec_weight: [1.00,1.00,1.00,1.0,1.0,1.0] # running weight orientation, position 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,60] #[150.0, 2000.0, 30, 40]
vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True terminal: True
run_weight: 0.00 run_weight: 1.0
use_metric: True use_metric: True
link_pose_cfg: link_pose_cfg:
@@ -44,7 +44,7 @@ cost:
weight: [2000,50000.0,30,50] #[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 vec_convergence: [0.0,0.0] # orientation, position, orientation metric activation, position metric activation
terminal: True terminal: True
run_weight: 0.0 run_weight: 0.001
use_metric: True use_metric: True
cspace_cfg: cspace_cfg:
@@ -54,7 +54,9 @@ cost:
bound_cfg: bound_cfg:
weight: [5000.0, 50000.0, 50000.0,50000.0] # needs to be 3 values weight: [5000.0, 50000.0, 50000.0,50000.0] # needs to be 3 values
#weight: [000.0, 000.0, 000.0,000.0]
smooth_weight: [0.0,10000.0,10.0, 0.0] # [vel, acc, jerk, alpha_v-not-used] smooth_weight: [0.0,10000.0,10.0, 0.0] # [vel, acc, jerk, alpha_v-not-used]
#smooth_weight: [0.0,0000.0,0.0, 0.0] # [vel, acc, jerk, alpha_v-not-used]
run_weight_velocity: 0.00 run_weight_velocity: 0.00
run_weight_acceleration: 1.0 run_weight_acceleration: 1.0
run_weight_jerk: 1.0 run_weight_jerk: 1.0
@@ -79,9 +81,9 @@ cost:
lbfgs: lbfgs:
n_iters: 175 #175 n_iters: 125 #175
inner_iters: 25 inner_iters: 25
cold_start_n_iters: 150 cold_start_n_iters: null
min_iters: 25 min_iters: 25
line_search_scale: [0.01,0.3,0.7,1.0] #[0.01,0.2, 0.3,0.5,0.7,0.9, 1.0] # line_search_scale: [0.01,0.3,0.7,1.0] #[0.01,0.2, 0.3,0.5,0.7,0.9, 1.0] #
fixed_iters: True fixed_iters: True
@@ -91,7 +93,7 @@ lbfgs:
epsilon: 0.01 epsilon: 0.01
history: 15 history: 15
use_cuda_graph: True use_cuda_graph: True
n_envs: 1 n_problems: 1
store_debug: False store_debug: False
use_cuda_kernel: True use_cuda_kernel: True
stable_mode: True stable_mode: True

7
src/curobo/content/configs/task/particle_ik.yml
View File

@@ -31,6 +31,7 @@ cost:
vec_convergence: [0.00, 0.000] # orientation, position vec_convergence: [0.00, 0.000] # orientation, position
terminal: False terminal: False
use_metric: True use_metric: True
run_weight: 1.0
link_pose_cfg: link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
weight: [30, 50, 10, 10] #[20.0, 100.0] weight: [30, 50, 10, 10] #[20.0, 100.0]
@@ -56,7 +57,7 @@ mppi:
init_cov : 1.0 #0.15 #.5 #.5 init_cov : 1.0 #0.15 #.5 #.5
gamma : 1.0 gamma : 1.0
n_iters : 4 n_iters : 4
cold_start_n_iters: 4 cold_start_n_iters: null
step_size_mean : 0.9 step_size_mean : 0.9
step_size_cov : 0.1 step_size_cov : 0.1
beta : 0.01 beta : 0.01
@@ -69,12 +70,12 @@ mppi:
sample_mode : 'BEST' sample_mode : 'BEST'
base_action : 'REPEAT' base_action : 'REPEAT'
squash_fn : 'CLAMP' squash_fn : 'CLAMP'
n_envs : 1 n_problems : 1
use_cuda_graph : True use_cuda_graph : True
seed : 0 seed : 0
store_debug : False store_debug : False
random_mean : False random_mean : False
sample_per_env : True sample_per_problem: True
sync_cuda_time : False sync_cuda_time : False
use_coo_sparse : True use_coo_sparse : True
sample_params: sample_params:

4
src/curobo/content/configs/task/particle_mpc.yml
View File

@@ -89,12 +89,12 @@ mppi:
sample_mode : 'BEST' sample_mode : 'BEST'
base_action : 'REPEAT' base_action : 'REPEAT'
squash_fn : 'CLAMP' squash_fn : 'CLAMP'
n_envs : 1 n_problems : 1
use_cuda_graph : True use_cuda_graph : True
seed : 0 seed : 0
store_debug : False store_debug : False
random_mean : True random_mean : True
sample_per_env : False sample_per_problem: False
sync_cuda_time : True sync_cuda_time : True
use_coo_sparse : True use_coo_sparse : True
sample_params: sample_params:

16
src/curobo/content/configs/task/particle_trajopt.yml
View File

@@ -30,17 +30,17 @@ model:
cost: cost:
pose_cfg: pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
run_vec_weight: [1.0,1.0,1.0,1.0,1.0,1.0] # running weight run_vec_weight: [0.0,0.0,0.0,0.0,0.0,0.0] # running weight
weight: [250.0, 5000.0, 20, 20] weight: [250.0, 5000.0, 20, 20]
vec_convergence: [0.0,0.0,1000.0,1000.0] vec_convergence: [0.0,0.0,1000.0,1000.0]
terminal: True terminal: True
run_weight: 0.00 run_weight: 1.0
use_metric: True use_metric: True
link_pose_cfg: link_pose_cfg:
vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] vec_weight: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
run_vec_weight: [1.0,1.0,1.0,1.0,1.0,1.0] # running weight run_vec_weight: [0.00,0.00,0.00,0.0,0.0,0.0] # running weight orientation, position
weight: [250.0, 5000.0, 40, 40] weight: [0.0, 5000.0, 40, 40]
vec_convergence: [0.0,0.0,1000.0,1000.0] vec_convergence: [0.0,0.0,1000.0,1000.0]
terminal: True terminal: True
run_weight: 0.00 run_weight: 0.00
@@ -54,10 +54,12 @@ cost:
bound_cfg: bound_cfg:
weight: [0.1, 0.1,0.0,0.0] weight: [0.1, 0.1,0.0,0.0]
activation_distance: [0.0,0.0,0.0,0.0] #-#0.01 activation_distance: [0.0,0.0,0.0,0.0] #-#0.01
#smooth_weight: [0.0,100.0,1.0,0.0]
smooth_weight: [0.0,20.0,0.0,0.0] smooth_weight: [0.0,20.0,0.0,0.0]
run_weight_velocity: 0.0 run_weight_velocity: 0.0
run_weight_acceleration: 1.0 run_weight_acceleration: 1.0
run_weight_jerk: 1.0 run_weight_jerk: 1.0
null_space_weight: [0.00]
primitive_collision_cfg: primitive_collision_cfg:
weight: 5000.0 weight: 5000.0
@@ -80,7 +82,7 @@ mppi:
init_cov : 0.5 init_cov : 0.5
gamma : 1.0 gamma : 1.0
n_iters : 2 n_iters : 2
cold_start_n_iters: 2 cold_start_n_iters: null
step_size_mean : 0.9 step_size_mean : 0.9
step_size_cov : 0.01 step_size_cov : 0.01
beta : 0.01 beta : 0.01
@@ -93,12 +95,12 @@ mppi:
sample_mode : 'BEST' sample_mode : 'BEST'
base_action : 'REPEAT' base_action : 'REPEAT'
squash_fn : 'CLAMP' squash_fn : 'CLAMP'
n_envs : 1 n_problems : 1
use_cuda_graph : True use_cuda_graph : True
seed : 0 seed : 0
store_debug : False store_debug : False
random_mean : True random_mean : True
sample_per_env : True sample_per_problem: True
sync_cuda_time : False sync_cuda_time : False
use_coo_sparse : True use_coo_sparse : True
sample_params: sample_params:

8
src/curobo/cuda_robot_model/cuda_robot_generator.py
View File

@@ -145,6 +145,9 @@ class CudaRobotGeneratorConfig:
#: cspace config #: cspace config
cspace: Union[None, CSpaceConfig, Dict[str, List[Any]]] = None cspace: Union[None, CSpaceConfig, Dict[str, List[Any]]] = None
#: Enable loading meshes from kinematics parser.
load_meshes: bool = False
def __post_init__(self): def __post_init__(self):
# add root path: # add root path:
# Check if an external asset path is provided: # Check if an external asset path is provided:
@@ -283,7 +286,10 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
) )
else: else:
self._kinematics_parser = UrdfKinematicsParser( self._kinematics_parser = UrdfKinematicsParser(
self.urdf_path, mesh_root=self.asset_root_path, extra_links=self.extra_links self.urdf_path,
mesh_root=self.asset_root_path,
extra_links=self.extra_links,
load_meshes=self.load_meshes,
) )
if self.lock_joints is None: if self.lock_joints is None:

8
src/curobo/cuda_robot_model/cuda_robot_model.py
View File

@@ -30,7 +30,7 @@ from curobo.cuda_robot_model.types import (
SelfCollisionKinematicsConfig, SelfCollisionKinematicsConfig,
) )
from curobo.curobolib.kinematics import get_cuda_kinematics from curobo.curobolib.kinematics import get_cuda_kinematics
from curobo.geom.types import Sphere from curobo.geom.types import Mesh, Sphere
from curobo.types.base import TensorDeviceType from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose from curobo.types.math import Pose
from curobo.util.logger import log_error from curobo.util.logger import log_error
@@ -390,4 +390,8 @@ class CudaRobotModel(CudaRobotModelConfig):
return self.kinematics_config.base_link return self.kinematics_config.base_link
def update_kinematics_config(self, new_kin_config: KinematicsTensorConfig): def update_kinematics_config(self, new_kin_config: KinematicsTensorConfig):
self.kinematics_config.copy(new_kin_config) self.kinematics_config.copy_(new_kin_config)
@property
def retract_config(self):
return self.kinematics_config.cspace.retract_config

95
src/curobo/cuda_robot_model/types.py
View File

@@ -23,7 +23,7 @@ from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose from curobo.types.math import Pose
from curobo.types.state import JointState from curobo.types.state import JointState
from curobo.types.tensor import T_DOF from curobo.types.tensor import T_DOF
from curobo.util.tensor_util import copy_if_not_none from curobo.util.tensor_util import clone_if_not_none, copy_if_not_none
class JointType(Enum): class JointType(Enum):
@@ -75,20 +75,28 @@ class JointLimits:
self.tensor_args, self.tensor_args,
) )
def copy_(self, new_jl: JointLimits):
self.joint_names = new_jl.joint_names.copy()
self.position.copy_(new_jl.position)
self.velocity.copy_(new_jl.velocity)
self.acceleration.copy_(new_jl.acceleration)
self.effort = copy_if_not_none(new_jl.effort, self.effort)
return self
@dataclass @dataclass
class CSpaceConfig: class CSpaceConfig:
joint_names: List[str] joint_names: List[str]
retract_config: Optional[T_DOF] = None retract_config: Optional[T_DOF] = None
cspace_distance_weight: Optional[T_DOF] = None cspace_distance_weight: Optional[T_DOF] = None
null_space_weight: Optional[T_DOF] = None # List[float] null_space_weight: Optional[T_DOF] = None
tensor_args: TensorDeviceType = TensorDeviceType() tensor_args: TensorDeviceType = TensorDeviceType()
max_acceleration: Union[float, List[float]] = 10.0 max_acceleration: Union[float, List[float]] = 10.0
max_jerk: Union[float, List[float]] = 500.0 max_jerk: Union[float, List[float]] = 500.0
velocity_scale: Union[float, List[float]] = 1.0 velocity_scale: Union[float, List[float]] = 1.0
acceleration_scale: Union[float, List[float]] = 1.0 acceleration_scale: Union[float, List[float]] = 1.0
jerk_scale: Union[float, List[float]] = 1.0 jerk_scale: Union[float, List[float]] = 1.0
position_limit_clip: Union[float, List[float]] = 0.01 position_limit_clip: Union[float, List[float]] = 0.0
def __post_init__(self): def __post_init__(self):
if self.retract_config is not None: if self.retract_config is not None:
@@ -147,12 +155,31 @@ class CSpaceConfig:
joint_names = [self.joint_names[n] for n in new_index] joint_names = [self.joint_names[n] for n in new_index]
self.joint_names = joint_names self.joint_names = joint_names
def copy_(self, new_config: CSpaceConfig):
self.joint_names = new_config.joint_names.copy()
self.retract_config = copy_if_not_none(new_config.retract_config, self.retract_config)
self.null_space_weight = copy_if_not_none(
new_config.null_space_weight, self.null_space_weight
)
self.cspace_distance_weight = copy_if_not_none(
new_config.cspace_distance_weight, self.cspace_distance_weight
)
self.tensor_args = self.tensor_args
self.max_jerk = copy_if_not_none(new_config.max_jerk, self.max_jerk)
self.max_acceleration = copy_if_not_none(new_config.max_acceleration, self.max_acceleration)
self.velocity_scale = copy_if_not_none(new_config.velocity_scale, self.velocity_scale)
self.acceleration_scale = copy_if_not_none(
new_config.acceleration_scale, self.acceleration_scale
)
self.jerk_scale = copy_if_not_none(new_config.jerk_scale, self.jerk_scale)
return self
def clone(self) -> CSpaceConfig: def clone(self) -> CSpaceConfig:
return CSpaceConfig( return CSpaceConfig(
joint_names=self.joint_names.copy(), joint_names=self.joint_names.copy(),
retract_config=copy_if_not_none(self.retract_config), retract_config=clone_if_not_none(self.retract_config),
null_space_weight=copy_if_not_none(self.null_space_weight), null_space_weight=clone_if_not_none(self.null_space_weight),
cspace_distance_weight=copy_if_not_none(self.cspace_distance_weight), cspace_distance_weight=clone_if_not_none(self.cspace_distance_weight),
tensor_args=self.tensor_args, tensor_args=self.tensor_args,
max_jerk=self.max_jerk.clone(), max_jerk=self.max_jerk.clone(),
max_acceleration=self.max_acceleration.clone(), max_acceleration=self.max_acceleration.clone(),
@@ -215,12 +242,48 @@ class KinematicsTensorConfig:
cspace: Optional[CSpaceConfig] = None cspace: Optional[CSpaceConfig] = None
base_link: str = "base_link" base_link: str = "base_link"
ee_link: str = "ee_link" ee_link: str = "ee_link"
reference_link_spheres: Optional[torch.Tensor] = None
def __post_init__(self): def __post_init__(self):
if self.cspace is None and self.joint_limits is not None: if self.cspace is None and self.joint_limits is not None:
self.load_cspace_cfg_from_kinematics() self.load_cspace_cfg_from_kinematics()
if self.joint_limits is not None and self.cspace is not None: if self.joint_limits is not None and self.cspace is not None:
self.joint_limits = self.cspace.scale_joint_limits(self.joint_limits) self.joint_limits = self.cspace.scale_joint_limits(self.joint_limits)
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):
self.fixed_transforms.copy_(new_config.fixed_transforms)
self.link_map.copy_(new_config.link_map)
self.joint_map.copy_(new_config.joint_map)
self.joint_map_type.copy_(new_config.joint_map_type)
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)
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:
self.link_sphere_idx_map.copy_(new_config.link_sphere_idx_map)
if new_config.link_name_to_idx_map is not None and self.link_name_to_idx_map is not None:
self.link_name_to_idx_map = new_config.link_name_to_idx_map.copy()
if (
new_config.reference_link_spheres is not None
and self.reference_link_spheres is not None
):
self.reference_link_spheres.copy_(new_config.reference_link_spheres)
self.base_link = new_config.base_link
self.ee_idx = new_config.ee_idx
self.ee_link = new_config.ee_link
self.debug = new_config.debug
self.cspace.copy_(new_config.cspace)
self.n_dof = new_config.n_dof
self.non_fixed_joint_names = new_config.non_fixed_joint_names
self.joint_names = new_config.joint_names
self.link_names = new_config.link_names
self.mesh_link_names = new_config.mesh_link_names
self.total_spheres = new_config.total_spheres
return self
def load_cspace_cfg_from_kinematics(self): def load_cspace_cfg_from_kinematics(self):
retract_config = ( retract_config = (
@@ -270,6 +333,13 @@ class KinematicsTensorConfig:
link_sphere_index = self.get_sphere_index_from_link_name(link_name) link_sphere_index = self.get_sphere_index_from_link_name(link_name)
return self.link_spheres[link_sphere_index, :] return self.link_spheres[link_sphere_index, :]
def get_reference_link_spheres(
self,
link_name: str,
):
link_sphere_index = self.get_sphere_index_from_link_name(link_name)
return self.reference_link_spheres[link_sphere_index, :]
def attach_object( def attach_object(
self, self,
sphere_radius: Optional[float] = None, sphere_radius: Optional[float] = None,
@@ -332,6 +402,19 @@ class KinematicsTensorConfig:
""" """
return self.get_link_spheres(link_name).shape[0] return self.get_link_spheres(link_name).shape[0]
def disable_link_spheres(self, link_name: str):
if link_name not in self.link_name_to_idx_map.keys():
raise ValueError(link_name + " not found in spheres")
curr_spheres = self.get_link_spheres(link_name)
curr_spheres[:, 3] = -100.0
self.update_link_spheres(link_name, curr_spheres)
def enable_link_spheres(self, link_name: str):
if link_name not in self.link_name_to_idx_map.keys():
raise ValueError(link_name + " not found in spheres")
curr_spheres = self.get_reference_link_spheres(link_name)
self.update_link_spheres(link_name, curr_spheres)
@dataclass @dataclass
class SelfCollisionKinematicsConfig: class SelfCollisionKinematicsConfig:

1
src/curobo/cuda_robot_model/urdf_kinematics_parser.py
View File

@@ -183,6 +183,7 @@ class UrdfKinematicsParser(KinematicsParser):
else: else:
# convert to list: # convert to list:
mesh_pose = Pose.from_matrix(mesh_pose).to_list() mesh_pose = Pose.from_matrix(mesh_pose).to_list()
return CuroboMesh( return CuroboMesh(
name=link_name, name=link_name,
pose=mesh_pose, pose=mesh_pose,

10
src/curobo/cuda_robot_model/usd_kinematics_parser.py
View File

@@ -13,7 +13,6 @@ from typing import Dict, List, Optional, Tuple
# Third Party # Third Party
import numpy as np import numpy as np
from pxr import Usd, UsdPhysics
# CuRobo # CuRobo
from curobo.cuda_robot_model.kinematics_parser import KinematicsParser, LinkParams from curobo.cuda_robot_model.kinematics_parser import KinematicsParser, LinkParams
@@ -22,6 +21,15 @@ from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose from curobo.types.math import Pose
from curobo.util.logger import log_error from curobo.util.logger import log_error
try:
# Third Party
from pxr import Usd, UsdPhysics
except ImportError:
raise ImportError(
"usd-core failed to import, install with pip install usd-core"
+ " NOTE: Do not install this if using with ISAAC SIM."
)
class UsdKinematicsParser(KinematicsParser): class UsdKinematicsParser(KinematicsParser):
"""An experimental kinematics parser from USD. """An experimental kinematics parser from USD.

127
src/curobo/curobolib/cpp/geom_cuda.cpp
View File

@@ -16,43 +16,55 @@
// CUDA forward declarations // CUDA forward declarations
std::vector<torch::Tensor> self_collision_distance( std::vector<torch::Tensor>self_collision_distance(
torch::Tensor out_distance, torch::Tensor out_vec, torch::Tensor out_distance,
torch::Tensor out_vec,
torch::Tensor sparse_index, torch::Tensor sparse_index,
const torch::Tensor robot_spheres, // batch_size x n_spheres x 4 const torch::Tensor robot_spheres, // batch_size x n_spheres x 4
const torch::Tensor collision_offset, // n_spheres x n_spheres const torch::Tensor collision_offset, // n_spheres x n_spheres
const torch::Tensor weight, const torch::Tensor weight,
const torch::Tensor collision_matrix, // n_spheres x n_spheres const torch::Tensor collision_matrix, // n_spheres x n_spheres
const torch::Tensor thread_locations, const int locations_size, const torch::Tensor thread_locations,
const int batch_size, const int nspheres, const bool compute_grad = false, const int locations_size,
const int batch_size,
const int nspheres,
const bool compute_grad = false,
const int ndpt = 8, // Does this need to match template? const int ndpt = 8, // Does this need to match template?
const bool debug = false); const bool debug = false);
// CUDA forward declarations // CUDA forward declarations
std::vector<torch::Tensor> swept_sphere_obb_clpt( std::vector<torch::Tensor>swept_sphere_obb_clpt(
const torch::Tensor sphere_position, // batch_size, 3 const torch::Tensor sphere_position, // batch_size, 3
torch::Tensor distance, // batch_size, 1 torch::Tensor distance, // batch_size, 1
torch::Tensor torch::Tensor
closest_point, // batch size, 4 -> written out as x,y,z,0 for gradient closest_point, // batch size, 4 -> written out as x,y,z,0 for gradient
torch::Tensor sparsity_idx, const torch::Tensor weight, torch::Tensor sparsity_idx,
const torch::Tensor activation_distance, const torch::Tensor speed_dt, const torch::Tensor weight,
const torch::Tensor activation_distance,
const torch::Tensor speed_dt,
const torch::Tensor obb_accel, // n_boxes, 4, 4 const torch::Tensor obb_accel, // n_boxes, 4, 4
const torch::Tensor obb_bounds, // n_boxes, 3 const torch::Tensor obb_bounds, // n_boxes, 3
const torch::Tensor obb_pose, // n_boxes, 4, 4 const torch::Tensor obb_pose, // n_boxes, 4, 4
const torch::Tensor obb_enable, // n_boxes, 4, const torch::Tensor obb_enable, // n_boxes, 4,
const torch::Tensor n_env_obb, // n_boxes, 4, 4 const torch::Tensor n_env_obb, // n_boxes, 4, 4
const torch::Tensor env_query_idx, // n_boxes, 4, 4 const torch::Tensor env_query_idx, // n_boxes, 4, 4
const int max_nobs, const int batch_size, const int horizon, const int max_nobs,
const int n_spheres, const int sweep_steps, const bool enable_speed_metric, const int batch_size,
const bool transform_back, const bool compute_distance, const int horizon,
const int n_spheres,
const int sweep_steps,
const bool enable_speed_metric,
const bool transform_back,
const bool compute_distance,
const bool use_batch_env); const bool use_batch_env);
std::vector<torch::Tensor> std::vector<torch::Tensor>
sphere_obb_clpt(const torch::Tensor sphere_position, // batch_size, 4 sphere_obb_clpt(const torch::Tensor sphere_position, // batch_size, 4
torch::Tensor distance, torch::Tensor distance,
torch::Tensor closest_point, // batch size, 3 torch::Tensor closest_point, // batch size, 3
torch::Tensor sparsity_idx, const torch::Tensor weight, torch::Tensor sparsity_idx,
const torch::Tensor weight,
const torch::Tensor activation_distance, const torch::Tensor activation_distance,
const torch::Tensor obb_accel, // n_boxes, 4, 4 const torch::Tensor obb_accel, // n_boxes, 4, 4
const torch::Tensor obb_bounds, // n_boxes, 3 const torch::Tensor obb_bounds, // n_boxes, 3
@@ -60,48 +72,65 @@ sphere_obb_clpt(const torch::Tensor sphere_position, // batch_size, 4
const torch::Tensor obb_enable, // n_boxes, 4, 4 const torch::Tensor obb_enable, // n_boxes, 4, 4
const torch::Tensor n_env_obb, // n_boxes, 4, 4 const torch::Tensor n_env_obb, // n_boxes, 4, 4
const torch::Tensor env_query_idx, // n_boxes, 4, 4 const torch::Tensor env_query_idx, // n_boxes, 4, 4
const int max_nobs, const int batch_size, const int horizon, const int max_nobs,
const int n_spheres, const bool transform_back, const int batch_size,
const bool compute_distance, const bool use_batch_env); const int horizon,
const int n_spheres,
const bool transform_back,
const bool compute_distance,
const bool use_batch_env);
std::vector<torch::Tensor> pose_distance( std::vector<torch::Tensor>pose_distance(
torch::Tensor out_distance, torch::Tensor out_position_distance, torch::Tensor out_distance,
torch::Tensor out_position_distance,
torch::Tensor out_rotation_distance, torch::Tensor out_rotation_distance,
torch::Tensor distance_p_vector, // batch size, 3 torch::Tensor distance_p_vector, // batch size, 3
torch::Tensor distance_q_vector, // batch size, 4 torch::Tensor distance_q_vector, // batch size, 4
torch::Tensor out_gidx, torch::Tensor out_gidx,
const torch::Tensor current_position, // batch_size, 3 const torch::Tensor current_position, // batch_size, 3
const torch::Tensor goal_position, // n_boxes, 3 const torch::Tensor goal_position, // n_boxes, 3
const torch::Tensor current_quat, const torch::Tensor goal_quat, const torch::Tensor current_quat,
const torch::Tensor goal_quat,
const torch::Tensor vec_weight, // n_boxes, 4, 4 const torch::Tensor vec_weight, // n_boxes, 4, 4
const torch::Tensor weight, // n_boxes, 4, 4 const torch::Tensor weight, // n_boxes, 4, 4
const torch::Tensor vec_convergence, const torch::Tensor run_weight, const torch::Tensor vec_convergence,
const torch::Tensor run_vec_weight, const torch::Tensor batch_pose_idx, const torch::Tensor run_weight,
const int batch_size, const int horizon, const int mode, const torch::Tensor run_vec_weight,
const int num_goals = 1, const bool compute_grad = false, const torch::Tensor offset_waypoint,
const bool write_distance = true, const bool use_metric = false); const torch::Tensor offset_tstep_fraction,
const torch::Tensor batch_pose_idx,
const int batch_size,
const int horizon,
const int mode,
const int num_goals = 1,
const bool compute_grad = false,
const bool write_distance = true,
const bool use_metric = false,
const bool project_distance = true);
std::vector<torch::Tensor> std::vector<torch::Tensor>
backward_pose_distance(torch::Tensor out_grad_p, torch::Tensor out_grad_q, backward_pose_distance(torch::Tensor out_grad_p,
torch::Tensor out_grad_q,
const torch::Tensor grad_distance, // batch_size, 3 const torch::Tensor grad_distance, // batch_size, 3
const torch::Tensor grad_p_distance, // n_boxes, 3 const torch::Tensor grad_p_distance, // n_boxes, 3
const torch::Tensor grad_q_distance, const torch::Tensor grad_q_distance,
const torch::Tensor pose_weight, const torch::Tensor pose_weight,
const torch::Tensor grad_p_vec, // n_boxes, 4, 4 const torch::Tensor grad_p_vec, // n_boxes, 4, 4
const torch::Tensor grad_q_vec, const int batch_size, const torch::Tensor grad_q_vec,
const int batch_size,
const bool use_distance = false); const bool use_distance = false);
// C++ interface // C++ interface
// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. // 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_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) \ #define CHECK_CONTIGUOUS(x) \
AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") AT_ASSERTM(x.is_contiguous(), # x " must be contiguous")
#define CHECK_INPUT(x) \ #define CHECK_INPUT(x) \
CHECK_CUDA(x); \ CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x) CHECK_CONTIGUOUS(x)
std::vector<torch::Tensor> self_collision_distance_wrapper( std::vector<torch::Tensor>self_collision_distance_wrapper(
torch::Tensor out_distance, torch::Tensor out_vec, torch::Tensor out_distance, torch::Tensor out_vec,
torch::Tensor sparse_index, torch::Tensor sparse_index,
const torch::Tensor robot_spheres, // batch_size x n_spheres x 4 const torch::Tensor robot_spheres, // batch_size x n_spheres x 4
@@ -110,8 +139,8 @@ std::vector<torch::Tensor> self_collision_distance_wrapper(
const torch::Tensor collision_matrix, // n_spheres const torch::Tensor collision_matrix, // n_spheres
const torch::Tensor thread_locations, const int thread_locations_size, const torch::Tensor thread_locations, const int thread_locations_size,
const int batch_size, const int nspheres, const bool compute_grad = false, const int batch_size, const int nspheres, const bool compute_grad = false,
const int ndpt = 8, const bool debug = false) { const int ndpt = 8, const bool debug = false)
{
CHECK_INPUT(out_distance); CHECK_INPUT(out_distance);
CHECK_INPUT(out_vec); CHECK_INPUT(out_vec);
CHECK_INPUT(robot_spheres); CHECK_INPUT(robot_spheres);
@@ -128,7 +157,7 @@ std::vector<torch::Tensor> self_collision_distance_wrapper(
thread_locations_size, batch_size, nspheres, compute_grad, ndpt, debug); thread_locations_size, batch_size, nspheres, compute_grad, ndpt, debug);
} }
std::vector<torch::Tensor> sphere_obb_clpt_wrapper( std::vector<torch::Tensor>sphere_obb_clpt_wrapper(
const torch::Tensor sphere_position, // batch_size, 4 const torch::Tensor sphere_position, // batch_size, 4
torch::Tensor distance, torch::Tensor distance,
@@ -143,9 +172,10 @@ std::vector<torch::Tensor> sphere_obb_clpt_wrapper(
const torch::Tensor env_query_idx, // n_boxes, 4, 4 const torch::Tensor env_query_idx, // n_boxes, 4, 4
const int max_nobs, const int batch_size, const int horizon, const int max_nobs, const int batch_size, const int horizon,
const int n_spheres, const bool transform_back, const bool compute_distance, const int n_spheres, const bool transform_back, const bool compute_distance,
const bool use_batch_env) { const bool use_batch_env)
{
const at::cuda::OptionalCUDAGuard guard(sphere_position.device()); const at::cuda::OptionalCUDAGuard guard(sphere_position.device());
CHECK_INPUT(distance); CHECK_INPUT(distance);
CHECK_INPUT(closest_point); CHECK_INPUT(closest_point);
CHECK_INPUT(sphere_position); CHECK_INPUT(sphere_position);
@@ -159,7 +189,8 @@ std::vector<torch::Tensor> sphere_obb_clpt_wrapper(
n_env_obb, env_query_idx, max_nobs, batch_size, horizon, n_spheres, n_env_obb, env_query_idx, max_nobs, batch_size, horizon, n_spheres,
transform_back, compute_distance, use_batch_env); transform_back, compute_distance, use_batch_env);
} }
std::vector<torch::Tensor> swept_sphere_obb_clpt_wrapper(
std::vector<torch::Tensor>swept_sphere_obb_clpt_wrapper(
const torch::Tensor sphere_position, // batch_size, 4 const torch::Tensor sphere_position, // batch_size, 4
torch::Tensor distance, torch::Tensor distance,
torch::Tensor closest_point, // batch size, 3 torch::Tensor closest_point, // batch size, 3
@@ -174,8 +205,10 @@ std::vector<torch::Tensor> swept_sphere_obb_clpt_wrapper(
const int max_nobs, const int batch_size, const int horizon, const int max_nobs, const int batch_size, const int horizon,
const int n_spheres, const int sweep_steps, const bool enable_speed_metric, const int n_spheres, const int sweep_steps, const bool enable_speed_metric,
const bool transform_back, const bool compute_distance, const bool transform_back, const bool compute_distance,
const bool use_batch_env) { const bool use_batch_env)
{
const at::cuda::OptionalCUDAGuard guard(sphere_position.device()); const at::cuda::OptionalCUDAGuard guard(sphere_position.device());
CHECK_INPUT(distance); CHECK_INPUT(distance);
CHECK_INPUT(closest_point); CHECK_INPUT(closest_point);
CHECK_INPUT(sphere_position); CHECK_INPUT(sphere_position);
@@ -187,7 +220,8 @@ std::vector<torch::Tensor> swept_sphere_obb_clpt_wrapper(
env_query_idx, max_nobs, batch_size, horizon, n_spheres, sweep_steps, env_query_idx, max_nobs, batch_size, horizon, n_spheres, sweep_steps,
enable_speed_metric, transform_back, compute_distance, use_batch_env); enable_speed_metric, transform_back, compute_distance, use_batch_env);
} }
std::vector<torch::Tensor> pose_distance_wrapper(
std::vector<torch::Tensor>pose_distance_wrapper(
torch::Tensor out_distance, torch::Tensor out_position_distance, torch::Tensor out_distance, torch::Tensor out_position_distance,
torch::Tensor out_rotation_distance, torch::Tensor out_rotation_distance,
torch::Tensor distance_p_vector, // batch size, 3 torch::Tensor distance_p_vector, // batch size, 3
@@ -199,10 +233,12 @@ std::vector<torch::Tensor> pose_distance_wrapper(
const torch::Tensor vec_weight, // n_boxes, 4, 4 const torch::Tensor vec_weight, // n_boxes, 4, 4
const torch::Tensor weight, const torch::Tensor vec_convergence, const torch::Tensor weight, const torch::Tensor vec_convergence,
const torch::Tensor run_weight, const torch::Tensor run_vec_weight, const torch::Tensor run_weight, const torch::Tensor run_vec_weight,
const torch::Tensor offset_waypoint, const torch::Tensor offset_tstep_fraction,
const torch::Tensor batch_pose_idx, const int batch_size, const int horizon, const torch::Tensor batch_pose_idx, const int batch_size, const int horizon,
const int mode, const int num_goals = 1, const bool compute_grad = false, const int mode, const int num_goals = 1, const bool compute_grad = false,
const bool write_distance = false, const bool use_metric = false) { const bool write_distance = false, const bool use_metric = false,
const bool project_distance = true)
{
// at::cuda::DeviceGuard guard(angle.device()); // at::cuda::DeviceGuard guard(angle.device());
CHECK_INPUT(out_distance); CHECK_INPUT(out_distance);
CHECK_INPUT(out_position_distance); CHECK_INPUT(out_position_distance);
@@ -214,24 +250,30 @@ std::vector<torch::Tensor> pose_distance_wrapper(
CHECK_INPUT(current_quat); CHECK_INPUT(current_quat);
CHECK_INPUT(goal_quat); CHECK_INPUT(goal_quat);
CHECK_INPUT(batch_pose_idx); CHECK_INPUT(batch_pose_idx);
CHECK_INPUT(offset_waypoint);
CHECK_INPUT(offset_tstep_fraction);
const at::cuda::OptionalCUDAGuard guard(current_position.device()); const at::cuda::OptionalCUDAGuard guard(current_position.device());
return pose_distance( return pose_distance(
out_distance, out_position_distance, out_rotation_distance, out_distance, out_position_distance, out_rotation_distance,
distance_p_vector, distance_q_vector, out_gidx, current_position, distance_p_vector, distance_q_vector, out_gidx, current_position,
goal_position, current_quat, goal_quat, vec_weight, weight, goal_position, current_quat, goal_quat, vec_weight, weight,
vec_convergence, run_weight, run_vec_weight, batch_pose_idx, batch_size, vec_convergence, run_weight, run_vec_weight,
horizon, mode, num_goals, compute_grad, write_distance, use_metric); offset_waypoint,
offset_tstep_fraction,
batch_pose_idx, batch_size,
horizon, mode, num_goals, compute_grad, write_distance, use_metric, project_distance);
} }
std::vector<torch::Tensor> backward_pose_distance_wrapper( std::vector<torch::Tensor>backward_pose_distance_wrapper(
torch::Tensor out_grad_p, torch::Tensor out_grad_q, torch::Tensor out_grad_p, torch::Tensor out_grad_q,
const torch::Tensor grad_distance, // batch_size, 3 const torch::Tensor grad_distance, // batch_size, 3
const torch::Tensor grad_p_distance, // n_boxes, 3 const torch::Tensor grad_p_distance, // n_boxes, 3
const torch::Tensor grad_q_distance, const torch::Tensor pose_weight, const torch::Tensor grad_q_distance, const torch::Tensor pose_weight,
const torch::Tensor grad_p_vec, // n_boxes, 4, 4 const torch::Tensor grad_p_vec, // n_boxes, 4, 4
const torch::Tensor grad_q_vec, const int batch_size, const torch::Tensor grad_q_vec, const int batch_size,
const bool use_distance) { const bool use_distance)
{
CHECK_INPUT(out_grad_p); CHECK_INPUT(out_grad_p);
CHECK_INPUT(out_grad_q); CHECK_INPUT(out_grad_q);
CHECK_INPUT(grad_distance); CHECK_INPUT(grad_distance);
@@ -245,7 +287,8 @@ std::vector<torch::Tensor> backward_pose_distance_wrapper(
pose_weight, grad_p_vec, grad_q_vec, batch_size, use_distance); pose_weight, grad_p_vec, grad_q_vec, batch_size, use_distance);
} }
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
m.def("pose_distance", &pose_distance_wrapper, "Pose Distance (curobolib)"); m.def("pose_distance", &pose_distance_wrapper, "Pose Distance (curobolib)");
m.def("pose_distance_backward", &backward_pose_distance_wrapper, m.def("pose_distance_backward", &backward_pose_distance_wrapper,
"Pose Distance Backward (curobolib)"); "Pose Distance Backward (curobolib)");

75
src/curobo/curobolib/cpp/kinematics_fused_cuda.cpp
View File

@@ -18,39 +18,56 @@
std::vector<torch::Tensor> std::vector<torch::Tensor>
matrix_to_quaternion(torch::Tensor out_quat, matrix_to_quaternion(torch::Tensor out_quat,
const torch::Tensor in_rot // batch_size, 3 const torch::Tensor in_rot // batch_size, 3
); );
std::vector<torch::Tensor> kin_fused_forward( std::vector<torch::Tensor>kin_fused_forward(
torch::Tensor link_pos, torch::Tensor link_quat, torch::Tensor link_pos,
torch::Tensor batch_robot_spheres, torch::Tensor global_cumul_mat, torch::Tensor link_quat,
const torch::Tensor joint_vec, const torch::Tensor fixed_transform, torch::Tensor batch_robot_spheres,
const torch::Tensor robot_spheres, const torch::Tensor link_map, torch::Tensor global_cumul_mat,
const torch::Tensor joint_map, const torch::Tensor joint_map_type, const torch::Tensor joint_vec,
const torch::Tensor store_link_map, const torch::Tensor link_sphere_map, const torch::Tensor fixed_transform,
const int batch_size, const int n_spheres, 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 bool use_global_cumul = false);
std::vector<torch::Tensor> kin_fused_backward_16t( std::vector<torch::Tensor>kin_fused_backward_16t(
torch::Tensor grad_out, const torch::Tensor grad_nlinks_pos, torch::Tensor grad_out,
const torch::Tensor grad_nlinks_quat, const torch::Tensor grad_spheres, const torch::Tensor grad_nlinks_pos,
const torch::Tensor global_cumul_mat, const torch::Tensor joint_vec, const torch::Tensor grad_nlinks_quat,
const torch::Tensor fixed_transform, const torch::Tensor robot_spheres, const torch::Tensor grad_spheres,
const torch::Tensor link_map, const torch::Tensor joint_map, const torch::Tensor global_cumul_mat,
const torch::Tensor joint_map_type, const torch::Tensor store_link_map, const torch::Tensor joint_vec,
const torch::Tensor link_sphere_map, const torch::Tensor link_chain_map, const torch::Tensor fixed_transform,
const int batch_size, const int n_spheres, const bool sparsity_opt = true, 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 bool use_global_cumul = false);
// C++ interface // C++ interface
// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. // 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_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) \ #define CHECK_CONTIGUOUS(x) \
AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") AT_ASSERTM(x.is_contiguous(), # x " must be contiguous")
#define CHECK_INPUT(x) \ #define CHECK_INPUT(x) \
CHECK_CUDA(x); \ CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x) CHECK_CONTIGUOUS(x)
std::vector<torch::Tensor> kin_forward_wrapper( std::vector<torch::Tensor>kin_forward_wrapper(
torch::Tensor link_pos, torch::Tensor link_quat, torch::Tensor link_pos, torch::Tensor link_quat,
torch::Tensor batch_robot_spheres, torch::Tensor global_cumul_mat, torch::Tensor batch_robot_spheres, torch::Tensor global_cumul_mat,
const torch::Tensor joint_vec, const torch::Tensor fixed_transform, const torch::Tensor joint_vec, const torch::Tensor fixed_transform,
@@ -58,8 +75,8 @@ std::vector<torch::Tensor> kin_forward_wrapper(
const torch::Tensor joint_map, const torch::Tensor joint_map_type, 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 store_link_map, const torch::Tensor link_sphere_map,
const int batch_size, const int n_spheres, const int batch_size, const int n_spheres,
const bool use_global_cumul = false) { const bool use_global_cumul = false)
{
const at::cuda::OptionalCUDAGuard guard(joint_vec.device()); const at::cuda::OptionalCUDAGuard guard(joint_vec.device());
// TODO: add check input // TODO: add check input
@@ -69,7 +86,7 @@ std::vector<torch::Tensor> kin_forward_wrapper(
store_link_map, link_sphere_map, batch_size, n_spheres, use_global_cumul); store_link_map, link_sphere_map, batch_size, n_spheres, use_global_cumul);
} }
std::vector<torch::Tensor> kin_backward_wrapper( std::vector<torch::Tensor>kin_backward_wrapper(
torch::Tensor grad_out, const torch::Tensor grad_nlinks_pos, torch::Tensor grad_out, const torch::Tensor grad_nlinks_pos,
const torch::Tensor grad_nlinks_quat, const torch::Tensor grad_spheres, 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 global_cumul_mat, const torch::Tensor joint_vec,
@@ -78,7 +95,8 @@ std::vector<torch::Tensor> kin_backward_wrapper(
const torch::Tensor joint_map_type, const torch::Tensor store_link_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 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 int batch_size, const int n_spheres, const bool sparsity_opt = true,
const bool use_global_cumul = false) { const bool use_global_cumul = false)
{
const at::cuda::OptionalCUDAGuard guard(joint_vec.device()); const at::cuda::OptionalCUDAGuard guard(joint_vec.device());
return kin_fused_backward_16t( return kin_fused_backward_16t(
@@ -91,14 +109,17 @@ std::vector<torch::Tensor> kin_backward_wrapper(
std::vector<torch::Tensor> std::vector<torch::Tensor>
matrix_to_quaternion_wrapper(torch::Tensor out_quat, matrix_to_quaternion_wrapper(torch::Tensor out_quat,
const torch::Tensor in_rot // batch_size, 3 const torch::Tensor in_rot // batch_size, 3
) { )
{
const at::cuda::OptionalCUDAGuard guard(in_rot.device()); const at::cuda::OptionalCUDAGuard guard(in_rot.device());
CHECK_INPUT(in_rot); CHECK_INPUT(in_rot);
CHECK_INPUT(out_quat); CHECK_INPUT(out_quat);
return matrix_to_quaternion(out_quat, in_rot); return matrix_to_quaternion(out_quat, in_rot);
} }
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
m.def("forward", &kin_forward_wrapper, "Kinematics fused forward (CUDA)"); m.def("forward", &kin_forward_wrapper, "Kinematics fused forward (CUDA)");
m.def("backward", &kin_backward_wrapper, "Kinematics fused backward (CUDA)"); m.def("backward", &kin_backward_wrapper, "Kinematics fused backward (CUDA)");
m.def("matrix_to_quaternion", &matrix_to_quaternion_wrapper, m.def("matrix_to_quaternion", &matrix_to_quaternion_wrapper,

843
src/curobo/curobolib/cpp/kinematics_fused_kernel.cu
View File

File diff suppressed because it is too large Load Diff

75
src/curobo/curobolib/cpp/lbfgs_step_cuda.cpp
View File

@@ -15,35 +15,58 @@
#include <c10/cuda/CUDAGuard.h> #include <c10/cuda/CUDAGuard.h>
// CUDA forward declarations // CUDA forward declarations
std::vector<torch::Tensor> reduce_cuda(torch::Tensor vec, torch::Tensor vec2, std::vector<torch::Tensor>reduce_cuda(torch::Tensor vec,
torch::Tensor vec2,
torch::Tensor rho_buffer, torch::Tensor rho_buffer,
torch::Tensor sum, const int batch_size, torch::Tensor sum,
const int v_dim, const int m); const int batch_size,
const int v_dim,
const int m);
std::vector<torch::Tensor> std::vector<torch::Tensor>
lbfgs_step_cuda(torch::Tensor step_vec, torch::Tensor rho_buffer, lbfgs_step_cuda(torch::Tensor step_vec,
torch::Tensor y_buffer, torch::Tensor s_buffer, torch::Tensor rho_buffer,
torch::Tensor grad_q, const float epsilon, const int batch_size, torch::Tensor y_buffer,
const int m, const int v_dim); torch::Tensor s_buffer,
torch::Tensor grad_q,
const float epsilon,
const int batch_size,
const int m,
const int v_dim);
std::vector<torch::Tensor> std::vector<torch::Tensor>
lbfgs_update_cuda(torch::Tensor rho_buffer, torch::Tensor y_buffer, lbfgs_update_cuda(torch::Tensor rho_buffer,
torch::Tensor s_buffer, torch::Tensor q, torch::Tensor grad_q, torch::Tensor y_buffer,
torch::Tensor x_0, torch::Tensor grad_0, const int batch_size, torch::Tensor s_buffer,
const int m, const int v_dim); torch::Tensor q,
torch::Tensor grad_q,
torch::Tensor x_0,
torch::Tensor grad_0,
const int batch_size,
const int m,
const int v_dim);
std::vector<torch::Tensor> std::vector<torch::Tensor>
lbfgs_cuda_fuse(torch::Tensor step_vec, torch::Tensor rho_buffer, lbfgs_cuda_fuse(torch::Tensor step_vec,
torch::Tensor y_buffer, torch::Tensor s_buffer, torch::Tensor q, torch::Tensor rho_buffer,
torch::Tensor grad_q, torch::Tensor x_0, torch::Tensor grad_0, torch::Tensor y_buffer,
const float epsilon, const int batch_size, const int m, torch::Tensor s_buffer,
const int v_dim, const bool stable_mode); torch::Tensor q,
torch::Tensor grad_q,
torch::Tensor x_0,
torch::Tensor grad_0,
const float epsilon,
const int batch_size,
const int m,
const int v_dim,
const bool stable_mode);
// C++ interface // C++ interface
// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. // 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_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) \ #define CHECK_CONTIGUOUS(x) \
AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") AT_ASSERTM(x.is_contiguous(), # x " must be contiguous")
#define CHECK_INPUT(x) \ #define CHECK_INPUT(x) \
CHECK_CUDA(x); \ CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x) CHECK_CONTIGUOUS(x)
@@ -52,8 +75,8 @@ std::vector<torch::Tensor>
lbfgs_step_call(torch::Tensor step_vec, torch::Tensor rho_buffer, lbfgs_step_call(torch::Tensor step_vec, torch::Tensor rho_buffer,
torch::Tensor y_buffer, torch::Tensor s_buffer, torch::Tensor y_buffer, torch::Tensor s_buffer,
torch::Tensor grad_q, const float epsilon, const int batch_size, torch::Tensor grad_q, const float epsilon, const int batch_size,
const int m, const int v_dim) { const int m, const int v_dim)
{
CHECK_INPUT(step_vec); CHECK_INPUT(step_vec);
CHECK_INPUT(rho_buffer); CHECK_INPUT(rho_buffer);
CHECK_INPUT(y_buffer); CHECK_INPUT(y_buffer);
@@ -69,7 +92,8 @@ std::vector<torch::Tensor>
lbfgs_update_call(torch::Tensor rho_buffer, torch::Tensor y_buffer, lbfgs_update_call(torch::Tensor rho_buffer, torch::Tensor y_buffer,
torch::Tensor s_buffer, torch::Tensor q, torch::Tensor grad_q, torch::Tensor s_buffer, torch::Tensor q, torch::Tensor grad_q,
torch::Tensor x_0, torch::Tensor grad_0, const int batch_size, torch::Tensor x_0, torch::Tensor grad_0, const int batch_size,
const int m, const int v_dim) { const int m, const int v_dim)
{
CHECK_INPUT(rho_buffer); CHECK_INPUT(rho_buffer);
CHECK_INPUT(y_buffer); CHECK_INPUT(y_buffer);
CHECK_INPUT(s_buffer); CHECK_INPUT(s_buffer);
@@ -86,7 +110,8 @@ lbfgs_update_call(torch::Tensor rho_buffer, torch::Tensor y_buffer,
std::vector<torch::Tensor> std::vector<torch::Tensor>
reduce_cuda_call(torch::Tensor vec, torch::Tensor vec2, reduce_cuda_call(torch::Tensor vec, torch::Tensor vec2,
torch::Tensor rho_buffer, torch::Tensor sum, torch::Tensor rho_buffer, torch::Tensor sum,
const int batch_size, const int v_dim, const int m) { const int batch_size, const int v_dim, const int m)
{
const at::cuda::OptionalCUDAGuard guard(sum.device()); const at::cuda::OptionalCUDAGuard guard(sum.device());
return reduce_cuda(vec, vec2, rho_buffer, sum, batch_size, v_dim, m); return reduce_cuda(vec, vec2, rho_buffer, sum, batch_size, v_dim, m);
@@ -97,7 +122,8 @@ lbfgs_call(torch::Tensor step_vec, torch::Tensor rho_buffer,
torch::Tensor y_buffer, torch::Tensor s_buffer, torch::Tensor q, torch::Tensor y_buffer, torch::Tensor s_buffer, torch::Tensor q,
torch::Tensor grad_q, torch::Tensor x_0, torch::Tensor grad_0, torch::Tensor grad_q, torch::Tensor x_0, torch::Tensor grad_0,
const float epsilon, const int batch_size, const int m, const float epsilon, const int batch_size, const int m,
const int v_dim, const bool stable_mode) { const int v_dim, const bool stable_mode)
{
CHECK_INPUT(step_vec); CHECK_INPUT(step_vec);
CHECK_INPUT(rho_buffer); CHECK_INPUT(rho_buffer);
CHECK_INPUT(y_buffer); CHECK_INPUT(y_buffer);
@@ -113,7 +139,8 @@ lbfgs_call(torch::Tensor step_vec, torch::Tensor rho_buffer,
stable_mode); stable_mode);
} }
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
m.def("step", &lbfgs_step_call, "L-BFGS step (CUDA)"); m.def("step", &lbfgs_step_call, "L-BFGS step (CUDA)");
m.def("update", &lbfgs_update_call, "L-BFGS Update (CUDA)"); m.def("update", &lbfgs_update_call, "L-BFGS Update (CUDA)");
m.def("forward", &lbfgs_call, "L-BFGS Update + Step (CUDA)"); m.def("forward", &lbfgs_call, "L-BFGS Update + Step (CUDA)");

490
src/curobo/curobolib/cpp/lbfgs_step_kernel.cu
View File

File diff suppressed because it is too large Load Diff

56
src/curobo/curobolib/cpp/line_search_cuda.cpp
View File

@@ -18,49 +18,67 @@
// CUDA forward declarations // CUDA forward declarations
std::vector<torch::Tensor> std::vector<torch::Tensor>
update_best_cuda(torch::Tensor best_cost, torch::Tensor best_q, update_best_cuda(torch::Tensor best_cost,
torch::Tensor best_q,
torch::Tensor best_iteration, torch::Tensor best_iteration,
torch::Tensor current_iteration, torch::Tensor current_iteration,
const torch::Tensor cost, const torch::Tensor cost,
const torch::Tensor q, const int d_opt, const int cost_s1, const torch::Tensor q,
const int cost_s2, const int iteration, const int d_opt,
const int cost_s1,
const int cost_s2,
const int iteration,
const float delta_threshold, const float delta_threshold,
const float relative_threshold = 0.999); const float relative_threshold = 0.999);
std::vector<torch::Tensor> line_search_cuda( std::vector<torch::Tensor>line_search_cuda(
// torch::Tensor m, // torch::Tensor m,
torch::Tensor best_x, torch::Tensor best_c, torch::Tensor best_grad, torch::Tensor best_x,
const torch::Tensor g_x, const torch::Tensor x_set, torch::Tensor best_c,
const torch::Tensor step_vec, const torch::Tensor c_0, torch::Tensor best_grad,
const torch::Tensor alpha_list, const torch::Tensor c_idx, const float c_1, const torch::Tensor g_x,
const float c_2, const bool strong_wolfe, const bool approx_wolfe, const torch::Tensor x_set,
const int l1, const int l2, const int batchsize); const torch::Tensor step_vec,
const torch::Tensor c_0,
const torch::Tensor alpha_list,
const torch::Tensor c_idx,
const float c_1,
const float c_2,
const bool strong_wolfe,
const bool approx_wolfe,
const int l1,
const int l2,
const int batchsize);
// C++ interface // C++ interface
// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. // 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_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) \ #define CHECK_CONTIGUOUS(x) \
AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") AT_ASSERTM(x.is_contiguous(), # x " must be contiguous")
#define CHECK_INPUT(x) \ #define CHECK_INPUT(x) \
CHECK_CUDA(x); \ CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x) CHECK_CONTIGUOUS(x)
std::vector<torch::Tensor> line_search_call( std::vector<torch::Tensor>line_search_call(
// torch::Tensor m, // torch::Tensor m,
torch::Tensor best_x, torch::Tensor best_c, torch::Tensor best_grad, torch::Tensor best_x, torch::Tensor best_c, torch::Tensor best_grad,
const torch::Tensor g_x, const torch::Tensor x_set, const torch::Tensor g_x, const torch::Tensor x_set,
const torch::Tensor step_vec, const torch::Tensor c_0, const torch::Tensor step_vec, const torch::Tensor c_0,
const torch::Tensor alpha_list, const torch::Tensor c_idx, const float c_1, const torch::Tensor alpha_list, const torch::Tensor c_idx, const float c_1,
const float c_2, const bool strong_wolfe, const bool approx_wolfe, const float c_2, const bool strong_wolfe, const bool approx_wolfe,
const int l1, const int l2, const int batchsize) { const int l1, const int l2, const int batchsize)
{
CHECK_INPUT(g_x); CHECK_INPUT(g_x);
CHECK_INPUT(x_set); CHECK_INPUT(x_set);
CHECK_INPUT(step_vec); CHECK_INPUT(step_vec);
CHECK_INPUT(c_0); CHECK_INPUT(c_0);
CHECK_INPUT(alpha_list); CHECK_INPUT(alpha_list);
CHECK_INPUT(c_idx); CHECK_INPUT(c_idx);
// CHECK_INPUT(m); // CHECK_INPUT(m);
CHECK_INPUT(best_x); CHECK_INPUT(best_x);
CHECK_INPUT(best_c); CHECK_INPUT(best_c);
@@ -82,8 +100,8 @@ update_best_call(torch::Tensor best_cost, torch::Tensor best_q,
const torch::Tensor q, const int d_opt, const int cost_s1, const torch::Tensor q, const int d_opt, const int cost_s1,
const int cost_s2, const int iteration, const int cost_s2, const int iteration,
const float delta_threshold, const float delta_threshold,
const float relative_threshold=0.999) { const float relative_threshold = 0.999)
{
CHECK_INPUT(best_cost); CHECK_INPUT(best_cost);
CHECK_INPUT(best_q); CHECK_INPUT(best_q);
CHECK_INPUT(cost); CHECK_INPUT(cost);
@@ -96,8 +114,8 @@ update_best_call(torch::Tensor best_cost, torch::Tensor best_q,
cost_s1, cost_s2, iteration, delta_threshold, relative_threshold); cost_s1, cost_s2, iteration, delta_threshold, relative_threshold);
} }
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
m.def("update_best", &update_best_call, "Update Best (CUDA)"); m.def("update_best", &update_best_call, "Update Best (CUDA)");
m.def("line_search", &line_search_call, "Line search (CUDA)"); m.def("line_search", &line_search_call, "Line search (CUDA)");
} }

203
src/curobo/curobolib/cpp/line_search_kernel.cu
View File

@@ -37,55 +37,72 @@
#define FULL_MASK 0xffffffff #define FULL_MASK 0xffffffff
namespace Curobo { namespace Curobo
{
namespace Optimization { namespace Optimization
{
template <typename scalar_t, typename psum_t> template<typename scalar_t, typename psum_t>
__inline__ __device__ void reduce(scalar_t v, int m, unsigned mask, __inline__ __device__ void reduce(scalar_t v, int m, unsigned mask,
psum_t *data, scalar_t *result) { psum_t *data, scalar_t *result)
{
psum_t val = v; psum_t val = v;
val += __shfl_down_sync(mask, val, 1); val += __shfl_down_sync(mask, val, 1);
val += __shfl_down_sync(mask, val, 2); val += __shfl_down_sync(mask, val, 2);
val += __shfl_down_sync(mask, val, 4); val += __shfl_down_sync(mask, val, 4);
val += __shfl_down_sync(mask, val, 8); val += __shfl_down_sync(mask, val, 8);
val += __shfl_down_sync(mask, val, 16); val += __shfl_down_sync(mask, val, 16);
// int leader = __ffs(mask) – 1; // select a leader lane // int leader = __ffs(mask) – 1; // select a leader lane
int leader = 0; int leader = 0;
if (threadIdx.x % 32 == leader) {
if (m <= 32) { if (threadIdx.x % 32 == leader)
{
if (m <= 32)
{
result[0] = (scalar_t)val; result[0] = (scalar_t)val;
} else { }
else
{
data[(threadIdx.x + 1) / 32] = val; data[(threadIdx.x + 1) / 32] = val;
} }
} }
if (m > 32) {
if (m > 32)
{
__syncthreads(); __syncthreads();
int elems = (m + 31) / 32; int elems = (m + 31) / 32;
assert(elems <= 32); assert(elems <= 32);
unsigned mask2 = __ballot_sync(FULL_MASK, threadIdx.x < elems); unsigned mask2 = __ballot_sync(FULL_MASK, threadIdx.x < elems);
if (threadIdx.x < elems) { // only the first warp will do this work
if (threadIdx.x < elems) // only the first warp will do this work
{
psum_t val2 = data[threadIdx.x % 32]; psum_t val2 = data[threadIdx.x % 32];
int shift = 1; int shift = 1;
for (int i = elems - 1; i > 0; i /= 2) {
for (int i = elems - 1; i > 0; i /= 2)
{
val2 += __shfl_down_sync(mask2, val2, shift); val2 += __shfl_down_sync(mask2, val2, shift);
shift *= 2; shift *= 2;
} }
// int leader = __ffs(mask2) – 1; // select a leader lane // int leader = __ffs(mask2) – 1; // select a leader lane
int leader = 0; int leader = 0;
if (threadIdx.x % 32 == leader) {
if (threadIdx.x % 32 == leader)
{
result[0] = (scalar_t)val2; result[0] = (scalar_t)val2;
} }
} }
} }
__syncthreads(); __syncthreads();
} }
// Launched with l2 threads/block and batchsize blocks
template<typename scalar_t, typename psum_t>
__global__ void line_search_kernel(
// Launched with l2 threads/block and batchsize blocks
template <typename scalar_t, typename psum_t>
__global__ void line_search_kernel(
// int64_t *m_idx, // 4x1x1 // int64_t *m_idx, // 4x1x1
scalar_t *best_x, // 4x280 scalar_t *best_x, // 4x280
scalar_t *best_c, // 4x1 scalar_t *best_c, // 4x1
@@ -101,15 +118,16 @@ __global__ void line_search_kernel(
const int l1, // 6 const int l1, // 6
const int l2, // 280 const int l2, // 280
const int batchsize) // 4 const int batchsize) // 4
{ {
int batch = blockIdx.x; int batch = blockIdx.x;
__shared__ psum_t data[32]; __shared__ psum_t data[32];
__shared__ scalar_t result[32]; __shared__ scalar_t result[32];
assert(l1 <= 32); assert(l1 <= 32);
unsigned mask = __ballot_sync(FULL_MASK, threadIdx.x < l2); unsigned mask = __ballot_sync(FULL_MASK, threadIdx.x < l2);
if (threadIdx.x >= l2) { if (threadIdx.x >= l2)
{
return; return;
} }
@@ -117,7 +135,8 @@ __global__ void line_search_kernel(
// g_step = g0 @ step_vec_T // g_step = g0 @ step_vec_T
// g_x @ step_vec_T // g_x @ step_vec_T
for (int i = 0; i < l1; i++) { for (int i = 0; i < l1; i++)
{
reduce(g_x[batch * l1 * l2 + l2 * i + threadIdx.x] * sv_elem, l2, mask, reduce(g_x[batch * l1 * l2 + l2 * i + threadIdx.x] * sv_elem, l2, mask,
&data[0], &result[i]); &data[0], &result[i]);
} }
@@ -128,7 +147,9 @@ __global__ void line_search_kernel(
bool wolfe_1 = false; bool wolfe_1 = false;
bool wolfe = false; bool wolfe = false;
bool condition = threadIdx.x < l1; bool condition = threadIdx.x < l1;
if (condition) {
if (condition)
{
// scalar_t alpha_list_elem = alpha_list[batch*l1 + threadIdx.x]; // scalar_t alpha_list_elem = alpha_list[batch*l1 + threadIdx.x];
scalar_t alpha_list_elem = alpha_list[threadIdx.x]; scalar_t alpha_list_elem = alpha_list[threadIdx.x];
@@ -138,9 +159,13 @@ __global__ void line_search_kernel(
// condition 2: // condition 2:
bool wolfe_2; bool wolfe_2;
if (strong_wolfe) {
wolfe_2 = abs(result[threadIdx.x]) <= c_2 * abs(result[0]); if (strong_wolfe)
} else { {
wolfe_2 = abs(result[threadIdx.x]) <= c_2 *abs(result[0]);
}
else
{
wolfe_2 = result[threadIdx.x] >= c_2 * result[0]; wolfe_2 = result[threadIdx.x] >= c_2 * result[0];
} }
@@ -153,30 +178,40 @@ __global__ void line_search_kernel(
__syncthreads(); __syncthreads();
__shared__ int idx_shared; __shared__ int idx_shared;
if (threadIdx.x == 0) {
if (threadIdx.x == 0)
{
int m_id = 0; int m_id = 0;
int m1_id = 0; int m1_id = 0;
scalar_t max1 = step_success[0]; scalar_t max1 = step_success[0];
scalar_t max2 = step_success_w1[0]; scalar_t max2 = step_success_w1[0];
for (int i = 1; i < l1; i++) {
if (max1 < step_success[i]) { for (int i = 1; i < l1; i++)
{
if (max1 < step_success[i])
{
max1 = step_success[i]; max1 = step_success[i];
m_id = i; m_id = i;
} }
if (max2 < step_success_w1[i]) {
if (max2 < step_success_w1[i])
{
max2 = step_success_w1[i]; max2 = step_success_w1[i];
m1_id = i; m1_id = i;
} }
} }
if (!approx_wolfe) { if (!approx_wolfe)
{
// m_idx = torch.where(m_idx == 0, m1_idx, m_idx) // m_idx = torch.where(m_idx == 0, m1_idx, m_idx)
if (m_id == 0) { if (m_id == 0)
{
m_id = m1_id; m_id = m1_id;
} }
// m_idx[m_idx == 0] = 1 // m_idx[m_idx == 0] = 1
if (m_id == 0) { if (m_id == 0)
{
m_id = 1; m_id = 1;
} }
} }
@@ -191,21 +226,27 @@ __global__ void line_search_kernel(
// idx_shared contains index in l1 // idx_shared contains index in l1
// //
__syncthreads(); __syncthreads();
if (threadIdx.x < l2) {
if (threadIdx.x == 0) { if (threadIdx.x < l2)
{
if (threadIdx.x == 0)
{
// printf("block: %d, idx_shared: %d\n", batch, idx_shared); // printf("block: %d, idx_shared: %d\n", batch, idx_shared);
} }
best_x[batch * l2 + threadIdx.x] = x_set[idx_shared * l2 + threadIdx.x]; best_x[batch * l2 + threadIdx.x] = x_set[idx_shared * l2 + threadIdx.x];
best_grad[batch * l2 + threadIdx.x] = g_x[idx_shared * l2 + threadIdx.x]; best_grad[batch * l2 + threadIdx.x] = g_x[idx_shared * l2 + threadIdx.x];
} }
if (threadIdx.x == 0) {
if (threadIdx.x == 0)
{
best_c[batch] = c[idx_shared]; best_c[batch] = c[idx_shared];
} }
} }
// Launched with l2 threads/block and #blocks = batchsize
template<typename scalar_t, typename psum_t>
__global__ void line_search_kernel_mask(
// Launched with l2 threads/block and #blocks = batchsize
template <typename scalar_t, typename psum_t>
__global__ void line_search_kernel_mask(
// int64_t *m_idx, // 4x1x1 // int64_t *m_idx, // 4x1x1
scalar_t *best_x, // 4x280 scalar_t *best_x, // 4x280
scalar_t *best_c, // 4x1 scalar_t *best_c, // 4x1
@@ -221,15 +262,16 @@ __global__ void line_search_kernel_mask(
const int l1, // 6 const int l1, // 6
const int l2, // 280 const int l2, // 280
const int batchsize) // 4 const int batchsize) // 4
{ {
int batch = blockIdx.x; int batch = blockIdx.x;
__shared__ psum_t data[32]; __shared__ psum_t data[32];
__shared__ scalar_t result[32]; __shared__ scalar_t result[32];
assert(l1 <= 32); assert(l1 <= 32);
unsigned mask = __ballot_sync(FULL_MASK, threadIdx.x < l2); unsigned mask = __ballot_sync(FULL_MASK, threadIdx.x < l2);
if (threadIdx.x >= l2) { if (threadIdx.x >= l2)
{
return; return;
} }
@@ -237,7 +279,8 @@ __global__ void line_search_kernel_mask(
// g_step = g0 @ step_vec_T // g_step = g0 @ step_vec_T
// g_x @ step_vec_T // g_x @ step_vec_T
for (int i = 0; i < l1; i++) { for (int i = 0; i < l1; i++)
{
reduce(g_x[batch * l1 * l2 + l2 * i + threadIdx.x] * sv_elem, l2, mask, reduce(g_x[batch * l1 * l2 + l2 * i + threadIdx.x] * sv_elem, l2, mask,
&data[0], &result[i]); &data[0], &result[i]);
} }
@@ -248,7 +291,9 @@ __global__ void line_search_kernel_mask(
bool wolfe_1 = false; bool wolfe_1 = false;
bool wolfe = false; bool wolfe = false;
bool condition = threadIdx.x < l1; bool condition = threadIdx.x < l1;
if (condition) {
if (condition)
{
scalar_t alpha_list_elem = alpha_list[threadIdx.x]; scalar_t alpha_list_elem = alpha_list[threadIdx.x];
// scalar_t alpha_list_elem = alpha_list[batch*l1 + threadIdx.x]; // scalar_t alpha_list_elem = alpha_list[batch*l1 + threadIdx.x];
@@ -259,9 +304,13 @@ __global__ void line_search_kernel_mask(
// condition 2: // condition 2:
bool wolfe_2; bool wolfe_2;
if (strong_wolfe) {
wolfe_2 = abs(result[threadIdx.x]) <= c_2 * abs(result[0]); if (strong_wolfe)
} else { {
wolfe_2 = abs(result[threadIdx.x]) <= c_2 *abs(result[0]);
}
else
{
wolfe_2 = result[threadIdx.x] >= c_2 * result[0]; wolfe_2 = result[threadIdx.x] >= c_2 * result[0];
} }
@@ -286,22 +335,35 @@ __global__ void line_search_kernel_mask(
__syncthreads(); __syncthreads();
__shared__ int idx_shared; __shared__ int idx_shared;
if (threadIdx.x == 0) {
if (!approx_wolfe) { if (threadIdx.x == 0)
if (id == 32) { // msk is zero {
if (!approx_wolfe)
{
if (id == 32) // msk is zero
{
id = id1; id = id1;
} }
if (id == 0) { // bit 0 is set
if (id == 0) // bit 0 is set
{
id = id1; id = id1;
} }
if (id == 32) { // msk is zero
if (id == 32) // msk is zero
{
id = 1; id = 1;
} }
if (id == 0) {
if (id == 0)
{
id = 1; id = 1;
} }
} else { }
if (id == 32) { // msk is zero else
{
if (id == 32) // msk is zero
{
id = 0; id = 0;
} }
} }
@@ -349,28 +411,34 @@ __global__ void line_search_kernel_mask(
// one index per batch is computed // one index per batch is computed
//////////////////////////////////// ////////////////////////////////////
__syncthreads(); __syncthreads();
if (threadIdx.x < l2) {
if (threadIdx.x == 0) { if (threadIdx.x < l2)
{
if (threadIdx.x == 0)
{
// printf("block: %d, idx_shared: %d\n", batch, idx_shared); // printf("block: %d, idx_shared: %d\n", batch, idx_shared);
} }
best_x[batch * l2 + threadIdx.x] = x_set[idx_shared * l2 + threadIdx.x]; best_x[batch * l2 + threadIdx.x] = x_set[idx_shared * l2 + threadIdx.x];
best_grad[batch * l2 + threadIdx.x] = g_x[idx_shared * l2 + threadIdx.x]; best_grad[batch * l2 + threadIdx.x] = g_x[idx_shared * l2 + threadIdx.x];
} }
if (threadIdx.x == 0) {
if (threadIdx.x == 0)
{
best_c[batch] = c[idx_shared]; best_c[batch] = c[idx_shared];
} }
} }
} // namespace Optimization
} // namespace Optimization
} // namespace Curobo } // namespace Curobo
std::vector<torch::Tensor> line_search_cuda( std::vector<torch::Tensor>line_search_cuda(
// torch::Tensor m_idx, // torch::Tensor m_idx,
torch::Tensor best_x, torch::Tensor best_c, torch::Tensor best_grad, torch::Tensor best_x, torch::Tensor best_c, torch::Tensor best_grad,
const torch::Tensor g_x, const torch::Tensor x_set, const torch::Tensor g_x, const torch::Tensor x_set,
const torch::Tensor step_vec, const torch::Tensor c_0, const torch::Tensor step_vec, const torch::Tensor c_0,
const torch::Tensor alpha_list, const torch::Tensor c_idx, const float c_1, const torch::Tensor alpha_list, const torch::Tensor c_idx, const float c_1,
const float c_2, const bool strong_wolfe, const bool approx_wolfe, const float c_2, const bool strong_wolfe, const bool approx_wolfe,
const int l1, const int l2, const int batchsize) { const int l1, const int l2, const int batchsize)
{
using namespace Curobo::Optimization; using namespace Curobo::Optimization;
assert(l2 <= 1024); assert(l2 <= 1024);
@@ -383,7 +451,8 @@ std::vector<torch::Tensor> line_search_cuda(
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
g_x.scalar_type(), "line_search_cu", ([&] { g_x.scalar_type(), "line_search_cu", ([&] {
line_search_kernel_mask<scalar_t, scalar_t> line_search_kernel_mask<scalar_t, scalar_t>
<<<blocksPerGrid, threadsPerBlock, 0, stream>>>( << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
// m_idx.data_ptr<int>(), // m_idx.data_ptr<int>(),
best_x.data_ptr<scalar_t>(), best_c.data_ptr<scalar_t>(), best_x.data_ptr<scalar_t>(), best_c.data_ptr<scalar_t>(),
best_grad.data_ptr<scalar_t>(), g_x.data_ptr<scalar_t>(), best_grad.data_ptr<scalar_t>(), g_x.data_ptr<scalar_t>(),
@@ -394,5 +463,5 @@ std::vector<torch::Tensor> line_search_cuda(
})); }));
C10_CUDA_KERNEL_LAUNCH_CHECK(); C10_CUDA_KERNEL_LAUNCH_CHECK();
return {best_x, best_c, best_grad}; return { best_x, best_c, best_grad };
} }

1070
src/curobo/curobolib/cpp/pose_distance_kernel.cu
View File

File diff suppressed because it is too large Load Diff

325
src/curobo/curobolib/cpp/self_collision_kernel.cu
View File

@@ -24,26 +24,30 @@
#define GOALSET 2 #define GOALSET 2
#define BATCH_GOALSET 3 #define BATCH_GOALSET 3
namespace Curobo { namespace Curobo
{
namespace Geometry { namespace Geometry
{
template <typename scalar_t> __inline__ __device__ scalar_t relu(scalar_t var) { template<typename scalar_t>__inline__ __device__ scalar_t relu(scalar_t var)
{
if (var < 0) if (var < 0)
return 0; return 0;
else else
return var; return var;
} }
template <typename scalar_t> template<typename scalar_t>
__global__ void self_collision_distance_kernel( __global__ void self_collision_distance_kernel(
scalar_t *out_distance, // batch x 1 scalar_t *out_distance, // batch x 1
scalar_t *out_vec, // batch x nspheres x 4 scalar_t *out_vec, // batch x nspheres x 4
const scalar_t *robot_spheres, // batch x nspheres x 4 const scalar_t *robot_spheres, // batch x nspheres x 4
const scalar_t *collision_threshold, const int batch_size, const scalar_t *collision_threshold, const int batch_size,
const int nspheres, const scalar_t *weight, const bool write_grad = false) { const int nspheres, const scalar_t *weight, const bool write_grad = false)
{
const int batch_idx = blockDim.x * blockIdx.x + threadIdx.x; const int batch_idx = blockDim.x * blockIdx.x + threadIdx.x;
if (batch_idx >= batch_size) {
if (batch_idx >= batch_size)
{
return; return;
} }
float r_diff, distance; float r_diff, distance;
@@ -51,22 +55,33 @@ __global__ void self_collision_distance_kernel(
float3 sph1, sph2, dist_vec; float3 sph1, sph2, dist_vec;
int sph1_idx = -1; int sph1_idx = -1;
int sph2_idx = -1; int sph2_idx = -1;
// iterate over spheres: // iterate over spheres:
for (int i = 0; i < nspheres; i++) { for (int i = 0; i < nspheres; i++)
{
sph1 = *(float3 *)&robot_spheres[batch_idx * nspheres * 4 + i * 4]; sph1 = *(float3 *)&robot_spheres[batch_idx * nspheres * 4 + i * 4];
for (int j = i + 1; j < nspheres; j++) {
for (int j = i + 1; j < nspheres; j++)
{
r_diff = collision_threshold[i * nspheres + j]; r_diff = collision_threshold[i * nspheres + j];
if (isinf(r_diff)) {
if (isinf(r_diff))
{
continue; continue;
} }
sph2 = *(float3 *)&robot_spheres[batch_idx * nspheres * 4 + j * 4]; sph2 = *(float3 *)&robot_spheres[batch_idx * nspheres * 4 + j * 4];
// compute sphere distance: // compute sphere distance:
distance = relu(r_diff - length(sph1 - sph2)); distance = relu(r_diff - length(sph1 - sph2));
if (distance > max_penetration) {
if (distance > max_penetration)
{
max_penetration = distance; max_penetration = distance;
sph1_idx = i; sph1_idx = i;
sph2_idx = j; sph2_idx = j;
if (write_grad) {
if (write_grad)
{
dist_vec = (sph1 - sph2) / distance; dist_vec = (sph1 - sph2) / distance;
} }
} }
@@ -74,31 +89,33 @@ __global__ void self_collision_distance_kernel(
} }
// write out pose distance: // write out pose distance:
if (max_penetration > 0) { if (max_penetration > 0)
{
out_distance[batch_idx] = weight[0] * max_penetration; out_distance[batch_idx] = weight[0] * max_penetration;
if (write_grad) {
if (write_grad)
{
*(float3 *)&out_vec[batch_idx * nspheres * 4 + sph1_idx * 4] = *(float3 *)&out_vec[batch_idx * nspheres * 4 + sph1_idx * 4] =
weight[0] * -1 * dist_vec; weight[0] * -1 * dist_vec;
*(float3 *)&out_vec[batch_idx * nspheres * 4 + sph2_idx * 4] = *(float3 *)&out_vec[batch_idx * nspheres * 4 + sph2_idx * 4] =
weight[0] * dist_vec; weight[0] * dist_vec;
} }
} }
} }
typedef struct { typedef struct {
float d; float d;
int16_t i; int16_t i;
int16_t j; int16_t j;
} dist_t; } dist_t;
///////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////// // n warps per row
// n warps per row // ndpt rows per warp
// ndpt rows per warp ///////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////// template<typename scalar_t>
template <typename scalar_t> __global__ void self_collision_distance_kernel4(
__global__ void self_collision_distance_kernel4(
scalar_t *out_distance, // batch x 1 scalar_t *out_distance, // batch x 1
scalar_t *out_vec, // batch x nspheres x 3 scalar_t *out_vec, // batch x nspheres x 3
const scalar_t *robot_spheres, // batch x nspheres x 3 const scalar_t *robot_spheres, // batch x nspheres x 3
@@ -107,31 +124,36 @@ __global__ void self_collision_distance_kernel4(
const int ndpt, // number of distances to be computed per thread const int ndpt, // number of distances to be computed per thread
const int nwpr, // number of warps per row const int nwpr, // number of warps per row
const scalar_t *weight, uint8_t *sparse_index, const scalar_t *weight, uint8_t *sparse_index,
const bool write_grad = false) { const bool write_grad = false)
{
int batch_idx = blockIdx.x; int batch_idx = blockIdx.x;
int warp_idx = threadIdx.x / 32; int warp_idx = threadIdx.x / 32;
int i = ndpt * (warp_idx / nwpr); // starting row number for this warp int i = ndpt * (warp_idx / nwpr); // starting row number for this warp
int j = (warp_idx % nwpr) * 32; // starting column number for this warp int j = (warp_idx % nwpr) * 32; // starting column number for this warp
dist_t max_d = {.d = 0.0, .i = 0, .j = 0}; dist_t max_d = { .d = 0.0, .i = 0, .j = 0 };
__shared__ dist_t max_darr[32]; __shared__ dist_t max_darr[32];
// Optimization: About 1/3 of the warps will have no work. // Optimization: About 1/3 of the warps will have no work.
// We compute distances only when i<j. If i is never <j in this warp, we have // We compute distances only when i<j. If i is never <j in this warp, we have
// no work // no work
if (i > j + 31) { // this warp has no work if (i > j + 31) // this warp has no work
{
max_darr[warp_idx] = max_d; max_darr[warp_idx] = max_d;
return; return;
} }
// load robot_spheres to shared memory // load robot_spheres to shared memory
extern __shared__ float4 __rs_shared[]; extern __shared__ float4 __rs_shared[];
if (threadIdx.x < nspheres) {
if (threadIdx.x < nspheres)
{
float4 sph = *(float4 *)&robot_spheres[4 * (batch_idx * nspheres + threadIdx.x)]; float4 sph = *(float4 *)&robot_spheres[4 * (batch_idx * nspheres + threadIdx.x)];
//float4 sph = make_float4(robot_spheres[3 * (batch_idx * nspheres + threadIdx.x)],
//robot_spheres[3 * (batch_idx * nspheres + threadIdx.x) + 1], // float4 sph = make_float4(robot_spheres[3 * (batch_idx * nspheres + threadIdx.x)],
//robot_spheres[3 * (batch_idx * nspheres + threadIdx.x) + 2], // robot_spheres[3 * (batch_idx * nspheres + threadIdx.x) + 1],
//robot_spheres_radius[threadIdx.x]) ; // robot_spheres[3 * (batch_idx * nspheres + threadIdx.x) + 2],
// robot_spheres_radius[threadIdx.x]) ;
sph.w += offsets[threadIdx.x]; sph.w += offsets[threadIdx.x];
__rs_shared[threadIdx.x] = sph; __rs_shared[threadIdx.x] = sph;
} }
@@ -148,17 +170,23 @@ __global__ void self_collision_distance_kernel4(
j = j + threadIdx.x % 32; // column number for this thread j = j + threadIdx.x % 32; // column number for this thread
float4 sph2; float4 sph2;
if (j < nspheres) {
if (j < nspheres)
{
sph2 = __rs_shared[j]; // we need not load sph2 in every iteration. sph2 = __rs_shared[j]; // we need not load sph2 in every iteration.
#pragma unroll 16 #pragma unroll 16
for (int k = 0; k < ndpt; k++, i++) { // increment i also here
if (i < nspheres && j > i) {
// check if self collision is allowed here:
if (coll_matrix[i * nspheres + j] == 1) {
for (int k = 0; k < ndpt; k++, i++) // increment i also here
{
if ((i < nspheres) && (j > i))
{
// check if self collision is allowed here:
if (coll_matrix[i * nspheres + j] == 1)
{
float4 sph1 = __rs_shared[i]; float4 sph1 = __rs_shared[i];
if (sph1.w <= 0.0 || sph2.w <= 0.0)
if ((sph1.w <= 0.0) || (sph2.w <= 0.0))
{ {
continue; continue;
} }
@@ -166,11 +194,14 @@ __global__ void self_collision_distance_kernel4(
float d = sqrt((sph1.x - sph2.x) * (sph1.x - sph2.x) + float d = sqrt((sph1.x - sph2.x) * (sph1.x - sph2.x) +
(sph1.y - sph2.y) * (sph1.y - sph2.y) + (sph1.y - sph2.y) * (sph1.y - sph2.y) +
(sph1.z - sph2.z) * (sph1.z - sph2.z)); (sph1.z - sph2.z) * (sph1.z - sph2.z));
//float distance = max((float)0.0, (float)(r_diff - d));
// float distance = max((float)0.0, (float)(r_diff - d));
float distance = r_diff - d; float distance = r_diff - d;
// printf("%d, %d: (%d, %d) %f new\n", blockIdx.x, threadIdx.x, i, j, // printf("%d, %d: (%d, %d) %f new\n", blockIdx.x, threadIdx.x, i, j,
// distance); // distance);
if (distance > max_d.d) { if (distance > max_d.d)
{
max_d.d = distance; max_d.d = distance;
max_d.i = i; max_d.i = i;
max_d.j = j; max_d.j = j;
@@ -179,6 +210,7 @@ __global__ void self_collision_distance_kernel4(
} }
} }
} }
// max_d has the result max for this thread // max_d has the result max for this thread
////////////////////////////////////////////////////// //////////////////////////////////////////////////////
@@ -189,15 +221,22 @@ __global__ void self_collision_distance_kernel4(
// Optimization: Skip the reduction if all the values are zero // Optimization: Skip the reduction if all the values are zero
unsigned zero_mask = __ballot_sync( unsigned zero_mask = __ballot_sync(
0xffffffff, max_d.d != 0.0); // we expect most values to be 0. So, 0xffffffff, max_d.d != 0.0); // we expect most values to be 0. So,
// zero_mask should be 0 in the common case. // zero_mask should be 0 in the common case.
if (zero_mask != 0) { // some of the values are non-zero if (zero_mask != 0) // some of the values are non-zero
{
unsigned mask = __ballot_sync(0xffffffff, threadIdx.x < blockDim.x); unsigned mask = __ballot_sync(0xffffffff, threadIdx.x < blockDim.x);
if (threadIdx.x < blockDim.x) {
if (threadIdx.x < blockDim.x)
{
// dist_t max_d = dist_sh[threadIdx.x]; // dist_t max_d = dist_sh[threadIdx.x];
for (int offset = 16; offset > 0; offset /= 2) { for (int offset = 16; offset > 0; offset /= 2)
{
uint64_t nd = __shfl_down_sync(mask, *(uint64_t *)&max_d, offset); uint64_t nd = __shfl_down_sync(mask, *(uint64_t *)&max_d, offset);
dist_t d_temp = *(dist_t *)&nd; dist_t d_temp = *(dist_t *)&nd;
if (d_temp.d > max_d.d) {
if (d_temp.d > max_d.d)
{
max_d = d_temp; max_d = d_temp;
} }
} }
@@ -205,15 +244,18 @@ __global__ void self_collision_distance_kernel4(
} }
// thread0 in the warp has the max_d for the warp // thread0 in the warp has the max_d for the warp
if (threadIdx.x % 32 == 0) { if (threadIdx.x % 32 == 0)
{
max_darr[warp_idx] = max_d; max_darr[warp_idx] = max_d;
// printf("threadIdx.x=%d, blockIdx.x=%d, max_d=%f\n", threadIdx.x, // printf("threadIdx.x=%d, blockIdx.x=%d, max_d=%f\n", threadIdx.x,
// blockIdx.x, max_d); // blockIdx.x, max_d);
} }
if (threadIdx.x < nspheres) { if (threadIdx.x < nspheres)
if (write_grad && sparse_index[batch_idx * nspheres + threadIdx.x] != 0) { {
if (write_grad && (sparse_index[batch_idx * nspheres + threadIdx.x] != 0))
{
*(float4 *)&out_vec[batch_idx * nspheres * 4 + threadIdx.x * 4] = *(float4 *)&out_vec[batch_idx * nspheres * 4 + threadIdx.x * 4] =
make_float4(0.0); make_float4(0.0);
sparse_index[batch_idx * nspheres + threadIdx.x] = 0; sparse_index[batch_idx * nspheres + threadIdx.x] = 0;
@@ -221,20 +263,28 @@ __global__ void self_collision_distance_kernel4(
} }
__syncthreads(); __syncthreads();
if (threadIdx.x == 0) { if (threadIdx.x == 0)
{
dist_t max_d = max_darr[0]; dist_t max_d = max_darr[0];
// TODO: This can be parallized // TODO: This can be parallized
for (int i = 1; i < (blockDim.x + 31) / 32; i++) { for (int i = 1; i < (blockDim.x + 31) / 32; i++)
if (max_darr[i].d > max_d.d) { {
if (max_darr[i].d > max_d.d)
{
max_d = max_darr[i]; max_d = max_darr[i];
} }
} }
////////////////////////////////////////////////////// //////////////////////////////////////////////////////
// Write out the final results // Write out the final results
////////////////////////////////////////////////////// //////////////////////////////////////////////////////
if (max_d.d != 0.0) { if (max_d.d != 0.0)
{
out_distance[batch_idx] = weight[0] * max_d.d; out_distance[batch_idx] = weight[0] * max_d.d;
if (write_grad) {
if (write_grad)
{
float3 sph1 = float3 sph1 =
*(float3 *)&robot_spheres[4 * (batch_idx * nspheres + max_d.i)]; *(float3 *)&robot_spheres[4 * (batch_idx * nspheres + max_d.i)];
float3 sph2 = float3 sph2 =
@@ -247,20 +297,24 @@ __global__ void self_collision_distance_kernel4(
sparse_index[batch_idx * nspheres + max_d.i] = 1; sparse_index[batch_idx * nspheres + max_d.i] = 1;
sparse_index[batch_idx * nspheres + max_d.j] = 1; sparse_index[batch_idx * nspheres + max_d.j] = 1;
} }
} else { }
else
{
out_distance[batch_idx] = 0; out_distance[batch_idx] = 0;
} }
} }
} }
template <typename scalar_t, int ndpt, int NBPB>
__global__ void self_collision_distance_kernel7( template<typename scalar_t, int ndpt, int NBPB>
__global__ void self_collision_distance_kernel7(
scalar_t *out_distance, // batch x 1 scalar_t *out_distance, // batch x 1
scalar_t *out_vec, // batch x nspheres x 3 scalar_t *out_vec, // batch x nspheres x 3
uint8_t *sparse_index, uint8_t *sparse_index,
const scalar_t *robot_spheres, // batch x nspheres x 3 const scalar_t *robot_spheres, // batch x nspheres x 3
const scalar_t *offsets, // nspheres const scalar_t *offsets, // nspheres
const scalar_t *weight, const int16_t *locations_, const int batch_size, const scalar_t *weight, const int16_t *locations_, const int batch_size,
const int nspheres, const bool write_grad = false) { const int nspheres, const bool write_grad = false)
{
uint32_t batch_idx = blockIdx.x * NBPB; uint32_t batch_idx = blockIdx.x * NBPB;
uint8_t nbpb = min(NBPB, batch_size - batch_idx); uint8_t nbpb = min(NBPB, batch_size - batch_idx);
@@ -272,17 +326,21 @@ __global__ void self_collision_distance_kernel7(
// sphere2[batch=0] sphere2[batch=1] sphere2[batch=2] sphere2[batch=4] // sphere2[batch=0] sphere2[batch=1] sphere2[batch=2] sphere2[batch=4]
// ... // ...
extern __shared__ float4 __rs_shared[]; extern __shared__ float4 __rs_shared[];
if (threadIdx.x < nspheres) { // threadIdx.x is sphere index
if (threadIdx.x < nspheres) // threadIdx.x is sphere index
{
#pragma unroll #pragma unroll
for (int l = 0; l < nbpb; l++) {
for (int l = 0; l < nbpb; l++)
{
float4 sph = *(float4 *)&robot_spheres[4 * ((batch_idx + l) * nspheres + threadIdx.x)]; float4 sph = *(float4 *)&robot_spheres[4 * ((batch_idx + l) * nspheres + threadIdx.x)];
//float4 sph = make_float4( // float4 sph = make_float4(
// robot_spheres[3 * ((batch_idx + l) * nspheres + threadIdx.x)], // robot_spheres[3 * ((batch_idx + l) * nspheres + threadIdx.x)],
// robot_spheres[3 * ((batch_idx + l) * nspheres + threadIdx.x) + 1], // robot_spheres[3 * ((batch_idx + l) * nspheres + threadIdx.x) + 1],
// robot_spheres[3 * ((batch_idx + l) * nspheres + threadIdx.x) + 2], // robot_spheres[3 * ((batch_idx + l) * nspheres + threadIdx.x) + 2],
// robot_spheres_radius[threadIdx.x] // robot_spheres_radius[threadIdx.x]
//); // );
sph.w += offsets[threadIdx.x]; sph.w += offsets[threadIdx.x];
__rs_shared[NBPB * threadIdx.x + l] = sph; __rs_shared[NBPB * threadIdx.x + l] = sph;
@@ -295,28 +353,34 @@ __global__ void self_collision_distance_kernel7(
// in registers (max_d). // in registers (max_d).
// Each thread computes upto ndpt distances. // Each thread computes upto ndpt distances.
////////////////////////////////////////////////////// //////////////////////////////////////////////////////
dist_t max_d[NBPB] = {{.d = 0.0, .i = 0, .j = 0}}; dist_t max_d[NBPB] = {{ .d = 0.0, .i = 0, .j = 0 } };
int16_t indices[ndpt * 2]; int16_t indices[ndpt * 2];
for (uint8_t i =0; i< ndpt * 2; i++)
for (uint8_t i = 0; i < ndpt * 2; i++)
{ {
indices[i] = locations_[(threadIdx.x) * 2 * ndpt + i]; indices[i] = locations_[(threadIdx.x) * 2 * ndpt + i];
} }
#pragma unroll #pragma unroll
for (uint8_t k = 0; k < ndpt; k++) {
for (uint8_t k = 0; k < ndpt; k++)
{
// We are iterating through ndpt pair of spheres across batch // We are iterating through ndpt pair of spheres across batch
// if we increase ndpt, then we can compute for more spheres? // if we increase ndpt, then we can compute for more spheres?
int i = indices[k * 2]; int i = indices[k * 2];
int j = indices[k * 2 + 1]; int j = indices[k * 2 + 1];
if (i == -1 || j == -1)
if ((i == -1) || (j == -1))
continue; continue;
#pragma unroll #pragma unroll
for (uint16_t l = 0; l < nbpb; l++) { // iterate through nbpb batches
for (uint16_t l = 0; l < nbpb; l++) // iterate through nbpb batches
{
float4 sph1 = __rs_shared[NBPB * i + l]; float4 sph1 = __rs_shared[NBPB * i + l];
float4 sph2 = __rs_shared[NBPB * j + l]; float4 sph2 = __rs_shared[NBPB * j + l];
if (sph1.w <= 0.0 || sph2.w <= 0.0)
if ((sph1.w <= 0.0) || (sph2.w <= 0.0))
{ {
continue; continue;
} }
@@ -326,13 +390,16 @@ __global__ void self_collision_distance_kernel7(
(sph1.y - sph2.y) * (sph1.y - sph2.y) + (sph1.y - sph2.y) * (sph1.y - sph2.y) +
(sph1.z - sph2.z) * (sph1.z - sph2.z)); (sph1.z - sph2.z) * (sph1.z - sph2.z));
float f_diff = r_diff - d; float f_diff = r_diff - d;
if (f_diff > max_d[l].d) {
if (f_diff > max_d[l].d)
{
max_d[l].d = f_diff; max_d[l].d = f_diff;
max_d[l].i = i; max_d[l].i = i;
max_d[l].j = j; max_d[l].j = j;
} }
} }
} }
// max_d has the result max for this thread // max_d has the result max for this thread
////////////////////////////////////////////////////// //////////////////////////////////////////////////////
@@ -344,40 +411,54 @@ __global__ void self_collision_distance_kernel7(
__shared__ dist_t max_darr[32 * NBPB]; __shared__ dist_t max_darr[32 * NBPB];
#pragma unroll #pragma unroll
for (uint8_t l = 0; l < nbpb; l++) {
for (uint8_t l = 0; l < nbpb; l++)
{
// Perform warp-wide reductions // Perform warp-wide reductions
// Optimization: Skip the reduction if all the values are zero // Optimization: Skip the reduction if all the values are zero
unsigned zero_mask = __ballot_sync( unsigned zero_mask = __ballot_sync(
0xffffffff, 0xffffffff,
max_d[l].d != 0.0); // we expect most values to be 0. So, zero_mask max_d[l].d != 0.0); // we expect most values to be 0. So, zero_mask
// should be 0 in the common case. // should be 0 in the common case.
if (zero_mask != 0) { // some of the values are non-zero if (zero_mask != 0) // some of the values are non-zero
{
unsigned mask = __ballot_sync(0xffffffff, threadIdx.x < blockDim.x); unsigned mask = __ballot_sync(0xffffffff, threadIdx.x < blockDim.x);
if (threadIdx.x < blockDim.x) {
if (threadIdx.x < blockDim.x)
{
// dist_t max_d = dist_sh[threadIdx.x]; // dist_t max_d = dist_sh[threadIdx.x];
for (int offset = 16; offset > 0; offset /= 2) { for (int offset = 16; offset > 0; offset /= 2)
{
// the offset here is linked to ndpt? // the offset here is linked to ndpt?
uint64_t nd = __shfl_down_sync(mask, *(uint64_t *)&max_d[l], offset); uint64_t nd = __shfl_down_sync(mask, *(uint64_t *)&max_d[l], offset);
dist_t d_temp = *(dist_t *)&nd; dist_t d_temp = *(dist_t *)&nd;
if (d_temp.d > max_d[l].d) {
if (d_temp.d > max_d[l].d)
{
max_d[l] = d_temp; max_d[l] = d_temp;
} }
} }
} }
} }
// thread0 in the warp has the max_d for the warp // thread0 in the warp has the max_d for the warp
if (threadIdx.x % 32 == 0) { if (threadIdx.x % 32 == 0)
{
max_darr[(threadIdx.x / 32) + 32 * l] = max_d[l]; max_darr[(threadIdx.x / 32) + 32 * l] = max_d[l];
// printf("threadIdx.x=%d, blockIdx.x=%d, max_d=%f\n", threadIdx.x, // printf("threadIdx.x=%d, blockIdx.x=%d, max_d=%f\n", threadIdx.x,
// blockIdx.x, max_d); // blockIdx.x, max_d);
} }
} }
if (threadIdx.x < nspheres) { if (threadIdx.x < nspheres)
for (int l = 0; l < nbpb; l++) { {
for (int l = 0; l < nbpb; l++)
{
if (write_grad && if (write_grad &&
sparse_index[(batch_idx + l) * nspheres + threadIdx.x] != 0) { (sparse_index[(batch_idx + l) * nspheres + threadIdx.x] != 0))
{
*(float4 *)&out_vec[(batch_idx + l) * nspheres * 4 + threadIdx.x * 4] = *(float4 *)&out_vec[(batch_idx + l) * nspheres * 4 + threadIdx.x * 4] =
make_float4(0.0); make_float4(0.0);
sparse_index[(batch_idx + l) * nspheres + threadIdx.x] = 0; sparse_index[(batch_idx + l) * nspheres + threadIdx.x] = 0;
@@ -386,22 +467,32 @@ __global__ void self_collision_distance_kernel7(
} }
__syncthreads(); __syncthreads();
if (threadIdx.x == 0) { if (threadIdx.x == 0)
{
#pragma unroll #pragma unroll
for (uint8_t l = 0; l < nbpb; l++) {
for (uint8_t l = 0; l < nbpb; l++)
{
dist_t max_d = max_darr[l * 32]; dist_t max_d = max_darr[l * 32];
// TODO: This can be parallized // TODO: This can be parallized
for (int i = 1; i < (blockDim.x + 31) / 32; i++) { for (int i = 1; i < (blockDim.x + 31) / 32; i++)
if (max_darr[l * 32 + i].d > max_d.d) { {
if (max_darr[l * 32 + i].d > max_d.d)
{
max_d = max_darr[l * 32 + i]; max_d = max_darr[l * 32 + i];
} }
} }
////////////////////////////////////////////////////// //////////////////////////////////////////////////////
// Write out the final results // Write out the final results
////////////////////////////////////////////////////// //////////////////////////////////////////////////////
if (max_d.d != 0.0) { if (max_d.d != 0.0)
{
out_distance[batch_idx + l] = weight[0] * max_d.d; out_distance[batch_idx + l] = weight[0] * max_d.d;
if (write_grad) {
if (write_grad)
{
float3 sph1 = float3 sph1 =
*(float3 *)&robot_spheres[4 * *(float3 *)&robot_spheres[4 *
((batch_idx + l) * nspheres + max_d.i)]; ((batch_idx + l) * nspheres + max_d.i)];
@@ -416,19 +507,21 @@ __global__ void self_collision_distance_kernel7(
sparse_index[(batch_idx + l) * nspheres + max_d.i] = 1; sparse_index[(batch_idx + l) * nspheres + max_d.i] = 1;
sparse_index[(batch_idx + l) * nspheres + max_d.j] = 1; sparse_index[(batch_idx + l) * nspheres + max_d.j] = 1;
} }
} else { }
else
{
out_distance[batch_idx + l] = 0; out_distance[batch_idx + l] = 0;
} }
} }
} }
} }
} // namespace Geometry } // namespace Geometry
} // namespace Curobo } // namespace Curobo
// This is the best version so far. // This is the best version so far.
// It precomputes the start addresses per thread on the cpu. // It precomputes the start addresses per thread on the cpu.
// The rest is similar to the version above. // The rest is similar to the version above.
std::vector<torch::Tensor> self_collision_distance( std::vector<torch::Tensor>self_collision_distance(
torch::Tensor out_distance, torch::Tensor out_vec, torch::Tensor out_distance, torch::Tensor out_vec,
torch::Tensor sparse_index, torch::Tensor sparse_index,
const torch::Tensor robot_spheres, // batch_size x n_spheres x 3 const torch::Tensor robot_spheres, // batch_size x n_spheres x 3
@@ -437,7 +530,8 @@ std::vector<torch::Tensor> self_collision_distance(
const torch::Tensor thread_locations, const int thread_locations_size, const torch::Tensor thread_locations, const int thread_locations_size,
const int batch_size, const int nspheres, const bool compute_grad = false, const int batch_size, const int nspheres, const bool compute_grad = false,
const int ndpt = 8, // Does this need to match template? const int ndpt = 8, // Does this need to match template?
const bool experimental_kernel = false) { const bool experimental_kernel = false)
{
using namespace Curobo::Geometry; using namespace Curobo::Geometry;
// use efficient kernel based on number of threads: // use efficient kernel based on number of threads:
@@ -448,12 +542,14 @@ std::vector<torch::Tensor> self_collision_distance(
int threadsPerBlock = ((thread_locations_size / 2) + ndpt - 1) / int threadsPerBlock = ((thread_locations_size / 2) + ndpt - 1) /
ndpt; // location_size must be an even number. We store ndpt; // location_size must be an even number. We store
// i,j for each sphere pair. // i,j for each sphere pair.
//assert(threadsPerBlock/nbpb <=32); // assert(threadsPerBlock/nbpb <=32);
if (threadsPerBlock < 32*nbpb) if (threadsPerBlock < 32 * nbpb)
{ {
threadsPerBlock = 32 * nbpb; threadsPerBlock = 32 * nbpb;
} }
if (threadsPerBlock < nspheres) if (threadsPerBlock < nspheres)
{ {
threadsPerBlock = nspheres; threadsPerBlock = nspheres;
@@ -462,19 +558,17 @@ std::vector<torch::Tensor> self_collision_distance(
cudaStream_t stream = at::cuda::getCurrentCUDAStream(); cudaStream_t stream = at::cuda::getCurrentCUDAStream();
// ((threadsPerBlock >= nspheres))&& // ((threadsPerBlock >= nspheres))&&
if (experimental_kernel) { if (experimental_kernel)
{
int smemSize = nbpb * nspheres * sizeof(float4); int smemSize = nbpb * nspheres * sizeof(float4);
if (ndpt == 1) if (ndpt == 1)
{ {
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
robot_spheres.scalar_type(), "self_collision_distance", ([&] { robot_spheres.scalar_type(), "self_collision_distance", ([&] {
self_collision_distance_kernel7<scalar_t, 1, nbpb> self_collision_distance_kernel7<scalar_t, 1, nbpb>
<<<blocksPerGrid, threadsPerBlock, smemSize, stream>>>( << < blocksPerGrid, threadsPerBlock, smemSize, stream >> > (
out_distance.data_ptr<scalar_t>(), out_distance.data_ptr<scalar_t>(),
out_vec.data_ptr<scalar_t>(), out_vec.data_ptr<scalar_t>(),
sparse_index.data_ptr<uint8_t>(), sparse_index.data_ptr<uint8_t>(),
@@ -488,11 +582,10 @@ std::vector<torch::Tensor> self_collision_distance(
else if (ndpt == 2) else if (ndpt == 2)
{ {
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
robot_spheres.scalar_type(), "self_collision_distance", ([&] { robot_spheres.scalar_type(), "self_collision_distance", ([&] {
self_collision_distance_kernel7<scalar_t, 2, nbpb> self_collision_distance_kernel7<scalar_t, 2, nbpb>
<<<blocksPerGrid, threadsPerBlock, smemSize, stream>>>( << < blocksPerGrid, threadsPerBlock, smemSize, stream >> > (
out_distance.data_ptr<scalar_t>(), out_distance.data_ptr<scalar_t>(),
out_vec.data_ptr<scalar_t>(), out_vec.data_ptr<scalar_t>(),
sparse_index.data_ptr<uint8_t>(), sparse_index.data_ptr<uint8_t>(),
@@ -506,11 +599,10 @@ std::vector<torch::Tensor> self_collision_distance(
else if (ndpt == 4) else if (ndpt == 4)
{ {
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
robot_spheres.scalar_type(), "self_collision_distance", ([&] { robot_spheres.scalar_type(), "self_collision_distance", ([&] {
self_collision_distance_kernel7<scalar_t, 4, nbpb> self_collision_distance_kernel7<scalar_t, 4, nbpb>
<<<blocksPerGrid, threadsPerBlock, smemSize, stream>>>( << < blocksPerGrid, threadsPerBlock, smemSize, stream >> > (
out_distance.data_ptr<scalar_t>(), out_distance.data_ptr<scalar_t>(),
out_vec.data_ptr<scalar_t>(), out_vec.data_ptr<scalar_t>(),
sparse_index.data_ptr<uint8_t>(), sparse_index.data_ptr<uint8_t>(),
@@ -523,11 +615,10 @@ std::vector<torch::Tensor> self_collision_distance(
} }
else if (ndpt == 8) else if (ndpt == 8)
{ {
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
robot_spheres.scalar_type(), "self_collision_distance", ([&] { robot_spheres.scalar_type(), "self_collision_distance", ([&] {
self_collision_distance_kernel7<scalar_t, 8, nbpb> self_collision_distance_kernel7<scalar_t, 8, nbpb>
<<<blocksPerGrid, threadsPerBlock, smemSize, stream>>>( << < blocksPerGrid, threadsPerBlock, smemSize, stream >> > (
out_distance.data_ptr<scalar_t>(), out_distance.data_ptr<scalar_t>(),
out_vec.data_ptr<scalar_t>(), out_vec.data_ptr<scalar_t>(),
sparse_index.data_ptr<uint8_t>(), sparse_index.data_ptr<uint8_t>(),
@@ -540,12 +631,10 @@ std::vector<torch::Tensor> self_collision_distance(
} }
else if (ndpt == 32) else if (ndpt == 32)
{ {
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
robot_spheres.scalar_type(), "self_collision_distance", ([&] { robot_spheres.scalar_type(), "self_collision_distance", ([&] {
self_collision_distance_kernel7<scalar_t, 32, nbpb> self_collision_distance_kernel7<scalar_t, 32, nbpb>
<<<blocksPerGrid, threadsPerBlock, smemSize, stream>>>( << < blocksPerGrid, threadsPerBlock, smemSize, stream >> > (
out_distance.data_ptr<scalar_t>(), out_distance.data_ptr<scalar_t>(),
out_vec.data_ptr<scalar_t>(), out_vec.data_ptr<scalar_t>(),
sparse_index.data_ptr<uint8_t>(), sparse_index.data_ptr<uint8_t>(),
@@ -561,7 +650,7 @@ std::vector<torch::Tensor> self_collision_distance(
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
robot_spheres.scalar_type(), "self_collision_distance", ([&] { robot_spheres.scalar_type(), "self_collision_distance", ([&] {
self_collision_distance_kernel7<scalar_t, 64, nbpb> self_collision_distance_kernel7<scalar_t, 64, nbpb>
<<<blocksPerGrid, threadsPerBlock, smemSize, stream>>>( << < blocksPerGrid, threadsPerBlock, smemSize, stream >> > (
out_distance.data_ptr<scalar_t>(), out_distance.data_ptr<scalar_t>(),
out_vec.data_ptr<scalar_t>(), out_vec.data_ptr<scalar_t>(),
sparse_index.data_ptr<uint8_t>(), sparse_index.data_ptr<uint8_t>(),
@@ -577,7 +666,7 @@ std::vector<torch::Tensor> self_collision_distance(
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
robot_spheres.scalar_type(), "self_collision_distance", ([&] { robot_spheres.scalar_type(), "self_collision_distance", ([&] {
self_collision_distance_kernel7<scalar_t, 128, nbpb> self_collision_distance_kernel7<scalar_t, 128, nbpb>
<<<blocksPerGrid, threadsPerBlock, smemSize, stream>>>( << < blocksPerGrid, threadsPerBlock, smemSize, stream >> > (
out_distance.data_ptr<scalar_t>(), out_distance.data_ptr<scalar_t>(),
out_vec.data_ptr<scalar_t>(), out_vec.data_ptr<scalar_t>(),
sparse_index.data_ptr<uint8_t>(), sparse_index.data_ptr<uint8_t>(),
@@ -587,14 +676,13 @@ std::vector<torch::Tensor> self_collision_distance(
thread_locations.data_ptr<int16_t>(), batch_size, nspheres, thread_locations.data_ptr<int16_t>(), batch_size, nspheres,
compute_grad); compute_grad);
})); }));
} }
else if (ndpt == 512) else if (ndpt == 512)
{ {
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
robot_spheres.scalar_type(), "self_collision_distance", ([&] { robot_spheres.scalar_type(), "self_collision_distance", ([&] {
self_collision_distance_kernel7<scalar_t, 512, nbpb> self_collision_distance_kernel7<scalar_t, 512, nbpb>
<<<blocksPerGrid, threadsPerBlock, smemSize, stream>>>( << < blocksPerGrid, threadsPerBlock, smemSize, stream >> > (
out_distance.data_ptr<scalar_t>(), out_distance.data_ptr<scalar_t>(),
out_vec.data_ptr<scalar_t>(), out_vec.data_ptr<scalar_t>(),
sparse_index.data_ptr<uint8_t>(), sparse_index.data_ptr<uint8_t>(),
@@ -611,12 +699,14 @@ std::vector<torch::Tensor> self_collision_distance(
} }
} }
else { else
{
int ndpt_n = 32; // number of distances to be computed per thread int ndpt_n = 32; // number of distances to be computed per thread
int nwpr = (nspheres + 31) / 32; int nwpr = (nspheres + 31) / 32;
int warpsPerBlock = nwpr * ((nspheres + ndpt_n - 1) / ndpt_n); int warpsPerBlock = nwpr * ((nspheres + ndpt_n - 1) / ndpt_n);
threadsPerBlock = warpsPerBlock * 32; threadsPerBlock = warpsPerBlock * 32;
blocksPerGrid = batch_size; blocksPerGrid = batch_size;
// printf("Blocks: %d, threads/block:%d, ndpt=%d, nwpr=%d\n", blocksPerGrid, // printf("Blocks: %d, threads/block:%d, ndpt=%d, nwpr=%d\n", blocksPerGrid,
// threadsPerBlock, ndpt, nwpr); // threadsPerBlock, ndpt, nwpr);
@@ -626,7 +716,7 @@ std::vector<torch::Tensor> self_collision_distance(
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
robot_spheres.scalar_type(), "self_collision_distance", ([&] { robot_spheres.scalar_type(), "self_collision_distance", ([&] {
self_collision_distance_kernel4<scalar_t> self_collision_distance_kernel4<scalar_t>
<<<blocksPerGrid, threadsPerBlock, smemSize, stream>>>( << < blocksPerGrid, threadsPerBlock, smemSize, stream >> > (
out_distance.data_ptr<scalar_t>(), out_distance.data_ptr<scalar_t>(),
out_vec.data_ptr<scalar_t>(), out_vec.data_ptr<scalar_t>(),
robot_spheres.data_ptr<scalar_t>(), robot_spheres.data_ptr<scalar_t>(),
@@ -638,6 +728,5 @@ std::vector<torch::Tensor> self_collision_distance(
} }
C10_CUDA_KERNEL_LAUNCH_CHECK(); C10_CUDA_KERNEL_LAUNCH_CHECK();
return {out_distance, out_vec, sparse_index}; return { out_distance, out_vec, sparse_index };
} }

820
src/curobo/curobolib/cpp/sphere_obb_kernel.cu
View File

File diff suppressed because it is too large Load Diff

183
src/curobo/curobolib/cpp/tensor_step_cuda.cpp
View File

@@ -13,72 +13,120 @@
#include <c10/cuda/CUDAGuard.h> #include <c10/cuda/CUDAGuard.h>
#include <vector> #include <vector>
std::vector<torch::Tensor> step_position_clique( std::vector<torch::Tensor>step_position_clique(
torch::Tensor out_position, torch::Tensor out_velocity, torch::Tensor out_position,
torch::Tensor out_acceleration, torch::Tensor out_jerk, torch::Tensor out_velocity,
const torch::Tensor u_position, const torch::Tensor start_position, torch::Tensor out_acceleration,
const torch::Tensor start_velocity, const torch::Tensor start_acceleration, torch::Tensor out_jerk,
const torch::Tensor traj_dt, const int batch_size, const int horizon, const torch::Tensor u_position,
const int dof); const torch::Tensor start_position,
std::vector<torch::Tensor> step_position_clique2(
torch::Tensor out_position, torch::Tensor out_velocity,
torch::Tensor out_acceleration, torch::Tensor out_jerk,
const torch::Tensor u_position, const torch::Tensor start_position,
const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
const torch::Tensor traj_dt, const int batch_size, const int horizon,
const int dof, const int mode);
std::vector<torch::Tensor> step_position_clique2_idx(
torch::Tensor out_position, torch::Tensor out_velocity,
torch::Tensor out_acceleration, torch::Tensor out_jerk,
const torch::Tensor u_position, const torch::Tensor start_position,
const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
const torch::Tensor start_idx, const torch::Tensor traj_dt,
const int batch_size, const int horizon, const int dof, const int mode);
std::vector<torch::Tensor> backward_step_position_clique(
torch::Tensor out_grad_position, const torch::Tensor grad_position,
const torch::Tensor grad_velocity, const torch::Tensor grad_acceleration,
const torch::Tensor grad_jerk, const torch::Tensor traj_dt,
const int batch_size, const int horizon, const int dof);
std::vector<torch::Tensor> backward_step_position_clique2(
torch::Tensor out_grad_position, const torch::Tensor grad_position,
const torch::Tensor grad_velocity, const torch::Tensor grad_acceleration,
const torch::Tensor grad_jerk, const torch::Tensor traj_dt,
const int batch_size, const int horizon, const int dof, const int mode);
std::vector<torch::Tensor>
step_acceleration(torch::Tensor out_position, torch::Tensor out_velocity,
torch::Tensor out_acceleration, torch::Tensor out_jerk,
const torch::Tensor u_acc, const torch::Tensor start_position,
const torch::Tensor start_velocity, const torch::Tensor start_velocity,
const torch::Tensor start_acceleration, const torch::Tensor start_acceleration,
const torch::Tensor traj_dt, const int batch_size, const torch::Tensor traj_dt,
const int horizon, const int dof, const bool use_rk2 = true); const int batch_size,
const int horizon,
const int dof);
std::vector<torch::Tensor>step_position_clique2(
torch::Tensor out_position,
torch::Tensor out_velocity,
torch::Tensor out_acceleration,
torch::Tensor out_jerk,
const torch::Tensor u_position,
const torch::Tensor start_position,
const torch::Tensor start_velocity,
const torch::Tensor start_acceleration,
const torch::Tensor traj_dt,
const int batch_size,
const int horizon,
const int dof,
const int mode);
std::vector<torch::Tensor>step_position_clique2_idx(
torch::Tensor out_position,
torch::Tensor out_velocity,
torch::Tensor out_acceleration,
torch::Tensor out_jerk,
const torch::Tensor u_position,
const torch::Tensor start_position,
const torch::Tensor start_velocity,
const torch::Tensor start_acceleration,
const torch::Tensor start_idx,
const torch::Tensor traj_dt,
const int batch_size,
const int horizon,
const int dof,
const int mode);
std::vector<torch::Tensor> step_acceleration_idx( std::vector<torch::Tensor>backward_step_position_clique(
torch::Tensor out_position, torch::Tensor out_velocity, torch::Tensor out_grad_position,
torch::Tensor out_acceleration, torch::Tensor out_jerk, const torch::Tensor grad_position,
const torch::Tensor u_acc, const torch::Tensor start_position, const torch::Tensor grad_velocity,
const torch::Tensor start_velocity, const torch::Tensor start_acceleration, const torch::Tensor grad_acceleration,
const torch::Tensor start_idx, const torch::Tensor traj_dt, const torch::Tensor grad_jerk,
const int batch_size, const int horizon, const int dof, const torch::Tensor traj_dt,
const int batch_size,
const int horizon,
const int dof);
std::vector<torch::Tensor>backward_step_position_clique2(
torch::Tensor out_grad_position,
const torch::Tensor grad_position,
const torch::Tensor grad_velocity,
const torch::Tensor grad_acceleration,
const torch::Tensor grad_jerk,
const torch::Tensor traj_dt,
const int batch_size,
const int horizon,
const int dof,
const int mode);
std::vector<torch::Tensor>
step_acceleration(torch::Tensor out_position,
torch::Tensor out_velocity,
torch::Tensor out_acceleration,
torch::Tensor out_jerk,
const torch::Tensor u_acc,
const torch::Tensor start_position,
const torch::Tensor start_velocity,
const torch::Tensor start_acceleration,
const torch::Tensor traj_dt,
const int batch_size,
const int horizon,
const int dof,
const bool use_rk2 = true); const bool use_rk2 = true);
std::vector<torch::Tensor>step_acceleration_idx(
torch::Tensor out_position,
torch::Tensor out_velocity,
torch::Tensor out_acceleration,
torch::Tensor out_jerk,
const torch::Tensor u_acc,
const torch::Tensor start_position,
const torch::Tensor start_velocity,
const torch::Tensor start_acceleration,
const torch::Tensor start_idx,
const torch::Tensor traj_dt,
const int batch_size,
const int horizon,
const int dof,
const bool use_rk2 = true);
// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. // 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_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) \ #define CHECK_CONTIGUOUS(x) \
AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") AT_ASSERTM(x.is_contiguous(), # x " must be contiguous")
#define CHECK_INPUT(x) \ #define CHECK_INPUT(x) \
CHECK_CUDA(x); \ CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x) CHECK_CONTIGUOUS(x)
std::vector<torch::Tensor> step_position_clique_wrapper( std::vector<torch::Tensor>step_position_clique_wrapper(
torch::Tensor out_position, torch::Tensor out_velocity, torch::Tensor out_position, torch::Tensor out_velocity,
torch::Tensor out_acceleration, torch::Tensor out_jerk, torch::Tensor out_acceleration, torch::Tensor out_jerk,
const torch::Tensor u_position, const torch::Tensor start_position, const torch::Tensor u_position, const torch::Tensor start_position,
const torch::Tensor start_velocity, const torch::Tensor start_acceleration, const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
const torch::Tensor traj_dt, const int batch_size, const int horizon, const torch::Tensor traj_dt, const int batch_size, const int horizon,
const int dof) { const int dof)
{
const at::cuda::OptionalCUDAGuard guard(u_position.device()); const at::cuda::OptionalCUDAGuard guard(u_position.device());
assert(false); // not supported assert(false); // not supported
CHECK_INPUT(u_position); CHECK_INPUT(u_position);
CHECK_INPUT(out_position); CHECK_INPUT(out_position);
@@ -96,14 +144,15 @@ std::vector<torch::Tensor> step_position_clique_wrapper(
batch_size, horizon, dof); batch_size, horizon, dof);
} }
std::vector<torch::Tensor> step_position_clique2_wrapper( std::vector<torch::Tensor>step_position_clique2_wrapper(
torch::Tensor out_position, torch::Tensor out_velocity, torch::Tensor out_position, torch::Tensor out_velocity,
torch::Tensor out_acceleration, torch::Tensor out_jerk, torch::Tensor out_acceleration, torch::Tensor out_jerk,
const torch::Tensor u_position, const torch::Tensor start_position, const torch::Tensor u_position, const torch::Tensor start_position,
const torch::Tensor start_velocity, const torch::Tensor start_acceleration, const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
const torch::Tensor traj_dt, const int batch_size, const int horizon, const torch::Tensor traj_dt, const int batch_size, const int horizon,
const int dof, const int dof,
const int mode) { const int mode)
{
const at::cuda::OptionalCUDAGuard guard(u_position.device()); const at::cuda::OptionalCUDAGuard guard(u_position.device());
CHECK_INPUT(u_position); CHECK_INPUT(u_position);
@@ -122,14 +171,15 @@ std::vector<torch::Tensor> step_position_clique2_wrapper(
batch_size, horizon, dof, mode); batch_size, horizon, dof, mode);
} }
std::vector<torch::Tensor> step_position_clique2_idx_wrapper( std::vector<torch::Tensor>step_position_clique2_idx_wrapper(
torch::Tensor out_position, torch::Tensor out_velocity, torch::Tensor out_position, torch::Tensor out_velocity,
torch::Tensor out_acceleration, torch::Tensor out_jerk, torch::Tensor out_acceleration, torch::Tensor out_jerk,
const torch::Tensor u_position, const torch::Tensor start_position, const torch::Tensor u_position, const torch::Tensor start_position,
const torch::Tensor start_velocity, const torch::Tensor start_acceleration, const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
const torch::Tensor start_idx, const torch::Tensor traj_dt, const torch::Tensor start_idx, const torch::Tensor traj_dt,
const int batch_size, const int horizon, const int dof, const int batch_size, const int horizon, const int dof,
const int mode) { const int mode)
{
const at::cuda::OptionalCUDAGuard guard(u_position.device()); const at::cuda::OptionalCUDAGuard guard(u_position.device());
CHECK_INPUT(u_position); CHECK_INPUT(u_position);
@@ -149,12 +199,14 @@ std::vector<torch::Tensor> step_position_clique2_idx_wrapper(
batch_size, horizon, dof, mode); batch_size, horizon, dof, mode);
} }
std::vector<torch::Tensor> backward_step_position_clique_wrapper( std::vector<torch::Tensor>backward_step_position_clique_wrapper(
torch::Tensor out_grad_position, const torch::Tensor grad_position, torch::Tensor out_grad_position, const torch::Tensor grad_position,
const torch::Tensor grad_velocity, const torch::Tensor grad_acceleration, const torch::Tensor grad_velocity, const torch::Tensor grad_acceleration,
const torch::Tensor grad_jerk, const torch::Tensor traj_dt, const torch::Tensor grad_jerk, const torch::Tensor traj_dt,
const int batch_size, const int horizon, const int dof) { const int batch_size, const int horizon, const int dof)
{
const at::cuda::OptionalCUDAGuard guard(grad_position.device()); const at::cuda::OptionalCUDAGuard guard(grad_position.device());
assert(false); // not supported assert(false); // not supported
CHECK_INPUT(out_grad_position); CHECK_INPUT(out_grad_position);
CHECK_INPUT(grad_position); CHECK_INPUT(grad_position);
@@ -167,12 +219,14 @@ std::vector<torch::Tensor> backward_step_position_clique_wrapper(
out_grad_position, grad_position, grad_velocity, grad_acceleration, out_grad_position, grad_position, grad_velocity, grad_acceleration,
grad_jerk, traj_dt, batch_size, horizon, dof); grad_jerk, traj_dt, batch_size, horizon, dof);
} }
std::vector<torch::Tensor> backward_step_position_clique2_wrapper(
std::vector<torch::Tensor>backward_step_position_clique2_wrapper(
torch::Tensor out_grad_position, const torch::Tensor grad_position, torch::Tensor out_grad_position, const torch::Tensor grad_position,
const torch::Tensor grad_velocity, const torch::Tensor grad_acceleration, const torch::Tensor grad_velocity, const torch::Tensor grad_acceleration,
const torch::Tensor grad_jerk, const torch::Tensor traj_dt, const torch::Tensor grad_jerk, const torch::Tensor traj_dt,
const int batch_size, const int horizon, const int dof, const int batch_size, const int horizon, const int dof,
const int mode) { const int mode)
{
const at::cuda::OptionalCUDAGuard guard(grad_position.device()); const at::cuda::OptionalCUDAGuard guard(grad_position.device());
CHECK_INPUT(out_grad_position); CHECK_INPUT(out_grad_position);
@@ -187,13 +241,14 @@ std::vector<torch::Tensor> backward_step_position_clique2_wrapper(
grad_jerk, traj_dt, batch_size, horizon, dof, mode); grad_jerk, traj_dt, batch_size, horizon, dof, mode);
} }
std::vector<torch::Tensor> step_acceleration_wrapper( std::vector<torch::Tensor>step_acceleration_wrapper(
torch::Tensor out_position, torch::Tensor out_velocity, torch::Tensor out_position, torch::Tensor out_velocity,
torch::Tensor out_acceleration, torch::Tensor out_jerk, torch::Tensor out_acceleration, torch::Tensor out_jerk,
const torch::Tensor u_acc, const torch::Tensor start_position, const torch::Tensor u_acc, const torch::Tensor start_position,
const torch::Tensor start_velocity, const torch::Tensor start_acceleration, const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
const torch::Tensor traj_dt, const int batch_size, const int horizon, const torch::Tensor traj_dt, const int batch_size, const int horizon,
const int dof, const bool use_rk2 = true) { const int dof, const bool use_rk2 = true)
{
const at::cuda::OptionalCUDAGuard guard(u_acc.device()); const at::cuda::OptionalCUDAGuard guard(u_acc.device());
CHECK_INPUT(u_acc); CHECK_INPUT(u_acc);
@@ -212,14 +267,15 @@ std::vector<torch::Tensor> step_acceleration_wrapper(
dof, use_rk2); dof, use_rk2);
} }
std::vector<torch::Tensor> step_acceleration_idx_wrapper( std::vector<torch::Tensor>step_acceleration_idx_wrapper(
torch::Tensor out_position, torch::Tensor out_velocity, torch::Tensor out_position, torch::Tensor out_velocity,
torch::Tensor out_acceleration, torch::Tensor out_jerk, torch::Tensor out_acceleration, torch::Tensor out_jerk,
const torch::Tensor u_acc, const torch::Tensor start_position, const torch::Tensor u_acc, const torch::Tensor start_position,
const torch::Tensor start_velocity, const torch::Tensor start_acceleration, const torch::Tensor start_velocity, const torch::Tensor start_acceleration,
const torch::Tensor start_idx, const torch::Tensor traj_dt, const torch::Tensor start_idx, const torch::Tensor traj_dt,
const int batch_size, const int horizon, const int dof, const int batch_size, const int horizon, const int dof,
const bool use_rk2 = true) { const bool use_rk2 = true)
{
const at::cuda::OptionalCUDAGuard guard(u_acc.device()); const at::cuda::OptionalCUDAGuard guard(u_acc.device());
CHECK_INPUT(u_acc); CHECK_INPUT(u_acc);
@@ -239,7 +295,8 @@ std::vector<torch::Tensor> step_acceleration_idx_wrapper(
batch_size, horizon, dof, use_rk2); batch_size, horizon, dof, use_rk2);
} }
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
m.def("step_position", &step_position_clique_wrapper, m.def("step_position", &step_position_clique_wrapper,
"Tensor Step Position (curobolib)"); "Tensor Step Position (curobolib)");
m.def("step_position2", &step_position_clique2_wrapper, m.def("step_position2", &step_position_clique2_wrapper,

970
src/curobo/curobolib/cpp/tensor_step_kernel.cu
View File

File diff suppressed because it is too large Load Diff

55
src/curobo/curobolib/cpp/update_best_kernel.cu
View File

@@ -26,14 +26,14 @@
#include <cub/cub.cuh> #include <cub/cub.cuh>
#include <math.h> #include <math.h>
namespace Curobo { namespace Curobo
{
namespace Optimization { namespace Optimization
{
// We launch with d_opt*cost_s1 threads. // We launch with d_opt*cost_s1 threads.
// We assume that cost_s2 is always 1. // We assume that cost_s2 is always 1.
template <typename scalar_t> template<typename scalar_t>
__global__ void update_best_kernel(scalar_t *best_cost, // 200x1 __global__ void update_best_kernel(scalar_t *best_cost, // 200x1
scalar_t *best_q, // 200x7 scalar_t *best_q, // 200x7
int16_t *best_iteration, // 200 x 1 int16_t *best_iteration, // 200 x 1
int16_t *current_iteration, // 1 int16_t *current_iteration, // 1
@@ -45,23 +45,29 @@ __global__ void update_best_kernel(scalar_t *best_cost, // 200x1
const int iteration, const int iteration,
const float delta_threshold, const float delta_threshold,
const float relative_threshold) // 1 const float relative_threshold) // 1
{ {
int tid = blockIdx.x * blockDim.x + threadIdx.x; int tid = blockIdx.x * blockDim.x + threadIdx.x;
int size = cost_s1 * d_opt; // size of best_q int size = cost_s1 * d_opt; // size of best_q
if (tid >= size) {
if (tid >= size)
{
return; return;
} }
const int cost_idx = tid / d_opt; const int cost_idx = tid / d_opt;
const float cost_new = cost[cost_idx]; const float cost_new = cost[cost_idx];
const float best_cost_in = best_cost[cost_idx]; const float best_cost_in = best_cost[cost_idx];
const bool change = (best_cost_in - cost_new) > delta_threshold && cost_new < best_cost_in * relative_threshold; const bool change = (best_cost_in - cost_new) > delta_threshold &&
if (change) { cost_new < best_cost_in * relative_threshold;
if (change)
{
best_q[tid] = q[tid]; // update best_q best_q[tid] = q[tid]; // update best_q
if (tid % d_opt == 0) { if (tid % d_opt == 0)
best_cost[cost_idx] = cost_new ; // update best_cost {
//best_iteration[cost_idx] = curr_iter + iteration; // best_cost[cost_idx] = cost_new; // update best_cost
// best_iteration[cost_idx] = curr_iter + iteration; //
// this tensor keeps track of whether the cost reduced by at least // this tensor keeps track of whether the cost reduced by at least
// delta_threshold. // delta_threshold.
// here iteration is the last_best parameter. // here iteration is the last_best parameter.
@@ -80,14 +86,13 @@ __global__ void update_best_kernel(scalar_t *best_cost, // 200x1
} }
} }
//.if (tid == 0) // .if (tid == 0)
//{ // {
// curr_iter += 1; // curr_iter += 1;
// current_iteration[0] = curr_iter; // current_iteration[0] = curr_iter;
//} // }
}
} } // namespace Optimization
} // namespace Optimization
} // namespace Curobo } // namespace Curobo
std::vector<torch::Tensor> std::vector<torch::Tensor>
@@ -98,12 +103,14 @@ update_best_cuda(torch::Tensor best_cost, torch::Tensor best_q,
const torch::Tensor q, const int d_opt, const int cost_s1, const torch::Tensor q, const int d_opt, const int cost_s1,
const int cost_s2, const int iteration, const int cost_s2, const int iteration,
const float delta_threshold, const float delta_threshold,
const float relative_threshold = 0.999) { const float relative_threshold = 0.999)
{
using namespace Curobo::Optimization; using namespace Curobo::Optimization;
const int threadsPerBlock = 128; const int threadsPerBlock = 128;
const int cost_size = cost_s1 * d_opt; const int cost_size = cost_s1 * d_opt;
assert(cost_s2 == 1); // assumption assert(cost_s2 == 1); // assumption
const int blocksPerGrid = (cost_size + threadsPerBlock - 1) / threadsPerBlock; const int blocksPerGrid = (cost_size + threadsPerBlock - 1) / threadsPerBlock;
// printf("cost_s1=%d, d_opt=%d, blocksPerGrid=%d\n", cost_s1, d_opt, // printf("cost_s1=%d, d_opt=%d, blocksPerGrid=%d\n", cost_s1, d_opt,
// blocksPerGrid); // blocksPerGrid);
@@ -112,7 +119,7 @@ update_best_cuda(torch::Tensor best_cost, torch::Tensor best_q,
AT_DISPATCH_FLOATING_TYPES( AT_DISPATCH_FLOATING_TYPES(
cost.scalar_type(), "update_best_cu", ([&] { cost.scalar_type(), "update_best_cu", ([&] {
update_best_kernel<scalar_t> update_best_kernel<scalar_t>
<<<blocksPerGrid, threadsPerBlock, 0, stream>>>( << < blocksPerGrid, threadsPerBlock, 0, stream >> > (
best_cost.data_ptr<scalar_t>(), best_q.data_ptr<scalar_t>(), best_cost.data_ptr<scalar_t>(), best_q.data_ptr<scalar_t>(),
best_iteration.data_ptr<int16_t>(), best_iteration.data_ptr<int16_t>(),
current_iteration.data_ptr<int16_t>(), current_iteration.data_ptr<int16_t>(),
@@ -122,5 +129,5 @@ update_best_cuda(torch::Tensor best_cost, torch::Tensor best_q,
})); }));
C10_CUDA_KERNEL_LAUNCH_CHECK(); C10_CUDA_KERNEL_LAUNCH_CHECK();
return {best_cost, best_q, best_iteration}; return { best_cost, best_q, best_iteration };
} }

331
src/curobo/curobolib/geom.py
View File

@@ -153,6 +153,8 @@ def get_pose_distance(
vec_convergence, vec_convergence,
run_weight, run_weight,
run_vec_weight, run_vec_weight,
offset_waypoint,
offset_tstep_fraction,
batch_pose_idx, batch_pose_idx,
batch_size, batch_size,
horizon, horizon,
@@ -161,6 +163,7 @@ def get_pose_distance(
write_grad=False, write_grad=False,
write_distance=False, write_distance=False,
use_metric=False, use_metric=False,
project_distance=True,
): ):
if batch_pose_idx.shape[0] != batch_size: if batch_pose_idx.shape[0] != batch_size:
raise ValueError("Index buffer size is different from batch size") raise ValueError("Index buffer size is different from batch size")
@@ -181,6 +184,8 @@ def get_pose_distance(
vec_convergence, vec_convergence,
run_weight, run_weight,
run_vec_weight, run_vec_weight,
offset_waypoint,
offset_tstep_fraction,
batch_pose_idx, batch_pose_idx,
batch_size, batch_size,
horizon, horizon,
@@ -189,6 +194,7 @@ def get_pose_distance(
write_grad, write_grad,
write_distance, write_distance,
use_metric, use_metric,
project_distance,
) )
out_distance = r[0] out_distance = r[0]
@@ -229,6 +235,331 @@ def get_pose_distance_backward(
return r[0], r[1] return r[0], r[1]
@torch.jit.script
def backward_PoseError_jit(grad_g_dist, grad_out_distance, weight, g_vec):
grad_vec = grad_g_dist + (grad_out_distance * weight)
grad = 1.0 * (grad_vec).unsqueeze(-1) * g_vec
return grad
# full method:
@torch.jit.script
def backward_full_PoseError_jit(
grad_out_distance, grad_g_dist, grad_r_err, p_w, q_w, g_vec_p, g_vec_q
):
p_grad = (grad_g_dist + (grad_out_distance * p_w)).unsqueeze(-1) * g_vec_p
q_grad = (grad_r_err + (grad_out_distance * q_w)).unsqueeze(-1) * g_vec_q
# p_grad = ((grad_out_distance * p_w)).unsqueeze(-1) * g_vec_p
# q_grad = ((grad_out_distance * q_w)).unsqueeze(-1) * g_vec_q
return p_grad, q_grad
class PoseErrorDistance(torch.autograd.Function):
@staticmethod
def forward(
ctx,
current_position,
goal_position,
current_quat,
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
offset_waypoint,
offset_tstep_fraction,
batch_pose_idx,
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
out_p_grad,
out_q_grad,
batch_size,
horizon,
mode, # =PoseErrorType.BATCH_GOAL.value,
num_goals,
use_metric, # =False,
project_distance, # =True,
):
# out_distance = current_position[..., 0].detach().clone() * 0.0
# out_position_distance = out_distance.detach().clone()
# out_rotation_distance = out_distance.detach().clone()
# out_vec = (
# torch.cat((current_position.detach().clone(), current_quat.detach().clone()), dim=-1)
# * 0.0
# )
# out_idx = out_distance.clone().to(dtype=torch.long)
(
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
) = get_pose_distance(
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
current_position.contiguous(),
goal_position,
current_quat.contiguous(),
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
offset_waypoint,
offset_tstep_fraction,
batch_pose_idx,
batch_size,
horizon,
mode,
num_goals,
current_position.requires_grad,
True,
use_metric,
project_distance,
)
ctx.save_for_backward(out_p_vec, out_r_vec, weight, out_p_grad, out_q_grad)
return out_distance, out_position_distance, out_rotation_distance, out_idx # .view(-1,1)
@staticmethod
def backward(ctx, grad_out_distance, grad_g_dist, grad_r_err, grad_out_idx):
(g_vec_p, g_vec_q, weight, out_grad_p, out_grad_q) = ctx.saved_tensors
pos_grad = None
quat_grad = None
batch_size = g_vec_p.shape[0] * g_vec_p.shape[1]
if ctx.needs_input_grad[0] and ctx.needs_input_grad[2]:
pos_grad, quat_grad = get_pose_distance_backward(
out_grad_p,
out_grad_q,
grad_out_distance.contiguous(),
grad_g_dist.contiguous(),
grad_r_err.contiguous(),
weight,
g_vec_p,
g_vec_q,
batch_size,
use_distance=True,
)
elif ctx.needs_input_grad[0]:
pos_grad = backward_PoseError_jit(grad_g_dist, grad_out_distance, weight[1], g_vec_p)
elif ctx.needs_input_grad[2]:
quat_grad = backward_PoseError_jit(grad_r_err, grad_out_distance, weight[0], g_vec_q)
return (
pos_grad,
None,
quat_grad,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
)
class PoseError(torch.autograd.Function):
@staticmethod
def forward(
ctx,
current_position: torch.Tensor,
goal_position: torch.Tensor,
current_quat: torch.Tensor,
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
offset_waypoint,
offset_tstep_fraction,
batch_pose_idx,
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
out_p_grad,
out_q_grad,
batch_size,
horizon,
mode,
num_goals,
use_metric,
project_distance,
return_loss,
):
"""Compute error in pose
_extended_summary_
Args:
ctx: _description_
current_position: _description_
goal_position: _description_
current_quat: _description_
goal_quat: _description_
vec_weight: _description_
weight: _description_
vec_convergence: _description_
run_weight: _description_
run_vec_weight: _description_
offset_waypoint: _description_
offset_tstep_fraction: _description_
batch_pose_idx: _description_
out_distance: _description_
out_position_distance: _description_
out_rotation_distance: _description_
out_p_vec: _description_
out_r_vec: _description_
out_idx: _description_
out_p_grad: _description_
out_q_grad: _description_
batch_size: _description_
horizon: _description_
mode: _description_
num_goals: _description_
use_metric: _description_
project_distance: _description_
return_loss: _description_
Returns:
_description_
"""
# out_distance = current_position[..., 0].detach().clone() * 0.0
# out_position_distance = out_distance.detach().clone()
# out_rotation_distance = out_distance.detach().clone()
# out_vec = (
# torch.cat((current_position.detach().clone(), current_quat.detach().clone()), dim=-1)
# * 0.0
# )
# out_idx = out_distance.clone().to(dtype=torch.long)
ctx.return_loss = return_loss
(
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
) = get_pose_distance(
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
current_position.contiguous(),
goal_position,
current_quat.contiguous(),
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
offset_waypoint,
offset_tstep_fraction,
batch_pose_idx,
batch_size,
horizon,
mode,
num_goals,
current_position.requires_grad,
False,
use_metric,
project_distance,
)
ctx.save_for_backward(out_p_vec, out_r_vec)
return out_distance
@staticmethod
def backward(ctx, grad_out_distance): # , grad_g_dist, grad_r_err, grad_out_idx):
pos_grad = None
quat_grad = None
if ctx.needs_input_grad[0] and ctx.needs_input_grad[2]:
(g_vec_p, g_vec_q) = ctx.saved_tensors
pos_grad = g_vec_p
quat_grad = g_vec_q
if ctx.return_loss:
pos_grad = pos_grad * grad_out_distance.unsqueeze(1)
quat_grad = quat_grad * grad_out_distance.unsqueeze(1)
elif ctx.needs_input_grad[0]:
(g_vec_p, g_vec_q) = ctx.saved_tensors
pos_grad = g_vec_p
if ctx.return_loss:
pos_grad = pos_grad * grad_out_distance.unsqueeze(1)
elif ctx.needs_input_grad[2]:
(g_vec_p, g_vec_q) = ctx.saved_tensors
quat_grad = g_vec_q
if ctx.return_loss:
quat_grad = quat_grad * grad_out_distance.unsqueeze(1)
return (
pos_grad,
None,
quat_grad,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
)
class SdfSphereOBB(torch.autograd.Function): class SdfSphereOBB(torch.autograd.Function):
@staticmethod @staticmethod
def forward( def forward(

118
src/curobo/geom/cv.py Normal file
View File

@@ -0,0 +1,118 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Third Party
import torch
@torch.jit.script
def project_depth_to_pointcloud(depth_image: torch.Tensor, intrinsics_matrix: torch.Tensor):
"""Projects numpy depth image to point cloud.
Args:
np_depth_image: numpy array float, shape (h, w).
intrinsics array: numpy array float, 3x3 intrinsics matrix.
Returns:
array of float (h, w, 3)
"""
fx = intrinsics_matrix[0, 0]
fy = intrinsics_matrix[1, 1]
cx = intrinsics_matrix[0, 2]
cy = intrinsics_matrix[1, 2]
height = depth_image.shape[0]
width = depth_image.shape[1]
input_x = torch.arange(width, dtype=torch.float32, device=depth_image.device)
input_y = torch.arange(height, dtype=torch.float32, device=depth_image.device)
input_x, input_y = torch.meshgrid(input_x, input_y, indexing="ij")
input_x, input_y = input_x.T, input_y.T
input_z = depth_image
output_x = (input_x * input_z - cx * input_z) / fx
output_y = (input_y * input_z - cy * input_z) / fy
raw_pc = torch.stack([output_x, output_y, input_z], -1)
return raw_pc
@torch.jit.script
def get_projection_rays(height: int, width: int, intrinsics_matrix: torch.Tensor):
"""Projects numpy depth image to point cloud.
Args:
np_depth_image: numpy array float, shape (h, w).
intrinsics array: numpy array float, 3x3 intrinsics matrix.
Returns:
array of float (h, w, 3)
"""
fx = intrinsics_matrix[:, 0, 0]
fy = intrinsics_matrix[:, 1, 1]
cx = intrinsics_matrix[:, 0, 2]
cy = intrinsics_matrix[:, 1, 2]
input_x = torch.arange(width, dtype=torch.float32, device=intrinsics_matrix.device)
input_y = torch.arange(height, dtype=torch.float32, device=intrinsics_matrix.device)
input_x, input_y = torch.meshgrid(input_x, input_y, indexing="ij")
input_x, input_y = input_x.T, input_y.T
input_x = input_x.unsqueeze(0).repeat(intrinsics_matrix.shape[0], 1, 1)
input_y = input_y.unsqueeze(0).repeat(intrinsics_matrix.shape[0], 1, 1)
input_z = torch.ones(
(intrinsics_matrix.shape[0], height, width),
device=intrinsics_matrix.device,
dtype=torch.float32,
)
output_x = (input_x - cx) / fx
output_y = (input_y - cy) / fy
rays = torch.stack([output_x, output_y, input_z], -1).reshape(
intrinsics_matrix.shape[0], width * height, 3
)
rays = rays * (1.0 / 1000.0)
return rays
@torch.jit.script
def project_pointcloud_to_depth(
pointcloud: torch.Tensor,
output_image: torch.Tensor,
):
"""Projects pointcloud to depth image
Args:
np_depth_image: numpy array float, shape (h, w).
intrinsics array: numpy array float, 3x3 intrinsics matrix.
Returns:
array of float (h, w)
"""
width, height = output_image.shape
output_image = output_image.view(-1)
output_image[:] = pointcloud[:, 2]
output_image = output_image.view(width, height)
return output_image
@torch.jit.script
def project_depth_using_rays(
depth_image: torch.Tensor, rays: torch.Tensor, filter_origin: bool = False
):
if filter_origin:
depth_image = torch.where(depth_image < 0.01, 0, depth_image)
depth_image = depth_image.view(depth_image.shape[0], -1, 1).contiguous()
points = depth_image * rays
return points

19
src/curobo/geom/sdf/world.py
View File

@@ -364,8 +364,12 @@ class WorldPrimitiveCollision(WorldCollision):
if self.cache is not None and "obb" in self.cache and self.cache["obb"] not in [None, 0]: if self.cache is not None and "obb" in self.cache and self.cache["obb"] not in [None, 0]:
self._create_obb_cache(self.cache["obb"]) self._create_obb_cache(self.cache["obb"])
def load_collision_model(self, world_config: WorldConfig, env_idx=0): def load_collision_model(
self._load_collision_model_in_cache(world_config, env_idx) self, world_config: WorldConfig, env_idx=0, fix_cache_reference: bool = False
):
self._load_collision_model_in_cache(
world_config, env_idx, fix_cache_reference=fix_cache_reference
)
def load_batch_collision_model(self, world_config_list: List[WorldConfig]): def load_batch_collision_model(self, world_config_list: List[WorldConfig]):
"""Load a batch of collision environments from a list of world configs. """Load a batch of collision environments from a list of world configs.
@@ -436,7 +440,9 @@ class WorldPrimitiveCollision(WorldCollision):
) )
self.collision_types["primitive"] = True self.collision_types["primitive"] = True
def _load_collision_model_in_cache(self, world_config: WorldConfig, env_idx: int = 0): def _load_collision_model_in_cache(
self, world_config: WorldConfig, env_idx: int = 0, fix_cache_reference: bool = False
):
cube_objs = world_config.cuboid cube_objs = world_config.cuboid
max_obb = len(cube_objs) max_obb = len(cube_objs)
self.world_model = world_config self.world_model = world_config
@@ -444,8 +450,11 @@ class WorldPrimitiveCollision(WorldCollision):
log_info("No OBB objs") log_info("No OBB objs")
return return
if self._cube_tensor_list is None or self._cube_tensor_list[0].shape[1] < max_obb: if self._cube_tensor_list is None or self._cube_tensor_list[0].shape[1] < max_obb:
if not fix_cache_reference:
log_info("Creating Obb cache" + str(max_obb)) log_info("Creating Obb cache" + str(max_obb))
self._create_obb_cache(max_obb) self._create_obb_cache(max_obb)
else:
log_error("number of OBB is larger than collision cache, create larger cache.")
# load as a batch: # load as a batch:
pose_batch = [c.pose for c in cube_objs] pose_batch = [c.pose for c in cube_objs]
@@ -835,7 +844,9 @@ class WorldPrimitiveCollision(WorldCollision):
return dist return dist
def clear_cache(self): def clear_cache(self):
pass if self._cube_tensor_list is not None:
self._cube_tensor_list[2][:] = 1
self._env_n_obbs[:] = 0
def get_voxels_in_bounding_box( def get_voxels_in_bounding_box(
self, self,

11
src/curobo/geom/sdf/world_blox.py
View File

@@ -55,7 +55,7 @@ class WorldBloxCollision(WorldMeshCollision):
self._blox_voxel_sizes = [0.02] self._blox_voxel_sizes = [0.02]
super().__init__(config) super().__init__(config)
def load_collision_model(self, world_model: WorldConfig): def load_collision_model(self, world_model: WorldConfig, fix_cache_reference: bool = False):
# load nvblox mesh # load nvblox mesh
if len(world_model.blox) > 0: if len(world_model.blox) > 0:
# check if there is a mapper instance: # check if there is a mapper instance:
@@ -109,15 +109,17 @@ class WorldBloxCollision(WorldMeshCollision):
self._blox_names = names self._blox_names = names
self.collision_types["blox"] = True self.collision_types["blox"] = True
return super().load_collision_model(world_model) return super().load_collision_model(world_model, fix_cache_reference=fix_cache_reference)
def clear_cache(self): def clear_cache(self):
self._blox_mapper.clear() self._blox_mapper.clear()
self._blox_mapper.update_hashmaps()
super().clear_cache() super().clear_cache()
def clear_blox_layer(self, layer_name: str): def clear_blox_layer(self, layer_name: str):
index = self._blox_names.index(layer_name) index = self._blox_names.index(layer_name)
self._blox_mapper.clear(index) self._blox_mapper.clear(index)
self._blox_mapper.update_hashmaps()
def _get_blox_sdf( def _get_blox_sdf(
self, self,
@@ -147,6 +149,7 @@ class WorldBloxCollision(WorldMeshCollision):
speed_dt, speed_dt,
sweep_steps, sweep_steps,
enable_speed_metric, enable_speed_metric,
return_loss=False,
): ):
d = self._blox_mapper.query_sphere_trajectory_sdf_cost( d = self._blox_mapper.query_sphere_trajectory_sdf_cost(
query_spheres, query_spheres,
@@ -160,6 +163,8 @@ class WorldBloxCollision(WorldMeshCollision):
self._blox_tensor_list[1], self._blox_tensor_list[1],
sweep_steps, sweep_steps,
enable_speed_metric, enable_speed_metric,
return_loss,
use_experimental=False,
) )
return d return d
@@ -279,6 +284,7 @@ class WorldBloxCollision(WorldMeshCollision):
speed_dt=speed_dt, speed_dt=speed_dt,
sweep_steps=sweep_steps, sweep_steps=sweep_steps,
enable_speed_metric=enable_speed_metric, enable_speed_metric=enable_speed_metric,
return_loss=return_loss,
) )
if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and ( if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
@@ -332,6 +338,7 @@ class WorldBloxCollision(WorldMeshCollision):
speed_dt=speed_dt, speed_dt=speed_dt,
sweep_steps=sweep_steps, sweep_steps=sweep_steps,
enable_speed_metric=enable_speed_metric, enable_speed_metric=enable_speed_metric,
return_loss=return_loss,
) )
if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and ( if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (

12
src/curobo/geom/sdf/world_mesh.py
View File

@@ -72,7 +72,11 @@ class WorldMeshCollision(WorldPrimitiveCollision):
return super()._init_cache() return super()._init_cache()
def load_collision_model( def load_collision_model(
self, world_model: WorldConfig, env_idx: int = 0, load_obb_obs: bool = True self,
world_model: WorldConfig,
env_idx: int = 0,
load_obb_obs: bool = True,
fix_cache_reference: bool = False,
): ):
max_nmesh = len(world_model.mesh) max_nmesh = len(world_model.mesh)
if max_nmesh > 0: if max_nmesh > 0:
@@ -91,14 +95,16 @@ class WorldMeshCollision(WorldPrimitiveCollision):
self.collision_types["mesh"] = True self.collision_types["mesh"] = True
if load_obb_obs: if load_obb_obs:
super().load_collision_model(world_model, env_idx) super().load_collision_model(
world_model, env_idx, fix_cache_reference=fix_cache_reference
)
else: else:
self.world_model = world_model self.world_model = world_model
def load_batch_collision_model(self, world_config_list: List[WorldConfig]): def load_batch_collision_model(self, world_config_list: List[WorldConfig]):
max_nmesh = max([len(x.mesh) for x in world_config_list]) max_nmesh = max([len(x.mesh) for x in world_config_list])
if self._mesh_tensor_list is None or self._mesh_tensor_list[0].shape[1] < max_nmesh: if self._mesh_tensor_list is None or self._mesh_tensor_list[0].shape[1] < max_nmesh:
log_info("Creating new Mesh cache: " + str(max_nmesh)) log_warn("Creating new Mesh cache: " + str(max_nmesh))
self._create_mesh_cache(max_nmesh) self._create_mesh_cache(max_nmesh)
for env_idx, world_model in enumerate(world_config_list): for env_idx, world_model in enumerate(world_config_list):

113
src/curobo/geom/transform.py
View File

@@ -228,6 +228,61 @@ def pose_inverse(
return out_position, out_quaternion return out_position, out_quaternion
@wp.kernel
def compute_pose_inverse(
position: wp.array(dtype=wp.vec3),
quat: wp.array(dtype=wp.vec4),
out_position: wp.array(dtype=wp.vec3),
out_quat: wp.array(dtype=wp.vec4),
): # b pose_1 and b pose_2, compute pose_1 * pose_2
b_idx = wp.tid()
# read data:
in_position = position[b_idx]
in_quat = quat[b_idx]
# read point
# create a transform from a vector/quaternion:
t_1 = wp.transform(in_position, wp.quaternion(in_quat[1], in_quat[2], in_quat[3], in_quat[0]))
t_3 = wp.transform_inverse(t_1)
# write pt:
out_q = wp.transform_get_rotation(t_3)
out_v = wp.vec4()
out_v[0] = out_q[3] # out_q[3]
out_v[1] = out_q[0] # [0]
out_v[2] = out_q[1] # wp.extract(out_q, 1)
out_v[3] = out_q[2] # wp.extract(out_q, 2)
out_position[b_idx] = wp.transform_get_translation(t_3)
out_quat[b_idx] = out_v
@wp.kernel
def compute_matrix_to_quat(
in_mat: wp.array(dtype=wp.mat33),
out_quat: wp.array(dtype=wp.vec4),
):
# b pose_1 and b pose_2, compute pose_1 * pose_2
b_idx = wp.tid()
# read data:
in_m = in_mat[b_idx]
# read point
# create a transform from a vector/quaternion:
out_q = wp.quat_from_matrix(in_m)
out_v = wp.vec4()
out_v[0] = out_q[3] # wp.extract(out_q, 3)
out_v[1] = out_q[0] # wp.extract(out_q, 0)
out_v[2] = out_q[1] # wp.extract(out_q, 1)
out_v[3] = out_q[2] # wp.extract(out_q, 2)
# write pt:
out_quat[b_idx] = out_v
@wp.kernel @wp.kernel
def compute_transform_point( def compute_transform_point(
position: wp.array(dtype=wp.vec3), position: wp.array(dtype=wp.vec3),
@@ -331,37 +386,6 @@ def compute_batch_pose_multipy(
out_quat[b_idx] = out_v out_quat[b_idx] = out_v
@wp.kernel
def compute_pose_inverse(
position: wp.array(dtype=wp.vec3),
quat: wp.array(dtype=wp.vec4),
out_position: wp.array(dtype=wp.vec3),
out_quat: wp.array(dtype=wp.vec4),
): # b pose_1 and b pose_2, compute pose_1 * pose_2
b_idx = wp.tid()
# read data:
in_position = position[b_idx]
in_quat = quat[b_idx]
# read point
# create a transform from a vector/quaternion:
t_1 = wp.transform(in_position, wp.quaternion(in_quat[1], in_quat[2], in_quat[3], in_quat[0]))
t_3 = wp.transform_inverse(t_1)
# write pt:
out_q = wp.transform_get_rotation(t_3)
out_v = wp.vec4()
out_v[0] = out_q[3]
out_v[1] = out_q[0]
out_v[2] = out_q[1]
out_v[3] = out_q[2]
out_position[b_idx] = wp.transform_get_translation(t_3)
out_quat[b_idx] = out_v
@wp.kernel @wp.kernel
def compute_quat_to_matrix( def compute_quat_to_matrix(
quat: wp.array(dtype=wp.vec4), quat: wp.array(dtype=wp.vec4),
@@ -382,30 +406,6 @@ def compute_quat_to_matrix(
out_mat[b_idx] = m_1 out_mat[b_idx] = m_1
@wp.kernel
def compute_matrix_to_quat(
in_mat: wp.array(dtype=wp.mat33),
out_quat: wp.array(dtype=wp.vec4),
):
# b pose_1 and b pose_2, compute pose_1 * pose_2
b_idx = wp.tid()
# read data:
in_m = in_mat[b_idx]
# read point
# create a transform from a vector/quaternion:
out_q = wp.quat_from_matrix(in_m)
out_v = wp.vec4()
out_v[0] = out_q[3]
out_v[1] = out_q[0]
out_v[2] = out_q[1]
out_v[3] = out_q[2]
# write pt:
out_quat[b_idx] = out_v
@wp.kernel @wp.kernel
def compute_pose_multipy( def compute_pose_multipy(
position: wp.array(dtype=wp.vec3), position: wp.array(dtype=wp.vec3),
@@ -962,6 +962,7 @@ class PoseInverse(torch.autograd.Function):
adj_quaternion, adj_quaternion,
) )
ctx.b = b ctx.b = b
wp.launch( wp.launch(
kernel=compute_pose_inverse, kernel=compute_pose_inverse,
dim=b, dim=b,
@@ -1071,6 +1072,7 @@ class QuatToMatrix(torch.autograd.Function):
adj_quaternion, adj_quaternion,
) )
ctx.b = b ctx.b = b
wp.launch( wp.launch(
kernel=compute_quat_to_matrix, kernel=compute_quat_to_matrix,
dim=b, dim=b,
@@ -1153,6 +1155,7 @@ class MatrixToQuaternion(torch.autograd.Function):
adj_mat, adj_mat,
) )
ctx.b = b ctx.b = b
wp.launch( wp.launch(
kernel=compute_matrix_to_quat, kernel=compute_matrix_to_quat,
dim=b, dim=b,

135
src/curobo/geom/types.py
View File

@@ -12,7 +12,7 @@ from __future__ import annotations
# Standard Library # Standard Library
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence from typing import Any, Dict, List, Optional, Sequence, Union
# Third Party # Third Party
import numpy as np import numpy as np
@@ -22,6 +22,7 @@ import trimesh
# CuRobo # CuRobo
from curobo.geom.sphere_fit import SphereFitType, fit_spheres_to_mesh from curobo.geom.sphere_fit import SphereFitType, fit_spheres_to_mesh
from curobo.types.base import TensorDeviceType from curobo.types.base import TensorDeviceType
from curobo.types.camera import CameraObservation
from curobo.types.math import Pose from curobo.types.math import Pose
from curobo.util.logger import log_error, log_warn from curobo.util.logger import log_error, log_warn
from curobo.util_file import get_assets_path, join_path from curobo.util_file import get_assets_path, join_path
@@ -60,7 +61,7 @@ class Obstacle:
tensor_args: TensorDeviceType = TensorDeviceType() tensor_args: TensorDeviceType = TensorDeviceType()
def get_trimesh_mesh(self, process: bool = True) -> trimesh.Trimesh: def get_trimesh_mesh(self, process: bool = True, process_color: bool = True) -> trimesh.Trimesh:
"""Create a trimesh instance from the obstacle representation. """Create a trimesh instance from the obstacle representation.
Args: Args:
@@ -74,8 +75,11 @@ class Obstacle:
""" """
raise NotImplementedError raise NotImplementedError
def save_as_mesh(self, file_path: str): def save_as_mesh(self, file_path: str, transform_with_pose: bool = False):
mesh_scene = self.get_trimesh_mesh() mesh_scene = self.get_trimesh_mesh()
if transform_with_pose:
mesh_scene.apply_transform(self.get_transform_matrix())
mesh_scene.export(file_path) mesh_scene.export(file_path)
def get_cuboid(self) -> Cuboid: def get_cuboid(self) -> Cuboid:
@@ -238,7 +242,7 @@ class Cuboid(Obstacle):
if self.pose is None: if self.pose is None:
log_error("Cuboid Obstacle requires Pose") log_error("Cuboid Obstacle requires Pose")
def get_trimesh_mesh(self, process: bool = True): def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
m = trimesh.creation.box(extents=self.dims) m = trimesh.creation.box(extents=self.dims)
if self.color is not None: if self.color is not None:
color_visual = trimesh.visual.color.ColorVisuals( color_visual = trimesh.visual.color.ColorVisuals(
@@ -254,7 +258,7 @@ class Capsule(Obstacle):
base: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0]) base: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
tip: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0]) tip: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
def get_trimesh_mesh(self, process: bool = True): def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
height = self.tip[2] - self.base[2] height = self.tip[2] - self.base[2]
if ( if (
height < 0 height < 0
@@ -280,7 +284,7 @@ class Cylinder(Obstacle):
radius: float = 0.0 radius: float = 0.0
height: float = 0.0 height: float = 0.0
def get_trimesh_mesh(self, process: bool = True): def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
m = trimesh.creation.cylinder(radius=self.radius, height=self.height) m = trimesh.creation.cylinder(radius=self.radius, height=self.height)
if self.color is not None: if self.color is not None:
color_visual = trimesh.visual.color.ColorVisuals( color_visual = trimesh.visual.color.ColorVisuals(
@@ -304,7 +308,7 @@ class Sphere(Obstacle):
if self.pose is not None: if self.pose is not None:
self.position = self.pose[:3] self.position = self.pose[:3]
def get_trimesh_mesh(self, process: bool = True): def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
m = trimesh.creation.icosphere(radius=self.radius) m = trimesh.creation.icosphere(radius=self.radius)
if self.color is not None: if self.color is not None:
color_visual = trimesh.visual.color.ColorVisuals( color_visual = trimesh.visual.color.ColorVisuals(
@@ -356,9 +360,9 @@ class Sphere(Obstacle):
@dataclass @dataclass
class Mesh(Obstacle): class Mesh(Obstacle):
file_path: Optional[str] = None file_path: Optional[str] = None
file_string: Optional[ file_string: Optional[str] = (
str None # storing full mesh as a string, loading from this is not implemented yet.
] = None # storing full mesh as a string, loading from this is not implemented yet. )
urdf_path: Optional[str] = None # useful for visualization in isaac gym. urdf_path: Optional[str] = None # useful for visualization in isaac gym.
vertices: Optional[List[List[float]]] = None vertices: Optional[List[List[float]]] = None
faces: Optional[List[int]] = None faces: Optional[List[int]] = None
@@ -375,13 +379,13 @@ class Mesh(Obstacle):
self.vertices = np.ravel(self.scale) * self.vertices self.vertices = np.ravel(self.scale) * self.vertices
self.scale = None self.scale = None
def get_trimesh_mesh(self, process: bool = True): def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
# load mesh from filepath or verts and faces: # load mesh from filepath or verts and faces:
if self.file_path is not None: if self.file_path is not None:
m = trimesh.load(self.file_path, process=process, force="mesh") m = trimesh.load(self.file_path, process=process, force="mesh")
if isinstance(m, trimesh.Scene): if isinstance(m, trimesh.Scene):
m = m.dump(concatenate=True) m = m.dump(concatenate=True)
if isinstance(m.visual, trimesh.visual.texture.TextureVisuals): if process_color and isinstance(m.visual, trimesh.visual.texture.TextureVisuals):
m.visual = m.visual.to_color() m.visual = m.visual.to_color()
else: else:
m = trimesh.Trimesh( m = trimesh.Trimesh(
@@ -467,7 +471,7 @@ class BloxMap(Obstacle):
name=self.name + "_mesh", file_path=self.mesh_file_path, pose=self.pose name=self.name + "_mesh", file_path=self.mesh_file_path, pose=self.pose
) )
def get_trimesh_mesh(self, process: bool = True): def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
if self.mesh is not None: if self.mesh is not None:
return self.mesh.get_trimesh_mesh(process) return self.mesh.get_trimesh_mesh(process)
else: else:
@@ -475,6 +479,91 @@ class BloxMap(Obstacle):
return None return None
@dataclass
class PointCloud(Obstacle):
points: Union[torch.Tensor, np.ndarray, List[List[float]]] = None
points_features: Union[torch.Tensor, np.ndarray, List[List[float]], None] = None
def __post_init__(self):
if self.scale is not None and self.points is not None:
self.points = np.ravel(self.scale) * self.points
self.scale = None
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
points = self.points
if isinstance(points, torch.Tensor):
points = points.view(-1, 3).cpu().numpy()
if isinstance(points, list):
points = np.ndarray(points)
mesh = Mesh.from_pointcloud(points, pose=self.pose)
return mesh.get_trimesh_mesh()
def get_mesh_data(self, process: bool = True):
verts = faces = None
m = self.get_trimesh_mesh(process=process)
verts = m.vertices.view(np.ndarray)
faces = m.faces
return verts, faces
@staticmethod
def from_camera_observation(
camera_obs: CameraObservation, name: str = "pc_obstacle", pose: Optional[List[float]] = None
):
return PointCloud(name=name, pose=pose, points=camera_obs.get_pointcloud())
def get_bounding_spheres(
self,
n_spheres: int = 1,
surface_sphere_radius: float = 0.002,
fit_type: SphereFitType = SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE,
voxelize_method: str = "ray",
pre_transform_pose: Optional[Pose] = None,
tensor_args: TensorDeviceType = TensorDeviceType(),
) -> List[Sphere]:
"""Compute n spheres that fits in the volume of the object.
Args:
n: number of spheres
Returns:
spheres
"""
# sample points in pointcloud:
# mesh = self.get_trimesh_mesh()
# pts, n_radius = fit_spheres_to_mesh(
# mesh, n_spheres, surface_sphere_radius, fit_type, voxelize_method=voxelize_method
# )
obj_pose = Pose.from_list(self.pose, tensor_args)
# transform object:
# transform points:
if pre_transform_pose is not None:
obj_pose = pre_transform_pose.multiply(obj_pose) # convert object pose to another frame
if pts is None or len(pts) == 0:
log_warn("spheres could not be fit!, adding one sphere at origin")
pts = np.zeros((1, 3))
pts[0, :] = mesh.centroid
n_radius = [0.02]
obj_pose = Pose.from_list([0, 0, 0, 1, 0, 0, 0], tensor_args)
points_cuda = tensor_args.to_device(pts)
pts = obj_pose.transform_points(points_cuda).cpu().view(-1, 3).numpy()
new_spheres = [
Sphere(
name="sph_" + str(i),
pose=[pts[i, 0], pts[i, 1], pts[i, 2], 1, 0, 0, 0],
radius=n_radius[i],
)
for i in range(pts.shape[0])
]
return new_spheres
@dataclass @dataclass
class WorldConfig(Sequence): class WorldConfig(Sequence):
"""Representation of World for use in CuRobo.""" """Representation of World for use in CuRobo."""
@@ -664,23 +753,25 @@ class WorldConfig(Sequence):
) )
@staticmethod @staticmethod
def get_scene_graph(current_world: WorldConfig): def get_scene_graph(current_world: WorldConfig, process_color: bool = True):
m_world = WorldConfig.create_mesh_world(current_world) m_world = WorldConfig.create_mesh_world(current_world)
mesh_scene = trimesh.scene.scene.Scene(base_frame="world") mesh_scene = trimesh.scene.scene.Scene(base_frame="world_origin")
mesh_list = m_world mesh_list = m_world
for m in mesh_list: for m in mesh_list:
mesh_scene.add_geometry( mesh_scene.add_geometry(
m.get_trimesh_mesh(), m.get_trimesh_mesh(process_color=process_color),
geom_name=m.name, geom_name=m.name,
parent_node_name="world", parent_node_name="world_origin",
transform=m.get_transform_matrix(), transform=m.get_transform_matrix(),
) )
return mesh_scene return mesh_scene
@staticmethod @staticmethod
def create_merged_mesh_world(current_world: WorldConfig, process: bool = True): def create_merged_mesh_world(
mesh_scene = WorldConfig.get_scene_graph(current_world) current_world: WorldConfig, process: bool = True, process_color: bool = True
):
mesh_scene = WorldConfig.get_scene_graph(current_world, process_color=process_color)
mesh_scene = mesh_scene.dump(concatenate=True) mesh_scene = mesh_scene.dump(concatenate=True)
new_mesh = Mesh( new_mesh = Mesh(
vertices=mesh_scene.vertices.view(np.ndarray), vertices=mesh_scene.vertices.view(np.ndarray),
@@ -702,8 +793,10 @@ class WorldConfig(Sequence):
def get_collision_check_world(self, mesh_process: bool = False): def get_collision_check_world(self, mesh_process: bool = False):
return WorldConfig.create_collision_support_world(self, process=mesh_process) return WorldConfig.create_collision_support_world(self, process=mesh_process)
def save_world_as_mesh(self, file_path: str, save_as_scene_graph=False): def save_world_as_mesh(
mesh_scene = WorldConfig.get_scene_graph(self) self, file_path: str, save_as_scene_graph=False, process_color: bool = True
):
mesh_scene = WorldConfig.get_scene_graph(self, process_color=process_color)
if save_as_scene_graph: if save_as_scene_graph:
mesh_scene = mesh_scene.dump(concatenate=True) mesh_scene = mesh_scene.dump(concatenate=True)

9
src/curobo/graph/graph_nx.py
View File

@@ -10,10 +10,19 @@
# #
# Standard Library
import random
# Third Party # Third Party
import networkx as nx import networkx as nx
import numpy as np
import torch import torch
# This is needed to get deterministic results from networkx.
# Note: it has to be set in global space.
np.random.seed(2)
random.seed(2)
class NetworkxGraph(object): class NetworkxGraph(object):
def __init__(self): def __init__(self):

12
src/curobo/opt/__init__.py
View File

@@ -8,3 +8,15 @@
# without an express license agreement from NVIDIA CORPORATION or # without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited. # its affiliates is strictly prohibited.
# #
""" Optimization module containing several numerical solvers.
Base for an opimization solver is at :class:`opt_base.Optimizer`. cuRobo provides two base classes
for implementing two popular ways to optimize, (1) using particles
with :class:`particle.particle_opt_base.ParticleOptBase` and (2) using Newton/Quasi-Newton solvers
with :class:`newton.newton_base.NewtonOptBase`. :class:`newton.newton_base.NewtonOptBase` contains
implementations of several line search schemes. Note that these line search schemes are approximate
as cuRobo tries different line search magnitudes in parallel and chooses the largest that satisfies
line search conditions.
"""

2
src/curobo/opt/newton/__init__.py
View File

@@ -10,5 +10,5 @@
# #
""" """
This module contains code for cuda accelerated kinematics. This module contains Newton/Quasi-Newton solvers.
""" """

6
src/curobo/opt/newton/lbfgs.py
View File

@@ -95,9 +95,9 @@ class LBFGSOpt(NewtonOptBase, LBFGSOptConfig):
self.step_q_buffer[:] = 0.0 self.step_q_buffer[:] = 0.0
return super().reset() return super().reset()
def update_nenvs(self, n_envs): def update_nproblems(self, n_problems):
self.init_hessian(b=n_envs) self.init_hessian(b=n_problems)
return super().update_nenvs(n_envs) return super().update_nproblems(n_problems)
def init_hessian(self, b=1): def init_hessian(self, b=1):
self.x_0 = torch.zeros( self.x_0 = torch.zeros(

37
src/curobo/opt/newton/newton_base.py
View File

@@ -79,7 +79,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
self.outer_iters = math.ceil(self.n_iters / self.inner_iters) self.outer_iters = math.ceil(self.n_iters / self.inner_iters)
# create line search # create line search
self.update_nenvs(self.n_envs) self.update_nproblems(self.n_problems)
self.reset() self.reset()
@@ -135,7 +135,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
if self.store_debug: if self.store_debug:
self.debug.append(q.view(-1, self.action_horizon, self.d_action).clone()) self.debug.append(q.view(-1, self.action_horizon, self.d_action).clone())
with profiler.record_function("newton_base/init_opt"): with profiler.record_function("newton_base/init_opt"):
q = q.view(self.n_envs, self.action_horizon * self.d_action) q = q.view(self.n_problems, self.action_horizon * self.d_action)
grad_q = q.detach() * 0.0 grad_q = q.detach() * 0.0
# run opt graph # run opt graph
if not self.cu_opt_init: if not self.cu_opt_init:
@@ -150,7 +150,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
if check_convergence(self.best_iteration, self.current_iteration, self.last_best): if check_convergence(self.best_iteration, self.current_iteration, self.last_best):
break break
best_q = best_q.view(self.n_envs, self.action_horizon, self.d_action) best_q = best_q.view(self.n_problems, self.action_horizon, self.d_action)
return best_q return best_q
def reset(self): def reset(self):
@@ -171,9 +171,6 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
if self.store_debug: if self.store_debug:
self.debug.append(self.best_q.view(-1, self.action_horizon, self.d_action).clone()) self.debug.append(self.best_q.view(-1, self.action_horizon, self.d_action).clone())
self.debug_cost.append(self.best_cost.detach().view(-1, 1).clone()) self.debug_cost.append(self.best_cost.detach().view(-1, 1).clone())
# self.debug.append(q.view(-1, self.action_horizon, self.d_action).clone())
# self.debug_cost.append(cost_n.detach().view(-1, 1).clone())
# print(grad_q)
return self.best_q.detach(), self.best_cost.detach(), q.detach(), grad_q.detach() return self.best_q.detach(), self.best_cost.detach(), q.detach(), grad_q.detach()
@@ -222,11 +219,11 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
def _compute_cost_gradient(self, x): def _compute_cost_gradient(self, x):
x_n = x.detach().requires_grad_(True) x_n = x.detach().requires_grad_(True)
x_in = x_n.view( x_in = x_n.view(
self.n_envs * self.num_particles, self.action_horizon, self.rollout_fn.d_action self.n_problems * self.num_particles, self.action_horizon, self.rollout_fn.d_action
) )
trajectories = self.rollout_fn(x_in) # x_n = (batch*line_search_scale) x horizon x d_action trajectories = self.rollout_fn(x_in) # x_n = (batch*line_search_scale) x horizon x d_action
cost = torch.sum( cost = torch.sum(
trajectories.costs.view(self.n_envs, self.num_particles, self.horizon), trajectories.costs.view(self.n_problems, self.num_particles, self.horizon),
dim=-1, dim=-1,
keepdim=True, keepdim=True,
) )
@@ -235,7 +232,7 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
return ( return (
cost, cost,
g_x, g_x,
) # cost: [n_envs, n_particles, 1], g_x: [n_envs, n_particles, horizon*d_action] ) # cost: [n_problems, n_particles, 1], g_x: [n_problems, n_particles, horizon*d_action]
def _wolfe_line_search(self, x, step_direction): def _wolfe_line_search(self, x, step_direction):
# x_set = get_x_set_jit(step_direction, x, self.alpha_list, self.action_lows, self.action_highs) # x_set = get_x_set_jit(step_direction, x, self.alpha_list, self.action_lows, self.action_highs)
@@ -455,36 +452,40 @@ class NewtonOptBase(Optimizer, NewtonOptConfig):
mask_q = mask.unsqueeze(-1).expand(-1, self.d_opt) mask_q = mask.unsqueeze(-1).expand(-1, self.d_opt)
self.best_q.copy_(torch.where(mask_q, q, self.best_q)) self.best_q.copy_(torch.where(mask_q, q, self.best_q))
def update_nenvs(self, n_envs): def update_nproblems(self, n_problems):
self.l_vec = torch.ones( self.l_vec = torch.ones(
(n_envs, self.num_particles, 1), (n_problems, self.num_particles, 1),
device=self.tensor_args.device, device=self.tensor_args.device,
dtype=self.tensor_args.dtype, dtype=self.tensor_args.dtype,
) )
self.best_cost = ( self.best_cost = (
torch.ones((n_envs, 1), device=self.tensor_args.device, dtype=self.tensor_args.dtype) torch.ones(
(n_problems, 1), device=self.tensor_args.device, dtype=self.tensor_args.dtype
)
* 5000000.0 * 5000000.0
) )
self.best_q = torch.zeros( self.best_q = torch.zeros(
(n_envs, self.d_opt), device=self.tensor_args.device, dtype=self.tensor_args.dtype (n_problems, self.d_opt), device=self.tensor_args.device, dtype=self.tensor_args.dtype
) )
self.best_grad_q = torch.zeros( self.best_grad_q = torch.zeros(
(n_envs, 1, self.d_opt), device=self.tensor_args.device, dtype=self.tensor_args.dtype (n_problems, 1, self.d_opt),
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
) )
# create list: # create list:
self.alpha_list = self.line_scale.repeat(n_envs, 1, 1) self.alpha_list = self.line_scale.repeat(n_problems, 1, 1)
self.zero_alpha_list = self.alpha_list[:, :, 0:1].contiguous() self.zero_alpha_list = self.alpha_list[:, :, 0:1].contiguous()
h = self.alpha_list.shape[1] h = self.alpha_list.shape[1]
self.c_idx = torch.arange( self.c_idx = torch.arange(
0, n_envs * h, step=(h), device=self.tensor_args.device, dtype=torch.long 0, n_problems * h, step=(h), device=self.tensor_args.device, dtype=torch.long
) )
self.best_iteration = torch.zeros( self.best_iteration = torch.zeros(
(n_envs), device=self.tensor_args.device, dtype=torch.int16 (n_problems), device=self.tensor_args.device, dtype=torch.int16
) )
self.current_iteration = torch.zeros((1), device=self.tensor_args.device, dtype=torch.int16) self.current_iteration = torch.zeros((1), device=self.tensor_args.device, dtype=torch.int16)
self.cu_opt_init = False self.cu_opt_init = False
super().update_nenvs(n_envs) super().update_nproblems(n_problems)
def _initialize_opt_iters_graph(self, q, grad_q, shift_steps): def _initialize_opt_iters_graph(self, q, grad_q, shift_steps):
if self.use_cuda_graph: if self.use_cuda_graph:

272
src/curobo/opt/opt_base.py
View File

@@ -8,6 +8,9 @@
# without an express license agreement from NVIDIA CORPORATION or # without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited. # its affiliates is strictly prohibited.
# #
""" Base module for Optimization. """
from __future__ import annotations
# Standard Library # Standard Library
import time import time
from abc import abstractmethod from abc import abstractmethod
@@ -28,31 +31,75 @@ from curobo.util.torch_utils import is_cuda_graph_available
@dataclass @dataclass
class OptimizerConfig: class OptimizerConfig:
"""Configuration for an :meth:`Optimizer`."""
#: Number of optimization variables per timestep.
d_action: int d_action: int
#: Lower bound for optimization variables.
action_lows: List[float] action_lows: List[float]
#: Higher bound for optimization variables
action_highs: List[float] action_highs: List[float]
horizon: int
n_iters: int #:
rollout_fn: RolloutBase
tensor_args: TensorDeviceType
use_cuda_graph: bool
store_debug: bool
debug_info: Any
cold_start_n_iters: int
num_particles: Union[int, None]
n_envs: int
sync_cuda_time: bool
use_coo_sparse: bool
action_horizon: int action_horizon: int
#: Number of timesteps in optimization state, total variables = d_action * horizon
horizon: int
#: Number of iterations to run optimization
n_iters: int
#: Number of iterations to run optimization during the first instance. Setting to None will
#: use n_iters. This parameter is useful in MPC like settings where we need to run many
#: iterations during initialization (cold start) and then run only few iterations (warm start).
cold_start_n_iters: Union[int, None]
#: Rollout function to use for computing cost, given optimization variables.
rollout_fn: RolloutBase
#: Tensor device to use for optimization.
tensor_args: TensorDeviceType
#: Capture optimization iteration in a cuda graph and replay graph instead of eager execution.
#: Enabling this can make optimization 10x faster. But changing control flow, tensor
#: shapes, or problem type is not allowed.
use_cuda_graph: bool
#: Record debugging data such as optimization variables, and cost at every iteration. Enabling
#: this will disable cuda graph.
store_debug: bool
#: Use this to record additional attributes from rollouts.
debug_info: Any
#: Number of parallel problems to optimize.
n_problems: int
#: Number of particles to use per problem. Common optimization solvers use many particles to
#: optimize a single problem. E.g., MPPI rolls out many parallel samples and computes a weighted
#: mean. In cuRobo, Quasi-Newton solvers use particles to run many line search magnitudes.
#: Total optimization batch size = n_problems * num_particles.
num_particles: Union[int, None]
#: Synchronize device before computing solver time.
sync_cuda_time: bool
#: Matmul with a Sparse tensor is used to create particles for each problem index to save memory
#: and compute. Some versions of pytorch do not support coo sparse, specifically during
#: torch profile runs. Set this to False to use a standard tensor.
use_coo_sparse: bool
def __post_init__(self): def __post_init__(self):
object.__setattr__(self, "action_highs", self.tensor_args.to_device(self.action_highs)) object.__setattr__(self, "action_highs", self.tensor_args.to_device(self.action_highs))
object.__setattr__(self, "action_lows", self.tensor_args.to_device(self.action_lows)) object.__setattr__(self, "action_lows", self.tensor_args.to_device(self.action_lows))
# check cuda graph version:
if self.use_cuda_graph: if self.use_cuda_graph:
self.use_cuda_graph = is_cuda_graph_available() self.use_cuda_graph = is_cuda_graph_available()
if self.num_particles is None: if self.num_particles is None:
self.num_particles = 1 self.num_particles = 1
if self.cold_start_n_iters is None:
self.cold_start_n_iters = self.n_iters
@staticmethod @staticmethod
def create_data_dict( def create_data_dict(
@@ -61,6 +108,17 @@ class OptimizerConfig:
tensor_args: TensorDeviceType = TensorDeviceType(), tensor_args: TensorDeviceType = TensorDeviceType(),
child_dict: Optional[Dict] = None, child_dict: Optional[Dict] = None,
): ):
"""Helper function to create dictionary from optimizer parameters and rollout class.
Args:
data_dict: optimizer parameters dictionary.
rollout_fn: rollout function.
tensor_args: tensor cuda device.
child_dict: new dictionary where parameters will be stored.
Returns:
Dictionary with parameters to create a :meth:`OptimizerConfig`
"""
if child_dict is None: if child_dict is None:
child_dict = deepcopy(data_dict) child_dict = deepcopy(data_dict)
child_dict["d_action"] = rollout_fn.d_action child_dict["d_action"] = rollout_fn.d_action
@@ -78,17 +136,33 @@ class OptimizerConfig:
class Optimizer(OptimizerConfig): class Optimizer(OptimizerConfig):
def __init__(self, config: Optional[OptimizerConfig] = None) -> None: def __init__(self, config: Optional[OptimizerConfig] = None) -> None:
"""Base optimization solver class
Args:
config: Initialized with parameters from a dataclass.
"""
if config is not None: if config is not None:
super().__init__(**vars(config)) super().__init__(**vars(config))
self.opt_dt = 0.0 self.opt_dt = 0.0
self.COLD_START = True self.COLD_START = True
self.update_nenvs(self.n_envs) self.update_nproblems(self.n_problems)
self._batch_goal = None self._batch_goal = None
self._rollout_list = None self._rollout_list = None
self.debug = [] self.debug = []
self.debug_cost = [] self.debug_cost = []
def optimize(self, opt_tensor: torch.Tensor, shift_steps=0, n_iters=None) -> torch.Tensor: def optimize(self, opt_tensor: torch.Tensor, shift_steps=0, n_iters=None) -> torch.Tensor:
"""Find a solution through optimization given the initial values for variables.
Args:
opt_tensor: Initial value of optimization variables.
Shape: [n_problems, action_horizon, d_action]
shift_steps: Shift variables along action_horizon. Useful in mpc warm-start setting.
n_iters: Override number of iterations to run optimization.
Returns:
Optimized values returned as a tensor of shape [n_problems, action_horizon, d_action].
"""
if self.COLD_START: if self.COLD_START:
n_iters = self.cold_start_n_iters n_iters = self.cold_start_n_iters
self.COLD_START = False self.COLD_START = False
@@ -99,17 +173,12 @@ class Optimizer(OptimizerConfig):
self.opt_dt = time.time() - st_time self.opt_dt = time.time() - st_time
return out return out
@abstractmethod
def _optimize(self, opt_tensor: torch.Tensor, shift_steps=0, n_iters=None) -> torch.Tensor:
pass
def _shift(self, shift_steps=0):
"""
Shift the variables in the solver to hotstart the next timestep
"""
return
def update_params(self, goal: Goal): def update_params(self, goal: Goal):
"""Update parameters in the :meth:`curobo.rollout.rollout_base.RolloutBase` instance.
Args:
goal: parameters to update rollout instance.
"""
with profiler.record_function("OptBase/batch_goal"): with profiler.record_function("OptBase/batch_goal"):
if self._batch_goal is not None: if self._batch_goal is not None:
self._batch_goal.copy_(goal, update_idx_buffers=True) # why True? self._batch_goal.copy_(goal, update_idx_buffers=True) # why True?
@@ -118,80 +187,37 @@ class Optimizer(OptimizerConfig):
self.rollout_fn.update_params(self._batch_goal) self.rollout_fn.update_params(self._batch_goal)
def reset(self): def reset(self):
""" """Reset optimizer."""
Reset the controller
"""
self.rollout_fn.reset() self.rollout_fn.reset()
self.debug = [] self.debug = []
self.debug_cost = [] self.debug_cost = []
# self.COLD_START = True
def update_nenvs(self, n_envs): def update_nproblems(self, n_problems: int):
assert n_envs > 0 """Update the number of problems that need to be optimized.
self._update_env_kernel(n_envs, self.num_particles)
self.n_envs = n_envs
def _update_env_kernel(self, n_envs, num_particles): Args:
log_info( n_problems: number of problems.
"Updating env kernel [n_envs: " """
+ str(n_envs) assert n_problems > 0
+ " , num_particles: " self._update_problem_kernel(n_problems, self.num_particles)
+ str(num_particles) self.n_problems = n_problems
+ " ]"
)
self.env_col = torch.arange( def get_nproblem_tensor(self, x):
0, n_envs, step=1, dtype=torch.long, device=self.tensor_args.device """This function takes an input tensor of shape (n_problem,....) and converts it into
) (n_particles,...).
self.n_select_ = torch.tensor(
[x * n_envs + x for x in range(n_envs)],
device=self.tensor_args.device,
dtype=torch.long,
)
# create sparse tensor:
sparse_indices = []
for i in range(n_envs):
sparse_indices.extend([[i * num_particles + x, i] for x in range(num_particles)])
sparse_values = torch.ones(len(sparse_indices))
sparse_indices = torch.tensor(sparse_indices)
if self.use_coo_sparse:
self.env_kernel_ = torch.sparse_coo_tensor(
sparse_indices.t(),
sparse_values,
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
)
else:
self.env_kernel_ = torch.sparse_coo_tensor(
sparse_indices.t(),
sparse_values,
device="cpu",
dtype=self.tensor_args.dtype,
)
self.env_kernel_ = self.env_kernel_.to_dense().to(device=self.tensor_args.device)
self._env_seeds = self.num_particles
def get_nenv_tensor(self, x):
"""This function takes an input tensor of shape (n_env,....) and converts it into
(n_particles,...)
""" """
# if x.shape[0] != self.n_envs:
# x_env = x.unsqueeze(0).repeat(self.n_envs, 1)
# else:
# x_env = x
# create a tensor # create a tensor
nx_env = self.env_kernel_ @ x nx_problem = self.problem_kernel_ @ x
return nx_env return nx_problem
def reset_seed(self): def reset_seed(self):
"""Reset seeds."""
return True return True
def reset_cuda_graph(self): def reset_cuda_graph(self):
"""Reset CUDA Graphs. This does not work, workaround is to create a new instance."""
if self.use_cuda_graph: if self.use_cuda_graph:
self.cu_opt_init = False self.cu_opt_init = False
else: else:
@@ -199,7 +225,87 @@ class Optimizer(OptimizerConfig):
self._batch_goal = None self._batch_goal = None
self.rollout_fn.reset_cuda_graph() self.rollout_fn.reset_cuda_graph()
def reset_shape(self):
"""Reset any flags in rollout class. Useful to reinitialize tensors for a new shape."""
self.rollout_fn.reset_shape()
self._batch_goal = None
def get_all_rollout_instances(self) -> List[RolloutBase]: def get_all_rollout_instances(self) -> List[RolloutBase]:
"""Get all instances of Rollout class in the optimizer."""
if self._rollout_list is None: if self._rollout_list is None:
self._rollout_list = [self.rollout_fn] self._rollout_list = [self.rollout_fn]
return self._rollout_list return self._rollout_list
@abstractmethod
def _optimize(self, opt_tensor: torch.Tensor, shift_steps=0, n_iters=None) -> torch.Tensor:
"""Implement this function in a derived class containing the solver.
Args:
opt_tensor: Initial value of optimization variables.
Shape: [n_problems, action_horizon, d_action]
shift_steps: Shift variables along action_horizon. Useful in mpc warm-start setting.
n_iters: Override number of iterations to run optimization.
Returns:
Optimized variables in tensor shape [action_horizon, d_action].
"""
return opt_tensor
@abstractmethod
def _shift(self, shift_steps=0):
"""Shift the variables in the solver to hotstart the next timestep.
Args:
shift_steps: Number of timesteps to shift.
"""
return
def _update_problem_kernel(self, n_problems: int, num_particles: int):
"""Update matrix used to map problem index to number of particles.
Args:
n_problems: Number of optimization problems.
num_particles: Number of particles per problem.
"""
log_info(
"Updating problem kernel [n_problems: "
+ str(n_problems)
+ " , num_particles: "
+ str(num_particles)
+ " ]"
)
self.problem_col = torch.arange(
0, n_problems, step=1, dtype=torch.long, device=self.tensor_args.device
)
self.n_select_ = torch.tensor(
[x * n_problems + x for x in range(n_problems)],
device=self.tensor_args.device,
dtype=torch.long,
)
# create sparse tensor:
sparse_indices = []
for i in range(n_problems):
sparse_indices.extend([[i * num_particles + x, i] for x in range(num_particles)])
sparse_values = torch.ones(len(sparse_indices))
sparse_indices = torch.tensor(sparse_indices)
if self.use_coo_sparse:
self.problem_kernel_ = torch.sparse_coo_tensor(
sparse_indices.t(),
sparse_values,
device=self.tensor_args.device,
dtype=self.tensor_args.dtype,
)
else:
self.problem_kernel_ = torch.sparse_coo_tensor(
sparse_indices.t(),
sparse_values,
device="cpu",
dtype=self.tensor_args.dtype,
)
self.problem_kernel_ = self.problem_kernel_.to_dense().to(
device=self.tensor_args.device
)
self._problem_seeds = self.num_particles

62
src/curobo/opt/particle/parallel_mppi.py
View File

@@ -65,7 +65,7 @@ class ParallelMPPIConfig(ParticleOptConfig):
alpha: float alpha: float
gamma: float gamma: float
kappa: float kappa: float
sample_per_env: bool sample_per_problem: bool
def __post_init__(self): def __post_init__(self):
self.init_cov = self.tensor_args.to_device(self.init_cov).unsqueeze(0) self.init_cov = self.tensor_args.to_device(self.init_cov).unsqueeze(0)
@@ -264,7 +264,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
actions = trajectories.actions actions = trajectories.actions
total_costs = self._compute_total_cost(costs) total_costs = self._compute_total_cost(costs)
# Let's reshape to n_envs now: # Let's reshape to n_problems now:
# first find the means before doing exponential utility: # first find the means before doing exponential utility:
w = self._exp_util(total_costs) w = self._exp_util(total_costs)
@@ -272,7 +272,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
# Update best action # Update best action
if self.sample_mode == SampleMode.BEST: if self.sample_mode == SampleMode.BEST:
best_idx = torch.argmax(w, dim=-1) best_idx = torch.argmax(w, dim=-1)
self.best_traj.copy_(actions[self.env_col, best_idx]) self.best_traj.copy_(actions[self.problem_col, best_idx])
if self.store_rollouts and self.visual_traj is not None: if self.store_rollouts and self.visual_traj is not None:
vis_seq = getattr(trajectories.state, self.visual_traj) vis_seq = getattr(trajectories.state, self.visual_traj)
@@ -281,7 +281,9 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
self.top_idx = top_idx self.top_idx = top_idx
top_trajs = torch.index_select(vis_seq, 0, top_idx[0]) top_trajs = torch.index_select(vis_seq, 0, top_idx[0])
for i in range(1, top_idx.shape[0]): for i in range(1, top_idx.shape[0]):
trajs = torch.index_select(vis_seq, 0, top_idx[i] + (self.particles_per_env * i)) trajs = torch.index_select(
vis_seq, 0, top_idx[i] + (self.particles_per_problem * i)
)
top_trajs = torch.cat((top_trajs, trajs), dim=0) top_trajs = torch.cat((top_trajs, trajs), dim=0)
if self.top_trajs is None or top_trajs.shape != self.top_trajs: if self.top_trajs is None or top_trajs.shape != self.top_trajs:
self.top_trajs = top_trajs self.top_trajs = top_trajs
@@ -317,13 +319,15 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
act_seq = self.mean_action.unsqueeze(-3) + scaled_delta act_seq = self.mean_action.unsqueeze(-3) + scaled_delta
cat_list = [act_seq] cat_list = [act_seq]
if self.neg_per_env > 0: if self.neg_per_problem > 0:
neg_action = -1.0 * self.mean_action neg_action = -1.0 * self.mean_action
neg_act_seqs = neg_action.unsqueeze(-3).expand(-1, self.neg_per_env, -1, -1) neg_act_seqs = neg_action.unsqueeze(-3).expand(-1, self.neg_per_problem, -1, -1)
cat_list.append(neg_act_seqs) cat_list.append(neg_act_seqs)
if self.null_per_env > 0: if self.null_per_problem > 0:
cat_list.append( cat_list.append(
self.null_act_seqs[: self.null_per_env].unsqueeze(0).expand(self.n_envs, -1, -1, -1) self.null_act_seqs[: self.null_per_problem]
.unsqueeze(0)
.expand(self.n_problems, -1, -1, -1)
) )
act_seq = torch.cat( act_seq = torch.cat(
@@ -343,8 +347,8 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
def update_init_mean(self, init_mean): def update_init_mean(self, init_mean):
# update mean: # update mean:
# init_mean = init_mean.clone() # init_mean = init_mean.clone()
if init_mean.shape[0] != self.n_envs: if init_mean.shape[0] != self.n_problems:
init_mean = init_mean.expand(self.n_envs, -1, -1) init_mean = init_mean.expand(self.n_problems, -1, -1)
if not copy_tensor(init_mean, self.mean_action): if not copy_tensor(init_mean, self.mean_action):
self.mean_action = init_mean.clone() self.mean_action = init_mean.clone()
if not copy_tensor(init_mean, self.best_traj): if not copy_tensor(init_mean, self.best_traj):
@@ -353,27 +357,27 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
def reset_mean(self): def reset_mean(self):
with profiler.record_function("mppi/reset_mean"): with profiler.record_function("mppi/reset_mean"):
if self.random_mean: if self.random_mean:
mean = self.mean_lib.get_samples([self.n_envs]) mean = self.mean_lib.get_samples([self.n_problems])
self.update_init_mean(mean) self.update_init_mean(mean)
else: else:
self.update_init_mean(self.init_mean) self.update_init_mean(self.init_mean)
def reset_covariance(self): def reset_covariance(self):
with profiler.record_function("mppi/reset_cov"): with profiler.record_function("mppi/reset_cov"):
# init_cov can either be a single value, or n_envs x 1 or n_envs x d_action # init_cov can either be a single value, or n_problems x 1 or n_problems x d_action
if self.cov_type == CovType.SIGMA_I: if self.cov_type == CovType.SIGMA_I:
# init_cov can either be a single value, or n_envs x 1 # init_cov can either be a single value, or n_problems x 1
self.cov_action = self.init_cov self.cov_action = self.init_cov
if self.init_cov.shape[0] != self.n_envs: if self.init_cov.shape[0] != self.n_problems:
self.cov_action = self.init_cov.unsqueeze(0).expand(self.n_envs, -1) self.cov_action = self.init_cov.unsqueeze(0).expand(self.n_problems, -1)
self.inv_cov_action = 1.0 / self.cov_action self.inv_cov_action = 1.0 / self.cov_action
a = torch.sqrt(self.cov_action) a = torch.sqrt(self.cov_action)
if not copy_tensor(a, self.scale_tril): if not copy_tensor(a, self.scale_tril):
self.scale_tril = a self.scale_tril = a
elif self.cov_type == CovType.DIAG_A: elif self.cov_type == CovType.DIAG_A:
# init_cov can either be a single value, or n_envs x 1 or n_envs x 7 # init_cov can either be a single value, or n_problems x 1 or n_problems x 7
init_cov = self.init_cov.clone() init_cov = self.init_cov.clone()
# if(init_cov.shape[-1] != self.d_action): # if(init_cov.shape[-1] != self.d_action):
@@ -382,8 +386,8 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
if len(init_cov.shape) == 2 and init_cov.shape[-1] != self.d_action: if len(init_cov.shape) == 2 and init_cov.shape[-1] != self.d_action:
init_cov = init_cov.expand(-1, self.d_action) init_cov = init_cov.expand(-1, self.d_action)
init_cov = init_cov.unsqueeze(1) init_cov = init_cov.unsqueeze(1)
if init_cov.shape[0] != self.n_envs: if init_cov.shape[0] != self.n_problems:
init_cov = init_cov.expand(self.n_envs, -1, -1) init_cov = init_cov.expand(self.n_problems, -1, -1)
if not copy_tensor(init_cov.clone(), self.cov_action): if not copy_tensor(init_cov.clone(), self.cov_action):
self.cov_action = init_cov.clone() self.cov_action = init_cov.clone()
self.inv_cov_action = 1.0 / self.cov_action self.inv_cov_action = 1.0 / self.cov_action
@@ -523,16 +527,18 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
n_iters = 1 n_iters = 1
else: else:
n_iters = self.n_iters n_iters = self.n_iters
if self.sample_per_env: if self.sample_per_problem:
s_set = ( s_set = (
self.sample_lib.get_samples( self.sample_lib.get_samples(
sample_shape=[self.sampled_particles_per_env * self.n_envs * n_iters], sample_shape=[
self.sampled_particles_per_problem * self.n_problems * n_iters
],
base_seed=self.seed, base_seed=self.seed,
) )
.view( .view(
n_iters, n_iters,
self.n_envs, self.n_problems,
self.sampled_particles_per_env, self.sampled_particles_per_problem,
self.action_horizon, self.action_horizon,
self.d_action, self.d_action,
) )
@@ -540,13 +546,17 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
) )
else: else:
s_set = self.sample_lib.get_samples( s_set = self.sample_lib.get_samples(
sample_shape=[n_iters * (self.sampled_particles_per_env)], sample_shape=[n_iters * (self.sampled_particles_per_problem)],
base_seed=self.seed, base_seed=self.seed,
) )
s_set = s_set.view( s_set = s_set.view(
n_iters, 1, self.sampled_particles_per_env, self.action_horizon, self.d_action n_iters,
1,
self.sampled_particles_per_problem,
self.action_horizon,
self.d_action,
) )
s_set = s_set.repeat(1, self.n_envs, 1, 1, 1).clone() s_set = s_set.repeat(1, self.n_problems, 1, 1, 1).clone()
s_set[:, :, -1, :, :] = 0.0 s_set[:, :, -1, :, :] = 0.0
if not copy_tensor(s_set, self._sample_set): if not copy_tensor(s_set, self._sample_set):
log_info("ParallelMPPI: Updating sample set") log_info("ParallelMPPI: Updating sample set")
@@ -575,7 +585,7 @@ class ParallelMPPI(ParticleOptBase, ParallelMPPIConfig):
Parameters Parameters
---------- ----------
state : dict or np.ndarray state : dict or np.ndarray
Initial state to set the simulation env to Initial state to set the simulation problem to
""" """
return super().generate_rollouts(init_act) return super().generate_rollouts(init_act)

35
src/curobo/opt/particle/particle_opt_base.py
View File

@@ -1,4 +1,3 @@
#!/usr/bin/env python
# #
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# #
@@ -95,7 +94,7 @@ class ParticleOptBase(Optimizer, ParticleOptConfig):
self.trajectories = None self.trajectories = None
self.cu_opt_init = False self.cu_opt_init = False
self.info = ParticleOptInfo() self.info = ParticleOptInfo()
self.update_num_particles_per_env(self.num_particles) self.update_num_particles_per_problem(self.num_particles)
@abstractmethod @abstractmethod
def _get_action_seq(self, mode=SampleMode): def _get_action_seq(self, mode=SampleMode):
@@ -172,7 +171,7 @@ class ParticleOptBase(Optimizer, ParticleOptConfig):
Parameters Parameters
---------- ----------
state : dict or np.ndarray state : dict or np.ndarray
Initial state to set the simulation env to Initial state to set the simulation problem to
""" """
act_seq = self.sample_actions(init_act) act_seq = self.sample_actions(init_act)
@@ -259,10 +258,10 @@ class ParticleOptBase(Optimizer, ParticleOptConfig):
# generate random simulated trajectories # generate random simulated trajectories
trajectory = self.generate_rollouts() trajectory = self.generate_rollouts()
trajectory.actions = trajectory.actions.view( trajectory.actions = trajectory.actions.view(
self.n_envs, self.particles_per_env, self.action_horizon, self.d_action self.n_problems, self.particles_per_problem, self.action_horizon, self.d_action
) )
trajectory.costs = trajectory.costs.view( trajectory.costs = trajectory.costs.view(
self.n_envs, self.particles_per_env, self.horizon self.n_problems, self.particles_per_problem, self.horizon
) )
with profiler.record_function("mpc/mppi/update_distribution"): with profiler.record_function("mpc/mppi/update_distribution"):
self._update_distribution(trajectory) self._update_distribution(trajectory)
@@ -275,26 +274,26 @@ class ParticleOptBase(Optimizer, ParticleOptConfig):
curr_action_seq = self._get_action_seq(mode=self.sample_mode) curr_action_seq = self._get_action_seq(mode=self.sample_mode)
return curr_action_seq return curr_action_seq
def update_nenvs(self, n_envs): def update_nproblems(self, n_problems):
assert n_envs > 0 assert n_problems > 0
self.total_num_particles = n_envs * self.num_particles self.total_num_particles = n_problems * self.num_particles
self.cu_opt_init = False self.cu_opt_init = False
super().update_nenvs(n_envs) super().update_nproblems(n_problems)
def update_num_particles_per_env(self, num_particles_per_env): def update_num_particles_per_problem(self, num_particles_per_problem):
self.null_per_env = round(int(self.null_act_frac * num_particles_per_env * 0.5)) self.null_per_problem = round(int(self.null_act_frac * num_particles_per_problem * 0.5))
self.neg_per_env = ( self.neg_per_problem = (
round(int(self.null_act_frac * num_particles_per_env)) - self.null_per_env round(int(self.null_act_frac * num_particles_per_problem)) - self.null_per_problem
) )
self.sampled_particles_per_env = ( self.sampled_particles_per_problem = (
num_particles_per_env - self.null_per_env - self.neg_per_env num_particles_per_problem - self.null_per_problem - self.neg_per_problem
) )
self.particles_per_env = num_particles_per_env self.particles_per_problem = num_particles_per_problem
if self.null_per_env > 0: if self.null_per_problem > 0:
self.null_act_seqs = torch.zeros( self.null_act_seqs = torch.zeros(
self.null_per_env, self.null_per_problem,
self.action_horizon, self.action_horizon,
self.d_action, self.d_action,
device=self.tensor_args.device, device=self.tensor_args.device,

3
src/curobo/opt/particle/particle_opt_utils.py
View File

@@ -121,7 +121,7 @@ def get_stomp_cov_jit(
else: else:
A[k * horizon + i, k * horizon + index] = fd_array[j + 3] A[k * horizon + i, k * horizon + index] = fd_array[j + 3]
R = torch.matmul(A.transpose(-2, -1), A) R = torch.matmul(A.transpose(-2, -1).clone(), A.clone())
M = torch.inverse(R) M = torch.inverse(R)
scaled_M = (1 / horizon) * M / (torch.max(torch.abs(M), dim=1)[0].unsqueeze(0)) scaled_M = (1 / horizon) * M / (torch.max(torch.abs(M), dim=1)[0].unsqueeze(0))
cov = M / torch.max(torch.abs(M)) cov = M / torch.max(torch.abs(M))
@@ -132,7 +132,6 @@ def get_stomp_cov_jit(
scale_tril = torch.linalg.cholesky(cov) scale_tril = torch.linalg.cholesky(cov)
else: else:
scale_tril = cov scale_tril = cov
""" """
k = 0 k = 0
act_cov_matrix = cov[k * horizon:k * horizon + horizon, k * horizon:k * horizon + horizon] act_cov_matrix = cov[k * horizon:k * horizon + horizon, k * horizon:k * horizon + horizon]

21
src/curobo/rollout/arm_base.py
View File

@@ -39,7 +39,7 @@ from curobo.rollout.rollout_base import Goal, RolloutBase, RolloutConfig, Rollou
from curobo.types.base import TensorDeviceType from curobo.types.base import TensorDeviceType
from curobo.types.robot import CSpaceConfig, RobotConfig from curobo.types.robot import CSpaceConfig, RobotConfig
from curobo.types.state import JointState from curobo.types.state import JointState
from curobo.util.logger import log_info, log_warn from curobo.util.logger import log_error, log_info, log_warn
from curobo.util.tensor_util import cat_sum from curobo.util.tensor_util import cat_sum
@@ -366,6 +366,7 @@ class ArmBase(RolloutBase, ArmBaseConfig):
) )
cost_list.append(coll_cost) cost_list.append(coll_cost)
if return_list: if return_list:
return cost_list return cost_list
cost = cat_sum(cost_list) cost = cat_sum(cost_list)
return cost return cost
@@ -424,6 +425,7 @@ class ArmBase(RolloutBase, ArmBaseConfig):
out_metrics = self.constraint_fn(state) out_metrics = self.constraint_fn(state)
out_metrics.state = state out_metrics.state = state
out_metrics = self.convergence_fn(state, out_metrics) out_metrics = self.convergence_fn(state, out_metrics)
out_metrics.cost = self.cost_fn(state)
return out_metrics return out_metrics
def get_metrics_cuda_graph(self, state: JointState): def get_metrics_cuda_graph(self, state: JointState):
@@ -451,6 +453,8 @@ class ArmBase(RolloutBase, ArmBaseConfig):
with torch.cuda.graph(self.cu_metrics_graph, stream=s): with torch.cuda.graph(self.cu_metrics_graph, stream=s):
self._cu_out_metrics = self.get_metrics(self._cu_metrics_state_in) self._cu_out_metrics = self.get_metrics(self._cu_metrics_state_in)
self._metrics_cuda_graph_init = True self._metrics_cuda_graph_init = True
if self._cu_metrics_state_in.position.shape != state.position.shape:
log_error("cuda graph changed")
self._cu_metrics_state_in.copy_(state) self._cu_metrics_state_in.copy_(state)
self.cu_metrics_graph.replay() self.cu_metrics_graph.replay()
out_metrics = self._cu_out_metrics out_metrics = self._cu_out_metrics
@@ -462,17 +466,6 @@ class ArmBase(RolloutBase, ArmBaseConfig):
): ):
if out_metrics is None: if out_metrics is None:
out_metrics = RolloutMetrics() out_metrics = RolloutMetrics()
if (
self.convergence_cfg.null_space_cfg is not None
and self.null_convergence.enabled
and self._goal_buffer.batch_retract_state_idx is not None
):
out_metrics.cost = self.null_convergence.forward_target_idx(
self._goal_buffer.retract_state,
state.state_seq.position,
self._goal_buffer.batch_retract_state_idx,
)
return out_metrics return out_metrics
def _get_augmented_state(self, state: JointState) -> KinematicModelState: def _get_augmented_state(self, state: JointState) -> KinematicModelState:
@@ -688,9 +681,11 @@ class ArmBase(RolloutBase, ArmBaseConfig):
act_seq = self.dynamics_model.init_action_mean.unsqueeze(0).repeat(self.batch_size, 1, 1) act_seq = self.dynamics_model.init_action_mean.unsqueeze(0).repeat(self.batch_size, 1, 1)
return act_seq return act_seq
def reset_cuda_graph(self): def reset_shape(self):
self._goal_idx_update = True self._goal_idx_update = True
super().reset_shape()
def reset_cuda_graph(self):
super().reset_cuda_graph() super().reset_cuda_graph()
def get_action_from_state(self, state: JointState): def get_action_from_state(self, state: JointState):

94
src/curobo/rollout/arm_reacher.py
View File

@@ -20,16 +20,16 @@ import torch.autograd.profiler as profiler
from curobo.geom.sdf.world import WorldCollision from curobo.geom.sdf.world import WorldCollision
from curobo.rollout.cost.cost_base import CostConfig from curobo.rollout.cost.cost_base import CostConfig
from curobo.rollout.cost.dist_cost import DistCost, DistCostConfig from curobo.rollout.cost.dist_cost import DistCost, DistCostConfig
from curobo.rollout.cost.pose_cost import PoseCost, PoseCostConfig from curobo.rollout.cost.pose_cost import PoseCost, PoseCostConfig, PoseCostMetric
from curobo.rollout.cost.straight_line_cost import StraightLineCost from curobo.rollout.cost.straight_line_cost import StraightLineCost
from curobo.rollout.cost.zero_cost import ZeroCost from curobo.rollout.cost.zero_cost import ZeroCost
from curobo.rollout.dynamics_model.kinematic_model import KinematicModelState from curobo.rollout.dynamics_model.kinematic_model import KinematicModelState
from curobo.rollout.rollout_base import Goal, RolloutMetrics from curobo.rollout.rollout_base import Goal, RolloutMetrics
from curobo.types.base import TensorDeviceType from curobo.types.base import TensorDeviceType
from curobo.types.robot import RobotConfig from curobo.types.robot import RobotConfig
from curobo.types.tensor import T_BValue_float from curobo.types.tensor import T_BValue_float, T_BValue_int
from curobo.util.helpers import list_idx_if_not_none from curobo.util.helpers import list_idx_if_not_none
from curobo.util.logger import log_info from curobo.util.logger import log_error, log_info
from curobo.util.tensor_util import cat_max, cat_sum from curobo.util.tensor_util import cat_max, cat_sum
# Local Folder # Local Folder
@@ -42,6 +42,8 @@ class ArmReacherMetrics(RolloutMetrics):
position_error: Optional[T_BValue_float] = None position_error: Optional[T_BValue_float] = None
rotation_error: Optional[T_BValue_float] = None rotation_error: Optional[T_BValue_float] = None
pose_error: Optional[T_BValue_float] = None pose_error: Optional[T_BValue_float] = None
goalset_index: Optional[T_BValue_int] = None
null_space_error: Optional[T_BValue_float] = None
def __getitem__(self, idx): def __getitem__(self, idx):
d_list = [ d_list = [
@@ -53,6 +55,8 @@ class ArmReacherMetrics(RolloutMetrics):
self.position_error, self.position_error,
self.rotation_error, self.rotation_error,
self.pose_error, self.pose_error,
self.goalset_index,
self.null_space_error,
] ]
idx_vals = list_idx_if_not_none(d_list, idx) idx_vals = list_idx_if_not_none(d_list, idx)
return ArmReacherMetrics(*idx_vals) return ArmReacherMetrics(*idx_vals)
@@ -65,10 +69,14 @@ class ArmReacherMetrics(RolloutMetrics):
constraint=None if self.constraint is None else self.constraint.clone(), constraint=None if self.constraint is None else self.constraint.clone(),
feasible=None if self.feasible is None else self.feasible.clone(), feasible=None if self.feasible is None else self.feasible.clone(),
state=None if self.state is None else self.state, state=None if self.state is None else self.state,
cspace_error=None if self.cspace_error is None else self.cspace_error, cspace_error=None if self.cspace_error is None else self.cspace_error.clone(),
position_error=None if self.position_error is None else self.position_error, position_error=None if self.position_error is None else self.position_error.clone(),
rotation_error=None if self.rotation_error is None else self.rotation_error, rotation_error=None if self.rotation_error is None else self.rotation_error.clone(),
pose_error=None if self.pose_error is None else self.pose_error, pose_error=None if self.pose_error is None else self.pose_error.clone(),
goalset_index=None if self.goalset_index is None else self.goalset_index.clone(),
null_space_error=(
None if self.null_space_error is None else self.null_space_error.clone()
),
) )
@@ -254,6 +262,7 @@ class ArmReacher(ArmBase, ArmReacherConfig):
goal_cost = self.goal_cost.forward( goal_cost = self.goal_cost.forward(
ee_pos_batch, ee_quat_batch, self._goal_buffer ee_pos_batch, ee_quat_batch, self._goal_buffer
) )
cost_list.append(goal_cost) cost_list.append(goal_cost)
with profiler.record_function("cost/link_poses"): with profiler.record_function("cost/link_poses"):
if self._goal_buffer.links_goal_pose is not None and self.cost_cfg.pose_cfg is not None: if self._goal_buffer.links_goal_pose is not None and self.cost_cfg.pose_cfg is not None:
@@ -296,36 +305,21 @@ class ArmReacher(ArmBase, ArmReacherConfig):
g_dist, g_dist,
) )
# cost += z_acc
cost_list.append(z_acc) cost_list.append(z_acc)
# print(self.cost_cfg.zero_jerk_cfg) if self.cost_cfg.zero_jerk_cfg is not None and self.zero_jerk_cost.enabled:
if (
self.cost_cfg.zero_jerk_cfg is not None
and self.zero_jerk_cost.enabled
# and g_dist is not None
):
# jerk = self.dynamics_model._aux_matrix @ state_batch.acceleration
z_jerk = self.zero_jerk_cost.forward( z_jerk = self.zero_jerk_cost.forward(
state_batch.jerk, state_batch.jerk,
g_dist, g_dist,
) )
cost_list.append(z_jerk) cost_list.append(z_jerk)
# cost += z_jerk
if ( if self.cost_cfg.zero_vel_cfg is not None and self.zero_vel_cost.enabled:
self.cost_cfg.zero_vel_cfg is not None
and self.zero_vel_cost.enabled
# and g_dist is not None
):
z_vel = self.zero_vel_cost.forward( z_vel = self.zero_vel_cost.forward(
state_batch.velocity, state_batch.velocity,
g_dist, g_dist,
) )
# cost += z_vel
# print(z_vel.shape)
cost_list.append(z_vel) cost_list.append(z_vel)
cost = cat_sum(cost_list) cost = cat_sum(cost_list)
# print(cost[:].T)
return cost return cost
def convergence_fn( def convergence_fn(
@@ -350,6 +344,7 @@ class ArmReacher(ArmBase, ArmReacherConfig):
) = self.pose_convergence.forward_out_distance( ) = self.pose_convergence.forward_out_distance(
state.ee_pos_seq, state.ee_quat_seq, self._goal_buffer state.ee_pos_seq, state.ee_quat_seq, self._goal_buffer
) )
out_metrics.goalset_index = self.pose_convergence.goalset_index_buffer # .clone()
if ( if (
self._goal_buffer.links_goal_pose is not None self._goal_buffer.links_goal_pose is not None
and self.convergence_cfg.pose_cfg is not None and self.convergence_cfg.pose_cfg is not None
@@ -389,6 +384,17 @@ class ArmReacher(ArmBase, ArmReacherConfig):
True, True,
) )
if (
self.convergence_cfg.null_space_cfg is not None
and self.null_convergence.enabled
and self._goal_buffer.batch_retract_state_idx is not None
):
out_metrics.null_space_error = self.null_convergence.forward_target_idx(
self._goal_buffer.retract_state,
state.state_seq.position,
self._goal_buffer.batch_retract_state_idx,
)
return out_metrics return out_metrics
def update_params( def update_params(
@@ -420,3 +426,43 @@ class ArmReacher(ArmBase, ArmReacherConfig):
else: else:
self.dist_cost.disable_cost() self.dist_cost.disable_cost()
self.cspace_convergence.disable_cost() self.cspace_convergence.disable_cost()
def get_pose_costs(self, include_link_pose: bool = False, include_convergence: bool = True):
pose_costs = [self.goal_cost]
if include_convergence:
pose_costs += [self.pose_convergence]
if include_link_pose:
log_error("Not implemented yet")
return pose_costs
def update_pose_cost_metric(
self,
metric: PoseCostMetric,
):
pose_costs = self.get_pose_costs()
if metric.hold_partial_pose:
if metric.hold_vec_weight is None:
log_error("hold_vec_weight is required")
[x.hold_partial_pose(metric.hold_vec_weight) for x in pose_costs]
if metric.release_partial_pose:
[x.release_partial_pose() for x in pose_costs]
if metric.reach_partial_pose:
if metric.reach_vec_weight is None:
log_error("reach_vec_weight is required")
[x.reach_partial_pose(metric.reach_vec_weight) for x in pose_costs]
if metric.reach_full_pose:
[x.reach_full_pose() for x in pose_costs]
pose_costs = self.get_pose_costs(include_convergence=False)
if metric.remove_offset_waypoint:
[x.remove_offset_waypoint() for x in pose_costs]
if metric.offset_position is not None or metric.offset_rotation is not None:
[
x.update_offset_waypoint(
offset_position=metric.offset_position,
offset_rotation=metric.offset_rotation,
offset_tstep_fraction=metric.offset_tstep_fraction,
)
for x in pose_costs
]

2
src/curobo/rollout/cost/bound_cost.py
View File

@@ -257,13 +257,13 @@ class BoundCost(CostBase, BoundCostConfig):
return cost return cost
def update_dt(self, dt: Union[float, torch.Tensor]): def update_dt(self, dt: Union[float, torch.Tensor]):
# return super().update_dt(dt)
if self.cost_type == BoundCostType.BOUNDS_SMOOTH: if self.cost_type == BoundCostType.BOUNDS_SMOOTH:
v_scale = dt / self._dt v_scale = dt / self._dt
a_scale = v_scale**2 a_scale = v_scale**2
j_scale = v_scale**3 j_scale = v_scale**3
self.smooth_weight[1] *= a_scale self.smooth_weight[1] *= a_scale
self.smooth_weight[2] *= j_scale self.smooth_weight[2] *= j_scale
return super().update_dt(dt) return super().update_dt(dt)

619
src/curobo/rollout/cost/pose_cost.py
View File

@@ -8,19 +8,23 @@
# without an express license agreement from NVIDIA CORPORATION or # without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited. # its affiliates is strictly prohibited.
# #
from __future__ import annotations
# Standard Library # Standard Library
import math
from dataclasses import dataclass from dataclasses import dataclass
from enum import Enum from enum import Enum
from typing import List, Optional from typing import List, Optional
# Third Party # Third Party
import torch import torch
from torch.autograd import Function
# CuRobo # CuRobo
from curobo.curobolib.geom import get_pose_distance, get_pose_distance_backward from curobo.curobolib.geom import PoseError, PoseErrorDistance
from curobo.rollout.rollout_base import Goal from curobo.rollout.rollout_base import Goal
from curobo.types.base import TensorDeviceType
from curobo.types.math import OrientationError, Pose from curobo.types.math import OrientationError, Pose
from curobo.util.logger import log_error
# Local Folder # Local Folder
from .cost_base import CostBase, CostConfig from .cost_base import CostBase, CostConfig
@@ -37,7 +41,11 @@ class PoseErrorType(Enum):
class PoseCostConfig(CostConfig): class PoseCostConfig(CostConfig):
cost_type: PoseErrorType = PoseErrorType.BATCH_GOAL cost_type: PoseErrorType = PoseErrorType.BATCH_GOAL
use_metric: bool = False use_metric: bool = False
project_distance: bool = True
run_vec_weight: Optional[List[float]] = None run_vec_weight: Optional[List[float]] = None
use_projected_distance: bool = True
offset_waypoint: List[float] = None
offset_tstep_fraction: float = -1.0
def __post_init__(self): def __post_init__(self):
if self.run_vec_weight is not None: if self.run_vec_weight is not None:
@@ -54,392 +62,85 @@ class PoseCostConfig(CostConfig):
self.vec_convergence = torch.zeros( self.vec_convergence = torch.zeros(
2, device=self.tensor_args.device, dtype=self.tensor_args.dtype 2, device=self.tensor_args.device, dtype=self.tensor_args.dtype
) )
if self.offset_waypoint is None:
self.offset_waypoint = [0, 0, 0, 0, 0, 0]
if self.run_weight is None:
self.run_weight = 1
self.offset_waypoint = self.tensor_args.to_device(self.offset_waypoint)
if isinstance(self.offset_tstep_fraction, float):
self.offset_tstep_fraction = self.tensor_args.to_device([self.offset_tstep_fraction])
return super().__post_init__() return super().__post_init__()
@torch.jit.script @dataclass
def backward_PoseError_jit(grad_g_dist, grad_out_distance, weight, g_vec): class PoseCostMetric:
grad_vec = grad_g_dist + (grad_out_distance * weight) hold_partial_pose: bool = False
grad = 1.0 * (grad_vec).unsqueeze(-1) * g_vec release_partial_pose: bool = False
return grad hold_vec_weight: Optional[torch.Tensor] = None
reach_partial_pose: bool = False
reach_full_pose: bool = False
reach_vec_weight: Optional[torch.Tensor] = None
offset_position: Optional[torch.Tensor] = None
offset_rotation: Optional[torch.Tensor] = None
offset_tstep_fraction: float = -1.0
remove_offset_waypoint: bool = False
def clone(self):
# full method: return PoseCostMetric(
@torch.jit.script hold_partial_pose=self.hold_partial_pose,
def backward_full_PoseError_jit( release_partial_pose=self.release_partial_pose,
grad_out_distance, grad_g_dist, grad_r_err, p_w, q_w, g_vec_p, g_vec_q hold_vec_weight=None if self.hold_vec_weight is None else self.hold_vec_weight.clone(),
): reach_partial_pose=self.reach_partial_pose,
p_grad = (grad_g_dist + (grad_out_distance * p_w)).unsqueeze(-1) * g_vec_p reach_full_pose=self.reach_full_pose,
q_grad = (grad_r_err + (grad_out_distance * q_w)).unsqueeze(-1) * g_vec_q reach_vec_weight=(
# p_grad = ((grad_out_distance * p_w)).unsqueeze(-1) * g_vec_p None if self.reach_vec_weight is None else self.reach_vec_weight.clone()
# q_grad = ((grad_out_distance * q_w)).unsqueeze(-1) * g_vec_q ),
offset_position=None if self.offset_position is None else self.offset_position.clone(),
return p_grad, q_grad offset_rotation=None if self.offset_rotation is None else self.offset_rotation.clone(),
offset_tstep_fraction=self.offset_tstep_fraction,
remove_offset_waypoint=self.remove_offset_waypoint,
class PoseErrorDistance(Function):
@staticmethod
def forward(
ctx,
current_position,
goal_position,
current_quat,
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
batch_pose_idx,
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
out_p_grad,
out_q_grad,
batch_size,
horizon,
mode=PoseErrorType.BATCH_GOAL.value,
num_goals=1,
use_metric=False,
):
# out_distance = current_position[..., 0].detach().clone() * 0.0
# out_position_distance = out_distance.detach().clone()
# out_rotation_distance = out_distance.detach().clone()
# out_vec = (
# torch.cat((current_position.detach().clone(), current_quat.detach().clone()), dim=-1)
# * 0.0
# )
# out_idx = out_distance.clone().to(dtype=torch.long)
(
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
) = get_pose_distance(
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
current_position.contiguous(),
goal_position,
current_quat.contiguous(),
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
batch_pose_idx,
batch_size,
horizon,
mode,
num_goals,
current_position.requires_grad,
True,
use_metric,
)
ctx.save_for_backward(out_p_vec, out_r_vec, weight, out_p_grad, out_q_grad)
return out_distance, out_position_distance, out_rotation_distance, out_idx # .view(-1,1)
@staticmethod
def backward(ctx, grad_out_distance, grad_g_dist, grad_r_err, grad_out_idx):
(g_vec_p, g_vec_q, weight, out_grad_p, out_grad_q) = ctx.saved_tensors
pos_grad = None
quat_grad = None
batch_size = g_vec_p.shape[0] * g_vec_p.shape[1]
if ctx.needs_input_grad[0] and ctx.needs_input_grad[2]:
pos_grad, quat_grad = get_pose_distance_backward(
out_grad_p,
out_grad_q,
grad_out_distance.contiguous(),
grad_g_dist.contiguous(),
grad_r_err.contiguous(),
weight,
g_vec_p,
g_vec_q,
batch_size,
use_distance=True,
)
# pos_grad, quat_grad = backward_full_PoseError_jit(
# grad_out_distance,
# grad_g_dist, grad_r_err, p_w, q_w, g_vec_p, g_vec_q
# )
elif ctx.needs_input_grad[0]:
pos_grad = backward_PoseError_jit(grad_g_dist, grad_out_distance, p_w, g_vec_p)
# grad_vec = grad_g_dist + (grad_out_distance * weight[1])
# pos_grad = 1.0 * (grad_vec).unsqueeze(-1) * g_vec[..., 4:]
elif ctx.needs_input_grad[2]:
quat_grad = backward_PoseError_jit(grad_r_err, grad_out_distance, q_w, g_vec_q)
# grad_vec = grad_r_err + (grad_out_distance * weight[0])
# quat_grad = 1.0 * (grad_vec).unsqueeze(-1) * g_vec[..., :4]
return (
pos_grad,
None,
quat_grad,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
) )
@classmethod
def create_grasp_approach_metric(
cls,
offset_position: float = 0.5,
linear_axis: int = 2,
tstep_fraction: float = 0.6,
tensor_args: TensorDeviceType = TensorDeviceType(),
) -> PoseCostMetric:
"""Enables moving to a pregrasp and then locked orientation movement to final grasp.
class PoseLoss(Function): Since this is added as a cost, the trajectory will not reach the exact offset, instead it
@staticmethod will try to take a blended path to the final grasp without stopping at the offset.
def forward(
ctx,
current_position,
goal_position,
current_quat,
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
batch_pose_idx,
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
out_p_grad,
out_q_grad,
batch_size,
horizon,
mode=PoseErrorType.BATCH_GOAL.value,
num_goals=1,
use_metric=False,
):
# out_distance = current_position[..., 0].detach().clone() * 0.0
# out_position_distance = out_distance.detach().clone()
# out_rotation_distance = out_distance.detach().clone()
# out_vec = (
# torch.cat((current_position.detach().clone(), current_quat.detach().clone()), dim=-1)
# * 0.0
# )
# out_idx = out_distance.clone().to(dtype=torch.long)
( Args:
out_distance, offset_position: offset in meters.
out_position_distance, linear_axis: specifies the x y or z axis.
out_rotation_distance, tstep_fraction: specifies the timestep fraction to start activating this constraint.
out_p_vec, tensor_args: cuda device.
out_r_vec,
out_idx,
) = get_pose_distance(
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
current_position.contiguous(),
goal_position,
current_quat.contiguous(),
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
batch_pose_idx,
batch_size,
horizon,
mode,
num_goals,
current_position.requires_grad,
False,
use_metric,
)
ctx.save_for_backward(out_p_vec, out_r_vec)
return out_distance
@staticmethod Returns:
def backward(ctx, grad_out_distance): # , grad_g_dist, grad_r_err, grad_out_idx): cost metric.
pos_grad = None """
quat_grad = None hold_vec_weight = tensor_args.to_device([1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
if ctx.needs_input_grad[0] and ctx.needs_input_grad[2]: hold_vec_weight[3 + linear_axis] = 0.0
(g_vec_p, g_vec_q) = ctx.saved_tensors offset_position_vec = tensor_args.to_device([0.0, 0.0, 0.0])
pos_grad = g_vec_p * grad_out_distance.unsqueeze(1) offset_position_vec[linear_axis] = offset_position
quat_grad = g_vec_q * grad_out_distance.unsqueeze(1) return cls(
pos_grad = pos_grad.unsqueeze(-2) hold_partial_pose=True,
quat_grad = quat_grad.unsqueeze(-2) hold_vec_weight=hold_vec_weight,
elif ctx.needs_input_grad[0]: offset_position=offset_position_vec,
(g_vec_p, g_vec_q) = ctx.saved_tensors offset_tstep_fraction=tstep_fraction,
pos_grad = g_vec_p * grad_out_distance.unsqueeze(1)
pos_grad = pos_grad.unsqueeze(-2)
elif ctx.needs_input_grad[2]:
(g_vec_p, g_vec_q) = ctx.saved_tensors
quat_grad = g_vec_q * grad_out_distance.unsqueeze(1)
quat_grad = quat_grad.unsqueeze(-2)
return (
pos_grad,
None,
quat_grad,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
) )
@classmethod
class PoseError(Function): def reset_metric(cls) -> PoseCostMetric:
@staticmethod return PoseCostMetric(
def forward( remove_offset_waypoint=True,
ctx, reach_full_pose=True,
current_position, release_partial_pose=True,
goal_position,
current_quat,
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
batch_pose_idx,
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
out_p_grad,
out_q_grad,
batch_size,
horizon,
mode=PoseErrorType.BATCH_GOAL.value,
num_goals=1,
use_metric=False,
):
# out_distance = current_position[..., 0].detach().clone() * 0.0
# out_position_distance = out_distance.detach().clone()
# out_rotation_distance = out_distance.detach().clone()
# out_vec = (
# torch.cat((current_position.detach().clone(), current_quat.detach().clone()), dim=-1)
# * 0.0
# )
# out_idx = out_distance.clone().to(dtype=torch.long)
(
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
) = get_pose_distance(
out_distance,
out_position_distance,
out_rotation_distance,
out_p_vec,
out_r_vec,
out_idx,
current_position.contiguous(),
goal_position,
current_quat.contiguous(),
goal_quat,
vec_weight,
weight,
vec_convergence,
run_weight,
run_vec_weight,
batch_pose_idx,
batch_size,
horizon,
mode,
num_goals,
current_position.requires_grad,
False,
use_metric,
)
ctx.save_for_backward(out_p_vec, out_r_vec)
return out_distance
@staticmethod
def backward(ctx, grad_out_distance): # , grad_g_dist, grad_r_err, grad_out_idx):
pos_grad = None
quat_grad = None
if ctx.needs_input_grad[0] and ctx.needs_input_grad[2]:
(g_vec_p, g_vec_q) = ctx.saved_tensors
pos_grad = g_vec_p
quat_grad = g_vec_q
elif ctx.needs_input_grad[0]:
(g_vec_p, g_vec_q) = ctx.saved_tensors
pos_grad = g_vec_p
elif ctx.needs_input_grad[2]:
(g_vec_p, g_vec_q) = ctx.saved_tensors
quat_grad = g_vec_q
return (
pos_grad,
None,
quat_grad,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
) )
@@ -449,13 +150,88 @@ class PoseCost(CostBase, PoseCostConfig):
CostBase.__init__(self) CostBase.__init__(self)
self.rot_weight = self.vec_weight[0:3] self.rot_weight = self.vec_weight[0:3]
self.pos_weight = self.vec_weight[3:6] self.pos_weight = self.vec_weight[3:6]
self._vec_convergence = self.tensor_args.to_device(self.vec_convergence) self._vec_convergence = self.tensor_args.to_device(self.vec_convergence)
self._batch_size = 0 self._batch_size = 0
self._horizon = 0 self._horizon = 0
def update_metric(self, metric: PoseCostMetric):
if metric.hold_partial_pose:
if metric.hold_vec_weight is None:
log_error("hold_vec_weight is required")
self.hold_partial_pose(metric.hold_vec_weight)
if metric.release_partial_pose:
self.release_partial_pose()
if metric.reach_partial_pose:
if metric.reach_vec_weight is None:
log_error("reach_vec_weight is required")
self.reach_partial_pose(metric.reach_vec_weight)
if metric.reach_full_pose:
self.reach_full_pose()
if metric.remove_offset_waypoint:
self.remove_offset_waypoint()
if metric.offset_position is not None or metric.offset_rotation is not None:
self.update_offset_waypoint(
offset_position=self.offset_position,
offset_rotation=self.offset_rotation,
offset_tstep_fraction=self.offset_tstep_fraction,
)
def hold_partial_pose(self, run_vec_weight: torch.Tensor):
self.run_vec_weight.copy_(run_vec_weight)
def release_partial_pose(self):
self.run_vec_weight[:] = 0.0
def reach_partial_pose(self, vec_weight: torch.Tensor):
self.vec_weight[:] = vec_weight
def reach_full_pose(self):
self.vec_weight[:] = 1.0
def update_offset_waypoint(
self,
offset_position: Optional[torch.Tensor] = None,
offset_rotation: Optional[torch.Tensor] = None,
offset_tstep_fraction: float = 0.75,
):
if offset_position is not None:
self.offset_waypoint[3:].copy_(offset_position)
if offset_rotation is not None:
self.offset_waypoint[:3].copy_(offset_rotation)
self.offset_tstep_fraction[:] = offset_tstep_fraction
if self._horizon <= 0:
print(self.weight)
log_error(
"Updating offset waypoint is only possible after initializing motion gen"
+ " run motion_gen.warmup() before adding offset_waypoint"
)
self.update_run_weight(run_tstep_fraction=offset_tstep_fraction)
def remove_offset_waypoint(self):
self.offset_tstep_fraction[:] = -1.0
self.update_run_weight()
def update_run_weight(
self, run_tstep_fraction: float = 0.0, run_weight: Optional[float] = None
):
if self._horizon == 1:
return
if run_weight is None:
run_weight = self.run_weight
active_steps = math.floor(self._horizon * run_tstep_fraction)
self._run_weight_vec[:, :active_steps] = 0
self._run_weight_vec[:, active_steps:-1] = run_weight
def update_batch_size(self, batch_size, horizon): def update_batch_size(self, batch_size, horizon):
if batch_size != self._batch_size or horizon != self._horizon: if batch_size != self._batch_size or horizon != self._horizon:
# print(self.weight)
# print(batch_size, horizon, self._batch_size, self._horizon)
# batch_size = b*h # batch_size = b*h
self.out_distance = torch.zeros( self.out_distance = torch.zeros(
(batch_size, horizon), device=self.tensor_args.device, dtype=self.tensor_args.dtype (batch_size, horizon), device=self.tensor_args.device, dtype=self.tensor_args.dtype
@@ -493,12 +269,16 @@ class PoseCost(CostBase, PoseCostConfig):
self._run_weight_vec = torch.ones( self._run_weight_vec = torch.ones(
(1, horizon), device=self.tensor_args.device, dtype=self.tensor_args.dtype (1, horizon), device=self.tensor_args.device, dtype=self.tensor_args.dtype
) )
if self.terminal and self.run_weight is not None: if self.terminal and self.run_weight is not None and horizon > 1:
self._run_weight_vec[:, :-1] *= self.run_weight self._run_weight_vec[:, :-1] = self.run_weight
self._batch_size = batch_size self._batch_size = batch_size
self._horizon = horizon self._horizon = horizon
@property
def goalset_index_buffer(self):
return self.out_idx
def _forward_goal_distribution(self, ee_pos_batch, ee_rot_batch, ee_goal_pos, ee_goal_rot): def _forward_goal_distribution(self, ee_pos_batch, ee_rot_batch, ee_goal_pos, ee_goal_rot):
ee_goal_pos = ee_goal_pos.unsqueeze(1) ee_goal_pos = ee_goal_pos.unsqueeze(1)
ee_goal_pos = ee_goal_pos.unsqueeze(1) ee_goal_pos = ee_goal_pos.unsqueeze(1)
@@ -563,13 +343,13 @@ class PoseCost(CostBase, PoseCostConfig):
def _update_cost_type(self, ee_goal_pos, ee_pos_batch, num_goals): def _update_cost_type(self, ee_goal_pos, ee_pos_batch, num_goals):
d_g = len(ee_goal_pos.shape) d_g = len(ee_goal_pos.shape)
b_sze = ee_goal_pos.shape[0] b_sze = ee_goal_pos.shape[0]
if d_g == 2 and b_sze == 1: # 1, 3 if d_g == 2 and b_sze == 1: # [1, 3]
self.cost_type = PoseErrorType.SINGLE_GOAL self.cost_type = PoseErrorType.SINGLE_GOAL
elif d_g == 2 and b_sze == ee_pos_batch.shape[0]: # b, 3 elif d_g == 2 and b_sze > 1: # [b, 3]
self.cost_type = PoseErrorType.BATCH_GOAL self.cost_type = PoseErrorType.BATCH_GOAL
elif d_g == 3: elif d_g == 3 and b_sze == 1: # [1, goalset, 3]
self.cost_type = PoseErrorType.GOALSET self.cost_type = PoseErrorType.GOALSET
elif len(ee_goal_pos.shape) == 4 and b_sze == ee_pos_bath.shape[0]: elif d_g == 3 and b_sze > 1: # [b, goalset,3]
self.cost_type = PoseErrorType.BATCH_GOALSET self.cost_type = PoseErrorType.BATCH_GOALSET
def forward_out_distance( def forward_out_distance(
@@ -599,6 +379,8 @@ class PoseCost(CostBase, PoseCostConfig):
self._vec_convergence, self._vec_convergence,
self._run_weight_vec, self._run_weight_vec,
self.run_vec_weight, self.run_vec_weight,
self.offset_waypoint,
self.offset_tstep_fraction,
goal.batch_pose_idx, goal.batch_pose_idx,
self.out_distance, self.out_distance,
self.out_position_distance, self.out_position_distance,
@@ -613,7 +395,9 @@ class PoseCost(CostBase, PoseCostConfig):
self.cost_type.value, self.cost_type.value,
num_goals, num_goals,
self.use_metric, self.use_metric,
self.project_distance,
) )
# print(self.out_idx.shape, self.out_idx[:,-1])
# print(goal.batch_pose_idx.shape) # print(goal.batch_pose_idx.shape)
cost = distance # .view(b, h)#.clone() cost = distance # .view(b, h)#.clone()
r_err = r_err # .view(b, h) r_err = r_err # .view(b, h)
@@ -632,38 +416,12 @@ class PoseCost(CostBase, PoseCostConfig):
ee_goal_rot = goal_pose.quaternion ee_goal_rot = goal_pose.quaternion
num_goals = goal_pose.n_goalset num_goals = goal_pose.n_goalset
self._update_cost_type(ee_goal_pos, ee_pos_batch, num_goals) self._update_cost_type(ee_goal_pos, ee_pos_batch, num_goals)
# print(self.cost_type)
b, h, _ = ee_pos_batch.shape b, h, _ = ee_pos_batch.shape
self.update_batch_size(b, h) self.update_batch_size(b, h)
# return self.out_distance # return self.out_distance
# print(b,h, ee_goal_pos.shape) # print(b,h, ee_goal_pos.shape)
if self.return_loss:
distance = PoseLoss.apply(
ee_pos_batch,
ee_goal_pos,
ee_rot_batch, # .view(-1, 4).contiguous(),
ee_goal_rot,
self.vec_weight,
self.weight,
self._vec_convergence,
self._run_weight_vec,
self.run_vec_weight,
goal.batch_pose_idx,
self.out_distance,
self.out_position_distance,
self.out_rotation_distance,
self.out_p_vec,
self.out_q_vec,
self.out_idx,
self.out_p_grad,
self.out_q_grad,
b,
h,
self.cost_type.value,
num_goals,
self.use_metric,
)
else:
distance = PoseError.apply( distance = PoseError.apply(
ee_pos_batch, ee_pos_batch,
ee_goal_pos, ee_goal_pos,
@@ -674,6 +432,8 @@ class PoseCost(CostBase, PoseCostConfig):
self._vec_convergence, self._vec_convergence,
self._run_weight_vec, self._run_weight_vec,
self.run_vec_weight, self.run_vec_weight,
self.offset_waypoint,
self.offset_tstep_fraction,
goal.batch_pose_idx, goal.batch_pose_idx,
self.out_distance, self.out_distance,
self.out_position_distance, self.out_position_distance,
@@ -688,9 +448,14 @@ class PoseCost(CostBase, PoseCostConfig):
self.cost_type.value, self.cost_type.value,
num_goals, num_goals,
self.use_metric, self.use_metric,
self.project_distance,
self.return_loss,
) )
cost = distance cost = distance
# if link_name is None and cost.shape[0]==8:
# print(ee_pos_batch[...,-1].squeeze())
# print(cost.shape)
return cost return cost
def forward_pose( def forward_pose(
@@ -708,33 +473,7 @@ class PoseCost(CostBase, PoseCostConfig):
b = query_pose.position.shape[0] b = query_pose.position.shape[0]
h = query_pose.position.shape[1] h = query_pose.position.shape[1]
num_goals = 1 num_goals = 1
if self.return_loss:
distance = PoseLoss.apply(
query_pose.position.unsqueeze(1),
ee_goal_pos,
query_pose.quaternion.unsqueeze(1),
ee_goal_quat,
self.vec_weight,
self.weight,
self._vec_convergence,
self._run_weight_vec,
self.run_vec_weight,
batch_pose_idx,
self.out_distance,
self.out_position_distance,
self.out_rotation_distance,
self.out_p_vec,
self.out_q_vec,
self.out_idx,
self.out_p_grad,
self.out_q_grad,
b,
h,
self.cost_type.value,
num_goals,
self.use_metric,
)
else:
distance = PoseError.apply( distance = PoseError.apply(
query_pose.position.unsqueeze(1), query_pose.position.unsqueeze(1),
ee_goal_pos, ee_goal_pos,
@@ -745,6 +484,8 @@ class PoseCost(CostBase, PoseCostConfig):
self._vec_convergence, self._vec_convergence,
self._run_weight_vec, self._run_weight_vec,
self.run_vec_weight, self.run_vec_weight,
self.offset_waypoint,
self.offset_tstep_fraction,
batch_pose_idx, batch_pose_idx,
self.out_distance, self.out_distance,
self.out_position_distance, self.out_position_distance,
@@ -759,5 +500,7 @@ class PoseCost(CostBase, PoseCostConfig):
self.cost_type.value, self.cost_type.value,
num_goals, num_goals,
self.use_metric, self.use_metric,
self.project_distance,
self.return_loss,
) )
return distance return distance

11
src/curobo/rollout/cost/stop_cost.py
View File

@@ -68,9 +68,14 @@ class StopCost(CostBase, StopCostConfig):
self.max_vel = (sum_matrix @ delta_vel).unsqueeze(-1) self.max_vel = (sum_matrix @ delta_vel).unsqueeze(-1)
def forward(self, vels): def forward(self, vels):
vel_abs = torch.abs(vels) cost = velocity_cost(vels, self.weight, self.max_vel)
vel_abs = torch.nn.functional.relu(vel_abs - self.max_vel) return cost
cost = self.weight * (torch.sum(vel_abs**2, dim=-1))
@torch.jit.script
def velocity_cost(vels, weight, max_vel):
vel_abs = torch.abs(vels)
vel_abs = torch.nn.functional.relu(vel_abs - max_vel[: vels.shape[1]])
cost = weight * (torch.sum(vel_abs**2, dim=-1))
return cost return cost

10
src/curobo/rollout/rollout_base.py
View File

@@ -23,7 +23,7 @@ import torch
# CuRobo # CuRobo
from curobo.types.base import TensorDeviceType from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose from curobo.types.math import Pose
from curobo.types.robot import CSpaceConfig, State, JointState from curobo.types.robot import CSpaceConfig, State
from curobo.types.tensor import ( from curobo.types.tensor import (
T_BDOF, T_BDOF,
T_DOF, T_DOF,
@@ -33,6 +33,7 @@ from curobo.types.tensor import (
T_BValue_float, T_BValue_float,
) )
from curobo.util.helpers import list_idx_if_not_none from curobo.util.helpers import list_idx_if_not_none
from curobo.util.logger import log_info
from curobo.util.sample_lib import HaltonGenerator from curobo.util.sample_lib import HaltonGenerator
from curobo.util.tensor_util import copy_tensor from curobo.util.tensor_util import copy_tensor
@@ -235,6 +236,7 @@ class Goal(Sequence):
batch_retract_state_idx=self.batch_retract_state_idx, batch_retract_state_idx=self.batch_retract_state_idx,
batch_goal_state_idx=self.batch_goal_state_idx, batch_goal_state_idx=self.batch_goal_state_idx,
links_goal_pose=self.links_goal_pose, links_goal_pose=self.links_goal_pose,
n_goalset=self.n_goalset,
) )
def _tensor_repeat_seeds(self, tensor, num_seeds): def _tensor_repeat_seeds(self, tensor, num_seeds):
@@ -353,7 +355,7 @@ class Goal(Sequence):
def _copy_tensor(self, ref_buffer, buffer): def _copy_tensor(self, ref_buffer, buffer):
if buffer is not None: if buffer is not None:
if ref_buffer is not None: if ref_buffer is not None and buffer.shape == ref_buffer.shape:
if not copy_tensor(buffer, ref_buffer): if not copy_tensor(buffer, ref_buffer):
ref_buffer = buffer.clone() ref_buffer = buffer.clone()
else: else:
@@ -553,6 +555,10 @@ class RolloutBase:
self._rollout_constraint_cuda_graph_init = False self._rollout_constraint_cuda_graph_init = False
if self.cu_rollout_constraint_graph is not None: if self.cu_rollout_constraint_graph is not None:
self.cu_rollout_constraint_graph.reset() self.cu_rollout_constraint_graph.reset()
self.reset_shape()
def reset_shape(self):
pass
@abstractmethod @abstractmethod
def get_action_from_state(self, state: State): def get_action_from_state(self, state: State):

64
src/curobo/types/camera.py
View File

@@ -16,10 +16,15 @@ from typing import List, Optional
# Third Party # Third Party
import torch import torch
from torch.profiler import record_function
# CuRobo # CuRobo
from curobo.geom.cv import (
get_projection_rays,
project_depth_using_rays,
project_pointcloud_to_depth,
)
from curobo.types.math import Pose from curobo.types.math import Pose
from curobo.util.logger import log_error, log_warn
@dataclass @dataclass
@@ -30,25 +35,82 @@ class CameraObservation:
depth_image: Optional[torch.Tensor] = None depth_image: Optional[torch.Tensor] = None
image_segmentation: Optional[torch.Tensor] = None image_segmentation: Optional[torch.Tensor] = None
projection_matrix: Optional[torch.Tensor] = None projection_matrix: Optional[torch.Tensor] = None
projection_rays: Optional[torch.Tensor] = None
resolution: Optional[List[int]] = None resolution: Optional[List[int]] = None
pose: Optional[Pose] = None pose: Optional[Pose] = None
intrinsics: Optional[torch.Tensor] = None intrinsics: Optional[torch.Tensor] = None
timestamp: float = 0.0 timestamp: float = 0.0
def filter_depth(self, distance: float = 0.01):
self.depth_image = torch.where(self.depth_image < distance, 0, self.depth_image)
@property @property
def shape(self): def shape(self):
return self.rgb_image.shape return self.rgb_image.shape
@record_function("camera/copy_")
def copy_(self, new_data: CameraObservation): def copy_(self, new_data: CameraObservation):
if self.rgb_image is not None:
self.rgb_image.copy_(new_data.rgb_image) self.rgb_image.copy_(new_data.rgb_image)
if self.depth_image is not None:
self.depth_image.copy_(new_data.depth_image) self.depth_image.copy_(new_data.depth_image)
if self.image_segmentation is not None:
self.image_segmentation.copy_(new_data.image_segmentation) self.image_segmentation.copy_(new_data.image_segmentation)
if self.projection_matrix is not None:
self.projection_matrix.copy_(new_data.projection_matrix) self.projection_matrix.copy_(new_data.projection_matrix)
if self.projection_rays is not None:
self.projection_rays.copy_(new_data.projection_rays)
if self.pose is not None:
self.pose.copy_(new_data.pose)
self.resolution = new_data.resolution self.resolution = new_data.resolution
@record_function("camera/clone")
def clone(self):
return CameraObservation(
depth_image=self.depth_image.clone() if self.depth_image is not None else None,
rgb_image=self.rgb_image.clone() if self.rgb_image is not None else None,
intrinsics=self.intrinsics.clone() if self.intrinsics is not None else None,
resolution=self.resolution,
pose=self.pose.clone() if self.pose is not None else None,
)
def to(self, device: torch.device): def to(self, device: torch.device):
if self.rgb_image is not None: if self.rgb_image is not None:
self.rgb_image = self.rgb_image.to(device=device) self.rgb_image = self.rgb_image.to(device=device)
if self.depth_image is not None: if self.depth_image is not None:
self.depth_image = self.depth_image.to(device=device) self.depth_image = self.depth_image.to(device=device)
return self return self
def get_pointcloud(self):
if self.projection_rays is None:
self.update_projection_rays()
depth_image = self.depth_image
if len(self.depth_image.shape) == 2:
depth_image = self.depth_image.unsqueeze(0)
point_cloud = project_depth_using_rays(depth_image, self.projection_rays)
return point_cloud
def get_image_from_pointcloud(self, pointcloud, output_image: Optional[torch.Tensor] = None):
if output_image is None:
output_image = torch.zeros(
(self.depth_image.shape[0], self.depth_image.shape[1]),
dtype=pointcloud.dtype,
device=pointcloud.device,
)
depth_image = project_pointcloud_to_depth(pointcloud, output_image=output_image)
return depth_image
def update_projection_rays(self):
intrinsics = self.intrinsics
if len(intrinsics.shape) == 2:
intrinsics = intrinsics.unsqueeze(0)
project_rays = get_projection_rays(
self.depth_image.shape[-2], self.depth_image.shape[-1], intrinsics
)
if self.projection_rays is None:
self.projection_rays = project_rays
self.projection_rays.copy_(project_rays)

16
src/curobo/types/math.py
View File

@@ -19,6 +19,7 @@ import numpy as np
import torch import torch
import torch.autograd.profiler as profiler import torch.autograd.profiler as profiler
from torch.autograd import Function from torch.autograd import Function
from torch.profiler import record_function
# CuRobo # CuRobo
from curobo.geom.transform import ( from curobo.geom.transform import (
@@ -33,7 +34,7 @@ from curobo.geom.transform import (
from curobo.types.base import TensorDeviceType from curobo.types.base import TensorDeviceType
from curobo.util.helpers import list_idx_if_not_none from curobo.util.helpers import list_idx_if_not_none
from curobo.util.logger import log_error, log_info, log_warn from curobo.util.logger import log_error, log_info, log_warn
from curobo.util.tensor_util import copy_if_not_none, copy_tensor from curobo.util.tensor_util import clone_if_not_none, copy_tensor
# Local Folder # Local Folder
from .tensor import T_BPosition, T_BQuaternion, T_BRotation from .tensor import T_BPosition, T_BQuaternion, T_BRotation
@@ -256,10 +257,10 @@ class Pose(Sequence):
def clone(self): def clone(self):
return Pose( return Pose(
position=copy_if_not_none(self.position), position=clone_if_not_none(self.position),
quaternion=copy_if_not_none(self.quaternion), quaternion=clone_if_not_none(self.quaternion),
normalize_rotation=False, normalize_rotation=False,
# rotation=copy_if_not_none(self.rotation), # rotation=clone_if_not_none(self.rotation),
) )
def to(self, tensor_args: TensorDeviceType): def to(self, tensor_args: TensorDeviceType):
@@ -408,6 +409,7 @@ class Pose(Sequence):
gpt_out, gpt_out,
) )
@record_function("math/pose/transform_points")
def batch_transform_points( def batch_transform_points(
self, self,
points: torch.Tensor, points: torch.Tensor,
@@ -432,6 +434,12 @@ class Pose(Sequence):
def shape(self): def shape(self):
return self.position.shape return self.position.shape
def compute_offset_pose(self, offset: Pose) -> Pose:
return self.multiply(offset)
def compute_local_pose(self, world_pose: Pose) -> Pose:
return self.inverse().multiply(world_pose)
def quat_multiply(q1, q2, q_res): def quat_multiply(q1, q2, q_res):
a_w = q1[..., 0] a_w = q1[..., 0]

22
src/curobo/types/state.py
View File

@@ -25,7 +25,7 @@ from curobo.types.tensor import T_BDOF, T_DOF
from curobo.util.logger import log_error, log_info, log_warn from curobo.util.logger import log_error, log_info, log_warn
from curobo.util.tensor_util import ( from curobo.util.tensor_util import (
check_tensor_shapes, check_tensor_shapes,
copy_if_not_none, clone_if_not_none,
copy_tensor, copy_tensor,
fd_tensor, fd_tensor,
tensor_repeat_seeds, tensor_repeat_seeds,
@@ -121,12 +121,14 @@ class JointState(State):
def repeat_seeds(self, num_seeds: int): def repeat_seeds(self, num_seeds: int):
return JointState( return JointState(
position=tensor_repeat_seeds(self.position, num_seeds), position=tensor_repeat_seeds(self.position, num_seeds),
velocity=tensor_repeat_seeds(self.velocity, num_seeds) velocity=(
if self.velocity is not None tensor_repeat_seeds(self.velocity, num_seeds) if self.velocity is not None else None
else None, ),
acceleration=tensor_repeat_seeds(self.acceleration, num_seeds) acceleration=(
tensor_repeat_seeds(self.acceleration, num_seeds)
if self.acceleration is not None if self.acceleration is not None
else None, else None
),
joint_names=self.joint_names, joint_names=self.joint_names,
) )
@@ -153,10 +155,10 @@ class JointState(State):
if self.joint_names is not None: if self.joint_names is not None:
j_names = self.joint_names.copy() j_names = self.joint_names.copy()
return JointState( return JointState(
position=copy_if_not_none(self.position), position=clone_if_not_none(self.position),
velocity=copy_if_not_none(self.velocity), velocity=clone_if_not_none(self.velocity),
acceleration=copy_if_not_none(self.acceleration), acceleration=clone_if_not_none(self.acceleration),
jerk=copy_if_not_none(self.jerk), jerk=clone_if_not_none(self.jerk),
joint_names=j_names, joint_names=j_names,
tensor_args=self.tensor_args, tensor_args=self.tensor_args,
) )

64
src/curobo/util/metrics.py Normal file
View File

@@ -0,0 +1,64 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Standard Library
from dataclasses import dataclass
from typing import List, Optional
# Third Party
import numpy as np
try:
# Third Party
from robometrics.statistics import (
Statistic,
TrajectoryGroupMetrics,
TrajectoryMetrics,
percent_true,
)
except ImportError:
raise ImportError(
"Benchmarking library not found, pip install "
+ '"robometrics[evaluator] @ git+https://github.com/fishbotics/robometrics.git"'
)
@dataclass
class CuroboMetrics(TrajectoryMetrics):
time: float = np.inf
cspace_path_length: float = 0.0
perception_success: bool = False
perception_interpolated_success: bool = False
jerk: float = np.inf
@dataclass
class CuroboGroupMetrics(TrajectoryGroupMetrics):
time: float = np.inf
cspace_path_length: Optional[Statistic] = None
perception_success: float = 0.0
perception_interpolated_success: float = 0.0
jerk: float = np.inf
@classmethod
def from_list(cls, group: List[CuroboMetrics]):
unskipped = [m for m in group if not m.skip]
successes = [m for m in unskipped if m.success]
data = super().from_list(group)
data.time = Statistic.from_list([m.time for m in successes])
data.cspace_path_length = Statistic.from_list([m.cspace_path_length for m in successes])
data.perception_success = percent_true([m.perception_success for m in group])
data.perception_interpolated_success = percent_true(
[m.perception_interpolated_success for m in group]
)
data.jerk = Statistic.from_list([m.jerk for m in successes])
return data

23
src/curobo/util/sample_lib.py
View File

@@ -111,7 +111,7 @@ class HaltonSampleLib(BaseSampleLib):
super().__init__(sample_config) super().__init__(sample_config)
# create halton generator: # create halton generator:
self.halton_generator = HaltonGenerator( self.halton_generator = HaltonGenerator(
self.ndims, seed=self.seed, tensor_args=self.tensor_args self.d_action, seed=self.seed, tensor_args=self.tensor_args
) )
def get_samples(self, sample_shape, base_seed=None, filter_smooth=False, **kwargs): def get_samples(self, sample_shape, base_seed=None, filter_smooth=False, **kwargs):
@@ -122,8 +122,10 @@ class HaltonSampleLib(BaseSampleLib):
seed = self.seed if base_seed is None else base_seed seed = self.seed if base_seed is None else base_seed
self.sample_shape = sample_shape self.sample_shape = sample_shape
self.seed = seed self.seed = seed
self.samples = self.halton_generator.get_gaussian_samples(sample_shape[0]) self.samples = self.halton_generator.get_gaussian_samples(
self.samples = self.samples.view(self.samples.shape[0], self.horizon, self.d_action) sample_shape[0] * self.horizon
)
self.samples = self.samples.view(sample_shape[0], self.horizon, self.d_action)
if filter_smooth: if filter_smooth:
self.samples = self.filter_smooth(self.samples) self.samples = self.filter_smooth(self.samples)
@@ -301,7 +303,7 @@ class StompSampleLib(BaseSampleLib):
self.filter_coeffs = None self.filter_coeffs = None
self.halton_generator = HaltonGenerator( self.halton_generator = HaltonGenerator(
self.ndims, seed=self.seed, tensor_args=self.tensor_args self.d_action, seed=self.seed, tensor_args=self.tensor_args
) )
def get_samples(self, sample_shape, base_seed=None, **kwargs): def get_samples(self, sample_shape, base_seed=None, **kwargs):
@@ -314,7 +316,10 @@ class StompSampleLib(BaseSampleLib):
# self.seed = seed # self.seed = seed
# torch.manual_seed(self.seed) # torch.manual_seed(self.seed)
halton_samples = self.halton_generator.get_gaussian_samples(sample_shape[0]) halton_samples = self.halton_generator.get_gaussian_samples(
sample_shape[0] * self.horizon
).view(sample_shape[0], self.horizon * self.d_action)
halton_samples = ( halton_samples = (
self.stomp_scale_tril.unsqueeze(0) @ halton_samples.unsqueeze(-1) self.stomp_scale_tril.unsqueeze(0) @ halton_samples.unsqueeze(-1)
).squeeze(-1) ).squeeze(-1)
@@ -415,7 +420,7 @@ class HaltonGenerator:
up_bounds=[1], up_bounds=[1],
low_bounds=[0], low_bounds=[0],
seed=123, seed=123,
store_buffer: Optional[int] = 1000, store_buffer: Optional[int] = 2000,
): ):
self._seed = seed self._seed = seed
self.tensor_args = tensor_args self.tensor_args = tensor_args
@@ -499,7 +504,7 @@ class HaltonGenerator:
@profiler.record_function("halton_generator/gaussian_samples") @profiler.record_function("halton_generator/gaussian_samples")
def get_gaussian_samples(self, num_samples, variance=1.0): def get_gaussian_samples(self, num_samples, variance=1.0):
std_dev = np.sqrt(variance) std_dev = np.sqrt(variance)
uniform_samples = self.get_samples(num_samples) uniform_samples = self.get_samples(num_samples, False)
gaussian_halton_samples = gaussian_transform( gaussian_halton_samples = gaussian_transform(
uniform_samples, self.proj_mat, self.i_mat, std_dev uniform_samples, self.proj_mat, self.i_mat, std_dev
) )
@@ -508,7 +513,7 @@ class HaltonGenerator:
@torch.jit.script @torch.jit.script
def gaussian_transform( def gaussian_transform(
uniform_samples: torch.Tensor, proj_mat: torch.Tensor, i_mat: torch.Tensor, variance: float uniform_samples: torch.Tensor, proj_mat: torch.Tensor, i_mat: torch.Tensor, std_dev: float
): ):
"""Compute a guassian transform of uniform samples. """Compute a guassian transform of uniform samples.
@@ -525,7 +530,7 @@ def gaussian_transform(
# these values. # these values.
changed_samples = 1.99 * uniform_samples - 0.99 changed_samples = 1.99 * uniform_samples - 0.99
gaussian_halton_samples = proj_mat * torch.erfinv(changed_samples) gaussian_halton_samples = proj_mat * torch.erfinv(changed_samples)
i_mat = i_mat * variance i_mat = i_mat * std_dev
gaussian_halton_samples = torch.matmul(gaussian_halton_samples, i_mat) gaussian_halton_samples = torch.matmul(gaussian_halton_samples, i_mat)
return gaussian_halton_samples return gaussian_halton_samples

20
src/curobo/util/tensor_util.py
View File

@@ -31,11 +31,29 @@ def copy_tensor(new_tensor: torch.Tensor, mem_tensor: torch.Tensor):
return False return False
def copy_if_not_none(x): def copy_if_not_none(new_tensor, ref_tensor):
"""Clones x if it's not None. """Clones x if it's not None.
TODO: Rename this to clone_if_not_none TODO: Rename this to clone_if_not_none
Args:
x (torch.Tensor): _description_
Returns:
_type_: _description_
"""
if ref_tensor is not None and new_tensor is not None:
ref_tensor.copy_(new_tensor)
elif ref_tensor is None and new_tensor is not None:
ref_tensor = new_tensor
return ref_tensor
def clone_if_not_none(x):
"""Clones x if it's not None.
Args: Args:
x (torch.Tensor): _description_ x (torch.Tensor): _description_

7
src/curobo/util/torch_utils.py
View File

@@ -34,3 +34,10 @@ def is_cuda_graph_available():
log_warn("WARNING: Disabling CUDA Graph as pytorch < 1.10") log_warn("WARNING: Disabling CUDA Graph as pytorch < 1.10")
return False return False
return True return True
def is_torch_compile_available():
if version.parse(torch.__version__) < version.parse("2.0"):
log_warn("WARNING: Disabling compile as pytorch < 2.0")
return False
return True

10
src/curobo/util/usd_helper.py
View File

@@ -16,7 +16,6 @@ from typing import Dict, List, Optional, Union
# Third Party # Third Party
import numpy as np import numpy as np
import torch import torch
from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics, UsdShade
from tqdm import tqdm from tqdm import tqdm
# CuRobo # CuRobo
@@ -47,6 +46,15 @@ from curobo.util_file import (
) )
from curobo.wrap.reacher.motion_gen import MotionGenResult from curobo.wrap.reacher.motion_gen import MotionGenResult
try:
# Third Party
from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics, UsdShade
except ImportError:
raise ImportError(
"usd-core failed to import, install with pip install usd-core"
+ " NOTE: Do not install this if using with ISAAC SIM."
)
def set_prim_translate(prim, translation): def set_prim_translate(prim, translation):
UsdGeom.Xformable(prim).AddTranslateOp().Set(Gf.Vec3d(translation)) UsdGeom.Xformable(prim).AddTranslateOp().Set(Gf.Vec3d(translation))

2
src/curobo/util/warp.py
View File

@@ -24,5 +24,5 @@ def init_warp(quiet=True, tensor_args: TensorDeviceType = TensorDeviceType()):
# wp.config.enable_backward = True # wp.config.enable_backward = True
wp.init() wp.init()
wp.force_load(wp.device_from_torch(tensor_args.device)) # wp.force_load(wp.device_from_torch(tensor_args.device))
return True return True

204
src/curobo/wrap/model/robot_segmenter.py Normal file
View File

@@ -0,0 +1,204 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Standard Library
from typing import Dict, Optional, Tuple, Union
# Third Party
import torch
from torch.profiler import record_function
# CuRobo
from curobo.geom.cv import (
get_projection_rays,
project_depth_using_rays,
)
from curobo.geom.types import PointCloud
from curobo.types.base import TensorDeviceType
from curobo.types.camera import CameraObservation
from curobo.types.math import Pose
from curobo.types.robot import RobotConfig
from curobo.types.state import JointState
from curobo.util.logger import log_error
from curobo.util_file import get_robot_configs_path, join_path, load_yaml
from curobo.wrap.model.robot_world import RobotWorld, RobotWorldConfig
class RobotSegmenter:
def __init__(
self,
robot_world: RobotWorld,
distance_threshold: float = 0.05,
use_cuda_graph: bool = True,
):
self._robot_world = robot_world
self._projection_rays = None
self.ready = False
self._out_points_buffer = None
self._out_gp = None
self._out_gq = None
self._out_gpt = None
self._cu_graph = None
self._use_cuda_graph = use_cuda_graph
self.tensor_args = robot_world.tensor_args
self.distance_threshold = distance_threshold
@staticmethod
def from_robot_file(
robot_file: Union[str, Dict],
collision_sphere_buffer: Optional[float],
distance_threshold: float = 0.05,
use_cuda_graph: bool = True,
tensor_args: TensorDeviceType = TensorDeviceType(),
):
robot_dict = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
if collision_sphere_buffer is not None:
robot_dict["kinematics"]["collision_sphere_buffer"] = collision_sphere_buffer
robot_cfg = RobotConfig.from_dict(robot_dict, tensor_args=tensor_args)
config = RobotWorldConfig.load_from_config(
robot_cfg,
None,
collision_activation_distance=0.0,
tensor_args=tensor_args,
)
robot_world = RobotWorld(config)
return RobotSegmenter(
robot_world, distance_threshold=distance_threshold, use_cuda_graph=use_cuda_graph
)
def get_pointcloud_from_depth(self, camera_obs: CameraObservation):
if self._projection_rays is None:
self.update_camera_projection(camera_obs)
depth_image = camera_obs.depth_image
if len(depth_image.shape) == 2:
depth_image = depth_image.unsqueeze(0)
points = project_depth_using_rays(depth_image, self._projection_rays)
return points
def update_camera_projection(self, camera_obs: CameraObservation):
intrinsics = camera_obs.intrinsics
if len(intrinsics.shape) == 2:
intrinsics = intrinsics.unsqueeze(0)
project_rays = get_projection_rays(
camera_obs.depth_image.shape[-2], camera_obs.depth_image.shape[-1], intrinsics
)
if self._projection_rays is None:
self._projection_rays = project_rays
self._projection_rays.copy_(project_rays)
self.ready = True
@record_function("robot_segmenter/get_robot_mask")
def get_robot_mask(
self,
camera_obs: CameraObservation,
joint_state: JointState,
):
"""
Assumes 1 robot and batch of depth images, batch of poses
"""
if len(camera_obs.depth_image.shape) != 3:
log_error("Send depth image as (batch, height, width)")
active_js = self._robot_world.get_active_js(joint_state)
mask, filtered_image = self.get_robot_mask_from_active_js(camera_obs, active_js)
return mask, filtered_image
def get_robot_mask_from_active_js(
self, camera_obs: CameraObservation, active_joint_state: JointState
):
q = active_joint_state.position
mask, filtered_image = self._call_op(camera_obs, q)
return mask, filtered_image
def _create_cg_graph(self, cam_obs, q):
self._cu_cam_obs = cam_obs.clone()
self._cu_q = q.clone()
s = torch.cuda.Stream(device=self.tensor_args.device)
s.wait_stream(torch.cuda.current_stream(device=self.tensor_args.device))
with torch.cuda.stream(s):
for _ in range(3):
self._cu_out, self._cu_filtered_out = self._mask_op(self._cu_cam_obs, self._cu_q)
torch.cuda.current_stream(device=self.tensor_args.device).wait_stream(s)
self._cu_graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(self._cu_graph, stream=s):
self._cu_out, self._cu_filtered_out = self._mask_op(
self._cu_cam_obs,
self._cu_q,
)
def _call_op(self, cam_obs, q):
if self._use_cuda_graph:
if self._cu_graph is None:
self._create_cg_graph(cam_obs, q)
self._cu_cam_obs.copy_(cam_obs)
self._cu_q.copy_(q)
self._cu_graph.replay()
return self._cu_out.clone(), self._cu_filtered_out.clone()
return self._mask_op(cam_obs, q)
@record_function("robot_segmenter/_mask_op")
def _mask_op(self, camera_obs, q):
if len(q.shape) == 1:
q = q.unsqueeze(0)
points = self.get_pointcloud_from_depth(camera_obs)
camera_to_robot = camera_obs.pose
if self._out_points_buffer is None:
self._out_points_buffer = points.clone()
if self._out_gpt is None:
self._out_gpt = torch.zeros((points.shape[0], points.shape[1], 3), device=points.device)
if self._out_gp is None:
self._out_gp = torch.zeros((camera_to_robot.position.shape[0], 3), device=points.device)
if self._out_gq is None:
self._out_gq = torch.zeros(
(camera_to_robot.quaternion.shape[0], 4), device=points.device
)
points_in_robot_frame = camera_to_robot.batch_transform_points(
points,
out_buffer=self._out_points_buffer,
gp_out=self._out_gp,
gq_out=self._out_gq,
gpt_out=self._out_gpt,
)
out_points = points_in_robot_frame
dist = self._robot_world.get_point_robot_distance(out_points, q)
mask, filtered_image = mask_image(camera_obs.depth_image, dist, self.distance_threshold)
return mask, filtered_image
@torch.jit.script
def mask_image(
image: torch.Tensor, distance: torch.Tensor, distance_threshold: float
) -> Tuple[torch.Tensor, torch.Tensor]:
distance = distance.view(
image.shape[0],
image.shape[1],
image.shape[2],
)
mask = torch.logical_and((image > 0.0), (distance > -distance_threshold))
filtered_image = torch.where(mask, 0, image)
return mask, filtered_image

53
src/curobo/wrap/model/robot_world.py
View File

@@ -34,6 +34,7 @@ from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose from curobo.types.math import Pose
from curobo.types.robot import RobotConfig from curobo.types.robot import RobotConfig
from curobo.types.state import JointState from curobo.types.state import JointState
from curobo.util.logger import log_error
from curobo.util.sample_lib import HaltonGenerator from curobo.util.sample_lib import HaltonGenerator
from curobo.util.warp import init_warp from curobo.util.warp import init_warp
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
@@ -60,8 +61,8 @@ class RobotWorldConfig:
world_model: Union[None, str, Dict, WorldConfig, List[WorldConfig], List[str]] = None, world_model: Union[None, str, Dict, WorldConfig, List[WorldConfig], List[str]] = None,
tensor_args: TensorDeviceType = TensorDeviceType(), tensor_args: TensorDeviceType = TensorDeviceType(),
n_envs: int = 1, n_envs: int = 1,
n_meshes: int = 10, n_meshes: int = 50,
n_cuboids: int = 10, n_cuboids: int = 50,
collision_activation_distance: float = 0.2, collision_activation_distance: float = 0.2,
self_collision_activation_distance: float = 0.0, self_collision_activation_distance: float = 0.0,
max_collision_distance: float = 1.0, max_collision_distance: float = 1.0,
@@ -74,6 +75,8 @@ class RobotWorldConfig:
if isinstance(robot_config, str): if isinstance(robot_config, str):
robot_config = load_yaml(join_path(get_robot_configs_path(), robot_config))["robot_cfg"] robot_config = load_yaml(join_path(get_robot_configs_path(), robot_config))["robot_cfg"]
if isinstance(robot_config, Dict): if isinstance(robot_config, Dict):
if "robot_cfg" in robot_config:
robot_config = robot_config["robot_cfg"]
robot_config = RobotConfig.from_dict(robot_config, tensor_args) robot_config = RobotConfig.from_dict(robot_config, tensor_args)
kinematics = CudaRobotModel(robot_config.kinematics) kinematics = CudaRobotModel(robot_config.kinematics)
@@ -178,8 +181,11 @@ class RobotWorld(RobotWorldConfig):
def __init__(self, config: RobotWorldConfig) -> None: def __init__(self, config: RobotWorldConfig) -> None:
RobotWorldConfig.__init__(self, **vars(config)) RobotWorldConfig.__init__(self, **vars(config))
self._batch_pose_idx = None self._batch_pose_idx = None
self._camera_projection_rays = None
def get_kinematics(self, q: torch.Tensor) -> CudaRobotModelState: def get_kinematics(self, q: torch.Tensor) -> CudaRobotModelState:
if len(q.shape) == 1:
log_error("q should be of shape [b, dof]")
state = self.kinematics.get_state(q) state = self.kinematics.get_state(q)
return state return state
@@ -344,6 +350,37 @@ class RobotWorld(RobotWorldConfig):
distance = self.pose_cost.forward_pose(x_des, x_current, self._batch_pose_idx) distance = self.pose_cost.forward_pose(x_des, x_current, self._batch_pose_idx)
return distance return distance
def get_point_robot_distance(self, points: torch.Tensor, q: torch.Tensor):
"""Compute distance from the robot at q joint configuration to points (e.g., pointcloud)
Args:
points: [b,n,3]
q: [1, dof]
Returns:
distance: [b,1] Positive is in collision with robot
NOTE: This currently does not support batched robot but can be done easily.
"""
if len(q.shape) == 1:
log_error("q should be of shape [b, dof]")
kin_state = self.get_kinematics(q)
b, n = None, None
if len(points.shape) == 3:
b, n, _ = points.shape
points = points.view(b * n, 3)
pt_distance = point_robot_distance(kin_state.link_spheres_tensor, points)
if b is not None:
pt_distance = pt_distance.view(b, n)
return pt_distance
def get_active_js(self, full_js: JointState):
active_jnames = self.kinematics.joint_names
out_js = full_js.get_ordered_joint_state(active_jnames)
return out_js
@torch.jit.script @torch.jit.script
def sum_mask(d1, d2, d3): def sum_mask(d1, d2, d3):
@@ -357,3 +394,15 @@ def mask(d1, d2, d3):
d_total = d1 + d2 + d3 d_total = d1 + d2 + d3
d_mask = d_total == 0.0 d_mask = d_total == 0.0
return d_mask return d_mask
@torch.jit.script
def point_robot_distance(link_spheres_tensor, points):
robot_spheres = link_spheres_tensor.view(1, -1, 4).contiguous()
robot_radius = robot_spheres[:, :, 3]
points = points.unsqueeze(1)
sph_distance = (
torch.linalg.norm(points - robot_spheres[:, :, :3], dim=-1) - robot_radius
) # b, n_spheres
pt_distance = torch.max(-1 * sph_distance, dim=-1)[0]
return pt_distance

3
src/curobo/wrap/reacher/evaluator.py
View File

@@ -52,7 +52,8 @@ def smooth_cost(abs_acc, abs_jerk, opt_dt):
# jerk = torch.max(torch.max(abs_jerk, dim=-1)[0], dim=-1)[0] # jerk = torch.max(torch.max(abs_jerk, dim=-1)[0], dim=-1)[0]
jerk = torch.mean(torch.max(abs_jerk, dim=-1)[0], dim=-1) jerk = torch.mean(torch.max(abs_jerk, dim=-1)[0], dim=-1)
mean_acc = torch.mean(torch.max(abs_acc, dim=-1)[0], dim=-1) # [0] mean_acc = torch.mean(torch.max(abs_acc, dim=-1)[0], dim=-1) # [0]
a = (jerk * 0.001) + opt_dt + (mean_acc * 0.01) a = (jerk * 0.001) + 5.0 * opt_dt + (mean_acc * 0.01)
return a return a

61
src/curobo/wrap/reacher/ik_solver.py
View File

@@ -26,6 +26,7 @@ from curobo.opt.newton.newton_base import NewtonOptBase, NewtonOptConfig
from curobo.opt.particle.parallel_es import ParallelES, ParallelESConfig from curobo.opt.particle.parallel_es import ParallelES, ParallelESConfig
from curobo.opt.particle.parallel_mppi import ParallelMPPI, ParallelMPPIConfig from curobo.opt.particle.parallel_mppi import ParallelMPPI, ParallelMPPIConfig
from curobo.rollout.arm_reacher import ArmReacher, ArmReacherConfig from curobo.rollout.arm_reacher import ArmReacher, ArmReacherConfig
from curobo.rollout.cost.pose_cost import PoseCostMetric
from curobo.rollout.rollout_base import Goal, RolloutBase, RolloutMetrics from curobo.rollout.rollout_base import Goal, RolloutBase, RolloutMetrics
from curobo.types.base import TensorDeviceType from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose from curobo.types.math import Pose
@@ -56,6 +57,7 @@ class IKSolverConfig:
world_coll_checker: Optional[WorldCollision] = None world_coll_checker: Optional[WorldCollision] = None
sample_rejection_ratio: int = 50 sample_rejection_ratio: int = 50
tensor_args: TensorDeviceType = TensorDeviceType() tensor_args: TensorDeviceType = TensorDeviceType()
use_cuda_graph: bool = True
@staticmethod @staticmethod
@profiler.record_function("ik_solver/load_from_robot_config") @profiler.record_function("ik_solver/load_from_robot_config")
@@ -72,12 +74,12 @@ class IKSolverConfig:
base_cfg_file: str = "base_cfg.yml", base_cfg_file: str = "base_cfg.yml",
particle_file: str = "particle_ik.yml", particle_file: str = "particle_ik.yml",
gradient_file: str = "gradient_ik.yml", gradient_file: str = "gradient_ik.yml",
use_cuda_graph: Optional[bool] = None, use_cuda_graph: bool = True,
self_collision_check: bool = True, self_collision_check: bool = True,
self_collision_opt: bool = True, self_collision_opt: bool = True,
grad_iters: Optional[int] = None, grad_iters: Optional[int] = None,
use_particle_opt: bool = True, use_particle_opt: bool = True,
collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.PRIMITIVE, collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.MESH,
sync_cuda_time: Optional[bool] = None, sync_cuda_time: Optional[bool] = None,
use_gradient_descent: bool = False, use_gradient_descent: bool = False,
collision_cache: Optional[Dict[str, int]] = None, collision_cache: Optional[Dict[str, int]] = None,
@@ -90,6 +92,7 @@ class IKSolverConfig:
regularization: bool = True, regularization: bool = True,
collision_activation_distance: Optional[float] = None, collision_activation_distance: Optional[float] = None,
high_precision: bool = False, high_precision: bool = False,
project_pose_to_goal_frame: bool = True,
): ):
if position_threshold <= 0.001: if position_threshold <= 0.001:
high_precision = True high_precision = True
@@ -116,6 +119,9 @@ class IKSolverConfig:
base_config_data["convergence"]["cspace_cfg"]["weight"] = 0.0 base_config_data["convergence"]["cspace_cfg"]["weight"] = 0.0
config_data["cost"]["bound_cfg"]["null_space_weight"] = 0.0 config_data["cost"]["bound_cfg"]["null_space_weight"] = 0.0
grad_config_data["cost"]["bound_cfg"]["null_space_weight"] = 0.0 grad_config_data["cost"]["bound_cfg"]["null_space_weight"] = 0.0
if isinstance(robot_cfg, str):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_cfg))["robot_cfg"]
if ee_link_name is not None: if ee_link_name is not None:
if isinstance(robot_cfg, RobotConfig): if isinstance(robot_cfg, RobotConfig):
raise NotImplementedError("ee link cannot be changed after creating RobotConfig") raise NotImplementedError("ee link cannot be changed after creating RobotConfig")
@@ -123,8 +129,6 @@ class IKSolverConfig:
robot_cfg.kinematics.ee_link = ee_link_name robot_cfg.kinematics.ee_link = ee_link_name
else: else:
robot_cfg["kinematics"]["ee_link"] = ee_link_name robot_cfg["kinematics"]["ee_link"] = ee_link_name
if isinstance(robot_cfg, str):
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_cfg))["robot_cfg"]
if isinstance(robot_cfg, dict): if isinstance(robot_cfg, dict):
robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args) robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
@@ -160,8 +164,8 @@ class IKSolverConfig:
grad_config_data["cost"]["self_collision_cfg"]["weight"] = 0.0 grad_config_data["cost"]["self_collision_cfg"]["weight"] = 0.0
if grad_iters is not None: if grad_iters is not None:
grad_config_data["lbfgs"]["n_iters"] = grad_iters grad_config_data["lbfgs"]["n_iters"] = grad_iters
config_data["mppi"]["n_envs"] = 1 config_data["mppi"]["n_problems"] = 1
grad_config_data["lbfgs"]["n_envs"] = 1 grad_config_data["lbfgs"]["n_problems"] = 1
grad_cfg = ArmReacherConfig.from_dict( grad_cfg = ArmReacherConfig.from_dict(
robot_cfg, robot_cfg,
grad_config_data["model"], grad_config_data["model"],
@@ -241,6 +245,7 @@ class IKSolverConfig:
world_coll_checker=world_coll_checker, world_coll_checker=world_coll_checker,
rollout_fn=aux_rollout, rollout_fn=aux_rollout,
tensor_args=tensor_args, tensor_args=tensor_args,
use_cuda_graph=use_cuda_graph,
) )
return ik_cfg return ik_cfg
@@ -259,6 +264,7 @@ class IKResult(Sequence):
error: T_BValue_float error: T_BValue_float
solve_time: float solve_time: float
debug_info: Optional[Any] = None debug_info: Optional[Any] = None
goalset_index: Optional[torch.Tensor] = None
def __getitem__(self, idx): def __getitem__(self, idx):
success = self.success[idx] success = self.success[idx]
@@ -272,6 +278,7 @@ class IKResult(Sequence):
position_error=self.position_error[idx], position_error=self.position_error[idx],
rotation_error=self.rotation_error[idx], rotation_error=self.rotation_error[idx],
debug_info=self.debug_info, debug_info=self.debug_info,
goalset_index=None if self.goalset_index is None else self.goalset_index[idx],
) )
def __len__(self): def __len__(self):
@@ -317,7 +324,7 @@ class IKSolver(IKSolverConfig):
self.solver.rollout_fn.retract_state.unsqueeze(0) self.solver.rollout_fn.retract_state.unsqueeze(0)
) )
self.dof = self.solver.safety_rollout.d_action self.dof = self.solver.safety_rollout.d_action
self._col = torch.arange(0, 1, device=self.tensor_args.device, dtype=torch.long) self._col = None # torch.arange(0, 1, device=self.tensor_args.device, dtype=torch.long)
# self.fixed_seeds = self.solver.safety_rollout.sample_random_actions(100 * 200) # self.fixed_seeds = self.solver.safety_rollout.sample_random_actions(100 * 200)
# create random seeder: # create random seeder:
@@ -355,10 +362,14 @@ class IKSolver(IKSolverConfig):
self._goal_buffer, self._goal_buffer,
self.tensor_args, self.tensor_args,
) )
# print("Goal:", self._goal_buffer.links_goal_pose)
if update_reference: if update_reference:
self.solver.update_nenvs(self._solve_state.get_ik_batch_size()) self.reset_shape()
if self.use_cuda_graph and self._col is not None:
log_error("changing goal type, breaking previous cuda graph.")
self.reset_cuda_graph() self.reset_cuda_graph()
self.solver.update_nproblems(self._solve_state.get_ik_batch_size())
self._goal_buffer.current_state = self.init_state.repeat_seeds(goal_pose.batch) self._goal_buffer.current_state = self.init_state.repeat_seeds(goal_pose.batch)
self._col = torch.arange( self._col = torch.arange(
0, 0,
@@ -676,6 +687,8 @@ class IKSolver(IKSolverConfig):
if newton_iters is not None: if newton_iters is not None:
self.solver.newton_optimizer.outer_iters = self.og_newton_iters self.solver.newton_optimizer.outer_iters = self.og_newton_iters
ik_result = self.get_result(num_seeds, result, goal_buffer.goal_pose, return_seeds) ik_result = self.get_result(num_seeds, result, goal_buffer.goal_pose, return_seeds)
if ik_result.goalset_index is not None:
ik_result.goalset_index[ik_result.goalset_index >= goal_pose.n_goalset] = 0
return ik_result return ik_result
@@ -684,15 +697,18 @@ class IKSolver(IKSolverConfig):
self, num_seeds: int, result: WrapResult, goal_pose: Pose, return_seeds: int self, num_seeds: int, result: WrapResult, goal_pose: Pose, return_seeds: int
) -> IKResult: ) -> IKResult:
success = self.get_success(result.metrics, num_seeds=num_seeds) success = self.get_success(result.metrics, num_seeds=num_seeds)
if result.metrics.cost is not None: # if result.metrics.cost is not None:
result.metrics.pose_error += result.metrics.cost # result.metrics.pose_error += result.metrics.cost * 0.0001
if result.metrics.null_space_error is not None:
result.metrics.pose_error += result.metrics.null_space_error
if result.metrics.cspace_error is not None: if result.metrics.cspace_error is not None:
result.metrics.pose_error += result.metrics.cspace_error result.metrics.pose_error += result.metrics.cspace_error
q_sol, success, position_error, rotation_error, total_error = get_result( q_sol, success, position_error, rotation_error, total_error, goalset_index = get_result(
result.metrics.pose_error, result.metrics.pose_error,
result.metrics.position_error, result.metrics.position_error,
result.metrics.rotation_error, result.metrics.rotation_error,
result.metrics.goalset_index,
success, success,
result.action.position, result.action.position,
self._col, self._col,
@@ -717,6 +733,7 @@ class IKSolver(IKSolverConfig):
solve_time=result.solve_time, solve_time=result.solve_time,
error=total_error, error=total_error,
debug_info={"solver": result.debug}, debug_info={"solver": result.debug},
goalset_index=goalset_index,
) )
return ik_result return ik_result
@@ -959,6 +976,10 @@ class IKSolver(IKSolverConfig):
self.solver.reset_cuda_graph() self.solver.reset_cuda_graph()
self.rollout_fn.reset_cuda_graph() self.rollout_fn.reset_cuda_graph()
def reset_shape(self):
self.solver.reset_shape()
self.rollout_fn.reset_shape()
def attach_object_to_robot( def attach_object_to_robot(
self, self,
sphere_radius: float, sphere_radius: float,
@@ -977,6 +998,17 @@ class IKSolver(IKSolverConfig):
def get_retract_config(self): def get_retract_config(self):
return self.rollout_fn.dynamics_model.retract_config return self.rollout_fn.dynamics_model.retract_config
def update_pose_cost_metric(
self,
metric: PoseCostMetric,
):
rollouts = self.get_all_rollout_instances()
[
rollout.update_pose_cost_metric(metric)
for rollout in rollouts
if isinstance(rollout, ArmReacher)
]
@torch.jit.script @torch.jit.script
def get_success( def get_success(
@@ -1001,6 +1033,7 @@ def get_result(
pose_error, pose_error,
position_error, position_error,
rotation_error, rotation_error,
goalset_index: Union[torch.Tensor, None],
success, success,
sol_position, sol_position,
col, col,
@@ -1018,4 +1051,6 @@ def get_result(
position_error = position_error[idx].view(batch_size, return_seeds) position_error = position_error[idx].view(batch_size, return_seeds)
rotation_error = rotation_error[idx].view(batch_size, return_seeds) rotation_error = rotation_error[idx].view(batch_size, return_seeds)
total_error = position_error + rotation_error total_error = position_error + rotation_error
return q_sol, success, position_error, rotation_error, total_error if goalset_index is not None:
goalset_index = goalset_index[idx].view(batch_size, return_seeds)
return q_sol, success, position_error, rotation_error, total_error, goalset_index

613
src/curobo/wrap/reacher/motion_gen.py
View File

File diff suppressed because it is too large Load Diff

4
src/curobo/wrap/reacher/mpc.py
View File

@@ -80,7 +80,7 @@ class MpcSolverConfig:
task_file = "particle_mpc.yml" task_file = "particle_mpc.yml"
config_data = load_yaml(join_path(get_task_configs_path(), task_file)) config_data = load_yaml(join_path(get_task_configs_path(), task_file))
config_data["mppi"]["n_envs"] = 1 config_data["mppi"]["n_problems"] = 1
if step_dt is not None: if step_dt is not None:
config_data["model"]["dt_traj_params"]["base_dt"] = step_dt config_data["model"]["dt_traj_params"]["base_dt"] = step_dt
if particle_opt_iters is not None: if particle_opt_iters is not None:
@@ -238,7 +238,7 @@ class MpcSolver(MpcSolverConfig):
self.tensor_args, self.tensor_args,
) )
if update_reference: if update_reference:
self.solver.update_nenvs(self._solve_state.get_batch_size()) self.solver.update_nproblems(self._solve_state.get_batch_size())
self.reset() self.reset()
self.reset_cuda_graph() self.reset_cuda_graph()
self._col = torch.arange( self._col = torch.arange(

93
src/curobo/wrap/reacher/trajopt.py
View File

@@ -30,6 +30,7 @@ from curobo.opt.newton.newton_base import NewtonOptBase, NewtonOptConfig
from curobo.opt.particle.parallel_es import ParallelES, ParallelESConfig from curobo.opt.particle.parallel_es import ParallelES, ParallelESConfig
from curobo.opt.particle.parallel_mppi import ParallelMPPI, ParallelMPPIConfig from curobo.opt.particle.parallel_mppi import ParallelMPPI, ParallelMPPIConfig
from curobo.rollout.arm_reacher import ArmReacher, ArmReacherConfig from curobo.rollout.arm_reacher import ArmReacher, ArmReacherConfig
from curobo.rollout.cost.pose_cost import PoseCostMetric
from curobo.rollout.dynamics_model.integration_utils import ( from curobo.rollout.dynamics_model.integration_utils import (
action_interpolate_kernel, action_interpolate_kernel,
interpolate_kernel, interpolate_kernel,
@@ -39,7 +40,7 @@ from curobo.types.base import TensorDeviceType
from curobo.types.robot import JointState, RobotConfig from curobo.types.robot import JointState, RobotConfig
from curobo.types.tensor import T_BDOF, T_DOF, T_BValue_bool, T_BValue_float from curobo.types.tensor import T_BDOF, T_DOF, T_BValue_bool, T_BValue_float
from curobo.util.helpers import list_idx_if_not_none from curobo.util.helpers import list_idx_if_not_none
from curobo.util.logger import log_info, log_warn from curobo.util.logger import log_error, log_info, log_warn
from curobo.util.trajectory import ( from curobo.util.trajectory import (
InterpolateType, InterpolateType,
calculate_dt_no_clamp, calculate_dt_no_clamp,
@@ -78,6 +79,7 @@ class TrajOptSolverConfig:
trim_steps: Optional[List[int]] = None trim_steps: Optional[List[int]] = None
store_debug_in_result: bool = False store_debug_in_result: bool = False
optimize_dt: bool = True optimize_dt: bool = True
use_cuda_graph: bool = True
@staticmethod @staticmethod
@profiler.record_function("trajopt_config/load_from_robot_config") @profiler.record_function("trajopt_config/load_from_robot_config")
@@ -98,14 +100,14 @@ class TrajOptSolverConfig:
interpolation_type: InterpolateType = InterpolateType.LINEAR_CUDA, interpolation_type: InterpolateType = InterpolateType.LINEAR_CUDA,
interpolation_steps: int = 10000, interpolation_steps: int = 10000,
interpolation_dt: float = 0.01, interpolation_dt: float = 0.01,
use_cuda_graph: Optional[bool] = None, use_cuda_graph: bool = True,
self_collision_check: bool = False, self_collision_check: bool = False,
self_collision_opt: bool = True, self_collision_opt: bool = True,
grad_trajopt_iters: Optional[int] = None, grad_trajopt_iters: Optional[int] = None,
num_seeds: int = 2, num_seeds: int = 2,
seed_ratio: Dict[str, int] = {"linear": 1.0, "bias": 0.0, "start": 0.0, "end": 0.0}, seed_ratio: Dict[str, int] = {"linear": 1.0, "bias": 0.0, "start": 0.0, "end": 0.0},
use_particle_opt: bool = True, use_particle_opt: bool = True,
collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.PRIMITIVE, collision_checker_type: Optional[CollisionCheckerType] = CollisionCheckerType.MESH,
traj_evaluator_config: TrajEvaluatorConfig = TrajEvaluatorConfig(), traj_evaluator_config: TrajEvaluatorConfig = TrajEvaluatorConfig(),
traj_evaluator: Optional[TrajEvaluator] = None, traj_evaluator: Optional[TrajEvaluator] = None,
minimize_jerk: bool = True, minimize_jerk: bool = True,
@@ -128,6 +130,7 @@ class TrajOptSolverConfig:
state_finite_difference_mode: Optional[str] = None, state_finite_difference_mode: Optional[str] = None,
filter_robot_command: bool = False, filter_robot_command: bool = False,
optimize_dt: bool = True, optimize_dt: bool = True,
project_pose_to_goal_frame: bool = True,
): ):
# NOTE: Don't have default optimize_dt, instead read from a configuration file. # NOTE: Don't have default optimize_dt, instead read from a configuration file.
# use default values, disable environment collision checking # use default values, disable environment collision checking
@@ -163,6 +166,16 @@ class TrajOptSolverConfig:
if traj_tsteps is None: if traj_tsteps is None:
traj_tsteps = grad_config_data["model"]["horizon"] traj_tsteps = grad_config_data["model"]["horizon"]
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
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 config_data["model"]["horizon"] = traj_tsteps
grad_config_data["model"]["horizon"] = traj_tsteps grad_config_data["model"]["horizon"] = traj_tsteps
if minimize_jerk is not None: if minimize_jerk is not None:
@@ -212,8 +225,8 @@ class TrajOptSolverConfig:
if not self_collision_opt: if not self_collision_opt:
config_data["cost"]["self_collision_cfg"]["weight"] = 0.0 config_data["cost"]["self_collision_cfg"]["weight"] = 0.0
grad_config_data["cost"]["self_collision_cfg"]["weight"] = 0.0 grad_config_data["cost"]["self_collision_cfg"]["weight"] = 0.0
config_data["mppi"]["n_envs"] = 1 config_data["mppi"]["n_problems"] = 1
grad_config_data["lbfgs"]["n_envs"] = 1 grad_config_data["lbfgs"]["n_problems"] = 1
if fixed_iters is not None: if fixed_iters is not None:
grad_config_data["lbfgs"]["fixed_iters"] = fixed_iters grad_config_data["lbfgs"]["fixed_iters"] = fixed_iters
@@ -256,8 +269,9 @@ class TrajOptSolverConfig:
world_coll_checker=world_coll_checker, world_coll_checker=world_coll_checker,
tensor_args=tensor_args, tensor_args=tensor_args,
) )
arm_rollout_mppi = None
with profiler.record_function("trajopt_config/create_rollouts"): with profiler.record_function("trajopt_config/create_rollouts"):
if use_particle_opt:
arm_rollout_mppi = ArmReacher(cfg) arm_rollout_mppi = ArmReacher(cfg)
arm_rollout_grad = ArmReacher(grad_cfg) arm_rollout_grad = ArmReacher(grad_cfg)
@@ -266,20 +280,22 @@ class TrajOptSolverConfig:
aux_rollout = ArmReacher(safety_cfg) aux_rollout = ArmReacher(safety_cfg)
interpolate_rollout = ArmReacher(safety_cfg) interpolate_rollout = ArmReacher(safety_cfg)
if trajopt_dt is not None: if trajopt_dt is not None:
if arm_rollout_mppi is not None:
arm_rollout_mppi.update_traj_dt(dt=trajopt_dt) arm_rollout_mppi.update_traj_dt(dt=trajopt_dt)
aux_rollout.update_traj_dt(dt=trajopt_dt) aux_rollout.update_traj_dt(dt=trajopt_dt)
arm_rollout_grad.update_traj_dt(dt=trajopt_dt) arm_rollout_grad.update_traj_dt(dt=trajopt_dt)
arm_rollout_safety.update_traj_dt(dt=trajopt_dt) arm_rollout_safety.update_traj_dt(dt=trajopt_dt)
if arm_rollout_mppi is not None:
config_dict = ParallelMPPIConfig.create_data_dict( config_dict = ParallelMPPIConfig.create_data_dict(
config_data["mppi"], arm_rollout_mppi, tensor_args config_data["mppi"], arm_rollout_mppi, tensor_args
) )
parallel_mppi = None
if use_es is not None and use_es: if use_es is not None and use_es:
mppi_cfg = ParallelESConfig(**config_dict) mppi_cfg = ParallelESConfig(**config_dict)
if es_learning_rate is not None: if es_learning_rate is not None:
mppi_cfg.learning_rate = es_learning_rate mppi_cfg.learning_rate = es_learning_rate
parallel_mppi = ParallelES(mppi_cfg) parallel_mppi = ParallelES(mppi_cfg)
else: elif use_particle_opt:
mppi_cfg = ParallelMPPIConfig(**config_dict) mppi_cfg = ParallelMPPIConfig(**config_dict)
parallel_mppi = ParallelMPPI(mppi_cfg) parallel_mppi = ParallelMPPI(mppi_cfg)
@@ -307,7 +323,7 @@ class TrajOptSolverConfig:
cfg = WrapConfig( cfg = WrapConfig(
safety_rollout=arm_rollout_safety, safety_rollout=arm_rollout_safety,
optimizers=opt_list, optimizers=opt_list,
compute_metrics=not evaluate_interpolated_trajectory, compute_metrics=True, # not evaluate_interpolated_trajectory,
use_cuda_graph_metrics=grad_config_data["lbfgs"]["use_cuda_graph"], use_cuda_graph_metrics=grad_config_data["lbfgs"]["use_cuda_graph"],
sync_cuda_time=sync_cuda_time, sync_cuda_time=sync_cuda_time,
) )
@@ -337,6 +353,7 @@ class TrajOptSolverConfig:
trim_steps=trim_steps, trim_steps=trim_steps,
store_debug_in_result=store_debug_in_result, store_debug_in_result=store_debug_in_result,
optimize_dt=optimize_dt, optimize_dt=optimize_dt,
use_cuda_graph=use_cuda_graph,
) )
return trajopt_cfg return trajopt_cfg
@@ -360,6 +377,7 @@ class TrajResult(Sequence):
optimized_dt: Optional[torch.Tensor] = None optimized_dt: Optional[torch.Tensor] = None
raw_solution: Optional[JointState] = None raw_solution: Optional[JointState] = None
raw_action: Optional[torch.Tensor] = None raw_action: Optional[torch.Tensor] = None
goalset_index: Optional[torch.Tensor] = None
def __getitem__(self, idx): def __getitem__(self, idx):
# position_error = rotation_error = cspace_error = path_buffer_last_tstep = None # position_error = rotation_error = cspace_error = path_buffer_last_tstep = None
@@ -372,6 +390,7 @@ class TrajResult(Sequence):
self.position_error, self.position_error,
self.rotation_error, self.rotation_error,
self.cspace_error, self.cspace_error,
self.goalset_index,
] ]
idx_vals = list_idx_if_not_none(d_list, idx) idx_vals = list_idx_if_not_none(d_list, idx)
@@ -388,6 +407,7 @@ class TrajResult(Sequence):
position_error=idx_vals[3], position_error=idx_vals[3],
rotation_error=idx_vals[4], rotation_error=idx_vals[4],
cspace_error=idx_vals[5], cspace_error=idx_vals[5],
goalset_index=idx_vals[6],
) )
def __len__(self): def __len__(self):
@@ -405,7 +425,7 @@ class TrajOptSolver(TrajOptSolverConfig):
3, int(self.action_horizon / 2), self.tensor_args 3, int(self.action_horizon / 2), self.tensor_args
).unsqueeze(0) ).unsqueeze(0)
assert self.action_horizon / 2 != 0.0 assert self.action_horizon / 2 != 0.0
self.solver.update_nenvs(self.num_seeds) self.solver.update_nproblems(self.num_seeds)
self._max_joint_vel = ( self._max_joint_vel = (
self.solver.safety_rollout.state_bounds.velocity.view(2, self.dof)[1, :].reshape( self.solver.safety_rollout.state_bounds.velocity.view(2, self.dof)[1, :].reshape(
1, 1, self.dof 1, 1, self.dof
@@ -469,12 +489,18 @@ class TrajOptSolver(TrajOptSolverConfig):
self._goal_buffer, self._goal_buffer,
self.tensor_args, self.tensor_args,
) )
if update_reference: if update_reference:
self.solver.update_nenvs(self._solve_state.get_batch_size()) if self.use_cuda_graph and self._col is not None:
log_error("changing goal type, breaking previous cuda graph.")
self.reset_cuda_graph() self.reset_cuda_graph()
self.solver.update_nproblems(self._solve_state.get_batch_size())
self._col = torch.arange( self._col = torch.arange(
0, goal.batch, device=self.tensor_args.device, dtype=torch.long 0, goal.batch, device=self.tensor_args.device, dtype=torch.long
) )
self.reset_shape()
return self._goal_buffer return self._goal_buffer
def solve_any( def solve_any(
@@ -586,6 +612,8 @@ class TrajOptSolver(TrajOptSolverConfig):
num_seeds, num_seeds,
solve_state.batch_mode, solve_state.batch_mode,
) )
if traj_result.goalset_index is not None:
traj_result.goalset_index[traj_result.goalset_index >= goal.goal_pose.n_goalset] = 0
if newton_iters is not None: if newton_iters is not None:
self.solver.newton_optimizer.outer_iters = self._og_newton_iters self.solver.newton_optimizer.outer_iters = self._og_newton_iters
self.solver.newton_optimizer.fixed_iters = self._og_newton_fixed_iters self.solver.newton_optimizer.fixed_iters = self._og_newton_fixed_iters
@@ -839,16 +867,18 @@ class TrajOptSolver(TrajOptSolverConfig):
if self.evaluate_interpolated_trajectory: if self.evaluate_interpolated_trajectory:
with profiler.record_function("trajopt/evaluate_interpolated"): with profiler.record_function("trajopt/evaluate_interpolated"):
if self.use_cuda_graph_metrics: if self.use_cuda_graph_metrics:
result.metrics = self.interpolate_rollout.get_metrics_cuda_graph( metrics = self.interpolate_rollout.get_metrics_cuda_graph(interpolated_trajs)
interpolated_trajs
)
else: else:
result.metrics = self.interpolate_rollout.get_metrics(interpolated_trajs) metrics = self.interpolate_rollout.get_metrics(interpolated_trajs)
result.metrics.feasible = metrics.feasible
result.metrics.position_error = metrics.position_error
result.metrics.rotation_error = metrics.rotation_error
result.metrics.cspace_error = metrics.cspace_error
result.metrics.goalset_index = metrics.goalset_index
st_time = time.time() st_time = time.time()
feasible = torch.all(result.metrics.feasible, dim=-1) feasible = torch.all(result.metrics.feasible, dim=-1)
# if self.num_seeds == 1:
# print(result.metrics)
if result.metrics.position_error is not None: if result.metrics.position_error is not None:
converge = torch.logical_and( converge = torch.logical_and(
result.metrics.position_error[..., -1] <= self.position_threshold, result.metrics.position_error[..., -1] <= self.position_threshold,
@@ -877,10 +907,10 @@ class TrajOptSolver(TrajOptSolverConfig):
optimized_dt=opt_dt, optimized_dt=opt_dt,
raw_solution=result.action, raw_solution=result.action,
raw_action=result.raw_action, raw_action=result.raw_action,
goalset_index=result.metrics.goalset_index,
) )
else: else:
# get path length: # get path length:
# max_vel =
if self.evaluate_interpolated_trajectory: if self.evaluate_interpolated_trajectory:
smooth_label, smooth_cost = self.traj_evaluator.evaluate_interpolated_smootheness( smooth_label, smooth_cost = self.traj_evaluator.evaluate_interpolated_smootheness(
interpolated_trajs, interpolated_trajs,
@@ -896,7 +926,6 @@ class TrajOptSolver(TrajOptSolverConfig):
self.solver.rollout_fn.dynamics_model.cspace_distance_weight, self.solver.rollout_fn.dynamics_model.cspace_distance_weight,
self._velocity_bounds, self._velocity_bounds,
) )
# print(smooth_label, success, self._velocity_bounds.shape, self.solver.rollout_fn.dynamics_model.cspace_distance_weight)
with profiler.record_function("trajopt/best_select"): with profiler.record_function("trajopt/best_select"):
success[~smooth_label] = False success[~smooth_label] = False
@@ -907,7 +936,8 @@ class TrajOptSolver(TrajOptSolverConfig):
convergence_error = result.metrics.cspace_error[..., -1] convergence_error = result.metrics.cspace_error[..., -1]
else: else:
raise ValueError("convergence check requires either goal_pose or goal_state") raise ValueError("convergence check requires either goal_pose or goal_state")
error = convergence_error + smooth_cost running_cost = torch.mean(result.metrics.cost, dim=-1) * 0.0001
error = convergence_error + smooth_cost + running_cost
error[~success] += 10000.0 error[~success] += 10000.0
if batch_mode: if batch_mode:
idx = torch.argmin(error.view(goal.batch, num_seeds), dim=-1) idx = torch.argmin(error.view(goal.batch, num_seeds), dim=-1)
@@ -923,13 +953,16 @@ class TrajOptSolver(TrajOptSolverConfig):
best_act_seq = result.action[idx] best_act_seq = result.action[idx]
best_raw_action = result.raw_action[idx] best_raw_action = result.raw_action[idx]
interpolated_traj = interpolated_trajs[idx] interpolated_traj = interpolated_trajs[idx]
position_error = rotation_error = cspace_error = None goalset_index = position_error = rotation_error = cspace_error = None
if result.metrics.position_error is not None: if result.metrics.position_error is not None:
position_error = result.metrics.position_error[idx, -1] position_error = result.metrics.position_error[idx, -1]
if result.metrics.rotation_error is not None: if result.metrics.rotation_error is not None:
rotation_error = result.metrics.rotation_error[idx, -1] rotation_error = result.metrics.rotation_error[idx, -1]
if result.metrics.cspace_error is not None: if result.metrics.cspace_error is not None:
cspace_error = result.metrics.cspace_error[idx, -1] cspace_error = result.metrics.cspace_error[idx, -1]
if result.metrics.goalset_index is not None:
goalset_index = result.metrics.goalset_index[idx, -1]
opt_dt = opt_dt[idx] opt_dt = opt_dt[idx]
if self.sync_cuda_time: if self.sync_cuda_time:
torch.cuda.synchronize() torch.cuda.synchronize()
@@ -965,6 +998,7 @@ class TrajOptSolver(TrajOptSolverConfig):
optimized_dt=opt_dt, optimized_dt=opt_dt,
raw_solution=best_act_seq, raw_solution=best_act_seq,
raw_action=best_raw_action, raw_action=best_raw_action,
goalset_index=goalset_index,
) )
return traj_result return traj_result
@@ -999,7 +1033,6 @@ class TrajOptSolver(TrajOptSolverConfig):
return self.solve_batch_goalset( return self.solve_batch_goalset(
goal, seed_traj, use_nn_seed, return_all_solutions, num_seeds, seed_success goal, seed_traj, use_nn_seed, return_all_solutions, num_seeds, seed_success
) )
return traj_result
def get_linear_seed(self, start_state, goal_state): def get_linear_seed(self, start_state, goal_state):
start_q = start_state.position.view(-1, 1, self.dof) start_q = start_state.position.view(-1, 1, self.dof)
@@ -1173,6 +1206,11 @@ class TrajOptSolver(TrajOptSolverConfig):
self.interpolate_rollout.reset_cuda_graph() self.interpolate_rollout.reset_cuda_graph()
self.rollout_fn.reset_cuda_graph() self.rollout_fn.reset_cuda_graph()
def reset_shape(self):
self.solver.reset_shape()
self.interpolate_rollout.reset_shape()
self.rollout_fn.reset_shape()
@property @property
def kinematics(self) -> CudaRobotModel: def kinematics(self) -> CudaRobotModel:
return self.rollout_fn.dynamics_model.robot_model return self.rollout_fn.dynamics_model.robot_model
@@ -1205,3 +1243,14 @@ class TrajOptSolver(TrajOptSolverConfig):
def get_full_js(self, active_js: JointState) -> JointState: def get_full_js(self, active_js: JointState) -> JointState:
return self.rollout_fn.get_full_dof_from_solution(active_js) return self.rollout_fn.get_full_dof_from_solution(active_js)
def update_pose_cost_metric(
self,
metric: PoseCostMetric,
):
rollouts = self.get_all_rollout_instances()
[
rollout.update_pose_cost_metric(metric)
for rollout in rollouts
if isinstance(rollout, ArmReacher)
]

124
src/curobo/wrap/reacher/types.py
View File

@@ -13,7 +13,7 @@ from __future__ import annotations
# Standard Library # Standard Library
from dataclasses import dataclass from dataclasses import dataclass
from enum import Enum from enum import Enum
from typing import Dict, List, Optional from typing import Dict, List, Optional, Union
# CuRobo # CuRobo
from curobo.rollout.rollout_base import Goal from curobo.rollout.rollout_base import Goal
@@ -126,16 +126,26 @@ class ReacherSolveState:
if ( if (
current_solve_state is None current_solve_state is None
or current_goal_buffer is None or current_goal_buffer is None
or current_solve_state != solve_state
or (current_goal_buffer.retract_state is None and retract_config is not None) or (current_goal_buffer.retract_state is None and retract_config is not None)
or (current_goal_buffer.goal_state is None and goal_state is not None) or (current_goal_buffer.goal_state is None and goal_state is not None)
or (current_goal_buffer.links_goal_pose is None and link_poses is not None) or (current_goal_buffer.links_goal_pose is None and link_poses is not None)
): ):
update_reference = True
elif current_solve_state != solve_state:
new_goal_pose = get_padded_goalset(
solve_state, current_solve_state, current_goal_buffer, goal_pose
)
if new_goal_pose is not None:
goal_pose = new_goal_pose
else:
update_reference = True
if update_reference:
current_solve_state = solve_state current_solve_state = solve_state
current_goal_buffer = solve_state.create_goal_buffer( current_goal_buffer = solve_state.create_goal_buffer(
goal_pose, goal_state, retract_config, link_poses, tensor_args goal_pose, goal_state, retract_config, link_poses, tensor_args
) )
update_reference = True
else: else:
current_goal_buffer.goal_pose.copy_(goal_pose) current_goal_buffer.goal_pose.copy_(goal_pose)
if retract_config is not None: if retract_config is not None:
@@ -145,6 +155,7 @@ class ReacherSolveState:
if link_poses is not None: if link_poses is not None:
for k in link_poses.keys(): 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])
return current_solve_state, current_goal_buffer, update_reference return current_solve_state, current_goal_buffer, update_reference
def update_goal( def update_goal(
@@ -155,17 +166,26 @@ class ReacherSolveState:
tensor_args: TensorDeviceType = TensorDeviceType(), tensor_args: TensorDeviceType = TensorDeviceType(),
): ):
solve_state = self solve_state = self
update_reference = False update_reference = False
if ( if (
current_solve_state is None current_solve_state is None
or current_goal_buffer is None or current_goal_buffer is None
or current_solve_state != solve_state
or (current_goal_buffer.goal_state is None and goal.goal_state is not None) or (current_goal_buffer.goal_state is None and goal.goal_state is not None)
or (current_goal_buffer.goal_state is not None and goal.goal_state is None) or (current_goal_buffer.goal_state is not None and goal.goal_state is None)
): ):
# TODO: Check for change in update idx buffers, currently we assume update_reference = True
# that solve_state captures difference in idx buffers elif current_solve_state != solve_state:
new_goal_pose = get_padded_goalset(
solve_state, current_solve_state, current_goal_buffer, goal.goal_pose
)
if new_goal_pose is not None:
goal = goal.clone()
goal.goal_pose = new_goal_pose
else:
update_reference = True
if update_reference:
current_solve_state = solve_state current_solve_state = solve_state
current_goal_buffer = goal.create_index_buffers( current_goal_buffer = goal.create_index_buffers(
solve_state.batch_size, solve_state.batch_size,
@@ -174,7 +194,6 @@ class ReacherSolveState:
solve_state.num_seeds, solve_state.num_seeds,
tensor_args, tensor_args,
) )
update_reference = True
else: else:
current_goal_buffer.copy_(goal, update_idx_buffers=False) current_goal_buffer.copy_(goal, update_idx_buffers=False)
return current_solve_state, current_goal_buffer, update_reference return current_solve_state, current_goal_buffer, update_reference
@@ -185,3 +204,92 @@ class MotionGenSolverState:
solve_type: ReacherSolveType solve_type: ReacherSolveType
ik_solve_state: ReacherSolveState ik_solve_state: ReacherSolveState
trajopt_solve_state: ReacherSolveState trajopt_solve_state: ReacherSolveState
def get_padded_goalset(
solve_state: ReacherSolveState,
current_solve_state: ReacherSolveState,
current_goal_buffer: Goal,
new_goal_pose: Pose,
) -> Union[Pose, None]:
if (
current_solve_state.solve_type == ReacherSolveType.GOALSET
and solve_state.solve_type == ReacherSolveType.SINGLE
):
# convert single goal to goal set
# solve_state.solve_type = ReacherSolveType.GOALSET
# solve_state.n_goalset = current_solve_state.n_goalset
goal_pose = current_goal_buffer.goal_pose.clone()
goal_pose.position[:] = new_goal_pose.position
goal_pose.quaternion[:] = new_goal_pose.quaternion
return goal_pose
elif (
current_solve_state.solve_type == ReacherSolveType.BATCH_GOALSET
and solve_state.solve_type == ReacherSolveType.BATCH
and new_goal_pose.n_goalset <= current_solve_state.n_goalset
and new_goal_pose.batch == current_solve_state.batch_size
):
goal_pose = current_goal_buffer.goal_pose.clone()
if len(new_goal_pose.position.shape) == 2:
new_goal_pose = new_goal_pose.unsqueeze(1)
goal_pose.position[..., :, :] = new_goal_pose.position
goal_pose.quaternion[..., :, :] = new_goal_pose.quaternion
return goal_pose
elif (
current_solve_state.solve_type == ReacherSolveType.BATCH_ENV_GOALSET
and solve_state.solve_type == ReacherSolveType.BATCH_ENV
and new_goal_pose.n_goalset <= current_solve_state.n_goalset
and new_goal_pose.batch == current_solve_state.batch_size
):
goal_pose = current_goal_buffer.goal_pose.clone()
if len(new_goal_pose.position.shape) == 2:
new_goal_pose = new_goal_pose.unsqueeze(1)
goal_pose.position[..., :, :] = new_goal_pose.position
goal_pose.quaternion[..., :, :] = new_goal_pose.quaternion
return goal_pose
elif (
current_solve_state.solve_type == ReacherSolveType.GOALSET
and solve_state.solve_type == ReacherSolveType.GOALSET
and new_goal_pose.n_goalset <= current_solve_state.n_goalset
):
goal_pose = current_goal_buffer.goal_pose.clone()
goal_pose.position[..., : new_goal_pose.n_goalset, :] = new_goal_pose.position
goal_pose.quaternion[..., : new_goal_pose.n_goalset, :] = new_goal_pose.quaternion
goal_pose.position[..., new_goal_pose.n_goalset :, :] = new_goal_pose.position[..., :1, :]
goal_pose.quaternion[..., new_goal_pose.n_goalset :, :] = new_goal_pose.quaternion[
..., :1, :
]
return goal_pose
elif (
current_solve_state.solve_type == ReacherSolveType.BATCH_GOALSET
and solve_state.solve_type == ReacherSolveType.BATCH_GOALSET
and new_goal_pose.n_goalset <= current_solve_state.n_goalset
and new_goal_pose.batch == current_solve_state.batch_size
):
goal_pose = current_goal_buffer.goal_pose.clone()
goal_pose.position[..., : new_goal_pose.n_goalset, :] = new_goal_pose.position
goal_pose.quaternion[..., : new_goal_pose.n_goalset, :] = new_goal_pose.quaternion
goal_pose.position[..., new_goal_pose.n_goalset :, :] = new_goal_pose.position[..., :1, :]
goal_pose.quaternion[..., new_goal_pose.n_goalset :, :] = new_goal_pose.quaternion[
..., :1, :
]
return goal_pose
elif (
current_solve_state.solve_type == ReacherSolveType.BATCH_ENV_GOALSET
and solve_state.solve_type == ReacherSolveType.BATCH_ENV_GOALSET
and new_goal_pose.n_goalset <= current_solve_state.n_goalset
and new_goal_pose.batch == current_solve_state.batch_size
):
goal_pose = current_goal_buffer.goal_pose.clone()
goal_pose.position[..., : new_goal_pose.n_goalset, :] = new_goal_pose.position
goal_pose.quaternion[..., : new_goal_pose.n_goalset, :] = new_goal_pose.quaternion
goal_pose.position[..., new_goal_pose.n_goalset :, :] = new_goal_pose.position[..., :1, :]
goal_pose.quaternion[..., new_goal_pose.n_goalset :, :] = new_goal_pose.quaternion[
..., :1, :
]
return goal_pose
return None

16
src/curobo/wrap/wrap_base.py
View File

@@ -54,7 +54,7 @@ class WrapBase(WrapConfig):
def __init__(self, config: Optional[WrapConfig] = None): def __init__(self, config: Optional[WrapConfig] = None):
if config is not None: if config is not None:
WrapConfig.__init__(self, **vars(config)) WrapConfig.__init__(self, **vars(config))
self.n_envs = 1 self.n_problems = 1
self.opt_dt = 0 self.opt_dt = 0
self._rollout_list = None self._rollout_list = None
self._opt_rollouts = None self._opt_rollouts = None
@@ -83,11 +83,11 @@ class WrapBase(WrapConfig):
debug_list.append(opt.debug_cost) debug_list.append(opt.debug_cost)
return debug_list return debug_list
def update_nenvs(self, n_envs): def update_nproblems(self, n_problems):
if n_envs != self.n_envs: if n_problems != self.n_problems:
self.n_envs = n_envs self.n_problems = n_problems
for opt in self.optimizers: for opt in self.optimizers:
opt.update_nenvs(self.n_envs) opt.update_nproblems(self.n_problems)
def update_params(self, goal: Goal): def update_params(self, goal: Goal):
with profiler.record_function("wrap_base/safety/update_params"): with profiler.record_function("wrap_base/safety/update_params"):
@@ -117,6 +117,12 @@ class WrapBase(WrapConfig):
opt.reset_cuda_graph() opt.reset_cuda_graph()
self._init_solver = False self._init_solver = False
def reset_shape(self):
self.safety_rollout.reset_shape()
for opt in self.optimizers:
opt.reset_shape()
self._init_solver = False
@property @property
def rollout_fn(self): def rollout_fn(self):
return self.safety_rollout return self.safety_rollout

14
tests/goal_test.py
View File

@@ -47,6 +47,9 @@ def test_repeat_seeds():
out_r_v = torch.zeros((b, 4), device=tensor_args.device, dtype=tensor_args.dtype) out_r_v = torch.zeros((b, 4), device=tensor_args.device, dtype=tensor_args.dtype)
out_idx = torch.zeros((b, 1), device=tensor_args.device, dtype=torch.int32) out_idx = torch.zeros((b, 1), device=tensor_args.device, dtype=torch.int32)
vec_weight = torch.ones((6), device=tensor_args.device, dtype=tensor_args.dtype) vec_weight = torch.ones((6), device=tensor_args.device, dtype=tensor_args.dtype)
offset_waypoint = torch.zeros((6), device=tensor_args.device, dtype=tensor_args.dtype)
offset_tstep_fraction = torch.ones((1), device=tensor_args.device, dtype=tensor_args.dtype)
weight = tensor_args.to_device([1, 1, 1, 1]) weight = tensor_args.to_device([1, 1, 1, 1])
vec_convergence = tensor_args.to_device([0, 0]) vec_convergence = tensor_args.to_device([0, 0])
run_weight = tensor_args.to_device([1]) run_weight = tensor_args.to_device([1])
@@ -66,6 +69,8 @@ def test_repeat_seeds():
vec_convergence, vec_convergence,
run_weight, run_weight,
vec_weight.clone() * 0.0, vec_weight.clone() * 0.0,
offset_waypoint,
offset_tstep_fraction,
g.batch_pose_idx, g.batch_pose_idx,
start_pose.position.shape[0], start_pose.position.shape[0],
1, 1,
@@ -74,6 +79,7 @@ def test_repeat_seeds():
False, False,
False, False,
True, True,
True,
) )
assert torch.sum(r[0]).item() == 0.0 assert torch.sum(r[0]).item() == 0.0
@@ -112,6 +118,9 @@ def test_horizon_repeat_seeds():
out_r_v = torch.zeros((b, h, 4), device=tensor_args.device, dtype=tensor_args.dtype) out_r_v = torch.zeros((b, h, 4), device=tensor_args.device, dtype=tensor_args.dtype)
out_idx = torch.zeros((b, h, 1), device=tensor_args.device, dtype=torch.int32) out_idx = torch.zeros((b, h, 1), device=tensor_args.device, dtype=torch.int32)
vec_weight = torch.ones((6), device=tensor_args.device, dtype=tensor_args.dtype) vec_weight = torch.ones((6), device=tensor_args.device, dtype=tensor_args.dtype)
offset_waypoint = torch.zeros((6), device=tensor_args.device, dtype=tensor_args.dtype)
offset_tstep_fraction = torch.ones((1), device=tensor_args.device, dtype=tensor_args.dtype)
weight = tensor_args.to_device([1, 1, 1, 1]) weight = tensor_args.to_device([1, 1, 1, 1])
vec_convergence = tensor_args.to_device([0, 0]) vec_convergence = tensor_args.to_device([0, 0])
run_weight = torch.zeros((1, h), device=tensor_args.device) run_weight = torch.zeros((1, h), device=tensor_args.device)
@@ -132,6 +141,8 @@ def test_horizon_repeat_seeds():
vec_convergence, vec_convergence,
run_weight, run_weight,
vec_weight.clone() * 0.0, vec_weight.clone() * 0.0,
offset_waypoint,
offset_tstep_fraction,
g.batch_pose_idx, g.batch_pose_idx,
start_pose.position.shape[0], start_pose.position.shape[0],
h, h,
@@ -140,5 +151,6 @@ def test_horizon_repeat_seeds():
True, True,
False, False,
False, False,
True,
) )
assert torch.sum(r[0]).item() == 0.0 assert torch.sum(r[0]).item() < 1e-6

2
tests/ik_config_test.py
View File

@@ -146,5 +146,5 @@ def test_eval(config, expected):
result = ik_solver.solve_single(goal) result = ik_solver.solve_single(goal)
success = result.success success = result.success
if expected is not -100: if expected != -100:
assert success.item() == expected assert success.item() == expected

32
tests/kinematics_test.py
View File

@@ -169,3 +169,35 @@ def test_locked_joints_kinematics():
assert torch.linalg.norm(out.ee_position - out_locked.ee_position) < 1e-5 assert torch.linalg.norm(out.ee_position - out_locked.ee_position) < 1e-5
assert torch.linalg.norm(out.ee_quaternion - out_locked.ee_quaternion) < 1e-5 assert torch.linalg.norm(out.ee_quaternion - out_locked.ee_quaternion) < 1e-5
assert torch.linalg.norm(out.link_spheres_tensor - out_locked.link_spheres_tensor) < 1e-5 assert torch.linalg.norm(out.link_spheres_tensor - out_locked.link_spheres_tensor) < 1e-5
def test_franka_toggle_link_collision(cfg):
tensor_args = TensorDeviceType()
robot_model = CudaRobotModel(cfg)
q_test = torch.as_tensor([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **vars(tensor_args)).view(1, -1)
sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("panda_link5")
state = robot_model.get_state(q_test.clone())
attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
# check if all values are equal to position and quaternion
link_radius = attached_spheres[..., 3].clone()
assert torch.count_nonzero(link_radius <= 0.0) == 0
robot_model.kinematics_config.disable_link_spheres("panda_link5")
state = robot_model.get_state(q_test.clone())
attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
sph_radius = attached_spheres[..., 3].clone()
assert torch.count_nonzero(sph_radius < 0.0)
robot_model.kinematics_config.enable_link_spheres("panda_link5")
state = robot_model.get_state(q_test.clone())
radius = link_radius - state.link_spheres_tensor[:, sph_idx, 3]
assert torch.count_nonzero(radius == 0.0)

51
tests/motion_gen_api_test.py
View File

@@ -14,11 +14,10 @@ import pytest
import torch import torch
# CuRobo # CuRobo
from curobo.geom.types import WorldConfig
from curobo.types.base import TensorDeviceType from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose from curobo.types.math import Pose
from curobo.types.robot import JointState, RobotConfig from curobo.types.robot import JointState
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml from curobo.util_file import get_robot_configs_path, join_path, load_yaml
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
@@ -60,3 +59,49 @@ def test_motion_gen_attach_obstacle_offset(motion_gen):
start_state, [obstacle], world_objects_pose_offset=offset_pose start_state, [obstacle], world_objects_pose_offset=offset_pose
) )
assert True assert True
def test_motion_gen_lock_js_update():
tensor_args = TensorDeviceType()
world_file = "collision_table.yml"
robot_file = "franka_mobile.yml"
robot_config = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
robot_config["kinematics"]["lock_joints"] = {"base_x": 0.0, "base_y": 0.0, "base_z": 0.0}
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_config,
world_file,
tensor_args,
use_cuda_graph=True,
)
motion_gen_instance = MotionGen(motion_gen_config)
motion_gen_instance.warmup()
retract_cfg = motion_gen_instance.get_retract_config()
start_state = JointState.from_position(retract_cfg.view(1, -1))
kin_state = motion_gen_instance.compute_kinematics(start_state)
ee_pose = kin_state.ee_pose.clone()
# test motion gen:
plan_start = start_state.clone()
plan_start.position[..., :-2] += 0.1
result = motion_gen_instance.plan_single(plan_start, ee_pose)
assert result.success.item()
lock_js = {"base_x": 1.0, "base_y": 0.0, "base_z": 0.0}
motion_gen_instance.update_locked_joints(lock_js, robot_config)
kin_state_new = motion_gen_instance.compute_kinematics(start_state)
ee_pose_shift = kin_state_new.ee_pose.clone()
assert torch.norm(ee_pose.position[..., 0] - ee_pose_shift.position[..., 0]).item() == 1.0
assert torch.norm(ee_pose.position[..., 1:] - ee_pose_shift.position[..., 1:]).item() == 0.0
# test motion gen with new lock state:
result = motion_gen_instance.plan_single(plan_start, ee_pose_shift)
assert result.success.item()
result = motion_gen_instance.plan_single(
plan_start, ee_pose, MotionGenPlanConfig(max_attempts=3)
)
assert result.success.item() == False

314
tests/motion_gen_constrained_test.py Normal file
View File

@@ -0,0 +1,314 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Third Party
import pytest
import torch
# CuRobo
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import JointState
from curobo.wrap.reacher.motion_gen import (
MotionGen,
MotionGenConfig,
MotionGenPlanConfig,
PoseCostMetric,
)
@pytest.fixture(scope="module")
def motion_gen(request):
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,
use_cuda_graph=True,
project_pose_to_goal_frame=request.param,
)
motion_gen_instance = MotionGen(motion_gen_config)
motion_gen_instance.warmup(enable_graph=False, warmup_js_trajopt=False)
return motion_gen_instance
@pytest.mark.parametrize(
"motion_gen",
[
(True),
(False),
],
indirect=True,
)
def test_approach_grasp_pose(motion_gen):
# run full pose planning
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = state.ee_pose.clone()
goal_pose.position[0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
m_config = MotionGenPlanConfig(max_attempts=1)
result = motion_gen.plan_single(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 1
# run grasp pose planning:
m_config.pose_cost_metric = PoseCostMetric.create_grasp_approach_metric()
result = motion_gen.plan_single(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 1
@pytest.mark.parametrize(
"motion_gen",
[
(True),
(False),
],
indirect=True,
)
def test_reach_only_position(motion_gen):
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = state.ee_pose.clone()
goal_pose.position[0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
m_config = MotionGenPlanConfig(
max_attempts=1,
pose_cost_metric=PoseCostMetric(
reach_partial_pose=True,
reach_vec_weight=motion_gen.tensor_args.to_device([0, 0, 0, 1, 1, 1]),
),
)
result = motion_gen.plan_single(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 1
reached_state = result.optimized_plan[-1]
reached_pose = motion_gen.compute_kinematics(reached_state).ee_pose.clone()
assert goal_pose.angular_distance(reached_pose) > 0.0
assert goal_pose.linear_distance(reached_pose) <= 0.005
@pytest.mark.parametrize(
"motion_gen",
[
(True),
(False),
],
indirect=True,
)
def test_reach_only_orientation(motion_gen):
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = state.ee_pose.clone()
goal_pose.position[0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
m_config = MotionGenPlanConfig(
max_attempts=1,
pose_cost_metric=PoseCostMetric(
reach_partial_pose=True,
reach_vec_weight=motion_gen.tensor_args.to_device([1, 1, 1, 0, 0, 0]),
),
)
result = motion_gen.plan_single(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 1
reached_state = result.optimized_plan[-1]
reached_pose = motion_gen.compute_kinematics(reached_state).ee_pose.clone()
assert goal_pose.linear_distance(reached_pose) > 0.0
assert goal_pose.angular_distance(reached_pose) < 0.05
@pytest.mark.parametrize(
"motion_gen",
[
(True),
(False),
],
indirect=True,
)
def test_hold_orientation(motion_gen):
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = state.ee_pose.clone()
goal_pose.position[0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
start_pose = motion_gen.compute_kinematics(start_state).ee_pose.clone()
goal_pose.quaternion = start_pose.quaternion.clone()
m_config = MotionGenPlanConfig(
max_attempts=1,
pose_cost_metric=PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([1, 1, 1, 0, 0, 0]),
),
)
result = motion_gen.plan_single(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 1
traj_pose = motion_gen.compute_kinematics(result.optimized_plan).ee_pose.clone()
# assert goal_pose.linear_distance(traj_pose) > 0.0
goal_pose = goal_pose.repeat(traj_pose.shape[0])
assert torch.max(goal_pose.angular_distance(traj_pose)) < 0.05
@pytest.mark.parametrize(
"motion_gen",
[
(True),
(False),
],
indirect=True,
)
def test_hold_position(motion_gen):
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = state.ee_pose.clone()
goal_pose.position[0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
start_pose = motion_gen.compute_kinematics(start_state).ee_pose.clone()
goal_pose.position = start_pose.position.clone()
m_config = MotionGenPlanConfig(
max_attempts=1,
pose_cost_metric=PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([0, 0, 0, 1, 1, 1]),
),
)
result = motion_gen.plan_single(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 1
traj_pose = motion_gen.compute_kinematics(result.optimized_plan).ee_pose.clone()
goal_pose = goal_pose.repeat(traj_pose.shape[0])
assert torch.max(goal_pose.linear_distance(traj_pose)) < 0.005
@pytest.mark.parametrize(
"motion_gen",
[
(False),
(True),
],
indirect=True,
)
def test_hold_partial_pose(motion_gen):
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = state.ee_pose.clone()
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
start_pose = motion_gen.compute_kinematics(start_state).ee_pose.clone()
goal_pose.position = start_pose.position.clone()
goal_pose.quaternion = start_pose.quaternion.clone()
if motion_gen.project_pose_to_goal_frame:
offset_pose = Pose.from_list([0, 0.1, 0, 1, 0, 0, 0])
goal_pose = goal_pose.multiply(offset_pose)
else:
goal_pose.position[0, 1] += 0.2
m_config = MotionGenPlanConfig(
max_attempts=1,
pose_cost_metric=PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([1, 1, 1, 1, 0, 1]),
),
)
result = motion_gen.plan_single(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 1
traj_pose = motion_gen.compute_kinematics(result.optimized_plan).ee_pose.clone()
goal_pose = goal_pose.repeat(traj_pose.shape[0])
if motion_gen.project_pose_to_goal_frame:
traj_pose = goal_pose.compute_local_pose(traj_pose)
traj_pose.position[:, 1] = 0.0
assert torch.max(traj_pose.position) < 0.005
else:
goal_pose.position[:, 1] = 0.0
traj_pose.position[:, 1] = 0.0
assert torch.max(goal_pose.linear_distance(traj_pose)) < 0.005
@pytest.mark.parametrize(
"motion_gen",
[
(False),
(True),
],
indirect=True,
)
def test_hold_partial_pose_fail(motion_gen):
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = state.ee_pose.clone()
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
start_pose = motion_gen.compute_kinematics(start_state).ee_pose.clone()
goal_pose.position = start_pose.position.clone()
goal_pose.quaternion = start_pose.quaternion.clone()
if motion_gen.project_pose_to_goal_frame:
offset_pose = Pose.from_list([0, 0.1, 0.1, 1, 0, 0, 0])
goal_pose = goal_pose.multiply(offset_pose)
else:
goal_pose.position[0, 1] += 0.2
goal_pose.position[0, 0] += 0.2
m_config = MotionGenPlanConfig(
max_attempts=1,
pose_cost_metric=PoseCostMetric(
hold_partial_pose=True,
hold_vec_weight=motion_gen.tensor_args.to_device([1, 1, 1, 1, 0, 1]),
),
)
result = motion_gen.plan_single(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 0
assert result.valid_query == False

201
tests/motion_gen_goalset_test.py Normal file
View File

@@ -0,0 +1,201 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Third Party
import pytest
import torch
# CuRobo
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import JointState
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
@pytest.fixture(scope="function")
def motion_gen():
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,
use_cuda_graph=True,
)
motion_gen_instance = MotionGen(motion_gen_config)
return motion_gen_instance
@pytest.fixture(scope="function")
def motion_gen_batch():
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,
use_cuda_graph=True,
)
motion_gen_instance = MotionGen(motion_gen_config)
return motion_gen_instance
def test_goalset_padded(motion_gen):
# run goalset planning
motion_gen.warmup(n_goalset=10)
motion_gen.reset()
m_config = MotionGenPlanConfig(False, True)
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.repeat(10, 1).view(1, -1, 3),
quaternion=state.ee_quat_seq.repeat(10, 1).view(1, -1, 4),
)
goal_pose.position[0, 0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
result = motion_gen.plan_goalset(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 1
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.repeat(2, 1).view(1, -1, 3),
quaternion=state.ee_quat_seq.repeat(2, 1).view(1, -1, 4),
)
goal_pose.position[0, 0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
result = motion_gen.plan_goalset(start_state, goal_pose, m_config.clone())
# run goalset with less goals:
# run single planning
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(state.ee_pos_seq, quaternion=state.ee_quat_seq)
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
result = motion_gen.plan_single(start_state, goal_pose, m_config.clone())
assert torch.count_nonzero(result.success) == 1
def test_batch_goalset_padded(motion_gen_batch):
motion_gen = motion_gen_batch
motion_gen.warmup(n_goalset=10, batch=3, enable_graph=False)
# run goalset planning
motion_gen.reset()
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.repeat(3 * 3, 1).view(3, -1, 3).contiguous(),
quaternion=state.ee_quat_seq.repeat(3 * 3, 1).view(3, -1, 4).contiguous(),
)
goal_pose.position[0, 1, 1] = 0.2
goal_pose.position[1, 0, 1] = 0.2
goal_pose.position[2, 1, 1] = 0.2
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.2).repeat_seeds(3)
m_config = MotionGenPlanConfig(False, True, max_attempts=10, enable_graph_attempt=20)
result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config.clone())
# get final solutions:
assert torch.count_nonzero(result.success) == result.success.shape[0]
reached_state = motion_gen.compute_kinematics(result.optimized_plan.trim_trajectory(-1))
#
goal_position = torch.cat(
[
goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0)
for x in range(len(result.goalset_index))
]
)
assert result.goalset_index is not None
assert torch.max(torch.norm(goal_position - reached_state.ee_pos_seq, dim=-1)) < 0.005
# run goalset with less goals:
motion_gen.reset()
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.repeat(6, 1).view(3, -1, 3).contiguous(),
quaternion=state.ee_quat_seq.repeat(6, 1).view(3, -1, 4).contiguous(),
)
goal_pose.position[0, 1, 1] = 0.2
goal_pose.position[1, 0, 1] = 0.2
goal_pose.position[2, 1, 1] = 0.2
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.2).repeat_seeds(3)
result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config)
# get final solutions:
assert torch.count_nonzero(result.success) == result.success.shape[0]
reached_state = motion_gen.compute_kinematics(result.optimized_plan.trim_trajectory(-1))
#
goal_position = torch.cat(
[
goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0)
for x in range(len(result.goalset_index))
]
)
assert result.goalset_index is not None
assert torch.max(torch.norm(goal_position - reached_state.ee_pos_seq, dim=-1)) < 0.005
# run single planning
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()
).repeat_seeds(3)
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(3)
goal_pose.position[1, 0] -= 0.1
result = motion_gen.plan_batch(start_state, goal_pose.clone(), m_config)
assert torch.count_nonzero(result.success) == 3
# get final solutions:
q = result.optimized_plan.trim_trajectory(-1).squeeze(1)
reached_state = motion_gen.compute_kinematics(q)
assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005

53
tests/motion_gen_module_test.py
View File

@@ -23,7 +23,7 @@ from curobo.util_file import get_robot_configs_path, get_world_configs_path, joi
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
@pytest.fixture(scope="function") @pytest.fixture(scope="module")
def motion_gen(): def motion_gen():
tensor_args = TensorDeviceType() tensor_args = TensorDeviceType()
world_file = "collision_table.yml" world_file = "collision_table.yml"
@@ -38,7 +38,7 @@ def motion_gen():
return motion_gen_instance return motion_gen_instance
@pytest.fixture(scope="function") @pytest.fixture(scope="module")
def motion_gen_batch_env(): def motion_gen_batch_env():
tensor_args = TensorDeviceType() tensor_args = TensorDeviceType()
world_files = ["collision_table.yml", "collision_test.yml"] world_files = ["collision_table.yml", "collision_test.yml"]
@@ -101,6 +101,7 @@ def test_motion_gen_goalset(motion_gen):
state.ee_pos_seq.repeat(2, 1).view(1, -1, 3), state.ee_pos_seq.repeat(2, 1).view(1, -1, 3),
quaternion=state.ee_quat_seq.repeat(2, 1).view(1, -1, 4), quaternion=state.ee_quat_seq.repeat(2, 1).view(1, -1, 4),
) )
goal_pose.position[0, 0, 0] -= 0.1
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3) start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
@@ -121,6 +122,13 @@ def test_motion_gen_goalset(motion_gen):
< 0.005 < 0.005
) )
assert result.goalset_index is not None
assert (
torch.norm(goal_pose.position[:, result.goalset_index, :] - reached_state.ee_pos_seq)
< 0.005
)
def test_motion_gen_batch_goalset(motion_gen): def test_motion_gen_batch_goalset(motion_gen):
motion_gen.reset() motion_gen.reset()
@@ -130,29 +138,36 @@ def test_motion_gen_batch_goalset(motion_gen):
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1))) state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose( goal_pose = Pose(
state.ee_pos_seq.repeat(4, 1).view(2, -1, 3), state.ee_pos_seq.repeat(6, 1).view(3, -1, 3).clone(),
quaternion=state.ee_quat_seq.repeat(4, 1).view(2, -1, 4), quaternion=state.ee_quat_seq.repeat(6, 1).view(3, -1, 4).clone(),
) )
goal_pose.position[0, 1, 1] = 0.2
goal_pose.position[1, 0, 1] = 0.2
goal_pose.position[2, 1, 1] = 0.2
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(2) start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.2).repeat_seeds(3)
m_config = MotionGenPlanConfig(False, True, num_trajopt_seeds=10) m_config = MotionGenPlanConfig(False, True, num_trajopt_seeds=10, max_attempts=1)
result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config) result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config)
# get final solutions: # get final solutions:
assert torch.count_nonzero(result.success) > 0 assert torch.count_nonzero(result.success) == result.success.shape[0]
reached_state = motion_gen.compute_kinematics(result.optimized_plan.trim_trajectory(-1)) reached_state = motion_gen.compute_kinematics(result.optimized_plan.trim_trajectory(-1))
assert ( #
torch.min( goal_position = torch.cat(
torch.norm(goal_pose.position[:, 0, :] - reached_state.ee_pos_seq), [
torch.norm(goal_pose.position[:, 1, :] - reached_state.ee_pos_seq), goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0)
) for x in range(len(result.goalset_index))
< 0.005 ]
) )
assert result.goalset_index is not None
assert torch.max(torch.norm(goal_position - reached_state.ee_pos_seq, dim=-1)) < 0.005
def test_motion_gen_batch(motion_gen): def test_motion_gen_batch(motion_gen):
motion_gen.reset() motion_gen.reset()
@@ -188,7 +203,6 @@ def test_motion_gen_batch(motion_gen):
], ],
) )
def test_motion_gen_batch_graph(motion_gen_str: str, interpolation: InterpolateType, request): def test_motion_gen_batch_graph(motion_gen_str: str, interpolation: InterpolateType, request):
# return
motion_gen = request.getfixturevalue(motion_gen_str) motion_gen = request.getfixturevalue(motion_gen_str)
motion_gen.graph_planner.interpolation_type = interpolation motion_gen.graph_planner.interpolation_type = interpolation
@@ -275,6 +289,17 @@ def test_motion_gen_batch_env_goalset(motion_gen_batch_env):
< 0.005 < 0.005
) )
goal_position = torch.cat(
[
goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0)
for x in range(len(result.goalset_index))
]
)
assert result.goalset_index is not None
assert torch.max(torch.norm(goal_position - reached_state.ee_pos_seq, dim=-1)) < 0.005
@pytest.mark.parametrize( @pytest.mark.parametrize(
"motion_gen_str,enable_graph", "motion_gen_str,enable_graph",

68
tests/motion_gen_speed_test.py Normal file
View File

@@ -0,0 +1,68 @@
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Third Party
import pytest
import torch
# CuRobo
from curobo.geom.types import WorldConfig
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import JointState, RobotConfig
from curobo.util.trajectory import InterpolateType
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
@pytest.fixture(scope="function")
def motion_gen(request):
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,
velocity_scale=request.param,
interpolation_steps=10000,
interpolation_dt=0.02,
)
motion_gen_instance = MotionGen(motion_gen_config)
return motion_gen_instance
@pytest.mark.parametrize(
"motion_gen",
[
(1.0),
(0.75),
(0.5),
(0.25),
(0.15),
(0.1),
],
indirect=True,
)
def test_motion_gen_velocity_scale(motion_gen):
retract_cfg = motion_gen.get_retract_config()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(state.ee_pos_seq, quaternion=state.ee_quat_seq)
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
m_config = MotionGenPlanConfig(False, True, max_attempts=10)
result = motion_gen.plan_single(start_state, goal_pose, m_config)
assert torch.count_nonzero(result.success) == 1

2
tests/mpc_test.py
View File

@@ -130,7 +130,7 @@ def test_mpc_single(mpc_str, expected, request):
if result.metrics.pose_error.item() < 0.05: if result.metrics.pose_error.item() < 0.05:
converged = True converged = True
break break
if tstep > 100: if tstep > 200:
break break
assert converged == expected assert converged == expected

Some files were not shown because too many files have changed in this diff Show More