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

Add planning to grasp API

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2024-11-22 14:15:18 -08:00
parent 18e9ebd35f
commit 36ea382dab
38 changed files with 939 additions and 535 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -146,7 +146,6 @@ namespace Curobo
}
}
template<bool project_distance>
__device__ __forceinline__ void
compute_pose_distance_vector(float *result_vec,
const float3 goal_position,
@@ -156,7 +155,8 @@ namespace Curobo
const float *vec_weight,
const float3 offset_position,
const float3 offset_rotation,
const bool reach_offset)
const bool reach_offset,
const bool project_distance)
{
// project current position to goal frame:
float3 error_position = make_float3(0, 0, 0);
@@ -253,7 +253,7 @@ namespace Curobo
}
}
template<bool project_distance, bool use_metric>
template<bool use_metric>
__device__ __forceinline__ void
compute_pose_distance(float *distance_vec, float& distance, float& position_distance,
float& rotation_distance, const float3 current_position,
@@ -265,17 +265,19 @@ namespace Curobo
const float r_alpha,
const float3 offset_position,
const float3 offset_rotation,
const bool reach_offset)
const bool reach_offset,
const bool project_distance)
{
compute_pose_distance_vector<project_distance>(&distance_vec[0],
goal_position,
goal_quat,
current_position,
current_quat,
&vec_weight[0],
offset_position,
offset_rotation,
reach_offset);
compute_pose_distance_vector(&distance_vec[0],
goal_position,
goal_quat,
current_position,
current_quat,
&vec_weight[0],
offset_position,
offset_rotation,
reach_offset,
project_distance);
position_distance = 0;
rotation_distance = 0;
@@ -394,7 +396,7 @@ namespace Curobo
*(float3 *)&out_grad_q[batch_idx * 4 + 1] = g_q;
}
template<typename scalar_t, bool write_distance, bool project_distance, bool use_metric>
template<typename scalar_t, bool write_distance, bool use_metric>
__global__ void goalset_pose_distance_kernel(
scalar_t *out_distance, scalar_t *out_position_distance,
scalar_t *out_rotation_distance, scalar_t *out_p_vec, scalar_t *out_q_vec,
@@ -405,7 +407,9 @@ namespace Curobo
const scalar_t *run_weight, const scalar_t *run_vec_weight,
const scalar_t *offset_waypoint,
const scalar_t *offset_tstep_fraction,
const int32_t *batch_pose_idx, const int mode, const int num_goals,
const int32_t *batch_pose_idx,
const uint8_t *project_distance_tensor,
const int mode, const int num_goals,
const int batch_size, const int horizon, const bool write_grad = false)
{
const int t_idx = (blockDim.x * blockIdx.x + threadIdx.x);
@@ -416,7 +420,7 @@ namespace Curobo
{
return;
}
const bool project_distance = project_distance_tensor[0];
// read current pose:
float3 position =
*(float3 *)&current_position[batch_idx * horizon * 3 + h_idx * 3];
@@ -511,7 +515,7 @@ namespace Curobo
float4 gq4 = *(float4 *)&goal_quat[(offset + k) * 4];
l_goal_quat = make_float4(gq4.y, gq4.z, gq4.w, gq4.x);
compute_pose_distance<project_distance, use_metric>(&distance_vec[0],
compute_pose_distance<use_metric>(&distance_vec[0],
pose_distance,
position_distance,
rotation_distance,
@@ -531,7 +535,8 @@ namespace Curobo
r_w_alpha,
offset_position,
offset_rotation,
reach_offset);
reach_offset,
project_distance);
if (pose_distance <= best_distance)
{
@@ -657,10 +662,10 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
const torch::Tensor offset_waypoint,
const torch::Tensor offset_tstep_fraction,
const torch::Tensor batch_pose_idx, // batch_size, 1
const torch::Tensor project_distance,
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)
const bool write_distance = true, const bool use_metric = false)
{
using namespace Curobo::Pose;
@@ -684,8 +689,7 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (project_distance)
{
if (use_metric)
{
if (write_distance)
@@ -693,7 +697,7 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
AT_DISPATCH_FLOATING_TYPES(
current_position.scalar_type(), "batch_pose_distance", ([&] {
goalset_pose_distance_kernel
<scalar_t, true, true, true><< < blocksPerGrid, threadsPerBlock, 0,
<scalar_t, true, true><< < blocksPerGrid, threadsPerBlock, 0,
stream >> > (
out_distance.data_ptr<scalar_t>(),
out_position_distance.data_ptr<scalar_t>(),
@@ -711,7 +715,9 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
run_vec_weight.data_ptr<scalar_t>(),
offset_waypoint.data_ptr<scalar_t>(),
offset_tstep_fraction.data_ptr<scalar_t>(),
batch_pose_idx.data_ptr<int32_t>(), mode, num_goals,
batch_pose_idx.data_ptr<int32_t>(),
project_distance.data_ptr<uint8_t>(),
mode, num_goals,
batch_size, horizon, compute_grad);
}));
}
@@ -720,7 +726,7 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
AT_DISPATCH_FLOATING_TYPES(
current_position.scalar_type(), "batch_pose_distance", ([&] {
goalset_pose_distance_kernel
<scalar_t, false, true, true><< < blocksPerGrid, threadsPerBlock, 0,
<scalar_t, false, true><< < blocksPerGrid, threadsPerBlock, 0,
stream >> > (
out_distance.data_ptr<scalar_t>(),
out_position_distance.data_ptr<scalar_t>(),
@@ -738,7 +744,9 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
run_vec_weight.data_ptr<scalar_t>(),
offset_waypoint.data_ptr<scalar_t>(),
offset_tstep_fraction.data_ptr<scalar_t>(),
batch_pose_idx.data_ptr<int32_t>(), mode, num_goals,
batch_pose_idx.data_ptr<int32_t>(),
project_distance.data_ptr<uint8_t>(),
mode, num_goals,
batch_size, horizon, compute_grad);
}));
}
@@ -749,7 +757,7 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
{
AT_DISPATCH_FLOATING_TYPES(
current_position.scalar_type(), "batch_pose_distance", ([&] {
goalset_pose_distance_kernel<scalar_t, true, true, false>
goalset_pose_distance_kernel<scalar_t, true, false>
<< < blocksPerGrid, threadsPerBlock, 0, stream >> > (
out_distance.data_ptr<scalar_t>(),
out_position_distance.data_ptr<scalar_t>(),
@@ -767,7 +775,9 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
run_vec_weight.data_ptr<scalar_t>(),
offset_waypoint.data_ptr<scalar_t>(),
offset_tstep_fraction.data_ptr<scalar_t>(),
batch_pose_idx.data_ptr<int32_t>(), mode, num_goals,
batch_pose_idx.data_ptr<int32_t>(),
project_distance.data_ptr<uint8_t>(),
mode, num_goals,
batch_size, horizon, compute_grad);
}));
}
@@ -775,7 +785,7 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
{
AT_DISPATCH_FLOATING_TYPES(
current_position.scalar_type(), "batch_pose_distance", ([&] {
goalset_pose_distance_kernel<scalar_t, false, true, false>
goalset_pose_distance_kernel<scalar_t, false, false>
<< < blocksPerGrid, threadsPerBlock, 0, stream >> > (
out_distance.data_ptr<scalar_t>(),
out_position_distance.data_ptr<scalar_t>(),
@@ -793,127 +803,15 @@ pose_distance(torch::Tensor out_distance, torch::Tensor out_position_distance,
run_vec_weight.data_ptr<scalar_t>(),
offset_waypoint.data_ptr<scalar_t>(),
offset_tstep_fraction.data_ptr<scalar_t>(),
batch_pose_idx.data_ptr<int32_t>(), mode, num_goals,
batch_pose_idx.data_ptr<int32_t>(),
project_distance.data_ptr<uint8_t>(),
mode, num_goals,
batch_size, horizon, compute_grad);
}));
}
}
}
else
{
if (use_metric)
{
if (write_distance)
{
AT_DISPATCH_FLOATING_TYPES(
current_position.scalar_type(), "batch_pose_distance", ([&] {
goalset_pose_distance_kernel
<scalar_t, true, false, true><< < blocksPerGrid, threadsPerBlock, 0,
stream >> > (
out_distance.data_ptr<scalar_t>(),
out_position_distance.data_ptr<scalar_t>(),
out_rotation_distance.data_ptr<scalar_t>(),
distance_p_vector.data_ptr<scalar_t>(),
distance_q_vector.data_ptr<scalar_t>(),
out_gidx.data_ptr<int32_t>(),
current_position.data_ptr<scalar_t>(),
goal_position.data_ptr<scalar_t>(),
current_quat.data_ptr<scalar_t>(),
goal_quat.data_ptr<scalar_t>(),
vec_weight.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
vec_convergence.data_ptr<scalar_t>(),
run_weight.data_ptr<scalar_t>(),
run_vec_weight.data_ptr<scalar_t>(),
offset_waypoint.data_ptr<scalar_t>(),
offset_tstep_fraction.data_ptr<scalar_t>(),
batch_pose_idx.data_ptr<int32_t>(), mode, num_goals,
batch_size, horizon, compute_grad);
}));
}
else
{
AT_DISPATCH_FLOATING_TYPES(
current_position.scalar_type(), "batch_pose_distance", ([&] {
goalset_pose_distance_kernel
<scalar_t, false, false, true><< < blocksPerGrid, threadsPerBlock, 0,
stream >> > (
out_distance.data_ptr<scalar_t>(),
out_position_distance.data_ptr<scalar_t>(),
out_rotation_distance.data_ptr<scalar_t>(),
distance_p_vector.data_ptr<scalar_t>(),
distance_q_vector.data_ptr<scalar_t>(),
out_gidx.data_ptr<int32_t>(),
current_position.data_ptr<scalar_t>(),
goal_position.data_ptr<scalar_t>(),
current_quat.data_ptr<scalar_t>(),
goal_quat.data_ptr<scalar_t>(),
vec_weight.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
vec_convergence.data_ptr<scalar_t>(),
run_weight.data_ptr<scalar_t>(),
run_vec_weight.data_ptr<scalar_t>(),
offset_waypoint.data_ptr<scalar_t>(),
offset_tstep_fraction.data_ptr<scalar_t>(),
batch_pose_idx.data_ptr<int32_t>(), mode, num_goals,
batch_size, horizon, compute_grad);
}));
}
}
else
{
if (write_distance)
{
AT_DISPATCH_FLOATING_TYPES(
current_position.scalar_type(), "batch_pose_distance", ([&] {
goalset_pose_distance_kernel<scalar_t, true, false, false>
<< < blocksPerGrid, threadsPerBlock, 0, stream >> > (
out_distance.data_ptr<scalar_t>(),
out_position_distance.data_ptr<scalar_t>(),
out_rotation_distance.data_ptr<scalar_t>(),
distance_p_vector.data_ptr<scalar_t>(),
distance_q_vector.data_ptr<scalar_t>(),
out_gidx.data_ptr<int32_t>(),
current_position.data_ptr<scalar_t>(),
goal_position.data_ptr<scalar_t>(),
current_quat.data_ptr<scalar_t>(),
goal_quat.data_ptr<scalar_t>(),
vec_weight.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
vec_convergence.data_ptr<scalar_t>(),
run_weight.data_ptr<scalar_t>(),
run_vec_weight.data_ptr<scalar_t>(),
offset_waypoint.data_ptr<scalar_t>(),
offset_tstep_fraction.data_ptr<scalar_t>(),
batch_pose_idx.data_ptr<int32_t>(), mode, num_goals,
batch_size, horizon, compute_grad);
}));
}
else
{
AT_DISPATCH_FLOATING_TYPES(
current_position.scalar_type(), "batch_pose_distance", ([&] {
goalset_pose_distance_kernel<scalar_t, false, false, false>
<< < blocksPerGrid, threadsPerBlock, 0, stream >> > (
out_distance.data_ptr<scalar_t>(),
out_position_distance.data_ptr<scalar_t>(),
out_rotation_distance.data_ptr<scalar_t>(),
distance_p_vector.data_ptr<scalar_t>(),
distance_q_vector.data_ptr<scalar_t>(),
out_gidx.data_ptr<int32_t>(),
current_position.data_ptr<scalar_t>(),
goal_position.data_ptr<scalar_t>(),
current_quat.data_ptr<scalar_t>(),
goal_quat.data_ptr<scalar_t>(),
vec_weight.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(),
vec_convergence.data_ptr<scalar_t>(),
run_weight.data_ptr<scalar_t>(),
run_vec_weight.data_ptr<scalar_t>(),
offset_waypoint.data_ptr<scalar_t>(),
offset_tstep_fraction.data_ptr<scalar_t>(),
batch_pose_idx.data_ptr<int32_t>(), mode, num_goals,
batch_size, horizon, compute_grad);
}));
}
}
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
return { out_distance, out_position_distance, out_rotation_distance,