NaN values in Augmented Lagrangian for inequality constraints with penalty coefficient = 0

[juan43ramirez]

Bug

Using the Augmented Lagrangian formulation for an inequality-constrained problem with a penalty coefficient of 0 leads to NaN values in the Lagrangian—this is caused by an underlying division by zero.

Steps

Simple example that reproduces the error:

import torch
import cooper
import cooper.formulations
import cooper.optim

class SimpleCMP(cooper.ConstrainedMinimizationProblem):
    def __init__(self):

        super().__init__()

        multiplier = cooper.multipliers.DenseMultiplier(num_constraints=1)
        penalty_coefficient = cooper.penalty_coefficients.DensePenaltyCoefficient(
            init=torch.tensor([0.0])
        )

        self.constraint = cooper.Constraint(
            constraint_type=cooper.ConstraintType.INEQUALITY,
            formulation_type=cooper.formulations.AugmentedLagrangian,
            multiplier=multiplier,
            penalty_coefficient=penalty_coefficient,
        )

    def compute_cmp_state(self, x: torch.Tensor) -> cooper.CMPState:
        constraint_state = cooper.ConstraintState(violation=x - 1.0)
        return cooper.CMPState(
            loss=x ** 2,
            observed_constraints={self.constraint: constraint_state},
        )

if __name__ == "__main__":

    x = torch.nn.Parameter(torch.tensor(2.0))
    cmp = SimpleCMP()

    primal_optimizer = torch.optim.SGD([x], lr=0.01)
    dual_optimizer = torch.optim.SGD(cmp.dual_parameters(), lr=0.01, maximize=True)

    cooper_optimizer = cooper.optim.AlternatingPrimalDualOptimizer(
        cmp=cmp,
        primal_optimizers=primal_optimizer,
        dual_optimizers=dual_optimizer,
    )

    # The line below leads to a nan primal and lagrangian
    rollout = cooper_optimizer.roll(compute_cmp_state_kwargs={"x": x})
    primal_lagrangian = rollout.primal_lagrangian_store.lagrangian

    assert torch.isnan(primal_lagrangian)

Expected behavior

When the penalty coefficient is zero, the formulation should default to the standard Lagrangian contribution—this behavior is already the case for equality constraints.

Context

Fixing this requires editing the following lines to avoid dividing by the penalty coefficient when it is zero.

Note that this may also be an opportunity to optimize for numerical precision, as very small penalty coefficients can lead to instability (dividing a small numerator by a small denominator).

https://github.com/cooper-org/cooper/blob/30b158e0262f7073c7710bd1c4996e462af4ff35/src/cooper/formulations/utils.py#L155-L159