Support continuous (unbounded) joints properly (#59)

sea-bass · web-flow · commit 3cf3437909e4 · 2023-11-15T10:26:45.000-05:00
diff --git a/include/pick_ik/robot.hpp b/include/pick_ik/robot.hpp
@@ -13,12 +13,27 @@ namespace pick_ik {
 
 struct Robot {
     struct Variable {
-        double clip_min, clip_max;
-        double span;
-        double min;
-        double max;
+        /// @brief Min, max, and middle position values of the variable.
+        double min, max, mid;
+
+        /// @brief Whether the variable's position is bounded.
+        bool bounded;
+
+        /// @brief The half-span (min - max) / 2.0 of the variable, or a default value if unbounded.
+        double half_span;
+
         double max_velocity_rcp;
         double minimal_displacement_factor;
+
+        /// @brief Generates a valid variable value given an optional initial value (for unbounded
+        /// joints).
+        auto generate_valid_value(double init_val = 0.0) const -> double;
+
+        /// @brief Returns true if a value is valid given the variable bounds.
+        auto is_valid(double val) const -> bool;
+
+        /// @brief Clamps a configuration to joint limits.
+        auto clamp_to_limits(double val) const -> double;
     };
     std::vector<Variable> variables;
 
@@ -27,8 +42,11 @@ struct Robot {
                      moveit::core::JointModelGroup const* jmg,
                      std::vector<size_t> tip_link_indices) -> Robot;
 
-    /** @brief Returns a random valid configuration. */
-    auto get_random_valid_configuration() const -> std::vector<double>;
+    /**
+     * @brief Sets a variable vector to a random configuration.
+     * @details Here, "valid" denotes that the joint values are with their specified limits.
+     */
+    auto set_random_valid_configuration(std::vector<double>& config) const -> void;
 
     /** @brief Check is a configuration is valid. */
     auto is_valid_configuration(std::vector<double> const& config) const -> bool;
diff --git a/src/goal.cpp b/src/goal.cpp
@@ -94,7 +94,7 @@ auto make_center_joints_cost_fn(Robot robot) -> CostFn {
         assert(robot.variables.size() == active_positions.size());
         for (size_t i = 0; i < active_positions.size(); ++i) {
             auto const& variable = robot.variables[i];
-            if (variable.clip_max == std::numeric_limits<double>::max()) {
+            if (!variable.bounded) {
                 continue;
             }
 
@@ -113,16 +113,16 @@ auto make_avoid_joint_limits_cost_fn(Robot robot) -> CostFn {
         assert(robot.variables.size() == active_positions.size());
         for (size_t i = 0; i < active_positions.size(); ++i) {
             auto const& variable = robot.variables[i];
-            if (variable.clip_max == std::numeric_limits<double>::max()) {
+            if (!variable.bounded) {
                 continue;
             }
 
             auto const position = active_positions[i];
             auto const weight = variable.minimal_displacement_factor;
-            auto const mid = (variable.min + variable.max) * 0.5;
-            auto const span = variable.span;
-            sum +=
-                std::pow(std::fmax(0.0, std::fabs(position - mid) * 2.0 - span * 0.5) * weight, 2);
+            sum += std::pow(
+                std::fmax(0.0, std::fabs(position - variable.mid) * 2.0 - variable.half_span) *
+                    weight,
+                2);
         }
         return sum;
     };
diff --git a/src/ik_gradient.cpp b/src/ik_gradient.cpp
@@ -25,7 +25,6 @@ auto step(GradientIk& self, Robot const& robot, CostFn const& cost_fn, double st
     auto const count = self.local.size();
 
     // compute gradient direction
-    self.working = self.local;
     for (size_t i = 0; i < count; ++i) {
         // test negative displacement
         self.working[i] = self.local[i] - step_size;
@@ -55,8 +54,6 @@ auto step(GradientIk& self, Robot const& robot, CostFn const& cost_fn, double st
                    [&](auto value) { return value * f; });
 
     // initialize line search
-    self.working = self.local;
-
     for (size_t i = 0; i < count; ++i) {
         self.working[i] = self.local[i] - self.gradient[i];
     }
@@ -78,9 +75,9 @@ auto step(GradientIk& self, Robot const& robot, CostFn const& cost_fn, double st
     // apply optimization step
     // (move along gradient direction by estimated step size)
     for (size_t i = 0; i < count; ++i) {
-        self.working[i] = std::clamp(self.local[i] - self.gradient[i] * joint_diff,
-                                     robot.variables[i].clip_min,
-                                     robot.variables[i].clip_max);
+        auto const& var = robot.variables[i];
+        auto updated_value = self.local[i] - self.gradient[i] * joint_diff;
+        self.working[i] = var.clamp_to_limits(updated_value);
     }
 
     // Always accept the solution and continue
diff --git a/src/ik_memetic.cpp b/src/ik_memetic.cpp
@@ -96,7 +96,10 @@ void MemeticIk::initPopulation(Robot const& robot,
     std::vector<double> const zero_grad(robot.variables.size(), 0.0);
     population_.resize(params_.population_size);
     for (size_t i = 0; i < params_.elite_size; ++i) {
-        auto const genotype = (i == 0) ? initial_guess : robot.get_random_valid_configuration();
+        auto genotype = initial_guess;
+        if (i > 0) {
+            robot.set_random_valid_configuration(genotype);
+        }
         population_[i] = Individual{genotype, cost_fn(genotype), 1.0, zero_grad};
     }
 
@@ -105,6 +108,10 @@ void MemeticIk::initPopulation(Robot const& robot,
         population_[i] = Individual{initial_guess, 0.0, 1.0, zero_grad};
     }
 
+    // Initialize fitnesses and extinctions
+    for (auto& individual : population_) {
+        individual.fitness = cost_fn(individual.genes);
+    }
     computeExtinctions();
     previous_fitness_.reset();
 }
@@ -149,11 +156,11 @@ void MemeticIk::reproduce(Robot const& robot, CostFn const& cost_fn) {
 
                 // Mutate
                 if (rsl::uniform_real(0.0, 1.0) < mutation_prob) {
-                    gene += extinction * joint.span * rsl::uniform_real(-0.5, 0.5);
+                    gene += extinction * joint.half_span * rsl::uniform_real(-1.0, 1.0);
                 }
 
                 // Clamp to valid joint values
-                gene = std::clamp(gene, joint.clip_min, joint.clip_max);
+                gene = robot.variables[j_idx].clamp_to_limits(gene);
 
                 // Approximate gradient
                 population_[i].gradient[j_idx] = gene - original_gene;
@@ -173,7 +180,7 @@ void MemeticIk::reproduce(Robot const& robot, CostFn const& cost_fn) {
 
         } else {
             // If the mating pool is empty, roll a new population member randomly.
-            population_[i].genes = robot.get_random_valid_configuration();
+            robot.set_random_valid_configuration(population_[i].genes);
             population_[i].fitness = cost_fn(population_[i].genes);
             for (auto& g : population_[i].gradient) {
                 g = 0.0;
diff --git a/src/pick_ik_plugin.cpp b/src/pick_ik_plugin.cpp
@@ -179,7 +179,7 @@ class PickIKPlugin : public kinematics::KinematicsBase {
             RCLCPP_WARN(
                 LOGGER,
                 "Initial guess exceeds joint limits. Regenerating a random valid configuration.");
-            init_state = robot_.get_random_valid_configuration();
+            robot_.set_random_valid_configuration(init_state);
         }
 
         // Optimize until a valid solution is found or we have timed out.
@@ -309,15 +309,11 @@ class PickIKPlugin : public kinematics::KinematicsBase {
             if (found_valid_solution || timeout_elapsed) {
                 done_optimizing = true;
             } else {
-                init_state = robot_.get_random_valid_configuration();
+                robot_.set_random_valid_configuration(init_state);
                 remaining_timeout -= total_optim_time.count();
             }
         }
 
-        if (!robot_.is_valid_configuration(solution)) {
-            std::cout << "INVALID SOLUTION!" << std::endl;
-        }
-
         return found_valid_solution;
     }
 
diff --git a/src/robot.cpp b/src/robot.cpp
@@ -13,8 +13,34 @@
 #include <moveit/robot_model/robot_model.h>
 #include <moveit/robot_state/robot_state.h>
 
+namespace {
+constexpr double kUnboundedVariableHalfSpan = M_PI;
+constexpr double kUnboundedJointSampleSpread = M_PI;
+}  // namespace
+
 namespace pick_ik {
 
+auto Robot::Variable::generate_valid_value(double init_val /* = 0.0*/) const -> double {
+    if (bounded) {
+        return rsl::uniform_real(min, max);
+    } else {
+        return rsl::uniform_real(init_val - kUnboundedJointSampleSpread,
+                                 init_val + kUnboundedJointSampleSpread);
+    }
+}
+
+auto Robot::Variable::is_valid(double val) const -> bool {
+    return (!bounded) || (val <= max && val >= min);
+}
+
+auto Robot::Variable::clamp_to_limits(double val) const -> double {
+    if (bounded) {
+        return std::clamp(val, min, max);
+    } else {
+        return std::clamp(val, val - half_span, val + half_span);
+    }
+}
+
 auto Robot::from(std::shared_ptr<moveit::core::RobotModel const> const& model,
                  moveit::core::JointModelGroup const* jmg,
                  std::vector<size_t> tip_link_indices) -> Robot {
@@ -33,17 +59,11 @@ auto Robot::from(std::shared_ptr<moveit::core::RobotModel const> const& model,
 
         auto var = Variable{};
 
-        bool bounded = bounds.position_bounded_;
-
+        var.bounded = bounds.position_bounded_;
         var.min = bounds.min_position_;
         var.max = bounds.max_position_;
-
-        var.clip_min = bounded ? var.min : std::numeric_limits<double>::lowest();
-        var.clip_max = bounded ? var.max : std::numeric_limits<double>::max();
-
-        var.span = var.max - var.min;
-
-        if (!(var.span >= 0 && var.span < std::numeric_limits<double>::max())) var.span = 1;
+        var.mid = 0.5 * (var.min + var.max);
+        var.half_span = var.bounded ? (var.max - var.min) / 2.0 : kUnboundedVariableHalfSpan;
 
         auto const max_velocity = bounds.max_velocity_;
         var.max_velocity_rcp = max_velocity > 0.0 ? 1.0 / max_velocity : 0.0;
@@ -64,22 +84,20 @@ auto Robot::from(std::shared_ptr<moveit::core::RobotModel const> const& model,
     return robot;
 }
 
-auto Robot::get_random_valid_configuration() const -> std::vector<double> {
-    std::vector<double> config;
+auto Robot::set_random_valid_configuration(std::vector<double>& config) const -> void {
     auto const num_vars = variables.size();
-    config.reserve(num_vars);
+    if (config.size() != num_vars) {
+        config.resize(num_vars);
+    }
     for (size_t idx = 0; idx < num_vars; ++idx) {
-        auto const var = variables[idx];
-        config.push_back(rsl::uniform_real(var.clip_min, var.clip_max));
+        config[idx] = variables[idx].generate_valid_value(config[idx]);
     }
-    return config;
 }
 
 auto Robot::is_valid_configuration(std::vector<double> const& config) const -> bool {
     auto const num_vars = variables.size();
     for (size_t idx = 0; idx < num_vars; ++idx) {
-        auto const var = variables[idx];
-        if (config[idx] > var.clip_max || config[idx] < var.clip_min) {
+        if (!variables[idx].is_valid(config[idx])) {
             return false;
         }
     }

Original file line number	Diff line number	Diff line change
`@@ -179,7 +179,7 @@ class PickIKPlugin : public kinematics::KinematicsBase {`
`179`	`179`	`RCLCPP_WARN(`
`180`	`180`	`LOGGER,`
`181`	`181`	`"Initial guess exceeds joint limits. Regenerating a random valid configuration.");`
`182`		`- init_state = robot_.get_random_valid_configuration();`
	`182`	`+ robot_.set_random_valid_configuration(init_state);`
`183`	`183`	`}`
`184`	`184`
`185`	`185`	`// Optimize until a valid solution is found or we have timed out.`
`@@ -309,15 +309,11 @@ class PickIKPlugin : public kinematics::KinematicsBase {`
`309`	`309`	`if (found_valid_solution \|\| timeout_elapsed) {`
`310`	`310`	`done_optimizing = true;`
`311`	`311`	`} else {`
`312`		`- init_state = robot_.get_random_valid_configuration();`
	`312`	`+ robot_.set_random_valid_configuration(init_state);`
`313`	`313`	`remaining_timeout -= total_optim_time.count();`
`314`	`314`	`}`
`315`	`315`	`}`
`316`	`316`
`317`		`- if (!robot_.is_valid_configuration(solution)) {`
`318`		`- std::cout << "INVALID SOLUTION!" << std::endl;`
`319`		`- }`
`320`		`-`
`321`	`317`	`return found_valid_solution;`
`322`	`318`	`}`
`323`	`319`