WIP augmented Lagrangian solver for adding equality constraints

[bamos]

This has some examples and a prototype @dishank-b and I created to help us get in equality constraints with an augmented Lagrangian method as described in this slide deck. It's appealing to have this method in Theseus because it relies mostly on the unconstrained GN/LM NLLS solvers we already have to solve for the update with the constraints moved into the objective as a penalty (via augmenting the Lagnangian). I created a scratch directory (that we can delete before squashing and merging this PR) with file augmented_lagrangian.py that implements an initial prototype of the method along with 2 examples in quadratic.py and car.py that call into it.

The quadratic example (quadratic.py)

This just makes sure we can solve min_x ||x||^2 s.t. x_0 = 1, which has an analytic solution of x^\star = [1., 0] and would make a good unit test once we've better-integrated this.

$ ./quadratic.py 
=== [0] mu: 1.00e+00, ||g||: 5.00e-01
{'x': tensor([[0.5000, 0.0000]])}
=== [1] mu: 1.00e+00, ||g||: 2.50e-01
{'x': tensor([[0.7500, 0.0000]])}
=== [2] mu: 2.00e+00, ||g||: 8.33e-02
{'x': tensor([[0.9167, 0.0000]])}
=== [3] mu: 4.00e+00, ||g||: 1.67e-02
{'x': tensor([[0.9833, 0.0000]])}
=== [4] mu: 4.00e+00, ||g||: 3.33e-03
{'x': tensor([[0.9967, 0.0000]])}
=== [5] mu: 4.00e+00, ||g||: 6.67e-04
{'x': tensor([[0.9993, 0.0000]])}
=== [6] mu: 4.00e+00, ||g||: 1.33e-04
{'x': tensor([[0.9999, 0.0000]])}
=== [7] mu: 4.00e+00, ||g||: 2.67e-05
{'x': tensor([[1.0000, 0.0000]])}
=== [8] mu: 4.00e+00, ||g||: 5.30e-06
{'x': tensor([[1.0000, 0.0000]])}
=== [9] mu: 4.00e+00, ||g||: 1.07e-06
{'x': tensor([[1.0000, 0.0000]])}

Car example (car.py)

Solves this optimization problem:

We can reproduce this in a simple setting that goes from the origin to the upper right, but in general it doesn't seem as stable as what Vandenberghe has implemented:

Discussion on remaining design choices

I still don't see the best way of integrating the prototype in augmented_lagrangian.py into Theseus. The rough changes are:

1. The addition of an equality_constraint_fn (g in the slides) that should equal zero.
2. Augmenting the cost with the term from function
3. Creating an outer loop that calls into the unconstrained optimizer on the augmented objective

One option is to create a new AugmentedLagrangian optimizer class that adds the constraints, updates the cost, and then iteratively calls into an unconstrained solver. This seems appealing as the existing optimizer could be kept as a child and called into.

Another design choice is how the user should specify the equality constraints. I've currently made the equality_constraint_fn have the same interface as a user-specified err_fn when using AutoDiffCostFunction.

TODOs once design choices are resolved

• Move the quadratic example to a unit test
• Add batching support and add a test for that too
• Add a more general way of augmenting the cost function beyond assuming it's an AutoDiffCostFunction. I think it could be done similar to https://github.com/facebookresearch/theseus/blob/main/theseus/core/robust_cost_function.py
• Add docs and an example for solving equality-constrained problems
• Try to stabilize the car example some more and get it to work on some of the harder settings
• Think about adding any other constrained NLLS examples
• Work out the details for differentiation/the backward pass and add examples/unit tests for that
• Delete scratch directory

Even more future directions

• Think about adding support for inequality constraints
• Is it a limitation to ask the user to provide a single function returning all of the equality constraints?

[bamos]

@mhmukadam @luisenp - do you think it makes sense to continue and try to move this prototype into an AugmentedLagrangian optimizer class with the details above, or should we proceed in another way?

[luisenp]

Thanks for working on this @bamos and @dishank-b, this is great!

I left some comments for discussion on the overall design. I think it's also OK if for a first draft we manually set up all the different terms in the augmented lagrangian to test things out (e.g., just add any extra terms necessary when creating the objective). Once we settle on a good representation for the equality constraints and multipliers, we can figure out how to automate this with an AugmentedLagrangian class, which is a suggestion that I like.

[luisenp]

I'm thinking it might be better (cleaner and more efficient) to introduce extra cost function terms for the constraints terms in the augmented Lagrangian, rather than extend to a new function with a larger error dimension. That is, the second row here would have its own cost function terms:

and Theseus already minimizes for the sum of squares internally. In terms of code, I think this should just involve moving the err_augmented term to its own CostFunction object. I see the following advantages:

• If the original error provides jacobian implementation, we can still use them rather than deferring to torch autograd once we put the original error in the augmented error function.
• More computationally efficient for sparse solvers, because the block jacobians should be of smaller size.
• If the user provides jacobians dg(x)/dx, it should be easier to manually compute the jacobians of the augmented Lagrangian terms compared to the current approach.
• The resulting code will be cleaner, but this is maybe a more subjective point.

One disadvantage of this approach is that we don't allow objectives from being modified (costs added/deleted) after the optimizer is created, but we should be able to get around this with an AugmentedLagrangianclass in a number of ways. I think for now, we can proof of concept things by manually adding the extra costs at the beginning of the example, just to get a feel for how things would look like. We can then figure out how to automate this process later.

[luisenp]

For this, I like the RobustCost-style wrapper, as you mentioned in the TODOs. Something like

equality_constraint = th.EqualityConstraint(th.AutoDiffCostFunction(...))

Following on the other comment above, a class like AugmentedLagrangian could then iterate over all cost functions in the objective, and add extra cost terms for any costs of type EqualityConstraint that computes the err_augmented term corresponding to each equality constraint.