Growing IRIS-ZO Regions Along a Parametrized Subspace (#22558)

cohnt · web-flow · commit 76b7b12fb39e · 2025-02-12T15:16:11.000-05:00
diff --git a/planning/iris/iris_zo.cc b/planning/iris/iris_zo.cc
@@ -84,7 +84,8 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
                    const HPolyhedron& domain, const IrisZoOptions& options) {
   auto start = std::chrono::high_resolution_clock::now();
   const int num_threads_to_use =
-      checker.SupportsParallelChecking()
+      checker.SupportsParallelChecking() &&
+              options.get_parameterization_is_threadsafe()
           ? std::min(options.parallelism.num_threads(),
                      checker.num_allocated_contexts())
           : 1;
@@ -100,14 +101,30 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
   // Prevent directly terminating if the ellipsoid is too large.
   double previous_volume = 0;
 
-  const int dim = starting_ellipsoid.ambient_dimension();
-  int current_num_faces = domain.A().rows();
+  const int ambient_dimension = checker.plant().num_positions();
+  const int parameterized_dimension =
+      options.get_parameterization_dimension().value_or(ambient_dimension);
 
   DRAKE_THROW_UNLESS(num_threads_to_use > 0);
-  DRAKE_THROW_UNLESS(domain.ambient_dimension() == dim);
+  DRAKE_THROW_UNLESS(starting_ellipsoid.ambient_dimension() ==
+                     parameterized_dimension);
+  DRAKE_THROW_UNLESS(domain.ambient_dimension() == parameterized_dimension);
   DRAKE_THROW_UNLESS(domain.IsBounded());
   DRAKE_THROW_UNLESS(domain.PointInSet(current_ellipsoid_center));
 
+  const int computed_ambient_dimension =
+      options.get_parameterization()(starting_ellipsoid_center).size();
+  if (computed_ambient_dimension != ambient_dimension) {
+    throw std::runtime_error(fmt::format(
+        "The plant has {} positions, but the given parameterization "
+        "returned a point with the wrong dimension (its size was "
+        "{}) when called on {}.",
+        ambient_dimension, computed_ambient_dimension,
+        fmt_eigen(starting_ellipsoid_center.transpose())));
+  }
+
+  int current_num_faces = domain.A().rows();
+
   if (options.max_iterations_separating_planes <= 0) {
     throw std::runtime_error(
         "The maximum number of iterations for separating planes "
@@ -116,7 +133,7 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
   VPolytope cvxh_vpoly;
   if (options.containment_points.has_value()) {
     cvxh_vpoly = VPolytope(options.containment_points.value());
-    DRAKE_THROW_UNLESS(domain.ambient_dimension() ==
+    DRAKE_THROW_UNLESS(parameterized_dimension ==
                        options.containment_points->rows());
 
     constexpr float kPointInSetTol = 1e-5;
@@ -133,7 +150,8 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
 
     for (int col = 0; col < options.containment_points->cols(); ++col) {
       Eigen::VectorXd conf = options.containment_points->col(col);
-      cont_vec.emplace_back(conf);
+      cont_vec.emplace_back(options.get_parameterization()(conf));
+      DRAKE_ASSERT(cont_vec.back().size() == ambient_dimension);
     }
 
     std::vector<uint8_t> containment_point_col_free =
@@ -147,11 +165,14 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
   }
   // For debugging visualization.
   Eigen::Vector3d point_to_draw = Eigen::Vector3d::Zero();
-  if (options.meshcat && dim <= 3) {
+  if (options.meshcat && ambient_dimension <= 3) {
     std::string path = "seedpoint";
     options.meshcat->SetObject(path, Sphere(0.06),
                                geometry::Rgba(0.1, 1, 1, 1.0));
-    point_to_draw.head(dim) = current_ellipsoid_center;
+    Eigen::VectorXd conf_ambient =
+        options.get_parameterization()(current_ellipsoid_center);
+    DRAKE_ASSERT(conf_ambient.size() == ambient_dimension);
+    point_to_draw.head(ambient_dimension) = conf_ambient;
     options.meshcat->SetTransform(path, RigidTransform<double>(point_to_draw));
   }
 
@@ -174,7 +195,8 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
     log()->info("IrisZo worst case test requires {} samples.", N_max);
   }
 
-  std::vector<Eigen::VectorXd> particles(N_max, Eigen::VectorXd::Zero(dim));
+  std::vector<Eigen::VectorXd> particles(
+      N_max, Eigen::VectorXd::Zero(parameterized_dimension));
 
   int iteration = 0;
   HPolyhedron P = domain;
@@ -184,7 +206,7 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
   // TODO(wernerpe): Find a better solution than hardcoding 300.
   constexpr int kNumFacesToPreAllocate = 300;
   Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> A(
-      P.A().rows() + kNumFacesToPreAllocate, dim);
+      P.A().rows() + kNumFacesToPreAllocate, parameterized_dimension);
   Eigen::VectorXd b(P.A().rows() + kNumFacesToPreAllocate);
 
   while (true) {
@@ -228,15 +250,24 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
                                           options.mixing_steps);
       }
 
-      // Copy top slice of particles, due to collision checker only accepting
-      // vectors of configurations.
-      // TODO(wernerpe): Remove this copy operation.
-      std::vector<Eigen::VectorXd> particles_step(particles.begin(),
-                                                  particles.begin() + N_k);
+      // Copy top slice of particles, applying thet parameterization function to
+      // each one, due to collision checker only accepting vectors of
+      // configurations.
+      // TODO(wernerpe, cohnt): Remove this copy operation.
+      // TODO(cohnt): Consider parallelizing the parameterization calls, in case
+      // it's an expensive operation.
+      std::vector<Eigen::VectorXd> ambient_particles(N_k);
+      std::transform(particles.begin(), particles.begin() + N_k,
+                     ambient_particles.begin(), options.get_parameterization());
+
+      for (int i = 0; i < ssize(ambient_particles); ++i) {
+        // Only run this check in debug mode, because it's expensive.
+        DRAKE_ASSERT(ambient_particles[i].size() == ambient_dimension);
+      }
 
       // Find all particles in collision
       std::vector<uint8_t> particle_col_free =
-          checker.CheckConfigsCollisionFree(particles_step,
+          checker.CheckConfigsCollisionFree(ambient_particles,
                                             options.parallelism);
       std::vector<Eigen::VectorXd> particles_in_collision;
       int number_particles_in_collision_unadaptive_test = 0;
@@ -280,34 +311,45 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
 
       // For each particle in collision, we run a bisection search to find a
       // configuration on the boundary of the obstacle.
-      const auto particle_update_work =
-          [&checker, &particles_in_collision_updated, &particles_in_collision,
-           &current_ellipsoid_center,
-           &options](const int thread_num, const int64_t index) {
-            const int point_idx = static_cast<int>(index);
-            auto start_point = particles_in_collision[point_idx];
-
-            Eigen::VectorXd current_point = start_point;
-            Eigen::VectorXd curr_pt_lower = current_ellipsoid_center;
-
-            // Update current point using a fixed number of bisection steps.
-            if (!checker.CheckConfigCollisionFree(curr_pt_lower, thread_num)) {
-              current_point = curr_pt_lower;
+      const auto particle_update_work = [&checker,
+                                         &particles_in_collision_updated,
+                                         &particles_in_collision,
+                                         &current_ellipsoid_center,
+                                         &options](const int thread_num,
+                                                   const int64_t index) {
+        const int point_idx = static_cast<int>(index);
+        auto start_point = particles_in_collision[point_idx];
+
+        Eigen::VectorXd current_point = start_point;
+        Eigen::VectorXd curr_pt_lower = current_ellipsoid_center;
+
+        // Update current point using a fixed number of bisection steps.
+        Eigen::VectorXd current_point_ambient =
+            options.get_parameterization()(curr_pt_lower);
+        DRAKE_ASSERT(current_point_ambient.size() ==
+                     checker.plant().num_positions());
+        if (!checker.CheckConfigCollisionFree(current_point_ambient,
+                                              thread_num)) {
+          current_point = curr_pt_lower;
+        } else {
+          Eigen::VectorXd curr_pt_upper = current_point;
+          for (int i = 0; i < options.bisection_steps; ++i) {
+            Eigen::VectorXd query = 0.5 * (curr_pt_upper + curr_pt_lower);
+            Eigen::VectorXd query_ambient =
+                options.get_parameterization()(query);
+            DRAKE_ASSERT(query_ambient.size() ==
+                         checker.plant().num_positions());
+            if (checker.CheckConfigCollisionFree(query_ambient, thread_num)) {
+              curr_pt_lower = query;
             } else {
-              Eigen::VectorXd curr_pt_upper = current_point;
-              for (int i = 0; i < options.bisection_steps; ++i) {
-                Eigen::VectorXd query = 0.5 * (curr_pt_upper + curr_pt_lower);
-                if (checker.CheckConfigCollisionFree(query, thread_num)) {
-                  curr_pt_lower = query;
-                } else {
-                  curr_pt_upper = query;
-                  current_point = query;
-                }
-              }
+              curr_pt_upper = query;
+              current_point = query;
             }
+          }
+        }
 
-            particles_in_collision_updated[point_idx] = current_point;
-          };
+        particles_in_collision_updated[point_idx] = current_point;
+      };
 
       // Update all particles in parallel.
       DynamicParallelForIndexLoop(DegreeOfParallelism(num_threads_to_use), 0,
@@ -378,15 +420,18 @@ HPolyhedron IrisZo(const planning::CollisionChecker& checker,
           }
 
           // Fill in meshcat if added for debugging.
-          if (options.meshcat && dim <= 3) {
+          if (options.meshcat && ambient_dimension <= 3) {
             for (int pt_to_draw = 0; pt_to_draw < number_particles_in_collision;
                  ++pt_to_draw) {
               std::string path = fmt::format(
                   "face_pt/iteration{:02}/sepit{:02}/{:03}/pt", iteration,
                   num_iterations_separating_planes, current_num_faces);
               options.meshcat->SetObject(path, Sphere(0.03),
                                          geometry::Rgba(1, 1, 0.1, 1.0));
-              point_to_draw.head(dim) = nearest_particle;
+              Eigen::VectorXd ambient_particle =
+                  options.get_parameterization()(nearest_particle);
+              DRAKE_ASSERT(ambient_particle.size() == ambient_dimension);
+              point_to_draw.head(ambient_dimension) = ambient_particle;
               options.meshcat->SetTransform(
                   path, RigidTransform<double>(point_to_draw));
             }
diff --git a/planning/iris/iris_zo.h b/planning/iris/iris_zo.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <filesystem>
+#include <functional>
 #include <memory>
 #include <optional>
 #include <utility>
@@ -22,7 +23,10 @@ namespace planning {
  * @experimental
  * @see IrisZo for more details.
  **/
-struct IrisZoOptions {
+class IrisZoOptions {
+ public:
+  DRAKE_DEFAULT_COPY_AND_MOVE_AND_ASSIGN(IrisZoOptions);
+
   /** Passes this object to an Archive.
   Refer to @ref yaml_serialization "YAML Serialization" for background.
   Note: This only serializes options that are YAML built-in types. */
@@ -92,7 +96,8 @@ struct IrisZoOptions {
 
   /** Number of threads to use when updating the particles. If the user requests
    * more threads than the CollisionChecker supports, that number of threads
-   * will be used instead. */
+   * will be used instead. However, see also `parameterization_is_threadsafe`.
+   */
   Parallelism parallelism{Parallelism::Max()};
 
   /** Enables print statements indicating the progress of IrisZo. */
@@ -128,6 +133,62 @@ struct IrisZoOptions {
   currently and when the
   configuration space is <= 3 dimensional.*/
   std::shared_ptr<geometry::Meshcat> meshcat{};
+
+  typedef std::function<Eigen::VectorXd(const Eigen::VectorXd&)>
+      ParameterizationFunction;
+
+  /** Ordinarily, IRIS-ZO grows collision free regions in the robot's
+   * configuration space C. This allows the user to specify a function f:Q→C ,
+   * and grow the region in Q instead. The function should be a map R^m to
+   * R^n, where n is the dimension of the plant configuration space, determined
+   * via `checker.plant().num_positions()` and m is `parameterization_dimension`
+   * if specified. The user must provide `parameterization`, which is the
+   * function f, `parameterization_is_threadsafe`, which is whether or not
+   * `parametrization` can be called concurrently, and
+   * `parameterization_dimension`, the dimension of the input space Q. */
+  void set_parameterization(const ParameterizationFunction& parameterization,
+                            bool parameterization_is_threadsafe,
+                            int parameterization_dimension) {
+    parameterization_ = parameterization;
+    parameterization_is_threadsafe_ = parameterization_is_threadsafe;
+    parameterization_dimension_ = parameterization_dimension;
+  }
+
+  /** Get the parameterization function.
+   * @note If the user has not specified this with `set_parameterization()`,
+   * then the default value of `parameterization_` is the identity function,
+   * indicating that the regions should be grown in the full configuration space
+   * (in the standard coordinate system). */
+  const ParameterizationFunction& get_parameterization() const {
+    return parameterization_;
+  }
+
+  /** Returns whether or not the user has specified the parameterization to be
+   * threadsafe.
+   * @note The default `parameterization_` is the identity function, which is
+   * threadsafe. */
+  bool get_parameterization_is_threadsafe() const {
+    return parameterization_is_threadsafe_;
+  }
+
+  /** Returns what the user has specified as the input dimension for the
+   * parameterization function, or std::nullopt if it has not been set. A
+   * std::nullopt value indicates that
+   * IrisZo should use the ambient configuration space dimension as the input
+   * dimension to the parameterization. */
+  std::optional<int> get_parameterization_dimension() const {
+    return parameterization_dimension_;
+  }
+
+ private:
+  bool parameterization_is_threadsafe_{true};
+
+  std::optional<int> parameterization_dimension_{std::nullopt};
+
+  std::function<Eigen::VectorXd(const Eigen::VectorXd&)> parameterization_{
+      [](const Eigen::VectorXd& q) -> Eigen::VectorXd {
+        return q;
+      }};
 };
 
 /** The IRIS-ZO (Iterative Regional Inflation by Semidefinite programming - Zero
@@ -160,7 +221,9 @@ box representing joint limits (e.g. from HPolyhedron::MakeBox).
 fraction, confidence level, and various algorithmic settings.
 
 The @p starting_ellipsoid and @p domain must describe elements in the same
-ambient dimension as the configuration space of the robot.
+ambient dimension as the configuration space of the robot, unless a
+parameterization is specified (in which case, they must match
+`options.parameterization_dimension`).
 @return A HPolyhedron representing the computed collision-free region in
 configuration space.
 @ingroup robot_planning
diff --git a/planning/iris/test/iris_zo_test.cc b/planning/iris/test/iris_zo_test.cc
