[SymForce] Expose optimization status

So we don't just tell you if we early exited or not

Topic: sym-opt-status
Relative: sky-enum-name-array
GitOrigin-RevId: 33044218c30d42b0cb10d27237b93a55a082fb74

README.md

@@ -266,6 +266,12 @@ Now run the optimization! This returns an [`Optimizer.Result`](https://symforce.
 result = optimizer.optimize(initial_values)
 ```
 
+We can check that the optimization succeeded, and look at the final error:
+```python
+assert result.status == Optimizer.Status.SUCCESS
+print(result.error())
+```
+
 Let's visualize what the optimizer did. The orange circles represent the fixed landmarks, the blue
 circles represent the robot, and the dotted lines represent the bearing measurements.
 
@@ -503,6 +509,7 @@ values.Set('e', sym::kDefaultEpsilond);
 // Optimize!
 const auto stats = optimizer.Optimize(values);
 
+std::cout << "Exit status: " << stats.status << std::endl;
 std::cout << "Optimized values:" << values << std::endl;
 ```
 

lcmtypes/symforce.lcm

@@ -215,16 +215,37 @@ struct sparse_matrix_structure_t {
   int64_t shape[2];
 }
 
+enum optimization_status_t {
+  // Uninitialized enum value
+  INVALID = 0,
+  // The optimization converged successfully
+  SUCCESS = 1,
+  // We hit the iteration limit before converging
+  HIT_ITERATION_LIMIT = 2,
+  // The solver failed to converge for some reason (other than hitting the iteration limit)
+  FAILED = 3,
+}
+
+enum levenberg_marquardt_solver_failure_reason_t : int32_t {
+  // Uninitialized enum value
+  INVALID = 0,
+  // We could not increase lambda high enough to make progress
+  LAMBDA_OUT_OF_BOUNDS = 1,
+}
+
 // Debug stats for a full optimization run
 struct optimization_stats_t {
   optimization_iteration_t iterations[];
 
   // Index into iterations of the best iteration (containing the optimal Values)
   int32_t best_index;
 
-  // Did the optimization early exit? (either because it converged, or because it could not find a
-  // good step)
-  boolean early_exited;
+  // What was the result of the optimization? (did it converge, fail, etc.)
+  optimization_status_t status;
+
+  // If status == FAILED, why?  This should be cast to the Optimizer::FailureReason enum
+  // for the nonlinear solver you used.
+  int32_t failure_reason;
 
   // The sparsity pattern of the jacobian, filled out if debug_stats=true
   sparse_matrix_structure_t jacobian_sparsity;

symforce/examples/bundle_adjustment/run_bundle_adjustment.cc

@@ -175,9 +175,11 @@ void RunBundleAdjustment() {
   spdlog::info("Lambda: {}", last_iter.current_lambda);
   spdlog::info("Initial error: {}", first_iter.new_error);
   spdlog::info("Final error: {}", best_iter.new_error);
+  spdlog::info("Status: {}", stats.status);
 
   // Check successful convergence
   SYM_ASSERT(best_iter.new_error < 10);
+  SYM_ASSERT(stats.status == sym::optimization_status_t::SUCCESS);
 }
 
 }  // namespace bundle_adjustment

symforce/examples/bundle_adjustment_fixed_size/run_bundle_adjustment.cc

@@ -211,9 +211,11 @@ void RunBundleAdjustment() {
   spdlog::info("Lambda: {}", last_iter.current_lambda);
   spdlog::info("Initial error: {}", first_iter.new_error);
   spdlog::info("Final error: {}", best_iter.new_error);
+  spdlog::info("Status: {}", stats.status);
 
   // Check successful convergence
   SYM_ASSERT(best_iter.new_error < 10);
+  SYM_ASSERT(stats.status == sym::optimization_status_t::SUCCESS);
 }
 
 }  // namespace bundle_adjustment_fixed_size

symforce/examples/robot_2d_localization/robot_2d_localization.py

@@ -124,6 +124,7 @@ def main() -> None:
     # Print some values
     print(f"Num iterations: {len(result.iterations) - 1}")
     print(f"Final error: {result.error():.6f}")
+    print(f"Status: {result.status}")
 
     for i, pose in enumerate(result.optimized_values["poses"]):
         print(f"Pose {i}: t = {pose.position()}, heading = {pose.rotation().to_tangent()[0]}")

symforce/examples/robot_2d_localization/run_localization.cc

@@ -85,6 +85,7 @@ void RunLocalization() {
   SYM_ASSERT(iteration_stats.size() == 9);
   SYM_ASSERT(6.39635 < first_iter.new_error && first_iter.new_error < 6.39637);
   SYM_ASSERT(best_iter.new_error < 0.00022003);
+  SYM_ASSERT(stats.status == sym::optimization_status_t::SUCCESS);
 
   const sym::Pose2d expected_p0({0.477063, 0.878869, -0.583038, -0.824491});
   const sym::Pose2d expected_p1({0.65425, 0.756279, 1.01671, -0.238356});

symforce/examples/robot_3d_localization/robot_3d_localization.py

@@ -251,6 +251,7 @@ def main() -> None:
     # Print some values
     print(f"Num iterations: {len(result.iterations) - 1}")
     print(f"Final error: {result.error():.6f}")
+    print(f"Status: {result.status}")
 
     for i, pose in enumerate(result.optimized_values["world_T_body"]):
         print(f"world_T_body {i}: t = {pose.position()}, R = {pose.rotation().to_tangent()}")

symforce/examples/robot_3d_localization/run_dynamic_size.cc

@@ -105,6 +105,7 @@ void RunDynamic() {
   spdlog::info("Lambda: {}", last_iter.current_lambda);
   spdlog::info("Initial error: {}", first_iter.new_error);
   spdlog::info("Final error: {}", best_iter.new_error);
+  spdlog::info("Status: {}", stats.status);
 }
 
 // Explicit template specializations for floats and doubles

symforce/examples/robot_3d_localization/run_fixed_size.cc

@@ -82,6 +82,7 @@ void RunFixed() {
   spdlog::info("Lambda: {}", last_iter.current_lambda);
   spdlog::info("Initial error: {}", first_iter.new_error);
   spdlog::info("Final error: {}", best_iter.new_error);
+  spdlog::info("Status: {}", stats.status);
 }
 
 template sym::Factor<double> BuildFixedFactor<double>();

symforce/opt/gnc_optimizer.h

@@ -73,9 +73,10 @@ class GncOptimizer : public BaseOptimizerType {
     // Iterate.
     BaseOptimizer::Optimize(values, num_iterations, populate_best_linearization, stats);
     while (static_cast<int>(stats.iterations.size()) < num_iterations) {
-      // NOTE(aaron): We shouldn't be here unless the optimization early exited (i.e. we had
-      // iterations left)
-      SYM_ASSERT(stats.early_exited);
+      if (stats.status != optimization_status_t::SUCCESS) {
+        // NOTE(aaron): The previous optimization did not converge, so do not continue
+        return;
+      }
 
       if (!updating_gnc) {
         return;

symforce/opt/levenberg_marquardt_solver.h

@@ -8,6 +8,7 @@
 #include <Eigen/Dense>
 #include <Eigen/Sparse>
 
+#include <lcmtypes/sym/levenberg_marquardt_solver_failure_reason_t.hpp>
 #include <lcmtypes/sym/optimization_stats_t.hpp>
 #include <lcmtypes/sym/optimizer_params_t.hpp>
 
@@ -126,6 +127,7 @@ class LevenbergMarquardtSolver {
   using MatrixType = typename LinearSolverType::MatrixType;
   using StateType = internal::LevenbergMarquardtState<MatrixType>;
   using LinearizationType = Linearization<MatrixType>;
+  using FailureReason = levenberg_marquardt_solver_failure_reason_t;
 
   // Function that evaluates the objective function and produces a quadratic approximation of
   // it by linearizing a least-squares residual.
@@ -172,9 +174,11 @@ class LevenbergMarquardtSolver {
 
   void UpdateParams(const optimizer_params_t& p);
 
-  // Run one iteration of the optimization. Returns true if the optimization should early exit.
-  bool Iterate(const LinearizeFunc& func, OptimizationStats<Scalar>& stats,
-               const bool debug_stats = false, const bool include_jacobians = false);
+  // Run one iteration of the optimization. Returns the optimization status, which will be empty if
+  // the optimization should not exit yet.
+  optional<std::pair<optimization_status_t, FailureReason>> Iterate(
+      const LinearizeFunc& func, OptimizationStats<Scalar>& stats, const bool debug_stats = false,
+      const bool include_jacobians = false);
 
   const Values<Scalar>& GetBestValues() const {
     SYM_ASSERT(state_.BestIsValid());

symforce/opt/levenberg_marquardt_solver.tcc

@@ -141,9 +141,11 @@ void LevenbergMarquardtSolver<ScalarType, LinearSolverType>::UpdateParams(
 }
 
 template <typename ScalarType, typename LinearSolverType>
-bool LevenbergMarquardtSolver<ScalarType, LinearSolverType>::Iterate(
-    const LinearizeFunc& func, OptimizationStats<Scalar>& stats, const bool debug_stats,
-    const bool include_jacobians) {
+optional<std::pair<optimization_status_t, levenberg_marquardt_solver_failure_reason_t>>
+LevenbergMarquardtSolver<ScalarType, LinearSolverType>::Iterate(const LinearizeFunc& func,
+                                                                OptimizationStats<Scalar>& stats,
+                                                                const bool debug_stats,
+                                                                const bool include_jacobians) {
   SYM_TIME_SCOPE("LM<{}>::Iterate()", id_);
 
   // new -> init
@@ -234,9 +236,12 @@ bool LevenbergMarquardtSolver<ScalarType, LinearSolverType>::Iterate(
     spdlog::warn("LM<{}> Encountered non-finite error: {}", id_, new_error);
   }
 
+  optional<std::pair<optimization_status_t, FailureReason>> status{};
+
   // Early exit if the reduction in error is too small.
-  bool should_early_exit =
-      (relative_reduction > 0) && (relative_reduction < p_.early_exit_min_reduction);
+  if ((relative_reduction > 0) && (relative_reduction < p_.early_exit_min_reduction)) {
+    status = {optimization_status_t::SUCCESS, {}};
+  }
 
   {
     SYM_TIME_SCOPE("LM<{}>: accept_update bookkeeping", id_);
@@ -251,7 +256,9 @@ bool LevenbergMarquardtSolver<ScalarType, LinearSolverType>::Iterate(
     }
 
     // If we didn't accept the update and lambda is maxed out, just exit.
-    should_early_exit |= (!accept_update && current_lambda_ >= p_.lambda_upper_bound);
+    if (!accept_update && current_lambda_ >= p_.lambda_upper_bound) {
+      status = {optimization_status_t::FAILED, FailureReason::LAMBDA_OUT_OF_BOUNDS};
+    }
 
     if (!accept_update) {
       current_lambda_ *= p_.lambda_up_factor;
@@ -277,7 +284,7 @@ bool LevenbergMarquardtSolver<ScalarType, LinearSolverType>::Iterate(
     iteration_stats.update_accepted = accept_update;
   }
 
-  return should_early_exit;
+  return status;
 }
 
 template <typename ScalarType, typename LinearSolverType>

symforce/opt/optimization_stats.h

@@ -21,9 +21,12 @@ struct OptimizationStats {
   // Index into iterations of the best iteration (containing the optimal Values)
   int32_t best_index{0};
 
-  // Did the optimization early exit? (either because it converged, or because it could not find a
-  // good step)
-  bool early_exited{false};
+  // What was the result of the optimization?
+  optimization_status_t status{};
+
+  // If status == FAILED, why?  This should be cast to the Optimizer::FailureReason enum for the
+  // nonlinear solver you used.
+  int32_t failure_reason{};
 
   // The linearization at best_index (at optimized_values), filled out if
   // populate_best_linearization=true
@@ -37,7 +40,7 @@ struct OptimizationStats {
   sparse_matrix_structure_t cholesky_factor_sparsity;
 
   optimization_stats_t GetLcmType() const {
-    return optimization_stats_t(iterations, best_index, early_exited, jacobian_sparsity,
+    return optimization_stats_t(iterations, best_index, status, failure_reason, jacobian_sparsity,
                                 linear_solver_ordering, cholesky_factor_sparsity);
   }
 
@@ -48,7 +51,8 @@ struct OptimizationStats {
     iterations.reserve(num_iterations);
 
     best_index = {};
-    early_exited = {};
+    status = {};
+    failure_reason = {};
     best_linearization = {};
     jacobian_sparsity = {};
     linear_solver_ordering = {};

symforce/opt/optimizer.h

@@ -72,6 +72,7 @@ class Optimizer {
  public:
   using Scalar = ScalarType;
   using NonlinearSolver = NonlinearSolverType;
+  using FailureReason = typename NonlinearSolver::FailureReason;
 
   /**
    * Base constructor

symforce/opt/optimizer.py

@@ -13,7 +13,11 @@
 import numpy as np
 
 from lcmtypes.sym._index_entry_t import index_entry_t
+from lcmtypes.sym._levenberg_marquardt_solver_failure_reason_t import (
+    levenberg_marquardt_solver_failure_reason_t,
+)
 from lcmtypes.sym._optimization_iteration_t import optimization_iteration_t
+from lcmtypes.sym._optimization_status_t import optimization_status_t
 from lcmtypes.sym._optimizer_params_t import optimizer_params_t
 from lcmtypes.sym._sparse_matrix_structure_t import sparse_matrix_structure_t
 from lcmtypes.sym._values_t import values_t
@@ -106,6 +110,9 @@ class Params:
         early_exit_min_reduction: float = 1e-6
         enable_bold_updates: bool = False
 
+    Status = optimization_status_t
+    FailureReason = levenberg_marquardt_solver_failure_reason_t
+
     @dataclass
     class Result:
         """
@@ -127,9 +134,11 @@ class Result:
             to be the last iteration, if the optimizer tried additional steps which did not reduce
             the error
 
-        early_exited:
-            Did the optimization early exit?  This can happen because it converged successfully,
-            of because it was unable to make progress
+        status:
+            What was the result of the optimization? (did it converge, fail, etc.)
+
+        failure_reason:
+            If status == FAILED, why?
 
         best_linearization:
             The linearization at best_index (at optimized_values), filled out if
@@ -161,8 +170,12 @@ def best_index(self) -> int:
             return self._stats.best_index
 
         @cached_property
-        def early_exited(self) -> bool:
-            return self._stats.early_exited
+        def status(self) -> optimization_status_t:
+            return self._stats.status
+
+        @cached_property
+        def failure_reason(self) -> levenberg_marquardt_solver_failure_reason_t:
+            return Optimizer.FailureReason(self._stats.failure_reason)
 
         @cached_property
         def best_linearization(self) -> T.Optional[cc_sym.Linearization]:

symforce/opt/optimizer.tcc

@@ -287,18 +287,32 @@ void Optimizer<ScalarType, NonlinearSolverType>::IterateToConvergence(
     Values<Scalar>& values, const int num_iterations, const bool populate_best_linearization,
     OptimizationStats<Scalar>& stats) {
   SYM_TIME_SCOPE("Optimizer<{}>::IterateToConvergence", name_);
-  bool optimization_early_exited = false;
+  SYM_ASSERT(num_iterations > 0, "num_iterations must be positive, got {}", num_iterations);
 
   // Iterate
-  for (int i = 0; i < num_iterations; i++) {
-    const bool should_early_exit =
+  int i;
+  for (i = 0; i < num_iterations; i++) {
+    const auto maybe_status_and_failure_reason =
         nonlinear_solver_.Iterate(linearize_func_, stats, debug_stats_, include_jacobians_);
-    if (should_early_exit) {
-      optimization_early_exited = true;
+    if (maybe_status_and_failure_reason) {
+      const auto& status_and_failure_reason = maybe_status_and_failure_reason.value();
+
+      SYM_ASSERT(status_and_failure_reason.first != optimization_status_t::INVALID,
+                 "NonlinearSolver::Iterate should never return INVALID");
+      SYM_ASSERT(status_and_failure_reason.first != optimization_status_t::HIT_ITERATION_LIMIT,
+                 "NonlinearSolver::Iterate should never return HIT_ITERATION_LIMIT");
+
+      stats.status = status_and_failure_reason.first;
+      stats.failure_reason = status_and_failure_reason.second.int_value();
       break;
     }
   }
 
+  if (i == num_iterations) {
+    stats.status = optimization_status_t::HIT_ITERATION_LIMIT;
+    stats.failure_reason = {};
+  }
+
   {
     SYM_TIME_SCOPE("Optimizer<{}>::CopyValuesAndLinearization", name_);
     // Save best results
@@ -316,8 +330,6 @@ void Optimizer<ScalarType, NonlinearSolverType>::IterateToConvergence(
     const auto& linearization = nonlinear_solver_.GetBestLinearization();
     stats.jacobian_sparsity = GetSparseStructure(linearization.jacobian);
   }
-
-  stats.early_exited = optimization_early_exited;
 }
 
 template <typename ScalarType, typename NonlinearSolverType>