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

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

@@ -16,43 +16,55 @@
// CUDA forward declarations
std::vector<torch::Tensor> self_collision_distance(
torch::Tensor out_distance, torch::Tensor out_vec,
torch::Tensor sparse_index,
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 weight,
const torch::Tensor collision_matrix, // n_spheres x n_spheres
const torch::Tensor thread_locations, 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 bool debug = false);
std::vector<torch::Tensor>self_collision_distance(
torch::Tensor out_distance,
torch::Tensor out_vec,
torch::Tensor sparse_index,
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 weight,
const torch::Tensor collision_matrix, // n_spheres x n_spheres
const torch::Tensor thread_locations,
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 bool debug = false);
// CUDA forward declarations
std::vector<torch::Tensor> swept_sphere_obb_clpt(
const torch::Tensor sphere_position, // batch_size, 3
torch::Tensor distance, // batch_size, 1
torch::Tensor
closest_point, // batch size, 4 -> written out as x,y,z,0 for gradient
torch::Tensor sparsity_idx, 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_bounds, // n_boxes, 3
const torch::Tensor obb_pose, // n_boxes, 4, 4
const torch::Tensor obb_enable, // n_boxes, 4,
const torch::Tensor n_env_obb, // 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 n_spheres, const int sweep_steps, const bool enable_speed_metric,
const bool transform_back, const bool compute_distance,
const bool use_batch_env);
std::vector<torch::Tensor>swept_sphere_obb_clpt(
const torch::Tensor sphere_position, // batch_size, 3
torch::Tensor distance, // batch_size, 1
torch::Tensor
closest_point, // batch size, 4 -> written out as x,y,z,0 for gradient
torch::Tensor sparsity_idx,
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_bounds, // n_boxes, 3
const torch::Tensor obb_pose, // n_boxes, 4, 4
const torch::Tensor obb_enable, // n_boxes, 4,
const torch::Tensor n_env_obb, // 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 n_spheres,
const int sweep_steps,
const bool enable_speed_metric,
const bool transform_back,
const bool compute_distance,
const bool use_batch_env);
std::vector<torch::Tensor>
sphere_obb_clpt(const torch::Tensor sphere_position, // batch_size, 4
torch::Tensor distance,
torch::Tensor closest_point, // batch size, 3
torch::Tensor sparsity_idx, const torch::Tensor weight,
torch::Tensor distance,
torch::Tensor closest_point, // batch size, 3
torch::Tensor sparsity_idx,
const torch::Tensor weight,
const torch::Tensor activation_distance,
const torch::Tensor obb_accel, // n_boxes, 4, 4
const torch::Tensor obb_bounds, // n_boxes, 3
@@ -60,58 +72,75 @@ sphere_obb_clpt(const torch::Tensor sphere_position, // batch_size, 4
const torch::Tensor obb_enable, // 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 int max_nobs, const int batch_size, const int horizon,
const int n_spheres, const bool transform_back,
const bool compute_distance, const bool use_batch_env);
const int max_nobs,
const int batch_size,
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(
torch::Tensor out_distance, torch::Tensor out_position_distance,
torch::Tensor out_rotation_distance,
torch::Tensor distance_p_vector, // batch size, 3
torch::Tensor distance_q_vector, // batch size, 4
torch::Tensor out_gidx,
const torch::Tensor current_position, // batch_size, 3
const torch::Tensor goal_position, // n_boxes, 3
const torch::Tensor current_quat, const torch::Tensor goal_quat,
const torch::Tensor vec_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 run_vec_weight, 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);
std::vector<torch::Tensor>pose_distance(
torch::Tensor out_distance,
torch::Tensor out_position_distance,
torch::Tensor out_rotation_distance,
torch::Tensor distance_p_vector, // batch size, 3
torch::Tensor distance_q_vector, // batch size, 4
torch::Tensor out_gidx,
const torch::Tensor current_position, // batch_size, 3
const torch::Tensor goal_position, // n_boxes, 3
const torch::Tensor current_quat,
const torch::Tensor goal_quat,
const torch::Tensor vec_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 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 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>
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_p_distance, // n_boxes, 3
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_q_vec, const int batch_size,
const bool use_distance = false);
const torch::Tensor grad_p_vec, // n_boxes, 4, 4
const torch::Tensor grad_q_vec,
const int batch_size,
const bool use_distance = false);
// C++ interface
// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4.
#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) \
AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) \
CHECK_CUDA(x); \
#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), # x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) \
AT_ASSERTM(x.is_contiguous(), # x " must be contiguous")
#define CHECK_INPUT(x) \
CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x)
std::vector<torch::Tensor> self_collision_distance_wrapper(
torch::Tensor out_distance, torch::Tensor out_vec,
torch::Tensor sparse_index,
const torch::Tensor robot_spheres, // batch_size x n_spheres x 4
const torch::Tensor collision_offset, // n_spheres
const torch::Tensor weight,
const torch::Tensor collision_matrix, // n_spheres
const torch::Tensor thread_locations, const int thread_locations_size,
const int batch_size, const int nspheres, const bool compute_grad = false,
const int ndpt = 8, const bool debug = false) {
std::vector<torch::Tensor>self_collision_distance_wrapper(
torch::Tensor out_distance, torch::Tensor out_vec,
torch::Tensor sparse_index,
const torch::Tensor robot_spheres, // batch_size x n_spheres x 4
const torch::Tensor collision_offset, // n_spheres
const torch::Tensor weight,
const torch::Tensor collision_matrix, // n_spheres
const torch::Tensor thread_locations, const int thread_locations_size,
const int batch_size, const int nspheres, const bool compute_grad = false,
const int ndpt = 8, const bool debug = false)
{
CHECK_INPUT(out_distance);
CHECK_INPUT(out_vec);
CHECK_INPUT(robot_spheres);
@@ -123,29 +152,30 @@ std::vector<torch::Tensor> self_collision_distance_wrapper(
const at::cuda::OptionalCUDAGuard guard(robot_spheres.device());
return self_collision_distance(
out_distance, out_vec, sparse_index, robot_spheres,
collision_offset, weight, collision_matrix, thread_locations,
thread_locations_size, batch_size, nspheres, compute_grad, ndpt, debug);
out_distance, out_vec, sparse_index, robot_spheres,
collision_offset, weight, collision_matrix, thread_locations,
thread_locations_size, batch_size, nspheres, compute_grad, ndpt, debug);
}
std::vector<torch::Tensor> sphere_obb_clpt_wrapper(
const torch::Tensor sphere_position, // batch_size, 4
torch::Tensor distance,
torch::Tensor closest_point, // batch size, 3
torch::Tensor sparsity_idx, const torch::Tensor weight,
const torch::Tensor activation_distance,
const torch::Tensor obb_accel, // n_boxes, 4, 4
const torch::Tensor obb_bounds, // n_boxes, 3
const torch::Tensor obb_pose, // 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 env_query_idx, // n_boxes, 4, 4
const int max_nobs, const int batch_size, const int horizon,
const int n_spheres, const bool transform_back, const bool compute_distance,
const bool use_batch_env) {
std::vector<torch::Tensor>sphere_obb_clpt_wrapper(
const torch::Tensor sphere_position, // batch_size, 4
torch::Tensor distance,
torch::Tensor closest_point, // batch size, 3
torch::Tensor sparsity_idx, const torch::Tensor weight,
const torch::Tensor activation_distance,
const torch::Tensor obb_accel, // n_boxes, 4, 4
const torch::Tensor obb_bounds, // n_boxes, 3
const torch::Tensor obb_pose, // 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 env_query_idx, // n_boxes, 4, 4
const int max_nobs, const int batch_size, const int horizon,
const int n_spheres, const bool transform_back, const bool compute_distance,
const bool use_batch_env)
{
const at::cuda::OptionalCUDAGuard guard(sphere_position.device());
CHECK_INPUT(distance);
CHECK_INPUT(closest_point);
CHECK_INPUT(sphere_position);
@@ -154,55 +184,61 @@ std::vector<torch::Tensor> sphere_obb_clpt_wrapper(
CHECK_INPUT(activation_distance);
CHECK_INPUT(obb_accel);
return sphere_obb_clpt(
sphere_position, distance, closest_point, sparsity_idx, weight,
activation_distance, obb_accel, obb_bounds, obb_pose, obb_enable,
n_env_obb, env_query_idx, max_nobs, batch_size, horizon, n_spheres,
transform_back, compute_distance, use_batch_env);
sphere_position, distance, closest_point, sparsity_idx, weight,
activation_distance, obb_accel, obb_bounds, obb_pose, obb_enable,
n_env_obb, env_query_idx, max_nobs, batch_size, horizon, n_spheres,
transform_back, compute_distance, use_batch_env);
}
std::vector<torch::Tensor> swept_sphere_obb_clpt_wrapper(
const torch::Tensor sphere_position, // batch_size, 4
torch::Tensor distance,
torch::Tensor closest_point, // batch size, 3
torch::Tensor sparsity_idx, 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_bounds, // n_boxes, 3
const torch::Tensor obb_pose, // 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 env_query_idx, // n_boxes, 4, 4
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 bool transform_back, const bool compute_distance,
const bool use_batch_env) {
std::vector<torch::Tensor>swept_sphere_obb_clpt_wrapper(
const torch::Tensor sphere_position, // batch_size, 4
torch::Tensor distance,
torch::Tensor closest_point, // batch size, 3
torch::Tensor sparsity_idx, 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_bounds, // n_boxes, 3
const torch::Tensor obb_pose, // 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 env_query_idx, // n_boxes, 4, 4
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 bool transform_back, const bool compute_distance,
const bool use_batch_env)
{
const at::cuda::OptionalCUDAGuard guard(sphere_position.device());
CHECK_INPUT(distance);
CHECK_INPUT(closest_point);
CHECK_INPUT(sphere_position);
return swept_sphere_obb_clpt(
sphere_position,
distance, closest_point, sparsity_idx, weight, activation_distance,
speed_dt, obb_accel, obb_bounds, obb_pose, obb_enable, n_env_obb,
env_query_idx, max_nobs, batch_size, horizon, n_spheres, sweep_steps,
enable_speed_metric, transform_back, compute_distance, use_batch_env);
sphere_position,
distance, closest_point, sparsity_idx, weight, activation_distance,
speed_dt, obb_accel, obb_bounds, obb_pose, obb_enable, n_env_obb,
env_query_idx, max_nobs, batch_size, horizon, n_spheres, sweep_steps,
enable_speed_metric, transform_back, compute_distance, use_batch_env);
}
std::vector<torch::Tensor> pose_distance_wrapper(
torch::Tensor out_distance, torch::Tensor out_position_distance,
torch::Tensor out_rotation_distance,
torch::Tensor distance_p_vector, // batch size, 3
torch::Tensor distance_q_vector, // batch size, 4
torch::Tensor out_gidx,
const torch::Tensor current_position, // batch_size, 3
const torch::Tensor goal_position, // n_boxes, 3
const torch::Tensor current_quat, const torch::Tensor goal_quat,
const torch::Tensor vec_weight, // n_boxes, 4, 4
const torch::Tensor weight, const torch::Tensor vec_convergence,
const torch::Tensor run_weight, const torch::Tensor run_vec_weight,
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 = false, const bool use_metric = false) {
std::vector<torch::Tensor>pose_distance_wrapper(
torch::Tensor out_distance, torch::Tensor out_position_distance,
torch::Tensor out_rotation_distance,
torch::Tensor distance_p_vector, // batch size, 3
torch::Tensor distance_q_vector, // batch size, 4
torch::Tensor out_gidx,
const torch::Tensor current_position, // batch_size, 3
const torch::Tensor goal_position, // n_boxes, 3
const torch::Tensor current_quat, const torch::Tensor goal_quat,
const torch::Tensor vec_weight, // n_boxes, 4, 4
const torch::Tensor weight, const torch::Tensor vec_convergence,
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 int mode, const int num_goals = 1, const bool compute_grad = false,
const bool write_distance = false, const bool use_metric = false,
const bool project_distance = true)
{
// at::cuda::DeviceGuard guard(angle.device());
CHECK_INPUT(out_distance);
CHECK_INPUT(out_position_distance);
@@ -214,24 +250,30 @@ std::vector<torch::Tensor> pose_distance_wrapper(
CHECK_INPUT(current_quat);
CHECK_INPUT(goal_quat);
CHECK_INPUT(batch_pose_idx);
CHECK_INPUT(offset_waypoint);
CHECK_INPUT(offset_tstep_fraction);
const at::cuda::OptionalCUDAGuard guard(current_position.device());
return pose_distance(
out_distance, out_position_distance, out_rotation_distance,
distance_p_vector, distance_q_vector, out_gidx, current_position,
goal_position, current_quat, goal_quat, vec_weight, weight,
vec_convergence, run_weight, run_vec_weight, batch_pose_idx, batch_size,
horizon, mode, num_goals, compute_grad, write_distance, use_metric);
out_distance, out_position_distance, out_rotation_distance,
distance_p_vector, distance_q_vector, out_gidx, 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, batch_size,
horizon, mode, num_goals, compute_grad, write_distance, use_metric, project_distance);
}
std::vector<torch::Tensor> backward_pose_distance_wrapper(
torch::Tensor out_grad_p, torch::Tensor out_grad_q,
const torch::Tensor grad_distance, // batch_size, 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_p_vec, // n_boxes, 4, 4
const torch::Tensor grad_q_vec, const int batch_size,
const bool use_distance) {
std::vector<torch::Tensor>backward_pose_distance_wrapper(
torch::Tensor out_grad_p, torch::Tensor out_grad_q,
const torch::Tensor grad_distance, // batch_size, 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_p_vec, // n_boxes, 4, 4
const torch::Tensor grad_q_vec, const int batch_size,
const bool use_distance)
{
CHECK_INPUT(out_grad_p);
CHECK_INPUT(out_grad_q);
CHECK_INPUT(grad_distance);
@@ -241,18 +283,19 @@ std::vector<torch::Tensor> backward_pose_distance_wrapper(
const at::cuda::OptionalCUDAGuard guard(grad_distance.device());
return backward_pose_distance(
out_grad_p, out_grad_q, grad_distance, grad_p_distance, grad_q_distance,
pose_weight, grad_p_vec, grad_q_vec, batch_size, use_distance);
out_grad_p, out_grad_q, grad_distance, grad_p_distance, grad_q_distance,
pose_weight, grad_p_vec, grad_q_vec, batch_size, use_distance);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("pose_distance", &pose_distance_wrapper, "Pose Distance (curobolib)");
m.def("pose_distance_backward", &backward_pose_distance_wrapper,
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
m.def("pose_distance", &pose_distance_wrapper, "Pose Distance (curobolib)");
m.def("pose_distance_backward", &backward_pose_distance_wrapper,
"Pose Distance Backward (curobolib)");
m.def("closest_point", &sphere_obb_clpt_wrapper,
m.def("closest_point", &sphere_obb_clpt_wrapper,
"Closest Point OBB(curobolib)");
m.def("swept_closest_point", &swept_sphere_obb_clpt_wrapper,
m.def("swept_closest_point", &swept_sphere_obb_clpt_wrapper,
"Swept Closest Point OBB(curobolib)");
m.def("self_collision_distance", &self_collision_distance_wrapper,