source_current/src/ffolayer/ffoqp.py

import torch
import torch.nn as nn
from torch.optim import Optimizer
from torch import Tensor
from torch.autograd import Function

import numpy as np
import scipy
import time
import cvxpy
# import solvers
# from qpthlocal.solvers.pdipm import batch as pdipm_b
# from qpthlocal.solvers.pdipm import spbatch as pdipm_spb
# from qpthlocal.solvers.cvxpy import forward_single_np
from .utils import forward_single_np
from enum import Enum
from .utils import extract_nBatch, expandParam
from typing import cast, List, Optional, Union

# from cvxpylayers.torch import CvxpyLayer

# class QPSolvers(Enum):
#     PDIPM_BATCHED = 1
#     CVXPY = 2

# class ffoqp(torch.nn.Module):
#     def __init__(self, eps=1e-12, verbose=0, notImprovedLim=3, maxiter=20, solver=None, lamb=100):
#         super(ffoqp, self).__init__()
#         self.eps = eps
#         self.verbose = verbose
#         self.notImprovedLim = notImprovedLim
#         self.maxiter = maxiter
#         self.solver = solver if solver is not None else QPSolvers.CVXPY
#         self.lamb = lamb

def ffoqp(eps=1e-12, verbose=0, notImprovedLim=3, maxIter=20, lamb=100, check_Q_spd=True,
          solver='GUROBI', solver_opts={"verbose": False}):

    class QPFunctionFn(torch.autograd.Function):
        @staticmethod
        def forward(ctx, Q_, p_, G_, h_, A_, b_):
            # p_ = p_ + 1/lamb * torch.randn_like(p_)
            start_time = time.time()
            nBatch = extract_nBatch(Q_, p_, G_, h_, A_, b_)
            Q, _ = expandParam(Q_, nBatch, 3)
            p, _ = expandParam(p_, nBatch, 2)
            G, _ = expandParam(G_, nBatch, 3)
            h, _ = expandParam(h_, nBatch, 2)
            A, _ = expandParam(A_, nBatch, 3)
            b, _ = expandParam(b_, nBatch, 2)

            if check_Q_spd:
                try:
                    torch.linalg.cholesky(Q)
                except:
                    raise RuntimeError('Q is not SPD.')

            _, nineq, nz = G.size()
            neq = A.size(1) if A.nelement() > 0 else 0
            assert(neq > 0 or nineq > 0)
            ctx.neq, ctx.nineq, ctx.nz = neq, nineq, nz

            # if solver == QPSolvers.PDIPM_BATCHED:
            #     ctx.Q_LU, ctx.S_LU, ctx.R = pdipm_b.pre_factor_kkt(Q, G, A)
            #     zhats, nus, lams, slacks = pdipm_b.forward(
            #         Q, p, G, h, A, b, ctx.Q_LU, ctx.S_LU, ctx.R,
            #         eps, verbose, notImprovedLim, maxIter)
            # elif solver == QPSolvers.CVXPY:
            vals = torch.Tensor(nBatch).type_as(Q)
            zhats = torch.Tensor(nBatch, ctx.nz).type_as(Q)
            lams = torch.Tensor(nBatch, ctx.nineq).type_as(Q)
            nus = torch.Tensor(nBatch, ctx.neq).type_as(Q) \
                if ctx.neq > 0 else torch.Tensor()
            slacks = torch.Tensor(nBatch, ctx.nineq).type_as(Q)

            for i in range(nBatch):
                Ai, bi = (A[i], b[i]) if neq > 0 else (None, None)
                vals[i], zhati, nui, lami, si = forward_single_np(
                    *[x.cpu().numpy() if x is not None else None
                    for x in (Q[i], p[i], G[i], h[i], Ai, bi)],
                    solver=solver, solver_opts=solver_opts)
                # if zhati[0] is None:
                #     import IPython, sys; IPython.embed(); sys.exit(-1)
                zhats[i] = torch.Tensor(zhati)
                lams[i] = torch.Tensor(lami)
                slacks[i] = torch.Tensor(si)
                if neq > 0:
                    nus[i] = torch.Tensor(nui)

            ctx.vals = vals
            ctx.lams = lams
            ctx.nus = nus
            ctx.slacks = slacks

            # else:
            #     raise NotImplementedError("Solver not implemented")

            # ctx.vals = vals
            ctx.lams = lams
            ctx.nus = nus
            ctx.slacks = slacks

            ctx.save_for_backward(zhats, lams, nus, Q_, p_, G_, h_, A_, b_)
            # print('value', vals)
            # print('solution', zhats)
            return zhats

        @staticmethod
        def backward(ctx, grad_output):
            # Backward pass to compute gradients with respect to inputs
            zhats, lams, nus, Q_, p_, G_, h_, A_, b_ = ctx.saved_tensors
            lams = torch.clamp(lams, min=0)

            nBatch = extract_nBatch(Q_, p_, G_, h_, A_, b_)
            # Formulate a different QP to solve
            # L = f + \lamb * (g + lams * h - g^*) + \lamb^2 * |h_+|^2
            Q, Q_e = expandParam(Q_, nBatch, 3)
            p, p_e = expandParam(p_, nBatch, 2)
            G, G_e = expandParam(G_, nBatch, 3)
            h, h_e = expandParam(h_, nBatch, 2)
            A, A_e = expandParam(A_, nBatch, 3)
            b, b_e = expandParam(b_, nBatch, 2)

            Q, p, G, h, A, b = Q.to(zhats.device), p.to(zhats.device), G.to(zhats.device), h.to(zhats.device), A.to(zhats.device), b.to(zhats.device)

            # Running gradient descent for a few iterations
            _, nineq, nz = G.size()
            neq = A.size(1) if A.nelement() > 0 else 0

            delta_directions = grad_output.unsqueeze(-1)
            zhats = zhats.unsqueeze(-1).detach()

            # Iterative solution
            # print('newzhat shape:', zhats.shape)
            # print('Q shape:', Q.shape)
            iterative = False
            if iterative:
                with torch.enable_grad():
                    newzhat = zhats.clone().detach().requires_grad_(True)
                    optimizer = torch.optim.Adam([newzhat], lr=1e-3)
                    gd_maxiter = 1 # int(np.sqrt(lamb) * 100)
                    # print('dual solutions', lams)
                    for i in range(gd_maxiter):
                        objectives = (0.5 * newzhat.transpose(-1,-2) @ Q.detach() @ newzhat + (p.detach().unsqueeze(1) + delta_directions.transpose(-1,-2) / lamb) @ newzhat).squeeze(-1,-2)
                        violations = G.detach() @ newzhat - h.unsqueeze(-1)
                        active_constraints = (lams > 1e-5).unsqueeze(-1).float()
                        ineq_penalties = lams.unsqueeze(1) @ violations + 0.5 * lamb * torch.sum((violations * active_constraints) ** 2, dim=(-1,-2))
                        if neq > 0:
                            eq_penalties = nus.unsqueeze(1) @ (A.detach() @ newzhat.unsqueeze(-1) - b.detach().unsqueeze(-1))
                        else:
                            eq_penalties = 0
                        # print('obj, vio, active_constraints, lamb, ineq_penality shape:', objectives.shape, violations.shape, active_constraints.shape, lamb, ineq_penalties.shape)
                        lagrangians = objectives + ineq_penalties + eq_penalties
                        loss = torch.sum(lagrangians)
                        # print('Iteration {}: loss: {}, obj: {}, violation: {}'.format(i, loss, objectives.mean(), violations.mean()))
                        loss.backward()
                        optimizer.step()
                        optimizer.zero_grad()

                # print('new zhat', newzhats.detach())

                # Clampping the dual solutions
                # lams = torch.clamp(lams, max=100)

            else:
                # Deterministic solution by solving a new QP
                temperature = 10
                start_time = time.time()
                active_constraints = torch.tanh(lams * temperature).unsqueeze(-1)
                # active_constraints = (lams > 1e-3).unsqueeze(-1).float()
                G_active = G * active_constraints
                h_active = h.unsqueeze(-1) * active_constraints
                newQ = Q + lamb * G_active.transpose(-1,-2) @ G_active # + torch.eye(nz).repeat(nBatch, 1, 1).to(Q.device)
                newp = p.unsqueeze(-1) + delta_directions / lamb - lamb * G_active.transpose(-1,-2) @ h_active + G.transpose(-1,-2) @ lams.unsqueeze(-1) # - zhats
                # print('newQ, newp shape:', newQ.shape, newp.shape)
                # print('A, b shape:', A.shape, b.shape)
                if neq > 0:
                    newQ = torch.cat((newQ, lamb * A), dim=1)
                    newp = torch.cat((newp, - lamb * b.unsqueeze(-1)), dim=1)
                # print('condition number:', torch.linalg.cond(newQ))
                # newzhat = torch.linalg.solve(newQ, -newp)
                newzhat = torch.linalg.lstsq(newQ, -newp, driver='gels').solution
                # newzhat = - newQ.pinverse() @ newp
            # print('prediction', p)
            # print('solution distance:', torch.linalg.norm(newzhat - zhats))
            # print(lams)
            # print(zhats)
            # print(newzhat)
            # print('solution shape:', newzhat.shape, zhats.shape)
            # print('solution max distance:', torch.max(torch.abs(newzhat - zhats)))
            # assert False
            # print('newQ, p, newp, z shape:', newQ.shape, p.shape, newp.shape, newzhats.shape)

            # Computing the gradients of the Lagrangians
            start_time = time.time()
            with torch.enable_grad():
                Q_torch = Q.detach().clone().requires_grad_(True)
                p_torch = p.detach().clone().requires_grad_(True)
                G_torch = G.detach().clone().requires_grad_(True)
                h_torch = h.detach().clone().requires_grad_(True)
                A_torch = A.detach().clone().requires_grad_(True)
                b_torch = b.detach().clone().requires_grad_(True)
               
                upper_level_objectives = (delta_directions.transpose(-1,-2) @ newzhat).squeeze(-1,-2)
                objectives = (0.5 * newzhat.transpose(-1,-2) @ Q_torch @ newzhat + p_torch.unsqueeze(1) @ newzhat).squeeze(-1,-2) # 1/2 * z^T Q z + p^T z
                optimal_objectives = (0.5 * zhats.transpose(-1,-2) @ Q_torch @ zhats + p_torch.unsqueeze(1) @ zhats).squeeze(-1,-2) # 1/2 * z*^T Q z* + p^T z*
                violations = G_torch @ newzhat - h.unsqueeze(-1) # G z - h
                active_constraints = torch.tanh(lams * temperature).unsqueeze(-1).float()
                ineq_penalties = lams.unsqueeze(1) @ violations + 0.5 * lamb * torch.sum((violations * active_constraints) ** 2, dim=(-1,-2))
                if neq > 0:
                    eq_violations = A_torch @ newzhat - b_torch.unsqueeze(-1)
                    eq_penalties = nus.unsqueeze(1) @ eq_violations # + 0.5 * lamb * torch.sum(eq_violations ** 2, dim=(-1,-2))
                    # print(eq_violations)
                else:
                    eq_penalties = 0
                # print('obj, vio, active_constraints, lamb, ineq_penality shape:', objectives.shape, violations.shape, active_constraints.shape, lamb, ineq_penalties.shape)
                lagrangians = upper_level_objectives / lamb + objectives - optimal_objectives + ineq_penalties # + eq_penalties
                loss = torch.sum(lagrangians) * lamb
                loss.backward()

                Q_grad = Q_torch.grad
                p_grad = p_torch.grad
                G_grad = G_torch.grad
                h_grad = h_torch.grad
                A_grad = A_torch.grad
                b_grad = b_torch.grad

            return (Q_grad, p_grad, G_grad, h_grad, A_grad, b_grad)  # (None,) * len(ctx.saved_tensors)

    return QPFunctionFn.apply

def to_numpy(x):
    # convert torch tensor to numpy array
    return x.cpu().detach().double().numpy()