Differentiable Parameters

OPTIMake supports computing the gradient of the optimal solution with respect to parameters (parameter sensitivity). After a parameter is declared as differentiable during modeling, the generated code additionally provides a backward interface, which can be used to:

• Compute the local feedback gain of MPC, e.g., the derivative of the first control input with respect to the initial state $\partial u_0 / \partial x_0$
• Analyze the sensitivity of the optimal solution with respect to parameters such as reference trajectories and weights
• Use the solver as a differentiable optimization layer and learn parameters in the model by gradient descent (e.g., objective weights and constraint parameters)

Principle

Denote the optimization variables by $v$, the differentiable parameters by $\theta$, and the optimal solution by $v^*(\theta)$. For a scalar function $c(v)$ of interest, the user provides its gradient $\nabla_v c$ with respect to the optimal solution, and the backward interface computes the gradient of $c$ with respect to the differentiable parameters (vector-Jacobian product):

$$\nabla_\theta c = \left(\frac{\partial v^*}{\partial \theta}\right)^T \nabla_v c$$

backward is based on implicit differentiation of the KKT conditions at the optimal solution. It does not backpropagate through the iterations, line search, or other steps of the solving process, and the computational cost of each call is approximately that of one KKT system solve.

info

If $c$ also depends directly on the parameters $\theta$ (i.e., $c(v, \theta)$), the user needs to add the partial derivative of $c$ with respect to $\theta$ with $v$ held fixed.

Modeling

Set differentiable=True in the parameter or parameters interface to declare a parameter as differentiable. Both stage-independent and stage-dependent parameters are supported:

from optimake import *

prob = multi_stage_problem('tracking', 20)

p = prob.variable('p')
v = prob.variable('v')
a = prob.variable('a')

# differentiable parameters
p0, v0 = prob.parameters(['p0', 'v0'], differentiable=True)
weight = prob.parameter('weight', differentiable=True)
p_ref = prob.parameter('p_ref', stage_dependent=True, differentiable=True)

prob.objective(general_objective(weight * (p - p_ref)**2 + 0.1 * a**2))
prob.equality(differential_equation(state=[p, v], state_dot=[v, a],
                                    stepsize=0.1, discretization_method='erk4'))
prob.start_equality(general_equality([p - p0, v - v0]))

option = codegen_option()
option.default_tolerance_level = 'high'
code_generator().codegen(prob, option)

• differentiable defaults to False. Differentiable parameters are still ordinary parameters; how they are set in param and param_stage and their indices are unchanged
• As long as the model contains differentiable parameters, the backward interface is generated during code generation; no additional code generation option is needed
• Both the 'pdipm' and 'sqp' solving algorithms are supported

Generated Interface

Taking the tracking model above as an example, the generated content related to differentiable parameters is as follows:

Name	Description
problem.grad_var[N][DIM_VAR]	Input, $\nabla_v c$. The variable order is the same as output.primal.var, indexed by TRACKING_VAR_*
output.grad_param[DIM_DP]	Output, the gradient of $c$ with respect to the stage-independent differentiable parameters
output.grad_param_stage[N][DIM_DPS]	Output, the gradient of $c$ with respect to the stage-dependent differentiable parameters of each stage
TRACKING_DIM_DP, TRACKING_DIM_DPS	Number of stage-independent and stage-dependent differentiable parameters
TRACKING_DIFF_<NAME>	Index of a stage-independent differentiable parameter in grad_param, e.g., TRACKING_DIFF_WEIGHT
TRACKING_DIFF_STAGE_<NAME>	Index of a stage-dependent differentiable parameter in grad_param_stage[i], e.g., TRACKING_DIFF_STAGE_P_REF

warning

The gradient arrays contain only differentiable parameters, and their indices differ from those in param/param_stage. When reading gradients, use the DIFF macros (e.g., TRACKING_DIFF_WEIGHT) instead of the PARAM macros (e.g., TRACKING_PARAM_WEIGHT). When the number of differentiable parameters of one kind is 0, the corresponding grad_param or grad_param_stage is not generated.

The backward interface is defined as follows:

int tracking_backward(const Tracking_Problem *const prob,
                      const Tracking_Option *const option,
                      Tracking_WorkSpace *const workspace,
                      Tracking_Output *const output, int refactor);

• prob, option, workspace: The same as those used in the preceding solve, where prob->grad_var is the input $\nabla_v c$
• output: grad_param and grad_param_stage are written; solution results in output such as the primal variables are not modified
• refactor:
  • 1: Recompute the derivatives and refactorize the KKT system at the final solution, then compute the gradient
  • 0: Reuse the existing KKT factorization (from the preceding solve or the most recent call with refactor=1) and solve the linear system only once
• Return value: 1 indicates success; 0 indicates failure, in which case the gradient arrays are set to zero and the reason for failure is printed

Usage

The calling sequence is: solve → set grad_var → backward → read the gradients. The following example computes the gradients of the first control input $a_0$ and the position at the last stage $p_{N-1}$ with respect to the differentiable parameters:

int solve_status = tracking_solve(&problem, &option, &workspace, &output);
if (solve_status == 1) {
    /* c = a_0 */
    memset(problem.grad_var, 0, sizeof(problem.grad_var));
    problem.grad_var[0][TRACKING_VAR_A] = 1.0;
    /* use refactor = 1 for the first call after each solve */
    if (tracking_backward(&problem, &option, &workspace, &output, 1) == 1) {
        double da0_dp0 = output.grad_param[TRACKING_DIFF_P0];
        double da0_dweight = output.grad_param[TRACKING_DIFF_WEIGHT];
        double da0_dpref_5 = output.grad_param_stage[5][TRACKING_DIFF_STAGE_P_REF];
    }

    /* c = p_{N-1}, reuse the factorization with refactor = 0 after changing grad_var */
    memset(problem.grad_var, 0, sizeof(problem.grad_var));
    problem.grad_var[problem.valid_dim_N - 1][TRACKING_VAR_P] = 1.0;
    if (tracking_backward(&problem, &option, &workspace, &output, 0) == 1) {
        double dpN_dp0 = output.grad_param[TRACKING_DIFF_P0];
    }
}

info

Each call to backward yields the gradient of one scalar function, i.e., one row of the Jacobian matrix. To obtain the full Jacobian matrix (e.g., $\partial u_0 / \partial x_0$), set grad_var to the unit vector of each component in turn and call backward multiple times: use refactor=1 for the first call and refactor=0 for the rest.

Notes

• Solve status: The gradient is meaningful only at a converged optimal solution, so it is recommended to use the result of backward only when solve_status is 1. backward itself does not check the solve status; for other solve statuses, the returned value is only a sensitivity estimate of the current local system
• Calling order: Between solve and backward, do not modify the problem (including parameters and valid_dim_N), option, or workspace, except for grad_var. After the parameters change, solve again before calling backward
• Output: Each call to backward overwrites grad_param and grad_param_stage, and solve does not update these two arrays. The gradients of stages beyond valid_dim_N and of differentiable parameters not used in the model are 0
• Choice of refactor: The KKT factorization retained at the end of the solve usually comes from before the last iterate update, so using refactor=0 directly may introduce additional error. It is recommended to use refactor=1 for the first call after each solve, and then use refactor=0 when changing grad_var for the same solution
• Workspace: backward uses the workspace as computation space but does not modify the primal/dual variables in it, so the warm start of the next solve is not affected
• Accuracy: The accuracy of the gradient depends on the solution accuracy, so a higher solution accuracy is recommended (e.g., the code generation option default_tolerance_level = 'high'). When 'gauss-newton' or 'bfgs' is used as the Hessian approximation, the gradient is computed with the corresponding approximate Hessian and is therefore an approximate gradient; 'exact' is recommended when high accuracy is required
• Non-differentiable points: The derivative may not be unique where an inequality constraint switches between active and inactive, or where a soft constraint switches before and after soft constraint hardening. When a variable is restricted by an active bound, its derivative with respect to the parameters is 0. Deactivated constraints do not contribute to the gradient
• Unsupported cases: The following cases cause an error during code generation
  • Problems defined with external C/C++ functions
  • differential_equation with the 'irk4' discretization method
  • The resolution of an occupancy grid depending on differentiable parameters