From e3122015503084e949b066a5f66cedbe8370d240 Mon Sep 17 00:00:00 2001
From: Henning Kayser <henningkayser@picknik.ai>
Date: Wed, 31 Jan 2024 15:41:18 +0000
Subject: [PATCH] Implements a collision clearance cost based on distance
 checks

---
 .../include/stomp_moveit/cost_functions.hpp   | 43 +++++++++++++++++++
 .../stomp_moveit_planning_context.hpp         |  2 +
 moveit_planners/stomp/res/stomp_moveit.yaml   |  5 +++
 .../src/stomp_moveit_planning_context.cpp     | 18 +++++---
 4 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/moveit_planners/stomp/include/stomp_moveit/cost_functions.hpp b/moveit_planners/stomp/include/stomp_moveit/cost_functions.hpp
index 705f796cba..cafbc62002 100644
--- a/moveit_planners/stomp/include/stomp_moveit/cost_functions.hpp
+++ b/moveit_planners/stomp/include/stomp_moveit/cost_functions.hpp
@@ -215,6 +215,49 @@ CostFn getCollisionCostFunction(const std::shared_ptr<const planning_scene::Plan
   return getCostFunctionFromStateValidator(collision_validator_fn, COL_CHECK_DISTANCE);
 }
 
+/**
+ * Creates a cost function for distances to collisions of group states in the planning scene.
+ *
+ * @param planning_scene    The planning scene instance to use for distance checking
+ * @param group             The group to use for computing link transforms from joint positions
+ *
+ * @return                  Cost function that computes smooth costs for close-to-collision path segments
+ */
+CostFn getCollisionDistanceCostFunction(const std::shared_ptr<const planning_scene::PlanningScene>& planning_scene,
+                                        const moveit::core::JointModelGroup* group, double min_distance = 0.03)
+{
+  const auto& joints = group ? group->getActiveJointModels() : planning_scene->getRobotModel()->getActiveJointModels();
+
+  StateValidatorFn distance_check_fn = [=](const Eigen::VectorXd& positions) {
+    static moveit::core::RobotState state(planning_scene->getCurrentState());
+
+    // Update robot state values
+    setJointPositions(positions, joints, state);
+    state.update();
+
+    // NOTE: if we need support for filtering links or attached objects, we can do this using DistanceRequests
+    return min_distance - planning_scene->distanceToCollisionUnpadded(state);
+  };
+
+  CostFn cost_fn = [=](const Eigen::MatrixXd& values, Eigen::VectorXd& costs, bool& validity) {
+    // Assume zero cost and valid trajectory from the start
+    costs.setZero(values.cols());
+    validity = true;
+
+    // Iterate over sample waypoints and check for collision distances
+    for (int timestep = 0; timestep < values.cols(); ++timestep)
+    {
+      Eigen::VectorXd current = values.col(timestep);
+      costs(timestep) += distance_check_fn(current);
+      validity = validity && costs(timestep) <= 0.0;
+    }
+
+    return true;
+  };
+
+  return cost_fn;
+}
+
 /**
  * Creates a cost function for binary constraint checks applied to group states.
  * This function uses a StateValidatorFn for computing smooth penalty costs from binary
diff --git a/moveit_planners/stomp/include/stomp_moveit/stomp_moveit_planning_context.hpp b/moveit_planners/stomp/include/stomp_moveit/stomp_moveit_planning_context.hpp
index e7a33c5aec..0417b5958d 100644
--- a/moveit_planners/stomp/include/stomp_moveit/stomp_moveit_planning_context.hpp
+++ b/moveit_planners/stomp/include/stomp_moveit/stomp_moveit_planning_context.hpp
@@ -67,6 +67,8 @@ class StompPlanningContext : public planning_interface::PlanningContext
   void setPathPublisher(const std::shared_ptr<rclcpp::Publisher<visualization_msgs::msg::MarkerArray>>& path_publisher);
   std::shared_ptr<rclcpp::Publisher<visualization_msgs::msg::MarkerArray>> getPathPublisher();
 
+  const stomp_moveit::Params& getParams() const;
+
 private:
   const stomp_moveit::Params params_;
   std::shared_ptr<stomp::Stomp> stomp_;
diff --git a/moveit_planners/stomp/res/stomp_moveit.yaml b/moveit_planners/stomp/res/stomp_moveit.yaml
index 30f382b845..3dbea7abe8 100644
--- a/moveit_planners/stomp/res/stomp_moveit.yaml
+++ b/moveit_planners/stomp/res/stomp_moveit.yaml
@@ -52,6 +52,11 @@ stomp_moveit:
     description: "Weight of the acceleration costs to be applied for the general cost calculation. A value of 0 disables this cost.",
     default_value: 0.1,
   }
+  maximize_clearance: {
+    type: bool,
+    description: "Experimental: Enables a cost function that attempts to maximize the distance to collision.",
+    default_value: false,
+  }
   delta_t: {
     type: double,
     description: "Assumed time change between consecutive points - used for computing control costs",
diff --git a/moveit_planners/stomp/src/stomp_moveit_planning_context.cpp b/moveit_planners/stomp/src/stomp_moveit_planning_context.cpp
index f5895ed0d3..0ce125b7ab 100644
--- a/moveit_planners/stomp/src/stomp_moveit_planning_context.cpp
+++ b/moveit_planners/stomp/src/stomp_moveit_planning_context.cpp
@@ -159,21 +159,22 @@ stomp::TaskPtr createStompTask(const stomp::StompConfiguration& config, StompPla
   // Cost, noise and filter functions are provided for planning.
   // TODO(henningkayser): parameterize cost penalties
   using namespace stomp_moveit;
-  CostFn cost_fn;
+  std::vector<CostFn> cost_functions = { costs::getCollisionCostFunction(planning_scene, group,
+                                                                         1.0 /* collision penalty */) };
   if (!constraints.empty())
   {
-    cost_fn = costs::sum({ costs::getCollisionCostFunction(planning_scene, group, 1.0 /* collision penalty */),
-                           costs::getConstraintsCostFunction(planning_scene, group, constraints.getAllConstraints(),
-                                                             1.0 /* constraint penalty */) });
+    cost_functions.push_back(costs::getConstraintsCostFunction(planning_scene, group, constraints.getAllConstraints(),
+                                                               1.0 /* constraint penalty */));
   }
-  else
+  if (context.getParams().maximize_clearance)
   {
-    cost_fn = costs::getCollisionCostFunction(planning_scene, group, 1.0 /* collision penalty */);
+    cost_functions.push_back(costs::getCollisionDistanceCostFunction(planning_scene, group));
   }
 
   // TODO(henningkayser): parameterize stddev
   const std::vector<double> stddev(group->getActiveJointModels().size(), 0.1);
   auto noise_generator_fn = noise::getNormalDistributionGenerator(num_timesteps, stddev);
+  auto cost_fn = costs::sum(cost_functions);
   auto filter_fn =
       filters::chain({ filters::simpleSmoothingMatrix(num_timesteps), filters::enforcePositionBounds(group) });
   auto iteration_callback_fn =
@@ -301,6 +302,11 @@ void StompPlanningContext::clear()
 {
 }
 
+const stomp_moveit::Params& StompPlanningContext::getParams() const
+{
+  return params_;
+}
+
 void StompPlanningContext::setPathPublisher(
     const std::shared_ptr<rclcpp::Publisher<visualization_msgs::msg::MarkerArray>>& path_publisher)
 {