[wip] feat: support custom projection

action.sh

@@ -1,4 +1,4 @@
-frag="debug"
+frag="release"
 if [ $frag = "debug" ]; then
     if [ ! -d "build_debug" ]; then
         mkdir "build_debug"
@@ -12,3 +12,4 @@ elif [ $frag = "release" ]; then
 fi
 cmake -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DCMAKE_BUILD_TYPE=$frag ..
 make -j4
+ln -sf $PWD/_omplpy.cpython-38-x86_64-linux-gnu.so ../python/ompl/

python/ompl/__init__.py

@@ -9,10 +9,12 @@
 
 VectorLike = Union[np.ndarray, Sequence[float]]
 IsValidFunc = Callable[[List[float]], bool]
+ProjectFunc = Callable[[], np.ndarray]
 
 
 class Algorithm(Enum):
     BKPIECE1 = "BKPIECE1"
+    IKPIECE1 = "IKPIECE1"
     KPIECE1 = "KPIECE1"
     LBKPIECE1 = "LBKPIECE1"
     RRTConnect = "RRTConnect"
@@ -256,7 +258,7 @@ def __init__(
 class ConstrainedPlanner(_OMPLPlannerBase):
     def __init__(
         self,
-        eq_const: Callable[[np.ndarray], Tuple[np.ndarray, np.ndarray]],
+        eq_const: Optional[Callable[[np.ndarray], Tuple[np.ndarray, np.ndarray]]],
         lb: VectorLike,
         ub: VectorLike,
         is_valid: IsValidFunc,
@@ -265,6 +267,9 @@ def __init__(
         algo: Algorithm = Algorithm.RRTConnect,
         algo_range: Optional[float] = None,
         cs_type: ConstStateType = ConstStateType.PROJECTION,
+        f_and_project: Optional[
+            Tuple[Callable[[np.ndarray], np.ndarray], ProjectFunc]
+        ] = None,
     ):
 
         lb = np.array(lb)
@@ -276,26 +281,35 @@ def __init__(
             dim = len(lb)
             validation_box = np.array([validation_box] * dim)
 
-        class ConstraintFunction:
-            jac_cache: Optional[np.ndarray] = None
+        if f_and_project is not None:
+            f, f_project = f_and_project
+            f_jac = None
+        else:
 
-            def __call__(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
-                return eq_const(x)
+            class ConstraintFunction:
+                jac_cache: Optional[np.ndarray] = None
 
-            def f(self, x: np.ndarray) -> List[float]:
-                val, jac = self.__call__(x)
-                self.jac_cache = jac
-                return val.tolist()
+                def __call__(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+                    return eq_const(x)
 
-            def jac(self, x: np.ndarray) -> List[List[float]]:
-                assert self.jac_cache is not None
-                return self.jac_cache.tolist()
+                def f(self, x: np.ndarray) -> List[float]:
+                    val, jac = self.__call__(x)
+                    self.jac_cache = jac
+                    return val.tolist()
 
-        const_fn = ConstraintFunction()
+                def jac(self, x: np.ndarray) -> List[List[float]]:
+                    assert self.jac_cache is not None
+                    return self.jac_cache.tolist()
+
+            const_fn = ConstraintFunction()
+            f = const_fn.f
+            f_jac = const_fn.jac
+            f_project = None
 
         self._planner = _omplpy._ConstrainedPlanner(
-            const_fn.f,
-            const_fn.jac,
+            f,
+            f_jac,
+            f_project,
             lb,
             ub,
             is_valid,

src/higher.hpp

@@ -38,6 +38,7 @@
 #include <ompl/util/PPM.h>
 #include <ompl/util/Time.h>
 
+#include <Eigen/Dense>
 #include <boost/filesystem.hpp>
 #include <cassert>
 #include <cstddef>
@@ -77,6 +78,7 @@ namespace ot = ompl::tools;
 // pybind fail to convert numpy to eigen in callback case
 using ConstFn = std::function<std::vector<double>(std::vector<double>)>;
 using ConstJacFn = std::function<std::vector<std::vector<double>>(std::vector<double>)>;
+using ProjectFn = std::function<bool(Eigen::Ref<Eigen::VectorXd>)>;
 
 template <typename T>
 std::shared_ptr<T> create_algorithm(const ob::SpaceInformationPtr si, std::optional<double> range)
@@ -194,6 +196,8 @@ class BoxMotionValidator : public ob::MotionValidator
   std::vector<double> width_;
 };
 
+std::shared_ptr<ob::Constraint> global_fuck_constraint_;
+
 std::optional<std::vector<double>> split_geodesic_with_box(const ob::State* s1,
                                                            const ob::State* s2,
                                                            const ob::SpaceInformation* si,
@@ -257,6 +261,9 @@ std::optional<std::vector<double>> split_geodesic_with_box(const ob::State* s1,
       if (not_evaluated_yet) {
         const auto s_new = si->allocState();
         space->interpolate(s1, s2, fraction, s_new);
+        if (!global_fuck_constraint_->isSatisfied(s_new)) {
+          return {};
+        }
         if (!si->isValid(s_new)) {
           return {};
         }
@@ -311,8 +318,8 @@ class GeodesicBoxMotionValidator : public ob::MotionValidator
 class ConstraintWrap : public ob::Constraint
 {
  public:
-  ConstraintWrap(const ConstFn& f, const ConstJacFn& jac, size_t dim_ambient, size_t dim_constraint)
-      : ob::Constraint(dim_ambient, dim_constraint), f_(f), jac_(jac)
+  ConstraintWrap(ConstFn f, ConstJacFn jac, size_t dim_ambient, size_t dim_constraint)
+      : ob::Constraint(dim_ambient, dim_constraint), f_(std::move(f)), jac_(std::move(jac))
   {
   }
 
@@ -342,6 +349,41 @@ class ConstraintWrap : public ob::Constraint
   ConstJacFn jac_;
 };
 
+class ConstraintWrapCustomProjection : public ob::Constraint
+{
+ public:
+  ConstraintWrapCustomProjection(ConstFn f,
+                                 ProjectFn f_project,
+                                 size_t dim_ambient,
+                                 size_t dim_constraint)
+      : ob::Constraint(dim_ambient, dim_constraint),
+        f_(std::move(f)),
+        f_project_(std::move(f_project))
+  {
+  }
+
+  void function(const Eigen::Ref<const Eigen::VectorXd>& x,
+                Eigen::Ref<Eigen::VectorXd> out) const override
+  {
+    std::vector<double> xvec(x.data(), x.data() + x.size());
+    const auto yvec = f_(xvec);
+    for (size_t i = 0; i < yvec.size(); ++i) {
+      out[i] = yvec[i];
+    }
+  }
+
+  void jacobian(const Eigen::Ref<const Eigen::VectorXd>& x,
+                Eigen::Ref<Eigen::MatrixXd> out) const override
+  {
+    throw std::runtime_error("should not reach here");
+  }
+
+  bool project(Eigen::Ref<Eigen::VectorXd> x) const override { return this->f_project_(x); }
+
+  ConstFn f_;
+  ProjectFn f_project_;
+};
+
 template <bool Constrained>
 struct CollisionAwareSpaceInformation {
   void resetCount() { this->is_valid_call_count_ = 0; }
@@ -399,8 +441,9 @@ enum ConstStateType { PROJECTION, ATLAS, TANGENT };
 
 struct ConstrainedCollisoinAwareSpaceInformation : public CollisionAwareSpaceInformation<true> {
   ConstrainedCollisoinAwareSpaceInformation(
-      const ConstFn& f_const,
-      const ConstJacFn& jac_const,
+      ConstFn f_const,
+      ConstJacFn jac_const,
+      ProjectFn f_project,
       const std::vector<double>& lb,
       const std::vector<double>& ub,
       const std::function<bool(std::vector<double>)>& is_valid,
@@ -411,8 +454,17 @@ struct ConstrainedCollisoinAwareSpaceInformation : public CollisionAwareSpaceInf
   {
     size_t dim_ambient = lb.size();
     size_t dim_constraint = f_const(lb).size();  // detect by dummy input
-    std::shared_ptr<ob::Constraint> constraint =
-        std::make_shared<ConstraintWrap>(f_const, jac_const, dim_ambient, dim_constraint);
+    std::shared_ptr<ob::Constraint> constraint;
+    if (jac_const == nullptr && f_project != nullptr) {
+      constraint = std::make_shared<ConstraintWrapCustomProjection>(
+          std::move(f_const), std::move(f_project), dim_ambient, dim_constraint);
+    } else if (jac_const != nullptr && f_project == nullptr) {
+      constraint = std::make_shared<ConstraintWrap>(
+          std::move(f_const), std::move(jac_const), dim_ambient, dim_constraint);
+    } else {
+      throw std::runtime_error("either f_project or jac_const should be nullptr");
+    }
+    global_fuck_constraint_ = constraint;
     const auto space = bound2space(lb, ub);
 
     ob::ConstrainedStateSpacePtr css;
@@ -517,6 +569,8 @@ struct PlannerBase {
       return create_algorithm<og::BKPIECE1>(space_info, range);
     } else if (name.compare("KPIECE1") == 0) {
       return create_algorithm<og::KPIECE1>(space_info, range);
+    } else if (name.compare("IKPIECE1") == 0) {
+      return create_algorithm<og::IKPIECE1>(space_info, range);
     } else if (name.compare("RRT") == 0) {
       return create_algorithm<og::RRT>(space_info, range);
     } else if (name.compare("RRTConnect") == 0) {
@@ -535,8 +589,9 @@ struct PlannerBase {
 };
 
 struct ConstrainedPlanner : public PlannerBase<true> {
-  ConstrainedPlanner(const ConstFn& f_const,
-                     const ConstJacFn& jac_const,
+  ConstrainedPlanner(ConstFn f_const,
+                     ConstJacFn jac_const,
+                     ProjectFn f_project,
                      const std::vector<double>& lb,
                      const std::vector<double>& ub,
                      const std::function<bool(std::vector<double>)>& is_valid,
@@ -547,9 +602,26 @@ struct ConstrainedPlanner : public PlannerBase<true> {
                      ConstStateType cs_type = ConstStateType::PROJECTION)
       : PlannerBase<true>{nullptr, nullptr}
   {
-    csi_ = std::make_unique<ConstrainedCollisoinAwareSpaceInformation>(
-        f_const, jac_const, lb, ub, is_valid, max_is_valid_call, box_width, cs_type);
+    csi_ = std::make_unique<ConstrainedCollisoinAwareSpaceInformation>(std::move(f_const),
+                                                                       std::move(jac_const),
+                                                                       f_project,
+                                                                       lb,
+                                                                       ub,
+                                                                       is_valid,
+                                                                       max_is_valid_call,
+                                                                       box_width,
+                                                                       cs_type);
     const auto algo = get_algorithm(algo_name, range);
+    if (algo_name == "IKPIECE1") {
+      /// dyn pointer cast to std::shared_ptr<og::IKPIECE1>
+      auto ikpiece1 = std::dynamic_pointer_cast<og::IKPIECE1>(algo);
+      ikpiece1->set_project_fn(f_project);
+      Eigen::VectorXd motion_step_box(box_width.size());
+      for (size_t i = 0; i < box_width.size(); ++i) {
+        motion_step_box(i) = box_width[i];
+      }
+      ikpiece1->set_motion_step_box(motion_step_box);
+    }
 
     setup_ = std::make_unique<og::SimpleSetup>(csi_->si_);
     setup_->setPlanner(algo);

src/wrapper.cpp

@@ -1,3 +1,4 @@
+#include <pybind11/eigen.h>
 #include <pybind11/functional.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
@@ -18,8 +19,9 @@ PYBIND11_MODULE(_omplpy, m)
       .value("TANGENT", ConstStateType::TANGENT);
 
   py::class_<ConstrainedPlanner>(m, "_ConstrainedPlanner")
-      .def(py::init<const ConstFn&,
-                    const ConstJacFn&,
+      .def(py::init<ConstFn,
+                    ConstJacFn,
+                    ProjectFn,
                     std::vector<double>,
                     std::vector<double>,
                     std::function<bool(std::vector<double>)>,

tests/test_constrained_planner.py

@@ -36,6 +36,12 @@ def eq_const(x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
     return y, jac
 
 
+def project(x: np.ndarray) -> bool:
+    dist = np.linalg.norm(x)
+    x /= dist
+    return True
+
+
 def is_valid(vec: List[float]) -> bool:
     x, y, z = vec
     if abs(x) < 0.2 and abs(z) < 0.8:
@@ -44,27 +50,39 @@ def is_valid(vec: List[float]) -> bool:
 
 
 def test_constrained_planner(visualize: bool = False):
-    planner = ConstrainedPlanner(
+    planner_numelical = ConstrainedPlanner(
         eq_const, [-2, -2, -2], [2, 2, 2], is_valid, 10000, 0.1, algo_range=0.3
     )
+    f_and_project = (lambda x: eq_const(x)[0], project)
+    planner_with_custom_project = ConstrainedPlanner(
+        None,
+        [-2, -2, -2],
+        [2, 2, 2],
+        is_valid,
+        10000,
+        0.1,
+        algo_range=0.3,
+        f_and_project=f_and_project,
+    )
     start = np.array([-1, 0.0, 0.0])
     goal = np.array([1, 0.0, 0.0])
-    traj = planner.solve(start, goal, False)
-
-    assert traj is not None
-    for p in traj:
-        val = eq_const(p)[0]
-        assert np.all(val < 1e-3)
-    assert np.linalg.norm(traj[0] - start) < 1e-3
-    assert np.linalg.norm(traj[-1] - goal) < 1e-3
-
-    if visualize:
-        fig = plt.figure()
-        ax = fig.add_subplot(projection="3d")
-        plt_sphere([[0, 0, 0]], [0.97], fig, ax)
-        X = np.array(traj)
-        ax.scatter(X[:, 0], X[:, 1], X[:, 2])
-        plt.show()
+    for planner in [planner_numelical, planner_with_custom_project]:
+        traj = planner.solve(start, goal, False)
+
+        assert traj is not None
+        for p in traj:
+            val = eq_const(p)[0]
+            assert np.all(val < 1e-3)
+        assert np.linalg.norm(traj[0] - start) < 1e-3
+        assert np.linalg.norm(traj[-1] - goal) < 1e-3
+
+        if visualize:
+            fig = plt.figure()
+            ax = fig.add_subplot(projection="3d")
+            plt_sphere([[0, 0, 0]], [0.97], fig, ax)
+            X = np.array(traj)
+            ax.scatter(X[:, 0], X[:, 1], X[:, 2])
+            plt.show()
 
 
 if __name__ == "__main__":