Using model as output prediction constraint in MPC leads to incorrect input

[cmichiel]

Hey!

First of all, this package looks amazing and the application is really interesting. I hope it gets widely adopted.

Further, I am trying to implement an MPC controller where I am using a neural network as the prediction model. The model is trained on past input/output data and future inputs such that the output can be predicted in a multistep fashion. For the network I am using a RBF neural network with one layer and RBF activation functions. I have validated the model by analytically constructing the network in Matlab and running the controller and it manages to reference track. I also checked the conversion from Pytorch to L4Casadi and gave it an input that I know that works and it gave the correct output.

I have tried replacing the mathematical description of the system and then it works so I am suspecting that the problem is with l4Casadi and how I define the constraint with the L4Casadi model. Is there a specific way the input should be provided?

Model Code:

class ModelKoopman(nn.Module):
    def __init__(self, Tini, N, nbasis, input_dim, output_dim, KoopmanBasis):
        super().__init__()
        self.Tini = Tini
        self.N = N
        self.nbasis = nbasis
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.in_linear_features = (2 * Tini - 1)*input_dim
        self.out_linear_features = (nbasis)
        self.in_nl_features = (2 * Tini - 1+N)*input_dim
        self.out_nl_features = (nbasis)
        self.in_2_features = (nbasis)+N
        self.out_2_features = (N)*output_dim

        self.basis_funcKoopman = getattr(rbf_gauss.RBF_gaussian,KoopmanBasis)
        self.RBFKoopman= rbf_gauss.RBF_gaussian(self.in_linear_features, self.out_linear_features,self.basis_funcKoopman)
        self.l_2 = nn.Linear(self.in_2_features, self.out_2_features, bias=False)
        #nn.init.uniform_(self.l_2.weight)

    def forward(self, data):
        data_ini = data[:, 0 : (2 * self.Tini - 1)*self.input_dim]
        data_f   = data[:, (2 * self.Tini - 1)*self.input_dim :]
        x1 = self.RBFKoopman(data_ini)
        x1 = torch.cat((x1,data_f),1)
        x = self.l_2(x1)
        return x

MPC Setup Code:

def setup_SPC( n_y, n_u, N, Tini, Q, R, Rd, S):
    """Sets up the SPC controller.

        Args:
        n_y: Number of outputs.
        n_u: Number of inputs.
        N: Prediction horizon.
        Tini: Initial state length.
        Q: State tracking cost matrix (n_y x n_y).
        R: Input cost matrix (n_u x n_u).
        Rd: Input rate of change cost matrix (n_u x n_u).
        S: Terminal state cost

    Returns:
        S_spc: The CasADi solver.
        opt_x_num: Numerical instance of optimization variables (for warm-starting).
        opt_p_num: Numerical instance of parameters.
        lbx: Lower bounds on variables.
        ubx: Upper bounds on variables.
    """
    
    # Create optimization variables:
    opt_x = struct_symMX([
        entry('y_N', shape=(n_y), repeat=N),
        entry('u_N', shape=(n_u), repeat=N)
    ])

    # Create parameters of the optimization problem
    opt_p = struct_symMX([
        entry('u_Tini', shape=(n_u), repeat=Tini-1),
        entry('y_Tini', shape=(n_y), repeat=Tini),
        entry('ref', shape=(n_y), repeat=N),
        entry('u_prev', shape=(n_u)),
    ])

    # Create numerical instances of the structures (holding all zeros as entries)
    opt_x_num = opt_x(0)
    opt_p_num = opt_p(0)

    xmean = np.load("Train_mean.npy")
    xstd = np.load("Train_std.npy")
    
    
    # Define the objective function
      # Define the objective function
    obj = 0
    for k in range(N):
        y_k = opt_x['y_N', k]
        # State tracking cost
        obj += Q * cs.sumsqr(y_k - opt_p['ref', k])
        # Input cost
        obj += R * cs.sumsqr(opt_x['u_N', k])
        # Input rate of change cost (using u_prev for the first step)
        if k == 0:
            obj += Rd * cs.sumsqr(opt_x['u_N', k] - opt_p['u_prev'])
        else:
            obj += Rd * cs.sumsqr(opt_x['u_N', k] - opt_x['u_N', k-1])

    # # Terminal state cost
    obj += S * cs.sumsqr(opt_x['y_N',N-1] - opt_p['ref', N-1]) # Correct indexing for terminal cost


    # Create the constraints:
    inputs = (cs.vertcat(*opt_p['u_Tini'], *opt_p['y_Tini'],  *opt_x['u_N'])-xmean)/xstd
    y_sym = cs.MX.sym("yN", N, 1)
    input_sym = cs.MX.sym("uN", 2*Tini-1+N, 1)
    model = ModelKoopman(3, N, 40, 1, 1, "gaussian")
    model.load_state_dict(torch.load("RBF_Params_Koopman_Tini3_nbasisKoopman40_N10_KoopmanBasis_gaussian",weights_only=True))
   L4Model = l4c.L4CasADi(model,device='cpu',batched=True, name='y_expr', generate_jac_jac=True, generate_adj1=False, generate_jac_adj1=False)
    f = cs.Function("xdot", [y_sym, input_sym], [y_sym.T -  L4Model (input_sym.T)])
    cons = []
    yN = cs.vertcat(*opt_x['y_N'])

    # Create lower and upper bound structures and set all values to plus/minus infinity.
    lbx = opt_x(-np.inf)
    ubx = opt_x(np.inf)

    # Set only bounds on u_N
    lbx['u_N'] = -15
    ubx['u_N'] = 15

    opts = {'ipopt.print_level': 0, 'print_time': 0}
    #opts = {'fatrop.print_level': 0, 'print_time': 0}
    # Create Optim
    nlp = {'x': opt_x, 'f':obj, 'g': f(yN, inputs), 'p': opt_p}
    S_spc = nlpsol('S', 'ipopt', nlp, opts)
    
    return S_spc, opt_x_num, opt_p_num, lbx, ubx

[Tim-Salzmann]

Hi,

Thanks for reaching out. An easy way to sanity check if L4CasADi as a framework is the problem here would be to replace the learned mode with the mathematical description of the system but in PyTorch with L4CasADi. I am unsure if this is what you did already or if you replaced it with the mathematical description of the system in CasADi.

Should this result in the correct behavior, then the problem lies with the interaction of the model and optimizer rather than the L4CasADi framework, which simply exposes the model to CasADi / the optimizer. Keep in mind that without guidance the optimizer could query the model (and its derivatives) at any point in input space. Theoretically these points can be far away from where it was trained. This can result in undefined behavior where the optimizer diverges.

Let me know if this helps.

Best
Tim